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

zenml/zen_server/deploy/helm/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v2
 name: zenml
-version: "0.58.2"
+version: "0.62.0"
 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.58.2
-helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.58.2
+# example command for version 0.62.0
+helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.62.0
 ```
 
 Note: Ensure you have OCI support enabled in your Helm client and that you are authenticated with Amazon ECR.

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -21,7 +21,7 @@ from pydantic import BaseModel, Field
 from zenml.config.server_config import ServerConfiguration
 from zenml.exceptions import SubscriptionUpgradeRequiredError
 from zenml.logger import get_logger
-from zenml.zen_server.cloud_utils import ZenMLCloudSession
+from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.feature_gate.feature_gate_interface import (
     FeatureGateInterface,
 )
@@ -59,8 +59,12 @@ class RawUsageEvent(BaseModel):
     )
 
 
-class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
-    """Feature Gate interface definition."""
+class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
+    """ZenML Cloud Feature Gate implementation."""
+
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
 
     def check_entitlement(self, resource: ResourceType) -> None:
         """Checks if a user is entitled to create a resource.
@@ -72,7 +76,7 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
             SubscriptionUpgradeRequiredError: in case a subscription limit is reached
         """
         try:
-            response = self._get(
+            response = self._connection.get(
                 endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
             )
         except SubscriptionUpgradeRequiredError:
@@ -110,7 +114,9 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
                 "resource_id": str(resource_id),
             },
         ).model_dump()
-        response = self._post(endpoint=USAGE_EVENT_ENDPOINT, data=data)
+        response = self._connection.post(
+            endpoint=USAGE_EVENT_ENDPOINT, data=data
+        )
         if response.status_code != 200:
             logger.error(
                 "Usage report not accepted by upstream backend. "

zenml/zen_server/pipeline_deployment/utils.py

@@ -17,7 +17,7 @@ from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
 )
-from zenml.enums import StackComponentType, StoreType
+from zenml.enums import ExecutionStatus, StackComponentType, StoreType
 from zenml.integrations.utils import get_integration_for_module
 from zenml.models import (
     CodeReferenceRequest,
@@ -26,6 +26,7 @@ from zenml.models import (
     PipelineDeploymentRequest,
     PipelineDeploymentResponse,
     PipelineRunResponse,
+    PipelineRunUpdate,
     StackResponse,
 )
 from zenml.new.pipelines.run_utils import (
@@ -88,8 +89,6 @@ def run_pipeline(
     ensure_async_orchestrator(deployment=deployment_request, stack=stack)
 
     new_deployment = zen_store().create_deployment(deployment_request)
-    placeholder_run = create_placeholder_run(deployment=new_deployment)
-    assert placeholder_run
 
     if auth_context.access_token:
         token = auth_context.access_token
@@ -125,52 +124,62 @@ def run_pipeline(
         deployment_id=new_deployment.id
     )
 
-    def _task() -> None:
-        pypi_requirements, apt_packages = get_requirements_for_stack(
-            stack=stack
-        )
+    placeholder_run = create_placeholder_run(deployment=new_deployment)
+    assert placeholder_run
 
-        if build.python_version:
-            version_info = version.parse(build.python_version)
-            python_version = f"{version_info.major}.{version_info.minor}"
-        else:
-            python_version = (
-                f"{sys.version_info.major}.{sys.version_info.minor}"
+    def _task() -> None:
+        try:
+            pypi_requirements, apt_packages = get_requirements_for_stack(
+                stack=stack
             )
 
-        dockerfile = generate_dockerfile(
-            pypi_requirements=pypi_requirements,
-            apt_packages=apt_packages,
-            zenml_version=zenml_version,
-            python_version=python_version,
-        )
+            if build.python_version:
+                version_info = version.parse(build.python_version)
+                python_version = f"{version_info.major}.{version_info.minor}"
+            else:
+                python_version = (
+                    f"{sys.version_info.major}.{sys.version_info.minor}"
+                )
+
+            dockerfile = generate_dockerfile(
+                pypi_requirements=pypi_requirements,
+                apt_packages=apt_packages,
+                zenml_version=zenml_version,
+                python_version=python_version,
+            )
 
-        image_hash = generate_image_hash(dockerfile=dockerfile)
+            image_hash = generate_image_hash(dockerfile=dockerfile)
 
-        runner_image = workload_manager().build_and_push_image(
-            workload_id=new_deployment.id,
-            dockerfile=dockerfile,
-            image_name=f"{RUNNER_IMAGE_REPOSITORY}:{image_hash}",
-            sync=True,
-        )
+            runner_image = workload_manager().build_and_push_image(
+                workload_id=new_deployment.id,
+                dockerfile=dockerfile,
+                image_name=f"{RUNNER_IMAGE_REPOSITORY}:{image_hash}",
+                sync=True,
+            )
 
-        workload_manager().log(
-            workload_id=new_deployment.id,
-            message="Starting pipeline deployment.",
-        )
-        workload_manager().run(
-            workload_id=new_deployment.id,
-            image=runner_image,
-            command=command,
-            arguments=args,
-            environment=environment,
-            timeout_in_seconds=30,
-            sync=True,
-        )
-        workload_manager().log(
-            workload_id=new_deployment.id,
-            message="Pipeline deployed successfully.",
-        )
+            workload_manager().log(
+                workload_id=new_deployment.id,
+                message="Starting pipeline deployment.",
+            )
+            workload_manager().run(
+                workload_id=new_deployment.id,
+                image=runner_image,
+                command=command,
+                arguments=args,
+                environment=environment,
+                timeout_in_seconds=30,
+                sync=True,
+            )
+            workload_manager().log(
+                workload_id=new_deployment.id,
+                message="Pipeline deployed successfully.",
+            )
+        except Exception:
+            zen_store().update_run(
+                run_id=placeholder_run.id,
+                run_update=PipelineRunUpdate(status=ExecutionStatus.FAILED),
+            )
+            raise
 
     if background_tasks:
         background_tasks.add_task(_task)
@@ -357,8 +366,12 @@ def apply_run_config(
         step_config = StepConfiguration.model_validate(step_config_dict)
 
         if update := run_config.steps.get(invocation_id):
+            update_dict = update.model_dump()
+            # Get rid of deprecated name to prevent overriding the step name
+            # with `None`.
+            update_dict.pop("name", None)
             step_config = pydantic_utils.update_model(
-                step_config, update=update
+                step_config, update=update_dict
             )
         steps[invocation_id] = Step(spec=step.spec, config=step_config)
 

zenml/zen_server/rbac/utils.py

@@ -103,7 +103,9 @@ def dehydrate_response_model(
         )
 
     dehydrated_values = {}
-    for key, value in dict(model).items():
+    # See `get_subresources_for_model(...)` for a detailed explanation why we
+    # need to use `model.__iter__()` here
+    for key, value in model.__iter__():
         dehydrated_values[key] = _dehydrate_value(
             value, permissions=permissions
         )
@@ -484,7 +486,13 @@ def get_subresources_for_model(
     """
     resources = set()
 
-    for value in dict(model).values():
+    # We don't want to use `model.model_dump()` here as that recursively
+    # converts models to dicts, but we want to preserve those classes for
+    # the recursive `_get_subresources_for_value` calls.
+    # We previously used `dict(model)` here, but that lead to issues with
+    # models overwriting `__getattr__`, this `model.__iter__()` has the same
+    # results though.
+    for _, value in model.__iter__():
         resources.update(_get_subresources_for_value(value))
 
     return resources

zenml/zen_server/rbac/zenml_cloud_rbac.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Dict, List, Set, Tuple
 
-from zenml.zen_server.cloud_utils import ZenMLCloudSession
+from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.rbac.models import Action, Resource
 from zenml.zen_server.rbac.rbac_interface import RBACInterface
 from zenml.zen_server.utils import server_config
@@ -74,9 +74,13 @@ def _convert_from_cloud_resource(cloud_resource: str) -> Resource:
         return Resource(type=resource_type_and_id)
 
 
-class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
+class ZenMLCloudRBAC(RBACInterface):
     """RBAC implementation that uses the ZenML Pro Management Plane as a backend."""
 
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
+
     def check_permissions(
         self, user: "UserResponse", resources: Set[Resource], action: Action
     ) -> Dict[Resource, bool]:
@@ -110,7 +114,9 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             ],
             "action": str(action),
         }
-        response = self._get(endpoint=PERMISSIONS_ENDPOINT, params=params)
+        response = self._connection.get(
+            endpoint=PERMISSIONS_ENDPOINT, params=params
+        )
         value = response.json()
 
         assert isinstance(value, dict)
@@ -147,7 +153,7 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             "resource": _convert_to_cloud_resource(resource),
             "action": str(action),
         }
-        response = self._get(
+        response = self._connection.get(
             endpoint=ALLOWED_RESOURCE_IDS_ENDPOINT, params=params
         )
         response_json = response.json()
@@ -177,4 +183,4 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             "resource": _convert_to_cloud_resource(resource),
             "actions": [str(action) for action in actions],
         }
-        self._post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)
+        self._connection.post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)

zenml/zen_server/routers/devices_endpoints.py

@@ -27,7 +27,7 @@ from zenml.constants import (
     DEVICES,
     VERSION_1,
 )
-from zenml.enums import OAuthDeviceStatus
+from zenml.enums import OAuthDeviceStatus, OnboardingStep
 from zenml.models import (
     OAuthDeviceFilter,
     OAuthDeviceInternalUpdate,
@@ -270,6 +270,9 @@ def verify_authorized_device(
             update=update,
         )
 
+    store.update_onboarding_state(
+        completed_steps={OnboardingStep.DEVICE_VERIFIED}
+    )
     return device_model
 
 

zenml/zen_server/routers/server_endpoints.py

@@ -13,12 +13,19 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from typing import Optional
+from typing import List, Optional
 
 from fastapi import APIRouter, Security
 
 import zenml
-from zenml.constants import ACTIVATE, API, INFO, SERVER_SETTINGS, VERSION_1
+from zenml.constants import (
+    ACTIVATE,
+    API,
+    INFO,
+    ONBOARDING_STATE,
+    SERVER_SETTINGS,
+    VERSION_1,
+)
 from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
@@ -64,6 +71,26 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    ONBOARDING_STATE,
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def get_onboarding_state(
+    _: AuthContext = Security(authorize),
+) -> List[str]:
+    """Get the onboarding state of the server.
+
+    Returns:
+        The onboarding state of the server.
+    """
+    return zen_store().get_onboarding_state()
+
+
 # We don't have any concrete value that tells us whether a server is a cloud
 # tenant, so we use `external_server_id` as the best proxy option.
 # For cloud tenants, we don't add these endpoints as the server settings don't

zenml/zen_server/routers/service_connectors_endpoints.py

@@ -21,6 +21,7 @@ from fastapi import APIRouter, Depends, Security
 from zenml.constants import (
     API,
     SERVICE_CONNECTOR_CLIENT,
+    SERVICE_CONNECTOR_FULL_STACK,
     SERVICE_CONNECTOR_TYPES,
     SERVICE_CONNECTOR_VERIFY,
     SERVICE_CONNECTORS,
@@ -35,6 +36,13 @@ from zenml.models import (
     ServiceConnectorTypeModel,
     ServiceConnectorUpdate,
 )
+from zenml.models.v2.misc.full_stack import (
+    ServiceConnectorInfo,
+    ServiceConnectorResourcesInfo,
+)
+from zenml.service_connectors.service_connector_utils import (
+    get_resources_options_from_resource_model_for_full_stack,
+)
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
@@ -393,3 +401,52 @@ def get_service_connector_type(
         The requested service connector type.
     """
     return zen_store().get_service_connector_type(connector_type)
+
+
+@router.post(
+    SERVICE_CONNECTOR_FULL_STACK,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_resources_based_on_service_connector_info(
+    connector_info: Optional[ServiceConnectorInfo] = None,
+    connector_uuid: Optional[UUID] = None,
+    _: AuthContext = Security(authorize),
+) -> ServiceConnectorResourcesInfo:
+    """Gets the list of resources that a service connector can access.
+
+    Args:
+        connector_info: The service connector info.
+        connector_uuid: The service connector uuid.
+
+    Returns:
+        The list of resources that the service connector configuration has
+        access to and consumable from UI/CLI.
+
+    Raises:
+        ValueError: If both connector_info and connector_uuid are provided.
+        ValueError: If neither connector_info nor connector_uuid are provided.
+    """
+    if connector_info is not None and connector_uuid is not None:
+        raise ValueError(
+            "Only one of connector_info or connector_uuid must be provided."
+        )
+    if connector_info is None and connector_uuid is None:
+        raise ValueError(
+            "Either connector_info or connector_uuid must be provided."
+        )
+
+    if connector_info is not None:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+        )
+    elif connector_uuid is not None:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR,
+            action=Action.READ,
+            resource_id=connector_uuid,
+        )
+
+    return get_resources_options_from_resource_model_for_full_stack(
+        connector_details=connector_info or connector_uuid  # type: ignore[arg-type]
+    )

zenml/zen_server/routers/stack_deployment_endpoints.py

@@ -0,0 +1,158 @@
+#  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.
+"""Endpoint definitions for stack deployments."""
+
+import datetime
+from typing import Optional
+
+from fastapi import APIRouter, Request, Security
+
+from zenml.constants import (
+    API,
+    CONFIG,
+    INFO,
+    STACK,
+    STACK_DEPLOYMENT,
+    STACK_DEPLOYMENT_API_TOKEN_EXPIRATION,
+    VERSION_1,
+)
+from zenml.enums import StackDeploymentProvider
+from zenml.models import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
+from zenml.stack_deployments.utils import get_stack_deployment_class
+from zenml.zen_server.auth import AuthContext, authorize
+from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.utils import verify_permission
+from zenml.zen_server.utils import (
+    handle_exceptions,
+)
+
+router = APIRouter(
+    prefix=API + VERSION_1 + STACK_DEPLOYMENT,
+    tags=["stacks"],
+    responses={401: error_response, 403: error_response},
+)
+
+
+@router.get(
+    INFO,
+)
+@handle_exceptions
+def get_stack_deployment_info(
+    provider: StackDeploymentProvider,
+    _: AuthContext = Security(authorize),
+) -> StackDeploymentInfo:
+    """Get information about a stack deployment provider.
+
+    Args:
+        provider: The stack deployment provider.
+
+    Returns:
+        Information about the stack deployment provider.
+    """
+    stack_deployment_class = get_stack_deployment_class(provider)
+    return stack_deployment_class.get_deployment_info()
+
+
+@router.get(
+    CONFIG,
+)
+@handle_exceptions
+def get_stack_deployment_config(
+    request: Request,
+    provider: StackDeploymentProvider,
+    stack_name: str,
+    location: Optional[str] = None,
+    auth_context: AuthContext = Security(authorize),
+) -> StackDeploymentConfig:
+    """Return the URL to deploy the ZenML stack to the specified cloud provider.
+
+    Args:
+        request: The FastAPI request object.
+        provider: The stack deployment provider.
+        stack_name: The name of the stack.
+        location: The location where the stack should be deployed.
+        auth_context: The authentication context.
+
+    Returns:
+        The cloud provider console URL where the stack will be deployed and
+        the configuration for the stack deployment.
+    """
+    verify_permission(
+        resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+    )
+    verify_permission(
+        resource_type=ResourceType.STACK_COMPONENT,
+        action=Action.CREATE,
+    )
+    verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE)
+
+    stack_deployment_class = get_stack_deployment_class(provider)
+    # Get the base server URL used to call this FastAPI endpoint
+    url = request.url.replace(path="").replace(query="")
+    # Use HTTPS for the URL
+    url = url.replace(scheme="https")
+
+    token = auth_context.access_token
+    assert token is not None
+
+    # A new API token is generated for the stack deployment
+    expires = datetime.datetime.utcnow() + datetime.timedelta(
+        minutes=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION
+    )
+    api_token = token.encode(expires=expires)
+
+    return stack_deployment_class(
+        stack_name=stack_name,
+        location=location,
+        zenml_server_url=str(url),
+        zenml_server_api_token=api_token,
+    ).get_deployment_config()
+
+
+@router.get(
+    STACK,
+)
+@handle_exceptions
+def get_deployed_stack(
+    provider: StackDeploymentProvider,
+    stack_name: str,
+    location: Optional[str] = None,
+    date_start: Optional[datetime.datetime] = None,
+    _: AuthContext = Security(authorize),
+) -> Optional[DeployedStack]:
+    """Return a matching ZenML stack that was deployed and registered.
+
+    Args:
+        provider: The stack deployment provider.
+        stack_name: The name of the stack.
+        location: The location where the stack should be deployed.
+        date_start: The date when the deployment started.
+
+    Returns:
+        The ZenML stack that was deployed and registered or None if the stack
+        was not found.
+    """
+    stack_deployment_class = get_stack_deployment_class(provider)
+    return stack_deployment_class(
+        stack_name=stack_name,
+        location=location,
+        # These fields are not needed for this operation
+        zenml_server_url="",
+        zenml_server_api_token="",
+    ).get_stack(date_start=date_start)

zenml/zen_server/routers/steps_endpoints.py

@@ -117,7 +117,8 @@ def create_run_step(
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
     verify_permission_for_model(pipeline_run, action=Action.UPDATE)
 
-    return zen_store().create_run_step(step_run=step)
+    step_response = zen_store().create_run_step(step_run=step)
+    return dehydrate_response_model(step_response)
 
 
 @router.get(

zenml/zen_server/routers/workspaces_endpoints.py

@@ -22,6 +22,7 @@ from zenml.constants import (
     API,
     ARTIFACTS,
     CODE_REPOSITORIES,
+    FULL_STACK,
     GET_OR_CREATE,
     MODEL_VERSIONS,
     MODELS,
@@ -51,6 +52,7 @@ from zenml.models import (
     ComponentFilter,
     ComponentRequest,
     ComponentResponse,
+    FullStackRequest,
     ModelRequest,
     ModelResponse,
     ModelVersionArtifactRequest,
@@ -349,6 +351,68 @@ def create_stack(
     )
 
 
+@router.post(
+    WORKSPACES + "/{workspace_name_or_id}" + FULL_STACK,
+    response_model=StackResponse,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_full_stack(
+    workspace_name_or_id: Union[str, UUID],
+    full_stack: FullStackRequest,
+    auth_context: AuthContext = Security(authorize),
+) -> StackResponse:
+    """Creates a stack for a particular workspace.
+
+    Args:
+        workspace_name_or_id: Name or ID of the workspace.
+        full_stack: Stack to register.
+        auth_context: Authentication context.
+
+    Returns:
+        The created stack.
+    """
+    workspace = zen_store().get_workspace(workspace_name_or_id)
+
+    is_connector_create_needed = False
+    for connector_id_or_info in full_stack.service_connectors:
+        if isinstance(connector_id_or_info, UUID):
+            service_connector = zen_store().get_service_connector(
+                connector_id_or_info, hydrate=False
+            )
+            verify_permission_for_model(
+                model=service_connector, action=Action.READ
+            )
+        else:
+            is_connector_create_needed = True
+    if is_connector_create_needed:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+        )
+
+    is_component_create_needed = False
+    for component_id_or_info in full_stack.components.values():
+        if isinstance(component_id_or_info, UUID):
+            component = zen_store().get_stack_component(
+                component_id_or_info, hydrate=False
+            )
+            verify_permission_for_model(model=component, action=Action.READ)
+        else:
+            is_component_create_needed = True
+    if is_component_create_needed:
+        verify_permission(
+            resource_type=ResourceType.STACK_COMPONENT,
+            action=Action.CREATE,
+        )
+
+    verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE)
+
+    full_stack.user = auth_context.user.id
+    full_stack.workspace = workspace.id
+
+    return zen_store().create_full_stack(full_stack)
+
+
 @router.get(
     WORKSPACES + "/{workspace_name_or_id}" + STACK_COMPONENTS,
     response_model=Page[ComponentResponse],