soorena62
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
from typing import Any, Dict
|
|
from workflow.task import Task, ExecutionContext, Status, TaskEvent
|
|
|
|
|
|
class Module(Task):
|
|
"""
|
|
Leaf in the Composite Pattern.
|
|
Performs actual computation.
|
|
Users override ONLY run().
|
|
"""
|
|
|
|
def __init__(self, name: str):
|
|
super().__init__(name)
|
|
self._in_ports: Dict[str, None] = {}
|
|
self._out_ports: Dict[str, None] = {}
|
|
|
|
# Ports (Module-only)
|
|
def add_in_port(self, name: str):
|
|
self._in_ports[name] = None
|
|
|
|
def add_out_port(self, name: str):
|
|
self._out_ports[name] = None
|
|
|
|
# Execution
|
|
def run(self, context: ExecutionContext):
|
|
"""
|
|
Final execution method.
|
|
Subclasses override THIS method only.
|
|
No compute(), no run_node(), no extra execution methods.
|
|
"""
|
|
|
|
# Logging + Events + Status
|
|
self.log(f"Starting module: {self.name}")
|
|
context.log(f"[{self.name}] start")
|
|
self._emit(TaskEvent.BEFORE_RUN)
|
|
self._set_status(Status.RUNNING)
|
|
|
|
try:
|
|
# Read inputs from ExecutionContex
|
|
inputs: Dict[str, Any] = {}
|
|
for port_name in self._in_ports.keys():
|
|
key = f"{self.name}.{port_name}"
|
|
inputs[port_name] = context.get_asset(key)
|
|
|
|
result = self._user_run(context, inputs)
|
|
|
|
# Write outputs
|
|
if isinstance(result, dict):
|
|
for out_name, value in result.items():
|
|
key = f"{self.name}.{out_name}"
|
|
context.put_asset(key, value)
|
|
else:
|
|
context.put_asset(f"{self.name}.result", result)
|
|
|
|
# Metadata
|
|
context.add_metadata(f"{self.name}.status", "completed")
|
|
|
|
# Status + Events + Logging
|
|
self._set_status(Status.COMPLETED)
|
|
self._emit(TaskEvent.COMPLETED, result)
|
|
self.log(f"Module completed: {self.name}")
|
|
context.log(f"[{self.name}] completed")
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
self._set_status(Status.FAILED)
|
|
self._emit(TaskEvent.ERROR, str(e))
|
|
self.log(f"Module failed: {self.name} — {e}")
|
|
context.log(f"[{self.name}] failed: {e}")
|
|
raise
|
|
|
|
# Internal wrapper for user logic
|
|
def _user_run(self, context: ExecutionContext, inputs: Dict[str, Any]):
|
|
"""
|
|
This is the ONLY method users override.
|
|
"""
|
|
raise NotImplementedError(
|
|
f"Module subclass '{self.name}' must implement run(context, inputs)"
|
|
) |