Add an option to set the Dataflow container base repository (Python, Java, Go)

CHANGES.md

@@ -82,6 +82,7 @@
 * (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)).
 * (Python) Added `AwsSecret` (`AwsSecretsManager` / `type:awssecret`) to retrieve secrets from AWS Secrets Manager, configured from `AwsOptions` (`--aws_*`) with optional per-secret connection overrides ([#40528](https://github.com/apache/beam/issues/40528)).
 * (Go) Added `schema.RegisterLogicalTypeConversion` and `schema.RegisterLogicalTypeCoder` to register portable logical types in one call, and Go types for the standard `beam:logical_type:date:v1` and `beam:logical_type:micros_instant:v1` logical types, so schemas from other SDKs with these types can be used in cross-language pipelines ([#39684](https://github.com/apache/beam/issues/39684)).
+* (Python, Java, Go) Added `--dataflow_container_base_repository` (Python, Go) and `--dataflowContainerBaseRepository` (Java) to set the repository (e.g. `docker.example.com/my-prefix/cloud-dataflow/v1beta3`) from which the Dataflow runner pulls its default SDK container images, including the default images of cross-language Apache Beam environments. Defaults to `gcr.io/cloud-dataflow/v1beta3` ([#40569](https://github.com/apache/beam/issues/40569)).
 
 ## Breaking Changes
 

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -169,6 +169,7 @@
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.CharMatcher;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
@@ -1041,9 +1042,7 @@ protected RunnerApi.Pipeline applySdkEnvironmentOverrides(
           String imageAndTag =
               normalizeDataflowImageAndTag(
                   containerImage.substring(containerImage.lastIndexOf("/")));
-          containerImage =
-              DataflowRunnerInfo.getDataflowRunnerInfo().getContainerImageBaseRepository()
-                  + imageAndTag;
+          containerImage = getContainerImageBaseRepository(options) + imageAndTag;
         }
         environmentBuilder.setPayload(
             RunnerApi.DockerPayload.newBuilder()
@@ -2722,14 +2721,30 @@ static String getContainerImageForJob(DataflowPipelineOptions options) {
 
   /** Construct the default Dataflow container full image URL. */
   static String getDefaultContainerImageUrl(DataflowPipelineOptions options) {
-    DataflowRunnerInfo dataflowRunnerInfo = DataflowRunnerInfo.getDataflowRunnerInfo();
     return String.format(
         "%s/%s:%s",
-        dataflowRunnerInfo.getContainerImageBaseRepository(),
+        getContainerImageBaseRepository(options),
         getDefaultContainerImageNameForJob(options),
         getDefaultContainerVersion(options));
   }
 
+  /**
+   * Returns the base repository of the default Dataflow container images, without a trailing slash.
+   * Uses {@code --dataflowContainerBaseRepository} if set, and the repository configured in {@link
+   * DataflowRunnerInfo} otherwise.
+   */
+  @VisibleForTesting
+  static String getContainerImageBaseRepository(DataflowPipelineOptions options) {
+    String baseRepository = options.getDataflowContainerBaseRepository();
+    if (baseRepository != null) {
+      String trimmed = CharMatcher.is('/').trimTrailingFrom(baseRepository.trim());
+      if (!trimmed.isEmpty()) {
+        return trimmed;
+      }
+    }
+    return DataflowRunnerInfo.getDataflowRunnerInfo().getContainerImageBaseRepository();
+  }
+
   /**
    * Construct the default Dataflow container image name based on pipeline type and Java version.
    */

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java

@@ -136,6 +136,23 @@ public String getAlgorithm() {
 
   void setSdkContainerImage(String value);
 
+  /**
+   * Base repository (registry host and path prefix) of the default Dataflow container images, e.g.
+   * {@code docker.example.com/my-prefix/cloud-dataflow/v1beta3}. If unset, the default repository
+   * of the Dataflow runner is used.
+   */
+  @Description(
+      "Base repository (registry host and path prefix) from which the Dataflow runner pulls its "
+          + "default SDK container images, e.g. "
+          + "docker.example.com/my-prefix/cloud-dataflow/v1beta3. The runner appends the image "
+          + "name and tag, e.g. <base>/beam_java17_sdk:<version>. This is also used for the "
+          + "default images of cross-language Apache Beam environments. Ignored for the Java SDK "
+          + "image if sdkContainerImage is set. If not set, the Google Cloud Dataflow default "
+          + "(gcr.io/cloud-dataflow/v1beta3) is used.")
+  @Nullable String getDataflowContainerBaseRepository();
+
+  void setDataflowContainerBaseRepository(@Nullable String value);
+
   /**
    * GCE <a href="https://cloud.google.com/compute/docs/networking">network</a> for launching
    * workers.

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java

@@ -34,6 +34,7 @@
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.matchesRegex;
 import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.startsWith;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -1311,6 +1312,20 @@ public void testApplySdkEnvironmentOverridesRcByDefault() throws IOException {
     assertThat(runner.applySdkEnvironmentOverrides(pipeline, options), equalTo(expectedPipeline));
   }
 
+  @Test
+  public void testApplySdkEnvironmentOverridesWithCustomBaseRepository() throws IOException {
+    DataflowPipelineOptions options = buildPipelineOptions();
+    options.setDataflowContainerBaseRepository(
+        "docker.example.com/my-prefix/cloud-dataflow/v1beta3/");
+    String dockerHubPythonContainerUrl = "apache/beam_python3.11_sdk:2.68.0rc2";
+    String customPythonContainerUrl =
+        "docker.example.com/my-prefix/cloud-dataflow/v1beta3/beam_python3.11_sdk:2.68.0";
+    DataflowRunner runner = DataflowRunner.fromOptions(options);
+    RunnerApi.Pipeline pipeline = containerUrlToPipeline(dockerHubPythonContainerUrl);
+    RunnerApi.Pipeline expectedPipeline = containerUrlToPipeline(customPythonContainerUrl);
+    assertThat(runner.applySdkEnvironmentOverrides(pipeline, options), equalTo(expectedPipeline));
+  }
+
   @Test
   public void testStageArtifactWithoutStagedName() throws IOException {
     DataflowPipelineOptions options = buildPipelineOptions();
@@ -2106,6 +2121,44 @@ public void testGetContainerImageForJobFromOptionWithPlaceholder() {
     }
   }
 
+  @Test
+  public void testGetContainerImageBaseRepositoryDefault() {
+    DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class);
+    String defaultBaseRepository =
+        DataflowRunnerInfo.getDataflowRunnerInfo().getContainerImageBaseRepository();
+    assertThat(
+        DataflowRunner.getContainerImageBaseRepository(options), equalTo(defaultBaseRepository));
+    options.setDataflowContainerBaseRepository("");
+    assertThat(
+        DataflowRunner.getContainerImageBaseRepository(options), equalTo(defaultBaseRepository));
+    assertThat(getContainerImageForJob(options), startsWith(defaultBaseRepository + "/"));
+  }
+
+  @Test
+  public void testGetContainerImageForJobWithCustomBaseRepository() {
+    String baseRepository = "docker.example.com/my-prefix/cloud-dataflow/v1beta3";
+    DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class);
+    String defaultImage = getContainerImageForJob(options);
+    String defaultBaseRepository =
+        DataflowRunnerInfo.getDataflowRunnerInfo().getContainerImageBaseRepository();
+
+    for (String value : new String[] {baseRepository, baseRepository + "/"}) {
+      options.setDataflowContainerBaseRepository(value);
+      assertThat(DataflowRunner.getContainerImageBaseRepository(options), equalTo(baseRepository));
+      assertThat(
+          getContainerImageForJob(options),
+          equalTo(baseRepository + defaultImage.substring(defaultBaseRepository.length())));
+    }
+  }
+
+  @Test
+  public void testSdkContainerImageTakesPrecedenceOverBaseRepository() {
+    DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class);
+    options.setDataflowContainerBaseRepository("docker.example.com/prefix");
+    options.setSdkContainerImage("some-container");
+    assertThat(getContainerImageForJob(options), equalTo("some-container"));
+  }
+
   @Test
   public void testStreamingWriteWithNoShardingReturnsNewTransform() {
     PipelineOptions options = TestPipeline.testingPipelineOptions();

sdks/go/pkg/beam/runners/dataflow/dataflow.go

@@ -45,11 +45,13 @@ import (
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/protox"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow/dataflowlib"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcsx"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/hooks/perf"
+	"google.golang.org/protobuf/proto"
 )
 
 // TODO(herohde) 5/16/2017: the Dataflow flags should match the other SDKs.
@@ -59,6 +61,7 @@ var (
 	stagingLocation                = flag.String("staging_location", "", "GCS staging location (required).")
 	workerHarnessImage             = flag.String("worker_harness_container_image", "", "Worker harness container image (optional). Deprecated in favor of the sdk_container_image flag.")
 	image                          = flag.String("sdk_container_image", "", "Worker harness container image (optional).")
+	containerBaseRepository        = flag.String("dataflow_container_base_repository", "", "Base repository (registry host and path prefix) from which Dataflow pulls the default SDK container images, e.g. docker.example.com/my-prefix/cloud-dataflow/v1beta3 (optional). If set, it is also used for the default images of cross-language Apache Beam environments. Defaults to "+defaultContainerBaseRepository+".")
 	labels                         = flag.String("labels", "", "JSON-formatted map[string]string of job labels (optional).")
 	serviceAccountEmail            = flag.String("service_account_email", "", "Service account email (optional).")
 	numWorkers                     = flag.Int64("num_workers", 0, "Number of workers (optional).")
@@ -237,6 +240,16 @@ func Execute(ctx context.Context, p *beam.Pipeline) (beam.PipelineResult, error)
 	if err != nil {
 		return nil, errors.WithContext(err, "applying container image overrides")
 	}
+	// Only rewrite cross-language Apache Beam images when a base repository is
+	// set explicitly, to keep the default behavior of pulling them from Docker
+	// Hub unchanged. Images already rewritten by the overrides above no longer
+	// match apache/beam* and are left as is.
+	if strings.TrimSpace(*containerBaseRepository) != "" {
+		err = updateBeamContainerImagesForDataflow(model, getContainerBaseRepository(), opts.ContainerImage)
+		if err != nil {
+			return nil, errors.WithContext(err, "updating cross-language container images")
+		}
+	}
 
 	if *dryRun {
 		log.Info(ctx, "Dry-run: not submitting job!")
@@ -468,13 +481,72 @@ func getContainerImage(ctx context.Context) string {
 			// to have a dev image.
 			// However, RC versions should automatically be picked up, since
 			// they are never tagged the RC number, just the main version.
-			return "gcr.io/cloud-dataflow/v1beta3/beam_go_sdk:" + core.SdkVersion
+			return getContainerBaseRepository() + "/beam_go_sdk:" + core.SdkVersion
 		}
 		return envConfig
 	}
 	panic(fmt.Sprintf("Unsupported environment %v", urn))
 }
 
+// defaultContainerBaseRepository is the repository of the default Dataflow
+// container images, used when --dataflow_container_base_repository is unset.
+const defaultContainerBaseRepository = "gcr.io/cloud-dataflow/v1beta3"
+
+// getContainerBaseRepository returns the base repository of the default
+// Dataflow container images, without a trailing slash.
+func getContainerBaseRepository() string {
+	if repo := strings.TrimRight(strings.TrimSpace(*containerBaseRepository), "/"); repo != "" {
+		return repo
+	}
+	return defaultContainerBaseRepository
+}
+
+// updateBeamContainerImagesForDataflow rewrites the images of Docker
+// environments that use Apache Beam images from Docker Hub (apache/beam_*),
+// typically from cross-language transforms, to the same image in
+// baseRepository. Release candidate suffixes are trimmed from the tag, since
+// Dataflow images are only tagged with the release version. The image of the
+// pipeline's own SDK environment (sdkImage) is left unchanged.
+func updateBeamContainerImagesForDataflow(p *pipepb.Pipeline, baseRepository, sdkImage string) error {
+	for _, env := range p.GetComponents().GetEnvironments() {
+		if env.GetUrn() != graphx.URNEnvDocker {
+			continue
+		}
+		var payload pipepb.DockerPayload
+		if err := proto.Unmarshal(env.GetPayload(), &payload); err != nil {
+			return err
+		}
+		img := payload.GetContainerImage()
+		if !strings.HasPrefix(img, beamContainerImagePrefix) || img == sdkImage {
+			continue
+		}
+		payload.ContainerImage = baseRepository + "/" + trimRCTag(img[strings.LastIndex(img, "/")+1:])
+		pl, err := proto.Marshal(&payload)
+		if err != nil {
+			return err
+		}
+		env.Payload = pl
+	}
+	return nil
+}
+
+// beamContainerImagePrefix is the prefix of Apache Beam images on Docker Hub.
+const beamContainerImagePrefix = "apache/beam"
+
+// trimRCTag removes a release candidate suffix (e.g. "rc1") from the tag of
+// an image name such as "beam_java11_sdk:2.60.0rc1".
+func trimRCTag(imageAndTag string) string {
+	tagIdx := strings.LastIndex(imageAndTag, ":")
+	if tagIdx < 0 {
+		return imageAndTag
+	}
+	tag := imageAndTag[tagIdx:]
+	if rcIdx := strings.LastIndex(strings.ToLower(tag), "rc"); rcIdx > 0 {
+		return imageAndTag[:tagIdx] + tag[:rcIdx]
+	}
+	return imageAndTag
+}
+
 func checkSoftDeletePolicyEnabled(ctx context.Context, bucketName string, locationName string) {
 	bucket, _, err := gcsx.ParseObject(bucketName)
 	if err != nil {

sdks/go/pkg/beam/runners/dataflow/dataflow_test.go

@@ -24,9 +24,12 @@ import (
 	"testing"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx"
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow/dataflowlib"
+	"google.golang.org/protobuf/proto"
 )
 
 func TestDontUseFlagAsPipelineOption(t *testing.T) {
@@ -364,6 +367,67 @@ func TestGetJobOptions_DockerGCROverride(t *testing.T) {
 	}
 }
 
+func TestGetContainerImage_CustomBaseRepository(t *testing.T) {
+	const base = "docker.example.com/my-prefix/cloud-dataflow/v1beta3"
+	for _, value := range []string{base, base + "/", " " + base + " "} {
+		resetGlobals()
+		*jobopts.EnvironmentType = "docker"
+		*containerBaseRepository = value
+
+		if got, want := getContainerImage(context.Background()), base+"/beam_go_sdk:"+core.SdkVersion; got != want {
+			t.Errorf("getContainerImage() with base repository %q = %q, want %q", value, got, want)
+		}
+	}
+}
+
+func TestGetContainerImage_SdkContainerImageTakesPrecedenceOverBaseRepository(t *testing.T) {
+	resetGlobals()
+	*jobopts.EnvironmentType = "docker"
+	*containerBaseRepository = "docker.example.com/prefix"
+	*image = "testContainerImageOverride"
+
+	if got, want := getContainerImage(context.Background()), "testContainerImageOverride"; got != want {
+		t.Fatalf("getContainerImage() = %q, want %q", got, want)
+	}
+}
+
+func TestUpdateBeamContainerImagesForDataflow(t *testing.T) {
+	const base = "docker.example.com/my-prefix/cloud-dataflow/v1beta3"
+	const sdkImage = "apache/beam_go_sdk:2.60.0"
+	tests := map[string]struct {
+		urn, image, want string
+	}{
+		"xlang":      {graphx.URNEnvDocker, "apache/beam_java11_sdk:2.60.0", base + "/beam_java11_sdk:2.60.0"},
+		"xlangRC":    {graphx.URNEnvDocker, "apache/beam_python3.11_sdk:2.60.0rc1", base + "/beam_python3.11_sdk:2.60.0"},
+		"sdkImage":   {graphx.URNEnvDocker, sdkImage, sdkImage},
+		"nonBeam":    {graphx.URNEnvDocker, "other_org/beam_java11_sdk:2.60.0", "other_org/beam_java11_sdk:2.60.0"},
+		"overridden": {graphx.URNEnvDocker, "gcr.io/my-project/beam_java11_sdk:2.60.0", "gcr.io/my-project/beam_java11_sdk:2.60.0"},
+		"notDocker":  {"beam:env:process:v1", "apache/beam_java11_sdk:2.60.0", "apache/beam_java11_sdk:2.60.0"},
+	}
+	envs := map[string]*pipepb.Environment{}
+	for name, test := range tests {
+		pl, err := proto.Marshal(&pipepb.DockerPayload{ContainerImage: test.image})
+		if err != nil {
+			t.Fatal(err)
+		}
+		envs[name] = &pipepb.Environment{Urn: test.urn, Payload: pl}
+	}
+	p := &pipepb.Pipeline{Components: &pipepb.Components{Environments: envs}}
+
+	if err := updateBeamContainerImagesForDataflow(p, base, sdkImage); err != nil {
+		t.Fatalf("updateBeamContainerImagesForDataflow() returned error %v", err)
+	}
+	for name, test := range tests {
+		var payload pipepb.DockerPayload
+		if err := proto.Unmarshal(p.GetComponents().GetEnvironments()[name].GetPayload(), &payload); err != nil {
+			t.Fatal(err)
+		}
+		if got := payload.GetContainerImage(); got != test.want {
+			t.Errorf("%s: container image = %q, want %q", name, got, test.want)
+		}
+	}
+}
+
 func TestGetJobOptions_TransformMapping(t *testing.T) {
 	resetGlobals()
 	*stagingLocation = "gs://testStagingLocation"
@@ -584,6 +648,7 @@ func resetGlobals() {
 	*gcpopts.Project = ""
 	*gcpopts.Region = ""
 	*image = ""
+	*containerBaseRepository = ""
 	*jobopts.EnvironmentType = ""
 	*jobopts.EnvironmentConfig = ""
 	*jobopts.Experiments = ""

sdks/python/apache_beam/options/pipeline_options.py

@@ -1583,6 +1583,19 @@ def _add_argparse_args(cls, parser):
             'separated by a comma where first value gives a regex to '
             'identify the container image to override and the second value '
             'gives the replacement container image.'))
+    parser.add_argument(
+        '--dataflow_container_base_repository',
+        default=None,
+        help=(
+            'Base repository (registry host and path prefix) from which the '
+            'Dataflow runner pulls its default SDK container images, e.g. '
+            'docker.example.com/my-project/cloud-dataflow/v1beta3. The '
+            'runner appends the image name and tag, e.g. '
+            '<base>/beam_python3.11_sdk:<version>. This is also used for the '
+            'default images of cross-language Apache Beam environments. '
+            'Ignored for the local SDK image if --sdk_container_image is set. '
+            'If not set, the Google Cloud Dataflow default '
+            '(gcr.io/cloud-dataflow/v1beta3) is used.'))
     parser.add_argument(
         '--default_sdk_harness_log_level',
         default=None,

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -410,6 +410,13 @@ def run_pipeline(self, pipeline, options, pipeline_proto=None):
         # if prebuild_sdk_container_engine is specified we will build a new sdk
         # container image with dependencies pre-installed and use that image,
         # instead of using the inferred default container image.
+        worker_options = options.view_as(WorkerOptions)
+        if (worker_options.dataflow_container_base_repository and
+            not worker_options.sdk_container_image):
+          # Build on top of the Dataflow image in the configured repository
+          # instead of the Apache Beam image on Docker Hub.
+          worker_options.sdk_container_image = (
+              apiclient.get_container_image_from_options(options))
         self._default_environment = (
             environments.DockerEnvironment.from_options(options))
         options.view_as(WorkerOptions).sdk_container_image = (