Split futures

Author: schustmi

src/zenml/execution/pipeline/dynamic/outputs.py

@@ -14,19 +14,13 @@
 """Dynamic pipeline execution outputs."""
 
 from concurrent.futures import Future
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Tuple,
-    Union,
-)
+from typing import Any, List, Tuple, Union
 
 from zenml.logger import get_logger
 from zenml.models import (
     ArtifactVersionResponse,
 )
 
-
 logger = get_logger(__name__)
 
 
@@ -37,18 +31,61 @@ class OutputArtifact(ArtifactVersionResponse):
     step_name: str
 
 
-StepRunOutputs = Union[
-    None, OutputArtifact, Tuple[OutputArtifact, ...]
-]
+StepRunOutputs = Union[None, OutputArtifact, Tuple[OutputArtifact, ...]]
+
+
+class ArtifactFuture:
+    """Future for a step run output artifact."""
+
+    def __init__(
+        self, wrapped: Future[StepRunOutputs], invocation_id: str, index: int
+    ) -> None:
+        """Initialize the future.
+
+        Args:
+            wrapped: The wrapped future object.
+            invocation_id: The invocation ID of the step run.
+        """
+        self._wrapped = wrapped
+        self._invocation_id = invocation_id
+        self._index = index
+
+    def _wait(self) -> None:
+        self._wrapped.result()
+
+    def result(self) -> OutputArtifact:
+        """Get the step run output artifact.
+
+        Returns:
+            The step run output artifact.
+        """
+        result = self._wrapped.result()
+        if isinstance(result, OutputArtifact):
+            return result
+        elif isinstance(result, tuple):
+            return result[self._index]
+        else:
+            raise RuntimeError(
+                f"Step {self._invocation_id} returned an invalid output: {result}"
+            )
+
+    def load(self) -> Any:
+        """Load the step run output artifact data.
+
+        Returns:
+            The step run output artifact data.
+        """
+        return self.result().load()
 
 
-# TODO: maybe one future per artifact? But for a step that doesn't return anything, the user wouldn't have a future to wait for.
-# Or that step returns a future that returns None? Would be similar to a python function.
 class StepRunOutputsFuture:
     """Future for a step run output."""
 
     def __init__(
-        self, wrapped: Future[StepRunOutputs], invocation_id: str
+        self,
+        wrapped: Future[StepRunOutputs],
+        invocation_id: str,
+        output_keys: List[str],
     ) -> None:
         """Initialize the future.
 
@@ -58,12 +95,12 @@ def __init__(
         """
         self._wrapped = wrapped
         self._invocation_id = invocation_id
+        self._output_keys = output_keys
 
-    def wait(self) -> None:
-        """Wait for the future to complete."""
+    def _wait(self) -> None:
         self._wrapped.result()
 
-    def result(self) -> StepRunOutputs:
+    def artifacts(self) -> StepRunOutputs:
         """Get the step run output artifacts.
 
         Returns:
@@ -80,7 +117,7 @@ def load(self) -> Any:
         Returns:
             The step run output artifact data.
         """
-        result = self.result()
+        result = self.artifacts()
 
         if result is None:
             return None
@@ -90,3 +127,39 @@ def load(self) -> Any:
             return tuple(item.load() for item in result)
         else:
             raise ValueError(f"Invalid step run output: {result}")
+
+    def __getitem__(self, key: Union[str, int]) -> ArtifactFuture:
+        if isinstance(key, str):
+            index = self._output_keys.index(key)
+        elif isinstance(key, int):
+            index = key
+        else:
+            raise ValueError(f"Invalid key type: {type(key)}")
+
+        if index > len(self._output_keys):
+            raise IndexError(f"Index out of range: {index}")
+
+        return ArtifactFuture(
+            wrapped=self._wrapped,
+            invocation_id=self._invocation_id,
+            index=index,
+        )
+
+    def __iter__(self) -> Any:
+        if not self._output_keys:
+            raise ValueError(
+                f"Step {self._invocation_id} does not return any outputs."
+            )
+
+        for index in range(len(self._output_keys)):
+            yield ArtifactFuture(
+                wrapped=self._wrapped,
+                invocation_id=self._invocation_id,
+                index=index,
+            )
+
+    def __len__(self) -> int:
+        return len(self._output_keys)
+
+
+StepRunFuture = Union[ArtifactFuture, StepRunOutputsFuture]

src/zenml/execution/pipeline/dynamic/runner.py

@@ -35,7 +35,9 @@
 from zenml.config.compiler import Compiler
 from zenml.config.step_configurations import Step
 from zenml.execution.pipeline.dynamic.outputs import (
+    ArtifactFuture,
     OutputArtifact,
+    StepRunFuture,
     StepRunOutputs,
     StepRunOutputsFuture,
 )
@@ -153,6 +155,9 @@ def run_pipeline(self) -> None:
                 self._orchestrator.run_init_hook(snapshot=self._snapshot)
                 try:
                     # TODO: step logging isn't threadsafe
+                    # TODO: what should be allowed as pipeline returns?
+                    #  (artifacts, json serializable, anything?)
+                    #  how do we show it in the UI?
                     self.pipeline._call_entrypoint(**pipeline_parameters)
                 except:
                     publish_failed_pipeline_run(run.id)
