Update uuid to be unique for a specific execution

xlml/apis/task.py

@@ -152,7 +152,9 @@ def post_process(self) -> DAGNode:
       A DAG node that executes the post process.
     """
     with TaskGroup(group_id="post_process") as group:
-      process_id = metric.generate_process_id.override(retries=0)()
+      process_id = metric.generate_process_id.override(retries=0)(
+          self.task_test_config.benchmark_id
+      )
       metric.process_metrics.override(retries=0)(
           process_id,
           self.task_test_config,
@@ -243,7 +245,9 @@ def post_process(self) -> DAGNode:
       A DAG node that executes the post process.
     """
     with TaskGroup(group_id="post_process") as group:
-      process_id = metric.generate_process_id.override(retries=0)()
+      process_id = metric.generate_process_id.override(retries=0)(
+          self.task_test_config.benchmark_id
+      )
       metric.process_metrics.override(retries=0)(
           process_id,
           self.task_test_config,
@@ -360,7 +364,9 @@ def post_process(self) -> DAGNode:
       A DAG node that executes the post process.
     """
     with TaskGroup(group_id="post_process") as group:
-      process_id = metric.generate_process_id.override(retries=0)()
+      process_id = metric.generate_process_id.override(retries=0)(
+          self.task_test_config.benchmark_id
+      )
       metric.process_metrics.override(retries=0)(
           process_id,
           self.task_test_config,

xlml/utils/bigquery.py

@@ -26,9 +26,9 @@
 from google.cloud import bigquery
 
 
-BENCHMARK_BQ_JOB_TABLE_NAME = "job_history"
-BENCHMARK_BQ_METRIC_TABLE_NAME = "metric_history"
-BENCHMARK_BQ_METADATA_TABLE_NAME = "metadata_history"
+BQ_JOB_TABLE_NAME = "job_history"
+BQ_METRIC_TABLE_NAME = "metric_history"
+BQ_METADATA_TABLE_NAME = "metadata_history"
 
 
 @dataclasses.dataclass
@@ -82,9 +82,7 @@ def __init__(
   ):
     self.project = google.auth.default()[1] if project is None else project
     self.database = (
-        metric_config.DatasetOption.BENCHMARK_DATASET.value
-        if database is None
-        else database
+        metric_config.DatasetOption.XLML_DATASET.value if database is None else database
     )
     self.client = bigquery.Client(
         project=project,
@@ -95,15 +93,15 @@ def __init__(
 
   @property
   def job_history_table_id(self):
-    return ".".join((self.project, self.database, BENCHMARK_BQ_JOB_TABLE_NAME))
+    return ".".join((self.project, self.database, BQ_JOB_TABLE_NAME))
 
   @property
   def metric_history_table_id(self):
-    return ".".join((self.project, self.database, BENCHMARK_BQ_METRIC_TABLE_NAME))
+    return ".".join((self.project, self.database, BQ_METRIC_TABLE_NAME))
 
   @property
   def metadata_history_table_id(self):
-    return ".".join((self.project, self.database, BENCHMARK_BQ_METADATA_TABLE_NAME))
+    return ".".join((self.project, self.database, BQ_METADATA_TABLE_NAME))
 
   def is_valid_metric(self, value: float):
     """Check if float metric is valid for BigQuery table."""

xlml/utils/bigquery_test.py

@@ -66,14 +66,14 @@ def test_is_valid_metric(self, x: float, expected_value: bool):
   @mock.patch.object(google.auth, "default", return_value=["mock", "mock_project"])
   @mock.patch.object(bigquery.Client, "get_table", return_value="mock_table")
   @mock.patch.object(bigquery.Client, "insert_rows", return_value=["there is an error"])
-  def test_insert_failure(self, default, get_table, insert_rows):
+  def test_insert_failure(self, insert_rows, get_table, default):
     bq_metric = test_bigquery.BigQueryMetricClient()
     self.assertRaises(RuntimeError, bq_metric.insert, self.test_runs)
 
   @mock.patch.object(google.auth, "default", return_value=["mock", "mock_project"])
   @mock.patch.object(bigquery.Client, "get_table", return_value="mock_table")
   @mock.patch.object(bigquery.Client, "insert_rows", return_value=[])
-  def test_insert_success(self, default, get_table, insert_rows):
+  def test_insert_success(self, insert_rows, get_table, default):
     bq_metric = test_bigquery.BigQueryMetricClient()
     bq_metric.insert(self.test_runs)
 

xlml/utils/metric.py

@@ -346,13 +346,18 @@ def generate_row_uuid(base_id: str, index: int) -> str:
 
 
 @task(trigger_rule="all_done")
-def generate_process_id() -> str:
-  """Generate a process id that will be a base id for uuid of test runs.
+def generate_process_id(benchmark_id: str) -> str:
+  """Generate a process id that will be a base id for single/multiple test run(s).
+
+  Args:
+    benchmark_id: The unique key for metrics generated by the test.
 
   Returns:
-    A random uuid.
+    An id based on benchmark_id and Airflow run_id.
   """
-  return str(uuid.uuid4())
+  context = get_current_context()
+  id = str(f"{benchmark_id}_{context['run_id']}")
+  return hashlib.sha256(id.encode("utf-8")).hexdigest()
 
 
 def is_valid_entry() -> bool: