add gpu info

Author: zhongwen666

rock/actions/sandbox/response.py

@@ -58,6 +58,8 @@ class SandboxStatusResponse(BaseModel):
     namespace: str | None = None
     cpus: float | None = None
     memory: str | None = None
+    num_gpus: float | None = None
+    accelerator_type: str | None = None
     disk: str | None = None
     disk_limit_rootfs: str | None = Field(default=None, deprecated="Use 'disk' instead")
     state: State | None = None

rock/actions/sandbox/sandbox_info.py

@@ -21,6 +21,8 @@ class SandboxInfo(TypedDict, total=False):
     create_user_gray_flag: bool
     cpus: float
     memory: str
+    num_gpus: float
+    accelerator_type: str
     disk: str
     create_time: str
     start_time: str

rock/admin/proto/response.py

@@ -34,6 +34,8 @@ class SandboxStatusResponse(BaseModel):
     namespace: str | None = None
     cpus: float | None = None
     memory: str | None = None
+    num_gpus: float | None = None
+    accelerator_type: str | None = None
     disk: str | None = None
     disk_limit_rootfs: str | None = Field(default=None, deprecated="Use 'disk' instead")
     start_time: str | None = None
@@ -64,6 +66,8 @@ def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusRespons
             namespace=sandbox_info.get("namespace"),
             cpus=sandbox_info.get("cpus"),
             memory=sandbox_info.get("memory"),
+            num_gpus=sandbox_info.get("num_gpus"),
+            accelerator_type=sandbox_info.get("accelerator_type"),
             disk=sandbox_info.get("disk"),
             disk_limit_rootfs=sandbox_info.get("disk"),
             start_time=sandbox_info.get("start_time"),

rock/sandbox/sandbox_manager.py

@@ -168,6 +168,8 @@ async def start_async(
             "image": docker_deployment_config.image,
             "cpus": docker_deployment_config.cpus,
             "memory": docker_deployment_config.memory,
+            "num_gpus": docker_deployment_config.num_gpus,
+            "accelerator_type": docker_deployment_config.accelerator_type,
         }
         if docker_deployment_config.disk is not None:
             sandbox_info["disk"] = docker_deployment_config.disk
@@ -447,6 +449,8 @@ async def get_status(self, sandbox_id, include_all_states: bool = False) -> Sand
             namespace=sandbox_info.get("namespace"),
             cpus=sandbox_info.get("cpus"),
             memory=sandbox_info.get("memory"),
+            num_gpus=sandbox_info.get("num_gpus"),
+            accelerator_type=sandbox_info.get("accelerator_type"),
             disk=sandbox_info.get("disk"),
             disk_limit_rootfs=sandbox_info.get("disk"),
             start_time=sandbox_info.get("start_time"),

rock/sandbox/service/opensandbox_proxy_service.py

@@ -250,6 +250,8 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) ->
             namespace=info.get("namespace"),
             cpus=info.get("cpus"),
             memory=info.get("memory"),
+            num_gpus=info.get("num_gpus"),
+            accelerator_type=info.get("accelerator_type"),
             disk=info.get("disk"),
             disk_limit_rootfs=info.get("disk"),
             start_time=info.get("start_time"),

rock/sandbox/service/sandbox_proxy_service.py

@@ -1135,6 +1135,8 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) ->
             namespace=info.get("namespace"),
             cpus=info.get("cpus"),
             memory=info.get("memory"),
+            num_gpus=info.get("num_gpus"),
+            accelerator_type=info.get("accelerator_type"),
             disk=info.get("disk"),
             disk_limit_rootfs=info.get("disk"),
             start_time=info.get("start_time"),

tests/unit/admin/proto/test_sandbox_response.py

@@ -118,6 +118,18 @@ def test_from_sandbox_info_with_none_disk(self):
         response = SandboxStatusResponse.from_sandbox_info(sandbox_info)
         assert response.disk is None
 
