chore: Add sample for custom training job with PSC-I through aiplatform_v1beta1.

PiperOrigin-RevId: 699196687

Author: vertex-sdk-bot

samples/model-builder/conftest.py

@@ -16,6 +16,7 @@
 from unittest.mock import patch
 
 from google.cloud import aiplatform
+from google.cloud import aiplatform_v1beta1
 import vertexai
 from vertexai.resources import preview as preview_resources
 import pytest
@@ -368,6 +369,34 @@ def mock_run_custom_package_training_job(mock_custom_package_training_job):
         yield mock
 
 
+@pytest.fixture
+def mock_job_service_client_v1beta1():
+    mock = MagicMock(aiplatform_v1beta1.JobServiceClient)
+    yield mock
+
+
+@pytest.fixture
+def mock_get_job_service_client_v1beta1(mock_job_service_client_v1beta1):
+    with patch.object(aiplatform_v1beta1, "JobServiceClient") as mock:
+        mock.return_value = mock_job_service_client_v1beta1
+        yield mock
+
+
+@pytest.fixture
+def mock_create_custom_job_v1beta1(mock_job_service_client_v1beta1):
+    with patch.object(
+        mock_job_service_client_v1beta1, "create_custom_job"
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_get_create_custom_job_request_v1beta1():
+    with patch.object(aiplatform_v1beta1, "CreateCustomJobRequest") as mock:
+        mock.return_value = mock_custom_job
+        yield mock
+
+
 @pytest.fixture
 def mock_custom_job():
     mock = MagicMock(aiplatform.CustomJob)

samples/model-builder/create_custom_job_psci_sample.py

@@ -0,0 +1,65 @@
+# Copyright 2024 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
+#
+#     https://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.
+
+#  [START aiplatform_sdk_create_custom_job_psci_sample]
+from google.cloud import aiplatform
+from google.cloud import aiplatform_v1beta1
+
+
+def create_custom_job_psci_sample(
+    project: str,
+    location: str,
+    bucket: str,
+    display_name: str,
+    machine_type: str,
+    replica_count: int,
+    image_uri: str,
+    network_attachment_name: str,
+):
+    """Custom training job sample with PSC-I through aiplatform_v1beta1."""
+    aiplatform.init(project=project, location=location, staging_bucket=bucket)
+
+    client_options = {"api_endpoint": f"{location}-aiplatform.googleapis.com"}
+
+    client = aiplatform_v1beta1.JobServiceClient(client_options=client_options)
+
+    request = aiplatform_v1beta1.CreateCustomJobRequest(
+        parent=f"projects/{project}/locations/{location}",
+        custom_job=aiplatform_v1beta1.CustomJob(
+            display_name=display_name,
+            job_spec=aiplatform_v1beta1.CustomJobSpec(
+                worker_pool_specs=[
+                    aiplatform_v1beta1.WorkerPoolSpec(
+                        machine_spec=aiplatform_v1beta1.MachineSpec(
+                            machine_type=machine_type,
+                        ),
+                        replica_count=replica_count,
+                        container_spec=aiplatform_v1beta1.ContainerSpec(
+                            image_uri=image_uri,
+                        ),
+                    )
+                ],
+                psc_interface_config=aiplatform_v1beta1.PscInterfaceConfig(
+                    network_attachment=network_attachment_name,
+                ),
+            ),
+        ),
+    )
+
+    response = client.create_custom_job(request=request)
+
+    return response
+
+
+#  [END aiplatform_sdk_create_custom_job_psci_sample]

samples/model-builder/create_custom_job_psci_sample_test.py

@@ -0,0 +1,77 @@
+# Copyright 2024 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
+#
+#     https://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 create_custom_job_psci_sample
+from google.cloud import aiplatform_v1beta1
+import test_constants as constants
+
+
+def test_create_custom_job_psci_sample(
+    mock_sdk_init,
+    mock_get_job_service_client_v1beta1,
+    mock_get_create_custom_job_request_v1beta1,
+    mock_create_custom_job_v1beta1,
+):
+    """Custom training job sample with PSC-I through aiplatform_v1beta1."""
+    create_custom_job_psci_sample.create_custom_job_psci_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        bucket=constants.STAGING_BUCKET,
+        display_name=constants.DISPLAY_NAME,
+        machine_type=constants.MACHINE_TYPE,
+        replica_count=1,
+        image_uri=constants.TRAIN_IMAGE,
+        network_attachment_name=constants.NETWORK_ATTACHMENT_NAME,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        staging_bucket=constants.STAGING_BUCKET,
+    )
+
+    mock_get_job_service_client_v1beta1.assert_called_once_with(
+        client_options={
+            "api_endpoint": f"{constants.LOCATION}-aiplatform.googleapis.com"
+        }
+    )
+
+    mock_get_create_custom_job_request_v1beta1.assert_called_once_with(
+        parent=f"projects/{constants.PROJECT}/locations/{constants.LOCATION}",
+        custom_job=aiplatform_v1beta1.CustomJob(
+            display_name=constants.DISPLAY_NAME,
+            job_spec=aiplatform_v1beta1.CustomJobSpec(
+                worker_pool_specs=[
+                    aiplatform_v1beta1.WorkerPoolSpec(
+                        machine_spec=aiplatform_v1beta1.MachineSpec(
+                            machine_type=constants.MACHINE_TYPE,
+                        ),
+                        replica_count=constants.REPLICA_COUNT,
+                        container_spec=aiplatform_v1beta1.ContainerSpec(
+                            image_uri=constants.TRAIN_IMAGE,
+                        ),
+                    )
+                ],
+                psc_interface_config=aiplatform_v1beta1.PscInterfaceConfig(
+                    network_attachment=constants.NETWORK_ATTACHMENT_NAME,
+                ),
+            ),
+        ),
+    )
+
+    request = aiplatform_v1beta1.CreateCustomJobRequest(
+        mock_get_create_custom_job_request_v1beta1.return_value
+    )
+
+    mock_create_custom_job_v1beta1.assert_called_once_with(request=request)

samples/model-builder/test_constants.py

@@ -114,6 +114,7 @@
 MACHINE_TYPE = "n1-standard-4"
 ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED"
 ACCELERATOR_COUNT = 0
+NETWORK_ATTACHMENT_NAME = "network-attachment-name"
 
 # Model constants
 MODEL_RESOURCE_NAME = f"{PARENT}/models/1234"