Feature/sandbox_info #190

rock/actions/sandbox/request.py

@@ -74,3 +74,6 @@ class UploadRequest(BaseModel):
 
     target_path: str
     """Remote file path to upload to."""
+
+    oss_threshold: int = 1 * 1024 * 1024
+    """File size threshold (in bytes) for using OSS upload. Default is 1 MB."""

rock/actions/sandbox/response.py

@@ -2,23 +2,25 @@
 
 from pydantic import BaseModel, Field
 
+from rock._codes import codes
 
-class IsAliveResponse(BaseModel):
-    """Response to the is_alive request.
 
-    You can test the result with bool().
-    """
+class SandboxResult(BaseModel):
+    code: codes | None = None
+    # 向前兼容
+    exit_code: int | None = None
+    failure_reason: str | None = None
+    message: str | None = None
 
-    is_alive: bool
 
-    message: str = ""
-    """Error message if is_alive is False."""
+class IsAliveResponse(SandboxResult):
+    is_alive: bool
 
     def __bool__(self) -> bool:
         return self.is_alive
 
 
-class SandboxStatusResponse(BaseModel):
+class SandboxStatusResponse(SandboxResult):
     sandbox_id: str = None
     status: dict = None
     port_mapping: dict = None
@@ -30,30 +32,28 @@ class SandboxStatusResponse(BaseModel):
     swe_rex_version: str | None = None
     user_id: str | None = None
     experiment_id: str | None = None
+    cpus: float | None = None
+    memory: str | None = None
 
 
-class CommandResponse(BaseModel):
+class CommandResponse(SandboxResult):
     stdout: str = ""
     stderr: str = ""
-    exit_code: int | None = None
 
 
-class WriteFileResponse(BaseModel):
+class WriteFileResponse(SandboxResult):
     success: bool = False
-    message: str = ""
 
 
-class OssSetupResponse(BaseModel):
+class OssSetupResponse(SandboxResult):
     success: bool = False
-    message: str = ""
 
 
-class ExecuteBashSessionResponse(BaseModel):
+class ExecuteBashSessionResponse(SandboxResult):
     success: bool = False
-    message: str = ""
 
 
-class CreateBashSessionResponse(BaseModel):
+class CreateBashSessionResponse(SandboxResult):
     output: str = ""
 
     session_type: Literal["bash"] = "bash"
@@ -63,40 +63,38 @@ class CreateBashSessionResponse(BaseModel):
 """Union type for all create session responses. Do not use this directly."""
 
 
-class BashObservation(BaseModel):
+class BashObservation(SandboxResult):
     session_type: Literal["bash"] = "bash"
     output: str = ""
-    exit_code: int | None = None
     failure_reason: str = ""
     expect_string: str = ""
 
 
 Observation = BashObservation
 
 
-class CloseBashSessionResponse(BaseModel):
+class CloseBashSessionResponse(SandboxResult):
     session_type: Literal["bash"] = "bash"
 
 
 CloseSessionResponse = Annotated[CloseBashSessionResponse, Field(discriminator="session_type")]
 """Union type for all close session responses. Do not use this directly."""
 
 
-class ReadFileResponse(BaseModel):
+class ReadFileResponse(SandboxResult):
     content: str = ""
     """Content of the file as a string."""
 
 
-class UploadResponse(BaseModel):
+class UploadResponse(SandboxResult):
     success: bool = False
-    message: str = ""
     file_name: str = ""
 
 
 FileUploadResponse = UploadResponse
 
 
-class CloseResponse(BaseModel):
+class CloseResponse(SandboxResult):
     """Response for close operations."""
 
     pass

rock/admin/proto/response.py

@@ -19,3 +19,5 @@ class SandboxStatusResponse(BaseModel):
     swe_rex_version: str | None = None
     user_id: str | None = None
     experiment_id: str | None = None
+    cpus: float | None = None
+    memory: str | None = None

rock/config.py

@@ -74,6 +74,14 @@ class RuntimeConfig:
     python_env_path: str = field(default_factory=lambda: env_vars.ROCK_PYTHON_ENV_PATH)
     envhub_db_url: str = field(default_factory=lambda: env_vars.ROCK_ENVHUB_DB_URL)
 
