feat: Allow setting default service account

PiperOrigin-RevId: 559266585

Author: Ark-kun

google/cloud/aiplatform/initializer.py

@@ -98,6 +98,7 @@ def __init__(self):
  self._credentials = None
  self._encryption_spec_key_name = None
  self._network = None
+ self._service_account = None
 
  def init(
  self,
@@ -113,6 +114,7 @@ def init(
  credentials: Optional[auth_credentials.Credentials] = None,
  encryption_spec_key_name: Optional[str] = None,
  network: Optional[str] = None,
+ service_account: Optional[str] = None,
  ):
  """Updates common initialization parameters with provided options.
 
@@ -155,6 +157,12 @@ def init(
  Private services access must already be configured for the network.
  If specified, all eligible jobs and resources created will be peered
  with this VPC.
+ service_account (str):
+ Optional. The service account used to launch jobs and deploy models.
+ Jobs that use service_account: BatchPredictionJob, CustomJob,
+ PipelineJob, HyperparameterTuningJob, CustomTrainingJob,
+ CustomPythonPackageTrainingJob, CustomContainerTrainingJob,
+ ModelEvaluationJob.
  Raises:
  ValueError:
  If experiment_description is provided but experiment is not.
@@ -194,6 +202,8 @@ def init(
  self._encryption_spec_key_name = encryption_spec_key_name
  if network is not None:
  self._network = network
+ if service_account is not None:
+ self._service_account = service_account
 
  if experiment:
  metadata._experiment_tracker.set_experiment(
@@ -297,6 +307,11 @@ def network(self) -> Optional[str]:
  """Default Compute Engine network to peer to, if provided."""
  return self._network
 
+ @property
+ def service_account(self) -> Optional[str]:
+ """Default service account, if provided."""
+ return self._service_account
+
  @property
  def experiment_name(self) -> Optional[str]:
  """Default experiment name, if provided."""

google/cloud/aiplatform/jobs.py

@@ -761,6 +761,7 @@ def create(
  )
  gapic_batch_prediction_job.explanation_spec = explanation_spec
 
+ service_account = service_account or initializer.global_config.service_account
  if service_account:
  gapic_batch_prediction_job.service_account = service_account
 
@@ -1693,6 +1694,7 @@ def run(
  `restart_job_on_worker_restart` to False.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  self._run(
  service_account=service_account,
@@ -1880,6 +1882,8 @@ def submit(
  raise ValueError(
  "'experiment' is required since you've enabled autolog in 'from_local_script'."
  )
+
+ service_account = service_account or initializer.global_config.service_account
  if service_account:
  self._gca_resource.job_spec.service_account = service_account
 
@@ -2356,6 +2360,7 @@ def run(
  `restart_job_on_worker_restart` to False.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  self._run(
  service_account=service_account,

google/cloud/aiplatform/model_evaluation/model_evaluation_job.py

@@ -278,6 +278,7 @@ def submit(
  Returns:
  (ModelEvaluationJob): Instantiated represnetation of the model evaluation job.
  """
+ service_account = service_account or initializer.global_config.service_account
 
  if isinstance(model_name, aiplatform.Model):
  model_resource_name = model_name.versioned_resource_name

google/cloud/aiplatform/models.py

@@ -1096,6 +1096,7 @@ def _deploy_call(
  to the resource project.
  Users deploying the Model must have the `iam.serviceAccounts.actAs`
  permission on this service account.
+ If not specified, uses the service account set in aiplatform.init.
  explanation_spec (aiplatform.explain.ExplanationSpec):
  Optional. Specification of Model explanation.
  metadata (Sequence[Tuple[str, str]]):
@@ -1120,6 +1121,8 @@ def _deploy_call(
  is not 0 or 100.
  """
 
+ service_account = service_account or initializer.global_config.service_account
+
  max_replica_count = max(min_replica_count, max_replica_count)
 
  if bool(accelerator_type) != bool(accelerator_count):

google/cloud/aiplatform/pipeline_job_schedules.py

@@ -226,6 +226,7 @@ def _create(
  if max_concurrent_run_count:
  self._gca_resource.max_concurrent_run_count = max_concurrent_run_count
 
+ service_account = service_account or initializer.global_config.service_account
  network = network or initializer.global_config.network
 
  if service_account:

google/cloud/aiplatform/pipeline_jobs.py

@@ -382,6 +382,7 @@ def submit(
  current Experiment Run.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  if service_account:
  self._gca_resource.service_account = service_account

google/cloud/aiplatform/training_jobs.py

@@ -3223,6 +3223,7 @@ def run(
  produce a Vertex AI Model.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  worker_pool_specs, managed_model = self._prepare_and_validate_run(
  model_display_name=model_display_name,
@@ -4579,6 +4580,7 @@ def run(
  were not provided in constructor.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  worker_pool_specs, managed_model = self._prepare_and_validate_run(
  model_display_name=model_display_name,
@@ -7348,6 +7350,7 @@ def run(
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
+ If not specified, uses the service account set in aiplatform.init.
  network (str):
  The full name of the Compute Engine network to which the job
  should be peered. For example, projects/12345/global/networks/myVPC.
@@ -7501,6 +7504,7 @@ def run(
  produce a Vertex AI Model.
  """
  network = network or initializer.global_config.network
+ service_account = service_account or initializer.global_config.service_account
 
  worker_pool_specs, managed_model = self._prepare_and_validate_run(
  model_display_name=model_display_name,

google/cloud/aiplatform/utils/gcs_utils.py

@@ -217,6 +217,7 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
  """
  project = project or initializer.global_config.project
  location = location or initializer.global_config.location
+ service_account = service_account or initializer.global_config.service_account
  credentials = credentials or initializer.global_config.credentials
 
  output_artifacts_gcs_dir = (

samples/model-builder/init_sample.py

@@ -25,6 +25,7 @@ def init_sample(
  staging_bucket: Optional[str] = None,
  credentials: Optional[auth_credentials.Credentials] = None,
  encryption_spec_key_name: Optional[str] = None,
+ service_account: Optional[str] = None,
 ):
 
  from google.cloud import aiplatform
@@ -36,6 +37,7 @@ def init_sample(
  staging_bucket=staging_bucket,
  credentials=credentials,
  encryption_spec_key_name=encryption_spec_key_name,
+ service_account=service_account,
  )
 
 

samples/model-builder/init_sample_test.py

@@ -26,6 +26,7 @@ def test_init_sample(mock_sdk_init):
  staging_bucket=constants.STAGING_BUCKET,
  credentials=constants.CREDENTIALS,
  encryption_spec_key_name=constants.ENCRYPTION_SPEC_KEY_NAME,
+ service_account=constants.SERVICE_ACCOUNT,
  )
 
  mock_sdk_init.assert_called_once_with(
@@ -35,4 +36,5 @@ def test_init_sample(mock_sdk_init):
  staging_bucket=constants.STAGING_BUCKET,
  credentials=constants.CREDENTIALS,
  encryption_spec_key_name=constants.ENCRYPTION_SPEC_KEY_NAME,
+ service_account=constants.SERVICE_ACCOUNT,
  )