feat: Support custom predictor Docker image builds on non-x86 architectures (#2115)

TeoZosa · abcdefgs0324 · web-flow · commit 87dd5c0c2ebb · 2025-03-06T13:03:38.000-08:00
* Enforce Linux `x86_64` Docker image builds Under the assumption that only machine types with x86 processors are supported for prediction and custom training. * Extract Docker build platform arg to method parameter While currently only x86 processors are supported for prediction and custom training, this will allow users to control this behavior should that ever change in the future. Additionally, it allows users to, e.g., override the `TARGETOS` component of the `TARGETPLATFORM`. * Change platform default arg to `None` To enforce the flag is set by users as opposed to providing a universal default. - See: #2115 (comment) * Make platform configurable from `LocalModel.build_cpr_model()` To enable the flag to be set by users (e.g., to build images on non-x86 architectures). * Fix docstring for `platform` param Resolves: - #2115 (comment) - #2115 (comment) * Test platform parameter in `test_build_cpr_model_upload_and_deploy()` Resolves: - #2115 (review) * Fix tests Resolves (partially): - #2115 (comment) * Test specifying platform in local model builds Resolves (partially): - #2115 (comment) * Test other platform strings in local model builds --------- Co-authored-by: Chun-Hsiang Wang <chunhsiang@google.com>
diff --git a/google/cloud/aiplatform/docker_utils/build.py b/google/cloud/aiplatform/docker_utils/build.py
@@ -418,6 +418,7 @@ def build_image(
     pip_command: str = "pip",
     python_command: str = "python",
     no_cache: bool = True,
+    platform: Optional[str] = None,
     **kwargs,
 ) -> Image:
     """Builds a Docker image.
@@ -459,6 +460,10 @@ def build_image(
             reduces the image building time. See
             https://docs.docker.com/develop/develop-images/dockerfile_best-practices/#leverage-build-cache
             for more details.
+        platform (str):
+            Optional. The target platform for the Docker image build. See
+            https://docs.docker.com/build/building/multi-platform/#building-multi-platform-images
+            for more details.
         **kwargs:
             Other arguments to pass to underlying method that generates the Dockerfile.
 
@@ -472,9 +477,14 @@ def build_image(
 
     tag_options = ["-t", output_image_name]
     cache_args = ["--no-cache"] if no_cache else []
+    platform_args = ["--platform", platform] if platform is not None else []
 
     command = (
-        ["docker", "build"] + cache_args + tag_options + ["--rm", "-f-", host_workdir]
+        ["docker", "build"]
+        + cache_args
+        + platform_args
+        + tag_options
+        + ["--rm", "-f-", host_workdir]
     )
 
     requirements_relative_path = _get_relative_path_to_workdir(
diff --git a/google/cloud/aiplatform/prediction/local_model.py b/google/cloud/aiplatform/prediction/local_model.py
@@ -246,6 +246,7 @@ def build_cpr_model(
         requirements_path: Optional[str] = None,
         extra_packages: Optional[List[str]] = None,
         no_cache: bool = False,
+        platform: Optional[str] = None,
     ) -> "LocalModel":
         """Builds a local model from a custom predictor.
 
@@ -274,6 +275,7 @@ def build_cpr_model(
                 predictor=$CUSTOM_PREDICTOR_CLASS,
                 requirements_path="./user_src_dir/requirements.txt",
                 extra_packages=["./user_src_dir/user_code/custom_package.tar.gz"],
+                platform="linux/amd64", # i.e., if you're building on a non-x86 machine
             )
 
         In the built image, user provided files will be copied as follows:
@@ -340,6 +342,10 @@ def build_cpr_model(
                 reduces the image building time. See
                 https://docs.docker.com/develop/develop-images/dockerfile_best-practices/#leverage-build-cache
                 for more details.
+            platform (str):
+                Optional. The target platform for the Docker image build. See
+                https://docs.docker.com/build/building/multi-platform/#building-multi-platform-images
+                for more details.
 
         Returns:
             local model: Instantiated representation of the local model.
@@ -391,6 +397,7 @@ def build_cpr_model(
             pip_command="pip3" if is_prebuilt_prediction_image else "pip",
             python_command="python3" if is_prebuilt_prediction_image else "python",
             no_cache=no_cache,
+            platform=platform,
         )
 
         container_spec = gca_model_compat.ModelContainerSpec(
diff --git a/tests/system/aiplatform/test_prediction_cpr.py b/tests/system/aiplatform/test_prediction_cpr.py
@@ -49,7 +49,8 @@ class TestPredictionCpr(e2e_base.TestEndToEnd):
 
     _temp_prefix = "temp-vertex-sdk-e2e-prediction-cpr"
 
-    def test_build_cpr_model_upload_and_deploy(self, shared_state, caplog):
+    @pytest.mark.parametrize("platform", [None, "linux/amd64"])
+    def test_build_cpr_model_upload_and_deploy(self, shared_state, caplog, platform):
         """Creates a CPR model from custom predictor, uploads it and deploys."""
 
         caplog.set_level(logging.INFO)
@@ -61,6 +62,7 @@ def test_build_cpr_model_upload_and_deploy(self, shared_state, caplog):
             _IMAGE_URI,
             predictor=SklearnPredictor,
             requirements_path=os.path.join(_USER_CODE_DIR, _REQUIREMENTS_FILE),
+            platform=platform,
         )
 
         with local_model.deploy_to_local_endpoint(
diff --git a/tests/unit/aiplatform/test_prediction.py b/tests/unit/aiplatform/test_prediction.py
@@ -1304,6 +1304,7 @@ class {predictor_class}:
             pip_command="pip",
             python_command="python",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_fails_handler_is_none(
@@ -1418,6 +1419,7 @@ class {handler_class}:
             pip_command="pip",
             python_command="python",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_with_custom_handler_and_predictor_is_none(
@@ -1472,6 +1474,7 @@ class {handler_class}:
             pip_command="pip",
             python_command="python",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_creates_and_get_localmodel_base_is_prebuilt(
@@ -1527,6 +1530,7 @@ class {predictor_class}:
             pip_command="pip3",
             python_command="python3",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_creates_and_get_localmodel_with_requirements_path(
@@ -1584,6 +1588,7 @@ class {predictor_class}:
             pip_command="pip",
             python_command="python",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_creates_and_get_localmodel_with_extra_packages(
@@ -1641,6 +1646,7 @@ class {predictor_class}:
             pip_command="pip",
             python_command="python",
             no_cache=False,
+            platform=None,
         )
 
     def test_build_cpr_model_creates_and_get_localmodel_no_cache(
@@ -1695,6 +1701,70 @@ class {predictor_class}:
             pip_command="pip",
             python_command="python",
             no_cache=no_cache,
+            platform=None,
+        )
+
+    @pytest.mark.parametrize(
+        "platform",
+        [
+            None,
+            "linux/amd64",
+            "some_arbitrary_platform_value_that_will_by_validated_by_docker_build_command",
+        ],
+    )
+    def test_build_cpr_model_creates_and_get_localmodel_platform(
+        self,
+        tmp_path,
+        inspect_source_from_class_mock_predictor_only,
+        is_prebuilt_prediction_container_uri_is_false_mock,
+        build_image_mock,
+        platform,
+    ):
+        src_dir = tmp_path / _TEST_SRC_DIR
+        src_dir.mkdir()
+        predictor = src_dir / _TEST_PREDICTOR_FILE
+        predictor.write_text(
+            textwrap.dedent(
+                """
+                class {predictor_class}:
+                    pass
+                """
+            ).format(predictor_class=_TEST_PREDICTOR_CLASS)
+        )
+        my_predictor = self._load_module(_TEST_PREDICTOR_CLASS, str(predictor))
+
+        local_model = LocalModel.build_cpr_model(
+            str(src_dir), _TEST_OUTPUT_IMAGE, predictor=my_predictor, platform=platform
+        )
+
+        assert local_model.serving_container_spec.image_uri == _TEST_OUTPUT_IMAGE
+        assert local_model.serving_container_spec.predict_route == DEFAULT_PREDICT_ROUTE
+        assert local_model.serving_container_spec.health_route == DEFAULT_HEALTH_ROUTE
+        inspect_source_from_class_mock_predictor_only.assert_called_once_with(
+            my_predictor, str(src_dir)
+        )
+        is_prebuilt_prediction_container_uri_is_false_mock.assert_called_once_with(
+            _DEFAULT_BASE_IMAGE
+        )
+        build_image_mock.assert_called_once_with(
+            _DEFAULT_BASE_IMAGE,
+            str(src_dir),
+            _TEST_OUTPUT_IMAGE,
+            python_module=_DEFAULT_PYTHON_MODULE,
+            requirements_path=None,
+            extra_requirements=_DEFAULT_SDK_REQUIREMENTS,
+            extra_packages=None,
+            exposed_ports=[DEFAULT_HTTP_PORT],
+            environment_variables={
+                "HANDLER_MODULE": _DEFAULT_HANDLER_MODULE,
+                "HANDLER_CLASS": _DEFAULT_HANDLER_CLASS,
+                "PREDICTOR_MODULE": f"{_TEST_SRC_DIR}.{_TEST_PREDICTOR_FILE_STEM}",
+                "PREDICTOR_CLASS": _TEST_PREDICTOR_CLASS,
+            },
+            pip_command="pip",
+            python_command="python",
+            no_cache=False,
+            platform=platform,
         )
 
     def test_deploy_to_local_endpoint(
