feat: Component Status updates (#152)

# Summary
Closes #148.
 
# Changes
* Add a new `Status` enum to `plugboard.schemas`.
* Add a `status` property to the `Component`.
* Removes from flakiness from the tuner tests.

Author: toby-coleman

plugboard/component/component.py

@@ -18,6 +18,7 @@
     UnrecognisedEventError,
     ValidationError,
 )
+from plugboard.schemas.state import Status
 from plugboard.state import StateBackend
 from plugboard.utils import DI, ClassRegistry, ExportMixin, is_on_ray_worker
 
@@ -73,6 +74,8 @@ def __init__(
             namespace=self.name,
             component=self,
         )
+        self.status = Status.CREATED
+        self._is_running = False
         self._field_inputs: dict[str, _t.Any] = {}
         self._field_inputs_ready: bool = False
 
@@ -88,6 +91,12 @@ def __init_subclass__(cls, *args: _t.Any, **kwargs: _t.Any) -> None:
         # Configure IO last in case it fails in case of components with dynamic io args
         cls._configure_io()
 
+    async def _set_status(self, status: Status, publish: bool = True) -> None:
+        """Sets the status of the component and optionaly publishes it to the state backend."""
+        self.status = status
+        if publish and self._state and self._state_is_connected:
+            await self._state.upsert_component(self)
+
     @classmethod
     def _configure_io(cls) -> None:
         # Get all parent classes that are Component subclasses
@@ -203,8 +212,7 @@ def _handle_init_wrapper(self) -> _t.Callable:
         async def _wrapper() -> None:
             with self._job_id_ctx():
                 await self._init()
-                if self._state is not None and self._state_is_connected:
-                    await self._state.upsert_component(self)
+                await self._set_status(Status.INIT)
 
         return _wrapper
 
@@ -233,14 +241,21 @@ def _handle_step_wrapper(self) -> _t.Callable:
         @wraps(self.step)
         async def _wrapper() -> None:
             with self._job_id_ctx():
+                await self._set_status(Status.RUNNING, publish=not self._is_running)
                 await self.io.read()
                 await self._handle_events()
                 self._bind_inputs()
                 if self._can_step:
-                    await self._step()
+                    try:
+                        await self._step()
+                    except Exception as e:
+                        await self._set_status(Status.FAILED)
+                        self._logger.exception("Component step failed")
+                        raise e
                 self._bind_outputs()
                 await self.io.write()
                 self._field_inputs_ready = False
+                await self._set_status(Status.WAITING, publish=not self._is_running)
 
         return _wrapper
 
@@ -304,14 +319,21 @@ async def _stop_event_handler(self, event: StopEvent) -> None:
             await self.io.close()
         except IOStreamClosedError:
             pass
+        await self._set_status(Status.STOPPED)
 
     async def run(self) -> None:
         """Executes component logic for all steps to completion."""
-        while True:
-            try:
-                await self.step()
-            except IOStreamClosedError:
-                break
+        self._is_running = True
+        await self._set_status(Status.RUNNING)
+        try:
+            while True:
+                try:
+                    await self.step()
+                except IOStreamClosedError:
+                    break
+            await self._set_status(Status.COMPLETED)
+        finally:
+            self._is_running = False
 
     async def destroy(self) -> None:
         """Performs tear-down actions for `Component`."""
@@ -327,6 +349,7 @@ def dict(self) -> dict[str, _t.Any]:  # noqa: D102
         return {
             "id": self.id,
             "name": self.name,
+            "status": self.status,
             **field_data,
             "exports": {name: getattr(self, name, None) for name in self.exports or []},
         }

plugboard/process/ray_process.py

@@ -76,6 +76,7 @@ async def _update_component_attributes(self) -> None:
                     **state[str(IODirection.INPUT)],
                     **state[str(IODirection.OUTPUT)],
                     **state["exports"],
+                    "status": state["status"],
                 }
             )
 

plugboard/schemas/__init__.py

@@ -19,7 +19,7 @@
 from .entities import Entity
 from .io import IODirection
 from .process import ProcessArgsDict, ProcessArgsSpec, ProcessSpec
