Pending Changes for Active Stack

Author: preyasshah9

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[str] = None
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

src/zenml/pipelines/pipeline_definition.py

@@ -92,6 +92,7 @@
     source_utils,
     yaml_utils,
 )
+from zenml.utils.stack_utils import temporary_active_stack
 from zenml.utils.string_utils import format_name_template
 
 if TYPE_CHECKING:
@@ -565,19 +566,26 @@ def build(
             if settings:
                 compile_args["settings"] = settings
 
-            deployment, _, _ = self._compile(**compile_args)
-            pipeline_id = self._register().id
-
-            local_repo = code_repository_utils.find_active_code_repository()
-            code_repository = build_utils.verify_local_repository_context(
-                deployment=deployment, local_repo_context=local_repo
+            _from_config_file = self._parse_config_file(
+                config_path=config_path,
+                matcher=list(PipelineRunConfiguration.model_fields.keys()),
             )
+            run_config = PipelineRunConfiguration(**_from_config_file)
 
-            return build_utils.create_pipeline_build(
-                deployment=deployment,
-                pipeline_id=pipeline_id,
-                code_repository=code_repository,
-            )
+            with temporary_active_stack(stack_name_or_id=run_config.stack):
+                deployment, _, _ = self._compile(**compile_args)
+                pipeline_id = self._register().id
+
+                local_repo = code_repository_utils.find_active_code_repository()
+                code_repository = build_utils.verify_local_repository_context(
+                    deployment=deployment, local_repo_context=local_repo
+                )
+
+                return build_utils.create_pipeline_build(
+                    deployment=deployment,
+                    pipeline_id=pipeline_id,
+                    code_repository=code_repository,
+                )
 
     def _create_deployment(
         self,
@@ -800,35 +808,42 @@ def _run(
         logger.info(f"Initiating a new run for the pipeline: `{self.name}`.")
 
         with track_handler(AnalyticsEvent.RUN_PIPELINE) as analytics_handler:
-            stack = Client().active_stack
-            deployment = self._create_deployment(**self._run_args)
+            _from_config_file = self._parse_config_file(
+                config_path=self._run_args.get("config_path"),
+                matcher=list(PipelineRunConfiguration.model_fields.keys()),
+            )
+            run_config = PipelineRunConfiguration(**_from_config_file)
 
-            self.log_pipeline_deployment_metadata(deployment)
-            run = create_placeholder_run(deployment=deployment)
+            with temporary_active_stack(stack_name_or_id=run_config.stack):
+                stack = Client().active_stack
+                deployment = self._create_deployment(**self._run_args)
 
-            analytics_handler.metadata = self._get_pipeline_analytics_metadata(
-                deployment=deployment,
-                stack=stack,
-                run_id=run.id if run else None,
-            )
+                self.log_pipeline_deployment_metadata(deployment)
+                run = create_placeholder_run(deployment=deployment)
 
-            if run:
-                run_url = dashboard_utils.get_run_url(run)
-                if run_url:
-                    logger.info(f"Dashboard URL for Pipeline Run: {run_url}")
-                else:
-                    logger.info(
-                        "You can visualize your pipeline runs in the `ZenML "
-                        "Dashboard`. In order to try it locally, please run "
-                        "`zenml login --local`."
-                    )
+                analytics_handler.metadata = self._get_pipeline_analytics_metadata(
+                    deployment=deployment,
+                    stack=stack,
+                    run_id=run.id if run else None,
+                )
 
-            deploy_pipeline(
-                deployment=deployment, stack=stack, placeholder_run=run
-            )
-            if run:
-                return Client().get_pipeline_run(run.id)
-            return None
+                if run:
+                    run_url = dashboard_utils.get_run_url(run)
+                    if run_url:
+                        logger.info(f"Dashboard URL for Pipeline Run: {run_url}")
+                    else:
+                        logger.info(
+                            "You can visualize your pipeline runs in the `ZenML "
+                            "Dashboard`. In order to try it locally, please run "
+                            "`zenml login --local`."
+                        )
+
+                deploy_pipeline(
+                    deployment=deployment, stack=stack, placeholder_run=run
+                )
+                if run:
+                    return Client().get_pipeline_run(run.id)
+                return None
 
     @staticmethod
     def log_pipeline_deployment_metadata(

src/zenml/pipelines/run_utils.py

@@ -1,3 +1,16 @@
+#  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.
 """Utility functions for running pipelines."""
 
 import time
@@ -248,6 +261,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 = [