Add All Folders

Dagster/API_R&D/api/module.py

@@ -0,0 +1,89 @@
+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