zenml-nightly 0.63.0.dev20240801__py3-none-any.whl → 0.64.0.dev20240811__py3-none-any.whl

zenml/config/source.py

@@ -42,6 +42,7 @@ class SourceType(Enum):
     INTERNAL = "internal"
     DISTRIBUTION_PACKAGE = "distribution_package"
     CODE_REPOSITORY = "code_repository"
+    NOTEBOOK = "notebook"
     UNKNOWN = "unknown"
 
 
@@ -229,6 +230,63 @@ class CodeRepositorySource(Source):
         return value
 
 
+class NotebookSource(Source):
+    """Source representing an object defined in a notebook.
+
+    Attributes:
+        code_path: Path where the notebook cell code for this source is
+            uploaded.
+        replacement_module: Name of the module from which this source should
+            be loaded in case the code is not running in a notebook.
+    """
+
+    code_path: Optional[str] = None
+    replacement_module: Optional[str] = None
+    type: SourceType = SourceType.NOTEBOOK
+
+    # Private attribute that is used to store the code but should not be
+    # serialized
+    _cell_code: Optional[str] = None
+
+    @field_validator("type")
+    @classmethod
+    def _validate_type(cls, value: SourceType) -> SourceType:
+        """Validate the source type.
+
+        Args:
+            value: The source type.
+
+        Raises:
+            ValueError: If the source type is not `NOTEBOOK`.
+
+        Returns:
+            The source type.
+        """
+        if value != SourceType.NOTEBOOK:
+            raise ValueError("Invalid source type.")
+
+        return value
+
+    @field_validator("module")
+    @classmethod
+    def _validate_module(cls, value: str) -> str:
+        """Validate the module.
+
+        Args:
+            value: The module.
+
+        Raises:
+            ValueError: If the module is not `__main__`.
+
+        Returns:
+            The module.
+        """
+        if value != "__main__":
+            raise ValueError("Invalid module for notebook source.")
+
+        return value
+
+
 def convert_source(source: Any) -> Any:
     """Converts an old source string to a source object.
 

zenml/constants.py

@@ -165,7 +165,6 @@ ENV_ZENML_SKIP_IMAGE_BUILDER_DEFAULT = "ZENML_SKIP_IMAGE_BUILDER_DEFAULT"
 ENV_ZENML_REQUIRES_CODE_DOWNLOAD = "ZENML_REQUIRES_CODE_DOWNLOAD"
 ENV_ZENML_SERVER = "ZENML_SERVER"
 ENV_ZENML_LOCAL_SERVER = "ZENML_LOCAL_SERVER"
-ENV_ZENML_HUB_URL = "ZENML_HUB_URL"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
@@ -279,7 +278,6 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT = "nosniff"
 _csp_script_src_urls = ["https://widgets-v3.featureos.app"]
 _csp_connect_src_urls = [
     "https://sdkdocs.zenml.io",
-    "https://hubapi.zenml.io",
     "https://analytics.zenml.io",
 ]
 _csp_img_src_urls = [
@@ -315,6 +313,7 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS = (
 )
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
 DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD = False
+DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS = 30
 
 # Configurations to decide which resources report their usage and check for
 # entitlement in the case of a cloud deployment. Expected Format is this:
@@ -494,3 +493,11 @@ BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
 
 
 STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
+
+# ZenML Pro
+ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT = (
+    "\nHINT: Since you are trying to communicate with the ZenML Pro Tenant, "
+    "please make sure that your tenant is in RUNNING state on your "
+    "Organization page. If the tenant is PAUSED you can `Resume` it via UI "
+    "and try again."
+)

zenml/entrypoints/base_entrypoint_configuration.py

@@ -27,10 +27,15 @@ from zenml.constants import (
     handle_bool_env_var,
 )
 from zenml.logger import get_logger
-from zenml.utils import code_repository_utils, source_utils, uuid_utils
+from zenml.utils import (
+    code_repository_utils,
+    code_utils,
+    source_utils,
+    uuid_utils,
+)
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import CodeReferenceResponse, PipelineDeploymentResponse
 
 logger = get_logger(__name__)
 DEFAULT_ENTRYPOINT_COMMAND = [
@@ -198,7 +203,7 @@ class BaseEntrypointConfiguration(ABC):
 
         Raises:
             RuntimeError: If the current environment requires code download
-                but the deployment does not have an associated code reference.
+                but the deployment does not have a reference to any code.
         """
         requires_code_download = handle_bool_env_var(
             ENV_ZENML_REQUIRES_CODE_DOWNLOAD
@@ -207,17 +212,33 @@ class BaseEntrypointConfiguration(ABC):
         if not requires_code_download:
             return
 
-        code_reference = deployment.code_reference
-        if not code_reference:
+        if code_reference := deployment.code_reference:
+            self.download_code_from_code_repository(
+                code_reference=code_reference
+            )
+        elif code_path := deployment.code_path:
+            self.download_code_from_artifact_store(code_path=code_path)
+        else:
             raise RuntimeError(
-                "Code download required but no code reference provided."
+                "Code download required but no code reference or path provided."
             )
 
+        logger.info("Code download finished.")
+
+    def download_code_from_code_repository(
+        self, code_reference: "CodeReferenceResponse"
+    ) -> None:
+        """Download code from a code repository.
+
+        Args:
+            code_reference: The reference to the code.
+        """
         logger.info(
             "Downloading code from code repository `%s` (commit `%s`).",
             code_reference.code_repository.name,
             code_reference.commit,
         )
+
         model = Client().get_code_repository(code_reference.code_repository.id)
         repo = BaseCodeRepository.from_model(model)
         code_repo_root = os.path.abspath("code")
@@ -234,10 +255,34 @@ class BaseEntrypointConfiguration(ABC):
         code_repository_utils.set_custom_local_repository(
             root=code_repo_root, commit=code_reference.commit, repo=repo
         )
-        # Add downloaded file directory to python path
+
         sys.path.insert(0, download_dir)
+        os.chdir(download_dir)
 
-        logger.info("Code download finished.")
+    def download_code_from_artifact_store(self, code_path: str) -> None:
+        """Download code from the artifact store.
+
+        Args:
+            code_path: Path where the code is stored.
+        """
+        logger.info(
+            "Downloading code from artifact store path `%s`.", code_path
+        )
+
+        # Do not remove this line, we need to instantiate the artifact store to
+        # register the filesystem needed for the file download
+        _ = Client().active_stack.artifact_store
+
+        extract_dir = os.path.abspath("code")
+        os.makedirs(extract_dir)
+
+        code_utils.download_and_extract_code(
+            code_path=code_path, extract_dir=extract_dir
+        )
+
+        source_utils.set_custom_source_root(extract_dir)
+        sys.path.insert(0, extract_dir)
+        os.chdir(extract_dir)
 
     @abstractmethod
     def run(self) -> None:

zenml/enums.py

@@ -171,7 +171,6 @@ class CliCategories(StrEnum):
 
     STACK_COMPONENTS = "Stack Components"
     MODEL_DEPLOYMENT = "Model Deployment"
-    HUB = "ZenML Hub"
     INTEGRATIONS = "Integrations"
     MANAGEMENT_TOOLS = "Management Tools"
     MODEL_CONTROL_PLANE = "Model Control Plane"
@@ -303,6 +302,7 @@ class EnvironmentType(StrEnum):
     NOTEBOOK = "notebook"
     PAPERSPACE = "paperspace"
     WSL = "wsl"
+    LIGHTNING_AI_STUDIO = "lightning_ai_studio"
 
 
 class ModelStages(StrEnum):

zenml/environment.py

@@ -15,7 +15,6 @@
 
 import os
 import platform
-from importlib.util import find_spec
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, cast
 
@@ -70,6 +69,8 @@ def get_environment() -> str:
         return EnvironmentType.BITBUCKET_CI
     elif Environment.in_ci():
         return EnvironmentType.GENERIC_CI
+    elif Environment.in_lightning_ai_studio():
+        return EnvironmentType.LIGHTNING_AI_STUDIO
     elif Environment.in_docker():
         return EnvironmentType.DOCKER
     elif Environment.in_container():
@@ -257,15 +258,17 @@ class Environment(metaclass=SingletonMetaClass):
         if Environment.in_google_colab():
             return True
 
-        if find_spec("IPython") is not None:
-            from IPython import get_ipython
+        try:
+            ipython = get_ipython()  # type: ignore[name-defined]
+        except NameError:
+            return False
 
-            if get_ipython().__class__.__name__ in [
-                "TerminalInteractiveShell",
-                "ZMQInteractiveShell",
-                "DatabricksShell",
-            ]:
-                return True
+        if ipython.__class__.__name__ in [
+            "TerminalInteractiveShell",
+            "ZMQInteractiveShell",
+            "DatabricksShell",
+        ]:
+            return True
         return False
 
     @staticmethod
@@ -339,6 +342,19 @@ class Environment(metaclass=SingletonMetaClass):
         """
         return "microsoft-standard" in platform.uname().release
 
+    @staticmethod
+    def in_lightning_ai_studio() -> bool:
+        """If the current Python process is running in Lightning.ai studios.
+
+        Returns:
+            `True` if the current Python process is running in Lightning.ai studios,
+            `False` otherwise.
+        """
+        return (
+            "LIGHTNING_CLOUD_URL" in os.environ
+            and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
+        )
+
     def register_component(
         self, component: "BaseEnvironmentComponent"
     ) -> "BaseEnvironmentComponent":

zenml/image_builders/base_image_builder.py

@@ -119,7 +119,7 @@ class BaseImageBuilder(StackComponent, ABC):
 
         hash_ = hashlib.sha1()  # nosec
         with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as f:
-            build_context.write_archive(f, gzip=True)
+            build_context.write_archive(f, use_gzip=True)
 
             while True:
                 data = f.read(64 * 1024)

zenml/image_builders/build_context.py

@@ -14,18 +14,18 @@
 """Image build context."""
 
 import os
-from pathlib import Path
-from typing import IO, Dict, List, Optional, Set, Tuple, cast
+from typing import IO, Dict, List, Optional, Set, cast
 
 from zenml.constants import REPOSITORY_DIRECTORY_NAME
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import io_utils, string_utils
+from zenml.utils.archivable import Archivable
 
 logger = get_logger(__name__)
 
 
-class BuildContext:
+class BuildContext(Archivable):
     """Image build context.
 
     This class is responsible for creating an archive of the files needed to
@@ -45,9 +45,9 @@ class BuildContext:
                 given, a file called `.dockerignore` in the build context root
                 directory will be used instead if it exists.
         """
+        super().__init__()
         self._root = root
         self._dockerignore_file = dockerignore_file
-        self._extra_files: Dict[str, str] = {}
 
     @property
     def dockerignore_file(self) -> Optional[str]:
@@ -68,70 +68,26 @@ class BuildContext:
 
         return None
 
-    def add_file(self, source: str, destination: str) -> None:
-        """Adds a file to the build context.
-
-        Args:
-            source: The source of the file to add. This can either be a path
-                or the file content.
-            destination: The path inside the build context where the file
-                should be added.
-        """
-        if fileio.exists(source):
-            with fileio.open(source) as f:
-                self._extra_files[destination] = f.read()
-        else:
-            self._extra_files[destination] = source
-
-    def add_directory(self, source: str, destination: str) -> None:
-        """Adds a directory to the build context.
-
-        Args:
-            source: Path to the directory.
-            destination: The path inside the build context where the directory
-                should be added.
-
-        Raises:
-            ValueError: If `source` does not point to a directory.
-        """
-        if not fileio.isdir(source):
-            raise ValueError(
-                f"Can't add directory {source} to the build context as it "
-                "does not exist or is not a directory."
-            )
-
-        for dir, _, files in fileio.walk(source):
-            dir_path = Path(fileio.convert_to_str(dir))
-            for file_name in files:
-                file_name = fileio.convert_to_str(file_name)
-                file_source = dir_path / file_name
-                file_destination = (
-                    Path(destination)
-                    / dir_path.relative_to(source)
-                    / file_name
-                )
-
-                with file_source.open("r") as f:
-                    self._extra_files[file_destination.as_posix()] = f.read()
-
-    def write_archive(self, output_file: IO[bytes], gzip: bool = True) -> None:
+    def write_archive(
+        self, output_file: IO[bytes], use_gzip: bool = True
+    ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            gzip: Whether to use `gzip` to compress the file.
+            use_gzip: Whether to use `gzip` to compress the file.
         """
         from docker.utils import build as docker_build_utils
 
-        files = self._get_files()
-        extra_files = self._get_extra_files()
+        files = self.get_files()
+        extra_files = self.get_extra_files()
 
         context_archive = docker_build_utils.create_archive(
             fileobj=output_file,
             root=self._root,
-            files=sorted(files),
-            gzip=gzip,
-            extra_files=extra_files,
+            files=sorted(files.keys()),
+            gzip=use_gzip,
+            extra_files=list(extra_files.items()),
         )
 
         build_context_size = os.path.getsize(context_archive.name)
@@ -151,33 +107,30 @@ class BuildContext:
                 os.path.join(self._root, ".dockerignore"),
             )
 
