split gem api from deployment service and sink ray operations to ray service

Author: FangwenDave

rock/admin/core/ray_service.py

@@ -1,3 +1,4 @@
+import asyncio
 import ray
 import time
 
@@ -34,6 +35,31 @@ def increment_ray_request_count(self):
 
     def get_ray_rwlock(self):
         return self._ray_rwlock
+
+    async def async_ray_get_actor(self, sandbox_id: str):
+        """Async wrapper for ray.get_actor() using asyncio.to_thread for non-blocking execution."""
+        self.increment_ray_request_count()
+        try:
+            result = await asyncio.to_thread(ray.get_actor, sandbox_id, namespace=self._config.namespace)
+        except ValueError as e:
+            logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e)
+            raise e
+        except Exception as e:
+            logger.error("ray get actor failed", exc_info=e)
+            error_msg = str(e.args[0]) if len(e.args) > 0 else f"ray get actor failed, {str(e)}"
+            raise Exception(error_msg)
+        return result
+
+    async def async_ray_get(self, ray_future: ray.ObjectRef):
+        """Async wrapper for ray.get() using asyncio.to_thread for non-blocking execution."""
+        self.increment_ray_request_count()
+        try:
+            result = await asyncio.to_thread(ray.get, ray_future, timeout=60)
+        except Exception as e:
+            logger.error("ray get failed", exc_info=e)
+            error_msg = str(e.args[0]) if len(e.args) > 0 else f"ray get failed, {str(e)}"
+            raise Exception(error_msg)
+        return result
 
 
     def _setup_ray_reconnect_scheduler(self):

rock/sandbox/gem_manager.py

@@ -16,11 +16,13 @@
 from rock.config import RockConfig
 from rock.deployments.config import DockerDeploymentConfig
 from rock.sandbox.sandbox_manager import SandboxManager
+from rock.sandbox.service.env_service import RayEnvService
 from rock.utils.providers import RedisProvider
 from rock.admin.core.ray_service import RayService
 
 
 class GemManager(SandboxManager):
+    _env_service: RayEnvService
     def __init__(
         self,
         rock_config: RockConfig,
@@ -30,6 +32,7 @@ def __init__(
         enable_runtime_auto_clear: bool = False,
     ):
         super().__init__(rock_config, redis_provider, ray_namespace, ray_service, enable_runtime_auto_clear)
+        self._env_service = RayEnvService(ray_namespace=ray_namespace, ray_service=ray_service)
 
     async def env_make(self, env_id: str) -> EnvMakeResponse:
         config = DockerDeploymentConfig(image=env_vars.ROCK_ENVHUB_DEFAULT_DOCKER_IMAGE)
@@ -51,22 +54,22 @@ async def wait_until_alive(sandbox_id: str, interval: float = 1.0):
         except asyncio.TimeoutError:
             raise Exception("Sandbox startup timeout after 300s")
 
-        make_response = await self._deployment_service.env_make(
+        make_response = await self._env_service.env_make(
             EnvMakeRequest(
                 env_id=env_id,
                 sandbox_id=sandbox_start_response.sandbox_id,
             )
         )
         return make_response
-    
+
     async def env_step(self, request: EnvStepRequest) -> EnvStepResponse:
-        return await self._deployment_service.env_step(request)
+        return await self._env_service.env_step(request)
 
     async def env_reset(self, request: EnvResetRequest) -> EnvResetResponse:
-        return await self._deployment_service.env_reset(request)
+        return await self._env_service.env_reset(request)
 
     async def env_close(self, request: EnvCloseRequest) -> EnvCloseResponse:
-        return await self._deployment_service.env_close(request)
+        return await self._env_service.env_close(request)
 
     async def env_list(self, sandbox_id: str) -> EnvListResponse:
-        return await self._deployment_service.env_list(sandbox_id)
+        return await self._env_service.env_list(sandbox_id)

rock/sandbox/service/deployment_service.py

@@ -1,7 +1,5 @@
 from abc import abstractmethod
-import asyncio
-from rock.actions.envs.request import EnvCloseRequest, EnvMakeRequest, EnvResetRequest, EnvStepRequest
-from rock.actions.envs.response import EnvCloseResponse, EnvListResponse, EnvMakeResponse, EnvResetResponse, EnvStepResponse
+
 from rock.actions.sandbox.response import CommandResponse, State, SystemResourceMetrics
 from rock.actions.sandbox.sandbox_info import SandboxInfo
 from rock.admin.core.ray_service import RayService
@@ -50,26 +48,6 @@ async def get_sandbox_statistics(self, sandbox_id: str):
     async def commit(self, sandbox_id: str, image_tag: str, username: str, password: str) -> CommandResponse:
         ...
 
-    @abstractmethod
-    async def env_step(self, *args, **kwargs):
-        ...
-
-    @abstractmethod
-    async def env_make(self, *args, **kwargs):
-        ...
-    
-    @abstractmethod
-    async def env_reset(self, *args, **kwargs):
-        ...
-
-    @abstractmethod
-    async def env_list(self, *args, **kwargs):
-        ...
-
-    @abstractmethod
-    async def env_close(self, *args, **kwargs):
-        ...
-
     @abstractmethod
     async def collect_system_resource_metrics(self) -> SystemResourceMetrics:
         ...
@@ -85,36 +63,10 @@ def _get_actor_name(self, sandbox_id):
 
     async def is_deployment_alive(self, sandbox_id) -> bool:
         try:
-            actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
+            actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
         except ValueError:
             return False
-        return await self.async_ray_get(actor.is_alive.remote())
-    
-    async def async_ray_get_actor(self, sandbox_id: str):
-        """Async wrapper for ray.get_actor() using asyncio.to_thread for non-blocking execution."""
-        self._ray_service.increment_ray_request_count()
-        try:
-            actor_name = self._get_actor_name(sandbox_id)
-            result = await asyncio.to_thread(ray.get_actor, actor_name, namespace=self._ray_namespace)
-        except ValueError as e:
-            logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e)
-            raise e
-        except Exception as e:
-            logger.error("ray get actor failed", exc_info=e)
-            error_msg = str(e.args[0]) if len(e.args) > 0 else f"ray get actor failed, {str(e)}"
-            raise Exception(error_msg)
-        return result
-
-    async def async_ray_get(self, ray_future: ray.ObjectRef):
-        """Async wrapper for ray.get() using asyncio.to_thread for non-blocking execution."""
-        self._ray_service.increment_ray_request_count()
-        try:
-            result = await asyncio.to_thread(ray.get, ray_future, timeout=60)
-        except Exception as e:
-            logger.error("ray get failed", exc_info=e)
-            error_msg = str(e.args[0]) if len(e.args) > 0 else f"ray get failed, {str(e)}"
-            raise Exception(error_msg)
-        return result
+        return await self._ray_service.async_ray_get(actor.is_alive.remote())
 
     async def submit(self, config: DockerDeploymentConfig, user_info: dict) -> SandboxInfo:
         async with self._ray_service.get_ray_rwlock().read_lock():
@@ -127,7 +79,7 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict) -> Sandb
             sandbox_actor.set_user_id.remote(user_id)
             sandbox_actor.set_experiment_id.remote(experiment_id)
             sandbox_actor.set_namespace.remote(namespace)
-            sandbox_info: SandboxInfo = await self.async_ray_get(sandbox_actor.sandbox_info.remote())
+            sandbox_info: SandboxInfo = await self._ray_service.async_ray_get(sandbox_actor.sandbox_info.remote())
             sandbox_info["user_id"] = user_id
             sandbox_info["experiment_id"] = experiment_id
             sandbox_info["namespace"] = namespace
@@ -155,78 +107,44 @@ def _generate_actor_options(self, config: DockerDeploymentConfig) -> dict:
 
     async def stop(self, sandbox_id: str):
         async with self._ray_service.get_ray_rwlock().read_lock():
