Add Active Stack as part of configuration YAML

docs/book/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines.md

@@ -132,6 +132,7 @@ python run.py
 <summary>An example of a generated YAML configuration template</summary>
 
 ```yaml
+active_stack: Optional[str]
 build: Union[PipelineBuildBase, UUID, NoneType]
 enable_artifact_metadata: Optional[bool]
 enable_artifact_visualization: Optional[bool]
@@ -330,6 +331,12 @@ These are boolean flags for various configurations:
 * `enable_cache`: Utilize [caching](../../starter-guide/cache-previous-executions.md) or not.
 * `enable_step_logs`: Enable tracking [step logs](managing-steps.md#enable-or-disable-logs-storing).
 
+### `active_stack` name or ID
+
+The name or the UUID of the `active stack` to use for this
+pipeline. If specified, the active stack is set for the duration of the pipeline execution and restored upon 
+completion. If not specified, the current active stack is used.
+
 ### `build` ID
 
 The UUID of the [`build`](../infrastructure-management/containerize-your-pipeline.md) to use for this pipeline. If specified, Docker image building is skipped for remote orchestrators, and the Docker image specified in this build is used.

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

@@ -75,6 +75,7 @@
     prepare_model_versions,
 )
 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,
@@ -536,7 +537,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,
@@ -619,7 +622,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,
@@ -1014,6 +1019,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,
@@ -1437,3 +1444,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,33 @@ 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,
+        exception_type: Optional[Type[BaseException]],
+        exception_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        """Get a stack component flavor by name and type from a ZenStore.
+
+        Args:
+            exception_type: Type of the exception that was raised.
+                None if no exception.
+            exception_value: Type of exception that was raised.
+                e.g., divide_by_zero error. None if no exception.
+            traceback: Traceback report. None if no exception.
+        """
+        if self._default_stack.id != Client().active_stack.id:
+            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()