feat: Add max_wait_duration option to custom jobs.

PiperOrigin-RevId: 671592468

Author: vertex-sdk-bot

google/cloud/aiplatform/jobs.py

@@ -2215,6 +2215,7 @@ def run(
         disable_retries: bool = False,
         persistent_resource_id: Optional[str] = None,
         scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
+        max_wait_duration: Optional[int] = None,
     ) -> None:
         """Run this configured CustomJob.
 
@@ -2285,6 +2286,10 @@ def run(
                 PersistentResource, otherwise, the job will be rejected.
             scheduling_strategy (gca_custom_job_compat.Scheduling.Strategy):
                 Optional. Indicates the job scheduling strategy.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 1 day.
         """
         network = network or initializer.global_config.network
         service_account = service_account or initializer.global_config.service_account
@@ -2303,6 +2308,7 @@ def run(
             disable_retries=disable_retries,
             persistent_resource_id=persistent_resource_id,
             scheduling_strategy=scheduling_strategy,
+            max_wait_duration=max_wait_duration,
         )
 
     @base.optional_sync()
@@ -2321,6 +2327,7 @@ def _run(
         disable_retries: bool = False,
         persistent_resource_id: Optional[str] = None,
         scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
+        max_wait_duration: Optional[int] = None,
     ) -> None:
         """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -2389,6 +2396,10 @@ def _run(
                 PersistentResource, otherwise, the job will be rejected.
             scheduling_strategy (gca_custom_job_compat.Scheduling.Strategy):
                 Optional. Indicates the job scheduling strategy.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 1 day.
         """
         self.submit(
             service_account=service_account,
@@ -2403,6 +2414,7 @@ def _run(
             disable_retries=disable_retries,
             persistent_resource_id=persistent_resource_id,
             scheduling_strategy=scheduling_strategy,
+            max_wait_duration=max_wait_duration,
         )
 
         self._block_until_complete()
@@ -2422,6 +2434,7 @@ def submit(
         disable_retries: bool = False,
         persistent_resource_id: Optional[str] = None,
         scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
+        max_wait_duration: Optional[int] = None,
     ) -> None:
         """Submit the configured CustomJob.
 
@@ -2487,6 +2500,10 @@ def submit(
                 PersistentResource, otherwise, the job will be rejected.
             scheduling_strategy (gca_custom_job_compat.Scheduling.Strategy):
                 Optional. Indicates the job scheduling strategy.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 1 day.
 
         Raises:
             ValueError:
@@ -2514,13 +2531,20 @@ def submit(
             or restart_job_on_worker_restart
             or disable_retries
             or scheduling_strategy
+            or max_wait_duration
         ):
             timeout = duration_pb2.Duration(seconds=timeout) if timeout else None
+            max_wait_duration = (
+                duration_pb2.Duration(seconds=max_wait_duration)
+                if max_wait_duration
+                else None
+            )
             self._gca_resource.job_spec.scheduling = gca_custom_job_compat.Scheduling(
                 timeout=timeout,
                 restart_job_on_worker_restart=restart_job_on_worker_restart,
                 disable_retries=disable_retries,
                 strategy=scheduling_strategy,
+                max_wait_duration=max_wait_duration,
             )
 
         if enable_web_access:
@@ -2886,6 +2910,7 @@ def run(
         create_request_timeout: Optional[float] = None,
         disable_retries: bool = False,
         scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
+        max_wait_duration: Optional[int] = None,  # seconds
     ) -> None:
         """Run this configured CustomJob.
 
@@ -2936,6 +2961,10 @@ def run(
                 `restart_job_on_worker_restart` to False.
             scheduling_strategy (gca_custom_job_compat.Scheduling.Strategy):
                 Optional. Indicates the job scheduling strategy.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 1 day.
         """
         network = network or initializer.global_config.network
         service_account = service_account or initializer.global_config.service_account
@@ -2951,6 +2980,7 @@ def run(
             create_request_timeout=create_request_timeout,
             disable_retries=disable_retries,
             scheduling_strategy=scheduling_strategy,
+            max_wait_duration=max_wait_duration,
         )
 
     @base.optional_sync()
@@ -2966,6 +2996,7 @@ def _run(
         create_request_timeout: Optional[float] = None,
         disable_retries: bool = False,
         scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
+        max_wait_duration: Optional[int] = None,  # seconds
     ) -> None:
         """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -3014,6 +3045,10 @@ def _run(
                 `restart_job_on_worker_restart` to False.
             scheduling_strategy (gca_custom_job_compat.Scheduling.Strategy):
                 Optional. Indicates the job scheduling strategy.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 1 day.
         """
         if service_account:
             self._gca_resource.trial_job_spec.service_account = service_account
@@ -3025,15 +3060,22 @@ def _run(
             timeout
             or restart_job_on_worker_restart
             or disable_retries
+            or max_wait_duration
             or scheduling_strategy
         ):
-            duration = duration_pb2.Duration(seconds=timeout) if timeout else None
+            timeout = duration_pb2.Duration(seconds=timeout) if timeout else None
+            max_wait_duration = (
+                duration_pb2.Duration(seconds=max_wait_duration)
+                if max_wait_duration
+                else None
+            )
             self._gca_resource.trial_job_spec.scheduling = (
                 gca_custom_job_compat.Scheduling(
-                    timeout=duration,
+                    timeout=timeout,
                     restart_job_on_worker_restart=restart_job_on_worker_restart,
                     disable_retries=disable_retries,
                     strategy=scheduling_strategy,
+                    max_wait_duration=max_wait_duration,
                 )
             )
 

google/cloud/aiplatform/preview/jobs.py

@@ -266,6 +266,7 @@ def submit(
         tensorboard: Optional[str] = None,
         create_request_timeout: Optional[float] = None,
         disable_retries: bool = False,
+        max_wait_duration: Optional[int] = None,
     ) -> None:
         """Submit the configured CustomJob.
 
@@ -322,6 +323,10 @@ def submit(
                 Indicates if the job should retry for internal errors after the
                 job starts running. If True, overrides
                 `restart_job_on_worker_restart` to False.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 30 minutes.
 
         Raises:
             ValueError:
@@ -342,12 +347,23 @@ def submit(
         if network:
             self._gca_resource.job_spec.network = network
 
-        if timeout or restart_job_on_worker_restart or disable_retries:
+        if (
+            timeout
+            or restart_job_on_worker_restart
+            or disable_retries
+            or max_wait_duration
+        ):
             timeout = duration_pb2.Duration(seconds=timeout) if timeout else None
+            max_wait_duration = (
+                duration_pb2.Duration(seconds=max_wait_duration)
+                if max_wait_duration
+                else None
+            )
             self._gca_resource.job_spec.scheduling = gca_custom_job_compat.Scheduling(
                 timeout=timeout,
                 restart_job_on_worker_restart=restart_job_on_worker_restart,
                 disable_retries=disable_retries,
+                max_wait_duration=max_wait_duration,
             )
 
         if enable_web_access:
@@ -741,6 +757,7 @@ def _run(
         sync: bool = True,
         create_request_timeout: Optional[float] = None,
         disable_retries: bool = False,
+        max_wait_duration: Optional[int] = None,
     ) -> None:
         """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -787,20 +804,35 @@ def _run(
                 Indicates if the job should retry for internal errors after the
                 job starts running. If True, overrides
                 `restart_job_on_worker_restart` to False.
+            max_wait_duration (int):
+                This is the maximum duration that a job will wait for the
+                requested resources to be provisioned in seconds. If set to 0,
+                the job will wait indefinitely. The default is 30 minutes.
         """
         if service_account:
             self._gca_resource.trial_job_spec.service_account = service_account
 
         if network:
             self._gca_resource.trial_job_spec.network = network
 
-        if timeout or restart_job_on_worker_restart or disable_retries:
-            duration = duration_pb2.Duration(seconds=timeout) if timeout else None
+        if (
+            timeout
+            or restart_job_on_worker_restart
+            or disable_retries
+            or max_wait_duration
+        ):
+            timeout = duration_pb2.Duration(seconds=timeout) if timeout else None
+            max_wait_duration = (
+                duration_pb2.Duration(seconds=max_wait_duration)
+                if max_wait_duration
+                else None
+            )
             self._gca_resource.trial_job_spec.scheduling = (
                 gca_custom_job_compat.Scheduling(
-                    timeout=duration,
+                    timeout=timeout,
                     restart_job_on_worker_restart=restart_job_on_worker_restart,
                     disable_retries=disable_retries,
+                    max_wait_duration=max_wait_duration,
                 )
             )