zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.61.0.dev20240712__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.61.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.61.0
+helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.61.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/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/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/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],

zenml/zen_server/zen_server_api.py

@@ -63,6 +63,7 @@ from zenml.zen_server.routers import (
     service_connectors_endpoints,
     service_endpoints,
     stack_components_endpoints,
+    stack_deployment_endpoints,
     stacks_endpoints,
     steps_endpoints,
     tags_endpoints,
@@ -290,6 +291,7 @@ app.include_router(service_accounts_endpoints.router)
 app.include_router(service_connectors_endpoints.router)
 app.include_router(service_connectors_endpoints.types_router)
 app.include_router(service_endpoints.router)
+app.include_router(stack_deployment_endpoints.router)
 app.include_router(stacks_endpoints.router)
 app.include_router(stack_components_endpoints.router)
 app.include_router(stack_components_endpoints.types_router)

zenml/zen_stores/migrations/utils.py

@@ -364,7 +364,7 @@ class MigrationUtils(BaseModel):
                         # Convert column values to the correct type
                         for column in table.columns:
                             # Blob columns are stored as binary strings
-                            if column.type.python_type == bytes and isinstance(
+                            if column.type.python_type is bytes and isinstance(
                                 row[column.name], str
                             ):
                                 # Convert the string to bytes

zenml/zen_stores/migrations/versions/0.60.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.60.0].
+
+Revision ID: 0.60.0
+Revises: 25155145c545
+Create Date: 2024-06-25 14:26:48.867005
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.60.0"
+down_revision = "25155145c545"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/0.61.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.61.0].
+
+Revision ID: 0.61.0
+Revises: 0d707865f404
+Create Date: 2024-07-08 15:52:54.765307
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.61.0"
+down_revision = "0d707865f404"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py

@@ -0,0 +1,30 @@
+"""adding labels to stacks [0d707865f404].
+
+Revision ID: 0d707865f404
+Revises: 0.60.0
+Create Date: 2024-07-04 16:10:07.709184
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0d707865f404"
+down_revision = "0.60.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("labels", sa.LargeBinary(), nullable=True)
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_column("labels")