Allow override beam version for PythonExternalTransform via pipeline …

…option (#31691)

sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java

@@ -486,12 +486,10 @@ public OutputT expand(InputT input) {
         OutputT output = null;
         int port = PythonService.findAvailablePort();
         PipelineOptionsFactory.register(PythonExternalTransformOptions.class);
-        boolean useTransformService =
-            input
-                .getPipeline()
-                .getOptions()
-                .as(PythonExternalTransformOptions.class)
-                .getUseTransformService();
+        PythonExternalTransformOptions options =
+            input.getPipeline().getOptions().as(PythonExternalTransformOptions.class);
+        boolean useTransformService = options.getUseTransformService();
+        @Nullable String customBeamRequirement = options.getCustomBeamRequirement();
         boolean pythonAvailable = isPythonAvailable();
         boolean dockerAvailable = isDockerAvailable();
 
@@ -557,6 +555,9 @@ public OutputT expand(InputT input) {
               new PythonService(
                       "apache_beam.runners.portability.expansion_service_main", args.build())
                   .withExtraPackages(extraPackages);
+          if (!Strings.isNullOrEmpty(customBeamRequirement)) {
+            service = service.withCustomBeamRequirement(customBeamRequirement);
+          }
           try (AutoCloseable p = service.start()) {
             // allow more time waiting for the port ready for transient expansion service setup.
             PythonService.waitForPort("localhost", port, 60000);

sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformOptions.java

@@ -29,4 +29,14 @@ public interface PythonExternalTransformOptions extends PipelineOptions {
   boolean getUseTransformService();
 
   void setUseTransformService(boolean useTransformService);
+
+  @Description("Custom Beam version for bootstrap Beam venv.")
+  String getCustomBeamRequirement();
+
+  /**
+   * Set custom Beam version for bootstrap Beam venv.
+   *
+   * <p>For example: 2.50.0rc1, "/path/to/apache-beam.whl"
+   */
+  void setCustomBeamRequirement(String customBeamRequirement);
 }

sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonService.java

@@ -41,13 +41,16 @@ public class PythonService {
   private static final Logger LOG = LoggerFactory.getLogger(PythonService.class);
 
   private final String module;
+  private String beamRequirement;
   private final List<String> args;
   private final List<String> extraPackages;
 
   public PythonService(String module, List<String> args, List<String> extraPackages) {
     this.module = module;
     this.args = args;
     this.extraPackages = extraPackages;
+    this.beamRequirement =
+        getMatchingStablePythonSDKVersion(ReleaseInfo.getReleaseInfo().getSdkVersion());
   }
 
   public PythonService(String module, List<String> args) {
@@ -72,6 +75,18 @@ public PythonService withExtraPackages(List<String> extraPackages) {
         ImmutableList.<String>builder().addAll(this.extraPackages).addAll(extraPackages).build());
   }
 
+  /**
+   * Override the Beam version to be installed in the service environment.
+   *
+   * @param customBeamRequirement the custom Beam requirement, can be a version (e.g. 2.57.0) or a
+   *     path (e.g. /path/to/apache-beam.whl).
+   * @return this instance where Beam version overriden in place.
+   */
+  public PythonService withCustomBeamRequirement(String customBeamRequirement) {
+    this.beamRequirement = customBeamRequirement;
+    return this;
+  }
+
   @SuppressWarnings("argument")
   public AutoCloseable start() throws IOException, InterruptedException {
     File bootstrapScript = File.createTempFile("bootstrap_beam_venv", ".py");
@@ -82,9 +97,7 @@ public AutoCloseable start() throws IOException, InterruptedException {
     List<String> bootstrapCommand = new ArrayList<>();
     bootstrapCommand.add(whichPython());
     bootstrapCommand.add(bootstrapScript.getAbsolutePath());
-    bootstrapCommand.add(
-        "--beam_version="
-            + getMatchingStablePythonSDKVersion(ReleaseInfo.getReleaseInfo().getSdkVersion()));
+    bootstrapCommand.add("--beam_version=" + beamRequirement);
     if (!extraPackages.isEmpty()) {
       bootstrapCommand.add("--extra_packages=" + String.join(";", extraPackages));
     }