zenml-nightly 0.82.1.dev20250527__py3-none-any.whl → 0.83.0.dev20250529__py3-none-any.whl

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -14,12 +14,12 @@
 """Implementation of the Skypilot base VM orchestrator."""
 
 import os
-import re
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
 from uuid import uuid4
 
 import sky
+from sky import StatusRefreshMode
 
 from zenml.entrypoints import PipelineEntrypointConfiguration
 from zenml.enums import StackComponentType
@@ -31,6 +31,15 @@ from zenml.integrations.skypilot.flavors.skypilot_orchestrator_base_vm_config im
 from zenml.integrations.skypilot.orchestrators.skypilot_orchestrator_entrypoint_configuration import (
     SkypilotOrchestratorEntrypointConfiguration,
 )
+from zenml.integrations.skypilot.utils import (
+    create_docker_run_command,
+    prepare_docker_setup,
+    prepare_launch_kwargs,
+    prepare_resources_kwargs,
+    prepare_task_kwargs,
+    sanitize_cluster_name,
+    sky_job_get,
+)
 from zenml.logger import get_logger
 from zenml.orchestrators import (
     ContainerizedOrchestrator,
@@ -252,32 +261,21 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
         entrypoint_str = " ".join(command)
         arguments_str = " ".join(args)
 
-        task_envs = environment
-        docker_environment_str = " ".join(
-            f"-e {k}={v}" for k, v in environment.items()
-        )
-        custom_run_args = " ".join(settings.docker_run_args)
-        if custom_run_args:
-            custom_run_args += " "
-
-        instance_type = settings.instance_type or self.DEFAULT_INSTANCE_TYPE
+        task_envs = environment.copy()
 
         # Set up credentials
         self.setup_credentials()
 
-        # Guaranteed by stack validation
-        assert stack is not None and stack.container_registry is not None
+        # Prepare Docker setup
+        setup, docker_creds_envs = prepare_docker_setup(
+            container_registry_uri=stack.container_registry.config.uri,
+            credentials=stack.container_registry.credentials,
+            use_sudo=True,  # Base orchestrator uses sudo
+        )
 
-        if docker_creds := stack.container_registry.credentials:
-            docker_username, docker_password = docker_creds
-            setup = (
-                f"sudo docker login --username $DOCKER_USERNAME --password "
-                f"$DOCKER_PASSWORD {stack.container_registry.config.uri}"
-            )
-            task_envs["DOCKER_USERNAME"] = docker_username
-            task_envs["DOCKER_PASSWORD"] = docker_password
-        else:
-            setup = None
+        # Update task_envs with Docker credentials
+        if docker_creds_envs:
+            task_envs.update(docker_creds_envs)
 
         # Run the entire pipeline
 
@@ -291,45 +289,49 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                 down = False
                 idle_minutes_to_autostop = None
             else:
-                run_command = f"sudo docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}"
+                run_command = create_docker_run_command(
+                    image=image,
+                    entrypoint_str=entrypoint_str,
+                    arguments_str=arguments_str,
+                    environment=task_envs,
+                    docker_run_args=settings.docker_run_args,
+                    use_sudo=True,  # Base orchestrator uses sudo
+                )
                 down = settings.down
                 idle_minutes_to_autostop = settings.idle_minutes_to_autostop
-            task = sky.Task(
-                run=run_command,
+
+            # Create the Task with all parameters and task settings
+            task_kwargs = prepare_task_kwargs(
+                settings=settings,
+                run_command=run_command,
                 setup=setup,
-                envs=task_envs,
+                task_envs=task_envs,
+                task_name=f"{orchestrator_run_name}",
             )
+
+            task = sky.Task(**task_kwargs)
             logger.debug(f"Running run: {run_command}")
 
-            task = task.set_resources(
-                sky.Resources(
-                    cloud=self.cloud,
-                    instance_type=instance_type,
-                    cpus=settings.cpus,
-                    memory=settings.memory,
-                    accelerators=settings.accelerators,
-                    accelerator_args=settings.accelerator_args,
-                    use_spot=settings.use_spot,
-                    job_recovery=settings.job_recovery,
-                    region=settings.region,
-                    zone=settings.zone,
-                    image_id=image
-                    if isinstance(self.cloud, sky.clouds.Kubernetes)
-                    else settings.image_id,
-                    disk_size=settings.disk_size,
-                    disk_tier=settings.disk_tier,
-                )
+            # Set resources with all parameters and resource settings
+            resources_kwargs = prepare_resources_kwargs(
+                cloud=self.cloud,
+                settings=settings,
+                default_instance_type=self.DEFAULT_INSTANCE_TYPE,
+                kubernetes_image=image
+                if isinstance(self.cloud, sky.clouds.Kubernetes)
+                else None,
             )
-            # Do not detach run if logs are being streamed
-            # Otherwise, the logs will not be streamed after the task is submitted
-            # Could also be a parameter in the settings to control this behavior
-            detach_run = not settings.stream_logs
+
+            task = task.set_resources(sky.Resources(**resources_kwargs))
 
             launch_new_cluster = True
             if settings.cluster_name:
-                cluster_info = sky.status(
-                    refresh=True, cluster_names=settings.cluster_name
+                status_request_id = sky.status(
+                    refresh=StatusRefreshMode.AUTO,
+                    cluster_names=[settings.cluster_name],
                 )
+                cluster_info = sky.stream_and_get(status_request_id)
+
                 if cluster_info:
                     logger.info(
                         f"Found existing cluster {settings.cluster_name}. Reusing..."
@@ -342,7 +344,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     )
                     cluster_name = settings.cluster_name
             else:
-                cluster_name = self.sanitize_cluster_name(
+                cluster_name = sanitize_cluster_name(
                     f"{orchestrator_run_name}"
                 )
                 logger.info(
@@ -350,33 +352,55 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                 )
 
             if launch_new_cluster:
-                sky.launch(
+                # Prepare launch parameters with additional launch settings
+                launch_kwargs = prepare_launch_kwargs(
+                    settings=settings,
+                    down=down,
+                    idle_minutes_to_autostop=idle_minutes_to_autostop,
+                )
+                logger.info(
+                    f"Launching the task on a new cluster: {cluster_name}"
+                )
+                launch_job_id = sky.launch(
                     task,
                     cluster_name,
-                    retry_until_up=settings.retry_until_up,
-                    idle_minutes_to_autostop=idle_minutes_to_autostop,
-                    down=down,
-                    stream_logs=settings.stream_logs,
-                    backend=None,
-                    detach_setup=True,
-                    detach_run=detach_run,
+                    **launch_kwargs,
                 )
+                sky_job_get(launch_job_id, settings.stream_logs, cluster_name)
+
             else:
-                # Make sure the cluster is up -
-                # If the cluster is already up, this will not do anything
-                sky.start(
+                # Prepare exec parameters with additional launch settings
+                exec_kwargs = {
+                    "down": down,
+                    "backend": None,
+                    **settings.launch_settings,  # Can reuse same settings for exec
+                }
+
+                # Remove None values to avoid overriding SkyPilot defaults
+                exec_kwargs = {
+                    k: v for k, v in exec_kwargs.items() if v is not None
+                }
+
+                # Make sure the cluster is up
+                start_request_id = sky.start(
                     settings.cluster_name,
                     down=down,
                     idle_minutes_to_autostop=idle_minutes_to_autostop,
                     retry_until_up=settings.retry_until_up,
                 )
-                sky.exec(
+                sky.stream_and_get(start_request_id)
+
+                logger.info(
+                    f"Executing the task on the cluster: {settings.cluster_name}"
+                )
+                exec_job_id = sky.exec(
                     task,
-                    settings.cluster_name,
-                    down=down,
-                    stream_logs=settings.stream_logs,
-                    backend=None,
-                    detach_run=detach_run,
+                    cluster_name=settings.cluster_name,
+                    **exec_kwargs,
+                )
+                assert settings.cluster_name is not None
+                sky_job_get(
+                    exec_job_id, settings.stream_logs, settings.cluster_name
                 )
 
         except Exception as e:
@@ -386,19 +410,3 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
         finally:
             # Unset the service connector AWS profile ENV variable
             self.prepare_environment_variable(set=False)
-
-    def sanitize_cluster_name(self, name: str) -> str:
-        """Sanitize the value to be used in a cluster name.
-
-        Args:
-            name: Arbitrary input cluster name.
-
-        Returns:
-            Sanitized cluster name.
-        """
-        name = re.sub(
-            r"[^a-z0-9-]", "-", name.lower()
-        )  # replaces any character that is not a lowercase letter, digit, or hyphen with a hyphen
-        name = re.sub(r"^[-]+", "", name)  # trim leading hyphens
-        name = re.sub(r"[-]+$", "", name)  # trim trailing hyphens
-        return name

zenml/integrations/skypilot/orchestrators/skypilot_orchestrator_entrypoint.py

@@ -32,8 +32,20 @@ from zenml.integrations.skypilot.orchestrators.skypilot_base_vm_orchestrator imp
     ENV_ZENML_SKYPILOT_ORCHESTRATOR_RUN_ID,
     SkypilotBaseOrchestrator,
 )
+from zenml.integrations.skypilot.utils import (
+    create_docker_run_command,
+    prepare_docker_setup,
+    prepare_launch_kwargs,
+    prepare_resources_kwargs,
+    prepare_task_kwargs,
+    sanitize_cluster_name,
+    sky_job_get,
+)
 from zenml.logger import get_logger
-from zenml.orchestrators.dag_runner import ThreadedDagRunner
+from zenml.orchestrators.dag_runner import NodeStatus, ThreadedDagRunner
+from zenml.orchestrators.publish_utils import (
+    publish_failed_pipeline_run,
+)
 from zenml.orchestrators.utils import get_config_environment_vars
 
 logger = get_logger(__name__)
@@ -65,212 +77,228 @@ def main() -> None:
         TypeError: If the active stack's orchestrator is not an instance of
             SkypilotBaseOrchestrator.
         ValueError: If the active stack's container registry is None.
+        Exception: If the orchestration or one of the steps fails.
     """
     # Log to the container's stdout so it can be streamed by the client.
     logger.info("Skypilot orchestrator VM started.")
 
     # Parse / extract args.
     args = parse_args()
-
     orchestrator_run_id = socket.gethostname()
 
-    deployment = Client().get_deployment(args.deployment_id)
-
-    pipeline_dag = {
-        step_name: step.spec.upstream_steps
-        for step_name, step in deployment.step_configurations.items()
-    }
-    step_command = StepEntrypointConfiguration.get_entrypoint_command()
-    entrypoint_str = " ".join(step_command)
-
-    active_stack = Client().active_stack
-
-    orchestrator = active_stack.orchestrator
-    if not isinstance(orchestrator, SkypilotBaseOrchestrator):
-        raise TypeError(
-            "The active stack's orchestrator is not an instance of SkypilotBaseOrchestrator."
-        )
-
-    # Set up credentials
-    orchestrator.setup_credentials()
-
-    # Set the service connector AWS profile ENV variable
-    orchestrator.prepare_environment_variable(set=True)
+    run = None
 
-    # get active container registry
-    container_registry = active_stack.container_registry
-    if container_registry is None:
-        raise ValueError("Container registry cannot be None.")
+    try:
+        deployment = Client().get_deployment(args.deployment_id)
 
-    if docker_creds := container_registry.credentials:
-        docker_username, docker_password = docker_creds
-        setup = (
-            f"docker login --username $DOCKER_USERNAME --password "
-            f"$DOCKER_PASSWORD {container_registry.config.uri}"
-        )
-        task_envs = {
-            "DOCKER_USERNAME": docker_username,
-            "DOCKER_PASSWORD": docker_password,
+        pipeline_dag = {
+            step_name: step.spec.upstream_steps
+            for step_name, step in deployment.step_configurations.items()
         }
-    else:
-        setup = None
-        task_envs = None
-
-    unique_resource_configs: Dict[str, str] = {}
-    for step_name, step in deployment.step_configurations.items():
-        settings = cast(
-            SkypilotBaseOrchestratorSettings,
-            orchestrator.get_settings(step),
-        )
-        # Handle both str and Dict[str, int] types for accelerators
-        if isinstance(settings.accelerators, dict):
-            accelerators_hashable = frozenset(settings.accelerators.items())
-        elif isinstance(settings.accelerators, str):
-            accelerators_hashable = frozenset({(settings.accelerators, 1)})
-        else:
-            accelerators_hashable = None
-        resource_config = (
-            settings.instance_type,
-            settings.cpus,
-            settings.memory,
-            settings.disk_size,  # Assuming disk_size is part of the settings
-            settings.disk_tier,  # Assuming disk_tier is part of the settings
-            settings.use_spot,
-            settings.job_recovery,
-            settings.region,
-            settings.zone,
-            accelerators_hashable,
-        )
-        cluster_name_parts = [
-            orchestrator.sanitize_cluster_name(str(part))
-            for part in resource_config
-            if part is not None
-        ]
-        cluster_name = f"cluster-{orchestrator_run_id}" + "-".join(
-            cluster_name_parts
-        )
-        unique_resource_configs[step_name] = cluster_name
-
-    run = Client().list_pipeline_runs(
-        sort_by="asc:created",
-        size=1,
-        deployment_id=args.deployment_id,
-        status=ExecutionStatus.INITIALIZING,
-    )[0]
+        step_command = StepEntrypointConfiguration.get_entrypoint_command()
+        entrypoint_str = " ".join(step_command)
 
-    logger.info("Fetching pipeline run: %s", run.id)
+        active_stack = Client().active_stack
 
-    def run_step_on_skypilot_vm(step_name: str) -> None:
-        """Run a pipeline step in a separate Skypilot VM.
+        orchestrator = active_stack.orchestrator
+        if not isinstance(orchestrator, SkypilotBaseOrchestrator):
+            raise TypeError(
+                "The active stack's orchestrator is not an instance of SkypilotBaseOrchestrator."
+            )
 
-        Args:
-            step_name: Name of the step.
-        """
-        cluster_name = unique_resource_configs[step_name]
+        # Set up credentials
+        orchestrator.setup_credentials()
 
-        image = SkypilotBaseOrchestrator.get_image(
-            deployment=deployment, step_name=step_name
-        )
+        # Set the service connector AWS profile ENV variable
+        orchestrator.prepare_environment_variable(set=True)
 
-        step_args = StepEntrypointConfiguration.get_entrypoint_arguments(
-            step_name=step_name, deployment_id=deployment.id
-        )
-        arguments_str = " ".join(step_args)
+        # get active container registry
+        container_registry = active_stack.container_registry
+        if container_registry is None:
+            raise ValueError("Container registry cannot be None.")
 
-        step = deployment.step_configurations[step_name]
-        settings = cast(
-            SkypilotBaseOrchestratorSettings,
-            orchestrator.get_settings(step),
+        # Prepare Docker setup
+        setup, task_envs = prepare_docker_setup(
+            container_registry_uri=container_registry.config.uri,
+            credentials=container_registry.credentials,
+            use_sudo=False,  # Entrypoint doesn't use sudo
         )
-        env = get_config_environment_vars()
-        env[ENV_ZENML_SKYPILOT_ORCHESTRATOR_RUN_ID] = orchestrator_run_id
 
-        docker_environment_str = " ".join(
-            f"-e {k}={v}" for k, v in env.items()
-        )
-        custom_run_args = " ".join(settings.docker_run_args)
-        if custom_run_args:
-            custom_run_args += " "
-
-        # Set up the task
-        run_command = f"docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}"
-        task_name = f"{deployment.id}-{step_name}-{time.time()}"
-        task = sky.Task(
-            run=run_command,
-            setup=setup,
-            envs=task_envs,
-            name=task_name,
-        )
-        task = task.set_resources(
-            sky.Resources(
-                cloud=orchestrator.cloud,
-                instance_type=settings.instance_type
-                or orchestrator.DEFAULT_INSTANCE_TYPE,
-                cpus=settings.cpus,
-                memory=settings.memory,
-                disk_size=settings.disk_size,
-                disk_tier=settings.disk_tier,
-                accelerators=settings.accelerators,
-                accelerator_args=settings.accelerator_args,
-                use_spot=settings.use_spot,
-                job_recovery=settings.job_recovery,
-                region=settings.region,
-                zone=settings.zone,
-                image_id=settings.image_id,
+        unique_resource_configs: Dict[str, str] = {}
+        for step_name, step in deployment.step_configurations.items():
+            settings = cast(
+                SkypilotBaseOrchestratorSettings,
+                orchestrator.get_settings(step),
             )
-        )
+            # Handle both str and Dict[str, int] types for accelerators
+            if isinstance(settings.accelerators, dict):
+                accelerators_hashable = frozenset(
+                    settings.accelerators.items()
+                )
+            elif isinstance(settings.accelerators, str):
+                accelerators_hashable = frozenset({(settings.accelerators, 1)})
+            else:
+                accelerators_hashable = None
+            resource_config = (
+                settings.instance_type,
+                settings.cpus,
+                settings.memory,
+                settings.disk_size,  # Assuming disk_size is part of the settings
+                settings.disk_tier,  # Assuming disk_tier is part of the settings
+                settings.use_spot,
+                settings.job_recovery,
+                settings.region,
+                settings.zone,
+                accelerators_hashable,
+            )
+            cluster_name_parts = [
+                sanitize_cluster_name(str(part))
+                for part in resource_config
+                if part is not None
+            ]
+            cluster_name = f"cluster-{orchestrator_run_id}" + "-".join(
+                cluster_name_parts
+            )
+            unique_resource_configs[step_name] = cluster_name
 
-        sky.launch(
-            task,
-            cluster_name,
-            retry_until_up=settings.retry_until_up,
-            idle_minutes_to_autostop=settings.idle_minutes_to_autostop,
-            down=settings.down,
-            stream_logs=settings.stream_logs,
-            detach_setup=True,
-            detach_run=True,
-        )
+        run = Client().list_pipeline_runs(
+            sort_by="asc:created",
+            size=1,
+            deployment_id=args.deployment_id,
+            status=ExecutionStatus.INITIALIZING,
+        )[0]
 
-        # Wait for pod to finish.
-        logger.info(f"Waiting for pod of step `{step_name}` to start...")
+        logger.info("Fetching pipeline run: %s", run.id)
 
-        current_run = Client().get_pipeline_run(run.id)
+        def run_step_on_skypilot_vm(step_name: str) -> None:
+            """Run a pipeline step in a separate Skypilot VM.
 
-        step_is_finished = False
-        while not step_is_finished:
-            time.sleep(10)
-            current_run = Client().get_pipeline_run(run.id)
+            Args:
+                step_name: Name of the step.
+
+            Raises:
+                Exception: If the step execution fails.
+            """
+            logger.info(f"Running step `{step_name}` on a VM...")
             try:
-                step_is_finished = current_run.steps[
-                    step_name
-                ].status.is_finished
-            except KeyError:
-                # Step is not yet in the run, so we wait for it to appear
-                continue
-
-        # Pop the resource configuration for this step
-        unique_resource_configs.pop(step_name)
-
-        if cluster_name in unique_resource_configs.values():
-            # If there are more steps using this configuration, skip deprovisioning the cluster
-            logger.info(
-                f"Resource configuration for cluster '{cluster_name}' "
-                "is used by subsequent steps. Skipping the deprovisioning of "
-                "the cluster."
-            )
-        else:
-            # If there are no more steps using this configuration, down the cluster
-            logger.info(
-                f"Resource configuration for cluster '{cluster_name}' "
-                "is not used by subsequent steps. deprovisioning the cluster."
-            )
-            sky.down(cluster_name)
+                cluster_name = unique_resource_configs[step_name]
+
+                image = SkypilotBaseOrchestrator.get_image(
+                    deployment=deployment, step_name=step_name
+                )
+
+                step_args = (
+                    StepEntrypointConfiguration.get_entrypoint_arguments(
+                        step_name=step_name, deployment_id=deployment.id
+                    )
+                )
+                arguments_str = " ".join(step_args)
+
+                step = deployment.step_configurations[step_name]
+                settings = cast(
+                    SkypilotBaseOrchestratorSettings,
+                    orchestrator.get_settings(step),
+                )
+                env = get_config_environment_vars()
+                env[ENV_ZENML_SKYPILOT_ORCHESTRATOR_RUN_ID] = (
+                    orchestrator_run_id
+                )
+
+                # Create the Docker run command
+                run_command = create_docker_run_command(
+                    image=image,
+                    entrypoint_str=entrypoint_str,
+                    arguments_str=arguments_str,
+                    environment=env,
+                    docker_run_args=settings.docker_run_args,
+                    use_sudo=False,  # Entrypoint doesn't use sudo
+                )
+
+                task_name = f"{deployment.id}-{step_name}-{time.time()}"
+
+                # Create task kwargs
+                task_kwargs = prepare_task_kwargs(
+                    settings=settings,
+                    run_command=run_command,
+                    setup=setup,
+                    task_envs=task_envs,
+                    task_name=task_name,
+                )
+
+                task = sky.Task(**task_kwargs)
+
+                # Set resources
+                resources_kwargs = prepare_resources_kwargs(
+                    cloud=orchestrator.cloud,
+                    settings=settings,
+                    default_instance_type=orchestrator.DEFAULT_INSTANCE_TYPE,
+                )
+
+                task = task.set_resources(sky.Resources(**resources_kwargs))
+
+                # Prepare launch parameters
+                launch_kwargs = prepare_launch_kwargs(
+                    settings=settings,
+                )
+
+                # sky.launch now returns a request ID (async). Capture it so we can
+                # optionally stream logs and block until completion when desired.
+                launch_request_id = sky.launch(
+                    task,
+                    cluster_name,
+                    **launch_kwargs,
+                )
+                sky_job_get(launch_request_id, True, cluster_name)
+
+                # Pop the resource configuration for this step
+                unique_resource_configs.pop(step_name)
+
+                if cluster_name in unique_resource_configs.values():
+                    # If there are more steps using this configuration, skip deprovisioning the cluster
+                    logger.info(
+                        f"Resource configuration for cluster '{cluster_name}' "
+                        "is used by subsequent steps. Skipping the deprovisioning of "
+                        "the cluster."
+                    )
+                else:
+                    # If there are no more steps using this configuration, down the cluster
+                    logger.info(
+                        f"Resource configuration for cluster '{cluster_name}' "
+                        "is not used by subsequent steps. deprovisioning the cluster."
+                    )
+                    down_request_id = sky.down(cluster_name)
+                    # Wait for the cluster to be terminated
+                    sky.stream_and_get(down_request_id)
+
+                logger.info(
+                    f"Running step `{step_name}` on a VM is completed."
+                )
+
+            except Exception as e:
+                logger.error(f"Failed while launching step `{step_name}`: {e}")
+                raise
+
+        dag_runner = ThreadedDagRunner(
+            dag=pipeline_dag, run_fn=run_step_on_skypilot_vm
+        )
+        dag_runner.run()
+
+        failed_nodes = []
+        for node in dag_runner.nodes:
+            if dag_runner.node_states[node] == NodeStatus.FAILED:
+                failed_nodes.append(node)
 
-        logger.info(f"Running step `{step_name}` on a VM is completed.")
+        if failed_nodes:
+            raise Exception(f"One or more steps failed: {failed_nodes}")
 
-    ThreadedDagRunner(dag=pipeline_dag, run_fn=run_step_on_skypilot_vm).run()
+    except Exception as e:
+        logger.error(f"Orchestrator failed: {e}")
 
-    logger.info("Orchestration VM provisioned.")
+        # Try to mark the pipeline run as failed
+        if run:
+            publish_failed_pipeline_run(run.id)
+            logger.info("Marked pipeline run as failed in ZenML.")
+        raise
 
 
 if __name__ == "__main__":