zenml-nightly 0.66.0.dev20240925__tar.gz → 0.66.0.dev20240926__tar.gz

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.66.0.dev20240925
+Version: 0.66.0.dev20240926
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.66.0.dev20240925"
+version = "0.66.0.dev20240926"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.66.0.dev20240926/src/zenml/VERSION

@@ -0,0 +1 @@
+0.66.0.dev20240926

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/zenml/integrations/__init__.py

@@ -69,6 +69,7 @@ from zenml.integrations.skypilot_aws import SkypilotAWSIntegration  # noqa
 from zenml.integrations.skypilot_gcp import SkypilotGCPIntegration  # noqa
 from zenml.integrations.skypilot_azure import SkypilotAzureIntegration  # noqa
 from zenml.integrations.skypilot_lambda import SkypilotLambdaIntegration  # noqa
+from zenml.integrations.skypilot_kubernetes import SkypilotKubernetesIntegration  # noqa
 from zenml.integrations.slack import SlackIntegration  # noqa
 from zenml.integrations.spark import SparkIntegration  # noqa
 from zenml.integrations.tekton import TektonIntegration  # noqa

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/zenml/integrations/constants.py

@@ -64,6 +64,7 @@ SKYPILOT_AWS = "skypilot_aws"
 SKYPILOT_GCP = "skypilot_gcp"
 SKYPILOT_AZURE = "skypilot_azure"
 SKYPILOT_LAMBDA = "skypilot_lambda"
+SKYPILOT_KUBERNETES = "skypilot_kubernetes"
 SLACK = "slack"
 SPARK = "spark"
 TEKTON = "tekton"

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/zenml/integrations/deepchecks/__init__.py

