Merge pull request #23051: Fixes RunInference test failure

apache · Sep 7, 2022 · c78fc9a · c78fc9a
2 parents de35279 + 62dd4b4
commit c78fc9a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/...s/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java b/...s/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
@@ -286,6 +286,9 @@ public PythonExternalTransform<InputT, OutputT> withOutputCoder(Coder<?> outputC
    * @return updated wrapper for the cross-language transform.
    */
   public PythonExternalTransform<InputT, OutputT> withExtraPackages(List<String> extraPackages) {
+    if (extraPackages.isEmpty()) {
+      return this;
+    }
     Preconditions.checkState(
         Strings.isNullOrEmpty(expansionService),
         "Extra packages only apply to auto-started expansion service.");

diff --git a/...s/python/src/main/java/org/apache/beam/sdk/extensions/python/transforms/RunInference.java b/...s/python/src/main/java/org/apache/beam/sdk/extensions/python/transforms/RunInference.java
@@ -130,6 +130,9 @@ public RunInference<OutputT> withKwarg(String key, Object arg) {
   /**
    * Specifies any extra Pypi packages required by the RunInference model handler.
    *
+   * <p>This should only be specified when using the default expansion service, i.e. when not using
+   * {@link #withExpansionService(String)} to provide an expansion service.
+   *
    * <p>For model handlers provided by Beam Python SDK, the implementation will automatically try to
    * infer correct packages needed, so this may be omitted.
    *
@@ -192,7 +195,7 @@ public PCollection<OutputT> expand(PCollection<?> input) {
       outputCoder = (Coder<OutputT>) KvCoder.of(keyCoder, RowCoder.of(schema));
     }
 
-    if (this.extraPackages.isEmpty()) {
+    if (this.expansionService.isEmpty() && this.extraPackages.isEmpty()) {
       this.extraPackages.addAll(inferExtraPackagesFromModelHandler());
     }