-            actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-            await self.async_ray_get(actor.stop.remote())
+            actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+            await self._ray_service.async_ray_get(actor.stop.remote())
             logger.info(f"run time stop over {sandbox_id}")
             ray.kill(actor)
 
     async def get_status(self, sandbox_id: str) -> SandboxInfo:
         async with self._ray_service.get_ray_rwlock().read_lock():
-            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())
+            actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+            sandbox_info: SandboxInfo = await self._ray_service.async_ray_get(actor.sandbox_info.remote())
+            remote_status: ServiceStatus = await self._ray_service.async_ray_get(actor.get_status.remote())
             sandbox_info["phases"] = remote_status.phases
             sandbox_info["port_mapping"] = remote_status.get_port_mapping()
-            alive = await self.async_ray_get(actor.is_alive.remote())
+            alive = await self._ray_service.async_ray_get(actor.is_alive.remote())
             sandbox_info["alive"] = alive.is_alive
             if alive.is_alive:
                 sandbox_info["state"] = State.RUNNING
             return sandbox_info
 
     async def get_mount(self, sandbox_id: str):
         with self._ray_service.get_ray_rwlock().read_lock():
-            actor = await self.async_ray_get_actor(sandbox_id)
-            result = await self.async_ray_get(actor.get_mount.remote())
+            actor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+            result = await self._ray_service.async_ray_get(actor.get_mount.remote())
             logger.info(f"get_mount: {result}")
             return result
 
     async def get_sandbox_statistics(self, sandbox_id: str):
         async with self._ray_service.get_ray_rwlock().read_lock():
-            actor = await self.async_ray_get_actor(sandbox_id)
-            result = await self.async_ray_get(actor.get_sandbox_statistics.remote())
+            actor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+            result = await self._ray_service.async_ray_get(actor.get_sandbox_statistics.remote())
             logger.info(f"get_sandbox_statistics: {result}")
             return result
 
     async def commit(self, sandbox_id) -> CommandResponse:
         with self._ray_service.get_ray_rwlock().read_lock():
-            actor = await self.async_ray_get_actor(sandbox_id)
-            result = await self.async_ray_get(actor.commit.remote())
+            actor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+            result = await self._ray_service.async_ray_get(actor.commit.remote())
             logger.info(f"commit: {result}")
             return result
-    
-    async def env_step(self, request: EnvStepRequest) -> EnvStepResponse:
-        sandbox_id = request.sandbox_id
-        actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-        result = await self.async_ray_get(actor.env_step.remote(request))
-        logger.info(f"env_step: {result}")
-        return result
-    
-    async def env_make(self, request: EnvMakeRequest) -> EnvMakeResponse:
-        sandbox_id = request.sandbox_id
-        actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-        result = await self.async_ray_get(actor.env_make.remote(request))
-        logger.info(f"env_make: {result}")
-        return result
-
-    async def env_reset(self, request: EnvResetRequest) -> EnvResetResponse:
-        sandbox_id = request.sandbox_id
-        actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-        result = await self.async_ray_get(actor.env_reset.remote(request))
-        logger.info(f"env_reset: {result}")
-        return result
-
-    async def env_close(self, request: EnvCloseRequest) -> EnvCloseResponse:
-        sandbox_id = request.sandbox_id
-        actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-        result = await self.async_ray_get(actor.env_close.remote(request))
-        logger.info(f"env_close: {result}")
-        return result
-
-    async def env_list(self, sandbox_id) -> EnvListResponse:
-        actor: SandboxActor = await self.async_ray_get_actor(sandbox_id)
-        result = await self.async_ray_get(actor.env_list.remote())
-        logger.info(f"env_list: {result}")
-        return result
 
     async def collect_system_resource_metrics(self) -> SystemResourceMetrics:
         """Collect system resource metrics"""

rock/sandbox/service/env_service.py

