diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index 97f20177a5..8d8583f850 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -4514,7 +4514,7 @@ def run( training_filter_split: Optional[str] = None, validation_filter_split: Optional[str] = None, test_filter_split: Optional[str] = None, - budget_milli_node_hours: int = 1000, + budget_milli_node_hours: Optional[int] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, disable_early_stopping: bool = False, @@ -4580,18 +4580,26 @@ def run( single DataItem is matched by more than one of the FilterSplit filters, then it is assigned to the first set that applies to it in the training, validation, test order. This is ignored if Dataset is not provided. - budget_milli_node_hours: int = 1000 + budget_milli_node_hours (int): Optional. The train budget of creating this Model, expressed in milli node hours i.e. 1,000 value in this field means 1 node hour. + + Defaults by `prediction_type`: + + `classification` - For Cloud models the budget must be: 8,000 - 800,000 + milli node hours (inclusive). The default value is 192,000 which + represents one day in wall time, assuming 8 nodes are used. + `object_detection` - For Cloud models the budget must be: 20,000 - 900,000 + milli node hours (inclusive). The default value is 216,000 which represents + one day in wall time, assuming 9 nodes are used. + The training cost of the model will not exceed this budget. The final cost will be attempted to be close to the budget, though may end up being (even) noticeably smaller - at the backend's discretion. This especially may happen when further model training ceases to provide - any improvements. - If the budget is set to a value known to be insufficient to train a - Model for the given training set, the training won't be attempted and + any improvements. If the budget is set to a value known to be insufficient to + train a Model for the given training set, the training won't be attempted and will error. - The minimum value is 1000 and the maximum is 72000. model_display_name (str): Optional. The display name of the managed Vertex AI Model. The name can be up to 128 characters long and can be consist of any UTF-8 diff --git a/tests/unit/aiplatform/test_automl_image_training_jobs.py b/tests/unit/aiplatform/test_automl_image_training_jobs.py index 7f092f12d1..330fe3d0ea 100644 --- a/tests/unit/aiplatform/test_automl_image_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_image_training_jobs.py @@ -34,7 +34,7 @@ _TEST_DISPLAY_NAME = "test-display-name" _TEST_METADATA_SCHEMA_URI_IMAGE = schema.dataset.metadata.image -_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS = 1000 +_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS = 7500 _TEST_TRAINING_DISABLE_EARLY_STOPPING = True _TEST_MODEL_TYPE_ICN = "CLOUD" # Image Classification default _TEST_MODEL_TYPE_IOD = "CLOUD_HIGH_ACCURACY_1" # Image Object Detection default @@ -493,6 +493,7 @@ def test_splits_fraction( training_fraction_split=_TEST_FRACTION_SPLIT_TRAINING, validation_fraction_split=_TEST_FRACTION_SPLIT_VALIDATION, test_fraction_split=_TEST_FRACTION_SPLIT_TEST, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, ) @@ -560,6 +561,7 @@ def test_splits_filter( training_filter_split=_TEST_FILTER_SPLIT_TRAINING, validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, test_filter_split=_TEST_FILTER_SPLIT_TEST, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, ) @@ -624,6 +626,7 @@ def test_splits_default( model_from_job = job.run( dataset=mock_dataset_image, model_display_name=_TEST_MODEL_DISPLAY_NAME, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, )