More fixes

Author: schustmi

src/zenml/orchestrators/step_launcher.py

@@ -459,21 +459,26 @@ def _run_step(
                     step_operator_name=step_operator_name,
                     step_run_info=step_run_info,
                 )
+            elif not self._snapshot.is_dynamic:
+                self._run_step_in_current_thread(
+                    pipeline_run=pipeline_run,
+                    step_run=step_run,
+                    step_run_info=step_run_info,
+                    input_artifacts=step_run.regular_inputs,
+                    output_artifact_uris=output_artifact_uris,
+                )
             else:
                 from zenml.pipelines.dynamic.runner import (
                     should_run_in_process,
                 )
 
-                should_run_out_of_process = (
-                    self._snapshot.is_dynamic
-                    and self._step.config.in_process is False
-                )
-
                 if should_run_in_process(
                     self._step,
                     self._snapshot.pipeline_configuration.docker_settings,
                 ):
-                    if should_run_out_of_process:
+                    if self._step.config.in_process is False:
+                        # The step was configured to run out of process, but
+                        # the orchestrator doesn't support it.
                         logger.warning(
                             "The %s does not support running dynamic out of "
                             "process steps. Running step `%s` in current "

src/zenml/pipelines/dynamic/context.py

@@ -1,5 +1,5 @@
 import contextvars
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Self
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Self, cast
 
 if TYPE_CHECKING:
     from zenml.models import PipelineRunResponse, PipelineSnapshotResponse
@@ -8,20 +8,38 @@
 
 
 class BaseContext:
+    """Base context class."""
+
     __context_var__: ClassVar[contextvars.ContextVar[Self]]
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self) -> None:
+        """Initialize the context."""
         self._token: Optional[contextvars.Token[Any]] = None
 
     @classmethod
     def get(cls: type[Self]) -> Optional[Self]:
-        return cls.__context_var__.get(None)
+        """Get the active context for the current thread.
+
+        Returns:
+            The active context for the current thread.
+        """
+        return cast(Optional[Self], cls.__context_var__.get(None))
 
     def __enter__(self) -> Self:
+        """Enter the context.
+
+        Returns:
+            The context object.
+        """
         self._token = self.__context_var__.set(self)
         return self
 
     def __exit__(self, *_: Any) -> None:
+        """Exit the context.
+
+        Raises:
+            RuntimeError: If the context has not been entered.
+        """
         if not self._token:
             raise RuntimeError(
                 f"Can't exit {self.__class__.__name__} because it has not been "
@@ -31,6 +49,8 @@ def __exit__(self, *_: Any) -> None:
 
 
 class DynamicPipelineRunContext(BaseContext):
+    """Dynamic pipeline run context."""
+
     __context_var__ = contextvars.ContextVar("dynamic_pipeline_run_context")
 
     def __init__(
@@ -40,6 +60,14 @@ def __init__(
         run: "PipelineRunResponse",
         runner: "DynamicPipelineRunner",
     ) -> None:
+        """Initialize the dynamic pipeline run context.
+
+        Args:
+            pipeline: The dynamic pipeline that is being executed.
+            snapshot: The snapshot of the pipeline.
+            run: The pipeline run.
+            runner: The dynamic pipeline runner.
+        """
         super().__init__()
         self._pipeline = pipeline
         self._snapshot = snapshot
@@ -63,6 +91,15 @@ def runner(self) -> "DynamicPipelineRunner":
         return self._runner
 
     def __enter__(self) -> Self:
+        """Enter the dynamic pipeline run context.
+
+        Raises:
+            RuntimeError: If the dynamic pipeline run context has already been
+                entered.
+
+        Returns:
+            The dynamic pipeline run context object.
+        """
         if self._token is not None:
             raise RuntimeError(
                 "Calling a pipeline within a dynamic pipeline is not allowed."

src/zenml/pipelines/dynamic/pipeline_definition.py

@@ -236,7 +236,7 @@ def __call__(
         """
         if should_prevent_pipeline_execution():
             logger.info("Preventing execution of pipeline '%s'.", self.name)
-            return
+            return None
 
         stack = Client().active_stack
         if not stack.orchestrator.supports_dynamic_pipelines:
@@ -286,7 +286,7 @@ def _compute_output_schema(self) -> Optional[Dict[str, Any]]:
         """
         try:
             outputs = parse_return_type_annotations(self.entrypoint)
-            model_fields = {
+            model_fields: Dict[str, Any] = {
                 name: (output.resolved_annotation, ...)
                 for name, output in outputs.items()
             }

src/zenml/pipelines/dynamic/runner.py

@@ -39,6 +39,7 @@
 from zenml.pipelines.run_utils import create_placeholder_run
 from zenml.stack import Stack
 from zenml.steps.entrypoint_function_utils import StepArtifact
+from zenml.steps.utils import OutputSignature
 from zenml.utils import source_utils
 
 if TYPE_CHECKING:
@@ -51,6 +52,8 @@
 
 
 class DynamicStepRunOutput(ArtifactVersionResponse):
+    """Dynamic step run output artifact."""
+
     output_name: str
     step_name: str
 
@@ -61,17 +64,41 @@ class DynamicStepRunOutput(ArtifactVersionResponse):
 
 
 class StepRunOutputsFuture:
-    def __init__(self, wrapped: Future[StepRunOutputs], invocation_id: str):
+    """Future for a step run output."""
+
+    def __init__(
+        self, wrapped: Future[StepRunOutputs], invocation_id: str
+    ) -> 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
 
     def wait(self) -> None:
+        """Wait for the future to complete."""
         self._wrapped.result()
 
     def result(self) -> StepRunOutputs:
+        """Get the step run output artifacts.
+
+        Returns:
+            The step run output artifacts.
+        """
         return self._wrapped.result()
 
     def load(self) -> Any:
+        """Get the step run output artifact data.
+
+        Raises:
+            ValueError: If the step run output is invalid.
+
+        Returns:
+            The step run output artifact data.
+        """
         result = self.result()
 
         if result is None:
@@ -153,6 +180,7 @@ def run_pipeline(self) -> None:
             ):
                 self._orchestrator.run_init_hook(snapshot=self._snapshot)
                 try:
+                    # TODO: step logging isn't threadsafe
                     self.pipeline._call_entrypoint(**pipeline_parameters)
                 except:
                     publish_failed_pipeline_run(run.id)
@@ -240,6 +268,20 @@ def _prepare_step_run(
         "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
     ] = None,
 ) -> Tuple[Dict[str, Any], Set[str]]:
+    """Prepare a step run.
+
+    Args:
+        step: The step to prepare.
+        args: The arguments for the step function.
+        kwargs: The keyword arguments for the step function.
+        after: The step run output futures to wait for.
+
+    Returns:
+        A tuple containing the inputs and the upstream steps.
+
+    Raises:
+        ValueError: If an invalid step function input was passed.
+    """
     upstream_steps = set()
 
     if isinstance(after, StepRunOutputsFuture):
@@ -250,7 +292,7 @@ def _prepare_step_run(
             item.wait()
             upstream_steps.add(item._invocation_id)
 
-    def _await_and_validate_input(input: Any):
+    def _await_and_validate_input(input: Any) -> Any:
         if isinstance(input, StepRunOutputsFuture):
             input = input.result()
 
@@ -260,15 +302,16 @@ def _await_and_validate_input(input: Any):
             and isinstance(input[0], DynamicStepRunOutput)
         ):
             raise ValueError(
-                "Passing multiple step run outputs to another step is not allowed."
+                "Passing multiple step run outputs to another step is not "
+                "allowed."
             )
 
         if isinstance(input, DynamicStepRunOutput):
             upstream_steps.add(input.step_name)
 
         return input
 
-    args = [_await_and_validate_input(arg) for arg in args]
+    args = tuple(_await_and_validate_input(arg) for arg in args)
     kwargs = {
         key: _await_and_validate_input(value) for key, value in kwargs.items()
     }
@@ -295,7 +338,7 @@ def _compile_step(
             input_artifacts[name] = StepArtifact(
                 invocation_id=value.step_name,
                 output_name=value.output_name,
-                annotation=Any,
+                annotation=OutputSignature(resolved_annotation=Any),
                 pipeline=pipeline,
             )
         elif isinstance(value, (ArtifactVersionResponse, ExternalArtifact)):
@@ -335,6 +378,18 @@ def _run_step_sync(
     orchestrator_run_id: str,
     retry: bool = False,
 ) -> StepRunResponse:
+    """Run a step in the active thread.
+
+    Args:
+        snapshot: The snapshot.
+        step: The step to run.
+        orchestrator_run_id: The orchestrator run ID.
+        retry: Whether to retry the step if it fails.
+
+    Returns:
+        The step run response.
+    """
+
     def _launch_step() -> StepRunResponse:
         launcher = StepLauncher(
             snapshot=snapshot,
@@ -383,6 +438,14 @@ def _launch_step() -> StepRunResponse:
 
 
 def _load_step_outputs(step_run_id: UUID) -> StepRunOutputs:
+    """Load the outputs of a step run.
+
+    Args:
+        step_run_id: The ID of the step run.
+
+    Returns:
+        The outputs of the step run.
+    """
     step_run = Client().zen_store.get_run_step(step_run_id)
 
     def _convert_output_artifact(
@@ -410,6 +473,15 @@ def _convert_output_artifact(
 def _should_retry_locally(
     step: "Step", pipeline_docker_settings: "DockerSettings"
 ) -> bool:
+    """Determine if a step should be retried locally.
+
+    Args:
+        step: The step.
+        pipeline_docker_settings: The Docker settings of the parent pipeline.
+
+    Returns:
+        Whether the step should be retried locally.
+    """
     if step.config.step_operator:
         return True
 
@@ -425,6 +497,15 @@ def _should_retry_locally(
 def should_run_in_process(
     step: "Step", pipeline_docker_settings: "DockerSettings"
 ) -> bool:
+    """Determine if a step should be run in process.
+
+    Args:
+        step: The step.
+        pipeline_docker_settings: The Docker settings of the parent pipeline.
+
+    Returns:
+        Whether the step should be run in process.
+    """
     if step.config.step_operator:
         return False
 
@@ -448,6 +529,16 @@ def get_config_template(
     step: "BaseStep",
     pipeline: "DynamicPipeline",
 ) -> Optional["Step"]:
+    """Get the config template for a step executed in a dynamic pipeline.
+
+    Args:
+        snapshot: The snapshot of the pipeline.
+        step: The step to get the config template for.
+        pipeline: The dynamic pipeline that the step is being executed in.
+
+    Returns:
+        The config template for the step.
+    """
     for index, step_ in enumerate(pipeline.depends_on):
         if step_._static_id == step._static_id:
             break

src/zenml/pipelines/run_utils.py

@@ -2,7 +2,16 @@
 
 import time
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Set,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -60,7 +69,7 @@ def should_prevent_pipeline_execution() -> bool:
 
 
 @contextmanager
-def prevent_pipeline_execution():
+def prevent_pipeline_execution() -> Generator[None, None, None]:
     """Context manager to prevent pipeline execution."""
     with env_utils.temporary_environment(
         {ENV_ZENML_PREVENT_PIPELINE_EXECUTION: "True"}

src/zenml/steps/base_step.py

@@ -30,6 +30,7 @@
     Type,
     TypeVar,
     Union,
+    cast,
 )
 from uuid import UUID
 
@@ -489,6 +490,12 @@ def __call__(
         from zenml.pipelines.pipeline_definition import Pipeline
 
         if context := DynamicPipelineRunContext.get():
+            after = cast(
+                Union[
+                    "StepRunOutputsFuture", Sequence["StepRunOutputsFuture"], None
+                ],
+                after,
+            )
             return context.runner.run_step_sync(self, id, args, kwargs, after)
 
         if not Pipeline.ACTIVE_PIPELINE: