zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__py3-none-any.whl

zenml/zen_server/deploy/deployment.py

@@ -21,19 +21,27 @@ from zenml.enums import ServerProviderType
 from zenml.services.service_status import ServiceState
 
 
-class ServerDeploymentConfig(BaseModel):
-    """Generic server deployment configuration.
+class LocalServerDeploymentConfig(BaseModel):
+    """Generic local server deployment configuration.
 
-    All server deployment configurations should inherit from this class and
-    handle extra attributes as provider specific attributes.
+    All local server deployment configurations should inherit from this class
+    and handle extra attributes as provider specific attributes.
 
     Attributes:
-        name: Name of the server deployment.
         provider: The server provider type.
     """
 
-    name: str
     provider: ServerProviderType
+
+    @property
+    def url(self) -> Optional[str]:
+        """Get the configured server URL.
+
+        Returns:
+            The configured server URL.
+        """
+        return None
+
     model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
@@ -44,13 +52,13 @@ class ServerDeploymentConfig(BaseModel):
     )
 
 
-class ServerDeploymentStatus(BaseModel):
-    """Server deployment status.
+class LocalServerDeploymentStatus(BaseModel):
+    """Local server deployment status.
 
     Ideally this should convey the following information:
 
     * whether the server's deployment is managed by this client (i.e. if
-    the server was deployed with `zenml up`)
+    the server was deployed with `zenml login --local`)
     * for a managed deployment, the status of the deployment/tear-down, e.g.
     not deployed, deploying, running, deleting, deployment timeout/error,
     tear-down timeout/error etc.
@@ -72,7 +80,7 @@ class ServerDeploymentStatus(BaseModel):
     ca_crt: Optional[str] = None
 
 
-class ServerDeployment(BaseModel):
+class LocalServerDeployment(BaseModel):
     """Server deployment.
 
     Attributes:
@@ -80,8 +88,8 @@ class ServerDeployment(BaseModel):
         status: The server deployment status.
     """
 
-    config: ServerDeploymentConfig
-    status: Optional[ServerDeploymentStatus] = None
+    config: LocalServerDeploymentConfig
+    status: Optional[LocalServerDeploymentStatus] = None
 
     @property
     def is_running(self) -> bool:

zenml/zen_server/deploy/docker/docker_provider.py

@@ -14,7 +14,7 @@
 """Zen Server docker deployer implementation."""
 
 import shutil
-from typing import ClassVar, List, Optional, Tuple, Type, cast
+from typing import ClassVar, Optional, Tuple, Type, cast
 from uuid import uuid4
 
 from zenml.enums import ServerProviderType
@@ -31,7 +31,7 @@ from zenml.services import (
     ServiceEndpointProtocol,
 )
 from zenml.zen_server.deploy.base_provider import BaseServerProvider
-from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
+from zenml.zen_server.deploy.deployment import LocalServerDeploymentConfig
 from zenml.zen_server.deploy.docker.docker_zen_server import (
     DOCKER_ZENML_SERVER_DEFAULT_TIMEOUT,
     ZEN_SERVER_HEALTHCHECK_URL_PATH,
@@ -47,14 +47,14 @@ class DockerServerProvider(BaseServerProvider):
     """Docker ZenML server provider."""
 
     TYPE: ClassVar[ServerProviderType] = ServerProviderType.DOCKER
-    CONFIG_TYPE: ClassVar[Type[ServerDeploymentConfig]] = (
+    CONFIG_TYPE: ClassVar[Type[LocalServerDeploymentConfig]] = (
         DockerServerDeploymentConfig
     )
 
     @classmethod
     def _get_service_configuration(
         cls,
-        server_config: ServerDeploymentConfig,
+        server_config: LocalServerDeploymentConfig,
     ) -> Tuple[
         ServiceConfig,
         ServiceEndpointConfig,
@@ -75,7 +75,7 @@ class DockerServerProvider(BaseServerProvider):
                 root_runtime_path=DockerZenServer.config_path(),
                 singleton=True,
                 image=server_config.image,
-                name=server_config.name,
+                name=ServerProviderType.DOCKER.value,
                 server=server_config,
             ),
             ContainerServiceEndpointConfig(
@@ -91,7 +91,7 @@ class DockerServerProvider(BaseServerProvider):
 
     def _create_service(
         self,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
         """Create, start and return the docker ZenML server deployment service.
