WIP: Fix Vertex AI scheduler 64-character limit for long pipeline names (#2262)

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,

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],

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.