diff --git a/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py b/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py index b5aa67e63ee..985baaa0d72 100644 --- a/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +++ b/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py @@ -29,6 +29,7 @@ # permissions and limitations under the License. """Implementation of the VertexAI orchestrator.""" +import hashlib import os import re import types @@ -111,6 +112,7 @@ StepOperatorEntrypointConfiguration, ) from zenml.utils.io_utils import get_global_config_directory +from zenml.utils.string_utils import append_random_suffix if TYPE_CHECKING: from zenml.config.base_settings import BaseSettings @@ -128,6 +130,23 @@ STEP_JOB_NAME_METADATA_KEY = "job_name" +def sanitize_vertex_job_name(job_name: str) -> str: + """Sanitizes the Vertex AI job name to comply with GCP requirements.""" + job_name = job_name.lower() + job_name = job_name.replace("_", "-").replace(" ", "-") + job_name = re.sub(r"[^a-z0-9\-]", "", job_name) + + if not job_name or not job_name[0].isalpha(): + job_name = f"j-{job_name}" + + return append_random_suffix( + original_string=job_name, + suffix_length=5, + max_length=64, + separator="-" + ) + + def _clean_pipeline_name(pipeline_name: str) -> str: """Clean pipeline name to be a valid Vertex AI Pipeline name. @@ -529,7 +548,7 @@ def _create_dynamic_pipeline() -> Any: step_name_to_dynamic_component[step_name] = component @dsl.pipeline( # type: ignore[misc] - display_name=orchestrator_run_name, + display_name=sanitize_vertex_job_name(orchestrator_run_name), ) def dynamic_pipeline() -> None: """Dynamic pipeline.""" @@ -880,7 +899,7 @@ def _upload_and_run_pipeline( # Instantiate the Vertex AI Pipelines job run = aiplatform.PipelineJob( - display_name=pipeline_name, + display_name=sanitize_vertex_job_name(pipeline_name), template_path=pipeline_file_path, job_id=job_id, pipeline_root=self._pipeline_root, diff --git a/src/zenml/orchestrators/utils.py b/src/zenml/orchestrators/utils.py index 25f6c296967..b159f20d745 100644 --- a/src/zenml/orchestrators/utils.py +++ b/src/zenml/orchestrators/utils.py @@ -35,6 +35,7 @@ from zenml.enums import APITokenType, AuthScheme, StackComponentType, StoreType from zenml.logger import get_logger from zenml.stack import StackComponent +from zenml.utils.string_utils import append_random_suffix logger = get_logger(__name__) @@ -64,24 +65,23 @@ def get_orchestrator_run_name( Returns: The orchestrator run name. """ - suffix_length = 32 - pipeline_name = f"{pipeline_name}_" + if max_length and max_length < 8: + raise ValueError( + "Maximum length for orchestrator run name must be 8 or above." + ) + suffix_length = 32 if max_length: - if max_length < 8: - raise ValueError( - "Maximum length for orchestrator run name must be 8 or above." - ) - - # Make sure we always have a certain suffix to guarantee no overlap - # with other runs - suffix_length = min(32, max(8, max_length - len(pipeline_name))) - pipeline_name = pipeline_name[: (max_length - suffix_length)] - - suffix = "".join(random.choices("0123456789abcdef", k=suffix_length)) - - return f"{pipeline_name}{suffix}" - + # The -1 accounts for the "_" separator added by the utility function + suffix_length = min(32, max(8, max_length - len(pipeline_name) - 1)) + + return append_random_suffix( + original_string=pipeline_name, + suffix_length=suffix_length, + max_length=max_length, + separator="", + charset="0123456789abcdef" + ) def is_setting_enabled( is_enabled_on_step: Optional[bool], diff --git a/src/zenml/utils/string_utils.py b/src/zenml/utils/string_utils.py index 3cee230e1e2..1a3cc7e42a4 100644 --- a/src/zenml/utils/string_utils.py +++ b/src/zenml/utils/string_utils.py @@ -27,6 +27,40 @@ V = TypeVar("V", bound=Any) +def append_random_suffix( + original_string: str, + suffix_length: int, + max_length: Optional[int] = None, + separator: str = "-", + charset: str = string.ascii_lowercase + string.digits +) -> str: + """Appends a random suffix to a string, truncating the original if necessary. + + Args: + original_string: The base string to modify. + suffix_length: The length of the random suffix to generate. + max_length: Maximum allowed length of the final merged string. + separator: The string used to join the original string and suffix. + charset: The pool of characters to use for generating the random suffix. + + Returns: + The merged string with the random suffix appended. + """ + suffix = "".join(random.choices(charset, k=suffix_length)) + + if max_length is not None: + allowed_original_length = max_length - len(separator) - suffix_length + + if allowed_original_length > 0: + original_string = original_string[:allowed_original_length] + else: + return suffix[:max_length] + + if original_string: + return f"{original_string}{separator}{suffix}" + return suffix + + def get_human_readable_time(seconds: float) -> str: """Convert seconds into a human-readable string.