Skip to content

Commit 62b4c4c

Browse files
Ark-kunArk-kun
committed
feat: Added the PipelineJob.submit_from_pipeline_func method
1 parent f5e8ee1 commit 62b4c4c

3 files changed

Lines changed: 221 additions & 2 deletions

File tree

‎google/cloud/aiplatform/pipeline_jobs.py‎

Lines changed: 147 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import logging
2020
import time
2121
import re
22-
from typing import Any, Dict, List, Optional
22+
import tempfile
23+
from typing import Any, Callable, Dict, List, Optional
2324

2425
from google.auth import credentials as auth_credentials
2526
from google.cloud.aiplatform import base
2627
from google.cloud.aiplatform import initializer
2728
from google.cloud.aiplatform import utils
29+
from google.cloud.aiplatform.utils import gcs_utils
2830
from google.cloud.aiplatform.utils import yaml_utils
2931
from google.cloud.aiplatform.utils import pipeline_utils
3032
from google.protobuf import json_format
@@ -615,3 +617,147 @@ def clone(
615617
)
616618

617619
return cloned
620+
621+
@staticmethod
622+
def submit_from_pipeline_func(
623+
# Parameters for the PipelineJob constructor
624+
pipeline_func: Callable,
625+
arguments: Optional[Dict[str, Any]] = None,
626+
output_artifacts_gcs_dir: Optional[str] = None,
627+
enable_caching: Optional[bool] = None,
628+
context_name: Optional[str] = "pipeline",
629+
display_name: Optional[str] = None,
630+
labels: Optional[Dict[str, str]] = None,
631+
job_id: Optional[str] = None,
632+
# Parameters for the PipelineJob.submit method
633+
service_account: Optional[str] = None,
634+
network: Optional[str] = None,
635+
create_request_timeout: Optional[float] = None,
636+
# Parameters for the Vertex SDK
637+
project: Optional[str] = None,
638+
location: Optional[str] = None,
639+
credentials: Optional[auth_credentials.Credentials] = None,
640+
encryption_spec_key_name: Optional[str] = None,
641+
) -> "PipelineJob":
642+
"""Creates pipelineJob from a pipeline function and submits it for execution.
643+
644+
Args:
645+
pipeline_func (Callable):
646+
Required. A pipeline function to compile.
647+
A pipeline function creates instances of components and connects
648+
component inputs to outputs.
649+
arguments (Dict[str, Any]):
650+
Optional. The mapping from runtime parameter names to its values that
651+
control the pipeline run.
652+
output_artifacts_gcs_dir (str):
653+
Optional. The GCS location of the pipeline outputs.
654+
A GCS bucket for artifacts will be created if not specified.
655+
enable_caching (bool):
656+
Optional. Whether to turn on caching for the run.
657+
658+
If this is not set, defaults to the compile time settings, which
659+
are True for all tasks by default, while users may specify
660+
different caching options for individual tasks.
661+
662+
If this is set, the setting applies to all tasks in the pipeline.
663+
664+
Overrides the compile time settings.
665+
context_name (str):
666+
Optional. The name of metadata context. Used for cached execution reuse.
667+
display_name (str):
668+
Optional. The user-defined name of this Pipeline.
669+
labels (Dict[str, str]):
670+
Optional. The user defined metadata to organize PipelineJob.
671+
job_id (str):
672+
Optional. The unique ID of the job run.
673+
If not specified, pipeline name + timestamp will be used.
674+
675+
service_account (str):
676+
Optional. Specifies the service account for workload run-as account.
677+
Users submitting jobs must have act-as permission on this run-as account.
678+
network (str):
679+
Optional. The full name of the Compute Engine network to which the job
680+
should be peered. For example, projects/12345/global/networks/myVPC.
681+
682+
Private services access must already be configured for the network.
683+
If left unspecified, the job is not peered with any network.
684+
create_request_timeout (float):
685+
Optional. The timeout for the create request in seconds.
686+
687+
project (str):
688+
Optional. The project that you want to run this PipelineJob in. If not set,
689+
the project set in aiplatform.init will be used.
690+
location (str):
691+
Optional. Location to create PipelineJob. If not set,
692+
location set in aiplatform.init will be used.
693+
credentials (auth_credentials.Credentials):
694+
Optional. Custom credentials to use to create this PipelineJob.
695+
Overrides credentials set in aiplatform.init.
696+
encryption_spec_key_name (str):
697+
Optional. The Cloud KMS resource identifier of the customer
698+
managed encryption key used to protect the job. Has the
699+
form:
700+
``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
701+
The key needs to be in the same region as where the compute
702+
resource is created.
703+
704+
If this is set, then all
705+
resources created by the PipelineJob will
706+
be encrypted with the provided encryption key.
707+
708+
Overrides encryption_spec_key_name set in aiplatform.init.
709+
710+
Returns:
711+
A Vertex AI PipelineJob.
712+
713+
Raises:
714+
ValueError: If job_id or labels have incorrect format.
715+
"""
716+
717+
# Importing the KFP module here to prevent import errors when the kfp package is not installed.
718+
from kfp.v2 import compiler as compiler_v2
719+
720+
if not output_artifacts_gcs_dir:
721+
output_artifacts_gcs_dir = (
722+
gcs_utils.create_gcs_directory_for_pipeline_artifacts(
723+
service_account=service_account,
724+
project=project,
725+
location=location,
726+
credentials=credentials,
727+
)
728+
)
729+
730+
automatic_display_name = (
731+
pipeline_func.__name__.replace("_", " ")
732+
+ " "
733+
+ datetime.datetime.now().isoformat(sep=" ")
734+
)
735+
display_name = display_name or automatic_display_name
736+
job_id = job_id or re.sub(
737+
r"[^-a-z0-9]", "-", automatic_display_name.lower()
738+
).strip("-")
739+
pipeline_file = tempfile.mktemp(suffix=".json")
740+
compiler_v2.Compiler().compile(
741+
pipeline_func=pipeline_func,
742+
pipeline_name=context_name,
743+
package_path=pipeline_file,
744+
)
745+
pipeline_job = PipelineJob(
746+
template_path=pipeline_file,
747+
parameter_values=arguments,
748+
pipeline_root=output_artifacts_gcs_dir,
749+
enable_caching=enable_caching,
750+
display_name=display_name,
751+
job_id=job_id,
752+
labels=labels,
753+
project=project,
754+
location=location,
755+
credentials=credentials,
756+
encryption_spec_key_name=encryption_spec_key_name,
757+
)
758+
pipeline_job.submit(
759+
service_account=service_account,
760+
network=network,
761+
create_request_timeout=create_request_timeout,
762+
)
763+
return pipeline_job

