diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index 70431c9565..01756f668b 100644 --- a/samples/model-builder/conftest.py +++ b/samples/model-builder/conftest.py @@ -147,6 +147,18 @@ def mock_custom_training_job(): yield mock +@pytest.fixture +def mock_custom_container_training_job(): + mock = MagicMock(aiplatform.training_jobs.CustomContainerTrainingJob) + yield mock + + +@pytest.fixture +def mock_custom_package_training_job(): + mock = MagicMock(aiplatform.training_jobs.CustomPythonPackageTrainingJob) + yield mock + + @pytest.fixture def mock_image_training_job(): mock = MagicMock(aiplatform.training_jobs.AutoMLImageTrainingJob) @@ -198,15 +210,41 @@ def mock_run_automl_image_training_job(mock_image_training_job): @pytest.fixture -def mock_init_custom_training_job(): - with patch.object(aiplatform.training_jobs.CustomTrainingJob, "__init__") as mock: - mock.return_value = None +def mock_get_custom_training_job(mock_custom_training_job): + with patch.object(aiplatform, "CustomTrainingJob") as mock: + mock.return_value = mock_custom_training_job + yield mock + + +@pytest.fixture +def mock_get_custom_container_training_job(mock_custom_container_training_job): + with patch.object(aiplatform, "CustomContainerTrainingJob") as mock: + mock.return_value = mock_custom_container_training_job + yield mock + + +@pytest.fixture +def mock_get_custom_package_training_job(mock_custom_package_training_job): + with patch.object(aiplatform, "CustomPythonPackageTrainingJob") as mock: + mock.return_value = mock_custom_package_training_job + yield mock + + +@pytest.fixture +def mock_run_custom_training_job(mock_custom_training_job): + with patch.object(mock_custom_training_job, "run") as mock: + yield mock + + +@pytest.fixture +def mock_run_custom_container_training_job(mock_custom_container_training_job): + with patch.object(mock_custom_container_training_job, "run") as mock: yield mock @pytest.fixture -def mock_run_custom_training_job(): - with patch.object(aiplatform.training_jobs.CustomTrainingJob, "run") as mock: +def mock_run_custom_package_training_job(mock_custom_package_training_job): + with patch.object(mock_custom_package_training_job, "run") as mock: yield mock diff --git a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py b/samples/model-builder/create_training_pipeline_custom_container_job_sample.py similarity index 75% rename from samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py rename to samples/model-builder/create_training_pipeline_custom_container_job_sample.py index 7d7dc6357c..1c4b717487 100644 --- a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_container_job_sample.py @@ -17,18 +17,18 @@ from google.cloud import aiplatform -# [START aiplatform_sdk_create_training_pipeline_custom_job_sample] -def create_training_pipeline_custom_training_managed_dataset_sample( +# [START aiplatform_sdk_create_training_pipeline_custom_container_job_sample] +def create_training_pipeline_custom_container_job_sample( project: str, location: str, + staging_bucket: str, display_name: str, - script_path: str, container_uri: str, model_serving_container_image_uri: str, - dataset_id: int, + dataset_id: Optional[str] = None, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, - replica_count: int = 0, + replica_count: int = 1, machine_type: str = "n1-standard-4", accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED", accelerator_count: int = 0, @@ -37,19 +37,19 @@ def create_training_pipeline_custom_training_managed_dataset_sample( test_fraction_split: float = 0.1, sync: bool = True, ): - aiplatform.init(project=project, location=location) + aiplatform.init(project=project, location=location, staging_bucket=staging_bucket) - job = aiplatform.CustomTrainingJob( + job = aiplatform.CustomContainerTrainingJob( display_name=display_name, - script_path=script_path, container_uri=container_uri, model_serving_container_image_uri=model_serving_container_image_uri, ) - my_image_ds = aiplatform.ImageDataset(dataset_id) + # This example uses an ImageDataset, but you can use another type + dataset = aiplatform.ImageDataset(dataset_id) if dataset_id else None model = job.run( - dataset=my_image_ds, + dataset=dataset, model_display_name=model_display_name, args=args, replica_count=replica_count, @@ -70,4 +70,4 @@ def create_training_pipeline_custom_training_managed_dataset_sample( return model -# [END aiplatform_sdk_create_training_pipeline_custom_job_sample] +# [END aiplatform_sdk_create_training_pipeline_custom_container_job_sample] diff --git a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py b/samples/model-builder/create_training_pipeline_custom_container_job_test.py similarity index 76% rename from samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py rename to samples/model-builder/create_training_pipeline_custom_container_job_test.py index 4197f658b1..3762c931a8 100644 --- a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py +++ b/samples/model-builder/create_training_pipeline_custom_container_job_test.py @@ -13,24 +13,24 @@ # limitations under the License. -import create_training_pipeline_custom_training_managed_dataset_sample +import create_training_pipeline_custom_container_job_sample import test_constants as constants -def test_create_training_pipeline_custom_job_sample( +def test_create_training_pipeline_custom_container_job_sample( mock_sdk_init, mock_image_dataset, - mock_init_custom_training_job, - mock_run_custom_training_job, mock_get_image_dataset, + mock_get_custom_container_training_job, + mock_run_custom_container_training_job, ): - create_training_pipeline_custom_training_managed_dataset_sample.create_training_pipeline_custom_training_managed_dataset_sample( + create_training_pipeline_custom_container_job_sample.create_training_pipeline_custom_container_job_sample( project=constants.PROJECT, location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, display_name=constants.DISPLAY_NAME, args=constants.ARGS, - script_path=constants.SCRIPT_PATH, container_uri=constants.CONTAINER_URI, model_serving_container_image_uri=constants.CONTAINER_URI, dataset_id=constants.RESOURCE_ID, @@ -44,25 +44,25 @@ def test_create_training_pipeline_custom_job_sample( test_fraction_split=constants.TEST_FRACTION_SPLIT, ) - mock_get_image_dataset.assert_called_once_with(constants.RESOURCE_ID) - mock_sdk_init.assert_called_once_with( - project=constants.PROJECT, location=constants.LOCATION + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, ) - mock_init_custom_training_job.assert_called_once_with( + + mock_get_custom_container_training_job.assert_called_once_with( display_name=constants.DISPLAY_NAME, - script_path=constants.SCRIPT_PATH, container_uri=constants.CONTAINER_URI, model_serving_container_image_uri=constants.CONTAINER_URI, ) - mock_run_custom_training_job.assert_called_once_with( + mock_run_custom_container_training_job.assert_called_once_with( dataset=mock_image_dataset, model_display_name=constants.DISPLAY_NAME_2, - args=constants.ARGS, replica_count=constants.REPLICA_COUNT, machine_type=constants.MACHINE_TYPE, accelerator_type=constants.ACCELERATOR_TYPE, accelerator_count=constants.ACCELERATOR_COUNT, + args=constants.ARGS, training_fraction_split=constants.TRAINING_FRACTION_SPLIT, validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT, test_fraction_split=constants.TEST_FRACTION_SPLIT, diff --git a/samples/model-builder/create_training_pipeline_custom_job_sample.py b/samples/model-builder/create_training_pipeline_custom_job_sample.py index cea8d25cde..175bee8a6e 100644 --- a/samples/model-builder/create_training_pipeline_custom_job_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_job_sample.py @@ -21,10 +21,12 @@ def create_training_pipeline_custom_job_sample( project: str, location: str, + staging_bucket: str, display_name: str, script_path: str, container_uri: str, model_serving_container_image_uri: str, + dataset_id: Optional[str] = None, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, replica_count: int = 0, @@ -36,7 +38,7 @@ def create_training_pipeline_custom_job_sample( test_fraction_split: float = 0.1, sync: bool = True, ): - aiplatform.init(project=project, location=location) + aiplatform.init(project=project, location=location, staging_bucket=staging_bucket) job = aiplatform.CustomTrainingJob( display_name=display_name, @@ -45,7 +47,11 @@ def create_training_pipeline_custom_job_sample( model_serving_container_image_uri=model_serving_container_image_uri, ) + # This example uses an ImageDataset, but you can use another type + dataset = aiplatform.ImageDataset(dataset_id) if dataset_id else None + model = job.run( + dataset=dataset, model_display_name=model_display_name, args=args, replica_count=replica_count, diff --git a/samples/model-builder/create_training_pipeline_custom_job_test.py b/samples/model-builder/create_training_pipeline_custom_job_test.py index f01fd94a9a..b1d35b7360 100644 --- a/samples/model-builder/create_training_pipeline_custom_job_test.py +++ b/samples/model-builder/create_training_pipeline_custom_job_test.py @@ -18,17 +18,23 @@ def test_create_training_pipeline_custom_job_sample( - mock_sdk_init, mock_init_custom_training_job, mock_run_custom_training_job, + mock_sdk_init, + mock_image_dataset, + mock_get_image_dataset, + mock_get_custom_training_job, + mock_run_custom_training_job, ): create_training_pipeline_custom_job_sample.create_training_pipeline_custom_job_sample( project=constants.PROJECT, location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, display_name=constants.DISPLAY_NAME, args=constants.ARGS, script_path=constants.SCRIPT_PATH, container_uri=constants.CONTAINER_URI, model_serving_container_image_uri=constants.CONTAINER_URI, + dataset_id=constants.RESOURCE_ID, model_display_name=constants.DISPLAY_NAME_2, replica_count=constants.REPLICA_COUNT, machine_type=constants.MACHINE_TYPE, @@ -40,15 +46,18 @@ def test_create_training_pipeline_custom_job_sample( ) mock_sdk_init.assert_called_once_with( - project=constants.PROJECT, location=constants.LOCATION + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, ) - mock_init_custom_training_job.assert_called_once_with( + mock_get_custom_training_job.assert_called_once_with( display_name=constants.DISPLAY_NAME, script_path=constants.SCRIPT_PATH, container_uri=constants.CONTAINER_URI, model_serving_container_image_uri=constants.CONTAINER_URI, ) mock_run_custom_training_job.assert_called_once_with( + dataset=mock_image_dataset, model_display_name=constants.DISPLAY_NAME_2, replica_count=constants.REPLICA_COUNT, machine_type=constants.MACHINE_TYPE, diff --git a/samples/model-builder/create_training_pipeline_custom_package_job_sample.py b/samples/model-builder/create_training_pipeline_custom_package_job_sample.py new file mode 100644 index 0000000000..9550d91acb --- /dev/null +++ b/samples/model-builder/create_training_pipeline_custom_package_job_sample.py @@ -0,0 +1,77 @@ +# Copyright 2021 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. + +from typing import List, Optional, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_create_training_pipeline_custom_package_job_sample] +def create_training_pipeline_custom_package_job_sample( + project: str, + location: str, + staging_bucket: str, + display_name: str, + python_package_gcs_uri: str, + python_module_name: str, + container_uri: str, + model_serving_container_image_uri: str, + dataset_id: Optional[str] = None, + model_display_name: Optional[str] = None, + args: Optional[List[Union[str, float, int]]] = None, + replica_count: int = 1, + machine_type: str = "n1-standard-4", + accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED", + accelerator_count: int = 0, + training_fraction_split: float = 0.8, + validation_fraction_split: float = 0.1, + test_fraction_split: float = 0.1, + sync: bool = True, +): + aiplatform.init(project=project, location=location, staging_bucket=staging_bucket) + + job = aiplatform.CustomPythonPackageTrainingJob( + display_name=display_name, + python_package_gcs_uri=python_package_gcs_uri, + python_module_name=python_module_name, + container_uri=container_uri, + model_serving_container_image_uri=model_serving_container_image_uri, + ) + + # This example uses an ImageDataset, but you can use another type + dataset = aiplatform.ImageDataset(dataset_id) if dataset_id else None + + model = job.run( + dataset=dataset, + model_display_name=model_display_name, + args=args, + replica_count=replica_count, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + training_fraction_split=training_fraction_split, + validation_fraction_split=validation_fraction_split, + test_fraction_split=test_fraction_split, + sync=sync, + ) + + model.wait() + + print(model.display_name) + print(model.resource_name) + print(model.uri) + return model + + +# [END aiplatform_sdk_create_training_pipeline_custom_package_job_sample] diff --git a/samples/model-builder/create_training_pipeline_custom_package_job_test.py b/samples/model-builder/create_training_pipeline_custom_package_job_test.py new file mode 100644 index 0000000000..1156adc8b3 --- /dev/null +++ b/samples/model-builder/create_training_pipeline_custom_package_job_test.py @@ -0,0 +1,75 @@ +# Copyright 2021 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_training_pipeline_custom_package_job_sample +import test_constants as constants + + +def test_create_training_pipeline_custom_package_job_sample( + mock_sdk_init, + mock_image_dataset, + mock_get_image_dataset, + mock_get_custom_package_training_job, + mock_run_custom_package_training_job, +): + + create_training_pipeline_custom_package_job_sample.create_training_pipeline_custom_package_job_sample( + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, + display_name=constants.DISPLAY_NAME, + python_package_gcs_uri=constants.PYTHON_PACKAGE_GCS_URI, + python_module_name=constants.PYTHON_MODULE_NAME, + container_uri=constants.CONTAINER_URI, + args=constants.ARGS, + model_serving_container_image_uri=constants.CONTAINER_URI, + dataset_id=constants.RESOURCE_ID, + model_display_name=constants.DISPLAY_NAME_2, + replica_count=constants.REPLICA_COUNT, + machine_type=constants.MACHINE_TYPE, + accelerator_type=constants.ACCELERATOR_TYPE, + accelerator_count=constants.ACCELERATOR_COUNT, + training_fraction_split=constants.TRAINING_FRACTION_SPLIT, + validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT, + test_fraction_split=constants.TEST_FRACTION_SPLIT, + ) + + mock_sdk_init.assert_called_once_with( + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, + ) + + mock_get_custom_package_training_job.assert_called_once_with( + display_name=constants.DISPLAY_NAME, + python_package_gcs_uri=constants.PYTHON_PACKAGE_GCS_URI, + python_module_name=constants.PYTHON_MODULE_NAME, + container_uri=constants.CONTAINER_URI, + model_serving_container_image_uri=constants.CONTAINER_URI, + ) + + mock_run_custom_package_training_job.assert_called_once_with( + dataset=mock_image_dataset, + model_display_name=constants.DISPLAY_NAME_2, + replica_count=constants.REPLICA_COUNT, + machine_type=constants.MACHINE_TYPE, + accelerator_type=constants.ACCELERATOR_TYPE, + accelerator_count=constants.ACCELERATOR_COUNT, + args=constants.ARGS, + training_fraction_split=constants.TRAINING_FRACTION_SPLIT, + validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT, + test_fraction_split=constants.TEST_FRACTION_SPLIT, + sync=True, + ) diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py index 641fa1c490..69da01dbd8 100644 --- a/samples/model-builder/test_constants.py +++ b/samples/model-builder/test_constants.py @@ -94,7 +94,7 @@ SCRIPT_PATH = "task.py" CONTAINER_URI = "gcr.io/my_project/my_image:latest" ARGS = ["--tfds", "tf_flowers:3.*.*"] -REPLICA_COUNT = 0 +REPLICA_COUNT = 1 MACHINE_TYPE = "n1-standard-4" ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED" ACCELERATOR_COUNT = 0 @@ -177,3 +177,17 @@ "households": "270.0", "median_income": "3.014700", } +MODEL_SERVING_CONTAINER_COMMAND = (["/usr/bin/tensorflow_model_server"],) +MODEL_SERVING_CONTAINER_ARGS = ( + [ + f"--model_name={MODEL_NAME}", + "--model_base_path=$(AIP_STORAGE_URI)", + "--rest_api_port=8080", + "--port=8500", + "--file_system_poll_wait_seconds=31540000", + ], +) +PYTHON_PACKAGE_GCS_URI = ( + "gs://bucket3/custom-training-python-package/my_app/trainer-0.1.tar.gz" +) +PYTHON_MODULE_NAME = "trainer.task"