From 38e2d5f8060f1d677b4e6def409420363c6c2509 Mon Sep 17 00:00:00 2001 From: Gabriel Mechali Date: Mon, 20 Jul 2026 11:06:23 -0400 Subject: [PATCH 1/2] Add preproc stage to ingestion workflow ingstion helper --- pipeline/workflow/ingestion-helper/clients/spanner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index b389e65c..b27e748b 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -38,6 +38,7 @@ class IngestionState(str, Enum): class IngestionStage(str, Enum): DATAFLOW = "dataflow" POSTPROCESSING = "postprocessing" + PREPROCESSING = "preprocessing" class SpannerClient: @@ -323,7 +324,7 @@ def update_ingestion_history(self, Args: workflow_id: The ID of the workflow. status: The status of the ingestion stage (PENDING, RUNNING, SUCCESS, FAILURE). - stage: The stage of the ingestion (dataflow, postprocessing, etc.). + stage: The stage of the ingestion (preprocessing, dataflow, postprocessing, etc.). job_id: The Dataflow job ID. ingested_imports: List of ingested import names. metrics: A dictionary containing metrics about the ingestion. From 651eddbd30e53331e97acf27e5cee7db5faf7da8 Mon Sep 17 00:00:00 2001 From: Gabriel Mechali Date: Thu, 23 Jul 2026 11:06:46 -0400 Subject: [PATCH 2/2] checking in spanner hange for preproc --- .../clients/fakeworkflow.yaml | 310 ++++++++++++++++++ .../ingestion-helper/clients/spanner.py | 31 +- .../ingestion-helper/clients/spanner_test.py | 26 ++ 3 files changed, 358 insertions(+), 9 deletions(-) create mode 100644 pipeline/workflow/ingestion-helper/clients/fakeworkflow.yaml diff --git a/pipeline/workflow/ingestion-helper/clients/fakeworkflow.yaml b/pipeline/workflow/ingestion-helper/clients/fakeworkflow.yaml new file mode 100644 index 00000000..63d00be2 --- /dev/null +++ b/pipeline/workflow/ingestion-helper/clients/fakeworkflow.yaml @@ -0,0 +1,310 @@ + main: + params: [input] + steps: + - init: + assign: + - project_id: 'datcom-website-dev' + - workflow_id: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + - version: '${"version-" + string(int(sys.now()))}' + - bucket_name: '${text.split(input.tempLocation, "/")[2]}' + - execution_error: null + - current_stage: "dataflow" + - lock_timeout: 82800 + - run_embeddings: true + - run_postproc: true + - postprocessing_result: null + - embedding_result: null + - decoded_imports: '${json.decode(input.importList)}' + - num_imports: '${len(decoded_imports)}' + - combined_import_name: "" + - imports_status_list: [] + - imports_history_list: [] + - imports_version_list: [] + - launch_params: + projectId: '${project_id}' + spannerInstanceId: '${input.spannerInstanceId}' + spannerDatabaseId: '${input.spannerDatabaseId}' + importList: '${input.importList}' + tempLocation: '${input.tempLocation}' + stagingLocation: '${input.tempLocation}' + forceCombineNodes: 'true' + isBaseDc: 'false' + - build_lists: + for: + value: imp + in: ${decoded_imports} + steps: + - append_lists: + assign: + - item_gcs_path: '${"gs://" + bucket_name + "/ingestion/internal/metadata/" + imp.importName + "/" + version}' + - status_item: {} + - status_item.importName: '${imp.importName}' + - status_item.latestVersion: '${item_gcs_path}' + - status_item.graphPath: '${imp.graphPath}' + - status_item.status: "STAGING" + - history_item: {} + - history_item.importName: '${imp.importName}' + - history_item.latestVersion: '${version}' + - imports_status_list: '${list.concat(imports_status_list, status_item)}' + - imports_history_list: '${list.concat(imports_history_list, history_item)}' + - imports_version_list: '${list.concat(imports_version_list, imp.importName)}' + - clean_item: '${text.replace_all(text.replace_all(text.to_lower(imp.importName), "/", "-"), "_", "-")}' + - join_prefix: '${if(combined_import_name == "", "", "-")}' + - combined_import_name: '${combined_import_name + join_prefix + clean_item}' + - set_dataflow_job_name: + assign: + - import_name_len: '${len(combined_import_name)}' + - substring_end: '${if(import_name_len < 35, import_name_len, 35)}' + - sanitized_short_import: '${text.substring(combined_import_name, 0, substring_end)}' + - dataflow_job_name: '${"gmechali-june22-" + sanitized_short_import + "-" + string(int(sys.now()))}' + - acquire_lock: + try: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/database/lock/acquire' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + timeout: '${lock_timeout}' + result: lock_result + retry: + predicate: '${http.default_retry_predicate}' + max_retries: 60 # Approx 5 hours + backoff: + initial_delay: 60 + max_delay: 300 # Max 5 minutes retry interval + multiplier: 2 + - update_history_pending: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'PENDING' + stage: 'dataflow' + importList: '${imports_history_list}' + result: history_pending_result + - process_ingestion: + try: + steps: + - set_import_staging: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/status' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + imports: '${imports_status_list}' + status: "STAGING" + result: staging_result + - run_flex_template: + call: googleapis.dataflow.v1b3.projects.locations.flexTemplates.launch + args: + projectId: '${project_id}' + location: '${input.region}' + body: + launchParameter: + jobName: '${dataflow_job_name}' + containerSpecGcsPath: 'gs://datcom-templates/templates/flex/ingestion-stable.json' + parameters: '${launch_params}' + environment: + serviceAccountEmail: 'gmechali-june22-dc-ing-df-sa@datcom-website-dev.iam.gserviceaccount.com' + tempLocation: '${input.tempLocation}' + stagingLocation: '${input.tempLocation}' + result: launch_result + - get_job_id: + assign: + - job_id: '${launch_result.job.id}' + - update_history_running: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'RUNNING' + stage: 'dataflow' + jobId: '${job_id}' + result: history_running_result + - poll_job: + steps: + - get_status: + call: googleapis.dataflow.v1b3.projects.locations.jobs.get + args: + projectId: '${project_id}' + location: '${input.region}' + jobId: '${job_id}' + result: job_status + - check_terminal: + switch: + - condition: '${job_status.currentState in ["JOB_STATE_DONE", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED", "JOB_STATE_UPDATED", "JOB_STATE_DRAINED"]}' + next: check_success + - wait_and_retry: + call: sys.sleep + args: + seconds: 60 + next: poll_job + - check_success: + switch: + - condition: '${job_status.currentState == "JOB_STATE_DONE"}' + next: update_history_dataflow_success + - update_history_dataflow_success: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'SUCCESS' + stage: 'dataflow' + jobId: '${job_id}' + importList: '${imports_history_list}' + result: history_dataflow_success_result + next: ingestion_postprocessing + - fail_on_job_status: + raise: '${ "Dataflow job failed with state: " + job_status.currentState }' + - ingestion_postprocessing: + steps: + - set_stage_postproc: + assign: + - current_stage: "postprocessing" + - update_history_postproc_running: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'RUNNING' + stage: 'postprocessing' + result: history_postproc_running_result + - run_postprocessing_parallel: + parallel: + shared: [postprocessing_result, embedding_result] + branches: + - postprocessing_branch: + steps: + - check_postproc_flag: + switch: + - condition: '${not run_postproc}' + next: end_postproc_branch + - run_postprocessings: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/aggregation/run' + auth: + type: OIDC + body: + importList: '${json.decode(input.importList)}' + result: postprocessing_result + - end_postproc_branch: + assign: + - dummy: true + + - embeddings_branch: + steps: + - check_embeddings_flag: + switch: + - condition: '${not run_embeddings}' + next: end_embeddings_branch + - run_embeddings: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/embeddings/ingest' + auth: + type: OIDC + body: + enableEmbeddings: true + timeout: 1800 + result: embedding_result + - end_embeddings_branch: + assign: + - dummy: true + - promote_version: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/version' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + imports: '${imports_version_list}' + version: '${version}' + comment: '${"Auto-promoted by workflow " + workflow_id}' + result: promote_result + - update_ingestion_history_step: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-status' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + jobId: '${job_id}' + status: 'SUCCESS' + importList: '${imports_history_list}' + result: history_result + - update_history_postproc_success: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'SUCCESS' + stage: null + result: history_postproc_success_result + except: + as: e + steps: + - update_history_failure: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'FAILURE' + stage: '${current_stage}' + jobId: '${default(job_id, "N/A")}' + importList: '${imports_history_list}' + result: history_failure_result + - capture_error: + assign: + - execution_error: '${e}' + - release_lock_step: + call: http.post + args: + url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/database/lock/release' + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + result: release_lock_result + - fail_workflow: + switch: + - condition: '${execution_error != null}' + raise: '${execution_error}' + + - restart_service: + call: googleapis.run.v2.projects.locations.services.patch + args: + name: "projects/datcom-website-dev/locations/us-central1/services/gmechali-june22-dc-datacommons-service" + updateMask: "template.labels" + body: + template: + labels: + restarted-at: '${string(int(sys.now()))}' + + + - return_result: + return: '${launch_result}' diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index b27e748b..d4957041 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -336,26 +336,39 @@ def update_ingestion_history(self, def _update(transaction: Transaction): status_str = status.value if hasattr(status, 'value') else status + stage_str = stage.value if stage and hasattr(stage, 'value') else stage + columns = ["WorkflowExecutionID", "Status"] values = [workflow_id, status_str] - if stage: - stage_str = stage.value if hasattr(stage, 'value') else stage + if stage_str: columns.append("Stage") values.append(stage_str) - if status == IngestionState.PENDING: - columns.append("CreationTimestamp") - values.append(spanner.COMMIT_TIMESTAMP) + is_base_dc = os.environ.get('IS_BASE_DC', 'true').lower() == 'true' + if not is_base_dc: + if stage == IngestionStage.PREPROCESSING: + columns.append("CreationTimestamp") + values.append(spanner.COMMIT_TIMESTAMP) + else: + if status == IngestionState.PENDING: + columns.append("CreationTimestamp") + values.append(spanner.COMMIT_TIMESTAMP) # The statements below allow us to construct a partial update, only for the fields that are set. - if status in (IngestionState.SUCCESS, IngestionState.FAILURE, - IngestionState.RETRY): + terminal_statuses = { + IngestionState.SUCCESS.value, + IngestionState.FAILURE.value, + IngestionState.RETRY.value, + } + if status_str in terminal_statuses: columns.append("CompletionTimestamp") values.append(spanner.COMMIT_TIMESTAMP) columns.append("IngestionFailure") - values.append(status in (IngestionState.FAILURE, - IngestionState.RETRY)) + values.append(status_str in { + IngestionState.FAILURE.value, + IngestionState.RETRY.value, + }) if job_id: columns.append("DataflowJobID") diff --git a/pipeline/workflow/ingestion-helper/clients/spanner_test.py b/pipeline/workflow/ingestion-helper/clients/spanner_test.py index df7ed6f8..31d89231 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner_test.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner_test.py @@ -453,6 +453,32 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertEqual(kwargs['columns'], ["WorkflowExecutionID", "Status", "Stage", "CreationTimestamp", "IngestedImports"]) self.assertEqual(kwargs['values'], [["wf-123", "PENDING", "dataflow", spanner.COMMIT_TIMESTAMP, ["import1", "import2"]]]) + @patch('google.cloud.spanner.Client') + def test_update_ingestion_history_preprocessing(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") + client.update_ingestion_history( + workflow_id="wf-123", + status=IngestionState.RUNNING, + stage=IngestionStage.PREPROCESSING, + ingested_imports=["import1"] + ) + + mock_transaction.insert_or_update.assert_called_once() + _, kwargs = mock_transaction.insert_or_update.call_args + self.assertEqual(kwargs['table'], 'IngestionHistory') + self.assertEqual(kwargs['columns'], ["WorkflowExecutionID", "Status", "Stage", "CreationTimestamp", "IngestedImports"]) + self.assertEqual(kwargs['values'], [["wf-123", "RUNNING", "preprocessing", spanner.COMMIT_TIMESTAMP, ["import1"]]]) + @patch('google.cloud.spanner.Client') def test_update_ingestion_history_running(self, mock_spanner_client): mock_instance = MagicMock()