@@ -0,0 +1,74 @@
+from abc import ABC, abstractmethod
+
+from rock.actions.envs.request import EnvCloseRequest, EnvMakeRequest, EnvResetRequest, EnvStepRequest
+from rock.actions.envs.response import EnvCloseResponse, EnvListResponse, EnvMakeResponse, EnvResetResponse, EnvStepResponse
+from rock.admin.core.ray_service import RayService
+from rock.logger import init_logger
+from rock.sandbox.sandbox_actor import SandboxActor
+
+logger = init_logger(__name__)
+
+
+class AbstractEnvService(ABC):
+    @abstractmethod
+    async def env_step(self, request: EnvStepRequest) -> EnvStepResponse:
+        ...
+
+    @abstractmethod
+    async def env_make(self, request: EnvMakeRequest) -> EnvMakeResponse:
+        ...
+    
+    @abstractmethod
+    async def env_reset(self, request: EnvResetRequest) -> EnvResetResponse:
+        ...
+
+    @abstractmethod
+    async def env_close(self, request: EnvCloseRequest) -> EnvCloseResponse:
+        ...
+    
+    @abstractmethod
+    async def env_list(self, sandbox_id) -> EnvListResponse:
+        ...
+
+
+class RayEnvService(AbstractEnvService):
+    def __init__(self, ray_namespace: str, ray_service: RayService):
+        self._ray_namespace = ray_namespace
+        self._ray_service = ray_service
+
+    def _get_actor_name(self, sandbox_id):
+        return sandbox_id
+
+    async def env_step(self, request: EnvStepRequest) -> EnvStepResponse:
+        sandbox_id = request.sandbox_id
+        actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+        result = await self._ray_service.async_ray_get(actor.env_step.remote(request))
+        logger.info(f"env_step: {result}")
+        return result
+
+    async def env_make(self, request: EnvMakeRequest) -> EnvMakeResponse:
+        sandbox_id = request.sandbox_id
+        actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+        result = await self._ray_service.async_ray_get(actor.env_make.remote(request))
+        logger.info(f"env_make: {result}")
+        return result
+
+    async def env_reset(self, request: EnvResetRequest) -> EnvResetResponse:
+        sandbox_id = request.sandbox_id
+        actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+        result = await self._ray_service.async_ray_get(actor.env_reset.remote(request))
+        logger.info(f"env_reset: {result}")
+        return result
+
+    async def env_close(self, request: EnvCloseRequest) -> EnvCloseResponse:
+        sandbox_id = request.sandbox_id
+        actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+        result = await self._ray_service.async_ray_get(actor.env_close.remote(request))
+        logger.info(f"env_close: {result}")
+        return result
+
+    async def env_list(self, sandbox_id) -> EnvListResponse:
+        actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id))
+        result = await self._ray_service.async_ray_get(actor.env_list.remote())
+        logger.info(f"env_list: {result}")
+        return result

tests/unit/sandbox/service/test_deployment_service.py

@@ -7,7 +7,7 @@
 async def test_get_actor_not_exist_raises_value_error(ray_deployment_service):
     sandbox_id = "unknown"
     with pytest.raises(Exception) as exc_info:
-        await ray_deployment_service.async_ray_get_actor(sandbox_id)
+        await ray_deployment_service._ray_service.async_ray_get_actor(sandbox_id)
     assert exc_info.type == ValueError
 
 

tests/unit/sandbox/test_sandbox_manager.py

@@ -67,7 +67,7 @@ async def test_ray_actor_is_alive(sandbox_manager):
 
     assert await wait_sandbox_instance_alive(sandbox_manager, response.sandbox_id)
 
-    sandbox_actor = await sandbox_manager._deployment_service.async_ray_get_actor(response.sandbox_id)
+    sandbox_actor = await sandbox_manager._deployment_service._ray_service.async_ray_get_actor(response.sandbox_id)
     ray.kill(sandbox_actor)
 
     assert not await sandbox_manager._deployment_service.is_deployment_alive(response.sandbox_id)