From b3252be1a9796aa387d88f804551ff04f4e970b3 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Sun, 25 Oct 2026 18:28:46 -0700 Subject: [PATCH 1/5] Resolve staged sklearn model path on the worker in RunInference xlang tests --- .../transforms/xlang/inference/inference.go | 33 +++++++++- .../transforms/RunInferenceTransformTest.java | 60 +++++++++++-------- 2 files changed, 68 insertions(+), 25 deletions(-) diff --git a/sdks/go/test/integration/transforms/xlang/inference/inference.go b/sdks/go/test/integration/transforms/xlang/inference/inference.go index 7fd89874fa94..f0a7a5530a56 100644 --- a/sdks/go/test/integration/transforms/xlang/inference/inference.go +++ b/sdks/go/test/integration/transforms/xlang/inference/inference.go @@ -22,8 +22,37 @@ 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" ) +// stagedModelName is the name under which the test expansion service +// (sdks/python/apache_beam/runners/portability/expansion_service_test.py) +// stages the test model with the Python environment. +const stagedModelName = "sklearn_model" + +// stagedSklearnModelHandler is a SklearnModelHandlerNumpy whose model_uri is +// the name of a file staged with the Python environment. Where the Python SDK +// harness puts staged files is only known on the worker, so the full path is +// resolved in load_model(), after the harness has set SEMI_PERSISTENT_DIRECTORY. +const stagedSklearnModelHandler = `import os +from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy + +class StagedSklearnModelHandlerNumpy(SklearnModelHandlerNumpy): + def __init__(self, model_uri, *args, **kwargs): + super().__init__(model_uri, *args, **kwargs) + self._staged_name = model_uri + + def load_model(self): + staged_dir = os.environ.get('SEMI_PERSISTENT_DIRECTORY') + if not staged_dir: + raise RuntimeError( + 'SEMI_PERSISTENT_DIRECTORY is not set, cannot locate staged file ' + + self._staged_name) + self.update_model_path( + os.path.join(staged_dir, 'staged', self._staged_name)) + return super().load_model() +` + func SklearnInference(expansionAddr string) *beam.Pipeline { p, s := beam.NewPipelineWithRoot() @@ -39,7 +68,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(stagedModelName) + model.ModelHandlerProvider = python.CallableSource(stagedSklearnModelHandler) + outCol := model.RunInference(s, input, inference.WithExpansionAddr(expansionAddr)) passert.Equals(s, outCol, output[0], output[1]) return p } diff --git a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java index 2d1b520d731a..f7644147428d 100644 --- a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java +++ b/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; @@ -40,13 +39,42 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +/** + * Tests the Java RunInference wrapper against the Python test expansion service. + * + *

