From c722a33142d4972eef47f5b7679f6c3ca6478d4c Mon Sep 17 00:00:00 2001 From: Vishal Gupta Date: Fri, 17 Jul 2026 04:58:20 +0000 Subject: [PATCH 1/2] Add per-provenance counters to dataflow and spanner --- .../pipeline/GraphIngestionPipeline.java | 25 ++++++---- .../ingestion-helper/clients/spanner.py | 46 ++++++++++++++----- .../ingestion-helper/clients/spanner_test.py | 38 +++++++++++++++ .../ingestion-helper/utils/imports.py | 45 ++++++++++++++---- 4 files changed, 122 insertions(+), 32 deletions(-) diff --git a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java index 1efbe88b..ae1a9857 100644 --- a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java +++ b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java @@ -33,14 +33,6 @@ public class GraphIngestionPipeline { private static final Logger LOGGER = LoggerFactory.getLogger(GraphIngestionPipeline.class); private static final Counter nodeInvalidTypeCounter = Metrics.counter(GraphIngestionPipeline.class, "mcf_nodes_without_type"); - private static final Counter nodeCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_node_count"); - private static final Counter edgeCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_edge_count"); - private static final Counter obsCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_observation_count"); - private static final Counter timeSeriesCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_timeseries_count"); // List of imports that require node combination. private static final List IMPORTS_TO_COMBINE = List.of("Schema", "Place", "Provenance"); @@ -137,6 +129,13 @@ private static void processImport( boolean isBaseDc = options.getIsBaseDc(); String provenance = ProvenanceUtils.getProvenanceDcid(importName, isBaseDc); + Counter nodeCounter = Metrics.counter(GraphIngestionPipeline.class, "node_count:" + importName); + Counter edgeCounter = Metrics.counter(GraphIngestionPipeline.class, "edge_count:" + importName); + Counter obsCounter = + Metrics.counter(GraphIngestionPipeline.class, "observation_count:" + importName); + Counter timeSeriesCounter = + Metrics.counter(GraphIngestionPipeline.class, "timeseries_count:" + importName); + // 1. Prepare Deletes: // Generate mutations to delete existing data for this import/provenance. // Create a dummy signal if deletes are skipped, so downstream dependencies are satisfied @@ -162,7 +161,8 @@ private static void processImport( PipelineUtils.InputFormat format = PipelineUtils.resolveFormat(graphPath); if (format == PipelineUtils.InputFormat.TFRECORD) { TfRecordProcessingResult result = - processTfRecordImport(pipeline, importName, graphPath, isBaseDc); + processTfRecordImport( + pipeline, importName, graphPath, isBaseDc, obsCounter, timeSeriesCounter); writeToSpanner( pipeline, spannerClient, @@ -326,7 +326,12 @@ private static class TfRecordProcessingResult { } private static TfRecordProcessingResult processTfRecordImport( - Pipeline pipeline, String importName, String graphPath, boolean isBaseDc) { + Pipeline pipeline, + String importName, + String graphPath, + boolean isBaseDc, + Counter obsCounter, + Counter timeSeriesCounter) { PCollection optGraph = PipelineUtils.readOptimizedMcfGraph(importName, graphPath, pipeline); PCollection uniqueSeries = diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index 593b581d..a99af194 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -62,7 +62,8 @@ def __init__(self, embedding_table: str = "NodeEmbedding", embedding_index: str = "NodeEmbeddingIndex", embedding_label_index: str = "NodeEmbeddingLabelIndex", - emulator_host: str = None): + emulator_host: str = None, + is_base_dc: bool = True): """Initializes a Spanner client and connects to a specific database.""" client_options = {"api_endpoint": "spanner.googleapis.com"} credentials = None @@ -94,6 +95,7 @@ def __init__(self, self.embedding_table = embedding_table self.embedding_index = embedding_index self.embedding_label_index = embedding_label_index + self.is_base_dc = is_base_dc if not models: models = self._DEFAULT_MODELS @@ -401,6 +403,9 @@ def update_import_version_history(self, logging.info( f"Updating ImportVersionHistory table for workflow {workflow_id}") + m = metrics if metrics else {} + import_metrics = m.get('import_metrics', {}) + def _insert(transaction: Transaction): columns = [ "ImportName", "Version", "UpdateTimestamp", @@ -408,17 +413,29 @@ def _insert(transaction: Transaction): "NodeCount", "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] - m = metrics if metrics else {} version_history_values = [] for import_json in import_list_json: + import_name = import_json.get('importName') + if not import_name: + continue + + short_name = import_name.split(':')[-1] + counts = ( + import_metrics.get(short_name) or + import_metrics.get(import_name) or + (import_json if all(k in import_json and import_json[k] is not None for k in ('node_count', 'edge_count', 'obs_count', 'ts_count')) else {}) or + (import_json if all(k in import_json and import_json[k] is not None for k in ('nodeCount', 'edgeCount', 'obsCount', 'tsCount')) else {}) + ) + version_history_values.append([ - import_json['importName'], import_json['latestVersion'], + import_name, import_json.get('latestVersion'), spanner.COMMIT_TIMESTAMP, workflow_id, status, m.get('execution_time'), - m.get('node_count'), - m.get('edge_count'), - m.get('obs_count'), - m.get('ts_count'), "ingestion-workflow:" + workflow_id + counts.get('node_count') or counts.get('nodeCount'), + counts.get('edge_count') or counts.get('edgeCount'), + counts.get('obs_count') or counts.get('obsCount'), + counts.get('ts_count') or counts.get('tsCount'), + "ingestion-workflow:" + workflow_id ]) if version_history_values: @@ -512,6 +529,12 @@ def update_version_history(self, import_name = import_name.split(':')[-1] logging.info(f"Updating version history for {import_name} to {version}") + m = metrics if metrics else {} + node_count = m.get('node_count') + edge_count = m.get('edge_count') + obs_count = m.get('obs_count') + ts_count = m.get('ts_count') + def _record(transaction: Transaction): columns = [ "ImportName", "Version", "UpdateTimestamp", @@ -519,15 +542,14 @@ def _record(transaction: Transaction): "NodeCount", "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] - m = metrics if metrics else {} values = [[ import_name, version, spanner.COMMIT_TIMESTAMP, workflow_id, status, m.get('execution_time'), - m.get('node_count'), - m.get('edge_count'), - m.get('obs_count'), - m.get('ts_count'), comment + node_count, + edge_count, + obs_count, + ts_count, comment ]] transaction.insert(table="ImportVersionHistory", columns=columns, diff --git a/pipeline/workflow/ingestion-helper/clients/spanner_test.py b/pipeline/workflow/ingestion-helper/clients/spanner_test.py index bba47135..fb547d23 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner_test.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner_test.py @@ -702,7 +702,45 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): None, None, None, None, None, None, "ingestion-workflow:wf-789" ]]) + @patch('google.cloud.spanner.Client') + def test_update_import_version_history_per_import_counts(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_transaction = MagicMock() + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_transaction, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + client = SpannerClient("project", "instance", "database") + + import_list = [ + {"importName": "import1", "latestVersion": "v1.0"}, + {"importName": "import2", "latestVersion": "v2.0"} + ] + metrics = { + 'execution_time': 120, + 'import_metrics': { + "import1": {'node_count': 10, 'edge_count': 20, 'ts_count': 5, 'obs_count': 15}, + "import2": {'node_count': 30, 'edge_count': 40, 'ts_count': 25, 'obs_count': 35}, + } + } + client.update_import_version_history( + import_list, workflow_id="wf-123", status="SUCCESS", metrics=metrics + ) + + mock_transaction.insert.assert_called_once() + _, kwargs = mock_transaction.insert.call_args + self.assertEqual(kwargs['table'], 'ImportVersionHistory') + self.assertEqual(kwargs['values'], [ + ["import1", "v1.0", spanner.COMMIT_TIMESTAMP, "wf-123", "SUCCESS", 120, 10, 20, 15, 5, "ingestion-workflow:wf-123"], + ["import2", "v2.0", spanner.COMMIT_TIMESTAMP, "wf-123", "SUCCESS", 120, 30, 40, 35, 25, "ingestion-workflow:wf-123"] + ]) + if __name__ == '__main__': unittest.main() + diff --git a/pipeline/workflow/ingestion-helper/utils/imports.py b/pipeline/workflow/ingestion-helper/utils/imports.py index caed3ec8..de729119 100644 --- a/pipeline/workflow/ingestion-helper/utils/imports.py +++ b/pipeline/workflow/ingestion-helper/utils/imports.py @@ -127,7 +127,8 @@ def get_ingestion_metrics(project_id, location, job_id): job_id: The Dataflow job ID. Returns: - A dictionary containing 'obs_count', 'node_count', 'edge_count', 'ts_count', and 'execution_time'. + A dictionary containing 'obs_count', 'node_count', 'edge_count', 'ts_count', 'execution_time', + and 'import_metrics'. """ dataflow = build('dataflow', 'v1b3', cache_discovery=False) # Fetch Dataflow metrics @@ -136,6 +137,7 @@ def get_ingestion_metrics(project_id, location, job_id): obs_count = 0 ts_count = 0 execution_time = 0 + import_metrics = {} if project_id and job_id: try: # Fetch Job details for execution time @@ -159,14 +161,36 @@ def get_ingestion_metrics(project_id, location, job_id): jobId=job_id).execute() for metric in metrics.get('metrics', []): name = metric['name']['name'] - if name == 'graph_node_count': - node_count += int(metric['scalar']) - elif name == 'graph_edge_count': - edge_count += int(metric['scalar']) - elif name == 'graph_observation_count': - obs_count += int(metric['scalar']) - elif name == 'graph_timeseries_count': - ts_count += int(metric['scalar']) + scalar = int(metric.get('scalar', 0)) + if ':' in name: + metric_type, imp_name = name.split(':', 1) + else: + metric_type, imp_name = name, None + + if imp_name and imp_name not in import_metrics: + import_metrics[imp_name] = { + 'node_count': 0, + 'edge_count': 0, + 'obs_count': 0, + 'ts_count': 0, + } + + if metric_type == 'node_count': + node_count += scalar + if imp_name: + import_metrics[imp_name]['node_count'] = scalar + elif metric_type == 'edge_count': + edge_count += scalar + if imp_name: + import_metrics[imp_name]['edge_count'] = scalar + elif metric_type == 'observation_count': + obs_count += scalar + if imp_name: + import_metrics[imp_name]['obs_count'] = scalar + elif metric_type == 'timeseries_count': + ts_count += scalar + if imp_name: + import_metrics[imp_name]['ts_count'] = scalar except HttpError as e: logging.error( f"Error fetching dataflow metrics for job {job_id}: {e}") @@ -175,5 +199,6 @@ def get_ingestion_metrics(project_id, location, job_id): 'node_count': node_count, 'edge_count': edge_count, 'ts_count': ts_count, - 'execution_time': execution_time + 'execution_time': execution_time, + 'import_metrics': import_metrics, } From d7af80f70780eb4cfc49938b386838d4962d9186 Mon Sep 17 00:00:00 2001 From: Vishal Gupta Date: Tue, 21 Jul 2026 07:13:41 +0000 Subject: [PATCH 2/2] Address reviewer comments --- .../ingestion/pipeline/GraphIngestionPipeline.java | 13 +++++++++---- .../workflow/ingestion-helper/clients/spanner.py | 11 ++--------- .../workflow/ingestion-helper/routes/database.py | 2 +- pipeline/workflow/ingestion-helper/utils/imports.py | 7 +++---- 4 files changed, 15 insertions(+), 18 deletions(-) diff --git a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java index ae1a9857..6c74942a 100644 --- a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java +++ b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java @@ -129,12 +129,17 @@ private static void processImport( boolean isBaseDc = options.getIsBaseDc(); String provenance = ProvenanceUtils.getProvenanceDcid(importName, isBaseDc); - Counter nodeCounter = Metrics.counter(GraphIngestionPipeline.class, "node_count:" + importName); - Counter edgeCounter = Metrics.counter(GraphIngestionPipeline.class, "edge_count:" + importName); + String shortName = + importName.contains(":") + ? importName.substring(importName.lastIndexOf(':') + 1) + : importName; + + Counter nodeCounter = Metrics.counter(GraphIngestionPipeline.class, "node_count:" + shortName); + Counter edgeCounter = Metrics.counter(GraphIngestionPipeline.class, "edge_count:" + shortName); Counter obsCounter = - Metrics.counter(GraphIngestionPipeline.class, "observation_count:" + importName); + Metrics.counter(GraphIngestionPipeline.class, "observation_count:" + shortName); Counter timeSeriesCounter = - Metrics.counter(GraphIngestionPipeline.class, "timeseries_count:" + importName); + Metrics.counter(GraphIngestionPipeline.class, "timeseries_count:" + shortName); // 1. Prepare Deletes: // Generate mutations to delete existing data for this import/provenance. diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index a99af194..5c1a874c 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -62,8 +62,7 @@ def __init__(self, embedding_table: str = "NodeEmbedding", embedding_index: str = "NodeEmbeddingIndex", embedding_label_index: str = "NodeEmbeddingLabelIndex", - emulator_host: str = None, - is_base_dc: bool = True): + emulator_host: str = None): """Initializes a Spanner client and connects to a specific database.""" client_options = {"api_endpoint": "spanner.googleapis.com"} credentials = None @@ -95,7 +94,6 @@ def __init__(self, self.embedding_table = embedding_table self.embedding_index = embedding_index self.embedding_label_index = embedding_label_index - self.is_base_dc = is_base_dc if not models: models = self._DEFAULT_MODELS @@ -420,12 +418,7 @@ def _insert(transaction: Transaction): continue short_name = import_name.split(':')[-1] - counts = ( - import_metrics.get(short_name) or - import_metrics.get(import_name) or - (import_json if all(k in import_json and import_json[k] is not None for k in ('node_count', 'edge_count', 'obs_count', 'ts_count')) else {}) or - (import_json if all(k in import_json and import_json[k] is not None for k in ('nodeCount', 'edgeCount', 'obsCount', 'tsCount')) else {}) - ) + counts = import_metrics.get(short_name) or import_json version_history_values.append([ import_name, import_json.get('latestVersion'), diff --git a/pipeline/workflow/ingestion-helper/routes/database.py b/pipeline/workflow/ingestion-helper/routes/database.py index d5a89f12..da439462 100644 --- a/pipeline/workflow/ingestion-helper/routes/database.py +++ b/pipeline/workflow/ingestion-helper/routes/database.py @@ -58,7 +58,7 @@ def acquire_ingestion_lock(req: LockAcquireRequest, spanner: SpannerClient = Dep status_ok = spanner.acquire_lock(req.workflowId, req.timeout) if not status_ok: raise HTTPException( - status_code=400, + status_code=503, detail=f"Failed to acquire lock: Lock already held or acquisition timed out for workflow {req.workflowId}" ) return BaseResponse(status=ResponseStatus.OK) diff --git a/pipeline/workflow/ingestion-helper/utils/imports.py b/pipeline/workflow/ingestion-helper/utils/imports.py index de729119..f9364847 100644 --- a/pipeline/workflow/ingestion-helper/utils/imports.py +++ b/pipeline/workflow/ingestion-helper/utils/imports.py @@ -162,10 +162,9 @@ def get_ingestion_metrics(project_id, location, job_id): for metric in metrics.get('metrics', []): name = metric['name']['name'] scalar = int(metric.get('scalar', 0)) - if ':' in name: - metric_type, imp_name = name.split(':', 1) - else: - metric_type, imp_name = name, None + match = re.match(r"^([^:]+)(?::(.+))?$", name) + metric_type = match.group(1) if match else name + imp_name = match.group(2).split(':')[-1] if match and match.group(2) else None if imp_name and imp_name not in import_metrics: import_metrics[imp_name] = {