feat: Support custom containers in CustomJob.from_local_script (#1483)

* feat: update from_local_script()

* fix from_local_script()

* fix: update unit tests

* fix lint failed

* fix: add system test

* fix: from_local_script

* fix: from_local_script

* fix: system test

google/cloud/aiplatform/jobs.py

@@ -1265,7 +1265,12 @@ def from_local_script(
             script_path (str):
                 Required. Local path to training script.
             container_uri (str):
-                Required: Uri of the training container image to use for custom job.
+                Required. Uri of the training container image to use for custom job.
+                Support images in Artifact Registry, Container Registry, or Docker Hub.
+                Vertex AI provides a wide range of executor images with pre-installed
+                packages to meet users' various use cases. See the list of `pre-built containers
+                for training <https://cloud.google.com/vertex-ai/docs/training/pre-built-containers>`.
+                If not using image from this list, please make sure python3 and pip3 are installed in your container.
             args (Optional[Sequence[str]]):
                 Optional. Command line arguments to be passed to the Python task.
             requirements (Sequence[str]):
@@ -1388,7 +1393,10 @@ def from_local_script(
                 spec["container_spec"] = {
                     "image_uri": reduction_server_container_uri,
                 }
-            else:
+            ## check if the container is pre-built
+            elif ("docker.pkg.dev/vertex-ai/" in container_uri) or (
+                "gcr.io/cloud-aiplatform/" in container_uri
+            ):
                 spec["python_package_spec"] = {
                     "executor_image_uri": container_uri,
                     "python_module": python_packager.module_name,
@@ -1403,6 +1411,30 @@ def from_local_script(
                         {"name": key, "value": value}
                         for key, value in environment_variables.items()
                     ]
+            else:
+                command = [
+                    "sh",
+                    "-c",
+                    "pip install --upgrade pip && "
+                    + f"pip3 install -q --user {package_gcs_uri} && ".replace(
+                        "gs://", "/gcs/"
+                    )
+                    + f"python3 -m {python_packager.module_name}",
+                ]
+
+                spec["container_spec"] = {
+                    "image_uri": container_uri,
+                    "command": command,
+                }
+
+                if args:
+                    spec["container_spec"]["args"] = args
+
+                if environment_variables:
+                    spec["container_spec"]["env"] = [
+                        {"name": key, "value": value}
+                        for key, value in environment_variables.items()
+                    ]
 
         return cls(
             display_name=display_name,

tests/system/aiplatform/test_custom_job.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+
+import pytest
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform.compat.types import job_state as gca_job_state
+from tests.system.aiplatform import e2e_base
+
+_PREBUILT_CONTAINER_IMAGE = "gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest"
+_CUSTOM_CONTAINER_IMAGE = "python:3.8"
+
+_DIR_NAME = os.path.dirname(os.path.abspath(__file__))
+_LOCAL_TRAINING_SCRIPT_PATH = os.path.join(
+    _DIR_NAME, "test_resources/custom_job_script.py"
+)
+
+
+@pytest.mark.usefixtures(
+    "prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources"
+)
+class TestCustomJob(e2e_base.TestEndToEnd):
+
+    _temp_prefix = "temp-vertex-sdk-custom-job"
+
+    def test_from_local_script_prebuilt_container(self, shared_state):
+        shared_state["resources"] = []
+
+        aiplatform.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+            staging_bucket=shared_state["staging_bucket_name"],
+        )
+
+        display_name = self._make_display_name("custom-job")
+
+        custom_job = aiplatform.CustomJob.from_local_script(
+            display_name=display_name,
+            script_path=_LOCAL_TRAINING_SCRIPT_PATH,
+            container_uri=_PREBUILT_CONTAINER_IMAGE,
+        )
+        custom_job.run()
+
+        shared_state["resources"].append(custom_job)
+
+        assert custom_job.state == gca_job_state.JobState.JOB_STATE_SUCCEEDED
+
+    def test_from_local_script_custom_container(self, shared_state):
+
+        aiplatform.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+            staging_bucket=shared_state["staging_bucket_name"],
+        )
+
+        display_name = self._make_display_name("custom-job")
+
+        custom_job = aiplatform.CustomJob.from_local_script(
+            display_name=display_name,
+            script_path=_LOCAL_TRAINING_SCRIPT_PATH,
+            container_uri=_CUSTOM_CONTAINER_IMAGE,
+        )
+        custom_job.run()
+
+        shared_state["resources"].append(custom_job)
+
+        assert custom_job.state == gca_job_state.JobState.JOB_STATE_SUCCEEDED

tests/system/aiplatform/test_resources/custom_job_script.py

@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+print("Test CustomJob script.")