@@ -171,9 +176,7 @@ def launch_step(
         id: Optional[str],
         args: Tuple[Any],
         kwargs: Dict[str, Any],
-        after: Union[
-            "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
-        ] = None,
+        after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
         concurrent: Literal[False] = False,
     ) -> StepRunOutputs: ...
 
@@ -184,9 +187,7 @@ def launch_step(
         id: Optional[str],
         args: Tuple[Any],
         kwargs: Dict[str, Any],
-        after: Union[
-            "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
-        ] = None,
+        after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
         concurrent: Literal[True] = True,
     ) -> "StepRunOutputsFuture": ...
 
@@ -196,9 +197,7 @@ def launch_step(
         id: Optional[str],
         args: Tuple[Any],
         kwargs: Dict[str, Any],
-        after: Union[
-            "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
-        ] = None,
+        after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
         concurrent: bool = False,
     ) -> Union[StepRunOutputs, "StepRunOutputsFuture"]:
         """Launch a step.
@@ -240,8 +239,11 @@ def _launch() -> StepRunOutputs:
         if concurrent:
             ctx = contextvars.copy_context()
             future = self._executor.submit(ctx.run, _launch)
+            compiled_step.config.outputs
             step_run_future = StepRunOutputsFuture(
-                wrapped=future, invocation_id=compiled_step.spec.invocation_id
+                wrapped=future,
+                invocation_id=compiled_step.spec.invocation_id,
+                output_keys=list(compiled_step.config.outputs),
             )
             self._futures.append(step_run_future)
             return step_run_future
@@ -251,7 +253,7 @@ def _launch() -> StepRunOutputs:
     def await_all_step_run_futures(self) -> None:
         """Await all step run output futures."""
         for future in self._futures:
-            future.wait()
+            future.artifacts()
         self._futures = []
 
 
@@ -262,9 +264,7 @@ def compile_dynamic_step_invocation(
     id: Optional[str],
     args: Tuple[Any],
     kwargs: Dict[str, Any],
-    after: Union[
-        "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
-    ] = None,
+    after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
 ) -> "Step":
     """Compile a dynamic step invocation.
 
@@ -282,27 +282,25 @@ def compile_dynamic_step_invocation(
     """
     upstream_steps = set()
 
-    if isinstance(after, StepRunOutputsFuture):
-        after.wait()
+    if isinstance(after, StepRunFuture):
+        after._wait()
         upstream_steps.add(after._invocation_id)
     elif isinstance(after, Sequence):
         for item in after:
-            item.wait()
+            item._wait()
             upstream_steps.add(item._invocation_id)
 
     def _await_and_validate_input(input: Any) -> Any:
         if isinstance(input, StepRunOutputsFuture):
-            input = input.result()
+            if len(input._output_keys) != 1:
+                raise ValueError(
+                    "Passing multiple step run outputs to another step is not "
+                    "allowed."
+                )
+            input = input.artifacts()
 
-        if (
-            input
-            and isinstance(input, tuple)
-            and isinstance(input[0], OutputArtifact)
-        ):
-            raise ValueError(
-                "Passing multiple step run outputs to another step is not "
-                "allowed."
-            )
+        if isinstance(input, ArtifactFuture):
+            input = input.result()
 
         if isinstance(input, OutputArtifact):
             upstream_steps.add(input.step_name)

src/zenml/orchestrators/step_runner.py

@@ -219,6 +219,7 @@ def run(
             step_failed = False
             try:
                 if (
+                    # TODO: do we need to disable this for dynamic pipelines?
                     pipeline_run.snapshot
                     and self._stack.orchestrator.run_init_cleanup_at_step_level
                 ):

src/zenml/steps/base_step.py

@@ -94,7 +94,10 @@
         Mapping[str, Sequence["MaterializerClassOrSource"]],
     ]
 
-    from zenml.execution.pipeline.dynamic.outputs import StepRunOutputsFuture
+    from zenml.execution.pipeline.dynamic.outputs import (
+        StepRunFuture,
+        StepRunOutputsFuture,
+    )
 
 
 logger = get_logger(__name__)
@@ -468,8 +471,8 @@ def __call__(
         after: Union[
             str,
             StepArtifact,
-            "StepRunOutputsFuture",
-            Sequence[Union[str, StepArtifact, "StepRunOutputsFuture"]],
+            "StepRunFuture",
+            Sequence[Union[str, StepArtifact, "StepRunFuture"]],
             None,
         ] = None,
         **kwargs: Any,
@@ -512,8 +515,8 @@ def __call__(
         if run_context := DynamicPipelineRunContext.get():
             after = cast(
                 Union[
-                    "StepRunOutputsFuture",
-                    Sequence["StepRunOutputsFuture"],
+                    "StepRunFuture",
+                    Sequence["StepRunFuture"],
                     None,
                 ],
                 after,
@@ -627,9 +630,7 @@ def submit(
         self,
         *args: Any,
         id: Optional[str] = None,
-        after: Union[
-            "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
-        ] = None,
+        after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
         **kwargs: Any,
     ) -> "StepRunOutputsFuture":
         from zenml.execution.pipeline.dynamic.run_context import (