@@ -35,7 +35,7 @@ class DeepchecksIntegration(Integration):
 
     NAME = DEEPCHECKS
     REQUIREMENTS = [
-        "deepchecks[vision]>=0.18.0",
+        "deepchecks[vision]~=0.18.0",
         "torchvision>=0.14.0",
         "opencv-python==4.5.5.64",  # pin to same version
         "opencv-python-headless==4.5.5.64",  # pin to same version

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/zenml/integrations/deepchecks/data_validators/deepchecks_data_validator.py

@@ -17,6 +17,7 @@ from typing import (
     Any,
     ClassVar,
     Dict,
+    List,
     Optional,
     Sequence,
     Tuple,
@@ -28,9 +29,8 @@ import pandas as pd
 from deepchecks.core.checks import BaseCheck
 from deepchecks.core.suite import SuiteResult
 from deepchecks.tabular import Dataset as TabularData
+from deepchecks.tabular import ModelComparisonSuite
 from deepchecks.tabular import Suite as TabularSuite
-
-# not part of deepchecks.tabular.checks
 from deepchecks.tabular.suites import full_suite as full_tabular_suite
 from deepchecks.vision import Suite as VisionSuite
 from deepchecks.vision import VisionData
@@ -102,7 +102,7 @@ class DeepchecksDataValidator(BaseDataValidator):
         comparison_dataset: Optional[
             Union[pd.DataFrame, DataLoader[Any]]
         ] = None,
-        model: Optional[Union[ClassifierMixin, Module]] = None,
+        models: Optional[List[Union[ClassifierMixin, Module]]] = None,
         check_list: Optional[Sequence[str]] = None,
         dataset_kwargs: Dict[str, Any] = {},
         check_kwargs: Dict[str, Dict[str, Any]] = {},
@@ -123,7 +123,7 @@ class DeepchecksDataValidator(BaseDataValidator):
                 validation.
             comparison_dataset: Optional secondary (comparison) dataset argument
                 used during comparison checks.
-            model: Optional model argument used during validation.
+            models: Optional model argument used during validation.
             check_list: Optional list of ZenML Deepchecks check identifiers
                 specifying the list of Deepchecks checks to be performed.
             dataset_kwargs: Additional keyword arguments to be passed to the
@@ -149,6 +149,7 @@ class DeepchecksDataValidator(BaseDataValidator):
         # arguments and the check list.
         is_tabular = False
         is_vision = False
+        is_multi_model = False
         for dataset in [reference_dataset, comparison_dataset]:
             if dataset is None:
                 continue
@@ -163,7 +164,18 @@ class DeepchecksDataValidator(BaseDataValidator):
                     f"data and {str(DataLoader)} for computer vision data."
                 )
 
-        if model:
+        if models:
+            # if there's more than one models, we should set the
+            # is_multi_model to True
+            if len(models) > 1:
+                is_multi_model = True
+            # if the models are of different types, raise an error
+            # only the same type of models can be used for comparison
+            if len(set(type(model) for model in models)) > 1:
+                raise TypeError(
+                    "Models used for comparison checks must be of the same type."
+                )
+            model = models[0]
             if isinstance(model, ClassifierMixin):
                 is_tabular = True
             elif isinstance(model, Module):
@@ -190,8 +202,18 @@ class DeepchecksDataValidator(BaseDataValidator):
         if not check_list:
             # default to executing all the checks listed in the supplied
             # checks enum type if a custom check list is not supplied
+            # don't include the TABULAR_PERFORMANCE_BIAS check enum value
+            # as it requires a protected feature name to be set
+            checks_to_exclude = [
+                DeepchecksModelValidationCheck.TABULAR_PERFORMANCE_BIAS
+            ]
+            check_enum_values = [
+                check.value
+                for check in check_enum
+                if check not in checks_to_exclude
+            ]
             tabular_checks, vision_checks = cls._split_checks(
-                check_enum.values()
+                check_enum_values
             )
             if is_tabular:
                 check_list = tabular_checks
@@ -254,6 +276,10 @@ class DeepchecksDataValidator(BaseDataValidator):
             suite_class = VisionSuite
             full_suite = full_vision_suite()
 
+        # if is_multi_model is True, we need to use the ModelComparisonSuite
+        if is_multi_model:
+            suite_class = ModelComparisonSuite
+
         train_dataset = dataset_class(reference_dataset, **dataset_kwargs)
         test_dataset = None
         if comparison_dataset is not None:
@@ -294,13 +320,28 @@ class DeepchecksDataValidator(BaseDataValidator):
                     continue
                 condition_method(**condition_kwargs)
 
-            suite.add(check)
-        return suite.run(
-            train_dataset=train_dataset,
-            test_dataset=test_dataset,
-            model=model,
-            **run_kwargs,
-        )
+            # if the check is supported by the suite, add it
+            if isinstance(check, suite.supported_checks()):
+                suite.add(check)
+            else:
+                logger.warning(
+                    f"Check {check_name} is not supported by the {suite_class} "
+                    "suite. Ignoring the check."
+                )
+
+        if isinstance(suite, ModelComparisonSuite):
+            return suite.run(
+                models=models,
+                train_datasets=train_dataset,
+                test_datasets=test_dataset,
+            )
+        else:
+            return suite.run(
+                train_dataset=train_dataset,
+                test_dataset=test_dataset,
+                model=models[0] if models else None,
+                **run_kwargs,
+            )
 
     def data_validation(
         self,
@@ -444,7 +485,7 @@ class DeepchecksDataValidator(BaseDataValidator):
             check_enum=check_enum,
             reference_dataset=dataset,
             comparison_dataset=comparison_dataset,
-            model=model,
+            models=[model],
             check_list=check_list,
             dataset_kwargs=dataset_kwargs,
             check_kwargs=check_kwargs,

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/zenml/integrations/deepchecks/validation_checks.py

@@ -153,8 +153,8 @@ class DeepchecksDataIntegrityCheck(DeepchecksValidationCheck):
 
     This list reflects the set of data integrity checks provided by Deepchecks:
 
-      * [for tabular data](https://docs.deepchecks.com/en/stable/checks_gallery/tabular.html#data-integrity)
-      * [for computer vision](https://docs.deepchecks.com/en/stable/checks_gallery/vision.html#data-integrity)
+      * [for tabular data](https://docs.deepchecks.com/stable/tabular/auto_checks/data_integrity/index.html)
+      * [for computer vision](https://docs.deepchecks.com/stable/vision/auto_checks/data_integrity/index.html)
 
     All these checks inherit from `deepchecks.tabular.SingleDatasetCheck` or
     `deepchecks.vision.SingleDatasetCheck` and require a single dataset as input.
@@ -176,6 +176,9 @@ class DeepchecksDataIntegrityCheck(DeepchecksValidationCheck):
     TABULAR_FEATURE_LABEL_CORRELATION = source_utils.resolve(
         tabular_checks.FeatureLabelCorrelation
     ).import_path
+    TABULAR_IDENTIFIER_LABEL_CORRELATION = source_utils.resolve(
+        tabular_checks.IdentifierLabelCorrelation
+    ).import_path
     TABULAR_IS_SINGLE_VALUE = source_utils.resolve(
         tabular_checks.IsSingleValue
     ).import_path
@@ -197,6 +200,12 @@ class DeepchecksDataIntegrityCheck(DeepchecksValidationCheck):
     TABULAR_STRING_MISMATCH = source_utils.resolve(
         tabular_checks.StringMismatch
     ).import_path
+    TABULAR_CLASS_IMBALANCE = source_utils.resolve(
+        tabular_checks.ClassImbalance
+    ).import_path
+    TABULAR_PERCENT_OF_NULLS = source_utils.resolve(
+        tabular_checks.PercentOfNulls
+    ).import_path
 
     VISION_IMAGE_PROPERTY_OUTLIERS = source_utils.resolve(
         vision_checks.ImagePropertyOutliers
@@ -204,6 +213,9 @@ class DeepchecksDataIntegrityCheck(DeepchecksValidationCheck):
     VISION_LABEL_PROPERTY_OUTLIERS = source_utils.resolve(
         vision_checks.LabelPropertyOutliers
     ).import_path
+    VISION_PROPERTY_LABEL_CORRELATION = source_utils.resolve(
+        vision_checks.PropertyLabelCorrelation
+    ).import_path
 
 
 class DeepchecksDataDriftCheck(DeepchecksValidationCheck):
@@ -246,19 +258,37 @@ class DeepchecksDataDriftCheck(DeepchecksValidationCheck):
     TABULAR_TRAIN_TEST_FEATURE_DRIFT = source_utils.resolve(
         tabular_checks.TrainTestFeatureDrift
     ).import_path
+    TABULAR_FEATURE_DRIFT = source_utils.resolve(
+        tabular_checks.FeatureDrift
+    ).import_path
     TABULAR_TRAIN_TEST_LABEL_DRIFT = source_utils.resolve(
         tabular_checks.TrainTestLabelDrift
     ).import_path
+    TABULAR_LABEL_DRIFT = source_utils.resolve(
+        tabular_checks.LabelDrift
+    ).import_path
     TABULAR_TRAIN_TEST_SAMPLES_MIX = source_utils.resolve(
         tabular_checks.TrainTestSamplesMix
     ).import_path
     TABULAR_WHOLE_DATASET_DRIFT = source_utils.resolve(
         tabular_checks.WholeDatasetDrift
     ).import_path
+    TABULAR_NEW_CATEGORY_TRAIN_TEST = source_utils.resolve(
+        tabular_checks.NewCategoryTrainTest
+    ).import_path
+    TABULAR_MULTIVARIATE_DRIFT = source_utils.resolve(
+        tabular_checks.MultivariateDrift
+    ).import_path
 
+    VISION_PROPERTY_LABEL_CORRELATION_CHANGE = source_utils.resolve(
+        vision_checks.PropertyLabelCorrelationChange
+    ).import_path
     VISION_HEATMAP_COMPARISON = source_utils.resolve(
         vision_checks.HeatmapComparison
     ).import_path
+    VISION_LABEL_DRIFT = source_utils.resolve(
+        vision_checks.LabelDrift
+    ).import_path
     VISION_IMAGE_DATASET_DRIFT = source_utils.resolve(
         vision_checks.ImageDatasetDrift
     ).import_path
@@ -268,9 +298,6 @@ class DeepchecksDataDriftCheck(DeepchecksValidationCheck):
     VISION_NEW_LABELS = source_utils.resolve(
         vision_checks.NewLabels
     ).import_path
-    VISION_TRAIN_TEST_LABEL_DRIFT = source_utils.resolve(
-        vision_checks.TrainTestLabelDrift
-    ).import_path
 
 
 class DeepchecksModelValidationCheck(DeepchecksValidationCheck):
@@ -296,6 +323,12 @@ class DeepchecksModelValidationCheck(DeepchecksValidationCheck):
     TABULAR_MODEL_INFERENCE_TIME = source_utils.resolve(
         tabular_checks.ModelInferenceTime
     ).import_path
+    TABULAR_MODEL_INFO = source_utils.resolve(
+        tabular_checks.ModelInfo
+    ).import_path
+    TABULAR_PERFORMANCE_BIAS = source_utils.resolve(
+        tabular_checks.model_evaluation.PerformanceBias
+    ).import_path
     TABULAR_REGRESSION_ERROR_DISTRIBUTION = source_utils.resolve(
         tabular_checks.RegressionErrorDistribution
     ).import_path
@@ -308,6 +341,18 @@ class DeepchecksModelValidationCheck(DeepchecksValidationCheck):
     TABULAR_SEGMENT_PERFORMANCE = source_utils.resolve(
         tabular_checks.SegmentPerformance
     ).import_path
+    TABULAR_WEAK_SEGMENT_PERFORMANCE = source_utils.resolve(
+        tabular_checks.WeakSegmentsPerformance
+    ).import_path
+    TABULAR_SINGLE_DATASET_PERFORMANCE = source_utils.resolve(
+        tabular_checks.SingleDatasetPerformance
+    ).import_path
+    TABULAR_TRAIN_TEST_PERFORMANCE = source_utils.resolve(
+        tabular_checks.TrainTestPerformance
+    ).import_path
+    TABULAR_MULTI_MODEL_PERFORMANCE_REPORT = source_utils.resolve(
+        tabular_checks.MultiModelPerformanceReport
+    ).import_path
 
     VISION_CONFUSION_MATRIX_REPORT = source_utils.resolve(
         vision_checks.ConfusionMatrixReport
@@ -318,6 +363,12 @@ class DeepchecksModelValidationCheck(DeepchecksValidationCheck):
     VISION_MEAN_AVERAGE_RECALL_REPORT = source_utils.resolve(
         vision_checks.MeanAverageRecallReport
     ).import_path
+    VISION_SINGLE_DATASET_PERFORMANCE = source_utils.resolve(
+        vision_checks.SingleDatasetPerformance
+    ).import_path
+    VISION_WEAK_SEGMENT_PERFORMANCE = source_utils.resolve(
+        vision_checks.WeakSegmentsPerformance
+    ).import_path
 
 
 class DeepchecksModelDriftCheck(DeepchecksValidationCheck):
@@ -343,6 +394,9 @@ class DeepchecksModelDriftCheck(DeepchecksValidationCheck):
     TABULAR_TRAIN_TEST_PREDICTION_DRIFT = source_utils.resolve(
         tabular_checks.TrainTestPredictionDrift
     ).import_path
+    TABULAR_PREDICTION_DRIFT = source_utils.resolve(
+        tabular_checks.PredictionDrift
+    ).import_path
     TABULAR_UNUSED_FEATURES = source_utils.resolve(
         tabular_checks.UnusedFeatures
     ).import_path
@@ -356,3 +410,6 @@ class DeepchecksModelDriftCheck(DeepchecksValidationCheck):
     VISION_TRAIN_TEST_PREDICTION_DRIFT = source_utils.resolve(
         vision_checks.TrainTestPredictionDrift
     ).import_path
+    VISION_PREDICTION_DRIFT = source_utils.resolve(
+        vision_checks.PredictionDrift
+    ).import_path

{zenml_nightly-0.66.0.dev20240925 → zenml_nightly-0.66.0.dev20240926}/src/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_nightly-0.66.0.dev20240926/src/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_nightly-0.66.0.dev20240926/src/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",
+]

zenml_nightly-0.66.0.dev20240926/src/zenml/integrations/skypilot_kubernetes/flavors/skypilot_orchestrator_kubernetes_vm_flavor.py

@@ -0,0 +1,125 @@
+#  Copyright (c) ZenML GmbH 2023. 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 orchestrator Kubernetes flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
+from zenml.integrations.skypilot.flavors.skypilot_orchestrator_base_vm_config import (
+    SkypilotBaseOrchestratorConfig,
+    SkypilotBaseOrchestratorSettings,
+)
+from zenml.integrations.skypilot_kubernetes import (
+    SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR,
+)
+from zenml.logger import get_logger
+from zenml.models import ServiceConnectorRequirements
+from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.skypilot_kubernetes.orchestrators import (
+        SkypilotKubernetesOrchestrator,
+    )
+
+
+logger = get_logger(__name__)
+
+
+class SkypilotKubernetesOrchestratorSettings(SkypilotBaseOrchestratorSettings):
+    """Skypilot orchestrator settings."""
+
+
+class SkypilotKubernetesOrchestratorConfig(
+    SkypilotBaseOrchestratorConfig, SkypilotKubernetesOrchestratorSettings
+):
+    """Skypilot orchestrator config."""
+
+
+class SkypilotKubernetesOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Flavor for the Skypilot Kubernetes orchestrator."""
+
+    @property
+    def name(self) -> str:
+        """Name of the orchestrator flavor.
+
+        Returns:
+            Name of the orchestrator flavor.
+        """
+        return SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE,
+        )
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/kubernetes-skypilot.png"
+
+    @property
+    def config_class(self) -> Type[BaseOrchestratorConfig]:
+        """Config class for the base orchestrator flavor.
+
+        Returns:
+            The config class.
+        """
+        return SkypilotKubernetesOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["SkypilotKubernetesOrchestrator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            Implementation class for this flavor.
+        """
+        from zenml.integrations.skypilot_kubernetes.orchestrators import (
+            SkypilotKubernetesOrchestrator,
+        )
+
+        return SkypilotKubernetesOrchestrator