‎google/cloud/aiplatform/utils/gcs_utils.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from google.cloud import storage
2626

2727
from google.cloud.aiplatform import initializer
28-
28+
from google.cloud.aiplatform.utils import resource_manager_utils
2929

3030
_logger = logging.getLogger(__name__)
3131

@@ -163,3 +163,50 @@ def stage_local_data_in_gcs(
163163
)
164164

165165
return staged_data_uri
166+
167+
168+
def create_gcs_directory_for_pipeline_artifacts(
169+
service_account: Optional[str] = None,
170+
project: Optional[str] = None,
171+
location: Optional[str] = None,
172+
credentials: Optional[auth_credentials.Credentials] = None,
173+
):
174+
project = project or initializer.global_config.project
175+
location = location or initializer.global_config.location
176+
credentials = credentials or initializer.global_config.credentials
177+
178+
pipelines_bucket_name = project + "-vertex-pipelines-" + location
179+
output_artifacts_gcs_dir = "gs://" + pipelines_bucket_name + "/output_artifacts/"
180+
# Creating the bucket if needed
181+
storage_client = storage.Client(
182+
project=project,
183+
credentials=credentials,
184+
)
185+
pipelines_bucket = storage.Bucket(
186+
client=storage_client,
187+
name=pipelines_bucket_name,
188+
)
189+
if not pipelines_bucket.exists():
190+
_logger.info(f'Creating GCS bucket for Vertex Pipelines "{pipelines_bucket_name}"')
191+
pipelines_bucket = storage_client.create_bucket(
192+
bucket_or_name=pipelines_bucket,
193+
project=project,
194+
location=location,
195+
)
196+
# Giving the service account read and write access to teh new bucket
197+
# Workaround for error: "Failed to create pipeline job. Error: Service account `[email protected]`
198+
# does not have `[storage.objects.get, storage.objects.create]` IAM permission(s) to the bucket `xxxxxxxx-vertex-pipelines-us-central1`.
199+
# Please either copy the files to the Google Cloud Storage bucket owned by your project, or grant the required IAM permission(s) to the service account."
200+
if not service_account:
201+
# Getting the project number to use in service account
202+
project_number = resource_manager_utils.get_project_number(project)
203+
service_account = f"{project_number}[email protected]"
204+
bucket_iam_policy = pipelines_bucket.get_iam_policy()
205+
bucket_iam_policy.setdefault("roles/storage.objectCreator", set()).add(
206+
f"serviceAccount:{service_account}"
207+
)
208+
bucket_iam_policy.setdefault("roles/storage.objectViewer", set()).add(
209+
f"serviceAccount:{service_account}"
210+
)
211+
pipelines_bucket.set_iam_policy(bucket_iam_policy)
212+
return output_artifacts_gcs_dir

‎google/cloud/aiplatform/utils/resource_manager_utils.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,3 +48,29 @@ def get_project_id(
4848
project = projects_client.get_project(name=f"projects/{project_number}")
4949

5050
return project.project_id
51+
52+
def get_project_number(
53+
project_id: str,
54+
credentials: Optional[auth_credentials.Credentials] = None,
55+
) -> str:
56+
"""Gets project ID given the project number
57+
58+
Args:
59+
project_id (str):
60+
Required. Google Cloud project unique ID.
61+
credentials: The custom credentials to use when making API calls.
62+
Optional. If not provided, default credentials will be used.
63+
64+
Returns:
65+
str - The automatically generated unique numerical identifier for your GCP project.
66+
67+
"""
68+
69+
credentials = credentials or initializer.global_config.credentials
70+
71+
projects_client = resourcemanager.ProjectsClient(credentials=credentials)
72+
73+
project = projects_client.get_project(name=f"projects/{project_id}")
74+
project_number = project.name.split("/", 1)[1]
75+
76+
return project_number

0 commit comments

Comments
 (0)