webui: add capabilities support

Signed-off-by: Cloaked Clock <213850831+cloakedclock@users.noreply.github.com>

Author: cloakedclock

scaler/protocol/capnp/message.capnp

@@ -150,7 +150,8 @@ struct StateTask {
     functionName @1 :Data;
     state @2 :CommonType.TaskState;
     worker @3 :Data;
-    metadata @4 :Data;
+    capabilities @4 :List(CommonType.TaskCapability);
+    metadata @5 :Data;
 }
 
 struct StateGraphTask {

scaler/protocol/python/message.py

@@ -588,20 +588,32 @@ def state(self) -> TaskState:
     def worker(self) -> WorkerID:
         return WorkerID(self._msg.worker)
 
+    @property
+    def capabilities(self) -> Dict[str, int]:
+        return {capability.name: capability.value for capability in self._msg.capabilities}
+
     @property
     def metadata(self) -> bytes:
         return self._msg.metadata
 
     @staticmethod
     def new_msg(
-        task_id: TaskID, function_name: bytes, task_state: TaskState, worker: WorkerID, metadata: bytes = b""
+        task_id: TaskID,
+        function_name: bytes,
+        task_state: TaskState,
+        worker: WorkerID,
+        capabilities: Dict[str, int],
+        metadata: bytes = b"",
     ) -> "StateTask":
         return StateTask(
             _message.StateTask(
                 taskId=bytes(task_id),
                 functionName=function_name,
                 state=task_state.value,
                 worker=bytes(worker),
+                capabilities=[
+                    TaskCapability.new_msg(name, value).get_message() for name, value in capabilities.items()
+                ],
                 metadata=metadata,
             )
         )

scaler/scheduler/controllers/task_controller.py

@@ -332,7 +332,10 @@ async def __send_task_cancel_confirm_to_client(self, task_cancel_confirm: TaskCa
     async def __send_monitor(self, task_id: TaskID, function_name: bytes, metadata: bytes = b""):
         worker = self._worker_controller.get_worker_by_task_id(task_id)
         task_state = self._task_state_manager.get_state_machine(task_id).current_state()
-        await self._binder_monitor.send(StateTask.new_msg(task_id, function_name, task_state, worker, metadata))
+        capabilities = self._task_id_to_task[task_id].capabilities if task_id in self._task_id_to_task else {}
+        await self._binder_monitor.send(
+            StateTask.new_msg(task_id, function_name, task_state, worker, capabilities, metadata)
+        )
 
     async def __routing(self, task_id: TaskID, transition: TaskTransition, **kwargs):
         state_machine = self._task_state_manager.on_transition(task_id, transition)

scaler/ui/memory_window.py

@@ -69,16 +69,15 @@ def setup_memory_chart(self, settings: Settings):
             self._plot = ui.plotly(self._figure).classes("w-full h-full")
             self._settings = settings
 
-    def handle_task_state(self, state: StateTask):
+    def handle_task_state(self, state_task: StateTask):
         """
         Only completed tasks have profiling data.
         Use this data to fill in history.
         """
-
-        if state.metadata == b"":
+        if state_task.metadata == b"":
             return
 
-        profile_result = ProfileResult.deserialize(state.metadata)
+        profile_result = ProfileResult.deserialize(state_task.metadata)
 
         worker_memory = profile_result.memory_peak
         worker_duration = profile_result.duration_s