@@ -142,7 +142,7 @@ class DockerServerProvider(BaseServerProvider):
     def _update_service(
         self,
         service: BaseService,
-        config: ServerDeploymentConfig,
+        config: LocalServerDeploymentConfig,
         timeout: Optional[int] = None,
     ) -> BaseService:
         """Update the docker ZenML server deployment service.
@@ -244,12 +244,9 @@ class DockerServerProvider(BaseServerProvider):
         service.stop(timeout)
         shutil.rmtree(DockerZenServer.config_path())
 
-    def _get_service(self, server_name: str) -> BaseService:
+    def _get_service(self) -> BaseService:
         """Get the docker ZenML server deployment service.
 
-        Args:
-            server_name: The server deployment name.
-
         Returns:
             The service instance.
 
@@ -260,27 +257,11 @@ class DockerServerProvider(BaseServerProvider):
         if service is None:
             raise KeyError("The docker ZenML server is not deployed.")
 
-        if service.config.name != server_name:
-            raise KeyError(
-                "The docker ZenML server is deployed but with a different name."
-            )
-
         return service
 
-    def _list_services(self) -> List[BaseService]:
-        """Get all service instances for all deployed ZenML servers.
-
-        Returns:
-            A list of service instances.
-        """
-        service = DockerZenServer.get_service()
-        if service:
-            return [service]
-        return []
-
     def _get_deployment_config(
         self, service: BaseService
-    ) -> ServerDeploymentConfig:
+    ) -> LocalServerDeploymentConfig:
         """Recreate the server deployment configuration from a service instance.
 
         Args:

zenml/zen_server/deploy/docker/docker_zen_server.py

@@ -22,16 +22,20 @@ import zenml
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
+    DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
     ENV_ZENML_ANALYTICS_OPT_IN,
     ENV_ZENML_CONFIG_PATH,
     ENV_ZENML_DISABLE_DATABASE_MIGRATION,
     ENV_ZENML_LOCAL_STORES_PATH,
+    ENV_ZENML_SERVER,
+    ENV_ZENML_SERVER_AUTH_SCHEME,
     ENV_ZENML_SERVER_AUTO_ACTIVATE,
     ENV_ZENML_SERVER_DEPLOYMENT_TYPE,
+    ENV_ZENML_USER_ID,
     LOCAL_STORES_DIRECTORY_NAME,
     ZEN_SERVER_ENTRYPOINT,
 )
-from zenml.enums import StoreType
+from zenml.enums import AuthScheme, StoreType
 from zenml.logger import get_logger
 from zenml.models import ServerDeploymentType
 from zenml.services import (
@@ -46,7 +50,7 @@ from zenml.services.container.container_service import (
     SERVICE_CONTAINER_PATH,
 )
 from zenml.utils.io_utils import get_global_config_directory
-from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
+from zenml.zen_server.deploy.deployment import LocalServerDeploymentConfig
 
 logger = get_logger(__name__)
 
@@ -57,7 +61,7 @@ DOCKER_ZENML_SERVER_DEFAULT_IMAGE = (
 DOCKER_ZENML_SERVER_DEFAULT_TIMEOUT = 60
 
 
-class DockerServerDeploymentConfig(ServerDeploymentConfig):
+class DockerServerDeploymentConfig(LocalServerDeploymentConfig):
     """Docker server deployment configuration.
 
     Attributes:
@@ -69,6 +73,15 @@ class DockerServerDeploymentConfig(ServerDeploymentConfig):
     image: str = DOCKER_ZENML_SERVER_DEFAULT_IMAGE
     store: Optional[StoreConfiguration] = None
 
+    @property
+    def url(self) -> Optional[str]:
+        """Get the configured server URL.
+
+        Returns:
+            The configured server URL.
+        """
+        return f"http://{DEFAULT_LOCAL_SERVICE_IP_ADDRESS}:{self.port}"
+
     model_config = ConfigDict(extra="forbid")
 
 
@@ -155,12 +168,15 @@ class DockerZenServer(ContainerService):
         gc = GlobalConfiguration()
 
         cmd, env = super()._get_container_cmd()
