Add Active Stack as part of configuration YAML

src/zenml/config/pipeline_run_configuration.py

@@ -29,6 +29,7 @@ class PipelineRunConfiguration(
 ):
     """Class for pipeline run configurations."""
 
+    active_stack: Optional[str] = None
     run_name: Optional[str] = None
     enable_cache: Optional[bool] = None
     enable_artifact_metadata: Optional[bool] = None

src/zenml/new/pipelines/pipeline.py

@@ -73,6 +73,7 @@
 from zenml.new.pipelines.model_utils import NewModelRequest
 from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Stack
+from zenml.stack.utils import stack_context
 from zenml.steps import BaseStep
 from zenml.steps.entrypoint_function_utils import (
     StepArtifact,
@@ -535,7 +536,9 @@ def build(
         Returns:
             The build output.
         """
-        with track_handler(event=AnalyticsEvent.BUILD_PIPELINE):
+        with track_handler(
+            event=AnalyticsEvent.BUILD_PIPELINE
+        ), stack_context():
             self._prepare_if_possible()
             deployment, pipeline_spec, _, _ = self._compile(
                 config_path=config_path,
@@ -621,7 +624,9 @@ def _run(
 
         logger.info(f"Initiating a new run for the pipeline: `{self.name}`.")
 
-        with track_handler(AnalyticsEvent.RUN_PIPELINE) as analytics_handler:
+        with track_handler(
+            AnalyticsEvent.RUN_PIPELINE
+        ) as analytics_handler, stack_context():
             deployment, pipeline_spec, schedule, build = self._compile(
                 config_path=config_path,
                 run_name=run_name,
@@ -1152,6 +1157,8 @@ def _compile(
         # Update with the values in code so they take precedence
         run_config = pydantic_utils.update_model(run_config, update=update)
 
+        self._update_stack_from_config(run_config)
+
         deployment, pipeline_spec = Compiler().compile(
             pipeline=self,
             stack=Client().active_stack,
@@ -1564,3 +1571,14 @@ def _prepare_if_possible(self) -> None:
                 )
             else:
                 self.prepare()
+
+    def _update_stack_from_config(
+        self, run_configuration: PipelineRunConfiguration
+    ) -> None:
+        """Activate the stack from the pipeline run configuration if one is given.
+
+        Args:
+            run_configuration: The run configuration for this pipeline.
+        """
+        if run_configuration.active_stack is not None:
+            Client().activate_stack(run_configuration.active_stack)

src/zenml/stack/utils.py

@@ -13,7 +13,8 @@
 #  permissions and limitations under the License.
 """Util functions for handling stacks, components, and flavors."""
 
-from typing import Any, Dict, Optional
+from types import TracebackType
+from typing import Any, Dict, Optional, Type
 
 from zenml.client import Client
 from zenml.enums import StackComponentType, StoreType
@@ -139,3 +140,25 @@ def get_flavor_by_name_and_type_from_zen_store(
             f"'{component_type}' exists."
         )
     return flavors[0]
+
+
+class stack_context:
+    """Context handler to reset the original active stack."""
+
+    def __init__(self) -> None:
+        """Constructor for stack_context saves active stack."""
+        self._default_stack = Client().active_stack
+
+    def __enter__(self) -> None:
+        """Enters in the stack context."""
+        pass
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        """Restores the original active stack."""
+        if self._default_stack.id != Client().active_stack:
+            Client().activate_stack(self._default_stack.id)

tests/integration/functional/cli/test_pipeline.py

@@ -368,6 +368,39 @@ def test_pipeline_run_with_config_file(clean_client: "Client", tmp_path):
     assert runs[0].name == "custom_run_name"
 
 
+def test_pipeline_run_with_different_stack_in_config_file(
+    clean_client: "Client", tmp_path
+):
+    """Tests that the run command works with a run config file with an active stack defined."""
+    runner = CliRunner()
+    run_command = cli.commands["pipeline"].commands["run"]
+
+    pipeline_id = pipeline_instance.register().id
+
+    components = {
+        key: components[0].id
+        for key, components in Client().active_stack_model.components.items()
+    }
+    new_stack = Client().create_stack(name="new", components=components)
+
+    config_path = tmp_path / "config.yaml"
+    run_config = PipelineRunConfiguration(
+        run_name="custom_run_name", active_stack=str(new_stack.id)
+    )
+    config_path.write_text(run_config.yaml())
+
+    result = runner.invoke(
+        run_command, [pipeline_instance.name, "--config", str(config_path)]
+    )
+    assert result.exit_code == 0
+
+    runs = Client().list_pipeline_runs(pipeline_id=pipeline_id)
+    assert len(runs) == 1
+    assert runs[0].name == "custom_run_name"
+    assert runs[0].stack.id == new_stack.id
+    assert Client().active_stack.id != new_stack.id
+
+
 def test_pipeline_run_with_different_stack(clean_client: "Client"):
     """Tests that the run command works with a different stack."""
     runner = CliRunner()