This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 11e1ce2c1d3 Add an option to set the Dataflow container base
repository (Python, Java, Go) (#40690)
11e1ce2c1d3 is described below
commit 11e1ce2c1d3894f9479799ce8c646fce835c6c9f
Author: Shunping Huang <[email protected]>
AuthorDate: Tue Oct 20 19:46:53 2026 -0400
Add an option to set the Dataflow container base repository (Python, Java,
Go) (#40690)
* [Python] Add --dataflow_container_base_repository option
Let users point the Dataflow runner's default SDK container images, and
the rewritten cross-language apache/beam images, at a base repository
other than gcr.io/cloud-dataflow/v1beta3. This lets jobs run in Trusted
Partner Cloud (TPC) universes, which have their own registries. If the
option is unset, behavior is unchanged.
* [Java] Add --dataflowContainerBaseRepository option
* [Go] Add --dataflow_container_base_repository flag
* [Python] Use dataflow_container_base_repository for prebuilt SDK base
image
* [Java] Mark dataflowContainerBaseRepository as @Nullable
---
CHANGES.md | 1 +
.../beam/runners/dataflow/DataflowRunner.java | 25 ++++++--
.../options/DataflowPipelineWorkerPoolOptions.java | 17 +++++
.../beam/runners/dataflow/DataflowRunnerTest.java | 53 +++++++++++++++
sdks/go/pkg/beam/runners/dataflow/dataflow.go | 74 ++++++++++++++++++++-
sdks/go/pkg/beam/runners/dataflow/dataflow_test.go | 65 +++++++++++++++++++
.../python/apache_beam/options/pipeline_options.py | 13 ++++
.../runners/dataflow/dataflow_runner.py | 7 ++
.../runners/dataflow/dataflow_runner_test.py | 50 +++++++++++++++
.../runners/dataflow/internal/apiclient.py | 31 +++++++--
.../runners/dataflow/internal/apiclient_test.py | 75 ++++++++++++++++++++++
11 files changed, 401 insertions(+), 10 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index d06940d12cf..d044e08105b 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 f1b0c89090d..796aaa960a4 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.sdk.values.WindowingStrategy;
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 @@ public class DataflowRunner extends
PipelineRunner<DataflowPipelineJob> {
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 @@ public class DataflowRunner extends
PipelineRunner<DataflowPipelineJob> {
/** 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 bd2b968ff89..e8cc7ddc8e9 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 interface DataflowPipelineWorkerPoolOptions extends
GcpOptions, FileStagi
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.
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 b196aa1714c..02e70133293 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.is;
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 class DataflowRunnerTest implements Serializable {
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 class DataflowRunnerTest implements Serializable {
}
}
+ @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 ce5e99b0a8d..87226404fcd 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 83d3c010843..5348eef6e1a 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 dca07da5146..da3aae4e6a1 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1583,6 +1583,19 @@ class WorkerOptions(PipelineOptions):
'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,
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index 389e51b5728..58bcbb5c334 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 @@ class DataflowRunner(PipelineRunner):
# 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 c10902cf1fc..c68f053c7eb 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 @@ import apache_beam.transforms as ptransform
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 @@ class DataflowRunnerTest(unittest.TestCase,
ExtraAssertionsMixin):
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 8a478f3fb63..7cb4e3fff93 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 @@ class DataflowApplicationClient(object):
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 @@ class DataflowApplicationClient(object):
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 @@ class DataflowApplicationClient(object):
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 b82c7f55067..368ac5dbe38 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 @@ class UtilTest(unittest.TestCase):
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])?$'
