fix: Update milli node_hours for image training (#663)

Co-authored-by: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com>

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

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,
         )