fix: step needs to be registered to a workflow before it can be run

Author: mrkaye97

hatchet_sdk/v2/workflows.py

@@ -1,4 +1,5 @@
 import asyncio
+from abc import abstractmethod
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
@@ -141,13 +142,23 @@ def __init__(
         self.concurrency__max_runs = concurrency__max_runs
         self.concurrency__limit_strategy = concurrency__limit_strategy
 
+
+class RegisteredStep(Step[R]):
+    def __init__(
+        self,
+        workflow: "BaseWorkflowImpl",
+        step: Step[R],
+    ) -> None:
+        self.workflow = workflow
+        self.step = step
+
     def call(self, ctx: Context) -> R:
         if self.is_async_function:
             raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
 
         sync_fn = self.fn
         if is_sync_fn(sync_fn):
-            return sync_fn(None, ctx)
+            return sync_fn(self.workflow, ctx)
 
         raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
 
@@ -160,7 +171,7 @@ async def acall(self, ctx: Context) -> R:
         async_fn = self.fn
 
         if is_async_fn(async_fn):
-            return await async_fn(None, ctx)
+            return await async_fn(self.workflow, ctx)
 
         raise TypeError(f"{self.name} is not an async function. Use `call` instead.")
 

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 Step
+    from hatchet_sdk.v2.workflows import RegisteredStep
 
 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, "Step[Any]"]
+    action_registry: dict[str, "RegisteredStep[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 Step
+    from hatchet_sdk.v2.workflows import RegisteredStep
 
 
 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, "Step[T]"] = {},
+        action_registry: dict[str, "RegisteredStep[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, "Step[T]"] = action_registry
+        self.action_registry: dict[str, "RegisteredStep[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: "Step[T]", action: Action
+        self, context: Context, step: "RegisteredStep[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,7 +233,7 @@ def thread_action_func(
     async def async_wrapped_action_func(
         self,
         context: Context,
-        step: "Step[T]",
+        step: "RegisteredStep[T]",
         action: Action,
         run_id: str,
     ) -> T:

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 Step, StepType
+from hatchet_sdk.v2.workflows import RegisteredStep, 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, Step[Any]] = {}
+        self.action_registry: dict[str, RegisteredStep[Any]] = {}
         self.validator_registry: dict[str, WorkflowValidator] = {}
 
         self.killing: bool = False
@@ -124,7 +124,9 @@ def register_workflow(self, workflow: Union["BaseWorkflowImpl", Any]) -> None:
 
         for step in workflow.steps:
             action_name = workflow.create_action_name(namespace, step)
-            self.action_registry[action_name] = step
+            self.action_registry[action_name] = RegisteredStep(
+                workflow=workflow, step=step
+            )
             return_type = get_type_hints(step.fn).get("return")
 
             self.validator_registry[action_name] = WorkflowValidator(