+        env[ENV_ZENML_SERVER] = "true"
         env[ENV_ZENML_CONFIG_PATH] = os.path.join(
             SERVICE_CONTAINER_PATH,
             SERVICE_CONTAINER_GLOBAL_CONFIG_DIR,
         )
+        env[ENV_ZENML_SERVER_AUTH_SCHEME] = AuthScheme.NO_AUTH.value
         env[ENV_ZENML_SERVER_DEPLOYMENT_TYPE] = ServerDeploymentType.DOCKER
         env[ENV_ZENML_ANALYTICS_OPT_IN] = str(gc.analytics_opt_in)
+        env[ENV_ZENML_USER_ID] = str(gc.user_id)
 
         # Set the local stores path to the same path used by the client (mounted
         # in the container by the super class). This ensures that the server's

zenml/zen_server/deploy/helm/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v2
 name: zenml
-version: "0.68.0"
+version: "0.68.1"
 description: Open source MLOps framework for portable production ready ML pipelines
 keywords:
 - mlops

zenml/zen_server/deploy/helm/README.md

@@ -20,8 +20,8 @@ ZenML is an open-source MLOps framework designed to help you create robust, main
 To install the ZenML chart directly from Amazon ECR, use the following command:
 
 ```bash
-# example command for version 0.68.0
-helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.68.0
+# example command for version 0.68.1
+helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.68.1
 ```
 
 Note: Ensure you have OCI support enabled in your Helm client and that you are authenticated with Amazon ECR.

zenml/zen_server/exceptions.py

