zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  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.
@@ -16,12 +16,21 @@
 import os
 from collections import defaultdict
 from tempfile import TemporaryDirectory, mkdtemp
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 from datasets import Dataset, load_from_disk
 from datasets.dataset_dict import DatasetDict
 
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.pandas_materializer import PandasMaterializer
@@ -33,6 +42,31 @@ if TYPE_CHECKING:
 DEFAULT_DATASET_DIR = "hf_datasets"
 
 
+def extract_repo_name(checksum_str: str) -> Optional[str]:
+    """Extracts the repo name from the checksum string.
+
+    An example of a checksum_str is:
+    "hf://datasets/nyu-mll/glue@bcdcba79d07bc864c1c254ccfcedcce55bcc9a8c/mrpc/train-00000-of-00001.parquet"
+    and the expected output is "nyu-mll/glue".
+
+    Args:
+        checksum_str: The checksum_str to extract the repo name from.
+
+    Returns:
+        str: The extracted repo name.
+    """
+    dataset = None
+    try:
+        parts = checksum_str.split("/")
+        if len(parts) >= 4:
+            # Case: nyu-mll/glue
+            dataset = f"{parts[3]}/{parts[4].split('@')[0]}"
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return dataset
+
+
 class HFDatasetMaterializer(BaseMaterializer):
     """Materializer to read data to and from huggingface datasets."""
 
@@ -103,3 +137,54 @@ class HFDatasetMaterializer(BaseMaterializer):
                     metadata[key][dataset_name] = value
             return dict(metadata)
         raise ValueError(f"Unsupported type {type(ds)}")
+
+    def save_visualizations(
+        self, ds: Union[Dataset, DatasetDict]
+    ) -> Dict[str, VisualizationType]:
+        """Save visualizations for the dataset.
+
+        Args:
+            ds: The Dataset or DatasetDict to visualize.
+
+        Returns:
+            A dictionary mapping visualization paths to their types.
+
+        Raises:
+            ValueError: If the given object is not a `Dataset` or `DatasetDict`.
+        """
+        visualizations = {}
+
+        if isinstance(ds, Dataset):
+            datasets = {"default": ds}
+        elif isinstance(ds, DatasetDict):
+            datasets = ds
+        else:
+            raise ValueError(f"Unsupported type {type(ds)}")
+
+        for name, dataset in datasets.items():
+            # Generate a unique identifier for the dataset
+            if dataset.info.download_checksums:
+                dataset_id = extract_repo_name(
+                    [x for x in dataset.info.download_checksums.keys()][0]
+                )
+                if dataset_id:
+                    # Create the iframe HTML
+                    html = f"""
+                    <iframe
+                    src="https://huggingface.co/datasets/{dataset_id}/embed/viewer"
+                    frameborder="0"
+                    width="100%"
+                    height="560px"
+                    ></iframe>
+                    """
+
+                    # Save the HTML to a file
+                    visualization_path = os.path.join(
+                        self.uri, f"{name}_viewer.html"
+                    )
+                    with fileio.open(visualization_path, "w") as f:
+                        f.write(html)
+
+                    visualizations[visualization_path] = VisualizationType.HTML
+
+        return visualizations

zenml/integrations/huggingface/steps/accelerate_runner.py

