fix(BA-934): Resolve GQL Agent live_stat field (#3928)

Author: fregataa

changes/3928.fix.md

@@ -0,0 +1 @@
+Fix GQL Agent `live_stat` resolver to properly parse UUID keys in JSON data as strings

src/ai/backend/common/msgpack.py

@@ -7,9 +7,10 @@
 import os
 import pickle
 import uuid
+from collections.abc import Mapping
 from decimal import Decimal
 from pathlib import PosixPath, PurePosixPath
-from typing import Any
+from typing import Any, Callable, Optional, Protocol
 
 import msgpack as _msgpack
 import temporenc
@@ -59,29 +60,45 @@ def _default(obj: object) -> _msgpack.ExtType:
     raise TypeError(f"Unknown type: {obj!r} ({type(obj)})")
 
 
-def _ext_hook(code: int, data: bytes) -> Any:
-    match code:
-        case ExtTypes.UUID:
-            return uuid.UUID(bytes=data)
-        case ExtTypes.DATETIME:
-            return temporenc.unpackb(data).datetime()
-        case ExtTypes.DECIMAL:
-            return pickle.loads(data)
-        case ExtTypes.POSIX_PATH:
-            return PosixPath(os.fsdecode(data))
-        case ExtTypes.PURE_POSIX_PATH:
-            return PurePosixPath(os.fsdecode(data))
-        case ExtTypes.ENUM:
-            return pickle.loads(data)
-        case ExtTypes.RESOURCE_SLOT:
-            return pickle.loads(data)
-        case ExtTypes.BACKENDAI_BINARY_SIZE:
-            return pickle.loads(data)
-        case ExtTypes.IMAGE_REF:
-            return pickle.loads(data)
-    return _msgpack.ExtType(code, data)
+class ExtFunc(Protocol):
+    def __call__(self, data: bytes) -> Any:
+        pass
 
 
+_DEFAULT_EXT_HOOK: Mapping[ExtTypes, ExtFunc] = {
+    ExtTypes.UUID: lambda data: uuid.UUID(bytes=data),
+    ExtTypes.DATETIME: lambda data: temporenc.unpackb(data).datetime(),
+    ExtTypes.DECIMAL: lambda data: pickle.loads(data),
+    ExtTypes.POSIX_PATH: lambda data: PosixPath(os.fsdecode(data)),
+    ExtTypes.PURE_POSIX_PATH: lambda data: PurePosixPath(os.fsdecode(data)),
+    ExtTypes.ENUM: lambda data: pickle.loads(data),
+    ExtTypes.RESOURCE_SLOT: lambda data: pickle.loads(data),
+    ExtTypes.BACKENDAI_BINARY_SIZE: lambda data: pickle.loads(data),
+    ExtTypes.IMAGE_REF: lambda data: pickle.loads(data),
+}
+
+
+class _Deserializer:
+    def __init__(self, mapping: Optional[Mapping[int, ExtFunc]] = None):
+        self._ext_hook: dict[int, ExtFunc] = {}
+        mapping = mapping or {}
+        self._ext_hook = {**mapping}
+        for ext_type, func in _DEFAULT_EXT_HOOK.items():
+            if ext_type not in self._ext_hook:
+                self._ext_hook[ext_type] = func
+
+    @property
+    def ext_hook(self) -> Callable[[int, bytes], Any]:
+        def _hook_callable(code: int, data: bytes) -> Any:
+            if code in self._ext_hook:
+                return self._ext_hook[code](data)
+            return _msgpack.ExtType(code, data)
+
+        return _hook_callable
+
+
+uuid_to_str: Mapping[int, ExtFunc] = {ExtTypes.UUID: lambda data: str(uuid.UUID(bytes=data))}
+
 DEFAULT_PACK_OPTS = {
     "use_bin_type": True,  # bytes -> bin type (default for Python 3)
     "strict_types": True,  # do not serialize subclasses using superclasses
@@ -92,7 +109,7 @@ def _ext_hook(code: int, data: bytes) -> Any:
     "raw": False,  # assume str as UTF-8 (default for Python 3)
     "strict_map_key": False,  # allow using UUID as map keys
     "use_list": False,  # array -> tuple
-    "ext_hook": _ext_hook,
+    "ext_hook": _Deserializer().ext_hook,
 }
 
 
@@ -104,6 +121,10 @@ def packb(data: Any, **kwargs) -> bytes:
     return ret
 
 
-def unpackb(packed: bytes, **kwargs) -> Any:
+def unpackb(
+    packed: bytes, ext_hook_mapping: Optional[Mapping[int, ExtFunc]] = None, **kwargs
+) -> Any:
     opts = {**DEFAULT_UNPACK_OPTS, **kwargs}
+    if ext_hook_mapping is not None:
+        opts["ext_hook"] = _Deserializer(ext_hook_mapping).ext_hook
     return _msgpack.unpackb(packed, **opts)

src/ai/backend/manager/models/gql_models/agent.py

@@ -238,7 +238,7 @@ async def _pipe_builder(r: Redis):
         ret = []
         for stat in await redis_helper.execute(ctx.redis_stat, _pipe_builder):
             if stat is not None:
-                ret.append(msgpack.unpackb(stat))
+                ret.append(msgpack.unpackb(stat, ext_hook_mapping=msgpack.uuid_to_str))
             else:
                 ret.append(None)
 
@@ -602,7 +602,7 @@ async def _pipe_builder(r: Redis):
         ret = []
         for stat in await redis_helper.execute(ctx.redis_stat, _pipe_builder):
             if stat is not None:
-                ret.append(msgpack.unpackb(stat))
+                ret.append(msgpack.unpackb(stat, ext_hook_mapping=msgpack.uuid_to_str))
             else:
                 ret.append(None)
 

tests/common/test_msgpack.py

@@ -56,6 +56,16 @@ def test_msgpack_uuid_as_map_key():
     assert unpacked[device_id] == 1234
 
 
+def test_msgpack_uuid_to_str():
+    device_id = uuid.uuid4()
+    str_device_id = str(device_id)
+    data = {device_id: 1234}
+    packed = msgpack.packb(data)
+    unpacked = msgpack.unpackb(packed, ext_hook_mapping=msgpack.uuid_to_str)
+    assert isinstance(next(iter(unpacked.keys())), str)
+    assert unpacked[str_device_id] == 1234
+
+
 def test_msgpack_datetime():
     now = datetime.now(tzutc())
     data = {"timestamp": now}