@@ -20,6 +20,7 @@ from pydantic import BaseModel
 
 from zenml.exceptions import (
     AuthorizationException,
+    CredentialsNotValid,
     DoesNotExistException,
     DuplicateRunNameError,
     EntityExistsError,
@@ -78,6 +79,7 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     # 403 Forbidden
     (IllegalOperationError, 403),
     # 401 Unauthorized
+    (CredentialsNotValid, 401),
     (AuthorizationException, 401),
     # 402 Payment required
     (SubscriptionUpgradeRequiredError, 402),
@@ -253,4 +255,13 @@ def exception_from_response(
 
         exception = default_exc
 
+    # There is one special case where we want to return a specific exception:
+    # 401 Unauthorized exceptions thrown directly by FastAPI in the course of
+    # authentication are interpreted as AuthorizationException, but we want to
+    # return CredentialsNotValid instead.
+    if response.status_code == 401:
+        if not isinstance(exception(), CredentialsNotValid):
+            if response.headers.get("WWW-Authenticate"):
+                return CredentialsNotValid(exc_msg)
+
     return exception(exc_msg)

zenml/zen_server/jwt.py

@@ -25,7 +25,7 @@ from uuid import UUID
 import jwt
 from pydantic import BaseModel
 
-from zenml.exceptions import AuthorizationException
+from zenml.exceptions import CredentialsNotValid
 from zenml.logger import get_logger
 from zenml.zen_server.utils import server_config
 
@@ -71,7 +71,7 @@ class JWTToken(BaseModel):
             The decoded JWT access token.
 
         Raises:
-            AuthorizationException: If the token is invalid.
+            CredentialsNotValid: If the token is invalid.
         """
         config = server_config()
 
@@ -87,18 +87,18 @@ class JWTToken(BaseModel):
             )
             claims = cast(Dict[str, Any], claims_data)
         except jwt.PyJWTError as e:
-            raise AuthorizationException(f"Invalid JWT token: {e}") from e
+            raise CredentialsNotValid(f"Invalid JWT token: {e}") from e
 
         subject: str = claims.pop("sub", "")
         if not subject:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is missing"
             )
 
         try:
             user_id = UUID(subject)
         except ValueError:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is not a valid UUID"
             )
 
@@ -107,7 +107,7 @@ class JWTToken(BaseModel):
             try:
                 device_id = UUID(claims.pop("device_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the device_id claim is not a valid "
                     "UUID"
                 )
@@ -117,7 +117,7 @@ class JWTToken(BaseModel):
             try:
                 api_key_id = UUID(claims.pop("api_key_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the api_key_id claim is not a valid "
                     "UUID"
                 )
@@ -127,7 +127,7 @@ class JWTToken(BaseModel):
             try:
                 pipeline_id = UUID(claims.pop("pipeline_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the pipeline_id claim is not a valid "
                     "UUID"
                 )
@@ -137,7 +137,7 @@ class JWTToken(BaseModel):
             try:
                 schedule_id = UUID(claims.pop("schedule_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the schedule_id claim is not a valid "
                     "UUID"
                 )

zenml/zen_server/routers/auth_endpoints.py

@@ -116,6 +116,8 @@ class OAuthLoginRequestForm:
         Raises:
             HTTPException: If the request is invalid.
         """
+        config = server_config()
+
         if not grant_type:
             # Detect the grant type from the form data
             if username is not None:
@@ -124,8 +126,21 @@ class OAuthLoginRequestForm:
                 self.grant_type = OAuthGrantTypes.ZENML_API_KEY
             elif device_code:
                 self.grant_type = OAuthGrantTypes.OAUTH_DEVICE_CODE
-            else:
+            elif config.auth_scheme == AuthScheme.EXTERNAL:
                 self.grant_type = OAuthGrantTypes.ZENML_EXTERNAL
+            elif config.auth_scheme in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
+                # For no auth and basic HTTP auth schemes, we also allow the
+                # password grant type to be used for backwards compatibility
+                self.grant_type = OAuthGrantTypes.OAUTH_PASSWORD
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Invalid request: grant type is required.",
+                )
         else:
             if grant_type not in OAuthGrantTypes.values():
                 logger.info(
@@ -137,10 +152,15 @@ class OAuthLoginRequestForm:
                 )
             self.grant_type = OAuthGrantTypes(grant_type)
 
-        config = server_config()
-
         if self.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
-            if config.auth_scheme != AuthScheme.OAUTH2_PASSWORD_BEARER:
+            # For the no auth and basic HTTP auth schemes, we also allow the
+            # password grant type to be used for compatibility with other
+            # auth schemes
+            if config.auth_scheme not in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
                 logger.info(
                     f"Request with unsupported grant type: {self.grant_type}"
                 )
@@ -280,6 +300,7 @@ def token(
     Raises:
         ValueError: If the grant type is invalid.
     """
+    config = server_config()
     if auth_form_data.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
         auth_context = authenticate_credentials(
             user_name_or_id=auth_form_data.username,
@@ -297,8 +318,6 @@ def token(
         )
 
     elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_EXTERNAL:
-        config = server_config()
-
         assert config.external_cookie_name is not None
         assert config.external_login_url is not None
 
@@ -399,8 +418,11 @@ def device_authorization(
     config = server_config()
     store = zen_store()
 
-    # Use this opportunity to delete expired devices
-    store.delete_expired_authorized_devices()
+    try:
+        # Use this opportunity to delete expired devices
+        store.delete_expired_authorized_devices()
+    except Exception:
+        logger.exception("Failed to delete expired devices")
 
     # Fetch additional details about the client from the user-agent header
     user_agent_header = request.headers.get("User-Agent")

zenml/zen_server/routers/stack_components_endpoints.py

@@ -141,7 +141,7 @@ def update_stack_component(
         existing_component = zen_store().get_stack_component(component_id)
         validate_stack_component_config(
             configuration_dict=component_update.configuration,
-            flavor_name=existing_component.flavor,
+            flavor=existing_component.flavor_name,
             component_type=existing_component.type,
             zen_store=zen_store(),
             # We allow custom flavors to fail import on the server side.

zenml/zen_server/routers/workspaces_endpoints.py

@@ -453,7 +453,7 @@ def create_stack_component(
 
     validate_stack_component_config(
         configuration_dict=component.configuration,
-        flavor_name=component.flavor,
+        flavor=component.flavor,
         component_type=component.type,
         zen_store=zen_store(),
         # We allow custom flavors to fail import on the server side.

zenml/zen_server/template_execution/runner_entrypoint_configuration.py

@@ -17,9 +17,7 @@ from zenml.client import Client
 from zenml.entrypoints.base_entrypoint_configuration import (
     BaseEntrypointConfiguration,
 )
-from zenml.pipelines.run_utils import (
-    deploy_pipeline,
-)
+from zenml.pipelines.run_utils import deploy_pipeline, get_placeholder_run
 
 
 class RunnerEntrypointConfiguration(BaseEntrypointConfiguration):
@@ -36,4 +34,9 @@ class RunnerEntrypointConfiguration(BaseEntrypointConfiguration):
         stack = Client().active_stack
         assert deployment.stack and stack.id == deployment.stack.id
 
-        deploy_pipeline(deployment=deployment, stack=stack)
+        placeholder_run = get_placeholder_run(deployment_id=deployment.id)
+        deploy_pipeline(
+            deployment=deployment,
+            stack=stack,
+            placeholder_run=placeholder_run,
+        )

zenml/zen_server/template_execution/utils.py

@@ -3,7 +3,7 @@
 import copy
 import hashlib
 import sys
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional
 from uuid import UUID
 
 from fastapi import BackgroundTasks
@@ -22,11 +22,9 @@ from zenml.constants import (
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
 )
 from zenml.enums import ExecutionStatus, StackComponentType, StoreType
-from zenml.integrations.utils import get_integration_for_module
 from zenml.logger import get_logger
 from zenml.models import (
     CodeReferenceRequest,
-    ComponentResponse,
     FlavorFilter,
     PipelineDeploymentRequest,
     PipelineDeploymentResponse,
@@ -43,7 +41,7 @@ from zenml.pipelines.run_utils import (
     validate_stack_is_runnable_from_server,
 )
 from zenml.stack.flavor import Flavor
-from zenml.utils import dict_utils, settings_utils
+from zenml.utils import dict_utils, requirements_utils, settings_utils
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
@@ -151,8 +149,8 @@ def run_template(
     assert placeholder_run
 
     def _task() -> None:
-        pypi_requirements, apt_packages = get_requirements_for_stack(
-            stack=stack
+        pypi_requirements, apt_packages = (
+            requirements_utils.get_requirements_for_stack(stack=stack)
         )
 
         if build.python_version:
@@ -248,7 +246,7 @@ def ensure_async_orchestrator(
     """
     orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
     flavors = zen_store().list_flavors(
-        FlavorFilter(name=orchestrator.flavor, type=orchestrator.type)
+        FlavorFilter(name=orchestrator.flavor_name, type=orchestrator.type)
     )
     flavor = Flavor.from_model(flavors[0])
 
@@ -266,59 +264,6 @@ def ensure_async_orchestrator(
         )
 
 
-def get_requirements_for_stack(
-    stack: StackResponse,
-) -> Tuple[List[str], List[str]]:
-    """Get requirements for a stack model.
-
-    Args:
-        stack: The stack for which to get the requirements.
-
-    Returns:
-        Tuple of PyPI and APT requirements of the stack.
-    """
-    pypi_requirements: Set[str] = set()
-    apt_packages: Set[str] = set()
-
-    for component_list in stack.components.values():
-        assert len(component_list) == 1
-        component = component_list[0]
-        (
-            component_pypi_requirements,
-            component_apt_packages,
-        ) = get_requirements_for_component(component=component)
-        pypi_requirements = pypi_requirements.union(
-            component_pypi_requirements
-        )
-        apt_packages = apt_packages.union(component_apt_packages)
-
-    return sorted(pypi_requirements), sorted(apt_packages)
-
-
-def get_requirements_for_component(
-    component: ComponentResponse,
-) -> Tuple[List[str], List[str]]:
-    """Get requirements for a component model.
-
-    Args:
-        component: The component for which to get the requirements.
-
-    Returns:
-        Tuple of PyPI and APT requirements of the component.
-    """
-    flavors = zen_store().list_flavors(
-        FlavorFilter(name=component.flavor, type=component.type)
-    )
-    assert len(flavors) == 1
-    flavor_source = flavors[0].source
-    integration = get_integration_for_module(module_name=flavor_source)
-
-    if integration:
-        return integration.get_requirements(), integration.APT_PACKAGES
-    else:
-        return [], []
-
-
 def generate_image_hash(dockerfile: str) -> str:
     """Generate a hash of the Dockerfile.
 
@@ -512,7 +457,7 @@ def get_pipeline_run_analytics_metadata(
     own_stack = stack_creator and stack_creator.id == deployment.user.id
 
     stack_metadata = {
-        component_type.value: component_list[0].flavor
+        component_type.value: component_list[0].flavor_name
         for component_type, component_list in stack.components.items()
     }