zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__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,18 +22,20 @@ import zenml
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
-    DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD,
+    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_SERVER_USE_LEGACY_DASHBOARD,
+    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 (
@@ -48,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__)
 
@@ -59,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:
@@ -70,7 +72,15 @@ class DockerServerDeploymentConfig(ServerDeploymentConfig):
     port: int = 8238
     image: str = DOCKER_ZENML_SERVER_DEFAULT_IMAGE
     store: Optional[StoreConfiguration] = None
-    use_legacy_dashboard: bool = DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD
+
+    @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")
 
@@ -158,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
@@ -174,9 +187,6 @@ class DockerZenServer(ContainerService):
             LOCAL_STORES_DIRECTORY_NAME,
         )
         env[ENV_ZENML_DISABLE_DATABASE_MIGRATION] = "True"
-        env[ENV_ZENML_SERVER_USE_LEGACY_DASHBOARD] = str(
-            self.config.server.use_legacy_dashboard
-        )
         env[ENV_ZENML_SERVER_AUTO_ACTIVATE] = "True"
 
         return cmd, env

zenml/zen_server/deploy/helm/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v2
 name: zenml
-version: "0.66.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.66.0
-helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.66.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/feature_gate/zenml_cloud_feature_gate.py

@@ -110,7 +110,7 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
             feature=resource,
             total=1 if not is_decrement else -1,
             metadata={
-                "tenant_id": str(server_config.external_server_id),
+                "tenant_id": str(server_config.get_external_server_id()),
                 "resource_id": str(resource_id),
             },
         ).model_dump()

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/artifact_version_endpoints.py

@@ -196,7 +196,7 @@ def delete_artifact_version(
 
 
 @artifact_version_router.delete(
-    "/",
+    "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

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/runs_endpoints.py

@@ -20,7 +20,6 @@ from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import (
     API,
-    GRAPH,
     PIPELINE_CONFIGURATION,
     REFRESH,
     RUNS,
@@ -29,7 +28,6 @@ from zenml.constants import (
     VERSION_1,
 )
 from zenml.enums import ExecutionStatus, StackComponentType
-from zenml.lineage_graph.lineage_graph import LineageGraph
 from zenml.logger import get_logger
 from zenml.models import (
     Page,
@@ -206,32 +204,6 @@ def delete_run(
     )
 
 
-@router.get(
-    "/{run_id}" + GRAPH,
-    response_model=LineageGraph,
-    responses={401: error_response, 404: error_response, 422: error_response},
-)
-@handle_exceptions
-def get_run_dag(
-    run_id: UUID,
-    _: AuthContext = Security(authorize),
-) -> LineageGraph:
-    """Get the DAG for a given pipeline run.
-
-    Args:
-        run_id: ID of the pipeline run to use to get the DAG.
-
-    Returns:
-        The DAG for a given pipeline run.
-    """
-    run = verify_permissions_and_get_entity(
-        id=run_id, get_method=zen_store().get_run, hydrate=True
-    )
-    graph = LineageGraph()
-    graph.generate_run_nodes_and_edges(run)
-    return graph
-
-
 @router.get(
     "/{run_id}" + STEPS,
     response_model=Page[StepRunResponse],

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.new.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,
+        )