soorena62
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
from api.task import Task, Status, TaskEvent
|
|
from api.assets import Asset
|
|
from api.execution_context import ExecutionContext
|
|
from api.io_port import CompatibilityException, InPort, OutPort
|
|
|
|
|
|
class Module(Task):
|
|
"""
|
|
Leaf node: atomic executable unit.
|
|
Executes using its input ports and produces assets on output ports.
|
|
"""
|
|
|
|
def add_in_port(self, port: InPort):
|
|
port.parent_task = self
|
|
self.in_ports.append(port)
|
|
|
|
def add_out_port(self, port: OutPort):
|
|
port.parent_task = self
|
|
self.out_ports.append(port)
|
|
|
|
# Hooks
|
|
def before_run(self, context: ExecutionContext):
|
|
pass
|
|
|
|
def after_run(self, context: ExecutionContext):
|
|
pass
|
|
|
|
def on_error(self, error: Exception):
|
|
pass
|
|
|
|
# Validation
|
|
def validate(self):
|
|
super().validate()
|
|
|
|
# Main execution entry point
|
|
def run(self, context: ExecutionContext):
|
|
self.before_run(context)
|
|
|
|
if not self.check_inputs_ready(context):
|
|
self._set_status(Status.PENDING)
|
|
return
|
|
|
|
self._set_status(Status.RUNNING)
|
|
self._emit(TaskEvent.ON_START, {"module": self.name})
|
|
|
|
try:
|
|
# Gather input data
|
|
consumed_data = {}
|
|
for port in self.in_ports:
|
|
if not port.is_ready(context):
|
|
raise CompatibilityException(
|
|
f"Input {port.full_name()} is not ready"
|
|
)
|
|
asset = context.get(port)
|
|
consumed_data[port.name] = asset.data
|
|
|
|
# Execute module logic
|
|
result = self.execute(consumed_data)
|
|
|
|
if not isinstance(result, dict):
|
|
raise ValueError(f"Module '{self.name}' must return dict")
|
|
|
|
# Store outputs
|
|
for out_port in self.out_ports:
|
|
if out_port.name not in result:
|
|
raise ValueError(
|
|
f"Output '{out_port.name}' missing in module '{self.name}' result"
|
|
)
|
|
|
|
asset = Asset(
|
|
data=result[out_port.name],
|
|
dtype=out_port.dtype
|
|
)
|
|
|
|
context.put(out_port, asset)
|
|
|
|
self._set_status(Status.COMPLETED)
|
|
self._emit(TaskEvent.ON_FINISH, {"module": self.name})
|
|
self.after_run(context)
|
|
|
|
except Exception as e:
|
|
self._set_status(Status.FAILED)
|
|
self.on_error(e)
|
|
self._emit(TaskEvent.ON_ERROR, {"error": str(e)})
|
|
raise e
|
|
|
|
# Default execute
|
|
def execute(self, inputs):
|
|
return inputs
|