zenml-nightly 0.66.0.dev20240924__py3-none-any.whl → 0.66.0.dev20240926__py3-none-any.whl

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -32,10 +32,12 @@
 import os
 import re
 import types
+import urllib
 from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    Iterator,
     List,
     Optional,
     Tuple,
@@ -46,15 +48,18 @@ from uuid import UUID
 
 from google.api_core import exceptions as google_exceptions
 from google.cloud import aiplatform
+from google.cloud.aiplatform_v1.types import PipelineState
 from kfp import dsl
 from kfp.compiler import Compiler
 
 from zenml.config.resource_settings import ResourceSettings
 from zenml.constants import (
+    METADATA_ORCHESTRATOR_LOGS_URL,
+    METADATA_ORCHESTRATOR_RUN_ID,
     METADATA_ORCHESTRATOR_URL,
 )
 from zenml.entrypoints import StepEntrypointConfiguration
-from zenml.enums import StackComponentType
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.integrations.gcp import GCP_ARTIFACT_STORE_FLAVOR
 from zenml.integrations.gcp.constants import (
     GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
@@ -77,7 +82,11 @@ from zenml.utils.io_utils import get_global_config_directory
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
-    from zenml.models import PipelineDeploymentResponse, ScheduleResponse
+    from zenml.models import (
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+        ScheduleResponse,
+    )
     from zenml.stack import Stack
 
 logger = get_logger(__name__)
@@ -245,8 +254,8 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             ):
                 logger.warning(
                     "Vertex orchestrator only uses schedules with the "
-                    "`cron_expression` property, with optional `start_time` and/or `end_time`. "
-                    "All other properties are ignored."
+                    "`cron_expression` property, with optional `start_time` "
+                    "and/or `end_time`. All other properties are ignored."
                 )
             if deployment.schedule.cron_expression is None:
                 raise ValueError(
@@ -302,7 +311,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> Any:
+    ) -> Iterator[Dict[str, MetadataType]]:
         """Creates a KFP JSON pipeline.
 
         # noqa: DAR402
@@ -337,12 +346,15 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 environment.
 
         Raises:
-            ValueError: If the attribute `pipeline_root` is not set and it
+            ValueError: If the attribute `pipeline_root` is not set, and it
                 can be not generated using the path of the artifact store in the
                 stack because it is not a
                 `zenml.integrations.gcp.artifact_store.GCPArtifactStore`. Also gets
                 raised if attempting to schedule pipeline run without using the
                 `zenml.integrations.gcp.artifact_store.GCPArtifactStore`.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
         """
         orchestrator_run_name = get_orchestrator_run_name(
             pipeline_name=deployment.pipeline_configuration.name
@@ -556,15 +568,15 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         )
 
         # Using the Google Cloud AIPlatform client, upload and execute the
-        # pipeline
-        # on the Vertex AI Pipelines service.
-        self._upload_and_run_pipeline(
+        # pipeline on the Vertex AI Pipelines service.
+        if metadata := self._upload_and_run_pipeline(
             pipeline_name=deployment.pipeline_configuration.name,
             pipeline_file_path=pipeline_file_path,
             run_name=orchestrator_run_name,
             settings=settings,
             schedule=deployment.schedule,
-        )
+        ):
+            yield from metadata
 
     def _upload_and_run_pipeline(
         self,
@@ -573,7 +585,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         run_name: str,
         settings: VertexOrchestratorSettings,
         schedule: Optional["ScheduleResponse"] = None,
-    ) -> None:
+    ) -> Iterator[Dict[str, MetadataType]]:
         """Uploads and run the pipeline on the Vertex AI Pipelines service.
 
         Args:
@@ -585,7 +597,11 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             schedule: The schedule the pipeline will run on.
 
         Raises:
-            RuntimeError: If the Vertex Orchestrator fails to provision or any other Runtime errors
+            RuntimeError: If the Vertex Orchestrator fails to provision or any
+                other Runtime errors.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
         """
         # We have to replace the hyphens in the run name with underscores
         # and lower case the string, because the Vertex AI Pipelines service
