zenml-nightly 0.61.0.dev20240712__py3-none-any.whl → 0.62.0.dev20240727__py3-none-any.whl

zenml/constants.py

@@ -173,6 +173,7 @@ ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES"
 )
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
+ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
@@ -375,6 +376,7 @@ SERVICE_CONNECTOR_CLIENT = "/client"
 SERVICE_CONNECTOR_RESOURCES = "/resources"
 SERVICE_CONNECTOR_TYPES = "/service_connector_types"
 SERVICE_CONNECTOR_VERIFY = "/verify"
+SERVICE_CONNECTOR_FULL_STACK = "/full_stack_resources"
 MODELS = "/models"
 MODEL_VERSIONS = "/model_versions"
 MODEL_VERSION_ARTIFACTS = "/model_version_artifacts"
@@ -392,6 +394,7 @@ STEPS = "/steps"
 TAGS = "/tags"
 TRIGGERS = "/triggers"
 TRIGGER_EXECUTIONS = "/trigger_executions"
+ONBOARDING_STATE = "/onboarding_state"
 USERS = "/users"
 URL = "/url"
 VERSION_1 = "/v1"

zenml/enums.py

@@ -386,8 +386,24 @@ class PluginSubType(StrEnum):
     PIPELINE_RUN = "pipeline_run"
 
 
+class OnboardingStep(StrEnum):
+    """All onboarding steps."""
+
+    DEVICE_VERIFIED = "device_verified"
+    PIPELINE_RUN = "pipeline_run"
+    STARTER_SETUP_COMPLETED = "starter_setup_completed"
+    STACK_WITH_REMOTE_ORCHESTRATOR_CREATED = (
+        "stack_with_remote_orchestrator_created"
+    )
+    PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR = (
+        "pipeline_run_with_remote_orchestrator"
+    )
+    PRODUCTION_SETUP_COMPLETED = "production_setup_completed"
+
+
 class StackDeploymentProvider(StrEnum):
     """All possible stack deployment providers."""
 
     AWS = "aws"
     GCP = "gcp"
+    AZURE = "azure"

zenml/integrations/__init__.py

@@ -23,6 +23,7 @@ from zenml.integrations.aws import AWSIntegration  # noqa
 from zenml.integrations.azure import AzureIntegration  # noqa
 from zenml.integrations.bentoml import BentoMLIntegration  # noqa
 from zenml.integrations.bitbucket import BitbucketIntegration  # noqa
+from zenml.integrations.databricks import DatabricksIntegration  # noqa
 from zenml.integrations.comet import CometIntegration  # noqa
 from zenml.integrations.deepchecks import DeepchecksIntegration  # noqa
 from zenml.integrations.discord import DiscordIntegration  # noqa

zenml/integrations/azure/__init__.py

@@ -42,8 +42,8 @@ class AzureIntegration(Integration):
         "adlfs>=2021.10.0",
         "azure-keyvault-keys",
         "azure-keyvault-secrets",
-        "azure-identity==1.10.0",
-        "azureml-core==1.54.0.post1",
+        "azure-identity",
+        "azureml-core==1.56.0",
         "azure-mgmt-containerservice>=20.0.0",
         "azure-storage-blob==12.17.0",  # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056
         "kubernetes",

zenml/integrations/constants.py

@@ -22,6 +22,7 @@ BENTOML = "bentoml"
 BITBUCKET = "bitbucket"
 COMET = "comet"
 DASH = "dash"
+DATABRICKS = "databricks"
 DEEPCHECKS = "deepchecks"
 DISCORD = "discord"
 EVIDENTLY = "evidently"

