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 } diff --git a/sdks/go/test/integration/transforms/xlang/inference/inference.go b/sdks/go/test/integration/transforms/xlang/inference/inference.go index 7fd89874fa94..e492b7e7e78e 100644 --- a/sdks/go/test/integration/transforms/xlang/inference/inference.go +++ b/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 } 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..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 @@ -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.>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> 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 +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.create());