@@ -593,8 +609,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         job_id = _clean_pipeline_name(run_name)
 
         # Get the credentials that would be used to create the Vertex AI
-        # Pipelines
-        # job.
+        # Pipelines job.
         credentials, project_id = self._get_authentication()
 
         # Instantiate the Vertex AI Pipelines job
@@ -629,7 +644,8 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         try:
             if schedule:
                 logger.info(
-                    "Scheduling job using native Vertex AI Pipelines scheduling..."
+                    "Scheduling job using native Vertex AI Pipelines "
+                    "scheduling..."
                 )
                 run.create_schedule(
                     display_name=schedule.name,
@@ -645,13 +661,12 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     "No schedule detected. Creating one-off Vertex job..."
                 )
                 logger.info(
-                    "Submitting pipeline job with job_id `%s` to Vertex AI Pipelines "
-                    "service.",
+                    "Submitting pipeline job with job_id `%s` to Vertex AI "
+                    "Pipelines service.",
                     job_id,
                 )
 
                 # Submit the job to Vertex AI Pipelines service.
-
                 run.submit(
                     service_account=self.config.workload_service_account,
                     network=self.config.network,
@@ -661,6 +676,9 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     run._dashboard_uri(),
                 )
 
+                # Yield metadata based on the generated job object
+                yield from self.compute_metadata(run)
+
                 if settings.synchronous:
                     logger.info(
                         "Waiting for the Vertex AI Pipelines job to finish..."
@@ -738,6 +756,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             The dynamic component with the resource settings applied.
         """
         # Set optional CPU, RAM and GPU constraints for the pipeline
+        cpu_limit = None
         if resource_settings:
             cpu_limit = resource_settings.cpu_count or self.config.cpu_limit
 
@@ -778,3 +797,173 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 )
 
         return dynamic_component
+
+    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: The run that was executed by this orchestrator.
+
+        Returns:
+            the actual status of the pipeline job.
+
+        Raises:
+            AssertionError: If the run was not executed by to this orchestrator.
+            ValueError: If it fetches an unknown state or if we can not fetch
+                the orchestrator run ID.
+        """
+        # Make sure that the stack exists and is accessible
+        if run.stack is None:
+            raise ValueError(
+                "The stack that the run was executed on is not available "
+                "anymore."
+            )
+
+        # Make sure that the run belongs to this orchestrator
+        assert (
+            self.id
+            == run.stack.components[StackComponentType.ORCHESTRATOR][0].id
+        )
+
+        # Initialize the Vertex client
+        credentials, project_id = self._get_authentication()
+        aiplatform.init(
+            project=project_id,
+            location=self.config.location,
+            credentials=credentials,
+        )
+
+        # Fetch the status of the PipelineJob
+        if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+        elif run.orchestrator_run_id is not None:
+            run_id = run.orchestrator_run_id
+        else:
+            raise ValueError(
+                "Can not find the orchestrator run ID, thus can not fetch "
+                "the status."
+            )
+        status = aiplatform.PipelineJob.get(run_id).state
+
+        # Map the potential outputs to ZenML ExecutionStatus. Potential values:
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/describe_pipeline_execution.html#
+        if status in [PipelineState.PIPELINE_STATE_UNSPECIFIED]:
+            return run.status
+        elif status in [
+            PipelineState.PIPELINE_STATE_QUEUED,
+            PipelineState.PIPELINE_STATE_PENDING,
+        ]:
+            return ExecutionStatus.INITIALIZING
+        elif status in [
+            PipelineState.PIPELINE_STATE_RUNNING,
+            PipelineState.PIPELINE_STATE_PAUSED,
+        ]:
+            return ExecutionStatus.RUNNING
+        elif status in [PipelineState.PIPELINE_STATE_SUCCEEDED]:
+            return ExecutionStatus.COMPLETED
+
+        elif status in [
+            PipelineState.PIPELINE_STATE_FAILED,
+            PipelineState.PIPELINE_STATE_CANCELLING,
+            PipelineState.PIPELINE_STATE_CANCELLED,
+        ]:
+            return ExecutionStatus.FAILED
+        else:
+            raise ValueError("Unknown status for the pipeline job.")
+
+    def compute_metadata(
+        self, job: aiplatform.PipelineJob
+    ) -> Iterator[Dict[str, MetadataType]]:
+        """Generate run metadata based on the corresponding Vertex PipelineJob.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
+        """
+        metadata: Dict[str, MetadataType] = {}
+
+        # Orchestrator Run ID
+        if run_id := self._compute_orchestrator_run_id(job):
+            metadata[METADATA_ORCHESTRATOR_RUN_ID] = run_id
+
+        # URL to the Vertex's pipeline view
+        if orchestrator_url := self._compute_orchestrator_url(job):
+            metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_url)
+
+        # URL to the corresponding Logs Explorer page
+        if logs_url := self._compute_orchestrator_logs_url(job):
+            metadata[METADATA_ORCHESTRATOR_LOGS_URL] = Uri(logs_url)
+
+        yield metadata
+
+    @staticmethod
+    def _compute_orchestrator_url(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Generate the Orchestrator Dashboard URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+             the URL to the dashboard view in Vertex.
+        """
+        try:
+            return str(job._dashboard_uri())
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline url: {e}"
+            )
+            return None
+
+    @staticmethod
+    def _compute_orchestrator_logs_url(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Generate the Logs Explorer URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+            the URL querying the pipeline logs in Logs Explorer on GCP.
+        """
+        try:
+            base_url = "https://console.cloud.google.com/logs/query"
+            query = f"""
+             resource.type="aiplatform.googleapis.com/PipelineJob"
+             resource.labels.pipeline_job_id="{job.job_id}"
+             """
+            encoded_query = urllib.parse.quote(query)
+            return f"{base_url}?project={job.project}&query={encoded_query}"
+
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the logs url: {e}"
+            )
+            return None
+
+    @staticmethod
+    def _compute_orchestrator_run_id(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Fetch the Orchestrator Run ID upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+            the Execution ID of the run in Vertex.
+        """
+        try:
+            if job.job_id:
+                return str(job.job_id)
+
+            return None
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline run ID: {e}"
+            )
+            return None

zenml/integrations/lightning/__init__.py

@@ -28,7 +28,7 @@ class LightningIntegration(Integration):
     """Definition of Lightning Integration for ZenML."""
 
     NAME = LIGHTNING
-    REQUIREMENTS = ["lightning-sdk"]
+    REQUIREMENTS = ["lightning-sdk>=0.1.17"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py

@@ -85,6 +85,15 @@ class LightningOrchestratorConfig(
         """
         return self.synchronous
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return False
+
 
 class LightningOrchestratorFlavor(BaseOrchestratorFlavor):
     """Lightning orchestrator flavor."""

zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

@@ -103,20 +103,29 @@ class LightningOrchestrator(WheeledOrchestrator):
 
         Args:
             deployment: The pipeline deployment to prepare or run.
+
+        Raises:
+            ValueError: If the user id and api key or username and organization
         """
         settings = cast(
             LightningOrchestratorSettings, self.get_settings(deployment)
         )
-        if settings.user_id:
-            os.environ["LIGHTNING_USER_ID"] = settings.user_id
-        if settings.api_key:
-            os.environ["LIGHTNING_API_KEY"] = settings.api_key
+        if not settings.user_id or not settings.api_key:
+            raise ValueError(
+                "Lightning orchestrator requires `user_id` and `api_key` both to be set in the settings."
+            )
+        os.environ["LIGHTNING_USER_ID"] = settings.user_id
+        os.environ["LIGHTNING_API_KEY"] = settings.api_key
         if settings.username:
             os.environ["LIGHTNING_USERNAME"] = settings.username
+        elif settings.organization:
+            os.environ["LIGHTNING_ORG"] = settings.organization
+        else:
+            raise ValueError(
+                "Lightning orchestrator requires either `username` or `organization` to be set in the settings."
+            )
         if settings.teamspace:
             os.environ["LIGHTNING_TEAMSPACE"] = settings.teamspace
-        if settings.organization:
-            os.environ["LIGHTNING_ORG"] = settings.organization
 
     @property
     def config(self) -> LightningOrchestratorConfig:
@@ -267,9 +276,7 @@ class LightningOrchestrator(WheeledOrchestrator):
         ) as code_file:
             code_archive.write_archive(code_file)
             code_path = code_file.name
-
         filename = f"{orchestrator_run_name}.tar.gz"
-
         # Construct the env variables for the pipeline
         env_vars = environment.copy()
         orchestrator_run_id = str(uuid4())
@@ -392,9 +399,7 @@ class LightningOrchestrator(WheeledOrchestrator):
                 f"Installing requirements: {pipeline_requirements_to_string}"
             )
             studio.run(f"uv pip install {pipeline_requirements_to_string}")
-            studio.run(
-                "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-            )
+            studio.run("pip install zenml -y")
 
             for custom_command in settings.custom_commands or []:
                 studio.run(
@@ -488,9 +493,7 @@ class LightningOrchestrator(WheeledOrchestrator):
             )
             studio.run("pip install uv")
             studio.run(f"uv pip install {requirements}")
-            studio.run(
-                "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-            )
+            studio.run("pip install zenml -y")
             # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
             for command in settings.custom_commands or []:
                 output = studio.run(
@@ -563,9 +566,7 @@ class LightningOrchestrator(WheeledOrchestrator):
         )
         studio.run("pip install uv")
         studio.run(f"uv pip install {details['requirements']}")
-        studio.run(
-            "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-        )
+        studio.run("pip install zenml -y")
         # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
         for command in custom_commands or []:
             output = studio.run(

zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

@@ -166,9 +166,7 @@ def main() -> None:
         f"uv pip install {pipeline_requirements_to_string}"
     )
     logger.info(output)
-    output = main_studio.run(
-        "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-    )
+    output = main_studio.run("pip install zenml -y")
     logger.info(output)
 
     for command in pipeline_settings.custom_commands or []:
@@ -250,9 +248,7 @@ def main() -> None:
                     f"uv pip install {step_requirements_to_string}"
                 )
                 logger.info(output)
-                output = studio.run(
-                    "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-                )
+                output = studio.run("pip install zenml -y")
                 logger.info(output)
                 for command in step_settings.custom_commands or []:
                     output = studio.run(

zenml/integrations/mlflow/steps/mlflow_registry.py

@@ -146,6 +146,8 @@ def mlflow_register_model_step(
         metadata.zenml_pipeline_run_uuid = pipeline_run_uuid
     if metadata.zenml_workspace is None:
         metadata.zenml_workspace = zenml_workspace
+    if getattr(metadata, "mlflow_run_id", None) is None:
+        setattr(metadata, "mlflow_run_id", mlflow_run_id)
 
     # Register model version
     model_version = model_registry.register_model_version(

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -250,6 +250,7 @@ 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()
         )
@@ -271,13 +272,10 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                 f"sudo docker login --username $DOCKER_USERNAME --password "
                 f"$DOCKER_PASSWORD {stack.container_registry.config.uri}"
             )
-            task_envs = {
-                "DOCKER_USERNAME": docker_username,
-                "DOCKER_PASSWORD": docker_password,
-            }
+            task_envs["DOCKER_USERNAME"] = docker_username
+            task_envs["DOCKER_PASSWORD"] = docker_password
         else:
             setup = None
-            task_envs = None
 
         # Run the entire pipeline
 
@@ -285,15 +283,22 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
         self.prepare_environment_variable(set=True)
 
         try:
+            if isinstance(self.cloud, sky.clouds.Kubernetes):
+                run_command = f"${{VIRTUAL_ENV:+$VIRTUAL_ENV/bin/}}{entrypoint_str} {arguments_str}"
+                setup = None
+                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}"
+                down = settings.down
+                idle_minutes_to_autostop = settings.idle_minutes_to_autostop
             task = sky.Task(
-                run=f"sudo docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}",
+                run=run_command,
                 setup=setup,
                 envs=task_envs,
             )
-            logger.debug(
-                f"Running run: sudo docker run --rm {custom_run_args}{docker_environment_str} {image} {entrypoint_str} {arguments_str}"
-            )
-            logger.debug(f"Running run: {setup}")
+            logger.debug(f"Running run: {run_command}")
+
             task = task.set_resources(
                 sky.Resources(
                     cloud=self.cloud,
@@ -306,15 +311,24 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     job_recovery=settings.job_recovery,
                     region=settings.region,
                     zone=settings.zone,
-                    image_id=settings.image_id,
+                    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 the cluster name
-            cluster_name = settings.cluster_name
-            if cluster_name is None:
+            if settings.cluster_name:
+                sky.exec(
+                    task,
+                    settings.cluster_name,
+                    down=down,
+                    stream_logs=settings.stream_logs,
+                    backend=None,
+                    detach_run=True,
+                )
+            else:
                 # Find existing cluster
                 for i in sky.status(refresh=True):
                     if isinstance(
@@ -324,21 +338,19 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                         logger.info(
                             f"Found existing cluster {cluster_name}. Reusing..."
                         )
-            if cluster_name is None:
                 cluster_name = self.sanitize_cluster_name(
                     f"{orchestrator_run_name}"
                 )
-
-            # Launch the cluster
-            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,
-            )
+                # Launch the cluster
+                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,
+                    detach_setup=True,
+                )
 
         except Exception as e:
             logger.error(f"Pipeline run failed: {e}")

zenml/integrations/skypilot_kubernetes/__init__.py

@@ -0,0 +1,52 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Initialization of the Skypilot Kubernetes integration for ZenML.
+
+The Skypilot integration sub-module powers an alternative to the local
+orchestrator for a remote orchestration of ZenML pipelines on VMs.
+"""
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    SKYPILOT_KUBERNETES,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR = "vm_kubernetes"
+
+
+class SkypilotKubernetesIntegration(Integration):
+    """Definition of Skypilot Kubernetes Integration for ZenML."""
+
+    NAME = SKYPILOT_KUBERNETES
+    # all 0.6.x versions of skypilot[kubernetes] are compatible
+    REQUIREMENTS = ["skypilot[kubernetes]~=0.6.1"]
+    APT_PACKAGES = ["openssh-client", "rsync"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Skypilot Kubernetes integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.skypilot_kubernetes.flavors import (
+            SkypilotKubernetesOrchestratorFlavor,
+        )
+
+        return [SkypilotKubernetesOrchestratorFlavor]
+
+
+SkypilotKubernetesIntegration.check_installation()

zenml/integrations/skypilot_kubernetes/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Skypilot integration flavor for Skypilot Kubernetes orchestrator."""
+
+from zenml.integrations.skypilot_kubernetes.flavors.skypilot_orchestrator_kubernetes_vm_flavor import (
+    SkypilotKubernetesOrchestratorConfig,
+    SkypilotKubernetesOrchestratorFlavor,
+    SkypilotKubernetesOrchestratorSettings,
+)
+
+__all__ = [
+    "SkypilotKubernetesOrchestratorConfig",
+    "SkypilotKubernetesOrchestratorFlavor",
+    "SkypilotKubernetesOrchestratorSettings",
+]