[DCP Preprocessor] Move the Preprocessor to inside the workflow - Part 2

infra/dcp/main.tf

@@ -138,9 +138,9 @@ locals {
     helper_service_image              = var.ingestion_helper_service_image
 
     # Dataflow Network Configuration
-    dataflow_ip_configuration     = var.ingestion_dataflow_ip_configuration
-    dataflow_subnetwork           = var.ingestion_dataflow_subnetwork
-    dataflow_template_gcs_path    = var.ingestion_dataflow_template_gcs_path
+    dataflow_ip_configuration  = var.ingestion_dataflow_ip_configuration
+    dataflow_subnetwork        = var.ingestion_dataflow_subnetwork
+    dataflow_template_gcs_path = var.ingestion_dataflow_template_gcs_path
   }
 }
 

infra/dcp/modules/datacommons_services/variables.tf

@@ -71,7 +71,7 @@ variable "mcp_instructions_path" {
 }
 
 variable "resolve_with_spanner_embeddings" {
-  type    = bool
+  type = bool
 }
 
 # =============================================================================

infra/dcp/modules/ingestion/preprocessing_job/main.tf

@@ -114,11 +114,3 @@ resource "google_secret_manager_secret_iam_member" "preprocessing_maps_key_acces
   role      = "roles/secretmanager.secretAccessor"
   member    = "serviceAccount:${google_service_account.preprocessing_sa.email}"
 }
-
-resource "google_project_iam_member" "preprocessing_workflow_invoker" {
-  project = var.project_id
-  role    = "roles/workflows.invoker"
-  member  = "serviceAccount:${google_service_account.preprocessing_sa.email}"
-}
-
-

infra/dcp/modules/ingestion/workflow/main.tf

@@ -32,6 +32,7 @@ resource "google_workflows_workflow" "ingestion_orchestrator" {
     embeddings_timeout             = var.embeddings_timeout
     clean_namespace_prefix         = local.clean_namespace_prefix
     enable_redis_cache_clearing    = var.enable_redis_cache_clearing
+    preprocessing_job_name         = var.preprocessing_job_name
   })
 }
 

infra/dcp/modules/ingestion/workflow/variables.tf

@@ -59,6 +59,12 @@ variable "enable_redis_cache_clearing" {
   default     = false
 }
 
+variable "preprocessing_job_name" {
+  type        = string
+  description = "Name of the ingestion preprocessing Cloud Run job"
+  default     = ""
+}
+
 variable "ingestion_artifacts_path" {
   type        = string
   description = "Path where pre-processed files are placed for the next stage"

infra/dcp/modules/ingestion/workflow/workflow.yaml

@@ -7,6 +7,7 @@ main:
           - workflow_id: '$${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}'
           - version: '$${"version-" + string(int(sys.now()))}'
           - bucket_name: '$${text.split(input.tempLocation, "/")[2]}'
+          - temp_path_prefix: '$${text.replace_all(input.tempLocation, "gs://" + bucket_name + "/", "")}'
           - execution_error: null
           - postprocessing_result: null
           - embedding_result: null
@@ -17,21 +18,50 @@ main:
           - max_lock_retries: $${int(lock_timeout / 120)}
           - run_embeddings: ${enable_embeddings_generation}
           - run_postproc: ${enable_bigquery_postprocessing}
-          - decoded_imports: '$${json.decode(input.importList)}'
-          - num_imports: '$${len(decoded_imports)}'
+          - decoded_imports: []
+          - num_imports: 0
           - combined_import_name: ""
+          - prep_job_name: '${preprocessing_job_name}'
+          - comma_separated_imports: ""
           - imports_status_list: []
           - imports_history_list: []
           - imports_version_list: []
           - launch_params:
               projectId: '$${project_id}'
-              spannerInstanceId: '$${input.spannerInstanceId}'
-              spannerDatabaseId: '$${input.spannerDatabaseId}'
-              importList: '$${input.importList}'
+              spannerInstanceId: '$${default(map.get(input, "spannerInstanceId"), "")}'
+              spannerDatabaseId: '$${default(map.get(input, "spannerDatabaseId"), "")}'
+              importList: '$${default(map.get(input, "importList"), "")}'
               tempLocation: '$${input.tempLocation}'
               stagingLocation: '$${input.tempLocation}'
               forceCombineNodes: 'true'
               isBaseDc: 'false'
+    - normalize_imports:
+        switch:
+          - condition: '$${default(map.get(input, "importList"), "") != ""}'
+            assign:
+              - decoded_imports: '$${json.decode(input.importList)}'
+          - condition: '$${len(default(map.get(input, "imports"), [])) > 0}'
+            steps:
+              - init_decoded:
+                  assign:
+                    - decoded_imports: []
+              - build_decoded_loop:
+                  for:
+                    value: imp_str
+                    in: '$${input.imports}'
+                    steps:
+                      - add_to_decoded:
+                          assign:
+                            - new_item:
+                                importName: '$${imp_str}'
+                                graphPath: ""
+                            - decoded_imports: '$${list.concat(decoded_imports, new_item)}'
+          - condition: true
+            assign:
+              - decoded_imports: []
+    - set_num_imports:
+        assign:
+          - num_imports: '$${len(decoded_imports)}'
     - build_lists:
         for:
           value: imp
@@ -48,13 +78,41 @@ main:
                   - 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}'