The test expansion service attaches a small sklearn model to the Python environment as a + * staged file named {@code sklearn_model} (see + * sdks/python/apache_beam/runners/portability/expansion_service_test.py). Where the Python SDK + * harness puts staged files is only known on the worker, so the tests pass the staged file name and + * resolve the full path when the model is loaded. + */ @RunWith(JUnit4.class) public class RunInferenceTransformTest extends BaseExternalTest { + private static final String STAGED_MODEL_NAME = "sklearn_model"; + + // A SklearnModelHandlerNumpy whose model_uri is the name of a file staged with the Python + // environment. load_model() runs on the worker, after the SDK harness has set + // SEMI_PERSISTENT_DIRECTORY. + 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 __init__(self, model_uri, *args, **kwargs):\n" + + " super().__init__(model_uri, *args, **kwargs)\n" + + " self._staged_name = model_uri\n" + + " def load_model(self):\n" + + " staged_dir = os.environ.get('SEMI_PERSISTENT_DIRECTORY')\n" + + " if not staged_dir:\n" + + " raise RuntimeError(\n" + + " 'SEMI_PERSISTENT_DIRECTORY is not set, cannot locate staged file '\n" + + " + self._staged_name)\n" + + " self.update_model_path(\n" + + " os.path.join(staged_dir, 'staged', self._staged_name))\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 +86,17 @@ public void testRunInference() { .apply(Create.>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", STAGED_MODEL_NAME) .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 +112,6 @@ public KV> apply(Iterable 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 +125,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", STAGED_MODEL_NAME) .withExpansionService(expansionAddr)) .apply(Values.create()); From 842ab1264f24cfef675115196ec083c761bcf609 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Sun, 25 Oct 2026 18:28:47 -0700 Subject: [PATCH 2/5] Trigger XVR postcommits --- .github/trigger_files/beam_PostCommit_XVR_Direct.json | 2 +- .../beam_PostCommit_XVR_JavaUsingPython_Dataflow.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_XVR_Direct.json b/.github/trigger_files/beam_PostCommit_XVR_Direct.json index 49cff1ed291e..6a84ff06e55f 100644 --- a/.github/trigger_files/beam_PostCommit_XVR_Direct.json +++ b/.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" } diff --git a/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json b/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json index 7ab7bcd9a9c6..37dd25bf9029 100644 --- a/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json +++ b/.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 } From bf4247037da2d36cc5f32ad8a03660777882fba5 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Sun, 25 Oct 2026 18:43:33 -0700 Subject: [PATCH 3/5] Simplify staged model handler in RunInference xlang tests --- .../transforms/xlang/inference/inference.go | 30 +++++----------- .../transforms/RunInferenceTransformTest.java | 34 +++++-------------- 2 files changed, 17 insertions(+), 47 deletions(-) diff --git a/sdks/go/test/integration/transforms/xlang/inference/inference.go b/sdks/go/test/integration/transforms/xlang/inference/inference.go index f0a7a5530a56..e793e4cdab1e 100644 --- a/sdks/go/test/integration/transforms/xlang/inference/inference.go +++ b/sdks/go/test/integration/transforms/xlang/inference/inference.go @@ -25,31 +25,19 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/xlang/python" ) -// stagedModelName is the name under which the test expansion service -// (sdks/python/apache_beam/runners/portability/expansion_service_test.py) -// stages the test model with the Python environment. -const stagedModelName = "sklearn_model" - -// stagedSklearnModelHandler is a SklearnModelHandlerNumpy whose model_uri is -// the name of a file staged with the Python environment. Where the Python SDK -// harness puts staged files is only known on the worker, so the full path is -// resolved in load_model(), after the harness has set SEMI_PERSISTENT_DIRECTORY. +// stagedSklearnModelHandler loads the model the test expansion service stages +// with the Python environment as "sklearn_model" +// (sdks/python/apache_beam/runners/portability/expansion_service_test.py). +// 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 __init__(self, model_uri, *args, **kwargs): - super().__init__(model_uri, *args, **kwargs) - self._staged_name = model_uri - def load_model(self): - staged_dir = os.environ.get('SEMI_PERSISTENT_DIRECTORY') - if not staged_dir: - raise RuntimeError( - 'SEMI_PERSISTENT_DIRECTORY is not set, cannot locate staged file ' - + self._staged_name) - self.update_model_path( - os.path.join(staged_dir, 'staged', self._staged_name)) + self.update_model_path(os.path.join( + os.environ['SEMI_PERSISTENT_DIRECTORY'], 'staged', + os.path.basename(self._model_uri))) return super().load_model() ` @@ -68,7 +56,7 @@ func SklearnInference(expansionAddr string) *beam.Pipeline { Inference: 1, }, } - model := inference.SklearnModel(stagedModelName) + 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]) diff --git a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java index f7644147428d..4d2b3fa2b08e 100644 --- a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java +++ b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java @@ -39,37 +39,19 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** - * Tests the Java RunInference wrapper against the Python test expansion service. - * - *

The test expansion service attaches a small sklearn model to the Python environment as a - * staged file named {@code sklearn_model} (see - * sdks/python/apache_beam/runners/portability/expansion_service_test.py). Where the Python SDK - * harness puts staged files is only known on the worker, so the tests pass the staged file name and - * resolve the full path when the model is loaded. - */ @RunWith(JUnit4.class) public class RunInferenceTransformTest extends BaseExternalTest { - private static final String STAGED_MODEL_NAME = "sklearn_model"; - - // A SklearnModelHandlerNumpy whose model_uri is the name of a file staged with the Python - // environment. load_model() runs on the worker, after the SDK harness has set - // SEMI_PERSISTENT_DIRECTORY. + // The test expansion service stages the test model with the Python environment as + // "sklearn_model" (sdks/python/apache_beam/runners/portability/expansion_service_test.py). + // 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 __init__(self, model_uri, *args, **kwargs):\n" - + " super().__init__(model_uri, *args, **kwargs)\n" - + " self._staged_name = model_uri\n" + " def load_model(self):\n" - + " staged_dir = os.environ.get('SEMI_PERSISTENT_DIRECTORY')\n" - + " if not staged_dir:\n" - + " raise RuntimeError(\n" - + " 'SEMI_PERSISTENT_DIRECTORY is not set, cannot locate staged file '\n" - + " + self._staged_name)\n" - + " self.update_model_path(\n" - + " os.path.join(staged_dir, 'staged', self._staged_name))\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 @@ -87,7 +69,7 @@ public void testRunInference() { .setCoder(IterableCoder.of(VarLongCoder.of())) .apply( RunInference.of(STAGED_MODEL_HANDLER, schema) - .withKwarg("model_uri", STAGED_MODEL_NAME) + .withKwarg("model_uri", "sklearn_model") .withExpansionService(expansionAddr)); PAssert.that(col).containsInAnyOrder(row0, row1); } @@ -125,7 +107,7 @@ public void testRunInferenceWithKVs() { .setCoder(KvCoder.of(VarLongCoder.of(), IterableCoder.of(VarLongCoder.of()))) .apply( RunInference.ofKVs(getModelLoaderScriptWithKVs(), schema, VarLongCoder.of()) - .withKwarg("model_uri", STAGED_MODEL_NAME) + .withKwarg("model_uri", "sklearn_model") .withExpansionService(expansionAddr)) .apply(Values.create()); From 7d688407a93ce672c8194f19926c5cc306afb984 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Sun, 25 Oct 2026 19:03:52 -0700 Subject: [PATCH 4/5] Drop file reference from staged model handler comments --- .../integration/transforms/xlang/inference/inference.go | 6 ++---- .../python/transforms/RunInferenceTransformTest.java | 4 ++-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/sdks/go/test/integration/transforms/xlang/inference/inference.go b/sdks/go/test/integration/transforms/xlang/inference/inference.go index e793e4cdab1e..e492b7e7e78e 100644 --- a/sdks/go/test/integration/transforms/xlang/inference/inference.go +++ b/sdks/go/test/integration/transforms/xlang/inference/inference.go @@ -26,10 +26,8 @@ import ( ) // stagedSklearnModelHandler loads the model the test expansion service stages -// with the Python environment as "sklearn_model" -// (sdks/python/apache_beam/runners/portability/expansion_service_test.py). -// Staged files are only reachable on the worker, so the path is resolved in -// load_model(). +// 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 diff --git a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java index 4d2b3fa2b08e..2bc56c438555 100644 --- a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java +++ b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/transforms/RunInferenceTransformTest.java @@ -42,8 +42,8 @@ @RunWith(JUnit4.class) public class RunInferenceTransformTest extends BaseExternalTest { // The test expansion service stages the test model with the Python environment as - // "sklearn_model" (sdks/python/apache_beam/runners/portability/expansion_service_test.py). - // Staged files are only reachable on the worker, so the path is resolved in load_model(). + // "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" From 9acd720bfd2a1f17025a00bd936f5c7aa77a560c Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Mon, 26 Oct 2026 09:42:15 -0700 Subject: [PATCH 5/5] Retrigger tests