feat: Add disable_retries option to custom jobs.

PiperOrigin-RevId: 557870565

Author: vertex-sdk-bot

google/cloud/aiplatform/jobs.py

@@ -1629,6 +1629,7 @@ def run(
  tensorboard: Optional[str] = None,
  sync: bool = True,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Run this configured CustomJob.
 
@@ -1686,6 +1687,10 @@ def run(
  will unblock and it will be executed in a concurrent Future.
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
  """
  network = network or initializer.global_config.network
 
@@ -1700,6 +1705,7 @@ def run(
  tensorboard=tensorboard,
  sync=sync,
  create_request_timeout=create_request_timeout,
+ disable_retries=disable_retries,
  )
 
  @base.optional_sync()
@@ -1715,6 +1721,7 @@ def _run(
  tensorboard: Optional[str] = None,
  sync: bool = True,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -1770,6 +1777,10 @@ def _run(
  will unblock and it will be executed in a concurrent Future.
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
  """
  self.submit(
  service_account=service_account,
@@ -1781,6 +1792,7 @@ def _run(
  experiment_run=experiment_run,
  tensorboard=tensorboard,
  create_request_timeout=create_request_timeout,
+ disable_retries=disable_retries,
  )
 
  self._block_until_complete()
@@ -1797,6 +1809,7 @@ def submit(
  experiment_run: Optional[Union["aiplatform.ExperimentRun", str]] = None,
  tensorboard: Optional[str] = None,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Submit the configured CustomJob.
 
@@ -1849,6 +1862,10 @@ def submit(
  https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
 
  Raises:
  ValueError:
@@ -1869,11 +1886,12 @@ def submit(
  if network:
  self._gca_resource.job_spec.network = network
 
- if timeout or restart_job_on_worker_restart:
+ if timeout or restart_job_on_worker_restart or disable_retries:
  timeout = duration_pb2.Duration(seconds=timeout) if timeout 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,
  )
 
  if enable_web_access:
@@ -2287,6 +2305,7 @@ def run(
  tensorboard: Optional[str] = None,
  sync: bool = True,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Run this configured CustomJob.
 
@@ -2331,6 +2350,10 @@ def run(
  will unblock and it will be executed in a concurrent Future.
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
  """
  network = network or initializer.global_config.network
 
@@ -2343,6 +2366,7 @@ def run(
  tensorboard=tensorboard,
  sync=sync,
  create_request_timeout=create_request_timeout,
+ disable_retries=disable_retries,
  )
 
  @base.optional_sync()
@@ -2356,6 +2380,7 @@ def _run(
  tensorboard: Optional[str] = None,
  sync: bool = True,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -2398,19 +2423,24 @@ def _run(
  will unblock and it will be executed in a concurrent Future.
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
  """
  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:
+ if timeout or restart_job_on_worker_restart or disable_retries:
  duration = duration_pb2.Duration(seconds=timeout) if timeout else None
  self._gca_resource.trial_job_spec.scheduling = (
  gca_custom_job_compat.Scheduling(
  timeout=duration,
  restart_job_on_worker_restart=restart_job_on_worker_restart,
+ disable_retries=disable_retries,
  )
  )
 

google/cloud/aiplatform/preview/jobs.py

@@ -238,6 +238,7 @@ def submit(
  experiment_run: Optional[Union["aiplatform.ExperimentRun", str]] = None,
  tensorboard: Optional[str] = None,
  create_request_timeout: Optional[float] = None,
+ disable_retries: bool = False,
  ) -> None:
  """Submit the configured CustomJob.
 
@@ -290,6 +291,10 @@ def submit(
  https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
  create_request_timeout (float):
  Optional. The timeout for the create request in seconds.
+ disable_retries (bool):
+ Indicates if the job should retry for internal errors after the
+ job starts running. If True, overrides
+ `restart_job_on_worker_restart` to False.
 
  Raises:
  ValueError:
@@ -310,11 +315,12 @@ def submit(
  if network:
  self._gca_resource.job_spec.network = network
 
- if timeout or restart_job_on_worker_restart:
+ if timeout or restart_job_on_worker_restart or disable_retries:
  timeout = duration_pb2.Duration(seconds=timeout) if timeout 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,
  )
 
  if enable_web_access: