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

zenml/zen_server/template_execution/utils.py

@@ -42,7 +42,7 @@ from zenml.pipelines.run_utils import (
 )
 from zenml.stack.flavor import Flavor
 from zenml.utils import dict_utils, requirements_utils, settings_utils
-from zenml.zen_server.auth import AuthContext
+from zenml.zen_server.auth import AuthContext, generate_access_token
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
 )
@@ -111,17 +111,6 @@ def run_template(
 
     new_deployment = zen_store().create_deployment(deployment_request)
 
-    if auth_context.access_token:
-        token = auth_context.access_token
-        token.pipeline_id = deployment_request.pipeline
-
-        # We create a non-expiring token to make sure its active for the entire
-        # duration of the pipeline run
-        api_token = token.encode(expires=None)
-    else:
-        assert auth_context.encoded_access_token
-        api_token = auth_context.encoded_access_token
-
     server_url = server_config().server_url
     if not server_url:
         raise RuntimeError(
@@ -130,6 +119,18 @@ def run_template(
     assert build.zenml_version
     zenml_version = build.zenml_version
 
+    placeholder_run = create_placeholder_run(deployment=new_deployment)
+    assert placeholder_run
+
+    # We create an API token scoped to the pipeline run
+    api_token = generate_access_token(
+        user_id=auth_context.user.id,
+        pipeline_run_id=placeholder_run.id,
+        # Keep the original API key or device scopes, if any
+        api_key=auth_context.api_key,
+        device=auth_context.device,
+    ).access_token
+
     environment = {
         ENV_ZENML_ACTIVE_WORKSPACE_ID: str(new_deployment.workspace.id),
         ENV_ZENML_ACTIVE_STACK_ID: str(stack.id),
@@ -145,9 +146,6 @@ def run_template(
         deployment_id=new_deployment.id
     )
 
-    placeholder_run = create_placeholder_run(deployment=new_deployment)
-    assert placeholder_run
-
     def _task() -> None:
         pypi_requirements, apt_packages = (
             requirements_utils.get_requirements_for_stack(stack=stack)
@@ -350,7 +348,7 @@ def deployment_request_from_template(
     )
 
     step_config_dict_base = pipeline_configuration.model_dump(
-        exclude={"name", "parameters"}
+        exclude={"name", "parameters", "tags"}
     )
     steps = {}
     for invocation_id, step in deployment.step_configurations.items():

zenml/zen_server/utils.py

@@ -43,6 +43,7 @@ from zenml.enums import StoreType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
+from zenml.zen_server.cache import MemoryCache
 from zenml.zen_server.deploy.deployment import (
     LocalServerDeployment,
 )
@@ -67,6 +68,7 @@ _rbac: Optional[RBACInterface] = None
 _feature_gate: Optional[FeatureGateInterface] = None
 _workload_manager: Optional[WorkloadManagerInterface] = None
 _plugin_flavor_registry: Optional[PluginFlavorRegistry] = None
+_memcache: Optional[MemoryCache] = None
 
 
 def zen_store() -> "SqlZenStore":
@@ -222,6 +224,31 @@ def initialize_zen_store() -> None:
     _zen_store = zen_store_
 
 
+def initialize_memcache(max_capacity: int, default_expiry: int) -> None:
+    """Initialize the memory cache.
+
+    Args:
+        max_capacity: The maximum capacity of the cache.
+        default_expiry: The default expiry time in seconds.
+    """
+    global _memcache
+    _memcache = MemoryCache(max_capacity, default_expiry)
+
+
+def memcache() -> MemoryCache:
+    """Return the memory cache.
+
+    Returns:
+        The memory cache.
+
+    Raises:
+        RuntimeError: If the memory cache is not initialized.
+    """
+    if _memcache is None:
+        raise RuntimeError("Memory cache not initialized")
+    return _memcache
+
+
 _server_config: Optional[ServerConfiguration] = None
 
 

zenml/zen_server/zen_server_api.py

@@ -95,6 +95,7 @@ from zenml.zen_server.secure_headers import (
 )
 from zenml.zen_server.utils import (
     initialize_feature_gate,
+    initialize_memcache,
     initialize_plugins,
     initialize_rbac,
     initialize_workload_manager,
@@ -335,9 +336,10 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
 @app.on_event("startup")
 def initialize() -> None:
     """Initialize the ZenML server."""
+    cfg = server_config()
     # Set the maximum number of worker threads
     to_thread.current_default_thread_limiter().total_tokens = (
-        server_config().thread_pool_size
+        cfg.thread_pool_size
     )
     # IMPORTANT: these need to be run before the fastapi app starts, to avoid
     # race conditions
@@ -347,6 +349,7 @@ def initialize() -> None:
     initialize_workload_manager()
     initialize_plugins()
     initialize_secure_headers()
+    initialize_memcache(cfg.memcache_max_capacity, cfg.memcache_default_expiry)
 
 
 DASHBOARD_REDIRECT_URL = None

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

@@ -0,0 +1,23 @@
+"""Release [0.70.0].
+
+Revision ID: 0.70.0
+Revises: 904464ea4041
+Create Date: 2024-11-12 12:20:05.537646
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.70.0"
+down_revision = "904464ea4041"
+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/rest_zen_store.py

@@ -110,6 +110,7 @@ from zenml.constants import (
     WORKSPACES,
 )
 from zenml.enums import (
+    APITokenType,
     OAuthGrantTypes,
     StackDeploymentProvider,
     StoreType,
@@ -3872,17 +3873,16 @@ class RestZenStore(BaseZenStore):
 
     def get_api_token(
         self,
-        pipeline_id: Optional[UUID] = None,
         schedule_id: Optional[UUID] = None,
-        expires_minutes: Optional[int] = None,
+        pipeline_run_id: Optional[UUID] = None,
+        step_run_id: Optional[UUID] = None,
     ) -> str:
         """Get an API token for a workload.
 
         Args:
-            pipeline_id: The ID of the pipeline to get a token for.
             schedule_id: The ID of the schedule to get a token for.
-            expires_minutes: The number of minutes for which the token should
-                be valid. If not provided, the token will be valid indefinitely.
+            pipeline_run_id: The ID of the pipeline run to get a token for.
+            step_run_id: The ID of the step run to get a token for.
 
         Returns:
             The API token.
@@ -3890,13 +3890,16 @@ class RestZenStore(BaseZenStore):
         Raises:
             ValueError: if the server response is not valid.
         """
-        params: Dict[str, Any] = {}
-        if pipeline_id:
-            params["pipeline_id"] = pipeline_id
+        params: Dict[str, Any] = {
+            # Python clients may only request workload tokens.
+            "token_type": APITokenType.WORKLOAD.value,
+        }
         if schedule_id:
             params["schedule_id"] = schedule_id
-        if expires_minutes:
-            params["expires_minutes"] = expires_minutes
+        if pipeline_run_id:
+            params["pipeline_run_id"] = pipeline_run_id
+        if step_run_id:
+            params["step_run_id"] = step_run_id
         response_body = self.get(API_TOKEN, params=params)
         if not isinstance(response_body, str):
             raise ValueError(

{zenml_nightly-0.68.1.dev20241112.dist-info → zenml_nightly-0.70.0.dev20241114.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.68.1.dev20241112
+Version: 0.70.0.dev20241114
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -474,7 +474,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.68.1 is out. Check out the release notes
+    🎉 Version 0.70.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.