From f2036736f1c2d450a65e754c725bf1f41ec6a082 Mon Sep 17 00:00:00 2001 From: ivanmkc Date: Mon, 26 Apr 2021 15:28:16 -0400 Subject: [PATCH] Ran linter --- samples/model-builder/conftest.py | 16 ++++++----- ...ng_pipeline_custom_container_job_sample.py | 9 ++++--- ...ning_pipeline_custom_container_job_test.py | 6 +++-- ...ate_training_pipeline_custom_job_sample.py | 11 +++++--- ...reate_training_pipeline_custom_job_test.py | 4 +-- ...ning_pipeline_custom_package_job_sample.py | 9 ++++--- ...aining_pipeline_custom_package_job_test.py | 8 +++--- ..._custom_training_managed_dataset_sample.py | 14 ++++++---- ...ne_custom_training_managed_dataset_test.py | 2 +- samples/model-builder/test_constants.py | 27 +++++++++++-------- 10 files changed, 62 insertions(+), 44 deletions(-) diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index b454644974..169b497859 100644 --- a/samples/model-builder/conftest.py +++ b/samples/model-builder/conftest.py @@ -154,9 +154,7 @@ def mock_run_automl_image_training_job(): @pytest.fixture def mock_init_custom_training_job(): - with patch.object( - aiplatform.training_jobs.CustomTrainingJob, "__init__" - ) as mock: + with patch.object(aiplatform.training_jobs.CustomTrainingJob, "__init__") as mock: mock.return_value = None yield mock @@ -174,11 +172,13 @@ def mock_init_custom_container_training_job(): ) as mock: mock.return_value = None yield mock - + @pytest.fixture def mock_run_custom_container_training_job(): - with patch.object(aiplatform.training_jobs.CustomContainerTrainingJob, "run") as mock: + with patch.object( + aiplatform.training_jobs.CustomContainerTrainingJob, "run" + ) as mock: yield mock @@ -189,11 +189,13 @@ def mock_init_custom_package_training_job(): ) as mock: mock.return_value = None yield mock - + @pytest.fixture def mock_run_custom_package_training_job(): - with patch.object(aiplatform.training_jobs.CustomPythonPackageTrainingJob, "run") as mock: + with patch.object( + aiplatform.training_jobs.CustomPythonPackageTrainingJob, "run" + ) as mock: yield mock diff --git a/samples/model-builder/create_training_pipeline_custom_container_job_sample.py b/samples/model-builder/create_training_pipeline_custom_container_job_sample.py index 0f63cc36f2..56a6494bc9 100644 --- a/samples/model-builder/create_training_pipeline_custom_container_job_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_container_job_sample.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud import aiplatform from typing import List, Optional, Union +from google.cloud import aiplatform + # [START aiplatform_sdk_create_training_pipeline_custom_container_job_sample] def create_training_pipeline_custom_container_job_sample( @@ -22,8 +23,8 @@ def create_training_pipeline_custom_container_job_sample( location: str, staging_bucket: str, display_name: str, - container_uri: str, - model_serving_container_image_uri: str, + container_uri: str, + model_serving_container_image_uri: str, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, replica_count: int = 1, @@ -52,7 +53,7 @@ def create_training_pipeline_custom_container_job_sample( accelerator_count=accelerator_count, training_fraction_split=training_fraction_split, validation_fraction_split=validation_fraction_split, - test_fraction_split=test_fraction_split, + test_fraction_split=test_fraction_split, sync=sync, ) diff --git a/samples/model-builder/create_training_pipeline_custom_container_job_test.py b/samples/model-builder/create_training_pipeline_custom_container_job_test.py index 4d1cc2ac5e..2037b8f74e 100644 --- a/samples/model-builder/create_training_pipeline_custom_container_job_test.py +++ b/samples/model-builder/create_training_pipeline_custom_container_job_test.py @@ -42,9 +42,11 @@ def test_create_training_pipeline_custom_container_job_sample( ) mock_sdk_init.assert_called_once_with( - project=constants.PROJECT, location=constants.LOCATION, staging_bucket=constants.STAGING_BUCKET + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, ) - + mock_init_custom_container_training_job.assert_called_once_with( display_name=constants.DISPLAY_NAME, container_uri=constants.CONTAINER_URI, 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 830ea736b6..cea8d25cde 100644 --- a/samples/model-builder/create_training_pipeline_custom_job_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_job_sample.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud import aiplatform from typing import List, Optional, Union +from google.cloud import aiplatform + # [START aiplatform_sdk_create_training_pipeline_custom_job_sample] def create_training_pipeline_custom_job_sample( @@ -23,7 +24,7 @@ def create_training_pipeline_custom_job_sample( display_name: str, script_path: str, container_uri: str, - model_serving_container_image_uri: str, + model_serving_container_image_uri: str, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, replica_count: int = 0, @@ -37,10 +38,12 @@ def create_training_pipeline_custom_job_sample( ): aiplatform.init(project=project, location=location) - job = aiplatform.CustomTrainingJob(display_name=display_name, + job = aiplatform.CustomTrainingJob( + display_name=display_name, script_path=script_path, container_uri=container_uri, - model_serving_container_image_uri=model_serving_container_image_uri) + model_serving_container_image_uri=model_serving_container_image_uri, + ) model = job.run( model_display_name=model_display_name, 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 3378b782dc..f01fd94a9a 100644 --- a/samples/model-builder/create_training_pipeline_custom_job_test.py +++ b/samples/model-builder/create_training_pipeline_custom_job_test.py @@ -18,9 +18,7 @@ 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_init_custom_training_job, mock_run_custom_training_job, ): create_training_pipeline_custom_job_sample.create_training_pipeline_custom_job_sample( 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 index f27a0bbc3a..627a1b5dab 100644 --- a/samples/model-builder/create_training_pipeline_custom_package_job_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_package_job_sample.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud import aiplatform 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( @@ -24,8 +25,8 @@ def create_training_pipeline_custom_package_job_sample( display_name: str, python_package_gcs_uri: str, python_module_name: str, - container_uri: str, - model_serving_container_image_uri: str, + container_uri: str, + model_serving_container_image_uri: str, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, replica_count: int = 1, @@ -56,7 +57,7 @@ def create_training_pipeline_custom_package_job_sample( accelerator_count=accelerator_count, training_fraction_split=training_fraction_split, validation_fraction_split=validation_fraction_split, - test_fraction_split=test_fraction_split, + test_fraction_split=test_fraction_split, sync=sync, ) 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 index fdb3bb3dae..cf9daf0ac4 100644 --- a/samples/model-builder/create_training_pipeline_custom_package_job_test.py +++ b/samples/model-builder/create_training_pipeline_custom_package_job_test.py @@ -29,7 +29,7 @@ def test_create_training_pipeline_custom_package_job_sample( 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, + python_module_name=constants.PYTHON_MODULE_NAME, container_uri=constants.CONTAINER_URI, args=constants.ARGS, model_serving_container_image_uri=constants.CONTAINER_URI, @@ -44,9 +44,11 @@ def test_create_training_pipeline_custom_package_job_sample( ) mock_sdk_init.assert_called_once_with( - project=constants.PROJECT, location=constants.LOCATION, staging_bucket=constants.STAGING_BUCKET + project=constants.PROJECT, + location=constants.LOCATION, + staging_bucket=constants.STAGING_BUCKET, ) - + mock_init_custom_package_training_job.assert_called_once_with( display_name=constants.DISPLAY_NAME, python_package_gcs_uri=constants.PYTHON_PACKAGE_GCS_URI, diff --git a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py b/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py index 89b6f5bcb2..7d7dc6357c 100644 --- a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py +++ b/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud import aiplatform from typing import List, Optional, Union +from google.cloud import aiplatform + + # [START aiplatform_sdk_create_training_pipeline_custom_job_sample] def create_training_pipeline_custom_training_managed_dataset_sample( project: str, @@ -23,7 +25,7 @@ def create_training_pipeline_custom_training_managed_dataset_sample( script_path: str, container_uri: str, model_serving_container_image_uri: str, - dataset_id: int, + dataset_id: int, model_display_name: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, replica_count: int = 0, @@ -37,10 +39,12 @@ def create_training_pipeline_custom_training_managed_dataset_sample( ): aiplatform.init(project=project, location=location) - job = aiplatform.CustomTrainingJob(display_name=display_name, + job = aiplatform.CustomTrainingJob( + display_name=display_name, script_path=script_path, container_uri=container_uri, - model_serving_container_image_uri=model_serving_container_image_uri) + model_serving_container_image_uri=model_serving_container_image_uri, + ) my_image_ds = aiplatform.ImageDataset(dataset_id) @@ -54,7 +58,7 @@ def create_training_pipeline_custom_training_managed_dataset_sample( accelerator_count=accelerator_count, training_fraction_split=training_fraction_split, validation_fraction_split=validation_fraction_split, - test_fraction_split=test_fraction_split, + test_fraction_split=test_fraction_split, sync=sync, ) diff --git a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py b/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py index a97e84ed43..4197f658b1 100644 --- a/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py +++ b/samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py @@ -62,7 +62,7 @@ def test_create_training_pipeline_custom_job_sample( replica_count=constants.REPLICA_COUNT, machine_type=constants.MACHINE_TYPE, accelerator_type=constants.ACCELERATOR_TYPE, - accelerator_count=constants.ACCELERATOR_COUNT, + 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, diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py index bfe9c607c0..003e6ab22d 100644 --- a/samples/model-builder/test_constants.py +++ b/samples/model-builder/test_constants.py @@ -60,14 +60,19 @@ MACHINE_TYPE = "n1-standard-4" ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED" ACCELERATOR_COUNT = 0 -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'], -MODEL_SERVING_CONTAINER_PREDICT_ROUTE=f'/v1/models/{MODEL_NAME}:predict', -MODEL_SERVING_CONTAINER_HEALTH_ROUTE=f'/v1/models/{MODEL_NAME}' -PYTHON_PACKAGE_GCS_URI="gs://bucket3/custom-training-python-package/my_app/trainer-0.1.tar.gz" -PYTHON_MODULE_NAME="trainer.task" \ No newline at end of file +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", + ], +) +MODEL_SERVING_CONTAINER_PREDICT_ROUTE = (f"/v1/models/{MODEL_NAME}:predict",) +MODEL_SERVING_CONTAINER_HEALTH_ROUTE = f"/v1/models/{MODEL_NAME}" +PYTHON_PACKAGE_GCS_URI = ( + "gs://bucket3/custom-training-python-package/my_app/trainer-0.1.tar.gz" +) +PYTHON_MODULE_NAME = "trainer.task"