Automatically Switch to upload_graph when the graph is large

apache · Sep 22, 2023 · ed10fa1 · ed10fa1
1 parent aa2bf58
commit ed10fa1
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/...le-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/...le-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
@@ -1331,6 +1331,14 @@ public DataflowPipelineJob run(Pipeline pipeline) {
       hooks.modifyEnvironmentBeforeSubmission(newJob.getEnvironment());
     }
 
+    // enable upload_graph when the graph is too large
+    if ((Utf8.encodedLength(newJob.toString()) >= CREATE_JOB_REQUEST_LIMIT_BYTES) &&
+        !hasExperiment(options, "upload_graph")) {
+      List<String> experiments = new ArrayList<>(options.getExperiments());
+      experiments.add("upload_graph");
+      options.setExperiments(ImmutableList.copyOf(experiments));
+    }
+
     // Upload the job to GCS and remove the graph object from the API call.  The graph
     // will be downloaded from GCS by the service.
     if (hasExperiment(options, "upload_graph")) {