Use protobuf JSON formatting for worker metadata

Author: dlwh

lib/marin/src/marin/mcp/babysitter.py

@@ -11,6 +11,7 @@
 from pathlib import Path
 from typing import Any
 
+from google.protobuf import json_format
 from iris.cli.bug_report import gather_bug_report
 from iris.cli.job import build_job_summary
 from iris.cli.token_store import cluster_name_from_url, load_any_token, load_token
@@ -30,6 +31,7 @@
 DEFAULT_PROFILE_SECONDS = 1
 MAX_LIST_JOBS_PAGE_SIZE = 500
 DEFAULT_LIST_JOBS_LIMIT = 100
+_PROTO_TO_DICT_OPTIONS = dict(preserving_proto_field_name=True)
 
 _ZEPHYR_PROGRESS_RE = re.compile(
     r"\[(?P<stage>[^\]]+)\]\s+"
@@ -183,16 +185,12 @@ def job_status_to_json(job: job_pb2.JobStatus, tasks: Iterable[job_pb2.TaskStatu
     }
 
 
-def _attribute_value_to_json(value) -> Any:
-    kind = value.WhichOneof("value")
-    if kind is None:
-        return None
-    return getattr(value, kind)
+def _worker_metadata_to_json(metadata: job_pb2.WorkerMetadata) -> dict[str, Any]:
+    return json_format.MessageToDict(metadata, **_PROTO_TO_DICT_OPTIONS)
 
 
 def worker_status_to_json(worker: controller_pb2.Controller.WorkerHealthStatus) -> dict[str, Any]:
     """Serialize Iris worker health into stable JSON."""
-    metadata = worker.metadata
     return {
         "worker_id": worker.worker_id,
         "healthy": bool(worker.healthy),
@@ -201,23 +199,7 @@ def worker_status_to_json(worker: controller_pb2.Controller.WorkerHealthStatus)
         "running_job_ids": list(worker.running_job_ids),
         "address": worker.address,
         "status_message": worker.status_message,
-        "metadata": {
-            "hostname": metadata.hostname,
-            "ip_address": metadata.ip_address,
-            "cpu_count": int(metadata.cpu_count),
-            "memory_bytes": int(metadata.memory_bytes),
-            "disk_bytes": int(metadata.disk_bytes),
-            "device": _device_config_to_json(metadata.device) if metadata.HasField("device") else _cpu_device_json(),
-            "tpu_name": metadata.tpu_name,
-            "tpu_worker_id": metadata.tpu_worker_id,
-            "gpu_count": int(metadata.gpu_count),
-            "gpu_name": metadata.gpu_name,
-            "gpu_memory_mb": int(metadata.gpu_memory_mb),
-            "gce_instance_name": metadata.gce_instance_name,
-            "gce_zone": metadata.gce_zone,
-            "git_hash": metadata.git_hash,
-            "attributes": {key: _attribute_value_to_json(value) for key, value in metadata.attributes.items()},
-        },
+        "metadata": _worker_metadata_to_json(worker.metadata),
     }
 
 

tests/mcp/test_babysitter.py

@@ -12,6 +12,7 @@
     classify_diagnosis,
     parse_zephyr_progress,
     task_status_to_json,
+    worker_status_to_json,
 )
 
 
@@ -91,6 +92,40 @@ def test_task_status_json_includes_attempts_timestamps_and_usage():
     assert payload["attempts"][1]["exit_code"] == 137
 
 
+def test_worker_status_json_uses_protobuf_metadata_serializer():
+    worker = controller_pb2.Controller.WorkerHealthStatus(
+        worker_id="worker-a",
+        healthy=True,
+        consecutive_failures=1,
+        last_heartbeat=_timestamp(12_345),
+        running_job_ids=["/alice/train"],
+        address="worker-a:10001",
+        status_message="healthy",
+        metadata=job_pb2.WorkerMetadata(
+            hostname="worker-a",
+            memory_bytes=123_456_789_012,
+            device=job_pb2.DeviceConfig(gpu=job_pb2.GpuDevice(variant="H100", count=8)),
+            tpu_worker_hostnames="worker-a,worker-b",
+            attributes={
+                "region": job_pb2.AttributeValue(string_value="us-central1"),
+                "slots": job_pb2.AttributeValue(int_value=4),
+            },
+        ),
+    )
+
+    payload = worker_status_to_json(worker)
+
+    assert payload["last_heartbeat_ms"] == 12_345
+    assert payload["metadata"]["hostname"] == "worker-a"
+    assert payload["metadata"]["memory_bytes"] == "123456789012"
+    assert payload["metadata"]["device"] == {"gpu": {"variant": "H100", "count": 8}}
+    assert payload["metadata"]["tpu_worker_hostnames"] == "worker-a,worker-b"
+    assert payload["metadata"]["attributes"] == {
+        "region": {"string_value": "us-central1"},
+        "slots": {"int_value": "4"},
+    }
+
+
 def test_job_summary_payload_preserves_summary_task_fields():
     job = job_pb2.JobStatus(
         job_id="/alice/train",