diff --git a/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample.py b/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample.py index bdaf170f10..05e4a4d7f6 100644 --- a/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample.py +++ b/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample.py @@ -25,6 +25,7 @@ def create_batch_prediction_job_dedicated_resources_sample( job_display_name: str, gcs_source: Union[str, Sequence[str]], gcs_destination: str, + instances_format: str = "jsonl", machine_type: str = "n1-standard-2", accelerator_count: int = 1, accelerator_type: Union[str, aiplatform_v1.AcceleratorType] = "NVIDIA_TESLA_K80", @@ -40,6 +41,7 @@ def create_batch_prediction_job_dedicated_resources_sample( job_display_name=job_display_name, gcs_source=gcs_source, gcs_destination_prefix=gcs_destination, + instances_format=instances_format, machine_type=machine_type, accelerator_count=accelerator_count, accelerator_type=accelerator_type, diff --git a/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample_test.py b/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample_test.py index 109d368023..dda8abe05f 100644 --- a/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample_test.py +++ b/samples/model-builder/create_batch_prediction_job_dedicated_resources_sample_test.py @@ -31,6 +31,7 @@ def test_create_batch_prediction_job_sample( job_display_name=constants.DISPLAY_NAME, gcs_source=constants.GCS_SOURCES, gcs_destination=constants.GCS_DESTINATION, + instances_format=constants.INSTANCES_FORMAT, machine_type=constants.ACCELERATOR_TYPE, accelerator_count=constants.ACCELERATOR_COUNT, accelerator_type=constants.ACCELERATOR_TYPE, @@ -46,6 +47,7 @@ def test_create_batch_prediction_job_sample( job_display_name=constants.DISPLAY_NAME, gcs_source=constants.GCS_SOURCES, gcs_destination_prefix=constants.GCS_DESTINATION, + instances_format=constants.INSTANCES_FORMAT, machine_type=constants.ACCELERATOR_TYPE, accelerator_count=constants.ACCELERATOR_COUNT, accelerator_type=constants.ACCELERATOR_TYPE, diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py index 196ceed0ce..471b1392a3 100644 --- a/samples/model-builder/test_constants.py +++ b/samples/model-builder/test_constants.py @@ -52,6 +52,7 @@ GCS_SOURCES = ["gs://bucket1/source1.jsonl", "gs://bucket7/source4.jsonl"] BIGQUERY_SOURCE = "bq://bigquery-public-data.ml_datasets.iris" GCS_DESTINATION = "gs://bucket3/output-dir/" +INSTANCES_FORMAT = "jsonl" TRAINING_FRACTION_SPLIT = 0.7 TEST_FRACTION_SPLIT = 0.15