Resolve staged sklearn model path on the worker in RunInference xlang tests

.github/trigger_files/beam_PostCommit_XVR_Direct.json

@@ -1,4 +1,4 @@
 {
-  "modification": 3,
+  "modification": 4,
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 2
+    "modification": 3
 }

sdks/go/test/integration/transforms/xlang/inference/inference.go

@@ -22,8 +22,23 @@ import (
 	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/universal"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/xlang/inference"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/xlang/python"
 )
 
+// stagedSklearnModelHandler loads the model the test expansion service stages
+// with the Python environment as "sklearn_model". Staged files are only
+// reachable on the worker, so the path is resolved in load_model().
+const stagedSklearnModelHandler = `import os
+from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy
+
+class StagedSklearnModelHandlerNumpy(SklearnModelHandlerNumpy):
+  def load_model(self):
+    self.update_model_path(os.path.join(
+        os.environ['SEMI_PERSISTENT_DIRECTORY'], 'staged',
+        os.path.basename(self._model_uri)))
+    return super().load_model()
+`
+
 func SklearnInference(expansionAddr string) *beam.Pipeline {
 	p, s := beam.NewPipelineWithRoot()
 
@@ -39,7 +54,9 @@ func SklearnInference(expansionAddr string) *beam.Pipeline {
 			Inference: 1,
 		},
 	}
-	outCol := inference.SklearnModel("/tmp/staged/sklearn_model").RunInference(s, input, inference.WithExpansionAddr(expansionAddr))
+	model := inference.SklearnModel("sklearn_model")
+	model.ModelHandlerProvider = python.CallableSource(stagedSklearnModelHandler)
+	outCol := model.RunInference(s, input, inference.WithExpansionAddr(expansionAddr))
 	passert.Equals(s, outCol, output[0], output[1])
 	return p
 }

sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java

@@ -19,7 +19,6 @@
 
 import java.util.Arrays;
 import java.util.List;
-import java.util.Optional;
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.VarLongCoder;
@@ -42,11 +41,22 @@
 
 @RunWith(JUnit4.class)
 public class RunInferenceTransformTest extends BaseExternalTest {
+  // The test expansion service stages the test model with the Python environment as
+  // "sklearn_model". Staged files are only reachable on the worker, so the path is resolved
+  // in load_model().
+  private static final String STAGED_MODEL_HANDLER =
+      "import os\n"
+          + "from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy\n"
+          + "class StagedSklearnModelHandlerNumpy(SklearnModelHandlerNumpy):\n"
+          + "  def load_model(self):\n"
+          + "    self.update_model_path(os.path.join(\n"
+          + "        os.environ['SEMI_PERSISTENT_DIRECTORY'], 'staged',\n"
+          + "        os.path.basename(self._model_uri)))\n"
+          + "    return super().load_model()\n";
+
   @Test
   @Category({ValidatesRunner.class, UsesPythonExpansionService.class})
   public void testRunInference() {
-    String stagingLocation =
-        Optional.ofNullable(System.getProperty("semiPersistDir")).orElse("/tmp");
     Schema schema =
         Schema.of(
             Schema.Field.of("example", Schema.FieldType.array(Schema.FieldType.INT64)),
@@ -58,25 +68,17 @@ public void testRunInference() {
             .apply(Create.<Iterable<Long>>of(Arrays.asList(0L, 0L), Arrays.asList(1L, 1L)))
             .setCoder(IterableCoder.of(VarLongCoder.of()))
             .apply(
-                RunInference.of(
-                        "apache_beam.ml.inference.sklearn_inference.SklearnModelHandlerNumpy",
-                        schema)
-                    .withKwarg(
-                        // The test expansion service creates the test model and saves it to the
-                        // returning external environment as a dependency.
-                        // (sdks/python/apache_beam/runners/portability/expansion_service_test.py)
-                        // The dependencies for Python SDK harness are supposed to be staged to
-                        // $SEMI_PERSIST_DIR/staged directory.
-                        "model_uri", String.format("%s/staged/sklearn_model", stagingLocation))
+                RunInference.of(STAGED_MODEL_HANDLER, schema)
+                    .withKwarg("model_uri", "sklearn_model")
                     .withExpansionService(expansionAddr));
     PAssert.that(col).containsInAnyOrder(row0, row1);
   }
 
   private String getModelLoaderScriptWithKVs() {
-    String s = "from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy\n";
+    String s = STAGED_MODEL_HANDLER;
     s = s + "from apache_beam.ml.inference.base import KeyedModelHandler\n";
     s = s + "def get_model_handler(model_uri):\n";
-    s = s + "  return KeyedModelHandler(SklearnModelHandlerNumpy(model_uri))\n";
+    s = s + "  return KeyedModelHandler(StagedSklearnModelHandlerNumpy(model_uri))\n";
 
     return s;
   }
@@ -92,8 +94,6 @@ public KV<Long, Iterable<Long>> apply(Iterable<Long> input) {
   @Test
   @Category({ValidatesRunner.class, UsesPythonExpansionService.class})
   public void testRunInferenceWithKVs() {
-    String stagingLocation =
-        Optional.ofNullable(System.getProperty("semiPersistDir")).orElse("/tmp");
     Schema schema =
         Schema.of(
             Schema.Field.of("example", Schema.FieldType.array(Schema.FieldType.INT64)),
@@ -107,13 +107,7 @@ public void testRunInferenceWithKVs() {
             .setCoder(KvCoder.of(VarLongCoder.of(), IterableCoder.of(VarLongCoder.of())))
             .apply(
                 RunInference.ofKVs(getModelLoaderScriptWithKVs(), schema, VarLongCoder.of())
-                    .withKwarg(
-                        // The test expansion service creates the test model and saves it to the
-                        // returning external environment as a dependency.
-                        // (sdks/python/apache_beam/runners/portability/expansion_service_test.py)
-                        // The dependencies for Python SDK harness are supposed to be staged to
-                        // $SEMI_PERSIST_DIR/staged directory.
-                        "model_uri", String.format("%s/staged/sklearn_model", stagingLocation))
+                    .withKwarg("model_uri", "sklearn_model")
                     .withExpansionService(expansionAddr))
             .apply(Values.<Row>create());