Skip to content

Commit

Permalink
Ran linter
Browse files Browse the repository at this point in the history
  • Loading branch information
@ivanmkc
ivanmkc committed Apr 26, 2021
1 parent 70c57d8 commit f203673
Show file tree
Hide file tree
Showing 10 changed files with 62 additions and 44 deletions.
16 changes: 9 additions & 7 deletions samples/model-builder/conftest.py
View file
Expand Up @@ -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

Expand All @@ -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


Expand All @@ -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


Expand Down
9 changes: 5 additions & 4 deletions samples/model-builder/create_training_pipeline_custom_container_job_sample.py
View file
Expand Up @@ -12,18 +12,19 @@
# 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(
project: str,
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,
Expand Down Expand Up @@ -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,
)

Expand Down
6 changes: 4 additions & 2 deletions samples/model-builder/create_training_pipeline_custom_container_job_test.py
View file
Expand Up @@ -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,
Expand Down
11 changes: 7 additions & 4 deletions samples/model-builder/create_training_pipeline_custom_job_sample.py
View file
Expand Up @@ -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(
Expand All @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 1 addition & 3 deletions samples/model-builder/create_training_pipeline_custom_job_test.py
View file
Expand Up @@ -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(
Expand Down
9 changes: 5 additions & 4 deletions samples/model-builder/create_training_pipeline_custom_package_job_sample.py
View file
Expand Up @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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,
)

Expand Down
8 changes: 5 additions & 3 deletions samples/model-builder/create_training_pipeline_custom_package_job_test.py
View file
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
14 changes: 9 additions & 5 deletions samples/model-builder/create_training_pipeline_custom_training_managed_dataset_sample.py
View file
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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)

Expand All @@ -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,
)

Expand Down
2 changes: 1 addition & 1 deletion samples/model-builder/create_training_pipeline_custom_training_managed_dataset_test.py
View file
Expand Up @@ -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,
Expand Down
27 changes: 16 additions & 11 deletions samples/model-builder/test_constants.py
View file
Expand Up @@ -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"
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"

0 comments on commit f203673

Please sign in to comment.