Merge pull request #23051: Fixes RunInference test failure

chamikaramj · web-flow · commit c78fc9a0dd72 · 2022-09-06T22:22:56.000-07:00
diff --git a/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java b/sdks/java/extensions/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/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/transforms/RunInference.java b/sdks/java/extensions/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());
     }
 

Original file line number	Diff line number	Diff line change
`@@ -130,6 +130,9 @@ public RunInference<OutputT> withKwarg(String key, Object arg) {`
`130`	`130`	`/**`
`131`	`131`	`* Specifies any extra Pypi packages required by the RunInference model handler.`
`132`	`132`	`*`
	`133`	`+ * <p>This should only be specified when using the default expansion service, i.e. when not using`
	`134`	`+ * {@link #withExpansionService(String)} to provide an expansion service.`
	`135`	`+ *`
`133`	`136`	`* <p>For model handlers provided by Beam Python SDK, the implementation will automatically try to`
`134`	`137`	`* infer correct packages needed, so this may be omitted.`
`135`	`138`	`*`
`@@ -192,7 +195,7 @@ public PCollection<OutputT> expand(PCollection<?> input) {`
`192`	`195`	`outputCoder = (Coder<OutputT>) KvCoder.of(keyCoder, RowCoder.of(schema));`
`193`	`196`	`}`
`194`	`197`
`195`		`- if (this.extraPackages.isEmpty()) {`
	`198`	`+ if (this.expansionService.isEmpty() && this.extraPackages.isEmpty()) {`
`196`	`199`	`this.extraPackages.addAll(inferExtraPackagesFromModelHandler());`
`197`	`200`	`}`
`198`	`201`