zenml-nightly 0.68.1.dev20241111__py3-none-any.whl → 0.70.0.dev20241114__py3-none-any.whl

zenml/stack_deployments/gcp_stack_deployment.py

@@ -70,11 +70,12 @@ and are aware of any potential costs:
 
 - A GCS bucket registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/gcp).
 - A Google Artifact Registry registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/gcp).
-- Vertex AI registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/vertex).
+- Vertex AI registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/vertex)
+and as a [ZenML step operator](https://docs.zenml.io/stack-components/step-operators/vertex).
 - GCP Cloud Build registered as a [ZenML image builder](https://docs.zenml.io/stack-components/image-builders/gcp).
 - A GCP Service Account with the minimum necessary permissions to access the
 above resources.
-- An GCP Service Account access key used to give access to ZenML to connect to
+- A GCP Service Account access key used to give access to ZenML to connect to
 the above resources through a [ZenML service connector](https://docs.zenml.io/how-to/auth-management/gcp-service-connector).
 
 The Deployment Manager deployment will automatically create a GCP Service
@@ -259,14 +260,30 @@ GCP project and to clean up the resources created by the stack by using
         )
 
         if self.deployment_type == STACK_DEPLOYMENT_TERRAFORM:
-            config = f"""module "zenml_stack" {{
+            config = f"""terraform {{
+    required_providers {{
+        google = {{
+            source  = "hashicorp/google"
+        }}
+        zenml = {{
+            source = "zenml-io/zenml"
+        }}
+    }}
+}}
+
+provider "google" {{
+    region  = "{self.location or "europe-west3"}"
+    project = your GCP project name
+}}
+
+provider "zenml" {{
+    server_url = "{self.zenml_server_url}"
+    api_token = "{self.zenml_server_api_token}"
+}}
+
+module "zenml_stack" {{
     source  = "zenml-io/zenml-stack/gcp"
 
-    project_id = "my-gcp-project"
-    region = "{self.location or "europe-west3"}"
-    zenml_server_url = "{self.zenml_server_url}"
-    zenml_api_key = ""
-    zenml_api_token = "{self.zenml_server_api_token}"
     zenml_stack_name = "{self.stack_name}"
     zenml_stack_deployment = "{self.deployment_type}"
 }}

zenml/stack_deployments/stack_deployment.py

@@ -219,16 +219,14 @@ class ZenMLCloudStackDeployment(BaseModel):
             if stack.labels.get("zenml:deployment") != self.deployment_type:
                 continue
 
-            artifact_store = stack.components[
-                StackComponentType.ARTIFACT_STORE
-            ][0]
+            orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
 
-            if not artifact_store.connector:
+            if not orchestrator.connector:
                 continue
 
             return DeployedStack(
                 stack=stack,
-                service_connector=artifact_store.connector,
+                service_connector=orchestrator.connector,
             )
 
         return None

zenml/steps/utils.py

@@ -442,6 +442,11 @@ def log_step_metadata(
             from within a step or if no pipeline name or ID is provided and
             the function is not called from within a step.
     """
+    logger.warning(
+        "The `log_step_metadata` function is deprecated and will soon be "
+        "removed. Please use `log_metadata` instead."
+    )
+
     step_context = None
     if not step_name:
         with contextlib.suppress(RuntimeError):

zenml/utils/metadata_utils.py

@@ -0,0 +1,335 @@
+#  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.
+"""Utility functions to handle metadata for ZenML entities."""
+
+import contextlib
+from typing import Dict, Optional, Union, overload
+from uuid import UUID
+
+from zenml.client import Client
+from zenml.enums import MetadataResourceTypes
+from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.steps.step_context import get_step_context
+
+logger = get_logger(__name__)
+
+
+@overload
+def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_name: str,
+    artifact_version: Optional[str] = None,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_name: str,
+    model_version: str,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_name: str,
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+def log_metadata(
+    metadata: Dict[str, MetadataType],
+    # Parameters to manually log metadata for steps and runs
+    step_id: Optional[UUID] = None,
+    step_name: Optional[str] = None,
+    run_id_name_or_prefix: Optional[Union[UUID, str]] = None,
+    # Parameters to manually log metadata for artifacts
+    artifact_version_id: Optional[UUID] = None,
+    artifact_name: Optional[str] = None,
+    artifact_version: Optional[str] = None,
+    # Parameters to manually log metadata for models
+    model_version_id: Optional[UUID] = None,
+    model_name: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> None:
+    """Logs metadata for various resource types in a generalized way.
+
+    Args:
+        metadata: The metadata to log.
+        step_id: The ID of the step.
+        step_name: The name of the step.
+        run_id_name_or_prefix: The id, name or prefix of the run
+        artifact_version_id: The ID of the artifact version
+        artifact_name: The name of the artifact.
+        artifact_version: The version of the artifact.
+        model_version_id: The ID of the model version.
+        model_name: The name of the model.
+        model_version: The version of the model
+
+    Raises:
+        ValueError: If no identifiers are provided and the function is not
+            called from within a step.
+    """
+    client = Client()
+
+    # If a step name is provided, we need a run_id_name_or_prefix and will log
+    # metadata for the steps pipeline and model accordingly.
+    if step_name is not None and run_id_name_or_prefix is not None:
+        run_model = client.get_pipeline_run(
+            name_id_or_prefix=run_id_name_or_prefix
+        )
+        step_model = run_model.steps[step_name]
+
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=run_model.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If a step is identified by id, fetch it directly through the client,
+    # follow a similar procedure and log metadata for its pipeline and model
+    # as well.
+    elif step_id is not None:
+        step_model = client.get_run_step(step_run_id=step_id)
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.pipeline_run_id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If a pipeline run id is identified, we need to log metadata to it and its
+    # model as well.
+    elif run_id_name_or_prefix is not None:
+        run_model = client.get_pipeline_run(
+            name_id_or_prefix=run_id_name_or_prefix
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=run_model.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        if run_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=run_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If the user provides a model name and version, we use to model abstraction
+    # to fetch the model version and attach the corresponding metadata to it.
+    elif model_name is not None and model_version is not None:
+        from zenml import Model
+
+        mv = Model(name=model_name, version=model_version)
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=mv.id,
+            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        )
+
+    # If the user provides a model version id, we use the client to fetch it and
+    # attach the metadata to it.
+    elif model_version_id is not None:
+        model_version_id = client.get_model_version(
+            model_version_name_or_number_or_id=model_version_id
+        ).id
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=model_version_id,
+            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        )
+
+    # If the user provides an artifact name, there are three possibilities. If
+    # an artifact version is also provided with the name, we use both to fetch
+    # the artifact version and use it to log the metadata. If no version is
+    # provided, if the function is called within a step we search the artifacts
+    # of the step if not we fetch the latest version and attach the metadata
+    # to the latest version.
+    elif artifact_name is not None:
+        if artifact_version:
+            artifact_version_model = client.get_artifact_version(
+                name_id_or_prefix=artifact_name, version=artifact_version
+            )
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=artifact_version_model.id,
+                resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+            )
+        else:
+            step_context = None
+            with contextlib.suppress(RuntimeError):
+                step_context = get_step_context()
+
+            if step_context:
+                step_context.add_output_metadata(
+                    metadata=metadata, output_name=artifact_name
+                )
+            else:
+                artifact_version_model = client.get_artifact_version(
+                    name_id_or_prefix=artifact_name
+                )
+                client.create_run_metadata(
+                    metadata=metadata,
+                    resource_id=artifact_version_model.id,
+                    resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+                )
+
+    # If the user directly provides an artifact_version_id, we use the client to
+    # fetch is and attach the metadata accordingly.
+    elif artifact_version_id is not None:
+        artifact_version_model = client.get_artifact_version(
+            name_id_or_prefix=artifact_version_id,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=artifact_version_model.id,
+            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+        )
+
+    # If every additional value is None, that means we are calling it bare bones
+    # and this call needs to happen during a step execution. We will use the
+    # step context to fetch the step, run and possibly the model version and
+    # attach the metadata accordingly.
+    elif all(
+        v is None
+        for v in [
+            step_id,
+            step_name,
+            run_id_name_or_prefix,
+            artifact_version_id,
+            artifact_name,
+            artifact_version,
+            model_version_id,
+            model_name,
+            model_version,
+        ]
+    ):
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            raise ValueError(
+                "You are calling 'log_metadata()' outside of a step execution. "
+                "If you would like to add metadata to a ZenML entity outside "
+                "of the step execution, please provide the required "
+                "identifiers."
+            )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_context.pipeline_run.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_context.step_run.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_context.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_context.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    else:
+        raise ValueError(
+            """
+            Unsupported way to call the `log_metadata`. Possible combinations "
+            include:
+            
+            # Inside a step
+            # Logs the metadata to the step, its run and possibly its model
+            log_metadata(metadata={})
+            
+            # Manually logging for a step
+            # Logs the metadata to the step, its run and possibly its model
+            log_metadata(metadata={}, step_name=..., run_id_name_or_prefix=...)
+            log_metadata(metadata={}, step_id=...)
+            
+            # Manually logging for a run 
+            # Logs the metadata to the run, possibly its model
+            log_metadata(metadata={}, run_id_name_or_prefix=...)
+            
+            # Manually logging for a model
+            log_metadata(metadata={}, model_name=..., model_version=...)
+            log_metadata(metadata={}, model_version_id=...)
+            
+            # Manually logging for an artifact
+            log_metadata(metadata={}, artifact_name=...)  # inside a step 
+            log_metadata(metadata={}, artifact_name=..., artifact_version=...)
+            log_metadata(metadata={}, artifact_version_id=...)
+            """
+        )

zenml/zen_server/auth.py

@@ -14,13 +14,13 @@
 """Authentication module for ZenML server."""
 
 from contextvars import ContextVar
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Callable, Optional, Union
 from urllib.parse import urlencode
 from uuid import UUID
 
 import requests
-from fastapi import Depends
+from fastapi import Depends, Response
 from fastapi.security import (
     HTTPBasic,
     HTTPBasicCredentials,
@@ -37,7 +37,7 @@ from zenml.constants import (
     LOGIN,
     VERSION_1,
 )
-from zenml.enums import AuthScheme, OAuthDeviceStatus
+from zenml.enums import AuthScheme, ExecutionStatus, OAuthDeviceStatus
 from zenml.exceptions import (
     AuthorizationException,
     CredentialsNotValid,
@@ -51,11 +51,13 @@ from zenml.models import (
     ExternalUserModel,
     OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
+    OAuthTokenResponse,
     UserAuthModel,
     UserRequest,
     UserResponse,
     UserUpdate,
 )
+from zenml.zen_server.cache import cache_result
 from zenml.zen_server.exceptions import http_exception_from_error
 from zenml.zen_server.jwt import JWTToken
 from zenml.zen_server.utils import server_config, zen_store
@@ -329,6 +331,148 @@ def authenticate_credentials(
                 ),
             )
 
+        if decoded_token.schedule_id:
+            # If the token contains a schedule ID, we need to check if the
+            # schedule still exists in the database. We use a cached version
+            # of the schedule active status to avoid unnecessary database
+            # queries.
+
+            @cache_result(expiry=30)
+            def get_schedule_active(schedule_id: UUID) -> Optional[bool]:
+                """Get the active status of a schedule.
+
+                Args:
+                    schedule_id: The schedule ID.
+
+                Returns:
+                    The schedule active status or None if the schedule does not
+                    exist.
+                """
+                try:
+                    schedule = zen_store().get_schedule(
+                        schedule_id, hydrate=False
+                    )
+                except KeyError:
+                    return False
+
+                return schedule.active
+
+            schedule_active = get_schedule_active(decoded_token.schedule_id)
+            if schedule_active is None:
+                error = (
+                    f"Authentication error: error retrieving token schedule "
+                    f"{decoded_token.schedule_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if not schedule_active:
+                error = (
+                    f"Authentication error: schedule {decoded_token.schedule_id} "
+                    "is not active"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.pipeline_run_id:
+            # If the token contains a pipeline run ID, we need to check if the
+            # pipeline run exists in the database and the pipeline run has
+            # not concluded. We use a cached version of the pipeline run status
+            # to avoid unnecessary database queries.
+
+            @cache_result(expiry=30)
+            def get_pipeline_run_status(
+                pipeline_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a pipeline run.
+
+                Args:
+                    pipeline_run_id: The pipeline run ID.
+
+                Returns:
+                    The pipeline run status or None if the pipeline run does not
+                    exist.
+                """
+                try:
+                    pipeline_run = zen_store().get_run(
+                        pipeline_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return pipeline_run.status
+
+            pipeline_run_status = get_pipeline_run_status(
+                decoded_token.pipeline_run_id
+            )
+            if pipeline_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token pipeline run "
+                    f"{decoded_token.pipeline_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if pipeline_run_status in [
+                ExecutionStatus.FAILED,
+                ExecutionStatus.COMPLETED,
+            ]:
+                error = (
+                    f"The execution of pipeline run "
+                    f"{decoded_token.pipeline_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.step_run_id:
+            # If the token contains a step run ID, we need to check if the
+            # step run exists in the database and the step run has not concluded.
+            # We use a cached version of the step run status to avoid unnecessary
+            # database queries.
+
+            @cache_result(expiry=30)
+            def get_step_run_status(
+                step_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a step run.
+
+                Args:
+                    step_run_id: The step run ID.
+
+                Returns:
+                    The step run status or None if the step run does not exist.
+                """
+                try:
+                    step_run = zen_store().get_run_step(
+                        step_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return step_run.status
+
+            step_run_status = get_step_run_status(decoded_token.step_run_id)
+            if step_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token step run "
+                    f"{decoded_token.step_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if step_run_status in [
+                ExecutionStatus.FAILED,
+                ExecutionStatus.COMPLETED,
+            ]:
+                error = (
+                    f"The execution of step run "
+                    f"{decoded_token.step_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
         auth_context = AuthContext(
             user=user_model,
             access_token=decoded_token,
@@ -660,6 +804,80 @@ def authenticate_api_key(
     return AuthContext(user=user_model, api_key=internal_api_key)
 
 
+def generate_access_token(
+    user_id: UUID,
+    response: Optional[Response] = None,
+    device: Optional[OAuthDeviceInternalResponse] = None,
+    api_key: Optional[APIKeyInternalResponse] = None,
+    expires_in: Optional[int] = None,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
+) -> OAuthTokenResponse:
+    """Generates an access token for the given user.
+
+    Args:
+        user_id: The ID of the user.
+        response: The FastAPI response object.
+        device: The device used for authentication.
+        api_key: The service account API key used for authentication.
+        expires_in: The number of seconds until the token expires.
+        schedule_id: The ID of the schedule to scope the token to.
+        pipeline_run_id: The ID of the pipeline run to scope the token to.
+        step_run_id: The ID of the step run to scope the token to.
+
+    Returns:
+        An authentication response with an access token.
+    """
+    config = server_config()
+
+    # If the expiration time is not supplied, the JWT tokens are set to expire
+    # according to the values configured in the server config. Device tokens are
+    # handled separately from regular user tokens.
+    expires: Optional[datetime] = None
+    if expires_in:
+        expires = datetime.utcnow() + timedelta(seconds=expires_in)
+    elif device:
+        # If a device was used for authentication, the token will expire
+        # at the same time as the device.
+        expires = device.expires
+        if expires:
+            expires_in = max(
+                int(expires.timestamp() - datetime.utcnow().timestamp()), 0
+            )
+    elif config.jwt_token_expire_minutes:
+        expires = datetime.utcnow() + timedelta(
+            minutes=config.jwt_token_expire_minutes
+        )
+        expires_in = config.jwt_token_expire_minutes * 60
+
+    access_token = JWTToken(
+        user_id=user_id,
+        device_id=device.id if device else None,
+        api_key_id=api_key.id if api_key else None,
+        schedule_id=schedule_id,
+        pipeline_run_id=pipeline_run_id,
+        step_run_id=step_run_id,
+    ).encode(expires=expires)
+
+    if not device and response:
+        # Also set the access token as an HTTP only cookie in the response
+        response.set_cookie(
+            key=config.get_auth_cookie_name(),
+            value=access_token,
+            httponly=True,
+            samesite="lax",
+            max_age=config.jwt_token_expire_minutes * 60
+            if config.jwt_token_expire_minutes
+            else None,
+            domain=config.auth_cookie_domain,
+        )
+
+    return OAuthTokenResponse(
+        access_token=access_token, expires_in=expires_in, token_type="bearer"
+    )
+
+
 def http_authentication(
     credentials: HTTPBasicCredentials = Depends(HTTPBasic()),
 ) -> AuthContext: