Pending Changes for Active Stack

Author: Preyas Shah

src/zenml/cli/pipeline.py

@@ -35,6 +35,7 @@
 )
 from zenml.pipelines.pipeline_definition import Pipeline
 from zenml.utils import source_utils, uuid_utils
+from zenml.utils.stack_utils import temporary_active_stack
 from zenml.utils.yaml_utils import write_yaml
 
 logger = get_logger(__name__)
@@ -190,7 +191,7 @@ def build_pipeline(
             "your source code root."
         )
 
-    with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+    with temporary_active_stack(stack_name_or_id=stack_name_or_id):
         pipeline_instance = _import_pipeline(source=source)
 
         pipeline_instance = pipeline_instance.with_options(
@@ -276,7 +277,7 @@ def run_pipeline(
             "your source code root."
         )
 
-    with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+    with temporary_active_stack(stack_name_or_id=stack_name_or_id):
         pipeline_instance = _import_pipeline(source=source)
 
         build: Union[str, PipelineBuildBase, None] = None
@@ -353,7 +354,7 @@ def create_run_template(
             "init` at your source code root."
         )
 
-    with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+    with temporary_active_stack(stack_name_or_id=stack_name_or_id):
         pipeline_instance = _import_pipeline(source=source)
 
         pipeline_instance = pipeline_instance.with_options(

src/zenml/cli/utils.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Utility functions for the CLI."""
 
-import contextlib
 import json
 import os
 import platform
@@ -27,7 +26,6 @@
     Any,
     Callable,
     Dict,
-    Iterator,
     List,
     NoReturn,
     Optional,
@@ -83,8 +81,6 @@
 from zenml.utils.time_utils import expires_in
 
 if TYPE_CHECKING:
-    from uuid import UUID
-
     from rich.text import Text
 
     from zenml.enums import ExecutionStatus
@@ -2468,33 +2464,6 @@ def wrapper(function: F) -> F:
     return inner_decorator
 
 
-@contextlib.contextmanager
-def temporary_active_stack(
-    stack_name_or_id: Union["UUID", str, None] = None,
-) -> Iterator["Stack"]:
-    """Contextmanager to temporarily activate a stack.
-
-    Args:
-        stack_name_or_id: The name or ID of the stack to activate. If not given,
-            this contextmanager will not do anything.
-
-    Yields:
-        The active stack.
-    """
-    from zenml.client import Client
-
-    try:
-        if stack_name_or_id:
-            old_stack_id = Client().active_stack_model.id
-            Client().activate_stack(stack_name_or_id)
-        else:
-            old_stack_id = None
-        yield Client().active_stack
-    finally:
-        if old_stack_id:
-            Client().activate_stack(old_stack_id)
-
-
 def get_package_information(
     package_names: Optional[List[str]] = None,
 ) -> Dict[str, str]:

src/zenml/config/pipeline_run_configuration.py

@@ -26,6 +26,7 @@
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
 from zenml.models import PipelineBuildBase
+from zenml.stack.stack import Stack
 from zenml.utils import pydantic_utils
 
 
@@ -46,6 +47,7 @@ class PipelineRunConfiguration(
     steps: Dict[str, StepConfigurationUpdate] = {}
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     tags: Optional[List[str]] = None
+    stack: Optional[Stack] = None
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

src/zenml/pipelines/run_utils.py

@@ -248,6 +248,11 @@ def validate_run_config_is_runnable_from_server(
             "Can't set schedule when running pipeline via Rest API."
         )
 
+    if run_configuration.stack:
+        raise ValueError(
+            "Can't switch stack when running pipeline via Rest API."
+        )
+
     if run_configuration.settings.get("docker"):
         raise ValueError(
             "Can't set DockerSettings when running pipeline via Rest API."

src/zenml/utils/stack_utils.py

@@ -0,0 +1,53 @@
+#  Copyright (c) ZenML GmbH 2023. 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.
+"""Utilities for stack."""
+
+import contextlib
+from typing import (
+    TYPE_CHECKING,
+    Iterator,
+    Union,
+)
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from zenml.stack import Stack
+
+
+@contextlib.contextmanager
+def temporary_active_stack(
+    stack_name_or_id: Union["UUID", str, None] = None,
+) -> Iterator["Stack"]:
+    """Contextmanager to temporarily activate a stack.
+
+    Args:
+        stack_name_or_id: The name or ID of the stack to activate. If not given,
+            this contextmanager will not do anything.
+
+    Yields:
+        The active stack.
+    """
+    from zenml.client import Client
+
+    try:
+        if stack_name_or_id:
+            old_stack_id = Client().active_stack_model.id
+            Client().activate_stack(stack_name_or_id)
+        else:
+            old_stack_id = None
+        yield Client().active_stack
+    finally:
+        if old_stack_id:
+            Client().activate_stack(old_stack_id)

tests/integration/functional/cli/utils.py

@@ -18,7 +18,6 @@
 from zenml.cli import cli
 from zenml.cli.utils import (
     parse_name_and_extra_arguments,
-    temporary_active_stack,
 )
 from zenml.client import Client
 from zenml.models import (
@@ -27,6 +26,7 @@
     UserResponse,
     WorkspaceResponse,
 )
+from zenml.utils.stack_utils import temporary_active_stack
 from zenml.utils.string_utils import random_str
 
 SAMPLE_CUSTOM_ARGUMENTS = [