zenml/integrations/databricks/__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 Databricks integration for ZenML."""
+
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    DATABRICKS,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+DATABRICKS_ORCHESTRATOR_FLAVOR = "databricks"
+DATABRICKS_MODEL_DEPLOYER_FLAVOR = "databricks"
+DATABRICKS_SERVICE_ARTIFACT = "databricks_deployment_service"
+
+
+class DatabricksIntegration(Integration):
+    """Definition of Databricks Integration for ZenML."""
+
+    NAME = DATABRICKS
+    REQUIREMENTS = ["databricks-sdk==0.28.0"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Databricks integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.databricks.flavors import (
+            DatabricksOrchestratorFlavor,
+            DatabricksModelDeployerFlavor,
+        )
+
+        return [
+            DatabricksOrchestratorFlavor,
+            DatabricksModelDeployerFlavor,
+        ]
+
+DatabricksIntegration.check_installation()

zenml/integrations/databricks/flavors/__init__.py

@@ -0,0 +1,30 @@
+#  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.
+"""Databricks integration flavors."""
+
+from zenml.integrations.databricks.flavors.databricks_orchestrator_flavor import (
+    DatabricksOrchestratorConfig,
+    DatabricksOrchestratorFlavor,
+)
+from zenml.integrations.databricks.flavors.databricks_model_deployer_flavor import (
+    DatabricksModelDeployerConfig,
+    DatabricksModelDeployerFlavor,
+)
+
+__all__ = [
+    "DatabricksOrchestratorFlavor",
+    "DatabricksOrchestratorConfig",
+    "DatabricksModelDeployerFlavor",
+    "DatabricksModelDeployerConfig",
+]

zenml/integrations/databricks/flavors/databricks_model_deployer_flavor.py

@@ -0,0 +1,118 @@
+#  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.
+"""Databricks model deployer flavor."""
+
+from typing import TYPE_CHECKING, Dict, Optional, Type
+
+from pydantic import BaseModel
+
+from zenml.integrations.databricks import DATABRICKS_MODEL_DEPLOYER_FLAVOR
+from zenml.model_deployers.base_model_deployer import (
+    BaseModelDeployerConfig,
+    BaseModelDeployerFlavor,
+)
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (
+        DatabricksModelDeployer,
+    )
+
+
+class DatabricksBaseConfig(BaseModel):
+    """Databricks Inference Endpoint configuration."""
+
+    workload_size: str
+    scale_to_zero_enabled: bool = False
+    env_vars: Optional[Dict[str, str]] = None
+    workload_type: Optional[str] = None
+    endpoint_secret_name: Optional[str] = None
+
+
+class DatabricksModelDeployerConfig(BaseModelDeployerConfig):
+    """Configuration for the Databricks model deployer.
+
+    Attributes:
+        host: Databricks host.
+        secret_name: Secret name to use for authentication.
+        client_id: Databricks client id.
+        client_secret: Databricks client secret.
+    """
+
+    host: str
+    secret_name: Optional[str] = None
+    client_id: Optional[str] = SecretField(default=None)
+    client_secret: Optional[str] = SecretField(default=None)
+
+
+class DatabricksModelDeployerFlavor(BaseModelDeployerFlavor):
+    """Databricks Endpoint model deployer flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return DATABRICKS_MODEL_DEPLOYER_FLAVOR
+
+    @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/model_deployer/databricks.png"
+
+    @property
+    def config_class(self) -> Type[DatabricksModelDeployerConfig]:
+        """Returns `DatabricksModelDeployerConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return DatabricksModelDeployerConfig
+
+    @property
+    def implementation_class(self) -> Type["DatabricksModelDeployer"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (
+            DatabricksModelDeployer,
+        )
+
+        return DatabricksModelDeployer

zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py

@@ -0,0 +1,147 @@
+#  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.
+"""Databricks orchestrator base config and settings."""
+
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.databricks import DATABRICKS_ORCHESTRATOR_FLAVOR
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestratorConfig
+from zenml.orchestrators.base_orchestrator import BaseOrchestratorFlavor
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.databricks.orchestrators import (
+        DatabricksOrchestrator,
+    )
+
+
+logger = get_logger(__name__)
+
+
+class DatabricksOrchestratorSettings(BaseSettings):
+    """Databricks orchestrator base settings.
+
+    Attributes:
+        spark_version: Spark version.
+        num_workers: Number of workers.
+        node_type_id: Node type id.
+        policy_id: Policy id.
+        autotermination_minutes: Autotermination minutes.
+        autoscale: Autoscale.
+        single_user_name: Single user name.
+        spark_conf: Spark configuration.
+        spark_env_vars: Spark environment variables.
+        schedule_timezone: Schedule timezone.
+    """
+
+    # Resources
+    spark_version: Optional[str] = None
+    num_workers: Optional[int] = None
+    node_type_id: Optional[str] = None
+    policy_id: Optional[str] = None
+    autotermination_minutes: Optional[int] = None
+    autoscale: Tuple[int, int] = (2, 3)
+    single_user_name: Optional[str] = None
+    spark_conf: Optional[Dict[str, str]] = None
+    spark_env_vars: Optional[Dict[str, str]] = None
+    schedule_timezone: Optional[str] = None
+
+
+class DatabricksOrchestratorConfig(
+    BaseOrchestratorConfig, DatabricksOrchestratorSettings
+):
+    """Databricks orchestrator base config.
+
+    Attributes:
+        host: Databricks host.
+        client_id: Databricks client id.
+        client_secret: Databricks client secret.
+    """
+
+    host: str
+    client_id: str = SecretField(default=None)
+    client_secret: str = SecretField(default=None)
+
+    @property
+    def is_local(self) -> bool:
+        """Checks if this stack component is running locally.
+
+        Returns:
+            True if this config is for a local component, False otherwise.
+        """
+        return False
+
+
+class DatabricksOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Databricks orchestrator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return DATABRICKS_ORCHESTRATOR_FLAVOR
+
+    @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/databricks.png"
+
+    @property
+    def config_class(self) -> Type[DatabricksOrchestratorConfig]:
+        """Returns `KubeflowOrchestratorConfig` config class.
+
+        Returns:
+                The config class.
+        """
+        return DatabricksOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["DatabricksOrchestrator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.databricks.orchestrators import (
+            DatabricksOrchestrator,
+        )
+
+        return DatabricksOrchestrator

zenml/integrations/databricks/model_deployers/__init__.py

@@ -0,0 +1,20 @@
+#  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:
+#
+#       http://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 Databricks model deployers."""
+
+from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (  # noqa
+    DatabricksModelDeployer,
+)
+
+__all__ = ["DatabricksModelDeployer"]