+                  - comma_separated_imports: '$${if(comma_separated_imports == "", imp.importName, comma_separated_imports + "," + imp.importName)}'
     - 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: '$${"${clean_namespace_prefix}" + sanitized_short_import + "-" + string(int(sys.now()))}'
 
+    - run_preprocessing:
+        call: launch_and_poll_preprocessing_job
+        args:
+          job_name: '$${prep_job_name}'
+          workflow_id: '$${workflow_id}'
+          temp_location: '$${input.tempLocation}'
+          imports_arg: '$${comma_separated_imports}'
+          region: '$${input.region}'
+          project_id: '${project_id}'
+        result: prep_res
+
+    - read_handshake_file:
+        call: googleapis.storage.v1.objects.get
+        args:
+          bucket: '$${bucket_name}'
+          object: '$${text.url_encode(temp_path_prefix + "/datacommons/ingestion_records/" + workflow_id + ".json")}'
+          alt: 'media'
+        result: handshake_data
+
+    - update_launch_params:
+        assign:
+          - launch_params.spannerInstanceId: '$${handshake_data.spannerInstanceId}'
+          - launch_params.spannerDatabaseId: '$${handshake_data.spannerDatabaseId}'
+          - launch_params.importList: '$${handshake_data.importList}'
+          - launch_params.tempLocation: '$${handshake_data.tempLocation}'
+          - launch_params.stagingLocation: '$${handshake_data.tempLocation}'
+
     - try_acquire_lock:
         try:
           call: http.post
@@ -108,7 +166,7 @@ main:
                     workflowId: '$${workflow_id}'
                     imports: '$${imports_status_list}'
                     status: 'STAGING'
-            
+
             - launch_dataflow:
                 call: launch_and_poll_dataflow
                 args:
@@ -398,3 +456,55 @@ launch_and_poll_dataflow:
         return:
           job_id: $${job_id}
           launch_result: $${launch_result}
+
+launch_and_poll_preprocessing_job:
+  params: [job_name, workflow_id, temp_location, imports_arg, region, project_id]
+  steps:
+    - run_cloud_run_job:
+        call: googleapis.run.v2.projects.locations.jobs.run
+        args:
+          name: '$${"projects/" + project_id + "/locations/" + region + "/jobs/" + job_name}'
+          body:
+            overrides:
+              containerOverrides:
+                - env:
+                    - name: 'WORKFLOW_EXECUTION_ID'
+                      value: '$${workflow_id}'
+                    - name: 'TEMP_LOCATION'
+                      value: '$${temp_location}'
+                    - name: 'DATA_RUN_MODE'
+                      value: 'dcpbridge'
+                  args: ['$${"--imports=" + imports_arg}']
+        result: run_res
+
+    - assign_execution_name:
+        assign:
+          - execution_name: '$${run_res.metadata.name}'
+
+    - poll_execution:
+        steps:
+          - get_execution:
+              call: googleapis.run.v2.projects.locations.jobs.executions.get
+              args:
+                name: '$${execution_name}'
+              result: exec_status
+          - check_exec_terminal:
+              switch:
+                - condition: '$${default(map.get(exec_status, "succeededCount"), 0) > 0 or default(map.get(exec_status, "failedCount"), 0) > 0 or default(map.get(exec_status, "cancelledCount"), 0) > 0}'
+                  next: check_exec_success
+          - wait_for_execution:
+              call: sys.sleep
+              args:
+                seconds: 30
+              next: poll_execution
+
+    - check_exec_success:
+        switch:
+          - condition: '$${default(map.get(exec_status, "succeededCount"), 0) > 0}'
+            next: return_success
+
+    - fail_on_exec_status:
+        raise: '$${"Preprocessing Cloud Run job failed or was cancelled: " + exec_status.name}'
+
+    - return_success:
+        return: '$${exec_status}'