refactor: move non-management methods back to proxy service (#433)

Author: FangwenDave

pyproject.toml

@@ -46,7 +46,8 @@ admin = [
     "boto3",
     "ray[default]==2.43.0",
     "pip",
-    "cryptography==39.0.1"
+    "cryptography==39.0.1",
+    "fakeredis[json]",
 ]
 
 rocklet = [
@@ -111,7 +112,6 @@ test = [
     "pytest-trio",
     "pytest-twisted",
     "pytest-env",
-    "fakeredis[json]",
 ]
 
 [tool.setuptools.packages.find]

rock/admin/entrypoints/sandbox_api.py

@@ -13,6 +13,7 @@
     UploadResponse,
     WriteFileResponse,
 )
+from rock.actions.response import ResponseStatus
 from rock.admin.proto.request import (
     SandboxBashAction,
     SandboxCloseBashSessionRequest,
@@ -103,7 +104,10 @@ async def create_session(request: SandboxCreateBashSessionRequest) -> RockRespon
 @sandbox_router.post("/run_in_session")
 @handle_exceptions(error_message="run in session failed")
 async def run(action: SandboxBashAction) -> RockResponse[BashObservation]:
-    return RockResponse(result=await sandbox_manager.run_in_session(action))
+    result = await sandbox_manager.run_in_session(action)
+    if result.exit_code is not None and result.exit_code == -1:
+        return RockResponse(status=ResponseStatus.FAILED, error=result.failure_reason)
+    return RockResponse(result=result)
 
 
 @sandbox_router.post("/close_session")

rock/admin/main.py

@@ -54,6 +54,11 @@ async def lifespan(app: FastAPI):
     # init redis provider
     if args.env == "local":
         redis_provider = None
+    elif args.env == "test":
+        from fakeredis import aioredis
+
+        redis_provider = RedisProvider(host=None, port=None, password="")
+        redis_provider.client = aioredis.FakeRedis(decode_responses=True)
     else:
         redis_provider = RedisProvider(
             host=rock_config.redis.host,

rock/sandbox/sandbox_manager.py

@@ -37,6 +37,7 @@
 from rock.sandbox.base_manager import BaseManager
 from rock.sandbox.operator.abstract import AbstractOperator
 from rock.sandbox.sandbox_actor import SandboxActor
+from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService
 from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError
 from rock.utils import (
     EAGLE_EYE_TRACE_ID,
@@ -71,6 +72,7 @@ def __init__(
         self._ray_namespace = ray_namespace
         self._operator = operator
         self._aes_encrypter = AESEncryption()
+        self._proxy_service = SandboxProxyService(rock_config=rock_config, redis_provider=redis_provider)
         if redis_provider:
             self._operator.set_redis_provider(redis_provider)
         logger.info("sandbox service init success")
@@ -318,63 +320,28 @@ async def get_remote_status(self, sandbox_id: str, host_ip: str) -> ServiceStatu
         return ServiceStatus()
 
     async def create_session(self, request: CreateSessionRequest) -> CreateBashSessionResponse:
-        actor_name = self.deployment_manager.get_actor_name(request.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {request.sandbox_id} not found to create session")
-        await self._update_expire_time(request.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.create_session.remote(request))
+        return await self._proxy_service.create_session(request)
 
     @monitor_sandbox_operation()
     async def run_in_session(self, action: Action) -> BashObservation:
-        actor_name = self.deployment_manager.get_actor_name(action.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {action.sandbox_id} not found to run in session")
-        await self._update_expire_time(action.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.run_in_session.remote(action))
+        return await self._proxy_service.run_in_session(action)
 
     async def close_session(self, request: CloseBashSessionRequest) -> CloseBashSessionResponse:
-        actor_name = self.deployment_manager.get_actor_name(request.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {request.sandbox_id} not found to close session")
-        await self._update_expire_time(request.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.close_session.remote(request))
+        return await self._proxy_service.close_session(request)
 
     async def execute(self, command: Command) -> CommandResponse:
-        actor_name = self.deployment_manager.get_actor_name(command.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {command.sandbox_id} not found to execute")
-        await self._update_expire_time(command.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.execute.remote(command))
+        return await self._proxy_service.execute(command)
 
     async def read_file(self, request: ReadFileRequest) -> ReadFileResponse:
-        actor_name = self.deployment_manager.get_actor_name(request.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {request.sandbox_id} not found to read file")
-        await self._update_expire_time(request.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.read_file.remote(request))
+        return await self._proxy_service.read_file(request)
 
     @monitor_sandbox_operation()
     async def write_file(self, request: WriteFileRequest) -> WriteFileResponse:
-        actor_name = self.deployment_manager.get_actor_name(request.sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {request.sandbox_id} not found to write file")
-        await self._update_expire_time(request.sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.write_file.remote(request))
+        return await self._proxy_service.write_file(request)
 
     @monitor_sandbox_operation()
     async def upload(self, file: UploadFile, target_path: str, sandbox_id: str) -> UploadResponse:
-        actor_name = self.deployment_manager.get_actor_name(sandbox_id)
-        sandbox_actor = await self._ray_service.async_ray_get_actor(actor_name, self._ray_namespace)
-        if sandbox_actor is None:
-            raise Exception(f"sandbox {sandbox_id} not found to upload file")
-        await self._update_expire_time(sandbox_id)
-        return await self._ray_service.async_ray_get(sandbox_actor.upload.remote(file, target_path))
+        return await self._proxy_service.upload(file, target_path, sandbox_id)
 
     async def _is_expired(self, sandbox_id):
         timeout_dict = await self._redis_provider.json_get(timeout_sandbox_key(sandbox_id), "$")

tests/integration/conftest.py

@@ -106,7 +106,7 @@ def admin_remote_server():
         [
             "admin",
             "--env",
-            "local",
+            "test",
             "--role",
             "admin",
             "--port",