+    def test_from_sandbox_info_with_gpu(self):
+        sandbox_info = {
+            "sandbox_id": "test-sandbox",
+            "num_gpus": 0.5,
+            "accelerator_type": "A100",
+        }
+
+        response = SandboxStatusResponse.from_sandbox_info(sandbox_info)
+
+        assert response.num_gpus == 0.5
+        assert response.accelerator_type == "A100"
+
 
 # ---- actions/sandbox/response.SandboxStatusResponse tests ----
 
@@ -148,3 +160,11 @@ def test_actions_status_response_auto_transition_fields(self):
 
         assert response.auto_archive_time == "2026-01-01T00:10:00+00:00"
         assert response.auto_delete_time == "2026-01-01T00:20:00+00:00"
+
+    def test_actions_status_response_gpu_fields(self):
+        from rock.actions.sandbox.response import SandboxStatusResponse as ActionStatusResponse
+
+        response = ActionStatusResponse(num_gpus=2, accelerator_type="H20")
+
+        assert response.num_gpus == 2
+        assert response.accelerator_type == "H20"

tests/unit/sandbox/test_get_status_include_all_states.py

@@ -119,3 +119,16 @@ async def test_get_status_sets_auto_stop_time_from_timeout(self, sandbox_manager
         result = await sandbox_manager.get_status("sandbox-1")
 
         assert result.auto_stop_time == "2286-11-21T01:46:39+08:00"
+
+    @pytest.mark.asyncio
+    async def test_get_status_returns_gpu_info(self, sandbox_manager, mock_operator, mock_meta_store):
+        sandbox_info = _make_sandbox_info(state=State.RUNNING)
+        sandbox_info["num_gpus"] = 0.5
+        sandbox_info["accelerator_type"] = "A100"
+        mock_meta_store.get = AsyncMock(return_value=sandbox_info)
+        mock_operator.get_status = AsyncMock(return_value=_make_sandbox_info(state=State.RUNNING))
+
+        result = await sandbox_manager.get_status("sandbox-1")
+
+        assert result.num_gpus == 0.5
+        assert result.accelerator_type == "A100"

tests/unit/sandbox/test_proxy_get_status.py

@@ -134,6 +134,18 @@ async def test_running_sandbox_returns_phases_and_port_mapping(
         }
         assert result.port_mapping == {22555: 22555, 8080: 8080}
 
+    async def test_get_status_returns_gpu_info(self, proxy_service, mock_meta_store, mock_rpc_client):
+        info = _make_meta_info(state=State.RUNNING)
+        info.update(num_gpus=2, accelerator_type="H20")
+        mock_meta_store.get.return_value = info
+        mock_rpc_client.post.side_effect = Exception("connection refused")
+        mock_rpc_client.get.side_effect = Exception("connection refused")
+
+        result = await proxy_service.get_status("sandbox-1")
+
+        assert result.num_gpus == 2
+        assert result.accelerator_type == "H20"
+
     async def test_pending_to_running_triggers_meta_update(self, proxy_service, mock_meta_store, mock_rpc_client):
         mock_meta_store.get.return_value = _make_meta_info(state=State.PENDING, host_ip="10.0.0.1")
         mock_rpc_client.post.side_effect = [

tests/unit/sandbox/test_proxy_get_status_route.py

@@ -44,6 +44,8 @@ async def test_get_status_success(self, client, mock_service):
             state=State.RUNNING,
             is_alive=True,
             host_ip="10.0.0.1",
+            num_gpus=2,
+            accelerator_type="H20",
         )
 
         resp = await client.get("/get_status", params={"sandbox_id": "sandbox-1"})
@@ -53,6 +55,8 @@ async def test_get_status_success(self, client, mock_service):
         assert body["result"]["sandbox_id"] == "sandbox-1"
         assert body["result"]["state"] == State.RUNNING
         assert body["result"]["is_alive"] is True
+        assert body["result"]["num_gpus"] == 2
+        assert body["result"]["accelerator_type"] == "H20"
         mock_service.get_status.assert_awaited_once_with("sandbox-1", False)
 
     async def test_get_status_passes_include_all_states(self, client, mock_service):