fix: clean up step registration

Author: mrkaye97

hatchet_sdk/v2/workflows.py

@@ -116,6 +116,7 @@ def __init__(
         self,
         fn: Callable[[Any, Context], R] | Callable[[Any, Context], Awaitable[R]],
         type: StepType,
+        workflow: Union["BaseWorkflow", None] = None,
         name: str = "",
         timeout: str = "60m",
         parents: list[str] = [],
@@ -129,6 +130,7 @@ def __init__(
     ) -> None:
         self.fn = fn
         self.is_async_function = is_async_fn(fn)
+        self.workflow = workflow
 
         self.type = type
         self.timeout = timeout
@@ -142,44 +144,38 @@ def __init__(
         self.concurrency__max_runs = concurrency__max_runs
         self.concurrency__limit_strategy = concurrency__limit_strategy
 
-
-class RegisteredStep(Generic[R]):
-    def __init__(
-        self,
-        workflow: "BaseWorkflow",
-        step: Step[R],
-    ) -> None:
-        self.workflow = workflow
-        self.step = step
-
     def call(self, ctx: Context) -> R:
-        if self.step.is_async_function:
-            raise TypeError(
-                f"{self.step.name} is not a sync function. Use `acall` instead."
-            )
+        if not self.is_registered:
+            raise ValueError("Only steps that have been registered can be called.")
+
+        if self.is_async_function:
+            raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
 
-        sync_fn = self.step.fn
+        sync_fn = self.fn
         if is_sync_fn(sync_fn):
             return sync_fn(self.workflow, ctx)
 
-        raise TypeError(
-            f"{self.step.name} is not a sync function. Use `acall` instead."
-        )
+        raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
 
     async def acall(self, ctx: Context) -> R:
-        if not self.step.is_async_function:
+        if not self.is_registered:
+            raise ValueError("Only steps that have been registered can be called.")
+
+        if not self.is_async_function:
             raise TypeError(
-                f"{self.step.name} is not an async function. Use `call` instead."
+                f"{self.name} is not an async function. Use `call` instead."
             )
 
-        async_fn = self.step.fn
+        async_fn = self.fn
 
         if is_async_fn(async_fn):
             return await async_fn(self.workflow, ctx)
 
-        raise TypeError(
-            f"{self.step.name} is not an async function. Use `call` instead."
-        )
+        raise TypeError(f"{self.name} is not an async function. Use `call` instead.")
+
+    @property
+    def is_registered(self) -> bool:
+        return self.workflow is not None
 
 
 class WorkflowDeclaration(Generic[TWorkflowInput]):

hatchet_sdk/worker/runner/run_loop_manager.py

@@ -15,7 +15,7 @@
 from hatchet_sdk.worker.runner.utils.capture_logs import capture_logs
 
 if TYPE_CHECKING:
-    from hatchet_sdk.v2.workflows import RegisteredStep
+    from hatchet_sdk.v2.workflows import Step
 
 STOP_LOOP_TYPE = Literal["STOP_LOOP"]
 STOP_LOOP: STOP_LOOP_TYPE = "STOP_LOOP"
@@ -27,7 +27,7 @@
 @dataclass
 class WorkerActionRunLoopManager:
     name: str
-    action_registry: dict[str, "RegisteredStep[Any]"]
+    action_registry: dict[str, "Step[Any]"]
     validator_registry: dict[str, WorkflowValidator]
     max_runs: int | None
     config: ClientConfig

hatchet_sdk/worker/runner/runner.py

@@ -40,7 +40,7 @@
 T = TypeVar("T")
 
 if TYPE_CHECKING:
-    from hatchet_sdk.v2.workflows import RegisteredStep
+    from hatchet_sdk.v2.workflows import Step
 
 
 class WorkerStatus(Enum):
@@ -57,7 +57,7 @@ def __init__(
         event_queue: "Queue[Any]",
         max_runs: int | None = None,
         handle_kill: bool = True,
-        action_registry: dict[str, "RegisteredStep[T]"] = {},
+        action_registry: dict[str, "Step[T]"] = {},
         validator_registry: dict[str, WorkflowValidator] = {},
         config: ClientConfig = ClientConfig(),
         labels: dict[str, str | int] = {},
@@ -69,7 +69,7 @@ def __init__(
         self.max_runs = max_runs
         self.tasks: dict[str, asyncio.Task[Any]] = {}  # Store run ids and futures
         self.contexts: dict[str, Context] = {}  # Store run ids and contexts
-        self.action_registry: dict[str, "RegisteredStep[T]"] = action_registry
+        self.action_registry: dict[str, "Step[T]"] = action_registry
         self.validator_registry = validator_registry
 
         self.event_queue = event_queue
@@ -216,7 +216,7 @@ def inner_callback(task: asyncio.Task[Any]) -> None:
 
     ## TODO: Stricter type hinting here
     def thread_action_func(
-        self, context: Context, step: "RegisteredStep[T]", action: Action
+        self, context: Context, step: "Step[T]", action: Action
     ) -> T:
         if action.step_run_id is not None and action.step_run_id != "":
             self.threads[action.step_run_id] = current_thread()
@@ -233,15 +233,15 @@ def thread_action_func(
     async def async_wrapped_action_func(
         self,
         context: Context,
-        step: "RegisteredStep[T]",
+        step: "Step[T]",
         action: Action,
         run_id: str,
     ) -> T:
         wr.set(context.workflow_run_id())
         sr.set(context.step_run_id)
 
         try:
-            if step.step.is_async_function:
+            if step.is_async_function:
                 return await step.acall(context)
             else:
                 pfunc = functools.partial(

hatchet_sdk/worker/worker.py

@@ -25,7 +25,7 @@
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.types import WorkflowValidator
 from hatchet_sdk.utils.typing import is_basemodel_subclass
-from hatchet_sdk.v2.workflows import RegisteredStep, StepType
+from hatchet_sdk.v2.workflows import Step, StepType
 from hatchet_sdk.worker.action_listener_process import (
     ActionEvent,
     worker_action_listener_process,
@@ -74,7 +74,7 @@ def __init__(
 
         self.client: Client
 
-        self.action_registry: dict[str, RegisteredStep[Any]] = {}
+        self.action_registry: dict[str, Step[Any]] = {}
         self.validator_registry: dict[str, WorkflowValidator] = {}
 
         self.killing: bool = False
@@ -123,10 +123,10 @@ def register_workflow(self, workflow: Union["BaseWorkflow", Any]) -> None:
             sys.exit(1)
 
         for step in workflow.steps:
+            step.workflow = workflow
+
             action_name = workflow.create_action_name(namespace, step)
-            self.action_registry[action_name] = RegisteredStep(
-                workflow=workflow, step=step
-            )
+            self.action_registry[action_name] = step
             return_type = get_type_hints(step.fn).get("return")
 
             self.validator_registry[action_name] = WorkflowValidator(