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

zenml/utils/pipeline_docker_image_builder.py

@@ -66,6 +66,7 @@ DEFAULT_DOCKER_PARENT_IMAGE = (
     f"zenmldocker/zenml:{zenml.__version__}-"
     f"py{sys.version_info.major}.{sys.version_info.minor}"
 )
+DEFAULT_ZENML_DOCKER_REPOSITORY = "zenml"
 
 PIP_DEFAULT_ARGS = {
     "no-cache-dir": None,
@@ -198,10 +199,15 @@ class PipelineDockerImageBuilder:
                 # We will build an additional image on top of this one later
                 # to include user files and/or install requirements. The image
                 # we build now will be used as the parent for the next build.
-                user_image_name = (
-                    f"{docker_settings.target_repository}:"
-                    f"{tag}-intermediate-build"
-                )
+                repository = docker_settings.target_repository
+                if not repository:
+                    if container_registry:
+                        repository = (
+                            container_registry.config.default_repository
+                        )
+
+                repository = repository or DEFAULT_ZENML_DOCKER_REPOSITORY
+                user_image_name = f"{repository}:" f"{tag}-intermediate-build"
                 if push and container_registry:
                     user_image_name = (
                         f"{container_registry.config.uri}/{user_image_name}"
@@ -291,6 +297,14 @@ class PipelineDockerImageBuilder:
             if docker_settings.install_stack_requirements:
                 apt_packages += stack.apt_packages
 
+            # include apt packages from all required integrations
+            for integration in docker_settings.required_integrations:
+                # get the integration
+                integration_cls = integration_registry.integrations[
+                    integration
+                ]
+                apt_packages += integration_cls.APT_PACKAGES
+
             if apt_packages:
                 logger.info(
                     "Including apt packages: %s",
@@ -366,7 +380,14 @@ class PipelineDockerImageBuilder:
         Returns:
             The docker image name.
         """
-        target_image_name = f"{docker_settings.target_repository}:{tag}"
+        repository = docker_settings.target_repository
+        if not repository:
+            if container_registry:
+                repository = container_registry.config.default_repository
+
+        repository = repository or DEFAULT_ZENML_DOCKER_REPOSITORY
+
+        target_image_name = f"{repository}:{tag}"
         if container_registry:
             target_image_name = (
                 f"{container_registry.config.uri}/{target_image_name}"
@@ -460,69 +481,6 @@ class PipelineDockerImageBuilder:
                     "- Including python packages from local environment"
                 )
 
-        # Generate/Read requirements file for user-defined requirements
-        if isinstance(docker_settings.requirements, str):
-            path = os.path.abspath(docker_settings.requirements)
-            try:
-                user_requirements = io_utils.read_file_contents_as_string(path)
-            except FileNotFoundError as e:
-                raise FileNotFoundError(
-                    f"Requirements file {path} does not exist."
-                ) from e
-            if log:
-                logger.info(
-                    "- Including user-defined requirements from file `%s`",
-                    path,
-                )
-        elif isinstance(docker_settings.requirements, List):
-            user_requirements = "\n".join(docker_settings.requirements)
-            if log:
-                logger.info(
-                    "- Including user-defined requirements: %s",
-                    ", ".join(f"`{r}`" for r in docker_settings.requirements),
-                )
-        else:
-            user_requirements = None
-
-        if user_requirements:
-            requirements_files.append(
-                (".zenml_user_requirements", user_requirements, [])
-            )
-
-        # Generate requirements file for all required integrations
-        integration_requirements = set(
-            itertools.chain.from_iterable(
-                integration_registry.select_integration_requirements(
-                    integration_name=integration,
-                    target_os=OperatingSystemType.LINUX,
-                )
-                for integration in docker_settings.required_integrations
-            )
-        )
-
-        if docker_settings.install_stack_requirements:
-            integration_requirements.update(stack.requirements())
-            if code_repository:
-                integration_requirements.update(code_repository.requirements)
-
-        if integration_requirements:
-            integration_requirements_list = sorted(integration_requirements)
-            integration_requirements_file = "\n".join(
-                integration_requirements_list
-            )
-            requirements_files.append(
-                (
-                    ".zenml_integration_requirements",
-                    integration_requirements_file,
-                    [],
-                )
-            )
-            if log:
-                logger.info(
-                    "- Including integration requirements: %s",
-                    ", ".join(f"`{r}`" for r in integration_requirements_list),
-                )
-
         # Generate requirements files for all ZenML Hub plugins
         if docker_settings.required_hub_plugins:
             (
@@ -560,6 +518,85 @@ class PipelineDockerImageBuilder:
                         ", ".join(f"`{r}`" for r in hub_pypi_requirements),
                     )
 
+        if docker_settings.install_stack_requirements:
+            stack_requirements = stack.requirements()
+            if code_repository:
+                stack_requirements.update(code_repository.requirements)
+
+            if stack_requirements:
+                stack_requirements_list = sorted(stack_requirements)
+                stack_requirements_file = "\n".join(stack_requirements_list)
+                requirements_files.append(
+                    (
+                        ".zenml_stack_integration_requirements",
+                        stack_requirements_file,
+                        [],
+                    )
+                )
+                if log:
+                    logger.info(
+                        "- Including stack requirements: %s",
+                        ", ".join(f"`{r}`" for r in stack_requirements_list),
+                    )
+
+        # Generate requirements file for all required integrations
+        integration_requirements = set(
+            itertools.chain.from_iterable(
+                integration_registry.select_integration_requirements(
+                    integration_name=integration,
+                    target_os=OperatingSystemType.LINUX,
+                )
+                for integration in docker_settings.required_integrations
+            )
+        )
+
+        if integration_requirements:
+            integration_requirements_list = sorted(integration_requirements)
+            integration_requirements_file = "\n".join(
+                integration_requirements_list
+            )
+            requirements_files.append(
+                (
+                    ".zenml_integration_requirements",
+                    integration_requirements_file,
+                    [],
+                )
+            )
+            if log:
+                logger.info(
+                    "- Including integration requirements: %s",
+                    ", ".join(f"`{r}`" for r in integration_requirements_list),
+                )
+
+        # Generate/Read requirements file for user-defined requirements
+        if isinstance(docker_settings.requirements, str):
+            path = os.path.abspath(docker_settings.requirements)
+            try:
+                user_requirements = io_utils.read_file_contents_as_string(path)
+            except FileNotFoundError as e:
+                raise FileNotFoundError(
+                    f"Requirements file {path} does not exist."
+                ) from e
+            if log:
+                logger.info(
+                    "- Including user-defined requirements from file `%s`",
+                    path,
+                )
+        elif isinstance(docker_settings.requirements, List):
+            user_requirements = "\n".join(docker_settings.requirements)
+            if log:
+                logger.info(
+                    "- Including user-defined requirements: %s",
+                    ", ".join(f"`{r}`" for r in docker_settings.requirements),
+                )
+        else:
+            user_requirements = None
+
+        if user_requirements:
+            requirements_files.append(
+                (".zenml_user_requirements", user_requirements, [])
+            )
+
         return requirements_files
 
     @staticmethod

zenml/utils/pydantic_utils.py

@@ -45,7 +45,7 @@ def update_model(
     original: M,
     update: Union["BaseModel", Dict[str, Any]],
     recursive: bool = True,
-    exclude_none: bool = True,
+    exclude_none: bool = False,
 ) -> M:
     """Updates a pydantic model.
 
@@ -53,8 +53,7 @@ def update_model(
         original: The model to update.
         update: The update values.
         recursive: If `True`, dictionary values will be updated recursively.
-        exclude_none: If `True`, `None` values in the update dictionary
-            will be removed.
+        exclude_none: If `True`, `None` values in the update will be removed.
 
     Returns:
         The updated model.
@@ -67,7 +66,9 @@ def update_model(
         else:
             update_dict = update
     else:
-        update_dict = update.model_dump(exclude_unset=True)
+        update_dict = update.model_dump(
+            exclude_unset=True, exclude_none=exclude_none
+        )
 
     original_dict = original.model_dump(exclude_unset=True)
     if recursive:
@@ -75,7 +76,7 @@ def update_model(
     else:
         values = {**original_dict, **update_dict}
 
-    return original.__class__(**values)
+    return original.__class__.model_validate(values)
 
 
 class TemplateGenerator:

zenml/utils/source_utils.py

@@ -38,6 +38,7 @@ from zenml.config.source import (
     Source,
     SourceType,
 )
+from zenml.constants import ENV_ZENML_CUSTOM_SOURCE_ROOT
 from zenml.environment import Environment
 from zenml.logger import get_logger
 
@@ -57,7 +58,9 @@ BuiltinFunctionTypeSource = Source(
     type=SourceType.BUILTIN,
 )
 
-_CUSTOM_SOURCE_ROOT: Optional[str] = None
+_CUSTOM_SOURCE_ROOT: Optional[str] = os.getenv(
+    ENV_ZENML_CUSTOM_SOURCE_ROOT, None
+)
 
 
 def load(source: Union[Source, str]) -> Any:

zenml/zen_server/cloud_utils.py

@@ -12,6 +12,8 @@ from zenml.zen_server.utils import server_config
 
 ZENML_CLOUD_RBAC_ENV_PREFIX = "ZENML_CLOUD_"
 
+_cloud_connection: Optional["ZenMLCloudConnection"] = None
+
 
 class ZenMLCloudConfiguration(BaseModel):
     """ZenML Pro RBAC configuration."""
@@ -59,7 +61,7 @@ class ZenMLCloudConfiguration(BaseModel):
     )
 
 
-class ZenMLCloudSession:
+class ZenMLCloudConnection:
     """Class to use for communication between server and control plane."""
 
     def __init__(self) -> None:
@@ -67,7 +69,7 @@ class ZenMLCloudSession:
         self._config = ZenMLCloudConfiguration.from_environment()
         self._session: Optional[requests.Session] = None
 
-    def _get(
+    def get(
         self, endpoint: str, params: Optional[Dict[str, Any]]
     ) -> requests.Response:
         """Send a GET request using the active session.
@@ -105,7 +107,7 @@ class ZenMLCloudSession:
 
         return response
 
-    def _post(
+    def post(
         self,
         endpoint: str,
         params: Optional[Dict[str, Any]] = None,
@@ -217,3 +219,16 @@ class ZenMLCloudSession:
             raise RuntimeError("Could not fetch auth token from auth0.")
 
         return str(access_token)
+
+
+def cloud_connection() -> ZenMLCloudConnection:
+    """Return the initialized cloud connection.
+
+    Returns:
+        The cloud connection.
+    """
+    global _cloud_connection
+    if _cloud_connection is None:
+        _cloud_connection = ZenMLCloudConnection()
+
+    return _cloud_connection

zenml/zen_server/dashboard/assets/{404-CDPQCl4D.js → 404-B_YdvmwS.js}

@@ -1 +1 @@
-import{j as e}from"./@tanstack-CEbkxrhX.js";import{h as s,r as t}from"./index-QORVVTMN.js";import{E as r}from"./EmptyState-BkooiGtL.js";import{H as a}from"./help-B8rqCvqn.js";import{L as o}from"./@react-router-DYovave8.js";import"./@radix-C9DBgJhe.js";import"./@reactflow-CHBapDaj.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-CFOkMR_E.js";import{h as s,r as t}from"./index-BczVOqUf.js";import{E as r}from"./EmptyState-BMLnFVlB.js";import{H as a}from"./help-bu_DgLKI.js";import{L as o}from"./@react-router-CO-OsFwI.js";import"./@tanstack-DYiOyJUL.js";import"./@reactflow-l_1hUr1S.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};