-from .state import StateBackendArgsDict, StateBackendArgsSpec, StateBackendSpec
+from .state import StateBackendArgsDict, StateBackendArgsSpec, StateBackendSpec, Status
 from .tune import (
     Direction,
     ObjectiveSpec,
@@ -55,6 +55,7 @@
     "StateBackendSpec",
     "StateBackendArgsDict",
     "StateBackendArgsSpec",
+    "Status",
     "TuneArgsDict",
     "TuneArgsSpec",
     "TuneSpec",

plugboard/schemas/state.py

@@ -1,6 +1,7 @@
 """Provides `StateBackendSpec` class."""
 
 from datetime import datetime, timezone
+from enum import StrEnum
 import typing as _t
 
 from pydantic import Field
@@ -12,6 +13,33 @@
 DEFAULT_STATE_BACKEND_CLS_PATH: str = "plugboard.state.DictStateBackend"
 
 
+class Status(StrEnum):
+    """`Status` describes the status of either a `Component` or a `Process`.
+
+    Attributes:
+        CREATED: The `Component` or `Process` has been created but not yet started.
+        INIT: The `Component` or `Process` has been initialised but has not started running.
+        RUNNING: The `Component` or `Process` is currently running.
+        WAITING: The `Component` or `Process` is waiting for input.
+        COMPLETED: The `Component` or `Process` has completed successfully.
+        FAILED: The `Component` or `Process` has failed.
+        STOPPED: The `Component` or `Process` has been cancelled or stopped.
+    """
+
+    CREATED = "created"
+    INIT = "init"
+    RUNNING = "running"
+    WAITING = "waiting"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    STOPPED = "stopped"
+
+    @property
+    def is_terminal(self) -> bool:
+        """Returns whether the status is terminal."""
+        return self in {self.COMPLETED, self.FAILED, self.STOPPED}
+
+
 class StateBackendArgsDict(_t.TypedDict):
     """`TypedDict` of the [`StateBackend`][plugboard.state.StateBackend] constructor arguments."""
 

tests/integration/test_process_stop_event.py

@@ -11,7 +11,7 @@
 from plugboard.connector import AsyncioConnector, Connector, ConnectorBuilder, RabbitMQConnector
 from plugboard.events import EventConnectorBuilder, StopEvent
 from plugboard.process import LocalProcess, Process, RayProcess
-from plugboard.schemas import ConnectorSpec
+from plugboard.schemas import ConnectorSpec, Status
 from tests.conftest import ComponentTestHelper, zmq_connector_cls
 
 
@@ -98,6 +98,7 @@ async def stop_after() -> None:
 
         for c in components:
             assert c.is_initialised
+            assert c.status == Status.INIT
 
         async with asyncio.TaskGroup() as tg:
             tg.create_task(process.run())
@@ -116,6 +117,7 @@ async def stop_after() -> None:
         for c in [comp_b1, comp_b2, comp_b3, comp_b4, comp_b5]:
             assert c.is_finished
             assert c.step_count == pytest.approx(iters_before_stop, abs=stop_tolerance)
+            assert c.status == Status.STOPPED
 
         # A performs n+1 full steps and is interrupted on step n+2 before a final update of out_1,
         # hence n+2.

tests/integration/test_state_backend.py

@@ -8,7 +8,7 @@
 from plugboard.component import Component, IOController
 from plugboard.connector import AsyncioConnector, Connector
 from plugboard.process import LocalProcess
-from plugboard.schemas import ConnectorSpec
+from plugboard.schemas import ConnectorSpec, Status
 from plugboard.state import StateBackend
 from tests.conftest import ComponentTestHelper
 from tests.integration.conftest import (
@@ -208,4 +208,8 @@ async def test_state_backend_process_init(
     assert await state_backend.get_connector(conn_1.id) == conn_1.dict()
     assert await state_backend.get_connector(conn_2.id) == conn_2.dict()
 
+    # All components must report their INIT status to the StateBackend
+    for c in B_components:
+        assert (await state_backend.get_component(c.id))["status"] == Status.INIT
+
     await process.destroy()

tests/integration/test_tuner.py

@@ -116,7 +116,8 @@ async def test_multi_objective_tune(config: dict, ray_ctx: None) -> None:
     assert not [t for t in result if t.error]
     # Results must contain two objectives and correct optimimum must be found
     # The best result must be a list of two results
-    assert best_result[0].config["a.iters"] == 2
-    assert best_result[1].config["b.factor"] == -1
-    assert best_result[0].metrics["c.in_1"] == 1
-    assert best_result[1].metrics["b.out_1"] == -1
+    assert len(best_result) == 2
+    assert all(r.config["a.iters"] == 2 for r in best_result)
+    assert -1 in set(r.config["b.factor"] for r in best_result)
+    assert -1 in set(r.metrics["b.out_1"] for r in best_result)
+    assert 1 in set(r.metrics["c.in_1"] for r in best_result)

tests/unit/test_component.py

@@ -7,7 +7,7 @@
 
 from plugboard.component import Component, IOController as IO
 from plugboard.connector import AsyncioConnector
-from plugboard.schemas import ConnectorSpec
+from plugboard.schemas import ComponentArgsDict, ConnectorSpec, Status
 
 
 class A(Component):
@@ -17,6 +17,23 @@ async def step(self) -> None:
         self.c = {"a": self.a, "b": self.b}
 
 
+class S(Component):
+    io = IO(inputs=[], outputs=["component_status"])
+
+    def __init__(self, max_iters: int = 5, **kwargs: _t.Unpack[ComponentArgsDict]) -> None:
+        super().__init__(**kwargs)
+        self._remaining_iters = max_iters
+        self.raise_exception = False
+
+    async def step(self) -> None:
+        self._remaining_iters -= 1
+        if self.raise_exception:
+            raise ValueError("This is a test exception.")
+        self.component_status = str(self.status)
+        if self._remaining_iters <= 0:
+            await self.io.close()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "initial_values", [{"a": [-1], "b": [-2]}, {"a": [-2]}, {"a": [-2, -1]}, {}]
@@ -49,3 +66,32 @@ async def test_component_initial_values(initial_values: dict[str, _t.Iterable])
                 assert component.c.get(field) == input_idx - n_init[field]
 
     await component.io.close()
+
+
+@pytest.mark.asyncio
+async def test_component_status() -> None:
+    """Tests the status of a `Component` across its lifecycle."""
+    component = S(name="status-component")
+    assert component.status == Status.CREATED, "Component should start with CREATED status"
+
+    await component.init()
+    assert component.status == Status.INIT, "Component should be INIT after init"
+
+    component.raise_exception = False
+    await component.step()
+    assert component.status == Status.WAITING, "Component should be WAITING after step"
+    assert component.component_status == "running", "Status should be running during step"
+
+    component.raise_exception = True
+    try:
+        await component.step()
+    except ValueError:
+        pass
+    assert component.status == Status.FAILED, "Component should be FAILED after exception"
+
+    run_component = S(name="run-status-component")
+    await run_component.init()
+
+    run_component.raise_exception = False
+    await run_component.run()
+    assert run_component.status == Status.COMPLETED, "Component should be COMPLETED after run"