From a4a6e533908a14646bc80d632e4d6c45264c3097 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=87=E5=96=9C?= Date: Mon, 22 Dec 2025 16:53:46 +0800 Subject: [PATCH 1/5] feat: add mem and cpu support to sandbox --- rock/actions/sandbox/request.py | 3 ++ rock/actions/sandbox/response.py | 2 + rock/admin/proto/response.py | 2 + rock/config.py | 8 ++++ rock/deployments/ray.py | 24 +++++----- rock/sdk/sandbox/client.py | 14 ++++-- rock/utils/format.py | 27 +++++++++++ tests/unit/utils/test_format.py | 81 ++++++++++++++++++++++++++++++++ 8 files changed, 144 insertions(+), 17 deletions(-) create mode 100644 rock/utils/format.py create mode 100644 tests/unit/utils/test_format.py diff --git a/rock/actions/sandbox/request.py b/rock/actions/sandbox/request.py index 7c471eb3c..e0781ffc5 100644 --- a/rock/actions/sandbox/request.py +++ b/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.""" diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index b2cd459c3..cbd67e94d 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -30,6 +30,8 @@ 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): diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index 63868c66a..e3c23aa72 100644 --- a/rock/admin/proto/response.py +++ b/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 diff --git a/rock/config.py b/rock/config.py index 7146419d0..16d8ef848 100644 --- a/rock/config.py +++ b/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( diff --git a/rock/deployments/ray.py b/rock/deployments/ray.py index 4039ebee5..5e3a1214c 100644 --- a/rock/deployments/ray.py +++ b/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 InvalidParameterRockException +from rock.utils.format import parse_memory_size logger = init_logger(__name__) @@ -21,15 +23,13 @@ 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: + # TODO: check upper limit from runtime_config + actor_options["num_cpus"] = self._config.cpus + actor_options["memory"] = parse_memory_size(self._config.memory) + sandbox_actor = SandboxActor.options(**actor_options).remote(self.config) + 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}") diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 71eeed6ca..1662491ed 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -38,7 +38,7 @@ WriteFileResponse, ) from rock.sdk.common.constants import PID_PREFIX, PID_SUFFIX, RunModeType -from rock.sdk.common.exceptions import InvalidParameterRockException +from rock.sdk.common.exceptions import InternalServerRockError, InvalidParameterRockException from rock.sdk.sandbox.agent.base import Agent from rock.sdk.sandbox.config import SandboxConfig, SandboxGroupConfig from rock.sdk.sandbox.remote_user import LinuxRemoteUser, RemoteUser @@ -240,7 +240,7 @@ async def create_session(self, create_session_request: CreateBashSessionRequest) logging.debug(f"Create session response: {response}") if "Success" != response.get("status"): - raise Exception(f"Failed to execute command: {response}") + raise InternalServerRockError(f"Failed to execute command: {response}") result: dict = response.get("result") # type: ignore return CreateBashSessionResponse(**result) @@ -539,14 +539,18 @@ def _build_nohup_detached_message( return "\n".join(lines) async def upload(self, request: UploadRequest) -> UploadResponse: - return await self.upload_by_path(file_path=request.source_path, target_path=request.target_path) + return await self.upload_by_path( + file_path=request.source_path, target_path=request.target_path, oss_threshold=request.oss_threshold + ) - async def upload_by_path(self, file_path: str | Path, target_path: str) -> UploadResponse: + async def upload_by_path( + self, file_path: str | Path, target_path: str, oss_threshold: int = 1 * 1024 * 1024 + ) -> UploadResponse: path_str = file_path file_path = Path(file_path) if not file_path.exists(): return UploadResponse(success=False, message=f"File not found: {file_path}") - if env_vars.ROCK_OSS_ENABLE and os.path.getsize(file_path) > 1024 * 1024 * 1: + if env_vars.ROCK_OSS_ENABLE and os.path.getsize(file_path) > oss_threshold: return await self._upload_via_oss(path_str, target_path) url = f"{self._url}/upload" headers = self._build_headers() diff --git a/rock/utils/format.py b/rock/utils/format.py new file mode 100644 index 000000000..9bc4eaaf3 --- /dev/null +++ b/rock/utils/format.py @@ -0,0 +1,27 @@ +def serialize_model(model): + return model.model_dump() if hasattr(model, "model_dump") else model.dict() + + +def parse_memory_size(size_str: str) -> int: + size_str = size_str.strip().lower() + units = { + "b": 1, + "k": 1024, + "kb": 1024, + "m": 1024**2, + "mb": 1024**2, + "g": 1024**3, + "gb": 1024**3, + "t": 1024**4, + "tb": 1024**4, + } + import re + + match = re.match(r"^(\d+(?:\.\d+)?)\s*([a-z]+)?$", size_str) + if not match: + raise ValueError(f"Invalid memory size format: {size_str}") + number = float(match.group(1)) + unit = match.group(2) or "b" + if unit not in units: + raise ValueError(f"Unknown memory unit: {unit}") + return int(number * units[unit]) diff --git a/tests/unit/utils/test_format.py b/tests/unit/utils/test_format.py new file mode 100644 index 000000000..799b39a09 --- /dev/null +++ b/tests/unit/utils/test_format.py @@ -0,0 +1,81 @@ +import pytest + +from rock.utils.format import parse_memory_size + + +def test_bytes_without_unit(): + assert parse_memory_size("100") == 100 + assert parse_memory_size("0") == 0 + assert parse_memory_size("1024") == 1024 + + +def test_bytes_with_b_unit(): + assert parse_memory_size("100b") == 100 + assert parse_memory_size("100B") == 100 + assert parse_memory_size("0b") == 0 + + +def test_kilobytes(): + assert parse_memory_size("1k") == 1024 + assert parse_memory_size("1K") == 1024 + assert parse_memory_size("1kb") == 1024 + assert parse_memory_size("1KB") == 1024 + assert parse_memory_size("2k") == 2048 + + +def test_megabytes(): + assert parse_memory_size("1m") == 1024**2 + assert parse_memory_size("1M") == 1024**2 + assert parse_memory_size("1mb") == 1024**2 + assert parse_memory_size("1MB") == 1024**2 + assert parse_memory_size("2m") == 2 * 1024**2 + + +def test_gigabytes(): + assert parse_memory_size("1g") == 1024**3 + assert parse_memory_size("1G") == 1024**3 + assert parse_memory_size("1gb") == 1024**3 + assert parse_memory_size("1GB") == 1024**3 + assert parse_memory_size("2g") == 2 * 1024**3 + + +def test_terabytes(): + assert parse_memory_size("1t") == 1024**4 + assert parse_memory_size("1T") == 1024**4 + assert parse_memory_size("1tb") == 1024**4 + assert parse_memory_size("1TB") == 1024**4 + + +def test_decimal_values(): + assert parse_memory_size("1.5k") == int(1.5 * 1024) + assert parse_memory_size("2.5m") == int(2.5 * 1024**2) + assert parse_memory_size("0.5g") == int(0.5 * 1024**3) + + +def test_whitespace_handling(): + assert parse_memory_size(" 100 ") == 100 + assert parse_memory_size(" 1k ") == 1024 + assert parse_memory_size("1 k") == 1024 + assert parse_memory_size(" 1 mb ") == 1024**2 + + +def test_invalid_format(): + with pytest.raises(ValueError, match="Invalid memory size format"): + parse_memory_size("abc") + with pytest.raises(ValueError, match="Invalid memory size format"): + parse_memory_size("1.2.3k") + with pytest.raises(ValueError, match="Invalid memory size format"): + parse_memory_size("") + + +def test_unknown_unit(): + with pytest.raises(ValueError, match="Unknown memory unit"): + parse_memory_size("100x") + with pytest.raises(ValueError, match="Unknown memory unit"): + parse_memory_size("100pb") + + +def test_edge_cases(): + assert parse_memory_size("0.0") == 0 + assert parse_memory_size("0.0k") == 0 + assert parse_memory_size("1000") == 1000 From 0821a1805a0f1c9a42a8178afd87ae0b354c8c3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=87=E5=96=9C?= Date: Mon, 22 Dec 2025 17:31:46 +0800 Subject: [PATCH 2/5] fix: add spec check --- rock/actions/response.py | 3 ++ rock/deployments/ray.py | 16 +++++++++-- rock/sdk/sandbox/client.py | 57 +++++++++++++++++++++++++++----------- rock/utils/exception.py | 10 +++++++ 4 files changed, 67 insertions(+), 19 deletions(-) diff --git a/rock/actions/response.py b/rock/actions/response.py index 3fef5afd2..7cc0a8c06 100644 --- a/rock/actions/response.py +++ b/rock/actions/response.py @@ -3,10 +3,13 @@ from pydantic import BaseModel +import rock + class ResponseStatus(str, Enum): SUCCESS = "Success" FAILED = "Failed" + code: rock.codes | None = None class BaseResponse(BaseModel): diff --git a/rock/deployments/ray.py b/rock/deployments/ray.py index 5e3a1214c..d90d3e1da 100644 --- a/rock/deployments/ray.py +++ b/rock/deployments/ray.py @@ -4,7 +4,7 @@ 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 InvalidParameterRockException +from rock.sdk.common.exceptions import BadRequestRockError, InvalidParameterRockException from rock.utils.format import parse_memory_size logger = init_logger(__name__) @@ -25,9 +25,19 @@ async def _create_sandbox_actor(self, actor_name: str): """Create sandbox actor instance""" actor_options = {"name": actor_name, "lifetime": "detached"} try: - # TODO: check upper limit from runtime_config + 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"] = parse_memory_size(self._config.memory) + actor_options["memory"] = memory sandbox_actor = SandboxActor.options(**actor_options).remote(self.config) return sandbox_actor except ValueError as e: diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 1662491ed..b65c7d7d4 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -13,7 +13,8 @@ from httpx import ReadTimeout from typing_extensions import deprecated -from rock import env_vars +import rock +from rock import env_vars, raise_for_code from rock.actions import ( AbstractSandbox, Action, @@ -135,6 +136,9 @@ async def start(self): logging.debug(f"Start sandbox response: {response}") if "Success" != response.get("status"): + code: rock.codes = response.get("code", None) + if code is not None: + raise_for_code(code, f"Failed to start container: {response}") raise Exception(f"Failed to start sandbox: {response}") self._sandbox_id = response.get("result").get("sandbox_id") self._host_name = response.get("result").get("host_name") @@ -328,9 +332,11 @@ async def arun( ) -> Observation: """ Asynchronously run a command in the sandbox environment. + This method supports two execution modes: - NORMAL: Execute command synchronously and wait for completion - NOHUP: Execute command in background using nohup, suitable for long-running tasks + Args: cmd (str): The command to execute in the sandbox session (str, optional): The session identifier to run the command in. @@ -340,34 +346,53 @@ async def arun( wait_interval (int, optional): Interval in seconds between process completion checks for nohup mode. Minimum value is 5 seconds. Defaults to 10. mode (RunModeType, optional): Execution mode - either "normal" or "nohup". - Defaults to RunMode.NORMAL. + Defaults to RunMode.NORMAL.value. response_limited_bytes_in_nohup (int | None, optional): Maximum bytes to read from nohup output file. If None, reads entire output. Only applies to nohup mode. Defaults to None. nohup_command_timeout (int, optional): Timeout in seconds for the nohup command submission itself. Defaults to 60. + Returns: Observation: Command execution result containing output, exit code, and failure reason if any. - For normal mode: Returns immediate execution result - For nohup mode: Returns result after process completion or timeout + Raises: - InvalidParameterRockException: If an unsupported run mode is provided - ReadTimeout: If command execution times out (nohup mode) - Exception: For other execution failures in nohup mode + BadRequestRockError: If client's invalid parameters are provided. + InternalServerRockError: If ROCK Server errors occur, such as Sandbox timeout, client should retry. + Examples: # Normal synchronous execution result = await sandbox.arun("ls -la") + # Background execution with nohup - result = await sandbox.arun( - "python long_running_script.py", - mode="nohup", - wait_timeout=600 - ) - # Limited output reading in nohup mode - result = await sandbox.arun( - "generate_large_output.sh", - mode="nohup", - response_limited_bytes_in_nohup=1024 - ) + while retry_times < retry_limit: + try: + observation: Observation = await sandbox.arun("python long_running_script.py", mode="nohup") + if observation.exit_code != 0: + logging.warning( + f"Command failed with exit code {observation.exit_code}, " + f"output: {observation.output}, failure_reason: {observation.failure_reason}" + ) + return observation + except RockException as e: + if rock.codes.is_server_error(e.code): + if retry_times >= retry_limit: + logging.error(f"All {retry_limit} attempts failed") + raise e + else: + retry_times += 1 + logging.error( + f"Server error occurred, code: {e.code}, message: {e.code.get_reason_phrase()}, " + f"exception: {str(e)}, will not retry, times: {retry_times}." + ) + await asyncio.sleep(2) + continue + else: + logging.error( + f"Non-retriable error occurred, code: {e.code}, message: {e.code.get_reason_phrase()}, exception: {str(e)}." + ) + raise e """ if mode not in (RunMode.NORMAL, RunMode.NOHUP): raise InvalidParameterRockException(f"Unsupported arun mode: {mode}") diff --git a/rock/utils/exception.py b/rock/utils/exception.py index e6dd93ab2..2e15d9f06 100644 --- a/rock/utils/exception.py +++ b/rock/utils/exception.py @@ -1,6 +1,7 @@ import functools import logging +from rock import RockException from rock.actions import ResponseStatus, RockResponse logger = logging.getLogger(__name__) @@ -21,6 +22,15 @@ def decorator(func): async def wrapper(*args, **kwargs): try: return await func(*args, **kwargs) + except RockException as e: + logging.error(f"RockException in {func.__name__}: {str(e)}", exc_info=True) + return RockException( + status=ResponseStatus.FAILED, + message=error_message, + error=str(e), + code=e.code, + result=None, + ) except Exception as e: logger.error(f"Error in {func.__name__}: {str(e)}", exc_info=True) return RockResponse(status=ResponseStatus.FAILED, message=error_message, error=str(e), result=None) From 3914dd552f8cf6f8d581edc06098fc17a81161f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=87=E5=96=9C?= Date: Mon, 22 Dec 2025 19:54:22 +0800 Subject: [PATCH 3/5] refactor: update imports for consistency across modules Co-developed-by: Aone Copilot --- rock/actions/response.py | 4 ++-- rock/deployments/ray.py | 2 +- rock/sdk/common/exceptions.py | 22 +++++++++++----------- rock/utils/exception.py | 4 ++-- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/rock/actions/response.py b/rock/actions/response.py index 7cc0a8c06..0bea8790f 100644 --- a/rock/actions/response.py +++ b/rock/actions/response.py @@ -3,13 +3,13 @@ from pydantic import BaseModel -import rock +from rock._codes import codes class ResponseStatus(str, Enum): SUCCESS = "Success" FAILED = "Failed" - code: rock.codes | None = None + code: codes | None = None class BaseResponse(BaseModel): diff --git a/rock/deployments/ray.py b/rock/deployments/ray.py index d90d3e1da..52f742b46 100644 --- a/rock/deployments/ray.py +++ b/rock/deployments/ray.py @@ -38,7 +38,7 @@ async def _create_sandbox_actor(self, actor_name: str): ) actor_options["num_cpus"] = self._config.cpus actor_options["memory"] = memory - sandbox_actor = SandboxActor.options(**actor_options).remote(self.config) + 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) diff --git a/rock/sdk/common/exceptions.py b/rock/sdk/common/exceptions.py index 402d8e383..b01db750c 100644 --- a/rock/sdk/common/exceptions.py +++ b/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) diff --git a/rock/utils/exception.py b/rock/utils/exception.py index 2e15d9f06..bab27db4a 100644 --- a/rock/utils/exception.py +++ b/rock/utils/exception.py @@ -1,8 +1,8 @@ import functools import logging -from rock import RockException -from rock.actions import ResponseStatus, RockResponse +from rock.actions.response import ResponseStatus, RockResponse +from rock.sdk.common.exceptions import RockException logger = logging.getLogger(__name__) From 3949afcd86c0a4d56a33eb0908fb67022f192712 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=87=E5=96=9C?= Date: Fri, 26 Dec 2025 14:48:31 +0800 Subject: [PATCH 4/5] refactor: clean up response models and imports Co-developed-by: Aone Copilot --- rock/actions/response.py | 3 --- rock/actions/sandbox/response.py | 45 ++++++++++++++------------------ rock/rocklet/local_sandbox.py | 8 +++++- rock/sdk/sandbox/client.py | 13 ++++++--- 4 files changed, 36 insertions(+), 33 deletions(-) diff --git a/rock/actions/response.py b/rock/actions/response.py index 0bea8790f..3fef5afd2 100644 --- a/rock/actions/response.py +++ b/rock/actions/response.py @@ -3,13 +3,10 @@ from pydantic import BaseModel -from rock._codes import codes - class ResponseStatus(str, Enum): SUCCESS = "Success" FAILED = "Failed" - code: codes | None = None class BaseResponse(BaseModel): diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index cbd67e94d..b52dac3c6 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -2,23 +2,25 @@ from pydantic import BaseModel, Field +import rock -class IsAliveResponse(BaseModel): - """Response to the is_alive request. - You can test the result with bool(). - """ +class SandboxResult(BaseModel): + code: rock.codes | None = None + # 向前兼容 + exit_code: int | None = None + failure_reason: str | None = None + message: str = "" - 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 @@ -34,28 +36,24 @@ class SandboxStatusResponse(BaseModel): 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" @@ -65,18 +63,16 @@ 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" @@ -84,21 +80,20 @@ class CloseBashSessionResponse(BaseModel): """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 diff --git a/rock/rocklet/local_sandbox.py b/rock/rocklet/local_sandbox.py index 4177b57a5..b0635b854 100644 --- a/rock/rocklet/local_sandbox.py +++ b/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.""" diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index b65c7d7d4..205153491 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -13,7 +13,6 @@ from httpx import ReadTimeout from typing_extensions import deprecated -import rock from rock import env_vars, raise_for_code from rock.actions import ( AbstractSandbox, @@ -38,6 +37,7 @@ WriteFileRequest, WriteFileResponse, ) +from rock.actions.sandbox.response import SandboxResult from rock.sdk.common.constants import PID_PREFIX, PID_SUFFIX, RunModeType from rock.sdk.common.exceptions import InternalServerRockError, InvalidParameterRockException from rock.sdk.sandbox.agent.base import Agent @@ -136,9 +136,10 @@ async def start(self): logging.debug(f"Start sandbox response: {response}") if "Success" != response.get("status"): - code: rock.codes = response.get("code", None) - if code is not None: - raise_for_code(code, f"Failed to start container: {response}") + result = response.get("result", None) + if result is not None: + rock_result = SandboxResult(**result) + raise_for_code(rock_result.code, f"Failed to start container: {response}") raise Exception(f"Failed to start sandbox: {response}") self._sandbox_id = response.get("result").get("sandbox_id") self._host_name = response.get("result").get("host_name") @@ -244,6 +245,10 @@ async def create_session(self, create_session_request: CreateBashSessionRequest) logging.debug(f"Create session response: {response}") if "Success" != response.get("status"): + result = response.get("result", None) + if result is not None: + rock_result = SandboxResult(**result) + raise_for_code(rock_result.code, f"Failed to create session: {response}") raise InternalServerRockError(f"Failed to execute command: {response}") result: dict = response.get("result") # type: ignore return CreateBashSessionResponse(**result) From fd40960f34e597674c1f6def618450b7f62652bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=87=E5=96=9C?= Date: Fri, 26 Dec 2025 15:20:09 +0800 Subject: [PATCH 5/5] fix: resolve circular import issue in SandboxResult Change import from 'import rock' to 'from rock._codes import codes' to avoid circular dependency when rock module is being initialized. --- rock/actions/sandbox/response.py | 7 ++++--- tests/unit/sdk/test_arun_nohup.py | 3 +-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index b52dac3c6..ed48b1de2 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -2,15 +2,15 @@ from pydantic import BaseModel, Field -import rock +from rock._codes import codes class SandboxResult(BaseModel): - code: rock.codes | None = None + code: codes | None = None # 向前兼容 exit_code: int | None = None failure_reason: str | None = None - message: str = "" + message: str | None = None class IsAliveResponse(SandboxResult): @@ -66,6 +66,7 @@ class CreateBashSessionResponse(SandboxResult): class BashObservation(SandboxResult): session_type: Literal["bash"] = "bash" output: str = "" + failure_reason: str = "" expect_string: str = "" diff --git a/tests/unit/sdk/test_arun_nohup.py b/tests/unit/sdk/test_arun_nohup.py index 28daab44f..723d0171b 100644 --- a/tests/unit/sdk/test_arun_nohup.py +++ b/tests/unit/sdk/test_arun_nohup.py @@ -1,8 +1,8 @@ import types import pytest - from httpx import ReadTimeout + from rock.actions.sandbox.response import Observation from rock.sdk.common.constants import PID_PREFIX, PID_SUFFIX from rock.sdk.sandbox.client import Sandbox @@ -241,4 +241,3 @@ async def fake_wait(self, pid, session, wait_timeout, wait_interval): assert result.exit_code == 0 assert result.output == "full-log" assert any(cmd.startswith("cat ") for cmd in executed_commands) -