fix test case: add rock auth in get status

Author: FangwenDave

rock/actions/sandbox/sandbox_info.py

@@ -20,6 +20,7 @@ class SandboxInfo(TypedDict, total=False):
     create_user_gray_flag: bool
     cpus: float
     memory: str
+    alive: bool
 
 
 class SandboxListItem(SandboxInfo):

rock/sandbox/sandbox_manager.py

@@ -25,6 +25,8 @@
 from rock.utils.format import parse_memory_size
 from rock.utils.providers import RedisProvider
 from rock.admin.core.ray_service import RayService
+from rock.rocklet import __version__ as swe_version
+from rock.sandbox import __version__ as gateway_version
 
 logger = init_logger(__name__)
 
@@ -58,7 +60,6 @@ async def submit(self, config: DeploymentConfig, user_info: dict = {}):
         async with self._ray_service.get_ray_rwlock().read_lock():
             deployment_config: DeploymentConfig = await self.deployment_manager.init_config(config)
             sandbox_id = deployment_config.container_name
-            # deployment: AbstractDeployment = deployment_config.get_deployment()
             self.validate_sandbox_spec(self.rock_config.runtime, config)
             self._sandbox_meta[sandbox_id] = {"image": deployment_config.image}
             sandbox_info: SandboxInfo = await self._deployment_service.submit(deployment_config, user_info)
@@ -121,26 +122,46 @@ async def _clear_redis_keys(self, sandbox_id):
     @monitor_sandbox_operation()
     async def get_status(self, sandbox_id) -> SandboxStatusResponse:
         async with self._ray_service.get_ray_rwlock().read_lock():
-            response: SandboxStatusResponse = await self._deployment_service.get_status(sandbox_id)
-            sandbox_info: SandboxInfo = self.get_info_from_response(response)
+            deployment_info: SandboxInfo = await self._deployment_service.get_status(sandbox_id)
+            sandbox_info: SandboxInfo = None
             if self._redis_provider:
+                sandbox_info = await self.build_sandbox_info_from_redis(sandbox_id)
+                if sandbox_info is None:
+                    sandbox_info = deployment_info
+                else:
+                    sandbox_info["state"] = deployment_info.get("state")
                 await self._redis_provider.json_set(alive_sandbox_key(sandbox_id), "$", sandbox_info)
                 await self._update_expire_time(sandbox_id)
-                # logger.info(f"sandbox {sandbox_id} status is {remote_status}, write to redis")
-            return response
-
-    def get_info_from_response(self, response: SandboxStatusResponse) -> SandboxInfo:
-        return SandboxInfo(
-            host_name=response.host_name,
-            host_ip=response.host_ip,
-            user_id=response.user_id,
-            experiment_id=response.experiment_id,
-            namespace=response.namespace,
-            sandbox_id=response.sandbox_id,
-            cpus=response.cpus,
-            memory=response.memory,
-            port_mapping=response.port_mapping,
-        )
+                remote_info = {k: v for k, v in deployment_info.items() if k in ['status', 'port_mapping', 'alive']}
+                sandbox_info.update(remote_info)
+                logger.info(f"sandbox {sandbox_id} status is {sandbox_info}, write to redis")
+            else:
+                sandbox_info = deployment_info
+
+            return SandboxStatusResponse(
+                sandbox_id=sandbox_id,
+                status=sandbox_info.get("status"),
+                state=sandbox_info.get("state"),
+                port_mapping=sandbox_info.get("port_mapping"),
+                host_name=sandbox_info.get("host_name"),
+                host_ip=sandbox_info.get("host_ip"),
+                is_alive=sandbox_info.get("alive"),
+                image=sandbox_info.get("image"),
+                swe_rex_version=swe_version,
+                gateway_version=gateway_version,
+                user_id=sandbox_info.get("user_id"),
+                experiment_id=sandbox_info.get("experiment_id"),
+                namespace=sandbox_info.get("namespace"),
+                cpus=sandbox_info.get("cpus"),
+                memory=sandbox_info.get("memory"),
+            )
+
+    async def build_sandbox_info_from_redis(self, sandbox_id: str) -> SandboxInfo | None:
+        if self._redis_provider:
+            sandbox_status = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$")
+            if sandbox_status and len(sandbox_status) > 0:
+                return sandbox_status[0]
+        return None
 
     async def create_session(self, request: CreateSessionRequest) -> CreateBashSessionResponse:
         sandbox_actor = await self._deployment_service.async_ray_get_actor(request.sandbox_id)

rock/sandbox/service/deployment_service.py

@@ -14,8 +14,6 @@
 from rock.sandbox.sandbox_actor import SandboxActor
 from rock.sdk.common.exceptions import BadRequestRockError
 from rock.utils.format import parse_memory_size
-from rock.rocklet import __version__ as swe_version
-from rock.sandbox import __version__ as gateway_version
 
 logger = init_logger(__name__)
 
@@ -26,12 +24,12 @@ async def get_deployment(self, sandbox_id: str) -> AbstractDeployment:
         ...
 
     @abstractmethod
-    async def submit(self, config: DeploymentConfig, user_info: dict) -> SandboxStartResponse:
+    async def submit(self, config: DeploymentConfig, user_info: dict) -> SandboxInfo:
         """Get status of sandbox."""
         ...
 
     @abstractmethod
-    async def get_status(self, *args, **kwargs) -> SandboxStatusResponse:
+    async def get_status(self, *args, **kwargs) -> SandboxInfo:
         """Get status of sandbox."""
         ...
 
@@ -129,30 +127,17 @@ async def stop(self, sandbox_id: str):
         logger.info(f"run time stop over {sandbox_id}")
         ray.kill(actor)
 
-    async def get_status(self, sandbox_id: str) -> SandboxStatusResponse:
+    async def get_status(self, sandbox_id: str) -> SandboxInfo:
         actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
         sandbox_info: SandboxInfo = await self.async_ray_get(actor.sandbox_info.remote())
         remote_status: ServiceStatus = await self.async_ray_get(actor.get_status.remote())
+        sandbox_info["status"] = remote_status.phases
+        sandbox_info["port_mapping"] = remote_status.get_port_mapping()
         alive = await self.async_ray_get(actor.is_alive.remote())
+        sandbox_info["alive"] = alive.is_alive
         if alive.is_alive:
             sandbox_info["state"] = State.RUNNING
-        return SandboxStatusResponse(
-            sandbox_id=sandbox_id,
-            status=remote_status.phases,
-            port_mapping=remote_status.get_port_mapping(),
-            host_name=sandbox_info.get("host_name"),
-            host_ip=sandbox_info.get("host_ip"),
-            is_alive=alive.is_alive,
-            image=sandbox_info.get("image"),
-            swe_rex_version=swe_version,
-            gateway_version=gateway_version,
-            user_id=sandbox_info.get("user_id"),
-            experiment_id=sandbox_info.get("experiment_id"),
-            namespace=sandbox_info.get("namespace"),
-            cpus=sandbox_info.get("cpus"),
-            memory=sandbox_info.get("memory"),
-            state=sandbox_info.get("state"),
-        )
+        return sandbox_info
 
     async def get_mount(self, sandbox_id: str):
         actor = await self.async_ray_get_actor(sandbox_id)