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