zenml-nightly 0.66.0.dev20240923__py3-none-any.whl → 0.66.0.dev20240928__py3-none-any.whl

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -19,6 +19,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    Iterator,
     List,
     Optional,
     Tuple,
@@ -46,8 +47,11 @@ from azure.identity import DefaultAzureCredential
 
 from zenml.config.base_settings import BaseSettings
 from zenml.config.step_configurations import Step
-from zenml.constants import METADATA_ORCHESTRATOR_URL
-from zenml.enums import StackComponentType
+from zenml.constants import (
+    METADATA_ORCHESTRATOR_RUN_ID,
+    METADATA_ORCHESTRATOR_URL,
+)
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.integrations.azure.azureml_utils import create_or_get_compute
 from zenml.integrations.azure.flavors.azureml import AzureMLComputeTypes
 from zenml.integrations.azure.flavors.azureml_orchestrator_flavor import (
@@ -65,7 +69,7 @@ from zenml.stack import StackValidator
 from zenml.utils.string_utils import b64_encode
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
     from zenml.stack import Stack
 
 logger = get_logger(__name__)
@@ -199,7 +203,7 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> None:
+    ) -> Iterator[Dict[str, MetadataType]]:
         """Prepares or runs a pipeline on AzureML.
 
         Args:
@@ -210,6 +214,9 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
 
         Raises:
             RuntimeError: If the creation of the schedule fails.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
         """
         # Authentication
         if connector := self.get_connector():
@@ -379,6 +386,10 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         else:
             job = ml_client.jobs.create_or_update(pipeline_job)
             logger.info(f"Pipeline {run_name} has been started.")
+
+            # Yield metadata based on the generated job object
+            yield from self.compute_metadata(job)
+
             assert job.services is not None
             assert job.name is not None
 
@@ -428,3 +439,145 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
                 f"job: {e}"
             )
             return {}
+
+    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 execution.
+
+        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 AzureML client
+        if connector := self.get_connector():
+            credentials = connector.connect()
+        else:
+            credentials = DefaultAzureCredential()
+
+        ml_client = MLClient(
+            credential=credentials,
+            subscription_id=self.config.subscription_id,
+            resource_group_name=self.config.resource_group,
+            workspace_name=self.config.workspace,
+        )
+
+        # 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 = ml_client.jobs.get(run_id).status
+
+        # Map the potential outputs to ZenML ExecutionStatus. Potential values:
+        # https://learn.microsoft.com/en-us/python/api/azure-ai-ml/azure.ai.ml.entities.pipelinejob?view=azure-python#azure-ai-ml-entities-pipelinejob-status
+        if status in [
+            "NotStarted",
+            "Starting",
+            "Provisioning",
+            "Preparing",
+            "Queued",
+        ]:
+            return ExecutionStatus.INITIALIZING
+        elif status in ["Running", "Finalizing"]:
+            return ExecutionStatus.RUNNING
+        elif status in [
+            "CancelRequested",
+            "Failed",
+            "Canceled",
+            "NotResponding",
+        ]:
+            return ExecutionStatus.FAILED
+        elif status in ["Completed"]:
+            return ExecutionStatus.COMPLETED
+        else:
+            raise ValueError("Unknown status for the pipeline job.")
+
+    def compute_metadata(self, job: Any) -> Iterator[Dict[str, MetadataType]]:
+        """Generate run metadata based on the generated AzureML PipelineJob.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
+        """
+        # Metadata
+        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 AzureML's pipeline view
+        if orchestrator_url := self._compute_orchestrator_url(job):
+            metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_url)
+
+        yield metadata
+
+    @staticmethod
+    def _compute_orchestrator_url(job: Any) -> Optional[str]:
+        """Generate the Orchestrator Dashboard URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+             the URL to the dashboard view in AzureML.
+        """
+        try:
+            if job.studio_url:
+                return str(job.studio_url)
+
+            return None
+
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline url: {e}"
+            )
+            return None
+
+    @staticmethod
+    def _compute_orchestrator_run_id(job: Any) -> Optional[str]:
+        """Generate the Orchestrator Dashboard URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+             the URL to the dashboard view in AzureML.
+        """
+        try:
+            if job.name:
+                return str(job.name)
+
+            return None
+
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline run ID: {e}"
+            )
+            return None

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/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/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/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