zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.70.0.dev20241116__py3-none-any.whl

zenml/orchestrators/utils.py

@@ -28,16 +28,16 @@ from zenml.constants import (
     ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
     ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
-    PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
 from zenml.enums import AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
 from zenml.stack import StackComponent
 from zenml.utils.string_utils import format_name_template
 
+logger = get_logger(__name__)
+
 if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-    from zenml.models import PipelineDeploymentResponse
 
 
 def get_orchestrator_run_name(pipeline_name: str) -> str:
@@ -80,16 +80,23 @@ def is_setting_enabled(
 
 
 def get_config_environment_vars(
-    deployment: Optional["PipelineDeploymentResponse"] = None,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
 ) -> Dict[str, str]:
     """Gets environment variables to set for mirroring the active config.
 
-    If a pipeline deployment is given, the environment variables will be set to
-    include a newly generated API token valid for the duration of the pipeline
-    run instead of the API token from the global config.
+    If a schedule ID, pipeline run ID or step run ID is given, and the current
+    client is not authenticated to a server with an API key, the environment
+    variables will be updated to include a newly generated workload API token
+    that will be valid for the duration of the schedule, pipeline run, or step
+    run instead of the current API token used to authenticate the client.
 
     Args:
-        deployment: Optional deployment to use for the environment variables.
+        schedule_id: Optional schedule ID to use to generate a new API token.
+        pipeline_run_id: Optional pipeline run ID to use to generate a new API
+            token.
+        step_run_id: Optional step run ID to use to generate a new API token.
 
     Returns:
         Environment variable dict.
@@ -107,34 +114,46 @@ def get_config_environment_vars(
     ):
         credentials_store = get_credentials_store()
         url = global_config.store_configuration.url
-        api_key = credentials_store.get_api_key(url)
         api_token = credentials_store.get_token(url, allow_expired=False)
-        if api_key:
-            environment_vars[ENV_ZENML_STORE_PREFIX + "API_KEY"] = api_key
-        elif deployment:
-            # When connected to an authenticated ZenML server, if a pipeline
-            # deployment is supplied, we need to fetch an API token that will be
-            # valid for the duration of the pipeline run.
+        if schedule_id or pipeline_run_id or step_run_id:
+            # When connected to an authenticated ZenML server, if a schedule ID,
+            # pipeline run ID or step run ID is supplied, we need to fetch a new
+            # workload API token scoped to the schedule, pipeline run or step
+            # run.
             assert isinstance(global_config.zen_store, RestZenStore)
-            pipeline_id: Optional[UUID] = None
-            if deployment.pipeline:
-                pipeline_id = deployment.pipeline.id
-            schedule_id: Optional[UUID] = None
-            expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
-            if deployment.schedule:
-                schedule_id = deployment.schedule.id
-                # If a schedule is given, this is a long running pipeline that
-                # should not have an API token that expires.
-                expires_minutes = None
+
+            # If only a schedule is given, the pipeline run credentials will
+            # be valid for the entire duration of the schedule.
+            api_key = credentials_store.get_api_key(url)
+            if not api_key and not pipeline_run_id and not step_run_id:
+                logger.warning(
+                    "An API token without an expiration time will be generated "
+                    "and used to run this pipeline on a schedule. This is very "
+                    "insecure because the API token will be valid for the "
+                    "entire lifetime of the schedule and can be used to access "
+                    "your user account if accidentally leaked. When deploying "
+                    "a pipeline on a schedule, it is strongly advised to use a "
+                    "service account API key to authenticate to the ZenML "
+                    "server instead of your regular user account. For more "
+                    "information, see "
+                    "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+                )
+
+            # The schedule, pipeline run or step run credentials are scoped to
+            # the schedule, pipeline run or step run and will only be valid for
+            # the duration of the schedule/pipeline run/step run.
             new_api_token = global_config.zen_store.get_api_token(
-                pipeline_id=pipeline_id,
                 schedule_id=schedule_id,
-                expires_minutes=expires_minutes,
+                pipeline_run_id=pipeline_run_id,
+                step_run_id=step_run_id,
             )
+
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
                 new_api_token
             )
         elif api_token:
+            # For all other cases, the pipeline run environment is configured
+            # with the current access token.
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
                 api_token.access_token
             )

zenml/stack_deployments/aws_stack_deployment.py

@@ -71,7 +71,8 @@ of any potential costs:
 
 - An S3 bucket registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/s3).
 - An ECR repository registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/aws).
-- Sagemaker registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/sagemaker).
+- Sagemaker registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/sagemaker)
+as well as a [ZenML step operator](https://docs.zenml.io/stack-components/step-operators/sagemaker).
 - An IAM user and IAM role with the minimum necessary permissions to access the
 above resources.
 - An AWS access key used to give access to ZenML to connect to the above
@@ -257,13 +258,29 @@ console.
 
         config: Optional[str] = None
         if self.deployment_type == STACK_DEPLOYMENT_TERRAFORM:
-            config = f"""module "zenml_stack" {{
-    source  = "zenml-io/zenml-stack/aws"
+            config = f"""terraform {{
+    required_providers {{
+        aws = {{
+            source  = "hashicorp/aws"
+        }}
+        zenml = {{
+            source = "zenml-io/zenml"
+        }}
+    }}
+}}
 
+provider "aws" {{
     region = "{self.location or "eu-central-1"}"
-    zenml_server_url = "{self.zenml_server_url}"
-    zenml_api_key = ""
-    zenml_api_token = "{self.zenml_server_api_token}"
+}}
+
+provider "zenml" {{
+    server_url = "{self.zenml_server_url}"
+    api_token = "{self.zenml_server_api_token}"
+}}
+
+module "zenml_stack" {{
+    source  = "zenml-io/zenml-stack/aws"
+
     zenml_stack_name = "{self.stack_name}"
     zenml_stack_deployment = "{self.deployment_type}"
 }}

zenml/stack_deployments/azure_stack_deployment.py

@@ -259,14 +259,37 @@ ZenML's access to your Azure subscription.
             The configuration or script to deploy the ZenML stack to the
             specified cloud provider.
         """
-        config = f"""module "zenml_stack" {{
+        config = f"""terraform {{
+    required_providers {{
+        azurerm = {{
+            source  = "hashicorp/azurerm"
+        }}
+        azuread = {{
+            source  = "hashicorp/azuread"
+        }}
+        zenml = {{
+            source = "zenml-io/zenml"
+        }}
+    }}
+}}
+
+provider "azurerm" {{
+    features {{
+        resource_group {{
+            prevent_deletion_if_contains_resources = false
+        }}
+    }}
+}}
+
+provider "zenml" {{
+    server_url = "{self.zenml_server_url}"
+    api_token = "{self.zenml_server_api_token}"
+}}
+
+module "zenml_stack" {{
     source  = "zenml-io/zenml-stack/azure"
 
     location = "{self.location or "eastus"}"
-    orchestrator = "azureml"
-    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/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/base_step.py

@@ -327,12 +327,12 @@ class BaseStep:
             The artifacts, external artifacts, model version artifacts/metadata and parameters for the step.
         """
         from zenml.artifacts.external_artifact import ExternalArtifact
+        from zenml.metadata.lazy_load import LazyRunMetadataResponse
         from zenml.model.lazy_load import ModelVersionDataLazyLoader
         from zenml.models.v2.core.artifact_version import (
             ArtifactVersionResponse,
             LazyArtifactVersionResponse,
         )
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         signature = inspect.signature(self.entrypoint, follow_wrapped=True)
 

zenml/steps/entrypoint_function_utils.py

@@ -32,6 +32,7 @@ from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.lazy_load import LazyRunMetadataResponse
 from zenml.steps.utils import (
     OutputSignature,
     parse_return_type_annotations,
@@ -136,10 +137,7 @@ class EntrypointFunctionDefinition(NamedTuple):
             UnmaterializedArtifact,
         )
         from zenml.client_lazy_loader import ClientLazyLoader
-        from zenml.models import (
-            ArtifactVersionResponse,
-            RunMetadataResponse,
-        )
+        from zenml.models import ArtifactVersionResponse
 
         if key not in self.inputs:
             raise KeyError(
@@ -154,8 +152,8 @@ class EntrypointFunctionDefinition(NamedTuple):
                 StepArtifact,
                 ExternalArtifact,
                 ArtifactVersionResponse,
-                RunMetadataResponse,
                 ClientLazyLoader,
+                LazyRunMetadataResponse,
             ),
         ):
             # If we were to do any type validation for artifacts here, we

zenml/steps/step_context.py

@@ -35,11 +35,12 @@ if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType
     from zenml.model.model import Model
     from zenml.models import (
-        ArtifactVersionResponse,
         PipelineResponse,
         PipelineRunResponse,
         StepRunResponse,
     )
+    from zenml.models.v2.core.step_run import StepRunInputResponse
+
 
 logger = get_logger(__name__)
 
@@ -191,7 +192,7 @@ class StepContext(metaclass=SingletonMetaClass):
         return self.model_version.to_model_class()
 
     @property
-    def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def inputs(self) -> Dict[str, "StepRunInputResponse"]:
         """Returns the input artifacts of the current step.
 
         Returns:

zenml/steps/utils.py

@@ -26,7 +26,11 @@ from typing_extensions import Annotated
 
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
-from zenml.enums import ExecutionStatus, MetadataResourceTypes
+from zenml.enums import (
+    ArtifactSaveType,
+    ExecutionStatus,
+    MetadataResourceTypes,
+)
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
@@ -438,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):
@@ -547,8 +556,10 @@ def run_as_single_step_pipeline(
     # 4. Load output artifacts
     step_run = next(iter(run.steps.values()))
     outputs = [
-        step_run.outputs[output_name].load()
+        artifact_version.load()
         for output_name in step_run.config.outputs.keys()
+        for artifact_version in step_run.outputs[output_name]
+        if artifact_version.save_type == ArtifactSaveType.STEP_OUTPUT
     ]
 
     if len(outputs) == 0: