Use ParamSpec for more clarity in the typing of BaseStep

Author: simonpernas

src/zenml/integrations/huggingface/steps/accelerate_runner.py

@@ -17,7 +17,17 @@
 """Step function to run any ZenML step using Accelerate."""
 
 import functools
-from typing import Any, Callable, Dict, Optional, TypeVar, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    ParamSpec,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 import cloudpickle as pickle
 from accelerate.commands.launch import (
@@ -32,13 +42,27 @@
 
 logger = get_logger(__name__)
 F = TypeVar("F", bound=Callable[..., Any])
-T = TypeVar("T", bound=BaseStep[Any])
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+@overload
+def run_with_accelerate(
+    **accelerate_launch_kwargs: Any,
+) -> Callable[[BaseStep[P, R]], BaseStep[P, R]]: ...
+
+
+@overload
+def run_with_accelerate(
+    step_function_top_level: BaseStep[P, R],
+    /,
+) -> BaseStep[P, R]: ...
 
 
 def run_with_accelerate(
-    step_function_top_level: Optional[T] = None,
+    step_function_top_level: Optional[BaseStep[P, R]] = None,
     **accelerate_launch_kwargs: Any,
-) -> Union[Callable[[T], T], T]:
+) -> Union[Callable[[BaseStep[P, R]], BaseStep[P, R]], BaseStep[P, R]]:
     """Run a function with accelerate.
 
     Accelerate package: https://huggingface.co/docs/accelerate/en/index
@@ -71,9 +95,10 @@ def training_pipeline(some_param: int, ...):
         The accelerate-enabled version of the step.
     """
 
-    def _decorator(step_function: T) -> T:
+    def _decorator(step_function: BaseStep[P, R]) -> BaseStep[P, R]:
         def _wrapper(
-            entrypoint: F, accelerate_launch_kwargs: Dict[str, Any]
+            entrypoint: F,
+            accelerate_launch_kwargs: Dict[str, Any],
         ) -> F:
             @functools.wraps(entrypoint)
             def inner(*args: Any, **kwargs: Any) -> Any:

src/zenml/integrations/whylogs/steps/whylogs_profiler.py

@@ -58,7 +58,7 @@ def get_whylogs_profiler_step(
     dataset_timestamp: Optional[datetime.datetime] = None,
     dataset_id: Optional[str] = None,
     enable_whylabs: bool = True,
-) -> BaseStep[Any]:
+) -> BaseStep[..., Any]:
     """Shortcut function to create a new instance of the WhylogsProfilerStep step.
 
     The returned WhylogsProfilerStep can be used in a pipeline to generate a

src/zenml/orchestrators/step_run_utils.py

@@ -185,7 +185,7 @@ def _get_docstring_and_source_code_from_step_instance(
         """
         from zenml.steps.base_step import BaseStep
 
-        step_instance = BaseStep[Any].load_from_source(step.spec.source)
+        step_instance = BaseStep[..., Any].load_from_source(step.spec.source)
 
         docstring = step_instance.docstring
         if docstring and len(docstring) > TEXT_FIELD_MAX_LENGTH:

src/zenml/orchestrators/step_runner.py

@@ -302,15 +302,17 @@ def _evaluate_artifact_names_in_collections(
             for d in collections:
                 d[name] = d.pop(k)
 
-    def _load_step(self) -> "BaseStep[Any]":
+    def _load_step(self) -> "BaseStep[..., Any]":
         """Load the step instance.
 
         Returns:
             The step instance.
         """
         from zenml.steps import BaseStep
 
-        step_instance = BaseStep[Any].load_from_source(self._step.spec.source)
+        step_instance = BaseStep[..., Any].load_from_source(
+            self._step.spec.source
+        )
         step_instance = copy.deepcopy(step_instance)
         step_instance._configuration = self._step.config
         return step_instance

src/zenml/pipelines/pipeline_definition.py

@@ -1138,7 +1138,7 @@ def _compute_unique_identifier(self, pipeline_spec: PipelineSpec) -> str:
 
     def add_step_invocation(
         self,
-        step: "BaseStep[Any]",
+        step: "BaseStep[..., Any]",
         input_artifacts: Dict[str, StepArtifact],
         external_artifacts: Dict[
             str, Union["ExternalArtifact", "ArtifactVersionResponse"]
@@ -1208,7 +1208,7 @@ def add_step_invocation(
 
     def _compute_invocation_id(
         self,
-        step: "BaseStep[Any]",
+        step: "BaseStep[..., Any]",
         custom_id: Optional[str] = None,
         allow_suffix: bool = True,
     ) -> str:

src/zenml/steps/base_step.py

@@ -16,17 +16,17 @@
 import copy
 import hashlib
 import inspect
+from abc import abstractmethod
 from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     Dict,
     Generic,
     List,
     Mapping,
     Optional,
-    Protocol,
+    ParamSpec,
     Sequence,
     Tuple,
     Type,
@@ -93,24 +93,11 @@
 
 logger = get_logger(__name__)
 
-T = TypeVar("T", bound="BaseStep[Any]")
-F = TypeVar("F", bound=Callable[..., Any])
+P = ParamSpec("P")
+R = TypeVar("R")
 
 
-class _AbstractEntrypoint(Protocol[F]):
-    entrypoint: F
-    """Abstract method for core step logic.
-
-        Args:
-            *args: Positional arguments passed to the step.
-            **kwargs: Keyword arguments passed to the step.
-
-        Returns:
-            The output of the step.
-    """
-
-
-class BaseStep(Generic[F], _AbstractEntrypoint[F]):
+class BaseStep(Generic[P, R]):
     """Abstract base class for all ZenML steps."""
 
     def __init__(
@@ -227,8 +214,20 @@ def __init__(
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
 
+    @abstractmethod
+    def entrypoint(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Abstract method for core step logic.
+
+        Args:
+            *args: Positional arguments passed to the step.
+            **kwargs: Keyword arguments passed to the step.
+
+        Returns:
+            The output of the step.
+        """
+
     @classmethod
-    def load_from_source(cls, source: Union[Source, str]) -> "BaseStep[F]":
+    def load_from_source(cls, source: Union[Source, str]) -> "BaseStep[P, R]":
         """Loads a step from source.
 
         Args:
@@ -585,7 +584,7 @@ def configuration(self) -> "PartialStepConfiguration":
         return self._configuration
 
     def configure(
-        self: T,
+        self: "BaseStep[P,R]",
         enable_cache: Optional[bool] = None,
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
@@ -604,7 +603,7 @@ def configure(
         merge: bool = True,
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
-    ) -> T:
+    ) -> "BaseStep[P,R]":
         """Configures the step.
 
         Configuration merging example:
@@ -737,7 +736,7 @@ def with_options(
         model: Optional["Model"] = None,
         merge: bool = True,
         substitutions: Optional[Dict[str, str]] = None,
-    ) -> "BaseStep[F]":
+    ) -> "BaseStep[P, R]":
         """Copies the step and applies the given configurations.
 
         Args:
@@ -793,7 +792,7 @@ def with_options(
         )
         return step_copy
 
-    def copy(self) -> "BaseStep[F]":
+    def copy(self) -> "BaseStep[P, R]":
         """Copies the step.
 
         Returns:

src/zenml/steps/decorated_step.py

@@ -19,7 +19,7 @@
 from zenml.steps import BaseStep
 
 
-class _DecoratedStep(BaseStep[Any]):
+class _DecoratedStep(BaseStep[..., Any]):
     """Internal BaseStep subclass used by the step decorator."""
 
     @property

src/zenml/steps/step_decorator.py

@@ -20,6 +20,7 @@
     Dict,
     Mapping,
     Optional,
+    ParamSpec,
     Sequence,
     Type,
     TypeVar,
@@ -47,13 +48,15 @@
         Mapping[str, Sequence[MaterializerClassOrSource]],
     ]
     F = TypeVar("F", bound=Callable[..., Any])
+    P = ParamSpec("P")
+    R = TypeVar("R")
 
 
 logger = get_logger(__name__)
 
 
 @overload
-def step(_func: "F") -> "BaseStep[F]": ...
+def step(_func: Callable[P, R]) -> "BaseStep[P,R]": ...
 
 
 @overload
@@ -74,11 +77,11 @@ def step(
     model: Optional["Model"] = None,
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
-) -> Callable[["F"], "BaseStep[F]"]: ...
+) -> Callable[[Callable[P, R]], "BaseStep[P,R]"]: ...
 
 
 def step(
-    _func: Optional["F"] = None,
+    _func: Optional[Callable[P, R]] = None,
     *,
     name: Optional[str] = None,
     enable_cache: Optional[bool] = None,
@@ -95,7 +98,7 @@ def step(
     model: Optional["Model"] = None,
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
-) -> Union["BaseStep[F]", Callable[["F"], "BaseStep[F]"]]:
+) -> Union["BaseStep[P,R]", Callable[[Callable[P, R]], "BaseStep[P,R]"]]:
     """Decorator to create a ZenML step.
 
     Args:
@@ -132,10 +135,10 @@ def step(
         The step instance.
     """
 
-    def inner_decorator(func: "F") -> "BaseStep[F]":
+    def inner_decorator(func: Callable[P, R]) -> "BaseStep[P,R]":
         from zenml.steps.decorated_step import _DecoratedStep
 
-        class_: Type["BaseStep[F]"] = type(
+        class_: Type["BaseStep[P,R]"] = type(
             func.__name__,
             (_DecoratedStep,),
             {

src/zenml/steps/step_invocation.py

@@ -33,7 +33,7 @@ class StepInvocation:
     def __init__(
         self,
         id: str,
-        step: "BaseStep[Any]",
+        step: "BaseStep[..., Any]",
         input_artifacts: Dict[str, "StepArtifact"],
         external_artifacts: Dict[
             str, Union["ExternalArtifact", "ArtifactVersionResponse"]

src/zenml/steps/utils.py

@@ -24,7 +24,9 @@
     Callable,
     Dict,
     Optional,
+    ParamSpec,
     Tuple,
+    TypeVar,
     Union,
 )
 from uuid import UUID
@@ -49,6 +51,9 @@
 if TYPE_CHECKING:
     from zenml.steps import BaseStep
 
+    P = ParamSpec("P")
+    R = TypeVar("R")
+
 
 logger = get_logger(__name__)
 
@@ -499,7 +504,7 @@ def log_step_metadata(
 
 
 def run_as_single_step_pipeline(
-    __step: "BaseStep[Any]", *args: Any, **kwargs: Any
+    __step: "BaseStep[P,R]", *args: Any, **kwargs: Any
 ) -> Any:
     """Runs the step as a single step pipeline.