-    def _get_files(self) -> Set[str]:
-        """Gets all non-ignored files in the build context root directory.
+    def get_files(self) -> Dict[str, str]:
+        """Gets all regular files that should be included in the archive.
 
         Returns:
-            All build context files.
+            A dict {path_in_archive: path_on_filesystem} for all regular files
+            in the archive.
         """
         if self._root:
-            exclude_patterns = self._get_exclude_patterns()
             from docker.utils import build as docker_build_utils
 
-            return cast(
+            exclude_patterns = self._get_exclude_patterns()
+
+            archive_paths = cast(
                 Set[str],
                 docker_build_utils.exclude_paths(
                     self._root, patterns=exclude_patterns
                 ),
             )
+            return {
+                archive_path: os.path.join(self._root, archive_path)
+                for archive_path in archive_paths
+            }
         else:
-            return set()
-
-    def _get_extra_files(self) -> List[Tuple[str, str]]:
-        """Gets all extra files of the build context.
-
-        Returns:
-            A tuple (path, file_content) for all extra files in the build
-            context.
-        """
-        return list(self._extra_files.items())
+            return {}
 
     def _get_exclude_patterns(self) -> List[str]:
         """Gets all exclude patterns from the dockerignore file.

zenml/integrations/azure/__init__.py

@@ -27,6 +27,7 @@ from zenml.stack import Flavor
 
 AZURE_ARTIFACT_STORE_FLAVOR = "azure"
 AZUREML_STEP_OPERATOR_FLAVOR = "azureml"
+AZUREML_ORCHESTRATOR_FLAVOR = "azureml"
 
 # Service connector constants
 AZURE_CONNECTOR_TYPE = "azure"
@@ -47,6 +48,7 @@ class AzureIntegration(Integration):
         "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",
+        "azure-ai-ml==1.18.0"
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
@@ -65,11 +67,13 @@ class AzureIntegration(Integration):
         from zenml.integrations.azure.flavors import (
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         )
 
         return [
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         ]
 
 

zenml/integrations/azure/flavors/__init__.py

@@ -16,10 +16,17 @@
 from zenml.integrations.azure.flavors.azure_artifact_store_flavor import (
     AzureArtifactStoreConfig,
     AzureArtifactStoreFlavor,
+
 )
 from zenml.integrations.azure.flavors.azureml_step_operator_flavor import (
     AzureMLStepOperatorConfig,
     AzureMLStepOperatorFlavor,
+    AzureMLStepOperatorSettings,
+)
+from zenml.integrations.azure.flavors.azureml_orchestrator_flavor import (
+    AzureMLOrchestratorConfig,
+    AzureMLOrchestratorFlavor,
+    AzureMLOrchestratorSettings,
 )
 
 __all__ = [
@@ -27,4 +34,8 @@ __all__ = [
     "AzureArtifactStoreConfig",
     "AzureMLStepOperatorFlavor",
     "AzureMLStepOperatorConfig",
+    "AzureMLStepOperatorSettings",
+    "AzureMLOrchestratorFlavor",
+    "AzureMLOrchestratorConfig",
+    "AzureMLOrchestratorSettings",
 ]