@@ -111,15 +111,9 @@ def run_with_accelerate(
                     if isinstance(v, bool):
                         if v:
                             commands.append(f"--{k}")
-                    elif isinstance(v, str):
-                        commands += [f"--{k}", '"{v}"']
                     elif type(v) in (list, tuple, set):
                         for each in v:
-                            commands.append(f"--{k}")
-                            if isinstance(each, str):
-                                commands.append(f'"{each}"')
-                            else:
-                                commands.append(f"{each}")
+                            commands += [f"--{k}", f"{each}"]
                     else:
                         commands += [f"--{k}", f"{v}"]
 

zenml/integrations/integration.py

@@ -60,6 +60,7 @@ class Integration(metaclass=IntegrationMeta):
 
     REQUIREMENTS: List[str] = []
     APT_PACKAGES: List[str] = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL: List[str] = []
 
     @classmethod
     def check_installation(cls) -> bool:
@@ -143,6 +144,29 @@ class Integration(metaclass=IntegrationMeta):
         """
         return cls.REQUIREMENTS
 
+    @classmethod
+    def get_uninstall_requirements(
+        cls, target_os: Optional[str] = None
+    ) -> List[str]:
+        """Method to get the uninstall requirements for the integration.
+
+        Args:
+            target_os: The target operating system to get the requirements for.
+
+        Returns:
+            A list of requirements.
+        """
+        ret = []
+        for each in cls.get_requirements(target_os=target_os):
+            is_ignored = False
+            for ignored in cls.REQUIREMENTS_IGNORED_ON_UNINSTALL:
+                if each.startswith(ignored):
+                    is_ignored = True
+                    break
+            if not is_ignored:
+                ret.append(each)
+        return ret
+
     @classmethod
     def activate(cls) -> None:
         """Abstract method to activate the integration."""

zenml/integrations/kubeflow/__init__.py

@@ -31,6 +31,9 @@ class KubeflowIntegration(Integration):
 
     NAME = KUBEFLOW
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]  # Only 1.x version that supports pyyaml 6
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by GCP as well
+    ]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py

@@ -149,7 +149,7 @@ class KubeflowOrchestratorConfig(
 
     kubeflow_hostname: Optional[str] = None
     kubeflow_namespace: str = "kubeflow"
-    kubernetes_context: Optional[str]  # TODO: Potential setting
+    kubernetes_context: Optional[str] = None  # TODO: Potential setting
 
     @model_validator(mode="before")
     @classmethod

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -433,7 +433,6 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
             The dynamic container component.
         """
 
-        @dsl.container_component  # type: ignore[misc]
         def dynamic_container_component() -> dsl.ContainerSpec:
             """Dynamic container component.
 

zenml/integrations/kubernetes/__init__.py

@@ -31,7 +31,9 @@ class KubernetesIntegration(Integration):
 
     NAME = KUBERNETES
     REQUIREMENTS = ["kubernetes>=21.7,<26"]
-
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by many others
+    ]
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
         """Declare the stack component flavors for the Kubernetes integration.

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -225,8 +225,11 @@ def wait_pod(
             response = core_api.read_namespaced_pod_log(
                 name=pod_name,
                 namespace=namespace,
+                _preload_content=False,
             )
-            logs = response.splitlines()
+            raw_data = response.data
+            decoded_log = raw_data.decode("utf-8", errors="replace")
+            logs = decoded_log.splitlines()
             if len(logs) > logged_lines:
                 for line in logs[logged_lines:]:
                     logger.info(line)

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -38,7 +38,6 @@ from kubernetes import config as k8s_config
 
 from zenml.config.base_settings import BaseSettings
 from zenml.enums import StackComponentType
-from zenml.environment import Environment
 from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import (
     KubernetesOrchestratorConfig,
     KubernetesOrchestratorSettings,
@@ -335,19 +334,8 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 environment.
 
         Raises:
-            RuntimeError: If trying to run from a Jupyter notebook.
+            RuntimeError: If the Kubernetes orchestrator is not configured.
         """
-        # First check whether the code is running in a notebook.
-        if Environment.in_notebook():
-            raise RuntimeError(
-                "The Kubernetes orchestrator cannot run pipelines in a notebook "
-                "environment. The reason is that it is non-trivial to create "
-                "a Docker image of a notebook. Please consider refactoring "
-                "your notebook cells into separate scripts in a Python module "
-                "and run the code outside of a notebook when using this "
-                "orchestrator."
-            )
-
         for step_name, step in deployment.step_configurations.items():
             if self.requires_resources_in_orchestration_environment(step):
                 logger.warning(

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -139,24 +139,42 @@ def build_pod_manifest(
         ],
         security_context=security_context,
     )
+    image_pull_secrets = []
+    if pod_settings:
+        image_pull_secrets = [
+            k8s_client.V1LocalObjectReference(name=name)
+            for name in pod_settings.image_pull_secrets
+        ]
 
     pod_spec = k8s_client.V1PodSpec(
         containers=[container_spec],
         restart_policy="Never",
+        image_pull_secrets=image_pull_secrets,
     )
 
     if service_account_name is not None:
         pod_spec.service_account_name = service_account_name
 
+    labels = {}
+
     if pod_settings:
         add_pod_settings(pod_spec, pod_settings)
 
-    pod_metadata = k8s_client.V1ObjectMeta(
-        name=pod_name,
-        labels={
+        # Add pod_settings.labels to the labels
+        if pod_settings.labels:
+            labels.update(pod_settings.labels)
+
+    # Add run_name and pipeline_name to the labels
+    labels.update(
+        {
             "run": run_name,
             "pipeline": pipeline_name,
-        },
+        }
+    )
+
+    pod_metadata = k8s_client.V1ObjectMeta(
+        name=pod_name,
+        labels=labels,
     )
 
     if pod_settings and pod_settings.annotations:

zenml/integrations/kubernetes/pod_settings.py

@@ -33,6 +33,8 @@ class KubernetesPodSettings(BaseSettings):
         volumes: Volumes to mount in the pod.
         volume_mounts: Volume mounts to apply to the pod containers.
         host_ipc: Whether to enable host IPC for the pod.
+        image_pull_secrets: Image pull secrets to use for the pod.
+        labels: Labels to apply to the pod.
     """
 
     node_selectors: Dict[str, str] = {}
@@ -43,6 +45,8 @@ class KubernetesPodSettings(BaseSettings):
     volumes: List[Dict[str, Any]] = []
     volume_mounts: List[Dict[str, Any]] = []
     host_ipc: bool = False
+    image_pull_secrets: List[str] = []
+    labels: Dict[str, str] = {}
 
     @field_validator("volumes", mode="before")
     @classmethod

zenml/integrations/label_studio/annotators/label_studio_annotator.py

@@ -527,6 +527,7 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,
+                _,
             ) = artifact_store.get_credentials()
 
             if aws_access_key_id and aws_secret_access_key:

zenml/integrations/langchain/__init__.py

@@ -30,6 +30,7 @@ class LangchainIntegration(Integration):
         "pyyaml>=6.0.1",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyyaml","tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/lightgbm/__init__.py

@@ -22,6 +22,7 @@ class LightGBMIntegration(Integration):
 
     NAME = LIGHTGBM
     REQUIREMENTS = ["lightgbm>=1.0.0"]
+    APT_PACKAGES = ["libgomp1"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/mlflow/__init__.py

@@ -33,7 +33,7 @@ class MlflowIntegration(Integration):
     NAME = MLFLOW
 
     REQUIREMENTS = [
-        "mlflow>=2.1.1,<=2.14.1",
+        "mlflow>=2.1.1,<=2.14.2",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed
@@ -45,9 +45,11 @@ class MlflowIntegration(Integration):
         # This downgrades pydantic to v1 even though mlflow does not have
         # any issues with v2. This is why we have to pin it here so a downgrade
         # will not happen.
-        "pydantic>=2.7.0,<2.8.0"
+        "pydantic>=2.7.0,<2.8.0",
     ]
 
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["python-rapidjson", "pydantic"]
+
     @classmethod
     def activate(cls) -> None:
         """Activate the MLflow integration."""

zenml/integrations/mlflow/model_registries/mlflow_model_registry.py

@@ -371,7 +371,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
             self.get_model(name=name)
         except KeyError:
             logger.info(
-                f"No registered model with name {name} found. Creating a new"
+                f"No registered model with name {name} found. Creating a new "
                 "registered model."
             )
             self.register_model(
@@ -615,7 +615,11 @@ class MLFlowModelRegistry(BaseModelRegistry):
         for mlflow_model_version in mlflow_model_versions:
             # check if given MlFlow model version matches the given request
             # before casting it
-            if stage and not mlflow_model_version.current_stage == str(stage):
+            if (
+                stage
+                and not ModelVersionStage(mlflow_model_version.current_stage)
+                == stage
+            ):
                 continue
             if created_after and not (
                 mlflow_model_version.creation_timestamp

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -260,7 +260,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
             )
 
         if self.endpoint.prediction_url is not None:
-            if type(request) == pd.DataFrame:
+            if type(request) is pd.DataFrame:
                 response = requests.post(  # nosec
                     self.endpoint.prediction_url,
                     json={"instances": request.to_dict("records")},

zenml/integrations/neural_prophet/__init__.py

@@ -26,6 +26,7 @@ class NeuralProphetIntegration(Integration):
         "holidays>=0.4.1,<0.25.0",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/polars/__init__.py

@@ -25,6 +25,7 @@ class PolarsIntegration(Integration):
         "polars>=0.19.5",
         "pyarrow>=12.0.0"
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyarrow"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/prodigy/__init__.py

@@ -29,6 +29,7 @@ class ProdigyIntegration(Integration):
         "prodigy",
         "urllib3<2",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["urllib3"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/pycaret/__init__.py

@@ -28,6 +28,12 @@ class PyCaretIntegration(Integration):
         "catboost",
         "lightgbm",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "scikit-learn",
+        "xgboost",
+        "catboost",
+        "lightgbm",
+    ]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/registry.py

@@ -124,6 +124,43 @@ class IntegrationRegistry(object):
                 )
             ]
 
+    def select_uninstall_requirements(
+        self,
+        integration_name: Optional[str] = None,
+        target_os: Optional[str] = None,
+    ) -> List[str]:
+        """Select the uninstall requirements for a given integration or all integrations.
+
+        Args:
+            integration_name: Name of the integration to check.
+            target_os: Target OS for the requirements.
+
+        Returns:
+            List of requirements for the integration uninstall.
+
+        Raises:
+            KeyError: If the integration is not found.
+        """
+        if integration_name:
+            if integration_name in self.list_integration_names:
+                return self._integrations[
+                    integration_name
+                ].get_uninstall_requirements(target_os=target_os)
+            else:
+                raise KeyError(
+                    f"Version {integration_name} does not exist. "
+                    f"Currently the following integrations are implemented. "
+                    f"{self.list_integration_names}"
+                )
+        else:
+            return [
+                requirement
+                for name in self.list_integration_names
+                for requirement in self._integrations[
+                    name
+                ].get_uninstall_requirements(target_os=target_os)
+            ]
+
     def is_installed(self, integration_name: Optional[str] = None) -> bool:
         """Checks if all requirements for an integration are installed.
 

zenml/integrations/s3/artifact_stores/s3_artifact_store.py

@@ -122,7 +122,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
     def get_credentials(
         self,
-    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
         """Gets authentication credentials.
 
         If an authentication secret is configured, the secret values are
@@ -130,8 +130,8 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         attributes.
 
         Returns:
-            Tuple (key, secret, token) of credentials used to authenticate with
-            the S3 filesystem.
+            Tuple (key, secret, token, region) of credentials used to
+            authenticate with the S3 filesystem.
 
         Raises:
             RuntimeError: If the AWS connector behaves unexpectedly.
@@ -151,6 +151,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 credentials.access_key,
                 credentials.secret_key,
                 credentials.token,
+                client.meta.region_name,
             )
 
         secret = self.get_typed_authentication_secret(
@@ -161,9 +162,10 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 secret.aws_access_key_id,
                 secret.aws_secret_access_key,
                 secret.aws_session_token,
+                None,
             )
         else:
-            return self.config.key, self.config.secret, self.config.token
+            return self.config.key, self.config.secret, self.config.token, None
 
     @property
     def filesystem(self) -> ZenMLS3Filesystem:
@@ -176,13 +178,22 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         if self._filesystem and not self.connector_has_expired():
             return self._filesystem
 
-        key, secret, token = self.get_credentials()
+        key, secret, token, region = self.get_credentials()
+
+        # Use the region from the connector if available, otherwise some
+        # remote workloads (e.g. Sagemaker) might not work correctly because
+        # they look for the bucket in the wrong region
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+        if self.config.client_kwargs:
+            client_kwargs.update(self.config.client_kwargs)
 
         self._filesystem = ZenMLS3Filesystem(
             key=key,
             secret=secret,
             token=token,
-            client_kwargs=self.config.client_kwargs,
+            client_kwargs=client_kwargs,
             config_kwargs=self.config.config_kwargs,
             s3_additional_kwargs=self.config.s3_additional_kwargs,
         )

zenml/integrations/seldon/__init__.py

@@ -33,6 +33,7 @@ class SeldonIntegration(Integration):
     REQUIREMENTS = [
         "kubernetes==18.20.0",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/seldon/model_deployers/seldon_model_deployer.py

@@ -290,6 +290,7 @@ class SeldonModelDeployer(BaseModelDeployer):
             assert isinstance(artifact_store, S3ArtifactStore)
 
             (
+                region_name,
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -43,7 +43,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
             `{'tpu_vm': True, 'runtime_version': 'tpu-vm-base'}` for TPUs.
         use_spot: whether to use spot instances. If None, defaults to
             False.
-        spot_recovery: the spot recovery strategy to use for the managed
+        job_recovery: the spot recovery strategy to use for the managed
             spot to recover the cluster from preemption. Refer to
             `recovery_strategy module <https://github.com/skypilot-org/skypilot/blob/master/sky/spot/recovery_strategy.py>`__ # pylint: disable=line-too-long
             for more details.
@@ -103,7 +103,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
     )
     accelerator_args: Optional[Dict[str, str]] = None
     use_spot: Optional[bool] = None
-    spot_recovery: Optional[str] = None
+    job_recovery: Optional[str] = None
     region: Optional[str] = None
     zone: Optional[str] = None
     image_id: Union[Dict[str, str], str, None] = Field(

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -303,7 +303,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     accelerators=settings.accelerators,
                     accelerator_args=settings.accelerator_args,
                     use_spot=settings.use_spot,
-                    spot_recovery=settings.spot_recovery,
+                    job_recovery=settings.job_recovery,
                     region=settings.region,
                     zone=settings.zone,
                     image_id=settings.image_id,

zenml/integrations/skypilot/orchestrators/skypilot_orchestrator_entrypoint.py

@@ -136,7 +136,7 @@ def main() -> None:
             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.spot_recovery,
+            settings.job_recovery,
             settings.region,
             settings.zone,
             accelerators_hashable,
@@ -213,7 +213,7 @@ def main() -> None:
                 accelerators=settings.accelerators,
                 accelerator_args=settings.accelerator_args,
                 use_spot=settings.use_spot,
-                spot_recovery=settings.spot_recovery,
+                job_recovery=settings.job_recovery,
                 region=settings.region,
                 zone=settings.zone,
                 image_id=settings.image_id,

zenml/integrations/skypilot_aws/__init__.py

@@ -31,7 +31,8 @@ class SkypilotAWSIntegration(Integration):
     """Definition of Skypilot AWS Integration for ZenML."""
 
     NAME = SKYPILOT_AWS
-    REQUIREMENTS = ["skypilot[aws]<=0.5.0"]
+    # all 0.6.x versions of skypilot[aws] are compatible
+    REQUIREMENTS = ["skypilot[aws]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_azure/__init__.py

@@ -19,13 +19,11 @@ orchestrator for a remote orchestration of ZenML pipelines on VMs.
 from typing import List, Type
 
 from zenml.integrations.constants import (
-
     SKYPILOT_AZURE,
 )
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
-
 SKYPILOT_AZURE_ORCHESTRATOR_FLAVOR = "vm_azure"
 
 
@@ -33,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]<=0.5.0"]
+    REQUIREMENTS = ["skypilot-nightly[azure]==1.0.0.dev20240716"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod