feat: Add PSC interface config support for Custom Training Jobs

PiperOrigin-RevId: 770333532

Author: vertex-sdk-bot

google/cloud/aiplatform/jobs.py

@@ -45,6 +45,7 @@
  model_deployment_monitoring_job as gca_model_deployment_monitoring_job_compat,
  job_state_v1beta1 as gca_job_state_v1beta1,
  model_monitoring_v1beta1 as gca_model_monitoring_v1beta1,
+ service_networking as gca_service_networking,
 ) # TODO(b/242108750): remove temporary logic once model monitoring for batch prediction is GA
 
 from google.cloud.aiplatform.constants import base as constants
@@ -2236,6 +2237,9 @@ def run(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> None:
  """Run this configured CustomJob.
 
@@ -2310,6 +2314,9 @@ def run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
  """
  network = network or initializer.global_config.network
  service_account = service_account or initializer.global_config.service_account
@@ -2329,6 +2336,7 @@ def run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  @base.optional_sync()
@@ -2348,6 +2356,9 @@ def _run(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> None:
  """Helper method to ensure network synchronization and to run the configured CustomJob.
 
@@ -2420,6 +2431,9 @@ def _run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
  """
  self.submit(
  service_account=service_account,
@@ -2435,6 +2449,7 @@ def _run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  self._block_until_complete()
@@ -2455,6 +2470,9 @@ def submit(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> None:
  """Submit the configured CustomJob.
 
@@ -2524,6 +2542,9 @@ def submit(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Raises:
  ValueError:
@@ -2546,6 +2567,9 @@ def submit(
  if network:
  self._gca_resource.job_spec.network = network
 
+ if psc_interface_config:
+ self._gca_resource.job_spec.psc_interface_config = psc_interface_config
+
  if (
  timeout
  or restart_job_on_worker_restart

google/cloud/aiplatform/training_jobs.py

@@ -45,6 +45,7 @@
  training_pipeline as gca_training_pipeline,
  study as gca_study_compat,
  custom_job as gca_custom_job_compat,
+ service_networking as gca_service_networking,
 )
 
 from google.cloud.aiplatform.utils import _timestamped_gcs_dir
@@ -1553,6 +1554,9 @@ def _prepare_training_task_inputs_and_output_dir(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Tuple[Dict, str]:
  """Prepares training task inputs and output directory for custom job.
 
@@ -1617,6 +1621,8 @@ def _prepare_training_task_inputs_and_output_dir(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. The PSC interface config for the job.
  Returns:
  Training task inputs and Output directory for custom job.
  """
@@ -1645,6 +1651,8 @@ def _prepare_training_task_inputs_and_output_dir(
  training_task_inputs["enable_dashboard_access"] = enable_dashboard_access
  if persistent_resource_id:
  training_task_inputs["persistent_resource_id"] = persistent_resource_id
+ if psc_interface_config:
+ training_task_inputs["psc_interface_config"] = psc_interface_config
 
  if (
  timeout
@@ -3055,6 +3063,9 @@ def run(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Runs the custom training job.
 
@@ -3433,6 +3444,9 @@ def run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  The trained Vertex AI model resource or None if the training
@@ -3504,6 +3518,7 @@ def run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  def submit(
@@ -3564,6 +3579,9 @@ def submit(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Submits the custom training job without blocking until completion.
 
@@ -3887,6 +3905,9 @@ def submit(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -3958,6 +3979,7 @@ def submit(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  @base.optional_sync(construct_object_on_arg="managed_model")
@@ -4007,6 +4029,9 @@ def _run(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Packages local script and launches training_job.
 
@@ -4209,6 +4234,8 @@ def _run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. The PSC interface config for the job.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -4265,6 +4292,7 @@ def _run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  model = self._run_job(
@@ -4596,6 +4624,9 @@ def run(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Runs the custom training job.
 
@@ -4912,6 +4943,9 @@ def run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -4982,6 +5016,7 @@ def run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  def submit(
@@ -5042,6 +5077,9 @@ def submit(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Submits the custom training job without blocking until completion.
 
@@ -5358,6 +5396,9 @@ def submit(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -5428,6 +5469,7 @@ def submit(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  @base.optional_sync(construct_object_on_arg="managed_model")
@@ -5476,6 +5518,9 @@ def _run(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Packages local script and launches training_job.
  Args:
@@ -5674,6 +5719,9 @@ def _run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -5724,6 +5772,7 @@ def _run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  model = self._run_job(
@@ -7755,6 +7804,9 @@ def run(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Runs the custom training job.
 
@@ -8072,6 +8124,9 @@ def run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -8137,6 +8192,7 @@ def run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  @base.optional_sync(construct_object_on_arg="managed_model")
@@ -8184,6 +8240,9 @@ def _run(
  persistent_resource_id: Optional[str] = None,
  scheduling_strategy: Optional[gca_custom_job_compat.Scheduling.Strategy] = None,
  max_wait_duration: Optional[int] = None,
+ psc_interface_config: Optional[
+ gca_service_networking.PscInterfaceConfig
+ ] = None,
  ) -> Optional[models.Model]:
  """Packages local script and launches training_job.
 
@@ -8367,6 +8426,9 @@ def _run(
  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.
+ psc_interface_config (gca_service_networking.PscInterfaceConfig):
+ Optional. Configuration for Private Service Connect interface
+ used for training.
 
  Returns:
  model: The trained Vertex AI Model resource or None if training did not
@@ -8417,6 +8479,7 @@ def _run(
  persistent_resource_id=persistent_resource_id,
  scheduling_strategy=scheduling_strategy,
  max_wait_duration=max_wait_duration,
+ psc_interface_config=psc_interface_config,
  )
 
  model = self._run_job(

tests/unit/aiplatform/constants.py

@@ -222,6 +222,7 @@ class TrainingJobConstants:
  )
  _TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
  _TEST_SPOT_STRATEGY = custom_job.Scheduling.Strategy.SPOT
+ _TEST_PSC_INTERFACE_CONFIG = {"network_attachment": "network_attachment_value"}
 
  def create_tpu_job_proto(tpu_version):
  worker_pool_spec = (