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

* 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 <[email protected]>

Author: TeoZosa

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(

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(

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(

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(