diff --git a/CHANGES.md b/CHANGES.md
index d06940d12cf4..d044e08105bf 100644
--- a/CHANGES.md
+++ b/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
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
index f1b0c89090d3..796aaa960a4c 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
+++ b/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.
*/
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java
index bd2b968ff893..e8cc7ddc8e91 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java
+++ b/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. /beam_java17_sdk:. 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 network for launching
* workers.
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
index b196aa1714c6..02e701332939 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
+++ b/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();
diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflow.go b/sdks/go/pkg/beam/runners/dataflow/dataflow.go
index ce5e99b0a8de..87226404fcd7 100644
--- a/sdks/go/pkg/beam/runners/dataflow/dataflow.go
+++ b/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 {
diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go
index 83d3c0108439..5348eef6e1ac 100644
--- a/sdks/go/pkg/beam/runners/dataflow/dataflow_test.go
+++ b/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 = ""
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
index dca07da5146e..da3aae4e6a17 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/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. '
+ '/beam_python3.11_sdk:. 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,
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index 389e51b5728a..58bcbb5c334e 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/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 = (
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
index c10902cf1fc6..c68f053c7eb0 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
@@ -27,6 +27,7 @@
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import WorkerOptions
from apache_beam.pipeline import AppliedPTransform
from apache_beam.pipeline import Pipeline
from apache_beam.portability import common_urns
@@ -361,6 +362,55 @@ def test_environment_override_translation_sdk_container_image(self):
dependencies=environments.python_sdk_dependencies(options=options) +
main_session_dep))
+ def _run_prebuild_and_get_base_image(self, extra_properties):
+ self.default_properties.extend(
+ ['--prebuild_sdk_container_engine=local_docker'] + extra_properties)
+ options = PipelineOptions(self.default_properties)
+ options.view_as(DebugOptions).add_experiment(
+ 'disable_logging_submission_environment')
+ base_images = []
+
+ def fake_build(pipeline_options):
+ # The builder uses sdk_container_image as the base image if set.
+ base_images.append(
+ pipeline_options.view_as(WorkerOptions).sdk_container_image)
+ return 'prebuilt/beam_python_prebuilt_sdk:tag'
+
+ remote_runner = DataflowRunner()
+ with mock.patch('apache_beam.runners.portability.sdk_container_builder.'
+ 'SdkContainerImageBuilder.build_container_image',
+ side_effect=fake_build):
+ with Pipeline(remote_runner, options=options) as p:
+ p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
+
+ from apache_beam.utils import proto_utils
+ images = [
+ proto_utils.parse_Bytes(
+ env.payload, beam_runner_api_pb2.DockerPayload).container_image for
+ env in remote_runner.proto_pipeline.components.environments.values()
+ ]
+ self.assertEqual(images, ['prebuilt/beam_python_prebuilt_sdk:tag'])
+ self.assertEqual(len(base_images), 1)
+ return base_images[0]
+
+ def test_prebuild_uses_dataflow_container_base_repository(self):
+ base = 'docker.example.com/my-prefix/cloud-dataflow/v1beta3'
+ base_image = self._run_prebuild_and_get_base_image(
+ ['--dataflow_container_base_repository=' + base])
+ self.assertTrue(
+ base_image.startswith(base + '/beam_python'),
+ 'unexpected base image %s' % base_image)
+
+ def test_prebuild_without_base_repository_keeps_default_base_image(self):
+ self.assertIsNone(self._run_prebuild_and_get_base_image([]))
+
+ def test_prebuild_sdk_container_image_takes_precedence(self):
+ base_image = self._run_prebuild_and_get_base_image([
+ '--dataflow_container_base_repository=docker.example.com/prefix',
+ '--sdk_container_image=FOO'
+ ])
+ self.assertEqual(base_image, 'FOO')
+
def test_remote_runner_translation(self):
remote_runner = DataflowRunner()
with Pipeline(remote_runner,
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
index 8a478f3fb633..7cb4e3fff932 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -789,7 +789,8 @@ def create_job(self, job: Job):
return None
@staticmethod
- def _update_container_image_for_dataflow(beam_container_image_url):
+ def _update_container_image_for_dataflow(
+ beam_container_image_url, pipeline_options=None):
# By default Dataflow pipelines use containers hosted in Dataflow GCR
# instead of Docker Hub.
image_suffix = beam_container_image_url.rsplit('/', 1)[1]
@@ -799,7 +800,9 @@ def _update_container_image_for_dataflow(beam_container_image_url):
if len(check_rc) == 2:
image_suffix = image_suffix[:-2 - len(check_rc[1])]
- return names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/' + image_suffix
+ return (
+ get_container_base_repository_from_options(pipeline_options) + '/' +
+ image_suffix)
@staticmethod
def _apply_sdk_environment_overrides(
@@ -831,7 +834,7 @@ def _apply_sdk_environment_overrides(
new_container_image != current_sdk_container_image):
new_container_image = (
DataflowApplicationClient._update_container_image_for_dataflow(
- docker_payload.container_image))
+ docker_payload.container_image, pipeline_options))
if not new_container_image:
raise ValueError(
@@ -1200,6 +1203,26 @@ def _get_container_image_tag():
return base_version
+def get_container_base_repository_from_options(pipeline_options):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ Args:
+ pipeline_options (PipelineOptions): A container for pipeline options.
+ May be None, in which case the default repository is returned.
+
+ Returns:
+ str: Base repository of the default Dataflow container images, without
+ a trailing slash.
+ """
+ base_repository = None
+ if pipeline_options is not None:
+ base_repository = pipeline_options.view_as(
+ WorkerOptions).dataflow_container_base_repository
+ if base_repository:
+ base_repository = base_repository.strip().rstrip('/')
+ return base_repository or names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY
+
+
def get_container_image_from_options(pipeline_options):
"""For internal use only; no backwards-compatibility guarantees.
@@ -1215,7 +1238,7 @@ def get_container_image_from_options(pipeline_options):
# Dataflow Legacy and Portable Runner exist in different repositories.
# Set to legacy format, override if Dataflow Portable Runner
- container_repo = names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY
+ container_repo = get_container_base_repository_from_options(pipeline_options)
image_name = '{repository}/beam_python{major}.{minor}_sdk'.format(
repository=container_repo,
major=sys.version_info[0],
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
index b82c7f550673..368ac5dbe381 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
@@ -520,6 +520,81 @@ def test_pipeline_sdk_not_overridden_prime(self):
self._verify_pipeline_sdk_not_overridden(pipeline_options)
+ def test_default_container_base_repository(self):
+ pipeline_options = PipelineOptions(['--temp_location', 'gs://any/temp'])
+ self.assertEqual(
+ apiclient.get_container_base_repository_from_options(pipeline_options),
+ names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY)
+ self.assertEqual(
+ apiclient.get_container_base_repository_from_options(None),
+ names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY)
+ self.assertTrue(
+ apiclient.get_container_image_from_options(pipeline_options).startswith(
+ names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python'))
+
+ def test_custom_container_base_repository(self):
+ base = 'docker.example.com/my-prefix/cloud-dataflow/v1beta3'
+ for value in [base, base + '/']:
+ pipeline_options = PipelineOptions([
+ '--temp_location',
+ 'gs://any/temp',
+ '--dataflow_container_base_repository=' + value
+ ])
+ self.assertEqual(
+ apiclient.get_container_base_repository_from_options(
+ pipeline_options),
+ base)
+ self.assertEqual(
+ apiclient.get_container_image_from_options(pipeline_options),
+ '%s/beam_python%d.%d_sdk:%s' % (
+ base,
+ sys.version_info[0],
+ sys.version_info[1],
+ apiclient._get_required_container_version()))
+
+ def test_sdk_container_image_takes_precedence_over_base_repository(self):
+ pipeline_options = PipelineOptions([
+ '--temp_location',
+ 'gs://any/temp',
+ '--dataflow_container_base_repository=docker.example.com/prefix',
+ '--sdk_container_image=dummy_prefix/dummy_name:dummy_tag'
+ ])
+ self.assertEqual(
+ apiclient.get_container_image_from_options(pipeline_options),
+ 'dummy_prefix/dummy_name:dummy_tag')
+
+ def test_custom_container_base_repository_for_xlang_environments(self):
+ base = 'docker.example.com/my-prefix/cloud-dataflow/v1beta3'
+ pipeline_options = PipelineOptions([
+ '--experiments=use_runner_v2',
+ '--temp_location',
+ 'gs://any-location/temp',
+ '--dataflow_container_base_repository=' + base
+ ])
+ pipeline = Pipeline(options=pipeline_options)
+ pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned
+
+ dummy_env = DockerEnvironment(
+ container_image='apache/beam_java11_sdk:2.00.0rc1')
+ proto_pipeline, _ = pipeline.to_runner_api(
+ return_context=True, default_environment=dummy_env)
+
+ # Accessing non-public method for testing.
+ apiclient.DataflowApplicationClient._apply_sdk_environment_overrides(
+ proto_pipeline, {}, pipeline_options)
+
+ from apache_beam.utils import proto_utils
+ images = [
+ proto_utils.parse_Bytes(
+ env.payload, beam_runner_api_pb2.DockerPayload).container_image
+ for env in proto_pipeline.components.environments.values()
+ ]
+ self.assertIn(base + '/beam_java11_sdk:2.00.0', images)
+ self.assertFalse(
+ any(
+ image.startswith(names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY)
+ for image in images))
+
def test_invalid_default_job_name(self):
# Regexp for job names in dataflow.
regexp = '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$'