+    @dataclass
+    class StandardSpec:
+        memory: str = "8g"
+        cpus: int = 2
+
+    # standard_spec: StandardSpec = field(default_factory=lambda: StandardSpec())
+    # max_allowed_spec: StandardSpec = field(default_factory=lambda: StandardSpec(cpus=16, memory="64g"))
+
     def __post_init__(self) -> None:
         if not self.python_env_path:
             raise Exception(

rock/deployments/ray.py

@@ -4,6 +4,8 @@
 from rock.deployments.docker import DockerDeployment
 from rock.logger import init_logger
 from rock.sandbox.sandbox_actor import SandboxActor
+from rock.sdk.common.exceptions import BadRequestRockError, InvalidParameterRockException
+from rock.utils.format import parse_memory_size
 
 logger = init_logger(__name__)
 
@@ -21,15 +23,23 @@ async def creator_actor(self, actor_name: str):
 
     async def _create_sandbox_actor(self, actor_name: str):
         """Create sandbox actor instance"""
-        if self.config.actor_resource and self.config.actor_resource_num:
-            sandbox_actor = SandboxActor.options(
-                name=actor_name,
-                resources={self.config.actor_resource: self.config.actor_resource_num},
-                lifetime="detached",
-            ).remote(self._config, self)
-        else:
-            sandbox_actor = SandboxActor.options(
-                name=actor_name,
-                lifetime="detached",
-            ).remote(self._config, self)
-        return sandbox_actor
+        actor_options = {"name": actor_name, "lifetime": "detached"}
+        try:
+            memory = parse_memory_size(self._config.memory)
+            # TODO: refine max allowed spec check
+            max_memory = parse_memory_size("8g")
+            if self._config.cpus > 16:
+                raise BadRequestRockError(
+                    f"Requested CPUs {self._config.cpus} exceed the maximum allowed {self._config.max_allowed_spec.cpus}"
+                )
+            if memory > max_memory:
+                raise BadRequestRockError(
+                    f"Requested memory {self._config.memory} exceed the maximum allowed {self._config.max_allowed_spec.memory}"
+                )
+            actor_options["num_cpus"] = self._config.cpus
+            actor_options["memory"] = memory
+            sandbox_actor = SandboxActor.options(**actor_options).remote(self.config, self)
+            return sandbox_actor
+        except ValueError as e:
+            logger.warning(f"Invalid memory size: {self._config.memory}", exc_info=e)
+            raise InvalidParameterRockException(f"Invalid memory size: {self._config.memory}")

rock/rocklet/local_sandbox.py

@@ -17,6 +17,7 @@
 import psutil
 from typing_extensions import Self
 
+import rock
 from rock.actions import (
     AbstractSandbox,
     BashObservation,
@@ -419,7 +420,12 @@ async def create_session(self, request: CreateSessionRequest) -> CreateSessionRe
             msg = f"unknown session type: {request!r}"
             raise ValueError(msg)
         self.sessions[request.session] = session
-        return await session.start()
+        try:
+            return await session.start()
+        except Exception as ex:
+            self.command_logger.error(f"[create_session error]:{str(ex)}", exc_info=True)
+            response = CreateSessionResponse(code=rock.codes.COMMAND_ERROR, failure_reason=str(ex), exit_code=1)
+            return response
 
     async def run_in_session(self, action: Action) -> Observation:
         """Runs a command in a session."""

rock/sdk/common/exceptions.py

@@ -1,12 +1,12 @@
-import rock
+from rock._codes import codes
 from rock.actions.response import RockResponse
 from rock.utils.deprecated import deprecated
 
 
 class RockException(Exception):
-    _code: rock.codes = None
+    _code: codes = None
 
-    def __init__(self, message, code: rock.codes = None):
+    def __init__(self, message, code: codes = None):
         super().__init__(message)
         self._code = code
 
@@ -22,29 +22,29 @@ def __init__(self, message):
 
 
 class BadRequestRockError(RockException):
-    def __init__(self, message, code: rock.codes = rock.codes.BAD_REQUEST):
+    def __init__(self, message, code: codes = codes.BAD_REQUEST):
         super().__init__(message, code)
 
 
 class InternalServerRockError(RockException):
-    def __init__(self, message, code: rock.codes = rock.codes.INTERNAL_SERVER_ERROR):
+    def __init__(self, message, code: codes = codes.INTERNAL_SERVER_ERROR):
         super().__init__(message, code)
 
 
 class CommandRockError(RockException):
-    def __init__(self, message, code: rock.codes = rock.codes.COMMAND_ERROR):
+    def __init__(self, message, code: codes = codes.COMMAND_ERROR):
         super().__init__(message, code)
 
 
-def raise_for_code(code: rock.codes, message: str):
-    if code is None or rock.codes.is_success(code):
+def raise_for_code(code: codes, message: str):
+    if code is None or codes.is_success(code):
         return
 
-    if rock.codes.is_client_error(code):
+    if codes.is_client_error(code):
         raise BadRequestRockError(message)
-    if rock.codes.is_server_error(code):
+    if codes.is_server_error(code):
         raise InternalServerRockError(message)
-    if rock.codes.is_command_error(code):
+    if codes.is_command_error(code):
         raise CommandRockError(message)
 
     raise RockException(message, code=code)