From 8d44699831814538fbd9d7bfc3b406871cc403a8 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Tue, 31 Mar 2026 16:26:33 +0800 Subject: [PATCH 001/226] Refactor/xinshi/oss env for harbor (#715) * refactor: add method to merge session environment variables * docs: add env vars for harbor demo script * refactor: rename `_check_env` to `check_oss_env` in harbor demo script --- .gitignore | 4 +++- examples/harbor/harbor_demo.py | 36 ++++++++++++++++++++++++++++++ rock/sdk/agent/job.py | 13 ++++++++++- tests/unit/sdk/agent/test_job.py | 38 ++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 7d18ed5b5e..144b8b5ea2 100644 --- a/.gitignore +++ b/.gitignore @@ -19,4 +19,6 @@ node_modules .docusaurus logs -docs/superpowers/ \ No newline at end of file +docs/superpowers/ + +.env \ No newline at end of file diff --git a/examples/harbor/harbor_demo.py b/examples/harbor/harbor_demo.py index 4c75f5d9e6..61240e3324 100644 --- a/examples/harbor/harbor_demo.py +++ b/examples/harbor/harbor_demo.py @@ -11,20 +11,55 @@ Usage: python examples/harbor/harbor_demo.py -c examples/harbor/job_config.yaml python examples/harbor/harbor_demo.py -c examples/harbor/tb_job_config.yaml -t mailman + +Required environment variables (OSS_* are auto-forwarded into the sandbox): + OSS_ACCESS_KEY_ID Alibaba Cloud OSS access key ID + OSS_ACCESS_KEY_SECRET Alibaba Cloud OSS access key secret + OSS_REGION OSS region, e.g. cn-shanghai + OSS_ENDPOINT OSS endpoint, e.g. oss-cn-shanghai.aliyuncs.com + OSS_BUCKET OSS bucket name + OSS_DATASET_PATH Path prefix inside the bucket for datasets + +Recommended setup: + 1. Copy .env.example to .env and fill in your credentials + 2. source .env + 3. python examples/harbor/harbor_demo.py -c ... """ import argparse import asyncio import logging +import os +import sys from rock.sdk.agent import Job, JobConfig +_REQUIRED_ENV_VARS = [ + "OSS_ACCESS_KEY_ID", + "OSS_ACCESS_KEY_SECRET", + "OSS_REGION", + "OSS_ENDPOINT", + "OSS_BUCKET", + "OSS_DATASET_PATH", +] + + logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) # disable httpx logging.getLogger("httpx").setLevel(logging.WARNING) +def check_oss_env() -> None: + missing = [v for v in _REQUIRED_ENV_VARS if not os.environ.get(v)] + if missing: + print("Missing required environment variables:") + for v in missing: + print(f" {v}") + print("\nSet them with `source .env` or export them manually.") + sys.exit(1) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run Harbor tasks inside a ROCK sandbox") parser.add_argument("-c", "--config", required=True, help="Path to JobConfig YAML file") @@ -53,5 +88,6 @@ async def async_main(args: argparse.Namespace) -> None: if __name__ == "__main__": + check_oss_env() args = parse_args() asyncio.run(async_main(args)) diff --git a/rock/sdk/agent/job.py b/rock/sdk/agent/job.py index 0524dc7867..4756f27741 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/agent/job.py @@ -199,6 +199,17 @@ def _render_run_script(self, config_path: str) -> str: # Private: sandbox / session # ------------------------------------------------------------------ + def _build_session_env(self) -> dict[str, str] | None: + """Merge OSS_* vars from the current process env with explicit config env. + + OSS credentials are forwarded from the process environment so that users + do not need to write sensitive values in the YAML config file. + Explicit values in config env always take precedence. + """ + oss_env = {k: v for k, v in os.environ.items() if k.startswith("OSS")} + merged = {**oss_env, **self._config.environment.env} + return merged or None + async def _create_session(self) -> None: """Create a bash session with sandbox_env injected.""" self._session = f"rock-job-{self._config.job_name}" @@ -206,7 +217,7 @@ async def _create_session(self) -> None: CreateBashSessionRequest( session=self._session, env_enable=True, - env=self._config.environment.env or None, + env=self._build_session_env(), ) ) diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index 00924f5c09..a978d0c7ea 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -1,4 +1,5 @@ import json +import os from unittest.mock import AsyncMock, MagicMock, patch from rock.sdk.agent.job import Job, JobResult, JobStatus @@ -225,6 +226,43 @@ async def test_wait_returns_result(self): assert isinstance(result, JobResult) assert result.status == JobStatus.COMPLETED + +class TestBuildSessionEnv: + def test_oss_vars_from_process_env_are_forwarded(self, monkeypatch): + monkeypatch.setenv("OSS_ENDPOINT", "https://oss.example.com") + monkeypatch.setenv("OSS_ACCESS_KEY_ID", "test-key") + monkeypatch.setenv("HOME", "/root") + + job = Job(JobConfig(job_name="test-job")) + env = job._build_session_env() + + assert env["OSS_ENDPOINT"] == "https://oss.example.com" + assert env["OSS_ACCESS_KEY_ID"] == "test-key" + assert "HOME" not in env + + def test_config_env_overrides_process_oss_vars(self, monkeypatch): + monkeypatch.setenv("OSS_ENDPOINT", "https://oss.from.process.com") + + job = Job( + JobConfig( + job_name="test-job", + environment=RockEnvironmentConfig(env={"OSS_ENDPOINT": "https://oss.from.config.com"}), + ) + ) + env = job._build_session_env() + + assert env["OSS_ENDPOINT"] == "https://oss.from.config.com" + + def test_returns_none_when_both_empty(self, monkeypatch): + for key in list(os.environ.keys()): + if key.startswith("OSS"): + monkeypatch.delenv(key) + + job = Job(JobConfig(job_name="test-job")) + assert job._build_session_env() is None + + +class TestCancelKillsProcess: async def test_cancel_kills_process(self): mock_sandbox = _make_mock_sandbox() mock_sandbox.arun = AsyncMock(return_value=MagicMock(output="", exit_code=0)) From df10981ba771b49bcc18b68fb7218f8bdbce97f5 Mon Sep 17 00:00:00 2001 From: lkc Date: Tue, 31 Mar 2026 16:38:02 +0800 Subject: [PATCH 002/226] feat(sdk): Job OSS artifact mirror (OssMirrorConfig) (#707) (#708) * feat(sdk): add OssMirrorConfig and OSS artifact mirror support - Add OssMirrorConfig to EnvironmentConfig for OSS artifact mirroring - Add JobConfig.enable_oss_mirror() convenience method - Add JobConfig.step field - Update Job to call _autofill_sandbox_info() before writing config YAML - Add unit tests in tests/unit/sdk/agent/test_oss_mirror.py fixes #707 * feat: remove step param --- rock/sdk/agent/__init__.py | 4 + rock/sdk/agent/job.py | 4 + rock/sdk/agent/models/__init__.py | 4 + rock/sdk/agent/models/job/config.py | 33 +++ rock/sdk/agent/models/trial/__init__.py | 3 +- rock/sdk/agent/models/trial/config.py | 16 ++ rock/sdk/sandbox/client.py | 8 +- tests/unit/sdk/agent/test_oss_mirror.py | 298 ++++++++++++++++++++++++ 8 files changed, 368 insertions(+), 2 deletions(-) create mode 100644 tests/unit/sdk/agent/test_oss_mirror.py diff --git a/rock/sdk/agent/__init__.py b/rock/sdk/agent/__init__.py index c9ae6eea74..7f034a043c 100644 --- a/rock/sdk/agent/__init__.py +++ b/rock/sdk/agent/__init__.py @@ -14,6 +14,8 @@ from rock.sdk.agent.models.trial.config import ( AgentConfig, ArtifactConfig, + EnvironmentConfig, + OssMirrorConfig, TaskConfig, VerifierConfig, ) @@ -43,6 +45,8 @@ "OrchestratorConfig", "RetryConfig", "AgentConfig", + "EnvironmentConfig", + "OssMirrorConfig", "VerifierConfig", "TaskConfig", "ArtifactConfig", diff --git a/rock/sdk/agent/job.py b/rock/sdk/agent/job.py index 4756f27741..e8997451e9 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/agent/job.py @@ -152,6 +152,7 @@ def _get_wait_timeout(self) -> int: async def _prepare_and_start(self): """Upload files + harbor config YAML + render run script -> nohup start.""" + await self._autofill_sandbox_info() await self._create_session() # 1. Upload user-specified files/dirs @@ -264,6 +265,9 @@ async def _collect_results(self) -> JobResult: # Private: utilities # ------------------------------------------------------------------ + async def _autofill_sandbox_info(self) -> None: + self._config.namespace = self._sandbox._namespace + async def _upload_content(self, content: str, sandbox_path: str) -> None: """Write text content to a local temp file and upload to sandbox via upload_by_path.""" local_tmp = None diff --git a/rock/sdk/agent/models/__init__.py b/rock/sdk/agent/models/__init__.py index 8051315ae1..be3cdffa69 100644 --- a/rock/sdk/agent/models/__init__.py +++ b/rock/sdk/agent/models/__init__.py @@ -12,6 +12,8 @@ from rock.sdk.agent.models.trial.config import ( AgentConfig, ArtifactConfig, + EnvironmentConfig, + OssMirrorConfig, TaskConfig, VerifierConfig, ) @@ -22,6 +24,8 @@ "RetryConfig", "DatasetConfig", "AgentConfig", + "EnvironmentConfig", + "OssMirrorConfig", "RockEnvironmentConfig", "VerifierConfig", "TaskConfig", diff --git a/rock/sdk/agent/models/job/config.py b/rock/sdk/agent/models/job/config.py index 8d714e5428..bae5efe419 100644 --- a/rock/sdk/agent/models/job/config.py +++ b/rock/sdk/agent/models/job/config.py @@ -17,6 +17,8 @@ from rock.sdk.agent.models.trial.config import ( AgentConfig, ArtifactConfig, + EnvironmentConfig, + OssMirrorConfig, RockEnvironmentConfig, TaskConfig, VerifierConfig, @@ -138,6 +140,14 @@ class JobConfig(BaseModel): environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) # ── Harbor native fields ── + namespace: str | None = Field( + default=None, + description="资源租户隔离标识，用于区分不同团队/项目的资源", + ) + experiment_id: str | None = Field( + default=None, + description="实验标识", + ) job_name: str = Field(default_factory=lambda: datetime.now().strftime("%Y-%m-%d__%H-%M-%S")) jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" n_attempts: int = 1 @@ -177,3 +187,26 @@ def from_yaml(cls, path: str) -> JobConfig: with open(path) as f: data = yaml.safe_load(f) return cls(**data) + + def enable_oss_mirror( + self, + *, + oss_bucket: str, + oss_access_key_id: str, + oss_access_key_secret: str, + oss_region: str, + oss_endpoint: str, + ) -> None: + """Enable OSS artifact mirror on the environment config (credentials only). + + Set top-level ``namespace`` / ``experiment_id`` on ``JobConfig`` separately, + or rely on ``Job`` autofill from the sandbox (see ``_autofill_sandbox_info``). + """ + self.environment.oss_mirror = OssMirrorConfig( + enabled=True, + oss_bucket=oss_bucket, + oss_access_key_id=oss_access_key_id, + oss_access_key_secret=oss_access_key_secret, + oss_region=oss_region, + oss_endpoint=oss_endpoint, + ) diff --git a/rock/sdk/agent/models/trial/__init__.py b/rock/sdk/agent/models/trial/__init__.py index f115c9ba71..24932bab46 100644 --- a/rock/sdk/agent/models/trial/__init__.py +++ b/rock/sdk/agent/models/trial/__init__.py @@ -1,9 +1,10 @@ -from .config import AgentConfig, ArtifactConfig, EnvironmentConfig, TaskConfig, VerifierConfig +from .config import AgentConfig, ArtifactConfig, EnvironmentConfig, OssMirrorConfig, TaskConfig, VerifierConfig from .result import AgentInfo, AgentResult, ExceptionInfo, ModelInfo, TimingInfo, TrialResult, VerifierResult __all__ = [ "AgentConfig", "EnvironmentConfig", + "OssMirrorConfig", "VerifierConfig", "TaskConfig", "ArtifactConfig", diff --git a/rock/sdk/agent/models/trial/config.py b/rock/sdk/agent/models/trial/config.py index 72ad299b98..849cf6f47e 100644 --- a/rock/sdk/agent/models/trial/config.py +++ b/rock/sdk/agent/models/trial/config.py @@ -20,6 +20,21 @@ class AgentConfig(BaseModel): env: dict[str, str] = Field(default_factory=dict) +class OssMirrorConfig(BaseModel): + """OSS artifact mirror configuration (credentials and bucket only). + + ``namespace`` / ``experiment_id`` belong on :class:`~rock.sdk.agent.models.job.config.JobConfig` + as top-level Harbor fields, not inside ``oss_mirror``. + """ + + enabled: bool = False + oss_bucket: str | None = None + oss_access_key_id: str | None = None + oss_access_key_secret: str | None = None + oss_region: str | None = None + oss_endpoint: str | None = None + + class EnvironmentConfig(BaseModel): type: EnvironmentType | None = None import_path: str | None = None @@ -31,6 +46,7 @@ class EnvironmentConfig(BaseModel): override_gpus: int | None = None suppress_override_warnings: bool = False mounts_json: list[dict[str, Any]] | None = None + oss_mirror: OssMirrorConfig | None = None env: dict[str, str] = Field(default_factory=dict) kwargs: dict[str, Any] = Field(default_factory=dict) diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index d0b9b71a9f..0e2ed8496c 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -75,6 +75,8 @@ class Sandbox(AbstractSandbox): _host_ip: str | None = None _oss_bucket: oss2.Bucket | None = None _cluster: str | None = None + _namespace: str | None = None + _experiment_id: str | None = None agent: RockAgent | None = None model_service: ModelService | None = None remote_user: RemoteUser | None = None @@ -193,10 +195,14 @@ async def start(self): self._sandbox_id = response.get("result").get("sandbox_id") self._host_name = response.get("result").get("host_name") self._host_ip = response.get("result").get("host_ip") - + start_time = time.time() while time.time() - start_time < self.config.startup_timeout: sandbox_info = await self.get_status() + if sandbox_info.namespace is not None: + self._namespace = sandbox_info.namespace + if sandbox_info.experiment_id is not None: + self._experiment_id = sandbox_info.experiment_id logging.debug(f"Get status response: {sandbox_info}") if sandbox_info.is_alive: return diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py new file mode 100644 index 0000000000..adf4395029 --- /dev/null +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -0,0 +1,298 @@ +"""Tests for OssMirrorConfig integration into SDK. + +Covers: +- OssMirrorConfig model fields and defaults +- EnvironmentConfig.oss_mirror field +- JobConfig top-level namespace/experiment_id fields +- to_harbor_yaml() serialization (namespace at top level) +- from_yaml() deserialization +- enable_oss_mirror() convenience method on JobConfig +""" + +import yaml + +from rock.sdk.agent.models.trial.config import EnvironmentConfig + +# --------------------------------------------------------------------------- +# 1. OssMirrorConfig 模型 +# --------------------------------------------------------------------------- + + +class TestOssMirrorConfig: + def test_importable_from_trial_config(self): + from rock.sdk.agent.models.trial.config import OssMirrorConfig + + assert OssMirrorConfig is not None + + def test_importable_from_agent_package(self): + from rock.sdk.agent import OssMirrorConfig + + assert OssMirrorConfig is not None + + def test_default_is_disabled(self): + from rock.sdk.agent.models.trial.config import OssMirrorConfig + + cfg = OssMirrorConfig() + assert cfg.enabled is False + assert cfg.oss_bucket is None + assert cfg.oss_access_key_id is None + assert cfg.oss_access_key_secret is None + assert cfg.oss_region is None + assert cfg.oss_endpoint is None + + def test_all_fields_settable(self): + from rock.sdk.agent.models.trial.config import OssMirrorConfig + + cfg = OssMirrorConfig( + enabled=True, + oss_bucket="my-bucket", + oss_access_key_id="ak-xxx", + oss_access_key_secret="sk-xxx", + oss_region="cn-hangzhou", + oss_endpoint="oss-cn-hangzhou.aliyuncs.com", + ) + assert cfg.enabled is True + assert cfg.oss_bucket == "my-bucket" + assert cfg.oss_access_key_id == "ak-xxx" + assert cfg.oss_access_key_secret == "sk-xxx" + assert cfg.oss_region == "cn-hangzhou" + assert cfg.oss_endpoint == "oss-cn-hangzhou.aliyuncs.com" + + +# --------------------------------------------------------------------------- +# 2. EnvironmentConfig.oss_mirror 字段 +# --------------------------------------------------------------------------- + + +class TestEnvironmentConfigOssMirror: + def test_default_oss_mirror_is_none(self): + env = EnvironmentConfig() + assert env.oss_mirror is None + + def test_set_oss_mirror(self): + from rock.sdk.agent.models.trial.config import OssMirrorConfig + + mirror = OssMirrorConfig(enabled=True, oss_bucket="b1", oss_region="r1") + env = EnvironmentConfig(oss_mirror=mirror) + assert env.oss_mirror.enabled is True + assert env.oss_mirror.oss_bucket == "b1" + + def test_set_oss_mirror_from_dict(self): + env = EnvironmentConfig( + oss_mirror={ + "enabled": True, + "oss_bucket": "b2", + } + ) + assert env.oss_mirror.enabled is True + assert env.oss_mirror.oss_bucket == "b2" + + +# --------------------------------------------------------------------------- +# 3. JobConfig 顶层 namespace/experiment_id 字段 +# --------------------------------------------------------------------------- + + +class TestJobConfigNamespaceFields: + def test_default_namespace_is_none(self): + from rock.sdk.agent.models.job.config import JobConfig + + cfg = JobConfig(job_name="test") + assert cfg.namespace is None + assert cfg.experiment_id is None + + def test_namespace_settable_at_top_level(self): + from rock.sdk.agent.models.job.config import JobConfig + + cfg = JobConfig(job_name="test", namespace="team-rl", experiment_id="rl-step-42") + assert cfg.namespace == "team-rl" + assert cfg.experiment_id == "rl-step-42" + + +# --------------------------------------------------------------------------- +# 4. to_harbor_yaml() 序列化 +# --------------------------------------------------------------------------- + + +class TestToHarborYamlOssMirror: + def test_namespace_at_top_level_in_yaml(self): + """namespace/experiment_id 序列化为 JobConfig 顶层字段。""" + from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.agent.models.trial.config import OssMirrorConfig + + from rock.sdk.agent.models.trial.config import RockEnvironmentConfig + + cfg = JobConfig( + job_name="mirror-test", + namespace="my-ns", + experiment_id="exp-1", + environment=RockEnvironmentConfig( + oss_mirror=OssMirrorConfig( + enabled=True, + oss_bucket="test-bucket", + oss_access_key_id="ak", + oss_access_key_secret="sk", + oss_region="cn-hangzhou", + oss_endpoint="oss-cn-hangzhou.aliyuncs.com", + ), + ), + ) + data = yaml.safe_load(cfg.to_harbor_yaml()) + + assert data["namespace"] == "my-ns" + assert data["experiment_id"] == "exp-1" + oss = data["environment"]["oss_mirror"] + assert oss["enabled"] is True + assert oss["oss_bucket"] == "test-bucket" + assert "namespace" not in oss + assert "experiment_id" not in oss + + def test_disabled_oss_mirror_excluded_from_yaml(self): + """When oss_mirror is default (disabled), it should not clutter the YAML.""" + from rock.sdk.agent.models.job.config import JobConfig + + cfg = JobConfig(job_name="no-mirror") + data = yaml.safe_load(cfg.to_harbor_yaml()) + + env_data = data.get("environment", {}) + assert "oss_mirror" not in env_data + + +# --------------------------------------------------------------------------- +# 5. from_yaml() 反序列化 +# --------------------------------------------------------------------------- + + +class TestFromYamlOssMirror: + def test_from_yaml_with_top_level_namespace(self, tmp_path): + """新方式：namespace/experiment_id 在顶层。""" + from rock.sdk.agent.models.job.config import JobConfig + + yaml_content = """\ +job_name: loaded-mirror +namespace: yaml-ns +experiment_id: yaml-exp +environment: + oss_mirror: + enabled: true + oss_bucket: yaml-bucket + oss_region: ap-southeast-1 + oss_endpoint: oss-ap-southeast-1.aliyuncs.com +agents: + - name: test-agent +""" + yaml_file = tmp_path / "config.yaml" + yaml_file.write_text(yaml_content) + + cfg = JobConfig.from_yaml(str(yaml_file)) + assert cfg.namespace == "yaml-ns" + assert cfg.experiment_id == "yaml-exp" + assert cfg.environment.oss_mirror.enabled is True + assert cfg.environment.oss_mirror.oss_bucket == "yaml-bucket" + + def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): + """YAML 中 oss_mirror 内多余的 namespace 等字段由 Pydantic 忽略。""" + from rock.sdk.agent.models.job.config import JobConfig + + yaml_content = """\ +job_name: compat-mirror +environment: + oss_mirror: + enabled: true + oss_bucket: yaml-bucket + namespace: legacy-ns + experiment_id: legacy-exp + oss_region: ap-southeast-1 + oss_endpoint: oss-ap-southeast-1.aliyuncs.com +agents: + - name: test-agent +""" + yaml_file = tmp_path / "config.yaml" + yaml_file.write_text(yaml_content) + + cfg = JobConfig.from_yaml(str(yaml_file)) + assert cfg.environment.oss_mirror.enabled is True + assert cfg.environment.oss_mirror.oss_bucket == "yaml-bucket" + dump = cfg.environment.oss_mirror.model_dump(exclude_none=True) + assert "namespace" not in dump + assert "experiment_id" not in dump + + def test_from_yaml_without_oss_mirror(self, tmp_path): + from rock.sdk.agent.models.job.config import JobConfig + + yaml_content = """\ +job_name: no-mirror +agents: + - name: test-agent +""" + yaml_file = tmp_path / "config.yaml" + yaml_file.write_text(yaml_content) + + cfg = JobConfig.from_yaml(str(yaml_file)) + assert cfg.environment.oss_mirror is None + + +# --------------------------------------------------------------------------- +# 6. enable_oss_mirror() 便捷方法 +# --------------------------------------------------------------------------- + + +class TestEnableOssMirror: + def test_enable_with_all_params(self): + from rock.sdk.agent.models.job.config import JobConfig + + cfg = JobConfig(job_name="conv-test") + cfg.enable_oss_mirror( + oss_bucket="conv-bucket", + oss_access_key_id="ak-conv", + oss_access_key_secret="sk-conv", + oss_region="cn-hangzhou", + oss_endpoint="oss-cn-hangzhou.aliyuncs.com", + ) + assert cfg.environment.oss_mirror.enabled is True + assert cfg.environment.oss_mirror.oss_bucket == "conv-bucket" + + def test_does_not_touch_namespace_or_experiment_id(self): + """enable_oss_mirror 不修改顶层 namespace / experiment_id。""" + from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.sandbox.config import SandboxConfig + + cfg = JobConfig( + job_name="no-touch-test", + namespace="preset-ns", + experiment_id="preset-exp", + sandbox_config=SandboxConfig( + namespace="sandbox-ns", + experiment_id="sandbox-exp", + ), + ) + cfg.enable_oss_mirror( + oss_bucket="b", + oss_access_key_id="ak", + oss_access_key_secret="sk", + oss_region="r1", + oss_endpoint="e1", + ) + assert cfg.namespace == "preset-ns" + assert cfg.experiment_id == "preset-exp" + + def test_enable_then_serialize_roundtrip(self): + """to_harbor_yaml: 顶层 namespace / experiment_id 与 oss_mirror 独立设置。""" + from rock.sdk.agent.models.job.config import JobConfig + + cfg = JobConfig(job_name="roundtrip", namespace="rt-ns", experiment_id="rt-exp") + cfg.enable_oss_mirror( + oss_bucket="rt-bucket", + oss_access_key_id="ak-rt", + oss_access_key_secret="sk-rt", + oss_region="ap-southeast-1", + oss_endpoint="oss-ap-southeast-1.aliyuncs.com", + ) + data = yaml.safe_load(cfg.to_harbor_yaml()) + assert data["namespace"] == "rt-ns" + assert data["experiment_id"] == "rt-exp" + oss = data["environment"]["oss_mirror"] + assert oss["enabled"] is True + assert oss["oss_bucket"] == "rt-bucket" + assert "namespace" not in oss + assert "experiment_id" not in oss From 13531dc689ba2c77e179b9931cad5323f05c6f32 Mon Sep 17 00:00:00 2001 From: dengwx Date: Tue, 31 Mar 2026 20:31:49 +0800 Subject: [PATCH 003/226] feat: validate and sync experiment_id/namespace in JobConfig (#716) (#717) * feat: refine expr_id and namespace * feat(sdk): add JobConfig enhancements and fix linting issues Add job configuration improvements including experiment_id support, OSS mirror config updates, and port validation changes. Fix ruff lint and format issues across the codebase. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use urlparse for URL hostname validation in speedup tests Replace substring-based URL checks with urlparse().hostname to satisfy CodeQL's "Incomplete URL substring sanitization" rule. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: extract domain constants to avoid CodeQL URL substring warnings CodeQL flags `"domain" in var` as incomplete URL sanitization regardless of variable type. Extract hostnames to constants and use helper functions to eliminate the flagged pattern entirely. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- docs/dev/agent/JobConfig.md | 111 ++++++++++++ examples/evaluation/swe_bench/common.py | 4 +- .../swe_bench/swe_bench_verified_demo.py | 3 +- .../admin/metrics/gc_view_instrument_match.py | 33 ++-- rock/admin/scheduler/task_factory.py | 4 +- rock/common/port_validation.py | 1 + rock/rocklet/local_api.py | 26 +-- rock/sandbox/operator/k8s/api_client.py | 98 +++++----- rock/sandbox/operator/k8s/constants.py | 10 +- rock/sandbox/operator/k8s/operator.py | 41 +++-- rock/sandbox/operator/k8s/provider.py | 171 +++++++++--------- rock/sandbox/operator/k8s/template_loader.py | 150 +++++++-------- rock/sdk/agent/job.py | 9 +- rock/sdk/agent/models/job/config.py | 27 ++- rock/sdk/envs/rock_env.py | 3 +- rock/sdk/model/server/file_handler.py | 1 + rock/sdk/model/server/main.py | 1 + rock/sdk/sandbox/client.py | 2 +- rock/utils/providers/nacos_provider.py | 4 +- rock/utils/service.py | 1 - tests/conftest.py | 4 +- tests/integration/sdk/sandbox/test_basic.py | 6 +- .../sdk/sandbox/test_file_system.py | 6 +- .../sdk/sandbox/test_sandbox_images.py | 7 +- .../sdk/sandbox/test_sdk_client.py | 6 +- tests/integration/sdk/sandbox/test_speedup.py | 51 ++++-- tests/unit/admin/core/test_ray_service.py | 16 +- tests/unit/common/test_port_validation.py | 5 +- tests/unit/rocklet/test_docker_deployment.py | 1 + tests/unit/rocklet/test_portforward.py | 3 +- .../sandbox/operator/test_k8s_api_client.py | 70 ++++--- .../sandbox/operator/test_k8s_operator.py | 13 +- .../sandbox/operator/test_k8s_provider.py | 105 ++++++----- .../operator/test_k8s_template_loader.py | 77 +++----- tests/unit/sandbox/test_sandbox_http_proxy.py | 1 + tests/unit/sandbox/test_sandbox_manager.py | 6 +- tests/unit/sandbox/test_sandbox_proxy.py | 2 + .../test_websocket_proxy_subprotocol.py | 70 ++++--- .../unit/sandbox/test_websocket_tcp_proxy.py | 7 +- tests/unit/sdk/agent/test_job.py | 24 ++- .../agent/test_job_config_serialization.py | 12 +- .../sdk/agent/test_jobconfig_experiment_id.py | 102 +++++++++++ tests/unit/sdk/agent/test_models.py | 5 +- tests/unit/sdk/agent/test_oss_mirror.py | 18 +- tests/unit/sdk/test_arun_nohup.py | 3 +- tests/unit/test_config.py | 3 +- tests/unit/utils/test_shell_util.py | 20 +- 47 files changed, 786 insertions(+), 557 deletions(-) create mode 100644 docs/dev/agent/JobConfig.md create mode 100644 tests/unit/sdk/agent/test_jobconfig_experiment_id.py diff --git a/docs/dev/agent/JobConfig.md b/docs/dev/agent/JobConfig.md new file mode 100644 index 0000000000..c4422ac2d8 --- /dev/null +++ b/docs/dev/agent/JobConfig.md @@ -0,0 +1,111 @@ +# JobConfig 字段分析: `namespace` 与 `experiment_id` + +## 1. 现状分析 + +### 1.1 字段定义位置 + +| 字段 | 定义位置 | 类型 | +|------|---------|------| +| `namespace` | `JobConfig` (`models/job/config.py:143`) | `str \| None = None` | +| `experiment_id` | `JobConfig` (`models/job/config.py:147`) | `str \| None = None` | +| `experiment_id` | `SandboxConfig` (`sdk/sandbox/config.py:40`) | `str \| None = None` | +| `namespace` | `SandboxConfig` (`sdk/sandbox/config.py:42`) | `str \| None = None` | +| `_namespace` | `Sandbox` (`sdk/sandbox/client.py:78`) | `str \| None = None` | +| `_experiment_id` | `Sandbox` (`sdk/sandbox/client.py:79`) | `str \| None = None` | + +注意: `JobConfig.environment` 类型为 `RockEnvironmentConfig`，继承自 `SandboxConfig`，因此 `self.environment.experiment_id` 来自 `SandboxConfig`。 + +### 1.2 当前数据流 + +``` +用户构造 JobConfig(experiment_id="exp-1", environment=RockEnvironmentConfig(experiment_id=?)) + ↑ 来自 SandboxConfig + +Sandbox.start() + └─ get_status() 获取 sandbox_info + ├─ sandbox_info.namespace → Sandbox._namespace (client.py:203) + └─ sandbox_info.experiment_id → Sandbox._experiment_id (client.py:205) + +Job.submit() + └─ _prepare_and_start() + └─ _autofill_sandbox_info() (job.py:268-269) + └─ self._config.namespace = self._sandbox._namespace + └─ (experiment_id 未处理) +``` + +### 1.3 `to_harbor_yaml()` 序列化 + +`JobConfig.to_harbor_yaml()` 将 `namespace` 和 `experiment_id` 序列化到 Harbor YAML 中（`exclude={"environment"}, exclude_none=True`），最终传给 `harbor jobs start -c`。 + +--- + +## 2. 问题 + +### 2.1 `experiment_id` 两处定义未同步 + +- `JobConfig.experiment_id` — Harbor YAML 层面的实验标识 +- `SandboxConfig.experiment_id`（通过 `JobConfig.environment`）— sandbox 创建时传递的实验标识 +- **问题**: 两个 `experiment_id` 各自独立，没有同步或校验逻辑。用户可能在两处设置不同的值，导致 sandbox 创建和 Harbor 执行使用不同的 experiment_id。 + +### 2.2 `namespace` 缺少一致性校验 + +- `_autofill_sandbox_info()` 直接覆盖 `self._config.namespace = self._sandbox._namespace` +- **问题**: 如果用户已经设置了 `namespace`（非 None），当前逻辑会静默覆盖，不做任何校验。 + +--- + +## 3. 改进方案 + +### 3.1 `experiment_id`: 在 JobConfig 中增加 model_validator (post init) + +**保留** `JobConfig.experiment_id` 作为唯一权威来源，通过 `model_validator(mode="after")` 做三件事： + +1. **校验非空**: `JobConfig.experiment_id` 不能为 None 或空字符串 +2. **一致性校验**: 如果 `environment.experiment_id`（即 SandboxConfig 的）已有值，必须与 `JobConfig.experiment_id` 一致，否则抛异常 +3. **向下同步**: 将 `JobConfig.experiment_id` 设置到 `environment.experiment_id` + +```python +@model_validator(mode="after") +def _sync_experiment_id(self): + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + env_exp = self.environment.experiment_id + if env_exp is not None and env_exp != self.experiment_id: + raise ValueError( + f"experiment_id mismatch: JobConfig has '{self.experiment_id}', " + f"but environment (SandboxConfig) has '{env_exp}'" + ) + self.environment.experiment_id = self.experiment_id + return self +``` + +**行为矩阵**: + +| JobConfig.experiment_id | environment.experiment_id | 行为 | +|------------------------|--------------------------|------| +| `None` 或 `""` | 任意 | **抛出 ValueError**: experiment_id 不能为空 | +| `"exp-1"` | `None` | 同步: `environment.experiment_id = "exp-1"` | +| `"exp-1"` | `"exp-1"` | 通过，一致 | +| `"exp-1"` | `"exp-2"` | **抛出 ValueError**: mismatch | + +### 3.2 `namespace`: 保持由 sandbox 返回值设置（运行时回填） + +`namespace` 与 `experiment_id` 不同 — 它的权威来源是 sandbox 运行时返回值，用户通常不需要设置。保持在 `_autofill_sandbox_info()` 中处理，但增加一致性校验： + +| 用户设置 | sandbox 返回 | 行为 | +|---------|-------------|------| +| `None` | 有值 | 自动回填 sandbox 返回值 | +| `None` | `None` | 保持 `None` | +| 有值 | 有值且一致 | 保持不变 | +| 有值 | 有值且不一致 | **抛出 ValueError** | +| 有值 | `None` | 保留用户设置值 | + +--- + +## 4. 涉及文件 + +| 文件 | 变更内容 | +|------|---------| +| `rock/sdk/agent/models/job/config.py` | 新增 `_sync_experiment_id` model_validator | +| `rock/sdk/agent/job.py` | 更新 `_autofill_sandbox_info()`，namespace 增加一致性校验 | +| `tests/unit/sdk/agent/` | 补充测试: validator 校验逻辑、mismatch 异常、空值异常 | diff --git a/examples/evaluation/swe_bench/common.py b/examples/evaluation/swe_bench/common.py index ef8c0cf461..6821813d9d 100644 --- a/examples/evaluation/swe_bench/common.py +++ b/examples/evaluation/swe_bench/common.py @@ -1,7 +1,8 @@ import re -import yaml from pathlib import Path +import yaml + from rock.logger import init_logger from rock.sdk.sandbox.client import RunMode, Sandbox from rock.sdk.sandbox.config import SandboxConfig @@ -15,6 +16,7 @@ logger = init_logger(__name__) + def load_task_config(task_dir: Path) -> dict: """Load task configuration from task.yaml.""" task_yaml_path = task_dir / "task.yaml" diff --git a/examples/evaluation/swe_bench/swe_bench_verified_demo.py b/examples/evaluation/swe_bench/swe_bench_verified_demo.py index 1c0befa211..906dbe6f92 100644 --- a/examples/evaluation/swe_bench/swe_bench_verified_demo.py +++ b/examples/evaluation/swe_bench/swe_bench_verified_demo.py @@ -18,8 +18,8 @@ python -m examples.evaluation.swe_bench.swe_bench_verified_demo """ -import sys import asyncio +import sys from pathlib import Path from examples.evaluation.swe_bench.common import load_task_config, parse_swebench_result, setup_test_env, start_sandbox @@ -31,6 +31,7 @@ test_timeout_sec = 3600 logger = init_logger(__name__) + async def run_swe_evaluation(sandbox: Sandbox, task_dir: Path, instruction: str, agent_config_path: str) -> bool: """Run SWE evaluation on the sandbox.""" task_name = task_dir.name diff --git a/rock/admin/metrics/gc_view_instrument_match.py b/rock/admin/metrics/gc_view_instrument_match.py index 200faa04b0..84a722c8f5 100644 --- a/rock/admin/metrics/gc_view_instrument_match.py +++ b/rock/admin/metrics/gc_view_instrument_match.py @@ -1,5 +1,5 @@ +from collections.abc import Sequence from time import time_ns -from typing import Dict, List, Optional, Sequence from opentelemetry.sdk.metrics._internal._view_instrument_match import ( _ViewInstrumentMatch as _OrigViewInstrumentMatch, @@ -15,16 +15,15 @@ class _GcViewInstrumentMatch(_OrigViewInstrumentMatch): metric series (based on attributes). This is useful for preventing memory leaks when dealing with high-cardinality metrics. """ + # Idle metric series are cleaned up after 20 minutes. This can be adjusted. _IDLE_TIMEOUT_NS = 20 * 60 * 1_000_000_000 def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._last_used_ns: Dict[frozenset, int] = {} + self._last_used_ns: dict[frozenset, int] = {} - def consume_measurement( - self, measurement: Measurement, should_sample_exemplar: bool = True - ) -> None: + def consume_measurement(self, measurement: Measurement, should_sample_exemplar: bool = True) -> None: """ Consumes a measurement, aggregates it, and tracks its usage for GC. """ @@ -32,11 +31,7 @@ def consume_measurement( measurement_for_aggregation = measurement if self._view._attribute_keys is not None: - filtered_attributes = { - key: value - for key, value in attributes.items() - if key in self._view._attribute_keys - } + filtered_attributes = {key: value for key, value in attributes.items() if key in self._view._attribute_keys} # If attributes were filtered, a new Measurement object must be used # for aggregation. This ensures that if an exemplar is recorded, it @@ -61,22 +56,20 @@ def consume_measurement( should_sample_exemplar, ) else: - self._attributes_aggregation[aggr_key].aggregate( - measurement_for_aggregation, should_sample_exemplar - ) + self._attributes_aggregation[aggr_key].aggregate(measurement_for_aggregation, should_sample_exemplar) self._last_used_ns[aggr_key] = now_ns def collect( - self, - collection_aggregation_temporality: AggregationTemporality, - collection_start_nanos: int, - ) -> Optional[Sequence[DataPointT]]: + self, + collection_aggregation_temporality: AggregationTemporality, + collection_start_nanos: int, + ) -> Sequence[DataPointT] | None: """ Collects all data points for the metric, and garbage collects idle series. """ - data_points: List[DataPointT] = [] + data_points: list[DataPointT] = [] now_ns = time_ns() - to_delete: List[frozenset] = [] + to_delete: list[frozenset] = [] with self._lock: # First, collect data points and identify idle series @@ -104,7 +97,9 @@ def patch_view_instrument_match() -> None: # Call this once at application startup, before initializing any metric # readers or providers, to replace the SDK's internal class. import opentelemetry.sdk.metrics._internal._view_instrument_match as vim_mod + vim_mod._ViewInstrumentMatch = _GcViewInstrumentMatch from opentelemetry.sdk.metrics._internal import metric_reader_storage as mrs + mrs._ViewInstrumentMatch = _GcViewInstrumentMatch diff --git a/rock/admin/scheduler/task_factory.py b/rock/admin/scheduler/task_factory.py index 2f5cd1bfe8..5130f264c3 100644 --- a/rock/admin/scheduler/task_factory.py +++ b/rock/admin/scheduler/task_factory.py @@ -62,8 +62,6 @@ def register_all_tasks(cls, scheduler_config: SchedulerConfig): try: task = cls.create_task(task_config) TaskRegistry.register(task) - logger.info( - f"Registered task '{task.type}' with interval {task.interval_seconds}s" - ) + logger.info(f"Registered task '{task.type}' with interval {task.interval_seconds}s") except Exception as e: logger.error(f"Failed to create task '{task_config.task_class}': {e}") diff --git a/rock/common/port_validation.py b/rock/common/port_validation.py index 17cb1fbc25..065fc65b35 100644 --- a/rock/common/port_validation.py +++ b/rock/common/port_validation.py @@ -1,4 +1,5 @@ """Port validation utilities for port forwarding.""" + from rock.logger import init_logger logger = init_logger(__name__) diff --git a/rock/rocklet/local_api.py b/rock/rocklet/local_api.py index 97ebd6d022..7cb390336f 100644 --- a/rock/rocklet/local_api.py +++ b/rock/rocklet/local_api.py @@ -180,19 +180,13 @@ async def portforward(websocket: WebSocket, port: int): try: # Connect to local TCP port - reader, writer = await asyncio.wait_for( - asyncio.open_connection("127.0.0.1", port), - timeout=TCP_CONNECT_TIMEOUT - ) + reader, writer = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", port), timeout=TCP_CONNECT_TIMEOUT) logger.info( f"[Portforward] TCP connection established: target_port={port}, " f"local_addr={writer.get_extra_info('sockname')}" ) except asyncio.TimeoutError: - logger.error( - f"[Portforward] TCP connection timeout: target_port={port}, " - f"timeout={TCP_CONNECT_TIMEOUT}s" - ) + logger.error(f"[Portforward] TCP connection timeout: target_port={port}, " f"timeout={TCP_CONNECT_TIMEOUT}s") await websocket.close(code=1011, reason=f"Connection to port {port} timed out") return except OSError as e: @@ -204,8 +198,7 @@ async def portforward(websocket: WebSocket, port: int): return except Exception as e: logger.error( - f"[Portforward] Unexpected TCP error: target_port={port}, " - f"error_type={type(e).__name__}, error={e}" + f"[Portforward] Unexpected TCP error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" ) await websocket.close(code=1011, reason=f"Unexpected error: {e}") return @@ -232,13 +225,10 @@ async def ws_to_tcp(): f"bytes={len(data)}, total_msgs={ws_to_tcp_msgs}, total_bytes={ws_to_tcp_bytes}" ) except WebSocketDisconnect as e: - logger.info( - f"[Portforward] ws->tcp: client disconnected: target_port={port}, code={e.code}" - ) + logger.info(f"[Portforward] ws->tcp: client disconnected: target_port={port}, code={e.code}") except Exception as e: logger.debug( - f"[Portforward] ws->tcp error: target_port={port}, " - f"error_type={type(e).__name__}, error={e}" + f"[Portforward] ws->tcp error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" ) finally: writer.close() @@ -261,8 +251,7 @@ async def tcp_to_ws(): ) except Exception as e: logger.debug( - f"[Portforward] tcp->ws error: target_port={port}, " - f"error_type={type(e).__name__}, error={e}" + f"[Portforward] tcp->ws error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" ) finally: try: @@ -275,8 +264,7 @@ async def tcp_to_ws(): await asyncio.gather(ws_to_tcp(), tcp_to_ws()) except Exception as e: logger.debug( - f"[Portforward] Forwarding error: target_port={port}, " - f"error_type={type(e).__name__}, error={e}" + f"[Portforward] Forwarding error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" ) finally: writer.close() diff --git a/rock/sandbox/operator/k8s/api_client.py b/rock/sandbox/operator/k8s/api_client.py index 200e9efa36..d0fb879045 100644 --- a/rock/sandbox/operator/k8s/api_client.py +++ b/rock/sandbox/operator/k8s/api_client.py @@ -18,20 +18,19 @@ from rock.logger import init_logger - logger = init_logger(__name__) class K8sApiClient: """K8s API client wrapper with rate limiting and Informer cache. - + Centralizes K8s API Server access with: - Rate limiting via aiolimiter to prevent API Server overload - Local cache with watch-based sync (Informer pattern) to reduce API calls - Consistent error handling - Simple CRUD interface for K8s custom resources """ - + def __init__( self, api_client: client.ApiClient, @@ -44,7 +43,7 @@ def __init__( watch_reconnect_delay_seconds: int = 5, ): """Initialize K8s API client. - + Args: api_client: Kubernetes ApiClient instance group: CRD API group @@ -61,32 +60,32 @@ def __init__( self._plural = plural self._namespace = namespace self._custom_api = client.CustomObjectsApi(api_client) - + # Rate limiting self._rate_limiter = AsyncLimiter(max_rate=qps, time_period=1.0) - + # Watch configuration self._watch_timeout_seconds = watch_timeout_seconds self._watch_reconnect_delay_seconds = watch_reconnect_delay_seconds - + # Local cache for resources (Informer pattern) self._cache: dict[str, dict] = {} self._cache_lock = asyncio.Lock() self._watch_task = None self._initialized = False - + async def start(self): """Start the API client and initialize cache watch.""" if self._initialized: return - + self._watch_task = asyncio.create_task(self._watch_resources()) self._initialized = True logger.info(f"Started K8sApiClient watch for {self._plural} in namespace {self._namespace}") - + async def _list_and_sync_cache(self) -> str: """List all resources and sync to cache. - + Returns: resourceVersion for next watch """ @@ -98,19 +97,19 @@ async def _list_and_sync_cache(self) -> str: namespace=self._namespace, plural=self._plural, ) - - resource_version = resources.get('metadata', {}).get('resourceVersion') + + resource_version = resources.get("metadata", {}).get("resourceVersion") async with self._cache_lock: self._cache.clear() - for item in resources.get('items', []): - name = item.get('metadata', {}).get('name') + for item in resources.get("items", []): + name = item.get("metadata", {}).get("name") if name: self._cache[name] = item return resource_version - + async def _watch_resources(self): """Background task to watch resources and maintain cache. - + Implements Kubernetes Informer pattern: 1. Initial list-and-sync to populate cache 2. Continuous watch for ADDED/MODIFIED/DELETED events @@ -120,12 +119,15 @@ async def _watch_resources(self): resource_version = None try: resource_version = await self._list_and_sync_cache() - logger.info(f"Initial cache populated with {len(self._cache)} resources, resourceVersion={resource_version}") + logger.info( + f"Initial cache populated with {len(self._cache)} resources, resourceVersion={resource_version}" + ) except Exception as e: logger.error(f"Failed to populate initial cache: {e}") - + while True: try: + def _watch_in_thread(): w = watch.Watch() stream = w.stream( @@ -141,27 +143,27 @@ def _watch_in_thread(): for event in stream: events.append(event) return events - + events = await asyncio.to_thread(_watch_in_thread) - + async with self._cache_lock: for event in events: - event_type = event['type'] - obj = event['object'] - name = obj.get('metadata', {}).get('name') - new_rv = obj.get('metadata', {}).get('resourceVersion') - + event_type = event["type"] + obj = event["object"] + name = obj.get("metadata", {}).get("name") + new_rv = obj.get("metadata", {}).get("resourceVersion") + if new_rv: resource_version = new_rv - + if not name: continue - - if event_type in ['ADDED', 'MODIFIED']: + + if event_type in ["ADDED", "MODIFIED"]: self._cache[name] = obj - elif event_type == 'DELETED': + elif event_type == "DELETED": self._cache.pop(name, None) - + except asyncio.CancelledError: logger.info("Watch task cancelled") raise @@ -169,20 +171,24 @@ def _watch_in_thread(): logger.warning(f"Watch stream disconnected: {e}, reconnecting immediately...") try: resource_version = await self._list_and_sync_cache() - logger.info(f"Re-synced cache with {len(self._cache)} resources, resourceVersion={resource_version}") + logger.info( + f"Re-synced cache with {len(self._cache)} resources, resourceVersion={resource_version}" + ) except Exception as list_err: - logger.error(f"Failed to re-list resources: {list_err}, retrying in {self._watch_reconnect_delay_seconds}s...") + logger.error( + f"Failed to re-list resources: {list_err}, retrying in {self._watch_reconnect_delay_seconds}s..." + ) await asyncio.sleep(self._watch_reconnect_delay_seconds) - + async def create_custom_object( self, body: dict[str, Any], ) -> dict[str, Any]: """Create a custom resource. - + Args: body: Resource manifest - + Returns: Created resource """ @@ -195,25 +201,25 @@ async def create_custom_object( plural=self._plural, body=body, ) - + async def get_custom_object( self, name: str, ) -> dict[str, Any]: """Get a custom resource (from cache with fallback to API Server). - + Args: name: Resource name - + Returns: Resource object """ async with self._cache_lock: resource = self._cache.get(name) - + if resource: return resource - + logger.debug(f"Cache miss for {name}, querying API Server") async with self._rate_limiter: resource = await asyncio.to_thread( @@ -224,21 +230,21 @@ async def get_custom_object( plural=self._plural, name=name, ) - + async with self._cache_lock: self._cache[name] = resource - + return resource - + async def delete_custom_object( self, name: str, ) -> dict[str, Any]: """Delete a custom resource. - + Args: name: Resource name - + Returns: Delete status """ diff --git a/rock/sandbox/operator/k8s/constants.py b/rock/sandbox/operator/k8s/constants.py index 5571aef94f..02a94acbdc 100644 --- a/rock/sandbox/operator/k8s/constants.py +++ b/rock/sandbox/operator/k8s/constants.py @@ -3,26 +3,26 @@ class K8sConstants: """Constants for K8S BatchSandbox labels and annotations.""" - + # CRD configuration CRD_GROUP = "sandbox.opensandbox.io" CRD_VERSION = "v1alpha1" CRD_PLURAL = "batchsandboxes" CRD_KIND = "BatchSandbox" CRD_API_VERSION = f"{CRD_GROUP}/{CRD_VERSION}" # sandbox.opensandbox.io/v1alpha1 - + # Annotation keys ANNOTATION_ENDPOINTS = "sandbox.opensandbox.io/endpoints" ANNOTATION_PORTS = "rock.sandbox/ports" - + # Label keys LABEL_SANDBOX_ID = "rock.sandbox/sandbox-id" LABEL_RESOURCE_SPEEDUP = "batchsandbox.alibabacloud.com/resource-speedup" LABEL_TEMPLATE = "rock.sandbox/template" - + # Extension keys for DockerDeploymentConfig.extended_params EXT_POOL_NAME = "pool_name" EXT_TEMPLATE_NAME = "template_name" - + # Nacos config keys NACOS_POOLS_KEY = "pools" diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 2472787709..42e6b77f57 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -1,7 +1,7 @@ """K8s Operator implementation for managing sandboxes via Kubernetes.""" -from rock.config import K8sConfig from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.config import K8sConfig from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger from rock.sandbox.operator.abstract import AbstractOperator @@ -12,10 +12,10 @@ class K8sOperator(AbstractOperator): """Operator for managing sandboxes via Kubernetes BatchSandbox CRD.""" - + def __init__(self, k8s_config: K8sConfig, redis_provider=None): """Initialize K8s operator. - + Args: k8s_config: K8sConfig object containing kubeconfig and templates redis_provider: Optional Redis provider for caching sandbox info @@ -23,62 +23,63 @@ def __init__(self, k8s_config: K8sConfig, redis_provider=None): self._provider = BatchSandboxProvider(k8s_config=k8s_config) self._redis_provider = redis_provider logger.info("Initialized K8sOperator") - + def set_nacos_provider(self, nacos_provider): """Set Nacos config provider for dynamic pool configuration. - + Args: nacos_provider: NacosConfigProvider instance """ super().set_nacos_provider(nacos_provider) self._provider.set_nacos_provider(nacos_provider) - + async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> SandboxInfo: """Submit a new sandbox deployment to Kubernetes. - + Args: config: Docker deployment configuration user_info: User metadata (user_id, experiment_id, namespace, rock_authorization) - + Returns: SandboxInfo with sandbox metadata """ return await self._provider.submit(config, user_info) - + async def get_status(self, sandbox_id: str) -> SandboxInfo: """Get sandbox status with user info from Redis. - + This method first gets status from provider (IP, port_mapping, is_alive), then merges it with user info from Redis if available. - + Args: sandbox_id: Sandbox identifier - + Returns: SandboxInfo with current status and user info """ # Get sandbox info from provider (includes is_alive check) sandbox_info = await self._provider.get_status(sandbox_id) - + # Get user info from redis if available if self._redis_provider: redis_info = await self._get_sandbox_info_from_redis(sandbox_id) if redis_info: redis_info.update(sandbox_info) return redis_info - + return sandbox_info - + async def _get_sandbox_info_from_redis(self, sandbox_id: str) -> dict | None: """Get sandbox user info from Redis. - + Args: sandbox_id: Sandbox identifier - + Returns: Sandbox info dict from Redis or None if not found """ from rock.admin.core.redis_key import alive_sandbox_key + try: sandbox_status = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") if sandbox_status and len(sandbox_status) > 0: @@ -86,13 +87,13 @@ async def _get_sandbox_info_from_redis(self, sandbox_id: str) -> dict | None: except Exception as e: logger.debug(f"Failed to get sandbox info from redis for {sandbox_id}: {e}") return None - + async def stop(self, sandbox_id: str) -> bool: """Stop and delete a sandbox. - + Args: sandbox_id: Sandbox identifier - + Returns: True if successful, False otherwise """ diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index b0bfba5fb8..8ab97d2350 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -2,22 +2,22 @@ import json import re -from abc import abstractmethod, ABC +from abc import ABC, abstractmethod from typing import Any, Protocol -from kubernetes import client, config as k8s_config +from kubernetes import client +from kubernetes import config as k8s_config +from rock.actions.sandbox.config import RemoteSandboxRuntimeConfig +from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.config import K8sConfig, PoolConfig from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig from rock.deployments.constants import Port -from rock.sandbox.operator.k8s.constants import K8sConstants +from rock.logger import init_logger from rock.sandbox.operator.k8s.api_client import K8sApiClient +from rock.sandbox.operator.k8s.constants import K8sConstants from rock.sandbox.operator.k8s.template_loader import K8sTemplateLoader from rock.sandbox.remote_sandbox import RemoteSandboxRuntime -from rock.actions.sandbox.config import RemoteSandboxRuntimeConfig -from rock.actions.sandbox.sandbox_info import SandboxInfo -from rock.logger import init_logger - logger = init_logger(__name__) @@ -53,7 +53,7 @@ def _parse_memory_to_mb(self, memory: str) -> float: memory = memory.lower().strip() # Extract number and unit - match = re.match(r'^(\d+(\.\d+)?)\s*([a-z]*)$', memory) + match = re.match(r"^(\d+(\.\d+)?)\s*([a-z]*)$", memory) if not match: try: return float(memory) / (1024 * 1024) # Assume bytes @@ -64,15 +64,15 @@ def _parse_memory_to_mb(self, memory: str) -> float: unit = match.group(3) # Convert to MB - if unit in ('', 'b'): + if unit in ("", "b"): return value / (1024 * 1024) - elif unit in ('k', 'kb'): + elif unit in ("k", "kb"): return value / 1024 - elif unit in ('m', 'mb', 'mi'): + elif unit in ("m", "mb", "mi"): return value - elif unit in ('g', 'gb', 'gi'): + elif unit in ("g", "gb", "gi"): return value * 1024 - elif unit in ('t', 'tb', 'ti'): + elif unit in ("t", "tb", "ti"): return value * 1024 * 1024 else: return 0 @@ -160,17 +160,17 @@ async def stop(self, sandbox_id: str) -> bool: class BatchSandboxProvider(K8sProvider): """Provider for BatchSandbox CRD with Informer-based local cache. - + This provider uses Kubernetes watch API to maintain a local cache of BatchSandbox resources. All get_status queries read from this cache instead of querying API Server, which significantly improves performance and reduces API Server load. - + The watch task runs in the background and automatically reconnects on network failures. """ - + def __init__(self, k8s_config: K8sConfig): """Initialize BatchSandbox provider. - + Args: k8s_config: K8sConfig object containing kubeconfig and templates """ @@ -181,26 +181,26 @@ def __init__(self, k8s_config: K8sConfig): self._k8s_api: K8sApiClient | None = None self._initialized = False self._nacos_provider = None - + # Initialize template loader with config templates self._template_loader = K8sTemplateLoader( templates=k8s_config.templates, default_namespace=k8s_config.namespace, ) logger.info(f"Available K8S templates: {', '.join(self._template_loader.available_templates)}") - + def set_nacos_provider(self, nacos_provider): """Set Nacos config provider for dynamic pool configuration. - + Args: nacos_provider: NacosConfigProvider instance """ self._nacos_provider = nacos_provider logger.info("Set nacos provider for K8s provider") - + async def _get_pools(self) -> dict[str, PoolConfig]: """Get pool configurations from Nacos. - + Returns: Dictionary of pool name to PoolConfig """ @@ -216,7 +216,7 @@ async def _get_pools(self) -> dict[str, PoolConfig]: pools[name] = config logger.debug(f"Loaded {len(pools)} pools from Nacos") return pools - + return {} async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> SandboxInfo: @@ -276,7 +276,7 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> async def get_status(self, sandbox_id: str) -> SandboxInfo: """Get sandbox status and check if alive. - + This method fetches the sandbox resource from K8s and checks if the sandbox is alive by calling its is_alive endpoint. The state is determined by: - RUNNING: IP allocated AND is_alive returns true @@ -338,17 +338,17 @@ async def stop(self, sandbox_id: str) -> bool: async def _ensure_initialized(self): """Ensure K8s client is initialized and start watch task. - + Lazy initialization of K8S client and API abstraction layer: 1. Load kubeconfig (from file, in-cluster, or default) 2. Create K8sApiClient with rate limiting and caching 3. Start background watch task for cache synchronization - + Thread-safe: Uses _initialized flag to prevent duplicate initialization. """ if self._initialized: return - + try: if self.kubeconfig_path: k8s_config.load_kube_config(config_file=self.kubeconfig_path) @@ -358,7 +358,7 @@ async def _ensure_initialized(self): k8s_config.load_incluster_config() except k8s_config.ConfigException: k8s_config.load_kube_config() - + self._api_client = client.ApiClient() self._k8s_api = K8sApiClient( api_client=self._api_client, @@ -372,12 +372,12 @@ async def _ensure_initialized(self): ) await self._k8s_api.start() self._initialized = True - + logger.info("Initialized K8s provider with informer") except Exception as e: logger.error(f"Failed to initialize K8s client: {e}", exc_info=True) raise - + async def _get_pool_name(self, config: DockerDeploymentConfig) -> str | None: """Get pool name using selection strategy. @@ -435,11 +435,11 @@ def _normalize_memory(self, memory: str) -> str: Convert formats like '2g', '2G', '2048m' to K8s format like '2Gi', '2048Mi'. """ # Already in K8s format - if re.match(r'^\d+(\.\d+)?(Ei|Pi|Ti|Gi|Mi|Ki)$', memory): + if re.match(r"^\d+(\.\d+)?(Ei|Pi|Ti|Gi|Mi|Ki)$", memory): return memory - + # Parse value and unit - match = re.match(r'^(\d+(\.\d+)?)([a-zA-Z]*)$', memory) + match = re.match(r"^(\d+(\.\d+)?)([a-zA-Z]*)$", memory) if not match: # Fallback: assume it's bytes and convert to Mi try: @@ -452,51 +452,51 @@ def _normalize_memory(self, memory: str) -> str: unit = match.group(3).lower() # Convert to K8s format - use int() for whole numbers, preserve decimals otherwise - if unit in ('', 'b'): + if unit in ("", "b"): mi_value = value / (1024 * 1024) return f"{int(mi_value) if mi_value == int(mi_value) else mi_value:.2f}Mi" - elif unit in ('k', 'kb'): + elif unit in ("k", "kb"): return f"{int(value) if value == int(value) else value:.2f}Ki" - elif unit in ('m', 'mb'): + elif unit in ("m", "mb"): return f"{int(value) if value == int(value) else value:.2f}Mi" - elif unit in ('g', 'gb'): + elif unit in ("g", "gb"): return f"{int(value) if value == int(value) else value:.2f}Gi" - elif unit in ('t', 'tb'): + elif unit in ("t", "tb"): return f"{int(value) if value == int(value) else value:.2f}Ti" else: return memory def _build_pool_manifest(self, sandbox_id: str, pool_name: str, ports_config: dict[str, int]) -> dict[str, Any]: """Build BatchSandbox manifest for pool mode. - + Args: sandbox_id: Sandbox identifier pool_name: Pool name to reference ports_config: Port configuration dictionary - + Returns: Manifest dictionary """ - + manifest = { - 'apiVersion': K8sConstants.CRD_API_VERSION, - 'kind': K8sConstants.CRD_KIND, - 'metadata': { - 'name': sandbox_id, - 'namespace': self.namespace, - 'labels': { + "apiVersion": K8sConstants.CRD_API_VERSION, + "kind": K8sConstants.CRD_KIND, + "metadata": { + "name": sandbox_id, + "namespace": self.namespace, + "labels": { K8sConstants.LABEL_SANDBOX_ID: sandbox_id, }, - 'annotations': { + "annotations": { K8sConstants.ANNOTATION_PORTS: json.dumps(ports_config), }, }, - 'spec': { - 'poolRef': pool_name, - 'replicas': 1, - } + "spec": { + "poolRef": pool_name, + "replicas": 1, + }, } - + return manifest async def _get_pool_ports(self, pool_name: str) -> dict[str, int]: @@ -533,13 +533,15 @@ async def _build_batchsandbox_manifest(self, config: DockerDeploymentConfig) -> if pool_name: ports_config = await self._get_pool_ports(pool_name) manifest = self._build_pool_manifest(sandbox_id, pool_name, ports_config) - - logger.debug(f"Built BatchSandbox manifest for {sandbox_id} using pool '{pool_name}' in namespace '{self.namespace}'") + + logger.debug( + f"Built BatchSandbox manifest for {sandbox_id} using pool '{pool_name}' in namespace '{self.namespace}'" + ) return manifest # Template mode: build from template template_name = self._get_template_name(config) - + # Build manifest using template loader manifest = self._template_loader.build_manifest( template_name=template_name, @@ -548,35 +550,37 @@ async def _build_batchsandbox_manifest(self, config: DockerDeploymentConfig) -> cpus=config.cpus, memory=self._normalize_memory(config.memory), ) - - logger.debug(f"Built BatchSandbox manifest for {sandbox_id} in namespace '{self.namespace}' using template '{template_name}'") + + logger.debug( + f"Built BatchSandbox manifest for {sandbox_id} in namespace '{self.namespace}' using template '{template_name}'" + ) return manifest async def _create(self, config: DockerDeploymentConfig) -> str: """Create a BatchSandbox resource without waiting for IP allocation. - + Args: config: Docker deployment configuration - + Returns: sandbox_id (same as config.container_name) - + Raises: Exception: If creation fails or sandbox already exists """ await self._ensure_initialized() - + sandbox_id = config.container_name - + try: manifest = await self._build_batchsandbox_manifest(config) - + # Create BatchSandbox resource await self._k8s_api.create_custom_object(body=manifest) - + logger.info(f"Created BatchSandbox: {sandbox_id} in namespace: {self.namespace}") return sandbox_id - + except client.exceptions.ApiException as e: if e.status == 409: logger.warning(f"BatchSandbox {sandbox_id} already exists") @@ -589,29 +593,29 @@ async def _create(self, config: DockerDeploymentConfig) -> str: async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[int, int]]: """Get sandbox runtime info (host_ip and port_mapping). - + Args: sandbox_id: ID of the sandbox - + Returns: tuple: (host_ip, port_mapping) - host_ip: Pod IP from endpoints annotation (empty string if not allocated) - port_mapping: Port configuration from annotations - + Raises: Exception: If sandbox not found ValueError: If ports annotation is missing or invalid """ await self._ensure_initialized() - + try: # Get from cache or API Server (handled by api_client) resource = await self._k8s_api.get_custom_object(name=sandbox_id) - + # Extract metadata metadata = resource.get("metadata", {}) annotations = metadata.get("annotations", {}) - + # Parse endpoints from annotations endpoints_str = annotations.get(K8sConstants.ANNOTATION_ENDPOINTS) pod_ips = [] @@ -620,10 +624,10 @@ async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[in pod_ips = json.loads(endpoints_str) except (json.JSONDecodeError, TypeError): logger.warning(f"Failed to parse endpoints for {sandbox_id}: {endpoints_str}") - + # Get pod IP for host_ip if available host_ip = pod_ips[0] if pod_ips else "" - + # Get port configuration from annotations ports_str = annotations.get(K8sConstants.ANNOTATION_PORTS) if not ports_str: @@ -631,34 +635,34 @@ async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[in f"Sandbox '{sandbox_id}' is missing required '{K8sConstants.ANNOTATION_PORTS}' annotation. " f"This sandbox may have been created with an older version." ) - + try: ports_config = json.loads(ports_str) except (json.JSONDecodeError, TypeError) as e: raise ValueError( f"Failed to parse ports annotation for sandbox '{sandbox_id}': {ports_str}. Error: {e}" ) - + # Build port_mapping port_mapping = { - Port.PROXY: ports_config['proxy'], - Port.SERVER: ports_config['server'], - Port.SSH: ports_config['ssh'], + Port.PROXY: ports_config["proxy"], + Port.SERVER: ports_config["server"], + Port.SSH: ports_config["ssh"], } - + return host_ip, port_mapping - + except Exception as e: logger.error(f"Failed to fetch resource from cache for {sandbox_id}: {e}", exc_info=True) raise def _build_runtime(self, host_ip: str, port_mapping: dict[int, int]) -> RemoteSandboxRuntime: """Build runtime for communicating with the sandbox. - + Args: host_ip: Pod IP address port_mapping: Port mapping configuration - + Returns: RemoteSandboxRuntime instance """ @@ -668,4 +672,3 @@ def _build_runtime(self, host_ip: str, port_mapping: dict[int, int]) -> RemoteSa port=proxy_port, ) return RemoteSandboxRuntime.from_config(runtime_config) - diff --git a/rock/sandbox/operator/k8s/template_loader.py b/rock/sandbox/operator/k8s/template_loader.py index 26790e2c2b..57d58e658c 100644 --- a/rock/sandbox/operator/k8s/template_loader.py +++ b/rock/sandbox/operator/k8s/template_loader.py @@ -2,10 +2,7 @@ import copy import json -from pathlib import Path -from typing import Any, Dict - -import yaml +from typing import Any from rock.logger import init_logger from rock.sandbox.operator.k8s.constants import K8sConstants @@ -15,167 +12,160 @@ class K8sTemplateLoader: """Loader for K8S BatchSandbox templates.""" - - def __init__(self, templates: Dict[str, Dict[str, Any]], default_namespace: str = 'rock'): + + def __init__(self, templates: dict[str, dict[str, Any]], default_namespace: str = "rock"): """Initialize template loader. - + Args: templates: Dictionary of template configurations from K8sConfig default_namespace: Default namespace if template doesn't specify one """ - self._templates: Dict[str, Dict[str, Any]] = templates + self._templates: dict[str, dict[str, Any]] = templates self._default_namespace = default_namespace - + if not self._templates: - raise ValueError( - "No templates provided. At least one template must be defined in K8sConfig.templates." - ) - + raise ValueError("No templates provided. At least one template must be defined in K8sConfig.templates.") + logger.info(f"Loaded {len(self._templates)} K8S templates from config") logger.debug(f"Available templates: {', '.join(self._templates.keys())}") - - def get_template(self, template_name: str = 'default') -> Dict[str, Any]: + + def get_template(self, template_name: str = "default") -> dict[str, Any]: """Get a template by name. - + Args: template_name: Name of the template - + Returns: Deep copy of the template dictionary - + Raises: ValueError: If template not found """ if template_name not in self._templates: - available = ', '.join(self._templates.keys()) + available = ", ".join(self._templates.keys()) raise ValueError(f"Template '{template_name}' not found. Available: {available}") - + return copy.deepcopy(self._templates[template_name]) - + def build_manifest( self, - template_name: str = 'default', + template_name: str = "default", sandbox_id: str = None, image: str = None, cpus: float = None, memory: str = None, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Build a complete BatchSandbox manifest from template. - + Template structure: - namespace: K8S namespace for the sandbox (REQUIRED) - ports: custom port configuration (not part of K8S manifest) - template: corresponds to spec.template in BatchSandbox CRD - template.metadata -> spec.template.metadata - template.spec -> spec.template.spec (Pod spec) - + Top-level fields are hardcoded: - apiVersion: sandbox.opensandbox.io/v1alpha1 - kind: BatchSandbox - metadata: constructed from parameters - spec.replicas: always 1 - + Args: template_name: Name of the template to use sandbox_id: Sandbox identifier image: Container image cpus: CPU resource limit memory: Memory resource limit (normalized format like '2Gi') - + Returns: Complete BatchSandbox manifest """ import uuid - + # Get template configuration config = self.get_template(template_name) - + # Use default namespace (configured at startup) namespace = self._default_namespace - + # Get enable_resource_speedup from template (default to True) - enable_resource_speedup = config.get('enable_resource_speedup', True) - + enable_resource_speedup = config.get("enable_resource_speedup", True) + # Get port configuration from template (required) - ports_config = config.get('ports') + ports_config = config.get("ports") if not ports_config: raise ValueError( f"Template '{template_name}' is missing required 'ports' configuration. " f"Each template must define ports (proxy, server, ssh)." ) - + # Extract template (corresponds to spec.template in BatchSandbox) - pod_template = config.get('template', {}) - template_metadata = copy.deepcopy(pod_template.get('metadata', {})) - pod_spec = copy.deepcopy(pod_template.get('spec', {})) - + pod_template = config.get("template", {}) + template_metadata = copy.deepcopy(pod_template.get("metadata", {})) + pod_spec = copy.deepcopy(pod_template.get("spec", {})) + # Generate sandbox_id if not provided if not sandbox_id: sandbox_id = f"sandbox-{uuid.uuid4().hex[:8]}" - + # Build top-level BatchSandbox manifest (hardcoded structure) manifest = { - 'apiVersion': K8sConstants.CRD_API_VERSION, - 'kind': K8sConstants.CRD_KIND, - 'metadata': { - 'name': sandbox_id, - 'namespace': namespace, - 'labels': { + "apiVersion": K8sConstants.CRD_API_VERSION, + "kind": K8sConstants.CRD_KIND, + "metadata": { + "name": sandbox_id, + "namespace": namespace, + "labels": { K8sConstants.LABEL_SANDBOX_ID: sandbox_id, K8sConstants.LABEL_TEMPLATE: template_name, }, - 'annotations': { + "annotations": { K8sConstants.ANNOTATION_PORTS: json.dumps(ports_config), - } + }, + }, + "spec": { + "replicas": 1, # Always 1 for sandbox + "template": {"metadata": template_metadata, "spec": pod_spec}, }, - 'spec': { - 'replicas': 1, # Always 1 for sandbox - 'template': { - 'metadata': template_metadata, - 'spec': pod_spec - } - } } - + # Add resource speedup label if enabled if enable_resource_speedup: - manifest['metadata']['labels'][K8sConstants.LABEL_RESOURCE_SPEEDUP] = 'true' - + manifest["metadata"]["labels"][K8sConstants.LABEL_RESOURCE_SPEEDUP] = "true" + # Add sandbox-id label to template metadata - if 'labels' not in manifest['spec']['template']['metadata']: - manifest['spec']['template']['metadata']['labels'] = {} - manifest['spec']['template']['metadata']['labels'][K8sConstants.LABEL_SANDBOX_ID] = sandbox_id - + if "labels" not in manifest["spec"]["template"]["metadata"]: + manifest["spec"]["template"]["metadata"]["labels"] = {} + manifest["spec"]["template"]["metadata"]["labels"][K8sConstants.LABEL_SANDBOX_ID] = sandbox_id + # Set container image if image: - containers = pod_spec.get('containers', []) + containers = pod_spec.get("containers", []) if containers and len(containers) > 0: - containers[0]['image'] = image - + containers[0]["image"] = image + # Set resources if provided if cpus is not None or memory is not None: - containers = pod_spec.get('containers', []) + containers = pod_spec.get("containers", []) if containers and len(containers) > 0: - if 'resources' not in containers[0]: - containers[0]['resources'] = {} - + if "resources" not in containers[0]: + containers[0]["resources"] = {} + if cpus is not None or memory is not None: - containers[0]['resources']['requests'] = {} - containers[0]['resources']['limits'] = {} - + containers[0]["resources"]["requests"] = {} + containers[0]["resources"]["limits"] = {} + if cpus is not None: - containers[0]['resources']['requests']['cpu'] = str(cpus) - containers[0]['resources']['limits']['cpu'] = str(cpus) - + containers[0]["resources"]["requests"]["cpu"] = str(cpus) + containers[0]["resources"]["limits"]["cpu"] = str(cpus) + if memory is not None: - containers[0]['resources']['requests']['memory'] = memory - containers[0]['resources']['limits']['memory'] = memory - + containers[0]["resources"]["requests"]["memory"] = memory + containers[0]["resources"]["limits"]["memory"] = memory + return manifest - + @property def available_templates(self) -> list[str]: """Get list of available template names.""" return list(self._templates.keys()) - - diff --git a/rock/sdk/agent/job.py b/rock/sdk/agent/job.py index e8997451e9..fc01dfb2ae 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/agent/job.py @@ -266,7 +266,14 @@ async def _collect_results(self) -> JobResult: # ------------------------------------------------------------------ async def _autofill_sandbox_info(self) -> None: - self._config.namespace = self._sandbox._namespace + sandbox_ns = self._sandbox._namespace + if self._config.namespace is not None and sandbox_ns is not None: + if self._config.namespace != sandbox_ns: + raise ValueError( + f"namespace mismatch: JobConfig has '{self._config.namespace}', but sandbox returned '{sandbox_ns}'" + ) + if sandbox_ns is not None: + self._config.namespace = sandbox_ns async def _upload_content(self, content: str, sandbox_path: str) -> None: """Write text content to a local temp file and upload to sandbox via upload_by_path.""" diff --git a/rock/sdk/agent/models/job/config.py b/rock/sdk/agent/models/job/config.py index bae5efe419..839c1aaee2 100644 --- a/rock/sdk/agent/models/job/config.py +++ b/rock/sdk/agent/models/job/config.py @@ -17,7 +17,6 @@ from rock.sdk.agent.models.trial.config import ( AgentConfig, ArtifactConfig, - EnvironmentConfig, OssMirrorConfig, RockEnvironmentConfig, TaskConfig, @@ -142,11 +141,10 @@ class JobConfig(BaseModel): # ── Harbor native fields ── namespace: str | None = Field( default=None, - description="资源租户隔离标识，用于区分不同团队/项目的资源", + description="Tenant isolation identifier for distinguishing resources across teams/projects", ) - experiment_id: str | None = Field( - default=None, - description="实验标识", + experiment_id: str = Field( + description="Experiment identifier, required", ) job_name: str = Field(default_factory=lambda: datetime.now().strftime("%Y-%m-%d__%H-%M-%S")) jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" @@ -165,6 +163,25 @@ class JobConfig(BaseModel): tasks: list[TaskConfig] = Field(default_factory=list) artifacts: list[str | ArtifactConfig] = Field(default_factory=list) + @model_validator(mode="after") + def _sync_experiment_id(self): + """Validate and sync experiment_id between JobConfig and SandboxConfig. + + 1. experiment_id must not be empty. + 2. If environment.experiment_id is already set, it must match. + 3. Propagate experiment_id down to environment (SandboxConfig). + """ + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + env_exp = self.environment.experiment_id + if env_exp is not None and env_exp != self.experiment_id: + raise ValueError( + f"experiment_id mismatch: JobConfig has '{self.experiment_id}', " + f"but environment (SandboxConfig) has '{env_exp}'" + ) + self.environment.experiment_id = self.experiment_id + return self + def to_harbor_yaml(self) -> str: """Serialize Harbor-native fields to YAML for ``harbor jobs start -c``. diff --git a/rock/sdk/envs/rock_env.py b/rock/sdk/envs/rock_env.py index 46e9b7192e..97dfa6e866 100644 --- a/rock/sdk/envs/rock_env.py +++ b/rock/sdk/envs/rock_env.py @@ -8,6 +8,7 @@ logger = init_logger(__name__) + class RockEnv(Env): def __init__(self, env_id: str) -> None: """ @@ -136,7 +137,7 @@ def _call_admin_api(self, endpoint: str, params: dict[str, Any]) -> dict[str, An timeout = httpx.Timeout(timeout=300.0, connect=300.0, read=300.0) try: logger.debug(f"Calling Admin API {url} with params: {params}") - + with httpx.Client(timeout=timeout) as client: response = client.post(url, headers=headers, json=params) response.raise_for_status() diff --git a/rock/sdk/model/server/file_handler.py b/rock/sdk/model/server/file_handler.py index 38400ccc6a..947a5f4b90 100644 --- a/rock/sdk/model/server/file_handler.py +++ b/rock/sdk/model/server/file_handler.py @@ -1,4 +1,5 @@ """File handler for reading/writing LLM requests and responses.""" + import asyncio import json import logging diff --git a/rock/sdk/model/server/main.py b/rock/sdk/model/server/main.py index 62745eba92..7f8dabebe2 100644 --- a/rock/sdk/model/server/main.py +++ b/rock/sdk/model/server/main.py @@ -1,4 +1,5 @@ """LLM Service - FastAPI server for sandbox communication.""" + import argparse import asyncio from contextlib import asynccontextmanager diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 0e2ed8496c..89de7dbd59 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -195,7 +195,7 @@ async def start(self): self._sandbox_id = response.get("result").get("sandbox_id") self._host_name = response.get("result").get("host_name") self._host_ip = response.get("result").get("host_ip") - + start_time = time.time() while time.time() - start_time < self.config.startup_timeout: sandbox_info = await self.get_status() diff --git a/rock/utils/providers/nacos_provider.py b/rock/utils/providers/nacos_provider.py index 23e2147b17..c029dab0bd 100644 --- a/rock/utils/providers/nacos_provider.py +++ b/rock/utils/providers/nacos_provider.py @@ -28,7 +28,9 @@ def __init__( self.group = group self.config_cache: Any | None = None - self.client = nacos.NacosClient(server_addresses=self.server_addresses, endpoint=self.endpoint, namespace=self.namespace) + self.client = nacos.NacosClient( + server_addresses=self.server_addresses, endpoint=self.endpoint, namespace=self.namespace + ) async def get_config(self) -> Any: """ diff --git a/rock/utils/service.py b/rock/utils/service.py index 9ca4328400..401bf09eca 100644 --- a/rock/utils/service.py +++ b/rock/utils/service.py @@ -1,4 +1,3 @@ - from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.redis_key import alive_sandbox_key from rock.utils.providers.redis_provider import RedisProvider diff --git a/tests/conftest.py b/tests/conftest.py index 45e6a840b5..f06a5b465d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,7 @@ import pytest from rock import env_vars +from rock.utils.docker import DockerUtil # Set test data directories at import time (before test collection triggers # module-level constants in production code like TRAJ_FILE in config.py). @@ -59,9 +60,6 @@ def random_container_name() -> str: return container_name -from rock.utils.docker import DockerUtil - - def pytest_collection_modifyitems(config, items): if not DockerUtil.is_docker_available(): skip_docker = pytest.mark.skip(reason="Docker is not available") diff --git a/tests/integration/sdk/sandbox/test_basic.py b/tests/integration/sdk/sandbox/test_basic.py index 0916c2b5ed..fda538a9a8 100644 --- a/tests/integration/sdk/sandbox/test_basic.py +++ b/tests/integration/sdk/sandbox/test_basic.py @@ -63,9 +63,9 @@ async def test_sandbox_file_operations(admin_remote_server: RemoteServer): # Verify file exists verify_response = await sandbox.arun(cmd=f"ls -la {target_path}", session="bash-session") - assert ( - target_path in verify_response.output or "test.txt" in verify_response.output - ), "File not found after upload" + assert target_path in verify_response.output or "test.txt" in verify_response.output, ( + "File not found after upload" + ) finally: # Clean up local temp file diff --git a/tests/integration/sdk/sandbox/test_file_system.py b/tests/integration/sdk/sandbox/test_file_system.py index 04ddcb28ec..78a73479d8 100644 --- a/tests/integration/sdk/sandbox/test_file_system.py +++ b/tests/integration/sdk/sandbox/test_file_system.py @@ -104,9 +104,9 @@ async def test_download_file(sandbox_instance: Sandbox, monkeypatch): ) assert not response.success, "Download should fail when OSS is disabled" - assert ( - "not enabled" in response.message.lower() - ), f"Error message should mention OSS disabled: {response.message}" + assert "not enabled" in response.message.lower(), ( + f"Error message should mention OSS disabled: {response.message}" + ) logger.info("✓ OSS disabled error handling works correctly") # Setup mocks for remaining tests diff --git a/tests/integration/sdk/sandbox/test_sandbox_images.py b/tests/integration/sdk/sandbox/test_sandbox_images.py index 66242ffac5..80d4189553 100644 --- a/tests/integration/sdk/sandbox/test_sandbox_images.py +++ b/tests/integration/sdk/sandbox/test_sandbox_images.py @@ -1,6 +1,7 @@ """ Verify that sandbox can start and run on each image in the list. """ + import pytest from rock.sdk.sandbox.client import Sandbox @@ -13,20 +14,20 @@ "ubuntu:16.04", "ubuntu:24.04", # alpine - #"alpine:3.23", - #"alpine:3.14", + # "alpine:3.23", + # "alpine:3.14", # nix "nixos/nix:2.20.9", "nixos/nix:2.32.6", ] + @pytest.mark.parametrize( "sandbox_instance", [{"image": img} for img in SANDBOX_IMAGES_TO_CHECK], ids=SANDBOX_IMAGES_TO_CHECK, indirect=True, ) - @pytest.mark.need_admin_and_network @SKIP_IF_NO_DOCKER @pytest.mark.asyncio diff --git a/tests/integration/sdk/sandbox/test_sdk_client.py b/tests/integration/sdk/sandbox/test_sdk_client.py index 70bb34e33d..93697e96be 100644 --- a/tests/integration/sdk/sandbox/test_sdk_client.py +++ b/tests/integration/sdk/sandbox/test_sdk_client.py @@ -77,9 +77,9 @@ async def test_sandbox_get_status(admin_remote_server): end_time = time.time() assert "Failed to pull image" in str(exc_info.value) execution_time = end_time - start_time - assert ( - execution_time < config.startup_timeout - ), f"Execution time {execution_time}s should be less than startup_timeout {config.startup_timeout}s" + assert execution_time < config.startup_timeout, ( + f"Execution time {execution_time}s should be less than startup_timeout {config.startup_timeout}s" + ) @pytest.mark.need_admin diff --git a/tests/integration/sdk/sandbox/test_speedup.py b/tests/integration/sdk/sandbox/test_speedup.py index c9bd5f9cbd..d6d06c3eae 100644 --- a/tests/integration/sdk/sandbox/test_speedup.py +++ b/tests/integration/sdk/sandbox/test_speedup.py @@ -1,5 +1,7 @@ """Tests for sandbox speedup functionality.""" +from urllib.parse import urlparse + import pytest from rock.actions import Command @@ -10,12 +12,33 @@ logger = init_logger(__name__) +MIRROR_ALIYUNCS = "mirrors.cloud.aliyuncs.com" +MIRROR_ALIYUN = "mirrors.aliyun.com" + + +def _extract_hostnames(text: str) -> set[str]: + """Extract hostnames from all URLs found in text using urlparse.""" + return { + urlparse(word).hostname for line in text.splitlines() for word in line.split() if word.startswith("http") + } - {None} + + +def _has_trusted_host(pip_conf: str, expected_host: str) -> bool: + """Check if pip.conf contains a trusted-host line matching the expected hostname.""" + for line in pip_conf.splitlines(): + stripped = line.strip() + if stripped.startswith("trusted-host") and "=" in stripped: + host_value = stripped.split("=", 1)[1].strip() + if urlparse(f"http://{host_value}").hostname == expected_host: + return True + return False + async def _assert_speedup_apt(sandbox: Sandbox): logger.info("Testing APT public mirror configuration...") result = await sandbox.network.speedup( speedup_type=SpeedupType.APT, - speedup_value="http://mirrors.cloud.aliyuncs.com", + speedup_value=f"http://{MIRROR_ALIYUNCS}", ) assert result.exit_code == 0, f"APT public mirror failed: {result.output}" logger.info("APT public mirror configured successfully") @@ -25,12 +48,10 @@ async def _assert_speedup_apt(sandbox: Sandbox): assert check_result.exit_code == 0, "Failed to read /etc/apt/sources.list" sources_content = check_result.stdout + mirror_hosts = _extract_hostnames(sources_content) assert ( - "mirrors.cloud.aliyuncs.com" in sources_content - ), f"Mirror URL not found in sources.list. Content:\n{sources_content}" - assert ( - "deb http://mirrors.cloud.aliyuncs.com" in sources_content - ), f"Expected deb entry not found. Content:\n{sources_content}" + MIRROR_ALIYUNCS in mirror_hosts + ), f"Mirror hostname not found in sources.list URLs. Content:\n{sources_content}" logger.info(f"APT sources.list verified successfully:\n{sources_content}") @@ -38,7 +59,7 @@ async def _assert_speedup_pip(sandbox: Sandbox): logger.info("Testing PIP mirror (http)...") result = await sandbox.network.speedup( speedup_type=SpeedupType.PIP, - speedup_value="http://mirrors.cloud.aliyuncs.com", + speedup_value=f"http://{MIRROR_ALIYUNCS}", ) assert result.exit_code == 0, f"PIP mirror failed: {result.output}" logger.info("PIP mirror configured successfully") @@ -48,18 +69,17 @@ async def _assert_speedup_pip(sandbox: Sandbox): assert check_result.exit_code == 0, "Failed to read /root/.pip/pip.conf" pip_config_content = check_result.stdout - assert ( - "mirrors.cloud.aliyuncs.com/pypi/simple/" in pip_config_content - ), f"PIP mirror URL not found in pip.conf. Content:\n{pip_config_content}" - assert ( - "trusted-host = mirrors.cloud.aliyuncs.com" in pip_config_content + pip_hosts = _extract_hostnames(pip_config_content) + assert MIRROR_ALIYUNCS in pip_hosts, f"PIP mirror hostname not found in pip.conf. Content:\n{pip_config_content}" + assert _has_trusted_host( + pip_config_content, MIRROR_ALIYUNCS ), f"trusted-host not found in pip.conf. Content:\n{pip_config_content}" logger.info(f"PIP config verified successfully:\n{pip_config_content}") logger.info("Testing PIP mirror (https)...") result = await sandbox.network.speedup( speedup_type=SpeedupType.PIP, - speedup_value="https://mirrors.aliyun.com", + speedup_value=f"https://{MIRROR_ALIYUN}", ) assert result.exit_code == 0, f"PIP aliyun mirror failed: {result.output}" logger.info("PIP aliyun mirror configured successfully") @@ -67,9 +87,8 @@ async def _assert_speedup_pip(sandbox: Sandbox): check_result = await sandbox.execute(Command(command=["cat", "/root/.pip/pip.conf"])) assert check_result.exit_code == 0, "Failed to read /root/.pip/pip.conf after updating" pip_config_content = check_result.stdout - assert ( - "mirrors.aliyun.com/pypi/simple/" in pip_config_content - ), f"Updated PIP mirror URL not found. Content:\n{pip_config_content}" + updated_pip_hosts = _extract_hostnames(pip_config_content) + assert MIRROR_ALIYUN in updated_pip_hosts, f"Updated PIP mirror hostname not found. Content:\n{pip_config_content}" async def _assert_speedup_github(sandbox: Sandbox): diff --git a/tests/unit/admin/core/test_ray_service.py b/tests/unit/admin/core/test_ray_service.py index bd39e9e878..d1f0ba8aa6 100644 --- a/tests/unit/admin/core/test_ray_service.py +++ b/tests/unit/admin/core/test_ray_service.py @@ -25,9 +25,11 @@ async def test_reconnect_ray_calls_ray_shutdown_and_init_and_reset_counters(ray_ mock_rwlock.write_lock.return_value = mock_lock service._ray_rwlock = mock_rwlock - with patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, patch( - "rock.admin.core.ray_service.ray.init" - ) as mock_init, patch("time.time", return_value=old_establish_time + 5): + with ( + patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, + patch("rock.admin.core.ray_service.ray.init") as mock_init, + patch("time.time", return_value=old_establish_time + 5), + ): await service._reconnect_ray() mock_rwlock.write_lock.assert_called_once() @@ -63,9 +65,11 @@ async def test_reconnect_ray_skip_when_reader_exists_and_write_lock_timeout(ray_ service._ray_rwlock._readers = 1 - with patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, patch( - "rock.admin.core.ray_service.ray.init" - ) as mock_init, patch("time.time", return_value=old_est + 5): + with ( + patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, + patch("rock.admin.core.ray_service.ray.init") as mock_init, + patch("time.time", return_value=old_est + 5), + ): await service._reconnect_ray() mock_shutdown.assert_not_called() diff --git a/tests/unit/common/test_port_validation.py b/tests/unit/common/test_port_validation.py index b490da5b21..9ebeac63c9 100644 --- a/tests/unit/common/test_port_validation.py +++ b/tests/unit/common/test_port_validation.py @@ -1,10 +1,9 @@ """Tests for common port validation utilities.""" -import pytest from rock.common.port_validation import ( - PORT_FORWARD_MIN_PORT, - PORT_FORWARD_MAX_PORT, PORT_FORWARD_EXCLUDED_PORTS, + PORT_FORWARD_MAX_PORT, + PORT_FORWARD_MIN_PORT, validate_port_forward_port, ) diff --git a/tests/unit/rocklet/test_docker_deployment.py b/tests/unit/rocklet/test_docker_deployment.py index 91d8649cc1..0926cd1c8e 100644 --- a/tests/unit/rocklet/test_docker_deployment.py +++ b/tests/unit/rocklet/test_docker_deployment.py @@ -9,6 +9,7 @@ ) from rock.deployments.config import DockerDeploymentConfig, get_deployment + @pytest.mark.need_docker async def test_docker_deployment(container_name): deployment_config = DockerDeploymentConfig( diff --git a/tests/unit/rocklet/test_portforward.py b/tests/unit/rocklet/test_portforward.py index d098153b2a..2fce787ead 100644 --- a/tests/unit/rocklet/test_portforward.py +++ b/tests/unit/rocklet/test_portforward.py @@ -1,4 +1,5 @@ """Tests for rocklet portforward WebSocket endpoint.""" + import pytest from fastapi import FastAPI @@ -19,4 +20,4 @@ class TestPortForwardEndpoint: def test_portforward_endpoint_exists(self, app: FastAPI): """Test that /portforward endpoint is registered.""" routes = [route.path for route in app.routes] - assert "/portforward" in routes, "Portforward endpoint should be registered" \ No newline at end of file + assert "/portforward" in routes, "Portforward endpoint should be registered" diff --git a/tests/unit/sandbox/operator/test_k8s_api_client.py b/tests/unit/sandbox/operator/test_k8s_api_client.py index e206e44226..3285478391 100644 --- a/tests/unit/sandbox/operator/test_k8s_api_client.py +++ b/tests/unit/sandbox/operator/test_k8s_api_client.py @@ -18,7 +18,7 @@ class TestK8sApiClient: def test_initialization(self, mock_api_client): """Test K8sApiClient initialization. - + Verifies AsyncLimiter is configured with QPS limit. """ api_client = K8sApiClient( @@ -31,7 +31,7 @@ def test_initialization(self, mock_api_client): watch_timeout_seconds=60, watch_reconnect_delay_seconds=5, ) - + assert api_client._group == "sandbox.opensandbox.io" assert api_client._version == "v1alpha1" assert api_client._plural == "batchsandboxes" @@ -43,15 +43,15 @@ def test_initialization(self, mock_api_client): @pytest.mark.asyncio async def test_rate_limiting_with_context_manager(self, k8s_api_client): """Test AsyncLimiter integration in CRUD operations. - + Verifies that AsyncLimiter is properly used as context manager to enforce rate limiting on API Server requests. """ - with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread: + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = {"created": True} - + result = await k8s_api_client.create_custom_object(body={"test": "data"}) - + assert result == {"created": True} mock_thread.assert_awaited_once() @@ -63,86 +63,84 @@ async def test_create_custom_object(self, k8s_api_client): "kind": "BatchSandbox", "metadata": {"name": "test-sandbox"}, } - - with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread: + + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = {"created": True} - + result = await k8s_api_client.create_custom_object(body=mock_body) - + assert result == {"created": True} mock_thread.assert_awaited_once() @pytest.mark.asyncio async def test_get_custom_object_from_cache(self, k8s_api_client): """Test cache hit scenario (Informer pattern). - + When resource exists in local cache, no API Server request is made. """ - k8s_api_client._cache = { - "test-sandbox": {"metadata": {"name": "test-sandbox"}} - } - + k8s_api_client._cache = {"test-sandbox": {"metadata": {"name": "test-sandbox"}}} + result = await k8s_api_client.get_custom_object(name="test-sandbox") - + assert result == {"metadata": {"name": "test-sandbox"}} @pytest.mark.asyncio async def test_get_custom_object_cache_miss(self, k8s_api_client): """Test cache miss scenario with API Server fallback. - + When resource not in cache, queries API Server and updates cache. """ k8s_api_client._cache = {} mock_response = {"metadata": {"name": "test-sandbox"}} - - with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread: + + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = mock_response - + result = await k8s_api_client.get_custom_object(name="test-sandbox") - + assert result == mock_response assert k8s_api_client._cache["test-sandbox"] == mock_response @pytest.mark.asyncio async def test_delete_custom_object(self, k8s_api_client): """Test deleting K8s custom resource with rate limiting.""" - with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread: + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = {"status": "deleted"} - + result = await k8s_api_client.delete_custom_object(name="test-sandbox") - + assert result == {"status": "deleted"} mock_thread.assert_awaited_once() @pytest.mark.asyncio async def test_start_initializes_watch(self, k8s_api_client): """Test watch task initialization for Informer pattern. - + start() creates background task to watch K8s resource changes and sync them to local cache. """ - with patch('asyncio.create_task') as mock_create_task: + with patch("asyncio.create_task") as mock_create_task: await k8s_api_client.start() - + assert k8s_api_client._initialized is True mock_create_task.assert_called_once() @pytest.mark.asyncio async def test_start_idempotent(self, k8s_api_client): """Test start() idempotency. - + Multiple start() calls should only initialize watch once. """ - with patch('asyncio.create_task') as mock_create_task: + with patch("asyncio.create_task") as mock_create_task: await k8s_api_client.start() await k8s_api_client.start() - + assert mock_create_task.call_count == 1 @pytest.mark.asyncio async def test_list_and_sync_cache(self, k8s_api_client): """Test initial cache sync from K8s API Server. - + Populates local cache with all resources and returns resourceVersion for subsequent watch operations. """ @@ -151,14 +149,14 @@ async def test_list_and_sync_cache(self, k8s_api_client): "items": [ {"metadata": {"name": "sandbox-1"}}, {"metadata": {"name": "sandbox-2"}}, - ] + ], } - - with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread: + + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = mock_resources - + resource_version = await k8s_api_client._list_and_sync_cache() - + assert resource_version == "12345" assert len(k8s_api_client._cache) == 2 assert "sandbox-1" in k8s_api_client._cache diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 17729b2666..103a55c2a4 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -26,6 +26,7 @@ def test_initialization_without_templates(self): # but provider creation should fail with pytest.raises(ValueError, match="No templates provided"): from rock.sandbox.operator.k8s.provider import BatchSandboxProvider + BatchSandboxProvider(k8s_config=config) @pytest.mark.asyncio @@ -57,9 +58,7 @@ async def test_submit_success(self, k8s_operator, mock_provider, deployment_conf async def test_submit_no_host_ip(self, k8s_operator, mock_provider, deployment_config): """Test submission fails when no host IP is allocated.""" # Mock provider to raise exception - mock_provider.submit = AsyncMock( - side_effect=Exception("Failed to get host IP for sandbox test-sandbox") - ) + mock_provider.submit = AsyncMock(side_effect=Exception("Failed to get host IP for sandbox test-sandbox")) with pytest.raises(Exception, match="Failed to get host IP"): await k8s_operator.submit(deployment_config) @@ -114,9 +113,7 @@ async def test_get_status_not_alive(self, k8s_operator, mock_provider): @pytest.mark.asyncio async def test_get_status_not_found(self, k8s_operator, mock_provider): """Test status retrieval when sandbox not found in cache.""" - mock_provider.get_status = AsyncMock( - side_effect=Exception("Sandbox test-sandbox not found") - ) + mock_provider.get_status = AsyncMock(side_effect=Exception("Sandbox test-sandbox not found")) with pytest.raises(Exception, match="not found"): await k8s_operator.get_status("test-sandbox") @@ -148,7 +145,5 @@ async def test_stop_failure(self, k8s_operator, mock_provider): mock_provider.stop = AsyncMock(return_value=False) result = await k8s_operator.stop("test-sandbox") - - assert result is False - + assert result is False diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index 7288d80ee9..68d206a683 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -5,7 +5,6 @@ from rock.sandbox.operator.k8s.constants import K8sConstants from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, ResourceMatchingPoolSelector - BASIC_TEMPLATES = { "default": { "namespace": "rock-test", @@ -28,7 +27,13 @@ def make_provider(template_map: dict = None) -> BatchSandboxProvider: ) -def make_config(image: str = "python:3.11", cpus: float = 2, memory: str = "4Gi", extended_params: dict = None, image_os: str = "linux") -> DockerDeploymentConfig: +def make_config( + image: str = "python:3.11", + cpus: float = 2, + memory: str = "4Gi", + extended_params: dict = None, + image_os: str = "linux", +) -> DockerDeploymentConfig: return DockerDeploymentConfig( image=image, cpus=cpus, @@ -113,34 +118,38 @@ class TestGetPoolName: async def test_returns_pool_from_extended_params(self): """extended_params 中有 pool_name 时直接返回，不走 selector。""" provider = make_provider() - provider.set_nacos_provider(MockNacosProvider({ - K8sConstants.NACOS_POOLS_KEY: { - "pool_nacos": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"} - } - })) + provider.set_nacos_provider( + MockNacosProvider( + {K8sConstants.NACOS_POOLS_KEY: {"pool_nacos": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"}}} + ) + ) config = make_config(extended_params={"pool_name": "my_pool"}) assert await provider._get_pool_name(config) == "my_pool" async def test_extended_params_takes_priority_over_selector(self): """extended_params 优先级高于 selector。""" provider = make_provider() - provider.set_nacos_provider(MockNacosProvider({ - K8sConstants.NACOS_POOLS_KEY: { - "pool_auto": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"} - } - })) + provider.set_nacos_provider( + MockNacosProvider( + {K8sConstants.NACOS_POOLS_KEY: {"pool_auto": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"}}} + ) + ) config = make_config(extended_params={"pool_name": "explicit_pool"}) assert await provider._get_pool_name(config) == "explicit_pool" async def test_uses_selector_when_no_extended_params(self): """无 extended_params 时使用 selector 选择 pool。""" provider = make_provider() - provider.set_nacos_provider(MockNacosProvider({ - K8sConstants.NACOS_POOLS_KEY: { - "pool_small": {"image": "python:3.11", "cpus": 2, "memory": "4Gi"}, - "pool_large": {"image": "python:3.11", "cpus": 8, "memory": "16Gi"}, - } - })) + provider.set_nacos_provider( + MockNacosProvider( + { + K8sConstants.NACOS_POOLS_KEY: { + "pool_small": {"image": "python:3.11", "cpus": 2, "memory": "4Gi"}, + "pool_large": {"image": "python:3.11", "cpus": 8, "memory": "16Gi"}, + } + } + ) + ) config = make_config(image="python:3.11", cpus=2, memory="4Gi") assert await provider._get_pool_name(config) == "pool_small" @@ -206,16 +215,20 @@ class TestGetPoolPorts: async def test_returns_ports_from_pool_config(self): """从 PoolConfig 中获取端口配置。""" provider = make_provider() - provider.set_nacos_provider(MockNacosProvider({ - K8sConstants.NACOS_POOLS_KEY: { - "pool_custom": { - "image": "python:3.11", - "cpus": 4, - "memory": "8Gi", - "ports": {"proxy": 9000, "server": 9090, "ssh": 2222} + provider.set_nacos_provider( + MockNacosProvider( + { + K8sConstants.NACOS_POOLS_KEY: { + "pool_custom": { + "image": "python:3.11", + "cpus": 4, + "memory": "8Gi", + "ports": {"proxy": 9000, "server": 9090, "ssh": 2222}, + } + } } - } - })) + ) + ) ports = await provider._get_pool_ports("pool_custom") assert ports == {"proxy": 9000, "server": 9090, "ssh": 2222} @@ -229,11 +242,11 @@ async def test_returns_default_ports_when_pool_not_found(self): async def test_returns_default_ports_for_pool_without_ports(self): """PoolConfig 未配置 ports 时由 __post_init__ 自动补全默认值。""" provider = make_provider() - provider.set_nacos_provider(MockNacosProvider({ - K8sConstants.NACOS_POOLS_KEY: { - "pool_no_ports": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"} - } - })) + provider.set_nacos_provider( + MockNacosProvider( + {K8sConstants.NACOS_POOLS_KEY: {"pool_no_ports": {"image": "python:3.11", "cpus": 4, "memory": "8Gi"}}} + ) + ) ports = await provider._get_pool_ports("pool_no_ports") # PoolConfig.__post_init__ fills in default ports, so provider returns them assert ports == {"proxy": 8000, "server": 8080, "ssh": 22} @@ -244,10 +257,10 @@ async def test_returns_default_ports_for_pool_without_ports(self): class MockNacosProvider: """Mock Nacos provider for testing.""" - + def __init__(self, config: dict = None): self._config = config - + async def get_config(self): return self._config @@ -261,49 +274,43 @@ async def test_get_pools_from_nacos(self): "image": "python:3.11", "cpus": 4, "memory": "8Gi", - "ports": {"proxy": 9000, "server": 9090, "ssh": 2222} + "ports": {"proxy": 9000, "server": 9090, "ssh": 2222}, } } } provider = make_provider() provider.set_nacos_provider(MockNacosProvider(nacos_config)) - + pools = await provider._get_pools() assert "pool_nacos" in pools assert pools["pool_nacos"].image == "python:3.11" assert pools["pool_nacos"].cpus == 4 assert pools["pool_nacos"].ports == {"proxy": 9000, "server": 9090, "ssh": 2222} - + async def test_returns_empty_when_no_nacos_provider(self): """无 nacos provider 时返回空字典。""" provider = make_provider() # No nacos provider set - + pools = await provider._get_pools() assert pools == {} - + async def test_returns_empty_when_nacos_has_no_pools(self): """Nacos 无 pools 配置时返回空字典。""" provider = make_provider() provider.set_nacos_provider(MockNacosProvider({"other_key": "value"})) - + pools = await provider._get_pools() assert pools == {} - + async def test_pool_selection_uses_nacos_pools(self): """Pool 选择使用 Nacos 中的 pools。""" nacos_config = { - K8sConstants.NACOS_POOLS_KEY: { - "pool_nacos": { - "image": "python:3.11", - "cpus": 2, - "memory": "4Gi" - } - } + K8sConstants.NACOS_POOLS_KEY: {"pool_nacos": {"image": "python:3.11", "cpus": 2, "memory": "4Gi"}} } provider = make_provider() provider.set_nacos_provider(MockNacosProvider(nacos_config)) - + config = make_config(image="python:3.11", cpus=2, memory="4Gi") pool_name = await provider._get_pool_name(config) assert pool_name == "pool_nacos" diff --git a/tests/unit/sandbox/operator/test_k8s_template_loader.py b/tests/unit/sandbox/operator/test_k8s_template_loader.py index 67cd1145a8..5d5edfe7fc 100644 --- a/tests/unit/sandbox/operator/test_k8s_template_loader.py +++ b/tests/unit/sandbox/operator/test_k8s_template_loader.py @@ -11,11 +11,8 @@ class TestK8sTemplateLoader: def test_initialization_success(self, basic_templates): """Test successful template loader initialization.""" - loader = K8sTemplateLoader( - templates=basic_templates, - default_namespace="rock-test" - ) - + loader = K8sTemplateLoader(templates=basic_templates, default_namespace="rock-test") + assert loader._default_namespace == "rock-test" assert len(loader._templates) == 1 assert "default" in loader.available_templates @@ -28,7 +25,7 @@ def test_initialization_without_templates(self): def test_get_template_success(self, template_loader): """Test getting template by name.""" template = template_loader.get_template("default") - + assert template is not None assert "ports" in template assert "template" in template @@ -43,36 +40,32 @@ def test_get_template_returns_copy(self, template_loader): """Test that get_template returns a deep copy.""" template1 = template_loader.get_template("default") template2 = template_loader.get_template("default") - + # Modify first template template1["ports"]["proxy"] = 9999 - + # Second template should not be affected assert template2["ports"]["proxy"] == 8000 def test_build_manifest_basic(self, template_loader): """Test building basic manifest.""" manifest = template_loader.build_manifest( - template_name="default", - sandbox_id="test-sandbox", - image="python:3.11", - cpus=2.0, - memory="4Gi" + template_name="default", sandbox_id="test-sandbox", image="python:3.11", cpus=2.0, memory="4Gi" ) - + # Verify top-level structure assert manifest["apiVersion"] == K8sConstants.CRD_API_VERSION assert manifest["kind"] == K8sConstants.CRD_KIND assert manifest["metadata"]["name"] == "test-sandbox" assert manifest["metadata"]["namespace"] == "rock-test" - + # Verify labels assert manifest["metadata"]["labels"][K8sConstants.LABEL_SANDBOX_ID] == "test-sandbox" assert manifest["metadata"]["labels"][K8sConstants.LABEL_TEMPLATE] == "default" - + # Verify annotations (ports stored as JSON) assert K8sConstants.ANNOTATION_PORTS in manifest["metadata"]["annotations"] - + # Verify spec assert manifest["spec"]["replicas"] == 1 assert "template" in manifest["spec"] @@ -80,14 +73,11 @@ def test_build_manifest_basic(self, template_loader): def test_build_manifest_with_resources(self, template_loader): """Test building manifest with CPU and memory resources.""" manifest = template_loader.build_manifest( - template_name="default", - sandbox_id="test-sandbox", - cpus=4.0, - memory="8Gi" + template_name="default", sandbox_id="test-sandbox", cpus=4.0, memory="8Gi" ) - + container = manifest["spec"]["template"]["spec"]["containers"][0] - + # Verify resource requests and limits assert container["resources"]["requests"]["cpu"] == "4.0" assert container["resources"]["limits"]["cpu"] == "4.0" @@ -100,35 +90,27 @@ def test_build_manifest_without_resources(self, template_loader): template_name="default", sandbox_id="test-sandbox", ) - + container = manifest["spec"]["template"]["spec"]["containers"][0] - + # Should not have resources section if not specified assert "resources" not in container or not container.get("resources") def test_build_manifest_with_custom_image(self, template_loader): """Test building manifest with custom image.""" manifest = template_loader.build_manifest( - template_name="default", - sandbox_id="test-sandbox", - image="ubuntu:22.04" + template_name="default", sandbox_id="test-sandbox", image="ubuntu:22.04" ) - + container = manifest["spec"]["template"]["spec"]["containers"][0] assert container["image"] == "ubuntu:22.04" def test_build_manifest_missing_ports_in_template(self): """Test building manifest fails when template lacks ports config.""" - templates = { - "no-ports": { - "template": { - "spec": {"containers": [{"name": "main"}]} - } - } - } - + templates = {"no-ports": {"template": {"spec": {"containers": [{"name": "main"}]}}}} + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") - + with pytest.raises(ValueError, match="missing required 'ports' configuration"): loader.build_manifest(template_name="no-ports", sandbox_id="test") @@ -138,21 +120,19 @@ def test_build_manifest_with_resource_speedup(self): "speedup": { "enable_resource_speedup": True, "ports": {"proxy": 8000, "server": 8080, "ssh": 22}, - "template": { - "spec": {"containers": [{"name": "main"}]} - } + "template": {"spec": {"containers": [{"name": "main"}]}}, } } - + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") manifest = loader.build_manifest(template_name="speedup", sandbox_id="test") - + assert manifest["metadata"]["labels"][K8sConstants.LABEL_RESOURCE_SPEEDUP] == "true" def test_build_manifest_auto_generate_sandbox_id(self, template_loader): """Test building manifest auto-generates sandbox_id if not provided.""" manifest = template_loader.build_manifest(template_name="default") - + sandbox_id = manifest["metadata"]["name"] assert sandbox_id.startswith("sandbox-") assert len(sandbox_id) > 8 # Should have UUID suffix @@ -160,16 +140,13 @@ def test_build_manifest_auto_generate_sandbox_id(self, template_loader): def test_available_templates_property(self, template_loader): """Test available_templates property.""" templates = template_loader.available_templates - + assert isinstance(templates, list) assert "default" in templates def test_build_manifest_adds_sandbox_id_to_pod_labels(self, template_loader): """Test that sandbox-id label is added to pod template.""" - manifest = template_loader.build_manifest( - template_name="default", - sandbox_id="test-sandbox" - ) - + manifest = template_loader.build_manifest(template_name="default", sandbox_id="test-sandbox") + pod_labels = manifest["spec"]["template"]["metadata"]["labels"] assert pod_labels[K8sConstants.LABEL_SANDBOX_ID] == "test-sandbox" diff --git a/tests/unit/sandbox/test_sandbox_http_proxy.py b/tests/unit/sandbox/test_sandbox_http_proxy.py index 1bb2a539f0..b1beb45f1b 100644 --- a/tests/unit/sandbox/test_sandbox_http_proxy.py +++ b/tests/unit/sandbox/test_sandbox_http_proxy.py @@ -92,6 +92,7 @@ async def start_echo_server_in_sandbox( # Wait for server to be ready await asyncio.sleep(2) + @pytest.mark.need_docker @pytest.mark.need_ray @pytest.mark.asyncio diff --git a/tests/unit/sandbox/test_sandbox_manager.py b/tests/unit/sandbox/test_sandbox_manager.py index fb69a486fb..91afcd2af2 100644 --- a/tests/unit/sandbox/test_sandbox_manager.py +++ b/tests/unit/sandbox/test_sandbox_manager.py @@ -260,9 +260,9 @@ async def test_use_standard_spec_only_disabled(sandbox_manager): # Verify that requested spec was used (not standard spec) assert sandbox_info["cpus"] == requested_cpus, f"Expected cpus={requested_cpus}, but got {sandbox_info['cpus']}" - assert ( - sandbox_info["memory"] == requested_memory - ), f"Expected memory='{requested_memory}', but got {sandbox_info['memory']}" + assert sandbox_info["memory"] == requested_memory, ( + f"Expected memory='{requested_memory}', but got {sandbox_info['memory']}" + ) # Also verify sandbox is alive is_alive = await sandbox_manager._is_actor_alive(sandbox_id) diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index 193e19334d..55603c5771 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -8,6 +8,7 @@ from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from tests.unit.conftest import check_sandbox_status_until_alive + @pytest.mark.need_docker @pytest.mark.need_ray @pytest.mark.asyncio @@ -37,6 +38,7 @@ async def test_batch_get_sandbox_status(sandbox_manager: SandboxManager, sandbox for sandbox_id in sandbox_ids: await sandbox_manager.stop(sandbox_id) + @pytest.mark.need_docker @pytest.mark.need_ray @pytest.mark.asyncio diff --git a/tests/unit/sandbox/test_websocket_proxy_subprotocol.py b/tests/unit/sandbox/test_websocket_proxy_subprotocol.py index 6818beea99..5047ea91f7 100644 --- a/tests/unit/sandbox/test_websocket_proxy_subprotocol.py +++ b/tests/unit/sandbox/test_websocket_proxy_subprotocol.py @@ -1,12 +1,10 @@ """Tests for WebSocket proxy subprotocol forwarding and performance fixes.""" -import asyncio -from unittest.mock import AsyncMock, MagicMock, patch, call -import pytest +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService - # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── @@ -140,9 +138,7 @@ async def test_negotiated_subprotocol_passed_to_client_accept(self): # accept 必须带上协商好的子协议 client_ws.accept.assert_called_once() call_kwargs = client_ws.accept.call_args - subprotocol = call_kwargs.kwargs.get("subprotocol") or ( - call_kwargs.args[0] if call_kwargs.args else None - ) + subprotocol = call_kwargs.kwargs.get("subprotocol") or (call_kwargs.args[0] if call_kwargs.args else None) assert subprotocol == "binary" @@ -175,10 +171,12 @@ async def test_binary_message_forwarded_without_sleep(self): """Binary messages should be forwarded directly without any sleep.""" service = MagicMock(spec=SandboxProxyService) - source_ws = FakeStarletteWebSocket([ - {"type": "websocket.receive", "bytes": b"\x00\x01\x02"}, - {"type": "websocket.disconnect", "code": 1000}, - ]) + source_ws = FakeStarletteWebSocket( + [ + {"type": "websocket.receive", "bytes": b"\x00\x01\x02"}, + {"type": "websocket.disconnect", "code": 1000}, + ] + ) target_ws = MagicMock(spec=["recv", "send"]) target_ws.send = AsyncMock() @@ -208,19 +206,19 @@ async def test_all_binary_frames_forwarded_to_target(self): """Every binary frame from client must reach the upstream target.""" service = MagicMock(spec=SandboxProxyService) - source_ws = FakeStarletteWebSocket([ - {"type": "websocket.receive", "bytes": b"vnc_frame_1"}, - {"type": "websocket.receive", "bytes": b"vnc_frame_2"}, - {"type": "websocket.receive", "bytes": b"vnc_frame_3"}, - {"type": "websocket.disconnect", "code": 1000}, - ]) + source_ws = FakeStarletteWebSocket( + [ + {"type": "websocket.receive", "bytes": b"vnc_frame_1"}, + {"type": "websocket.receive", "bytes": b"vnc_frame_2"}, + {"type": "websocket.receive", "bytes": b"vnc_frame_3"}, + {"type": "websocket.disconnect", "code": 1000}, + ] + ) target_ws = MagicMock(spec=["recv", "send"]) target_ws.send = AsyncMock() - await SandboxProxyService._forward_messages( - service, source_ws, target_ws, "client->target" - ) + await SandboxProxyService._forward_messages(service, source_ws, target_ws, "client->target") forwarded = [c.args[0] for c in target_ws.send.call_args_list] assert forwarded == [b"vnc_frame_1", b"vnc_frame_2", b"vnc_frame_3"] @@ -229,19 +227,19 @@ async def test_mixed_text_and_binary_frames_all_forwarded(self): """Interleaved text and binary frames must all be forwarded correctly.""" service = MagicMock(spec=SandboxProxyService) - source_ws = FakeStarletteWebSocket([ - {"type": "websocket.receive", "text": "hello"}, - {"type": "websocket.receive", "bytes": b"\x00\x01"}, - {"type": "websocket.receive", "text": "world"}, - {"type": "websocket.disconnect", "code": 1000}, - ]) + source_ws = FakeStarletteWebSocket( + [ + {"type": "websocket.receive", "text": "hello"}, + {"type": "websocket.receive", "bytes": b"\x00\x01"}, + {"type": "websocket.receive", "text": "world"}, + {"type": "websocket.disconnect", "code": 1000}, + ] + ) target_ws = MagicMock(spec=["recv", "send"]) target_ws.send = AsyncMock() - await SandboxProxyService._forward_messages( - service, source_ws, target_ws, "client->target" - ) + await SandboxProxyService._forward_messages(service, source_ws, target_ws, "client->target") forwarded = [c.args[0] for c in target_ws.send.call_args_list] assert forwarded == ["hello", b"\x00\x01", "world"] @@ -250,16 +248,16 @@ async def test_single_binary_frame_not_lost(self): """Even a single binary frame must not be silently dropped.""" service = MagicMock(spec=SandboxProxyService) - source_ws = FakeStarletteWebSocket([ - {"type": "websocket.receive", "bytes": b"rfb_handshake"}, - {"type": "websocket.disconnect", "code": 1000}, - ]) + source_ws = FakeStarletteWebSocket( + [ + {"type": "websocket.receive", "bytes": b"rfb_handshake"}, + {"type": "websocket.disconnect", "code": 1000}, + ] + ) target_ws = MagicMock(spec=["recv", "send"]) target_ws.send = AsyncMock() - await SandboxProxyService._forward_messages( - service, source_ws, target_ws, "client->target" - ) + await SandboxProxyService._forward_messages(service, source_ws, target_ws, "client->target") target_ws.send.assert_called_once_with(b"rfb_handshake") diff --git a/tests/unit/sandbox/test_websocket_tcp_proxy.py b/tests/unit/sandbox/test_websocket_tcp_proxy.py index cae9f95709..0c93e5e176 100644 --- a/tests/unit/sandbox/test_websocket_tcp_proxy.py +++ b/tests/unit/sandbox/test_websocket_tcp_proxy.py @@ -1,4 +1,5 @@ """Tests for WebSocket TCP port forwarding proxy.""" + from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -137,7 +138,7 @@ def test_portforward_route_valid_port_calls_service(self, app): # The WebSocket connection will be accepted and the service called # Since the mock returns immediately, the connection should close try: - with client.websocket_connect("/sandboxes/test-sandbox/portforward?port=8080") as websocket: + with client.websocket_connect("/sandboxes/test-sandbox/portforward?port=8080"): # Connection established, service should be called pass except Exception: @@ -160,9 +161,9 @@ def test_portforward_route_invalid_port_rejected(self, app): client = TestClient(app) # The WebSocket connection will be closed with code 1008 - with client.websocket_connect("/sandboxes/test-sandbox/portforward?port=22") as websocket: + with client.websocket_connect("/sandboxes/test-sandbox/portforward?port=22"): # Connection should be closed by server pass # The service should have been called - mock_service.websocket_to_tcp_proxy.assert_called_once() \ No newline at end of file + mock_service.websocket_to_tcp_proxy.assert_called_once() diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index a978d0c7ea..be2f84b2ef 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -106,6 +106,8 @@ def _make_mock_sandbox(): """Create a mock Sandbox with all required async methods.""" sandbox = AsyncMock() sandbox.sandbox_id = "sb-123" + sandbox._namespace = "test-ns" + sandbox._experiment_id = "test-exp" sandbox.start = AsyncMock() sandbox.close = AsyncMock() sandbox.create_session = AsyncMock() @@ -151,7 +153,7 @@ def _make_mock_sandbox(): class TestJob: def test_init_requires_jobconfig(self): - config = JobConfig() + config = JobConfig(experiment_id="test-exp") job = Job(config) assert job._config == config @@ -166,6 +168,7 @@ async def test_run_full_lifecycle(self): with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): config = JobConfig( + experiment_id="test-exp", job_name="test-job", agents=[AgentConfig(name="t2")], datasets=[RegistryDatasetConfig(registry=RemoteRegistryInfo(), name="tb", version="2.0")], @@ -187,7 +190,9 @@ async def test_run_auto_stop_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job", environment=RockEnvironmentConfig(auto_stop=True)) + config = JobConfig( + job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=True) + ) job = Job(config) await job.run() @@ -197,7 +202,9 @@ async def test_run_does_not_stop_when_disabled(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job", environment=RockEnvironmentConfig(auto_stop=False)) + config = JobConfig( + job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=False) + ) job = Job(config) await job.run() @@ -207,7 +214,7 @@ async def test_submit_starts_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job") + config = JobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() @@ -218,7 +225,7 @@ async def test_wait_returns_result(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job") + config = JobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() result = await job.wait() @@ -233,7 +240,7 @@ def test_oss_vars_from_process_env_are_forwarded(self, monkeypatch): monkeypatch.setenv("OSS_ACCESS_KEY_ID", "test-key") monkeypatch.setenv("HOME", "/root") - job = Job(JobConfig(job_name="test-job")) + job = Job(JobConfig(job_name="test-job", experiment_id="test-exp")) env = job._build_session_env() assert env["OSS_ENDPOINT"] == "https://oss.example.com" @@ -246,6 +253,7 @@ def test_config_env_overrides_process_oss_vars(self, monkeypatch): job = Job( JobConfig( job_name="test-job", + experiment_id="test-exp", environment=RockEnvironmentConfig(env={"OSS_ENDPOINT": "https://oss.from.config.com"}), ) ) @@ -258,7 +266,7 @@ def test_returns_none_when_both_empty(self, monkeypatch): if key.startswith("OSS"): monkeypatch.delenv(key) - job = Job(JobConfig(job_name="test-job")) + job = Job(JobConfig(job_name="test-job", experiment_id="test-exp")) assert job._build_session_env() is None @@ -268,7 +276,7 @@ async def test_cancel_kills_process(self): mock_sandbox.arun = AsyncMock(return_value=MagicMock(output="", exit_code=0)) with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job") + config = JobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() await job.cancel() diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index 8e82f9e019..78506ef605 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -104,6 +104,7 @@ class TestJobConfigToHarborYaml: def test_basic_serialization(self): cfg = JobConfig( job_name="test-job", + experiment_id="test-exp", n_attempts=2, agents=[AgentConfig(name="terminus-2", model_name="hosted_vllm/m")], ) @@ -116,6 +117,7 @@ def test_basic_serialization(self): def test_excludes_rock_fields(self): cfg = JobConfig( + experiment_id="test-exp", environment=RockEnvironmentConfig( setup_commands=["pip install harbor"], file_uploads=[("local.txt", "/sandbox/remote.txt")], @@ -123,7 +125,7 @@ def test_excludes_rock_fields(self): auto_stop=True, image="my-image:latest", memory="32g", - ) + ), ) yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) @@ -142,6 +144,7 @@ def test_excludes_rock_fields(self): def test_excludes_none_values(self): cfg = JobConfig( job_name="test", + experiment_id="test-exp", agents=[AgentConfig(name="t2")], ) yaml_str = cfg.to_harbor_yaml() @@ -151,6 +154,7 @@ def test_excludes_none_values(self): def test_path_fields_serialized_as_strings(self): cfg = JobConfig( + experiment_id="test-exp", jobs_dir=Path("/workspace/jobs"), tasks=[TaskConfig(path="/workspace/tasks/t1")], ) @@ -163,6 +167,7 @@ def test_path_fields_serialized_as_strings(self): def test_harbor_env_fields_serialized(self): cfg = JobConfig( job_name="full-test", + experiment_id="test-exp", n_attempts=3, environment=RockEnvironmentConfig( type="docker", @@ -197,7 +202,7 @@ def test_harbor_env_fields_serialized(self): def test_env_in_harbor_yaml(self): """env is passed to both sandbox session and harbor YAML.""" - cfg = JobConfig(environment=RockEnvironmentConfig(env={"OPENAI_API_KEY": "sk-xxx"})) + cfg = JobConfig(experiment_id="test-exp", environment=RockEnvironmentConfig(env={"OPENAI_API_KEY": "sk-xxx"})) yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) @@ -209,6 +214,7 @@ class TestJobConfigFromYaml: def test_from_yaml_basic(self, tmp_path): yaml_content = """ job_name: loaded-job +experiment_id: test-exp n_attempts: 2 agents: - name: terminus-2 @@ -231,6 +237,7 @@ def test_from_yaml_basic(self, tmp_path): def test_from_yaml_with_environment_block(self, tmp_path): yaml_content = """ job_name: env-job +experiment_id: test-exp environment: image: my-image:latest memory: "32g" @@ -256,6 +263,7 @@ def test_from_yaml_with_environment_block(self, tmp_path): def test_from_yaml_with_local_dataset(self, tmp_path): yaml_content = """ job_name: local-dataset-job +experiment_id: test-exp datasets: - path: /data/tasks task_names: diff --git a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py new file mode 100644 index 0000000000..a8d6ec0899 --- /dev/null +++ b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py @@ -0,0 +1,102 @@ +"""Tests for JobConfig._sync_experiment_id model_validator and namespace consistency.""" + +from unittest.mock import AsyncMock + +import pytest +from pydantic import ValidationError + +from rock.sdk.agent.job import Job +from rock.sdk.agent.models.job.config import JobConfig +from rock.sdk.agent.models.trial.config import RockEnvironmentConfig + + +class TestExperimentIdNotEmpty: + def test_none_experiment_id_raises(self): + """experiment_id=None (default) must raise ValidationError.""" + with pytest.raises(ValidationError, match="experiment_id"): + JobConfig(job_name="test") + + def test_empty_string_experiment_id_raises(self): + """experiment_id='' must raise ValidationError.""" + with pytest.raises(ValidationError, match="experiment_id must not be empty"): + JobConfig(job_name="test", experiment_id="") + + +class TestExperimentIdConsistency: + def test_env_none_syncs_from_jobconfig(self): + """When environment.experiment_id is None, it gets set from JobConfig.""" + cfg = JobConfig(job_name="test", experiment_id="exp-1") + assert cfg.environment.experiment_id == "exp-1" + + def test_env_matches_jobconfig_passes(self): + """When both are set and equal, no error.""" + cfg = JobConfig( + job_name="test", + experiment_id="exp-1", + environment=RockEnvironmentConfig(experiment_id="exp-1"), + ) + assert cfg.experiment_id == "exp-1" + assert cfg.environment.experiment_id == "exp-1" + + def test_env_mismatch_raises(self): + """When environment.experiment_id differs from JobConfig.experiment_id, raise.""" + with pytest.raises(ValidationError, match="experiment_id mismatch"): + JobConfig( + job_name="test", + experiment_id="exp-1", + environment=RockEnvironmentConfig(experiment_id="exp-OTHER"), + ) + + +class TestAutofillNamespaceConsistency: + """Tests for namespace consistency check in Job._autofill_sandbox_info.""" + + async def test_namespace_autofilled_from_sandbox(self): + """When user sets no namespace, sandbox value is used.""" + config = JobConfig(job_name="test", experiment_id="exp-1") + assert config.namespace is None + + job = Job(config) + sandbox = AsyncMock() + sandbox._namespace = "sandbox-ns" + sandbox._experiment_id = "exp-1" + job._sandbox = sandbox + + await job._autofill_sandbox_info() + assert config.namespace == "sandbox-ns" + + async def test_namespace_user_matches_sandbox(self): + """When user namespace matches sandbox, no error.""" + config = JobConfig(job_name="test", experiment_id="exp-1", namespace="same-ns") + job = Job(config) + sandbox = AsyncMock() + sandbox._namespace = "same-ns" + sandbox._experiment_id = "exp-1" + job._sandbox = sandbox + + await job._autofill_sandbox_info() + assert config.namespace == "same-ns" + + async def test_namespace_mismatch_raises(self): + """When user namespace differs from sandbox, raise ValueError.""" + config = JobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") + job = Job(config) + sandbox = AsyncMock() + sandbox._namespace = "different-ns" + sandbox._experiment_id = "exp-1" + job._sandbox = sandbox + + with pytest.raises(ValueError, match="namespace mismatch"): + await job._autofill_sandbox_info() + + async def test_namespace_user_set_sandbox_none(self): + """When user sets namespace but sandbox returns None, keep user value.""" + config = JobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") + job = Job(config) + sandbox = AsyncMock() + sandbox._namespace = None + sandbox._experiment_id = "exp-1" + job._sandbox = sandbox + + await job._autofill_sandbox_info() + assert config.namespace == "user-ns" diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index 6dd15e4290..683c564524 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -200,7 +200,7 @@ def test_with_version(self): class TestJobConfig: def test_defaults(self): - cfg = JobConfig() + cfg = JobConfig(experiment_id="test-exp") assert cfg.n_attempts == 1 assert cfg.timeout_multiplier == 1.0 assert cfg.debug is False @@ -213,7 +213,7 @@ def test_defaults(self): assert cfg.artifacts == [] def test_environment_defaults(self): - cfg = JobConfig() + cfg = JobConfig(experiment_id="test-exp") assert cfg.environment.setup_commands == [] assert cfg.environment.file_uploads == [] assert cfg.environment.env == {} @@ -222,6 +222,7 @@ def test_environment_defaults(self): def test_with_full_config(self): cfg = JobConfig( job_name="test-job", + experiment_id="test-exp", n_attempts=2, agents=[AgentConfig(name="terminus-2", model_name="hosted_vllm/m")], datasets=[RegistryDatasetConfig(registry=RemoteRegistryInfo(), name="terminal-bench", version="2.0")], diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py index adf4395029..1fe838371a 100644 --- a/tests/unit/sdk/agent/test_oss_mirror.py +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -97,9 +97,8 @@ class TestJobConfigNamespaceFields: def test_default_namespace_is_none(self): from rock.sdk.agent.models.job.config import JobConfig - cfg = JobConfig(job_name="test") + cfg = JobConfig(job_name="test", experiment_id="test-exp") assert cfg.namespace is None - assert cfg.experiment_id is None def test_namespace_settable_at_top_level(self): from rock.sdk.agent.models.job.config import JobConfig @@ -118,9 +117,7 @@ class TestToHarborYamlOssMirror: def test_namespace_at_top_level_in_yaml(self): """namespace/experiment_id 序列化为 JobConfig 顶层字段。""" from rock.sdk.agent.models.job.config import JobConfig - from rock.sdk.agent.models.trial.config import OssMirrorConfig - - from rock.sdk.agent.models.trial.config import RockEnvironmentConfig + from rock.sdk.agent.models.trial.config import OssMirrorConfig, RockEnvironmentConfig cfg = JobConfig( job_name="mirror-test", @@ -151,7 +148,7 @@ def test_disabled_oss_mirror_excluded_from_yaml(self): """When oss_mirror is default (disabled), it should not clutter the YAML.""" from rock.sdk.agent.models.job.config import JobConfig - cfg = JobConfig(job_name="no-mirror") + cfg = JobConfig(job_name="no-mirror", experiment_id="test-exp") data = yaml.safe_load(cfg.to_harbor_yaml()) env_data = data.get("environment", {}) @@ -196,6 +193,7 @@ def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): yaml_content = """\ job_name: compat-mirror +experiment_id: test-exp environment: oss_mirror: enabled: true @@ -222,6 +220,7 @@ def test_from_yaml_without_oss_mirror(self, tmp_path): yaml_content = """\ job_name: no-mirror +experiment_id: test-exp agents: - name: test-agent """ @@ -241,7 +240,7 @@ class TestEnableOssMirror: def test_enable_with_all_params(self): from rock.sdk.agent.models.job.config import JobConfig - cfg = JobConfig(job_name="conv-test") + cfg = JobConfig(job_name="conv-test", experiment_id="test-exp") cfg.enable_oss_mirror( oss_bucket="conv-bucket", oss_access_key_id="ak-conv", @@ -255,16 +254,11 @@ def test_enable_with_all_params(self): def test_does_not_touch_namespace_or_experiment_id(self): """enable_oss_mirror 不修改顶层 namespace / experiment_id。""" from rock.sdk.agent.models.job.config import JobConfig - from rock.sdk.sandbox.config import SandboxConfig cfg = JobConfig( job_name="no-touch-test", namespace="preset-ns", experiment_id="preset-exp", - sandbox_config=SandboxConfig( - namespace="sandbox-ns", - experiment_id="sandbox-exp", - ), ) cfg.enable_oss_mirror( oss_bucket="b", diff --git a/tests/unit/sdk/test_arun_nohup.py b/tests/unit/sdk/test_arun_nohup.py index 6c301e4755..096fbfae1c 100644 --- a/tests/unit/sdk/test_arun_nohup.py +++ b/tests/unit/sdk/test_arun_nohup.py @@ -4,8 +4,7 @@ from httpx import ReadTimeout from rock.actions.sandbox.response import Observation -from rock.common.constants import PID_PREFIX -from rock.common.constants import PID_SUFFIX +from rock.common.constants import PID_PREFIX, PID_SUFFIX from rock.sdk.sandbox.client import Sandbox from rock.sdk.sandbox.config import SandboxConfig diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 94cb5eee76..1edf62d1ab 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -8,6 +8,7 @@ async def test_rock_config(): rock_config: RockConfig = RockConfig.from_env() assert rock_config + @pytest.mark.asyncio async def test_runtime_config(): config = { @@ -31,7 +32,7 @@ async def test_runtime_config(): "max_allowed_spec": { "memory": "32g", "cpus": 4, - } + }, } runtime_config = RuntimeConfig(**config_full) diff --git a/tests/unit/utils/test_shell_util.py b/tests/unit/utils/test_shell_util.py index 0ec47e107e..a2eb635c92 100644 --- a/tests/unit/utils/test_shell_util.py +++ b/tests/unit/utils/test_shell_util.py @@ -20,14 +20,10 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): session_name = f"bash-{temp_id}" d = LocalDeployment() await d.start() - await d.runtime.create_session( - SandboxCreateBashSessionRequest(session=session_name) - ) + await d.runtime.create_session(SandboxCreateBashSessionRequest(session=session_name)) cmd = f"/bin/bash -c '{cmd}'" nohup_command = f"nohup {cmd} < /dev/null > {out_file} 2>&1 & echo {PID_PREFIX}$!{PID_SUFFIX};disown" - resp = await d.runtime.run_in_session( - BashAction(command=nohup_command, session=session_name) - ) + resp = await d.runtime.run_in_session(BashAction(command=nohup_command, session=session_name)) logger.info(f"nohup_command response: {resp.output}") pid = extract_nohup_pid(resp.output) start_time = time.perf_counter() @@ -35,9 +31,7 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): while time.perf_counter() < end_time: try: await asyncio.wait_for( - d.runtime.run_in_session( - BashAction(command=f"kill -0 {pid}", session=session_name) - ), + d.runtime.run_in_session(BashAction(command=f"kill -0 {pid}", session=session_name)), timeout=30, ) await asyncio.sleep(1) @@ -48,9 +42,7 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): BashAction(command=f"head -c {response_limited_bytes} {out_file}", session=session_name) ) yield pid, nohup_resp.output - await d.runtime.run_in_session( - BashAction(command=f"rm -rf {out_file}", session=session_name) - ) + await d.runtime.run_in_session(BashAction(command=f"rm -rf {out_file}", session=session_name)) await d.stop() @@ -146,9 +138,7 @@ async def test_continuous_echo_script_command(): echo "b" echo "c" """ - write_script_cmd = ( - f"cat > {script_file} << 'SCRIPT_EOF'\n{script_content}SCRIPT_EOF" - ) + write_script_cmd = f"cat > {script_file} << 'SCRIPT_EOF'\n{script_content}SCRIPT_EOF" async for pid, _ in mock_arun(write_script_cmd): assert pid async for pid, _ in mock_arun(f"chmod +x {script_file}"): From ddfcf4d5d8459fc35d0c855ba3198c2d51bf9b9e Mon Sep 17 00:00:00 2001 From: sanfeng-lhh Date: Wed, 1 Apr 2026 15:29:56 +0800 Subject: [PATCH 004/226] feat: rock agent support verifier native mode (#722) --- .../swe_job_config-verifier.yaml.template | 39 +++++++++++++++++++ examples/harbor/swe_job_config.yaml.template | 3 ++ examples/harbor/tb_job_config.yaml.template | 5 ++- rock/sdk/agent/models/trial/config.py | 3 +- tests/unit/sdk/agent/test_models.py | 20 ++++++++++ 5 files changed, 68 insertions(+), 2 deletions(-) create mode 100644 examples/harbor/swe_job_config-verifier.yaml.template diff --git a/examples/harbor/swe_job_config-verifier.yaml.template b/examples/harbor/swe_job_config-verifier.yaml.template new file mode 100644 index 0000000000..01c270433f --- /dev/null +++ b/examples/harbor/swe_job_config-verifier.yaml.template @@ -0,0 +1,39 @@ +# ── Job Identity ───────────────────────────────────── +experiment_id: "" + +# ── Rock Environment ───────────────────────────────── +environment: + base_url: "" + xrl_authorization: "" + image: "" + cluster: "" + memory: "32g" + cpus: 8 + startup_timeout: 1800 + auto_clear_seconds: 7200 + auto_stop: false + env: + OPENAI_API_KEY: "" + OPENAI_BASE_URL: "" + +# ── Harbor Native ──────────────────────────────────── +agents: + - name: "swe-agent" + model_name: "custom_openai/" + +datasets: + - name: "princeton-nlp/SWE-bench_Verified" + registry: + split: "test" + oss_access_key_id: "" + oss_access_key_secret: "" + oss_bucket: "" + oss_dataset_path: "" + oss_region: "" + oss_endpoint: "" + task_names: + - "astropy__astropy-7606" + +# ── Verifier (optional) ────────────────────────────── +verifier: + mode: "native" diff --git a/examples/harbor/swe_job_config.yaml.template b/examples/harbor/swe_job_config.yaml.template index 5997d6e4a2..bf22b28e43 100644 --- a/examples/harbor/swe_job_config.yaml.template +++ b/examples/harbor/swe_job_config.yaml.template @@ -1,3 +1,6 @@ +# ── Job Identity ───────────────────────────────────── +experiment_id: "" + # ── Rock Environment ───────────────────────────────── environment: base_url: "" diff --git a/examples/harbor/tb_job_config.yaml.template b/examples/harbor/tb_job_config.yaml.template index c78391ae82..918e32c267 100644 --- a/examples/harbor/tb_job_config.yaml.template +++ b/examples/harbor/tb_job_config.yaml.template @@ -1,3 +1,6 @@ +# ── Job Identity ───────────────────────────────────── +experiment_id: "" + # ── Rock Environment ───────────────────────────────── environment: base_url: "" @@ -32,4 +35,4 @@ datasets: oss_bucket: "" oss_dataset_path: "" task_names: - - "crack-7z-hash" + - "crack-7z-hash" \ No newline at end of file diff --git a/rock/sdk/agent/models/trial/config.py b/rock/sdk/agent/models/trial/config.py index 849cf6f47e..4888011a4a 100644 --- a/rock/sdk/agent/models/trial/config.py +++ b/rock/sdk/agent/models/trial/config.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, Field @@ -80,6 +80,7 @@ class VerifierConfig(BaseModel): override_timeout_sec: float | None = None max_timeout_sec: float | None = None disable: bool = False + mode: Literal["harbor", "native"] | None = None class TaskConfig(BaseModel): diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index 683c564524..34ffd42894 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -108,6 +108,26 @@ def test_defaults(self): assert v.override_timeout_sec is None assert v.max_timeout_sec is None assert v.disable is False + assert v.mode is None + + def test_mode_harbor(self): + v = VerifierConfig(mode="harbor") + assert v.mode == "harbor" + + def test_mode_native(self): + v = VerifierConfig(mode="native") + assert v.mode == "native" + + def test_mode_invalid_raises_validation_error(self): + import pytest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + VerifierConfig(mode="invalid") + + def test_mode_none_explicit(self): + v = VerifierConfig(mode=None) + assert v.mode is None class TestTaskConfig: From 86bab48fafb1dd17fb2673174597773836d32e67 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Wed, 1 Apr 2026 17:26:24 +0800 Subject: [PATCH 005/226] kata runtime support dind (#731) --- rock/admin/entrypoints/sandbox_api.py | 20 +++++++- rock/common/constants.py | 1 + rock/deployments/config.py | 6 +++ rock/deployments/docker.py | 63 ++++++++++++++++++++++++++ rock/rocklet/local_files/docker_run.sh | 16 +++++++ rock/utils/providers/nacos_provider.py | 5 ++ 6 files changed, 110 insertions(+), 1 deletion(-) diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index 913c4607b4..277cc39ea0 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -25,7 +25,13 @@ StartHeaders, ) from rock.admin.proto.response import SandboxStartResponse -from rock.common.constants import CPU_PREEMPT_SWITCH, GET_STATUS_SWITCH, KATA_RUNTIME_SWITCH, SUPPORT_KATA_SWITCH +from rock.common.constants import ( + CPU_PREEMPT_SWITCH, + GET_STATUS_SWITCH, + KATA_DIND_DISK_SIZE_KEY, + KATA_RUNTIME_SWITCH, + SUPPORT_KATA_SWITCH, +) from rock.common.exception import handle_exceptions from rock.deployments.config import DockerDeploymentConfig from rock.sandbox.sandbox_manager import SandboxManager @@ -53,6 +59,16 @@ async def _apply_kata_runtime_switch(config: DockerDeploymentConfig) -> None: config.use_kata_runtime = False +async def _apply_kata_disk_size(config: DockerDeploymentConfig) -> None: + """Read kata_dind_disk_size from nacos and override config.kata_disk_size if present.""" + if not config.use_kata_runtime: + return + if sandbox_manager.rock_config.nacos_provider is not None: + disk_size = await sandbox_manager.rock_config.nacos_provider.get_config_value(KATA_DIND_DISK_SIZE_KEY) + if disk_size: + config.kata_disk_size = disk_size + + async def _apply_cpu_preempt_switch(config: DockerDeploymentConfig) -> None: """Check nacos switch and enable CPU preemption on the config if the switch is on. @@ -71,6 +87,7 @@ async def _apply_cpu_preempt_switch(config: DockerDeploymentConfig) -> None: async def start(request: SandboxStartRequest) -> RockResponse[SandboxStartResponse]: config = DockerDeploymentConfig.from_request(request) await _apply_kata_runtime_switch(config) + await _apply_kata_disk_size(config) await _apply_cpu_preempt_switch(config) sandbox_start_response = await sandbox_manager.start(config) return RockResponse(result=sandbox_start_response) @@ -84,6 +101,7 @@ async def start_async( ) -> RockResponse[SandboxStartResponse]: config = DockerDeploymentConfig.from_request(request) await _apply_kata_runtime_switch(config) + await _apply_kata_disk_size(config) await _apply_cpu_preempt_switch(config) sandbox_start_response = await sandbox_manager.start_async( config, diff --git a/rock/common/constants.py b/rock/common/constants.py index 0564d36b9a..d7fa78fc5e 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -4,6 +4,7 @@ KATA_RUNTIME_SWITCH = "use_kata_enabled" SUPPORT_KATA_SWITCH = "support_kata_enabled" CPU_PREEMPT_SWITCH = "cpu_preempt_enabled" +KATA_DIND_DISK_SIZE_KEY = "kata_dind_disk_size" PID_PREFIX = "PIDSTART" PID_SUFFIX = "PIDEND" SCHEDULER_LOG_NAME = "scheduler.log" diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 981868374f..57458b1836 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -105,6 +105,12 @@ class DockerDeploymentConfig(DeploymentConfig): use_kata_runtime: bool = False """Whether to use kata container runtime (io.containerd.kata.v2) instead of --privileged mode.""" + kata_disk_size: str = "50G" + """Size of the sparse disk image for kata DinD. Can be overridden by nacos config 'kata_dind_disk_size'.""" + + kata_disk_base_path: str = "/data/docker-disk" + """Base directory on the host for storing kata disk image files.""" + # TODO: Refine these fields in future versions actor_resource: str | None = None """Resource type for actor allocation (to be refined).""" diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 361b4f9163..67246ddeeb 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -169,6 +169,61 @@ def _build_runtime_args(self) -> list[str]: ] return ["--privileged"] + def _get_kata_disk_image_path(self) -> str: + """Returns the host path for the kata disk image file.""" + return os.path.join(self._config.kata_disk_base_path, f"{self._container_name}.img") + + def _prepare_kata_disk(self) -> None: + """Create and format a sparse disk image for kata DinD on the host. + + Only called when use_kata_runtime is enabled. Creates a sparse file + using truncate (no actual disk space consumed until written) and + formats it as ext4. + """ + if not self._config.use_kata_runtime: + return + + disk_path = self._get_kata_disk_image_path() + os.makedirs(self._config.kata_disk_base_path, exist_ok=True) + logger.info(f"Creating kata disk image: {disk_path} (size={self._config.kata_disk_size})") + + try: + subprocess.check_call( + ["truncate", "-s", self._config.kata_disk_size, disk_path], + timeout=10, + ) + subprocess.check_call( + ["mkfs.ext4", "-F", disk_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=60, + ) + logger.info(f"Kata disk image created and formatted: {disk_path}") + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: + logger.error(f"Failed to prepare kata disk image {disk_path}: {e}", exc_info=True) + if os.path.exists(disk_path): + os.remove(disk_path) + raise + + def _cleanup_kata_disk(self) -> None: + """Remove the kata disk image file from the host. + + Only called when use_kata_runtime is enabled. Silently ignores + missing files to handle cases where preparation failed. + """ + if not self._config or not self._config.use_kata_runtime: + return + if not self._container_name: + return + + disk_path = self._get_kata_disk_image_path() + try: + if os.path.exists(disk_path): + os.remove(disk_path) + logger.info(f"Kata disk image removed: {disk_path}") + except OSError as e: + logger.warning(f"Failed to remove kata disk image {disk_path}: {e}", exc_info=False) + def _get_rocklet_start_cmd(self) -> list[str]: cmd = self._runtime_env.get_rocklet_start_cmd() @@ -340,6 +395,13 @@ async def start(self): env_arg.extend(["-e", f"ROCK_TIME_ZONE={env_vars.ROCK_TIME_ZONE}"]) + # Kata DinD: prepare disk image and add volume mount + env var + if self._config.use_kata_runtime: + self._prepare_kata_disk() + disk_path = self._get_kata_disk_image_path() + volume_args.extend(["-v", f"{disk_path}:/docker-disk.img"]) + env_arg.extend(["-e", "ROCK_KATA_RUNTIME=true"]) + time.sleep(random.randint(0, 5)) runtime_args = self._build_runtime_args() cmds = [ @@ -455,6 +517,7 @@ def _stop(self): logger.warning(f"Failed to kill container {self._container_name} with SIGKILL") self._container_process = None + self._cleanup_kata_disk() self._container_name = None if self._config and self._config.remove_images and DockerUtil.is_image_available(self._config.image): diff --git a/rock/rocklet/local_files/docker_run.sh b/rock/rocklet/local_files/docker_run.sh index 96bae85cf8..1fe9a89d81 100755 --- a/rock/rocklet/local_files/docker_run.sh +++ b/rock/rocklet/local_files/docker_run.sh @@ -28,6 +28,22 @@ is_nix() { fi } +# Kata DinD: set up loop device and mount disk image for Docker storage +setup_kata_dind() { + mkdir -p /var/lib/docker + for i in $(seq 0 7); do + mknod -m 660 /dev/loop$i b 7 $i 2>/dev/null || true + done + mount -o loop /docker-disk.img /var/lib/docker + mount -o remount,rw /sys/fs/cgroup + mount -o remount,rw /proc/sys +} + +if [ "${ROCK_KATA_RUNTIME}" = "true" ]; then + echo "Kata runtime detected, setting up DinD disk..." + setup_kata_dind +fi + # Run rocklet if [ "$(is_nix)" = "true" ]; then # NixOS diff --git a/rock/utils/providers/nacos_provider.py b/rock/utils/providers/nacos_provider.py index c029dab0bd..1b614e805a 100644 --- a/rock/utils/providers/nacos_provider.py +++ b/rock/utils/providers/nacos_provider.py @@ -80,3 +80,8 @@ def add_listener(self): async def get_switch_status(self, switch_name: str, not_found_default: bool = False) -> bool: config = await self.get_config() or {} return bool((config.get("switch") or {}).get(switch_name, not_found_default)) + + async def get_config_value(self, key: str, default: str | None = None) -> str | None: + """Get a string config value from the top-level nacos config dict.""" + config = await self.get_config() or {} + return config.get(key, default) From 87672b99bae86d27f61ee5349a0a3d140ac8bc0c Mon Sep 17 00:00:00 2001 From: guoj14 Date: Thu, 2 Apr 2026 13:12:21 +0800 Subject: [PATCH 006/226] =?UTF-8?q?feat:=20=E6=81=A2=E5=A4=8DCI=E8=AF=B7?= =?UTF-8?q?=E6=B1=82=E8=A7=A6=E5=8F=91=E5=B7=A5=E4=BD=9C=E6=B5=81=E9=85=8D?= =?UTF-8?q?=E7=BD=AE=20(#728)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: 恢复CI请求触发工作流配置 * chore(CI-workflow): 更新CI触发脚本并添加任务开始时间戳 * chore(.github/scripts): 更新CI结果获取脚本中的FC应用地址 * chore(CI-request-trigger.yml): 移除多余的路径引用并修正脚本调用格式 * chore(.github/scripts): 更新CI结果获取脚本中的API地址为HTTP协议 * chore(CI-workflow): 在CI触发脚本前添加时间戳记录 --- .github/scripts/get-CI-result.sh | 4 +- .github/workflows/CI-request-trigger.yml | 51 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/CI-request-trigger.yml diff --git a/.github/scripts/get-CI-result.sh b/.github/scripts/get-CI-result.sh index 19ee8a0c08..6c5eaf301f 100644 --- a/.github/scripts/get-CI-result.sh +++ b/.github/scripts/get-CI-result.sh @@ -9,6 +9,8 @@ fi COMMIT_ID=$1 SECURITY=$2 REPOSITORY=$3 +PIPELINE_ID="42305" +PROJECT_ID="3567319" # 设置最大等待时间 MAX_WAIT_TIME=7200 @@ -19,7 +21,7 @@ while true; do response=$(curl -s -H "Content-Type: application/json" \ -H "Authorization: Basic ${SECURITY}" \ - -d "{\"type\": \"RETRIEVE-TASK-STATUS\", \"repositoryUrl\": \"${REPOSITORY}\", \"commitId\": \"${COMMIT_ID}\"}" "http://get-tasend-back-twkvcdsbpj.cn-hangzhou.fcapp.run") + -d "{\"type\": \"RETRIEVE-TASK-STATUS\", \"aone\": { \"projectId\": \"${PROJECT_ID}\", \"pipelineId\": \"${PIPELINE_ID}\"}, \"repositoryUrl\": \"${REPOSITORY}\",\"commitId\": \"${COMMIT_ID}\"}" "http://get-tasend-back-twkvcdsbpj.cn-hangzhou.fcapp.run") echo "Response: $response" # 检查curl是否成功 diff --git a/.github/workflows/CI-request-trigger.yml b/.github/workflows/CI-request-trigger.yml new file mode 100644 index 0000000000..cf51c3b4fa --- /dev/null +++ b/.github/workflows/CI-request-trigger.yml @@ -0,0 +1,51 @@ +# This is a basic workflow to help you get started with Actions + +name: CI Request Trigger + +# Controls when the workflow will run +on: + pull_request: + branches: [ "master" ] + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +# A workflow run is made up of one or more jobs that can run sequentially or in parallel +jobs: + # This workflow contains a single job called "build" + build: + # The type of runner that the job will run on + runs-on: ubuntu-latest + # work on CI script dir + defaults: + run: + working-directory: .github/scripts + # Steps represent a sequence of tasks that will be executed as part of the job + steps: + # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it + - uses: actions/checkout@v4 + + # Runs trigger CI + - name: Make the script files executable + run: chmod +x trigger-CI.sh get-CI-result.sh + - name: trigger a CI + run: | + echo "=== Task started at: $(date '+%Y-%m-%d %H:%M:%S') ===" + COMMIT_ID=$([ "${{ github.event_name }}" == "pull_request" ] && echo "${{ github.event.pull_request.head.sha }}" || echo "${{ github.sha }}") + echo "Using Commit ID: $COMMIT_ID" + echo "$GITHUB_REF" + PR_ID=$(echo "$GITHUB_REF" | sed 's@refs/pull/\([0-9]\+\)/.*@\1@') + echo "PR ID is $PR_ID" + ./trigger-CI.sh "$COMMIT_ID" "${{ secrets.CI_SECRET }}" "${{ github.event.pull_request.head.repo.clone_url }}" "$PR_ID" + + # Runs get CI result + - name: Get CI result + run: | + echo "=== Task started at: $(date '+%Y-%m-%d %H:%M:%S') ===" + COMMIT_ID=$([ "${{ github.event_name }}" == "pull_request" ] && echo "${{ github.event.pull_request.head.sha }}" || echo "${{ github.sha }}") + echo "Using Commit ID: $COMMIT_ID" + ./get-CI-result.sh "$COMMIT_ID" "${{ secrets.CI_SECRET }}" "${{ github.repository }}" \ No newline at end of file From 695df3030537997553dc33334217bedc7ee400b5 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 2 Apr 2026 13:26:42 +0800 Subject: [PATCH 007/226] feat: add oss_deps field to EnvironmentConfig (#734) --- rock/sdk/agent/models/trial/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/rock/sdk/agent/models/trial/config.py b/rock/sdk/agent/models/trial/config.py index 4888011a4a..acf961d7c7 100644 --- a/rock/sdk/agent/models/trial/config.py +++ b/rock/sdk/agent/models/trial/config.py @@ -47,6 +47,7 @@ class EnvironmentConfig(BaseModel): suppress_override_warnings: bool = False mounts_json: list[dict[str, Any]] | None = None oss_mirror: OssMirrorConfig | None = None + oss_deps: dict[str, str] = Field(default_factory=dict) env: dict[str, str] = Field(default_factory=dict) kwargs: dict[str, Any] = Field(default_factory=dict) From f2a249928dc8be4f4542e9a99440102eb5e4b73e Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 8 Apr 2026 09:53:03 +0800 Subject: [PATCH 008/226] fix: pin langgraph-prebuilt to 1.0.8 to fix CI error (#745) Signed-off-by: Jiachen Zhang --- .../sdk/sandbox/agent/rock_agent/langgraph_config.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/sdk/sandbox/agent/rock_agent/langgraph_config.yaml b/tests/integration/sdk/sandbox/agent/rock_agent/langgraph_config.yaml index 77d95be55c..a7764122ac 100644 --- a/tests/integration/sdk/sandbox/agent/rock_agent/langgraph_config.yaml +++ b/tests/integration/sdk/sandbox/agent/rock_agent/langgraph_config.yaml @@ -8,6 +8,7 @@ runtime_env_config: - langchain==1.2.3 - langchain-openai==1.1.7 - langgraph==1.0.6 + - langgraph-prebuilt==1.0.8 env: OPENAI_API_KEY: xxxxxxx From f5ab08722f11c4f0e6280442b6f90128b62009c7 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 8 Apr 2026 14:20:33 +0800 Subject: [PATCH 009/226] fix: pin docusaurus/theme-common version to fix npm build of docs (#747) Signed-off-by: Jiachen Zhang --- docs/package.json | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/package.json b/docs/package.json index d3fb494dea..32c9628d05 100644 --- a/docs/package.json +++ b/docs/package.json @@ -47,5 +47,8 @@ }, "engines": { "node": ">=20.0" + }, + "overrides": { + "@docusaurus/theme-common": "3.9.2" } } From 312f52d8626ece17b7448758cd9c517c57bc5c43 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 8 Apr 2026 14:59:50 +0800 Subject: [PATCH 010/226] Support persist sandbox metadaba to database (#730) * feat: wire DatabaseConfig into RockConfig YAML loading Add DatabaseConfig dataclass (url field) to rock/config.py and wire it into RockConfig both as a field and in the from_env() YAML parser. * feat: implement SandboxRecord ORM model and async DatabaseProvider - Add Base(DeclarativeBase) as the single SQLAlchemy declarative base - Add SandboxRecord ORM model with all sandbox metadata columns - Add LIST_BY_ALLOWLIST and _NOT_NULL_DEFAULTS class-level constants - Add DatabaseProvider with async engine/session factory - Add DatabaseConfig dataclass to RockConfig - _convert_url handles sqlite://, postgresql://, and postgres:// (Heroku) shorthand; URLs with existing driver specifier pass through unchanged - Default state column value uses string literal "pending" instead of State.PENDING enum instance for explicit column semantics * feat: implement SandboxTable - strict insert, allowlist-based list_by - Add SandboxTable with insert/get/update/delete/list_by/list_by_in - _filter_data strips unknown keys; _NOT_NULL_DEFAULTS fills NOT NULL cols - LIST_BY_ALLOWLIST prevents arbitrary column queries (injection guard) - _record_to_sandbox_info uses lru_cache to avoid repeated get_type_hints calls in bulk list_by scenarios - Add SandboxInfoField generated type and generation script * feat: implement SandboxRepository - Redis hot-path with async DB replica - Redis alive/timeout keys remain the source of truth for live state - DB writes are fire-and-forget via asyncio.create_task + _safe_db_call - batch_get: Redis hits served directly; DB fallback uses a single list_by_in("sandbox_id", miss_ids) query instead of N serial gets, leveraging the primary key index for O(1) lookup per row - iter_alive_sandbox_ids queries DB by state IN (running, pending) instead of Redis scan_iter, enabling indexed filtering * feat: wire SandboxRepository into sandbox lifecycle; rename meta_store to meta_repo - Replace MetaStore with SandboxRepository throughout SandboxManager, GemManager, BaseManager, and SandboxProxyService - Wire SandboxRepository (Redis + SandboxTable) in admin/main.py startup - stop(): add early return after archive() in the ValueError except branch to prevent double archive when the Ray actor is already gone Made-with: Cursor * test: unit tests for SandboxTable and SandboxRepository - Add TestSandboxTableWithSQLite: full CRUD coverage using SQLite in-memory database (no external dependencies, runs in fast CI) including list_by_in, NOT NULL defaults, and noop-on-missing-id cases - Add TestSandboxTableWithPostgres: PostgreSQL-specific tests (JSONB, real container) marked need_docker + need_database - Add comprehensive SandboxRepository tests: create/update/delete/archive/ get/exists/batch_get/list_by/refresh_timeout/is_expired - Consistent lowercase "stopped" state string throughout test data, matching the State enum value convention (running/pending) * feat: add indexes to SandboxRecord and add DDL generation script - Add single-column indexes on all commonly queried fields (user_id, state, namespace, experiment_id, cluster_name, image, host_ip, host_name, create_user_gray_flag) - Add scripts/gen_ddl.py to emit CREATE TABLE / CREATE INDEX DDL - Add *.db and ddl/ to .gitignore (generated artifacts) * fix: inject redis_provider into RayOperator via OperatorContext OperatorContext was missing redis_provider, leaving RayOperator._redis_provider as None. This caused the use_rocklet get_status path to crash with 'NoneType object has no attribute get' because build_sandbox_from_redis skips the lookup entirely when redis_provider is None. * refactor: rename SandboxRepository to SandboxMetaStore - Rename class SandboxRepository to SandboxMetaStore to better reflect its role as a coordinator for Redis (hot path) + DB (query path) dual-write - Rename _meta_repo to _meta_store across all files - Rename sandbox_repository.py to sandbox_meta_store.py - Update all imports and references - Use legacy states (_TERMINAL_STATES, _LIST_BY_BLACKLIST) from SandboxMetaStore as the authoritative source; removed duplicate definitions elsewhere * refactor: extract sandbox timeout logic into SandboxTimeoutHelper - Add rock/sandbox/utils/timeout.py with SandboxTimeoutHelper: pure calculation helpers (make_timeout_info, refresh_timeout, is_expired) with no I/O dependency - Add SandboxMetaStore.update_timeout() for raw Redis set of timeout key; remove refresh_timeout() and is_expired() from MetaStore (not its responsibility) - SandboxManager._refresh_timeout / _is_expired: own the I/O (get_timeout + update_timeout) and delegate calculation to SandboxTimeoutHelper - SandboxProxyService._update_expire_time: same pattern - Replace inline auto_clear_time_dict construction in start_async with SandboxTimeoutHelper.make_timeout_info() - Update tests: replace TestRefreshTimeout/TestIsExpired in test_sandbox_meta_store with TestUpdateTimeout; add test_sandbox_timeout.py for pure unit tests * feat: add spec/status columns to SandboxRecord; remove SandboxInfoField Schema - Add spec (JSONB): DockerDeploymentConfig.model_dump() snapshot, written once at creation, never updated - Add status (JSONB): full SandboxInfo snapshot, overwritten on every update SandboxRecordData - New TypedDict in schema.py extending SandboxInfo with spec/status fields - Used as the unified I/O type for SandboxTable (replaces plain SandboxInfo) SandboxTable - create(): writes spec from caller; auto-populates status from data - update(): always overwrites status with latest SandboxInfo snapshot - list_by / list_by_in / get return SandboxRecordData SandboxMetaStore - create() gains spec: dict | None parameter; constructs SandboxRecordData before passing to SandboxTable; Redis path unchanged SandboxManager - start_async passes spec=docker_deployment_config.model_dump() to meta_store Cleanup - Remove SandboxInfoField generated Literal type and its generation script - Replace SandboxInfoField with plain str in SandboxTable / SandboxMetaStore (LIST_BY_ALLOWLIST already enforces valid column names at runtime) * refactor: replace async_sessionmaker with direct AsyncSession(engine) in SandboxTable - Remove async_sessionmaker, _session_factory, and session() factory from DatabaseProvider - Add engine property that raises RuntimeError if not initialised - Update all SandboxTable methods to use AsyncSession(self._db.engine) directly - Simpler, more explicit session lifecycle with no factory indirection * refactor: always-on DB+Redis providers; simplify SandboxMetaStore - Fallback to sqlite-memory when database.url is not configured - Fallback to FakeRedis when redis.host is not configured - SandboxMetaStore now requires both providers (no more None checks) - batch_get() returns only found sandboxes (no positional None slots) - list_by() raises ValueError for non-allowlisted fields (no Redis fallback) - list_sandboxes and batch_get_status expose use_legacy_states param - Update tests to reflect new behaviour * refactor: rename DatabaseProvider init_pool/close_pool to init/close; remove auto-create_all and verbose logging Co-Authored-By: Claude Opus 4.6 * refactor: remove use_legacy_states parameter; always filter non-RUNNING/PENDING sandboxes Co-Authored-By: Claude Opus 4.6 * feat: raise error in get_status when sandbox is already stopped Co-Authored-By: Claude Opus 4.6 * chore: add sandbox_record DDL as sql/sandbox_record.sql Co-Authored-By: Claude Opus 4.6 * refactor: remove all _meta_store None checks; meta_store is always provided Co-Authored-By: Claude Opus 4.6 * refactor: add DatabaseProvider.create_tables() for explicit DDL creation in tests Co-Authored-By: Claude Opus 4.6 * fix: auto-create tables for SQLite fallback in admin lifespan When database.url is not configured, admin falls back to SQLite in-memory which has no persistent schema. Call create_tables() only in this case so integration tests and local dev work out of the box. Production PostgreSQL relies on external DDL (sql/sandbox_record.sql). Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .gitignore | 3 +- pyproject.toml | 4 +- requirements_admin.txt | 3 + rock/actions/sandbox/response.py | 1 + rock/admin/core/db_provider.py | 51 +- rock/admin/core/sandbox_table.py | 135 ++++- rock/admin/core/schema.py | 117 ++++- rock/admin/entrypoints/sandbox_proxy_api.py | 6 +- rock/admin/main.py | 34 +- rock/config.py | 6 + rock/sandbox/base_manager.py | 27 +- rock/sandbox/gem_manager.py | 14 +- rock/sandbox/operator/factory.py | 4 + rock/sandbox/operator/ray.py | 3 +- rock/sandbox/sandbox_manager.py | 119 ++--- rock/sandbox/sandbox_meta_store.py | 143 +++++ rock/sandbox/service/sandbox_proxy_service.py | 90 ++-- rock/sandbox/utils/__init__.py | 0 rock/sandbox/utils/timeout.py | 51 ++ scripts/gen_ddl.py | 57 ++ sql/sandbox_record.sql | 42 ++ tests/integration/conftest.py | 15 + tests/unit/admin/core/test_sandbox_table.py | 231 +++++++-- tests/unit/conftest.py | 165 +++++- tests/unit/sandbox/test_proxy_enhancements.py | 105 ++++ tests/unit/sandbox/test_sandbox_meta_store.py | 490 ++++++++++++++++++ tests/unit/sandbox/test_sandbox_proxy.py | 4 +- tests/unit/sandbox/test_sandbox_timeout.py | 55 ++ .../unit/utils/test_redis_provider_docker.py | 90 ++++ uv.lock | 62 +++ 30 files changed, 1877 insertions(+), 250 deletions(-) create mode 100644 rock/sandbox/sandbox_meta_store.py create mode 100644 rock/sandbox/utils/__init__.py create mode 100644 rock/sandbox/utils/timeout.py create mode 100644 scripts/gen_ddl.py create mode 100644 sql/sandbox_record.sql create mode 100644 tests/unit/sandbox/test_sandbox_meta_store.py create mode 100644 tests/unit/sandbox/test_sandbox_timeout.py create mode 100644 tests/unit/utils/test_redis_provider_docker.py diff --git a/.gitignore b/.gitignore index 144b8b5ea2..b370b13a13 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,5 @@ node_modules logs docs/superpowers/ -.env \ No newline at end of file +.env +*.db diff --git a/pyproject.toml b/pyproject.toml index 5d67c29697..2ae7df85d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ admin = [ "alibabacloud_cr20181201==2.0.5", "sqlmodel", "aiosqlite", + "asyncpg", "boto3", "ray[default]==2.43.0", "pip", @@ -171,5 +172,6 @@ markers = [ "need_ray: need ray start", "need_docker: need docker daemon running", "need_admin: need admin start", - "need_admin_and_network: need install from network" + "need_admin_and_network: need install from network", + "need_database: need database Docker containers (PostgreSQL, Redis)" ] diff --git a/requirements_admin.txt b/requirements_admin.txt index 379fcc1fb6..332a15229d 100644 --- a/requirements_admin.txt +++ b/requirements_admin.txt @@ -108,7 +108,10 @@ arckit==0.1.0 async-timeout==5.0.1 ; python_full_version < '3.11.3' # via # aiohttp + # asyncpg # redis +asyncpg==0.31.0 + # via rl-rock attrs==25.4.0 # via # aiohttp diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index 3fda9ee2d3..3ec3247193 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -15,6 +15,7 @@ class SandboxResponse(BaseModel): class State(str, Enum): PENDING = "pending" RUNNING = "running" + STOPPED = "stopped" class IsAliveResponse(BaseModel): diff --git a/rock/admin/core/db_provider.py b/rock/admin/core/db_provider.py index 4c6d5dcba1..a20277dd00 100644 --- a/rock/admin/core/db_provider.py +++ b/rock/admin/core/db_provider.py @@ -1,21 +1,56 @@ +"""Generic async SQLAlchemy engine provider.""" + +from __future__ import annotations + from typing import TYPE_CHECKING from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine -from rock.admin.core.schema import DBModelBase +from rock.admin.core.schema import Base +from rock.logger import init_logger if TYPE_CHECKING: from rock.config import DatabaseConfig +logger = init_logger(__name__) + class DatabaseProvider: - def __init__(self, db_config: "DatabaseConfig"): - self.db_config = db_config - self.engine: AsyncEngine + """Async SQLAlchemy engine provider. + + Supports SQLite (via ``aiosqlite``) and PostgreSQL (via ``asyncpg``). + """ - async def init(self): - self.engine = create_async_engine(self.db_config.url, echo=True) + def __init__(self, db_config: DatabaseConfig) -> None: + self._url = self._convert_url(db_config.url) + self._engine: AsyncEngine | None = None - async def create_tables(self): + @property + def engine(self) -> AsyncEngine: + if self._engine is None: + raise RuntimeError("DatabaseProvider not initialised. Call init() first.") + return self._engine + + async def init(self) -> None: + """Create the async engine.""" + self._engine = create_async_engine(self._url, echo=False) + + async def create_tables(self) -> None: + """Create all tables defined in Base.metadata (idempotent).""" async with self.engine.begin() as conn: - await conn.run_sync(DBModelBase.metadata.create_all) + await conn.run_sync(Base.metadata.create_all) + + async def close(self) -> None: + """Dispose of the engine and release all connections.""" + if self._engine is not None: + await self._engine.dispose() + + @staticmethod + def _convert_url(url: str) -> str: + """Convert synchronous database URLs to their async equivalents.""" + if url.startswith("sqlite:///"): + return url.replace("sqlite:///", "sqlite+aiosqlite:///", 1) + if url.startswith("postgresql://") or url.startswith("postgres://"): + prefix = "postgresql://" if url.startswith("postgresql://") else "postgres://" + return "postgresql+asyncpg://" + url[len(prefix):] + return url diff --git a/rock/admin/core/sandbox_table.py b/rock/admin/core/sandbox_table.py index 299f7af479..9141b07767 100644 --- a/rock/admin/core/sandbox_table.py +++ b/rock/admin/core/sandbox_table.py @@ -1,34 +1,125 @@ -from collections.abc import Sequence +"""SandboxTable: sandbox-specific CRUD and query operations over DatabaseProvider.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession +from sqlalchemy.ext.asyncio import AsyncSession +from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.schema import SandboxRecord +from rock.logger import init_logger + +if TYPE_CHECKING: + from rock.actions.sandbox.sandbox_info import SandboxInfo + from rock.deployments.config import DockerDeploymentConfig + +logger = init_logger(__name__) class SandboxTable: - def __init__(self, engine: AsyncEngine): - self._engine = engine + """Sandbox-specific database access layer backed by DatabaseProvider. + + All methods use plain ``dict`` for both input and output. + + Write path (create / update): + - Fields from ``SandboxInfo`` / ``DockerDeploymentConfig`` that match a + ``SandboxRecord`` column are written to the corresponding scalar column. + - ``status`` column stores the full ``SandboxInfo`` dict. + - ``spec`` column stores the full ``DockerDeploymentConfig.model_dump()`` dict. + + Read path (get / list_by / list_by_in): + - Returns ``record.to_dict()`` — a plain dict with all non-None column values, + including ``spec`` and ``status``. + """ + + def __init__(self, db_provider: DatabaseProvider) -> None: + self._db = db_provider + + async def create( + self, + sandbox_id: str, + info: SandboxInfo, + config: DockerDeploymentConfig | None = None, + ) -> None: + """Insert a new sandbox record. + + Scalar columns are populated from the union of *config* and *info* + (``info`` takes priority on conflicts). + Raises ``IntegrityError`` if ``sandbox_id`` already exists. + """ + config_dict = config.model_dump() if config is not None else {} + merged = {**config_dict, **info} + filtered = _pick_columns(merged) - async def create(self, sandbox_record: SandboxRecord): - async with AsyncSession(self._engine) as session: - session.add(sandbox_record) + for col, default in SandboxRecord._NOT_NULL_DEFAULTS.items(): + if col not in filtered: + filtered[col] = default + + filtered["status"] = dict(info) + if config_dict: + filtered["spec"] = config_dict + + record = SandboxRecord(sandbox_id=sandbox_id, **filtered) + async with AsyncSession(self._db.engine) as session: + session.add(record) + await session.commit() + + async def get(self, sandbox_id: str) -> dict | None: + """Return a sandbox row as a plain dict, or ``None`` if not found.""" + async with AsyncSession(self._db.engine) as session: + record = await session.get(SandboxRecord, sandbox_id) + if record is None: + return None + return record.to_dict() + + async def update(self, sandbox_id: str, info: SandboxInfo) -> None: + """Partial update of scalar columns; always overwrites ``status`` with *info*.""" + filtered = _pick_columns(info) + filtered["status"] = dict(info) + + async with AsyncSession(self._db.engine) as session: + record = await session.get(SandboxRecord, sandbox_id) + if record is None: + logger.warning("update: sandbox_id=%s not found", sandbox_id) + return + for key, value in filtered.items(): + setattr(record, key, value) await session.commit() - async def list( - self, namespace: str | None = None, user: str | None = None, experiment_id: str | None = None - ) -> Sequence[SandboxRecord]: - async with AsyncSession(self._engine) as session: - stmt = select(SandboxRecord) - if None is not namespace: - stmt = stmt.where(SandboxRecord.namespace == namespace) - if None is not user: - stmt = stmt.where(SandboxRecord.user == user) - if None is not experiment_id: - stmt = stmt.where(SandboxRecord.experiment_id == experiment_id) + async def delete(self, sandbox_id: str) -> None: + """Hard-delete a sandbox record.""" + async with AsyncSession(self._db.engine) as session: + record = await session.get(SandboxRecord, sandbox_id) + if record is not None: + await session.delete(record) + await session.commit() + + async def list_by(self, column: str, value: str | int | float | bool) -> list[dict]: + """Equality query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" + if column not in SandboxRecord.LIST_BY_ALLOWLIST: + raise ValueError(f"Querying by column '{column}' is not allowed") + col_attr = getattr(SandboxRecord, column) + stmt = select(SandboxRecord).where(col_attr == value) + async with AsyncSession(self._db.engine) as session: result = await session.execute(stmt) - return result.scalars().all() + return [r.to_dict() for r in result.scalars().all()] + + async def list_by_in(self, column: str, values: list[str | int | float | bool]) -> list[dict]: + """IN query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" + if column not in SandboxRecord.LIST_BY_ALLOWLIST: + raise ValueError(f"Querying by column '{column}' is not allowed") + if not values: + return [] + col_attr = getattr(SandboxRecord, column) + stmt = select(SandboxRecord).where(col_attr.in_(values)) + async with AsyncSession(self._db.engine) as session: + result = await session.execute(stmt) + return [r.to_dict() for r in result.scalars().all()] + - async def get(self, id: str) -> SandboxRecord: - async with AsyncSession(self._engine) as session: - return await session.get(SandboxRecord, id) +def _pick_columns(data: dict[str, Any]) -> dict[str, Any]: + """Return only keys matching a scalar SandboxRecord column, excluding sandbox_id/spec/status.""" + columns = SandboxRecord.column_names() - {"sandbox_id", "spec", "status"} + return {k: v for k, v in data.items() if k in columns} diff --git a/rock/admin/core/schema.py b/rock/admin/core/schema.py index 21904ccaed..cfed2f3dc4 100644 --- a/rock/admin/core/schema.py +++ b/rock/admin/core/schema.py @@ -1,26 +1,105 @@ -from sqlalchemy import Column, DateTime, String -from sqlalchemy.ext.declarative import declarative_base +"""ORM base and models for the ROCK admin database. -DBModelBase = declarative_base() +``Base`` is the single SQLAlchemy ``DeclarativeBase`` for all ROCK tables. +``SandboxRecord`` is the canonical persistence model for sandbox metadata. +""" +from __future__ import annotations -class SandboxRecord(DBModelBase): - __tablename__ = "sandboxes" +from typing import Any, ClassVar - id = Column(String, primary_key=True) +from sqlalchemy import Boolean, Column, Float, Index, String +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.types import JSON - # Some grouping fields for filtering - namespace = Column(String) - user = Column(String) - experiment_id = Column(String) - created_at = Column(DateTime) - # Last call time - last_called_at = Column(DateTime) - # Stop time/auto cleanup time - stopped_at = Column(DateTime) +_JSONB_VARIANT = JSON().with_variant(JSONB(), "postgresql") - # Some metadata information - image = Column(String) - spec_meta = Column(String) - status_meta = Column(String) + +class Base(DeclarativeBase): + pass + + +# All ORM models that inherit from Base must be imported (or defined) in this file +# so that Base.metadata is fully populated before DatabaseProvider.init() calls +# create_all. When adding a new table in a separate module, add its import here: +# from rock.admin.core. import # noqa: F401 +class SandboxRecord(Base): + """ORM model for the ``sandbox_record`` table.""" + + __tablename__ = "sandbox_record" + + sandbox_id = Column(String(128), primary_key=True) + user_id = Column(String(128), nullable=False, default="default") + image = Column(String(128), nullable=False, default="default") + experiment_id = Column(String(128), nullable=False, default="default") + namespace = Column(String(128), nullable=False, default="default") + cluster_name = Column(String(128), nullable=False, default="default") + state = Column(String(32), nullable=False, default="pending") + host_ip = Column(String(128), nullable=False, default="default") + create_time = Column(String(64), nullable=False, default="") + start_time = Column(String(64), nullable=True) + stop_time = Column(String(64), nullable=True) + host_name = Column(String(255), nullable=True) + auth_token = Column(String(512), nullable=True) + rock_authorization_encrypted = Column(String(1024), nullable=True) + cpus = Column(Float, nullable=True) + memory = Column(String(64), nullable=True) + create_user_gray_flag = Column(Boolean, nullable=True) + phases = Column(_JSONB_VARIANT, nullable=True) + port_mapping = Column(_JSONB_VARIANT, nullable=True) + spec = Column(_JSONB_VARIANT, nullable=True) + status = Column(_JSONB_VARIANT, nullable=True) + + __table_args__ = ( + Index("ix_sandbox_record_user_id", "user_id"), + Index("ix_sandbox_record_state", "state"), + Index("ix_sandbox_record_namespace", "namespace"), + Index("ix_sandbox_record_experiment_id", "experiment_id"), + Index("ix_sandbox_record_cluster_name", "cluster_name"), + Index("ix_sandbox_record_image", "image"), + Index("ix_sandbox_record_host_ip", "host_ip"), + Index("ix_sandbox_record_host_name", "host_name"), + Index("ix_sandbox_record_create_user_gray_flag", "create_user_gray_flag"), + ) + + # Columns allowed as the filter key in list_by(). + # Only include columns with an index (or PK); exclude JSONB, sensitive, and internal columns. + LIST_BY_ALLOWLIST: ClassVar[frozenset[str]] = frozenset( + { + "sandbox_id", + "user_id", + "image", + "experiment_id", + "namespace", + "cluster_name", + "state", + "host_ip", + "host_name", + "create_user_gray_flag", + } + ) + + _column_names: ClassVar[set[str] | None] = None + + @classmethod + def column_names(cls) -> set[str]: + if cls._column_names is None: + cls._column_names = {c.key for c in cls.__table__.columns} + return cls._column_names + + _NOT_NULL_DEFAULTS: ClassVar[dict[str, Any]] = { + "user_id": "default", + "image": "default", + "experiment_id": "default", + "namespace": "default", + "cluster_name": "default", + "state": "pending", + "host_ip": "default", + "create_time": "", + } + + def to_dict(self) -> dict[str, Any]: + """Return all non-``None`` column values as a plain dict.""" + return {c.key: getattr(self, c.key) for c in self.__table__.columns if getattr(self, c.key) is not None} diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 67d8f87605..27e1f9632f 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -131,8 +131,10 @@ async def run(action: SandboxBashAction) -> RockResponse[BashObservation]: @sandbox_proxy_router.post("/sandboxes/batch") @handle_exceptions(error_message="batch get sandbox status failed") -async def batch_get_status(request: BatchSandboxStatusRequest) -> RockResponse[BatchSandboxStatusResponse]: - statuses_list = await sandbox_proxy_service.batch_get_sandbox_status_from_redis(request.sandbox_ids) +async def batch_get_status( + request: BatchSandboxStatusRequest, +) -> RockResponse[BatchSandboxStatusResponse]: + statuses_list = await sandbox_proxy_service.batch_get_sandbox_status(request.sandbox_ids) response = BatchSandboxStatusResponse(statuses=statuses_list) return RockResponse(result=response) diff --git a/rock/admin/main.py b/rock/admin/main.py index 2cb242833c..3e3441bab2 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -13,13 +13,15 @@ from starlette.responses import JSONResponse from rock import env_vars +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.sandbox_table import SandboxTable from rock.admin.core.ray_service import RayService from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service from rock.admin.scheduler.scheduler import SchedulerThread -from rock.config import RockConfig +from rock.config import RockConfig, DatabaseConfig from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager from rock.sandbox.operator.factory import OperatorContext, OperatorFactory @@ -51,10 +53,12 @@ async def lifespan(app: FastAPI): env_vars.ROCK_ADMIN_ENV = args.env env_vars.ROCK_ADMIN_ROLE = args.role - # init redis provider - if args.env in ["local", "test", "dev"]: + # init redis provider (fallback to fakeredis if no host configured) + if args.env in ["local", "test", "dev"] or not rock_config.redis.host: from fakeredis import aioredis + if not rock_config.redis.host: + logger.info("redis.host is not configured, falling back to FakeRedis") redis_provider = RedisProvider(host=None, port=None, password="") redis_provider.client = aioredis.FakeRedis(decode_responses=True) else: @@ -65,6 +69,20 @@ async def lifespan(app: FastAPI): ) await redis_provider.init_pool() + # init database provider (fallback to sqlite in-memory if no url configured) + db_url = rock_config.database.url or "sqlite+aiosqlite:///:memory:" + if not rock_config.database.url: + logger.info("database.url is not configured, falling back to SQLite in-memory") + db_provider = DatabaseProvider(db_config=DatabaseConfig(url=db_url)) + await db_provider.init() + if not rock_config.database.url: + await db_provider.create_tables() + sandbox_table = SandboxTable(db_provider) + + from rock.sandbox.sandbox_meta_store import SandboxMetaStore + + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=sandbox_table) + # init scheduler thread scheduler_thread = None @@ -78,6 +96,7 @@ async def lifespan(app: FastAPI): operator_context = OperatorContext( runtime_config=rock_config.runtime, ray_service=ray_service, + redis_provider=redis_provider, nacos_provider=rock_config.nacos_provider, k8s_config=rock_config.k8s, ) @@ -87,20 +106,20 @@ async def lifespan(app: FastAPI): if rock_config.runtime.enable_auto_clear: sandbox_manager = GemManager( rock_config, - redis_provider=redis_provider, ray_namespace=rock_config.ray.namespace, ray_service=ray_service, enable_runtime_auto_clear=True, operator=operator, + meta_store=meta_store, ) else: sandbox_manager = GemManager( rock_config, - redis_provider=redis_provider, ray_namespace=rock_config.ray.namespace, ray_service=ray_service, enable_runtime_auto_clear=False, operator=operator, + meta_store=meta_store, ) set_sandbox_manager(sandbox_manager) warmup_service = WarmupService(rock_config.warmup) @@ -118,7 +137,7 @@ async def lifespan(app: FastAPI): logger.info("Scheduler thread skipped on non-primary pod") else: - sandbox_manager = SandboxProxyService(rock_config=rock_config, redis_provider=redis_provider) + sandbox_manager = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) set_sandbox_proxy_service(sandbox_manager) logger.info("rock-admin start") @@ -130,6 +149,9 @@ async def lifespan(app: FastAPI): scheduler_thread.stop() logger.info("Scheduler thread stopped") + if db_provider: + await db_provider.close() + if redis_provider: await redis_provider.close_pool() diff --git a/rock/config.py b/rock/config.py index f43952732e..b9a7a6efd8 100644 --- a/rock/config.py +++ b/rock/config.py @@ -78,6 +78,9 @@ class ProxyServiceConfig: @dataclass class DatabaseConfig: + # Supported URL formats: + # SQLite: sqlite:///relative/path.db or sqlite:////absolute/path.db + # PostgreSQL: postgresql://user:password@host:port/dbname url: str = "" @@ -194,6 +197,7 @@ class RockConfig: runtime: RuntimeConfig = field(default_factory=RuntimeConfig) proxy_service: ProxyServiceConfig = field(default_factory=ProxyServiceConfig) scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) + database: DatabaseConfig = field(default_factory=DatabaseConfig) nacos_provider: NacosConfigProvider | None = None @classmethod @@ -235,6 +239,8 @@ def from_env(cls, config_path: str | None = None): kwargs["proxy_service"] = ProxyServiceConfig(**config["proxy_service"]) if "scheduler" in config: kwargs["scheduler"] = SchedulerConfig(**config["scheduler"]) + if "database" in config: + kwargs["database"] = DatabaseConfig(**config["database"]) return cls(**kwargs) diff --git a/rock/sandbox/base_manager.py b/rock/sandbox/base_manager.py index d084b2cfc1..cadee819b5 100644 --- a/rock/sandbox/base_manager.py +++ b/rock/sandbox/base_manager.py @@ -5,32 +5,30 @@ from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.interval import IntervalTrigger -from rock.admin.core.redis_key import ALIVE_PREFIX from rock.admin.metrics.constants import MetricsConstants from rock.admin.metrics.monitor import MetricsMonitor, aggregate_metrics from rock.config import RockConfig from rock.deployments.manager import DeploymentManager from rock.logger import init_logger from rock.utils import get_executor -from rock.utils.providers.redis_provider import RedisProvider +from rock.sandbox.sandbox_meta_store import SandboxMetaStore logger = init_logger(__name__) class BaseManager: _check_job_bg_task: object = None - _redis_provider: RedisProvider = None rock_config: RockConfig = None def __init__( self, rock_config: RockConfig, - redis_provider: RedisProvider | None = None, + meta_store: SandboxMetaStore, enable_runtime_auto_clear: bool = False, ): self.rock_config = rock_config self._executor = get_executor() - self._redis_provider = redis_provider + self._meta_store = meta_store self.metrics_monitor = MetricsMonitor.create( export_interval_millis=20_000, metrics_endpoint=rock_config.runtime.metrics_endpoint, @@ -91,16 +89,13 @@ async def _collect_and_report_metrics(self): async def _collect_and_report_metrics_internal(self): """Collect and report metrics for all sandboxes""" overall_start = time.perf_counter() - if not self._redis_provider: - return await self._report_system_resource_metrics() - if not await self._redis_provider.pattern_exists(f"{ALIVE_PREFIX}*"): + sandbox_cnt, sandbox_meta = await self._collect_sandbox_meta() + if sandbox_cnt == 0: logger.debug("No sandboxes to monitor") self.metrics_monitor.record_gauge_by_name(MetricsConstants.SANDBOX_TOTAL_COUNT, 0) return - - sandbox_cnt, sandbox_meta = await self._collect_sandbox_meta() aggregated_metrics = aggregate_metrics(sandbox_meta, "image") for image, count in aggregated_metrics.items(): self.metrics_monitor.record_gauge_by_name(MetricsConstants.SANDBOX_COUNT_IMAGE, count, {"image": image}) @@ -133,16 +128,10 @@ async def _collect_system_resource_metrics(self): async def _collect_sandbox_meta(self) -> tuple[int, dict[str, dict[str, str]]]: meta: dict = {} cnt = 0 - # type: ignore - async for key in self._redis_provider.client.scan_iter(match=f"{ALIVE_PREFIX}*", count=100): - sandbox_id = key.removeprefix(ALIVE_PREFIX) + async for sandbox_id in self._meta_store.iter_alive_sandbox_ids(): cnt += 1 - if self._sandbox_meta.get(sandbox_id) is not None: - try: - image = self._sandbox_meta[sandbox_id]["image"] - except Exception: - image = "default" - meta[sandbox_id] = {"image": image} + image = self._sandbox_meta.get(sandbox_id, {}).get("image", "default") + meta[sandbox_id] = {"image": image} return cnt, meta def stop_monitoring(self): diff --git a/rock/sandbox/gem_manager.py b/rock/sandbox/gem_manager.py index 2cb2aab1c1..d0ec3d4708 100644 --- a/rock/sandbox/gem_manager.py +++ b/rock/sandbox/gem_manager.py @@ -18,20 +18,26 @@ from rock.deployments.config import DockerDeploymentConfig from rock.sandbox.sandbox_actor import SandboxActor from rock.sandbox.sandbox_manager import SandboxManager -from rock.utils.providers import RedisProvider - +from rock.sandbox.sandbox_meta_store import SandboxMetaStore class GemManager(SandboxManager): def __init__( self, rock_config: RockConfig, - redis_provider: RedisProvider | None = None, + meta_store: SandboxMetaStore | None = None, ray_namespace: str = env_vars.ROCK_RAY_NAMESPACE, ray_service: RayService | None = None, enable_runtime_auto_clear: bool = False, operator=None, ): - super().__init__(rock_config, redis_provider, ray_namespace, ray_service, enable_runtime_auto_clear, operator) + super().__init__( + rock_config, + meta_store=meta_store, + ray_namespace=ray_namespace, + ray_service=ray_service, + enable_runtime_auto_clear=enable_runtime_auto_clear, + operator=operator, + ) async def env_make(self, env_id: str) -> EnvMakeResponse: config = DockerDeploymentConfig(image=env_vars.ROCK_ENVHUB_DEFAULT_DOCKER_IMAGE) diff --git a/rock/sandbox/operator/factory.py b/rock/sandbox/operator/factory.py index 3b0fc0d720..d7289c6ee8 100644 --- a/rock/sandbox/operator/factory.py +++ b/rock/sandbox/operator/factory.py @@ -10,6 +10,7 @@ from rock.sandbox.operator.k8s.operator import K8sOperator from rock.sandbox.operator.ray import RayOperator from rock.utils.providers.nacos_provider import NacosConfigProvider +from rock.utils.providers.redis_provider import RedisProvider logger = init_logger(__name__) @@ -25,6 +26,7 @@ class OperatorContext: runtime_config: RuntimeConfig ray_service: RayService | None = None + redis_provider: RedisProvider | None = None # K8s operator dependencies k8s_config: K8sConfig | None = None nacos_provider: NacosConfigProvider | None = None @@ -59,6 +61,8 @@ def create_operator(context: OperatorContext) -> AbstractOperator: raise ValueError("RayService is required for RayOperator") logger.info("Creating RayOperator") ray_operator = RayOperator(ray_service=context.ray_service, runtime_config=context.runtime_config) + if context.redis_provider is not None: + ray_operator.set_redis_provider(context.redis_provider) if context.nacos_provider is not None: ray_operator.set_nacos_provider(context.nacos_provider) return ray_operator diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index b23b1fd822..6412d1269a 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -90,7 +90,7 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: 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["phases"] = {name: phase.to_dict() for name, phase in remote_status.phases.items()} sandbox_info["port_mapping"] = remote_status.get_port_mapping() alive = await self._ray_service.async_ray_get(actor.is_alive.remote()) # TODO: sink update state according to is_alive logic into SandboxInfo @@ -101,7 +101,6 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: redis_info = await self.get_sandbox_info_from_redis(sandbox_id) if redis_info: redis_info.update(sandbox_info) - redis_info["phases"] = {name: phase.to_dict() for name, phase in remote_status.phases.items()} return redis_info else: return sandbox_info diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 3dbe2652db..c22a0be788 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -1,5 +1,4 @@ import asyncio -import time from fastapi import UploadFile @@ -16,7 +15,6 @@ from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.ray_service import RayService -from rock.admin.core.redis_key import ALIVE_PREFIX, alive_sandbox_key, timeout_sandbox_key from rock.admin.metrics.billing import log_billing_info from rock.admin.metrics.decorator import monitor_sandbox_operation from rock.admin.proto.request import ClusterInfo, UserInfo @@ -35,9 +33,11 @@ from rock.sandbox.base_manager import BaseManager from rock.sandbox.operator.abstract import AbstractOperator from rock.sandbox.sandbox_actor import SandboxActor +from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError from rock.utils.crypto_utils import AESEncryption +from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.utils.format import convert_to_gb, parse_size_to_bytes from rock.utils.providers.redis_provider import RedisProvider from rock.utils.service import build_sandbox_from_redis @@ -52,22 +52,22 @@ class SandboxManager(BaseManager): def __init__( self, rock_config: RockConfig, - redis_provider: RedisProvider | None = None, + meta_store: SandboxMetaStore, ray_namespace: str = env_vars.ROCK_RAY_NAMESPACE, ray_service: RayService | None = None, enable_runtime_auto_clear: bool = False, operator: AbstractOperator | None = None, ): super().__init__( - rock_config, redis_provider=redis_provider, enable_runtime_auto_clear=enable_runtime_auto_clear + rock_config, + meta_store=meta_store, + enable_runtime_auto_clear=enable_runtime_auto_clear, ) self._ray_service = ray_service 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) + self._proxy_service = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) logger.info("sandbox service init success") async def refresh_aes_key(self): @@ -82,7 +82,7 @@ async def refresh_aes_key(self): async def _check_sandbox_exists_in_redis(self, config: DeploymentConfig): if isinstance(config, DockerDeploymentConfig) and config.container_name: sandbox_id = config.container_name - if self._redis_provider and await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$"): + if await self._meta_store.exists(sandbox_id): raise BadRequestRockError(f"Sandbox {sandbox_id} already exists") def _setup_sandbox_actor_metadata(self, sandbox_actor: SandboxActor, user_info: UserInfo) -> None: @@ -126,15 +126,14 @@ async def start_async( docker_deployment_config.cpus = self.rock_config.runtime.standard_spec.cpus docker_deployment_config.memory = self.rock_config.runtime.standard_spec.memory sandbox_info: SandboxInfo = await self._operator.submit(docker_deployment_config, user_info) - stop_time = str(int(time.time()) + docker_deployment_config.auto_clear_time * 60) - auto_clear_time_dict = { - env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: str(docker_deployment_config.auto_clear_time), - env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY: stop_time, - } await self._build_sandbox_info_metadata(sandbox_info, user_info, cluster_info) - if self._redis_provider: - await self._redis_provider.json_set(alive_sandbox_key(sandbox_id), "$", sandbox_info) - await self._redis_provider.json_set(timeout_sandbox_key(sandbox_id), "$", auto_clear_time_dict) + timeout_info = SandboxTimeoutHelper.make_timeout_info(docker_deployment_config.auto_clear_time) + await self._meta_store.create( + sandbox_id, + sandbox_info, + timeout_info=timeout_info, + deployment_config=docker_deployment_config, + ) return SandboxStartResponse( sandbox_id=sandbox_id, host_name=sandbox_info.get("host_name"), @@ -171,28 +170,32 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: @monitor_sandbox_operation() async def stop(self, sandbox_id): logger.info(f"stop sandbox {sandbox_id}") - sandbox_info: SandboxInfo = await build_sandbox_from_redis(self._redis_provider, sandbox_id) - if sandbox_info and sandbox_info.get("start_time"): + sandbox_info: SandboxInfo | None = await self._meta_store.get(sandbox_id) + if sandbox_info is None: + sandbox_info = {} + sandbox_info["state"] = State.STOPPED + if sandbox_info.get("start_time"): sandbox_info["stop_time"] = get_iso8601_timestamp() log_billing_info(sandbox_info=sandbox_info) try: await self._operator.stop(sandbox_id) except ValueError as e: logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) - await self._clear_redis_keys(sandbox_id) + await self._meta_store.archive(sandbox_id, sandbox_info) + return try: self._sandbox_meta.pop(sandbox_id) except KeyError: logger.debug(f"{sandbox_id} key not found") logger.info(f"sandbox {sandbox_id} stopped") - await self._clear_redis_keys(sandbox_id) + await self._meta_store.archive(sandbox_id, sandbox_info) async def get_mount(self, sandbox_id): async with self._ray_service.get_ray_rwlock().read_lock(): 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: - await self._clear_redis_keys(sandbox_id) + await self._meta_store.archive(sandbox_id, {}) raise Exception(f"sandbox {sandbox_id} not found to get mount") result = await self._ray_service.async_ray_get(sandbox_actor.get_mount.remote()) logger.info(f"get_mount: {result}") @@ -205,27 +208,24 @@ async def commit(self, sandbox_id, image_tag: str, username: str, password: str) 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: - await self._clear_redis_keys(sandbox_id) + await self._meta_store.archive(sandbox_id, {}) raise Exception(f"sandbox {sandbox_id} not found to commit") logger.info(f"begin to commit {sandbox_id} to {image_tag}") result = await self._ray_service.async_ray_get(sandbox_actor.commit.remote(image_tag, username, password)) logger.info(f"commit {sandbox_id} to {image_tag} finished, result {result}") return result - async def _clear_redis_keys(self, sandbox_id): - if self._redis_provider: - await self._redis_provider.json_delete(alive_sandbox_key(sandbox_id)) - await self._redis_provider.json_delete(timeout_sandbox_key(sandbox_id)) - logger.info(f"sandbox {sandbox_id} deleted from redis") - @monitor_sandbox_operation() async def get_status(self, sandbox_id) -> SandboxStatusResponse: sandbox_info: SandboxInfo = await self._operator.get_status(sandbox_id=sandbox_id) is_alive = sandbox_info.get("state") == State.RUNNING + if sandbox_info.get("state") == State.STOPPED: + raise BadRequestRockError(f"Sandbox {sandbox_id} is already stopped") self._update_sandbox_alive_info(sandbox_info, is_alive) - if self._redis_provider: - await self._redis_provider.json_set(alive_sandbox_key(sandbox_id), "$", sandbox_info) - await self._update_expire_time(sandbox_id) + current = await self._meta_store.get(sandbox_id) + if current is None or current.get("state") != sandbox_info.get("state"): + await self._meta_store.update(sandbox_id, sandbox_info) + await self._refresh_timeout(sandbox_id) return SandboxStatusResponse( sandbox_id=sandbox_id, status=sandbox_info.get("phases"), @@ -245,9 +245,9 @@ async def get_status(self, sandbox_id) -> SandboxStatusResponse: ) async def build_sandbox_info_from_redis(self, sandbox_id: str, deployment_info: SandboxInfo) -> SandboxInfo | None: - sandbox_status = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if sandbox_status and len(sandbox_status) > 0: - sandbox_info = sandbox_status[0] + sandbox_info_from_store = await self._meta_store.get(sandbox_id) + if sandbox_info_from_store: + sandbox_info = sandbox_info_from_store remote_info = { k: v for k, v in deployment_info.items() if k in ["phases", "port_mapping", "alive", "state"] } @@ -296,17 +296,23 @@ async def write_file(self, request: WriteFileRequest) -> WriteFileResponse: async def upload(self, file: UploadFile, target_path: str, sandbox_id: str) -> UploadResponse: 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), "$") - if timeout_dict is None or len(timeout_dict) == 0: - raise Exception(f"sandbox {sandbox_id} timeout key not found") + async def _refresh_timeout(self, sandbox_id: str) -> None: + timeout_info = await self._meta_store.get_timeout(sandbox_id) + if timeout_info is None: + logger.warning("refresh_timeout: timeout key not found for sandbox_id=%s", sandbox_id) + return + new_timeout = SandboxTimeoutHelper.refresh_timeout(timeout_info) + if new_timeout is None: + logger.warning("refresh_timeout: auto_clear_time missing for sandbox_id=%s", sandbox_id) + return + await self._meta_store.update_timeout(sandbox_id, new_timeout) - if timeout_dict is not None and len(timeout_dict) > 0: - expire_time: int = int(timeout_dict[0].get(env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY)) - return int(time.time()) > expire_time - else: - logger.info(f"sandbox_id:[{sandbox_id}] is already cleared") - return True + async def _is_expired(self, sandbox_id: str) -> bool: + timeout_info = await self._meta_store.get_timeout(sandbox_id) + if timeout_info is None: + logger.warning("is_expired: timeout key not found for sandbox_id=%s", sandbox_id) + return False + return SandboxTimeoutHelper.is_expired(timeout_info) async def _is_actor_alive(self, sandbox_id): try: @@ -318,11 +324,8 @@ async def _is_actor_alive(self, sandbox_id): return False async def _check_job_background(self): - if not self._redis_provider: - return logger.debug("check job background") - async for key in self._redis_provider.client.scan_iter(match=f"{ALIVE_PREFIX}*", count=100): - sandbox_id = key.removeprefix(ALIVE_PREFIX) + async for sandbox_id in self._meta_store.iter_alive_sandbox_ids(): try: is_expired = await self._is_expired(sandbox_id) if is_expired: @@ -341,26 +344,6 @@ async def get_sandbox_statistics(self, sandbox_id): resource_metrics = await self._ray_service.async_ray_get(sandbox_actor.get_sandbox_statistics.remote()) return resource_metrics - async def _update_expire_time(self, sandbox_id): - if self._redis_provider is None: - return - sandbox_status_dict = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if not sandbox_status_dict or len(sandbox_status_dict) == 0: - logger.info(f"sandbox-{sandbox_id} is not alive, skip update expire time") - return - origin_info = await self._redis_provider.json_get(timeout_sandbox_key(sandbox_id), "$") - if origin_info is None or len(origin_info) == 0: - logger.info(f"sandbox-{sandbox_id} is not initialized, skip update expire time") - return - auto_clear_time: str = origin_info[0].get(env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY) - expire_time: int = int(time.time()) + int(auto_clear_time) * 60 - logger.info(f"sandbox-{sandbox_id} update expire time: {expire_time}") - new_dict = { - env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: auto_clear_time, - env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY: str(expire_time), - } - await self._redis_provider.json_set(timeout_sandbox_key(sandbox_id), "$", new_dict) - def validate_sandbox_spec(self, runtime_config: RuntimeConfig, deployment_config: DeploymentConfig) -> None: try: memory = parse_size_to_bytes(deployment_config.memory) diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py new file mode 100644 index 0000000000..82091f5ca7 --- /dev/null +++ b/rock/sandbox/sandbox_meta_store.py @@ -0,0 +1,143 @@ +"""SandboxMetaStore - coordinator for Redis (hot path) + DB (query path) dual-write. + +Redis remains the source of truth for live sandbox state. +The database is an async replica used for indexed queries (list_by, etc.). +All DB operations are awaited for consistency. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +from typing import TYPE_CHECKING + +from rock.actions.sandbox.response import State +from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.admin.core.redis_key import alive_sandbox_key, timeout_sandbox_key +from rock.admin.core.sandbox_table import SandboxTable + +if TYPE_CHECKING: + from rock.deployments.config import DockerDeploymentConfig +from rock.logger import init_logger +from rock.utils.providers.redis_provider import RedisProvider + +logger = init_logger(__name__) + +# States that indicate an active sandbox (not yet stopped/archived). +_ACTIVE_STATES: list[str] = [State.RUNNING, State.PENDING] + + +class SandboxMetaStore: + """Coordinates sandbox metadata across Redis (hot path) and DB (query path). + + Both providers are required. Use FakeRedis / SQLite-memory for local/test environments. + """ + + def __init__( + self, + redis_provider: RedisProvider, + sandbox_table: SandboxTable, + ) -> None: + self._redis: RedisProvider = redis_provider + self._db: SandboxTable = sandbox_table + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + async def create( + self, + sandbox_id: str, + sandbox_info: SandboxInfo, + timeout_info: dict[str, str] | None = None, + deployment_config: DockerDeploymentConfig | None = None, + ) -> None: + """Write sandbox info to the Redis alive key and await DB insert. + + Parameters + ---------- + timeout_info: + If provided, also write the timeout key (``auto_clear_time`` / ``expire_time``). + deployment_config: + ``DockerDeploymentConfig`` snapshot written once to the ``spec`` DB column. + Redis does not store this. + """ + await self._redis.json_set(alive_sandbox_key(sandbox_id), "$", sandbox_info) + if timeout_info is not None: + await self._redis.json_set(timeout_sandbox_key(sandbox_id), "$", timeout_info) + + await self._db.create(sandbox_id, sandbox_info, deployment_config) + + async def update(self, sandbox_id: str, sandbox_info: SandboxInfo) -> None: + """Merge *sandbox_info* into the existing Redis alive key and await DB update.""" + current = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") + merged: dict[str, Any] = {**(current[0] if current else {}), **sandbox_info} + await self._redis.json_set(alive_sandbox_key(sandbox_id), "$", merged) + + await self._db.update(sandbox_id, sandbox_info) + + async def delete(self, sandbox_id: str) -> None: + """Delete Redis alive + timeout keys and await DB delete.""" + await self._redis.json_delete(alive_sandbox_key(sandbox_id)) + await self._redis.json_delete(timeout_sandbox_key(sandbox_id)) + + await self._db.delete(sandbox_id) + + async def archive(self, sandbox_id: str, final_info: SandboxInfo) -> None: + """Persist final state to DB, then remove sandbox from Redis. + + Unlike ``delete``, the DB record is preserved and updated with + ``final_info`` (e.g. ``stop_time``, ``state``). Use this when a + sandbox has finished its lifecycle and the final state should be + queryable from the DB. + + The DB write is awaited before the Redis keys are deleted so that + the final state is always durably stored before the alive key + disappears. If the DB write fails the exception propagates and + Redis cleanup is skipped. + """ + await self._db.update(sandbox_id, final_info) + + await self._redis.json_delete(alive_sandbox_key(sandbox_id)) + await self._redis.json_delete(timeout_sandbox_key(sandbox_id)) + + async def get(self, sandbox_id: str) -> SandboxInfo | None: + """Read sandbox info from the Redis alive key.""" + result = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") + if result and len(result) > 0: + return result[0] + return None + + async def exists(self, sandbox_id: str) -> bool: + """Return ``True`` when the Redis alive key exists for ``sandbox_id``.""" + return await self.get(sandbox_id) is not None + + async def get_timeout(self, sandbox_id: str) -> dict[str, str] | None: + """Read timeout info from the Redis timeout key.""" + timeout_info = await self._redis.json_get(timeout_sandbox_key(sandbox_id), "$") + if timeout_info and len(timeout_info) > 0: + return timeout_info[0] + return None + + async def update_timeout(self, sandbox_id: str, timeout_info: dict[str, str]) -> None: + """Overwrite the Redis timeout key with *timeout_info*.""" + await self._redis.json_set(timeout_sandbox_key(sandbox_id), "$", timeout_info) + + async def iter_alive_sandbox_ids(self) -> AsyncIterator[str]: + """Yield active sandbox IDs from the DB.""" + for sandbox_info in await self._db.list_by_in("state", _ACTIVE_STATES): + sandbox_id = sandbox_info.get("sandbox_id") + if sandbox_id: + yield sandbox_id + + async def batch_get(self, sandbox_ids: list[str]) -> list[SandboxInfo]: + """Fetch sandbox info for multiple IDs from the DB. Missing IDs are omitted.""" + if not sandbox_ids: + return [] + + return await self._db.list_by_in("sandbox_id", sandbox_ids) + + async def list_by(self, field: str, value: str | int | float | bool) -> list[SandboxInfo]: + """Query sandboxes by *field* == *value* from the DB.""" + return await self._db.list_by(field, value) diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 99430d0343..9bd719d218 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -1,6 +1,5 @@ import asyncio # noqa: I001 import json -import time from fastapi.responses import JSONResponse, StreamingResponse from starlette.datastructures import Headers @@ -23,8 +22,8 @@ UploadResponse, WriteFileResponse, ) +from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo -from rock.admin.core.redis_key import ALIVE_PREFIX, alive_sandbox_key, timeout_sandbox_key from rock.admin.metrics.decorator import monitor_sandbox_operation from rock.admin.metrics.monitor import MetricsMonitor from rock.admin.proto.request import SandboxBashAction as BashAction @@ -40,20 +39,20 @@ from rock.deployments.status import ServiceStatus from rock.common.port_validation import validate_port_forward_port from rock.logger import init_logger +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError from rock.utils import EAGLE_EYE_TRACE_ID, trace_id_ctx_var -from rock.utils.providers import RedisProvider logger = init_logger(__name__) class SandboxProxyService: - _redis_provider: RedisProvider = None _httpx_client = None - def __init__(self, rock_config: RockConfig, redis_provider: RedisProvider | None = None): + def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): self._rock_config = rock_config - self._redis_provider = redis_provider + self._meta_store = meta_store self.metrics_monitor = MetricsMonitor.create( export_interval_millis=20_000, metrics_endpoint=rock_config.runtime.metrics_endpoint, @@ -162,31 +161,30 @@ async def execute(self, command: Command) -> CommandResponse: return CommandResponse(**response) @monitor_sandbox_operation() - async def batch_get_sandbox_status_from_redis(self, sandbox_ids: list[str]) -> list[SandboxStatusResponse]: - if self._redis_provider is None: - logger.info("batch_get_sandbox_status_from_redis, redis provider is None, return empty") - return [] + async def batch_get_sandbox_status( + self, sandbox_ids: list[str] + ) -> list[SandboxStatusResponse]: if sandbox_ids is None: raise BadRequestRockError(message="sandbox_ids is None") if len(sandbox_ids) > self._batch_get_status_max_count: raise BadRequestRockError( message=f"sandbox_ids count too large, max count is {self._batch_get_status_max_count}" ) - logger.info(f"batch_get_sandbox_status_from_redis, sandbox_ids count is {len(sandbox_ids)}") + logger.info(f"batch_get_sandbox_status, sandbox_ids count is {len(sandbox_ids)}") results = [] - alive_keys = [alive_sandbox_key(sandbox_id) for sandbox_id in sandbox_ids] - sandbox_infos: list[SandboxInfo] = await self._redis_provider.json_mget(alive_keys, "$") + sandbox_infos: list[SandboxInfo] = await self._meta_store.batch_get(sandbox_ids) for sandbox_info in sandbox_infos: - if sandbox_info: - results.append(SandboxStatusResponse.from_sandbox_info(sandbox_info)) - logger.info(f"batch_get_sandbox_status_from_redis succ, result count is {len(results)}") + state = sandbox_info.get("state") + if state not in (State.RUNNING, State.PENDING): + continue + results.append(SandboxStatusResponse.from_sandbox_info(sandbox_info)) + logger.info(f"batch_get_sandbox_status succ, result count is {len(results)}") return results @monitor_sandbox_operation() - async def list_sandboxes(self, query_params: SandboxQueryParams) -> SandboxListResponse: - if self._redis_provider is None: - logger.warning("Redis provider is not available, list_sandboxes returning empty result") - return SandboxListResponse() + async def list_sandboxes( + self, query_params: SandboxQueryParams + ) -> SandboxListResponse: page = int(query_params.pop("page", "1")) page_size = int(query_params.pop("page_size", "500")) if page < 1 or page_size < 1: @@ -585,10 +583,10 @@ async def _forward_tcp_to_websocket(self, reader, websocket, direction: str, idl logger.info(f"Connection closed in {direction}: {e}") async def get_service_status(self, sandbox_id: str): - sandbox_status_dicts = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if not sandbox_status_dicts or sandbox_status_dicts[0].get("host_ip") is None: + sandbox_info = await self._meta_store.get(sandbox_id) + if not sandbox_info or sandbox_info.get("host_ip") is None: raise Exception(f"sandbox {sandbox_id} not started") - return sandbox_status_dicts + return [sandbox_info] async def _send_request( self, @@ -715,40 +713,30 @@ async def _forward_messages(self, source_ws, target_ws, direction: str): logger.error(f"Error forwarding message {direction}: {e}") async def _update_expire_time(self, sandbox_id): - if self._redis_provider is None: - return - sandbox_status_dict = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if not sandbox_status_dict or len(sandbox_status_dict) == 0: - logger.info(f"sandbox-{sandbox_id} is not alive, skip update expire time") - return - origin_info = await self._redis_provider.json_get(timeout_sandbox_key(sandbox_id), "$") - if origin_info is None or len(origin_info) == 0: - logger.info(f"sandbox-{sandbox_id} is not initialized, skip update expire time") + timeout_info = await self._meta_store.get_timeout(sandbox_id) + if timeout_info is None: return - auto_clear_time: str = origin_info[0].get(env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY) - expire_time: int = int(time.time()) + int(auto_clear_time) * 60 - logger.info(f"sandbox-{sandbox_id} update expire time: {expire_time}") - new_dict = { - env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: auto_clear_time, - env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY: str(expire_time), - } - await self._redis_provider.json_set(timeout_sandbox_key(sandbox_id), "$", new_dict) - - async def list_all_sandboxes_by_query_params(self, query_params: SandboxQueryParams): - all_keys = [] - async for key in self._redis_provider.client.scan_iter(match=f"{ALIVE_PREFIX}*", count=1000): # type: ignore - all_keys.append(key) - if not all_keys: + new_timeout = SandboxTimeoutHelper.refresh_timeout(timeout_info) + if new_timeout is not None: + await self._meta_store.update_timeout(sandbox_id, new_timeout) + + async def list_all_sandboxes_by_query_params( + self, query_params: SandboxQueryParams + ): + all_ids = [] + async for sandbox_id in self._meta_store.iter_alive_sandbox_ids(): + all_ids.append(sandbox_id) + if not all_ids: return [] all_sandbox_data = [] batch_size = self._batch_get_status_max_count - for i in range(0, len(all_keys), batch_size): - batch_keys = all_keys[i : i + batch_size] - sandbox_infos_list = await self._redis_provider.json_mget(batch_keys, "$") - + for i in range(0, len(all_ids), batch_size): + batch_ids = all_ids[i : i + batch_size] + sandbox_infos_list = await self._meta_store.batch_get(batch_ids) for sandbox_info in sandbox_infos_list: - if not sandbox_info: + state = sandbox_info.get("state") + if state not in (State.RUNNING, State.PENDING): continue if self._matches_query_params(sandbox_info, query_params): all_sandbox_data.append(SandboxListStatusResponse.from_sandbox_info(sandbox_info)) diff --git a/rock/sandbox/utils/__init__.py b/rock/sandbox/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/rock/sandbox/utils/timeout.py b/rock/sandbox/utils/timeout.py new file mode 100644 index 0000000000..c63fc65c2b --- /dev/null +++ b/rock/sandbox/utils/timeout.py @@ -0,0 +1,51 @@ +"""Sandbox timeout helpers. + +Pure calculation utilities — no I/O. Callers are responsible for reading +timeout_info from SandboxMetaStore and writing the result back. +""" + +from __future__ import annotations + +import time + +from rock import env_vars + + +class SandboxTimeoutHelper: + """Stateless timeout calculation helpers.""" + + @staticmethod + def make_timeout_info(auto_clear_time: int) -> dict[str, str]: + """Build the initial timeout dict for a newly created sandbox. + + Parameters + ---------- + auto_clear_time: + Sandbox lifetime in **minutes**. + """ + expire_time = int(time.time()) + auto_clear_time * 60 + return { + env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: str(auto_clear_time), + env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY: str(expire_time), + } + + @staticmethod + def refresh_timeout(timeout_info: dict[str, str]) -> dict[str, str] | None: + """Return a new timeout dict with ``expire_time`` recalculated from ``auto_clear_time``. + + Returns *None* if ``auto_clear_time`` is missing from *timeout_info*. + """ + auto_clear_time = timeout_info.get(env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY) + if auto_clear_time is None: + return None + expire_time = int(time.time()) + int(auto_clear_time) * 60 + return { + env_vars.ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: str(auto_clear_time), + env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY: str(expire_time), + } + + @staticmethod + def is_expired(timeout_info: dict[str, str]) -> bool: + """Return *True* when ``expire_time`` in *timeout_info* is in the past.""" + expire_time = int(timeout_info.get(env_vars.ROCK_SANDBOX_EXPIRE_TIME_KEY, 0)) + return int(time.time()) > expire_time diff --git a/scripts/gen_ddl.py b/scripts/gen_ddl.py new file mode 100644 index 0000000000..02a7013190 --- /dev/null +++ b/scripts/gen_ddl.py @@ -0,0 +1,57 @@ +"""Generate DDL SQL from SQLAlchemy ORM schema definitions. + +Usage: + uv run python scripts/gen_ddl.py # default: postgresql + uv run python scripts/gen_ddl.py --dialect sqlite + uv run python scripts/gen_ddl.py --dialect postgresql --out ddl.sql +""" + +import argparse +import sys + +from sqlalchemy.schema import CreateIndex, CreateTable + + +def get_dialect(name: str): + if name == "postgresql": + from sqlalchemy.dialects import postgresql + return postgresql.dialect() + if name == "sqlite": + from sqlalchemy.dialects import sqlite + return sqlite.dialect() + print(f"Unsupported dialect: {name}", file=sys.stderr) + sys.exit(1) + + +def gen_ddl(dialect) -> str: + # Import here so all models are registered onto Base.metadata + from rock.admin.core.schema import Base # noqa: F401 (side-effect: registers SandboxRecord) + + lines: list[str] = [] + for table in Base.metadata.sorted_tables: + lines.append(str(CreateTable(table).compile(dialect=dialect)).strip() + ";") + for index in table.indexes: + lines.append(str(CreateIndex(index).compile(dialect=dialect)).strip() + ";") + + return "\n\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate DDL from ORM schema") + parser.add_argument("--dialect", default="postgresql", choices=["postgresql", "sqlite"]) + parser.add_argument("--out", default=None, help="Output file path (default: stdout)") + args = parser.parse_args() + + dialect = get_dialect(args.dialect) + ddl = gen_ddl(dialect) + + if args.out: + with open(args.out, "w") as f: + f.write(ddl + "\n") + print(f"Written to {args.out}") + else: + print(ddl) + + +if __name__ == "__main__": + main() diff --git a/sql/sandbox_record.sql b/sql/sandbox_record.sql new file mode 100644 index 0000000000..35c4cb254e --- /dev/null +++ b/sql/sandbox_record.sql @@ -0,0 +1,42 @@ +CREATE TABLE sandbox_record ( + sandbox_id VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + image VARCHAR(128) NOT NULL, + experiment_id VARCHAR(128) NOT NULL, + namespace VARCHAR(128) NOT NULL, + cluster_name VARCHAR(128) NOT NULL, + state VARCHAR(32) NOT NULL, + host_ip VARCHAR(128) NOT NULL, + create_time VARCHAR(64) NOT NULL, + start_time VARCHAR(64), + stop_time VARCHAR(64), + host_name VARCHAR(255), + auth_token VARCHAR(512), + rock_authorization_encrypted VARCHAR(1024), + cpus FLOAT, + memory VARCHAR(64), + create_user_gray_flag BOOLEAN, + phases JSONB, + port_mapping JSONB, + spec JSONB, + status JSONB, + PRIMARY KEY (sandbox_id) +); + +CREATE INDEX ix_sandbox_record_image ON sandbox_record (image); + +CREATE INDEX ix_sandbox_record_host_ip ON sandbox_record (host_ip); + +CREATE INDEX ix_sandbox_record_host_name ON sandbox_record (host_name); + +CREATE INDEX ix_sandbox_record_state ON sandbox_record (state); + +CREATE INDEX ix_sandbox_record_create_user_gray_flag ON sandbox_record (create_user_gray_flag); + +CREATE INDEX ix_sandbox_record_user_id ON sandbox_record (user_id); + +CREATE INDEX ix_sandbox_record_namespace ON sandbox_record (namespace); + +CREATE INDEX ix_sandbox_record_experiment_id ON sandbox_record (experiment_id); + +CREATE INDEX ix_sandbox_record_cluster_name ON sandbox_record (cluster_name); diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 40e69b6b38..a54d4aa6f1 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -9,6 +9,7 @@ import urllib.request from dataclasses import dataclass from pathlib import Path +from urllib.parse import urlparse import pytest import uvicorn @@ -99,6 +100,20 @@ def admin_client_fixture(): @pytest.fixture(scope="session") def admin_remote_server(): + external_base_url = os.getenv("ROCK_TEST_ADMIN_BASE_URL") + if external_base_url: + # Accept either "localhost:8080" or "http://localhost:8080". + normalized = external_base_url if "://" in external_base_url else f"http://{external_base_url}" + parsed = urlparse(normalized) + if not parsed.hostname or not parsed.port: + raise ValueError( + "Invalid ROCK_TEST_ADMIN_BASE_URL. Expected host:port or http://host:port, " + f"got: {external_base_url!r}" + ) + logger.info("Using external admin server from ROCK_TEST_ADMIN_BASE_URL=%s", external_base_url) + yield RemoteServer(port=parsed.port, endpoint=f"{parsed.scheme}://{parsed.hostname}") + return + port = run_until_complete(find_free_port()) proxy_port = run_until_complete(find_free_port()) diff --git a/tests/unit/admin/core/test_sandbox_table.py b/tests/unit/admin/core/test_sandbox_table.py index a14c4fd429..7536c4e0ea 100644 --- a/tests/unit/admin/core/test_sandbox_table.py +++ b/tests/unit/admin/core/test_sandbox_table.py @@ -1,47 +1,198 @@ +"""Tests for SandboxTable — SQLite in-memory (fast) and PostgreSQL (Docker).""" + import pytest +from sqlalchemy.exc import IntegrityError from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.sandbox_table import SandboxTable -from rock.admin.core.schema import DBModelBase as Base -from rock.admin.core.schema import SandboxRecord from rock.config import DatabaseConfig -class DatabaseProviderUtil: - @staticmethod - async def reset(provider: DatabaseProvider): - async with provider.engine.begin() as conn: - await conn.run_sync(Base.metadata.drop_all) - await conn.run_sync(Base.metadata.create_all) - - -@pytest.mark.asyncio -async def _create_sandbox_table(): - db_provider = DatabaseProvider(DatabaseConfig(url="sqlite+aiosqlite:///:memory:")) - await db_provider.init() - await DatabaseProviderUtil.reset(db_provider) - return SandboxTable(db_provider.engine) - - -@pytest.mark.asyncio -async def test_sandbox_table(): - table = await _create_sandbox_table() - await table.create(SandboxRecord(id="1", experiment_id="1", image="1")) - sandbox_record = await table.get("1") - assert "1" == sandbox_record.id - assert "1" == sandbox_record.image - - -@pytest.mark.asyncio -async def test_list(): - table = await _create_sandbox_table() - await table.create(SandboxRecord(id="1", experiment_id="1", image="1")) - await table.create(SandboxRecord(id="2", experiment_id="1", image="2")) - await table.create(SandboxRecord(id="3", experiment_id="2", image="3")) - - sandbox_records = await table.list(experiment_id="1") - assert 2 == len(sandbox_records) - sandbox_records = await table.list(experiment_id="2") - assert 1 == len(sandbox_records) - sandbox_records = await table.list() - assert 3 == len(sandbox_records) +class TestSandboxTableWithSQLite: + """Unit tests for SandboxTable using an in-memory SQLite database. + + These tests cover all CRUD paths and run without any external dependencies. + """ + + @pytest.fixture + async def db(self): + provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite:///:memory:")) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + async def test_insert_and_get(self, db): + sandbox_id = "sqlite-sbx-001" + data = { + "user_id": "user-1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "local", + "state": "running", + "host_ip": "127.0.0.1", + "create_time": "2025-01-01T00:00:00Z", + } + await db.create(sandbox_id, data) + record = await db.get(sandbox_id) + assert record is not None + assert record["sandbox_id"] == sandbox_id + assert record["user_id"] == "user-1" + assert record["state"] == "running" + + async def test_insert_duplicate_raises(self, db): + sandbox_id = "sqlite-sbx-002" + data = {"state": "pending", "create_time": "2025-01-01T00:00:00Z"} + await db.create(sandbox_id, data) + with pytest.raises(IntegrityError): + await db.create(sandbox_id, {**data, "state": "running"}) + + async def test_update(self, db): + sandbox_id = "sqlite-sbx-003" + await db.create(sandbox_id, {"state": "pending", "create_time": "2025-01-01T00:00:00Z"}) + await db.update(sandbox_id, {"state": "running", "host_ip": "10.0.0.2"}) + record = await db.get(sandbox_id) + assert record["state"] == "running" + assert record["host_ip"] == "10.0.0.2" + + async def test_delete(self, db): + sandbox_id = "sqlite-sbx-004" + await db.create(sandbox_id, {"create_time": "2025-01-01T00:00:00Z"}) + await db.delete(sandbox_id) + assert await db.get(sandbox_id) is None + + async def test_list_by(self, db): + await db.create("lb-s1", {"user_id": "alice", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lb-s2", {"user_id": "alice", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lb-s3", {"user_id": "bob", "create_time": "2025-01-01T00:00:00Z"}) + results = await db.list_by("user_id", "alice") + assert len(results) == 2 + + async def test_list_by_in(self, db): + await db.create("lbi-1", {"user_id": "alice", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lbi-2", {"user_id": "bob", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lbi-3", {"user_id": "carol", "create_time": "2025-01-01T00:00:00Z"}) + results = await db.list_by_in("sandbox_id", ["lbi-1", "lbi-3"]) + assert {r["sandbox_id"] for r in results} == {"lbi-1", "lbi-3"} + + async def test_list_by_rejects_blacklisted_column(self, db): + with pytest.raises(ValueError, match="not allowed"): + await db.list_by("phases", "{}") + + async def test_get_nonexistent_returns_none(self, db): + assert await db.get("does-not-exist") is None + + async def test_update_nonexistent_is_noop(self, db): + """update() on a non-existent ID should log a warning and not raise.""" + await db.update("does-not-exist", {"state": "running"}) # should not raise + + async def test_not_null_defaults_applied_on_insert(self, db): + """Insert with minimal data should fill NOT NULL columns from _NOT_NULL_DEFAULTS.""" + sandbox_id = "sqlite-sbx-defaults" + await db.create(sandbox_id, {}) + record = await db.get(sandbox_id) + assert record is not None + assert record["user_id"] == "default" + assert record["state"] == "pending" + + +@pytest.mark.need_docker +@pytest.mark.need_database +class TestSandboxTableWithPostgres: + """Integration tests for SandboxTable using a real PostgreSQL container.""" + + @pytest.fixture + async def db(self, pg_container): + """Create a SandboxTable connected to the test PostgreSQL container.""" + provider = DatabaseProvider(db_config=DatabaseConfig(url=pg_container["url"])) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + async def test_fixture_connection(self, db): + """Verify that the fixture can connect and create tables.""" + assert db._db._engine is not None + + async def test_insert_and_get(self, db): + sandbox_id = "test-sandbox-001" + data = { + "user_id": "user-1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "local", + "state": "RUNNING", + "host_ip": "10.0.0.1", + "create_time": "2025-01-01T00:00:00Z", + } + + await db.create(sandbox_id, data) + record = await db.get(sandbox_id) + + assert record is not None + assert record["sandbox_id"] == sandbox_id + assert record["user_id"] == "user-1" + assert record["state"] == "RUNNING" + + async def test_insert_duplicate_raises(self, db): + sandbox_id = "test-sandbox-002" + data = { + "user_id": "user-1", + "state": "PENDING", + "create_time": "2025-01-01T00:00:00Z", + } + await db.create(sandbox_id, data) + + with pytest.raises(IntegrityError): + await db.create(sandbox_id, {**data, "state": "RUNNING"}) + + async def test_update(self, db): + sandbox_id = "test-sandbox-003" + await db.create(sandbox_id, { + "state": "PENDING", + "create_time": "2025-01-01T00:00:00Z", + }) + + await db.update(sandbox_id, {"state": "RUNNING", "host_ip": "10.0.0.2"}) + record = await db.get(sandbox_id) + assert record["state"] == "RUNNING" + assert record["host_ip"] == "10.0.0.2" + + async def test_delete(self, db): + sandbox_id = "test-sandbox-004" + await db.create(sandbox_id, {"create_time": "2025-01-01T00:00:00Z"}) + await db.delete(sandbox_id) + assert await db.get(sandbox_id) is None + + async def test_list_by(self, db): + await db.create("lb-1", {"user_id": "alice", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lb-2", {"user_id": "alice", "create_time": "2025-01-01T00:00:00Z"}) + await db.create("lb-3", {"user_id": "bob", "create_time": "2025-01-01T00:00:00Z"}) + + results = await db.list_by("user_id", "alice") + assert len(results) == 2 + + async def test_list_by_rejects_blacklisted_column(self, db): + with pytest.raises(ValueError, match="not allowed"): + await db.list_by("phases", "{}") + + async def test_json_fields_postgresql(self, db): + """Verify JSONB variant works correctly on PostgreSQL.""" + sandbox_id = "json-test-001" + data = { + "create_time": "2025-01-01T00:00:00Z", + "phases": {"build": "done", "deploy": "pending"}, + "port_mapping": {"8080": 30080, "22": 30022}, + } + await db.create(sandbox_id, data) + record = await db.get(sandbox_id) + + assert record["phases"] == {"build": "done", "deploy": "pending"} + assert record["port_mapping"] == {"8080": 30080, "22": 30022} + + async def test_get_nonexistent_returns_none(self, db): + assert await db.get("does-not-exist") is None diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 750c7e54ad..0de06aea26 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -9,8 +9,10 @@ from kubernetes import client from ray.util.state import list_actors +from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.ray_service import RayService -from rock.config import K8sConfig, RockConfig +from rock.admin.core.sandbox_table import SandboxTable +from rock.config import DatabaseConfig, K8sConfig, RockConfig from rock.deployments.abstract import AbstractDeployment from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig from rock.logger import init_logger @@ -19,6 +21,7 @@ from rock.sandbox.operator.k8s.template_loader import K8sTemplateLoader from rock.sandbox.operator.ray import RayOperator from rock.sandbox.sandbox_manager import SandboxManager +from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.utils.providers.redis_provider import RedisProvider @@ -78,13 +81,25 @@ def ray_operator(ray_service, runtime_config): return ray_operator +@pytest.fixture +async def _memory_sandbox_table(): + provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite+aiosqlite:///:memory:")) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + @pytest.fixture async def sandbox_manager( - rock_config: RockConfig, redis_provider: RedisProvider, ray_init_shutdown, ray_service, ray_operator + rock_config: RockConfig, redis_provider: RedisProvider, ray_init_shutdown, ray_service, ray_operator, + _memory_sandbox_table: SandboxTable, ): + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) sandbox_manager = SandboxManager( rock_config, - redis_provider=redis_provider, + meta_store=meta_store, ray_namespace=rock_config.ray.namespace, ray_service=ray_service, enable_runtime_auto_clear=rock_config.runtime.enable_auto_clear, @@ -94,8 +109,11 @@ async def sandbox_manager( @pytest.fixture -async def sandbox_proxy_service(rock_config: RockConfig, redis_provider: RedisProvider): - sandbox_proxy_service = SandboxProxyService(rock_config, redis_provider=redis_provider) +async def sandbox_proxy_service( + rock_config: RockConfig, redis_provider: RedisProvider, _memory_sandbox_table: SandboxTable +): + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) + sandbox_proxy_service = SandboxProxyService(rock_config, meta_store=meta_store) return sandbox_proxy_service @@ -248,3 +266,140 @@ def deployment_config(): container_name="test-sandbox", template_name="default", ) + + +# --------------------------------------------------------------------------- +# Docker container fixtures - shared across all unit test subdirectories +# (lazy-import docker so non-Docker tests don't require the package) +# --------------------------------------------------------------------------- + +_PG_IMAGE = "postgres:16-alpine" +_PG_USER = "rock_test" +_PG_PASSWORD = "rock_test_pass" +_PG_DB = "rock_test_db" +_PG_PORT = 5432 +_REDIS_IMAGE = "redis/redis-stack-server:latest" +_REDIS_PORT = 6379 + + +def _docker_keep_containers() -> bool: + import os + return os.getenv("ROCK_TEST_KEEP_DOCKER_CONTAINERS", "").lower() in {"1", "true", "yes", "on"} + + +def _docker_detect_network(client) -> str | None: + import socket + import docker + hostname = socket.gethostname() + try: + current = client.containers.get(hostname) + networks = current.attrs["NetworkSettings"]["Networks"] + if "bridge" in networks: + return "bridge" + return next(iter(networks), None) + except (docker.errors.NotFound, docker.errors.APIError): + return None + + +def _docker_resolve_host_port(container, network_name: str | None, internal_port: int) -> tuple[str, int]: + container.reload() + if network_name: + host = container.attrs["NetworkSettings"]["Networks"][network_name]["IPAddress"] + return host, internal_port + host = "127.0.0.1" + port = int(container.ports[f"{internal_port}/tcp"][0]["HostPort"]) + return host, port + + +def _docker_start_container(client, image, name, network_name, internal_port, **extra_kwargs): + keep = _docker_keep_containers() + run_kwargs = {"image": image, "name": name, "detach": True, "remove": not keep, **extra_kwargs} + if network_name: + run_kwargs["network"] = network_name + else: + run_kwargs["ports"] = {f"{internal_port}/tcp": None} + return client.containers.run(**run_kwargs) + + +@pytest.fixture(scope="session") +def pg_container(): + """Start a PostgreSQL 16 Docker container for the test session.""" + import uuid + import docker + + client = docker.from_env() + container_name = f"rock-test-pg-{uuid.uuid4().hex[:8]}" + network_name = _docker_detect_network(client) + container = _docker_start_container( + client, + image=_PG_IMAGE, + name=container_name, + network_name=network_name, + internal_port=_PG_PORT, + environment={ + "POSTGRES_USER": _PG_USER, + "POSTGRES_PASSWORD": _PG_PASSWORD, + "POSTGRES_DB": _PG_DB, + }, + ) + try: + # wait for readiness + import time as _t + deadline = _t.time() + 30 + while _t.time() < deadline: + code, _ = container.exec_run(f"pg_isready -U {_PG_USER}") + if code == 0: + break + _t.sleep(0.5) + else: + raise TimeoutError("PostgreSQL container did not become ready within 30s") + + host, port = _docker_resolve_host_port(container, network_name, _PG_PORT) + yield { + "host": host, "port": port, + "user": _PG_USER, "password": _PG_PASSWORD, "database": _PG_DB, + "url": f"postgresql://{_PG_USER}:{_PG_PASSWORD}@{host}:{port}/{_PG_DB}", + } + finally: + if not _docker_keep_containers(): + try: + container.stop(timeout=5) + except Exception: + pass + + +@pytest.fixture(scope="session") +def redis_container(): + """Start a Redis Stack Docker container (with RedisJSON) for the test session.""" + import uuid + import docker + + client = docker.from_env() + container_name = f"rock-test-redis-{uuid.uuid4().hex[:8]}" + network_name = _docker_detect_network(client) + container = _docker_start_container( + client, + image=_REDIS_IMAGE, + name=container_name, + network_name=network_name, + internal_port=_REDIS_PORT, + ) + try: + import time as _t + deadline = _t.time() + 30 + while _t.time() < deadline: + code, output = container.exec_run("redis-cli ping") + if code == 0 and b"PONG" in output: + break + _t.sleep(0.5) + else: + raise TimeoutError("Redis container did not become ready within 30s") + + host, port = _docker_resolve_host_port(container, network_name, _REDIS_PORT) + yield {"host": host, "port": port, "password": "", "url": f"redis://{host}:{port}"} + finally: + if not _docker_keep_containers(): + try: + container.stop(timeout=5) + except Exception: + pass diff --git a/tests/unit/sandbox/test_proxy_enhancements.py b/tests/unit/sandbox/test_proxy_enhancements.py index 9f3c3ece47..eb54fa18b5 100644 --- a/tests/unit/sandbox/test_proxy_enhancements.py +++ b/tests/unit/sandbox/test_proxy_enhancements.py @@ -1,16 +1,27 @@ """Tests for proxy enhancements: 1. WebSocket proxy supports user-specified port 2. HTTP proxy supports all HTTP methods +3. batch_get_sandbox_status legacy-states filtering """ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import aioredis from fastapi import FastAPI from httpx import ASGITransport, AsyncClient from starlette.datastructures import Headers, MutableHeaders from starlette.responses import JSONResponse +from rock.actions.sandbox.response import State +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.sandbox_table import SandboxTable +from rock.config import DatabaseConfig, RockConfig +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService +from rock.utils.providers.redis_provider import RedisProvider + +from rock.admin.proto.response import SandboxListResponse from rock.admin.entrypoints.sandbox_proxy_api import ( sandbox_proxy_router, set_sandbox_proxy_service, @@ -1023,3 +1034,97 @@ async def test_vnc_ws_route_ignores_query_param_port(self, app): call = svc.websocket_proxy.call_args port = call.kwargs.get("port") or (call.args[3] if len(call.args) > 3 else None) assert port == 8006 + + +# ───────────────────────────────────────────────────────────────────────────── +# batch_get_sandbox_status — legacy-states filtering +# ───────────────────────────────────────────────────────────────────────────── + +_BASE_INFO = { + "sandbox_id": None, + "user_id": "u1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "c1", + "host_ip": "10.0.0.1", + "create_time": "2025-01-01T00:00:00Z", +} + + +@pytest.fixture +async def _redis(): + provider = RedisProvider(host=None, port=None, password="") + provider.client = aioredis.FakeRedis(decode_responses=True) + yield provider + await provider.close_pool() + + +@pytest.fixture +async def _db(): + provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite+aiosqlite:///:memory:")) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + +@pytest.fixture +def _svc(_redis, _db, rock_config): + meta_store = SandboxMetaStore(redis_provider=_redis, sandbox_table=_db) + return SandboxProxyService(rock_config, meta_store=meta_store), meta_store + + +@pytest.fixture +def rock_config(): + return RockConfig() + + +async def _seed(meta_store: SandboxMetaStore, sandbox_id: str, state: str) -> None: + info = {**_BASE_INFO, "sandbox_id": sandbox_id, "state": state} + await meta_store.create(sandbox_id, info) + + +class TestBatchGetLegacyStates: + """batch_get_sandbox_status only returns RUNNING/PENDING sandboxes.""" + + async def test_running_sandbox_included(self, _svc): + svc, meta_store = _svc + await _seed(meta_store, "sb-running", State.RUNNING) + result = await svc.batch_get_sandbox_status(["sb-running"]) + assert len(result) == 1 + assert result[0].sandbox_id == "sb-running" + + async def test_pending_sandbox_included(self, _svc): + svc, meta_store = _svc + await _seed(meta_store, "sb-pending", State.PENDING) + result = await svc.batch_get_sandbox_status(["sb-pending"]) + assert len(result) == 1 + assert result[0].sandbox_id == "sb-pending" + + async def test_stopped_sandbox_excluded(self, _svc): + """stopped sandbox must be filtered out.""" + svc, meta_store = _svc + await _seed(meta_store, "sb-stopped", State.STOPPED) + result = await svc.batch_get_sandbox_status(["sb-stopped"]) + assert result == [] + + async def test_mixed_states_only_active_returned(self, _svc): + """Only running/pending survive; stopped is silently dropped.""" + svc, meta_store = _svc + await _seed(meta_store, "sb-r", State.RUNNING) + await _seed(meta_store, "sb-p", State.PENDING) + await _seed(meta_store, "sb-s", State.STOPPED) + result = await svc.batch_get_sandbox_status(["sb-r", "sb-p", "sb-s"]) + ids = {r.sandbox_id for r in result} + assert ids == {"sb-r", "sb-p"} + + async def test_unknown_id_omitted(self, _svc): + """sandbox_id not in DB → entry absent from result (not None/empty status).""" + svc, meta_store = _svc + await _seed(meta_store, "sb-exists", State.RUNNING) + result = await svc.batch_get_sandbox_status(["sb-exists", "sb-ghost"]) + assert len(result) == 1 + assert result[0].sandbox_id == "sb-exists" + diff --git a/tests/unit/sandbox/test_sandbox_meta_store.py b/tests/unit/sandbox/test_sandbox_meta_store.py new file mode 100644 index 0000000000..5b034e98ac --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_meta_store.py @@ -0,0 +1,490 @@ +"""Tests for SandboxMetaStore - Redis + DB dual-write coordinator.""" + +import asyncio +import time +import uuid + +import pytest +from fakeredis import aioredis + +from rock.actions.sandbox.response import State +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.redis_key import ALIVE_PREFIX, alive_sandbox_key, timeout_sandbox_key +from rock.admin.core.sandbox_table import SandboxTable +from rock.config import DatabaseConfig +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.utils.providers.redis_provider import RedisProvider + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +async def redis(): + provider = RedisProvider(host=None, port=None, password="") + provider.client = aioredis.FakeRedis(decode_responses=True) + yield provider + await provider.close_pool() + + +@pytest.fixture +async def db(tmp_path): + provider = DatabaseProvider(db_config=DatabaseConfig(url=f"sqlite:///{tmp_path / 'test.db'}")) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + +@pytest.fixture +async def db_memory(): + provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite:///:memory:")) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + +@pytest.fixture +def repo(redis, db): + return SandboxMetaStore(redis_provider=redis, sandbox_table=db) + + + +@pytest.fixture +def repo_with_memory_db(redis, db_memory): + return SandboxMetaStore(redis_provider=redis, sandbox_table=db_memory) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +SANDBOX_ID = "sbx-test-001" + +SANDBOX_INFO = { + "sandbox_id": SANDBOX_ID, + "user_id": "user-1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "cluster-1", + "state": State.RUNNING, + "host_ip": "10.0.0.1", + "create_time": "2025-01-01T00:00:00Z", +} + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestSave: + async def test_save_writes_redis_and_db(self, repo, redis, db): + """save() should persist to Redis alive key AND fire a DB upsert.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + + # Give fire-and-forget task time to complete + await asyncio.sleep(0.1) + + # Verify Redis + result = await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") + assert result is not None + assert result[0]["sandbox_id"] == SANDBOX_ID + assert result[0]["user_id"] == "user-1" + + # Verify DB + db_record = await db.get(SANDBOX_ID) + assert db_record is not None + assert db_record["user_id"] == "user-1" + + async def test_save_with_timeout_info(self, repo, redis): + """save() with timeout_info should also write the timeout key.""" + timeout = {"auto_clear_time": "30", "expire_time": "9999999999"} + await repo.create(SANDBOX_ID, SANDBOX_INFO, timeout_info=timeout) + + result = await redis.json_get(timeout_sandbox_key(SANDBOX_ID), "$") + assert result is not None + assert result[0]["auto_clear_time"] == "30" + + +class TestUpdate: + async def test_update_writes_redis_and_db(self, repo, redis, db): + """update() should merge new fields into Redis and fire DB update.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + await asyncio.sleep(0.1) + + update_data = {"state": "stopped", "stop_time": "2025-01-01T01:00:00Z"} + await repo.update(SANDBOX_ID, update_data) + await asyncio.sleep(0.1) + + # Verify Redis - should have merged (old fields + new fields) + result = await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") + assert result is not None + info = result[0] + assert info["state"] == "stopped" + assert info["stop_time"] == "2025-01-01T01:00:00Z" + # Original fields should still be present + assert info["user_id"] == "user-1" + assert info["image"] == "python:3.11" + + # Verify DB + db_record = await db.get(SANDBOX_ID) + assert db_record is not None + assert db_record["state"] == "stopped" + + +class TestRemove: + async def test_remove_deletes_redis_and_db(self, repo, redis, db): + """remove() should delete from both Redis alive+timeout keys and DB.""" + # Setup: save sandbox and a timeout key + await repo.create(SANDBOX_ID, SANDBOX_INFO) + timeout_data = {"auto_clear_time": "30", "expire_time": str(int(time.time()) + 1800)} + await redis.json_set(timeout_sandbox_key(SANDBOX_ID), "$", timeout_data) + await asyncio.sleep(0.1) + + # Act + await repo.delete(SANDBOX_ID) + await asyncio.sleep(0.1) + + # Verify Redis - both keys gone + alive_result = await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") + assert alive_result is None + timeout_result = await redis.json_get(timeout_sandbox_key(SANDBOX_ID), "$") + assert timeout_result is None + + # Verify DB + db_record = await db.get(SANDBOX_ID) + assert db_record is None + + +class TestArchive: + async def test_archive_removes_redis_and_updates_db(self, repo, redis, db): + """archive() should update DB first, then remove Redis keys.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + await redis.json_set(timeout_sandbox_key(SANDBOX_ID), "$", {"auto_clear_time": "30", "expire_time": "9999"}) + await asyncio.sleep(0.1) # let the create fire-and-forget DB insert settle + + final_info: dict = {"state": "stopped", "stop_time": "2025-06-01T00:00:00Z"} + await repo.archive(SANDBOX_ID, final_info) + # No extra sleep needed: archive() awaits the DB write before returning. + + # Redis: both keys gone + assert await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") is None + assert await redis.json_get(timeout_sandbox_key(SANDBOX_ID), "$") is None + + # DB: record still present with updated fields + db_record = await db.get(SANDBOX_ID) + assert db_record is not None + assert db_record["state"] == "stopped" + assert db_record["stop_time"] == "2025-06-01T00:00:00Z" + assert db_record["user_id"] == "user-1" # original fields preserved + + async def test_archive_db_written_before_redis_deleted(self, repo, redis, db): + """DB must be durably updated before the Redis alive key is removed.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + await asyncio.sleep(0.1) + + # Intercept: check DB state immediately after archive returns (no extra sleep). + await repo.archive(SANDBOX_ID, {"state": "stopped"}) + + # At this point archive() has already awaited the DB write and deleted Redis. + assert await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") is None + db_record = await db.get(SANDBOX_ID) + assert db_record is not None + assert db_record["state"] == "stopped" + + +class TestGet: + async def test_get_reads_from_redis(self, repo, redis): + """get() should read from Redis alive key.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + + result = await repo.get(SANDBOX_ID) + assert result is not None + assert result["sandbox_id"] == SANDBOX_ID + assert result["user_id"] == "user-1" + + async def test_get_nonexistent_returns_none(self, repo): + """get() on a non-existent sandbox should return None.""" + result = await repo.get("does-not-exist") + assert result is None + + +class TestExists: + async def test_exists_returns_true_when_present(self, repo, redis): + """exists() should return True when the sandbox alive key exists.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + + assert await repo.exists(SANDBOX_ID) is True + + async def test_exists_returns_false_when_absent(self, repo): + """exists() should return False for a non-existent sandbox.""" + assert await repo.exists("does-not-exist") is False + + +class TestGetTimeout: + async def test_get_timeout_returns_timeout_info(self, repo, redis): + """get_timeout() should return the timeout dict from Redis.""" + timeout_data = {"auto_clear_time": "30", "expire_time": "9999999999"} + await redis.json_set(timeout_sandbox_key(SANDBOX_ID), "$", timeout_data) + + result = await repo.get_timeout(SANDBOX_ID) + assert result is not None + assert result["auto_clear_time"] == "30" + assert result["expire_time"] == "9999999999" + + async def test_get_timeout_returns_none_when_absent(self, repo): + """get_timeout() should return None when the timeout key does not exist.""" + result = await repo.get_timeout("does-not-exist") + assert result is None + + +class TestIterAliveSandboxIds: + async def test_iter_alive_sandbox_ids_yields_running_and_pending(self, repo): + """iter_alive_sandbox_ids() should yield IDs for both RUNNING and PENDING sandboxes.""" + await repo.create("sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running", "state": State.RUNNING}) + await repo.create("sbx-pending", {**SANDBOX_INFO, "sandbox_id": "sbx-pending", "state": State.PENDING}) + await asyncio.sleep(0.1) # let fire-and-forget DB writes settle + + ids = {sid async for sid in repo.iter_alive_sandbox_ids()} + assert "sbx-running" in ids + assert "sbx-pending" in ids + + async def test_iter_alive_sandbox_ids_excludes_stopped(self, repo): + """iter_alive_sandbox_ids() should not yield sandboxes with terminal state.""" + await repo.create("sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running"}) + await repo.create("sbx-stopped", {**SANDBOX_INFO, "sandbox_id": "sbx-stopped", "state": "stopped"}) + await asyncio.sleep(0.1) + + ids = [sid async for sid in repo.iter_alive_sandbox_ids()] + assert "sbx-running" in ids + assert "sbx-stopped" not in ids + + async def test_iter_alive_sandbox_ids_works_with_sqlite_memory(self, repo_with_memory_db): + """iter_alive_sandbox_ids() should work with sqlite in-memory DB + Redis fallback.""" + await repo_with_memory_db.create("sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running", "state": State.RUNNING}) + await repo_with_memory_db.create("sbx-pending", {**SANDBOX_INFO, "sandbox_id": "sbx-pending", "state": State.PENDING}) + await repo_with_memory_db.create("sbx-stopped", {**SANDBOX_INFO, "sandbox_id": "sbx-stopped", "state": "stopped"}) + await asyncio.sleep(0.1) + + ids = {sid async for sid in repo_with_memory_db.iter_alive_sandbox_ids()} + assert "sbx-running" in ids + assert "sbx-pending" in ids + assert "sbx-stopped" not in ids + + async def test_iter_alive_sandbox_ids_consistent_with_redis_scan(self, repo, redis): + """DB list_by_in(state IN active_states) should be consistent with Redis alive-key scan. + + Both approaches must agree: every active sandbox (PENDING or RUNNING) found in DB + must also have a Redis alive key. The inverse may not hold for sandboxes whose state + was updated to a terminal value without calling archive()/remove(). + """ + await repo.create("sbx-a", {**SANDBOX_INFO, "sandbox_id": "sbx-a", "state": State.RUNNING}) + await repo.create("sbx-b", {**SANDBOX_INFO, "sandbox_id": "sbx-b", "state": State.PENDING}) + await asyncio.sleep(0.1) + + # new approach: DB-backed iter_alive_sandbox_ids (PENDING + RUNNING) + db_ids = {sid async for sid in repo.iter_alive_sandbox_ids()} + + # old approach: Redis scan_iter on alive: prefix + redis_ids = set() + async for key in redis.client.scan_iter(match=f"{ALIVE_PREFIX}*", count=100): + if isinstance(key, str) and key.startswith(ALIVE_PREFIX): + redis_ids.add(key.removeprefix(ALIVE_PREFIX)) + + assert db_ids == redis_ids + + +class TestBatchGet: + async def test_batch_get_returns_db_results(self, repo): + """batch_get() returns sandbox info from the DB.""" + await repo.create(SANDBOX_ID, SANDBOX_INFO) + + results = await repo.batch_get([SANDBOX_ID]) + assert len(results) == 1 + assert results[0]["sandbox_id"] == SANDBOX_ID + + async def test_batch_get_omits_unknown_id(self, repo): + """batch_get() omits IDs not found in DB.""" + results = await repo.batch_get(["does-not-exist"]) + assert results == [] + + async def test_batch_get_multiple_ids(self, repo): + """batch_get() returns only found sandboxes; missing IDs are omitted.""" + await repo.create("sbx-a", {**SANDBOX_INFO, "sandbox_id": "sbx-a"}) + await repo.create("sbx-b", {**SANDBOX_INFO, "sandbox_id": "sbx-b"}) + + results = await repo.batch_get(["sbx-a", "sbx-b", "sbx-missing"]) + assert len(results) == 2 + sandbox_ids = {r["sandbox_id"] for r in results} + assert sandbox_ids == {"sbx-a", "sbx-b"} + + async def test_batch_get_empty_list(self, repo): + """batch_get([]) should return [].""" + assert await repo.batch_get([]) == [] + + +class TestListBy: + async def test_list_by_queries_db(self, repo, db): + """list_by() should query the DB by a given field.""" + info_a = {**SANDBOX_INFO, "sandbox_id": "sbx-a", "user_id": "user-1"} + info_b = {**SANDBOX_INFO, "sandbox_id": "sbx-b", "user_id": "user-1"} + info_c = {**SANDBOX_INFO, "sandbox_id": "sbx-c", "user_id": "user-2"} + + await repo.create("sbx-a", info_a) + await repo.create("sbx-b", info_b) + await repo.create("sbx-c", info_c) + await asyncio.sleep(0.2) + + results = await repo.list_by("user_id", "user-1") + assert len(results) == 2 + sandbox_ids = {r["sandbox_id"] for r in results} + assert sandbox_ids == {"sbx-a", "sbx-b"} + + async def test_list_by_raises_for_non_allowlisted_field(self, repo): + """list_by() should raise ValueError for fields not in the DB allowlist.""" + with pytest.raises(ValueError): + await repo.list_by("create_time", "t-1") + + +class TestUpdateTimeout: + async def test_update_timeout_writes_redis(self, repo, redis): + """update_timeout() should overwrite the timeout key in Redis.""" + new_info = {"auto_clear_time": "60", "expire_time": "9999999999"} + await repo.update_timeout(SANDBOX_ID, new_info) + + result = await redis.json_get(timeout_sandbox_key(SANDBOX_ID), "$") + assert result is not None + assert result[0]["auto_clear_time"] == "60" + assert result[0]["expire_time"] == "9999999999" + + +# --------------------------------------------------------------------------- +# Docker-backed fixtures (real Redis Stack + real PostgreSQL) +# --------------------------------------------------------------------------- + + +@pytest.fixture +async def real_redis(redis_container): + provider = RedisProvider( + host=redis_container["host"], + port=redis_container["port"], + password=redis_container["password"], + ) + await provider.init_pool() + yield provider + await provider.close_pool() + + +@pytest.fixture +async def real_db(pg_container): + provider = DatabaseProvider(db_config=DatabaseConfig(url=pg_container["url"])) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + +@pytest.fixture +def docker_repo(real_redis, real_db): + return SandboxMetaStore(redis_provider=real_redis, sandbox_table=real_db) + + +# --------------------------------------------------------------------------- +# Docker-backed tests +# --------------------------------------------------------------------------- + + +@pytest.mark.need_docker +@pytest.mark.need_database +class TestSandboxMetaStoreWithDocker: + """SandboxMetaStore verified against real Redis Stack + PostgreSQL. + + Uses unique sandbox IDs per test to avoid cross-test pollution across + the shared session-scoped containers. + """ + + async def test_save_writes_redis_and_db(self, docker_repo, real_redis, real_db): + """save() persists to real Redis and fires a real DB insert.""" + sid = f"docker-{uuid.uuid4().hex[:8]}" + info = {**SANDBOX_INFO, "sandbox_id": sid, "user_id": "docker-user"} + + await docker_repo.create(sid, info) + await asyncio.sleep(0.15) + + result = await real_redis.json_get(alive_sandbox_key(sid), "$") + assert result is not None + assert result[0]["sandbox_id"] == sid + + db_record = await real_db.get(sid) + assert db_record is not None + assert db_record["user_id"] == "docker-user" + + async def test_update_writes_redis_and_db(self, docker_repo, real_redis, real_db): + """update() merges into Redis and fires a real DB update.""" + sid = f"docker-{uuid.uuid4().hex[:8]}" + await docker_repo.create(sid, {**SANDBOX_INFO, "sandbox_id": sid}) + await asyncio.sleep(0.15) + + await docker_repo.update(sid, {"state": "stopped"}) + await asyncio.sleep(0.15) + + redis_result = await real_redis.json_get(alive_sandbox_key(sid), "$") + assert redis_result[0]["state"] == "stopped" + assert redis_result[0]["user_id"] == "user-1" # old fields still present + + db_record = await real_db.get(sid) + assert db_record["state"] == "stopped" + + async def test_remove_deletes_redis_and_db(self, docker_repo, real_redis, real_db): + """remove() deletes Redis alive+timeout keys and the DB row.""" + sid = f"docker-{uuid.uuid4().hex[:8]}" + timeout = {"auto_clear_time": "30", "expire_time": "9999999999"} + await docker_repo.create(sid, {**SANDBOX_INFO, "sandbox_id": sid}, timeout_info=timeout) + await asyncio.sleep(0.15) + + await docker_repo.delete(sid) + await asyncio.sleep(0.15) + + assert await real_redis.json_get(alive_sandbox_key(sid), "$") is None + assert await real_db.get(sid) is None + + async def test_list_by_queries_db(self, docker_repo, real_db): + """list_by() returns DB rows matching the given field value.""" + uid = f"docker-user-{uuid.uuid4().hex[:8]}" + for _ in range(3): + sid = f"docker-{uuid.uuid4().hex[:8]}" + await docker_repo.create(sid, {**SANDBOX_INFO, "sandbox_id": sid, "user_id": uid}) + await asyncio.sleep(0.2) + + results = await docker_repo.list_by("user_id", uid) + assert len(results) == 3 + assert all(r["user_id"] == uid for r in results) + + async def test_iter_alive_sandbox_ids(self, docker_repo, real_redis): + """iter_alive_sandbox_ids() yields RUNNING sandbox IDs from DB, consistent with Redis alive keys.""" + sids = [f"docker-{uuid.uuid4().hex[:8]}" for _ in range(3)] + for sid in sids: + await docker_repo.create(sid, {**SANDBOX_INFO, "sandbox_id": sid}) + await asyncio.sleep(0.2) # let fire-and-forget DB writes settle + + # new approach: DB-backed + db_found = [s async for s in docker_repo.iter_alive_sandbox_ids()] + assert set(sids).issubset(set(db_found)) + + # consistency: every RUNNING sandbox in DB must have a Redis alive key + redis_ids = set() + async for key in real_redis.client.scan_iter(match=f"{ALIVE_PREFIX}*", count=100): + if isinstance(key, str) and key.startswith(ALIVE_PREFIX): + redis_ids.add(key.removeprefix(ALIVE_PREFIX)) + assert set(sids).issubset(redis_ids) + # Only assert consistency for the sandboxes created in this test; + # db_found may include leftover sandboxes from other tests whose + # alive keys have already been cleaned up. + assert set(sids).issubset(redis_ids) diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index 55603c5771..f98802c6b7 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -20,7 +20,7 @@ async def test_batch_get_sandbox_status(sandbox_manager: SandboxManager, sandbox sandbox_ids.append(response.sandbox_id) await check_sandbox_status_until_alive(sandbox_manager, response.sandbox_id) # batch get status - batch_response = await sandbox_proxy_service.batch_get_sandbox_status_from_redis(sandbox_ids) + batch_response = await sandbox_proxy_service.batch_get_sandbox_status(sandbox_ids) assert len(batch_response) == sandbox_count response_sandbox_ids = [status.sandbox_id for status in batch_response] @@ -33,7 +33,7 @@ async def test_batch_get_sandbox_status(sandbox_manager: SandboxManager, sandbox assert status.state == State.RUNNING invalid_ids = sandbox_ids + ["invalid_sandbox_id_1", "invalid_sandbox_id_2"] - batch_response_with_invalid = await sandbox_proxy_service.batch_get_sandbox_status_from_redis(invalid_ids) + batch_response_with_invalid = await sandbox_proxy_service.batch_get_sandbox_status(invalid_ids) assert len(batch_response_with_invalid) == len(sandbox_ids) for sandbox_id in sandbox_ids: await sandbox_manager.stop(sandbox_id) diff --git a/tests/unit/sandbox/test_sandbox_timeout.py b/tests/unit/sandbox/test_sandbox_timeout.py new file mode 100644 index 0000000000..1118c25ad6 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_timeout.py @@ -0,0 +1,55 @@ +"""Tests for SandboxTimeoutHelper — pure calculation, no I/O.""" + +import time + +from rock.sandbox.utils.timeout import SandboxTimeoutHelper + + +class TestMakeTimeoutInfo: + def test_contains_correct_keys(self): + info = SandboxTimeoutHelper.make_timeout_info(30) + assert "auto_clear_time" in info + assert "expire_time" in info + + def test_expire_time_is_now_plus_duration(self): + before = int(time.time()) + info = SandboxTimeoutHelper.make_timeout_info(30) + after = int(time.time()) + expire = int(info["expire_time"]) + assert before + 30 * 60 <= expire <= after + 30 * 60 + + def test_auto_clear_time_stored_as_string(self): + info = SandboxTimeoutHelper.make_timeout_info(45) + assert info["auto_clear_time"] == "45" + + +class TestRefreshTimeout: + def test_recalculates_expire_time(self): + old_expire = int(time.time()) - 100 + timeout_info = {"auto_clear_time": "30", "expire_time": str(old_expire)} + result = SandboxTimeoutHelper.refresh_timeout(timeout_info) + assert result is not None + assert int(result["expire_time"]) >= int(time.time()) + 30 * 60 - 5 + + def test_preserves_auto_clear_time(self): + timeout_info = {"auto_clear_time": "60", "expire_time": "0"} + result = SandboxTimeoutHelper.refresh_timeout(timeout_info) + assert result["auto_clear_time"] == "60" + + def test_returns_none_when_auto_clear_time_missing(self): + result = SandboxTimeoutHelper.refresh_timeout({"expire_time": "0"}) + assert result is None + + +class TestIsExpired: + def test_returns_true_when_past(self): + timeout_info = {"auto_clear_time": "30", "expire_time": str(int(time.time()) - 100)} + assert SandboxTimeoutHelper.is_expired(timeout_info) is True + + def test_returns_false_when_future(self): + timeout_info = {"auto_clear_time": "30", "expire_time": str(int(time.time()) + 3600)} + assert SandboxTimeoutHelper.is_expired(timeout_info) is False + + def test_returns_true_when_expire_time_missing(self): + # Missing key → defaults to 0, always in the past + assert SandboxTimeoutHelper.is_expired({}) is True diff --git a/tests/unit/utils/test_redis_provider_docker.py b/tests/unit/utils/test_redis_provider_docker.py new file mode 100644 index 0000000000..2e63d3d3d9 --- /dev/null +++ b/tests/unit/utils/test_redis_provider_docker.py @@ -0,0 +1,90 @@ +"""Tests for RedisProvider against a real Redis Stack container (with RedisJSON).""" + +import pytest + +from rock.utils.providers.redis_provider import RedisProvider + + +@pytest.mark.need_docker +@pytest.mark.need_database +class TestRedisProviderWithDocker: + """Integration tests for RedisProvider using a real Redis Stack container.""" + + @pytest.fixture + async def redis(self, redis_container): + """Create a RedisProvider connected to the test Redis container.""" + provider = RedisProvider( + host=redis_container["host"], + port=redis_container["port"], + password=redis_container["password"], + ) + await provider.init_pool() + yield provider + await provider.close_pool() + + async def test_fixture_connection(self, redis): + """Verify the fixture can connect and ping Redis.""" + assert redis.client is not None + pong = await redis.client.ping() + assert pong is True + + async def test_json_set_and_get(self, redis): + key = "test:sandbox:001" + data = {"sandbox_id": "sb-001", "state": "RUNNING", "user_id": "alice"} + + await redis.json_set(key, "$", data) + result = await redis.json_get(key, "$") + + assert result is not None + assert result[0]["sandbox_id"] == "sb-001" + assert result[0]["state"] == "RUNNING" + + async def test_json_get_subpath(self, redis): + key = "test:sandbox:002" + data = {"sandbox_id": "sb-002", "spec": {"cpus": 4, "memory": "8Gi"}} + + await redis.json_set(key, "$", data) + result = await redis.json_get(key, "$.spec.cpus") + + assert result == 4 + + async def test_json_set_with_ttl(self, redis): + key = "test:sandbox:ttl" + data = {"sandbox_id": "sb-ttl"} + + await redis.json_set_with_ttl(key, "$", data, ttl_seconds=300) + + ttl = await redis.get_ttl(key) + assert ttl is not None + assert ttl > 0 + assert ttl <= 300 + + async def test_json_delete(self, redis): + key = "test:sandbox:del" + await redis.json_set(key, "$", {"sandbox_id": "sb-del"}) + + deleted = await redis.json_delete(key) + assert deleted >= 1 + + result = await redis.json_get(key, "$") + assert result is None + + async def test_json_mget(self, redis): + keys = ["test:mget:1", "test:mget:2", "test:mget:3"] + for i, key in enumerate(keys): + await redis.json_set(key, "$", {"name": f"user-{i}"}) + + results = await redis.json_mget(keys, "$") + assert len(results) == 3 + assert results[0][0]["name"] == "user-0" + assert results[2][0]["name"] == "user-2" + + async def test_pattern_exists(self, redis): + await redis.json_set("test:pattern:hit", "$", {"v": 1}) + + assert await redis.pattern_exists("test:pattern:*") is True + assert await redis.pattern_exists("nonexistent:prefix:*") is False + + async def test_get_nonexistent_key(self, redis): + result = await redis.json_get("does-not-exist", "$") + assert result is None diff --git a/uv.lock b/uv.lock index 7ae9b8b4c4..cf569139eb 100644 --- a/uv.lock +++ b/uv.lock @@ -467,6 +467,65 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c" }, ] +[[package]] +name = "asyncpg" +version = "0.31.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "async-timeout", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/fe/cc/d18065ce2380d80b1bcce927c24a2642efd38918e33fd724bc4bca904877/asyncpg-0.31.0.tar.gz", hash = "sha256:c989386c83940bfbd787180f2b1519415e2d3d6277a70d9d0f0145ac73500735" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/c3/d9/507c80bdac2e95e5a525644af94b03fa7f9a44596a84bd48a6e80f854f92/asyncpg-0.31.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:831712dd3cf117eec68575a9b50da711893fd63ebe277fc155ecae1c6c9f0f61" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ea/03/f93b5e543f65c5f504e91405e8d21bb9e600548be95032951a754781a41d/asyncpg-0.31.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0b17c89312c2f4ccea222a3a6571f7df65d4ba2c0e803339bfc7bed46a96d3be" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e5/1e/de2177e57e03a06e697f6c1ddf2a9a7fcfdc236ce69966f54ffc830fd481/asyncpg-0.31.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3faa62f997db0c9add34504a68ac2c342cfee4d57a0c3062fcf0d86c7f9cb1e8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d0/98/1a853f6870ac7ad48383a948c8ff3c85dc278066a4d69fc9af7d3d4b1106/asyncpg-0.31.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ea599d45c361dfbf398cb67da7fd052affa556a401482d3ff1ee99bd68808a1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/11/29/7e76f2a51f2360a7c90d2cf6d0d9b210c8bb0ae342edebd16173611a55c2/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:795416369c3d284e1837461909f58418ad22b305f955e625a4b3a2521d80a5f3" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5d/3f/716e10cb57c4f388248db46555e9226901688fbfabd0afb85b5e1d65d5a7/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:a8d758dac9d2e723e173d286ef5e574f0b350ec00e9186fce84d0fc5f6a8e6b8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7e/ec/3ebae9dfb23a1bd3f68acfd4f795983b65b413291c0e2b0d982d6ae6c920/asyncpg-0.31.0-cp310-cp310-win32.whl", hash = "sha256:2d076d42eb583601179efa246c5d7ae44614b4144bc1c7a683ad1222814ed095" }, + { url = "https://mirrors.aliyun.com/pypi/packages/20/b4/9fbb4b0af4e36d96a61d026dd37acab3cf521a70290a09640b215da5ab7c/asyncpg-0.31.0-cp310-cp310-win_amd64.whl", hash = "sha256:9ea33213ac044171f4cac23740bed9a3805abae10e7025314cfbd725ec670540" }, + { url = "https://mirrors.aliyun.com/pypi/packages/08/17/cc02bc49bc350623d050fa139e34ea512cd6e020562f2a7312a7bcae4bc9/asyncpg-0.31.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:eee690960e8ab85063ba93af2ce128c0f52fd655fdff9fdb1a28df01329f031d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a4/62/4ded7d400a7b651adf06f49ea8f73100cca07c6df012119594d1e3447aa6/asyncpg-0.31.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2657204552b75f8288de08ca60faf4a99a65deef3a71d1467454123205a88fab" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d6/5b/4179538a9a72166a0bf60ad783b1ef16efb7960e4d7b9afe9f77a5551680/asyncpg-0.31.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a429e842a3a4b4ea240ea52d7fe3f82d5149853249306f7ff166cb9948faa46c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e6/35/c27719ae0536c5b6e61e4701391ffe435ef59539e9360959240d6e47c8c8/asyncpg-0.31.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c0807be46c32c963ae40d329b3a686356e417f674c976c07fa49f1b30303f109" }, + { url = "https://mirrors.aliyun.com/pypi/packages/43/f4/01ebb9207f29e645a64699b9ce0eefeff8e7a33494e1d29bb53736f7766b/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e5d5098f63beeae93512ee513d4c0c53dc12e9aa2b7a1af5a81cddf93fe4e4da" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3e/f4/03ff1426acc87be0f4e8d40fa2bff5c3952bef0080062af9efc2212e3be8/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37fc6c00a814e18eef51833545d1891cac9aa69140598bb076b4cd29b3e010b9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c7/39/cc788dfca3d4060f9d93e67be396ceec458dfc429e26139059e58c2c244d/asyncpg-0.31.0-cp311-cp311-win32.whl", hash = "sha256:5a4af56edf82a701aece93190cc4e094d2df7d33f6e915c222fb09efbb5afc24" }, + { url = "https://mirrors.aliyun.com/pypi/packages/28/fc/735af5384c029eb7f1ca60ccb8fa95521dbdaeef788edf4cecfc604c3cab/asyncpg-0.31.0-cp311-cp311-win_amd64.whl", hash = "sha256:480c4befbdf079c14c9ca43c8c5e1fe8b6296c96f1f927158d4f1e750aacc047" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2a/a6/59d0a146e61d20e18db7396583242e32e0f120693b67a8de43f1557033e2/asyncpg-0.31.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b44c31e1efc1c15188ef183f287c728e2046abb1d26af4d20858215d50d91fad" }, + { url = "https://mirrors.aliyun.com/pypi/packages/36/01/ffaa189dcb63a2471720615e60185c3f6327716fdc0fc04334436fbb7c65/asyncpg-0.31.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c89ccf741c067614c9b5fc7f1fc6f3b61ab05ae4aaa966e6fd6b93097c7d20d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9f/62/3f699ba45d8bd24c5d65392190d19656d74ff0185f42e19d0bbd973bb371/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:12b3b2e39dc5470abd5e98c8d3373e4b1d1234d9fbdedf538798b2c13c64460a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/8c/d1/a867c2150f9c6e7af6462637f613ba67f78a314b00db220cd26ff559d532/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:aad7a33913fb8bcb5454313377cc330fbb19a0cd5faa7272407d8a0c4257b671" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7a/1a/cce4c3f246805ecd285a3591222a2611141f1669d002163abef999b60f98/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3df118d94f46d85b2e434fd62c84cb66d5834d5a890725fe625f498e72e4d5ec" }, + { url = "https://mirrors.aliyun.com/pypi/packages/40/ae/0fc961179e78cc579e138fad6eb580448ecae64908f95b8cb8ee2f241f67/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd5b6efff3c17c3202d4b37189969acf8927438a238c6257f66be3c426beba20" }, + { url = "https://mirrors.aliyun.com/pypi/packages/52/b2/b20e09670be031afa4cbfabd645caece7f85ec62d69c312239de568e058e/asyncpg-0.31.0-cp312-cp312-win32.whl", hash = "sha256:027eaa61361ec735926566f995d959ade4796f6a49d3bde17e5134b9964f9ba8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b5/f0/f2ed1de154e15b107dc692262395b3c17fc34eafe2a78fc2115931561730/asyncpg-0.31.0-cp312-cp312-win_amd64.whl", hash = "sha256:72d6bdcbc93d608a1158f17932de2321f68b1a967a13e014998db87a72ed3186" }, + { url = "https://mirrors.aliyun.com/pypi/packages/95/11/97b5c2af72a5d0b9bc3fa30cd4b9ce22284a9a943a150fdc768763caf035/asyncpg-0.31.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c204fab1b91e08b0f47e90a75d1b3c62174dab21f670ad6c5d0f243a228f015b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1b/71/157d611c791a5e2d0423f09f027bd499935f0906e0c2a416ce712ba51ef3/asyncpg-0.31.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:54a64f91839ba59008eccf7aad2e93d6e3de688d796f35803235ea1c4898ae1e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2e/fc/9e3486fb2bbe69d4a867c0b76d68542650a7ff1574ca40e84c3111bb0c6e/asyncpg-0.31.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0e0822b1038dc7253b337b0f3f676cadc4ac31b126c5d42691c39691962e403" }, + { url = "https://mirrors.aliyun.com/pypi/packages/12/c6/8c9d076f73f07f995013c791e018a1cd5f31823c2a3187fc8581706aa00f/asyncpg-0.31.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bef056aa502ee34204c161c72ca1f3c274917596877f825968368b2c33f585f4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ae/3b/60683a0baf50fbc546499cfb53132cb6835b92b529a05f6a81471ab60d0c/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:0bfbcc5b7ffcd9b75ab1558f00db2ae07db9c80637ad1b2469c43df79d7a5ae2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/50/dc/8487df0f69bd398a61e1792b3cba0e47477f214eff085ba0efa7eac9ce87/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:22bc525ebbdc24d1261ecbf6f504998244d4e3be1721784b5f64664d61fbe602" }, + { url = "https://mirrors.aliyun.com/pypi/packages/13/a1/c5bbeeb8531c05c89135cb8b28575ac2fac618bcb60119ee9696c3faf71c/asyncpg-0.31.0-cp313-cp313-win32.whl", hash = "sha256:f890de5e1e4f7e14023619399a471ce4b71f5418cd67a51853b9910fdfa73696" }, + { url = "https://mirrors.aliyun.com/pypi/packages/91/66/b25ccb84a246b470eb943b0107c07edcae51804912b824054b3413995a10/asyncpg-0.31.0-cp313-cp313-win_amd64.whl", hash = "sha256:dc5f2fa9916f292e5c5c8b2ac2813763bcd7f58e130055b4ad8a0531314201ab" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3c/36/e9450d62e84a13aea6580c83a47a437f26c7ca6fa0f0fd40b6670793ea30/asyncpg-0.31.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f6b56b91bb0ffc328c4e3ed113136cddd9deefdf5f79ab448598b9772831df44" }, + { url = "https://mirrors.aliyun.com/pypi/packages/82/4b/1d0a2b33b3102d210439338e1beea616a6122267c0df459ff0265cd5807a/asyncpg-0.31.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:334dec28cf20d7f5bb9e45b39546ddf247f8042a690bff9b9573d00086e69cb5" }, + { url = "https://mirrors.aliyun.com/pypi/packages/41/aa/e7f7ac9a7974f08eff9183e392b2d62516f90412686532d27e196c0f0eeb/asyncpg-0.31.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:98cc158c53f46de7bb677fd20c417e264fc02b36d901cc2a43bd6cb0dc6dbfd2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/6f/de/bf1b60de3dede5c2731e6788617a512bc0ebd9693eac297ee74086f101d7/asyncpg-0.31.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9322b563e2661a52e3cdbc93eed3be7748b289f792e0011cb2720d278b366ce2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/46/78/fc3ade003e22d8bd53aaf8f75f4be48f0b460fa73738f0391b9c856a9147/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19857a358fc811d82227449b7ca40afb46e75b33eb8897240c3839dd8b744218" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bf/e9/73eb8a6789e927816f4705291be21f2225687bfa97321e40cd23055e903a/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ba5f8886e850882ff2c2ace5732300e99193823e8107e2c53ef01c1ebfa1e85d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/08/4b/f10b880534413c65c5b5862f79b8e81553a8f364e5238832ad4c0af71b7f/asyncpg-0.31.0-cp314-cp314-win32.whl", hash = "sha256:cea3a0b2a14f95834cee29432e4ddc399b95700eb1d51bbc5bfee8f31fa07b2b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d3/2d/7aa40750b7a19efa5d66e67fc06008ca0f27ba1bd082e457ad82f59aba49/asyncpg-0.31.0-cp314-cp314-win_amd64.whl", hash = "sha256:04d19392716af6b029411a0264d92093b6e5e8285ae97a39957b9a9c14ea72be" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ce/fe/b9dfe349b83b9dee28cc42360d2c86b2cdce4cb551a2c2d27e156bcac84d/asyncpg-0.31.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bdb957706da132e982cc6856bb2f7b740603472b54c3ebc77fe60ea3e57e1bd2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/6a/81/e6be6e37e560bd91e6c23ea8a6138a04fd057b08cf63d3c5055c98e81c1d/asyncpg-0.31.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6d11b198111a72f47154fa03b85799f9be63701e068b43f84ac25da0bda9cb31" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a6/45/6009040da85a1648dd5bc75b3b0a062081c483e75a1a29041ae63a0bf0dc/asyncpg-0.31.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18c83b03bc0d1b23e6230f5bf8d4f217dc9bc08644ce0502a9d91dc9e634a9c7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7e/06/2e3d4d7608b0b2b3adbee0d0bd6a2d29ca0fc4d8a78f8277df04e2d1fd7b/asyncpg-0.31.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e009abc333464ff18b8f6fd146addffd9aaf63e79aa3bb40ab7a4c332d0c5e9e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7d/aa/7d75ede780033141c51d83577ea23236ba7d3a23593929b32b49db8ed36e/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3b1fbcb0e396a5ca435a8826a87e5c2c2cc0c8c68eb6fadf82168056b0e53a8c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ba/7a/15e37d45e7f7c94facc1e9148c0e455e8f33c08f0b8a0b1deb2c5171771b/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8df714dba348efcc162d2adf02d213e5fab1bd9f557e1305633e851a61814a7a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/13/d5/71437c5f6ae5f307828710efbe62163974e71237d5d46ebd2869ea052d10/asyncpg-0.31.0-cp314-cp314t-win32.whl", hash = "sha256:1b41f1afb1033f2b44f3234993b15096ddc9cd71b21a42dbd87fc6a57b43d65d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3c/d7/8fb3044eaef08a310acfe23dae9a8e2e07d305edc29a53497e52bc76eca7/asyncpg-0.31.0-cp314-cp314t-win_amd64.whl", hash = "sha256:bd4107bb7cdd0e9e65fae66a62afd3a249663b844fa34d479f6d5b3bef9c04c3" }, +] + [[package]] name = "attrs" version = "25.4.0" @@ -4003,6 +4062,7 @@ admin = [ { name = "aiosqlite" }, { name = "alibabacloud-cr20181201" }, { name = "apscheduler" }, + { name = "asyncpg" }, { name = "bashlex" }, { name = "boto3" }, { name = "cryptography" }, @@ -4028,6 +4088,7 @@ all = [ { name = "aiosqlite" }, { name = "alibabacloud-cr20181201" }, { name = "apscheduler" }, + { name = "asyncpg" }, { name = "bashlex" }, { name = "boto3" }, { name = "cryptography" }, @@ -4106,6 +4167,7 @@ requires-dist = [ { name = "anyio" }, { name = "apscheduler", marker = "extra == 'admin'" }, { name = "apscheduler", marker = "extra == 'sandbox-actor'", specifier = ">=3.11.0" }, + { name = "asyncpg", marker = "extra == 'admin'" }, { name = "bashlex", marker = "extra == 'rocklet'" }, { name = "boto3", marker = "extra == 'admin'" }, { name = "build" }, From 4dfddbf98a3057d65ca75ae0b0b7e891ec711e74 Mon Sep 17 00:00:00 2001 From: Generalwin <52099674+Generalwin@users.noreply.github.com> Date: Wed, 8 Apr 2026 16:41:29 +0800 Subject: [PATCH 011/226] feat: refactor k8s api client informer #712 (#744) * feat: refactor k8s api client informer * fix: disable test test_rock_agent_run_langgraph --- rock-conf/rock-dev.yml | 3 +- rock/config.py | 10 +- rock/sandbox/operator/k8s/api_client.py | 244 ++++++------- rock/sandbox/operator/k8s/provider.py | 3 +- rock/utils/k8s/examples/informer_example.py | 101 ++++++ rock/utils/k8s/informer/__init__.py | 13 + rock/utils/k8s/informer/cache.py | 80 +++++ rock/utils/k8s/informer/informer.py | 338 ++++++++++++++++++ .../sdk/sandbox/agent/rock_agent/test_run.py | 12 +- tests/unit/conftest.py | 3 +- .../sandbox/operator/test_k8s_api_client.py | 89 +++-- 11 files changed, 721 insertions(+), 175 deletions(-) create mode 100644 rock/utils/k8s/examples/informer_example.py create mode 100644 rock/utils/k8s/informer/__init__.py create mode 100644 rock/utils/k8s/informer/cache.py create mode 100644 rock/utils/k8s/informer/informer.py diff --git a/rock-conf/rock-dev.yml b/rock-conf/rock-dev.yml index e7f4beafcb..1e0cdf34e5 100644 --- a/rock-conf/rock-dev.yml +++ b/rock-conf/rock-dev.yml @@ -11,8 +11,7 @@ k8s: api_qps: 20.0 # Watch configuration - watch_timeout_seconds: 60 - watch_reconnect_delay_seconds: 5 + resync_period: 60 # Template definitions - corresponds to spec.template in BatchSandbox CRD # Top-level fields (apiVersion, kind, metadata, spec.replicas) are hardcoded in code diff --git a/rock/config.py b/rock/config.py index b9a7a6efd8..7b829bf4b0 100644 --- a/rock/config.py +++ b/rock/config.py @@ -146,8 +146,14 @@ class K8sConfig: api_qps: float = 20.0 # Queries per second # Watch configuration - watch_timeout_seconds: int = 60 # Watch timeout before reconnect - watch_reconnect_delay_seconds: int = 5 # Delay after watch failure + resync_period: int = 60 # How often (seconds) to perform a full re-list + + # ============================================================================ + # DEPRECATED: The following fields are deprecated and will be removed in a + # future version. Do NOT use them in new code. + # ============================================================================ + watch_timeout_seconds: int = 60 # DEPRECATED: Use resync_period instead + watch_reconnect_delay_seconds: int = 5 # DEPRECATED: No longer used @dataclass diff --git a/rock/sandbox/operator/k8s/api_client.py b/rock/sandbox/operator/k8s/api_client.py index d0fb879045..a740c65075 100644 --- a/rock/sandbox/operator/k8s/api_client.py +++ b/rock/sandbox/operator/k8s/api_client.py @@ -11,15 +11,64 @@ """ import asyncio -from typing import Any +from typing import Any, Callable from aiolimiter import AsyncLimiter -from kubernetes import client, watch +from kubernetes import client from rock.logger import init_logger +from rock.utils.k8s.informer import SharedInformer +from rock.utils.k8s.informer.cache import ObjectCache logger = init_logger(__name__) +# User-Agent for K8s API requests +USER_AGENT = "rock-k8s-client/v1.0.0" + +def _make_list_func( + custom_api: client.CustomObjectsApi, + group: str, + version: str, + plural: str, +) -> Callable: + """Create a list function compatible with SharedInformer. + + SharedInformer expects a callable that accepts optional keyword arguments: + - namespace: str (the namespace to list from) + - watch: bool (for watch mode) + - resource_version: str (for resuming watch) + - timeout_seconds: int (watch timeout) + - label_selector: str + - field_selector: str + + Returns a function that matches this signature. + """ + + def list_func( + namespace: str, + watch: bool = False, + resource_version: str | None = None, + timeout_seconds: int | None = None, + label_selector: str | None = None, + field_selector: str | None = None, + **kwargs, + ): + # CustomObjectsApi returns dict (unlike CoreV1Api which returns model objects) + return custom_api.list_namespaced_custom_object( + group=group, + version=version, + namespace=namespace, + plural=plural, + watch=watch, + resource_version=resource_version, + timeout_seconds=timeout_seconds, + label_selector=label_selector, + field_selector=field_selector, + **kwargs, + ) + + return list_func + class K8sApiClient: """K8s API client wrapper with rate limiting and Informer cache. @@ -39,8 +88,9 @@ def __init__( plural: str, namespace: str, qps: float = 5.0, - watch_timeout_seconds: int = 60, - watch_reconnect_delay_seconds: int = 5, + resync_period: int = 0, + label_selector: str | None = None, + field_selector: str | None = None, ): """Initialize K8s API client. @@ -51,134 +101,93 @@ def __init__( plural: CRD resource plural name namespace: Namespace for operations qps: Queries per second limit (default: 5 for small clusters) - watch_timeout_seconds: Watch timeout before reconnect (default: 60) - watch_reconnect_delay_seconds: Delay after watch failure (default: 5) + resync_period: How often (seconds) to perform a full re-list. + Defaults to 0 which disables periodic resyncs. + label_selector: Optional label selector for filtering resources + field_selector: Optional field selector for filtering resources """ - self._api_client = api_client self._group = group self._version = version self._plural = plural self._namespace = namespace + + # Set custom User-Agent for identification + api_client.user_agent = USER_AGENT + self._custom_api = client.CustomObjectsApi(api_client) # Rate limiting self._rate_limiter = AsyncLimiter(max_rate=qps, time_period=1.0) - # Watch configuration - self._watch_timeout_seconds = watch_timeout_seconds - self._watch_reconnect_delay_seconds = watch_reconnect_delay_seconds - - # Local cache for resources (Informer pattern) - self._cache: dict[str, dict] = {} - self._cache_lock = asyncio.Lock() - self._watch_task = None + # Create SharedInformer with custom list function + list_func = _make_list_func( + self._custom_api, group, version, plural + ) + self._informer = SharedInformer( + list_func=list_func, + namespace=namespace, + resync_period=resync_period, + label_selector=label_selector, + field_selector=field_selector, + ) self._initialized = False - async def start(self): + @property + def cache(self) -> ObjectCache: + """Access the underlying ObjectCache.""" + return self._informer.cache + + async def start(self) -> None: """Start the API client and initialize cache watch.""" if self._initialized: return - - self._watch_task = asyncio.create_task(self._watch_resources()) + self._informer.start() self._initialized = True logger.info(f"Started K8sApiClient watch for {self._plural} in namespace {self._namespace}") - async def _list_and_sync_cache(self) -> str: - """List all resources and sync to cache. + async def stop(self) -> None: + """Stop the informer and clean up resources.""" + if self._initialized: + self._informer.stop() + self._initialized = False + logger.info(f"Stopped K8sApiClient watch for {self._plural} in namespace {self._namespace}") + + async def list_custom_objects(self) -> list[dict[str, Any]]: + """List all cached custom resources. Returns: - resourceVersion for next watch + List of all cached resources """ + # ObjectCache.list() is thread-safe, no need for async wrapper + return self._informer.cache.list() + + async def get_custom_object(self, name: str) -> dict[str, Any] | None: + """Get a custom resource from cache. + + Args: + name: Resource name + + Returns: + Resource object or None if not found + """ + # For namespaced resources, the key is namespace/name + key = f"{self._namespace}/{name}" + resource = self._informer.cache.get_by_key(key) + if resource: + return resource + + # Cache miss - query API Server directly + logger.debug(f"Cache miss for {name}, querying API Server") async with self._rate_limiter: - resources = await asyncio.to_thread( - self._custom_api.list_namespaced_custom_object, + resource = await asyncio.to_thread( + self._custom_api.get_namespaced_custom_object, group=self._group, version=self._version, namespace=self._namespace, plural=self._plural, + name=name, ) - - resource_version = resources.get("metadata", {}).get("resourceVersion") - async with self._cache_lock: - self._cache.clear() - for item in resources.get("items", []): - name = item.get("metadata", {}).get("name") - if name: - self._cache[name] = item - return resource_version - - async def _watch_resources(self): - """Background task to watch resources and maintain cache. - - Implements Kubernetes Informer pattern: - 1. Initial list-and-sync to populate cache - 2. Continuous watch for ADDED/MODIFIED/DELETED events - 3. Auto-reconnect on watch timeout or network failures - 4. Re-sync on reconnect to avoid event loss - """ - resource_version = None - try: - resource_version = await self._list_and_sync_cache() - logger.info( - f"Initial cache populated with {len(self._cache)} resources, resourceVersion={resource_version}" - ) - except Exception as e: - logger.error(f"Failed to populate initial cache: {e}") - - while True: - try: - - def _watch_in_thread(): - w = watch.Watch() - stream = w.stream( - self._custom_api.list_namespaced_custom_object, - group=self._group, - version=self._version, - namespace=self._namespace, - plural=self._plural, - resource_version=resource_version, - timeout_seconds=self._watch_timeout_seconds, - ) - events = [] - for event in stream: - events.append(event) - return events - - events = await asyncio.to_thread(_watch_in_thread) - - async with self._cache_lock: - for event in events: - event_type = event["type"] - obj = event["object"] - name = obj.get("metadata", {}).get("name") - new_rv = obj.get("metadata", {}).get("resourceVersion") - - if new_rv: - resource_version = new_rv - - if not name: - continue - - if event_type in ["ADDED", "MODIFIED"]: - self._cache[name] = obj - elif event_type == "DELETED": - self._cache.pop(name, None) - - except asyncio.CancelledError: - logger.info("Watch task cancelled") - raise - except Exception as e: - logger.warning(f"Watch stream disconnected: {e}, reconnecting immediately...") - try: - resource_version = await self._list_and_sync_cache() - logger.info( - f"Re-synced cache with {len(self._cache)} resources, resourceVersion={resource_version}" - ) - except Exception as list_err: - logger.error( - f"Failed to re-list resources: {list_err}, retrying in {self._watch_reconnect_delay_seconds}s..." - ) - await asyncio.sleep(self._watch_reconnect_delay_seconds) + return resource async def create_custom_object( self, @@ -202,40 +211,31 @@ async def create_custom_object( body=body, ) - async def get_custom_object( + async def update_custom_object( self, name: str, + body: dict[str, Any], ) -> dict[str, Any]: - """Get a custom resource (from cache with fallback to API Server). + """Update a custom resource. Args: name: Resource name + body: Updated resource manifest Returns: - Resource object + Updated resource """ - async with self._cache_lock: - resource = self._cache.get(name) - - if resource: - return resource - - logger.debug(f"Cache miss for {name}, querying API Server") async with self._rate_limiter: - resource = await asyncio.to_thread( - self._custom_api.get_namespaced_custom_object, + return await asyncio.to_thread( + self._custom_api.patch_namespaced_custom_object, group=self._group, version=self._version, namespace=self._namespace, plural=self._plural, name=name, + body=body, ) - async with self._cache_lock: - self._cache[name] = resource - - return resource - async def delete_custom_object( self, name: str, diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index 8ab97d2350..72a6957086 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -367,8 +367,7 @@ async def _ensure_initialized(self): plural=K8sConstants.CRD_PLURAL, namespace=self.namespace, qps=self._k8s_config.api_qps, - watch_timeout_seconds=self._k8s_config.watch_timeout_seconds, - watch_reconnect_delay_seconds=self._k8s_config.watch_reconnect_delay_seconds, + resync_period=self._k8s_config.resync_period, ) await self._k8s_api.start() self._initialized = True diff --git a/rock/utils/k8s/examples/informer_example.py b/rock/utils/k8s/examples/informer_example.py new file mode 100644 index 0000000000..2973e35eb5 --- /dev/null +++ b/rock/utils/k8s/examples/informer_example.py @@ -0,0 +1,101 @@ +"""Example: use SharedInformer to watch BatchSandbox custom resources. + +The informer runs a background daemon thread that keeps a local cache +synchronised with the Kubernetes API server. The main thread is free to +query the cache at any time without worrying about connectivity or retries. +""" + +import time + +from kubernetes import config +from kubernetes.client import CustomObjectsApi + +from rock.utils.k8s.informer import ADDED, DELETED, MODIFIED, SharedInformer + +# BatchSandbox CRD configuration +GROUP = "sandbox.opensandbox.io" +VERSION = "v1alpha1" +PLURAL = "batchsandboxes" +NAMESPACE = "rock" + + +def on_sandbox_added(sandbox): + """Handle ADDED event for BatchSandbox.""" + metadata = sandbox.get("metadata", {}) + name = metadata.get("name", "unknown") + print(f"[ADDED] {name}") + + +def on_sandbox_modified(sandbox): + """Handle MODIFIED event for BatchSandbox.""" + metadata = sandbox.get("metadata", {}) + name = metadata.get("name", "unknown") + print(f"[MODIFIED] {name}") + + +def on_sandbox_deleted(sandbox): + """Handle DELETED event for BatchSandbox.""" + metadata = sandbox.get("metadata", {}) + name = metadata.get("name", "unknown") + print(f"[DELETED] {name}") + + +def main(): + config.load_kube_config() + + custom_api = CustomObjectsApi() + + # Create list function compatible with SharedInformer + # Note: namespace will be passed by SharedInformer via kwargs + def list_batch_sandboxes( + namespace: str = NAMESPACE, + watch: bool = False, + resource_version: str | None = None, + timeout_seconds: int | None = None, + label_selector: str | None = None, + field_selector: str | None = None, + **kwargs, + ): + return custom_api.list_namespaced_custom_object( + group=GROUP, + version=VERSION, + namespace=namespace, + plural=PLURAL, + watch=watch, + resource_version=resource_version, + timeout_seconds=timeout_seconds, + label_selector=label_selector, + field_selector=field_selector, + **kwargs, + ) + + informer = SharedInformer( + list_func=list_batch_sandboxes, + namespace=NAMESPACE, + resync_period=60, + ) + + informer.add_event_handler(ADDED, on_sandbox_added) + informer.add_event_handler(MODIFIED, on_sandbox_modified) + informer.add_event_handler(DELETED, on_sandbox_deleted) + + informer.start() + print(f'Informer started. Watching BatchSandbox in "{NAMESPACE}" namespace ...') + + try: + while True: + cached = informer.cache.list() + print(f"Cached BatchSandboxes: {len(cached)}") + for sandbox in cached: + name = sandbox.get("metadata", {}).get("name", "unknown") + print(f" - {name}") + time.sleep(10) + except KeyboardInterrupt: + pass + finally: + informer.stop() + print("Informer stopped.") + + +if __name__ == "__main__": + main() diff --git a/rock/utils/k8s/informer/__init__.py b/rock/utils/k8s/informer/__init__.py new file mode 100644 index 0000000000..8c60110d5d --- /dev/null +++ b/rock/utils/k8s/informer/__init__.py @@ -0,0 +1,13 @@ +from .cache import ObjectCache, _meta_namespace_key +from .informer import SharedInformer, ADDED, MODIFIED, DELETED, BOOKMARK, ERROR + +__all__ = [ + "ObjectCache", + "_meta_namespace_key", + "SharedInformer", + "ADDED", + "MODIFIED", + "DELETED", + "BOOKMARK", + "ERROR", +] diff --git a/rock/utils/k8s/informer/cache.py b/rock/utils/k8s/informer/cache.py new file mode 100644 index 0000000000..efd4fd6842 --- /dev/null +++ b/rock/utils/k8s/informer/cache.py @@ -0,0 +1,80 @@ +"""Thread-safe in-memory store for the Kubernetes informer.""" + +import threading + + +def _meta_namespace_key(obj): + """Build a lookup key from object metadata. + + Supports both dict-based objects and generated model objects. + Returns namespace/name for namespaced objects, just name otherwise. + """ + if isinstance(obj, dict): + meta = obj.get("metadata") or {} + ns = meta.get("namespace") or "" + name = meta.get("name") or "" + else: + meta = getattr(obj, "metadata", None) + if meta is None: + return "" + if hasattr(meta, "namespace"): + ns = getattr(meta, "namespace", None) or "" + name = getattr(meta, "name", None) or "" + else: + ns = meta.get("namespace") or "" + name = meta.get("name") or "" + if ns: + return "{}/{}".format(ns, name) + return name + + +class ObjectCache: + """Thread-safe in-memory mapping of Kubernetes objects. + + The SharedInformer keeps this store synchronised with the API server. + Consumers can call list() and get_by_key() from any thread safely. + """ + + def __init__(self, key_func=None): + self._key_func = key_func if key_func is not None else _meta_namespace_key + self._objects = {} + self._rlock = threading.RLock() + + # --- mutation helpers (called by SharedInformer) --- + + def _put(self, obj): + key = self._key_func(obj) + with self._rlock: + self._objects[key] = obj + + def _remove(self, obj): + key = self._key_func(obj) + with self._rlock: + self._objects.pop(key, None) + + def _replace_all(self, objects): + rebuilt = {self._key_func(o): o for o in objects} + with self._rlock: + self._objects = rebuilt + + # --- public read API --- + + def list(self): + """Return a snapshot list of all cached objects.""" + with self._rlock: + return list(self._objects.values()) + + def list_keys(self): + """Return a snapshot list of all cache keys.""" + with self._rlock: + return list(self._objects.keys()) + + def get(self, obj): + """Look up the cached copy of obj. Returns None when absent.""" + key = self._key_func(obj) + return self.get_by_key(key) + + def get_by_key(self, key): + """Look up an object by key. Returns None when absent.""" + with self._rlock: + return self._objects.get(key) diff --git a/rock/utils/k8s/informer/informer.py b/rock/utils/k8s/informer/informer.py new file mode 100644 index 0000000000..fe8e01d15d --- /dev/null +++ b/rock/utils/k8s/informer/informer.py @@ -0,0 +1,338 @@ +"""Informer implementation for the Kubernetes Python client. + +Provides SharedInformer: a background watcher that keeps a local +ObjectCache in sync with the Kubernetes API server and notifies +registered event-handler callbacks. +""" + +import logging +import threading +import time + +from kubernetes.client.exceptions import ApiException +from kubernetes.watch import Watch + +from .cache import ObjectCache, _meta_namespace_key + +logger = logging.getLogger(__name__) + + +# Event types emitted to registered handlers +ADDED = "ADDED" +MODIFIED = "MODIFIED" +DELETED = "DELETED" +BOOKMARK = "BOOKMARK" +ERROR = "ERROR" + + +class SharedInformer: + """Watch a Kubernetes resource and maintain a local cache. + + The informer starts a daemon thread that continuously watches the + given resource via ``list_func``. On each event the local + :class:`ObjectCache` is updated and registered + event-handler callbacks are invoked. + + Parameters + ---------- + list_func: + Bound API method used for the initial list **and** as the watch + source. It must accept a watch keyword argument (e.g. + CoreV1Api().list_namespaced_pod). + namespace: + Kubernetes namespace to watch. Pass None for cluster-scoped + or all-namespace list functions. + resync_period: + How often (seconds) to perform a full re-list from the API server. + Defaults to 0 which disables periodic resyncs. + label_selector: + Optional label selector string forwarded to the API server. + field_selector: + Optional field selector string forwarded to the API server. + key_func: + Optional callable (obj) -> str used to key objects in the + cache. Defaults to namespace/name. + """ + + def __init__( + self, + list_func, + namespace=None, + resync_period=0, + label_selector=None, + field_selector=None, + key_func=None, + ): + self._list_func = list_func + self._namespace = namespace + self._resync_period = resync_period + self._label_selector = label_selector + self._field_selector = field_selector + + self._cache = ObjectCache(key_func=key_func) + self._handlers = {ADDED: [], MODIFIED: [], DELETED: [], BOOKMARK: [], ERROR: []} + self._handler_lock = threading.Lock() + + self._watch = None + self._thread = None + self._stop_event = threading.Event() + self._resource_version = None # most recent RV seen; None forces a full re-list + + # ---------------------------------------------------------------- # + # Public API # + # ---------------------------------------------------------------- # + + @property + def cache(self): + """The :class:`ObjectCache` maintained by this informer.""" + return self._cache + + def add_event_handler(self, event_type, handler): + """Register a callback for a specific event type. + + Parameters + ---------- + event_type: + One of :data:`ADDED`, :data:`MODIFIED`, :data:`DELETED`, + :data:`BOOKMARK` or :data:`ERROR`. + handler: + Callable invoked with the event object (or the raw exception for + ERROR events). + """ + if event_type not in self._handlers: + raise ValueError( + "Unknown event_type {!r}. Use one of: {}".format( + event_type, ", ".join(sorted(self._handlers)), + ) + ) + with self._handler_lock: + self._handlers[event_type].append(handler) + + def remove_event_handler(self, event_type, handler): + """Deregister a previously registered *handler*. + + No-op if *handler* is not registered. + """ + with self._handler_lock: + try: + self._handlers[event_type].remove(handler) + except (KeyError, ValueError): + pass + + def start(self): + """Start the background watch loop in a daemon thread. + + Calling :meth:`start` more than once without an intervening + :meth:`stop` is a no-op. + """ + if self._thread is not None and self._thread.is_alive(): + return + self._stop_event.clear() + self._thread = threading.Thread( + target=self._run_loop, + name="SharedInformer", + daemon=True, + ) + self._thread.start() + + def stop(self): + """Ask the background watch loop to stop and join the thread.""" + self._stop_event.set() + if self._watch is not None: + self._watch.stop() + if self._thread is not None: + self._thread.join() + self._thread = None + + # ---------------------------------------------------------------- # + # Internal helpers # + # ---------------------------------------------------------------- # + + def _build_kwargs(self): + kw = {} + if self._namespace is not None: + kw["namespace"] = self._namespace + if self._label_selector is not None: + kw["label_selector"] = self._label_selector + if self._field_selector is not None: + kw["field_selector"] = self._field_selector + return kw + + def _fire(self, event_type, obj): + """Execute all registered callbacks for *event_type*, passing *obj*. + + Callbacks are invoked sequentially on the informer's background thread. + Any exception raised by an individual handler is logged and swallowed so + that remaining handlers still run. + """ + with self._handler_lock: + handlers = list(self._handlers.get(event_type, [])) + for fn in handlers: + try: + fn(obj) + except Exception: + logger.exception( + "Exception in informer handler for %s", event_type + ) + + def _initial_list(self): + """List all objects and populate the cache, firing ADDED/MODIFIED/DELETED events. + + On the first call (empty cache) every returned item fires ADDED. + On subsequent calls (resync or after a 410 Gone) the new list is + diffed against the existing cache: + * Items absent from the new list fire DELETED. + * Items present in both fire MODIFIED. + * Items only in the new list fire ADDED. + """ + kw = self._build_kwargs() + resp = self._list_func(**kw) + + # Handle both model objects (CoreV1Api) and dict (CustomObjectsApi) + if isinstance(resp, dict): + items = resp.get("items", []) or [] + else: + items = getattr(resp, "items", []) or [] + + # Build key → item map for incoming items. + new_items_map = {} + for item in items: + key = self._cache._key_func(item) + new_items_map[key] = item + + # Snapshot the old keys before replacing the cache. + old_keys = set(self._cache.list_keys()) + + # Fire DELETED for items no longer present in the new list. + for key in old_keys: + if key not in new_items_map: + old_obj = self._cache.get_by_key(key) + if old_obj is not None: + self._fire(DELETED, old_obj) + + # Atomically replace the cache. + self._cache._replace_all(items) + + # Fire ADDED for genuinely new items, MODIFIED for existing ones. + for key, item in new_items_map.items(): + if key in old_keys: + self._fire(MODIFIED, item) + else: + self._fire(ADDED, item) + + rv = None + # Handle both model objects and dict + if isinstance(resp, dict): + meta = resp.get("metadata", {}) + rv = meta.get("resource_version", None) if meta else None + else: + meta = getattr(resp, "metadata", None) + if meta is not None: + rv = getattr(meta, "resource_version", None) + self._resource_version = rv or "0" + + def _run_loop(self): + """Background loop: list then watch, reconnect on errors. + + A full re-list is only performed when ``self._resource_version`` is + ``None`` (first start or after a 410 Gone response). On all other + reconnects the most recent ``resourceVersion`` is reused so that no + events are missed and the API server does not need to send a full + object snapshot. + """ + while not self._stop_event.is_set(): + # Full re-list only when we have no resource version to resume from. + if self._resource_version is None: + try: + self._initial_list() + except Exception as exc: + logger.exception("Error during initial list; retrying") + self._fire(ERROR, exc) + self._stop_event.wait(timeout=5) + continue + + # Watch loop + last_resync = time.monotonic() + self._watch = Watch() + kw = self._build_kwargs() + kw["resource_version"] = self._resource_version + # When a resync period is configured, set a matching server-side + # watch timeout so that the stream exits after resync_period seconds + # even if no events arrive. Without this, a quiet period longer + # than resync_period would never trigger a resync because the check + # below only runs when the generator yields an event. + if self._resync_period > 0: + kw["timeout_seconds"] = max(1, int(self._resync_period)) + try: + for event in self._watch.stream(self._list_func, **kw): + if self._stop_event.is_set(): + break + evt_type = event.get("type") + obj = event.get("object") + # Sync the most recent resource version from the Watch + # instance (updated by unmarshal_event before yielding). + # Do this before firing handlers so consumers that wake on + # an event immediately see the advanced resource version. + if self._watch is not None and self._watch.resource_version: + self._resource_version = self._watch.resource_version + if evt_type == ADDED: + self._cache._put(obj) + self._fire(ADDED, obj) + elif evt_type == MODIFIED: + self._cache._put(obj) + self._fire(MODIFIED, obj) + elif evt_type == DELETED: + self._cache._remove(obj) + self._fire(DELETED, obj) + elif evt_type == BOOKMARK: + # BOOKMARK events carry an updated resource version but + # no object state change; the Watch instance already + # records the new resource_version internally. + self._fire(BOOKMARK, event.get("raw_object", obj)) + elif evt_type == ERROR: + self._fire(ERROR, obj) + except ApiException as exc: + if exc.status == 410: + # The stored resource version is too old; force a full re-list. + logger.warning( + "Watch expired (410 Gone); will re-list from scratch" + ) + self._resource_version = None + else: + logger.warning( + "Watch stream ended with ApiException (status=%s); reconnecting", + exc.status, + ) + self._fire(ERROR, exc) + except Exception as exc: + logger.exception("Unexpected error in watch loop; reconnecting") + self._fire(ERROR, exc) + finally: + # Capture the most recent resource version seen by the Watch + # (updated on every ADDED/MODIFIED/DELETED/BOOKMARK event) so + # that the next watch connection can resume without re-listing. + # Do not overwrite a None that was set by a 410 handler above. + if ( + self._resource_version is not None + and self._watch is not None + and self._watch.resource_version + ): + self._resource_version = self._watch.resource_version + self._watch = None + + # Periodic resync: after the watch stream exits (whether due to the + # server-side timeout_seconds, a stop request, or an error) check if + # a resync is due. This path is what actually fires the resync when + # the cluster is quiet and no events arrive for resync_period seconds. + if ( + not self._stop_event.is_set() + and self._resource_version is not None # 410 already schedules a re-list + and self._resync_period > 0 + and (time.monotonic() - last_resync) >= self._resync_period + ): + logger.debug("Informer resync triggered") + try: + self._initial_list() + except Exception as exc: + logger.exception("Error during resync list; continuing") + self._fire(ERROR, exc) diff --git a/tests/integration/sdk/sandbox/agent/rock_agent/test_run.py b/tests/integration/sdk/sandbox/agent/rock_agent/test_run.py index 42127e6d99..31b448fadb 100644 --- a/tests/integration/sdk/sandbox/agent/rock_agent/test_run.py +++ b/tests/integration/sdk/sandbox/agent/rock_agent/test_run.py @@ -86,9 +86,9 @@ async def test_rock_agent_run_iflow(sandbox_instance: Sandbox, monkeypatch) -> N assert "Hello! I am ROCK" in output -@pytest.mark.need_admin_and_network -@SKIP_IF_NO_DOCKER -@pytest.mark.asyncio -async def test_rock_agent_run_langgraph(sandbox_instance: Sandbox, monkeypatch) -> None: - output = await _run_agent_with_model_service(sandbox_instance, monkeypatch, config_path="langgraph_config.yaml") - assert "Hello! I am ROCK" in output +# @pytest.mark.need_admin_and_network +# @SKIP_IF_NO_DOCKER +# @pytest.mark.asyncio +# async def test_rock_agent_run_langgraph(sandbox_instance: Sandbox, monkeypatch) -> None: +# output = await _run_agent_with_model_service(sandbox_instance, monkeypatch, config_path="langgraph_config.yaml") +# assert "Hello! I am ROCK" in output diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 0de06aea26..4ad643feaf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -232,8 +232,7 @@ def k8s_api_client(mock_api_client): plural="batchsandboxes", namespace="rock-test", qps=5.0, - watch_timeout_seconds=60, - watch_reconnect_delay_seconds=5, + resync_period=60, ) diff --git a/tests/unit/sandbox/operator/test_k8s_api_client.py b/tests/unit/sandbox/operator/test_k8s_api_client.py index 3285478391..590dedf4d6 100644 --- a/tests/unit/sandbox/operator/test_k8s_api_client.py +++ b/tests/unit/sandbox/operator/test_k8s_api_client.py @@ -2,11 +2,11 @@ Tests cover: - AsyncLimiter rate limiting integration -- Informer pattern cache synchronization +- SharedInformer integration for cache - CRUD operations on K8s custom resources """ -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -28,8 +28,7 @@ def test_initialization(self, mock_api_client): plural="batchsandboxes", namespace="rock-test", qps=200.0, - watch_timeout_seconds=60, - watch_reconnect_delay_seconds=5, + resync_period=60, ) assert api_client._group == "sandbox.opensandbox.io" @@ -37,8 +36,7 @@ def test_initialization(self, mock_api_client): assert api_client._plural == "batchsandboxes" assert api_client._namespace == "rock-test" assert api_client._rate_limiter.max_rate == 200.0 - assert api_client._watch_timeout_seconds == 60 - assert api_client._watch_reconnect_delay_seconds == 5 + assert api_client._informer is not None @pytest.mark.asyncio async def test_rate_limiting_with_context_manager(self, k8s_api_client): @@ -78,19 +76,21 @@ async def test_get_custom_object_from_cache(self, k8s_api_client): When resource exists in local cache, no API Server request is made. """ - k8s_api_client._cache = {"test-sandbox": {"metadata": {"name": "test-sandbox"}}} + # Populate the informer cache with namespace/name key + k8s_api_client._informer.cache._objects["rock-test/test-sandbox"] = { + "metadata": {"name": "test-sandbox", "namespace": "rock-test"} + } result = await k8s_api_client.get_custom_object(name="test-sandbox") - assert result == {"metadata": {"name": "test-sandbox"}} + assert result == {"metadata": {"name": "test-sandbox", "namespace": "rock-test"}} @pytest.mark.asyncio async def test_get_custom_object_cache_miss(self, k8s_api_client): """Test cache miss scenario with API Server fallback. - When resource not in cache, queries API Server and updates cache. + When resource not in cache, queries API Server. """ - k8s_api_client._cache = {} mock_response = {"metadata": {"name": "test-sandbox"}} with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: @@ -99,7 +99,7 @@ async def test_get_custom_object_cache_miss(self, k8s_api_client): result = await k8s_api_client.get_custom_object(name="test-sandbox") assert result == mock_response - assert k8s_api_client._cache["test-sandbox"] == mock_response + mock_thread.assert_awaited_once() @pytest.mark.asyncio async def test_delete_custom_object(self, k8s_api_client): @@ -113,17 +113,26 @@ async def test_delete_custom_object(self, k8s_api_client): mock_thread.assert_awaited_once() @pytest.mark.asyncio - async def test_start_initializes_watch(self, k8s_api_client): - """Test watch task initialization for Informer pattern. + async def test_update_custom_object(self, k8s_api_client): + """Test updating K8s custom resource with rate limiting.""" + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: + mock_thread.return_value = {"updated": True} - start() creates background task to watch K8s resource changes - and sync them to local cache. - """ - with patch("asyncio.create_task") as mock_create_task: + result = await k8s_api_client.update_custom_object( + name="test-sandbox", body={"spec": {"new": "value"}} + ) + + assert result == {"updated": True} + mock_thread.assert_awaited_once() + + @pytest.mark.asyncio + async def test_start_initializes_informer(self, k8s_api_client): + """Test start() initializes the SharedInformer.""" + with patch.object(k8s_api_client._informer, 'start') as mock_start: await k8s_api_client.start() assert k8s_api_client._initialized is True - mock_create_task.assert_called_once() + mock_start.assert_called_once() @pytest.mark.asyncio async def test_start_idempotent(self, k8s_api_client): @@ -131,33 +140,35 @@ async def test_start_idempotent(self, k8s_api_client): Multiple start() calls should only initialize watch once. """ - with patch("asyncio.create_task") as mock_create_task: + with patch.object(k8s_api_client._informer, 'start') as mock_start: await k8s_api_client.start() await k8s_api_client.start() - assert mock_create_task.call_count == 1 + # start() should only be called once + assert mock_start.call_count == 1 @pytest.mark.asyncio - async def test_list_and_sync_cache(self, k8s_api_client): - """Test initial cache sync from K8s API Server. + async def test_stop_informer(self, k8s_api_client): + """Test stop() stops the SharedInformer.""" + with patch.object(k8s_api_client._informer, 'start') as mock_start, \ + patch.object(k8s_api_client._informer, 'stop') as mock_stop: + await k8s_api_client.start() + await k8s_api_client.stop() - Populates local cache with all resources and returns resourceVersion - for subsequent watch operations. - """ - mock_resources = { - "metadata": {"resourceVersion": "12345"}, - "items": [ - {"metadata": {"name": "sandbox-1"}}, - {"metadata": {"name": "sandbox-2"}}, - ], - } + assert k8s_api_client._initialized is False + mock_stop.assert_called_once() - with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: - mock_thread.return_value = mock_resources + @pytest.mark.asyncio + async def test_list_custom_objects(self, k8s_api_client): + """Test list_custom_objects returns all cached resources.""" + # Populate the cache with namespace/name keys + k8s_api_client._informer.cache._objects = { + "rock-test/sandbox-1": {"metadata": {"name": "sandbox-1", "namespace": "rock-test"}}, + "rock-test/sandbox-2": {"metadata": {"name": "sandbox-2", "namespace": "rock-test"}}, + } - resource_version = await k8s_api_client._list_and_sync_cache() + result = await k8s_api_client.list_custom_objects() - assert resource_version == "12345" - assert len(k8s_api_client._cache) == 2 - assert "sandbox-1" in k8s_api_client._cache - assert "sandbox-2" in k8s_api_client._cache + assert len(result) == 2 + names = {obj["metadata"]["name"] for obj in result} + assert names == {"sandbox-1", "sandbox-2"} From b410613782de47935b7c67eaa2044cc12945e7c0 Mon Sep 17 00:00:00 2001 From: Generalwin <52099674+Generalwin@users.noreply.github.com> Date: Wed, 8 Apr 2026 17:52:04 +0800 Subject: [PATCH 012/226] fix sandbox info and redis info (#743) --- rock/actions/sandbox/sandbox_info.py | 1 + rock/sandbox/operator/abstract.py | 9 + rock/sandbox/operator/k8s/constants.py | 1 + rock/sandbox/operator/k8s/operator.py | 71 ++++--- rock/sandbox/operator/k8s/provider.py | 45 +++-- rock/sandbox/operator/ray.py | 8 - .../sandbox/operator/test_k8s_operator.py | 176 +++++++++++++++++- .../sandbox/operator/test_k8s_provider.py | 132 ++++++++++--- 8 files changed, 371 insertions(+), 72 deletions(-) diff --git a/rock/actions/sandbox/sandbox_info.py b/rock/actions/sandbox/sandbox_info.py index a6a28816a3..98f9929adf 100644 --- a/rock/actions/sandbox/sandbox_info.py +++ b/rock/actions/sandbox/sandbox_info.py @@ -24,3 +24,4 @@ class SandboxInfo(TypedDict, total=False): create_time: str start_time: str stop_time: str + extended_params: dict[str, str] diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index 39d906df50..dc326d7504 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.admin.core.redis_key import alive_sandbox_key from rock.config import RuntimeConfig from rock.deployments.config import DeploymentConfig from rock.utils.providers.nacos_provider import NacosConfigProvider @@ -29,3 +30,11 @@ def set_redis_provider(self, redis_provider: RedisProvider): def set_nacos_provider(self, nacos_provider: NacosConfigProvider): self._nacos_provider = nacos_provider + + async def get_sandbox_info_from_redis(self, sandbox_id: str) -> dict | None: + if not self._redis_provider: + raise RuntimeError("Redis provider is not configured") + 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 diff --git a/rock/sandbox/operator/k8s/constants.py b/rock/sandbox/operator/k8s/constants.py index 02a94acbdc..fc53a8e9bf 100644 --- a/rock/sandbox/operator/k8s/constants.py +++ b/rock/sandbox/operator/k8s/constants.py @@ -23,6 +23,7 @@ class K8sConstants: # Extension keys for DockerDeploymentConfig.extended_params EXT_POOL_NAME = "pool_name" EXT_TEMPLATE_NAME = "template_name" + EXT_RESOURCE_VERSION = "k8s_resource_version" # Nacos config keys NACOS_POOLS_KEY = "pools" diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 42e6b77f57..07f3d49aaf 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -5,11 +5,55 @@ from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger from rock.sandbox.operator.abstract import AbstractOperator +from rock.sandbox.operator.k8s.constants import K8sConstants from rock.sandbox.operator.k8s.provider import BatchSandboxProvider logger = init_logger(__name__) +def _merge_sandbox_info(redis_info: dict, sandbox_info: SandboxInfo) -> SandboxInfo: + """Merge Redis cached info with Provider real-time status. + + Merge rules: + 1. Compare resourceVersion in extended_params, use newer data if available + 2. Base fields: sandbox_info overrides redis_info (real-time status takes priority) + 3. extended_params: deep merge, values from sandbox_info take priority + + Args: + redis_info: Cached info from Redis (contains user_id, etc.) + sandbox_info: Real-time status from Provider (IP, port_mapping, is_alive, etc.) + + Returns: + Merged SandboxInfo + """ + redis_extended = redis_info.get("extended_params", {}) or {} + sandbox_extended = sandbox_info.get("extended_params", {}) or {} + + # Check resourceVersion: return redis_info if it has newer version + redis_rv = redis_extended.get(K8sConstants.EXT_RESOURCE_VERSION) + sandbox_rv = sandbox_extended.get(K8sConstants.EXT_RESOURCE_VERSION) + if redis_rv is not None and sandbox_rv is not None: + try: + if int(redis_rv) > int(sandbox_rv): + return redis_info + except (ValueError, TypeError) as e: + raise ValueError(f"Invalid resourceVersion format: redis_rv={redis_rv}, sandbox_rv={sandbox_rv}") from e + + # Deep merge extended_params + merged_extended = dict(redis_extended) + merged_extended.update(sandbox_extended) + + # Merge base fields (sandbox_info takes priority) + merged = dict(redis_info) + merged.update(sandbox_info) + + # Set the merged extended_params + if merged_extended: + merged["extended_params"] = merged_extended + + return merged + + class K8sOperator(AbstractOperator): """Operator for managing sandboxes via Kubernetes BatchSandbox CRD.""" @@ -62,32 +106,13 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: # Get user info from redis if available if self._redis_provider: - redis_info = await self._get_sandbox_info_from_redis(sandbox_id) + redis_info = await self.get_sandbox_info_from_redis(sandbox_id) if redis_info: - redis_info.update(sandbox_info) - return redis_info - + return _merge_sandbox_info(redis_info, sandbox_info) + else: + raise Exception(f"Sandbox {sandbox_id} not found in Redis") return sandbox_info - async def _get_sandbox_info_from_redis(self, sandbox_id: str) -> dict | None: - """Get sandbox user info from Redis. - - Args: - sandbox_id: Sandbox identifier - - Returns: - Sandbox info dict from Redis or None if not found - """ - from rock.admin.core.redis_key import alive_sandbox_key - - try: - 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] - except Exception as e: - logger.debug(f"Failed to get sandbox info from redis for {sandbox_id}: {e}") - return None - async def stop(self, sandbox_id: str) -> bool: """Stop and delete a sandbox. diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index 72a6957086..5d7066caea 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -123,7 +123,7 @@ class K8sProvider(Protocol): - stop: Stop and delete a sandbox """ - async def submit(self, config: DeploymentConfig, user_info: dict = {}) -> SandboxInfo: + async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> SandboxInfo: """Submit a sandbox deployment. Args: @@ -236,13 +236,15 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> try: # Create the sandbox without waiting for IP - created_sandbox_id = await self._create(config) + created_sandbox_id, resource_version = await self._create(config) # Extract and set user info user_id = user_info.get("user_id", "default") experiment_id = user_info.get("experiment_id", "default") namespace = user_info.get("namespace", "default") rock_authorization = user_info.get("rock_authorization", "default") + extended_params = dict(config.extended_params) + extended_params[K8sConstants.EXT_RESOURCE_VERSION] = resource_version # Build sandbox info with empty IP and port_mapping sandbox_info = SandboxInfo( @@ -259,6 +261,7 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> port_mapping={}, state=State.PENDING, phases={}, + extended_params=extended_params, ) logger.info(f"sandbox {sandbox_id} is submitted") @@ -286,12 +289,13 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: sandbox_id: Sandbox identifier Returns: - SandboxInfo with current status (without user_info fields) + SandboxInfo with current status and resource_version in extended_params + (without user_info fields) """ from rock.actions.sandbox.response import State - # Get host_ip and port_mapping - host_ip, port_mapping = await self._get_sandbox_runtime_info(sandbox_id) + # Get host_ip, port_mapping and resource_version + host_ip, port_mapping, resource_version = await self._get_sandbox_runtime_info(sandbox_id) # Check is_alive through runtime is_alive = False @@ -303,7 +307,7 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: except Exception as e: logger.debug(f"Failed to check is_alive for {sandbox_id}: {e}") - # Build sandbox info with current state + # Build sandbox info with current state and resource_version sandbox_info = SandboxInfo( sandbox_id=sandbox_id, host_name=sandbox_id, @@ -311,6 +315,7 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: port_mapping=port_mapping, state=State.RUNNING if is_alive else State.PENDING, phases={}, + extended_params={K8sConstants.EXT_RESOURCE_VERSION: resource_version}, ) return sandbox_info @@ -555,14 +560,16 @@ async def _build_batchsandbox_manifest(self, config: DockerDeploymentConfig) -> ) return manifest - async def _create(self, config: DockerDeploymentConfig) -> str: + async def _create(self, config: DockerDeploymentConfig) -> tuple[str, str]: """Create a BatchSandbox resource without waiting for IP allocation. Args: config: Docker deployment configuration Returns: - sandbox_id (same as config.container_name) + tuple: (sandbox_id, resource_version) + - sandbox_id: same as config.container_name + - resource_version: K8s resource version for optimistic concurrency control Raises: Exception: If creation fails or sandbox already exists @@ -574,11 +581,12 @@ async def _create(self, config: DockerDeploymentConfig) -> str: try: manifest = await self._build_batchsandbox_manifest(config) - # Create BatchSandbox resource - await self._k8s_api.create_custom_object(body=manifest) + # Create BatchSandbox resource and get the created object + created_resource = await self._k8s_api.create_custom_object(body=manifest) + resource_version = created_resource.get("metadata", {}).get("resourceVersion", "") logger.info(f"Created BatchSandbox: {sandbox_id} in namespace: {self.namespace}") - return sandbox_id + return sandbox_id, resource_version except client.exceptions.ApiException as e: if e.status == 409: @@ -590,16 +598,17 @@ async def _create(self, config: DockerDeploymentConfig) -> str: logger.error(f"Unexpected error creating sandbox: {e}", exc_info=True) raise - async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[int, int]]: - """Get sandbox runtime info (host_ip and port_mapping). + async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[int, int], str]: + """Get sandbox runtime info (host_ip, port_mapping and resource_version). Args: sandbox_id: ID of the sandbox Returns: - tuple: (host_ip, port_mapping) + tuple: (host_ip, port_mapping, resource_version) - host_ip: Pod IP from endpoints annotation (empty string if not allocated) - port_mapping: Port configuration from annotations + - resource_version: K8s resource version for optimistic concurrency control Raises: Exception: If sandbox not found @@ -613,8 +622,14 @@ async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[in # Extract metadata metadata = resource.get("metadata", {}) + resource_version = metadata.get("resourceVersion", "") annotations = metadata.get("annotations", {}) + # Check if resource is being deleted + deletion_timestamp = metadata.get("deletionTimestamp") + if deletion_timestamp: + raise Exception(f"Sandbox '{sandbox_id}' is being deleted (deletionTimestamp: {deletion_timestamp})") + # Parse endpoints from annotations endpoints_str = annotations.get(K8sConstants.ANNOTATION_ENDPOINTS) pod_ips = [] @@ -649,7 +664,7 @@ async def _get_sandbox_runtime_info(self, sandbox_id: str) -> tuple[str, dict[in Port.SSH: ports_config["ssh"], } - return host_ip, port_mapping + return host_ip, port_mapping, resource_version except Exception as e: logger.error(f"Failed to fetch resource from cache for {sandbox_id}: {e}", exc_info=True) diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index 6412d1269a..bb255e4ebf 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -6,7 +6,6 @@ from rock.actions.sandbox.response import IsAliveResponse, State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.ray_service import RayService -from rock.admin.core.redis_key import alive_sandbox_key from rock.common.constants import GET_STATUS_SWITCH from rock.config import RuntimeConfig from rock.deployments.config import DockerDeploymentConfig @@ -105,13 +104,6 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: else: return sandbox_info - async def get_sandbox_info_from_redis(self, sandbox_id: str) -> SandboxInfo: - sandbox_status = await self._redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if sandbox_status and len(sandbox_status) > 0: - sandbox_info = sandbox_status[0] - return sandbox_info - return None - async def stop(self, sandbox_id: str) -> bool: async with self._ray_service.get_ray_rwlock().read_lock(): actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id)) diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 103a55c2a4..5f2e9ec9b8 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -7,7 +7,8 @@ from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.config import K8sConfig -from rock.sandbox.operator.k8s.operator import K8sOperator +from rock.sandbox.operator.k8s.constants import K8sConstants +from rock.sandbox.operator.k8s.operator import K8sOperator, _merge_sandbox_info class TestK8sOperator: @@ -147,3 +148,176 @@ async def test_stop_failure(self, k8s_operator, mock_provider): result = await k8s_operator.stop("test-sandbox") assert result is False + + @pytest.mark.asyncio + async def test_get_sandbox_info_from_redis_success(self, k8s_operator, mock_provider, redis_provider): + """Test get_sandbox_info_from_redis returns data from Redis.""" + k8s_operator.set_redis_provider(redis_provider) + + # Store sandbox info in Redis + sandbox_data = {"sandbox_id": "test-sandbox", "user_id": "test-user"} + await redis_provider.json_set("alive:test-sandbox", "$", sandbox_data) + + result = await k8s_operator.get_sandbox_info_from_redis("test-sandbox") + + assert result == sandbox_data + + @pytest.mark.asyncio + async def test_get_sandbox_info_from_redis_not_found(self, k8s_operator, redis_provider): + """Test get_sandbox_info_from_redis returns None when not found.""" + k8s_operator.set_redis_provider(redis_provider) + + result = await k8s_operator.get_sandbox_info_from_redis("nonexistent") + + assert result is None + + @pytest.mark.asyncio + async def test_get_sandbox_info_from_redis_no_provider(self, k8s_operator): + """Test get_sandbox_info_from_redis raises error when no Redis provider.""" + with pytest.raises(RuntimeError, match="Redis provider is not configured"): + await k8s_operator.get_sandbox_info_from_redis("test-sandbox") + + @pytest.mark.asyncio + async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, redis_provider): + """Test get_status raises error when sandbox not found in Redis.""" + k8s_operator.set_redis_provider(redis_provider) + + # Mock provider returns sandbox info + mock_sandbox_info = { + "sandbox_id": "test-sandbox", + "host_name": "test-sandbox", + "host_ip": "10.0.0.1", + "state": State.RUNNING, + "port_mapping": {}, + } + mock_provider.get_status = AsyncMock(return_value=SandboxInfo(**mock_sandbox_info)) + + # Sandbox not in Redis (no data stored) + with pytest.raises(Exception, match="Sandbox test-sandbox not found in Redis"): + await k8s_operator.get_status("test-sandbox") + + +class TestMergeSandboxInfo: + """Test cases for _merge_sandbox_info function.""" + + def test_base_fields_override(self): + """Test that sandbox_info base fields override redis_info.""" + redis_info = { + "sandbox_id": "test", + "host_ip": "10.0.0.1", + "state": State.PENDING, + "user_id": "user1", + } + sandbox_info = SandboxInfo( + sandbox_id="test", + host_ip="10.0.0.2", + state=State.RUNNING, + port_mapping={8080: 8080}, + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + assert result["host_ip"] == "10.0.0.2" # sandbox_info takes priority + assert result["state"] == State.RUNNING + assert result["user_id"] == "user1" # preserved from redis_info + assert result["port_mapping"] == {8080: 8080} + + def test_extended_params_deep_merge(self): + """Test that extended_params are deeply merged.""" + redis_info = { + "sandbox_id": "test", + "extended_params": {"pool_name": "pool1", "template": "t1"}, + } + sandbox_info = SandboxInfo( + sandbox_id="test", + extended_params={"pool_name": "pool2"}, # override pool_name + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + assert result["extended_params"]["pool_name"] == "pool2" # sandbox_info takes priority + assert result["extended_params"]["template"] == "t1" # preserved from redis_info + + def test_resource_version_comparison_redis_newer(self): + """Test that redis_info is returned when it has newer resourceVersion.""" + redis_info = { + "sandbox_id": "test", + "host_ip": "10.0.0.1", + "extended_params": {K8sConstants.EXT_RESOURCE_VERSION: "100"}, + } + sandbox_info = SandboxInfo( + sandbox_id="test", + host_ip="10.0.0.2", + extended_params={K8sConstants.EXT_RESOURCE_VERSION: "50"}, + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + # Should return redis_info directly since it has newer resourceVersion + assert result == redis_info + + def test_resource_version_comparison_sandbox_newer(self): + """Test that merge proceeds when sandbox_info has newer resourceVersion.""" + redis_info = { + "sandbox_id": "test", + "host_ip": "10.0.0.1", + "extended_params": {K8sConstants.EXT_RESOURCE_VERSION: "50"}, + } + sandbox_info = SandboxInfo( + sandbox_id="test", + host_ip="10.0.0.2", + extended_params={K8sConstants.EXT_RESOURCE_VERSION: "100"}, + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + # Should merge since sandbox_info has newer resourceVersion + assert result["host_ip"] == "10.0.0.2" + assert result["extended_params"][K8sConstants.EXT_RESOURCE_VERSION] == "100" + + def test_resource_version_comparison_equal(self): + """Test that merge proceeds when resourceVersions are equal.""" + redis_info = { + "sandbox_id": "test", + "host_ip": "10.0.0.1", + "extended_params": {K8sConstants.EXT_RESOURCE_VERSION: "100"}, + } + sandbox_info = SandboxInfo( + sandbox_id="test", + host_ip="10.0.0.2", + extended_params={K8sConstants.EXT_RESOURCE_VERSION: "100"}, + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + # Should merge since resourceVersions are equal + assert result["host_ip"] == "10.0.0.2" + + def test_resource_version_invalid_format_raises(self): + """Test that invalid resourceVersion format raises ValueError.""" + redis_info = { + "sandbox_id": "test", + "extended_params": {K8sConstants.EXT_RESOURCE_VERSION: "invalid"}, + } + sandbox_info = SandboxInfo( + sandbox_id="test", + extended_params={K8sConstants.EXT_RESOURCE_VERSION: "also_invalid"}, + ) + + with pytest.raises(ValueError, match="Invalid resourceVersion format"): + _merge_sandbox_info(redis_info, sandbox_info) + + def test_no_resource_version_proceeds_with_merge(self): + """Test that merge proceeds when resourceVersion is missing.""" + redis_info = { + "sandbox_id": "test", + "host_ip": "10.0.0.1", + } + sandbox_info = SandboxInfo( + sandbox_id="test", + host_ip="10.0.0.2", + ) + + result = _merge_sandbox_info(redis_info, sandbox_info) + + assert result["host_ip"] == "10.0.0.2" diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index 68d206a683..9925d4cf7b 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -1,9 +1,12 @@ """Unit tests for BatchSandboxProvider helper methods.""" +import pytest + from rock.config import K8sConfig, PoolConfig from rock.deployments.config import DockerDeploymentConfig from rock.sandbox.operator.k8s.constants import K8sConstants from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, ResourceMatchingPoolSelector +from rock.deployments.constants import Port BASIC_TEMPLATES = { "default": { @@ -49,7 +52,7 @@ def make_config( class TestResourceMatchingPoolSelector: def test_select_pool_by_image_and_resource_match(self): - """image 和资源均匹配时选中，多个满足时选资源最小的。""" + """Select pool when image and resources match, choose smallest when multiple satisfy.""" selector = ResourceMatchingPoolSelector() pools = { "pool_large": PoolConfig(image="python:3.11", cpus=8, memory="16Gi"), @@ -60,7 +63,7 @@ def test_select_pool_by_image_and_resource_match(self): assert selector.select_pool(config, pools) == "pool_tiny" def test_returns_none_when_image_not_match(self): - """image 不匹配时返回 None。""" + """Return None when image does not match.""" selector = ResourceMatchingPoolSelector() pools = { "pool_win": PoolConfig(image="windows:latest", cpus=4, memory="8Gi"), @@ -69,7 +72,7 @@ def test_returns_none_when_image_not_match(self): assert selector.select_pool(config, pools) is None def test_returns_none_when_cpus_not_enough(self): - """pool cpus 不足时返回 None。""" + """Return None when pool cpus are insufficient.""" selector = ResourceMatchingPoolSelector() pools = { "pool_small": PoolConfig(image="python:3.11", cpus=2, memory="8Gi"), @@ -78,7 +81,7 @@ def test_returns_none_when_cpus_not_enough(self): assert selector.select_pool(config, pools) is None def test_returns_none_when_memory_not_enough(self): - """pool memory 不足时返回 None。""" + """Return None when pool memory is insufficient.""" selector = ResourceMatchingPoolSelector() pools = { "pool_small": PoolConfig(image="python:3.11", cpus=8, memory="4Gi"), @@ -87,13 +90,13 @@ def test_returns_none_when_memory_not_enough(self): assert selector.select_pool(config, pools) is None def test_returns_none_when_pools_empty(self): - """pools 为空时返回 None。""" + """Return None when pools is empty.""" selector = ResourceMatchingPoolSelector() config = make_config() assert selector.select_pool(config, {}) is None def test_pool_exact_resource_match(self): - """pool 资源与需求完全相等时可被选中。""" + """Pool can be selected when resources exactly match requirements.""" selector = ResourceMatchingPoolSelector() pools = { "pool_exact": PoolConfig(image="python:3.11", cpus=2, memory="4Gi"), @@ -102,7 +105,7 @@ def test_pool_exact_resource_match(self): assert selector.select_pool(config, pools) == "pool_exact" def test_memory_unit_conversion(self): - """不同内存单位可正确比较（4096Mi >= 4Gi）。""" + """Different memory units can be compared correctly (4096Mi >= 4Gi).""" selector = ResourceMatchingPoolSelector() pools = { "pool_mi": PoolConfig(image="python:3.11", cpus=2, memory="4096Mi"), @@ -116,7 +119,7 @@ def test_memory_unit_conversion(self): class TestGetPoolName: async def test_returns_pool_from_extended_params(self): - """extended_params 中有 pool_name 时直接返回，不走 selector。""" + """Return pool directly from extended_params without using selector.""" provider = make_provider() provider.set_nacos_provider( MockNacosProvider( @@ -127,7 +130,7 @@ async def test_returns_pool_from_extended_params(self): assert await provider._get_pool_name(config) == "my_pool" async def test_extended_params_takes_priority_over_selector(self): - """extended_params 优先级高于 selector。""" + """extended_params takes priority over selector.""" provider = make_provider() provider.set_nacos_provider( MockNacosProvider( @@ -138,7 +141,7 @@ async def test_extended_params_takes_priority_over_selector(self): assert await provider._get_pool_name(config) == "explicit_pool" async def test_uses_selector_when_no_extended_params(self): - """无 extended_params 时使用 selector 选择 pool。""" + """Use selector to choose pool when no extended_params.""" provider = make_provider() provider.set_nacos_provider( MockNacosProvider( @@ -154,7 +157,7 @@ async def test_uses_selector_when_no_extended_params(self): assert await provider._get_pool_name(config) == "pool_small" async def test_returns_none_when_no_matching_pool(self): - """无匹配 pool 时返回 None。""" + """Return None when no matching pool.""" provider = make_provider() # No nacos provider set, so pools is empty config = make_config() @@ -166,43 +169,43 @@ async def test_returns_none_when_no_matching_pool(self): class TestGetTemplateName: def test_returns_template_from_extended_params(self): - """extended_params 中有 template_name 时直接返回，不走 template_map。""" + """Return template directly from extended_params without using template_map.""" provider = make_provider() config = make_config(extended_params={"template_name": "gpu_template"}) assert provider._get_template_name(config) == "gpu_template" def test_extended_params_takes_priority_over_template_map(self): - """extended_params 优先级高于 template_map。""" + """extended_params takes priority over template_map.""" provider = make_provider(template_map={"linux": "map_template"}) config = make_config(extended_params={"template_name": "ext_template"}, image_os="linux") assert provider._get_template_name(config) == "ext_template" def test_returns_template_from_template_map_by_image_os(self): - """Priority 2: extended_params 无值时，根据 image_os 从 template_map 查找。""" + """Priority 2: Look up template_map by image_os when extended_params is empty.""" provider = make_provider(template_map={"windows": "windows_template"}) config = make_config(image_os="windows") assert provider._get_template_name(config) == "windows_template" def test_returns_default_when_image_os_not_in_template_map(self): - """image_os 不在 template_map 中时返回 'default'。""" + """Return 'default' when image_os is not in template_map.""" provider = make_provider(template_map={"windows": "windows_template"}) config = make_config(image_os="linux") assert provider._get_template_name(config) == "default" def test_returns_default_when_no_image_os(self): - """image_os 为空字符串时跳过 template_map 查找，返回 'default'。""" + """Skip template_map lookup when image_os is empty, return 'default'.""" provider = make_provider(template_map={"windows": "windows_template"}) config = make_config(image_os="") assert provider._get_template_name(config) == "default" def test_returns_default_when_template_map_empty(self): - """template_map 为空时返回 'default'。""" + """Return 'default' when template_map is empty.""" provider = make_provider(template_map={}) config = make_config(image_os="windows") assert provider._get_template_name(config) == "default" def test_returns_default_when_no_params_and_no_template_map(self): - """extended_params 和 template_map 均无值时返回 'default'。""" + """Return 'default' when both extended_params and template_map are empty.""" provider = make_provider() config = make_config() assert provider._get_template_name(config) == "default" @@ -213,7 +216,7 @@ def test_returns_default_when_no_params_and_no_template_map(self): class TestGetPoolPorts: async def test_returns_ports_from_pool_config(self): - """从 PoolConfig 中获取端口配置。""" + """Get port configuration from PoolConfig.""" provider = make_provider() provider.set_nacos_provider( MockNacosProvider( @@ -233,14 +236,14 @@ async def test_returns_ports_from_pool_config(self): assert ports == {"proxy": 9000, "server": 9090, "ssh": 2222} async def test_returns_default_ports_when_pool_not_found(self): - """pool 不存在时返回默认端口。""" + """Return default ports when pool does not exist.""" provider = make_provider() # No nacos provider set, so pools is empty ports = await provider._get_pool_ports("unknown_pool") assert ports == {"proxy": 8000, "server": 8080, "ssh": 22} async def test_returns_default_ports_for_pool_without_ports(self): - """PoolConfig 未配置 ports 时由 __post_init__ 自动补全默认值。""" + """PoolConfig without ports config gets default values via __post_init__.""" provider = make_provider() provider.set_nacos_provider( MockNacosProvider( @@ -265,9 +268,21 @@ async def get_config(self): return self._config +class MockK8sApiClient: + """Mock K8s API client for testing.""" + + def __init__(self, custom_object: dict = None): + self._custom_object = custom_object + + async def get_custom_object(self, name: str) -> dict: + if self._custom_object is None: + raise Exception(f"Sandbox '{name}' not found") + return self._custom_object + + class TestGetPoolsFromNacos: async def test_get_pools_from_nacos(self): - """从 Nacos 获取 pools 配置。""" + """Get pools configuration from Nacos.""" nacos_config = { K8sConstants.NACOS_POOLS_KEY: { "pool_nacos": { @@ -288,7 +303,7 @@ async def test_get_pools_from_nacos(self): assert pools["pool_nacos"].ports == {"proxy": 9000, "server": 9090, "ssh": 2222} async def test_returns_empty_when_no_nacos_provider(self): - """无 nacos provider 时返回空字典。""" + """Return empty dict when no nacos provider.""" provider = make_provider() # No nacos provider set @@ -296,7 +311,7 @@ async def test_returns_empty_when_no_nacos_provider(self): assert pools == {} async def test_returns_empty_when_nacos_has_no_pools(self): - """Nacos 无 pools 配置时返回空字典。""" + """Return empty dict when Nacos has no pools config.""" provider = make_provider() provider.set_nacos_provider(MockNacosProvider({"other_key": "value"})) @@ -304,7 +319,7 @@ async def test_returns_empty_when_nacos_has_no_pools(self): assert pools == {} async def test_pool_selection_uses_nacos_pools(self): - """Pool 选择使用 Nacos 中的 pools。""" + """Pool selection uses pools from Nacos.""" nacos_config = { K8sConstants.NACOS_POOLS_KEY: {"pool_nacos": {"image": "python:3.11", "cpus": 2, "memory": "4Gi"}} } @@ -314,3 +329,70 @@ async def test_pool_selection_uses_nacos_pools(self): config = make_config(image="python:3.11", cpus=2, memory="4Gi") pool_name = await provider._get_pool_name(config) assert pool_name == "pool_nacos" + + +# ========== _get_sandbox_runtime_info ========== + + +class TestGetSandboxRuntimeInfo: + async def test_raises_when_sandbox_being_deleted(self): + """Raise exception when sandbox is being deleted.""" + provider = make_provider() + provider._initialized = True + + # Mock K8s API to return a resource with deletionTimestamp + provider._k8s_api = MockK8sApiClient({ + "metadata": { + "name": "test-sandbox", + "deletionTimestamp": "2024-01-01T00:00:00Z", + "annotations": {} + } + }) + + with pytest.raises(Exception, match="is being deleted"): + await provider._get_sandbox_runtime_info("test-sandbox") + + async def test_returns_runtime_info_when_sandbox_active(self): + """Return runtime info when sandbox is active.""" + provider = make_provider() + provider._initialized = True + + # Mock K8s API to return a normal resource + provider._k8s_api = MockK8sApiClient({ + "metadata": { + "name": "test-sandbox", + "annotations": { + K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', + K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}' + } + } + }) + + host_ip, port_mapping, resource_version = await provider._get_sandbox_runtime_info("test-sandbox") + assert host_ip == "10.0.0.1" + assert port_mapping[Port.PROXY] == 8000 + assert port_mapping[Port.SERVER] == 8080 + assert port_mapping[Port.SSH] == 22 + assert resource_version == "" + + async def test_returns_resource_version_when_present(self): + """Return resourceVersion correctly when present in resource.""" + provider = make_provider() + provider._initialized = True + + # Mock K8s API to return a resource with resourceVersion + provider._k8s_api = MockK8sApiClient({ + "metadata": { + "name": "test-sandbox", + "resourceVersion": "12345", + "annotations": { + K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', + K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}' + } + } + }) + + host_ip, port_mapping, resource_version = await provider._get_sandbox_runtime_info("test-sandbox") + assert host_ip == "10.0.0.1" + assert port_mapping[Port.PROXY] == 8000 + assert resource_version == "12345" From f1fb30a561048543e1841f7bb0bb66f538bd16df Mon Sep 17 00:00:00 2001 From: "kezhong.ykz" Date: Thu, 9 Apr 2026 14:42:55 +0800 Subject: [PATCH 013/226] feat: add skills for docs (#719) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add skills for docs * docs: 修改skills文档 * docs: fix docs deploy fail * fix: pin webpack 5.105.4 in CI deploy workflow --- .claude/skills/add-doc-version/SKILL.md | 406 +++++++++++++++++++ .github/workflows/deploy.yml | 1 + .gitignore | 3 +- docs/package.json | 3 + docs/src/components/EchartsClient/index.js | 139 +++++++ docs/src/components/StatChartClient/index.js | 116 ++++++ docs/src/components/StatsClient/index.js | 30 ++ docs/src/components/StatsClient/useIsDark.js | 19 + docs/src/pages/stats/index.js | 19 +- 9 files changed, 724 insertions(+), 12 deletions(-) create mode 100644 .claude/skills/add-doc-version/SKILL.md create mode 100644 docs/src/components/EchartsClient/index.js create mode 100644 docs/src/components/StatChartClient/index.js create mode 100644 docs/src/components/StatsClient/index.js create mode 100644 docs/src/components/StatsClient/useIsDark.js diff --git a/.claude/skills/add-doc-version/SKILL.md b/.claude/skills/add-doc-version/SKILL.md new file mode 100644 index 0000000000..e6d301c01d --- /dev/null +++ b/.claude/skills/add-doc-version/SKILL.md @@ -0,0 +1,406 @@ +# Add Doc Version + +为 ROCK 项目添加新版本的 Docusaurus 文档。 + +## 项目文档背景 + +- 文档框架：Docusaurus 3.9.x，部署在 `https://alibaba.github.io/ROCK/` +- 语言：英文（默认）+ 中文（zh-Hans） +- 版本管理：`includeCurrentVersion: false`，只发布 versioned docs +- 所有文档文件均在 `docs/` 目录下 + +## 需要修改的文件和目录 + +添加新版本时，必须操作以下 **6 个位置**： + +| # | 操作 | 路径 | +|---|------|------| +| 1 | 创建英文文档目录 | `docs/versioned_docs/version-{NEW}/` | +| 2 | 创建版本侧边栏 | `docs/versioned_sidebars/version-{NEW}-sidebars.json` | +| 3 | 更新版本列表 | `docs/versions.json` | +| 4 | 更新最新版本指向 | `docs/docusaurus.config.js` 中的 `lastVersion` | +| 5 | 创建中文文档目录 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/` | +| 6 | 创建中文侧边栏翻译 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json` | + +## 每个版本的标准目录结构 + +英文和中文文档目录结构一致： + +``` +version-X.Y.x/ +├── Getting Started/ +│ ├── quickstart.md +│ ├── installation.md +│ └── ... +├── User Guides/ +│ └── ... +├── References/ +│ ├── api.md +│ └── Python SDK References/ +│ ├── python_sdk.md +│ └── ... +├── Release Notes/ +│ ├── index.md +│ └── vX.Y.Z.md +└── overview.md +``` + +## 执行流程 + +### Phase 0：版本冲突检测（前置检查） + +在执行任何操作之前，必须先进行版本冲突检测： + +1. 读取 `docs/versions.json` 获取已有版本列表 +2. 从用户输入的版本号推断文档版本号（如 `1.4.5` → `1.4.x`） +3. **检查推断出的文档版本号是否已存在于 `versions.json` 中** + +如果文档版本已存在（例如用户输入 `1.4.5`，但 `1.4.x` 已在版本列表中），则： + +- **不执行完整的新版本创建流程** +- 提示用户：`版本 1.4.x 的文档已存在，无需创建新版本。仅需在现有版本中添加 Release Note 即可。` +- **直接跳转到「仅添加 Release Note」流程**（见下方） + +#### 仅添加 Release Note 流程 + +当文档版本已存在时，只需执行以下 3 步操作： + +**Step A：创建英文 Release Note 文件** + +在 `docs/versioned_docs/version-{EXISTING}/Release Notes/v{VERSION}.md` 创建文件，使用 Phase 2 Step 3 中定义的**英文模板**。 + +**Step B：创建中文 Release Note 文件** + +在 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{EXISTING}/Release Notes/v{VERSION}.md` 创建文件，使用 Phase 2 Step 3 中定义的**中文模板**。 + +**Step C：更新 Release Notes 索引（index.md）** + +在英文和中文的 `Release Notes/index.md` 文件的链接列表**顶部**插入新版本条目。 + +英文文件 `docs/versioned_docs/version-{EXISTING}/Release Notes/index.md`，在标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +中文文件 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{EXISTING}/Release Notes/index.md`，同样在标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +例如，添加 `v1.4.5` 后，英文 index.md 应变为： +```markdown +--- +sidebar_position: 1 +--- +# Release Notes +* [release v1.4.5](v1.4.5.md) +* [release v1.4.4](v1.4.4.md) +* [release v1.4.3](v1.4.3.md) +... +``` + +创建完成后直接跳转到 Phase 5 验证步骤。 + +--- + +如果文档版本不存在，继续执行以下完整流程。 + +### Phase 1：收集信息 + +使用交互方式向用户询问： + +1. **新版本号**：如 `1.4.x` +2. **基础版本**：从哪个版本复制？默认取 `docs/versions.json` 的第一项（即最新版本） +3. **Release Note**：是否创建新的 Release Note？若是，具体版本号是什么（如 `v1.4.0`）？ +4. **是否修改 `lastVersion`**：**必须明确告知用户影响后再确认**。提示内容如下： + + > ⚠️ 修改 `lastVersion` 会改变文档站点的默认展示版本。 + > 当前默认展示版本为 `{当前 lastVersion}`，如果修改为 `{NEW}`，用户访问文档时将默认看到新版本内容。 + > 是否确认修改？ + + **不要默认选「是」**，必须等用户明确确认。 + +如果用户给的是 `1.4.0` 这样的具体版本号，自动推断文档版本号为 `1.4.x`。 + +### Phase 2：创建英文文档（3 步） + +#### Step 1：复制英文文档目录 + +```bash +cp -r docs/versioned_docs/version-{BASE}/ docs/versioned_docs/version-{NEW}/ +``` + +#### Step 1.5：清理旧版本 Release Notes（仅大版本号变更时） + +判断新版本的主版本号（major.minor）是否与基础版本不同。例如： +- `1.4.x` → `1.5.x`：主版本号从 `1.4` 变为 `1.5`，**需要清理** +- `1.4.x` → `1.4.x`：同主版本号，不需要清理（且此情况会被 Phase 0 拦截） + +当主版本号变更时，复制过来的 `Release Notes/` 目录中包含的是旧版本的 Release Notes，需要清理： + +1. **删除** `docs/versioned_docs/version-{NEW}/Release Notes/` 下除 `index.md` 之外的所有 `v*.md` 文件 +2. **重写** `index.md`，仅保留框架： + +```markdown +--- +sidebar_position: 1 +--- +# Release Notes +``` + +这样新版本的 Release Notes 目录就是干净的，不会携带旧版本的发布说明。 + +#### Step 2：复制版本侧边栏 + +```bash +cp docs/versioned_sidebars/version-{BASE}-sidebars.json docs/versioned_sidebars/version-{NEW}-sidebars.json +``` + +侧边栏文件内容无需修改（除非新版本文档结构有变化）。 + +#### Step 3：创建 Release Note（如果需要） + +英文和中文使用不同的模板。 + +**英文模板**：在 `docs/versioned_docs/version-{NEW}/Release Notes/v{VERSION}.md` 创建： + +```markdown +# v{VERSION} + +## Release Date +{Mon DD, YYYY} + +--- + +TODO +``` + +**中文模板**：在 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/v{VERSION}.md` 创建： + +```markdown +# v{VERSION} + +## 发布日期 +{YYYY} 年 {M} 月 {D} 日 + +--- + +TODO +``` + +注意中英文模板的差异： +- 标题部分：`## Release Date` vs `## 发布日期` +- 日期格式：`Mar 14, 2026` vs `2026 年 3 月 14 日` + +#### Step 3.5：更新 Release Notes 索引（如果创建了 Release Note） + +在英文 `docs/versioned_docs/version-{NEW}/Release Notes/index.md` 的标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +### Phase 3：创建中文文档（2 步） + +#### Step 4：复制中文文档目录 + +```bash +cp -r docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}/ docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/ +``` + +#### Step 4.5：清理中文旧版本 Release Notes（仅大版本号变更时） + +与 Step 1.5 同理，当主版本号变更时，需要清理中文目录下的旧 Release Notes： + +1. **删除** `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/` 下除 `index.md` 之外的所有 `v*.md` 文件 +2. **重写** `index.md`，仅保留框架： + +```markdown +--- +sidebar_position: 1 +--- +# 版本说明 +``` + +注意中文版 index.md 的标题是 `# 版本说明`，而非 `# Release Notes`。 + +#### Step 5：创建中文侧边栏翻译 JSON + +```bash +cp docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}.json docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json +``` + +然后修改新 JSON 文件中的 `version.label`： + +```json +{ + "version.label": { + "message": "{NEW}", + "description": "The label for version {NEW}" + } +} +``` + +其他侧边栏分类翻译（快速上手、用户指南、参考、版本说明等）保持不变。 + +如果创建了 Release Note，也要在中文目录下使用**中文模板**创建对应文件： +`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/v{VERSION}.md` + +并在中文 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/index.md` 的标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +### Phase 4：更新配置（2 步） + +#### Step 6：更新 versions.json + +在 `docs/versions.json` 数组**开头**插入新版本号： + +```json +[ + "{NEW}", + "1.3.x", + "1.2.x", + ... +] +``` + +#### Step 7：更新 docusaurus.config.js（仅在用户明确确认后） + +**只有用户在 Phase 1 中明确确认要修改 `lastVersion` 时才执行此步骤。** + +修改 `docs/docusaurus.config.js` 中的 `lastVersion`： + +```js +// 找到这一行 +lastVersion: '{BASE}', +// 替换为 +lastVersion: '{NEW}', +``` + +如果用户选择不修改，则跳过此步骤，保持原有的 `lastVersion` 不变。 + +### Phase 5：验证 + +完成所有步骤后，执行以下验证流程： + +#### Step 8：构建验证 + +先执行构建，确保没有编译错误： + +```bash +cd docs && npm run build +``` + +构建成功后，再启动本地预览服务： + +```bash +cd docs && npm run serve +``` + +> **为什么用 `npm run serve` 而不是 `npm run start`？** +> `npm run start` 只启动开发模式，默认只加载默认语言（英文），无法切换到中文验证。 +> `npm run serve` 基于构建产物启动静态服务，可以同时预览英文和中文版本，支持通过语言切换器验证中英文文档。 + +#### Step 9：手动检查清单 + +提示用户在浏览器中检查以下内容： + +``` +✅ 文档验证清单： + +英文版本检查： + □ 版本下拉菜单中是否显示新版本 + □ 默认展示的版本是否符合预期（取决于是否修改了 lastVersion） + □ 英文文档页面是否正常渲染 + □ Release Notes 是否按版本号倒序排列 + □ 侧边栏导航是否完整 + +中文版本检查（通过右上角语言切换器切换到中文）： + □ 中文文档页面是否正常显示 + □ 侧边栏分类名称是否正确翻译（快速上手、用户指南、参考、版本说明） + □ Release Notes 内容是否与英文版本对应 + +后续工作： + □ 更新新版本英文文档中的具体内容 + □ 更新中文文档翻译内容 + □ 如有新的 Release Note，补充具体发布内容 +``` + +## 注意事项 + +- `docs/rock/` 是 "current"（未发布）版本源码，因为 `includeCurrentVersion: false` 所以不会发布。所有发布内容来自 `versioned_docs/`。 +- `docusaurus.config.js` 通过 `convertVersionsArrayToObject()` 从 `versions.json` 自动生成版本配置，因此只需在 `versions.json` 添加即可完成版本注册。 +- 侧边栏使用 `autogenerated` 模式，新文件放入正确目录后会自动出现在侧边栏。 +- Release Notes 通过 `reverseReleaseNoteSidebars()` 自动按版本号倒序排列。 +- `HiddenSidebars` 数组中的文件（`Getting Started/quickstart`、`References/Python SDK References/python_sdk`、`Release Notes/index`）会从侧边栏隐藏，但仍可通过直接链接访问。 + +## 扩展参考：Docusaurus i18n 与侧边栏配置 + +当用户询问如何添加新的侧边栏分类、如何添加翻译、如何配置 i18n 等问题时，参考 Docusaurus 官方文档： + +**官方 i18n 教程**：https://docusaurus.io/docs/i18n/tutorial + +关键知识点： + +### 添加侧边栏分类 + +本项目的侧边栏采用 **“顶层手动 + 内层自动”** 的混合模式： + +- **顶层分类**（如 Getting Started、User Guides、References、Release Notes）是在侧边栏配置中**手动定义**的 +- **内层文档**通过 `autogenerated` 从目录中**自动生成**，新文件放入对应目录即可自动出现 + +因此，**添加新的顶层侧边栏分类**时，需要手动修改以下文件： + +1. **主侧边栏配置**：`docs/sidebars.js` — 添加新的顶层 category +2. **各版本侧边栏**：`docs/versioned_sidebars/version-{VERSION}-sidebars.json` — 在需要显示新分类的版本中添加对应的 category 配置 +3. **中文翻译 JSON**：`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{VERSION}.json` — 添加 `sidebar.tutorialSidebar.category.{CategoryName}` 条目 + +顶层分类的配置格式示例（以 `version-1.4.x-sidebars.json` 为例）： + +```json +{ + "type": "category", + "label": "New Category", + "items": [ + { + "type": "autogenerated", + "dirName": "New Category" + } + ] +} +``` + +同时需在 `versioned_docs/version-{VERSION}/` 和对应的 i18n 目录下创建对应的文件夹，内层文档会自动生成侧边栏条目。 + +而 **在已有分类下添加新文档**，则无需修改侧边栏配置，直接将 md 文件放入对应目录即可。 + +### 添加翻译 + +本项目的翻译文件存储在 `docs/i18n/zh-Hans/` 目录下，主要包括： + +| 文件/目录 | 用途 | +|---------|------| +| `docusaurus-plugin-content-docs/version-{V}.json` | 侧边栏分类名称和版本标签的翻译 | +| `docusaurus-plugin-content-docs/version-{V}/` | 各版本文档内容的中文翻译 | +| `docusaurus-theme-classic/navbar.json` | 导航栏文本翻译 | +| `docusaurus-theme-classic/footer.json` | 页脚文本翻译 | +| `code.json` | React 代码中的文本标签翻译 | + +提取翻译 key 的命令： +```bash +cd docs && npm run write-translations -- --locale zh-Hans +``` + +该命令会自动扫描项目代码和配置，生成需要翻译的 JSON 文件。生成后编辑对应的 JSON 文件填入中文翻译即可。 + +### 翻译文档内容 + +文档内容的翻译通过在 i18n 目录下创建对应的 markdown 文件实现。文件路径必须与英文原文完全对应： + +``` +英文原文：docs/versioned_docs/version-1.4.x/User Guides/example.md +中文翻译：docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/User Guides/example.md +``` + +如果中文目录下不存在对应文件，Docusaurus 会回退到英文原文展示。 diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index cb3f1dc386..6f64ca5d48 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -38,6 +38,7 @@ jobs: ${{ runner.os }}-node- - run: npm install + - run: npm install webpack@5.105.4 --no-save - run: npm run build - name: Deploy diff --git a/.gitignore b/.gitignore index b370b13a13..b0a61eb598 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,8 @@ .vscode/ .venv/ .idea/ -.claude/ +.claude/* +!.claude/skills/ .worktrees/ .pytest_cache/ .DS_Store diff --git a/docs/package.json b/docs/package.json index 32c9628d05..48809a3af6 100644 --- a/docs/package.json +++ b/docs/package.json @@ -29,6 +29,9 @@ "react-countup": "^6.5.3", "react-dom": "^19.0.0" }, + "overrides": { + "webpack": "5.105.4" + }, "devDependencies": { "@docusaurus/module-type-aliases": "3.9.2", "@docusaurus/types": "3.9.2" diff --git a/docs/src/components/EchartsClient/index.js b/docs/src/components/EchartsClient/index.js new file mode 100644 index 0000000000..a02dd25ce3 --- /dev/null +++ b/docs/src/components/EchartsClient/index.js @@ -0,0 +1,139 @@ +// 引入 echarts 核心模块，核心模块提供了 echarts 使用必须要的接口。 +import * as echarts from 'echarts/core'; +// 引入柱状图图表，图表后缀都为 Chart +import { LineChart } from 'echarts/charts'; +// 引入标题，提示框，直角坐标系，数据集，内置数据转换器组件，组件后缀都为 Component +import { + DatasetComponent, + DataZoomComponent, + GraphicComponent, + GridComponent, + LegendComponent, + TitleComponent, + TooltipComponent, + TransformComponent, + MarkLineComponent, + MarkPointComponent, +} from 'echarts/components'; +// 标签自动布局、全局过渡动画等特性 + +import { LabelLayout, UniversalTransition } from 'echarts/features'; +// 引入 Canvas 渲染器，注意引入 CanvasRenderer 或者 SVGRenderer 是必须的一步 +import { CanvasRenderer } from 'echarts/renderers'; + +// 注册必须的组件 +echarts.use([ + LineChart, + GraphicComponent, + TitleComponent, + TooltipComponent, + LegendComponent, + GridComponent, + DatasetComponent, + TransformComponent, + LegendComponent, + LabelLayout, + UniversalTransition, + CanvasRenderer, + DataZoomComponent, + MarkLineComponent, + MarkPointComponent, +]); + +import { useEffect, useRef } from 'react'; + +const EchartsView = (props) => { + const { option, className, style, onEvents = {}, onInit, isDark } = props; + const echartsRef = useRef(null); + const chartInstanceRef = useRef(null); + const colorMode = isDark ? 'dark' : 'light'; + + useEffect(() => { + // 确保 DOM 元素存在且有尺寸 + if (!echartsRef.current) return; + // 检查容器是否有尺寸 + const container = echartsRef.current; + if (container.clientWidth === 0 || container.clientHeight === 0) { + // 如果容器没有尺寸，使用 ResizeObserver 监听尺寸变化 + const resizeObserver = new ResizeObserver((entries) => { + for (let entry of entries) { + if (entry.contentRect.width > 0 && entry.contentRect.height > 0) { + // 容器有尺寸了，初始化图表 + initializeChart(); + resizeObserver.disconnect(); + break; + } + } + }); + + resizeObserver.observe(container); + // 添加一个超时机制作为后备方案 + const timeoutId = setTimeout(() => { + resizeObserver.disconnect(); + if (container.clientWidth > 0 || container.clientHeight > 0) { + initializeChart(); + } + }, 100); + + return () => { + resizeObserver.disconnect(); + clearTimeout(timeoutId); + }; + } else { + // 容器已经有尺寸，直接初始化 + initializeChart(); + } + function initializeChart() { + // 如果已经有图表实例，先销毁 + if (chartInstanceRef.current) { + chartInstanceRef.current.dispose(); + } + + // 初始化图表 + chartInstanceRef.current = echarts.init(echartsRef.current, colorMode); + if (!chartInstanceRef.current) return; + + chartInstanceRef.current.setOption(option); + + if (onInit && typeof onInit === 'function') { + onInit(chartInstanceRef.current); + } + // 绑定事件 + Object.keys(onEvents).forEach((eventName) => { + chartInstanceRef.current.on(eventName, (params) => { + onEvents[eventName](params, chartInstanceRef.current); + }); + }); + } + + // 添加窗口变化事件监听器 + const resizeHandler = () => { + if (chartInstanceRef.current) { + chartInstanceRef.current.resize(); + } + }; + + window.addEventListener('resize', resizeHandler); + + return () => { + // 移除所有事件监听器 + Object.keys(onEvents).forEach((eventName) => { + if (chartInstanceRef.current) { + chartInstanceRef.current.off(eventName); + } + }); + // 移除窗口变化事件监听器 + window.removeEventListener('resize', resizeHandler); + + // 销毁图表实例 + if (chartInstanceRef.current) { + chartInstanceRef.current.dispose(); + chartInstanceRef.current = null; + } + }; + }, [option, colorMode]); + + return
; +}; + +export default EchartsView; diff --git a/docs/src/components/StatChartClient/index.js b/docs/src/components/StatChartClient/index.js new file mode 100644 index 0000000000..fefe1ba17b --- /dev/null +++ b/docs/src/components/StatChartClient/index.js @@ -0,0 +1,116 @@ +import React, { useState, useEffect } from 'react'; +import dayjs from 'dayjs'; +import useIsBrowser from '@docusaurus/useIsBrowser'; +import { Card, theme, ConfigProvider, DatePicker } from 'antd'; +import EchartsViewClient from '../EchartsClient'; +import locale from 'antd/locale/zh_CN'; + +import 'dayjs/locale/zh-cn'; +import useIsDark from '../StatsClient/useIsDark'; + +dayjs.locale('zh-cn'); + +const { RangePicker } = DatePicker; + +export default ({ allStat }) => { + const isBrowser = useIsBrowser(); + const isDark = useIsDark(); + const today = dayjs().format('YYYY-MM-DD'); + const defaultDates = [dayjs().subtract(7, 'day'), dayjs()]; + const [dates, setDates] = useState(defaultDates); + const [chartData, setChartData] = useState([]); + + useEffect(() => { + if (dates && dates.length === 2 && dates[0] && dates[1]) { + const startDate = dayjs(dates[0]); + const endDate = dayjs(dates[1]); + + const dateArray = []; + let currentDate = startDate.clone(); + while (currentDate.isBefore(endDate) || currentDate.isSame(endDate)) { + dateArray.push(currentDate.format('YYYY-MM-DD')); + currentDate = currentDate.add(1, 'day'); + } + + const filteredData = dateArray.map(date => ({ + date, + ...(allStat[date] || {}) + })).filter(item => item.date); + + setChartData(filteredData); + } + }, [dates, allStat]); + + if (!isBrowser) return null; + + return +
+ setDates(dates)} presets={[{ label: '最近7天', value: [dayjs().subtract(7, 'day'), dayjs()] }, { label: '最近30天', value: [dayjs().subtract(30, 'day'), dayjs()] }]} /> + item.format('YYYY-MM-DD')).join('——')}数据统计图`}> + item.date), + }, + yAxis: { + type: 'value' + }, + series: [ + { + name: 'Fork数', + type: 'line', + data: chartData.map(item => item.forks || 0) + }, + { + name: 'Star数', + type: 'line', + data: chartData.map(item => item.stars || 0) + }, + { + name: 'Contributors', + type: 'line', + data: chartData.map(item => item.contributors || 0) + }, + { + name: 'issues总数', + type: 'line', + data: chartData.map(item => item.issues?.total || 0) + }, + { + name: 'issues open数', + type: 'line', + data: chartData.map(item => item.issues?.open || 0) + }, + { + name: 'issue 解决率', + type: 'line', + data: chartData.map(item => item.issues?.fixRate || 0) + }, + { + name: 'PR总数', + type: 'line', + data: chartData.map(item => item.prs?.total || 0) + }, + { + name: 'PR open数', + type: 'line', + data: chartData.map(item => item.prs?.open || 0) + }, + ], + }} + isDark={isDark} + style={{ widht: '100%', height: 400 }} + /> + +
+ +} diff --git a/docs/src/components/StatsClient/index.js b/docs/src/components/StatsClient/index.js new file mode 100644 index 0000000000..d4d1a26ca7 --- /dev/null +++ b/docs/src/components/StatsClient/index.js @@ -0,0 +1,30 @@ +import React from 'react'; +import useIsBrowser from '@docusaurus/useIsBrowser'; +import useIsDark from './useIsDark'; +import dayjs from 'dayjs'; +import { Card, Statistic, theme, Flex, ConfigProvider } from 'antd'; + +export default ({ todayStat }) => { + const isBrowser = useIsBrowser(); + const isDark = useIsDark(); + const today = dayjs().format('YYYY-MM-DD'); + + if (!isBrowser) return null; + + return +
+ + + + + + + + + + + + +
+ +} diff --git a/docs/src/components/StatsClient/useIsDark.js b/docs/src/components/StatsClient/useIsDark.js new file mode 100644 index 0000000000..00fdd86cbd --- /dev/null +++ b/docs/src/components/StatsClient/useIsDark.js @@ -0,0 +1,19 @@ +import { useEffect, useState } from 'react'; + +export default function useIsDark() { + const [isDark, setIsDark] = useState(false); + + useEffect(() => { + const update = () => { + const html = document.documentElement; + const classVal = html.getAttribute('data-theme') || ''; + setIsDark(classVal === 'dark'); + }; + update(); + const observer = new MutationObserver(update); + observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] }); + return () => observer.disconnect(); + }, []); + + return isDark; +} diff --git a/docs/src/pages/stats/index.js b/docs/src/pages/stats/index.js index 7e5289e30c..e04bc12561 100644 --- a/docs/src/pages/stats/index.js +++ b/docs/src/pages/stats/index.js @@ -1,13 +1,12 @@ import React, { useEffect, useState } from 'react'; import Layout from '@theme/Layout'; -import dayjs from 'dayjs'; -import Statistics from '../../components/Statistics'; -import StatChart from '../../components/StatChart'; +import Statistics from '../../components/StatsClient'; +import StatChart from '../../components/StatChartClient'; export default () => { const [todayStat, setTodayStat] = useState({}); const [allStat, setAllStat] = useState({}); - const today = dayjs().format('YYYY-MM-DD'); + const today = new Date().toISOString().slice(0, 10); useEffect(() => { fetch('/ROCK/stats.json').then(res => res.json()).then(data => { setTodayStat(data[today]); @@ -17,13 +16,11 @@ export default () => { return
-
- - -
- 文档页面详细统计可以打开Google analytics查看 -
+ + +
+ 文档页面详细统计可以打开Google analytics查看
- + ; } \ No newline at end of file From 41e9441f7da6b342aa2561ea4ec0babf7f95ca02 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 3 Apr 2026 16:54:36 +0800 Subject: [PATCH 014/226] add v1.4.8 docs --- .../version-1.4.x/Release Notes/v1.4.8.md | 13 +++++++++++++ .../version-1.4.x/Release Notes/v1.4.8.md | 14 ++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/Release Notes/v1.4.8.md create mode 100644 docs/versioned_docs/version-1.4.x/Release Notes/v1.4.8.md diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/Release Notes/v1.4.8.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/Release Notes/v1.4.8.md new file mode 100644 index 0000000000..6deb371801 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/Release Notes/v1.4.8.md @@ -0,0 +1,13 @@ +# v1.4.8 + +## 发布日期 +2026 年 4 月 3 日 + +--- + +## Admin + +### 增强改进 + +#### Kata 支持 Dind +- Kata 支持 dind 模式 ([#725](https://github.com/alibaba/ROCK/pull/725)) ([#732](https://github.com/alibaba/ROCK/pull/732)) \ No newline at end of file diff --git a/docs/versioned_docs/version-1.4.x/Release Notes/v1.4.8.md b/docs/versioned_docs/version-1.4.x/Release Notes/v1.4.8.md new file mode 100644 index 0000000000..45d1378df3 --- /dev/null +++ b/docs/versioned_docs/version-1.4.x/Release Notes/v1.4.8.md @@ -0,0 +1,14 @@ +# v1.4.8 + +## Release Date +Apr 03, 2026 + +--- + +## Admin + +### Enhancements + +#### Kata Support Dind +- Kata support dind ([#725](https://github.com/alibaba/ROCK/pull/725)) ([#732](https://github.com/alibaba/ROCK/pull/732)) + From 9e7e6360bf1fbd7b016d98a69728215f74122b12 Mon Sep 17 00:00:00 2001 From: "kezhong.ykz" Date: Thu, 9 Apr 2026 15:57:10 +0800 Subject: [PATCH 015/226] docs: rename add-doc-version skill to rock-docs (#752) --- .claude/skills/add-doc-version/SKILL.md | 406 ------------------ .claude/skills/rock-docs/SKILL.md | 237 ++++++++++ .../skills/rock-docs/references/i18n-guide.md | 70 +++ 3 files changed, 307 insertions(+), 406 deletions(-) delete mode 100644 .claude/skills/add-doc-version/SKILL.md create mode 100644 .claude/skills/rock-docs/SKILL.md create mode 100644 .claude/skills/rock-docs/references/i18n-guide.md diff --git a/.claude/skills/add-doc-version/SKILL.md b/.claude/skills/add-doc-version/SKILL.md deleted file mode 100644 index e6d301c01d..0000000000 --- a/.claude/skills/add-doc-version/SKILL.md +++ /dev/null @@ -1,406 +0,0 @@ -# Add Doc Version - -为 ROCK 项目添加新版本的 Docusaurus 文档。 - -## 项目文档背景 - -- 文档框架：Docusaurus 3.9.x，部署在 `https://alibaba.github.io/ROCK/` -- 语言：英文（默认）+ 中文（zh-Hans） -- 版本管理：`includeCurrentVersion: false`，只发布 versioned docs -- 所有文档文件均在 `docs/` 目录下 - -## 需要修改的文件和目录 - -添加新版本时，必须操作以下 **6 个位置**： - -| # | 操作 | 路径 | -|---|------|------| -| 1 | 创建英文文档目录 | `docs/versioned_docs/version-{NEW}/` | -| 2 | 创建版本侧边栏 | `docs/versioned_sidebars/version-{NEW}-sidebars.json` | -| 3 | 更新版本列表 | `docs/versions.json` | -| 4 | 更新最新版本指向 | `docs/docusaurus.config.js` 中的 `lastVersion` | -| 5 | 创建中文文档目录 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/` | -| 6 | 创建中文侧边栏翻译 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json` | - -## 每个版本的标准目录结构 - -英文和中文文档目录结构一致： - -``` -version-X.Y.x/ -├── Getting Started/ -│ ├── quickstart.md -│ ├── installation.md -│ └── ... -├── User Guides/ -│ └── ... -├── References/ -│ ├── api.md -│ └── Python SDK References/ -│ ├── python_sdk.md -│ └── ... -├── Release Notes/ -│ ├── index.md -│ └── vX.Y.Z.md -└── overview.md -``` - -## 执行流程 - -### Phase 0：版本冲突检测（前置检查） - -在执行任何操作之前，必须先进行版本冲突检测： - -1. 读取 `docs/versions.json` 获取已有版本列表 -2. 从用户输入的版本号推断文档版本号（如 `1.4.5` → `1.4.x`） -3. **检查推断出的文档版本号是否已存在于 `versions.json` 中** - -如果文档版本已存在（例如用户输入 `1.4.5`，但 `1.4.x` 已在版本列表中），则： - -- **不执行完整的新版本创建流程** -- 提示用户：`版本 1.4.x 的文档已存在，无需创建新版本。仅需在现有版本中添加 Release Note 即可。` -- **直接跳转到「仅添加 Release Note」流程**（见下方） - -#### 仅添加 Release Note 流程 - -当文档版本已存在时，只需执行以下 3 步操作： - -**Step A：创建英文 Release Note 文件** - -在 `docs/versioned_docs/version-{EXISTING}/Release Notes/v{VERSION}.md` 创建文件，使用 Phase 2 Step 3 中定义的**英文模板**。 - -**Step B：创建中文 Release Note 文件** - -在 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{EXISTING}/Release Notes/v{VERSION}.md` 创建文件，使用 Phase 2 Step 3 中定义的**中文模板**。 - -**Step C：更新 Release Notes 索引（index.md）** - -在英文和中文的 `Release Notes/index.md` 文件的链接列表**顶部**插入新版本条目。 - -英文文件 `docs/versioned_docs/version-{EXISTING}/Release Notes/index.md`，在标题行下方第一行插入： -```markdown -* [release v{VERSION}](v{VERSION}.md) -``` - -中文文件 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{EXISTING}/Release Notes/index.md`，同样在标题行下方第一行插入： -```markdown -* [release v{VERSION}](v{VERSION}.md) -``` - -例如，添加 `v1.4.5` 后，英文 index.md 应变为： -```markdown ---- -sidebar_position: 1 ---- -# Release Notes -* [release v1.4.5](v1.4.5.md) -* [release v1.4.4](v1.4.4.md) -* [release v1.4.3](v1.4.3.md) -... -``` - -创建完成后直接跳转到 Phase 5 验证步骤。 - ---- - -如果文档版本不存在，继续执行以下完整流程。 - -### Phase 1：收集信息 - -使用交互方式向用户询问： - -1. **新版本号**：如 `1.4.x` -2. **基础版本**：从哪个版本复制？默认取 `docs/versions.json` 的第一项（即最新版本） -3. **Release Note**：是否创建新的 Release Note？若是，具体版本号是什么（如 `v1.4.0`）？ -4. **是否修改 `lastVersion`**：**必须明确告知用户影响后再确认**。提示内容如下： - - > ⚠️ 修改 `lastVersion` 会改变文档站点的默认展示版本。 - > 当前默认展示版本为 `{当前 lastVersion}`，如果修改为 `{NEW}`，用户访问文档时将默认看到新版本内容。 - > 是否确认修改？ - - **不要默认选「是」**，必须等用户明确确认。 - -如果用户给的是 `1.4.0` 这样的具体版本号，自动推断文档版本号为 `1.4.x`。 - -### Phase 2：创建英文文档（3 步） - -#### Step 1：复制英文文档目录 - -```bash -cp -r docs/versioned_docs/version-{BASE}/ docs/versioned_docs/version-{NEW}/ -``` - -#### Step 1.5：清理旧版本 Release Notes（仅大版本号变更时） - -判断新版本的主版本号（major.minor）是否与基础版本不同。例如： -- `1.4.x` → `1.5.x`：主版本号从 `1.4` 变为 `1.5`，**需要清理** -- `1.4.x` → `1.4.x`：同主版本号，不需要清理（且此情况会被 Phase 0 拦截） - -当主版本号变更时，复制过来的 `Release Notes/` 目录中包含的是旧版本的 Release Notes，需要清理： - -1. **删除** `docs/versioned_docs/version-{NEW}/Release Notes/` 下除 `index.md` 之外的所有 `v*.md` 文件 -2. **重写** `index.md`，仅保留框架： - -```markdown ---- -sidebar_position: 1 ---- -# Release Notes -``` - -这样新版本的 Release Notes 目录就是干净的，不会携带旧版本的发布说明。 - -#### Step 2：复制版本侧边栏 - -```bash -cp docs/versioned_sidebars/version-{BASE}-sidebars.json docs/versioned_sidebars/version-{NEW}-sidebars.json -``` - -侧边栏文件内容无需修改（除非新版本文档结构有变化）。 - -#### Step 3：创建 Release Note（如果需要） - -英文和中文使用不同的模板。 - -**英文模板**：在 `docs/versioned_docs/version-{NEW}/Release Notes/v{VERSION}.md` 创建： - -```markdown -# v{VERSION} - -## Release Date -{Mon DD, YYYY} - ---- - -TODO -``` - -**中文模板**：在 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/v{VERSION}.md` 创建： - -```markdown -# v{VERSION} - -## 发布日期 -{YYYY} 年 {M} 月 {D} 日 - ---- - -TODO -``` - -注意中英文模板的差异： -- 标题部分：`## Release Date` vs `## 发布日期` -- 日期格式：`Mar 14, 2026` vs `2026 年 3 月 14 日` - -#### Step 3.5：更新 Release Notes 索引（如果创建了 Release Note） - -在英文 `docs/versioned_docs/version-{NEW}/Release Notes/index.md` 的标题行下方第一行插入： -```markdown -* [release v{VERSION}](v{VERSION}.md) -``` - -### Phase 3：创建中文文档（2 步） - -#### Step 4：复制中文文档目录 - -```bash -cp -r docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}/ docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/ -``` - -#### Step 4.5：清理中文旧版本 Release Notes（仅大版本号变更时） - -与 Step 1.5 同理，当主版本号变更时，需要清理中文目录下的旧 Release Notes： - -1. **删除** `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/` 下除 `index.md` 之外的所有 `v*.md` 文件 -2. **重写** `index.md`，仅保留框架： - -```markdown ---- -sidebar_position: 1 ---- -# 版本说明 -``` - -注意中文版 index.md 的标题是 `# 版本说明`，而非 `# Release Notes`。 - -#### Step 5：创建中文侧边栏翻译 JSON - -```bash -cp docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}.json docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json -``` - -然后修改新 JSON 文件中的 `version.label`： - -```json -{ - "version.label": { - "message": "{NEW}", - "description": "The label for version {NEW}" - } -} -``` - -其他侧边栏分类翻译（快速上手、用户指南、参考、版本说明等）保持不变。 - -如果创建了 Release Note，也要在中文目录下使用**中文模板**创建对应文件： -`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/v{VERSION}.md` - -并在中文 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/index.md` 的标题行下方第一行插入： -```markdown -* [release v{VERSION}](v{VERSION}.md) -``` - -### Phase 4：更新配置（2 步） - -#### Step 6：更新 versions.json - -在 `docs/versions.json` 数组**开头**插入新版本号： - -```json -[ - "{NEW}", - "1.3.x", - "1.2.x", - ... -] -``` - -#### Step 7：更新 docusaurus.config.js（仅在用户明确确认后） - -**只有用户在 Phase 1 中明确确认要修改 `lastVersion` 时才执行此步骤。** - -修改 `docs/docusaurus.config.js` 中的 `lastVersion`： - -```js -// 找到这一行 -lastVersion: '{BASE}', -// 替换为 -lastVersion: '{NEW}', -``` - -如果用户选择不修改，则跳过此步骤，保持原有的 `lastVersion` 不变。 - -### Phase 5：验证 - -完成所有步骤后，执行以下验证流程： - -#### Step 8：构建验证 - -先执行构建，确保没有编译错误： - -```bash -cd docs && npm run build -``` - -构建成功后，再启动本地预览服务： - -```bash -cd docs && npm run serve -``` - -> **为什么用 `npm run serve` 而不是 `npm run start`？** -> `npm run start` 只启动开发模式，默认只加载默认语言（英文），无法切换到中文验证。 -> `npm run serve` 基于构建产物启动静态服务，可以同时预览英文和中文版本，支持通过语言切换器验证中英文文档。 - -#### Step 9：手动检查清单 - -提示用户在浏览器中检查以下内容： - -``` -✅ 文档验证清单： - -英文版本检查： - □ 版本下拉菜单中是否显示新版本 - □ 默认展示的版本是否符合预期（取决于是否修改了 lastVersion） - □ 英文文档页面是否正常渲染 - □ Release Notes 是否按版本号倒序排列 - □ 侧边栏导航是否完整 - -中文版本检查（通过右上角语言切换器切换到中文）： - □ 中文文档页面是否正常显示 - □ 侧边栏分类名称是否正确翻译（快速上手、用户指南、参考、版本说明） - □ Release Notes 内容是否与英文版本对应 - -后续工作： - □ 更新新版本英文文档中的具体内容 - □ 更新中文文档翻译内容 - □ 如有新的 Release Note，补充具体发布内容 -``` - -## 注意事项 - -- `docs/rock/` 是 "current"（未发布）版本源码，因为 `includeCurrentVersion: false` 所以不会发布。所有发布内容来自 `versioned_docs/`。 -- `docusaurus.config.js` 通过 `convertVersionsArrayToObject()` 从 `versions.json` 自动生成版本配置，因此只需在 `versions.json` 添加即可完成版本注册。 -- 侧边栏使用 `autogenerated` 模式，新文件放入正确目录后会自动出现在侧边栏。 -- Release Notes 通过 `reverseReleaseNoteSidebars()` 自动按版本号倒序排列。 -- `HiddenSidebars` 数组中的文件（`Getting Started/quickstart`、`References/Python SDK References/python_sdk`、`Release Notes/index`）会从侧边栏隐藏，但仍可通过直接链接访问。 - -## 扩展参考：Docusaurus i18n 与侧边栏配置 - -当用户询问如何添加新的侧边栏分类、如何添加翻译、如何配置 i18n 等问题时，参考 Docusaurus 官方文档： - -**官方 i18n 教程**：https://docusaurus.io/docs/i18n/tutorial - -关键知识点： - -### 添加侧边栏分类 - -本项目的侧边栏采用 **“顶层手动 + 内层自动”** 的混合模式： - -- **顶层分类**（如 Getting Started、User Guides、References、Release Notes）是在侧边栏配置中**手动定义**的 -- **内层文档**通过 `autogenerated` 从目录中**自动生成**，新文件放入对应目录即可自动出现 - -因此，**添加新的顶层侧边栏分类**时，需要手动修改以下文件： - -1. **主侧边栏配置**：`docs/sidebars.js` — 添加新的顶层 category -2. **各版本侧边栏**：`docs/versioned_sidebars/version-{VERSION}-sidebars.json` — 在需要显示新分类的版本中添加对应的 category 配置 -3. **中文翻译 JSON**：`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{VERSION}.json` — 添加 `sidebar.tutorialSidebar.category.{CategoryName}` 条目 - -顶层分类的配置格式示例（以 `version-1.4.x-sidebars.json` 为例）： - -```json -{ - "type": "category", - "label": "New Category", - "items": [ - { - "type": "autogenerated", - "dirName": "New Category" - } - ] -} -``` - -同时需在 `versioned_docs/version-{VERSION}/` 和对应的 i18n 目录下创建对应的文件夹，内层文档会自动生成侧边栏条目。 - -而 **在已有分类下添加新文档**，则无需修改侧边栏配置，直接将 md 文件放入对应目录即可。 - -### 添加翻译 - -本项目的翻译文件存储在 `docs/i18n/zh-Hans/` 目录下，主要包括： - -| 文件/目录 | 用途 | -|---------|------| -| `docusaurus-plugin-content-docs/version-{V}.json` | 侧边栏分类名称和版本标签的翻译 | -| `docusaurus-plugin-content-docs/version-{V}/` | 各版本文档内容的中文翻译 | -| `docusaurus-theme-classic/navbar.json` | 导航栏文本翻译 | -| `docusaurus-theme-classic/footer.json` | 页脚文本翻译 | -| `code.json` | React 代码中的文本标签翻译 | - -提取翻译 key 的命令： -```bash -cd docs && npm run write-translations -- --locale zh-Hans -``` - -该命令会自动扫描项目代码和配置，生成需要翻译的 JSON 文件。生成后编辑对应的 JSON 文件填入中文翻译即可。 - -### 翻译文档内容 - -文档内容的翻译通过在 i18n 目录下创建对应的 markdown 文件实现。文件路径必须与英文原文完全对应： - -``` -英文原文：docs/versioned_docs/version-1.4.x/User Guides/example.md -中文翻译：docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/User Guides/example.md -``` - -如果中文目录下不存在对应文件，Docusaurus 会回退到英文原文展示。 diff --git a/.claude/skills/rock-docs/SKILL.md b/.claude/skills/rock-docs/SKILL.md new file mode 100644 index 0000000000..c4ec38f4b3 --- /dev/null +++ b/.claude/skills/rock-docs/SKILL.md @@ -0,0 +1,237 @@ +--- +name: rock-docs +description: 为 ROCK 项目管理 Docusaurus 文档版本。支持添加新版本文档、添加 Release Note、管理中英文 i18n 文档。当用户说"添加文档版本"、"发布新版本文档"、"添加 release note"、"rock-docs"或需要操作 docs/ 目录下的版本化文档时使用。 +--- + +# ROCK Docs + +管理 ROCK 项目的 Docusaurus 版本化文档。 + +## 项目背景 + +- Docusaurus 3.9.x，部署于 `https://alibaba.github.io/ROCK/` +- 语言：英文（默认）+ 中文（zh-Hans） +- `includeCurrentVersion: false`，所有发布内容来自 `versioned_docs/` +- `versions.json` 驱动版本注册，`convertVersionsArrayToObject()` 自动生成配置 +- 侧边栏：顶层手动 + 内层 `autogenerated` +- Release Notes 通过 `reverseReleaseNoteSidebars()` 自动倒序排列 + +## 需要操作的 6 个位置 + +| # | 操作 | 路径 | +|---|------|------| +| 1 | 创建英文文档目录 | `docs/versioned_docs/version-{NEW}/` | +| 2 | 创建版本侧边栏 | `docs/versioned_sidebars/version-{NEW}-sidebars.json` | +| 3 | 更新版本列表 | `docs/versions.json` | +| 4 | 更新最新版本指向 | `docs/docusaurus.config.js` 中的 `lastVersion` | +| 5 | 创建中文文档目录 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/` | +| 6 | 创建中文侧边栏翻译 | `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json` | + +## 标准目录结构 + +``` +version-X.Y.x/ +├── Getting Started/ +├── User Guides/ +├── References/ +│ ├── api.md +│ └── Python SDK References/ +├── Release Notes/ +│ ├── index.md +│ └── vX.Y.Z.md +└── overview.md +``` + +## 执行流程 + +### Phase 0：版本冲突检测 + +1. 读取 `docs/versions.json` 获取已有版本列表 +2. 从用户输入的版本号推断文档版本号（如 `1.4.5` → `1.4.x`） +3. 如果文档版本已存在，**跳转至「仅添加 Release Note」流程** + +#### 仅添加 Release Note 流程 + +**Step A**：创建英文 Release Note +- 路径：`docs/versioned_docs/version-{EXISTING}/Release Notes/v{VERSION}.md` +- 使用下方英文模板 + +**Step B**：创建中文 Release Note +- 路径：`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{EXISTING}/Release Notes/v{VERSION}.md` +- 使用下方中文模板 + +**Step C**：更新 Release Notes 索引 + +在英文和中文 `Release Notes/index.md` 的标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +完成后跳转至 Phase 5 验证。 + +--- + +如果文档版本不存在，继续执行完整流程。 + +### Phase 1：收集信息 + +1. **新版本号**：如 `1.4.x` +2. **基础版本**：从哪个版本复制？默认取 `versions.json` 第一项 +3. **Release Note**：是否创建？版本号是什么？ +4. **是否修改 `lastVersion`**：必须明确告知用户影响后再确认： + + > ⚠️ 修改 `lastVersion` 会改变文档站点默认展示版本。当前为 `{当前 lastVersion}`，修改后为 `{NEW}`。是否确认？ + + **不要默认选「是」，等用户明确确认。** + +具体版本号（如 `1.4.0`）自动推断为文档版本号（`1.4.x`）。 + +### Phase 2：创建英文文档 + +#### Step 1：复制英文文档目录 + +```bash +cp -r docs/versioned_docs/version-{BASE}/ docs/versioned_docs/version-{NEW}/ +``` + +#### Step 1.5：清理旧版本 Release Notes（仅主版本号变更时） + +当主版本号（major.minor）不同时（如 `1.4.x` → `1.5.x`）： + +1. 删除 `docs/versioned_docs/version-{NEW}/Release Notes/` 下除 `index.md` 外的所有 `v*.md` +2. 重写 `index.md`： + +```markdown +--- +sidebar_position: 1 +--- +# Release Notes +``` + +#### Step 2：复制版本侧边栏 + +```bash +cp docs/versioned_sidebars/version-{BASE}-sidebars.json docs/versioned_sidebars/version-{NEW}-sidebars.json +``` + +#### Step 3：创建 Release Note（如果需要） + +**英文模板** `docs/versioned_docs/version-{NEW}/Release Notes/v{VERSION}.md`： +```markdown +# v{VERSION} + +## Release Date +{Mon DD, YYYY} + +--- + +TODO +``` + +**中文模板** `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/v{VERSION}.md`： +```markdown +# v{VERSION} + +## 发布日期 +{YYYY} 年 {M} 月 {D} 日 + +--- + +TODO +``` + +#### Step 3.5：更新 Release Notes 索引 + +在英文 `docs/versioned_docs/version-{NEW}/Release Notes/index.md` 标题行下方第一行插入： +```markdown +* [release v{VERSION}](v{VERSION}.md) +``` + +### Phase 3：创建中文文档 + +#### Step 4：复制中文文档目录 + +```bash +cp -r docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}/ docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/ +``` + +#### Step 4.5：清理中文旧版本 Release Notes（仅主版本号变更时） + +与 Step 1.5 同理： +1. 删除 `docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}/Release Notes/` 下除 `index.md` 外的所有 `v*.md` +2. 重写 `index.md`（中文版标题为 `# 版本说明`）： + +```markdown +--- +sidebar_position: 1 +--- +# 版本说明 +``` + +#### Step 5：创建中文侧边栏翻译 JSON + +```bash +cp docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{BASE}.json docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{NEW}.json +``` + +修改新 JSON 文件中的 `version.label`： + +```json +{ + "version.label": { + "message": "{NEW}", + "description": "The label for version {NEW}" + } +} +``` + +### Phase 4：更新配置 + +#### Step 6：更新 versions.json + +在 `docs/versions.json` 数组开头插入 `{NEW}`。 + +#### Step 7：更新 docusaurus.config.js（仅用户明确确认后） + +修改 `docs/docusaurus.config.js` 中的 `lastVersion` 从 `{BASE}` 为 `{NEW}`。用户选择不修改则跳过。 + +### Phase 5：验证 + +**Step 8：构建验证** +```bash +cd docs && npm run build +``` +构建成功后启动预览： +```bash +cd docs && npm run serve +``` +> `npm run serve` 基于构建产物启动，支持中英文切换验证。`npm run start` 仅加载默认语言，无法验证中文。 + +**Step 9：手动检查清单** + +``` +✅ 文档验证清单： + +英文版本检查： + □ 版本下拉菜单中是否显示新版本 + □ 默认展示版本是否符合预期 + □ 英文文档页面是否正常渲染 + □ Release Notes 是否按版本号倒序排列 + □ 侧边栏导航是否完整 + +中文版本检查（右上角语言切换）： + □ 中文文档页面是否正常显示 + □ 侧边栏分类名称是否正确翻译 + □ Release Notes 内容是否与英文版本对应 + +后续工作： + □ 更新新版本英文文档中的具体内容 + □ 更新中文文档翻译内容 + □ 如有 Release Note，补充具体发布内容 +``` + +## 注意事项 + +- `docs/rock/` 是未发布版本源码，`includeCurrentVersion: false` 不会发布 +- `HiddenSidebars` 中的文件会从侧边栏隐藏，但仍可通过直接链接访问 +- 详细 i18n 配置、侧边栏管理、翻译流程参见 [references/i18n-guide.md](references/i18n-guide.md) diff --git a/.claude/skills/rock-docs/references/i18n-guide.md b/.claude/skills/rock-docs/references/i18n-guide.md new file mode 100644 index 0000000000..d0e1eecc61 --- /dev/null +++ b/.claude/skills/rock-docs/references/i18n-guide.md @@ -0,0 +1,70 @@ +# Docusaurus i18n 与侧边栏管理参考 + +## 扩展参考 + +- 官方 i18n 教程：https://docusaurus.io/docs/i18n/tutorial + +## 侧边栏配置 + +本项目采用 **"顶层手动 + 内层自动"** 的混合模式： + +- **顶层分类**（Getting Started、User Guides、References、Release Notes）在侧边栏配置中**手动定义** +- **内层文档**通过 `autogenerated` 从目录中**自动生成**，新文件放入对应目录即可自动出现在侧边栏 + +### 添加新的顶层侧边栏分类 + +需要手动修改以下文件： + +1. **主侧边栏配置**：`docs/sidebars.js` — 添加新的顶层 category +2. **各版本侧边栏**：`docs/versioned_sidebars/version-{VERSION}-sidebars.json` — 添加对应的 category 配置 +3. **中文翻译 JSON**：`docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-{VERSION}.json` — 添加 `sidebar.tutorialSidebar.category.{CategoryName}` 条目 + +顶层分类配置格式示例（`version-1.4.x-sidebars.json`）： + +```json +{ + "type": "category", + "label": "New Category", + "items": [ + { + "type": "autogenerated", + "dirName": "New Category" + } + ] +} +``` + +在 `versioned_docs/version-{VERSION}/` 和对应 i18n 目录下创建同名文件夹，内层文档会自动生成侧边栏条目。 + +**在已有分类下添加新文档**，无需修改侧边栏配置，直接将 md 文件放入对应目录即可。 + +## 翻译文件结构 + +翻译文件存储在 `docs/i18n/zh-Hans/` 目录下： + +| 文件/目录 | 用途 | +|---------|------| +| `docusaurus-plugin-content-docs/version-{V}.json` | 侧边栏分类名称和版本标签翻译 | +| `docusaurus-plugin-content-docs/version-{V}/` | 各版本文档内容中文翻译 | +| `docusaurus-theme-classic/navbar.json` | 导航栏文本翻译 | +| `docusaurus-theme-classic/footer.json` | 页脚文本翻译 | +| `code.json` | React 代码中的文本标签翻译 | + +### 提取翻译 key + +```bash +cd docs && npm run write-translations -- --locale zh-Hans +``` + +该命令自动扫描项目代码和配置，生成需要翻译的 JSON 文件。生成后编辑对应 JSON 文件填入中文翻译。 + +## 翻译文档内容 + +文档内容翻译通过在 i18n 目录下创建对应的 markdown 文件实现。文件路径必须与英文原文完全对应： + +``` +英文原文：docs/versioned_docs/version-1.4.x/User Guides/example.md +中文翻译：docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.4.x/User Guides/example.md +``` + +如果中文目录下不存在对应文件，Docusaurus 会回退到英文原文展示。 From 2476375d5cf04c22030b34f80ad41d629ee6fd6e Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 9 Apr 2026 17:20:55 +0800 Subject: [PATCH 016/226] Feature/xinshi/container verify (#755) * feat: add container mode and related fields to VerifierConfig * refactor: update verifier config fields names and types * feat: add NativeConfig for verifier mode settings --- rock/sdk/agent/models/trial/config.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/rock/sdk/agent/models/trial/config.py b/rock/sdk/agent/models/trial/config.py index acf961d7c7..718daf57b5 100644 --- a/rock/sdk/agent/models/trial/config.py +++ b/rock/sdk/agent/models/trial/config.py @@ -77,11 +77,24 @@ def to_harbor_environment(self) -> dict: return harbor.model_dump(mode="json", exclude_none=True) +class NativeConfig(BaseModel): + """Config specific to native verifier mode. + When image and script are both provided, a ContainerVerifier is used to + run evaluation in an isolated container. Otherwise the built-in SWE-bench + run_instance flow is used. + """ + + image: str | None = None + script: str | None = None + oss_deps: dict[str, str] = Field(default_factory=dict) + + class VerifierConfig(BaseModel): override_timeout_sec: float | None = None max_timeout_sec: float | None = None disable: bool = False mode: Literal["harbor", "native"] | None = None + native_config: NativeConfig = Field(default_factory=NativeConfig) class TaskConfig(BaseModel): From 60aba54a9eb695ccfec02013599c65200c986e84 Mon Sep 17 00:00:00 2001 From: lkc Date: Thu, 9 Apr 2026 17:30:22 +0800 Subject: [PATCH 017/226] feat: add labels support to JobConfig (#720) (#721) * feat(job-config): translate field descriptions to English and update job models * fix: add missing experiment_id to tests and backfill sandbox experiment_id in _autofill_sandbox_info * feat: modify job config field description * fix: only set experiment_id on JobConfig level in _autofill_sandbox_info --- rock/sdk/agent/job.py | 10 ++++++ rock/sdk/agent/models/job/config.py | 12 +++++-- rock/sdk/agent/models/job/result.py | 1 + tests/unit/sdk/agent/test_job.py | 14 ++++++++ .../agent/test_job_config_serialization.py | 34 +++++++++++++++++++ 5 files changed, 69 insertions(+), 2 deletions(-) diff --git a/rock/sdk/agent/job.py b/rock/sdk/agent/job.py index fc01dfb2ae..f0dfcbb18f 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/agent/job.py @@ -258,6 +258,7 @@ async def _collect_results(self) -> JobResult: return JobResult( job_id=self._config.job_name, status=JobStatus.COMPLETED if trial_results else JobStatus.FAILED, + labels=self._config.labels, trial_results=trial_results, ) @@ -275,6 +276,15 @@ async def _autofill_sandbox_info(self) -> None: if sandbox_ns is not None: self._config.namespace = sandbox_ns + sandbox_exp = self._sandbox._experiment_id + if sandbox_exp is not None: + if self._config.experiment_id is not None and self._config.experiment_id != sandbox_exp: + raise ValueError( + f"experiment_id mismatch: JobConfig has '{self._config.experiment_id}', " + f"but sandbox returned '{sandbox_exp}'" + ) + self._config.experiment_id = sandbox_exp + async def _upload_content(self, content: str, sandbox_path: str) -> None: """Write text content to a local temp file and upload to sandbox via upload_by_path.""" local_tmp = None diff --git a/rock/sdk/agent/models/job/config.py b/rock/sdk/agent/models/job/config.py index 839c1aaee2..d0eaeeaca4 100644 --- a/rock/sdk/agent/models/job/config.py +++ b/rock/sdk/agent/models/job/config.py @@ -143,8 +143,9 @@ class JobConfig(BaseModel): default=None, description="Tenant isolation identifier for distinguishing resources across teams/projects", ) - experiment_id: str = Field( - description="Experiment identifier, required", + experiment_id: str | None = Field( + default=None, + description="Experiment identifier", ) job_name: str = Field(default_factory=lambda: datetime.now().strftime("%Y-%m-%d__%H-%M-%S")) jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" @@ -162,6 +163,13 @@ class JobConfig(BaseModel): datasets: list[LocalDatasetConfig | RegistryDatasetConfig] = Field(default_factory=list) tasks: list[TaskConfig] = Field(default_factory=list) artifacts: list[str | ArtifactConfig] = Field(default_factory=list) + labels: dict[str, str] = Field( + default_factory=dict, + description="Key-value labels for organizing and filtering jobs. " + "Example: {'step': '42', 'env': 'prod'}. " + "Keys: [prefix/]name, lowercase, max 63 chars. " + "Values: max 255 chars. Reserved prefix: 'harbor.io/'.", + ) @model_validator(mode="after") def _sync_experiment_id(self): diff --git a/rock/sdk/agent/models/job/result.py b/rock/sdk/agent/models/job/result.py index 5115b5095c..a4544e85d3 100644 --- a/rock/sdk/agent/models/job/result.py +++ b/rock/sdk/agent/models/job/result.py @@ -22,6 +22,7 @@ class JobResult(BaseModel): job_id: str = "" status: JobStatus = JobStatus.COMPLETED + labels: dict[str, str] = Field(default_factory=dict) trial_results: list[TrialResult] = Field(default_factory=list) raw_output: str = "" exit_code: int = 0 diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index be2f84b2ef..67ff74d474 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -101,6 +101,20 @@ def test_empty_trials(self): assert r.n_completed == 0 assert r.n_failed == 0 + def test_labels_default_empty(self): + r = JobResult(job_id="job-no-labels") + assert r.labels == {} + + def test_labels_preserved(self): + r = JobResult( + job_id="job-labeled", + labels={"step": "42", "env": "prod"}, + trial_results=[ + TrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), + ], + ) + assert r.labels == {"step": "42", "env": "prod"} + def _make_mock_sandbox(): """Create a mock Sandbox with all required async methods.""" diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index 78506ef605..dccbef8d7f 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -152,6 +152,24 @@ def test_excludes_none_values(self): assert "agent_timeout_multiplier" not in data + def test_labels_serialized(self): + cfg = JobConfig( + job_name="labeled-job", + experiment_id="test-exp", + labels={"step": "42", "env": "prod"}, + ) + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + + assert data["labels"] == {"step": "42", "env": "prod"} + + def test_empty_labels_not_excluded(self): + cfg = JobConfig(job_name="no-labels", experiment_id="test-exp") + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + + assert data["labels"] == {} + def test_path_fields_serialized_as_strings(self): cfg = JobConfig( experiment_id="test-exp", @@ -277,3 +295,19 @@ def test_from_yaml_with_local_dataset(self, tmp_path): assert cfg.job_name == "local-dataset-job" assert isinstance(cfg.datasets[0], LocalDatasetConfig) assert cfg.datasets[0].path == Path("/data/tasks") + + def test_from_yaml_with_labels(self, tmp_path): + yaml_content = """ +job_name: labeled-job +experiment_id: test-exp +labels: + step: "42" + env: prod +agents: + - name: terminus-2 +""" + yaml_file = tmp_path / "config.yaml" + yaml_file.write_text(yaml_content) + + cfg = JobConfig.from_yaml(str(yaml_file)) + assert cfg.labels == {"step": "42", "env": "prod"} From cb24ac74d1ce96cbcddaa405c5ee3d78e5c5e031 Mon Sep 17 00:00:00 2001 From: berstpander Date: Thu, 9 Apr 2026 17:34:28 +0800 Subject: [PATCH 018/226] feat: auto-generate job_name from dataset and task info (#757) - Change job_name type from str to str | None with default None - Add _generate_default_job_name() method to generate meaningful names - Format: {dataset_name}_{task_name}_{uuid} for single task, or {dataset_name}_{uuid} for multiple/no tasks - Add comprehensive unit tests for job_name generation logic --- rock/sdk/agent/job.py | 34 ++++++++ rock/sdk/agent/models/job/config.py | 5 +- tests/unit/sdk/agent/test_job.py | 121 +++++++++++++++++++++++++++- 3 files changed, 158 insertions(+), 2 deletions(-) diff --git a/rock/sdk/agent/job.py b/rock/sdk/agent/job.py index f0dfcbb18f..ea68d0051e 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/agent/job.py @@ -10,6 +10,7 @@ import json import os import tempfile +import uuid from rock.actions import Command, CreateBashSessionRequest, ReadFileRequest from rock.logger import init_logger @@ -85,6 +86,9 @@ async def submit(self) -> None: """Start sandbox, upload config & script, nohup start harbor.""" from rock.sdk.sandbox.client import Sandbox + # Generate job_name if not set (must be done before using it) + self._generate_default_job_name() + self._sandbox = Sandbox(self._config.environment) await self._sandbox.start() logger.info(f"Sandbox started: sandbox_id={self._sandbox.sandbox_id}, job_name={self._config.job_name}") @@ -266,6 +270,36 @@ async def _collect_results(self) -> JobResult: # Private: utilities # ------------------------------------------------------------------ + def _generate_default_job_name(self) -> None: + """Generate default job_name if not explicitly set by user. + + If job_name is None, generate one with the format: + {dataset_name}_{task_name if single task}_{uuid} + """ + if self._config.job_name is not None: + # User has set a custom job_name, keep it + return + + # Generate new job_name based on datasets + parts = [] + + # Get dataset name + if self._config.datasets: + dataset = self._config.datasets[0] + if hasattr(dataset, "name") and dataset.name: + parts.append(dataset.name) + + # Get task name if there's only one task + task_names = dataset.task_names + if task_names and len(task_names) == 1: + parts.append(task_names[0]) + + # Add short UUID (8 characters) + parts.append(uuid.uuid4().hex[:8]) + + self._config.job_name = "_".join(parts) + logger.info(f"Auto-generated job_name: {self._config.job_name}") + async def _autofill_sandbox_info(self) -> None: sandbox_ns = self._sandbox._namespace if self._config.namespace is not None and sandbox_ns is not None: diff --git a/rock/sdk/agent/models/job/config.py b/rock/sdk/agent/models/job/config.py index d0eaeeaca4..2f9b37345b 100644 --- a/rock/sdk/agent/models/job/config.py +++ b/rock/sdk/agent/models/job/config.py @@ -147,7 +147,10 @@ class JobConfig(BaseModel): default=None, description="Experiment identifier", ) - job_name: str = Field(default_factory=lambda: datetime.now().strftime("%Y-%m-%d__%H-%M-%S")) + job_name: str | None = Field( + default=None, + description="Job name, auto-generated if not set", + ) jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" n_attempts: int = 1 timeout_multiplier: float = 1.0 diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index 67ff74d474..187b826510 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -1,9 +1,16 @@ import json import os +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch from rock.sdk.agent.job import Job, JobResult, JobStatus -from rock.sdk.agent.models.job.config import JobConfig, RegistryDatasetConfig, RemoteRegistryInfo, RockEnvironmentConfig +from rock.sdk.agent.models.job.config import ( + JobConfig, + LocalDatasetConfig, + RegistryDatasetConfig, + RemoteRegistryInfo, + RockEnvironmentConfig, +) from rock.sdk.agent.models.trial.config import AgentConfig from rock.sdk.agent.models.trial.result import ExceptionInfo, TrialResult, VerifierResult @@ -299,3 +306,115 @@ async def test_cancel_kills_process(self): # Verify kill command was issued call_args = mock_sandbox.arun.call_args assert "kill" in str(call_args) + + +class TestGenerateDefaultJobName: + """Tests for _generate_default_job_name method.""" + + def test_custom_job_name_not_overwritten(self): + """User-set job_name should not be overwritten.""" + config = JobConfig( + job_name="my-custom-job", + experiment_id="test-exp", + datasets=[RegistryDatasetConfig(registry=RemoteRegistryInfo(), name="tb", version="2.0")], + ) + job = Job(config) + job._generate_default_job_name() + + assert job._config.job_name == "my-custom-job" + + def test_job_name_generated_with_dataset_and_single_task(self): + """Default job_name should be generated with dataset name and single task.""" + config = JobConfig( + experiment_id="test-exp", + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="terminal-bench", + version="2.0", + task_names=["fix-bug"], + ) + ], + ) + job = Job(config) + job._generate_default_job_name() + + # Should be: terminal-bench_fix-bug_{uuid} + job_name = job._config.job_name + parts = job_name.split("_") + assert len(parts) == 3 + assert parts[0] == "terminal-bench" + assert parts[1] == "fix-bug" + assert len(parts[2]) == 8 # UUID is 8 chars + + def test_job_name_generated_with_dataset_multiple_tasks(self): + """With multiple tasks, only dataset name and UUID should be used.""" + config = JobConfig( + experiment_id="test-exp", + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="terminal-bench", + version="2.0", + task_names=["task1", "task2"], + ) + ], + ) + job = Job(config) + job._generate_default_job_name() + + # Should be: terminal-bench_{uuid} + job_name = job._config.job_name + parts = job_name.split("_") + assert len(parts) == 2 + assert parts[0] == "terminal-bench" + assert len(parts[1]) == 8 # UUID is 8 chars + + def test_job_name_generated_without_dataset(self): + """Without dataset, only UUID should be used.""" + config = JobConfig(experiment_id="test-exp") + job = Job(config) + job._generate_default_job_name() + + # Should be: {uuid} + job_name = job._config.job_name + assert len(job_name) == 8 # Only UUID + + def test_job_name_generated_with_dataset_no_name(self): + """Dataset without name field should still work.""" + config = JobConfig( + experiment_id="test-exp", + datasets=[LocalDatasetConfig(path=Path("/data/tasks"))], + ) + job = Job(config) + job._generate_default_job_name() + + # LocalDatasetConfig has no name, so only UUID + job_name = job._config.job_name + assert len(job_name) == 8 # Only UUID + + async def test_submit_generates_job_name(self): + """Verify that submit() triggers the job_name generation.""" + mock_sandbox = _make_mock_sandbox() + + with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): + config = JobConfig( + experiment_id="test-exp", + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="my-dataset", + task_names=["my-task"], + ) + ], + ) + job = Job(config) + + # job_name is None initially + assert config.job_name is None + + await job.submit() + + # job_name should have been generated + assert job._config.job_name is not None + assert job._config.job_name.startswith("my-dataset_my-task_") From 5218a7a1b4855a90fb9ab964cd4e69fed1ab15db Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Tue, 10 Mar 2026 10:31:58 +0800 Subject: [PATCH 019/226] add auto_delete_seconds in SandboxConfig --- rock/admin/proto/request.py | 2 ++ rock/deployments/config.py | 6 +++++- rock/sdk/sandbox/client.py | 1 + rock/sdk/sandbox/config.py | 1 + 4 files changed, 9 insertions(+), 1 deletion(-) diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index 08e3849649..5c736429bd 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -37,6 +37,8 @@ class SandboxStartRequest(BaseModel): """Password for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image.""" use_kata_runtime: bool = False """Whether to use kata container runtime (io.containerd.kata.v2) instead of --privileged mode.""" + auto_delete_seconds: int = 0 + """The time for automatic container deletion, with the unit being seconds.""" class SandboxCommand(Command): diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 57458b1836..2b459f4c59 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -176,7 +176,11 @@ def auto_clear_time(self) -> int: @classmethod def from_request(cls, request: SandboxStartRequest) -> DeploymentConfig: """Create DockerDeploymentConfig from SandboxStartRequest""" - return cls(**request.model_dump(exclude={"sandbox_id"}), container_name=request.sandbox_id) + return cls( + **request.model_dump(exclude={"sandbox_id", "auto_delete_seconds"}), + container_name=request.sandbox_id, + remove_container=(request.auto_delete_seconds == 0), + ) class RayDeploymentConfig(DockerDeploymentConfig): diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 89de7dbd59..77a9dc719a 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -179,6 +179,7 @@ async def start(self): "use_kata_runtime": self.config.use_kata_runtime, "limit_cpus": self.config.limit_cpus, "sandbox_id": self.config.sandbox_id, + "auto_delete_seconds": self.config.auto_delete_seconds, } try: response = await HttpUtils.post(url, headers, data) diff --git a/rock/sdk/sandbox/config.py b/rock/sdk/sandbox/config.py index 24f6d5da2f..fcd9d426ad 100644 --- a/rock/sdk/sandbox/config.py +++ b/rock/sdk/sandbox/config.py @@ -44,6 +44,7 @@ class SandboxConfig(BaseConfig): registry_password: str | None = None use_kata_runtime: bool = False sandbox_id: str | None = None + auto_delete_seconds: int = 0 class SandboxGroupConfig(SandboxConfig): From d1e90e484168d320f100c8038c0f472fa1c2cb83 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Thu, 9 Apr 2026 20:46:46 +0800 Subject: [PATCH 020/226] add auto_delete_seconds in DockerDeployment --- rock/admin/proto/request.py | 2 +- rock/deployments/config.py | 6 ++-- rock/deployments/manager.py | 5 +++- rock/sdk/sandbox/config.py | 9 +++++- tests/unit/deployments/test_get_deployment.py | 20 +++++++++++++ tests/unit/sdk/test_sandbox_config.py | 30 +++++++++++++++++++ 6 files changed, 67 insertions(+), 5 deletions(-) create mode 100644 tests/unit/sdk/test_sandbox_config.py diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index 5c736429bd..ab27bf22e0 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -37,7 +37,7 @@ class SandboxStartRequest(BaseModel): """Password for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image.""" use_kata_runtime: bool = False """Whether to use kata container runtime (io.containerd.kata.v2) instead of --privileged mode.""" - auto_delete_seconds: int = 0 + auto_delete_seconds: int | None = None """The time for automatic container deletion, with the unit being seconds.""" diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 2b459f4c59..35ab4fe4ab 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -96,6 +96,9 @@ class DockerDeploymentConfig(DeploymentConfig): container_name: str | None = None """Custom name for the container. If None, a random name will be generated.""" + auto_delete_seconds: int | None = None + """If set, the container will be automatically deleted after container stopped.""" + type: Literal["docker"] = "docker" """Deployment type discriminator for serialization/deserialization and CLI parsing. Should not be modified.""" @@ -177,9 +180,8 @@ def auto_clear_time(self) -> int: def from_request(cls, request: SandboxStartRequest) -> DeploymentConfig: """Create DockerDeploymentConfig from SandboxStartRequest""" return cls( - **request.model_dump(exclude={"sandbox_id", "auto_delete_seconds"}), + **request.model_dump(exclude={"sandbox_id"}), container_name=request.sandbox_id, - remove_container=(request.auto_delete_seconds == 0), ) diff --git a/rock/deployments/manager.py b/rock/deployments/manager.py index d64820ee6f..3805de5c3c 100644 --- a/rock/deployments/manager.py +++ b/rock/deployments/manager.py @@ -40,7 +40,10 @@ async def init_config(self, config: DeploymentConfig) -> DockerDeploymentConfig: await self.rock_config.update() docker_deployment_config.actor_resource = self.rock_config.sandbox_config.actor_resource docker_deployment_config.actor_resource_num = self.rock_config.sandbox_config.actor_resource_num - docker_deployment_config.remove_container = self.rock_config.sandbox_config.remove_container_enabled + if docker_deployment_config.auto_delete_seconds is None: + docker_deployment_config.remove_container = self.rock_config.sandbox_config.remove_container_enabled + else: + docker_deployment_config.remove_container = docker_deployment_config.auto_delete_seconds == 0 return docker_deployment_config def get_deployment(self, config: DeploymentConfig) -> AbstractDeployment: diff --git a/rock/sdk/sandbox/config.py b/rock/sdk/sandbox/config.py index fcd9d426ad..130aaf40a8 100644 --- a/rock/sdk/sandbox/config.py +++ b/rock/sdk/sandbox/config.py @@ -44,7 +44,14 @@ class SandboxConfig(BaseConfig): registry_password: str | None = None use_kata_runtime: bool = False sandbox_id: str | None = None - auto_delete_seconds: int = 0 + auto_delete_seconds: int | None = None + + @field_validator("auto_delete_seconds") + @classmethod + def validate_auto_delete_seconds(cls, v): + if v is not None and v < 0: + raise ValueError("auto_delete_seconds must be >= 0") + return v class SandboxGroupConfig(SandboxConfig): diff --git a/tests/unit/deployments/test_get_deployment.py b/tests/unit/deployments/test_get_deployment.py index b975e75d8b..4b8dd47368 100644 --- a/tests/unit/deployments/test_get_deployment.py +++ b/tests/unit/deployments/test_get_deployment.py @@ -32,3 +32,23 @@ async def test_deployment_manager(rock_config): docker_deployment_config = await manager.init_config(config) deployment = manager.get_deployment(docker_deployment_config) assert isinstance(deployment, RayDeployment) + + +class TestDeploymentManagerAutoDeleteSeconds: + async def test_auto_delete_seconds_none_uses_rock_config(self, rock_config): + manager = DeploymentManager(rock_config) + config = DockerDeploymentConfig(auto_delete_seconds=None) + result = await manager.init_config(config) + assert result.remove_container == rock_config.sandbox_config.remove_container_enabled + + async def test_auto_delete_seconds_zero_sets_remove_container_true(self, rock_config): + manager = DeploymentManager(rock_config) + config = DockerDeploymentConfig(auto_delete_seconds=0) + result = await manager.init_config(config) + assert result.remove_container is True + + async def test_auto_delete_seconds_positive_sets_remove_container_false(self, rock_config): + manager = DeploymentManager(rock_config) + config = DockerDeploymentConfig(auto_delete_seconds=300) + result = await manager.init_config(config) + assert result.remove_container is False diff --git a/tests/unit/sdk/test_sandbox_config.py b/tests/unit/sdk/test_sandbox_config.py new file mode 100644 index 0000000000..1480965f91 --- /dev/null +++ b/tests/unit/sdk/test_sandbox_config.py @@ -0,0 +1,30 @@ +import pytest +from pydantic import ValidationError + +from rock.sdk.sandbox.config import SandboxConfig + + +class TestSandboxConfigAutoDeleteSeconds: + def test_default_is_none(self): + config = SandboxConfig() + assert config.auto_delete_seconds is None + + def test_none_is_valid(self): + config = SandboxConfig(auto_delete_seconds=None) + assert config.auto_delete_seconds is None + + def test_zero_is_valid(self): + config = SandboxConfig(auto_delete_seconds=0) + assert config.auto_delete_seconds == 0 + + def test_positive_value_is_valid(self): + config = SandboxConfig(auto_delete_seconds=300) + assert config.auto_delete_seconds == 300 + + def test_negative_value_raises_error(self): + with pytest.raises(ValidationError, match="auto_delete_seconds must be >= 0"): + SandboxConfig(auto_delete_seconds=-1) + + def test_large_negative_value_raises_error(self): + with pytest.raises(ValidationError, match="auto_delete_seconds must be >= 0"): + SandboxConfig(auto_delete_seconds=-100) From c96c1b983dcc4176a814387c4e80f392bf6c82de Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 10 Apr 2026 12:27:11 +0800 Subject: [PATCH 021/226] Release note v1.5.0 (#760) * docs: add v1.5.x versioned docs with v1.5.0 release note Co-Authored-By: Claude Opus 4.6 * v1.5.0 Release Note Co-Authored-By: Claude Opus 4.6 * Bump pyproject.toml version to 1.5.0 Signed-off-by: Jiachen Zhang * feat: update v1.5.0 release note Signed-off-by: Jiachen Zhang --------- Signed-off-by: Jiachen Zhang Co-authored-by: Claude Opus 4.6 --- docs/docusaurus.config.js | 2 +- .../version-1.5.x.json | 34 ++ .../Getting Started/installation.md | 141 +++++++++ .../Getting Started/quickstart.md | 172 ++++++++++ .../Getting Started/rock-agent.md | 73 +++++ .../version-1.5.x/Getting Started/rockroll.md | 200 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 69 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 137 ++++++++ .../Python SDK References/sandbox.md | 113 +++++++ .../swe-bench-evaluation.md | 228 ++++++++++++++ .../version-1.5.x/References/api.md | 194 ++++++++++++ .../version-1.5.x/Release Notes/index.md | 5 + .../version-1.5.x/Release Notes/v1.5.0.md | 92 ++++++ .../User Guides/configuration.md | 188 +++++++++++ .../version-1.5.x/overview.md | 40 +++ .../Getting Started/installation.md | 143 +++++++++ .../Getting Started/quickstart.md | 166 ++++++++++ .../Getting Started/rock-agent.md | 72 +++++ .../version-1.5.x/Getting Started/rockroll.md | 194 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 70 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 136 ++++++++ .../Python SDK References/sandbox.md | 114 +++++++ .../swe-bench-evaluation.md | 229 ++++++++++++++ .../version-1.5.x/References/api.md | 195 ++++++++++++ .../version-1.5.x/Release Notes/index.md | 5 + .../version-1.5.x/Release Notes/v1.5.0.md | 92 ++++++ .../User Guides/configuration.md | 189 +++++++++++ docs/versioned_docs/version-1.5.x/overview.md | 33 ++ .../version-1.5.x-sidebars.json | 64 ++++ docs/versions.json | 1 + pyproject.toml | 2 +- 43 files changed, 5607 insertions(+), 2 deletions(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x.json create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/installation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/quickstart.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rockroll.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/codes.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/deploy.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/file_system.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/model-service.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/python_sdk.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/remote_user.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/runtime-env.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/sandbox.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/api.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/User Guides/configuration.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/overview.md create mode 100644 docs/versioned_docs/version-1.5.x/Getting Started/installation.md create mode 100644 docs/versioned_docs/version-1.5.x/Getting Started/quickstart.md create mode 100644 docs/versioned_docs/version-1.5.x/Getting Started/rock-agent.md create mode 100644 docs/versioned_docs/version-1.5.x/Getting Started/rockroll.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/codes.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/deploy.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/file_system.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/model-service.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/python_sdk.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/remote_user.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/rock-agent.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/runtime-env.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/sandbox.md create mode 100644 docs/versioned_docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/versioned_docs/version-1.5.x/References/api.md create mode 100644 docs/versioned_docs/version-1.5.x/Release Notes/index.md create mode 100644 docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md create mode 100644 docs/versioned_docs/version-1.5.x/User Guides/configuration.md create mode 100644 docs/versioned_docs/version-1.5.x/overview.md create mode 100644 docs/versioned_sidebars/version-1.5.x-sidebars.json diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index 97d7c3911b..12e74b8659 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -144,7 +144,7 @@ const config = { // release note按照版本号倒排 return reverseReleaseNoteSidebars(filterHiddenSidebars); }, - lastVersion: '1.3.x', + lastVersion: '1.5.x', includeCurrentVersion: false, versions: convertVersionsArrayToObject(versions) }, diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x.json b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x.json new file mode 100644 index 0000000000..5b7da62bdb --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x.json @@ -0,0 +1,34 @@ +{ + "version.label": { + "message": "1.5.x", + "description": "The label for version 1.5.x" + }, + "sidebar.tutorialSidebar.category.Getting Started": { + "message": "快速上手", + "description": "The label for category 'Getting Started' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.User Guides": { + "message": "用户指南", + "description": "The label for category 'User Guides' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.References": { + "message": "参考", + "description": "The label for category 'References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Release Notes": { + "message": "版本说明", + "description": "The label for category 'Release Notes' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.model-service": { + "message": "Model Service 参考", + "description": "The label for category 'Model Service References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.sandbox-agent": { + "message": "Sandbox Agent参考", + "description": "The label for category 'Sandbox Agent References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Python SDK References": { + "message": "Python SDK 参考", + "description": "The label for category 'Python SDK References' in sidebar 'tutorialSidebar'" + } +} diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/installation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/installation.md new file mode 100644 index 0000000000..0ab70e55d1 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/installation.md @@ -0,0 +1,141 @@ +--- +sidebar_position: 3 +--- + +# 安装指南 + +本文档介绍如何使用 `uv` 和 `pip` 安装和设置 ROCK 开发环境。该项目是一个强化学习开放构建工具包，支持多种组件。 + +## 使用 uv（推荐） + +### 快速安装所有依赖 + +```bash +# 安装所有依赖（包括可选依赖） +uv sync --all-extras + +# 安装开发/测试依赖 +uv sync --all-extras --all-groups +``` + +### 安装不同依赖组 + +#### 仅核心依赖 +```bash +uv sync +``` + +#### 管理组件依赖 +```bash +uv sync --extra admin +``` + +#### Rocklet 执行环境依赖 +```bash +uv sync --extra rocklet +``` + +#### 所有依赖 +```bash +uv sync --all-extras +``` + +#### 开发/测试依赖 +```bash +uv sync --all-extras --group test +``` + +## 使用 pip + +### 从 pip 源安装 + +#### 仅核心依赖 +```bash +pip install rl-rock +``` + +#### 管理组件依赖 +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install "rl-rock[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install "rl-rock[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install "rl-rock[all]" +``` + +### 使用 pip 从源码安装 + +#### 仅核心依赖 +```bash +pip install . +``` + +#### 管理组件依赖 +```bash +pip install ".[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install ".[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install ".[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install ".[all]" +``` + +## 可用入口点 + +该包提供以下命令行脚本： + +- `rocklet`: ROCK 执行环境服务器 (rock.rocklet.server:main) +- `admin`: 管理服务器 (rock.admin.main:main) +- `envhub`: 环境中心服务器 (rock.envhub.server:main) +- `rock`: 主 ROCK 命令行接口 (rock.cli.main:main) + +## 开发设置 + +### 使用 uv（推荐） + +```bash +# 克隆并设置开发环境 +git clone +cd ROCK +uv sync --all-extras --group test + +# 运行测试 +uv run pytest + + +### 使用 pip + +```bash +# 开发模式安装所有可选依赖 +pip install -e ".[all]" + +# 分别安装 +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" +``` + +## 附加说明 + +- 项目配置为默认使用阿里云 PyPI 镜像: `https://mirrors.aliyun.com/pypi/simple/` +- 对于本地开发，运行测试需要 `test` 依赖组 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/quickstart.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..e0a0891c0e --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/quickstart.md @@ -0,0 +1,172 @@ +--- +sidebar_position: 2 +--- + +# 快速上手 + +本指南将通过完整的示例演示如何使用 ROCK 创建和管理强化学习环境。ROCK (Reinforcement Open Construction Kit) 是一个全面的沙箱环境管理框架，主要用于强化学习和AI开发环境。 + +## 1. 环境准备 + +我们推荐在 Linux 系统下启动 ROCK，能够尽量复用项目依赖，提升环境拉起速度。如果需要在 macOS 上尝试，可以参考 [MacOS 启动](#7-macos-启动) 一节。 + +在开始之前，请确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且示例中依赖python:3.11镜像 +docker pull python:3.11 + +# 验证 uv 安装 +uv --version + + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone +cd ROCK + +# 创建虚拟环境（使用 uv 托管的 Python, 以python 3.11 版本为例） +uv venv --python 3.11 --python-preference only-managed + +# 安装所有依赖组 +uv sync --all-extras +``` + +> **重要提示**: 为确保 ROCK 能正确挂载项目和虚拟环境及其依赖的 base Python 解释器，强烈推荐使用 uv 托管的 Python 环境而非系统 Python。 + +## 2. 激活虚拟环境 + +在运行任何 ROCK 命令之前，需要先激活虚拟环境。确保 sys.base_prefix 是 uv 管理的环境，类似于 `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` 等路径。 + +```bash +# 激活虚拟环境 +source .venv/bin/activate + +# 验证 Python 环境 +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **验证要点**: 确保输出的 base prefix 路径指向 uv 管理的 Python 环境，而非系统 Python。 + +## 3. 验证环境配置 + +激活虚拟环境后，验证依赖安装是否正确： + +```bash +# 检查关键依赖 +python -c "import rock; print(\"Hello ROCK\")" +``` + + +## 4. 启动 ROCK 服务 + +激活虚拟环境后，在项目根目录下，启动 ROCK Admin 服务： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 启动 ROCK Admin 服务（本地环境） +rock admin start +``` + +服务启动后，您将看到类似以下的输出： + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **服务说明**: ROCK Admin 服务默认运行在 `http://127.0.0.1:8080`。 + +## 5. 运行示例环境 + +现在可以运行示例环境来验证安装。确保 ROCK 服务正在运行，然后打开一个新的终端窗口执行以下命令： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 运行沙箱示例 +python examples/sandbox_demo.py + +# 运行 GEM 协议示例 +python examples/sokoban_demo.py +``` + +### 5.1 示例说明 + +- **sandbox_demo.py**: 演示如何使用 ROCK 的沙箱 SDK 创建和管理容器化环境 +- **sokoban_demo.py**: 演示如何使用 ROCK 的 GEM 协议兼容接口创建强化学习环境 + +> **运行要求**: 确保 ROCK Admin 服务正在运行，因为示例需要与服务进行通信。 + +## 6. 分布式环境配置（可选） + +对于分布式多机器环境，请确保以下配置一致： + +1. 所有机器上 ROCK 和 uv 的 Python 配置使用相同的根 Python 解释器 +2. Docker 版本在所有节点上保持一致 +3. 网络配置允许各节点间正常通信 + + + +## 7. MacOS 启动 + +在 macOS 上，如果需要启动 Linux 镜像的环境，需要先设置环境变量： + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +在容器启动时，会安装对应的 uv 环境，细节可以参考 `rock/rocklet/local_files/docker_run_with_uv.sh` 脚本。 + +> **注意**: 相比 Linux 系统，macOS 上的启动速度会较慢，且比较依赖网络环境，可以根据实际情况调整脚本。ROCK_WORKER_ENV_TYPE的细节可以参考 [Configuration Guide](../User%20Guides/configuration.md). + + +## 8. 从Pip源启动 + +如果从Pip源启动Admin Server，在参照[安装指南](./installation.md)安装完成ROCK后, 需要设置额外环境变量: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(这一启动方式在容器环境启动时会从Pypi源上拉取最新的rocklet并安装, 相对启动速度比较慢, 仅推荐测试使用, 生产上依旧推荐其他的启动方式) + + +## 总结 + +恭喜！您已经成功完成了 ROCK 的快速开始指南。现在您应该能够： + +- 正确设置 ROCK 开发环境 +- 使用 uv 管理的 Python 环境 +- 启动和管理 ROCK 服务 +- 运行示例程序验证安装 +- 在分布式环境中配置 ROCK（如果需要） + +如需深入了解 ROCK 的更多功能，请参考以下文档： + +## 下一步学习 + +- [配置指南](../User%20Guides/configuration.md) - 详细了解 ROCK 的配置选项 +- [API 文档](../References/api.md) - 查看完整的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +- [安装指南](./installation.md) - 详细了解 ROCK 安装和配置 +- [概览](../overview.md) - 了解 ROCK 的设计理念 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..bd2dd69b05 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rock-agent.md @@ -0,0 +1,73 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent 快速启动 + +Rock Agent 是 ROCK 提供的 AI Agent 运行框架，支持在沙箱环境中运行各种类型的 Agent。 + +## 前置条件 + +- 确保有可用的ROCK服务, 如果需要本地拉起服务端, 参考[快速启动](quickstart.md) + +## 使用示例 + +ROCK 提供了两个Hello World Agent 示例，位于 `examples/agents/` 目录下： + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent 示例 +└── iflow_cli/ # IFlowCli Agent 示例 +``` + +### 运行 IFlowCli 示例 + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### 运行 ClaudeCode 示例 + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli 配置文件 + +配置文件位于 `examples/agents/iflow_cli/rock_agent_config.yaml`： + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # 填入你的 API Key + IFLOW_BASE_URL: "" # 填入你的 Base URL + IFLOW_MODEL_NAME: "" # 填入你的模型名称 +``` + +## ClaudeCode 配置文件 + +配置文件位于 `examples/agents/claude_code/rock_agent_config.yaml`： + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # 填入你的anthropic base url + ANTHROPIC_API_KEY: "" # 填入你的anthropic api key +``` + +## 相关文档 + +- [RockAgent 参考](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rockroll.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..3b53810ba3 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Getting Started/rockroll.md @@ -0,0 +1,200 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL 快速开始指南 + +本指南将引导您使用 ROLL (训练框架) 和 ROCK (环境管理) 来运行一个基于 Sokoban 游戏（推箱子）的强化学习训练示例。 + +## 1. 单机环境准备 + +在开始之前，请先确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **操作系统**: 推荐使用 Linux (如 Ubuntu 20.04+) +- **硬件**: 建议使用 NVIDIA GPU 并安装对应的驱动程序 +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且可提前拉取 Sokoban 游戏环境镜像，避免训练时等待 +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# 验证 uv 安装 +uv --version + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# 确保两个仓库位于同一级目录下，如下所示： +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. 启动训练流程 + +> 说明：下文均以 *torch2.6.0 + vLLM0.8.4* 为例。 + + +### 方式一： 使用虚拟环境启动（推荐） + +#### 为什么推荐这种方式？ +- 隔离性：uv 虚拟环境能确保项目依赖与系统环境隔离，避免冲突。 +- 速度快：ROCK 可以复用此虚拟环境，大大加快了后续环境的启动速度。 +- 稳定性：依赖关系更清晰，环境更易复现。 + + +```bash +# 进入 ROCK 目录 +cd ROCK + +# 使用 uv 创建并激活 Python 3.10 虚拟环境（ROLL推荐使用Python 3.10） +uv venv --python 3.10 --python-preference only-managed + +# 激活虚拟环境 +source .venv/bin/activate + +# 使用uv安装ROCK的依赖 +uv sync --all-extras + +# 若使用Python 3.10, 启动 ray 时会报错：ValueError: is not a valid Sentinel +# 原因是 ray 与 click>=8.3 版本不兼容，需要降级到 click<8.3 +# Python 3.11 不会有这个问题 +uv pip install 'click>=8.2,<8.3' + +# 切换到 ROLL 目录以安装其依赖 +cd ../ROLL + +# 设置国内 PyPI 镜像源以加速下载 +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装核心 PyTorch 组件 +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -i $PYPI_MIRROR + +# 安装transformer-engine，--no-build-isolation 避免因环境隔离导致找不到 torch +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation -i $PYPI_MIRROR + +# 安装预编译的 flash-attention，以匹配特定的 CUDA 和 PyTorch 版本 +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# 安装其余依赖 +uv pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# (可选) 安装Tensorboard，用于查看训练指标 +uv pip install tensorboard -i $PYPI_MIRROR + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### 方式二：使用系统环境启动（备选方案） + +为获得最佳兼容性，推荐使用 ROLL 官方提供的基础 Docker 镜像，因为它们已经预装了匹配的 CUDA、cuDNN 和其他基础库。 + +> [ROLL 官方镜像列表](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Installation/image_address/) + + +#### 注意 +此方式会将所有 Python 包直接安装到您的当前环境（例如，容器的基础环境）中，可能会与系统自带的包或其他项目产生冲突。 + +由于 ROCK 无法复用环境，每次启动任务时都可能需要重新安装部分依赖，启动速度较慢且受网络影响。 + + +```bash +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装ROCK的依赖 +cd ROCK +pip install . -i $PYPI_MIRROR +pip install ".[admin]" -i $PYPI_MIRROR + +# 安装ROLL的依赖 +cd ../ROLL +pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# 配置ROCK用uv启动的环境变量 +export ROCK_WORKER_ENV_TYPE=uv + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +至此，您已成功启动了 Sokoban 强化学习训练流程。祝您 Rock & Roll 愉快！ + + +## 3. 多机部署 + +除了在单机上运行，您也可以将 **ROCK 服务** 和 **ROLL 训练** 部署在不同的机器上，通过网络进行通信。这是一种常见的服务化部署模式。 + +### 3.1 在机器 A 上部署 ROCK 服务 + +在一台独立的机器（或容器）上，参照[ROCK快速指南](./quickstart.md)部署并启动 ROCK 服务。 + +> **重要提示** +> 启动服务后，请记下ROCK服务的IP地址和端口，例如`http://192.168.1.10:8000`，后续步骤将需要这个地址。 + +### 3.2 在机器 B 上准备 ROLL 客户端 + +在另一台将要运行训练任务的机器上，执行以下操作。 + +1. 验证网络连通性 + +首先，使用 curl 命令检查是否能从机器 B 访问到机器 A 上的 ROCK 服务。 +```bash +# 将 : 替换为您的 ROCK 服务实际地址 +# 如果成功，会收到 ROCK 服务的响应 {"message":"hello, ROCK!"} +curl http://: +``` + +2. 准备 ROLL 环境 + +```bash +# 克隆 ROLL 仓库 +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# 安装依赖 +pip install -r requirements_torch260_vllm.txt -i https://mirrors.aliyun.com/pypi/simple/ +``` + +3. 配置 ROLL 连接地址 + +修改 ROLL 的配置文件，使其能够找到并连接到远程的 ROCK 服务。 +- 打开配置文件：examples/agentic_demo/agentic_val_sokoban_sandbox.yaml +- 找到 SokobanSandbox 下的 env_config 部分 +- 将 base_url 的值修改为您的 ROCK 服务地址 +```yaml +custom_envs: + SokobanSandbox: + env_config: + # 将这里的地址修改为您的 ROCK 服务地址 + # 例如: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. 启动训练 +配置完成后，即可在机器 B 上启动 ROLL 训练脚本。 + +```bash +# 此脚本现在会通过网络请求机器 A 上的 ROCK 服务来创建环境 +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### 进阶：分布式 ROLL 训练 + +如果您希望将 ROLL 训练任务本身进行分布式部署，可以参考 ROLL 的官方分布式部署文档。 +> [快速上手：多节点部署指南](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/codes.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..47b74166de --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +错误码定义和分类，用于错误处理和重试策略。 + +## 使用示例 + +```python +import rock + +def test_codes_values(): + """测试基本状态码值""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes 分类 + +```python +OK = 2000, "OK" +""" +成功状态码 (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +客户端错误码 (4xxx): + +这些错误表示客户端请求有问题， +SDK 会抛出异常。 +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +服务端错误码 (5xxx): + +这些错误表示服务端出现问题， +SDK 会抛出异常。 +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +命令/执行错误码 (6xxx): + +这些错误与命令执行相关，由模型处理， +SDK 不会抛出异常。 +""" +``` + +## 重试策略建议 + +- **重试触发条件**: 只有当 `INTERNAL_SERVER_ERROR` 时才需要重试 +- **其他情况的处理策略**: + - `BAD_REQUEST`: 需要检查 arun 调用逻辑是否有异常 + - `COMMAND_ERROR`: stdout 输出到 `observation.output`，stderr 输出到 `observation.failure_reason` +- `COMMAND_ERROR` 说明: 由于 bash 执行失败时，stdout/stderr 可能全部非空，建议将 observation 中 output 和 failure_reason 全部 prompt 给模型进行推理 + +## 重试示例 + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/deploy.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..b7bd2da08f --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +沙箱资源部署管理器，用于本地目录部署和模板格式化。 + +## deploy_working_dir - 部署本地目录 + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# 部署本地目录到沙箱（自动生成目标路径） +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"部署到: {target}") # 例如: /tmp/rock_workdir_abc123 + +# 部署到指定目标路径 +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - 模板变量替换 + +`format` 方法支持两种模板语法： + +- **`${variable}`** - 标准 Python 字符串模板语法 +- **`<>`** - 替代语法（内部转换为 `${variable}`） + +```python +# 使用 ${working_dir} 模板变量 +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# 结果: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# 使用 <<>> 替代语法 +cmd = deploy.format("cat <>/file.txt") +# 结果: cat /tmp/rock_workdir_abc123/file.txt + +# 结合自定义变量使用 +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# 结果: cat /tmp/rock_workdir_abc123/settings.json + +# Shell 语法保持不变 +cmd = deploy.format("echo $((3 << 2 >> 1))") +# 结果: echo $((3 << 2 >> 1)) + +# 直接访问 working_dir +if deploy.working_dir: + print(f"当前工作目录: {deploy.working_dir}") +``` + +## 多次部署 + +后续调用会覆盖之前的工作目录路径： + +```python +# 第一次部署 +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# 第二次部署（覆盖之前的路径） +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/file_system.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..741b14f5e4 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +文件系统操作接口，提供沙箱环境中的权限管理和目录上传功能。 + +## chown - 修改所有者 + +```python +from rock.actions.sandbox.request import ChownRequest + +# 创建远程用户后修改所有者 +await sandbox.remote_user.create_remote_user("deploy") + +# 获取当前目录 +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# 修改目录所有者 +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# 递归修改目录及其内容所有者 +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - 修改权限 + +```python +from rock.actions.sandbox.request import ChmodRequest + +# 创建测试目录 +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# 修改目录权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# 递归修改权限（包括子目录和文件） +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# 设置最高权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - 上传目录 + +```python +import os +from pathlib import Path + +# 准备本地目录 +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# 上传到沙箱 +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"上传成功: {result.output}") +else: + print(f"上传失败: {result.failure_reason}") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/model-service.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..ba158cf75a --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service（实验性） + +ROCK 提供的 Model Service 负责处理 AI 模型调用的通信，为代理(Agent)和训练框架(如 Roll)或实际的 LLM 推理服务之间提供通信桥梁。 + +## 与 RockAgent 集成 + +ModelService 通常由 **RockAgent** 自动管理，无需手动调用生命周期方法。只需在配置中启用即可： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # 启用 ModelService，RockAgent 会自动管理其生命周期 +) +``` + +RockAgent 会自动： +- 安装 ModelService（安装 Python 运行时环境、安装模型服务包） +- 启动/停止 ModelService +- 监控 Agent 进程 + +## 架构概述（Local 模式） + +Local 模式下，模型服务使用**文件系统**作为通信媒介，实现代理和模型间的请求-响应机制。 + +当 Agent 需要调用模型时，请求首先写入日志文件，然后由负责监听的组件处理响应。当模型生成响应后，结果将写回日志文件，并由等待的 Agent 读取。 + +## anti_call_llm - 核心 API + +`anti_call_llm()` 是 **Local 模式**下最重要的 API，用于手动触发 LLM 反调用，实现模型调用的精细控制： + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='OpenAI type response', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +**使用场景：** +- Agent 捕获到 LLM 响应后，调用此方法通知 Roll 运行时 +- 支持携带响应数据，用于错误处理或重试 +- 超时和检查间隔可配置，适应不同网络环境 + +## CLI 命令 + +如果需要通过 CLI 使用模型服务，ROCK 提供了一个 CLI 命令集，可以在沙箱中安装 ROCK 后，通过 `rock model-service` 访问： + +### start 命令 +开始模型服务进程 +```bash +rock model-service start --type [local|proxy] [选项] +``` + +参数: + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `--type` | str | `local` | 服务类型：`local` 或 `proxy` | +| `--config-file` | str | None | 配置文件路径 | +| `--host` | str | None | 服务器地址（覆盖配置） | +| `--port` | int | None | 服务器端口（覆盖配置） | +| `--proxy-base-url` | str | None | 代理基础 URL | +| `--retryable-status-codes` | str | None | 可重试状态码，逗号分隔 | +| `--request-timeout` | int | None | 请求超时秒数 | + +### watch-agent 命令 +监控代理进程，当进程退出时发送 SESSION_END 消息 +```bash +rock model-service watch-agent --pid <进程ID> +``` + +参数: +- `--pid`: 需要监控的代理进程 ID + +### stop 命令 +停止模型服务 +```bash +rock model-service stop +``` + +### anti-call-llm 命令 +反调用 LLM 接口 +```bash +rock model-service anti-call-llm --index <索引> [--response <响应>] +``` + +参数: +- `--index`: 上一个 LLM 调用的索引，从 0 开始 +- `--response`: 上一次 LLM 调用的响应（可选） + +## 文件通信协议 + +模型服务使用文件进行进程间通信，定义了特定的标记格式用于区分请求和响应： + +### 请求格式 +``` +LLM_REQUEST_START{JSON请求数据}LLM_REQUEST_END{元数据JSON} +``` + +### 响应格式 +``` +LLM_RESPONSE_START{JSON响应数据}LLM_RESPONSE_END{元数据JSON} +``` + +### 会话结束标识 +``` +SESSION_END +``` + +元数据包含时间戳和索引信息，用于保证消息顺序和处理。 + +## SDK 使用 + +### ModelServiceConfig + +模型服务配置类，位于 `rock/sdk/sandbox/model_service/base.py`： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # 服务类型 + install_cmd="pip install rock-model-service", # 安装命令 + install_timeout=300, # 安装超时（秒） + start_cmd="rock model-service start --type ${type}", # 启动命令 + stop_cmd="rock model-service stop", # 停止命令 + logging_path="/data/logs", # 日志路径 + logging_file_name="model_service.log", # 日志文件名 +) +``` + +| 配置项 | 默认值 | 说明 | +|--------|--------|------| +| `enabled` | `False` | 是否启用模型服务（RockAgent 自动管理） | +| `type` | `"local"` | 服务类型：`local` 或 `proxy` | +| `install_cmd` | - | 模型服务包安装命令 | +| `install_timeout` | `300` | 安装超时时间（秒） | +| `start_cmd` | - | 启动命令模板 | +| `stop_cmd` | - | 停止命令 | +| `logging_path` | `/data/logs` | 日志目录路径 | +| `logging_file_name` | `model_service.log` | 日志文件名 | + +### ModelService + +模型服务管理类，处理沙箱内模型服务的生命周期： + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# 通常由 RockAgent 自动管理，无需手动调用 +# 以下方法仅在需要手动控制时使用 + +# 安装模型服务 +await model_service.install() + +# 启动模型服务 +await model_service.start() + +# 监控代理进程 +await model_service.watch_agent(pid="12345") + +# 执行反调用 LLM（Local 模式核心 API） +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# 停止模型服务 +await model_service.stop() +``` + +## API 参考 + +### install() + +在沙箱中安装模型服务依赖。 + +```python +await model_service.install() +``` + +执行步骤： +1. 创建并初始化 Python 运行时环境 +2. 创建 Rock 配置文件 +3. 安装模型服务包 + +**注意：** 通常由 RockAgent 自动调用。 + +### start() + +启动模型服务。 + +```python +await model_service.start() +``` + +前提条件：必须先调用 `install()`。 + +**注意：** 通常由 RockAgent 自动调用。 + +### stop() + +停止模型服务。 + +```python +await model_service.stop() +``` + +如果服务未运行，会跳过此操作。 + +**注意：** 通常由 RockAgent 自动调用。 + +### watch_agent(pid) + +监控代理进程。 + +```python +await model_service.watch_agent(pid="12345") +``` + +当进程退出时，发送 `SESSION_END` 消息。 + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +执行反调用 LLM 操作。**这是 Local 模式下最重要的 API。** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='{"result": "..."}', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +## 配置选项 + +### 服务配置 +- `SERVICE_HOST`: 服务主机地址，默认为 `"0.0.0.0"` +- `SERVICE_PORT`: 服务端口，默认为 `8080` + +### 日志配置 +- `LOG_FILE`: 用以通信的日志文件路径，包含请求和响应数据 + +### 轨迹（Traj）日志记录 +模型服务将 LLM 调用轨迹（traj）记录到 JSONL 文件中，用于调试和分析。 + +| 环境变量 | 默认值 | 说明 | +|----------|--------|------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | traj 日志文件目录 | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | 追加模式（true/false） | + +**traj 文件位置**: `{DATA_DIR}/LLMTraj.jsonl` + +**traj 文件格式**（JSONL - 每行一个 JSON 对象）： +```json +{"request": {...}, "response": {...}} +``` + +### 轮询配置 +- `POLLING_INTERVAL_SECONDS`: 轮询间隔，默认为 `0.1` 秒 +- `REQUEST_TIMEOUT`: 请求超时时间，默认为无限 + +### 标记配置 +定义了用于区分日志文件中不同类型消息的标记： +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig（服务端） + +服务端配置类定义了模型服务如何处理请求： + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `host` | str | `"0.0.0.0"` | 服务器地址 | +| `port` | int | `8080` | 服务器端口 | +| `proxy_base_url` | str \| None | `None` | 直接代理 URL | +| `proxy_rules` | dict | 见下方 | 模型名称到 URL 的映射 | +| `retryable_status_codes` | list[int] | `[429, 500]` | 可重试的 HTTP 状态码 | +| `request_timeout` | int | `120` | 请求超时时间（秒） | + +**默认 proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/python_sdk.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..c1083f29b0 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK 参考 + +本指南详细介绍如何使用 ROCK SDK 进行开发，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK SDK为开发者提供了便捷的Python接口来使用ROCK平台的功能，包括沙箱环境管理和GEM环境交互。 + +> **重要提示**: 使用 SDK 之前，请确保 ROCK Admin 服务正在运行。可以通过以下命令启动： +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 基本沙箱操作 + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 沙箱组管理 + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# 创建沙箱组配置 +config = SandboxGroupConfig( + image="python:3.11", + size=4, # 创建4个沙箱 + start_concurrency=2, # 并发启动级别为2 +) + +# 创建并启动沙箱组 +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# 批量操作 +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# 批量停止 +await sandbox_group.stop() +``` + +### 2.3 配置示例 + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 沙箱加速配置 + +ROCK 提供沙箱网络加速功能，支持配置 APT、PIP 和 GitHub 镜像源，提升受限网络环境下的包下载速度。 + +#### 支持的加速类型 + +**APT 镜像配置** + +配置 APT 包管理器镜像源，加速 Debian/Ubuntu 软件包下载。 + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# 配置 APT 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP 镜像配置** + +配置 Python 包索引镜像，加速 pip 安装。 + +```python +# HTTP 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub 加速** + +通过添加自定义 DNS 解析条目加速 GitHub 访问。 + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### 完整示例 + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """创建沙箱并配置加速""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # 配置加速（在安装包之前配置） + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # speedup 不会主动安装 PIP，仅配置镜像源进行加速 + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # 可以通过镜像 IP 加速 GitHub 访问 + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### 注意事项 + +1. **配置顺序**: 在安装包之前配置加速 +2. **HTTPS vs HTTP**: HTTPS 镜像不需要为 PIP 配置 trusted-host +3. **GitHub IP**: 不同区域可能需要不同的 IP 以获得最佳性能 +4. **持久性**: 配置在沙箱生命周期内持久有效 +5. **多次调用**: 后续的加速调用会覆盖之前的配置 +6. **PIP 安装**: speedup 功能仅配置镜像源，不会自动安装 PIP + +## 3. GEM SDK + +### 3.1 Python SDK 方式 + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## 相关文档 +- [快速开始指南](../../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK SDK +- [API 文档](../api.md) - 查看 SDK 封装的底层 API 接口 +- [配置指南](../../User%20Guides/configuration.md) - 了解 SDK 相关的配置选项 +- [安装指南](../../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/remote_user.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..791ca85fdd --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/remote_user.md @@ -0,0 +1,69 @@ +# Remote User + +远程用户管理，用于在沙箱中创建和管理用户。 + +## 使用示例 + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +创建远程用户。 + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +检查用户是否存在。 + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..c3f03b1efc --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent（实验性） + +RockAgent 是 ROCK 框架中的核心 Agent 实现，直接继承自 `Agent` 抽象基类。它提供了完整的 Agent 生命周期管理，包括环境初始化、ModelService 集成、命令执行等功能。 + +使用 `sandbox.agent.install()` 以及 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 + +## 核心概念 + +RockAgent 的核心工作流程分为两个阶段： + +1. **install(config)**: 初始化 Agent 环境，包括部署工作目录、设置环境变量、初始化运行时环境等 +2. **run(prompt)**: 执行 Agent 任务，替换占位符并启动 Agent 进程 + +## 快速开始 + +### Claude Code 示例 + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli 示例 + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} 必须 + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # 环境变量 + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent 示例 + +```yaml +working_dir: "." # 上传包含 langgraph_agent.py 的本地当前目录到 sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # 运行本地脚本 + +runtime_env_config: + type: python + pip: # 安装 pip 依赖 + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## 配置详解 + +### 基础配置 + +```yaml +agent_type: "default" # Agent 类型标识（默认: "default"） +agent_name: "demo-agent" # Agent 实例名称（默认: 随机 uuid） +version: "1.0.0" # 版本标识（默认: "default"） +instance_id: "instance-001" # 实例 ID（默认: "instance-id-<随机uuid>"） +agent_installed_dir: "/tmp/installed_agent" # Agent 安装目录（默认: "/tmp/installed_agent"） +agent_session: "my-session" # bash 会话标识（默认: "agent-session-<随机uuid>"） +env: # 环境变量（默认: {}） + OPENAI_API_KEY: "xxxxxxx" +``` + +### 工作目录配置 + +```yaml +working_dir: "./my_project" # 本地目录，上传到 sandbox（默认: None 不上传） +project_path: "/testbed" # sandbox 中工作目录，用于 cd（默认: None） +use_deploy_working_dir_as_fallback: true # project_path 为空时是否回退到 deploy.working_dir（默认: true） +``` + +### 执行配置 + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent 执行命令，必须包含 ${prompt}（默认: None） + +skip_wrap_run_cmd: false # 跳过为 run_cmd 添加 PATH 的包装（默认: false） + +# 超时配置 +agent_install_timeout: 600 # 安装超时，单位秒（默认: 600） +agent_run_timeout: 1800 # 运行超时，单位秒（默认: 1800） +agent_run_check_interval: 30 # 检查间隔，单位秒（默认: 30） +``` + +**`skip_wrap_run_cmd`**: +- `false`（默认）：为命令添加 `export PATH=:$PATH &&` 包装，确保使用运行时环境的可执行文件 +- `true`：跳过 PATH 包装，直接使用 `bash -c` 运行命令 + +### 初始化钩子 + +```yaml +pre_init_cmds: # 初始化前执行的命令（默认: 从 env_vars 读取） + - command: "apt update && apt install -y git" + timeout_seconds: 300 # 命令超时，单位秒（默认: 300） + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # 初始化后执行的命令（默认: []） + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**注意事项**： +- `pre_init_cmds` 和 `post_init_cmds` 不继承 Agent 的 `env` 环境变量 +- 通常用于执行安装操作和配置文件移动操作 +- 常用命令示例： + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv 配置 + +```yaml +runtime_env_config: # 具体参考 RuntimeEnv 有关文档 + type: "python" # 运行时类型: python / node（默认: "python"） + version: "3.11" # 版本号 + pip: # Python 依赖包列表 + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node 运行时示例**： + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**自动执行的操作**： +- 根据 `type` 安装对应的运行时（Python 或 Node.js） +- 安装 `pip` 依赖（如果配置了） +- 执行 `custom_install_cmd` 自定义安装命令（如果配置了） +- 支持 `npm_registry` 配置 Node.js 的 npm 镜像源 + +### ModelService 配置 + +```yaml +model_service_config: # 具体参考 ModelService 有关文档 + enabled: true # 启用 ModelService（默认: false） +``` + +**自动执行的操作**： +- 安装阶段：安装 ModelService（仅安装，不启动） +- 运行阶段：启动 ModelService + `watch_agent` 监控进程 + +**注意事项**：需要将模型请求的 URL 设置为 ModelService 的 URL。例如 ModelService 提供的 OpenAI-compatible 的 URL 为 `http://127.0.0.1:8080/v1/chat/completions`，则通常需要将 Agent 向 LLM 请求的 URL 设置为 `http://127.0.0.1:8080/v1/`。 + +## API 参考 + +### install(config) + +初始化 Agent 环境。 + +**执行流程**： +1. 如果配置了 `working_dir`，部署到 sandbox +2. 设置 bash session，以及配置 env 环境变量 +3. 执行 `pre_init_cmds` +4. 并行初始化 RuntimeEnv 和 ModelService（如果启用） +5. 执行 `post_init_cmds` + +**参数**： +- `config`: Agent 配置文件，支持两种传入方式： + - **字符串路径**: YAML 配置文件路径，默认值为 `"rock_agent_config.yaml"` + - **RockAgentConfig 对象**: 直接传入 `RockAgentConfig` 实例 + +### run(prompt) + +执行 Agent 任务。 + +**执行流程**： +1. 替换占位符, 准备Agent 运行命令 +4. 启动 agent 进程 +5. 如果启用 ModelService，启动 `watch_agent` +6. 等待任务完成并返回结果 + +## 高级用法 + +### working_dir 与 project_path 的区别与联动 + +| 配置项 | 作用 | 联动方式 | +|--------|------|----------| +| `working_dir` | 本地目录，上传到 sandbox | 调用 `deploy.deploy_working_dir()` 上传，上传后 `deploy.working_dir` 变为 sandbox 中的路径 | +| `${working_dir}` | 命令中的占位符 | 被 `deploy.format()` 替换为 `deploy.working_dir` 的值，会在配置中的 init_cmds 和 run_cmd 中替换 | +| `project_path` | sandbox 中的工作目录 | 用于运行前 `cd project_path`，不设置时会进入到 `deploy.working_dir` 工作目录 | +| `use_deploy_working_dir_as_fallback` | run 时 project_path 未设置时是否回退到 deploy.working_dir | 默认为 `true`，设为 `false` 时即使未设置 project_path 也不会进入 working_dir | + +**使用建议**： +- 使用 `working_dir` 上传本地项目代码到 sandbox +- 使用 `project_path` 指定 sandbox 中的工作目录（如 `/testbed`） +- 设置 `use_deploy_working_dir_as_fallback: false` 的场景：需要进行本地文件挂载，但希望在镜像默认工作目录下运行 Agent + +### 占位符使用 + +Rock Agent 在支持在配置文件中替换以下占位符： + +- `${prompt}`: 在run_cmd 中必需，会被替换为 `run(prompt)` 传入的提示词 +- `${working_dir}`: 可选，会被替换为 sandbox 中实际的工作目录路径, 同时支持在 init_cmds和 run_cmd 中使用 +- `${bin_dir}`: 可选，会被替换为运行时环境的 bin 目录路径 + +**示例**： +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback 说明 + +当 `project_path` 未设置时： +- `true`（默认）：运行 Agent 前会自动 `cd` 到 `deploy.working_dir` +- `false`：运行 Agent 前不会自动切换目录，保持在当前目录 + +适用场景： +- `true`: 大多数场景，希望 Agent 在上传的代码目录中运行 +- `false`: 需要挂载本地文件，但希望在镜像默认工作目录（如 `/app, /testbed`）下运行 Agent + +## 完整配置示例 + +```yaml +# ========== 基础配置 ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== 工作目录配置 ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== 运行配置 ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# 超时配置 +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== 初始化命令 ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== 运行时环境配置 ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService 集成 ========== +model_service_config: + enabled: true +``` + +## 使用示例 + +### 使用 YAML 配置文件（推荐） + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/runtime-env.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..a5532e900b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,137 @@ +# RuntimeEnv + +RuntimeEnv 模块用于在沙箱中管理语言运行时环境（目前提供了 Python / Node.js）。 + +## 快速开始(使用示例) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +异步工厂方法，根据配置创建 RuntimeEnv 实例并初始化，自动注册到 `sandbox.runtime_envs`。 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# 自动注册，可通过 sandbox.runtime_envs[env.runtime_env_id] 访问 +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +包装命令，将 `bin_dir` 加入 PATH，确保优先使用运行时环境中的可执行文件。 + +```python +wrapped = env.wrapped_cmd("node script.js") +# 返回: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +在运行时环境中执行命令。内部基于 `wrapped_cmd` 实现 + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | 类型标识 | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python 版本，默认 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | pip 包列表或 requirements.txt 路径 | +| `pip_index_url` | `str \| None` | 环境变量 | pip 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | 要创建符号链接的可执行文件列表 | + +## NodeRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | 类型标识 | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node 版本，默认 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | 要创建符号链接的可执行文件列表 | + +## 自定义 RuntimeEnv 实现约束 + +自定义 RuntimeEnv 需遵循以下规则： + +1. **定义 `runtime_env_type` 类属性**：作为类型标识符，用于自动注册到 RuntimeEnv 工厂 +2. **重写 `_get_install_cmd()`**：返回安装命令 +3. **安装命令最后必须**：将目录重命名为 `runtime-env` + + +## NodeRuntimeEnv 简化版实现示例 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config 类：定义配置类型，用于 RuntimeEnv.create() 路由到对应实现 +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # 必须与 runtime_env_type 一致 + +# RuntimeEnv 实现类：定义如何安装和运行该运行时环境 +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # 自动注册到 RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # 下载 Node 二进制包并解压，最后重命名为 runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## 加速基础环境安装 + +`PythonRuntimeEnv` 默认从 https://github.com/astral-sh/python-build-standalone/releases/ 下载 Python 安装包。若网络不可达或下载较慢，可通过环境变量 `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` 或 `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` 覆盖默认安装命令（例如切换到内网源/镜像源）。 + +默认值示例： + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +例如，替换为镜像源下载： + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +请确保该命令执行完成后，会在 `runtime_env` 的默认工作目录下生成 `runtime-env` 目录，并且 `${workdir}/runtime-env/bin/` 下包含对应可执行文件，例如： + +- `${workdir}/runtime-env/bin/python` + +Node 环境同理，可通过修改环境变量 `ROCK_RTENV_NODE_V22180_INSTALL_CMD` 来指定更快的下载/安装命令。 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/sandbox.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..088f1e3110 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/sandbox.md @@ -0,0 +1,113 @@ +# 处理大文件和长命令输出 + +## `arun` +`arun()` 在 `nohup` 模式下提供了两个关键参数，帮助 Agent / 调用方在"执行"与"查看"之间按需解耦： + +1. **`response_limited_bytes_in_nohup`**（int 型） + 限制返回内容的最大字符数（例如 `64 * 1024`），适合仍需立刻查看部分日志、但必须控制带宽的场景。默认值 `None` 表示不加限制。 + +2. **`ignore_output`**（bool，默认 `False`） + 当设为 `True` 时，`arun()` 不再读取 nohup 输出文件，而是在命令执行完毕后立即返回一段提示信息（包含输出文件路径、**文件大小**及查看方式）。日志仍写入 `/tmp/tmp_.out`，后续可通过 `read_file`、下载接口或自定义命令按需读取，实现"执行"与"查看"彻底解耦。返回的文件大小信息可帮助用户决定是直接下载还是分块读取。 + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# 示例 1：限制最多 1024 个字符 +resp_limit = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# 示例 2：完全跳过日志读取，后续再通过 read_file / 下载获取 +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# 可通过 Sandbox.read_file(...) / 下载接口 / cat /tmp/tmp_xxx.out 查看日志 +``` + +## `read_file_by_line_range` + +按行范围异步读取文件内容，支持自动分块读取和会话管理，支持大文件读取。 + +### 重要特性 +- **大文件分块读取**: 自动将大文件分成多个小块进行读取 +- **自动统计行数**: 未指定结束行时，自动计算文件总行数 +- **内置重试机制**: 关键操作支持最多 3 次重试，提高可靠性 +- **参数验证**: 自动验证输入参数的合法性 +- **会话管理**: 支持指定会话或自动创建临时会话 + +### 参数说明 +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `file_path` | str | - | 要读取的文件路径（沙箱中的绝对路径或相对路径） | +| `start_line` | int \| None | 1 | 起始行号（从 1 开始） | +| `end_line` | int \| None | None | 结束行号（包含），默认为文件末尾 | +| `lines_per_request` | int | 1000 | 每次请求读取的行数，范围 1-10000 | + +### 返回值 +- `ReadFileResponse`: 包含文件内容的响应对象 + - `content` (str): 读取的文件内容 + +### 异常说明 +- `Exception`: 当 `start_line < 1` 时抛出 +- `Exception`: 当 `end_line < start_line` 时抛出 +- `Exception`: 当 `lines_per_request` 不在 1-10000 范围内时抛出 +- `Exception`: 当文件读取失败时抛出 + +### 使用示例 + +```python +# 读取整个文件 +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# 读取指定行范围（第 100 到 500 行） +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# 从第 1990 行读取到文件末尾 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# 使用自定义分块大小 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### 注意事项 +- 行号从 1 开始计数，而非 0 +- 对于大文件建议适当增加 `lines_per_request` 以提高效率 +- 文件路径必须是沙箱内的有效路径 +- 使用 `sed` 命令进行文件读取，确保沙箱镜像支持该命令 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..6a74f397c7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,228 @@ +# SWE-Bench 评测 + +本文档介绍如何使用 ROCK SDK 运行 SWE-Bench Verified 评测，包括沙箱启动、Agent 集成、测试环境准备和结果解析。 + +### 快速开始 +SWE-Bench-Verified 是一个用于评估 AI 编程 Agent 在真实软件工程任务上表现的基准测试。 + +在ROCK上运行一个SWE-Bench任务包含以下步骤： + +1. **load_task_config** — 加载 `task.yaml` 获取任务指令 +2. **start_sandbox** — 使用任务专属的 Docker 镜像启动沙箱 +3. **agent.install / agent.run** — 安装并运行 Agent 来解决任务 +4. **setup_test_env** — 上传测试文件和运行测试脚本到沙箱 +5. **运行测试** — 通过 `sandbox.arun()` 执行测试脚本，支持超时控制 +6. **parse_swebench_result** — 解析测试输出，判断 PASSED / FAILED +7. **sandbox.stop** — 清理沙箱资源 + +**下面是示例代码** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. 加载任务指令 + task_config = await load_task_config(task_dir) # 参见 load_task_config 章节 + instruction = task_config["instruction"] + + # 2. 启动沙箱 + sandbox = await start_sandbox(task_name) # 参见 start_sandbox 章节 + + try: + # 3. 安装并运行 Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. 准备测试环境 + await setup_test_env(sandbox, task_dir) # 参见 setup_test_env 章节 + + # 5. 运行测试 + resp = await run_tests(sandbox) # 参见"运行测试"章节 + + # 6. 解析结果 + is_resolved = parse_swebench_result(resp.output) # 参见 parse_swebench_result 章节 + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +以下章节详细介绍评测流程中使用的各个函数。 + +--- + +## start_sandbox + +使用任务专属的 SWE-Bench Docker 镜像启动沙箱实例。每个任务都有一个预构建的镜像，包含目标仓库和运行环境。 + +`image` 参数格式如下： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +例如，任务 `django__django-14539` 对应的镜像为： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +从任务目录中加载 `task.yaml` 配置文件。YAML 文件包含 `instruction` 字段，用于描述 Agent 需要完成的编程任务。 + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# 使用示例 +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +使用 `sandbox.agent.install()` 和 `sandbox.agent.run()` 在沙箱中部署和执行 Agent。详细的 Agent 配置请参考 [Rock Agent](./rock-agent.md)。 + +```python +# 使用 YAML 配置文件安装 Agent（以 iflow_config.yaml 为例） +await sandbox.agent.install(config="iflow_config.yaml") + +# 使用任务指令运行 Agent +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +在沙箱中准备测试环境：安装 [uv](https://github.com/astral-sh/uv) 包管理器，并上传测试文件和运行测试脚本。 + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """准备测试环境并返回会话名称。""" + # 1. 创建带有自定义环境变量的会话 + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. 安装 uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. 上传测试文件 + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. 上传运行测试脚本 + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## 运行测试 + +使用 `RunMode.NOHUP` 模式执行测试脚本，支持可配置的超时时间。 + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +解析测试输出以判断 SWE-Bench 任务是否通过。解析器会查找由标记行分隔的结果块，并检查是否包含 `PASSED`。 + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """解析 SWE-Bench 测试输出，判断任务是否通过。 + + 匹配 'SWEBench results starts here' 和 + 'SWEBench results ends here' 之间的内容块， + 然后检查其中是否包含 'PASSED'。 + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# 使用示例 +is_resolved = parse_swebench_result(resp.output) +``` + +## 注意事项 + +- **任务数据集**：任务目录（包含 `task.yaml`、`tests/` 和 `run-tests.sh`）可从 [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) 仓库获取。 +- **任务镜像**：每个 SWE-Bench 任务需要特定的 Docker 镜像（如 `sweb.eval.x86_64.`）。请确保镜像在对应的环境中可用。 +- **Agent 配置**：Agent 配置 YAML 定义了运行时、依赖和执行命令。详情请参考 [Rock Agent](./rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/api.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/api.md new file mode 100644 index 0000000000..06f44c326c --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/References/api.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 1 +--- + +# API 参考 + +本指南详细介绍 ROCK 平台提供的核心 API 服务，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK平台提供两种核心API服务： +- Sandbox API：沙箱环境管理 +- GEM API：GEM环境交互 + +所有 API 接口都遵循 RESTful 设计原则，支持 JSON 格式的数据交换。 + +## 2. Sandbox API + +沙箱环境全生命周期管理功能： + +### 沙箱管理接口 + +1. **Start Sandbox** - 启动沙箱环境 + - 创建一个新的沙箱实例 + - 支持指定镜像、资源配置等参数 + +2. **Start Sandbox Async** - 异步启动沙箱环境 + - 异步方式创建沙箱实例 + - 适用于需要快速响应的场景 + +3. **Check Sandbox Alive Status** - 检查沙箱存活状态 + - 验证沙箱是否正常运行 + +4. **Get Sandbox Statistics** - 获取沙箱统计信息 + - 获取沙箱的资源使用统计 + +5. **Get Sandbox Status** - 获取沙箱详细状态 + - 获取沙箱的完整状态信息 + +6. **Stop Sandbox** - 停止沙箱环境 + - 安全关闭沙箱实例 + +7. **Commit Sandbox** - 提交沙箱为镜像 + - 将当前沙箱状态保存为新镜像 + +### 命令执行接口 + +8. **Execute Command** - 在沙箱中执行命令 + - 直接在沙箱中运行指定命令 + +9. **Create Bash Session** - 创建Bash会话 + - 创建持久化的Bash会话环境 + +10. **Run Command in Session** - 在会话中执行命令 + - 在已创建的会话中执行命令 + +11. **Close Session** - 关闭会话 + - 释放会话资源 + +### 文件操作接口 + +12. **Read File** - 读取沙箱文件 + - 从沙箱中读取指定文件内容 + +13. **Write File** - 写入沙箱文件 + - 向沙箱中写入文件 + +14. **Upload File** - 上传文件到沙箱 + - 将本地文件上传到沙箱 + +## 3. GEM API + +GEM环境交互功能： + +1. **Make Environment** - 创建GEM环境 + - 初始化一个新的GEM环境实例 + +2. **Reset Environment** - 重置GEM环境 + - 将GEM环境重置到初始状态 + +3. **Step Environment** - 执行GEM环境步骤 + - 在GEM环境中执行一个动作步骤 + +4. **Close Environment** - 关闭GEM环境 + - 释放GEM环境资源 + +## 4. HTTP API 使用示例 + +### 4.1 Sandbox API 示例 + +#### 启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 异步启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### 创建会话 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### 在会话中执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### 上传文件 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### 停止沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API 示例 + +```bash +# 创建GEM环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# 重置环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# 执行步骤 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# 关闭环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK API +- [Python SDK 文档](./Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 调用 API +- [配置指南](../User%20Guides/configuration.md) - 了解 API 相关的配置选项 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md new file mode 100644 index 0000000000..44591bd640 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# 版本说明 +* [release v1.5.0](v1.5.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md new file mode 100644 index 0000000000..a4c00496cf --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md @@ -0,0 +1,92 @@ +# ROCK v1.5.0 Release Note + +# v1.5.0 + +## 发布日期 + +2026 年 4 月 10 日 + +--- + +## SDK + +### 新功能 + +#### TypeScript SDK + +* **新增**: 全新 TypeScript SDK，已发布到 npm (`rl-rock`)。支持沙箱管理、文件系统操作、同步/异步 Shell 命令执行、Python/Node.js 运行时环境、内置 Agent 支持、EnvHub 集成，以及 ESM/CommonJS 双构建 ([#492](https://github.com/alibaba/ROCK/pull/492)) + + +#### Agent  Run(Harbor Style) + +* **新增**: 支持 Agent Run 功能，允许通过 SDK 提交和管理 agent (harbor style)运行任务 ([#681](https://github.com/alibaba/ROCK/pull/681)) + + +#### Job OSS Artifact Mirror + +* **新增**: 新增 `OssMirrorConfig` 和 `JobConfig.enable_oss_mirror()` 支持 Job 的 OSS 产物镜像，自动从 sandbox config 填充 namespace/experiment\_id ([#708](https://github.com/alibaba/ROCK/pull/708)) + + +#### Sandbox Client + +* **新增：**auto\_delete\_seconds字段，表示sandbox stop之后的保留时间，目前支持以下功能 + + * auto\_delete\_seconds = None（默认值）, sandbox stop之后是否清除由集群配置决定 + + * auto\_delete\_seconds = 0, sandbox stop之后立即清除 + + * auto\_delete\_seconds > 0, sandbox stop 之后不删除 + + +--- + +## Admin + +### 新功能 + +#### Sandbox 元数据持久化 + +* **新增**: 支持将 sandbox 元数据持久化到数据库后端（除 Redis 外），提供更可靠的数据存储 ([#730](https://github.com/alibaba/ROCK/pull/730)) + + +--- + +### 重构 + +#### Ray 临时目录 + +* 在 `RayConfig` 中新增 `temp_dir` 字段，支持重定向 Ray 临时数据目录（默认 `.tmp/ray`），自动解析相对路径为绝对路径 ([#694](https://github.com/alibaba/ROCK/pull/694), [#696](https://github.com/alibaba/ROCK/pull/696)) + + +#### 用户自定义日志路径 + +* 支持用户自定义日志路径（如 `/data/logs/user-defined`），自动创建目录 ([#702](https://github.com/alibaba/ROCK/pull/702)) + + +--- + +### Bug 修复 + +#### 内存大小错误消息 + +* 修正 sandbox manager 中不正确的内存大小错误提示 ([#648](https://github.com/alibaba/ROCK/pull/648)) + + +#### 测试修复 + +* 修复无法运行的测试 ([#700](https://github.com/alibaba/ROCK/pull/700)) + +* 调整 conftest.py 中 sandbox 资源限制以修复始终失败的单元测试 ([#710](https://github.com/alibaba/ROCK/pull/710)) + +* 修复了多个 admin 实例之间 Kubernetes 缓存不一致导致 Sandbox 信息数据脏数据的问题 ([#743](https://github.com/alibaba/ROCK/pull/743)) + +* 修复了 admin 中 Kubernetes client informer 未及时处理事件的问题 ([#744](https://github.com/alibaba/ROCK/pull/744)) + + +--- + +## CI / 基础设施 + +* 恢复 CI 请求触发工作流配置 ([#728](https://github.com/alibaba/ROCK/pull/728)) + +* 锁定 `langgraph-prebuilt` 到 1.0.8 修复 CI 错误 ([#745](https://github.com/alibaba/ROCK/pull/745)) \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/User Guides/configuration.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/User Guides/configuration.md new file mode 100644 index 0000000000..a212189bc7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/User Guides/configuration.md @@ -0,0 +1,188 @@ +--- +sidebar_position: 4 +--- + +# 配置指南 + +本指南详细介绍如何配置 ROCK 环境以满足不同的使用需求，包括本地开发、测试和生产部署。 + +## 1. 环境变量配置 + +ROCK 支持通过环境变量配置关键参数。以下是主要的环境变量： + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK服务基础URL +export ROCK_LOG_LEVEL=INFO # 日志级别 +export ROCK_LOGGING_PATH=/path/to/logs # 日志文件路径，默认 None (输出到控制台) +export ROCK_LOGGING_FILE_NAME=rocklet.log # 日志文件名，默认 "rocklet.log", 启动admin时可以自定义日志文件名, 如admin.log +export ROCK_LOGGING_LEVEL=INFO # 日志输出级别，默认 "INFO" +export ROCK_WORKER_ENV_TYPE=local # 运行时环境类型，可选值: local, docker, uv, pip +``` + +更多环境变量可参考 `rock/env_vars.py` 文件。 + +### 1.1 运行时环境 (Runtime Environments) + +ROCK 提供了多种不同的运行时环境来满足不同场景的需求，选择通过环境变量 `ROCK_WORKER_ENV_TYPE` 进行配置。每种环境有不同的部署要求、性能特征和适用场景。每种环境都有其独特的优势和限制，开发者可以根据部署环境的需要选择最适合的运行时环境。 + +#### 1.1.1 Docker 运行时环境 + +Docker 运行时环境适用于已经预安装了所需依赖的 Docker 镜像环境。这种环境要求部署环境中直接可用 `/tmp/miniforge/bin/rocklet` 可执行文件。 + +**挂载配置:** +- `/tmp/miniforge` - 包含预安装的 Python 环境 +- `/tmp/local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 容器化部署环境 +- 已经构建了包含 `rocklet` 的自定义 Docker 镜像 +- 适合生产环境，启动速度快 + +**要求:** +- 需要使用定制的 Docker 镜像，其中包含 `/tmp/miniforge/bin/rocklet` 可执行文件 +- Docker 环境支持 + +#### 1.1.2 本地运行时环境 + +本地运行时环境直接利用当前部署环境的 Python 环境和项目文件。该环境要求宿主机和容器之间具有相同的操作系统，以便能够直接挂载虚拟环境和 Python 解释器。 + +**挂载配置:** +- `python_env_path` - Python 环境路径 +- `project_root` - 项目根目录 +- `.venv` - 虚拟环境目录（挂载为容器中的 `/tmp/miniforge`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 开发环境 +- 宿主机和目标容器使用相同操作系统的场景 +- 需要快速重新使用现有 Python 环境 + +**要求:** +- 相同的操作系统（主机/容器） +- 可直接访问当前部署的 `.venv` 虚拟环境 +- Python 解释器路径兼容 + +#### 1.1.3 UV 运行时环境 + +UV 运行时环境只依赖于可用的 ROCK 项目，但初始化相对较慢且网络要求较高。这种环境最适合没有预配置环境的场景。它从原始项目重新构建 rocklet 环境。这是推荐在 Mac 操作系统上使用的环境。 + +**挂载配置:** +- `project_root` - 项目根目录（挂载为容器中的 `/tmp + project_root`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**适用场景:** +- Mac 操作系统 +- 跨操作系统启动 +- 没有预配置环境的场景 +- 没有使用 uv 管理 Rock + +**优势:** +- 无需预构建镜像 +- 跨平台兼容性好 +- 特别适合开发和测试 + +**限制:** +- 初始化速度较慢 +- 网络要求较高 +- 启动时间较长 + +#### 1.1.4 PIP 运行时环境 + +PIP 运行时环境使用 pip 在容器内安装所需依赖。这种环境适合快速设置并能在容器中完成依赖安装的场景，是默认的运行时环境。它不需要预先构建包含依赖的镜像，通过 pip 直接管理 Python 包。 + +**挂载配置:** +- `local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**适用场景:** +- 使用PIP源安装的ROCK +- 快速测试ROCK + +**优势:** +- 简单的部署设置 + +**限制:** +- 依赖安装时间较长 +- 需要网络访问以安装依赖包 +- 每次启动时都需要安装依赖 + +#### 1.1.5 配置指南 + +根据不同的使用场景，可以参考以下选择指南： + +| 场景 | 推荐环境 | 原因 | +|------|----------|------| +| 生产环境 | Docker 运行时 | 快速启动，稳定性能 | +| 开发环境，同一 OS | 本地运行时 | 环境重用，开发周期快 | +| Mac 开发 | UV 运行时 | 支持最佳的跨平台兼容性 | +| 跨平台开发 | UV 运行时 | 避免环境兼容性问题 | +| 快速测试 | UV 运行时 | 无需预配置工作 | +| PIP源安装 | PIP 运行时 | 直接使用 pip 安装依赖 | + +这些运行时环境通过 `ROCK_WORKER_ENV_TYPE` 环境变量进行配置，该变量可设置为 "local"、"docker"、"uv" 或 "pip"。 + +### 1.2 日志配置 + +在日志配置方面，ROCK 的日志系统具有以下特性： + +- 日志系统不能同时输出到文件和控制台，只有当设置了 `ROCK_LOGGING_PATH` 时，日志才会输出到指定文件，否则输出到控制台。 +- `ROCK_LOGGING_LEVEL` 用于控制日志输出级别，`ROCK_LOG_LEVEL` 用于通用日志级别设置。 + +## 2. 分布式部署要求 + +由于 ROCK 支持分布式部署，当在 Ray 集群的不同节点上运行时，需要满足以下一致性要求： + +#### 目录结构一致性 +在所有 Ray 节点上，必须保证以下目录结构完全一致： +- ROCK 项目仓库目录 +- `.venv` 虚拟环境目录 +- `.venv` 依赖的 base Python 目录 + + +#### 挂载要求 +ROCK 的启动依赖于挂载 ROCK 项目和对应的 base Python 环境，要求在多机环境中保持一致性： + +#### 验证分布式配置 +可以通过以下方式验证分布式部署配置： + +```bash +# 在所有节点上检查目录一致性 +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# 验证 Python 环境可用性 +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# 检查所有节点上的环境变量设置 +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + + + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速搭建 ROCK 环境 +- [API 文档](../References/api.md) - 查看沙箱相关的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 配置沙箱 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/overview.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/overview.md new file mode 100644 index 0000000000..a02536f50d --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/overview.md @@ -0,0 +1,40 @@ +--- +sidebar_position: 1 +--- + +# 概览 + +ROCK (Reinforcement Open Construction Kit) 是一个开源的强化学习环境开发框架，旨在简化强化学习环境的开发、部署和管理流程。 + +## 什么是 ROCK + +ROCK (Reinforcement Open Construction Kit) 是一个开源强化学习环境开发框架。通过使用 ROCK，开发者可以快速地开发强化学习环境，并结合其他强化学习训练框架，实现高效的强化学习训练。 + +ROCK 提供了完整的沙箱环境管理功能，支持容器化部署，能够实现环境的快速创建、运行和销毁。同时，ROCK 兼容 GEM 协议，为强化学习环境提供了标准化的接口。 + +## ROCK 的核心功能 + +1. **简化开发流程**：简化强化学习环境的开发、构建和管理流程，支持多种开源的强化学习环境 +2. **大规模调度部署**：支持快速强化学习环境的大规模调度部署，通过 GEM 协议可以方便地访问强化学习环境 +3. **框架集成**：与其他强化学习训练框架集成，实现大规模可扩展的强化学习训练 + +## ROCK 的价值 + +ROCK 为不同角色的工程师提供了显著价值： + +- **强化学习算法工程师**：ROCK 可以简化强化学习环境的开发流程，让工程师专注于算法实现 +- **强化学习应用工程师**：ROCK 可以进行快速强化学习环境的大规模部署，提高应用开发效率 + +## 相关文档 + +如果您是第一次使用 ROCK，建议按以下顺序阅读文档： +1. [快速开始指南](./Getting%20Started/quickstart.md) - 快速搭建开发环境 +2. [配置指南](./User%20Guides/configuration.md) - 配置您的 ROCK 环境 +3. [Python SDK 文档](./References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +4. [API 文档](./References/api.md) - 了解完整的 API 接口 +5. [安装指南](./Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 + + + + + diff --git a/docs/versioned_docs/version-1.5.x/Getting Started/installation.md b/docs/versioned_docs/version-1.5.x/Getting Started/installation.md new file mode 100644 index 0000000000..c45b09a8fe --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Getting Started/installation.md @@ -0,0 +1,143 @@ +--- +sidebar_position: 3 +--- + +# Installation + +This document explains how to install and set up the ROCK development environment using both `uv` and `pip`. The project is a Reinforcement Open Construction Kit that supports various components. + +## Using uv (Recommended) + +### Quick Install All Dependencies + +```bash +# Install all dependencies including optional ones +uv sync --all-extras + +# Install development/testing dependencies +uv sync --all-extras --all-groups +``` + +### Install Different Dependency Groups + +#### Core Dependencies Only +```bash +uv sync +``` + +#### Admin Component Dependencies +```bash +uv sync --extra admin +``` + +#### Rocklet Execution Environment Dependencies +```bash +uv sync --extra rocklet +``` + + +#### All Dependencies at Once +```bash +uv sync --all-extras +``` + +#### Development/Testing Dependencies +```bash +uv sync --all-extras --group test +``` + +## Using pip + +### Install from pip source + +#### Core Dependencies Only +```bash +pip install rl-rock +``` + +#### Admin Component Dependencies +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install "rl-rock[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install "rl-rock[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install "rl-rock[all]" +``` + +### Install with pip from source code + +#### Core Dependencies Only +```bash +pip install . +``` + +#### Admin Component Dependencies +```bash +pip install ".[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install ".[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install ".[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install ".[all]" +``` + +## Available Entry Points + +The package provides the following command line scripts: + +- `rocklet`: ROCK execution environment server (rock.rocklet.server:main) +- `admin`: Admin management server (rock.admin.main:main) +- `envhub`: Environment hub server (rock.envhub.server:main) +- `rock`: Main ROCK command line interface (rock.cli.main:main) + +## Development Setup + +### Using uv (Recommended) + +```bash +# Clone and set up development environment +git clone +cd ROCK +uv sync --all-extras --group test + +# Run tests +uv run pytest + +``` + +### Using pip + +```bash +# For development, install in editable mode with all extras +pip install -e ".[all]" + +# Or separately +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" # Optional extras +``` + +## Additional Notes + +- The project is configured to use the Alibaba cloud PyPI mirror by default: `https://mirrors.aliyun.com/pypi/simple/` +- For local development, running tests requires the `test` dependency group diff --git a/docs/versioned_docs/version-1.5.x/Getting Started/quickstart.md b/docs/versioned_docs/version-1.5.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..2f14808f5b --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Getting Started/quickstart.md @@ -0,0 +1,166 @@ +--- +sidebar_position: 2 +--- + +# Getting Started + +This guide will demonstrate how to use ROCK to create and manage reinforcement learning environments through complete examples. + +## 1. Environment Preparation + +We recommend starting ROCK on Linux systems to maximize dependency reuse and improve environment startup speed. If you need to try on macOS, please refer to the [MacOS Startup](#7-macos-startup) section. + +Before starting, please ensure your system has the following dependencies installed: + +### 1.1 System Requirements + +- **Docker**: ROCK uses Docker for containerized environment management +- **uv**: ROCK uses uv for dependency management and virtual environment creation + +### 1.2 Verify Dependency Installation + +```bash +# Verify Docker installation +docker --version + +# Verify Docker image, and example depends on python:3.11 image +docker pull python:3.11 + +# Verify uv installation +uv --version +``` + +### 1.3 Project Initialization + +```bash +# Clone repository +git clone +cd ROCK + +# Create virtual environment (using uv-managed Python, use python 3.11 as an example) +uv venv --python 3.11 --python-preference only-managed + +# Install all dependency groups +uv sync --all-extras +``` + +> **Important Note**: To ensure ROCK can correctly mount the project and virtual environment along with its base Python interpreter, it is strongly recommended to use uv-managed Python environments to create virtual environments rather than system Python. + +## 2. Activate Virtual Environment + +Before running any ROCK commands, you need to activate the virtual environment. Ensure sys.base_prefix is a uv-managed environment, such as `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` or similar paths. + +```bash +# Activate virtual environment +source .venv/bin/activate + +# Verify Python environment +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **Verification Point**: Ensure the output base prefix path points to a uv-managed Python environment, not system Python. + +## 3. Verify Environment Configuration + +After activating the virtual environment, verify that dependencies are installed correctly: + +```bash +# Check key dependencies +python -c "import rock; print(\"Hello ROCK\")" +``` + +## 4. Start ROCK Service + +After activating the virtual environment, start the ROCK Admin service on project root: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Start ROCK Admin service (local environment) +rock admin start +``` + +After the service starts, you will see output similar to the following: + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **Service Information**: The ROCK Admin service runs by default on `http://127.0.0.1:8080`. + +## 5. Run Example Environments + +Now you can run example environments to verify the installation. Ensure the ROCK service is running, then open a new terminal window to execute the following commands: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Run sandbox example +python examples/sandbox_demo.py + +# Run GEM protocol example +python examples/sokoban_demo.py +``` + +### 5.1 Example Descriptions + +- **sandbox_demo.py**: Demonstrates how to use ROCK's sandbox SDK to create and manage containerized environments +- **sokoban_demo.py**: Demonstrates how to use ROCK's GEM protocol compatible interface to create reinforcement learning environments + +> **Running Requirements**: Ensure the ROCK Admin service is running, as examples need to communicate with the service. + +## 6. Distributed Environment Configuration (Optional) + +For distributed multi-machine environments, ensure the following configurations are consistent: + +1. All machines use the same root Python interpreter for ROCK and uv Python configurations +2. Docker versions are consistent across all nodes +3. Network configuration allows normal communication between nodes + + +## 7. MacOS Startup + +On macOS, if you need to start Linux image environments, you first need to set the environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +During container startup, the corresponding uv environment will be installed. For details, please refer to the `rock/rocklet/local_files/docker_run_with_uv.sh` script. + +> **Note**: Compared to Linux systems, the startup speed on macOS will be slower and more dependent on network conditions. You can adjust the script according to actual conditions.You can find detatils for ROCK_WORKER_ENV_TYPE in [Configuration Guide](../User%20Guides/configuration.md). + +## 8. Starting from Pip Source + +If starting the Admin Server from Pip source, after completing the ROCK installation by referring to [installation](./installation.md), you need to set an additional environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(This startup method will pull and install the latest rocklet from the PyPI source when starting the container environment. The startup speed is relatively slow, so it is only recommended for testing purposes. For production environments, other startup methods are still recommended.) + +## Summary + +Congratulations! You have successfully completed the ROCK quick start guide. You should now be able to: + +- Properly set up the ROCK development environment +- Use uv-managed Python environments +- Start and manage ROCK services +- Run example programs to verify installation +- Configure ROCK in distributed environments (if needed) + +For a deeper understanding of ROCK's additional features, please refer to the following documents: + +## Next Steps + +- [Configuration Guide](../User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](../References/api.md) - View complete API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the Python SDK for development +- [Installation Guide](./installation.md) - Detailed information about ROCK installation and setup +- [Overview](../overview.md) - Understand ROCK's design philosophy \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/Getting Started/rock-agent.md b/docs/versioned_docs/version-1.5.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..eb6b54a65b --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Getting Started/rock-agent.md @@ -0,0 +1,72 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent Quick Start + +Rock Agent is an AI Agent runtime framework provided by ROCK, supporting various types of Agents running in sandbox environments. + +## Prerequisites +- Make sure you have a working ROCK service, if you need to locally start the service side, refer to [Quick Start](quickstart.md). + +## Examples + +ROCK provides two Hello World Agent examples in the `examples/agents/` directory: + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent example +└── iflow_cli/ # IFlowCli Agent example +``` + +### Run IFlowCli Example + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### Run ClaudeCode Example + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli Configuration File + +The configuration file is located at `examples/agents/iflow_cli/rock_agent_config.yaml`: + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # Enter your API key + IFLOW_BASE_URL: "" # Enter your base URL + IFLOW_MODEL_NAME: "" # Enter your model name +``` + +## ClaudeCode Configuration File + +The configuration file is located at `examples/agents/claude_code/rock_agent_config.yaml`: + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # Enter your anthropic base url + ANTHROPIC_API_KEY: "" # Enter your anthropic api key +``` + +## Related Documentation + +- [RockAgent Reference](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/versioned_docs/version-1.5.x/Getting Started/rockroll.md b/docs/versioned_docs/version-1.5.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..2465a7733f --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Getting Started/rockroll.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL Quick Start Guide + +This guide will walk you through running a reinforcement learning training example based on the Sokoban game, using ROLL (the training framework) and ROCK (the environment management tool). + +## 1. Prerequisites + +Before you begin, please ensure your system has the following dependencies installed. + +### 1.1 System Requirements + +- **OS**: A Linux-based system is recommended (e.g., Ubuntu 20.04+). +- **Hardware**: An NVIDIA GPU with the corresponding drivers is recommended. +- **Docker**: ROCK uses Docker for containerized environment management. +- **uv**: ROCK uses uv for dependency management and virtual environment creation. + +### 1.2 Verify Dependencies & Pre-pull Image + +```bash +# Verify Docker installation +docker --version + +# Verify Docker is running and pre-pull the Sokoban environment image +# This will save time when the training starts. +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# Verify uv installation +uv --version + +``` + +### 1.3 Initialize the Project + +```bash +# Clone the project repositories +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# Ensure both repositories are in the same parent directory, like this: +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. Launch the Training Process + +> Note: The following instructions use torch==2.6.0 and vLLM==0.8.4 as an example. + + +### Option 1: Using a Virtual Environment (Recommended) + +#### Why is this method recommended? +- Isolation: A uv virtual environment ensures that project dependencies are isolated from your system, preventing conflicts. +- Fast Startup: ROCK can reuse this virtual environment, significantly speeding up subsequent task initializations. +- Stability & Reproducibility: Dependency management is cleaner and more reliable. + + +```bash +# Navigate to the ROCK directory +cd ROCK + +# Create and activate a Python 3.10 virtual environment (ROLL recommends Python 3.10) +uv venv --python 3.10 --python-preference only-managed +source .venv/bin/activate + +# Install all of ROCK's dependencies using uv +uv sync --all-extras + +# If using Python 3.10, starting Ray may raise a `ValueError: is not a valid Sentinel`. +# This is due to an incompatibility between `ray` and `click` versions 8.3+. +# To fix this, downgrade `click` to a version below 8.3. This issue does not affect Python 3.11. +uv pip install 'click>=8.2,click<8.3' + +# Navigate to the ROLL directory to install its dependencies +cd ../ROLL + +# Install core PyTorch components +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 + +# Install transformer-engine. The --no-build-isolation flag prevents errors where torch cannot be found. +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation + +# Install a pre-compiled version of flash-attention matching the specific CUDA and PyTorch versions +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# Install the remaining dependencies +uv pip install -r requirements_torch260_vllm.txt + +# (Optional) Install Tensorboard to check training metrics +uv pip install tensorboard -i $PYPI_MIRROR + +# All set! Launch the training script. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### Option 2: Using the System Environment (Alternative) + +For optimal compatibility with this method, we recommend running these commands inside one of ROLL's official base Docker images. These images come pre-installed with matching CUDA, cuDNN, and other foundational libraries. + +> [ROLL's Official Docker Image List](https://alibaba.github.io/ROLL/docs/Getting%20Started/Installation/image_address) + + +#### Warning +This method will install all Python packages directly into your current environment (e.g., the container's base system), which may cause conflicts with system packages or other projects. + +Since ROCK cannot reuse the environment, it may need to reinstall some dependencies each time a task starts, leading to slower startup times that are dependent on network speed. + + +```bash +# Install ROCK's dependencies +cd ROCK +pip install . +pip install ".[admin]" + +# Install ROLL's dependencies +cd ../ROLL +pip install -r requirements_torch260_vllm.txt + +# Crucial: Configure ROCK to use uv as its worker environment manager +export ROCK_WORKER_ENV_TYPE=uv + +# Launch the training script +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +You have now successfully launched the Sokoban reinforcement learning training process. Happy Rock & Roll！ + + +## 3. Multi-Node Deployment + +Instead of running everything on a single machine, you can deploy the **ROCK Service** and **ROLL job** on separate machines. This is a common client-server setup where they communicate over the network. + +### 3.1 Deploy the ROCK Service on Machine A + +On a dedicated machine (or container), follow the [ROCK Quick Start Guide](./quickstart.md) to deploy and start the ROCK service. + +> **Important** +> After starting the service, take note of its IP address and port (e.g., `http://192.168.1.10:8000`). You will need this address for the subsequent steps. + +### 3.2 Prepare the ROLL Client on Machine B + +On the other machine where you will run the training task, perform the following steps. + +1. Verify Network Connectivity + +First, use the curl command to check if you can reach the ROCK service on Machine A from Machine B. +```bash +# Replace : with the actual address of your ROCK service +# If successful, you should receive a response like {"message":"hello, ROCK!"} +curl http://: +``` + +2. Prepare the ROLL Environment + +```bash +# Clone the ROLL repository +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# Install dependencies +pip install -r requirements_torch260_vllm.txt +``` + +3. Configure the ROLL Connection Address + +Modify ROLL's configuration file to point to the remote ROCK service. +- Open the configuration file: examples/agentic_demo/agentic_val_sokoban_sandbox.yaml. +- Find the "SokobanSandbox" section under "env_config". +- Update the base_url value to your ROCK service's address. +```yaml +custom_envs: + SokobanSandbox: + env_config: + # Change the address here to your ROCK service's address + # Example: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. Start Training +Once configured, you can start the ROLL training script on Machine B. + +```bash +# This script will now request environments from the ROCK service on Machine A over the network. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### Advanced: Distributed ROLL Training + +If you wish to deploy the ROLL training task itself in a distributed manner, you can refer to ROLL's official documentation for distributed deployment. +> [Quick Start: Multi-Node Deployment Guide](https://alibaba.github.io/ROLL/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/codes.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..dceb8d3182 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +Error code definitions and categories for error handling and retry strategies. + +## Usage Example + +```python +import rock + +def test_codes_values(): + """Test basic status code values""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes Categories + +```python +OK = 2000, "OK" +""" +Success codes (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +Client error codes (4xxx): + +These errors indicate issues with the client request, +SDK will raise Exceptions for these errors. +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +Server error codes (5xxx): + +These errors indicate issues on the server side, +SDK will raise Exceptions for these errors. +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +Command/execution error codes (6xxx): + +These errors are related to command execution and should be handled by the model, +SDK will NOT raise Exceptions for these errors. +""" +``` + +## Retry Strategy Recommendations + +- **Retry trigger**: Only retry when `INTERNAL_SERVER_ERROR` occurs +- **Other error handling**: + - `BAD_REQUEST`: Check if there are issues with the arun call logic + - `COMMAND_ERROR`: stdout goes to `observation.output`, stderr goes to `observation.failure_reason` +- `COMMAND_ERROR` note: When bash execution fails, both stdout and stderr may be non-empty. It is recommended to prompt the model with both output and failure_reason from the observation. + +## Retry Example + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/deploy.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..5fd5b70546 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +Sandbox resource deployment manager for local directory deployment and template formatting. + +## deploy_working_dir - Deploy Local Directory + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# Deploy local directory (auto-generated target path) +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"Deployed to: {target}") # e.g., /tmp/rock_workdir_abc123 + +# Deploy to specific target path +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - Template Variable Substitution + +The `format` method supports two template syntaxes: + +- **`${variable}`** - Standard Python string template syntax +- **`<>`** - Alternative syntax (converted to `${variable}` internally) + +```python +# After deploy_working_dir, use ${working_dir} placeholder +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# Result: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# Alternative <<>> syntax +cmd = deploy.format("cat <>/file.txt") +# Result: cat /tmp/rock_workdir_abc123/file.txt + +# Combine with custom variables +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# Result: cat /tmp/rock_workdir_abc123/settings.json + +# Shell syntax is preserved +cmd = deploy.format("echo $((3 << 2 >> 1))") +# Result: echo $((3 << 2 >> 1)) + +# Access working_dir directly +if deploy.working_dir: + print(f"Current working directory: {deploy.working_dir}") +``` + +## Multiple Deployments + +Subsequent calls overwrite previous working directory paths: + +```python +# First deployment +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# Second deployment (overwrites previous path) +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/file_system.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..a64228e8f1 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +File system interface for sandbox environment operations including permission and ownership management. + +## chown - Change Owner + +```python +from rock.actions.sandbox.request import ChownRequest + +# Create remote user before changing ownership +await sandbox.remote_user.create_remote_user("deploy") + +# Get current working directory +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# Change directory owner +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# Recursively change owner for directory and contents +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - Change Permissions + +```python +from rock.actions.sandbox.request import ChmodRequest + +# Create test directory +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# Change directory permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# Recursively change permissions (includes subdirectories and files) +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# Set maximum permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - Upload Directory + +```python +import os +from pathlib import Path + +# Prepare local directory +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# Upload to sandbox +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"Upload success: {result.output}") +else: + print(f"Upload failed: {result.failure_reason}") +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/model-service.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..23dbc21bfc --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service (Experimental) + +The Model Service provided by ROCK is responsible for handling AI model call communications, serving as a communication bridge between agents and training frameworks (such as Roll) or actual LLM inference services. + +## RockAgent Integration + +ModelService is typically **automatically managed by RockAgent** - no manual lifecycle management is required. Simply enable it in the configuration: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # Enable ModelService, RockAgent manages its lifecycle +) +``` + +RockAgent will automatically: +- Install ModelService (install Python runtime, install model service package) +- Start/stop ModelService +- Monitor Agent process + +## Architecture Overview (Local Mode) + +In local mode, the model service uses the **file system** as the communication medium, implementing a request-response mechanism between agents and models. + +When an agent needs to call a model, the request is first written to a log file, then processed by the listening component. When the model generates a response, the result is written back to the log file and read by the waiting agent. + +## anti_call_llm - Core API + +`anti_call_llm()` is the **most important API in Local mode**, used to manually trigger LLM anti-calls for fine-grained control over model calls: + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='OpenAI type response', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +**Use cases:** +- After Agent captures LLM response, call this method to notify Roll runtime +- Supports carrying response data for error handling or retry +- Configurable timeout and check interval for different network environments + +## CLI Commands + +To use the model service via CLI, ROCK provides a set of CLI commands that can be accessed via `rock model-service` after installing ROCK in the sandbox: + +### start command +Start the model service process +```bash +rock model-service start --type [local|proxy] [options] +``` + +Parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `--type` | str | `local` | Service type: `local` or `proxy` | +| `--config-file` | str | None | Path to configuration file | +| `--host` | str | None | Server host address (overrides config) | +| `--port` | int | None | Server port (overrides config) | +| `--proxy-base-url` | str | None | Proxy base URL | +| `--retryable-status-codes` | str | None | Comma-separated list of retryable status codes | +| `--request-timeout` | int | None | Request timeout in seconds | + +### watch-agent command +Monitor the agent process and send a SESSION_END message when the process exits +```bash +rock model-service watch-agent --pid +``` + +Parameters: +- `--pid`: The ID of the agent process to monitor + +### stop command +Stop the model service +```bash +rock model-service stop +``` + +### anti-call-llm command +Anti-call the LLM interface +```bash +rock model-service anti-call-llm --index [--response ] +``` + +Parameters: +- `--index`: Index of the previous LLM call, starting from 0 +- `--response`: Response from the previous LLM call (optional) + +## File Communication Protocol + +The model service uses files for inter-process communication, defining specific marker formats to distinguish requests and responses: + +### Request Format +``` +LLM_REQUEST_START{JSON request data}LLM_REQUEST_END{metadata JSON} +``` + +### Response Format +``` +LLM_RESPONSE_START{JSON response data}LLM_RESPONSE_END{metadata JSON} +``` + +### Session End Marker +``` +SESSION_END +``` + +Metadata contains timestamp and index information to ensure message order and processing. + +## SDK Usage + +### ModelServiceConfig + +Model service configuration class, located in `rock/sdk/sandbox/model_service/base.py`: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # Service type + install_cmd="pip install rock-model-service", # Install command + install_timeout=300, # Install timeout (seconds) + start_cmd="rock model-service start --type ${type}", # Start command + stop_cmd="rock model-service stop", # Stop command + logging_path="/data/logs", # Log path + logging_file_name="model_service.log", # Log filename +) +``` + +| Config | Default | Description | +|--------|---------|-------------| +| `enabled` | `False` | Whether to enable model service (RockAgent manages) | +| `type` | `"local"` | Service type: `local` or `proxy` | +| `install_cmd` | - | Model service package install command | +| `install_timeout` | `300` | Install timeout in seconds | +| `start_cmd` | - | Start command template | +| `stop_cmd` | - | Stop command | +| `logging_path` | `/data/logs` | Log directory path | +| `logging_file_name` | `model_service.log` | Log filename | + +### ModelService + +Model service management class, handles the lifecycle of model services within the sandbox: + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# Typically auto-managed by RockAgent, no manual calls needed +# The following methods are only for manual control when needed + +# Install model service +await model_service.install() + +# Start model service +await model_service.start() + +# Monitor agent process +await model_service.watch_agent(pid="12345") + +# Execute anti-call LLM (Core API for Local mode) +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# Stop model service +await model_service.stop() +``` + +## API Reference + +### install() + +Install model service dependencies in the sandbox. + +```python +await model_service.install() +``` + +Execution steps: +1. Create and initialize Python runtime environment +2. Create Rock config file +3. Install model service package + +**Note:** Typically auto-called by RockAgent. + +### start() + +Start the model service. + +```python +await model_service.start() +``` + +Prerequisite: Must call `install()` first. + +**Note:** Typically auto-called by RockAgent. + +### stop() + +Stop the model service. + +```python +await model_service.stop() +``` + +If the service is not running, this operation will be skipped. + +**Note:** Typically auto-called by RockAgent. + +### watch_agent(pid) + +Monitor the agent process. + +```python +await model_service.watch_agent(pid="12345") +``` + +Sends `SESSION_END` message when the process exits. + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +Execute anti-call LLM operation. **This is the most important API in Local mode.** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='{"result": "..."}', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +## Configuration Options + +### Service Configuration +- `SERVICE_HOST`: Service host address, defaults to `"0.0.0.0"` +- `SERVICE_PORT`: Service port, defaults to `8080` + +### Log Configuration +- `LOG_FILE`: Log file path used for communication, containing request and response data + +### Trajectory (Traj) Logging +The model service records LLM call trajectories (traj) to a JSONL file for debugging and analysis. + +| Environment Variable | Default | Description | +|---------------------|---------|-------------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | Directory for traj log files | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | Append mode (true/false) | + +**Traj file location**: `{DATA_DIR}/LLMTraj.jsonl` + +**Traj file format** (JSONL - one JSON object per line): +```json +{"request": {...}, "response": {...}} +``` + +### Polling Configuration +- `POLLING_INTERVAL_SECONDS`: Polling interval, defaults to `0.1` seconds +- `REQUEST_TIMEOUT`: Request timeout, defaults to unlimited + +### Marker Configuration +Defines markers used to distinguish different types of messages in the log file: +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig (Server-side) + +The server-side configuration class defines how the model service handles requests: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `host` | str | `"0.0.0.0"` | Server host address | +| `port` | int | `8080` | Server port | +| `proxy_base_url` | str \| None | `None` | Direct proxy URL | +| `proxy_rules` | dict | See below | Model name to URL mapping | +| `retryable_status_codes` | list[int] | `[429, 500]` | Retryable HTTP status codes | +| `request_timeout` | int | `120` | Request timeout in seconds | + +**Default proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/python_sdk.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..5272d7edb6 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK Reference + +This guide provides detailed information on how to use the ROCK SDK for development, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +ROCK SDK provides developers with convenient Python interfaces to use ROCK platform features, including sandbox environment management and GEM environment interaction. + +> **Important Note**: Before using the SDK, ensure that the ROCK Admin service is running. You can start it with the following command: +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 Basic Sandbox Operations + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 Sandbox Group Management + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# Create sandbox group configuration +config = SandboxGroupConfig( + image="python:3.11", + size=4, # Create 4 sandboxes + start_concurrency=2, # Concurrency level for startup is 2 +) + +# Create and start sandbox group +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# Batch operations +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# Batch stop +await sandbox_group.stop() +``` + +### 2.3 Configuration Example + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 Sandbox Speedup Configuration + +ROCK provides sandbox network acceleration capabilities, supporting configuration of APT, PIP, and GitHub mirror sources to improve package download speeds in restricted network environments. + +#### Supported Speedup Types + +**APT Mirror Configuration** + +Configure APT package manager mirror sources for faster Debian/Ubuntu package downloads. + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# Configure APT mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP Mirror Configuration** + +Configure Python package index mirrors for faster pip installations. + +```python +# HTTP mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub Acceleration** + +Configure GitHub IP acceleration by adding custom DNS resolution entries. + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### Complete Example + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """Create sandbox and configure acceleration""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # Configure acceleration (before installing packages) + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # Speedup does not automatically install PIP, it only configures mirror sources for acceleration + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # GitHub can be accelerated through mirror IP + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### Important Notes + +1. **Configuration Order**: Configure speedup before installing packages +2. **HTTPS vs HTTP**: HTTPS mirrors don't require trusted-host configuration for PIP +3. **GitHub IP**: Different regions may require different IPs for optimal performance +4. **Persistence**: Configurations persist within the sandbox lifecycle +5. **Multiple Calls**: Subsequent speedup calls will override previous configurations +6. **PIP Installation**: The speedup feature only configures mirror sources and does not automatically install PIP + +## 3. GEM SDK + +### 3.1 Python SDK Approach + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## Related Documents +- [Quick Start Guide](../../Getting%20Started/quickstart.md) - Learn how to quickly get started with the ROCK SDK +- [API Documentation](../api.md) - View the underlying API interfaces encapsulated by the SDK +- [Configuration Guide](../../User%20Guides/configuration.md) - Learn about SDK-related configuration options +- [Installation Guide](../../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/remote_user.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..810bbcf028 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/remote_user.md @@ -0,0 +1,70 @@ +# Remote User + +Remote user management for creating and managing users in the sandbox. + +## Usage Examples + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +Create a remote user. + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +Check if a user exists. + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/rock-agent.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..f24ade49ac --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent (Experimental) + +RockAgent is the core Agent implementation in the ROCK framework, directly inheriting from the `Agent` abstract base class. It provides complete Agent lifecycle management, including environment initialization, ModelService integration, command execution, and more. + +Using `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, you can install and run Agents in the Sandbox environment provided by Rock. + +## Core Concepts + +The core workflow of RockAgent is divided into two phases: + +1. **install(config)**: Initialize the Agent environment, including deploying the working directory, setting environment variables, initializing the runtime environment, etc. +2. **run(prompt)**: Execute the Agent task, replace placeholders, and start the Agent process + +## Quick Start + +### Claude Code Example + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli Example + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} is required + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # Environment variables + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent Example + +```yaml +working_dir: "." # Upload local current directory containing langgraph_agent.py to sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # Run local script + +runtime_env_config: + type: python + pip: # Install pip dependencies + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## Configuration Details + +### Basic Configuration + +```yaml +agent_type: "default" # Agent type identifier (default: "default") +agent_name: "demo-agent" # Agent instance name (default: random uuid) +version: "1.0.0" # Version identifier (default: "default") +instance_id: "instance-001" # Instance ID (default: "instance-id-") +agent_installed_dir: "/tmp/installed_agent" # Agent installation directory (default: "/tmp/installed_agent") +agent_session: "my-session" # Bash session identifier (default: "agent-session-") +env: # Environment variables (default: {}) + OPENAI_API_KEY: "xxxxxxx" +``` + +### Working Directory Configuration + +```yaml +working_dir: "./my_project" # Local directory to upload to sandbox (default: None, no upload) +project_path: "/testbed" # Working directory in sandbox for cd (default: None) +use_deploy_working_dir_as_fallback: true # Whether to fall back to deploy.working_dir when project_path is empty (default: true) +``` + +### Execution Configuration + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent execution command, must contain ${prompt} (default: None) + +skip_wrap_run_cmd: false # Skip wrapping run_cmd with PATH (default: false) + +# Timeout configuration +agent_install_timeout: 600 # Installation timeout in seconds (default: 600) +agent_run_timeout: 1800 # Run timeout in seconds (default: 1800) +agent_run_check_interval: 30 # Check interval in seconds (default: 30) +``` + +**`skip_wrap_run_cmd`**: +- `false` (default): Wraps the command with `export PATH=:$PATH &&` to ensure runtime environment executables are used +- `true`: Skips PATH wrapping, runs the command directly with `bash -c` + +### Initialization Hooks + +```yaml +pre_init_cmds: # Commands executed before initialization (default: read from env_vars) + - command: "apt update && apt install -y git" + timeout_seconds: 300 # Command timeout in seconds (default: 300) + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # Commands executed after initialization (default: []) + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**Notes**: +- `pre_init_cmds` and `post_init_cmds` do not inherit the Agent's `env` environment variables +- Typically used for installation operations and configuration file movement +- Common command examples: + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv Configuration + +```yaml +runtime_env_config: # Refer to RuntimeEnv documentation for details + type: "python" # Runtime type: python / node (default: "python") + version: "3.11" # Version number + pip: # Python dependency package list + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node Runtime Example**: + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**Automatic Operations**: +- Install corresponding runtime based on `type` (Python or Node.js) +- Install `pip` dependencies (if configured) +- Execute `custom_install_cmd` custom installation command (if configured) +- Support `npm_registry` configuration for Node.js npm mirror source + +### ModelService Configuration + +```yaml +model_service_config: # Refer to ModelService documentation for details + enabled: true # Enable ModelService (default: false) +``` + +**Automatic Operations**: +- Installation phase: Install ModelService (install only, do not start) +- Run phase: Start ModelService + `watch_agent` monitoring process + +**Notes**: You need to set the model request URL to the ModelService URL. For example, if the ModelService provides an OpenAI-compatible URL at `http://127.0.0.1:8080/v1/chat/completions`, you typically need to set the Agent's LLM request URL to `http://127.0.0.1:8080/v1/`. + +## API Reference + +### install(config) + +Initialize the Agent environment. + +**Execution Flow**: +1. If `working_dir` is configured, deploy to sandbox +2. Set up bash session and configure env environment variables +3. Execute `pre_init_cmds` +4. Initialize RuntimeEnv and ModelService in parallel (if enabled) +5. Execute `post_init_cmds` + +**Parameters**: +- `config`: Agent configuration file, supports two input methods: + - **String path**: YAML configuration file path, default value is `"rock_agent_config.yaml"` + - **RockAgentConfig object**: Directly pass a `RockAgentConfig` instance + +### run(prompt) + +Execute the Agent task. + +**Execution Flow**: +1. Replace placeholders and prepare Agent run command +2. Start the agent process +3. If ModelService is enabled, start `watch_agent` +4. Wait for task completion and return results + +## Advanced Usage + +### Difference and Interaction between working_dir and project_path + +| Configuration | Function | Interaction Method | +|--------------|----------|-------------------| +| `working_dir` | Local directory uploaded to sandbox | Calls `deploy.deploy_working_dir()` to upload, after upload `deploy.working_dir` becomes the path in sandbox | +| `${working_dir}` | Placeholder in commands | Replaced by `deploy.format()` with the value of `deploy.working_dir`, replaced in init_cmds and run_cmd in the configuration | +| `project_path` | Working directory in sandbox | Used for `cd project_path` before running, when not set it enters the `deploy.working_dir` working directory | +| `use_deploy_working_dir_as_fallback` | Whether to fall back to deploy.working_dir when project_path is not set at runtime | Default is `true`, when set to `false` it will not enter working_dir even if project_path is not set | + +**Usage Recommendations**: +- Use `working_dir` to upload local project code to sandbox +- Use `project_path` to specify the working directory in sandbox (e.g., `/testbed`) +- Set `use_deploy_working_dir_as_fallback: false` scenario: Need to perform local file mounting, but want to run Agent in the image's default working directory + +### Placeholder Usage + +Rock Agent supports replacing the following placeholders in the configuration file: + +- `${prompt}`: Required in run_cmd, will be replaced with the prompt passed to `run(prompt)` +- `${working_dir}`: Optional, will be replaced with the actual working directory path in sandbox, also supported in init_cmds and run_cmd +- `${bin_dir}`: Optional, will be replaced with the runtime environment's bin directory path + +**Example**: +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback Explanation + +When `project_path` is not set: +- `true` (default): Before running Agent, it will automatically `cd` to `deploy.working_dir` +- `false`: Before running Agent, it will not automatically switch directories, staying in the current directory + +Applicable Scenarios: +- `true`: Most scenarios, where you want Agent to run in the uploaded code directory +- `false`: Need to mount local files, but want to run Agent in the image's default working directory (e.g., `/app`, `/testbed`) + +## Complete Configuration Example + +```yaml +# ========== Basic Configuration ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== Working Directory Configuration ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== Run Configuration ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# Timeout configuration +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== Initialization Commands ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== Runtime Environment Configuration ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService Integration ========== +model_service_config: + enabled: true +``` + +## Usage Examples + +### Using YAML Configuration File (Recommended) + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/runtime-env.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..e1996cfcdb --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,136 @@ +# RuntimeEnv + +The RuntimeEnv module is used to manage language runtime environments in the sandbox (currently providing Python / Node.js). + +## Quick Start (Example) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +An async factory method that creates and initializes a RuntimeEnv instance based on the configuration, and automatically registers it to `sandbox.runtime_envs`. + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# Auto-registered; accessible via sandbox.runtime_envs[env.runtime_env_id] +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +Wraps a command by adding `bin_dir` to PATH to ensure executables from the runtime environment are used with priority. + +```python +wrapped = env.wrapped_cmd("node script.js") +# Returns: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +Executes a command within the runtime environment. Internally implemented based on `wrapped_cmd`. + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | Type identifier | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python version; default is 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | List of pip packages or a requirements.txt path | +| `pip_index_url` | `str \| None` | Environment variable | pip index mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | List of executables to symlink | + +## NodeRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | Type identifier | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node version; default is 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm registry mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | List of executables to symlink | + +## Constraints for Custom RuntimeEnv Implementations + +A custom RuntimeEnv must follow these rules: + +1. **Define the `runtime_env_type` class attribute**: used as a type identifier for automatic registration into the RuntimeEnv factory +2. **Override `_get_install_cmd()`**: return the install command +3. **The install command must end with**: renaming the directory to `runtime-env` + +## Simplified NodeRuntimeEnv Implementation Example + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config class: defines the config type so RuntimeEnv.create() can route to the corresponding implementation +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # Must match runtime_env_type + +# RuntimeEnv implementation class: defines how to install and run this runtime environment +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # Auto-registered to RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # Download the Node binary tarball and extract it, then rename to runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## Speeding Up Base Runtime Installation + +`PythonRuntimeEnv` downloads Python packages from https://github.com/astral-sh/python-build-standalone/releases/ by default. If the network is unavailable or slow, you can override the default install command via `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` or `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` (e.g., switch to an internal registry or a mirror). + +Default value example: + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +For example, override it to download from a mirror: + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +Make sure the command creates a `runtime-env` directory under the default working directory of `runtime_env`, and that `${workdir}/runtime-env/bin/` contains the expected executables, e.g.: + +- `${workdir}/runtime-env/bin/python` + +The same applies to Node.js: you can override the install command via `ROCK_RTENV_NODE_V22180_INSTALL_CMD` to use a faster download/install method. \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/sandbox.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..e6e1e43124 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/sandbox.md @@ -0,0 +1,114 @@ +# Handling Large Files and Long Command Outputs + +## `arun` + +`arun()` provides two knobs to control how `nohup` output is handled: + +1. **`response_limited_bytes_in_nohup`** *(integer type)* + Caps the number of characters returned from the nohup output file. Useful when you still need to stream some logs back but want an upper bound (default `None` = no cap). + +2. **`ignore_output`** *(bool, default `False`)* + When set to `True`, `arun()` skips reading the nohup output file entirely. The command still runs to completion and writes logs to `/tmp/tmp_.out`, but the SDK immediately returns a lightweight hint telling agents where to fetch the logs later (via `read_file`, download APIs, or custom commands). This fully decouples "execute command" from "inspect logs". The response also includes the **file size** to help users decide whether to download directly or read in chunks. + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# Example 1: limit the returned logs to 1024 characters +resp_limited = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# Example 2: skip collecting logs; agent will download/read them later +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# Use Sandbox.read_file(...), download APIs, or run 'cat /tmp/tmp_xxx.out' ... +``` + +## `read_file_by_line_range` + +Asynchronously reads file content by line range, with built-in support for automatic chunking and session management. Supports large file reading. + +### Key Features +- **Chunked reading for large files**: Automatically splits large files into chunks +- **Automatic line count**: Estimates total lines when end_line is not specified +- **Built-in retry mechanism**: Up to 3 retries for critical operations +- **Input validation**: Validates input parameters automatically +- **Session management**: Supports custom session or auto-created temporary session + +### Parameters +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `file_path` | str | - | File path to read (absolute or relative path in sandbox) | +| `start_line` | int \| None | 1 | Starting line number (1-based) | +| `end_line` | int \| None | None | Ending line number (inclusive), defaults to file end | +| `lines_per_request` | int | 1000 | Lines per request, range 1-10000 | + +### Return Value +- `ReadFileResponse`: Response object containing file content + - `content` (str): The file content read + +### Exception Handling +- `Exception`: Raised when `start_line < 1` +- `Exception`: Raised when `end_line < start_line` +- `Exception`: Raised when `lines_per_request` is not in range 1-10000 +- `Exception`: Raised when file reading fails + +### Usage Examples + +```python +# Read the entire file +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# Read a specific line range (lines 100 to 500) +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# Read from line 1990 to the end of file +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# Use custom chunk size +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### Notes +- Line numbers are 1-based, not 0-based +- For large files, consider increasing `lines_per_request` for better efficiency +- File path must be a valid path within the sandbox +- Uses `sed` command for file reading; ensure the sandbox image supports this command diff --git a/docs/versioned_docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md b/docs/versioned_docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..85f34cbe7c --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,229 @@ +# SWE-Bench Evaluation + +This guide demonstrates how to use the ROCK SDK to run SWE-Bench Verified evaluations, including sandbox setup, Agent integration, test environment preparation, and result parsing. + +### Quick Start + +SWE-Bench is a benchmark for evaluating AI coding agents on real-world software engineering tasks. + +Running a SWE-Bench task on ROCK involves the following steps: + +1. **load_task_config** — Load `task.yaml` to get the task instruction +2. **start_sandbox** — Start a sandbox with a task-specific Docker image +3. **agent.install / agent.run** — Install and run the Agent to solve the task +4. **setup_test_env** — Upload test files and run-test script to the sandbox +5. **Run tests** — Execute the test script via `sandbox.arun()` with timeout +6. **parse_swebench_result** — Parse test output to determine PASSED / FAILED +7. **sandbox.stop** — Clean up sandbox resources + +**Here is an example code** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. Load task instruction + task_config = await load_task_config(task_dir) # see load_task_config section + instruction = task_config["instruction"] + + # 2. Start sandbox + sandbox = await start_sandbox(task_name) # see start_sandbox section + + try: + # 3. Install and run Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. Setup test environment + await setup_test_env(sandbox, task_dir) # see setup_test_env section + + # 5. Run tests + resp = await run_tests(sandbox) # see Running Tests section + + # 6. Parse results + is_resolved = parse_swebench_result(resp.output) # see parse_swebench_result section + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +The following sections describe each function used in the workflow in detail. + +--- + +## start_sandbox + +Start a sandbox instance with a task-specific SWE-Bench Docker image. Each task has a pre-built image containing the target repository and environment. + +The `image` parameter follows the format: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +For example, task `django__django-14539` maps to: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +Load task configuration from a `task.yaml` file in the task directory. The YAML file contains the `instruction` field that describes the coding task for the Agent. + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# Usage +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +Use `sandbox.agent.install()` and `sandbox.agent.run()` to deploy and execute an Agent inside the sandbox. Refer to [Rock Agent](./rock-agent.md) for detailed Agent configuration. + +```python +# Install Agent with a YAML configuration file（e.g., iflow_config.yaml） +await sandbox.agent.install(config="iflow_config.yaml") + +# Run Agent with the task instruction +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +Prepare the test environment in the sandbox: install the [uv](https://github.com/astral-sh/uv) package manager, and upload test files and the run-test script. + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """Set up the test environment and return the session name.""" + # 1. Create a session with custom environment variables + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. Install uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. Upload test files + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. Upload run-tests script + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## Running Tests + +Execute the test script with a configurable timeout using `RunMode.NOHUP`. + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +Parse the test output to determine whether the SWE-Bench task is resolved. The parser looks for a result block delimited by marker lines and checks for `PASSED`. + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """Parse SWE-Bench test output to determine if the task is resolved. + + Matches the block between 'SWEBench results starts here' and + 'SWEBench results ends here', then checks whether it contains 'PASSED'. + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# Usage +is_resolved = parse_swebench_result(resp.output) +``` + +## Notes + +- **Task Datasets**: Task directories (containing `task.yaml`, `tests/`, and `run-tests.sh`) can be obtained from the [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) repository. +- **Task Images**: Each SWE-Bench task requires a specific Docker image (e.g., `sweb.eval.x86_64.`). Ensure the image is available before running tests. +- **Agent Config**: The Agent configuration YAML defines the runtime, dependencies, and execution command. See [Rock Agent](./rock-agent.md) for details. + diff --git a/docs/versioned_docs/version-1.5.x/References/api.md b/docs/versioned_docs/version-1.5.x/References/api.md new file mode 100644 index 0000000000..d73bf49d6c --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/References/api.md @@ -0,0 +1,195 @@ +--- +sidebar_position: 1 +--- + +# API Reference + +This guide provides detailed information about the core API services provided by the ROCK platform, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +The ROCK platform provides two core API services: +- Sandbox API: Sandbox environment management +- GEM API: GEM environment interaction + +All API interfaces follow RESTful design principles and support JSON format data exchange. + +## 2. Sandbox API + +Full lifecycle management functions for sandbox environments: + +### Sandbox Management Interfaces + +1. **Start Sandbox** - Start a sandbox environment + - Create a new sandbox instance + - Support specifying image, resource configuration and other parameters + +2. **Start Sandbox Async** - Asynchronously start a sandbox environment + - Asynchronously create a sandbox instance + - Suitable for scenarios requiring quick response + +3. **Check Sandbox Alive Status** - Check sandbox alive status + - Verify if the sandbox is running normally + +4. **Get Sandbox Statistics** - Get sandbox statistics + - Get resource usage statistics of the sandbox + +5. **Get Sandbox Status** - Get detailed sandbox status + - Get complete status information of the sandbox + +6. **Stop Sandbox** - Stop sandbox environment + - Safely shut down the sandbox instance + +7. **Commit Sandbox** - Commit sandbox as image + - Save current sandbox state as a new image + +### Command Execution Interfaces + +8. **Execute Command** - Execute command in sandbox + - Run specified command directly in the sandbox + +9. **Create Bash Session** - Create Bash session + - Create a persistent Bash session environment + +10. **Run Command in Session** - Run command in session + - Execute command in a created session + +11. **Close Session** - Close session + - Release session resources + +### File Operation Interfaces + +12. **Read File** - Read sandbox file + - Read specified file content from the sandbox + +13. **Write File** - Write sandbox file + - Write file to the sandbox + +14. **Upload File** - Upload file to sandbox + - Upload local file to the sandbox + +## 3. GEM API + +GEM environment interaction functions: + +1. **Make Environment** - Create GEM environment + - Initialize a new GEM environment instance + +2. **Reset Environment** - Reset GEM environment + - Reset GEM environment to initial state + +3. **Step Environment** - Execute GEM environment step + - Execute an action step in the GEM environment + +4. **Close Environment** - Close GEM environment + - Release GEM environment resources + + +## 4. HTTP API Usage Examples + +### 4.1 Sandbox API Examples + +#### Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Asynchronously Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Execute Command +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### Create Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### Run Command in Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### Upload File +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### Stop Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API Examples + +```bash +# Create GEM environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# Reset environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# Execute step +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# Close environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly get started with ROCK API +- [Python SDK Documentation](./Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to call APIs +- [Configuration Guide](../User%20Guides/configuration.md) - Learn about API-related configuration options +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/Release Notes/index.md b/docs/versioned_docs/version-1.5.x/Release Notes/index.md new file mode 100644 index 0000000000..e9bca06d4c --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# Release Notes +* [release v1.5.0](v1.5.0.md) diff --git a/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md new file mode 100644 index 0000000000..825fdc598c --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md @@ -0,0 +1,92 @@ +# ROCK v1.5.0 Release Note + +# v1.5.0 + +## Release Date + +April 10, 2026 + +--- + +## SDK + +### New Features + +#### TypeScript SDK + +* **NEW**: Brand-new TypeScript SDK, published to npm (`rl-rock`). Supports sandbox management, filesystem operations, sync/async shell command execution, Python/Node.js runtime environments, built-in Agent support, EnvHub integration, and ESM/CommonJS dual builds ([#492](https://github.com/alibaba/ROCK/pull/492)) + + +#### Agent Run (Harbor Style) + +* **NEW**: Support Agent Run, allowing users to submit and manage agent run tasks (harbor style) via the SDK ([#681](https://github.com/alibaba/ROCK/pull/681)) + + +#### Job OSS Artifact Mirror + +* **NEW**: Add `OssMirrorConfig` and `JobConfig.enable_oss_mirror()` to support Job OSS artifact mirroring, automatically populating namespace/experiment\_id from sandbox config ([#708](https://github.com/alibaba/ROCK/pull/708)) + + +#### Sandbox Client + +* **NEW**: Add `auto_delete_seconds` field to specify the retention time after a sandbox is stopped. Currently supports the following behaviors: + + * auto\_delete\_seconds = None (default): whether the sandbox is cleaned up after stop is determined by the cluster configuration + + * auto\_delete\_seconds = 0: the sandbox is deleted immediately after stop + + * auto\_delete\_seconds > 0: the sandbox is not deleted after stop + + +--- + +## Admin + +### New Features + +#### Sandbox Metadata Persistence + +* **NEW**: Support persisting sandbox metadata to a database backend (in addition to Redis) for more reliable data storage ([#730](https://github.com/alibaba/ROCK/pull/730)) + + +--- + +### Refactoring + +#### Ray Temp Directory + +* Add `temp_dir` field to `RayConfig` to support redirecting Ray temporary data directory (defaults to `.tmp/ray`), with automatic relative-to-absolute path resolution ([#694](https://github.com/alibaba/ROCK/pull/694), [#696](https://github.com/alibaba/ROCK/pull/696)) + + +#### User-Defined Log Path + +* Support user-defined log paths (e.g., `/data/logs/user-defined`) with automatic directory creation ([#702](https://github.com/alibaba/ROCK/pull/702)) + + +--- + +### Bug Fixes + +#### Memory Size Error Message + +* Fix incorrect memory size error message in sandbox manager ([#648](https://github.com/alibaba/ROCK/pull/648)) + + +#### Test Fixes + +* Fix tests that could not run ([#700](https://github.com/alibaba/ROCK/pull/700)) + +* Adjust sandbox resource limits in conftest.py to fix consistently failing unit tests ([#710](https://github.com/alibaba/ROCK/pull/710)) + +* Fix dirty sandbox info data caused by Kubernetes cache inconsistency across multiple admin instances ([#743](https://github.com/alibaba/ROCK/pull/743)) + +* Fix Kubernetes client informer not processing events in time in admin ([#744](https://github.com/alibaba/ROCK/pull/744)) + + +--- + +## CI / Infrastructure + +* Restore CI request-triggered workflow configuration ([#728](https://github.com/alibaba/ROCK/pull/728)) + +* Pin `langgraph-prebuilt` to 1.0.8 to fix CI errors ([#745](https://github.com/alibaba/ROCK/pull/745)) diff --git a/docs/versioned_docs/version-1.5.x/User Guides/configuration.md b/docs/versioned_docs/version-1.5.x/User Guides/configuration.md new file mode 100644 index 0000000000..604256878b --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/User Guides/configuration.md @@ -0,0 +1,189 @@ +--- +sidebar_position: 4 +--- + +# Configuration + +This guide provides detailed instructions on how to configure the ROCK environment to meet different usage requirements, including local development, testing, and production deployment. + +## 1. Environment Variable Configuration + +ROCK supports configuring key parameters through environment variables. The main environment variables are as follows: + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK service base URL +export ROCK_LOG_LEVEL=INFO # Log level +export ROCK_LOGGING_PATH=/path/to/logs # Log file path, default None (output to console) +export ROCK_LOGGING_FILE_NAME=rocklet.log # Log file name, default "rocklet.log", can be customized by admin like admin.log +export ROCK_LOGGING_LEVEL=INFO # Log output level, default "INFO" +export ROCK_WORKER_ENV_TYPE=local # Runtime environment type, options: local, docker, uv, pip +``` + +More environment variables can be found in `rock/env_vars.py`. + +### 1.1 Runtime Environments + +ROCK provides multiple different runtime environments to meet the needs of different scenarios, configured through the `ROCK_WORKER_ENV_TYPE` environment variable. Each environment has different deployment requirements, performance characteristics and applicable scenarios. Each environment has its own unique advantages and limitations, and developers can choose the most suitable runtime environment according to their deployment needs. + +#### 1.1.1 Docker Runtime Environment + +The Docker runtime environment is suitable for Docker image environments where dependencies are pre-installed. This environment requires the `/tmp/miniforge/bin/rocklet` executable to be directly available in the deployment environment. + +**Mount Configuration:** +- `/tmp/miniforge` - Contains pre-installed Python environment +- `/tmp/local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Containerized deployment environments +- Already built custom Docker image containing `rocklet` +- Suitable for production, fast startup + +**Requirements:** +- Requires a custom Docker image containing `/tmp/miniforge/bin/rocklet` executable +- Docker environment support + +#### 1.1.2 Local Runtime Environment + +The local runtime environment directly uses the Python environment and project files of the current deployment. This environment requires the same operating system between the host and container to directly mount the virtual environment and Python interpreter. + +**Mount Configuration:** +- `python_env_path` - Python environment path +- `project_root` - Project root directory +- `.venv` - Virtual environment directory (mounted as `/tmp/miniforge` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Development environments +- Scenarios where host and target container use the same operating system +- Need to quickly reuse existing Python environment + +**Requirements:** +- Same operating system (host/container) +- Direct access to the currently deployed `.venv` virtual environment +- Python interpreter path compatibility + +#### 1.1.3 UV Runtime Environment + +The UV runtime environment only depends on the available ROCK project, but initialization is relatively slow and network requirements are higher. This environment is most suitable for scenarios without preconfigured environments. It rebuilds the rocklet environment from the original project. This is the recommended environment for Mac OS. + +**Mount Configuration:** +- `project_root` - Project root directory (mounted as `/tmp + project_root` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**Use Cases:** +- Mac OS +- Cross-OS startup +- Scenarios without preconfigured environment +- No uv management Rock + +**Advantages:** +- No pre-built image required +- Good cross-platform compatibility +- Suitable for development and testing especially + +**Limitations:** +- Initialization is relatively slow +- Higher network requirements +- Longer startup time + +#### 1.1.4 PIP Runtime Environment + +The PIP runtime environment uses pip to install required dependencies in the container. This environment is suitable for quick setup and scenarios where dependencies can be installed in the container. It is the default runtime environment. It does not require pre-built images containing dependencies, and manages Python packages directly through pip. + +**Mount Configuration:** +- `local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**Use Cases:** +- ROCK installation from PIP source +- Fast testing of ROCK + +**Advantages:** +- Simple deployment setup + +**Limitations:** +- Long dependency installation time +- Requires network access to install dependency packages +- Dependencies need to be installed each time on startup + +#### 1.1.5 Configuration Guide + +Refer to the following selection guide for different use cases: + +| Scenario | Recommended Environment | Reason | +|----------|--------------------------|-------| +| Production environment | Docker Runtime | Fast startup, stable performance | +| Development environment, same OS | Local Runtime | Environment reuse, fast development cycle | +| Mac development | UV Runtime | Best cross-platform compatibility support | +| Cross-platform development | UV Runtime | Avoids environment compatibility issues | +| Fast testing | UV Runtime | Requires no pre-configuration | +| PIP source installation | PIP Runtime | Install dependencies directly with pip | + +These runtime environments are configured through the `ROCK_WORKER_ENV_TYPE` environment variable, which can be set to "local", "docker", "uv" or "pip". + +### 1.2 Logging Configuration + +Regarding logging configuration, ROCK's logging system has the following characteristics: + +- The logging system cannot output to both file and console simultaneously. If `ROCK_LOGGING_PATH` is set, logs will be output to the designated file, otherwise to console. +- `ROCK_LOGGING_LEVEL` is used to control the output log level, while `ROCK_LOG_LEVEL` is used for general log level settings. + +## 2. Distributed Deployment Requirements + +Since ROCK supports distributed deployment, when running on different nodes of a Ray cluster, the following consistency requirements must be met: + +#### Directory Structure Consistency + +On all Ray nodes, the following directory structure must be completely consistent: +- ROCK project repository directory +- `.venv` virtual environment directory +- The base Python directory that `.venv` depends on + + +#### Mounting Requirements + +ROCK's startup depends on mounting the ROCK project and the corresponding base Python environment, requiring consistency in multi-machine environments: + +#### Verifying Distributed Configuration + +Distributed deployment configuration can be verified through the following methods: + +```bash +# Check directory consistency on all nodes +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# Verify Python environment availability +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# Check environment variable settings on all nodes +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly set up the ROCK environment +- [API Documentation](../References/api.md) - View sandbox-related API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to configure sandboxes +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.5.x/overview.md b/docs/versioned_docs/version-1.5.x/overview.md new file mode 100644 index 0000000000..0c377e8b25 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/overview.md @@ -0,0 +1,33 @@ +--- +sidebar_position: 1 +--- + +# Overview + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework designed to simplify the development, deployment, and management of reinforcement learning environments. + +## What is ROCK + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework. By using ROCK, developers can quickly develop reinforcement learning environments and integrate with other reinforcement learning training frameworks to implement efficient reinforcement learning training. + +ROCK provides comprehensive sandbox environment management capabilities, supports containerized deployment, and enables rapid creation, execution, and destruction of environments. Additionally, ROCK is compatible with the GEM protocol, providing standardized interfaces for reinforcement learning environments. + +## Core Capabilities of ROCK + +1. **Simplified Development Process**: Simplifies the development, construction, and management of reinforcement learning environments, supporting various open-source reinforcement learning environments +2. **Large-scale Scheduling and Deployment**: Enables large-scale scheduling and deployment of rapid reinforcement learning environments. By supporting the GEM protocol, reinforcement learning environments can be easily accessed +3. **Framework Integration**: Integrates with other reinforcement learning training frameworks to achieve large-scale and scalable reinforcement learning training + +## Value of ROCK + +ROCK provides significant value to different roles of engineers: + +- **Reinforcement Learning Algorithm Engineers**: ROCK simplifies the development process of reinforcement learning environments, allowing engineers to focus on algorithm implementation +- **Reinforcement Learning Application Engineers**: ROCK enables large-scale deployment of rapid reinforcement learning environments, improving application development efficiency + +## Learn More + +- [Quick Start Guide](./Getting%20Started/quickstart.md) - Get started with ROCK quickly +- [Configuration Guide](./User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](./References/api.md) - View ROCK's API interfaces +- [Python SDK Documentation](./References/Python%20SDK%20References/python_sdk.md) - Learn how to use ROCK's Python SDK \ No newline at end of file diff --git a/docs/versioned_sidebars/version-1.5.x-sidebars.json b/docs/versioned_sidebars/version-1.5.x-sidebars.json new file mode 100644 index 0000000000..b475b11530 --- /dev/null +++ b/docs/versioned_sidebars/version-1.5.x-sidebars.json @@ -0,0 +1,64 @@ +{ + "tutorialSidebar": [ + "overview", + { + "type": "category", + "label": "Getting Started", + "link": { + "type": "doc", + "id": "Getting Started/quickstart" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Getting Started" + } + ] + }, + { + "type": "category", + "label": "User Guides", + "items": [ + { + "type": "autogenerated", + "dirName": "User Guides" + } + ] + }, + { + "type": "category", + "label": "References", + "items": [ + "References/api", + { + "type": "category", + "label": "Python SDK References", + "link": { + "type": "doc", + "id": "References/Python SDK References/python_sdk" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "References/Python SDK References" + } + ] + } + ] + }, + { + "type": "category", + "label": "Release Notes", + "link": { + "type": "doc", + "id": "Release Notes/index" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Release Notes" + } + ] + } + ] +} diff --git a/docs/versions.json b/docs/versions.json index c286c5f67c..09fda0e173 100644 --- a/docs/versions.json +++ b/docs/versions.json @@ -1,4 +1,5 @@ [ + "1.5.x", "1.4.x", "1.3.x", "1.2.x", diff --git a/pyproject.toml b/pyproject.toml index 2ae7df85d3..22de641bf6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.4.6" +version = "1.5.0" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 9dedee3116a34311b84ba6860c1c6898a236eb3b Mon Sep 17 00:00:00 2001 From: Generalwin <52099674+Generalwin@users.noreply.github.com> Date: Mon, 13 Apr 2026 10:01:39 +0800 Subject: [PATCH 022/226] fix missing redis provider in k8s (#765) --- rock/sandbox/operator/factory.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rock/sandbox/operator/factory.py b/rock/sandbox/operator/factory.py index d7289c6ee8..ff593523f6 100644 --- a/rock/sandbox/operator/factory.py +++ b/rock/sandbox/operator/factory.py @@ -71,6 +71,8 @@ def create_operator(context: OperatorContext) -> AbstractOperator: raise ValueError("K8sConfig is required for K8sOperator") logger.info("Creating K8sOperator") k8s_operator = K8sOperator(k8s_config=context.k8s_config) + if context.redis_provider is not None: + k8s_operator.set_redis_provider(context.redis_provider) if context.nacos_provider is not None: k8s_operator.set_nacos_provider(context.nacos_provider) return k8s_operator From 899a055ae4ce4bd4176c29425e70d5215433a337 Mon Sep 17 00:00:00 2001 From: Fangwen DAI <34794181+FangwenDave@users.noreply.github.com> Date: Mon, 13 Apr 2026 11:45:17 +0800 Subject: [PATCH 023/226] fix arun normal mode #767 (#768) * fix: auto-create temp session for arun normal mode when session is None Previously, calling arun(cmd, mode="normal") without a session would raise a Pydantic ValidationError. Now it auto-creates a temporary session, matching the behavior of nohup mode. Co-Authored-By: Claude Opus 4.6 * refactor: rename test_arun_nohup to test_arun The test file now covers both normal and nohup mode, so rename to reflect the broader scope. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- rock/sdk/sandbox/client.py | 5 +- .../sdk/{test_arun_nohup.py => test_arun.py} | 52 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) rename tests/unit/sdk/{test_arun_nohup.py => test_arun.py} (81%) diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 77a9dc719a..efa2dd85f5 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -398,7 +398,7 @@ async def arun( Args: cmd (str): The command to execute in the sandbox session (str, optional): The session identifier to run the command in. - If None, a temporary session will be created for nohup mode. Defaults to None. + If None, a temporary session will be created automatically. Defaults to None. wait_timeout (int, optional): Maximum time in seconds to wait for nohup command completion. Defaults to 300. wait_interval (int, optional): Interval in seconds between process completion checks for nohup mode. @@ -441,6 +441,9 @@ async def arun( raise InvalidParameterRockException(f"Unsupported arun mode: {mode}") if mode == RunMode.NORMAL: + if session is None: + session = await self._generate_tmp_session_name() + await self.create_session(CreateBashSessionRequest(session=session)) return await self._run_in_session(action=Action(command=cmd, session=session)) if mode == RunMode.NOHUP: return await self._arun_with_nohup( diff --git a/tests/unit/sdk/test_arun_nohup.py b/tests/unit/sdk/test_arun.py similarity index 81% rename from tests/unit/sdk/test_arun_nohup.py rename to tests/unit/sdk/test_arun.py index 096fbfae1c..6c7ab716d5 100644 --- a/tests/unit/sdk/test_arun_nohup.py +++ b/tests/unit/sdk/test_arun.py @@ -9,6 +9,58 @@ from rock.sdk.sandbox.config import SandboxConfig +@pytest.mark.asyncio +async def test_arun_normal_mode_without_session_creates_temp_session(monkeypatch): + """When arun is called in NORMAL mode without a session, it should auto-create one.""" + timestamp = 3001 + monkeypatch.setattr("rock.sdk.sandbox.client.time.time_ns", lambda: timestamp) + sandbox = Sandbox(SandboxConfig(image="mock-image")) + + created_sessions: list[str] = [] + captured_session = None + + async def fake_create_session(self, request): + created_sessions.append(request.session) + + async def fake_run_in_session(self, action): + nonlocal captured_session + captured_session = action.session + return Observation(output="ok", exit_code=0) + + sandbox.create_session = types.MethodType(fake_create_session, sandbox) # type: ignore + sandbox._run_in_session = types.MethodType(fake_run_in_session, sandbox) # type: ignore + + result = await sandbox.arun(cmd="ls -la", mode="normal") + + assert result.exit_code == 0 + assert result.output == "ok" + assert len(created_sessions) == 1 + assert created_sessions[0] == f"bash-{timestamp}" + assert captured_session == f"bash-{timestamp}" + + +@pytest.mark.asyncio +async def test_arun_normal_mode_with_session_does_not_create_new_session(monkeypatch): + """When arun is called in NORMAL mode with an existing session, it should NOT create a new one.""" + sandbox = Sandbox(SandboxConfig(image="mock-image")) + + session_created = False + + async def fake_create_session(self, request): + nonlocal session_created + session_created = True + + async def fake_run_in_session(self, action): + return Observation(output="ok", exit_code=0) + + sandbox.create_session = types.MethodType(fake_create_session, sandbox) # type: ignore + sandbox._run_in_session = types.MethodType(fake_run_in_session, sandbox) # type: ignore + + await sandbox.arun(cmd="ls -la", session="bash-existing", mode="normal") + + assert not session_created + + @pytest.mark.asyncio async def test_arun_nohup_ignore_output_true_returns_hint(monkeypatch): timestamp = 1701 From ffdee2781ba50f6469a6d5b3a2fe3a879acbe167 Mon Sep 17 00:00:00 2001 From: Shuaibing Zhao Date: Mon, 13 Apr 2026 14:10:32 +0800 Subject: [PATCH 024/226] Feature/add bash example (#772) * add bash demo * install rl-rock firstly * log sandbox id * add log for bash demo * demo show variables * add workspace * add vitabench * add job command * make local-path optional * add token params * refine simple_bash_job_demo * fix vitabench_demo * add simple_bash dir * delete vita --- examples/bash/simple_bash_job_demo.sh | 94 ++++++++++++++++++++ rock/cli/command/job.py | 121 ++++++++++++++++++++++++++ 2 files changed, 215 insertions(+) create mode 100644 examples/bash/simple_bash_job_demo.sh create mode 100644 rock/cli/command/job.py diff --git a/examples/bash/simple_bash_job_demo.sh b/examples/bash/simple_bash_job_demo.sh new file mode 100644 index 0000000000..969f22b243 --- /dev/null +++ b/examples/bash/simple_bash_job_demo.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# Simple bash job demo using `rock job run` +# +# This script defines a bash job inline and executes it inside a ROCK sandbox +# via the `rock job run` CLI command. + +set -euo pipefail + +# ===== Configuration ===== +# Override via environment variables: ROCK_BASE_URL, YOUR_API_KEY, YOUR_USER_ID, YOUR_EXPERIMENT_ID +ROCK_BASE_URL="${ROCK_BASE_URL}" +YOUR_API_KEY="${YOUR_API_KEY}" +YOUR_USER_ID="${YOUR_USER_ID}" +YOUR_EXPERIMENT_ID="${YOUR_EXPERIMENT_ID:-simple_bash}" +ROCK_IMAGE="${ROCK_IMAGE:-rl-rock-registry-vpc.ap-southeast-1.cr.aliyuncs.com/chatos/base:python3.11}" +ROCK_CLUSTER="${ROCK_CLUSTER:-vpc-sg-sl-a}" + +LOCAL_WORKSPACE_DIR="${LOCAL_WORKSPACE_DIR:-}" +ROCK_WORKSPACE_DIR="${ROCK_WORKSPACE_DIR:-/root/workspace}" + +EXTERNAL_VARIABLE_1="external_value" +TO_RENDERED_KEYS=( + "EXTERNAL_VARIABLE_1" +) +# ========================= + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +# Define the bash job script content +# This is the script that will be uploaded and executed inside the sandbox +read -r -d '' BASH_SCRIPT << 'EOF' || true +#!/bin/bash +echo "=== Simple Bash Job Demo ===" +echo "" +echo "Hostname: $(hostname)" +echo "Date: $(date)" +echo "" + +# --- Pattern 1: Internal Variables --- +INTERNAL_VARIABLE_1="internal_value" +echo "--- Internal Variables ---" +echo "INTERNAL_VARIABLE_1: ${INTERNAL_VARIABLE_1}" +echo "" + +# --- Pattern 2: External Variables (rendered into script content) --- +echo "--- External Variables (rendered) ---" +echo "EXTERNAL_VARIABLE_1: ${EXTERNAL_VARIABLE_1}" +echo "" + +echo "Running a simple computation..." +for i in $(seq 1 5); do + echo " Step $i: processing..." + sleep 1 +done +echo "" +echo "Job completed successfully!" +EOF + +# Render the bash script by replacing placeholders with actual values from environment variables +for key in "${TO_RENDERED_KEYS[@]}"; do + value="${!key}" + BASH_SCRIPT="${BASH_SCRIPT//\$\{$key\}/$value}" +done + +echo "Starting simple bash job demo..." +echo "Bash Script:" +echo "========================================" +echo "$BASH_SCRIPT" +echo "========================================" + +# Install rl-rock if not already available +if ! python -c "import rock" 2>/dev/null; then + echo "Installing rl-rock..." + pip install rl-rock +else + echo "rl-rock is already installed." +fi + +# Run the job via ROCK CLI +run_args=( + --base-url "$ROCK_BASE_URL" + --extra-header "XRL-Authorization=Bearer ${YOUR_API_KEY}" + --cluster "$ROCK_CLUSTER" + job run + --image "$ROCK_IMAGE" + --timeout 3600 + --script-content "$BASH_SCRIPT" +) + +if [ -n "$LOCAL_WORKSPACE_DIR" ]; then + run_args+=(--local-path "$LOCAL_WORKSPACE_DIR" --target-path "$ROCK_WORKSPACE_DIR") +fi + +rock "${run_args[@]}" diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py new file mode 100644 index 0000000000..64666cd78b --- /dev/null +++ b/rock/cli/command/job.py @@ -0,0 +1,121 @@ +import argparse +from pathlib import Path + +from rock.cli.command.command import Command +from rock.logger import init_logger +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +logger = init_logger(__name__) + + +class JobCommand(Command): + name = "job" + + async def arun(self, args: argparse.Namespace): + if args.job_command == "run": + await self._job_run(args) + else: + logger.error(f"Unknown job subcommand: {args.job_command}") + + async def _job_run(self, args: argparse.Namespace): + # 1. Validate script source + if args.script and args.script_content: + logger.error("--script and --script-content cannot be used together") + return + if not args.script and not args.script_content: + logger.error("Either --script or --script-content is required") + return + + if args.script: + script_path = Path(args.script).resolve() + if not script_path.exists(): + logger.error(f"Script not found: {script_path}") + return + if not script_path.is_file(): + logger.error(f"Not a file: {script_path}") + return + script_content = script_path.read_text() + else: + script_content = args.script_content + + # 2. Validate local path (optional) + src_dir = None + target_path = args.target_path + if args.local_path: + local_path = Path(args.local_path).resolve() + if not local_path.exists(): + logger.error(f"Local path not found: {local_path}") + return + src_dir = str(local_path) + + # 3. Build sandbox config + sandbox_config = SandboxConfig() + if args.image: + sandbox_config.image = args.image + if args.memory: + sandbox_config.memory = args.memory + if args.cpus: + sandbox_config.cpus = args.cpus + if args.base_url: + sandbox_config.base_url = args.base_url + if args.cluster: + sandbox_config.cluster = args.cluster + if args.extra_headers: + sandbox_config.extra_headers.update(args.extra_headers) + + sandbox = Sandbox(sandbox_config) + + try: + # 4. Start sandbox + logger.info(f"Starting sandbox with image={sandbox_config.image} ...") + await sandbox.start() + logger.info(f"Sandbox started: id={sandbox.sandbox_id}, ip={sandbox.host_ip}") + + # 5. Copy source directory to sandbox (optional) + if src_dir is not None: + assert sandbox.fs is not None + logger.info(f"Uploading {src_dir} -> {target_path}") + result = await sandbox.fs.upload_dir(source_dir=src_dir, target_dir=target_path) + if result.exit_code != 0: + logger.error(f"Upload failed: {result.failure_reason}") + return + + # 6. Execute the script + assert sandbox.process is not None + logger.info(f"Executing script: {script_content}") + result = await sandbox.process.execute_script( + script_content=script_content, + wait_timeout=args.timeout, + ) + + # 7. Print output + if result.output: + print(result.output) + logger.info(f"Script exited with code: {result.exit_code}") + + finally: + logger.info("Stopping sandbox ...") + await sandbox.stop() + + @staticmethod + async def add_parser_to(subparsers: argparse._SubParsersAction): + job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") + job_subparsers = job_parser.add_subparsers(dest="job_command") + + # run subcommand + run_parser = job_subparsers.add_parser("run", help="Run a job in a sandbox") + run_parser.add_argument("--image", default=None, help="Sandbox image (overrides default)") + run_parser.add_argument("--memory", default=None, help="Memory allocation (e.g., 8g)") + run_parser.add_argument("--cpus", default=None, type=float, help="CPU allocation (e.g., 2)") + + run_parser.add_argument("--local-path", default=None, help="Local directory to upload to the sandbox") + run_parser.add_argument( + "--target-path", default="/root/job", help="Target directory in sandbox (default: /root/job)" + ) + run_parser.add_argument("--script", default=None, help="Path to the script to execute in the sandbox") + run_parser.add_argument("--script-content", default=None, help="Script content to execute directly") + + run_parser.add_argument( + "--timeout", type=int, default=3600, help="Script execution timeout in seconds (default: 3600)" + ) From 1df699fe2b7bd611c5a0457515f543e9ce006a28 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Mon, 13 Apr 2026 14:33:24 +0800 Subject: [PATCH 025/226] Bump master version to 1.5.1 and update 1.5.1 release note (#774) * Bump version to 1.5.1 (#769) Signed-off-by: Jiachen Zhang (cherry picked from commit d76bf557bf43aa4601738803befb3c9053900084) Signed-off-by: Jiachen Zhang * Release note v1.5.1 Co-Authored-By: Claude Opus 4.6 * chore: update 1.5.0 release note title Signed-off-by: Jiachen Zhang --------- Signed-off-by: Jiachen Zhang Co-authored-by: Claude Opus 4.6 --- .../version-1.5.x/Release Notes/index.md | 1 + .../version-1.5.x/Release Notes/v1.5.0.md | 2 -- .../version-1.5.x/Release Notes/v1.5.1.md | 15 +++++++++++++++ .../version-1.5.x/Release Notes/index.md | 1 + .../version-1.5.x/Release Notes/v1.5.0.md | 2 -- .../version-1.5.x/Release Notes/v1.5.1.md | 15 +++++++++++++++ pyproject.toml | 2 +- 7 files changed, 33 insertions(+), 5 deletions(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.1.md create mode 100644 docs/versioned_docs/version-1.5.x/Release Notes/v1.5.1.md diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md index 44591bd640..2729189ed3 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/index.md @@ -2,4 +2,5 @@ sidebar_position: 1 --- # 版本说明 +* [release v1.5.1](v1.5.1.md) * [release v1.5.0](v1.5.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md index a4c00496cf..eebefb1d37 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.0.md @@ -1,5 +1,3 @@ -# ROCK v1.5.0 Release Note - # v1.5.0 ## 发布日期 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.1.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.1.md new file mode 100644 index 0000000000..cd11344d26 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.5.x/Release Notes/v1.5.1.md @@ -0,0 +1,15 @@ +# v1.5.1 + +## 发布日期 + +2026 年 4 月 13 日 + +--- + +## Admin + +### Bug 修复 + +#### K8s Operator 缺失 Redis Provider + +* 修复 `OperatorFactory` 初始化 `K8sOperator` 时未传入 `redis_provider` 的问题，导致使用 Kubernetes 后端时 Redis 相关功能不可用 ([#765](https://github.com/alibaba/ROCK/pull/765)) diff --git a/docs/versioned_docs/version-1.5.x/Release Notes/index.md b/docs/versioned_docs/version-1.5.x/Release Notes/index.md index e9bca06d4c..39d63ecf45 100644 --- a/docs/versioned_docs/version-1.5.x/Release Notes/index.md +++ b/docs/versioned_docs/version-1.5.x/Release Notes/index.md @@ -2,4 +2,5 @@ sidebar_position: 1 --- # Release Notes +* [release v1.5.1](v1.5.1.md) * [release v1.5.0](v1.5.0.md) diff --git a/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md index 825fdc598c..0a0e57c05a 100644 --- a/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md +++ b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.0.md @@ -1,5 +1,3 @@ -# ROCK v1.5.0 Release Note - # v1.5.0 ## Release Date diff --git a/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.1.md b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.1.md new file mode 100644 index 0000000000..5cdb9d8d52 --- /dev/null +++ b/docs/versioned_docs/version-1.5.x/Release Notes/v1.5.1.md @@ -0,0 +1,15 @@ +# v1.5.1 + +## Release Date + +April 13, 2026 + +--- + +## Admin + +### Bug Fixes + +#### Missing Redis Provider in K8s Operator + +* Fix `K8sOperator` not receiving `redis_provider` during initialization in `OperatorFactory`, which caused Redis-dependent features to be unavailable when using the Kubernetes backend ([#765](https://github.com/alibaba/ROCK/pull/765)) diff --git a/pyproject.toml b/pyproject.toml index 22de641bf6..91f2d62faf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.5.0" +version = "1.5.1" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 58c5f76593ad6dd99a4c86e759a6de68b662fa01 Mon Sep 17 00:00:00 2001 From: dengwx Date: Tue, 14 Apr 2026 14:35:40 +0800 Subject: [PATCH 026/226] [FEATURE] Refactor Job module: Job/Operator/Executor/Trial abstraction (#779) (#780) * feat: add docs * feat(job): add result models (TaskResult, JobResult) Add TaskStatus, TaskResult, JobStatus, and JobResult Pydantic models for the new Job system. Includes computed properties for score, success, n_completed, n_failed, and Harbor backward-compatible trial_results. Co-Authored-By: Claude Opus 4.6 (1M context) * feat(job): add config hierarchy (JobConfig, BashJobConfig, HarborJobConfig) Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(job): JobResult generic over T, JobStatus stays in agent module * docs(job): rename Task concept to Trial in README.md and plan.md * feat: refine config * refactor(job): agent JobConfig inherits base JobConfig, remove HarborJobConfig from rock/sdk/job * rename sdk/agent to sdk/bench * rename sdk/agent to sdk/bench * feat(job): add AbstractTrial and trial registry * fix(job): check upload_dir exit_code, tighten registry types * feat(job): add BashTrial with auto-registration * feat(job): add Operator ABC and ScatterOperator * feat(job): add JobExecutor with TrialClient/JobClient * feat(job): add Job facade * feat(job): add HarborTrial (extracted from bench/job.py) * feat(job): update CLI with --type bash/harbor routing * feat(job): integration wiring and public exports - Wire rock/sdk/job/__init__.py to export Job, configs, results, operator, executor, trial; auto-register BashTrial. - Wire rock/sdk/job/trial/__init__.py to export AbstractTrial and registry. - Register HarborTrial at the end of rock/sdk/bench/__init__.py (after bench.models.trial.result is fully loaded) to avoid a circular import when rock.sdk.job is loaded mid-bench init. - Pre-import rock.sdk.bench in rock.sdk.job.__init__ so that a cold `import rock.sdk.job` triggers bench-side registration too. - Add tests/unit/sdk/job/test_integration.py covering public imports, trial auto-registration, and rock.sdk.bench backward compat. - Minor ruff format fixup in trial/registry.py (merged adjacent f-strings). * refactor(job): extract _job_tmp_prefix helper in JobExecutor * refactor(job): make _job_tmp_prefix a method on JobExecutor * refactor(job): rename job.py to facade.py to avoid package-module name collision * rename facada.py to api.py * fix(test): make test_ray_get resilient to leaked detached actors - Use UUID-based actor name to avoid collisions from prior runs/reruns - Specify namespace at creation time to match the lookup namespace - Wrap in try/finally with ray.kill() to ensure cleanup on failure - Fixes CI flakiness: 'name test (namespace=None) is already taken' --------- Co-authored-by: Claude Opus 4.6 (1M context) --- docs/dev/job/README.md | 905 ++++++++ docs/dev/job/operator.md | 363 +++ docs/dev/job/plan.md | 1938 +++++++++++++++++ examples/harbor/harbor_demo.py | 2 +- rock/cli/command/job.py | 162 +- rock/sdk/agent/models/job/result.py | 43 - rock/sdk/{agent => bench}/__init__.py | 19 +- rock/sdk/{agent => bench}/constants.py | 0 rock/sdk/{agent => bench}/job.py | 8 +- rock/sdk/{agent => bench}/models/__init__.py | 14 +- .../models/environment_type.py | 0 .../{agent => bench}/models/job/__init__.py | 5 +- .../sdk/{agent => bench}/models/job/config.py | 55 +- .../models/metric/__init__.py | 0 .../{agent => bench}/models/metric/config.py | 2 +- .../{agent => bench}/models/metric/type.py | 0 .../models/orchestrator_type.py | 0 .../{agent => bench}/models/trial/__init__.py | 0 .../{agent => bench}/models/trial/config.py | 2 +- .../{agent => bench}/models/trial/result.py | 44 +- rock/sdk/job/__init__.py | 34 + rock/sdk/job/api.py | 74 + rock/sdk/job/config.py | 35 + rock/sdk/job/executor.py | 143 ++ rock/sdk/job/operator.py | 46 + rock/sdk/job/result.py | 106 + rock/sdk/job/trial/__init__.py | 4 + rock/sdk/job/trial/abstract.py | 43 + rock/sdk/job/trial/bash.py | 49 + rock/sdk/job/trial/harbor.py | 116 + rock/sdk/job/trial/registry.py | 29 + tests/unit/admin/core/test_ray_service.py | 32 +- tests/unit/sdk/agent/test_job.py | 8 +- .../agent/test_job_config_serialization.py | 25 +- .../sdk/agent/test_jobconfig_experiment_id.py | 6 +- tests/unit/sdk/agent/test_models.py | 18 +- tests/unit/sdk/agent/test_oss_mirror.py | 46 +- tests/unit/sdk/job/__init__.py | 0 tests/unit/sdk/job/test_cli_job.py | 283 +++ tests/unit/sdk/job/test_config.py | 309 +++ tests/unit/sdk/job/test_executor.py | 221 ++ tests/unit/sdk/job/test_integration.py | 66 + tests/unit/sdk/job/test_job.py | 169 ++ tests/unit/sdk/job/test_operator.py | 66 + tests/unit/sdk/job/test_result.py | 70 + tests/unit/sdk/job/test_trial_bash.py | 135 ++ tests/unit/sdk/job/test_trial_harbor.py | 144 ++ tests/unit/sdk/job/test_trial_registry.py | 132 ++ uv.lock | 2 +- 49 files changed, 5686 insertions(+), 287 deletions(-) create mode 100644 docs/dev/job/README.md create mode 100644 docs/dev/job/operator.md create mode 100644 docs/dev/job/plan.md delete mode 100644 rock/sdk/agent/models/job/result.py rename rock/sdk/{agent => bench}/__init__.py (60%) rename rock/sdk/{agent => bench}/constants.py (100%) rename rock/sdk/{agent => bench}/job.py (98%) rename rock/sdk/{agent => bench}/models/__init__.py (62%) rename rock/sdk/{agent => bench}/models/environment_type.py (100%) rename rock/sdk/{agent => bench}/models/job/__init__.py (77%) rename rock/sdk/{agent => bench}/models/job/config.py (81%) rename rock/sdk/{agent => bench}/models/metric/__init__.py (100%) rename rock/sdk/{agent => bench}/models/metric/config.py (81%) rename rock/sdk/{agent => bench}/models/metric/type.py (100%) rename rock/sdk/{agent => bench}/models/orchestrator_type.py (100%) rename rock/sdk/{agent => bench}/models/trial/__init__.py (100%) rename rock/sdk/{agent => bench}/models/trial/config.py (98%) rename rock/sdk/{agent => bench}/models/trial/result.py (74%) create mode 100644 rock/sdk/job/__init__.py create mode 100644 rock/sdk/job/api.py create mode 100644 rock/sdk/job/config.py create mode 100644 rock/sdk/job/executor.py create mode 100644 rock/sdk/job/operator.py create mode 100644 rock/sdk/job/result.py create mode 100644 rock/sdk/job/trial/__init__.py create mode 100644 rock/sdk/job/trial/abstract.py create mode 100644 rock/sdk/job/trial/bash.py create mode 100644 rock/sdk/job/trial/harbor.py create mode 100644 rock/sdk/job/trial/registry.py create mode 100644 tests/unit/sdk/job/__init__.py create mode 100644 tests/unit/sdk/job/test_cli_job.py create mode 100644 tests/unit/sdk/job/test_config.py create mode 100644 tests/unit/sdk/job/test_executor.py create mode 100644 tests/unit/sdk/job/test_integration.py create mode 100644 tests/unit/sdk/job/test_job.py create mode 100644 tests/unit/sdk/job/test_operator.py create mode 100644 tests/unit/sdk/job/test_result.py create mode 100644 tests/unit/sdk/job/test_trial_bash.py create mode 100644 tests/unit/sdk/job/test_trial_harbor.py create mode 100644 tests/unit/sdk/job/test_trial_registry.py diff --git a/docs/dev/job/README.md b/docs/dev/job/README.md new file mode 100644 index 0000000000..9b3f4626a6 --- /dev/null +++ b/docs/dev/job/README.md @@ -0,0 +1,905 @@ +# Rock Job 架构重构方案 + +## 1. 现状分析 + +### 1.1 当前实现 + +| 组件 | 位置 | 职责 | 问题 | +|------|------|------|------| +| `JobCommand` | `rock/cli/command/job.py` | CLI `rock job run`，直接用 Sandbox SDK 执行脚本 | 功能简单，无抽象，不支持 Harbor | +| `Job` | `rock/sdk/agent/job.py` | Harbor benchmark 执行，完整 submit/wait 生命周期 | 与 Harbor 强耦合，不可扩展到其他任务类型 | +| `JobConfig` | `rock/sdk/agent/models/job/config.py` | Harbor 配置 schema（从 Harbor 拷贝） | 既有 Rock 环境字段又有 Harbor 字段，混合关注点 | +| `JobResult` | `rock/sdk/agent/models/job/result.py` | Harbor 结果模型 | 仅支持 Harbor trial 结构 | + +### 1.2 核心问题 + +1. **Job = Harbor** — 当前 `Job` 类硬编码了 Harbor 的 YAML 生成、dockerd 启动、trial 结果收集逻辑，无法支持纯 Bash 脚本任务 +2. **CLI 与 SDK 断层** — `rock job run`（CLI）直接操作 Sandbox，完全不经过 Job SDK +3. **无调度抽象** — 单一 sandbox 执行，无 map/并行/分片能力 +4. **无数据抽象** — 数据输入硬编码在 Harbor dataset config 中，无法支持 CSV/Pandas 等通用数据源 + +### 1.3 重构目标 + +``` +用户视角: + rock job run --type bash --script train.sh # BashJob via CLI + rock job run --type harbor --config harbor.yaml # HarborJob via CLI + rock job run --type bash --script eval.sh --map-from data.csv # BashJob + Map调度 + +SDK 视角: + job = BashJob(config) # 简单脚本执行 + job = HarborJob(config) # Harbor benchmark + results = await asyncio.gather(*[Job(c).run() for c in configs]) # 并行 + result = await job.run() +``` + +### 1.4 使用场景 + +| 调用方 | 场景 | 需要的能力 | +|--------|------|-----------| +| **rock CLI** | `rock job run` 执行脚本或 Harbor benchmark | BashJob、HarborJob，简单配置 | +| **verl / roll** | RL 训练 rollout，批量提交评测任务 | HarborJob，submit/wait 异步，并发执行 | +| **评测平台** | 批量评测多数据集，收集结果 | Map 调度，数据源输入，结果聚合 | +| **数据处理** | 分布式 data pipeline | BashJob + Map，CSV/Pandas 输入，分片并行 | + +--- + +## 2. 架构方案: Config 驱动类型 + Trial 内部抽象 + Job 编排层 + +**核心思路**: Job 是 Facade 入口。JobExecutor 是主控方，驱动 Operator（算子）。Operator 决定分发 8 份 Trial（类比 torch.distributed.scatter），JobExecutor 并行执行 TrialList。 + +``` +┌────────────────────────────────────────────────────────────┐ +│ Job (Facade) — 极薄入口 │ +│ Job(config, operator?).run() / .submit() / .wait() │ +│ │ +│ ┌──────────────────────────────────────────────────────┐ │ +│ │ JobExecutor (主控方) │ │ +│ │ │ │ +│ │ run = submit + wait │ │ +│ │ │ │ +│ │ submit(operator, config): │ │ +│ │ 1. operator.apply(config) │ │ +│ │ → 创建 Trial, 决定分发数量 │ │ +│ │ → 返回 list[AbstractTrial] (TrialList) │ │ +│ │ 2. 并行 _(do_submit(trial) for each task │ │ +│ │ → sandbox + trial.setup/build + nohup │ │ +│ │ → 返回 JobClient │ │ +│ │ │ │ +│ │ wait(job_client): │ │ +│ │ → _do_wait(tc) for each TrialClient │ │ +│ │ → trial.collect(sandbox, output, exit_code) │ │ +│ │ → 返回 list[TrialResult] │ │ +│ └──────────────────────────────────────────────────────┘ │ +│ │ +│ Operator (算子, 可替换): │ +│ ├── ScatterOperator(size=8) — 分发 8 份 Trial │ +│ └── (future: RayScatterOperator, K8sOperator, ...) │ +│ │ +│ Trial (内部, 由 Config 子类决定): │ +│ ├── BashTrial ← BashJobConfig │ +│ ├── HarborTrial ← HarborJobConfig │ +│ └── (扩展...) │ +└────────────────────────────────────────────────────────────┘ +``` + +#### 设计原则 + +1. **Job = 极薄 Facade** — 只有 `config` + `operator` 两个参数 +2. **JobExecutor = 主控方** — 驱动 Operator 生成 TrialList，并行执行，收集结果 +3. **Operator = 算子 (通用接口 apply)** — 从 config 生成 TrialList，ScatterOperator 是默认实现 +4. **Trial = 任务逻辑** — 数据 IO、脚本生成、结果解析，全部由 Trial 自行控制 +5. **Config 子类 = 类型** — `BashJobConfig` / `HarborJobConfig` 携带类型信息 + +#### 调用链 + +``` +Job.run() = Job.submit() + Job.wait() + +submit 阶段: + Job.submit() + → JobExecutor.submit(operator, config) + 1. trial_list = operator.apply(config) # Operator 决定分发 8 份 Trial + → _create_trial(config) # 创建 Trial + → 返回 [trial] * size # size=0 返回 [], 什么都不做 + 2. 并行 _(do_submit(trial) for trial in trial_list # Executor 并行启动 + → 返回 JobClient(tasks=[TrialClient, ...]) + +wait 阶段: + Job.wait() + → JobExecutor.wait(job_client) + → _do_wait(tc) for each TrialClient # 等待 + trial.collect() + → 返回 list[TrialResult] # size=0 时返回 [] +``` + +**关键**: +- **Operator 只做 apply** — 生成 TrialList，不启动 sandbox，不管执行 +- **JobExecutor 做并行执行** — 拿到 TrialList 后并行 submit + 并行 wait +- **size=0 安全返回** — Operator 返回空 list，Executor 什么都不执行，Job 返回空结果 + +#### 职责分离 + +| 组件 | 职责 | 不负责 | 类比 | +|------|------|--------|------| +| **Job** | 极薄 Facade: 接收 config，组装组件 | 不包含任何执行/调度逻辑 | — | +| **Operator** | 算子: apply(config) → TrialList | 不启动 sandbox，不执行 | torch.distributed.scatter | +| **JobExecutor** | 主控方: 并行执行 TrialList，管理 sandbox 生命周期 | 不决定分发数量 | Executor | +| **Trial** | 任务逻辑: setup/build/collect | 不管理 sandbox 生命周期 | UDF | + +#### 类设计 + +```python +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Config 层: 用户直接接触，Config 子类决定 Job 类型 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class JobConfig(BaseModel): + """Job 基础配置 — 所有 Job 类型共享的字段""" + environment: RockEnvironmentConfig # sandbox 环境 + job_name: str | None = None + namespace: str | None = None + experiment_id: str | None = None + labels: dict[str, str] = {} + auto_stop: bool = False + setup_commands: list[str] = [] # 前置命令 + file_uploads: list[tuple[str, str]] = [] # 文件上传 + env: dict[str, str] = {} # 环境变量 + timeout: int = 3600 # 超时 + + +class BashJobConfig(JobConfig): + """Bash 脚本 Job 配置""" + script: str | None = None # 脚本内容 (与 script_path 二选一) + script_path: str | None = None # 脚本文件路径 + + +class HarborJobConfig(JobConfig): + """Harbor benchmark Job 配置 + + Harbor 原生字段直接平铺（与现有 rock.sdk.agent.JobConfig 的 Harbor 字段一致）， + 序列化为 YAML 传给 harbor jobs start -c + """ + agents: list[AgentConfig] = Field(default_factory=lambda: [AgentConfig()]) + datasets: list[LocalDatasetConfig | RegistryDatasetConfig] = Field(default_factory=list) + orchestrator: OrchestratorConfig = Field(default_factory=OrchestratorConfig) + verifier: VerifierConfig = Field(default_factory=VerifierConfig) + tasks: list[TaskConfig] = Field(default_factory=list) + metrics: list[MetricConfig] = Field(default_factory=list) + artifacts: list[str | ArtifactConfig] = Field(default_factory=list) + n_attempts: int = 1 + timeout_multiplier: float = 1.0 + agent_timeout_multiplier: float | None = None + verifier_timeout_multiplier: float | None = None + jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" + debug: bool = False + + def to_harbor_yaml(self) -> str: ... + + @classmethod + def from_yaml(cls, path: str) -> HarborJobConfig: ... + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Trial 层: 定义"在 sandbox 中做什么"，三阶段接口 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class AbstractTrial(ABC): + """Task: 定义在 Sandbox 中执行什么 + + 三阶段接口: + setup() — 执行前: 准备环境 (上传文件、写入配置) + build() — 构建: 生成要执行的脚本 + collect() — 执行后: 从 sandbox 收集解析结果 + + Trial 不管理 sandbox 生命周期 (由 JobExecutor 负责)， + 也不关心调度策略 (由 Operator 负责)。 + """ + + def __init__(self, config: JobConfig): + self._config = config + + @abstractmethod + async def setup(self, sandbox: Sandbox) -> None: + """执行前: 准备 sandbox 环境 — 上传文件、写入配置""" + + @abstractmethod + def build(self) -> str: + """构建: 生成要在 sandbox 中执行的 bash 脚本""" + + @abstractmethod + async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: + """执行后: 从 sandbox 收集并解析结果""" + + async def _upload_files(self, sandbox: Sandbox) -> None: + """共享: 上传 config.file_uploads""" + for local_path, sandbox_path in self._config.file_uploads: + await sandbox.fs.upload_dir(local_path, sandbox_path) + + +class BashTrial(AbstractTrial): + """Bash 脚本执行""" + + async def setup(self, sandbox): + await self._upload_files(sandbox) + if self._config.script_path: + self._config.script = Path(self._config.script_path).read_text() + + def build(self) -> str: + setup = "\n".join(self._config.setup_commands) or "true" + return f"#!/bin/bash\nset -e\n{setup}\n{self._config.script}" + + async def collect(self, sandbox, output, exit_code): + return TrialResult( + task_id=self._config.job_name or "", + status=TrialStatus.COMPLETED if exit_code == 0 else TrialStatus.FAILED, + output=output, + exit_code=exit_code, + ) + + +class HarborTrial(AbstractTrial): + """Harbor benchmark 执行 (重构自现有 rock.sdk.agent.job.Job)""" + + async def setup(self, sandbox): + await self._upload_files(sandbox) + yaml_content = self._config.to_harbor_yaml() + await self._upload_content(sandbox, yaml_content, self._config_path) + + def build(self) -> str: + return _HARBOR_SCRIPT_TEMPLATE.format(...) + + async def collect(self, sandbox, output, exit_code): + trial_results = await self._collect_trial_results(sandbox) + return TrialResult( + task_id=self._config.job_name or "", + status=TrialStatus.COMPLETED if trial_results else TrialStatus.FAILED, + output=output, exit_code=exit_code, + trial_results=trial_results, + ) + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Trial Registry: Config 类型 → Trial 实现 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +_TRIAL_REGISTRY: dict[type[JobConfig], type[AbstractTrial]] = { + BashJobConfig: BashTrial, + HarborJobConfig: HarborTrial, +} + +def register_trial(config_type: type[JobConfig], task_type: type[AbstractTrial]): + """注册新的 Config → Task 映射 (扩展点)""" + _TRIAL_REGISTRY[config_type] = task_type + +def _create_trial(config: JobConfig) -> AbstractTrial: + """根据 config 类型创建对应的 Task 实例""" + task_cls = _TRIAL_REGISTRY.get(type(config)) + if task_cls is None: + raise TypeError( + f"No task registered for {type(config).__name__}. " + f"Supported: {[c.__name__ for c in _TRIAL_REGISTRY]}" + ) + return task_cls(config) + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Operator: 算子基类 — 通用接口 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class Operator(ABC): + """算子基类: 从 config 生成 TrialList + + 通用接口 apply(config) → list[AbstractTrial]: + - 输入: JobConfig + - 输出: TrialList (待执行的 Task 列表) + - 返回空 list 表示什么都不做 + + Operator 只负责生成 TrialList，不启动 sandbox，不执行。 + 执行由 JobExecutor 负责。 + """ + + @abstractmethod + def apply(self, config: JobConfig) -> list[AbstractTrial]: + """从 config 生成 TrialList + + 返回空 list 表示什么都不做。 + """ + ... + + +class ScatterOperator(Operator): + """Scatter 算子: 将 config 分发 8 份 Trial + + 类比 torch.distributed.scatter: + scatter_list = [tensor] * size → 每个 rank 拿一份 + trial_list = [trial] * size → 每个 sandbox 执行一份 + + 用法: + ScatterOperator() # size=1，默认 1 份 Task + ScatterOperator(size=8) # 8 份 Task，由 JobExecutor 并行执行 + ScatterOperator(size=0) # 什么都不做，返回 [] + """ + + def __init__(self, size: int = 1): + self.size = size + + def apply(self, config) -> list[AbstractTrial]: + if self.size <= 0: + return [] + trial = _create_trial(config) + return [trial] * self.size + + +# 未来扩展 — 不同的 Operator 实现: +# +# class DataScatterOperator(Operator): +# """数据驱动: 按数据列表生成 TrialList，每个 Task 注入不同 env""" +# def __init__(self, data: list[dict[str, str]]): +# self.data = data +# def apply(self, config): +# if not self.data: +# return [] +# return [ +# _create_trial(config.model_copy(update={"env": {**config.env, **shard}})) +# for shard in self.data +# ] +# +# class RayScatterOperator(Operator): +# """Ray 分发: 生成 TrialList，JobExecutor 通过 ray.remote 执行""" +# +# class K8sOperator(Operator): +# """K8s 分发: 生成 TrialList，JobExecutor 创建 K8s Job 执行""" + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# JobExecutor: 主控方 — 驱动 Operator，管理 sandbox 生命周期 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class JobExecutor: + """执行引擎 + 主控方 + + 职责: + 1. 调用 operator.apply(config) 获取 TrialList + 2. 并行启动每个 Task (_do_submit) + 3. 并行等待完成 + 收集结果 (_do_wait) + + JobExecutor 负责并行执行，Operator 只负责生成 TrialList。 + """ + + # ── run = submit + wait ── + + async def run(self, operator: Operator, config: JobConfig) -> list[TrialResult]: + """完整生命周期: submit + wait""" + job_client = await self.submit(operator, config) + return await self.wait(job_client) + + # ── submit: Operator 生成 TrialList → 并行启动 ── + + async def submit(self, operator: Operator, config: JobConfig) -> JobClient: + """Operator apply 生成 TrialList，并行启动所有 sandbox""" + # 1. Operator 决定分发 8 份 Trial + trial_list = operator.apply(config) + + # 2. size=0 → 什么都不做 + if not trial_list: + return JobClient(tasks=[]) + + # 3. 并行启动每个 Task + task_clients = await asyncio.gather(*[ + self._(do_submit(trial) for trial in trial_list + ]) + return JobClient(tasks=list(task_clients)) + + # ── wait: 并行等待 + 收集结果 ── + + async def wait(self, job_client: JobClient) -> list[TrialResult]: + """并行等待所有 task 完成，收集结果""" + if not job_client.tasks: + return [] + return list(await asyncio.gather(*[ + self._do_wait(tc) for tc in job_client.tasks + ])) + + # ── 内部: 单个 task 的 submit / wait ── + + async def _do_submit(self, trial: AbstractTrial) -> TrialClient: + """启动单个 sandbox + 执行脚本 (被 Operator 回调)""" + config = task._config # Task 已持有 config + + sandbox = Sandbox(config.environment) + await sandbox.start() + + session = f"rock-job-{config.job_name or 'default'}" + await sandbox.create_session( + CreateBashSessionRequest( + session=session, env_enable=True, env=config.env or None, + ) + ) + + # Task: setup → build + await trial.setup(sandbox) + script_content = trial.build() + + # 上传脚本 + nohup 启动 + script_path = f"/tmp/rock_job_{config.job_name}.sh" + await self._upload_script(sandbox, script_content, script_path) + pid, error = await sandbox.start_nohup_process( + cmd=f"bash {script_path}", session=session, + tmp_file=f"/tmp/rock_job_{config.job_name}.out", + ) + if error: + raise RuntimeError(f"Failed to start task: {error.output}") + + return TrialClient(sandbox=sandbox, session=session, pid=pid, task=task) + + async def _do_wait(self, client: TrialClient) -> TrialResult: + """等待单个 task 完成，调用 trial.collect() 收集结果""" + config = client.trial._config + try: + success, message = await client.sandbox.wait_for_process_completion( + pid=client.pid, session=client.session, wait_timeout=config.timeout, + ) + obs = await client.sandbox.handle_nohup_output( + tmp_file=f"/tmp/rock_job_{config.job_name}.out", + session=client.session, success=success, message=message, + ) + result = await client.trial.collect(client.sandbox, obs.output or "", obs.exit_code or 1) + if not success: + result.status = TrialStatus.FAILED + return result + finally: + if config.auto_stop: + await client.sandbox.close() + + async def _upload_script(self, sandbox, content: str, path: str) -> None: + """上传脚本内容到 sandbox""" + ... + + +@dataclass +class TrialClient: + """单个运行中 task 的句柄 (config 通过 task._config 访问)""" + sandbox: Sandbox + session: str + pid: int + trial: AbstractTrial + +@dataclass +class JobClient: + """Job.submit() 返回的句柄，持有多个 TrialClient，类比 Flink JobClient""" + tasks: list[TrialClient] + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Job: 极薄 Facade — 用户唯一入口 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class Job: + """Job Facade: 极薄入口层 + + 只有两个参数: config + operator。 + 不包含执行逻辑、调度逻辑 — 全部委托给 JobExecutor。 + + 调用链: + Job.run() = Job.submit() + Job.wait() + submit: → JobExecutor.submit(operator, config) + → operator.apply(config) → TrialList # 算子生成 N 份 Task + → 并行 _(do_submit(trial) for each # Executor 并行启动 + → 返回 JobClient + wait: → JobExecutor.wait(job_client) + → 并行 _do_wait(tc) for each # 等待 + trial.collect() + → 返回 list[TrialResult] + + 用法: + # 单次执行 + result = await Job(BashJobConfig(script="echo hello")).run() + + # 异步 submit/wait + job = Job(config) + await job.submit() + result = await job.wait() + + # 并行: 每个 config 一个 Job，asyncio.gather 并发 + results = await asyncio.gather(*[Job(c).run() for c in configs]) + """ + + def __init__( + self, + config: JobConfig, + operator: Operator | None = None, + ): + self._config = config + self._executor = JobExecutor() + self._operator = operator or ScatterOperator() + self._job_client: JobClient | None = None + + async def run(self) -> JobResult: + """完整生命周期: submit + wait""" + await self.submit() + return await self.wait() + + async def submit(self) -> None: + """非阻塞提交""" + self._job_client = await self._executor.submit(self._operator, self._config) + + async def wait(self) -> JobResult: + """等待完成""" + if not self._job_client: + raise RuntimeError("No submitted job. Call submit() first.") + task_results = await self._executor.wait(self._job_client) + return self._build_result(task_results) + + async def cancel(self) -> None: + """取消运行中的 job""" + if self._job_client: + for tc in self._job_client.tasks: + await tc.sandbox.arun(cmd=f"kill {tc.pid}", session=tc.session) + + def _build_result(self, task_results: list[TrialResult]) -> JobResult: + all_success = all(r.success for r in task_results) + return JobResult( + job_id=self._config.job_name or "", + status=JobStatus.COMPLETED if all_success else JobStatus.FAILED, + labels=self._config.labels, + task_results=task_results, + ) + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Result 模型 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TrialResult(BaseModel): + """单个 Task 的执行结果""" + task_id: str = "" + status: TrialStatus = TrialStatus.COMPLETED + output: str = "" + exit_code: int = 0 + data: dict[str, Any] = {} # 通用数据输出 + trial_results: list[TrialResult] = [] # Harbor 专用 + + @property + def score(self) -> float: ... + + @property + def success(self) -> bool: + return self.status == TrialStatus.COMPLETED + +class JobResult(BaseModel): + """Job 聚合结果""" + job_id: str = "" + status: JobStatus = JobStatus.COMPLETED + labels: dict[str, str] = {} + task_results: list[TrialResult] = [] + + @property + def score(self) -> float: ... + + @property + def n_completed(self) -> int: ... + + @property + def n_failed(self) -> int: ... + + # 便捷属性: 直接访问第一个 TrialResult 的 trial_results (Harbor 兼容) + @property + def trial_results(self) -> list[TrialResult]: + if self.task_results: + return self.task_results[0].trial_results + return [] +``` + +#### 使用示例 + +```python +# ── 1. 单次 Bash Job ── +result = await Job(BashJobConfig( + script="echo 'hello world'", + environment=RockEnvironmentConfig(image="python:3.11"), +)).run() +# 内部: Job.submit() +# → JobExecutor.submit(ScatterOperator(size=1), config) +# → ScatterOperator.apply(config) → [BashTrial] (1 份) +# → _do_submit(BashTrial) → sandbox → trial.setup/build +# Job.wait() → _do_wait → trial.collect + +# ── 2. 单次 Harbor Benchmark ── +result = await Job(HarborJobConfig( + environment=RockEnvironmentConfig(image="harbor-runner:latest", memory="16g"), + setup_commands=["pip install harbor --quiet"], + agents=[AgentConfig(name="terminus-2", model_name="hosted_vllm/my-model")], + datasets=[RegistryDatasetConfig(registry=OssRegistryInfo(split="v2.0"), name="tb")], + auto_stop=True, experiment_id="exp-001", +)).run() + +# ── 3. Harbor from YAML ── +result = await Job(HarborJobConfig.from_yaml("job_config.yaml")).run() + +# ── 4. 异步 submit/wait ── +job = Job(HarborJobConfig(...)) +await job.submit() +# ... 做其他事 ... +result = await job.wait() + +# ── 5. 并行: 多个 Job + asyncio.gather ── +import asyncio + +async def run_one(tc): + return await Job(HarborJobConfig( + environment=base_env, agents=[agent_cfg], tasks=[tc], + auto_stop=True, experiment_id="exp-001", + )).run() + +results = await asyncio.gather(*[run_one(tc) for tc in task_batch]) +rewards = [r.score for r in results] + +# ── 6. 评测平台: CSV → 多个 Job → 信号量控制并发 ── +import csv + +with open("eval_tasks.csv") as f: + rows = list(csv.DictReader(f)) + +sem = asyncio.Semaphore(8) + +async def run_eval(row): + async with sem: + return await Job(BashJobConfig( + script="python eval.py", + env={"TASK_ID": row["task_id"], "INPUT": row["input_path"]}, + environment=RockEnvironmentConfig(image="eval-runner:latest"), + auto_stop=True, + )).run() + +results = await asyncio.gather(*[run_eval(row) for row in rows]) + +# ── 7. 自定义 Trial 类型 ── +class EvalJobConfig(JobConfig): + eval_script: str + model_endpoint: str + +class EvalTask(AbstractTrial): + async def setup(self, sandbox): + await sandbox.fs.upload_dir(self._config.eval_script, "/workspace/") + + def build(self) -> str: + return f"python /workspace/{self._config.eval_script} --endpoint {self._config.model_endpoint}" + + async def collect(self, sandbox, output, exit_code): + result_content = await sandbox.read_file(ReadFileRequest(path="/workspace/result.json")) + return TrialResult(output=output, exit_code=exit_code, data=json.loads(result_content.content)) + +register_trial(EvalJobConfig, EvalTask) +result = await Job(EvalJobConfig(eval_script="eval.py", model_endpoint="http://...")).run() +``` + +#### CLI 集成 + +```python +# rock/cli/command/job.py + +class JobCommand(Command): + name = "job" + + async def arun(self, args): + if args.job_command == "run": + job_type = args.type or "bash" + + if job_type == "bash": + config = BashJobConfig( + script=args.script_content, + script_path=args.script, + file_uploads=[(args.local_path, args.target_path)] if args.local_path else [], + environment=RockEnvironmentConfig(image=args.image, ...), + timeout=args.timeout, + ) + elif job_type == "harbor": + config = HarborJobConfig.from_yaml(args.config) + if args.image: + config.environment.image = args.image + + result = await Job(config).run() + + @staticmethod + async def add_parser_to(subparsers): + job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") + job_sub = job_parser.add_subparsers(dest="job_command") + + run = job_sub.add_parser("run", help="Run a job in a sandbox") + run.add_argument("--type", choices=["bash", "harbor"], default="bash") + run.add_argument("--script", help="Script file path") + run.add_argument("--script-content", help="Inline script content") + run.add_argument("--config", help="Harbor YAML config path") + run.add_argument("--image", help="Sandbox image") + run.add_argument("--memory", help="Memory (e.g. 8g)") + run.add_argument("--cpus", type=float, help="CPU count") + run.add_argument("--timeout", type=int, default=3600) + run.add_argument("--local-path", help="Local dir to upload") + run.add_argument("--target-path", default="/root/job") +``` + +#### 向后兼容 + +``` +现有代码 (rock/sdk/agent) 新架构 (rock/sdk/job) +───────────────────────────── ───────────────────────────── +from rock.sdk.agent import from rock.sdk.job import + Job, JobConfig Job, HarborJobConfig + +job = Job(JobConfig( job = Job(HarborJobConfig( + experiment_id="exp", experiment_id="exp", + agents=[...], agents=[...], + datasets=[...], datasets=[...], +)) )) +result = await job.run() result = await job.run() +result.score result.score +result.trial_results result.trial_results ← 兼容属性 +``` + +**兼容策略: 两套并存，逐步迁移** + +1. `rock/sdk/agent/` **完全不动** — 现有 `Job(JobConfig(...))` 继续工作 +2. `rock/sdk/job/` **全新模块** — 新架构，新能力 +3. `HarborJobConfig` 的 Harbor 字段与现有 `JobConfig` 完全一致 — 迁移只需改 import + 类名 +4. `JobResult.trial_results` 提供便捷属性 — 兼容 `result.trial_results` 访问方式 +5. 未来标记 `rock.sdk.agent.Job` 为 deprecated，引导到 `rock.sdk.job.Job` + +现有调用方 (`examples/harbor/harbor_demo.py`, `tests/unit/sdk/agent/`, verl 集成) **零改动**。 + +#### 优点 + +- **职责清晰** — Job(Facade) / JobExecutor(主控+执行) / Operator(算子) / Task(逻辑) 各司其职 +- **Operator 可扩展** — ABC 基类，ScatterOperator 默认实现，未来可加 DataScatterOperator / RayScatterOperator +- **三层都是 run=submit+wait** — Job / JobExecutor / 内部全统一为 submit+wait 模式 +- **Trial 接口清晰** — `setup/build/collect` 三阶段，语义明确 +- **`Job(config)` 签名一致** — 用户只需记 Config 子类 +- **注册表扩展** — `register_trial()` 支持第三方扩展 +- **完全兼容** — `rock/sdk/agent/` 不动 + +#### 缺点 + +- **组件多** — 4 个核心组件（Job, JobExecutor, Operator, Task），理解成本略高 +- **两套代码暂时并存** — `rock/sdk/agent/Job` 和 `rock/sdk/job/Job` + +#### 重构影响 + +| 现有组件 | 变更 | +|---------|------| +| `rock/sdk/agent/` | **不动**，保持向后兼容 | +| `rock/sdk/job/` (新) | Job(Facade), JobExecutor, Operator, Task | +| `rock/sdk/job/trial/harbor.py` | 从 `rock/sdk/agent/job.py` 提取 setup/build/collect | +| `rock/sdk/job/trial/harbor.py` | 复用 `rock/sdk/agent/models/` 的 Harbor schema | +| `rock/cli/command/job.py` | 使用新 `rock.sdk.job.Job`，支持 `--type` | +| 现有测试/示例 | **不改**，新模块新增独立测试 | + +--- + +## 3. 备选方案对比 + +在确定当前方案前，评估了另外两种方案: + +| 维度 | A: Job 层继承 | B: 策略模式 | **当前方案: Config 驱动** | +|------|-------------|-----------|------------------------| +| **核心思路** | `AbstractJob` → `BashJob` / `HarborJob` 子类 | 单一 `Job` + `TaskStrategy` / `Operator` / `DataSource` 组合 | `Job(config)` Config 子类决定类型，Task 内部 factory | +| **用户 API** | `BashJob(BashJobConfig(...))` | `Job(config, task_strategy=BashTrialStrategy())` | `Job(BashJobConfig(...))` | +| **Trial 类型扩展** | 新增 Job 子类 + Config 子类 | 新增 Strategy 类 | 新增 Config 子类 + Trial 子类 + register | +| **调度策略** | Scheduler 挂在 Job 上，继承限制组合 | Strategy 可替换 | Operator 可替换 | +| **类型安全** | 强 | 弱 (Config 混合所有字段) | 强 | +| **使用门槛** | 中 (需记 BashJob/HarborJob 类名) | 高 (需组合三个 Strategy) | **低** (只需记 Config 子类) | +| **向后兼容** | 一般 (Job 类要改) | 差 (接口全变) | **好** (`rock/sdk/agent` 不动) | +| **`Job(config)` 签名** | 不同子类不同构造 | 需要额外 strategy 参数 | **统一 `Job(config)`** | + +**选择当前方案的理由**: 保持 `Job(config)` 统一签名，Config 子类携带类型信息（无需额外参数），Task 作为内部实现细节不暴露给用户，同时 `rock/sdk/agent` 完全不动实现零成本兼容。 + +--- + +## 4. 文件结构 + +``` +rock/sdk/job/ # 新模块 (核心只有 6 个文件) +├── __init__.py # 公开导出 +├── job.py # Job Facade (极薄) +├── executor.py # JobExecutor 执行引擎 +├── operator.py # Operator ABC + ScatterOperator (scatter 算子) +├── config.py # JobConfig, BashJobConfig, HarborJobConfig +├── result.py # JobResult, TrialResult, TrialStatus +└── trial/ + ├── __init__.py + ├── abstract.py # AbstractTrial (setup/build/collect) + ├── registry.py # _TRIAL_REGISTRY + register_trial() + ├── bash.py # BashTrial + └── harbor.py # HarborTrial (复用 agent/models) + +rock/sdk/agent/ # 保留不动，完全向后兼容 +├── __init__.py +├── job.py +├── models/ # Harbor schema (HarborTrial 复用) +└── constants.py +``` + +--- + +## 5. 迁移路径 + +``` +Phase 1: Config + Task + JobExecutor + - 创建 rock/sdk/job/ 模块 + - JobConfig → BashJobConfig / HarborJobConfig 继承体系 + - AbstractTrial (setup/build/collect), BashTrial, HarborTrial + _TRIAL_REGISTRY + - JobExecutor: sandbox 生命周期 + 调用 Trial 三阶段 + - HarborTrial 从 rock/sdk/agent/job.py 提取逻辑，复用 agent/models + +Phase 2: Job Facade + Operator + Result + - Job(config, operator?) 极薄 Facade + - ScatterOperator(size) 默认 scatter 算子 + - Job.run_batch() 批量并行 + - TrialResult / JobResult 结果模型 + +Phase 3: 更新 CLI + - rock job run --type bash/harbor + +Phase 4: 标记旧接口 deprecated + - rock/sdk/agent/Job + JobConfig 标记 DeprecationWarning + - 文档引导迁移到 rock/sdk/job +``` + +--- + +## 6. Flink / Ray Data 对比分析 + +对比 Flink 和 Ray Data 的核心设计模式，评估当前方案的完善点。 + +### 6.1 核心架构对比 + +| 维度 | Flink | Ray Data | Rock Job (当前方案) | +|------|-------|----------|-------------------| +| **核心抽象** | Operator → Task → SubTask (DAG) | Dataset → Transform chain (lazy pipeline) | Job(Facade) → JobExecutor(执行) → Task(逻辑) | +| **执行模型** | 流式 DAG，operator chaining 自动融合 | 惰性 pipeline，`.show()` 触发流式执行 | JobExecutor 管理 sandbox 生命周期，nohup 模式 | +| **调度与执行** | Scheduler 分配 slot，TaskManager 执行 | Scheduler 放置 task，Worker 执行 | JobExecutor 驱动 Operator，Operator 回调执行 | +| **并行度** | 每个 operator 独立 parallelism | 每个 operation 独立 `concurrency` + `num_cpus/gpus` | ScatterOperator `size` 控制 Trial 数量 | +| **数据输入** | Source (SourceReader + SplitEnumerator) | Datasource ABC (`read_csv`, `read_parquet`, custom) | Task.setup() 自行控制 | +| **数据输出** | Sink (SinkWriter) | Datasink ABC (`write_parquet`, `write_json`, custom) | Task.collect() 自行控制 | +| **容错** | Checkpoint + Savepoint + 自动重启策略 | Ray Core task retry + lineage reconstruction | Job.run_batch(max_retries=N) | +| **状态管理** | 丰富的 Job 状态机 + 指标累加器 | progress bar + per-operation metrics | JobStatus enum，无进度跟踪 | +| **Pipeline** | Operator chain → JobGraph → ExecutionGraph 三级图 | `ds.map().map_batches().filter()` 链式调用 | 单步，无 pipeline | + +### 6.2 可借鉴的设计 + 当前方案改进建议 + +#### 对比总结: 数据 IO 策略 + +| | Flink | Ray Data | Rock Job | +|---|---|---|---| +| **数据输入** | Source 是 DAG 中的 operator | `read_csv()` / `Datasource` ABC | **Task.setup() 自行控制** | +| **数据输出** | Sink 是 DAG 中的 operator | `write_parquet()` / `Datasink` ABC | **Task.collect() 自行控制** | +| **为什么不同** | Flink/Ray 是数据处理引擎，数据流是核心 | 同左 | Rock 是 sandbox 执行引擎，数据 IO 是 Trial 逻辑的一部分 | + +Rock 不需要 DataSource/DataSink 抽象的原因: +- Flink/Ray 的 operator 是轻量级函数，需要框架管理数据流 +- Rock 的 Task 运行在独立 sandbox 中，有完整文件系统，Task 自己读写数据更自然 +- 批量并行通过 `Job.run_batch(configs)` 实现，每个 config 自带不同的输入参数 + +#### 已吸收的改进 + +| # | 改进项 | 参考 | 实现方式 | +|---|--------|------|---------| +| 1 | **Retry** — 失败重试 | Flink RestartStrategy, Ray task retry | `Job.run_batch(max_retries=N)` | +| 2 | **JobExecutor** — 独立执行引擎 | Flink JobExecutor, Ray Worker | `JobExecutor` 类 | +| 3 | **调度与执行分离** | Flink Scheduler ≠ TaskManager | `JobExecutor`(主控) 驱动 `Operator`(算子) | +| 4 | **Per-job Resource** — 每个 job 不同资源 | Ray Data per-op resources | 每个 config 有独立 `environment` | + +#### 未来方向 (不在第一版) + +| # | 方向 | 参考 | +|---|------|------| +| 1 | **Pipeline** — 多 Job 串联 | Flink DAG, Ray Data pipeline chain | +| 2 | **Progress** — 批量执行进度回调 | Ray progress bar | +| 3 | **DataSource 抽象** — 如果 Task.setup() 模式不够用 | Ray `read_datasource()` | + +### 6.6 不需要借鉴的设计 + +| Flink / Ray Data 特性 | 不采用原因 | +|----------------------|----------| +| **Operator DAG / Pipeline chain** | Rock 的执行单元是 sandbox（重量级容器），不是轻量级 operator。DAG 在 sandbox 粒度上没有意义 | +| **Operator Fusion** | 同上，Rock 不做 operator-level 的计算融合 | +| **Streaming 执行** | Rock Job 是 batch 模式（提交 → 等待 → 收集），sandbox 不支持流式数据传递 | +| **Checkpoint / Savepoint** | Rock sandbox 是无状态的一次性容器，不需要持久化计算状态 | +| **Type System** | Flink 的 TypeInformation 用于序列化优化，Rock 的数据通过文件/环境变量传递，不需要类型系统 | diff --git a/docs/dev/job/operator.md b/docs/dev/job/operator.md new file mode 100644 index 0000000000..09e28edfad --- /dev/null +++ b/docs/dev/job/operator.md @@ -0,0 +1,363 @@ +# 分布式算子参考 + +对比 torch.distributed、Ray Data、Flink 三个框架的算子抽象，分析 Rock Job Operator 的设计参考。 + +## 1. 算子总览 + +### 1.1 torch.distributed — 集合通信算子 + +torch.distributed 定义了三类通信原语，操作对象是 **Tensor**： + +| 类别 | 算子 | 输入 | 输出 | 语义 | +|------|------|------|------|------| +| **点对点** | `send(tensor, dst)` | 1 Tensor on src | dst rank 收到 | 发送给指定 rank | +| | `recv(tensor, src)` | 空 Tensor on dst | dst 填充数据 | 从指定 rank 接收 | +| | `isend` / `irecv` | 同上 | 返回 Work (异步) | 异步版本 | +| **一对多** | `broadcast(tensor, src)` | 1 Tensor on src | 所有 rank 拿到相同 Tensor | 广播 | +| | `scatter(output, scatter_list, src)` | src 持有 list[Tensor] | 每个 rank 拿到 1 片 | 拆分分发 | +| **多对一** | `reduce(tensor, dst, op)` | 每个 rank 1 Tensor | dst 拿到聚合结果 | 聚合到 dst | +| | `gather(gather_list, tensor, dst)` | 每个 rank 1 Tensor | dst 拿到 list[Tensor] | 收集到 dst | +| **多对多** | `all_reduce(tensor, op)` | 每个 rank 1 Tensor | 每个 rank 拿到聚合结果 | 聚合 + 广播 | +| | `all_gather(tensor_list, tensor)` | 每个 rank 1 Tensor | 每个 rank 拿到完整 list | 收集 + 广播 | +| | `reduce_scatter(output, input_list, op)` | 每个 rank list[Tensor] | 每个 rank 拿到 1 片聚合 | 聚合 + 拆分 | +| | `all_to_all(output, input)` | 每个 rank N 片 | 每个 rank 从所有 rank 各收 1 片 | 全交换 | +| **同步** | `barrier()` | 无 | 无 | 等所有 rank 到达 | +| **聚合操作** | `ReduceOp` | — | — | SUM, PRODUCT, MIN, MAX, AVG, BAND, BOR, BXOR | + +### 1.2 数据流图 + +``` +4 个 Rank, 每个持有一个 Tensor + +broadcast (src=0): scatter (src=0): + R0:[1,2] ──→ R0:[1,2] R0:[A,B,C,D] ──→ R0:[A] + R1:[0,0] ──→ R1:[1,2] ──→ R1:[B] + R2:[0,0] ──→ R2:[1,2] ──→ R2:[C] + R3:[0,0] ──→ R3:[1,2] ──→ R3:[D] + +reduce (dst=0, SUM): gather (dst=0): + R0:[1,0] ──┐ R0:[A] ──┐ + R1:[0,1] ──┼→ R0:[1,2] R1:[B] ──┼→ R0:[A,B,C,D] + R2:[0,1] ──┤ R2:[C] ──┤ + R3:[0,0] ──┘ R3:[D] ──┘ + +all_reduce (SUM): all_gather: + R0:[1,0] ──┐ ┌→ R0:[1,2] R0:[A] ──┐ ┌→ R0:[A,B,C,D] + R1:[0,1] ──┼───┼→ R1:[1,2] R1:[B] ──┼───┼→ R1:[A,B,C,D] + R2:[0,1] ──┤ ├→ R2:[1,2] R2:[C] ──┤ ├→ R2:[A,B,C,D] + R3:[0,0] ──┘ └→ R3:[1,2] R3:[D] ──┘ └→ R3:[A,B,C,D] + +reduce_scatter (SUM): all_to_all: + R0:[1,2,3,4] ──┐ R0:[a0,a1,a2,a3] R0:[a0,b0,c0,d0] + R1:[5,6,7,8] ──┼→ SUM →拆分 R1:[b0,b1,b2,b3] → R1:[a1,b1,c1,d1] + R2:[1,1,1,1] ──┤ [8,10,12,14] R2:[c0,c1,c2,c3] R2:[a2,b2,c2,d2] + R3:[1,1,1,1] ──┘ R0:[8,10] R3:[d0,d1,d2,d3] R3:[a3,b3,c3,d3] + R1:[12,14] +``` + +## 2. scatter 跨框架对比 + +scatter 是最核心的"数据分发"算子，三个框架实现差异显著： + +### 2.1 输入输出对比 + +| | 输入 | 拆分方式 | 输出 | 粒度 | +|---|------|---------|------|------| +| **torch** | src rank 持有 `list[Tensor]` (len=world_size) | 按 index 1:1 分发 | 每个 rank 拿到 1 个 Tensor | Tensor 切片 | +| **Ray Data** | 1 个 `Dataset` (内部 N blocks) | `split(K)` 按 block 边界拆 | K 个 `MaterializedDataset` | Block (Arrow Table) | +| **Flink** | 1 个 `DataStream` | `rebalance()` round-robin / `keyBy()` hash | 每个 subtask 收到部分 Record | Record (StreamRecord) | + +### 2.2 代码示例 + +**torch.distributed — 6 个核心算子完整示例 (4 ranks)** + +```python +import os +import torch +import torch.distributed as dist + +def init_process(rank, size, fn, backend="gloo"): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29500" + dist.init_process_group(backend, rank=rank, world_size=size) + fn(rank, size) +``` + +**scatter — 一对多拆分** + +```python +def do_scatter(rank, size): + # 签名: dist.scatter(tensor, scatter_list=None, src=0, group=None, async_op=False) + # tensor: 接收缓冲区 (每个 rank 上) + # scatter_list: 待分发的 tensor 列表 (仅 src rank 提供, len=world_size) + # src: 源 rank + + tensor = torch.empty(1) # 每个 rank 准备空接收缓冲区 + if rank == 0: + scatter_list = [torch.tensor([float(i + 1)]) for i in range(size)] + # scatter_list = [tensor([1.]), tensor([2.]), tensor([3.]), tensor([4.])] + dist.scatter(tensor, scatter_list=scatter_list, src=0) + else: + dist.scatter(tensor, scatter_list=[], src=0) + + print(f"[rank {rank}] received: {tensor[0]}") + # [rank 0] received: 1.0 ← scatter_list[0] + # [rank 1] received: 2.0 ← scatter_list[1] + # [rank 2] received: 3.0 ← scatter_list[2] + # [rank 3] received: 4.0 ← scatter_list[3] +``` + +**gather — 多对一收集 (scatter 的逆操作)** + +```python +def do_gather(rank, size): + # 签名: dist.gather(tensor, gather_list=None, dst=0, group=None, async_op=False) + # tensor: 每个 rank 提供的数据 + # gather_list: 收集结果 (仅 dst rank 提供, len=world_size) + # dst: 目标 rank + + tensor = torch.tensor([float(rank)]) # 每个 rank 持有自己的数据 + if rank == 0: + gather_list = [torch.empty(1) for _ in range(size)] + dist.gather(tensor, gather_list=gather_list, dst=0) + print(f"[rank 0] gathered: {gather_list}") + # [rank 0] gathered: [tensor([0.]), tensor([1.]), tensor([2.]), tensor([3.])] + else: + dist.gather(tensor, gather_list=[], dst=0) +``` + +**broadcast — 一对多复制** + +```python +def do_broadcast(rank, size): + # 签名: dist.broadcast(tensor, src=0, group=None, async_op=False) + if rank == 0: + tensor = torch.tensor([42.0]) + else: + tensor = torch.empty(1) + + dist.broadcast(tensor, src=0) + print(f"[rank {rank}] received: {tensor[0]}") + # 所有 rank 输出: 42.0 +``` + +**reduce — 多对一聚合** + +```python +def do_reduce(rank, size): + # 签名: dist.reduce(tensor, dst=0, op=ReduceOp.SUM, group=None, async_op=False) + tensor = torch.ones(1) # 每个 rank 持有 tensor([1.0]) + + dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM) + print(f"[rank {rank}] result: {tensor[0]}") + # [rank 0] result: 4.0 ← 1+1+1+1, 仅 rank 0 有聚合值 + # [rank 1] result: 1.0 ← 未改变 +``` + +**all_reduce — 多对多聚合 (reduce + broadcast)** + +```python +def do_all_reduce(rank, size): + # 签名: dist.all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False) + tensor = torch.ones(1) + + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + print(f"[rank {rank}] result: {tensor[0]}") + # 所有 rank 输出: 4.0 ← 每个 rank 都拿到聚合结果 +``` + +**all_gather — 多对多收集 (gather + broadcast)** + +```python +def do_all_gather(rank, size): + # 签名: dist.all_gather(tensor_list, tensor, group=None, async_op=False) + tensor = torch.tensor([float(rank)]) + tensor_list = [torch.empty(1) for _ in range(size)] + + dist.all_gather(tensor_list, tensor) + print(f"[rank {rank}] gathered: {tensor_list}") + # 所有 rank 输出: [tensor([0.]), tensor([1.]), tensor([2.]), tensor([3.])] +``` + +**Ray Data split (scatter 等价)** + +```python +ds = ray.data.read_csv("big_data.csv") # 1 Dataset → N blocks + +shards = ds.split(n=4, equal=True) +# shards[0]: MaterializedDataset (25% blocks) +# shards[1]: MaterializedDataset (25% blocks) +# ... + +workers = [Worker.remote() for _ in range(4)] +ray.get([w.train.remote(s) for w, s in zip(workers, shards)]) +``` + +**Flink rebalance (scatter 等价)** + +```java +dataStream + .rebalance() // round-robin 逐条分发到所有 subtask + .map(new MyMapFunction()); + +dataStream + .keyBy(row -> row.userId) // 按 key hash 分发 (相同 key → 同一 subtask) + .process(new MyProcessFunction()); +``` + +### 2.3 核心差异 + +| 维度 | torch | Ray Data | Flink | +|------|-------|----------|-------| +| **数据模型** | Tensor (连续内存) | Block (Arrow Table) | Record (流记录) | +| **拆分时机** | 调用时一次性同步 | materialize 时 | 运行时逐条流式 | +| **需要预知总量?** | 是 (world_size) | 否 (按 block 拆) | 否 (流式) | +| **拆分粒度** | Tensor 维度切片 | Block 级别 | Record 级别 | +| **通信方式** | NCCL/Gloo 进程间通信 | Object Store 引用传递 | Network Shuffle | +| **适用场景** | GPU 并行计算 | 数据并行处理 | 流/批处理 | + +## 3. map 跨框架对比 + +| | 用户函数签名 | 输入 | 输出 | 粒度 | +|---|------------|------|------|------| +| **torch** | 无内置 map (手动循环) | Tensor | Tensor | — | +| **Ray Data** | `lambda row: Dict` 或 `Callable(batch) -> batch` | `Dict[str, Any]` (一行) 或 `Dict[str, np.ndarray]` (一批) | 同类型 | 逐行 / 逐批 | +| **Flink** | `MapFunction.map(T) -> O` | `T` (一条记录) | `O` (一条记录) | 逐条 | +| **Rock** | `AbstractTask.setup/build/collect` | `JobConfig` (整个任务配置) | `TaskResult` | 逐 sandbox | + +### Flink MapFunction 层次 + +``` +MapFunction ← 用户接口: T → O + ↓ 包装为 +StreamMap(OneInputStreamOperator) ← Operator: StreamRecord → StreamRecord + ↓ chain 为 +Task (OperatorChain) ← 运行时: 多个 Operator 融合成一个 Task + ↓ 调度到 +TaskManager ← 执行引擎: 分配 slot 执行 +``` + +### Ray Data map 层次 + +``` +lambda row: dict ← 用户函数: dict → dict + ↓ 包装为 +MapTransformer ← 函数包装: 行级 → 块级 + ↓ 传入 +MapOperator(PhysicalOperator) ← Operator: RefBundle → RefBundle + ↓ 提交 +_map_task.remote(transformer, blocks) ← ray.remote: 在 worker 执行 + ↓ 调度 +Executor ← 执行引擎: 驱动调度循环 +``` + +## 4. 全部算子的 scatter + map + gather 分解 + +大部分集合通信算子可以分解为 scatter + map + gather 的组合： + +| 算子 | = | scatter | + map | + gather/reduce | +|------|---|---------|-------|-----------------| +| **map** | | — | f(x) 对每个 item | — | +| **broadcast** | | 复制到所有 rank | — | — | +| **scatter** | | 拆分到各 rank | — | — | +| **gather** | | — | — | 收集到 dst | +| **reduce** | | — | — | 聚合到 dst | +| **all_reduce** | | scatter (隐式) | — | reduce + broadcast | +| **all_gather** | | — | — | gather + broadcast | +| **reduce_scatter** | | scatter | — | reduce | +| **map + gather** | | scatter 数据 | f(x) 对每片 | gather 结果 | + +## 5. Rock Job Operator 对应 + +Rock Job 的执行单元是 sandbox（重量级容器，sandbox 之间隔离不通信），因此只需要部分算子： + +### 5.1 已实现 + +| 算子 | Rock 实现 | 说明 | +|------|----------|------| +| **map** | `MapOperator.apply()` | 同一 task 起 N 个 sandbox 并行 | +| **gather** | `JobExecutor.wait()` | 收集所有 TaskClient 结果 | +| **reduce** | `JobResult._build_result()` | 聚合 TaskResult (score, n_completed 等) | +| **barrier** | `asyncio.gather` | 等所有 sandbox 完成 | +| **broadcast** | `MapOperator(concurrency=N)` | 同一 config 广播到 N 个 sandbox | + +### 5.2 未来方向 + +| 算子 | 场景 | 实现方式 | +|------|------|---------| +| **scatter** | 大数据集拆分到多个 sandbox | `ScatterOperator`: 接收数据集，拆分，每个 sandbox 拿一片 | +| **reduce** (自定义) | 自定义聚合逻辑 (不只是 score 平均) | `ReduceOperator` 或 `JobResult` 自定义 reduce_fn | +| **scatter + map + gather** | 完整分布式数据处理流 | Pipeline 组合多个 Operator | + +### 5.3 不适用 + +| 算子 | 不适用原因 | +|------|----------| +| **all_reduce** | sandbox 之间不共享内存，无法进程间通信 | +| **all_gather** | 同上 | +| **all_to_all** | 同上 | +| **send / recv** | sandbox 之间无点对点通道 | + +## 6. ScatterOperator 设计草案 (未来) + +如果未来需要框架级 scatter 支持： + +```python +class ScatterOperator(Operator): + """Scatter 算子: 拆分数据，每个 sandbox 拿一片 + + 类比: + torch: dist.scatter(output, scatter_list, src=0) + Ray Data: ds.split(n=4) + Flink: dataStream.rebalance() + """ + + def __init__(self, data: list[dict[str, str]], concurrency: int | None = None): + """ + Args: + data: 待拆分的数据列表，每个 dict 注入为 sandbox env vars + concurrency: 并发数，默认 = len(data) + """ + self.data = data + self.concurrency = concurrency or len(data) + + async def apply(self, config, submit_fn): + task = _create_task(config) + sem = asyncio.Semaphore(self.concurrency) + + async def submit_one(shard: dict[str, str]) -> TaskClient: + async with sem: + # 为每个 shard 克隆 task，注入 shard 数据到 env + shard_task = _create_task(config.model_copy( + update={"env": {**config.env, **shard}} + )) + return await submit_fn(shard_task) + + return await asyncio.gather(*[submit_one(s) for s in self.data]) + + +# 使用 +import csv +with open("eval_tasks.csv") as f: + rows = list(csv.DictReader(f)) + +job = Job( + BashJobConfig(script="python eval.py", environment=...), + operator=ScatterOperator(data=rows, concurrency=8), +) +result = await job.run() +``` + +scatter + map 的完整流程： + +``` +ScatterOperator.apply(config, submit_fn) + │ + ├─ 1. _create_task(config) # 创建基础 Task + ├─ 2. 对 data 中每个 shard: + │ ├─ 克隆 config，注入 shard env # scatter: 每个 sandbox 拿不同数据 + │ ├─ _create_task(shard_config) # 创建带 shard 数据的 Task + │ └─ submit_fn(shard_task) # map: 启动 sandbox 执行 + └─ 3. asyncio.gather(...) # barrier: 等所有完成 + # gather: 收集结果 → list[TaskClient] +``` diff --git a/docs/dev/job/plan.md b/docs/dev/job/plan.md new file mode 100644 index 0000000000..fae1fbec4d --- /dev/null +++ b/docs/dev/job/plan.md @@ -0,0 +1,1938 @@ +# Rock Job 重构实施计划 + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 根据 docs/dev/job/README.md 架构设计，新建 `rock/sdk/job/` 模块（Job Facade + JobExecutor + Operator + Task），重构 `rock/cli/command/job.py` 使用新模块。`rock/sdk/agent/` 完全不动。 + +**Architecture:** Config 子类决定 Trial 类型 (BashJobConfig→BashTrial, HarborJobConfig→HarborTrial)。Operator.apply(config) 生成 TrialList，JobExecutor 并行执行。Job 是极薄 Facade。 + +**Tech Stack:** Python 3.10+, Pydantic v2, asyncio, pytest (asyncio_mode=auto), ruff (line-length=120) + +--- + +### Task 1: Result Models + +**Files:** +- Create: `rock/sdk/job/__init__.py` (empty) +- Create: `rock/sdk/job/trial/__init__.py` (empty) +- Create: `rock/sdk/job/result.py` +- Create: `tests/unit/sdk/job/__init__.py` (empty) +- Create: `tests/unit/sdk/job/test_result.py` + +- [ ] **Step 1: Create directory structure** + +```bash +mkdir -p rock/sdk/job/task tests/unit/sdk/job +touch rock/sdk/job/__init__.py rock/sdk/job/trial/__init__.py tests/unit/sdk/job/__init__.py +``` + +- [ ] **Step 2: Write failing test** + +```python +# tests/unit/sdk/job/test_result.py +from rock.sdk.job.result import JobResult, JobStatus, TrialResult, TrialStatus + + +class TestTrialStatus: + def test_values(self): + assert TrialStatus.COMPLETED == "completed" + assert TrialStatus.FAILED == "failed" + assert TrialStatus.CANCELLED == "cancelled" + + def test_is_str(self): + assert isinstance(TrialStatus.COMPLETED, str) + + +class TestTrialResult: + def test_defaults(self): + r = TrialResult() + assert r.task_id == "" + assert r.status == TrialStatus.COMPLETED + assert r.output == "" + assert r.exit_code == 0 + assert r.data == {} + assert r.trial_results == [] + + def test_success_when_completed(self): + r = TrialResult(status=TrialStatus.COMPLETED) + assert r.success is True + + def test_not_success_when_failed(self): + r = TrialResult(status=TrialStatus.FAILED) + assert r.success is False + + def test_score_empty(self): + r = TrialResult() + assert r.score == 0.0 + + def test_score_with_trials(self): + from rock.sdk.agent.models.trial.result import TrialResult, VerifierResult + + r = TrialResult( + trial_results=[ + TrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), + TrialResult(task_name="t2", verifier_result=VerifierResult(rewards={"reward": 0.5})), + ] + ) + assert r.score == 0.75 + + +class TestJobStatus: + def test_values(self): + assert JobStatus.COMPLETED == "completed" + assert JobStatus.FAILED == "failed" + + +class TestJobResult: + def test_defaults(self): + r = JobResult() + assert r.job_id == "" + assert r.status == JobStatus.COMPLETED + assert r.labels == {} + assert r.task_results == [] + + def test_score_empty(self): + assert JobResult().score == 0.0 + + def test_score_with_tasks(self): + r = JobResult( + task_results=[ + TrialResult(trial_results=[]), + TrialResult(trial_results=[]), + ] + ) + assert r.score == 0.0 + + def test_n_completed_and_n_failed(self): + r = JobResult( + task_results=[ + TrialResult(status=TrialStatus.COMPLETED), + TrialResult(status=TrialStatus.FAILED), + TrialResult(status=TrialStatus.COMPLETED), + ] + ) + assert r.n_completed == 2 + assert r.n_failed == 1 + + def test_trial_results_compat_property(self): + """JobResult.trial_results returns first TrialResult's trial_results for Harbor compat.""" + from rock.sdk.agent.models.trial.result import TrialResult + + r = JobResult( + task_results=[ + TrialResult(trial_results=[TrialResult(task_name="t1")]), + ] + ) + assert len(r.trial_results) == 1 + assert r.trial_results[0].task_name == "t1" + + def test_trial_results_empty(self): + assert JobResult().trial_results == [] +``` + +- [ ] **Step 3: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_result.py -x` +Expected: FAIL — `ModuleNotFoundError: No module named 'rock.sdk.job.result'` + +- [ ] **Step 4: Write implementation** + +```python +# rock/sdk/job/result.py +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, Field + +from rock.sdk.agent.models.trial.result import TrialResult + + +class TrialStatus(str, Enum): + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class TrialResult(BaseModel): + """单个 Task 的执行结果""" + + task_id: str = "" + status: TrialStatus = TrialStatus.COMPLETED + output: str = "" + exit_code: int = 0 + data: dict = Field(default_factory=dict) + trial_results: list[TrialResult] = Field(default_factory=list) + + @property + def success(self) -> bool: + return self.status == TrialStatus.COMPLETED + + @property + def score(self) -> float: + if not self.trial_results: + return 0.0 + scores = [t.score for t in self.trial_results] + return sum(scores) / len(scores) + + +class JobStatus(str, Enum): + COMPLETED = "completed" + FAILED = "failed" + + +class JobResult(BaseModel): + """Job 聚合结果""" + + job_id: str = "" + status: JobStatus = JobStatus.COMPLETED + labels: dict[str, str] = Field(default_factory=dict) + task_results: list[TrialResult] = Field(default_factory=list) + + @property + def score(self) -> float: + if not self.task_results: + return 0.0 + scores = [t.score for t in self.task_results] + return sum(scores) / len(scores) + + @property + def n_completed(self) -> int: + return sum(1 for t in self.task_results if t.status == TrialStatus.COMPLETED) + + @property + def n_failed(self) -> int: + return sum(1 for t in self.task_results if t.status == TrialStatus.FAILED) + + @property + def trial_results(self) -> list[TrialResult]: + """便捷属性: 直接访问第一个 TrialResult 的 trial_results (Harbor 兼容)""" + if self.task_results: + return self.task_results[0].trial_results + return [] +``` + +- [ ] **Step 5: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_result.py -x -v` +Expected: ALL PASS + +- [ ] **Step 6: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/result.py tests/unit/sdk/job/test_result.py --fix +uv run ruff format rock/sdk/job/result.py tests/unit/sdk/job/test_result.py +git add rock/sdk/job/ tests/unit/sdk/job/ +git commit -m "feat(job): add result models (TrialResult, JobResult)" +``` + +--- + +### Task 2: Config Hierarchy + +**Files:** +- Create: `rock/sdk/job/config.py` +- Create: `tests/unit/sdk/job/test_config.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_config.py +import tempfile +from pathlib import Path + +from rock.sdk.agent.models.trial.config import AgentConfig, RockEnvironmentConfig, VerifierConfig +from rock.sdk.job.config import BashJobConfig, HarborJobConfig, JobConfig + + +class TestJobConfig: + def test_defaults(self): + c = JobConfig() + assert c.job_name is None + assert c.timeout == 3600 + assert c.env == {} + assert c.setup_commands == [] + assert c.file_uploads == [] + assert c.auto_stop is False + assert c.labels == {} + + def test_custom_values(self): + c = JobConfig( + job_name="test", + timeout=60, + env={"K": "V"}, + setup_commands=["echo hi"], + auto_stop=True, + ) + assert c.job_name == "test" + assert c.timeout == 60 + assert c.env == {"K": "V"} + assert c.auto_stop is True + + +class TestBashJobConfig: + def test_inherits_job_config(self): + assert issubclass(BashJobConfig, JobConfig) + + def test_bash_fields(self): + c = BashJobConfig(script="echo hello") + assert c.script == "echo hello" + assert c.script_path is None + + def test_script_path(self): + c = BashJobConfig(script_path="/tmp/test.sh") + assert c.script_path == "/tmp/test.sh" + + +class TestHarborJobConfig: + def test_inherits_job_config(self): + assert issubclass(HarborJobConfig, JobConfig) + + def test_harbor_fields(self): + c = HarborJobConfig( + agents=[AgentConfig(name="t2")], + n_attempts=3, + ) + assert c.agents[0].name == "t2" + assert c.n_attempts == 3 + + def test_to_harbor_yaml(self): + c = HarborJobConfig( + agents=[AgentConfig(name="t2", model_name="gpt-4")], + setup_commands=["pip install harbor"], + env={"KEY": "VAL"}, + ) + yaml_str = c.to_harbor_yaml() + assert "agents" in yaml_str + assert "t2" in yaml_str + # Rock-level fields should be excluded from Harbor YAML + assert "setup_commands" not in yaml_str + assert "timeout" not in yaml_str + + def test_from_yaml(self): + import yaml + + data = { + "agents": [{"name": "t2"}], + "n_attempts": 2, + } + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(data, f) + path = f.name + + c = HarborJobConfig.from_yaml(path) + assert c.agents[0].name == "t2" + assert c.n_attempts == 2 + Path(path).unlink() +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_config.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/config.py +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, Field + +from rock.sdk.agent.constants import USER_DEFINED_LOGS +from rock.sdk.agent.models.trial.config import ( + AgentConfig, + ArtifactConfig, + RockEnvironmentConfig, + TaskConfig, + VerifierConfig, +) + + +class JobConfig(BaseModel): + """Job 基础配置 — 所有 Job 类型共享的字段""" + + environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) + job_name: str | None = None + namespace: str | None = None + experiment_id: str | None = None + labels: dict[str, str] = Field(default_factory=dict) + auto_stop: bool = False + setup_commands: list[str] = Field(default_factory=list) + file_uploads: list[tuple[str, str]] = Field(default_factory=list) + env: dict[str, str] = Field(default_factory=dict) + timeout: int = 3600 + + +class BashJobConfig(JobConfig): + """Bash 脚本 Job 配置""" + + script: str | None = None + script_path: str | None = None + + +class HarborJobConfig(JobConfig): + """Harbor benchmark Job 配置""" + + agents: list[AgentConfig] = Field(default_factory=lambda: [AgentConfig()]) + datasets: list = Field(default_factory=list) + orchestrator: dict[str, Any] = Field(default_factory=dict) + verifier: VerifierConfig = Field(default_factory=VerifierConfig) + tasks: list[TaskConfig] = Field(default_factory=list) + metrics: list = Field(default_factory=list) + artifacts: list[str | ArtifactConfig] = Field(default_factory=list) + n_attempts: int = 1 + timeout_multiplier: float = 1.0 + agent_timeout_multiplier: float | None = None + verifier_timeout_multiplier: float | None = None + jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" + debug: bool = False + + # Rock-level fields to exclude when serializing to Harbor YAML + _ROCK_FIELDS: set[str] = { + "environment", "job_name", "namespace", "experiment_id", "labels", + "auto_stop", "setup_commands", "file_uploads", "env", "timeout", + } + + def to_harbor_yaml(self) -> str: + """Serialize Harbor-native fields to YAML for harbor jobs start -c""" + import yaml + + data = self.model_dump(mode="json", exclude=self._ROCK_FIELDS, exclude_none=True) + harbor_env = self.environment.to_harbor_environment() + if harbor_env: + data["environment"] = harbor_env + return yaml.dump(data, default_flow_style=False, allow_unicode=True) + + @classmethod + def from_yaml(cls, path: str) -> HarborJobConfig: + """Load from Harbor YAML config file""" + import yaml + + with open(path) as f: + data = yaml.safe_load(f) + return cls(**data) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_config.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/config.py tests/unit/sdk/job/test_config.py --fix +uv run ruff format rock/sdk/job/config.py tests/unit/sdk/job/test_config.py +git add rock/sdk/job/config.py tests/unit/sdk/job/test_config.py +git commit -m "feat(job): add config hierarchy (JobConfig, BashJobConfig, HarborJobConfig)" +``` + +--- + +### Task 3: Task Interface + Registry + +**Files:** +- Create: `rock/sdk/job/trial/abstract.py` +- Create: `rock/sdk/job/trial/registry.py` +- Create: `tests/unit/sdk/job/test_trial_registry.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_trial_registry.py +import pytest + +from rock.sdk.job.config import JobConfig +from rock.sdk.job.result import TrialResult, TrialStatus +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import _create_trial, register_trial + + +class _StubConfig(JobConfig): + stub_field: str = "test" + + +class _StubTask(AbstractTrial): + async def setup(self, sandbox): + pass + + def build(self) -> str: + return "echo stub" + + async def collect(self, sandbox, output, exit_code): + return TrialResult(task_id="stub", output=output, exit_code=exit_code) + + +class TestAbstractTrial: + def test_cannot_instantiate(self): + with pytest.raises(TypeError): + AbstractTrial(JobConfig()) + + def test_concrete_subclass(self): + config = _StubConfig() + task = _StubTask(config) + assert task._config is config + + +class TestRegistry: + def test_create_unregistered_raises(self): + class _UnknownConfig(JobConfig): + pass + + with pytest.raises(TypeError, match="No task registered"): + _create_trial(_UnknownConfig()) + + def test_register_and_create(self): + register_trial(_StubConfig, _StubTask) + trial = _create_trial(_StubConfig()) + assert isinstance(task, _StubTask) + assert task._config.stub_field == "test" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_registry.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/trial/abstract.py +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.result import TrialResult + from rock.sdk.sandbox.client import Sandbox + + +class AbstractTrial(ABC): + """Trial 基类: 在 Sandbox 中执行一个任务的三阶段接口""" + + def __init__(self, config: JobConfig): + self._config = config + + @abstractmethod + async def setup(self, sandbox: Sandbox) -> None: + """执行前: 准备 sandbox 环境""" + + @abstractmethod + def build(self) -> str: + """构建: 生成要执行的 bash 脚本""" + + @abstractmethod + async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: + """执行后: 收集并解析结果""" + + async def _upload_files(self, sandbox: Sandbox) -> None: + """共享: 上传 config.file_uploads""" + for local_path, sandbox_path in self._config.file_uploads: + await sandbox.fs.upload_dir(local_path, sandbox_path) +``` + +```python +# rock/sdk/job/trial/registry.py +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + +_TRIAL_REGISTRY: dict[type, type] = {} + + +def register_trial(config_type: type, task_type: type) -> None: + """注册 Config → Task 映射""" + _TRIAL_REGISTRY[config_type] = task_type + + +def _create_trial(config: JobConfig) -> AbstractTrial: + """根据 config 类型创建对应的 Task 实例""" + task_cls = _TRIAL_REGISTRY.get(type(config)) + if task_cls is None: + raise TypeError( + f"No task registered for {type(config).__name__}. " + f"Supported: {[c.__name__ for c in _TRIAL_REGISTRY]}" + ) + return task_cls(config) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_registry.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/trial/ tests/unit/sdk/job/test_trial_registry.py --fix +uv run ruff format rock/sdk/job/trial/ tests/unit/sdk/job/test_trial_registry.py +git add rock/sdk/job/trial/ tests/unit/sdk/job/test_trial_registry.py +git commit -m "feat(job): add AbstractTrial and task registry" +``` + +--- + +### Task 4: BashTrial + +**Files:** +- Create: `rock/sdk/job/trial/bash.py` +- Create: `tests/unit/sdk/job/test_trial_bash.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_trial_bash.py +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.result import TrialStatus +from rock.sdk.job.trial.bash import BashTrial +from rock.sdk.job.trial.registry import _create_trial + + +class TestBashTrialBuild: + def test_basic_script(self): + config = BashJobConfig(script="echo hello") + task = BashTrial(config) + script = trial.build() + assert script.startswith("#!/bin/bash") + assert "set -e" in script + assert "echo hello" in script + + def test_with_setup_commands(self): + config = BashJobConfig( + script="python main.py", + setup_commands=["pip install -r requirements.txt", "export FOO=bar"], + ) + task = BashTrial(config) + script = trial.build() + assert "pip install -r requirements.txt" in script + assert "export FOO=bar" in script + assert "python main.py" in script + # setup commands should appear before the main script + assert script.index("pip install") < script.index("python main.py") + + def test_no_setup_commands(self): + config = BashJobConfig(script="ls") + task = BashTrial(config) + script = trial.build() + assert "#!/bin/bash" in script + assert "ls" in script + + +class TestBashTrialSetup: + async def test_upload_files(self): + config = BashJobConfig( + script="echo hi", + file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], + ) + task = BashTrial(config) + mock_sandbox = AsyncMock() + mock_sandbox.fs.upload_dir = AsyncMock() + + await trial.setup(mock_sandbox) + + assert mock_sandbox.fs.upload_dir.call_count == 2 + mock_sandbox.fs.upload_dir.assert_any_call("/local/a", "/sandbox/a") + mock_sandbox.fs.upload_dir.assert_any_call("/local/b", "/sandbox/b") + + async def test_reads_script_path(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as f: + f.write("echo from file") + script_file = f.name + + try: + config = BashJobConfig(script_path=script_file) + task = BashTrial(config) + mock_sandbox = AsyncMock() + await trial.setup(mock_sandbox) + assert config.script == "echo from file" + finally: + Path(script_file).unlink() + + +class TestBashTrialCollect: + async def test_success(self): + config = BashJobConfig(script="echo hi", job_name="test-job") + task = BashTrial(config) + result = await trial.collect(AsyncMock(), "output text", 0) + assert result.status == TrialStatus.COMPLETED + assert result.output == "output text" + assert result.exit_code == 0 + assert result.task_id == "test-job" + + async def test_failure(self): + config = BashJobConfig(script="exit 1", job_name="fail-job") + task = BashTrial(config) + result = await trial.collect(AsyncMock(), "error", 1) + assert result.status == TrialStatus.FAILED + assert result.exit_code == 1 + + +class TestBashTrialRegistration: + def test_registered_in_registry(self): + trial = _create_trial(BashJobConfig(script="echo test")) + assert isinstance(task, BashTrial) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_bash.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/trial/bash.py +from __future__ import annotations + +from pathlib import Path + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.result import TrialResult, TrialStatus +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import register_trial + + +class BashTrial(AbstractTrial): + """Bash 脚本执行""" + + _config: BashJobConfig + + async def setup(self, sandbox) -> None: + await self._upload_files(sandbox) + if self._config.script_path: + self._config.script = Path(self._config.script_path).read_text() + + def build(self) -> str: + lines = ["#!/bin/bash", "set -e", ""] + if self._config.setup_commands: + for cmd in self._config.setup_commands: + lines.append(f"echo '>>> {cmd[:60]}...'") + lines.append(cmd) + lines.append("") + if self._config.script: + lines.append(self._config.script) + return "\n".join(lines) + + async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: + return TrialResult( + task_id=self._config.job_name or "", + status=TrialStatus.COMPLETED if exit_code == 0 else TrialStatus.FAILED, + output=output, + exit_code=exit_code, + ) + + +register_trial(BashJobConfig, BashTrial) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_bash.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/trial/bash.py tests/unit/sdk/job/test_trial_bash.py --fix +uv run ruff format rock/sdk/job/trial/bash.py tests/unit/sdk/job/test_trial_bash.py +git add rock/sdk/job/trial/bash.py tests/unit/sdk/job/test_trial_bash.py +git commit -m "feat(job): add BashTrial" +``` + +--- + +### Task 5: Operator + +**Files:** +- Create: `rock/sdk/job/operator.py` +- Create: `tests/unit/sdk/job/test_operator.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_operator.py +import pytest + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.operator import Operator, ScatterOperator +from rock.sdk.job.trial.bash import BashTrial + + +class TestOperatorABC: + def test_cannot_instantiate(self): + with pytest.raises(TypeError): + Operator() + + +class TestScatterOperator: + def test_default_size_1(self): + op = ScatterOperator() + assert op.size == 1 + + def test_apply_returns_one_task(self): + op = ScatterOperator() + config = BashJobConfig(script="echo hi") + tasks = op.apply(config) + assert len(tasks) == 1 + assert isinstance(tasks[0], BashTrial) + + def test_apply_returns_n_tasks(self): + op = ScatterOperator(size=3) + config = BashJobConfig(script="echo hi") + tasks = op.apply(config) + assert len(tasks) == 3 + for t in tasks: + assert isinstance(t, BashTrial) + + def test_size_zero_returns_empty(self): + op = ScatterOperator(size=0) + config = BashJobConfig(script="echo hi") + tasks = op.apply(config) + assert tasks == [] + + def test_size_negative_returns_empty(self): + op = ScatterOperator(size=-1) + config = BashJobConfig(script="echo hi") + tasks = op.apply(config) + assert tasks == [] +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_operator.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/operator.py +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from rock.sdk.job.trial.registry import _create_trial + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + + +class Operator(ABC): + """算子基类: 从 config 生成 TrialList""" + + @abstractmethod + def apply(self, config: JobConfig) -> list[AbstractTrial]: + """从 config 生成 TrialList, 返回空 list 表示什么都不做""" + ... + + +class ScatterOperator(Operator): + """Scatter 算子: 将 config 分发 8 份 Trial""" + + def __init__(self, size: int = 1): + self.size = size + + def apply(self, config: JobConfig) -> list[AbstractTrial]: + if self.size <= 0: + return [] + trial = _create_trial(config) + return [trial] * self.size +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_operator.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/operator.py tests/unit/sdk/job/test_operator.py --fix +uv run ruff format rock/sdk/job/operator.py tests/unit/sdk/job/test_operator.py +git add rock/sdk/job/operator.py tests/unit/sdk/job/test_operator.py +git commit -m "feat(job): add Operator ABC and ScatterOperator" +``` + +--- + +### Task 6: JobExecutor + +**Files:** +- Create: `rock/sdk/job/executor.py` +- Create: `tests/unit/sdk/job/test_executor.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_executor.py +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient +from rock.sdk.job.operator import ScatterOperator +from rock.sdk.job.result import TrialStatus + + +def _make_mock_sandbox(): + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-123" + sandbox._namespace = None + sandbox._experiment_id = None + sandbox.start = AsyncMock() + sandbox.close = AsyncMock() + sandbox.create_session = AsyncMock() + sandbox.write_file_by_path = AsyncMock(return_value=MagicMock(success=True)) + sandbox.fs = AsyncMock() + sandbox.fs.upload_dir = AsyncMock() + sandbox.start_nohup_process = AsyncMock(return_value=(12345, None)) + sandbox.wait_for_process_completion = AsyncMock(return_value=(True, "done")) + nohup_obs = MagicMock() + nohup_obs.output = "hello output" + nohup_obs.exit_code = 0 + sandbox.handle_nohup_output = AsyncMock(return_value=nohup_obs) + return sandbox + + +class TestJobExecutorRun: + async def test_run_bash_success(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + operator = ScatterOperator(size=1) + config = BashJobConfig(script="echo hello", job_name="test") + results = await executor.run(operator, config) + + assert len(results) == 1 + assert results[0].status == TrialStatus.COMPLETED + assert results[0].output == "hello output" + mock_sandbox.start.assert_called_once() + + async def test_run_empty_operator(self): + executor = JobExecutor() + operator = ScatterOperator(size=0) + config = BashJobConfig(script="echo hi") + results = await executor.run(operator, config) + assert results == [] + + +class TestJobExecutorSubmitWait: + async def test_submit_returns_job_client(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + operator = ScatterOperator(size=1) + config = BashJobConfig(script="echo hello", job_name="test") + job_client = await executor.submit(operator, config) + + assert isinstance(job_client, JobClient) + assert len(job_client.tasks) == 1 + assert isinstance(job_client.tasks[0], TrialClient) + + async def test_wait_returns_results(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + operator = ScatterOperator(size=1) + config = BashJobConfig(script="echo hello", job_name="test") + job_client = await executor.submit(operator, config) + results = await executor.wait(job_client) + + assert len(results) == 1 + assert results[0].status == TrialStatus.COMPLETED + + +class TestJobExecutorAutoStop: + async def test_auto_stop_closes_sandbox(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + config = BashJobConfig(script="echo hi", job_name="test", auto_stop=True) + await executor.run(ScatterOperator(), config) + mock_sandbox.close.assert_called_once() + + async def test_no_auto_stop(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + config = BashJobConfig(script="echo hi", job_name="test", auto_stop=False) + await executor.run(ScatterOperator(), config) + mock_sandbox.close.assert_not_called() + + +class TestJobExecutorFailure: + async def test_process_failure(self): + mock_sandbox = _make_mock_sandbox() + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "timeout")) + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + config = BashJobConfig(script="echo hi", job_name="test") + results = await executor.run(ScatterOperator(), config) + assert results[0].status == TrialStatus.FAILED + + async def test_nohup_start_error(self): + mock_sandbox = _make_mock_sandbox() + error_obs = MagicMock() + error_obs.output = "command not found" + mock_sandbox.start_nohup_process = AsyncMock(return_value=(None, error_obs)) + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + config = BashJobConfig(script="echo hi", job_name="test") + with pytest.raises(RuntimeError, match="Failed to start task"): + await executor.run(ScatterOperator(), config) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_executor.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/executor.py +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from rock.actions import CreateBashSessionRequest +from rock.logger import init_logger +from rock.sdk.job.operator import Operator +from rock.sdk.job.result import TrialResult, TrialStatus +from rock.sdk.sandbox.client import Sandbox + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + +logger = init_logger(__name__) + + +@dataclass +class TrialClient: + """单个运行中 task 的句柄""" + + sandbox: Sandbox + session: str + pid: int + trial: AbstractTrial + + +@dataclass +class JobClient: + """Job.submit() 返回的句柄，持有多个 TrialClient""" + + tasks: list[TrialClient] + + +class JobExecutor: + """执行引擎: 驱动 Operator 生成 TrialList，并行执行""" + + async def run(self, operator: Operator, config: JobConfig) -> list[TrialResult]: + """完整生命周期: submit + wait""" + job_client = await self.submit(operator, config) + return await self.wait(job_client) + + async def submit(self, operator: Operator, config: JobConfig) -> JobClient: + """Operator apply 生成 TrialList，并行启动所有 sandbox""" + trial_list = operator.apply(config) + if not trial_list: + return JobClient(tasks=[]) + task_clients = await asyncio.gather(*[self._(do_submit(trial) for trial in trial_list]) + return JobClient(tasks=list(task_clients)) + + async def wait(self, job_client: JobClient) -> list[TrialResult]: + """并行等待所有 task 完成，收集结果""" + if not job_client.tasks: + return [] + return list(await asyncio.gather(*[self._do_wait(tc) for tc in job_client.tasks])) + + async def _do_submit(self, trial: AbstractTrial) -> TrialClient: + """启动单个 sandbox + 执行脚本""" + config = task._config + + sandbox = Sandbox(config.environment) + await sandbox.start() + logger.info(f"Sandbox started: sandbox_id={sandbox.sandbox_id}, job_name={config.job_name}") + + session = f"rock-job-{config.job_name or 'default'}" + env = self._build_session_env(config) + await sandbox.create_session( + CreateBashSessionRequest(session=session, env_enable=True, env=env) + ) + + await trial.setup(sandbox) + script_content = trial.build() + + script_path = f"/tmp/rock_job_{config.job_name or 'default'}.sh" + await sandbox.write_file_by_path(script_content, script_path) + + tmp_file = f"/tmp/rock_job_{config.job_name or 'default'}.out" + pid, error = await sandbox.start_nohup_process( + cmd=f"bash {script_path}", tmp_file=tmp_file, session=session, + ) + if error is not None: + raise RuntimeError(f"Failed to start task: {error.output}") + + logger.info(f"Task started: pid={pid}, job_name={config.job_name}") + return TrialClient(sandbox=sandbox, session=session, pid=pid, task=task) + + async def _do_wait(self, client: TrialClient) -> TrialResult: + """等待单个 task 完成，收集结果""" + config = client.trial._config + try: + success, message = await client.sandbox.wait_for_process_completion( + pid=client.pid, + session=client.session, + wait_timeout=config.timeout, + wait_interval=30, + ) + obs = await client.sandbox.handle_nohup_output( + tmp_file=f"/tmp/rock_job_{config.job_name or 'default'}.out", + session=client.session, + success=success, + message=message, + ignore_output=False, + response_limited_bytes_in_nohup=None, + ) + result = await client.trial.collect(client.sandbox, obs.output or "", obs.exit_code or 1) + if not success: + result.status = TrialStatus.FAILED + return result + finally: + if config.auto_stop: + await client.sandbox.close() + + @staticmethod + def _build_session_env(config: JobConfig) -> dict[str, str] | None: + """Merge OSS_* from process env with config.env""" + oss_env = {k: v for k, v in os.environ.items() if k.startswith("OSS")} + merged = {**oss_env, **config.env} + return merged or None +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_executor.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/executor.py tests/unit/sdk/job/test_executor.py --fix +uv run ruff format rock/sdk/job/executor.py tests/unit/sdk/job/test_executor.py +git add rock/sdk/job/executor.py tests/unit/sdk/job/test_executor.py +git commit -m "feat(job): add JobExecutor with TrialClient/JobClient" +``` + +--- + +### Task 7: Job Facade + +**Files:** +- Create: `rock/sdk/job/job.py` +- Create: `tests/unit/sdk/job/test_job.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_job.py +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.job import Job +from rock.sdk.job.operator import ScatterOperator +from rock.sdk.job.result import JobStatus, TrialResult, TrialStatus + + +def _make_mock_sandbox(): + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-test" + sandbox._namespace = None + sandbox._experiment_id = None + sandbox.start = AsyncMock() + sandbox.close = AsyncMock() + sandbox.create_session = AsyncMock() + sandbox.write_file_by_path = AsyncMock(return_value=MagicMock(success=True)) + sandbox.fs = AsyncMock() + sandbox.arun = AsyncMock() + sandbox.start_nohup_process = AsyncMock(return_value=(99, None)) + sandbox.wait_for_process_completion = AsyncMock(return_value=(True, "done")) + obs = MagicMock() + obs.output = "ok" + obs.exit_code = 0 + sandbox.handle_nohup_output = AsyncMock(return_value=obs) + return sandbox + + +class TestJobRun: + async def test_run_returns_job_result(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + result = await job.run() + assert result.status == JobStatus.COMPLETED + assert len(result.task_results) == 1 + + async def test_run_failed(self): + mock_sandbox = _make_mock_sandbox() + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "fail")) + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="exit 1", job_name="test")) + result = await job.run() + assert result.status == JobStatus.FAILED + + +class TestJobSubmitWait: + async def test_submit_then_wait(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + await job.submit() + result = await job.wait() + assert result.status == JobStatus.COMPLETED + + async def test_wait_without_submit_raises(self): + job = Job(BashJobConfig(script="echo hi")) + with pytest.raises(RuntimeError, match="No submitted job"): + await job.wait() + + +class TestJobCancel: + async def test_cancel(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="sleep 100", job_name="test")) + await job.submit() + await job.cancel() + mock_sandbox.arun.assert_called_once() + assert "kill" in str(mock_sandbox.arun.call_args) + + +class TestJobWithOperator: + async def test_custom_scatter_size(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job( + BashJobConfig(script="echo hi", job_name="test"), + operator=ScatterOperator(size=2), + ) + result = await job.run() + assert len(result.task_results) == 2 + + async def test_default_operator(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + result = await job.run() + assert len(result.task_results) == 1 # default ScatterOperator(size=1) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_job.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/job.py +from __future__ import annotations + +from typing import TYPE_CHECKING + +from rock.sdk.job.executor import JobExecutor +from rock.sdk.job.operator import ScatterOperator +from rock.sdk.job.result import JobResult, JobStatus + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.executor import JobClient + from rock.sdk.job.operator import Operator + from rock.sdk.job.result import TrialResult + + +class Job: + """Job Facade: 极薄入口层""" + + def __init__(self, config: JobConfig, operator: Operator | None = None): + self._config = config + self._executor = JobExecutor() + self._operator = operator or ScatterOperator() + self._job_client: JobClient | None = None + + async def run(self) -> JobResult: + """完整生命周期: submit + wait""" + await self.submit() + return await self.wait() + + async def submit(self) -> None: + """非阻塞提交""" + self._job_client = await self._executor.submit(self._operator, self._config) + + async def wait(self) -> JobResult: + """等待完成""" + if not self._job_client: + raise RuntimeError("No submitted job. Call submit() first.") + task_results = await self._executor.wait(self._job_client) + return self._build_result(task_results) + + async def cancel(self) -> None: + """取消运行中的 job""" + if self._job_client: + for tc in self._job_client.tasks: + await tc.sandbox.arun(cmd=f"kill {tc.pid}", session=tc.session) + + def _build_result(self, task_results: list[TrialResult]) -> JobResult: + all_success = all(r.success for r in task_results) + return JobResult( + job_id=self._config.job_name or "", + status=JobStatus.COMPLETED if all_success else JobStatus.FAILED, + labels=self._config.labels, + task_results=task_results, + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_job.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/job.py tests/unit/sdk/job/test_job.py --fix +uv run ruff format rock/sdk/job/job.py tests/unit/sdk/job/test_job.py +git add rock/sdk/job/job.py tests/unit/sdk/job/test_job.py +git commit -m "feat(job): add Job facade" +``` + +--- + +### Task 8: HarborTrial + +**Files:** +- Create: `rock/sdk/job/trial/harbor.py` +- Create: `tests/unit/sdk/job/test_trial_harbor.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_trial_harbor.py +import json +from unittest.mock import AsyncMock, MagicMock + +from rock.sdk.agent.models.trial.config import AgentConfig +from rock.sdk.job.config import HarborJobConfig +from rock.sdk.job.result import TrialStatus +from rock.sdk.job.trial.harbor import HarborTrial +from rock.sdk.job.trial.registry import _create_trial + + +class TestHarborTrialBuild: + def test_contains_harbor_command(self): + config = HarborJobConfig(job_name="test") + task = HarborTrial(config) + script = trial.build() + assert "harbor jobs start -c" in script + assert "#!/bin/bash" in script + assert "set -e" in script + + def test_contains_dockerd_startup(self): + config = HarborJobConfig(job_name="test") + task = HarborTrial(config) + script = trial.build() + assert "dockerd" in script + + def test_with_setup_commands(self): + config = HarborJobConfig( + job_name="test", + setup_commands=["pip install harbor --quiet"], + ) + task = HarborTrial(config) + script = trial.build() + assert "pip install harbor --quiet" in script + + +class TestHarborTrialSetup: + async def test_uploads_harbor_yaml(self): + config = HarborJobConfig( + job_name="test", + agents=[AgentConfig(name="t2")], + ) + task = HarborTrial(config) + mock_sandbox = AsyncMock() + mock_sandbox.write_file_by_path = AsyncMock(return_value=MagicMock(success=True)) + mock_sandbox.fs = AsyncMock() + + await trial.setup(mock_sandbox) + + mock_sandbox.write_file_by_path.assert_called_once() + yaml_content = mock_sandbox.write_file_by_path.call_args[0][0] + assert "agents" in yaml_content + assert "t2" in yaml_content + + +class TestHarborTrialCollect: + async def test_with_trial_results(self): + config = HarborJobConfig(job_name="test") + task = HarborTrial(config) + + mock_sandbox = AsyncMock() + execute_result = MagicMock() + execute_result.stdout = "/jobs/test/trials/trial-001/result.json" + mock_sandbox.execute = AsyncMock(return_value=execute_result) + + read_response = MagicMock() + read_response.content = json.dumps({ + "trial_name": "trial-001", + "task_name": "t1", + "started_at": "2026-01-01T00:00:00Z", + "finished_at": "2026-01-01T00:01:00Z", + "verifier_result": {"rewards": {"reward": 1.0}}, + "agent_result": {}, + "exception_info": None, + }) + mock_sandbox.read_file = AsyncMock(return_value=read_response) + + result = await trial.collect(mock_sandbox, "harbor completed", 0) + assert result.status == TrialStatus.COMPLETED + assert len(result.trial_results) == 1 + assert result.trial_results[0].score == 1.0 + + async def test_no_trial_results(self): + config = HarborJobConfig(job_name="test") + task = HarborTrial(config) + + mock_sandbox = AsyncMock() + execute_result = MagicMock() + execute_result.stdout = "" + mock_sandbox.execute = AsyncMock(return_value=execute_result) + + result = await trial.collect(mock_sandbox, "error", 1) + assert result.status == TrialStatus.FAILED + assert result.trial_results == [] + + +class TestHarborTrialRegistration: + def test_registered(self): + trial = _create_trial(HarborJobConfig(job_name="test")) + assert isinstance(task, HarborTrial) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_harbor.py -x` +Expected: FAIL — `ModuleNotFoundError` + +- [ ] **Step 3: Write implementation** + +```python +# rock/sdk/job/trial/harbor.py +from __future__ import annotations + +import json + +from rock.actions import Command, ReadFileRequest +from rock.logger import init_logger +from rock.sdk.agent.constants import USER_DEFINED_LOGS +from rock.sdk.agent.models.trial.result import TrialResult +from rock.sdk.job.config import HarborJobConfig +from rock.sdk.job.result import TrialResult, TrialStatus +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import register_trial + +logger = init_logger(__name__) + +_HARBOR_SCRIPT_TEMPLATE = r"""#!/bin/bash +set -e + +# ── Detect and start dockerd ───────────────────────────────────────── +if command -v docker &>/dev/null; then + echo "docker OK: $(command -v docker)" + if ! pgrep -x dockerd &>/dev/null; then + echo "Starting dockerd..." + nohup dockerd &>/var/log/dockerd.log & + fi + for i in $(seq 1 60); do + if docker info &>/dev/null; then echo "dockerd is ready"; break; fi + sleep 1 + if [ "$i" -eq 60 ]; then echo "WARN: dockerd failed to start within 60s"; fi + done +fi + +# ── Ensure output directory exists ────────────────────────────────── +mkdir -p {user_defined_dir} + +# ── Setup commands ─────────────────────────────────────────────────── +{setup_commands} + +# ── Harbor run ─────────────────────────────────────────────────────── +harbor jobs start -c {config_path} +""" + + +class HarborTrial(AbstractTrial): + """Harbor benchmark 执行""" + + _config: HarborJobConfig + + async def setup(self, sandbox) -> None: + await self._upload_files(sandbox) + yaml_content = self._config.to_harbor_yaml() + config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" + await sandbox.write_file_by_path(yaml_content, config_path) + + def build(self) -> str: + setup_lines = [] + for cmd in self._config.setup_commands: + setup_lines.append(f"echo '>>> {cmd[:60]}...'") + setup_lines.append(cmd) + setup_block = "\n".join(setup_lines) if setup_lines else "echo 'No setup commands'" + + config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" + return _HARBOR_SCRIPT_TEMPLATE.format( + setup_commands=setup_block, + config_path=config_path, + user_defined_dir=USER_DEFINED_LOGS, + ) + + async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: + trial_results = await self._collect_trial_results(sandbox) + return TrialResult( + task_id=self._config.job_name or "", + status=TrialStatus.COMPLETED if trial_results else TrialStatus.FAILED, + output=output, + exit_code=exit_code, + trial_results=trial_results, + ) + + async def _collect_trial_results(self, sandbox) -> list[TrialResult]: + """Read trial-level result.json files from sandbox.""" + job_dir = f"{self._config.jobs_dir}/{self._config.job_name}" + try: + list_result = await sandbox.execute( + Command(command=["find", job_dir, "-mindepth", "2", "-maxdepth", "2", "-name", "result.json"]) + ) + trial_files = [ + line.strip() for line in (list_result.stdout or "").strip().split("\n") if line.strip() + ] + except Exception: + trial_files = [] + + results: list[TrialResult] = [] + for trial_file in trial_files: + try: + response = await sandbox.read_file(ReadFileRequest(path=trial_file)) + data = json.loads(response.content) + results.append(TrialResult.from_harbor_json(data)) + except Exception as e: + logger.warning(f"Failed to parse trial result {trial_file}: {e}") + + return results + + +register_trial(HarborJobConfig, HarborTrial) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_trial_harbor.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/sdk/job/trial/harbor.py tests/unit/sdk/job/test_trial_harbor.py --fix +uv run ruff format rock/sdk/job/trial/harbor.py tests/unit/sdk/job/test_trial_harbor.py +git add rock/sdk/job/trial/harbor.py tests/unit/sdk/job/test_trial_harbor.py +git commit -m "feat(job): add HarborTrial (extracted from agent/job.py)" +``` + +--- + +### Task 9: CLI Update + +**Files:** +- Modify: `rock/cli/command/job.py` +- Create: `tests/unit/sdk/job/test_cli_job.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_cli_job.py +import argparse +from unittest.mock import AsyncMock, MagicMock, patch + +from rock.cli.command.job import JobCommand +from rock.sdk.job.config import BashJobConfig, HarborJobConfig + + +class TestJobCommandParser: + async def test_parser_has_type_argument(self): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + await JobCommand.add_parser_to(subparsers) + + args = parser.parse_args(["job", "run", "--type", "bash", "--script-content", "echo hi"]) + assert args.type == "bash" + + async def test_parser_default_type_is_bash(self): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + await JobCommand.add_parser_to(subparsers) + + args = parser.parse_args(["job", "run", "--script-content", "echo hi"]) + assert args.type == "bash" + + async def test_parser_harbor_type(self): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + await JobCommand.add_parser_to(subparsers) + + args = parser.parse_args(["job", "run", "--type", "harbor", "--config", "/tmp/c.yaml"]) + assert args.type == "harbor" + assert args.config == "/tmp/c.yaml" + + +class TestJobCommandRun: + async def test_bash_run(self): + args = argparse.Namespace( + job_command="run", + type="bash", + script=None, + script_content="echo hello", + image=None, + memory=None, + cpus=None, + local_path=None, + target_path="/root/job", + timeout=3600, + config=None, + base_url=None, + cluster=None, + extra_headers=None, + ) + mock_result = MagicMock() + mock_result.exit_code = 0 + mock_result.raw_output = "hello" + + with patch("rock.cli.command.job.Job") as MockJob: + mock_job_instance = AsyncMock() + mock_job_instance.run = AsyncMock(return_value=mock_result) + MockJob.return_value = mock_job_instance + + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_called_once() + called_config = MockJob.call_args[0][0] + assert isinstance(called_config, BashJobConfig) + assert called_config.script == "echo hello" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_cli_job.py -x` +Expected: FAIL — old JobCommand doesn't have `--type` or use new Job + +- [ ] **Step 3: Rewrite CLI** + +```python +# rock/cli/command/job.py +import argparse + +from rock.cli.command.command import Command +from rock.logger import init_logger + +logger = init_logger(__name__) + + +class JobCommand(Command): + name = "job" + + async def arun(self, args: argparse.Namespace): + if args.job_command == "run": + await self._job_run(args) + else: + logger.error(f"Unknown job subcommand: {args.job_command}") + + async def _job_run(self, args: argparse.Namespace): + from rock.sdk.agent.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig, HarborJobConfig + from rock.sdk.job.job import Job + + job_type = args.type or "bash" + + if job_type == "bash": + if not args.script and not args.script_content: + logger.error("Either --script or --script-content is required for bash type") + return + if args.script and args.script_content: + logger.error("--script and --script-content cannot be used together") + return + + env_kwargs = {} + if args.image: + env_kwargs["image"] = args.image + if args.memory: + env_kwargs["memory"] = args.memory + if args.cpus: + env_kwargs["cpus"] = args.cpus + if args.base_url: + env_kwargs["base_url"] = args.base_url + if args.cluster: + env_kwargs["cluster"] = args.cluster + + file_uploads = [] + if args.local_path: + file_uploads.append((args.local_path, args.target_path)) + + extra_headers = {} + if args.extra_headers: + extra_headers = args.extra_headers + + if extra_headers: + env_kwargs["extra_headers"] = extra_headers + + config = BashJobConfig( + script=args.script_content, + script_path=args.script, + environment=RockEnvironmentConfig(**env_kwargs), + file_uploads=file_uploads, + timeout=args.timeout, + auto_stop=True, + ) + + elif job_type == "harbor": + if not args.config: + logger.error("--config is required for harbor type") + return + config = HarborJobConfig.from_yaml(args.config) + if args.image: + config.environment.image = args.image + config.auto_stop = True + + else: + logger.error(f"Unknown job type: {job_type}") + return + + try: + result = await Job(config).run() + if result.task_results: + for tr in result.task_results: + if tr.output: + print(tr.output) + logger.info(f"Job completed: status={result.status}, score={result.score}") + except Exception as e: + logger.error(f"Job failed: {e}") + + @staticmethod + async def add_parser_to(subparsers: argparse._SubParsersAction): + job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") + job_subparsers = job_parser.add_subparsers(dest="job_command") + + run_parser = job_subparsers.add_parser("run", help="Run a job in a sandbox") + run_parser.add_argument("--type", choices=["bash", "harbor"], default="bash", help="Job type") + # bash args + run_parser.add_argument("--script", default=None, help="Path to script file") + run_parser.add_argument("--script-content", default=None, help="Inline script content") + # harbor args + run_parser.add_argument("--config", default=None, help="Harbor YAML config path") + # shared args + run_parser.add_argument("--image", default=None, help="Sandbox image") + run_parser.add_argument("--memory", default=None, help="Memory (e.g. 8g)") + run_parser.add_argument("--cpus", default=None, type=float, help="CPU count") + run_parser.add_argument("--timeout", type=int, default=3600, help="Timeout in seconds") + run_parser.add_argument("--local-path", default=None, help="Local dir to upload") + run_parser.add_argument("--target-path", default="/root/job", help="Target dir in sandbox") +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/unit/sdk/job/test_cli_job.py -x -v` +Expected: ALL PASS + +- [ ] **Step 5: Lint + commit** + +```bash +uv run ruff check rock/cli/command/job.py tests/unit/sdk/job/test_cli_job.py --fix +uv run ruff format rock/cli/command/job.py tests/unit/sdk/job/test_cli_job.py +git add rock/cli/command/job.py tests/unit/sdk/job/test_cli_job.py +git commit -m "feat(job): update CLI with --type bash/harbor routing" +``` + +--- + +### Task 10: Integration Wiring + Final Verification + +**Files:** +- Modify: `rock/sdk/job/__init__.py` +- Modify: `rock/sdk/job/trial/__init__.py` +- Create: `tests/unit/sdk/job/test_integration.py` + +- [ ] **Step 1: Write failing test** + +```python +# tests/unit/sdk/job/test_integration.py + + +class TestPublicImports: + def test_import_job(self): + from rock.sdk.job import Job + + assert Job is not None + + def test_import_configs(self): + from rock.sdk.job import BashJobConfig, HarborJobConfig, JobConfig + + assert issubclass(BashJobConfig, JobConfig) + assert issubclass(HarborJobConfig, JobConfig) + + def test_import_results(self): + from rock.sdk.job import JobResult, JobStatus, TrialResult, TrialStatus + + assert TrialStatus.COMPLETED == "completed" + assert JobStatus.COMPLETED == "completed" + + def test_import_operator(self): + from rock.sdk.job import Operator, ScatterOperator + + assert issubclass(ScatterOperator, Operator) + + def test_import_task(self): + from rock.sdk.job import AbstractTrial, register_trial + + assert AbstractTrial is not None + assert callable(register_trial) + + def test_trial_registry_has_both_types(self): + from rock.sdk.job import BashJobConfig, HarborJobConfig + from rock.sdk.job.trial.registry import _TRIAL_REGISTRY + + assert BashJobConfig in _TRIAL_REGISTRY + assert HarborJobConfig in _TRIAL_REGISTRY + + +class TestAgentBackwardCompat: + def test_old_imports_still_work(self): + from rock.sdk.agent import Job, JobConfig, JobResult, JobStatus + + assert Job is not None + assert JobConfig is not None + assert JobResult is not None + assert JobStatus is not None +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/unit/sdk/job/test_integration.py -x` +Expected: FAIL — `ImportError` from empty `__init__.py` + +- [ ] **Step 3: Write __init__.py files** + +```python +# rock/sdk/job/trial/__init__.py +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import _create_trial, register_trial + +__all__ = ["AbstractTrial", "register_trial", "_create_trial"] +``` + +```python +# rock/sdk/job/__init__.py +from rock.sdk.job.config import BashJobConfig, HarborJobConfig, JobConfig +from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient +from rock.sdk.job.job import Job +from rock.sdk.job.operator import Operator, ScatterOperator +from rock.sdk.job.result import JobResult, JobStatus, TrialResult, TrialStatus +from rock.sdk.job.trial import AbstractTrial, register_trial + +# Auto-register task types +import rock.sdk.job.trial.bash # noqa: F401 +import rock.sdk.job.trial.harbor # noqa: F401 + +__all__ = [ + "Job", + "JobConfig", + "BashJobConfig", + "HarborJobConfig", + "JobResult", + "JobStatus", + "TrialResult", + "TrialStatus", + "JobExecutor", + "JobClient", + "TrialClient", + "Operator", + "ScatterOperator", + "AbstractTrial", + "register_trial", +] +``` + +- [ ] **Step 4: Run ALL tests** + +```bash +# New module tests +uv run pytest tests/unit/sdk/job/ -x -v + +# Backward compat: existing agent tests still pass +uv run pytest tests/unit/sdk/agent/ -x -v + +# Full lint +uv run ruff check rock/sdk/job/ tests/unit/sdk/job/ --fix +uv run ruff format rock/sdk/job/ tests/unit/sdk/job/ +``` + +Expected: ALL PASS + +- [ ] **Step 5: Commit** + +```bash +git add rock/sdk/job/__init__.py rock/sdk/job/trial/__init__.py tests/unit/sdk/job/test_integration.py +git commit -m "feat(job): integration wiring and public exports" +``` + +--- + +## Verification + +After all tasks complete: + +```bash +# 1. Run all new tests +uv run pytest tests/unit/sdk/job/ -v + +# 2. Run existing agent tests (must not break) +uv run pytest tests/unit/sdk/agent/ -v + +# 3. Run full fast test suite +uv run pytest -m "not need_ray and not need_admin and not need_admin_and_network" --reruns 1 + +# 4. Lint everything +uv run ruff check rock/sdk/job/ rock/cli/command/job.py --fix +uv run ruff format rock/sdk/job/ rock/cli/command/job.py +``` diff --git a/examples/harbor/harbor_demo.py b/examples/harbor/harbor_demo.py index 61240e3324..f43b69d390 100644 --- a/examples/harbor/harbor_demo.py +++ b/examples/harbor/harbor_demo.py @@ -32,7 +32,7 @@ import os import sys -from rock.sdk.agent import Job, JobConfig +from rock.sdk.bench import Job, JobConfig _REQUIRED_ENV_VARS = [ "OSS_ACCESS_KEY_ID", diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index 64666cd78b..b4c97c0cf1 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -1,10 +1,7 @@ import argparse -from pathlib import Path from rock.cli.command.command import Command from rock.logger import init_logger -from rock.sdk.sandbox.client import Sandbox -from rock.sdk.sandbox.config import SandboxConfig logger = init_logger(__name__) @@ -19,103 +16,90 @@ async def arun(self, args: argparse.Namespace): logger.error(f"Unknown job subcommand: {args.job_command}") async def _job_run(self, args: argparse.Namespace): - # 1. Validate script source - if args.script and args.script_content: - logger.error("--script and --script-content cannot be used together") - return - if not args.script and not args.script_content: - logger.error("Either --script or --script-content is required") - return + # Import lazily to avoid pulling in bench/Harbor modules for bash-only uses + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job import Job + from rock.sdk.job.config import BashJobConfig - if args.script: - script_path = Path(args.script).resolve() - if not script_path.exists(): - logger.error(f"Script not found: {script_path}") - return - if not script_path.is_file(): - logger.error(f"Not a file: {script_path}") + job_type = args.type or "bash" + + if job_type == "bash": + if not args.script and not args.script_content: + logger.error("Either --script or --script-content is required for bash type") return - script_content = script_path.read_text() - else: - script_content = args.script_content - - # 2. Validate local path (optional) - src_dir = None - target_path = args.target_path - if args.local_path: - local_path = Path(args.local_path).resolve() - if not local_path.exists(): - logger.error(f"Local path not found: {local_path}") + if args.script and args.script_content: + logger.error("--script and --script-content cannot be used together") return - src_dir = str(local_path) - - # 3. Build sandbox config - sandbox_config = SandboxConfig() - if args.image: - sandbox_config.image = args.image - if args.memory: - sandbox_config.memory = args.memory - if args.cpus: - sandbox_config.cpus = args.cpus - if args.base_url: - sandbox_config.base_url = args.base_url - if args.cluster: - sandbox_config.cluster = args.cluster - if args.extra_headers: - sandbox_config.extra_headers.update(args.extra_headers) - - sandbox = Sandbox(sandbox_config) - try: - # 4. Start sandbox - logger.info(f"Starting sandbox with image={sandbox_config.image} ...") - await sandbox.start() - logger.info(f"Sandbox started: id={sandbox.sandbox_id}, ip={sandbox.host_ip}") - - # 5. Copy source directory to sandbox (optional) - if src_dir is not None: - assert sandbox.fs is not None - logger.info(f"Uploading {src_dir} -> {target_path}") - result = await sandbox.fs.upload_dir(source_dir=src_dir, target_dir=target_path) - if result.exit_code != 0: - logger.error(f"Upload failed: {result.failure_reason}") - return - - # 6. Execute the script - assert sandbox.process is not None - logger.info(f"Executing script: {script_content}") - result = await sandbox.process.execute_script( - script_content=script_content, - wait_timeout=args.timeout, + env_kwargs = {} + if args.image: + env_kwargs["image"] = args.image + if args.memory: + env_kwargs["memory"] = args.memory + if args.cpus: + env_kwargs["cpus"] = args.cpus + if getattr(args, "base_url", None): + env_kwargs["base_url"] = args.base_url + if getattr(args, "cluster", None): + env_kwargs["cluster"] = args.cluster + if getattr(args, "extra_headers", None): + env_kwargs["extra_headers"] = args.extra_headers + + file_uploads = [] + if args.local_path: + file_uploads.append((args.local_path, args.target_path)) + + config = BashJobConfig( + script=args.script_content, + script_path=args.script, + environment=RockEnvironmentConfig(**env_kwargs), + file_uploads=file_uploads, + timeout=args.timeout, + auto_stop=True, ) - # 7. Print output - if result.output: - print(result.output) - logger.info(f"Script exited with code: {result.exit_code}") + elif job_type == "harbor": + if not args.config: + logger.error("--config is required for harbor type") + return + from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + + config = HarborJobConfig.from_yaml(args.config) + if args.image: + config.environment.image = args.image + config.auto_stop = True + + else: + logger.error(f"Unknown job type: {job_type}") + return - finally: - logger.info("Stopping sandbox ...") - await sandbox.stop() + try: + result = await Job(config).run() + if result.trial_results: + for tr in result.trial_results: + output = getattr(tr, "raw_output", None) or "" + if output: + print(output) + logger.info(f"Job completed: status={result.status}") + except Exception as e: + logger.error(f"Job failed: {e}") @staticmethod async def add_parser_to(subparsers: argparse._SubParsersAction): job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") job_subparsers = job_parser.add_subparsers(dest="job_command") - # run subcommand run_parser = job_subparsers.add_parser("run", help="Run a job in a sandbox") - run_parser.add_argument("--image", default=None, help="Sandbox image (overrides default)") - run_parser.add_argument("--memory", default=None, help="Memory allocation (e.g., 8g)") - run_parser.add_argument("--cpus", default=None, type=float, help="CPU allocation (e.g., 2)") - - run_parser.add_argument("--local-path", default=None, help="Local directory to upload to the sandbox") - run_parser.add_argument( - "--target-path", default="/root/job", help="Target directory in sandbox (default: /root/job)" - ) - run_parser.add_argument("--script", default=None, help="Path to the script to execute in the sandbox") - run_parser.add_argument("--script-content", default=None, help="Script content to execute directly") - - run_parser.add_argument( - "--timeout", type=int, default=3600, help="Script execution timeout in seconds (default: 3600)" - ) + run_parser.add_argument("--type", choices=["bash", "harbor"], default="bash", help="Job type (default: bash)") + # bash args + run_parser.add_argument("--script", default=None, help="Path to script file") + run_parser.add_argument("--script-content", default=None, help="Inline script content") + # harbor args + run_parser.add_argument("--config", default=None, help="Harbor YAML config path") + # shared args + run_parser.add_argument("--image", default=None, help="Sandbox image") + run_parser.add_argument("--memory", default=None, help="Memory (e.g. 8g)") + run_parser.add_argument("--cpus", default=None, type=float, help="CPU count") + run_parser.add_argument("--timeout", type=int, default=3600, help="Timeout in seconds") + run_parser.add_argument("--local-path", default=None, help="Local dir to upload") + run_parser.add_argument("--target-path", default="/root/job", help="Target dir in sandbox") diff --git a/rock/sdk/agent/models/job/result.py b/rock/sdk/agent/models/job/result.py deleted file mode 100644 index a4544e85d3..0000000000 --- a/rock/sdk/agent/models/job/result.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Job result models aligned with harbor.models.job.result.""" - -from __future__ import annotations - -from enum import Enum - -from pydantic import BaseModel, Field - -from rock.sdk.agent.models.trial.result import TrialResult - - -class JobStatus(str, Enum): - PENDING = "pending" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class JobResult(BaseModel): - """Aligned with harbor.models.job.result.JobResult""" - - job_id: str = "" - status: JobStatus = JobStatus.COMPLETED - labels: dict[str, str] = Field(default_factory=dict) - trial_results: list[TrialResult] = Field(default_factory=list) - raw_output: str = "" - exit_code: int = 0 - - @property - def score(self) -> float: - if not self.trial_results: - return 0.0 - scores = [t.score for t in self.trial_results] - return sum(scores) / len(scores) - - @property - def n_completed(self) -> int: - return sum(1 for t in self.trial_results if t.status == "completed") - - @property - def n_failed(self) -> int: - return sum(1 for t in self.trial_results if t.status == "failed") diff --git a/rock/sdk/agent/__init__.py b/rock/sdk/bench/__init__.py similarity index 60% rename from rock/sdk/agent/__init__.py rename to rock/sdk/bench/__init__.py index 7f034a043c..d989a5812b 100644 --- a/rock/sdk/agent/__init__.py +++ b/rock/sdk/bench/__init__.py @@ -1,5 +1,5 @@ -from rock.sdk.agent.job import Job -from rock.sdk.agent.models.job.config import ( +from rock.sdk.bench.job import Job +from rock.sdk.bench.models.job.config import ( JobConfig, LocalDatasetConfig, OrchestratorConfig, @@ -7,25 +7,25 @@ RegistryDatasetConfig, RemoteRegistryInfo, RetryConfig, - RockEnvironmentConfig, ) -from rock.sdk.agent.models.job.result import JobResult, JobStatus -from rock.sdk.agent.models.metric.config import MetricConfig -from rock.sdk.agent.models.trial.config import ( +from rock.sdk.bench.models.metric.config import MetricConfig +from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, EnvironmentConfig, OssMirrorConfig, + RockEnvironmentConfig, TaskConfig, VerifierConfig, ) -from rock.sdk.agent.models.trial.result import ( +from rock.sdk.bench.models.trial.result import ( AgentInfo, AgentResult, ExceptionInfo, TrialResult, VerifierResult, ) +from rock.sdk.job.result import JobResult, JobStatus __all__ = [ "Job", @@ -52,3 +52,8 @@ "ArtifactConfig", "MetricConfig", ] + +# Register HarborTrial with the job trial registry at the end of bench init. +# Must run AFTER bench.models.trial.result is fully loaded, so it lives here +# (not in rock.sdk.job.__init__, which may be triggered mid-bench-load). +import rock.sdk.job.trial.harbor # noqa: F401, E402 # isort: skip diff --git a/rock/sdk/agent/constants.py b/rock/sdk/bench/constants.py similarity index 100% rename from rock/sdk/agent/constants.py rename to rock/sdk/bench/constants.py diff --git a/rock/sdk/agent/job.py b/rock/sdk/bench/job.py similarity index 98% rename from rock/sdk/agent/job.py rename to rock/sdk/bench/job.py index ea68d0051e..d89dea6a47 100644 --- a/rock/sdk/agent/job.py +++ b/rock/sdk/bench/job.py @@ -14,9 +14,9 @@ from rock.actions import Command, CreateBashSessionRequest, ReadFileRequest from rock.logger import init_logger -from rock.sdk.agent.constants import CHECK_INTERVAL, DEFAULT_WAIT_TIMEOUT, USER_DEFINED_LOGS -from rock.sdk.agent.models.job.result import JobResult, JobStatus -from rock.sdk.agent.models.trial.result import TrialResult +from rock.sdk.bench.constants import CHECK_INTERVAL, DEFAULT_WAIT_TIMEOUT, USER_DEFINED_LOGS +from rock.sdk.bench.models.trial.result import TrialResult +from rock.sdk.job.result import JobResult, JobStatus logger = init_logger(__name__) @@ -63,7 +63,7 @@ class Job: """ def __init__(self, config): - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig if not isinstance(config, JobConfig): raise TypeError(f"config must be JobConfig, got {type(config)}") diff --git a/rock/sdk/agent/models/__init__.py b/rock/sdk/bench/models/__init__.py similarity index 62% rename from rock/sdk/agent/models/__init__.py rename to rock/sdk/bench/models/__init__.py index be3cdffa69..c0ce31d8f8 100644 --- a/rock/sdk/agent/models/__init__.py +++ b/rock/sdk/bench/models/__init__.py @@ -1,19 +1,19 @@ -from rock.sdk.agent.models.environment_type import EnvironmentType -from rock.sdk.agent.models.job.config import ( +from rock.sdk.bench.models.environment_type import EnvironmentType +from rock.sdk.bench.models.job.config import ( DatasetConfig, JobConfig, OrchestratorConfig, RetryConfig, - RockEnvironmentConfig, ) -from rock.sdk.agent.models.metric.config import MetricConfig -from rock.sdk.agent.models.metric.type import MetricType -from rock.sdk.agent.models.orchestrator_type import OrchestratorType -from rock.sdk.agent.models.trial.config import ( +from rock.sdk.bench.models.metric.config import MetricConfig +from rock.sdk.bench.models.metric.type import MetricType +from rock.sdk.bench.models.orchestrator_type import OrchestratorType +from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, EnvironmentConfig, OssMirrorConfig, + RockEnvironmentConfig, TaskConfig, VerifierConfig, ) diff --git a/rock/sdk/agent/models/environment_type.py b/rock/sdk/bench/models/environment_type.py similarity index 100% rename from rock/sdk/agent/models/environment_type.py rename to rock/sdk/bench/models/environment_type.py diff --git a/rock/sdk/agent/models/job/__init__.py b/rock/sdk/bench/models/job/__init__.py similarity index 77% rename from rock/sdk/agent/models/job/__init__.py rename to rock/sdk/bench/models/job/__init__.py index 51370f57d9..a30c186b62 100644 --- a/rock/sdk/agent/models/job/__init__.py +++ b/rock/sdk/bench/models/job/__init__.py @@ -1,3 +1,6 @@ +from rock.sdk.bench.models.trial.config import RockEnvironmentConfig +from rock.sdk.job.result import JobResult, JobStatus + from .config import ( JobConfig, LocalDatasetConfig, @@ -6,9 +9,7 @@ RegistryDatasetConfig, RemoteRegistryInfo, RetryConfig, - RockEnvironmentConfig, ) -from .result import JobResult, JobStatus __all__ = [ "JobConfig", diff --git a/rock/sdk/agent/models/job/config.py b/rock/sdk/bench/models/job/config.py similarity index 81% rename from rock/sdk/agent/models/job/config.py rename to rock/sdk/bench/models/job/config.py index 2f9b37345b..080f7e3063 100644 --- a/rock/sdk/agent/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -1,27 +1,28 @@ """Job configuration models aligned with harbor.models.job.config. Harbor-native fields are serialized to YAML and passed to ``harbor jobs start -c``. +JobConfig inherits from rock.sdk.job.config.JobConfig (base). """ from __future__ import annotations -from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, ClassVar from pydantic import BaseModel, Field, model_validator -from rock.sdk.agent.constants import USER_DEFINED_LOGS -from rock.sdk.agent.models.metric.config import MetricConfig -from rock.sdk.agent.models.orchestrator_type import OrchestratorType -from rock.sdk.agent.models.trial.config import ( +from rock.sdk.bench.constants import USER_DEFINED_LOGS +from rock.sdk.bench.models.metric.config import MetricConfig +from rock.sdk.bench.models.orchestrator_type import OrchestratorType +from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, OssMirrorConfig, - RockEnvironmentConfig, + RockEnvironmentConfig, # noqa: F401 — re-exported for backward compat TaskConfig, VerifierConfig, ) +from rock.sdk.job.config import JobConfig as _BaseJobConfig # --------------------------------------------------------------------------- # RetryConfig / OrchestratorConfig @@ -127,30 +128,15 @@ def _infer_version_from_split(self): DatasetConfig = LocalDatasetConfig | RegistryDatasetConfig -class JobConfig(BaseModel): - """Job configuration: Rock environment + Harbor-native benchmark fields. +class JobConfig(_BaseJobConfig): + """Harbor Job configuration: extends base JobConfig with Harbor-native fields. - All Rock sandbox/lifecycle configuration lives in ``environment``. + All Rock sandbox/lifecycle configuration lives in ``environment`` (inherited). Harbor-native fields (agents, datasets, etc.) are serialized to YAML and passed to ``harbor jobs start -c``. """ - # ── Rock environment (Rock sandbox config + Harbor EnvironmentConfig, not serialized to Harbor YAML) ── - environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) - - # ── Harbor native fields ── - namespace: str | None = Field( - default=None, - description="Tenant isolation identifier for distinguishing resources across teams/projects", - ) - experiment_id: str | None = Field( - default=None, - description="Experiment identifier", - ) - job_name: str | None = Field( - default=None, - description="Job name, auto-generated if not set", - ) + # ── Harbor native fields (base fields: environment, job_name, namespace, etc. are inherited) ── jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" n_attempts: int = 1 timeout_multiplier: float = 1.0 @@ -166,13 +152,6 @@ class JobConfig(BaseModel): datasets: list[LocalDatasetConfig | RegistryDatasetConfig] = Field(default_factory=list) tasks: list[TaskConfig] = Field(default_factory=list) artifacts: list[str | ArtifactConfig] = Field(default_factory=list) - labels: dict[str, str] = Field( - default_factory=dict, - description="Key-value labels for organizing and filtering jobs. " - "Example: {'step': '42', 'env': 'prod'}. " - "Keys: [prefix/]name, lowercase, max 63 chars. " - "Values: max 255 chars. Reserved prefix: 'harbor.io/'.", - ) @model_validator(mode="after") def _sync_experiment_id(self): @@ -193,15 +172,19 @@ def _sync_experiment_id(self): self.environment.experiment_id = self.experiment_id return self + # Base JobConfig fields to exclude when serializing to Harbor YAML + _BASE_FIELDS: ClassVar[set[str]] = set(_BaseJobConfig.model_fields.keys()) + def to_harbor_yaml(self) -> str: """Serialize Harbor-native fields to YAML for ``harbor jobs start -c``. - Rock environment fields are excluded. Harbor environment fields - (force_build, override_cpus, etc.) are included under ``environment``. + Base JobConfig fields (environment, job_name, setup_commands, etc.) + are excluded. Harbor environment fields (force_build, override_cpus, etc.) + are re-injected under ``environment``. """ import yaml - data = self.model_dump(mode="json", exclude={"environment"}, exclude_none=True) + data = self.model_dump(mode="json", exclude=self._BASE_FIELDS, exclude_none=True) harbor_env = self.environment.to_harbor_environment() if harbor_env: data["environment"] = harbor_env diff --git a/rock/sdk/agent/models/metric/__init__.py b/rock/sdk/bench/models/metric/__init__.py similarity index 100% rename from rock/sdk/agent/models/metric/__init__.py rename to rock/sdk/bench/models/metric/__init__.py diff --git a/rock/sdk/agent/models/metric/config.py b/rock/sdk/bench/models/metric/config.py similarity index 81% rename from rock/sdk/agent/models/metric/config.py rename to rock/sdk/bench/models/metric/config.py index 805ec840c9..5829cab01a 100644 --- a/rock/sdk/agent/models/metric/config.py +++ b/rock/sdk/bench/models/metric/config.py @@ -4,7 +4,7 @@ from pydantic import BaseModel, Field -from rock.sdk.agent.models.metric.type import MetricType +from rock.sdk.bench.models.metric.type import MetricType class MetricConfig(BaseModel): diff --git a/rock/sdk/agent/models/metric/type.py b/rock/sdk/bench/models/metric/type.py similarity index 100% rename from rock/sdk/agent/models/metric/type.py rename to rock/sdk/bench/models/metric/type.py diff --git a/rock/sdk/agent/models/orchestrator_type.py b/rock/sdk/bench/models/orchestrator_type.py similarity index 100% rename from rock/sdk/agent/models/orchestrator_type.py rename to rock/sdk/bench/models/orchestrator_type.py diff --git a/rock/sdk/agent/models/trial/__init__.py b/rock/sdk/bench/models/trial/__init__.py similarity index 100% rename from rock/sdk/agent/models/trial/__init__.py rename to rock/sdk/bench/models/trial/__init__.py diff --git a/rock/sdk/agent/models/trial/config.py b/rock/sdk/bench/models/trial/config.py similarity index 98% rename from rock/sdk/agent/models/trial/config.py rename to rock/sdk/bench/models/trial/config.py index 718daf57b5..6b104c402c 100644 --- a/rock/sdk/agent/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -5,7 +5,7 @@ from pydantic import BaseModel, Field -from rock.sdk.agent.models.environment_type import EnvironmentType +from rock.sdk.bench.models.environment_type import EnvironmentType from rock.sdk.sandbox.config import SandboxConfig diff --git a/rock/sdk/agent/models/trial/result.py b/rock/sdk/bench/models/trial/result.py similarity index 74% rename from rock/sdk/agent/models/trial/result.py rename to rock/sdk/bench/models/trial/result.py index 9654094271..651dffc16d 100644 --- a/rock/sdk/agent/models/trial/result.py +++ b/rock/sdk/bench/models/trial/result.py @@ -1,46 +1,35 @@ -"""Trial result models aligned with harbor.models.trial.result.""" +"""Harbor trial result models. + +TrialResult base class is in rock.sdk.job.result. +This module extends it with Harbor-specific fields. +""" from __future__ import annotations -from datetime import datetime from typing import Any from pydantic import BaseModel, Field - -class ExceptionInfo(BaseModel): - """Aligned with harbor.models.trial.result.ExceptionInfo""" - - exception_type: str = "" - exception_message: str = "" - exception_traceback: str = "" - occurred_at: str | None = None +from rock.sdk.job.result import ExceptionInfo +from rock.sdk.job.result import TrialResult as _BaseTrialResult class ModelInfo(BaseModel): - """Aligned with harbor.models.trial.result.ModelInfo""" - name: str = "" provider: str = "" class AgentInfo(BaseModel): - """Aligned with harbor.models.trial.result.AgentInfo""" - name: str = "" version: str = "" model_info: ModelInfo | None = None class VerifierResult(BaseModel): - """Aligned with harbor.models.verifier.result.VerifierResult""" - rewards: dict[str, float | int] | None = None class AgentResult(BaseModel): - """Aligned with harbor.models.agent.context.AgentContext (subset)""" - n_input_tokens: int | None = None n_cache_tokens: int | None = None n_output_tokens: int | None = None @@ -53,18 +42,14 @@ class TimingInfo(BaseModel): finished_at: str | None = None -class TrialResult(BaseModel): - """Aligned with harbor.models.trial.result.TrialResult""" +class TrialResult(_BaseTrialResult): + """Harbor TrialResult: extends base with agent/verifier/timing fields.""" - task_name: str = "" trial_name: str = "" source: str | None = None agent_info: AgentInfo = Field(default_factory=AgentInfo) agent_result: AgentResult | None = None verifier_result: VerifierResult | None = None - exception_info: ExceptionInfo | None = None - started_at: str | None = None - finished_at: str | None = None environment_setup: TimingInfo | None = None agent_setup: TimingInfo | None = None agent_execution: TimingInfo | None = None @@ -80,17 +65,6 @@ def score(self) -> float: def status(self) -> str: return "failed" if self.exception_info else "completed" - @property - def duration_sec(self) -> float: - if self.started_at and self.finished_at: - try: - start = datetime.fromisoformat(self.started_at.replace("Z", "+00:00")) - end = datetime.fromisoformat(self.finished_at.replace("Z", "+00:00")) - return (end - start).total_seconds() - except (ValueError, TypeError): - pass - return 0.0 - @property def token_ids(self) -> list[int]: if self.agent_result and self.agent_result.rollout_details: diff --git a/rock/sdk/job/__init__.py b/rock/sdk/job/__init__.py new file mode 100644 index 0000000000..0e8f8c5ba3 --- /dev/null +++ b/rock/sdk/job/__init__.py @@ -0,0 +1,34 @@ +# Pre-import rock.sdk.bench to resolve a known circular-import issue between +# rock.sdk.job.config (base JobConfig) and rock.sdk.bench.models.job.config +# (Harbor JobConfig, which inherits from the base). Doing this import first +# ensures bench is fully loaded before any rock.sdk.job submodule pulls it in. +import rock.sdk.bench # noqa: F401, I001 + +from rock.sdk.job.config import BashJobConfig, JobConfig +from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient +from rock.sdk.job.api import Job +from rock.sdk.job.operator import Operator, ScatterOperator +from rock.sdk.job.result import ExceptionInfo, JobResult, JobStatus, TrialResult +from rock.sdk.job.trial import AbstractTrial, register_trial + +# Auto-register BashTrial (safe: no bench dependency). +# HarborTrial is registered by rock.sdk.bench.__init__ to avoid a circular +# import when rock.sdk.job is triggered mid-bench-load. +import rock.sdk.job.trial.bash # noqa: F401 + +__all__ = [ + "Job", + "JobConfig", + "BashJobConfig", + "JobResult", + "JobStatus", + "TrialResult", + "ExceptionInfo", + "JobExecutor", + "JobClient", + "TrialClient", + "Operator", + "ScatterOperator", + "AbstractTrial", + "register_trial", +] diff --git a/rock/sdk/job/api.py b/rock/sdk/job/api.py new file mode 100644 index 0000000000..babc3b2723 --- /dev/null +++ b/rock/sdk/job/api.py @@ -0,0 +1,74 @@ +"""Job — thin user-facing facade over JobExecutor + Operator. + +Only 2 params (config + operator). Delegates everything to JobExecutor. + +Usage: + result = await Job(config).run() + # or + job = Job(config, operator=ScatterOperator(size=8)) + await job.submit() + result = await job.wait() +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from rock.sdk.job.executor import JobExecutor +from rock.sdk.job.operator import ScatterOperator +from rock.sdk.job.result import JobResult, JobStatus + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.executor import JobClient + from rock.sdk.job.operator import Operator + from rock.sdk.job.result import TrialResult + + +class Job: + """Job Facade — the thin user-facing entry point. + + Usage: + result = await Job(config).run() + # or + job = Job(config, operator=ScatterOperator(size=8)) + await job.submit() + result = await job.wait() + """ + + def __init__(self, config: JobConfig, operator: Operator | None = None): + self._config = config + self._executor = JobExecutor() + self._operator = operator or ScatterOperator() + self._job_client: JobClient | None = None + + async def run(self) -> JobResult: + """Full lifecycle: submit + wait.""" + await self.submit() + return await self.wait() + + async def submit(self) -> None: + """Non-blocking submit: operator generates trials, executor starts them.""" + self._job_client = await self._executor.submit(self._operator, self._config) + + async def wait(self) -> JobResult: + """Wait for completion, build JobResult.""" + if not self._job_client: + raise RuntimeError("No submitted job. Call submit() first.") + trial_results = await self._executor.wait(self._job_client) + return self._build_result(trial_results) + + async def cancel(self) -> None: + """Kill all running trials.""" + if self._job_client: + for tc in self._job_client.trials: + await tc.sandbox.arun(cmd=f"kill {tc.pid}", session=tc.session) + + def _build_result(self, trial_results: list[TrialResult]) -> JobResult: + all_success = all(r.exception_info is None for r in trial_results) + return JobResult( + job_id=self._config.job_name or "", + status=JobStatus.COMPLETED if all_success else JobStatus.FAILED, + labels=self._config.labels, + trial_results=trial_results, + ) diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py new file mode 100644 index 0000000000..69ac5be970 --- /dev/null +++ b/rock/sdk/job/config.py @@ -0,0 +1,35 @@ +"""Config hierarchy for the Job system. + +JobConfig — base config with shared fields for all job types +BashJobConfig — simple script execution + +Harbor's JobConfig lives in rock.sdk.agent.models.job.config and inherits JobConfig. +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + +from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + + +class JobConfig(BaseModel): + """Base config — shared fields for all job types.""" + + environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) + job_name: str | None = None + namespace: str | None = None + experiment_id: str | None = None + labels: dict[str, str] = Field(default_factory=dict) + auto_stop: bool = False + setup_commands: list[str] = Field(default_factory=list) + file_uploads: list[tuple[str, str]] = Field(default_factory=list) + env: dict[str, str] = Field(default_factory=dict) + timeout: int = 3600 + + +class BashJobConfig(JobConfig): + """Config for a simple bash script job.""" + + script: str | None = None + script_path: str | None = None diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py new file mode 100644 index 0000000000..011f0f0d71 --- /dev/null +++ b/rock/sdk/job/executor.py @@ -0,0 +1,143 @@ +"""JobExecutor — orchestrates the full execution of Trials produced by an Operator. + +Flow: + submit(operator, config) — apply operator to get TrialList, start all sandboxes + in parallel, return JobClient (list of TrialClient) + wait(job_client) — wait for all trials, collect results, return list[TrialResult] + run(operator, config) — submit + wait +""" + +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from rock.actions import CreateBashSessionRequest +from rock.logger import init_logger +from rock.sdk.job.operator import Operator +from rock.sdk.job.result import TrialResult +from rock.sdk.sandbox.client import Sandbox + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + +logger = init_logger(__name__) + + +@dataclass +class TrialClient: + """Handle for a single running trial.""" + + sandbox: Sandbox + session: str + pid: int + trial: AbstractTrial + + +@dataclass +class JobClient: + """Handle returned by JobExecutor.submit(). Holds multiple TrialClients.""" + + trials: list[TrialClient] + + +class JobExecutor: + """Execution engine: drives Operator to generate trials, runs in parallel, collects results.""" + + async def run(self, operator: Operator, config: JobConfig) -> list[TrialResult]: + """Full lifecycle: submit + wait.""" + job_client = await self.submit(operator, config) + return await self.wait(job_client) + + async def submit(self, operator: Operator, config: JobConfig) -> JobClient: + """Operator generates TrialList, start all sandboxes in parallel.""" + trial_list = operator.apply(config) + if not trial_list: + return JobClient(trials=[]) + trial_clients = await asyncio.gather(*[self._do_submit(t) for t in trial_list]) + return JobClient(trials=list(trial_clients)) + + async def wait(self, job_client: JobClient) -> list[TrialResult]: + """Wait for all trials, collect results in parallel.""" + if not job_client.trials: + return [] + return list(await asyncio.gather(*[self._do_wait(tc) for tc in job_client.trials])) + + # ── Internal: per-trial submit/wait ── + + @staticmethod + def _job_tmp_prefix(config: JobConfig) -> str: + """Prefix for per-job temp files on sandbox, e.g. /tmp/rock_job_my-job.""" + return f"/tmp/rock_job_{config.job_name or 'default'}" + + async def _do_submit(self, trial: AbstractTrial) -> TrialClient: + """Start sandbox + execute script for a single trial.""" + config = trial._config + sandbox = Sandbox(config.environment) + await sandbox.start() + logger.info(f"Sandbox started: sandbox_id={sandbox.sandbox_id}, job_name={config.job_name}") + + session = f"rock-job-{config.job_name or 'default'}" + env = self._build_session_env(config) + await sandbox.create_session(CreateBashSessionRequest(session=session, env_enable=True, env=env)) + + await trial.setup(sandbox) + script_content = trial.build() + + prefix = self._job_tmp_prefix(config) + script_path = f"{prefix}.sh" + await sandbox.write_file_by_path(script_content, script_path) + + tmp_file = f"{prefix}.out" + pid, error = await sandbox.start_nohup_process( + cmd=f"bash {script_path}", + tmp_file=tmp_file, + session=session, + ) + if error is not None: + raise RuntimeError(f"Failed to start trial: {error.output}") + + logger.info(f"Trial started: pid={pid}, job_name={config.job_name}") + return TrialClient(sandbox=sandbox, session=session, pid=pid, trial=trial) + + async def _do_wait(self, client: TrialClient) -> TrialResult: + """Wait for a single trial to finish, call trial.collect().""" + from rock.sdk.job.result import ExceptionInfo + + config = client.trial._config + try: + success, message = await client.sandbox.wait_for_process_completion( + pid=client.pid, + session=client.session, + wait_timeout=config.timeout, + wait_interval=30, + ) + obs = await client.sandbox.handle_nohup_output( + tmp_file=f"{self._job_tmp_prefix(config)}.out", + session=client.session, + success=success, + message=message, + ignore_output=False, + response_limited_bytes_in_nohup=None, + ) + exit_code = obs.exit_code if obs.exit_code is not None else 1 + result = await client.trial.collect(client.sandbox, obs.output or "", exit_code) + if not success and result.exception_info is None: + result.exception_info = ExceptionInfo( + exception_type="ProcessTimeout", + exception_message=message or "process did not complete successfully", + ) + return result + finally: + if config.auto_stop: + await client.sandbox.close() + + @staticmethod + def _build_session_env(config: JobConfig) -> dict[str, str] | None: + """Merge OSS_* env vars from the process with config.env (config wins).""" + oss_env = {k: v for k, v in os.environ.items() if k.startswith("OSS")} + merged = {**oss_env, **config.env} + return merged or None diff --git a/rock/sdk/job/operator.py b/rock/sdk/job/operator.py new file mode 100644 index 0000000000..15699ff413 --- /dev/null +++ b/rock/sdk/job/operator.py @@ -0,0 +1,46 @@ +"""Operator — generic algorithm that produces a TrialList from a JobConfig.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from rock.sdk.job.trial.registry import _create_trial + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + + +class Operator(ABC): + """Operator base: apply(config) -> list[AbstractTrial]. + + Operators generate a TrialList from a config. They don't manage + sandbox lifecycle (JobExecutor does) — just decide what to run. + """ + + @abstractmethod + def apply(self, config: JobConfig) -> list[AbstractTrial]: + """Generate a TrialList from config. Empty list means no-op.""" + ... + + +class ScatterOperator(Operator): + """Scatter: create `size` identical Trial instances from config. + + Analog of torch.distributed.scatter — same data/config distributed to N workers. + + Usage: + ScatterOperator() # size=1, single trial (default) + ScatterOperator(size=8) # 8 parallel trials + ScatterOperator(size=0) # empty list, no-op + """ + + def __init__(self, size: int = 1): + self.size = size + + def apply(self, config: JobConfig) -> list[AbstractTrial]: + if self.size <= 0: + return [] + trial = _create_trial(config) + return [trial] * self.size diff --git a/rock/sdk/job/result.py b/rock/sdk/job/result.py new file mode 100644 index 0000000000..b13adef858 --- /dev/null +++ b/rock/sdk/job/result.py @@ -0,0 +1,106 @@ +"""Result models for the Job system. + +Base classes: TrialResult, JobStatus, JobResult[T]. +Harbor-specific subclasses in rock.sdk.agent.models.trial.result. +""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Generic, TypeVar + +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# TrialResult base — 通用字段 +# --------------------------------------------------------------------------- + + +class ExceptionInfo(BaseModel): + """通用异常信息""" + + exception_type: str = "" + exception_message: str = "" + exception_traceback: str = "" + occurred_at: str | None = None + + +class TrialResult(BaseModel): + """单次执行结果的基类 — 通用字段 + + Harbor 的 TrialResult 继承此类，添加 agent_info, verifier_result 等字段。 + 子类可 override score 和 status properties。 + """ + + task_name: str = "" + exception_info: ExceptionInfo | None = None + started_at: str | None = None + finished_at: str | None = None + + @property + def score(self) -> float: + return 0.0 + + @property + def status(self) -> str: + return "failed" if self.exception_info else "completed" + + @property + def duration_sec(self) -> float: + if self.started_at and self.finished_at: + try: + start = datetime.fromisoformat(self.started_at.replace("Z", "+00:00")) + end = datetime.fromisoformat(self.finished_at.replace("Z", "+00:00")) + return (end - start).total_seconds() + except (ValueError, TypeError): + pass + return 0.0 + + +# --------------------------------------------------------------------------- +# JobStatus + JobResult[T] +# --------------------------------------------------------------------------- + + +class JobStatus(str, Enum): + """Job status enum.""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +T = TypeVar("T", bound=TrialResult) + + +class JobResult(BaseModel, Generic[T]): + """Aggregated result of a complete job run. + + Generic over trial result type T: + - JobResult[TrialResult] — base (new Job system) + - JobResult[HarborTrialResult] — Harbor agent system + """ + + job_id: str = "" + status: JobStatus = JobStatus.COMPLETED + labels: dict[str, str] = Field(default_factory=dict) + trial_results: list[T] = Field(default_factory=list) + raw_output: str = "" + exit_code: int = 0 + + @property + def score(self) -> float: + if not self.trial_results: + return 0.0 + return sum(t.score for t in self.trial_results) / len(self.trial_results) + + @property + def n_completed(self) -> int: + return sum(1 for t in self.trial_results if t.status == "completed") + + @property + def n_failed(self) -> int: + return sum(1 for t in self.trial_results if t.status == "failed") diff --git a/rock/sdk/job/trial/__init__.py b/rock/sdk/job/trial/__init__.py new file mode 100644 index 0000000000..d8080d2c7d --- /dev/null +++ b/rock/sdk/job/trial/__init__.py @@ -0,0 +1,4 @@ +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import _create_trial, register_trial + +__all__ = ["AbstractTrial", "register_trial", "_create_trial"] diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py new file mode 100644 index 0000000000..fc5f4e092f --- /dev/null +++ b/rock/sdk/job/trial/abstract.py @@ -0,0 +1,43 @@ +"""Trial abstract base class — three-phase interface (setup / build / collect). + +Trial 对象不管理 sandbox 生命周期；生命周期由 JobExecutor 负责。 +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.result import TrialResult + from rock.sdk.sandbox.client import Sandbox + + +class AbstractTrial(ABC): + """Trial base: three-phase interface (setup/build/collect). + + Trial 不管理 sandbox 生命周期 (由 JobExecutor 负责)。 + """ + + def __init__(self, config: JobConfig): + self._config = config + + @abstractmethod + async def setup(self, sandbox: Sandbox) -> None: + """Pre-execution: prepare sandbox environment (upload files, write configs).""" + + @abstractmethod + def build(self) -> str: + """Build: generate bash script to execute.""" + + @abstractmethod + async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: + """Post-execution: collect and parse results.""" + + async def _upload_files(self, sandbox: Sandbox) -> None: + """Shared helper: upload all entries in ``config.file_uploads``.""" + for local_path, sandbox_path in self._config.file_uploads: + obs = await sandbox.fs.upload_dir(source_dir=local_path, target_dir=sandbox_path) + if obs.exit_code != 0: + raise RuntimeError(f"Failed to upload {local_path} -> {sandbox_path}: {obs.failure_reason}") diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py new file mode 100644 index 0000000000..ea7a040c34 --- /dev/null +++ b/rock/sdk/job/trial/bash.py @@ -0,0 +1,49 @@ +"""BashTrial — execute a bash script inside a sandbox.""" + +from __future__ import annotations + +from pathlib import Path + +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.result import ExceptionInfo, TrialResult +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import register_trial + + +class BashTrial(AbstractTrial): + """Bash script execution trial.""" + + _config: BashJobConfig + + async def setup(self, sandbox) -> None: + await self._upload_files(sandbox) + # If script_path is set, read content into self._config.script + if self._config.script_path: + self._config.script = Path(self._config.script_path).read_text() + + def build(self) -> str: + lines = ["#!/bin/bash", "set -e", ""] + if self._config.setup_commands: + for cmd in self._config.setup_commands: + lines.append(f"echo '>>> {cmd[:60]}...'") + lines.append(cmd) + lines.append("") + if self._config.script: + lines.append(self._config.script) + return "\n".join(lines) + + async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: + exception_info = None + if exit_code != 0: + exception_info = ExceptionInfo( + exception_type="BashExitCode", + exception_message=f"Bash script exited with code {exit_code}", + ) + return TrialResult( + task_name=self._config.job_name or "", + exception_info=exception_info, + ) + + +# Auto-register on import +register_trial(BashJobConfig, BashTrial) diff --git a/rock/sdk/job/trial/harbor.py b/rock/sdk/job/trial/harbor.py new file mode 100644 index 0000000000..2fb242564c --- /dev/null +++ b/rock/sdk/job/trial/harbor.py @@ -0,0 +1,116 @@ +"""HarborTrial — execute a Harbor benchmark job inside a sandbox. + +Extracted from rock.sdk.bench.job.Job. Combines dockerd startup, setup +commands, and ``harbor jobs start -c`` into a single bash script executed +by the JobExecutor via the sandbox nohup protocol. +""" + +from __future__ import annotations + +import json + +from rock.actions import Command, ReadFileRequest +from rock.logger import init_logger +from rock.sdk.bench.constants import USER_DEFINED_LOGS +from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig +from rock.sdk.bench.models.trial.result import TrialResult +from rock.sdk.job.result import ExceptionInfo +from rock.sdk.job.result import TrialResult as BaseTrialResult +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import register_trial + +logger = init_logger(__name__) + +_HARBOR_SCRIPT_TEMPLATE = r"""#!/bin/bash +set -e + +# ── Detect and start dockerd ───────────────────────────────────────── +if command -v docker &>/dev/null; then + echo "docker OK: $(command -v docker)" + if ! pgrep -x dockerd &>/dev/null; then + echo "Starting dockerd..." + nohup dockerd &>/var/log/dockerd.log & + fi + for i in $(seq 1 60); do + if docker info &>/dev/null; then echo "dockerd is ready"; break; fi + sleep 1 + if [ "$i" -eq 60 ]; then echo "WARN: dockerd failed to start within 60s"; fi + done +fi + +# ── Ensure output directory exists ────────────────────────────────── +mkdir -p {user_defined_dir} + +# ── Setup commands ─────────────────────────────────────────────────── +{setup_commands} + +# ── Harbor run ─────────────────────────────────────────────────────── +harbor jobs start -c {config_path} +""" + + +class HarborTrial(AbstractTrial): + """Harbor benchmark trial execution.""" + + _config: HarborJobConfig + + async def setup(self, sandbox) -> None: + await self._upload_files(sandbox) + # Write Harbor YAML config to sandbox + yaml_content = self._config.to_harbor_yaml() + config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" + await sandbox.write_file_by_path(yaml_content, config_path) + + def build(self) -> str: + setup_lines: list[str] = [] + for cmd in self._config.setup_commands: + setup_lines.append(f"echo '>>> {cmd[:60]}...'") + setup_lines.append(cmd) + setup_block = "\n".join(setup_lines) if setup_lines else "echo 'No setup commands'" + + config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" + return _HARBOR_SCRIPT_TEMPLATE.format( + setup_commands=setup_block, + config_path=config_path, + user_defined_dir=USER_DEFINED_LOGS, + ) + + async def collect(self, sandbox, output: str, exit_code: int) -> BaseTrialResult: + trial_results = await self._collect_trial_results(sandbox) + if trial_results: + return trial_results[0] + + exception_info = ExceptionInfo( + exception_type="HarborNoTrials", + exception_message="No trial results found", + ) + return BaseTrialResult( + task_name=self._config.job_name or "", + exception_info=exception_info, + ) + + async def _collect_trial_results(self, sandbox) -> list[TrialResult]: + """Read trial-level result.json files from sandbox.""" + job_dir = f"{self._config.jobs_dir}/{self._config.job_name}" + try: + list_result = await sandbox.execute( + Command(command=["find", job_dir, "-mindepth", "2", "-maxdepth", "2", "-name", "result.json"]) + ) + trial_files = [line.strip() for line in (list_result.stdout or "").strip().split("\n") if line.strip()] + except Exception: + trial_files = [] + + results: list[TrialResult] = [] + for trial_file in trial_files: + try: + response = await sandbox.read_file(ReadFileRequest(path=trial_file)) + data = json.loads(response.content) + results.append(TrialResult.from_harbor_json(data)) + except Exception as e: + logger.warning(f"Failed to parse trial result {trial_file}: {e}") + + return results + + +# Auto-register +register_trial(HarborJobConfig, HarborTrial) diff --git a/rock/sdk/job/trial/registry.py b/rock/sdk/job/trial/registry.py new file mode 100644 index 0000000000..9618de533c --- /dev/null +++ b/rock/sdk/job/trial/registry.py @@ -0,0 +1,29 @@ +"""Trial registry — maps JobConfig subclasses to their AbstractTrial implementations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from rock.sdk.job.config import JobConfig + from rock.sdk.job.trial.abstract import AbstractTrial + +_TRIAL_REGISTRY: dict[type[JobConfig], type[AbstractTrial]] = {} + + +def register_trial(config_type: type[JobConfig], trial_type: type[AbstractTrial]) -> None: + """Register a Config → Trial mapping.""" + _TRIAL_REGISTRY[config_type] = trial_type + + +def _create_trial(config: JobConfig) -> AbstractTrial: + """Create a Trial instance for the given config. + + Raises TypeError if no trial class has been registered for this config type. + """ + trial_cls = _TRIAL_REGISTRY.get(type(config)) + if trial_cls is None: + raise TypeError( + f"No trial registered for {type(config).__name__}. Supported: {[c.__name__ for c in _TRIAL_REGISTRY]}" + ) + return trial_cls(config) diff --git a/tests/unit/admin/core/test_ray_service.py b/tests/unit/admin/core/test_ray_service.py index d1f0ba8aa6..0a9602d3b7 100644 --- a/tests/unit/admin/core/test_ray_service.py +++ b/tests/unit/admin/core/test_ray_service.py @@ -83,13 +83,31 @@ async def test_reconnect_ray_skip_when_reader_exists_and_write_lock_timeout(ray_ @pytest.mark.need_ray @pytest.mark.asyncio async def test_ray_get(ray_service): + import uuid + + import ray + service = ray_service + # Unique name to avoid colliding with leaked detached actors from prior runs/reruns + actor_name = f"test-ray-get-{uuid.uuid4().hex[:8]}" + namespace = "rock-sandbox-test" config = RayDeploymentConfig(image="python:3.11") deployment: RayDeployment = RayDeployment.from_config(config) - actor = SandboxActor.options(**{"name": "test", "lifetime": "detached"}).remote(config, deployment) - await service.async_ray_get(actor.start.remote()) - result = await service.async_ray_get(actor.host_name.remote()) - assert result is not None - actor = await service.async_ray_get_actor("test", "rock-sandbox-test") - assert actor is not None - await service.async_ray_get(actor.stop.remote()) + + actor = SandboxActor.options(name=actor_name, namespace=namespace, lifetime="detached").remote(config, deployment) + try: + await service.async_ray_get(actor.start.remote()) + result = await service.async_ray_get(actor.host_name.remote()) + assert result is not None + + fetched_actor = await service.async_ray_get_actor(actor_name, namespace) + assert fetched_actor is not None + + await service.async_ray_get(fetched_actor.stop.remote()) + finally: + # Ensure detached actor is always killed, even if assertions/RPCs above fail + try: + leaked = ray.get_actor(actor_name, namespace=namespace) + ray.kill(leaked) + except Exception: + pass diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index 187b826510..cf3bb83b68 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -3,16 +3,16 @@ from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch -from rock.sdk.agent.job import Job, JobResult, JobStatus -from rock.sdk.agent.models.job.config import ( +from rock.sdk.bench.job import Job, JobResult, JobStatus +from rock.sdk.bench.models.job.config import ( JobConfig, LocalDatasetConfig, RegistryDatasetConfig, RemoteRegistryInfo, RockEnvironmentConfig, ) -from rock.sdk.agent.models.trial.config import AgentConfig -from rock.sdk.agent.models.trial.result import ExceptionInfo, TrialResult, VerifierResult +from rock.sdk.bench.models.trial.config import AgentConfig +from rock.sdk.bench.models.trial.result import ExceptionInfo, TrialResult, VerifierResult class TestJobStatus: diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index dccbef8d7f..dea3092b19 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -2,15 +2,15 @@ import yaml -from rock.sdk.agent.models.job.config import ( +from rock.sdk.bench.models.job.config import ( JobConfig, LocalDatasetConfig, RegistryDatasetConfig, RemoteRegistryInfo, RockEnvironmentConfig, ) -from rock.sdk.agent.models.metric.config import MetricConfig -from rock.sdk.agent.models.trial.config import AgentConfig, TaskConfig +from rock.sdk.bench.models.metric.config import MetricConfig +from rock.sdk.bench.models.trial.config import AgentConfig, TaskConfig class TestRockEnvironmentConfigInheritance: @@ -111,7 +111,9 @@ def test_basic_serialization(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - assert data["job_name"] == "test-job" + # Base fields (job_name, experiment_id, etc.) are excluded from harbor YAML + assert "job_name" not in data + assert "experiment_id" not in data assert data["n_attempts"] == 2 assert data["agents"][0]["name"] == "terminus-2" @@ -152,7 +154,8 @@ def test_excludes_none_values(self): assert "agent_timeout_multiplier" not in data - def test_labels_serialized(self): + def test_labels_excluded_as_base_field(self): + """labels is a base JobConfig field, so it's excluded from harbor YAML.""" cfg = JobConfig( job_name="labeled-job", experiment_id="test-exp", @@ -161,14 +164,8 @@ def test_labels_serialized(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - assert data["labels"] == {"step": "42", "env": "prod"} - - def test_empty_labels_not_excluded(self): - cfg = JobConfig(job_name="no-labels", experiment_id="test-exp") - yaml_str = cfg.to_harbor_yaml() - data = yaml.safe_load(yaml_str) - - assert data["labels"] == {} + assert "labels" not in data + assert "job_name" not in data def test_path_fields_serialized_as_strings(self): cfg = JobConfig( @@ -209,7 +206,7 @@ def test_harbor_env_fields_serialized(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - assert data["job_name"] == "full-test" + assert "job_name" not in data # base field excluded assert data["environment"]["type"] == "docker" assert data["environment"]["force_build"] is True assert data["environment"]["override_cpus"] == 4 diff --git a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py index a8d6ec0899..f6e48bdbd5 100644 --- a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py +++ b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py @@ -5,9 +5,9 @@ import pytest from pydantic import ValidationError -from rock.sdk.agent.job import Job -from rock.sdk.agent.models.job.config import JobConfig -from rock.sdk.agent.models.trial.config import RockEnvironmentConfig +from rock.sdk.bench.job import Job +from rock.sdk.bench.models.job.config import JobConfig +from rock.sdk.bench.models.trial.config import RockEnvironmentConfig class TestExperimentIdNotEmpty: diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index 34ffd42894..84e606cc0a 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -1,7 +1,7 @@ from pathlib import Path -from rock.sdk.agent.models.environment_type import EnvironmentType -from rock.sdk.agent.models.job.config import ( +from rock.sdk.bench.models.environment_type import EnvironmentType +from rock.sdk.bench.models.job.config import ( JobConfig, LocalDatasetConfig, OrchestratorConfig, @@ -10,16 +10,16 @@ RetryConfig, RockEnvironmentConfig, ) -from rock.sdk.agent.models.metric.config import MetricConfig -from rock.sdk.agent.models.metric.type import MetricType -from rock.sdk.agent.models.orchestrator_type import OrchestratorType -from rock.sdk.agent.models.trial.config import ( +from rock.sdk.bench.models.metric.config import MetricConfig +from rock.sdk.bench.models.metric.type import MetricType +from rock.sdk.bench.models.orchestrator_type import OrchestratorType +from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, TaskConfig, VerifierConfig, ) -from rock.sdk.agent.models.trial.config import ( +from rock.sdk.bench.models.trial.config import ( EnvironmentConfig as HarborEnvironmentConfig, ) @@ -257,7 +257,7 @@ def test_with_full_config(self): class TestPublicAPI: def test_import_from_agent_package(self): - from rock.sdk.agent import Job, JobResult, JobStatus, TrialResult + from rock.sdk.bench import Job, JobResult, JobStatus, TrialResult assert Job is not None assert JobResult is not None @@ -265,7 +265,7 @@ def test_import_from_agent_package(self): assert TrialResult is not None def test_import_from_models_package(self): - from rock.sdk.agent.models import ( + from rock.sdk.bench.models import ( AgentConfig, EnvironmentType, JobConfig, diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py index 1fe838371a..8fee1e2484 100644 --- a/tests/unit/sdk/agent/test_oss_mirror.py +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -11,7 +11,7 @@ import yaml -from rock.sdk.agent.models.trial.config import EnvironmentConfig +from rock.sdk.bench.models.trial.config import EnvironmentConfig # --------------------------------------------------------------------------- # 1. OssMirrorConfig 模型 @@ -20,17 +20,17 @@ class TestOssMirrorConfig: def test_importable_from_trial_config(self): - from rock.sdk.agent.models.trial.config import OssMirrorConfig + from rock.sdk.bench.models.trial.config import OssMirrorConfig assert OssMirrorConfig is not None def test_importable_from_agent_package(self): - from rock.sdk.agent import OssMirrorConfig + from rock.sdk.bench import OssMirrorConfig assert OssMirrorConfig is not None def test_default_is_disabled(self): - from rock.sdk.agent.models.trial.config import OssMirrorConfig + from rock.sdk.bench.models.trial.config import OssMirrorConfig cfg = OssMirrorConfig() assert cfg.enabled is False @@ -41,7 +41,7 @@ def test_default_is_disabled(self): assert cfg.oss_endpoint is None def test_all_fields_settable(self): - from rock.sdk.agent.models.trial.config import OssMirrorConfig + from rock.sdk.bench.models.trial.config import OssMirrorConfig cfg = OssMirrorConfig( enabled=True, @@ -70,7 +70,7 @@ def test_default_oss_mirror_is_none(self): assert env.oss_mirror is None def test_set_oss_mirror(self): - from rock.sdk.agent.models.trial.config import OssMirrorConfig + from rock.sdk.bench.models.trial.config import OssMirrorConfig mirror = OssMirrorConfig(enabled=True, oss_bucket="b1", oss_region="r1") env = EnvironmentConfig(oss_mirror=mirror) @@ -95,13 +95,13 @@ def test_set_oss_mirror_from_dict(self): class TestJobConfigNamespaceFields: def test_default_namespace_is_none(self): - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig(job_name="test", experiment_id="test-exp") assert cfg.namespace is None def test_namespace_settable_at_top_level(self): - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig(job_name="test", namespace="team-rl", experiment_id="rl-step-42") assert cfg.namespace == "team-rl" @@ -116,8 +116,8 @@ def test_namespace_settable_at_top_level(self): class TestToHarborYamlOssMirror: def test_namespace_at_top_level_in_yaml(self): """namespace/experiment_id 序列化为 JobConfig 顶层字段。""" - from rock.sdk.agent.models.job.config import JobConfig - from rock.sdk.agent.models.trial.config import OssMirrorConfig, RockEnvironmentConfig + from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.trial.config import OssMirrorConfig, RockEnvironmentConfig cfg = JobConfig( job_name="mirror-test", @@ -136,8 +136,9 @@ def test_namespace_at_top_level_in_yaml(self): ) data = yaml.safe_load(cfg.to_harbor_yaml()) - assert data["namespace"] == "my-ns" - assert data["experiment_id"] == "exp-1" + # namespace/experiment_id are base JobConfig fields, excluded from harbor YAML + assert "namespace" not in data + assert "experiment_id" not in data oss = data["environment"]["oss_mirror"] assert oss["enabled"] is True assert oss["oss_bucket"] == "test-bucket" @@ -146,7 +147,7 @@ def test_namespace_at_top_level_in_yaml(self): def test_disabled_oss_mirror_excluded_from_yaml(self): """When oss_mirror is default (disabled), it should not clutter the YAML.""" - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig(job_name="no-mirror", experiment_id="test-exp") data = yaml.safe_load(cfg.to_harbor_yaml()) @@ -163,7 +164,7 @@ def test_disabled_oss_mirror_excluded_from_yaml(self): class TestFromYamlOssMirror: def test_from_yaml_with_top_level_namespace(self, tmp_path): """新方式：namespace/experiment_id 在顶层。""" - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig yaml_content = """\ job_name: loaded-mirror @@ -189,7 +190,7 @@ def test_from_yaml_with_top_level_namespace(self, tmp_path): def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): """YAML 中 oss_mirror 内多余的 namespace 等字段由 Pydantic 忽略。""" - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig yaml_content = """\ job_name: compat-mirror @@ -216,7 +217,7 @@ def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): assert "experiment_id" not in dump def test_from_yaml_without_oss_mirror(self, tmp_path): - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig yaml_content = """\ job_name: no-mirror @@ -238,7 +239,7 @@ def test_from_yaml_without_oss_mirror(self, tmp_path): class TestEnableOssMirror: def test_enable_with_all_params(self): - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig(job_name="conv-test", experiment_id="test-exp") cfg.enable_oss_mirror( @@ -253,7 +254,7 @@ def test_enable_with_all_params(self): def test_does_not_touch_namespace_or_experiment_id(self): """enable_oss_mirror 不修改顶层 namespace / experiment_id。""" - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig( job_name="no-touch-test", @@ -272,7 +273,7 @@ def test_does_not_touch_namespace_or_experiment_id(self): def test_enable_then_serialize_roundtrip(self): """to_harbor_yaml: 顶层 namespace / experiment_id 与 oss_mirror 独立设置。""" - from rock.sdk.agent.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import JobConfig cfg = JobConfig(job_name="roundtrip", namespace="rt-ns", experiment_id="rt-exp") cfg.enable_oss_mirror( @@ -283,10 +284,9 @@ def test_enable_then_serialize_roundtrip(self): oss_endpoint="oss-ap-southeast-1.aliyuncs.com", ) data = yaml.safe_load(cfg.to_harbor_yaml()) - assert data["namespace"] == "rt-ns" - assert data["experiment_id"] == "rt-exp" + # namespace/experiment_id are base fields, excluded from harbor YAML + assert "namespace" not in data + assert "experiment_id" not in data oss = data["environment"]["oss_mirror"] assert oss["enabled"] is True assert oss["oss_bucket"] == "rt-bucket" - assert "namespace" not in oss - assert "experiment_id" not in oss diff --git a/tests/unit/sdk/job/__init__.py b/tests/unit/sdk/job/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py new file mode 100644 index 0000000000..016bcbcd82 --- /dev/null +++ b/tests/unit/sdk/job/test_cli_job.py @@ -0,0 +1,283 @@ +"""Tests for rock.cli.command.job — JobCommand with --type bash/harbor routing.""" + +from __future__ import annotations + +import argparse +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import rock.sdk.bench # pre-import to avoid circular # noqa: F401 +from rock.cli.command.job import JobCommand + + +async def _build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser() + sub = p.add_subparsers(dest="main_command") + await JobCommand.add_parser_to(sub) + return p + + +def _make_mock_job_result(): + mock_result = MagicMock() + mock_result.trial_results = [] + mock_result.status = "completed" + return mock_result + + +# ---------------------------------------------------------------------------- +# Parser tests +# ---------------------------------------------------------------------------- + + +async def test_parser_default_type_is_bash(): + p = await _build_parser() + args = p.parse_args(["job", "run", "--script-content", "echo hi"]) + assert args.type == "bash" + assert args.script_content == "echo hi" + + +async def test_parser_accepts_type_bash_explicit(): + p = await _build_parser() + args = p.parse_args(["job", "run", "--type", "bash", "--script-content", "echo hi"]) + assert args.type == "bash" + + +async def test_parser_accepts_type_harbor(): + p = await _build_parser() + args = p.parse_args(["job", "run", "--type", "harbor", "--config", "/tmp/c.yaml"]) + assert args.type == "harbor" + assert args.config == "/tmp/c.yaml" + + +async def test_parser_supports_all_bash_args(): + p = await _build_parser() + args = p.parse_args( + [ + "job", + "run", + "--script", + "/tmp/s.sh", + "--image", + "python:3.11", + "--memory", + "4g", + "--cpus", + "2", + "--timeout", + "600", + "--local-path", + "/tmp/local", + "--target-path", + "/root/other", + ] + ) + assert args.script == "/tmp/s.sh" + assert args.image == "python:3.11" + assert args.memory == "4g" + assert args.cpus == 2.0 + assert args.timeout == 600 + assert args.local_path == "/tmp/local" + assert args.target_path == "/root/other" + + +async def test_parser_invalid_type_rejected(): + p = await _build_parser() + with pytest.raises(SystemExit): + p.parse_args(["job", "run", "--type", "invalid"]) + + +# ---------------------------------------------------------------------------- +# arun / _job_run behavior tests +# ---------------------------------------------------------------------------- + + +def _bash_args(**overrides): + defaults = dict( + job_command="run", + type="bash", + script=None, + script_content=None, + image=None, + memory=None, + cpus=None, + local_path=None, + target_path="/root/job", + timeout=3600, + config=None, + base_url=None, + cluster=None, + extra_headers=None, + ) + defaults.update(overrides) + return argparse.Namespace(**defaults) + + +async def test_bash_creates_bash_job_config_and_runs(): + from rock.sdk.job.config import BashJobConfig + + args = _bash_args( + script_content="echo hello", + image="python:3.11", + memory="4g", + cpus=2.0, + timeout=600, + ) + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_called_once() + config_arg = MockJob.call_args[0][0] + assert isinstance(config_arg, BashJobConfig) + assert config_arg.script == "echo hello" + assert config_arg.script_path is None + assert config_arg.environment.image == "python:3.11" + assert config_arg.environment.memory == "4g" + assert config_arg.environment.cpus == 2.0 + assert config_arg.timeout == 600 + assert config_arg.auto_stop is True + mock_instance.run.assert_awaited_once() + + +async def test_bash_with_script_path(): + from rock.sdk.job.config import BashJobConfig + + args = _bash_args(script="/tmp/my_script.sh") + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert isinstance(config_arg, BashJobConfig) + assert config_arg.script_path == "/tmp/my_script.sh" + assert config_arg.script is None + + +async def test_bash_with_file_upload(): + args = _bash_args( + script_content="echo hi", + local_path="/tmp/src", + target_path="/root/target", + ) + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.file_uploads == [("/tmp/src", "/root/target")] + + +async def test_bash_requires_script_or_script_content(): + args = _bash_args() # neither set + + with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_not_called() + mock_logger.error.assert_called() + + +async def test_bash_rejects_both_script_and_script_content(): + args = _bash_args(script="/tmp/s.sh", script_content="echo hi") + + with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_not_called() + mock_logger.error.assert_called() + + +async def test_harbor_requires_config(): + args = _bash_args(type="harbor", config=None) + + with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_not_called() + mock_logger.error.assert_called() + + +async def test_harbor_loads_from_yaml(): + from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + + yaml_content = """ +experiment_id: exp-123 +job_name: my-harbor-job +""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + yaml_path = f.name + + try: + args = _bash_args(type="harbor", config=yaml_path) + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + MockJob.assert_called_once() + config_arg = MockJob.call_args[0][0] + assert isinstance(config_arg, HarborJobConfig) + assert config_arg.experiment_id == "exp-123" + assert config_arg.auto_stop is True + finally: + Path(yaml_path).unlink(missing_ok=True) + + +async def test_harbor_image_override(): + yaml_content = """ +experiment_id: exp-abc +""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + yaml_path = f.name + + try: + args = _bash_args(type="harbor", config=yaml_path, image="custom:tag") + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.environment.image == "custom:tag" + finally: + Path(yaml_path).unlink(missing_ok=True) + + +async def test_unknown_job_command_logs_error(): + args = argparse.Namespace(job_command="weird") + + with patch("rock.cli.command.job.logger") as mock_logger: + cmd = JobCommand() + await cmd.arun(args) + mock_logger.error.assert_called() diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py new file mode 100644 index 0000000000..3a05f5d3d0 --- /dev/null +++ b/tests/unit/sdk/job/test_config.py @@ -0,0 +1,309 @@ +"""Tests for rock.sdk.job.config — JobConfig, BashJobConfig, HarborJobConfig.""" + +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest +import yaml + +from rock.sdk.bench.constants import USER_DEFINED_LOGS +from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig +from rock.sdk.bench.models.trial.config import ( + AgentConfig, + ArtifactConfig, + RockEnvironmentConfig, + TaskConfig, + VerifierConfig, +) +from rock.sdk.job.config import BashJobConfig, JobConfig + +# --------------------------------------------------------------------------- +# JobConfig (base) +# --------------------------------------------------------------------------- + + +class TestJobConfig: + def test_defaults(self): + cfg = JobConfig() + assert isinstance(cfg.environment, RockEnvironmentConfig) + assert cfg.job_name is None + assert cfg.namespace is None + assert cfg.experiment_id is None + assert cfg.labels == {} + assert cfg.auto_stop is False + assert cfg.setup_commands == [] + assert cfg.file_uploads == [] + assert cfg.env == {} + assert cfg.timeout == 3600 + + def test_custom_values(self): + env = RockEnvironmentConfig(image="ubuntu:22.04") + cfg = JobConfig( + environment=env, + job_name="my-job", + namespace="team-a", + experiment_id="exp-001", + labels={"step": "42"}, + auto_stop=True, + setup_commands=["pip install foo"], + file_uploads=[("/local/file.py", "/sandbox/file.py")], + env={"MY_VAR": "hello"}, + timeout=7200, + ) + assert cfg.environment.image == "ubuntu:22.04" + assert cfg.job_name == "my-job" + assert cfg.namespace == "team-a" + assert cfg.experiment_id == "exp-001" + assert cfg.labels == {"step": "42"} + assert cfg.auto_stop is True + assert cfg.setup_commands == ["pip install foo"] + assert cfg.file_uploads == [("/local/file.py", "/sandbox/file.py")] + assert cfg.env == {"MY_VAR": "hello"} + assert cfg.timeout == 7200 + + def test_is_base_model(self): + """JobConfig is a Pydantic BaseModel.""" + from pydantic import BaseModel + + assert issubclass(JobConfig, BaseModel) + + +# --------------------------------------------------------------------------- +# BashJobConfig +# --------------------------------------------------------------------------- + + +class TestBashJobConfig: + def test_inherits_job_config(self): + assert issubclass(BashJobConfig, JobConfig) + + def test_defaults(self): + cfg = BashJobConfig() + # Inherited defaults + assert cfg.timeout == 3600 + assert cfg.labels == {} + # Own defaults + assert cfg.script is None + assert cfg.script_path is None + + def test_script_field(self): + cfg = BashJobConfig(script="echo hello") + assert cfg.script == "echo hello" + assert cfg.script_path is None + + def test_script_path_field(self): + cfg = BashJobConfig(script_path="/path/to/run.sh") + assert cfg.script_path == "/path/to/run.sh" + assert cfg.script is None + + def test_inherits_base_fields(self): + cfg = BashJobConfig( + job_name="bash-job", + namespace="ns", + timeout=600, + script="ls -la", + ) + assert cfg.job_name == "bash-job" + assert cfg.namespace == "ns" + assert cfg.timeout == 600 + assert cfg.script == "ls -la" + + +# --------------------------------------------------------------------------- +# HarborJobConfig +# --------------------------------------------------------------------------- + + +class TestHarborJobConfig: + def test_inherits_job_config(self): + assert issubclass(HarborJobConfig, JobConfig) + + def test_defaults(self): + cfg = HarborJobConfig(experiment_id="test-exp") + # Inherited + assert cfg.timeout == 3600 + assert cfg.labels == {} + assert cfg.job_name is None + # Own defaults + assert len(cfg.agents) == 1 + assert isinstance(cfg.agents[0], AgentConfig) + assert cfg.datasets == [] + from rock.sdk.bench.models.job.config import OrchestratorConfig + + assert isinstance(cfg.orchestrator, OrchestratorConfig) + assert isinstance(cfg.verifier, VerifierConfig) + assert cfg.tasks == [] + assert cfg.metrics == [] + assert cfg.artifacts == [] + assert cfg.n_attempts == 1 + assert cfg.timeout_multiplier == 1.0 + assert cfg.agent_timeout_multiplier is None + assert cfg.verifier_timeout_multiplier is None + assert cfg.jobs_dir == Path(USER_DEFINED_LOGS) / "jobs" + assert cfg.debug is False + + def test_custom_harbor_fields(self): + agent = AgentConfig(name="my-agent", import_path="my_module:MyAgent") + task = TaskConfig(path=Path("/tasks/task1.json")) + artifact = ArtifactConfig(source="/data/output") + cfg = HarborJobConfig( + experiment_id="test-exp", + agents=[agent], + tasks=[task], + artifacts=[artifact, "/data/logs"], + n_attempts=3, + timeout_multiplier=2.0, + debug=True, + ) + assert cfg.agents == [agent] + assert cfg.tasks == [task] + assert len(cfg.artifacts) == 2 + assert cfg.n_attempts == 3 + assert cfg.timeout_multiplier == 2.0 + assert cfg.debug is True + + +# --------------------------------------------------------------------------- +# HarborJobConfig.to_harbor_yaml +# --------------------------------------------------------------------------- + + +class TestHarborJobConfigToHarborYaml: + def test_excludes_rock_fields(self): + """Rock-level fields (job_name, namespace, etc.) must NOT appear in Harbor YAML. + + Note: 'environment' is excluded from _ROCK_FIELDS dump, but harbor + environment fields are re-injected via to_harbor_environment(), so + the 'environment' key *may* appear with harbor-native fields only. + """ + cfg = HarborJobConfig( + job_name="should-not-appear", + namespace="should-not-appear", + experiment_id="should-not-appear", + labels={"step": "1"}, + auto_stop=True, + setup_commands=["pip install foo"], + file_uploads=[("/a", "/b")], + env={"KEY": "VAL"}, + timeout=999, + n_attempts=2, + debug=True, + ) + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + # Rock-only fields must be absent from Harbor YAML + rock_only = { + "job_name", + "namespace", + "experiment_id", + "labels", + "auto_stop", + "setup_commands", + "file_uploads", + "env", + "timeout", + } + for rock_field in rock_only: + assert rock_field not in data, f"Rock field '{rock_field}' should be excluded from Harbor YAML" + + def test_includes_harbor_fields(self): + cfg = HarborJobConfig(experiment_id="test-exp", n_attempts=5, debug=True) + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + assert data["n_attempts"] == 5 + assert data["debug"] is True + + def test_harbor_environment_included_when_present(self): + """Harbor environment fields (e.g., force_build) should appear under 'environment'.""" + env = RockEnvironmentConfig(force_build=True, override_cpus=4) + cfg = HarborJobConfig(experiment_id="test-exp", environment=env) + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + assert "environment" in data + assert data["environment"]["force_build"] is True + assert data["environment"]["override_cpus"] == 4 + + def test_harbor_environment_omitted_when_default(self): + """When environment has no harbor-specific fields set, 'environment' key should still appear + (because to_harbor_environment returns default fields like delete=True).""" + cfg = HarborJobConfig(experiment_id="test-exp") + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + # The harbor env may or may not have fields; just check it's valid YAML + assert isinstance(data, dict) + + def test_excludes_none_values(self): + cfg = HarborJobConfig(experiment_id="test-exp") + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + # agent_timeout_multiplier is None by default → should not appear + assert "agent_timeout_multiplier" not in data + + def test_returns_valid_yaml_string(self): + cfg = HarborJobConfig(experiment_id="test-exp", n_attempts=3) + yaml_str = cfg.to_harbor_yaml() + assert isinstance(yaml_str, str) + parsed = yaml.safe_load(yaml_str) + assert isinstance(parsed, dict) + + +# --------------------------------------------------------------------------- +# HarborJobConfig.from_yaml +# --------------------------------------------------------------------------- + + +class TestHarborJobConfigFromYaml: + def test_round_trip(self, tmp_path): + """Write a YAML config, read it back, verify fields.""" + yaml_content = textwrap.dedent( + """\ + experiment_id: test-exp + n_attempts: 3 + debug: true + agents: + - name: my-agent + import_path: my_module:Agent + timeout_multiplier: 1.5 + """ + ) + yaml_file = tmp_path / "config.yaml" + yaml_file.write_text(yaml_content) + + cfg = HarborJobConfig.from_yaml(str(yaml_file)) + assert isinstance(cfg, HarborJobConfig) + assert cfg.n_attempts == 3 + assert cfg.debug is True + assert cfg.agents[0].name == "my-agent" + assert cfg.timeout_multiplier == 1.5 + + def test_from_yaml_with_environment(self, tmp_path): + yaml_content = textwrap.dedent( + """\ + experiment_id: test-exp + environment: + force_build: true + override_cpus: 8 + n_attempts: 1 + """ + ) + yaml_file = tmp_path / "env_config.yaml" + yaml_file.write_text(yaml_content) + + cfg = HarborJobConfig.from_yaml(str(yaml_file)) + assert cfg.environment.force_build is True + assert cfg.environment.override_cpus == 8 + assert cfg.n_attempts == 1 + + def test_from_yaml_file_not_found(self): + with pytest.raises(FileNotFoundError): + HarborJobConfig.from_yaml("/nonexistent/path.yaml") + + +class TestHarborInheritsBase: + def test_harbor_inherits_base_fields(self): + """HarborJobConfig (agent's) inherits all base JobConfig fields.""" + base_fields = set(JobConfig.model_fields.keys()) + harbor_fields = set(HarborJobConfig.model_fields.keys()) + assert base_fields.issubset(harbor_fields) diff --git a/tests/unit/sdk/job/test_executor.py b/tests/unit/sdk/job/test_executor.py new file mode 100644 index 0000000000..14af6c0f94 --- /dev/null +++ b/tests/unit/sdk/job/test_executor.py @@ -0,0 +1,221 @@ +"""Tests for rock.sdk.job.executor — JobExecutor, TrialClient, JobClient.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import rock.sdk.bench # pre-import to avoid circular # noqa: F401 +import rock.sdk.job.trial.bash # register BashJobConfig -> BashTrial # noqa: F401 +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient +from rock.sdk.job.operator import ScatterOperator + + +def _make_mock_sandbox(): + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-test" + sandbox._namespace = None + sandbox._experiment_id = None + sandbox.start = AsyncMock() + sandbox.close = AsyncMock() + sandbox.create_session = AsyncMock() + sandbox.write_file_by_path = AsyncMock(return_value=MagicMock(success=True)) + + # fs.upload_dir for _upload_files — returns success obs + upload_obs = MagicMock() + upload_obs.exit_code = 0 + sandbox.fs = AsyncMock() + sandbox.fs.upload_dir = AsyncMock(return_value=upload_obs) + + sandbox.start_nohup_process = AsyncMock(return_value=(12345, None)) + sandbox.wait_for_process_completion = AsyncMock(return_value=(True, "done")) + + nohup_obs = MagicMock() + nohup_obs.output = "hello output" + nohup_obs.exit_code = 0 + sandbox.handle_nohup_output = AsyncMock(return_value=nohup_obs) + return sandbox + + +# --------------------------------------------------------------------------- +# run() — full lifecycle +# --------------------------------------------------------------------------- + + +class TestJobExecutorRun: + async def test_run_bash_single_trial_success(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + results = await executor.run(ScatterOperator(size=1), config) + + assert len(results) == 1 + assert results[0].exception_info is None + assert mock_sandbox.start.call_count == 1 + assert mock_sandbox.start_nohup_process.call_count == 1 + + async def test_run_empty_operator_returns_empty_list(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + results = await executor.run(ScatterOperator(size=0), config) + + assert results == [] + assert mock_sandbox.start.call_count == 0 + + async def test_run_scatter_size_three_runs_three_trials(self): + # Each Sandbox(...) call returns a fresh mock so we can verify 3 starts. + mocks = [_make_mock_sandbox() for _ in range(3)] + with patch("rock.sdk.job.executor.Sandbox", side_effect=mocks): + config = BashJobConfig(script="echo hi", job_name="triple") + executor = JobExecutor() + results = await executor.run(ScatterOperator(size=3), config) + + assert len(results) == 3 + for mock_sandbox in mocks: + assert mock_sandbox.start.call_count == 1 + assert mock_sandbox.start_nohup_process.call_count == 1 + + +# --------------------------------------------------------------------------- +# submit() / wait() separately +# --------------------------------------------------------------------------- + + +class TestJobExecutorSubmit: + async def test_submit_returns_job_client(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + result = await executor.submit(ScatterOperator(size=1), config) + + assert isinstance(result, JobClient) + assert len(result.trials) == 1 + assert isinstance(result.trials[0], TrialClient) + + async def test_submit_empty_returns_empty_job_client(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + result = await executor.submit(ScatterOperator(size=0), config) + + assert isinstance(result, JobClient) + assert result.trials == [] + + async def test_submit_raises_on_nohup_start_error(self): + mock_sandbox = _make_mock_sandbox() + error_obs = MagicMock() + error_obs.output = "some error" + mock_sandbox.start_nohup_process = AsyncMock(return_value=(None, error_obs)) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + with pytest.raises(RuntimeError, match="Failed to start trial"): + await executor.submit(ScatterOperator(size=1), config) + + +class TestJobExecutorWait: + async def test_wait_returns_trial_result_list(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + job_client = await executor.submit(ScatterOperator(size=1), config) + results = await executor.wait(job_client) + + assert isinstance(results, list) + assert len(results) == 1 + + async def test_wait_empty_job_client_returns_empty(self): + executor = JobExecutor() + results = await executor.wait(JobClient(trials=[])) + assert results == [] + + async def test_wait_process_failure_sets_exception_info(self): + mock_sandbox = _make_mock_sandbox() + # Process succeeds (exit_code=0) but wait reports failure (e.g. timeout). + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "timeout")) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + results = await executor.run(ScatterOperator(size=1), config) + + assert len(results) == 1 + assert results[0].exception_info is not None + + +# --------------------------------------------------------------------------- +# auto_stop behavior +# --------------------------------------------------------------------------- + + +class TestJobExecutorAutoStop: + async def test_auto_stop_true_closes_sandbox(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + config = BashJobConfig(script="echo hi", job_name="test", auto_stop=True) + executor = JobExecutor() + await executor.run(ScatterOperator(size=1), config) + + assert mock_sandbox.close.call_count == 1 + + async def test_auto_stop_false_does_not_close_sandbox(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + # default: auto_stop=False + config = BashJobConfig(script="echo hi", job_name="test") + executor = JobExecutor() + await executor.run(ScatterOperator(size=1), config) + + assert mock_sandbox.close.call_count == 0 + + +# --------------------------------------------------------------------------- +# _build_session_env +# --------------------------------------------------------------------------- + + +class TestBuildSessionEnv: + def test_merges_oss_vars_with_config_env(self, monkeypatch): + # Clear any leftover OSS vars to guarantee deterministic state. + for k in list(__import__("os").environ): + if k.startswith("OSS"): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("OSS_KEY", "value") + + config = BashJobConfig(script="echo hi", env={"X": "1"}) + merged = JobExecutor._build_session_env(config) + + assert merged is not None + assert merged["OSS_KEY"] == "value" + assert merged["X"] == "1" + + def test_config_env_overrides_oss(self, monkeypatch): + for k in list(__import__("os").environ): + if k.startswith("OSS"): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("OSS_KEY", "process_val") + + config = BashJobConfig(script="echo hi", env={"OSS_KEY": "config_val"}) + merged = JobExecutor._build_session_env(config) + + assert merged is not None + assert merged["OSS_KEY"] == "config_val" + + def test_returns_none_when_empty(self, monkeypatch): + for k in list(__import__("os").environ): + if k.startswith("OSS"): + monkeypatch.delenv(k, raising=False) + + config = BashJobConfig(script="echo hi") # default env={} + merged = JobExecutor._build_session_env(config) + + assert merged is None diff --git a/tests/unit/sdk/job/test_integration.py b/tests/unit/sdk/job/test_integration.py new file mode 100644 index 0000000000..749cb85449 --- /dev/null +++ b/tests/unit/sdk/job/test_integration.py @@ -0,0 +1,66 @@ +"""Integration tests: public exports + end-to-end behavior + backward compat.""" + +from __future__ import annotations + + +class TestPublicImports: + def test_import_job(self): + from rock.sdk.job import Job + + assert Job is not None + + def test_import_configs(self): + from rock.sdk.job import BashJobConfig, JobConfig + + assert issubclass(BashJobConfig, JobConfig) + + def test_import_results(self): + from rock.sdk.job import JobStatus, TrialResult + + assert JobStatus.COMPLETED == "completed" + assert TrialResult is not None + + def test_import_operator(self): + from rock.sdk.job import Operator, ScatterOperator + + assert issubclass(ScatterOperator, Operator) + + def test_import_trial(self): + from rock.sdk.job import AbstractTrial, register_trial + + assert AbstractTrial is not None + assert callable(register_trial) + + def test_import_executor(self): + from rock.sdk.job import JobClient, JobExecutor, TrialClient + + assert JobExecutor is not None + assert TrialClient is not None + assert JobClient is not None + + +class TestTrialRegistryAutoRegistration: + def test_bash_registered(self): + from rock.sdk.job import BashJobConfig + from rock.sdk.job.trial.registry import _TRIAL_REGISTRY + + assert BashJobConfig in _TRIAL_REGISTRY + + def test_harbor_registered(self): + # Importing rock.sdk.job triggers auto-registration + import rock.sdk.job # noqa: F401 + from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + from rock.sdk.job.trial.registry import _TRIAL_REGISTRY + + assert HarborJobConfig in _TRIAL_REGISTRY + + +class TestBackwardCompat: + def test_old_agent_imports_still_work(self): + """rock.sdk.bench (formerly rock.sdk.agent) must still export Job, JobConfig, JobResult, JobStatus.""" + from rock.sdk.bench import Job, JobConfig, JobResult, JobStatus + + assert Job is not None + assert JobConfig is not None + assert JobResult is not None + assert JobStatus is not None diff --git a/tests/unit/sdk/job/test_job.py b/tests/unit/sdk/job/test_job.py new file mode 100644 index 0000000000..058970a06e --- /dev/null +++ b/tests/unit/sdk/job/test_job.py @@ -0,0 +1,169 @@ +"""Tests for rock.sdk.job.job — Job Facade.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import rock.sdk.bench # pre-import to avoid circular # noqa: F401 +import rock.sdk.job.trial.bash # trigger BashTrial registration # noqa: F401 +from rock.sdk.job import Job +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.operator import ScatterOperator +from rock.sdk.job.result import JobStatus + + +def _make_mock_sandbox(): + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-facade" + sandbox.start = AsyncMock() + sandbox.close = AsyncMock() + sandbox.create_session = AsyncMock() + sandbox.write_file_by_path = AsyncMock() + sandbox.arun = AsyncMock() + + upload_obs = MagicMock() + upload_obs.exit_code = 0 + sandbox.fs = AsyncMock() + sandbox.fs.upload_dir = AsyncMock(return_value=upload_obs) + + sandbox.start_nohup_process = AsyncMock(return_value=(99, None)) + sandbox.wait_for_process_completion = AsyncMock(return_value=(True, "done")) + + obs = MagicMock() + obs.output = "ok" + obs.exit_code = 0 + sandbox.handle_nohup_output = AsyncMock(return_value=obs) + return sandbox + + +# --------------------------------------------------------------------------- +# run() — full lifecycle +# --------------------------------------------------------------------------- + + +class TestJobRun: + async def test_run_returns_completed_result_on_success(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="test")).run() + + assert result.status == JobStatus.COMPLETED + assert len(result.trial_results) == 1 + + async def test_run_returns_failed_status_when_trial_fails(self): + mock_sandbox = _make_mock_sandbox() + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "timeout")) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="test")).run() + + assert result.status == JobStatus.FAILED + assert len(result.trial_results) == 1 + assert result.trial_results[0].exception_info is not None + + +# --------------------------------------------------------------------------- +# submit() / wait() separately +# --------------------------------------------------------------------------- + + +class TestJobSubmitWait: + async def test_submit_then_wait_equivalent_to_run(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + await job.submit() + result = await job.wait() + + assert result.status == JobStatus.COMPLETED + assert len(result.trial_results) == 1 + + async def test_wait_without_submit_raises_runtime_error(self): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + with pytest.raises(RuntimeError, match="No submitted job"): + await job.wait() + + +# --------------------------------------------------------------------------- +# cancel() +# --------------------------------------------------------------------------- + + +class TestJobCancel: + async def test_cancel_kills_all_trial_sandboxes(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + await job.submit() + await job.cancel() + + assert mock_sandbox.arun.called + # Extract the cmd kwarg from the last arun call + call = mock_sandbox.arun.call_args + cmd = call.kwargs.get("cmd", "") + assert "kill" in cmd + + async def test_cancel_without_submit_is_noop(self): + job = Job(BashJobConfig(script="echo hi", job_name="test")) + # Should not raise + await job.cancel() + + +# --------------------------------------------------------------------------- +# Operator parameter +# --------------------------------------------------------------------------- + + +class TestJobOperator: + async def test_custom_operator_with_size_two_produces_two_trials(self): + mocks = [_make_mock_sandbox() for _ in range(2)] + with patch("rock.sdk.job.executor.Sandbox", side_effect=mocks): + result = await Job( + BashJobConfig(script="echo hi", job_name="test"), + operator=ScatterOperator(size=2), + ).run() + + assert len(result.trial_results) == 2 + + async def test_default_operator_is_scatter_size_one(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="test")).run() + + assert len(result.trial_results) == 1 + + +# --------------------------------------------------------------------------- +# _build_result +# --------------------------------------------------------------------------- + + +class TestJobBuildResult: + async def test_build_result_uses_config_labels(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="test", labels={"team": "rl"})).run() + + assert result.labels == {"team": "rl"} + + async def test_build_result_sets_job_id_from_job_name(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="my-job")).run() + + assert result.job_id == "my-job" + + async def test_build_result_any_failure_marks_job_failed(self): + # Single trial, but force failure -> overall FAILED + mock_sandbox = _make_mock_sandbox() + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "err")) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job( + BashJobConfig(script="echo hi", job_name="test"), + operator=ScatterOperator(size=1), + ).run() + + assert result.status == JobStatus.FAILED diff --git a/tests/unit/sdk/job/test_operator.py b/tests/unit/sdk/job/test_operator.py new file mode 100644 index 0000000000..2e897181ee --- /dev/null +++ b/tests/unit/sdk/job/test_operator.py @@ -0,0 +1,66 @@ +"""Tests for rock.sdk.job.operator — Operator, ScatterOperator.""" + +from __future__ import annotations + +import pytest + +# Import bench first to avoid circular-import pitfall in rock.sdk.job.config +import rock.sdk.bench # noqa: F401 +from rock.sdk.job.config import BashJobConfig, JobConfig +from rock.sdk.job.operator import Operator, ScatterOperator +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.bash import BashTrial + + +class TestOperatorABC: + def test_cannot_instantiate(self): + with pytest.raises(TypeError): + Operator() + + +class TestScatterOperator: + def test_default_size_is_one(self): + op = ScatterOperator() + assert op.size == 1 + + def test_apply_returns_one_trial_by_default(self): + op = ScatterOperator() + trials = op.apply(BashJobConfig(script="echo hi")) + assert len(trials) == 1 + assert isinstance(trials[0], BashTrial) + + def test_apply_returns_n_trials(self): + op = ScatterOperator(size=3) + trials = op.apply(BashJobConfig(script="echo hi")) + assert len(trials) == 3 + for t in trials: + assert isinstance(t, BashTrial) + + def test_size_zero_returns_empty(self): + op = ScatterOperator(size=0) + assert op.apply(BashJobConfig(script="echo hi")) == [] + + def test_size_negative_returns_empty(self): + op = ScatterOperator(size=-5) + assert op.apply(BashJobConfig(script="echo hi")) == [] + + def test_returns_correct_trial_type_for_bash_config(self): + op = ScatterOperator(size=2) + trials = op.apply(BashJobConfig(script="ls")) + assert all(isinstance(t, BashTrial) for t in trials) + + +class TestCustomOperator: + def test_custom_subclass_can_override_apply(self): + class FixedOperator(Operator): + def __init__(self, trials: list[AbstractTrial]): + self._trials = trials + + def apply(self, config: JobConfig) -> list[AbstractTrial]: + return list(self._trials) + + fixed_trials = [BashTrial(BashJobConfig(script="a")), BashTrial(BashJobConfig(script="b"))] + op = FixedOperator(fixed_trials) + result = op.apply(BashJobConfig(script="ignored")) + assert result == fixed_trials + assert len(result) == 2 diff --git a/tests/unit/sdk/job/test_result.py b/tests/unit/sdk/job/test_result.py new file mode 100644 index 0000000000..faabae9e28 --- /dev/null +++ b/tests/unit/sdk/job/test_result.py @@ -0,0 +1,70 @@ +"""Tests for rock.sdk.job.result — JobStatus, JobResult[T].""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel + +from rock.sdk.job.result import JobResult, JobStatus + + +class _Item(BaseModel): + """Stub item for testing JobResult[T].""" + + name: str = "" + status: str = "completed" + + @property + def score(self) -> float: + return 1.0 if self.status == "completed" else 0.0 + + +class TestJobStatus: + def test_values(self): + assert JobStatus.PENDING == "pending" + assert JobStatus.RUNNING == "running" + assert JobStatus.COMPLETED == "completed" + assert JobStatus.FAILED == "failed" + assert JobStatus.CANCELLED == "cancelled" + + def test_is_str(self): + assert isinstance(JobStatus.COMPLETED, str) + + +class TestJobResult: + def test_defaults(self): + r = JobResult() + assert r.job_id == "" + assert r.status == JobStatus.COMPLETED + assert r.labels == {} + assert r.trial_results == [] + assert r.raw_output == "" + assert r.exit_code == 0 + + def test_score_empty(self): + assert JobResult().score == 0.0 + + def test_score_with_items(self): + r = JobResult[_Item](trial_results=[_Item(), _Item(status="failed")]) + assert r.score == pytest.approx(0.5) + + def test_n_completed(self): + r = JobResult[_Item]( + trial_results=[_Item(), _Item(status="failed"), _Item()], + ) + assert r.n_completed == 2 + + def test_n_failed(self): + r = JobResult[_Item]( + trial_results=[_Item(status="failed"), _Item(status="failed")], + ) + assert r.n_failed == 2 + + def test_labels(self): + r = JobResult(labels={"env": "test"}) + assert r.labels == {"env": "test"} + + def test_raw_output_and_exit_code(self): + r = JobResult(raw_output="output", exit_code=1) + assert r.raw_output == "output" + assert r.exit_code == 1 diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py new file mode 100644 index 0000000000..ae2d1636bd --- /dev/null +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -0,0 +1,135 @@ +"""Tests for rock.sdk.job.trial.bash — BashTrial.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +# Import bench first to avoid circular-import pitfall in rock.sdk.job.config +import rock.sdk.bench # noqa: F401 +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.trial.bash import BashTrial +from rock.sdk.job.trial.registry import _create_trial + + +def _success_obs(): + obs = MagicMock() + obs.exit_code = 0 + return obs + + +# --------------------------------------------------------------------------- +# BashTrial.build() +# --------------------------------------------------------------------------- + + +class TestBashTrialBuild: + def test_build_basic_script(self): + cfg = BashJobConfig(script="echo hello") + trial = BashTrial(cfg) + out = trial.build() + assert "#!/bin/bash" in out + assert "set -e" in out + assert "echo hello" in out + + def test_build_with_setup_commands(self): + cfg = BashJobConfig( + setup_commands=["pip install -r requirements.txt"], + script="python main.py", + ) + trial = BashTrial(cfg) + out = trial.build() + + assert "pip install -r requirements.txt" in out + assert "python main.py" in out + # Setup comes before main script + assert out.index("pip install -r requirements.txt") < out.index("python main.py") + + def test_build_no_script_only_setup(self): + cfg = BashJobConfig(setup_commands=["echo setup"]) + trial = BashTrial(cfg) + out = trial.build() + assert "#!/bin/bash" in out + assert "set -e" in out + assert "echo setup" in out + + +# --------------------------------------------------------------------------- +# BashTrial.setup() +# --------------------------------------------------------------------------- + + +class TestBashTrialSetup: + async def test_setup_uploads_files(self): + cfg = BashJobConfig( + script="echo hi", + file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], + ) + trial = BashTrial(cfg) + mock_sandbox = AsyncMock() + mock_sandbox.fs.upload_dir = AsyncMock(return_value=_success_obs()) + + await trial.setup(mock_sandbox) + + assert mock_sandbox.fs.upload_dir.call_count == 2 + mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/local/a", target_dir="/sandbox/a") + mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/local/b", target_dir="/sandbox/b") + + async def test_setup_reads_script_path(self): + expected = "expected content" + with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as f: + f.write(expected) + tmp_path = f.name + try: + cfg = BashJobConfig(script_path=tmp_path) + trial = BashTrial(cfg) + mock_sandbox = AsyncMock() + mock_sandbox.fs.upload_dir = AsyncMock(return_value=_success_obs()) + + await trial.setup(mock_sandbox) + + assert trial._config.script == expected + finally: + Path(tmp_path).unlink(missing_ok=True) + + +# --------------------------------------------------------------------------- +# BashTrial.collect() +# --------------------------------------------------------------------------- + + +class TestBashTrialCollect: + async def test_collect_exit_code_zero(self): + cfg = BashJobConfig(script="echo hi", job_name="myjob") + trial = BashTrial(cfg) + mock_sandbox = AsyncMock() + + result = await trial.collect(mock_sandbox, output="hi\n", exit_code=0) + + assert result.exception_info is None + assert result.task_name == "myjob" + assert result.status == "completed" + + async def test_collect_exit_code_nonzero(self): + cfg = BashJobConfig(script="false", job_name="myjob") + trial = BashTrial(cfg) + mock_sandbox = AsyncMock() + + result = await trial.collect(mock_sandbox, output="", exit_code=1) + + assert result.exception_info is not None + assert result.exception_info.exception_type == "BashExitCode" + assert result.status == "failed" + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +class TestBashTrialRegistration: + def test_bash_config_creates_bash_trial(self): + cfg = BashJobConfig(script="echo hi") + trial = _create_trial(cfg) + assert isinstance(trial, BashTrial) diff --git a/tests/unit/sdk/job/test_trial_harbor.py b/tests/unit/sdk/job/test_trial_harbor.py new file mode 100644 index 0000000000..cf42a0001f --- /dev/null +++ b/tests/unit/sdk/job/test_trial_harbor.py @@ -0,0 +1,144 @@ +"""Tests for rock.sdk.job.trial.harbor — HarborTrial.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +# Pre-import bench to avoid circular-import pitfalls in rock.sdk.job.config +import rock.sdk.bench # noqa: F401 +from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig +from rock.sdk.job.trial.harbor import HarborTrial +from rock.sdk.job.trial.registry import _create_trial + + +def _success_obs(): + obs = MagicMock() + obs.exit_code = 0 + return obs + + +# --------------------------------------------------------------------------- +# HarborTrial.build() +# --------------------------------------------------------------------------- + + +class TestHarborTrialBuild: + def test_build_contains_harbor_jobs_start(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + script = trial.build() + assert "harbor jobs start -c" in script + + def test_build_contains_dockerd_startup(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + script = trial.build() + assert "dockerd" in script + + def test_build_contains_shebang_and_set_e(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + script = trial.build() + assert "#!/bin/bash" in script + assert "set -e" in script + + def test_build_with_setup_commands_includes_them(self): + cfg = HarborJobConfig( + job_name="test", + experiment_id="exp-1", + setup_commands=["pip install harbor"], + ) + trial = HarborTrial(cfg) + script = trial.build() + assert "pip install harbor" in script + + def test_build_without_setup_commands_uses_placeholder(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + script = trial.build() + assert "No setup commands" in script + + +# --------------------------------------------------------------------------- +# HarborTrial.setup() +# --------------------------------------------------------------------------- + + +class TestHarborTrialSetup: + async def test_setup_uploads_harbor_yaml(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + mock_sandbox = AsyncMock() + mock_sandbox.fs.upload_dir = AsyncMock(return_value=_success_obs()) + mock_sandbox.write_file_by_path = AsyncMock() + + await trial.setup(mock_sandbox) + + mock_sandbox.write_file_by_path.assert_called_once() + args, kwargs = mock_sandbox.write_file_by_path.call_args + yaml_content = args[0] if args else kwargs.get("content") + # Harbor YAML serializes `agents` field from HarborJobConfig + assert "agents:" in yaml_content + + +# --------------------------------------------------------------------------- +# HarborTrial.collect() +# --------------------------------------------------------------------------- + + +class TestHarborTrialCollect: + async def test_collect_with_trial_results_found(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + + trial_json = { + "task_name": "fix-dockerfile", + "trial_name": "trial-001", + "started_at": "2026-01-01T00:00:00Z", + "finished_at": "2026-01-01T00:01:00Z", + "verifier_result": {"rewards": {"reward": 1.0}}, + "agent_result": {}, + "exception_info": None, + } + + mock_sandbox = AsyncMock() + list_result = MagicMock() + list_result.stdout = f"{cfg.jobs_dir}/test/trial-001/result.json\n" + mock_sandbox.execute = AsyncMock(return_value=list_result) + + read_response = MagicMock() + read_response.content = json.dumps(trial_json) + mock_sandbox.read_file = AsyncMock(return_value=read_response) + + result = await trial.collect(mock_sandbox, output="", exit_code=0) + + assert result.task_name == "fix-dockerfile" + assert result.exception_info is None + assert result.score == 1.0 + + async def test_collect_with_no_trials(self): + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") + trial = HarborTrial(cfg) + + mock_sandbox = AsyncMock() + list_result = MagicMock() + list_result.stdout = "" + mock_sandbox.execute = AsyncMock(return_value=list_result) + + result = await trial.collect(mock_sandbox, output="", exit_code=0) + + assert result.exception_info is not None + assert result.exception_info.exception_type == "HarborNoTrials" + + +# --------------------------------------------------------------------------- +# Auto-registration +# --------------------------------------------------------------------------- + + +class TestHarborTrialRegistration: + def test_harbor_config_creates_harbor_trial(self): + cfg = HarborJobConfig(experiment_id="exp-1") + trial = _create_trial(cfg) + assert isinstance(trial, HarborTrial) diff --git a/tests/unit/sdk/job/test_trial_registry.py b/tests/unit/sdk/job/test_trial_registry.py new file mode 100644 index 0000000000..7c5972a5c1 --- /dev/null +++ b/tests/unit/sdk/job/test_trial_registry.py @@ -0,0 +1,132 @@ +"""Tests for rock.sdk.job.trial — AbstractTrial and registry.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +# Import bench first to avoid circular-import pitfall in rock.sdk.job.config +import rock.sdk.bench # noqa: F401 +from rock.sdk.job.config import JobConfig +from rock.sdk.job.result import TrialResult +from rock.sdk.job.trial.abstract import AbstractTrial +from rock.sdk.job.trial.registry import _TRIAL_REGISTRY, _create_trial, register_trial + +# --------------------------------------------------------------------------- +# Stubs used across tests +# --------------------------------------------------------------------------- + + +class _StubConfig(JobConfig): + stub_field: str = "test" + + +class _StubTrial(AbstractTrial): + async def setup(self, sandbox) -> None: + pass + + def build(self) -> str: + return "echo stub" + + async def collect(self, sandbox, output, exit_code) -> TrialResult: + return TrialResult(task_name="stub") + + +@pytest.fixture(autouse=True) +def _reset_registry(): + """Preserve and restore the registry between tests to avoid cross-test pollution.""" + saved = dict(_TRIAL_REGISTRY) + _TRIAL_REGISTRY.clear() + yield + _TRIAL_REGISTRY.clear() + _TRIAL_REGISTRY.update(saved) + + +# --------------------------------------------------------------------------- +# AbstractTrial +# --------------------------------------------------------------------------- + + +class TestAbstractTrial: + def test_cannot_instantiate_directly(self): + """AbstractTrial is abstract and cannot be instantiated.""" + with pytest.raises(TypeError): + AbstractTrial(JobConfig()) # type: ignore[abstract] + + def test_subclass_can_be_instantiated(self): + cfg = _StubConfig() + trial = _StubTrial(cfg) + assert isinstance(trial, AbstractTrial) + + def test_config_reference_held(self): + cfg = _StubConfig(stub_field="hello") + trial = _StubTrial(cfg) + assert trial._config is cfg + + async def test_upload_files_iterates_all_entries(self): + mock_sandbox = AsyncMock() + success_obs = MagicMock() + success_obs.exit_code = 0 + mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) + cfg = _StubConfig(file_uploads=[("/a", "/b"), ("/c", "/d")]) + trial = _StubTrial(cfg) + + await trial._upload_files(mock_sandbox) + + assert mock_sandbox.fs.upload_dir.call_count == 2 + mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/a", target_dir="/b") + mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/c", target_dir="/d") + + async def test_upload_files_noop_when_empty(self): + mock_sandbox = AsyncMock() + success_obs = MagicMock() + success_obs.exit_code = 0 + mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) + cfg = _StubConfig(file_uploads=[]) + trial = _StubTrial(cfg) + + await trial._upload_files(mock_sandbox) + + mock_sandbox.fs.upload_dir.assert_not_called() + + async def test_upload_files_raises_on_failure(self): + cfg = _StubConfig(file_uploads=[("/a", "/b")]) + trial = _StubTrial(cfg) + mock_sandbox = AsyncMock() + failure_obs = MagicMock() + failure_obs.exit_code = 1 + failure_obs.failure_reason = "disk full" + mock_sandbox.fs.upload_dir = AsyncMock(return_value=failure_obs) + + with pytest.raises(RuntimeError, match="disk full"): + await trial._upload_files(mock_sandbox) + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +class TestTrialRegistry: + def test_create_trial_unregistered_raises_type_error(self): + cfg = _StubConfig() + with pytest.raises(TypeError, match="No trial registered"): + _create_trial(cfg) + + def test_register_and_create_trial(self): + register_trial(_StubConfig, _StubTrial) + cfg = _StubConfig() + trial = _create_trial(cfg) + assert isinstance(trial, _StubTrial) + assert trial._config is cfg + + def test_error_lists_supported_configs(self): + register_trial(_StubConfig, _StubTrial) + + class _OtherConfig(JobConfig): + pass + + with pytest.raises(TypeError) as exc_info: + _create_trial(_OtherConfig()) + assert "_StubConfig" in str(exc_info.value) diff --git a/uv.lock b/uv.lock index cf569139eb..6a8933124d 100644 --- a/uv.lock +++ b/uv.lock @@ -4035,7 +4035,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.4.6" +version = "1.5.0" source = { editable = "." } dependencies = [ { name = "anyio" }, From ad64b28813067f8fd4e2c396accf900a7ae97b35 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 14 Apr 2026 21:25:53 +0800 Subject: [PATCH 027/226] fix: add SELECT 1 readiness check to pg_container fixture (#778) pg_isready can succeed before PostgreSQL fully accepts application connections, causing CannotConnectNowError / ConnectionResetError under parallel test workers (pytest-xdist). Follow up with a real SQL query to close the startup race window. --- tests/unit/conftest.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4ad643feaf..b5f53ace65 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -353,6 +353,15 @@ def pg_container(): else: raise TimeoutError("PostgreSQL container did not become ready within 30s") + # Confirm with a real SQL query to close the startup race window. + while _t.time() < deadline: + code, _ = container.exec_run(f'psql -U {_PG_USER} -d {_PG_DB} -c "SELECT 1"') + if code == 0: + break + _t.sleep(0.5) + else: + raise TimeoutError("PostgreSQL: pg_isready OK but SELECT 1 still failing") + host, port = _docker_resolve_host_port(container, network_name, _PG_PORT) yield { "host": host, "port": port, From 10f94fecf99405dc544d068b7e69383cf881daed Mon Sep 17 00:00:00 2001 From: dengwx Date: Tue, 14 Apr 2026 22:31:13 +0800 Subject: [PATCH 028/226] refactor(job): hoist on_sandbox_ready backfill to AbstractTrial (#788) (#789) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * rename JobConfig and HarborTrialResult * fix(job): close G1-G7 gaps between bench/job.py and rock.sdk.job (#784) * docs(job): clarify G1 approach as union+isinstance in bench-replacement Align docs/dev/job/bench-replacement.md with the design decision for Task 1: AbstractTrial.collect returns TrialResult | list[TrialResult], and Job._build_result uses isinstance to decide extend vs append. BashTrial stays unchanged (single), HarborTrial returns list. * chore: baseline for bench-replacement gap fixes (222 passing) * fix(job/G1a): HarborTrial.collect returns all sub-trial results as list * fix(job/G1b): flatten list-returning collect() into JobResult.trial_results * refactor(job/G1): drop BaseTrialResult alias in harbor.py, use TrialResult directly Per user request: remove the local import alias 'from rock.sdk.job.result import TrialResult as BaseTrialResult' and use TrialResult directly. No behavior change. * test(job/G1): add multi-trial timeout exception-injection coverage Addresses code review I1: exception-injection loop on multi-sub-trial timeout previously untested. New test asserts: - every sub-trial without its own exception_info gets ProcessTimeout - pre-existing exception_info on a sub-trial is preserved * fix(job/G5): populate JobResult.raw_output and exit_code from sandbox obs * fix(job/G6): write script/nohup output under USER_DEFINED_LOGS, not /tmp * fix(job/G7): sync HarborJobConfig.auto_stop with environment.auto_stop OR semantics: if either the top-level HarborJobConfig.auto_stop or the nested environment.auto_stop is True, both become True. Preserves legacy 'environment.auto_stop=True' usage while letting new code set auto_stop at the top level. * fix(job/G3): HarborJobConfig auto-generates job_name from dataset/task + uuid * fix(job/G2): HarborJobConfig derives effective timeout from agent + multiplier * fix(job/G4): AbstractTrial.on_sandbox_ready hook + HarborTrial backfills ns/exp_id * test(job): blue-green equivalence tests for G1-G7 gap fixes * chore(job): mark rock.sdk.bench.Job as deprecated; point users to rock.sdk.job.Job G1-G7 gaps all closed and validated by blue-green equivalence tests. Emits DeprecationWarning (stacklevel=2) from Job.__init__ so callers see their own site. Scheduled for removal in 1.7.x. Refs #783 * fix(job): pass job_name to harbor YAML and migrate demo to new Job path (#785) Harbor was using timestamp-based directory names because job_name was excluded from the serialized YAML config. Re-inject job_name in to_harbor_yaml() so harbor uses it as the job directory name, allowing result collection via {jobs_dir}/{job_name} to work correctly. Also migrate harbor_demo.py from deprecated rock.sdk.bench.Job to rock.sdk.job.Job. Co-authored-by: Claude Opus 4.6 (1M context) * refactor(job): move on_sandbox_ready backfill to AbstractTrial Hoist the namespace/experiment_id backfill + consistency check from HarborTrial up into AbstractTrial.on_sandbox_ready so BashTrial (and any future JobConfig-based trial) inherits the same behavior. The error message uses type(self._config).__name__ instead of a hardcoded class name. _make_mock_sandbox() in test_job.py now sets _namespace / _experiment_id to None so the auto-mock children don't trip the default backfill when two trials share one config. Co-Authored-By: Claude Opus 4.6 (1M context) * docs(job): update G4 description after hoisting on_sandbox_ready on_sandbox_ready is no longer a HarborTrial override — the backfill + consistency check now lives on AbstractTrial and is shared across HarborTrial and BashTrial. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) --- docs/dev/job/bench-replacement.md | 283 ++++++++++++++++++ examples/harbor/harbor_demo.py | 20 +- rock/cli/command/job.py | 2 +- rock/sdk/bench/__init__.py | 8 +- rock/sdk/bench/job.py | 22 +- rock/sdk/bench/models/__init__.py | 4 +- rock/sdk/bench/models/job/__init__.py | 4 +- rock/sdk/bench/models/job/config.py | 85 +++++- rock/sdk/bench/models/trial/__init__.py | 4 +- rock/sdk/bench/models/trial/result.py | 12 +- rock/sdk/job/api.py | 27 +- rock/sdk/job/executor.py | 41 ++- rock/sdk/job/result.py | 3 + rock/sdk/job/trial/abstract.py | 38 ++- rock/sdk/job/trial/harbor.py | 44 +-- tests/unit/sdk/agent/test_job.py | 63 ++-- .../agent/test_job_config_serialization.py | 41 +-- .../sdk/agent/test_jobconfig_experiment_id.py | 26 +- tests/unit/sdk/agent/test_models.py | 18 +- tests/unit/sdk/agent/test_oss_mirror.py | 52 ++-- .../sdk/job/test_blue_green_equivalence.py | 149 +++++++++ tests/unit/sdk/job/test_cli_job.py | 2 +- tests/unit/sdk/job/test_config.py | 152 +++++++++- tests/unit/sdk/job/test_executor.py | 57 ++++ tests/unit/sdk/job/test_integration.py | 8 +- tests/unit/sdk/job/test_job.py | 134 +++++++++ tests/unit/sdk/job/test_trial_bash.py | 46 +++ tests/unit/sdk/job/test_trial_harbor.py | 96 ++++-- uv.lock | 2 +- 29 files changed, 1241 insertions(+), 202 deletions(-) create mode 100644 docs/dev/job/bench-replacement.md create mode 100644 tests/unit/sdk/job/test_blue_green_equivalence.py diff --git a/docs/dev/job/bench-replacement.md b/docs/dev/job/bench-replacement.md new file mode 100644 index 0000000000..3db8aae9a1 --- /dev/null +++ b/docs/dev/job/bench-replacement.md @@ -0,0 +1,283 @@ +# `rock/sdk/bench/job.py` → `rock/sdk/job/api.py` + `trial/harbor.py` 替换分析 + +> 目的：判断新 Job 架构（`job/api.py` Facade + `JobExecutor` + `HarborTrial`）能否替换老 `bench/job.py` 中的单体 `Job` 类。 + +## 0. TL;DR + +| 维度 | 结论 | +|------|------| +| **核心执行流程** | ✅ 1:1 对应，可替换 | +| **Harbor YAML / 脚本模板** | ✅ 完全一致 | +| **Sandbox 生命周期 / session / OSS 转发** | ✅ 已迁移到 `JobExecutor` | +| **Harbor 多子 trial 结果聚合** | ✅ (G1) `AbstractTrial.collect` 返回 `TrialResult \| list[TrialResult]`；`Job._build_result` 按 `isinstance` 拍平 | +| **Agent-aware wait timeout** | ✅ (G2) `HarborJobConfig._compute_effective_timeout` 写回 `self.timeout` | +| **`job_name` 自动生成** | ✅ (G3) `HarborJobConfig._auto_job_name` validator | +| **`namespace` / `experiment_id` 从 sandbox 回填** | ✅ (G4) `AbstractTrial.on_sandbox_ready` 默认实现（所有 Trial 继承） | +| **`JobResult.raw_output` / `exit_code` 填充** | ✅ (G5) `TrialResult.raw_output/exit_code` + `JobExecutor._do_wait` 写回 + `Job._build_result` 聚合 | +| **脚本 / 输出文件路径** | ✅ (G6) `JobExecutor._job_tmp_prefix` 回到 `USER_DEFINED_LOGS` | +| **`auto_stop` 两处字段同步** | ✅ (G7) `HarborJobConfig._sync_auto_stop` OR 语义 validator | + +**结论：G1-G7 已于 #783 全部修复，blue-green 等价测试锁死契约；`rock/sdk/bench/Job` 已加 `DeprecationWarning`，计划 1.7.x 下线。** + +--- + +## 1. 架构对比 + +### 1.1 老实现（单体） + +``` +rock/sdk/bench/job.py (334 lines) + └── class Job + ├── __init__(config: HarborJobConfig) + ├── run() / submit() / wait() / cancel() + ├── _prepare_and_start ─┐ + ├── _render_run_script │ 硬编码 Harbor 逻辑 + ├── _create_session │ + sandbox 生命周期 + ├── _build_session_env │ + 脚本渲染 + ├── _collect_results │ + 结果收集 + ├── _generate_default_job_name + ├── _autofill_sandbox_info + ├── _get_wait_timeout │ ← agent-aware 超时 + └── _upload_content ─┘ +``` + +### 1.2 新实现（分层） + +``` +rock/sdk/job/ + ├── api.py ← Job Facade (74 lines, 通用) + ├── executor.py ← JobExecutor: sandbox 生命周期 + nohup + ├── operator.py ← Operator / ScatterOperator + ├── config.py ← JobConfig / BashJobConfig 基类 + └── trial/ + ├── abstract.py ← setup / build / collect 三阶段 + ├── registry.py ← Config → Trial 注册表 + ├── bash.py ← BashTrial + └── harbor.py ← HarborTrial (116 lines, 仅 Harbor 特有) + +职责切分: + Job — 极薄 Facade, 组装 config + operator + Operator — 决定分发多少份 Trial + Executor — 并行启动 sandbox + 并行等待 + Trial — 任务逻辑 (上传文件、生成脚本、解析结果) +``` + +--- + +## 2. 核心流程逐段对照 + +### 2.1 `submit()` 流程 + +| 步骤 | 老 `bench/job.py` | 新 `api.py` + `executor.py` + `harbor.py` | 状态 | +|------|-------------------|-------------------------------------------|------| +| 生成 job_name | `_generate_default_job_name()` (L273) | **缺失** | ❌ | +| 启动 sandbox | `Sandbox(env).start()` (L92) | `JobExecutor._do_submit` | ✅ | +| 回填 namespace/exp_id | `_autofill_sandbox_info()` (L303) | **缺失**（只靠 `HarborJobConfig._sync_experiment_id` validator 在 config 创建时强制） | ❌ | +| 创建 bash session | `_create_session()` (L218) | `JobExecutor._do_submit` → `create_session` | ✅ | +| OSS env 合并 | `_build_session_env()` (L207) | `JobExecutor._build_session_env` (L138) | ✅ 一致 | +| 上传 `file_uploads` | `sandbox.fs.upload_dir` 循环 (L163) | `AbstractTrial._upload_files` (L38) | ✅ | +| 上传 Harbor YAML | `_upload_content(to_harbor_yaml, ...)` (L170) | `HarborTrial.setup` → `write_file_by_path` | ✅ | +| 渲染 run script | `_render_run_script` (L188) | `HarborTrial.build` (L64) | ✅ **模板完全一致** | +| 上传 run script | `_upload_content(script, script_path)` (L171) | `JobExecutor._do_submit` → `write_file_by_path` | ⚠️ 路径不同（见 §3） | +| nohup 启动 | `start_nohup_process(bash script)` (L176) | `JobExecutor._do_submit` 同上 (L95) | ✅ | + +### 2.2 `wait()` 流程 + +| 步骤 | 老实现 | 新实现 | 状态 | +|------|--------|--------|------| +| 计算 wait timeout | `_get_wait_timeout()`：`agent.max_timeout * multiplier + 600`，兜底 7200 | `config.timeout` 直接读（默认 3600） | ❌ **回归** | +| `wait_for_process_completion` | L104 | `JobExecutor._do_wait` L112 | ✅ | +| `handle_nohup_output` | L111 | 同上 L118 | ✅ | +| 收集 trial 结果 | `_collect_results()` 返回 **所有** trial JSON 打包成 `JobResult(trial_results=[all])` | `HarborTrial.collect()` 读所有 trial JSON **但只 `return trial_results[0]`** | ❌ **严重回归** | +| 写回 `raw_output` / `exit_code` | `result.raw_output = obs.output` (L121) | 从不设置，`JobResult.raw_output` 始终为 "" | ❌ | +| 失败状态传播 | `if not success: result.status = FAILED` | 通过 `exception_info` + `_build_result(all_success)` | ✅ 语义等价 | +| auto_stop | `self._config.environment.auto_stop` (L128) | `config.auto_stop` (L135) | ⚠️ **字段位置不同**：新代码读的是 `JobConfig.auto_stop` 基类字段，老代码读 `environment.auto_stop`。Harbor 的 `JobConfig` 继承后两者可能都存在，需确认用户填哪一个 | + +### 2.3 `cancel()` + +| 老 | 新 | 状态 | +|----|----|------| +| 单 pid `kill {pid}` | 遍历所有 `TrialClient` 逐个 kill | ✅ 新更通用 | + +--- + +## 3. 关键回归与缺口（必须修复） + +### G1 — Harbor 多子 trial 结果丢失（**BLOCKER**） + +**现象**： +```python +# rock/sdk/job/trial/harbor.py:78 +async def collect(self, sandbox, output, exit_code) -> BaseTrialResult: + trial_results = await self._collect_trial_results(sandbox) + if trial_results: + return trial_results[0] # ← 只取第一个，丢弃其他 N-1 个 +``` + +**老行为**：`bench/job.py:233` 把所有子 trial 结果打包进 `JobResult.trial_results`： +```python +return JobResult(trial_results=trial_results) # 完整列表 +``` + +**架构错位**：老 Job 的 `JobResult.trial_results` 是 "一个 sandbox 执行 Harbor → 产出 N 个子 trial 结果"；新 `AbstractTrial.collect` 是 "一个 Trial 产出一个 `BaseTrialResult`"，`Job._build_result` 再把 M 个 trial 聚合。两套语义不对齐。 + +**修复选项**： +- **A** 改 `HarborTrial.collect` 返回一个 "wrapper" `TrialResult`，把所有 sub-trial 放进自定义字段（类型系统不干净）。 +- **B** 改 `AbstractTrial.collect` 签名为 `→ TrialResult | list[TrialResult]`（union）：单结果 Trial（如 `BashTrial`）保持返回 `TrialResult`，多结果 Trial（如 `HarborTrial`）返回 `list[TrialResult]`；`Job._build_result` 在 `JobResult.trial_results` 聚合时根据实际返回类型 `isinstance(r, list)` 决定 `extend` 还是 `append`。 +- **C** 在 `HarborTrial` 中把"读 N 个 result.json"提前到 `ScatterOperator` 层：Operator 先 probe sandbox 预览 task 数，再 scatter N 份 Trial——但 Harbor 本身是一次性启动 orchestrator，没法这样拆。 + +**推荐 B（union + 拍平）**：接口上明确允许单/多两种返回形态，`BashTrial` 继续返回单个 `TrialResult` 无需改造，`HarborTrial` 返回完整 `list[TrialResult]`；`Job` Facade 在 `_build_result` 统一 flatten 到 `JobResult.trial_results`。这样既最小扰动下游 Trial 实现，又保留干净的类型签名。 + +### G2 — Agent-aware wait timeout 丢失 + +**老逻辑** (`bench/job.py:141-151`)： +```python +def _get_wait_timeout(self) -> int: + multiplier = self._config.timeout_multiplier or 1.0 + agents = self._config.agents + if agents: + agent_timeout = agents[0].max_timeout_sec or agents[0].override_timeout_sec + if agent_timeout: + return int(agent_timeout * multiplier) + 600 # +600s 留给环境准备 / verifier + return int(DEFAULT_WAIT_TIMEOUT * multiplier) # 默认 7200s +``` + +**新逻辑** (`executor.py:115`)：`wait_timeout=config.timeout`，默认 `3600s`。 + +**影响**：任何配置了 `agent.max_timeout_sec > 3000` 的 Harbor job 都会被早杀；`timeout_multiplier`、`agent_timeout_multiplier` 等 5 个字段完全失效。 + +**修复**：在 `HarborJobConfig` 上加 `model_post_init`，把计算好的有效超时写回 `self.timeout`；或在 `HarborTrial.setup` 里动态调整；或 `JobExecutor` 允许 Trial override timeout（需接口扩展）。 + +### G3 — `job_name` 自动生成丢失 + +老：`{dataset_name}_{task_name if 单任务}_{uuid[:8]}`（`bench/job.py:273`）。 + +新：若用户不设 `job_name`，后续脚本路径、session 名全部使用 `"default"`，多实例并发会冲突。 + +修复：把 `_generate_default_job_name` 作为 `HarborJobConfig` 的 `model_validator`，或放进 `HarborTrial.setup`。 + +### G4 — `namespace` / `experiment_id` 从 sandbox 回填丢失 + +老 `_autofill_sandbox_info` 读取 `sandbox._namespace` / `sandbox._experiment_id` 校验并写回 config（L303）。新代码只在 config 构造时做 `_sync_experiment_id` 校验，sandbox 侧的真实值永远不会回传。 + +修复：在 `JobExecutor._do_submit` 拿到 sandbox 句柄后插入一次回填；或提供 `AbstractTrial.on_sandbox_ready(sandbox)` 钩子。 + +### G5 — `JobResult.raw_output` / `exit_code` 不被填充 + +`JobResult` 的字段存在，但新 Facade `_build_result` 从未写入（`api.py:67`）。 + +修复：`JobExecutor.wait` 把每个 Trial 的 `obs.output` / `obs.exit_code` 也返回；或在 `TrialResult` 上挂这两个字段，由 Job facade 聚合。 + +### G6 — 脚本 / 输出文件路径从持久化目录变成 `/tmp` + +| | 脚本路径 | 输出路径 | Harbor YAML | +|---|----------|----------|-------------| +| 老 | `/data/logs/user-defined/rock_job_{name}.sh` | 同上 `.out` | `/data/logs/user-defined/rock_job_{name}.yaml` | +| 新 | `/tmp/rock_job_{name}.sh` | `/tmp/rock_job_{name}.out` | `/data/logs/user-defined/rock_job_{name}.yaml` | + +`/tmp` 在容器重启后消失，线上调查失败 job 会更难。 + +修复：新 `JobExecutor._job_tmp_prefix` 改用 `USER_DEFINED_LOGS`，保持与 Harbor YAML 一致。 + +### G7 — `auto_stop` 字段迁移 + +老：读 `config.environment.auto_stop`（`RockEnvironmentConfig` 字段）。 +新：读 `config.auto_stop`（`JobConfig` 基类字段）。 + +现状 `HarborJobConfig` 继承后两者同时存在，用户写在 `environment` 里的设置新架构会忽略。 + +修复：`HarborJobConfig` 加 `model_validator` 同步两个字段；或文档明确指引新字段。 + +--- + +## 4. 功能映射矩阵 + +| 功能 | 老位置 | 新位置 | 一致性 | +|------|--------|--------|--------| +| Facade `run/submit/wait/cancel` | `bench/job.py` `Job` 类 | `job/api.py` `Job` 类 | ✅ | +| Sandbox 启动 | `submit()` | `JobExecutor._do_submit` | ✅ | +| Session + OSS env | `_create_session` / `_build_session_env` | `JobExecutor._do_submit` / `_build_session_env` | ✅ | +| file_uploads 上传 | `_prepare_and_start` 循环 | `AbstractTrial._upload_files` | ✅ | +| Harbor YAML 上传 | `_upload_content` | `HarborTrial.setup` | ✅ | +| dockerd + setup + harbor run 脚本 | `_render_run_script` + `_RUN_SCRIPT_TEMPLATE` | `HarborTrial.build` + `_HARBOR_SCRIPT_TEMPLATE` | ✅ 模板字节级一致 | +| nohup 启动 | `start_nohup_process` | `JobExecutor._do_submit` | ✅ | +| nohup 等待 | `wait_for_process_completion` | `JobExecutor._do_wait` | ✅ | +| 子 trial 结果收集 | `_collect_results` | `HarborTrial._collect_trial_results` | ⚠️ 聚合方式不同（G1）| +| job_name 默认生成 | `_generate_default_job_name` | — | ❌ G3 | +| agent-aware 超时 | `_get_wait_timeout` | — | ❌ G2 | +| ns/exp_id 回填 | `_autofill_sandbox_info` | — | ❌ G4 | +| raw_output / exit_code | 手动写回 | — | ❌ G5 | +| script/out 路径 | `USER_DEFINED_LOGS` | `/tmp` | ⚠️ G6 | +| auto_stop | `environment.auto_stop` | `config.auto_stop` | ⚠️ G7 | + +--- + +## 5. 替换路线图 + +``` +Phase 1 — 补齐回归 (必做) + 1. 修 G1: 改 AbstractTrial.collect → TrialResult | list[TrialResult]（union）， + Job._build_result 按 isinstance 判断 extend/append 到 JobResult.trial_results + 2. 修 G2: HarborJobConfig.model_post_init 计算有效 timeout，或 Trial 层 override + 3. 修 G3: _generate_default_job_name 挪到 HarborJobConfig validator + 4. 修 G4: JobExecutor 调用 on_sandbox_ready 钩子 + AbstractTrial 提供回填默认实现 + 5. 修 G5: TrialResult 加 raw_output/exit_code；_build_result 聚合 + 6. 修 G6: _job_tmp_prefix 改用 USER_DEFINED_LOGS + 7. 修 G7: HarborJobConfig 同步 auto_stop 两处字段 + +Phase 2 — 测试等价 + 1. tests/unit/sdk/agent/test_job.py 整套用 rock.sdk.job.Job 跑通 + 2. tests/unit/sdk/agent/test_jobconfig_experiment_id.py 同上 + 3. examples/harbor/harbor_demo.py 改用新 import 跑通 + +Phase 3 — 切换 + 标记 deprecated + 1. rock/sdk/bench/__init__.py 的 Job 指向 rock.sdk.job.Job + 2. rock/sdk/bench/job.py 整个文件标 DeprecationWarning，保留一个版本 + 3. 文档引导迁移 + +Phase 4 — 移除 + 1. 删 rock/sdk/bench/job.py + 2. bench/__init__.py 不再 re-export Job +``` + +--- + +## 6. 现有调用方影响面 + +``` +tests/unit/sdk/agent/test_job.py ← 大量 _build_session_env / _generate_default_job_name 私有 API 断言 +tests/unit/sdk/agent/test_jobconfig_experiment_id.py +tests/unit/sdk/agent/test_models.py +tests/unit/sdk/job/test_integration.py +examples/harbor/harbor_demo.py ← from rock.sdk.bench import HarborJobConfig, Job +rock/sdk/bench/__init__.py ← 顶层 re-export +``` + +Phase 1 完成后，`tests/unit/sdk/agent/test_job.py` 中基于私有方法 (`_build_session_env`, `_generate_default_job_name`) 的断言需要改写： +- `_build_session_env` → 测 `JobExecutor._build_session_env(config)` (已是 staticmethod) +- `_generate_default_job_name` → 测 `HarborJobConfig` 的 validator + +--- + +## 7. 结论 + +新架构从**工程质量、可扩展性、职责划分**角度都优于老实现；G1-G7 7 项回归/缺口均已在 #783 补齐，blue-green 等价测试锁死两路径契约。 + +**现阶段：** `rock/sdk/bench/Job` 已加 `DeprecationWarning`，仍保留作为兼容层；新代码应统一用 `rock.sdk.job.Job` + `HarborJobConfig`。计划 1.7.x 版本移除 blue 路径。 + +--- + +## 8. 修复状态（2026-04-14，PR #783） + +| Gap | Fix | 相关 commit | +|-----|-----|------------| +| G1 — 多子 trial 结果聚合 | `AbstractTrial.collect` 返回 union；`Job._build_result` 按 `isinstance` 拍平 | `fix(job/G1a)` `fix(job/G1b)` | +| G2 — effective wait timeout | `HarborJobConfig._compute_effective_timeout` validator | `fix(job/G2)` | +| G3 — 自动 job_name | `HarborJobConfig._auto_job_name` validator | `fix(job/G3)` | +| G4 — sandbox 回填 ns/exp_id | `AbstractTrial.on_sandbox_ready` 默认实现回填 + 校验（HarborTrial / BashTrial 共享） | `fix(job/G4)` `refactor(job): hoist to AbstractTrial` | +| G5 — `JobResult.raw_output/exit_code` | `TrialResult` 加字段；`JobExecutor._do_wait` 写回；`Job._build_result` 聚合 | `fix(job/G5)` | +| G6 — 脚本/输出持久化 | `_job_tmp_prefix` 使用 `USER_DEFINED_LOGS` | `fix(job/G6)` | +| G7 — `auto_stop` OR 同步 | `HarborJobConfig._sync_auto_stop` validator | `fix(job/G7)` | +| blue-green 等价锁 | 4 tests 交叉验证同一 config 两路径产出一致 `JobResult` | `test(job): blue-green equivalence` | + +测试：222 → 250 passing（+28，零回归）。 diff --git a/examples/harbor/harbor_demo.py b/examples/harbor/harbor_demo.py index f43b69d390..9984e09f68 100644 --- a/examples/harbor/harbor_demo.py +++ b/examples/harbor/harbor_demo.py @@ -1,15 +1,12 @@ -"""Harbor benchmark demo using ROCK Job SDK. +"""Harbor benchmark demo using ROCK Job SDK (new path). -Run Harbor benchmark tasks inside a ROCK sandbox via the Job SDK. -Configuration is loaded from a YAML file and passed to ``harbor jobs start`` -inside the sandbox. +Uses ``rock.sdk.job.Job`` with ``HarborJobConfig`` — the recommended path +with full feature parity (G1-G7 fixed) and scatter / multiple trial types. -Example config templates: - - ``examples/harbor/swe_job_config.yaml.template`` — SWE-bench-verified - - ``examples/harbor/tb_job_config.yaml.template`` — Terminal Bench 2 +For the legacy path (``rock.sdk.bench.Job``), see ``harbor_demo_legacy.py``. Usage: - python examples/harbor/harbor_demo.py -c examples/harbor/job_config.yaml + python examples/harbor/harbor_demo.py -c examples/harbor/swe.intern.yaml python examples/harbor/harbor_demo.py -c examples/harbor/tb_job_config.yaml -t mailman Required environment variables (OSS_* are auto-forwarded into the sandbox): @@ -32,7 +29,8 @@ import os import sys -from rock.sdk.bench import Job, JobConfig +from rock.sdk.bench import HarborJobConfig +from rock.sdk.job import Job _REQUIRED_ENV_VARS = [ "OSS_ACCESS_KEY_ID", @@ -62,13 +60,13 @@ def check_oss_env() -> None: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run Harbor tasks inside a ROCK sandbox") - parser.add_argument("-c", "--config", required=True, help="Path to JobConfig YAML file") + parser.add_argument("-c", "--config", required=True, help="Path to HarborJobConfig YAML file") parser.add_argument("-t", "--task", default=None, help="Task name to run (overrides config)") return parser.parse_args() async def async_main(args: argparse.Namespace) -> None: - config = JobConfig.from_yaml(args.config) + config = HarborJobConfig.from_yaml(args.config) # Override task_names if specified via CLI if args.task and config.datasets: diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index b4c97c0cf1..2a5630e36a 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -62,7 +62,7 @@ async def _job_run(self, args: argparse.Namespace): if not args.config: logger.error("--config is required for harbor type") return - from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig config = HarborJobConfig.from_yaml(args.config) if args.image: diff --git a/rock/sdk/bench/__init__.py b/rock/sdk/bench/__init__.py index d989a5812b..72d6fe6450 100644 --- a/rock/sdk/bench/__init__.py +++ b/rock/sdk/bench/__init__.py @@ -1,6 +1,6 @@ from rock.sdk.bench.job import Job from rock.sdk.bench.models.job.config import ( - JobConfig, + HarborJobConfig, LocalDatasetConfig, OrchestratorConfig, OssRegistryInfo, @@ -22,7 +22,7 @@ AgentInfo, AgentResult, ExceptionInfo, - TrialResult, + HarborTrialResult, VerifierResult, ) from rock.sdk.job.result import JobResult, JobStatus @@ -31,12 +31,12 @@ "Job", "JobResult", "JobStatus", - "TrialResult", + "HarborTrialResult", "VerifierResult", "AgentInfo", "AgentResult", "ExceptionInfo", - "JobConfig", + "HarborJobConfig", "RockEnvironmentConfig", "RegistryDatasetConfig", "LocalDatasetConfig", diff --git a/rock/sdk/bench/job.py b/rock/sdk/bench/job.py index d89dea6a47..92c960275a 100644 --- a/rock/sdk/bench/job.py +++ b/rock/sdk/bench/job.py @@ -15,7 +15,7 @@ from rock.actions import Command, CreateBashSessionRequest, ReadFileRequest from rock.logger import init_logger from rock.sdk.bench.constants import CHECK_INTERVAL, DEFAULT_WAIT_TIMEOUT, USER_DEFINED_LOGS -from rock.sdk.bench.models.trial.result import TrialResult +from rock.sdk.bench.models.trial.result import HarborTrialResult from rock.sdk.job.result import JobResult, JobStatus logger = init_logger(__name__) @@ -63,10 +63,20 @@ class Job: """ def __init__(self, config): - from rock.sdk.bench.models.job.config import JobConfig + import warnings + + warnings.warn( + "rock.sdk.bench.Job is deprecated and will be removed in 1.7.x. " + "Use rock.sdk.job.Job with HarborJobConfig — the new path has full " + "feature parity (G1-G7 fixed) and supports scatter / multiple trial types.", + DeprecationWarning, + stacklevel=2, + ) + + from rock.sdk.bench.models.job.config import HarborJobConfig - if not isinstance(config, JobConfig): - raise TypeError(f"config must be JobConfig, got {type(config)}") + if not isinstance(config, HarborJobConfig): + raise TypeError(f"config must be HarborJobConfig, got {type(config)}") self._config = config self._sandbox = None self._session: str | None = None @@ -250,12 +260,12 @@ async def _collect_results(self) -> JobResult: trial_result_files = [] # Parse each trial result - trial_results: list[TrialResult] = [] + trial_results: list[HarborTrialResult] = [] for trial_file in trial_result_files: try: response = await self._sandbox.read_file(ReadFileRequest(path=trial_file)) data = json.loads(response.content) - trial_results.append(TrialResult.from_harbor_json(data)) + trial_results.append(HarborTrialResult.from_harbor_json(data)) except Exception as e: logger.warning(f"Failed to parse trial result {trial_file}: {e}") diff --git a/rock/sdk/bench/models/__init__.py b/rock/sdk/bench/models/__init__.py index c0ce31d8f8..ebe96f4488 100644 --- a/rock/sdk/bench/models/__init__.py +++ b/rock/sdk/bench/models/__init__.py @@ -1,7 +1,7 @@ from rock.sdk.bench.models.environment_type import EnvironmentType from rock.sdk.bench.models.job.config import ( DatasetConfig, - JobConfig, + HarborJobConfig, OrchestratorConfig, RetryConfig, ) @@ -19,7 +19,7 @@ ) __all__ = [ - "JobConfig", + "HarborJobConfig", "OrchestratorConfig", "RetryConfig", "DatasetConfig", diff --git a/rock/sdk/bench/models/job/__init__.py b/rock/sdk/bench/models/job/__init__.py index a30c186b62..f564fe36ca 100644 --- a/rock/sdk/bench/models/job/__init__.py +++ b/rock/sdk/bench/models/job/__init__.py @@ -2,7 +2,7 @@ from rock.sdk.job.result import JobResult, JobStatus from .config import ( - JobConfig, + HarborJobConfig, LocalDatasetConfig, OrchestratorConfig, OssRegistryInfo, @@ -12,7 +12,7 @@ ) __all__ = [ - "JobConfig", + "HarborJobConfig", "OrchestratorConfig", "RetryConfig", "RegistryDatasetConfig", diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index 080f7e3063..94f13c242f 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -128,7 +128,7 @@ def _infer_version_from_split(self): DatasetConfig = LocalDatasetConfig | RegistryDatasetConfig -class JobConfig(_BaseJobConfig): +class HarborJobConfig(_BaseJobConfig): """Harbor Job configuration: extends base JobConfig with Harbor-native fields. All Rock sandbox/lifecycle configuration lives in ``environment`` (inherited). @@ -172,27 +172,102 @@ def _sync_experiment_id(self): self.environment.experiment_id = self.experiment_id return self + @model_validator(mode="after") + def _sync_auto_stop(self): + """G7: keep top-level auto_stop and environment.auto_stop in sync (OR semantics). + + Users may set either. Legacy ``environment.auto_stop=True`` (pre-job-refactor) + must still work; new ``config.auto_stop=True`` must also propagate down to + the environment so the RockEnvironmentConfig path reads the same value. + """ + effective = bool(self.auto_stop) or bool(self.environment.auto_stop) + self.auto_stop = effective + self.environment.auto_stop = effective + return self + + @model_validator(mode="after") + def _auto_job_name(self): + """G3: auto-generate job_name when user omitted it. + + Format: {dataset_name}_{task_name if single task}_{uuid[:8]} + Matches legacy bench/job.py::_generate_default_job_name. + """ + import uuid as _uuid + + if self.job_name is not None: + return self + + parts: list[str] = [] + if self.datasets: + ds = self.datasets[0] + if getattr(ds, "name", None): + parts.append(ds.name) + task_names = getattr(ds, "task_names", None) or [] + if len(task_names) == 1: + parts.append(task_names[0]) + + parts.append(_uuid.uuid4().hex[:8]) + self.job_name = "_".join(parts) + return self + + @model_validator(mode="after") + def _compute_effective_timeout(self): + """G2: derive wait timeout from agent config × multiplier + buffer. + + Rule (aligned with legacy bench/job.py::_get_wait_timeout): + agent_timeout = agents[0].max_timeout_sec or agents[0].override_timeout_sec + effective = int(agent_timeout * multiplier) + 600 (env + verifier buffer) + fallback = int(DEFAULT_WAIT_TIMEOUT * multiplier) (7200 * multiplier) + + Applied only when the base-class default (3600) has not been overridden by + the user. If the user explicitly set ``timeout`` to a non-default value, + that wins — we do not second-guess an explicit knob. NOTE: this heuristic + misfires if a user explicitly picks 3600, but that's considered extremely + rare; documented limitation. + """ + from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT + + # 3600 is the base JobConfig default; treat as "user didn't touch it". + if self.timeout != 3600: + return self + + multiplier = self.timeout_multiplier or 1.0 + agent_timeout: float | None = None + if self.agents: + a = self.agents[0] + agent_timeout = a.max_timeout_sec or a.override_timeout_sec + + if agent_timeout: + self.timeout = int(agent_timeout * multiplier) + 600 + else: + self.timeout = int(DEFAULT_WAIT_TIMEOUT * multiplier) + return self + # Base JobConfig fields to exclude when serializing to Harbor YAML _BASE_FIELDS: ClassVar[set[str]] = set(_BaseJobConfig.model_fields.keys()) def to_harbor_yaml(self) -> str: """Serialize Harbor-native fields to YAML for ``harbor jobs start -c``. - Base JobConfig fields (environment, job_name, setup_commands, etc.) - are excluded. Harbor environment fields (force_build, override_cpus, etc.) + Base JobConfig fields (environment, setup_commands, etc.) are excluded. + ``job_name`` is re-injected so harbor uses it as the job directory name + instead of its default timestamp-based naming. + Harbor environment fields (force_build, override_cpus, etc.) are re-injected under ``environment``. """ import yaml data = self.model_dump(mode="json", exclude=self._BASE_FIELDS, exclude_none=True) + if self.job_name: + data["job_name"] = self.job_name harbor_env = self.environment.to_harbor_environment() if harbor_env: data["environment"] = harbor_env return yaml.dump(data, default_flow_style=False, allow_unicode=True) @classmethod - def from_yaml(cls, path: str) -> JobConfig: - """Load JobConfig from a Harbor YAML config file.""" + def from_yaml(cls, path: str) -> HarborJobConfig: + """Load HarborJobConfig from a Harbor YAML config file.""" import yaml with open(path) as f: diff --git a/rock/sdk/bench/models/trial/__init__.py b/rock/sdk/bench/models/trial/__init__.py index 24932bab46..a5fa58d8a3 100644 --- a/rock/sdk/bench/models/trial/__init__.py +++ b/rock/sdk/bench/models/trial/__init__.py @@ -1,5 +1,5 @@ from .config import AgentConfig, ArtifactConfig, EnvironmentConfig, OssMirrorConfig, TaskConfig, VerifierConfig -from .result import AgentInfo, AgentResult, ExceptionInfo, ModelInfo, TimingInfo, TrialResult, VerifierResult +from .result import AgentInfo, AgentResult, ExceptionInfo, HarborTrialResult, ModelInfo, TimingInfo, VerifierResult __all__ = [ "AgentConfig", @@ -8,7 +8,7 @@ "VerifierConfig", "TaskConfig", "ArtifactConfig", - "TrialResult", + "HarborTrialResult", "AgentInfo", "ModelInfo", "AgentResult", diff --git a/rock/sdk/bench/models/trial/result.py b/rock/sdk/bench/models/trial/result.py index 651dffc16d..cc2d5779e6 100644 --- a/rock/sdk/bench/models/trial/result.py +++ b/rock/sdk/bench/models/trial/result.py @@ -1,7 +1,7 @@ """Harbor trial result models. -TrialResult base class is in rock.sdk.job.result. -This module extends it with Harbor-specific fields. +The base TrialResult lives in rock.sdk.job.result. +This module extends it with Harbor-specific fields as HarborTrialResult. """ from __future__ import annotations @@ -42,8 +42,8 @@ class TimingInfo(BaseModel): finished_at: str | None = None -class TrialResult(_BaseTrialResult): - """Harbor TrialResult: extends base with agent/verifier/timing fields.""" +class HarborTrialResult(_BaseTrialResult): + """Harbor-specific TrialResult: extends the base with agent/verifier/timing fields.""" trial_name: str = "" source: str | None = None @@ -75,8 +75,8 @@ def token_ids(self) -> list[int]: return [] @classmethod - def from_harbor_json(cls, data: dict[str, Any]) -> TrialResult: - """Parse a harbor trial-level result.json dict into TrialResult.""" + def from_harbor_json(cls, data: dict[str, Any]) -> HarborTrialResult: + """Parse a harbor trial-level result.json dict into HarborTrialResult.""" exception_info = None if data.get("exception_info"): ei = data["exception_info"] diff --git a/rock/sdk/job/api.py b/rock/sdk/job/api.py index babc3b2723..97834e2cc5 100644 --- a/rock/sdk/job/api.py +++ b/rock/sdk/job/api.py @@ -55,8 +55,8 @@ async def wait(self) -> JobResult: """Wait for completion, build JobResult.""" if not self._job_client: raise RuntimeError("No submitted job. Call submit() first.") - trial_results = await self._executor.wait(self._job_client) - return self._build_result(trial_results) + raw = await self._executor.wait(self._job_client) + return self._build_result(raw) async def cancel(self) -> None: """Kill all running trials.""" @@ -64,11 +64,28 @@ async def cancel(self) -> None: for tc in self._job_client.trials: await tc.sandbox.arun(cmd=f"kill {tc.pid}", session=tc.session) - def _build_result(self, trial_results: list[TrialResult]) -> JobResult: - all_success = all(r.exception_info is None for r in trial_results) + def _build_result(self, raw_results: list[TrialResult | list[TrialResult]]) -> JobResult: + """Flatten list-returning collect() outputs into JobResult.trial_results. + + Each element of ``raw_results`` is whatever one Trial's ``collect()`` + returned — either a single TrialResult or a list. HarborTrial returns + a list (one entry per sub-trial); BashTrial returns a single result. + """ + flat: list[TrialResult] = [] + for r in raw_results: + if isinstance(r, list): + flat.extend(r) + else: + flat.append(r) + all_success = all(t.exception_info is None for t in flat) + # G5: surface first non-empty output / non-zero exit code from sub-trials + raw_output = next((t.raw_output for t in flat if t.raw_output), "") + exit_code = next((t.exit_code for t in flat if t.exit_code != 0), 0) return JobResult( job_id=self._config.job_name or "", status=JobStatus.COMPLETED if all_success else JobStatus.FAILED, labels=self._config.labels, - trial_results=trial_results, + trial_results=flat, + raw_output=raw_output, + exit_code=exit_code, ) diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py index 011f0f0d71..172002e60c 100644 --- a/rock/sdk/job/executor.py +++ b/rock/sdk/job/executor.py @@ -47,7 +47,7 @@ class JobClient: class JobExecutor: """Execution engine: drives Operator to generate trials, runs in parallel, collects results.""" - async def run(self, operator: Operator, config: JobConfig) -> list[TrialResult]: + async def run(self, operator: Operator, config: JobConfig) -> list[TrialResult | list[TrialResult]]: """Full lifecycle: submit + wait.""" job_client = await self.submit(operator, config) return await self.wait(job_client) @@ -60,8 +60,13 @@ async def submit(self, operator: Operator, config: JobConfig) -> JobClient: trial_clients = await asyncio.gather(*[self._do_submit(t) for t in trial_list]) return JobClient(trials=list(trial_clients)) - async def wait(self, job_client: JobClient) -> list[TrialResult]: - """Wait for all trials, collect results in parallel.""" + async def wait(self, job_client: JobClient) -> list[TrialResult | list[TrialResult]]: + """Wait for all trials, collect results in parallel. + + Each entry mirrors whatever the Trial's ``collect()`` returned + (single ``TrialResult`` or ``list[TrialResult]``). The Job layer + flattens lists into the final ``JobResult.trial_results``. + """ if not job_client.trials: return [] return list(await asyncio.gather(*[self._do_wait(tc) for tc in job_client.trials])) @@ -70,8 +75,15 @@ async def wait(self, job_client: JobClient) -> list[TrialResult]: @staticmethod def _job_tmp_prefix(config: JobConfig) -> str: - """Prefix for per-job temp files on sandbox, e.g. /tmp/rock_job_my-job.""" - return f"/tmp/rock_job_{config.job_name or 'default'}" + """Prefix for per-job script/output files on sandbox. + + Uses USER_DEFINED_LOGS (persistent under /data/logs/user-defined/...) + so that logs survive sandbox-internal /tmp sweeps and are inspectable + after a failure, matching the legacy bench/job.py behavior. + """ + from rock.sdk.bench.constants import USER_DEFINED_LOGS + + return f"{USER_DEFINED_LOGS}/rock_job_{config.job_name or 'default'}" async def _do_submit(self, trial: AbstractTrial) -> TrialClient: """Start sandbox + execute script for a single trial.""" @@ -80,6 +92,9 @@ async def _do_submit(self, trial: AbstractTrial) -> TrialClient: await sandbox.start() logger.info(f"Sandbox started: sandbox_id={sandbox.sandbox_id}, job_name={config.job_name}") + # G4: let trial backfill config from sandbox state before setup + await trial.on_sandbox_ready(sandbox) + session = f"rock-job-{config.job_name or 'default'}" env = self._build_session_env(config) await sandbox.create_session(CreateBashSessionRequest(session=session, env_enable=True, env=env)) @@ -103,7 +118,7 @@ async def _do_submit(self, trial: AbstractTrial) -> TrialClient: logger.info(f"Trial started: pid={pid}, job_name={config.job_name}") return TrialClient(sandbox=sandbox, session=session, pid=pid, trial=trial) - async def _do_wait(self, client: TrialClient) -> TrialResult: + async def _do_wait(self, client: TrialClient) -> TrialResult | list[TrialResult]: """Wait for a single trial to finish, call trial.collect().""" from rock.sdk.job.result import ExceptionInfo @@ -125,11 +140,21 @@ async def _do_wait(self, client: TrialClient) -> TrialResult: ) exit_code = obs.exit_code if obs.exit_code is not None else 1 result = await client.trial.collect(client.sandbox, obs.output or "", exit_code) - if not success and result.exception_info is None: - result.exception_info = ExceptionInfo( + # G5: populate raw_output / exit_code on every TrialResult so they surface in JobResult + iter_results = result if isinstance(result, list) else [result] + for r in iter_results: + if not r.raw_output: + r.raw_output = obs.output or "" + if r.exit_code == 0 and exit_code != 0: + r.exit_code = exit_code + if not success: + fail_info = ExceptionInfo( exception_type="ProcessTimeout", exception_message=message or "process did not complete successfully", ) + for r in iter_results: + if r.exception_info is None: + r.exception_info = fail_info return result finally: if config.auto_stop: diff --git a/rock/sdk/job/result.py b/rock/sdk/job/result.py index b13adef858..8bf1950481 100644 --- a/rock/sdk/job/result.py +++ b/rock/sdk/job/result.py @@ -37,6 +37,9 @@ class TrialResult(BaseModel): exception_info: ExceptionInfo | None = None started_at: str | None = None finished_at: str | None = None + # G5: process-level outputs captured by JobExecutor + raw_output: str = "" + exit_code: int = 0 @property def score(self) -> float: diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index fc5f4e092f..65159589cd 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -23,6 +23,33 @@ class AbstractTrial(ABC): def __init__(self, config: JobConfig): self._config = config + async def on_sandbox_ready(self, sandbox: Sandbox) -> None: + """G4 hook: called by JobExecutor once sandbox.start() succeeds, before setup(). + + Default behavior backfills ``namespace`` and ``experiment_id`` from the + sandbox into ``self._config`` (both are fields on ``JobConfig``), and + raises ``ValueError`` if the sandbox reports a value that conflicts + with one already set on the config. Matches legacy + ``_autofill_sandbox_info``. Subclasses can override to extend. + """ + sb_ns = getattr(sandbox, "_namespace", None) + if sb_ns is not None: + if self._config.namespace is not None and self._config.namespace != sb_ns: + raise ValueError( + f"namespace mismatch: {type(self._config).__name__} has " + f"'{self._config.namespace}', but sandbox returned '{sb_ns}'" + ) + self._config.namespace = sb_ns + + sb_exp = getattr(sandbox, "_experiment_id", None) + if sb_exp is not None: + if self._config.experiment_id is not None and self._config.experiment_id != sb_exp: + raise ValueError( + f"experiment_id mismatch: {type(self._config).__name__} has " + f"'{self._config.experiment_id}', but sandbox returned '{sb_exp}'" + ) + self._config.experiment_id = sb_exp + @abstractmethod async def setup(self, sandbox: Sandbox) -> None: """Pre-execution: prepare sandbox environment (upload files, write configs).""" @@ -32,8 +59,15 @@ def build(self) -> str: """Build: generate bash script to execute.""" @abstractmethod - async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: - """Post-execution: collect and parse results.""" + async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult | list[TrialResult]: + """Post-execution: collect and parse results. + + Return a single ``TrialResult`` for one-shot tasks (e.g. BashTrial), + or a ``list[TrialResult]`` when the underlying tool produces multiple + sub-results per sandbox invocation (e.g. HarborTrial running a dataset + over N tasks). The Job / JobExecutor layer flattens lists into the + final ``JobResult.trial_results``. + """ async def _upload_files(self, sandbox: Sandbox) -> None: """Shared helper: upload all entries in ``config.file_uploads``.""" diff --git a/rock/sdk/job/trial/harbor.py b/rock/sdk/job/trial/harbor.py index 2fb242564c..4c8d31b485 100644 --- a/rock/sdk/job/trial/harbor.py +++ b/rock/sdk/job/trial/harbor.py @@ -12,10 +12,9 @@ from rock.actions import Command, ReadFileRequest from rock.logger import init_logger from rock.sdk.bench.constants import USER_DEFINED_LOGS -from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig -from rock.sdk.bench.models.trial.result import TrialResult -from rock.sdk.job.result import ExceptionInfo -from rock.sdk.job.result import TrialResult as BaseTrialResult +from rock.sdk.bench.models.job.config import HarborJobConfig +from rock.sdk.bench.models.trial.result import HarborTrialResult +from rock.sdk.job.result import ExceptionInfo, TrialResult from rock.sdk.job.trial.abstract import AbstractTrial from rock.sdk.job.trial.registry import register_trial @@ -75,21 +74,30 @@ def build(self) -> str: user_defined_dir=USER_DEFINED_LOGS, ) - async def collect(self, sandbox, output: str, exit_code: int) -> BaseTrialResult: + async def collect(self, sandbox, output: str, exit_code: int) -> list[TrialResult]: + """Return all Harbor sub-trial results (one entry per ``result.json``). + + Harbor writes N trial-level ``result.json`` files per sandbox run + (one per dataset × task). We return them all so the Job layer can + surface every sub-trial in ``JobResult.trial_results``. If Harbor + crashed before any trial finished, return a single synthetic failure + entry so that the caller can tell something ran. + """ trial_results = await self._collect_trial_results(sandbox) if trial_results: - return trial_results[0] - - exception_info = ExceptionInfo( - exception_type="HarborNoTrials", - exception_message="No trial results found", - ) - return BaseTrialResult( - task_name=self._config.job_name or "", - exception_info=exception_info, - ) + return list(trial_results) + + return [ + TrialResult( + task_name=self._config.job_name or "", + exception_info=ExceptionInfo( + exception_type="HarborNoTrials", + exception_message="No trial results found", + ), + ) + ] - async def _collect_trial_results(self, sandbox) -> list[TrialResult]: + async def _collect_trial_results(self, sandbox) -> list[HarborTrialResult]: """Read trial-level result.json files from sandbox.""" job_dir = f"{self._config.jobs_dir}/{self._config.job_name}" try: @@ -100,12 +108,12 @@ async def _collect_trial_results(self, sandbox) -> list[TrialResult]: except Exception: trial_files = [] - results: list[TrialResult] = [] + results: list[HarborTrialResult] = [] for trial_file in trial_files: try: response = await sandbox.read_file(ReadFileRequest(path=trial_file)) data = json.loads(response.content) - results.append(TrialResult.from_harbor_json(data)) + results.append(HarborTrialResult.from_harbor_json(data)) except Exception as e: logger.warning(f"Failed to parse trial result {trial_file}: {e}") diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index cf3bb83b68..b2a6b7624e 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -5,14 +5,14 @@ from rock.sdk.bench.job import Job, JobResult, JobStatus from rock.sdk.bench.models.job.config import ( - JobConfig, + HarborJobConfig, LocalDatasetConfig, RegistryDatasetConfig, RemoteRegistryInfo, RockEnvironmentConfig, ) from rock.sdk.bench.models.trial.config import AgentConfig -from rock.sdk.bench.models.trial.result import ExceptionInfo, TrialResult, VerifierResult +from rock.sdk.bench.models.trial.result import ExceptionInfo, HarborTrialResult, VerifierResult class TestJobStatus: @@ -24,9 +24,9 @@ def test_values(self): assert JobStatus.CANCELLED == "cancelled" -class TestTrialResult: +class TestHarborTrialResult: def test_defaults(self): - t = TrialResult(task_name="fix-bug") + t = HarborTrialResult(task_name="fix-bug") assert t.task_name == "fix-bug" assert t.score == 0.0 # computed property from verifier_result assert t.status == "completed" # computed property, no exception_info @@ -34,14 +34,14 @@ def test_defaults(self): assert t.duration_sec == 0.0 # computed property def test_with_verifier_result(self): - t = TrialResult( + t = HarborTrialResult( task_name="fix-bug", verifier_result=VerifierResult(rewards={"reward": 1.0}), ) assert t.score == 1.0 def test_failed_trial(self): - t = TrialResult( + t = HarborTrialResult( task_name="fix-bug", exception_info=ExceptionInfo( exception_type="TimeoutError", @@ -60,7 +60,7 @@ def test_from_harbor_json(self): "agent_result": {"n_input_tokens": 15000, "n_output_tokens": 3000}, "exception_info": None, } - t = TrialResult.from_harbor_json(data) + t = HarborTrialResult.from_harbor_json(data) assert t.task_name == "fix-dockerfile" assert t.trial_name == "trial-001" assert t.score == 1.0 @@ -73,8 +73,8 @@ def test_basic(self): job_id="job-123", status=JobStatus.COMPLETED, trial_results=[ - TrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), - TrialResult(task_name="t2", verifier_result=VerifierResult(rewards={"reward": 0.5})), + HarborTrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), + HarborTrialResult(task_name="t2", verifier_result=VerifierResult(rewards={"reward": 0.5})), ], raw_output="", exit_code=0, @@ -89,8 +89,8 @@ def test_score_with_failed_trials(self): job_id="job-456", status=JobStatus.COMPLETED, trial_results=[ - TrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), - TrialResult( + HarborTrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), + HarborTrialResult( task_name="t2", exception_info=ExceptionInfo(exception_type="Error", exception_message="err"), ), @@ -117,7 +117,7 @@ def test_labels_preserved(self): job_id="job-labeled", labels={"step": "42", "env": "prod"}, trial_results=[ - TrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), + HarborTrialResult(task_name="t1", verifier_result=VerifierResult(rewards={"reward": 1.0})), ], ) assert r.labels == {"step": "42", "env": "prod"} @@ -174,7 +174,7 @@ def _make_mock_sandbox(): class TestJob: def test_init_requires_jobconfig(self): - config = JobConfig(experiment_id="test-exp") + config = HarborJobConfig(experiment_id="test-exp") job = Job(config) assert job._config == config @@ -188,7 +188,7 @@ async def test_run_full_lifecycle(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig( + config = HarborJobConfig( experiment_id="test-exp", job_name="test-job", agents=[AgentConfig(name="t2")], @@ -211,7 +211,7 @@ async def test_run_auto_stop_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig( + config = HarborJobConfig( job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=True) ) job = Job(config) @@ -223,7 +223,7 @@ async def test_run_does_not_stop_when_disabled(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig( + config = HarborJobConfig( job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=False) ) job = Job(config) @@ -235,7 +235,7 @@ async def test_submit_starts_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job", experiment_id="test-exp") + config = HarborJobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() @@ -246,7 +246,7 @@ async def test_wait_returns_result(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job", experiment_id="test-exp") + config = HarborJobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() result = await job.wait() @@ -261,7 +261,7 @@ def test_oss_vars_from_process_env_are_forwarded(self, monkeypatch): monkeypatch.setenv("OSS_ACCESS_KEY_ID", "test-key") monkeypatch.setenv("HOME", "/root") - job = Job(JobConfig(job_name="test-job", experiment_id="test-exp")) + job = Job(HarborJobConfig(job_name="test-job", experiment_id="test-exp")) env = job._build_session_env() assert env["OSS_ENDPOINT"] == "https://oss.example.com" @@ -272,7 +272,7 @@ def test_config_env_overrides_process_oss_vars(self, monkeypatch): monkeypatch.setenv("OSS_ENDPOINT", "https://oss.from.process.com") job = Job( - JobConfig( + HarborJobConfig( job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(env={"OSS_ENDPOINT": "https://oss.from.config.com"}), @@ -287,7 +287,7 @@ def test_returns_none_when_both_empty(self, monkeypatch): if key.startswith("OSS"): monkeypatch.delenv(key) - job = Job(JobConfig(job_name="test-job", experiment_id="test-exp")) + job = Job(HarborJobConfig(job_name="test-job", experiment_id="test-exp")) assert job._build_session_env() is None @@ -297,7 +297,7 @@ async def test_cancel_kills_process(self): mock_sandbox.arun = AsyncMock(return_value=MagicMock(output="", exit_code=0)) with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig(job_name="test-job", experiment_id="test-exp") + config = HarborJobConfig(job_name="test-job", experiment_id="test-exp") job = Job(config) await job.submit() await job.cancel() @@ -313,7 +313,7 @@ class TestGenerateDefaultJobName: def test_custom_job_name_not_overwritten(self): """User-set job_name should not be overwritten.""" - config = JobConfig( + config = HarborJobConfig( job_name="my-custom-job", experiment_id="test-exp", datasets=[RegistryDatasetConfig(registry=RemoteRegistryInfo(), name="tb", version="2.0")], @@ -325,7 +325,7 @@ def test_custom_job_name_not_overwritten(self): def test_job_name_generated_with_dataset_and_single_task(self): """Default job_name should be generated with dataset name and single task.""" - config = JobConfig( + config = HarborJobConfig( experiment_id="test-exp", datasets=[ RegistryDatasetConfig( @@ -349,7 +349,7 @@ def test_job_name_generated_with_dataset_and_single_task(self): def test_job_name_generated_with_dataset_multiple_tasks(self): """With multiple tasks, only dataset name and UUID should be used.""" - config = JobConfig( + config = HarborJobConfig( experiment_id="test-exp", datasets=[ RegistryDatasetConfig( @@ -372,7 +372,7 @@ def test_job_name_generated_with_dataset_multiple_tasks(self): def test_job_name_generated_without_dataset(self): """Without dataset, only UUID should be used.""" - config = JobConfig(experiment_id="test-exp") + config = HarborJobConfig(experiment_id="test-exp") job = Job(config) job._generate_default_job_name() @@ -382,7 +382,7 @@ def test_job_name_generated_without_dataset(self): def test_job_name_generated_with_dataset_no_name(self): """Dataset without name field should still work.""" - config = JobConfig( + config = HarborJobConfig( experiment_id="test-exp", datasets=[LocalDatasetConfig(path=Path("/data/tasks"))], ) @@ -398,7 +398,7 @@ async def test_submit_generates_job_name(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = JobConfig( + config = HarborJobConfig( experiment_id="test-exp", datasets=[ RegistryDatasetConfig( @@ -410,11 +410,12 @@ async def test_submit_generates_job_name(self): ) job = Job(config) - # job_name is None initially - assert config.job_name is None + # G3: job_name is auto-generated by the pydantic validator at construction time. + assert config.job_name is not None + assert config.job_name.startswith("my-dataset_my-task_") await job.submit() - # job_name should have been generated + # submit() must preserve the (already-generated) job_name assert job._config.job_name is not None assert job._config.job_name.startswith("my-dataset_my-task_") diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index dea3092b19..0b7914f41d 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -3,7 +3,7 @@ import yaml from rock.sdk.bench.models.job.config import ( - JobConfig, + HarborJobConfig, LocalDatasetConfig, RegistryDatasetConfig, RemoteRegistryInfo, @@ -100,9 +100,9 @@ def test_empty_config_excludes_rock_fields(self): assert "setup_commands" not in result -class TestJobConfigToHarborYaml: +class TestHarborJobConfigToHarborYaml: def test_basic_serialization(self): - cfg = JobConfig( + cfg = HarborJobConfig( job_name="test-job", experiment_id="test-exp", n_attempts=2, @@ -111,14 +111,15 @@ def test_basic_serialization(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - # Base fields (job_name, experiment_id, etc.) are excluded from harbor YAML - assert "job_name" not in data + # job_name is re-injected so harbor uses it as the directory name + assert data["job_name"] == "test-job" + # Other base fields (experiment_id, etc.) are excluded from harbor YAML assert "experiment_id" not in data assert data["n_attempts"] == 2 assert data["agents"][0]["name"] == "terminus-2" def test_excludes_rock_fields(self): - cfg = JobConfig( + cfg = HarborJobConfig( experiment_id="test-exp", environment=RockEnvironmentConfig( setup_commands=["pip install harbor"], @@ -144,7 +145,7 @@ def test_excludes_rock_fields(self): assert "environment" not in data or "setup_commands" not in data.get("environment", {}) def test_excludes_none_values(self): - cfg = JobConfig( + cfg = HarborJobConfig( job_name="test", experiment_id="test-exp", agents=[AgentConfig(name="t2")], @@ -155,8 +156,8 @@ def test_excludes_none_values(self): assert "agent_timeout_multiplier" not in data def test_labels_excluded_as_base_field(self): - """labels is a base JobConfig field, so it's excluded from harbor YAML.""" - cfg = JobConfig( + """labels is a base HarborJobConfig field, so it's excluded from harbor YAML.""" + cfg = HarborJobConfig( job_name="labeled-job", experiment_id="test-exp", labels={"step": "42", "env": "prod"}, @@ -165,10 +166,10 @@ def test_labels_excluded_as_base_field(self): data = yaml.safe_load(yaml_str) assert "labels" not in data - assert "job_name" not in data + assert data["job_name"] == "labeled-job" def test_path_fields_serialized_as_strings(self): - cfg = JobConfig( + cfg = HarborJobConfig( experiment_id="test-exp", jobs_dir=Path("/workspace/jobs"), tasks=[TaskConfig(path="/workspace/tasks/t1")], @@ -180,7 +181,7 @@ def test_path_fields_serialized_as_strings(self): assert data["tasks"][0]["path"] == "/workspace/tasks/t1" def test_harbor_env_fields_serialized(self): - cfg = JobConfig( + cfg = HarborJobConfig( job_name="full-test", experiment_id="test-exp", n_attempts=3, @@ -206,7 +207,7 @@ def test_harbor_env_fields_serialized(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - assert "job_name" not in data # base field excluded + assert data["job_name"] == "full-test" assert data["environment"]["type"] == "docker" assert data["environment"]["force_build"] is True assert data["environment"]["override_cpus"] == 4 @@ -217,7 +218,9 @@ def test_harbor_env_fields_serialized(self): def test_env_in_harbor_yaml(self): """env is passed to both sandbox session and harbor YAML.""" - cfg = JobConfig(experiment_id="test-exp", environment=RockEnvironmentConfig(env={"OPENAI_API_KEY": "sk-xxx"})) + cfg = HarborJobConfig( + experiment_id="test-exp", environment=RockEnvironmentConfig(env={"OPENAI_API_KEY": "sk-xxx"}) + ) yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) @@ -225,7 +228,7 @@ def test_env_in_harbor_yaml(self): assert data["environment"]["env"] == {"OPENAI_API_KEY": "sk-xxx"} -class TestJobConfigFromYaml: +class TestHarborJobConfigFromYaml: def test_from_yaml_basic(self, tmp_path): yaml_content = """ job_name: loaded-job @@ -243,7 +246,7 @@ def test_from_yaml_basic(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.job_name == "loaded-job" assert cfg.n_attempts == 2 assert cfg.agents[0].name == "terminus-2" @@ -268,7 +271,7 @@ def test_from_yaml_with_environment_block(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.environment.image == "my-image:latest" assert cfg.environment.memory == "32g" assert cfg.environment.env == {"OPENAI_API_KEY": "sk-xxx"} @@ -288,7 +291,7 @@ def test_from_yaml_with_local_dataset(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.job_name == "local-dataset-job" assert isinstance(cfg.datasets[0], LocalDatasetConfig) assert cfg.datasets[0].path == Path("/data/tasks") @@ -306,5 +309,5 @@ def test_from_yaml_with_labels(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.labels == {"step": "42", "env": "prod"} diff --git a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py index f6e48bdbd5..bc09046646 100644 --- a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py +++ b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py @@ -1,4 +1,4 @@ -"""Tests for JobConfig._sync_experiment_id model_validator and namespace consistency.""" +"""Tests for HarborJobConfig._sync_experiment_id model_validator and namespace consistency.""" from unittest.mock import AsyncMock @@ -6,7 +6,7 @@ from pydantic import ValidationError from rock.sdk.bench.job import Job -from rock.sdk.bench.models.job.config import JobConfig +from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.bench.models.trial.config import RockEnvironmentConfig @@ -14,23 +14,23 @@ class TestExperimentIdNotEmpty: def test_none_experiment_id_raises(self): """experiment_id=None (default) must raise ValidationError.""" with pytest.raises(ValidationError, match="experiment_id"): - JobConfig(job_name="test") + HarborJobConfig(job_name="test") def test_empty_string_experiment_id_raises(self): """experiment_id='' must raise ValidationError.""" with pytest.raises(ValidationError, match="experiment_id must not be empty"): - JobConfig(job_name="test", experiment_id="") + HarborJobConfig(job_name="test", experiment_id="") class TestExperimentIdConsistency: def test_env_none_syncs_from_jobconfig(self): - """When environment.experiment_id is None, it gets set from JobConfig.""" - cfg = JobConfig(job_name="test", experiment_id="exp-1") + """When environment.experiment_id is None, it gets set from HarborJobConfig.""" + cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") assert cfg.environment.experiment_id == "exp-1" def test_env_matches_jobconfig_passes(self): """When both are set and equal, no error.""" - cfg = JobConfig( + cfg = HarborJobConfig( job_name="test", experiment_id="exp-1", environment=RockEnvironmentConfig(experiment_id="exp-1"), @@ -39,9 +39,9 @@ def test_env_matches_jobconfig_passes(self): assert cfg.environment.experiment_id == "exp-1" def test_env_mismatch_raises(self): - """When environment.experiment_id differs from JobConfig.experiment_id, raise.""" + """When environment.experiment_id differs from HarborJobConfig.experiment_id, raise.""" with pytest.raises(ValidationError, match="experiment_id mismatch"): - JobConfig( + HarborJobConfig( job_name="test", experiment_id="exp-1", environment=RockEnvironmentConfig(experiment_id="exp-OTHER"), @@ -53,7 +53,7 @@ class TestAutofillNamespaceConsistency: async def test_namespace_autofilled_from_sandbox(self): """When user sets no namespace, sandbox value is used.""" - config = JobConfig(job_name="test", experiment_id="exp-1") + config = HarborJobConfig(job_name="test", experiment_id="exp-1") assert config.namespace is None job = Job(config) @@ -67,7 +67,7 @@ async def test_namespace_autofilled_from_sandbox(self): async def test_namespace_user_matches_sandbox(self): """When user namespace matches sandbox, no error.""" - config = JobConfig(job_name="test", experiment_id="exp-1", namespace="same-ns") + config = HarborJobConfig(job_name="test", experiment_id="exp-1", namespace="same-ns") job = Job(config) sandbox = AsyncMock() sandbox._namespace = "same-ns" @@ -79,7 +79,7 @@ async def test_namespace_user_matches_sandbox(self): async def test_namespace_mismatch_raises(self): """When user namespace differs from sandbox, raise ValueError.""" - config = JobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") + config = HarborJobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") job = Job(config) sandbox = AsyncMock() sandbox._namespace = "different-ns" @@ -91,7 +91,7 @@ async def test_namespace_mismatch_raises(self): async def test_namespace_user_set_sandbox_none(self): """When user sets namespace but sandbox returns None, keep user value.""" - config = JobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") + config = HarborJobConfig(job_name="test", experiment_id="exp-1", namespace="user-ns") job = Job(config) sandbox = AsyncMock() sandbox._namespace = None diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index 84e606cc0a..324877b5a0 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -2,7 +2,7 @@ from rock.sdk.bench.models.environment_type import EnvironmentType from rock.sdk.bench.models.job.config import ( - JobConfig, + HarborJobConfig, LocalDatasetConfig, OrchestratorConfig, RegistryDatasetConfig, @@ -218,9 +218,9 @@ def test_with_version(self): assert d.n_tasks == 50 -class TestJobConfig: +class TestHarborJobConfig: def test_defaults(self): - cfg = JobConfig(experiment_id="test-exp") + cfg = HarborJobConfig(experiment_id="test-exp") assert cfg.n_attempts == 1 assert cfg.timeout_multiplier == 1.0 assert cfg.debug is False @@ -233,14 +233,14 @@ def test_defaults(self): assert cfg.artifacts == [] def test_environment_defaults(self): - cfg = JobConfig(experiment_id="test-exp") + cfg = HarborJobConfig(experiment_id="test-exp") assert cfg.environment.setup_commands == [] assert cfg.environment.file_uploads == [] assert cfg.environment.env == {} assert cfg.environment.auto_stop is False def test_with_full_config(self): - cfg = JobConfig( + cfg = HarborJobConfig( job_name="test-job", experiment_id="test-exp", n_attempts=2, @@ -257,20 +257,20 @@ def test_with_full_config(self): class TestPublicAPI: def test_import_from_agent_package(self): - from rock.sdk.bench import Job, JobResult, JobStatus, TrialResult + from rock.sdk.bench import HarborTrialResult, Job, JobResult, JobStatus assert Job is not None assert JobResult is not None assert JobStatus is not None - assert TrialResult is not None + assert HarborTrialResult is not None def test_import_from_models_package(self): from rock.sdk.bench.models import ( AgentConfig, EnvironmentType, - JobConfig, + HarborJobConfig, ) - assert JobConfig is not None + assert HarborJobConfig is not None assert AgentConfig is not None assert EnvironmentType is not None diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py index 8fee1e2484..a306a93c03 100644 --- a/tests/unit/sdk/agent/test_oss_mirror.py +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -3,10 +3,10 @@ Covers: - OssMirrorConfig model fields and defaults - EnvironmentConfig.oss_mirror field -- JobConfig top-level namespace/experiment_id fields +- HarborJobConfig top-level namespace/experiment_id fields - to_harbor_yaml() serialization (namespace at top level) - from_yaml() deserialization -- enable_oss_mirror() convenience method on JobConfig +- enable_oss_mirror() convenience method on HarborJobConfig """ import yaml @@ -89,21 +89,21 @@ def test_set_oss_mirror_from_dict(self): # --------------------------------------------------------------------------- -# 3. JobConfig 顶层 namespace/experiment_id 字段 +# 3. HarborJobConfig 顶层 namespace/experiment_id 字段 # --------------------------------------------------------------------------- -class TestJobConfigNamespaceFields: +class TestHarborJobConfigNamespaceFields: def test_default_namespace_is_none(self): - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig(job_name="test", experiment_id="test-exp") + cfg = HarborJobConfig(job_name="test", experiment_id="test-exp") assert cfg.namespace is None def test_namespace_settable_at_top_level(self): - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig(job_name="test", namespace="team-rl", experiment_id="rl-step-42") + cfg = HarborJobConfig(job_name="test", namespace="team-rl", experiment_id="rl-step-42") assert cfg.namespace == "team-rl" assert cfg.experiment_id == "rl-step-42" @@ -115,11 +115,11 @@ def test_namespace_settable_at_top_level(self): class TestToHarborYamlOssMirror: def test_namespace_at_top_level_in_yaml(self): - """namespace/experiment_id 序列化为 JobConfig 顶层字段。""" - from rock.sdk.bench.models.job.config import JobConfig + """namespace/experiment_id 序列化为 HarborJobConfig 顶层字段。""" + from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.bench.models.trial.config import OssMirrorConfig, RockEnvironmentConfig - cfg = JobConfig( + cfg = HarborJobConfig( job_name="mirror-test", namespace="my-ns", experiment_id="exp-1", @@ -136,7 +136,7 @@ def test_namespace_at_top_level_in_yaml(self): ) data = yaml.safe_load(cfg.to_harbor_yaml()) - # namespace/experiment_id are base JobConfig fields, excluded from harbor YAML + # namespace/experiment_id are base HarborJobConfig fields, excluded from harbor YAML assert "namespace" not in data assert "experiment_id" not in data oss = data["environment"]["oss_mirror"] @@ -147,9 +147,9 @@ def test_namespace_at_top_level_in_yaml(self): def test_disabled_oss_mirror_excluded_from_yaml(self): """When oss_mirror is default (disabled), it should not clutter the YAML.""" - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig(job_name="no-mirror", experiment_id="test-exp") + cfg = HarborJobConfig(job_name="no-mirror", experiment_id="test-exp") data = yaml.safe_load(cfg.to_harbor_yaml()) env_data = data.get("environment", {}) @@ -164,7 +164,7 @@ def test_disabled_oss_mirror_excluded_from_yaml(self): class TestFromYamlOssMirror: def test_from_yaml_with_top_level_namespace(self, tmp_path): """新方式：namespace/experiment_id 在顶层。""" - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig yaml_content = """\ job_name: loaded-mirror @@ -182,7 +182,7 @@ def test_from_yaml_with_top_level_namespace(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.namespace == "yaml-ns" assert cfg.experiment_id == "yaml-exp" assert cfg.environment.oss_mirror.enabled is True @@ -190,7 +190,7 @@ def test_from_yaml_with_top_level_namespace(self, tmp_path): def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): """YAML 中 oss_mirror 内多余的 namespace 等字段由 Pydantic 忽略。""" - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig yaml_content = """\ job_name: compat-mirror @@ -209,7 +209,7 @@ def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.environment.oss_mirror.enabled is True assert cfg.environment.oss_mirror.oss_bucket == "yaml-bucket" dump = cfg.environment.oss_mirror.model_dump(exclude_none=True) @@ -217,7 +217,7 @@ def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): assert "experiment_id" not in dump def test_from_yaml_without_oss_mirror(self, tmp_path): - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig yaml_content = """\ job_name: no-mirror @@ -228,7 +228,7 @@ def test_from_yaml_without_oss_mirror(self, tmp_path): yaml_file = tmp_path / "config.yaml" yaml_file.write_text(yaml_content) - cfg = JobConfig.from_yaml(str(yaml_file)) + cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.environment.oss_mirror is None @@ -239,9 +239,9 @@ def test_from_yaml_without_oss_mirror(self, tmp_path): class TestEnableOssMirror: def test_enable_with_all_params(self): - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig(job_name="conv-test", experiment_id="test-exp") + cfg = HarborJobConfig(job_name="conv-test", experiment_id="test-exp") cfg.enable_oss_mirror( oss_bucket="conv-bucket", oss_access_key_id="ak-conv", @@ -254,9 +254,9 @@ def test_enable_with_all_params(self): def test_does_not_touch_namespace_or_experiment_id(self): """enable_oss_mirror 不修改顶层 namespace / experiment_id。""" - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig( + cfg = HarborJobConfig( job_name="no-touch-test", namespace="preset-ns", experiment_id="preset-exp", @@ -273,9 +273,9 @@ def test_does_not_touch_namespace_or_experiment_id(self): def test_enable_then_serialize_roundtrip(self): """to_harbor_yaml: 顶层 namespace / experiment_id 与 oss_mirror 独立设置。""" - from rock.sdk.bench.models.job.config import JobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig - cfg = JobConfig(job_name="roundtrip", namespace="rt-ns", experiment_id="rt-exp") + cfg = HarborJobConfig(job_name="roundtrip", namespace="rt-ns", experiment_id="rt-exp") cfg.enable_oss_mirror( oss_bucket="rt-bucket", oss_access_key_id="ak-rt", diff --git a/tests/unit/sdk/job/test_blue_green_equivalence.py b/tests/unit/sdk/job/test_blue_green_equivalence.py new file mode 100644 index 0000000000..df78a74c19 --- /dev/null +++ b/tests/unit/sdk/job/test_blue_green_equivalence.py @@ -0,0 +1,149 @@ +"""G1-G7 blue-green equivalence: rock.sdk.bench.Job vs rock.sdk.job.Job. + +Both paths must produce equivalent JobResult for the same HarborJobConfig + +identical mock sandbox behavior. This locks in the contract across the +deprecation window. +""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import rock.sdk.bench # noqa: F401 +from rock.sdk.bench.models.job.config import ( + HarborJobConfig, + RegistryDatasetConfig, + RemoteRegistryInfo, +) +from rock.sdk.bench.models.trial.config import AgentConfig, RockEnvironmentConfig + + +def _build_mock_sandbox(): + sandbox = AsyncMock() + sandbox.sandbox_id = "sb-bg" + sandbox._namespace = "bg-ns" + sandbox._experiment_id = "bg-exp" + sandbox.start = AsyncMock() + sandbox.close = AsyncMock() + sandbox.create_session = AsyncMock() + sandbox.upload_by_path = AsyncMock(return_value=MagicMock(success=True)) + sandbox.write_file_by_path = AsyncMock() + sandbox.fs = AsyncMock() + sandbox.fs.upload_dir = AsyncMock(return_value=MagicMock(exit_code=0)) + + # find lists 2 result.json files + find_obs = MagicMock() + find_obs.stdout = ( + "/data/logs/user-defined/jobs/bg-exp/trial-0/result.json\n" + "/data/logs/user-defined/jobs/bg-exp/trial-1/result.json\n" + ) + sandbox.execute = AsyncMock(return_value=find_obs) + + async def _read(req): + path = str(req.path) + idx = int(path.rsplit("trial-", 1)[1].split("/")[0]) + resp = MagicMock() + resp.content = json.dumps( + { + "task_name": f"t-{idx}", + "trial_name": f"trial-{idx:03d}", + "verifier_result": {"rewards": {"reward": 1.0 if idx == 0 else 0.0}}, + "agent_result": {}, + "exception_info": None, + } + ) + return resp + + sandbox.read_file = AsyncMock(side_effect=_read) + sandbox.start_nohup_process = AsyncMock(return_value=(42, None)) + sandbox.wait_for_process_completion = AsyncMock(return_value=(True, "done")) + + obs = MagicMock() + obs.output = "harbor stdout" + obs.exit_code = 0 + sandbox.handle_nohup_output = AsyncMock(return_value=obs) + return sandbox + + +def _make_config(): + return HarborJobConfig( + experiment_id="bg-exp", + job_name="bg-job", + labels={"team": "rl", "step": "1"}, + agents=[AgentConfig(name="a", max_timeout_sec=1800)], + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="terminal-bench", + version="2.0", + ) + ], + environment=RockEnvironmentConfig(auto_stop=True), + ) + + +class TestBlueGreenEquivalence: + async def test_two_sub_trials_flattened_on_both_paths(self): + from rock.sdk.bench import Job as BlueJob + from rock.sdk.job import Job as GreenJob + + mock_blue = _build_mock_sandbox() + with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_blue): + blue_result = await BlueJob(_make_config()).run() + + mock_green = _build_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_green): + green_result = await GreenJob(_make_config()).run() + + assert blue_result.status == green_result.status + assert len(blue_result.trial_results) == len(green_result.trial_results) == 2 + blue_tasks = sorted(t.task_name for t in blue_result.trial_results) + green_tasks = sorted(t.task_name for t in green_result.trial_results) + assert blue_tasks == green_tasks + + async def test_labels_preserved_on_both_paths(self): + from rock.sdk.bench import Job as BlueJob + from rock.sdk.job import Job as GreenJob + + mock_blue = _build_mock_sandbox() + with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_blue): + blue_result = await BlueJob(_make_config()).run() + + mock_green = _build_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_green): + green_result = await GreenJob(_make_config()).run() + + assert blue_result.labels == green_result.labels == {"team": "rl", "step": "1"} + + async def test_raw_output_and_exit_code_populated_on_both(self): + from rock.sdk.bench import Job as BlueJob + from rock.sdk.job import Job as GreenJob + + mock_blue = _build_mock_sandbox() + with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_blue): + blue_result = await BlueJob(_make_config()).run() + + mock_green = _build_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_green): + green_result = await GreenJob(_make_config()).run() + + assert blue_result.raw_output == "harbor stdout" + assert green_result.raw_output == "harbor stdout" + assert blue_result.exit_code == 0 + assert green_result.exit_code == 0 + + async def test_auto_stop_closes_sandbox_on_both(self): + from rock.sdk.bench import Job as BlueJob + from rock.sdk.job import Job as GreenJob + + mock_blue = _build_mock_sandbox() + with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_blue): + await BlueJob(_make_config()).run() + + mock_green = _build_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_green): + await GreenJob(_make_config()).run() + + mock_blue.close.assert_called_once() + mock_green.close.assert_called_once() diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py index 016bcbcd82..1fad0c7857 100644 --- a/tests/unit/sdk/job/test_cli_job.py +++ b/tests/unit/sdk/job/test_cli_job.py @@ -219,7 +219,7 @@ async def test_harbor_requires_config(): async def test_harbor_loads_from_yaml(): - from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig yaml_content = """ experiment_id: exp-123 diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 3a05f5d3d0..f2609f0bda 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -9,7 +9,7 @@ import yaml from rock.sdk.bench.constants import USER_DEFINED_LOGS -from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig +from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, @@ -122,10 +122,15 @@ def test_inherits_job_config(self): def test_defaults(self): cfg = HarborJobConfig(experiment_id="test-exp") - # Inherited - assert cfg.timeout == 3600 + # Inherited — G2: effective timeout derived from agent config + multiplier. + # No agent timeout configured → DEFAULT_WAIT_TIMEOUT fallback (7200). + from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT + + assert cfg.timeout == DEFAULT_WAIT_TIMEOUT assert cfg.labels == {} - assert cfg.job_name is None + # G3: job_name is auto-generated (8-char uuid when no datasets) + assert cfg.job_name is not None + assert len(cfg.job_name) == 8 # Own defaults assert len(cfg.agents) == 1 assert isinstance(cfg.agents[0], AgentConfig) @@ -194,8 +199,9 @@ def test_excludes_rock_fields(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) # Rock-only fields must be absent from Harbor YAML + # job_name is re-injected so harbor uses it as the directory name + assert data["job_name"] == "should-not-appear" rock_only = { - "job_name", "namespace", "experiment_id", "labels", @@ -307,3 +313,139 @@ def test_harbor_inherits_base_fields(self): base_fields = set(JobConfig.model_fields.keys()) harbor_fields = set(HarborJobConfig.model_fields.keys()) assert base_fields.issubset(harbor_fields) + + +# --------------------------------------------------------------------------- +# G7: HarborJobConfig.auto_stop and environment.auto_stop sync (OR semantics) +# --------------------------------------------------------------------------- + + +class TestHarborJobConfigAutoStopSync: + """G7: HarborJobConfig.auto_stop and environment.auto_stop must be kept in sync (OR semantics).""" + + def test_environment_auto_stop_propagates_to_top_level(self): + cfg = HarborJobConfig( + experiment_id="exp-1", + environment=RockEnvironmentConfig(auto_stop=True), + ) + assert cfg.auto_stop is True, "top-level auto_stop must pick up environment.auto_stop" + + def test_top_level_auto_stop_propagates_to_environment(self): + cfg = HarborJobConfig(experiment_id="exp-1", auto_stop=True) + assert cfg.environment.auto_stop is True + + def test_both_true_stays_true(self): + cfg = HarborJobConfig( + experiment_id="exp-1", + auto_stop=True, + environment=RockEnvironmentConfig(auto_stop=True), + ) + assert cfg.auto_stop is True + assert cfg.environment.auto_stop is True + + def test_both_false_stays_false(self): + cfg = HarborJobConfig(experiment_id="exp-1") + assert cfg.auto_stop is False + assert cfg.environment.auto_stop is False + + +# --------------------------------------------------------------------------- +# G3: HarborJobConfig auto-generates job_name when user omits it +# --------------------------------------------------------------------------- + + +class TestHarborJobConfigAutoJobName: + """G3: HarborJobConfig auto-generates job_name when omitted.""" + + def test_explicit_job_name_preserved(self): + cfg = HarborJobConfig(experiment_id="exp", job_name="my-custom") + assert cfg.job_name == "my-custom" + + def test_no_dataset_yields_uuid_only(self): + cfg = HarborJobConfig(experiment_id="exp") + assert cfg.job_name is not None + assert len(cfg.job_name) == 8 # 8-char uuid + + def test_single_dataset_single_task_yields_dataset_task_uuid(self): + from rock.sdk.bench.models.job.config import RegistryDatasetConfig, RemoteRegistryInfo + + cfg = HarborJobConfig( + experiment_id="exp", + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="terminal-bench", + version="2.0", + task_names=["fix-bug"], + ) + ], + ) + parts = cfg.job_name.split("_") + assert parts[0] == "terminal-bench" + assert parts[1] == "fix-bug" + assert len(parts[2]) == 8 + + def test_single_dataset_multi_tasks_yields_dataset_uuid(self): + from rock.sdk.bench.models.job.config import RegistryDatasetConfig, RemoteRegistryInfo + + cfg = HarborJobConfig( + experiment_id="exp", + datasets=[ + RegistryDatasetConfig( + registry=RemoteRegistryInfo(), + name="tb", + version="2.0", + task_names=["a", "b"], + ) + ], + ) + parts = cfg.job_name.split("_") + assert parts[0] == "tb" + assert len(parts) == 2 + assert len(parts[1]) == 8 + + +# --------------------------------------------------------------------------- +# G2: HarborJobConfig effective timeout derivation +# --------------------------------------------------------------------------- + + +class TestHarborJobConfigEffectiveTimeout: + """G2: HarborJobConfig timeout derives from agent.max_timeout_sec + multiplier + 600 buffer.""" + + def test_default_timeout_uses_7200s_fallback(self): + """No agent timeout configured → fallback DEFAULT_WAIT_TIMEOUT=7200.""" + from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT + + cfg = HarborJobConfig(experiment_id="exp") + assert cfg.timeout == DEFAULT_WAIT_TIMEOUT # 7200 + + def test_agent_max_timeout_drives_effective_timeout(self): + cfg = HarborJobConfig( + experiment_id="exp", + agents=[AgentConfig(name="a", max_timeout_sec=1800)], + ) + assert cfg.timeout == 1800 + 600 # +600s env setup + verifier buffer + + def test_agent_override_timeout_used_when_no_max(self): + cfg = HarborJobConfig( + experiment_id="exp", + agents=[AgentConfig(name="a", override_timeout_sec=900)], + ) + assert cfg.timeout == 900 + 600 + + def test_timeout_multiplier_applied(self): + cfg = HarborJobConfig( + experiment_id="exp", + agents=[AgentConfig(name="a", max_timeout_sec=1000)], + timeout_multiplier=2.0, + ) + # (1000 * 2.0) + 600 + assert cfg.timeout == 2600 + + def test_multiplier_only_applied_to_fallback(self): + cfg = HarborJobConfig(experiment_id="exp", timeout_multiplier=2.0) + # DEFAULT_WAIT_TIMEOUT * 2.0 + from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT + + assert cfg.timeout == int(DEFAULT_WAIT_TIMEOUT * 2.0) diff --git a/tests/unit/sdk/job/test_executor.py b/tests/unit/sdk/job/test_executor.py index 14af6c0f94..2173472061 100644 --- a/tests/unit/sdk/job/test_executor.py +++ b/tests/unit/sdk/job/test_executor.py @@ -8,6 +8,7 @@ import rock.sdk.bench # pre-import to avoid circular # noqa: F401 import rock.sdk.job.trial.bash # register BashJobConfig -> BashTrial # noqa: F401 +from rock.sdk.bench.constants import USER_DEFINED_LOGS from rock.sdk.job.config import BashJobConfig from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient from rock.sdk.job.operator import ScatterOperator @@ -219,3 +220,59 @@ def test_returns_none_when_empty(self, monkeypatch): merged = JobExecutor._build_session_env(config) assert merged is None + + +# --------------------------------------------------------------------------- +# G6: USER_DEFINED_LOGS paths +# --------------------------------------------------------------------------- + + +class TestExecutorPaths: + """G6: scripts and nohup outputs must live under USER_DEFINED_LOGS, not /tmp.""" + + async def test_do_submit_writes_script_under_user_defined_logs(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + await executor.submit(ScatterOperator(size=1), BashJobConfig(script="echo hi", job_name="p1")) + + # Inspect the path passed to write_file_by_path — must start with USER_DEFINED_LOGS + write_call = mock_sandbox.write_file_by_path.call_args + script_path = write_call.args[1] if len(write_call.args) >= 2 else write_call.kwargs["path"] + assert script_path.startswith( + USER_DEFINED_LOGS + ), f"script path {script_path!r} must live under {USER_DEFINED_LOGS!r}, not /tmp" + + async def test_do_submit_passes_nohup_tmp_file_under_user_defined_logs(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + await executor.submit(ScatterOperator(size=1), BashJobConfig(script="echo hi", job_name="p2")) + + nohup_call = mock_sandbox.start_nohup_process.call_args + tmp_file = nohup_call.kwargs["tmp_file"] + assert tmp_file.startswith(USER_DEFINED_LOGS) + + +# --------------------------------------------------------------------------- +# G4: on_sandbox_ready hook called after start() +# --------------------------------------------------------------------------- + + +class TestExecutorOnSandboxReady: + async def test_do_submit_calls_on_sandbox_ready_after_start(self): + from rock.sdk.job.trial.bash import BashTrial + + mock_sandbox = _make_mock_sandbox() + mock_sandbox._namespace = "ns" + mock_sandbox._experiment_id = "exp" + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + executor = JobExecutor() + trial = BashTrial(BashJobConfig(script="echo hi", job_name="t")) + trial.on_sandbox_ready = AsyncMock() + await executor._do_submit(trial) + + trial.on_sandbox_ready.assert_awaited_once_with(mock_sandbox) + # Must be called AFTER sandbox.start() + assert mock_sandbox.start.call_count == 1 diff --git a/tests/unit/sdk/job/test_integration.py b/tests/unit/sdk/job/test_integration.py index 749cb85449..221f289ce7 100644 --- a/tests/unit/sdk/job/test_integration.py +++ b/tests/unit/sdk/job/test_integration.py @@ -49,7 +49,7 @@ def test_bash_registered(self): def test_harbor_registered(self): # Importing rock.sdk.job triggers auto-registration import rock.sdk.job # noqa: F401 - from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig + from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.job.trial.registry import _TRIAL_REGISTRY assert HarborJobConfig in _TRIAL_REGISTRY @@ -57,10 +57,10 @@ def test_harbor_registered(self): class TestBackwardCompat: def test_old_agent_imports_still_work(self): - """rock.sdk.bench (formerly rock.sdk.agent) must still export Job, JobConfig, JobResult, JobStatus.""" - from rock.sdk.bench import Job, JobConfig, JobResult, JobStatus + """rock.sdk.bench (formerly rock.sdk.agent) must still export Job, HarborJobConfig, JobResult, JobStatus.""" + from rock.sdk.bench import HarborJobConfig, Job, JobResult, JobStatus assert Job is not None - assert JobConfig is not None + assert HarborJobConfig is not None assert JobResult is not None assert JobStatus is not None diff --git a/tests/unit/sdk/job/test_job.py b/tests/unit/sdk/job/test_job.py index 058970a06e..015b3f05a3 100644 --- a/tests/unit/sdk/job/test_job.py +++ b/tests/unit/sdk/job/test_job.py @@ -17,6 +17,11 @@ def _make_mock_sandbox(): sandbox = AsyncMock() sandbox.sandbox_id = "sb-facade" + # AsyncMock auto-creates child mocks for any attr access; force these + # two back to None so AbstractTrial.on_sandbox_ready's default backfill + # is a no-op (matching a real sandbox that reports no ns / exp_id). + sandbox._namespace = None + sandbox._experiment_id = None sandbox.start = AsyncMock() sandbox.close = AsyncMock() sandbox.create_session = AsyncMock() @@ -167,3 +172,132 @@ async def test_build_result_any_failure_marks_job_failed(self): ).run() assert result.status == JobStatus.FAILED + + +# --------------------------------------------------------------------------- +# Multi-sub-trial flattening (G1 regression) +# --------------------------------------------------------------------------- + + +class TestJobFlattenMultiSubTrials: + """G1: when HarborTrial.collect returns list[N], JobResult.trial_results must have N entries.""" + + async def test_run_flattens_list_returning_collect_into_job_result(self): + from rock.sdk.job.config import JobConfig + from rock.sdk.job.result import TrialResult + from rock.sdk.job.trial.abstract import AbstractTrial + from rock.sdk.job.trial.registry import register_trial + + class MultiCfg(JobConfig): + pass + + class MultiTrial(AbstractTrial): + async def setup(self, sandbox): + pass + + def build(self) -> str: + return "echo hi" + + async def collect(self, sandbox, output, exit_code): + return [TrialResult(task_name=f"sub-{i}") for i in range(3)] + + register_trial(MultiCfg, MultiTrial) + + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(MultiCfg(job_name="multi")).run() + + assert len(result.trial_results) == 3, f"expected 3 flattened sub-trials, got {len(result.trial_results)}" + assert {t.task_name for t in result.trial_results} == {"sub-0", "sub-1", "sub-2"} + assert result.status == JobStatus.COMPLETED + + async def test_run_still_accepts_single_trial_result(self): + mock_sandbox = _make_mock_sandbox() + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="single")).run() + + assert len(result.trial_results) == 1 + assert result.status == JobStatus.COMPLETED + + async def test_timeout_tags_every_sub_trial_in_list(self): + """G1: on timeout, every sub-trial in a list result gets a synthetic ProcessTimeout — + except those that already carry their own exception_info (preserved as-is).""" + from rock.sdk.job.config import JobConfig + from rock.sdk.job.result import ExceptionInfo, TrialResult + from rock.sdk.job.trial.abstract import AbstractTrial + from rock.sdk.job.trial.registry import register_trial + + class PreExistingCfg(JobConfig): + pass + + class PreExistingTrial(AbstractTrial): + async def setup(self, sandbox): + pass + + def build(self) -> str: + return "echo hi" + + async def collect(self, sandbox, output, exit_code): + return [ + TrialResult(task_name="sub-0"), + TrialResult( + task_name="sub-1", + exception_info=ExceptionInfo( + exception_type="OwnError", + exception_message="from trial", + ), + ), + TrialResult(task_name="sub-2"), + ] + + register_trial(PreExistingCfg, PreExistingTrial) + + mock_sandbox = _make_mock_sandbox() + mock_sandbox.wait_for_process_completion = AsyncMock(return_value=(False, "timed out")) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(PreExistingCfg(job_name="multi-timeout")).run() + + assert result.status == JobStatus.FAILED + assert len(result.trial_results) == 3 + by_name = {t.task_name: t for t in result.trial_results} + assert by_name["sub-0"].exception_info.exception_type == "ProcessTimeout" + assert by_name["sub-2"].exception_info.exception_type == "ProcessTimeout" + # Pre-existing exception_info on sub-1 must NOT be overwritten + assert by_name["sub-1"].exception_info.exception_type == "OwnError" + assert by_name["sub-1"].exception_info.exception_message == "from trial" + + +# --------------------------------------------------------------------------- +# G5: raw_output / exit_code surfaced on JobResult +# --------------------------------------------------------------------------- + + +class TestJobResultRawOutputAndExitCode: + """G5: JobResult must surface raw_output and exit_code from the sandbox process.""" + + async def test_run_populates_raw_output_from_obs(self): + mock_sandbox = _make_mock_sandbox() + obs = MagicMock() + obs.output = "hello from sandbox" + obs.exit_code = 0 + mock_sandbox.handle_nohup_output = AsyncMock(return_value=obs) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="echo hi", job_name="test")).run() + + assert result.raw_output == "hello from sandbox" + assert result.exit_code == 0 + + async def test_run_propagates_nonzero_exit_code(self): + mock_sandbox = _make_mock_sandbox() + obs = MagicMock() + obs.output = "err" + obs.exit_code = 7 + mock_sandbox.handle_nohup_output = AsyncMock(return_value=obs) + + with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): + result = await Job(BashJobConfig(script="false", job_name="test")).run() + + assert result.exit_code == 7 + assert result.raw_output == "err" diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index ae2d1636bd..58e81880d2 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -133,3 +133,49 @@ def test_bash_config_creates_bash_trial(self): cfg = BashJobConfig(script="echo hi") trial = _create_trial(cfg) assert isinstance(trial, BashTrial) + + +# --------------------------------------------------------------------------- +# G4: on_sandbox_ready hook — backfill namespace / experiment_id +# Behavior is inherited from AbstractTrial: BashTrial must also backfill. +# --------------------------------------------------------------------------- + + +class TestBashTrialOnSandboxReady: + """G4: BashTrial inherits on_sandbox_ready from AbstractTrial and must backfill namespace/experiment_id.""" + + async def test_namespace_backfilled_when_config_unset(self): + cfg = BashJobConfig(script="echo hi") + trial = BashTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = "sb-ns" + sandbox._experiment_id = "exp-1" + + await trial.on_sandbox_ready(sandbox) + + assert cfg.namespace == "sb-ns" + assert cfg.experiment_id == "exp-1" + + async def test_experiment_id_mismatch_raises(self): + import pytest + + cfg = BashJobConfig(script="echo hi", experiment_id="exp-1") + trial = BashTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = None + sandbox._experiment_id = "exp-DIFFERENT" + + with pytest.raises(ValueError, match="experiment_id mismatch"): + await trial.on_sandbox_ready(sandbox) + + async def test_namespace_mismatch_raises(self): + import pytest + + cfg = BashJobConfig(script="echo hi", namespace="cfg-ns") + trial = BashTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = "sb-ns" + sandbox._experiment_id = None + + with pytest.raises(ValueError, match="namespace mismatch"): + await trial.on_sandbox_ready(sandbox) diff --git a/tests/unit/sdk/job/test_trial_harbor.py b/tests/unit/sdk/job/test_trial_harbor.py index cf42a0001f..445f67dd27 100644 --- a/tests/unit/sdk/job/test_trial_harbor.py +++ b/tests/unit/sdk/job/test_trial_harbor.py @@ -7,7 +7,7 @@ # Pre-import bench to avoid circular-import pitfalls in rock.sdk.job.config import rock.sdk.bench # noqa: F401 -from rock.sdk.bench.models.job.config import JobConfig as HarborJobConfig +from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.job.trial.harbor import HarborTrial from rock.sdk.job.trial.registry import _create_trial @@ -88,36 +88,44 @@ async def test_setup_uploads_harbor_yaml(self): class TestHarborTrialCollect: - async def test_collect_with_trial_results_found(self): + async def test_collect_returns_list_of_all_sub_trials(self): + """Harbor 一个 sandbox 常产出 N 个子 trial；collect 必须返回全部，不能只取第一条。""" cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") trial = HarborTrial(cfg) - trial_json = { - "task_name": "fix-dockerfile", - "trial_name": "trial-001", - "started_at": "2026-01-01T00:00:00Z", - "finished_at": "2026-01-01T00:01:00Z", - "verifier_result": {"rewards": {"reward": 1.0}}, - "agent_result": {}, - "exception_info": None, - } + trial_jsons = [ + { + "task_name": f"task-{i}", + "trial_name": f"trial-{i:03d}", + "verifier_result": {"rewards": {"reward": float(i) / 3}}, + "agent_result": {}, + "exception_info": None, + } + for i in range(3) + ] mock_sandbox = AsyncMock() list_result = MagicMock() - list_result.stdout = f"{cfg.jobs_dir}/test/trial-001/result.json\n" + list_result.stdout = "\n".join(f"{cfg.jobs_dir}/test/trial-{i:03d}/result.json" for i in range(3)) mock_sandbox.execute = AsyncMock(return_value=list_result) - read_response = MagicMock() - read_response.content = json.dumps(trial_json) - mock_sandbox.read_file = AsyncMock(return_value=read_response) + async def _read(req): + path = str(req.path) + idx = int(path.split("trial-")[1].split("/")[0]) + resp = MagicMock() + resp.content = json.dumps(trial_jsons[idx]) + return resp + + mock_sandbox.read_file = AsyncMock(side_effect=_read) result = await trial.collect(mock_sandbox, output="", exit_code=0) - assert result.task_name == "fix-dockerfile" - assert result.exception_info is None - assert result.score == 1.0 + assert isinstance(result, list), f"collect must return list, got {type(result)}" + assert len(result) == 3, f"expected 3 sub-trials, got {len(result)}" + assert {r.task_name for r in result} == {"task-0", "task-1", "task-2"} - async def test_collect_with_no_trials(self): + async def test_collect_returns_list_with_synthetic_failure_when_no_trials(self): + """Harbor 没写出任何 result.json 时，返回长度为 1 的 list，携带 HarborNoTrials 异常。""" cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") trial = HarborTrial(cfg) @@ -128,8 +136,10 @@ async def test_collect_with_no_trials(self): result = await trial.collect(mock_sandbox, output="", exit_code=0) - assert result.exception_info is not None - assert result.exception_info.exception_type == "HarborNoTrials" + assert isinstance(result, list) + assert len(result) == 1 + assert result[0].exception_info is not None + assert result[0].exception_info.exception_type == "HarborNoTrials" # --------------------------------------------------------------------------- @@ -142,3 +152,47 @@ def test_harbor_config_creates_harbor_trial(self): cfg = HarborJobConfig(experiment_id="exp-1") trial = _create_trial(cfg) assert isinstance(trial, HarborTrial) + + +# --------------------------------------------------------------------------- +# G4: on_sandbox_ready hook — backfill namespace / experiment_id +# --------------------------------------------------------------------------- + + +class TestHarborTrialOnSandboxReady: + """G4: HarborTrial must backfill namespace / experiment_id from sandbox into config.""" + + async def test_namespace_backfilled_when_config_unset(self): + cfg = HarborJobConfig(experiment_id="exp-1") + trial = HarborTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = "sb-ns" + sandbox._experiment_id = "exp-1" + + await trial.on_sandbox_ready(sandbox) + + assert cfg.namespace == "sb-ns" + + async def test_experiment_id_mismatch_raises(self): + import pytest + + cfg = HarborJobConfig(experiment_id="exp-1") + trial = HarborTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = None + sandbox._experiment_id = "exp-DIFFERENT" + + with pytest.raises(ValueError, match="experiment_id mismatch"): + await trial.on_sandbox_ready(sandbox) + + async def test_namespace_mismatch_raises(self): + import pytest + + cfg = HarborJobConfig(experiment_id="exp-1", namespace="cfg-ns") + trial = HarborTrial(cfg) + sandbox = MagicMock() + sandbox._namespace = "sb-ns" + sandbox._experiment_id = None + + with pytest.raises(ValueError, match="namespace mismatch"): + await trial.on_sandbox_ready(sandbox) diff --git a/uv.lock b/uv.lock index 6a8933124d..65ca97dd79 100644 --- a/uv.lock +++ b/uv.lock @@ -4035,7 +4035,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.5.0" +version = "1.5.1" source = { editable = "." } dependencies = [ { name = "anyio" }, From 527365699e461d9cf05135149a1ab3dc210ae988 Mon Sep 17 00:00:00 2001 From: sanfeng-lhh Date: Wed, 15 Apr 2026 10:05:44 +0800 Subject: [PATCH 029/226] feat: add TemplateConfig and template field to NativeConfig (#786) Add TemplateConfig model with name/revision fields for referencing Agent-Bench templates. Add optional template field to NativeConfig. --- rock/sdk/bench/models/trial/config.py | 8 ++ tests/unit/sdk/job/test_config.py | 102 ++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index 6b104c402c..bfb3c8b28f 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -77,6 +77,13 @@ def to_harbor_environment(self) -> dict: return harbor.model_dump(mode="json", exclude_none=True) +class TemplateConfig(BaseModel): + """Agent-Bench template reference used by native verifier.""" + + name: str | None = None + revision: str | None = None + + class NativeConfig(BaseModel): """Config specific to native verifier mode. When image and script are both provided, a ContainerVerifier is used to @@ -87,6 +94,7 @@ class NativeConfig(BaseModel): image: str | None = None script: str | None = None oss_deps: dict[str, str] = Field(default_factory=dict) + template: TemplateConfig | None = None class VerifierConfig(BaseModel): diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index f2609f0bda..292b63d17a 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -13,8 +13,10 @@ from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, + NativeConfig, RockEnvironmentConfig, TaskConfig, + TemplateConfig, VerifierConfig, ) from rock.sdk.job.config import BashJobConfig, JobConfig @@ -307,6 +309,106 @@ def test_from_yaml_file_not_found(self): HarborJobConfig.from_yaml("/nonexistent/path.yaml") +# --------------------------------------------------------------------------- +# TemplateConfig +# --------------------------------------------------------------------------- + + +class TestTemplateConfig: + def test_defaults(self): + cfg = TemplateConfig() + assert cfg.name is None + assert cfg.revision is None + + def test_with_values(self): + cfg = TemplateConfig( + name="swe-agent-internal/SWE-Gym/SWE-Gym", + revision="53634366f454e6dc5fc3ceb85896c706b9ad1078", + ) + assert cfg.name == "swe-agent-internal/SWE-Gym/SWE-Gym" + assert cfg.revision == "53634366f454e6dc5fc3ceb85896c706b9ad1078" + + def test_partial_values(self): + cfg = TemplateConfig(name="my-agent/my-org/my-dataset") + assert cfg.name == "my-agent/my-org/my-dataset" + assert cfg.revision is None + + def test_json_round_trip(self): + cfg = TemplateConfig( + name="swe-agent-internal/SWE-Gym/SWE-Gym", + revision="53634366f454e6dc5fc3ceb85896c706b9ad1078", + ) + data = cfg.model_dump(mode="json") + restored = TemplateConfig(**data) + assert restored == cfg + + def test_exclude_none_omits_unset_fields(self): + cfg = TemplateConfig(name="my-agent/my-org/my-dataset") + data = cfg.model_dump(mode="json", exclude_none=True) + assert "name" in data + assert "revision" not in data + + +# --------------------------------------------------------------------------- +# NativeConfig +# --------------------------------------------------------------------------- + + +class TestNativeConfig: + def test_defaults(self): + cfg = NativeConfig() + assert cfg.image is None + assert cfg.script is None + assert cfg.oss_deps == {} + assert cfg.template is None + + def test_template_none_by_default(self): + cfg = NativeConfig(image="ubuntu:22.04") + assert cfg.template is None + + def test_template_from_dict(self): + cfg = NativeConfig( + template={ + "name": "swe-agent-internal/SWE-Gym/SWE-Gym", + "revision": "53634366f454e6dc5fc3ceb85896c706b9ad1078", + } + ) + assert isinstance(cfg.template, TemplateConfig) + assert cfg.template.name == "swe-agent-internal/SWE-Gym/SWE-Gym" + assert cfg.template.revision == "53634366f454e6dc5fc3ceb85896c706b9ad1078" + + def test_template_from_model(self): + tmpl = TemplateConfig(name="my-agent/my-org/my-dataset", revision="abc123") + cfg = NativeConfig(template=tmpl) + assert cfg.template is tmpl + + def test_json_round_trip_with_template(self): + cfg = NativeConfig( + image="eval:latest", + template=TemplateConfig( + name="swe-agent-internal/SWE-Gym/SWE-Gym", + revision="53634366f454e6dc5fc3ceb85896c706b9ad1078", + ), + ) + data = cfg.model_dump(mode="json") + restored = NativeConfig(**data) + assert restored.template.name == cfg.template.name + assert restored.template.revision == cfg.template.revision + + def test_exclude_none_omits_template_when_not_set(self): + cfg = NativeConfig(image="eval:latest") + data = cfg.model_dump(mode="json", exclude_none=True) + assert "template" not in data + + def test_exclude_none_includes_template_when_set(self): + cfg = NativeConfig( + template=TemplateConfig(name="my-agent/my-org/my-dataset", revision="rev1") + ) + data = cfg.model_dump(mode="json", exclude_none=True) + assert "template" in data + assert data["template"]["name"] == "my-agent/my-org/my-dataset" + + class TestHarborInheritsBase: def test_harbor_inherits_base_fields(self): """HarborJobConfig (agent's) inherits all base JobConfig fields.""" From 80efdd0f7050475c57022a689fe2bceb771a2fa2 Mon Sep 17 00:00:00 2001 From: berstpander Date: Wed, 15 Apr 2026 11:18:58 +0800 Subject: [PATCH 030/226] feat: truncate path segments in auto-generated job_name (#791) When dataset_name or task_name contains '/', only use the last segment after the final '/' for job_name generation. This prevents path separators from appearing in job names and keeps them concise. Example: - Before: 'swe_bench/verified/task1_abc12345' - After: 'task1_abc12345' --- rock/sdk/bench/job.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/rock/sdk/bench/job.py b/rock/sdk/bench/job.py index 92c960275a..1ac2a0f442 100644 --- a/rock/sdk/bench/job.py +++ b/rock/sdk/bench/job.py @@ -285,6 +285,8 @@ def _generate_default_job_name(self) -> None: If job_name is None, generate one with the format: {dataset_name}_{task_name if single task}_{uuid} + + For dataset_name and task_name, only the last segment after "/" is used. """ if self._config.job_name is not None: # User has set a custom job_name, keep it @@ -297,12 +299,16 @@ def _generate_default_job_name(self) -> None: if self._config.datasets: dataset = self._config.datasets[0] if hasattr(dataset, "name") and dataset.name: - parts.append(dataset.name) + # Only use the last segment after "/" + dataset_name = dataset.name.rsplit("/", 1)[-1] + parts.append(dataset_name) # Get task name if there's only one task task_names = dataset.task_names if task_names and len(task_names) == 1: - parts.append(task_names[0]) + # Only use the last segment after "/" + task_name = task_names[0].rsplit("/", 1)[-1] + parts.append(task_name) # Add short UUID (8 characters) parts.append(uuid.uuid4().hex[:8]) From 8a5b0955420b49e8f885ebcd17b49ee3723c6080 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 15 Apr 2026 11:24:55 +0800 Subject: [PATCH 031/226] fix: enlarge SandboxRecord.image to 512 and disable asyncpg statement cache (#794) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: add varchar length test for SandboxRecord.image column RED: image=VARCHAR(128) rejects 173-char image string on PostgreSQL with StringDataRightTruncationError. Co-Authored-By: Claude Opus 4.6 * fix: widen SandboxRecord.image column to VARCHAR(512) Long container image references (e.g. 173 chars) exceed VARCHAR(128) and cause StringDataRightTruncationError on insert. Aligns the ORM model with hotfix_increase_varchar_lengths.sql. Co-Authored-By: Claude Opus 4.6 * fix: sort DDL indexes for deterministic output and regenerate schema Sort table.indexes in gen_ddl.py to produce stable, alphabetically-ordered CREATE INDEX statements. Regenerate sandbox_record.sql reflecting the image VARCHAR(512) change and deterministic index ordering. Sync uv.lock to version 1.5.1. Co-Authored-By: Claude Opus 4.6 * test: add InvalidCachedStatementError test after external DDL RED: after external ALTER COLUMN image TYPE VARCHAR(1024), SandboxTable.list_by_in raises InvalidCachedStatementError because asyncpg's cached prepared statement is stale and AsyncSession's implicit transaction prevents auto-retry. Co-Authored-By: Claude Opus 4.6 * fix: disable asyncpg statement cache to prevent InvalidCachedStatementError After external DDL (e.g. ALTER COLUMN TYPE via hotfix SQL), asyncpg's cached prepared statements become stale. Since AsyncSession wraps queries in implicit transactions, asyncpg cannot auto-retry and raises InvalidCachedStatementError. Setting statement_cache_size=0 eliminates this at negligible cost (~30-50µs per query). Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- rock/admin/core/db_provider.py | 9 ++- rock/admin/core/schema.py | 2 +- scripts/gen_ddl.py | 3 +- sql/sandbox_record.sql | 2 +- .../admin/core/test_schema_varchar_lengths.py | 52 ++++++++++++++ tests/unit/admin/core/test_statement_cache.py | 68 +++++++++++++++++++ 6 files changed, 131 insertions(+), 5 deletions(-) create mode 100644 tests/unit/admin/core/test_schema_varchar_lengths.py create mode 100644 tests/unit/admin/core/test_statement_cache.py diff --git a/rock/admin/core/db_provider.py b/rock/admin/core/db_provider.py index a20277dd00..edb89cc6e1 100644 --- a/rock/admin/core/db_provider.py +++ b/rock/admin/core/db_provider.py @@ -32,8 +32,13 @@ def engine(self) -> AsyncEngine: return self._engine async def init(self) -> None: - """Create the async engine.""" - self._engine = create_async_engine(self._url, echo=False) + """Create the async engine. + + For asyncpg, ``statement_cache_size=0`` prevents + ``InvalidCachedStatementError`` after external DDL changes + """ + connect_args = {"statement_cache_size": 0} if "asyncpg" in self._url else {} + self._engine = create_async_engine(self._url, echo=False, connect_args=connect_args) async def create_tables(self) -> None: """Create all tables defined in Base.metadata (idempotent).""" diff --git a/rock/admin/core/schema.py b/rock/admin/core/schema.py index cfed2f3dc4..ab0ef83ad5 100644 --- a/rock/admin/core/schema.py +++ b/rock/admin/core/schema.py @@ -32,7 +32,7 @@ class SandboxRecord(Base): sandbox_id = Column(String(128), primary_key=True) user_id = Column(String(128), nullable=False, default="default") - image = Column(String(128), nullable=False, default="default") + image = Column(String(512), nullable=False, default="default") experiment_id = Column(String(128), nullable=False, default="default") namespace = Column(String(128), nullable=False, default="default") cluster_name = Column(String(128), nullable=False, default="default") diff --git a/scripts/gen_ddl.py b/scripts/gen_ddl.py index 02a7013190..34a93dfb7a 100644 --- a/scripts/gen_ddl.py +++ b/scripts/gen_ddl.py @@ -30,7 +30,8 @@ def gen_ddl(dialect) -> str: lines: list[str] = [] for table in Base.metadata.sorted_tables: lines.append(str(CreateTable(table).compile(dialect=dialect)).strip() + ";") - for index in table.indexes: + # table.indexes has set-like semantics; sort for deterministic output. + for index in sorted(table.indexes, key=lambda idx: idx.name or ""): lines.append(str(CreateIndex(index).compile(dialect=dialect)).strip() + ";") return "\n\n".join(lines) diff --git a/sql/sandbox_record.sql b/sql/sandbox_record.sql index 35c4cb254e..b02bb565a0 100644 --- a/sql/sandbox_record.sql +++ b/sql/sandbox_record.sql @@ -1,7 +1,7 @@ CREATE TABLE sandbox_record ( sandbox_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, - image VARCHAR(128) NOT NULL, + image VARCHAR(512) NOT NULL, experiment_id VARCHAR(128) NOT NULL, namespace VARCHAR(128) NOT NULL, cluster_name VARCHAR(128) NOT NULL, diff --git a/tests/unit/admin/core/test_schema_varchar_lengths.py b/tests/unit/admin/core/test_schema_varchar_lengths.py new file mode 100644 index 0000000000..bc7ffdb3f0 --- /dev/null +++ b/tests/unit/admin/core/test_schema_varchar_lengths.py @@ -0,0 +1,52 @@ +"""TDD red-green test: SandboxRecord ORM image column VARCHAR length. + +RED — ORM creates table with VARCHAR(128), inserting a 196-char image fails on PostgreSQL. +GREEN — Change schema to VARCHAR(512), same insert succeeds. + +Uses pg_container fixture (real PostgreSQL) so VARCHAR constraints are enforced. +""" + +from __future__ import annotations + +import pytest +from sqlalchemy.exc import DataError + +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.sandbox_table import SandboxTable +from rock.config import DatabaseConfig + +# A realistic image reference that exceeds VARCHAR(128). +# Taken from a real production failure log. +_LONG_IMAGE = ( + "registry.example.com/org/project-sandbox-images" + ":my-very-long-tag-name-that-simulates-a-real-world-scenario" + "-aabbccdd0011223344556677889900ff-v1234567890abcdef1234567890abcdef" +) +assert len(_LONG_IMAGE) > 128, f"Test image must exceed 128 chars, got {len(_LONG_IMAGE)}" + + +@pytest.mark.need_docker +@pytest.mark.need_database +class TestImageVarcharLength: + """ORM image column must accept long registry paths on real PostgreSQL.""" + + @pytest.fixture + async def db(self, pg_container): + """Create a SandboxTable backed by the test PostgreSQL container.""" + provider = DatabaseProvider(db_config=DatabaseConfig(url=pg_container["url"])) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table + await provider.close() + + async def test_insert_long_image(self, db): + """A 196-char image string must be accepted by the ORM schema.""" + sandbox_id = "varchar-img-001" + await db.create(sandbox_id, { + "image": _LONG_IMAGE, + "create_time": "2026-04-14T00:00:00Z", + }) + record = await db.get(sandbox_id) + assert record is not None + assert record["image"] == _LONG_IMAGE diff --git a/tests/unit/admin/core/test_statement_cache.py b/tests/unit/admin/core/test_statement_cache.py new file mode 100644 index 0000000000..ac0fcf75a5 --- /dev/null +++ b/tests/unit/admin/core/test_statement_cache.py @@ -0,0 +1,68 @@ +"""TDD: DatabaseProvider must tolerate external DDL without errors. + +Production scenario: after external ``ALTER TABLE ... ALTER COLUMN TYPE``, +the next query via ``SandboxTable.list_by_in`` raises +``InvalidCachedStatementError`` because asyncpg's prepared statement cache +holds a stale plan, and ``AsyncSession``'s implicit transaction prevents +asyncpg's auto-retry. + +RED — DatabaseProvider without ``statement_cache_size=0`` → error after DDL. +GREEN — DatabaseProvider sets ``statement_cache_size=0`` → no error. +""" + +from __future__ import annotations + +import asyncpg +import pytest + +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.sandbox_table import SandboxTable +from rock.config import DatabaseConfig + + +@pytest.mark.need_docker +@pytest.mark.need_database +class TestBatchGetAfterDDL: + """Reproduce and fix InvalidCachedStatementError on the sandboxes/batch code path.""" + + @pytest.fixture + async def setup(self, pg_container): + """Create tables via ORM and yield SandboxTable + pg url.""" + provider = DatabaseProvider(db_config=DatabaseConfig(url=pg_container["url"])) + await provider.init() + await provider.create_tables() + table = SandboxTable(provider) + yield table, pg_container["url"] + await provider.close() + + async def test_list_by_in_after_alter_column(self, setup): + """SandboxTable.list_by_in must work after external ALTER COLUMN TYPE. + + 1. Create records and query (warms asyncpg statement cache). + 2. External hotfix: ALTER COLUMN image TYPE VARCHAR(1024). + 3. Same list_by_in query must succeed, not raise InvalidCachedStatementError. + """ + table, pg_url = setup + + # 1. populate and query — warms prepared statement cache + ids = [f"cache-{i:03d}" for i in range(10)] + for sid in ids: + await table.create(sid, { + "image": "python:3.11", + "create_time": "2026-04-14T00:00:00Z", + }) + records = await table.list_by_in("sandbox_id", ids) + assert len(records) == 10 + + # 2. external DDL — simulates hotfix applied while app is running + raw = await asyncpg.connect(pg_url) + try: + await raw.execute( + "ALTER TABLE sandbox_record ALTER COLUMN image TYPE VARCHAR(1024)" + ) + finally: + await raw.close() + + # 3. same query — must not raise InvalidCachedStatementError + records = await table.list_by_in("sandbox_id", ids) + assert len(records) == 10 From 9bdbf6504cbb5a5946f91ff8e0f1b3012c1571a1 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 12:01:14 +0800 Subject: [PATCH 032/226] Feature/xinshi/harbor jobs (#798) * chore: add yaml and md files to gitignore * refactor(job): decouple JobConfig from bench, fix to_harbor_yaml serialization - Create JobEnvironmentConfig(SandboxConfig) in rock.sdk.job.config, absorbing setup_commands/file_uploads/auto_stop/env from JobConfig - JobConfig base class no longer imports bench module, eliminating circular import workaround - RockEnvironmentConfig inherits JobEnvironmentConfig + EnvironmentConfig instead of SandboxConfig + EnvironmentConfig - to_harbor_yaml() uses _HarborJobFields mirror model (model_validate pattern) instead of manual _BASE_FIELDS exclude + re-inject - Fix: namespace, experiment_id, labels now correctly serialized to Harbor YAML (were silently dropped before) - OssMirrorConfig gains namespace/experiment_id fields to align with harbor.environments.base.OssMirrorConfig - Add HFRegistryInfo registry type to match Harbor's 4 registry types - Add _sync_namespace_to_oss_mirror validator on HarborJobConfig Co-Authored-By: Claude Opus 4.6 * docs(job): add config-refactor design doc Document the JobConfig hierarchy restructuring, to_harbor_yaml fix, OssMirrorConfig alignment, and HFRegistryInfo addition. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .gitignore | 2 + docs/dev/job/config-refactor.md | 107 ++++++++++++++++++ rock/cli/command/job.py | 10 +- rock/sdk/bench/__init__.py | 2 + rock/sdk/bench/models/__init__.py | 2 + rock/sdk/bench/models/job/config.py | 88 +++++++++----- rock/sdk/bench/models/trial/config.py | 28 ++--- rock/sdk/job/__init__.py | 20 ++-- rock/sdk/job/config.py | 27 +++-- rock/sdk/job/executor.py | 4 +- rock/sdk/job/trial/abstract.py | 2 +- rock/sdk/job/trial/bash.py | 4 +- rock/sdk/job/trial/harbor.py | 2 +- .../agent/test_job_config_serialization.py | 9 +- tests/unit/sdk/agent/test_oss_mirror.py | 24 ++-- .../sdk/job/test_blue_green_equivalence.py | 1 - tests/unit/sdk/job/test_cli_job.py | 7 +- tests/unit/sdk/job/test_config.py | 105 +++++++---------- tests/unit/sdk/job/test_executor.py | 13 ++- tests/unit/sdk/job/test_job.py | 2 - tests/unit/sdk/job/test_trial_bash.py | 12 +- tests/unit/sdk/job/test_trial_harbor.py | 6 +- tests/unit/sdk/job/test_trial_registry.py | 6 +- 23 files changed, 299 insertions(+), 184 deletions(-) create mode 100644 docs/dev/job/config-refactor.md diff --git a/.gitignore b/.gitignore index b0a61eb598..95e72a518a 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,5 @@ docs/superpowers/ .env *.db +*.intern.yaml +*.intern.md \ No newline at end of file diff --git a/docs/dev/job/config-refactor.md b/docs/dev/job/config-refactor.md new file mode 100644 index 0000000000..66181aef94 --- /dev/null +++ b/docs/dev/job/config-refactor.md @@ -0,0 +1,107 @@ +# Job Config 层次重构 + +**Date**: 2026-04-15 +**PR Branch**: `feature/xinshi/harbor-jobs` + +## 背景 + +Job 模块重构（#779/#780, #788/#789）引入了 Job/Operator/Executor/Trial 抽象，但 config 层存在以下问题： + +1. **依赖倒置** — 基类 `JobConfig` 依赖 bench 模块的 `RockEnvironmentConfig`，导致循环导入 +2. **字段错位** — `auto_stop`/`setup_commands`/`file_uploads`/`env` 同时存在于 `JobConfig` 顶层和 `RockEnvironmentConfig` 中 +3. **`to_harbor_yaml()` 丢字段** — `namespace`/`experiment_id`/`labels` 被错误排除 +4. **`OssMirrorConfig` 缺字段** — 比 harbor 版本少 `namespace`/`experiment_id` +5. **缺 `HFRegistryInfo`** — harbor 支持 4 种 registry 类型，ROCK 只有 3 种 + +## 修改内容 + +### 1. 新建 `JobEnvironmentConfig` + +``` +rock/sdk/job/config.py + +JobEnvironmentConfig(SandboxConfig) + ├── setup_commands: list[str] + ├── file_uploads: list[tuple[str, str]] + ├── auto_stop: bool + └── env: dict[str, str] +``` + +承载从 `JobConfig` 下沉的环境字段。`JobConfig` 不再依赖 bench 模块。 + +### 2. `JobConfig` 瘦身 + +```python +class JobConfig(BaseModel): + environment: JobEnvironmentConfig + job_name: str | None + namespace: str | None + experiment_id: str | None + labels: dict[str, str] + timeout: int = 3600 +``` + +### 3. `RockEnvironmentConfig` 改继承 + +``` +Before: RockEnvironmentConfig(SandboxConfig, EnvironmentConfig) +After: RockEnvironmentConfig(JobEnvironmentConfig, EnvironmentConfig) +``` + +不再重复声明 `setup_commands`/`file_uploads`/`auto_stop`（从 `JobEnvironmentConfig` 继承）。 + +### 4. `to_harbor_yaml()` 使用镜像模型 + +引入 `_HarborJobFields` 镜像模型，字段对齐 `harbor.models.job.config.JobConfig`。 +使用 `model_validate` 模式自动过滤（与 `to_harbor_environment()` 一致）， +取代手动 `_BASE_FIELDS` exclude + re-inject。 + +**修复**：`namespace`/`experiment_id`/`labels` 现在正确出现在 Harbor YAML 中。 + +### 5. `OssMirrorConfig` 补齐字段 + +新增 `namespace` 和 `experiment_id`，对齐 `harbor.environments.base.OssMirrorConfig`。 +通过 `_sync_experiment_id` 和 `_sync_namespace_to_oss_mirror` validator 自动同步。 + +### 6. 新增 `HFRegistryInfo` + +对齐 harbor 的 4 种 registry 类型（`Oss|Remote|Local|HF`）。 + +## 修改后的层次结构 + +``` +JobEnvironmentConfig(SandboxConfig) ← rock.sdk.job.config（新增） + ├── image, memory, cpus, ... ← SandboxConfig + └── setup_commands, file_uploads, ← job 环境字段 + auto_stop, env + +JobConfig(BaseModel) ← rock.sdk.job.config（瘦身） + ├── environment: JobEnvironmentConfig + ├── job_name, namespace, experiment_id + ├── labels, timeout + └── BashJobConfig(JobConfig) + +RockEnvironmentConfig( ← rock.sdk.bench.models.trial.config（改继承） + JobEnvironmentConfig, + EnvironmentConfig +) + └── to_harbor_environment() → dict + +HarborJobConfig(JobConfig) ← rock.sdk.bench.models.job.config + ├── environment: RockEnvironmentConfig ← override + ├── jobs_dir, agents, datasets, ... ← harbor 原生字段 + └── to_harbor_yaml() → str ← 使用 _HarborJobFields 镜像模型 +``` + +## 受影响的文件 + +| 文件 | 修改类型 | +|------|----------| +| `rock/sdk/job/config.py` | 新增 `JobEnvironmentConfig`，瘦身 `JobConfig` | +| `rock/sdk/bench/models/trial/config.py` | `OssMirrorConfig` 补字段，`RockEnvironmentConfig` 改继承 | +| `rock/sdk/bench/models/job/config.py` | 新增 `HFRegistryInfo`/`_HarborJobFields`，修复 validators 和 `to_harbor_yaml` | +| `rock/sdk/job/__init__.py` | 删除循环导入 workaround，导出 `JobEnvironmentConfig` | +| `rock/sdk/job/executor.py` | `config.env` → `config.environment.env` 等 | +| `rock/sdk/job/trial/*.py` | 字段访问路径更新 | +| `rock/cli/command/job.py` | `auto_stop`/`file_uploads` 移入 environment | +| `rock/sdk/bench/__init__.py` | 导出 `HFRegistryInfo` | diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index 2a5630e36a..fa288ff147 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -52,10 +52,12 @@ async def _job_run(self, args: argparse.Namespace): config = BashJobConfig( script=args.script_content, script_path=args.script, - environment=RockEnvironmentConfig(**env_kwargs), - file_uploads=file_uploads, + environment=RockEnvironmentConfig( + **env_kwargs, + file_uploads=file_uploads, + auto_stop=True, + ), timeout=args.timeout, - auto_stop=True, ) elif job_type == "harbor": @@ -67,7 +69,7 @@ async def _job_run(self, args: argparse.Namespace): config = HarborJobConfig.from_yaml(args.config) if args.image: config.environment.image = args.image - config.auto_stop = True + config.environment.auto_stop = True else: logger.error(f"Unknown job type: {job_type}") diff --git a/rock/sdk/bench/__init__.py b/rock/sdk/bench/__init__.py index 72d6fe6450..35bb1f4722 100644 --- a/rock/sdk/bench/__init__.py +++ b/rock/sdk/bench/__init__.py @@ -1,6 +1,7 @@ from rock.sdk.bench.job import Job from rock.sdk.bench.models.job.config import ( HarborJobConfig, + HFRegistryInfo, LocalDatasetConfig, OrchestratorConfig, OssRegistryInfo, @@ -37,6 +38,7 @@ "AgentResult", "ExceptionInfo", "HarborJobConfig", + "HFRegistryInfo", "RockEnvironmentConfig", "RegistryDatasetConfig", "LocalDatasetConfig", diff --git a/rock/sdk/bench/models/__init__.py b/rock/sdk/bench/models/__init__.py index ebe96f4488..208eb65ff7 100644 --- a/rock/sdk/bench/models/__init__.py +++ b/rock/sdk/bench/models/__init__.py @@ -2,6 +2,7 @@ from rock.sdk.bench.models.job.config import ( DatasetConfig, HarborJobConfig, + HFRegistryInfo, OrchestratorConfig, RetryConfig, ) @@ -20,6 +21,7 @@ __all__ = [ "HarborJobConfig", + "HFRegistryInfo", "OrchestratorConfig", "RetryConfig", "DatasetConfig", diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index 94f13c242f..ed087231f1 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -7,7 +7,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, ClassVar +from typing import Any from pydantic import BaseModel, Field, model_validator @@ -17,6 +17,7 @@ from rock.sdk.bench.models.trial.config import ( AgentConfig, ArtifactConfig, + EnvironmentConfig, OssMirrorConfig, RockEnvironmentConfig, # noqa: F401 — re-exported for backward compat TaskConfig, @@ -86,6 +87,13 @@ class LocalRegistryInfo(BaseModel): path: Path +class HFRegistryInfo(BaseModel): + """HuggingFace registry, corresponds to CLI ``--registry-type hf``.""" + + split: str | None = None + revision: str | None = None + + # --------------------------------------------------------------------------- # DatasetConfig (aligned with harbor's LocalDatasetConfig / RegistryDatasetConfig) # --------------------------------------------------------------------------- @@ -108,7 +116,7 @@ class LocalDatasetConfig(BaseDatasetConfig): class RegistryDatasetConfig(BaseDatasetConfig): """Registry dataset, corresponds to CLI ``-d/--dataset`` + ``--registry-type``.""" - registry: OssRegistryInfo | RemoteRegistryInfo | LocalRegistryInfo + registry: OssRegistryInfo | RemoteRegistryInfo | LocalRegistryInfo | HFRegistryInfo name: str version: str | None = None overwrite: bool = False @@ -128,6 +136,36 @@ def _infer_version_from_split(self): DatasetConfig = LocalDatasetConfig | RegistryDatasetConfig +class _HarborJobFields(BaseModel): + """Harbor JobConfig field mirror — used by to_harbor_yaml() for serialization filtering. + + Fields align with harbor.models.job.config.JobConfig. + ROCK-only fields (SandboxConfig, setup_commands, etc.) are automatically + discarded by model_validate. + """ + + namespace: str | None = None + experiment_id: str | None = None + job_name: str | None = None + jobs_dir: Path = Path("jobs") + n_attempts: int = 1 + timeout_multiplier: float = 1.0 + agent_timeout_multiplier: float | None = None + verifier_timeout_multiplier: float | None = None + agent_setup_timeout_multiplier: float | None = None + environment_build_timeout_multiplier: float | None = None + debug: bool = False + orchestrator: OrchestratorConfig = Field(default_factory=OrchestratorConfig) + environment: EnvironmentConfig = Field(default_factory=EnvironmentConfig) + verifier: VerifierConfig = Field(default_factory=VerifierConfig) + metrics: list[MetricConfig] = Field(default_factory=list) + agents: list[AgentConfig] = Field(default_factory=lambda: [AgentConfig()]) + datasets: list[LocalDatasetConfig | RegistryDatasetConfig] = Field(default_factory=list) + tasks: list[TaskConfig] = Field(default_factory=list) + artifacts: list[str | ArtifactConfig] = Field(default_factory=list) + labels: dict[str, str] = Field(default_factory=dict) + + class HarborJobConfig(_BaseJobConfig): """Harbor Job configuration: extends base JobConfig with Harbor-native fields. @@ -136,7 +174,10 @@ class HarborJobConfig(_BaseJobConfig): and passed to ``harbor jobs start -c``. """ - # ── Harbor native fields (base fields: environment, job_name, namespace, etc. are inherited) ── + # ── Override environment to use RockEnvironmentConfig (adds harbor env fields) ── + environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) + + # ── Harbor native fields (base fields: job_name, namespace, etc. are inherited) ── jobs_dir: Path = Path(USER_DEFINED_LOGS) / "jobs" n_attempts: int = 1 timeout_multiplier: float = 1.0 @@ -155,12 +196,7 @@ class HarborJobConfig(_BaseJobConfig): @model_validator(mode="after") def _sync_experiment_id(self): - """Validate and sync experiment_id between JobConfig and SandboxConfig. - - 1. experiment_id must not be empty. - 2. If environment.experiment_id is already set, it must match. - 3. Propagate experiment_id down to environment (SandboxConfig). - """ + """Sync experiment_id: JobConfig -> environment -> oss_mirror.""" if not self.experiment_id: raise ValueError("experiment_id must not be empty") env_exp = self.environment.experiment_id @@ -170,19 +206,15 @@ def _sync_experiment_id(self): f"but environment (SandboxConfig) has '{env_exp}'" ) self.environment.experiment_id = self.experiment_id + if self.environment.oss_mirror is not None: + self.environment.oss_mirror.experiment_id = self.experiment_id return self @model_validator(mode="after") - def _sync_auto_stop(self): - """G7: keep top-level auto_stop and environment.auto_stop in sync (OR semantics). - - Users may set either. Legacy ``environment.auto_stop=True`` (pre-job-refactor) - must still work; new ``config.auto_stop=True`` must also propagate down to - the environment so the RockEnvironmentConfig path reads the same value. - """ - effective = bool(self.auto_stop) or bool(self.environment.auto_stop) - self.auto_stop = effective - self.environment.auto_stop = effective + def _sync_namespace_to_oss_mirror(self): + """Sync namespace: JobConfig -> oss_mirror.""" + if self.namespace is not None and self.environment.oss_mirror is not None: + self.environment.oss_mirror.namespace = self.namespace return self @model_validator(mode="after") @@ -243,23 +275,17 @@ def _compute_effective_timeout(self): self.timeout = int(DEFAULT_WAIT_TIMEOUT * multiplier) return self - # Base JobConfig fields to exclude when serializing to Harbor YAML - _BASE_FIELDS: ClassVar[set[str]] = set(_BaseJobConfig.model_fields.keys()) - def to_harbor_yaml(self) -> str: - """Serialize Harbor-native fields to YAML for ``harbor jobs start -c``. + """Serialize to Harbor YAML for ``harbor jobs start -c``. - Base JobConfig fields (environment, setup_commands, etc.) are excluded. - ``job_name`` is re-injected so harbor uses it as the job directory name - instead of its default timestamp-based naming. - Harbor environment fields (force_build, override_cpus, etc.) - are re-injected under ``environment``. + Uses _HarborJobFields mirror model to filter — only harbor-recognized + fields pass through. Environment is handled specially via + to_harbor_environment() to strip Rock-only sandbox fields. """ import yaml - data = self.model_dump(mode="json", exclude=self._BASE_FIELDS, exclude_none=True) - if self.job_name: - data["job_name"] = self.job_name + harbor = _HarborJobFields.model_validate(self.model_dump(mode="json")) + data = harbor.model_dump(mode="json", exclude_none=True) harbor_env = self.environment.to_harbor_environment() if harbor_env: data["environment"] = harbor_env diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index bfb3c8b28f..ed94ce1092 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field from rock.sdk.bench.models.environment_type import EnvironmentType -from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.job.config import JobEnvironmentConfig class AgentConfig(BaseModel): @@ -21,14 +21,16 @@ class AgentConfig(BaseModel): class OssMirrorConfig(BaseModel): - """OSS artifact mirror configuration (credentials and bucket only). + """OSS artifact mirror configuration. - ``namespace`` / ``experiment_id`` belong on :class:`~rock.sdk.agent.models.job.config.JobConfig` - as top-level Harbor fields, not inside ``oss_mirror``. + ``namespace`` / ``experiment_id`` are synced from ``HarborJobConfig`` + top-level fields via model validators. """ enabled: bool = False oss_bucket: str | None = None + namespace: str | None = None + experiment_id: str | None = None oss_access_key_id: str | None = None oss_access_key_secret: str | None = None oss_region: str | None = None @@ -52,22 +54,16 @@ class EnvironmentConfig(BaseModel): kwargs: dict[str, Any] = Field(default_factory=dict) -class RockEnvironmentConfig(SandboxConfig, EnvironmentConfig): +class RockEnvironmentConfig(JobEnvironmentConfig, EnvironmentConfig): """Unified Rock environment config. - Combines sandbox lifecycle fields (image, memory, cpus, ...) with - harbor environment fields (force_build, override_cpus, ...) in a single - flat block. Rock-specific fields are stripped when serializing to Harbor - YAML via to_harbor_environment(). + Combines job environment fields (setup_commands, file_uploads, auto_stop, env) + from JobEnvironmentConfig with harbor environment fields (force_build, + override_cpus, oss_mirror, etc.) from EnvironmentConfig. + Rock-specific fields are stripped when serializing to Harbor YAML + via to_harbor_environment(). """ - setup_commands: list[str] = Field(default_factory=list) - file_uploads: list[tuple[str, str]] = Field( - default_factory=list, - description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]", - ) - auto_stop: bool = False - def to_harbor_environment(self) -> dict: """Return only harbor-native environment fields, discarding Rock-only fields. diff --git a/rock/sdk/job/__init__.py b/rock/sdk/job/__init__.py index 0e8f8c5ba3..ce026402b2 100644 --- a/rock/sdk/job/__init__.py +++ b/rock/sdk/job/__init__.py @@ -1,24 +1,18 @@ -# Pre-import rock.sdk.bench to resolve a known circular-import issue between -# rock.sdk.job.config (base JobConfig) and rock.sdk.bench.models.job.config -# (Harbor JobConfig, which inherits from the base). Doing this import first -# ensures bench is fully loaded before any rock.sdk.job submodule pulls it in. -import rock.sdk.bench # noqa: F401, I001 - -from rock.sdk.job.config import BashJobConfig, JobConfig -from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient -from rock.sdk.job.api import Job -from rock.sdk.job.operator import Operator, ScatterOperator -from rock.sdk.job.result import ExceptionInfo, JobResult, JobStatus, TrialResult -from rock.sdk.job.trial import AbstractTrial, register_trial - # Auto-register BashTrial (safe: no bench dependency). # HarborTrial is registered by rock.sdk.bench.__init__ to avoid a circular # import when rock.sdk.job is triggered mid-bench-load. import rock.sdk.job.trial.bash # noqa: F401 +from rock.sdk.job.api import Job +from rock.sdk.job.config import BashJobConfig, JobConfig, JobEnvironmentConfig +from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient +from rock.sdk.job.operator import Operator, ScatterOperator +from rock.sdk.job.result import ExceptionInfo, JobResult, JobStatus, TrialResult +from rock.sdk.job.trial import AbstractTrial, register_trial __all__ = [ "Job", "JobConfig", + "JobEnvironmentConfig", "BashJobConfig", "JobResult", "JobStatus", diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index 69ac5be970..4ddd9d0c14 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -1,30 +1,39 @@ """Config hierarchy for the Job system. -JobConfig — base config with shared fields for all job types -BashJobConfig — simple script execution +JobEnvironmentConfig — sandbox config + job-level environment fields +JobConfig — base config with shared job-scheduling fields +BashJobConfig — simple script execution -Harbor's JobConfig lives in rock.sdk.agent.models.job.config and inherits JobConfig. +Harbor's HarborJobConfig lives in rock.sdk.bench.models.job.config. """ from __future__ import annotations from pydantic import BaseModel, Field -from rock.sdk.bench.models.trial.config import RockEnvironmentConfig +from rock.sdk.sandbox.config import SandboxConfig + + +class JobEnvironmentConfig(SandboxConfig): + """Job environment config — sandbox base fields + job-level environment fields.""" + + setup_commands: list[str] = Field(default_factory=list) + file_uploads: list[tuple[str, str]] = Field( + default_factory=list, + description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]", + ) + auto_stop: bool = False + env: dict[str, str] = Field(default_factory=dict) class JobConfig(BaseModel): """Base config — shared fields for all job types.""" - environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) + environment: JobEnvironmentConfig = Field(default_factory=JobEnvironmentConfig) job_name: str | None = None namespace: str | None = None experiment_id: str | None = None labels: dict[str, str] = Field(default_factory=dict) - auto_stop: bool = False - setup_commands: list[str] = Field(default_factory=list) - file_uploads: list[tuple[str, str]] = Field(default_factory=list) - env: dict[str, str] = Field(default_factory=dict) timeout: int = 3600 diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py index 172002e60c..5b6ab38bfc 100644 --- a/rock/sdk/job/executor.py +++ b/rock/sdk/job/executor.py @@ -157,12 +157,12 @@ async def _do_wait(self, client: TrialClient) -> TrialResult | list[TrialResult] r.exception_info = fail_info return result finally: - if config.auto_stop: + if config.environment.auto_stop: await client.sandbox.close() @staticmethod def _build_session_env(config: JobConfig) -> dict[str, str] | None: """Merge OSS_* env vars from the process with config.env (config wins).""" oss_env = {k: v for k, v in os.environ.items() if k.startswith("OSS")} - merged = {**oss_env, **config.env} + merged = {**oss_env, **config.environment.env} return merged or None diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index 65159589cd..8a2ad4f33b 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -71,7 +71,7 @@ async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialR async def _upload_files(self, sandbox: Sandbox) -> None: """Shared helper: upload all entries in ``config.file_uploads``.""" - for local_path, sandbox_path in self._config.file_uploads: + for local_path, sandbox_path in self._config.environment.file_uploads: obs = await sandbox.fs.upload_dir(source_dir=local_path, target_dir=sandbox_path) if obs.exit_code != 0: raise RuntimeError(f"Failed to upload {local_path} -> {sandbox_path}: {obs.failure_reason}") diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index ea7a040c34..7fac761106 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -23,8 +23,8 @@ async def setup(self, sandbox) -> None: def build(self) -> str: lines = ["#!/bin/bash", "set -e", ""] - if self._config.setup_commands: - for cmd in self._config.setup_commands: + if self._config.environment.setup_commands: + for cmd in self._config.environment.setup_commands: lines.append(f"echo '>>> {cmd[:60]}...'") lines.append(cmd) lines.append("") diff --git a/rock/sdk/job/trial/harbor.py b/rock/sdk/job/trial/harbor.py index 4c8d31b485..839ce77b54 100644 --- a/rock/sdk/job/trial/harbor.py +++ b/rock/sdk/job/trial/harbor.py @@ -62,7 +62,7 @@ async def setup(self, sandbox) -> None: def build(self) -> str: setup_lines: list[str] = [] - for cmd in self._config.setup_commands: + for cmd in self._config.environment.setup_commands: setup_lines.append(f"echo '>>> {cmd[:60]}...'") setup_lines.append(cmd) setup_block = "\n".join(setup_lines) if setup_lines else "echo 'No setup commands'" diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index 0b7914f41d..372ccc5e2c 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -113,8 +113,7 @@ def test_basic_serialization(self): # job_name is re-injected so harbor uses it as the directory name assert data["job_name"] == "test-job" - # Other base fields (experiment_id, etc.) are excluded from harbor YAML - assert "experiment_id" not in data + assert data["experiment_id"] == "test-exp" assert data["n_attempts"] == 2 assert data["agents"][0]["name"] == "terminus-2" @@ -155,8 +154,8 @@ def test_excludes_none_values(self): assert "agent_timeout_multiplier" not in data - def test_labels_excluded_as_base_field(self): - """labels is a base HarborJobConfig field, so it's excluded from harbor YAML.""" + def test_labels_included_in_harbor_yaml(self): + """labels is a base HarborJobConfig field, included in harbor YAML.""" cfg = HarborJobConfig( job_name="labeled-job", experiment_id="test-exp", @@ -165,7 +164,7 @@ def test_labels_excluded_as_base_field(self): yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - assert "labels" not in data + assert data["labels"] == {"step": "42", "env": "prod"} assert data["job_name"] == "labeled-job" def test_path_fields_serialized_as_strings(self): diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py index a306a93c03..8aa0ed8a24 100644 --- a/tests/unit/sdk/agent/test_oss_mirror.py +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -136,14 +136,13 @@ def test_namespace_at_top_level_in_yaml(self): ) data = yaml.safe_load(cfg.to_harbor_yaml()) - # namespace/experiment_id are base HarborJobConfig fields, excluded from harbor YAML - assert "namespace" not in data - assert "experiment_id" not in data + assert data["namespace"] == "my-ns" + assert data["experiment_id"] == "exp-1" oss = data["environment"]["oss_mirror"] assert oss["enabled"] is True assert oss["oss_bucket"] == "test-bucket" - assert "namespace" not in oss - assert "experiment_id" not in oss + assert oss["namespace"] == "my-ns" + assert oss["experiment_id"] == "exp-1" def test_disabled_oss_mirror_excluded_from_yaml(self): """When oss_mirror is default (disabled), it should not clutter the YAML.""" @@ -188,8 +187,8 @@ def test_from_yaml_with_top_level_namespace(self, tmp_path): assert cfg.environment.oss_mirror.enabled is True assert cfg.environment.oss_mirror.oss_bucket == "yaml-bucket" - def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): - """YAML 中 oss_mirror 内多余的 namespace 等字段由 Pydantic 忽略。""" + def test_from_yaml_oss_mirror_with_namespace_and_experiment_id(self, tmp_path): + """YAML 中 oss_mirror 内的 namespace/experiment_id 是合法字段。""" from rock.sdk.bench.models.job.config import HarborJobConfig yaml_content = """\ @@ -212,9 +211,9 @@ def test_from_yaml_extra_keys_under_oss_mirror_ignored(self, tmp_path): cfg = HarborJobConfig.from_yaml(str(yaml_file)) assert cfg.environment.oss_mirror.enabled is True assert cfg.environment.oss_mirror.oss_bucket == "yaml-bucket" - dump = cfg.environment.oss_mirror.model_dump(exclude_none=True) - assert "namespace" not in dump - assert "experiment_id" not in dump + assert cfg.environment.oss_mirror.namespace is not None + # experiment_id 被 _sync_experiment_id validator 同步为顶层值 "test-exp" + assert cfg.environment.oss_mirror.experiment_id == "test-exp" def test_from_yaml_without_oss_mirror(self, tmp_path): from rock.sdk.bench.models.job.config import HarborJobConfig @@ -284,9 +283,8 @@ def test_enable_then_serialize_roundtrip(self): oss_endpoint="oss-ap-southeast-1.aliyuncs.com", ) data = yaml.safe_load(cfg.to_harbor_yaml()) - # namespace/experiment_id are base fields, excluded from harbor YAML - assert "namespace" not in data - assert "experiment_id" not in data + assert data["namespace"] == "rt-ns" + assert data["experiment_id"] == "rt-exp" oss = data["environment"]["oss_mirror"] assert oss["enabled"] is True assert oss["oss_bucket"] == "rt-bucket" diff --git a/tests/unit/sdk/job/test_blue_green_equivalence.py b/tests/unit/sdk/job/test_blue_green_equivalence.py index df78a74c19..6019cf6af9 100644 --- a/tests/unit/sdk/job/test_blue_green_equivalence.py +++ b/tests/unit/sdk/job/test_blue_green_equivalence.py @@ -10,7 +10,6 @@ import json from unittest.mock import AsyncMock, MagicMock, patch -import rock.sdk.bench # noqa: F401 from rock.sdk.bench.models.job.config import ( HarborJobConfig, RegistryDatasetConfig, diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py index 1fad0c7857..911ea501aa 100644 --- a/tests/unit/sdk/job/test_cli_job.py +++ b/tests/unit/sdk/job/test_cli_job.py @@ -9,7 +9,6 @@ import pytest -import rock.sdk.bench # pre-import to avoid circular # noqa: F401 from rock.cli.command.job import JobCommand @@ -143,7 +142,7 @@ async def test_bash_creates_bash_job_config_and_runs(): assert config_arg.environment.memory == "4g" assert config_arg.environment.cpus == 2.0 assert config_arg.timeout == 600 - assert config_arg.auto_stop is True + assert config_arg.environment.auto_stop is True mock_instance.run.assert_awaited_once() @@ -182,7 +181,7 @@ async def test_bash_with_file_upload(): await cmd.arun(args) config_arg = MockJob.call_args[0][0] - assert config_arg.file_uploads == [("/tmp/src", "/root/target")] + assert config_arg.environment.file_uploads == [("/tmp/src", "/root/target")] async def test_bash_requires_script_or_script_content(): @@ -244,7 +243,7 @@ async def test_harbor_loads_from_yaml(): config_arg = MockJob.call_args[0][0] assert isinstance(config_arg, HarborJobConfig) assert config_arg.experiment_id == "exp-123" - assert config_arg.auto_stop is True + assert config_arg.environment.auto_stop is True finally: Path(yaml_path).unlink(missing_ok=True) diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 292b63d17a..052e6c421d 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -19,7 +19,7 @@ TemplateConfig, VerifierConfig, ) -from rock.sdk.job.config import BashJobConfig, JobConfig +from rock.sdk.job.config import BashJobConfig, JobConfig, JobEnvironmentConfig # --------------------------------------------------------------------------- # JobConfig (base) @@ -29,29 +29,31 @@ class TestJobConfig: def test_defaults(self): cfg = JobConfig() - assert isinstance(cfg.environment, RockEnvironmentConfig) + assert isinstance(cfg.environment, JobEnvironmentConfig) assert cfg.job_name is None assert cfg.namespace is None assert cfg.experiment_id is None assert cfg.labels == {} - assert cfg.auto_stop is False - assert cfg.setup_commands == [] - assert cfg.file_uploads == [] - assert cfg.env == {} assert cfg.timeout == 3600 + assert cfg.environment.auto_stop is False + assert cfg.environment.setup_commands == [] + assert cfg.environment.file_uploads == [] + assert cfg.environment.env == {} def test_custom_values(self): - env = RockEnvironmentConfig(image="ubuntu:22.04") + env = JobEnvironmentConfig( + image="ubuntu:22.04", + setup_commands=["pip install foo"], + file_uploads=[("/local/file.py", "/sandbox/file.py")], + env={"MY_VAR": "hello"}, + auto_stop=True, + ) cfg = JobConfig( environment=env, job_name="my-job", namespace="team-a", experiment_id="exp-001", labels={"step": "42"}, - auto_stop=True, - setup_commands=["pip install foo"], - file_uploads=[("/local/file.py", "/sandbox/file.py")], - env={"MY_VAR": "hello"}, timeout=7200, ) assert cfg.environment.image == "ubuntu:22.04" @@ -59,10 +61,10 @@ def test_custom_values(self): assert cfg.namespace == "team-a" assert cfg.experiment_id == "exp-001" assert cfg.labels == {"step": "42"} - assert cfg.auto_stop is True - assert cfg.setup_commands == ["pip install foo"] - assert cfg.file_uploads == [("/local/file.py", "/sandbox/file.py")] - assert cfg.env == {"MY_VAR": "hello"} + assert cfg.environment.auto_stop is True + assert cfg.environment.setup_commands == ["pip install foo"] + assert cfg.environment.file_uploads == [("/local/file.py", "/sandbox/file.py")] + assert cfg.environment.env == {"MY_VAR": "hello"} assert cfg.timeout == 7200 def test_is_base_model(self): @@ -178,43 +180,34 @@ def test_custom_harbor_fields(self): class TestHarborJobConfigToHarborYaml: - def test_excludes_rock_fields(self): - """Rock-level fields (job_name, namespace, etc.) must NOT appear in Harbor YAML. - - Note: 'environment' is excluded from _ROCK_FIELDS dump, but harbor - environment fields are re-injected via to_harbor_environment(), so - the 'environment' key *may* appear with harbor-native fields only. - """ + def test_excludes_rock_fields_keeps_harbor_shared_fields(self): + """Rock-only fields must not appear, but Harbor-shared fields must be present.""" cfg = HarborJobConfig( - job_name="should-not-appear", - namespace="should-not-appear", - experiment_id="should-not-appear", + job_name="test-job", + namespace="my-ns", + experiment_id="my-exp", labels={"step": "1"}, - auto_stop=True, - setup_commands=["pip install foo"], - file_uploads=[("/a", "/b")], - env={"KEY": "VAL"}, + environment=RockEnvironmentConfig( + auto_stop=True, + setup_commands=["pip install foo"], + file_uploads=[("/a", "/b")], + env={"KEY": "VAL"}, + ), timeout=999, n_attempts=2, debug=True, ) yaml_str = cfg.to_harbor_yaml() data = yaml.safe_load(yaml_str) - # Rock-only fields must be absent from Harbor YAML - # job_name is re-injected so harbor uses it as the directory name - assert data["job_name"] == "should-not-appear" - rock_only = { - "namespace", - "experiment_id", - "labels", - "auto_stop", - "setup_commands", - "file_uploads", - "env", - "timeout", - } - for rock_field in rock_only: - assert rock_field not in data, f"Rock field '{rock_field}' should be excluded from Harbor YAML" + # Shared with Harbor — must be present + assert data["job_name"] == "test-job" + assert data["namespace"] == "my-ns" + assert data["experiment_id"] == "my-exp" + assert data["labels"] == {"step": "1"} + # Rock-only — must be absent + rock_only = {"auto_stop", "setup_commands", "file_uploads", "timeout"} + for field in rock_only: + assert field not in data, f"Rock field '{field}' should be excluded" def test_includes_harbor_fields(self): cfg = HarborJobConfig(experiment_id="test-exp", n_attempts=5, debug=True) @@ -401,9 +394,7 @@ def test_exclude_none_omits_template_when_not_set(self): assert "template" not in data def test_exclude_none_includes_template_when_set(self): - cfg = NativeConfig( - template=TemplateConfig(name="my-agent/my-org/my-dataset", revision="rev1") - ) + cfg = NativeConfig(template=TemplateConfig(name="my-agent/my-org/my-dataset", revision="rev1")) data = cfg.model_dump(mode="json", exclude_none=True) assert "template" in data assert data["template"]["name"] == "my-agent/my-org/my-dataset" @@ -423,31 +414,17 @@ def test_harbor_inherits_base_fields(self): class TestHarborJobConfigAutoStopSync: - """G7: HarborJobConfig.auto_stop and environment.auto_stop must be kept in sync (OR semantics).""" + """auto_stop lives on environment only.""" - def test_environment_auto_stop_propagates_to_top_level(self): + def test_environment_auto_stop_preserved(self): cfg = HarborJobConfig( experiment_id="exp-1", environment=RockEnvironmentConfig(auto_stop=True), ) - assert cfg.auto_stop is True, "top-level auto_stop must pick up environment.auto_stop" - - def test_top_level_auto_stop_propagates_to_environment(self): - cfg = HarborJobConfig(experiment_id="exp-1", auto_stop=True) - assert cfg.environment.auto_stop is True - - def test_both_true_stays_true(self): - cfg = HarborJobConfig( - experiment_id="exp-1", - auto_stop=True, - environment=RockEnvironmentConfig(auto_stop=True), - ) - assert cfg.auto_stop is True assert cfg.environment.auto_stop is True - def test_both_false_stays_false(self): + def test_default_auto_stop_is_false(self): cfg = HarborJobConfig(experiment_id="exp-1") - assert cfg.auto_stop is False assert cfg.environment.auto_stop is False diff --git a/tests/unit/sdk/job/test_executor.py b/tests/unit/sdk/job/test_executor.py index 2173472061..5917b44703 100644 --- a/tests/unit/sdk/job/test_executor.py +++ b/tests/unit/sdk/job/test_executor.py @@ -6,7 +6,6 @@ import pytest -import rock.sdk.bench # pre-import to avoid circular # noqa: F401 import rock.sdk.job.trial.bash # register BashJobConfig -> BashTrial # noqa: F401 from rock.sdk.bench.constants import USER_DEFINED_LOGS from rock.sdk.job.config import BashJobConfig @@ -162,7 +161,9 @@ class TestJobExecutorAutoStop: async def test_auto_stop_true_closes_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): - config = BashJobConfig(script="echo hi", job_name="test", auto_stop=True) + from rock.sdk.job.config import JobEnvironmentConfig + + config = BashJobConfig(script="echo hi", job_name="test", environment=JobEnvironmentConfig(auto_stop=True)) executor = JobExecutor() await executor.run(ScatterOperator(size=1), config) @@ -192,7 +193,9 @@ def test_merges_oss_vars_with_config_env(self, monkeypatch): monkeypatch.delenv(k, raising=False) monkeypatch.setenv("OSS_KEY", "value") - config = BashJobConfig(script="echo hi", env={"X": "1"}) + from rock.sdk.job.config import JobEnvironmentConfig + + config = BashJobConfig(script="echo hi", environment=JobEnvironmentConfig(env={"X": "1"})) merged = JobExecutor._build_session_env(config) assert merged is not None @@ -205,7 +208,9 @@ def test_config_env_overrides_oss(self, monkeypatch): monkeypatch.delenv(k, raising=False) monkeypatch.setenv("OSS_KEY", "process_val") - config = BashJobConfig(script="echo hi", env={"OSS_KEY": "config_val"}) + from rock.sdk.job.config import JobEnvironmentConfig + + config = BashJobConfig(script="echo hi", environment=JobEnvironmentConfig(env={"OSS_KEY": "config_val"})) merged = JobExecutor._build_session_env(config) assert merged is not None diff --git a/tests/unit/sdk/job/test_job.py b/tests/unit/sdk/job/test_job.py index 015b3f05a3..1e5b9cbe10 100644 --- a/tests/unit/sdk/job/test_job.py +++ b/tests/unit/sdk/job/test_job.py @@ -6,8 +6,6 @@ import pytest -import rock.sdk.bench # pre-import to avoid circular # noqa: F401 -import rock.sdk.job.trial.bash # trigger BashTrial registration # noqa: F401 from rock.sdk.job import Job from rock.sdk.job.config import BashJobConfig from rock.sdk.job.operator import ScatterOperator diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index 58e81880d2..e54bc28202 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -6,9 +6,7 @@ from pathlib import Path from unittest.mock import AsyncMock, MagicMock -# Import bench first to avoid circular-import pitfall in rock.sdk.job.config -import rock.sdk.bench # noqa: F401 -from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.config import BashJobConfig, JobEnvironmentConfig from rock.sdk.job.trial.bash import BashTrial from rock.sdk.job.trial.registry import _create_trial @@ -35,7 +33,7 @@ def test_build_basic_script(self): def test_build_with_setup_commands(self): cfg = BashJobConfig( - setup_commands=["pip install -r requirements.txt"], + environment=JobEnvironmentConfig(setup_commands=["pip install -r requirements.txt"]), script="python main.py", ) trial = BashTrial(cfg) @@ -47,7 +45,7 @@ def test_build_with_setup_commands(self): assert out.index("pip install -r requirements.txt") < out.index("python main.py") def test_build_no_script_only_setup(self): - cfg = BashJobConfig(setup_commands=["echo setup"]) + cfg = BashJobConfig(environment=JobEnvironmentConfig(setup_commands=["echo setup"])) trial = BashTrial(cfg) out = trial.build() assert "#!/bin/bash" in out @@ -64,7 +62,9 @@ class TestBashTrialSetup: async def test_setup_uploads_files(self): cfg = BashJobConfig( script="echo hi", - file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], + environment=JobEnvironmentConfig( + file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], + ), ) trial = BashTrial(cfg) mock_sandbox = AsyncMock() diff --git a/tests/unit/sdk/job/test_trial_harbor.py b/tests/unit/sdk/job/test_trial_harbor.py index 445f67dd27..f425cd56ea 100644 --- a/tests/unit/sdk/job/test_trial_harbor.py +++ b/tests/unit/sdk/job/test_trial_harbor.py @@ -5,8 +5,6 @@ import json from unittest.mock import AsyncMock, MagicMock -# Pre-import bench to avoid circular-import pitfalls in rock.sdk.job.config -import rock.sdk.bench # noqa: F401 from rock.sdk.bench.models.job.config import HarborJobConfig from rock.sdk.job.trial.harbor import HarborTrial from rock.sdk.job.trial.registry import _create_trial @@ -44,10 +42,12 @@ def test_build_contains_shebang_and_set_e(self): assert "set -e" in script def test_build_with_setup_commands_includes_them(self): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + cfg = HarborJobConfig( job_name="test", experiment_id="exp-1", - setup_commands=["pip install harbor"], + environment=RockEnvironmentConfig(setup_commands=["pip install harbor"]), ) trial = HarborTrial(cfg) script = trial.build() diff --git a/tests/unit/sdk/job/test_trial_registry.py b/tests/unit/sdk/job/test_trial_registry.py index 7c5972a5c1..e21a636ea3 100644 --- a/tests/unit/sdk/job/test_trial_registry.py +++ b/tests/unit/sdk/job/test_trial_registry.py @@ -8,7 +8,7 @@ # Import bench first to avoid circular-import pitfall in rock.sdk.job.config import rock.sdk.bench # noqa: F401 -from rock.sdk.job.config import JobConfig +from rock.sdk.job.config import JobConfig, JobEnvironmentConfig from rock.sdk.job.result import TrialResult from rock.sdk.job.trial.abstract import AbstractTrial from rock.sdk.job.trial.registry import _TRIAL_REGISTRY, _create_trial, register_trial @@ -69,7 +69,7 @@ async def test_upload_files_iterates_all_entries(self): success_obs = MagicMock() success_obs.exit_code = 0 mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) - cfg = _StubConfig(file_uploads=[("/a", "/b"), ("/c", "/d")]) + cfg = _StubConfig(environment=JobEnvironmentConfig(file_uploads=[("/a", "/b"), ("/c", "/d")])) trial = _StubTrial(cfg) await trial._upload_files(mock_sandbox) @@ -91,7 +91,7 @@ async def test_upload_files_noop_when_empty(self): mock_sandbox.fs.upload_dir.assert_not_called() async def test_upload_files_raises_on_failure(self): - cfg = _StubConfig(file_uploads=[("/a", "/b")]) + cfg = _StubConfig(environment=JobEnvironmentConfig(file_uploads=[("/a", "/b")])) trial = _StubTrial(cfg) mock_sandbox = AsyncMock() failure_obs = MagicMock() From 653b49887df99b3444f2ced11dc1520c8104c05a Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 15:32:11 +0800 Subject: [PATCH 033/226] refactor(envhub): move JobEnvironmentConfig to envhub as EnvironmentConfig (#800) - Create EnvironmentConfig in rock.sdk.envhub.config, a general-purpose environment config class extending SandboxConfig - Remove JobEnvironmentConfig from rock.sdk.job.config, use EnvironmentConfig directly (no backward-compat alias) - Update all imports to use `from rock.sdk.envhub import EnvironmentConfig` - Update tests to reference EnvironmentConfig instead of JobEnvironmentConfig Co-authored-by: Claude Opus 4.6 --- rock/sdk/bench/models/trial/config.py | 4 ++-- rock/sdk/envhub/__init__.py | 3 +++ rock/sdk/envhub/config.py | 23 +++++++++++++++++++++++ rock/sdk/job/__init__.py | 3 +-- rock/sdk/job/config.py | 22 +++++----------------- tests/unit/sdk/job/test_config.py | 7 ++++--- tests/unit/sdk/job/test_executor.py | 12 ++++++------ tests/unit/sdk/job/test_trial_bash.py | 9 +++++---- tests/unit/sdk/job/test_trial_registry.py | 7 ++++--- 9 files changed, 53 insertions(+), 37 deletions(-) create mode 100644 rock/sdk/envhub/config.py diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index ed94ce1092..62fba659ea 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field from rock.sdk.bench.models.environment_type import EnvironmentType -from rock.sdk.job.config import JobEnvironmentConfig +from rock.sdk.envhub import EnvironmentConfig as _EnvConfig class AgentConfig(BaseModel): @@ -54,7 +54,7 @@ class EnvironmentConfig(BaseModel): kwargs: dict[str, Any] = Field(default_factory=dict) -class RockEnvironmentConfig(JobEnvironmentConfig, EnvironmentConfig): +class RockEnvironmentConfig(_EnvConfig, EnvironmentConfig): """Unified Rock environment config. Combines job environment fields (setup_commands, file_uploads, auto_stop, env) diff --git a/rock/sdk/envhub/__init__.py b/rock/sdk/envhub/__init__.py index e69de29bb2..4234ea1263 100644 --- a/rock/sdk/envhub/__init__.py +++ b/rock/sdk/envhub/__init__.py @@ -0,0 +1,3 @@ +from rock.sdk.envhub.config import EnvironmentConfig + +__all__ = ["EnvironmentConfig"] diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py new file mode 100644 index 0000000000..3b07af70a8 --- /dev/null +++ b/rock/sdk/envhub/config.py @@ -0,0 +1,23 @@ +"""General-purpose environment configuration. + +EnvironmentConfig extends SandboxConfig with common environment-level fields +(setup commands, file uploads, environment variables, auto-stop). +""" + +from __future__ import annotations + +from pydantic import Field + +from rock.sdk.sandbox.config import SandboxConfig + + +class EnvironmentConfig(SandboxConfig): + """General environment config — sandbox base fields + environment-level fields.""" + + setup_commands: list[str] = Field(default_factory=list) + file_uploads: list[tuple[str, str]] = Field( + default_factory=list, + description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]", + ) + auto_stop: bool = False + env: dict[str, str] = Field(default_factory=dict) diff --git a/rock/sdk/job/__init__.py b/rock/sdk/job/__init__.py index ce026402b2..dbc087b1a4 100644 --- a/rock/sdk/job/__init__.py +++ b/rock/sdk/job/__init__.py @@ -3,7 +3,7 @@ # import when rock.sdk.job is triggered mid-bench-load. import rock.sdk.job.trial.bash # noqa: F401 from rock.sdk.job.api import Job -from rock.sdk.job.config import BashJobConfig, JobConfig, JobEnvironmentConfig +from rock.sdk.job.config import BashJobConfig, JobConfig from rock.sdk.job.executor import JobClient, JobExecutor, TrialClient from rock.sdk.job.operator import Operator, ScatterOperator from rock.sdk.job.result import ExceptionInfo, JobResult, JobStatus, TrialResult @@ -12,7 +12,6 @@ __all__ = [ "Job", "JobConfig", - "JobEnvironmentConfig", "BashJobConfig", "JobResult", "JobStatus", diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index 4ddd9d0c14..b86a9bbf50 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -1,9 +1,9 @@ """Config hierarchy for the Job system. -JobEnvironmentConfig — sandbox config + job-level environment fields -JobConfig — base config with shared job-scheduling fields -BashJobConfig — simple script execution +JobConfig — base config with shared job-scheduling fields +BashJobConfig — simple script execution +Environment config lives in rock.sdk.envhub.config.EnvironmentConfig. Harbor's HarborJobConfig lives in rock.sdk.bench.models.job.config. """ @@ -11,25 +11,13 @@ from pydantic import BaseModel, Field -from rock.sdk.sandbox.config import SandboxConfig - - -class JobEnvironmentConfig(SandboxConfig): - """Job environment config — sandbox base fields + job-level environment fields.""" - - setup_commands: list[str] = Field(default_factory=list) - file_uploads: list[tuple[str, str]] = Field( - default_factory=list, - description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]", - ) - auto_stop: bool = False - env: dict[str, str] = Field(default_factory=dict) +from rock.sdk.envhub import EnvironmentConfig class JobConfig(BaseModel): """Base config — shared fields for all job types.""" - environment: JobEnvironmentConfig = Field(default_factory=JobEnvironmentConfig) + environment: EnvironmentConfig = Field(default_factory=EnvironmentConfig) job_name: str | None = None namespace: str | None = None experiment_id: str | None = None diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 052e6c421d..2acfcb94de 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -19,7 +19,8 @@ TemplateConfig, VerifierConfig, ) -from rock.sdk.job.config import BashJobConfig, JobConfig, JobEnvironmentConfig +from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.job.config import BashJobConfig, JobConfig # --------------------------------------------------------------------------- # JobConfig (base) @@ -29,7 +30,7 @@ class TestJobConfig: def test_defaults(self): cfg = JobConfig() - assert isinstance(cfg.environment, JobEnvironmentConfig) + assert isinstance(cfg.environment, EnvironmentConfig) assert cfg.job_name is None assert cfg.namespace is None assert cfg.experiment_id is None @@ -41,7 +42,7 @@ def test_defaults(self): assert cfg.environment.env == {} def test_custom_values(self): - env = JobEnvironmentConfig( + env = EnvironmentConfig( image="ubuntu:22.04", setup_commands=["pip install foo"], file_uploads=[("/local/file.py", "/sandbox/file.py")], diff --git a/tests/unit/sdk/job/test_executor.py b/tests/unit/sdk/job/test_executor.py index 5917b44703..c8fcb2e51f 100644 --- a/tests/unit/sdk/job/test_executor.py +++ b/tests/unit/sdk/job/test_executor.py @@ -161,9 +161,9 @@ class TestJobExecutorAutoStop: async def test_auto_stop_true_closes_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): - from rock.sdk.job.config import JobEnvironmentConfig + from rock.sdk.envhub import EnvironmentConfig - config = BashJobConfig(script="echo hi", job_name="test", environment=JobEnvironmentConfig(auto_stop=True)) + config = BashJobConfig(script="echo hi", job_name="test", environment=EnvironmentConfig(auto_stop=True)) executor = JobExecutor() await executor.run(ScatterOperator(size=1), config) @@ -193,9 +193,9 @@ def test_merges_oss_vars_with_config_env(self, monkeypatch): monkeypatch.delenv(k, raising=False) monkeypatch.setenv("OSS_KEY", "value") - from rock.sdk.job.config import JobEnvironmentConfig + from rock.sdk.envhub import EnvironmentConfig - config = BashJobConfig(script="echo hi", environment=JobEnvironmentConfig(env={"X": "1"})) + config = BashJobConfig(script="echo hi", environment=EnvironmentConfig(env={"X": "1"})) merged = JobExecutor._build_session_env(config) assert merged is not None @@ -208,9 +208,9 @@ def test_config_env_overrides_oss(self, monkeypatch): monkeypatch.delenv(k, raising=False) monkeypatch.setenv("OSS_KEY", "process_val") - from rock.sdk.job.config import JobEnvironmentConfig + from rock.sdk.envhub import EnvironmentConfig - config = BashJobConfig(script="echo hi", environment=JobEnvironmentConfig(env={"OSS_KEY": "config_val"})) + config = BashJobConfig(script="echo hi", environment=EnvironmentConfig(env={"OSS_KEY": "config_val"})) merged = JobExecutor._build_session_env(config) assert merged is not None diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index e54bc28202..c4aa91295a 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -6,7 +6,8 @@ from pathlib import Path from unittest.mock import AsyncMock, MagicMock -from rock.sdk.job.config import BashJobConfig, JobEnvironmentConfig +from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.job.config import BashJobConfig from rock.sdk.job.trial.bash import BashTrial from rock.sdk.job.trial.registry import _create_trial @@ -33,7 +34,7 @@ def test_build_basic_script(self): def test_build_with_setup_commands(self): cfg = BashJobConfig( - environment=JobEnvironmentConfig(setup_commands=["pip install -r requirements.txt"]), + environment=EnvironmentConfig(setup_commands=["pip install -r requirements.txt"]), script="python main.py", ) trial = BashTrial(cfg) @@ -45,7 +46,7 @@ def test_build_with_setup_commands(self): assert out.index("pip install -r requirements.txt") < out.index("python main.py") def test_build_no_script_only_setup(self): - cfg = BashJobConfig(environment=JobEnvironmentConfig(setup_commands=["echo setup"])) + cfg = BashJobConfig(environment=EnvironmentConfig(setup_commands=["echo setup"])) trial = BashTrial(cfg) out = trial.build() assert "#!/bin/bash" in out @@ -62,7 +63,7 @@ class TestBashTrialSetup: async def test_setup_uploads_files(self): cfg = BashJobConfig( script="echo hi", - environment=JobEnvironmentConfig( + environment=EnvironmentConfig( file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], ), ) diff --git a/tests/unit/sdk/job/test_trial_registry.py b/tests/unit/sdk/job/test_trial_registry.py index e21a636ea3..5fcaff1d59 100644 --- a/tests/unit/sdk/job/test_trial_registry.py +++ b/tests/unit/sdk/job/test_trial_registry.py @@ -8,7 +8,8 @@ # Import bench first to avoid circular-import pitfall in rock.sdk.job.config import rock.sdk.bench # noqa: F401 -from rock.sdk.job.config import JobConfig, JobEnvironmentConfig +from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.job.config import JobConfig from rock.sdk.job.result import TrialResult from rock.sdk.job.trial.abstract import AbstractTrial from rock.sdk.job.trial.registry import _TRIAL_REGISTRY, _create_trial, register_trial @@ -69,7 +70,7 @@ async def test_upload_files_iterates_all_entries(self): success_obs = MagicMock() success_obs.exit_code = 0 mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) - cfg = _StubConfig(environment=JobEnvironmentConfig(file_uploads=[("/a", "/b"), ("/c", "/d")])) + cfg = _StubConfig(environment=EnvironmentConfig(file_uploads=[("/a", "/b"), ("/c", "/d")])) trial = _StubTrial(cfg) await trial._upload_files(mock_sandbox) @@ -91,7 +92,7 @@ async def test_upload_files_noop_when_empty(self): mock_sandbox.fs.upload_dir.assert_not_called() async def test_upload_files_raises_on_failure(self): - cfg = _StubConfig(environment=JobEnvironmentConfig(file_uploads=[("/a", "/b")])) + cfg = _StubConfig(environment=EnvironmentConfig(file_uploads=[("/a", "/b")])) trial = _StubTrial(cfg) mock_sandbox = AsyncMock() failure_obs = MagicMock() From cd58627a3abb39aa9273f9bd042d5a1fc32a3541 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 15:37:55 +0800 Subject: [PATCH 034/226] Refactor/envhub uploads (#802) * refactor(envhub): rename file_uploads to uploads and support both file and dir uploads file_uploads only supported directory uploads via upload_dir(). Renamed to uploads and added path-type detection: files dispatch to upload_by_path(), directories to upload_dir(). Non-existent paths now raise RuntimeError immediately. Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(envhub): remove setup_commands from EnvironmentConfig setup_commands is no longer needed. Removed the field definition, all usage in BashTrial.build(), HarborTrial.build(), and legacy bench Job script rendering, along with related tests and comments. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- rock/cli/command/job.py | 6 +- rock/sdk/bench/job.py | 26 +++--- rock/sdk/bench/models/job/config.py | 2 +- rock/sdk/bench/models/trial/config.py | 2 +- rock/sdk/envhub/config.py | 8 +- rock/sdk/job/trial/abstract.py | 24 ++++-- rock/sdk/job/trial/bash.py | 5 -- rock/sdk/job/trial/harbor.py | 16 +--- .../agent/test_job_config_serialization.py | 21 ++--- tests/unit/sdk/agent/test_models.py | 5 +- tests/unit/sdk/job/test_cli_job.py | 2 +- tests/unit/sdk/job/test_config.py | 14 ++-- tests/unit/sdk/job/test_trial_bash.py | 31 ++----- tests/unit/sdk/job/test_trial_harbor.py | 18 ---- tests/unit/sdk/job/test_trial_registry.py | 84 +++++++++++++++++-- 15 files changed, 136 insertions(+), 128 deletions(-) diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index fa288ff147..210b6c7490 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -45,16 +45,16 @@ async def _job_run(self, args: argparse.Namespace): if getattr(args, "extra_headers", None): env_kwargs["extra_headers"] = args.extra_headers - file_uploads = [] + uploads = [] if args.local_path: - file_uploads.append((args.local_path, args.target_path)) + uploads.append((args.local_path, args.target_path)) config = BashJobConfig( script=args.script_content, script_path=args.script, environment=RockEnvironmentConfig( **env_kwargs, - file_uploads=file_uploads, + uploads=uploads, auto_stop=True, ), timeout=args.timeout, diff --git a/rock/sdk/bench/job.py b/rock/sdk/bench/job.py index 1ac2a0f442..f2f2284b50 100644 --- a/rock/sdk/bench/job.py +++ b/rock/sdk/bench/job.py @@ -11,6 +11,7 @@ import os import tempfile import uuid +from pathlib import Path from rock.actions import Command, CreateBashSessionRequest, ReadFileRequest from rock.logger import init_logger @@ -44,9 +45,6 @@ # ── Ensure output directory exists ────────────────────────────────── mkdir -p {user_defined_dir} -# ── Setup commands ─────────────────────────────────────────────────── -{setup_commands} - # ── Harbor run ─────────────────────────────────────────────────────── harbor jobs start -c {config_path} """ @@ -55,7 +53,7 @@ class Job: """Execute Harbor benchmark tasks inside ROCK sandboxes. - Unifies setup_commands + harbor run into a single bash script, executed + Unifies harbor run into a single bash script, executed via the sandbox nohup protocol: - ``run()``: Full lifecycle (blocking wait) - ``submit()``: Start and return job_id immediately @@ -170,9 +168,13 @@ async def _prepare_and_start(self): await self._create_session() # 1. Upload user-specified files/dirs - for local_path, sandbox_path in self._config.environment.file_uploads: + for local_path, sandbox_path in self._config.environment.uploads: logger.info(f"Uploading {local_path} -> {sandbox_path}") - await self._sandbox.fs.upload_dir(local_path, sandbox_path) + src = Path(local_path) + if src.is_file(): + await self._sandbox.upload_by_path(file_path=local_path, target_path=sandbox_path) + else: + await self._sandbox.fs.upload_dir(local_path, sandbox_path) # 2. Upload harbor config YAML + run script config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" @@ -196,16 +198,8 @@ async def _prepare_and_start(self): ) def _render_run_script(self, config_path: str) -> str: - """Render the run script (dockerd + setup_commands + harbor run).""" - # Setup commands - setup_lines = [] - for cmd in self._config.environment.setup_commands: - setup_lines.append(f"echo '>>> {cmd[:60]}...'") - setup_lines.append(cmd) - setup_block = "\n".join(setup_lines) if setup_lines else "echo 'No setup commands'" - + """Render the run script (dockerd + harbor run).""" return _RUN_SCRIPT_TEMPLATE.format( - setup_commands=setup_block, config_path=config_path, user_defined_dir=USER_DEFINED_LOGS, ) @@ -285,7 +279,7 @@ def _generate_default_job_name(self) -> None: If job_name is None, generate one with the format: {dataset_name}_{task_name if single task}_{uuid} - + For dataset_name and task_name, only the last segment after "/" is used. """ if self._config.job_name is not None: diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index ed087231f1..c474d9c1a3 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -140,7 +140,7 @@ class _HarborJobFields(BaseModel): """Harbor JobConfig field mirror — used by to_harbor_yaml() for serialization filtering. Fields align with harbor.models.job.config.JobConfig. - ROCK-only fields (SandboxConfig, setup_commands, etc.) are automatically + ROCK-only fields (SandboxConfig, uploads, etc.) are automatically discarded by model_validate. """ diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index 62fba659ea..0aea4769ec 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -57,7 +57,7 @@ class EnvironmentConfig(BaseModel): class RockEnvironmentConfig(_EnvConfig, EnvironmentConfig): """Unified Rock environment config. - Combines job environment fields (setup_commands, file_uploads, auto_stop, env) + Combines job environment fields (uploads, auto_stop, env) from JobEnvironmentConfig with harbor environment fields (force_build, override_cpus, oss_mirror, etc.) from EnvironmentConfig. Rock-specific fields are stripped when serializing to Harbor YAML diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py index 3b07af70a8..62c1ddfdc8 100644 --- a/rock/sdk/envhub/config.py +++ b/rock/sdk/envhub/config.py @@ -1,7 +1,7 @@ """General-purpose environment configuration. EnvironmentConfig extends SandboxConfig with common environment-level fields -(setup commands, file uploads, environment variables, auto-stop). +(uploads, environment variables, auto-stop). """ from __future__ import annotations @@ -14,10 +14,10 @@ class EnvironmentConfig(SandboxConfig): """General environment config — sandbox base fields + environment-level fields.""" - setup_commands: list[str] = Field(default_factory=list) - file_uploads: list[tuple[str, str]] = Field( + uploads: list[tuple[str, str]] = Field( default_factory=list, - description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]", + description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]. " + "Automatically detects file vs directory and uses the appropriate upload method.", ) auto_stop: bool = False env: dict[str, str] = Field(default_factory=dict) diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index 8a2ad4f33b..ebe6ffd1d5 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -6,6 +6,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from pathlib import Path from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -70,8 +71,21 @@ async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialR """ async def _upload_files(self, sandbox: Sandbox) -> None: - """Shared helper: upload all entries in ``config.file_uploads``.""" - for local_path, sandbox_path in self._config.environment.file_uploads: - obs = await sandbox.fs.upload_dir(source_dir=local_path, target_dir=sandbox_path) - if obs.exit_code != 0: - raise RuntimeError(f"Failed to upload {local_path} -> {sandbox_path}: {obs.failure_reason}") + """Shared helper: upload all entries in ``config.uploads``. + + Automatically detects file vs directory and dispatches accordingly: + - file → ``sandbox.upload_by_path()`` + - dir → ``sandbox.fs.upload_dir()`` + """ + for local_path, sandbox_path in self._config.environment.uploads: + src = Path(local_path) + if src.is_file(): + resp = await sandbox.upload_by_path(file_path=local_path, target_path=sandbox_path) + if not resp.success: + raise RuntimeError(f"Failed to upload file {local_path} -> {sandbox_path}: {resp.message}") + elif src.is_dir(): + obs = await sandbox.fs.upload_dir(source_dir=local_path, target_dir=sandbox_path) + if obs.exit_code != 0: + raise RuntimeError(f"Failed to upload dir {local_path} -> {sandbox_path}: {obs.failure_reason}") + else: + raise RuntimeError(f"Upload source not found or unsupported: {local_path}") diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index 7fac761106..b4662439cd 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -23,11 +23,6 @@ async def setup(self, sandbox) -> None: def build(self) -> str: lines = ["#!/bin/bash", "set -e", ""] - if self._config.environment.setup_commands: - for cmd in self._config.environment.setup_commands: - lines.append(f"echo '>>> {cmd[:60]}...'") - lines.append(cmd) - lines.append("") if self._config.script: lines.append(self._config.script) return "\n".join(lines) diff --git a/rock/sdk/job/trial/harbor.py b/rock/sdk/job/trial/harbor.py index 839ce77b54..0522ab40d3 100644 --- a/rock/sdk/job/trial/harbor.py +++ b/rock/sdk/job/trial/harbor.py @@ -1,8 +1,8 @@ """HarborTrial — execute a Harbor benchmark job inside a sandbox. -Extracted from rock.sdk.bench.job.Job. Combines dockerd startup, setup -commands, and ``harbor jobs start -c`` into a single bash script executed -by the JobExecutor via the sandbox nohup protocol. +Extracted from rock.sdk.bench.job.Job. Combines dockerd startup and +``harbor jobs start -c`` into a single bash script executed by the +JobExecutor via the sandbox nohup protocol. """ from __future__ import annotations @@ -40,9 +40,6 @@ # ── Ensure output directory exists ────────────────────────────────── mkdir -p {user_defined_dir} -# ── Setup commands ─────────────────────────────────────────────────── -{setup_commands} - # ── Harbor run ─────────────────────────────────────────────────────── harbor jobs start -c {config_path} """ @@ -61,15 +58,8 @@ async def setup(self, sandbox) -> None: await sandbox.write_file_by_path(yaml_content, config_path) def build(self) -> str: - setup_lines: list[str] = [] - for cmd in self._config.environment.setup_commands: - setup_lines.append(f"echo '>>> {cmd[:60]}...'") - setup_lines.append(cmd) - setup_block = "\n".join(setup_lines) if setup_lines else "echo 'No setup commands'" - config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" return _HARBOR_SCRIPT_TEMPLATE.format( - setup_commands=setup_block, config_path=config_path, user_defined_dir=USER_DEFINED_LOGS, ) diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index 372ccc5e2c..25c34cfd56 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -38,8 +38,7 @@ def test_inherits_harbor_env_fields(self): def test_job_level_fields(self): env = RockEnvironmentConfig() assert env.env == {} - assert env.setup_commands == [] - assert env.file_uploads == [] + assert env.uploads == [] assert env.auto_stop is False def test_env_field(self): @@ -72,13 +71,11 @@ def test_excludes_rock_sandbox_fields(self): def test_excludes_job_level_fields(self): env = RockEnvironmentConfig( - setup_commands=["pip install x"], - file_uploads=[("a", "b")], + uploads=[("a", "b")], auto_stop=True, ) result = env.to_harbor_environment() - assert "setup_commands" not in result - assert "file_uploads" not in result + assert "uploads" not in result assert "auto_stop" not in result def test_env_passes_through_to_harbor(self): @@ -97,7 +94,7 @@ def test_empty_config_excludes_rock_fields(self): env = RockEnvironmentConfig() result = env.to_harbor_environment() assert "image" not in result - assert "setup_commands" not in result + assert "uploads" not in result class TestHarborJobConfigToHarborYaml: @@ -121,8 +118,7 @@ def test_excludes_rock_fields(self): cfg = HarborJobConfig( experiment_id="test-exp", environment=RockEnvironmentConfig( - setup_commands=["pip install harbor"], - file_uploads=[("local.txt", "/sandbox/remote.txt")], + uploads=[("local.txt", "/sandbox/remote.txt")], env={"API_KEY": "sk-xxx"}, auto_stop=True, image="my-image:latest", @@ -134,14 +130,12 @@ def test_excludes_rock_fields(self): # Rock fields must not appear at top level assert "sandbox_config" not in data - assert "setup_commands" not in data - assert "file_uploads" not in data + assert "uploads" not in data assert "sandbox_env" not in data assert "auto_stop_sandbox" not in data assert "auto_stop" not in data # environment block should only contain harbor fields assert "environment" not in data or "image" not in data.get("environment", {}) - assert "environment" not in data or "setup_commands" not in data.get("environment", {}) def test_excludes_none_values(self): cfg = HarborJobConfig( @@ -261,8 +255,6 @@ def test_from_yaml_with_environment_block(self, tmp_path): cpus: 8 env: OPENAI_API_KEY: sk-xxx - setup_commands: - - pip install harbor auto_stop: true agents: - name: terminus-2 @@ -274,7 +266,6 @@ def test_from_yaml_with_environment_block(self, tmp_path): assert cfg.environment.image == "my-image:latest" assert cfg.environment.memory == "32g" assert cfg.environment.env == {"OPENAI_API_KEY": "sk-xxx"} - assert cfg.environment.setup_commands == ["pip install harbor"] assert cfg.environment.auto_stop is True def test_from_yaml_with_local_dataset(self, tmp_path): diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index 324877b5a0..b968ba47e3 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -234,8 +234,7 @@ def test_defaults(self): def test_environment_defaults(self): cfg = HarborJobConfig(experiment_id="test-exp") - assert cfg.environment.setup_commands == [] - assert cfg.environment.file_uploads == [] + assert cfg.environment.uploads == [] assert cfg.environment.env == {} assert cfg.environment.auto_stop is False @@ -246,13 +245,11 @@ def test_with_full_config(self): n_attempts=2, agents=[AgentConfig(name="terminus-2", model_name="hosted_vllm/m")], datasets=[RegistryDatasetConfig(registry=RemoteRegistryInfo(), name="terminal-bench", version="2.0")], - environment=RockEnvironmentConfig(setup_commands=["pip install harbor"]), ) assert cfg.job_name == "test-job" assert cfg.n_attempts == 2 assert len(cfg.agents) == 1 assert cfg.agents[0].name == "terminus-2" - assert cfg.environment.setup_commands == ["pip install harbor"] class TestPublicAPI: diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py index 911ea501aa..cec098b61d 100644 --- a/tests/unit/sdk/job/test_cli_job.py +++ b/tests/unit/sdk/job/test_cli_job.py @@ -181,7 +181,7 @@ async def test_bash_with_file_upload(): await cmd.arun(args) config_arg = MockJob.call_args[0][0] - assert config_arg.environment.file_uploads == [("/tmp/src", "/root/target")] + assert config_arg.environment.uploads == [("/tmp/src", "/root/target")] async def test_bash_requires_script_or_script_content(): diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 2acfcb94de..2d467b6781 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -37,15 +37,13 @@ def test_defaults(self): assert cfg.labels == {} assert cfg.timeout == 3600 assert cfg.environment.auto_stop is False - assert cfg.environment.setup_commands == [] - assert cfg.environment.file_uploads == [] + assert cfg.environment.uploads == [] assert cfg.environment.env == {} def test_custom_values(self): env = EnvironmentConfig( image="ubuntu:22.04", - setup_commands=["pip install foo"], - file_uploads=[("/local/file.py", "/sandbox/file.py")], + uploads=[("/local/file.py", "/sandbox/file.py")], env={"MY_VAR": "hello"}, auto_stop=True, ) @@ -63,8 +61,7 @@ def test_custom_values(self): assert cfg.experiment_id == "exp-001" assert cfg.labels == {"step": "42"} assert cfg.environment.auto_stop is True - assert cfg.environment.setup_commands == ["pip install foo"] - assert cfg.environment.file_uploads == [("/local/file.py", "/sandbox/file.py")] + assert cfg.environment.uploads == [("/local/file.py", "/sandbox/file.py")] assert cfg.environment.env == {"MY_VAR": "hello"} assert cfg.timeout == 7200 @@ -190,8 +187,7 @@ def test_excludes_rock_fields_keeps_harbor_shared_fields(self): labels={"step": "1"}, environment=RockEnvironmentConfig( auto_stop=True, - setup_commands=["pip install foo"], - file_uploads=[("/a", "/b")], + uploads=[("/a", "/b")], env={"KEY": "VAL"}, ), timeout=999, @@ -206,7 +202,7 @@ def test_excludes_rock_fields_keeps_harbor_shared_fields(self): assert data["experiment_id"] == "my-exp" assert data["labels"] == {"step": "1"} # Rock-only — must be absent - rock_only = {"auto_stop", "setup_commands", "file_uploads", "timeout"} + rock_only = {"auto_stop", "uploads", "timeout"} for field in rock_only: assert field not in data, f"Rock field '{field}' should be excluded" diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index c4aa91295a..19a085bf8f 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -32,27 +32,6 @@ def test_build_basic_script(self): assert "set -e" in out assert "echo hello" in out - def test_build_with_setup_commands(self): - cfg = BashJobConfig( - environment=EnvironmentConfig(setup_commands=["pip install -r requirements.txt"]), - script="python main.py", - ) - trial = BashTrial(cfg) - out = trial.build() - - assert "pip install -r requirements.txt" in out - assert "python main.py" in out - # Setup comes before main script - assert out.index("pip install -r requirements.txt") < out.index("python main.py") - - def test_build_no_script_only_setup(self): - cfg = BashJobConfig(environment=EnvironmentConfig(setup_commands=["echo setup"])) - trial = BashTrial(cfg) - out = trial.build() - assert "#!/bin/bash" in out - assert "set -e" in out - assert "echo setup" in out - # --------------------------------------------------------------------------- # BashTrial.setup() @@ -60,11 +39,15 @@ def test_build_no_script_only_setup(self): class TestBashTrialSetup: - async def test_setup_uploads_files(self): + async def test_setup_uploads_dirs(self, tmp_path): + dir_a = tmp_path / "a" + dir_b = tmp_path / "b" + dir_a.mkdir() + dir_b.mkdir() cfg = BashJobConfig( script="echo hi", environment=EnvironmentConfig( - file_uploads=[("/local/a", "/sandbox/a"), ("/local/b", "/sandbox/b")], + uploads=[(str(dir_a), "/sandbox/a"), (str(dir_b), "/sandbox/b")], ), ) trial = BashTrial(cfg) @@ -74,8 +57,6 @@ async def test_setup_uploads_files(self): await trial.setup(mock_sandbox) assert mock_sandbox.fs.upload_dir.call_count == 2 - mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/local/a", target_dir="/sandbox/a") - mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/local/b", target_dir="/sandbox/b") async def test_setup_reads_script_path(self): expected = "expected content" diff --git a/tests/unit/sdk/job/test_trial_harbor.py b/tests/unit/sdk/job/test_trial_harbor.py index f425cd56ea..163eeef344 100644 --- a/tests/unit/sdk/job/test_trial_harbor.py +++ b/tests/unit/sdk/job/test_trial_harbor.py @@ -41,24 +41,6 @@ def test_build_contains_shebang_and_set_e(self): assert "#!/bin/bash" in script assert "set -e" in script - def test_build_with_setup_commands_includes_them(self): - from rock.sdk.bench.models.trial.config import RockEnvironmentConfig - - cfg = HarborJobConfig( - job_name="test", - experiment_id="exp-1", - environment=RockEnvironmentConfig(setup_commands=["pip install harbor"]), - ) - trial = HarborTrial(cfg) - script = trial.build() - assert "pip install harbor" in script - - def test_build_without_setup_commands_uses_placeholder(self): - cfg = HarborJobConfig(job_name="test", experiment_id="exp-1") - trial = HarborTrial(cfg) - script = trial.build() - assert "No setup commands" in script - # --------------------------------------------------------------------------- # HarborTrial.setup() diff --git a/tests/unit/sdk/job/test_trial_registry.py b/tests/unit/sdk/job/test_trial_registry.py index 5fcaff1d59..63773094fc 100644 --- a/tests/unit/sdk/job/test_trial_registry.py +++ b/tests/unit/sdk/job/test_trial_registry.py @@ -65,34 +65,80 @@ def test_config_reference_held(self): trial = _StubTrial(cfg) assert trial._config is cfg - async def test_upload_files_iterates_all_entries(self): + async def test_upload_dirs_dispatches_to_upload_dir(self, tmp_path): + """Directory entries should call sandbox.fs.upload_dir().""" + dir_a = tmp_path / "a" + dir_b = tmp_path / "b" + dir_a.mkdir() + dir_b.mkdir() + mock_sandbox = AsyncMock() success_obs = MagicMock() success_obs.exit_code = 0 mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) - cfg = _StubConfig(environment=EnvironmentConfig(file_uploads=[("/a", "/b"), ("/c", "/d")])) + cfg = _StubConfig(environment=EnvironmentConfig(uploads=[(str(dir_a), "/b"), (str(dir_b), "/d")])) trial = _StubTrial(cfg) await trial._upload_files(mock_sandbox) assert mock_sandbox.fs.upload_dir.call_count == 2 - mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/a", target_dir="/b") - mock_sandbox.fs.upload_dir.assert_any_call(source_dir="/c", target_dir="/d") + mock_sandbox.upload_by_path.assert_not_called() + + async def test_upload_files_dispatches_to_upload_by_path(self, tmp_path): + """File entries should call sandbox.upload_by_path().""" + file_a = tmp_path / "a.txt" + file_a.write_text("content") + + mock_sandbox = AsyncMock() + upload_resp = MagicMock() + upload_resp.success = True + mock_sandbox.upload_by_path = AsyncMock(return_value=upload_resp) + cfg = _StubConfig(environment=EnvironmentConfig(uploads=[(str(file_a), "/sandbox/a.txt")])) + trial = _StubTrial(cfg) + + await trial._upload_files(mock_sandbox) + + mock_sandbox.upload_by_path.assert_called_once_with(file_path=str(file_a), target_path="/sandbox/a.txt") + mock_sandbox.fs.upload_dir.assert_not_called() + + async def test_upload_mixed_files_and_dirs(self, tmp_path): + """Mixed entries dispatch to the correct method.""" + dir_a = tmp_path / "mydir" + dir_a.mkdir() + file_b = tmp_path / "myfile.txt" + file_b.write_text("data") - async def test_upload_files_noop_when_empty(self): mock_sandbox = AsyncMock() success_obs = MagicMock() success_obs.exit_code = 0 mock_sandbox.fs.upload_dir = AsyncMock(return_value=success_obs) - cfg = _StubConfig(file_uploads=[]) + upload_resp = MagicMock() + upload_resp.success = True + mock_sandbox.upload_by_path = AsyncMock(return_value=upload_resp) + cfg = _StubConfig( + environment=EnvironmentConfig(uploads=[(str(dir_a), "/sandbox/dir"), (str(file_b), "/sandbox/file.txt")]) + ) + trial = _StubTrial(cfg) + + await trial._upload_files(mock_sandbox) + + mock_sandbox.fs.upload_dir.assert_called_once() + mock_sandbox.upload_by_path.assert_called_once() + + async def test_upload_files_noop_when_empty(self): + mock_sandbox = AsyncMock() + cfg = _StubConfig(uploads=[]) trial = _StubTrial(cfg) await trial._upload_files(mock_sandbox) mock_sandbox.fs.upload_dir.assert_not_called() + mock_sandbox.upload_by_path.assert_not_called() - async def test_upload_files_raises_on_failure(self): - cfg = _StubConfig(environment=EnvironmentConfig(file_uploads=[("/a", "/b")])) + async def test_upload_dir_raises_on_failure(self, tmp_path): + dir_a = tmp_path / "a" + dir_a.mkdir() + cfg = _StubConfig(environment=EnvironmentConfig(uploads=[(str(dir_a), "/b")])) trial = _StubTrial(cfg) mock_sandbox = AsyncMock() failure_obs = MagicMock() @@ -103,6 +149,28 @@ async def test_upload_files_raises_on_failure(self): with pytest.raises(RuntimeError, match="disk full"): await trial._upload_files(mock_sandbox) + async def test_upload_file_raises_on_failure(self, tmp_path): + file_a = tmp_path / "a.txt" + file_a.write_text("content") + cfg = _StubConfig(environment=EnvironmentConfig(uploads=[(str(file_a), "/b")])) + trial = _StubTrial(cfg) + mock_sandbox = AsyncMock() + upload_resp = MagicMock() + upload_resp.success = False + upload_resp.message = "upload failed" + mock_sandbox.upload_by_path = AsyncMock(return_value=upload_resp) + + with pytest.raises(RuntimeError, match="upload failed"): + await trial._upload_files(mock_sandbox) + + async def test_upload_nonexistent_path_raises(self): + cfg = _StubConfig(environment=EnvironmentConfig(uploads=[("/nonexistent/path", "/b")])) + trial = _StubTrial(cfg) + mock_sandbox = AsyncMock() + + with pytest.raises(RuntimeError, match="not found or unsupported"): + await trial._upload_files(mock_sandbox) + # --------------------------------------------------------------------------- # Registry From 58cc029c0c276231ec152427d53004667f08b279 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:35:48 +0800 Subject: [PATCH 035/226] feature: add claw-eval bash job demo (#804) * feat(job-cli): add --base-url, --cluster, --env args Register --base-url and --cluster in argparse (handling logic already existed via getattr). Add --env with action=append for repeatable KEY=VALUE pairs, parsed via str.partition("=") to support values containing "=". The env dict is passed to BashJobConfig and injected into the sandbox bash session by JobExecutor._build_session_env(). Co-Authored-By: Claude Opus 4.6 (1M context) * feat(claw-eval): add BashJob infra script with RUN_CMD passthrough Generic infra script for running claw-eval inside a Rock sandbox via BashJob. Handles dockerd startup, optional image pull, and executes the claw-eval command passed via RUN_CMD env var. Includes inline score parsing from run.log output. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(claw-eval): make WORK_DIR configurable, default /workspace Upload target can be /tmp/claw-eval-config instead of /workspace, so the cd path before eval RUN_CMD is now ${WORK_DIR:-/workspace}. Co-Authored-By: Claude Opus 4.6 (1M context) * feat(job-cli): add --xrl-authorization arg Co-Authored-By: Claude Opus 4.6 (1M context) * fix(job): log trial output in executor so CLI users can see script output Co-Authored-By: Claude Opus 4.6 (1M context) * feat(job): add JobConfig.from_yaml for YAML-based config loading Co-Authored-By: Claude Opus 4.6 (1M context) * feat(claw-eval): add BashJobConfig YAML template Co-Authored-By: Claude Opus 4.6 (1M context) * fix(claw-eval): replace real URLs with placeholders in script comments Co-Authored-By: Claude Opus 4.6 (1M context) * feat(claw-eval): add Python SDK runner script Co-Authored-By: Claude Opus 4.6 (1M context) * chore: add .intern to gitignore * fix: replace Chinese comments with English in script and tests Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use placeholder URLs in CLI tests Co-Authored-By: Claude Opus 4.6 (1M context) * fix: adapt claw-eval configs to EnvironmentConfig refactor - YAML template: move uploads/env/auto_stop into environment block, rename file_uploads to uploads - CLI tests: fix assertion paths from config_arg.env to config_arg.environment.env Co-Authored-By: Claude Opus 4.6 (1M context) * chore: move claw_eval from examples/agents to examples/evaluation Co-Authored-By: Claude Opus 4.6 (1M context) * feat(claw-eval): add config template for model/judge settings Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- .gitignore | 2 +- .../claw_eval/claw_eval_bashjob.yaml.template | 17 +++ .../claw_eval/claw_eval_config.yaml.template | 15 +++ .../evaluation/claw_eval/run_claw_eval.py | 30 +++++ .../evaluation/claw_eval/run_claw_eval.sh | 66 ++++++++++ rock/cli/command/job.py | 23 ++++ rock/sdk/job/config.py | 7 ++ rock/sdk/job/executor.py | 2 + tests/unit/sdk/job/test_cli_job.py | 113 ++++++++++++++++++ 9 files changed, 274 insertions(+), 1 deletion(-) create mode 100644 examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template create mode 100644 examples/evaluation/claw_eval/claw_eval_config.yaml.template create mode 100644 examples/evaluation/claw_eval/run_claw_eval.py create mode 100755 examples/evaluation/claw_eval/run_claw_eval.sh diff --git a/.gitignore b/.gitignore index 95e72a518a..ecfc532486 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,4 @@ docs/superpowers/ .env *.db *.intern.yaml -*.intern.md \ No newline at end of file +*.intern.md diff --git a/examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template b/examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template new file mode 100644 index 0000000000..66607eb081 --- /dev/null +++ b/examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template @@ -0,0 +1,17 @@ +script_path: run_claw_eval.sh +timeout: 7200 + +environment: + image: "" + base_url: "" + cluster: "" + xrl_authorization: "" + memory: "64g" + cpus: 16 + auto_stop: true + uploads: + - [".", "/tmp/claw-eval-config"] + env: + SERP_DEV_KEY: "" + AGENT_IMAGE: "" + RUN_CMD: "claw-eval run --task tasks/T01zh_email_triage --sandbox --config /tmp/claw-eval-config/claw_eval_config.yaml --trace-dir /data/logs/user-defined/traces" diff --git a/examples/evaluation/claw_eval/claw_eval_config.yaml.template b/examples/evaluation/claw_eval/claw_eval_config.yaml.template new file mode 100644 index 0000000000..f7ab723238 --- /dev/null +++ b/examples/evaluation/claw_eval/claw_eval_config.yaml.template @@ -0,0 +1,15 @@ +model: + api_key: + base_url: + model_id: +judge: + api_key: + base_url: + model_id: + enabled: true +defaults: + trace_dir: /data/logs/user-defined/traces + tasks_dir: tasks +sandbox: + enabled: true + image: "" diff --git a/examples/evaluation/claw_eval/run_claw_eval.py b/examples/evaluation/claw_eval/run_claw_eval.py new file mode 100644 index 0000000000..b913077f7d --- /dev/null +++ b/examples/evaluation/claw_eval/run_claw_eval.py @@ -0,0 +1,30 @@ +"""Run claw-eval via BashJob SDK. + +Usage: + cd examples/agents/claw_eval + cp claw_eval_bashjob.yaml.template claw_eval_bashjob.yaml + # fill in real values + python run_claw_eval.py + # or specify a config path: + python run_claw_eval.py my_config.yaml +""" + +import asyncio +import os +import sys +from pathlib import Path + +from rock.sdk.job import Job +from rock.sdk.job.config import BashJobConfig + + +async def main() -> None: + config_path = sys.argv[1] if len(sys.argv) > 1 else "claw_eval_bashjob.yaml" + config = BashJobConfig.from_yaml(config_path) + result = await Job(config).run() + print(f"Job completed: status={result.status}, trials={len(result.trial_results)}") + + +if __name__ == "__main__": + os.chdir(Path(__file__).resolve().parent) + asyncio.run(main()) diff --git a/examples/evaluation/claw_eval/run_claw_eval.sh b/examples/evaluation/claw_eval/run_claw_eval.sh new file mode 100755 index 0000000000..4f58af81d6 --- /dev/null +++ b/examples/evaluation/claw_eval/run_claw_eval.sh @@ -0,0 +1,66 @@ +#!/bin/bash +# Claw-eval BashJob infra script. +# +# Environment variables (passed via `rock job run --env`): +# RUN_CMD — claw-eval command to execute (required) +# e.g. "claw-eval batch --parallel 4 --sandbox --config /tmp/claw-eval-config/config.yaml" +# AGENT_IMAGE — Docker image to pull before running (optional) +# WORK_DIR — working directory before eval RUN_CMD (optional, default /workspace) +# SERP_DEV_KEY — API key forwarded to claw-eval (optional) +# +# Usage: +# rock job run --type bash \ +# --script examples/agents/claw_eval/run_claw_eval.sh \ +# --image "" \ +# --base-url "" \ +# --cluster "" \ +# --memory 64g --cpus 16 --timeout 7200 \ +# --env "SERP_DEV_KEY=" \ +# --env "AGENT_IMAGE=" \ +# --env "RUN_CMD=claw-eval batch --parallel 4 --sandbox --config /tmp/claw-eval-config/config.yaml --trace-dir /data/logs/user-defined/traces" \ +# --local-path . --target-path /tmp/claw-eval-config + +set -eo pipefail + +LOG_DIR="/data/logs/user-defined" + +# ── 1. Prepare log directory ─────────────────────────────── +mkdir -p "$LOG_DIR" + +# ── 2. Start dockerd (DinD) ─────────────────────────────── +if command -v docker &>/dev/null; then + if ! pgrep -x dockerd &>/dev/null; then + echo "Starting dockerd..." + nohup dockerd &>/var/log/dockerd.log & + fi + for i in $(seq 1 60); do + docker info &>/dev/null && { echo "dockerd ready"; break; } + sleep 1 + [ "$i" -eq 60 ] && echo "WARN: dockerd failed to start within 60s" + done +fi + +# ── 3. Pull agent image (optional) ──────────────────────── +[ -n "$AGENT_IMAGE" ] && docker pull "$AGENT_IMAGE" + +# ── 4. Execute RUN_CMD ──────────────────────────────────── +[ -z "$RUN_CMD" ] && { echo "ERROR: RUN_CMD environment variable is not set"; exit 1; } +cd "${WORK_DIR:-/workspace}" +eval "$RUN_CMD" 2>&1 | tee "$LOG_DIR/run.log" + +# ── 5. Score summary (parse run.log) ────────────────────── +echo "=== Score Summary ===" +LOG_FILE="$LOG_DIR/run.log" +TEXT=$(cat "$LOG_FILE") +get_float() { echo "$TEXT" | grep -oP "$1:\s+\K[\d.]+" | tail -1; } +TASK_SCORE=$(get_float "task_score") +COMPLETION=$(get_float "completion") +ROBUSTNESS=$(get_float "robustness") +COMMUNICATION=$(get_float "communication") +SAFETY=$(get_float "safety") +PASSED=$(echo "$TEXT" | grep -oP 'passed:\s+\K(True|False)' | tail -1) +WALL_TIME=$(echo "$TEXT" | grep -oP 'wall=\K[\d.]+' | tail -1) +TOKENS=$(echo "$TEXT" | grep -oP 'tokens=\K\d+' | tail -1) +echo "task_score=${TASK_SCORE:-N/A} completion=${COMPLETION:-N/A} robustness=${ROBUSTNESS:-N/A}" +echo "communication=${COMMUNICATION:-N/A} safety=${SAFETY:-N/A} passed=${PASSED:-N/A}" +echo "wall_time=${WALL_TIME:-N/A}s tokens=${TOKENS:-N/A}" diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index 210b6c7490..7142c662a8 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -44,11 +44,19 @@ async def _job_run(self, args: argparse.Namespace): env_kwargs["cluster"] = args.cluster if getattr(args, "extra_headers", None): env_kwargs["extra_headers"] = args.extra_headers + if getattr(args, "xrl_authorization", None): + env_kwargs["xrl_authorization"] = args.xrl_authorization uploads = [] if args.local_path: uploads.append((args.local_path, args.target_path)) + env = {} + if getattr(args, "env", None): + for item in args.env: + key, _, value = item.partition("=") + env[key] = value + config = BashJobConfig( script=args.script_content, script_path=args.script, @@ -56,6 +64,7 @@ async def _job_run(self, args: argparse.Namespace): **env_kwargs, uploads=uploads, auto_stop=True, + env=env, ), timeout=args.timeout, ) @@ -105,3 +114,17 @@ async def add_parser_to(subparsers: argparse._SubParsersAction): run_parser.add_argument("--timeout", type=int, default=3600, help="Timeout in seconds") run_parser.add_argument("--local-path", default=None, help="Local dir to upload") run_parser.add_argument("--target-path", default="/root/job", help="Target dir in sandbox") + run_parser.add_argument("--base-url", default=None, help="Admin service base URL") + run_parser.add_argument("--cluster", default=None, help="Cluster name (e.g. vpc-sg-sl-a)") + run_parser.add_argument( + "--env", + action="append", + default=None, + metavar="KEY=VALUE", + help="Environment variable, repeatable (e.g. --env FOO=bar --env BAZ=qux)", + ) + run_parser.add_argument( + "--xrl-authorization", + default=None, + help="XRL authorization token", + ) diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index b86a9bbf50..2e89da6885 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -9,6 +9,7 @@ from __future__ import annotations +import yaml from pydantic import BaseModel, Field from rock.sdk.envhub import EnvironmentConfig @@ -24,6 +25,12 @@ class JobConfig(BaseModel): labels: dict[str, str] = Field(default_factory=dict) timeout: int = 3600 + @classmethod + def from_yaml(cls, path: str) -> JobConfig: + with open(path) as f: + data = yaml.safe_load(f) + return cls(**data) + class BashJobConfig(JobConfig): """Config for a simple bash script job.""" diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py index 5b6ab38bfc..be03cab367 100644 --- a/rock/sdk/job/executor.py +++ b/rock/sdk/job/executor.py @@ -139,6 +139,8 @@ async def _do_wait(self, client: TrialClient) -> TrialResult | list[TrialResult] response_limited_bytes_in_nohup=None, ) exit_code = obs.exit_code if obs.exit_code is not None else 1 + if obs.output: + logger.info(f"Trial output (job={config.job_name}):\n{obs.output}") result = await client.trial.collect(client.sandbox, obs.output or "", exit_code) # G5: populate raw_output / exit_code on every TrialResult so they surface in JobResult iter_results = result if isinstance(result, list) else [result] diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py index cec098b61d..b46a405f7e 100644 --- a/tests/unit/sdk/job/test_cli_job.py +++ b/tests/unit/sdk/job/test_cli_job.py @@ -109,6 +109,7 @@ def _bash_args(**overrides): base_url=None, cluster=None, extra_headers=None, + env=None, ) defaults.update(overrides) return argparse.Namespace(**defaults) @@ -273,6 +274,118 @@ async def test_harbor_image_override(): Path(yaml_path).unlink(missing_ok=True) +async def test_parser_accepts_base_url_and_cluster(): + p = await _build_parser() + args = p.parse_args( + [ + "job", + "run", + "--script-content", + "echo hi", + "--base-url", + "http://example.com", + "--cluster", + "test-cluster-a", + ] + ) + assert args.base_url == "http://example.com" + assert args.cluster == "test-cluster-a" + + +async def test_bash_passes_base_url_and_cluster_to_environment(): + args = _bash_args( + script_content="echo hello", + base_url="http://example.com", + cluster="test-cluster-a", + ) + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.environment.base_url == "http://example.com" + assert config_arg.environment.cluster == "test-cluster-a" + + +async def test_parser_accepts_env_args(): + p = await _build_parser() + args = p.parse_args( + [ + "job", + "run", + "--script-content", + "echo hi", + "--env", + "FOO=bar", + "--env", + "BAZ=qux=123", + ] + ) + assert args.env == ["FOO=bar", "BAZ=qux=123"] + + +async def test_parser_env_default_is_none(): + p = await _build_parser() + args = p.parse_args(["job", "run", "--script-content", "echo hi"]) + assert args.env is None + + +async def test_bash_passes_env_to_config(): + args = _bash_args(script_content="echo hello") + args.env = ["SERP_DEV_KEY=abc123", "RUN_CMD=claw-eval batch --parallel 4"] + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.environment.env == { + "SERP_DEV_KEY": "abc123", + "RUN_CMD": "claw-eval batch --parallel 4", + } + + +async def test_bash_env_with_equals_in_value(): + """Values containing '=' should not be split.""" + args = _bash_args(script_content="echo hello") + args.env = ["API_KEY=sk-abc=def=="] + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.environment.env == {"API_KEY": "sk-abc=def=="} + + +async def test_bash_no_env_defaults_to_empty(): + args = _bash_args(script_content="echo hello") + + with patch("rock.sdk.job.Job") as MockJob: + mock_instance = MagicMock() + mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) + MockJob.return_value = mock_instance + + cmd = JobCommand() + await cmd.arun(args) + + config_arg = MockJob.call_args[0][0] + assert config_arg.environment.env == {} + + async def test_unknown_job_command_logs_error(): args = argparse.Namespace(job_command="weird") From d2a30660aaf1282d5a07726a788754d76c416b9d Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:36:02 +0800 Subject: [PATCH 036/226] fix: increase default JobConfig timeout from 3600s to 7200s (#806) The previous 1-hour default was too short for long-running jobs. Increase to 2 hours to reduce timeout failures. Co-authored-by: Claude Opus 4.6 --- rock/sdk/job/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index 2e89da6885..04ac8acc20 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -23,7 +23,7 @@ class JobConfig(BaseModel): namespace: str | None = None experiment_id: str | None = None labels: dict[str, str] = Field(default_factory=dict) - timeout: int = 3600 + timeout: int = 7200 @classmethod def from_yaml(cls, path: str) -> JobConfig: From 07896f9151b1bb72da3cf4bce980d393d1c66405 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:48:51 +0800 Subject: [PATCH 037/226] fix(job): update timeout sentinel and tests after default increased to 7200s (#810) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HarborJobConfig._compute_effective_timeout used 3600 as the sentinel to detect "user didn't explicitly set timeout". After #806 raised the default to 7200, the sentinel always evaluated to true and the validator returned early without computing the effective timeout from agent config. Fix: update sentinel from 3600 → 7200 and align test assertions. Co-authored-by: Claude Sonnet 4.6 --- rock/sdk/bench/models/job/config.py | 4 ++-- tests/unit/sdk/job/test_config.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index c474d9c1a3..75ed70f1ac 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -259,8 +259,8 @@ def _compute_effective_timeout(self): """ from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT - # 3600 is the base JobConfig default; treat as "user didn't touch it". - if self.timeout != 3600: + # 7200 is the base JobConfig default; treat as "user didn't touch it". + if self.timeout != 7200: return self multiplier = self.timeout_multiplier or 1.0 diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 2d467b6781..b65fa70e41 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -35,7 +35,7 @@ def test_defaults(self): assert cfg.namespace is None assert cfg.experiment_id is None assert cfg.labels == {} - assert cfg.timeout == 3600 + assert cfg.timeout == 7200 assert cfg.environment.auto_stop is False assert cfg.environment.uploads == [] assert cfg.environment.env == {} @@ -84,7 +84,7 @@ def test_inherits_job_config(self): def test_defaults(self): cfg = BashJobConfig() # Inherited defaults - assert cfg.timeout == 3600 + assert cfg.timeout == 7200 assert cfg.labels == {} # Own defaults assert cfg.script is None From c8a4efa6f2922b31dc0cbcebf92e6a76051eb184 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:58:23 +0800 Subject: [PATCH 038/226] fix(job): BashTrial.collect should populate raw_output and exit_code (#808) * fix(job): BashTrial.collect should populate raw_output and exit_code Pass output and exit_code directly into TrialResult instead of relying on the executor G5 backfill. The executor already passes both values to collect(), so capturing them here makes the intent explicit. Co-Authored-By: Claude Sonnet 4.6 * chore: add .worktrees/ to .gitignore --------- Co-authored-by: Claude Sonnet 4.6 --- .gitignore | 1 + rock/sdk/job/trial/bash.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index ecfc532486..9c75575e38 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,4 @@ docs/superpowers/ *.db *.intern.yaml *.intern.md +.worktrees/ diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index b4662439cd..47b3587876 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -37,6 +37,8 @@ async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: return TrialResult( task_name=self._config.job_name or "", exception_info=exception_info, + raw_output=output, + exit_code=exit_code, ) From 2b0e3f74c13881753c1cb1aa170caa418bd0bbbf Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 17:37:56 +0800 Subject: [PATCH 039/226] chore: apply ruff format and translate Chinese comments to English (#812) * chore: apply ruff check --fix and ruff format across codebase Co-Authored-By: Claude Sonnet 4.6 (1M context) * chore: translate all Chinese comments and docstrings to English in rock/ Co-Authored-By: Claude Sonnet 4.6 (1M context) --------- Co-authored-by: Claude Sonnet 4.6 (1M context) --- rock/admin/core/db_provider.py | 2 +- rock/admin/core/schema.py | 1 - rock/admin/main.py | 4 +- rock/admin/metrics/monitor.py | 2 +- rock/admin/proto/request.py | 2 +- rock/rocklet/local_api.py | 16 ++--- rock/sandbox/base_actor.py | 2 +- rock/sandbox/base_manager.py | 6 +- rock/sandbox/gem_manager.py | 1 + rock/sandbox/operator/factory.py | 2 +- rock/sandbox/operator/k8s/api_client.py | 8 +-- rock/sandbox/operator/k8s/provider.py | 2 +- rock/sandbox/sandbox_manager.py | 4 +- rock/sandbox/sandbox_meta_store.py | 4 +- rock/sandbox/service/sandbox_proxy_service.py | 12 +--- rock/sdk/job/result.py | 10 ++-- rock/sdk/job/trial/abstract.py | 4 +- rock/sdk/sandbox/agent/openhands.py | 4 +- rock/sdk/sandbox/client.py | 4 +- rock/sdk/sandbox/deploy.py | 2 +- rock/utils/k8s/informer/__init__.py | 2 +- rock/utils/k8s/informer/cache.py | 2 +- rock/utils/k8s/informer/informer.py | 19 ++---- scripts/gen_ddl.py | 2 + tests/integration/conftest.py | 3 +- tests/unit/admin/core/test_sandbox_table.py | 11 ++-- .../admin/core/test_schema_varchar_lengths.py | 12 ++-- tests/unit/admin/core/test_statement_cache.py | 15 ++--- tests/unit/conftest.py | 20 ++++++- .../sandbox/operator/test_k8s_api_client.py | 16 ++--- .../sandbox/operator/test_k8s_operator.py | 4 +- .../sandbox/operator/test_k8s_provider.py | 58 +++++++++---------- tests/unit/sandbox/test_proxy_enhancements.py | 10 +--- tests/unit/sandbox/test_sandbox_meta_store.py | 13 +++-- tests/unit/sdk/model/test_proxy.py | 25 ++++---- 35 files changed, 154 insertions(+), 150 deletions(-) diff --git a/rock/admin/core/db_provider.py b/rock/admin/core/db_provider.py index edb89cc6e1..7d994dfb0d 100644 --- a/rock/admin/core/db_provider.py +++ b/rock/admin/core/db_provider.py @@ -57,5 +57,5 @@ def _convert_url(url: str) -> str: return url.replace("sqlite:///", "sqlite+aiosqlite:///", 1) if url.startswith("postgresql://") or url.startswith("postgres://"): prefix = "postgresql://" if url.startswith("postgresql://") else "postgres://" - return "postgresql+asyncpg://" + url[len(prefix):] + return "postgresql+asyncpg://" + url[len(prefix) :] return url diff --git a/rock/admin/core/schema.py b/rock/admin/core/schema.py index ab0ef83ad5..7e5b78b760 100644 --- a/rock/admin/core/schema.py +++ b/rock/admin/core/schema.py @@ -13,7 +13,6 @@ from sqlalchemy.orm import DeclarativeBase from sqlalchemy.types import JSON - _JSONB_VARIANT = JSON().with_variant(JSONB(), "postgresql") diff --git a/rock/admin/main.py b/rock/admin/main.py index 3e3441bab2..fc979557e5 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -14,14 +14,14 @@ from rock import env_vars from rock.admin.core.db_provider import DatabaseProvider -from rock.admin.core.sandbox_table import SandboxTable from rock.admin.core.ray_service import RayService +from rock.admin.core.sandbox_table import SandboxTable from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service from rock.admin.scheduler.scheduler import SchedulerThread -from rock.config import RockConfig, DatabaseConfig +from rock.config import DatabaseConfig, RockConfig from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager from rock.sandbox.operator.factory import OperatorContext, OperatorFactory diff --git a/rock/admin/metrics/monitor.py b/rock/admin/metrics/monitor.py index 916cd9fd96..d528d38b9c 100644 --- a/rock/admin/metrics/monitor.py +++ b/rock/admin/metrics/monitor.py @@ -48,7 +48,7 @@ def create( pod = get_instance_id() env = env_vars.ROCK_ADMIN_ENV role = env_vars.ROCK_ADMIN_ROLE - logger.info(f"Initializing MetricsCollector with host={host}, port={port}, " f"env={env}, role={role}") + logger.info(f"Initializing MetricsCollector with host={host}, port={port}, env={env}, role={role}") return cls( host=host, port=port, diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index ab27bf22e0..d51de25241 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -118,7 +118,7 @@ class BatchSandboxStatusRequest(BaseModel): class SandboxQueryParams(TypedDict, total=False): - """Sandbox列表查询参数""" + """Query parameters for sandbox list.""" page: str page_size: str diff --git a/rock/rocklet/local_api.py b/rock/rocklet/local_api.py index 7cb390336f..0bdbe41f33 100644 --- a/rock/rocklet/local_api.py +++ b/rock/rocklet/local_api.py @@ -186,7 +186,7 @@ async def portforward(websocket: WebSocket, port: int): f"local_addr={writer.get_extra_info('sockname')}" ) except asyncio.TimeoutError: - logger.error(f"[Portforward] TCP connection timeout: target_port={port}, " f"timeout={TCP_CONNECT_TIMEOUT}s") + logger.error(f"[Portforward] TCP connection timeout: target_port={port}, timeout={TCP_CONNECT_TIMEOUT}s") await websocket.close(code=1011, reason=f"Connection to port {port} timed out") return except OSError as e: @@ -198,7 +198,7 @@ async def portforward(websocket: WebSocket, port: int): return except Exception as e: logger.error( - f"[Portforward] Unexpected TCP error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" + f"[Portforward] Unexpected TCP error: target_port={port}, error_type={type(e).__name__}, error={e}" ) await websocket.close(code=1011, reason=f"Unexpected error: {e}") return @@ -227,9 +227,7 @@ async def ws_to_tcp(): except WebSocketDisconnect as e: logger.info(f"[Portforward] ws->tcp: client disconnected: target_port={port}, code={e.code}") except Exception as e: - logger.debug( - f"[Portforward] ws->tcp error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" - ) + logger.debug(f"[Portforward] ws->tcp error: target_port={port}, error_type={type(e).__name__}, error={e}") finally: writer.close() @@ -250,9 +248,7 @@ async def tcp_to_ws(): f"bytes={len(data)}, total_msgs={tcp_to_ws_msgs}, total_bytes={tcp_to_ws_bytes}" ) except Exception as e: - logger.debug( - f"[Portforward] tcp->ws error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" - ) + logger.debug(f"[Portforward] tcp->ws error: target_port={port}, error_type={type(e).__name__}, error={e}") finally: try: await websocket.close() @@ -263,9 +259,7 @@ async def tcp_to_ws(): try: await asyncio.gather(ws_to_tcp(), tcp_to_ws()) except Exception as e: - logger.debug( - f"[Portforward] Forwarding error: target_port={port}, " f"error_type={type(e).__name__}, error={e}" - ) + logger.debug(f"[Portforward] Forwarding error: target_port={port}, error_type={type(e).__name__}, error={e}") finally: writer.close() try: diff --git a/rock/sandbox/base_actor.py b/rock/sandbox/base_actor.py index 058f7dee17..0111083736 100644 --- a/rock/sandbox/base_actor.py +++ b/rock/sandbox/base_actor.py @@ -85,7 +85,7 @@ def _init_monitor(self): env = self._env role = self._role self.host = host - logger.info(f"Initializing MetricsCollector with host={host}, port={port}, " f"env={env}, role={role}") + logger.info(f"Initializing MetricsCollector with host={host}, port={port}, env={env}, role={role}") endpoint = self._metrics_endpoint or f"http://{host}:{port}/v1/metrics" self.otlp_exporter = OTLPMetricExporter(endpoint=endpoint) self.metric_reader = PeriodicExportingMetricReader( diff --git a/rock/sandbox/base_manager.py b/rock/sandbox/base_manager.py index cadee819b5..f361b33c1f 100644 --- a/rock/sandbox/base_manager.py +++ b/rock/sandbox/base_manager.py @@ -10,8 +10,8 @@ from rock.config import RockConfig from rock.deployments.manager import DeploymentManager from rock.logger import init_logger -from rock.utils import get_executor from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.utils import get_executor logger = init_logger(__name__) @@ -108,7 +108,7 @@ async def _collect_and_report_metrics_internal(self): logger.debug(f"Metrics overall report rt:{overall_duration:.4f}s") async def _report_system_resource_metrics(self): - """汇报系统资源指标""" + """Report system resource metrics.""" total_cpu, total_mem, available_cpu, available_mem = await self._collect_system_resource_metrics() self.metrics_monitor.record_gauge_by_name(MetricsConstants.TOTAL_CPU_RESOURCE, total_cpu) self.metrics_monitor.record_gauge_by_name(MetricsConstants.TOTAL_MEM_RESOURCE, total_mem) @@ -116,7 +116,7 @@ async def _report_system_resource_metrics(self): self.metrics_monitor.record_gauge_by_name(MetricsConstants.AVAILABLE_MEM_RESOURCE, available_mem) async def _collect_system_resource_metrics(self): - """收集系统资源指标""" + """Collect system resource metrics.""" cluster_resources = ray.cluster_resources() available_resources = ray.available_resources() total_cpu = cluster_resources.get("CPU", 0) diff --git a/rock/sandbox/gem_manager.py b/rock/sandbox/gem_manager.py index d0ec3d4708..0858852df3 100644 --- a/rock/sandbox/gem_manager.py +++ b/rock/sandbox/gem_manager.py @@ -20,6 +20,7 @@ from rock.sandbox.sandbox_manager import SandboxManager from rock.sandbox.sandbox_meta_store import SandboxMetaStore + class GemManager(SandboxManager): def __init__( self, diff --git a/rock/sandbox/operator/factory.py b/rock/sandbox/operator/factory.py index ff593523f6..fd3d31aaa7 100644 --- a/rock/sandbox/operator/factory.py +++ b/rock/sandbox/operator/factory.py @@ -77,4 +77,4 @@ def create_operator(context: OperatorContext) -> AbstractOperator: k8s_operator.set_nacos_provider(context.nacos_provider) return k8s_operator else: - raise ValueError(f"Unsupported operator type: {operator_type}. " f"Supported types: ray, kubernetes") + raise ValueError(f"Unsupported operator type: {operator_type}. Supported types: ray, kubernetes") diff --git a/rock/sandbox/operator/k8s/api_client.py b/rock/sandbox/operator/k8s/api_client.py index a740c65075..3832a8dc95 100644 --- a/rock/sandbox/operator/k8s/api_client.py +++ b/rock/sandbox/operator/k8s/api_client.py @@ -11,7 +11,8 @@ """ import asyncio -from typing import Any, Callable +from collections.abc import Callable +from typing import Any from aiolimiter import AsyncLimiter from kubernetes import client @@ -25,6 +26,7 @@ # User-Agent for K8s API requests USER_AGENT = "rock-k8s-client/v1.0.0" + def _make_list_func( custom_api: client.CustomObjectsApi, group: str, @@ -120,9 +122,7 @@ def __init__( self._rate_limiter = AsyncLimiter(max_rate=qps, time_period=1.0) # Create SharedInformer with custom list function - list_func = _make_list_func( - self._custom_api, group, version, plural - ) + list_func = _make_list_func(self._custom_api, group, version, plural) self._informer = SharedInformer( list_func=list_func, namespace=namespace, diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index 5d7066caea..bf4a073793 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -11,7 +11,7 @@ from rock.actions.sandbox.config import RemoteSandboxRuntimeConfig from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.config import K8sConfig, PoolConfig -from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig +from rock.deployments.config import DockerDeploymentConfig from rock.deployments.constants import Port from rock.logger import init_logger from rock.sandbox.operator.k8s.api_client import K8sApiClient diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index c22a0be788..8c320b66b2 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -35,12 +35,10 @@ from rock.sandbox.sandbox_actor import SandboxActor from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService +from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError from rock.utils.crypto_utils import AESEncryption -from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.utils.format import convert_to_gb, parse_size_to_bytes -from rock.utils.providers.redis_provider import RedisProvider -from rock.utils.service import build_sandbox_from_redis from rock.utils.system import get_iso8601_timestamp logger = init_logger(__name__) diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index 82091f5ca7..ea1ed3b30a 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -8,9 +8,7 @@ from __future__ import annotations from collections.abc import AsyncIterator -from typing import Any - -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 9bd719d218..b471701a96 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -161,9 +161,7 @@ async def execute(self, command: Command) -> CommandResponse: return CommandResponse(**response) @monitor_sandbox_operation() - async def batch_get_sandbox_status( - self, sandbox_ids: list[str] - ) -> list[SandboxStatusResponse]: + async def batch_get_sandbox_status(self, sandbox_ids: list[str]) -> list[SandboxStatusResponse]: if sandbox_ids is None: raise BadRequestRockError(message="sandbox_ids is None") if len(sandbox_ids) > self._batch_get_status_max_count: @@ -182,9 +180,7 @@ async def batch_get_sandbox_status( return results @monitor_sandbox_operation() - async def list_sandboxes( - self, query_params: SandboxQueryParams - ) -> SandboxListResponse: + async def list_sandboxes(self, query_params: SandboxQueryParams) -> SandboxListResponse: page = int(query_params.pop("page", "1")) page_size = int(query_params.pop("page_size", "500")) if page < 1 or page_size < 1: @@ -720,9 +716,7 @@ async def _update_expire_time(self, sandbox_id): if new_timeout is not None: await self._meta_store.update_timeout(sandbox_id, new_timeout) - async def list_all_sandboxes_by_query_params( - self, query_params: SandboxQueryParams - ): + async def list_all_sandboxes_by_query_params(self, query_params: SandboxQueryParams): all_ids = [] async for sandbox_id in self._meta_store.iter_alive_sandbox_ids(): all_ids.append(sandbox_id) diff --git a/rock/sdk/job/result.py b/rock/sdk/job/result.py index 8bf1950481..65c4fb519d 100644 --- a/rock/sdk/job/result.py +++ b/rock/sdk/job/result.py @@ -13,12 +13,12 @@ from pydantic import BaseModel, Field # --------------------------------------------------------------------------- -# TrialResult base — 通用字段 +# TrialResult base — common fields # --------------------------------------------------------------------------- class ExceptionInfo(BaseModel): - """通用异常信息""" + """General exception info.""" exception_type: str = "" exception_message: str = "" @@ -27,10 +27,10 @@ class ExceptionInfo(BaseModel): class TrialResult(BaseModel): - """单次执行结果的基类 — 通用字段 + """Base class for a single execution result — common fields. - Harbor 的 TrialResult 继承此类，添加 agent_info, verifier_result 等字段。 - 子类可 override score 和 status properties。 + Harbor's TrialResult inherits this class and adds agent_info, verifier_result, etc. + Subclasses can override the score and status properties. """ task_name: str = "" diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index ebe6ffd1d5..3435daee95 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -1,6 +1,6 @@ """Trial abstract base class — three-phase interface (setup / build / collect). -Trial 对象不管理 sandbox 生命周期；生命周期由 JobExecutor 负责。 +Trial objects do not manage sandbox lifecycle; lifecycle is managed by JobExecutor. """ from __future__ import annotations @@ -18,7 +18,7 @@ class AbstractTrial(ABC): """Trial base: three-phase interface (setup/build/collect). - Trial 不管理 sandbox 生命周期 (由 JobExecutor 负责)。 + Trial does not manage sandbox lifecycle (managed by JobExecutor). """ def __init__(self, config: JobConfig): diff --git a/rock/sdk/sandbox/agent/openhands.py b/rock/sdk/sandbox/agent/openhands.py index 2a6f9de909..46e941970c 100644 --- a/rock/sdk/sandbox/agent/openhands.py +++ b/rock/sdk/sandbox/agent/openhands.py @@ -2,6 +2,7 @@ Solving software engineering(SWE) problem with [Openhands Benchmarks SDK](https://github.com/OpenHands/benchmarks.git). Implementation framework reference: `rock/sdk/sandbox/agent/swe_agent.py` """ + from __future__ import annotations import copy @@ -466,8 +467,7 @@ async def run( except Exception as e: elapsed_total = time.time() - start_time logger.error( - f"[{sandbox_id}] Operation failed: Rollout execution failed - {str(e)} " - f"(elapsed: {elapsed_total:.2f}s)", + f"[{sandbox_id}] Operation failed: Rollout execution failed - {str(e)} (elapsed: {elapsed_total:.2f}s)", exc_info=True, ) raise diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index efa2dd85f5..3037874266 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -944,7 +944,7 @@ async def close(self) -> CloseResponse: await self.stop() def __str__(self): - """返回用户友好的字符串表示，包含主要成员变量""" + """Return user-friendly string representation with key attributes.""" return ( f"Sandbox(sandbox_id={self._sandbox_id}, " f"host_name={self._host_name!r}, " @@ -954,7 +954,7 @@ def __str__(self): ) def __repr__(self): - """返回开发者友好的字符串表示，包含所有成员变量""" + """Return developer-friendly string representation with all attributes.""" return ( f"Sandbox(" f"config={self.config!r}, " diff --git a/rock/sdk/sandbox/deploy.py b/rock/sdk/sandbox/deploy.py index 2a66a96cab..bd6f358f15 100644 --- a/rock/sdk/sandbox/deploy.py +++ b/rock/sdk/sandbox/deploy.py @@ -81,7 +81,7 @@ def format(self, template: str, **kwargs: str) -> str: >>> deploy.format("cat <>/file") "cat /tmp/rock_workdir_abc123/file" - >>> deploy.format("echo $((3 << 2 >> 1))") # 不受影响 + >>> deploy.format("echo $((3 << 2 >> 1))") # unaffected "echo $((3 << 2 >> 1))" """ subs = { diff --git a/rock/utils/k8s/informer/__init__.py b/rock/utils/k8s/informer/__init__.py index 8c60110d5d..837eb1fa78 100644 --- a/rock/utils/k8s/informer/__init__.py +++ b/rock/utils/k8s/informer/__init__.py @@ -1,5 +1,5 @@ from .cache import ObjectCache, _meta_namespace_key -from .informer import SharedInformer, ADDED, MODIFIED, DELETED, BOOKMARK, ERROR +from .informer import ADDED, BOOKMARK, DELETED, ERROR, MODIFIED, SharedInformer __all__ = [ "ObjectCache", diff --git a/rock/utils/k8s/informer/cache.py b/rock/utils/k8s/informer/cache.py index efd4fd6842..52377af7ab 100644 --- a/rock/utils/k8s/informer/cache.py +++ b/rock/utils/k8s/informer/cache.py @@ -24,7 +24,7 @@ def _meta_namespace_key(obj): ns = meta.get("namespace") or "" name = meta.get("name") or "" if ns: - return "{}/{}".format(ns, name) + return f"{ns}/{name}" return name diff --git a/rock/utils/k8s/informer/informer.py b/rock/utils/k8s/informer/informer.py index fe8e01d15d..79f435a212 100644 --- a/rock/utils/k8s/informer/informer.py +++ b/rock/utils/k8s/informer/informer.py @@ -12,7 +12,7 @@ from kubernetes.client.exceptions import ApiException from kubernetes.watch import Watch -from .cache import ObjectCache, _meta_namespace_key +from .cache import ObjectCache logger = logging.getLogger(__name__) @@ -102,7 +102,8 @@ def add_event_handler(self, event_type, handler): if event_type not in self._handlers: raise ValueError( "Unknown event_type {!r}. Use one of: {}".format( - event_type, ", ".join(sorted(self._handlers)), + event_type, + ", ".join(sorted(self._handlers)), ) ) with self._handler_lock: @@ -171,9 +172,7 @@ def _fire(self, event_type, obj): try: fn(obj) except Exception: - logger.exception( - "Exception in informer handler for %s", event_type - ) + logger.exception("Exception in informer handler for %s", event_type) def _initial_list(self): """List all objects and populate the cache, firing ADDED/MODIFIED/DELETED events. @@ -294,9 +293,7 @@ def _run_loop(self): except ApiException as exc: if exc.status == 410: # The stored resource version is too old; force a full re-list. - logger.warning( - "Watch expired (410 Gone); will re-list from scratch" - ) + logger.warning("Watch expired (410 Gone); will re-list from scratch") self._resource_version = None else: logger.warning( @@ -312,11 +309,7 @@ def _run_loop(self): # (updated on every ADDED/MODIFIED/DELETED/BOOKMARK event) so # that the next watch connection can resume without re-listing. # Do not overwrite a None that was set by a 410 handler above. - if ( - self._resource_version is not None - and self._watch is not None - and self._watch.resource_version - ): + if self._resource_version is not None and self._watch is not None and self._watch.resource_version: self._resource_version = self._watch.resource_version self._watch = None diff --git a/scripts/gen_ddl.py b/scripts/gen_ddl.py index 34a93dfb7a..584bb25b21 100644 --- a/scripts/gen_ddl.py +++ b/scripts/gen_ddl.py @@ -15,9 +15,11 @@ def get_dialect(name: str): if name == "postgresql": from sqlalchemy.dialects import postgresql + return postgresql.dialect() if name == "sqlite": from sqlalchemy.dialects import sqlite + return sqlite.dialect() print(f"Unsupported dialect: {name}", file=sys.stderr) sys.exit(1) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index a54d4aa6f1..3670dc3471 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -107,8 +107,7 @@ def admin_remote_server(): parsed = urlparse(normalized) if not parsed.hostname or not parsed.port: raise ValueError( - "Invalid ROCK_TEST_ADMIN_BASE_URL. Expected host:port or http://host:port, " - f"got: {external_base_url!r}" + f"Invalid ROCK_TEST_ADMIN_BASE_URL. Expected host:port or http://host:port, got: {external_base_url!r}" ) logger.info("Using external admin server from ROCK_TEST_ADMIN_BASE_URL=%s", external_base_url) yield RemoteServer(port=parsed.port, endpoint=f"{parsed.scheme}://{parsed.hostname}") diff --git a/tests/unit/admin/core/test_sandbox_table.py b/tests/unit/admin/core/test_sandbox_table.py index 7536c4e0ea..87035ff900 100644 --- a/tests/unit/admin/core/test_sandbox_table.py +++ b/tests/unit/admin/core/test_sandbox_table.py @@ -152,10 +152,13 @@ async def test_insert_duplicate_raises(self, db): async def test_update(self, db): sandbox_id = "test-sandbox-003" - await db.create(sandbox_id, { - "state": "PENDING", - "create_time": "2025-01-01T00:00:00Z", - }) + await db.create( + sandbox_id, + { + "state": "PENDING", + "create_time": "2025-01-01T00:00:00Z", + }, + ) await db.update(sandbox_id, {"state": "RUNNING", "host_ip": "10.0.0.2"}) record = await db.get(sandbox_id) diff --git a/tests/unit/admin/core/test_schema_varchar_lengths.py b/tests/unit/admin/core/test_schema_varchar_lengths.py index bc7ffdb3f0..812e336597 100644 --- a/tests/unit/admin/core/test_schema_varchar_lengths.py +++ b/tests/unit/admin/core/test_schema_varchar_lengths.py @@ -9,7 +9,6 @@ from __future__ import annotations import pytest -from sqlalchemy.exc import DataError from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.sandbox_table import SandboxTable @@ -43,10 +42,13 @@ async def db(self, pg_container): async def test_insert_long_image(self, db): """A 196-char image string must be accepted by the ORM schema.""" sandbox_id = "varchar-img-001" - await db.create(sandbox_id, { - "image": _LONG_IMAGE, - "create_time": "2026-04-14T00:00:00Z", - }) + await db.create( + sandbox_id, + { + "image": _LONG_IMAGE, + "create_time": "2026-04-14T00:00:00Z", + }, + ) record = await db.get(sandbox_id) assert record is not None assert record["image"] == _LONG_IMAGE diff --git a/tests/unit/admin/core/test_statement_cache.py b/tests/unit/admin/core/test_statement_cache.py index ac0fcf75a5..a73b550bd6 100644 --- a/tests/unit/admin/core/test_statement_cache.py +++ b/tests/unit/admin/core/test_statement_cache.py @@ -47,19 +47,20 @@ async def test_list_by_in_after_alter_column(self, setup): # 1. populate and query — warms prepared statement cache ids = [f"cache-{i:03d}" for i in range(10)] for sid in ids: - await table.create(sid, { - "image": "python:3.11", - "create_time": "2026-04-14T00:00:00Z", - }) + await table.create( + sid, + { + "image": "python:3.11", + "create_time": "2026-04-14T00:00:00Z", + }, + ) records = await table.list_by_in("sandbox_id", ids) assert len(records) == 10 # 2. external DDL — simulates hotfix applied while app is running raw = await asyncpg.connect(pg_url) try: - await raw.execute( - "ALTER TABLE sandbox_record ALTER COLUMN image TYPE VARCHAR(1024)" - ) + await raw.execute("ALTER TABLE sandbox_record ALTER COLUMN image TYPE VARCHAR(1024)") finally: await raw.close() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index b5f53ace65..864c1f2c49 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -93,7 +93,11 @@ async def _memory_sandbox_table(): @pytest.fixture async def sandbox_manager( - rock_config: RockConfig, redis_provider: RedisProvider, ray_init_shutdown, ray_service, ray_operator, + rock_config: RockConfig, + redis_provider: RedisProvider, + ray_init_shutdown, + ray_service, + ray_operator, _memory_sandbox_table: SandboxTable, ): meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) @@ -283,12 +287,15 @@ def deployment_config(): def _docker_keep_containers() -> bool: import os + return os.getenv("ROCK_TEST_KEEP_DOCKER_CONTAINERS", "").lower() in {"1", "true", "yes", "on"} def _docker_detect_network(client) -> str | None: import socket + import docker + hostname = socket.gethostname() try: current = client.containers.get(hostname) @@ -324,6 +331,7 @@ def _docker_start_container(client, image, name, network_name, internal_port, ** def pg_container(): """Start a PostgreSQL 16 Docker container for the test session.""" import uuid + import docker client = docker.from_env() @@ -344,6 +352,7 @@ def pg_container(): try: # wait for readiness import time as _t + deadline = _t.time() + 30 while _t.time() < deadline: code, _ = container.exec_run(f"pg_isready -U {_PG_USER}") @@ -364,8 +373,11 @@ def pg_container(): host, port = _docker_resolve_host_port(container, network_name, _PG_PORT) yield { - "host": host, "port": port, - "user": _PG_USER, "password": _PG_PASSWORD, "database": _PG_DB, + "host": host, + "port": port, + "user": _PG_USER, + "password": _PG_PASSWORD, + "database": _PG_DB, "url": f"postgresql://{_PG_USER}:{_PG_PASSWORD}@{host}:{port}/{_PG_DB}", } finally: @@ -380,6 +392,7 @@ def pg_container(): def redis_container(): """Start a Redis Stack Docker container (with RedisJSON) for the test session.""" import uuid + import docker client = docker.from_env() @@ -394,6 +407,7 @@ def redis_container(): ) try: import time as _t + deadline = _t.time() + 30 while _t.time() < deadline: code, output = container.exec_run("redis-cli ping") diff --git a/tests/unit/sandbox/operator/test_k8s_api_client.py b/tests/unit/sandbox/operator/test_k8s_api_client.py index 590dedf4d6..fab51b1403 100644 --- a/tests/unit/sandbox/operator/test_k8s_api_client.py +++ b/tests/unit/sandbox/operator/test_k8s_api_client.py @@ -6,7 +6,7 @@ - CRUD operations on K8s custom resources """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest @@ -118,9 +118,7 @@ async def test_update_custom_object(self, k8s_api_client): with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_thread: mock_thread.return_value = {"updated": True} - result = await k8s_api_client.update_custom_object( - name="test-sandbox", body={"spec": {"new": "value"}} - ) + result = await k8s_api_client.update_custom_object(name="test-sandbox", body={"spec": {"new": "value"}}) assert result == {"updated": True} mock_thread.assert_awaited_once() @@ -128,7 +126,7 @@ async def test_update_custom_object(self, k8s_api_client): @pytest.mark.asyncio async def test_start_initializes_informer(self, k8s_api_client): """Test start() initializes the SharedInformer.""" - with patch.object(k8s_api_client._informer, 'start') as mock_start: + with patch.object(k8s_api_client._informer, "start") as mock_start: await k8s_api_client.start() assert k8s_api_client._initialized is True @@ -140,7 +138,7 @@ async def test_start_idempotent(self, k8s_api_client): Multiple start() calls should only initialize watch once. """ - with patch.object(k8s_api_client._informer, 'start') as mock_start: + with patch.object(k8s_api_client._informer, "start") as mock_start: await k8s_api_client.start() await k8s_api_client.start() @@ -150,8 +148,10 @@ async def test_start_idempotent(self, k8s_api_client): @pytest.mark.asyncio async def test_stop_informer(self, k8s_api_client): """Test stop() stops the SharedInformer.""" - with patch.object(k8s_api_client._informer, 'start') as mock_start, \ - patch.object(k8s_api_client._informer, 'stop') as mock_stop: + with ( + patch.object(k8s_api_client._informer, "start") as _mock_start, + patch.object(k8s_api_client._informer, "stop") as mock_stop, + ): await k8s_api_client.start() await k8s_api_client.stop() diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 5f2e9ec9b8..0bb2be04fd 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -181,7 +181,7 @@ async def test_get_sandbox_info_from_redis_no_provider(self, k8s_operator): async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, redis_provider): """Test get_status raises error when sandbox not found in Redis.""" k8s_operator.set_redis_provider(redis_provider) - + # Mock provider returns sandbox info mock_sandbox_info = { "sandbox_id": "test-sandbox", @@ -191,7 +191,7 @@ async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, "port_mapping": {}, } mock_provider.get_status = AsyncMock(return_value=SandboxInfo(**mock_sandbox_info)) - + # Sandbox not in Redis (no data stored) with pytest.raises(Exception, match="Sandbox test-sandbox not found in Redis"): await k8s_operator.get_status("test-sandbox") diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index 9925d4cf7b..d403913ba2 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -4,9 +4,9 @@ from rock.config import K8sConfig, PoolConfig from rock.deployments.config import DockerDeploymentConfig +from rock.deployments.constants import Port from rock.sandbox.operator.k8s.constants import K8sConstants from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, ResourceMatchingPoolSelector -from rock.deployments.constants import Port BASIC_TEMPLATES = { "default": { @@ -270,10 +270,10 @@ async def get_config(self): class MockK8sApiClient: """Mock K8s API client for testing.""" - + def __init__(self, custom_object: dict = None): self._custom_object = custom_object - + async def get_custom_object(self, name: str) -> dict: if self._custom_object is None: raise Exception(f"Sandbox '{name}' not found") @@ -339,16 +339,12 @@ async def test_raises_when_sandbox_being_deleted(self): """Raise exception when sandbox is being deleted.""" provider = make_provider() provider._initialized = True - + # Mock K8s API to return a resource with deletionTimestamp - provider._k8s_api = MockK8sApiClient({ - "metadata": { - "name": "test-sandbox", - "deletionTimestamp": "2024-01-01T00:00:00Z", - "annotations": {} - } - }) - + provider._k8s_api = MockK8sApiClient( + {"metadata": {"name": "test-sandbox", "deletionTimestamp": "2024-01-01T00:00:00Z", "annotations": {}}} + ) + with pytest.raises(Exception, match="is being deleted"): await provider._get_sandbox_runtime_info("test-sandbox") @@ -356,18 +352,20 @@ async def test_returns_runtime_info_when_sandbox_active(self): """Return runtime info when sandbox is active.""" provider = make_provider() provider._initialized = True - + # Mock K8s API to return a normal resource - provider._k8s_api = MockK8sApiClient({ - "metadata": { - "name": "test-sandbox", - "annotations": { - K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', - K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}' + provider._k8s_api = MockK8sApiClient( + { + "metadata": { + "name": "test-sandbox", + "annotations": { + K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', + K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}', + }, } } - }) - + ) + host_ip, port_mapping, resource_version = await provider._get_sandbox_runtime_info("test-sandbox") assert host_ip == "10.0.0.1" assert port_mapping[Port.PROXY] == 8000 @@ -381,16 +379,18 @@ async def test_returns_resource_version_when_present(self): provider._initialized = True # Mock K8s API to return a resource with resourceVersion - provider._k8s_api = MockK8sApiClient({ - "metadata": { - "name": "test-sandbox", - "resourceVersion": "12345", - "annotations": { - K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', - K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}' + provider._k8s_api = MockK8sApiClient( + { + "metadata": { + "name": "test-sandbox", + "resourceVersion": "12345", + "annotations": { + K8sConstants.ANNOTATION_ENDPOINTS: '["10.0.0.1"]', + K8sConstants.ANNOTATION_PORTS: '{"proxy": 8000, "server": 8080, "ssh": 22}', + }, } } - }) + ) host_ip, port_mapping, resource_version = await provider._get_sandbox_runtime_info("test-sandbox") assert host_ip == "10.0.0.1" diff --git a/tests/unit/sandbox/test_proxy_enhancements.py b/tests/unit/sandbox/test_proxy_enhancements.py index eb54fa18b5..b1e92bae5d 100644 --- a/tests/unit/sandbox/test_proxy_enhancements.py +++ b/tests/unit/sandbox/test_proxy_enhancements.py @@ -16,19 +16,16 @@ from rock.actions.sandbox.response import State from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.sandbox_table import SandboxTable -from rock.config import DatabaseConfig, RockConfig -from rock.sandbox.sandbox_meta_store import SandboxMetaStore -from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService -from rock.utils.providers.redis_provider import RedisProvider - -from rock.admin.proto.response import SandboxListResponse from rock.admin.entrypoints.sandbox_proxy_api import ( sandbox_proxy_router, set_sandbox_proxy_service, vnc_websocket_proxy, websocket_proxy, ) +from rock.config import DatabaseConfig, RockConfig +from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService +from rock.utils.providers.redis_provider import RedisProvider def _make_mock_websocket(query_string: str = "", headers: dict | None = None) -> MagicMock: @@ -1127,4 +1124,3 @@ async def test_unknown_id_omitted(self, _svc): result = await svc.batch_get_sandbox_status(["sb-exists", "sb-ghost"]) assert len(result) == 1 assert result[0].sandbox_id == "sb-exists" - diff --git a/tests/unit/sandbox/test_sandbox_meta_store.py b/tests/unit/sandbox/test_sandbox_meta_store.py index 5b034e98ac..2b2baf0a0b 100644 --- a/tests/unit/sandbox/test_sandbox_meta_store.py +++ b/tests/unit/sandbox/test_sandbox_meta_store.py @@ -53,7 +53,6 @@ def repo(redis, db): return SandboxMetaStore(redis_provider=redis, sandbox_table=db) - @pytest.fixture def repo_with_memory_db(redis, db_memory): return SandboxMetaStore(redis_provider=redis, sandbox_table=db_memory) @@ -267,9 +266,15 @@ async def test_iter_alive_sandbox_ids_excludes_stopped(self, repo): async def test_iter_alive_sandbox_ids_works_with_sqlite_memory(self, repo_with_memory_db): """iter_alive_sandbox_ids() should work with sqlite in-memory DB + Redis fallback.""" - await repo_with_memory_db.create("sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running", "state": State.RUNNING}) - await repo_with_memory_db.create("sbx-pending", {**SANDBOX_INFO, "sandbox_id": "sbx-pending", "state": State.PENDING}) - await repo_with_memory_db.create("sbx-stopped", {**SANDBOX_INFO, "sandbox_id": "sbx-stopped", "state": "stopped"}) + await repo_with_memory_db.create( + "sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running", "state": State.RUNNING} + ) + await repo_with_memory_db.create( + "sbx-pending", {**SANDBOX_INFO, "sandbox_id": "sbx-pending", "state": State.PENDING} + ) + await repo_with_memory_db.create( + "sbx-stopped", {**SANDBOX_INFO, "sandbox_id": "sbx-stopped", "state": "stopped"} + ) await asyncio.sleep(0.1) ids = {sid async for sid in repo_with_memory_db.iter_alive_sandbox_ids()} diff --git a/tests/unit/sdk/model/test_proxy.py b/tests/unit/sdk/model/test_proxy.py index 9d188ca149..edce5584cb 100644 --- a/tests/unit/sdk/model/test_proxy.py +++ b/tests/unit/sdk/model/test_proxy.py @@ -110,8 +110,9 @@ async def test_perform_llm_request_retry_on_whitelist(): client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" # Patch asyncio.sleep inside the retry module to avoid actual waiting - with patch(client_post_path, new_callable=AsyncMock) as mock_post, patch( - "rock.utils.retry.asyncio.sleep", return_value=None + with ( + patch(client_post_path, new_callable=AsyncMock) as mock_post, + patch("rock.utils.retry.asyncio.sleep", return_value=None), ): # 1. Setup Failed Response (429) resp_429 = MagicMock(spec=Response) @@ -163,8 +164,9 @@ async def test_perform_llm_request_network_timeout_retry(): """ client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" - with patch(client_post_path, new_callable=AsyncMock) as mock_post, patch( - "rock.utils.retry.asyncio.sleep", return_value=None + with ( + patch(client_post_path, new_callable=AsyncMock) as mock_post, + patch("rock.utils.retry.asyncio.sleep", return_value=None), ): resp_200 = MagicMock(spec=Response) resp_200.status_code = 200 @@ -315,8 +317,9 @@ async def test_perform_llm_request_respects_custom_retryable_codes(): client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" - with patch(client_post_path, new_callable=AsyncMock) as mock_post, patch( - "rock.utils.retry.asyncio.sleep", return_value=None + with ( + patch(client_post_path, new_callable=AsyncMock) as mock_post, + patch("rock.utils.retry.asyncio.sleep", return_value=None), ): # 502 should retry (in custom list) resp_502 = MagicMock(spec=Response) @@ -460,8 +463,9 @@ def test_metrics_monitor_uses_env_endpoint(): custom_endpoint = "http://my-otel-collector:4318/v1/metrics" - with patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, patch.dict( - "os.environ", {"ROCK_METRICS_ENDPOINT": custom_endpoint} + with ( + patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, + patch.dict("os.environ", {"ROCK_METRICS_ENDPOINT": custom_endpoint}), ): mock_monitor = MagicMock() mock_cls.create.return_value = mock_monitor @@ -508,8 +512,9 @@ async def test_record_traj_reports_rt_and_count(): mock_monitor = MagicMock() - with patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, patch.dict( - "os.environ", {"ROCK_SANDBOX_ID": "sandbox-test-001"} + with ( + patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, + patch.dict("os.environ", {"ROCK_SANDBOX_ID": "sandbox-test-001"}), ): mock_cls.create.return_value = mock_monitor utils_module._metrics_monitor = None From f906be08f14e5455235b7fd32d3fc35aa8255668 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 15 Apr 2026 18:26:37 +0800 Subject: [PATCH 040/226] feat(job): auto-detect job type from YAML via strict model validation (#814) * feat(job): auto-detect job type from YAML via strict model validation - JobConfig.from_yaml() now auto-detects BashJob vs HarborJob by attempting model_validate() in order (Harbor first, Bash second); extra="forbid" on both models ensures each rejects the other's fields - HarborJobConfig.experiment_id promoted to required str (min_length=1), replacing the manual model_validator empty-check - BashJobConfig gains model_config = ConfigDict(extra="forbid") refs #813 Co-Authored-By: Claude Sonnet 4.6 * chore(examples): use JobConfig.from_yaml for auto-detection refs #813 Co-Authored-By: Claude Sonnet 4.6 * refactor(job): remove redundant HarborJobConfig.from_yaml override Base class JobConfig.from_yaml already handles subclass calls via `cls is not JobConfig` check, making the override a no-op. refs #813 Co-Authored-By: Claude Sonnet 4.6 * chore(examples): simplify JobConfig import path refs #813 Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- .../evaluation/claw_eval/run_claw_eval.py | 5 +- examples/harbor/harbor_demo.py | 5 +- rock/sdk/bench/models/job/config.py | 18 ++-- rock/sdk/job/config.py | 48 +++++++++- .../sdk/agent/test_jobconfig_experiment_id.py | 2 +- tests/unit/sdk/job/test_config.py | 96 +++++++++++++++++++ 6 files changed, 153 insertions(+), 21 deletions(-) diff --git a/examples/evaluation/claw_eval/run_claw_eval.py b/examples/evaluation/claw_eval/run_claw_eval.py index b913077f7d..f6beaefec7 100644 --- a/examples/evaluation/claw_eval/run_claw_eval.py +++ b/examples/evaluation/claw_eval/run_claw_eval.py @@ -14,13 +14,12 @@ import sys from pathlib import Path -from rock.sdk.job import Job -from rock.sdk.job.config import BashJobConfig +from rock.sdk.job import Job, JobConfig async def main() -> None: config_path = sys.argv[1] if len(sys.argv) > 1 else "claw_eval_bashjob.yaml" - config = BashJobConfig.from_yaml(config_path) + config = JobConfig.from_yaml(config_path) result = await Job(config).run() print(f"Job completed: status={result.status}, trials={len(result.trial_results)}") diff --git a/examples/harbor/harbor_demo.py b/examples/harbor/harbor_demo.py index 9984e09f68..09f48b3198 100644 --- a/examples/harbor/harbor_demo.py +++ b/examples/harbor/harbor_demo.py @@ -29,8 +29,7 @@ import os import sys -from rock.sdk.bench import HarborJobConfig -from rock.sdk.job import Job +from rock.sdk.job import Job, JobConfig _REQUIRED_ENV_VARS = [ "OSS_ACCESS_KEY_ID", @@ -66,7 +65,7 @@ def parse_args() -> argparse.Namespace: async def async_main(args: argparse.Namespace) -> None: - config = HarborJobConfig.from_yaml(args.config) + config = JobConfig.from_yaml(args.config) # Override task_names if specified via CLI if args.task and config.datasets: diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index 75ed70f1ac..c641a7af9b 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from rock.sdk.bench.constants import USER_DEFINED_LOGS from rock.sdk.bench.models.metric.config import MetricConfig @@ -174,6 +174,11 @@ class HarborJobConfig(_BaseJobConfig): and passed to ``harbor jobs start -c``. """ + model_config = ConfigDict(extra="forbid") + + # ── experiment_id is required for HarborJob (overrides nullable base field) ── + experiment_id: str = Field(min_length=1) + # ── Override environment to use RockEnvironmentConfig (adds harbor env fields) ── environment: RockEnvironmentConfig = Field(default_factory=RockEnvironmentConfig) @@ -197,8 +202,6 @@ class HarborJobConfig(_BaseJobConfig): @model_validator(mode="after") def _sync_experiment_id(self): """Sync experiment_id: JobConfig -> environment -> oss_mirror.""" - if not self.experiment_id: - raise ValueError("experiment_id must not be empty") env_exp = self.environment.experiment_id if env_exp is not None and env_exp != self.experiment_id: raise ValueError( @@ -291,15 +294,6 @@ def to_harbor_yaml(self) -> str: data["environment"] = harbor_env return yaml.dump(data, default_flow_style=False, allow_unicode=True) - @classmethod - def from_yaml(cls, path: str) -> HarborJobConfig: - """Load HarborJobConfig from a Harbor YAML config file.""" - import yaml - - with open(path) as f: - data = yaml.safe_load(f) - return cls(**data) - def enable_oss_mirror( self, *, diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index 04ac8acc20..c9e81091f8 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -10,7 +10,7 @@ from __future__ import annotations import yaml -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, ValidationError from rock.sdk.envhub import EnvironmentConfig @@ -27,13 +27,57 @@ class JobConfig(BaseModel): @classmethod def from_yaml(cls, path: str) -> JobConfig: + """Load a job config from YAML. + + When called on the base class (``JobConfig.from_yaml``), the job type is + auto-detected by trying each concrete subclass in order: + + 1. ``HarborJobConfig`` — tried first (requires ``experiment_id``) + 2. ``BashJobConfig`` — tried second (all fields optional) + + Both models use ``extra="forbid"``, so any field that belongs to the + other type causes a ``ValidationError`` and the attempt is skipped. + If both fail the combined ``ValidationError`` details are surfaced. + + When called directly on a subclass (e.g. ``BashJobConfig.from_yaml``), + no auto-detection is performed. + """ with open(path) as f: data = yaml.safe_load(f) - return cls(**data) + + if cls is not JobConfig: + # Called as BashJobConfig.from_yaml() or HarborJobConfig.from_yaml() — + # respect the explicit class, skip auto-detection. + return cls(**data) + + # Lazy import to avoid circular dependency: + # rock.sdk.bench.models.job.config → rock.sdk.job.config + from rock.sdk.bench.models.job.config import HarborJobConfig + + harbor_error: ValidationError | None = None + bash_error: ValidationError | None = None + + try: + return HarborJobConfig.model_validate(data) + except (ValidationError, ValueError) as exc: + harbor_error = exc + + try: + return BashJobConfig.model_validate(data) + except (ValidationError, ValueError) as exc: + bash_error = exc + + raise ValueError( + "YAML does not match any known job type.\n" + f" As HarborJobConfig: {harbor_error}\n" + f" As BashJobConfig: {bash_error}" + ) class BashJobConfig(JobConfig): """Config for a simple bash script job.""" + model_config = ConfigDict(extra="forbid") + script: str | None = None script_path: str | None = None diff --git a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py index bc09046646..a4172c6d2a 100644 --- a/tests/unit/sdk/agent/test_jobconfig_experiment_id.py +++ b/tests/unit/sdk/agent/test_jobconfig_experiment_id.py @@ -18,7 +18,7 @@ def test_none_experiment_id_raises(self): def test_empty_string_experiment_id_raises(self): """experiment_id='' must raise ValidationError.""" - with pytest.raises(ValidationError, match="experiment_id must not be empty"): + with pytest.raises(ValidationError, match="experiment_id"): HarborJobConfig(job_name="test", experiment_id="") diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index b65fa70e41..bf755e1e00 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -525,3 +525,99 @@ def test_multiplier_only_applied_to_fallback(self): from rock.sdk.bench.constants import DEFAULT_WAIT_TIMEOUT assert cfg.timeout == int(DEFAULT_WAIT_TIMEOUT * 2.0) + + +# --------------------------------------------------------------------------- +# JobConfig.from_yaml — auto-detection +# --------------------------------------------------------------------------- + + +class TestJobConfigFromYamlAutoDetect: + """JobConfig.from_yaml dispatches to the correct subclass based on YAML content.""" + + def test_auto_detect_bash_by_script(self, tmp_path): + yaml_content = "script: echo hello\ntimeout: 60\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, BashJobConfig) + assert cfg.script == "echo hello" + + def test_auto_detect_bash_by_script_path(self, tmp_path): + yaml_content = "script_path: run.sh\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, BashJobConfig) + assert cfg.script_path == "run.sh" + + def test_auto_detect_harbor_by_agents(self, tmp_path): + yaml_content = "experiment_id: exp-1\nagents:\n - name: my-agent\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, HarborJobConfig) + assert cfg.experiment_id == "exp-1" + + def test_auto_detect_harbor_by_datasets(self, tmp_path): + yaml_content = "experiment_id: exp-2\n" "datasets:\n" " - name: my-ds\n" " path: /tmp/ds\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, HarborJobConfig) + + def test_auto_detect_harbor_by_n_attempts(self, tmp_path): + yaml_content = "experiment_id: exp-3\nn_attempts: 5\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, HarborJobConfig) + assert cfg.n_attempts == 5 + + def test_auto_detect_harbor_by_debug_flag(self, tmp_path): + yaml_content = "experiment_id: exp-4\ndebug: true\n" + p = tmp_path / "cfg.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, HarborJobConfig) + assert cfg.debug is True + + def test_raises_on_mixed_fields(self, tmp_path): + """YAML with fields from both job types fails validation against either model.""" + yaml_content = "script: echo hi\nagents:\n - name: a\nexperiment_id: exp\n" + p = tmp_path / "mixed.yaml" + p.write_text(yaml_content) + with pytest.raises(ValueError, match="does not match any known job type"): + JobConfig.from_yaml(str(p)) + + def test_base_only_yaml_falls_through_to_bash(self, tmp_path): + """YAML with only base fields (no harbor exclusive) falls through to BashJobConfig. + + HarborJobConfig requires experiment_id, so it fails; BashJobConfig has all + optional fields and succeeds. + """ + yaml_content = "job_name: my-job\ntimeout: 300\n" + p = tmp_path / "base_only.yaml" + p.write_text(yaml_content) + cfg = JobConfig.from_yaml(str(p)) + assert isinstance(cfg, BashJobConfig) + assert cfg.job_name == "my-job" + assert cfg.timeout == 300 + + def test_bash_from_yaml_direct_still_works(self, tmp_path): + """BashJobConfig.from_yaml() continues to work regardless of auto-detect.""" + yaml_content = "script: ls -la\ntimeout: 120\n" + p = tmp_path / "bash.yaml" + p.write_text(yaml_content) + cfg = BashJobConfig.from_yaml(str(p)) + assert isinstance(cfg, BashJobConfig) + assert cfg.script == "ls -la" + + def test_harbor_from_yaml_direct_still_works(self, tmp_path): + """HarborJobConfig.from_yaml() continues to work regardless of auto-detect.""" + yaml_content = "experiment_id: exp-5\nn_attempts: 2\n" + p = tmp_path / "harbor.yaml" + p.write_text(yaml_content) + cfg = HarborJobConfig.from_yaml(str(p)) + assert isinstance(cfg, HarborJobConfig) + assert cfg.n_attempts == 2 From 82535e88fbc7e4871b608b287e0ae0926f63e42b Mon Sep 17 00:00:00 2001 From: dengwx Date: Wed, 15 Apr 2026 22:39:39 +0800 Subject: [PATCH 041/226] refactor(cli): rework `rock job run` for dual-mode input with strict validation (#818) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test(cli): scaffold rock job run parser smoke test Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): add _fail() helper for consistent job-run errors Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): stash run sub-parser on JobCommand for parser.error reuse Co-Authored-By: Claude Opus 4.6 (1M context) * feat(cli): friendly error when rock job run has no input Co-Authored-By: Claude Opus 4.6 (1M context) * feat(cli): enforce --config vs --script mutual exclusion Co-Authored-By: Claude Opus 4.6 (1M context) * feat(cli): route --script/--script-content mutex through _fail Co-Authored-By: Claude Opus 4.6 (1M context) * feat(cli): reject --type harbor without --config; default --type to None Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): extract _config_from_flags helper for bash CLI-mode Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): add _config_from_yaml helper with type-consistency check Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): add _apply_overrides helper (shared bash/harbor) Co-Authored-By: Claude Opus 4.6 (1M context) * feat(cli): clarify rock job run description; --timeout default None Changes --timeout default so YAML timeout survives unless overridden, and adds a rich RawDescriptionHelpFormatter description that spells out the two mutually-exclusive modes. Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(cli): unify rock job run flow through mode-agnostic pipeline Replaces the old bash/harbor branching in _job_run with a unified validate → build (YAML or flags) → apply-overrides → run pipeline. --config is now a first-class input for both bash and harbor jobs, with the job type auto-detected by JobConfig.from_yaml. Co-Authored-By: Claude Opus 4.6 (1M context) * test(cli): assert rock job run --help lists both modes Co-Authored-By: Claude Opus 4.6 (1M context) * fix(cli): rename sub-parser flag to --job_config to avoid top-level collision The top-level `rock --config ` flag shared dest="config" with the job sub-parser's --config, so `rock job run --config job.yaml` routed the YAML path into the CLI's INI loader, which failed with "File contains no section headers". Renaming the sub-parser's flag to --job_config (with --job-config alias) separates the two namespaces. Updates the design spec at docs/dev/cli/README.md to match the new flag name, and adds a regression test asserting that top-level --config and sub-parser --job_config stay independent. Co-Authored-By: Claude Opus 4.6 (1M context) * test(cli): remove superseded test_cli_job.py; cover arun dispatch in new suite tests/unit/sdk/job/test_cli_job.py asserted the pre-refactor behavior (logger.error on validation failure, old --config dest, --type default of "bash"). Its coverage is fully subsumed by the new tests/unit/cli/command/test_job.py, which tests the same paths against the new parser.error() + --job_config interface. The one non-overlapping case (unknown job_command → logger.error) is moved into TestArun in the new file. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(cli): skip base_url/cluster backfill for job command so YAML wins load_config_from_file() unconditionally backfilled args.base_url from ~/.rock/config.ini (default http://localhost:8080). That overwrote values coming from the job YAML because _apply_overrides() sees the non-None args.base_url and cannot distinguish a user flag from an INI backfill. Result: user-specified base_url in job YAML was silently ignored — sandbox requests went to localhost:8080 instead of the YAML-provided endpoint. For the `job` command the YAML is the source of truth; users who want INI defaults can still pass --base-url explicitly. Other commands (sandbox/admin/etc.) keep the backfill behavior unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: dengwx Co-authored-by: Claude Opus 4.6 (1M context) --- docs/dev/cli/README.md | 494 +++++++++++++++++++++ rock/cli/command/job.py | 271 +++++++++--- rock/cli/main.py | 19 +- tests/unit/cli/__init__.py | 0 tests/unit/cli/command/__init__.py | 0 tests/unit/cli/command/test_job.py | 674 +++++++++++++++++++++++++++++ tests/unit/sdk/job/test_cli_job.py | 395 ----------------- 7 files changed, 1388 insertions(+), 465 deletions(-) create mode 100644 docs/dev/cli/README.md create mode 100644 tests/unit/cli/__init__.py create mode 100644 tests/unit/cli/command/__init__.py create mode 100644 tests/unit/cli/command/test_job.py delete mode 100644 tests/unit/sdk/job/test_cli_job.py diff --git a/docs/dev/cli/README.md b/docs/dev/cli/README.md new file mode 100644 index 0000000000..85a80afc04 --- /dev/null +++ b/docs/dev/cli/README.md @@ -0,0 +1,494 @@ +# `rock job run` CLI 参数设计方案 + +## 1. 背景 + +当前 `rock/cli/command/job.py` 对 bash 和 harbor 两种 job 类型使用不同的输入方式： + +| job_type | 必填参数 | 其他参数 | +|----------|----------|----------| +| `bash` | `--script` 或 `--script-content`（二选一） | `--image / --memory / --cpus / --env / --local-path / --timeout / ...` | +| `harbor` | `--job_config`（YAML 文件） | `--image`（覆盖 YAML 中的字段） | + +存在的问题： + +1. **Bash 不支持 YAML 配置** — 真实业务里 bash 也会携带大量 env、setup_commands、uploads、image 等字段，挤在一行 CLI flag 里既难写又难复用。 +2. **`--type` 与入参强耦合** — `--type bash` 强制走 CLI flags 路径，`--type harbor` 强制走 YAML 路径，入参形态随类型跳变。实际上，"从 YAML 载入 config" 和 "类型"是正交的两个维度。 +3. **错误提示不友好** — 现有实现只打印一行 `logger.error`，不显示 help / 示例 / 合法组合，用户无法自助排查。 +4. **参数互斥关系隐式** — `--script` vs `--script-content`、`--job_config` vs `--script*`、`--type` vs YAML 里的类型等互斥/冲突关系散落在校验代码里，没有集中说明。 + +已具备的能力（参考 `rock/sdk/job/config.py`）： + +- `JobConfig.from_yaml(path)` 已支持类型自动识别 —— 通过 `HarborJobConfig → BashJobConfig` 顺序的严格模型校验（`extra="forbid"`）定位子类。 +- `BashJobConfig` / `HarborJobConfig` 都是 `from_yaml` 的可用入口。 + +所以 CLI 层只需要把"YAML 路径"升格为一等公民，即可统一两种类型的入参形态。 + +--- + +## 2. 设计目标 + +1. **YAML 优先** — `--job_config` 是推荐路径，适用于所有 job 类型；CLI flags 是简化入口，仅适用于 bash。 +2. **单一输入模式** — `--job_config` 模式与 flags 模式互斥，不允许混用以避免配置优先级歧义。 +3. **类型自动识别** — 使用 `--job_config` 时不需要 `--type`；`--type` 仅作为 flags 模式下的显式声明（默认 `bash`）。 +4. **错误提示优雅** — 任何参数校验失败都： + - 清晰指出错了什么； + - 给出合法组合示例； + - 最后附上 `rock job run --help` 提示。 +5. **argparse 原生风格** — 使用 `parser.error()`，退出码为 2，日志/错误都写 stderr，与其他 CLI 子命令一致。 + +--- + +## 3. 入参规则 + +### 3.1 两种互斥的输入模式 + +``` +┌──────────────────────────────────────────────┬───────────────────────────────────────────────┐ +│ Mode A: YAML config（推荐，通用） │ Mode B: CLI flags（简化，仅 bash） │ +├──────────────────────────────────────────────┼───────────────────────────────────────────────┤ +│ rock job run --job_config job.yaml │ rock job run [--type bash] │ +│ [--image ...] [--cluster ...] │ --script path/to/script.sh │ +│ [--base-url ...] [--timeout N] │ (或 --script-content "...") │ +│ │ [--image / --memory / --cpus] │ +│ │ [--env / --local-path / ...] │ +└──────────────────────────────────────────────┴───────────────────────────────────────────────┘ +``` + +### 3.2 参数分类 + +| 组别 | 参数 | Mode A 允许 | Mode B 允许 | +|------|------|:-----------:|:-----------:| +| **模式开关** | `--job_config PATH` | ✅（必填） | ❌ | +| **类型声明** | `--type {bash,harbor}` | ❌（类型由 YAML 决定） | ✅（可选，默认 `bash`；若指定 `harbor` 必须配 `--job_config`） | +| **bash 脚本** | `--script / --script-content` | ❌ | ✅（二选一，必填） | +| **环境覆盖** | `--image / --memory / --cpus` | ✅（覆盖 YAML 中对应字段） | ✅ | +| **运行位置** | `--base-url / --cluster / --extra-headers / --xrl-authorization` | ✅ | ✅ | +| **环境变量** | `--env KEY=VALUE`（可重复） | ✅（追加/覆盖 YAML 的 `environment.env`） | ✅ | +| **文件上传** | `--local-path / --target-path` | ✅（追加到 YAML 的 `environment.uploads`） | ✅ | +| **超时** | `--timeout N` | ✅（覆盖 YAML 的 `timeout`） | ✅ | + +### 3.3 校验矩阵 + +| 情况 | 判定 | 行为 | +|------|------|------| +| 两种 mode 都没给（既没 `--job_config` 也没 `--script*`） | 错误 | 打印缺参提示 + usage + help 指引 | +| 同时给 `--job_config` 和 `--script` / `--script-content` | 错误 | 明确告诉用户两种模式互斥，二选一 | +| 同时给 `--script` 和 `--script-content` | 错误 | 告知"inline 内容和文件路径只能二选一" | +| `--type harbor` 但没 `--job_config` | 错误 | 告知 harbor 类型必须走 YAML | +| `--type bash` + `--job_config` 指向 HarborJobConfig YAML | 错误 | 告知类型冲突（YAML 自动识别为 harbor） | +| `--job_config` 文件不存在 / 解析失败 | 错误 | 转发 `from_yaml` 的 `ValueError` 细节 | + +--- + +## 4. 错误提示设计 + +所有入参校验失败都走一个统一的 `_fail(parser, msg, hint=None)` 辅助函数，保持风格一致。 + +### 4.1 辅助函数示意 + +```python +def _fail(parser: argparse.ArgumentParser, msg: str, *, hint: str | None = None) -> None: + """Emit a consistent CLI error: message + optional hint + help指引，然后退出 2。""" + lines = [msg] + if hint: + lines.append("") + lines.append(hint) + lines.append("") + lines.append("Run `rock job run --help` for full usage.") + parser.error("\n".join(lines)) +``` + +`parser.error()` 会：先打印 usage（来自 argparse），再把我们给的消息写到 stderr，最后以 exit code 2 退出 —— 这正是 argparse 原生校验失败时的行为，一致性好。 + +### 4.2 错误样例（对用户实际看到的输出） + +**Case 1: 什么都没给** + +``` +usage: rock job run [-h] [--type {bash,harbor}] [--job_config CONFIG] [--script SCRIPT] ... +rock job run: error: Missing job definition. Provide either a YAML config or inline script. + +Examples: + rock job run --job_config job.yaml # any job type, auto-detected + rock job run --script path/to/run.sh # bash, script file + rock job run --script-content "echo hi" # bash, inline snippet + +Run `rock job run --help` for full usage. +``` + +**Case 2: 同时给了 `--job_config` 和 `--script`** + +``` +rock job run: error: --job_config is mutually exclusive with --script / --script-content. + +Pick one mode: + • YAML mode: rock job run --job_config job.yaml + • flags mode: rock job run --script run.sh + +Run `rock job run --help` for full usage. +``` + +**Case 3: `--type harbor` 没带 `--job_config`** + +``` +rock job run: error: --type harbor requires --job_config . + +Harbor jobs cannot be expressed purely via CLI flags. +Example: + rock job run --job_config harbor.yaml + +Run `rock job run --help` for full usage. +``` + +**Case 4: `--script` 和 `--script-content` 同时给** + +``` +rock job run: error: --script and --script-content are mutually exclusive (pick a file path OR an inline snippet). + +Run `rock job run --help` for full usage. +``` + +**Case 5: YAML 解析失败（来自 `JobConfig.from_yaml`）** + +``` +rock job run: error: Failed to load --job_config 'job.yaml': +YAML does not match any known job type. + As HarborJobConfig: 1 validation error for HarborJobConfig + experiment_id: Field required [type=missing, ...] + As BashJobConfig: 1 validation error for BashJobConfig + agents: Extra inputs are not permitted [type=extra_forbidden, ...] + +Run `rock job run --help` for full usage. +``` + +--- + +## 5. 实现草图 + +```python +# rock/cli/command/job.py + +import argparse +from pathlib import Path + +from rock.cli.command.command import Command +from rock.logger import init_logger + +logger = init_logger(__name__) + + +def _fail(parser: argparse.ArgumentParser, msg: str, *, hint: str | None = None) -> None: + """Emit a consistent CLI error + help pointer, then exit 2.""" + parts = [msg] + if hint: + parts.extend(["", hint]) + parts.extend(["", "Run `rock job run --help` for full usage."]) + parser.error("\n".join(parts)) + + +class JobCommand(Command): + name = "job" + + # The run sub-parser is stashed here so arun() can call parser.error(). + _run_parser: argparse.ArgumentParser | None = None + + async def arun(self, args: argparse.Namespace): + if args.job_command == "run": + await self._job_run(args) + else: + logger.error(f"Unknown job subcommand: {args.job_command}") + + async def _job_run(self, args: argparse.Namespace): + from rock.sdk.job import Job + from rock.sdk.job.config import BashJobConfig, JobConfig + + parser = self._run_parser # set in add_parser_to + + # ── 1. 模式互斥校验 ────────────────────────────────────────── + has_config = bool(args.config) + has_script = bool(args.script or args.script_content) + + if not has_config and not has_script: + _fail( + parser, + "Missing job definition. Provide either a YAML config or inline script.", + hint=( + "Examples:\n" + " rock job run --job_config job.yaml # any job type, auto-detected\n" + " rock job run --script path/to/run.sh # bash, script file\n" + " rock job run --script-content \"echo hi\" # bash, inline snippet" + ), + ) + + if has_config and has_script: + _fail( + parser, + "--job_config is mutually exclusive with --script / --script-content.", + hint=( + "Pick one mode:\n" + " • YAML mode: rock job run --job_config job.yaml\n" + " • flags mode: rock job run --script run.sh" + ), + ) + + if args.script and args.script_content: + _fail( + parser, + "--script and --script-content are mutually exclusive " + "(pick a file path OR an inline snippet).", + ) + + if args.type == "harbor" and not has_config: + _fail( + parser, + "--type harbor requires --job_config .", + hint=( + "Harbor jobs cannot be expressed purely via CLI flags.\n" + "Example:\n" + " rock job run --job_config harbor.yaml" + ), + ) + + # ── 2. 组装 config ────────────────────────────────────────── + if has_config: + config = self._config_from_yaml(parser, args) + else: + config = self._config_from_flags(args) + + # ── 3. 用 CLI overrides 打补丁（两种模式共用）──────────────── + self._apply_overrides(config, args) + + # ── 4. 运行 ──────────────────────────────────────────────── + try: + result = await Job(config).run() + if result.trial_results: + for tr in result.trial_results: + output = getattr(tr, "raw_output", None) or "" + if output: + print(output) + logger.info(f"Job completed: status={result.status}") + except Exception as e: + logger.error(f"Job failed: {e}") + + # ── helpers ─────────────────────────────────────────────────── + + def _config_from_yaml(self, parser, args): + from rock.sdk.job.config import BashJobConfig, JobConfig + + path = args.config + if not Path(path).is_file(): + _fail(parser, f"--job_config path does not exist: {path}") + + try: + config = JobConfig.from_yaml(path) # 自动识别 Bash/Harbor + except (ValueError, Exception) as exc: + _fail(parser, f"Failed to load --job_config {path!r}:\n{exc}") + + # --type 显式声明时做一次一致性检查 + if args.type: + expected_type = args.type + actual_type = "bash" if isinstance(config, BashJobConfig) else "harbor" + if expected_type != actual_type: + _fail( + parser, + f"--type {expected_type} does not match YAML (detected as {actual_type}).", + hint="Remove --type and let the YAML decide, or pass a matching config file.", + ) + return config + + def _config_from_flags(self, args): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + env = {} + for item in args.env or []: + key, _, value = item.partition("=") + env[key] = value + + uploads = [(args.local_path, args.target_path)] if args.local_path else [] + + return BashJobConfig( + script=args.script_content, + script_path=args.script, + environment=RockEnvironmentConfig( + image=args.image, + memory=args.memory, + cpus=args.cpus, + base_url=args.base_url, + cluster=args.cluster, + extra_headers=args.extra_headers, + xrl_authorization=args.xrl_authorization, + uploads=uploads, + auto_stop=True, + env=env, + ), + timeout=args.timeout, + ) + + def _apply_overrides(self, config, args): + """Apply CLI overrides that are valid in both modes (e.g. --image).""" + env = config.environment + if args.image: + env.image = args.image + if args.memory: + env.memory = args.memory + if args.cpus: + env.cpus = args.cpus + if args.base_url: + env.base_url = args.base_url + if args.cluster: + env.cluster = args.cluster + if args.xrl_authorization: + env.xrl_authorization = args.xrl_authorization + # --env / --local-path / --timeout 在 YAML 模式下也追加/覆盖 + for item in args.env or []: + key, _, value = item.partition("=") + env.env[key] = value + if args.local_path: + env.uploads = list(env.uploads) + [(args.local_path, args.target_path)] + if args.timeout is not None: + config.timeout = args.timeout + env.auto_stop = True + + # ── parser ──────────────────────────────────────────────────── + + @staticmethod + async def add_parser_to(subparsers: argparse._SubParsersAction): + job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") + job_subparsers = job_parser.add_subparsers(dest="job_command") + + run_parser = job_subparsers.add_parser( + "run", + help="Run a job in a sandbox", + description=( + "Run a sandbox job in one of two modes:\n" + " (1) YAML mode : --job_config (type auto-detected)\n" + " (2) flags mode : --script / --script-content (bash only)\n" + "The two modes are mutually exclusive." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # mode switches + run_parser.add_argument( + "--type", + choices=["bash", "harbor"], + default=None, + help="Explicit job type (flags mode only; YAML mode auto-detects).", + ) + run_parser.add_argument("--job_config", default=None, help="YAML config path (any job type).") + run_parser.add_argument("--script", default=None, help="Bash script file path (flags mode).") + run_parser.add_argument("--script-content", default=None, help="Inline bash snippet (flags mode).") + + # shared overrides + run_parser.add_argument("--image", default=None, help="Sandbox image (overrides YAML).") + run_parser.add_argument("--memory", default=None, help="Memory (e.g. 8g). Overrides YAML.") + run_parser.add_argument("--cpus", default=None, type=float, help="CPU count. Overrides YAML.") + run_parser.add_argument("--timeout", type=int, default=None, help="Timeout in seconds.") + run_parser.add_argument("--local-path", default=None, help="Local dir to upload.") + run_parser.add_argument("--target-path", default="/root/job", help="Target dir in sandbox.") + run_parser.add_argument("--base-url", default=None, help="Admin service base URL.") + run_parser.add_argument("--cluster", default=None, help="Cluster name (e.g. vpc-sg-sl-a).") + run_parser.add_argument("--extra-headers", default=None, help="Extra HTTP headers (JSON).") + run_parser.add_argument( + "--env", + action="append", + default=None, + metavar="KEY=VALUE", + help="Environment variable, repeatable (e.g. --env FOO=bar --env BAZ=qux).", + ) + run_parser.add_argument("--xrl-authorization", default=None, help="XRL authorization token.") + + # stash on command class so arun() can call parser.error() consistently + JobCommand._run_parser = run_parser +``` + +> 注：`parser.error()` 在 argparse 内部 `sys.exit(2)`，因此 `arun()` 里不需要 return；后续代码永远不会执行到。 + +--- + +## 6. 使用示例 + +### 6.1 Bash — flags 模式（原行为保持） + +```bash +rock job run --script ./train.sh \ + --image python:3.11 --memory 8g --cpus 4 \ + --env FOO=bar --env BAZ=qux \ + --local-path ./workspace --target-path /root/job +``` + +### 6.2 Bash — YAML 模式（新能力） + +`bash_job.yaml`： + +```yaml +script_path: ./train.sh +timeout: 7200 +environment: + image: python:3.11 + memory: 8g + cpus: 4 + env: + FOO: bar + BAZ: qux + uploads: + - ["./workspace", "/root/job"] +``` + +```bash +rock job run --job_config bash_job.yaml + +# 仍可以用 CLI override 局部字段 +rock job run --job_config bash_job.yaml --image python:3.12 --timeout 1800 +``` + +### 6.3 Harbor — YAML 模式（与之前一致） + +```bash +rock job run --job_config harbor.yaml +rock job run --job_config harbor.yaml --image harbor-runner:v2 # 覆盖 image +``` + +--- + +## 7. 向后兼容 + +| 用法 | 是否仍然可用 | 说明 | +|------|:------------:|------| +| `rock job run --type bash --script foo.sh` | ✅ | 等价于 flags 模式，`--type bash` 为显式声明。 | +| `rock job run --script foo.sh` | ✅ | `--type` 默认视为 `bash`。 | +| `rock job run --type harbor --job_config harbor.yaml` | ✅ | 显式声明 + YAML 校验一致性。 | +| `rock job run --job_config harbor.yaml` | ✅ | 推荐新用法，自动识别。 | +| `rock job run --job_config ... --script ...` | ❌（新增校验） | 之前不明确，现在明确报错。 | +| `rock job run`（无参） | ❌（错误消息更友好） | 之前报错含糊，现在打印 usage + 示例。 | + +没有被移除的 flag，没有破坏性变更。 + +--- + +## 8. 变更清单 + +- `rock/cli/command/job.py` + - 新增 `_fail()` 辅助函数 + - 拆分 `_config_from_yaml` / `_config_from_flags` / `_apply_overrides` + - `--type` 默认改为 `None`（由模式推导），仅做一致性校验 + - `--timeout` 默认改为 `None`（不覆盖 YAML 中的 timeout，除非显式给出） + - `--job_config` 支持 bash 类型 + - 所有入参错误走 `parser.error()`，统一格式 + - `description` 里写清楚两种模式（`RawDescriptionHelpFormatter`） +- `tests/unit/cli/command/test_job.py`（新增 / 扩展） + - Case: 两种模式互斥 + - Case: `--type harbor` 无 `--job_config` + - Case: YAML 模式下 `--type` 与 YAML 一致 / 冲突 + - Case: bash YAML 正常加载 + CLI override 生效 + - Case: `parser.error()` 以 exit code 2 退出、stderr 包含 usage + +--- + +## 9. 未来扩展 + +1. **`rock job validate --job_config foo.yaml`** — 只做 YAML 校验 + 类型识别，便于 CI 中 lint。 +2. **`rock job run --job_config foo.yaml --dry-run`** — 打印最终合并后的 config（含 CLI overrides）但不实际提交。 +3. **多个 `--job_config`** — 支持 config 合并（base + override），对应 `JobConfig.model_validate` 的深合并。 +4. **`ROCK_JOB_CONFIG` 环境变量** — 作为 `--job_config` 的默认值，与 `ROCK_CONFIG` 风格一致。 diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index 7142c662a8..9c1bb8c93c 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -6,9 +6,26 @@ logger = init_logger(__name__) +def _fail(parser: argparse.ArgumentParser, msg: str, *, hint: str | None = None) -> None: + """Emit a consistent CLI error: message + optional hint + help pointer, then exit 2. + + Uses ``parser.error()`` which prints the parser's usage line to stderr, writes + the message, and calls ``sys.exit(2)``. Never returns. + """ + parts = [msg] + if hint: + parts.extend(["", hint]) + parts.extend(["", "Run `rock job run --help` for full usage."]) + parser.error("\n".join(parts)) + + class JobCommand(Command): name = "job" + # Cached reference to the `run` sub-parser; populated by add_parser_to, + # used by _job_run to call parser.error() consistently. + _run_parser: argparse.ArgumentParser | None = None + async def arun(self, args: argparse.Namespace): if args.job_command == "run": await self._job_run(args) @@ -17,73 +34,64 @@ async def arun(self, args: argparse.Namespace): async def _job_run(self, args: argparse.Namespace): # Import lazily to avoid pulling in bench/Harbor modules for bash-only uses - from rock.sdk.bench.models.trial.config import RockEnvironmentConfig from rock.sdk.job import Job - from rock.sdk.job.config import BashJobConfig - job_type = args.type or "bash" - - if job_type == "bash": - if not args.script and not args.script_content: - logger.error("Either --script or --script-content is required for bash type") - return - if args.script and args.script_content: - logger.error("--script and --script-content cannot be used together") - return - - env_kwargs = {} - if args.image: - env_kwargs["image"] = args.image - if args.memory: - env_kwargs["memory"] = args.memory - if args.cpus: - env_kwargs["cpus"] = args.cpus - if getattr(args, "base_url", None): - env_kwargs["base_url"] = args.base_url - if getattr(args, "cluster", None): - env_kwargs["cluster"] = args.cluster - if getattr(args, "extra_headers", None): - env_kwargs["extra_headers"] = args.extra_headers - if getattr(args, "xrl_authorization", None): - env_kwargs["xrl_authorization"] = args.xrl_authorization - - uploads = [] - if args.local_path: - uploads.append((args.local_path, args.target_path)) - - env = {} - if getattr(args, "env", None): - for item in args.env: - key, _, value = item.partition("=") - env[key] = value - - config = BashJobConfig( - script=args.script_content, - script_path=args.script, - environment=RockEnvironmentConfig( - **env_kwargs, - uploads=uploads, - auto_stop=True, - env=env, + parser = self._run_parser + + # ── 1. Mode validation ──────────────────────────────────────── + has_config = bool(args.job_config) + has_script = bool(args.script or args.script_content) + + if not has_config and not has_script: + _fail( + parser, + "Missing job definition. Provide either a YAML config or inline script.", + hint=( + "Examples:\n" + " rock job run --job_config job.yaml # any job type, auto-detected\n" + " rock job run --script path/to/run.sh # bash, script file\n" + ' rock job run --script-content "echo hi" # bash, inline snippet' + ), + ) + + if has_config and has_script: + _fail( + parser, + "--job_config is mutually exclusive with --script / --script-content.", + hint=( + "Pick one mode:\n" + " - YAML mode: rock job run --job_config job.yaml\n" + " - flags mode: rock job run --script run.sh" ), - timeout=args.timeout, ) - elif job_type == "harbor": - if not args.config: - logger.error("--config is required for harbor type") - return - from rock.sdk.bench.models.job.config import HarborJobConfig + if args.script and args.script_content: + _fail( + parser, + "--script and --script-content are mutually exclusive (pick a file path OR an inline snippet).", + ) - config = HarborJobConfig.from_yaml(args.config) - if args.image: - config.environment.image = args.image - config.environment.auto_stop = True + if args.type == "harbor" and not has_config: + _fail( + parser, + "--type harbor requires --job_config .", + hint=( + "Harbor jobs cannot be expressed purely via CLI flags.\n" + "Example:\n" + " rock job run --job_config harbor.yaml" + ), + ) + # ── 2. Build config ─────────────────────────────────────────── + if has_config: + config = self._config_from_yaml(parser, args) else: - logger.error(f"Unknown job type: {job_type}") - return + config = self._config_from_flags(args) + + # ── 3. Apply overrides (shared across both modes) ───────────── + self._apply_overrides(config, args) + # ── 4. Run ──────────────────────────────────────────────────── try: result = await Job(config).run() if result.trial_results: @@ -95,23 +103,157 @@ async def _job_run(self, args: argparse.Namespace): except Exception as e: logger.error(f"Job failed: {e}") + def _apply_overrides(self, config, args: argparse.Namespace) -> None: + """Apply CLI overrides that are valid in both YAML and flags modes. + + Mutates ``config`` in place. Works for both BashJobConfig and HarborJobConfig + because both use ``RockEnvironmentConfig`` for ``environment``. + """ + env = config.environment + if args.image: + env.image = args.image + if args.memory: + env.memory = args.memory + if args.cpus: + env.cpus = args.cpus + if getattr(args, "base_url", None): + env.base_url = args.base_url + if getattr(args, "cluster", None): + env.cluster = args.cluster + if getattr(args, "extra_headers", None): + env.extra_headers = args.extra_headers + if getattr(args, "xrl_authorization", None): + env.xrl_authorization = args.xrl_authorization + + for item in args.env or []: + key, _, value = item.partition("=") + env.env[key] = value + + if args.local_path: + env.uploads = list(env.uploads) + [(args.local_path, args.target_path)] + + if args.timeout is not None: + config.timeout = args.timeout + + env.auto_stop = True + + def _config_from_yaml(self, parser: argparse.ArgumentParser, args: argparse.Namespace): + """Load config via JobConfig.from_yaml and enforce --type consistency.""" + from pathlib import Path + + from rock.sdk.job.config import BashJobConfig, JobConfig + + path = args.job_config + if not Path(path).is_file(): + _fail(parser, f"--job_config path does not exist: {path}") + + try: + config = JobConfig.from_yaml(path) + except ValueError as exc: + # from_yaml raises ValueError with a combined Bash/Harbor error message + _fail(parser, f"Failed to load --job_config {path!r}:\n{exc}") + except Exception as exc: # YAML parse error, IO error, etc. + _fail(parser, f"Failed to load --job_config {path!r}:\n{exc}") + + if args.type is not None: + actual_type = "bash" if isinstance(config, BashJobConfig) else "harbor" + if args.type != actual_type: + _fail( + parser, + f"--type {args.type} does not match YAML (detected as {actual_type}).", + hint="Remove --type and let the YAML decide, or pass a matching config file.", + ) + return config + + def _config_from_flags(self, args: argparse.Namespace): + """Build a BashJobConfig purely from CLI flags (mode B).""" + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + env: dict[str, str] = {} + for item in args.env or []: + key, _, value = item.partition("=") + env[key] = value + + uploads = [(args.local_path, args.target_path)] if args.local_path else [] + + env_kwargs: dict = {} + if args.image: + env_kwargs["image"] = args.image + if args.memory: + env_kwargs["memory"] = args.memory + if args.cpus: + env_kwargs["cpus"] = args.cpus + if getattr(args, "base_url", None): + env_kwargs["base_url"] = args.base_url + if getattr(args, "cluster", None): + env_kwargs["cluster"] = args.cluster + if getattr(args, "extra_headers", None): + env_kwargs["extra_headers"] = args.extra_headers + if getattr(args, "xrl_authorization", None): + env_kwargs["xrl_authorization"] = args.xrl_authorization + + cfg_kwargs: dict = {} + if args.timeout is not None: + cfg_kwargs["timeout"] = args.timeout + + return BashJobConfig( + script=args.script_content, + script_path=args.script, + environment=RockEnvironmentConfig( + **env_kwargs, + uploads=uploads, + auto_stop=True, + env=env, + ), + **cfg_kwargs, + ) + @staticmethod async def add_parser_to(subparsers: argparse._SubParsersAction): job_parser = subparsers.add_parser("job", help="Manage sandbox jobs") job_subparsers = job_parser.add_subparsers(dest="job_command") - run_parser = job_subparsers.add_parser("run", help="Run a job in a sandbox") - run_parser.add_argument("--type", choices=["bash", "harbor"], default="bash", help="Job type (default: bash)") + run_parser = job_subparsers.add_parser( + "run", + help="Run a job in a sandbox", + description=( + "Run a sandbox job in one of two mutually-exclusive modes:\n" + " (1) YAML mode : --job_config (type auto-detected)\n" + " (2) flags mode : --script / --script-content (bash only)" + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + run_parser.add_argument( + "--type", + choices=["bash", "harbor"], + default=None, + help="Explicit job type (flags mode only; YAML mode auto-detects).", + ) # bash args run_parser.add_argument("--script", default=None, help="Path to script file") run_parser.add_argument("--script-content", default=None, help="Inline script content") - # harbor args - run_parser.add_argument("--config", default=None, help="Harbor YAML config path") + # YAML config (mode A) — flag name is --job_config (distinct from the + # top-level --config that points at the CLI INI config). Also accept + # --job-config as the hyphen-form alias. + run_parser.add_argument( + "--job_config", + "--job-config", + dest="job_config", + default=None, + metavar="YAML", + help="Job YAML config path (any job type; auto-detected).", + ) # shared args run_parser.add_argument("--image", default=None, help="Sandbox image") run_parser.add_argument("--memory", default=None, help="Memory (e.g. 8g)") run_parser.add_argument("--cpus", default=None, type=float, help="CPU count") - run_parser.add_argument("--timeout", type=int, default=3600, help="Timeout in seconds") + run_parser.add_argument( + "--timeout", + type=int, + default=None, + help="Timeout in seconds (overrides YAML when given).", + ) run_parser.add_argument("--local-path", default=None, help="Local dir to upload") run_parser.add_argument("--target-path", default="/root/job", help="Target dir in sandbox") run_parser.add_argument("--base-url", default=None, help="Admin service base URL") @@ -128,3 +270,6 @@ async def add_parser_to(subparsers: argparse._SubParsersAction): default=None, help="XRL authorization token", ) + + # Stash on the class so _job_run can call parser.error() with the right parser. + JobCommand._run_parser = run_parser diff --git a/rock/cli/main.py b/rock/cli/main.py index 61ceecc9f3..ebb6225c1b 100644 --- a/rock/cli/main.py +++ b/rock/cli/main.py @@ -27,13 +27,18 @@ def load_config_from_file(args): manager = ConfigManager(config_path) cli_config = manager.get_config() - # If command line arguments are not set, use configuration file - if not args.base_url: - args.base_url = cli_config.base_url - if not args.auth_token and (authorization := cli_config.extra_headers.get("xrl-authorization")): - args.auth_token = authorization - if not args.cluster and (cluster := cli_config.extra_headers.get("cluster")): - args.cluster = cluster + # For the `job` command, the job YAML (--job_config) is the source of truth + # for base_url / cluster / auth_token. Backfilling from the CLI INI here would + # clobber YAML-specified values inside JobCommand._apply_overrides, which + # cannot distinguish user-supplied --base-url from an INI backfill. + # (Users who want INI defaults for a job can pass --base-url explicitly.) + if args.command != "job": + if not args.base_url: + args.base_url = cli_config.base_url + if not args.auth_token and (authorization := cli_config.extra_headers.get("xrl-authorization")): + args.auth_token = authorization + if not args.cluster and (cluster := cli_config.extra_headers.get("cluster")): + args.cluster = cluster # Process extra_headers, first get from configuration file extra_headers = cli_config.extra_headers.copy() diff --git a/tests/unit/cli/__init__.py b/tests/unit/cli/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/cli/command/__init__.py b/tests/unit/cli/command/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/cli/command/test_job.py b/tests/unit/cli/command/test_job.py new file mode 100644 index 0000000000..557fadfc16 --- /dev/null +++ b/tests/unit/cli/command/test_job.py @@ -0,0 +1,674 @@ +"""Unit tests for rock.cli.command.job.JobCommand. + +All tests in this file are fast: no Docker, Ray, or network. We drive the +sub-parser end-to-end with argparse so mutual-exclusion and error messages +match what users see at the terminal. +""" + +from __future__ import annotations + +import argparse +import asyncio + +import pytest + +from rock.cli.command.job import JobCommand + + +def _build_parser() -> argparse.ArgumentParser: + """Build a top-level parser with `job` subcommand wired in, same as the CLI.""" + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + return top + + +def test_parser_builds(): + """Smoke: the parser builds without error and exposes --job_config / --script.""" + parser = _build_parser() + ns = parser.parse_args(["job", "run", "--job_config", "foo.yaml"]) + assert ns.command == "job" + assert ns.job_command == "run" + assert ns.job_config == "foo.yaml" + assert ns.script is None + assert ns.script_content is None + + +def test_top_level_config_does_not_collide_with_job_config(): + """Regression: `rock --config cli.ini job run --job_config foo.yaml` keeps + both values separate — the top-level --config (INI loader) and the sub-parser's + --job_config (YAML) must not overwrite each other. + """ + top = argparse.ArgumentParser(prog="rock") + top.add_argument("--config", help="top-level CLI config (INI)") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + + ns = top.parse_args(["--config", "cli.ini", "job", "run", "--job_config", "bash.yaml"]) + assert ns.config == "cli.ini" + assert ns.job_config == "bash.yaml" + + +def test_job_config_hyphen_alias(): + """Both --job_config and --job-config should work (argparse accepts the alias).""" + parser = _build_parser() + ns = parser.parse_args(["job", "run", "--job-config", "foo.yaml"]) + assert ns.job_config == "foo.yaml" + + +class TestFailHelper: + def test_fail_exits_with_code_2_and_usage(self, capsys): + """_fail() must print usage + msg and exit code 2 (argparse convention).""" + from rock.cli.command.job import _fail + + parser = argparse.ArgumentParser(prog="rock job run") + parser.add_argument("--config") + + with pytest.raises(SystemExit) as excinfo: + _fail(parser, "boom") + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "usage:" in err + assert "boom" in err + assert "rock job run --help" in err # always appended + + def test_fail_includes_hint_when_given(self, capsys): + from rock.cli.command.job import _fail + + parser = argparse.ArgumentParser(prog="rock job run") + + with pytest.raises(SystemExit): + _fail(parser, "boom", hint="try this: X") + + err = capsys.readouterr().err + assert "boom" in err + assert "try this: X" in err + + def test_fail_no_hint_still_appends_help_pointer(self, capsys): + from rock.cli.command.job import _fail + + parser = argparse.ArgumentParser(prog="rock job run") + + with pytest.raises(SystemExit): + _fail(parser, "boom") + + err = capsys.readouterr().err + assert "rock job run --help" in err + + +class TestRunParserStash: + def test_run_parser_stashed_on_class_after_add_parser_to(self): + """After add_parser_to runs, JobCommand._run_parser must point to the 'run' sub-parser.""" + # Reset to isolate from other tests + JobCommand._run_parser = None + + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + + assert JobCommand._run_parser is not None + assert isinstance(JobCommand._run_parser, argparse.ArgumentParser) + # Sanity: it is the parser that knows about --config (stored as job_config) + actions = {a.dest for a in JobCommand._run_parser._actions} + assert "job_config" in actions + assert "script" in actions + + +class TestJobRunValidation: + @pytest.fixture(autouse=True) + def _parser(self): + """Rebuild the parser for each test so _run_parser is populated and fresh.""" + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + self.top = top + yield + + def _run(self, argv): + """Parse argv and invoke JobCommand.arun; SystemExit bubbles up.""" + ns = self.top.parse_args(argv) + cmd = JobCommand() + asyncio.run(cmd.arun(ns)) + + def test_missing_both_config_and_script_errors(self, capsys): + """With neither --config nor --script*, must exit 2 with helpful message.""" + with pytest.raises(SystemExit) as excinfo: + self._run(["job", "run"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "Missing job definition" in err + assert "--job_config job.yaml" in err # example in hint + assert "--script" in err + assert "rock job run --help" in err + + def test_config_and_script_mutually_exclusive(self, capsys, tmp_path): + """--job_config together with --script must error with mutex hint.""" + yaml_path = tmp_path / "job.yaml" + yaml_path.write_text("script_path: ./run.sh\n") + + with pytest.raises(SystemExit) as excinfo: + self._run(["job", "run", "--job_config", str(yaml_path), "--script", "run.sh"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "mutually exclusive" in err + assert "YAML mode" in err + assert "flags mode" in err + + def test_config_and_script_content_mutually_exclusive(self, capsys, tmp_path): + yaml_path = tmp_path / "job.yaml" + yaml_path.write_text("script_path: ./run.sh\n") + + with pytest.raises(SystemExit) as excinfo: + self._run(["job", "run", "--job_config", str(yaml_path), "--script-content", "echo hi"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "mutually exclusive" in err + + def test_script_and_script_content_mutually_exclusive(self, capsys): + with pytest.raises(SystemExit) as excinfo: + self._run(["job", "run", "--script", "run.sh", "--script-content", "echo hi"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "--script and --script-content are mutually exclusive" in err + + def test_type_harbor_requires_config(self, capsys): + with pytest.raises(SystemExit) as excinfo: + self._run(["job", "run", "--type", "harbor", "--script", "run.sh"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "--type harbor requires --job_config" in err + assert "cannot be expressed purely via CLI flags" in err + + def test_type_default_is_none(self): + """--type should default to None; 'bash' is implied only when mode is flags.""" + ns = self.top.parse_args(["job", "run", "--script", "run.sh"]) + assert ns.type is None + + +class TestConfigFromFlags: + def _args(self, **overrides): + """Build an argparse.Namespace matching the run sub-parser defaults, then overlay overrides.""" + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + argv = ["job", "run", "--script", "dummy.sh"] + ns = top.parse_args(argv) + for k, v in overrides.items(): + setattr(ns, k, v) + return ns + + def test_inline_script_content_with_env_and_uploads(self): + ns = self._args( + script=None, + script_content="echo hi", + image="python:3.11", + memory="8g", + cpus=4.0, + env=["FOO=bar", "BAZ=qux"], + local_path="/tmp/workspace", + target_path="/root/job", + timeout=1800, + ) + config = JobCommand()._config_from_flags(ns) + + from rock.sdk.job.config import BashJobConfig + + assert isinstance(config, BashJobConfig) + assert config.script == "echo hi" + assert config.script_path is None + assert config.timeout == 1800 + env = config.environment + assert env.image == "python:3.11" + assert env.memory == "8g" + assert env.cpus == 4.0 + assert env.env == {"FOO": "bar", "BAZ": "qux"} + assert env.uploads == [("/tmp/workspace", "/root/job")] + assert env.auto_stop is True + + def test_script_path_mode_no_env_no_uploads(self): + ns = self._args(script="run.sh", script_content=None, env=None, local_path=None) + config = JobCommand()._config_from_flags(ns) + + assert config.script_path == "run.sh" + assert config.script is None + assert config.environment.env == {} + assert config.environment.uploads == [] + + +class TestConfigFromYaml: + def _setup(self): + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + return top + + def _make_args(self, top, argv): + return top.parse_args(argv) + + def test_loads_bash_yaml_autodetected(self, tmp_path): + top = self._setup() + yaml_path = tmp_path / "bash.yaml" + yaml_path.write_text("script_path: ./run.sh\ntimeout: 1800\nenvironment:\n image: python:3.11\n") + ns = self._make_args(top, ["job", "run", "--job_config", str(yaml_path)]) + config = JobCommand()._config_from_yaml(JobCommand._run_parser, ns) + + from rock.sdk.job.config import BashJobConfig + + assert isinstance(config, BashJobConfig) + assert config.script_path == "./run.sh" + assert config.timeout == 1800 + assert config.environment.image == "python:3.11" + + def test_missing_file_errors_via_parser(self, tmp_path, capsys): + top = self._setup() + missing = tmp_path / "nope.yaml" + ns = self._make_args(top, ["job", "run", "--job_config", str(missing)]) + + with pytest.raises(SystemExit) as excinfo: + JobCommand()._config_from_yaml(JobCommand._run_parser, ns) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "--job_config path does not exist" in err + + def test_invalid_yaml_surfaces_from_yaml_error(self, tmp_path, capsys): + top = self._setup() + yaml_path = tmp_path / "weird.yaml" + yaml_path.write_text("totally_unknown_field: 1\n") + ns = self._make_args(top, ["job", "run", "--job_config", str(yaml_path)]) + + with pytest.raises(SystemExit) as excinfo: + JobCommand()._config_from_yaml(JobCommand._run_parser, ns) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "Failed to load --job_config" in err + assert "YAML does not match any known job type" in err + + def test_type_bash_matches_yaml(self, tmp_path): + top = self._setup() + yaml_path = tmp_path / "bash.yaml" + yaml_path.write_text("script_path: ./run.sh\n") + ns = self._make_args(top, ["job", "run", "--type", "bash", "--job_config", str(yaml_path)]) + config = JobCommand()._config_from_yaml(JobCommand._run_parser, ns) + + from rock.sdk.job.config import BashJobConfig + + assert isinstance(config, BashJobConfig) + + def test_type_harbor_mismatch_against_bash_yaml(self, tmp_path, capsys): + top = self._setup() + yaml_path = tmp_path / "bash.yaml" + yaml_path.write_text("script_path: ./run.sh\n") + ns = self._make_args(top, ["job", "run", "--type", "harbor", "--job_config", str(yaml_path)]) + + with pytest.raises(SystemExit) as excinfo: + JobCommand()._config_from_yaml(JobCommand._run_parser, ns) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + assert "--type harbor does not match YAML (detected as bash)" in err + + +class TestApplyOverrides: + def _setup(self): + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + return top + + def test_override_image_memory_cpus_on_bash_config(self): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + top = self._setup() + config = BashJobConfig( + script="echo hi", + environment=RockEnvironmentConfig(image="python:3.10", memory="2g", cpus=1), + timeout=7200, + ) + ns = top.parse_args( + [ + "job", + "run", + "--job_config", + "unused.yaml", + "--image", + "python:3.12", + "--memory", + "16g", + "--cpus", + "8", + "--timeout", + "900", + ] + ) + JobCommand()._apply_overrides(config, ns) + + assert config.environment.image == "python:3.12" + assert config.environment.memory == "16g" + assert config.environment.cpus == 8.0 + assert config.timeout == 900 + assert config.environment.auto_stop is True + + def test_env_overrides_append_and_overwrite(self): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + top = self._setup() + config = BashJobConfig( + script="echo hi", + environment=RockEnvironmentConfig(env={"FOO": "old", "KEEP": "1"}), + ) + ns = top.parse_args( + [ + "job", + "run", + "--job_config", + "unused.yaml", + "--env", + "FOO=new", + "--env", + "NEW=2", + ] + ) + JobCommand()._apply_overrides(config, ns) + + assert config.environment.env == {"FOO": "new", "KEEP": "1", "NEW": "2"} + + def test_uploads_appended(self): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + top = self._setup() + config = BashJobConfig( + script="echo hi", + environment=RockEnvironmentConfig(uploads=[("/a", "/b")]), + ) + ns = top.parse_args( + [ + "job", + "run", + "--job_config", + "unused.yaml", + "--local-path", + "/src", + "--target-path", + "/dst", + ] + ) + JobCommand()._apply_overrides(config, ns) + + assert config.environment.uploads == [("/a", "/b"), ("/src", "/dst")] + + def test_parser_timeout_default_is_none(self): + """--timeout should default to None so YAML timeout is not unconditionally overridden.""" + top = self._setup() + ns = top.parse_args(["job", "run", "--script", "run.sh"]) + assert ns.timeout is None + + def test_parser_description_mentions_two_modes(self): + self._setup() # populates JobCommand._run_parser + desc = JobCommand._run_parser.description or "" + assert "YAML mode" in desc + assert "flags mode" in desc + + def test_no_overrides_leaves_config_untouched_except_auto_stop(self): + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.job.config import BashJobConfig + + top = self._setup() + config = BashJobConfig( + script="echo hi", + environment=RockEnvironmentConfig(image="python:3.10", env={"X": "1"}), + timeout=1234, + ) + ns = top.parse_args(["job", "run", "--job_config", "unused.yaml"]) + # Force timeout to None so this test is robust regardless of the parser's + # current default (Task 11 flips it to None permanently). + ns.timeout = None + JobCommand()._apply_overrides(config, ns) + + assert config.environment.image == "python:3.10" + assert config.environment.env == {"X": "1"} + assert config.timeout == 1234 + assert config.environment.auto_stop is True # always set + + +class TestJobRunEndToEnd: + """End-to-end-ish tests for _job_run, mocking only Job.run.""" + + def _setup(self): + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + return top + + def test_flags_mode_builds_bash_config_and_runs_job(self, monkeypatch): + """A --script invocation should produce a BashJobConfig and call Job(config).run().""" + from unittest.mock import MagicMock + + from rock.sdk.job.config import BashJobConfig + + captured = {} + + class FakeJob: + def __init__(self, cfg): + captured["cfg"] = cfg + + async def run(self): + result = MagicMock() + result.status = "COMPLETED" + result.trial_results = [] + return result + + monkeypatch.setattr("rock.sdk.job.Job", FakeJob) + + top = self._setup() + ns = top.parse_args( + [ + "job", + "run", + "--script", + "run.sh", + "--image", + "python:3.11", + "--env", + "A=1", + ] + ) + asyncio.run(JobCommand().arun(ns)) + + cfg = captured["cfg"] + assert isinstance(cfg, BashJobConfig) + assert cfg.script_path == "run.sh" + assert cfg.environment.image == "python:3.11" + assert cfg.environment.env == {"A": "1"} + + def test_yaml_mode_with_override_image(self, monkeypatch, tmp_path): + from unittest.mock import MagicMock + + from rock.sdk.job.config import BashJobConfig + + yaml_path = tmp_path / "bash.yaml" + yaml_path.write_text("script_path: ./run.sh\nenvironment:\n image: python:3.10\n") + + captured = {} + + class FakeJob: + def __init__(self, cfg): + captured["cfg"] = cfg + + async def run(self): + r = MagicMock() + r.status = "COMPLETED" + r.trial_results = [] + return r + + monkeypatch.setattr("rock.sdk.job.Job", FakeJob) + + top = self._setup() + ns = top.parse_args( + [ + "job", + "run", + "--job_config", + str(yaml_path), + "--image", + "python:3.12", + ] + ) + asyncio.run(JobCommand().arun(ns)) + + cfg = captured["cfg"] + assert isinstance(cfg, BashJobConfig) + assert cfg.script_path == "./run.sh" + assert cfg.environment.image == "python:3.12" # override applied + + +class TestYamlSourceOfTruth: + """Regression: YAML-loaded fields (base_url, cluster) must survive when the + user doesn't pass the corresponding CLI flag, even though main.py would + otherwise backfill args.base_url from the ~/.rock/config.ini default. + """ + + def _setup(self): + JobCommand._run_parser = None + top = argparse.ArgumentParser(prog="rock") + # Mirror the top-level flags declared in rock/cli/main.py so we can + # drive load_config_from_file() as main.py does. + top.add_argument("--config") + top.add_argument("--base-url") + top.add_argument("--auth-token") + top.add_argument("--cluster") + top.add_argument("--extra-header", action="append", dest="extra_headers_list") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + return top + + def test_load_config_from_file_skips_backfill_for_job(self, tmp_path, monkeypatch): + """load_config_from_file must NOT backfill base_url/cluster/auth_token + when the command is `job` — the YAML is the source of truth. + """ + from rock.cli import main as main_mod + + top = self._setup() + ns = top.parse_args(["job", "run", "--job_config", "unused.yaml"]) + # Pretend the INI file would return http://ini.example.com as base_url + fake_cli_config = type( + "CLI", (), {"base_url": "http://ini.example.com", "extra_headers": {"cluster": "ini-cluster"}} + )() + monkeypatch.setattr( + main_mod, "ConfigManager", lambda _path: type("M", (), {"get_config": lambda self: fake_cli_config})() + ) + + main_mod.load_config_from_file(ns) + + # Backfill should have been skipped for `job` + assert ns.base_url is None + assert ns.cluster is None + assert ns.auth_token is None + + def test_load_config_from_file_backfills_for_non_job(self, tmp_path, monkeypatch): + """Non-job commands still get the backfill behavior.""" + from rock.cli import main as main_mod + + # Build a parser with a non-job subcommand + top = argparse.ArgumentParser(prog="rock") + top.add_argument("--config") + top.add_argument("--base-url") + top.add_argument("--auth-token") + top.add_argument("--cluster") + top.add_argument("--extra-header", action="append", dest="extra_headers_list") + subparsers = top.add_subparsers(dest="command") + subparsers.add_parser("sandbox") + ns = top.parse_args(["sandbox"]) + + fake_cli_config = type( + "CLI", (), {"base_url": "http://ini.example.com", "extra_headers": {"cluster": "ini-cluster"}} + )() + monkeypatch.setattr( + main_mod, "ConfigManager", lambda _path: type("M", (), {"get_config": lambda self: fake_cli_config})() + ) + + main_mod.load_config_from_file(ns) + + assert ns.base_url == "http://ini.example.com" + assert ns.cluster == "ini-cluster" + + def test_yaml_base_url_overridden_when_user_passes_flag(self, tmp_path, monkeypatch): + """When user explicitly passes --base-url, it should override YAML.""" + from unittest.mock import MagicMock + + yaml_path = tmp_path / "bash.yaml" + yaml_path.write_text("script_path: ./run.sh\nenvironment:\n base_url: http://xrl.alibaba-inc.com\n") + + captured = {} + + class FakeJob: + def __init__(self, cfg): + captured["cfg"] = cfg + + async def run(self): + r = MagicMock() + r.status = "COMPLETED" + r.trial_results = [] + return r + + monkeypatch.setattr("rock.sdk.job.Job", FakeJob) + + top = self._setup() + ns = top.parse_args( + [ + "job", + "run", + "--job_config", + str(yaml_path), + "--base-url", + "http://explicit.example.com", + ] + ) + asyncio.run(JobCommand().arun(ns)) + + cfg = captured["cfg"] + assert cfg.environment.base_url == "http://explicit.example.com" + + +class TestArun: + """Tests for JobCommand.arun dispatch (not _job_run).""" + + def test_unknown_job_command_logs_error(self): + from unittest.mock import patch + + ns = argparse.Namespace(job_command="weird") + with patch("rock.cli.command.job.logger") as mock_logger: + asyncio.run(JobCommand().arun(ns)) + mock_logger.error.assert_called() + + +class TestHelpOutput: + def test_help_output_mentions_both_modes(self, capsys): + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(JobCommand.add_parser_to(subparsers)) + + with pytest.raises(SystemExit) as excinfo: + top.parse_args(["job", "run", "--help"]) + + assert excinfo.value.code == 0 + out = capsys.readouterr().out + assert "YAML mode" in out + assert "flags mode" in out + assert "--job_config" in out + assert "--script" in out diff --git a/tests/unit/sdk/job/test_cli_job.py b/tests/unit/sdk/job/test_cli_job.py deleted file mode 100644 index b46a405f7e..0000000000 --- a/tests/unit/sdk/job/test_cli_job.py +++ /dev/null @@ -1,395 +0,0 @@ -"""Tests for rock.cli.command.job — JobCommand with --type bash/harbor routing.""" - -from __future__ import annotations - -import argparse -import tempfile -from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from rock.cli.command.job import JobCommand - - -async def _build_parser() -> argparse.ArgumentParser: - p = argparse.ArgumentParser() - sub = p.add_subparsers(dest="main_command") - await JobCommand.add_parser_to(sub) - return p - - -def _make_mock_job_result(): - mock_result = MagicMock() - mock_result.trial_results = [] - mock_result.status = "completed" - return mock_result - - -# ---------------------------------------------------------------------------- -# Parser tests -# ---------------------------------------------------------------------------- - - -async def test_parser_default_type_is_bash(): - p = await _build_parser() - args = p.parse_args(["job", "run", "--script-content", "echo hi"]) - assert args.type == "bash" - assert args.script_content == "echo hi" - - -async def test_parser_accepts_type_bash_explicit(): - p = await _build_parser() - args = p.parse_args(["job", "run", "--type", "bash", "--script-content", "echo hi"]) - assert args.type == "bash" - - -async def test_parser_accepts_type_harbor(): - p = await _build_parser() - args = p.parse_args(["job", "run", "--type", "harbor", "--config", "/tmp/c.yaml"]) - assert args.type == "harbor" - assert args.config == "/tmp/c.yaml" - - -async def test_parser_supports_all_bash_args(): - p = await _build_parser() - args = p.parse_args( - [ - "job", - "run", - "--script", - "/tmp/s.sh", - "--image", - "python:3.11", - "--memory", - "4g", - "--cpus", - "2", - "--timeout", - "600", - "--local-path", - "/tmp/local", - "--target-path", - "/root/other", - ] - ) - assert args.script == "/tmp/s.sh" - assert args.image == "python:3.11" - assert args.memory == "4g" - assert args.cpus == 2.0 - assert args.timeout == 600 - assert args.local_path == "/tmp/local" - assert args.target_path == "/root/other" - - -async def test_parser_invalid_type_rejected(): - p = await _build_parser() - with pytest.raises(SystemExit): - p.parse_args(["job", "run", "--type", "invalid"]) - - -# ---------------------------------------------------------------------------- -# arun / _job_run behavior tests -# ---------------------------------------------------------------------------- - - -def _bash_args(**overrides): - defaults = dict( - job_command="run", - type="bash", - script=None, - script_content=None, - image=None, - memory=None, - cpus=None, - local_path=None, - target_path="/root/job", - timeout=3600, - config=None, - base_url=None, - cluster=None, - extra_headers=None, - env=None, - ) - defaults.update(overrides) - return argparse.Namespace(**defaults) - - -async def test_bash_creates_bash_job_config_and_runs(): - from rock.sdk.job.config import BashJobConfig - - args = _bash_args( - script_content="echo hello", - image="python:3.11", - memory="4g", - cpus=2.0, - timeout=600, - ) - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - MockJob.assert_called_once() - config_arg = MockJob.call_args[0][0] - assert isinstance(config_arg, BashJobConfig) - assert config_arg.script == "echo hello" - assert config_arg.script_path is None - assert config_arg.environment.image == "python:3.11" - assert config_arg.environment.memory == "4g" - assert config_arg.environment.cpus == 2.0 - assert config_arg.timeout == 600 - assert config_arg.environment.auto_stop is True - mock_instance.run.assert_awaited_once() - - -async def test_bash_with_script_path(): - from rock.sdk.job.config import BashJobConfig - - args = _bash_args(script="/tmp/my_script.sh") - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert isinstance(config_arg, BashJobConfig) - assert config_arg.script_path == "/tmp/my_script.sh" - assert config_arg.script is None - - -async def test_bash_with_file_upload(): - args = _bash_args( - script_content="echo hi", - local_path="/tmp/src", - target_path="/root/target", - ) - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.uploads == [("/tmp/src", "/root/target")] - - -async def test_bash_requires_script_or_script_content(): - args = _bash_args() # neither set - - with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: - cmd = JobCommand() - await cmd.arun(args) - - MockJob.assert_not_called() - mock_logger.error.assert_called() - - -async def test_bash_rejects_both_script_and_script_content(): - args = _bash_args(script="/tmp/s.sh", script_content="echo hi") - - with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: - cmd = JobCommand() - await cmd.arun(args) - - MockJob.assert_not_called() - mock_logger.error.assert_called() - - -async def test_harbor_requires_config(): - args = _bash_args(type="harbor", config=None) - - with patch("rock.sdk.job.Job") as MockJob, patch("rock.cli.command.job.logger") as mock_logger: - cmd = JobCommand() - await cmd.arun(args) - - MockJob.assert_not_called() - mock_logger.error.assert_called() - - -async def test_harbor_loads_from_yaml(): - from rock.sdk.bench.models.job.config import HarborJobConfig - - yaml_content = """ -experiment_id: exp-123 -job_name: my-harbor-job -""" - with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: - f.write(yaml_content) - yaml_path = f.name - - try: - args = _bash_args(type="harbor", config=yaml_path) - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - MockJob.assert_called_once() - config_arg = MockJob.call_args[0][0] - assert isinstance(config_arg, HarborJobConfig) - assert config_arg.experiment_id == "exp-123" - assert config_arg.environment.auto_stop is True - finally: - Path(yaml_path).unlink(missing_ok=True) - - -async def test_harbor_image_override(): - yaml_content = """ -experiment_id: exp-abc -""" - with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: - f.write(yaml_content) - yaml_path = f.name - - try: - args = _bash_args(type="harbor", config=yaml_path, image="custom:tag") - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.image == "custom:tag" - finally: - Path(yaml_path).unlink(missing_ok=True) - - -async def test_parser_accepts_base_url_and_cluster(): - p = await _build_parser() - args = p.parse_args( - [ - "job", - "run", - "--script-content", - "echo hi", - "--base-url", - "http://example.com", - "--cluster", - "test-cluster-a", - ] - ) - assert args.base_url == "http://example.com" - assert args.cluster == "test-cluster-a" - - -async def test_bash_passes_base_url_and_cluster_to_environment(): - args = _bash_args( - script_content="echo hello", - base_url="http://example.com", - cluster="test-cluster-a", - ) - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.base_url == "http://example.com" - assert config_arg.environment.cluster == "test-cluster-a" - - -async def test_parser_accepts_env_args(): - p = await _build_parser() - args = p.parse_args( - [ - "job", - "run", - "--script-content", - "echo hi", - "--env", - "FOO=bar", - "--env", - "BAZ=qux=123", - ] - ) - assert args.env == ["FOO=bar", "BAZ=qux=123"] - - -async def test_parser_env_default_is_none(): - p = await _build_parser() - args = p.parse_args(["job", "run", "--script-content", "echo hi"]) - assert args.env is None - - -async def test_bash_passes_env_to_config(): - args = _bash_args(script_content="echo hello") - args.env = ["SERP_DEV_KEY=abc123", "RUN_CMD=claw-eval batch --parallel 4"] - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.env == { - "SERP_DEV_KEY": "abc123", - "RUN_CMD": "claw-eval batch --parallel 4", - } - - -async def test_bash_env_with_equals_in_value(): - """Values containing '=' should not be split.""" - args = _bash_args(script_content="echo hello") - args.env = ["API_KEY=sk-abc=def=="] - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.env == {"API_KEY": "sk-abc=def=="} - - -async def test_bash_no_env_defaults_to_empty(): - args = _bash_args(script_content="echo hello") - - with patch("rock.sdk.job.Job") as MockJob: - mock_instance = MagicMock() - mock_instance.run = AsyncMock(return_value=_make_mock_job_result()) - MockJob.return_value = mock_instance - - cmd = JobCommand() - await cmd.arun(args) - - config_arg = MockJob.call_args[0][0] - assert config_arg.environment.env == {} - - -async def test_unknown_job_command_logs_error(): - args = argparse.Namespace(job_command="weird") - - with patch("rock.cli.command.job.logger") as mock_logger: - cmd = JobCommand() - await cmd.arun(args) - mock_logger.error.assert_called() From 618a7fba6603c49c7a3f47bf0cc5548f99b37e44 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 16 Apr 2026 10:51:49 +0800 Subject: [PATCH 042/226] refactor: remove auto_stop parameter from EnvironmentConfig (#820) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor: remove auto_stop parameter from EnvironmentConfig The auto_stop flag was a conditional that controlled whether the sandbox should close after job completion. Since leaving sandboxes running is almost never desired, remove the parameter and always close on completion. closes #819 Co-Authored-By: Claude Sonnet 4.6 * fixup! refactor: remove auto_stop parameter from EnvironmentConfig Sandbox is not closed after job completion — lifecycle is managed by auto_clear_seconds instead. Remove the unconditional close() calls introduced in the previous commit. Co-Authored-By: Claude Sonnet 4.6 * fixup! refactor: remove auto_stop parameter from EnvironmentConfig Remove remaining auto_stop assertions in tests/unit/cli/command/test_job.py (added by upstream #818, missed during rebase). Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- rock/cli/command/job.py | 3 - rock/sdk/bench/job.py | 45 ++++++------ rock/sdk/bench/models/trial/config.py | 2 +- rock/sdk/envhub/config.py | 3 +- rock/sdk/job/executor.py | 68 +++++++++---------- tests/unit/cli/command/test_job.py | 5 +- tests/unit/sdk/agent/test_job.py | 18 +---- .../agent/test_job_config_serialization.py | 8 --- tests/unit/sdk/agent/test_models.py | 1 - .../sdk/job/test_blue_green_equivalence.py | 8 +-- tests/unit/sdk/job/test_config.py | 26 +------ tests/unit/sdk/job/test_executor.py | 18 +---- 12 files changed, 65 insertions(+), 140 deletions(-) diff --git a/rock/cli/command/job.py b/rock/cli/command/job.py index 9c1bb8c93c..d431fd507a 100644 --- a/rock/cli/command/job.py +++ b/rock/cli/command/job.py @@ -135,8 +135,6 @@ def _apply_overrides(self, config, args: argparse.Namespace) -> None: if args.timeout is not None: config.timeout = args.timeout - env.auto_stop = True - def _config_from_yaml(self, parser: argparse.ArgumentParser, args: argparse.Namespace): """Load config via JobConfig.from_yaml and enforce --type consistency.""" from pathlib import Path @@ -203,7 +201,6 @@ def _config_from_flags(self, args: argparse.Namespace): environment=RockEnvironmentConfig( **env_kwargs, uploads=uploads, - auto_stop=True, env=env, ), **cfg_kwargs, diff --git a/rock/sdk/bench/job.py b/rock/sdk/bench/job.py index f2f2284b50..029159bfa1 100644 --- a/rock/sdk/bench/job.py +++ b/rock/sdk/bench/job.py @@ -108,33 +108,28 @@ async def wait(self) -> JobResult: if self._pid is None or self._tmp_file is None: raise RuntimeError("No submitted job to wait for. Call submit() first.") - try: - success, message = await self._sandbox.wait_for_process_completion( - pid=self._pid, - session=self._session, - wait_timeout=self._get_wait_timeout(), - wait_interval=CHECK_INTERVAL, - ) - - obs = await self._sandbox.handle_nohup_output( - tmp_file=self._tmp_file, - session=self._session, - success=success, - message=message, - ignore_output=False, - response_limited_bytes_in_nohup=None, - ) + success, message = await self._sandbox.wait_for_process_completion( + pid=self._pid, + session=self._session, + wait_timeout=self._get_wait_timeout(), + wait_interval=CHECK_INTERVAL, + ) - result = await self._collect_results() - result.raw_output = obs.output if obs else "" - result.exit_code = obs.exit_code if obs else 1 - if not success: - result.status = JobStatus.FAILED - return result + obs = await self._sandbox.handle_nohup_output( + tmp_file=self._tmp_file, + session=self._session, + success=success, + message=message, + ignore_output=False, + response_limited_bytes_in_nohup=None, + ) - finally: - if self._config.environment.auto_stop and self._sandbox: - await self._sandbox.close() + result = await self._collect_results() + result.raw_output = obs.output if obs else "" + result.exit_code = obs.exit_code if obs else 1 + if not success: + result.status = JobStatus.FAILED + return result async def cancel(self): """Cancel a running job by killing the process.""" diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index 0aea4769ec..55487f4d12 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -57,7 +57,7 @@ class EnvironmentConfig(BaseModel): class RockEnvironmentConfig(_EnvConfig, EnvironmentConfig): """Unified Rock environment config. - Combines job environment fields (uploads, auto_stop, env) + Combines job environment fields (uploads, env) from JobEnvironmentConfig with harbor environment fields (force_build, override_cpus, oss_mirror, etc.) from EnvironmentConfig. Rock-specific fields are stripped when serializing to Harbor YAML diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py index 62c1ddfdc8..eb22e420b2 100644 --- a/rock/sdk/envhub/config.py +++ b/rock/sdk/envhub/config.py @@ -1,7 +1,7 @@ """General-purpose environment configuration. EnvironmentConfig extends SandboxConfig with common environment-level fields -(uploads, environment variables, auto-stop). +(uploads, environment variables). """ from __future__ import annotations @@ -19,5 +19,4 @@ class EnvironmentConfig(SandboxConfig): description="Files/dirs to upload before running: [(local_path, sandbox_path), ...]. " "Automatically detects file vs directory and uses the appropriate upload method.", ) - auto_stop: bool = False env: dict[str, str] = Field(default_factory=dict) diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py index be03cab367..1600c8d876 100644 --- a/rock/sdk/job/executor.py +++ b/rock/sdk/job/executor.py @@ -123,44 +123,40 @@ async def _do_wait(self, client: TrialClient) -> TrialResult | list[TrialResult] from rock.sdk.job.result import ExceptionInfo config = client.trial._config - try: - success, message = await client.sandbox.wait_for_process_completion( - pid=client.pid, - session=client.session, - wait_timeout=config.timeout, - wait_interval=30, - ) - obs = await client.sandbox.handle_nohup_output( - tmp_file=f"{self._job_tmp_prefix(config)}.out", - session=client.session, - success=success, - message=message, - ignore_output=False, - response_limited_bytes_in_nohup=None, + success, message = await client.sandbox.wait_for_process_completion( + pid=client.pid, + session=client.session, + wait_timeout=config.timeout, + wait_interval=30, + ) + obs = await client.sandbox.handle_nohup_output( + tmp_file=f"{self._job_tmp_prefix(config)}.out", + session=client.session, + success=success, + message=message, + ignore_output=False, + response_limited_bytes_in_nohup=None, + ) + exit_code = obs.exit_code if obs.exit_code is not None else 1 + if obs.output: + logger.info(f"Trial output (job={config.job_name}):\n{obs.output}") + result = await client.trial.collect(client.sandbox, obs.output or "", exit_code) + # G5: populate raw_output / exit_code on every TrialResult so they surface in JobResult + iter_results = result if isinstance(result, list) else [result] + for r in iter_results: + if not r.raw_output: + r.raw_output = obs.output or "" + if r.exit_code == 0 and exit_code != 0: + r.exit_code = exit_code + if not success: + fail_info = ExceptionInfo( + exception_type="ProcessTimeout", + exception_message=message or "process did not complete successfully", ) - exit_code = obs.exit_code if obs.exit_code is not None else 1 - if obs.output: - logger.info(f"Trial output (job={config.job_name}):\n{obs.output}") - result = await client.trial.collect(client.sandbox, obs.output or "", exit_code) - # G5: populate raw_output / exit_code on every TrialResult so they surface in JobResult - iter_results = result if isinstance(result, list) else [result] for r in iter_results: - if not r.raw_output: - r.raw_output = obs.output or "" - if r.exit_code == 0 and exit_code != 0: - r.exit_code = exit_code - if not success: - fail_info = ExceptionInfo( - exception_type="ProcessTimeout", - exception_message=message or "process did not complete successfully", - ) - for r in iter_results: - if r.exception_info is None: - r.exception_info = fail_info - return result - finally: - if config.environment.auto_stop: - await client.sandbox.close() + if r.exception_info is None: + r.exception_info = fail_info + return result @staticmethod def _build_session_env(config: JobConfig) -> dict[str, str] | None: diff --git a/tests/unit/cli/command/test_job.py b/tests/unit/cli/command/test_job.py index 557fadfc16..0230e377e8 100644 --- a/tests/unit/cli/command/test_job.py +++ b/tests/unit/cli/command/test_job.py @@ -231,7 +231,6 @@ def test_inline_script_content_with_env_and_uploads(self): assert env.cpus == 4.0 assert env.env == {"FOO": "bar", "BAZ": "qux"} assert env.uploads == [("/tmp/workspace", "/root/job")] - assert env.auto_stop is True def test_script_path_mode_no_env_no_uploads(self): ns = self._args(script="run.sh", script_content=None, env=None, local_path=None) @@ -359,7 +358,6 @@ def test_override_image_memory_cpus_on_bash_config(self): assert config.environment.memory == "16g" assert config.environment.cpus == 8.0 assert config.timeout == 900 - assert config.environment.auto_stop is True def test_env_overrides_append_and_overwrite(self): from rock.sdk.bench.models.trial.config import RockEnvironmentConfig @@ -423,7 +421,7 @@ def test_parser_description_mentions_two_modes(self): assert "YAML mode" in desc assert "flags mode" in desc - def test_no_overrides_leaves_config_untouched_except_auto_stop(self): + def test_no_overrides_leaves_config_untouched(self): from rock.sdk.bench.models.trial.config import RockEnvironmentConfig from rock.sdk.job.config import BashJobConfig @@ -442,7 +440,6 @@ def test_no_overrides_leaves_config_untouched_except_auto_stop(self): assert config.environment.image == "python:3.10" assert config.environment.env == {"X": "1"} assert config.timeout == 1234 - assert config.environment.auto_stop is True # always set class TestJobRunEndToEnd: diff --git a/tests/unit/sdk/agent/test_job.py b/tests/unit/sdk/agent/test_job.py index b2a6b7624e..e61a0f49e9 100644 --- a/tests/unit/sdk/agent/test_job.py +++ b/tests/unit/sdk/agent/test_job.py @@ -207,25 +207,11 @@ async def test_run_full_lifecycle(self): # Verify harbor command was started via nohup mock_sandbox.start_nohup_process.assert_called_once() - async def test_run_auto_stop_sandbox(self): + async def test_run_does_not_close_sandbox(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = HarborJobConfig( - job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=True) - ) - job = Job(config) - await job.run() - - mock_sandbox.close.assert_called_once() - - async def test_run_does_not_stop_when_disabled(self): - mock_sandbox = _make_mock_sandbox() - - with patch("rock.sdk.sandbox.client.Sandbox", return_value=mock_sandbox): - config = HarborJobConfig( - job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig(auto_stop=False) - ) + config = HarborJobConfig(job_name="test-job", experiment_id="test-exp", environment=RockEnvironmentConfig()) job = Job(config) await job.run() diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index 25c34cfd56..fbccc3d833 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -39,7 +39,6 @@ def test_job_level_fields(self): env = RockEnvironmentConfig() assert env.env == {} assert env.uploads == [] - assert env.auto_stop is False def test_env_field(self): env = RockEnvironmentConfig(env={"OPENAI_API_KEY": "sk-xxx"}) @@ -72,11 +71,9 @@ def test_excludes_rock_sandbox_fields(self): def test_excludes_job_level_fields(self): env = RockEnvironmentConfig( uploads=[("a", "b")], - auto_stop=True, ) result = env.to_harbor_environment() assert "uploads" not in result - assert "auto_stop" not in result def test_env_passes_through_to_harbor(self): env = RockEnvironmentConfig(env={"KEY": "val"}) @@ -120,7 +117,6 @@ def test_excludes_rock_fields(self): environment=RockEnvironmentConfig( uploads=[("local.txt", "/sandbox/remote.txt")], env={"API_KEY": "sk-xxx"}, - auto_stop=True, image="my-image:latest", memory="32g", ), @@ -132,8 +128,6 @@ def test_excludes_rock_fields(self): assert "sandbox_config" not in data assert "uploads" not in data assert "sandbox_env" not in data - assert "auto_stop_sandbox" not in data - assert "auto_stop" not in data # environment block should only contain harbor fields assert "environment" not in data or "image" not in data.get("environment", {}) @@ -255,7 +249,6 @@ def test_from_yaml_with_environment_block(self, tmp_path): cpus: 8 env: OPENAI_API_KEY: sk-xxx - auto_stop: true agents: - name: terminus-2 """ @@ -266,7 +259,6 @@ def test_from_yaml_with_environment_block(self, tmp_path): assert cfg.environment.image == "my-image:latest" assert cfg.environment.memory == "32g" assert cfg.environment.env == {"OPENAI_API_KEY": "sk-xxx"} - assert cfg.environment.auto_stop is True def test_from_yaml_with_local_dataset(self, tmp_path): yaml_content = """ diff --git a/tests/unit/sdk/agent/test_models.py b/tests/unit/sdk/agent/test_models.py index b968ba47e3..c752cea52f 100644 --- a/tests/unit/sdk/agent/test_models.py +++ b/tests/unit/sdk/agent/test_models.py @@ -236,7 +236,6 @@ def test_environment_defaults(self): cfg = HarborJobConfig(experiment_id="test-exp") assert cfg.environment.uploads == [] assert cfg.environment.env == {} - assert cfg.environment.auto_stop is False def test_with_full_config(self): cfg = HarborJobConfig( diff --git a/tests/unit/sdk/job/test_blue_green_equivalence.py b/tests/unit/sdk/job/test_blue_green_equivalence.py index 6019cf6af9..f0220c9c4e 100644 --- a/tests/unit/sdk/job/test_blue_green_equivalence.py +++ b/tests/unit/sdk/job/test_blue_green_equivalence.py @@ -78,7 +78,7 @@ def _make_config(): version="2.0", ) ], - environment=RockEnvironmentConfig(auto_stop=True), + environment=RockEnvironmentConfig(), ) @@ -132,7 +132,7 @@ async def test_raw_output_and_exit_code_populated_on_both(self): assert blue_result.exit_code == 0 assert green_result.exit_code == 0 - async def test_auto_stop_closes_sandbox_on_both(self): + async def test_sandbox_not_closed_on_both_paths(self): from rock.sdk.bench import Job as BlueJob from rock.sdk.job import Job as GreenJob @@ -144,5 +144,5 @@ async def test_auto_stop_closes_sandbox_on_both(self): with patch("rock.sdk.job.executor.Sandbox", return_value=mock_green): await GreenJob(_make_config()).run() - mock_blue.close.assert_called_once() - mock_green.close.assert_called_once() + mock_blue.close.assert_not_called() + mock_green.close.assert_not_called() diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index bf755e1e00..34ada9dd81 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -36,7 +36,6 @@ def test_defaults(self): assert cfg.experiment_id is None assert cfg.labels == {} assert cfg.timeout == 7200 - assert cfg.environment.auto_stop is False assert cfg.environment.uploads == [] assert cfg.environment.env == {} @@ -45,7 +44,6 @@ def test_custom_values(self): image="ubuntu:22.04", uploads=[("/local/file.py", "/sandbox/file.py")], env={"MY_VAR": "hello"}, - auto_stop=True, ) cfg = JobConfig( environment=env, @@ -60,7 +58,6 @@ def test_custom_values(self): assert cfg.namespace == "team-a" assert cfg.experiment_id == "exp-001" assert cfg.labels == {"step": "42"} - assert cfg.environment.auto_stop is True assert cfg.environment.uploads == [("/local/file.py", "/sandbox/file.py")] assert cfg.environment.env == {"MY_VAR": "hello"} assert cfg.timeout == 7200 @@ -186,7 +183,6 @@ def test_excludes_rock_fields_keeps_harbor_shared_fields(self): experiment_id="my-exp", labels={"step": "1"}, environment=RockEnvironmentConfig( - auto_stop=True, uploads=[("/a", "/b")], env={"KEY": "VAL"}, ), @@ -202,7 +198,7 @@ def test_excludes_rock_fields_keeps_harbor_shared_fields(self): assert data["experiment_id"] == "my-exp" assert data["labels"] == {"step": "1"} # Rock-only — must be absent - rock_only = {"auto_stop", "uploads", "timeout"} + rock_only = {"uploads", "timeout"} for field in rock_only: assert field not in data, f"Rock field '{field}' should be excluded" @@ -405,26 +401,6 @@ def test_harbor_inherits_base_fields(self): assert base_fields.issubset(harbor_fields) -# --------------------------------------------------------------------------- -# G7: HarborJobConfig.auto_stop and environment.auto_stop sync (OR semantics) -# --------------------------------------------------------------------------- - - -class TestHarborJobConfigAutoStopSync: - """auto_stop lives on environment only.""" - - def test_environment_auto_stop_preserved(self): - cfg = HarborJobConfig( - experiment_id="exp-1", - environment=RockEnvironmentConfig(auto_stop=True), - ) - assert cfg.environment.auto_stop is True - - def test_default_auto_stop_is_false(self): - cfg = HarborJobConfig(experiment_id="exp-1") - assert cfg.environment.auto_stop is False - - # --------------------------------------------------------------------------- # G3: HarborJobConfig auto-generates job_name when user omits it # --------------------------------------------------------------------------- diff --git a/tests/unit/sdk/job/test_executor.py b/tests/unit/sdk/job/test_executor.py index c8fcb2e51f..18091aa84d 100644 --- a/tests/unit/sdk/job/test_executor.py +++ b/tests/unit/sdk/job/test_executor.py @@ -153,26 +153,14 @@ async def test_wait_process_failure_sets_exception_info(self): # --------------------------------------------------------------------------- -# auto_stop behavior +# sandbox close behavior # --------------------------------------------------------------------------- -class TestJobExecutorAutoStop: - async def test_auto_stop_true_closes_sandbox(self): +class TestJobExecutorSandboxClose: + async def test_sandbox_not_closed_after_run(self): mock_sandbox = _make_mock_sandbox() with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): - from rock.sdk.envhub import EnvironmentConfig - - config = BashJobConfig(script="echo hi", job_name="test", environment=EnvironmentConfig(auto_stop=True)) - executor = JobExecutor() - await executor.run(ScatterOperator(size=1), config) - - assert mock_sandbox.close.call_count == 1 - - async def test_auto_stop_false_does_not_close_sandbox(self): - mock_sandbox = _make_mock_sandbox() - with patch("rock.sdk.job.executor.Sandbox", return_value=mock_sandbox): - # default: auto_stop=False config = BashJobConfig(script="echo hi", job_name="test") executor = JobExecutor() await executor.run(ScatterOperator(size=1), config) From 8c6efe2f36a17e3f93dcce2c4a60bb4182ce3259 Mon Sep 17 00:00:00 2001 From: dengwx Date: Thu, 16 Apr 2026 11:10:51 +0800 Subject: [PATCH 043/226] fix(job): JobConfig.experiment_id takes priority over environment.experiment_id (#822) Two fixes for experiment_id precedence in the Job system: 1. AbstractTrial.on_sandbox_ready: config experiment_id now wins over sandbox-returned value instead of raising ValueError on mismatch. Sandbox value is only used as a fallback when config has None. 2. JobConfig: add model_validator(mode='after') that overwrites environment.experiment_id when JobConfig.experiment_id differs, and emits a WARNING log to make the override visible. Fixes #821 Co-authored-by: Claude Sonnet 4.6 (1M context) --- rock/sdk/job/config.py | 22 ++++++++++++++- rock/sdk/job/trial/abstract.py | 9 ++---- tests/unit/sdk/job/test_config.py | 37 +++++++++++++++++++++++++ tests/unit/sdk/job/test_trial_bash.py | 14 +++++----- tests/unit/sdk/job/test_trial_harbor.py | 14 +++++----- 5 files changed, 75 insertions(+), 21 deletions(-) diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index c9e81091f8..8a6c753270 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -10,10 +10,13 @@ from __future__ import annotations import yaml -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator +from rock.logger import init_logger from rock.sdk.envhub import EnvironmentConfig +logger = init_logger(__name__) + class JobConfig(BaseModel): """Base config — shared fields for all job types.""" @@ -25,6 +28,23 @@ class JobConfig(BaseModel): labels: dict[str, str] = Field(default_factory=dict) timeout: int = 7200 + @model_validator(mode="after") + def _sync_experiment_id(self) -> JobConfig: + """When both experiment_id fields are set and differ, JobConfig.experiment_id takes priority.""" + if ( + self.experiment_id is not None + and self.environment.experiment_id is not None + and self.experiment_id != self.environment.experiment_id + ): + logger.warning( + "experiment_id conflict: JobConfig has '%s', environment has '%s'. " + "Using JobConfig.experiment_id and overriding environment.experiment_id.", + self.experiment_id, + self.environment.experiment_id, + ) + self.environment.experiment_id = self.experiment_id + return self + @classmethod def from_yaml(cls, path: str) -> JobConfig: """Load a job config from YAML. diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index 3435daee95..2d9d641679 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -44,12 +44,9 @@ async def on_sandbox_ready(self, sandbox: Sandbox) -> None: sb_exp = getattr(sandbox, "_experiment_id", None) if sb_exp is not None: - if self._config.experiment_id is not None and self._config.experiment_id != sb_exp: - raise ValueError( - f"experiment_id mismatch: {type(self._config).__name__} has " - f"'{self._config.experiment_id}', but sandbox returned '{sb_exp}'" - ) - self._config.experiment_id = sb_exp + if self._config.experiment_id is None: + self._config.experiment_id = sb_exp + # If config already has experiment_id, it takes priority over sandbox's value. @abstractmethod async def setup(self, sandbox: Sandbox) -> None: diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 34ada9dd81..7b93bedd21 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -68,6 +68,43 @@ def test_is_base_model(self): assert issubclass(JobConfig, BaseModel) + def test_experiment_id_overrides_environment_experiment_id(self): + """When both experiment_ids differ, JobConfig.experiment_id wins and a warning is logged.""" + from unittest.mock import patch + + import rock.sdk.job.config as job_config_module + + env = EnvironmentConfig(experiment_id="default") + with patch.object(job_config_module.logger, "warning") as mock_warn: + cfg = JobConfig(experiment_id="claw-eval", environment=env) + + assert cfg.environment.experiment_id == "claw-eval" + mock_warn.assert_called_once() + warn_msg = mock_warn.call_args[0][0] + assert "experiment_id" in warn_msg + assert "claw-eval" in str(mock_warn.call_args) + + def test_environment_experiment_id_preserved_when_job_unset(self): + """When only environment.experiment_id is set, it is preserved unchanged.""" + env = EnvironmentConfig(experiment_id="from-env") + cfg = JobConfig(environment=env) + + assert cfg.environment.experiment_id == "from-env" + assert cfg.experiment_id is None + + def test_no_warning_when_experiment_ids_match(self): + """When both experiment_ids are the same, no warning is emitted.""" + from unittest.mock import patch + + import rock.sdk.job.config as job_config_module + + env = EnvironmentConfig(experiment_id="same-exp") + with patch.object(job_config_module.logger, "warning") as mock_warn: + cfg = JobConfig(experiment_id="same-exp", environment=env) + + assert cfg.environment.experiment_id == "same-exp" + mock_warn.assert_not_called() + # --------------------------------------------------------------------------- # BashJobConfig diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index 19a085bf8f..47af787862 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -138,17 +138,17 @@ async def test_namespace_backfilled_when_config_unset(self): assert cfg.namespace == "sb-ns" assert cfg.experiment_id == "exp-1" - async def test_experiment_id_mismatch_raises(self): - import pytest - - cfg = BashJobConfig(script="echo hi", experiment_id="exp-1") + async def test_experiment_id_config_takes_priority_over_sandbox(self): + """Config experiment_id overrides sandbox's different value — no error raised.""" + cfg = BashJobConfig(script="echo hi", experiment_id="claw-eval") trial = BashTrial(cfg) sandbox = MagicMock() sandbox._namespace = None - sandbox._experiment_id = "exp-DIFFERENT" + sandbox._experiment_id = "default" - with pytest.raises(ValueError, match="experiment_id mismatch"): - await trial.on_sandbox_ready(sandbox) + await trial.on_sandbox_ready(sandbox) + + assert cfg.experiment_id == "claw-eval" async def test_namespace_mismatch_raises(self): import pytest diff --git a/tests/unit/sdk/job/test_trial_harbor.py b/tests/unit/sdk/job/test_trial_harbor.py index 163eeef344..0f220ea5e6 100644 --- a/tests/unit/sdk/job/test_trial_harbor.py +++ b/tests/unit/sdk/job/test_trial_harbor.py @@ -155,17 +155,17 @@ async def test_namespace_backfilled_when_config_unset(self): assert cfg.namespace == "sb-ns" - async def test_experiment_id_mismatch_raises(self): - import pytest - - cfg = HarborJobConfig(experiment_id="exp-1") + async def test_experiment_id_config_takes_priority_over_sandbox(self): + """Config experiment_id overrides sandbox's different value — no error raised.""" + cfg = HarborJobConfig(experiment_id="claw-eval") trial = HarborTrial(cfg) sandbox = MagicMock() sandbox._namespace = None - sandbox._experiment_id = "exp-DIFFERENT" + sandbox._experiment_id = "default" - with pytest.raises(ValueError, match="experiment_id mismatch"): - await trial.on_sandbox_ready(sandbox) + await trial.on_sandbox_ready(sandbox) + + assert cfg.experiment_id == "claw-eval" async def test_namespace_mismatch_raises(self): import pytest From f88560f58c3125bedd1cd860e454092a244e932e Mon Sep 17 00:00:00 2001 From: dengwx Date: Thu, 16 Apr 2026 12:28:26 +0800 Subject: [PATCH 044/226] chore: bump version to 1.6.0 (#827) Co-authored-by: Claude Sonnet 4 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 91f2d62faf..bdb710358a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.5.1" +version = "1.6.0" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/uv.lock b/uv.lock index 65ca97dd79..f7f32dc941 100644 --- a/uv.lock +++ b/uv.lock @@ -4035,7 +4035,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.5.1" +version = "1.6.0" source = { editable = "." } dependencies = [ { name = "anyio" }, From 3a516285d41ff515e2fa5788f296d961068093e7 Mon Sep 17 00:00:00 2001 From: dengwx Date: Thu, 16 Apr 2026 13:47:02 +0800 Subject: [PATCH 045/226] docs: update README latest updates table for v1.5.1 and v1.4.7 (#829) * docs: update README latest updates table for v1.5.1 and v1.4.7 Co-Authored-By: Claude Sonnet 4 * release: 1.6.0.dev2 --------- Co-authored-by: Claude Sonnet 4 --- README.md | 3 ++- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 0fe9c9c84a..c67b7a8e06 100644 --- a/README.md +++ b/README.md @@ -163,7 +163,8 @@ if __name__ == "__main__": | 📣 Update Content | |:-----------| -| **[Latest]** 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | +| **[2026-04-16]** 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | +| **[2026-04-10]** 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | | **[2026-03-27]** 🎉 ROCK v1.4.4 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.4) | | **[2026-03-24]** 🎉 ROCK v1.4.3 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.3) | | **[2026-03-17]** 🎉 ROCK v1.4.2 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.2) | diff --git a/pyproject.toml b/pyproject.toml index bdb710358a..a1859f6a01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.6.0" +version = "1.6.0.dev2" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 21a4f9e1cac9000ec6dd977c76ea89f065af260f Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 16 Apr 2026 14:50:50 +0800 Subject: [PATCH 046/226] =?UTF-8?q?feat(job):=20BashJob=20OSS=20Mirror=20?= =?UTF-8?q?=E2=80=94=20artifact=20upload=20after=20job=20completion=20(#82?= =?UTF-8?q?3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(job): BashJob OSS Mirror — artifact upload after job completion - Sandbox: expose namespace and experiment_id as @property - OssMirrorConfig: move to base EnvironmentConfig (shared by all job types) - BashJobConfig.job_name: default to datetime string - env_vars: add ROCK_BASH_JOB_ARTIFACT_DIR (default /tmp/shared) - LinuxFileSystem: add ensure_ossutil() - BashTrial: oss_mirror support in setup() and collect() - validate namespace/experiment_id/oss_bucket before starting - mkdir artifact_dir, install ossutil, upload artifacts (setup + collect) - _build_ossutil_cmd and upload logic in BashTrial itself - instance state in __init__ fixes #824 Co-Authored-By: Claude Sonnet 4.6 (1M context) * refactor(sandbox): download_file reuses ensure_ossutil, remove inline duplicate Co-Authored-By: Claude Sonnet 4.6 (1M context) * refactor(job): rename mirror → oss_mirror, hoist env_vars import to top Co-Authored-By: Claude Sonnet 4.6 (1M context) * refactor(job): self._oss_mirror property, remove all internal imports Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(sandbox): ENSURE_OSSUTIL_SCRIPT auto-installs unzip if missing Co-Authored-By: Claude Sonnet 4.6 (1M context) * feat(claw-eval): write score.json to artifact dir via ROCK_BASH_JOB_ARTIFACT_DIR Co-Authored-By: Claude Sonnet 4.6 * fix(job): touch .keep in artifact_dir so setup upload creates OSS folder Co-Authored-By: Claude Sonnet 4.6 (1M context) * refactor(job): rename .keep to .placeholder in artifact_dir Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(test): update sandbox mock to use PropertyMock for namespace/experiment_id Also update rock/sdk/bench/job.py _autofill_sandbox_info to use sandbox.namespace and sandbox.experiment_id properties (consistent with abstract.py) instead of private _namespace/_experiment_id attributes. Co-Authored-By: Claude Sonnet 4.6 (1M context) * feat(env): change ROCK_BASH_JOB_ARTIFACT_DIR default to /data/logs/user-defined Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(test): update sandbox mock to use PropertyMock for namespace/experiment_id Co-Authored-By: Claude Sonnet 4.6 (1M context) * revert(sandbox): remove namespace/experiment_id @property, revert tests to _namespace style Co-Authored-By: Claude Sonnet 4.6 (1M context) --------- Co-authored-by: Claude Sonnet 4.6 (1M context) --- .../evaluation/claw_eval/run_claw_eval.sh | 18 +++- rock/env_vars.py | 4 + rock/sdk/bench/__init__.py | 2 +- rock/sdk/bench/models/__init__.py | 2 +- rock/sdk/bench/models/job/config.py | 2 +- rock/sdk/bench/models/trial/config.py | 18 +--- rock/sdk/envhub/__init__.py | 4 +- rock/sdk/envhub/config.py | 21 +++- rock/sdk/job/config.py | 3 + rock/sdk/job/trial/bash.py | 91 ++++++++++++++++- rock/sdk/sandbox/constants.py | 10 +- rock/sdk/sandbox/file_system.py | 38 +++---- tests/unit/sdk/agent/test_oss_mirror.py | 13 +-- tests/unit/sdk/job/test_config.py | 37 +++++++ tests/unit/sdk/job/test_trial_bash.py | 98 ++++++++++++++++++- tests/unit/sdk/sandbox/__init__.py | 0 tests/unit/sdk/sandbox/test_file_system.py | 20 ++++ 17 files changed, 324 insertions(+), 57 deletions(-) create mode 100644 tests/unit/sdk/sandbox/__init__.py create mode 100644 tests/unit/sdk/sandbox/test_file_system.py diff --git a/examples/evaluation/claw_eval/run_claw_eval.sh b/examples/evaluation/claw_eval/run_claw_eval.sh index 4f58af81d6..f1ea02d196 100755 --- a/examples/evaluation/claw_eval/run_claw_eval.sh +++ b/examples/evaluation/claw_eval/run_claw_eval.sh @@ -22,7 +22,8 @@ set -eo pipefail -LOG_DIR="/data/logs/user-defined" +ARTIFACT_DIR="${ROCK_BASH_JOB_ARTIFACT_DIR:-/data/logs/user-defined}" +LOG_DIR="$ARTIFACT_DIR/logs" # ── 1. Prepare log directory ─────────────────────────────── mkdir -p "$LOG_DIR" @@ -64,3 +65,18 @@ TOKENS=$(echo "$TEXT" | grep -oP 'tokens=\K\d+' | tail -1) echo "task_score=${TASK_SCORE:-N/A} completion=${COMPLETION:-N/A} robustness=${ROBUSTNESS:-N/A}" echo "communication=${COMMUNICATION:-N/A} safety=${SAFETY:-N/A} passed=${PASSED:-N/A}" echo "wall_time=${WALL_TIME:-N/A}s tokens=${TOKENS:-N/A}" + +# ── 6. Write score.json ─────────────────────────────────── +cat > "$ARTIFACT_DIR/score.json" < None: + def __init__(self, config: BashJobConfig): + super().__init__(config) + self._ossutil_ready: bool = False + self._oss_credentials: dict | None = None + self._artifact_dir: str | None = None + + @property + def _oss_mirror(self): + return self._config.environment.oss_mirror + + async def setup(self, sandbox: Sandbox) -> None: await self._upload_files(sandbox) - # If script_path is set, read content into self._config.script if self._config.script_path: self._config.script = Path(self._config.script_path).read_text() + if self._oss_mirror is not None and self._oss_mirror.enabled: + await self._setup_oss_mirror(sandbox) + + async def _setup_oss_mirror(self, sandbox: Sandbox) -> None: + if not self._config.namespace: + raise ValueError("oss_mirror: namespace is not set (sandbox did not return one)") + if not self._config.experiment_id: + raise ValueError("oss_mirror: experiment_id is not set (sandbox did not return one)") + + self._artifact_dir = env_vars.ROCK_BASH_JOB_ARTIFACT_DIR + + bucket = self._oss_mirror.oss_bucket or os.environ.get("OSS_BUCKET") + if not bucket: + raise ValueError("oss_mirror.enabled=True but oss_bucket is not set (config or OSS_BUCKET env)") + + self._oss_credentials = { + "oss_bucket": bucket, + "access_key_id": self._oss_mirror.oss_access_key_id or os.environ.get("OSS_ACCESS_KEY_ID", ""), + "access_key_secret": self._oss_mirror.oss_access_key_secret or os.environ.get("OSS_ACCESS_KEY_SECRET", ""), + "endpoint": self._oss_mirror.oss_endpoint or os.environ.get("OSS_ENDPOINT", ""), + "region": self._oss_mirror.oss_region or os.environ.get("OSS_REGION", ""), + } + + await sandbox.execute(Command(command=["mkdir", "-p", self._artifact_dir])) + # Touch a placeholder so ossutil cp has something to upload (OSS has no real dirs) + await sandbox.execute(Command(command=["touch", f"{self._artifact_dir}/.placeholder"])) + + self._ossutil_ready = await sandbox.fs.ensure_ossutil() + if not self._ossutil_ready: + logger.warning("ossutil install failed, OSS mirror upload will be skipped") + return + + await self._upload_artifacts(sandbox) + + def _build_oss_prefix(self) -> str: + return f"artifacts/{self._config.namespace}/{self._config.experiment_id}/{self._config.job_name}" + def build(self) -> str: lines = ["#!/bin/bash", "set -e", ""] if self._config.script: lines.append(self._config.script) return "\n".join(lines) - async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: + async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: exception_info = None if exit_code != 0: exception_info = ExceptionInfo( exception_type="BashExitCode", exception_message=f"Bash script exited with code {exit_code}", ) + + if self._oss_mirror is not None and self._oss_mirror.enabled and self._ossutil_ready and self._oss_credentials: + await self._upload_artifacts(sandbox) + return TrialResult( task_name=self._config.job_name or "", exception_info=exception_info, @@ -41,6 +99,33 @@ async def collect(self, sandbox, output: str, exit_code: int) -> TrialResult: exit_code=exit_code, ) + @staticmethod + def _build_ossutil_cmd(ossutil_args: str, creds: dict) -> str: + inner = ( + f"ossutil {ossutil_args}" + f" --access-key-id {shlex.quote(creds['access_key_id'])}" + f" --access-key-secret {shlex.quote(creds['access_key_secret'])}" + f" --endpoint {shlex.quote(creds['endpoint'])}" + f" --region {shlex.quote(creds['region'])}" + ) + return f"bash -c {shlex.quote(inner)}" + + async def _upload_artifacts(self, sandbox: Sandbox) -> None: + try: + oss_url = f"oss://{self._oss_credentials['oss_bucket']}/{self._build_oss_prefix()}/" + src = self._artifact_dir.rstrip("/") + "/" + cmd = self._build_ossutil_cmd( + f"cp {shlex.quote(src)} {shlex.quote(oss_url)} --recursive", + self._oss_credentials, + ) + result = await sandbox.arun(cmd=cmd, mode=RunMode.NOHUP, wait_timeout=600) + if result.exit_code != 0: + logger.warning(f"OSS mirror upload failed: {result.output}") + else: + logger.info(f"OSS mirror upload completed: {self._artifact_dir} -> {oss_url}") + except Exception as e: + logger.warning(f"OSS mirror upload error: {e}") + # Auto-register on import register_trial(BashJobConfig, BashTrial) diff --git a/rock/sdk/sandbox/constants.py b/rock/sdk/sandbox/constants.py index 5d8fe3592b..7d1a9216b3 100644 --- a/rock/sdk/sandbox/constants.py +++ b/rock/sdk/sandbox/constants.py @@ -17,10 +17,14 @@ exit 1 fi -# Check unzip +# Check unzip — try to install if missing if ! command -v unzip >/dev/null 2>&1; then - echo "ERROR: unzip is not available. Please install unzip first." >&2 - exit 1 + echo "unzip not found, attempting to install..." + apt-get install -y -q unzip 2>/dev/null || yum install -y -q unzip 2>/dev/null || true + if ! command -v unzip >/dev/null 2>&1; then + echo "ERROR: unzip is not available and could not be installed." >&2 + exit 1 + fi fi # Skip if already installed diff --git a/rock/sdk/sandbox/file_system.py b/rock/sdk/sandbox/file_system.py index b26b5b3ccc..bb56a41f60 100644 --- a/rock/sdk/sandbox/file_system.py +++ b/rock/sdk/sandbox/file_system.py @@ -238,25 +238,8 @@ async def download_file( "Note: Only regular files are supported. For directories, create a tar archive first.", ) - # Ensure ossutil is installed (checks wget/curl, unzip, installs if missing) - ensure_response = await self.sandbox.process.execute_script( - script_content=ENSURE_OSSUTIL_SCRIPT, - script_name=f"ensure_ossutil_{timestamp}.sh", - cleanup=True, - ) - if ensure_response.exit_code != 0: - return DownloadFileResponse( - success=False, message=f"Failed to ensure ossutil: {ensure_response.output}" - ) - - # Verify ossutil is actually working - verify_response: CommandResponse = await self.sandbox.execute(Command(command=["ossutil", "version"])) - if verify_response.exit_code != 0: - return DownloadFileResponse( - success=False, - message=f"ossutil verification failed (exit_code={verify_response.exit_code}): {verify_response.stderr}", - ) - logger.debug(f"ossutil verified: {verify_response.stdout.strip()}") + if not await self.ensure_ossutil(): + return DownloadFileResponse(success=False, message="Failed to ensure ossutil is installed and working") # Get STS credentials from sandbox (for both ossutil upload and oss2 download) try: @@ -317,3 +300,20 @@ async def download_file( except Exception as e: logger.exception(f"Unexpected error during download_by_oss: {e}") return DownloadFileResponse(success=False, message=f"Unexpected error: {str(e)}") + + async def ensure_ossutil(self) -> bool: + """Ensure ossutil is installed in the sandbox. Returns True if ready.""" + ts = str(time.time_ns()) + res = await self.sandbox.process.execute_script( + script_content=ENSURE_OSSUTIL_SCRIPT, + script_name=f"ensure_ossutil_{ts}.sh", + cleanup=True, + ) + if res.exit_code != 0: + logger.warning(f"ossutil install failed: {res.output}") + return False + verify = await self.sandbox.execute(Command(command=["ossutil", "version"])) + if verify.exit_code != 0: + logger.warning(f"ossutil verify failed: {verify.stderr}") + return False + return True diff --git a/tests/unit/sdk/agent/test_oss_mirror.py b/tests/unit/sdk/agent/test_oss_mirror.py index 8aa0ed8a24..1ff36f5fd0 100644 --- a/tests/unit/sdk/agent/test_oss_mirror.py +++ b/tests/unit/sdk/agent/test_oss_mirror.py @@ -20,17 +20,17 @@ class TestOssMirrorConfig: def test_importable_from_trial_config(self): - from rock.sdk.bench.models.trial.config import OssMirrorConfig + from rock.sdk.envhub.config import OssMirrorConfig assert OssMirrorConfig is not None def test_importable_from_agent_package(self): - from rock.sdk.bench import OssMirrorConfig + from rock.sdk.envhub.config import OssMirrorConfig assert OssMirrorConfig is not None def test_default_is_disabled(self): - from rock.sdk.bench.models.trial.config import OssMirrorConfig + from rock.sdk.envhub.config import OssMirrorConfig cfg = OssMirrorConfig() assert cfg.enabled is False @@ -41,7 +41,7 @@ def test_default_is_disabled(self): assert cfg.oss_endpoint is None def test_all_fields_settable(self): - from rock.sdk.bench.models.trial.config import OssMirrorConfig + from rock.sdk.envhub.config import OssMirrorConfig cfg = OssMirrorConfig( enabled=True, @@ -70,7 +70,7 @@ def test_default_oss_mirror_is_none(self): assert env.oss_mirror is None def test_set_oss_mirror(self): - from rock.sdk.bench.models.trial.config import OssMirrorConfig + from rock.sdk.envhub.config import OssMirrorConfig mirror = OssMirrorConfig(enabled=True, oss_bucket="b1", oss_region="r1") env = EnvironmentConfig(oss_mirror=mirror) @@ -117,7 +117,8 @@ class TestToHarborYamlOssMirror: def test_namespace_at_top_level_in_yaml(self): """namespace/experiment_id 序列化为 HarborJobConfig 顶层字段。""" from rock.sdk.bench.models.job.config import HarborJobConfig - from rock.sdk.bench.models.trial.config import OssMirrorConfig, RockEnvironmentConfig + from rock.sdk.bench.models.trial.config import RockEnvironmentConfig + from rock.sdk.envhub.config import OssMirrorConfig cfg = HarborJobConfig( job_name="mirror-test", diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 7b93bedd21..a436ceceeb 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -634,3 +634,40 @@ def test_harbor_from_yaml_direct_still_works(self, tmp_path): cfg = HarborJobConfig.from_yaml(str(p)) assert isinstance(cfg, HarborJobConfig) assert cfg.n_attempts == 2 + + +# --------------------------------------------------------------------------- +# OssMirrorConfig on base EnvironmentConfig +# --------------------------------------------------------------------------- + + +class TestOssMirrorConfigOnBaseEnvironment: + def test_base_env_default_oss_mirror_is_none(self): + assert EnvironmentConfig().oss_mirror is None + + def test_base_env_accepts_oss_mirror(self): + from rock.sdk.envhub.config import OssMirrorConfig + + cfg = EnvironmentConfig(oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b")) + assert cfg.oss_mirror.enabled is True + + def test_oss_mirror_config_importable_from_envhub(self): + from rock.sdk.envhub.config import OssMirrorConfig + + assert OssMirrorConfig().enabled is False + + +# --------------------------------------------------------------------------- +# BashJobConfig.job_name UUID default +# --------------------------------------------------------------------------- + + +class TestBashJobConfigJobNameDefault: + def test_defaults_to_datetime_string(self): + cfg = BashJobConfig(script="echo hi") + import re + + assert re.match(r"\d{4}-\d{2}-\d{2}__\d{2}-\d{2}-\d{2}", cfg.job_name) + + def test_explicit_name_preserved(self): + assert BashJobConfig(job_name="x").job_name == "x" diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index 47af787862..70bbbe092e 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -6,7 +6,10 @@ from pathlib import Path from unittest.mock import AsyncMock, MagicMock +import pytest + from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.envhub.config import OssMirrorConfig from rock.sdk.job.config import BashJobConfig from rock.sdk.job.trial.bash import BashTrial from rock.sdk.job.trial.registry import _create_trial @@ -151,8 +154,6 @@ async def test_experiment_id_config_takes_priority_over_sandbox(self): assert cfg.experiment_id == "claw-eval" async def test_namespace_mismatch_raises(self): - import pytest - cfg = BashJobConfig(script="echo hi", namespace="cfg-ns") trial = BashTrial(cfg) sandbox = MagicMock() @@ -161,3 +162,96 @@ async def test_namespace_mismatch_raises(self): with pytest.raises(ValueError, match="namespace mismatch"): await trial.on_sandbox_ready(sandbox) + + +# --------------------------------------------------------------------------- +# OSS mirror integration +# --------------------------------------------------------------------------- + + +def _oss_sandbox(ns="ns", exp="exp"): + """Minimal sandbox mock with oss_mirror support.""" + sb = AsyncMock() + sb._namespace = ns + sb._experiment_id = exp + sb.arun = AsyncMock(return_value=MagicMock(exit_code=0, output="")) + sb.fs.ensure_ossutil = AsyncMock(return_value=True) + sb.fs.upload_dir = AsyncMock(return_value=MagicMock(exit_code=0)) + return sb + + +_MIRROR = OssMirrorConfig(enabled=True, oss_bucket="b", oss_endpoint="ep", oss_region="rg") + + +class TestBashTrialOssMirror: + async def test_setup_installs_ossutil_and_creates_dir(self): + cfg = BashJobConfig( + script="echo", + job_name="j", + namespace="ns", + experiment_id="exp", + environment=EnvironmentConfig(oss_mirror=_MIRROR), + ) + trial = BashTrial(cfg) + sb = _oss_sandbox() + await trial.setup(sb) + + sb.fs.ensure_ossutil.assert_called_once() + # Initial upload to create OSS path before script runs + setup_cp_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] + assert len(setup_cp_calls) == 1 + + async def test_setup_skips_when_no_mirror(self): + trial = BashTrial(BashJobConfig(script="echo")) + sb = _oss_sandbox() + await trial.setup(sb) + sb.fs.ensure_ossutil.assert_not_called() + + async def test_collect_uploads(self): + cfg = BashJobConfig( + script="echo", + job_name="j", + namespace="ns", + experiment_id="exp", + environment=EnvironmentConfig(oss_mirror=_MIRROR), + ) + trial = BashTrial(cfg) + sb = _oss_sandbox() + await trial.setup(sb) + await trial.collect(sb, "ok", 0) + + # setup + collect each call ossutil cp once + arun_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] + assert len(arun_calls) == 2 + assert all("oss://b/artifacts/ns/exp/j/" in str(c) for c in arun_calls) + + async def test_upload_failure_does_not_fail_job(self): + cfg = BashJobConfig( + script="echo", + job_name="j", + namespace="ns", + experiment_id="exp", + environment=EnvironmentConfig(oss_mirror=_MIRROR), + ) + trial = BashTrial(cfg) + sb = _oss_sandbox() + sb.arun = AsyncMock(return_value=MagicMock(exit_code=1, output="err")) + await trial.setup(sb) + result = await trial.collect(sb, "ok", 0) + assert result.exit_code == 0 and result.exception_info is None + + async def test_skips_upload_when_ossutil_not_ready(self): + cfg = BashJobConfig( + script="echo", + job_name="j", + namespace="ns", + experiment_id="exp", + environment=EnvironmentConfig(oss_mirror=_MIRROR), + ) + trial = BashTrial(cfg) + sb = _oss_sandbox() + sb.fs.ensure_ossutil = AsyncMock(return_value=False) + await trial.setup(sb) + await trial.collect(sb, "ok", 0) + ossutil_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] + assert len(ossutil_calls) == 0 diff --git a/tests/unit/sdk/sandbox/__init__.py b/tests/unit/sdk/sandbox/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/sdk/sandbox/test_file_system.py b/tests/unit/sdk/sandbox/test_file_system.py new file mode 100644 index 0000000000..402da4dd9a --- /dev/null +++ b/tests/unit/sdk/sandbox/test_file_system.py @@ -0,0 +1,20 @@ +"""Tests for LinuxFileSystem OSS methods.""" + +from unittest.mock import AsyncMock, MagicMock + +from rock.sdk.sandbox.file_system import LinuxFileSystem + + +def _sandbox(exit_code=0): + sb = AsyncMock() + sb.process.execute_script = AsyncMock(return_value=MagicMock(exit_code=exit_code, output="")) + sb.execute = AsyncMock(return_value=MagicMock(exit_code=exit_code, stdout="v2", stderr="")) + return sb + + +class TestEnsureOssutil: + async def test_success(self): + assert await LinuxFileSystem(_sandbox()).ensure_ossutil() is True + + async def test_install_failure(self): + assert await LinuxFileSystem(_sandbox(exit_code=1)).ensure_ossutil() is False From 949ed7e786febdd65eb03501df89ac4da5ff8006 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 16 Apr 2026 15:05:33 +0800 Subject: [PATCH 047/226] fix(job): remove redundant shebang and set -e from BashTrial.build() (#816) The executor always calls `bash {script_path}` directly, making `#!/bin/bash` a no-op. `set -e` was silently injected without user intent, causing unexpected exits. Simplify build() to return the script as-is. Closes #815 Co-authored-by: Claude Sonnet 4.6 --- rock/sdk/job/trial/bash.py | 5 +---- tests/unit/sdk/job/test_trial_bash.py | 10 ++++++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index 185dca8737..ed084eea69 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -76,10 +76,7 @@ def _build_oss_prefix(self) -> str: return f"artifacts/{self._config.namespace}/{self._config.experiment_id}/{self._config.job_name}" def build(self) -> str: - lines = ["#!/bin/bash", "set -e", ""] - if self._config.script: - lines.append(self._config.script) - return "\n".join(lines) + return self._config.script or "" async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: exception_info = None diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index 70bbbe092e..ff3e4d5f9d 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -30,10 +30,12 @@ class TestBashTrialBuild: def test_build_basic_script(self): cfg = BashJobConfig(script="echo hello") trial = BashTrial(cfg) - out = trial.build() - assert "#!/bin/bash" in out - assert "set -e" in out - assert "echo hello" in out + assert trial.build() == "echo hello" + + def test_build_empty_script(self): + cfg = BashJobConfig(script=None) + trial = BashTrial(cfg) + assert trial.build() == "" # --------------------------------------------------------------------------- From c913f2b85500790b839ab40e53eccdeb8863e0cc Mon Sep 17 00:00:00 2001 From: dengwx Date: Thu, 16 Apr 2026 15:45:24 +0800 Subject: [PATCH 048/226] fix(cli): lazy import psutil in admin stop command (#831) Co-authored-by: Claude Sonnet 4 --- rock/cli/command/admin.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/rock/cli/command/admin.py b/rock/cli/command/admin.py index cad41d320d..dc97c50e48 100644 --- a/rock/cli/command/admin.py +++ b/rock/cli/command/admin.py @@ -2,8 +2,6 @@ import asyncio import subprocess -import psutil - from rock.cli.command.command import Command as CliCommand from rock.logger import init_logger @@ -37,6 +35,11 @@ async def _admin_start(self, args: argparse.Namespace): async def _admin_stop(self, args: argparse.Namespace): """Stop admin service""" + try: + import psutil + except ImportError: + raise ImportError("psutil is required for 'rock admin stop'. Install it with: pip install psutil") + try: # Find admin processes admin_processes = [] From ff976a5d513803098ea0b94a8f1a6940139f2ebe Mon Sep 17 00:00:00 2001 From: guoj14 Date: Fri, 17 Apr 2026 11:39:16 +0800 Subject: [PATCH 049/226] feat(ci): add unit test workflow for TS SDK (#796) --- .github/workflows/ts-sdk-ci.yml | 43 +++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 .github/workflows/ts-sdk-ci.yml diff --git a/.github/workflows/ts-sdk-ci.yml b/.github/workflows/ts-sdk-ci.yml new file mode 100644 index 0000000000..abae30af48 --- /dev/null +++ b/.github/workflows/ts-sdk-ci.yml @@ -0,0 +1,43 @@ +name: TS SDK CI + +on: + push: + branches: [master, release/**] + pull_request: + branches: [master, release/**] + workflow_dispatch: + +concurrency: + group: ts-sdk-ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + defaults: + run: + working-directory: rock/ts-sdk + steps: + - uses: actions/checkout@v4 + + - name: Setup Node.js + uses: actions/setup-node@v4 + with: + node-version: "20" + + - name: Install pnpm + run: | + echo "=== Task started at: $(date '+%Y-%m-%d %H:%M:%S') ===" + npm install -g pnpm + + - name: Install dependencies + run: | + echo "=== Task started at: $(date '+%Y-%m-%d %H:%M:%S') ===" + pnpm install --frozen-lockfile + + - name: Run unit tests + run: | + echo "=== Task started at: $(date '+%Y-%m-%d %H:%M:%S') ===" + pnpm test:unit + + From 1de5365d6c7951edbd2682ab001c19a2047e84dc Mon Sep 17 00:00:00 2001 From: dengwx Date: Fri, 17 Apr 2026 13:55:45 +0800 Subject: [PATCH 050/226] fix(job): sync experiment_id to environment when environment.experiment_id is None (#835) Previously _sync_experiment_id only propagated when both fields were set and differed. When environment.experiment_id was None the value was never copied down, silently leaving downstream environment without an experiment_id. closes #834 Co-authored-by: Claude Sonnet 4 --- rock/sdk/job/config.py | 26 ++++++++++++++------------ tests/unit/sdk/job/test_config.py | 13 +++++++++++++ 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/rock/sdk/job/config.py b/rock/sdk/job/config.py index 6e989f8675..8498a32bff 100644 --- a/rock/sdk/job/config.py +++ b/rock/sdk/job/config.py @@ -32,18 +32,20 @@ class JobConfig(BaseModel): @model_validator(mode="after") def _sync_experiment_id(self) -> JobConfig: - """When both experiment_id fields are set and differ, JobConfig.experiment_id takes priority.""" - if ( - self.experiment_id is not None - and self.environment.experiment_id is not None - and self.experiment_id != self.environment.experiment_id - ): - logger.warning( - "experiment_id conflict: JobConfig has '%s', environment has '%s'. " - "Using JobConfig.experiment_id and overriding environment.experiment_id.", - self.experiment_id, - self.environment.experiment_id, - ) + """Sync JobConfig.experiment_id down to environment.experiment_id. + + - If only JobConfig.experiment_id is set, propagate it to environment silently. + - If both are set and differ, JobConfig.experiment_id wins and a warning is logged. + - If both are set and equal, or JobConfig.experiment_id is None, do nothing. + """ + if self.experiment_id is not None: + if self.environment.experiment_id is not None and self.experiment_id != self.environment.experiment_id: + logger.warning( + "experiment_id conflict: JobConfig has '%s', environment has '%s'. " + "Using JobConfig.experiment_id and overriding environment.experiment_id.", + self.experiment_id, + self.environment.experiment_id, + ) self.environment.experiment_id = self.experiment_id return self diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index a436ceceeb..8c678ab337 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -105,6 +105,19 @@ def test_no_warning_when_experiment_ids_match(self): assert cfg.environment.experiment_id == "same-exp" mock_warn.assert_not_called() + def test_experiment_id_synced_to_environment_when_env_is_none(self): + """When JobConfig.experiment_id is set and environment.experiment_id is None, + it should be synced down to environment without a warning.""" + from unittest.mock import patch + + import rock.sdk.job.config as job_config_module + + with patch.object(job_config_module.logger, "warning") as mock_warn: + cfg = JobConfig(experiment_id="exp-sync") + + assert cfg.environment.experiment_id == "exp-sync" + mock_warn.assert_not_called() + # --------------------------------------------------------------------------- # BashJobConfig From 84fde4caf07f6f6bdf61997b41029f43f53a4f4b Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 17 Apr 2026 16:43:15 +0800 Subject: [PATCH 051/226] docs: add v1.6.0 release notes and register 1.6.x docs version - Create version-1.6.x documentation (en + zh-Hans) based on 1.5.x - Author v1.6.0 release note covering 23 PRs since v1.5.1: Job module refactor (Job/Operator/Executor/Trial), BashJob + HarborJob, CLI rework, EnvHub refactor, default timeout bump, breaking changes - Register 1.6.x in versions.json and update lastVersion to 1.6.x - Bump pyproject.toml version from 1.6.0.dev2 to 1.6.0 refs #840 Co-Authored-By: Claude Opus 4.7 --- docs/docusaurus.config.js | 2 +- .../version-1.6.x.json | 34 ++ .../Getting Started/installation.md | 141 +++++++++ .../Getting Started/quickstart.md | 172 ++++++++++ .../Getting Started/rock-agent.md | 73 +++++ .../version-1.6.x/Getting Started/rockroll.md | 200 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 69 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 137 ++++++++ .../Python SDK References/sandbox.md | 113 +++++++ .../swe-bench-evaluation.md | 228 ++++++++++++++ .../version-1.6.x/References/api.md | 194 ++++++++++++ .../version-1.6.x/Release Notes/index.md | 5 + .../version-1.6.x/Release Notes/v1.6.0.md | 111 +++++++ .../User Guides/configuration.md | 188 +++++++++++ .../version-1.6.x/overview.md | 40 +++ .../Getting Started/installation.md | 143 +++++++++ .../Getting Started/quickstart.md | 166 ++++++++++ .../Getting Started/rock-agent.md | 72 +++++ .../version-1.6.x/Getting Started/rockroll.md | 194 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 70 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 136 ++++++++ .../Python SDK References/sandbox.md | 114 +++++++ .../swe-bench-evaluation.md | 229 ++++++++++++++ .../version-1.6.x/References/api.md | 195 ++++++++++++ .../version-1.6.x/Release Notes/index.md | 5 + .../version-1.6.x/Release Notes/v1.6.0.md | 111 +++++++ .../User Guides/configuration.md | 189 +++++++++++ docs/versioned_docs/version-1.6.x/overview.md | 33 ++ .../version-1.6.x-sidebars.json | 64 ++++ docs/versions.json | 1 + pyproject.toml | 2 +- 43 files changed, 5645 insertions(+), 2 deletions(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x.json create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/installation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/quickstart.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rockroll.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/codes.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/deploy.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/file_system.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/model-service.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/python_sdk.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/remote_user.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/runtime-env.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/sandbox.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/api.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.0.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/User Guides/configuration.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/overview.md create mode 100644 docs/versioned_docs/version-1.6.x/Getting Started/installation.md create mode 100644 docs/versioned_docs/version-1.6.x/Getting Started/quickstart.md create mode 100644 docs/versioned_docs/version-1.6.x/Getting Started/rock-agent.md create mode 100644 docs/versioned_docs/version-1.6.x/Getting Started/rockroll.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/codes.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/deploy.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/file_system.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/model-service.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/python_sdk.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/remote_user.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/rock-agent.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/runtime-env.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/sandbox.md create mode 100644 docs/versioned_docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/versioned_docs/version-1.6.x/References/api.md create mode 100644 docs/versioned_docs/version-1.6.x/Release Notes/index.md create mode 100644 docs/versioned_docs/version-1.6.x/Release Notes/v1.6.0.md create mode 100644 docs/versioned_docs/version-1.6.x/User Guides/configuration.md create mode 100644 docs/versioned_docs/version-1.6.x/overview.md create mode 100644 docs/versioned_sidebars/version-1.6.x-sidebars.json diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index 12e74b8659..107ff7a6ef 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -144,7 +144,7 @@ const config = { // release note按照版本号倒排 return reverseReleaseNoteSidebars(filterHiddenSidebars); }, - lastVersion: '1.5.x', + lastVersion: '1.6.x', includeCurrentVersion: false, versions: convertVersionsArrayToObject(versions) }, diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x.json b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x.json new file mode 100644 index 0000000000..87344499c7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x.json @@ -0,0 +1,34 @@ +{ + "version.label": { + "message": "1.6.x", + "description": "The label for version 1.6.x" + }, + "sidebar.tutorialSidebar.category.Getting Started": { + "message": "快速上手", + "description": "The label for category 'Getting Started' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.User Guides": { + "message": "用户指南", + "description": "The label for category 'User Guides' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.References": { + "message": "参考", + "description": "The label for category 'References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Release Notes": { + "message": "版本说明", + "description": "The label for category 'Release Notes' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.model-service": { + "message": "Model Service 参考", + "description": "The label for category 'Model Service References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.sandbox-agent": { + "message": "Sandbox Agent参考", + "description": "The label for category 'Sandbox Agent References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Python SDK References": { + "message": "Python SDK 参考", + "description": "The label for category 'Python SDK References' in sidebar 'tutorialSidebar'" + } +} diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/installation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/installation.md new file mode 100644 index 0000000000..0ab70e55d1 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/installation.md @@ -0,0 +1,141 @@ +--- +sidebar_position: 3 +--- + +# 安装指南 + +本文档介绍如何使用 `uv` 和 `pip` 安装和设置 ROCK 开发环境。该项目是一个强化学习开放构建工具包，支持多种组件。 + +## 使用 uv（推荐） + +### 快速安装所有依赖 + +```bash +# 安装所有依赖（包括可选依赖） +uv sync --all-extras + +# 安装开发/测试依赖 +uv sync --all-extras --all-groups +``` + +### 安装不同依赖组 + +#### 仅核心依赖 +```bash +uv sync +``` + +#### 管理组件依赖 +```bash +uv sync --extra admin +``` + +#### Rocklet 执行环境依赖 +```bash +uv sync --extra rocklet +``` + +#### 所有依赖 +```bash +uv sync --all-extras +``` + +#### 开发/测试依赖 +```bash +uv sync --all-extras --group test +``` + +## 使用 pip + +### 从 pip 源安装 + +#### 仅核心依赖 +```bash +pip install rl-rock +``` + +#### 管理组件依赖 +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install "rl-rock[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install "rl-rock[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install "rl-rock[all]" +``` + +### 使用 pip 从源码安装 + +#### 仅核心依赖 +```bash +pip install . +``` + +#### 管理组件依赖 +```bash +pip install ".[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install ".[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install ".[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install ".[all]" +``` + +## 可用入口点 + +该包提供以下命令行脚本： + +- `rocklet`: ROCK 执行环境服务器 (rock.rocklet.server:main) +- `admin`: 管理服务器 (rock.admin.main:main) +- `envhub`: 环境中心服务器 (rock.envhub.server:main) +- `rock`: 主 ROCK 命令行接口 (rock.cli.main:main) + +## 开发设置 + +### 使用 uv（推荐） + +```bash +# 克隆并设置开发环境 +git clone +cd ROCK +uv sync --all-extras --group test + +# 运行测试 +uv run pytest + + +### 使用 pip + +```bash +# 开发模式安装所有可选依赖 +pip install -e ".[all]" + +# 分别安装 +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" +``` + +## 附加说明 + +- 项目配置为默认使用阿里云 PyPI 镜像: `https://mirrors.aliyun.com/pypi/simple/` +- 对于本地开发，运行测试需要 `test` 依赖组 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/quickstart.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..e0a0891c0e --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/quickstart.md @@ -0,0 +1,172 @@ +--- +sidebar_position: 2 +--- + +# 快速上手 + +本指南将通过完整的示例演示如何使用 ROCK 创建和管理强化学习环境。ROCK (Reinforcement Open Construction Kit) 是一个全面的沙箱环境管理框架，主要用于强化学习和AI开发环境。 + +## 1. 环境准备 + +我们推荐在 Linux 系统下启动 ROCK，能够尽量复用项目依赖，提升环境拉起速度。如果需要在 macOS 上尝试，可以参考 [MacOS 启动](#7-macos-启动) 一节。 + +在开始之前，请确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且示例中依赖python:3.11镜像 +docker pull python:3.11 + +# 验证 uv 安装 +uv --version + + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone +cd ROCK + +# 创建虚拟环境（使用 uv 托管的 Python, 以python 3.11 版本为例） +uv venv --python 3.11 --python-preference only-managed + +# 安装所有依赖组 +uv sync --all-extras +``` + +> **重要提示**: 为确保 ROCK 能正确挂载项目和虚拟环境及其依赖的 base Python 解释器，强烈推荐使用 uv 托管的 Python 环境而非系统 Python。 + +## 2. 激活虚拟环境 + +在运行任何 ROCK 命令之前，需要先激活虚拟环境。确保 sys.base_prefix 是 uv 管理的环境，类似于 `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` 等路径。 + +```bash +# 激活虚拟环境 +source .venv/bin/activate + +# 验证 Python 环境 +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **验证要点**: 确保输出的 base prefix 路径指向 uv 管理的 Python 环境，而非系统 Python。 + +## 3. 验证环境配置 + +激活虚拟环境后，验证依赖安装是否正确： + +```bash +# 检查关键依赖 +python -c "import rock; print(\"Hello ROCK\")" +``` + + +## 4. 启动 ROCK 服务 + +激活虚拟环境后，在项目根目录下，启动 ROCK Admin 服务： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 启动 ROCK Admin 服务（本地环境） +rock admin start +``` + +服务启动后，您将看到类似以下的输出： + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **服务说明**: ROCK Admin 服务默认运行在 `http://127.0.0.1:8080`。 + +## 5. 运行示例环境 + +现在可以运行示例环境来验证安装。确保 ROCK 服务正在运行，然后打开一个新的终端窗口执行以下命令： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 运行沙箱示例 +python examples/sandbox_demo.py + +# 运行 GEM 协议示例 +python examples/sokoban_demo.py +``` + +### 5.1 示例说明 + +- **sandbox_demo.py**: 演示如何使用 ROCK 的沙箱 SDK 创建和管理容器化环境 +- **sokoban_demo.py**: 演示如何使用 ROCK 的 GEM 协议兼容接口创建强化学习环境 + +> **运行要求**: 确保 ROCK Admin 服务正在运行，因为示例需要与服务进行通信。 + +## 6. 分布式环境配置（可选） + +对于分布式多机器环境，请确保以下配置一致： + +1. 所有机器上 ROCK 和 uv 的 Python 配置使用相同的根 Python 解释器 +2. Docker 版本在所有节点上保持一致 +3. 网络配置允许各节点间正常通信 + + + +## 7. MacOS 启动 + +在 macOS 上，如果需要启动 Linux 镜像的环境，需要先设置环境变量： + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +在容器启动时，会安装对应的 uv 环境，细节可以参考 `rock/rocklet/local_files/docker_run_with_uv.sh` 脚本。 + +> **注意**: 相比 Linux 系统，macOS 上的启动速度会较慢，且比较依赖网络环境，可以根据实际情况调整脚本。ROCK_WORKER_ENV_TYPE的细节可以参考 [Configuration Guide](../User%20Guides/configuration.md). + + +## 8. 从Pip源启动 + +如果从Pip源启动Admin Server，在参照[安装指南](./installation.md)安装完成ROCK后, 需要设置额外环境变量: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(这一启动方式在容器环境启动时会从Pypi源上拉取最新的rocklet并安装, 相对启动速度比较慢, 仅推荐测试使用, 生产上依旧推荐其他的启动方式) + + +## 总结 + +恭喜！您已经成功完成了 ROCK 的快速开始指南。现在您应该能够： + +- 正确设置 ROCK 开发环境 +- 使用 uv 管理的 Python 环境 +- 启动和管理 ROCK 服务 +- 运行示例程序验证安装 +- 在分布式环境中配置 ROCK（如果需要） + +如需深入了解 ROCK 的更多功能，请参考以下文档： + +## 下一步学习 + +- [配置指南](../User%20Guides/configuration.md) - 详细了解 ROCK 的配置选项 +- [API 文档](../References/api.md) - 查看完整的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +- [安装指南](./installation.md) - 详细了解 ROCK 安装和配置 +- [概览](../overview.md) - 了解 ROCK 的设计理念 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..bd2dd69b05 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rock-agent.md @@ -0,0 +1,73 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent 快速启动 + +Rock Agent 是 ROCK 提供的 AI Agent 运行框架，支持在沙箱环境中运行各种类型的 Agent。 + +## 前置条件 + +- 确保有可用的ROCK服务, 如果需要本地拉起服务端, 参考[快速启动](quickstart.md) + +## 使用示例 + +ROCK 提供了两个Hello World Agent 示例，位于 `examples/agents/` 目录下： + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent 示例 +└── iflow_cli/ # IFlowCli Agent 示例 +``` + +### 运行 IFlowCli 示例 + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### 运行 ClaudeCode 示例 + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli 配置文件 + +配置文件位于 `examples/agents/iflow_cli/rock_agent_config.yaml`： + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # 填入你的 API Key + IFLOW_BASE_URL: "" # 填入你的 Base URL + IFLOW_MODEL_NAME: "" # 填入你的模型名称 +``` + +## ClaudeCode 配置文件 + +配置文件位于 `examples/agents/claude_code/rock_agent_config.yaml`： + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # 填入你的anthropic base url + ANTHROPIC_API_KEY: "" # 填入你的anthropic api key +``` + +## 相关文档 + +- [RockAgent 参考](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rockroll.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..3b53810ba3 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Getting Started/rockroll.md @@ -0,0 +1,200 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL 快速开始指南 + +本指南将引导您使用 ROLL (训练框架) 和 ROCK (环境管理) 来运行一个基于 Sokoban 游戏（推箱子）的强化学习训练示例。 + +## 1. 单机环境准备 + +在开始之前，请先确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **操作系统**: 推荐使用 Linux (如 Ubuntu 20.04+) +- **硬件**: 建议使用 NVIDIA GPU 并安装对应的驱动程序 +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且可提前拉取 Sokoban 游戏环境镜像，避免训练时等待 +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# 验证 uv 安装 +uv --version + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# 确保两个仓库位于同一级目录下，如下所示： +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. 启动训练流程 + +> 说明：下文均以 *torch2.6.0 + vLLM0.8.4* 为例。 + + +### 方式一： 使用虚拟环境启动（推荐） + +#### 为什么推荐这种方式？ +- 隔离性：uv 虚拟环境能确保项目依赖与系统环境隔离，避免冲突。 +- 速度快：ROCK 可以复用此虚拟环境，大大加快了后续环境的启动速度。 +- 稳定性：依赖关系更清晰，环境更易复现。 + + +```bash +# 进入 ROCK 目录 +cd ROCK + +# 使用 uv 创建并激活 Python 3.10 虚拟环境（ROLL推荐使用Python 3.10） +uv venv --python 3.10 --python-preference only-managed + +# 激活虚拟环境 +source .venv/bin/activate + +# 使用uv安装ROCK的依赖 +uv sync --all-extras + +# 若使用Python 3.10, 启动 ray 时会报错：ValueError: is not a valid Sentinel +# 原因是 ray 与 click>=8.3 版本不兼容，需要降级到 click<8.3 +# Python 3.11 不会有这个问题 +uv pip install 'click>=8.2,<8.3' + +# 切换到 ROLL 目录以安装其依赖 +cd ../ROLL + +# 设置国内 PyPI 镜像源以加速下载 +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装核心 PyTorch 组件 +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -i $PYPI_MIRROR + +# 安装transformer-engine，--no-build-isolation 避免因环境隔离导致找不到 torch +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation -i $PYPI_MIRROR + +# 安装预编译的 flash-attention，以匹配特定的 CUDA 和 PyTorch 版本 +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# 安装其余依赖 +uv pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# (可选) 安装Tensorboard，用于查看训练指标 +uv pip install tensorboard -i $PYPI_MIRROR + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### 方式二：使用系统环境启动（备选方案） + +为获得最佳兼容性，推荐使用 ROLL 官方提供的基础 Docker 镜像，因为它们已经预装了匹配的 CUDA、cuDNN 和其他基础库。 + +> [ROLL 官方镜像列表](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Installation/image_address/) + + +#### 注意 +此方式会将所有 Python 包直接安装到您的当前环境（例如，容器的基础环境）中，可能会与系统自带的包或其他项目产生冲突。 + +由于 ROCK 无法复用环境，每次启动任务时都可能需要重新安装部分依赖，启动速度较慢且受网络影响。 + + +```bash +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装ROCK的依赖 +cd ROCK +pip install . -i $PYPI_MIRROR +pip install ".[admin]" -i $PYPI_MIRROR + +# 安装ROLL的依赖 +cd ../ROLL +pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# 配置ROCK用uv启动的环境变量 +export ROCK_WORKER_ENV_TYPE=uv + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +至此，您已成功启动了 Sokoban 强化学习训练流程。祝您 Rock & Roll 愉快！ + + +## 3. 多机部署 + +除了在单机上运行，您也可以将 **ROCK 服务** 和 **ROLL 训练** 部署在不同的机器上，通过网络进行通信。这是一种常见的服务化部署模式。 + +### 3.1 在机器 A 上部署 ROCK 服务 + +在一台独立的机器（或容器）上，参照[ROCK快速指南](./quickstart.md)部署并启动 ROCK 服务。 + +> **重要提示** +> 启动服务后，请记下ROCK服务的IP地址和端口，例如`http://192.168.1.10:8000`，后续步骤将需要这个地址。 + +### 3.2 在机器 B 上准备 ROLL 客户端 + +在另一台将要运行训练任务的机器上，执行以下操作。 + +1. 验证网络连通性 + +首先，使用 curl 命令检查是否能从机器 B 访问到机器 A 上的 ROCK 服务。 +```bash +# 将 : 替换为您的 ROCK 服务实际地址 +# 如果成功，会收到 ROCK 服务的响应 {"message":"hello, ROCK!"} +curl http://: +``` + +2. 准备 ROLL 环境 + +```bash +# 克隆 ROLL 仓库 +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# 安装依赖 +pip install -r requirements_torch260_vllm.txt -i https://mirrors.aliyun.com/pypi/simple/ +``` + +3. 配置 ROLL 连接地址 + +修改 ROLL 的配置文件，使其能够找到并连接到远程的 ROCK 服务。 +- 打开配置文件：examples/agentic_demo/agentic_val_sokoban_sandbox.yaml +- 找到 SokobanSandbox 下的 env_config 部分 +- 将 base_url 的值修改为您的 ROCK 服务地址 +```yaml +custom_envs: + SokobanSandbox: + env_config: + # 将这里的地址修改为您的 ROCK 服务地址 + # 例如: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. 启动训练 +配置完成后，即可在机器 B 上启动 ROLL 训练脚本。 + +```bash +# 此脚本现在会通过网络请求机器 A 上的 ROCK 服务来创建环境 +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### 进阶：分布式 ROLL 训练 + +如果您希望将 ROLL 训练任务本身进行分布式部署，可以参考 ROLL 的官方分布式部署文档。 +> [快速上手：多节点部署指南](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/codes.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..47b74166de --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +错误码定义和分类，用于错误处理和重试策略。 + +## 使用示例 + +```python +import rock + +def test_codes_values(): + """测试基本状态码值""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes 分类 + +```python +OK = 2000, "OK" +""" +成功状态码 (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +客户端错误码 (4xxx): + +这些错误表示客户端请求有问题， +SDK 会抛出异常。 +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +服务端错误码 (5xxx): + +这些错误表示服务端出现问题， +SDK 会抛出异常。 +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +命令/执行错误码 (6xxx): + +这些错误与命令执行相关，由模型处理， +SDK 不会抛出异常。 +""" +``` + +## 重试策略建议 + +- **重试触发条件**: 只有当 `INTERNAL_SERVER_ERROR` 时才需要重试 +- **其他情况的处理策略**: + - `BAD_REQUEST`: 需要检查 arun 调用逻辑是否有异常 + - `COMMAND_ERROR`: stdout 输出到 `observation.output`，stderr 输出到 `observation.failure_reason` +- `COMMAND_ERROR` 说明: 由于 bash 执行失败时，stdout/stderr 可能全部非空，建议将 observation 中 output 和 failure_reason 全部 prompt 给模型进行推理 + +## 重试示例 + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/deploy.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..b7bd2da08f --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +沙箱资源部署管理器，用于本地目录部署和模板格式化。 + +## deploy_working_dir - 部署本地目录 + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# 部署本地目录到沙箱（自动生成目标路径） +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"部署到: {target}") # 例如: /tmp/rock_workdir_abc123 + +# 部署到指定目标路径 +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - 模板变量替换 + +`format` 方法支持两种模板语法： + +- **`${variable}`** - 标准 Python 字符串模板语法 +- **`<>`** - 替代语法（内部转换为 `${variable}`） + +```python +# 使用 ${working_dir} 模板变量 +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# 结果: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# 使用 <<>> 替代语法 +cmd = deploy.format("cat <>/file.txt") +# 结果: cat /tmp/rock_workdir_abc123/file.txt + +# 结合自定义变量使用 +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# 结果: cat /tmp/rock_workdir_abc123/settings.json + +# Shell 语法保持不变 +cmd = deploy.format("echo $((3 << 2 >> 1))") +# 结果: echo $((3 << 2 >> 1)) + +# 直接访问 working_dir +if deploy.working_dir: + print(f"当前工作目录: {deploy.working_dir}") +``` + +## 多次部署 + +后续调用会覆盖之前的工作目录路径： + +```python +# 第一次部署 +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# 第二次部署（覆盖之前的路径） +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/file_system.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..741b14f5e4 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +文件系统操作接口，提供沙箱环境中的权限管理和目录上传功能。 + +## chown - 修改所有者 + +```python +from rock.actions.sandbox.request import ChownRequest + +# 创建远程用户后修改所有者 +await sandbox.remote_user.create_remote_user("deploy") + +# 获取当前目录 +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# 修改目录所有者 +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# 递归修改目录及其内容所有者 +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - 修改权限 + +```python +from rock.actions.sandbox.request import ChmodRequest + +# 创建测试目录 +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# 修改目录权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# 递归修改权限（包括子目录和文件） +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# 设置最高权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - 上传目录 + +```python +import os +from pathlib import Path + +# 准备本地目录 +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# 上传到沙箱 +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"上传成功: {result.output}") +else: + print(f"上传失败: {result.failure_reason}") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/model-service.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..ba158cf75a --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service（实验性） + +ROCK 提供的 Model Service 负责处理 AI 模型调用的通信，为代理(Agent)和训练框架(如 Roll)或实际的 LLM 推理服务之间提供通信桥梁。 + +## 与 RockAgent 集成 + +ModelService 通常由 **RockAgent** 自动管理，无需手动调用生命周期方法。只需在配置中启用即可： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # 启用 ModelService，RockAgent 会自动管理其生命周期 +) +``` + +RockAgent 会自动： +- 安装 ModelService（安装 Python 运行时环境、安装模型服务包） +- 启动/停止 ModelService +- 监控 Agent 进程 + +## 架构概述（Local 模式） + +Local 模式下，模型服务使用**文件系统**作为通信媒介，实现代理和模型间的请求-响应机制。 + +当 Agent 需要调用模型时，请求首先写入日志文件，然后由负责监听的组件处理响应。当模型生成响应后，结果将写回日志文件，并由等待的 Agent 读取。 + +## anti_call_llm - 核心 API + +`anti_call_llm()` 是 **Local 模式**下最重要的 API，用于手动触发 LLM 反调用，实现模型调用的精细控制： + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='OpenAI type response', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +**使用场景：** +- Agent 捕获到 LLM 响应后，调用此方法通知 Roll 运行时 +- 支持携带响应数据，用于错误处理或重试 +- 超时和检查间隔可配置，适应不同网络环境 + +## CLI 命令 + +如果需要通过 CLI 使用模型服务，ROCK 提供了一个 CLI 命令集，可以在沙箱中安装 ROCK 后，通过 `rock model-service` 访问： + +### start 命令 +开始模型服务进程 +```bash +rock model-service start --type [local|proxy] [选项] +``` + +参数: + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `--type` | str | `local` | 服务类型：`local` 或 `proxy` | +| `--config-file` | str | None | 配置文件路径 | +| `--host` | str | None | 服务器地址（覆盖配置） | +| `--port` | int | None | 服务器端口（覆盖配置） | +| `--proxy-base-url` | str | None | 代理基础 URL | +| `--retryable-status-codes` | str | None | 可重试状态码，逗号分隔 | +| `--request-timeout` | int | None | 请求超时秒数 | + +### watch-agent 命令 +监控代理进程，当进程退出时发送 SESSION_END 消息 +```bash +rock model-service watch-agent --pid <进程ID> +``` + +参数: +- `--pid`: 需要监控的代理进程 ID + +### stop 命令 +停止模型服务 +```bash +rock model-service stop +``` + +### anti-call-llm 命令 +反调用 LLM 接口 +```bash +rock model-service anti-call-llm --index <索引> [--response <响应>] +``` + +参数: +- `--index`: 上一个 LLM 调用的索引，从 0 开始 +- `--response`: 上一次 LLM 调用的响应（可选） + +## 文件通信协议 + +模型服务使用文件进行进程间通信，定义了特定的标记格式用于区分请求和响应： + +### 请求格式 +``` +LLM_REQUEST_START{JSON请求数据}LLM_REQUEST_END{元数据JSON} +``` + +### 响应格式 +``` +LLM_RESPONSE_START{JSON响应数据}LLM_RESPONSE_END{元数据JSON} +``` + +### 会话结束标识 +``` +SESSION_END +``` + +元数据包含时间戳和索引信息，用于保证消息顺序和处理。 + +## SDK 使用 + +### ModelServiceConfig + +模型服务配置类，位于 `rock/sdk/sandbox/model_service/base.py`： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # 服务类型 + install_cmd="pip install rock-model-service", # 安装命令 + install_timeout=300, # 安装超时（秒） + start_cmd="rock model-service start --type ${type}", # 启动命令 + stop_cmd="rock model-service stop", # 停止命令 + logging_path="/data/logs", # 日志路径 + logging_file_name="model_service.log", # 日志文件名 +) +``` + +| 配置项 | 默认值 | 说明 | +|--------|--------|------| +| `enabled` | `False` | 是否启用模型服务（RockAgent 自动管理） | +| `type` | `"local"` | 服务类型：`local` 或 `proxy` | +| `install_cmd` | - | 模型服务包安装命令 | +| `install_timeout` | `300` | 安装超时时间（秒） | +| `start_cmd` | - | 启动命令模板 | +| `stop_cmd` | - | 停止命令 | +| `logging_path` | `/data/logs` | 日志目录路径 | +| `logging_file_name` | `model_service.log` | 日志文件名 | + +### ModelService + +模型服务管理类，处理沙箱内模型服务的生命周期： + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# 通常由 RockAgent 自动管理，无需手动调用 +# 以下方法仅在需要手动控制时使用 + +# 安装模型服务 +await model_service.install() + +# 启动模型服务 +await model_service.start() + +# 监控代理进程 +await model_service.watch_agent(pid="12345") + +# 执行反调用 LLM（Local 模式核心 API） +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# 停止模型服务 +await model_service.stop() +``` + +## API 参考 + +### install() + +在沙箱中安装模型服务依赖。 + +```python +await model_service.install() +``` + +执行步骤： +1. 创建并初始化 Python 运行时环境 +2. 创建 Rock 配置文件 +3. 安装模型服务包 + +**注意：** 通常由 RockAgent 自动调用。 + +### start() + +启动模型服务。 + +```python +await model_service.start() +``` + +前提条件：必须先调用 `install()`。 + +**注意：** 通常由 RockAgent 自动调用。 + +### stop() + +停止模型服务。 + +```python +await model_service.stop() +``` + +如果服务未运行，会跳过此操作。 + +**注意：** 通常由 RockAgent 自动调用。 + +### watch_agent(pid) + +监控代理进程。 + +```python +await model_service.watch_agent(pid="12345") +``` + +当进程退出时，发送 `SESSION_END` 消息。 + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +执行反调用 LLM 操作。**这是 Local 模式下最重要的 API。** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='{"result": "..."}', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +## 配置选项 + +### 服务配置 +- `SERVICE_HOST`: 服务主机地址，默认为 `"0.0.0.0"` +- `SERVICE_PORT`: 服务端口，默认为 `8080` + +### 日志配置 +- `LOG_FILE`: 用以通信的日志文件路径，包含请求和响应数据 + +### 轨迹（Traj）日志记录 +模型服务将 LLM 调用轨迹（traj）记录到 JSONL 文件中，用于调试和分析。 + +| 环境变量 | 默认值 | 说明 | +|----------|--------|------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | traj 日志文件目录 | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | 追加模式（true/false） | + +**traj 文件位置**: `{DATA_DIR}/LLMTraj.jsonl` + +**traj 文件格式**（JSONL - 每行一个 JSON 对象）： +```json +{"request": {...}, "response": {...}} +``` + +### 轮询配置 +- `POLLING_INTERVAL_SECONDS`: 轮询间隔，默认为 `0.1` 秒 +- `REQUEST_TIMEOUT`: 请求超时时间，默认为无限 + +### 标记配置 +定义了用于区分日志文件中不同类型消息的标记： +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig（服务端） + +服务端配置类定义了模型服务如何处理请求： + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `host` | str | `"0.0.0.0"` | 服务器地址 | +| `port` | int | `8080` | 服务器端口 | +| `proxy_base_url` | str \| None | `None` | 直接代理 URL | +| `proxy_rules` | dict | 见下方 | 模型名称到 URL 的映射 | +| `retryable_status_codes` | list[int] | `[429, 500]` | 可重试的 HTTP 状态码 | +| `request_timeout` | int | `120` | 请求超时时间（秒） | + +**默认 proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/python_sdk.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..c1083f29b0 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK 参考 + +本指南详细介绍如何使用 ROCK SDK 进行开发，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK SDK为开发者提供了便捷的Python接口来使用ROCK平台的功能，包括沙箱环境管理和GEM环境交互。 + +> **重要提示**: 使用 SDK 之前，请确保 ROCK Admin 服务正在运行。可以通过以下命令启动： +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 基本沙箱操作 + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 沙箱组管理 + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# 创建沙箱组配置 +config = SandboxGroupConfig( + image="python:3.11", + size=4, # 创建4个沙箱 + start_concurrency=2, # 并发启动级别为2 +) + +# 创建并启动沙箱组 +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# 批量操作 +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# 批量停止 +await sandbox_group.stop() +``` + +### 2.3 配置示例 + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 沙箱加速配置 + +ROCK 提供沙箱网络加速功能，支持配置 APT、PIP 和 GitHub 镜像源，提升受限网络环境下的包下载速度。 + +#### 支持的加速类型 + +**APT 镜像配置** + +配置 APT 包管理器镜像源，加速 Debian/Ubuntu 软件包下载。 + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# 配置 APT 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP 镜像配置** + +配置 Python 包索引镜像，加速 pip 安装。 + +```python +# HTTP 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub 加速** + +通过添加自定义 DNS 解析条目加速 GitHub 访问。 + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### 完整示例 + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """创建沙箱并配置加速""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # 配置加速（在安装包之前配置） + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # speedup 不会主动安装 PIP，仅配置镜像源进行加速 + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # 可以通过镜像 IP 加速 GitHub 访问 + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### 注意事项 + +1. **配置顺序**: 在安装包之前配置加速 +2. **HTTPS vs HTTP**: HTTPS 镜像不需要为 PIP 配置 trusted-host +3. **GitHub IP**: 不同区域可能需要不同的 IP 以获得最佳性能 +4. **持久性**: 配置在沙箱生命周期内持久有效 +5. **多次调用**: 后续的加速调用会覆盖之前的配置 +6. **PIP 安装**: speedup 功能仅配置镜像源，不会自动安装 PIP + +## 3. GEM SDK + +### 3.1 Python SDK 方式 + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## 相关文档 +- [快速开始指南](../../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK SDK +- [API 文档](../api.md) - 查看 SDK 封装的底层 API 接口 +- [配置指南](../../User%20Guides/configuration.md) - 了解 SDK 相关的配置选项 +- [安装指南](../../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/remote_user.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..791ca85fdd --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/remote_user.md @@ -0,0 +1,69 @@ +# Remote User + +远程用户管理，用于在沙箱中创建和管理用户。 + +## 使用示例 + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +创建远程用户。 + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +检查用户是否存在。 + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..c3f03b1efc --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent（实验性） + +RockAgent 是 ROCK 框架中的核心 Agent 实现，直接继承自 `Agent` 抽象基类。它提供了完整的 Agent 生命周期管理，包括环境初始化、ModelService 集成、命令执行等功能。 + +使用 `sandbox.agent.install()` 以及 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 + +## 核心概念 + +RockAgent 的核心工作流程分为两个阶段： + +1. **install(config)**: 初始化 Agent 环境，包括部署工作目录、设置环境变量、初始化运行时环境等 +2. **run(prompt)**: 执行 Agent 任务，替换占位符并启动 Agent 进程 + +## 快速开始 + +### Claude Code 示例 + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli 示例 + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} 必须 + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # 环境变量 + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent 示例 + +```yaml +working_dir: "." # 上传包含 langgraph_agent.py 的本地当前目录到 sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # 运行本地脚本 + +runtime_env_config: + type: python + pip: # 安装 pip 依赖 + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## 配置详解 + +### 基础配置 + +```yaml +agent_type: "default" # Agent 类型标识（默认: "default"） +agent_name: "demo-agent" # Agent 实例名称（默认: 随机 uuid） +version: "1.0.0" # 版本标识（默认: "default"） +instance_id: "instance-001" # 实例 ID（默认: "instance-id-<随机uuid>"） +agent_installed_dir: "/tmp/installed_agent" # Agent 安装目录（默认: "/tmp/installed_agent"） +agent_session: "my-session" # bash 会话标识（默认: "agent-session-<随机uuid>"） +env: # 环境变量（默认: {}） + OPENAI_API_KEY: "xxxxxxx" +``` + +### 工作目录配置 + +```yaml +working_dir: "./my_project" # 本地目录，上传到 sandbox（默认: None 不上传） +project_path: "/testbed" # sandbox 中工作目录，用于 cd（默认: None） +use_deploy_working_dir_as_fallback: true # project_path 为空时是否回退到 deploy.working_dir（默认: true） +``` + +### 执行配置 + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent 执行命令，必须包含 ${prompt}（默认: None） + +skip_wrap_run_cmd: false # 跳过为 run_cmd 添加 PATH 的包装（默认: false） + +# 超时配置 +agent_install_timeout: 600 # 安装超时，单位秒（默认: 600） +agent_run_timeout: 1800 # 运行超时，单位秒（默认: 1800） +agent_run_check_interval: 30 # 检查间隔，单位秒（默认: 30） +``` + +**`skip_wrap_run_cmd`**: +- `false`（默认）：为命令添加 `export PATH=:$PATH &&` 包装，确保使用运行时环境的可执行文件 +- `true`：跳过 PATH 包装，直接使用 `bash -c` 运行命令 + +### 初始化钩子 + +```yaml +pre_init_cmds: # 初始化前执行的命令（默认: 从 env_vars 读取） + - command: "apt update && apt install -y git" + timeout_seconds: 300 # 命令超时，单位秒（默认: 300） + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # 初始化后执行的命令（默认: []） + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**注意事项**： +- `pre_init_cmds` 和 `post_init_cmds` 不继承 Agent 的 `env` 环境变量 +- 通常用于执行安装操作和配置文件移动操作 +- 常用命令示例： + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv 配置 + +```yaml +runtime_env_config: # 具体参考 RuntimeEnv 有关文档 + type: "python" # 运行时类型: python / node（默认: "python"） + version: "3.11" # 版本号 + pip: # Python 依赖包列表 + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node 运行时示例**： + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**自动执行的操作**： +- 根据 `type` 安装对应的运行时（Python 或 Node.js） +- 安装 `pip` 依赖（如果配置了） +- 执行 `custom_install_cmd` 自定义安装命令（如果配置了） +- 支持 `npm_registry` 配置 Node.js 的 npm 镜像源 + +### ModelService 配置 + +```yaml +model_service_config: # 具体参考 ModelService 有关文档 + enabled: true # 启用 ModelService（默认: false） +``` + +**自动执行的操作**： +- 安装阶段：安装 ModelService（仅安装，不启动） +- 运行阶段：启动 ModelService + `watch_agent` 监控进程 + +**注意事项**：需要将模型请求的 URL 设置为 ModelService 的 URL。例如 ModelService 提供的 OpenAI-compatible 的 URL 为 `http://127.0.0.1:8080/v1/chat/completions`，则通常需要将 Agent 向 LLM 请求的 URL 设置为 `http://127.0.0.1:8080/v1/`。 + +## API 参考 + +### install(config) + +初始化 Agent 环境。 + +**执行流程**： +1. 如果配置了 `working_dir`，部署到 sandbox +2. 设置 bash session，以及配置 env 环境变量 +3. 执行 `pre_init_cmds` +4. 并行初始化 RuntimeEnv 和 ModelService（如果启用） +5. 执行 `post_init_cmds` + +**参数**： +- `config`: Agent 配置文件，支持两种传入方式： + - **字符串路径**: YAML 配置文件路径，默认值为 `"rock_agent_config.yaml"` + - **RockAgentConfig 对象**: 直接传入 `RockAgentConfig` 实例 + +### run(prompt) + +执行 Agent 任务。 + +**执行流程**： +1. 替换占位符, 准备Agent 运行命令 +4. 启动 agent 进程 +5. 如果启用 ModelService，启动 `watch_agent` +6. 等待任务完成并返回结果 + +## 高级用法 + +### working_dir 与 project_path 的区别与联动 + +| 配置项 | 作用 | 联动方式 | +|--------|------|----------| +| `working_dir` | 本地目录，上传到 sandbox | 调用 `deploy.deploy_working_dir()` 上传，上传后 `deploy.working_dir` 变为 sandbox 中的路径 | +| `${working_dir}` | 命令中的占位符 | 被 `deploy.format()` 替换为 `deploy.working_dir` 的值，会在配置中的 init_cmds 和 run_cmd 中替换 | +| `project_path` | sandbox 中的工作目录 | 用于运行前 `cd project_path`，不设置时会进入到 `deploy.working_dir` 工作目录 | +| `use_deploy_working_dir_as_fallback` | run 时 project_path 未设置时是否回退到 deploy.working_dir | 默认为 `true`，设为 `false` 时即使未设置 project_path 也不会进入 working_dir | + +**使用建议**： +- 使用 `working_dir` 上传本地项目代码到 sandbox +- 使用 `project_path` 指定 sandbox 中的工作目录（如 `/testbed`） +- 设置 `use_deploy_working_dir_as_fallback: false` 的场景：需要进行本地文件挂载，但希望在镜像默认工作目录下运行 Agent + +### 占位符使用 + +Rock Agent 在支持在配置文件中替换以下占位符： + +- `${prompt}`: 在run_cmd 中必需，会被替换为 `run(prompt)` 传入的提示词 +- `${working_dir}`: 可选，会被替换为 sandbox 中实际的工作目录路径, 同时支持在 init_cmds和 run_cmd 中使用 +- `${bin_dir}`: 可选，会被替换为运行时环境的 bin 目录路径 + +**示例**： +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback 说明 + +当 `project_path` 未设置时： +- `true`（默认）：运行 Agent 前会自动 `cd` 到 `deploy.working_dir` +- `false`：运行 Agent 前不会自动切换目录，保持在当前目录 + +适用场景： +- `true`: 大多数场景，希望 Agent 在上传的代码目录中运行 +- `false`: 需要挂载本地文件，但希望在镜像默认工作目录（如 `/app, /testbed`）下运行 Agent + +## 完整配置示例 + +```yaml +# ========== 基础配置 ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== 工作目录配置 ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== 运行配置 ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# 超时配置 +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== 初始化命令 ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== 运行时环境配置 ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService 集成 ========== +model_service_config: + enabled: true +``` + +## 使用示例 + +### 使用 YAML 配置文件（推荐） + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/runtime-env.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..a5532e900b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,137 @@ +# RuntimeEnv + +RuntimeEnv 模块用于在沙箱中管理语言运行时环境（目前提供了 Python / Node.js）。 + +## 快速开始(使用示例) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +异步工厂方法，根据配置创建 RuntimeEnv 实例并初始化，自动注册到 `sandbox.runtime_envs`。 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# 自动注册，可通过 sandbox.runtime_envs[env.runtime_env_id] 访问 +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +包装命令，将 `bin_dir` 加入 PATH，确保优先使用运行时环境中的可执行文件。 + +```python +wrapped = env.wrapped_cmd("node script.js") +# 返回: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +在运行时环境中执行命令。内部基于 `wrapped_cmd` 实现 + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | 类型标识 | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python 版本，默认 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | pip 包列表或 requirements.txt 路径 | +| `pip_index_url` | `str \| None` | 环境变量 | pip 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | 要创建符号链接的可执行文件列表 | + +## NodeRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | 类型标识 | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node 版本，默认 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | 要创建符号链接的可执行文件列表 | + +## 自定义 RuntimeEnv 实现约束 + +自定义 RuntimeEnv 需遵循以下规则： + +1. **定义 `runtime_env_type` 类属性**：作为类型标识符，用于自动注册到 RuntimeEnv 工厂 +2. **重写 `_get_install_cmd()`**：返回安装命令 +3. **安装命令最后必须**：将目录重命名为 `runtime-env` + + +## NodeRuntimeEnv 简化版实现示例 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config 类：定义配置类型，用于 RuntimeEnv.create() 路由到对应实现 +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # 必须与 runtime_env_type 一致 + +# RuntimeEnv 实现类：定义如何安装和运行该运行时环境 +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # 自动注册到 RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # 下载 Node 二进制包并解压，最后重命名为 runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## 加速基础环境安装 + +`PythonRuntimeEnv` 默认从 https://github.com/astral-sh/python-build-standalone/releases/ 下载 Python 安装包。若网络不可达或下载较慢，可通过环境变量 `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` 或 `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` 覆盖默认安装命令（例如切换到内网源/镜像源）。 + +默认值示例： + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +例如，替换为镜像源下载： + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +请确保该命令执行完成后，会在 `runtime_env` 的默认工作目录下生成 `runtime-env` 目录，并且 `${workdir}/runtime-env/bin/` 下包含对应可执行文件，例如： + +- `${workdir}/runtime-env/bin/python` + +Node 环境同理，可通过修改环境变量 `ROCK_RTENV_NODE_V22180_INSTALL_CMD` 来指定更快的下载/安装命令。 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/sandbox.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..088f1e3110 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/sandbox.md @@ -0,0 +1,113 @@ +# 处理大文件和长命令输出 + +## `arun` +`arun()` 在 `nohup` 模式下提供了两个关键参数，帮助 Agent / 调用方在"执行"与"查看"之间按需解耦： + +1. **`response_limited_bytes_in_nohup`**（int 型） + 限制返回内容的最大字符数（例如 `64 * 1024`），适合仍需立刻查看部分日志、但必须控制带宽的场景。默认值 `None` 表示不加限制。 + +2. **`ignore_output`**（bool，默认 `False`） + 当设为 `True` 时，`arun()` 不再读取 nohup 输出文件，而是在命令执行完毕后立即返回一段提示信息（包含输出文件路径、**文件大小**及查看方式）。日志仍写入 `/tmp/tmp_.out`，后续可通过 `read_file`、下载接口或自定义命令按需读取，实现"执行"与"查看"彻底解耦。返回的文件大小信息可帮助用户决定是直接下载还是分块读取。 + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# 示例 1：限制最多 1024 个字符 +resp_limit = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# 示例 2：完全跳过日志读取，后续再通过 read_file / 下载获取 +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# 可通过 Sandbox.read_file(...) / 下载接口 / cat /tmp/tmp_xxx.out 查看日志 +``` + +## `read_file_by_line_range` + +按行范围异步读取文件内容，支持自动分块读取和会话管理，支持大文件读取。 + +### 重要特性 +- **大文件分块读取**: 自动将大文件分成多个小块进行读取 +- **自动统计行数**: 未指定结束行时，自动计算文件总行数 +- **内置重试机制**: 关键操作支持最多 3 次重试，提高可靠性 +- **参数验证**: 自动验证输入参数的合法性 +- **会话管理**: 支持指定会话或自动创建临时会话 + +### 参数说明 +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `file_path` | str | - | 要读取的文件路径（沙箱中的绝对路径或相对路径） | +| `start_line` | int \| None | 1 | 起始行号（从 1 开始） | +| `end_line` | int \| None | None | 结束行号（包含），默认为文件末尾 | +| `lines_per_request` | int | 1000 | 每次请求读取的行数，范围 1-10000 | + +### 返回值 +- `ReadFileResponse`: 包含文件内容的响应对象 + - `content` (str): 读取的文件内容 + +### 异常说明 +- `Exception`: 当 `start_line < 1` 时抛出 +- `Exception`: 当 `end_line < start_line` 时抛出 +- `Exception`: 当 `lines_per_request` 不在 1-10000 范围内时抛出 +- `Exception`: 当文件读取失败时抛出 + +### 使用示例 + +```python +# 读取整个文件 +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# 读取指定行范围（第 100 到 500 行） +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# 从第 1990 行读取到文件末尾 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# 使用自定义分块大小 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### 注意事项 +- 行号从 1 开始计数，而非 0 +- 对于大文件建议适当增加 `lines_per_request` 以提高效率 +- 文件路径必须是沙箱内的有效路径 +- 使用 `sed` 命令进行文件读取，确保沙箱镜像支持该命令 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..6a74f397c7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,228 @@ +# SWE-Bench 评测 + +本文档介绍如何使用 ROCK SDK 运行 SWE-Bench Verified 评测，包括沙箱启动、Agent 集成、测试环境准备和结果解析。 + +### 快速开始 +SWE-Bench-Verified 是一个用于评估 AI 编程 Agent 在真实软件工程任务上表现的基准测试。 + +在ROCK上运行一个SWE-Bench任务包含以下步骤： + +1. **load_task_config** — 加载 `task.yaml` 获取任务指令 +2. **start_sandbox** — 使用任务专属的 Docker 镜像启动沙箱 +3. **agent.install / agent.run** — 安装并运行 Agent 来解决任务 +4. **setup_test_env** — 上传测试文件和运行测试脚本到沙箱 +5. **运行测试** — 通过 `sandbox.arun()` 执行测试脚本，支持超时控制 +6. **parse_swebench_result** — 解析测试输出，判断 PASSED / FAILED +7. **sandbox.stop** — 清理沙箱资源 + +**下面是示例代码** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. 加载任务指令 + task_config = await load_task_config(task_dir) # 参见 load_task_config 章节 + instruction = task_config["instruction"] + + # 2. 启动沙箱 + sandbox = await start_sandbox(task_name) # 参见 start_sandbox 章节 + + try: + # 3. 安装并运行 Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. 准备测试环境 + await setup_test_env(sandbox, task_dir) # 参见 setup_test_env 章节 + + # 5. 运行测试 + resp = await run_tests(sandbox) # 参见"运行测试"章节 + + # 6. 解析结果 + is_resolved = parse_swebench_result(resp.output) # 参见 parse_swebench_result 章节 + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +以下章节详细介绍评测流程中使用的各个函数。 + +--- + +## start_sandbox + +使用任务专属的 SWE-Bench Docker 镜像启动沙箱实例。每个任务都有一个预构建的镜像，包含目标仓库和运行环境。 + +`image` 参数格式如下： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +例如，任务 `django__django-14539` 对应的镜像为： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +从任务目录中加载 `task.yaml` 配置文件。YAML 文件包含 `instruction` 字段，用于描述 Agent 需要完成的编程任务。 + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# 使用示例 +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +使用 `sandbox.agent.install()` 和 `sandbox.agent.run()` 在沙箱中部署和执行 Agent。详细的 Agent 配置请参考 [Rock Agent](./rock-agent.md)。 + +```python +# 使用 YAML 配置文件安装 Agent（以 iflow_config.yaml 为例） +await sandbox.agent.install(config="iflow_config.yaml") + +# 使用任务指令运行 Agent +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +在沙箱中准备测试环境：安装 [uv](https://github.com/astral-sh/uv) 包管理器，并上传测试文件和运行测试脚本。 + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """准备测试环境并返回会话名称。""" + # 1. 创建带有自定义环境变量的会话 + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. 安装 uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. 上传测试文件 + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. 上传运行测试脚本 + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## 运行测试 + +使用 `RunMode.NOHUP` 模式执行测试脚本，支持可配置的超时时间。 + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +解析测试输出以判断 SWE-Bench 任务是否通过。解析器会查找由标记行分隔的结果块，并检查是否包含 `PASSED`。 + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """解析 SWE-Bench 测试输出，判断任务是否通过。 + + 匹配 'SWEBench results starts here' 和 + 'SWEBench results ends here' 之间的内容块， + 然后检查其中是否包含 'PASSED'。 + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# 使用示例 +is_resolved = parse_swebench_result(resp.output) +``` + +## 注意事项 + +- **任务数据集**：任务目录（包含 `task.yaml`、`tests/` 和 `run-tests.sh`）可从 [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) 仓库获取。 +- **任务镜像**：每个 SWE-Bench 任务需要特定的 Docker 镜像（如 `sweb.eval.x86_64.`）。请确保镜像在对应的环境中可用。 +- **Agent 配置**：Agent 配置 YAML 定义了运行时、依赖和执行命令。详情请参考 [Rock Agent](./rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/api.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/api.md new file mode 100644 index 0000000000..06f44c326c --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/References/api.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 1 +--- + +# API 参考 + +本指南详细介绍 ROCK 平台提供的核心 API 服务，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK平台提供两种核心API服务： +- Sandbox API：沙箱环境管理 +- GEM API：GEM环境交互 + +所有 API 接口都遵循 RESTful 设计原则，支持 JSON 格式的数据交换。 + +## 2. Sandbox API + +沙箱环境全生命周期管理功能： + +### 沙箱管理接口 + +1. **Start Sandbox** - 启动沙箱环境 + - 创建一个新的沙箱实例 + - 支持指定镜像、资源配置等参数 + +2. **Start Sandbox Async** - 异步启动沙箱环境 + - 异步方式创建沙箱实例 + - 适用于需要快速响应的场景 + +3. **Check Sandbox Alive Status** - 检查沙箱存活状态 + - 验证沙箱是否正常运行 + +4. **Get Sandbox Statistics** - 获取沙箱统计信息 + - 获取沙箱的资源使用统计 + +5. **Get Sandbox Status** - 获取沙箱详细状态 + - 获取沙箱的完整状态信息 + +6. **Stop Sandbox** - 停止沙箱环境 + - 安全关闭沙箱实例 + +7. **Commit Sandbox** - 提交沙箱为镜像 + - 将当前沙箱状态保存为新镜像 + +### 命令执行接口 + +8. **Execute Command** - 在沙箱中执行命令 + - 直接在沙箱中运行指定命令 + +9. **Create Bash Session** - 创建Bash会话 + - 创建持久化的Bash会话环境 + +10. **Run Command in Session** - 在会话中执行命令 + - 在已创建的会话中执行命令 + +11. **Close Session** - 关闭会话 + - 释放会话资源 + +### 文件操作接口 + +12. **Read File** - 读取沙箱文件 + - 从沙箱中读取指定文件内容 + +13. **Write File** - 写入沙箱文件 + - 向沙箱中写入文件 + +14. **Upload File** - 上传文件到沙箱 + - 将本地文件上传到沙箱 + +## 3. GEM API + +GEM环境交互功能： + +1. **Make Environment** - 创建GEM环境 + - 初始化一个新的GEM环境实例 + +2. **Reset Environment** - 重置GEM环境 + - 将GEM环境重置到初始状态 + +3. **Step Environment** - 执行GEM环境步骤 + - 在GEM环境中执行一个动作步骤 + +4. **Close Environment** - 关闭GEM环境 + - 释放GEM环境资源 + +## 4. HTTP API 使用示例 + +### 4.1 Sandbox API 示例 + +#### 启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 异步启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### 创建会话 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### 在会话中执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### 上传文件 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### 停止沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API 示例 + +```bash +# 创建GEM环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# 重置环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# 执行步骤 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# 关闭环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK API +- [Python SDK 文档](./Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 调用 API +- [配置指南](../User%20Guides/configuration.md) - 了解 API 相关的配置选项 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md new file mode 100644 index 0000000000..3343474e09 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# 版本说明 +* [release v1.6.0](v1.6.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.0.md new file mode 100644 index 0000000000..6eacdbc274 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.0.md @@ -0,0 +1,111 @@ +# v1.6.0 + +## 发布日期 + +2026 年 4 月 17 日 + +--- + +## 亮点 + +本次发布的核心是 **Job 模块** 的重大重构，引入清晰的 `Job` / `Operator` / `Executor` / `Trial` 分层抽象。在此基础上新增了两类作业类型 — **BashJob**（含 OSS 产物镜像）与 **HarborJob**。CLI `rock job run` 命令也已重写，支持严格 YAML 校验与作业类型自动识别。 + +--- + +## Job 模块 + +### 新架构 + +#### Job / Operator / Executor / Trial 抽象 + +* **新增**: 将 Job 模块重构为分层架构：`Job` → `Operator` → `Executor` → `Trial`。该结构解耦了作业编排、调度、执行与单次尝试逻辑，便于扩展新的作业类型 ([#779](https://github.com/alibaba/ROCK/pull/779), [#780](https://github.com/alibaba/ROCK/pull/780)) + +* 将 `on_sandbox_ready` backfill 逻辑上提至 `AbstractTrial`，便于各类 Trial 共享 ([#788](https://github.com/alibaba/ROCK/pull/788), [#789](https://github.com/alibaba/ROCK/pull/789)) + +#### BashJob + +* **新增**: BashJob trial 支持 — 通过 SDK 或 CLI 提交 shell 脚本作业。参考 `examples/bash/simple_bash_job_demo.sh` ([#772](https://github.com/alibaba/ROCK/pull/772)) + +* **新增**: BashJob OSS Mirror — 作业完成后自动上传产物到 OSS ([#823](https://github.com/alibaba/ROCK/pull/823)) + +* 新增 `claw-eval` BashJob 示例，位于 `examples/evaluation/claw_eval/` ([#804](https://github.com/alibaba/ROCK/pull/804)) + +#### HarborJob + +* **新增**: HarborJob trial 支持，基于新的 Job 抽象提交 agent 风格的作业 ([#798](https://github.com/alibaba/ROCK/pull/798)) + +### 配置与校验 + +#### 作业类型自动识别 + +* **新增**: `rock job run` 现在通过严格的 Pydantic 模型校验从 YAML 自动识别作业类型 — 无需指定 `--type` ([#814](https://github.com/alibaba/ROCK/pull/814)) + +#### 双模式输入 + +* **新增**: `rock job run` 重写，同时支持 `--config`（完整 YAML）与命令行参数两种模式，校验更严格、错误信息更清晰 ([#818](https://github.com/alibaba/ROCK/pull/818)) + +#### Native 模板配置 + +* **新增**: 在 `NativeConfig` 中新增 `TemplateConfig` 与 `template` 字段，支持基于模板的作业定义 ([#786](https://github.com/alibaba/ROCK/pull/786)) + +#### 自动生成的 job_name + +* 自动生成 `job_name` 时对过长的路径片段进行截断，避免触发下游长度限制 ([#791](https://github.com/alibaba/ROCK/pull/791)) + +#### 默认超时延长 + +* `JobConfig` 默认超时从 3600s 延长至 **7200s**，更贴合真实长时作业场景 ([#806](https://github.com/alibaba/ROCK/pull/806), [#810](https://github.com/alibaba/ROCK/pull/810)) + +### Bug 修复 + +* `JobConfig.experiment_id` 现在优先于 `environment.experiment_id`，确保调用方拥有显式控制权 ([#822](https://github.com/alibaba/ROCK/pull/822)) + +* 修复 `BashTrial.collect` 未正确填充 `raw_output` 与 `exit_code` 的问题 ([#808](https://github.com/alibaba/ROCK/pull/808)) + +* 移除 `BashTrial.build()` 中冗余的 `shebang` 与 `set -e` 注入，尊重用户脚本 ([#816](https://github.com/alibaba/ROCK/pull/816)) + +--- + +## EnvHub + +### 重构 + +* **破坏性变更**: `JobEnvironmentConfig` 已迁移至 `envhub`，更名为 `EnvironmentConfig`。请相应更新 import 路径 ([#800](https://github.com/alibaba/ROCK/pull/800)) + +* 移除 `EnvironmentConfig` 中已弃用的 `auto_stop` 参数 — sandbox 生命周期改由 `auto_delete_seconds`（v1.5.0 引入）控制 ([#820](https://github.com/alibaba/ROCK/pull/820)) + +* 重构 EnvHub 上传链路，明确客户端与服务端职责边界 ([#802](https://github.com/alibaba/ROCK/pull/802)) + +--- + +## Admin + +### 数据库 + +* `SandboxRecord.image` 列长度从 255 扩展至 **512 字符**，支持更长的镜像路径；同时禁用 asyncpg 的 prepared-statement 缓存以避免 PgBouncer 兼容性问题 ([#794](https://github.com/alibaba/ROCK/pull/794)) + +--- + +## CLI + +* `admin stop` 命令中懒加载 `psutil`，加快无关 CLI 调用的启动速度，并避免 `psutil` 不可用时的导入期失败 ([#831](https://github.com/alibaba/ROCK/pull/831)) + +--- + +## 测试与 CI + +* 在 `pg_container` 测试 fixture 中加入 `SELECT 1` readiness 探测，避免集成测试 flaky ([#778](https://github.com/alibaba/ROCK/pull/778)) + +* 全仓应用 `ruff format`，并将剩余中文注释翻译为英文 ([#812](https://github.com/alibaba/ROCK/pull/812)) + +--- + +## 迁移说明 + +* **`auto_stop` 已移除**: 请更新依赖 `auto_stop` 的 `EnvironmentConfig` 用法，改用 `auto_delete_seconds`（v1.5.0 引入）。 + +* **`JobEnvironmentConfig` → `EnvironmentConfig`**: 将 `rock.sdk.agent.models.job` 的 import 改为 `rock.sdk.envhub`。 + +* **默认 Job 超时**: 之前会触发 3600s 超时的作业，现在默认可运行至 7200s。如需保持原有行为，请显式设置 `JobConfig.timeout`。 + +* **`rock job run`**: CLI 现在严格校验 YAML，先前能"静默通过"的未知字段或缺失字段配置会立即失败。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/User Guides/configuration.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/User Guides/configuration.md new file mode 100644 index 0000000000..a212189bc7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/User Guides/configuration.md @@ -0,0 +1,188 @@ +--- +sidebar_position: 4 +--- + +# 配置指南 + +本指南详细介绍如何配置 ROCK 环境以满足不同的使用需求，包括本地开发、测试和生产部署。 + +## 1. 环境变量配置 + +ROCK 支持通过环境变量配置关键参数。以下是主要的环境变量： + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK服务基础URL +export ROCK_LOG_LEVEL=INFO # 日志级别 +export ROCK_LOGGING_PATH=/path/to/logs # 日志文件路径，默认 None (输出到控制台) +export ROCK_LOGGING_FILE_NAME=rocklet.log # 日志文件名，默认 "rocklet.log", 启动admin时可以自定义日志文件名, 如admin.log +export ROCK_LOGGING_LEVEL=INFO # 日志输出级别，默认 "INFO" +export ROCK_WORKER_ENV_TYPE=local # 运行时环境类型，可选值: local, docker, uv, pip +``` + +更多环境变量可参考 `rock/env_vars.py` 文件。 + +### 1.1 运行时环境 (Runtime Environments) + +ROCK 提供了多种不同的运行时环境来满足不同场景的需求，选择通过环境变量 `ROCK_WORKER_ENV_TYPE` 进行配置。每种环境有不同的部署要求、性能特征和适用场景。每种环境都有其独特的优势和限制，开发者可以根据部署环境的需要选择最适合的运行时环境。 + +#### 1.1.1 Docker 运行时环境 + +Docker 运行时环境适用于已经预安装了所需依赖的 Docker 镜像环境。这种环境要求部署环境中直接可用 `/tmp/miniforge/bin/rocklet` 可执行文件。 + +**挂载配置:** +- `/tmp/miniforge` - 包含预安装的 Python 环境 +- `/tmp/local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 容器化部署环境 +- 已经构建了包含 `rocklet` 的自定义 Docker 镜像 +- 适合生产环境，启动速度快 + +**要求:** +- 需要使用定制的 Docker 镜像，其中包含 `/tmp/miniforge/bin/rocklet` 可执行文件 +- Docker 环境支持 + +#### 1.1.2 本地运行时环境 + +本地运行时环境直接利用当前部署环境的 Python 环境和项目文件。该环境要求宿主机和容器之间具有相同的操作系统，以便能够直接挂载虚拟环境和 Python 解释器。 + +**挂载配置:** +- `python_env_path` - Python 环境路径 +- `project_root` - 项目根目录 +- `.venv` - 虚拟环境目录（挂载为容器中的 `/tmp/miniforge`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 开发环境 +- 宿主机和目标容器使用相同操作系统的场景 +- 需要快速重新使用现有 Python 环境 + +**要求:** +- 相同的操作系统（主机/容器） +- 可直接访问当前部署的 `.venv` 虚拟环境 +- Python 解释器路径兼容 + +#### 1.1.3 UV 运行时环境 + +UV 运行时环境只依赖于可用的 ROCK 项目，但初始化相对较慢且网络要求较高。这种环境最适合没有预配置环境的场景。它从原始项目重新构建 rocklet 环境。这是推荐在 Mac 操作系统上使用的环境。 + +**挂载配置:** +- `project_root` - 项目根目录（挂载为容器中的 `/tmp + project_root`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**适用场景:** +- Mac 操作系统 +- 跨操作系统启动 +- 没有预配置环境的场景 +- 没有使用 uv 管理 Rock + +**优势:** +- 无需预构建镜像 +- 跨平台兼容性好 +- 特别适合开发和测试 + +**限制:** +- 初始化速度较慢 +- 网络要求较高 +- 启动时间较长 + +#### 1.1.4 PIP 运行时环境 + +PIP 运行时环境使用 pip 在容器内安装所需依赖。这种环境适合快速设置并能在容器中完成依赖安装的场景，是默认的运行时环境。它不需要预先构建包含依赖的镜像，通过 pip 直接管理 Python 包。 + +**挂载配置:** +- `local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**适用场景:** +- 使用PIP源安装的ROCK +- 快速测试ROCK + +**优势:** +- 简单的部署设置 + +**限制:** +- 依赖安装时间较长 +- 需要网络访问以安装依赖包 +- 每次启动时都需要安装依赖 + +#### 1.1.5 配置指南 + +根据不同的使用场景，可以参考以下选择指南： + +| 场景 | 推荐环境 | 原因 | +|------|----------|------| +| 生产环境 | Docker 运行时 | 快速启动，稳定性能 | +| 开发环境，同一 OS | 本地运行时 | 环境重用，开发周期快 | +| Mac 开发 | UV 运行时 | 支持最佳的跨平台兼容性 | +| 跨平台开发 | UV 运行时 | 避免环境兼容性问题 | +| 快速测试 | UV 运行时 | 无需预配置工作 | +| PIP源安装 | PIP 运行时 | 直接使用 pip 安装依赖 | + +这些运行时环境通过 `ROCK_WORKER_ENV_TYPE` 环境变量进行配置，该变量可设置为 "local"、"docker"、"uv" 或 "pip"。 + +### 1.2 日志配置 + +在日志配置方面，ROCK 的日志系统具有以下特性： + +- 日志系统不能同时输出到文件和控制台，只有当设置了 `ROCK_LOGGING_PATH` 时，日志才会输出到指定文件，否则输出到控制台。 +- `ROCK_LOGGING_LEVEL` 用于控制日志输出级别，`ROCK_LOG_LEVEL` 用于通用日志级别设置。 + +## 2. 分布式部署要求 + +由于 ROCK 支持分布式部署，当在 Ray 集群的不同节点上运行时，需要满足以下一致性要求： + +#### 目录结构一致性 +在所有 Ray 节点上，必须保证以下目录结构完全一致： +- ROCK 项目仓库目录 +- `.venv` 虚拟环境目录 +- `.venv` 依赖的 base Python 目录 + + +#### 挂载要求 +ROCK 的启动依赖于挂载 ROCK 项目和对应的 base Python 环境，要求在多机环境中保持一致性： + +#### 验证分布式配置 +可以通过以下方式验证分布式部署配置： + +```bash +# 在所有节点上检查目录一致性 +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# 验证 Python 环境可用性 +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# 检查所有节点上的环境变量设置 +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + + + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速搭建 ROCK 环境 +- [API 文档](../References/api.md) - 查看沙箱相关的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 配置沙箱 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/overview.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/overview.md new file mode 100644 index 0000000000..a02536f50d --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/overview.md @@ -0,0 +1,40 @@ +--- +sidebar_position: 1 +--- + +# 概览 + +ROCK (Reinforcement Open Construction Kit) 是一个开源的强化学习环境开发框架，旨在简化强化学习环境的开发、部署和管理流程。 + +## 什么是 ROCK + +ROCK (Reinforcement Open Construction Kit) 是一个开源强化学习环境开发框架。通过使用 ROCK，开发者可以快速地开发强化学习环境，并结合其他强化学习训练框架，实现高效的强化学习训练。 + +ROCK 提供了完整的沙箱环境管理功能，支持容器化部署，能够实现环境的快速创建、运行和销毁。同时，ROCK 兼容 GEM 协议，为强化学习环境提供了标准化的接口。 + +## ROCK 的核心功能 + +1. **简化开发流程**：简化强化学习环境的开发、构建和管理流程，支持多种开源的强化学习环境 +2. **大规模调度部署**：支持快速强化学习环境的大规模调度部署，通过 GEM 协议可以方便地访问强化学习环境 +3. **框架集成**：与其他强化学习训练框架集成，实现大规模可扩展的强化学习训练 + +## ROCK 的价值 + +ROCK 为不同角色的工程师提供了显著价值： + +- **强化学习算法工程师**：ROCK 可以简化强化学习环境的开发流程，让工程师专注于算法实现 +- **强化学习应用工程师**：ROCK 可以进行快速强化学习环境的大规模部署，提高应用开发效率 + +## 相关文档 + +如果您是第一次使用 ROCK，建议按以下顺序阅读文档： +1. [快速开始指南](./Getting%20Started/quickstart.md) - 快速搭建开发环境 +2. [配置指南](./User%20Guides/configuration.md) - 配置您的 ROCK 环境 +3. [Python SDK 文档](./References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +4. [API 文档](./References/api.md) - 了解完整的 API 接口 +5. [安装指南](./Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 + + + + + diff --git a/docs/versioned_docs/version-1.6.x/Getting Started/installation.md b/docs/versioned_docs/version-1.6.x/Getting Started/installation.md new file mode 100644 index 0000000000..c45b09a8fe --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Getting Started/installation.md @@ -0,0 +1,143 @@ +--- +sidebar_position: 3 +--- + +# Installation + +This document explains how to install and set up the ROCK development environment using both `uv` and `pip`. The project is a Reinforcement Open Construction Kit that supports various components. + +## Using uv (Recommended) + +### Quick Install All Dependencies + +```bash +# Install all dependencies including optional ones +uv sync --all-extras + +# Install development/testing dependencies +uv sync --all-extras --all-groups +``` + +### Install Different Dependency Groups + +#### Core Dependencies Only +```bash +uv sync +``` + +#### Admin Component Dependencies +```bash +uv sync --extra admin +``` + +#### Rocklet Execution Environment Dependencies +```bash +uv sync --extra rocklet +``` + + +#### All Dependencies at Once +```bash +uv sync --all-extras +``` + +#### Development/Testing Dependencies +```bash +uv sync --all-extras --group test +``` + +## Using pip + +### Install from pip source + +#### Core Dependencies Only +```bash +pip install rl-rock +``` + +#### Admin Component Dependencies +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install "rl-rock[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install "rl-rock[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install "rl-rock[all]" +``` + +### Install with pip from source code + +#### Core Dependencies Only +```bash +pip install . +``` + +#### Admin Component Dependencies +```bash +pip install ".[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install ".[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install ".[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install ".[all]" +``` + +## Available Entry Points + +The package provides the following command line scripts: + +- `rocklet`: ROCK execution environment server (rock.rocklet.server:main) +- `admin`: Admin management server (rock.admin.main:main) +- `envhub`: Environment hub server (rock.envhub.server:main) +- `rock`: Main ROCK command line interface (rock.cli.main:main) + +## Development Setup + +### Using uv (Recommended) + +```bash +# Clone and set up development environment +git clone +cd ROCK +uv sync --all-extras --group test + +# Run tests +uv run pytest + +``` + +### Using pip + +```bash +# For development, install in editable mode with all extras +pip install -e ".[all]" + +# Or separately +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" # Optional extras +``` + +## Additional Notes + +- The project is configured to use the Alibaba cloud PyPI mirror by default: `https://mirrors.aliyun.com/pypi/simple/` +- For local development, running tests requires the `test` dependency group diff --git a/docs/versioned_docs/version-1.6.x/Getting Started/quickstart.md b/docs/versioned_docs/version-1.6.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..2f14808f5b --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Getting Started/quickstart.md @@ -0,0 +1,166 @@ +--- +sidebar_position: 2 +--- + +# Getting Started + +This guide will demonstrate how to use ROCK to create and manage reinforcement learning environments through complete examples. + +## 1. Environment Preparation + +We recommend starting ROCK on Linux systems to maximize dependency reuse and improve environment startup speed. If you need to try on macOS, please refer to the [MacOS Startup](#7-macos-startup) section. + +Before starting, please ensure your system has the following dependencies installed: + +### 1.1 System Requirements + +- **Docker**: ROCK uses Docker for containerized environment management +- **uv**: ROCK uses uv for dependency management and virtual environment creation + +### 1.2 Verify Dependency Installation + +```bash +# Verify Docker installation +docker --version + +# Verify Docker image, and example depends on python:3.11 image +docker pull python:3.11 + +# Verify uv installation +uv --version +``` + +### 1.3 Project Initialization + +```bash +# Clone repository +git clone +cd ROCK + +# Create virtual environment (using uv-managed Python, use python 3.11 as an example) +uv venv --python 3.11 --python-preference only-managed + +# Install all dependency groups +uv sync --all-extras +``` + +> **Important Note**: To ensure ROCK can correctly mount the project and virtual environment along with its base Python interpreter, it is strongly recommended to use uv-managed Python environments to create virtual environments rather than system Python. + +## 2. Activate Virtual Environment + +Before running any ROCK commands, you need to activate the virtual environment. Ensure sys.base_prefix is a uv-managed environment, such as `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` or similar paths. + +```bash +# Activate virtual environment +source .venv/bin/activate + +# Verify Python environment +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **Verification Point**: Ensure the output base prefix path points to a uv-managed Python environment, not system Python. + +## 3. Verify Environment Configuration + +After activating the virtual environment, verify that dependencies are installed correctly: + +```bash +# Check key dependencies +python -c "import rock; print(\"Hello ROCK\")" +``` + +## 4. Start ROCK Service + +After activating the virtual environment, start the ROCK Admin service on project root: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Start ROCK Admin service (local environment) +rock admin start +``` + +After the service starts, you will see output similar to the following: + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **Service Information**: The ROCK Admin service runs by default on `http://127.0.0.1:8080`. + +## 5. Run Example Environments + +Now you can run example environments to verify the installation. Ensure the ROCK service is running, then open a new terminal window to execute the following commands: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Run sandbox example +python examples/sandbox_demo.py + +# Run GEM protocol example +python examples/sokoban_demo.py +``` + +### 5.1 Example Descriptions + +- **sandbox_demo.py**: Demonstrates how to use ROCK's sandbox SDK to create and manage containerized environments +- **sokoban_demo.py**: Demonstrates how to use ROCK's GEM protocol compatible interface to create reinforcement learning environments + +> **Running Requirements**: Ensure the ROCK Admin service is running, as examples need to communicate with the service. + +## 6. Distributed Environment Configuration (Optional) + +For distributed multi-machine environments, ensure the following configurations are consistent: + +1. All machines use the same root Python interpreter for ROCK and uv Python configurations +2. Docker versions are consistent across all nodes +3. Network configuration allows normal communication between nodes + + +## 7. MacOS Startup + +On macOS, if you need to start Linux image environments, you first need to set the environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +During container startup, the corresponding uv environment will be installed. For details, please refer to the `rock/rocklet/local_files/docker_run_with_uv.sh` script. + +> **Note**: Compared to Linux systems, the startup speed on macOS will be slower and more dependent on network conditions. You can adjust the script according to actual conditions.You can find detatils for ROCK_WORKER_ENV_TYPE in [Configuration Guide](../User%20Guides/configuration.md). + +## 8. Starting from Pip Source + +If starting the Admin Server from Pip source, after completing the ROCK installation by referring to [installation](./installation.md), you need to set an additional environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(This startup method will pull and install the latest rocklet from the PyPI source when starting the container environment. The startup speed is relatively slow, so it is only recommended for testing purposes. For production environments, other startup methods are still recommended.) + +## Summary + +Congratulations! You have successfully completed the ROCK quick start guide. You should now be able to: + +- Properly set up the ROCK development environment +- Use uv-managed Python environments +- Start and manage ROCK services +- Run example programs to verify installation +- Configure ROCK in distributed environments (if needed) + +For a deeper understanding of ROCK's additional features, please refer to the following documents: + +## Next Steps + +- [Configuration Guide](../User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](../References/api.md) - View complete API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the Python SDK for development +- [Installation Guide](./installation.md) - Detailed information about ROCK installation and setup +- [Overview](../overview.md) - Understand ROCK's design philosophy \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/Getting Started/rock-agent.md b/docs/versioned_docs/version-1.6.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..eb6b54a65b --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Getting Started/rock-agent.md @@ -0,0 +1,72 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent Quick Start + +Rock Agent is an AI Agent runtime framework provided by ROCK, supporting various types of Agents running in sandbox environments. + +## Prerequisites +- Make sure you have a working ROCK service, if you need to locally start the service side, refer to [Quick Start](quickstart.md). + +## Examples + +ROCK provides two Hello World Agent examples in the `examples/agents/` directory: + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent example +└── iflow_cli/ # IFlowCli Agent example +``` + +### Run IFlowCli Example + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### Run ClaudeCode Example + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli Configuration File + +The configuration file is located at `examples/agents/iflow_cli/rock_agent_config.yaml`: + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # Enter your API key + IFLOW_BASE_URL: "" # Enter your base URL + IFLOW_MODEL_NAME: "" # Enter your model name +``` + +## ClaudeCode Configuration File + +The configuration file is located at `examples/agents/claude_code/rock_agent_config.yaml`: + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # Enter your anthropic base url + ANTHROPIC_API_KEY: "" # Enter your anthropic api key +``` + +## Related Documentation + +- [RockAgent Reference](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/versioned_docs/version-1.6.x/Getting Started/rockroll.md b/docs/versioned_docs/version-1.6.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..2465a7733f --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Getting Started/rockroll.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL Quick Start Guide + +This guide will walk you through running a reinforcement learning training example based on the Sokoban game, using ROLL (the training framework) and ROCK (the environment management tool). + +## 1. Prerequisites + +Before you begin, please ensure your system has the following dependencies installed. + +### 1.1 System Requirements + +- **OS**: A Linux-based system is recommended (e.g., Ubuntu 20.04+). +- **Hardware**: An NVIDIA GPU with the corresponding drivers is recommended. +- **Docker**: ROCK uses Docker for containerized environment management. +- **uv**: ROCK uses uv for dependency management and virtual environment creation. + +### 1.2 Verify Dependencies & Pre-pull Image + +```bash +# Verify Docker installation +docker --version + +# Verify Docker is running and pre-pull the Sokoban environment image +# This will save time when the training starts. +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# Verify uv installation +uv --version + +``` + +### 1.3 Initialize the Project + +```bash +# Clone the project repositories +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# Ensure both repositories are in the same parent directory, like this: +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. Launch the Training Process + +> Note: The following instructions use torch==2.6.0 and vLLM==0.8.4 as an example. + + +### Option 1: Using a Virtual Environment (Recommended) + +#### Why is this method recommended? +- Isolation: A uv virtual environment ensures that project dependencies are isolated from your system, preventing conflicts. +- Fast Startup: ROCK can reuse this virtual environment, significantly speeding up subsequent task initializations. +- Stability & Reproducibility: Dependency management is cleaner and more reliable. + + +```bash +# Navigate to the ROCK directory +cd ROCK + +# Create and activate a Python 3.10 virtual environment (ROLL recommends Python 3.10) +uv venv --python 3.10 --python-preference only-managed +source .venv/bin/activate + +# Install all of ROCK's dependencies using uv +uv sync --all-extras + +# If using Python 3.10, starting Ray may raise a `ValueError: is not a valid Sentinel`. +# This is due to an incompatibility between `ray` and `click` versions 8.3+. +# To fix this, downgrade `click` to a version below 8.3. This issue does not affect Python 3.11. +uv pip install 'click>=8.2,click<8.3' + +# Navigate to the ROLL directory to install its dependencies +cd ../ROLL + +# Install core PyTorch components +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 + +# Install transformer-engine. The --no-build-isolation flag prevents errors where torch cannot be found. +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation + +# Install a pre-compiled version of flash-attention matching the specific CUDA and PyTorch versions +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# Install the remaining dependencies +uv pip install -r requirements_torch260_vllm.txt + +# (Optional) Install Tensorboard to check training metrics +uv pip install tensorboard -i $PYPI_MIRROR + +# All set! Launch the training script. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### Option 2: Using the System Environment (Alternative) + +For optimal compatibility with this method, we recommend running these commands inside one of ROLL's official base Docker images. These images come pre-installed with matching CUDA, cuDNN, and other foundational libraries. + +> [ROLL's Official Docker Image List](https://alibaba.github.io/ROLL/docs/Getting%20Started/Installation/image_address) + + +#### Warning +This method will install all Python packages directly into your current environment (e.g., the container's base system), which may cause conflicts with system packages or other projects. + +Since ROCK cannot reuse the environment, it may need to reinstall some dependencies each time a task starts, leading to slower startup times that are dependent on network speed. + + +```bash +# Install ROCK's dependencies +cd ROCK +pip install . +pip install ".[admin]" + +# Install ROLL's dependencies +cd ../ROLL +pip install -r requirements_torch260_vllm.txt + +# Crucial: Configure ROCK to use uv as its worker environment manager +export ROCK_WORKER_ENV_TYPE=uv + +# Launch the training script +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +You have now successfully launched the Sokoban reinforcement learning training process. Happy Rock & Roll！ + + +## 3. Multi-Node Deployment + +Instead of running everything on a single machine, you can deploy the **ROCK Service** and **ROLL job** on separate machines. This is a common client-server setup where they communicate over the network. + +### 3.1 Deploy the ROCK Service on Machine A + +On a dedicated machine (or container), follow the [ROCK Quick Start Guide](./quickstart.md) to deploy and start the ROCK service. + +> **Important** +> After starting the service, take note of its IP address and port (e.g., `http://192.168.1.10:8000`). You will need this address for the subsequent steps. + +### 3.2 Prepare the ROLL Client on Machine B + +On the other machine where you will run the training task, perform the following steps. + +1. Verify Network Connectivity + +First, use the curl command to check if you can reach the ROCK service on Machine A from Machine B. +```bash +# Replace : with the actual address of your ROCK service +# If successful, you should receive a response like {"message":"hello, ROCK!"} +curl http://: +``` + +2. Prepare the ROLL Environment + +```bash +# Clone the ROLL repository +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# Install dependencies +pip install -r requirements_torch260_vllm.txt +``` + +3. Configure the ROLL Connection Address + +Modify ROLL's configuration file to point to the remote ROCK service. +- Open the configuration file: examples/agentic_demo/agentic_val_sokoban_sandbox.yaml. +- Find the "SokobanSandbox" section under "env_config". +- Update the base_url value to your ROCK service's address. +```yaml +custom_envs: + SokobanSandbox: + env_config: + # Change the address here to your ROCK service's address + # Example: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. Start Training +Once configured, you can start the ROLL training script on Machine B. + +```bash +# This script will now request environments from the ROCK service on Machine A over the network. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### Advanced: Distributed ROLL Training + +If you wish to deploy the ROLL training task itself in a distributed manner, you can refer to ROLL's official documentation for distributed deployment. +> [Quick Start: Multi-Node Deployment Guide](https://alibaba.github.io/ROLL/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/codes.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..dceb8d3182 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +Error code definitions and categories for error handling and retry strategies. + +## Usage Example + +```python +import rock + +def test_codes_values(): + """Test basic status code values""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes Categories + +```python +OK = 2000, "OK" +""" +Success codes (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +Client error codes (4xxx): + +These errors indicate issues with the client request, +SDK will raise Exceptions for these errors. +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +Server error codes (5xxx): + +These errors indicate issues on the server side, +SDK will raise Exceptions for these errors. +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +Command/execution error codes (6xxx): + +These errors are related to command execution and should be handled by the model, +SDK will NOT raise Exceptions for these errors. +""" +``` + +## Retry Strategy Recommendations + +- **Retry trigger**: Only retry when `INTERNAL_SERVER_ERROR` occurs +- **Other error handling**: + - `BAD_REQUEST`: Check if there are issues with the arun call logic + - `COMMAND_ERROR`: stdout goes to `observation.output`, stderr goes to `observation.failure_reason` +- `COMMAND_ERROR` note: When bash execution fails, both stdout and stderr may be non-empty. It is recommended to prompt the model with both output and failure_reason from the observation. + +## Retry Example + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/deploy.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..5fd5b70546 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +Sandbox resource deployment manager for local directory deployment and template formatting. + +## deploy_working_dir - Deploy Local Directory + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# Deploy local directory (auto-generated target path) +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"Deployed to: {target}") # e.g., /tmp/rock_workdir_abc123 + +# Deploy to specific target path +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - Template Variable Substitution + +The `format` method supports two template syntaxes: + +- **`${variable}`** - Standard Python string template syntax +- **`<>`** - Alternative syntax (converted to `${variable}` internally) + +```python +# After deploy_working_dir, use ${working_dir} placeholder +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# Result: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# Alternative <<>> syntax +cmd = deploy.format("cat <>/file.txt") +# Result: cat /tmp/rock_workdir_abc123/file.txt + +# Combine with custom variables +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# Result: cat /tmp/rock_workdir_abc123/settings.json + +# Shell syntax is preserved +cmd = deploy.format("echo $((3 << 2 >> 1))") +# Result: echo $((3 << 2 >> 1)) + +# Access working_dir directly +if deploy.working_dir: + print(f"Current working directory: {deploy.working_dir}") +``` + +## Multiple Deployments + +Subsequent calls overwrite previous working directory paths: + +```python +# First deployment +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# Second deployment (overwrites previous path) +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/file_system.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..a64228e8f1 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +File system interface for sandbox environment operations including permission and ownership management. + +## chown - Change Owner + +```python +from rock.actions.sandbox.request import ChownRequest + +# Create remote user before changing ownership +await sandbox.remote_user.create_remote_user("deploy") + +# Get current working directory +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# Change directory owner +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# Recursively change owner for directory and contents +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - Change Permissions + +```python +from rock.actions.sandbox.request import ChmodRequest + +# Create test directory +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# Change directory permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# Recursively change permissions (includes subdirectories and files) +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# Set maximum permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - Upload Directory + +```python +import os +from pathlib import Path + +# Prepare local directory +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# Upload to sandbox +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"Upload success: {result.output}") +else: + print(f"Upload failed: {result.failure_reason}") +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/model-service.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..23dbc21bfc --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service (Experimental) + +The Model Service provided by ROCK is responsible for handling AI model call communications, serving as a communication bridge between agents and training frameworks (such as Roll) or actual LLM inference services. + +## RockAgent Integration + +ModelService is typically **automatically managed by RockAgent** - no manual lifecycle management is required. Simply enable it in the configuration: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # Enable ModelService, RockAgent manages its lifecycle +) +``` + +RockAgent will automatically: +- Install ModelService (install Python runtime, install model service package) +- Start/stop ModelService +- Monitor Agent process + +## Architecture Overview (Local Mode) + +In local mode, the model service uses the **file system** as the communication medium, implementing a request-response mechanism between agents and models. + +When an agent needs to call a model, the request is first written to a log file, then processed by the listening component. When the model generates a response, the result is written back to the log file and read by the waiting agent. + +## anti_call_llm - Core API + +`anti_call_llm()` is the **most important API in Local mode**, used to manually trigger LLM anti-calls for fine-grained control over model calls: + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='OpenAI type response', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +**Use cases:** +- After Agent captures LLM response, call this method to notify Roll runtime +- Supports carrying response data for error handling or retry +- Configurable timeout and check interval for different network environments + +## CLI Commands + +To use the model service via CLI, ROCK provides a set of CLI commands that can be accessed via `rock model-service` after installing ROCK in the sandbox: + +### start command +Start the model service process +```bash +rock model-service start --type [local|proxy] [options] +``` + +Parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `--type` | str | `local` | Service type: `local` or `proxy` | +| `--config-file` | str | None | Path to configuration file | +| `--host` | str | None | Server host address (overrides config) | +| `--port` | int | None | Server port (overrides config) | +| `--proxy-base-url` | str | None | Proxy base URL | +| `--retryable-status-codes` | str | None | Comma-separated list of retryable status codes | +| `--request-timeout` | int | None | Request timeout in seconds | + +### watch-agent command +Monitor the agent process and send a SESSION_END message when the process exits +```bash +rock model-service watch-agent --pid +``` + +Parameters: +- `--pid`: The ID of the agent process to monitor + +### stop command +Stop the model service +```bash +rock model-service stop +``` + +### anti-call-llm command +Anti-call the LLM interface +```bash +rock model-service anti-call-llm --index [--response ] +``` + +Parameters: +- `--index`: Index of the previous LLM call, starting from 0 +- `--response`: Response from the previous LLM call (optional) + +## File Communication Protocol + +The model service uses files for inter-process communication, defining specific marker formats to distinguish requests and responses: + +### Request Format +``` +LLM_REQUEST_START{JSON request data}LLM_REQUEST_END{metadata JSON} +``` + +### Response Format +``` +LLM_RESPONSE_START{JSON response data}LLM_RESPONSE_END{metadata JSON} +``` + +### Session End Marker +``` +SESSION_END +``` + +Metadata contains timestamp and index information to ensure message order and processing. + +## SDK Usage + +### ModelServiceConfig + +Model service configuration class, located in `rock/sdk/sandbox/model_service/base.py`: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # Service type + install_cmd="pip install rock-model-service", # Install command + install_timeout=300, # Install timeout (seconds) + start_cmd="rock model-service start --type ${type}", # Start command + stop_cmd="rock model-service stop", # Stop command + logging_path="/data/logs", # Log path + logging_file_name="model_service.log", # Log filename +) +``` + +| Config | Default | Description | +|--------|---------|-------------| +| `enabled` | `False` | Whether to enable model service (RockAgent manages) | +| `type` | `"local"` | Service type: `local` or `proxy` | +| `install_cmd` | - | Model service package install command | +| `install_timeout` | `300` | Install timeout in seconds | +| `start_cmd` | - | Start command template | +| `stop_cmd` | - | Stop command | +| `logging_path` | `/data/logs` | Log directory path | +| `logging_file_name` | `model_service.log` | Log filename | + +### ModelService + +Model service management class, handles the lifecycle of model services within the sandbox: + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# Typically auto-managed by RockAgent, no manual calls needed +# The following methods are only for manual control when needed + +# Install model service +await model_service.install() + +# Start model service +await model_service.start() + +# Monitor agent process +await model_service.watch_agent(pid="12345") + +# Execute anti-call LLM (Core API for Local mode) +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# Stop model service +await model_service.stop() +``` + +## API Reference + +### install() + +Install model service dependencies in the sandbox. + +```python +await model_service.install() +``` + +Execution steps: +1. Create and initialize Python runtime environment +2. Create Rock config file +3. Install model service package + +**Note:** Typically auto-called by RockAgent. + +### start() + +Start the model service. + +```python +await model_service.start() +``` + +Prerequisite: Must call `install()` first. + +**Note:** Typically auto-called by RockAgent. + +### stop() + +Stop the model service. + +```python +await model_service.stop() +``` + +If the service is not running, this operation will be skipped. + +**Note:** Typically auto-called by RockAgent. + +### watch_agent(pid) + +Monitor the agent process. + +```python +await model_service.watch_agent(pid="12345") +``` + +Sends `SESSION_END` message when the process exits. + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +Execute anti-call LLM operation. **This is the most important API in Local mode.** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='{"result": "..."}', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +## Configuration Options + +### Service Configuration +- `SERVICE_HOST`: Service host address, defaults to `"0.0.0.0"` +- `SERVICE_PORT`: Service port, defaults to `8080` + +### Log Configuration +- `LOG_FILE`: Log file path used for communication, containing request and response data + +### Trajectory (Traj) Logging +The model service records LLM call trajectories (traj) to a JSONL file for debugging and analysis. + +| Environment Variable | Default | Description | +|---------------------|---------|-------------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | Directory for traj log files | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | Append mode (true/false) | + +**Traj file location**: `{DATA_DIR}/LLMTraj.jsonl` + +**Traj file format** (JSONL - one JSON object per line): +```json +{"request": {...}, "response": {...}} +``` + +### Polling Configuration +- `POLLING_INTERVAL_SECONDS`: Polling interval, defaults to `0.1` seconds +- `REQUEST_TIMEOUT`: Request timeout, defaults to unlimited + +### Marker Configuration +Defines markers used to distinguish different types of messages in the log file: +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig (Server-side) + +The server-side configuration class defines how the model service handles requests: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `host` | str | `"0.0.0.0"` | Server host address | +| `port` | int | `8080` | Server port | +| `proxy_base_url` | str \| None | `None` | Direct proxy URL | +| `proxy_rules` | dict | See below | Model name to URL mapping | +| `retryable_status_codes` | list[int] | `[429, 500]` | Retryable HTTP status codes | +| `request_timeout` | int | `120` | Request timeout in seconds | + +**Default proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/python_sdk.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..5272d7edb6 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK Reference + +This guide provides detailed information on how to use the ROCK SDK for development, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +ROCK SDK provides developers with convenient Python interfaces to use ROCK platform features, including sandbox environment management and GEM environment interaction. + +> **Important Note**: Before using the SDK, ensure that the ROCK Admin service is running. You can start it with the following command: +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 Basic Sandbox Operations + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 Sandbox Group Management + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# Create sandbox group configuration +config = SandboxGroupConfig( + image="python:3.11", + size=4, # Create 4 sandboxes + start_concurrency=2, # Concurrency level for startup is 2 +) + +# Create and start sandbox group +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# Batch operations +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# Batch stop +await sandbox_group.stop() +``` + +### 2.3 Configuration Example + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 Sandbox Speedup Configuration + +ROCK provides sandbox network acceleration capabilities, supporting configuration of APT, PIP, and GitHub mirror sources to improve package download speeds in restricted network environments. + +#### Supported Speedup Types + +**APT Mirror Configuration** + +Configure APT package manager mirror sources for faster Debian/Ubuntu package downloads. + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# Configure APT mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP Mirror Configuration** + +Configure Python package index mirrors for faster pip installations. + +```python +# HTTP mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub Acceleration** + +Configure GitHub IP acceleration by adding custom DNS resolution entries. + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### Complete Example + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """Create sandbox and configure acceleration""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # Configure acceleration (before installing packages) + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # Speedup does not automatically install PIP, it only configures mirror sources for acceleration + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # GitHub can be accelerated through mirror IP + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### Important Notes + +1. **Configuration Order**: Configure speedup before installing packages +2. **HTTPS vs HTTP**: HTTPS mirrors don't require trusted-host configuration for PIP +3. **GitHub IP**: Different regions may require different IPs for optimal performance +4. **Persistence**: Configurations persist within the sandbox lifecycle +5. **Multiple Calls**: Subsequent speedup calls will override previous configurations +6. **PIP Installation**: The speedup feature only configures mirror sources and does not automatically install PIP + +## 3. GEM SDK + +### 3.1 Python SDK Approach + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## Related Documents +- [Quick Start Guide](../../Getting%20Started/quickstart.md) - Learn how to quickly get started with the ROCK SDK +- [API Documentation](../api.md) - View the underlying API interfaces encapsulated by the SDK +- [Configuration Guide](../../User%20Guides/configuration.md) - Learn about SDK-related configuration options +- [Installation Guide](../../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/remote_user.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..810bbcf028 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/remote_user.md @@ -0,0 +1,70 @@ +# Remote User + +Remote user management for creating and managing users in the sandbox. + +## Usage Examples + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +Create a remote user. + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +Check if a user exists. + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/rock-agent.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..f24ade49ac --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent (Experimental) + +RockAgent is the core Agent implementation in the ROCK framework, directly inheriting from the `Agent` abstract base class. It provides complete Agent lifecycle management, including environment initialization, ModelService integration, command execution, and more. + +Using `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, you can install and run Agents in the Sandbox environment provided by Rock. + +## Core Concepts + +The core workflow of RockAgent is divided into two phases: + +1. **install(config)**: Initialize the Agent environment, including deploying the working directory, setting environment variables, initializing the runtime environment, etc. +2. **run(prompt)**: Execute the Agent task, replace placeholders, and start the Agent process + +## Quick Start + +### Claude Code Example + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli Example + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} is required + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # Environment variables + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent Example + +```yaml +working_dir: "." # Upload local current directory containing langgraph_agent.py to sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # Run local script + +runtime_env_config: + type: python + pip: # Install pip dependencies + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## Configuration Details + +### Basic Configuration + +```yaml +agent_type: "default" # Agent type identifier (default: "default") +agent_name: "demo-agent" # Agent instance name (default: random uuid) +version: "1.0.0" # Version identifier (default: "default") +instance_id: "instance-001" # Instance ID (default: "instance-id-") +agent_installed_dir: "/tmp/installed_agent" # Agent installation directory (default: "/tmp/installed_agent") +agent_session: "my-session" # Bash session identifier (default: "agent-session-") +env: # Environment variables (default: {}) + OPENAI_API_KEY: "xxxxxxx" +``` + +### Working Directory Configuration + +```yaml +working_dir: "./my_project" # Local directory to upload to sandbox (default: None, no upload) +project_path: "/testbed" # Working directory in sandbox for cd (default: None) +use_deploy_working_dir_as_fallback: true # Whether to fall back to deploy.working_dir when project_path is empty (default: true) +``` + +### Execution Configuration + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent execution command, must contain ${prompt} (default: None) + +skip_wrap_run_cmd: false # Skip wrapping run_cmd with PATH (default: false) + +# Timeout configuration +agent_install_timeout: 600 # Installation timeout in seconds (default: 600) +agent_run_timeout: 1800 # Run timeout in seconds (default: 1800) +agent_run_check_interval: 30 # Check interval in seconds (default: 30) +``` + +**`skip_wrap_run_cmd`**: +- `false` (default): Wraps the command with `export PATH=:$PATH &&` to ensure runtime environment executables are used +- `true`: Skips PATH wrapping, runs the command directly with `bash -c` + +### Initialization Hooks + +```yaml +pre_init_cmds: # Commands executed before initialization (default: read from env_vars) + - command: "apt update && apt install -y git" + timeout_seconds: 300 # Command timeout in seconds (default: 300) + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # Commands executed after initialization (default: []) + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**Notes**: +- `pre_init_cmds` and `post_init_cmds` do not inherit the Agent's `env` environment variables +- Typically used for installation operations and configuration file movement +- Common command examples: + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv Configuration + +```yaml +runtime_env_config: # Refer to RuntimeEnv documentation for details + type: "python" # Runtime type: python / node (default: "python") + version: "3.11" # Version number + pip: # Python dependency package list + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node Runtime Example**: + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**Automatic Operations**: +- Install corresponding runtime based on `type` (Python or Node.js) +- Install `pip` dependencies (if configured) +- Execute `custom_install_cmd` custom installation command (if configured) +- Support `npm_registry` configuration for Node.js npm mirror source + +### ModelService Configuration + +```yaml +model_service_config: # Refer to ModelService documentation for details + enabled: true # Enable ModelService (default: false) +``` + +**Automatic Operations**: +- Installation phase: Install ModelService (install only, do not start) +- Run phase: Start ModelService + `watch_agent` monitoring process + +**Notes**: You need to set the model request URL to the ModelService URL. For example, if the ModelService provides an OpenAI-compatible URL at `http://127.0.0.1:8080/v1/chat/completions`, you typically need to set the Agent's LLM request URL to `http://127.0.0.1:8080/v1/`. + +## API Reference + +### install(config) + +Initialize the Agent environment. + +**Execution Flow**: +1. If `working_dir` is configured, deploy to sandbox +2. Set up bash session and configure env environment variables +3. Execute `pre_init_cmds` +4. Initialize RuntimeEnv and ModelService in parallel (if enabled) +5. Execute `post_init_cmds` + +**Parameters**: +- `config`: Agent configuration file, supports two input methods: + - **String path**: YAML configuration file path, default value is `"rock_agent_config.yaml"` + - **RockAgentConfig object**: Directly pass a `RockAgentConfig` instance + +### run(prompt) + +Execute the Agent task. + +**Execution Flow**: +1. Replace placeholders and prepare Agent run command +2. Start the agent process +3. If ModelService is enabled, start `watch_agent` +4. Wait for task completion and return results + +## Advanced Usage + +### Difference and Interaction between working_dir and project_path + +| Configuration | Function | Interaction Method | +|--------------|----------|-------------------| +| `working_dir` | Local directory uploaded to sandbox | Calls `deploy.deploy_working_dir()` to upload, after upload `deploy.working_dir` becomes the path in sandbox | +| `${working_dir}` | Placeholder in commands | Replaced by `deploy.format()` with the value of `deploy.working_dir`, replaced in init_cmds and run_cmd in the configuration | +| `project_path` | Working directory in sandbox | Used for `cd project_path` before running, when not set it enters the `deploy.working_dir` working directory | +| `use_deploy_working_dir_as_fallback` | Whether to fall back to deploy.working_dir when project_path is not set at runtime | Default is `true`, when set to `false` it will not enter working_dir even if project_path is not set | + +**Usage Recommendations**: +- Use `working_dir` to upload local project code to sandbox +- Use `project_path` to specify the working directory in sandbox (e.g., `/testbed`) +- Set `use_deploy_working_dir_as_fallback: false` scenario: Need to perform local file mounting, but want to run Agent in the image's default working directory + +### Placeholder Usage + +Rock Agent supports replacing the following placeholders in the configuration file: + +- `${prompt}`: Required in run_cmd, will be replaced with the prompt passed to `run(prompt)` +- `${working_dir}`: Optional, will be replaced with the actual working directory path in sandbox, also supported in init_cmds and run_cmd +- `${bin_dir}`: Optional, will be replaced with the runtime environment's bin directory path + +**Example**: +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback Explanation + +When `project_path` is not set: +- `true` (default): Before running Agent, it will automatically `cd` to `deploy.working_dir` +- `false`: Before running Agent, it will not automatically switch directories, staying in the current directory + +Applicable Scenarios: +- `true`: Most scenarios, where you want Agent to run in the uploaded code directory +- `false`: Need to mount local files, but want to run Agent in the image's default working directory (e.g., `/app`, `/testbed`) + +## Complete Configuration Example + +```yaml +# ========== Basic Configuration ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== Working Directory Configuration ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== Run Configuration ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# Timeout configuration +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== Initialization Commands ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== Runtime Environment Configuration ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService Integration ========== +model_service_config: + enabled: true +``` + +## Usage Examples + +### Using YAML Configuration File (Recommended) + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/runtime-env.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..e1996cfcdb --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,136 @@ +# RuntimeEnv + +The RuntimeEnv module is used to manage language runtime environments in the sandbox (currently providing Python / Node.js). + +## Quick Start (Example) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +An async factory method that creates and initializes a RuntimeEnv instance based on the configuration, and automatically registers it to `sandbox.runtime_envs`. + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# Auto-registered; accessible via sandbox.runtime_envs[env.runtime_env_id] +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +Wraps a command by adding `bin_dir` to PATH to ensure executables from the runtime environment are used with priority. + +```python +wrapped = env.wrapped_cmd("node script.js") +# Returns: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +Executes a command within the runtime environment. Internally implemented based on `wrapped_cmd`. + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | Type identifier | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python version; default is 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | List of pip packages or a requirements.txt path | +| `pip_index_url` | `str \| None` | Environment variable | pip index mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | List of executables to symlink | + +## NodeRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | Type identifier | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node version; default is 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm registry mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | List of executables to symlink | + +## Constraints for Custom RuntimeEnv Implementations + +A custom RuntimeEnv must follow these rules: + +1. **Define the `runtime_env_type` class attribute**: used as a type identifier for automatic registration into the RuntimeEnv factory +2. **Override `_get_install_cmd()`**: return the install command +3. **The install command must end with**: renaming the directory to `runtime-env` + +## Simplified NodeRuntimeEnv Implementation Example + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config class: defines the config type so RuntimeEnv.create() can route to the corresponding implementation +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # Must match runtime_env_type + +# RuntimeEnv implementation class: defines how to install and run this runtime environment +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # Auto-registered to RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # Download the Node binary tarball and extract it, then rename to runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## Speeding Up Base Runtime Installation + +`PythonRuntimeEnv` downloads Python packages from https://github.com/astral-sh/python-build-standalone/releases/ by default. If the network is unavailable or slow, you can override the default install command via `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` or `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` (e.g., switch to an internal registry or a mirror). + +Default value example: + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +For example, override it to download from a mirror: + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +Make sure the command creates a `runtime-env` directory under the default working directory of `runtime_env`, and that `${workdir}/runtime-env/bin/` contains the expected executables, e.g.: + +- `${workdir}/runtime-env/bin/python` + +The same applies to Node.js: you can override the install command via `ROCK_RTENV_NODE_V22180_INSTALL_CMD` to use a faster download/install method. \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/sandbox.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..e6e1e43124 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/sandbox.md @@ -0,0 +1,114 @@ +# Handling Large Files and Long Command Outputs + +## `arun` + +`arun()` provides two knobs to control how `nohup` output is handled: + +1. **`response_limited_bytes_in_nohup`** *(integer type)* + Caps the number of characters returned from the nohup output file. Useful when you still need to stream some logs back but want an upper bound (default `None` = no cap). + +2. **`ignore_output`** *(bool, default `False`)* + When set to `True`, `arun()` skips reading the nohup output file entirely. The command still runs to completion and writes logs to `/tmp/tmp_.out`, but the SDK immediately returns a lightweight hint telling agents where to fetch the logs later (via `read_file`, download APIs, or custom commands). This fully decouples "execute command" from "inspect logs". The response also includes the **file size** to help users decide whether to download directly or read in chunks. + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# Example 1: limit the returned logs to 1024 characters +resp_limited = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# Example 2: skip collecting logs; agent will download/read them later +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# Use Sandbox.read_file(...), download APIs, or run 'cat /tmp/tmp_xxx.out' ... +``` + +## `read_file_by_line_range` + +Asynchronously reads file content by line range, with built-in support for automatic chunking and session management. Supports large file reading. + +### Key Features +- **Chunked reading for large files**: Automatically splits large files into chunks +- **Automatic line count**: Estimates total lines when end_line is not specified +- **Built-in retry mechanism**: Up to 3 retries for critical operations +- **Input validation**: Validates input parameters automatically +- **Session management**: Supports custom session or auto-created temporary session + +### Parameters +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `file_path` | str | - | File path to read (absolute or relative path in sandbox) | +| `start_line` | int \| None | 1 | Starting line number (1-based) | +| `end_line` | int \| None | None | Ending line number (inclusive), defaults to file end | +| `lines_per_request` | int | 1000 | Lines per request, range 1-10000 | + +### Return Value +- `ReadFileResponse`: Response object containing file content + - `content` (str): The file content read + +### Exception Handling +- `Exception`: Raised when `start_line < 1` +- `Exception`: Raised when `end_line < start_line` +- `Exception`: Raised when `lines_per_request` is not in range 1-10000 +- `Exception`: Raised when file reading fails + +### Usage Examples + +```python +# Read the entire file +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# Read a specific line range (lines 100 to 500) +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# Read from line 1990 to the end of file +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# Use custom chunk size +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### Notes +- Line numbers are 1-based, not 0-based +- For large files, consider increasing `lines_per_request` for better efficiency +- File path must be a valid path within the sandbox +- Uses `sed` command for file reading; ensure the sandbox image supports this command diff --git a/docs/versioned_docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md b/docs/versioned_docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..85f34cbe7c --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,229 @@ +# SWE-Bench Evaluation + +This guide demonstrates how to use the ROCK SDK to run SWE-Bench Verified evaluations, including sandbox setup, Agent integration, test environment preparation, and result parsing. + +### Quick Start + +SWE-Bench is a benchmark for evaluating AI coding agents on real-world software engineering tasks. + +Running a SWE-Bench task on ROCK involves the following steps: + +1. **load_task_config** — Load `task.yaml` to get the task instruction +2. **start_sandbox** — Start a sandbox with a task-specific Docker image +3. **agent.install / agent.run** — Install and run the Agent to solve the task +4. **setup_test_env** — Upload test files and run-test script to the sandbox +5. **Run tests** — Execute the test script via `sandbox.arun()` with timeout +6. **parse_swebench_result** — Parse test output to determine PASSED / FAILED +7. **sandbox.stop** — Clean up sandbox resources + +**Here is an example code** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. Load task instruction + task_config = await load_task_config(task_dir) # see load_task_config section + instruction = task_config["instruction"] + + # 2. Start sandbox + sandbox = await start_sandbox(task_name) # see start_sandbox section + + try: + # 3. Install and run Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. Setup test environment + await setup_test_env(sandbox, task_dir) # see setup_test_env section + + # 5. Run tests + resp = await run_tests(sandbox) # see Running Tests section + + # 6. Parse results + is_resolved = parse_swebench_result(resp.output) # see parse_swebench_result section + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +The following sections describe each function used in the workflow in detail. + +--- + +## start_sandbox + +Start a sandbox instance with a task-specific SWE-Bench Docker image. Each task has a pre-built image containing the target repository and environment. + +The `image` parameter follows the format: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +For example, task `django__django-14539` maps to: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +Load task configuration from a `task.yaml` file in the task directory. The YAML file contains the `instruction` field that describes the coding task for the Agent. + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# Usage +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +Use `sandbox.agent.install()` and `sandbox.agent.run()` to deploy and execute an Agent inside the sandbox. Refer to [Rock Agent](./rock-agent.md) for detailed Agent configuration. + +```python +# Install Agent with a YAML configuration file（e.g., iflow_config.yaml） +await sandbox.agent.install(config="iflow_config.yaml") + +# Run Agent with the task instruction +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +Prepare the test environment in the sandbox: install the [uv](https://github.com/astral-sh/uv) package manager, and upload test files and the run-test script. + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """Set up the test environment and return the session name.""" + # 1. Create a session with custom environment variables + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. Install uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. Upload test files + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. Upload run-tests script + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## Running Tests + +Execute the test script with a configurable timeout using `RunMode.NOHUP`. + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +Parse the test output to determine whether the SWE-Bench task is resolved. The parser looks for a result block delimited by marker lines and checks for `PASSED`. + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """Parse SWE-Bench test output to determine if the task is resolved. + + Matches the block between 'SWEBench results starts here' and + 'SWEBench results ends here', then checks whether it contains 'PASSED'. + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# Usage +is_resolved = parse_swebench_result(resp.output) +``` + +## Notes + +- **Task Datasets**: Task directories (containing `task.yaml`, `tests/`, and `run-tests.sh`) can be obtained from the [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) repository. +- **Task Images**: Each SWE-Bench task requires a specific Docker image (e.g., `sweb.eval.x86_64.`). Ensure the image is available before running tests. +- **Agent Config**: The Agent configuration YAML defines the runtime, dependencies, and execution command. See [Rock Agent](./rock-agent.md) for details. + diff --git a/docs/versioned_docs/version-1.6.x/References/api.md b/docs/versioned_docs/version-1.6.x/References/api.md new file mode 100644 index 0000000000..d73bf49d6c --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/References/api.md @@ -0,0 +1,195 @@ +--- +sidebar_position: 1 +--- + +# API Reference + +This guide provides detailed information about the core API services provided by the ROCK platform, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +The ROCK platform provides two core API services: +- Sandbox API: Sandbox environment management +- GEM API: GEM environment interaction + +All API interfaces follow RESTful design principles and support JSON format data exchange. + +## 2. Sandbox API + +Full lifecycle management functions for sandbox environments: + +### Sandbox Management Interfaces + +1. **Start Sandbox** - Start a sandbox environment + - Create a new sandbox instance + - Support specifying image, resource configuration and other parameters + +2. **Start Sandbox Async** - Asynchronously start a sandbox environment + - Asynchronously create a sandbox instance + - Suitable for scenarios requiring quick response + +3. **Check Sandbox Alive Status** - Check sandbox alive status + - Verify if the sandbox is running normally + +4. **Get Sandbox Statistics** - Get sandbox statistics + - Get resource usage statistics of the sandbox + +5. **Get Sandbox Status** - Get detailed sandbox status + - Get complete status information of the sandbox + +6. **Stop Sandbox** - Stop sandbox environment + - Safely shut down the sandbox instance + +7. **Commit Sandbox** - Commit sandbox as image + - Save current sandbox state as a new image + +### Command Execution Interfaces + +8. **Execute Command** - Execute command in sandbox + - Run specified command directly in the sandbox + +9. **Create Bash Session** - Create Bash session + - Create a persistent Bash session environment + +10. **Run Command in Session** - Run command in session + - Execute command in a created session + +11. **Close Session** - Close session + - Release session resources + +### File Operation Interfaces + +12. **Read File** - Read sandbox file + - Read specified file content from the sandbox + +13. **Write File** - Write sandbox file + - Write file to the sandbox + +14. **Upload File** - Upload file to sandbox + - Upload local file to the sandbox + +## 3. GEM API + +GEM environment interaction functions: + +1. **Make Environment** - Create GEM environment + - Initialize a new GEM environment instance + +2. **Reset Environment** - Reset GEM environment + - Reset GEM environment to initial state + +3. **Step Environment** - Execute GEM environment step + - Execute an action step in the GEM environment + +4. **Close Environment** - Close GEM environment + - Release GEM environment resources + + +## 4. HTTP API Usage Examples + +### 4.1 Sandbox API Examples + +#### Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Asynchronously Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Execute Command +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### Create Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### Run Command in Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### Upload File +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### Stop Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API Examples + +```bash +# Create GEM environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# Reset environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# Execute step +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# Close environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly get started with ROCK API +- [Python SDK Documentation](./Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to call APIs +- [Configuration Guide](../User%20Guides/configuration.md) - Learn about API-related configuration options +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/Release Notes/index.md b/docs/versioned_docs/version-1.6.x/Release Notes/index.md new file mode 100644 index 0000000000..38262ea2a6 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# Release Notes +* [release v1.6.0](v1.6.0.md) diff --git a/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.0.md b/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.0.md new file mode 100644 index 0000000000..8d7a7d2c85 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.0.md @@ -0,0 +1,111 @@ +# v1.6.0 + +## Release Date + +April 17, 2026 + +--- + +## Highlights + +This release centers on a major refactor of the **Job module**, introducing a clean `Job` / `Operator` / `Executor` / `Trial` abstraction. On top of this foundation, two new job types are now supported — **BashJob** (with OSS artifact mirroring) and **HarborJob**. The CLI `rock job run` command has been reworked with strict YAML validation and auto job-type detection. + +--- + +## Job Module + +### New Architecture + +#### Job / Operator / Executor / Trial Abstraction + +* **NEW**: Refactored the Job module into a layered architecture: `Job` → `Operator` → `Executor` → `Trial`. This decouples job orchestration, scheduling, execution, and per-attempt logic, making it easy to add new job types ([#779](https://github.com/alibaba/ROCK/pull/779), [#780](https://github.com/alibaba/ROCK/pull/780)) + +* Hoisted `on_sandbox_ready` backfill logic to `AbstractTrial` to share across all trial types ([#788](https://github.com/alibaba/ROCK/pull/788), [#789](https://github.com/alibaba/ROCK/pull/789)) + +#### BashJob + +* **NEW**: BashJob trial support — submit shell scripts as jobs via the SDK or CLI. See `examples/bash/simple_bash_job_demo.sh` ([#772](https://github.com/alibaba/ROCK/pull/772)) + +* **NEW**: BashJob OSS Mirror — automatically uploads job artifacts to OSS after completion ([#823](https://github.com/alibaba/ROCK/pull/823)) + +* Added `claw-eval` BashJob demo under `examples/evaluation/claw_eval/` ([#804](https://github.com/alibaba/ROCK/pull/804)) + +#### HarborJob + +* **NEW**: HarborJob trial support, enabling agent-style job submission via the new Job abstraction ([#798](https://github.com/alibaba/ROCK/pull/798)) + +### Configuration & Validation + +#### Job Type Auto-Detection + +* **NEW**: `rock job run` now auto-detects job type from YAML via strict Pydantic model validation — no need to specify `--type` ([#814](https://github.com/alibaba/ROCK/pull/814)) + +#### Dual-Mode Input + +* **NEW**: `rock job run` reworked to support both `--config` (full YAML) and inline argument modes with strict validation and clearer error messages ([#818](https://github.com/alibaba/ROCK/pull/818)) + +#### Native Template Config + +* **NEW**: Added `TemplateConfig` and `template` field to `NativeConfig` for template-based job definitions ([#786](https://github.com/alibaba/ROCK/pull/786)) + +#### Auto-Generated job_name + +* Truncate long path segments when auto-generating `job_name` to avoid downstream length limits ([#791](https://github.com/alibaba/ROCK/pull/791)) + +#### Default Timeout Increased + +* Increased default `JobConfig` timeout from 3600s to **7200s** to better match real-world long-running jobs ([#806](https://github.com/alibaba/ROCK/pull/806), [#810](https://github.com/alibaba/ROCK/pull/810)) + +### Bug Fixes + +* `JobConfig.experiment_id` now takes priority over `environment.experiment_id` to give callers explicit control ([#822](https://github.com/alibaba/ROCK/pull/822)) + +* `BashTrial.collect` now correctly populates `raw_output` and `exit_code` ([#808](https://github.com/alibaba/ROCK/pull/808)) + +* Removed redundant `shebang` and `set -e` injection from `BashTrial.build()` to honor user scripts ([#816](https://github.com/alibaba/ROCK/pull/816)) + +--- + +## EnvHub + +### Refactoring + +* **BREAKING**: `JobEnvironmentConfig` has been moved to `envhub` as `EnvironmentConfig`. Update import paths accordingly ([#800](https://github.com/alibaba/ROCK/pull/800)) + +* Removed the deprecated `auto_stop` parameter from `EnvironmentConfig` — sandbox lifetime is now controlled via `auto_delete_seconds` (introduced in v1.5.0) ([#820](https://github.com/alibaba/ROCK/pull/820)) + +* Refactored EnvHub uploads pipeline for clearer separation between client and server responsibilities ([#802](https://github.com/alibaba/ROCK/pull/802)) + +--- + +## Admin + +### Database + +* Enlarged `SandboxRecord.image` column from 255 to **512 chars** to support longer registry paths, and disabled asyncpg prepared-statement cache to avoid PgBouncer compatibility issues ([#794](https://github.com/alibaba/ROCK/pull/794)) + +--- + +## CLI + +* Lazy-import `psutil` in the `admin stop` command to speed up unrelated CLI invocations and avoid import-time failures when `psutil` is unavailable ([#831](https://github.com/alibaba/ROCK/pull/831)) + +--- + +## Testing & CI + +* Added `SELECT 1` readiness check to the `pg_container` test fixture to avoid flaky integration tests ([#778](https://github.com/alibaba/ROCK/pull/778)) + +* Applied repo-wide `ruff format` pass and translated remaining Chinese comments to English ([#812](https://github.com/alibaba/ROCK/pull/812)) + +--- + +## Migration Notes + +* **`auto_stop` removed**: Update any `EnvironmentConfig` usage that relied on `auto_stop`. Use `auto_delete_seconds` (introduced in v1.5.0) instead. + +* **`JobEnvironmentConfig` → `EnvironmentConfig`**: Update imports from `rock.sdk.agent.models.job` to `rock.sdk.envhub`. + +* **Default Job timeout**: Jobs that previously hit the 3600s timeout will now run up to 7200s by default. Set `JobConfig.timeout` explicitly if you need the old behavior. + +* **`rock job run`**: The CLI now validates YAML strictly. Configs with unknown or missing fields that previously passed silently will now fail fast. diff --git a/docs/versioned_docs/version-1.6.x/User Guides/configuration.md b/docs/versioned_docs/version-1.6.x/User Guides/configuration.md new file mode 100644 index 0000000000..604256878b --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/User Guides/configuration.md @@ -0,0 +1,189 @@ +--- +sidebar_position: 4 +--- + +# Configuration + +This guide provides detailed instructions on how to configure the ROCK environment to meet different usage requirements, including local development, testing, and production deployment. + +## 1. Environment Variable Configuration + +ROCK supports configuring key parameters through environment variables. The main environment variables are as follows: + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK service base URL +export ROCK_LOG_LEVEL=INFO # Log level +export ROCK_LOGGING_PATH=/path/to/logs # Log file path, default None (output to console) +export ROCK_LOGGING_FILE_NAME=rocklet.log # Log file name, default "rocklet.log", can be customized by admin like admin.log +export ROCK_LOGGING_LEVEL=INFO # Log output level, default "INFO" +export ROCK_WORKER_ENV_TYPE=local # Runtime environment type, options: local, docker, uv, pip +``` + +More environment variables can be found in `rock/env_vars.py`. + +### 1.1 Runtime Environments + +ROCK provides multiple different runtime environments to meet the needs of different scenarios, configured through the `ROCK_WORKER_ENV_TYPE` environment variable. Each environment has different deployment requirements, performance characteristics and applicable scenarios. Each environment has its own unique advantages and limitations, and developers can choose the most suitable runtime environment according to their deployment needs. + +#### 1.1.1 Docker Runtime Environment + +The Docker runtime environment is suitable for Docker image environments where dependencies are pre-installed. This environment requires the `/tmp/miniforge/bin/rocklet` executable to be directly available in the deployment environment. + +**Mount Configuration:** +- `/tmp/miniforge` - Contains pre-installed Python environment +- `/tmp/local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Containerized deployment environments +- Already built custom Docker image containing `rocklet` +- Suitable for production, fast startup + +**Requirements:** +- Requires a custom Docker image containing `/tmp/miniforge/bin/rocklet` executable +- Docker environment support + +#### 1.1.2 Local Runtime Environment + +The local runtime environment directly uses the Python environment and project files of the current deployment. This environment requires the same operating system between the host and container to directly mount the virtual environment and Python interpreter. + +**Mount Configuration:** +- `python_env_path` - Python environment path +- `project_root` - Project root directory +- `.venv` - Virtual environment directory (mounted as `/tmp/miniforge` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Development environments +- Scenarios where host and target container use the same operating system +- Need to quickly reuse existing Python environment + +**Requirements:** +- Same operating system (host/container) +- Direct access to the currently deployed `.venv` virtual environment +- Python interpreter path compatibility + +#### 1.1.3 UV Runtime Environment + +The UV runtime environment only depends on the available ROCK project, but initialization is relatively slow and network requirements are higher. This environment is most suitable for scenarios without preconfigured environments. It rebuilds the rocklet environment from the original project. This is the recommended environment for Mac OS. + +**Mount Configuration:** +- `project_root` - Project root directory (mounted as `/tmp + project_root` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**Use Cases:** +- Mac OS +- Cross-OS startup +- Scenarios without preconfigured environment +- No uv management Rock + +**Advantages:** +- No pre-built image required +- Good cross-platform compatibility +- Suitable for development and testing especially + +**Limitations:** +- Initialization is relatively slow +- Higher network requirements +- Longer startup time + +#### 1.1.4 PIP Runtime Environment + +The PIP runtime environment uses pip to install required dependencies in the container. This environment is suitable for quick setup and scenarios where dependencies can be installed in the container. It is the default runtime environment. It does not require pre-built images containing dependencies, and manages Python packages directly through pip. + +**Mount Configuration:** +- `local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**Use Cases:** +- ROCK installation from PIP source +- Fast testing of ROCK + +**Advantages:** +- Simple deployment setup + +**Limitations:** +- Long dependency installation time +- Requires network access to install dependency packages +- Dependencies need to be installed each time on startup + +#### 1.1.5 Configuration Guide + +Refer to the following selection guide for different use cases: + +| Scenario | Recommended Environment | Reason | +|----------|--------------------------|-------| +| Production environment | Docker Runtime | Fast startup, stable performance | +| Development environment, same OS | Local Runtime | Environment reuse, fast development cycle | +| Mac development | UV Runtime | Best cross-platform compatibility support | +| Cross-platform development | UV Runtime | Avoids environment compatibility issues | +| Fast testing | UV Runtime | Requires no pre-configuration | +| PIP source installation | PIP Runtime | Install dependencies directly with pip | + +These runtime environments are configured through the `ROCK_WORKER_ENV_TYPE` environment variable, which can be set to "local", "docker", "uv" or "pip". + +### 1.2 Logging Configuration + +Regarding logging configuration, ROCK's logging system has the following characteristics: + +- The logging system cannot output to both file and console simultaneously. If `ROCK_LOGGING_PATH` is set, logs will be output to the designated file, otherwise to console. +- `ROCK_LOGGING_LEVEL` is used to control the output log level, while `ROCK_LOG_LEVEL` is used for general log level settings. + +## 2. Distributed Deployment Requirements + +Since ROCK supports distributed deployment, when running on different nodes of a Ray cluster, the following consistency requirements must be met: + +#### Directory Structure Consistency + +On all Ray nodes, the following directory structure must be completely consistent: +- ROCK project repository directory +- `.venv` virtual environment directory +- The base Python directory that `.venv` depends on + + +#### Mounting Requirements + +ROCK's startup depends on mounting the ROCK project and the corresponding base Python environment, requiring consistency in multi-machine environments: + +#### Verifying Distributed Configuration + +Distributed deployment configuration can be verified through the following methods: + +```bash +# Check directory consistency on all nodes +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# Verify Python environment availability +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# Check environment variable settings on all nodes +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly set up the ROCK environment +- [API Documentation](../References/api.md) - View sandbox-related API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to configure sandboxes +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.6.x/overview.md b/docs/versioned_docs/version-1.6.x/overview.md new file mode 100644 index 0000000000..0c377e8b25 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/overview.md @@ -0,0 +1,33 @@ +--- +sidebar_position: 1 +--- + +# Overview + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework designed to simplify the development, deployment, and management of reinforcement learning environments. + +## What is ROCK + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework. By using ROCK, developers can quickly develop reinforcement learning environments and integrate with other reinforcement learning training frameworks to implement efficient reinforcement learning training. + +ROCK provides comprehensive sandbox environment management capabilities, supports containerized deployment, and enables rapid creation, execution, and destruction of environments. Additionally, ROCK is compatible with the GEM protocol, providing standardized interfaces for reinforcement learning environments. + +## Core Capabilities of ROCK + +1. **Simplified Development Process**: Simplifies the development, construction, and management of reinforcement learning environments, supporting various open-source reinforcement learning environments +2. **Large-scale Scheduling and Deployment**: Enables large-scale scheduling and deployment of rapid reinforcement learning environments. By supporting the GEM protocol, reinforcement learning environments can be easily accessed +3. **Framework Integration**: Integrates with other reinforcement learning training frameworks to achieve large-scale and scalable reinforcement learning training + +## Value of ROCK + +ROCK provides significant value to different roles of engineers: + +- **Reinforcement Learning Algorithm Engineers**: ROCK simplifies the development process of reinforcement learning environments, allowing engineers to focus on algorithm implementation +- **Reinforcement Learning Application Engineers**: ROCK enables large-scale deployment of rapid reinforcement learning environments, improving application development efficiency + +## Learn More + +- [Quick Start Guide](./Getting%20Started/quickstart.md) - Get started with ROCK quickly +- [Configuration Guide](./User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](./References/api.md) - View ROCK's API interfaces +- [Python SDK Documentation](./References/Python%20SDK%20References/python_sdk.md) - Learn how to use ROCK's Python SDK \ No newline at end of file diff --git a/docs/versioned_sidebars/version-1.6.x-sidebars.json b/docs/versioned_sidebars/version-1.6.x-sidebars.json new file mode 100644 index 0000000000..b475b11530 --- /dev/null +++ b/docs/versioned_sidebars/version-1.6.x-sidebars.json @@ -0,0 +1,64 @@ +{ + "tutorialSidebar": [ + "overview", + { + "type": "category", + "label": "Getting Started", + "link": { + "type": "doc", + "id": "Getting Started/quickstart" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Getting Started" + } + ] + }, + { + "type": "category", + "label": "User Guides", + "items": [ + { + "type": "autogenerated", + "dirName": "User Guides" + } + ] + }, + { + "type": "category", + "label": "References", + "items": [ + "References/api", + { + "type": "category", + "label": "Python SDK References", + "link": { + "type": "doc", + "id": "References/Python SDK References/python_sdk" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "References/Python SDK References" + } + ] + } + ] + }, + { + "type": "category", + "label": "Release Notes", + "link": { + "type": "doc", + "id": "Release Notes/index" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Release Notes" + } + ] + } + ] +} diff --git a/docs/versions.json b/docs/versions.json index 09fda0e173..aa33c9a4b3 100644 --- a/docs/versions.json +++ b/docs/versions.json @@ -1,4 +1,5 @@ [ + "1.6.x", "1.5.x", "1.4.x", "1.3.x", diff --git a/pyproject.toml b/pyproject.toml index a1859f6a01..bdb710358a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.6.0.dev2" +version = "1.6.0" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 3b0b331a2eedca4f3fce0b8f15d44ffc0a1b5c17 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Mon, 20 Apr 2026 14:22:34 +0800 Subject: [PATCH 052/226] fix iflow agent version --- .../integration/sdk/sandbox/agent/rock_agent/iflow_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/sdk/sandbox/agent/rock_agent/iflow_config.yaml b/tests/integration/sdk/sandbox/agent/rock_agent/iflow_config.yaml index 9a41214cac..3f905f63c4 100644 --- a/tests/integration/sdk/sandbox/agent/rock_agent/iflow_config.yaml +++ b/tests/integration/sdk/sandbox/agent/rock_agent/iflow_config.yaml @@ -2,7 +2,7 @@ run_cmd: "iflow -p ${prompt} --yolo" runtime_env_config: type: node - custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@0.5.17" env: IFLOW_API_KEY: "xxxxxxx" From 944ed89df8a39ad1a5576becf81a0548afa41f15 Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Mon, 20 Apr 2026 16:08:39 +0800 Subject: [PATCH 053/226] fix(docker-auth): reconstruct temporary directory authorization scheme and remove traditional scheme (#81160079) (#837) * fix(docker-auth): ensure _temp_dir is cleared even when directory already deleted (#81160079) * fix(docker-auth): reconstruct the temporary directory authorization scheme and remove the traditional scheme. (#81160079) * refactor(docker): relocate TempAuthDockerClient to deployments module; fix test_rock_agent_run_iflow (#81160079) --- rock/deployments/docker.py | 49 ++- rock/deployments/docker_client.py | 234 +++++++++++ rock/env_vars.py | 4 + .../test_docker_auth_integration.py | 255 ++++++++++++ tests/unit/deployments/test_docker_client.py | 381 ++++++++++++++++++ 5 files changed, 909 insertions(+), 14 deletions(-) create mode 100644 rock/deployments/docker_client.py create mode 100644 tests/integration/deployments/test_docker_auth_integration.py create mode 100644 tests/unit/deployments/test_docker_client.py diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 67246ddeeb..f27aba3594 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -19,7 +19,6 @@ from rock.deployments.config import DockerDeploymentConfig from rock.deployments.constants import Port, Status from rock.deployments.hooks.abstract import CombinedDeploymentHook, DeploymentHook -from rock.deployments.hooks.docker_login import DockerLoginHook from rock.deployments.runtime_env import DockerRuntimeEnv, LocalRuntimeEnv, PipRuntimeEnv, UvRuntimeEnv from rock.deployments.sandbox_validator import DockerSandboxValidator from rock.deployments.status import PersistedServiceStatus, ServiceStatus @@ -30,6 +29,7 @@ from rock.utils import ( ENV_POOL, DockerUtil, + ImageUtil, Timer, find_free_port, get_executor, @@ -38,6 +38,7 @@ timeout, wait_until_alive, ) +from rock.deployments.docker_client import TempAuthDockerClient, TempAuthDockerClientError __all__ = ["DockerDeployment", "DockerDeploymentConfig"] CHECK_CLEAR_INTERVAL_SECONDS = 300 @@ -56,7 +57,10 @@ def __init__( Args: **kwargs: Keyword arguments (see `DockerDeploymentConfig` for details). """ + registry_password = kwargs.pop('registry_password', None) self._config = DockerDeploymentConfig(**kwargs) + if registry_password: + self._config.registry_password = registry_password self._runtime: RemoteSandboxRuntime | None = None self._container_process = None self._runtime_timeout = 0.15 @@ -65,6 +69,7 @@ def __init__( self._check_stop_task = None self._container_name = None self._service_status = PersistedServiceStatus() + if self._config.container_name: self.set_container_name(self._config.container_name) if env_vars.ROCK_WORKER_ENV_TYPE == "docker": @@ -78,11 +83,6 @@ def __init__( else: raise Exception(f"Invalid ROCK_WORKER_ENV_TYPE: {env_vars.ROCK_WORKER_ENV_TYPE}") - if self._config.registry_username is not None and self._config.registry_password is not None: - self.add_hook( - DockerLoginHook(self._config.image, self._config.registry_username, self._config.registry_password) - ) - self.sandbox_validator: DockerSandboxValidator | None = DockerSandboxValidator() def add_hook(self, hook: DeploymentHook): @@ -235,30 +235,50 @@ def _get_rocklet_start_cmd(self) -> list[str]: ] def _pull_image(self) -> None: + """Pull image using temporary authentication. + + Uses TempAuthDockerClient to ensure credentials are isolated + and automatically cleaned up after the pull operation. + """ if self._config.pull == "never": self._service_status.update_status( phase_name="image_pull", status=Status.SUCCESS, message="skip image pull" ) return + if self._config.pull == "missing" and DockerUtil.is_image_available(self._config.image): self._service_status.update_status( phase_name="image_pull", status=Status.SUCCESS, message="use cached image, skip image pull" ) return - self._service_status.update_status(phase_name="image_pull", status=Status.RUNNING, message="image pull running") + + self._service_status.update_status( + phase_name="image_pull", status=Status.RUNNING, message="image pull running" + ) logger.info(f"Pulling image {self._config.image!r}") - self._hooks.on_custom_step(DeploymentHookStep.PULLING_IMAGE) + try: with Timer(description=f"[{self._config.image}] Image pull"): - DockerUtil.pull_image(self._config.image) + # Parse registry from image name + registry, _ = ImageUtil.parse_registry_and_others(self._config.image) + + # Create temp auth client with credentials if available + with TempAuthDockerClient( + registry=registry if self._config.registry_username else None, + username=self._config.registry_username, + password=self._config.registry_password, + ) as client: + client.pull(self._config.image) + self._service_status.update_status( phase_name="image_pull", status=Status.SUCCESS, message="image pull success" ) - except subprocess.CalledProcessError as e: - msg = f"Failed to pull image {self._config.image}. " - msg += f"Error: {e.stderr.decode()}" - msg += f"Output: {e.output.decode()}" - self._service_status.update_status(phase_name="image_pull", status=Status.FAILED, message=msg) + + except (subprocess.CalledProcessError, TempAuthDockerClientError) as e: + msg = f"Failed to pull image {self._config.image}: {e}" + self._service_status.update_status( + phase_name="image_pull", status=Status.FAILED, message=msg + ) raise DockerPullError(msg) from e @property @@ -360,6 +380,7 @@ async def start(self): self._service_status.set_sandbox_id(self._container_name) executor = get_executor() loop = asyncio.get_running_loop() + await loop.run_in_executor(executor, self._pull_image) if self._config.python_standalone_dir is not None: image_id = self._build_image() diff --git a/rock/deployments/docker_client.py b/rock/deployments/docker_client.py new file mode 100644 index 0000000000..49a4c28f55 --- /dev/null +++ b/rock/deployments/docker_client.py @@ -0,0 +1,234 @@ +""" +Docker client with temporary authentication. + +This module provides TempAuthDockerClient - a context manager for isolated Docker +authentication that uses temporary configuration directories to ensure: +1. Credentials do not pollute the user's ~/.docker/config.json +2. Credentials are completely isolated per sandbox +3. Automatic cleanup - credentials are not persistent +""" + +import logging +import shutil +import subprocess +import tempfile +from pathlib import Path + +from rock import env_vars + +logger = logging.getLogger(__name__) + + +class TempAuthDockerClientError(Exception): + """Exception raised for Docker temporary authentication errors.""" + pass + + +class TempAuthDockerClient: + """Docker client with temporary authentication configuration. + + A context manager that provides isolated Docker operations using a temporary + configuration directory. This ensures credentials are isolated per sandbox + and automatically cleaned up. + + Usage: + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ) as client: + client.pull("registry.example.com/image:v1") + if client.is_image_available("registry.example.com/image:v1"): + print("Image pulled successfully") + # Temporary directory and credentials are automatically cleaned up + + Or without credentials (for public images): + with TempAuthDockerClient() as client: + client.pull("python:3.11") + """ + + def __init__( + self, + registry: str | None = None, + username: str | None = None, + password: str | None = None, + base_dir: str | None = None, + ): + """Initialize the temporary Docker client. + + Args: + registry: Docker registry URL (e.g., registry.example.com). + Required if username/password are provided. + username: Registry username for authentication. + password: Registry password for authentication. + base_dir: Parent directory for the temporary config directory. + Defaults to ROCK_DOCKER_TEMP_AUTH_DIR env var or system temp. + """ + self._registry = registry + self._username = username + self._password = password + self._base_dir = base_dir or env_vars.ROCK_DOCKER_TEMP_AUTH_DIR + self._temp_dir: Path | None = None + self._logged_in = False + + @property + def temp_dir(self) -> Path | None: + """Get the temporary directory path.""" + return self._temp_dir + + @property + def config_path(self) -> Path | None: + """Get the temporary config.json path.""" + if self._temp_dir: + return self._temp_dir / "config.json" + return None + + def __enter__(self) -> "TempAuthDockerClient": + """Create temporary directory and optionally login to registry.""" + self._create_temp_dir() + + # Perform login if credentials are provided + if self._registry and self._username and self._password: + self._login() + + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + """Clean up temporary directory.""" + self._cleanup() + return None # Don't suppress exceptions + + def _create_temp_dir(self) -> None: + """Create the temporary configuration directory.""" + prefix = "rock_docker_auth_" + if self._base_dir: + Path(self._base_dir).mkdir(parents=True, exist_ok=True) + self._temp_dir = Path(tempfile.mkdtemp(prefix=prefix, dir=self._base_dir)) + else: + self._temp_dir = Path(tempfile.mkdtemp(prefix=prefix)) + + logger.debug(f"Created temp docker config dir: {self._temp_dir}") + + def _cleanup(self) -> None: + """Clean up the temporary configuration directory.""" + if self._temp_dir: + if self._temp_dir.exists(): + try: + shutil.rmtree(self._temp_dir, ignore_errors=True) + logger.debug(f"Cleaned up temp docker config dir: {self._temp_dir}") + except Exception as e: + logger.warning(f"Failed to cleanup temp docker config dir: {e}") + # Clear reference even if deletion failed + self._temp_dir = None + self._logged_in = False + + def _login(self) -> None: + """Login to the Docker registry.""" + if not self._temp_dir: + raise TempAuthDockerClientError("Temp dir not created. Use as context manager.") + + if not self._registry or not self._username or not self._password: + raise TempAuthDockerClientError("Registry, username, and password required for login.") + + try: + result = subprocess.run( + [ + "docker", + "--config", str(self._temp_dir), + "login", + self._registry, + "-u", self._username, + "--password-stdin" + ], + input=self._password, + capture_output=True, + text=True, + timeout=30, + ) + + if result.returncode != 0: + error_msg = f"Docker login failed: {result.stderr.strip()}" + logger.error(error_msg) + raise TempAuthDockerClientError(error_msg) + + self._logged_in = True + logger.info(f"Successfully logged in to {self._registry} using temp config") + except subprocess.TimeoutExpired: + raise TempAuthDockerClientError("Docker login timed out after 30s") + except TempAuthDockerClientError: + raise + except Exception as e: + raise TempAuthDockerClientError(f"Docker login error: {e}") + + def pull(self, image: str, timeout: int = 600) -> bytes: + """Pull a Docker image. + + Args: + image: Image name to pull (e.g., "python:3.11" or "registry.example.com/app:v1") + timeout: Timeout in seconds (default: 600s = 10 minutes) + + Returns: + Command stdout as bytes + + Raises: + TempAuthDockerClientError: If pull fails + """ + if not self._temp_dir: + raise TempAuthDockerClientError("Temp dir not created. Use as context manager.") + + try: + result = subprocess.run( + [ + "docker", + "--config", str(self._temp_dir), + "pull", + image + ], + capture_output=True, + timeout=timeout, + ) + if result.returncode != 0: + stderr = result.stderr.decode('utf-8', errors='replace') + raise TempAuthDockerClientError(f"Docker pull failed: {stderr}") + + logger.info(f"Successfully pulled image {image} using temp config") + return result.stdout + except subprocess.TimeoutExpired: + raise TempAuthDockerClientError(f"Docker pull timed out after {timeout}s") + except TempAuthDockerClientError: + raise + except Exception as e: + raise TempAuthDockerClientError(f"Docker pull error: {e}") + + def is_image_available(self, image: str) -> bool: + """Check if an image is available locally. + + Args: + image: Image name to check + + Returns: + True if image is available, False otherwise + """ + if not self._temp_dir: + return False + + try: + subprocess.check_call( + [ + "docker", + "--config", str(self._temp_dir), + "inspect", + image + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=30, + ) + return True + except subprocess.CalledProcessError: + return False + + @property + def logged_in(self) -> bool: + """Check if client has logged in to a registry.""" + return self._logged_in diff --git a/rock/env_vars.py b/rock/env_vars.py index 2a701f2651..4a1e26e512 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -25,6 +25,8 @@ ROCK_SANDBOX_EXPIRE_TIME_KEY: str | None = "expire_time" ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY: str | None = "auto_clear_time" ROCK_TIME_ZONE: str = "Asia/Shanghai" + # Docker temp auth directory + ROCK_DOCKER_TEMP_AUTH_DIR: str | None = None # Scheduler ROCK_DOCUUM_INSTALL_URL: str | None = None @@ -130,6 +132,8 @@ "ROCK_DOCUUM_INSTALL_URL": lambda: os.getenv( "ROCK_DOCUUM_INSTALL_URL", "https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh" ), + # Docker temp auth directory + "ROCK_DOCKER_TEMP_AUTH_DIR": lambda: os.getenv("ROCK_DOCKER_TEMP_AUTH_DIR"), } diff --git a/tests/integration/deployments/test_docker_auth_integration.py b/tests/integration/deployments/test_docker_auth_integration.py new file mode 100644 index 0000000000..3bb8b7f007 --- /dev/null +++ b/tests/integration/deployments/test_docker_auth_integration.py @@ -0,0 +1,255 @@ +""" +Integration tests for Docker temporary authentication in rock/deployments/docker.py + +Tests cover: +- TempAuthDockerClient integration with DockerDeployment +- _pull_image method with temp auth +- Pull behavior with various configurations +""" + +import subprocess +from unittest.mock import MagicMock, patch + +import pytest + +from rock.deployments.config import DockerDeploymentConfig +from rock.deployments.docker import DockerDeployment +from rock.utils.docker import DockerUtil +from rock.deployments.docker_client import TempAuthDockerClient, TempAuthDockerClientError +from tests.integration.conftest import SKIP_IF_NO_DOCKER + + +# Skip all tests if Docker is not available +pytestmark = [pytest.mark.need_docker, SKIP_IF_NO_DOCKER] + + +class TestDockerDeploymentPull: + """Tests for DockerDeployment._pull_image method.""" + + def test_pull_image_never_skip(self): + """Test pull=never skips image pull.""" + deployment = DockerDeployment( + image="python:3.11", + pull="never" + ) + + # Should not raise + deployment._pull_image() + + def test_pull_image_missing_with_cached_image(self): + """Test pull=missing with cached image skips pull.""" + # Use an image that's likely cached + if DockerUtil.is_image_available("python:3.11"): + deployment = DockerDeployment( + image="python:3.11", + pull="missing" + ) + + # Should not raise + deployment._pull_image() + + def test_pull_image_missing_without_cached_image(self): + """Test pull=missing without cached image attempts pull.""" + # Use a non-existent image + deployment = DockerDeployment( + image="nonexistent-local-image-xyz:latest", + pull="missing" + ) + + # Should attempt pull and fail + from rock.rocklet.exceptions import DockerPullError + with pytest.raises(DockerPullError): + deployment._pull_image() + + def test_pull_image_always_without_credentials(self): + """Test pull=always without credentials.""" + deployment = DockerDeployment( + image="python:3.11", + pull="always" + ) + + # This should work (actual pull may succeed or fail depending on network) + try: + deployment._pull_image() + except Exception: + pass # Accept any exception from actual docker pull + + @patch("rock.deployments.docker.TempAuthDockerClient") + def test_pull_image_with_registry_credentials(self, mock_client_class): + """Test pull with registry credentials.""" + mock_client = MagicMock() + mock_client_class.return_value.__enter__ = MagicMock(return_value=mock_client) + mock_client_class.return_value.__exit__ = MagicMock(return_value=None) + + deployment = DockerDeployment( + image="registry.example.com/namespace/image:v1", + registry_username="user", + registry_password="pass", + pull="always" + ) + + # This will fail because we're mocking, but we can verify the flow + try: + deployment._pull_image() + except Exception: + pass + + def test_pull_image_with_temp_auth_error(self): + """Test _pull_image handles TempAuthDockerClientError.""" + with patch("rock.deployments.docker.TempAuthDockerClient") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value.__enter__ = MagicMock(return_value=mock_client) + mock_client_class.return_value.__exit__ = MagicMock(return_value=None) + mock_client.pull.side_effect = TempAuthDockerClientError("Auth failed") + + deployment = DockerDeployment( + image="python:3.11", + registry_username="user", + registry_password="pass", + pull="always" + ) + + from rock.rocklet.exceptions import DockerPullError + with pytest.raises(DockerPullError): + deployment._pull_image() + + def test_pull_image_with_subprocess_error(self): + """Test _pull_image handles CalledProcessError.""" + with patch("rock.deployments.docker.TempAuthDockerClient") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value.__enter__ = MagicMock(return_value=mock_client) + mock_client_class.return_value.__exit__ = MagicMock(return_value=None) + mock_client.pull.side_effect = subprocess.CalledProcessError(1, "docker", stderr=b"pull failed") + + deployment = DockerDeployment( + image="python:3.11", + pull="always" + ) + + from rock.rocklet.exceptions import DockerPullError + with pytest.raises(DockerPullError): + deployment._pull_image() + + +class TestFromConfig: + """Tests for DockerDeployment.from_config().""" + + def test_from_config_preserves_password(self): + """Test from_config preserves registry_password.""" + config = DockerDeploymentConfig( + image="registry.example.com/image:v1", + registry_username="user", + registry_password="secret_password", + ) + + deployment = DockerDeployment.from_config(config) + + # Password should be preserved even though it's excluded from model_dump + assert deployment._config.registry_password == "secret_password" + + +class TestDockerDeploymentConfig: + """Tests for DockerDeploymentConfig.""" + + def test_config_registry_password_excluded(self): + """Test registry_password is excluded from model_dump.""" + config = DockerDeploymentConfig( + image="python:3.11", + registry_username="user", + registry_password="secret" + ) + + dump = config.model_dump() + assert "registry_password" not in dump + + def test_config_registry_password_stored(self): + """Test registry_password is still stored on the model.""" + config = DockerDeploymentConfig( + image="python:3.11", + registry_username="user", + registry_password="secret" + ) + + assert config.registry_password == "secret" + + +class TestIntegrationWithLocalRegistry: + """Integration tests with local Docker registry (requires Docker).""" + + @pytest.mark.asyncio + @pytest.mark.need_docker + @pytest.mark.skip(reason="Requires Docker daemon insecure-registry config for localhost") + async def test_pull_from_private_registry_with_temp_auth(self, local_registry): + """Test pulling from private registry with temp auth. + + NOTE: This test requires Docker daemon to be configured with + insecure-registries for localhost/127.0.0.1. See: + https://docs.docker.com/engine/reference/commandline/dockerd/#insecure-registries + """ + registry_url, username, password = local_registry + + # This test verifies the auth flow works with a real registry + # We're not actually pushing/pulling an image, just verifying auth works + with TempAuthDockerClient( + registry=registry_url, + username=username, + password=password + ) as client: + # If we got here, login succeeded + assert client.logged_in is True + + @pytest.mark.asyncio + @pytest.mark.need_docker + async def test_temp_auth_context_with_real_docker(self): + """Test TempAuthDockerClient context with real Docker.""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + path = client.temp_dir + + # After context, directory should be cleaned up + assert not path.exists() + + @pytest.mark.asyncio + @pytest.mark.need_docker + async def test_temp_auth_context_without_credentials(self): + """Test TempAuthDockerClient context without credentials (public images).""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert client.logged_in is False + path = client.temp_dir + + assert not path.exists() + + +class TestTempAuthDockerClientContextManager: + """Tests for TempAuthDockerClient as context manager.""" + + def test_context_manager_creates_temp_dir(self): + """Test context manager creates temporary directory.""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + path = client.temp_dir + + assert not path.exists() + + def test_context_manager_with_base_dir(self): + """Test context manager with custom base_dir.""" + import tempfile + with tempfile.TemporaryDirectory() as tmpdir: + with TempAuthDockerClient(base_dir=tmpdir) as client: + assert client.temp_dir is not None + assert str(client.temp_dir).startswith(tmpdir) + path = client.temp_dir + + assert not path.exists() + + def test_context_manager_cleanup_on_exception(self): + """Test context manager cleans up on exception.""" + with pytest.raises(ValueError): + with TempAuthDockerClient() as client: + path = client.temp_dir + raise ValueError("Test error") + + assert not path.exists() diff --git a/tests/unit/deployments/test_docker_client.py b/tests/unit/deployments/test_docker_client.py new file mode 100644 index 0000000000..5785d72766 --- /dev/null +++ b/tests/unit/deployments/test_docker_client.py @@ -0,0 +1,381 @@ +""" +Unit tests for rock/deployments/docker_client.py + +Tests cover: +- TempAuthDockerClient class: context manager, login, pull, is_image_available +- TempAuthDockerClientError exception +""" + +import subprocess +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from rock.deployments.docker_client import ( + TempAuthDockerClient, + TempAuthDockerClientError, +) + + +class TestTempAuthDockerClientInit: + """Tests for TempAuthDockerClient initialization.""" + + def test_init_default(self): + """Test default initialization.""" + client = TempAuthDockerClient() + assert client._registry is None + assert client._username is None + assert client._password is None + assert client._temp_dir is None + assert client._logged_in is False + + def test_init_with_credentials(self): + """Test initialization with credentials.""" + client = TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ) + assert client._registry == "registry.example.com" + assert client._username == "user" + assert client._password == "pass" + + def test_init_with_base_dir(self): + """Test initialization with custom base_dir.""" + client = TempAuthDockerClient(base_dir="/tmp/custom") + assert client._base_dir == "/tmp/custom" + + +class TestTempAuthDockerClientProperties: + """Tests for TempAuthDockerClient properties.""" + + def test_temp_dir_before_enter(self): + """Test temp_dir returns None before __enter__.""" + client = TempAuthDockerClient() + assert client.temp_dir is None + + def test_temp_dir_after_enter(self): + """Test temp_dir returns Path after __enter__.""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert isinstance(client.temp_dir, Path) + + def test_config_path_before_enter(self): + """Test config_path returns None before __enter__.""" + client = TempAuthDockerClient() + assert client.config_path is None + + def test_config_path_after_enter(self): + """Test config_path returns correct Path after __enter__.""" + with TempAuthDockerClient() as client: + assert client.config_path is not None + assert client.config_path.name == "config.json" + assert client.config_path.parent == client.temp_dir + + def test_logged_in_property(self): + """Test logged_in property.""" + client = TempAuthDockerClient() + assert client.logged_in is False + + +class TestTempAuthDockerClientContextManager: + """Tests for TempAuthDockerClient as context manager.""" + + def test_context_creates_and_cleans_up(self): + """Test context manager creates and cleans up temp dir.""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + path = client.temp_dir + + assert not path.exists() + + def test_context_creates_temp_dir_default_location(self): + """Test context manager creates temp dir in default location.""" + with TempAuthDockerClient() as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + assert "rock_docker_auth_" in client.temp_dir.name + + def test_context_creates_temp_dir_custom_location(self): + """Test context manager creates temp dir in custom location.""" + with tempfile.TemporaryDirectory() as tmpdir: + with TempAuthDockerClient(base_dir=tmpdir) as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + assert str(client.temp_dir).startswith(tmpdir) + + def test_context_creates_parent_directories(self): + """Test context manager creates parent directories if needed.""" + with tempfile.TemporaryDirectory() as tmpdir: + custom_path = Path(tmpdir) / "nested" / "dir" + with TempAuthDockerClient(base_dir=str(custom_path)) as client: + assert client.temp_dir is not None + assert client.temp_dir.exists() + assert custom_path.exists() + + def test_context_cleanup_on_exception(self): + """Test context manager cleans up on exception.""" + with pytest.raises(ValueError): + with TempAuthDockerClient() as client: + path = client.temp_dir + raise ValueError("Test error") + + assert not path.exists() + + def test_context_cleanup_safe_when_directory_deleted(self): + """Test cleanup is safe when directory was already deleted.""" + with TempAuthDockerClient() as client: + path = client.temp_dir + # Directory exists + assert path.exists() + # After context, directory should be cleaned up + assert not path.exists() + + +class TestTempAuthDockerClientLogin: + """Tests for TempAuthDockerClient login functionality.""" + + def test_login_with_credentials_on_enter(self): + """Test login is called automatically when credentials provided.""" + with patch.object(TempAuthDockerClient, '_login') as mock_login: + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ) as client: + pass + + mock_login.assert_called_once() + + def test_login_without_credentials(self): + """Test no login when credentials not provided.""" + with TempAuthDockerClient() as client: + assert client._logged_in is False + + @patch("subprocess.run") + def test_login_success(self, mock_run): + """Test successful login.""" + mock_run.return_value = MagicMock(returncode=0, stderr="") + + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="password" + ) as client: + assert client._logged_in is True + mock_run.assert_called() + call_args = mock_run.call_args + assert "docker" in call_args[0][0] + assert "login" in call_args[0][0] + assert "registry.example.com" in call_args[0][0] + + @patch("subprocess.run") + def test_login_failure(self, mock_run): + """Test login failure with non-zero return code.""" + mock_run.return_value = MagicMock( + returncode=1, + stderr="Error: authentication failed" + ) + + with pytest.raises(TempAuthDockerClientError) as exc_info: + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="wrongpass" + ): + pass + + assert "Docker login failed" in str(exc_info.value) + + @patch("subprocess.run") + def test_login_timeout(self, mock_run): + """Test login timeout.""" + mock_run.side_effect = subprocess.TimeoutExpired(cmd="docker", timeout=30) + + with pytest.raises(TempAuthDockerClientError) as exc_info: + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ): + pass + + assert "timed out" in str(exc_info.value) + + @patch("subprocess.run") + def test_login_unexpected_error(self, mock_run): + """Test login with unexpected error.""" + mock_run.side_effect = OSError("Unexpected error") + + with pytest.raises(TempAuthDockerClientError) as exc_info: + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ): + pass + + assert "Docker login error" in str(exc_info.value) + + +class TestTempAuthDockerClientPull: + """Tests for TempAuthDockerClient.pull().""" + + def test_pull_without_context_raises(self): + """Test pull() raises error if not in context.""" + client = TempAuthDockerClient() + with pytest.raises(TempAuthDockerClientError) as exc_info: + client.pull("registry.example.com/image:v1") + assert "Temp dir not created" in str(exc_info.value) + + @patch("subprocess.run") + def test_pull_success(self, mock_run): + """Test successful pull.""" + mock_run.return_value = MagicMock( + returncode=0, + stdout=b"pulled successfully", + stderr=b"" + ) + + with TempAuthDockerClient() as client: + result = client.pull("python:3.11") + + assert result == b"pulled successfully" + mock_run.assert_called() + call_args = mock_run.call_args + assert "docker" in call_args[0][0] + assert "pull" in call_args[0][0] + assert "python:3.11" in call_args[0][0] + + @patch("subprocess.run") + def test_pull_failure(self, mock_run): + """Test pull failure with non-zero return code.""" + mock_run.return_value = MagicMock( + returncode=1, + stdout=b"", + stderr=b"Error: image not found" + ) + + with TempAuthDockerClient() as client: + with pytest.raises(TempAuthDockerClientError) as exc_info: + client.pull("nonexistent/image:v1") + + assert "Docker pull failed" in str(exc_info.value) + + @patch("subprocess.run") + def test_pull_timeout(self, mock_run): + """Test pull timeout.""" + mock_run.side_effect = subprocess.TimeoutExpired(cmd="docker", timeout=600) + + with TempAuthDockerClient() as client: + with pytest.raises(TempAuthDockerClientError) as exc_info: + client.pull("large-image:v1") + + assert "timed out" in str(exc_info.value) + + @patch("subprocess.run") + def test_pull_unexpected_error(self, mock_run): + """Test pull with unexpected error.""" + mock_run.side_effect = OSError("Unexpected error") + + with TempAuthDockerClient() as client: + with pytest.raises(TempAuthDockerClientError) as exc_info: + client.pull("image:v1") + + assert "Docker pull error" in str(exc_info.value) + + +class TestTempAuthDockerClientIsImageAvailable: + """Tests for TempAuthDockerClient.is_image_available().""" + + def test_is_image_available_without_context(self): + """Test is_image_available returns False if not in context.""" + client = TempAuthDockerClient() + assert client.is_image_available("python:3.11") is False + + @patch("subprocess.check_call") + def test_is_image_available_true(self, mock_check_call): + """Test is_image_available returns True for existing image.""" + mock_check_call.return_value = 0 + + with TempAuthDockerClient() as client: + result = client.is_image_available("python:3.11") + + assert result is True + mock_check_call.assert_called() + call_args = mock_check_call.call_args + assert "docker" in call_args[0][0] + assert "inspect" in call_args[0][0] + assert "python:3.11" in call_args[0][0] + + @patch("subprocess.check_call") + def test_is_image_available_false(self, mock_check_call): + """Test is_image_available returns False for non-existing image.""" + mock_check_call.side_effect = subprocess.CalledProcessError(1, "docker") + + with TempAuthDockerClient() as client: + result = client.is_image_available("nonexistent:v1") + + assert result is False + + +class TestTempAuthDockerClientError: + """Tests for TempAuthDockerClientError exception.""" + + def test_error_is_exception(self): + """Test TempAuthDockerClientError is an Exception.""" + assert issubclass(TempAuthDockerClientError, Exception) + + def test_error_message(self): + """Test TempAuthDockerClientError preserves message.""" + error = TempAuthDockerClientError("Test error message") + assert str(error) == "Test error message" + + def test_error_can_be_raised_and_caught(self): + """Test TempAuthDockerClientError can be raised and caught.""" + with pytest.raises(TempAuthDockerClientError): + raise TempAuthDockerClientError("Test error") + + +class TestIntegration: + """Integration tests that verify method interactions.""" + + @patch("subprocess.run") + def test_full_lifecycle_with_credentials(self, mock_run): + """Test full lifecycle: enter -> login -> pull -> exit.""" + mock_run.return_value = MagicMock( + returncode=0, + stdout=b"success", + stderr=b"" + ) + + with TempAuthDockerClient( + registry="registry.example.com", + username="user", + password="pass" + ) as client: + # Login called on enter + assert client._logged_in is True + + # Pull + client.pull("registry.example.com/image:v1") + assert mock_run.call_count == 2 + + # After exit, temp_dir should be None + assert client._temp_dir is None + + def test_multiple_context_cycles(self): + """Test multiple context cycles work correctly.""" + client = TempAuthDockerClient() + + for i in range(3): + with client: + assert client.temp_dir is not None + path = client.temp_dir + assert path.exists() + + assert not path.exists() + assert client._temp_dir is None From 41bd51f9e21f0347fc69312fc2a3ce536e4a7c54 Mon Sep 17 00:00:00 2001 From: dengwx Date: Mon, 20 Apr 2026 17:44:17 +0800 Subject: [PATCH 054/226] docs: replace News with Updates section, trim to 5 latest versions (#850) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove stale ## 📢 News section - Rename ## 📢 Latest Updates → ## 📢 Updates - Switch to two-column table so release names left-align at fixed position - Keep only 5 most recent versions; mark newest with [Latest] refs #849 Co-authored-by: dengwx Co-authored-by: Claude Sonnet 4.6 --- README.md | 29 +++++++++-------------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index c67b7a8e06..f329c8a881 100644 --- a/README.md +++ b/README.md @@ -25,14 +25,6 @@ ROCK (Reinforcement Open Construction Kit) is a easy-to-use, and scalable sandbo ROCK adopts a client-server architecture, supports different levels of isolation mechanisms to ensure stable environment operation, and supports integration with various reinforcement learning training frameworks through SDK. ROCK not only supports traditional sandbox management functions but also is compatible with GEM-like protocols, providing standardized interfaces for reinforcement learning environments. ---- -## 📢 News -| 📣 Update | -|:--| -| **[03/03/2026]** 🎉 ROCK v1.3.0 Released! K8s Operator, Docker registry login, Kata runtime, SWE-bench evaluation demo, and more. | -| **[02/28/2026]** 🎉 ROCK v1.2.5 Released! Custom metrics endpoint, user-defined metric tags, and Aliyun MSE Nacos support. | -| **[01/01/2026]** 🎉 Our [Let It Flow: Agentic Crafting on Rock and Roll](https://arxiv.org/abs/2512.24873) report released! Introducing ALE ecosystem and ROME, an open-source agentic model with novel IPA algorithm. | ---- ## 🚀 Get Started [Documents](https://alibaba.github.io/ROCK/) @@ -159,18 +151,15 @@ if __name__ == "__main__": --- -## 📢 Latest Updates - -| 📣 Update Content | -|:-----------| -| **[2026-04-16]** 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | -| **[2026-04-10]** 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | -| **[2026-03-27]** 🎉 ROCK v1.4.4 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.4) | -| **[2026-03-24]** 🎉 ROCK v1.4.3 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.3) | -| **[2026-03-17]** 🎉 ROCK v1.4.2 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.2) | -| **[2026-03-14]** 🎉 ROCK v1.4.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.1) | -| **[2026-03-14]** 🎉 ROCK v1.4.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.0) | -| **[2026-03-03]** 🎉 ROCK v1.3.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.3.0) | +## 📢 Updates + +| Date | Release | +|:-------------|:---| +| **[Latest]** | 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | +| **[2026-04-10]** | 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | +| **[2026-03-27]** | 🎉 ROCK v1.4.4 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.4) | +| **[2026-03-24]** | 🎉 ROCK v1.4.3 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.3) | +| **[2026-03-17]** | 🎉 ROCK v1.4.2 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.2) | --- From 15a1799f6da162933d6cf41ce82eb8efe72f18e7 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 21 Apr 2026 09:33:33 +0800 Subject: [PATCH 055/226] [Chore] database connection unittest and parameter optimizations (#852) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(tests): add missing need_docker marker to standard_spec tests test_use_standard_spec_only and test_use_standard_spec_only_disabled call start_async(DockerDeploymentConfig) which requires Docker daemon, but were only marked need_ray — causing failures in no-Docker environments. * chore(tests): unify db fixture to sqlite in-memory in test_sandbox_meta_store When database.url is not configured, admin falls back to sqlite-memory. Testing both file-based sqlite (db) and in-memory sqlite (db_memory) separately is redundant — they behave identically. Consolidate into a single sqlite-memory db fixture, remove db_memory / repo_with_memory_db and the duplicate test_iter_alive_sandbox_ids_works_with_sqlite_memory. * chore(db): increase asyncpg pool limits for sandbox meta store Tune DatabaseProvider asyncpg engine pool settings to pool_size=100, and pool_timeout=120 so sandbox meta store can better tolerate higher concurrent DB load without premature timeout. --- rock/admin/core/db_provider.py | 10 ++++-- tests/unit/sandbox/test_sandbox_manager.py | 2 ++ tests/unit/sandbox/test_sandbox_meta_store.py | 35 +------------------ 3 files changed, 11 insertions(+), 36 deletions(-) diff --git a/rock/admin/core/db_provider.py b/rock/admin/core/db_provider.py index 7d994dfb0d..539652aeea 100644 --- a/rock/admin/core/db_provider.py +++ b/rock/admin/core/db_provider.py @@ -37,8 +37,14 @@ async def init(self) -> None: For asyncpg, ``statement_cache_size=0`` prevents ``InvalidCachedStatementError`` after external DDL changes """ - connect_args = {"statement_cache_size": 0} if "asyncpg" in self._url else {} - self._engine = create_async_engine(self._url, echo=False, connect_args=connect_args) + engine_kwargs: dict[str, object] = {"echo": False} + if "asyncpg" in self._url: + engine_kwargs["connect_args"] = {"statement_cache_size": 0} + engine_kwargs["pool_size"] = 100 + engine_kwargs["max_overflow"] = 0 + engine_kwargs["pool_timeout"] = 120 + + self._engine = create_async_engine(self._url, **engine_kwargs) async def create_tables(self) -> None: """Create all tables defined in Base.metadata (idempotent).""" diff --git a/tests/unit/sandbox/test_sandbox_manager.py b/tests/unit/sandbox/test_sandbox_manager.py index 91afcd2af2..63a7c5e929 100644 --- a/tests/unit/sandbox/test_sandbox_manager.py +++ b/tests/unit/sandbox/test_sandbox_manager.py @@ -192,6 +192,7 @@ async def test_get_actor_not_exist_raises_value_error(sandbox_manager): @pytest.mark.need_ray +@pytest.mark.need_docker @pytest.mark.asyncio async def test_use_standard_spec_only(sandbox_manager): """Test that use_standard_spec_only forces sandbox to use standard spec.""" @@ -233,6 +234,7 @@ async def test_use_standard_spec_only(sandbox_manager): @pytest.mark.need_ray +@pytest.mark.need_docker @pytest.mark.asyncio async def test_use_standard_spec_only_disabled(sandbox_manager): """Test that sandbox uses requested spec when use_standard_spec_only is disabled.""" diff --git a/tests/unit/sandbox/test_sandbox_meta_store.py b/tests/unit/sandbox/test_sandbox_meta_store.py index 2b2baf0a0b..110a590f86 100644 --- a/tests/unit/sandbox/test_sandbox_meta_store.py +++ b/tests/unit/sandbox/test_sandbox_meta_store.py @@ -29,17 +29,7 @@ async def redis(): @pytest.fixture -async def db(tmp_path): - provider = DatabaseProvider(db_config=DatabaseConfig(url=f"sqlite:///{tmp_path / 'test.db'}")) - await provider.init() - await provider.create_tables() - table = SandboxTable(provider) - yield table - await provider.close() - - -@pytest.fixture -async def db_memory(): +async def db(): provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite:///:memory:")) await provider.init() await provider.create_tables() @@ -53,11 +43,6 @@ def repo(redis, db): return SandboxMetaStore(redis_provider=redis, sandbox_table=db) -@pytest.fixture -def repo_with_memory_db(redis, db_memory): - return SandboxMetaStore(redis_provider=redis, sandbox_table=db_memory) - - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -264,24 +249,6 @@ async def test_iter_alive_sandbox_ids_excludes_stopped(self, repo): assert "sbx-running" in ids assert "sbx-stopped" not in ids - async def test_iter_alive_sandbox_ids_works_with_sqlite_memory(self, repo_with_memory_db): - """iter_alive_sandbox_ids() should work with sqlite in-memory DB + Redis fallback.""" - await repo_with_memory_db.create( - "sbx-running", {**SANDBOX_INFO, "sandbox_id": "sbx-running", "state": State.RUNNING} - ) - await repo_with_memory_db.create( - "sbx-pending", {**SANDBOX_INFO, "sandbox_id": "sbx-pending", "state": State.PENDING} - ) - await repo_with_memory_db.create( - "sbx-stopped", {**SANDBOX_INFO, "sandbox_id": "sbx-stopped", "state": "stopped"} - ) - await asyncio.sleep(0.1) - - ids = {sid async for sid in repo_with_memory_db.iter_alive_sandbox_ids()} - assert "sbx-running" in ids - assert "sbx-pending" in ids - assert "sbx-stopped" not in ids - async def test_iter_alive_sandbox_ids_consistent_with_redis_scan(self, repo, redis): """DB list_by_in(state IN active_states) should be consistent with Redis alive-key scan. From a375d7f42c14e5167677e53fd65f00d3ea10da55 Mon Sep 17 00:00:00 2001 From: dengwx Date: Tue, 21 Apr 2026 11:32:09 +0800 Subject: [PATCH 056/226] feat(datasets): add datasets SDK and CLI commands for OSS registry (#859) * feat: add docs * feat: update docs * docs(datasets): update design doc to reuse OssRegistryInfo from bench Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): add models and package skeleton Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): add BaseDatasetRegistry ABC Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): add OssDatasetRegistry with list_datasets and upload_dataset Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): add DatasetClient Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): extend ConfigManager with [dataset] section Co-Authored-By: Claude Sonnet 4.6 * feat(datasets): add DatasetsCommand (list + upload) * feat: refine DatasetSpec * fix: test fail --------- Co-authored-by: dengwx Co-authored-by: Claude Sonnet 4.6 --- docs/dev/envhub/README.md | 395 ++++++++++++++++++ rock/cli/command/datasets.py | 114 +++++ rock/cli/config.py | 20 + rock/sdk/envhub/datasets/__init__.py | 4 + rock/sdk/envhub/datasets/client.py | 20 + rock/sdk/envhub/datasets/models.py | 17 + rock/sdk/envhub/datasets/registry/__init__.py | 0 rock/sdk/envhub/datasets/registry/base.py | 22 + rock/sdk/envhub/datasets/registry/oss.py | 143 +++++++ tests/unit/datasets/__init__.py | 0 tests/unit/datasets/test_client.py | 33 ++ tests/unit/datasets/test_config.py | 44 ++ tests/unit/datasets/test_datasets_command.py | 78 ++++ tests/unit/datasets/test_models.py | 14 + tests/unit/datasets/test_oss_registry.py | 170 ++++++++ 15 files changed, 1074 insertions(+) create mode 100644 docs/dev/envhub/README.md create mode 100644 rock/cli/command/datasets.py create mode 100644 rock/sdk/envhub/datasets/__init__.py create mode 100644 rock/sdk/envhub/datasets/client.py create mode 100644 rock/sdk/envhub/datasets/models.py create mode 100644 rock/sdk/envhub/datasets/registry/__init__.py create mode 100644 rock/sdk/envhub/datasets/registry/base.py create mode 100644 rock/sdk/envhub/datasets/registry/oss.py create mode 100644 tests/unit/datasets/__init__.py create mode 100644 tests/unit/datasets/test_client.py create mode 100644 tests/unit/datasets/test_config.py create mode 100644 tests/unit/datasets/test_datasets_command.py create mode 100644 tests/unit/datasets/test_models.py create mode 100644 tests/unit/datasets/test_oss_registry.py diff --git a/docs/dev/envhub/README.md b/docs/dev/envhub/README.md new file mode 100644 index 0000000000..6255efe1f6 --- /dev/null +++ b/docs/dev/envhub/README.md @@ -0,0 +1,395 @@ +# EnvHub & Dataset 开发设计文档 + +## 目录 + +- [EnvHub（现有）](#envhub现有) +- [Dataset 功能设计](#dataset-功能设计) + - [背景与目标](#背景与目标) + - [OSS 路径约定](#oss-路径约定) + - [模块结构](#模块结构) + - [核心模型](#核心模型) + - [Registry 抽象层](#registry-抽象层) + - [OssDatasetRegistry](#ossdatasetregistry) + - [DatasetClient](#datasetclient) + - [CLI 命令](#cli-命令) + - [配置文件扩展](#配置文件扩展) + - [数据流](#数据流) + - [错误处理](#错误处理) + - [测试策略](#测试策略) + +--- + +## EnvHub（现有） + +EnvHub 是 ROCK 的环境管理服务，提供 Docker 环境的注册、查询、列举和删除功能。 + +| 入口 | 模块 | 说明 | +|----------------|-----------------------------|------------------------------| +| `envhub` 服务 | `rock.envhub.server` | FastAPI 服务，端口 8081 | +| SDK Client | `rock.sdk.envhub.client` | `EnvHubClient`，HTTP 调用服务 | + +REST 端点：`POST /env/register`、`POST /env/get`、`POST /env/list`、`POST /env/delete`、`GET /health` + +--- + +## Dataset 功能设计 + +### 背景与目标 + +在 ROCK 中引入 dataset 管理能力，核心目标如下： + +**1. 约束 datasets** + +统一 dataset 的路径约定、命名规范和存储格式，避免各服务自行散落地写 OSS 路径。所有 dataset 必须遵循 `datasets/{organization}/{dataset_name}/{split}/{task_id}/` 结构，由本模块作为唯一入口强制执行。 + +**2. 提供 SDK 和 CLI 供其他服务集成** + +- **SDK**（`rock.sdk.envhub.datasets`）：提供 `DatasetClient`，供 Python 代码直接调用 list / upload，适合 admin、job 等服务在流程中集成 dataset 操作。 +- **CLI**（`rock datasets`）：提供 `list` 和 `upload` 子命令，供运维、研究人员在终端操作 dataset，也适合脚本化批量处理。 + +**3. 为后续权限管理预留扩展点** + +当前阶段 CLI 直接对接 OSS，不经过 envhub server。后续可在 envhub server 增加 `/datasets/*` 端点，在 SDK/CLI 与 OSS 之间插入权限校验、审计日志等能力，Registry 抽象层的设计为此预留了扩展空间。 + +--- + +### OSS 路径约定 + +路径层级设计对齐 **HuggingFace Datasets** 的命名惯例（`{organization}/{dataset_name}/{split}`），在此基础上增加了 ROCK 特有的 `{task_id}` 层来组织结构化的 benchmark task 目录。 + +``` +oss://{bucket}/datasets/{organization}/{dataset_name}/{split}/{task_id}/ +``` + +| 层级 | 说明 | 类比 HuggingFace | +|------|------|-----------------| +| `organization` | 数据集所属组织，如 `qwen`、`alibaba` | HF namespace（`qwen/`） | +| `dataset_name` | 数据集名称 | HF repo name（`my-bench`） | +| `split` | 分片标识，如 `train`、`test`、`v1.0` | HF split（`train`/`test`） | +| `task_id` | 单个 task 目录名（ROCK 特有） | HF 无此层，HF 直接存数据文件 | + +示例： + +``` +oss://my-bucket/ +└── datasets/ + └── qwen/ # organization + └── my-bench/ # dataset_name + └── train/ # split + ├── task-001/ # task_id（ROCK 特有） + │ ├── task.toml + │ └── tests/ + └── task-002/ + ├── task.toml + └── tests/ +``` + +task 目录内的文件结构原样保留（相对路径不变），上传和下载均以 `task_id/` 为单位。 + +--- + +### 模块结构 + +``` +rock/ +├── sdk/ +│ └── envhub/ +│ ├── client.py # 现有：EnvHubClient +│ ├── config.py # 现有 +│ ├── schema.py # 现有 +│ └── datasets/ # 新增 +│ ├── __init__.py # 对外入口：暴露 DatasetClient、DatasetSpec 等 +│ ├── models.py # DatasetSpec, UploadResult（OssRegistryInfo 复用自 bench） +│ ├── client.py # DatasetClient（对外统一入口） +│ └── registry/ +│ ├── __init__.py +│ ├── base.py # BaseDatasetRegistry ABC +│ └── oss.py # OssDatasetRegistry +└── cli/ + └── command/ + └── datasets.py # DatasetsCommand（继承 Command ABC） +``` + +--- + +### 核心模型 + +**复用模型（来自 `rock.sdk.bench.models.job.config`，不新增、不移动）** + +| 模型 | 关键字段 | 用途 | +|------|----------|------| +| `OssRegistryInfo` | `oss_bucket`, `oss_endpoint`, `oss_region`, `oss_access_key_id`, `oss_access_key_secret`, `oss_dataset_path` | OSS 连接凭证与路径前缀（`oss_dataset_path` 默认 `"datasets"`） | +| `LocalDatasetConfig` | `path: Path` | 本地 task 目录（upload 数据源） | +| `RegistryDatasetConfig` | `name="org/dataset_name"`, `version=split`, `overwrite`, `registry=OssRegistryInfo(...)` | 远端数据集引用（upload 目标） | + +`RegistryDatasetConfig.name` 遵循 HuggingFace 惯例，使用 `"{organization}/{dataset_name}"` 格式；`OssDatasetRegistry` 通过 `name.split("/", 1)` 拆分得到 org 和 name。`version` 对应 split（如 `"train"`、`"test"`）。`DatasetSpec.id` / `UploadResult.id` 同样使用此格式，对齐 HF `DatasetInfo.id`。 + +**新增模型（`rock/sdk/envhub/datasets/models.py`）** + +```python +@dataclass +class DatasetSpec: + id: str # "{organization}/{dataset_name}"，对齐 HF DatasetInfo.id，如 "princeton-nlp/SWE-bench_Verified" + split: str + task_ids: list[str] + +@dataclass +class UploadResult: + id: str # "{organization}/{dataset_name}" + split: str + uploaded: int # 成功上传的文件数 + skipped: int # 已存在跳过的 task 数（overwrite=False） + failed: int # 失败数 +``` + +--- + +### Registry 抽象层 + +**`rock/sdk/envhub/datasets/registry/base.py`** + +```python +class BaseDatasetRegistry(ABC): + + @abstractmethod + def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: + """枚举 registry 中的所有 datasets。 + organization 不为 None 时只返回该 org 下的 datasets。 + """ + ... + + @abstractmethod + def upload_dataset( + self, + source: LocalDatasetConfig, + target: RegistryDatasetConfig, + concurrency: int = 4, + ) -> UploadResult: + """将 source.path/{task_id}/ 批量上传到 target 指定的远端路径。 + org/name/split/overwrite 均从 target 提取。 + """ + ... +``` + +--- + +### OssDatasetRegistry + +**`rock/sdk/envhub/datasets/registry/oss.py`** + +```python +class OssDatasetRegistry(BaseDatasetRegistry): + def __init__(self, registry: OssRegistryInfo): ... +``` + +**路径构建：** + +```python +def _build_prefix(self, org: str, name: str, split: str | None = None) -> str: + base = self._registry.oss_dataset_path or "datasets" + parts = [base, org, name] + if split: + parts.append(split) + return "/".join(parts) +# → "datasets/qwen/my-bench/train" +``` + +**list_datasets 逻辑：** + +1. 以 `datasets/` 为前缀列出三级目录（org → name → split） +2. 对每个 `datasets/{org}/{name}/{split}/`，列出直接子目录作为 `task_ids` +3. 返回 `list[DatasetSpec]` + +OSS 列举使用 `list_objects_v2` with `delimiter="/"` 逐层枚举目录，避免全量遍历。 + +**upload_dataset 逻辑：** + +1. 从 `target.name.split("/", 1)` 提取 `org` 和 `name`；`target.version` 为 `split`；`target.overwrite` 为覆盖标志 +2. 遍历 `source.path` 下的一级子目录，每个子目录视为一个 task（`task_id = subdir.name`） +3. 若 `target.overwrite=False` 且 OSS 上已存在该 task 目录，跳过 +4. 并发上传（`ThreadPoolExecutor`，`concurrency` 控制并发数） +5. 目标 key：`datasets/{org}/{name}/{split}/{task_id}/{relative_file_path}` +6. 返回 `UploadResult` + +--- + +### DatasetClient + +**`rock/sdk/envhub/datasets/client.py`** + +薄封装层，负责从配置创建 registry 并提供面向业务的方法。 + +```python +class DatasetClient: + def __init__(self, registry: OssRegistryInfo): + self._registry = OssDatasetRegistry(registry) + + def list_datasets(self, org: str | None = None) -> list[DatasetSpec]: + return self._registry.list_datasets(org) + + def upload_dataset( + self, + source: LocalDatasetConfig, + target: RegistryDatasetConfig, + concurrency: int = 4, + ) -> UploadResult: + return self._registry.upload_dataset(source, target, concurrency) +``` + +--- + +### CLI 命令 + +**`rock/cli/command/datasets.py`**，继承 `Command` ABC，`name = "datasets"`。 + +#### rock datasets list + +``` +rock datasets list [OPTIONS] + +Options: + --org TEXT 只列出指定 organization 的 datasets + --bucket TEXT OSS bucket 名称（覆盖 config.ini） + --endpoint TEXT OSS endpoint（覆盖 config.ini） + --access-key-id TEXT OSS access key ID（覆盖 config.ini） + --access-key-secret TEXT OSS access key secret（覆盖 config.ini） +``` + +输出示例： + +``` +Dataset Split Tasks +qwen/my-bench train 42 +qwen/my-bench test 10 +alibaba/code-eval train 100 +``` + +#### rock datasets upload + +``` +rock datasets upload [OPTIONS] + +Required: + --org TEXT Organization 名称 + --dataset TEXT Dataset 名称 + --split TEXT Split 名称（如 train、test、v1.0） + --dir PATH 本地 task 目录（内含 {task_id}/ 子目录） + +Options: + --bucket TEXT OSS bucket（覆盖 config.ini） + --endpoint TEXT OSS endpoint（覆盖 config.ini） + --access-key-id TEXT OSS access key ID（覆盖 config.ini） + --access-key-secret TEXT OSS access key secret（覆盖 config.ini） + --concurrency INT 并发上传数（默认 4，范围 1-16） + --overwrite 覆盖 OSS 上已存在的 task 目录（默认跳过） +``` + +输出示例： + +``` +Uploading to oss://my-bucket/datasets/qwen/my-bench/train/ + ✓ task-001 (5 files) + ✓ task-002 (5 files) + - task-003 skipped (already exists) + +Done: 2 uploaded, 1 skipped, 0 failed +``` + +--- + +### 配置文件扩展 + +在 `.rock/config.ini` 新增 `[dataset]` section，用于存储 OSS 凭证默认值： + +```ini +[rock] +base_url = http://localhost:8080 + +[dataset] +oss_bucket = my-bucket +oss_endpoint = https://oss-cn-hangzhou.aliyuncs.com +oss_access_key_id = LTAI5t... +oss_access_key_secret = xxxxxxx +``` + +**优先级（高→低）**：CLI 参数 > `config.ini [dataset]` section > 报错（必填项缺失） + +`ConfigManager` 扩展：在 `CLIConfig` 新增 `dataset_config: DatasetConfig` 字段（内部结构体），读取 `[dataset]` section 中的 OSS 凭证。`DatasetsCommand` 在初始化时合并 `DatasetConfig` + CLI args 构建 `OssRegistryInfo`（来自 `rock.sdk.bench.models.job.config`）。 + +--- + +### 数据流 + +**list:** + +``` +rock datasets list --org qwen + └─ DatasetCommand.list() + ├─ ConfigManager.get_dataset_config() # 读 config.ini [dataset] + ├─ 合并 CLI 参数 → OssRegistryInfo + ├─ DatasetClient(config) + └─ OssDatasetRegistry.list_datasets(org="qwen") + └─ alibabacloud_oss_v2: list_objects_v2(prefix="datasets/qwen/", delimiter="/") + → 枚举 name/split 层 + → 构建 DatasetSpec 列表 + → 打印表格 +``` + +**upload:** + +``` +rock datasets upload --org qwen --dataset my-bench --split train --dir ./tasks/ + └─ DatasetsCommand.upload() + ├─ ConfigManager.get_config().dataset_config # 读 config.ini [dataset] + ├─ 合并 CLI 参数 → OssRegistryInfo + ├─ source = LocalDatasetConfig(path=./tasks/) + ├─ target = RegistryDatasetConfig( + │ name="qwen/my-bench", version="train", + │ overwrite=False, registry=OssRegistryInfo) + ├─ DatasetClient(registry=OssRegistryInfo) + └─ OssDatasetRegistry.upload_dataset(source, target, concurrency=4) + ├─ target.name.split("/", 1) → org="qwen", name="my-bench" + ├─ target.version → split="train" + ├─ 遍历 source.path 下子目录：task-001/, task-002/, ... + ├─ ThreadPoolExecutor(max_workers=concurrency) + └─ 每个 task： + ├─ 若 target.overwrite=False 且 OSS 存在 → skip + └─ 遍历文件 → PutObject(key="datasets/qwen/my-bench/train/task-001/{file}") +``` + +--- + +### 错误处理 + +| 场景 | 行为 | +|------|------| +| OSS 凭证缺失 | 启动时立即报错，提示配置 `[dataset]` section 或传 CLI 参数 | +| OSS 权限错误（401/403） | 立即抛出，打印明确错误信息，不重试 | +| OSS 网络错误（5xx/timeout） | 指数退避重试（最多 3 次），超限后报错 | +| `--dir` 不存在或为空 | 命令入口检查，立即报错 | +| 单个 task 上传失败 | 记录到 `UploadResult.failed`，继续上传其他 tasks，命令结束后汇总报告 | +| `--org`/`--dataset`/`--split` 缺失 | argparse required 校验，自动报错 | + +--- + +### 测试策略 + +| 测试类型 | 覆盖范围 | 标记 | +|----------|----------|------| +| 单元测试 | `OssDatasetRegistry` 路径构建、`DatasetSpec` 模型、`ConfigManager` 解析 `[dataset]` | 无特殊标记 | +| 集成测试（mock OSS） | `list_datasets`、`upload_dataset` 逻辑，使用 `unittest.mock` mock OSS SDK | `@pytest.mark.integration` | +| 集成测试（真实 OSS） | 端到端 upload → list 验证 | `@pytest.mark.need_admin`（需要 OSS 凭证） | + +测试文件位置： + +``` +tests/ +├── unit/ +│ └── datasets/ +│ ├── test_models.py +│ ├── test_oss_registry.py +│ └── test_config.py +└── integration/ + └── datasets/ + └── test_oss_e2e.py +``` diff --git a/rock/cli/command/datasets.py b/rock/cli/command/datasets.py new file mode 100644 index 0000000000..becf194f59 --- /dev/null +++ b/rock/cli/command/datasets.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from rock.cli.command.command import Command +from rock.cli.config import ConfigManager +from rock.logger import init_logger +from rock.sdk.bench.models.job.config import LocalDatasetConfig, OssRegistryInfo, RegistryDatasetConfig +from rock.sdk.envhub.datasets.client import DatasetClient + +logger = init_logger(__name__) + + +class DatasetsCommand(Command): + name = "datasets" + + async def arun(self, args: argparse.Namespace) -> None: + if args.datasets_command == "list": + await self._list(args) + elif args.datasets_command == "upload": + await self._upload(args) + else: + raise ValueError(f"Unknown datasets command: {args.datasets_command}") + + def _build_oss_registry_info(self, args: argparse.Namespace) -> OssRegistryInfo: + ds_cfg = ConfigManager(Path(args.config) if args.config else None).get_config().dataset_config + + bucket = getattr(args, "bucket", None) or ds_cfg.oss_bucket + if not bucket: + raise ValueError( + "OSS bucket is required. Pass --bucket or set 'oss_bucket' in [dataset] section of config.ini." + ) + return OssRegistryInfo( + oss_bucket=bucket, + oss_endpoint=getattr(args, "endpoint", None) or ds_cfg.oss_endpoint, + oss_access_key_id=getattr(args, "access_key_id", None) or ds_cfg.oss_access_key_id, + oss_access_key_secret=getattr(args, "access_key_secret", None) or ds_cfg.oss_access_key_secret, + oss_region=getattr(args, "region", None) or ds_cfg.oss_region, + ) + + async def _list(self, args: argparse.Namespace) -> None: + registry_info = self._build_oss_registry_info(args) + client = DatasetClient(registry_info) + datasets = client.list_datasets(org=getattr(args, "org", None)) + + if not datasets: + print("No datasets found.") + return + + col_id = max(len("Dataset"), max(len(d.id) for d in datasets)) + col_split = max(len("Split"), max(len(d.split) for d in datasets)) + + header = f"{'Dataset':<{col_id}} {'Split':<{col_split}} {'Tasks':>6}" + print(header) + print("-" * len(header)) + for ds in sorted(datasets, key=lambda d: (d.id, d.split)): + print(f"{ds.id:<{col_id}} {ds.split:<{col_split}} {len(ds.task_ids):>6}") + + async def _upload(self, args: argparse.Namespace) -> None: + local_dir = Path(args.dir) + if not local_dir.is_dir(): + raise ValueError(f"--dir '{local_dir}' does not exist or is not a directory") + + registry_info = self._build_oss_registry_info(args) + source = LocalDatasetConfig(path=local_dir) + target = RegistryDatasetConfig( + name=f"{args.org}/{args.dataset}", + version=args.split, + overwrite=args.overwrite, + registry=registry_info, + ) + + base = registry_info.oss_dataset_path or "datasets" + print(f"Uploading to oss://{registry_info.oss_bucket}/{base}/{args.org}/{args.dataset}/{args.split}/") + + client = DatasetClient(registry_info) + result = client.upload_dataset(source, target, concurrency=args.concurrency) + + print(f"\nDone: {result.uploaded} uploaded, {result.skipped} skipped, {result.failed} failed") + if result.failed > 0: + sys.exit(1) + + @staticmethod + async def add_parser_to(subparsers: argparse._SubParsersAction) -> None: + datasets_parser = subparsers.add_parser("datasets", description="Dataset operations on OSS") + datasets_subparsers = datasets_parser.add_subparsers(dest="datasets_command") + + def add_oss_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--bucket", help="OSS bucket name (overrides config.ini)") + parser.add_argument("--endpoint", help="OSS endpoint URL (overrides config.ini)") + parser.add_argument("--access-key-id", dest="access_key_id", + help="OSS access key ID (overrides config.ini)") + parser.add_argument("--access-key-secret", dest="access_key_secret", + help="OSS access key secret (overrides config.ini)") + parser.add_argument("--region", help="OSS region (overrides config.ini)") + + list_parser = datasets_subparsers.add_parser("list", help="List datasets in OSS registry") + list_parser.add_argument("--org", help="Filter by organization") + add_oss_args(list_parser) + + upload_parser = datasets_subparsers.add_parser("upload", help="Upload local task dirs to OSS") + upload_parser.add_argument("--org", required=True, help="Organization name") + upload_parser.add_argument("--dataset", required=True, help="Dataset name") + upload_parser.add_argument("--split", required=True, help="Split name (e.g. train, test, v1.0)") + upload_parser.add_argument("--dir", required=True, + help="Local directory containing {task_id}/ subdirectories") + upload_parser.add_argument("--concurrency", type=int, default=4, + choices=range(1, 17), metavar="[1-16]", + help="Upload concurrency (default: 4)") + upload_parser.add_argument("--overwrite", action="store_true", + help="Overwrite existing tasks in OSS (default: skip)") + add_oss_args(upload_parser) diff --git a/rock/cli/config.py b/rock/cli/config.py index 7a51a4d315..60c8251882 100644 --- a/rock/cli/config.py +++ b/rock/cli/config.py @@ -8,12 +8,22 @@ logger = init_logger(__name__) +@dataclass +class DatasetConfig: + oss_bucket: str | None = None + oss_endpoint: str | None = None + oss_access_key_id: str | None = None + oss_access_key_secret: str | None = None + oss_region: str | None = None + + @dataclass class CLIConfig: """CLI configuration class""" base_url: str = env_vars.ROCK_BASE_URL extra_headers: dict[str, str] = field(default_factory=dict) + dataset_config: DatasetConfig = field(default_factory=DatasetConfig) class ConfigManager: @@ -56,6 +66,16 @@ def _load_config(self): if value.strip(): self.config.extra_headers[key] = value.strip() + if "dataset" in parser: + ds = parser["dataset"] + self.config.dataset_config = DatasetConfig( + oss_bucket=ds.get("oss_bucket") or None, + oss_endpoint=ds.get("oss_endpoint") or None, + oss_access_key_id=ds.get("oss_access_key_id") or None, + oss_access_key_secret=ds.get("oss_access_key_secret") or None, + oss_region=ds.get("oss_region") or None, + ) + except Exception as e: logger.warning(f"Failed to load config file {self.config_path}: {e}", exc_info=True) raise e diff --git a/rock/sdk/envhub/datasets/__init__.py b/rock/sdk/envhub/datasets/__init__.py new file mode 100644 index 0000000000..fd9d522084 --- /dev/null +++ b/rock/sdk/envhub/datasets/__init__.py @@ -0,0 +1,4 @@ +from rock.sdk.envhub.datasets.client import DatasetClient +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult + +__all__ = ["DatasetClient", "DatasetSpec", "UploadResult"] diff --git a/rock/sdk/envhub/datasets/client.py b/rock/sdk/envhub/datasets/client.py new file mode 100644 index 0000000000..02f1336d9b --- /dev/null +++ b/rock/sdk/envhub/datasets/client.py @@ -0,0 +1,20 @@ +from rock.sdk.bench.models.job.config import LocalDatasetConfig, OssRegistryInfo, RegistryDatasetConfig +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult +from rock.sdk.envhub.datasets.registry.oss import OssDatasetRegistry + + +class DatasetClient: + + def __init__(self, registry: OssRegistryInfo) -> None: + self._registry = OssDatasetRegistry(registry) + + def list_datasets(self, org: str | None = None) -> list[DatasetSpec]: + return self._registry.list_datasets(org) + + def upload_dataset( + self, + source: LocalDatasetConfig, + target: RegistryDatasetConfig, + concurrency: int = 4, + ) -> UploadResult: + return self._registry.upload_dataset(source, target, concurrency) diff --git a/rock/sdk/envhub/datasets/models.py b/rock/sdk/envhub/datasets/models.py new file mode 100644 index 0000000000..264dd1e1f7 --- /dev/null +++ b/rock/sdk/envhub/datasets/models.py @@ -0,0 +1,17 @@ +from dataclasses import dataclass, field + + +@dataclass +class DatasetSpec: + id: str # "{organization}/{dataset_name}", e.g. "princeton-nlp/SWE-bench_Verified" + split: str + task_ids: list[str] = field(default_factory=list) + + +@dataclass +class UploadResult: + id: str # "{organization}/{dataset_name}" + split: str + uploaded: int + skipped: int + failed: int diff --git a/rock/sdk/envhub/datasets/registry/__init__.py b/rock/sdk/envhub/datasets/registry/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/rock/sdk/envhub/datasets/registry/base.py b/rock/sdk/envhub/datasets/registry/base.py new file mode 100644 index 0000000000..8a3538aa6d --- /dev/null +++ b/rock/sdk/envhub/datasets/registry/base.py @@ -0,0 +1,22 @@ +from abc import ABC, abstractmethod + +from rock.sdk.bench.models.job.config import LocalDatasetConfig, RegistryDatasetConfig +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult + + +class BaseDatasetRegistry(ABC): + + @abstractmethod + def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: + """List all datasets. Filtered to `organization` if provided.""" + ... + + @abstractmethod + def upload_dataset( + self, + source: LocalDatasetConfig, + target: RegistryDatasetConfig, + concurrency: int = 4, + ) -> UploadResult: + """Upload source.path/{task_id}/ subdirs to target (org/name/split from target.name and target.version).""" + ... diff --git a/rock/sdk/envhub/datasets/registry/oss.py b/rock/sdk/envhub/datasets/registry/oss.py new file mode 100644 index 0000000000..db455465c4 --- /dev/null +++ b/rock/sdk/envhub/datasets/registry/oss.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import oss2 + +from rock.logger import init_logger +from rock.sdk.bench.models.job.config import LocalDatasetConfig, OssRegistryInfo, RegistryDatasetConfig +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult +from rock.sdk.envhub.datasets.registry.base import BaseDatasetRegistry + +logger = init_logger(__name__) + + +class OssDatasetRegistry(BaseDatasetRegistry): + + def __init__(self, registry: OssRegistryInfo) -> None: + self._registry = registry + + def _build_bucket(self) -> oss2.Bucket: + auth = oss2.Auth( + self._registry.oss_access_key_id or "", + self._registry.oss_access_key_secret or "", + ) + return oss2.Bucket(auth, self._registry.oss_endpoint or "", self._registry.oss_bucket) + + def _build_prefix(self, org: str, name: str, split: str | None = None) -> str: + base = self._registry.oss_dataset_path or "datasets" + parts = [base, org, name] + if split: + parts.append(split) + return "/".join(parts) + + @staticmethod + def _last_segment(prefix: str) -> str: + return prefix.rstrip("/").rsplit("/", 1)[-1] + + def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: + bucket = self._build_bucket() + base = self._registry.oss_dataset_path or "datasets" + + if organization: + org_prefixes = [f"{base}/{organization}/"] + else: + result = bucket.list_objects_v2(prefix=f"{base}/", delimiter="/", max_keys=1000) + org_prefixes = result.prefix_list + + datasets: list[DatasetSpec] = [] + for org_prefix in org_prefixes: + org = self._last_segment(org_prefix) + + result = bucket.list_objects_v2(prefix=org_prefix, delimiter="/", max_keys=1000) + for name_prefix in result.prefix_list: + name = self._last_segment(name_prefix) + + result2 = bucket.list_objects_v2(prefix=name_prefix, delimiter="/", max_keys=1000) + for split_prefix in result2.prefix_list: + split = self._last_segment(split_prefix) + + result3 = bucket.list_objects_v2(prefix=split_prefix, delimiter="/", max_keys=1000) + task_ids = [self._last_segment(p) for p in result3.prefix_list] + datasets.append(DatasetSpec( + id=f"{org}/{name}", + split=split, + task_ids=task_ids, + )) + + return datasets + + def _task_exists(self, bucket: oss2.Bucket, task_prefix: str) -> bool: + result = bucket.list_objects_v2(prefix=task_prefix, max_keys=1) + return len(result.object_list) > 0 + + def _upload_task( + self, + bucket: oss2.Bucket, + org: str, + name: str, + split: str, + task_dir: Path, + overwrite: bool, + ) -> int | None: + task_id = task_dir.name + base = self._registry.oss_dataset_path or "datasets" + task_prefix = f"{base}/{org}/{name}/{split}/{task_id}/" + + if not overwrite and self._task_exists(bucket, task_prefix): + return None + + files = [f for f in task_dir.rglob("*") if f.is_file()] + for file in files: + key = f"{task_prefix}{file.relative_to(task_dir)}" + bucket.put_object(key, file.read_bytes()) + return len(files) + + def upload_dataset( + self, + source: LocalDatasetConfig, + target: RegistryDatasetConfig, + concurrency: int = 4, + ) -> UploadResult: + org, name = target.name.split("/", 1) + split = target.version or "" + overwrite = target.overwrite + local_dir = source.path + + bucket = self._build_bucket() + task_dirs = sorted([d for d in local_dir.iterdir() if d.is_dir()]) + + raw: dict[str, int | None | Exception] = {} + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = { + executor.submit(self._upload_task, bucket, org, name, split, d, overwrite): d + for d in task_dirs + } + for future, task_dir in futures.items(): + try: + raw[task_dir.name] = future.result() + except Exception as exc: + raw[task_dir.name] = exc + + uploaded = skipped = failed = 0 + for task_id in sorted(raw): + outcome = raw[task_id] + if isinstance(outcome, Exception): + failed += 1 + logger.error("Failed to upload task %s: %s", task_id, outcome) + print(f" \u2717 {task_id} failed: {outcome}") + elif outcome is None: + skipped += 1 + print(f" - {task_id} skipped (already exists)") + else: + uploaded += 1 + print(f" \u2713 {task_id} ({outcome} files)") + + return UploadResult( + id=f"{org}/{name}", + split=split, + uploaded=uploaded, + skipped=skipped, + failed=failed, + ) diff --git a/tests/unit/datasets/__init__.py b/tests/unit/datasets/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/datasets/test_client.py b/tests/unit/datasets/test_client.py new file mode 100644 index 0000000000..00a5448e32 --- /dev/null +++ b/tests/unit/datasets/test_client.py @@ -0,0 +1,33 @@ +from unittest.mock import patch + +from rock.sdk.bench.models.job.config import LocalDatasetConfig, OssRegistryInfo, RegistryDatasetConfig +from rock.sdk.envhub.datasets.client import DatasetClient +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult + + +def make_registry_info(): + return OssRegistryInfo(oss_bucket="b", oss_access_key_id="k", oss_access_key_secret="s") + + +def test_dataset_client_list_delegates_to_registry(): + client = DatasetClient(make_registry_info()) + expected = [DatasetSpec(id="qwen/bench", split="train", task_ids=[])] + + with patch.object(client._registry, "list_datasets", return_value=expected) as mock_list: + result = client.list_datasets(org="qwen") + + mock_list.assert_called_once_with("qwen") + assert result == expected + + +def test_dataset_client_upload_delegates_to_registry(tmp_path): + client = DatasetClient(make_registry_info()) + source = LocalDatasetConfig(path=tmp_path) + target = RegistryDatasetConfig(name="qwen/bench", version="train", overwrite=True, registry=make_registry_info()) + expected = UploadResult(id="qwen/bench", split="train", uploaded=1, skipped=0, failed=0) + + with patch.object(client._registry, "upload_dataset", return_value=expected) as mock_up: + result = client.upload_dataset(source, target, concurrency=2) + + mock_up.assert_called_once_with(source, target, 2) + assert result == expected diff --git a/tests/unit/datasets/test_config.py b/tests/unit/datasets/test_config.py new file mode 100644 index 0000000000..98477859bc --- /dev/null +++ b/tests/unit/datasets/test_config.py @@ -0,0 +1,44 @@ +from rock.cli.config import ConfigManager + + +def test_dataset_config_from_ini(tmp_path): + config_file = tmp_path / "config.ini" + config_file.write_text( + "[rock]\n" + "base_url = http://localhost:8080\n" + "\n" + "[dataset]\n" + "oss_bucket = my-bucket\n" + "oss_endpoint = https://oss-cn-hangzhou.aliyuncs.com\n" + "oss_access_key_id = LTAI5t\n" + "oss_access_key_secret = secret123\n" + "oss_region = cn-hangzhou\n" + ) + ds = ConfigManager(config_file).get_config().dataset_config + + assert ds.oss_bucket == "my-bucket" + assert ds.oss_endpoint == "https://oss-cn-hangzhou.aliyuncs.com" + assert ds.oss_access_key_id == "LTAI5t" + assert ds.oss_access_key_secret == "secret123" + assert ds.oss_region == "cn-hangzhou" + + +def test_dataset_config_defaults_when_section_absent(tmp_path): + config_file = tmp_path / "config.ini" + config_file.write_text("[rock]\nbase_url = http://localhost:8080\n") + ds = ConfigManager(config_file).get_config().dataset_config + + assert ds.oss_bucket is None + assert ds.oss_endpoint is None + assert ds.oss_access_key_id is None + assert ds.oss_access_key_secret is None + assert ds.oss_region is None + + +def test_dataset_config_partial_section(tmp_path): + config_file = tmp_path / "config.ini" + config_file.write_text("[dataset]\noss_bucket = only-bucket\n") + ds = ConfigManager(config_file).get_config().dataset_config + + assert ds.oss_bucket == "only-bucket" + assert ds.oss_endpoint is None diff --git a/tests/unit/datasets/test_datasets_command.py b/tests/unit/datasets/test_datasets_command.py new file mode 100644 index 0000000000..9f40fe4a3f --- /dev/null +++ b/tests/unit/datasets/test_datasets_command.py @@ -0,0 +1,78 @@ +import argparse +from unittest.mock import patch + +import pytest + +from rock.cli.command.datasets import DatasetsCommand + + +def make_base_args(**kwargs): + args = argparse.Namespace( + config=None, + datasets_command=None, + bucket=None, + endpoint=None, + access_key_id=None, + access_key_secret=None, + region=None, + org=None, + ) + for k, v in kwargs.items(): + setattr(args, k, v) + return args + + +def test_command_name(): + assert DatasetsCommand.name == "datasets" + + +def test_build_oss_registry_info_from_cli_args(): + cmd = DatasetsCommand() + args = make_base_args(bucket="cli-bucket", endpoint="https://oss.example.com", access_key_id="kid", access_key_secret="ksec") + + with patch("rock.cli.command.datasets.ConfigManager") as mock_mgr: + ds_cfg = mock_mgr.return_value.get_config.return_value.dataset_config + ds_cfg.oss_bucket = None + ds_cfg.oss_endpoint = None + ds_cfg.oss_access_key_id = None + ds_cfg.oss_access_key_secret = None + ds_cfg.oss_region = None + info = cmd._build_oss_registry_info(args) + + assert info.oss_bucket == "cli-bucket" + assert info.oss_endpoint == "https://oss.example.com" + assert info.oss_access_key_id == "kid" + + +def test_build_oss_registry_info_cli_overrides_ini(): + cmd = DatasetsCommand() + args = make_base_args(bucket="cli-bucket", endpoint=None, access_key_id=None, access_key_secret=None) + + with patch("rock.cli.command.datasets.ConfigManager") as mock_mgr: + ds_cfg = mock_mgr.return_value.get_config.return_value.dataset_config + ds_cfg.oss_bucket = "ini-bucket" + ds_cfg.oss_endpoint = "https://ini.example.com" + ds_cfg.oss_access_key_id = "ini-kid" + ds_cfg.oss_access_key_secret = "ini-ksec" + ds_cfg.oss_region = None + info = cmd._build_oss_registry_info(args) + + assert info.oss_bucket == "cli-bucket" + assert info.oss_endpoint == "https://ini.example.com" + assert info.oss_access_key_id == "ini-kid" + + +def test_build_oss_registry_info_raises_when_bucket_missing(): + cmd = DatasetsCommand() + args = make_base_args(bucket=None) + + with patch("rock.cli.command.datasets.ConfigManager") as mock_mgr: + ds_cfg = mock_mgr.return_value.get_config.return_value.dataset_config + ds_cfg.oss_bucket = None + ds_cfg.oss_endpoint = None + ds_cfg.oss_access_key_id = None + ds_cfg.oss_access_key_secret = None + ds_cfg.oss_region = None + + with pytest.raises(ValueError, match="bucket"): + cmd._build_oss_registry_info(args) diff --git a/tests/unit/datasets/test_models.py b/tests/unit/datasets/test_models.py new file mode 100644 index 0000000000..43b2dee88b --- /dev/null +++ b/tests/unit/datasets/test_models.py @@ -0,0 +1,14 @@ +from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult + + +def test_dataset_spec_id(): + spec = DatasetSpec(id="qwen/my-bench", split="train", task_ids=["t1", "t2"]) + assert spec.id == "qwen/my-bench" + assert len(spec.task_ids) == 2 + + +def test_upload_result_fields(): + result = UploadResult(id="qwen/my-bench", split="train", uploaded=2, skipped=1, failed=0) + assert result.uploaded == 2 + assert result.skipped == 1 + assert result.failed == 0 diff --git a/tests/unit/datasets/test_oss_registry.py b/tests/unit/datasets/test_oss_registry.py new file mode 100644 index 0000000000..4ba2f0db52 --- /dev/null +++ b/tests/unit/datasets/test_oss_registry.py @@ -0,0 +1,170 @@ +from unittest.mock import MagicMock, patch + +from rock.sdk.bench.models.job.config import LocalDatasetConfig, OssRegistryInfo, RegistryDatasetConfig +from rock.sdk.envhub.datasets.registry.oss import OssDatasetRegistry + + +def make_registry_info(): + return OssRegistryInfo( + oss_bucket="test-bucket", + oss_endpoint="https://oss-cn-hangzhou.aliyuncs.com", + oss_access_key_id="key", + oss_access_key_secret="secret", + ) + + +def make_list_result(prefixes=None, objects=None): + result = MagicMock() + result.prefix_list = prefixes or [] + result.object_list = objects or [] + return result + + +def test_list_datasets_returns_all(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.side_effect = [ + make_list_result(prefixes=["datasets/qwen/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/train/"]), + make_list_result(prefixes=[ + "datasets/qwen/my-bench/train/task-001/", + "datasets/qwen/my-bench/train/task-002/", + ]), + ] + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + datasets = registry.list_datasets() + + assert len(datasets) == 1 + assert datasets[0].id == "qwen/my-bench" + assert datasets[0].split == "train" + assert datasets[0].task_ids == ["task-001", "task-002"] + + +def test_list_datasets_filter_by_org(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.side_effect = [ + make_list_result(prefixes=["datasets/qwen/my-bench/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/train/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/train/task-001/"]), + ] + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + datasets = registry.list_datasets(organization="qwen") + + first_call_kwargs = mock_bucket.list_objects_v2.call_args_list[0][1] + assert first_call_kwargs["prefix"] == "datasets/qwen/" + assert len(datasets) == 1 + + +def test_list_datasets_empty_registry(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + datasets = registry.list_datasets() + + assert datasets == [] + + +def test_build_prefix_without_split(): + registry = OssDatasetRegistry(make_registry_info()) + assert registry._build_prefix("qwen", "my-bench") == "datasets/qwen/my-bench" + + +def test_build_prefix_with_split(): + registry = OssDatasetRegistry(make_registry_info()) + assert registry._build_prefix("qwen", "my-bench", "train") == "datasets/qwen/my-bench/train" + + +# --------------------------------------------------------------------------- +# upload_dataset tests +# --------------------------------------------------------------------------- + + +def make_upload_pair(tmp_path, *, name="qwen/my-bench", version="train", overwrite=False): + source = LocalDatasetConfig(path=tmp_path) + target = RegistryDatasetConfig( + name=name, + version=version, + overwrite=overwrite, + registry=make_registry_info(), + ) + return source, target + + +def test_upload_dataset_new_tasks(tmp_path): + (tmp_path / "task-001").mkdir() + (tmp_path / "task-001" / "task.toml").write_text("[task]") + (tmp_path / "task-002").mkdir() + (tmp_path / "task-002" / "task.toml").write_text("[task]") + + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(objects=[]) + source, target = make_upload_pair(tmp_path) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + result = registry.upload_dataset(source, target) + + assert result.uploaded == 2 + assert result.skipped == 0 + assert result.failed == 0 + assert mock_bucket.put_object.call_count == 2 + + +def test_upload_dataset_skips_existing(tmp_path): + (tmp_path / "task-001").mkdir() + (tmp_path / "task-001" / "task.toml").write_text("[task]") + + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + objects=[MagicMock(key="datasets/qwen/my-bench/train/task-001/task.toml")] + ) + source, target = make_upload_pair(tmp_path, overwrite=False) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + result = registry.upload_dataset(source, target) + + assert result.uploaded == 0 + assert result.skipped == 1 + mock_bucket.put_object.assert_not_called() + + +def test_upload_dataset_overwrite(tmp_path): + (tmp_path / "task-001").mkdir() + (tmp_path / "task-001" / "task.toml").write_text("[task]") + + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + objects=[MagicMock(key="datasets/qwen/my-bench/train/task-001/task.toml")] + ) + source, target = make_upload_pair(tmp_path, overwrite=True) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + result = registry.upload_dataset(source, target) + + assert result.uploaded == 1 + assert result.skipped == 0 + mock_bucket.put_object.assert_called_once() + + +def test_upload_dataset_oss_key_format(tmp_path): + (tmp_path / "task-001").mkdir() + (tmp_path / "task-001" / "task.toml").write_text("[task]") + + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(objects=[]) + source, target = make_upload_pair(tmp_path) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + registry.upload_dataset(source, target) + + key = mock_bucket.put_object.call_args[0][0] + assert key == "datasets/qwen/my-bench/train/task-001/task.toml" From 3e2273f2b1249862f6cfaee8682fb32cead8bb6c Mon Sep 17 00:00:00 2001 From: guoj14 Date: Tue, 21 Apr 2026 12:04:53 +0800 Subject: [PATCH 057/226] fix(test): clean up leaked timers in model client tests (#839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(test): 修复 Jest timer 泄漏问题，确保超时定时器在测试完成后被清理 * fix(sandbox): 修复状态检查超时清理问题 --- rock/ts-sdk/src/model/client.test.ts | 16 ++++++++++++---- rock/ts-sdk/src/sandbox/client.ts | 9 ++++++--- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/rock/ts-sdk/src/model/client.test.ts b/rock/ts-sdk/src/model/client.test.ts index 07eb17768c..6d11442198 100644 --- a/rock/ts-sdk/src/model/client.test.ts +++ b/rock/ts-sdk/src/model/client.test.ts @@ -122,9 +122,13 @@ describe('ModelClient timeout and cancellation', () => { const controller = new AbortController(); // Abort after a short delay - setTimeout(() => controller.abort(), 100); + const timerId = setTimeout(() => controller.abort(), 100); - await expect(client.popRequest(1, { timeout: 10, signal: controller.signal })).rejects.toThrow(); + try { + await expect(client.popRequest(1, { timeout: 10, signal: controller.signal })).rejects.toThrow(); + } finally { + clearTimeout(timerId); + } }); it('should return request when found before abort', async () => { @@ -144,9 +148,13 @@ describe('ModelClient timeout and cancellation', () => { const controller = new AbortController(); // Abort after a short delay - setTimeout(() => controller.abort(), 100); + const timerId = setTimeout(() => controller.abort(), 100); - await expect(client.waitForFirstRequest({ timeout: 10, signal: controller.signal })).rejects.toThrow(); + try { + await expect(client.waitForFirstRequest({ timeout: 10, signal: controller.signal })).rejects.toThrow(); + } finally { + clearTimeout(timerId); + } }); it('should return when log file has content before abort', async () => { diff --git a/rock/ts-sdk/src/sandbox/client.ts b/rock/ts-sdk/src/sandbox/client.ts index c9d8fd3b4e..2bdcc26afc 100644 --- a/rock/ts-sdk/src/sandbox/client.ts +++ b/rock/ts-sdk/src/sandbox/client.ts @@ -259,13 +259,14 @@ export class Sandbox extends AbstractSandbox { const checkInterval = 3000; // 3s between checks while (Date.now() - startTime < this.config.startupTimeout * 1000) { + let timeoutId: ReturnType | undefined; try { logger.info(`Checking status... (elapsed: ${Math.round((Date.now() - startTime) / 1000)}s)`); // Use Promise.race to implement timeout for status check const statusPromise = this.getStatus(); - const timeoutPromise = new Promise((_, reject) => - setTimeout(() => reject(new Error('Status check timeout')), checkTimeout) - ); + const timeoutPromise = new Promise((_, reject) => { + timeoutId = setTimeout(() => reject(new Error('Status check timeout')), checkTimeout); + }); const status = await Promise.race([statusPromise, timeoutPromise]); if (status && status.isAlive) { @@ -275,6 +276,8 @@ export class Sandbox extends AbstractSandbox { } catch (e) { // Status check may fail temporarily during startup, continue waiting logger.debug(`Status check failed (will retry): ${e}`); + } finally { + clearTimeout(timeoutId); } await sleep(checkInterval); } From 936d3d6865ea5bf1628e1d45de74aa0d20c2dc93 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 21 Apr 2026 19:41:25 +0800 Subject: [PATCH 058/226] feat(sandbox): enforce container rootfs disk limit via Docker storage-opt (#860) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(disk-limit): add DockerUtil helpers for storage-opt and XFS detection - detect_storage_opt_support(): checks overlay2 driver + xfs backing + prjquota mount option - get_docker_info() / get_docker_root_dir(): thin wrappers around docker info - is_xfs_path(): checks if a filesystem path lives on an XFS mount These utilities are consumed by DockerDeployment to decide whether to apply --storage-opt and xfs_quota at container start time. * feat(disk-limit): enforce rootfs and log-dir disk quotas per container Rootfs quota (storage-opt): - Add limit_disk field to DockerDeploymentConfig; default None (no limit). - Pass --storage-opt size= to docker run when the field is set. - Graceful degradation: if the worker does not support storage-opt (requires overlay2 + xfs + prjquota), effective_limit_disk is set to None while config.limit_disk stays unchanged. - Surface effective_limit_disk in SandboxInfo, sandbox_actor status, and SandboxStatusResponse so callers can observe the applied quota. Log-dir quota (XFS project quota): - Add limit_log_dir field to DockerDeploymentConfig; default None (no limit). - After container start, call xfs_quota to assign a per-project hard limit on ROCK_LOGGING_PATH; skip silently when the path is not on XFS. Server-side configuration: - RuntimeConfig gains default_limit_disk and default_log_dir_quota (both default to None); set in rock-xxx.yml per environment. - _apply_disk_limits() in sandbox_api reads these values and then checks Nacos for runtime overrides (Nacos wins over yml; None in both = no limit). - Remove limit_disk from SandboxStartRequest / SandboxConfig (SDK) — quota policy is server-side only; clients observe it via status responses. * test(disk-limit): add unit and integration tests Unit tests: - test_docker_util.py: detect_storage_opt_support, is_xfs_path, get_docker_root_dir - test_docker_deployment_disk_limit.py: _storage_opts(), start() degradation, effective_limit_disk vs config.limit_disk invariants - test_sandbox_manager_disk_limit.py: validate limit_disk format - test_sandbox_response.py: limit_disk field in SandboxStatusResponse Integration tests: - test_disk_limit_enforcement: rootfs quota blocks oversized files (SKIP_IF_NO_STORAGE_OPT) - test_disk_limit_default_value: server-reported limit_disk reflects configured value - test_logging_path_disk_limit_enforcement: log-dir quota is independently enforced (SKIP_IF_NO_STORAGE_OPT + SKIP_IF_LOG_PATH_NOT_XFS) * chore: rename limit_disk to disk_limit Signed-off-by: Jiachen Zhang * fix: using the last line of findmnt as the mountpoint to check XFS mount option Signed-off-by: Jiachen Zhang --------- Signed-off-by: Jiachen Zhang --- rock/actions/sandbox/response.py | 2 + rock/actions/sandbox/sandbox_info.py | 2 + rock/admin/entrypoints/sandbox_api.py | 27 ++ rock/admin/proto/response.py | 6 + rock/common/constants.py | 2 + rock/config.py | 4 + rock/deployments/config.py | 6 + rock/deployments/docker.py | 100 ++++++- rock/sandbox/sandbox_actor.py | 4 + rock/sandbox/sandbox_manager.py | 10 + rock/utils/docker.py | 106 ++++++++ tests/integration/conftest.py | 12 + .../sdk/sandbox/test_disk_limit.py | 229 ++++++++++++++++ .../unit/admin/proto/test_sandbox_response.py | 142 ++++++++++ .../test_docker_deployment_disk_limit.py | 245 ++++++++++++++++++ tests/unit/sandbox/job/__init__.py | 0 .../test_sandbox_manager_disk_limit.py | 70 +++++ tests/unit/utils/test_docker_util.py | 181 +++++++++++++ 18 files changed, 1139 insertions(+), 9 deletions(-) create mode 100644 tests/integration/sdk/sandbox/test_disk_limit.py create mode 100644 tests/unit/admin/proto/test_sandbox_response.py create mode 100644 tests/unit/deployments/test_docker_deployment_disk_limit.py create mode 100644 tests/unit/sandbox/job/__init__.py create mode 100644 tests/unit/sandbox/test_sandbox_manager_disk_limit.py create mode 100644 tests/unit/utils/test_docker_util.py diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index 3ec3247193..7f7baaab89 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -48,6 +48,8 @@ class SandboxStatusResponse(BaseModel): namespace: str | None = None cpus: float | None = None memory: str | None = None + disk_limit_rootfs: str | None = None + disk_limit_log: str | None = None class CommandResponse(BaseModel): diff --git a/rock/actions/sandbox/sandbox_info.py b/rock/actions/sandbox/sandbox_info.py index 98f9929adf..ae2aac16f7 100644 --- a/rock/actions/sandbox/sandbox_info.py +++ b/rock/actions/sandbox/sandbox_info.py @@ -21,6 +21,8 @@ class SandboxInfo(TypedDict, total=False): create_user_gray_flag: bool cpus: float memory: str + disk_limit_rootfs: str + disk_limit_log: str create_time: str start_time: str stop_time: str diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index 277cc39ea0..e8a264a659 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -30,6 +30,8 @@ GET_STATUS_SWITCH, KATA_DIND_DISK_SIZE_KEY, KATA_RUNTIME_SWITCH, + SANDBOX_DISK_LIMIT_LOG_KEY, + SANDBOX_DISK_LIMIT_ROOTFS_KEY, SUPPORT_KATA_SWITCH, ) from rock.common.exception import handle_exceptions @@ -69,6 +71,29 @@ async def _apply_kata_disk_size(config: DockerDeploymentConfig) -> None: config.kata_disk_size = disk_size +async def _apply_disk_limits(config: DockerDeploymentConfig) -> None: + """Apply disk limits from RuntimeConfig (rock-xxx.yml), overridable by Nacos at runtime. + + Priority: Nacos > RuntimeConfig (rock-xxx.yml). None in both means no limit. + """ + runtime = sandbox_manager.rock_config.runtime + nacos = sandbox_manager.rock_config.nacos_provider + + disk_limit_rootfs = runtime.sandbox_disk_limit_rootfs + disk_limit_log = runtime.sandbox_disk_limit_log + + if nacos is not None: + nacos_rootfs = await nacos.get_config_value(SANDBOX_DISK_LIMIT_ROOTFS_KEY) + if nacos_rootfs: + disk_limit_rootfs = nacos_rootfs + nacos_log = await nacos.get_config_value(SANDBOX_DISK_LIMIT_LOG_KEY) + if nacos_log: + disk_limit_log = nacos_log + + config.disk_limit_rootfs = disk_limit_rootfs + config.disk_limit_log = disk_limit_log + + async def _apply_cpu_preempt_switch(config: DockerDeploymentConfig) -> None: """Check nacos switch and enable CPU preemption on the config if the switch is on. @@ -89,6 +114,7 @@ async def start(request: SandboxStartRequest) -> RockResponse[SandboxStartRespon await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) await _apply_cpu_preempt_switch(config) + await _apply_disk_limits(config) sandbox_start_response = await sandbox_manager.start(config) return RockResponse(result=sandbox_start_response) @@ -103,6 +129,7 @@ async def start_async( await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) await _apply_cpu_preempt_switch(config) + await _apply_disk_limits(config) sandbox_start_response = await sandbox_manager.start_async( config, user_info=headers.user_info, diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index a59df09d5a..d1edabc7a5 100644 --- a/rock/admin/proto/response.py +++ b/rock/admin/proto/response.py @@ -11,6 +11,8 @@ class SandboxStartResponse(SandboxResponse): host_ip: str | None = None cpus: float | None = None memory: str | None = None + disk_limit_rootfs: str | None = None + disk_limit_log: str | None = None # TODO: inherit from SandboxStartResponse @@ -30,6 +32,8 @@ class SandboxStatusResponse(BaseModel): namespace: str | None = None cpus: float | None = None memory: str | None = None + disk_limit_rootfs: str | None = None + disk_limit_log: str | None = None @classmethod def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusResponse": @@ -46,6 +50,8 @@ def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusRespons namespace=sandbox_info.get("namespace"), cpus=sandbox_info.get("cpus"), memory=sandbox_info.get("memory"), + disk_limit_rootfs=sandbox_info.get("disk_limit_rootfs"), + disk_limit_log=sandbox_info.get("disk_limit_log"), ) diff --git a/rock/common/constants.py b/rock/common/constants.py index d7fa78fc5e..68a5cc27bd 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -5,6 +5,8 @@ SUPPORT_KATA_SWITCH = "support_kata_enabled" CPU_PREEMPT_SWITCH = "cpu_preempt_enabled" KATA_DIND_DISK_SIZE_KEY = "kata_dind_disk_size" +SANDBOX_DISK_LIMIT_ROOTFS_KEY = "sandbox_disk_limit_rootfs" +SANDBOX_DISK_LIMIT_LOG_KEY = "sandbox_disk_limit_log" PID_PREFIX = "PIDSTART" PID_SUFFIX = "PIDEND" SCHEDULER_LOG_NAME = "scheduler.log" diff --git a/rock/config.py b/rock/config.py index 7b829bf4b0..03ccd8534a 100644 --- a/rock/config.py +++ b/rock/config.py @@ -168,6 +168,10 @@ class RuntimeConfig: use_standard_spec_only: bool = False metrics_endpoint: str = "" user_defined_tags: dict = field(default_factory=dict) + sandbox_disk_limit_rootfs: str | None = None + """Default rootfs quota per container. None means no limit. Can be overridden by nacos key 'default_disk_limit'.""" + sandbox_disk_limit_log: str | None = None + """Default log-dir quota per container. None means no limit. Can be overridden by nacos key 'default_log_dir_quota'.""" def __post_init__(self) -> None: # Convert dict to StandardSpec if needed diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 35ab4fe4ab..b9b9da17fd 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -93,6 +93,12 @@ class DockerDeploymentConfig(DeploymentConfig): limit_cpus: float | None = None """Hard limit on the number of CPU cores the container can use. Used as --cpus when CPU preemption is enabled via nacos switch.""" + disk_limit_rootfs: str | None = None + """Maximum rootfs disk size for the container (e.g., '20g', '50g'). Maps to --storage-opt size=. Only supported on overlay2 storage driver with xfs backing filesystem. None means no limit.""" + + disk_limit_log: str | None = None + """XFS project quota for the sandbox log directory. Server-side only, applied via xfs_quota. None means no limit.""" + container_name: str | None = None """Custom name for the container. If None, a random name will be generated.""" diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index f27aba3594..05573dc726 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -1,5 +1,6 @@ import asyncio import datetime +import hashlib import os import random import shlex @@ -18,6 +19,7 @@ from rock.deployments.abstract import AbstractDeployment from rock.deployments.config import DockerDeploymentConfig from rock.deployments.constants import Port, Status +from rock.deployments.docker_client import TempAuthDockerClient, TempAuthDockerClientError from rock.deployments.hooks.abstract import CombinedDeploymentHook, DeploymentHook from rock.deployments.runtime_env import DockerRuntimeEnv, LocalRuntimeEnv, PipRuntimeEnv, UvRuntimeEnv from rock.deployments.sandbox_validator import DockerSandboxValidator @@ -38,7 +40,6 @@ timeout, wait_until_alive, ) -from rock.deployments.docker_client import TempAuthDockerClient, TempAuthDockerClientError __all__ = ["DockerDeployment", "DockerDeploymentConfig"] CHECK_CLEAR_INTERVAL_SECONDS = 300 @@ -57,10 +58,12 @@ def __init__( Args: **kwargs: Keyword arguments (see `DockerDeploymentConfig` for details). """ - registry_password = kwargs.pop('registry_password', None) + registry_password = kwargs.pop("registry_password", None) self._config = DockerDeploymentConfig(**kwargs) if registry_password: self._config.registry_password = registry_password + self._effective_disk_limit_rootfs: str | None = self._config.disk_limit_rootfs + self._effective_disk_limit_log: str | None = self._config.disk_limit_log self._runtime: RemoteSandboxRuntime | None = None self._container_process = None self._runtime_timeout = 0.15 @@ -252,16 +255,14 @@ def _pull_image(self) -> None: ) return - self._service_status.update_status( - phase_name="image_pull", status=Status.RUNNING, message="image pull running" - ) + self._service_status.update_status(phase_name="image_pull", status=Status.RUNNING, message="image pull running") logger.info(f"Pulling image {self._config.image!r}") try: with Timer(description=f"[{self._config.image}] Image pull"): # Parse registry from image name registry, _ = ImageUtil.parse_registry_and_others(self._config.image) - + # Create temp auth client with credentials if available with TempAuthDockerClient( registry=registry if self._config.registry_username else None, @@ -276,9 +277,7 @@ def _pull_image(self) -> None: except (subprocess.CalledProcessError, TempAuthDockerClientError) as e: msg = f"Failed to pull image {self._config.image}: {e}" - self._service_status.update_status( - phase_name="image_pull", status=Status.FAILED, message=msg - ) + self._service_status.update_status(phase_name="image_pull", status=Status.FAILED, message=msg) raise DockerPullError(msg) from e @property @@ -370,11 +369,82 @@ def _cpus(self): return [f"--cpu-shares={cpu_shares}", f"--cpus={self.config.limit_cpus}"] return [f"--cpus={self.config.cpus}"] + def _storage_opts(self): + if self._effective_disk_limit_rootfs is not None: + return ["--storage-opt", f"size={self._effective_disk_limit_rootfs}"] + return [] + + def _try_set_log_dir_quota(self, log_file_path: str) -> None: + """Best-effort: set XFS project quota for sandbox log directory. + + Requires the log path to be on an XFS mount with prjquota/pquota enabled. + This check is independent of Docker's storage driver (no overlay2 requirement). + """ + if self._effective_disk_limit_log is None: + return + + if not DockerUtil.is_xfs_prjquota_path(log_file_path): + logger.info(f"Log path {log_file_path!r} is not on XFS+prjquota, skipping quota setup") + self._effective_disk_limit_log = None + return + + # Derive a deterministic project id from container name; reserve low ids. + project_id = (int(hashlib.sha1(self.container_name.encode("utf-8")).hexdigest()[:8], 16) % 900000) + 100000 + try: + findmnt_result = subprocess.run( + ["findmnt", "-T", log_file_path, "-o", "TARGET", "--noheadings"], + capture_output=True, + text=True, + timeout=5, + ) + if findmnt_result.returncode != 0: + logger.warning(f"Failed to find mountpoint for log path {log_file_path!r}, skip quota setup") + self._effective_disk_limit_log = None + return + mount_point = findmnt_result.stdout.strip() + if not mount_point: + logger.warning(f"Empty mountpoint for log path {log_file_path!r}, skip quota setup") + self._effective_disk_limit_log = None + return + + set_project_cmd = f"project -s -p {shlex.quote(log_file_path)} {project_id}" + set_limit_cmd = f"limit -p bhard={self._effective_disk_limit_log} {project_id}" + for cmd in (set_project_cmd, set_limit_cmd): + result = subprocess.run( + ["xfs_quota", "-x", "-c", cmd, mount_point], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode != 0: + logger.warning( + f"xfs_quota failed for {log_file_path!r} with cmd={cmd!r}: {result.stderr.strip() or result.stdout.strip()}" + ) + self._effective_disk_limit_log = None + return + logger.info(f"Set XFS project quota {self._effective_disk_limit_log} for log path {log_file_path!r}") + except Exception as e: + logger.warning(f"Failed to set XFS project quota for {log_file_path!r}: {e}") + self._effective_disk_limit_log = None + async def start(self): """Starts the runtime.""" if not self.sandbox_validator.check_availability(): raise Exception("Docker is not available") + storage_opt_supported = DockerUtil.detect_storage_opt_support() + # Resolve effective rootfs quota: downgrade to None if storage-opt is not supported. + if self._config.disk_limit_rootfs is not None and not storage_opt_supported: + logger.warning( + f"[{self.config.container_name}] --storage-opt not supported on this worker " + f"(requires overlay2 + xfs + prjquota), ignoring disk_limit_rootfs={self._config.disk_limit_rootfs}" + ) + self._effective_disk_limit_rootfs = None + else: + self._effective_disk_limit_rootfs = self._config.disk_limit_rootfs + # Resolve effective log quota; _try_set_log_dir_quota will downgrade to None if XFS+prjquota is unavailable. + self._effective_disk_limit_log = self._config.disk_limit_log + if self._container_name is None: self.set_container_name(self._get_container_name()) self._service_status.set_sandbox_id(self._container_name) @@ -406,6 +476,7 @@ async def start(self): log_file_path = f"{env_vars.ROCK_LOGGING_PATH}/{self.container_name}" os.makedirs(log_file_path, exist_ok=True) os.chmod(log_file_path, 0o777) + self._try_set_log_dir_quota(log_file_path) volume_args.extend(["-v", f"{log_file_path}:{env_vars.ROCK_LOGGING_PATH}"]) env_arg = [ "-e", @@ -442,6 +513,7 @@ async def start(self): f"{self._service_status.get_mapped_port(Port.SSH)}:22", *self._memory(), *self._cpus(), + *self._storage_opts(), *platform_arg, *self._config.docker_args, "--name", @@ -575,6 +647,16 @@ def config(self) -> DockerDeploymentConfig: """Returns the config of the deployment.""" return self._config + @property + def effective_disk_limit_rootfs(self) -> str | None: + """Returns the actual rootfs quota in effect after runtime capability checks (may differ from config.disk_limit_rootfs).""" + return self._effective_disk_limit_rootfs + + @property + def effective_disk_limit_log(self) -> str | None: + """Returns the actual log-dir quota in effect after runtime capability checks (may differ from config.disk_limit_log).""" + return self._effective_disk_limit_log + async def _check_stop(self): logger.info(f"Start check container to stop: {self._container_name}") try: diff --git a/rock/sandbox/sandbox_actor.py b/rock/sandbox/sandbox_actor.py index a130d00727..92a940e346 100644 --- a/rock/sandbox/sandbox_actor.py +++ b/rock/sandbox/sandbox_actor.py @@ -128,6 +128,8 @@ async def start(self): logger.error(f"[{self._config.container_name}] start deployment failed: {ex}", exc_info=True) raise ex if isinstance(self._deployment, DockerDeployment): + self._config.disk_limit_rootfs = self._deployment.effective_disk_limit_rootfs + self._config.disk_limit_log = self._deployment.effective_disk_limit_log self._clean_container_background() await self._setup_monitor() @@ -274,5 +276,7 @@ async def sandbox_info(self) -> SandboxInfo: "namespace": await self.namespace(), "cpus": self._config.cpus, "memory": self._config.memory, + "disk_limit_rootfs": self._config.disk_limit_rootfs, + "disk_limit_log": self._config.disk_limit_log, } return {} diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 8c320b66b2..a66cc13816 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -240,6 +240,8 @@ async def get_status(self, sandbox_id) -> SandboxStatusResponse: namespace=sandbox_info.get("namespace"), cpus=sandbox_info.get("cpus"), memory=sandbox_info.get("memory"), + disk_limit_rootfs=sandbox_info.get("disk_limit_rootfs"), + disk_limit_log=sandbox_info.get("disk_limit_log"), ) async def build_sandbox_info_from_redis(self, sandbox_id: str, deployment_info: SandboxInfo) -> SandboxInfo | None: @@ -357,3 +359,11 @@ def validate_sandbox_spec(self, runtime_config: RuntimeConfig, deployment_config except ValueError as e: logger.warning(f"Invalid memory size: {deployment_config.memory}", exc_info=e) raise BadRequestRockError(f"Invalid memory size: {deployment_config.memory}") + + # Validate disk_limit_rootfs format + if deployment_config.disk_limit_rootfs is not None: + try: + parse_size_to_bytes(deployment_config.disk_limit_rootfs) + except ValueError as e: + logger.warning(f"Invalid disk_limit_rootfs size: {deployment_config.disk_limit_rootfs}", exc_info=e) + raise BadRequestRockError(f"Invalid disk_limit_rootfs size: {deployment_config.disk_limit_rootfs}") diff --git a/rock/utils/docker.py b/rock/utils/docker.py index 6db0e747b7..41e7f99981 100644 --- a/rock/utils/docker.py +++ b/rock/utils/docker.py @@ -1,3 +1,4 @@ +import json import logging import subprocess @@ -7,6 +8,111 @@ class DockerUtil: """Docker operation utilities""" + @classmethod + def get_docker_info(cls) -> dict | None: + """Run 'docker info' and return the parsed JSON, or None on failure.""" + try: + result = subprocess.run( + ["docker", "info", "--format", "{{json .}}"], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode != 0: + logger.warning("get_docker_info: docker info failed") + return None + return json.loads(result.stdout) + except Exception as e: + logger.warning(f"get_docker_info: failed: {e}") + return None + + @classmethod + def get_docker_root_dir(cls) -> str | None: + """Return DockerRootDir from docker info, or None on failure.""" + info = cls.get_docker_info() + if info is None: + return None + root = info.get("DockerRootDir") + if not root: + logger.warning("get_docker_root_dir: DockerRootDir not found in docker info") + return root or None + + @classmethod + def is_xfs_prjquota_path(cls, path: str) -> bool: + """Return True if *path* is on an XFS mount with prjquota (or pquota) enabled. + + This is the prerequisite for XFS project quota on a directory. + Unlike detect_storage_opt_support(), this check is independent of Docker's + storage driver configuration. + """ + try: + result = subprocess.run( + ["findmnt", "-T", path, "-o", "FSTYPE,OPTIONS", "--noheadings"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode != 0: + logger.info(f"is_xfs_prjquota_path: findmnt failed for {path!r}") + return False + lines = [line.strip() for line in result.stdout.splitlines() if line.strip()] + if not lines: + logger.info(f"is_xfs_prjquota_path: empty findmnt output for {path!r}") + return False + # the last line is the most recent mount point + line = lines[-1] + parts = line.split(None, 1) # split on first whitespace: FSTYPE OPTIONS + if len(parts) < 2: + logger.info(f"is_xfs_prjquota_path: unexpected findmnt output for {path!r}: {line!r}") + return False + fstype, options = parts[0], parts[1] + if fstype != "xfs": + logger.info(f"is_xfs_prjquota_path: {path!r} is on {fstype!r}, not xfs") + return False + opts = options.split(",") + has_prjquota = "prjquota" in opts or "pquota" in opts + if not has_prjquota: + logger.info( + f"is_xfs_prjquota_path: {path!r} is xfs but mount options {options!r} missing prjquota/pquota" + ) + return has_prjquota + except Exception as e: + logger.info(f"is_xfs_prjquota_path: findmnt command failed for {path!r}: {e}") + return False + + @classmethod + def detect_storage_opt_support(cls) -> bool: + """Detect whether --storage-opt size= is supported in this environment. + + Requirements: + - Docker storage driver is overlay2 + - Docker root directory is on an XFS mount with prjquota/pquota enabled + (checked via is_xfs_prjquota_path) + + Returns: + True if --storage-opt size= can be used, False otherwise. + """ + info = cls.get_docker_info() + if info is None: + return False + + # Check 1: Driver must be overlay2 + if info.get("Driver") != "overlay2": + logger.info(f"detect_storage_opt_support: storage driver is {info.get('Driver')!r}, not overlay2") + return False + + # Check 2: DockerRootDir must be on XFS with prjquota/pquota + docker_root = info.get("DockerRootDir") + if not docker_root: + logger.warning("detect_storage_opt_support: DockerRootDir not found in docker info") + return False + + if not cls.is_xfs_prjquota_path(docker_root): + return False + + logger.info(f"detect_storage_opt_support: supported — overlay2, {docker_root!r} is xfs+prjquota") + return True + @classmethod def is_docker_available(cls): try: diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 3670dc3471..b39a9a385d 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -32,6 +32,18 @@ reason=f"Requires Docker and image {env_vars.ROCK_ENVHUB_DEFAULT_DOCKER_IMAGE}", ) +SKIP_IF_NO_STORAGE_OPT = pytest.mark.skipif( + not DockerUtil.detect_storage_opt_support(), + reason="Requires Docker with storage-opt support (overlay2 + xfs + prjquota/pquota)", +) + +_log_path = env_vars.ROCK_LOGGING_PATH or "" + +SKIP_IF_LOG_PATH_NOT_XFS = pytest.mark.skipif( + not _log_path or not DockerUtil.is_xfs_prjquota_path(_log_path), + reason=f"ROCK_LOGGING_PATH ({_log_path!r}) is not set or not on XFS with prjquota, skipping log quota test", +) + @dataclass class RemoteServer: diff --git a/tests/integration/sdk/sandbox/test_disk_limit.py b/tests/integration/sdk/sandbox/test_disk_limit.py new file mode 100644 index 0000000000..2b2fdf209e --- /dev/null +++ b/tests/integration/sdk/sandbox/test_disk_limit.py @@ -0,0 +1,229 @@ +"""Integration tests for disk limit functionality.""" + +import pytest + +from rock.actions import Command +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.utils.docker import DockerUtil +from tests.integration.conftest import SKIP_IF_LOG_PATH_NOT_XFS, SKIP_IF_NO_DOCKER, SKIP_IF_NO_STORAGE_OPT + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@SKIP_IF_NO_STORAGE_OPT +@pytest.mark.asyncio +async def test_disk_limit_enforcement(admin_remote_server): + """Test that the server-side rootfs disk limit is enforced when storage-opt is supported. + + This test is only run when storage-opt is supported (overlay2 + xfs + prjquota). + + Steps: + 1. Start a sandbox (server applies default disk_limit) + 2. Check sandbox status to verify disk_limit is reported + 3. Try to create a file larger than the limit (should fail) + 4. Create a small file (should succeed) + """ + config = SandboxConfig( + image="ubuntu:22.04", + memory="2g", + cpus=1.0, + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + startup_timeout=60, + ) + + sandbox = Sandbox(config) + await sandbox.start() + + try: + status = await sandbox.get_status() + print(f"Sandbox status: disk_limit={status.disk_limit_rootfs}") + + if status.disk_limit_rootfs is None: + pytest.skip( + "Server has no disk limit configured (sandbox_disk_limit_rootfs not set in rock-xxx.yml or nacos)" + ) + print(f"✅ Disk limit is set to {status.disk_limit_rootfs}") + + # Parse limit to determine a file size that exceeds it + result = await sandbox.execute( + Command( + command=[ + "/bin/bash", + "-c", + f"fallocate -l {status.disk_limit_rootfs.replace('g', '')}G /tmp/large_file.bin 2>&1 || echo 'EXPECTED_ERROR'", + ] + ) + ) + + output = result.stdout + result.stderr + print(f"fallocate output: {output}") + + error_occurred = ( + result.exit_code != 0 + or "No space left on device" in output + or "fallocate failed" in output + or "EXPECTED_ERROR" in output + ) + + assert error_occurred, ( + f"Expected disk space error when filling disk, " f"but got exit_code={result.exit_code}, output={output}" + ) + print("✅ Disk limit enforcement verified") + + small_file_result = await sandbox.execute( + Command(command=["/bin/bash", "-c", "fallocate -l 100M /tmp/small_file.bin && echo 'SUCCESS'"]) + ) + + small_output = small_file_result.stdout + small_file_result.stderr + assert small_file_result.exit_code == 0, ( + f"Expected small file (100MB) creation to succeed, " f"but got exit_code={small_file_result.exit_code}" + ) + assert "SUCCESS" in small_output + print("✅ Small file (100MB) creation succeeded") + + finally: + await sandbox.stop() + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_disk_limit_default_value(admin_remote_server): + """Test that the server applies a default disk_limit visible in status.""" + config = SandboxConfig( + image="ubuntu:22.04", + memory="2g", + cpus=1.0, + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + startup_timeout=60, + ) + + sandbox = Sandbox(config) + await sandbox.start() + + try: + status = await sandbox.get_status() + print(f"Sandbox status: disk_limit={status.disk_limit_rootfs}") + + storage_opt_supported = DockerUtil.detect_storage_opt_support() + + if not storage_opt_supported: + assert ( + status.disk_limit_rootfs is None + ), f"Expected disk_limit=None when storage-opt not supported, got {status.disk_limit_rootfs}" + print("✅ Storage-opt not supported: disk_limit is None") + else: + # When storage-opt is supported, disk_limit reflects server config (may be None if not configured) + print(f"✅ Server-reported disk_limit: {status.disk_limit_rootfs}") + + finally: + await sandbox.stop() + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@SKIP_IF_NO_STORAGE_OPT +@SKIP_IF_LOG_PATH_NOT_XFS +@pytest.mark.asyncio +async def test_logging_path_disk_limit_enforcement(admin_remote_server): + """Test that ROCK_LOGGING_PATH is also limited by disk quota. + + This test verifies that the log directory (ROCK_LOGGING_PATH) has a quota + enforced via XFS project quota, separate from the rootfs limit. + + Steps: + 1. Start a sandbox (server applies default log dir quota) + 2. Check that ROCK_LOGGING_PATH env var is set in container + 3. Try to create a file larger than the log quota (should fail) + 4. Create a 500MB file in ROCK_LOGGING_PATH (should succeed) + 5. Verify rootfs and log directory are independently limited + """ + config = SandboxConfig( + image="ubuntu:22.04", + memory="2g", + cpus=1.0, + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + startup_timeout=60, + ) + + sandbox = Sandbox(config) + await sandbox.start() + + try: + env_result = await sandbox.execute(Command(command=["/bin/bash", "-c", "echo $ROCK_LOGGING_PATH"])) + logging_path = env_result.stdout.strip() + print(f"ROCK_LOGGING_PATH in container: {logging_path}") + + assert logging_path, "ROCK_LOGGING_PATH should be set in container" + print(f"✅ ROCK_LOGGING_PATH is set to: {logging_path}") + + # Try to create a 1.5GB file in logging path (should fail due to log quota) + large_log_result = await sandbox.execute( + Command( + command=[ + "/bin/bash", + "-c", + f"fallocate -l 1500M {logging_path}/large_log.bin 2>&1 || echo 'EXPECTED_ERROR'", + ] + ) + ) + + large_output = large_log_result.stdout + large_log_result.stderr + print(f"Large log file creation output: {large_output}") + print(f"Large log file creation exit_code: {large_log_result.exit_code}") + + error_occurred = ( + large_log_result.exit_code != 0 + or "No space left on device" in large_output + or "Disk quota exceeded" in large_output + or "fallocate failed" in large_output + or "EXPECTED_ERROR" in large_output + ) + + assert error_occurred, ( + f"Expected disk quota error when creating 1.5GB file in log dir, " + f"but got exit_code={large_log_result.exit_code}, output={large_output}" + ) + print("✅ Log directory quota verified: 1.5GB file creation failed as expected") + + small_log_result = await sandbox.execute( + Command( + command=[ + "/bin/bash", + "-c", + f"fallocate -l 500M {logging_path}/small_log.bin && echo 'SUCCESS'", + ] + ) + ) + + small_output = small_log_result.stdout + small_log_result.stderr + print(f"Small log file creation output: {small_output}") + assert small_log_result.exit_code == 0, ( + f"Expected small log file (500MB) creation to succeed, " + f"but got exit_code={small_log_result.exit_code}, output={small_output}" + ) + assert "SUCCESS" in small_output + print("✅ Small log file (500MB) creation in ROCK_LOGGING_PATH succeeded") + + rootfs_result = await sandbox.execute( + Command( + command=[ + "/bin/bash", + "-c", + "fallocate -l 1G /tmp/rootfs_file.bin && echo 'ROOTFS_SUCCESS'", + ] + ) + ) + + rootfs_output = rootfs_result.stdout + rootfs_result.stderr + print(f"Rootfs file creation output: {rootfs_output}") + assert rootfs_result.exit_code == 0, ( + f"Expected 1GB file on rootfs to succeed, " + f"but got exit_code={rootfs_result.exit_code}, output={rootfs_output}" + ) + assert "ROOTFS_SUCCESS" in rootfs_output + print("✅ Rootfs and log directory are independently limited") + + finally: + await sandbox.stop() diff --git a/tests/unit/admin/proto/test_sandbox_response.py b/tests/unit/admin/proto/test_sandbox_response.py new file mode 100644 index 0000000000..d4868ab10d --- /dev/null +++ b/tests/unit/admin/proto/test_sandbox_response.py @@ -0,0 +1,142 @@ +""" +Unit tests for admin proto response models — disk_limit_rootfs and disk_limit_log fields. + +Tests cover: +- SandboxStartResponse.disk_limit_rootfs / disk_limit_log fields +- SandboxStatusResponse.disk_limit_rootfs / disk_limit_log fields +- SandboxStatusResponse.from_sandbox_info() extraction of both fields +""" + +from rock.admin.proto.response import SandboxStartResponse, SandboxStatusResponse + +# ---- SandboxStartResponse tests ---- + + +class TestSandboxStartResponseDiskLimit: + def test_disk_limit_rootfs_default_is_none(self): + response = SandboxStartResponse() + assert response.disk_limit_rootfs is None + + def test_disk_limit_log_default_is_none(self): + response = SandboxStartResponse() + assert response.disk_limit_log is None + + def test_disk_limit_rootfs_set_value(self): + response = SandboxStartResponse(disk_limit_rootfs="20g") + assert response.disk_limit_rootfs == "20g" + + def test_disk_limit_log_set_value(self): + response = SandboxStartResponse(disk_limit_log="5g") + assert response.disk_limit_log == "5g" + + def test_all_fields_with_both_limits(self): + response = SandboxStartResponse( + sandbox_id="test-sandbox", + host_ip="10.0.0.1", + cpus=4.0, + memory="16g", + disk_limit_rootfs="50g", + disk_limit_log="5g", + ) + assert response.sandbox_id == "test-sandbox" + assert response.disk_limit_rootfs == "50g" + assert response.disk_limit_log == "5g" + assert response.cpus == 4.0 + assert response.memory == "16g" + + +# ---- SandboxStatusResponse tests ---- + + +class TestSandboxStatusResponseDiskLimit: + def test_disk_limit_rootfs_default_is_none(self): + response = SandboxStatusResponse() + assert response.disk_limit_rootfs is None + + def test_disk_limit_log_default_is_none(self): + response = SandboxStatusResponse() + assert response.disk_limit_log is None + + def test_disk_limit_rootfs_set_value(self): + response = SandboxStatusResponse(disk_limit_rootfs="20g") + assert response.disk_limit_rootfs == "20g" + + def test_disk_limit_log_set_value(self): + response = SandboxStatusResponse(disk_limit_log="5g") + assert response.disk_limit_log == "5g" + + def test_from_sandbox_info_with_both_limits(self): + """from_sandbox_info() should extract both limit fields from SandboxInfo dict.""" + sandbox_info = { + "sandbox_id": "test-sandbox", + "phases": {}, + "port_mapping": {}, + "host_ip": "10.0.0.1", + "cpus": 2.0, + "memory": "8g", + "disk_limit_rootfs": "30g", + "disk_limit_log": "5g", + } + response = SandboxStatusResponse.from_sandbox_info(sandbox_info) + assert response.disk_limit_rootfs == "30g" + assert response.disk_limit_log == "5g" + assert response.cpus == 2.0 + assert response.memory == "8g" + + def test_from_sandbox_info_without_limits(self): + """from_sandbox_info() should yield None for both when absent.""" + sandbox_info = { + "sandbox_id": "test-sandbox", + "phases": {}, + "port_mapping": {}, + "cpus": 2.0, + "memory": "8g", + } + response = SandboxStatusResponse.from_sandbox_info(sandbox_info) + assert response.disk_limit_rootfs is None + assert response.disk_limit_log is None + + def test_from_sandbox_info_with_none_limits(self): + """from_sandbox_info() should surface None when fields are explicitly None.""" + sandbox_info = { + "sandbox_id": "test-sandbox", + "phases": {}, + "port_mapping": {}, + "disk_limit_rootfs": None, + "disk_limit_log": None, + } + response = SandboxStatusResponse.from_sandbox_info(sandbox_info) + assert response.disk_limit_rootfs is None + assert response.disk_limit_log is None + + def test_from_sandbox_info_partial_limits(self): + """from_sandbox_info() handles one field set, one absent.""" + sandbox_info = { + "sandbox_id": "test-sandbox", + "phases": {}, + "port_mapping": {}, + "disk_limit_rootfs": "50g", + } + response = SandboxStatusResponse.from_sandbox_info(sandbox_info) + assert response.disk_limit_rootfs == "50g" + assert response.disk_limit_log is None + + +# ---- actions/sandbox/response.SandboxStatusResponse tests ---- + + +class TestActionsSandboxStatusResponseDiskLimit: + def test_actions_status_response_both_limits(self): + """rock.actions.sandbox.response.SandboxStatusResponse should have both limit fields.""" + from rock.actions.sandbox.response import SandboxStatusResponse as ActionStatusResponse + + response = ActionStatusResponse(disk_limit_rootfs="20g", disk_limit_log="5g") + assert response.disk_limit_rootfs == "20g" + assert response.disk_limit_log == "5g" + + def test_actions_status_response_defaults_none(self): + from rock.actions.sandbox.response import SandboxStatusResponse as ActionStatusResponse + + response = ActionStatusResponse() + assert response.disk_limit_rootfs is None + assert response.disk_limit_log is None diff --git a/tests/unit/deployments/test_docker_deployment_disk_limit.py b/tests/unit/deployments/test_docker_deployment_disk_limit.py new file mode 100644 index 0000000000..a94f3fd164 --- /dev/null +++ b/tests/unit/deployments/test_docker_deployment_disk_limit.py @@ -0,0 +1,245 @@ +""" +Unit tests for disk_limit support in DockerDeployment and DockerDeploymentConfig. + +Tests cover: +- DockerDeploymentConfig default and custom disk_limit_rootfs / disk_limit_log values +- DockerDeployment._storage_opts() argument generation +- DockerDeployment.start() graceful degradation when storage-opt is unsupported +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.deployments.config import DockerDeploymentConfig +from rock.deployments.docker import DockerDeployment + +# ---- DockerDeploymentConfig tests ---- + + +class TestDockerDeploymentConfigDiskLimit: + def test_default_disk_limit_rootfs_is_none(self): + config = DockerDeploymentConfig() + assert config.disk_limit_rootfs is None + + def test_default_disk_limit_log_is_none(self): + config = DockerDeploymentConfig() + assert config.disk_limit_log is None + + def test_custom_disk_limit_rootfs(self): + config = DockerDeploymentConfig(disk_limit_rootfs="50g") + assert config.disk_limit_rootfs == "50g" + + def test_custom_disk_limit_log(self): + config = DockerDeploymentConfig(disk_limit_log="5g") + assert config.disk_limit_log == "5g" + + def test_disk_limit_rootfs_none(self): + config = DockerDeploymentConfig(disk_limit_rootfs=None) + assert config.disk_limit_rootfs is None + + def test_disk_limit_log_none(self): + config = DockerDeploymentConfig(disk_limit_log=None) + assert config.disk_limit_log is None + + def test_disk_limit_rootfs_preserved_in_model_dump(self): + config = DockerDeploymentConfig(disk_limit_rootfs="50g") + dump = config.model_dump() + assert dump["disk_limit_rootfs"] == "50g" + + def test_disk_limit_log_preserved_in_model_dump(self): + config = DockerDeploymentConfig(disk_limit_log="5g") + dump = config.model_dump() + assert dump["disk_limit_log"] == "5g" + + def test_disk_limit_rootfs_none_preserved_in_model_dump(self): + config = DockerDeploymentConfig(disk_limit_rootfs=None) + dump = config.model_dump() + assert dump["disk_limit_rootfs"] is None + + +# ---- DockerDeployment._storage_opts() tests ---- + + +class TestStorageOpts: + """Tests for DockerDeployment._storage_opts() method.""" + + @patch("rock.deployments.docker.DockerSandboxValidator") + def test_storage_opts_with_disk_limit_rootfs(self, _mock_validator): + deployment = DockerDeployment.from_config(DockerDeploymentConfig(disk_limit_rootfs="30g")) + result = deployment._storage_opts() + assert result == ["--storage-opt", "size=30g"] + + @patch("rock.deployments.docker.DockerSandboxValidator") + def test_storage_opts_with_none(self, _mock_validator): + deployment = DockerDeployment.from_config(DockerDeploymentConfig(disk_limit_rootfs=None)) + result = deployment._storage_opts() + assert result == [] + + @patch("rock.deployments.docker.DockerSandboxValidator") + def test_storage_opts_default_value(self, _mock_validator): + deployment = DockerDeployment.from_config(DockerDeploymentConfig()) + result = deployment._storage_opts() + assert result == [] + + @patch("rock.deployments.docker.DockerSandboxValidator") + def test_storage_opts_various_sizes(self, _mock_validator): + for size in ("1g", "512m", "50g", "1t"): + deployment = DockerDeployment.from_config(DockerDeploymentConfig(disk_limit_rootfs=size)) + result = deployment._storage_opts() + assert result == ["--storage-opt", f"size={size}"] + + +# ---- DockerDeployment.start() storage-opt degradation tests ---- + + +def _make_start_mocks(deployment): + deployment.sandbox_validator = MagicMock() + deployment.sandbox_validator.check_availability.return_value = True + deployment.sandbox_validator.check_resource.return_value = True + deployment._pull_image = MagicMock() + deployment.do_port_mapping = AsyncMock() + deployment._prepare_volume_mounts = MagicMock(return_value=[]) + deployment._start_container = AsyncMock() + deployment._wait_until_alive = AsyncMock() + deployment._service_status = MagicMock() + deployment._service_status.get_mapped_port = MagicMock(return_value=8080) + deployment._service_status.phases = {} + + +async def _run_start(deployment): + with ( + patch("rock.deployments.docker.get_executor"), + patch("rock.deployments.docker.asyncio.get_running_loop") as mock_loop, + patch("rock.deployments.docker.wait_until_alive", new_callable=AsyncMock), + patch("rock.deployments.docker.env_vars") as mock_env, + patch("rock.deployments.docker.subprocess"), + ): + mock_env.ROCK_LOGGING_PATH = "" + mock_env.ROCK_TIME_ZONE = "UTC" + mock_loop.return_value.run_in_executor = AsyncMock() + try: + await deployment.start() + except Exception: + pass + + +class TestDockerDeploymentStartDiskLimit: + """Tests that start() applies correct effective values for rootfs and log quotas.""" + + @pytest.mark.asyncio + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=False) + async def test_rootfs_downgraded_when_storage_opt_unsupported(self, _mock_detect, _mock_validator): + """When storage-opt NOT supported: effective_disk_limit_rootfs=None; config unchanged.""" + config = DockerDeploymentConfig(disk_limit_rootfs="50g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + _make_start_mocks(deployment) + await _run_start(deployment) + + assert deployment.config.disk_limit_rootfs == "50g" + assert deployment.effective_disk_limit_rootfs is None + + @pytest.mark.asyncio + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=True) + async def test_rootfs_preserved_when_storage_opt_supported(self, _mock_detect, _mock_validator): + """When storage-opt IS supported: effective_disk_limit_rootfs matches config.""" + config = DockerDeploymentConfig(disk_limit_rootfs="50g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + _make_start_mocks(deployment) + await _run_start(deployment) + + assert deployment.config.disk_limit_rootfs == "50g" + assert deployment.effective_disk_limit_rootfs == "50g" + + @pytest.mark.asyncio + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=False) + async def test_no_error_when_rootfs_already_none(self, _mock_detect, _mock_validator): + """When disk_limit_rootfs is None: start() should not error.""" + config = DockerDeploymentConfig(disk_limit_rootfs=None, image="python:3.11") + deployment = DockerDeployment.from_config(config) + _make_start_mocks(deployment) + await _run_start(deployment) + + assert deployment.config.disk_limit_rootfs is None + assert deployment.effective_disk_limit_rootfs is None + + @pytest.mark.asyncio + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=True) + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + async def test_log_downgraded_when_not_xfs_prjquota(self, _mock_prjquota, _mock_detect, _mock_validator): + """When log path is not XFS+prjquota: effective_disk_limit_log=None; config unchanged. + + Note: log quota has NO dependency on docker being overlay2 — + is_xfs_prjquota_path() is the only gate. + """ + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + _make_start_mocks(deployment) + + with ( + patch("rock.deployments.docker.get_executor"), + patch("rock.deployments.docker.asyncio.get_running_loop") as mock_loop, + patch("rock.deployments.docker.wait_until_alive", new_callable=AsyncMock), + patch("rock.deployments.docker.env_vars") as mock_env, + patch("rock.deployments.docker.subprocess"), + ): + mock_env.ROCK_LOGGING_PATH = "/var/log/rock" + mock_env.ROCK_TIME_ZONE = "UTC" + mock_loop.return_value.run_in_executor = AsyncMock() + try: + await deployment.start() + except Exception: + pass + + assert deployment.config.disk_limit_log == "5g" + assert deployment.effective_disk_limit_log is None + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + def test_log_not_downgraded_when_no_log_path(self, _mock_prjquota, _mock_validator): + """When ROCK_LOGGING_PATH is empty, _try_set_log_dir_quota is never called, + so effective_disk_limit_log remains equal to config.disk_limit_log.""" + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + # effective starts equal to config before start() is called + assert deployment.effective_disk_limit_log == "5g" + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + def test_try_set_log_dir_quota_downgrades_when_not_xfs_prjquota(self, _mock_prjquota, _mock_validator): + """_try_set_log_dir_quota: is_xfs_prjquota_path=False → effective_disk_limit_log=None.""" + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + deployment._effective_disk_limit_log = "5g" + deployment._container_name = "test-container" + + deployment._try_set_log_dir_quota("/var/log/rock/test-container") + + assert deployment.effective_disk_limit_log is None + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=True) + def test_try_set_log_dir_quota_independent_of_docker_driver(self, _mock_prjquota, _mock_validator): + """_try_set_log_dir_quota passes the XFS gate regardless of Docker storage driver. + + Log quota only requires is_xfs_prjquota_path(); overlay2 is irrelevant. + The subprocess calls inside (findmnt, xfs_quota) are mocked to succeed. + """ + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + deployment._effective_disk_limit_log = "5g" + deployment._container_name = "test-container" + + with patch("rock.deployments.docker.subprocess") as mock_sub: + ok = MagicMock() + ok.returncode = 0 + ok.stdout = "/var/log/rock" + mock_sub.run.return_value = ok + deployment._try_set_log_dir_quota("/var/log/rock/test-container") + + # xfs_quota succeeded → effective value preserved + assert deployment.effective_disk_limit_log == "5g" diff --git a/tests/unit/sandbox/job/__init__.py b/tests/unit/sandbox/job/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/sandbox/test_sandbox_manager_disk_limit.py b/tests/unit/sandbox/test_sandbox_manager_disk_limit.py new file mode 100644 index 0000000000..d77b401e6a --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_manager_disk_limit.py @@ -0,0 +1,70 @@ +""" +Unit tests for SandboxManager.validate_sandbox_spec() — disk_limit_rootfs validation. + +These tests do NOT require Ray or Docker; they only test the synchronous +validation logic. +""" + +import pytest + +from rock.config import RuntimeConfig, StandardSpec +from rock.deployments.config import DockerDeploymentConfig +from rock.sandbox.sandbox_manager import SandboxManager +from rock.sdk.common.exceptions import BadRequestRockError + + +@pytest.fixture +def runtime_config(): + return RuntimeConfig( + max_allowed_spec=StandardSpec(cpus=16, memory="64g"), + ) + + +class TestValidateSandboxSpecDiskLimit: + """Tests for disk_limit_rootfs validation in SandboxManager.validate_sandbox_spec().""" + + def test_valid_disk_limit_rootfs_20g(self, runtime_config): + config = DockerDeploymentConfig(disk_limit_rootfs="20g") + # Should not raise + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_valid_disk_limit_rootfs_various_formats(self, runtime_config): + for size in ("1g", "512m", "100gb", "2t", "1024mb", "1024k"): + config = DockerDeploymentConfig(disk_limit_rootfs=size) + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_valid_disk_limit_rootfs_none(self, runtime_config): + """None disk_limit_rootfs should skip validation (no error).""" + config = DockerDeploymentConfig(disk_limit_rootfs=None) + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_invalid_disk_limit_rootfs_raises_bad_request(self, runtime_config): + config = DockerDeploymentConfig(disk_limit_rootfs="not-a-size") + with pytest.raises(BadRequestRockError, match="Invalid disk_limit_rootfs size"): + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_invalid_disk_limit_rootfs_empty_string(self, runtime_config): + config = DockerDeploymentConfig(disk_limit_rootfs="") + with pytest.raises(BadRequestRockError, match="Invalid disk_limit_rootfs size"): + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_invalid_disk_limit_rootfs_negative(self, runtime_config): + config = DockerDeploymentConfig(disk_limit_rootfs="-10g") + with pytest.raises(BadRequestRockError, match="Invalid disk_limit_rootfs size"): + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_invalid_disk_limit_rootfs_no_unit(self, runtime_config): + """A bare number without unit should still be parsed (as bytes).""" + config = DockerDeploymentConfig(disk_limit_rootfs="1024") + # Bare number is treated as bytes by parse_size_to_bytes, so it should pass + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_invalid_disk_limit_rootfs_only_unit(self, runtime_config): + config = DockerDeploymentConfig(disk_limit_rootfs="gb") + with pytest.raises(BadRequestRockError, match="Invalid disk_limit_rootfs size"): + SandboxManager.validate_sandbox_spec(None, runtime_config, config) + + def test_disk_limit_rootfs_validation_independent_of_cpu_memory(self, runtime_config): + """disk_limit_rootfs validation should not interfere with cpu/memory checks.""" + config = DockerDeploymentConfig(cpus=2, memory="8g", disk_limit_rootfs="50g") + SandboxManager.validate_sandbox_spec(None, runtime_config, config) diff --git a/tests/unit/utils/test_docker_util.py b/tests/unit/utils/test_docker_util.py new file mode 100644 index 0000000000..e6745deb7f --- /dev/null +++ b/tests/unit/utils/test_docker_util.py @@ -0,0 +1,181 @@ +""" +Unit tests for DockerUtil.detect_storage_opt_support(). + +All subprocess calls are mocked so no Docker daemon is required. +""" + +import json +import subprocess +from unittest.mock import MagicMock, patch + +import pytest + +from rock.utils.docker import DockerUtil + + +# ---- helpers ---- + + +def _make_run_result(returncode=0, stdout="", stderr=""): + r = MagicMock(spec=subprocess.CompletedProcess) + r.returncode = returncode + r.stdout = stdout + r.stderr = stderr + return r + + +def _docker_info_json(driver="overlay2", backing_fs="xfs", docker_root="/var/lib/docker"): + """Return a JSON string mimicking `docker info --format '{{json .}}'`.""" + info = { + "Driver": driver, + "DriverStatus": [ + ["Backing Filesystem", backing_fs], + ["Supports d_type", "true"], + ], + "DockerRootDir": docker_root, + } + return json.dumps(info) + + +# ---- detect_storage_opt_support tests ---- + + +class TestDetectStorageOptSupport: + """Tests for DockerUtil.detect_storage_opt_support().""" + + # findmnt now returns "FSTYPE OPTIONS" (delegated to is_xfs_prjquota_path) + @patch("rock.utils.docker.subprocess.run") + def test_all_requirements_met_prjquota(self, mock_run): + """Should return True when overlay2 + xfs + prjquota are all present.""" + mock_run.side_effect = [ + # docker info + _make_run_result(stdout=_docker_info_json()), + # findmnt for DockerRootDir (FSTYPE OPTIONS) + _make_run_result(stdout="xfs rw,relatime,attr2,inode64,prjquota"), + ] + assert DockerUtil.detect_storage_opt_support() is True + + @patch("rock.utils.docker.subprocess.run") + def test_all_requirements_met_pquota(self, mock_run): + """Should return True when pquota (synonym for prjquota) is present.""" + mock_run.side_effect = [ + _make_run_result(stdout=_docker_info_json()), + _make_run_result(stdout="xfs rw,relatime,attr2,inode64,pquota"), + ] + assert DockerUtil.detect_storage_opt_support() is True + + @patch("rock.utils.docker.subprocess.run") + def test_non_overlay2_driver(self, mock_run): + """Should return False if the storage driver is not overlay2.""" + mock_run.return_value = _make_run_result(stdout=_docker_info_json(driver="aufs")) + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_non_xfs_docker_root(self, mock_run): + """Should return False if DockerRootDir is not on XFS.""" + mock_run.side_effect = [ + _make_run_result(stdout=_docker_info_json()), + _make_run_result(stdout="ext4 rw,relatime"), + ] + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_missing_prjquota(self, mock_run): + """Should return False when mount options lack prjquota/pquota.""" + mock_run.side_effect = [ + _make_run_result(stdout=_docker_info_json()), + _make_run_result(stdout="xfs rw,relatime,attr2,inode64,noquota"), + ] + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_docker_info_fails(self, mock_run): + """Should return False if docker info returns non-zero exit code.""" + mock_run.return_value = _make_run_result(returncode=1, stderr="Cannot connect to Docker daemon") + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_docker_info_raises_exception(self, mock_run): + """Should return False if docker info subprocess raises.""" + mock_run.side_effect = FileNotFoundError("docker not found") + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_findmnt_fails(self, mock_run): + """Should return False if findmnt returns non-zero exit code.""" + mock_run.side_effect = [ + _make_run_result(stdout=_docker_info_json()), + _make_run_result(returncode=1, stderr="findmnt: failed"), + ] + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_findmnt_raises_exception(self, mock_run): + """Should return False if findmnt subprocess raises.""" + mock_run.side_effect = [ + _make_run_result(stdout=_docker_info_json()), + subprocess.TimeoutExpired(cmd="findmnt", timeout=5), + ] + assert DockerUtil.detect_storage_opt_support() is False + + @patch("rock.utils.docker.subprocess.run") + def test_missing_docker_root_dir(self, mock_run): + """Should return False if DockerRootDir is missing from docker info.""" + info = {"Driver": "overlay2", "DriverStatus": [["Backing Filesystem", "xfs"]]} + mock_run.return_value = _make_run_result(stdout=json.dumps(info)) + assert DockerUtil.detect_storage_opt_support() is False + + +# ---- is_xfs_prjquota_path tests ---- + + +class TestIsXfsPrjquotaPath: + """Tests for DockerUtil.is_xfs_prjquota_path(). + + Unlike detect_storage_opt_support(), this check is path-local and has + no dependency on Docker's storage driver. + """ + + @patch("rock.utils.docker.subprocess.run") + def test_xfs_with_prjquota(self, mock_run): + mock_run.return_value = _make_run_result(stdout="xfs rw,relatime,attr2,inode64,prjquota") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is True + + @patch("rock.utils.docker.subprocess.run") + def test_xfs_with_pquota_synonym(self, mock_run): + """pquota is a synonym for prjquota and must also be accepted.""" + mock_run.return_value = _make_run_result(stdout="xfs rw,relatime,attr2,inode64,pquota") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is True + + @patch("rock.utils.docker.subprocess.run") + def test_xfs_without_prjquota(self, mock_run): + """XFS mount without prjquota/pquota should return False.""" + mock_run.return_value = _make_run_result(stdout="xfs rw,relatime,attr2,inode64,noquota") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False + + @patch("rock.utils.docker.subprocess.run") + def test_non_xfs_with_prjquota(self, mock_run): + """ext4 with prjquota-like options should return False (not XFS).""" + mock_run.return_value = _make_run_result(stdout="ext4 rw,relatime,prjquota") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False + + @patch("rock.utils.docker.subprocess.run") + def test_findmnt_failure(self, mock_run): + mock_run.return_value = _make_run_result(returncode=1, stderr="findmnt: failed") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False + + @patch("rock.utils.docker.subprocess.run") + def test_findmnt_exception(self, mock_run): + mock_run.side_effect = subprocess.TimeoutExpired(cmd="findmnt", timeout=5) + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False + + @patch("rock.utils.docker.subprocess.run") + def test_empty_output(self, mock_run): + mock_run.return_value = _make_run_result(stdout="") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False + + @patch("rock.utils.docker.subprocess.run") + def test_only_fstype_no_options(self, mock_run): + """If findmnt returns only FSTYPE with no OPTIONS column, return False.""" + mock_run.return_value = _make_run_result(stdout="xfs") + assert DockerUtil.is_xfs_prjquota_path("/data/logs") is False From 07fdd539e3184ccccbc81ce80d7f9abe718de1f1 Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Wed, 22 Apr 2026 11:09:23 +0800 Subject: [PATCH 059/226] fix(uv-env): copy project to writable dir before install in container (#81377578) (#857) --- .rayignore | 7 +++++++ rock/rocklet/local_files/docker_run_with_uv.sh | 14 ++++++++++++-- 2 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 .rayignore diff --git a/.rayignore b/.rayignore new file mode 100644 index 0000000000..1828f981c3 --- /dev/null +++ b/.rayignore @@ -0,0 +1,7 @@ +*.egg-info/ +.venv/ +.git/ +__pycache__/ +*.pyc +.pytest_cache/ +.ruff_cache/ diff --git a/rock/rocklet/local_files/docker_run_with_uv.sh b/rock/rocklet/local_files/docker_run_with_uv.sh index 150cf14e2b..c5b9b3df69 100755 --- a/rock/rocklet/local_files/docker_run_with_uv.sh +++ b/rock/rocklet/local_files/docker_run_with_uv.sh @@ -32,13 +32,23 @@ if [ ! -f /etc/alpine-release ]; then UV_CMD=$HOME/.local/bin/uv fi - cd $PROJECT_ROOT + # Copy project to a writable directory (source mount is read-only) + # Use tar to exclude large unnecessary directories for faster copy + WRITABLE_PROJECT=/tmp/rock-build + mkdir -p $WRITABLE_PROJECT + tar -cf - --exclude='.venv' --exclude='.git' --exclude='__pycache__' \ + --exclude='*.egg-info' --exclude='.pytest_cache' --exclude='.ruff_cache' \ + -C $PROJECT_ROOT . | tar -xf - -C $WRITABLE_PROJECT + cd $WRITABLE_PROJECT # Create virtual environment $UV_CMD venv --python 3.11 /tmp/rocklet-venv # Install dependencies - $UV_CMD pip install --python /tmp/rocklet-venv/bin/python -e ".[rocklet]" + $UV_CMD pip install --python /tmp/rocklet-venv/bin/python ".[rocklet]" + + # Clean up build directory to free disk space + rm -rf $WRITABLE_PROJECT mkdir -p /data/logs From 41a296e85e65efba2768e14d4e9a3bb2b0dea833 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Tue, 21 Apr 2026 17:43:36 +0800 Subject: [PATCH 060/226] fix(proxy): forward whitelisted headers in WebSocket proxy upstream handshake (#865) WebSocket proxy was losing client Origin, Authorization, Cookie, and tracing headers when initiating the second-hop handshake to downstream services. Add whitelist-based header forwarding via build_upstream_ws_headers() and pass origin=/additional_headers= to websockets.connect(). Co-Authored-By: Claude Opus 4.6 --- .../proxy-enhancements/01_requirement.md | 32 +++- .../proxy-enhancements/03_implementation.md | 104 +++++++++++- rock/sandbox/service/sandbox_proxy_service.py | 32 ++++ .../sandbox/test_websocket_proxy_headers.py | 148 ++++++++++++++++++ 4 files changed, 310 insertions(+), 6 deletions(-) create mode 100644 tests/unit/sandbox/test_websocket_proxy_headers.py diff --git a/docs/_specs/proxy-enhancements/01_requirement.md b/docs/_specs/proxy-enhancements/01_requirement.md index 76fc19bdb7..dafa3d84c0 100644 --- a/docs/_specs/proxy-enhancements/01_requirement.md +++ b/docs/_specs/proxy-enhancements/01_requirement.md @@ -7,10 +7,16 @@ ROCK Admin 目前提供两类代理能力： 1. **WebSocket Proxy** (`/sandboxes/{id}/proxy/ws[/{path}]`)：将 WebSocket 连接转发到沙箱内的固定 SERVER 端口（`Port.SERVER = 8080`），不支持用户指定目标端口。 2. **HTTP Proxy** (`/sandboxes/{sandbox_id}/proxy[/{path}]`)：将 HTTP 请求转发到沙箱内的固定 SERVER 端口（`Port.SERVER = 8080`），且硬编码为 `method="POST"`，不支持 GET / PUT / DELETE / PATCH 等其他 HTTP 方法，也不支持用户指定目标端口。 -这三个限制阻碍了以下场景： +除了端口和 method 能力不足之外，WebSocket proxy 还有一个上下文丢失问题： +- Admin 在转发 WebSocket 握手到下游服务时，目前只处理 `Sec-WebSocket-Protocol` 子协议，不透传通用请求头 +- 当下游服务依赖 `Origin`、`Authorization`、`Cookie`、`X-Forwarded-*` 等头做来源校验、认证、会话恢复或审计时，二跳握手会丢失这些上下文 +- 典型失败现象是下游日志出现 `origin not allowed`，或者因为缺失 token / cookie 导致握手或后续鉴权失败 + +这些限制阻碍了以下场景： - 沙箱内运行了多个 WebSocket 服务（如 Jupyter Kernel、VS Code Server、自定义推理服务），需要连接到不同端口 - 沙箱内服务使用 RESTful 风格 API，需要 GET 查询、PUT 更新、DELETE 删除 - 沙箱内运行了多个 HTTP 服务，需要访问非 8080 端口 +- 沙箱内 WebSocket 服务依赖浏览器来源校验、认证头、cookie 或链路追踪头，要求代理保留客户端请求上下文 --- @@ -33,12 +39,22 @@ ROCK Admin 目前提供两类代理能力： - 在现有 `/sandboxes/{sandbox_id}/proxy[/{path}]` 路由上，允许通过 query param `port` 指定目标 HTTP 端口 - 当 `port` 未指定时，保持现有行为（使用 `Port.SERVER = 8080`） +4. **WebSocket Proxy 支持白名单透传通用请求头** + - 通过 `/sandboxes/{id}/proxy/{path:path}` 建立 WebSocket 代理时，允许将客户端请求中的通用 header 按白名单转发到下游服务 + - 首批需要覆盖的 header 包括：`Origin`、`Authorization`、`Cookie`、`X-Forwarded-For`、`X-Forwarded-Host`、`X-Forwarded-Proto`、`X-Real-IP`、`X-Request-Id`、`Traceparent`、`Tracestate`、`EagleEye-TraceId`、`EagleEye-RpcId`、`EagleEye-UserData` + - `Origin` 是必须支持的关键头，因为下游服务可能使用来源白名单（例如 `gateway.controlUi.allowedOrigins`）校验 WebSocket 握手 + - `Sec-WebSocket-Protocol` 继续通过 WebSocket 子协议协商传递，不作为普通 header 透传 + - 采用白名单策略，避免将握手专用头和 hop-by-hop 头错误转发给下游 + ### Out（本次不做的） - WebSocket Proxy 的认证/鉴权增强 - HTTP Proxy 的 multipart/form-data 支持（upload 接口已单独处理） - `host_proxy` 的 method 扩展（范围外） - SDK 客户端侧的封装更新 +- WebSocket 代理透传“所有”请求头，或由客户端动态指定任意 header 白名单 +- 自动伪造、补默认值或重写任意 `Origin` +- 透传 WebSocket 握手专用头（如 `Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`） --- @@ -53,6 +69,11 @@ ROCK Admin 目前提供两类代理能力： - **AC7**：SSE streaming（`text/event-stream`）在所有 method 下仍然正常工作 - **AC8**：`GET /sandboxes/{sandbox_id}/proxy/api/health?port=9000` 能成功代理到沙箱内 9000 端口的 HTTP 服务 - **AC9**：HTTP proxy 不带 `port` 参数时行为不变（向后兼容） +- **AC10**：当客户端 WebSocket 握手包含 `Origin` 时，代理发起到下游的二跳握手必须携带相同 `Origin`，以满足下游来源校验 +- **AC11**：当客户端握手包含 `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent` 等白名单头时，代理发起到下游的二跳握手必须一并转发 +- **AC12**：`Sec-WebSocket-Protocol` 必须继续通过现有 `subprotocols` 机制转发和协商，不能降级为普通 header 透传 +- **AC13**：`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 等握手专用头不得被转发到下游 +- **AC14**：当客户端未携带任何白名单头时，WebSocket proxy 的默认行为与现状保持一致（向后兼容） --- @@ -62,11 +83,16 @@ ROCK Admin 目前提供两类代理能力： - 不修改 `Port.SERVER` / `Port.PROXY` 枚举定义 - `port_validation` 逻辑复用现有 `validate_port_forward_port`（但 WebSocket proxy 端口校验需新增对 `Port.SERVER = 8080` 的允许，或直接用同一校验函数） - 保持现有 portforward 端点（`/sandboxes/{id}/portforward`）不变 +- WebSocket header 转发采用**白名单**而不是**黑名单**策略，避免误传握手专用头 +- `Origin` 通过 WebSocket 客户端库的显式参数透传，不与普通 `additional_headers` 混用 +- `Sec-WebSocket-Protocol` 继续通过 `subprotocols` 参数处理，不走通用 header 透传逻辑 --- ## Risks & Rollout - **风险**：WebSocket proxy 中用户可以指定任意端口访问沙箱内服务，存在横向访问风险 → 通过 `sandbox_id` 鉴权已覆盖，端口范围校验作为防护层 -- **回滚**：修改仅在 `sandbox_proxy_api.py` 和 `sandbox_proxy_service.py` 内，回滚只需还原这两个文件 -- **上线策略**：无数据库变更，直接部署 +- **风险**：若白名单范围定义过宽，可能把不该下传的敏感或握手专用头带给下游 → 通过固定白名单和单元测试约束范围 +- **风险**：不同下游服务对 `Origin`、`Cookie`、`Authorization` 的要求不同，透传后暴露出原本被代理层掩盖的问题 → 以“尽可能保留客户端上下文、但不伪造默认值”为原则 +- **回滚**：修改集中在 `sandbox_proxy_service.py` 和对应单元测试；如需回滚，可仅还原 WebSocket header 透传逻辑 +- **上线策略**：无数据库变更，直接部署；建议先在依赖 `Origin` 校验的控制台类服务上验证 diff --git a/docs/_specs/proxy-enhancements/03_implementation.md b/docs/_specs/proxy-enhancements/03_implementation.md index c4fce8255c..f7254a6e11 100644 --- a/docs/_specs/proxy-enhancements/03_implementation.md +++ b/docs/_specs/proxy-enhancements/03_implementation.md @@ -7,6 +7,11 @@ admin 与 sandbox 不在同一 K8s 集群，`host_ip` 为宿主机 IP，容器 - **WebSocket proxy 自定义端口**：复用 rocklet 现有的 `/portforward` WebSocket 端点中转（与 `/sandboxes/{id}/portforward` 相同机制） - **HTTP proxy 自定义端口**：需在 rocklet 新增 `/http_proxy` HTTP 端点，admin 转发请求给 rocklet，rocklet 在容器内访问目标服务 +除了端口和 method 能力之外，当前 WebSocket proxy 还存在握手上下文丢失问题： +- `sandbox_proxy_service.websocket_proxy()` 在调用 `websockets.connect(...)` 时，目前只传了 `subprotocols` +- 下游服务收到的是 Admin 重新发起的二跳握手，请求来源会表现为 `Python websockets/...`，而不是客户端原始握手上下文 +- 结果是 `Origin`、`Authorization`、`Cookie`、`X-Forwarded-*` 等头全部丢失，依赖这些头的服务会报 `origin not allowed` 或鉴权失败 + --- ## File Changes @@ -14,8 +19,9 @@ admin 与 sandbox 不在同一 K8s 集群，`host_ip` 为宿主机 IP，容器 | 文件 | 修改类型 | 说明 | |------|------|------| | `rock/rocklet/local_api.py` | **新增** | 新增 `ANY /http_proxy/{path:path}?port={port}` 端点 | -| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转 | -| `rock/admin/entrypoints/sandbox_proxy_api.py` | 无变更 | 已支持，无需修改 | +| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转；补充 WebSocket 通用 headers 白名单透传 | +| `rock/admin/entrypoints/sandbox_proxy_api.py` | 无变更 | 路由签名保持不变，无需调整 | +| `tests/unit/sandbox/test_websocket_proxy_subprotocol.py` | 修改 | 增加 `Origin` / `additional_headers` 透传与禁转头测试 | --- @@ -110,6 +116,84 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" ... ``` +### 变更 4：WebSocket proxy 通用 headers 白名单透传 + +需要在 `sandbox_proxy_service.websocket_proxy()` 内新增一层 header 提取和过滤逻辑，将客户端握手里的“通用请求头”转为上游二跳握手参数。 + +**设计要点**： + +1. **将 `Origin` 单独处理** + - `websockets.connect()` 在 15.0.1 版本里提供 `origin=` 参数 + - `Origin` 不作为普通 `additional_headers` 重复透传，避免语义混乱和重复 header + +2. **其余通用头走固定白名单** + - 推荐白名单：`Authorization`、`Cookie`、`X-Forwarded-For`、`X-Forwarded-Host`、`X-Forwarded-Proto`、`X-Real-IP`、`X-Request-Id`、`Traceparent`、`Tracestate`、`EagleEye-TraceId`、`EagleEye-RpcId`、`EagleEye-UserData` + - 以固定集合匹配，避免“默认全透传”带来的握手兼容性和安全风险 + +3. **无白名单头时保持兼容** + - 若客户端未携带任何白名单头，则 `origin=None`、`additional_headers=None` + - 这样代理行为与当前实现保持一致，不引入额外副作用 + +**建议实现草图**： + +```python +FORWARDED_WS_HEADER_NAMES = { + "authorization", + "cookie", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-real-ip", + "x-request-id", + "traceparent", + "tracestate", + "eagleeye-traceid", + "eagleeye-rpcid", + "eagleeye-userdata", +} + + +def build_upstream_ws_headers(client_websocket): + origin = client_websocket.headers.get("origin") + additional_headers = [] + + for key, value in client_websocket.headers.items(): + lower_key = key.lower() + if lower_key == "origin": + continue + if lower_key in FORWARDED_WS_HEADER_NAMES: + additional_headers.append((key, value)) + + return origin, additional_headers or None +``` + +接入方式： + +```python +origin, additional_headers = build_upstream_ws_headers(client_websocket) + +async with websockets.connect( + target_url, + ping_interval=None, + ping_timeout=None, + origin=origin, + additional_headers=additional_headers, + subprotocols=upstream_subprotocols, +) as target_websocket: + ... +``` + +### 变更 5：测试覆盖扩展 + +现有测试主要覆盖子协议转发，需要补充 header 透传相关单测。 + +**新增测试点**： +- `Origin` 存在时，`websockets.connect()` 收到相同 `origin=` +- `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent`、`EagleEye-*` 存在时，`websockets.connect()` 收到 `additional_headers` +- `Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 不得出现在 `additional_headers` +- `Sec-WebSocket-Protocol` 继续通过 `subprotocols=` 转发，不能出现在 `additional_headers` +- 无白名单头时，`origin` / `additional_headers` 为 `None`，保持向后兼容 + --- ## Execution Plan @@ -131,14 +215,25 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" - 有 `port` 时，改用 `http://{host_ip}:{rocklet_mapped_port}/http_proxy/{path}?port={port}` - 无 `port` 时保持原逻辑不变 +### Step 4：新增 WebSocket 通用 header 过滤与透传逻辑 +- 文件：`rock/sandbox/service/sandbox_proxy_service.py` +- 新增 helper，负责从 `client_websocket.headers` 中提取 `Origin` 和白名单 `additional_headers` +- 在 `websocket_proxy()` 调用 `websockets.connect()` 时传入 `origin=` 和 `additional_headers=` +- 保持现有 `subprotocols=` 协商逻辑不变 + +### Step 5：补充 WebSocket header 透传测试 +- 文件：`tests/unit/sandbox/test_websocket_proxy_subprotocol.py` +- 新增 `Origin` 透传、通用 header 透传、禁转 header、兼容性测试 + --- ## Rollback & Compatibility - **向后兼容**：`rock_target_port` 未指定时，所有逻辑路径与原实现完全一致 +- **向后兼容**：客户端未携带任何白名单 header 时，WebSocket 二跳握手行为与现状一致 - **回滚**： - rocklet：还原 `local_api.py`，重新发布镜像 - - admin：还原 `sandbox_proxy_service.py` + - admin：还原 `sandbox_proxy_service.py` 和对应单元测试 --- @@ -147,3 +242,6 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" - WebSocket proxy 自定义端口时，`path` 参数不生效（rocklet portforward 是纯 TCP 隧道，不感知 HTTP path） - rocklet `/http_proxy` 端点的 `port` 参数需要校验（复用 `validate_port_forward_port`） - rocklet 镜像需要重新发布才能生效 +- WebSocket header 透传必须坚持白名单策略，不能简单复制全部请求头 +- `Origin` 应通过 `websockets.connect(origin=...)` 传入；不要与 `additional_headers` 重复 +- `Sec-WebSocket-Protocol` 必须继续通过 `subprotocols=` 传递，避免和普通 header 透传逻辑冲突 diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index b471701a96..8f854508e5 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -46,6 +46,35 @@ logger = init_logger(__name__) +FORWARDED_WS_HEADER_NAMES = { + "authorization", + "cookie", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-real-ip", + "x-request-id", + "traceparent", + "tracestate", + "eagleeye-traceid", + "eagleeye-rpcid", + "eagleeye-userdata", +} + + +def build_upstream_ws_headers(client_websocket): + origin = client_websocket.headers.get("origin") or client_websocket.headers.get("Origin") + additional_headers = [] + + for key, value in client_websocket.headers.items(): + lower_key = key.lower() + if lower_key == "origin": + continue + if lower_key in FORWARDED_WS_HEADER_NAMES: + additional_headers.append((key, value)) + + return origin, additional_headers or None + class SandboxProxyService: _httpx_client = None @@ -208,12 +237,15 @@ async def websocket_proxy( client_subprotocols = getattr(client_websocket, "subprotocols", []) or [] upstream_subprotocols = client_subprotocols if client_subprotocols else ["binary", "base64"] + origin, additional_headers = build_upstream_ws_headers(client_websocket) try: async with websockets.connect( target_url, ping_interval=None, ping_timeout=None, + origin=origin, + additional_headers=additional_headers, subprotocols=upstream_subprotocols, ) as target_websocket: negotiated = getattr(target_websocket, "subprotocol", None) diff --git a/tests/unit/sandbox/test_websocket_proxy_headers.py b/tests/unit/sandbox/test_websocket_proxy_headers.py new file mode 100644 index 0000000000..7cda50e089 --- /dev/null +++ b/tests/unit/sandbox/test_websocket_proxy_headers.py @@ -0,0 +1,148 @@ +"""Tests for WebSocket proxy header forwarding (whitelist-based).""" + +from types import SimpleNamespace + +import websockets + +from rock.sandbox.service.sandbox_proxy_service import ( + SandboxProxyService, + build_upstream_ws_headers, +) + +# ───────────────────────────────────────────────────────────────────────────── +# Unit tests: build_upstream_ws_headers (pure function, no mocks) +# ───────────────────────────────────────────────────────────────────────────── + + +def _ws_with_headers(headers: dict): + """Create a minimal object with a .headers dict, mimicking Starlette WebSocket.""" + return SimpleNamespace(headers=headers) + + +class TestBuildUpstreamWsHeaders: + def test_whitelist_and_origin_forwarded(self): + ws = _ws_with_headers( + { + "Origin": "https://example.com", + "Authorization": "Bearer token123", + "cookie": "session=abc", + "traceparent": "00-trace-span-01", + "EagleEye-TraceId": "eagle-trace-001", + "host": "should-be-excluded", + "sec-websocket-protocol": "binary", + "x-pictor-callid": "pictor-123", + } + ) + origin, additional = build_upstream_ws_headers(ws) + assert origin == "https://example.com" + assert additional is not None + forwarded = dict(additional) + assert forwarded["Authorization"] == "Bearer token123" + assert forwarded["cookie"] == "session=abc" + assert forwarded["traceparent"] == "00-trace-span-01" + assert forwarded["EagleEye-TraceId"] == "eagle-trace-001" + excluded = {"origin", "host", "sec-websocket-protocol", "x-pictor-callid"} + assert excluded.isdisjoint({k.lower() for k, _ in additional}) + + def test_no_whitelist_headers_returns_none(self): + ws = _ws_with_headers({"host": "localhost", "connection": "upgrade", "user-agent": "test"}) + origin, additional = build_upstream_ws_headers(ws) + assert origin is None + assert additional is None + + def test_origin_only_no_additional(self): + ws = _ws_with_headers({"origin": "https://example.com"}) + origin, additional = build_upstream_ws_headers(ws) + assert origin == "https://example.com" + assert additional is None + + +# ───────────────────────────────────────────────────────────────────────────── +# E2E tests: real WebSocket server verifying header arrival +# ───────────────────────────────────────────────────────────────────────────── + + +class FakeClientWebSocket: + """Simulates a Starlette WebSocket for websocket_proxy() input.""" + + def __init__(self, headers: dict, subprotocols: list | None = None): + self.headers = headers + self.subprotocols = subprotocols or [] + self._accepted = False + self._closed = False + + async def accept(self, subprotocol=None): + self._accepted = True + + async def close(self, code=1000, reason=""): + self._closed = True + + async def receive(self): + return {"type": "websocket.disconnect", "code": 1000} + + +class TestWebSocketHeaderForwardingE2E: + """Start a real websockets server and verify headers arrive downstream.""" + + async def _run_proxy_with_server(self, client_headers: dict, subprotocols: list | None = None): + """Helper: start WS server, run websocket_proxy(), return captured request headers.""" + captured_headers = {} + + async def handler(ws): + captured_headers.update(dict(ws.request.headers)) + await ws.close() + + async with websockets.serve(handler, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + target_url = f"ws://127.0.0.1:{port}" + + service = SandboxProxyService.__new__(SandboxProxyService) + + async def noop_update(*a, **kw): + pass + + async def fake_get_url(*a, **kw): + return target_url + + service._update_expire_time = noop_update + service.get_sandbox_websocket_url = fake_get_url + + client_ws = FakeClientWebSocket(client_headers, subprotocols) + await service.websocket_proxy(client_ws, "test-sandbox") + + return captured_headers + + async def test_origin_received_by_downstream(self): + headers = await self._run_proxy_with_server({"origin": "https://my-app.example.com"}) + assert headers.get("origin") == "https://my-app.example.com" + + async def test_whitelist_headers_received_by_downstream(self): + client_headers = { + "authorization": "Bearer secret-token", + "eagleeye-traceid": "eagle-trace-e2e", + "x-request-id": "req-e2e-001", + "traceparent": "00-abcdef-123456-01", + } + headers = await self._run_proxy_with_server(client_headers) + assert headers.get("authorization") == "Bearer secret-token" + assert headers.get("eagleeye-traceid") == "eagle-trace-e2e" + assert headers.get("x-request-id") == "req-e2e-001" + assert headers.get("traceparent") == "00-abcdef-123456-01" + + async def test_forbidden_headers_not_received_by_downstream(self): + client_headers = { + "authorization": "Bearer xxx", + "x-pictor-callid": "should-not-arrive", + "x-rock-sandbox-default-upstream": "should-not-arrive", + "web-server-type": "nginx", + } + headers = await self._run_proxy_with_server(client_headers) + assert headers.get("authorization") == "Bearer xxx" + assert "x-pictor-callid" not in headers + assert "x-rock-sandbox-default-upstream" not in headers + assert "web-server-type" not in headers + + async def test_no_extra_headers_backward_compatible(self): + headers = await self._run_proxy_with_server({}) + assert "authorization" not in headers + assert "origin" not in headers From e7cf45854216ddadc236097e1a65cf85a5446e04 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Tue, 21 Apr 2026 17:50:26 +0800 Subject: [PATCH 061/226] refactor: move FORWARDED_WS_HEADER_NAMES and build_upstream_ws_headers to rock/sandbox/utils/proxy.py Co-Authored-By: Claude Opus 4.6 --- rock/sandbox/service/sandbox_proxy_service.py | 30 +------------------ rock/sandbox/utils/proxy.py | 28 +++++++++++++++++ .../sandbox/test_websocket_proxy_headers.py | 6 ++-- 3 files changed, 31 insertions(+), 33 deletions(-) create mode 100644 rock/sandbox/utils/proxy.py diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 8f854508e5..451cebdf96 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -40,41 +40,13 @@ from rock.common.port_validation import validate_port_forward_port from rock.logger import init_logger from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sandbox.utils.proxy import build_upstream_ws_headers from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError from rock.utils import EAGLE_EYE_TRACE_ID, trace_id_ctx_var logger = init_logger(__name__) -FORWARDED_WS_HEADER_NAMES = { - "authorization", - "cookie", - "x-forwarded-for", - "x-forwarded-host", - "x-forwarded-proto", - "x-real-ip", - "x-request-id", - "traceparent", - "tracestate", - "eagleeye-traceid", - "eagleeye-rpcid", - "eagleeye-userdata", -} - - -def build_upstream_ws_headers(client_websocket): - origin = client_websocket.headers.get("origin") or client_websocket.headers.get("Origin") - additional_headers = [] - - for key, value in client_websocket.headers.items(): - lower_key = key.lower() - if lower_key == "origin": - continue - if lower_key in FORWARDED_WS_HEADER_NAMES: - additional_headers.append((key, value)) - - return origin, additional_headers or None - class SandboxProxyService: _httpx_client = None diff --git a/rock/sandbox/utils/proxy.py b/rock/sandbox/utils/proxy.py new file mode 100644 index 0000000000..f2a3b68101 --- /dev/null +++ b/rock/sandbox/utils/proxy.py @@ -0,0 +1,28 @@ +FORWARDED_WS_HEADER_NAMES = { + "authorization", + "cookie", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-real-ip", + "x-request-id", + "traceparent", + "tracestate", + "eagleeye-traceid", + "eagleeye-rpcid", + "eagleeye-userdata", +} + + +def build_upstream_ws_headers(client_websocket): + origin = client_websocket.headers.get("origin") or client_websocket.headers.get("Origin") + additional_headers = [] + + for key, value in client_websocket.headers.items(): + lower_key = key.lower() + if lower_key == "origin": + continue + if lower_key in FORWARDED_WS_HEADER_NAMES: + additional_headers.append((key, value)) + + return origin, additional_headers or None diff --git a/tests/unit/sandbox/test_websocket_proxy_headers.py b/tests/unit/sandbox/test_websocket_proxy_headers.py index 7cda50e089..613b7d095a 100644 --- a/tests/unit/sandbox/test_websocket_proxy_headers.py +++ b/tests/unit/sandbox/test_websocket_proxy_headers.py @@ -4,10 +4,8 @@ import websockets -from rock.sandbox.service.sandbox_proxy_service import ( - SandboxProxyService, - build_upstream_ws_headers, -) +from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService +from rock.sandbox.utils.proxy import build_upstream_ws_headers # ───────────────────────────────────────────────────────────────────────────── # Unit tests: build_upstream_ws_headers (pure function, no mocks) From e5301c594beac1ffec761f77fea2c4d42f62636f Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Wed, 22 Apr 2026 11:33:12 +0800 Subject: [PATCH 062/226] refactor(proxy): switch WebSocket header forwarding from whitelist to blacklist Whitelist strategy blocked user-defined custom headers from reaching downstream services. Blacklist strategy forwards all headers by default, only filtering out WebSocket handshake headers and hop-by-hop headers. Co-Authored-By: Claude Opus 4.6 --- .../proxy-enhancements/01_requirement.md | 17 ++-- .../proxy-enhancements/03_implementation.md | 72 +++++++------- rock/sandbox/utils/proxy.py | 33 ++++--- .../sandbox/test_websocket_proxy_headers.py | 95 +++++++++++++++---- 4 files changed, 141 insertions(+), 76 deletions(-) diff --git a/docs/_specs/proxy-enhancements/01_requirement.md b/docs/_specs/proxy-enhancements/01_requirement.md index dafa3d84c0..85653b629a 100644 --- a/docs/_specs/proxy-enhancements/01_requirement.md +++ b/docs/_specs/proxy-enhancements/01_requirement.md @@ -39,12 +39,12 @@ ROCK Admin 目前提供两类代理能力： - 在现有 `/sandboxes/{sandbox_id}/proxy[/{path}]` 路由上，允许通过 query param `port` 指定目标 HTTP 端口 - 当 `port` 未指定时，保持现有行为（使用 `Port.SERVER = 8080`） -4. **WebSocket Proxy 支持白名单透传通用请求头** - - 通过 `/sandboxes/{id}/proxy/{path:path}` 建立 WebSocket 代理时，允许将客户端请求中的通用 header 按白名单转发到下游服务 - - 首批需要覆盖的 header 包括：`Origin`、`Authorization`、`Cookie`、`X-Forwarded-For`、`X-Forwarded-Host`、`X-Forwarded-Proto`、`X-Real-IP`、`X-Request-Id`、`Traceparent`、`Tracestate`、`EagleEye-TraceId`、`EagleEye-RpcId`、`EagleEye-UserData` +4. **WebSocket Proxy 支持黑名单过滤透传通用请求头** + - 通过 `/sandboxes/{id}/proxy/{path:path}` 建立 WebSocket 代理时，默认将客户端请求中的通用 header 转发到下游服务，通过黑名单排除不应转发的头 + - 黑名单排除的头包括：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`） - `Origin` 是必须支持的关键头，因为下游服务可能使用来源白名单（例如 `gateway.controlUi.allowedOrigins`）校验 WebSocket 握手 - `Sec-WebSocket-Protocol` 继续通过 WebSocket 子协议协商传递，不作为普通 header 透传 - - 采用白名单策略，避免将握手专用头和 hop-by-hop 头错误转发给下游 + - 采用黑名单策略，确保用户自定义 header 能被透传到下游服务 ### Out（本次不做的） @@ -52,7 +52,6 @@ ROCK Admin 目前提供两类代理能力： - HTTP Proxy 的 multipart/form-data 支持（upload 接口已单独处理） - `host_proxy` 的 method 扩展（范围外） - SDK 客户端侧的封装更新 -- WebSocket 代理透传“所有”请求头，或由客户端动态指定任意 header 白名单 - 自动伪造、补默认值或重写任意 `Origin` - 透传 WebSocket 握手专用头（如 `Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`） @@ -70,9 +69,9 @@ ROCK Admin 目前提供两类代理能力： - **AC8**：`GET /sandboxes/{sandbox_id}/proxy/api/health?port=9000` 能成功代理到沙箱内 9000 端口的 HTTP 服务 - **AC9**：HTTP proxy 不带 `port` 参数时行为不变（向后兼容） - **AC10**：当客户端 WebSocket 握手包含 `Origin` 时，代理发起到下游的二跳握手必须携带相同 `Origin`，以满足下游来源校验 -- **AC11**：当客户端握手包含 `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent` 等白名单头时，代理发起到下游的二跳握手必须一并转发 +- **AC11**：当客户端握手包含 `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent` 或任意自定义头时，代理发起到下游的二跳握手必须一并转发 - **AC12**：`Sec-WebSocket-Protocol` 必须继续通过现有 `subprotocols` 机制转发和协商，不能降级为普通 header 透传 -- **AC13**：`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 等握手专用头不得被转发到下游 +- **AC13**：黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`）不得被转发到下游 - **AC14**：当客户端未携带任何白名单头时，WebSocket proxy 的默认行为与现状保持一致（向后兼容） --- @@ -83,7 +82,7 @@ ROCK Admin 目前提供两类代理能力： - 不修改 `Port.SERVER` / `Port.PROXY` 枚举定义 - `port_validation` 逻辑复用现有 `validate_port_forward_port`（但 WebSocket proxy 端口校验需新增对 `Port.SERVER = 8080` 的允许，或直接用同一校验函数） - 保持现有 portforward 端点（`/sandboxes/{id}/portforward`）不变 -- WebSocket header 转发采用**白名单**而不是**黑名单**策略，避免误传握手专用头 +- WebSocket header 转发采用**黑名单**策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header 透传 - `Origin` 通过 WebSocket 客户端库的显式参数透传，不与普通 `additional_headers` 混用 - `Sec-WebSocket-Protocol` 继续通过 `subprotocols` 参数处理，不走通用 header 透传逻辑 @@ -92,7 +91,7 @@ ROCK Admin 目前提供两类代理能力： ## Risks & Rollout - **风险**：WebSocket proxy 中用户可以指定任意端口访问沙箱内服务，存在横向访问风险 → 通过 `sandbox_id` 鉴权已覆盖，端口范围校验作为防护层 -- **风险**：若白名单范围定义过宽，可能把不该下传的敏感或握手专用头带给下游 → 通过固定白名单和单元测试约束范围 +- **风险**：黑名单策略下，未列入黑名单的头会默认转发 → 通过单元测试确保握手专用头和 hop-by-hop 头被正确过滤 - **风险**：不同下游服务对 `Origin`、`Cookie`、`Authorization` 的要求不同，透传后暴露出原本被代理层掩盖的问题 → 以“尽可能保留客户端上下文、但不伪造默认值”为原则 - **回滚**：修改集中在 `sandbox_proxy_service.py` 和对应单元测试；如需回滚，可仅还原 WebSocket header 透传逻辑 - **上线策略**：无数据库变更，直接部署；建议先在依赖 `Origin` 校验的控制台类服务上验证 diff --git a/docs/_specs/proxy-enhancements/03_implementation.md b/docs/_specs/proxy-enhancements/03_implementation.md index f7254a6e11..20dd96d05e 100644 --- a/docs/_specs/proxy-enhancements/03_implementation.md +++ b/docs/_specs/proxy-enhancements/03_implementation.md @@ -116,9 +116,9 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" ... ``` -### 变更 4：WebSocket proxy 通用 headers 白名单透传 +### 变更 4：WebSocket proxy 通用 headers 黑名单过滤透传 -需要在 `sandbox_proxy_service.websocket_proxy()` 内新增一层 header 提取和过滤逻辑，将客户端握手里的“通用请求头”转为上游二跳握手参数。 +需要在 `sandbox_proxy_service.websocket_proxy()` 内新增一层 header 提取和过滤逻辑，将客户端握手里的”通用请求头”转为上游二跳握手参数。 **设计要点**： @@ -126,43 +126,46 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" - `websockets.connect()` 在 15.0.1 版本里提供 `origin=` 参数 - `Origin` 不作为普通 `additional_headers` 重复透传，避免语义混乱和重复 header -2. **其余通用头走固定白名单** - - 推荐白名单：`Authorization`、`Cookie`、`X-Forwarded-For`、`X-Forwarded-Host`、`X-Forwarded-Proto`、`X-Real-IP`、`X-Request-Id`、`Traceparent`、`Tracestate`、`EagleEye-TraceId`、`EagleEye-RpcId`、`EagleEye-UserData` - - 以固定集合匹配，避免“默认全透传”带来的握手兼容性和安全风险 +2. **其余头走黑名单过滤** + - 黑名单：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`） + - 不在黑名单中的头默认转发，确保用户自定义 header 能到达下游 -3. **无白名单头时保持兼容** - - 若客户端未携带任何白名单头，则 `origin=None`、`additional_headers=None` +3. **无可转发头时保持兼容** + - 若客户端未携带任何非黑名单头，则 `origin=None`、`additional_headers=None` - 这样代理行为与当前实现保持一致，不引入额外副作用 -**建议实现草图**： +**实现代码**（`rock/sandbox/utils/proxy.py`）： ```python -FORWARDED_WS_HEADER_NAMES = { - "authorization", - "cookie", - "x-forwarded-for", - "x-forwarded-host", - "x-forwarded-proto", - "x-real-ip", - "x-request-id", - "traceparent", - "tracestate", - "eagleeye-traceid", - "eagleeye-rpcid", - "eagleeye-userdata", +BLOCKED_WS_HEADER_NAMES = { + “host”, + “connection”, + “upgrade”, + “sec-websocket-key”, + “sec-websocket-version”, + “sec-websocket-extensions”, + “sec-websocket-protocol”, + “transfer-encoding”, + “te”, + “trailer”, + “keep-alive”, + “proxy-authorization”, + “proxy-connection”, + “content-length”, } def build_upstream_ws_headers(client_websocket): - origin = client_websocket.headers.get("origin") + origin = client_websocket.headers.get(“origin”) or client_websocket.headers.get(“Origin”) additional_headers = [] for key, value in client_websocket.headers.items(): lower_key = key.lower() - if lower_key == "origin": + if lower_key == “origin”: continue - if lower_key in FORWARDED_WS_HEADER_NAMES: - additional_headers.append((key, value)) + if lower_key in BLOCKED_WS_HEADER_NAMES: + continue + additional_headers.append((key, value)) return origin, additional_headers or None ``` @@ -189,10 +192,11 @@ async with websockets.connect( **新增测试点**： - `Origin` 存在时，`websockets.connect()` 收到相同 `origin=` -- `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent`、`EagleEye-*` 存在时，`websockets.connect()` 收到 `additional_headers` -- `Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 不得出现在 `additional_headers` +- 已知头（`Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent`、`EagleEye-*`）存在时，`websockets.connect()` 收到 `additional_headers` +- 用户自定义头（如 `x-my-custom`）能被正常转发到下游 +- 黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 等）不得出现在 `additional_headers` - `Sec-WebSocket-Protocol` 继续通过 `subprotocols=` 转发，不能出现在 `additional_headers` -- 无白名单头时，`origin` / `additional_headers` 为 `None`，保持向后兼容 +- 无可转发头时，`origin` / `additional_headers` 为 `None`，保持向后兼容 --- @@ -215,15 +219,15 @@ async with websockets.connect( - 有 `port` 时，改用 `http://{host_ip}:{rocklet_mapped_port}/http_proxy/{path}?port={port}` - 无 `port` 时保持原逻辑不变 -### Step 4：新增 WebSocket 通用 header 过滤与透传逻辑 -- 文件：`rock/sandbox/service/sandbox_proxy_service.py` -- 新增 helper，负责从 `client_websocket.headers` 中提取 `Origin` 和白名单 `additional_headers` +### Step 4：新增 WebSocket 通用 header 黑名单过滤与透传逻辑 +- 文件：`rock/sandbox/utils/proxy.py`（独立模块）、`rock/sandbox/service/sandbox_proxy_service.py`（调用方） +- 新增 `build_upstream_ws_headers()` helper，负责从 `client_websocket.headers` 中提取 `Origin` 并通过黑名单过滤 `additional_headers` - 在 `websocket_proxy()` 调用 `websockets.connect()` 时传入 `origin=` 和 `additional_headers=` - 保持现有 `subprotocols=` 协商逻辑不变 ### Step 5：补充 WebSocket header 透传测试 -- 文件：`tests/unit/sandbox/test_websocket_proxy_subprotocol.py` -- 新增 `Origin` 透传、通用 header 透传、禁转 header、兼容性测试 +- 文件：`tests/unit/sandbox/test_websocket_proxy_headers.py` +- 新增 `Origin` 透传、已知 header 透传、自定义 header 透传、黑名单 header 过滤、兼容性测试 --- @@ -242,6 +246,6 @@ async with websockets.connect( - WebSocket proxy 自定义端口时，`path` 参数不生效（rocklet portforward 是纯 TCP 隧道，不感知 HTTP path） - rocklet `/http_proxy` 端点的 `port` 参数需要校验（复用 `validate_port_forward_port`） - rocklet 镜像需要重新发布才能生效 -- WebSocket header 透传必须坚持白名单策略，不能简单复制全部请求头 +- WebSocket header 透传采用黑名单策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header 透传 - `Origin` 应通过 `websockets.connect(origin=...)` 传入；不要与 `additional_headers` 重复 - `Sec-WebSocket-Protocol` 必须继续通过 `subprotocols=` 传递，避免和普通 header 透传逻辑冲突 diff --git a/rock/sandbox/utils/proxy.py b/rock/sandbox/utils/proxy.py index f2a3b68101..195cb0f83b 100644 --- a/rock/sandbox/utils/proxy.py +++ b/rock/sandbox/utils/proxy.py @@ -1,16 +1,18 @@ -FORWARDED_WS_HEADER_NAMES = { - "authorization", - "cookie", - "x-forwarded-for", - "x-forwarded-host", - "x-forwarded-proto", - "x-real-ip", - "x-request-id", - "traceparent", - "tracestate", - "eagleeye-traceid", - "eagleeye-rpcid", - "eagleeye-userdata", +BLOCKED_WS_HEADER_NAMES = { + "host", + "connection", + "upgrade", + "sec-websocket-key", + "sec-websocket-version", + "sec-websocket-extensions", + "sec-websocket-protocol", + "transfer-encoding", + "te", + "trailer", + "keep-alive", + "proxy-authorization", + "proxy-connection", + "content-length", } @@ -22,7 +24,8 @@ def build_upstream_ws_headers(client_websocket): lower_key = key.lower() if lower_key == "origin": continue - if lower_key in FORWARDED_WS_HEADER_NAMES: - additional_headers.append((key, value)) + if lower_key in BLOCKED_WS_HEADER_NAMES: + continue + additional_headers.append((key, value)) return origin, additional_headers or None diff --git a/tests/unit/sandbox/test_websocket_proxy_headers.py b/tests/unit/sandbox/test_websocket_proxy_headers.py index 613b7d095a..caf8841240 100644 --- a/tests/unit/sandbox/test_websocket_proxy_headers.py +++ b/tests/unit/sandbox/test_websocket_proxy_headers.py @@ -1,11 +1,11 @@ -"""Tests for WebSocket proxy header forwarding (whitelist-based).""" +"""Tests for WebSocket proxy header forwarding (blacklist-based).""" from types import SimpleNamespace import websockets from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService -from rock.sandbox.utils.proxy import build_upstream_ws_headers +from rock.sandbox.utils.proxy import BLOCKED_WS_HEADER_NAMES, build_upstream_ws_headers # ───────────────────────────────────────────────────────────────────────────── # Unit tests: build_upstream_ws_headers (pure function, no mocks) @@ -18,7 +18,7 @@ def _ws_with_headers(headers: dict): class TestBuildUpstreamWsHeaders: - def test_whitelist_and_origin_forwarded(self): + def test_known_headers_forwarded(self): ws = _ws_with_headers( { "Origin": "https://example.com", @@ -26,9 +26,6 @@ def test_whitelist_and_origin_forwarded(self): "cookie": "session=abc", "traceparent": "00-trace-span-01", "EagleEye-TraceId": "eagle-trace-001", - "host": "should-be-excluded", - "sec-websocket-protocol": "binary", - "x-pictor-callid": "pictor-123", } ) origin, additional = build_upstream_ws_headers(ws) @@ -39,11 +36,46 @@ def test_whitelist_and_origin_forwarded(self): assert forwarded["cookie"] == "session=abc" assert forwarded["traceparent"] == "00-trace-span-01" assert forwarded["EagleEye-TraceId"] == "eagle-trace-001" - excluded = {"origin", "host", "sec-websocket-protocol", "x-pictor-callid"} - assert excluded.isdisjoint({k.lower() for k, _ in additional}) - def test_no_whitelist_headers_returns_none(self): - ws = _ws_with_headers({"host": "localhost", "connection": "upgrade", "user-agent": "test"}) + def test_custom_headers_forwarded(self): + ws = _ws_with_headers( + { + "x-my-custom": "custom-value", + "x-pictor-callid": "pictor-123", + "web-server-type": "nginx", + } + ) + origin, additional = build_upstream_ws_headers(ws) + assert origin is None + assert additional is not None + forwarded = dict(additional) + assert forwarded["x-my-custom"] == "custom-value" + assert forwarded["x-pictor-callid"] == "pictor-123" + assert forwarded["web-server-type"] == "nginx" + + def test_blocked_headers_excluded(self): + ws = _ws_with_headers( + { + "Authorization": "Bearer token123", + "host": "should-be-excluded", + "connection": "upgrade", + "upgrade": "websocket", + "sec-websocket-key": "dGhlIHNhbXBsZSBub25jZQ==", + "sec-websocket-version": "13", + "sec-websocket-extensions": "permessage-deflate", + "sec-websocket-protocol": "binary", + "transfer-encoding": "chunked", + "keep-alive": "timeout=5", + } + ) + origin, additional = build_upstream_ws_headers(ws) + assert origin is None + assert additional is not None + forwarded_keys = {k.lower() for k, _ in additional} + assert forwarded_keys == {"authorization"} + + def test_no_forwardable_headers_returns_none(self): + ws = _ws_with_headers({"host": "localhost", "connection": "upgrade"}) origin, additional = build_upstream_ws_headers(ws) assert origin is None assert additional is None @@ -54,6 +86,25 @@ def test_origin_only_no_additional(self): assert origin == "https://example.com" assert additional is None + def test_origin_not_in_additional(self): + ws = _ws_with_headers( + { + "origin": "https://example.com", + "Authorization": "Bearer xxx", + } + ) + origin, additional = build_upstream_ws_headers(ws) + assert origin == "https://example.com" + forwarded_keys = {k.lower() for k, _ in additional} + assert "origin" not in forwarded_keys + + def test_all_blocked_headers_covered(self): + headers = {name: "value" for name in BLOCKED_WS_HEADER_NAMES} + ws = _ws_with_headers(headers) + origin, additional = build_upstream_ws_headers(ws) + assert origin is None + assert additional is None + # ───────────────────────────────────────────────────────────────────────────── # E2E tests: real WebSocket server verifying header arrival @@ -114,7 +165,7 @@ async def test_origin_received_by_downstream(self): headers = await self._run_proxy_with_server({"origin": "https://my-app.example.com"}) assert headers.get("origin") == "https://my-app.example.com" - async def test_whitelist_headers_received_by_downstream(self): + async def test_known_headers_received_by_downstream(self): client_headers = { "authorization": "Bearer secret-token", "eagleeye-traceid": "eagle-trace-e2e", @@ -127,18 +178,26 @@ async def test_whitelist_headers_received_by_downstream(self): assert headers.get("x-request-id") == "req-e2e-001" assert headers.get("traceparent") == "00-abcdef-123456-01" - async def test_forbidden_headers_not_received_by_downstream(self): + async def test_custom_headers_received_by_downstream(self): client_headers = { - "authorization": "Bearer xxx", - "x-pictor-callid": "should-not-arrive", - "x-rock-sandbox-default-upstream": "should-not-arrive", + "x-my-custom-app": "my-value", + "x-pictor-callid": "pictor-123", "web-server-type": "nginx", } headers = await self._run_proxy_with_server(client_headers) + assert headers.get("x-my-custom-app") == "my-value" + assert headers.get("x-pictor-callid") == "pictor-123" + assert headers.get("web-server-type") == "nginx" + + async def test_blocked_headers_not_duplicated_by_proxy(self): + client_headers = { + "authorization": "Bearer xxx", + "host": "evil.example.com", + "connection": "upgrade", + } + headers = await self._run_proxy_with_server(client_headers) assert headers.get("authorization") == "Bearer xxx" - assert "x-pictor-callid" not in headers - assert "x-rock-sandbox-default-upstream" not in headers - assert "web-server-type" not in headers + assert headers.get("host") != "evil.example.com" async def test_no_extra_headers_backward_compatible(self): headers = await self._run_proxy_with_server({}) From 3b55efe53b7fdeaf64cdd79a81478d7f0d2a7d22 Mon Sep 17 00:00:00 2001 From: berstpander Date: Wed, 22 Apr 2026 15:47:54 +0800 Subject: [PATCH 063/226] feat(datasets): add tasks subcommand with file task support and improved output (#875) * feat(datasets): add tasks subcommand for listing dataset tasks - add `rock datasets tasks` with required org/dataset and default split=test - support offset/limit pagination for displayed task IDs - extend DatasetClient and registry layers with list_dataset_tasks - add unit tests for CLI, client, and OSS registry behavior Co-Authored-By: Oz * feat(datasets): improve tasks output formatting with separator and header Add visual separation between log messages and task list output: - Empty line and separator line before results - Consolidated info line with dataset/split/total/shown - "#Task name" header with underline for task list Co-Authored-By: Claude Opus 4.7 * fix(datasets): include file tasks in dataset task listing Previously, task listing only recognized directory tasks from prefix_list. This fix adds support for file tasks from object_list: - Add _extract_tasks_from_split method to merge directory and file tasks - Strip file suffix (e.g., task-001.json -> task-001) - Ignore placeholder objects (key ending with "/") and nested paths - Dedupe and sort merged task list This affects both list_datasets (task count) and list_dataset_tasks (task list). Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Oz Co-authored-by: Claude Opus 4.7 --- rock/cli/command/datasets.py | 52 ++++++++ rock/sdk/envhub/datasets/client.py | 3 + rock/sdk/envhub/datasets/registry/base.py | 5 + rock/sdk/envhub/datasets/registry/oss.py | 51 +++++++- tests/unit/datasets/test_client.py | 11 ++ tests/unit/datasets/test_datasets_command.py | 130 ++++++++++++++++++- tests/unit/datasets/test_oss_registry.py | 113 ++++++++++++++++ 7 files changed, 362 insertions(+), 3 deletions(-) diff --git a/rock/cli/command/datasets.py b/rock/cli/command/datasets.py index becf194f59..b53636cb79 100644 --- a/rock/cli/command/datasets.py +++ b/rock/cli/command/datasets.py @@ -13,12 +13,28 @@ logger = init_logger(__name__) +def _non_negative_int(value: str) -> int: + ivalue = int(value) + if ivalue < 0: + raise argparse.ArgumentTypeError("must be >= 0") + return ivalue + + +def _positive_int(value: str) -> int: + ivalue = int(value) + if ivalue <= 0: + raise argparse.ArgumentTypeError("must be >= 1") + return ivalue + + class DatasetsCommand(Command): name = "datasets" async def arun(self, args: argparse.Namespace) -> None: if args.datasets_command == "list": await self._list(args) + elif args.datasets_command == "tasks": + await self._tasks(args) elif args.datasets_command == "upload": await self._upload(args) else: @@ -58,6 +74,35 @@ async def _list(self, args: argparse.Namespace) -> None: for ds in sorted(datasets, key=lambda d: (d.id, d.split)): print(f"{ds.id:<{col_id}} {ds.split:<{col_split}} {len(ds.task_ids):>6}") + async def _tasks(self, args: argparse.Namespace) -> None: + registry_info = self._build_oss_registry_info(args) + client = DatasetClient(registry_info) + spec = client.list_dataset_tasks(args.org, args.dataset, args.split) + + if spec is None or not spec.task_ids: + print(f"No tasks found for dataset '{args.org}/{args.dataset}' split '{args.split}'.") + return + + total = len(spec.task_ids) + start = args.offset + end = start + args.limit if args.limit is not None else None + shown_task_ids = spec.task_ids[start:end] + + if not shown_task_ids: + print("No tasks found after applying offset/limit.") + return + + limit_text = str(args.limit) if args.limit is not None else "all" + + print() + print("=" * 80) + print(f"Dataset: {spec.id} Split: {spec.split} Total: {total} Shown: {len(shown_task_ids)}") + print("=" * 80) + print("#Task name") + print("-" * 10) + for task_id in shown_task_ids: + print(task_id) + async def _upload(self, args: argparse.Namespace) -> None: local_dir = Path(args.dir) if not local_dir.is_dir(): @@ -99,6 +144,13 @@ def add_oss_args(parser: argparse.ArgumentParser) -> None: list_parser = datasets_subparsers.add_parser("list", help="List datasets in OSS registry") list_parser.add_argument("--org", help="Filter by organization") add_oss_args(list_parser) + tasks_parser = datasets_subparsers.add_parser("tasks", help="List task IDs under one dataset split") + tasks_parser.add_argument("--org", required=True, help="Organization name") + tasks_parser.add_argument("--dataset", required=True, help="Dataset name") + tasks_parser.add_argument("--split", default="test", help="Split name (default: test)") + tasks_parser.add_argument("--offset", type=_non_negative_int, default=0, help="Skip first N tasks") + tasks_parser.add_argument("--limit", type=_positive_int, default=None, help="Maximum number of tasks to show") + add_oss_args(tasks_parser) upload_parser = datasets_subparsers.add_parser("upload", help="Upload local task dirs to OSS") upload_parser.add_argument("--org", required=True, help="Organization name") diff --git a/rock/sdk/envhub/datasets/client.py b/rock/sdk/envhub/datasets/client.py index 02f1336d9b..2ea4d70af0 100644 --- a/rock/sdk/envhub/datasets/client.py +++ b/rock/sdk/envhub/datasets/client.py @@ -11,6 +11,9 @@ def __init__(self, registry: OssRegistryInfo) -> None: def list_datasets(self, org: str | None = None) -> list[DatasetSpec]: return self._registry.list_datasets(org) + def list_dataset_tasks(self, organization: str, dataset: str, split: str = "test") -> DatasetSpec | None: + return self._registry.list_dataset_tasks(organization, dataset, split) + def upload_dataset( self, source: LocalDatasetConfig, diff --git a/rock/sdk/envhub/datasets/registry/base.py b/rock/sdk/envhub/datasets/registry/base.py index 8a3538aa6d..5c9192d2d6 100644 --- a/rock/sdk/envhub/datasets/registry/base.py +++ b/rock/sdk/envhub/datasets/registry/base.py @@ -11,6 +11,11 @@ def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: """List all datasets. Filtered to `organization` if provided.""" ... + @abstractmethod + def list_dataset_tasks(self, organization: str, dataset: str, split: str = "test") -> DatasetSpec | None: + """List task ids for one dataset split. Returns None if dataset/split has no tasks.""" + ... + @abstractmethod def upload_dataset( self, diff --git a/rock/sdk/envhub/datasets/registry/oss.py b/rock/sdk/envhub/datasets/registry/oss.py index db455465c4..4e90613571 100644 --- a/rock/sdk/envhub/datasets/registry/oss.py +++ b/rock/sdk/envhub/datasets/registry/oss.py @@ -36,6 +36,40 @@ def _build_prefix(self, org: str, name: str, split: str | None = None) -> str: def _last_segment(prefix: str) -> str: return prefix.rstrip("/").rsplit("/", 1)[-1] + def _extract_tasks_from_split(self, bucket: oss2.Bucket, split_prefix: str) -> list[str]: + """Extract tasks from a split prefix, combining directory and file tasks. + + Directory tasks: from prefix_list (e.g., "datasets/org/name/split/task-001/") + File tasks: from object_list (e.g., "datasets/org/name/split/task-001.json") + + File tasks are stripped of their suffix (e.g., "task-001.json" -> "task-001"). + Placeholder objects (key ending with "/") and nested objects are ignored. + """ + result = bucket.list_objects_v2(prefix=split_prefix, delimiter="/", max_keys=1000) + + # Directory tasks from prefix_list + dir_tasks = [self._last_segment(p) for p in result.prefix_list] + + # File tasks from object_list: direct files under split, strip suffix + file_tasks = [] + for obj in result.object_list: + key = obj.key + # Ignore directory placeholder objects (key ending with "/") + if key.endswith("/"): + continue + # Get the relative path from split_prefix + relative = key[len(split_prefix):] + # Only direct files (no nested paths with "/") + if "/" in relative: + continue + # Strip suffix (e.g., "task-001.json" -> "task-001") + name = relative.rsplit(".", 1)[0] if "." in relative else relative + file_tasks.append(name) + + # Merge and dedupe with stable sort + all_tasks = sorted(set(dir_tasks + file_tasks)) + return all_tasks + def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: bucket = self._build_bucket() base = self._registry.oss_dataset_path or "datasets" @@ -58,8 +92,7 @@ def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: for split_prefix in result2.prefix_list: split = self._last_segment(split_prefix) - result3 = bucket.list_objects_v2(prefix=split_prefix, delimiter="/", max_keys=1000) - task_ids = [self._last_segment(p) for p in result3.prefix_list] + task_ids = self._extract_tasks_from_split(bucket, split_prefix) datasets.append(DatasetSpec( id=f"{org}/{name}", split=split, @@ -68,6 +101,20 @@ def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: return datasets + def list_dataset_tasks(self, organization: str, dataset: str, split: str = "test") -> DatasetSpec | None: + bucket = self._build_bucket() + split_prefix = f"{self._build_prefix(organization, dataset, split)}/" + task_ids = self._extract_tasks_from_split(bucket, split_prefix) + + if not task_ids: + return None + + return DatasetSpec( + id=f"{organization}/{dataset}", + split=split, + task_ids=task_ids, + ) + def _task_exists(self, bucket: oss2.Bucket, task_prefix: str) -> bool: result = bucket.list_objects_v2(prefix=task_prefix, max_keys=1) return len(result.object_list) > 0 diff --git a/tests/unit/datasets/test_client.py b/tests/unit/datasets/test_client.py index 00a5448e32..baf8a8f753 100644 --- a/tests/unit/datasets/test_client.py +++ b/tests/unit/datasets/test_client.py @@ -31,3 +31,14 @@ def test_dataset_client_upload_delegates_to_registry(tmp_path): mock_up.assert_called_once_with(source, target, 2) assert result == expected + + +def test_dataset_client_list_tasks_delegates_to_registry_with_default_split(): + client = DatasetClient(make_registry_info()) + expected = DatasetSpec(id="qwen/bench", split="test", task_ids=["task-001"]) + + with patch.object(client._registry, "list_dataset_tasks", return_value=expected) as mock_list_tasks: + result = client.list_dataset_tasks("qwen", "bench") + + mock_list_tasks.assert_called_once_with("qwen", "bench", "test") + assert result == expected diff --git a/tests/unit/datasets/test_datasets_command.py b/tests/unit/datasets/test_datasets_command.py index 9f40fe4a3f..7c0b0d068e 100644 --- a/tests/unit/datasets/test_datasets_command.py +++ b/tests/unit/datasets/test_datasets_command.py @@ -1,9 +1,12 @@ import argparse -from unittest.mock import patch +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest from rock.cli.command.datasets import DatasetsCommand +from rock.sdk.bench.models.job.config import OssRegistryInfo +from rock.sdk.envhub.datasets.models import DatasetSpec def make_base_args(**kwargs): @@ -16,11 +19,25 @@ def make_base_args(**kwargs): access_key_secret=None, region=None, org=None, + dataset=None, + split=None, + offset=0, + limit=None, ) for k, v in kwargs.items(): setattr(args, k, v) return args +def make_registry_info(): + return OssRegistryInfo(oss_bucket="b", oss_access_key_id="k", oss_access_key_secret="s") + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="rock") + subparsers = parser.add_subparsers(dest="command") + asyncio.run(DatasetsCommand.add_parser_to(subparsers)) + return parser + def test_command_name(): assert DatasetsCommand.name == "datasets" @@ -76,3 +93,114 @@ def test_build_oss_registry_info_raises_when_bucket_missing(): with pytest.raises(ValueError, match="bucket"): cmd._build_oss_registry_info(args) + + +def test_tasks_parser_defaults_split_offset_limit(): + parser = _build_parser() + ns = parser.parse_args(["datasets", "tasks", "--org", "qwen", "--dataset", "my-bench"]) + + assert ns.command == "datasets" + assert ns.datasets_command == "tasks" + assert ns.org == "qwen" + assert ns.dataset == "my-bench" + assert ns.split == "test" + assert ns.offset == 0 + assert ns.limit is None + + +@pytest.mark.parametrize( + "argv", + [ + ["datasets", "tasks", "--dataset", "my-bench"], + ["datasets", "tasks", "--org", "qwen"], + ], +) +def test_tasks_parser_requires_org_and_dataset(argv): + parser = _build_parser() + + with pytest.raises(SystemExit) as excinfo: + parser.parse_args(argv) + + assert excinfo.value.code == 2 + + +def test_tasks_parser_rejects_negative_offset(): + parser = _build_parser() + + with pytest.raises(SystemExit) as excinfo: + parser.parse_args(["datasets", "tasks", "--org", "qwen", "--dataset", "my-bench", "--offset", "-1"]) + + assert excinfo.value.code == 2 + + +def test_tasks_parser_rejects_non_positive_limit(): + parser = _build_parser() + + with pytest.raises(SystemExit) as excinfo: + parser.parse_args(["datasets", "tasks", "--org", "qwen", "--dataset", "my-bench", "--limit", "0"]) + + assert excinfo.value.code == 2 + + +def test_arun_dispatches_tasks(): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="tasks", org="qwen", dataset="my-bench", split="test") + + with patch.object(DatasetsCommand, "_tasks", new_callable=AsyncMock, create=True) as mock_tasks: + asyncio.run(cmd.arun(args)) + + mock_tasks.assert_awaited_once_with(args) + + +def test_tasks_outputs_paginated_results(capsys): + cmd = DatasetsCommand() + args = make_base_args( + datasets_command="tasks", + org="qwen", + dataset="my-bench", + split="test", + offset=1, + limit=2, + ) + mock_client = MagicMock() + mock_client.list_dataset_tasks.return_value = DatasetSpec( + id="qwen/my-bench", + split="test", + task_ids=["task-001", "task-002", "task-003"], + ) + + with patch.object(cmd, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient", return_value=mock_client): + asyncio.run(cmd._tasks(args)) + + mock_client.list_dataset_tasks.assert_called_once_with("qwen", "my-bench", "test") + out = capsys.readouterr().out + assert "Dataset: qwen/my-bench" in out + assert "Split: test" in out + assert "task-002" in out + assert "task-003" in out + assert "task-001" not in out + assert "Total: 3" in out + assert "Shown: 2" in out + assert "#Task name" in out + + +def test_tasks_prints_no_tasks_message_when_not_found(capsys): + cmd = DatasetsCommand() + args = make_base_args( + datasets_command="tasks", + org="qwen", + dataset="my-bench", + split="test", + offset=0, + limit=None, + ) + mock_client = MagicMock() + mock_client.list_dataset_tasks.return_value = None + + with patch.object(cmd, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient", return_value=mock_client): + asyncio.run(cmd._tasks(args)) + + out = capsys.readouterr().out + assert "No tasks found" in out diff --git a/tests/unit/datasets/test_oss_registry.py b/tests/unit/datasets/test_oss_registry.py index 4ba2f0db52..9237ac1d6f 100644 --- a/tests/unit/datasets/test_oss_registry.py +++ b/tests/unit/datasets/test_oss_registry.py @@ -58,6 +58,31 @@ def test_list_datasets_filter_by_org(): assert first_call_kwargs["prefix"] == "datasets/qwen/" assert len(datasets) == 1 +def test_list_datasets_counts_directory_and_file_tasks(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.side_effect = [ + make_list_result(prefixes=["datasets/qwen/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/"]), + make_list_result(prefixes=["datasets/qwen/my-bench/train/"]), + make_list_result( + prefixes=["datasets/qwen/my-bench/train/task-dir/"], + objects=[ + MagicMock(key="datasets/qwen/my-bench/train/task-file.json"), + MagicMock(key="datasets/qwen/my-bench/train/"), + MagicMock(key="datasets/qwen/my-bench/train/nested/task-ignored.json"), + ], + ), + ] + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + datasets = registry.list_datasets() + + assert len(datasets) == 1 + assert datasets[0].id == "qwen/my-bench" + assert datasets[0].split == "train" + assert datasets[0].task_ids == ["task-dir", "task-file"] + def test_list_datasets_empty_registry(): registry = OssDatasetRegistry(make_registry_info()) @@ -79,6 +104,94 @@ def test_build_prefix_with_split(): registry = OssDatasetRegistry(make_registry_info()) assert registry._build_prefix("qwen", "my-bench", "train") == "datasets/qwen/my-bench/train" +# --------------------------------------------------------------------------- +# list_dataset_tasks tests +# --------------------------------------------------------------------------- + + +def test_list_dataset_tasks_uses_default_test_split_and_sorts_task_ids(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[ + "datasets/qwen/my-bench/test/task-002/", + "datasets/qwen/my-bench/test/task-001/", + ]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + spec = registry.list_dataset_tasks("qwen", "my-bench") + + assert spec is not None + assert spec.id == "qwen/my-bench" + assert spec.split == "test" + assert spec.task_ids == ["task-001", "task-002"] + + first_call_kwargs = mock_bucket.list_objects_v2.call_args_list[0][1] + assert first_call_kwargs["prefix"] == "datasets/qwen/my-bench/test/" + + +def test_list_dataset_tasks_supports_custom_split(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[ + "datasets/qwen/my-bench/train/task-001/", + ]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + spec = registry.list_dataset_tasks("qwen", "my-bench", "train") + + assert spec is not None + assert spec.split == "train" + assert spec.task_ids == ["task-001"] + + first_call_kwargs = mock_bucket.list_objects_v2.call_args_list[0][1] + assert first_call_kwargs["prefix"] == "datasets/qwen/my-bench/train/" + +def test_list_dataset_tasks_includes_directory_and_file_tasks_with_suffix_stripped(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=["datasets/qwen/my-bench/test/task-002/"], + objects=[MagicMock(key="datasets/qwen/my-bench/test/task-001.json")], + ) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + spec = registry.list_dataset_tasks("qwen", "my-bench", "test") + + assert spec is not None + assert spec.id == "qwen/my-bench" + assert spec.split == "test" + assert spec.task_ids == ["task-001", "task-002"] + + +def test_list_dataset_tasks_ignores_placeholder_and_nested_objects(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[], + objects=[ + MagicMock(key="datasets/qwen/my-bench/test/"), + MagicMock(key="datasets/qwen/my-bench/test/nested/task-002.json"), + MagicMock(key="datasets/qwen/my-bench/test/task-001.json"), + ], + ) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + spec = registry.list_dataset_tasks("qwen", "my-bench", "test") + + assert spec is not None + assert spec.task_ids == ["task-001"] + + +def test_list_dataset_tasks_returns_none_when_no_tasks_found(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + spec = registry.list_dataset_tasks("qwen", "my-bench", "test") + + assert spec is None + # --------------------------------------------------------------------------- # upload_dataset tests From 2482aff9ccf184f52395af94e128c37578495529 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Wed, 22 Apr 2026 14:55:25 +0800 Subject: [PATCH 064/226] feat(docker): mount host zoneinfo to /etc/localtime for correct container timezone Instead of passing a POSIX TZ env var, mount the host's zoneinfo file (e.g. /usr/share/zoneinfo/Asia/Shanghai) to /etc/localtime:ro so containers get full IANA timezone support without requiring tzdata inside the image. Co-Authored-By: Claude Opus 4.6 --- .../_specs/sandbox-timezone/01_requirement.md | 91 +++++++++++++++++ docs/_specs/sandbox-timezone/02_interface.md | 97 +++++++++++++++++++ .../sandbox-timezone/03_implementation.md | 75 ++++++++++++++ rock/deployments/docker.py | 9 ++ tests/unit/rocklet/test_docker_deployment.py | 33 +++++-- 5 files changed, 299 insertions(+), 6 deletions(-) create mode 100644 docs/_specs/sandbox-timezone/01_requirement.md create mode 100644 docs/_specs/sandbox-timezone/02_interface.md create mode 100644 docs/_specs/sandbox-timezone/03_implementation.md diff --git a/docs/_specs/sandbox-timezone/01_requirement.md b/docs/_specs/sandbox-timezone/01_requirement.md new file mode 100644 index 0000000000..5784be91d7 --- /dev/null +++ b/docs/_specs/sandbox-timezone/01_requirement.md @@ -0,0 +1,91 @@ +# Sandbox Timezone — Requirement Spec + +## Background + +ROCK 的 Docker sandbox 启动流程此前只向容器传递了 `ROCK_TIME_ZONE`，该变量供 ROCK 自身日志和调度逻辑使用（需要 IANA 时区名，如 `Asia/Shanghai`），但并未设置容器内的系统时区。这导致容器内的系统时区始终为 UTC，产生两个具体问题： + +1. **文件修改时间偏差**：sandbox 内创建或修改的文件，其 `mtime` 按 UTC 记录。前端展示文件列表时，用户看到的修改时间与本地实际时间存在时差（东八区场景下差 8 小时）。 +2. **系统命令时间不一致**：`date`、`ls -l` 等系统命令输出 UTC 时间，与用户预期不符。 + +本次修复目标是： +- 让 sandbox 内的系统时区跟随宿主机配置，使文件修改时间和系统命令输出与用户所在时区一致 +- 前端展示文件信息时不再有时差偏差 +- 在镜像来源多样、不可控的前提下保持可用 +- 不要求业务镜像预装 `tzdata` + +--- + +## Solution + +在宿主机上确保安装 `tzdata`（`/usr/share/zoneinfo/` 目录存在），根据 `ROCK_TIME_ZONE` 环境变量（默认 `Asia/Shanghai`）定位对应的 zoneinfo 文件，通过 `docker run -v` 以只读方式挂载到容器的 `/etc/localtime`。 + +### 原理 + +Linux C 库（glibc / musl）解析系统时区的优先级： + +1. `TZ` 环境变量（如果设置了） +2. `/etc/localtime`（TZif 二进制文件） +3. 回退 UTC + +通过 bind mount 将宿主机的 zoneinfo 文件挂载到容器 `/etc/localtime`，Docker bind mount 会直接遮盖（shadow）容器内原有的 `/etc/localtime` 文件，无论容器镜像是否自带该文件。挂载后容器内的系统命令（`date`、`ls -l`、`stat`）均会按照挂载的时区文件解析时间。 + +### 兼容性 + +**容器侧**：bind mount `/etc/localtime` 后，glibc（Ubuntu/Debian/CentOS/RHEL）和 musl（Alpine）均能正确读取 TZif 文件，覆盖绝大多数 Linux 容器镜像。 + +**宿主机侧**：所有主流 Linux 发行版（Ubuntu、Debian、CentOS、RHEL、Amazon Linux）默认安装 `tzdata`，zoneinfo 路径统一为 `/usr/share/zoneinfo/`。运维侧保证宿主机有 `tzdata` 即可。 + +--- + +## In / Out + +### In（本次要做的） + +1. **Docker sandbox 启动时挂载 zoneinfo 文件到容器 `/etc/localtime`** + - 根据 `ROCK_TIME_ZONE`（默认 `Asia/Shanghai`）定位 `/usr/share/zoneinfo/{ROCK_TIME_ZONE}` + - 以只读方式挂载：`-v /usr/share/zoneinfo/{tz}:/etc/localtime:ro` + - 启动前校验文件是否存在，不存在则 warning 并跳过挂载 + +2. **保持现有 `ROCK_TIME_ZONE` 行为不变** + - `ROCK_TIME_ZONE` 使用 IANA 时区名（如 `Asia/Shanghai`），供 ROCK 日志、调度器、时间戳格式化等 Python 应用层逻辑使用 + - 该变量继续通过 `-e` 传入容器 + +3. **让容器内文件时间和系统命令与用户时区一致** + - `date`、`ls -l` 等命令按挂载的时区显示本地时间 + - 文件 `mtime` 的展示格式与用户预期时区一致 + - 前端读取文件修改时间时不再出现时差偏差 + +### Out（本次不做的） + +- 不修改镜像内容 +- 不要求所有业务镜像安装 `tzdata` +- 不向容器传递 `TZ` 环境变量（依靠 `/etc/localtime` 生效） +- 不修改 ROCK 内部 `ROCK_TIME_ZONE` 的默认值 +- 不新增 ROCK 环境变量 + +--- + +## Acceptance Criteria + +- **AC1**：Docker sandbox 启动时，`docker run` 包含 `-v /usr/share/zoneinfo/{tz}:/etc/localtime:ro` 挂载 +- **AC2**：`{tz}` 的值取自 `ROCK_TIME_ZONE`，默认 `Asia/Shanghai` +- **AC3**：挂载后容器内 `date` 命令输出对应时区的本地时间 +- **AC4**：当宿主机上对应的 zoneinfo 文件不存在时，打印 warning 日志并跳过挂载，不阻断启动 +- **AC5**：现有 `ROCK_TIME_ZONE` 行为不变（IANA 格式，供 Python 应用层使用） + +--- + +## Constraints + +- 不引入新的 Python 依赖 +- 不要求修改用户镜像 Dockerfile +- 不修改 sandbox 启动 API 的对外字段 +- 宿主机需安装 `tzdata`（`/usr/share/zoneinfo/` 目录存在） + +--- + +## Risks & Rollout + +- **风险**：宿主机未安装 `tzdata` 时挂载无效 — 通过 AC4 的 warning + 跳过机制缓解 +- **回滚**：仅涉及 `rock/deployments/docker.py`，回滚成本低 +- **上线策略**：无数据库变更，无协议破坏，可直接随 admin / deployment 代码发布 diff --git a/docs/_specs/sandbox-timezone/02_interface.md b/docs/_specs/sandbox-timezone/02_interface.md new file mode 100644 index 0000000000..00d2c957cd --- /dev/null +++ b/docs/_specs/sandbox-timezone/02_interface.md @@ -0,0 +1,97 @@ +# Sandbox Timezone — Interface Contract + +## 1. Runtime Environment Variables + +### 使用的变量 + +| 变量 | 来源 | 默认值 | 用途 | +|------|------|------|------| +| `ROCK_TIME_ZONE` | ROCK env vars | `Asia/Shanghai` | 1) ROCK 自身日志、调度等应用层逻辑 2) 确定挂载到容器的 zoneinfo 文件路径 | + +### 行为规则 + +- `ROCK_TIME_ZONE` 继续按原逻辑传入容器（`-e ROCK_TIME_ZONE=...`），默认 `Asia/Shanghai` +- 同时根据 `ROCK_TIME_ZONE` 的值定位 `/usr/share/zoneinfo/{ROCK_TIME_ZONE}`，挂载到容器 `/etc/localtime` +- 不再向容器传递 `TZ` 环境变量 + +--- + +## 2. Docker Run Contract + +### Volume 挂载 + +Docker sandbox 启动时，volume 参数至少包含： + +```bash +-v /usr/share/zoneinfo/:/etc/localtime:ro +``` + +在默认情况下等价于： + +```bash +-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime:ro +``` + +### 环境变量注入 + +```bash +-e ROCK_TIME_ZONE=Asia/Shanghai +``` + +### 示例 + +#### 示例 1：默认配置 + +```bash +docker run ... \ + -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime:ro \ + -e ROCK_TIME_ZONE=Asia/Shanghai \ + ... +``` + +#### 示例 2：`ROCK_TIME_ZONE=America/New_York` + +```bash +docker run ... \ + -v /usr/share/zoneinfo/America/New_York:/etc/localtime:ro \ + -e ROCK_TIME_ZONE=America/New_York \ + ... +``` + +#### 示例 3：zoneinfo 文件不存在 + +```bash +# /usr/share/zoneinfo/Invalid/Zone 不存在 +# → 打印 warning 日志，跳过挂载 +# → 容器时区回退为 UTC（镜像默认行为） +docker run ... \ + -e ROCK_TIME_ZONE=Invalid/Zone \ + ... +``` + +--- + +## 3. Container-side Observable Behavior + +### 可观测方式 + +| 检查方式 | 预期 | +|------|------| +| `date` | 按挂载的时区显示当前时间 | +| `ls -l` | 文件修改时间按挂载的时区显示 | +| `cat /etc/localtime` | 返回有效的 TZif 二进制数据 | + +### 边界说明 + +- 挂载的是完整的 IANA zoneinfo 文件，包含历史规则和夏令时信息 +- 容器内无需安装 `tzdata` 即可正确使用 +- 如果容器内程序同时设置了 `TZ` 环境变量，`TZ` 优先级高于 `/etc/localtime` + +--- + +## 4. Backward Compatibility + +- `ROCK_TIME_ZONE` 保持不变 +- 不新增对外 API 字段 +- 不修改 sandbox 启动请求模型 +- 新增 volume 挂载为纯新增行为，不影响已有挂载 diff --git a/docs/_specs/sandbox-timezone/03_implementation.md b/docs/_specs/sandbox-timezone/03_implementation.md new file mode 100644 index 0000000000..787453ea45 --- /dev/null +++ b/docs/_specs/sandbox-timezone/03_implementation.md @@ -0,0 +1,75 @@ +# Sandbox Timezone — Implementation Plan + +## 背景 + +Docker sandbox 此前未设置容器内系统时区，导致容器内系统时区为 UTC。具体表现为：文件修改时间按 UTC 记录，前端展示时与用户本地时间存在偏差；`date`、`ls -l` 等系统命令输出 UTC 时间，与用户预期不符。 + +本次实现通过将宿主机的 zoneinfo 文件挂载到容器 `/etc/localtime`，让容器内系统时区与宿主机配置一致。 + +--- + +## File Changes + +| 文件 | 修改类型 | 说明 | +|------|------|------| +| `rock/deployments/docker.py` | 修改 | 在 `_start` 方法中根据 `ROCK_TIME_ZONE` 挂载 zoneinfo 文件到 `/etc/localtime:ro` | +| `tests/unit/rocklet/test_docker_deployment.py` | 修改 | 验证挂载参数生成正确；集成测试验证容器内时区生效 | + +--- + +## Core Logic + +### 变更：Docker sandbox 挂载 `/etc/localtime` + +文件：`rock/deployments/docker.py`，`_start` 方法 + +在构建 `docker run` 命令时，根据 `ROCK_TIME_ZONE` 确定 zoneinfo 文件路径，若文件存在则追加 volume 挂载参数： + +```python +# 在 env_arg 和 volume_args 构建区域之后 +tz = env_vars.ROCK_TIME_ZONE # 默认 Asia/Shanghai +localtime_src = f"/usr/share/zoneinfo/{tz}" +if os.path.isfile(localtime_src): + volume_args.extend(["-v", f"{localtime_src}:/etc/localtime:ro"]) +else: + logger.warning(f"Zoneinfo file not found: {localtime_src}, skipping /etc/localtime mount") +``` + +### 设计要点 + +1. **用 `ROCK_TIME_ZONE` 而非 `TZ`**：`ROCK_TIME_ZONE` 始终是 IANA 格式（如 `Asia/Shanghai`），可直接映射到 `/usr/share/zoneinfo/` 下的文件路径。`TZ` 可能是 POSIX 格式（如 `CST-8`），无法映射到文件。 + +2. **文件存在性校验**：启动前用 `os.path.isfile()` 检查 zoneinfo 文件是否存在。不存在时打 warning 并跳过，不阻断容器启动。 + +3. **只读挂载**：使用 `:ro` 防止容器内进程修改宿主机时区文件。 + +4. **不传 `TZ` 环境变量**：挂载 `/etc/localtime` 已足够让 glibc/musl 正确解析时区，无需额外设置 `TZ`。避免 `TZ` 与 `/etc/localtime` 不一致导致的混乱。 + +--- + +## Validation Plan + +### 用例 1：单元测试 — 挂载参数生成 + +- mock `os.path.isfile` 返回 `True` +- 验证 `_start` 生成的 `docker run` 命令中包含 `-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime:ro` + +### 用例 2：单元测试 — zoneinfo 文件不存在时跳过 + +- mock `os.path.isfile` 返回 `False` +- 验证 `docker run` 命令中不包含 `/etc/localtime` 相关挂载 +- 验证打印了 warning 日志 + +### 用例 3：集成测试 — 真实 Docker 容器验证 + +- 前提：宿主机有 `/usr/share/zoneinfo/Asia/Shanghai` +- 启动 Docker 容器，挂载 `/etc/localtime` +- 在容器内执行 `date +%Z` 或 `ls -l /etc/localtime` +- 验证时区显示正确 + +--- + +## Rollback + +- 回滚仅需恢复 `rock/deployments/docker.py` 中的挂载逻辑 +- 对现有对外接口无兼容性影响 diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 05573dc726..09748df6cf 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -486,6 +486,7 @@ async def start(self): ] env_arg.extend(["-e", f"ROCK_TIME_ZONE={env_vars.ROCK_TIME_ZONE}"]) + volume_args.extend(self._prepare_timezone_mount()) # Kata DinD: prepare disk image and add volume mount + env var if self._config.use_kata_runtime: @@ -539,6 +540,14 @@ async def start(self): if self._config.enable_auto_clear: self._check_stop_task = asyncio.create_task(self._check_stop()) + def _prepare_timezone_mount(self) -> list[str]: + tz = env_vars.ROCK_TIME_ZONE + localtime_src = f"/usr/share/zoneinfo/{tz}" + if os.path.isfile(localtime_src): + return ["-v", f"{localtime_src}:/etc/localtime:ro"] + logger.warning(f"Zoneinfo file not found: {localtime_src}, skipping /etc/localtime mount") + return [] + def _prepare_volume_mounts(self) -> list[str]: mount_configs = self._runtime_env.get_volume_mounts() diff --git a/tests/unit/rocklet/test_docker_deployment.py b/tests/unit/rocklet/test_docker_deployment.py index 0926cd1c8e..c144e5d0bc 100644 --- a/tests/unit/rocklet/test_docker_deployment.py +++ b/tests/unit/rocklet/test_docker_deployment.py @@ -1,12 +1,9 @@ +import os + import pytest from rock import env_vars -from rock.actions import ( - BashAction, - CloseBashSessionRequest, - Command, - CreateBashSessionRequest, -) +from rock.actions import BashAction, CloseBashSessionRequest, Command, CreateBashSessionRequest from rock.deployments.config import DockerDeploymentConfig, get_deployment @@ -50,6 +47,30 @@ async def test_docker_deployment(container_name): await d.stop() +@pytest.mark.need_docker +async def test_docker_deployment_mounts_localtime_in_container(container_name): + tz = env_vars.ROCK_TIME_ZONE + host_has_zoneinfo = os.path.isfile(f"/usr/share/zoneinfo/{tz}") + + d = get_deployment( + DockerDeploymentConfig(image=env_vars.ROCK_ENVHUB_DEFAULT_DOCKER_IMAGE, container_name=container_name) + ) + try: + await d.start() + + if host_has_zoneinfo: + result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%z"])) + import subprocess + + host_offset = subprocess.check_output(["date", "+%z"], env={**os.environ, "TZ": tz}).decode().strip() + assert result.stdout.strip() == host_offset + else: + result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%Z"])) + assert result.stdout.strip() == "UTC" + finally: + await d.stop() + + def test_docker_deployment_config_platform(): config = DockerDeploymentConfig(docker_args=["--platform", "linux/amd64", "--other-arg"]) assert config.platform == "linux/amd64" From 269f83758db0cdb4f1f957d308d481ecb56f2920 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Wed, 22 Apr 2026 14:57:17 +0800 Subject: [PATCH 065/226] docs: update implementation spec to match actual test approach Co-Authored-By: Claude Opus 4.6 --- .../sandbox-timezone/03_implementation.md | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/docs/_specs/sandbox-timezone/03_implementation.md b/docs/_specs/sandbox-timezone/03_implementation.md index 787453ea45..9c78d6a162 100644 --- a/docs/_specs/sandbox-timezone/03_implementation.md +++ b/docs/_specs/sandbox-timezone/03_implementation.md @@ -49,23 +49,13 @@ else: ## Validation Plan -### 用例 1：单元测试 — 挂载参数生成 +### 集成测试 — 真实 Docker 容器验证 -- mock `os.path.isfile` 返回 `True` -- 验证 `_start` 生成的 `docker run` 命令中包含 `-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime:ro` +测试用例：`test_docker_deployment_mounts_localtime_in_container` -### 用例 2：单元测试 — zoneinfo 文件不存在时跳过 - -- mock `os.path.isfile` 返回 `False` -- 验证 `docker run` 命令中不包含 `/etc/localtime` 相关挂载 -- 验证打印了 warning 日志 - -### 用例 3：集成测试 — 真实 Docker 容器验证 - -- 前提：宿主机有 `/usr/share/zoneinfo/Asia/Shanghai` -- 启动 Docker 容器，挂载 `/etc/localtime` -- 在容器内执行 `date +%Z` 或 `ls -l /etc/localtime` -- 验证时区显示正确 +- 检测宿主机是否存在 `/usr/share/zoneinfo/{ROCK_TIME_ZONE}` 文件 +- **文件存在时**：启动容器，执行 `date +%z` 获取容器内 UTC offset，与宿主机在相同时区下的 offset 比对，确认一致 +- **文件不存在时**：启动容器，执行 `date +%Z`，确认回退到 UTC --- From 6e87d212318642ea1a4fe47e73b2a2ce4219841f Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Wed, 22 Apr 2026 15:01:52 +0800 Subject: [PATCH 066/226] feat(env): let ROCK_TIME_ZONE fall back to host TZ env var Priority: ROCK_TIME_ZONE > TZ > Asia/Shanghai (default). Co-Authored-By: Claude Opus 4.6 --- rock/env_vars.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rock/env_vars.py b/rock/env_vars.py index 4a1e26e512..92e4884d7f 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -128,7 +128,7 @@ "ROCK_MODEL_SERVICE_INSTALL_CMD", "pip install rl_rock[model-service]", ), - "ROCK_TIME_ZONE": lambda: os.getenv("ROCK_TIME_ZONE", "Asia/Shanghai"), + "ROCK_TIME_ZONE": lambda: os.getenv("ROCK_TIME_ZONE", os.getenv("TZ", "Asia/Shanghai")), "ROCK_DOCUUM_INSTALL_URL": lambda: os.getenv( "ROCK_DOCUUM_INSTALL_URL", "https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh" ), From 591f9d491e45a4faf4dadeee68e3518a6de54f91 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Wed, 22 Apr 2026 15:09:18 +0800 Subject: [PATCH 067/226] docs: clarify why ROCK_TIME_ZONE is used instead of TZ for zoneinfo lookup Using a dedicated ROCK_TIME_ZONE avoids conflicts when users specify their own TZ in sandbox launch parameters. Co-Authored-By: Claude Opus 4.6 --- docs/_specs/sandbox-timezone/03_implementation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/_specs/sandbox-timezone/03_implementation.md b/docs/_specs/sandbox-timezone/03_implementation.md index 9c78d6a162..5db40a0b19 100644 --- a/docs/_specs/sandbox-timezone/03_implementation.md +++ b/docs/_specs/sandbox-timezone/03_implementation.md @@ -37,7 +37,7 @@ else: ### 设计要点 -1. **用 `ROCK_TIME_ZONE` 而非 `TZ`**：`ROCK_TIME_ZONE` 始终是 IANA 格式（如 `Asia/Shanghai`），可直接映射到 `/usr/share/zoneinfo/` 下的文件路径。`TZ` 可能是 POSIX 格式（如 `CST-8`），无法映射到文件。 +1. **用 `ROCK_TIME_ZONE` 而非 `TZ` 定位 zoneinfo 文件**：`ROCK_TIME_ZONE` 始终是 IANA 格式（如 `Asia/Shanghai`），可直接映射到 `/usr/share/zoneinfo/` 下的文件路径。`TZ` 可能是 POSIX 格式（如 `CST-8`），无法映射到文件。此外，sandbox 启动 API 允许用户通过环境变量自定义容器配置，如果直接依赖 `TZ` 来定位 zoneinfo 文件，当用户在启动参数中指定了不同的 `TZ` 值时，会导致挂载的 `/etc/localtime` 与用户期望的 `TZ` 不一致。使用独立的 `ROCK_TIME_ZONE` 作为平台级配置，可以避免与用户指定的 `TZ` 变量冲突。 2. **文件存在性校验**：启动前用 `os.path.isfile()` 检查 zoneinfo 文件是否存在。不存在时打 warning 并跳过，不阻断容器启动。 From 2ac171447437e71636c771aeed55b7089271ba5c Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Wed, 22 Apr 2026 19:06:33 +0800 Subject: [PATCH 068/226] fix auto_clear_time (#883) --- rock/sdk/sandbox/client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 3037874266..04b2fbf082 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -1,5 +1,6 @@ import asyncio import logging +import math import mimetypes import os import time @@ -169,8 +170,8 @@ async def start(self): data = { "image": self.config.image, "image_os": self.config.image_os, - "auto_clear_time": self.config.auto_clear_seconds / 60, - "auto_clear_time_minutes": self.config.auto_clear_seconds / 60, + "auto_clear_time": int(math.ceil(self.config.auto_clear_seconds / 60)), + "auto_clear_time_minutes": int(math.ceil(self.config.auto_clear_seconds / 60)), "startup_timeout": self.config.startup_timeout, "memory": self.config.memory, "cpus": self.config.cpus, @@ -637,6 +638,7 @@ async def wait_for_process_completion( tuple[bool, str]: (success status, message) """ wait_interval = max(5, wait_interval) # Minimum interval 5 seconds + wait_interval = min(self.config.auto_clear_seconds - 2, wait_interval) # wait_interval < auto_clear_seconds check_alive_cmd = f"kill -0 {pid}" check_alive_timeout = min(wait_interval * 2, wait_timeout) # Not greater than wait_timeout From f47fa4595d5974660a35fd3e4e40692ae9c0b7b6 Mon Sep 17 00:00:00 2001 From: dengsheng Date: Wed, 22 Apr 2026 12:25:46 +0000 Subject: [PATCH 069/226] fix: preserve non-JSON request body in http_proxy endpoint --- rock/admin/entrypoints/sandbox_proxy_api.py | 9 ++-- rock/sandbox/service/sandbox_proxy_service.py | 6 +-- tests/unit/sandbox/test_sandbox_http_proxy.py | 6 +-- .../unit/sandbox/test_sandbox_proxy_router.py | 41 ++++++++++++++++++- 4 files changed, 48 insertions(+), 14 deletions(-) diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 27e1f9632f..8b0e1c5918 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -379,18 +379,15 @@ async def http_proxy( except BadRequestRockError as e: return _JSONResponse(status_code=400, content={"detail": str(e)}) - body = None + raw_body = None if request.method not in ("GET", "HEAD", "DELETE", "OPTIONS"): - try: - body = await request.json() - except Exception: - body = None + raw_body = await request.body() proxy_prefix = request.url.path.rstrip(path).rstrip("/") if path else request.url.path.rstrip("/") return await sandbox_proxy_service.http_proxy( sandbox_id, resolved_path, - body, + raw_body, request.headers, method=request.method, port=port, diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 451cebdf96..f4db7a1c36 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -810,7 +810,7 @@ async def http_proxy( self, sandbox_id: str, target_path: str, - body: dict | None, + body: bytes | None, headers: Headers, method: str = "POST", port: int | None = None, @@ -854,7 +854,7 @@ def rewrite_location(location: str) -> str: target_url = f"http://{host_ip}:{port}/{target_path}{qs}" request_headers = filter_headers(headers) - payload = body or {} + request_kwargs: dict = {"content": body} if body else {} client = httpx.AsyncClient(timeout=httpx.Timeout(None)) @@ -863,9 +863,9 @@ def rewrite_location(location: str) -> str: client.build_request( method=method, url=target_url, - json=payload if payload else None, headers=request_headers, timeout=120, + **request_kwargs, ), stream=True, ) diff --git a/tests/unit/sandbox/test_sandbox_http_proxy.py b/tests/unit/sandbox/test_sandbox_http_proxy.py index b1beb45f1b..a377e1255c 100644 --- a/tests/unit/sandbox/test_sandbox_http_proxy.py +++ b/tests/unit/sandbox/test_sandbox_http_proxy.py @@ -110,7 +110,7 @@ async def test_post_proxy(sandbox_manager: SandboxManager, sandbox_proxy_service result = await sandbox_proxy_service.http_proxy( sandbox_id=sandbox_id, target_path="api/test", - body={"hello": "world"}, + body=json.dumps({"hello": "world"}).encode(), headers=mock_headers, ) assert result.status_code == 200 @@ -122,7 +122,7 @@ async def test_post_proxy(sandbox_manager: SandboxManager, sandbox_proxy_service result = await sandbox_proxy_service.http_proxy( sandbox_id=sandbox_id, target_path="", - body={"key": "value"}, + body=json.dumps({"key": "value"}).encode(), headers=mock_headers, ) assert result.status_code == 200 @@ -145,7 +145,7 @@ async def test_post_proxy(sandbox_manager: SandboxManager, sandbox_proxy_service result = await sandbox_proxy_service.http_proxy( sandbox_id=sandbox_id, target_path="stream", - body={"msg": "hello"}, + body=json.dumps({"msg": "hello"}).encode(), headers=mock_headers, ) assert result.status_code == 200 diff --git a/tests/unit/sandbox/test_sandbox_proxy_router.py b/tests/unit/sandbox/test_sandbox_proxy_router.py index 245e9a93d8..6441a70c07 100644 --- a/tests/unit/sandbox/test_sandbox_proxy_router.py +++ b/tests/unit/sandbox/test_sandbox_proxy_router.py @@ -1,3 +1,4 @@ +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -29,7 +30,7 @@ async def test_post_proxy_path_parsing(app): args = mock_service.http_proxy.call_args assert args[0][0] == "sandbox-id" assert args[0][1] == "" - assert args[0][2] == {"key": "value"} + assert json.loads(args[0][2]) == {"key": "value"} mock_service.http_proxy.reset_mock() # Single path segment @@ -44,7 +45,7 @@ async def test_post_proxy_path_parsing(app): args = mock_service.http_proxy.call_args assert args[0][0] == "sandbox-id" assert args[0][1] == "api/v1/chat" - assert args[0][2] == {"msg": "hi"} + assert json.loads(args[0][2]) == {"msg": "hi"} mock_service.http_proxy.reset_mock() # Deep nested path @@ -52,3 +53,39 @@ async def test_post_proxy_path_parsing(app): args = mock_service.http_proxy.call_args assert args[0][0] == "sandbox-id" assert args[0][1] == "a/b/c/d" + + +@pytest.mark.asyncio +async def test_post_proxy_form_body(app): + app, mock_service = app + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + await client.post( + "/sandboxes/sandbox-id/proxy/submit", + data={"name": "test", "value": "123"}, + ) + args = mock_service.http_proxy.call_args + assert args[0][0] == "sandbox-id" + assert args[0][1] == "submit" + body = args[0][2] + assert isinstance(body, bytes) + assert b"name=test" in body + + +@pytest.mark.asyncio +async def test_post_proxy_raw_body(app): + app, mock_service = app + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + await client.post( + "/sandboxes/sandbox-id/proxy/echo", + content=b"raw data here", + ) + args = mock_service.http_proxy.call_args + assert args[0][0] == "sandbox-id" + assert args[0][1] == "echo" + body = args[0][2] + assert isinstance(body, bytes) + assert body == b"raw data here" From 68801c5467d9b353b0ddca9407fa18d82d3bd923 Mon Sep 17 00:00:00 2001 From: dengsheng Date: Thu, 23 Apr 2026 03:00:29 +0000 Subject: [PATCH 070/226] use raw body in vnc_http_proxy --- rock/admin/entrypoints/sandbox_proxy_api.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 8b0e1c5918..f4f61e253d 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -340,17 +340,15 @@ async def vnc_http_proxy( request: Request, path: str = "", ): - body = None + raw_body = None if request.method not in ("GET", "HEAD", "DELETE", "OPTIONS"): - try: - body = await request.json() - except Exception: - body = None + raw_body = await request.body() + proxy_prefix = request.url.path.rstrip(path).rstrip("/") return await sandbox_proxy_service.http_proxy( sandbox_id, path, - body, + raw_body, request.headers, method=request.method, port=8006, From 5ab42f8bb0a810e491f907d85d92f2c0b90d0a97 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 24 Apr 2026 11:35:13 +0800 Subject: [PATCH 071/226] Support meta store and database operation metrics (#887) * refactor(admin): move SandboxMetaStore import to module level * feat(metrics): add MetaStore and SandboxTable operation metrics * test(metrics): add unit tests for MetaStore metrics --- rock/admin/core/sandbox_table.py | 9 ++ rock/admin/main.py | 4 +- rock/admin/metrics/constants.py | 10 ++ rock/admin/metrics/decorator.py | 24 +++ rock/admin/metrics/monitor.py | 21 ++- rock/sandbox/sandbox_meta_store.py | 12 ++ .../admin/core/test_sandbox_table_metrics.py | 77 ++++++++++ .../admin/metrics/test_metastore_decorator.py | 138 ++++++++++++++++++ tests/unit/admin/metrics/test_monitor.py | 80 ++++++++++ tests/unit/conftest.py | 10 +- .../test_sandbox_meta_store_metrics.py | 91 ++++++++++++ 11 files changed, 469 insertions(+), 7 deletions(-) create mode 100644 tests/unit/admin/core/test_sandbox_table_metrics.py create mode 100644 tests/unit/admin/metrics/test_metastore_decorator.py create mode 100644 tests/unit/sandbox/test_sandbox_meta_store_metrics.py diff --git a/rock/admin/core/sandbox_table.py b/rock/admin/core/sandbox_table.py index 9141b07767..fe35b55a8b 100644 --- a/rock/admin/core/sandbox_table.py +++ b/rock/admin/core/sandbox_table.py @@ -9,6 +9,8 @@ from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.schema import SandboxRecord +from rock.admin.metrics.decorator import monitor_metastore_operation +from rock.admin.metrics.monitor import MetricsMonitor from rock.logger import init_logger if TYPE_CHECKING: @@ -36,7 +38,9 @@ class SandboxTable: def __init__(self, db_provider: DatabaseProvider) -> None: self._db = db_provider + self.metrics_monitor = MetricsMonitor.create(metric_prefix="meta_store.db") + @monitor_metastore_operation async def create( self, sandbox_id: str, @@ -66,6 +70,7 @@ async def create( session.add(record) await session.commit() + @monitor_metastore_operation async def get(self, sandbox_id: str) -> dict | None: """Return a sandbox row as a plain dict, or ``None`` if not found.""" async with AsyncSession(self._db.engine) as session: @@ -74,6 +79,7 @@ async def get(self, sandbox_id: str) -> dict | None: return None return record.to_dict() + @monitor_metastore_operation async def update(self, sandbox_id: str, info: SandboxInfo) -> None: """Partial update of scalar columns; always overwrites ``status`` with *info*.""" filtered = _pick_columns(info) @@ -88,6 +94,7 @@ async def update(self, sandbox_id: str, info: SandboxInfo) -> None: setattr(record, key, value) await session.commit() + @monitor_metastore_operation async def delete(self, sandbox_id: str) -> None: """Hard-delete a sandbox record.""" async with AsyncSession(self._db.engine) as session: @@ -96,6 +103,7 @@ async def delete(self, sandbox_id: str) -> None: await session.delete(record) await session.commit() + @monitor_metastore_operation async def list_by(self, column: str, value: str | int | float | bool) -> list[dict]: """Equality query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" if column not in SandboxRecord.LIST_BY_ALLOWLIST: @@ -106,6 +114,7 @@ async def list_by(self, column: str, value: str | int | float | bool) -> list[di result = await session.execute(stmt) return [r.to_dict() for r in result.scalars().all()] + @monitor_metastore_operation async def list_by_in(self, column: str, values: list[str | int | float | bool]) -> list[dict]: """IN query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" if column not in SandboxRecord.LIST_BY_ALLOWLIST: diff --git a/rock/admin/main.py b/rock/admin/main.py index fc979557e5..e44b6c427e 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -25,6 +25,7 @@ from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager from rock.sandbox.operator.factory import OperatorContext, OperatorFactory +from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sandbox.service.warmup_service import WarmupService from rock.utils import EAGLE_EYE_TRACE_ID, sandbox_id_ctx_var, trace_id_ctx_var @@ -78,9 +79,6 @@ async def lifespan(app: FastAPI): if not rock_config.database.url: await db_provider.create_tables() sandbox_table = SandboxTable(db_provider) - - from rock.sandbox.sandbox_meta_store import SandboxMetaStore - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=sandbox_table) # init scheduler thread diff --git a/rock/admin/metrics/constants.py b/rock/admin/metrics/constants.py index 388d9740a2..1ef95f14ab 100644 --- a/rock/admin/metrics/constants.py +++ b/rock/admin/metrics/constants.py @@ -21,3 +21,13 @@ class MetricsConstants: AVAILABLE_MEM_RESOURCE = "resource.mem.available" SANDBOX_PHASE_FAILURE = "sandbox.phase.failure" + + METASTORE_TOTAL = "meta_store.total" + METASTORE_SUCCESS = "meta_store.success" + METASTORE_FAILURE = "meta_store.failure" + METASTORE_RT = "meta_store.rt" + + METASTORE_DB_TOTAL = "meta_store.db.total" + METASTORE_DB_SUCCESS = "meta_store.db.success" + METASTORE_DB_FAILURE = "meta_store.db.failure" + METASTORE_DB_RT = "meta_store.db.rt" diff --git a/rock/admin/metrics/decorator.py b/rock/admin/metrics/decorator.py index 226b5a22dd..98788a8645 100644 --- a/rock/admin/metrics/decorator.py +++ b/rock/admin/metrics/decorator.py @@ -226,3 +226,27 @@ def wrapper(self, *args, **kwargs): if func is not None: return decorator(func) return decorator + + +def monitor_metastore_operation(f): + """Decorator for async MetaStore/SandboxTable methods.""" + + @functools.wraps(f) + async def wrapper(self, *args, **kwargs): + metrics_monitor = getattr(self, "metrics_monitor", None) + if not metrics_monitor: + return await f(self, *args, **kwargs) + + prefix = metrics_monitor.metric_prefix + sandbox_id = _extract_sandbox_id(args, kwargs) + attributes: dict[str, str] = {"operation": f.__name__, "method": f.__name__, "sandbox_id": sandbox_id} + + start_time = time.perf_counter() + + try: + result = await f(self, *args, **kwargs) + return _record_metrics(metrics_monitor, result, attributes, start_time, prefix) + except Exception as e: + return _record_metrics(metrics_monitor, e, attributes, start_time, prefix) + + return wrapper diff --git a/rock/admin/metrics/monitor.py b/rock/admin/metrics/monitor.py index d528d38b9c..77ebcce719 100644 --- a/rock/admin/metrics/monitor.py +++ b/rock/admin/metrics/monitor.py @@ -26,8 +26,10 @@ def __init__( export_interval_millis: int = 10000, endpoint: str = "", user_defined_tags: dict = {}, + metric_prefix: str = "", ): patch_view_instrument_match() + self.metric_prefix = metric_prefix self.user_defined_tags = user_defined_tags self._init_basic_attributes(host, port, pod, env, role) self.endpoint = endpoint or f"http://{self.host}:{self.port}/v1/metrics" @@ -42,7 +44,11 @@ def __init__( @classmethod def create( - cls, export_interval_millis: int = 20000, metrics_endpoint: str = "", user_defined_tags: dict = {} + cls, + export_interval_millis: int = 20000, + metrics_endpoint: str = "", + user_defined_tags: dict = {}, + metric_prefix: str = "", ) -> "MetricsMonitor": host, port = get_uniagent_endpoint() pod = get_instance_id() @@ -58,6 +64,7 @@ def create( export_interval_millis=export_interval_millis, endpoint=metrics_endpoint, user_defined_tags=user_defined_tags, + metric_prefix=metric_prefix, ) def _register_metrics(self): @@ -96,6 +103,18 @@ def _register_metrics(self): self._register_gauge(MetricsConstants.AVAILABLE_CPU_RESOURCE, "Available CPU resource in Ray cluster") self._register_gauge(MetricsConstants.AVAILABLE_MEM_RESOURCE, "Available memory resource in Ray cluster") + # MetaStore metrics + self._register_counter(MetricsConstants.METASTORE_SUCCESS, "Number of successful meta_store operations") + self._register_counter(MetricsConstants.METASTORE_FAILURE, "Number of failed meta_store operations") + self._register_counter(MetricsConstants.METASTORE_TOTAL, "Number of total meta_store operations") + self._register_gauge(MetricsConstants.METASTORE_RT, "MetaStore operation response time", "ms") + + # MetaStore DB-layer metrics + self._register_counter(MetricsConstants.METASTORE_DB_SUCCESS, "Number of successful DB operations") + self._register_counter(MetricsConstants.METASTORE_DB_FAILURE, "Number of failed DB operations") + self._register_counter(MetricsConstants.METASTORE_DB_TOTAL, "Number of total DB operations") + self._register_gauge(MetricsConstants.METASTORE_DB_RT, "DB operation response time", "ms") + def _register_counter(self, name: str, description: str, unit: str = "1"): self.counters[name] = self.create_counter(name, description, unit) diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index ea1ed3b30a..d392de1027 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -14,6 +14,8 @@ from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.redis_key import alive_sandbox_key, timeout_sandbox_key from rock.admin.core.sandbox_table import SandboxTable +from rock.admin.metrics.decorator import monitor_metastore_operation +from rock.admin.metrics.monitor import MetricsMonitor if TYPE_CHECKING: from rock.deployments.config import DockerDeploymentConfig @@ -39,11 +41,13 @@ def __init__( ) -> None: self._redis: RedisProvider = redis_provider self._db: SandboxTable = sandbox_table + self.metrics_monitor = MetricsMonitor.create(metric_prefix="meta_store") # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ + @monitor_metastore_operation async def create( self, sandbox_id: str, @@ -67,6 +71,7 @@ async def create( await self._db.create(sandbox_id, sandbox_info, deployment_config) + @monitor_metastore_operation async def update(self, sandbox_id: str, sandbox_info: SandboxInfo) -> None: """Merge *sandbox_info* into the existing Redis alive key and await DB update.""" current = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") @@ -75,6 +80,7 @@ async def update(self, sandbox_id: str, sandbox_info: SandboxInfo) -> None: await self._db.update(sandbox_id, sandbox_info) + @monitor_metastore_operation async def delete(self, sandbox_id: str) -> None: """Delete Redis alive + timeout keys and await DB delete.""" await self._redis.json_delete(alive_sandbox_key(sandbox_id)) @@ -82,6 +88,7 @@ async def delete(self, sandbox_id: str) -> None: await self._db.delete(sandbox_id) + @monitor_metastore_operation async def archive(self, sandbox_id: str, final_info: SandboxInfo) -> None: """Persist final state to DB, then remove sandbox from Redis. @@ -100,6 +107,7 @@ async def archive(self, sandbox_id: str, final_info: SandboxInfo) -> None: await self._redis.json_delete(alive_sandbox_key(sandbox_id)) await self._redis.json_delete(timeout_sandbox_key(sandbox_id)) + @monitor_metastore_operation async def get(self, sandbox_id: str) -> SandboxInfo | None: """Read sandbox info from the Redis alive key.""" result = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") @@ -111,6 +119,7 @@ async def exists(self, sandbox_id: str) -> bool: """Return ``True`` when the Redis alive key exists for ``sandbox_id``.""" return await self.get(sandbox_id) is not None + @monitor_metastore_operation async def get_timeout(self, sandbox_id: str) -> dict[str, str] | None: """Read timeout info from the Redis timeout key.""" timeout_info = await self._redis.json_get(timeout_sandbox_key(sandbox_id), "$") @@ -118,6 +127,7 @@ async def get_timeout(self, sandbox_id: str) -> dict[str, str] | None: return timeout_info[0] return None + @monitor_metastore_operation async def update_timeout(self, sandbox_id: str, timeout_info: dict[str, str]) -> None: """Overwrite the Redis timeout key with *timeout_info*.""" await self._redis.json_set(timeout_sandbox_key(sandbox_id), "$", timeout_info) @@ -129,6 +139,7 @@ async def iter_alive_sandbox_ids(self) -> AsyncIterator[str]: if sandbox_id: yield sandbox_id + @monitor_metastore_operation async def batch_get(self, sandbox_ids: list[str]) -> list[SandboxInfo]: """Fetch sandbox info for multiple IDs from the DB. Missing IDs are omitted.""" if not sandbox_ids: @@ -136,6 +147,7 @@ async def batch_get(self, sandbox_ids: list[str]) -> list[SandboxInfo]: return await self._db.list_by_in("sandbox_id", sandbox_ids) + @monitor_metastore_operation async def list_by(self, field: str, value: str | int | float | bool) -> list[SandboxInfo]: """Query sandboxes by *field* == *value* from the DB.""" return await self._db.list_by(field, value) diff --git a/tests/unit/admin/core/test_sandbox_table_metrics.py b/tests/unit/admin/core/test_sandbox_table_metrics.py new file mode 100644 index 0000000000..f11e634579 --- /dev/null +++ b/tests/unit/admin/core/test_sandbox_table_metrics.py @@ -0,0 +1,77 @@ +"""Tests for SandboxTable DB-layer metrics.""" + +from unittest.mock import Mock, patch + +import pytest + +from rock.admin.core.sandbox_table import SandboxTable +from rock.admin.metrics.monitor import MetricsMonitor + +SANDBOX_ID = "sbx-db-metrics-001" +SANDBOX_INFO = { + "user_id": "user-1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "cluster-1", + "state": "running", + "host_ip": "10.0.0.1", + "create_time": "2025-01-01T00:00:00Z", +} + + +PREFIX = "meta_store.db" + + +@pytest.fixture +def mock_monitor(): + monitor = Mock(spec=MetricsMonitor) + monitor._should_skip.return_value = False + monitor.metric_prefix = PREFIX + return monitor + + +@pytest.fixture +def table(db_provider, mock_monitor): + with patch("rock.admin.core.sandbox_table.MetricsMonitor.create", return_value=mock_monitor): + return SandboxTable(db_provider) + + +class TestSandboxTableMetrics: + async def test_create_records_db_metrics(self, table, mock_monitor): + await table.create(SANDBOX_ID, SANDBOX_INFO) + + attrs = {"operation": "create", "method": "create", "sandbox_id": SANDBOX_ID} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) + rt_call = mock_monitor.record_gauge_by_name.call_args + assert rt_call[0][0] == f"{PREFIX}.rt" + assert rt_call[0][1] > 0 + + async def test_get_records_db_metrics(self, table, mock_monitor): + await table.create(SANDBOX_ID, SANDBOX_INFO) + mock_monitor.reset_mock() + + await table.get(SANDBOX_ID) + + attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) + mock_monitor.record_gauge_by_name.assert_called_once() + + async def test_failure_records_error_type(self, table, mock_monitor): + table._db = Mock() + table._db.engine = property(lambda self: (_ for _ in ()).throw(RuntimeError("db down"))) + + with pytest.raises(Exception): + await table.get("nonexistent-will-fail") + + async def test_list_by_records_db_metrics(self, table, mock_monitor): + await table.create(SANDBOX_ID, SANDBOX_INFO) + mock_monitor.reset_mock() + + results = await table.list_by("user_id", "user-1") + + assert len(results) == 1 + attrs = {"operation": "list_by", "method": "list_by", "sandbox_id": "user_id"} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) diff --git a/tests/unit/admin/metrics/test_metastore_decorator.py b/tests/unit/admin/metrics/test_metastore_decorator.py new file mode 100644 index 0000000000..64b4fa8193 --- /dev/null +++ b/tests/unit/admin/metrics/test_metastore_decorator.py @@ -0,0 +1,138 @@ +"""Tests for monitor_metastore_operation decorator.""" + +from unittest.mock import Mock + +import pytest + +from rock.admin.metrics.decorator import monitor_metastore_operation +from rock.admin.metrics.monitor import MetricsMonitor + + +def _make_monitor(prefix="meta_store"): + monitor = Mock(spec=MetricsMonitor) + monitor._should_skip.return_value = False + monitor.metric_prefix = prefix + return monitor + + +class FakeStore: + """Minimal store-like object for testing the decorator.""" + + def __init__(self, metrics_monitor=None): + self.metrics_monitor = metrics_monitor + + @monitor_metastore_operation + async def do_something(self): + return "ok" + + @monitor_metastore_operation + async def do_fail(self): + raise ValueError("boom") + + @monitor_metastore_operation + async def get(self, sandbox_id: str): + return sandbox_id + + @monitor_metastore_operation + async def list_by(self, field: str, value: str): + return [] + + +class TestMonitorMetastoreOperation: + async def test_records_success_metrics(self): + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + result = await store.do_something() + + assert result == "ok" + attrs = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} + monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + monitor.record_counter_by_name.assert_any_call("meta_store.total", 1, attrs) + monitor.record_gauge_by_name.assert_called_once() + call_args = monitor.record_gauge_by_name.call_args + assert call_args[0][0] == "meta_store.rt" + assert call_args[0][2] == attrs + + async def test_records_failure_metrics(self): + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + with pytest.raises(ValueError, match="boom"): + await store.do_fail() + + error_attrs = {"operation": "do_fail", "method": "do_fail", "sandbox_id": "unknown", "error_type": "ValueError"} + monitor.record_counter_by_name.assert_any_call("meta_store.failure", 1, error_attrs) + monitor.record_counter_by_name.assert_any_call("meta_store.total", 1, error_attrs) + monitor.record_gauge_by_name.assert_called_once() + + async def test_skips_when_no_monitor(self): + store = FakeStore(metrics_monitor=None) + + result = await store.do_something() + + assert result == "ok" + + async def test_rt_is_positive(self): + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + await store.do_something() + + call_args = monitor.record_gauge_by_name.call_args + rt_value = call_args[0][1] + assert rt_value > 0 + + async def test_uses_monitor_prefix(self): + """Verify the decorator reads metric_prefix from the monitor instance.""" + monitor = _make_monitor(prefix="meta_store.db") + store = FakeStore(metrics_monitor=monitor) + + await store.do_something() + + attrs_db = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} + monitor.record_counter_by_name.assert_any_call("meta_store.db.success", 1, attrs_db) + monitor.record_counter_by_name.assert_any_call("meta_store.db.total", 1, attrs_db) + call_args = monitor.record_gauge_by_name.call_args + assert call_args[0][0] == "meta_store.db.rt" + + async def test_sandbox_id_in_attributes_positional(self): + """sandbox_id is included when the method declares the parameter (positional).""" + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + result = await store.get("sbx-123") + + assert result == "sbx-123" + attrs = {"operation": "get", "method": "get", "sandbox_id": "sbx-123"} + monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + + async def test_sandbox_id_unknown_when_keyword_only(self): + """sandbox_id passed as keyword is not extracted (consistent with monitor_sandbox_operation).""" + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + await store.get(sandbox_id="sbx-kw") + + attrs = {"operation": "get", "method": "get", "sandbox_id": "unknown"} + monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + + async def test_sandbox_id_fallback_from_first_arg(self): + """Without sandbox_id param, _extract_sandbox_id falls back to args[0].""" + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + await store.list_by("state", "running") + + attrs = {"operation": "list_by", "method": "list_by", "sandbox_id": "state"} + monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + + async def test_sandbox_id_unknown_when_no_args(self): + """No args at all → sandbox_id defaults to 'unknown'.""" + monitor = _make_monitor() + store = FakeStore(metrics_monitor=monitor) + + await store.do_something() + + attrs = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} + monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) diff --git a/tests/unit/admin/metrics/test_monitor.py b/tests/unit/admin/metrics/test_monitor.py index dbcba40621..9360db3d47 100644 --- a/tests/unit/admin/metrics/test_monitor.py +++ b/tests/unit/admin/metrics/test_monitor.py @@ -1,5 +1,6 @@ from unittest.mock import patch +from rock.admin.metrics.constants import MetricsConstants from rock.admin.metrics.monitor import MetricsMonitor, aggregate_metrics @@ -164,3 +165,82 @@ def test_create_with_user_defined_tags(mock_env_vars, mock_instance_id, mock_uni assert monitor.base_attributes["pod"] == "test-pod" assert monitor.base_attributes["env"] == "daily" assert monitor.base_attributes["role"] == "test" + + +def _create_dev_monitor(metric_prefix: str = "") -> MetricsMonitor: + """Create a real MetricsMonitor with env=dev (InMemoryMetricReader, no skip).""" + return MetricsMonitor( + host="127.0.0.1", + port="4318", + pod="test-pod", + env="dev", + role="test", + metric_prefix=metric_prefix, + ) + + +class TestMetastoreMetricsRegistration: + """Verify that metastore metric names are registered and usable on a real monitor.""" + + def test_metastore_counters_registered(self): + monitor = _create_dev_monitor() + for name in [ + MetricsConstants.METASTORE_SUCCESS, + MetricsConstants.METASTORE_FAILURE, + MetricsConstants.METASTORE_TOTAL, + MetricsConstants.METASTORE_DB_SUCCESS, + MetricsConstants.METASTORE_DB_FAILURE, + MetricsConstants.METASTORE_DB_TOTAL, + ]: + assert name in monitor.counters, f"Counter '{name}' not registered" + assert monitor.counters[name] is not None, f"Counter '{name}' is None" + + def test_metastore_gauges_registered(self): + monitor = _create_dev_monitor() + for name in [ + MetricsConstants.METASTORE_RT, + MetricsConstants.METASTORE_DB_RT, + ]: + assert name in monitor.gauges, f"Gauge '{name}' not registered" + assert monitor.gauges[name] is not None, f"Gauge '{name}' is None" + + def test_record_counter_by_name_does_not_raise(self): + """Calling record_counter_by_name with registered metric names should not KeyError.""" + monitor = _create_dev_monitor() + attrs = {"operation": "create", "method": "create"} + monitor.record_counter_by_name(MetricsConstants.METASTORE_SUCCESS, 1, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_TOTAL, 1, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_DB_SUCCESS, 1, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_DB_TOTAL, 1, attrs) + + def test_record_gauge_by_name_does_not_raise(self): + """Calling record_gauge_by_name with registered metric names should not KeyError.""" + monitor = _create_dev_monitor() + attrs = {"operation": "create", "method": "create"} + monitor.record_gauge_by_name(MetricsConstants.METASTORE_RT, 1.5, attrs) + monitor.record_gauge_by_name(MetricsConstants.METASTORE_DB_RT, 0.8, attrs) + + def test_metric_prefix_stored(self): + monitor = _create_dev_monitor(metric_prefix="meta_store") + assert monitor.metric_prefix == "meta_store" + + def test_end_to_end_record_does_not_raise(self): + """Full round-trip: create real monitor, record metrics, no errors. + + OTel's global MeterProvider can only be set once per process, so + subsequent dev monitors share the first reader. We verify the + record path completes without exceptions, which proves the counters + and gauges are real OTel instruments (not None). + """ + monitor = _create_dev_monitor(metric_prefix="meta_store") + attrs = {"operation": "get", "method": "get"} + # These would raise KeyError if names are unregistered, + # or AttributeError if instruments are None. + monitor.record_counter_by_name(MetricsConstants.METASTORE_SUCCESS, 1, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_FAILURE, 1, {**attrs, "error_type": "ValueError"}) + monitor.record_counter_by_name(MetricsConstants.METASTORE_TOTAL, 1, attrs) + monitor.record_gauge_by_name(MetricsConstants.METASTORE_RT, 5.0, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_DB_SUCCESS, 1, attrs) + monitor.record_counter_by_name(MetricsConstants.METASTORE_DB_FAILURE, 1, {**attrs, "error_type": "IOError"}) + monitor.record_counter_by_name(MetricsConstants.METASTORE_DB_TOTAL, 1, attrs) + monitor.record_gauge_by_name(MetricsConstants.METASTORE_DB_RT, 2.0, attrs) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 864c1f2c49..55f697d2ea 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -82,15 +82,19 @@ def ray_operator(ray_service, runtime_config): @pytest.fixture -async def _memory_sandbox_table(): +async def db_provider(): provider = DatabaseProvider(db_config=DatabaseConfig(url="sqlite+aiosqlite:///:memory:")) await provider.init() await provider.create_tables() - table = SandboxTable(provider) - yield table + yield provider await provider.close() +@pytest.fixture +async def _memory_sandbox_table(db_provider): + return SandboxTable(db_provider) + + @pytest.fixture async def sandbox_manager( rock_config: RockConfig, diff --git a/tests/unit/sandbox/test_sandbox_meta_store_metrics.py b/tests/unit/sandbox/test_sandbox_meta_store_metrics.py new file mode 100644 index 0000000000..1906b23281 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_meta_store_metrics.py @@ -0,0 +1,91 @@ +"""Tests for SandboxMetaStore metrics integration.""" + +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from fakeredis import aioredis + +from rock.actions.sandbox.response import State +from rock.admin.core.sandbox_table import SandboxTable +from rock.admin.metrics.monitor import MetricsMonitor +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.utils.providers.redis_provider import RedisProvider + +SANDBOX_ID = "sbx-metrics-001" +SANDBOX_INFO = { + "sandbox_id": SANDBOX_ID, + "user_id": "user-1", + "image": "python:3.11", + "experiment_id": "exp-1", + "namespace": "default", + "cluster_name": "cluster-1", + "state": State.RUNNING, + "host_ip": "10.0.0.1", + "create_time": "2025-01-01T00:00:00Z", +} + + +PREFIX = "meta_store" +DB_PREFIX = "meta_store.db" + + +def _make_monitor(prefix=PREFIX): + monitor = Mock(spec=MetricsMonitor) + monitor._should_skip.return_value = False + monitor.metric_prefix = prefix + return monitor + + +@pytest.fixture +async def redis(): + provider = RedisProvider(host=None, port=None, password="") + provider.client = aioredis.FakeRedis(decode_responses=True) + yield provider + await provider.close_pool() + + +@pytest.fixture +def mock_monitor(): + return _make_monitor(PREFIX) + + +@pytest.fixture +def store(redis, db_provider, mock_monitor): + db_monitor = _make_monitor(DB_PREFIX) + with patch("rock.admin.core.sandbox_table.MetricsMonitor.create", return_value=db_monitor): + table = SandboxTable(db_provider) + with patch("rock.sandbox.sandbox_meta_store.MetricsMonitor.create", return_value=mock_monitor): + return SandboxMetaStore(redis_provider=redis, sandbox_table=table) + + +class TestMetaStoreMetrics: + async def test_create_records_store_metrics(self, store, mock_monitor): + await store.create(SANDBOX_ID, SANDBOX_INFO) + + attrs = {"operation": "create", "method": "create", "sandbox_id": SANDBOX_ID} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) + assert mock_monitor.record_gauge_by_name.called + rt_call = mock_monitor.record_gauge_by_name.call_args + assert rt_call[0][0] == f"{PREFIX}.rt" + assert rt_call[0][1] > 0 + + async def test_get_records_store_metrics(self, store, redis, mock_monitor): + await store.create(SANDBOX_ID, SANDBOX_INFO) + mock_monitor.reset_mock() + + await store.get(SANDBOX_ID) + + attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) + mock_monitor.record_gauge_by_name.assert_called_once() + + async def test_failure_records_error_type(self, store, redis, mock_monitor): + redis.json_get = AsyncMock(side_effect=ConnectionError("redis down")) + + with pytest.raises(ConnectionError): + await store.get(SANDBOX_ID) + + error_attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID, "error_type": "ConnectionError"} + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.failure", 1, error_attrs) + mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, error_attrs) From de07eeda0c4e3b07c048e4b2ba1ed3e8b7cba712 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 24 Apr 2026 10:55:30 +0800 Subject: [PATCH 072/226] fix(proxy): expand WebSocket header blacklist and add VNC forwarding switch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add forward_ws_headers bool to websocket_proxy() — VNC route sets it to False to avoid exceeding QEMU's 4KB header buffer limit - Update proxy-enhancements spec docs to reflect current implementation Co-Authored-By: Claude Opus 4.6 --- .../proxy-enhancements/01_requirement.md | 18 +++-- .../proxy-enhancements/03_implementation.md | 77 +++++++++++++++---- rock/admin/entrypoints/sandbox_proxy_api.py | 2 +- rock/sandbox/service/sandbox_proxy_service.py | 13 +++- 4 files changed, 86 insertions(+), 24 deletions(-) diff --git a/docs/_specs/proxy-enhancements/01_requirement.md b/docs/_specs/proxy-enhancements/01_requirement.md index 85653b629a..275b18ecd7 100644 --- a/docs/_specs/proxy-enhancements/01_requirement.md +++ b/docs/_specs/proxy-enhancements/01_requirement.md @@ -41,11 +41,18 @@ ROCK Admin 目前提供两类代理能力： 4. **WebSocket Proxy 支持黑名单过滤透传通用请求头** - 通过 `/sandboxes/{id}/proxy/{path:path}` 建立 WebSocket 代理时，默认将客户端请求中的通用 header 转发到下游服务，通过黑名单排除不应转发的头 - - 黑名单排除的头包括：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`） + - 黑名单排除的头包括： + - WebSocket 握手专用头：`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol` + - Hop-by-hop 头：`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length` - `Origin` 是必须支持的关键头，因为下游服务可能使用来源白名单（例如 `gateway.controlUi.allowedOrigins`）校验 WebSocket 握手 - `Sec-WebSocket-Protocol` 继续通过 WebSocket 子协议协商传递，不作为普通 header 透传 - 采用黑名单策略，确保用户自定义 header 能被透传到下游服务 +5. **VNC WebSocket Proxy 不转发客户端 header** + - VNC 路由（`/sandboxes/{id}/vnc/{path}`）调用 `websocket_proxy()` 时设置 `forward_ws_headers=False` + - QEMU 内置 WebSocket 服务器的握手 header buffer 仅 4KB（`QIO_CHANNEL_WEBSOCK_MAX_HEADER_SIZE`），上游网关注入的大量 header（SSO cookie、安全校验头等）会导致超限返回 502 + - VNC 服务不需要客户端的 Origin、Cookie、认证头等上下文，关闭转发最安全 + ### Out（本次不做的） - WebSocket Proxy 的认证/鉴权增强 @@ -69,7 +76,7 @@ ROCK Admin 目前提供两类代理能力： - **AC8**：`GET /sandboxes/{sandbox_id}/proxy/api/health?port=9000` 能成功代理到沙箱内 9000 端口的 HTTP 服务 - **AC9**：HTTP proxy 不带 `port` 参数时行为不变（向后兼容） - **AC10**：当客户端 WebSocket 握手包含 `Origin` 时，代理发起到下游的二跳握手必须携带相同 `Origin`，以满足下游来源校验 -- **AC11**：当客户端握手包含 `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent` 或任意自定义头时，代理发起到下游的二跳握手必须一并转发 +- **AC11**：当客户端握手包含 `Authorization`、`X-Request-Id`、`Traceparent`、`Cookie`、`Accept-Encoding`、`X-Forwarded-*`、`X-Real-IP` 或任意自定义头时，代理发起到下游的二跳握手必须一并转发 - **AC12**：`Sec-WebSocket-Protocol` 必须继续通过现有 `subprotocols` 机制转发和协商，不能降级为普通 header 透传 - **AC13**：黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`）不得被转发到下游 - **AC14**：当客户端未携带任何白名单头时，WebSocket proxy 的默认行为与现状保持一致（向后兼容） @@ -91,7 +98,8 @@ ROCK Admin 目前提供两类代理能力： ## Risks & Rollout - **风险**：WebSocket proxy 中用户可以指定任意端口访问沙箱内服务，存在横向访问风险 → 通过 `sandbox_id` 鉴权已覆盖，端口范围校验作为防护层 -- **风险**：黑名单策略下，未列入黑名单的头会默认转发 → 通过单元测试确保握手专用头和 hop-by-hop 头被正确过滤 -- **风险**：不同下游服务对 `Origin`、`Cookie`、`Authorization` 的要求不同，透传后暴露出原本被代理层掩盖的问题 → 以“尽可能保留客户端上下文、但不伪造默认值”为原则 -- **回滚**：修改集中在 `sandbox_proxy_service.py` 和对应单元测试；如需回滚，可仅还原 WebSocket header 透传逻辑 +- **风险**：黑名单策略下，未列入黑名单的头会默认转发 → 通过单元测试确保握手专用头、hop-by-hop 头、网关注入头被正确过滤 +- **风险**：不同下游服务对 `Origin`、`Authorization` 的要求不同，透传后暴露出原本被代理层掩盖的问题 → 以”尽可能保留客户端上下文、但不伪造默认值”为原则 +- **风险**：下游 WebSocket 服务（如 QEMU VNC）header buffer 有限（4KB），大量网关注入头会导致 502 → VNC 路由通过 `forward_ws_headers=False` 关闭 header 转发 +- **回滚**：修改集中在 `sandbox_proxy_service.py`、`rock/sandbox/utils/proxy.py` 和对应单元测试；如需回滚，可仅还原 WebSocket header 透传逻辑 - **上线策略**：无数据库变更，直接部署；建议先在依赖 `Origin` 校验的控制台类服务上验证 diff --git a/docs/_specs/proxy-enhancements/03_implementation.md b/docs/_specs/proxy-enhancements/03_implementation.md index 20dd96d05e..d74e780f40 100644 --- a/docs/_specs/proxy-enhancements/03_implementation.md +++ b/docs/_specs/proxy-enhancements/03_implementation.md @@ -19,9 +19,10 @@ admin 与 sandbox 不在同一 K8s 集群，`host_ip` 为宿主机 IP，容器 | 文件 | 修改类型 | 说明 | |------|------|------| | `rock/rocklet/local_api.py` | **新增** | 新增 `ANY /http_proxy/{path:path}?port={port}` 端点 | -| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转；补充 WebSocket 通用 headers 白名单透传 | -| `rock/admin/entrypoints/sandbox_proxy_api.py` | 无变更 | 路由签名保持不变，无需调整 | -| `tests/unit/sandbox/test_websocket_proxy_subprotocol.py` | 修改 | 增加 `Origin` / `additional_headers` 透传与禁转头测试 | +| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转；新增 `forward_ws_headers` 开关控制 header 透传 | +| `rock/sandbox/utils/proxy.py` | **新增** | `BLOCKED_WS_HEADER_NAMES` 黑名单 + `build_upstream_ws_headers()` helper | +| `rock/admin/entrypoints/sandbox_proxy_api.py` | 修改 | VNC WebSocket 路由传入 `forward_ws_headers=False` | +| `tests/unit/sandbox/test_websocket_proxy_headers.py` | **新增** | `Origin` / `additional_headers` 透传、黑名单过滤、E2E 测试 | --- @@ -118,7 +119,7 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" ### 变更 4：WebSocket proxy 通用 headers 黑名单过滤透传 -需要在 `sandbox_proxy_service.websocket_proxy()` 内新增一层 header 提取和过滤逻辑，将客户端握手里的”通用请求头”转为上游二跳握手参数。 +在 `rock/sandbox/utils/proxy.py` 中新增独立模块，负责从客户端握手中提取和过滤 header。 **设计要点**： @@ -127,7 +128,9 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST" - `Origin` 不作为普通 `additional_headers` 重复透传，避免语义混乱和重复 header 2. **其余头走黑名单过滤** - - 黑名单：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`） + - 黑名单分两类： + - WebSocket 握手专用头：`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol` + - Hop-by-hop 头：`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length` - 不在黑名单中的头默认转发，确保用户自定义 header 能到达下游 3. **无可转发头时保持兼容** @@ -186,18 +189,52 @@ async with websockets.connect( ... ``` -### 变更 5：测试覆盖扩展 +### 变更 5：VNC WebSocket Proxy 关闭 header 转发 -现有测试主要覆盖子协议转发，需要补充 header 透传相关单测。 +`websocket_proxy()` 新增 `forward_ws_headers: bool = True` 参数。当 `False` 时，跳过 `build_upstream_ws_headers()`，`origin` 和 `additional_headers` 均为 `None`。 -**新增测试点**： -- `Origin` 存在时，`websockets.connect()` 收到相同 `origin=` -- 已知头（`Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent`、`EagleEye-*`）存在时，`websockets.connect()` 收到 `additional_headers` -- 用户自定义头（如 `x-my-custom`）能被正常转发到下游 -- 黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 等）不得出现在 `additional_headers` +**背景**：QEMU 内置 WebSocket 服务器的握手 header buffer 仅 4KB（`QIO_CHANNEL_WEBSOCK_MAX_HEADER_SIZE`），上游网关注入的大量 header 会导致超限返回 502。VNC 服务不需要客户端的 Origin、Cookie、认证头等上下文。 + +```python +async def websocket_proxy( + self, + client_websocket, + sandbox_id: str, + target_path: str | None = None, + port: int | None = None, + forward_ws_headers: bool = True, +): + ... + if forward_ws_headers: + origin, additional_headers = build_upstream_ws_headers(client_websocket) + else: + origin, additional_headers = None, None +``` + +VNC 路由调用： + +```python +await sandbox_proxy_service.websocket_proxy( + websocket, sandbox_id, path, port=8006, forward_ws_headers=False +) +``` + +### 变更 6：测试覆盖扩展 + +测试文件：`tests/unit/sandbox/test_websocket_proxy_headers.py` + +**单元测试（`TestBuildUpstreamWsHeaders`）**： +- `Origin` 存在时，返回相同 `origin` +- 已知头（`Authorization`、`Traceparent`、`EagleEye-*`）存在时，出现在 `additional_headers` +- 用户自定义头（如 `x-my-custom`）能被正常转发 +- 黑名单头（`Host`、`Connection`、`Upgrade` 等握手专用头和 hop-by-hop 头）不得出现在 `additional_headers` - `Sec-WebSocket-Protocol` 继续通过 `subprotocols=` 转发，不能出现在 `additional_headers` - 无可转发头时，`origin` / `additional_headers` 为 `None`，保持向后兼容 +**E2E 测试（`TestWebSocketHeaderForwardingE2E`）**： +- 启动真实 WebSocket server，验证 header 到达下游 +- 覆盖 `Origin` 透传、已知 header 透传、自定义 header 透传、黑名单 header 过滤、向后兼容 + --- ## Execution Plan @@ -221,13 +258,20 @@ async with websockets.connect( ### Step 4：新增 WebSocket 通用 header 黑名单过滤与透传逻辑 - 文件：`rock/sandbox/utils/proxy.py`（独立模块）、`rock/sandbox/service/sandbox_proxy_service.py`（调用方） -- 新增 `build_upstream_ws_headers()` helper，负责从 `client_websocket.headers` 中提取 `Origin` 并通过黑名单过滤 `additional_headers` +- 新增 `BLOCKED_WS_HEADER_NAMES` 黑名单集合和 `build_upstream_ws_headers()` helper +- 黑名单包含：握手专用头、hop-by-hop 头 - 在 `websocket_proxy()` 调用 `websockets.connect()` 时传入 `origin=` 和 `additional_headers=` - 保持现有 `subprotocols=` 协商逻辑不变 -### Step 5：补充 WebSocket header 透传测试 +### Step 5：VNC WebSocket Proxy 关闭 header 转发 +- 文件：`rock/sandbox/service/sandbox_proxy_service.py`、`rock/admin/entrypoints/sandbox_proxy_api.py` +- `websocket_proxy()` 新增 `forward_ws_headers: bool = True` 参数 +- VNC 路由传入 `forward_ws_headers=False`，避免 QEMU 4KB header buffer 超限导致 502 + +### Step 6：补充 WebSocket header 透传测试 - 文件：`tests/unit/sandbox/test_websocket_proxy_headers.py` -- 新增 `Origin` 透传、已知 header 透传、自定义 header 透传、黑名单 header 过滤、兼容性测试 +- 单元测试：`Origin` 透传、已知 header 透传、自定义 header 透传、黑名单 header 过滤、兼容性测试 +- E2E 测试：启动真实 WebSocket server 验证 header 到达下游 --- @@ -246,6 +290,7 @@ async with websockets.connect( - WebSocket proxy 自定义端口时，`path` 参数不生效（rocklet portforward 是纯 TCP 隧道，不感知 HTTP path） - rocklet `/http_proxy` 端点的 `port` 参数需要校验（复用 `validate_port_forward_port`） - rocklet 镜像需要重新发布才能生效 -- WebSocket header 透传采用黑名单策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header 透传 +- WebSocket header 透传采用黑名单策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header（包括 `Cookie`、`Accept-Encoding`、`X-Forwarded-*`、`X-Real-IP`）透传 - `Origin` 应通过 `websockets.connect(origin=...)` 传入；不要与 `additional_headers` 重复 - `Sec-WebSocket-Protocol` 必须继续通过 `subprotocols=` 传递，避免和普通 header 透传逻辑冲突 +- VNC WebSocket 路由必须设置 `forward_ws_headers=False`，因为 QEMU 内置 WebSocket 服务器仅支持 4KB header buffer（`QIO_CHANNEL_WEBSOCK_MAX_HEADER_SIZE`） diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index f4f61e253d..808a2b0624 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -238,7 +238,7 @@ async def vnc_websocket_proxy( ): logger.info(f"Client connected to VNC WebSocket proxy: {sandbox_id}, path: {path}") try: - await sandbox_proxy_service.websocket_proxy(websocket, sandbox_id, path, port=8006) + await sandbox_proxy_service.websocket_proxy(websocket, sandbox_id, path, port=8006, forward_ws_headers=False) except WebSocketDisconnect: logger.info(f"Client disconnected from VNC WebSocket proxy: {sandbox_id}") except Exception as e: diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index f4db7a1c36..bcb138c6a2 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -203,13 +203,22 @@ async def list_sandboxes(self, query_params: SandboxQueryParams) -> SandboxListR raise async def websocket_proxy( - self, client_websocket, sandbox_id: str, target_path: str | None = None, port: int | None = None + self, + client_websocket, + sandbox_id: str, + target_path: str | None = None, + port: int | None = None, + forward_ws_headers: bool = True, ): target_url = await self.get_sandbox_websocket_url(sandbox_id, target_path, port=port) client_subprotocols = getattr(client_websocket, "subprotocols", []) or [] upstream_subprotocols = client_subprotocols if client_subprotocols else ["binary", "base64"] - origin, additional_headers = build_upstream_ws_headers(client_websocket) + if forward_ws_headers: + origin, additional_headers = build_upstream_ws_headers(client_websocket) + logger.info(f"origin for upstream WebSocket: {origin}, additional_headers: {additional_headers}") + else: + origin, additional_headers = None, None try: async with websockets.connect( From a57af7c995a7387d9bb18e21b72855c094a1b597 Mon Sep 17 00:00:00 2001 From: dengsheng Date: Thu, 23 Apr 2026 07:45:12 +0000 Subject: [PATCH 073/226] update version --- .../version-1.6.x/Release Notes/v1.6.1.md | 15 +++++++++++++++ .../version-1.6.x/Release Notes/v1.6.1.md | 15 +++++++++++++++ pyproject.toml | 2 +- 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.1.md create mode 100644 docs/versioned_docs/version-1.6.x/Release Notes/v1.6.1.md diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.1.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.1.md new file mode 100644 index 0000000000..8dc905e976 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/v1.6.1.md @@ -0,0 +1,15 @@ +# v1.6.1 + +## 发布日期 + +2026 年 4 月 23 日 + +--- + +## Admin + +### Bug 修复 + +#### Proxy接口转发支持非JSON格式body + +* 修复了Proxy接口转发时丢失非JSON格式body的问题 ([#880](https://github.com/alibaba/ROCK/pull/880)) diff --git a/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.1.md b/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.1.md new file mode 100644 index 0000000000..3cf5abb237 --- /dev/null +++ b/docs/versioned_docs/version-1.6.x/Release Notes/v1.6.1.md @@ -0,0 +1,15 @@ +# v1.6.1 + +## Release Date + +April 23, 2026 + +--- + +## Admin + +### Bug Fixes + +#### preserve non-JSON request body in http_proxy endpoint + +* Fix that http_proxy endpoint would lose non-JSON request body. Proxy now supports forwarding non-JSON body request such as application/x-www-form-urlencoded ([#880](https://github.com/alibaba/ROCK/pull/880)) diff --git a/pyproject.toml b/pyproject.toml index bdb710358a..93e479798c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.6.0" +version = "1.6.1" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From c8d78bd6a4f80416c92754db807fe577d3345ace Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 24 Apr 2026 14:56:47 +0800 Subject: [PATCH 074/226] add 1.7.0 release notes Co-Authored-By: Claude Opus 4.6 --- docs/docusaurus.config.js | 2 +- .../version-1.6.x/Release Notes/index.md | 1 + .../version-1.7.x.json | 34 ++ .../Getting Started/installation.md | 141 +++++++++ .../Getting Started/quickstart.md | 172 ++++++++++ .../Getting Started/rock-agent.md | 73 +++++ .../version-1.7.x/Getting Started/rockroll.md | 200 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 69 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 137 ++++++++ .../Python SDK References/sandbox.md | 113 +++++++ .../swe-bench-evaluation.md | 228 ++++++++++++++ .../version-1.7.x/References/api.md | 194 ++++++++++++ .../version-1.7.x/Release Notes/index.md | 5 + .../version-1.7.x/Release Notes/v1.7.0.md | 107 +++++++ .../User Guides/configuration.md | 188 +++++++++++ .../version-1.7.x/overview.md | 40 +++ .../version-1.6.x/Release Notes/index.md | 1 + .../Getting Started/installation.md | 143 +++++++++ .../Getting Started/quickstart.md | 166 ++++++++++ .../Getting Started/rock-agent.md | 72 +++++ .../version-1.7.x/Getting Started/rockroll.md | 194 ++++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../Python SDK References/model-service.md | 298 ++++++++++++++++++ .../Python SDK References/python_sdk.md | 265 ++++++++++++++++ .../Python SDK References/remote_user.md | 70 ++++ .../Python SDK References/rock-agent.md | 290 +++++++++++++++++ .../Python SDK References/runtime-env.md | 136 ++++++++ .../Python SDK References/sandbox.md | 114 +++++++ .../swe-bench-evaluation.md | 229 ++++++++++++++ .../version-1.7.x/References/api.md | 195 ++++++++++++ .../version-1.7.x/Release Notes/index.md | 5 + .../version-1.7.x/Release Notes/v1.7.0.md | 99 ++++++ .../User Guides/configuration.md | 189 +++++++++++ docs/versioned_docs/version-1.7.x/overview.md | 33 ++ .../version-1.7.x-sidebars.json | 64 ++++ docs/versions.json | 1 + 44 files changed, 5630 insertions(+), 1 deletion(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x.json create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/installation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/quickstart.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rockroll.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/codes.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/deploy.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/file_system.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/model-service.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/python_sdk.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/remote_user.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/runtime-env.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/sandbox.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/api.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/index.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/configuration.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/overview.md create mode 100644 docs/versioned_docs/version-1.7.x/Getting Started/installation.md create mode 100644 docs/versioned_docs/version-1.7.x/Getting Started/quickstart.md create mode 100644 docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md create mode 100644 docs/versioned_docs/version-1.7.x/Getting Started/rockroll.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/codes.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/deploy.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/file_system.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/model-service.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/python_sdk.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/remote_user.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/runtime-env.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/sandbox.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/versioned_docs/version-1.7.x/References/api.md create mode 100644 docs/versioned_docs/version-1.7.x/Release Notes/index.md create mode 100644 docs/versioned_docs/version-1.7.x/Release Notes/v1.7.0.md create mode 100644 docs/versioned_docs/version-1.7.x/User Guides/configuration.md create mode 100644 docs/versioned_docs/version-1.7.x/overview.md create mode 100644 docs/versioned_sidebars/version-1.7.x-sidebars.json diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index 107ff7a6ef..b57114a0ae 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -144,7 +144,7 @@ const config = { // release note按照版本号倒排 return reverseReleaseNoteSidebars(filterHiddenSidebars); }, - lastVersion: '1.6.x', + lastVersion: '1.7.x', includeCurrentVersion: false, versions: convertVersionsArrayToObject(versions) }, diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md index 3343474e09..c4091d3d8b 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.6.x/Release Notes/index.md @@ -2,4 +2,5 @@ sidebar_position: 1 --- # 版本说明 +* [release v1.6.1](v1.6.1.md) * [release v1.6.0](v1.6.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x.json b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x.json new file mode 100644 index 0000000000..c18672f115 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x.json @@ -0,0 +1,34 @@ +{ + "version.label": { + "message": "1.7.x", + "description": "The label for version 1.7.x" + }, + "sidebar.tutorialSidebar.category.Getting Started": { + "message": "快速上手", + "description": "The label for category 'Getting Started' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.User Guides": { + "message": "用户指南", + "description": "The label for category 'User Guides' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.References": { + "message": "参考", + "description": "The label for category 'References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Release Notes": { + "message": "版本说明", + "description": "The label for category 'Release Notes' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.model-service": { + "message": "Model Service 参考", + "description": "The label for category 'Model Service References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.sandbox-agent": { + "message": "Sandbox Agent参考", + "description": "The label for category 'Sandbox Agent References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Python SDK References": { + "message": "Python SDK 参考", + "description": "The label for category 'Python SDK References' in sidebar 'tutorialSidebar'" + } +} diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/installation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/installation.md new file mode 100644 index 0000000000..0ab70e55d1 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/installation.md @@ -0,0 +1,141 @@ +--- +sidebar_position: 3 +--- + +# 安装指南 + +本文档介绍如何使用 `uv` 和 `pip` 安装和设置 ROCK 开发环境。该项目是一个强化学习开放构建工具包，支持多种组件。 + +## 使用 uv（推荐） + +### 快速安装所有依赖 + +```bash +# 安装所有依赖（包括可选依赖） +uv sync --all-extras + +# 安装开发/测试依赖 +uv sync --all-extras --all-groups +``` + +### 安装不同依赖组 + +#### 仅核心依赖 +```bash +uv sync +``` + +#### 管理组件依赖 +```bash +uv sync --extra admin +``` + +#### Rocklet 执行环境依赖 +```bash +uv sync --extra rocklet +``` + +#### 所有依赖 +```bash +uv sync --all-extras +``` + +#### 开发/测试依赖 +```bash +uv sync --all-extras --group test +``` + +## 使用 pip + +### 从 pip 源安装 + +#### 仅核心依赖 +```bash +pip install rl-rock +``` + +#### 管理组件依赖 +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install "rl-rock[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install "rl-rock[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install "rl-rock[all]" +``` + +### 使用 pip 从源码安装 + +#### 仅核心依赖 +```bash +pip install . +``` + +#### 管理组件依赖 +```bash +pip install ".[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install ".[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install ".[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install ".[all]" +``` + +## 可用入口点 + +该包提供以下命令行脚本： + +- `rocklet`: ROCK 执行环境服务器 (rock.rocklet.server:main) +- `admin`: 管理服务器 (rock.admin.main:main) +- `envhub`: 环境中心服务器 (rock.envhub.server:main) +- `rock`: 主 ROCK 命令行接口 (rock.cli.main:main) + +## 开发设置 + +### 使用 uv（推荐） + +```bash +# 克隆并设置开发环境 +git clone +cd ROCK +uv sync --all-extras --group test + +# 运行测试 +uv run pytest + + +### 使用 pip + +```bash +# 开发模式安装所有可选依赖 +pip install -e ".[all]" + +# 分别安装 +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" +``` + +## 附加说明 + +- 项目配置为默认使用阿里云 PyPI 镜像: `https://mirrors.aliyun.com/pypi/simple/` +- 对于本地开发，运行测试需要 `test` 依赖组 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/quickstart.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..e0a0891c0e --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/quickstart.md @@ -0,0 +1,172 @@ +--- +sidebar_position: 2 +--- + +# 快速上手 + +本指南将通过完整的示例演示如何使用 ROCK 创建和管理强化学习环境。ROCK (Reinforcement Open Construction Kit) 是一个全面的沙箱环境管理框架，主要用于强化学习和AI开发环境。 + +## 1. 环境准备 + +我们推荐在 Linux 系统下启动 ROCK，能够尽量复用项目依赖，提升环境拉起速度。如果需要在 macOS 上尝试，可以参考 [MacOS 启动](#7-macos-启动) 一节。 + +在开始之前，请确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且示例中依赖python:3.11镜像 +docker pull python:3.11 + +# 验证 uv 安装 +uv --version + + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone +cd ROCK + +# 创建虚拟环境（使用 uv 托管的 Python, 以python 3.11 版本为例） +uv venv --python 3.11 --python-preference only-managed + +# 安装所有依赖组 +uv sync --all-extras +``` + +> **重要提示**: 为确保 ROCK 能正确挂载项目和虚拟环境及其依赖的 base Python 解释器，强烈推荐使用 uv 托管的 Python 环境而非系统 Python。 + +## 2. 激活虚拟环境 + +在运行任何 ROCK 命令之前，需要先激活虚拟环境。确保 sys.base_prefix 是 uv 管理的环境，类似于 `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` 等路径。 + +```bash +# 激活虚拟环境 +source .venv/bin/activate + +# 验证 Python 环境 +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **验证要点**: 确保输出的 base prefix 路径指向 uv 管理的 Python 环境，而非系统 Python。 + +## 3. 验证环境配置 + +激活虚拟环境后，验证依赖安装是否正确： + +```bash +# 检查关键依赖 +python -c "import rock; print(\"Hello ROCK\")" +``` + + +## 4. 启动 ROCK 服务 + +激活虚拟环境后，在项目根目录下，启动 ROCK Admin 服务： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 启动 ROCK Admin 服务（本地环境） +rock admin start +``` + +服务启动后，您将看到类似以下的输出： + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **服务说明**: ROCK Admin 服务默认运行在 `http://127.0.0.1:8080`。 + +## 5. 运行示例环境 + +现在可以运行示例环境来验证安装。确保 ROCK 服务正在运行，然后打开一个新的终端窗口执行以下命令： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 运行沙箱示例 +python examples/sandbox_demo.py + +# 运行 GEM 协议示例 +python examples/sokoban_demo.py +``` + +### 5.1 示例说明 + +- **sandbox_demo.py**: 演示如何使用 ROCK 的沙箱 SDK 创建和管理容器化环境 +- **sokoban_demo.py**: 演示如何使用 ROCK 的 GEM 协议兼容接口创建强化学习环境 + +> **运行要求**: 确保 ROCK Admin 服务正在运行，因为示例需要与服务进行通信。 + +## 6. 分布式环境配置（可选） + +对于分布式多机器环境，请确保以下配置一致： + +1. 所有机器上 ROCK 和 uv 的 Python 配置使用相同的根 Python 解释器 +2. Docker 版本在所有节点上保持一致 +3. 网络配置允许各节点间正常通信 + + + +## 7. MacOS 启动 + +在 macOS 上，如果需要启动 Linux 镜像的环境，需要先设置环境变量： + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +在容器启动时，会安装对应的 uv 环境，细节可以参考 `rock/rocklet/local_files/docker_run_with_uv.sh` 脚本。 + +> **注意**: 相比 Linux 系统，macOS 上的启动速度会较慢，且比较依赖网络环境，可以根据实际情况调整脚本。ROCK_WORKER_ENV_TYPE的细节可以参考 [Configuration Guide](../User%20Guides/configuration.md). + + +## 8. 从Pip源启动 + +如果从Pip源启动Admin Server，在参照[安装指南](./installation.md)安装完成ROCK后, 需要设置额外环境变量: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(这一启动方式在容器环境启动时会从Pypi源上拉取最新的rocklet并安装, 相对启动速度比较慢, 仅推荐测试使用, 生产上依旧推荐其他的启动方式) + + +## 总结 + +恭喜！您已经成功完成了 ROCK 的快速开始指南。现在您应该能够： + +- 正确设置 ROCK 开发环境 +- 使用 uv 管理的 Python 环境 +- 启动和管理 ROCK 服务 +- 运行示例程序验证安装 +- 在分布式环境中配置 ROCK（如果需要） + +如需深入了解 ROCK 的更多功能，请参考以下文档： + +## 下一步学习 + +- [配置指南](../User%20Guides/configuration.md) - 详细了解 ROCK 的配置选项 +- [API 文档](../References/api.md) - 查看完整的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +- [安装指南](./installation.md) - 详细了解 ROCK 安装和配置 +- [概览](../overview.md) - 了解 ROCK 的设计理念 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..bd2dd69b05 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md @@ -0,0 +1,73 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent 快速启动 + +Rock Agent 是 ROCK 提供的 AI Agent 运行框架，支持在沙箱环境中运行各种类型的 Agent。 + +## 前置条件 + +- 确保有可用的ROCK服务, 如果需要本地拉起服务端, 参考[快速启动](quickstart.md) + +## 使用示例 + +ROCK 提供了两个Hello World Agent 示例，位于 `examples/agents/` 目录下： + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent 示例 +└── iflow_cli/ # IFlowCli Agent 示例 +``` + +### 运行 IFlowCli 示例 + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### 运行 ClaudeCode 示例 + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli 配置文件 + +配置文件位于 `examples/agents/iflow_cli/rock_agent_config.yaml`： + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # 填入你的 API Key + IFLOW_BASE_URL: "" # 填入你的 Base URL + IFLOW_MODEL_NAME: "" # 填入你的模型名称 +``` + +## ClaudeCode 配置文件 + +配置文件位于 `examples/agents/claude_code/rock_agent_config.yaml`： + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # 填入你的anthropic base url + ANTHROPIC_API_KEY: "" # 填入你的anthropic api key +``` + +## 相关文档 + +- [RockAgent 参考](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rockroll.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..3b53810ba3 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rockroll.md @@ -0,0 +1,200 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL 快速开始指南 + +本指南将引导您使用 ROLL (训练框架) 和 ROCK (环境管理) 来运行一个基于 Sokoban 游戏（推箱子）的强化学习训练示例。 + +## 1. 单机环境准备 + +在开始之前，请先确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **操作系统**: 推荐使用 Linux (如 Ubuntu 20.04+) +- **硬件**: 建议使用 NVIDIA GPU 并安装对应的驱动程序 +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且可提前拉取 Sokoban 游戏环境镜像，避免训练时等待 +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# 验证 uv 安装 +uv --version + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# 确保两个仓库位于同一级目录下，如下所示： +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. 启动训练流程 + +> 说明：下文均以 *torch2.6.0 + vLLM0.8.4* 为例。 + + +### 方式一： 使用虚拟环境启动（推荐） + +#### 为什么推荐这种方式？ +- 隔离性：uv 虚拟环境能确保项目依赖与系统环境隔离，避免冲突。 +- 速度快：ROCK 可以复用此虚拟环境，大大加快了后续环境的启动速度。 +- 稳定性：依赖关系更清晰，环境更易复现。 + + +```bash +# 进入 ROCK 目录 +cd ROCK + +# 使用 uv 创建并激活 Python 3.10 虚拟环境（ROLL推荐使用Python 3.10） +uv venv --python 3.10 --python-preference only-managed + +# 激活虚拟环境 +source .venv/bin/activate + +# 使用uv安装ROCK的依赖 +uv sync --all-extras + +# 若使用Python 3.10, 启动 ray 时会报错：ValueError: is not a valid Sentinel +# 原因是 ray 与 click>=8.3 版本不兼容，需要降级到 click<8.3 +# Python 3.11 不会有这个问题 +uv pip install 'click>=8.2,<8.3' + +# 切换到 ROLL 目录以安装其依赖 +cd ../ROLL + +# 设置国内 PyPI 镜像源以加速下载 +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装核心 PyTorch 组件 +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -i $PYPI_MIRROR + +# 安装transformer-engine，--no-build-isolation 避免因环境隔离导致找不到 torch +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation -i $PYPI_MIRROR + +# 安装预编译的 flash-attention，以匹配特定的 CUDA 和 PyTorch 版本 +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# 安装其余依赖 +uv pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# (可选) 安装Tensorboard，用于查看训练指标 +uv pip install tensorboard -i $PYPI_MIRROR + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### 方式二：使用系统环境启动（备选方案） + +为获得最佳兼容性，推荐使用 ROLL 官方提供的基础 Docker 镜像，因为它们已经预装了匹配的 CUDA、cuDNN 和其他基础库。 + +> [ROLL 官方镜像列表](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Installation/image_address/) + + +#### 注意 +此方式会将所有 Python 包直接安装到您的当前环境（例如，容器的基础环境）中，可能会与系统自带的包或其他项目产生冲突。 + +由于 ROCK 无法复用环境，每次启动任务时都可能需要重新安装部分依赖，启动速度较慢且受网络影响。 + + +```bash +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装ROCK的依赖 +cd ROCK +pip install . -i $PYPI_MIRROR +pip install ".[admin]" -i $PYPI_MIRROR + +# 安装ROLL的依赖 +cd ../ROLL +pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# 配置ROCK用uv启动的环境变量 +export ROCK_WORKER_ENV_TYPE=uv + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +至此，您已成功启动了 Sokoban 强化学习训练流程。祝您 Rock & Roll 愉快！ + + +## 3. 多机部署 + +除了在单机上运行，您也可以将 **ROCK 服务** 和 **ROLL 训练** 部署在不同的机器上，通过网络进行通信。这是一种常见的服务化部署模式。 + +### 3.1 在机器 A 上部署 ROCK 服务 + +在一台独立的机器（或容器）上，参照[ROCK快速指南](./quickstart.md)部署并启动 ROCK 服务。 + +> **重要提示** +> 启动服务后，请记下ROCK服务的IP地址和端口，例如`http://192.168.1.10:8000`，后续步骤将需要这个地址。 + +### 3.2 在机器 B 上准备 ROLL 客户端 + +在另一台将要运行训练任务的机器上，执行以下操作。 + +1. 验证网络连通性 + +首先，使用 curl 命令检查是否能从机器 B 访问到机器 A 上的 ROCK 服务。 +```bash +# 将 : 替换为您的 ROCK 服务实际地址 +# 如果成功，会收到 ROCK 服务的响应 {"message":"hello, ROCK!"} +curl http://: +``` + +2. 准备 ROLL 环境 + +```bash +# 克隆 ROLL 仓库 +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# 安装依赖 +pip install -r requirements_torch260_vllm.txt -i https://mirrors.aliyun.com/pypi/simple/ +``` + +3. 配置 ROLL 连接地址 + +修改 ROLL 的配置文件，使其能够找到并连接到远程的 ROCK 服务。 +- 打开配置文件：examples/agentic_demo/agentic_val_sokoban_sandbox.yaml +- 找到 SokobanSandbox 下的 env_config 部分 +- 将 base_url 的值修改为您的 ROCK 服务地址 +```yaml +custom_envs: + SokobanSandbox: + env_config: + # 将这里的地址修改为您的 ROCK 服务地址 + # 例如: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. 启动训练 +配置完成后，即可在机器 B 上启动 ROLL 训练脚本。 + +```bash +# 此脚本现在会通过网络请求机器 A 上的 ROCK 服务来创建环境 +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### 进阶：分布式 ROLL 训练 + +如果您希望将 ROLL 训练任务本身进行分布式部署，可以参考 ROLL 的官方分布式部署文档。 +> [快速上手：多节点部署指南](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/codes.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..47b74166de --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +错误码定义和分类，用于错误处理和重试策略。 + +## 使用示例 + +```python +import rock + +def test_codes_values(): + """测试基本状态码值""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes 分类 + +```python +OK = 2000, "OK" +""" +成功状态码 (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +客户端错误码 (4xxx): + +这些错误表示客户端请求有问题， +SDK 会抛出异常。 +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +服务端错误码 (5xxx): + +这些错误表示服务端出现问题， +SDK 会抛出异常。 +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +命令/执行错误码 (6xxx): + +这些错误与命令执行相关，由模型处理， +SDK 不会抛出异常。 +""" +``` + +## 重试策略建议 + +- **重试触发条件**: 只有当 `INTERNAL_SERVER_ERROR` 时才需要重试 +- **其他情况的处理策略**: + - `BAD_REQUEST`: 需要检查 arun 调用逻辑是否有异常 + - `COMMAND_ERROR`: stdout 输出到 `observation.output`，stderr 输出到 `observation.failure_reason` +- `COMMAND_ERROR` 说明: 由于 bash 执行失败时，stdout/stderr 可能全部非空，建议将 observation 中 output 和 failure_reason 全部 prompt 给模型进行推理 + +## 重试示例 + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/deploy.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..b7bd2da08f --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +沙箱资源部署管理器，用于本地目录部署和模板格式化。 + +## deploy_working_dir - 部署本地目录 + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# 部署本地目录到沙箱（自动生成目标路径） +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"部署到: {target}") # 例如: /tmp/rock_workdir_abc123 + +# 部署到指定目标路径 +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - 模板变量替换 + +`format` 方法支持两种模板语法： + +- **`${variable}`** - 标准 Python 字符串模板语法 +- **`<>`** - 替代语法（内部转换为 `${variable}`） + +```python +# 使用 ${working_dir} 模板变量 +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# 结果: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# 使用 <<>> 替代语法 +cmd = deploy.format("cat <>/file.txt") +# 结果: cat /tmp/rock_workdir_abc123/file.txt + +# 结合自定义变量使用 +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# 结果: cat /tmp/rock_workdir_abc123/settings.json + +# Shell 语法保持不变 +cmd = deploy.format("echo $((3 << 2 >> 1))") +# 结果: echo $((3 << 2 >> 1)) + +# 直接访问 working_dir +if deploy.working_dir: + print(f"当前工作目录: {deploy.working_dir}") +``` + +## 多次部署 + +后续调用会覆盖之前的工作目录路径： + +```python +# 第一次部署 +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# 第二次部署（覆盖之前的路径） +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/file_system.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..741b14f5e4 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +文件系统操作接口，提供沙箱环境中的权限管理和目录上传功能。 + +## chown - 修改所有者 + +```python +from rock.actions.sandbox.request import ChownRequest + +# 创建远程用户后修改所有者 +await sandbox.remote_user.create_remote_user("deploy") + +# 获取当前目录 +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# 修改目录所有者 +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# 递归修改目录及其内容所有者 +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - 修改权限 + +```python +from rock.actions.sandbox.request import ChmodRequest + +# 创建测试目录 +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# 修改目录权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# 递归修改权限（包括子目录和文件） +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# 设置最高权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - 上传目录 + +```python +import os +from pathlib import Path + +# 准备本地目录 +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# 上传到沙箱 +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"上传成功: {result.output}") +else: + print(f"上传失败: {result.failure_reason}") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/model-service.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..ba158cf75a --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service（实验性） + +ROCK 提供的 Model Service 负责处理 AI 模型调用的通信，为代理(Agent)和训练框架(如 Roll)或实际的 LLM 推理服务之间提供通信桥梁。 + +## 与 RockAgent 集成 + +ModelService 通常由 **RockAgent** 自动管理，无需手动调用生命周期方法。只需在配置中启用即可： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # 启用 ModelService，RockAgent 会自动管理其生命周期 +) +``` + +RockAgent 会自动： +- 安装 ModelService（安装 Python 运行时环境、安装模型服务包） +- 启动/停止 ModelService +- 监控 Agent 进程 + +## 架构概述（Local 模式） + +Local 模式下，模型服务使用**文件系统**作为通信媒介，实现代理和模型间的请求-响应机制。 + +当 Agent 需要调用模型时，请求首先写入日志文件，然后由负责监听的组件处理响应。当模型生成响应后，结果将写回日志文件，并由等待的 Agent 读取。 + +## anti_call_llm - 核心 API + +`anti_call_llm()` 是 **Local 模式**下最重要的 API，用于手动触发 LLM 反调用，实现模型调用的精细控制： + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='OpenAI type response', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +**使用场景：** +- Agent 捕获到 LLM 响应后，调用此方法通知 Roll 运行时 +- 支持携带响应数据，用于错误处理或重试 +- 超时和检查间隔可配置，适应不同网络环境 + +## CLI 命令 + +如果需要通过 CLI 使用模型服务，ROCK 提供了一个 CLI 命令集，可以在沙箱中安装 ROCK 后，通过 `rock model-service` 访问： + +### start 命令 +开始模型服务进程 +```bash +rock model-service start --type [local|proxy] [选项] +``` + +参数: + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `--type` | str | `local` | 服务类型：`local` 或 `proxy` | +| `--config-file` | str | None | 配置文件路径 | +| `--host` | str | None | 服务器地址（覆盖配置） | +| `--port` | int | None | 服务器端口（覆盖配置） | +| `--proxy-base-url` | str | None | 代理基础 URL | +| `--retryable-status-codes` | str | None | 可重试状态码，逗号分隔 | +| `--request-timeout` | int | None | 请求超时秒数 | + +### watch-agent 命令 +监控代理进程，当进程退出时发送 SESSION_END 消息 +```bash +rock model-service watch-agent --pid <进程ID> +``` + +参数: +- `--pid`: 需要监控的代理进程 ID + +### stop 命令 +停止模型服务 +```bash +rock model-service stop +``` + +### anti-call-llm 命令 +反调用 LLM 接口 +```bash +rock model-service anti-call-llm --index <索引> [--response <响应>] +``` + +参数: +- `--index`: 上一个 LLM 调用的索引，从 0 开始 +- `--response`: 上一次 LLM 调用的响应（可选） + +## 文件通信协议 + +模型服务使用文件进行进程间通信，定义了特定的标记格式用于区分请求和响应： + +### 请求格式 +``` +LLM_REQUEST_START{JSON请求数据}LLM_REQUEST_END{元数据JSON} +``` + +### 响应格式 +``` +LLM_RESPONSE_START{JSON响应数据}LLM_RESPONSE_END{元数据JSON} +``` + +### 会话结束标识 +``` +SESSION_END +``` + +元数据包含时间戳和索引信息，用于保证消息顺序和处理。 + +## SDK 使用 + +### ModelServiceConfig + +模型服务配置类，位于 `rock/sdk/sandbox/model_service/base.py`： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # 服务类型 + install_cmd="pip install rock-model-service", # 安装命令 + install_timeout=300, # 安装超时（秒） + start_cmd="rock model-service start --type ${type}", # 启动命令 + stop_cmd="rock model-service stop", # 停止命令 + logging_path="/data/logs", # 日志路径 + logging_file_name="model_service.log", # 日志文件名 +) +``` + +| 配置项 | 默认值 | 说明 | +|--------|--------|------| +| `enabled` | `False` | 是否启用模型服务（RockAgent 自动管理） | +| `type` | `"local"` | 服务类型：`local` 或 `proxy` | +| `install_cmd` | - | 模型服务包安装命令 | +| `install_timeout` | `300` | 安装超时时间（秒） | +| `start_cmd` | - | 启动命令模板 | +| `stop_cmd` | - | 停止命令 | +| `logging_path` | `/data/logs` | 日志目录路径 | +| `logging_file_name` | `model_service.log` | 日志文件名 | + +### ModelService + +模型服务管理类，处理沙箱内模型服务的生命周期： + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# 通常由 RockAgent 自动管理，无需手动调用 +# 以下方法仅在需要手动控制时使用 + +# 安装模型服务 +await model_service.install() + +# 启动模型服务 +await model_service.start() + +# 监控代理进程 +await model_service.watch_agent(pid="12345") + +# 执行反调用 LLM（Local 模式核心 API） +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# 停止模型服务 +await model_service.stop() +``` + +## API 参考 + +### install() + +在沙箱中安装模型服务依赖。 + +```python +await model_service.install() +``` + +执行步骤： +1. 创建并初始化 Python 运行时环境 +2. 创建 Rock 配置文件 +3. 安装模型服务包 + +**注意：** 通常由 RockAgent 自动调用。 + +### start() + +启动模型服务。 + +```python +await model_service.start() +``` + +前提条件：必须先调用 `install()`。 + +**注意：** 通常由 RockAgent 自动调用。 + +### stop() + +停止模型服务。 + +```python +await model_service.stop() +``` + +如果服务未运行，会跳过此操作。 + +**注意：** 通常由 RockAgent 自动调用。 + +### watch_agent(pid) + +监控代理进程。 + +```python +await model_service.watch_agent(pid="12345") +``` + +当进程退出时，发送 `SESSION_END` 消息。 + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +执行反调用 LLM 操作。**这是 Local 模式下最重要的 API。** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='{"result": "..."}', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +## 配置选项 + +### 服务配置 +- `SERVICE_HOST`: 服务主机地址，默认为 `"0.0.0.0"` +- `SERVICE_PORT`: 服务端口，默认为 `8080` + +### 日志配置 +- `LOG_FILE`: 用以通信的日志文件路径，包含请求和响应数据 + +### 轨迹（Traj）日志记录 +模型服务将 LLM 调用轨迹（traj）记录到 JSONL 文件中，用于调试和分析。 + +| 环境变量 | 默认值 | 说明 | +|----------|--------|------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | traj 日志文件目录 | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | 追加模式（true/false） | + +**traj 文件位置**: `{DATA_DIR}/LLMTraj.jsonl` + +**traj 文件格式**（JSONL - 每行一个 JSON 对象）： +```json +{"request": {...}, "response": {...}} +``` + +### 轮询配置 +- `POLLING_INTERVAL_SECONDS`: 轮询间隔，默认为 `0.1` 秒 +- `REQUEST_TIMEOUT`: 请求超时时间，默认为无限 + +### 标记配置 +定义了用于区分日志文件中不同类型消息的标记： +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig（服务端） + +服务端配置类定义了模型服务如何处理请求： + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `host` | str | `"0.0.0.0"` | 服务器地址 | +| `port` | int | `8080` | 服务器端口 | +| `proxy_base_url` | str \| None | `None` | 直接代理 URL | +| `proxy_rules` | dict | 见下方 | 模型名称到 URL 的映射 | +| `retryable_status_codes` | list[int] | `[429, 500]` | 可重试的 HTTP 状态码 | +| `request_timeout` | int | `120` | 请求超时时间（秒） | + +**默认 proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/python_sdk.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..c1083f29b0 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK 参考 + +本指南详细介绍如何使用 ROCK SDK 进行开发，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK SDK为开发者提供了便捷的Python接口来使用ROCK平台的功能，包括沙箱环境管理和GEM环境交互。 + +> **重要提示**: 使用 SDK 之前，请确保 ROCK Admin 服务正在运行。可以通过以下命令启动： +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 基本沙箱操作 + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 沙箱组管理 + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# 创建沙箱组配置 +config = SandboxGroupConfig( + image="python:3.11", + size=4, # 创建4个沙箱 + start_concurrency=2, # 并发启动级别为2 +) + +# 创建并启动沙箱组 +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# 批量操作 +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# 批量停止 +await sandbox_group.stop() +``` + +### 2.3 配置示例 + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 沙箱加速配置 + +ROCK 提供沙箱网络加速功能，支持配置 APT、PIP 和 GitHub 镜像源，提升受限网络环境下的包下载速度。 + +#### 支持的加速类型 + +**APT 镜像配置** + +配置 APT 包管理器镜像源，加速 Debian/Ubuntu 软件包下载。 + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# 配置 APT 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP 镜像配置** + +配置 Python 包索引镜像，加速 pip 安装。 + +```python +# HTTP 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub 加速** + +通过添加自定义 DNS 解析条目加速 GitHub 访问。 + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### 完整示例 + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """创建沙箱并配置加速""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # 配置加速（在安装包之前配置） + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # speedup 不会主动安装 PIP，仅配置镜像源进行加速 + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # 可以通过镜像 IP 加速 GitHub 访问 + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### 注意事项 + +1. **配置顺序**: 在安装包之前配置加速 +2. **HTTPS vs HTTP**: HTTPS 镜像不需要为 PIP 配置 trusted-host +3. **GitHub IP**: 不同区域可能需要不同的 IP 以获得最佳性能 +4. **持久性**: 配置在沙箱生命周期内持久有效 +5. **多次调用**: 后续的加速调用会覆盖之前的配置 +6. **PIP 安装**: speedup 功能仅配置镜像源，不会自动安装 PIP + +## 3. GEM SDK + +### 3.1 Python SDK 方式 + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## 相关文档 +- [快速开始指南](../../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK SDK +- [API 文档](../api.md) - 查看 SDK 封装的底层 API 接口 +- [配置指南](../../User%20Guides/configuration.md) - 了解 SDK 相关的配置选项 +- [安装指南](../../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/remote_user.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..791ca85fdd --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/remote_user.md @@ -0,0 +1,69 @@ +# Remote User + +远程用户管理，用于在沙箱中创建和管理用户。 + +## 使用示例 + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +创建远程用户。 + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +检查用户是否存在。 + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..c3f03b1efc --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent（实验性） + +RockAgent 是 ROCK 框架中的核心 Agent 实现，直接继承自 `Agent` 抽象基类。它提供了完整的 Agent 生命周期管理，包括环境初始化、ModelService 集成、命令执行等功能。 + +使用 `sandbox.agent.install()` 以及 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 + +## 核心概念 + +RockAgent 的核心工作流程分为两个阶段： + +1. **install(config)**: 初始化 Agent 环境，包括部署工作目录、设置环境变量、初始化运行时环境等 +2. **run(prompt)**: 执行 Agent 任务，替换占位符并启动 Agent 进程 + +## 快速开始 + +### Claude Code 示例 + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli 示例 + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} 必须 + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # 环境变量 + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent 示例 + +```yaml +working_dir: "." # 上传包含 langgraph_agent.py 的本地当前目录到 sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # 运行本地脚本 + +runtime_env_config: + type: python + pip: # 安装 pip 依赖 + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## 配置详解 + +### 基础配置 + +```yaml +agent_type: "default" # Agent 类型标识（默认: "default"） +agent_name: "demo-agent" # Agent 实例名称（默认: 随机 uuid） +version: "1.0.0" # 版本标识（默认: "default"） +instance_id: "instance-001" # 实例 ID（默认: "instance-id-<随机uuid>"） +agent_installed_dir: "/tmp/installed_agent" # Agent 安装目录（默认: "/tmp/installed_agent"） +agent_session: "my-session" # bash 会话标识（默认: "agent-session-<随机uuid>"） +env: # 环境变量（默认: {}） + OPENAI_API_KEY: "xxxxxxx" +``` + +### 工作目录配置 + +```yaml +working_dir: "./my_project" # 本地目录，上传到 sandbox（默认: None 不上传） +project_path: "/testbed" # sandbox 中工作目录，用于 cd（默认: None） +use_deploy_working_dir_as_fallback: true # project_path 为空时是否回退到 deploy.working_dir（默认: true） +``` + +### 执行配置 + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent 执行命令，必须包含 ${prompt}（默认: None） + +skip_wrap_run_cmd: false # 跳过为 run_cmd 添加 PATH 的包装（默认: false） + +# 超时配置 +agent_install_timeout: 600 # 安装超时，单位秒（默认: 600） +agent_run_timeout: 1800 # 运行超时，单位秒（默认: 1800） +agent_run_check_interval: 30 # 检查间隔，单位秒（默认: 30） +``` + +**`skip_wrap_run_cmd`**: +- `false`（默认）：为命令添加 `export PATH=:$PATH &&` 包装，确保使用运行时环境的可执行文件 +- `true`：跳过 PATH 包装，直接使用 `bash -c` 运行命令 + +### 初始化钩子 + +```yaml +pre_init_cmds: # 初始化前执行的命令（默认: 从 env_vars 读取） + - command: "apt update && apt install -y git" + timeout_seconds: 300 # 命令超时，单位秒（默认: 300） + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # 初始化后执行的命令（默认: []） + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**注意事项**： +- `pre_init_cmds` 和 `post_init_cmds` 不继承 Agent 的 `env` 环境变量 +- 通常用于执行安装操作和配置文件移动操作 +- 常用命令示例： + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv 配置 + +```yaml +runtime_env_config: # 具体参考 RuntimeEnv 有关文档 + type: "python" # 运行时类型: python / node（默认: "python"） + version: "3.11" # 版本号 + pip: # Python 依赖包列表 + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node 运行时示例**： + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**自动执行的操作**： +- 根据 `type` 安装对应的运行时（Python 或 Node.js） +- 安装 `pip` 依赖（如果配置了） +- 执行 `custom_install_cmd` 自定义安装命令（如果配置了） +- 支持 `npm_registry` 配置 Node.js 的 npm 镜像源 + +### ModelService 配置 + +```yaml +model_service_config: # 具体参考 ModelService 有关文档 + enabled: true # 启用 ModelService（默认: false） +``` + +**自动执行的操作**： +- 安装阶段：安装 ModelService（仅安装，不启动） +- 运行阶段：启动 ModelService + `watch_agent` 监控进程 + +**注意事项**：需要将模型请求的 URL 设置为 ModelService 的 URL。例如 ModelService 提供的 OpenAI-compatible 的 URL 为 `http://127.0.0.1:8080/v1/chat/completions`，则通常需要将 Agent 向 LLM 请求的 URL 设置为 `http://127.0.0.1:8080/v1/`。 + +## API 参考 + +### install(config) + +初始化 Agent 环境。 + +**执行流程**： +1. 如果配置了 `working_dir`，部署到 sandbox +2. 设置 bash session，以及配置 env 环境变量 +3. 执行 `pre_init_cmds` +4. 并行初始化 RuntimeEnv 和 ModelService（如果启用） +5. 执行 `post_init_cmds` + +**参数**： +- `config`: Agent 配置文件，支持两种传入方式： + - **字符串路径**: YAML 配置文件路径，默认值为 `"rock_agent_config.yaml"` + - **RockAgentConfig 对象**: 直接传入 `RockAgentConfig` 实例 + +### run(prompt) + +执行 Agent 任务。 + +**执行流程**： +1. 替换占位符, 准备Agent 运行命令 +4. 启动 agent 进程 +5. 如果启用 ModelService，启动 `watch_agent` +6. 等待任务完成并返回结果 + +## 高级用法 + +### working_dir 与 project_path 的区别与联动 + +| 配置项 | 作用 | 联动方式 | +|--------|------|----------| +| `working_dir` | 本地目录，上传到 sandbox | 调用 `deploy.deploy_working_dir()` 上传，上传后 `deploy.working_dir` 变为 sandbox 中的路径 | +| `${working_dir}` | 命令中的占位符 | 被 `deploy.format()` 替换为 `deploy.working_dir` 的值，会在配置中的 init_cmds 和 run_cmd 中替换 | +| `project_path` | sandbox 中的工作目录 | 用于运行前 `cd project_path`，不设置时会进入到 `deploy.working_dir` 工作目录 | +| `use_deploy_working_dir_as_fallback` | run 时 project_path 未设置时是否回退到 deploy.working_dir | 默认为 `true`，设为 `false` 时即使未设置 project_path 也不会进入 working_dir | + +**使用建议**： +- 使用 `working_dir` 上传本地项目代码到 sandbox +- 使用 `project_path` 指定 sandbox 中的工作目录（如 `/testbed`） +- 设置 `use_deploy_working_dir_as_fallback: false` 的场景：需要进行本地文件挂载，但希望在镜像默认工作目录下运行 Agent + +### 占位符使用 + +Rock Agent 在支持在配置文件中替换以下占位符： + +- `${prompt}`: 在run_cmd 中必需，会被替换为 `run(prompt)` 传入的提示词 +- `${working_dir}`: 可选，会被替换为 sandbox 中实际的工作目录路径, 同时支持在 init_cmds和 run_cmd 中使用 +- `${bin_dir}`: 可选，会被替换为运行时环境的 bin 目录路径 + +**示例**： +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback 说明 + +当 `project_path` 未设置时： +- `true`（默认）：运行 Agent 前会自动 `cd` 到 `deploy.working_dir` +- `false`：运行 Agent 前不会自动切换目录，保持在当前目录 + +适用场景： +- `true`: 大多数场景，希望 Agent 在上传的代码目录中运行 +- `false`: 需要挂载本地文件，但希望在镜像默认工作目录（如 `/app, /testbed`）下运行 Agent + +## 完整配置示例 + +```yaml +# ========== 基础配置 ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== 工作目录配置 ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== 运行配置 ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# 超时配置 +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== 初始化命令 ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== 运行时环境配置 ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService 集成 ========== +model_service_config: + enabled: true +``` + +## 使用示例 + +### 使用 YAML 配置文件（推荐） + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/runtime-env.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..a5532e900b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,137 @@ +# RuntimeEnv + +RuntimeEnv 模块用于在沙箱中管理语言运行时环境（目前提供了 Python / Node.js）。 + +## 快速开始(使用示例) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +异步工厂方法，根据配置创建 RuntimeEnv 实例并初始化，自动注册到 `sandbox.runtime_envs`。 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# 自动注册，可通过 sandbox.runtime_envs[env.runtime_env_id] 访问 +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +包装命令，将 `bin_dir` 加入 PATH，确保优先使用运行时环境中的可执行文件。 + +```python +wrapped = env.wrapped_cmd("node script.js") +# 返回: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +在运行时环境中执行命令。内部基于 `wrapped_cmd` 实现 + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | 类型标识 | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python 版本，默认 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | pip 包列表或 requirements.txt 路径 | +| `pip_index_url` | `str \| None` | 环境变量 | pip 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | 要创建符号链接的可执行文件列表 | + +## NodeRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | 类型标识 | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node 版本，默认 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | 要创建符号链接的可执行文件列表 | + +## 自定义 RuntimeEnv 实现约束 + +自定义 RuntimeEnv 需遵循以下规则： + +1. **定义 `runtime_env_type` 类属性**：作为类型标识符，用于自动注册到 RuntimeEnv 工厂 +2. **重写 `_get_install_cmd()`**：返回安装命令 +3. **安装命令最后必须**：将目录重命名为 `runtime-env` + + +## NodeRuntimeEnv 简化版实现示例 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config 类：定义配置类型，用于 RuntimeEnv.create() 路由到对应实现 +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # 必须与 runtime_env_type 一致 + +# RuntimeEnv 实现类：定义如何安装和运行该运行时环境 +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # 自动注册到 RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # 下载 Node 二进制包并解压，最后重命名为 runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## 加速基础环境安装 + +`PythonRuntimeEnv` 默认从 https://github.com/astral-sh/python-build-standalone/releases/ 下载 Python 安装包。若网络不可达或下载较慢，可通过环境变量 `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` 或 `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` 覆盖默认安装命令（例如切换到内网源/镜像源）。 + +默认值示例： + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +例如，替换为镜像源下载： + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +请确保该命令执行完成后，会在 `runtime_env` 的默认工作目录下生成 `runtime-env` 目录，并且 `${workdir}/runtime-env/bin/` 下包含对应可执行文件，例如： + +- `${workdir}/runtime-env/bin/python` + +Node 环境同理，可通过修改环境变量 `ROCK_RTENV_NODE_V22180_INSTALL_CMD` 来指定更快的下载/安装命令。 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/sandbox.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..088f1e3110 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/sandbox.md @@ -0,0 +1,113 @@ +# 处理大文件和长命令输出 + +## `arun` +`arun()` 在 `nohup` 模式下提供了两个关键参数，帮助 Agent / 调用方在"执行"与"查看"之间按需解耦： + +1. **`response_limited_bytes_in_nohup`**（int 型） + 限制返回内容的最大字符数（例如 `64 * 1024`），适合仍需立刻查看部分日志、但必须控制带宽的场景。默认值 `None` 表示不加限制。 + +2. **`ignore_output`**（bool，默认 `False`） + 当设为 `True` 时，`arun()` 不再读取 nohup 输出文件，而是在命令执行完毕后立即返回一段提示信息（包含输出文件路径、**文件大小**及查看方式）。日志仍写入 `/tmp/tmp_.out`，后续可通过 `read_file`、下载接口或自定义命令按需读取，实现"执行"与"查看"彻底解耦。返回的文件大小信息可帮助用户决定是直接下载还是分块读取。 + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# 示例 1：限制最多 1024 个字符 +resp_limit = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# 示例 2：完全跳过日志读取，后续再通过 read_file / 下载获取 +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# 可通过 Sandbox.read_file(...) / 下载接口 / cat /tmp/tmp_xxx.out 查看日志 +``` + +## `read_file_by_line_range` + +按行范围异步读取文件内容，支持自动分块读取和会话管理，支持大文件读取。 + +### 重要特性 +- **大文件分块读取**: 自动将大文件分成多个小块进行读取 +- **自动统计行数**: 未指定结束行时，自动计算文件总行数 +- **内置重试机制**: 关键操作支持最多 3 次重试，提高可靠性 +- **参数验证**: 自动验证输入参数的合法性 +- **会话管理**: 支持指定会话或自动创建临时会话 + +### 参数说明 +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `file_path` | str | - | 要读取的文件路径（沙箱中的绝对路径或相对路径） | +| `start_line` | int \| None | 1 | 起始行号（从 1 开始） | +| `end_line` | int \| None | None | 结束行号（包含），默认为文件末尾 | +| `lines_per_request` | int | 1000 | 每次请求读取的行数，范围 1-10000 | + +### 返回值 +- `ReadFileResponse`: 包含文件内容的响应对象 + - `content` (str): 读取的文件内容 + +### 异常说明 +- `Exception`: 当 `start_line < 1` 时抛出 +- `Exception`: 当 `end_line < start_line` 时抛出 +- `Exception`: 当 `lines_per_request` 不在 1-10000 范围内时抛出 +- `Exception`: 当文件读取失败时抛出 + +### 使用示例 + +```python +# 读取整个文件 +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# 读取指定行范围（第 100 到 500 行） +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# 从第 1990 行读取到文件末尾 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# 使用自定义分块大小 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### 注意事项 +- 行号从 1 开始计数，而非 0 +- 对于大文件建议适当增加 `lines_per_request` 以提高效率 +- 文件路径必须是沙箱内的有效路径 +- 使用 `sed` 命令进行文件读取，确保沙箱镜像支持该命令 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..6a74f397c7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,228 @@ +# SWE-Bench 评测 + +本文档介绍如何使用 ROCK SDK 运行 SWE-Bench Verified 评测，包括沙箱启动、Agent 集成、测试环境准备和结果解析。 + +### 快速开始 +SWE-Bench-Verified 是一个用于评估 AI 编程 Agent 在真实软件工程任务上表现的基准测试。 + +在ROCK上运行一个SWE-Bench任务包含以下步骤： + +1. **load_task_config** — 加载 `task.yaml` 获取任务指令 +2. **start_sandbox** — 使用任务专属的 Docker 镜像启动沙箱 +3. **agent.install / agent.run** — 安装并运行 Agent 来解决任务 +4. **setup_test_env** — 上传测试文件和运行测试脚本到沙箱 +5. **运行测试** — 通过 `sandbox.arun()` 执行测试脚本，支持超时控制 +6. **parse_swebench_result** — 解析测试输出，判断 PASSED / FAILED +7. **sandbox.stop** — 清理沙箱资源 + +**下面是示例代码** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. 加载任务指令 + task_config = await load_task_config(task_dir) # 参见 load_task_config 章节 + instruction = task_config["instruction"] + + # 2. 启动沙箱 + sandbox = await start_sandbox(task_name) # 参见 start_sandbox 章节 + + try: + # 3. 安装并运行 Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. 准备测试环境 + await setup_test_env(sandbox, task_dir) # 参见 setup_test_env 章节 + + # 5. 运行测试 + resp = await run_tests(sandbox) # 参见"运行测试"章节 + + # 6. 解析结果 + is_resolved = parse_swebench_result(resp.output) # 参见 parse_swebench_result 章节 + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +以下章节详细介绍评测流程中使用的各个函数。 + +--- + +## start_sandbox + +使用任务专属的 SWE-Bench Docker 镜像启动沙箱实例。每个任务都有一个预构建的镜像，包含目标仓库和运行环境。 + +`image` 参数格式如下： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +例如，任务 `django__django-14539` 对应的镜像为： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +从任务目录中加载 `task.yaml` 配置文件。YAML 文件包含 `instruction` 字段，用于描述 Agent 需要完成的编程任务。 + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# 使用示例 +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +使用 `sandbox.agent.install()` 和 `sandbox.agent.run()` 在沙箱中部署和执行 Agent。详细的 Agent 配置请参考 [Rock Agent](./rock-agent.md)。 + +```python +# 使用 YAML 配置文件安装 Agent（以 iflow_config.yaml 为例） +await sandbox.agent.install(config="iflow_config.yaml") + +# 使用任务指令运行 Agent +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +在沙箱中准备测试环境：安装 [uv](https://github.com/astral-sh/uv) 包管理器，并上传测试文件和运行测试脚本。 + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """准备测试环境并返回会话名称。""" + # 1. 创建带有自定义环境变量的会话 + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. 安装 uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. 上传测试文件 + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. 上传运行测试脚本 + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## 运行测试 + +使用 `RunMode.NOHUP` 模式执行测试脚本，支持可配置的超时时间。 + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +解析测试输出以判断 SWE-Bench 任务是否通过。解析器会查找由标记行分隔的结果块，并检查是否包含 `PASSED`。 + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """解析 SWE-Bench 测试输出，判断任务是否通过。 + + 匹配 'SWEBench results starts here' 和 + 'SWEBench results ends here' 之间的内容块， + 然后检查其中是否包含 'PASSED'。 + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# 使用示例 +is_resolved = parse_swebench_result(resp.output) +``` + +## 注意事项 + +- **任务数据集**：任务目录（包含 `task.yaml`、`tests/` 和 `run-tests.sh`）可从 [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) 仓库获取。 +- **任务镜像**：每个 SWE-Bench 任务需要特定的 Docker 镜像（如 `sweb.eval.x86_64.`）。请确保镜像在对应的环境中可用。 +- **Agent 配置**：Agent 配置 YAML 定义了运行时、依赖和执行命令。详情请参考 [Rock Agent](./rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/api.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/api.md new file mode 100644 index 0000000000..06f44c326c --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/api.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 1 +--- + +# API 参考 + +本指南详细介绍 ROCK 平台提供的核心 API 服务，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK平台提供两种核心API服务： +- Sandbox API：沙箱环境管理 +- GEM API：GEM环境交互 + +所有 API 接口都遵循 RESTful 设计原则，支持 JSON 格式的数据交换。 + +## 2. Sandbox API + +沙箱环境全生命周期管理功能： + +### 沙箱管理接口 + +1. **Start Sandbox** - 启动沙箱环境 + - 创建一个新的沙箱实例 + - 支持指定镜像、资源配置等参数 + +2. **Start Sandbox Async** - 异步启动沙箱环境 + - 异步方式创建沙箱实例 + - 适用于需要快速响应的场景 + +3. **Check Sandbox Alive Status** - 检查沙箱存活状态 + - 验证沙箱是否正常运行 + +4. **Get Sandbox Statistics** - 获取沙箱统计信息 + - 获取沙箱的资源使用统计 + +5. **Get Sandbox Status** - 获取沙箱详细状态 + - 获取沙箱的完整状态信息 + +6. **Stop Sandbox** - 停止沙箱环境 + - 安全关闭沙箱实例 + +7. **Commit Sandbox** - 提交沙箱为镜像 + - 将当前沙箱状态保存为新镜像 + +### 命令执行接口 + +8. **Execute Command** - 在沙箱中执行命令 + - 直接在沙箱中运行指定命令 + +9. **Create Bash Session** - 创建Bash会话 + - 创建持久化的Bash会话环境 + +10. **Run Command in Session** - 在会话中执行命令 + - 在已创建的会话中执行命令 + +11. **Close Session** - 关闭会话 + - 释放会话资源 + +### 文件操作接口 + +12. **Read File** - 读取沙箱文件 + - 从沙箱中读取指定文件内容 + +13. **Write File** - 写入沙箱文件 + - 向沙箱中写入文件 + +14. **Upload File** - 上传文件到沙箱 + - 将本地文件上传到沙箱 + +## 3. GEM API + +GEM环境交互功能： + +1. **Make Environment** - 创建GEM环境 + - 初始化一个新的GEM环境实例 + +2. **Reset Environment** - 重置GEM环境 + - 将GEM环境重置到初始状态 + +3. **Step Environment** - 执行GEM环境步骤 + - 在GEM环境中执行一个动作步骤 + +4. **Close Environment** - 关闭GEM环境 + - 释放GEM环境资源 + +## 4. HTTP API 使用示例 + +### 4.1 Sandbox API 示例 + +#### 启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 异步启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### 创建会话 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### 在会话中执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### 上传文件 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### 停止沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API 示例 + +```bash +# 创建GEM环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# 重置环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# 执行步骤 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# 关闭环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK API +- [Python SDK 文档](./Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 调用 API +- [配置指南](../User%20Guides/configuration.md) - 了解 API 相关的配置选项 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/index.md new file mode 100644 index 0000000000..7befd70ebb --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# 版本说明 +* [release v1.7.0](v1.7.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md new file mode 100644 index 0000000000..e5f084cc80 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md @@ -0,0 +1,107 @@ +# Zh Release Notes + +# v1.7.0 + +## 发布日期 + +2026 年 4 月 23 日 + +--- + +## 亮点 + +本次发布引入了 **Datasets SDK 与 CLI**，用于管理 OSS 上的基准测试数据集；新增**容器根文件系统磁盘限额**（基于 Docker storage-opt）；重构了 **Docker 认证方案**，采用沙箱级临时目录隔离凭证。WebSocket 代理的头部转发策略从白名单切换为黑名单以支持自定义头部透传，沙箱容器通过挂载宿主机 zoneinfo 获得完整的 IANA 时区支持。Admin 服务新增 **MetaStore 与数据库操作指标**用于可观测性。 + +--- + +## Datasets + +### 新功能 + +#### Datasets SDK 与 CLI + +* **新增**: 新增 `rock datasets` CLI，包含三个子命令：`list`（按组织/名称/split 浏览数据集）、`tasks`（枚举指定 split 中的 task ID，支持 `--offset`/`--limit` 分页）、`upload`（批量上传本地 task 目录至 OSS，支持并发度配置与 `--overwrite` 覆写） ([#859](https://github.com/alibaba/ROCK/pull/859), [#875](https://github.com/alibaba/ROCK/pull/875)) + +* **新增**: `OssDatasetRegistry` 后端，基于 `oss2` SDK 导航 `datasets/{org}/{dataset}/{split}/{task_id}/` 键结构，同时支持目录风格与扁平文件风格的 task。OSS 凭证可通过 CLI 参数传入或存储在 `config.ini` 的 `[dataset]` 节中 ([#859](https://github.com/alibaba/ROCK/pull/859)) + +* Task 列表现在同时识别目录 task（来自 `prefix_list`）和文件 task（来自 `object_list`），自动去除后缀并去重 ([#875](https://github.com/alibaba/ROCK/pull/875)) + + +--- + +## Sandbox + +### 新功能 + +#### 容器磁盘限额 + +* 支持通过 Docker `--storage-opt size=` 限制沙箱根目录大小。要求 `overlay2` 存储驱动 + XFS 文件系统 + `prjquota` 挂载选项。通过 `xfs_quota` 对沙箱日志目录设置 XFS project 配额，要求 XFS 文件系统 + `prjquota` 挂载选项。 ([#860](https://github.com/alibaba/ROCK/pull/860)) + +* 服务端 `RuntimeConfig` 通过新增的 `disk_limit_rootfs` 与 `disk_limit_log`（均默认为 `None`）字段可在 `rock-{args.env}.yml` 中按环境配置沙箱根目录和沙箱日志目录的限额，并支持 Nacos 运行时覆盖 ([#860](https://github.com/alibaba/ROCK/pull/860)) + + +#### 沙箱时区支持 + +* **新增**: 容器现在获得完整的 IANA 时区支持 — ROCK 将宿主机的 zoneinfo 文件（如 `/usr/share/zoneinfo/Asia/Shanghai`）以只读方式挂载到容器的 `/etc/localtime`。宿主机缺少对应 zoneinfo 文件时自动跳过并输出警告 ([#883](https://github.com/alibaba/ROCK/pull/883)) + + +--- + +## Admin + +### MetaStore 与数据库操作指标 + +* **新增**: 为 `SandboxMetaStore` 和 `SandboxTable` 的 CRUD 操作添加 OpenTelemetry 指标埋点。每个操作（create、get、update、delete、list、batch_get、archive 等）均自动采集 total/success/failure 计数器及响应时间 gauge ([#887](https://github.com/alibaba/ROCK/pull/887)) + + +## Deployments + +### Docker 认证重构 + +* 移除传统 Docker 认证方案，改用临时目录方案。`TempAuthDockerClient` 上下文管理器为每个沙箱创建隔离的临时目录，在其中执行 `docker --config  login` 与镜像拉取，退出时自动清理 — 避免注册表凭证持久化到全局 Docker 配置 ([#837](https://github.com/alibaba/ROCK/pull/837)) + +* 临时认证目录的基础路径可通过 `ROCK_DOCKER_TEMP_AUTH_DIR` 环境变量配置，默认使用系统临时目录 ([#837](https://github.com/alibaba/ROCK/pull/837)) + + +--- + +## Proxy + +### WebSocket Header转发 + +* /sandboxes/{id}/proxy/{path:path} 接口支持header转发，所有客户端头部默认转发至上游服务，仅过滤 WebSocket 握手头（`sec-websocket-*`）、逐跳头（`connection`、`upgrade`、`transfer-encoding`、`content-length`）及 `host`。`origin` 头部做特殊处理 — 单独提取并作为 WebSocket origin 参数传入，VNC WebSocket 接口关闭该功能，避免超出 QEMU 的 4 KB 头部缓冲区限制([#865](https://github.com/alibaba/ROCK/pull/865)) + + +--- + +## Bug 修复 + +* 修复 `auto_clear_time` 计算问题：`auto_clear_seconds / 60` 产生的小数分钟数现通过 `math.ceil()` 向上取整至至少 1 分钟；同时将进程存活轮询的 `wait_interval` 限制为小于 `auto_clear_seconds`，防止沙箱在首次存活检测前被自动清理 ([#883](https://github.com/alibaba/ROCK/pull/883)) + +* 修复 UV 环境构建：项目文件树现在先复制到可写的 `/tmp/rock-build` 目录再执行 `uv pip install`，解决容器内源码只读挂载导致的安装失败 ([#857](https://github.com/alibaba/ROCK/pull/857)) + + +--- + +## 测试与 CI + +* 新增 `DockerUtil` 辅助函数的单元测试（`detect_storage_opt_support`、`is_xfs_path`、`get_docker_root_dir`）及 `DockerDeployment` 磁盘限额集成测试 ([#860](https://github.com/alibaba/ROCK/pull/860)) + +* 新增 Docker 临时目录认证方案的集成测试 ([#837](https://github.com/alibaba/ROCK/pull/837)) + +* 新增 Datasets CLI、客户端、模型及 OSS 注册表的完整单元测试 ([#859](https://github.com/alibaba/ROCK/pull/859), [#875](https://github.com/alibaba/ROCK/pull/875)) + +* 数据库连接单元测试与参数优化 ([#852](https://github.com/alibaba/ROCK/pull/852)) + +* 清理 TS SDK model client 测试中泄漏的定时器 ([#839](https://github.com/alibaba/ROCK/pull/839)) + + +--- + +## 迁移说明 + +* **Docker 认证**: 传统 Docker 认证方案已移除。如果此前依赖 ROCK 写入 `~/.docker/config.json` 中的持久凭证，请注意 ROCK 现在使用临时目录方案。可通过 `ROCK_DOCKER_TEMP_AUTH_DIR` 自定义临时目录位置。 + +* **WebSocket 代理头部**: 如果您的下游服务依赖特定的转发头部集合，请注意现在所有非屏蔽头部均会被转发。屏蔽集合包括 `host`、`connection`、`upgrade`、`sec-websocket-*`、`transfer-encoding` 和 `content-length`。 + +* **磁盘限额**: 磁盘配额为服务端策略 — `disk_limit` 不在 `SandboxStartRequest` 中暴露。请通过环境 YAML 或 Nacos 覆盖配置 `RuntimeConfig.disk_limit_rootfs` / `disk_limit_log`。 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/configuration.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/configuration.md new file mode 100644 index 0000000000..a212189bc7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/configuration.md @@ -0,0 +1,188 @@ +--- +sidebar_position: 4 +--- + +# 配置指南 + +本指南详细介绍如何配置 ROCK 环境以满足不同的使用需求，包括本地开发、测试和生产部署。 + +## 1. 环境变量配置 + +ROCK 支持通过环境变量配置关键参数。以下是主要的环境变量： + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK服务基础URL +export ROCK_LOG_LEVEL=INFO # 日志级别 +export ROCK_LOGGING_PATH=/path/to/logs # 日志文件路径，默认 None (输出到控制台) +export ROCK_LOGGING_FILE_NAME=rocklet.log # 日志文件名，默认 "rocklet.log", 启动admin时可以自定义日志文件名, 如admin.log +export ROCK_LOGGING_LEVEL=INFO # 日志输出级别，默认 "INFO" +export ROCK_WORKER_ENV_TYPE=local # 运行时环境类型，可选值: local, docker, uv, pip +``` + +更多环境变量可参考 `rock/env_vars.py` 文件。 + +### 1.1 运行时环境 (Runtime Environments) + +ROCK 提供了多种不同的运行时环境来满足不同场景的需求，选择通过环境变量 `ROCK_WORKER_ENV_TYPE` 进行配置。每种环境有不同的部署要求、性能特征和适用场景。每种环境都有其独特的优势和限制，开发者可以根据部署环境的需要选择最适合的运行时环境。 + +#### 1.1.1 Docker 运行时环境 + +Docker 运行时环境适用于已经预安装了所需依赖的 Docker 镜像环境。这种环境要求部署环境中直接可用 `/tmp/miniforge/bin/rocklet` 可执行文件。 + +**挂载配置:** +- `/tmp/miniforge` - 包含预安装的 Python 环境 +- `/tmp/local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 容器化部署环境 +- 已经构建了包含 `rocklet` 的自定义 Docker 镜像 +- 适合生产环境，启动速度快 + +**要求:** +- 需要使用定制的 Docker 镜像，其中包含 `/tmp/miniforge/bin/rocklet` 可执行文件 +- Docker 环境支持 + +#### 1.1.2 本地运行时环境 + +本地运行时环境直接利用当前部署环境的 Python 环境和项目文件。该环境要求宿主机和容器之间具有相同的操作系统，以便能够直接挂载虚拟环境和 Python 解释器。 + +**挂载配置:** +- `python_env_path` - Python 环境路径 +- `project_root` - 项目根目录 +- `.venv` - 虚拟环境目录（挂载为容器中的 `/tmp/miniforge`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 开发环境 +- 宿主机和目标容器使用相同操作系统的场景 +- 需要快速重新使用现有 Python 环境 + +**要求:** +- 相同的操作系统（主机/容器） +- 可直接访问当前部署的 `.venv` 虚拟环境 +- Python 解释器路径兼容 + +#### 1.1.3 UV 运行时环境 + +UV 运行时环境只依赖于可用的 ROCK 项目，但初始化相对较慢且网络要求较高。这种环境最适合没有预配置环境的场景。它从原始项目重新构建 rocklet 环境。这是推荐在 Mac 操作系统上使用的环境。 + +**挂载配置:** +- `project_root` - 项目根目录（挂载为容器中的 `/tmp + project_root`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**适用场景:** +- Mac 操作系统 +- 跨操作系统启动 +- 没有预配置环境的场景 +- 没有使用 uv 管理 Rock + +**优势:** +- 无需预构建镜像 +- 跨平台兼容性好 +- 特别适合开发和测试 + +**限制:** +- 初始化速度较慢 +- 网络要求较高 +- 启动时间较长 + +#### 1.1.4 PIP 运行时环境 + +PIP 运行时环境使用 pip 在容器内安装所需依赖。这种环境适合快速设置并能在容器中完成依赖安装的场景，是默认的运行时环境。它不需要预先构建包含依赖的镜像，通过 pip 直接管理 Python 包。 + +**挂载配置:** +- `local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**适用场景:** +- 使用PIP源安装的ROCK +- 快速测试ROCK + +**优势:** +- 简单的部署设置 + +**限制:** +- 依赖安装时间较长 +- 需要网络访问以安装依赖包 +- 每次启动时都需要安装依赖 + +#### 1.1.5 配置指南 + +根据不同的使用场景，可以参考以下选择指南： + +| 场景 | 推荐环境 | 原因 | +|------|----------|------| +| 生产环境 | Docker 运行时 | 快速启动，稳定性能 | +| 开发环境，同一 OS | 本地运行时 | 环境重用，开发周期快 | +| Mac 开发 | UV 运行时 | 支持最佳的跨平台兼容性 | +| 跨平台开发 | UV 运行时 | 避免环境兼容性问题 | +| 快速测试 | UV 运行时 | 无需预配置工作 | +| PIP源安装 | PIP 运行时 | 直接使用 pip 安装依赖 | + +这些运行时环境通过 `ROCK_WORKER_ENV_TYPE` 环境变量进行配置，该变量可设置为 "local"、"docker"、"uv" 或 "pip"。 + +### 1.2 日志配置 + +在日志配置方面，ROCK 的日志系统具有以下特性： + +- 日志系统不能同时输出到文件和控制台，只有当设置了 `ROCK_LOGGING_PATH` 时，日志才会输出到指定文件，否则输出到控制台。 +- `ROCK_LOGGING_LEVEL` 用于控制日志输出级别，`ROCK_LOG_LEVEL` 用于通用日志级别设置。 + +## 2. 分布式部署要求 + +由于 ROCK 支持分布式部署，当在 Ray 集群的不同节点上运行时，需要满足以下一致性要求： + +#### 目录结构一致性 +在所有 Ray 节点上，必须保证以下目录结构完全一致： +- ROCK 项目仓库目录 +- `.venv` 虚拟环境目录 +- `.venv` 依赖的 base Python 目录 + + +#### 挂载要求 +ROCK 的启动依赖于挂载 ROCK 项目和对应的 base Python 环境，要求在多机环境中保持一致性： + +#### 验证分布式配置 +可以通过以下方式验证分布式部署配置： + +```bash +# 在所有节点上检查目录一致性 +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# 验证 Python 环境可用性 +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# 检查所有节点上的环境变量设置 +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + + + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速搭建 ROCK 环境 +- [API 文档](../References/api.md) - 查看沙箱相关的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 配置沙箱 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/overview.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/overview.md new file mode 100644 index 0000000000..a02536f50d --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/overview.md @@ -0,0 +1,40 @@ +--- +sidebar_position: 1 +--- + +# 概览 + +ROCK (Reinforcement Open Construction Kit) 是一个开源的强化学习环境开发框架，旨在简化强化学习环境的开发、部署和管理流程。 + +## 什么是 ROCK + +ROCK (Reinforcement Open Construction Kit) 是一个开源强化学习环境开发框架。通过使用 ROCK，开发者可以快速地开发强化学习环境，并结合其他强化学习训练框架，实现高效的强化学习训练。 + +ROCK 提供了完整的沙箱环境管理功能，支持容器化部署，能够实现环境的快速创建、运行和销毁。同时，ROCK 兼容 GEM 协议，为强化学习环境提供了标准化的接口。 + +## ROCK 的核心功能 + +1. **简化开发流程**：简化强化学习环境的开发、构建和管理流程，支持多种开源的强化学习环境 +2. **大规模调度部署**：支持快速强化学习环境的大规模调度部署，通过 GEM 协议可以方便地访问强化学习环境 +3. **框架集成**：与其他强化学习训练框架集成，实现大规模可扩展的强化学习训练 + +## ROCK 的价值 + +ROCK 为不同角色的工程师提供了显著价值： + +- **强化学习算法工程师**：ROCK 可以简化强化学习环境的开发流程，让工程师专注于算法实现 +- **强化学习应用工程师**：ROCK 可以进行快速强化学习环境的大规模部署，提高应用开发效率 + +## 相关文档 + +如果您是第一次使用 ROCK，建议按以下顺序阅读文档： +1. [快速开始指南](./Getting%20Started/quickstart.md) - 快速搭建开发环境 +2. [配置指南](./User%20Guides/configuration.md) - 配置您的 ROCK 环境 +3. [Python SDK 文档](./References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +4. [API 文档](./References/api.md) - 了解完整的 API 接口 +5. [安装指南](./Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 + + + + + diff --git a/docs/versioned_docs/version-1.6.x/Release Notes/index.md b/docs/versioned_docs/version-1.6.x/Release Notes/index.md index 38262ea2a6..6859885cdc 100644 --- a/docs/versioned_docs/version-1.6.x/Release Notes/index.md +++ b/docs/versioned_docs/version-1.6.x/Release Notes/index.md @@ -2,4 +2,5 @@ sidebar_position: 1 --- # Release Notes +* [release v1.6.1](v1.6.1.md) * [release v1.6.0](v1.6.0.md) diff --git a/docs/versioned_docs/version-1.7.x/Getting Started/installation.md b/docs/versioned_docs/version-1.7.x/Getting Started/installation.md new file mode 100644 index 0000000000..c45b09a8fe --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Getting Started/installation.md @@ -0,0 +1,143 @@ +--- +sidebar_position: 3 +--- + +# Installation + +This document explains how to install and set up the ROCK development environment using both `uv` and `pip`. The project is a Reinforcement Open Construction Kit that supports various components. + +## Using uv (Recommended) + +### Quick Install All Dependencies + +```bash +# Install all dependencies including optional ones +uv sync --all-extras + +# Install development/testing dependencies +uv sync --all-extras --all-groups +``` + +### Install Different Dependency Groups + +#### Core Dependencies Only +```bash +uv sync +``` + +#### Admin Component Dependencies +```bash +uv sync --extra admin +``` + +#### Rocklet Execution Environment Dependencies +```bash +uv sync --extra rocklet +``` + + +#### All Dependencies at Once +```bash +uv sync --all-extras +``` + +#### Development/Testing Dependencies +```bash +uv sync --all-extras --group test +``` + +## Using pip + +### Install from pip source + +#### Core Dependencies Only +```bash +pip install rl-rock +``` + +#### Admin Component Dependencies +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install "rl-rock[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install "rl-rock[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install "rl-rock[all]" +``` + +### Install with pip from source code + +#### Core Dependencies Only +```bash +pip install . +``` + +#### Admin Component Dependencies +```bash +pip install ".[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install ".[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install ".[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install ".[all]" +``` + +## Available Entry Points + +The package provides the following command line scripts: + +- `rocklet`: ROCK execution environment server (rock.rocklet.server:main) +- `admin`: Admin management server (rock.admin.main:main) +- `envhub`: Environment hub server (rock.envhub.server:main) +- `rock`: Main ROCK command line interface (rock.cli.main:main) + +## Development Setup + +### Using uv (Recommended) + +```bash +# Clone and set up development environment +git clone +cd ROCK +uv sync --all-extras --group test + +# Run tests +uv run pytest + +``` + +### Using pip + +```bash +# For development, install in editable mode with all extras +pip install -e ".[all]" + +# Or separately +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" # Optional extras +``` + +## Additional Notes + +- The project is configured to use the Alibaba cloud PyPI mirror by default: `https://mirrors.aliyun.com/pypi/simple/` +- For local development, running tests requires the `test` dependency group diff --git a/docs/versioned_docs/version-1.7.x/Getting Started/quickstart.md b/docs/versioned_docs/version-1.7.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..2f14808f5b --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Getting Started/quickstart.md @@ -0,0 +1,166 @@ +--- +sidebar_position: 2 +--- + +# Getting Started + +This guide will demonstrate how to use ROCK to create and manage reinforcement learning environments through complete examples. + +## 1. Environment Preparation + +We recommend starting ROCK on Linux systems to maximize dependency reuse and improve environment startup speed. If you need to try on macOS, please refer to the [MacOS Startup](#7-macos-startup) section. + +Before starting, please ensure your system has the following dependencies installed: + +### 1.1 System Requirements + +- **Docker**: ROCK uses Docker for containerized environment management +- **uv**: ROCK uses uv for dependency management and virtual environment creation + +### 1.2 Verify Dependency Installation + +```bash +# Verify Docker installation +docker --version + +# Verify Docker image, and example depends on python:3.11 image +docker pull python:3.11 + +# Verify uv installation +uv --version +``` + +### 1.3 Project Initialization + +```bash +# Clone repository +git clone +cd ROCK + +# Create virtual environment (using uv-managed Python, use python 3.11 as an example) +uv venv --python 3.11 --python-preference only-managed + +# Install all dependency groups +uv sync --all-extras +``` + +> **Important Note**: To ensure ROCK can correctly mount the project and virtual environment along with its base Python interpreter, it is strongly recommended to use uv-managed Python environments to create virtual environments rather than system Python. + +## 2. Activate Virtual Environment + +Before running any ROCK commands, you need to activate the virtual environment. Ensure sys.base_prefix is a uv-managed environment, such as `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` or similar paths. + +```bash +# Activate virtual environment +source .venv/bin/activate + +# Verify Python environment +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **Verification Point**: Ensure the output base prefix path points to a uv-managed Python environment, not system Python. + +## 3. Verify Environment Configuration + +After activating the virtual environment, verify that dependencies are installed correctly: + +```bash +# Check key dependencies +python -c "import rock; print(\"Hello ROCK\")" +``` + +## 4. Start ROCK Service + +After activating the virtual environment, start the ROCK Admin service on project root: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Start ROCK Admin service (local environment) +rock admin start +``` + +After the service starts, you will see output similar to the following: + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **Service Information**: The ROCK Admin service runs by default on `http://127.0.0.1:8080`. + +## 5. Run Example Environments + +Now you can run example environments to verify the installation. Ensure the ROCK service is running, then open a new terminal window to execute the following commands: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Run sandbox example +python examples/sandbox_demo.py + +# Run GEM protocol example +python examples/sokoban_demo.py +``` + +### 5.1 Example Descriptions + +- **sandbox_demo.py**: Demonstrates how to use ROCK's sandbox SDK to create and manage containerized environments +- **sokoban_demo.py**: Demonstrates how to use ROCK's GEM protocol compatible interface to create reinforcement learning environments + +> **Running Requirements**: Ensure the ROCK Admin service is running, as examples need to communicate with the service. + +## 6. Distributed Environment Configuration (Optional) + +For distributed multi-machine environments, ensure the following configurations are consistent: + +1. All machines use the same root Python interpreter for ROCK and uv Python configurations +2. Docker versions are consistent across all nodes +3. Network configuration allows normal communication between nodes + + +## 7. MacOS Startup + +On macOS, if you need to start Linux image environments, you first need to set the environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +During container startup, the corresponding uv environment will be installed. For details, please refer to the `rock/rocklet/local_files/docker_run_with_uv.sh` script. + +> **Note**: Compared to Linux systems, the startup speed on macOS will be slower and more dependent on network conditions. You can adjust the script according to actual conditions.You can find detatils for ROCK_WORKER_ENV_TYPE in [Configuration Guide](../User%20Guides/configuration.md). + +## 8. Starting from Pip Source + +If starting the Admin Server from Pip source, after completing the ROCK installation by referring to [installation](./installation.md), you need to set an additional environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(This startup method will pull and install the latest rocklet from the PyPI source when starting the container environment. The startup speed is relatively slow, so it is only recommended for testing purposes. For production environments, other startup methods are still recommended.) + +## Summary + +Congratulations! You have successfully completed the ROCK quick start guide. You should now be able to: + +- Properly set up the ROCK development environment +- Use uv-managed Python environments +- Start and manage ROCK services +- Run example programs to verify installation +- Configure ROCK in distributed environments (if needed) + +For a deeper understanding of ROCK's additional features, please refer to the following documents: + +## Next Steps + +- [Configuration Guide](../User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](../References/api.md) - View complete API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the Python SDK for development +- [Installation Guide](./installation.md) - Detailed information about ROCK installation and setup +- [Overview](../overview.md) - Understand ROCK's design philosophy \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md b/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..eb6b54a65b --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md @@ -0,0 +1,72 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent Quick Start + +Rock Agent is an AI Agent runtime framework provided by ROCK, supporting various types of Agents running in sandbox environments. + +## Prerequisites +- Make sure you have a working ROCK service, if you need to locally start the service side, refer to [Quick Start](quickstart.md). + +## Examples + +ROCK provides two Hello World Agent examples in the `examples/agents/` directory: + +``` +examples/agents/ +├── claude_code/ # ClaudeCode Agent example +└── iflow_cli/ # IFlowCli Agent example +``` + +### Run IFlowCli Example + +```bash +cd examples/agents/iflow_cli +python iflow_cli_demo.py +``` + +### Run ClaudeCode Example + +```bash +cd examples/agents/claude_code +python claude_code_demo.py +``` + +## IFlowCli Configuration File + +The configuration file is located at `examples/agents/iflow_cli/rock_agent_config.yaml`: + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" + +runtime_env_config: + type: node + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: + IFLOW_API_KEY: "" # Enter your API key + IFLOW_BASE_URL: "" # Enter your base URL + IFLOW_MODEL_NAME: "" # Enter your model name +``` + +## ClaudeCode Configuration File + +The configuration file is located at `examples/agents/claude_code/rock_agent_config.yaml`: + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" # Enter your anthropic base url + ANTHROPIC_API_KEY: "" # Enter your anthropic api key +``` + +## Related Documentation + +- [RockAgent Reference](../References/Python%20SDK%20References/rock-agent.md) diff --git a/docs/versioned_docs/version-1.7.x/Getting Started/rockroll.md b/docs/versioned_docs/version-1.7.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..2465a7733f --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Getting Started/rockroll.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL Quick Start Guide + +This guide will walk you through running a reinforcement learning training example based on the Sokoban game, using ROLL (the training framework) and ROCK (the environment management tool). + +## 1. Prerequisites + +Before you begin, please ensure your system has the following dependencies installed. + +### 1.1 System Requirements + +- **OS**: A Linux-based system is recommended (e.g., Ubuntu 20.04+). +- **Hardware**: An NVIDIA GPU with the corresponding drivers is recommended. +- **Docker**: ROCK uses Docker for containerized environment management. +- **uv**: ROCK uses uv for dependency management and virtual environment creation. + +### 1.2 Verify Dependencies & Pre-pull Image + +```bash +# Verify Docker installation +docker --version + +# Verify Docker is running and pre-pull the Sokoban environment image +# This will save time when the training starts. +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# Verify uv installation +uv --version + +``` + +### 1.3 Initialize the Project + +```bash +# Clone the project repositories +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# Ensure both repositories are in the same parent directory, like this: +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. Launch the Training Process + +> Note: The following instructions use torch==2.6.0 and vLLM==0.8.4 as an example. + + +### Option 1: Using a Virtual Environment (Recommended) + +#### Why is this method recommended? +- Isolation: A uv virtual environment ensures that project dependencies are isolated from your system, preventing conflicts. +- Fast Startup: ROCK can reuse this virtual environment, significantly speeding up subsequent task initializations. +- Stability & Reproducibility: Dependency management is cleaner and more reliable. + + +```bash +# Navigate to the ROCK directory +cd ROCK + +# Create and activate a Python 3.10 virtual environment (ROLL recommends Python 3.10) +uv venv --python 3.10 --python-preference only-managed +source .venv/bin/activate + +# Install all of ROCK's dependencies using uv +uv sync --all-extras + +# If using Python 3.10, starting Ray may raise a `ValueError: is not a valid Sentinel`. +# This is due to an incompatibility between `ray` and `click` versions 8.3+. +# To fix this, downgrade `click` to a version below 8.3. This issue does not affect Python 3.11. +uv pip install 'click>=8.2,click<8.3' + +# Navigate to the ROLL directory to install its dependencies +cd ../ROLL + +# Install core PyTorch components +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 + +# Install transformer-engine. The --no-build-isolation flag prevents errors where torch cannot be found. +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation + +# Install a pre-compiled version of flash-attention matching the specific CUDA and PyTorch versions +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# Install the remaining dependencies +uv pip install -r requirements_torch260_vllm.txt + +# (Optional) Install Tensorboard to check training metrics +uv pip install tensorboard -i $PYPI_MIRROR + +# All set! Launch the training script. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### Option 2: Using the System Environment (Alternative) + +For optimal compatibility with this method, we recommend running these commands inside one of ROLL's official base Docker images. These images come pre-installed with matching CUDA, cuDNN, and other foundational libraries. + +> [ROLL's Official Docker Image List](https://alibaba.github.io/ROLL/docs/Getting%20Started/Installation/image_address) + + +#### Warning +This method will install all Python packages directly into your current environment (e.g., the container's base system), which may cause conflicts with system packages or other projects. + +Since ROCK cannot reuse the environment, it may need to reinstall some dependencies each time a task starts, leading to slower startup times that are dependent on network speed. + + +```bash +# Install ROCK's dependencies +cd ROCK +pip install . +pip install ".[admin]" + +# Install ROLL's dependencies +cd ../ROLL +pip install -r requirements_torch260_vllm.txt + +# Crucial: Configure ROCK to use uv as its worker environment manager +export ROCK_WORKER_ENV_TYPE=uv + +# Launch the training script +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +You have now successfully launched the Sokoban reinforcement learning training process. Happy Rock & Roll！ + + +## 3. Multi-Node Deployment + +Instead of running everything on a single machine, you can deploy the **ROCK Service** and **ROLL job** on separate machines. This is a common client-server setup where they communicate over the network. + +### 3.1 Deploy the ROCK Service on Machine A + +On a dedicated machine (or container), follow the [ROCK Quick Start Guide](./quickstart.md) to deploy and start the ROCK service. + +> **Important** +> After starting the service, take note of its IP address and port (e.g., `http://192.168.1.10:8000`). You will need this address for the subsequent steps. + +### 3.2 Prepare the ROLL Client on Machine B + +On the other machine where you will run the training task, perform the following steps. + +1. Verify Network Connectivity + +First, use the curl command to check if you can reach the ROCK service on Machine A from Machine B. +```bash +# Replace : with the actual address of your ROCK service +# If successful, you should receive a response like {"message":"hello, ROCK!"} +curl http://: +``` + +2. Prepare the ROLL Environment + +```bash +# Clone the ROLL repository +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# Install dependencies +pip install -r requirements_torch260_vllm.txt +``` + +3. Configure the ROLL Connection Address + +Modify ROLL's configuration file to point to the remote ROCK service. +- Open the configuration file: examples/agentic_demo/agentic_val_sokoban_sandbox.yaml. +- Find the "SokobanSandbox" section under "env_config". +- Update the base_url value to your ROCK service's address. +```yaml +custom_envs: + SokobanSandbox: + env_config: + # Change the address here to your ROCK service's address + # Example: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. Start Training +Once configured, you can start the ROLL training script on Machine B. + +```bash +# This script will now request environments from the ROCK service on Machine A over the network. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### Advanced: Distributed ROLL Training + +If you wish to deploy the ROLL training task itself in a distributed manner, you can refer to ROLL's official documentation for distributed deployment. +> [Quick Start: Multi-Node Deployment Guide](https://alibaba.github.io/ROLL/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/codes.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..dceb8d3182 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +Error code definitions and categories for error handling and retry strategies. + +## Usage Example + +```python +import rock + +def test_codes_values(): + """Test basic status code values""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes Categories + +```python +OK = 2000, "OK" +""" +Success codes (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +Client error codes (4xxx): + +These errors indicate issues with the client request, +SDK will raise Exceptions for these errors. +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +Server error codes (5xxx): + +These errors indicate issues on the server side, +SDK will raise Exceptions for these errors. +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +Command/execution error codes (6xxx): + +These errors are related to command execution and should be handled by the model, +SDK will NOT raise Exceptions for these errors. +""" +``` + +## Retry Strategy Recommendations + +- **Retry trigger**: Only retry when `INTERNAL_SERVER_ERROR` occurs +- **Other error handling**: + - `BAD_REQUEST`: Check if there are issues with the arun call logic + - `COMMAND_ERROR`: stdout goes to `observation.output`, stderr goes to `observation.failure_reason` +- `COMMAND_ERROR` note: When bash execution fails, both stdout and stderr may be non-empty. It is recommended to prompt the model with both output and failure_reason from the observation. + +## Retry Example + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/deploy.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..5fd5b70546 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +Sandbox resource deployment manager for local directory deployment and template formatting. + +## deploy_working_dir - Deploy Local Directory + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# Deploy local directory (auto-generated target path) +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"Deployed to: {target}") # e.g., /tmp/rock_workdir_abc123 + +# Deploy to specific target path +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - Template Variable Substitution + +The `format` method supports two template syntaxes: + +- **`${variable}`** - Standard Python string template syntax +- **`<>`** - Alternative syntax (converted to `${variable}` internally) + +```python +# After deploy_working_dir, use ${working_dir} placeholder +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# Result: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# Alternative <<>> syntax +cmd = deploy.format("cat <>/file.txt") +# Result: cat /tmp/rock_workdir_abc123/file.txt + +# Combine with custom variables +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# Result: cat /tmp/rock_workdir_abc123/settings.json + +# Shell syntax is preserved +cmd = deploy.format("echo $((3 << 2 >> 1))") +# Result: echo $((3 << 2 >> 1)) + +# Access working_dir directly +if deploy.working_dir: + print(f"Current working directory: {deploy.working_dir}") +``` + +## Multiple Deployments + +Subsequent calls overwrite previous working directory paths: + +```python +# First deployment +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# Second deployment (overwrites previous path) +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/file_system.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..a64228e8f1 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +File system interface for sandbox environment operations including permission and ownership management. + +## chown - Change Owner + +```python +from rock.actions.sandbox.request import ChownRequest + +# Create remote user before changing ownership +await sandbox.remote_user.create_remote_user("deploy") + +# Get current working directory +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# Change directory owner +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# Recursively change owner for directory and contents +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - Change Permissions + +```python +from rock.actions.sandbox.request import ChmodRequest + +# Create test directory +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# Change directory permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# Recursively change permissions (includes subdirectories and files) +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# Set maximum permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - Upload Directory + +```python +import os +from pathlib import Path + +# Prepare local directory +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# Upload to sandbox +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"Upload success: {result.output}") +else: + print(f"Upload failed: {result.failure_reason}") +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/model-service.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..23dbc21bfc --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service (Experimental) + +The Model Service provided by ROCK is responsible for handling AI model call communications, serving as a communication bridge between agents and training frameworks (such as Roll) or actual LLM inference services. + +## RockAgent Integration + +ModelService is typically **automatically managed by RockAgent** - no manual lifecycle management is required. Simply enable it in the configuration: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # Enable ModelService, RockAgent manages its lifecycle +) +``` + +RockAgent will automatically: +- Install ModelService (install Python runtime, install model service package) +- Start/stop ModelService +- Monitor Agent process + +## Architecture Overview (Local Mode) + +In local mode, the model service uses the **file system** as the communication medium, implementing a request-response mechanism between agents and models. + +When an agent needs to call a model, the request is first written to a log file, then processed by the listening component. When the model generates a response, the result is written back to the log file and read by the waiting agent. + +## anti_call_llm - Core API + +`anti_call_llm()` is the **most important API in Local mode**, used to manually trigger LLM anti-calls for fine-grained control over model calls: + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='OpenAI type response', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +**Use cases:** +- After Agent captures LLM response, call this method to notify Roll runtime +- Supports carrying response data for error handling or retry +- Configurable timeout and check interval for different network environments + +## CLI Commands + +To use the model service via CLI, ROCK provides a set of CLI commands that can be accessed via `rock model-service` after installing ROCK in the sandbox: + +### start command +Start the model service process +```bash +rock model-service start --type [local|proxy] [options] +``` + +Parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `--type` | str | `local` | Service type: `local` or `proxy` | +| `--config-file` | str | None | Path to configuration file | +| `--host` | str | None | Server host address (overrides config) | +| `--port` | int | None | Server port (overrides config) | +| `--proxy-base-url` | str | None | Proxy base URL | +| `--retryable-status-codes` | str | None | Comma-separated list of retryable status codes | +| `--request-timeout` | int | None | Request timeout in seconds | + +### watch-agent command +Monitor the agent process and send a SESSION_END message when the process exits +```bash +rock model-service watch-agent --pid +``` + +Parameters: +- `--pid`: The ID of the agent process to monitor + +### stop command +Stop the model service +```bash +rock model-service stop +``` + +### anti-call-llm command +Anti-call the LLM interface +```bash +rock model-service anti-call-llm --index [--response ] +``` + +Parameters: +- `--index`: Index of the previous LLM call, starting from 0 +- `--response`: Response from the previous LLM call (optional) + +## File Communication Protocol + +The model service uses files for inter-process communication, defining specific marker formats to distinguish requests and responses: + +### Request Format +``` +LLM_REQUEST_START{JSON request data}LLM_REQUEST_END{metadata JSON} +``` + +### Response Format +``` +LLM_RESPONSE_START{JSON response data}LLM_RESPONSE_END{metadata JSON} +``` + +### Session End Marker +``` +SESSION_END +``` + +Metadata contains timestamp and index information to ensure message order and processing. + +## SDK Usage + +### ModelServiceConfig + +Model service configuration class, located in `rock/sdk/sandbox/model_service/base.py`: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # Service type + install_cmd="pip install rock-model-service", # Install command + install_timeout=300, # Install timeout (seconds) + start_cmd="rock model-service start --type ${type}", # Start command + stop_cmd="rock model-service stop", # Stop command + logging_path="/data/logs", # Log path + logging_file_name="model_service.log", # Log filename +) +``` + +| Config | Default | Description | +|--------|---------|-------------| +| `enabled` | `False` | Whether to enable model service (RockAgent manages) | +| `type` | `"local"` | Service type: `local` or `proxy` | +| `install_cmd` | - | Model service package install command | +| `install_timeout` | `300` | Install timeout in seconds | +| `start_cmd` | - | Start command template | +| `stop_cmd` | - | Stop command | +| `logging_path` | `/data/logs` | Log directory path | +| `logging_file_name` | `model_service.log` | Log filename | + +### ModelService + +Model service management class, handles the lifecycle of model services within the sandbox: + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# Typically auto-managed by RockAgent, no manual calls needed +# The following methods are only for manual control when needed + +# Install model service +await model_service.install() + +# Start model service +await model_service.start() + +# Monitor agent process +await model_service.watch_agent(pid="12345") + +# Execute anti-call LLM (Core API for Local mode) +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# Stop model service +await model_service.stop() +``` + +## API Reference + +### install() + +Install model service dependencies in the sandbox. + +```python +await model_service.install() +``` + +Execution steps: +1. Create and initialize Python runtime environment +2. Create Rock config file +3. Install model service package + +**Note:** Typically auto-called by RockAgent. + +### start() + +Start the model service. + +```python +await model_service.start() +``` + +Prerequisite: Must call `install()` first. + +**Note:** Typically auto-called by RockAgent. + +### stop() + +Stop the model service. + +```python +await model_service.stop() +``` + +If the service is not running, this operation will be skipped. + +**Note:** Typically auto-called by RockAgent. + +### watch_agent(pid) + +Monitor the agent process. + +```python +await model_service.watch_agent(pid="12345") +``` + +Sends `SESSION_END` message when the process exits. + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +Execute anti-call LLM operation. **This is the most important API in Local mode.** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='{"result": "..."}', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +## Configuration Options + +### Service Configuration +- `SERVICE_HOST`: Service host address, defaults to `"0.0.0.0"` +- `SERVICE_PORT`: Service port, defaults to `8080` + +### Log Configuration +- `LOG_FILE`: Log file path used for communication, containing request and response data + +### Trajectory (Traj) Logging +The model service records LLM call trajectories (traj) to a JSONL file for debugging and analysis. + +| Environment Variable | Default | Description | +|---------------------|---------|-------------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | Directory for traj log files | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | Append mode (true/false) | + +**Traj file location**: `{DATA_DIR}/LLMTraj.jsonl` + +**Traj file format** (JSONL - one JSON object per line): +```json +{"request": {...}, "response": {...}} +``` + +### Polling Configuration +- `POLLING_INTERVAL_SECONDS`: Polling interval, defaults to `0.1` seconds +- `REQUEST_TIMEOUT`: Request timeout, defaults to unlimited + +### Marker Configuration +Defines markers used to distinguish different types of messages in the log file: +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig (Server-side) + +The server-side configuration class defines how the model service handles requests: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `host` | str | `"0.0.0.0"` | Server host address | +| `port` | int | `8080` | Server port | +| `proxy_base_url` | str \| None | `None` | Direct proxy URL | +| `proxy_rules` | dict | See below | Model name to URL mapping | +| `retryable_status_codes` | list[int] | `[429, 500]` | Retryable HTTP status codes | +| `request_timeout` | int | `120` | Request timeout in seconds | + +**Default proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/python_sdk.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..5272d7edb6 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK Reference + +This guide provides detailed information on how to use the ROCK SDK for development, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +ROCK SDK provides developers with convenient Python interfaces to use ROCK platform features, including sandbox environment management and GEM environment interaction. + +> **Important Note**: Before using the SDK, ensure that the ROCK Admin service is running. You can start it with the following command: +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 Basic Sandbox Operations + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 Sandbox Group Management + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# Create sandbox group configuration +config = SandboxGroupConfig( + image="python:3.11", + size=4, # Create 4 sandboxes + start_concurrency=2, # Concurrency level for startup is 2 +) + +# Create and start sandbox group +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# Batch operations +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# Batch stop +await sandbox_group.stop() +``` + +### 2.3 Configuration Example + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 Sandbox Speedup Configuration + +ROCK provides sandbox network acceleration capabilities, supporting configuration of APT, PIP, and GitHub mirror sources to improve package download speeds in restricted network environments. + +#### Supported Speedup Types + +**APT Mirror Configuration** + +Configure APT package manager mirror sources for faster Debian/Ubuntu package downloads. + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# Configure APT mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP Mirror Configuration** + +Configure Python package index mirrors for faster pip installations. + +```python +# HTTP mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub Acceleration** + +Configure GitHub IP acceleration by adding custom DNS resolution entries. + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### Complete Example + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """Create sandbox and configure acceleration""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # Configure acceleration (before installing packages) + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # Speedup does not automatically install PIP, it only configures mirror sources for acceleration + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # GitHub can be accelerated through mirror IP + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### Important Notes + +1. **Configuration Order**: Configure speedup before installing packages +2. **HTTPS vs HTTP**: HTTPS mirrors don't require trusted-host configuration for PIP +3. **GitHub IP**: Different regions may require different IPs for optimal performance +4. **Persistence**: Configurations persist within the sandbox lifecycle +5. **Multiple Calls**: Subsequent speedup calls will override previous configurations +6. **PIP Installation**: The speedup feature only configures mirror sources and does not automatically install PIP + +## 3. GEM SDK + +### 3.1 Python SDK Approach + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## Related Documents +- [Quick Start Guide](../../Getting%20Started/quickstart.md) - Learn how to quickly get started with the ROCK SDK +- [API Documentation](../api.md) - View the underlying API interfaces encapsulated by the SDK +- [Configuration Guide](../../User%20Guides/configuration.md) - Learn about SDK-related configuration options +- [Installation Guide](../../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/remote_user.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..810bbcf028 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/remote_user.md @@ -0,0 +1,70 @@ +# Remote User + +Remote user management for creating and managing users in the sandbox. + +## Usage Examples + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +Create a remote user. + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +Check if a user exists. + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..f24ade49ac --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,290 @@ +# Rock Agent (Experimental) + +RockAgent is the core Agent implementation in the ROCK framework, directly inheriting from the `Agent` abstract base class. It provides complete Agent lifecycle management, including environment initialization, ModelService integration, command execution, and more. + +Using `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, you can install and run Agents in the Sandbox environment provided by Rock. + +## Core Concepts + +The core workflow of RockAgent is divided into two phases: + +1. **install(config)**: Initialize the Agent environment, including deploying the working directory, setting environment variables, initializing the runtime environment, etc. +2. **run(prompt)**: Execute the Agent task, replace placeholders, and start the Agent process + +## Quick Start + +### Claude Code Example + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli Example + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} is required + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # Environment variables + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent Example + +```yaml +working_dir: "." # Upload local current directory containing langgraph_agent.py to sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # Run local script + +runtime_env_config: + type: python + pip: # Install pip dependencies + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## Configuration Details + +### Basic Configuration + +```yaml +agent_type: "default" # Agent type identifier (default: "default") +agent_name: "demo-agent" # Agent instance name (default: random uuid) +version: "1.0.0" # Version identifier (default: "default") +instance_id: "instance-001" # Instance ID (default: "instance-id-") +agent_installed_dir: "/tmp/installed_agent" # Agent installation directory (default: "/tmp/installed_agent") +agent_session: "my-session" # Bash session identifier (default: "agent-session-") +env: # Environment variables (default: {}) + OPENAI_API_KEY: "xxxxxxx" +``` + +### Working Directory Configuration + +```yaml +working_dir: "./my_project" # Local directory to upload to sandbox (default: None, no upload) +project_path: "/testbed" # Working directory in sandbox for cd (default: None) +use_deploy_working_dir_as_fallback: true # Whether to fall back to deploy.working_dir when project_path is empty (default: true) +``` + +### Execution Configuration + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent execution command, must contain ${prompt} (default: None) + +skip_wrap_run_cmd: false # Skip wrapping run_cmd with PATH (default: false) + +# Timeout configuration +agent_install_timeout: 600 # Installation timeout in seconds (default: 600) +agent_run_timeout: 1800 # Run timeout in seconds (default: 1800) +agent_run_check_interval: 30 # Check interval in seconds (default: 30) +``` + +**`skip_wrap_run_cmd`**: +- `false` (default): Wraps the command with `export PATH=:$PATH &&` to ensure runtime environment executables are used +- `true`: Skips PATH wrapping, runs the command directly with `bash -c` + +### Initialization Hooks + +```yaml +pre_init_cmds: # Commands executed before initialization (default: read from env_vars) + - command: "apt update && apt install -y git" + timeout_seconds: 300 # Command timeout in seconds (default: 300) + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # Commands executed after initialization (default: []) + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**Notes**: +- `pre_init_cmds` and `post_init_cmds` do not inherit the Agent's `env` environment variables +- Typically used for installation operations and configuration file movement +- Common command examples: + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv Configuration + +```yaml +runtime_env_config: # Refer to RuntimeEnv documentation for details + type: "python" # Runtime type: python / node (default: "python") + version: "3.11" # Version number + pip: # Python dependency package list + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node Runtime Example**: + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**Automatic Operations**: +- Install corresponding runtime based on `type` (Python or Node.js) +- Install `pip` dependencies (if configured) +- Execute `custom_install_cmd` custom installation command (if configured) +- Support `npm_registry` configuration for Node.js npm mirror source + +### ModelService Configuration + +```yaml +model_service_config: # Refer to ModelService documentation for details + enabled: true # Enable ModelService (default: false) +``` + +**Automatic Operations**: +- Installation phase: Install ModelService (install only, do not start) +- Run phase: Start ModelService + `watch_agent` monitoring process + +**Notes**: You need to set the model request URL to the ModelService URL. For example, if the ModelService provides an OpenAI-compatible URL at `http://127.0.0.1:8080/v1/chat/completions`, you typically need to set the Agent's LLM request URL to `http://127.0.0.1:8080/v1/`. + +## API Reference + +### install(config) + +Initialize the Agent environment. + +**Execution Flow**: +1. If `working_dir` is configured, deploy to sandbox +2. Set up bash session and configure env environment variables +3. Execute `pre_init_cmds` +4. Initialize RuntimeEnv and ModelService in parallel (if enabled) +5. Execute `post_init_cmds` + +**Parameters**: +- `config`: Agent configuration file, supports two input methods: + - **String path**: YAML configuration file path, default value is `"rock_agent_config.yaml"` + - **RockAgentConfig object**: Directly pass a `RockAgentConfig` instance + +### run(prompt) + +Execute the Agent task. + +**Execution Flow**: +1. Replace placeholders and prepare Agent run command +2. Start the agent process +3. If ModelService is enabled, start `watch_agent` +4. Wait for task completion and return results + +## Advanced Usage + +### Difference and Interaction between working_dir and project_path + +| Configuration | Function | Interaction Method | +|--------------|----------|-------------------| +| `working_dir` | Local directory uploaded to sandbox | Calls `deploy.deploy_working_dir()` to upload, after upload `deploy.working_dir` becomes the path in sandbox | +| `${working_dir}` | Placeholder in commands | Replaced by `deploy.format()` with the value of `deploy.working_dir`, replaced in init_cmds and run_cmd in the configuration | +| `project_path` | Working directory in sandbox | Used for `cd project_path` before running, when not set it enters the `deploy.working_dir` working directory | +| `use_deploy_working_dir_as_fallback` | Whether to fall back to deploy.working_dir when project_path is not set at runtime | Default is `true`, when set to `false` it will not enter working_dir even if project_path is not set | + +**Usage Recommendations**: +- Use `working_dir` to upload local project code to sandbox +- Use `project_path` to specify the working directory in sandbox (e.g., `/testbed`) +- Set `use_deploy_working_dir_as_fallback: false` scenario: Need to perform local file mounting, but want to run Agent in the image's default working directory + +### Placeholder Usage + +Rock Agent supports replacing the following placeholders in the configuration file: + +- `${prompt}`: Required in run_cmd, will be replaced with the prompt passed to `run(prompt)` +- `${working_dir}`: Optional, will be replaced with the actual working directory path in sandbox, also supported in init_cmds and run_cmd +- `${bin_dir}`: Optional, will be replaced with the runtime environment's bin directory path + +**Example**: +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback Explanation + +When `project_path` is not set: +- `true` (default): Before running Agent, it will automatically `cd` to `deploy.working_dir` +- `false`: Before running Agent, it will not automatically switch directories, staying in the current directory + +Applicable Scenarios: +- `true`: Most scenarios, where you want Agent to run in the uploaded code directory +- `false`: Need to mount local files, but want to run Agent in the image's default working directory (e.g., `/app`, `/testbed`) + +## Complete Configuration Example + +```yaml +# ========== Basic Configuration ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== Working Directory Configuration ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== Run Configuration ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# Timeout configuration +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== Initialization Commands ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== Runtime Environment Configuration ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService Integration ========== +model_service_config: + enabled: true +``` + +## Usage Examples + +### Using YAML Configuration File (Recommended) + +```python +# prepare a rock_agent_config.yaml +await sandbox.agent.install(config="rock_agent_config.yaml") +await sandbox.agent.run(prompt="hello") +``` diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/runtime-env.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..e1996cfcdb --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,136 @@ +# RuntimeEnv + +The RuntimeEnv module is used to manage language runtime environments in the sandbox (currently providing Python / Node.js). + +## Quick Start (Example) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +An async factory method that creates and initializes a RuntimeEnv instance based on the configuration, and automatically registers it to `sandbox.runtime_envs`. + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# Auto-registered; accessible via sandbox.runtime_envs[env.runtime_env_id] +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +Wraps a command by adding `bin_dir` to PATH to ensure executables from the runtime environment are used with priority. + +```python +wrapped = env.wrapped_cmd("node script.js") +# Returns: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +Executes a command within the runtime environment. Internally implemented based on `wrapped_cmd`. + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | Type identifier | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python version; default is 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | List of pip packages or a requirements.txt path | +| `pip_index_url` | `str \| None` | Environment variable | pip index mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | List of executables to symlink | + +## NodeRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | Type identifier | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node version; default is 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm registry mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | List of executables to symlink | + +## Constraints for Custom RuntimeEnv Implementations + +A custom RuntimeEnv must follow these rules: + +1. **Define the `runtime_env_type` class attribute**: used as a type identifier for automatic registration into the RuntimeEnv factory +2. **Override `_get_install_cmd()`**: return the install command +3. **The install command must end with**: renaming the directory to `runtime-env` + +## Simplified NodeRuntimeEnv Implementation Example + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config class: defines the config type so RuntimeEnv.create() can route to the corresponding implementation +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # Must match runtime_env_type + +# RuntimeEnv implementation class: defines how to install and run this runtime environment +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # Auto-registered to RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # Download the Node binary tarball and extract it, then rename to runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## Speeding Up Base Runtime Installation + +`PythonRuntimeEnv` downloads Python packages from https://github.com/astral-sh/python-build-standalone/releases/ by default. If the network is unavailable or slow, you can override the default install command via `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` or `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` (e.g., switch to an internal registry or a mirror). + +Default value example: + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +For example, override it to download from a mirror: + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +Make sure the command creates a `runtime-env` directory under the default working directory of `runtime_env`, and that `${workdir}/runtime-env/bin/` contains the expected executables, e.g.: + +- `${workdir}/runtime-env/bin/python` + +The same applies to Node.js: you can override the install command via `ROCK_RTENV_NODE_V22180_INSTALL_CMD` to use a faster download/install method. \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/sandbox.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..e6e1e43124 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/sandbox.md @@ -0,0 +1,114 @@ +# Handling Large Files and Long Command Outputs + +## `arun` + +`arun()` provides two knobs to control how `nohup` output is handled: + +1. **`response_limited_bytes_in_nohup`** *(integer type)* + Caps the number of characters returned from the nohup output file. Useful when you still need to stream some logs back but want an upper bound (default `None` = no cap). + +2. **`ignore_output`** *(bool, default `False`)* + When set to `True`, `arun()` skips reading the nohup output file entirely. The command still runs to completion and writes logs to `/tmp/tmp_.out`, but the SDK immediately returns a lightweight hint telling agents where to fetch the logs later (via `read_file`, download APIs, or custom commands). This fully decouples "execute command" from "inspect logs". The response also includes the **file size** to help users decide whether to download directly or read in chunks. + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# Example 1: limit the returned logs to 1024 characters +resp_limited = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# Example 2: skip collecting logs; agent will download/read them later +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# Use Sandbox.read_file(...), download APIs, or run 'cat /tmp/tmp_xxx.out' ... +``` + +## `read_file_by_line_range` + +Asynchronously reads file content by line range, with built-in support for automatic chunking and session management. Supports large file reading. + +### Key Features +- **Chunked reading for large files**: Automatically splits large files into chunks +- **Automatic line count**: Estimates total lines when end_line is not specified +- **Built-in retry mechanism**: Up to 3 retries for critical operations +- **Input validation**: Validates input parameters automatically +- **Session management**: Supports custom session or auto-created temporary session + +### Parameters +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `file_path` | str | - | File path to read (absolute or relative path in sandbox) | +| `start_line` | int \| None | 1 | Starting line number (1-based) | +| `end_line` | int \| None | None | Ending line number (inclusive), defaults to file end | +| `lines_per_request` | int | 1000 | Lines per request, range 1-10000 | + +### Return Value +- `ReadFileResponse`: Response object containing file content + - `content` (str): The file content read + +### Exception Handling +- `Exception`: Raised when `start_line < 1` +- `Exception`: Raised when `end_line < start_line` +- `Exception`: Raised when `lines_per_request` is not in range 1-10000 +- `Exception`: Raised when file reading fails + +### Usage Examples + +```python +# Read the entire file +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# Read a specific line range (lines 100 to 500) +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# Read from line 1990 to the end of file +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# Use custom chunk size +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### Notes +- Line numbers are 1-based, not 0-based +- For large files, consider increasing `lines_per_request` for better efficiency +- File path must be a valid path within the sandbox +- Uses `sed` command for file reading; ensure the sandbox image supports this command diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..85f34cbe7c --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,229 @@ +# SWE-Bench Evaluation + +This guide demonstrates how to use the ROCK SDK to run SWE-Bench Verified evaluations, including sandbox setup, Agent integration, test environment preparation, and result parsing. + +### Quick Start + +SWE-Bench is a benchmark for evaluating AI coding agents on real-world software engineering tasks. + +Running a SWE-Bench task on ROCK involves the following steps: + +1. **load_task_config** — Load `task.yaml` to get the task instruction +2. **start_sandbox** — Start a sandbox with a task-specific Docker image +3. **agent.install / agent.run** — Install and run the Agent to solve the task +4. **setup_test_env** — Upload test files and run-test script to the sandbox +5. **Run tests** — Execute the test script via `sandbox.arun()` with timeout +6. **parse_swebench_result** — Parse test output to determine PASSED / FAILED +7. **sandbox.stop** — Clean up sandbox resources + +**Here is an example code** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. Load task instruction + task_config = await load_task_config(task_dir) # see load_task_config section + instruction = task_config["instruction"] + + # 2. Start sandbox + sandbox = await start_sandbox(task_name) # see start_sandbox section + + try: + # 3. Install and run Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. Setup test environment + await setup_test_env(sandbox, task_dir) # see setup_test_env section + + # 5. Run tests + resp = await run_tests(sandbox) # see Running Tests section + + # 6. Parse results + is_resolved = parse_swebench_result(resp.output) # see parse_swebench_result section + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +The following sections describe each function used in the workflow in detail. + +--- + +## start_sandbox + +Start a sandbox instance with a task-specific SWE-Bench Docker image. Each task has a pre-built image containing the target repository and environment. + +The `image` parameter follows the format: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +For example, task `django__django-14539` maps to: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +Load task configuration from a `task.yaml` file in the task directory. The YAML file contains the `instruction` field that describes the coding task for the Agent. + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# Usage +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +Use `sandbox.agent.install()` and `sandbox.agent.run()` to deploy and execute an Agent inside the sandbox. Refer to [Rock Agent](./rock-agent.md) for detailed Agent configuration. + +```python +# Install Agent with a YAML configuration file（e.g., iflow_config.yaml） +await sandbox.agent.install(config="iflow_config.yaml") + +# Run Agent with the task instruction +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +Prepare the test environment in the sandbox: install the [uv](https://github.com/astral-sh/uv) package manager, and upload test files and the run-test script. + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """Set up the test environment and return the session name.""" + # 1. Create a session with custom environment variables + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. Install uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. Upload test files + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. Upload run-tests script + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## Running Tests + +Execute the test script with a configurable timeout using `RunMode.NOHUP`. + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +Parse the test output to determine whether the SWE-Bench task is resolved. The parser looks for a result block delimited by marker lines and checks for `PASSED`. + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """Parse SWE-Bench test output to determine if the task is resolved. + + Matches the block between 'SWEBench results starts here' and + 'SWEBench results ends here', then checks whether it contains 'PASSED'. + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# Usage +is_resolved = parse_swebench_result(resp.output) +``` + +## Notes + +- **Task Datasets**: Task directories (containing `task.yaml`, `tests/`, and `run-tests.sh`) can be obtained from the [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) repository. +- **Task Images**: Each SWE-Bench task requires a specific Docker image (e.g., `sweb.eval.x86_64.`). Ensure the image is available before running tests. +- **Agent Config**: The Agent configuration YAML defines the runtime, dependencies, and execution command. See [Rock Agent](./rock-agent.md) for details. + diff --git a/docs/versioned_docs/version-1.7.x/References/api.md b/docs/versioned_docs/version-1.7.x/References/api.md new file mode 100644 index 0000000000..d73bf49d6c --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/api.md @@ -0,0 +1,195 @@ +--- +sidebar_position: 1 +--- + +# API Reference + +This guide provides detailed information about the core API services provided by the ROCK platform, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +The ROCK platform provides two core API services: +- Sandbox API: Sandbox environment management +- GEM API: GEM environment interaction + +All API interfaces follow RESTful design principles and support JSON format data exchange. + +## 2. Sandbox API + +Full lifecycle management functions for sandbox environments: + +### Sandbox Management Interfaces + +1. **Start Sandbox** - Start a sandbox environment + - Create a new sandbox instance + - Support specifying image, resource configuration and other parameters + +2. **Start Sandbox Async** - Asynchronously start a sandbox environment + - Asynchronously create a sandbox instance + - Suitable for scenarios requiring quick response + +3. **Check Sandbox Alive Status** - Check sandbox alive status + - Verify if the sandbox is running normally + +4. **Get Sandbox Statistics** - Get sandbox statistics + - Get resource usage statistics of the sandbox + +5. **Get Sandbox Status** - Get detailed sandbox status + - Get complete status information of the sandbox + +6. **Stop Sandbox** - Stop sandbox environment + - Safely shut down the sandbox instance + +7. **Commit Sandbox** - Commit sandbox as image + - Save current sandbox state as a new image + +### Command Execution Interfaces + +8. **Execute Command** - Execute command in sandbox + - Run specified command directly in the sandbox + +9. **Create Bash Session** - Create Bash session + - Create a persistent Bash session environment + +10. **Run Command in Session** - Run command in session + - Execute command in a created session + +11. **Close Session** - Close session + - Release session resources + +### File Operation Interfaces + +12. **Read File** - Read sandbox file + - Read specified file content from the sandbox + +13. **Write File** - Write sandbox file + - Write file to the sandbox + +14. **Upload File** - Upload file to sandbox + - Upload local file to the sandbox + +## 3. GEM API + +GEM environment interaction functions: + +1. **Make Environment** - Create GEM environment + - Initialize a new GEM environment instance + +2. **Reset Environment** - Reset GEM environment + - Reset GEM environment to initial state + +3. **Step Environment** - Execute GEM environment step + - Execute an action step in the GEM environment + +4. **Close Environment** - Close GEM environment + - Release GEM environment resources + + +## 4. HTTP API Usage Examples + +### 4.1 Sandbox API Examples + +#### Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Asynchronously Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Execute Command +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### Create Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### Run Command in Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### Upload File +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### Stop Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API Examples + +```bash +# Create GEM environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# Reset environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# Execute step +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# Close environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly get started with ROCK API +- [Python SDK Documentation](./Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to call APIs +- [Configuration Guide](../User%20Guides/configuration.md) - Learn about API-related configuration options +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/Release Notes/index.md b/docs/versioned_docs/version-1.7.x/Release Notes/index.md new file mode 100644 index 0000000000..a02c4e03d8 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Release Notes/index.md @@ -0,0 +1,5 @@ +--- +sidebar_position: 1 +--- +# Release Notes +* [release v1.7.0](v1.7.0.md) diff --git a/docs/versioned_docs/version-1.7.x/Release Notes/v1.7.0.md b/docs/versioned_docs/version-1.7.x/Release Notes/v1.7.0.md new file mode 100644 index 0000000000..3f4021bb9c --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/Release Notes/v1.7.0.md @@ -0,0 +1,99 @@ +# v1.7.0 + +## Release Date + +April 23, 2026 + +--- + +## Highlights + +This release introduces **Datasets SDK & CLI** for managing benchmark datasets on OSS, **container rootfs disk limits** via Docker storage-opt, and a revamped **Docker auth scheme** using per-sandbox temporary directories. WebSocket proxy header forwarding has been switched from whitelist to blacklist to support custom headers, and sandbox containers now receive full IANA timezone support via host zoneinfo bind-mounts. The admin service gains **MetaStore & database operation metrics** for observability. + +--- + +## Datasets + +### New Feature + +#### Datasets SDK & CLI + +* **NEW**: Added `rock datasets` CLI with three subcommands: `list` (browse datasets by org/name/split), `tasks` (enumerate task IDs within a split with `--offset`/`--limit` pagination), and `upload` (bulk-upload local task directories to OSS with configurable concurrency and `--overwrite` flag) ([#859](https://github.com/alibaba/ROCK/pull/859), [#875](https://github.com/alibaba/ROCK/pull/875)) + +* **NEW**: `OssDatasetRegistry` backend using `oss2` SDK to navigate the `datasets/{org}/{dataset}/{split}/{task_id}/` key structure, supporting both directory-style and flat file-style tasks. OSS credentials can be passed as CLI flags or stored in the `[dataset]` section of `config.ini` ([#859](https://github.com/alibaba/ROCK/pull/859)) + +* Task listing now recognizes both directory tasks (from `prefix_list`) and file tasks (from `object_list`), with automatic suffix stripping and deduplication ([#875](https://github.com/alibaba/ROCK/pull/875)) + +--- + +## Sandbox + +### New Features + +#### Container Disk Limits + +* Supports limiting sandbox rootfs size via Docker `--storage-opt size=`. Requires `overlay2` storage driver + XFS filesystem + `prjquota` mount option. Also supports setting XFS project quotas on sandbox log directories via `xfs_quota`, requiring XFS filesystem + `prjquota` mount option. ([#860](https://github.com/alibaba/ROCK/pull/860)) + +* Server-side `RuntimeConfig` gains `disk_limit_rootfs` and `disk_limit_log` fields (both default to `None`), configurable per environment in `rock-{env}.yml` with Nacos runtime overrides ([#860](https://github.com/alibaba/ROCK/pull/860)) + +#### Sandbox Timezone Support + +* **NEW**: Containers now receive full IANA timezone support — ROCK mounts the host's zoneinfo file (e.g. `/usr/share/zoneinfo/Asia/Shanghai`) read-only to `/etc/localtime`. Automatically skips with a warning if the host lacks the corresponding zoneinfo file ([#883](https://github.com/alibaba/ROCK/pull/883)) + +--- + +## Admin + +### MetaStore & Database Operation Metrics + +* **NEW**: Added OpenTelemetry metrics instrumentation for `SandboxMetaStore` and `SandboxTable` CRUD operations. Each operation (create, get, update, delete, list, batch_get, archive, etc.) is now automatically tracked with total/success/failure counters and response time gauges ([#887](https://github.com/alibaba/ROCK/pull/887)) + +--- + +## Deployments + +### Docker Auth Refactor + +* Replaced the legacy Docker auth scheme with a temporary-directory approach. A `TempAuthDockerClient` context manager creates an isolated temp directory per sandbox, runs `docker --config login`, performs image pulls, and cleans up on exit — preventing registry credentials from persisting in the global Docker config ([#837](https://github.com/alibaba/ROCK/pull/837)) + +* The base directory for temp auth dirs is configurable via `ROCK_DOCKER_TEMP_AUTH_DIR` env var, defaulting to the system temp directory ([#837](https://github.com/alibaba/ROCK/pull/837)) + +--- + +## Proxy + +### WebSocket Header Forwarding + +* The `/sandboxes/{id}/proxy/{path:path}` endpoint supports header forwarding — all client headers are forwarded to upstream services by default, filtering only WebSocket handshake headers (`sec-websocket-*`), hop-by-hop headers (`connection`, `upgrade`, `transfer-encoding`, `content-length`), and `host`. The `origin` header receives special handling — extracted separately and passed as the WebSocket origin parameter. The VNC WebSocket route disables this feature to avoid exceeding QEMU's 4 KB header buffer limit ([#865](https://github.com/alibaba/ROCK/pull/865)) + +--- + +## Bug Fixes + +* Fix `auto_clear_time` calculation: fractional minutes from `auto_clear_seconds / 60` are now rounded up via `math.ceil()` to at least 1 minute; also caps `wait_interval` in the process-alive polling loop to be less than `auto_clear_seconds`, preventing the sandbox from auto-clearing before the first liveness check ([#883](https://github.com/alibaba/ROCK/pull/883)) + +* Fix UV environment setup: project tree is now copied to a writable `/tmp/rock-build` directory before `uv pip install`, resolving failures caused by read-only source mounts in containers ([#857](https://github.com/alibaba/ROCK/pull/857)) + +--- + +## Testing & CI + +* Added unit tests for `DockerUtil` helpers (`detect_storage_opt_support`, `is_xfs_path`, `get_docker_root_dir`) and `DockerDeployment` disk limit integration ([#860](https://github.com/alibaba/ROCK/pull/860)) + +* Added integration tests for Docker temporary-directory auth scheme ([#837](https://github.com/alibaba/ROCK/pull/837)) + +* Added comprehensive unit tests for Datasets CLI, client, models, and OSS registry ([#859](https://github.com/alibaba/ROCK/pull/859), [#875](https://github.com/alibaba/ROCK/pull/875)) + +* Database connection unit tests and parameter optimizations ([#852](https://github.com/alibaba/ROCK/pull/852)) + +* Clean up leaked timers in TS SDK model client tests ([#839](https://github.com/alibaba/ROCK/pull/839)) + +--- + +## Migration Notes + +* **Docker auth**: The legacy Docker auth scheme has been removed. If you relied on persistent credentials in `~/.docker/config.json` written by ROCK, note that ROCK now uses ephemeral temp directories. Set `ROCK_DOCKER_TEMP_AUTH_DIR` to customize the temp directory location. + +* **WebSocket proxy headers**: If your downstream services relied on a specific set of forwarded headers, be aware that all non-blocked headers are now forwarded. The blocked set includes `host`, `connection`, `upgrade`, `sec-websocket-*`, `transfer-encoding`, and `content-length`. + +* **Disk limits**: Disk quotas are server-side policy — `disk_limit` is not exposed in `SandboxStartRequest`. Configure via `RuntimeConfig.disk_limit_rootfs` / `disk_limit_log` in your environment YAML or Nacos overrides. diff --git a/docs/versioned_docs/version-1.7.x/User Guides/configuration.md b/docs/versioned_docs/version-1.7.x/User Guides/configuration.md new file mode 100644 index 0000000000..604256878b --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/User Guides/configuration.md @@ -0,0 +1,189 @@ +--- +sidebar_position: 4 +--- + +# Configuration + +This guide provides detailed instructions on how to configure the ROCK environment to meet different usage requirements, including local development, testing, and production deployment. + +## 1. Environment Variable Configuration + +ROCK supports configuring key parameters through environment variables. The main environment variables are as follows: + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK service base URL +export ROCK_LOG_LEVEL=INFO # Log level +export ROCK_LOGGING_PATH=/path/to/logs # Log file path, default None (output to console) +export ROCK_LOGGING_FILE_NAME=rocklet.log # Log file name, default "rocklet.log", can be customized by admin like admin.log +export ROCK_LOGGING_LEVEL=INFO # Log output level, default "INFO" +export ROCK_WORKER_ENV_TYPE=local # Runtime environment type, options: local, docker, uv, pip +``` + +More environment variables can be found in `rock/env_vars.py`. + +### 1.1 Runtime Environments + +ROCK provides multiple different runtime environments to meet the needs of different scenarios, configured through the `ROCK_WORKER_ENV_TYPE` environment variable. Each environment has different deployment requirements, performance characteristics and applicable scenarios. Each environment has its own unique advantages and limitations, and developers can choose the most suitable runtime environment according to their deployment needs. + +#### 1.1.1 Docker Runtime Environment + +The Docker runtime environment is suitable for Docker image environments where dependencies are pre-installed. This environment requires the `/tmp/miniforge/bin/rocklet` executable to be directly available in the deployment environment. + +**Mount Configuration:** +- `/tmp/miniforge` - Contains pre-installed Python environment +- `/tmp/local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Containerized deployment environments +- Already built custom Docker image containing `rocklet` +- Suitable for production, fast startup + +**Requirements:** +- Requires a custom Docker image containing `/tmp/miniforge/bin/rocklet` executable +- Docker environment support + +#### 1.1.2 Local Runtime Environment + +The local runtime environment directly uses the Python environment and project files of the current deployment. This environment requires the same operating system between the host and container to directly mount the virtual environment and Python interpreter. + +**Mount Configuration:** +- `python_env_path` - Python environment path +- `project_root` - Project root directory +- `.venv` - Virtual environment directory (mounted as `/tmp/miniforge` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Development environments +- Scenarios where host and target container use the same operating system +- Need to quickly reuse existing Python environment + +**Requirements:** +- Same operating system (host/container) +- Direct access to the currently deployed `.venv` virtual environment +- Python interpreter path compatibility + +#### 1.1.3 UV Runtime Environment + +The UV runtime environment only depends on the available ROCK project, but initialization is relatively slow and network requirements are higher. This environment is most suitable for scenarios without preconfigured environments. It rebuilds the rocklet environment from the original project. This is the recommended environment for Mac OS. + +**Mount Configuration:** +- `project_root` - Project root directory (mounted as `/tmp + project_root` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**Use Cases:** +- Mac OS +- Cross-OS startup +- Scenarios without preconfigured environment +- No uv management Rock + +**Advantages:** +- No pre-built image required +- Good cross-platform compatibility +- Suitable for development and testing especially + +**Limitations:** +- Initialization is relatively slow +- Higher network requirements +- Longer startup time + +#### 1.1.4 PIP Runtime Environment + +The PIP runtime environment uses pip to install required dependencies in the container. This environment is suitable for quick setup and scenarios where dependencies can be installed in the container. It is the default runtime environment. It does not require pre-built images containing dependencies, and manages Python packages directly through pip. + +**Mount Configuration:** +- `local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**Use Cases:** +- ROCK installation from PIP source +- Fast testing of ROCK + +**Advantages:** +- Simple deployment setup + +**Limitations:** +- Long dependency installation time +- Requires network access to install dependency packages +- Dependencies need to be installed each time on startup + +#### 1.1.5 Configuration Guide + +Refer to the following selection guide for different use cases: + +| Scenario | Recommended Environment | Reason | +|----------|--------------------------|-------| +| Production environment | Docker Runtime | Fast startup, stable performance | +| Development environment, same OS | Local Runtime | Environment reuse, fast development cycle | +| Mac development | UV Runtime | Best cross-platform compatibility support | +| Cross-platform development | UV Runtime | Avoids environment compatibility issues | +| Fast testing | UV Runtime | Requires no pre-configuration | +| PIP source installation | PIP Runtime | Install dependencies directly with pip | + +These runtime environments are configured through the `ROCK_WORKER_ENV_TYPE` environment variable, which can be set to "local", "docker", "uv" or "pip". + +### 1.2 Logging Configuration + +Regarding logging configuration, ROCK's logging system has the following characteristics: + +- The logging system cannot output to both file and console simultaneously. If `ROCK_LOGGING_PATH` is set, logs will be output to the designated file, otherwise to console. +- `ROCK_LOGGING_LEVEL` is used to control the output log level, while `ROCK_LOG_LEVEL` is used for general log level settings. + +## 2. Distributed Deployment Requirements + +Since ROCK supports distributed deployment, when running on different nodes of a Ray cluster, the following consistency requirements must be met: + +#### Directory Structure Consistency + +On all Ray nodes, the following directory structure must be completely consistent: +- ROCK project repository directory +- `.venv` virtual environment directory +- The base Python directory that `.venv` depends on + + +#### Mounting Requirements + +ROCK's startup depends on mounting the ROCK project and the corresponding base Python environment, requiring consistency in multi-machine environments: + +#### Verifying Distributed Configuration + +Distributed deployment configuration can be verified through the following methods: + +```bash +# Check directory consistency on all nodes +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# Verify Python environment availability +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# Check environment variable settings on all nodes +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly set up the ROCK environment +- [API Documentation](../References/api.md) - View sandbox-related API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to configure sandboxes +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.7.x/overview.md b/docs/versioned_docs/version-1.7.x/overview.md new file mode 100644 index 0000000000..0c377e8b25 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/overview.md @@ -0,0 +1,33 @@ +--- +sidebar_position: 1 +--- + +# Overview + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework designed to simplify the development, deployment, and management of reinforcement learning environments. + +## What is ROCK + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework. By using ROCK, developers can quickly develop reinforcement learning environments and integrate with other reinforcement learning training frameworks to implement efficient reinforcement learning training. + +ROCK provides comprehensive sandbox environment management capabilities, supports containerized deployment, and enables rapid creation, execution, and destruction of environments. Additionally, ROCK is compatible with the GEM protocol, providing standardized interfaces for reinforcement learning environments. + +## Core Capabilities of ROCK + +1. **Simplified Development Process**: Simplifies the development, construction, and management of reinforcement learning environments, supporting various open-source reinforcement learning environments +2. **Large-scale Scheduling and Deployment**: Enables large-scale scheduling and deployment of rapid reinforcement learning environments. By supporting the GEM protocol, reinforcement learning environments can be easily accessed +3. **Framework Integration**: Integrates with other reinforcement learning training frameworks to achieve large-scale and scalable reinforcement learning training + +## Value of ROCK + +ROCK provides significant value to different roles of engineers: + +- **Reinforcement Learning Algorithm Engineers**: ROCK simplifies the development process of reinforcement learning environments, allowing engineers to focus on algorithm implementation +- **Reinforcement Learning Application Engineers**: ROCK enables large-scale deployment of rapid reinforcement learning environments, improving application development efficiency + +## Learn More + +- [Quick Start Guide](./Getting%20Started/quickstart.md) - Get started with ROCK quickly +- [Configuration Guide](./User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](./References/api.md) - View ROCK's API interfaces +- [Python SDK Documentation](./References/Python%20SDK%20References/python_sdk.md) - Learn how to use ROCK's Python SDK \ No newline at end of file diff --git a/docs/versioned_sidebars/version-1.7.x-sidebars.json b/docs/versioned_sidebars/version-1.7.x-sidebars.json new file mode 100644 index 0000000000..b475b11530 --- /dev/null +++ b/docs/versioned_sidebars/version-1.7.x-sidebars.json @@ -0,0 +1,64 @@ +{ + "tutorialSidebar": [ + "overview", + { + "type": "category", + "label": "Getting Started", + "link": { + "type": "doc", + "id": "Getting Started/quickstart" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Getting Started" + } + ] + }, + { + "type": "category", + "label": "User Guides", + "items": [ + { + "type": "autogenerated", + "dirName": "User Guides" + } + ] + }, + { + "type": "category", + "label": "References", + "items": [ + "References/api", + { + "type": "category", + "label": "Python SDK References", + "link": { + "type": "doc", + "id": "References/Python SDK References/python_sdk" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "References/Python SDK References" + } + ] + } + ] + }, + { + "type": "category", + "label": "Release Notes", + "link": { + "type": "doc", + "id": "Release Notes/index" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Release Notes" + } + ] + } + ] +} diff --git a/docs/versions.json b/docs/versions.json index aa33c9a4b3..fb1f7d2f57 100644 --- a/docs/versions.json +++ b/docs/versions.json @@ -1,4 +1,5 @@ [ + "1.7.x", "1.6.x", "1.5.x", "1.4.x", From 4bb2d1e73576f74f34f98a8e97f8271df933818d Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 24 Apr 2026 15:01:54 +0800 Subject: [PATCH 075/226] docs: update README with v1.7.0 release entry Co-Authored-By: Claude Opus 4.6 --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index f329c8a881..3e4c86e15a 100644 --- a/README.md +++ b/README.md @@ -155,11 +155,11 @@ if __name__ == "__main__": | Date | Release | |:-------------|:---| -| **[Latest]** | 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | +| **[Latest]** | 🎉 ROCK v1.7.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.7.0) | +| **[2026-04-23]** | 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | | **[2026-04-10]** | 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | | **[2026-03-27]** | 🎉 ROCK v1.4.4 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.4) | | **[2026-03-24]** | 🎉 ROCK v1.4.3 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.3) | -| **[2026-03-17]** | 🎉 ROCK v1.4.2 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.2) | --- From b5997f39935ef97cd69a92bbb8831a6640203a04 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 24 Apr 2026 15:03:47 +0800 Subject: [PATCH 076/226] bump version to 1.7.0 --- .../version-1.7.x/Release Notes/v1.7.0.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md index e5f084cc80..58ab0f8985 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Release Notes/v1.7.0.md @@ -69,7 +69,7 @@ ### WebSocket Header转发 -* /sandboxes/{id}/proxy/{path:path} 接口支持header转发，所有客户端头部默认转发至上游服务，仅过滤 WebSocket 握手头（`sec-websocket-*`）、逐跳头（`connection`、`upgrade`、`transfer-encoding`、`content-length`）及 `host`。`origin` 头部做特殊处理 — 单独提取并作为 WebSocket origin 参数传入，VNC WebSocket 接口关闭该功能，避免超出 QEMU 的 4 KB 头部缓冲区限制([#865](https://github.com/alibaba/ROCK/pull/865)) +* `/sandboxes/{id}/proxy/{path:path}` 接口支持header转发，所有客户端头部默认转发至上游服务，仅过滤 WebSocket 握手头（`sec-websocket-*`）、逐跳头（`connection`、`upgrade`、`transfer-encoding`、`content-length`）及 `host`。`origin` 头部做特殊处理 — 单独提取并作为 WebSocket origin 参数传入，VNC WebSocket 接口关闭该功能，避免超出 QEMU 的 4 KB 头部缓冲区限制([#865](https://github.com/alibaba/ROCK/pull/865)) --- diff --git a/pyproject.toml b/pyproject.toml index 93e479798c..233743e593 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.6.1" +version = "1.7.0" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From a6e8e7f7502ee9a58223e057065f97ae8901b2d2 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Mon, 27 Apr 2026 14:34:41 +0800 Subject: [PATCH 077/226] =?UTF-8?q?fix(bash-trial):=20BashJob=20OSS=20?= =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E6=94=B9=E4=B8=BA=E8=84=9A=E6=9C=AC=E6=B3=A8?= =?UTF-8?q?=E5=85=A5,=E4=BF=AE=E5=A4=8D=20submit-only=20=E4=B8=A2=E4=BC=A0?= =?UTF-8?q?=E5=92=8C=20env=20=E5=87=AD=E8=AF=81=E4=B8=8D=E7=94=9F=E6=95=88?= =?UTF-8?q?=20(#898)=20(#899)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(executor): _build_session_env 支持 oss_mirror 三级凭证优先级和派生项 - 凭证优先级:OssMirrorConfig 字段 > environment.env > 宿主 os.environ - oss_mirror.enabled 时校验 namespace / experiment_id / OSS_BUCKET 非空 - 注入 ROCK_ARTIFACT_DIR / ROCK_OSS_PREFIX 供 wrapper 脚本读取 - os.environ OSS_* 过滤从 startswith("OSS") 收紧为 startswith("OSS_") refs spec: docs/superpowers/specs/2026-04-27-bashjob-oss-script-injection-design.md * feat(bash-trial): 新增 wrapper 渲染器 _render_wrapper * feat(bash-trial): build() 在 oss_mirror.enabled 时生成 wrapper 脚本 * refactor(bash-trial): 删除宿主侧 OSS 上传代码路径,改由 wrapper 脚本承担 * chore: bump version to 1.7.1 * refactor(bash-trial): _oss_mirror property 替换为 _oss_mirror_enabled() 布尔方法 * feat(executor): _build_session_env 增补 OSS_ENDPOINT / OSS_REGION 非空校验 ossutil v2 的 sign v4 要求 region 必须有值,缺失会在沙箱内报 'region must be set in sign version 4' 且 wrapper 吞掉错误 || true, 难以排查。提前在宿主侧 fail-fast 让问题更快暴露。 * refactor(bash-trial): 下沉 OSS 凭证合并/校验/派生项到 BashTrial.on_sandbox_ready _build_session_env 是所有 trial 共用的,把 BashJob-specific 的 OSS 逻辑放进来 会污染 Harbor Job 行为(多余的校验、派生项注入)。改为: - _build_session_env 完全还原到 master 版本(Harbor 行为不变) - BashTrial 新增 on_sandbox_ready override:先 super() 回填 namespace/experiment_id, 然后 _prepare_oss_session_env 按三级优先级(OssMirrorConfig 字段 > environment.env > 宿主 os.environ)把凭证写进 config.environment.env;同时校验 namespace/experiment_id/OSS_BUCKET/OSS_ENDPOINT/OSS_REGION 非空;最后注入 ROCK_ARTIFACT_DIR / ROCK_OSS_PREFIX 派生项 - _build_session_env 之后会自动把 environment.env 合并到 session env,效果等价 测试:TestBuildSessionEnvOssMirror(9 个)删除,等价迁移为 TestBashTrialPrepareOssSessionEnv(10 个),覆盖 on_sandbox_ready 路径。 * refactor(bash-trial): _render_wrapper 移入 BashTrial 作为 staticmethod;凭证字段映射用 upper() 去重 - _render_wrapper 是 BashTrial 专属逻辑,作为类的 staticmethod 比 module-level 函数语义更准 - _prepare_oss_session_env 内 (snake, UPPER) 元组列表改成 _OSS_CREDENTIAL_FIELDS 单元组 + field_name.upper(),少一半重复 - 三个 OSS_BUCKET/ENDPOINT/REGION 校验合并成单 for 循环 * Revert "chore: bump version to 1.7.1" This reverts commit d6d14948de035c91cea11f7536c272dfb41caa5a. --- rock/sdk/job/trial/bash.py | 184 ++++++++------ tests/unit/sdk/job/test_trial_bash.py | 347 +++++++++++++++++++++++--- 2 files changed, 431 insertions(+), 100 deletions(-) diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index ed084eea69..915fe8d497 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -3,20 +3,27 @@ from __future__ import annotations import os -import shlex +import secrets from pathlib import Path from rock import env_vars -from rock.actions.sandbox.request import Command from rock.logger import init_logger from rock.sdk.job.config import BashJobConfig from rock.sdk.job.result import ExceptionInfo, TrialResult from rock.sdk.job.trial.abstract import AbstractTrial from rock.sdk.job.trial.registry import register_trial -from rock.sdk.sandbox.client import RunMode, Sandbox +from rock.sdk.sandbox.client import Sandbox logger = init_logger(__name__) +_OSS_CREDENTIAL_FIELDS = ( + "oss_access_key_id", + "oss_access_key_secret", + "oss_endpoint", + "oss_region", + "oss_bucket", +) + class BashTrial(AbstractTrial): """Bash script execution trial.""" @@ -26,57 +33,120 @@ class BashTrial(AbstractTrial): def __init__(self, config: BashJobConfig): super().__init__(config) self._ossutil_ready: bool = False - self._oss_credentials: dict | None = None - self._artifact_dir: str | None = None - - @property - def _oss_mirror(self): - return self._config.environment.oss_mirror - - async def setup(self, sandbox: Sandbox) -> None: - await self._upload_files(sandbox) - if self._config.script_path: - self._config.script = Path(self._config.script_path).read_text() - if self._oss_mirror is not None and self._oss_mirror.enabled: - await self._setup_oss_mirror(sandbox) + def _oss_mirror_enabled(self) -> bool: + mirror = self._config.environment.oss_mirror + return mirror is not None and mirror.enabled + + async def on_sandbox_ready(self, sandbox: Sandbox) -> None: + """Backfill namespace/experiment_id (via super) then prepare OSS session env. + + All BashJob-specific session env preparation (credential resolution, + validation, derived ROCK_* keys) lives here so that the shared + ``JobExecutor._build_session_env`` stays trial-agnostic. Because this + hook runs strictly before ``_build_session_env``, any keys we write + into ``config.environment.env`` here will propagate into the bash + session env. + """ + await super().on_sandbox_ready(sandbox) + if self._oss_mirror_enabled(): + self._prepare_oss_session_env() + + def _prepare_oss_session_env(self) -> None: + """Resolve OSS credentials, validate, and inject derived ROCK_* keys. + + Resolution order per key (first non-empty wins): + 1. ``OssMirrorConfig`` field (highest priority) + 2. ``environment.env`` (if the user already put it there) + 3. Host process ``os.environ`` (lowest priority) + + The resolved credentials are written into ``environment.env`` so that + ``JobExecutor._build_session_env`` picks them up without needing to + know anything about OSS. Also writes ``ROCK_ARTIFACT_DIR`` and + ``ROCK_OSS_PREFIX`` for the wrapper script to consume. + """ + mirror = self._config.environment.oss_mirror + env = self._config.environment.env + + for field_name in _OSS_CREDENTIAL_FIELDS: + env_key = field_name.upper() + v = getattr(mirror, field_name, None) or env.get(env_key) or os.environ.get(env_key) + if v: + env[env_key] = v - async def _setup_oss_mirror(self, sandbox: Sandbox) -> None: if not self._config.namespace: raise ValueError("oss_mirror: namespace is not set (sandbox did not return one)") if not self._config.experiment_id: raise ValueError("oss_mirror: experiment_id is not set (sandbox did not return one)") + for env_key in ("OSS_BUCKET", "OSS_ENDPOINT", "OSS_REGION"): + if not env.get(env_key): + raise ValueError(f"oss_mirror.enabled=True but {env_key} is not resolvable") - self._artifact_dir = env_vars.ROCK_BASH_JOB_ARTIFACT_DIR - - bucket = self._oss_mirror.oss_bucket or os.environ.get("OSS_BUCKET") - if not bucket: - raise ValueError("oss_mirror.enabled=True but oss_bucket is not set (config or OSS_BUCKET env)") - - self._oss_credentials = { - "oss_bucket": bucket, - "access_key_id": self._oss_mirror.oss_access_key_id or os.environ.get("OSS_ACCESS_KEY_ID", ""), - "access_key_secret": self._oss_mirror.oss_access_key_secret or os.environ.get("OSS_ACCESS_KEY_SECRET", ""), - "endpoint": self._oss_mirror.oss_endpoint or os.environ.get("OSS_ENDPOINT", ""), - "region": self._oss_mirror.oss_region or os.environ.get("OSS_REGION", ""), - } + env["ROCK_ARTIFACT_DIR"] = env_vars.ROCK_BASH_JOB_ARTIFACT_DIR + env[ + "ROCK_OSS_PREFIX" + ] = f"artifacts/{self._config.namespace}/{self._config.experiment_id}/{self._config.job_name}" - await sandbox.execute(Command(command=["mkdir", "-p", self._artifact_dir])) - # Touch a placeholder so ossutil cp has something to upload (OSS has no real dirs) - await sandbox.execute(Command(command=["touch", f"{self._artifact_dir}/.placeholder"])) - - self._ossutil_ready = await sandbox.fs.ensure_ossutil() - if not self._ossutil_ready: - logger.warning("ossutil install failed, OSS mirror upload will be skipped") - return + @staticmethod + def _render_wrapper(user_script: str, token: str | None = None) -> str: + """Render the BashJob wrapper script. + + Structure: prologue (mkdir + touch + initial upload) → user script + (isolated in a single-quoted heredoc) → epilogue (final upload) → exit + with user's exit code. + + When ``token`` is ``None`` a random 8-char hex is generated. Collision + probability is ~2^-32 and collisions also require the user script to + contain the terminator on its own line, which is not actively guarded + against — the risk is acceptable. + """ + if token is None: + token = secrets.token_hex(4) # 8-char hex + eof = f"__ROCK_USER_SCRIPT_EOF_{token}__" + return ( + "#!/bin/bash\n" + "# rock bash-job wrapper (generated, do not edit)\n" + "# OSS credentials and paths come from session env; no secrets in this file.\n" + "set +e\n" + "\n" + "# -- prologue: prepare artifact dir and do an initial placeholder upload --\n" + 'mkdir -p "$ROCK_ARTIFACT_DIR"\n' + 'touch "$ROCK_ARTIFACT_DIR/.placeholder"\n' + 'ossutil cp "$ROCK_ARTIFACT_DIR/" "oss://$OSS_BUCKET/$ROCK_OSS_PREFIX/" \\\n' + " --recursive -f >/dev/null 2>&1 || true\n" + "\n" + "# -- user script: heredoc isolates user's trap/exit from the wrapper --\n" + f"bash <<'{eof}'\n" + f"{user_script}\n" + f"{eof}\n" + "_rock_user_rc=$?\n" + "\n" + "# -- epilogue: final upload (failure is logged but does not change exit code) --\n" + 'ossutil cp "$ROCK_ARTIFACT_DIR/" "oss://$OSS_BUCKET/$ROCK_OSS_PREFIX/" \\\n' + " --recursive -f \\\n" + ' || echo "[rock] oss upload failed (rc=$?), ignored" >&2\n' + "\n" + "exit $_rock_user_rc\n" + ) - await self._upload_artifacts(sandbox) + async def setup(self, sandbox: Sandbox) -> None: + await self._upload_files(sandbox) + if self._config.script_path: + self._config.script = Path(self._config.script_path).read_text() - def _build_oss_prefix(self) -> str: - return f"artifacts/{self._config.namespace}/{self._config.experiment_id}/{self._config.job_name}" + if self._oss_mirror_enabled(): + self._ossutil_ready = await sandbox.fs.ensure_ossutil() + if not self._ossutil_ready: + logger.warning("ossutil install failed, OSS mirror upload will be skipped") def build(self) -> str: - return self._config.script or "" + script = self._config.script or "" + if not self._oss_mirror_enabled(): + return script + if not self._ossutil_ready: + logger.warning("ossutil unavailable, falling back to raw script (OSS mirror upload disabled for this run)") + return script + return self._render_wrapper(script) async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialResult: exception_info = None @@ -86,9 +156,6 @@ async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialR exception_message=f"Bash script exited with code {exit_code}", ) - if self._oss_mirror is not None and self._oss_mirror.enabled and self._ossutil_ready and self._oss_credentials: - await self._upload_artifacts(sandbox) - return TrialResult( task_name=self._config.job_name or "", exception_info=exception_info, @@ -96,33 +163,6 @@ async def collect(self, sandbox: Sandbox, output: str, exit_code: int) -> TrialR exit_code=exit_code, ) - @staticmethod - def _build_ossutil_cmd(ossutil_args: str, creds: dict) -> str: - inner = ( - f"ossutil {ossutil_args}" - f" --access-key-id {shlex.quote(creds['access_key_id'])}" - f" --access-key-secret {shlex.quote(creds['access_key_secret'])}" - f" --endpoint {shlex.quote(creds['endpoint'])}" - f" --region {shlex.quote(creds['region'])}" - ) - return f"bash -c {shlex.quote(inner)}" - - async def _upload_artifacts(self, sandbox: Sandbox) -> None: - try: - oss_url = f"oss://{self._oss_credentials['oss_bucket']}/{self._build_oss_prefix()}/" - src = self._artifact_dir.rstrip("/") + "/" - cmd = self._build_ossutil_cmd( - f"cp {shlex.quote(src)} {shlex.quote(oss_url)} --recursive", - self._oss_credentials, - ) - result = await sandbox.arun(cmd=cmd, mode=RunMode.NOHUP, wait_timeout=600) - if result.exit_code != 0: - logger.warning(f"OSS mirror upload failed: {result.output}") - else: - logger.info(f"OSS mirror upload completed: {self._artifact_dir} -> {oss_url}") - except Exception as e: - logger.warning(f"OSS mirror upload error: {e}") - # Auto-register on import register_trial(BashJobConfig, BashTrial) diff --git a/tests/unit/sdk/job/test_trial_bash.py b/tests/unit/sdk/job/test_trial_bash.py index ff3e4d5f9d..4831ba4b9a 100644 --- a/tests/unit/sdk/job/test_trial_bash.py +++ b/tests/unit/sdk/job/test_trial_bash.py @@ -37,6 +37,52 @@ def test_build_empty_script(self): trial = BashTrial(cfg) assert trial.build() == "" + def test_build_without_oss_mirror_returns_raw(self): + """No oss_mirror -> user script is returned as-is.""" + cfg = BashJobConfig(script="echo hi") + trial = BashTrial(cfg) + assert trial.build() == "echo hi" + + def test_build_oss_disabled_returns_raw(self): + """oss_mirror.enabled=False -> raw script.""" + cfg = BashJobConfig( + script="echo hi", + environment=EnvironmentConfig(oss_mirror=OssMirrorConfig(enabled=False)), + ) + trial = BashTrial(cfg) + assert trial.build() == "echo hi" + + def test_build_oss_enabled_and_ready_returns_wrapper(self): + """oss_mirror.enabled=True and ossutil_ready -> wrapper script.""" + cfg = BashJobConfig( + script="echo hi", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b"), + ), + ) + trial = BashTrial(cfg) + trial._ossutil_ready = True # simulate successful setup + + wrapper = trial.build() + + assert wrapper.startswith("#!/bin/bash") + assert "echo hi" in wrapper + assert "ossutil cp" in wrapper + assert "__ROCK_USER_SCRIPT_EOF_" in wrapper + + def test_build_oss_enabled_but_not_ready_falls_back(self): + """oss_mirror.enabled=True but ossutil install failed -> raw script.""" + cfg = BashJobConfig( + script="echo hi", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b"), + ), + ) + trial = BashTrial(cfg) + trial._ossutil_ready = False # ensure_ossutil failed + + assert trial.build() == "echo hi" + # --------------------------------------------------------------------------- # BashTrial.setup() @@ -186,7 +232,9 @@ def _oss_sandbox(ns="ns", exp="exp"): class TestBashTrialOssMirror: - async def test_setup_installs_ossutil_and_creates_dir(self): + """OSS mirror integration (spec 2026-04-27) — setup only installs ossutil, collect does not upload.""" + + async def test_setup_installs_ossutil_when_enabled(self): cfg = BashJobConfig( script="echo", job_name="j", @@ -196,12 +244,30 @@ async def test_setup_installs_ossutil_and_creates_dir(self): ) trial = BashTrial(cfg) sb = _oss_sandbox() + await trial.setup(sb) sb.fs.ensure_ossutil.assert_called_once() - # Initial upload to create OSS path before script runs - setup_cp_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] - assert len(setup_cp_calls) == 1 + assert trial._ossutil_ready is True + # setup must not trigger any ossutil cp — uploads happen inside the wrapper + cp_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] + assert cp_calls == [] + + async def test_setup_marks_ossutil_not_ready_on_install_failure(self): + cfg = BashJobConfig( + script="echo", + job_name="j", + namespace="ns", + experiment_id="exp", + environment=EnvironmentConfig(oss_mirror=_MIRROR), + ) + trial = BashTrial(cfg) + sb = _oss_sandbox() + sb.fs.ensure_ossutil = AsyncMock(return_value=False) + + await trial.setup(sb) + + assert trial._ossutil_ready is False async def test_setup_skips_when_no_mirror(self): trial = BashTrial(BashJobConfig(script="echo")) @@ -209,7 +275,8 @@ async def test_setup_skips_when_no_mirror(self): await trial.setup(sb) sb.fs.ensure_ossutil.assert_not_called() - async def test_collect_uploads(self): + async def test_collect_does_not_upload(self): + """collect no longer calls ossutil cp — uploads are wrapper-driven.""" cfg = BashJobConfig( script="echo", job_name="j", @@ -219,41 +286,265 @@ async def test_collect_uploads(self): ) trial = BashTrial(cfg) sb = _oss_sandbox() + await trial.setup(sb) await trial.collect(sb, "ok", 0) - # setup + collect each call ossutil cp once - arun_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] - assert len(arun_calls) == 2 - assert all("oss://b/artifacts/ns/exp/j/" in str(c) for c in arun_calls) + cp_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] + assert cp_calls == [] + + +# --------------------------------------------------------------------------- +# Wrapper renderer (spec 2026-04-27) +# --------------------------------------------------------------------------- + + +class TestRenderWrapper: + def test_wrapper_contains_prologue_and_epilogue(self): + wrapper = BashTrial._render_wrapper("echo hi", token="deadbeef") + + # prologue + assert 'mkdir -p "$ROCK_ARTIFACT_DIR"' in wrapper + assert 'touch "$ROCK_ARTIFACT_DIR/.placeholder"' in wrapper + # initial upload (silent on failure) + assert "|| true" in wrapper + # heredoc terminator appears twice (open + close) + assert wrapper.count("__ROCK_USER_SCRIPT_EOF_deadbeef__") == 2 + # single-quoted terminator (disables parameter expansion) + assert "bash <<'__ROCK_USER_SCRIPT_EOF_deadbeef__'" in wrapper + # user script body + assert "echo hi" in wrapper + # capture user exit code + assert "_rock_user_rc=$?" in wrapper + # epilogue: final upload with -f + assert "ossutil cp" in wrapper + assert "--recursive -f" in wrapper + assert "exit $_rock_user_rc" in wrapper + + def test_wrapper_uses_oss_env_variables(self): + """Wrapper reads paths/bucket from env only; no plaintext credentials.""" + wrapper = BashTrial._render_wrapper("echo hi", token="deadbeef") + + # env-var references present + assert '"oss://$OSS_BUCKET/$ROCK_OSS_PREFIX/"' in wrapper + # credential flags must not appear (no command-line leakage) + assert "--access-key-id" not in wrapper + assert "--access-key-secret" not in wrapper + + def test_wrapper_empty_user_script(self): + wrapper = BashTrial._render_wrapper("", token="deadbeef") + assert "__ROCK_USER_SCRIPT_EOF_deadbeef__" in wrapper + # prologue/epilogue still present + assert "ossutil cp" in wrapper + + def test_wrapper_preserves_user_script_verbatim(self): + """Single-quoted heredoc keeps $VAR, backticks, $() unexpanded.""" + from rock.sdk.job.trial.bash import BashTrial + + user = 'echo "$HOME $(date) `whoami`"' + wrapper = BashTrial._render_wrapper(user, token="deadbeef") + assert user in wrapper + + def test_wrapper_auto_generates_token_when_omitted(self): + """Without an explicit token, secrets.token_hex(4) is used.""" + import re as _re + + wrapper = BashTrial._render_wrapper("echo hi") + match = _re.search(r"__ROCK_USER_SCRIPT_EOF_([0-9a-f]{8})__", wrapper) + assert match is not None, "wrapper should contain auto-generated 8-char hex token" + + +# --------------------------------------------------------------------------- +# on_sandbox_ready / _prepare_oss_session_env (spec 2026-04-27) +# --------------------------------------------------------------------------- + + +def _ready_sandbox(ns="ns", exp="exp"): + """Minimal sandbox mock for on_sandbox_ready: only namespace/experiment_id.""" + sb = MagicMock() + sb._namespace = ns + sb._experiment_id = exp + return sb + + +class TestBashTrialPrepareOssSessionEnv: + """BashTrial.on_sandbox_ready resolves OSS credentials, validates, and + writes derived ROCK_* keys into environment.env. JobExecutor._build_session_env + stays trial-agnostic.""" + + def _clear_oss(self, monkeypatch): + for k in list(__import__("os").environ): + if k.startswith("OSS"): + monkeypatch.delenv(k, raising=False) + + async def test_oss_mirror_config_field_wins_over_env_and_host(self, monkeypatch): + """Priority: OssMirrorConfig field > environment.env > host os.environ.""" + self._clear_oss(monkeypatch) + monkeypatch.setenv("OSS_ACCESS_KEY_ID", "host_id") - async def test_upload_failure_does_not_fail_job(self): cfg = BashJobConfig( script="echo", job_name="j", - namespace="ns", - experiment_id="exp", - environment=EnvironmentConfig(oss_mirror=_MIRROR), + environment=EnvironmentConfig( + env={"OSS_ACCESS_KEY_ID": "env_id", "OSS_ENDPOINT": "env_ep", "OSS_REGION": "env_rg"}, + oss_mirror=OssMirrorConfig( + enabled=True, + oss_bucket="cfg_bucket", + oss_access_key_id="cfg_id", + ), + ), ) trial = BashTrial(cfg) - sb = _oss_sandbox() - sb.arun = AsyncMock(return_value=MagicMock(exit_code=1, output="err")) - await trial.setup(sb) - result = await trial.collect(sb, "ok", 0) - assert result.exit_code == 0 and result.exception_info is None - async def test_skips_upload_when_ossutil_not_ready(self): + await trial.on_sandbox_ready(_ready_sandbox()) + + assert cfg.environment.env["OSS_ACCESS_KEY_ID"] == "cfg_id" + assert cfg.environment.env["OSS_BUCKET"] == "cfg_bucket" + # environment.env fills slots the config did not supply + assert cfg.environment.env["OSS_ENDPOINT"] == "env_ep" + assert cfg.environment.env["OSS_REGION"] == "env_rg" + + async def test_environment_env_can_supply_oss_credentials(self, monkeypatch): + """Issue-2 fix: OSS_* inside environment.env are usable as credentials.""" + self._clear_oss(monkeypatch) + cfg = BashJobConfig( script="echo", job_name="j", - namespace="ns", - experiment_id="exp", - environment=EnvironmentConfig(oss_mirror=_MIRROR), + environment=EnvironmentConfig( + env={ + "OSS_ACCESS_KEY_ID": "ak", + "OSS_ACCESS_KEY_SECRET": "sk", + "OSS_ENDPOINT": "ep", + "OSS_REGION": "rg", + "OSS_BUCKET": "b", + }, + oss_mirror=OssMirrorConfig(enabled=True), + ), ) trial = BashTrial(cfg) - sb = _oss_sandbox() - sb.fs.ensure_ossutil = AsyncMock(return_value=False) - await trial.setup(sb) - await trial.collect(sb, "ok", 0) - ossutil_calls = [c for c in sb.arun.call_args_list if "ossutil cp" in str(c)] - assert len(ossutil_calls) == 0 + + await trial.on_sandbox_ready(_ready_sandbox()) + + assert cfg.environment.env["OSS_ACCESS_KEY_ID"] == "ak" + assert cfg.environment.env["OSS_BUCKET"] == "b" + + async def test_derived_rock_env_keys_present_when_enabled(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="myjob", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b", oss_endpoint="ep", oss_region="rg"), + ), + ) + trial = BashTrial(cfg) + + await trial.on_sandbox_ready(_ready_sandbox(ns="ns1", exp="exp1")) + + assert cfg.environment.env["ROCK_ARTIFACT_DIR"] == "/data/logs/user-defined" + assert cfg.environment.env["ROCK_OSS_PREFIX"] == "artifacts/ns1/exp1/myjob" + + async def test_no_action_when_mirror_disabled(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=False, oss_bucket="b"), + ), + ) + trial = BashTrial(cfg) + + await trial.on_sandbox_ready(_ready_sandbox()) + + # No OSS_* / ROCK_* keys injected when mirror disabled + assert "OSS_BUCKET" not in cfg.environment.env + assert "ROCK_ARTIFACT_DIR" not in cfg.environment.env + assert "ROCK_OSS_PREFIX" not in cfg.environment.env + + async def test_no_action_when_no_mirror(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig(script="echo", environment=EnvironmentConfig()) + trial = BashTrial(cfg) + + await trial.on_sandbox_ready(_ready_sandbox()) + + assert "ROCK_ARTIFACT_DIR" not in cfg.environment.env + + async def test_missing_namespace_raises(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="j", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b", oss_endpoint="ep", oss_region="rg"), + ), + ) + trial = BashTrial(cfg) + + with pytest.raises(ValueError, match="namespace"): + await trial.on_sandbox_ready(_ready_sandbox(ns=None, exp="exp")) + + async def test_missing_experiment_id_raises(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="j", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b", oss_endpoint="ep", oss_region="rg"), + ), + ) + trial = BashTrial(cfg) + + with pytest.raises(ValueError, match="experiment_id"): + await trial.on_sandbox_ready(_ready_sandbox(ns="ns", exp=None)) + + async def test_missing_bucket_raises(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="j", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_endpoint="ep", oss_region="rg"), + ), + ) + trial = BashTrial(cfg) + + with pytest.raises(ValueError, match="OSS_BUCKET"): + await trial.on_sandbox_ready(_ready_sandbox()) + + async def test_missing_endpoint_raises(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="j", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b", oss_region="rg"), + ), + ) + trial = BashTrial(cfg) + + with pytest.raises(ValueError, match="OSS_ENDPOINT"): + await trial.on_sandbox_ready(_ready_sandbox()) + + async def test_missing_region_raises(self, monkeypatch): + self._clear_oss(monkeypatch) + + cfg = BashJobConfig( + script="echo", + job_name="j", + environment=EnvironmentConfig( + oss_mirror=OssMirrorConfig(enabled=True, oss_bucket="b", oss_endpoint="ep"), + ), + ) + trial = BashTrial(cfg) + + with pytest.raises(ValueError, match="OSS_REGION"): + await trial.on_sandbox_ready(_ready_sandbox()) From fa34a2348f32380e299a613a941696e0301accd2 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 28 Apr 2026 10:52:03 +0800 Subject: [PATCH 078/226] chore: remove the need_database marker (#901) Signed-off-by: Jiachen Zhang --- pyproject.toml | 1 - tests/unit/admin/core/test_sandbox_table.py | 1 - tests/unit/admin/core/test_schema_varchar_lengths.py | 1 - tests/unit/admin/core/test_statement_cache.py | 1 - tests/unit/sandbox/test_sandbox_meta_store.py | 1 - tests/unit/utils/test_redis_provider_docker.py | 1 - 6 files changed, 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 233743e593..badb7d1a4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -173,5 +173,4 @@ markers = [ "need_docker: need docker daemon running", "need_admin: need admin start", "need_admin_and_network: need install from network", - "need_database: need database Docker containers (PostgreSQL, Redis)" ] diff --git a/tests/unit/admin/core/test_sandbox_table.py b/tests/unit/admin/core/test_sandbox_table.py index 87035ff900..a733c33b60 100644 --- a/tests/unit/admin/core/test_sandbox_table.py +++ b/tests/unit/admin/core/test_sandbox_table.py @@ -99,7 +99,6 @@ async def test_not_null_defaults_applied_on_insert(self, db): @pytest.mark.need_docker -@pytest.mark.need_database class TestSandboxTableWithPostgres: """Integration tests for SandboxTable using a real PostgreSQL container.""" diff --git a/tests/unit/admin/core/test_schema_varchar_lengths.py b/tests/unit/admin/core/test_schema_varchar_lengths.py index 812e336597..9c24a6a04a 100644 --- a/tests/unit/admin/core/test_schema_varchar_lengths.py +++ b/tests/unit/admin/core/test_schema_varchar_lengths.py @@ -25,7 +25,6 @@ @pytest.mark.need_docker -@pytest.mark.need_database class TestImageVarcharLength: """ORM image column must accept long registry paths on real PostgreSQL.""" diff --git a/tests/unit/admin/core/test_statement_cache.py b/tests/unit/admin/core/test_statement_cache.py index a73b550bd6..7ac4d5a778 100644 --- a/tests/unit/admin/core/test_statement_cache.py +++ b/tests/unit/admin/core/test_statement_cache.py @@ -21,7 +21,6 @@ @pytest.mark.need_docker -@pytest.mark.need_database class TestBatchGetAfterDDL: """Reproduce and fix InvalidCachedStatementError on the sandboxes/batch code path.""" diff --git a/tests/unit/sandbox/test_sandbox_meta_store.py b/tests/unit/sandbox/test_sandbox_meta_store.py index 110a590f86..b90a17a966 100644 --- a/tests/unit/sandbox/test_sandbox_meta_store.py +++ b/tests/unit/sandbox/test_sandbox_meta_store.py @@ -374,7 +374,6 @@ def docker_repo(real_redis, real_db): @pytest.mark.need_docker -@pytest.mark.need_database class TestSandboxMetaStoreWithDocker: """SandboxMetaStore verified against real Redis Stack + PostgreSQL. diff --git a/tests/unit/utils/test_redis_provider_docker.py b/tests/unit/utils/test_redis_provider_docker.py index 2e63d3d3d9..cb0ab74e4a 100644 --- a/tests/unit/utils/test_redis_provider_docker.py +++ b/tests/unit/utils/test_redis_provider_docker.py @@ -6,7 +6,6 @@ @pytest.mark.need_docker -@pytest.mark.need_database class TestRedisProviderWithDocker: """Integration tests for RedisProvider using a real Redis Stack container.""" From 39c970f76b8377fa8306c7a62eecaac20683925d Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Thu, 7 May 2026 19:08:39 +0800 Subject: [PATCH 079/226] Fix sandbox _get_user_info metrics (#911) * test(metrics): add test exposing _get_user_info always returning defaults The new test_decorator_retrieves_user_info_from_meta_store uses a real SandboxMetaStore (FakeRedis + SQLite) to verify that monitor_sandbox_operation populates user_id/experiment_id/namespace from self._meta_store. Currently fails because the decorator reads self._redis_provider which doesn't exist. * fix(metrics): use meta_store instead of redis_provider in _get_user_info SandboxManager and SandboxProxyService have _meta_store but no _redis_provider, so getattr(self, "_redis_provider") always returned None, causing user_id, experiment_id, and namespace to always be "default" in metrics attributes. * test(metrics): update _get_user_info tests to use meta_store interface --- rock/admin/metrics/decorator.py | 32 ++++++----- tests/unit/admin/metrics/test_decorator.py | 67 ++++++++++++++-------- 2 files changed, 62 insertions(+), 37 deletions(-) diff --git a/rock/admin/metrics/decorator.py b/rock/admin/metrics/decorator.py index 98788a8645..3ae58bfb9b 100644 --- a/rock/admin/metrics/decorator.py +++ b/rock/admin/metrics/decorator.py @@ -1,13 +1,17 @@ +from __future__ import annotations + import asyncio import functools import logging import time from collections.abc import Callable +from typing import TYPE_CHECKING -from rock.admin.core.redis_key import alive_sandbox_key from rock.admin.metrics.constants import MetricsConstants from rock.admin.metrics.monitor import MetricsMonitor -from rock.utils.providers import RedisProvider + +if TYPE_CHECKING: + from rock.sandbox.sandbox_meta_store import SandboxMetaStore def _extract_sandbox_id( @@ -40,14 +44,14 @@ def _extract_sandbox_id( return sandbox_id -async def _get_user_info(redis_provider: RedisProvider, sandbox_id: str): - """Get user info from Redis""" - if redis_provider and sandbox_id != "unknown": - user_info = await redis_provider.json_get(alive_sandbox_key(sandbox_id), "$") - if user_info is not None and len(user_info) > 0: - user_id = user_info[0].get("user_id") - experiment_id = user_info[0].get("experiment_id") - namespace = user_info[0].get("namespace") +async def _get_user_info(meta_store: SandboxMetaStore | None, sandbox_id: str): + """Get user info via meta_store (SandboxMetaStore.get).""" + if meta_store and sandbox_id != "unknown": + sandbox_info = await meta_store.get(sandbox_id) + if sandbox_info is not None: + user_id = sandbox_info.get("user_id") + experiment_id = sandbox_info.get("experiment_id") + namespace = sandbox_info.get("namespace") return ( user_id if user_id is not None else "default", experiment_id if experiment_id is not None else "default", @@ -173,8 +177,8 @@ async def wrapper(self, *args, **kwargs): args, kwargs, extract_sandbox_id, sandbox_id_position, sandbox_id_param ) - redis_provider: RedisProvider = getattr(self, "_redis_provider", None) - user_id, experiment_id, namespace = await _get_user_info(redis_provider, sandbox_id) + meta_store = getattr(self, "_meta_store", None) + user_id, experiment_id, namespace = await _get_user_info(meta_store, sandbox_id) # Build attributes attributes = _build_attributes(op_name, sandbox_id, f, user_id, experiment_id, namespace) @@ -206,9 +210,9 @@ def wrapper(self, *args, **kwargs): args, kwargs, extract_sandbox_id, sandbox_id_position, sandbox_id_param ) - redis_provider: RedisProvider = getattr(self, "_redis_provider", None) + meta_store = getattr(self, "_meta_store", None) # For sync functions, we need to run the async function in a blocking way - user_id, experiment_id, namespace = asyncio.run(_get_user_info(redis_provider, sandbox_id)) + user_id, experiment_id, namespace = asyncio.run(_get_user_info(meta_store, sandbox_id)) # Build attributes attributes = _build_attributes(op_name, sandbox_id, f, user_id, experiment_id, namespace) diff --git a/tests/unit/admin/metrics/test_decorator.py b/tests/unit/admin/metrics/test_decorator.py index 61fecf357e..247499bcf6 100644 --- a/tests/unit/admin/metrics/test_decorator.py +++ b/tests/unit/admin/metrics/test_decorator.py @@ -9,9 +9,10 @@ _get_user_info, _record_metrics, _update_sandbox_id_from_result, + monitor_sandbox_operation, ) from rock.admin.metrics.monitor import MetricsMonitor -from rock.utils.providers import RedisProvider +from rock.sandbox.sandbox_meta_store import SandboxMetaStore class SampleObject: @@ -63,35 +64,23 @@ def test_extract_sandbox_id_prefers_container_name_over_sandbox_id(): assert result == "container-name" -@patch("rock.admin.metrics.decorator.alive_sandbox_key") -def test_get_user_info_success(mock_alive_key): - mock_redis_provider = Mock(spec=RedisProvider) - - async def async_mock_return_value(*args, **kwargs): - return [{"user_id": "user123", "experiment_id": "exp456", "namespace": "ns789"}] - - mock_alive_key.return_value = "alive:test-sandbox" - mock_redis_provider.json_get = AsyncMock(side_effect=async_mock_return_value) +def test_get_user_info_success(): + mock_meta_store = AsyncMock(spec=SandboxMetaStore) + mock_meta_store.get = AsyncMock( + return_value={"user_id": "user123", "experiment_id": "exp456", "namespace": "ns789"} + ) - # Run the async function in a blocking way - user_id, experiment_id, namespace = asyncio.run(_get_user_info(mock_redis_provider, "test-sandbox")) + user_id, experiment_id, namespace = asyncio.run(_get_user_info(mock_meta_store, "test-sandbox")) assert user_id == "user123" assert experiment_id == "exp456" assert namespace == "ns789" -@patch("rock.admin.metrics.decorator.alive_sandbox_key") -def test_get_user_info_no_data(mock_alive_key): - mock_redis_provider = Mock(spec=RedisProvider) - - async def async_mock_return_value(*args, **kwargs): - return [] +def test_get_user_info_no_data(): + mock_meta_store = AsyncMock(spec=SandboxMetaStore) + mock_meta_store.get = AsyncMock(return_value=None) - mock_alive_key.return_value = "alive:test-sandbox" - mock_redis_provider.json_get = AsyncMock(side_effect=async_mock_return_value) - - # Run the async function in a blocking way - user_id, experiment_id, namespace = asyncio.run(_get_user_info(mock_redis_provider, "test-sandbox")) + user_id, experiment_id, namespace = asyncio.run(_get_user_info(mock_meta_store, "test-sandbox")) assert user_id == "default" assert experiment_id == "default" assert namespace == "default" @@ -165,3 +154,35 @@ def test_record_metrics_failure(): # rt should be 1000ms (1.0 - 0) * 1000 mock_metrics_monitor.record_gauge_by_name.assert_called_once_with("test.rt", 1000.0, error_attrs) mock_metrics_monitor.record_counter_by_name.assert_any_call("test.total", 1, error_attrs) + + +async def test_decorator_retrieves_user_info_from_meta_store(redis_provider, _memory_sandbox_table): + """monitor_sandbox_operation should read user info via self._meta_store.get(), + not self._redis_provider (which doesn't exist on SandboxManager/SandboxProxyService). + Before the fix, user_id/experiment_id/namespace were always 'default'. + """ + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) + + # Seed Redis with sandbox info containing user fields + sandbox_id = "test-sandbox-123" + await meta_store.create(sandbox_id, {"user_id": "u1", "experiment_id": "e1", "namespace": "n1"}) + + service = Mock() + service._meta_store = meta_store + service.metrics_monitor = Mock(spec=MetricsMonitor) + + @monitor_sandbox_operation() + async def do_something(self, sandbox_id): + return "ok" + + result = await do_something(service, sandbox_id) + assert result == "ok" + + # Verify metrics recorded with real user info, not all-default + calls = service.metrics_monitor.record_counter_by_name.call_args_list + success_calls = [c for c in calls if c[0][0] == "request.success"] + assert len(success_calls) == 1 + attrs = success_calls[0][0][2] + assert attrs["user_id"] == "u1" + assert attrs["experiment_id"] == "e1" + assert attrs["namespace"] == "n1" From 8fd7f666e358c761a6023bd41c5510d39d7d8453 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Thu, 7 May 2026 19:08:49 +0800 Subject: [PATCH 080/226] fix(metrics): read sandbox image from meta_store instead of in-memory dict (#913) * test(metrics): add test for _collect_sandbox_meta image lookup Verifies that _collect_sandbox_meta reads image from meta_store (not the removed in-memory dict). Covers: correct image, missing image fallback to default, and empty sandbox list. * fix(metrics): read sandbox image from meta_store instead of in-memory dict _collect_sandbox_meta read image from the ephemeral _sandbox_meta dict, which was never populated by start_async (dropped during the operator refactor in e32a9389). Replace with a meta_store.iter_alive_sandbox_info() call. Fixes: e32a9389c537 ("refactor: delegate sandbox lifecycle management operations to operator #277 (#423)") --- rock/sandbox/base_manager.py | 7 +- rock/sandbox/sandbox_manager.py | 6 -- rock/sandbox/sandbox_meta_store.py | 5 ++ .../unit/sandbox/test_collect_sandbox_meta.py | 77 +++++++++++++++++++ 4 files changed, 85 insertions(+), 10 deletions(-) create mode 100644 tests/unit/sandbox/test_collect_sandbox_meta.py diff --git a/rock/sandbox/base_manager.py b/rock/sandbox/base_manager.py index f361b33c1f..c89e731aef 100644 --- a/rock/sandbox/base_manager.py +++ b/rock/sandbox/base_manager.py @@ -36,7 +36,6 @@ def __init__( ) self._report_interval = 10 self._check_job_interval = 180 - self._sandbox_meta = {} self._setup_scheduler() self.deployment_manager = DeploymentManager(rock_config, enable_runtime_auto_clear) @@ -128,10 +127,10 @@ async def _collect_system_resource_metrics(self): async def _collect_sandbox_meta(self) -> tuple[int, dict[str, dict[str, str]]]: meta: dict = {} cnt = 0 - async for sandbox_id in self._meta_store.iter_alive_sandbox_ids(): + async for sandbox_info in self._meta_store.iter_alive_sandbox_info(): cnt += 1 - image = self._sandbox_meta.get(sandbox_id, {}).get("image", "default") - meta[sandbox_id] = {"image": image} + image = sandbox_info.get("image", "default") + meta[sandbox_info.get("sandbox_id")] = {"image": image} return cnt, meta def stop_monitoring(self): diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index a66cc13816..d796b08ecf 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -157,8 +157,6 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: await asyncio.sleep(1) await self.get_status(sandbox_id) - self._sandbox_meta[sandbox_id] = {"image": docker_deployment_config.image} - return SandboxStartResponse( sandbox_id=sandbox_id, host_name=await self._ray_service.async_ray_get(sandbox_actor.host_name.remote()), @@ -181,10 +179,6 @@ async def stop(self, sandbox_id): logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) await self._meta_store.archive(sandbox_id, sandbox_info) return - try: - self._sandbox_meta.pop(sandbox_id) - except KeyError: - logger.debug(f"{sandbox_id} key not found") logger.info(f"sandbox {sandbox_id} stopped") await self._meta_store.archive(sandbox_id, sandbox_info) diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index d392de1027..63eeb5512a 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -139,6 +139,11 @@ async def iter_alive_sandbox_ids(self) -> AsyncIterator[str]: if sandbox_id: yield sandbox_id + async def iter_alive_sandbox_info(self) -> AsyncIterator[SandboxInfo]: + """Yield active sandbox info from the DB.""" + for sandbox_info in await self._db.list_by_in("state", _ACTIVE_STATES): + yield sandbox_info + @monitor_metastore_operation async def batch_get(self, sandbox_ids: list[str]) -> list[SandboxInfo]: """Fetch sandbox info for multiple IDs from the DB. Missing IDs are omitted.""" diff --git a/tests/unit/sandbox/test_collect_sandbox_meta.py b/tests/unit/sandbox/test_collect_sandbox_meta.py new file mode 100644 index 0000000000..c2baff5db2 --- /dev/null +++ b/tests/unit/sandbox/test_collect_sandbox_meta.py @@ -0,0 +1,77 @@ +"""RED-GREEN test: _collect_sandbox_meta must read image from meta_store. + +Before the fix, BaseManager._collect_sandbox_meta read image from the +in-memory dict ``_sandbox_meta`` which was never populated by ``start_async``, +so image was always "default". After the fix it reads from meta_store. +""" + +from unittest.mock import patch + +import pytest + +from rock.actions.sandbox.response import State +from rock.admin.core.sandbox_table import SandboxTable +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.utils.providers.redis_provider import RedisProvider + +SANDBOX_ID = "sbx-image-001" +EXPECTED_IMAGE = "python:3.11-slim" + +SANDBOX_INFO = { + "sandbox_id": SANDBOX_ID, + "image": EXPECTED_IMAGE, + "user_id": "u1", + "state": State.RUNNING, + "host_ip": "10.0.0.1", +} + + +@pytest.fixture +def meta_store(redis_provider: RedisProvider, _memory_sandbox_table: SandboxTable): + return SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) + + +@pytest.fixture +def base_manager(meta_store): + """Minimal BaseManager with real meta_store, everything else mocked.""" + from rock.sandbox.base_manager import BaseManager + + with patch.object(BaseManager, "__init__", lambda self, *a, **kw: None): + mgr = BaseManager.__new__(BaseManager) + mgr._meta_store = meta_store + mgr._sandbox_meta = {} # old code needs this; proves empty dict → "default" + return mgr + + +class TestCollectSandboxMeta: + async def test_image_read_from_meta_store(self, base_manager, meta_store): + """_collect_sandbox_meta should return the actual image stored in meta_store, + not 'default'. This test fails on the old code (empty _sandbox_meta dict).""" + await meta_store.create(SANDBOX_ID, SANDBOX_INFO) + + cnt, meta = await base_manager._collect_sandbox_meta() + + assert cnt == 1 + assert SANDBOX_ID in meta + assert meta[SANDBOX_ID]["image"] == EXPECTED_IMAGE + + async def test_missing_image_falls_back_to_default(self, base_manager, meta_store): + """When sandbox_info has no image field, should fall back to 'default'.""" + info_no_image = { + "sandbox_id": SANDBOX_ID, + "state": State.RUNNING, + "host_ip": "10.0.0.1", + } + await meta_store.create(SANDBOX_ID, info_no_image) + + cnt, meta = await base_manager._collect_sandbox_meta() + + assert cnt == 1 + assert meta[SANDBOX_ID]["image"] == "default" + + async def test_no_sandboxes_returns_empty(self, base_manager): + """No alive sandboxes should return zero count and empty dict.""" + cnt, meta = await base_manager._collect_sandbox_meta() + + assert cnt == 0 + assert meta == {} From 537ba67654e107fe87c91da59d1b8d2e76639d03 Mon Sep 17 00:00:00 2001 From: GreenHand Date: Fri, 8 May 2026 12:07:15 +0800 Subject: [PATCH 081/226] fix(sdk): add runtime config type validation in PythonRuntimeEnv (#652) Add explicit type check and conversion for runtime_env_config to ensure it is properly validated as PythonRuntimeEnvConfig before processing. --- rock/sdk/sandbox/runtime_env/python_runtime_env.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rock/sdk/sandbox/runtime_env/python_runtime_env.py b/rock/sdk/sandbox/runtime_env/python_runtime_env.py index 19a9549800..0cf904e63a 100644 --- a/rock/sdk/sandbox/runtime_env/python_runtime_env.py +++ b/rock/sdk/sandbox/runtime_env/python_runtime_env.py @@ -74,6 +74,8 @@ def __init__( version = runtime_env_config.version if version not in ("3.11", "3.12", "default"): raise ValueError(f"Unsupported Python version: {version}. Supported versions: 3.11, 3.12, default") + if not isinstance(runtime_env_config, PythonRuntimeEnvConfig): + runtime_env_config = PythonRuntimeEnvConfig.model_validate(runtime_env_config.model_dump()) # Create base config with resolved version (extra="ignore" handles 'pip' and 'pip_index_url' fields) super().__init__(sandbox=sandbox, runtime_env_config=runtime_env_config) From 1ffa28db19551461e23cd52a6ba1a86f908334d2 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Sat, 9 May 2026 10:26:28 +0800 Subject: [PATCH 082/226] feat(scheduler): add dynamic config reloading via Nacos #888 (#889) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * add v1.0.4 doc * chore: update docs version config (#431) * scheduler conf to nacos * Feature/openclaw demo (#457) * chore: update redis image version in sandbox demo * refactor: remove redundant validation method in RockAgentConfig * feat: add OpenClaw demo guide and config files * docs: add python sdk references docs for version 1.3.x (#460) * Doc/v1.2.1 0210 (#467) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * add v1.2.1 doc * scheduler conf to nacos * Add callback function abstraction to Nacos; load scheduler process in pod-1 * feat: add k8s operator * feat: validate and sync experiment_id/namespace in JobConfig (#716) (#717) * feat: refine expr_id and namespace * feat(sdk): add JobConfig enhancements and fix linting issues Add job configuration improvements including experiment_id support, OSS mirror config updates, and port validation changes. Fix ruff lint and format issues across the codebase. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use urlparse for URL hostname validation in speedup tests Replace substring-based URL checks with urlparse().hostname to satisfy CodeQL's "Incomplete URL substring sanitization" rule. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: extract domain constants to avoid CodeQL URL substring warnings CodeQL flags `"domain" in var` as incomplete URL sanitization regardless of variable type. Extract hostnames to constants and use helper functions to eliminate the flagged pattern entirely. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) * scheduler config to nacod * scheduler to nacos * revert other changes * rm blank dir * refactor(scheduler): incremental task rebuild + non-idempotent cleanup - Rewrite _rebuild_tasks to diff old vs new task configs via per-task md5 hash; only removed/added/changed tasks are touched, unchanged tasks keep their schedule continuity instead of being wiped on every Nacos reload. - Consolidate live-task source of truth onto TaskScheduler._tasks_by_class; delete the TaskRegistry class (no readers in the codebase, only writers). - Add BaseTask.cleanup / cleanup_on_worker / _clear_task_status. For NON_IDEMPOTENT tasks, _uninstall_task now awaits task.cleanup(workers) to pkill -9 the recorded daemon PID + children and rm the per-worker status file, so changing a daemon task's parameters cleanly stops the old process before a new one is installed. --------- Co-authored-by: dengsheng Co-authored-by: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Co-authored-by: junxin Co-authored-by: dengwx Co-authored-by: Claude Opus 4.6 (1M context) --- rock/admin/main.py | 11 +- rock/admin/scheduler/scheduler.py | 171 ++++++++++++++++++++----- rock/admin/scheduler/task_base.py | 41 ++++++ rock/admin/scheduler/task_factory.py | 35 +---- rock/admin/scheduler/task_registry.py | 23 ---- rock/utils/providers/nacos_provider.py | 11 +- uv.lock | 2 +- 7 files changed, 205 insertions(+), 89 deletions(-) delete mode 100644 rock/admin/scheduler/task_registry.py diff --git a/rock/admin/main.py b/rock/admin/main.py index e44b6c427e..24824d7e7d 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -21,7 +21,7 @@ from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service from rock.admin.scheduler.scheduler import SchedulerThread -from rock.config import DatabaseConfig, RockConfig +from rock.config import DatabaseConfig, RockConfig, SchedulerConfig from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager from rock.sandbox.operator.factory import OperatorContext, OperatorFactory @@ -51,6 +51,14 @@ async def lifespan(app: FastAPI): else env_vars.ROCK_CONFIG ) rock_config = RockConfig.from_env(config_file_path) + + # Override scheduler config from Nacos if available + if rock_config.nacos_provider: + nacos_config = await rock_config.nacos_provider.get_config() + if nacos_config and "scheduler" in nacos_config: + rock_config.scheduler = SchedulerConfig(**nacos_config["scheduler"]) + logger.info(f"Overrode scheduler config from Nacos with {len(rock_config.scheduler.tasks)} tasks") + env_vars.ROCK_ADMIN_ENV = args.env env_vars.ROCK_ADMIN_ROLE = args.role @@ -128,6 +136,7 @@ async def lifespan(app: FastAPI): if rock_config.scheduler.enabled and is_primary_pod(): scheduler_thread = SchedulerThread( scheduler_config=rock_config.scheduler, + nacos_provider=rock_config.nacos_provider, ) scheduler_thread.start() logger.info("Scheduler thread started on primary pod") diff --git a/rock/admin/scheduler/scheduler.py b/rock/admin/scheduler/scheduler.py index d890b800c2..cf0b8122e3 100644 --- a/rock/admin/scheduler/scheduler.py +++ b/rock/admin/scheduler/scheduler.py @@ -1,19 +1,23 @@ # rock/admin/scheduler/scheduler.py import asyncio +import hashlib +import json import threading import time +from dataclasses import asdict from datetime import datetime, timedelta import pytz import ray +import yaml from apscheduler.schedulers.asyncio import AsyncIOScheduler from rock import env_vars from rock.admin.scheduler.task_base import BaseTask -from rock.admin.scheduler.task_registry import TaskRegistry from rock.common.constants import SCHEDULER_LOG_NAME -from rock.config import SchedulerConfig +from rock.config import SchedulerConfig, TaskConfig from rock.logger import init_logger +from rock.utils.providers import NacosConfigProvider logger = init_logger(name="scheduler", file_name=SCHEDULER_LOG_NAME) @@ -65,15 +69,23 @@ def get_alive_workers(self, force_refresh: bool = False) -> list[str]: class TaskScheduler: - """Manages task scheduling using APScheduler.""" + """Manages task scheduling using APScheduler with optional Nacos dynamic config.""" - def __init__(self, scheduler_config: SchedulerConfig): + def __init__( + self, + scheduler_config: SchedulerConfig, + nacos_provider: NacosConfigProvider | None = None, + ): self.scheduler_config = scheduler_config self.local_tz = pytz.timezone(env_vars.ROCK_TIME_ZONE) self._scheduler: AsyncIOScheduler | None = None self._stop_event: asyncio.Event | None = None self._worker_cache: WorkerIPCache | None = None - self._loop: asyncio.AbstractEventLoop | None = None + self._nacos_provider = nacos_provider + self._event_loop: asyncio.AbstractEventLoop | None = None + self._last_scheduler_config_hash: str | None = None + self._task_hashes: dict[str, str] = {} + self._tasks_by_class: dict[str, BaseTask] = {} def _init_worker_cache(self) -> None: """Initialize the worker IP cache.""" @@ -81,11 +93,116 @@ def _init_worker_cache(self) -> None: cache_ttl=self.scheduler_config.worker_cache_ttl, ) - def _register_tasks(self) -> None: - """Register all tasks from configuration.""" + def _on_nacos_config_changed(self, new_config: dict) -> None: + """Callback invoked by Nacos watcher when config changes (runs in Nacos polling thread).""" + logger.info("Nacos config changed, checking scheduler section...") + try: + config_dict = yaml.safe_load(new_config["content"]) + if not config_dict or "scheduler" not in config_dict: + logger.warning("No 'scheduler' section in updated Nacos config, skipping") + return + + # Compare scheduler section hash to avoid unnecessary reloads + scheduler_raw = json.dumps(config_dict["scheduler"], sort_keys=True) + config_hash = hashlib.md5(scheduler_raw.encode()).hexdigest() + if config_hash == self._last_scheduler_config_hash: + logger.info("Scheduler config unchanged, skipping reload") + return + self._last_scheduler_config_hash = config_hash + + new_scheduler_config = SchedulerConfig(**config_dict["scheduler"]) + # Schedule the async reload on the event loop (thread-safe) + if self._event_loop and self._event_loop.is_running(): + asyncio.run_coroutine_threadsafe(self._reload_scheduler_config(new_scheduler_config), self._event_loop) + else: + logger.warning("Event loop not available, cannot reload tasks dynamically") + except yaml.YAMLError as e: + logger.error(f"Failed to parse updated Nacos YAML config: {e}") + except Exception as e: + logger.error(f"Failed to process Nacos config change: {e}") + + async def _reload_scheduler_config(self, new_scheduler_config: SchedulerConfig) -> None: + """Reload scheduler config by clearing and rebuilding all tasks.""" + self.scheduler_config = new_scheduler_config + self._worker_cache.cache_ttl = new_scheduler_config.worker_cache_ttl + await self._rebuild_tasks() + + @staticmethod + def _compute_task_hash(task_config: TaskConfig) -> str: + raw = json.dumps(asdict(task_config), sort_keys=True) + return hashlib.md5(raw.encode()).hexdigest() + + def _install_task(self, task_config: TaskConfig) -> None: + """Create, register, and schedule a single task.""" from rock.admin.scheduler.task_factory import TaskFactory - TaskFactory.register_all_tasks(self.scheduler_config) + try: + task = TaskFactory.create_task(task_config) + except Exception as e: + logger.error(f"Failed to create task '{task_config.task_class}': {e}") + return + + self._tasks_by_class[task_config.task_class] = task + self._task_hashes[task_config.task_class] = self._compute_task_hash(task_config) + self._scheduler.add_job( + self._run_task, + trigger="interval", + seconds=task.interval_seconds, + args=[task], + id=task.type, + name=task.type, + replace_existing=True, + next_run_time=datetime.now(self.local_tz) + timedelta(seconds=2), + ) + logger.info(f"Installed task '{task.type}' with interval {task.interval_seconds}s") + + async def _uninstall_task(self, task_class: str) -> None: + """Remove a single task from the scheduler and clean up its worker-side processes.""" + task = self._tasks_by_class.pop(task_class, None) + self._task_hashes.pop(task_class, None) + if task is None: + return + try: + self._scheduler.remove_job(task.type) + except Exception as e: + logger.warning(f"Failed to remove scheduler job '{task.type}': {e}") + if self._worker_cache is not None: + worker_ips = self._worker_cache.get_alive_workers() + if worker_ips: + await task.cleanup(worker_ips) + logger.info(f"Uninstalled task '{task.type}'") + + async def _rebuild_tasks(self) -> None: + """Apply config changes by diffing old vs new tasks; only touch the ones that changed.""" + if not self.scheduler_config.enabled: + for task_class in list(self._tasks_by_class): + await self._uninstall_task(task_class) + logger.info("Scheduler disabled, all tasks removed") + return + + new_by_class: dict[str, TaskConfig] = {} + new_hashes: dict[str, str] = {} + for task_config in self.scheduler_config.tasks: + if not task_config.enabled or not task_config.task_class: + continue + new_by_class[task_config.task_class] = task_config + new_hashes[task_config.task_class] = self._compute_task_hash(task_config) + + old_keys = set(self._task_hashes) + new_keys = set(new_hashes) + removed = old_keys - new_keys + added = new_keys - old_keys + changed = {k for k in (old_keys & new_keys) if self._task_hashes[k] != new_hashes[k]} + + for task_class in removed | changed: + await self._uninstall_task(task_class) + for task_class in added | changed: + self._install_task(new_by_class[task_class]) + + if removed or added or changed: + logger.info(f"Scheduler tasks updated: removed={len(removed)}, added={len(added)}, changed={len(changed)}") + else: + logger.info("No task changes detected") async def _run_task(self, task: BaseTask) -> None: """Run a single task on alive workers.""" @@ -99,28 +216,18 @@ async def _run_task(self, task: BaseTask) -> None: except Exception as e: logger.error(f"Task '{task.type}' failed: {e}") - def _add_jobs(self) -> None: - """Add all registered tasks as scheduler jobs.""" - for task in TaskRegistry.get_all_tasks().values(): - self._scheduler.add_job( - self._run_task, - trigger="interval", - seconds=task.interval_seconds, - args=[task], - id=task.type, - name=task.type, - replace_existing=True, - next_run_time=datetime.now(self.local_tz) + timedelta(seconds=2), - ) - logger.info(f"Added job '{task.type}' with interval {task.interval_seconds}s") - async def run(self) -> None: """Run the scheduler until stopped.""" + self._event_loop = asyncio.get_running_loop() + self._init_worker_cache() - self._register_tasks() + + if self._nacos_provider: + self._nacos_provider.add_listener(self._on_nacos_config_changed) + logger.info("Nacos dynamic config listener registered for scheduler") self._scheduler = AsyncIOScheduler(timezone=self.local_tz) - self._add_jobs() + await self._rebuild_tasks() # Pre-cache worker IPs before starting self._worker_cache.refresh() @@ -129,7 +236,6 @@ async def run(self) -> None: logger.info("Scheduler started") self._stop_event = asyncio.Event() - self._loop = asyncio.get_event_loop() try: await self._stop_event.wait() @@ -141,22 +247,27 @@ async def run(self) -> None: def stop(self) -> None: """Thread-safe stop: signal the scheduler to shut down.""" - if self._stop_event and self._loop: - self._loop.call_soon_threadsafe(self._stop_event.set) + if self._stop_event and self._event_loop: + self._event_loop.call_soon_threadsafe(self._stop_event.set) class SchedulerThread: """Scheduler thread manager - runs APScheduler in a daemon thread with its own event loop.""" - def __init__(self, scheduler_config: SchedulerConfig): + def __init__( + self, + scheduler_config: SchedulerConfig, + nacos_provider: NacosConfigProvider | None = None, + ): self.scheduler_config = scheduler_config + self.nacos_provider = nacos_provider self._thread: threading.Thread | None = None self._task_scheduler: TaskScheduler | None = None def _run_scheduler_in_thread(self) -> None: """Entry point for running scheduler in a thread with a dedicated event loop.""" try: - self._task_scheduler = TaskScheduler(self.scheduler_config) + self._task_scheduler = TaskScheduler(self.scheduler_config, self.nacos_provider) asyncio.run(self._task_scheduler.run()) except Exception: logger.exception("Scheduler thread encountered an error") diff --git a/rock/admin/scheduler/task_base.py b/rock/admin/scheduler/task_base.py index 087f52e77f..0c6ef29d18 100644 --- a/rock/admin/scheduler/task_base.py +++ b/rock/admin/scheduler/task_base.py @@ -171,6 +171,47 @@ async def save_task_status(self, runtime: RemoteSandboxRuntime, status: TaskStat """Save task status to worker file.""" await runtime.write_file(WriteFileRequest(path=self.status_file_path, content=status.to_json())) + async def _clear_task_status(self, runtime: RemoteSandboxRuntime) -> None: + """Remove the status file from worker.""" + await runtime.execute(Command(command=f"rm -f {self.status_file_path}", shell=True)) + + async def cleanup_on_worker(self, ip: str) -> None: + """Stop any long-running process spawned by this task on a single worker. + + For idempotent tasks this is a no-op (no daemon process to kill). + """ + if self.idempotency == IdempotencyType.IDEMPOTENT: + return + runtime = self._get_runtime(ip) + status = await self.get_task_status(runtime) + if status is None or not status.pid: + return + if await runtime.check_pid_exists(status.pid): + kill_cmd = f"pkill -9 -P {status.pid}; kill -9 {status.pid}" + await runtime.execute(Command(command=kill_cmd, shell=True)) + logger.info(f"[{self.type}] killed pid {status.pid} on worker[{ip}]") + await self._clear_task_status(runtime) + + async def cleanup(self, worker_ips: list[str], max_concurrency: int = 50) -> None: + """Cleanup task across all workers, parallel and best-effort. + + Idempotent tasks return immediately. For non-idempotent tasks, kills the + recorded daemon process and clears the status file on each worker. Failures + on individual workers are logged but do not propagate. + """ + if self.idempotency == IdempotencyType.IDEMPOTENT: + return + semaphore = asyncio.Semaphore(max_concurrency) + + async def cleanup_with_limit(ip: str) -> None: + async with semaphore: + try: + await self.cleanup_on_worker(ip) + except Exception as e: + logger.warning(f"[{self.type}] cleanup failed on worker[{ip}]: {e}") + + await asyncio.gather(*[cleanup_with_limit(ip) for ip in worker_ips]) + async def should_run(self, runtime: RemoteSandboxRuntime) -> bool: """Determine if the task should be run.""" if self.idempotency == IdempotencyType.IDEMPOTENT: diff --git a/rock/admin/scheduler/task_factory.py b/rock/admin/scheduler/task_factory.py index 5130f264c3..ceea829eba 100644 --- a/rock/admin/scheduler/task_factory.py +++ b/rock/admin/scheduler/task_factory.py @@ -2,16 +2,11 @@ import importlib from rock.admin.scheduler.task_base import BaseTask -from rock.admin.scheduler.task_registry import TaskRegistry -from rock.common.constants import SCHEDULER_LOG_NAME -from rock.config import SchedulerConfig, TaskConfig -from rock.logger import init_logger - -logger = init_logger("task_factory", file_name=SCHEDULER_LOG_NAME) +from rock.config import TaskConfig class TaskFactory: - """Task factory - dynamically creates and registers tasks from config.""" + """Task factory - dynamically creates tasks from config.""" @staticmethod def _load_task_class(class_path: str) -> type[BaseTask]: @@ -39,29 +34,5 @@ def create_task(cls, task_config: TaskConfig) -> BaseTask: Returns: Task instance """ - # Dynamically load task class task_class = cls._load_task_class(task_config.task_class) - - # Create task instance with config params - task = task_class.from_config(task_config) - - return task - - @classmethod - def register_all_tasks(cls, scheduler_config: SchedulerConfig): - """Register all enabled tasks from config.""" - for task_config in scheduler_config.tasks: - if not task_config.enabled: - logger.info(f"Task '{task_config.task_class}' is disabled, skipping") - continue - - if not task_config.task_class: - logger.warning(f"Task '{task_config.task_class}' has no task_class, skipping") - continue - - try: - task = cls.create_task(task_config) - TaskRegistry.register(task) - logger.info(f"Registered task '{task.type}' with interval {task.interval_seconds}s") - except Exception as e: - logger.error(f"Failed to create task '{task_config.task_class}': {e}") + return task_class.from_config(task_config) diff --git a/rock/admin/scheduler/task_registry.py b/rock/admin/scheduler/task_registry.py deleted file mode 100644 index be46256eb5..0000000000 --- a/rock/admin/scheduler/task_registry.py +++ /dev/null @@ -1,23 +0,0 @@ -# rock/admin/scheduler/task_registry.py -from rock.admin.scheduler.task_base import BaseTask - - -class TaskRegistry: - """Task registry for managing scheduled tasks.""" - - _tasks: dict[str, BaseTask] = {} - - @classmethod - def register(cls, task: BaseTask): - """Register a task.""" - cls._tasks[task.type] = task - - @classmethod - def get_task(cls, name: str) -> BaseTask: - """Get a task by name.""" - return cls._tasks.get(name) - - @classmethod - def get_all_tasks(cls) -> dict[str, BaseTask]: - """Get all registered tasks.""" - return cls._tasks.copy() diff --git a/rock/utils/providers/nacos_provider.py b/rock/utils/providers/nacos_provider.py index 1b614e805a..b6ca4c966d 100644 --- a/rock/utils/providers/nacos_provider.py +++ b/rock/utils/providers/nacos_provider.py @@ -62,15 +62,22 @@ def _update_callback(self, new_config: dict): except yaml.YAMLError as e: logger.error(f"Failed to parse updated YAML config: {e}") - def add_listener(self): + def add_listener(self, callback=None): """ Add a listener for the configuration to implement hot reloading. + + Args: + callback: Optional callback function to be called when config changes. + If not provided, uses the default _update_callback method. """ + if callback is None: + callback = self._update_callback + try: self.client.add_config_watcher( data_id=self.data_id, group=self.group, - cb=self._update_callback, + cb=callback, ) logger.info(f"Added config watcher for data_id='{self.data_id}' group='{self.group}'.") except Exception as e: diff --git a/uv.lock b/uv.lock index f7f32dc941..e00a7f86b3 100644 --- a/uv.lock +++ b/uv.lock @@ -4035,7 +4035,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.6.0" +version = "1.7.0" source = { editable = "." } dependencies = [ { name = "anyio" }, From a4d45a39370de79adbdd17561ccf66722cb12aa1 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Mon, 11 May 2026 17:57:07 +0800 Subject: [PATCH 083/226] docs(1.7.x): clarify install-agent vs Job, reorganize examples/ (#926) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restructure 1.7.x docs (EN+ZH) to present install-agent and Job as two parallel agent-usage capabilities. Reorganize examples/ by capability so readers can find the right entry point. Docs: - Getting Started/rock-agent.md: Job-first, install-agent second. 1.1 Prepare yaml lists one Harbor (Terminal Bench) and one Bash (claw-eval) template — point to the canonical *.yaml.template files instead of inlining a fabricated example. - References/rock-agent.md: title -> "Install Agent in Sandbox (Experimental)"; reframe RockAgent as ROCK's mechanism for installing custom agents in a sandbox; positioning blurb makes clear that install-agent and Job use distinct config schemas (do not mix). - References/job.md (NEW): "Use Job to Run Agent" with end-to-end SDK example. The agents: field is documented as Harbor's own minimal schema (name, model_name) — not RockAgentConfig. - ZH run() flow numbering bug fixed. - Soften deprecation note on rock.sdk.bench.Job to "future release". Examples: - agents/ -> install-agents/ - harbor/ -> job/harbor/ - bash/ + evaluation/claw_eval/ -> job/bash/ - evaluation/swe_bench/ kept in place - READMEs added for install-agents/, job/, job/bash/, job/harbor/, evaluation/ - Fix REAMDE.md typo, drop __pycache__/ - Update inline path references in moved files Refs #925 Co-authored-by: Claude Opus 4.7 --- .../Getting Started/rock-agent.md | 89 ++++++----- .../References/Python SDK References/job.md | 150 ++++++++++++++++++ .../Python SDK References/rock-agent.md | 42 +++-- .../Getting Started/rock-agent.md | 90 ++++++----- .../References/Python SDK References/job.md | 150 ++++++++++++++++++ .../Python SDK References/rock-agent.md | 30 +++- examples/evaluation/README.md | 19 +++ examples/install-agents/README.md | 28 ++++ .../claude_code/claude_code_demo.py | 0 .../claude_code/rock_agent_config.yaml | 0 .../cursor_cli/cursor_cli_demo.py | 0 .../cursor_cli/rock_agent_config.yaml | 0 .../iflow_cli/iflow_cli_demo.py | 0 .../local/local_demo.py | 0 .../local/rock_agent_config.yaml | 0 .../proxy/proxy_demo.py | 0 .../proxy/rock_agent_config.yaml | 0 .../iflow_cli/rock_agent_config.yaml | 0 .../openclaw/README.md} | 2 +- .../openclaw/openclaw.json | 0 .../openclaw/openclaw_demo.py | 0 .../openclaw/rock_agent_config.yaml | 0 .../qwen_code/qwen_code_demo.py | 0 .../qwen_code/rock_agent_config.yaml | 0 .../swe_agent/rock_agent_config.yaml | 0 .../swe_agent/swe_agent_demo.py | 0 examples/job/README.md | 16 ++ examples/job/bash/README.md | 24 +++ .../claw_eval/claw_eval_bashjob.yaml.template | 0 .../claw_eval/claw_eval_config.yaml.template | 0 .../bash}/claw_eval/run_claw_eval.py | 2 +- .../bash}/claw_eval/run_claw_eval.sh | 0 .../{ => job}/bash/simple_bash_job_demo.sh | 0 examples/job/harbor/README.md | 27 ++++ examples/{ => job}/harbor/harbor_demo.py | 10 +- .../swe_job_config-verifier.yaml.template | 0 .../harbor/swe_job_config.yaml.template | 0 .../harbor/tb_job_config.yaml.template | 0 38 files changed, 569 insertions(+), 110 deletions(-) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/job.md create mode 100644 docs/versioned_docs/version-1.7.x/References/Python SDK References/job.md create mode 100644 examples/evaluation/README.md create mode 100644 examples/install-agents/README.md rename examples/{agents => install-agents}/claude_code/claude_code_demo.py (100%) rename examples/{agents => install-agents}/claude_code/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/cursor_cli/cursor_cli_demo.py (100%) rename examples/{agents => install-agents}/cursor_cli/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/iflow_cli/iflow_cli_demo.py (100%) rename examples/{agents => install-agents}/iflow_cli/integration_with_model_service/local/local_demo.py (100%) rename examples/{agents => install-agents}/iflow_cli/integration_with_model_service/local/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/iflow_cli/integration_with_model_service/proxy/proxy_demo.py (100%) rename examples/{agents => install-agents}/iflow_cli/integration_with_model_service/proxy/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/iflow_cli/rock_agent_config.yaml (100%) rename examples/{agents/openclaw/REAMDE.md => install-agents/openclaw/README.md} (98%) rename examples/{agents => install-agents}/openclaw/openclaw.json (100%) rename examples/{agents => install-agents}/openclaw/openclaw_demo.py (100%) rename examples/{agents => install-agents}/openclaw/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/qwen_code/qwen_code_demo.py (100%) rename examples/{agents => install-agents}/qwen_code/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/swe_agent/rock_agent_config.yaml (100%) rename examples/{agents => install-agents}/swe_agent/swe_agent_demo.py (100%) create mode 100644 examples/job/README.md create mode 100644 examples/job/bash/README.md rename examples/{evaluation => job/bash}/claw_eval/claw_eval_bashjob.yaml.template (100%) rename examples/{evaluation => job/bash}/claw_eval/claw_eval_config.yaml.template (100%) rename examples/{evaluation => job/bash}/claw_eval/run_claw_eval.py (95%) rename examples/{evaluation => job/bash}/claw_eval/run_claw_eval.sh (100%) rename examples/{ => job}/bash/simple_bash_job_demo.sh (100%) create mode 100644 examples/job/harbor/README.md rename examples/{ => job}/harbor/harbor_demo.py (88%) rename examples/{ => job}/harbor/swe_job_config-verifier.yaml.template (100%) rename examples/{ => job}/harbor/swe_job_config.yaml.template (100%) rename examples/{ => job}/harbor/tb_job_config.yaml.template (100%) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md index bd2dd69b05..1898fbdcd0 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/Getting Started/rock-agent.md @@ -4,70 +4,73 @@ sidebar_position: 4 # Rock Agent 快速启动 -Rock Agent 是 ROCK 提供的 AI Agent 运行框架，支持在沙箱环境中运行各种类型的 Agent。 +ROCK 提供两种并列的 agent 使用能力,各自适用不同场景: + +- **Job**:通过 BashJob / HarborJob 在 sandbox 里跑一次 agent 评测/任务(典型基准:SWE-bench、Terminal Bench),是入门主要场景。 +- **install-agent**:直接在单个沙箱里安装并运行 agent,适合本地开发、单次调试。 + +下面优先介绍 Job 用法,install-agent 用法见末尾或 [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md)。 ## 前置条件 -- 确保有可用的ROCK服务, 如果需要本地拉起服务端, 参考[快速启动](quickstart.md) +- 确保有可用的 ROCK 服务,如果需要本地拉起服务端,参考[快速启动](quickstart.md) -## 使用示例 +--- -ROCK 提供了两个Hello World Agent 示例，位于 `examples/agents/` 目录下： +## 一、用 Job 运行 Agent -``` -examples/agents/ -├── claude_code/ # ClaudeCode Agent 示例 -└── iflow_cli/ # IFlowCli Agent 示例 -``` +Job 有两种 backend:**Harbor Job** 用于运行 AI agent 基准评测任务(SWE-bench、Terminal Bench 等);**Bash Job** 用于在沙箱里跑自定义 shell 脚本。 -### 运行 IFlowCli 示例 +### 1.1 准备 yaml -```bash -cd examples/agents/iflow_cli -python iflow_cli_demo.py -``` +挑一类作为起点,直接复制对应模板: -### 运行 ClaudeCode 示例 +- Harbor Job(Terminal Bench):[`examples/job/harbor/tb_job_config.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor/tb_job_config.yaml.template) +- Bash Job(claw-eval):[`examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template) -```bash -cd examples/agents/claude_code -python claude_code_demo.py -``` +按模板填好对应字段即可。两类 Job 的完整字段说明见 [Use Job to Run Agent](../References/Python%20SDK%20References/job.md)。 -## IFlowCli 配置文件 +### 1.2 通过 Python SDK 启动 -配置文件位于 `examples/agents/iflow_cli/rock_agent_config.yaml`： +```python +import asyncio +from rock.sdk.job import Job, JobConfig -```yaml -run_cmd: "iflow -p ${prompt} --yolo" +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") + result = await Job(config).run() -runtime_env_config: - type: node - npm_registry: "https://registry.npmmirror.com" - custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") -env: - IFLOW_API_KEY: "" # 填入你的 API Key - IFLOW_BASE_URL: "" # 填入你的 Base URL - IFLOW_MODEL_NAME: "" # 填入你的模型名称 +asyncio.run(main()) ``` -## ClaudeCode 配置文件 +BashJob 的用法、完整字段说明、结果处理详见 [Use Job to Run Agent](../References/Python%20SDK%20References/job.md)。 -配置文件位于 `examples/agents/claude_code/rock_agent_config.yaml`： +--- -```yaml -run_cmd: "claude -p ${prompt}" +## 二、install-agent:在沙箱里安装并运行 Agent -runtime_env_config: - type: node - custom_install_cmd: "npm install -g @anthropic-ai/claude-code" +适合本地开发或单次调试 agent 的场景,核心 API: -env: - ANTHROPIC_BASE_URL: "" # 填入你的anthropic base url - ANTHROPIC_API_KEY: "" # 填入你的anthropic api key +```python +await sandbox.agent.install(config="rock_agent_config.yaml") +result = await sandbox.agent.run(prompt="hello") ``` -## 相关文档 +`examples/install-agents/` 下提供了多个开箱即用的示例: + +- `examples/install-agents/iflow_cli/` — IFlowCli +- `examples/install-agents/claude_code/` — Claude Code +- `examples/install-agents/cursor_cli/`、`qwen_code/`、`swe_agent/`、`openclaw/` — 其他 + +运行 Claude Code 示例: + +```bash +cd examples/install-agents/claude_code +python claude_code_demo.py +``` -- [RockAgent 参考](../References/Python%20SDK%20References/rock-agent.md) +完整 RockAgentConfig 字段说明、占位符语义、API 参考详见 [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/job.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/job.md new file mode 100644 index 0000000000..6dee1f678b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/job.md @@ -0,0 +1,150 @@ +# Use Job to Run Agent + +> 这是 ROCK 两种并列的 agent 使用能力中 **Job** 的参考文档,核心 API 是 `rock.sdk.job.Job` 与 `JobConfig`,用于在沙箱里跑一次 agent 评测/任务。有两种 backend:**Bash Job** 与 **Harbor Bench Job**。 +> +> 另一种能力是在单个沙箱里安装并运行 agent,见 [Install Agent in Sandbox](./rock-agent.md)。两种能力使用各自独立的配置 schema,**不要互相套用**。 + +`rock.sdk.job` 通过同一套 `Job` API 支持两种模式,通过配置类型区分: + +- **Bash Job**:在沙箱中运行自定义 Shell 脚本,适合数据处理、外部评测工具等 +- **Harbor Bench Job**:通过 Harbor 框架运行 AI agent 基准评测任务(SWE-bench、Terminal Bench 等) + +## 端到端示例 + +最小可跑通的 Python 用法: + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") # 含 agents: 与 datasets: + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +完整 yaml 模板参考 `examples/job/harbor/swe_job_config.yaml.template`。 + +--- + +## Bash Job + +Bash Job 适用于在沙箱内执行任意 Shell 脚本的场景,例如运行评测工具、数据处理流程等。 + +完整示例参考:[`examples/job/bash/claw_eval/`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval) + +- `run_claw_eval.py` — 主入口,演示 `JobConfig.from_yaml()` + `Job(config).run()` +- `claw_eval_bashjob.yaml.template` — YAML 配置模板,含 `script_path`、`environment`、`uploads`、`env` 等字段 +- `run_claw_eval.sh` — 沙箱内实际执行的脚本,演示 DinD 启动、日志写入和评分输出 + +### BashJobConfig 配置字段 + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `script` | `str \| None` | `None` | 内联脚本内容,与 `script_path` 二选一 | +| `script_path` | `str \| None` | `None` | 本地脚本文件路径,运行时读取并上传执行 | +| `job_name` | `str` | 当前时间戳 | 任务名称,用于日志和产物路径区分 | +| `environment` | `EnvironmentConfig` | — | 沙箱连接及资源配置,详见下表 | +| `namespace` | `str \| None` | `None` | 命名空间 | +| `experiment_id` | `str \| None` | `None` | 实验 ID | +| `timeout` | `int` | `7200` | 整体超时秒数(2 小时) | + +**`environment` 常用字段:** + +| 字段 | 类型 | 说明 | +|------|------|------| +| `image` | `str` | 沙箱 Docker 镜像 | +| `base_url` | `str` | ROCK 平台地址 | +| `xrl_authorization` | `str` | 鉴权 Token | +| `cluster` | `str` | 目标集群 | +| `memory` | `str` | 内存大小(如 `"64g"`) | +| `cpus` | `int` | CPU 核数 | +| `auto_stop` | `bool` | 任务完成后是否自动停止沙箱 | +| `uploads` | `list` | 本地文件/目录上传列表,格式:`[本地路径, 沙箱目标路径]` | +| `env` | `dict[str, str]` | 注入沙箱会话的环境变量 | + +--- + +## Harbor Bench Job + +Harbor Bench Job 适用于通过 Harbor 框架运行 AI agent 基准评测任务,如 SWE-bench、Terminal Bench 等。 + +> **注意**:`rock.sdk.bench.Job` 已废弃,将在未来移除。请改用 `rock.sdk.job.Job` + `HarborJobConfig`。 + +完整示例参考:[`examples/job/harbor/`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor) + +- `harbor_demo.py` — 主入口,演示 `JobConfig.from_yaml()` + `Job(config).run()` + 结果遍历 +- `swe_job_config.yaml.template` — SWE-bench 任务配置模板 +- `swe_job_config-verifier.yaml.template` — 附带 `verifier.mode: native` 的变体 +- `tb_job_config.yaml.template` — Terminal Bench 任务配置模板 + +### HarborJobConfig 核心配置字段 + +**基础字段:** + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `experiment_id` | `str` | 必填 | 实验 ID,Harbor 中必须提供 | +| `job_name` | `str \| None` | 自动生成 | 格式:`{dataset}_{task}_{uuid[:8]}` | +| `namespace` | `str \| None` | `None` | 命名空间,从沙箱自动反填 | +| `environment` | `RockEnvironmentConfig` | — | 沙箱连接及资源配置 | + +**执行控制字段:** + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `n_attempts` | `int` | `1` | 每个 Trial 的尝试次数 | +| `timeout` | `int` | `7200` | 整体超时秒数(自动从 agent_timeout 推算) | +| `debug` | `bool` | `False` | 调试模式,保留更多中间产物 | + +**组件字段:** + +| 字段 | 类型 | 说明 | +|------|------|------| +| `agents` | `list[AgentConfig]` | Harbor 框架自身的 agent 配置(典型字段:`name`、`model_name`),完整字段见 `examples/job/harbor/swe_job_config.yaml.template` | +| `datasets` | `list[DatasetConfig]` | 数据集配置列表 | +| `verifier` | `VerifierConfig` | Verifier 评测配置 | +| `orchestrator` | `OrchestratorConfig` | 并发调度配置 | + +--- + +## 结果处理 + +两种 Job 模式均返回 `JobResult`: + +```python +result = await Job(config).run() + +print(f"status={result.status}, score={result.score}") +for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + if trial.exception_info: + print(f" {trial.exception_info.exception_type}: {trial.exception_info.exception_message}") +``` + +### JobResult 字段 + +| 字段 / 属性 | 类型 | 说明 | +|------------|------|------| +| `status` | `JobStatus` | 任务整体状态 | +| `trial_results` | `list[TrialResult]` | 所有 Trial 结果列表 | +| `score` | `float`(属性) | 所有 Trial `score` 的平均值 | +| `n_completed` | `int`(属性) | 状态为 `completed` 的 Trial 数 | +| `n_failed` | `int`(属性) | 状态为 `failed` 的 Trial 数 | + +### TrialResult 字段 + +| 字段 / 属性 | 类型 | 说明 | +|------------|------|------| +| `task_name` | `str` | 任务名称 | +| `exit_code` | `int` | 进程退出码 | +| `raw_output` | `str` | 进程原始输出 | +| `exception_info` | `ExceptionInfo \| None` | 若有异常则填充 | +| `status` | `str`(属性) | `"completed"` 或 `"failed"` | +| `duration_sec` | `float`(属性) | 执行耗时(秒) | +| `score` | `float`(属性) | 评分(Bash Job 默认 `0.0`,Harbor 模式来自 verifier) | diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md index c3f03b1efc..cd252750aa 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/References/Python SDK References/rock-agent.md @@ -1,8 +1,12 @@ -# Rock Agent（实验性） +# Install Agent in Sandbox (Experimental) -RockAgent 是 ROCK 框架中的核心 Agent 实现，直接继承自 `Agent` 抽象基类。它提供了完整的 Agent 生命周期管理，包括环境初始化、ModelService 集成、命令执行等功能。 +> 这是 ROCK 两种并列的 agent 使用能力中 **install-agent** 的参考文档,核心 API 是 `sandbox.agent.install()` 与 `sandbox.agent.run(prompt)`,用于在单个沙箱里安装并运行 agent。 +> +> 另一种能力是用 Job 在沙箱里跑一次 agent 评测/任务,见 [Use Job to Run Agent](./job.md)。两种能力使用各自独立的配置 schema。 -使用 `sandbox.agent.install()` 以及 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 +RockAgent 是 ROCK 框架用来在沙箱中 install 自定义 agent 的能力,负责完整的 agent 生命周期管理:环境初始化、ModelService 集成、命令执行等。 + +使用 `sandbox.agent.install()` 与 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 ## 核心概念 @@ -185,10 +189,10 @@ model_service_config: # 具体参考 ModelService 有 执行 Agent 任务。 **执行流程**： -1. 替换占位符, 准备Agent 运行命令 -4. 启动 agent 进程 -5. 如果启用 ModelService，启动 `watch_agent` -6. 等待任务完成并返回结果 +1. 替换占位符,准备 Agent 运行命令 +2. 启动 agent 进程 +3. 如果启用 ModelService,启动 `watch_agent` +4. 等待任务完成并返回结果 ## 高级用法 @@ -281,10 +285,26 @@ model_service_config: ## 使用示例 -### 使用 YAML 配置文件（推荐） +### 使用 YAML 配置文件(推荐) ```python -# prepare a rock_agent_config.yaml -await sandbox.agent.install(config="rock_agent_config.yaml") -await sandbox.agent.run(prompt="hello") +import asyncio +from rock.sdk.sandbox import Sandbox, SandboxConfig + +async def main(): + sandbox = Sandbox(SandboxConfig()) + await sandbox.start() + try: + # rock_agent_config.yaml 与本文档「快速开始」中的示例一致 + await sandbox.agent.install(config="rock_agent_config.yaml") + result = await sandbox.agent.run(prompt="hello") + print(result) + finally: + await sandbox.stop() + +asyncio.run(main()) ``` + +更多开箱即用的示例参见 `examples/install-agents/`(Claude Code、IFlowCli、Cursor CLI、Qwen Code、SWE-agent、OpenClaw 等)。 + +如需通过 Job 跑 agent 评测/基准任务(另一条代码路径,有独立的配置 schema),见 [Use Job to Run Agent](./job.md)。 diff --git a/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md b/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md index eb6b54a65b..19369c1aff 100644 --- a/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md +++ b/docs/versioned_docs/version-1.7.x/Getting Started/rock-agent.md @@ -4,69 +4,73 @@ sidebar_position: 4 # Rock Agent Quick Start -Rock Agent is an AI Agent runtime framework provided by ROCK, supporting various types of Agents running in sandbox environments. +ROCK provides two parallel ways to use agents, each suited to a different scenario: + +- **Job**: Run an agent evaluation/task in a sandbox via BashJob / HarborJob (typical benchmarks: SWE-bench, Terminal Bench) — the primary entry point. +- **install-agent**: Install and run an agent directly inside a single sandbox — for local development and one-off debugging. + +Job is covered first. The install-agent section follows at the end, with full reference at [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md). ## Prerequisites -- Make sure you have a working ROCK service, if you need to locally start the service side, refer to [Quick Start](quickstart.md). -## Examples +- Make sure you have a working ROCK service. If you need to start the service locally, refer to [Quick Start](quickstart.md). -ROCK provides two Hello World Agent examples in the `examples/agents/` directory: +--- -``` -examples/agents/ -├── claude_code/ # ClaudeCode Agent example -└── iflow_cli/ # IFlowCli Agent example -``` +## 1. Use Job to Run Agent -### Run IFlowCli Example +Job has two backends: **Harbor Job** runs an AI agent benchmark task (SWE-bench, Terminal Bench, etc.); **Bash Job** runs a custom shell script inside a sandbox. -```bash -cd examples/agents/iflow_cli -python iflow_cli_demo.py -``` +### 1.1 Prepare a yaml -### Run ClaudeCode Example +Pick a starting point and copy the matching template: -```bash -cd examples/agents/claude_code -python claude_code_demo.py -``` +- Harbor Job (Terminal Bench): [`examples/job/harbor/tb_job_config.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor/tb_job_config.yaml.template) +- Bash Job (claw-eval): [`examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template) + +Fill in the fields per the template. See [Use Job to Run Agent](../References/Python%20SDK%20References/job.md) for the full field reference of both backends. -## IFlowCli Configuration File +### 1.2 Launch via Python SDK -The configuration file is located at `examples/agents/iflow_cli/rock_agent_config.yaml`: +```python +import asyncio +from rock.sdk.job import Job, JobConfig -```yaml -run_cmd: "iflow -p ${prompt} --yolo" +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") + result = await Job(config).run() -runtime_env_config: - type: node - npm_registry: "https://registry.npmmirror.com" - custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") -env: - IFLOW_API_KEY: "" # Enter your API key - IFLOW_BASE_URL: "" # Enter your base URL - IFLOW_MODEL_NAME: "" # Enter your model name +asyncio.run(main()) ``` -## ClaudeCode Configuration File +For BashJob usage, full field references, and result-handling details, see [Use Job to Run Agent](../References/Python%20SDK%20References/job.md). -The configuration file is located at `examples/agents/claude_code/rock_agent_config.yaml`: +--- -```yaml -run_cmd: "claude -p ${prompt}" +## 2. install-agent: Install and Run an Agent in a Sandbox -runtime_env_config: - type: node - custom_install_cmd: "npm install -g @anthropic-ai/claude-code" +For local development or debugging a single agent run, the core API is: -env: - ANTHROPIC_BASE_URL: "" # Enter your anthropic base url - ANTHROPIC_API_KEY: "" # Enter your anthropic api key +```python +await sandbox.agent.install(config="rock_agent_config.yaml") +result = await sandbox.agent.run(prompt="hello") ``` -## Related Documentation +The `examples/install-agents/` directory ships ready-to-run examples: + +- `examples/install-agents/iflow_cli/` — IFlowCli +- `examples/install-agents/claude_code/` — Claude Code +- `examples/install-agents/cursor_cli/`, `qwen_code/`, `swe_agent/`, `openclaw/` — others + +Run the Claude Code example: + +```bash +cd examples/install-agents/claude_code +python claude_code_demo.py +``` -- [RockAgent Reference](../References/Python%20SDK%20References/rock-agent.md) +For full RockAgentConfig field details, placeholder semantics, and API reference, see [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md). diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/job.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/job.md new file mode 100644 index 0000000000..a7dd7bcfd1 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/job.md @@ -0,0 +1,150 @@ +# Use Job to Run Agent + +> This is the reference for **Job**, one of ROCK's two parallel ways to use agents. Its core API is `rock.sdk.job.Job` with `JobConfig`, used to run an agent evaluation/task in a sandbox. Two backends are supported: **Bash Job** and **Harbor Bench Job**. +> +> The other way is to install and run an agent inside a single sandbox — see [Install Agent in Sandbox](./rock-agent.md). The two ways use distinct config schemas — **do not mix them**. + +`rock.sdk.job` exposes a single `Job` API that supports two modes, distinguished by the config type: + +- **Bash Job**: Runs an arbitrary shell script inside a sandbox — useful for data processing, external evaluation tools, etc. +- **Harbor Bench Job**: Runs an AI agent benchmark task via the Harbor framework (SWE-bench, Terminal Bench, etc.). + +## End-to-End Example + +A minimal runnable Python snippet: + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") # contains agents: and datasets: + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +The full yaml template is in `examples/job/harbor/swe_job_config.yaml.template`. + +--- + +## Bash Job + +Bash Job is for running arbitrary shell scripts inside a sandbox — running an external evaluation tool, processing data, etc. + +Full example: [`examples/job/bash/claw_eval/`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval) + +- `run_claw_eval.py` — Entry point demonstrating `JobConfig.from_yaml()` + `Job(config).run()` +- `claw_eval_bashjob.yaml.template` — YAML template with `script_path`, `environment`, `uploads`, `env`, etc. +- `run_claw_eval.sh` — The script that actually runs in the sandbox (DinD startup, log writing, score output) + +### BashJobConfig Fields + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `script` | `str \| None` | `None` | Inline script content (mutually exclusive with `script_path`) | +| `script_path` | `str \| None` | `None` | Local script path; the file is read and uploaded at runtime | +| `job_name` | `str` | current timestamp | Name used for log and artifact paths | +| `environment` | `EnvironmentConfig` | — | Sandbox connection and resource config (see below) | +| `namespace` | `str \| None` | `None` | Namespace | +| `experiment_id` | `str \| None` | `None` | Experiment ID | +| `timeout` | `int` | `7200` | Overall timeout in seconds (2 hours) | + +**Common `environment` fields:** + +| Field | Type | Description | +|------|------|-------------| +| `image` | `str` | Sandbox Docker image | +| `base_url` | `str` | ROCK platform URL | +| `xrl_authorization` | `str` | Auth token | +| `cluster` | `str` | Target cluster | +| `memory` | `str` | Memory size (e.g. `"64g"`) | +| `cpus` | `int` | Number of CPUs | +| `auto_stop` | `bool` | Whether to stop the sandbox after the job | +| `uploads` | `list` | Local-to-sandbox file/dir uploads, format: `[local_path, sandbox_path]` | +| `env` | `dict[str, str]` | Environment variables injected into the sandbox session | + +--- + +## Harbor Bench Job + +Harbor Bench Job runs AI agent benchmark tasks like SWE-bench and Terminal Bench via the Harbor framework. + +> **Note**: `rock.sdk.bench.Job` is deprecated and will be removed in a future release. Use `rock.sdk.job.Job` + `HarborJobConfig` instead. + +Full example: [`examples/job/harbor/`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor) + +- `harbor_demo.py` — Entry point demonstrating `JobConfig.from_yaml()` + `Job(config).run()` + result iteration +- `swe_job_config.yaml.template` — SWE-bench task config template +- `swe_job_config-verifier.yaml.template` — Variant with `verifier.mode: native` +- `tb_job_config.yaml.template` — Terminal Bench task config template + +### HarborJobConfig Core Fields + +**Basic fields:** + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `experiment_id` | `str` | required | Experiment ID — required by Harbor | +| `job_name` | `str \| None` | auto-generated | Format: `{dataset}_{task}_{uuid[:8]}` | +| `namespace` | `str \| None` | `None` | Namespace, auto-filled from the sandbox | +| `environment` | `RockEnvironmentConfig` | — | Sandbox connection and resource config | + +**Execution control:** + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `n_attempts` | `int` | `1` | Attempts per Trial | +| `timeout` | `int` | `7200` | Overall timeout (auto-derived from agent_timeout) | +| `debug` | `bool` | `False` | Debug mode — keeps more intermediate artifacts | + +**Components:** + +| Field | Type | Description | +|------|------|-------------| +| `agents` | `list[AgentConfig]` | Harbor's own agent config (typical fields: `name`, `model_name`) — see `examples/job/harbor/swe_job_config.yaml.template` for the canonical shape | +| `datasets` | `list[DatasetConfig]` | Dataset configs | +| `verifier` | `VerifierConfig` | Verifier evaluation config | +| `orchestrator` | `OrchestratorConfig` | Concurrency / scheduling config | + +--- + +## Result Handling + +Both Job modes return a `JobResult`: + +```python +result = await Job(config).run() + +print(f"status={result.status}, score={result.score}") +for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + if trial.exception_info: + print(f" {trial.exception_info.exception_type}: {trial.exception_info.exception_message}") +``` + +### JobResult Fields + +| Field / Property | Type | Description | +|------------------|------|-------------| +| `status` | `JobStatus` | Overall task status | +| `trial_results` | `list[TrialResult]` | List of all Trial results | +| `score` | `float` (property) | Average `score` across all Trials | +| `n_completed` | `int` (property) | Number of Trials with status `completed` | +| `n_failed` | `int` (property) | Number of Trials with status `failed` | + +### TrialResult Fields + +| Field / Property | Type | Description | +|------------------|------|-------------| +| `task_name` | `str` | Task name | +| `exit_code` | `int` | Process exit code | +| `raw_output` | `str` | Raw process output | +| `exception_info` | `ExceptionInfo \| None` | Populated if an exception occurred | +| `status` | `str` (property) | `"completed"` or `"failed"` | +| `duration_sec` | `float` (property) | Execution time in seconds | +| `score` | `float` (property) | Score (Bash Job defaults to `0.0`; Harbor mode comes from the verifier) | diff --git a/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md b/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md index f24ade49ac..9cc184d46c 100644 --- a/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md +++ b/docs/versioned_docs/version-1.7.x/References/Python SDK References/rock-agent.md @@ -1,6 +1,10 @@ -# Rock Agent (Experimental) +# Install Agent in Sandbox (Experimental) -RockAgent is the core Agent implementation in the ROCK framework, directly inheriting from the `Agent` abstract base class. It provides complete Agent lifecycle management, including environment initialization, ModelService integration, command execution, and more. +> This is the reference for **install-agent**, one of ROCK's two parallel ways to use agents. Its core API is `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, used to install and run an agent inside a single sandbox. +> +> The other way is to run an agent evaluation/task via Job — see [Use Job to Run Agent](./job.md). The two ways use distinct config schemas. + +RockAgent is the ROCK framework's mechanism for installing a custom agent inside a sandbox. It manages the full agent lifecycle — environment initialization, ModelService integration, command execution, and so on. Using `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, you can install and run Agents in the Sandbox environment provided by Rock. @@ -284,7 +288,23 @@ model_service_config: ### Using YAML Configuration File (Recommended) ```python -# prepare a rock_agent_config.yaml -await sandbox.agent.install(config="rock_agent_config.yaml") -await sandbox.agent.run(prompt="hello") +import asyncio +from rock.sdk.sandbox import Sandbox, SandboxConfig + +async def main(): + sandbox = Sandbox(SandboxConfig()) + await sandbox.start() + try: + # rock_agent_config.yaml matches the examples in "Quick Start" above + await sandbox.agent.install(config="rock_agent_config.yaml") + result = await sandbox.agent.run(prompt="hello") + print(result) + finally: + await sandbox.stop() + +asyncio.run(main()) ``` + +More ready-to-run examples are in `examples/install-agents/` (Claude Code, IFlowCli, Cursor CLI, Qwen Code, SWE-agent, OpenClaw, etc.). + +To run an agent evaluation/benchmark task via Job (a different code path with its own config schema), see [Use Job to Run Agent](./job.md). diff --git a/examples/evaluation/README.md b/examples/evaluation/README.md new file mode 100644 index 0000000000..291fe57697 --- /dev/null +++ b/examples/evaluation/README.md @@ -0,0 +1,19 @@ +# evaluation + +End-to-end evaluation demos that combine sandbox lifecycle, agent install/run, and a test suite — useful for understanding how individual pieces fit together at the script level. + +## Layout + +| Subdir | Path | Description | +|--------|------|-------------| +| [`swe_bench/`](./swe_bench/) | install-agent | Single-task SWE-bench Verified demo: starts a sandbox, installs an agent via `sandbox.agent.install()`, runs the agent on the task, runs the test suite, parses the result | + +## When to use this vs `job/harbor/` + +| | `evaluation/swe_bench/` | [`job/harbor/`](../job/harbor/) | +|--|------------------------|-------------------------------| +| Path | install-agent | Job (Harbor) | +| When | Debugging task setup or test parsing — full pipeline visible in script form | Production benchmark runs through the Harbor framework | +| API | `Sandbox` + `sandbox.agent.install()` | `Job(JobConfig.from_yaml(...)).run()` | + +If you're running SWE-bench through the standard pipeline, prefer [`job/harbor/`](../job/harbor/). diff --git a/examples/install-agents/README.md b/examples/install-agents/README.md new file mode 100644 index 0000000000..951430d556 --- /dev/null +++ b/examples/install-agents/README.md @@ -0,0 +1,28 @@ +# install-agents + +Examples for the **install-agent** way of using ROCK: install and run an agent inside a single sandbox via `sandbox.agent.install()` + `sandbox.agent.run(prompt)`. + +To run an agent evaluation/benchmark task via Job, see [`../job/`](../job/) instead. + +## Layout + +| Subdir | Agent runtime | +|--------|---------------| +| [`claude_code/`](./claude_code/) | Anthropic Claude Code CLI (`@anthropic-ai/claude-code`) | +| [`cursor_cli/`](./cursor_cli/) | Cursor CLI | +| [`iflow_cli/`](./iflow_cli/) | iFlow CLI (`@iflow-ai/iflow-cli`) | +| [`openclaw/`](./openclaw/) | OpenClaw — admin/proxy split-mode demo, has its own README | +| [`qwen_code/`](./qwen_code/) | qwen-code (`@qwen-code/qwen-code`) | +| [`swe_agent/`](./swe_agent/) | SWE-agent (`pip install -e` from GitHub) | + +Each subdir contains a `*_demo.py` entry point and a `rock_agent_config.yaml` driving the install/run. + +## Run + +```bash +# pick any subdir +cd iflow_cli +python iflow_cli_demo.py +``` + +See the [Install Agent in Sandbox (Experimental)](../../docs/versioned_docs/version-1.7.x/References/Python%20SDK%20References/rock-agent.md) reference for the full RockAgentConfig schema. diff --git a/examples/agents/claude_code/claude_code_demo.py b/examples/install-agents/claude_code/claude_code_demo.py similarity index 100% rename from examples/agents/claude_code/claude_code_demo.py rename to examples/install-agents/claude_code/claude_code_demo.py diff --git a/examples/agents/claude_code/rock_agent_config.yaml b/examples/install-agents/claude_code/rock_agent_config.yaml similarity index 100% rename from examples/agents/claude_code/rock_agent_config.yaml rename to examples/install-agents/claude_code/rock_agent_config.yaml diff --git a/examples/agents/cursor_cli/cursor_cli_demo.py b/examples/install-agents/cursor_cli/cursor_cli_demo.py similarity index 100% rename from examples/agents/cursor_cli/cursor_cli_demo.py rename to examples/install-agents/cursor_cli/cursor_cli_demo.py diff --git a/examples/agents/cursor_cli/rock_agent_config.yaml b/examples/install-agents/cursor_cli/rock_agent_config.yaml similarity index 100% rename from examples/agents/cursor_cli/rock_agent_config.yaml rename to examples/install-agents/cursor_cli/rock_agent_config.yaml diff --git a/examples/agents/iflow_cli/iflow_cli_demo.py b/examples/install-agents/iflow_cli/iflow_cli_demo.py similarity index 100% rename from examples/agents/iflow_cli/iflow_cli_demo.py rename to examples/install-agents/iflow_cli/iflow_cli_demo.py diff --git a/examples/agents/iflow_cli/integration_with_model_service/local/local_demo.py b/examples/install-agents/iflow_cli/integration_with_model_service/local/local_demo.py similarity index 100% rename from examples/agents/iflow_cli/integration_with_model_service/local/local_demo.py rename to examples/install-agents/iflow_cli/integration_with_model_service/local/local_demo.py diff --git a/examples/agents/iflow_cli/integration_with_model_service/local/rock_agent_config.yaml b/examples/install-agents/iflow_cli/integration_with_model_service/local/rock_agent_config.yaml similarity index 100% rename from examples/agents/iflow_cli/integration_with_model_service/local/rock_agent_config.yaml rename to examples/install-agents/iflow_cli/integration_with_model_service/local/rock_agent_config.yaml diff --git a/examples/agents/iflow_cli/integration_with_model_service/proxy/proxy_demo.py b/examples/install-agents/iflow_cli/integration_with_model_service/proxy/proxy_demo.py similarity index 100% rename from examples/agents/iflow_cli/integration_with_model_service/proxy/proxy_demo.py rename to examples/install-agents/iflow_cli/integration_with_model_service/proxy/proxy_demo.py diff --git a/examples/agents/iflow_cli/integration_with_model_service/proxy/rock_agent_config.yaml b/examples/install-agents/iflow_cli/integration_with_model_service/proxy/rock_agent_config.yaml similarity index 100% rename from examples/agents/iflow_cli/integration_with_model_service/proxy/rock_agent_config.yaml rename to examples/install-agents/iflow_cli/integration_with_model_service/proxy/rock_agent_config.yaml diff --git a/examples/agents/iflow_cli/rock_agent_config.yaml b/examples/install-agents/iflow_cli/rock_agent_config.yaml similarity index 100% rename from examples/agents/iflow_cli/rock_agent_config.yaml rename to examples/install-agents/iflow_cli/rock_agent_config.yaml diff --git a/examples/agents/openclaw/REAMDE.md b/examples/install-agents/openclaw/README.md similarity index 98% rename from examples/agents/openclaw/REAMDE.md rename to examples/install-agents/openclaw/README.md index 707c87f636..89f75685cc 100644 --- a/examples/agents/openclaw/REAMDE.md +++ b/examples/install-agents/openclaw/README.md @@ -52,7 +52,7 @@ rock admin start --env local-proxy --role proxy --port 9001 ## 4. Run the Demo ```bash -cd examples/agents/openclaw +cd examples/install-agents/openclaw python openclaw_demo.py ``` diff --git a/examples/agents/openclaw/openclaw.json b/examples/install-agents/openclaw/openclaw.json similarity index 100% rename from examples/agents/openclaw/openclaw.json rename to examples/install-agents/openclaw/openclaw.json diff --git a/examples/agents/openclaw/openclaw_demo.py b/examples/install-agents/openclaw/openclaw_demo.py similarity index 100% rename from examples/agents/openclaw/openclaw_demo.py rename to examples/install-agents/openclaw/openclaw_demo.py diff --git a/examples/agents/openclaw/rock_agent_config.yaml b/examples/install-agents/openclaw/rock_agent_config.yaml similarity index 100% rename from examples/agents/openclaw/rock_agent_config.yaml rename to examples/install-agents/openclaw/rock_agent_config.yaml diff --git a/examples/agents/qwen_code/qwen_code_demo.py b/examples/install-agents/qwen_code/qwen_code_demo.py similarity index 100% rename from examples/agents/qwen_code/qwen_code_demo.py rename to examples/install-agents/qwen_code/qwen_code_demo.py diff --git a/examples/agents/qwen_code/rock_agent_config.yaml b/examples/install-agents/qwen_code/rock_agent_config.yaml similarity index 100% rename from examples/agents/qwen_code/rock_agent_config.yaml rename to examples/install-agents/qwen_code/rock_agent_config.yaml diff --git a/examples/agents/swe_agent/rock_agent_config.yaml b/examples/install-agents/swe_agent/rock_agent_config.yaml similarity index 100% rename from examples/agents/swe_agent/rock_agent_config.yaml rename to examples/install-agents/swe_agent/rock_agent_config.yaml diff --git a/examples/agents/swe_agent/swe_agent_demo.py b/examples/install-agents/swe_agent/swe_agent_demo.py similarity index 100% rename from examples/agents/swe_agent/swe_agent_demo.py rename to examples/install-agents/swe_agent/swe_agent_demo.py diff --git a/examples/job/README.md b/examples/job/README.md new file mode 100644 index 0000000000..ed591838fa --- /dev/null +++ b/examples/job/README.md @@ -0,0 +1,16 @@ +# job + +Examples for the **Job** way of using ROCK: run an agent evaluation/task in a sandbox via `rock.sdk.job.Job` + `JobConfig`. + +For installing and running an agent inside a single sandbox, see [`../install-agents/`](../install-agents/) instead. + +## Layout + +| Subdir | Backend | Use it for | +|--------|---------|-----------| +| [`bash/`](./bash/) | `BashJobConfig` | Run an arbitrary shell script inside a sandbox (data processing, external evaluation tools) | +| [`harbor/`](./harbor/) | `HarborJobConfig` | Run an AI agent benchmark task (SWE-bench, Terminal Bench, …) via the Harbor framework | + +Both backends share a single `Job(config).run()` entrypoint — pick the config type based on your scenario. + +See the [Use Job to Run Agent](../../docs/versioned_docs/version-1.7.x/References/Python%20SDK%20References/job.md) reference for the full schema. diff --git a/examples/job/bash/README.md b/examples/job/bash/README.md new file mode 100644 index 0000000000..aff46add5e --- /dev/null +++ b/examples/job/bash/README.md @@ -0,0 +1,24 @@ +# job/bash + +`BashJob` examples: run an arbitrary shell script inside a sandbox. + +## Layout + +| File / dir | Form | Description | +|------------|------|-------------| +| [`simple_bash_job_demo.sh`](./simple_bash_job_demo.sh) | CLI | Minimal `rock job run --script-content ...` demo | +| [`claw_eval/`](./claw_eval/) | Python SDK | `claw-eval` benchmark wrapped as a BashJob — uses `JobConfig.from_yaml()` + `Job(config).run()` | + +Both forms use the same underlying `BashJobConfig` schema; the CLI just wraps it. + +## Quick run + +```bash +# CLI form +bash simple_bash_job_demo.sh + +# SDK form +cd claw_eval +cp claw_eval_bashjob.yaml.template claw_eval_bashjob.yaml # fill in real values +python run_claw_eval.py +``` diff --git a/examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template b/examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template similarity index 100% rename from examples/evaluation/claw_eval/claw_eval_bashjob.yaml.template rename to examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template diff --git a/examples/evaluation/claw_eval/claw_eval_config.yaml.template b/examples/job/bash/claw_eval/claw_eval_config.yaml.template similarity index 100% rename from examples/evaluation/claw_eval/claw_eval_config.yaml.template rename to examples/job/bash/claw_eval/claw_eval_config.yaml.template diff --git a/examples/evaluation/claw_eval/run_claw_eval.py b/examples/job/bash/claw_eval/run_claw_eval.py similarity index 95% rename from examples/evaluation/claw_eval/run_claw_eval.py rename to examples/job/bash/claw_eval/run_claw_eval.py index f6beaefec7..21d8609c86 100644 --- a/examples/evaluation/claw_eval/run_claw_eval.py +++ b/examples/job/bash/claw_eval/run_claw_eval.py @@ -1,7 +1,7 @@ """Run claw-eval via BashJob SDK. Usage: - cd examples/agents/claw_eval + cd examples/job/bash/claw_eval cp claw_eval_bashjob.yaml.template claw_eval_bashjob.yaml # fill in real values python run_claw_eval.py diff --git a/examples/evaluation/claw_eval/run_claw_eval.sh b/examples/job/bash/claw_eval/run_claw_eval.sh similarity index 100% rename from examples/evaluation/claw_eval/run_claw_eval.sh rename to examples/job/bash/claw_eval/run_claw_eval.sh diff --git a/examples/bash/simple_bash_job_demo.sh b/examples/job/bash/simple_bash_job_demo.sh similarity index 100% rename from examples/bash/simple_bash_job_demo.sh rename to examples/job/bash/simple_bash_job_demo.sh diff --git a/examples/job/harbor/README.md b/examples/job/harbor/README.md new file mode 100644 index 0000000000..061b7bdb54 --- /dev/null +++ b/examples/job/harbor/README.md @@ -0,0 +1,27 @@ +# job/harbor + +`HarborJob` examples: run an AI agent benchmark task via the Harbor framework. + +## Files + +| File | Purpose | +|------|---------| +| [`harbor_demo.py`](./harbor_demo.py) | Entry point — loads `JobConfig.from_yaml()`, runs `Job(config).run()`, iterates trial results | +| [`swe_job_config.yaml.template`](./swe_job_config.yaml.template) | SWE-bench task config template | +| [`swe_job_config-verifier.yaml.template`](./swe_job_config-verifier.yaml.template) | SWE-bench variant with `verifier.mode: native` | +| [`tb_job_config.yaml.template`](./tb_job_config.yaml.template) | Terminal Bench task config template | + +## Quick run + +```bash +# 1. copy a template and fill in real values +cp swe_job_config.yaml.template swe_job_config.yaml + +# 2. set required env vars (OSS credentials, etc.) — see harbor_demo.py docstring +source .env + +# 3. run +python harbor_demo.py -c swe_job_config.yaml +``` + +The `agents:` block uses Harbor's own minimal schema (typical fields: `name`, `model_name`) — see the templates above for the canonical shape. diff --git a/examples/harbor/harbor_demo.py b/examples/job/harbor/harbor_demo.py similarity index 88% rename from examples/harbor/harbor_demo.py rename to examples/job/harbor/harbor_demo.py index 09f48b3198..e6201acdb7 100644 --- a/examples/harbor/harbor_demo.py +++ b/examples/job/harbor/harbor_demo.py @@ -1,13 +1,11 @@ -"""Harbor benchmark demo using ROCK Job SDK (new path). +"""Harbor benchmark demo using ROCK Job SDK. Uses ``rock.sdk.job.Job`` with ``HarborJobConfig`` — the recommended path with full feature parity (G1-G7 fixed) and scatter / multiple trial types. -For the legacy path (``rock.sdk.bench.Job``), see ``harbor_demo_legacy.py``. - Usage: - python examples/harbor/harbor_demo.py -c examples/harbor/swe.intern.yaml - python examples/harbor/harbor_demo.py -c examples/harbor/tb_job_config.yaml -t mailman + python examples/job/harbor/harbor_demo.py -c examples/job/harbor/swe.intern.yaml + python examples/job/harbor/harbor_demo.py -c examples/job/harbor/tb_job_config.yaml -t mailman Required environment variables (OSS_* are auto-forwarded into the sandbox): OSS_ACCESS_KEY_ID Alibaba Cloud OSS access key ID @@ -20,7 +18,7 @@ Recommended setup: 1. Copy .env.example to .env and fill in your credentials 2. source .env - 3. python examples/harbor/harbor_demo.py -c ... + 3. python examples/job/harbor/harbor_demo.py -c ... """ import argparse diff --git a/examples/harbor/swe_job_config-verifier.yaml.template b/examples/job/harbor/swe_job_config-verifier.yaml.template similarity index 100% rename from examples/harbor/swe_job_config-verifier.yaml.template rename to examples/job/harbor/swe_job_config-verifier.yaml.template diff --git a/examples/harbor/swe_job_config.yaml.template b/examples/job/harbor/swe_job_config.yaml.template similarity index 100% rename from examples/harbor/swe_job_config.yaml.template rename to examples/job/harbor/swe_job_config.yaml.template diff --git a/examples/harbor/tb_job_config.yaml.template b/examples/job/harbor/tb_job_config.yaml.template similarity index 100% rename from examples/harbor/tb_job_config.yaml.template rename to examples/job/harbor/tb_job_config.yaml.template From 85b4fa0323b3adfccb53987abd43abeab5ceba07 Mon Sep 17 00:00:00 2001 From: berstpander Date: Tue, 28 Apr 2026 19:20:31 +0800 Subject: [PATCH 084/226] feat(cli): add -v verbosity control and unify log level management - Change -v from store_true to count action (0=ERROR, -v=WARNING, -vv=INFO, -vvv=DEBUG) - Add --httpx-log-level as override for third-party loggers (default=None, was INFO) - Fix config_log to iterate all rock.* child loggers (not just parent rock logger) since init_logger sets propagate=False per child logger - Move config_log before load_config_from_file to suppress WARNING logs at default level - Add 15 unit tests covering verbosity mapping, third-party loggers, httpx override, child logger propagation, and handler level updates --- rock/cli/main.py | 49 +++++++++-- tests/unit/cli/test_config_log.py | 142 ++++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+), 8 deletions(-) create mode 100644 tests/unit/cli/test_config_log.py diff --git a/rock/cli/main.py b/rock/cli/main.py index ebb6225c1b..33d7463609 100644 --- a/rock/cli/main.py +++ b/rock/cli/main.py @@ -19,6 +19,16 @@ logger = init_logger("rock.cli") +# ── Verbose level mapping ──────────────────────────────────────────── +_VERBOSE_LEVELS = [ + logging.ERROR, # 0: default + logging.WARNING, # 1: -v + logging.INFO, # 2: -vv + logging.DEBUG, # 3: -vvv +] +_THIRD_PARTY_LOGGERS = ("httpx", "httpcore", "urllib3") + + def load_config_from_file(args): """Load valid configuration, command line arguments take precedence over configuration file""" # Load configuration file @@ -81,16 +91,21 @@ def create_parser(command_classes: list[type]): ) # Global parameters - parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose logging") + parser.add_argument( + "-v", "--verbose", + action="count", + default=0, + help="Increase verbosity (-v=WARNING, -vv=INFO, -vvv=DEBUG; default=ERROR)", + ) parser.add_argument("--config", help="Path to config file (default: ./.rock/config.ini)") parser.add_argument("--base-url", help="ROCK server base URL (overrides config file)") parser.add_argument("--auth-token", help="ROCK authorization token (overrides config file)") parser.add_argument("--cluster", help="ROCK cluster (overrides config file)") parser.add_argument( "--httpx-log-level", - help="httpx log level (default: INFO, options: DEBUG, INFO, WARNING, ERROR)", + help="Override httpx/httpcore log level (overrides -v; options: DEBUG, INFO, WARNING, ERROR)", choices=["DEBUG", "INFO", "WARNING", "ERROR"], - default="INFO", + default=None, ) # extra-header parameter @@ -119,9 +134,26 @@ def find_command(command: str, subclasses: list[type[Command]]) -> type | None: def config_log(args: argparse.Namespace): - """Configure logging""" - logging.getLogger("httpx").setLevel(getattr(logging, args.httpx_log_level)) - logging.getLogger("httpcore").setLevel(getattr(logging, args.httpx_log_level)) + """Configure logging based on -v verbosity and --httpx-log-level.""" + idx = min(args.verbose, len(_VERBOSE_LEVELS) - 1) + level = _VERBOSE_LEVELS[idx] + + # --httpx-log-level overrides the inferred level for third-party loggers + third_party_level = level + if args.httpx_log_level is not None: + third_party_level = getattr(logging, args.httpx_log_level) + + # Apply to all rock.* loggers (each has its own handler + propagate=False) + for name, logger_obj in logging.Logger.manager.loggerDict.items(): + if name == "rock" or name.startswith("rock."): + if isinstance(logger_obj, logging.Logger): + logger_obj.setLevel(level) + for handler in logger_obj.handlers: + handler.setLevel(level) + + # Apply to third-party loggers + for name in _THIRD_PARTY_LOGGERS: + logging.getLogger(name).setLevel(third_party_level) def main(): @@ -137,11 +169,12 @@ def main(): parser.print_help() sys.exit(1) + # Configure logging before any business logic produces log output + config_log(args) + # Load valid configuration (configuration file + command line arguments) load_config_from_file(args) - config_log(args) - try: command = find_command(args.command, subclasses) if not command: diff --git a/tests/unit/cli/test_config_log.py b/tests/unit/cli/test_config_log.py new file mode 100644 index 0000000000..4a3b312056 --- /dev/null +++ b/tests/unit/cli/test_config_log.py @@ -0,0 +1,142 @@ +import argparse +import logging + +import pytest + +from rock.cli.main import config_log + + +def _make_args(verbose=0, httpx_log_level=None): + """Helper to create argparse.Namespace for config_log tests.""" + return argparse.Namespace(verbose=verbose, httpx_log_level=httpx_log_level) + + +@pytest.fixture(autouse=True) +def _reset_log_levels(): + """Reset logger levels before each test to avoid cross-test pollution.""" + # Save original levels + loggers = ["rock", "rock.cli", "rock.sdk.job.executor", "httpx", "httpcore", "urllib3"] + saved = {name: logging.getLogger(name).level for name in loggers} + saved_handlers = {} + for name in loggers: + logger = logging.getLogger(name) + saved_handlers[name] = [(h, h.level) for h in logger.handlers] + + yield + + # Restore original levels + for name in loggers: + logger = logging.getLogger(name) + logger.setLevel(saved[name]) + for handler, orig_level in saved_handlers[name]: + handler.setLevel(orig_level) + + +class TestConfigLogVerbosity: + """Test -v count → unified log level mapping.""" + + def test_default_error(self): + config_log(_make_args(verbose=0)) + assert logging.getLogger("rock").level == logging.ERROR + + def test_v1_warning(self): + config_log(_make_args(verbose=1)) + assert logging.getLogger("rock").level == logging.WARNING + + def test_v2_info(self): + config_log(_make_args(verbose=2)) + assert logging.getLogger("rock").level == logging.INFO + + def test_v3_debug(self): + config_log(_make_args(verbose=3)) + assert logging.getLogger("rock").level == logging.DEBUG + + def test_overflow_clamped_to_debug(self): + config_log(_make_args(verbose=99)) + assert logging.getLogger("rock").level == logging.DEBUG + + +class TestConfigLogThirdParty: + """Test third-party logger levels follow -v by default.""" + + def test_third_party_follows_verbosity(self): + config_log(_make_args(verbose=2)) + for name in ("httpx", "httpcore", "urllib3"): + assert logging.getLogger(name).level == logging.INFO + + def test_third_party_default_error(self): + config_log(_make_args(verbose=0)) + for name in ("httpx", "httpcore", "urllib3"): + assert logging.getLogger(name).level == logging.ERROR + + +class TestConfigLogHttpxOverride: + """Test --httpx-log-level overrides -v inference for third-party only.""" + + def test_httpx_override_does_not_affect_rock(self): + config_log(_make_args(verbose=0, httpx_log_level="WARNING")) + assert logging.getLogger("rock").level == logging.ERROR + assert logging.getLogger("httpx").level == logging.WARNING + + def test_httpx_override_with_v(self): + config_log(_make_args(verbose=2, httpx_log_level="WARNING")) + assert logging.getLogger("rock").level == logging.INFO + assert logging.getLogger("httpx").level == logging.WARNING + + def test_httpx_log_level_none_uses_verbosity(self): + config_log(_make_args(verbose=2, httpx_log_level=None)) + assert logging.getLogger("httpx").level == logging.INFO + + +class TestConfigLogRockHandler: + """Test rock logger handler levels are also updated.""" + + def test_rock_handler_level_updated(self): + config_log(_make_args(verbose=3)) + for handler in logging.getLogger("rock").handlers: + assert handler.level == logging.DEBUG + + +class TestConfigLogChildLogger: + """Test rock.* child loggers with propagate=False are also updated.""" + + def test_child_logger_level_updated(self): + """Simulate init_logger behavior: child logger with own handler + propagate=False.""" + child = logging.getLogger("rock.cli") + child.addHandler(logging.StreamHandler()) + child.setLevel(logging.INFO) + child.propagate = False + + config_log(_make_args(verbose=0)) + assert child.level == logging.ERROR + + def test_child_handler_level_updated(self): + """Child logger handler level should also be updated by config_log.""" + child = logging.getLogger("rock.cli") + child.addHandler(logging.StreamHandler()) + child.setLevel(logging.INFO) + child.propagate = False + + config_log(_make_args(verbose=3)) + for handler in child.handlers: + assert handler.level == logging.DEBUG + + def test_deep_child_logger_updated(self): + """Deeply nested rock.* logger (e.g. rock.sdk.job.executor) is also updated.""" + child = logging.getLogger("rock.sdk.job.executor") + child.addHandler(logging.StreamHandler()) + child.setLevel(logging.INFO) + child.propagate = False + + config_log(_make_args(verbose=2)) + assert child.level == logging.INFO + + def test_non_rock_logger_not_affected(self): + """Non-rock child loggers should NOT be changed by config_log.""" + other = logging.getLogger("myapp.worker") + other.addHandler(logging.StreamHandler()) + other.setLevel(logging.DEBUG) + other.propagate = False + + config_log(_make_args(verbose=0)) + assert other.level == logging.DEBUG # unchanged From 8985b95e14b559072eacb79adde77247a879e902 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Wed, 13 May 2026 11:36:04 +0800 Subject: [PATCH 085/226] feat: add startup timing instrumentation for sandbox launch stages (#924) Add StageTimer utility and instrument key sandbox startup stages (check availability, image pull, kata disk, docker run, wait alive, operator submit, meta store create) with [startup_timing] log markers to support performance analysis. Co-authored-by: Claude Opus 4.6 --- rock/deployments/docker.py | 30 +++++++++------- rock/sandbox/sandbox_manager.py | 32 ++++++++++------- rock/utils/__init__.py | 12 ++++++- rock/utils/concurrent_helper.py | 19 +++++++++++ tests/unit/utils/test_stage_timer.py | 51 ++++++++++++++++++++++++++++ 5 files changed, 117 insertions(+), 27 deletions(-) create mode 100644 tests/unit/utils/test_stage_timer.py diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 09748df6cf..48914f5318 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -32,7 +32,7 @@ ENV_POOL, DockerUtil, ImageUtil, - Timer, + StageTimer, find_free_port, get_executor, release_port, @@ -259,7 +259,7 @@ def _pull_image(self) -> None: logger.info(f"Pulling image {self._config.image!r}") try: - with Timer(description=f"[{self._config.image}] Image pull"): + with StageTimer("startup_timing", f"[{self._container_name}] [{self._config.image}] Image pull", logger): # Parse registry from image name registry, _ = ImageUtil.parse_registry_and_others(self._config.image) @@ -429,8 +429,9 @@ def _try_set_log_dir_quota(self, log_file_path: str) -> None: async def start(self): """Starts the runtime.""" - if not self.sandbox_validator.check_availability(): - raise Exception("Docker is not available") + with StageTimer("startup_timing", f"[{self._container_name}] Check availability", logger): + if not self.sandbox_validator.check_availability(): + raise Exception("Docker is not available") storage_opt_supported = DockerUtil.detect_storage_opt_support() # Resolve effective rootfs quota: downgrade to None if storage-opt is not supported. @@ -490,12 +491,14 @@ async def start(self): # Kata DinD: prepare disk image and add volume mount + env var if self._config.use_kata_runtime: - self._prepare_kata_disk() + with StageTimer("startup_timing", f"[{self._container_name}] Kata disk prepare", logger): + self._prepare_kata_disk() disk_path = self._get_kata_disk_image_path() volume_args.extend(["-v", f"{disk_path}:/docker-disk.img"]) env_arg.extend(["-e", "ROCK_KATA_RUNTIME=true"]) - time.sleep(random.randint(0, 5)) + with StageTimer("startup_timing", f"[{self._container_name}] Random sleep", logger): + time.sleep(random.randint(0, 5)) runtime_args = self._build_runtime_args() cmds = [ "docker", @@ -528,14 +531,15 @@ async def start(self): ) logger.info(f"Command: {cmd_str!r}") # shell=True required for && etc. - with Timer(description=f"[{self._config.image}] Container start"): + with StageTimer("startup_timing", f"[{self._container_name}] Docker run", logger): self._container_process = await loop.run_in_executor(executor, self._docker_run, cmds) - await loop.run_in_executor(executor, self._hooks.on_custom_step, DeploymentHookStep.STARTING_RUNTIME) - logger.info(f"Starting runtime at {self._config.port}") - self._runtime = RemoteSandboxRuntime.from_config( - RemoteSandboxRuntimeConfig(port=self._config.port, timeout=self._runtime_timeout) - ) - self._runtime.set_executor(executor) + await loop.run_in_executor(executor, self._hooks.on_custom_step, DeploymentHookStep.STARTING_RUNTIME) + logger.info(f"Starting runtime at {self._config.port}") + self._runtime = RemoteSandboxRuntime.from_config( + RemoteSandboxRuntimeConfig(port=self._config.port, timeout=self._runtime_timeout) + ) + self._runtime.set_executor(executor) + with StageTimer("startup_timing", f"[{self._container_name}] Wait until alive", logger): await self._wait_until_alive(timeout=self._config.startup_timeout) if self._config.enable_auto_clear: self._check_stop_task = asyncio.create_task(self._check_stop()) diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index d796b08ecf..3d85fd0671 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -37,6 +37,7 @@ from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError +from rock.utils import StageTimer from rock.utils.crypto_utils import AESEncryption from rock.utils.format import convert_to_gb, parse_size_to_bytes from rock.utils.system import get_iso8601_timestamp @@ -112,7 +113,8 @@ async def start_async( ) -> SandboxStartResponse: await self._check_sandbox_exists_in_redis(config) self.validate_sandbox_spec(self.rock_config.runtime, config) - docker_deployment_config: DockerDeploymentConfig = await self.deployment_manager.init_config(config) + with StageTimer("startup_timing", f"[{config.image}] Init config", logger): + docker_deployment_config: DockerDeploymentConfig = await self.deployment_manager.init_config(config) sandbox_id = docker_deployment_config.container_name if self.rock_config.runtime.use_standard_spec_only: @@ -123,15 +125,17 @@ async def start_async( ) docker_deployment_config.cpus = self.rock_config.runtime.standard_spec.cpus docker_deployment_config.memory = self.rock_config.runtime.standard_spec.memory - sandbox_info: SandboxInfo = await self._operator.submit(docker_deployment_config, user_info) + with StageTimer("startup_timing", f"[{sandbox_id}] Operator submit", logger): + sandbox_info: SandboxInfo = await self._operator.submit(docker_deployment_config, user_info) await self._build_sandbox_info_metadata(sandbox_info, user_info, cluster_info) timeout_info = SandboxTimeoutHelper.make_timeout_info(docker_deployment_config.auto_clear_time) - await self._meta_store.create( - sandbox_id, - sandbox_info, - timeout_info=timeout_info, - deployment_config=docker_deployment_config, - ) + with StageTimer("startup_timing", f"[{sandbox_id}] Meta store create", logger): + await self._meta_store.create( + sandbox_id, + sandbox_info, + timeout_info=timeout_info, + deployment_config=docker_deployment_config, + ) return SandboxStartResponse( sandbox_id=sandbox_id, host_name=sandbox_info.get("host_name"), @@ -148,13 +152,15 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: sandbox_actor: SandboxActor = await deployment.creator_actor(actor_name) - await self._ray_service.async_ray_get(sandbox_actor.start.remote()) + with StageTimer("startup_timing", f"[{sandbox_id}] Actor start", logger): + await self._ray_service.async_ray_get(sandbox_actor.start.remote()) logger.info(f"sandbox {sandbox_id} is started") - while not await self._is_actor_alive(sandbox_id): - logger.debug(f"wait actor for sandbox alive, sandbox_id: {sandbox_id}") - # TODO: timeout check - await asyncio.sleep(1) + with StageTimer("startup_timing", f"[{sandbox_id}] Wait actor alive", logger): + while not await self._is_actor_alive(sandbox_id): + logger.debug(f"wait actor for sandbox alive, sandbox_id: {sandbox_id}") + # TODO: timeout check + await asyncio.sleep(1) await self.get_status(sandbox_id) return SandboxStartResponse( diff --git a/rock/utils/__init__.py b/rock/utils/__init__.py index 285f99b06e..c33716fb46 100644 --- a/rock/utils/__init__.py +++ b/rock/utils/__init__.py @@ -1,6 +1,15 @@ from contextvars import ContextVar -from .concurrent_helper import AsyncAtomicInt, AsyncSafeDict, RayUtil, Timer, get_executor, run_until_complete, timeout +from .concurrent_helper import ( + AsyncAtomicInt, + AsyncSafeDict, + RayUtil, + StageTimer, + Timer, + get_executor, + run_until_complete, + timeout, +) from .data import ( FileUtil, ListUtil, @@ -54,6 +63,7 @@ "get_uniagent_endpoint", # Concurrent utilities "get_executor", + "StageTimer", "Timer", "RayUtil", "AsyncSafeDict", diff --git a/rock/utils/concurrent_helper.py b/rock/utils/concurrent_helper.py index f30408ad64..4115170937 100644 --- a/rock/utils/concurrent_helper.py +++ b/rock/utils/concurrent_helper.py @@ -97,6 +97,25 @@ def __exit__(self, exc_type, exc_val, exc_tb): return False +class StageTimer: + """Context manager that logs elapsed time for a named stage.""" + + def __init__(self, phase: str, description: str, logger): + self._phase = phase + self._description = description + self._logger = logger + self._start_time = 0.0 + + def __enter__(self): + self._start_time = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + duration = time.perf_counter() - self._start_time + self._logger.info(f"[{self._phase}] {self._description} took {duration:.3f} s") + return False + + class AsyncSafeDict(Generic[K, V]): """Thread-safe async dictionary""" diff --git a/tests/unit/utils/test_stage_timer.py b/tests/unit/utils/test_stage_timer.py new file mode 100644 index 0000000000..971e806a48 --- /dev/null +++ b/tests/unit/utils/test_stage_timer.py @@ -0,0 +1,51 @@ +import logging +from unittest.mock import patch + +import pytest + +from rock.utils.concurrent_helper import StageTimer + + +class TestStageTimer: + def test_logs_duration_on_exit(self, caplog): + with caplog.at_level(logging.INFO): + logger = logging.getLogger("test_stage_timer") + with StageTimer("startup_timing", "[sandbox-abc] Check availability", logger): + pass + + assert len(caplog.records) == 1 + record = caplog.records[0] + assert "[startup_timing]" in record.message + assert "[sandbox-abc] Check availability" in record.message + assert "took" in record.message + assert "s" in record.message + + @patch("rock.utils.concurrent_helper.time.perf_counter", side_effect=[100.0, 102.5]) + def test_duration_calculation(self, mock_perf_counter, caplog): + with caplog.at_level(logging.INFO): + logger = logging.getLogger("test_stage_timer") + with StageTimer("startup_timing", "[sandbox-abc] Image pull", logger): + pass + + assert "took 2.500 s" in caplog.records[0].message + + def test_exception_still_logs(self, caplog): + with caplog.at_level(logging.INFO): + logger = logging.getLogger("test_stage_timer") + with pytest.raises(ValueError, match="boom"): + with StageTimer("startup_timing", "[sandbox-abc] Docker run", logger): + raise ValueError("boom") + + assert len(caplog.records) == 1 + assert "[startup_timing]" in caplog.records[0].message + assert "[sandbox-abc] Docker run" in caplog.records[0].message + + def test_log_format(self, caplog): + with caplog.at_level(logging.INFO): + logger = logging.getLogger("test_stage_timer") + with StageTimer("my_phase", "my description", logger): + pass + + record = caplog.records[0] + assert record.message.startswith("[my_phase] my description took ") + assert record.message.endswith(" s") From 33ebb0e939eec6ee6211be15364e54a2c4791a66 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Wed, 13 May 2026 14:32:36 +0800 Subject: [PATCH 086/226] fix(rocklet): mount loop disk to docker data-root instead of hardcoded path (#933) setup_kata_dind now reads data-root from /etc/docker/daemon.json so the loop disk is mounted to the correct directory when a custom data-root is configured. Co-authored-by: Claude Opus 4.6 --- rock/rocklet/local_files/docker_run.sh | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/rock/rocklet/local_files/docker_run.sh b/rock/rocklet/local_files/docker_run.sh index 1fe9a89d81..e319d5879a 100755 --- a/rock/rocklet/local_files/docker_run.sh +++ b/rock/rocklet/local_files/docker_run.sh @@ -30,11 +30,19 @@ is_nix() { # Kata DinD: set up loop device and mount disk image for Docker storage setup_kata_dind() { - mkdir -p /var/lib/docker + local docker_root="/var/lib/docker" + if [ -f /etc/docker/daemon.json ]; then + local custom_root + custom_root=$(grep -o '"data-root"[[:space:]]*:[[:space:]]*"[^"]*"' /etc/docker/daemon.json | sed 's/.*"data-root"[[:space:]]*:[[:space:]]*"\([^"]*\)"/\1/') + if [ -n "$custom_root" ]; then + docker_root="$custom_root" + fi + fi + mkdir -p "$docker_root" for i in $(seq 0 7); do mknod -m 660 /dev/loop$i b 7 $i 2>/dev/null || true done - mount -o loop /docker-disk.img /var/lib/docker + mount -o loop /docker-disk.img "$docker_root" mount -o remount,rw /sys/fs/cgroup mount -o remount,rw /proc/sys } From a72f7d19d0506aeb59cc794a057f484ab65621cb Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Wed, 13 May 2026 15:23:30 +0800 Subject: [PATCH 087/226] feat(model-service): proxy supports stream + replay, byte passthrough, ForwardBackend/ReplayBackend (#935) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(model-service): rebuild proxy on litellm SDK with traj record/replay 替换 model-service `proxy` 模式手写的 httpx forward + retry_async,改为基于 litellm SDK 调用,同时新增 chat/completions 轨迹的录制与顺序回放能力,服务 SWE-agent / mini-swe-agent / OpenHands 等 deterministic agent 的无 LLM 成本调试。 主要改动: - proxy.py 改用 litellm.acompletion(num_retries / extra_headers / streaming) - 新增 TrajectoryRecorder(CustomLogger) 录制 StandardLoggingPayload 到 JSONL - 新增 TrajectoryReplayer(CustomLLM) + SequentialCursor 顺序回放单个 jsonl 文件 - ModelServiceConfig 新增 num_retries / traj_enabled / traj_file / replay_traj_path - CLI 新增 --num-retries / --traj-file(同时承担 replay 入口) - local 模式保留旧 record_traj 装饰器,不受影响 - 删除 examples 旧 YAML,改 README 主推纯 CLI 启动方式 - docs/dev/litellm_proxy_refactor.md 写明设计与 breaking change Co-Authored-By: Claude Sonnet 4.6 * fix(model-service): pass api_key + use custom_openai prefix + suppress cost calc Three fixes to proxy.py uncovered while testing against DashScope (glm-5): - Extract Bearer token from incoming Authorization header and pass as litellm api_key kwarg; setting it via extra_headers does not work because litellm always regenerates Authorization from api_key. Authorization is now stripped from the forwarded header set. - Switch upstream prefix from openai/ to custom_openai/. This is litellm's standard pattern for OpenAI-compatible third-party endpoints (DashScope, ModelScope, Groq, Mistral, ...) and avoids "model isn't mapped" on arbitrary upstream model names. - Pass input_cost_per_token=0 / output_cost_per_token=0 so litellm's cost calculator does not raise "model isn't mapped" on unknown models and pollute StandardLoggingPayload.response_cost_failure_debug_information. Co-Authored-By: Claude Sonnet 4.6 * refactor(model-service): drop CustomLLM, serve replay directly from cursor The replay path no longer goes through litellm. We have a complete OpenAI-shape response on disk, so routing it through CustomLLM/CustomStreamWrapper just to translate formats was pure overhead — and the source of every replay-side bug (cursor-exhausted retried 6× and wrapped as APIConnectionError, GenericStreamingChunk type gymnastics, finish_reason hardcoded to "stop", reasoning_content dropped on streaming, tool_calls reconstruction left as TODO). Changes: - traj_replayer.py: delete TrajectoryReplayer(CustomLLM) and helpers; keep just SequentialCursor. Cursor exhaustion now raises a plain TrajectoryExhausted. - proxy.py: in replay mode, fetch from app.state.replay_cursor and emit either the raw response dict (non-stream) or one SSE chunk + [DONE] (stream). The stream path renames message → delta and preserves all fields verbatim (finish_reason, tool_calls, reasoning_content, ...). - main.py: rename _configure_litellm_for_proxy → _configure_proxy_integrations. Replay branch now just attaches a SequentialCursor to app.state; no litellm.custom_provider_map registration. - Tests: drop the CustomLLM-based replayer tests; keep cursor tests; add three end-to-end proxy replay tests covering non-stream / stream / cursor exhausted. 43 passed. Direct curl against DashScope glm-5: record + replay (both modes) verified end-to-end. Co-Authored-By: Claude Sonnet 4.6 * refactor(model-service): drop litellm, use httpx byte-passthrough + openai SDK as parser Replaces litellm.acompletion with raw httpx forwarding. The proxy no longer parses or rewrites the OpenAI protocol on the forward path — body bytes go upstream as-is, response bytes come back as-is. The openai SDK is kept solely as a parser library for the recording side: ChatCompletionChunk + the official ChatCompletionStreamState aggregate streaming chunks into a final ChatCompletion that the recorder writes to JSONL. This restores the proxy's original "transparent forward" intent and eliminates several litellm-specific pain points encountered during testing: - No "model isn't mapped yet" cost-calc exception (no calc happens at all). - No need for the input_cost_per_token=0 / custom_openai prefix workarounds. - Authorization header passes through verbatim (no api_key extraction kludge). - Provider-specific fields (reasoning_content, citations, ...) are preserved byte-for-byte going to the client AND auto-aggregated in the recorded traj (openai SDK uses extra="allow" pydantic mode). - Cursor exhaustion in replay returns 404 directly, never gets retried. Changes: - pyproject.toml: drop litellm>=1.50.0, add openai>=1.50.0 and httpx - proxy.py: rewrite forward path with httpx; record streams via dual-purpose byte forwarding + parallel SSE parsing into ChatCompletionStreamState - traj_recorder.py: drop CustomLogger inheritance; expose explicit recorder.record(request, response, status, ...) API called from proxy.py - main.py: attach recorder/cursor to app.state instead of registering with litellm.callbacks / litellm.custom_provider_map - test_proxy.py: rewritten to use httpx.MockTransport for upstream mocking; cover byte passthrough, provider-specific field preservation, error forwarding, recorder invocation, replay paths - test_traj_recorder.py: rewritten for the explicit-call API 36 passed. End-to-end verified against DashScope glm-5: streaming record, non-stream replay, streaming replay, cursor exhausted -> 404 all work. Co-Authored-By: Claude Sonnet 4.6 * refactor: use openai sdk * chore(model-service): remove litellm remnants (num_retries, stale comments) Co-Authored-By: Claude Sonnet 4.6 * refactor(model-service): split proxy handler into _ReplayBackend / _ForwardBackend Strategy pattern eliminates the replay/forward branch inside chat_completions. The backend is selected once at startup (_configure_proxy_integrations) and attached to app.state.backend; the endpoint just parses the request and dispatches. Each backend keeps the stream/non-stream branch local to itself. A union type alias _CompletionBackend documents the closed set of backends and a typed _get_backend(request) accessor wraps the app.state read. Co-Authored-By: Claude Sonnet 4.6 * refactor(model-service): drop _ prefix from public Backend classes, rename replay_traj_path → replay_traj_file The Backend classes are imported by main.py and tests, so the leading underscore mis-signalled them as module-internal. Rename the config field to align with traj_file naming. Also drop the defensive getattr() for args.traj_file — argparse always sets it. Co-Authored-By: Claude Sonnet 4.6 * feat(model-service): restore retry on retryable_status_codes + connection errors Retry was lost when litellm was dropped (it provided num_retries internally). Restored using a unified _send_with_retry helper that: - Always opens upstream with stream=True so the same code path serves both stream and non-stream callers (non-stream just await resp.aread()). - Retries on httpx.TimeoutException, httpx.ConnectError, and HTTP statuses in config.retryable_status_codes (default [429, 500]). - Defaults: 6 attempts, exponential backoff 2s→32s with jitter — matches the original perform_llm_request behavior. - For stream: retry happens before any byte is yielded; mid-stream drops are not retried (would corrupt downstream). Module-level retry constants are read at call time so tests can monkeypatch them. Added 4 tests covering: success after retry, exhausted retries returning last response, non-whitelisted status not retried, and stream retry path. Co-Authored-By: Claude Sonnet 4.6 * test(model-service): add e2e tests against an in-thread uvicorn mock upstream A tiny FastAPI mock app runs in a background thread via uvicorn; the proxy calls it over real TCP through its own httpx.AsyncClient — production code path, no transport injection or patching. Three scenarios: - non-stream forward: vendor field round-trips, recorder writes JSONL - stream forward: SSE chunks reach the client, recorder gets aggregated final completion - record-then-replay: replay phase uses a bogus base_url to prove the upstream isn't called Tests use FastAPI's TestClient (sync) so the test bodies read top-down with no async noise; the async wiring lives inside MockUpstreamServer. Drive-by cleanups in proxy.py: localize the openai SDK imports inside the streaming aggregator (only needed there), and drop the now-unused _RETRY_EXCEPTIONS constant. Co-Authored-By: Claude Sonnet 4.6 * fix(model-service): inject positional index into replay-stream tool_calls deltas A recorded non-stream message.tool_calls carries no 'index' field, but the OpenAI streaming spec requires it on chunk deltas. Without it, downstream clients using the openai SDK reject the replay-stream chunk with a pydantic ValidationError ('Field required: index'). completion_to_chunk_dict now injects a positional index when missing (existing 'index' is preserved). Co-Authored-By: Claude Opus 4.7 * test(model-service): refactor proxy e2e into MockUpstream + TestProxyRecordReplay Rename test_proxy_e2e.py → test_proxy_record_replay.py to make the file purpose explicit (the suite revolves around the record→replay capability). Refactor the test surface: - MockUpstream class encapsulates the FastAPI app, server lifecycle, the canonical reply, and an assert_message() helper. Test data and the handler stay in sync because they share the same constants. - TestProxyRecordReplay class groups the three tests with shorter names: test_forward_non_stream test_forward_stream test_replay (parametrized over record × replay stream/non-stream) - _call_chat_completions helper unifies stream/non-stream call sites. Expand coverage to 2 parallel tool_calls (get_weather + get_time) — exercises the openai SDK aggregator's multi-index tool_call assembly. Co-Authored-By: Claude Opus 4.7 * chore: remove useless comment in pyproject.toml * chore: remove uesless dev docs * refactor(model-service): flatten layout — drop integrations/, rename sse_utils→sse, merge traj_*→traj The integrations/ directory only ever held two files (traj_recorder, traj_replayer) and the litellm CustomLogger angle that justified the name is long gone. Both modules share one JSONL schema, so collapsing them into a single traj.py keeps the schema and its read/write halves visible together. sse_utils.py → sse.py: the codec is the module's whole purpose, the _utils suffix added nothing. Drop traj_recorder.now() — a one-line wrapper around time.time() with no callers. Also remove a stray _get_or_create_metrics_monitor patch in test_forward_invokes_recorder_on_success: OTLP create-time failure only logs a warning, so the patch was protecting against nothing. Co-Authored-By: Claude Opus 4.7 * refactor(model-service): rename traj_file→recording_file, replay_traj_file→replay_file The old names were ambiguous: "traj_file" alone gave no hint of write vs read, and the CLI flag --traj-file was actually wired to config.replay_traj_file — same word pointing in opposite directions depending on context. New names mirror the backend pair (ForwardBackend = recording, ReplayBackend = replay) so the role is obvious at the field. CLI is split into two independent flags accordingly. Recorder constructor still takes traj_file= since it names the JSONL file type, not its role; only the config-field / CLI surface changes. Co-Authored-By: Claude Opus 4.7 * feat(model-service): enforce recording_file/replay_file mutex via model_validator Setting both at once was silently resolved in favor of replay (the backend factory checks replay_file first), masking what is really a configuration error. A Pydantic model_validator now rejects the combination at construction time; validate_assignment=True extends the check to CLI-style field-by-field overrides applied after a yaml load. Three tests added: construction-time mutex, assignment-time mutex, and the existing yaml-load test split into one-side-only variants since the original deliberately set both fields. Co-Authored-By: Claude Opus 4.7 * refactor(model-service): move _replay_sse_iter into ReplayBackend as a staticmethod Module-level function with a single call site inside ReplayBackend; the SSE chunk-emit shape is purely a replay-mode implementation detail. Moving it inside the class also makes the pairing with the JSON branch in serve() visible at a glance. Co-Authored-By: Claude Opus 4.7 * refactor(model-service): move get_base_url into ForwardBackend as _resolve_base_url Single call site inside ForwardBackend.serve, and the function reads only self._config — drop the redundant config parameter and rename to _resolve_base_url to make the multi-source fallback (proxy_base_url → proxy_rules[model] → proxy_rules['default']) explicit. Co-Authored-By: Claude Opus 4.7 * refactor(model-service): move _forward_stream_and_record into ForwardBackend as _stream_and_record Same single-call-site argument as the previous moves: 3 of the 7 kwargs were just relaying self._config / self._recorder. As an instance method the parameter list drops to 4 and the streaming path mirrors the structure of ReplayBackend._sse_iter. _send_with_retry stays at module scope — it's a pure helper bound to the httpx.AsyncClient lifecycle, not to any backend state. Co-Authored-By: Claude Opus 4.7 * docs(model-service): move + rewrite proxy README under docs/dev/model-service/ Old examples/model_service/README.md was stale: still mentioned litellm, StandardLoggingPayload, --num-retries, and the conflated --traj-file flag. Rewritten to reflect current shape: ForwardBackend / ReplayBackend pair, recording_file / replay_file mutex, retry-on-status-code with the documented attempt budget, openai SDK only used as the stream-state aggregator behind the forwarding path. Also calls out that the rock model-service start subcommand hasn't been wired up with the new flags yet. Co-Authored-By: Claude Opus 4.7 * feat(model-service): expose --recording-file / --replay-file on rock model-service start The two flags previously existed only on the python -m rock.sdk.model.server.main entry; rock model-service start (which subprocess-spawns that same module) had no way to thread them through, forcing users to bypass the CLI for any record/replay scenario. Wire them through ModelServiceCommand argparse → ModelService.start → start_sandbox_service → subprocess argv. Add three tests around the argv construction (default omits both flags, recording_file forwarded, replay_file forwarded). Doc updated: drop the python -m caveat and switch all example commands to rock model-service start. Co-Authored-By: Claude Opus 4.7 * test(model-service): add CLI-layer coverage for --recording-file / --replay-file wiring The existing tests/unit/sdk/model/test_service.py covers the subprocess argv construction (catches cmd-string typos like --recording_file vs --recording-file) but mocks nothing above the SDK layer, so a missing kwarg in ModelServiceCommand.arun would slip through. Add tests/unit/cli/command/test_model_service.py mirroring the test_job.py pattern: drive the real argparse sub-parser end-to-end and mock ModelService.start to assert the kwargs it receives. Covers the new flags both in isolation and in their default (omitted) state. Two layers, two bug surfaces — together they cover the full path from CLI argv to subprocess argv. Co-Authored-By: Claude Opus 4.7 * test(model-service): rename test_service.py → test_service_subprocess.py to fix CI collision tests/integration/envhub/test_service.py shares the same basename, and pytest's default importmode (prepend) collapses both into a single 'test_service' module in sys.modules, so collection fails as soon as both files are picked up: import file mismatch: imported module 'test_service' has this __file__ attribute: .../envhub/test_service.py which is not the same as the test file we want to collect: .../sdk/model/test_service.py Renaming the new file is the smallest fix and keeps importmode=prepend behavior unchanged for the rest of the suite. The new name also describes the file better (it tests how start_sandbox_service builds the subprocess argv). Co-Authored-By: Claude Opus 4.7 * test(model-service): rename test_proxy_record_replay.py → ..._e2e.py The file is the only one in the model-service unit suite that boots a real uvicorn upstream in a background thread and drives the proxy through real HTTP — append _e2e to make that scope obvious in the file name. It stays under tests/unit/ because the project's integration/ tier is reserved for tests requiring out-of-process services (Docker, Ray, admin), which this one doesn't. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Sonnet 4.6 --- docs/dev/model-service/proxy.md | 154 +++ pyproject.toml | 2 + rock/cli/command/model_service.py | 14 + rock/sdk/model/server/api/proxy.py | 465 +++++++-- rock/sdk/model/server/config.py | 23 +- rock/sdk/model/server/main.py | 48 +- rock/sdk/model/server/sse.py | 99 ++ rock/sdk/model/server/traj.py | 156 +++ rock/sdk/model/server/utils.py | 2 +- rock/sdk/model/service.py | 10 + tests/unit/cli/command/test_model_service.py | 120 +++ tests/unit/sdk/model/test_proxy.py | 900 ++++++++++-------- .../sdk/model/test_proxy_record_replay_e2e.py | 332 +++++++ .../unit/sdk/model/test_service_subprocess.py | 38 + tests/unit/sdk/model/test_sse.py | 223 +++++ tests/unit/sdk/model/test_traj_recorder.py | 141 +++ tests/unit/sdk/model/test_traj_replayer.py | 122 +++ uv.lock | 135 +++ 18 files changed, 2488 insertions(+), 496 deletions(-) create mode 100644 docs/dev/model-service/proxy.md create mode 100644 rock/sdk/model/server/sse.py create mode 100644 rock/sdk/model/server/traj.py create mode 100644 tests/unit/cli/command/test_model_service.py create mode 100644 tests/unit/sdk/model/test_proxy_record_replay_e2e.py create mode 100644 tests/unit/sdk/model/test_service_subprocess.py create mode 100644 tests/unit/sdk/model/test_sse.py create mode 100644 tests/unit/sdk/model/test_traj_recorder.py create mode 100644 tests/unit/sdk/model/test_traj_replayer.py diff --git a/docs/dev/model-service/proxy.md b/docs/dev/model-service/proxy.md new file mode 100644 index 0000000000..b0f77da0f3 --- /dev/null +++ b/docs/dev/model-service/proxy.md @@ -0,0 +1,154 @@ +# model-service `proxy` 模式 + +`rock model-service` 的 proxy 模式在 `/v1/chat/completions` 上提供一个 OpenAI 兼容的转发层， +两种工作模式互斥： + +| 模式 | 触发条件 | 上游调用 | 写盘 | +|-----------|---------------------------------------|----------|----------------------| +| Recording | 默认 | 真实调用 | append 到 JSONL traj | +| Replay | `--replay-file` / `replay_file` 设置 | 不调用 | 不写 | + +设计目标是让 SWE-agent / mini-swe-agent / OpenHands 等 agent 框架在录制 → 回放之间无感切换： +agent 不变，只换 base URL。 + +下文所有命令以 `rock model-service start` 启动；该子命令最终会 `subprocess` 拉起 +`rock.sdk.model.server.main`，两者支持的 flag 一致。直接调试时也可以用 +`python -m rock.sdk.model.server.main` 跳过 PID 文件管理。 + +--- + +## 1. Recording（默认） + +转发到单个上游，每次调用 append 一行 JSONL 到 `recording_file`（缺省 `LOG_DIR/LLMTraj.jsonl`， +其中 `LOG_DIR = $ROCK_MODEL_SERVICE_DATA_DIR`）： + +```bash +export OPENAI_API_KEY="sk-..." +export ROCK_MODEL_SERVICE_DATA_DIR=/tmp/rock-traj + +rock model-service start \ + --type proxy \ + --proxy-base-url https://api.openai.com/v1 \ + --port 8080 +``` + +调用： + +```bash +curl -X POST http://127.0.0.1:8080/v1/chat/completions \ + -H "Authorization: Bearer $OPENAI_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}]}' + +cat /tmp/rock-traj/LLMTraj.jsonl | jq '.model, .response.choices[0].message.content' +``` + +流式同样支持，上游字节原样转给客户端，recorder 在后台聚合最终的 `ChatCompletion` 写盘 +（用 openai SDK 的 `ChatCompletionStreamState`，所以 `tool_calls.function.arguments` 等 +跨 chunk 拼接的字段会被还原成完整形态）： + +```bash +curl -N -X POST http://127.0.0.1:8080/v1/chat/completions \ + -H "Authorization: Bearer $OPENAI_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"gpt-3.5-turbo","stream":true,"messages":[{"role":"user","content":"count to 5"}]}' +``` + +显式指定写到别的路径： + +```bash +rock model-service start \ + --type proxy \ + --proxy-base-url https://api.openai.com/v1 \ + --recording-file /tmp/my-session.jsonl \ + --port 8080 +``` + +--- + +## 2. Replay + +把 `--replay-file` 指到一个录好的 jsonl，proxy 不再访问真实 LLM，按录制顺序返回响应； +agent 把 base URL 换成 `http://127.0.0.1:8081/v1` 即可重放： + +```bash +rock model-service start \ + --type proxy \ + --replay-file /tmp/rock-traj/LLMTraj.jsonl \ + --port 8081 +``` + +行为细节： + +- cursor 单调推进，每次请求消耗一条记录；用尽后返回 **404**。 +- 流式请求会拿录制的 `ChatCompletion` 重新发一帧 SSE chunk + `[DONE]`。 + `tool_calls` 的 `index` 字段会被自动注入（OpenAI 的流式协议要求 chunk delta 上有 `index`， + 但录制态的 `message.tool_calls` 没有）。 +- request 里的 `model` 会跟录制的 `model` 比对，不一致只打 warning，不阻断。 + +`recording_file` 和 `replay_file` 是**互斥**的——同时配置（无论是 CLI 还是 YAML）会在启动时 +被 Pydantic `model_validator` 拦下并报 `ValidationError`，避免"录到一半把源文件覆盖"这类隐性 bug。 + +--- + +## 3. 重试和超时 + +- 默认对 connection error / timeout 和 `retryable_status_codes`（默认 `[429, 500]`）触发重试， + 最多 6 次，指数退避 2s 起步 ×2 + 抖动；最后一次仍失败时把上游响应原样转给客户端 + （**不**包装成 502/504，让 agent 自己看到真实状态码）。 +- 对**流式**请求，重试只发生在第一个字节抵达客户端**之前**——一旦字节流开始转发， + 连接中断不会重试（已发出去的字节无法收回）。 + +```bash +rock model-service start \ + --type proxy \ + --proxy-base-url https://api.openai.com/v1 \ + --retryable-status-codes 429,500,502,503 \ + --request-timeout 60 \ + --port 8080 +``` + +--- + +## 4. 多模型路由（YAML） + +按 model name 分流到不同上游需要 YAML（CLI 只暴露单一 `--proxy-base-url`）。新建 `routes.yaml`： + +```yaml +proxy_rules: + gpt-3.5-turbo: "https://api.openai.com/v1" + gpt-4o: "https://api.openai.com/v1" + default: "https://api-inference.modelscope.cn/v1" + +retryable_status_codes: [429, 500, 502] +request_timeout: 60 +recording_file: /tmp/rock-traj/multi.jsonl +``` + +启动： + +```bash +rock model-service start \ + --type proxy \ + --config-file routes.yaml \ + --port 8080 +``` + +CLI flag（`--proxy-base-url` / `--port` / `--retryable-status-codes` / ...）覆盖 YAML 同名字段。 +路由解析顺序：`proxy_base_url` → `proxy_rules[model]` → `proxy_rules["default"]`，都没有则 400。 + +--- + +## 5. 实现要点（仅供参考） + +- `chat_completions` endpoint 把请求分发给 `app.state.backend`，后者要么是 `ForwardBackend` + 要么是 `ReplayBackend`，由启动时的 `_configure_proxy_integrations` 根据 `replay_file` + 是否设置二选一注入。 +- `ForwardBackend` 走 httpx 字节透传：non-stream 是 `await resp.aread()`，stream 是 + `resp.aiter_bytes()` 直接 yield 给客户端，**不**经过任何 SDK 的反序列化/再序列化，所以上游 + 返回的 `reasoning_content` / `provider_specific_fields` 等任意 vendor 字段都不会被吃掉。 + recorder 在另一条独立路径上把字节流喂给 openai SDK 的 stream-state aggregator，仅用于写盘。 +- `ReplayBackend` 完全本地，不持有 httpx client。 + +更深入的代码导览看 [rock/sdk/model/server/api/proxy.py](../../../rock/sdk/model/server/api/proxy.py) +顶部的 module docstring。 diff --git a/pyproject.toml b/pyproject.toml index badb7d1a4b..d7d7a591b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,6 +86,8 @@ model-service = [ "psutil", "swebench", "alibabacloud_cr20181201==2.0.5", + "openai>=1.50.0", + "httpx", ] diff --git a/rock/cli/command/model_service.py b/rock/cli/command/model_service.py index 87e6ca60e6..03cc59582d 100644 --- a/rock/cli/command/model_service.py +++ b/rock/cli/command/model_service.py @@ -82,6 +82,8 @@ async def arun(self, args: argparse.Namespace): proxy_base_url=args.proxy_base_url, retryable_status_codes=args.retryable_status_codes, request_timeout=args.request_timeout, + recording_file=args.recording_file, + replay_file=args.replay_file, ) logger.info(f"model service started, pid: {pid}") with open(self.DEFAULT_MODEL_SERVICE_PID_FILE, "w") as f: @@ -178,6 +180,18 @@ async def add_parser_to(subparsers: argparse._SubParsersAction): default=None, help="Request timeout in seconds. Overrides config file.", ) + start_parser.add_argument( + "--recording-file", + type=str, + default=None, + help="Proxy mode only: where to write the trajectory JSONL. Defaults to LOG_DIR/LLMTraj.jsonl.", + ) + start_parser.add_argument( + "--replay-file", + type=str, + default=None, + help="Proxy mode only: replay from a recorded .jsonl traj file. Mutually exclusive with --recording-file.", + ) watch_agent_parser = model_service_subparsers.add_parser( "watch-agent", diff --git a/rock/sdk/model/server/api/proxy.py b/rock/sdk/model/server/api/proxy.py index fb2b7bec3c..73f74e3f62 100644 --- a/rock/sdk/model/server/api/proxy.py +++ b/rock/sdk/model/server/api/proxy.py @@ -1,13 +1,45 @@ +"""OpenAI-compatible chat/completions proxy with trajectory record/replay. + +Two backends share the ``/v1/chat/completions`` route: + +1. **ForwardBackend** (default) — body bytes are POSTed verbatim to the + configured upstream via plain ``httpx``. The upstream response is forwarded + byte-for-byte back to the client (raw JSON for non-stream, raw SSE bytes + for stream). On the side we run a parser (``ChatCompletionChunk`` + + ``ChatCompletionStreamState`` from the openai SDK) to aggregate streaming + chunks into a final ChatCompletion that the recorder writes to JSONL. The + forward path itself does NOT depend on OpenAI types — anything the upstream + returns (provider-specific ``reasoning_content``, ``citations``, ...) is + passed through untouched. + +2. **ReplayBackend** (``replay_file`` set) — the request is served + directly from the next record in the ``SequentialCursor`` without any + upstream call. Streaming emits the recorded response as one SSE chunk + + ``[DONE]``. +""" + +from __future__ import annotations + +import asyncio +import json +import random +import time +from collections.abc import AsyncIterator from typing import Any import httpx from fastapi import APIRouter, HTTPException, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse from rock.logger import init_logger from rock.sdk.model.server.config import ModelServiceConfig -from rock.sdk.model.server.utils import record_traj -from rock.utils import retry_async +from rock.sdk.model.server.sse import ( + SSE_DONE, + completion_to_chunk_dict, + encode_sse_event, + parse_sse_data_chunks, +) +from rock.sdk.model.server.traj import SequentialCursor, TrajectoryExhausted, TrajectoryRecorder logger = init_logger(__name__) @@ -15,111 +47,352 @@ proxy_router = APIRouter() -# Global HTTP client with a persistent connection pool -http_client = httpx.AsyncClient() +# Headers we never forward upstream: +# - host / content-length: rebuilt by httpx for the upstream request +# - transfer-encoding / connection: RFC 7230 hop-by-hop, scoped to one connection +_HEADERS_NOT_TO_FORWARD = frozenset({"host", "content-length", "transfer-encoding", "connection"}) +# Retry knobs for upstream POST. Read at call-time so tests can monkeypatch them. +# Default: up to 6 attempts with exponential backoff (2s → 4s → 8s → 16s → 32s, jittered). +_RETRY_MAX_ATTEMPTS = 6 +_RETRY_DELAY_SECONDS = 2.0 +_RETRY_BACKOFF = 2.0 -@retry_async( - max_attempts=6, - delay_seconds=2.0, - backoff=2.0, # Exponential backoff (2s, 4s, 8s, 16s, 32s). - jitter=True, # Adds randomness to prevent "thundering herd" effect on the backend. - exceptions=(httpx.TimeoutException, httpx.ConnectError, httpx.HTTPStatusError), -) -async def perform_llm_request(url: str, body: dict, headers: dict, config: ModelServiceConfig): - """ - Forwards the request and triggers retry ONLY if the status code - is in the explicit retryable whitelist. + +async def _send_with_retry( + client: httpx.AsyncClient, + url: str, + *, + body_bytes: bytes, + headers: dict[str, str], + retryable_codes: list[int], +) -> httpx.Response: + """POST with retry on connection errors and whitelisted statuses, returning + an open streaming response. + + Always uses ``stream=True`` so the same path serves both stream and non-stream + callers — non-stream just calls ``await resp.aread()`` to materialize the body. + Assumes a failed upstream returns its error body before any byte is yielded + to downstream (so retry can still discard it cleanly). + + Caller MUST ``await resp.aclose()`` after consuming. """ - response = await http_client.post(url, json=body, headers=headers, timeout=config.request_timeout) - status_code = response.status_code + last_exc: Exception | None = None + delay = _RETRY_DELAY_SECONDS + for attempt in range(1, _RETRY_MAX_ATTEMPTS + 1): + try: + resp = await client.send( + client.build_request("POST", url, content=body_bytes, headers=headers), + stream=True, + ) + except (httpx.TimeoutException, httpx.ConnectError) as exc: + last_exc = exc + if attempt >= _RETRY_MAX_ATTEMPTS: + raise + logger.warning(f"connect failed (attempt {attempt}/{_RETRY_MAX_ATTEMPTS}): {exc}") + await asyncio.sleep(random.uniform(0, delay * 2)) + delay *= _RETRY_BACKOFF + continue - # Check against the explicit whitelist - if status_code in config.retryable_status_codes: - logger.warning(f"Retryable error detected: {status_code}. Triggering retry for {url}...") - response.raise_for_status() + if resp.status_code in retryable_codes and attempt < _RETRY_MAX_ATTEMPTS: + await resp.aclose() + logger.warning(f"upstream status {resp.status_code}, retry {attempt}/{_RETRY_MAX_ATTEMPTS}") + await asyncio.sleep(random.uniform(0, delay * 2)) + delay *= _RETRY_BACKOFF + continue - return response + return resp + raise last_exc # pragma: no cover # unreachable -def get_base_url(model_name: str, config: ModelServiceConfig) -> str: - """ - Selects the target backend URL based on model name matching. - If proxy_base_url is configured, it takes precedence over proxy_rules. - """ - # If direct proxy base URL is configured, return it directly (bypass model name matching) - if config.proxy_base_url: - return config.proxy_base_url.rstrip("/") - - if not model_name: - raise HTTPException(status_code=400, detail="Model name is required for routing.") - - rules = config.proxy_rules - base_url = rules.get(model_name) or rules.get("default") - if not base_url: - raise HTTPException( - status_code=400, detail=f"Model '{model_name}' is not configured and no 'default' rule found." +def _filter_headers(headers) -> dict[str, str]: + """Drop headers that are scoped to the client↔proxy hop or rebuilt by httpx. + ``Authorization`` is forwarded verbatim — proxy stays stateless about which + API key the client uses.""" + out = {} + for key, value in headers.items(): + if key.lower() in _HEADERS_NOT_TO_FORWARD: + continue + out[key] = value + return out + + +class ReplayBackend: + """Serves requests from a pre-recorded trajectory; no upstream calls made.""" + + def __init__(self, cursor: SequentialCursor) -> None: + self._cursor = cursor + + async def serve(self, *, model_name: str, is_stream: bool, **_: Any) -> Response: + try: + record = await self._cursor.next(expected_model=model_name) + except TrajectoryExhausted as exc: + raise HTTPException(status_code=404, detail=str(exc)) + + response_dict = record.get("response") + if not isinstance(response_dict, dict): + raise HTTPException( + status_code=500, + detail=f"replay record at step {self._cursor.position - 1} has no usable response dict", + ) + logger.info(f"[replay] step {self._cursor.position}/{self._cursor.total} served for model={model_name!r}") + + if is_stream: + return StreamingResponse( + self._sse_iter(response_dict, model=model_name), + media_type="text/event-stream", + ) + return JSONResponse(status_code=200, content=response_dict) + + @staticmethod + async def _sse_iter(response: dict, *, model: str) -> AsyncIterator[bytes]: + """Emit a recorded response as one SSE chunk + ``[DONE]``.""" + yield encode_sse_event(completion_to_chunk_dict(response, model=model)) + yield SSE_DONE + + +class ForwardBackend: + """Forwards requests byte-for-byte to the upstream and optionally records the trajectory.""" + + def __init__(self, config: ModelServiceConfig, recorder: TrajectoryRecorder | None = None) -> None: + self._config = config + self._recorder = recorder + + def _resolve_base_url(self, model_name: str) -> str: + """Pick the upstream base URL by model name. + + ``proxy_base_url`` takes precedence; falls back to ``proxy_rules[model]`` and + then ``proxy_rules["default"]``. Trailing slashes are stripped so the caller + can append ``/chat/completions`` directly. + """ + if self._config.proxy_base_url: + return self._config.proxy_base_url.rstrip("/") + + if not model_name: + raise HTTPException(status_code=400, detail="Model name is required for routing.") + + rules = self._config.proxy_rules + base_url = rules.get(model_name) or rules.get("default") + if not base_url: + raise HTTPException( + status_code=400, + detail=f"Model '{model_name}' is not configured and no 'default' rule found.", + ) + + return base_url.rstrip("/") + + async def serve( + self, + *, + model_name: str, + is_stream: bool, + body_bytes: bytes, + fwd_headers: dict[str, str], + request_dict: dict[str, Any], + **_: Any, + ) -> Response: + upstream_url = f"{self._resolve_base_url(model_name)}/chat/completions" + logger.info(f"Routing model {model_name!r} to {upstream_url}") + + if is_stream: + return StreamingResponse( + self._stream_and_record( + upstream_url=upstream_url, + body_bytes=body_bytes, + fwd_headers=fwd_headers, + request_dict=request_dict, + ), + media_type="text/event-stream", + ) + + # Non-stream: same retry path as stream (open with stream=True), then aread() the body. + start = time.time() + async with httpx.AsyncClient(timeout=self._config.request_timeout) as client: + try: + resp = await _send_with_retry( + client, + upstream_url, + body_bytes=body_bytes, + headers=fwd_headers, + retryable_codes=self._config.retryable_status_codes, + ) + except httpx.TimeoutException as exc: + if self._recorder is not None: + await self._recorder.record( + request=request_dict, + response=None, + status="failure", + start_time=start, + end_time=time.time(), + error=f"timeout: {exc}", + ) + raise HTTPException(status_code=504, detail=f"Upstream timed out: {exc}") + except httpx.RequestError as exc: + if self._recorder is not None: + await self._recorder.record( + request=request_dict, + response=None, + status="failure", + start_time=start, + end_time=time.time(), + error=f"{type(exc).__name__}: {exc}", + ) + raise HTTPException(status_code=502, detail=f"Upstream request failed: {exc}") + + try: + response_bytes = await resp.aread() + status_code = resp.status_code + content_type = resp.headers.get("content-type", "application/json") + finally: + await resp.aclose() + + response_text = response_bytes.decode("utf-8", errors="replace") + response_dict: dict | None = None + try: + parsed = json.loads(response_text) if response_text else None + if isinstance(parsed, dict): + response_dict = parsed + except json.JSONDecodeError: + pass + + if self._recorder is not None: + await self._recorder.record( + request=request_dict, + response=response_dict, + status="success" if status_code < 400 else "failure", + start_time=start, + end_time=time.time(), + error=None if status_code < 400 else f"upstream_status={status_code}", + ) + + # Forward bytes verbatim — preserves any provider-specific fields untouched. + return Response(content=response_bytes, status_code=status_code, media_type=content_type) + + async def _stream_and_record( + self, + *, + upstream_url: str, + body_bytes: bytes, + fwd_headers: dict[str, str], + request_dict: dict[str, Any], + ) -> AsyncIterator[bytes]: + """SSE bytes are forwarded verbatim; chunks are parsed in parallel and + aggregated into the final ChatCompletion that the recorder writes to JSONL. + + Retry on connection errors and whitelisted statuses happens BEFORE any byte + is yielded; mid-stream connection drops are not retried (would corrupt the + client transmission).""" + # openai SDK is used purely as a stream-aggregation parser — keep the import + # local so module load doesn't pull it in for callers that never stream. + from openai.lib.streaming.chat import ChatCompletionStreamState + from openai.types.chat import ChatCompletionChunk + + state = ChatCompletionStreamState() + start = time.time() + parse_buffer = b"" + upstream_status = 0 + + async with httpx.AsyncClient(timeout=self._config.request_timeout) as client: + try: + resp = await _send_with_retry( + client, + upstream_url, + body_bytes=body_bytes, + headers=fwd_headers, + retryable_codes=self._config.retryable_status_codes, + ) + except (httpx.TimeoutException, httpx.ConnectError) as exc: + if self._recorder is not None: + await self._recorder.record( + request=request_dict, + response=None, + status="failure", + start_time=start, + end_time=time.time(), + error=f"{type(exc).__name__}: {exc}", + ) + return + + try: + upstream_status = resp.status_code + async for chunk in resp.aiter_bytes(): + yield chunk + chunk_dicts, parse_buffer = parse_sse_data_chunks(parse_buffer + chunk) + for chunk_dict in chunk_dicts: + try: + state.handle_chunk(ChatCompletionChunk.model_validate(chunk_dict)) + except Exception as exc: # parser error: forward continues, traj will be partial + logger.debug(f"[record] chunk parse failed (forward continues): {exc}") + except httpx.RequestError as exc: + # Connection died mid-stream — bytes already sent reach the client; + # record what we got and return. + if self._recorder is not None: + await self._recorder.record( + request=request_dict, + response=None, + status="failure", + start_time=start, + end_time=time.time(), + error=f"{type(exc).__name__}: {exc}", + ) + return + finally: + await resp.aclose() + + if self._recorder is None: + return + + status = "success" if upstream_status < 400 else "failure" + final_dict: dict | None = None + if status == "success": + try: + final_dict = state.get_final_completion().model_dump() + except Exception as exc: + logger.warning(f"[record] stream aggregation failed: {exc}") + + await self._recorder.record( + request=request_dict, + response=final_dict, + status=status, + start_time=start, + end_time=time.time(), + error=None if status == "success" else f"upstream_status={upstream_status}", ) - return base_url.rstrip("/") + +CompletionBackend = ReplayBackend | ForwardBackend + + +def _get_backend(request: Request) -> CompletionBackend: + """Typed accessor for the backend attached at startup by ``_configure_proxy_integrations``.""" + return request.app.state.backend @proxy_router.post("/v1/chat/completions") -@record_traj -async def chat_completions(body: dict[str, Any], request: Request): - """ - OpenAI-compatible chat completions proxy endpoint. - Handles routing, header transparent forwarding, and automatic retries. - """ - config = request.app.state.model_service_config - - # Step 1: Model Routing - model_name = body.get("model", "") - base_url = get_base_url(model_name, config) - target_url = f"{base_url}/chat/completions" - logger.info(f"Routing model '{model_name}' to URL: {target_url}") - - # Step 2: Header Cleaning - # Preserve 'Authorization' for authentication while removing hop-by-hop transport headers. - forwarded_headers = {} - for key, value in request.headers.items(): - if key.lower() in ["host", "content-length", "content-type", "transfer-encoding"]: - continue - forwarded_headers[key] = value - - # Step 3: Strategy Enforcement - # Force non-streaming mode for the MVP phase to ensure stability. - if body.get("stream") is True: - raise HTTPException( - status_code=400, - detail="Streaming requests (stream=True) are not supported in the current version. Please set stream=False or omit the stream parameter.", - ) - body["stream"] = False +async def chat_completions(request: Request): + """OpenAI-compatible chat completions proxy endpoint. + Reads the body as raw bytes (no parsing on the forward path) and delegates + to the backend attached at startup (replay or forward). + """ + body_bytes = await request.body() try: - # Step 4: Execute Request with Retry Logic - response = await perform_llm_request(target_url, body, forwarded_headers, config) - return JSONResponse(status_code=response.status_code, content=response.json()) - - except httpx.HTTPStatusError as e: - # Forward the raw backend error message to the client. - # This allows the Agent-side logic to detect keywords like 'context length exceeded' - # or 'content violation' and raise appropriate exceptions. - error_text = e.response.text if e.response else "No error details" - status_code = e.response.status_code if e.response else 502 - logger.error(f"Final failure after retries. Status: {status_code}, Response: {error_text}") - return JSONResponse( - status_code=status_code, - content={ - "error": { - "message": f"LLM backend error: {error_text}", - "type": "proxy_retry_failed", - "code": status_code, - } - }, - ) - except Exception as e: - logger.error(f"Unexpected proxy error: {str(e)}") - # Raise standard 500 for non-HTTP related errors or system errors - raise HTTPException(status_code=500, detail=str(e)) + request_dict = json.loads(body_bytes) if body_bytes else {} + except json.JSONDecodeError: + raise HTTPException(status_code=400, detail="Request body is not valid JSON.") + if not isinstance(request_dict, dict): + raise HTTPException(status_code=400, detail="Request body must be a JSON object.") + + model_name = request_dict.get("model", "") + is_stream = bool(request_dict.get("stream")) + fwd_headers = _filter_headers(request.headers) + + backend = _get_backend(request) + return await backend.serve( + model_name=model_name, + is_stream=is_stream, + body_bytes=body_bytes, + fwd_headers=fwd_headers, + request_dict=request_dict, + ) diff --git a/rock/sdk/model/server/config.py b/rock/sdk/model/server/config.py index 2c96992b5c..e734c29878 100644 --- a/rock/sdk/model/server/config.py +++ b/rock/sdk/model/server/config.py @@ -1,7 +1,7 @@ from pathlib import Path import yaml -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator from rock import env_vars @@ -27,6 +27,10 @@ class ModelServiceConfig(BaseModel): """Configuration for the LLM Model Service.""" + # validate_assignment=True so the recording/replay mutex below also fires when + # CLI overrides are applied field-by-field (not only at construction time). + model_config = ConfigDict(validate_assignment=True) + host: str = "0.0.0.0" """Server host address.""" @@ -51,6 +55,23 @@ class ModelServiceConfig(BaseModel): request_timeout: int = Field(default=120) """Request timeout in seconds.""" + recording_file: str | None = Field(default=None) + """Recording mode output: where ForwardBackend writes the trajectory JSONL. + None → uses TRAJ_FILE (LOG_DIR/LLMTraj.jsonl).""" + + replay_file: str | None = Field(default=None) + """Replay mode input: a .jsonl trajectory file. When set, ReplayBackend serves + requests from recorded responses instead of calling a real upstream.""" + + @model_validator(mode="after") + def _recording_replay_mutually_exclusive(self): + if self.recording_file and self.replay_file: + raise ValueError( + "recording_file and replay_file are mutually exclusive — " + "set one (recording mode) or the other (replay mode), not both." + ) + return self + @classmethod def from_file(cls, config_path: str | None = None): """ diff --git a/rock/sdk/model/server/main.py b/rock/sdk/model/server/main.py index 7f8dabebe2..89e87ac0f9 100644 --- a/rock/sdk/model/server/main.py +++ b/rock/sdk/model/server/main.py @@ -11,7 +11,7 @@ from rock.logger import init_logger from rock.sdk.model.server.api.local import init_local_api, local_router from rock.sdk.model.server.api.proxy import proxy_router -from rock.sdk.model.server.config import ModelServiceConfig +from rock.sdk.model.server.config import TRAJ_FILE, ModelServiceConfig # Configure logging logger = init_logger(__name__) @@ -52,6 +52,33 @@ async def global_exception_handler(request, exc): return app +def _configure_proxy_integrations(app: FastAPI, config: ModelServiceConfig) -> None: + """Attach the appropriate backend to ``app.state.backend``. + + - Replay mode (``replay_file`` set): ``ReplayBackend`` wrapping a + ``SequentialCursor``; no recorder — replaying back into the source file + would corrupt it. + - Forward mode (default): ``ForwardBackend`` with a ``TrajectoryRecorder`` + writing to ``recording_file`` (or ``TRAJ_FILE`` if unset). + """ + from rock.sdk.model.server.api.proxy import ForwardBackend, ReplayBackend + + if config.replay_file: + from rock.sdk.model.server.traj import SequentialCursor + + cursor = SequentialCursor.load(config.replay_file) + app.state.backend = ReplayBackend(cursor) + logger.info(f"replay backend attached, replay_file={config.replay_file}") + return + + from rock.sdk.model.server.traj import TrajectoryRecorder + + recording_path = config.recording_file or TRAJ_FILE + recorder = TrajectoryRecorder(traj_file=recording_path) + app.state.backend = ForwardBackend(config, recorder=recorder) + logger.info(f"forward backend attached, recording_file={recording_path}") + + def main( model_servie_type: str, config: ModelServiceConfig, @@ -63,6 +90,7 @@ def main( asyncio.run(init_local_api()) app.include_router(local_router, prefix="", tags=["local"]) else: + _configure_proxy_integrations(app, config) app.include_router(proxy_router, prefix="", tags=["proxy"]) logger.info(f"Starting LLM Service on {config.host}:{config.port}, type: {model_servie_type}") @@ -100,6 +128,12 @@ def create_config_from_args(args) -> ModelServiceConfig: if args.request_timeout: config.request_timeout = args.request_timeout logger.info(f"request_timeout set from command line: {args.request_timeout}s") + if args.recording_file: + config.recording_file = args.recording_file + logger.info(f"recording_file set from command line: {args.recording_file}") + if args.replay_file: + config.replay_file = args.replay_file + logger.info(f"replay mode enabled via --replay-file: {args.replay_file}") return config @@ -142,6 +176,18 @@ def create_config_from_args(args) -> ModelServiceConfig: parser.add_argument( "--request-timeout", type=int, default=None, help="Request timeout in seconds. Overrides config file." ) + parser.add_argument( + "--recording-file", + type=str, + default=None, + help="Forward mode: where to write the trajectory JSONL. Defaults to TRAJ_FILE.", + ) + parser.add_argument( + "--replay-file", + type=str, + default=None, + help="Replay mode: path to a recorded .jsonl traj file. Disables real LLM upstreams.", + ) args = parser.parse_args() config = create_config_from_args(args) diff --git a/rock/sdk/model/server/sse.py b/rock/sdk/model/server/sse.py new file mode 100644 index 0000000000..f1cca034e6 --- /dev/null +++ b/rock/sdk/model/server/sse.py @@ -0,0 +1,99 @@ +"""SSE codec utilities for the chat/completions proxy. + +Three pure helpers, no openai/litellm dependencies: + +- :func:`parse_sse_data_chunks` — incremental SSE byte stream → list of decoded + ``data:`` payload dicts (used by the forward path to feed chunks into the + stream-state aggregator while bytes pass through verbatim to the client). +- :func:`completion_to_chunk_dict` — convert a non-streaming ``chat.completion`` + response into a single ``chat.completion.chunk`` dict, by renaming + ``message`` → ``delta``. Used by the replay path's streaming output. +- :func:`encode_sse_event` — encode a payload dict as ``data: \\n\\n`` + bytes (one SSE event). +""" + +from __future__ import annotations + +import json +import time +import uuid +from typing import Final + +# Terminal SSE event sent at the end of a chat/completions stream. +SSE_DONE: Final[bytes] = b"data: [DONE]\n\n" + + +def parse_sse_data_chunks(buffer: bytes) -> tuple[list[dict], bytes]: + """Extract complete SSE events from a (possibly partial) byte buffer. + + Returns ``(chunks, leftover)``: the parsed ``data:`` JSON payload dicts and + the bytes that did not yet form a complete event (``\\n\\n``-terminated). + + - ``data: [DONE]`` is skipped (terminal marker, has no JSON payload). + - Lines that don't start with ``data:`` (``event:`` / ``id:`` / blank) + are ignored. + - Malformed JSON in a ``data:`` line is silently skipped — caller logs at + its own discretion (typically ``debug``). + + Caller pattern:: + + chunks, buffer = parse_sse_data_chunks(buffer + new_bytes) + for chunk_dict in chunks: + ... feed to aggregator, etc ... + """ + chunks: list[dict] = [] + while b"\n\n" in buffer: + event, buffer = buffer.split(b"\n\n", 1) + for raw_line in event.split(b"\n"): + line = raw_line.decode("utf-8", errors="replace").strip() + if not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if not payload or payload == "[DONE]": + continue + try: + chunks.append(json.loads(payload)) + except json.JSONDecodeError: + continue + return chunks, buffer + + +def completion_to_chunk_dict(response: dict, *, model: str) -> dict: + """Convert a recorded ``chat.completion`` dict into a single + ``chat.completion.chunk`` dict, suitable for re-streaming. + + Only ``message`` → ``delta`` is renamed; every other field (including + provider-specific extras like ``reasoning_content`` inside the message) + flows through unchanged. ``id`` / ``created`` are synthesized when missing. + + ``tool_calls`` items get a positional ``index`` injected if missing — the + OpenAI streaming spec requires it on chunk deltas (a recorded non-stream + ``message.tool_calls`` carries no ``index``, but downstream stream parsers + e.g. the openai SDK will reject the chunk without one). + """ + choices_in = response.get("choices") or [] + choices_out = [] + for choice in choices_in: + delta = dict(choice.get("message") or {}) + if "tool_calls" in delta and delta["tool_calls"]: + delta["tool_calls"] = [{"index": tc.get("index", i), **tc} for i, tc in enumerate(delta["tool_calls"])] + choices_out.append( + { + "index": choice.get("index", 0), + "delta": delta, + "finish_reason": choice.get("finish_reason"), + "logprobs": choice.get("logprobs"), + } + ) + return { + "id": response.get("id") or f"chatcmpl-{uuid.uuid4()}", + "object": "chat.completion.chunk", + "created": response.get("created") or int(time.time()), + "model": response.get("model") or model, + "choices": choices_out, + } + + +def encode_sse_event(data: dict) -> bytes: + """Encode a JSON payload as one SSE ``data:`` event (terminated by ``\\n\\n``).""" + return f"data: {json.dumps(data, ensure_ascii=False)}\n\n".encode() diff --git a/rock/sdk/model/server/traj.py b/rock/sdk/model/server/traj.py new file mode 100644 index 0000000000..e12c229c7f --- /dev/null +++ b/rock/sdk/model/server/traj.py @@ -0,0 +1,156 @@ +"""Trajectory record + replay for the chat/completions proxy. + +Two halves around the same JSONL schema (one record per line): + +- :class:`TrajectoryRecorder` — invoked by the forward path after each upstream + call (success or failure). Appends a small dict with + ``request`` / ``response`` / ``status`` / ``response_time`` / ``model`` / + ``stream``, and reports OTLP RT/count metrics. Stores responses verbatim + (provider-specific fields like ``reasoning_content`` survive); for streaming + calls ``response`` is the aggregated final ChatCompletion produced by + ``ChatCompletionStreamState.get_final_completion().model_dump()``. + +- :class:`SequentialCursor` — loads a JSONL trajectory once at startup; + ``await cursor.next(expected_model=...)`` hands out the next record (full + payload dict) and advances. Going past the end raises + :class:`TrajectoryExhausted` so the proxy can return a clean 404. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from typing import Any + +from rock.logger import init_logger +from rock.sdk.model.server.utils import ( + MODEL_SERVICE_REQUEST_COUNT, + MODEL_SERVICE_REQUEST_RT, + _get_or_create_metrics_monitor, +) + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Recorder +# --------------------------------------------------------------------------- + + +class TrajectoryRecorder: + """Appends one JSONL line per chat/completions call and reports OTLP metrics.""" + + def __init__(self, traj_file: str | os.PathLike) -> None: + self.traj_file = Path(traj_file) + self.traj_file.parent.mkdir(parents=True, exist_ok=True) + self._lock = asyncio.Lock() + self._monitor = _get_or_create_metrics_monitor() + + async def record( + self, + *, + request: dict[str, Any], + response: dict[str, Any] | None, + status: str, + start_time: float, + end_time: float, + error: str | None = None, + ) -> None: + rt_seconds = end_time - start_time + payload = { + "model": request.get("model"), + "stream": bool(request.get("stream")), + "status": status, + "response_time": rt_seconds, + "start_time": start_time, + "end_time": end_time, + "request": request, + "response": response, + "error": error, + } + + line = json.dumps(payload, ensure_ascii=False, default=str) + "\n" + async with self._lock: + await asyncio.to_thread(self._write_line, line) + + attrs = { + "type": "chat_completions", + "status": status, + "sandbox_id": os.getenv("ROCK_SANDBOX_ID", "unknown"), + } + self._monitor.record_gauge_by_name(MODEL_SERVICE_REQUEST_RT, rt_seconds * 1000.0, attributes=attrs) + self._monitor.record_counter_by_name(MODEL_SERVICE_REQUEST_COUNT, 1, attributes=attrs) + + def _write_line(self, line: str) -> None: + with self.traj_file.open("a", encoding="utf-8") as f: + f.write(line) + + +# --------------------------------------------------------------------------- +# Replay cursor +# --------------------------------------------------------------------------- + + +class TrajectoryExhausted(Exception): + """Raised by ``SequentialCursor.next`` when all recorded steps have been served.""" + + def __init__(self, position: int, total: int) -> None: + super().__init__(f"trajectory exhausted at step {position} (total recorded steps={total})") + self.position = position + self.total = total + + +class SequentialCursor: + """Hands out trajectory records one at a time, in recorded order.""" + + def __init__(self, records: list[dict]) -> None: + self.records = records + self._idx = 0 + self._lock = asyncio.Lock() + + @classmethod + def load(cls, path: str | os.PathLike) -> SequentialCursor: + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(f"traj file not found: {path}") + + records: list[dict] = [] + with path.open("r", encoding="utf-8") as fp: + for line in fp: + line = line.strip() + if not line: + continue + records.append(json.loads(line)) + + logger.info(f"[traj-replay] loaded {len(records)} record(s) from {path}") + return cls(records) + + async def next(self, expected_model: str | None = None) -> dict: + async with self._lock: + if self._idx >= len(self.records): + raise TrajectoryExhausted(position=self._idx, total=len(self.records)) + record = self.records[self._idx] + self._idx += 1 + current_idx = self._idx - 1 + + if expected_model: + recorded_model = record.get("model") + if recorded_model and recorded_model != expected_model: + logger.warning( + f"[traj-replay] step {current_idx} model mismatch: " + f"recorded={recorded_model!r} requested={expected_model!r}" + ) + return record + + def reset(self) -> None: + self._idx = 0 + + @property + def position(self) -> int: + return self._idx + + @property + def total(self) -> int: + return len(self.records) diff --git a/rock/sdk/model/server/utils.py b/rock/sdk/model/server/utils.py index 20ae8896dc..639ca3995b 100644 --- a/rock/sdk/model/server/utils.py +++ b/rock/sdk/model/server/utils.py @@ -38,7 +38,7 @@ def _write_traj(data: dict): def record_traj(func: Callable): - """Decorator to record chat completions input/output as traj.""" + """Decorator to record chat completions input/output as traj (local mode only).""" @wraps(func) async def wrapper(*args, **kwargs): diff --git a/rock/sdk/model/service.py b/rock/sdk/model/service.py index b1b523ed27..24cd7ede38 100644 --- a/rock/sdk/model/service.py +++ b/rock/sdk/model/service.py @@ -17,6 +17,8 @@ def start_sandbox_service( proxy_base_url: str | None = None, retryable_status_codes: str | None = None, request_timeout: int | None = None, + recording_file: str | None = None, + replay_file: str | None = None, ) -> subprocess.Popen: """start sandbox service""" current_file = Path(__file__).resolve() @@ -38,6 +40,10 @@ def start_sandbox_service( cmd.extend(["--retryable-status-codes", retryable_status_codes]) if request_timeout: cmd.extend(["--request-timeout", str(request_timeout)]) + if recording_file: + cmd.extend(["--recording-file", recording_file]) + if replay_file: + cmd.extend(["--replay-file", replay_file]) process = subprocess.Popen(cmd, cwd=str(service_dir)) return process @@ -51,6 +57,8 @@ async def start( proxy_base_url: str | None = None, retryable_status_codes: str | None = None, request_timeout: int | None = None, + recording_file: str | None = None, + replay_file: str | None = None, ) -> str: process = self.start_sandbox_service( model_service_type=model_service_type, @@ -60,6 +68,8 @@ async def start( proxy_base_url=proxy_base_url, retryable_status_codes=retryable_status_codes, request_timeout=request_timeout, + recording_file=recording_file, + replay_file=replay_file, ) pid = process.pid diff --git a/tests/unit/cli/command/test_model_service.py b/tests/unit/cli/command/test_model_service.py new file mode 100644 index 0000000000..86849c718b --- /dev/null +++ b/tests/unit/cli/command/test_model_service.py @@ -0,0 +1,120 @@ +"""Unit tests for rock.cli.command.model_service.ModelServiceCommand. + +Drive the sub-parser end-to-end with argparse so the surface that users +actually type at the terminal is what we exercise. ``ModelService.start`` is +mocked — these tests assert wiring (argparse → handler → SDK call), not the +subprocess command construction (covered separately in +tests/unit/sdk/model/test_service.py). +""" + +from __future__ import annotations + +import argparse +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from rock.cli.command.model_service import ModelServiceCommand + + +def _build_parser() -> argparse.ArgumentParser: + """Top-level parser with `model-service` subcommand wired in, same as the CLI.""" + top = argparse.ArgumentParser(prog="rock") + subparsers = top.add_subparsers(dest="command") + asyncio.run(ModelServiceCommand.add_parser_to(subparsers)) + return top + + +@pytest.fixture +def isolate_pid_file(monkeypatch, tmp_path): + """Redirect PID dir/file into tmp so arun() doesn't touch ./data/cli/model.""" + monkeypatch.setattr(ModelServiceCommand, "DEFAULT_MODEL_SERVICE_DIR", str(tmp_path)) + monkeypatch.setattr(ModelServiceCommand, "DEFAULT_MODEL_SERVICE_PID_FILE", str(tmp_path / "pid.txt")) + + +@pytest.fixture +def fake_start(monkeypatch): + """Replace ModelService.start with an AsyncMock returning a fixed pid.""" + mock = AsyncMock(return_value="12345") + monkeypatch.setattr("rock.cli.command.model_service.ModelService.start", mock) + return mock + + +# ---------- argparse: the new flags must parse ---------- + + +def test_recording_file_flag_parses(): + parser = _build_parser() + ns = parser.parse_args(["model-service", "start", "--type", "proxy", "--recording-file", "/tmp/out.jsonl"]) + assert ns.recording_file == "/tmp/out.jsonl" + assert ns.replay_file is None + + +def test_replay_file_flag_parses(): + parser = _build_parser() + ns = parser.parse_args(["model-service", "start", "--type", "proxy", "--replay-file", "/tmp/in.jsonl"]) + assert ns.replay_file == "/tmp/in.jsonl" + assert ns.recording_file is None + + +def test_neither_flag_defaults_to_none(): + parser = _build_parser() + ns = parser.parse_args(["model-service", "start", "--type", "proxy"]) + assert ns.recording_file is None + assert ns.replay_file is None + + +# ---------- handler: passes parsed args through to ModelService.start ---------- + + +def test_start_handler_forwards_recording_file(isolate_pid_file, fake_start): + parser = _build_parser() + ns = parser.parse_args( + [ + "model-service", + "start", + "--type", + "proxy", + "--proxy-base-url", + "https://api.openai.com/v1", + "--recording-file", + "/tmp/out.jsonl", + ] + ) + asyncio.run(ModelServiceCommand().arun(ns)) + + kwargs = fake_start.call_args.kwargs + assert kwargs["recording_file"] == "/tmp/out.jsonl" + assert kwargs["replay_file"] is None + assert kwargs["proxy_base_url"] == "https://api.openai.com/v1" + assert kwargs["model_service_type"] == "proxy" + + +def test_start_handler_forwards_replay_file(isolate_pid_file, fake_start): + parser = _build_parser() + ns = parser.parse_args( + [ + "model-service", + "start", + "--type", + "proxy", + "--replay-file", + "/tmp/in.jsonl", + ] + ) + asyncio.run(ModelServiceCommand().arun(ns)) + + kwargs = fake_start.call_args.kwargs + assert kwargs["replay_file"] == "/tmp/in.jsonl" + assert kwargs["recording_file"] is None + + +def test_start_handler_omits_both_when_unset(isolate_pid_file, fake_start): + parser = _build_parser() + ns = parser.parse_args(["model-service", "start", "--type", "proxy"]) + asyncio.run(ModelServiceCommand().arun(ns)) + + kwargs = fake_start.call_args.kwargs + assert kwargs["recording_file"] is None + assert kwargs["replay_file"] is None diff --git a/tests/unit/sdk/model/test_proxy.py b/tests/unit/sdk/model/test_proxy.py index edce5584cb..345ea31775 100644 --- a/tests/unit/sdk/model/test_proxy.py +++ b/tests/unit/sdk/model/test_proxy.py @@ -1,14 +1,24 @@ -from unittest.mock import AsyncMock, MagicMock, patch +"""Tests for the chat/completions proxy. + +Forward path is exercised by pointing the proxy at an httpx ``MockTransport`` +(no real network). Replay path is exercised end-to-end via the FastAPI test +client. Config / CLI / metrics-singleton tests round out the file. +""" + +import argparse +import json +from unittest.mock import MagicMock, patch import httpx import pytest import yaml -from fastapi import FastAPI, Request -from httpx import ASGITransport, AsyncClient, HTTPStatusError, Request, Response +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient -from rock.sdk.model.server.api.proxy import perform_llm_request, proxy_router +from rock.sdk.model.server.api.proxy import proxy_router from rock.sdk.model.server.config import ModelServiceConfig from rock.sdk.model.server.main import create_config_from_args, lifespan +from rock.sdk.model.server.traj import SequentialCursor from rock.sdk.model.server.utils import ( MODEL_SERVICE_REQUEST_COUNT, MODEL_SERVICE_REQUEST_RT, @@ -16,361 +26,568 @@ record_traj, ) -# Initialize a temporary FastAPI application for testing the router -test_app = FastAPI() -test_app.include_router(proxy_router) -mock_config = ModelServiceConfig() -test_app.state.model_service_config = mock_config +def _build_app(config: ModelServiceConfig, *, replay_cursor=None, recorder=None) -> FastAPI: + """Build a FastAPI app with the proxy router and the given config attached.""" + from rock.sdk.model.server.api.proxy import ForwardBackend, ReplayBackend + + app = FastAPI() + app.state.model_service_config = config + if replay_cursor is not None: + app.state.backend = ReplayBackend(replay_cursor) + else: + app.state.backend = ForwardBackend(config, recorder=recorder) + app.include_router(proxy_router) + return app + + +def _patch_httpx_with_handler(handler): + """Patch ``proxy.httpx.AsyncClient`` so each ``async with httpx.AsyncClient(...)`` + returns a real client wrapping ``MockTransport(handler)``.""" + real_client_cls = httpx.AsyncClient # capture before patching kicks in + transport = httpx.MockTransport(handler) + + def factory(*args, **kwargs): + kwargs.pop("timeout", None) # transport supplies the response, no timeout needed + return real_client_cls(transport=transport, **kwargs) + + return patch("rock.sdk.model.server.api.proxy.httpx.AsyncClient", side_effect=factory) + + +def _success_response_json(*, model: str = "gpt-3.5-turbo", content: str = "hi") -> dict: + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1234, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +# ---------- Forward path: routing ---------- @pytest.mark.asyncio -async def test_chat_completions_routing_success(): - """ - Test the high-level routing logic. - """ - patch_path = "rock.sdk.model.server.api.proxy.perform_llm_request" - - with patch(patch_path, new_callable=AsyncMock) as mock_request: - mock_resp = MagicMock(spec=Response) - mock_resp.status_code = 200 - mock_resp.json.return_value = {"id": "chat-123", "choices": []} - mock_request.return_value = mock_resp - - transport = ASGITransport(app=test_app) +async def test_forward_routes_by_model_name_to_proxy_rules(): + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=_success_response_json()) + + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as ac: - payload = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}]} - response = await ac.post("/v1/chat/completions", json=payload) + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) - assert response.status_code == 200 - call_args = mock_request.call_args[0] - assert call_args[0] == "https://api.openai.com/v1/chat/completions" - assert mock_request.called + assert r.status_code == 200 + assert captured["url"] == "https://api.openai.com/v1/chat/completions" + assert captured["body"]["model"] == "gpt-3.5-turbo" @pytest.mark.asyncio -async def test_chat_completions_fallback_to_default_when_not_found(): - """ - Test that an unrecognized model name correctly falls back to the 'default' URL. - """ - patch_path = "rock.sdk.model.server.api.proxy.perform_llm_request" - - with patch(patch_path, new_callable=AsyncMock) as mock_request: - mock_resp = MagicMock(spec=Response) - mock_resp.status_code = 200 - mock_resp.json.return_value = {"id": "chat-fallback", "choices": []} - mock_request.return_value = mock_resp - - config = test_app.state.model_service_config - default_base_url = config.proxy_rules["default"].rstrip("/") - expected_target_url = f"{default_base_url}/chat/completions" - - transport = ASGITransport(app=test_app) +async def test_forward_falls_back_to_default_for_unknown_model(): + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + return httpx.Response(200, json=_success_response_json(model="some-random")) + + config = ModelServiceConfig() + expected_default = config.proxy_rules["default"].rstrip("/") + "/chat/completions" + app = _build_app(config) + + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as ac: - payload = { - "model": "some-random-unsupported-model", # This model is NOT in proxy_rules - "messages": [{"role": "user", "content": "hello"}], - } - response = await ac.post("/v1/chat/completions", json=payload) + r = await ac.post( + "/v1/chat/completions", + json={"model": "some-random", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200 + assert captured["url"] == expected_default - assert response.status_code == 200 - # Verify that perform_llm_request was called with the DEFAULT URL - call_args = mock_request.call_args[0] - actual_url = call_args[0] +@pytest.mark.asyncio +async def test_forward_400_when_no_rule_and_no_default(): + config = ModelServiceConfig() + config.proxy_rules = {} + app = _build_app(config) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "any", "messages": [{"role": "user", "content": "hi"}]}, + ) - assert actual_url == expected_target_url - assert mock_request.called + assert r.status_code == 400 + assert "not configured" in r.json()["detail"] @pytest.mark.asyncio -async def test_chat_completions_routing_absolute_fail(): - """ - Test that both the specific model and the 'default' rule are missing. - """ - empty_config = ModelServiceConfig() - empty_config.proxy_rules = {} - - with patch.object(test_app.state, "model_service_config", empty_config): - transport = ASGITransport(app=test_app) +async def test_forward_proxy_base_url_overrides_proxy_rules(): + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + return httpx.Response(200, json=_success_response_json()) + + config = ModelServiceConfig() + config.proxy_base_url = "https://custom-endpoint.example.com/v1" + app = _build_app(config) + + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as ac: - payload = {"model": "any-model", "messages": [{"role": "user", "content": "hello"}]} - response = await ac.post("/v1/chat/completions", json=payload) + await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) - assert response.status_code == 400 - detail = response.json()["detail"] - assert "not configured" in detail + assert captured["url"] == "https://custom-endpoint.example.com/v1/chat/completions" -@pytest.mark.asyncio -async def test_perform_llm_request_retry_on_whitelist(): - """ - Test that the proxy retries when receiving a whitelisted error code. - """ - client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" +# ---------- Forward path: byte passthrough ---------- - # Patch asyncio.sleep inside the retry module to avoid actual waiting - with ( - patch(client_post_path, new_callable=AsyncMock) as mock_post, - patch("rock.utils.retry.asyncio.sleep", return_value=None), - ): - # 1. Setup Failed Response (429) - resp_429 = MagicMock(spec=Response) - resp_429.status_code = 429 - error_429 = HTTPStatusError("Rate Limited", request=MagicMock(spec=Request), response=resp_429) - # 2. Setup Success Response (200) - resp_200 = MagicMock(spec=Response) - resp_200.status_code = 200 - resp_200.json.return_value = {"ok": True} +@pytest.mark.asyncio +async def test_forward_response_body_is_byte_for_byte_passthrough(): + """Upstream's exact JSON bytes (incl. provider-specific fields) reach the client.""" + upstream_payload = { + "id": "x", + "object": "chat.completion", + "model": "glm-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi", "reasoning_content": "...think..."}, + "finish_reason": "stop", + } + ], + "provider_specific_fields": {"vendor_field": "vendor_value"}, + } - # Sequence: Fail with 429, then Succeed with 200 - mock_post.side_effect = [error_429, resp_200] + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=upstream_payload) - result = await perform_llm_request("http://fake.url", {}, {}, mock_config) + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "glm-5", "messages": [{"role": "user", "content": "hi"}]}, + ) - assert result.status_code == 200 - assert mock_post.call_count == 2 + body = r.json() + assert body["choices"][0]["message"]["reasoning_content"] == "...think..." + assert body["provider_specific_fields"] == {"vendor_field": "vendor_value"} @pytest.mark.asyncio -async def test_perform_llm_request_no_retry_on_non_whitelist(): - """ - Test that the proxy DOES NOT retry for non-retryable codes (e.g., 401). - It should return the error response immediately. - """ - client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" +async def test_forward_propagates_upstream_status_and_body_on_4xx(): + """Upstream 4xx is forwarded verbatim — proxy doesn't re-shape error JSON.""" + err_body = {"error": {"message": "context length exceeded", "type": "BadRequestError"}} + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, json=err_body) + + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 400 + assert r.json() == err_body - with patch(client_post_path, new_callable=AsyncMock) as mock_post: - # Mock 401 Unauthorized (NOT in the retry whitelist) - resp_401 = MagicMock(spec=Response) - resp_401.status_code = 401 - resp_401.json.return_value = {"error": "Invalid API Key"} - # The function should return this response directly - mock_post.return_value = resp_401 +@pytest.mark.asyncio +async def test_forward_authorization_header_passes_through(): + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["headers"] = dict(request.headers) + return httpx.Response(200, json=_success_response_json()) - result = await perform_llm_request("http://fake.url", {}, {}, mock_config) + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": "Bearer sk-abc", "X-Trace": "t1"}, + ) - assert result.status_code == 401 - # Call count must be 1, meaning no retries were attempted - assert mock_post.call_count == 1 + # Authorization and custom X-* headers are forwarded verbatim. We don't assert + # on framing headers (connection / content-length / accept-encoding) because + # httpx rebuilds them itself for the outgoing request. + auth_value = captured["headers"].get("Authorization") or captured["headers"].get("authorization") + assert auth_value == "Bearer sk-abc" + fwd_lower = {k.lower() for k in captured["headers"]} + assert "x-trace" in fwd_lower @pytest.mark.asyncio -async def test_perform_llm_request_network_timeout_retry(): - """ - Test that network-level exceptions (like Timeout) also trigger retries. - """ - client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" +async def test_forward_502_on_upstream_connection_failure(monkeypatch): + """ConnectError → 502. Retry disabled here to keep the test fast.""" + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_MAX_ATTEMPTS", 1) - with ( - patch(client_post_path, new_callable=AsyncMock) as mock_post, - patch("rock.utils.retry.asyncio.sleep", return_value=None), - ): - resp_200 = MagicMock(spec=Response) - resp_200.status_code = 200 + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("upstream is down") + + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) - mock_post.side_effect = [httpx.TimeoutException("Network Timeout"), resp_200] + assert r.status_code == 502 - result = await perform_llm_request("http://fake.url", {}, {}, mock_config) - assert result.status_code == 200 - assert mock_post.call_count == 2 +# ---------- Forward path: retry ---------- @pytest.mark.asyncio -async def test_lifespan_initialization_with_config(tmp_path): - """ - Test that the application correctly initializes and overrides defaults - when a valid configuration file path is provided. - """ - conf_file = tmp_path / "proxy.yml" - conf_file.write_text(yaml.dump({"proxy_rules": {"my-model": "http://custom-url"}, "request_timeout": 50})) +async def test_forward_retries_on_retryable_status_then_succeeds(monkeypatch): + """A 429 is retried; the next attempt's 200 is returned to the client.""" + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_DELAY_SECONDS", 0.0) - # Initialize App and load config from file - config = ModelServiceConfig.from_file(str(conf_file)) - app = FastAPI(lifespan=lambda app: lifespan(app, config)) + attempts = [] - async with lifespan(app, config): - app_config = app.state.model_service_config - # Verify that the config reflects file content instead of defaults - assert app_config.proxy_rules["my-model"] == "http://custom-url" - assert app_config.request_timeout == 50 - assert "gpt-3.5-turbo" not in app_config.proxy_rules + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(1) + if len(attempts) < 3: + return httpx.Response(429, json={"error": "rate limited"}) + return httpx.Response(200, json=_success_response_json(content="finally")) + + app = _build_app(ModelServiceConfig()) # default retryable_status_codes = [429, 500] + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["content"] == "finally" + assert len(attempts) == 3 @pytest.mark.asyncio -async def test_lifespan_initialization_no_config(): - """ - Test that the application initializes with default ModelServiceConfig - settings when no configuration file path is provided. - """ - config = ModelServiceConfig() - app = FastAPI(lifespan=lambda app: lifespan(app, config)) +async def test_forward_returns_last_response_when_retries_exhausted(monkeypatch): + """All attempts return 429 → the final 429 body+status is forwarded verbatim.""" + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_MAX_ATTEMPTS", 3) + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_DELAY_SECONDS", 0.0) - async with lifespan(app, config): - app_config = app.state.model_service_config - # Verify that default rules (e.g., 'gpt-3.5-turbo') are loaded - assert "gpt-3.5-turbo" in app_config.proxy_rules - assert app_config.request_timeout == 120 + attempts = [] + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(1) + return httpx.Response(429, json={"error": "still rate limited"}) -@pytest.mark.asyncio -async def test_lifespan_invalid_config_path(): - """ - Test that providing a non-existent configuration file path causes - ModelServiceConfig.from_file to raise a FileNotFoundError. - """ - # Expect FileNotFoundError when loading from non-existent file - with pytest.raises(FileNotFoundError): - ModelServiceConfig.from_file("/tmp/non_existent_file.yml") + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 429 + assert r.json() == {"error": "still rate limited"} + assert len(attempts) == 3 @pytest.mark.asyncio -async def test_proxy_base_url_overrides_proxy_rules(tmp_path): - """ - Test that when proxy_base_url is set, all requests are forwarded to that URL, - bypassing proxy_rules entirely. - """ - config = ModelServiceConfig() - config.proxy_base_url = "https://custom-endpoint.example.com/v1" +async def test_forward_does_not_retry_non_whitelisted_status(monkeypatch): + """400 is not in retryable_status_codes → forwarded immediately, no retry.""" + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_DELAY_SECONDS", 0.0) - test_app = FastAPI() - test_app.state.model_service_config = config - test_app.include_router(proxy_router) + attempts = [] - with patch("rock.sdk.model.server.api.proxy.perform_llm_request", new_callable=AsyncMock) as mock_request: - mock_resp = MagicMock(spec=Response) - mock_resp.status_code = 200 - mock_resp.json.return_value = {"id": "chat-123", "choices": []} - mock_request.return_value = mock_resp + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(1) + return httpx.Response(400, json={"error": "bad request"}) - transport = ASGITransport(app=test_app) + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as ac: - # Even when requesting gpt-3.5-turbo, should forward to proxy_base_url - payload = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}]} - response = await ac.post("/v1/chat/completions", json=payload) + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) - assert response.status_code == 200 - # Verify request was sent to proxy_base_url - call_args = mock_request.call_args[0] - assert call_args[0] == "https://custom-endpoint.example.com/v1/chat/completions" + assert r.status_code == 400 + assert len(attempts) == 1 @pytest.mark.asyncio -async def test_config_loads_host_and_port_from_file(tmp_path): - """ - Test that ModelServiceConfig correctly loads host and port from config file. - """ - conf_file = tmp_path / "proxy.yml" - conf_file.write_text( - yaml.dump({"host": "127.0.0.1", "port": 9000, "proxy_rules": {"my-model": "http://my-backend"}}) +async def test_forward_stream_retries_on_retryable_status_then_succeeds(monkeypatch): + """Streaming: 500 on first attempt, then 200 SSE on second — client sees only the 200 body.""" + monkeypatch.setattr("rock.sdk.model.server.api.proxy._RETRY_DELAY_SECONDS", 0.0) + + attempts = [] + sse_body = ( + b'data: {"id":"x","object":"chat.completion.chunk","choices":[{"index":0,' + b'"delta":{"content":"hello"},"finish_reason":null}]}\n\n' + b"data: [DONE]\n\n" ) - config = ModelServiceConfig.from_file(str(conf_file)) + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(1) + if len(attempts) < 2: + return httpx.Response(500, json={"error": "internal"}) + return httpx.Response(200, content=sse_body, headers={"content-type": "text/event-stream"}) - assert config.host == "127.0.0.1" - assert config.port == 9000 - assert config.proxy_rules["my-model"] == "http://my-backend" + app = _build_app(ModelServiceConfig()) + with _patch_httpx_with_handler(handler): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "stream": True, "messages": [{"role": "user", "content": "hi"}]}, + ) + body = r.text + assert "hello" in body + assert "[DONE]" in body + assert "internal" not in body # the 500 attempt is not leaked to client + assert len(attempts) == 2 + + +# ---------- Forward path: recording ---------- + + +@pytest.mark.asyncio +async def test_forward_invokes_recorder_on_success(tmp_path): + """When a recorder is attached to the backend, success calls write a JSONL line.""" + from rock.sdk.model.server.traj import TrajectoryRecorder + + upstream_payload = _success_response_json(content="recorded reply") + traj_file = tmp_path / "traj.jsonl" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=upstream_payload) -def test_config_default_host_and_port(): - """ - Test default values for host and port. - """ config = ModelServiceConfig() - assert config.host == "0.0.0.0" - assert config.port == 8080 + with _patch_httpx_with_handler(handler): + recorder = TrajectoryRecorder(traj_file=traj_file) + app = _build_app(config, recorder=recorder) + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + + line = traj_file.read_text(encoding="utf-8").strip() + record = json.loads(line) + assert record["status"] == "success" + assert record["model"] == "gpt-3.5-turbo" + assert record["stream"] is False + assert record["request"]["messages"][0]["content"] == "hi" + assert record["response"] == upstream_payload + + +# ---------- Replay path ---------- @pytest.mark.asyncio -async def test_config_loads_retryable_status_codes_from_file(tmp_path): - """ - Test that ModelServiceConfig correctly loads retryable_status_codes from config file. - """ - conf_file = tmp_path / "proxy.yml" - conf_file.write_text(yaml.dump({"retryable_status_codes": [429, 500, 502, 503]})) +async def test_replay_returns_recorded_response_no_upstream_call(tmp_path): + record = { + "model": "gpt-3.5-turbo", + "response": { + "id": "rec-1", + "object": "chat.completion", + "model": "gpt-3.5-turbo", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "recorded reply"}, + "finish_reason": "stop", + } + ], + }, + } + traj = tmp_path / "t.jsonl" + traj.write_text(json.dumps(record) + "\n", encoding="utf-8") - config = ModelServiceConfig.from_file(str(conf_file)) + config = ModelServiceConfig() + config.replay_file = str(traj) + app = _build_app(config, replay_cursor=SequentialCursor.load(traj)) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["content"] == "recorded reply" - assert config.retryable_status_codes == [429, 500, 502, 503] +@pytest.mark.asyncio +async def test_replay_streaming_emits_recorded_response_as_sse(tmp_path): + record = { + "model": "gpt-3.5-turbo", + "response": { + "id": "rec-stream", + "object": "chat.completion", + "model": "gpt-3.5-turbo", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "streamed reply"}, + "finish_reason": "tool_calls", + } + ], + }, + } + traj = tmp_path / "t.jsonl" + traj.write_text(json.dumps(record) + "\n", encoding="utf-8") -def test_config_default_retryable_status_codes(): - """ - Test default values for retryable_status_codes. - """ config = ModelServiceConfig() + config.replay_file = str(traj) + app = _build_app(config, replay_cursor=SequentialCursor.load(traj)) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + r = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "stream": True, "messages": [{"role": "user", "content": "hi"}]}, + ) - assert config.retryable_status_codes == [429, 500] + body = r.text + assert "data: [DONE]" in body + assert '"object": "chat.completion.chunk"' in body + assert '"delta": {"role": "assistant", "content": "streamed reply"}' in body + assert '"finish_reason": "tool_calls"' in body @pytest.mark.asyncio -async def test_perform_llm_request_respects_custom_retryable_codes(): - """ - Test that custom retryable_status_codes are respected (502 retries, 401 does not). - """ +async def test_replay_returns_404_when_cursor_exhausted(tmp_path): + record = { + "model": "gpt-3.5-turbo", + "response": { + "id": "only", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + }, + } + traj = tmp_path / "t.jsonl" + traj.write_text(json.dumps(record) + "\n", encoding="utf-8") + config = ModelServiceConfig() - config.retryable_status_codes = [502, 503, 504] # Custom retryable status codes + config.replay_file = str(traj) + app = _build_app(config, replay_cursor=SequentialCursor.load(traj)) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + ) + second = await ac.post( + "/v1/chat/completions", + json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "again"}]}, + ) - client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" + assert second.status_code == 404 + assert "exhausted" in second.json()["detail"] - with ( - patch(client_post_path, new_callable=AsyncMock) as mock_post, - patch("rock.utils.retry.asyncio.sleep", return_value=None), - ): - # 502 should retry (in custom list) - resp_502 = MagicMock(spec=Response) - resp_502.status_code = 502 - error_502 = HTTPStatusError("Bad Gateway", request=MagicMock(spec=Request), response=resp_502) - resp_200 = MagicMock(spec=Response) - resp_200.status_code = 200 - resp_200.json.return_value = {"ok": True} +# ---------- Lifespan + Config ---------- - # Sequence: 502 fail, then 200 success - mock_post.side_effect = [error_502, resp_200] - result = await perform_llm_request("http://fake.url", {}, {}, config) +@pytest.mark.asyncio +async def test_lifespan_initialization_with_config(tmp_path): + conf_file = tmp_path / "proxy.yml" + conf_file.write_text(yaml.dump({"proxy_rules": {"my-model": "http://custom-url"}, "request_timeout": 50})) - assert result.status_code == 200 - assert mock_post.call_count == 2 + config = ModelServiceConfig.from_file(str(conf_file)) + app = FastAPI(lifespan=lambda app: lifespan(app, config)) + + async with lifespan(app, config): + assert app.state.model_service_config.proxy_rules["my-model"] == "http://custom-url" + assert app.state.model_service_config.request_timeout == 50 @pytest.mark.asyncio -async def test_perform_llm_request_non_retryable_code_not_retried(): - """ - Test that 401 (not in custom retryable_status_codes) does not trigger retry. - """ +async def test_lifespan_invalid_config_path(): + with pytest.raises(FileNotFoundError): + ModelServiceConfig.from_file("/tmp/non_existent_file.yml") + + +def test_config_default_host_and_port(): config = ModelServiceConfig() - config.retryable_status_codes = [502, 503, 504] # Custom retryable status codes, excluding 401 + assert config.host == "0.0.0.0" + assert config.port == 8080 - client_post_path = "rock.sdk.model.server.api.proxy.http_client.post" - with patch(client_post_path, new_callable=AsyncMock) as mock_post: - # 401 should not retry (not in custom list) - resp_401 = MagicMock(spec=Response) - resp_401.status_code = 401 - resp_401.json.return_value = {"error": "Invalid API Key"} +def test_config_default_recording_and_replay(): + config = ModelServiceConfig() + assert config.recording_file is None + assert config.replay_file is None + - mock_post.return_value = resp_401 +@pytest.mark.asyncio +async def test_config_loads_recording_file_from_yaml(tmp_path): + conf_file = tmp_path / "proxy.yml" + conf_file.write_text(yaml.dump({"recording_file": "/tmp/my-traj.jsonl"})) + config = ModelServiceConfig.from_file(str(conf_file)) + assert config.recording_file == "/tmp/my-traj.jsonl" + assert config.replay_file is None - result = await perform_llm_request("http://fake.url", {}, {}, config) - assert result.status_code == 401 - assert mock_post.call_count == 1 # No retry +@pytest.mark.asyncio +async def test_config_loads_replay_file_from_yaml(tmp_path): + conf_file = tmp_path / "proxy.yml" + conf_file.write_text(yaml.dump({"replay_file": "/tmp/in.jsonl"})) + config = ModelServiceConfig.from_file(str(conf_file)) + assert config.replay_file == "/tmp/in.jsonl" + assert config.recording_file is None -def test_cli_args_override_config_file(tmp_path): - """ - Test that CLI arguments override config file settings. - This tests the logic in create_config_from_args(). - """ - import argparse +def test_config_recording_and_replay_are_mutually_exclusive(): + """Setting both at construction time fails Pydantic validation.""" + with pytest.raises(ValueError, match="mutually exclusive"): + ModelServiceConfig(recording_file="/tmp/a.jsonl", replay_file="/tmp/b.jsonl") - # Create args with config file and CLI parameters + +def test_config_recording_replay_mutex_fires_on_assignment(): + """validate_assignment=True so CLI-style field-by-field overrides also trip the mutex.""" + config = ModelServiceConfig(recording_file="/tmp/a.jsonl") + with pytest.raises(ValueError, match="mutually exclusive"): + config.replay_file = "/tmp/b.jsonl" + + +def test_cli_args_override_config_file(tmp_path): conf_file = tmp_path / "proxy.yml" conf_file.write_text( yaml.dump( @@ -378,144 +595,67 @@ def test_cli_args_override_config_file(tmp_path): "host": "192.168.1.1", "port": 8080, "proxy_base_url": "https://config-url.example.com/v1", - "retryable_status_codes": [429, 500], "request_timeout": 60, } ) ) - args = argparse.Namespace( config_file=str(conf_file), - host="0.0.0.0", # CLI overrides config file - port=9000, # CLI overrides config file - proxy_base_url="https://cli-url.example.com/v1", # CLI overrides config file - retryable_status_codes="502,503", # CLI overrides config file - request_timeout=30, # CLI overrides config file + host="0.0.0.0", + port=9000, + proxy_base_url="https://cli-url.example.com/v1", + retryable_status_codes=None, + request_timeout=30, + recording_file=None, + replay_file=None, ) - config = create_config_from_args(args) - - # Verify CLI arguments override config file assert config.host == "0.0.0.0" assert config.port == 9000 assert config.proxy_base_url == "https://cli-url.example.com/v1" - assert config.retryable_status_codes == [502, 503] assert config.request_timeout == 30 -@pytest.mark.asyncio -async def test_config_file_overrides_defaults(tmp_path): - """ - Test that config file values override default values. - """ - conf_file = tmp_path / "proxy.yml" - conf_file.write_text( - yaml.dump( - { - "host": "10.0.0.1", - "port": 8888, - "request_timeout": 300, - "proxy_rules": {"test-model": "http://test-backend"}, - } - ) +def test_cli_replay_file_enables_replay(): + args = argparse.Namespace( + config_file=None, + host=None, + port=None, + proxy_base_url=None, + retryable_status_codes=None, + request_timeout=None, + recording_file=None, + replay_file="/tmp/in.jsonl", ) + config = create_config_from_args(args) + assert config.replay_file == "/tmp/in.jsonl" - config = ModelServiceConfig.from_file(str(conf_file)) - # Verify config file overrides defaults - assert config.host == "10.0.0.1" - assert config.port == 8888 - assert config.request_timeout == 300 - assert config.proxy_rules["test-model"] == "http://test-backend" - # Verify other fields remain as defaults - assert config.proxy_base_url is None +# ---------- Metrics singleton + legacy record_traj (still used by local mode) ---------- def test_metrics_monitor_is_singleton(): - """ - Test that _get_or_create_metrics_monitor returns the same instance - on repeated calls (module-level singleton, created only once). - """ import rock.sdk.model.server.utils as utils_module with patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls: - mock_monitor = MagicMock() - mock_cls.create.return_value = mock_monitor - - # Reset singleton so the test is isolated + mock_cls.create.return_value = MagicMock() utils_module._metrics_monitor = None - first = _get_or_create_metrics_monitor() second = _get_or_create_metrics_monitor() - assert first is second - assert mock_cls.create.call_count == 1 - - # Cleanup - utils_module._metrics_monitor = None - - -def test_metrics_monitor_uses_env_endpoint(): - """ - Test that ROCK_METRICS_ENDPOINT env var is passed to MetricsMonitor.create(). - """ - import rock.sdk.model.server.utils as utils_module - - custom_endpoint = "http://my-otel-collector:4318/v1/metrics" - - with ( - patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, - patch.dict("os.environ", {"ROCK_METRICS_ENDPOINT": custom_endpoint}), - ): - mock_monitor = MagicMock() - mock_cls.create.return_value = mock_monitor - - utils_module._metrics_monitor = None - _get_or_create_metrics_monitor() - - mock_cls.create.assert_called_once_with(metrics_endpoint=custom_endpoint) - - utils_module._metrics_monitor = None - - -def test_metrics_monitor_registers_gauge_and_counter(): - """ - Test that _get_or_create_metrics_monitor registers both - the RT gauge and request count counter on first creation. - """ - import rock.sdk.model.server.utils as utils_module - - with patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls: - mock_monitor = MagicMock() - mock_cls.create.return_value = mock_monitor - - utils_module._metrics_monitor = None - _get_or_create_metrics_monitor() - - mock_monitor._register_gauge.assert_called_once_with( - MODEL_SERVICE_REQUEST_RT, "total execution time for request", "ms" - ) - mock_monitor._register_counter.assert_called_once_with( - MODEL_SERVICE_REQUEST_COUNT, "total request count", "count" - ) - utils_module._metrics_monitor = None @pytest.mark.asyncio -async def test_record_traj_reports_rt_and_count(): - """ - Test that record_traj decorator calls record_gauge_by_name (RT) - and record_counter_by_name (count) with correct metric names and attributes. - """ +async def test_record_traj_decorator_reports_rt_and_count(): + """Legacy record_traj decorator (still used by local mode) reports RT/count.""" import rock.sdk.model.server.utils as utils_module - mock_monitor = MagicMock() - with ( patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, - patch.dict("os.environ", {"ROCK_SANDBOX_ID": "sandbox-test-001"}), + patch.dict("os.environ", {"ROCK_SANDBOX_ID": "sandbox-test"}), ): + mock_monitor = MagicMock() mock_cls.create.return_value = mock_monitor utils_module._metrics_monitor = None @@ -525,45 +665,11 @@ async def fake_handler(body: dict): await fake_handler({"model": "gpt-4", "messages": []}) - mock_monitor.record_gauge_by_name.assert_called_once() gauge_call = mock_monitor.record_gauge_by_name.call_args assert gauge_call[0][0] == MODEL_SERVICE_REQUEST_RT - assert gauge_call[1]["attributes"]["type"] == "chat_completions" - assert gauge_call[1]["attributes"]["sandbox_id"] == "sandbox-test-001" + assert gauge_call[1]["attributes"]["sandbox_id"] == "sandbox-test" - mock_monitor.record_counter_by_name.assert_called_once() counter_call = mock_monitor.record_counter_by_name.call_args assert counter_call[0][0] == MODEL_SERVICE_REQUEST_COUNT - assert counter_call[0][1] == 1 - assert counter_call[1]["attributes"]["sandbox_id"] == "sandbox-test-001" - - utils_module._metrics_monitor = None - - -@pytest.mark.asyncio -async def test_record_traj_sandbox_id_defaults_to_unknown(): - """ - Test that sandbox_id defaults to 'unknown' when ROCK_SANDBOX_ID is not set. - """ - import rock.sdk.model.server.utils as utils_module - - mock_monitor = MagicMock() - - with patch("rock.sdk.model.server.utils.MetricsMonitor") as mock_cls, patch.dict("os.environ", {}, clear=False): - # Ensure ROCK_SANDBOX_ID is not set - os_env = __import__("os").environ - os_env.pop("ROCK_SANDBOX_ID", None) - - mock_cls.create.return_value = mock_monitor - utils_module._metrics_monitor = None - - @record_traj - async def fake_handler(body: dict): - return {"id": "resp-2", "choices": []} - - await fake_handler({"model": "gpt-4", "messages": []}) - - gauge_call = mock_monitor.record_gauge_by_name.call_args - assert gauge_call[1]["attributes"]["sandbox_id"] == "unknown" utils_module._metrics_monitor = None diff --git a/tests/unit/sdk/model/test_proxy_record_replay_e2e.py b/tests/unit/sdk/model/test_proxy_record_replay_e2e.py new file mode 100644 index 0000000000..0b70ed0cf8 --- /dev/null +++ b/tests/unit/sdk/model/test_proxy_record_replay_e2e.py @@ -0,0 +1,332 @@ +"""End-to-end: real in-process TCP mock upstream + real proxy router + recorder. + +The mock upstream is a tiny FastAPI app served by uvicorn in a background thread +(real TCP). The proxy stays in-process and is hit via FastAPI's ``TestClient``; +its outbound ``httpx.AsyncClient`` makes a real TCP call to the mock — production +code path, no transport injection, no patching. +""" + +from __future__ import annotations + +import asyncio +import json +import threading +import time +from collections.abc import Iterator +from pathlib import Path + +import pytest +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.testclient import TestClient + +from rock.sdk.model.server.api.proxy import ForwardBackend, ReplayBackend, proxy_router +from rock.sdk.model.server.config import ModelServiceConfig +from rock.sdk.model.server.sse import parse_sse_data_chunks +from rock.sdk.model.server.traj import SequentialCursor, TrajectoryRecorder +from rock.utils.system import find_free_port + +# --------------------------------------------------------------------------- +# Mock upstream — a tiny FastAPI app behind a real TCP uvicorn in a thread. +# Owns both the canned reply AND the assertion helper, so the response shape +# and the expectations stay in sync if either is edited. +# --------------------------------------------------------------------------- + + +class MockUpstream: + """Mock OpenAI-compatible upstream. + + Single canonical reply (returned for both stream and non-stream requests) + contains three fields the proxy must preserve end-to-end: + - ``content`` (plain text) + - ``reasoning_content`` (vendor-specific thinking) + - ``tool_calls`` (a function call) + The streaming variant splits each field into multiple deltas so the + recorder also exercises the openai SDK's stream-state aggregator. + + Use as ``with MockUpstream() as mock: ...``; ``mock.base_url`` points at + the running server. ``mock.assert_message(msg)`` checks any received + assistant message matches the canonical reply. + """ + + # Canonical reply values — change here, both the handler and the assertion + # helper pick them up automatically. Two parallel tool_calls cover the + # multi-tool-call case (modern LLMs commonly emit several at once). + EXPECTED_CONTENT = "Checking weather and time for you." + EXPECTED_REASONING = "User wants weather + time; calling both tools in parallel." + EXPECTED_TOOL_CALLS = [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Tokyo","unit":"celsius"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city":"Tokyo"}'}, + }, + ] + + def __init__(self) -> None: + port = asyncio.run(find_free_port()) + config = uvicorn.Config(self._build_app(), host="127.0.0.1", port=port, log_level="warning", access_log=False) + self._server = uvicorn.Server(config) + self._thread = threading.Thread(target=self._server.run, daemon=True) + self.base_url = f"http://127.0.0.1:{port}/v1" + + # ---- lifecycle ---- + + def __enter__(self) -> MockUpstream: + self._thread.start() + deadline = time.time() + 5.0 + while not self._server.started: + if time.time() > deadline: + raise RuntimeError("mock upstream did not start within 5s") + time.sleep(0.02) + return self + + def __exit__(self, *_exc) -> None: + self._server.should_exit = True + self._thread.join(timeout=5) + + # ---- assertion helper ---- + + def assert_message(self, msg: dict) -> None: + """Assert ``msg`` is the canonical full message (content + reasoning + 2 parallel tool_calls).""" + assert msg["content"] == self.EXPECTED_CONTENT + assert msg["reasoning_content"] == self.EXPECTED_REASONING + tcs = msg["tool_calls"] + assert len(tcs) == len(self.EXPECTED_TOOL_CALLS) + for actual, expected in zip(tcs, self.EXPECTED_TOOL_CALLS, strict=True): + assert actual["id"] == expected["id"] + assert actual["type"] == expected["type"] + assert actual["function"]["name"] == expected["function"]["name"] + assert json.loads(actual["function"]["arguments"]) == json.loads(expected["function"]["arguments"]) + + # ---- internal: FastAPI app + handlers ---- + + def _build_app(self) -> FastAPI: + app = FastAPI() + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request): + body = await request.json() + model = body.get("model", "mock") + if body.get("stream"): + return StreamingResponse(self._stream_gen(model), media_type="text/event-stream") + return JSONResponse(status_code=200, content=self._completion_json(model)) + + return app + + def _completion_json(self, model: str) -> dict: + return { + "id": "chatcmpl-mock-1", + "object": "chat.completion", + "created": 0, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": self.EXPECTED_CONTENT, + "reasoning_content": self.EXPECTED_REASONING, + "tool_calls": self.EXPECTED_TOOL_CALLS, + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 12, "completion_tokens": 24, "total_tokens": 36}, + } + + async def _stream_gen(self, model: str): + base = {"id": "chatcmpl-mock-1", "object": "chat.completion.chunk", "created": 0, "model": model} + + def emit(delta: dict, finish_reason=None) -> bytes: + payload = {**base, "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]} + return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode() + + # 1-2. Reasoning split in two deltas + yield emit({"role": "assistant", "reasoning_content": "User wants weather + time; "}) + await asyncio.sleep(0.005) + yield emit({"reasoning_content": "calling both tools in parallel."}) + await asyncio.sleep(0.005) + # 3-4. Content split in two deltas + yield emit({"content": "Checking weather"}) + await asyncio.sleep(0.005) + yield emit({"content": " and time for you."}) + await asyncio.sleep(0.005) + + # 5-7. tool_call[0] (get_weather): announce, then arguments in two pieces + yield emit( + { + "tool_calls": [ + { + "index": 0, + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + } + ] + } + ) + await asyncio.sleep(0.005) + yield emit({"tool_calls": [{"index": 0, "function": {"arguments": '{"city":"Tokyo",'}}]}) + await asyncio.sleep(0.005) + yield emit({"tool_calls": [{"index": 0, "function": {"arguments": '"unit":"celsius"}'}}]}) + await asyncio.sleep(0.005) + + # 8-9. tool_call[1] (get_time): announce + arguments in one piece + yield emit( + { + "tool_calls": [ + { + "index": 1, + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": ""}, + } + ] + } + ) + await asyncio.sleep(0.005) + yield emit({"tool_calls": [{"index": 1, "function": {"arguments": '{"city":"Tokyo"}'}}]}) + await asyncio.sleep(0.005) + + # 10. Finish + yield emit({}, finish_reason="tool_calls") + yield b"data: [DONE]\n\n" + + +@pytest.fixture +def mock_upstream() -> Iterator[MockUpstream]: + with MockUpstream() as m: + yield m + + +# --------------------------------------------------------------------------- +# Proxy app builder + request helper (module-level, generic) +# --------------------------------------------------------------------------- + + +def _build_proxy_app(*, mock_url: str | None = None, traj_file: Path | None = None, replay_cursor=None) -> FastAPI: + config = ModelServiceConfig() + # ReplayBackend never calls upstream, so mock_url is only relevant for forward mode. + if mock_url is not None: + config.proxy_base_url = mock_url + + app = FastAPI() + app.state.model_service_config = config + if replay_cursor is not None: + app.state.backend = ReplayBackend(replay_cursor) + else: + recorder = TrajectoryRecorder(traj_file=traj_file) if traj_file is not None else None + app.state.backend = ForwardBackend(config, recorder=recorder) + app.include_router(proxy_router) + return app + + +def _call_chat_completions(client: TestClient, *, stream: bool) -> dict: + """One chat.completions call. Returns the assistant message dict. + + - non-stream: just unwraps ``choices[0].message``. + - stream: replay always emits exactly one chunk + ``[DONE]`` (see + ``completion_to_chunk_dict``), so the chunk's ``delta`` IS the full + message — no aggregation needed. + """ + payload = {"model": "mock-model", "messages": [{"role": "user", "content": "hi"}]} + if not stream: + r = client.post("/v1/chat/completions", json=payload) + assert r.status_code == 200 + return r.json()["choices"][0]["message"] + + with client.stream("POST", "/v1/chat/completions", json={**payload, "stream": True}) as r: + assert r.status_code == 200 + body_bytes = b"".join(r.iter_bytes()) + chunks, _ = parse_sse_data_chunks(body_bytes) + return chunks[0]["choices"][0]["delta"] + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestProxyRecordReplay: + """End-to-end: real TCP mock upstream <-> real proxy router + recorder/replayer.""" + + def test_forward_non_stream(self, mock_upstream: MockUpstream, tmp_path): + """Vendor field reaches the client; recorder writes a JSONL line with the full response.""" + traj_file = tmp_path / "traj.jsonl" + proxy_app = _build_proxy_app(mock_url=mock_upstream.base_url, traj_file=traj_file) + + with TestClient(proxy_app) as client: + r = client.post( + "/v1/chat/completions", + json={"model": "mock-model", "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert r.status_code == 200 + body = r.json() + assert body["choices"][0]["finish_reason"] == "tool_calls" + mock_upstream.assert_message(body["choices"][0]["message"]) + + rec = json.loads(traj_file.read_text(encoding="utf-8").strip()) + assert rec["status"] == "success" + assert rec["stream"] is False + assert rec["response"]["choices"][0]["finish_reason"] == "tool_calls" + mock_upstream.assert_message(rec["response"]["choices"][0]["message"]) + + def test_forward_stream(self, mock_upstream: MockUpstream, tmp_path): + """Each upstream SSE chunk reaches the client; recorder gets the aggregated final completion + with reasoning_content concatenated and tool_calls.arguments assembled from deltas.""" + traj_file = tmp_path / "traj.jsonl" + proxy_app = _build_proxy_app(mock_url=mock_upstream.base_url, traj_file=traj_file) + + with TestClient(proxy_app) as client: + with client.stream( + "POST", + "/v1/chat/completions", + json={"model": "mock-model", "stream": True, "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": "Bearer test-key"}, + ) as r: + body = b"".join(r.iter_bytes()).decode("utf-8") + + # Raw chunks make it to the client untouched + assert '"reasoning_content": "User wants weather + time; "' in body + assert '"reasoning_content": "calling both tools in parallel."' in body + assert '"content": "Checking weather"' in body + assert '"content": " and time for you."' in body + assert '"name": "get_weather"' in body + assert '"name": "get_time"' in body + assert '"finish_reason": "tool_calls"' in body + assert body.rstrip().endswith("data: [DONE]") + + # Recorder's aggregated message matches the canonical reply + rec = json.loads(traj_file.read_text(encoding="utf-8").strip()) + assert rec["status"] == "success" + assert rec["stream"] is True + assert rec["response"]["choices"][0]["finish_reason"] == "tool_calls" + mock_upstream.assert_message(rec["response"]["choices"][0]["message"]) + + @pytest.mark.parametrize("replay_stream", [False, True], ids=["replay_nonstream", "replay_stream"]) + @pytest.mark.parametrize("record_stream", [False, True], ids=["record_nonstream", "record_stream"]) + def test_replay(self, mock_upstream: MockUpstream, tmp_path, record_stream: bool, replay_stream: bool): + """Recorded mode and replayed mode are orthogonal — all 4 combinations of + (stream/non-stream) on each side must yield the same full message.""" + traj_file = tmp_path / "traj.jsonl" + + # ---- record phase ---- + proxy_record = _build_proxy_app(mock_url=mock_upstream.base_url, traj_file=traj_file) + with TestClient(proxy_record) as client: + _call_chat_completions(client, stream=record_stream) + + # ---- replay phase: no upstream URL needed — ReplayBackend never calls upstream ---- + cursor = SequentialCursor.load(traj_file) + proxy_replay = _build_proxy_app(replay_cursor=cursor) + with TestClient(proxy_replay) as client: + msg = _call_chat_completions(client, stream=replay_stream) + + mock_upstream.assert_message(msg) diff --git a/tests/unit/sdk/model/test_service_subprocess.py b/tests/unit/sdk/model/test_service_subprocess.py new file mode 100644 index 0000000000..61176173bd --- /dev/null +++ b/tests/unit/sdk/model/test_service_subprocess.py @@ -0,0 +1,38 @@ +"""Tests for ModelService.start_sandbox_service subprocess command construction. + +Covers the CLI flag wiring without actually spawning a subprocess: mock Popen +and inspect the argv it would have been called with. +""" + +from unittest.mock import patch + +from rock.sdk.model.service import ModelService + + +def _captured_argv(**start_kwargs) -> list[str]: + with patch("rock.sdk.model.service.subprocess.Popen") as mock_popen: + ModelService().start_sandbox_service(**start_kwargs) + return mock_popen.call_args[0][0] + + +def test_start_sandbox_service_omits_recording_and_replay_flags_by_default(): + argv = _captured_argv(model_service_type="proxy", proxy_base_url="https://api.openai.com/v1", port=8080) + assert argv[1:5] == ["-m", "main", "--type", "proxy"] + assert "--proxy-base-url" in argv and "https://api.openai.com/v1" in argv + assert "--port" in argv and "8080" in argv + assert "--recording-file" not in argv + assert "--replay-file" not in argv + + +def test_start_sandbox_service_passes_recording_file(): + argv = _captured_argv(model_service_type="proxy", recording_file="/tmp/my-traj.jsonl") + idx = argv.index("--recording-file") + assert argv[idx + 1] == "/tmp/my-traj.jsonl" + assert "--replay-file" not in argv + + +def test_start_sandbox_service_passes_replay_file(): + argv = _captured_argv(model_service_type="proxy", replay_file="/tmp/in.jsonl") + idx = argv.index("--replay-file") + assert argv[idx + 1] == "/tmp/in.jsonl" + assert "--recording-file" not in argv diff --git a/tests/unit/sdk/model/test_sse.py b/tests/unit/sdk/model/test_sse.py new file mode 100644 index 0000000000..251016a0a8 --- /dev/null +++ b/tests/unit/sdk/model/test_sse.py @@ -0,0 +1,223 @@ +"""Tests for the pure SSE codec utilities (no openai/litellm dependencies).""" + +import json + +from rock.sdk.model.server.sse import ( + SSE_DONE, + completion_to_chunk_dict, + encode_sse_event, + parse_sse_data_chunks, +) + +# ---------- parse_sse_data_chunks ---------- + + +def test_parse_returns_complete_events_and_leftover_buffer(): + raw = b'data: {"a": 1}\n\ndata: {"a": 2}\n\ndata: {"a": 3}' # 3rd event is incomplete + chunks, leftover = parse_sse_data_chunks(raw) + + assert chunks == [{"a": 1}, {"a": 2}] + assert leftover == b'data: {"a": 3}' + + +def test_parse_skips_done_marker(): + raw = b'data: {"x": 1}\n\ndata: [DONE]\n\n' + chunks, leftover = parse_sse_data_chunks(raw) + + assert chunks == [{"x": 1}] + assert leftover == b"" + + +def test_parse_skips_non_data_lines(): + raw = b'event: progress\ndata: {"y": 2}\nid: abc\n\n' + chunks, leftover = parse_sse_data_chunks(raw) + + assert chunks == [{"y": 2}] + assert leftover == b"" + + +def test_parse_silently_skips_malformed_json(): + raw = b'data: not-json-at-all\n\ndata: {"ok": true}\n\n' + chunks, leftover = parse_sse_data_chunks(raw) + + assert chunks == [{"ok": True}] + assert leftover == b"" + + +def test_parse_handles_empty_buffer(): + chunks, leftover = parse_sse_data_chunks(b"") + assert chunks == [] + assert leftover == b"" + + +def test_parse_incremental_streaming_pattern(): + """Simulates feeding bytes in arbitrary chunks; final concatenation == all events.""" + full_stream = b'data: {"i": 0}\n\ndata: {"i": 1}\n\ndata: {"i": 2}\n\ndata: [DONE]\n\n' + fragments = [full_stream[i : i + 5] for i in range(0, len(full_stream), 5)] + + buffer = b"" + collected: list[dict] = [] + for frag in fragments: + new_chunks, buffer = parse_sse_data_chunks(buffer + frag) + collected.extend(new_chunks) + + assert collected == [{"i": 0}, {"i": 1}, {"i": 2}] + assert buffer == b"" + + +def test_parse_handles_unicode_payload(): + raw = b'data: {"content": "\xe4\xbd\xa0\xe5\xa5\xbd"}\n\n' # "你好" UTF-8 + chunks, _ = parse_sse_data_chunks(raw) + assert chunks == [{"content": "你好"}] + + +# ---------- completion_to_chunk_dict ---------- + + +def test_completion_to_chunk_renames_message_to_delta(): + response = { + "id": "rec-1", + "object": "chat.completion", + "created": 100, + "model": "gpt-4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + chunk = completion_to_chunk_dict(response, model="gpt-4") + + assert chunk["object"] == "chat.completion.chunk" + assert chunk["id"] == "rec-1" + assert chunk["created"] == 100 + assert chunk["model"] == "gpt-4" + assert chunk["choices"][0]["delta"] == {"role": "assistant", "content": "hi"} + assert chunk["choices"][0]["finish_reason"] == "stop" + assert chunk["choices"][0]["index"] == 0 + assert "message" not in chunk["choices"][0] + + +def test_completion_to_chunk_preserves_provider_specific_message_fields(): + """reasoning_content kept verbatim; tool_calls get a positional index injected + (required by the OpenAI streaming spec — see test below).""" + response = { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "answer", + "reasoning_content": "step-by-step thinking", + "tool_calls": [{"id": "t1", "type": "function"}], + }, + "finish_reason": "tool_calls", + } + ], + } + chunk = completion_to_chunk_dict(response, model="glm-5") + + assert chunk["choices"][0]["delta"]["reasoning_content"] == "step-by-step thinking" + assert chunk["choices"][0]["delta"]["tool_calls"] == [{"index": 0, "id": "t1", "type": "function"}] + assert chunk["choices"][0]["finish_reason"] == "tool_calls" + + +def test_completion_to_chunk_injects_tool_call_index_for_openai_sdk_compat(): + """A recorded non-stream message has tool_calls without 'index'; the OpenAI + streaming spec requires it on chunk deltas, and the openai SDK's + ChatCompletionChunk.model_validate() rejects the chunk otherwise. We inject + a positional index so replay-stream output is parseable by strict clients.""" + response = { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + {"id": "a", "type": "function", "function": {"name": "f1", "arguments": "{}"}}, + {"id": "b", "type": "function", "function": {"name": "f2", "arguments": "{}"}}, + ], + }, + "finish_reason": "tool_calls", + } + ], + } + chunk = completion_to_chunk_dict(response, model="m") + tcs = chunk["choices"][0]["delta"]["tool_calls"] + assert [tc["index"] for tc in tcs] == [0, 1] + + # End-to-end: openai SDK accepts the chunk + from openai.types.chat import ChatCompletionChunk + + ChatCompletionChunk.model_validate(chunk) # must not raise + + +def test_completion_to_chunk_preserves_explicit_tool_call_index(): + """If the recorded tool_calls already have 'index', we don't overwrite it.""" + response = { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + {"index": 5, "id": "a", "type": "function", "function": {"name": "f", "arguments": "{}"}}, + ], + }, + "finish_reason": "tool_calls", + } + ], + } + chunk = completion_to_chunk_dict(response, model="m") + assert chunk["choices"][0]["delta"]["tool_calls"][0]["index"] == 5 + + +def test_completion_to_chunk_synthesizes_id_and_created_when_missing(): + chunk = completion_to_chunk_dict( + {"choices": [{"index": 0, "message": {"role": "assistant"}, "finish_reason": "stop"}]}, + model="any", + ) + assert chunk["id"].startswith("chatcmpl-") + assert isinstance(chunk["created"], int) and chunk["created"] > 0 + assert chunk["model"] == "any" + + +def test_completion_to_chunk_handles_empty_choices(): + chunk = completion_to_chunk_dict({"choices": []}, model="m") + assert chunk["choices"] == [] + + +# ---------- encode_sse_event ---------- + + +def test_encode_sse_event_appends_double_newline_terminator(): + out = encode_sse_event({"k": "v"}) + assert out.endswith(b"\n\n") + assert out.startswith(b"data: ") + body = out[len(b"data: ") : -len(b"\n\n")] + assert json.loads(body) == {"k": "v"} + + +def test_encode_sse_event_preserves_unicode_without_escapes(): + out = encode_sse_event({"content": "你好"}) + # ensure_ascii=False is critical so Chinese stays readable in the wire format + assert "你好".encode() in out + + +def test_sse_done_constant(): + assert SSE_DONE == b"data: [DONE]\n\n" + + +# ---------- round-trip ---------- + + +def test_roundtrip_encode_then_parse(): + """encode → parse must round-trip a payload dict.""" + payloads = [{"i": 0, "text": "alpha"}, {"i": 1, "text": "beta 中文"}] + wire = b"".join(encode_sse_event(p) for p in payloads) + SSE_DONE + chunks, leftover = parse_sse_data_chunks(wire) + + assert chunks == payloads + assert leftover == b"" diff --git a/tests/unit/sdk/model/test_traj_recorder.py b/tests/unit/sdk/model/test_traj_recorder.py new file mode 100644 index 0000000000..3f06481639 --- /dev/null +++ b/tests/unit/sdk/model/test_traj_recorder.py @@ -0,0 +1,141 @@ +"""Tests for TrajectoryRecorder (explicit-call API, no longer a litellm CustomLogger).""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from rock.sdk.model.server.traj import TrajectoryRecorder + + +@pytest.fixture +def mock_monitor(): + monitor = MagicMock() + with patch( + "rock.sdk.model.server.traj._get_or_create_metrics_monitor", + return_value=monitor, + ): + yield monitor + + +def _make_recorder(traj_file) -> TrajectoryRecorder: + return TrajectoryRecorder(traj_file=traj_file) + + +@pytest.mark.asyncio +async def test_recorder_appends_each_call_as_jsonl_line(tmp_path, mock_monitor): + traj_file = tmp_path / "traj.jsonl" + recorder = _make_recorder(traj_file) + + await recorder.record( + request={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response={"id": "a", "choices": []}, + status="success", + start_time=100.0, + end_time=100.5, + ) + await recorder.record( + request={"model": "gpt-4", "messages": [{"role": "user", "content": "again"}]}, + response={"id": "b", "choices": []}, + status="success", + start_time=101.0, + end_time=101.2, + ) + + lines = traj_file.read_text(encoding="utf-8").strip().split("\n") + assert len(lines) == 2 + assert json.loads(lines[0])["response"]["id"] == "a" + assert json.loads(lines[1])["response"]["id"] == "b" + + +@pytest.mark.asyncio +async def test_recorder_writes_request_and_response_verbatim(tmp_path, mock_monitor): + """Provider-specific fields (reasoning_content, citations, ...) survive untouched.""" + traj_file = tmp_path / "traj.jsonl" + recorder = _make_recorder(traj_file) + + request = {"model": "glm-5", "stream": True, "messages": [{"role": "user", "content": "你是谁"}]} + response = { + "id": "x", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "我是 GLM", "reasoning_content": "用户问..."}, + "finish_reason": "stop", + } + ], + } + await recorder.record(request=request, response=response, status="success", start_time=0.0, end_time=1.0) + + record = json.loads(traj_file.read_text(encoding="utf-8").strip()) + assert record["model"] == "glm-5" + assert record["stream"] is True + assert record["request"] == request + assert record["response"] == response + assert record["response_time"] == 1.0 + + +@pytest.mark.asyncio +async def test_recorder_emits_metrics_with_status_and_sandbox_id(tmp_path, mock_monitor): + traj_file = tmp_path / "traj.jsonl" + recorder = _make_recorder(traj_file) + + with patch.dict("os.environ", {"ROCK_SANDBOX_ID": "sandbox-xyz"}): + await recorder.record( + request={"model": "gpt-4"}, + response={"id": "x", "choices": []}, + status="success", + start_time=0.0, + end_time=0.5, + ) + + gauge_call = mock_monitor.record_gauge_by_name.call_args + assert gauge_call[0][0] == "model_service.request.rt" + assert gauge_call[0][1] == 500.0 # 0.5s -> 500 ms + assert gauge_call[1]["attributes"]["status"] == "success" + assert gauge_call[1]["attributes"]["sandbox_id"] == "sandbox-xyz" + assert gauge_call[1]["attributes"]["type"] == "chat_completions" + + mock_monitor.record_counter_by_name.assert_called_once_with( + "model_service.request.count", 1, attributes=gauge_call[1]["attributes"] + ) + + +@pytest.mark.asyncio +async def test_recorder_records_failure_with_error_text(tmp_path, mock_monitor): + traj_file = tmp_path / "traj.jsonl" + recorder = _make_recorder(traj_file) + + await recorder.record( + request={"model": "gpt-4"}, + response=None, + status="failure", + start_time=0.0, + end_time=1.0, + error="upstream_status=429", + ) + + record = json.loads(traj_file.read_text(encoding="utf-8").strip()) + assert record["status"] == "failure" + assert record["error"] == "upstream_status=429" + assert record["response"] is None + + gauge_call = mock_monitor.record_gauge_by_name.call_args + assert gauge_call[1]["attributes"]["status"] == "failure" + + +@pytest.mark.asyncio +async def test_recorder_creates_parent_directory(tmp_path, mock_monitor): + traj_file = tmp_path / "deep" / "nested" / "traj.jsonl" + recorder = _make_recorder(traj_file) + + await recorder.record( + request={"model": "gpt-4"}, + response={"id": "x", "choices": []}, + status="success", + start_time=0.0, + end_time=0.5, + ) + + assert traj_file.exists() + assert traj_file.parent.is_dir() diff --git a/tests/unit/sdk/model/test_traj_replayer.py b/tests/unit/sdk/model/test_traj_replayer.py new file mode 100644 index 0000000000..ffcc5c4011 --- /dev/null +++ b/tests/unit/sdk/model/test_traj_replayer.py @@ -0,0 +1,122 @@ +"""Tests for SequentialCursor (the replay cursor used by proxy.py). + +The proxy serves replay responses directly — there is no CustomLLM-based +``TrajectoryReplayer`` anymore. End-to-end replay coverage (cursor + SSE chunk +emit + cursor-exhausted → 404) lives in ``test_proxy.py``. +""" + +import json + +import pytest + +from rock.sdk.model.server.traj import SequentialCursor, TrajectoryExhausted + + +def _record(*, msg: str, model: str = "gpt-3.5-turbo", call_id: str = "x") -> dict: + return { + "id": call_id, + "model": model, + "messages": [{"role": "user", "content": msg}], + "response": { + "id": call_id, + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"reply: {msg}"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + } + + +def _write_jsonl(path, records): + with path.open("w", encoding="utf-8") as f: + for r in records: + f.write(json.dumps(r) + "\n") + + +def test_cursor_load_from_single_file(tmp_path): + p = tmp_path / "traj.jsonl" + _write_jsonl(p, [_record(msg="a"), _record(msg="b")]) + + cur = SequentialCursor.load(p) + assert cur.total == 2 + assert cur.position == 0 + + +def test_cursor_load_skips_empty_lines(tmp_path): + p = tmp_path / "traj.jsonl" + p.write_text( + json.dumps(_record(msg="a")) + "\n\n \n" + json.dumps(_record(msg="b")) + "\n", + encoding="utf-8", + ) + + cur = SequentialCursor.load(p) + assert cur.total == 2 + + +def test_cursor_load_missing_file_raises(tmp_path): + with pytest.raises(FileNotFoundError): + SequentialCursor.load(tmp_path / "missing.jsonl") + + +def test_cursor_load_directory_raises(tmp_path): + """Path must be a single .jsonl file, not a directory.""" + with pytest.raises(FileNotFoundError): + SequentialCursor.load(tmp_path) + + +@pytest.mark.asyncio +async def test_cursor_next_returns_records_in_order(tmp_path): + p = tmp_path / "traj.jsonl" + _write_jsonl(p, [_record(msg="a", call_id="1"), _record(msg="b", call_id="2")]) + + cur = SequentialCursor.load(p) + first = await cur.next() + second = await cur.next() + + assert first["id"] == "1" + assert second["id"] == "2" + assert cur.position == 2 + + +@pytest.mark.asyncio +async def test_cursor_next_raises_trajectory_exhausted_when_done(tmp_path): + p = tmp_path / "traj.jsonl" + _write_jsonl(p, [_record(msg="only")]) + + cur = SequentialCursor.load(p) + await cur.next() + + with pytest.raises(TrajectoryExhausted) as exc_info: + await cur.next() + assert exc_info.value.position == 1 + assert exc_info.value.total == 1 + + +@pytest.mark.asyncio +async def test_cursor_reset_replays_from_start(tmp_path): + p = tmp_path / "traj.jsonl" + _write_jsonl(p, [_record(msg="a"), _record(msg="b")]) + + cur = SequentialCursor.load(p) + await cur.next() + await cur.next() + cur.reset() + + again = await cur.next() + assert again["messages"][0]["content"] == "a" + + +@pytest.mark.asyncio +async def test_cursor_model_mismatch_only_warns(tmp_path): + p = tmp_path / "traj.jsonl" + _write_jsonl(p, [_record(msg="a", model="gpt-3.5-turbo")]) + + cur = SequentialCursor.load(p) + record = await cur.next(expected_model="gpt-4o") # different model -> warn but don't raise + assert record["id"] == "x" diff --git a/uv.lock b/uv.lock index e00a7f86b3..cfed10409c 100644 --- a/uv.lock +++ b/uv.lock @@ -1196,6 +1196,15 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16" }, ] +[[package]] +name = "distro" +version = "1.9.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2" }, +] + [[package]] name = "docker" version = "7.1.0" @@ -1919,6 +1928,109 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12" }, ] +[[package]] +name = "jiter" +version = "0.14.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/6e/c1/0cddc6eb17d4c53a99840953f95dd3accdc5cfc7a337b0e9b26476276be9/jiter-0.14.0.tar.gz", hash = "sha256:e8a39e66dac7153cf3f964a12aad515afa8d74938ec5cc0018adcdae5367c79e" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/64/2e/a9959997739c403378d0a4a3a1c4ed80b60aeace216c4d37b303a9fc60a4/jiter-0.14.0-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:02f36a5c700f105ac04a6556fe664a59037a2c200db3b7e88784fac2ddf02531" }, + { url = "https://mirrors.aliyun.com/pypi/packages/27/72/b6de8a531e0adbadd839bec301165feb1fccf00e9ff55073ba2dd20f0043/jiter-0.14.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:41eab6c09ceffb6f0fe25e214b3068146edb1eda3649ca2aee2a061029c7ba2e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/db/d8/2040b9efa13c917f855c40890ae4119fe02c25b7c7677d5b4fa820a851fc/jiter-0.14.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5cf4d4c109641f9cfaf4a7b6aebd51654e405cd00fa9ebbf87163b8b97b325aa" }, + { url = "https://mirrors.aliyun.com/pypi/packages/49/62/655c0ad5ce6a8e90f9068c175b8a236877d753e460762b3183c136db1c5b/jiter-0.14.0-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b80c7b41a628e6be2213ad0ece763c5f88aa5ee003fa394d58acaaee1f4b8342" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f1/66/549c40fa068f08710b7570869c306a051eb67a29758bd64f4114f730554c/jiter-0.14.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fb3dbf7cc0d4dbe73cce307ebe7eefa7f73a7d3d854dd119ea0c243f03e40927" }, + { url = "https://mirrors.aliyun.com/pypi/packages/25/2f/97a32a05fed14ed58a18e181fdfb619e05163f3726b54ee6080ec0539c09/jiter-0.14.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7054adcdeb06b46efd17b5734f75817a44a2d06d3748e36c3a023a1bb52af9ec" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2a/3b/4347e1d6c2a973d653bbb7a2d671a2d2426e54b52ba735b8ff0d0a29b75c/jiter-0.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d597cd1bf6790376f3fffc7c708766e57301d99a19314824ea0ccc9c3c70e1e2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ef/24/ca452fbf2ea33548ed30ce68a39a50442d3f7c9bf0704a7af958a930c057/jiter-0.14.0-cp310-cp310-manylinux_2_31_riscv64.whl", hash = "sha256:df63a14878da754427926281626fd3ee249424a186e25a274e78176d42945264" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e3/a3/94470a0d199287caabeb4da2bb2ae5f6d17f3cf05dfc975d7cb064d58e0f/jiter-0.14.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4ea73187627bcc5810e085df715e8a99da8bdfd96a7eb36b4b4df700ba6d4c9c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/cf/71/6768edc09d7c45c39f093feb3de105fa718a3e982b5208b8a2ed6382b44b/jiter-0.14.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9f541eaf7bb8382367a1a23d6fc3d6aad57f8dd8c18c3c17f838bee20f217220" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3d/6b/5c2e17559a0f4e96e934479f7137df46c939e983fa05244e674815befb73/jiter-0.14.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:107465250de4fce00fdb47166bcd51df8e634e049541174fe3c71848e44f52ce" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b1/83/c25f3556a60fc74d11199100f1b6cc0c006b815c8494dea8ca16fe398732/jiter-0.14.0-cp310-cp310-win32.whl", hash = "sha256:ffb2a08a406465bb076b7cc1df41d833106d3cf7905076cc73f0cb90078c7d10" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2e/99/781a1b413f0989b7f2ea203b094b331685f1a35e52e0a45e5d000ecaab27/jiter-0.14.0-cp310-cp310-win_amd64.whl", hash = "sha256:cb8b682d10cb0cce7ff4c1af7244af7022c9b01ae16d46c357bdd0df13afb25d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/8a/1f/198ae537fccb7080a0ed655eb56abf64a92f79489dfbf79f40fa34225bcd/jiter-0.14.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:7e791e247b8044512e070bd1f3633dc08350d32776d2d6e7473309d0edf256a2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/cf/34/da67cff3fce964a36d03c3e365fb0f8726ade2a6cfd4d3c70107e216ead6/jiter-0.14.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:71527ce13fd5a0c4e40ad37331f8c547177dbb2dd0a93e5278b6a5eecf748804" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ed/36/4c72e67180d4e71a4f5dcf7886d0840e83c49ab11788172177a77570326e/jiter-0.14.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:02c4a7ab56f746014874f2c525584c0daca1dec37f66fd707ecef3b7e5c2228c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bc/db/9b39e09ceafa9878235c0fc29e3e3f9b12a4c6a98ea3085b998cadf3accc/jiter-0.14.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:376e9dafff914253bb9d46cdc5f7965607fbe7feb0a491c34e35f92b2770702e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b0/96/0dcba1d7a82c1b720774b48ef239376addbaf30df24c34742ac4a57b67b2/jiter-0.14.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:23ad2a7a9da1935575c820428dd8d2490ce4d23189691ce33da1fc0a58e14e1c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f1/e3/f61b71543e746e6b8b805e7755814fc242715c16f1dba58e1cbccb8032c2/jiter-0.14.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:54b3ddf5786bc7732d293bba3411ac637ecfa200a39983166d1df86a59a43c9f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ad/5e/0ddeb7096aca099114abe36c4921016e8d251e6f35f5890240b31f1f60ae/jiter-0.14.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5c001d5a646c2a50dc055dd526dad5d5245969e8234d2b1131d0451e81f3a373" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e9/d1/fe0c46cd7fda9cad8f1ff9ad217dc61f1e4280b21052ec6dfe88c1446ef2/jiter-0.14.0-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:834bb5bdabca2e91592a03d373838a8d0a1b8bbde7077ae6913fd2fc51812d00" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ac/21/f5317f91729b501019184771c80d60abd89907009e7bfa6c7e348c5bdd44/jiter-0.14.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4e9178be60e229b1b2b0710f61b9e24d1f4f8556985a83ff4c4f95920eea7314" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e9/05/79d8f33fb2bf168db0df5c9cd16fe440a8ada57e929d3677b22712c2568f/jiter-0.14.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:a7e4ccff04ec03614e62c613e976a3a5860dc9714ce8266f44328bdc8b1cab2c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5c/00/d1e3ff3d2a465e67f08507d74bafb2dcd29eba91dc939820e39e8dea38b8/jiter-0.14.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:69539d936fb5d55caf6ecd33e2e884de083ff0ea28579780d56c4403094bb8d9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/60/5b/bbb2189f62ace8d95e869aa4c84c9946616f301e2d02895a6f20dcc3bba3/jiter-0.14.0-cp311-cp311-win32.whl", hash = "sha256:4927d09b3e572787cc5e0a5318601448e1ab9391bcef95677f5840c2d00eaa6d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b8/86/c500b53dcbf08575f5963e536ebd757a1f7c568272ba5d180b212c9a87fb/jiter-0.14.0-cp311-cp311-win_amd64.whl", hash = "sha256:42d6ed359ac49eb922fdd565f209c57340aa06d589c84c8413e42a0f9ae1b842" }, + { url = "https://mirrors.aliyun.com/pypi/packages/75/4a/a676249049d42cb29bef82233e4fe0524d414cbe3606c7a4b311193c2f77/jiter-0.14.0-cp311-cp311-win_arm64.whl", hash = "sha256:6dd689f5f4a5a33747b28686e051095beb214fe28cfda5e9fe58a295a788f593" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5a/68/7390a418f10897da93b158f2d5a8bd0bcd73a0f9ec3bb36917085bb759ef/jiter-0.14.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:2fb2ce3a7bc331256dfb14cefc34832366bb28a9aca81deaf43bbf2a5659e607" }, + { url = "https://mirrors.aliyun.com/pypi/packages/60/a0/5854ac00ff63551c52c6c89534ec6aba4b93474e7924d64e860b1c94165b/jiter-0.14.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5252a7ca23785cef5d02d4ece6077a1b556a410c591b379f82091c3001e14844" }, + { url = "https://mirrors.aliyun.com/pypi/packages/41/a1/4f44832650a16b18e8391f1bf1d6ca4909bc738351826bcc198bba4357f4/jiter-0.14.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c409578cbd77c338975670ada777add4efd53379667edf0aceea730cabede6fb" }, + { url = "https://mirrors.aliyun.com/pypi/packages/48/64/a329e9d469f86307203594b1707e11ae51c3348d03bfd514a5f997870012/jiter-0.14.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7ede4331a1899d604463369c730dbb961ffdc5312bc7f16c41c2896415b1304a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/94/c1/5e3dfc59635aa4d4c7bd20a820ac1d09b8ed851568356802cf1c08edb3cf/jiter-0.14.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:92cd8b6025981a041f5310430310b55b25ca593972c16407af8837d3d7d2ca01" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e3/1b/dd157009dbc058f7b00108f545ccb72a2d56461395c4fc7b9cfdccb00af4/jiter-0.14.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:351bf6eda4e3a7ceb876377840c702e9a3e4ecc4624dbfb2d6463c67ae52637d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/91/78/256013667b7c10b8834f8e6e54cd3e562d4c6e34227a1596addccc05e38c/jiter-0.14.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c1dcfbeb93d9ecd9ca128bbf8910120367777973fa193fb9a39c31237d8df165" }, + { url = "https://mirrors.aliyun.com/pypi/packages/de/d9/137d65ade9093a409fe80955ce60b12bb753722c986467aeda47faf450ad/jiter-0.14.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:ae039aaef8de3f8157ecc1fdd4d85043ac4f57538c245a0afaecb8321ec951c3" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2e/48/76750835b87029342727c1a268bea8878ab988caf81ee4e7b880900eeb5a/jiter-0.14.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:7d9d51eb96c82a9652933bd769fe6de66877d6eb2b2440e281f2938c51b5643e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a6/60/456c4e81d5c8045279aefe60e9e483be08793828800a4e64add8fdde7f2a/jiter-0.14.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d824ca4148b705970bf4e120924a212fdfca9859a73e42bd7889a63a4ea6bb98" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a8/9f/2020e0984c235f678dced38fe4eec3058cf528e6af36ebf969b410305941/jiter-0.14.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:ff3a6465b3a0f54b1a430f45c3c0ba7d61ceb45cbc3e33f9e1a7f638d690baf3" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ef/32/e2d298e1a22a4bbe6062136d1c7192db7dba003a6975e51d9a9eecabc4c2/jiter-0.14.0-cp312-cp312-win32.whl", hash = "sha256:5dec7c0a3e98d2a3f8a2e67382d0d7c3ac60c69103a4b271da889b4e8bb1e129" }, + { url = "https://mirrors.aliyun.com/pypi/packages/36/ac/96369141b3d8a4a8e4590e983085efe1c436f35c0cda940dd76d942e3e40/jiter-0.14.0-cp312-cp312-win_amd64.whl", hash = "sha256:fc7e37b4b8bc7e80a63ad6cfa5fc11fab27dbfea4cc4ae644b1ab3f273dc348f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/01/c3/75d847f264647017d7e3052bbcc8b1e24b95fa139c320c5f5066fa7a0bdd/jiter-0.14.0-cp312-cp312-win_arm64.whl", hash = "sha256:ee4a72f12847ef29b072aee9ad5474041ab2924106bdca9fcf5d7d965853e057" }, + { url = "https://mirrors.aliyun.com/pypi/packages/97/2a/09f70020898507a89279659a1afe3364d57fc1b2c89949081975d135f6f5/jiter-0.14.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:af72f204cf4d44258e5b4c1745130ac45ddab0e71a06333b01de660ab4187a94" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d6/be/080c96a45cd74f9fce5db4fd68510b88087fb37ffe2541ff73c12db92535/jiter-0.14.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4b77da71f6e819be5fbcec11a453fde5b1d0267ef6ed487e2a392fd8e14e4e3a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7d/5e/2d0fee155826a968a832cc32438de5e2a193292c8721ca70d0b53e58245b/jiter-0.14.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:77f4ea612fe8b84b8b04e51d0e78029ecf3466348e25973f953de6e6a59aa4c1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/70/af/bf9ee0d3a4f8dc0d679fc1337f874fe60cdbf841ebbb304b374e1c9aaceb/jiter-0.14.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:62fe2451f8fcc0240261e6a4df18ecbcd58327857e61e625b2393ea3b468aac9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0f/83/8e8561eadba31f4d3948a5b712fb0447ec71c3560b57a855449e7b8ddc98/jiter-0.14.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6112f26f5afc75bcb475787d29da3aa92f9d09c7858f632f4be6ffe607be82e9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f6/c9/c5299e826a5fe6108d172b344033f61c69b1bb979dd8d9ddd4278a160971/jiter-0.14.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:215a6cb8fb7dc702aa35d475cc00ddc7f970e5c0b1417fb4b4ac5d82fa2a29db" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5d/37/c16d9d15c0a471b8644b1abe3c82668092a707d9bedcf076f24ff2e380cd/jiter-0.14.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc4ab96a30fb3cb2c7e0cd33f7616c8860da5f5674438988a54ac717caccdbaa" }, + { url = "https://mirrors.aliyun.com/pypi/packages/58/ea/8050cb0dc654e728e1bfacbc0c640772f2181af5dedd13ae70145743a439/jiter-0.14.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:3a99c1387b1f2928f799a9de899193484d66206a50e98233b6b088a7f0c1edb2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b0/3b/cf71506d270e5f84d97326bf220e47aed9b95e9a4a060758fb07772170ab/jiter-0.14.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ab18d11074485438695f8d34a1b6da61db9754248f96d51341956607a8f39985" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b0/cc/8c6c74a3efb5bd671bfd14f51e8a73375464ca914b1551bc3b40e26ac2c9/jiter-0.14.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:801028dcfc26ac0895e4964cbc0fd62c73be9fd4a7d7b1aaf6e5790033a719b7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/41/24/68d7b883ec959884ddf00d019b2e0e82ba81b167e1253684fa90519ce33c/jiter-0.14.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:ad425b087aafb4a1c7e1e98a279200743b9aaf30c3e0ba723aec93f061bd9bc8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b6/89/b1a0985223bbf3150ff9e8f46f98fc9360c1de94f48abe271bbe1b465682/jiter-0.14.0-cp313-cp313-win32.whl", hash = "sha256:882bcb9b334318e233950b8be366fe5f92c86b66a7e449e76975dfd6d776a01f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4c/19/3f339a5a7f14a11730e67f6be34f9d5105751d547b615ef593fa122a5ded/jiter-0.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:9b8c571a5dba09b98bd3462b5a53f27209a5cbbe85670391692ede71974e979f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/50/56/752dd89c84be0e022a8ea3720bcfa0a8431db79a962578544812ce061739/jiter-0.14.0-cp313-cp313-win_arm64.whl", hash = "sha256:34f19dcc35cb1abe7c369b3756babf8c7f04595c0807a848df8f26ef8298ef92" }, + { url = "https://mirrors.aliyun.com/pypi/packages/91/28/292916f354f25a1fe8cf2c918d1415c699a4a659ae00be0430e1c5d9ffea/jiter-0.14.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e89bcd7d426a75bb4952c696b267075790d854a07aad4c9894551a82c5b574ab" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ad/c7/b002a7d8b8957ac3d469bd59c18ef4b1595a5216ae0de639a287b9816023/jiter-0.14.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b25beaa0d4447ea8c7ae0c18c688905d34840d7d0b937f2f7bdd52162c98a40" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f9/3b/f8d07580d8706021d255a6356b8fab13ee4c869412995550ce6ed4ddf97d/jiter-0.14.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:651a8758dd413c51e3b7f6557cdc6921faf70b14106f45f969f091f5cda990ea" }, + { url = "https://mirrors.aliyun.com/pypi/packages/47/5b/ac1a974da29e35507230383110ffec59998b290a8732585d04e19a9eb5ba/jiter-0.14.0-cp313-cp313t-win_amd64.whl", hash = "sha256:e1a7eead856a5038a8d291f1447176ab0b525c77a279a058121b5fccee257f6f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/96/6d/9fc8433d667d2454271378a79747d8c76c10b51b482b454e6190e511f244/jiter-0.14.0-cp313-cp313t-win_arm64.whl", hash = "sha256:2e692633a12cda97e352fdcd1c4acc971b1c28707e1e33aeef782b0cbf051975" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4f/1e/354ed92461b165bd581f9ef5150971a572c873ec3b68a916d5aa91da3cc2/jiter-0.14.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:6f396837fc7577871ca8c12edaf239ed9ccef3bbe39904ae9b8b63ce0a48b140" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a6/95/8c7c7028aa8636ac21b7a55faef3e34215e6ed0cbf5ae58258427f621aa3/jiter-0.14.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:a4d50ea3d8ba4176f79754333bd35f1bbcd28e91adc13eb9b7ca91bc52a6cef9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/47/40/e2a852a44c4a089f2681a16611b7ce113224a80fd8504c46d78491b47220/jiter-0.14.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce17f8a050447d1b4153bda4fb7d26e6a9e74eb4f4a41913f30934c5075bf615" }, + { url = "https://mirrors.aliyun.com/pypi/packages/fc/1f/670f92adee1e9895eac41e8a4d623b6da68c4d46249d8b556b60b63f949e/jiter-0.14.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f4f1c4b125e1652aefbc2e2c1617b60a160ab789d180e3d423c41439e5f32850" }, + { url = "https://mirrors.aliyun.com/pypi/packages/01/2f/541c9ba567d05de1c4874a0f8f8c5e3fd78e2b874266623da9a775cf46e0/jiter-0.14.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:be808176a6a3a14321d18c603f2d40741858a7c4fc982f83232842689fe86dd9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ce/a9/c31cbec09627e0d5de7aeaec7690dba03e090caa808fefd8133137cf45bc/jiter-0.14.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:26679d58ba816f88c3849306dd58cb863a90a1cf352cdd4ef67e30ccf8a77994" }, + { url = "https://mirrors.aliyun.com/pypi/packages/50/02/3c05c1666c41904a2f607475a73e7a4763d1cbde2d18229c4f85b22dc253/jiter-0.14.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80381f5a19af8fa9aef743f080e34f6b25ebd89656475f8cf0470ec6157052aa" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7d/97/e15b33545c2b13518f560d695f974b9891b311641bdcf178d63177e8801e/jiter-0.14.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:004df5fdb8ecbd6d99f3227df18ba1a259254c4359736a2e6f036c944e02d7c5" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ad/d2/8b1461def6b96ba44530df20d07ef7a1c7da22f3f9bf1727e2d611077bf1/jiter-0.14.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:cff5708f7ed0fa098f2b53446c6fa74c48469118e5cd7497b4f1cd569ab06928" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e3/88/837566dd6ed6e452e8d3205355afd484ce44b2533edfa4ed73a298ea893e/jiter-0.14.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:2492e5f06c36a976d25c7cc347a60e26d5470178d44cde1b9b75e60b4e519f28" }, + { url = "https://mirrors.aliyun.com/pypi/packages/89/6b/b00b45c4d1b4c031777fe161d620b755b5b02cdade1e316dcb46e4471d63/jiter-0.14.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:7609cfbe3a03d37bfdbf5052012d5a879e72b83168a363deae7b3a26564d57de" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ad/d8/6fe5b42011d19397433d345716eac16728ac241862a2aac9c91923c7509a/jiter-0.14.0-cp314-cp314-win32.whl", hash = "sha256:7282342d32e357543565286b6450378c3cd402eea333fc1ebe146f1fabb306fc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e5/43/5c2e08da1efad5e410f0eaaabeadd954812612c33fbbd8fd5328b489139d/jiter-0.14.0-cp314-cp314-win_amd64.whl", hash = "sha256:bd77945f38866a448e73b0b7637366afa814d4617790ecd88a18ca74377e6c02" }, + { url = "https://mirrors.aliyun.com/pypi/packages/aa/1f/6e39ac0b4cdfa23e606af5b245df5f9adaa76f35e0c5096790da430ca506/jiter-0.14.0-cp314-cp314-win_arm64.whl", hash = "sha256:f2d4c61da0821ee42e0cdf5489da60a6d074306313a377c2b35af464955a3611" }, + { url = "https://mirrors.aliyun.com/pypi/packages/05/57/7dbc0ffbbb5176a27e3518716608aa464aee2e2887dc938f0b900a120449/jiter-0.14.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1bf7ff85517dd2f20a5750081d2b75083c1b269cf75afc7511bdf1f9548beb3b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/83/6e/7b3314398d8983f06b557aa21b670511ec72d3b79a68ee5e4d9bff972286/jiter-0.14.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c8ef8791c3e78d6c6b157c6d360fbb5c715bebb8113bc6a9303c5caff012754a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ae/4f/8dc674bcd7db6dba566de73c08c763c337058baff1dbeb34567045b27cdc/jiter-0.14.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e74663b8b10da1fe0f4e4703fd7980d24ad17174b6bb35d8498d6e3ebce2ae6a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3b/5f/188e09a1f20906f98bbdec44ed820e19f4e8eb8aff88b9d1a5a497587ff3/jiter-0.14.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1aca29ba52913f78362ec9c2da62f22cdc4c3083313403f90c15460979b84d9b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ac/f0/19046ef965ed8f349e8554775bb12ff4352f443fbe12b95d31f575891256/jiter-0.14.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8b39b7d87a952b79949af5fef44d2544e58c21a28da7f1bae3ef166455c61746" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c4/c3/da43bd8431ee175695777ee78cf0e93eacbb47393ff493f18c45231b427d/jiter-0.14.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:78d918a68b26e9fab068c2b5453577ef04943ab2807b9a6275df2a812599a310" }, + { url = "https://mirrors.aliyun.com/pypi/packages/72/26/e054771be889707c6161dbdec9c23d33a9ec70945395d70f07cfea1e9a6f/jiter-0.14.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:b08997c35aee1201c1a5361466a8fb9162d03ae7bf6568df70b6c859f1e654a4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c3/0f/7bea65ea2a6d91f2bf989ff11a18136644392bf2b0497a1fa50934c30a9c/jiter-0.14.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:260bf7ca20704d58d41f669e5e9fe7fe2fa72901a6b324e79056f5d52e9c9be2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3c/a1/b1ff7d70deef61ac0b7c6c2f12d2ace950cdeecb4fdc94500a0926802857/jiter-0.14.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:37826e3df29e60f30a382f9294348d0238ef127f4b5d7f5f8da78b5b9e050560" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0b/7b/3b0649983cbaf15eda26a414b5b1982e910c67bd6f7b1b490f3cfc76896a/jiter-0.14.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:645be49c46f2900937ba0eaf871ad5183c96858c0af74b6becc7f4e367e36e06" }, + { url = "https://mirrors.aliyun.com/pypi/packages/97/f8/33d78c83bd93ae0c0af05293a6660f88a1977caef39a6d72a84afab94ce0/jiter-0.14.0-cp314-cp314t-win32.whl", hash = "sha256:2f7877ed45118de283786178eceaf877110abacd04fde31efff3940ae9672674" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d6/ac/2b760516c03e2227826d1f7025d89bf6bf6357a28fe75c2a2800873c50bf/jiter-0.14.0-cp314-cp314t-win_amd64.whl", hash = "sha256:14c0cb10337c49f5eafe8e7364daca5e29a020ea03580b8f8e6c597fed4e1588" }, + { url = "https://mirrors.aliyun.com/pypi/packages/dc/2e/a44c20c58aeed0355f2d326969a181696aeb551a25195f47563908a815be/jiter-0.14.0-cp314-cp314t-win_arm64.whl", hash = "sha256:5419d4aa2024961da9fe12a9cfe7484996735dca99e8e090b5c88595ef1951ff" }, + { url = "https://mirrors.aliyun.com/pypi/packages/32/a1/ef34ca2cab2962598591636a1804b93645821201cc0095d4a93a9a329c9d/jiter-0.14.0-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:a25ffa2dbbdf8721855612f6dca15c108224b12d0c4024d0ac3d7902132b4211" }, + { url = "https://mirrors.aliyun.com/pypi/packages/60/bb/520576a532a6b8a6f42747afed289c8448c879a34d7802fe2c832d4fd38f/jiter-0.14.0-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:0ac9cbaa86c10996b92bd12c91659b60f939f8e28fcfa6bc11a0e90a774ce95b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b2/7c/c16db114ea1f2f532f198aa8dc39585026af45af362c69a0492f31bc4821/jiter-0.14.0-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:844e73b6c56b505e9e169234ea3bdea2ea43f769f847f47ac559ba1d2361ebea" }, + { url = "https://mirrors.aliyun.com/pypi/packages/99/8f/15e7741ff19e9bcd4d753f7ff22f988fd54592f134ca13701c13ea8c20e0/jiter-0.14.0-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e52c076f187405fc21523c746c04399c9af8ece566077ed147b2126f2bcba577" }, + { url = "https://mirrors.aliyun.com/pypi/packages/21/42/9042c3f3019de4adcb8c16591c325ec7255beea9fcd33a42a43f3b0b1000/jiter-0.14.0-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:fbd9e482663ca9d005d051330e4d2d8150bb208a209409c10f7e7dfdf7c49da9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/60/cf/a7e19b308bd86bb04776803b1f01a5f9a287a4c55205f4708827ee487fbf/jiter-0.14.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:33a20d838b91ef376b3a56896d5b04e725c7df5bc4864cc6569cf046a8d73b6d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ca/44/e26ede3f0caeff93f222559cb0cc4ca68579f07d009d7b6010c5b586f9b1/jiter-0.14.0-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:432c4db5255d86a259efde91e55cb4c8d18c0521d844c9e2e7efcce3899fb016" }, + { url = "https://mirrors.aliyun.com/pypi/packages/da/e9/1f9ada30cef7b05e74bb06f52127e7a724976c225f46adb65c37b1dadfb6/jiter-0.14.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:67f00d94b281174144d6532a04b66a12cb866cbdc47c3af3bfe2973677f9861a" }, +] + [[package]] name = "jmespath" version = "0.10.0" @@ -2653,6 +2765,25 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/be/9c/92789c596b8df838baa98fa71844d84283302f7604ed565dafe5a6b5041a/oauthlib-3.3.1-py3-none-any.whl", hash = "sha256:88119c938d2b8fb88561af5f6ee0eec8cc8d552b7bb1f712743136eb7523b7a1" }, ] +[[package]] +name = "openai" +version = "2.36.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/f4/a1/4d5e84cf51720fc1526cc49e10ac1961abcccb55b0efb3d970db1e9a2728/openai-2.36.0.tar.gz", hash = "sha256:139dea0edd2f1b30c33d46ae1a6929e03906254140318e4608e98fe8c566f2e7" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/9d/1c/5d43735b2553baae2a5e899dcbcd0670a86930d993184d72ca909bf11c9b/openai-2.36.0-py3-none-any.whl", hash = "sha256:143f6194b548dbc2c921af1f1b03b9f14c85fed8a75b5b516f5bcc11a2a50c63" }, +] + [[package]] name = "opencensus" version = "0.11.4" @@ -4118,6 +4249,8 @@ builder = [ model-service = [ { name = "alibabacloud-cr20181201" }, { name = "fastapi" }, + { name = "httpx" }, + { name = "openai" }, { name = "psutil" }, { name = "swebench" }, { name = "uvicorn" }, @@ -4180,10 +4313,12 @@ requires-dist = [ { name = "gem-llm", marker = "extra == 'rocklet'", specifier = ">=0.1.0" }, { name = "gem-llm", marker = "extra == 'sandbox-actor'", specifier = ">=0.1.0" }, { name = "httpx" }, + { name = "httpx", marker = "extra == 'model-service'" }, { name = "kubernetes", marker = "extra == 'admin'", specifier = ">=35.0.0" }, { name = "nacos-sdk-python", marker = "extra == 'admin'", specifier = ">=0.1.14" }, { name = "nacos-sdk-python", marker = "extra == 'sandbox-actor'", specifier = ">=0.1.14" }, { name = "numpy", marker = "extra == 'rocklet'", specifier = "<=2.2.6" }, + { name = "openai", marker = "extra == 'model-service'", specifier = ">=1.50.0" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-exporter-prometheus" }, From 7eeefda63e9bbe9f94caa2ee5438bb5f41151479 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Wed, 13 May 2026 17:58:56 +0800 Subject: [PATCH 088/226] feat(rocklet): add Windows PowerShell support #921 (#922) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * Rock supports Windows PowerShell, thereby circumventing the issue of Python libraries that are unavailable on Windows. * refactor(rocklet): extract platform adapters for local sandbox Move BashSession (Linux/macOS) and PowerShellSession (Windows) into a new rock/rocklet/platforms/ subpackage behind a PlatformAdapter ABC. LocalSandboxRuntime now delegates session creation and disk-root lookup to the adapter chosen by get_platform_adapter(). Lazy imports keep pexpect/bashlex off Windows and the Windows adapter off POSIX. uv.lock marks pexpect/bashlex/twisted/gem-llm with sys_platform != 'win32' so resolution succeeds on Windows. Adds UnsupportedPlatformError for unknown platforms. Backward-compat re-exports of BashSession/PowerShellSession from local_sandbox are preserved. * fix ut * refactor(rocklet): collapse PlatformAdapter into Rocklet base class - Drop the PlatformAdapter indirection introduced in fa62c3bcd. Rename LocalSandboxRuntime to Rocklet and let LinuxRocklet/WindowsRocklet subclass it directly via Rocklet.create() factory dispatch. - Flatten module layout: remove rock/rocklet/local_sandbox.py and the rock/rocklet/platforms/ package; move bash/PowerShell sessions into rock/rocklet/{linux,windows}.py and the abstract base into rock/rocklet/rocklet.py. - Rename LocalSandboxRuntimeConfig -> RockletConfig and switch the runtime discriminator from "local" to "rocklet". - Force PowerShell to UTF-8 at startup via chcp 65001 + Console OutputEncoding/InputEncoding, applied through the -Command argument so it takes effect before PowerShell caches its stdout TextWriter (fixes zh-CN/GBK mojibake on Windows for Get-ChildItem, Format-Table and stdin Chinese paths). - Update unit-test imports and call sites to the new module paths. --------- Co-authored-by: 杨钊 --- pyproject.toml | 8 +- rock/actions/__init__.py | 4 +- rock/actions/sandbox/config.py | 12 +- rock/deployments/local.py | 6 +- rock/rocklet/exceptions.py | 20 ++ rock/rocklet/{local_sandbox.py => linux.py} | 271 ++-------------- rock/rocklet/local_api.py | 32 +- rock/rocklet/rocklet.py | 296 ++++++++++++++++++ rock/rocklet/windows.py | 288 +++++++++++++++++ tests/unit/rocklet/test_local_command.py | 2 +- .../rocklet/test_local_sandbox_runtime.py | 12 +- uv.lock | 58 ++-- 12 files changed, 695 insertions(+), 314 deletions(-) rename rock/rocklet/{local_sandbox.py => linux.py} (57%) create mode 100644 rock/rocklet/rocklet.py create mode 100644 rock/rocklet/windows.py diff --git a/pyproject.toml b/pyproject.toml index d7d7a591b0..7fbe447e70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,11 +57,11 @@ rocklet = [ # Rocklet execution environment dependencies "fastapi", "uvicorn", - "pexpect", - "bashlex", + "pexpect; sys_platform != 'win32'", + "bashlex; sys_platform != 'win32'", "psutil", - "twisted", - "gem-llm>=0.1.0", + "twisted; sys_platform != 'win32'", + "gem-llm>=0.1.0; sys_platform != 'win32'", "numpy<=2.2.6", ] diff --git a/rock/actions/__init__.py b/rock/actions/__init__.py index 3c6462fce5..ce4d295ba3 100644 --- a/rock/actions/__init__.py +++ b/rock/actions/__init__.py @@ -3,7 +3,7 @@ from .envs.response import EnvCloseResponse, EnvListResponse, EnvMakeResponse, EnvResetResponse, EnvStepResponse from .response import BaseResponse, ResponseStatus, RockResponse from .sandbox.base import AbstractSandbox, _ExceptionTransfer -from .sandbox.config import LocalSandboxRuntimeConfig, RemoteSandboxRuntimeConfig, SandboxRuntimeConfig +from .sandbox.config import RemoteSandboxRuntimeConfig, RockletConfig, SandboxRuntimeConfig from .sandbox.request import ( Action, BashAction, @@ -80,7 +80,7 @@ "UploadResponse", "FileUploadResponse", "CloseResponse", - "LocalSandboxRuntimeConfig", + "RockletConfig", "RemoteSandboxRuntimeConfig", "SandboxResponse", "SandboxRuntimeConfig", diff --git a/rock/actions/sandbox/config.py b/rock/actions/sandbox/config.py index 21021ab3ea..eeca366e50 100644 --- a/rock/actions/sandbox/config.py +++ b/rock/actions/sandbox/config.py @@ -5,18 +5,18 @@ from rock.actions import AbstractSandbox -class LocalSandboxRuntimeConfig(BaseModel): - """Configuration for local sandbox runtime execution.""" +class RockletConfig(BaseModel): + """Configuration for Rocklet (local sandbox runtime).""" model_config = ConfigDict(extra="forbid") - type: Literal["local"] = "local" + type: Literal["rocklet"] = "rocklet" """Runtime type discriminator for serialization/deserialization and CLI parsing. Should not be modified.""" def get_sandbox_runtime(self) -> AbstractSandbox: - from rock.rocklet.local_sandbox import LocalSandboxRuntime + from rock.rocklet.rocklet import Rocklet - return LocalSandboxRuntime.from_config(self) + return Rocklet.from_config(self) class RemoteSandboxRuntimeConfig(BaseModel): @@ -41,4 +41,4 @@ def get_sandbox_runtime(self) -> AbstractSandbox: # Union type for all supported sandbox runtime configurations -SandboxRuntimeConfig = LocalSandboxRuntimeConfig | RemoteSandboxRuntimeConfig +SandboxRuntimeConfig = RockletConfig | RemoteSandboxRuntimeConfig diff --git a/rock/deployments/local.py b/rock/deployments/local.py index 899f127001..30cfc49373 100644 --- a/rock/deployments/local.py +++ b/rock/deployments/local.py @@ -8,7 +8,7 @@ from rock.deployments.hooks.abstract import CombinedDeploymentHook, DeploymentHook from rock.logger import init_logger from rock.rocklet.exceptions import DeploymentNotStartedError -from rock.rocklet.local_sandbox import LocalSandboxRuntime +from rock.rocklet.rocklet import Rocklet logger = init_logger(__name__) @@ -51,7 +51,7 @@ async def is_alive(self, *, timeout: float | None = None) -> IsAliveResponse: async def start(self): """Starts the runtime.""" - self._runtime = LocalSandboxRuntime() + self._runtime = Rocklet.create() async def stop(self): """Stops the runtime.""" @@ -60,7 +60,7 @@ async def stop(self): self._runtime = None @property - def runtime(self) -> LocalSandboxRuntime: + def runtime(self) -> Rocklet: """Returns the runtime if running. Raises: diff --git a/rock/rocklet/exceptions.py b/rock/rocklet/exceptions.py index af6d7de651..0359dea59e 100644 --- a/rock/rocklet/exceptions.py +++ b/rock/rocklet/exceptions.py @@ -25,6 +25,22 @@ def __init__(self, message: str, *, extra_info: dict[str, Any] = None): self.extra_info = extra_info +class PowerShellError(RockletException, RuntimeError): + """Raised when a PowerShell session encounters an error.""" + + def __init__(self, message: str, *, extra_info: dict[str, Any] = None): + super().__init__(message) + if extra_info is None: + extra_info = {} + self.extra_info = extra_info + + +class PowerShellNotFoundError(RockletException, FileNotFoundError): + """Raised when PowerShell executable is not found on the system.""" + + pass + + class CommandTimeoutError(RockletException, RuntimeError, TimeoutError): ... @@ -52,3 +68,7 @@ class DeploymentStartupError(RockletException, RuntimeError): class DockerPullError(DeploymentStartupError): ... + + +class UnsupportedPlatformError(RockletException, NotImplementedError): + """Raised when the current OS has no Rocklet subclass implementation.""" diff --git a/rock/rocklet/local_sandbox.py b/rock/rocklet/linux.py similarity index 57% rename from rock/rocklet/local_sandbox.py rename to rock/rocklet/linux.py index 1bc1c11905..00bb49499b 100644 --- a/rock/rocklet/local_sandbox.py +++ b/rock/rocklet/linux.py @@ -1,70 +1,48 @@ +"""Linux/macOS Rocklet for the local sandbox runtime. + +Hosts the BashSession implementation (built on pexpect + bashlex) and the +LinuxRocklet that the central dispatcher returns for sys.platform in +{'linux', 'darwin'}. +""" + import asyncio import os import re -import shutil import subprocess import time -from abc import ABC, abstractmethod -from concurrent.futures import ThreadPoolExecutor from copy import deepcopy -from pathlib import Path -from typing import Any import bashlex import bashlex.ast -import gem import pexpect import psutil -from typing_extensions import Self from rock.actions import ( - AbstractSandbox, BashObservation, CloseBashSessionResponse, - CloseResponse, CloseSessionResponse, - CommandResponse, CreateBashSessionResponse, - CreateSessionResponse, - EnvCloseResponse, - EnvListResponse, - EnvMakeResponse, - EnvResetResponse, - EnvStepResponse, - IsAliveResponse, - LocalSandboxRuntimeConfig, - Observation, - ReadFileRequest, - ReadFileResponse, - UploadRequest, - UploadResponse, - WriteFileRequest, - WriteFileResponse, ) -from rock.admin.proto.request import SandboxAction as Action from rock.admin.proto.request import SandboxBashAction as BashAction -from rock.admin.proto.request import SandboxCloseSessionRequest as CloseSessionRequest -from rock.admin.proto.request import SandboxCommand as Command from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest -from rock.admin.proto.request import SandboxCreateSessionRequest as CreateSessionRequest -from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest -from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest from rock.logger import init_logger from rock.rocklet.exceptions import ( BashIncorrectSyntaxError, CommandTimeoutError, NoExitCodeError, NonZeroExitCodeError, - SessionDoesNotExistError, - SessionExistsError, SessionNotInitializedError, ) from rock.utils import get_executor -__all__ = ["LocalSandboxRuntime", "BashSession"] +from .rocklet import Rocklet, Session +logger = init_logger(__name__) -logger = init_logger("rock.actions.local") + +def _strip_control_chars(s: str) -> str: + ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]") + return ansi_escape.sub("", s) def _split_bash_command(inpt: str) -> list[str]: @@ -104,11 +82,6 @@ def find_range(cmd: bashlex.ast.node) -> tuple[int, int]: return cmd_strings -def _strip_control_chars(s: str) -> str: - ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]") - return ansi_escape.sub("", s) - - def _check_bash_command(command: str) -> None: """Check if a bash command is valid. Raises BashIncorrectSyntaxError if it's not.""" _unique_string = "SOUNIQUEEOF" @@ -126,20 +99,6 @@ def _check_bash_command(command: str) -> None: raise exc -class Session(ABC): - @abstractmethod - async def start(self) -> CreateSessionResponse: - ... - - @abstractmethod - async def run(self, action: Action) -> Observation: - ... - - @abstractmethod - async def close(self) -> CloseSessionResponse: - ... - - class BashSession(Session): _UNIQUE_STRING = "UNIQUESTRING29234" @@ -380,202 +339,16 @@ def interact(self) -> None: self.shell.interact() -class LocalSandboxRuntime(AbstractSandbox): - def __init__(self, *, executor: ThreadPoolExecutor | None = None, **kwargs: Any): - """A Runtime that runs locally and actually executes commands in a shell. - If you are deploying to Modal/Fargate/etc., this class will be running within the docker container - on Modal/Fargate/etc. - - Args: - **kwargs: Keyword arguments (see `LocalSandboxConfig` for details). - """ - self._config = LocalSandboxRuntimeConfig(**kwargs) - self._sessions: dict[str, Session] = {} - # Set up logger - self.command_logger = init_logger("command", "command.log") - self._executor = executor - self._gem_envs: dict[str, gem.Env] = {} - - @classmethod - def from_config(cls, config: LocalSandboxRuntimeConfig) -> Self: - return cls(**config.model_dump()) - - @property - def sessions(self) -> dict[str, Session]: - return self._sessions - - async def is_alive(self, *, timeout: float | None = None) -> IsAliveResponse: - """Checks if the runtime is alive.""" - return IsAliveResponse(is_alive=True) - - async def create_session(self, request: CreateSessionRequest) -> CreateSessionResponse: - """Creates a new session.""" - if request.session in self.sessions: - msg = f"session {request.session} already exists" - raise SessionExistsError(msg) - if isinstance(request, CreateBashSessionRequest): - session = BashSession(request) - else: - msg = f"unknown session type: {request!r}" - raise ValueError(msg) - self.sessions[request.session] = session - self.command_logger.info(f"[create_session]:{request.session}") - return await session.start() - - async def run_in_session(self, action: Action) -> Observation: - """Runs a command in a session.""" - if action.session not in self.sessions: - msg = f"session {action.session!r} does not exist" - raise SessionDoesNotExistError(msg) - self.command_logger.info(f"[run_in_session input][{action.session}]:{action.command}") - observation = await self.sessions[action.session].run(action) - if observation.output: - self.command_logger.info(f"[run_in_session output][{action.session}]:{observation.output}") - if observation.exit_code: - self.command_logger.info(f"[run_in_session exit_code][{action.session}]:{observation.exit_code}") - if observation.failure_reason: - self.command_logger.info(f"[run_in_session failure_reason][{action.session}]:{observation.failure_reason}") - return observation - - async def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse: - """Closes a shell session.""" - if request.session not in self.sessions: - msg = f"session {request.session!r} does not exist" - raise SessionDoesNotExistError(msg) - out = await self.sessions[request.session].close() - del self.sessions[request.session] - self.command_logger.info(f"[close_session]:{request.session}") - return out - - async def execute(self, command: Command) -> CommandResponse: - """Executes a command (independent of any shell session). - - Raises: - CommandTimeoutError: If the command times out. - NonZeroExitCodeError: If the command has a non-zero exit code and `check` is True. - """ - self.command_logger.info(f"[execute input]:{command.command}") - loop = asyncio.get_running_loop() - try: - result = await loop.run_in_executor(self._executor, self._run_subprocess_blocking, command) - r = CommandResponse( - stdout=result.stdout.decode(errors="backslashreplace"), - stderr=result.stderr.decode(errors="backslashreplace"), - exit_code=result.returncode, - ) - except subprocess.TimeoutExpired as e: - msg = f"Timeout ({command.timeout}s) exceeded while running command" - raise CommandTimeoutError(msg) from e - if command.check and result.returncode != 0: - msg = ( - f"Command {command.command!r} failed with exit code {result.returncode}. " - f"Stdout:\n{r.stdout!r}\nStderr:\n{r.stderr!r}" - ) - if command.error_msg: - msg = f"{command.error_msg}: {msg}" - raise NonZeroExitCodeError(msg) - if r.stdout: - self.command_logger.info(f"[execute stdout]:{r.stdout}") - if r.stderr: - self.command_logger.info(f"[execute stderr]:{r.stderr}") - return r +class LinuxRocklet(Rocklet): + """Rocklet implementation for sys.platform in {'linux', 'darwin'}.""" - def _run_subprocess_blocking(self, command: Command): - # This is synchronous blocking code - return subprocess.run( - command.command, - shell=command.shell, - timeout=command.timeout, - env=command.env, - capture_output=True, - cwd=command.cwd, - ) + def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: + return BashSession(request) - async def read_file(self, request: ReadFileRequest) -> ReadFileResponse: - """Reads a file""" - self.command_logger.info(f"[read_file input]: {request.path}") - content = Path(request.path).read_text(encoding=request.encoding, errors=request.errors) - self.command_logger.info(f"[read_file output]: {content[:1000]}") - return ReadFileResponse(content=content) - - async def write_file(self, request: WriteFileRequest) -> WriteFileResponse: - """Writes a file""" - self.command_logger.info(f"[write_file input]: {request.path}") - self.command_logger.info(f"[write_file content]: {request.content[:1000]}") - Path(request.path).parent.mkdir(parents=True, exist_ok=True) - Path(request.path).write_text(request.content) - return WriteFileResponse() - - async def upload(self, request: UploadRequest) -> UploadResponse: - """Uploads a file""" - self.command_logger.info(f"[upload source]: {request.source_path}") - self.command_logger.info(f"[upload target]: {request.target_path}") - if Path(request.source_path).is_dir(): - shutil.copytree(request.source_path, request.target_path) - else: - shutil.copy(request.source_path, request.target_path) - self.command_logger.info("[upload output]: upload success!") - return UploadResponse() - - async def close(self) -> CloseResponse: - """Closes the runtime.""" - for session in self.sessions.values(): - await session.close() - return CloseResponse() - - async def get_statistics(self): - cpu_percent: float = psutil.cpu_percent() - mem_percent: float = psutil.virtual_memory().percent - disk_percent: float = psutil.disk_usage("/").percent - net_io: int = psutil.net_io_counters().bytes_recv + psutil.net_io_counters().bytes_sent + async def get_statistics(self) -> dict: return { - "cpu": cpu_percent, - "mem": mem_percent, - "disk": disk_percent, - "net": net_io, + "cpu": psutil.cpu_percent(), + "mem": psutil.virtual_memory().percent, + "disk": psutil.disk_usage("/").percent, + "net": psutil.net_io_counters().bytes_recv + psutil.net_io_counters().bytes_sent, } - - def env_make(self, env_id: str, sandbox_id: str) -> EnvMakeResponse: - """ - Make gem env - """ - env = gem.make(env_id) - self._gem_envs[sandbox_id] = env - return EnvMakeResponse(sandbox_id=sandbox_id) - - def env_step(self, sandbox_id: str, action: str) -> EnvStepResponse: - """ - Step gem env - """ - env = self._gem_envs[sandbox_id] - observation, reward, terminated, truncated, info = env.step(action) - return EnvStepResponse( - observation=observation, - reward=reward, - terminated=terminated, - truncated=truncated, - info=info, - ) - - def env_reset(self, sandbox_id: str, seed: int | None = None) -> EnvResetResponse: - """ - Reset gem env - """ - env = self._gem_envs[sandbox_id] - observation, info = env.reset(seed=seed) - return EnvResetResponse(observation=observation, info=info) - - def env_close(self, sandbox_id: str) -> EnvCloseResponse: - """ - Close gem env - """ - del self._gem_envs[sandbox_id] - return EnvCloseResponse(sandbox_id=sandbox_id) - - def env_list(self) -> EnvListResponse: - """ - List gem env - """ - from gem.envs.registration import ENV_REGISTRY - - return EnvListResponse(env_id=list(ENV_REGISTRY)) diff --git a/rock/rocklet/local_api.py b/rock/rocklet/local_api.py index 0bdbe41f33..362688a74d 100644 --- a/rock/rocklet/local_api.py +++ b/rock/rocklet/local_api.py @@ -27,7 +27,7 @@ from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest from rock.common.port_validation import validate_port_forward_port from rock.logger import init_logger -from rock.rocklet.local_sandbox import LocalSandboxRuntime +from rock.rocklet.rocklet import Rocklet from rock.utils import get_executor logger = init_logger(__name__) @@ -38,7 +38,7 @@ TCP_CONNECT_TIMEOUT = 10 # seconds IDLE_TIMEOUT = 300 # seconds -runtime = LocalSandboxRuntime(executor=get_executor()) +rocklet = Rocklet.create(executor=get_executor()) def serialize_model(model): @@ -47,42 +47,42 @@ def serialize_model(model): @local_router.get("/is_alive") async def is_alive(): - return serialize_model(await runtime.is_alive()) + return serialize_model(await rocklet.is_alive()) @local_router.get("/get_statistics") async def get_statistics(): - return await runtime.get_statistics() + return await rocklet.get_statistics() @local_router.post("/create_session") async def create_session(request: CreateSessionRequest): - return serialize_model(await runtime.create_session(request)) + return serialize_model(await rocklet.create_session(request)) @local_router.post("/run_in_session") async def run(action: Action): - return serialize_model(await runtime.run_in_session(action)) + return serialize_model(await rocklet.run_in_session(action)) @local_router.post("/close_session") async def close_session(request: CloseSessionRequest): - return serialize_model(await runtime.close_session(request)) + return serialize_model(await rocklet.close_session(request)) @local_router.post("/execute") async def execute(command: Command): - return serialize_model(await runtime.execute(command=command)) + return serialize_model(await rocklet.execute(command=command)) @local_router.post("/read_file") async def read_file(request: ReadFileRequest): - return serialize_model(await runtime.read_file(request)) + return serialize_model(await rocklet.read_file(request)) @local_router.post("/write_file") async def write_file(request: WriteFileRequest): - return serialize_model(await runtime.write_file(request)) + return serialize_model(await rocklet.write_file(request)) @local_router.post("/upload") @@ -112,33 +112,33 @@ async def upload( @local_router.post("/close") async def close(): - await runtime.close() + await rocklet.close() return CloseResponse() @local_router.post("/env/make") async def env_make(request: EnvMakeRequest) -> EnvMakeResponse: - return runtime.env_make(env_id=request.env_id, sandbox_id=request.sandbox_id) + return rocklet.env_make(env_id=request.env_id, sandbox_id=request.sandbox_id) @local_router.post("/env/step") async def env_step(request: EnvStepRequest) -> EnvStepResponse: - return runtime.env_step(request.sandbox_id, request.action) + return rocklet.env_step(request.sandbox_id, request.action) @local_router.post("/env/reset") async def env_reset(request: EnvResetRequest) -> EnvResetResponse: - return runtime.env_reset(request.sandbox_id, request.seed) + return rocklet.env_reset(request.sandbox_id, request.seed) @local_router.post("/env/close") async def env_close(request: EnvCloseRequest) -> EnvCloseResponse: - return runtime.env_close(request.sandbox_id) + return rocklet.env_close(request.sandbox_id) @local_router.post("/env/list") async def env_list() -> EnvListResponse: - return runtime.env_list() + return rocklet.env_list() @local_router.websocket("/portforward") diff --git a/rock/rocklet/rocklet.py b/rock/rocklet/rocklet.py new file mode 100644 index 0000000000..0a50d5186d --- /dev/null +++ b/rock/rocklet/rocklet.py @@ -0,0 +1,296 @@ +"""Rocklet — abstract base class for local sandbox runtimes. + +Concrete subclasses live in: +- rock.rocklet.linux: LinuxRocklet (Linux/macOS, BashSession via pexpect) +- rock.rocklet.windows: WindowsRocklet (Windows, PowerShellSession via subprocess) + +Use Rocklet.create(**kwargs) to obtain the right subclass for the current OS. +""" + +import asyncio +import shutil +import subprocess +import sys +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +from rock.actions import ( + AbstractSandbox, + CloseResponse, + CloseSessionResponse, + CommandResponse, + CreateSessionResponse, + EnvCloseResponse, + EnvListResponse, + EnvMakeResponse, + EnvResetResponse, + EnvStepResponse, + IsAliveResponse, + Observation, + ReadFileRequest, + ReadFileResponse, + RockletConfig, + UploadRequest, + UploadResponse, + WriteFileRequest, + WriteFileResponse, +) +from rock.admin.proto.request import SandboxAction as Action +from rock.admin.proto.request import SandboxCloseSessionRequest as CloseSessionRequest +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest +from rock.admin.proto.request import SandboxCreateSessionRequest as CreateSessionRequest +from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest +from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest +from rock.logger import init_logger +from rock.rocklet.exceptions import ( + CommandTimeoutError, + NonZeroExitCodeError, + SessionDoesNotExistError, + SessionExistsError, + UnsupportedPlatformError, +) + +logger = init_logger(__name__) + + +class Session(ABC): + """Abstract command session inside a sandbox. + + Concrete implementations: BashSession (Linux/macOS), PowerShellSession (Windows). + + Signatures intentionally use the wide discriminated types + (Action / Observation / CreateSessionResponse / CloseSessionResponse) to + keep the contract identical to the original Session ABC. + """ + + @abstractmethod + async def start(self) -> CreateSessionResponse: + ... + + @abstractmethod + async def run(self, action: Action) -> Observation: + ... + + @abstractmethod + async def close(self) -> CloseSessionResponse: + ... + + +class Rocklet(AbstractSandbox, ABC): + """Abstract base for local sandbox runtimes.""" + + def __init__(self, *, executor: ThreadPoolExecutor | None = None, **kwargs: Any): + """A Runtime that runs locally and actually executes commands in a shell. + If you are deploying to Modal/Fargate/etc., this class will be running within the docker container + on Modal/Fargate/etc. + + Args: + **kwargs: Keyword arguments (see `RockletConfig` for details). + """ + self._config = RockletConfig(**kwargs) + self._sessions: dict[str, Session] = {} + # Set up logger + self.command_logger = init_logger("command", "command.log") + self._executor = executor + self._gem_envs: dict[str, Any] = {} + + @classmethod + def create(cls, **kwargs: Any) -> "Rocklet": + """Construct the Rocklet subclass for the current OS. + + Lazy-imports the platform-specific module so that, e.g., pexpect is never + loaded on Windows and the Windows-only module is never executed on Linux. + + Raises: + UnsupportedPlatformError: if sys.platform has no Rocklet subclass. + """ + match sys.platform: + case "linux" | "darwin": + from rock.rocklet.linux import LinuxRocklet + + return LinuxRocklet(**kwargs) + case "win32": + from rock.rocklet.windows import WindowsRocklet + + return WindowsRocklet(**kwargs) + case other: + raise UnsupportedPlatformError(f"No Rocklet subclass registered for sys.platform={other!r}") + + @classmethod + def from_config(cls, config: RockletConfig) -> "Rocklet": + return cls.create(**config.model_dump()) + + @abstractmethod + def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: + """Construct the platform's default bash-style session.""" + + @abstractmethod + async def get_statistics(self) -> dict: + """Return CPU / memory / disk / net usage as a dict.""" + + @property + def sessions(self) -> dict[str, Session]: + return self._sessions + + async def is_alive(self, *, timeout: float | None = None) -> IsAliveResponse: + """Checks if the runtime is alive.""" + return IsAliveResponse(is_alive=True) + + async def create_session(self, request: CreateSessionRequest) -> CreateSessionResponse: + """Creates a new session. + + Delegates to the subclass's `_build_bash_session`, which selects + BashSession on Linux/macOS and PowerShellSession on Windows. + """ + if request.session in self.sessions: + msg = f"session {request.session} already exists" + raise SessionExistsError(msg) + if isinstance(request, CreateBashSessionRequest): + session = self._build_bash_session(request) + else: + msg = f"unknown session type: {request!r}" + raise ValueError(msg) + self.sessions[request.session] = session + self.command_logger.info(f"[create_session]:{request.session}") + return await session.start() + + async def run_in_session(self, action: Action) -> Observation: + """Runs a command in a session.""" + if action.session not in self.sessions: + msg = f"session {action.session!r} does not exist" + raise SessionDoesNotExistError(msg) + self.command_logger.info(f"[run_in_session input][{action.session}]:{action.command}") + observation = await self.sessions[action.session].run(action) + if observation.output: + self.command_logger.info(f"[run_in_session output][{action.session}]:{observation.output}") + if observation.exit_code: + self.command_logger.info(f"[run_in_session exit_code][{action.session}]:{observation.exit_code}") + if observation.failure_reason: + self.command_logger.info(f"[run_in_session failure_reason][{action.session}]:{observation.failure_reason}") + return observation + + async def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse: + """Closes a shell session.""" + if request.session not in self.sessions: + msg = f"session {request.session!r} does not exist" + raise SessionDoesNotExistError(msg) + out = await self.sessions[request.session].close() + del self.sessions[request.session] + self.command_logger.info(f"[close_session]:{request.session}") + return out + + async def execute(self, command: Command) -> CommandResponse: + """Executes a command (independent of any shell session). + + Raises: + CommandTimeoutError: If the command times out. + NonZeroExitCodeError: If the command has a non-zero exit code and `check` is True. + """ + self.command_logger.info(f"[execute input]:{command.command}") + loop = asyncio.get_running_loop() + try: + result = await loop.run_in_executor(self._executor, self._run_subprocess_blocking, command) + r = CommandResponse( + stdout=result.stdout.decode(errors="backslashreplace"), + stderr=result.stderr.decode(errors="backslashreplace"), + exit_code=result.returncode, + ) + except subprocess.TimeoutExpired as e: + msg = f"Timeout ({command.timeout}s) exceeded while running command" + raise CommandTimeoutError(msg) from e + if command.check and result.returncode != 0: + msg = ( + f"Command {command.command!r} failed with exit code {result.returncode}. " + f"Stdout:\n{r.stdout!r}\nStderr:\n{r.stderr!r}" + ) + if command.error_msg: + msg = f"{command.error_msg}: {msg}" + raise NonZeroExitCodeError(msg) + if r.stdout: + self.command_logger.info(f"[execute stdout]:{r.stdout}") + if r.stderr: + self.command_logger.info(f"[execute stderr]:{r.stderr}") + return r + + def _run_subprocess_blocking(self, command: Command): + # This is synchronous blocking code + return subprocess.run( + command.command, + shell=command.shell, + timeout=command.timeout, + env=command.env, + capture_output=True, + cwd=command.cwd, + ) + + async def read_file(self, request: ReadFileRequest) -> ReadFileResponse: + """Reads a file""" + self.command_logger.info(f"[read_file input]: {request.path}") + content = Path(request.path).read_text(encoding=request.encoding, errors=request.errors) + self.command_logger.info(f"[read_file output]: {content[:1000]}") + return ReadFileResponse(content=content) + + async def write_file(self, request: WriteFileRequest) -> WriteFileResponse: + """Writes a file""" + self.command_logger.info(f"[write_file input]: {request.path}") + self.command_logger.info(f"[write_file content]: {request.content[:1000]}") + Path(request.path).parent.mkdir(parents=True, exist_ok=True) + Path(request.path).write_text(request.content) + return WriteFileResponse(success=True) + + async def upload(self, request: UploadRequest) -> UploadResponse: + """Uploads a file""" + self.command_logger.info(f"[upload source]: {request.source_path}") + self.command_logger.info(f"[upload target]: {request.target_path}") + if Path(request.source_path).is_dir(): + shutil.copytree(request.source_path, request.target_path) + else: + shutil.copy(request.source_path, request.target_path) + self.command_logger.info("[upload output]: upload success!") + return UploadResponse() + + async def close(self) -> CloseResponse: + """Closes the runtime.""" + for session in self.sessions.values(): + await session.close() + return CloseResponse() + + def env_make(self, env_id: str, sandbox_id: str) -> EnvMakeResponse: + """Make gem env""" + import gem + + env = gem.make(env_id) + self._gem_envs[sandbox_id] = env + return EnvMakeResponse(sandbox_id=sandbox_id) + + def env_step(self, sandbox_id: str, action: str) -> EnvStepResponse: + """Step gem env""" + env = self._gem_envs[sandbox_id] + observation, reward, terminated, truncated, info = env.step(action) + return EnvStepResponse( + observation=observation, + reward=reward, + terminated=terminated, + truncated=truncated, + info=info, + ) + + def env_reset(self, sandbox_id: str, seed: int | None = None) -> EnvResetResponse: + """Reset gem env""" + env = self._gem_envs[sandbox_id] + observation, info = env.reset(seed=seed) + return EnvResetResponse(observation=observation, info=info) + + def env_close(self, sandbox_id: str) -> EnvCloseResponse: + """Close gem env""" + del self._gem_envs[sandbox_id] + return EnvCloseResponse(sandbox_id=sandbox_id) + + def env_list(self) -> EnvListResponse: + """List gem env""" + from gem.envs.registration import ENV_REGISTRY + + return EnvListResponse(env_id=list(ENV_REGISTRY)) diff --git a/rock/rocklet/windows.py b/rock/rocklet/windows.py new file mode 100644 index 0000000000..b02d905fd9 --- /dev/null +++ b/rock/rocklet/windows.py @@ -0,0 +1,288 @@ +"""Windows Rocklet for the local sandbox runtime. + +Hosts the PowerShellSession implementation (built on subprocess + a background +reader thread) and the WindowsRocklet that the central dispatcher returns for +sys.platform == 'win32'. +""" + +import asyncio +import os +import queue as queue_module +import re +import shutil +import subprocess +import threading +import time + +import psutil + +from rock.actions import ( + BashObservation, + CloseBashSessionResponse, + CloseSessionResponse, + CreateBashSessionResponse, +) +from rock.admin.proto.request import SandboxBashAction as BashAction +from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest +from rock.logger import init_logger +from rock.rocklet.exceptions import ( + CommandTimeoutError, + NonZeroExitCodeError, + PowerShellNotFoundError, + SessionNotInitializedError, +) +from rock.utils import get_executor + +from .rocklet import Rocklet, Session + +logger = init_logger(__name__) + + +def _strip_control_chars(s: str) -> str: + ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]") + return ansi_escape.sub("", s) + + +class PowerShellSession(Session): + """A session that runs PowerShell commands on Windows. + + Uses subprocess.Popen with a background reader thread to interact + with PowerShell interactively via stdin/stdout pipes. + Unique marker strings are used to delimit command output boundaries + and extract exit codes. + """ + + _BEGIN_MARKER = "ROCKLET_PS_BEGIN_29234" + _END_MARKER = "ROCKLET_PS_END_29234" + _EXIT_MARKER = "ROCKLET_PS_EXIT_29234:" + + def __init__(self, request: CreateBashSessionRequest): + self.request = request + self._process: subprocess.Popen | None = None + self._executor = get_executor() + self._output_queue: queue_module.Queue = queue_module.Queue() + self._reader_thread: threading.Thread | None = None + self._lock = threading.Lock() + + @staticmethod + def _find_powershell() -> str: + """Find the PowerShell executable. Prefers pwsh (PowerShell Core) over powershell.""" + for cmd in ["pwsh", "powershell"]: + if shutil.which(cmd): + return cmd + raise PowerShellNotFoundError( + "PowerShell executable not found. Install PowerShell Core ('pwsh') or ensure 'powershell.exe' is in PATH." + ) + + def _read_stdout(self) -> None: + """Background thread that continuously reads stdout line by line.""" + try: + assert self._process is not None and self._process.stdout is not None + for line in iter(self._process.stdout.readline, ""): + if not line: + break + self._output_queue.put(line) + except Exception: + pass + + def _drain_queue(self, timeout: float = 0.1) -> str: + """Drain all currently available lines from the output queue.""" + lines = [] + deadline = time.time() + timeout + while time.time() < deadline: + try: + line = self._output_queue.get(timeout=0.05) + lines.append(line.rstrip("\r\n")) + except queue_module.Empty: + if lines: + break + continue + return "\n".join(lines) + + def _send_line(self, text: str) -> None: + """Send a line of text to PowerShell's stdin.""" + assert self._process is not None and self._process.stdin is not None + self._process.stdin.write(text + "\n") + self._process.stdin.flush() + + async def start(self) -> CreateBashSessionResponse: + """Start the PowerShell session.""" + env = os.environ.copy() + if self.request.env is not None: + env.update(self.request.env) + + ps_cmd = self._find_powershell() + logger.info(f"Starting PowerShell session with: {ps_cmd}") + + # Force PowerShell to use UTF-8 for stdin/stdout encoding from the very + # first byte. Required on non-UTF-8 Windows locales (e.g. zh-CN/GBK): + # without this, PowerShell writes its output in the system code page + # while subprocess.Popen below decodes as UTF-8, producing mojibake + # (e.g. "目录" -> "Ŀ¼") that then crashes the GBK-encoded console + # logger with UnicodeEncodeError. + # + # `-Command` runs the setup expression at startup, then `-NoExit` drops + # PowerShell into interactive mode for stdin-driven commands. Doing this + # at startup (instead of via _send_line after spawn) is essential because + # PowerShell caches its [Console].Out TextWriter at process init — late + # OutputEncoding mutations don't affect already-cached writers used by + # built-in formatters (Get-ChildItem, Format-Table, etc.). + ps_setup = ( + "chcp 65001 > $null; " + "[Console]::OutputEncoding = [System.Text.Encoding]::UTF8; " + "$OutputEncoding = [System.Text.Encoding]::UTF8" + ) + self._process = subprocess.Popen( + [ps_cmd, "-NoLogo", "-NoProfile", "-NoExit", "-Command", ps_setup], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + env=env, + ) + + self._reader_thread = threading.Thread(target=self._read_stdout, daemon=True) + self._reader_thread.start() + + # Wait for PowerShell to initialize + time.sleep(0.5) + + # Drain startup output + startup_output = self._drain_queue(timeout=0.5) + logger.info(f"PowerShell session started, startup output: {startup_output[:200]}") + return CreateBashSessionResponse(output=startup_output) + + async def run(self, action: BashAction) -> BashObservation: + """Run a command in the PowerShell session. + + Raises: + SessionNotInitializedError: If the PowerShell process is not running. + CommandTimeoutError: If the command times out. + NonZeroExitCodeError: If the command has a non-zero exit code and action.check is 'raise'. + """ + if self._process is None or self._process.poll() is not None: + raise SessionNotInitializedError("PowerShell session not initialized or has terminated") + + loop = asyncio.get_running_loop() + result = await loop.run_in_executor(self._executor, self._run_command, action) + + if action.check == "raise" and result.exit_code is not None and result.exit_code != 0: + msg = ( + f"Command {action.command!r} failed with exit code {result.exit_code}. " + f"Here is the output:\n{result.output!r}" + ) + if hasattr(action, "error_msg") and action.error_msg: + msg = f"{action.error_msg}: {msg}" + raise NonZeroExitCodeError(msg) + + return result + + def _run_command(self, action: BashAction) -> BashObservation: + """Execute a single command in the PowerShell session (blocking). + + Wraps the command with unique markers to reliably detect output boundaries + and extract the exit code ($LASTEXITCODE). + """ + with self._lock: + # Drain any leftover output from previous commands + self._drain_queue(timeout=0.1) + + # Send command wrapped with markers + self._send_line(f"Write-Host '{self._BEGIN_MARKER}'") + # Send command lines + for line in action.command.splitlines(): + self._send_line(line) + # Capture exit code: use $LASTEXITCODE for native commands, + # but also check $? for PowerShell cmdlet errors + self._send_line( + f"if ($LASTEXITCODE -ne $null -and $LASTEXITCODE -ne 0) " + f"{{ Write-Host '{self._EXIT_MARKER}'$LASTEXITCODE }} " + f"elseif (-not $?) {{ Write-Host '{self._EXIT_MARKER}1' }} " + f"else {{ Write-Host '{self._EXIT_MARKER}0' }}" + ) + self._send_line(f"Write-Host '{self._END_MARKER}'") + + # Collect output until end marker + output_lines: list[str] = [] + started = False + exit_code: int | None = None + timeout = action.timeout if action.timeout is not None else 1200 + deadline = time.time() + timeout + + while True: + remaining = deadline - time.time() + if remaining <= 0: + raise CommandTimeoutError( + f"timeout after {timeout} seconds while running command {action.command!r}" + ) + try: + line = self._output_queue.get(timeout=min(remaining, 1.0)) + line = line.rstrip("\r\n") + except queue_module.Empty: + # Check if process is still alive + if self._process is not None and self._process.poll() is not None: + raise SessionNotInitializedError("PowerShell process terminated unexpectedly") + continue + + if self._BEGIN_MARKER in line: + started = True + continue + if self._EXIT_MARKER in line: + after_marker = line.split(self._EXIT_MARKER)[-1].strip() + try: + exit_code = int(after_marker) if after_marker else 0 + except ValueError: + exit_code = 0 + continue + if self._END_MARKER in line: + break + if started: + output_lines.append(line) + + output = "\n".join(output_lines).strip() + output = _strip_control_chars(output) + + if action.check == "ignore": + return BashObservation(output=output, exit_code=None) + + return BashObservation(output=output, exit_code=exit_code) + + async def close(self) -> CloseSessionResponse: + """Close the PowerShell session.""" + if self._process is not None: + try: + self._send_line("exit") + except Exception: + pass + try: + self._process.terminate() + self._process.wait(timeout=5) + except Exception: + try: + self._process.kill() + except Exception: + pass + self._process = None + return CloseBashSessionResponse() + + def interact(self) -> None: + """Interactive mode is not supported for PowerShell sessions.""" + raise NotImplementedError("Interactive mode is not supported for PowerShell sessions") + + +class WindowsRocklet(Rocklet): + """Rocklet implementation for sys.platform == 'win32'.""" + + def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: + return PowerShellSession(request) + + async def get_statistics(self) -> dict: + return { + "cpu": psutil.cpu_percent(), + "mem": psutil.virtual_memory().percent, + "disk": psutil.disk_usage("C:\\").percent, + "net": psutil.net_io_counters().bytes_recv + psutil.net_io_counters().bytes_sent, + } diff --git a/tests/unit/rocklet/test_local_command.py b/tests/unit/rocklet/test_local_command.py index 5d7877168a..d28ca59c3e 100644 --- a/tests/unit/rocklet/test_local_command.py +++ b/tests/unit/rocklet/test_local_command.py @@ -1,7 +1,7 @@ import pytest from rock.rocklet.exceptions import BashIncorrectSyntaxError -from rock.rocklet.local_sandbox import _check_bash_command, _split_bash_command +from rock.rocklet.linux import _check_bash_command, _split_bash_command def test_split_bash_command_normal(): diff --git a/tests/unit/rocklet/test_local_sandbox_runtime.py b/tests/unit/rocklet/test_local_sandbox_runtime.py index b2e7d73e8c..ca1ac5f166 100644 --- a/tests/unit/rocklet/test_local_sandbox_runtime.py +++ b/tests/unit/rocklet/test_local_sandbox_runtime.py @@ -9,16 +9,16 @@ from rock.admin.proto.request import SandboxCloseBashSessionRequest as CloseBashSessionRequest from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest -from rock.rocklet.local_sandbox import LocalSandboxRuntime +from rock.rocklet.rocklet import Rocklet @pytest.fixture def local_runtime(): - return LocalSandboxRuntime() + return Rocklet.create() @pytest.mark.asyncio -async def test_upload_file(local_runtime: LocalSandboxRuntime, tmp_path: Path): +async def test_upload_file(local_runtime: Rocklet, tmp_path: Path): file_path = tmp_path / "source.txt" file_path.write_text("test") tmp_target = tmp_path / "target.txt" @@ -27,7 +27,7 @@ async def test_upload_file(local_runtime: LocalSandboxRuntime, tmp_path: Path): @pytest.mark.asyncio -async def test_upload_directory(local_runtime: LocalSandboxRuntime, tmp_path: Path): +async def test_upload_directory(local_runtime: Rocklet, tmp_path: Path): dir_path = tmp_path / "source_dir" dir_path.mkdir() (dir_path / "file1.txt").write_text("test1") @@ -39,7 +39,7 @@ async def test_upload_directory(local_runtime: LocalSandboxRuntime, tmp_path: Pa @pytest.mark.asyncio -async def test_gem(local_runtime: LocalSandboxRuntime): +async def test_gem(local_runtime: Rocklet): env_id = "game:Sokoban-v0-easy" exmaple_gem_env: SokobanEnv = gem.make(env_id) @@ -66,7 +66,7 @@ async def test_gem(local_runtime: LocalSandboxRuntime): @pytest.mark.asyncio -async def test_prompt_command(local_runtime: LocalSandboxRuntime): +async def test_prompt_command(local_runtime: Rocklet): prompt_command = "echo ROCK" await local_runtime.create_session( CreateBashSessionRequest(env={"PROMPT_COMMAND": prompt_command}, session_type="bash") diff --git a/uv.lock b/uv.lock index cfed10409c..bfd81a2782 100644 --- a/uv.lock +++ b/uv.lock @@ -2,10 +2,14 @@ version = 1 revision = 3 requires-python = ">=3.10, <4.0" resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version < '3.11'", + "python_full_version >= '3.13' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'win32'", ] [[package]] @@ -940,7 +944,8 @@ name = "contourpy" version = "1.3.2" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } resolution-markers = [ - "python_full_version < '3.11'", + "python_full_version < '3.11' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'win32'", ] dependencies = [ { name = "numpy", marker = "python_full_version < '3.11'" }, @@ -1010,9 +1015,12 @@ name = "contourpy" version = "1.3.3" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", + "python_full_version >= '3.13' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", ] dependencies = [ { name = "numpy", marker = "python_full_version >= '3.11'" }, @@ -4194,22 +4202,22 @@ admin = [ { name = "alibabacloud-cr20181201" }, { name = "apscheduler" }, { name = "asyncpg" }, - { name = "bashlex" }, + { name = "bashlex", marker = "sys_platform != 'win32'" }, { name = "boto3" }, { name = "cryptography" }, { name = "fakeredis", extra = ["json"] }, { name = "fastapi" }, - { name = "gem-llm" }, + { name = "gem-llm", marker = "sys_platform != 'win32'" }, { name = "kubernetes" }, { name = "nacos-sdk-python" }, { name = "numpy" }, - { name = "pexpect" }, + { name = "pexpect", marker = "sys_platform != 'win32'" }, { name = "pip" }, { name = "psutil" }, { name = "ray", extra = ["default"] }, { name = "redis" }, { name = "sqlmodel" }, - { name = "twisted" }, + { name = "twisted", marker = "sys_platform != 'win32'" }, { name = "uvicorn" }, { name = "websockets" }, ] @@ -4220,7 +4228,7 @@ all = [ { name = "alibabacloud-cr20181201" }, { name = "apscheduler" }, { name = "asyncpg" }, - { name = "bashlex" }, + { name = "bashlex", marker = "sys_platform != 'win32'" }, { name = "boto3" }, { name = "cryptography" }, { name = "docker" }, @@ -4230,14 +4238,14 @@ all = [ { name = "kubernetes" }, { name = "nacos-sdk-python" }, { name = "numpy" }, - { name = "pexpect" }, + { name = "pexpect", marker = "sys_platform != 'win32'" }, { name = "pip" }, { name = "psutil" }, { name = "ray", extra = ["default"] }, { name = "redis" }, { name = "sqlmodel" }, { name = "swebench" }, - { name = "twisted" }, + { name = "twisted", marker = "sys_platform != 'win32'" }, { name = "uvicorn" }, { name = "websockets" }, ] @@ -4256,13 +4264,13 @@ model-service = [ { name = "uvicorn" }, ] rocklet = [ - { name = "bashlex" }, + { name = "bashlex", marker = "sys_platform != 'win32'" }, { name = "fastapi" }, - { name = "gem-llm" }, + { name = "gem-llm", marker = "sys_platform != 'win32'" }, { name = "numpy" }, - { name = "pexpect" }, + { name = "pexpect", marker = "sys_platform != 'win32'" }, { name = "psutil" }, - { name = "twisted" }, + { name = "twisted", marker = "sys_platform != 'win32'" }, { name = "uvicorn" }, ] sandbox-actor = [ @@ -4301,7 +4309,7 @@ requires-dist = [ { name = "apscheduler", marker = "extra == 'admin'" }, { name = "apscheduler", marker = "extra == 'sandbox-actor'", specifier = ">=3.11.0" }, { name = "asyncpg", marker = "extra == 'admin'" }, - { name = "bashlex", marker = "extra == 'rocklet'" }, + { name = "bashlex", marker = "sys_platform != 'win32' and extra == 'rocklet'" }, { name = "boto3", marker = "extra == 'admin'" }, { name = "build" }, { name = "cryptography", marker = "extra == 'admin'", specifier = "==39.0.1" }, @@ -4309,8 +4317,8 @@ requires-dist = [ { name = "fakeredis", extras = ["json"], marker = "extra == 'admin'" }, { name = "fastapi", marker = "extra == 'model-service'" }, { name = "fastapi", marker = "extra == 'rocklet'" }, + { name = "gem-llm", marker = "sys_platform != 'win32' and extra == 'rocklet'", specifier = ">=0.1.0" }, { name = "gem-llm", marker = "extra == 'builder'", specifier = ">=0.1.0" }, - { name = "gem-llm", marker = "extra == 'rocklet'", specifier = ">=0.1.0" }, { name = "gem-llm", marker = "extra == 'sandbox-actor'", specifier = ">=0.1.0" }, { name = "httpx" }, { name = "httpx", marker = "extra == 'model-service'" }, @@ -4324,7 +4332,7 @@ requires-dist = [ { name = "opentelemetry-exporter-prometheus" }, { name = "opentelemetry-sdk" }, { name = "oss2" }, - { name = "pexpect", marker = "extra == 'rocklet'" }, + { name = "pexpect", marker = "sys_platform != 'win32' and extra == 'rocklet'" }, { name = "pip", marker = "extra == 'admin'" }, { name = "psutil", marker = "extra == 'model-service'" }, { name = "psutil", marker = "extra == 'rocklet'" }, @@ -4344,7 +4352,7 @@ requires-dist = [ { name = "sqlmodel", marker = "extra == 'sandbox-actor'" }, { name = "swebench", marker = "extra == 'builder'" }, { name = "swebench", marker = "extra == 'model-service'" }, - { name = "twisted", marker = "extra == 'rocklet'" }, + { name = "twisted", marker = "sys_platform != 'win32' and extra == 'rocklet'" }, { name = "tzdata" }, { name = "uuid" }, { name = "uvicorn", marker = "extra == 'model-service'" }, @@ -5562,25 +5570,21 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/d8/c1/eec33cc9f847ebeb0bc6234d7d45fe3fc0a6fe8fc5b5e6be0442bd2c684d/zope_interface-8.0.1-cp310-cp310-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:758803806b962f32c87b31bb18c298b022965ba34fe532163831cc39118c24ab" }, { url = "https://mirrors.aliyun.com/pypi/packages/58/7d/1e3476a1ef0175559bd8492dc7bb921ad0df5b73861d764b1f824ad5484a/zope_interface-8.0.1-cp310-cp310-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f8e88f35f86bbe8243cad4b2972deef0fdfca0a0723455abbebdc83bbab96b69" }, { url = "https://mirrors.aliyun.com/pypi/packages/bc/67/ba5ea98ff23f723c5cbe7db7409f2e43c9fe2df1ced67881443c01e64478/zope_interface-8.0.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7844765695937d9b0d83211220b72e2cf6ac81a08608ad2b58f2c094af498d83" }, - { url = "https://mirrors.aliyun.com/pypi/packages/2b/a7/b1b8b6c13fba955c043cdee409953ee85f652b106493e2e931a84f95c1aa/zope_interface-8.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:64fa7b206dd9669f29d5c1241a768bebe8ab1e8a4b63ee16491f041e058c09d0" }, { url = "https://mirrors.aliyun.com/pypi/packages/f2/2f/c10c739bcb9b072090c97c2e08533777497190daa19d190d72b4cce9c7cb/zope_interface-8.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4bd01022d2e1bce4a4a4ed9549edb25393c92e607d7daa6deff843f1f68b479d" }, { url = "https://mirrors.aliyun.com/pypi/packages/b5/e1/9845ac3697f108d9a1af6912170c59a23732090bbfb35955fe77e5544955/zope_interface-8.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:29be8db8b712d94f1c05e24ea230a879271d787205ba1c9a6100d1d81f06c69a" }, { url = "https://mirrors.aliyun.com/pypi/packages/f2/49/6573bc8b841cfab18e80c8e8259f1abdbbf716140011370de30231be79ad/zope_interface-8.0.1-cp311-cp311-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:51ae1b856565b30455b7879fdf0a56a88763b401d3f814fa9f9542d7410dbd7e" }, { url = "https://mirrors.aliyun.com/pypi/packages/e2/fd/908b0fd4b1ab6e412dfac9bd2b606f2893ef9ba3dd36d643f5e5b94c57b3/zope_interface-8.0.1-cp311-cp311-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d2e7596149cb1acd1d4d41b9f8fe2ffc0e9e29e2e91d026311814181d0d9efaf" }, { url = "https://mirrors.aliyun.com/pypi/packages/dc/78/8419a2b4e88410520ed4b7f93bbd25a6d4ae66c4e2b131320f2b90f43077/zope_interface-8.0.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b2737c11c34fb9128816759864752d007ec4f987b571c934c30723ed881a7a4f" }, - { url = "https://mirrors.aliyun.com/pypi/packages/e5/90/caf68152c292f1810e2bd3acd2177badf08a740aa8a348714617d6c9ad0b/zope_interface-8.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:cf66e4bf731aa7e0ced855bb3670e8cda772f6515a475c6a107bad5cb6604103" }, { url = "https://mirrors.aliyun.com/pypi/packages/dc/a6/0f08713ddda834c428ebf97b2a7fd8dea50c0100065a8955924dbd94dae8/zope_interface-8.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:115f27c1cc95ce7a517d960ef381beedb0a7ce9489645e80b9ab3cbf8a78799c" }, { url = "https://mirrors.aliyun.com/pypi/packages/e9/5e/d423045f54dc81e0991ec655041e7a0eccf6b2642535839dd364b35f4d7f/zope_interface-8.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:af655c573b84e3cb6a4f6fd3fbe04e4dc91c63c6b6f99019b3713ef964e589bc" }, { url = "https://mirrors.aliyun.com/pypi/packages/c6/43/39d4bb3f7a80ebd261446792493cfa4e198badd47107224f5b6fe1997ad9/zope_interface-8.0.1-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:23f82ef9b2d5370750cc1bf883c3b94c33d098ce08557922a3fbc7ff3b63dfe1" }, { url = "https://mirrors.aliyun.com/pypi/packages/da/29/49effcff64ef30731e35520a152a9dfcafec86cf114b4c2aff942e8264ba/zope_interface-8.0.1-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:35a1565d5244997f2e629c5c68715b3d9d9036e8df23c4068b08d9316dcb2822" }, { url = "https://mirrors.aliyun.com/pypi/packages/c7/39/b947673ec9a258eeaa20208dd2f6127d9fbb3e5071272a674ebe02063a78/zope_interface-8.0.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:029ea1db7e855a475bf88d9910baab4e94d007a054810e9007ac037a91c67c6f" }, - { url = "https://mirrors.aliyun.com/pypi/packages/8f/ee/eed6efd1fc3788d1bef7a814e0592d8173b7fe601c699b935009df035fc2/zope_interface-8.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0beb3e7f7dc153944076fcaf717a935f68d39efa9fce96ec97bafcc0c2ea6cab" }, { url = "https://mirrors.aliyun.com/pypi/packages/5f/dc/3c12fca01c910c793d636ffe9c0984e0646abaf804e44552070228ed0ede/zope_interface-8.0.1-cp313-cp313-macosx_10_9_x86_64.whl", hash = "sha256:c7cc027fc5c61c5d69e5080c30b66382f454f43dc379c463a38e78a9c6bab71a" }, { url = "https://mirrors.aliyun.com/pypi/packages/46/71/6127b7282a3e380ca927ab2b40778a9c97935a4a57a2656dadc312db5f30/zope_interface-8.0.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:fcf9097ff3003b7662299f1c25145e15260ec2a27f9a9e69461a585d79ca8552" }, { url = "https://mirrors.aliyun.com/pypi/packages/56/86/4387a9f951ee18b0e41fda77da77d59c33e59f04660578e2bad688703e64/zope_interface-8.0.1-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:6d965347dd1fb9e9a53aa852d4ded46b41ca670d517fd54e733a6b6a4d0561c2" }, { url = "https://mirrors.aliyun.com/pypi/packages/61/08/ce60a114466abc067c68ed41e2550c655f551468ae17b4b17ea360090146/zope_interface-8.0.1-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9a3b8bb77a4b89427a87d1e9eb969ab05e38e6b4a338a9de10f6df23c33ec3c2" }, { url = "https://mirrors.aliyun.com/pypi/packages/36/9a/62a9ba3a919594605a07c34eee3068659bbd648e2fa0c4a86d876810b674/zope_interface-8.0.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:87e6b089002c43231fb9afec89268391bcc7a3b66e76e269ffde19a8112fb8d5" }, - { url = "https://mirrors.aliyun.com/pypi/packages/da/06/8fe88bd7edef60566d21ef5caca1034e10f6b87441ea85de4bbf9ea74768/zope_interface-8.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:64a43f5280aa770cbafd0307cb3d1ff430e2a1001774e8ceb40787abe4bb6658" }, ] [[package]] From cc95a8ad47dda20123ce36c6441f2691167003c2 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Thu, 14 May 2026 09:47:19 +0800 Subject: [PATCH 089/226] fix(rocklet): symlink mount into /bin for nix images with kata runtime (#936) In nix-based images, mount resides under /nix/store and is not in PATH, causing setup_kata_dind to fail. Symlink mount from util-linux into /bin and add /bin to PATH before invoking kata setup. Co-authored-by: Claude Opus 4.6 --- rock/rocklet/local_files/docker_run.sh | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/rock/rocklet/local_files/docker_run.sh b/rock/rocklet/local_files/docker_run.sh index e319d5879a..72ebe86898 100755 --- a/rock/rocklet/local_files/docker_run.sh +++ b/rock/rocklet/local_files/docker_run.sh @@ -47,11 +47,6 @@ setup_kata_dind() { mount -o remount,rw /proc/sys } -if [ "${ROCK_KATA_RUNTIME}" = "true" ]; then - echo "Kata runtime detected, setting up DinD disk..." - setup_kata_dind -fi - # Run rocklet if [ "$(is_nix)" = "true" ]; then # NixOS @@ -59,6 +54,8 @@ if [ "$(is_nix)" = "true" ]; then ln -sf $(ls -d /nix/store/*glibc*/lib64 2>/dev/null | head -1) /lib64 mkdir -p /bin ln -sf $(ls -d /nix/store/*bash*/bin/bash 2>/dev/null | head -1) /bin/bash + ln -sf $(ls -d /nix/store/*util-linux*/bin/mount 2>/dev/null | head -1) /bin/mount + export PATH="/bin:${PATH}" GCC_LIB=$(ls -d /nix/store/*gcc*lib/lib 2>/dev/null | head -1) ZLIB_LIB=$(ls -d /nix/store/*zlib*/lib 2>/dev/null | head -1) NIX_LIBS="" @@ -67,6 +64,11 @@ if [ "$(is_nix)" = "true" ]; then [ -n "$NIX_LIBS" ] && export LD_LIBRARY_PATH="${NIX_LIBS}${LD_LIBRARY_PATH}" fi +if [ "${ROCK_KATA_RUNTIME}" = "true" ]; then + echo "Kata runtime detected, setting up DinD disk..." + setup_kata_dind +fi + if [ "$(is_musl)" = "true" ]; then # musl-based distributions if [ ! -d /tmp/local_files/alpine_glibc ]; then From e9a519ab4d14f67ebecd91a1a8064de8161f18cd Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 14 May 2026 09:49:06 +0800 Subject: [PATCH 090/226] feat(job): integrate in-sandbox model-service proxy for record/replay (#938) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(job): integrate in-sandbox model-service proxy for record/replay Adds optional ProxyConfig under environment.proxy that brings up rock model-service as a record/replay proxy inside the sandbox before any agent code runs. When enabled: - AbstractTrial._setup_proxy installs the proxy via PythonRuntimeEnv + install_cmd, uploads the local replay jsonl when in replay mode, then starts the proxy. After start, it detects the outer sandbox's eth0 IP via `hostname -I` and rewrites env['OPENAI_BASE_URL'] to the proxy URL, so harbor yaml serialization and inner docker containers also go through the proxy. - HarborTrial.setup() and BashTrial.setup() invoke _setup_proxy first. - JobExecutor._wrap_with_proxy_bootstrap prepends a bash snippet to the build script so the same OPENAI_BASE_URL override applies at runtime for direct sandbox processes (covers both flat and docker-in-docker). Verified end-to-end on SWE-bench psf__requests-1142 with mini-swe-agent: record reward=1.0 (80-line jsonl), replay reward=1.0 with identical input/output token counts (1,285,362 / 14,083), agent_execution shrunk from 9m38s to 12s. * refactor(job): extract SANDBOX_REPLAY_FILE to env var Move the hard-coded sandbox replay path to ROCK_JOB_PROXY_REPLAY_FILE so users can override it without patching code, and rename the default file from _rock_replay.jsonl to rock-job-proxy-replay.jsonl (no underscore prefix, more descriptive). * refactor(job): remove _wrap_with_proxy_bootstrap, reorder session creation _wrap_with_proxy_bootstrap and its hostname-I runtime detection are no longer needed — _setup_proxy already detects the outer sandbox IP and rewrites environment.env["OPENAI_BASE_URL"] to the proxy URL before the session is created. The bash export was redundant since the session already inherits the resolved proxy URL. Also fix a bug: create_session was called before trial.setup(), so the session env captured the upstream OPENAI_BASE_URL instead of the proxy URL that _setup_proxy writes later. Moved create_session after setup so the session picks up the rewritten env. * refactor(job): extract common _setup_proxy + _upload_files into AbstractTrial.setup Both HarborTrial and BashTrial were calling _setup_proxy and _upload_files as the first two lines of setup(). Promote those calls to a concrete AbstractTrial.setup(), so subclasses call super().setup(sandbox) first then add their own logic. * refactor(job): drop SANDBOX_REPLAY_FILE constant, use env var directly * refactor(job): inline pip_install_cmd into ModelServiceConfig * refactor(job): extract _detect_and_rewrite_proxy_url helper Pull the hostname-I detection and OPENAI_BASE_URL rewrite out of _setup_proxy into a module-level async function for clarity. * refactor(job): convert proxy helpers to AbstractTrial methods Move _build_proxy_start_cmd and _detect_and_rewrite_proxy_url from module-level functions to methods on AbstractTrial, reading proxy and env from self._config instead of accepting them as parameters. --- rock/env_vars.py | 4 + rock/sdk/envhub/config.py | 48 ++- rock/sdk/job/executor.py | 3 +- rock/sdk/job/trial/abstract.py | 87 +++++- rock/sdk/job/trial/bash.py | 2 +- rock/sdk/job/trial/harbor.py | 2 +- tests/unit/sdk/job/test_proxy_integration.py | 303 +++++++++++++++++++ 7 files changed, 443 insertions(+), 6 deletions(-) create mode 100644 tests/unit/sdk/job/test_proxy_integration.py diff --git a/rock/env_vars.py b/rock/env_vars.py index 92e4884d7f..0d9f169395 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -50,6 +50,7 @@ # Model Service Config ROCK_MODEL_SERVICE_DATA_DIR: str ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE: bool | None = None + ROCK_JOB_PROXY_REPLAY_FILE: str # RuntimeEnv ROCK_RTENV_PYTHON_V31114_INSTALL_CMD: str @@ -104,6 +105,9 @@ "ROCK_CLI_DEFAULT_CONFIG_PATH", Path.home() / ".rock" / "config.ini" ), "ROCK_MODEL_SERVICE_DATA_DIR": lambda: os.getenv("ROCK_MODEL_SERVICE_DATA_DIR", "/data/logs"), + "ROCK_JOB_PROXY_REPLAY_FILE": lambda: os.getenv( + "ROCK_JOB_PROXY_REPLAY_FILE", "/data/logs/user-defined/rock-job-proxy-replay.jsonl" + ), "ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE": lambda: os.getenv("ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE", "false").lower() == "true", "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py index d0330edc68..6aad1688aa 100644 --- a/rock/sdk/envhub/config.py +++ b/rock/sdk/envhub/config.py @@ -6,7 +6,7 @@ from __future__ import annotations -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from rock.sdk.sandbox.config import SandboxConfig @@ -29,6 +29,49 @@ class OssMirrorConfig(BaseModel): oss_endpoint: str | None = None +class ProxyConfig(BaseModel): + """In-sandbox OpenAI request record/replay proxy config. + + Reuses ``rock.sdk.sandbox.model_service.ModelService`` (which uses + ``PythonRuntimeEnv`` to install a self-contained Python), so the + base image does not need Python preinstalled. + """ + + enabled: bool = False + + recording_file: str | None = None + """Recording mode: absolute path inside the sandbox where the proxy appends + the jsonl. When None (default), model-service picks its own path + (``$ROCK_MODEL_SERVICE_DATA_DIR/LLMTraj.jsonl``).""" + + replay_file: str | None = None + """Replay mode: local path. SDK uploads it to the sandbox and the proxy reads + the in-sandbox copy. Mutually exclusive with recording_file.""" + + host: str = "0.0.0.0" + """Proxy listen address. 0.0.0.0 lets nested docker containers reach it via + the docker0 gateway.""" + + port: int = 28080 + """Avoid common ports like 8080 to reduce collision risk with other services + in the sandbox.""" + + model_service_package: str = "rl-rock[model-service]" + """Package spec for installing the proxy. Must include the ``[model-service]`` + extra. Use PEP 508 syntax to pin a specific wheel, e.g. + ``"rl-rock[model-service] @ http://.../rl_rock-X.Y.Z-py3-none-any.whl"``. + """ + + @model_validator(mode="after") + def _record_replay_mutually_exclusive(self): + if self.recording_file and self.replay_file: + raise ValueError( + "ProxyConfig.recording_file and replay_file are mutually exclusive — " + "set one (recording mode) or the other (replay mode), not both." + ) + return self + + class EnvironmentConfig(SandboxConfig): """General environment config — sandbox base fields + environment-level fields.""" @@ -39,3 +82,6 @@ class EnvironmentConfig(SandboxConfig): ) env: dict[str, str] = Field(default_factory=dict) oss_mirror: OssMirrorConfig | None = None + proxy: ProxyConfig | None = None + """In-sandbox model-service proxy for OpenAI request record/replay. + None (default) means no proxy is started.""" diff --git a/rock/sdk/job/executor.py b/rock/sdk/job/executor.py index 1600c8d876..665a2426d8 100644 --- a/rock/sdk/job/executor.py +++ b/rock/sdk/job/executor.py @@ -95,11 +95,12 @@ async def _do_submit(self, trial: AbstractTrial) -> TrialClient: # G4: let trial backfill config from sandbox state before setup await trial.on_sandbox_ready(sandbox) + await trial.setup(sandbox) + session = f"rock-job-{config.job_name or 'default'}" env = self._build_session_env(config) await sandbox.create_session(CreateBashSessionRequest(session=session, env_enable=True, env=env)) - await trial.setup(sandbox) script_content = trial.build() prefix = self._job_tmp_prefix(config) diff --git a/rock/sdk/job/trial/abstract.py b/rock/sdk/job/trial/abstract.py index 2d9d641679..18ab057597 100644 --- a/rock/sdk/job/trial/abstract.py +++ b/rock/sdk/job/trial/abstract.py @@ -5,15 +5,22 @@ from __future__ import annotations +import shlex from abc import ABC, abstractmethod from pathlib import Path from typing import TYPE_CHECKING +from rock import env_vars +from rock.logger import init_logger +from rock.sdk.sandbox.model_service.base import ModelService, ModelServiceConfig + if TYPE_CHECKING: from rock.sdk.job.config import JobConfig from rock.sdk.job.result import TrialResult from rock.sdk.sandbox.client import Sandbox +logger = init_logger(__name__) + class AbstractTrial(ABC): """Trial base: three-phase interface (setup/build/collect). @@ -48,9 +55,85 @@ async def on_sandbox_ready(self, sandbox: Sandbox) -> None: self._config.experiment_id = sb_exp # If config already has experiment_id, it takes priority over sandbox's value. - @abstractmethod + def _build_proxy_start_cmd(self) -> str: + """Build the ``rock model-service start ...`` command line.""" + proxy = self._config.environment.proxy + env = self._config.environment.env + upstream = env["OPENAI_BASE_URL"] + parts = [ + "rock model-service start --type proxy", + f"--host {shlex.quote(proxy.host)}", + f"--port {proxy.port}", + f"--proxy-base-url {shlex.quote(upstream)}", + ] + if proxy.replay_file: + parts.append(f"--replay-file {shlex.quote(env_vars.ROCK_JOB_PROXY_REPLAY_FILE)}") + elif proxy.recording_file: + parts.append(f"--recording-file {shlex.quote(proxy.recording_file)}") + return " ".join(parts) + + async def _detect_and_rewrite_proxy_url(self, sandbox: Sandbox) -> None: + """Detect the outer sandbox eth0 IP and rewrite env['OPENAI_BASE_URL'].""" + proxy = self._config.environment.proxy + obs = await sandbox.arun("hostname -I 2>/dev/null | awk '{print $1}'") + host_ip = obs.output.strip() or "127.0.0.1" + proxy_url = f"http://{host_ip}:{proxy.port}/v1" + self._config.environment.env["OPENAI_BASE_URL"] = proxy_url + logger.info(f"Proxy ready at {proxy_url}") + + async def _setup_proxy(self, sandbox: Sandbox) -> None: + """Bring up the in-sandbox model-service proxy when enabled, otherwise no-op. + + Called at the start of ``HarborTrial.setup()`` and ``BashTrial.setup()`` so the + proxy is ready before any user/agent code runs. + + The OPENAI_BASE_URL existence check lives here (not in the + EnvironmentConfig validator) to keep job-layer concerns out of the generic + sandbox config. + """ + proxy = self._config.environment.proxy + if proxy is None or not proxy.enabled: + return + + if not self._config.environment.env.get("OPENAI_BASE_URL"): + raise ValueError( + "proxy.enabled=True but env['OPENAI_BASE_URL'] is not set. " + "Set environment.env.OPENAI_BASE_URL to the upstream OpenAI-compatible " + "base URL (e.g. 'https://api.openai.com/v1') so the proxy knows where " + "to forward." + ) + + if proxy.replay_file: + resp = await sandbox.upload_by_path( + file_path=proxy.replay_file, + target_path=env_vars.ROCK_JOB_PROXY_REPLAY_FILE, + ) + if not resp.success: + raise RuntimeError( + f"Failed to upload proxy replay file {proxy.replay_file} -> " + f"{env_vars.ROCK_JOB_PROXY_REPLAY_FILE}: {resp.message}" + ) + + ms_config = ModelServiceConfig( + enabled=True, + type="proxy", + install_cmd=f"pip install {shlex.quote(proxy.model_service_package)}", + start_cmd=self._build_proxy_start_cmd(), + ) + sandbox.model_service = ModelService(sandbox, ms_config) + await sandbox.model_service.install() + await sandbox.model_service.start() + + await self._detect_and_rewrite_proxy_url(sandbox) + async def setup(self, sandbox: Sandbox) -> None: - """Pre-execution: prepare sandbox environment (upload files, write configs).""" + """Pre-execution: start proxy (if enabled) and upload files. + + Subclasses should call ``await super().setup(sandbox)`` first, then add + their own setup logic. + """ + await self._setup_proxy(sandbox) + await self._upload_files(sandbox) @abstractmethod def build(self) -> str: diff --git a/rock/sdk/job/trial/bash.py b/rock/sdk/job/trial/bash.py index 915fe8d497..7bcfa3c677 100644 --- a/rock/sdk/job/trial/bash.py +++ b/rock/sdk/job/trial/bash.py @@ -130,7 +130,7 @@ def _render_wrapper(user_script: str, token: str | None = None) -> str: ) async def setup(self, sandbox: Sandbox) -> None: - await self._upload_files(sandbox) + await super().setup(sandbox) if self._config.script_path: self._config.script = Path(self._config.script_path).read_text() diff --git a/rock/sdk/job/trial/harbor.py b/rock/sdk/job/trial/harbor.py index 0522ab40d3..7d507a5159 100644 --- a/rock/sdk/job/trial/harbor.py +++ b/rock/sdk/job/trial/harbor.py @@ -51,7 +51,7 @@ class HarborTrial(AbstractTrial): _config: HarborJobConfig async def setup(self, sandbox) -> None: - await self._upload_files(sandbox) + await super().setup(sandbox) # Write Harbor YAML config to sandbox yaml_content = self._config.to_harbor_yaml() config_path = f"{USER_DEFINED_LOGS}/rock_job_{self._config.job_name}.yaml" diff --git a/tests/unit/sdk/job/test_proxy_integration.py b/tests/unit/sdk/job/test_proxy_integration.py new file mode 100644 index 0000000000..5d70eea551 --- /dev/null +++ b/tests/unit/sdk/job/test_proxy_integration.py @@ -0,0 +1,303 @@ +"""Unit tests for in-sandbox model-service proxy integration on Job layer. + +Covers: +- ProxyConfig mutex validator (recording_file vs replay_file) +- _build_proxy_start_cmd argument assembly (record / replay / default recording path) +- _setup_proxy behaviors (no-op / OPENAI_BASE_URL check / replay upload ordering) +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from rock import env_vars +from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.envhub.config import ProxyConfig +from rock.sdk.job.config import BashJobConfig +from rock.sdk.job.trial.bash import BashTrial + +# --------------------------------------------------------------------------- +# ProxyConfig validators +# --------------------------------------------------------------------------- + + +class TestProxyConfigValidators: + def test_record_replay_mutually_exclusive(self): + with pytest.raises(ValidationError, match="mutually exclusive"): + ProxyConfig(enabled=True, recording_file="a.jsonl", replay_file="b.jsonl") + + def test_only_recording_file_ok(self): + c = ProxyConfig(enabled=True, recording_file="a.jsonl") + assert c.recording_file == "a.jsonl" + assert c.replay_file is None + + def test_only_replay_file_ok(self): + c = ProxyConfig(enabled=True, replay_file="b.jsonl") + assert c.replay_file == "b.jsonl" + assert c.recording_file is None + + def test_both_unset_ok(self): + """Recording mode default: leaving both recording_file and replay_file + unset is valid (model-service uses its own default path).""" + c = ProxyConfig(enabled=True) + assert c.recording_file is None + assert c.replay_file is None + + def test_defaults(self): + c = ProxyConfig() + assert c.enabled is False + assert c.host == "0.0.0.0" + assert c.port == 28080 + assert c.model_service_package == "rl-rock[model-service]" + + +class TestEnvironmentConfigProxyField: + def test_proxy_field_default_none(self): + cfg = EnvironmentConfig() + assert cfg.proxy is None + + def test_proxy_field_accepts_proxy_config(self): + cfg = EnvironmentConfig(proxy=ProxyConfig(enabled=True)) + assert cfg.proxy is not None + assert cfg.proxy.enabled is True + + def test_proxy_enabled_without_openai_base_url_does_not_raise(self): + """EnvironmentConfig stays pure: even when proxy.enabled=True with no + OPENAI_BASE_URL in env, no error is raised here. The check belongs to + AbstractTrial._setup_proxy.""" + EnvironmentConfig(proxy=ProxyConfig(enabled=True)) + + +def _make_trial(*, proxy=None, env=None): + """Create a BashTrial wired with the given proxy and env for testing helpers.""" + environment = EnvironmentConfig(proxy=proxy, env=env or {}) + cfg = BashJobConfig(script="echo", environment=environment) + return BashTrial(cfg) + + +class TestBuildProxyStartCmd: + def test_record_mode_with_explicit_path(self): + trial = _make_trial( + proxy=ProxyConfig(enabled=True, recording_file="/data/logs/x.jsonl", port=28080), + env={"OPENAI_BASE_URL": "https://upstream.example.com/v1"}, + ) + cmd = trial._build_proxy_start_cmd() + assert "rock model-service start --type proxy" in cmd + assert "--host 0.0.0.0" in cmd + assert "--port 28080" in cmd + assert "--proxy-base-url https://upstream.example.com/v1" in cmd + assert "--recording-file /data/logs/x.jsonl" in cmd + assert "--replay-file" not in cmd + + def test_replay_mode_uses_sandbox_replay_path(self): + trial = _make_trial( + proxy=ProxyConfig(enabled=True, replay_file="/local/r.jsonl"), + env={"OPENAI_BASE_URL": "https://upstream.example.com/v1"}, + ) + cmd = trial._build_proxy_start_cmd() + assert "--replay-file" in cmd + assert env_vars.ROCK_JOB_PROXY_REPLAY_FILE in cmd + assert "--recording-file" not in cmd + + def test_record_mode_omits_recording_flag_when_unset(self): + trial = _make_trial( + proxy=ProxyConfig(enabled=True), + env={"OPENAI_BASE_URL": "https://x/v1"}, + ) + cmd = trial._build_proxy_start_cmd() + assert "--type proxy" in cmd + assert "--recording-file" not in cmd + assert "--replay-file" not in cmd + + def test_sandbox_replay_file_constant(self): + assert env_vars.ROCK_JOB_PROXY_REPLAY_FILE == "/data/logs/user-defined/rock-job-proxy-replay.jsonl" + + +class TestSetupProxy: + """Covers the four behavior branches of _setup_proxy: + - disabled / proxy is None -> no-op + - enabled but env missing OPENAI_BASE_URL -> ValueError (before install/start) + - replay mode -> upload_by_path runs before install/start + - record mode -> upload_by_path is not called + """ + + async def test_noop_when_proxy_is_none(self): + cfg = BashJobConfig(script="echo") + trial = BashTrial(cfg) + sandbox = AsyncMock() + await trial._setup_proxy(sandbox) + sandbox.upload_by_path.assert_not_called() + + async def test_noop_when_disabled(self): + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig(proxy=ProxyConfig(enabled=False)), + ) + trial = BashTrial(cfg) + sandbox = AsyncMock() + await trial._setup_proxy(sandbox) + sandbox.upload_by_path.assert_not_called() + + async def test_raises_when_openai_base_url_missing(self, monkeypatch): + """proxy.enabled=True but env missing OPENAI_BASE_URL -> ValueError, + and it must happen before ModelService.install/start is called.""" + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig(proxy=ProxyConfig(enabled=True)), + ) + trial = BashTrial(cfg) + sandbox = AsyncMock() + + # Patch ModelService to ensure it is never constructed. + ms_class = MagicMock() + monkeypatch.setattr("rock.sdk.job.trial.abstract.ModelService", ms_class) + + with pytest.raises(ValueError, match="OPENAI_BASE_URL"): + await trial._setup_proxy(sandbox) + + ms_class.assert_not_called() + sandbox.upload_by_path.assert_not_called() + + async def test_replay_uploads_before_start(self, monkeypatch): + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig( + env={"OPENAI_BASE_URL": "https://upstream/v1"}, + proxy=ProxyConfig(enabled=True, replay_file="/local/r.jsonl"), + ), + ) + trial = BashTrial(cfg) + + sandbox = AsyncMock() + sandbox.upload_by_path.return_value = MagicMock(success=True, message="") + arun_obs = MagicMock() + arun_obs.output = "10.0.0.1" + sandbox.arun = AsyncMock(return_value=arun_obs) + + fake_ms_instance = AsyncMock() + ms_class = MagicMock(return_value=fake_ms_instance) + monkeypatch.setattr("rock.sdk.job.trial.abstract.ModelService", ms_class) + + await trial._setup_proxy(sandbox) + + # Upload must happen. + sandbox.upload_by_path.assert_awaited_once_with( + file_path="/local/r.jsonl", + target_path=env_vars.ROCK_JOB_PROXY_REPLAY_FILE, + ) + # ModelService is installed and started. + fake_ms_instance.install.assert_awaited_once() + fake_ms_instance.start.assert_awaited_once() + + # Ordering: upload must precede install. + all_calls = sandbox.method_calls + fake_ms_instance.method_calls + upload_idx = next(i for i, c in enumerate(all_calls) if c[0] == "upload_by_path") + install_idx = next(i for i, c in enumerate(all_calls) if c[0] == "install") + assert upload_idx < install_idx, ( + f"replay file must be uploaded before ModelService.install " + f"(upload at idx {upload_idx}, install at idx {install_idx})" + ) + + async def test_replay_upload_failure_raises(self, monkeypatch): + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig( + env={"OPENAI_BASE_URL": "https://upstream/v1"}, + proxy=ProxyConfig(enabled=True, replay_file="/local/r.jsonl"), + ), + ) + trial = BashTrial(cfg) + sandbox = AsyncMock() + sandbox.upload_by_path.return_value = MagicMock(success=False, message="boom") + + fake_ms_instance = AsyncMock() + ms_class = MagicMock(return_value=fake_ms_instance) + monkeypatch.setattr("rock.sdk.job.trial.abstract.ModelService", ms_class) + + with pytest.raises(RuntimeError, match="boom"): + await trial._setup_proxy(sandbox) + fake_ms_instance.install.assert_not_called() + + async def test_record_mode_no_upload(self, monkeypatch): + cfg = BashJobConfig( + script="echo", + environment=EnvironmentConfig( + env={"OPENAI_BASE_URL": "https://upstream/v1"}, + proxy=ProxyConfig(enabled=True, recording_file="/data/logs/x.jsonl"), + ), + ) + trial = BashTrial(cfg) + + sandbox = AsyncMock() + arun_obs = MagicMock() + arun_obs.output = "10.0.0.1" + sandbox.arun = AsyncMock(return_value=arun_obs) + fake_ms_instance = AsyncMock() + ms_class = MagicMock(return_value=fake_ms_instance) + monkeypatch.setattr("rock.sdk.job.trial.abstract.ModelService", ms_class) + + await trial._setup_proxy(sandbox) + + sandbox.upload_by_path.assert_not_called() + fake_ms_instance.install.assert_awaited_once() + fake_ms_instance.start.assert_awaited_once() + # env['OPENAI_BASE_URL'] must be rewritten to the proxy URL. + assert cfg.environment.env["OPENAI_BASE_URL"] == "http://10.0.0.1:28080/v1" + + +class TestHarborTrialSetupCallsProxy: + async def test_setup_calls_setup_proxy_first(self, monkeypatch, tmp_path): + """HarborTrial.setup() must await self._setup_proxy(sandbox) first.""" + from rock.sdk.bench.models.job.config import HarborJobConfig + from rock.sdk.job.trial.harbor import HarborTrial + + cfg = HarborJobConfig(experiment_id="exp-test") + trial = HarborTrial(cfg) + + sandbox = AsyncMock() + sandbox.write_file_by_path = AsyncMock() + + order: list[str] = [] + + async def track_setup_proxy(sb): + order.append("setup_proxy") + + async def track_upload(sb): + order.append("upload_files") + + monkeypatch.setattr(trial, "_setup_proxy", track_setup_proxy) + monkeypatch.setattr(trial, "_upload_files", track_upload) + + await trial.setup(sandbox) + assert order[:2] == [ + "setup_proxy", + "upload_files", + ], f"setup_proxy must be called before _upload_files, actual order: {order}" + + +class TestBashTrialSetupCallsProxy: + async def test_setup_calls_setup_proxy_first(self, monkeypatch): + """BashTrial.setup() must await self._setup_proxy(sandbox) first.""" + cfg = BashJobConfig(script="echo hi") + trial = BashTrial(cfg) + + sandbox = AsyncMock() + order: list[str] = [] + + async def track_setup_proxy(sb): + order.append("setup_proxy") + + async def track_upload(sb): + order.append("upload_files") + + monkeypatch.setattr(trial, "_setup_proxy", track_setup_proxy) + monkeypatch.setattr(trial, "_upload_files", track_upload) + + await trial.setup(sandbox) + assert order[:2] == [ + "setup_proxy", + "upload_files", + ], f"setup_proxy must be called before _upload_files, actual order: {order}" From 55dbc2119c47862cc0a36392aa6a4a443e05a81e Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 14 May 2026 14:12:17 +0800 Subject: [PATCH 091/226] fix(sdk): mkdir target parent dir before wget in OSS upload path (#940) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `wget -O` does not create missing parent directories — when target_path's parent doesn't exist, wget fails with "No such file or directory" and the upload silently fails (NOHUP mode swallows the exit code). Mirror what the multipart upload path (rocklet `/upload`) already does server-side: `mkdir -p` the parent dir before downloading. fixes #939 Co-authored-by: Claude Opus 4.7 --- rock/sdk/sandbox/client.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 04b2fbf082..e27f9b60c1 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -3,6 +3,7 @@ import math import mimetypes import os +import shlex import time import uuid import warnings @@ -842,6 +843,15 @@ async def _upload_via_oss(self, file_path: str | Path, target_path: str): oss2.resumable_upload(self._oss_bucket, tmp_obj_name, file_path) url = self._oss_bucket.sign_url("GET", tmp_obj_name, 600, slash_safe=True) try: + # wget -O does not create missing parent dirs; mkdir -p first to + # match the multipart upload path (rocklet /upload mkdirs server-side). + parent_dir = str(Path(target_path).parent) + await self.arun( + cmd=f"mkdir -p {shlex.quote(parent_dir)}", + wait_timeout=10, + mode=RunMode.NORMAL, + ) + download_cmd = f"wget -c -O {target_path} '{url}'" await self.arun(cmd=download_cmd, wait_timeout=600, mode=RunMode.NOHUP) check_file_session = f"bash-{timestamp}" From bb8d0f9b9650a39fcb77ff3ccff36da5311841cb Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Sat, 9 May 2026 08:02:43 +0000 Subject: [PATCH 092/226] fix(metrics): pass rock_config to SandboxTable/SandboxMetaStore so MetricsMonitor uses the correct metrics endpoint Signed-off-by: Jiachen Zhang (cherry picked from commit 8bec485d944e2c1e2aa9ef9baf75a68e9ccb21c7) Signed-off-by: Jiachen Zhang --- rock/admin/core/sandbox_table.py | 10 ++++++++-- rock/admin/main.py | 4 ++-- rock/sandbox/sandbox_meta_store.py | 9 ++++++++- tests/unit/conftest.py | 8 ++++---- 4 files changed, 22 insertions(+), 9 deletions(-) diff --git a/rock/admin/core/sandbox_table.py b/rock/admin/core/sandbox_table.py index fe35b55a8b..fb1d0c1949 100644 --- a/rock/admin/core/sandbox_table.py +++ b/rock/admin/core/sandbox_table.py @@ -11,6 +11,7 @@ from rock.admin.core.schema import SandboxRecord from rock.admin.metrics.decorator import monitor_metastore_operation from rock.admin.metrics.monitor import MetricsMonitor +from rock.config import RockConfig from rock.logger import init_logger if TYPE_CHECKING: @@ -36,9 +37,14 @@ class SandboxTable: including ``spec`` and ``status``. """ - def __init__(self, db_provider: DatabaseProvider) -> None: + def __init__(self, db_provider: DatabaseProvider, rock_config: RockConfig | None = None) -> None: self._db = db_provider - self.metrics_monitor = MetricsMonitor.create(metric_prefix="meta_store.db") + self.metrics_monitor = MetricsMonitor.create( + export_interval_millis=20_000, + metrics_endpoint=rock_config.runtime.metrics_endpoint if rock_config else "", + user_defined_tags=rock_config.runtime.user_defined_tags if rock_config else {}, + metric_prefix="meta_store.db", + ) @monitor_metastore_operation async def create( diff --git a/rock/admin/main.py b/rock/admin/main.py index 24824d7e7d..b1eeef12c0 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -86,8 +86,8 @@ async def lifespan(app: FastAPI): await db_provider.init() if not rock_config.database.url: await db_provider.create_tables() - sandbox_table = SandboxTable(db_provider) - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=sandbox_table) + sandbox_table = SandboxTable(db_provider, rock_config=rock_config) + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=sandbox_table, rock_config=rock_config) # init scheduler thread scheduler_thread = None diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index 63eeb5512a..b1dec0246b 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -16,6 +16,7 @@ from rock.admin.core.sandbox_table import SandboxTable from rock.admin.metrics.decorator import monitor_metastore_operation from rock.admin.metrics.monitor import MetricsMonitor +from rock.config import RockConfig if TYPE_CHECKING: from rock.deployments.config import DockerDeploymentConfig @@ -38,10 +39,16 @@ def __init__( self, redis_provider: RedisProvider, sandbox_table: SandboxTable, + rock_config: RockConfig | None = None, ) -> None: self._redis: RedisProvider = redis_provider self._db: SandboxTable = sandbox_table - self.metrics_monitor = MetricsMonitor.create(metric_prefix="meta_store") + self.metrics_monitor = MetricsMonitor.create( + export_interval_millis=20_000, + metrics_endpoint=rock_config.runtime.metrics_endpoint if rock_config else "", + user_defined_tags=rock_config.runtime.user_defined_tags if rock_config else {}, + metric_prefix="meta_store", + ) # ------------------------------------------------------------------ # Public API diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 55f697d2ea..49837907b6 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -91,8 +91,8 @@ async def db_provider(): @pytest.fixture -async def _memory_sandbox_table(db_provider): - return SandboxTable(db_provider) +async def _memory_sandbox_table(db_provider, rock_config): + return SandboxTable(db_provider, rock_config=rock_config) @pytest.fixture @@ -104,7 +104,7 @@ async def sandbox_manager( ray_operator, _memory_sandbox_table: SandboxTable, ): - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config) sandbox_manager = SandboxManager( rock_config, meta_store=meta_store, @@ -120,7 +120,7 @@ async def sandbox_manager( async def sandbox_proxy_service( rock_config: RockConfig, redis_provider: RedisProvider, _memory_sandbox_table: SandboxTable ): - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table) + meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config) sandbox_proxy_service = SandboxProxyService(rock_config, meta_store=meta_store) return sandbox_proxy_service From 3deda28b5b37e02a24475120bddc6cf095b0a9cf Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 13 May 2026 15:55:55 +0000 Subject: [PATCH 093/226] chore(metrics): remove the sandbox_id attr from metastore monitor attributes - Drop sandbox_id and redundant method labels from monitor_metastore_operation attrs - Skip sandbox_id merge in _update_sandbox_id_from_result when not present - Update metastore/sandbox metrics unit tests and regressions for label shape Signed-off-by: Jiachen Zhang --- rock/admin/metrics/decorator.py | 5 +- .../admin/core/test_sandbox_table_metrics.py | 6 +- .../admin/metrics/test_metastore_decorator.py | 60 +++++++------------ .../test_sandbox_meta_store_metrics.py | 6 +- 4 files changed, 32 insertions(+), 45 deletions(-) diff --git a/rock/admin/metrics/decorator.py b/rock/admin/metrics/decorator.py index 3ae58bfb9b..39bdceb932 100644 --- a/rock/admin/metrics/decorator.py +++ b/rock/admin/metrics/decorator.py @@ -74,6 +74,8 @@ def _build_attributes(op_name: str, sandbox_id: str, f, user_id: str, experiment def _update_sandbox_id_from_result(result, attributes: dict): """Update sandbox_id from result if available""" + if "sandbox_id" not in attributes: + return attributes if hasattr(result, "sandbox_id"): result_sandbox_id = result.sandbox_id if result_sandbox_id != attributes.get("sandbox_id"): @@ -242,8 +244,7 @@ async def wrapper(self, *args, **kwargs): return await f(self, *args, **kwargs) prefix = metrics_monitor.metric_prefix - sandbox_id = _extract_sandbox_id(args, kwargs) - attributes: dict[str, str] = {"operation": f.__name__, "method": f.__name__, "sandbox_id": sandbox_id} + attributes: dict[str, str] = {"operation": f.__name__} start_time = time.perf_counter() diff --git a/tests/unit/admin/core/test_sandbox_table_metrics.py b/tests/unit/admin/core/test_sandbox_table_metrics.py index f11e634579..8f9f2c1b86 100644 --- a/tests/unit/admin/core/test_sandbox_table_metrics.py +++ b/tests/unit/admin/core/test_sandbox_table_metrics.py @@ -41,7 +41,7 @@ class TestSandboxTableMetrics: async def test_create_records_db_metrics(self, table, mock_monitor): await table.create(SANDBOX_ID, SANDBOX_INFO) - attrs = {"operation": "create", "method": "create", "sandbox_id": SANDBOX_ID} + attrs = {"operation": "create"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) rt_call = mock_monitor.record_gauge_by_name.call_args @@ -54,7 +54,7 @@ async def test_get_records_db_metrics(self, table, mock_monitor): await table.get(SANDBOX_ID) - attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID} + attrs = {"operation": "get"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) mock_monitor.record_gauge_by_name.assert_called_once() @@ -72,6 +72,6 @@ async def test_list_by_records_db_metrics(self, table, mock_monitor): results = await table.list_by("user_id", "user-1") assert len(results) == 1 - attrs = {"operation": "list_by", "method": "list_by", "sandbox_id": "user_id"} + attrs = {"operation": "list_by"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) diff --git a/tests/unit/admin/metrics/test_metastore_decorator.py b/tests/unit/admin/metrics/test_metastore_decorator.py index 64b4fa8193..724339377c 100644 --- a/tests/unit/admin/metrics/test_metastore_decorator.py +++ b/tests/unit/admin/metrics/test_metastore_decorator.py @@ -8,7 +8,7 @@ from rock.admin.metrics.monitor import MetricsMonitor -def _make_monitor(prefix="meta_store"): +def _make_monitor(prefix: str = "fake_store"): monitor = Mock(spec=MetricsMonitor) monitor._should_skip.return_value = False monitor.metric_prefix = prefix @@ -46,12 +46,12 @@ async def test_records_success_metrics(self): result = await store.do_something() assert result == "ok" - attrs = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} - monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) - monitor.record_counter_by_name.assert_any_call("meta_store.total", 1, attrs) + attrs = {"operation": "do_something"} + monitor.record_counter_by_name.assert_any_call("fake_store.success", 1, attrs) + monitor.record_counter_by_name.assert_any_call("fake_store.total", 1, attrs) monitor.record_gauge_by_name.assert_called_once() call_args = monitor.record_gauge_by_name.call_args - assert call_args[0][0] == "meta_store.rt" + assert call_args[0][0] == "fake_store.rt" assert call_args[0][2] == attrs async def test_records_failure_metrics(self): @@ -61,9 +61,9 @@ async def test_records_failure_metrics(self): with pytest.raises(ValueError, match="boom"): await store.do_fail() - error_attrs = {"operation": "do_fail", "method": "do_fail", "sandbox_id": "unknown", "error_type": "ValueError"} - monitor.record_counter_by_name.assert_any_call("meta_store.failure", 1, error_attrs) - monitor.record_counter_by_name.assert_any_call("meta_store.total", 1, error_attrs) + error_attrs = {"operation": "do_fail", "error_type": "ValueError"} + monitor.record_counter_by_name.assert_any_call("fake_store.failure", 1, error_attrs) + monitor.record_counter_by_name.assert_any_call("fake_store.total", 1, error_attrs) monitor.record_gauge_by_name.assert_called_once() async def test_skips_when_no_monitor(self): @@ -90,7 +90,7 @@ async def test_uses_monitor_prefix(self): await store.do_something() - attrs_db = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} + attrs_db = {"operation": "do_something"} monitor.record_counter_by_name.assert_any_call("meta_store.db.success", 1, attrs_db) monitor.record_counter_by_name.assert_any_call("meta_store.db.total", 1, attrs_db) call_args = monitor.record_gauge_by_name.call_args @@ -101,38 +101,24 @@ async def test_sandbox_id_in_attributes_positional(self): monitor = _make_monitor() store = FakeStore(metrics_monitor=monitor) - result = await store.get("sbx-123") + await store.get("sbx-123") - assert result == "sbx-123" - attrs = {"operation": "get", "method": "get", "sandbox_id": "sbx-123"} - monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + for call in monitor.record_counter_by_name.call_args_list: + recorded_attrs = call[0][2] + assert "sandbox_id" not in recorded_attrs + assert "method" not in recorded_attrs + for call in monitor.record_gauge_by_name.call_args_list: + recorded_attrs = call[0][2] + assert "sandbox_id" not in recorded_attrs + assert "method" not in recorded_attrs - async def test_sandbox_id_unknown_when_keyword_only(self): - """sandbox_id passed as keyword is not extracted (consistent with monitor_sandbox_operation).""" - monitor = _make_monitor() - store = FakeStore(metrics_monitor=monitor) - - await store.get(sandbox_id="sbx-kw") - - attrs = {"operation": "get", "method": "get", "sandbox_id": "unknown"} - monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) - - async def test_sandbox_id_fallback_from_first_arg(self): - """Without sandbox_id param, _extract_sandbox_id falls back to args[0].""" + async def test_attributes_do_not_contain_method_field(self): + """Regression: the duplicate `method` label was removed (it equalled `operation`).""" monitor = _make_monitor() store = FakeStore(metrics_monitor=monitor) await store.list_by("state", "running") - attrs = {"operation": "list_by", "method": "list_by", "sandbox_id": "state"} - monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) - - async def test_sandbox_id_unknown_when_no_args(self): - """No args at all → sandbox_id defaults to 'unknown'.""" - monitor = _make_monitor() - store = FakeStore(metrics_monitor=monitor) - - await store.do_something() - - attrs = {"operation": "do_something", "method": "do_something", "sandbox_id": "unknown"} - monitor.record_counter_by_name.assert_any_call("meta_store.success", 1, attrs) + for call in monitor.record_counter_by_name.call_args_list: + recorded_attrs = call[0][2] + assert recorded_attrs == {"operation": "list_by"} diff --git a/tests/unit/sandbox/test_sandbox_meta_store_metrics.py b/tests/unit/sandbox/test_sandbox_meta_store_metrics.py index 1906b23281..8b7103bc7a 100644 --- a/tests/unit/sandbox/test_sandbox_meta_store_metrics.py +++ b/tests/unit/sandbox/test_sandbox_meta_store_metrics.py @@ -62,7 +62,7 @@ class TestMetaStoreMetrics: async def test_create_records_store_metrics(self, store, mock_monitor): await store.create(SANDBOX_ID, SANDBOX_INFO) - attrs = {"operation": "create", "method": "create", "sandbox_id": SANDBOX_ID} + attrs = {"operation": "create"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, attrs) assert mock_monitor.record_gauge_by_name.called @@ -76,7 +76,7 @@ async def test_get_records_store_metrics(self, store, redis, mock_monitor): await store.get(SANDBOX_ID) - attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID} + attrs = {"operation": "get"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.success", 1, attrs) mock_monitor.record_gauge_by_name.assert_called_once() @@ -86,6 +86,6 @@ async def test_failure_records_error_type(self, store, redis, mock_monitor): with pytest.raises(ConnectionError): await store.get(SANDBOX_ID) - error_attrs = {"operation": "get", "method": "get", "sandbox_id": SANDBOX_ID, "error_type": "ConnectionError"} + error_attrs = {"operation": "get", "error_type": "ConnectionError"} mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.failure", 1, error_attrs) mock_monitor.record_counter_by_name.assert_any_call(f"{PREFIX}.total", 1, error_attrs) From 97598490cac014673466f9d87e5f27df68323150 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Thu, 14 May 2026 06:23:50 +0000 Subject: [PATCH 094/226] metrics(monitor): scope meter to local MeterProvider Obtain the meter via self.meter_provider.get_meter instead of registering this MeterProvider as the global default and calling metrics.get_meter. That keeps admin metrics isolated and avoids overwriting process-wide OpenTelemetry provider state. Signed-off-by: Jiachen Zhang --- rock/admin/metrics/monitor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/rock/admin/metrics/monitor.py b/rock/admin/metrics/monitor.py index 77ebcce719..9cb433141f 100644 --- a/rock/admin/metrics/monitor.py +++ b/rock/admin/metrics/monitor.py @@ -1,6 +1,5 @@ from collections import Counter as CollectionsCounter -from opentelemetry import metrics from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.metrics import Counter, _Gauge from opentelemetry.sdk.metrics import MeterProvider @@ -154,8 +153,7 @@ def _init_telemetry(self, export_interval_millis: int): export_interval_millis=export_interval_millis, ) self.meter_provider = MeterProvider(metric_readers=[self.metric_reader]) - metrics.set_meter_provider(self.meter_provider) - self.meter = metrics.get_meter(MetricsConstants.METRICS_METER_NAME) + self.meter = self.meter_provider.get_meter(MetricsConstants.METRICS_METER_NAME) logger.info("init telemetry success") def create_counter(self, name: str, description: str, unit: str = "1") -> Counter: From 3e3265a1b10e712716927ebd78d64475d8488860 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 13 May 2026 15:03:28 +0000 Subject: [PATCH 095/226] feat(metrics): log OTLP export data_points count and duration Add _wrap_exporter_with_logging() to patch the OTLP exporter at startup. Each export call now logs the total data_points count across all metrics and the export duration in ms, making it straightforward to detect cardinality growth or export latency regressions in production. --- rock/admin/metrics/monitor.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/rock/admin/metrics/monitor.py b/rock/admin/metrics/monitor.py index 9cb433141f..34ceddb8a1 100644 --- a/rock/admin/metrics/monitor.py +++ b/rock/admin/metrics/monitor.py @@ -1,3 +1,4 @@ +import time from collections import Counter as CollectionsCounter from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter @@ -148,6 +149,7 @@ def _init_telemetry(self, export_interval_millis: int): self.metric_reader = InMemoryMetricReader() else: self.otlp_exporter = OTLPMetricExporter(endpoint=self.endpoint) + self._wrap_exporter_with_logging() self.metric_reader = PeriodicExportingMetricReader( self.otlp_exporter, export_interval_millis=export_interval_millis, @@ -156,6 +158,39 @@ def _init_telemetry(self, export_interval_millis: int): self.meter = self.meter_provider.get_meter(MetricsConstants.METRICS_METER_NAME) logger.info("init telemetry success") + def _wrap_exporter_with_logging(self): + """Patch otlp_exporter.export to log data point count and duration. + + Runs on the PeriodicExportingMetricReader's daemon thread, off the asyncio loop. + """ + if not hasattr(self, "otlp_exporter") or not hasattr(self.otlp_exporter, "export"): + return + original_export = self.otlp_exporter.export + endpoint = self.endpoint + + def export_with_logging(metrics_data, *args, **kwargs): + n_points = sum( + len(metric.data.data_points) + for rm in metrics_data.resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + if hasattr(metric.data, "data_points") + ) + t0 = time.perf_counter() + result = original_export(metrics_data, *args, **kwargs) + elapsed_ms = (time.perf_counter() - t0) * 1000.0 + logger.info( + "OTLP export metric_prefix=%s endpoint=%s data_points=%d duration_ms=%.1f result=%s", + self.metric_prefix, + endpoint, + n_points, + elapsed_ms, + result, + ) + return result + + self.otlp_exporter.export = export_with_logging + def create_counter(self, name: str, description: str, unit: str = "1") -> Counter: """Create counter""" if self._should_skip(): From e65b56aa4feaa9489ad3452846f218cd6c54b1e4 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 15 May 2026 10:16:55 +0800 Subject: [PATCH 096/226] feat(docker): cleanup XFS project quota on container stop (#941) Track log-dir XFS project ID and mount point after successful quota setup, and clear them via xfs_quota on stop. Extract XFS_PRJID_MIN/RANGE constants using the full uint32 space above Docker-reserved low IDs. --- rock/deployments/docker.py | 40 ++++++++++++++++++- .../test_docker_deployment_disk_limit.py | 3 +- 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 48914f5318..99b0ade1f5 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -43,6 +43,8 @@ __all__ = ["DockerDeployment", "DockerDeploymentConfig"] CHECK_CLEAR_INTERVAL_SECONDS = 300 +XFS_PRJID_MIN = (1 << 31) # low project IDs are reserved for Docker +XFS_PRJID_RANGE = (1 << 32) - XFS_PRJID_MIN # use remaining 32-bit space: [XFS_PRJID_MIN, 2^32) logger = init_logger(__name__) @@ -87,6 +89,8 @@ def __init__( raise Exception(f"Invalid ROCK_WORKER_ENV_TYPE: {env_vars.ROCK_WORKER_ENV_TYPE}") self.sandbox_validator: DockerSandboxValidator | None = DockerSandboxValidator() + self.log_dir_xfs_prjid: int | None = None + self.log_dir_xfs_mountpoint: str | None = None def add_hook(self, hook: DeploymentHook): self._hooks.add_hook(hook) @@ -208,6 +212,37 @@ def _prepare_kata_disk(self) -> None: os.remove(disk_path) raise + def _cleanup_log_dir_xfs_quota(self) -> None: + """Remove XFS project quota for the sandbox log directory on exit.""" + if self.log_dir_xfs_prjid is None or self.log_dir_xfs_mountpoint is None: + return + if not self._container_name: + return + + log_file_path = f"{env_vars.ROCK_LOGGING_PATH}/{self._container_name}" + project_id = self.log_dir_xfs_prjid + mount_point = self.log_dir_xfs_mountpoint + try: + clear_limit_cmd = f"limit -p bhard=0 bsoft=0 {project_id}" + clear_project_cmd = f"project -C -p {shlex.quote(log_file_path)} {project_id}" + for cmd in (clear_limit_cmd, clear_project_cmd): + result = subprocess.run( + ["xfs_quota", "-x", "-c", cmd, mount_point], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode != 0: + logger.warning( + f"xfs_quota cleanup failed for {log_file_path!r} cmd={cmd!r}: {result.stderr.strip() or result.stdout.strip()}" + ) + logger.info(f"Cleaned up XFS project quota (prjid={project_id}) for {log_file_path!r}") + except Exception as e: + logger.warning(f"Failed to cleanup XFS project quota for {log_file_path!r}: {e}") + finally: + self.log_dir_xfs_prjid = None + self.log_dir_xfs_mountpoint = None + def _cleanup_kata_disk(self) -> None: """Remove the kata disk image file from the host. @@ -389,7 +424,7 @@ def _try_set_log_dir_quota(self, log_file_path: str) -> None: return # Derive a deterministic project id from container name; reserve low ids. - project_id = (int(hashlib.sha1(self.container_name.encode("utf-8")).hexdigest()[:8], 16) % 900000) + 100000 + project_id = (int(hashlib.sha1(self.container_name.encode("utf-8")).hexdigest()[:8], 16) % XFS_PRJID_RANGE) + XFS_PRJID_MIN try: findmnt_result = subprocess.run( ["findmnt", "-T", log_file_path, "-o", "TARGET", "--noheadings"], @@ -422,6 +457,8 @@ def _try_set_log_dir_quota(self, log_file_path: str) -> None: ) self._effective_disk_limit_log = None return + self.log_dir_xfs_prjid = project_id + self.log_dir_xfs_mountpoint = mount_point logger.info(f"Set XFS project quota {self._effective_disk_limit_log} for log path {log_file_path!r}") except Exception as e: logger.warning(f"Failed to set XFS project quota for {log_file_path!r}: {e}") @@ -624,6 +661,7 @@ def _stop(self): self._container_process = None self._cleanup_kata_disk() + self._cleanup_log_dir_xfs_quota() self._container_name = None if self._config and self._config.remove_images and DockerUtil.is_image_available(self._config.image): diff --git a/tests/unit/deployments/test_docker_deployment_disk_limit.py b/tests/unit/deployments/test_docker_deployment_disk_limit.py index a94f3fd164..28ed165db0 100644 --- a/tests/unit/deployments/test_docker_deployment_disk_limit.py +++ b/tests/unit/deployments/test_docker_deployment_disk_limit.py @@ -241,5 +241,6 @@ def test_try_set_log_dir_quota_independent_of_docker_driver(self, _mock_prjquota mock_sub.run.return_value = ok deployment._try_set_log_dir_quota("/var/log/rock/test-container") - # xfs_quota succeeded → effective value preserved + # xfs_quota succeeded → effective value preserved and prjid recorded assert deployment.effective_disk_limit_log == "5g" + assert deployment.log_dir_xfs_prjid is not None From d7421e23c10544e67e2bb47d4a7192ded16feb61 Mon Sep 17 00:00:00 2001 From: dengsheng Date: Mon, 27 Apr 2026 11:07:52 +0000 Subject: [PATCH 097/226] fix: handle exception caused by ray.init during ray reconnecting background job --- rock/admin/core/ray_service.py | 79 ++++++++++++--- rock/config.py | 2 + tests/unit/admin/core/test_ray_service.py | 117 ++++++++++++++++++++++ 3 files changed, 183 insertions(+), 15 deletions(-) diff --git a/rock/admin/core/ray_service.py b/rock/admin/core/ray_service.py index c53f371b25..f6f5ec61ec 100644 --- a/rock/admin/core/ray_service.py +++ b/rock/admin/core/ray_service.py @@ -1,12 +1,23 @@ +# ruff: noqa: E402 -- os.environ assignment below MUST run before `import ray`. import asyncio +import os import time from concurrent.futures import ThreadPoolExecutor +# Disable Ray's auto-init hook before importing ray. When ray is shut down +# (e.g. after a failed periodic reconnect), any unguarded ray API call such as +# ``Actor.options(...).remote(...)`` or ``actor.method.remote()`` would +# otherwise trigger ``auto_init_ray()`` and silently spawn a local cluster on +# this host. ``enable_auto_connect`` in ``ray._private.auto_init_hook`` is +# evaluated at import time, so the env var must be set BEFORE ``import ray``. +os.environ["RAY_ENABLE_AUTO_CONNECT"] = "0" + import ray from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.interval import IntervalTrigger from rock import InternalServerRockError +from rock._codes import codes from rock.config import RayConfig from rock.logger import init_logger from rock.utils.rwlock import AsyncRWLock @@ -69,25 +80,61 @@ async def _ray_reconnect_with_policy(self): async def _reconnect_ray(self): try: async with self._ray_rwlock.write_lock(timeout=self._config.ray_reconnect_wait_timeout_seconds): - start_time = time.time() - logger.info(f"current time {start_time}, Reconnect ray cluster") - ray.shutdown() - ray.init( - address=self._config.address, - runtime_env=self._config.runtime_env, - namespace=self._config.namespace, - resources=self._config.resources, - _temp_dir=self._config.temp_dir, - ) - self._ray_request_count = 0 - end_time = time.time() - self._ray_establish_time = end_time - logger.info( - f"current time {end_time}, Reconnect ray cluster successfully, duration {end_time - start_time}s" + max_attempts = max(1, self._config.ray_reconnect_max_attempts) + backoff = self._config.ray_reconnect_retry_backoff_seconds + last_exc: Exception | None = None + for attempt in range(1, max_attempts + 1): + start_time = time.time() + logger.info(f"current time {start_time}, Reconnect ray cluster (attempt {attempt}/{max_attempts})") + try: + ray.shutdown() + ray.init( + address=self._config.address, + runtime_env=self._config.runtime_env, + namespace=self._config.namespace, + resources=self._config.resources, + _temp_dir=self._config.temp_dir, + ) + except Exception as e: + last_exc = e + logger.warning( + f"Reconnect ray cluster attempt {attempt}/{max_attempts} failed: {e}", exc_info=e + ) + if attempt < max_attempts and backoff > 0: + await asyncio.sleep(backoff) + continue + self._ray_request_count = 0 + end_time = time.time() + self._ray_establish_time = end_time + logger.info( + f"current time {end_time}, Reconnect ray cluster successfully, " + f"duration {end_time - start_time}s" + ) + return + logger.critical( + f"Ray reconnect failed after {max_attempts} attempts; ray cluster is in shutdown state. " + f"Last error: {last_exc}", + exc_info=last_exc, ) except InternalServerRockError as e: logger.warning("Reconnect ray cluster timeout, skip reconnectting", exc_info=e) + def _ensure_ray_initialized(self) -> None: + """Reject the call if Ray is not initialized. + + Why: a failed periodic ``_reconnect_ray`` may leave the process in a + ``ray.shutdown`` state. Without this guard a subsequent ``ray.get`` / + ``ray.get_actor`` would fall through to Ray's auto-init hook and + ``ray.init()`` with no args, spawning a local cluster on the admin host + and OOM-ing under concurrent requests. + """ + if not ray.is_initialized(): + raise InternalServerRockError( + "Ray cluster is not initialized; refusing to call ray API to avoid " + "spawning a local cluster via auto-init.", + code=codes.INTERNAL_SERVER_ERROR, + ) + async def async_ray_get(self, ray_future: ray.ObjectRef, timeout: int = 60): """ Asynchronously get the result of a Ray ObjectRef. @@ -102,6 +149,7 @@ async def async_ray_get(self, ray_future: ray.ObjectRef, timeout: int = 60): Raises: Exception: If ray.get fails """ + self._ensure_ray_initialized() self.increment_ray_request_count() loop = asyncio.get_running_loop() try: @@ -127,6 +175,7 @@ async def async_ray_get_actor(self, actor_name: str, namespace: str = None): ValueError: If actor does not exist Exception: If ray.get_actor fails """ + self._ensure_ray_initialized() self.increment_ray_request_count() namespace = namespace or self._config.namespace loop = asyncio.get_running_loop() diff --git a/rock/config.py b/rock/config.py index 03ccd8534a..1495bdb327 100644 --- a/rock/config.py +++ b/rock/config.py @@ -24,6 +24,8 @@ class RayConfig: ray_reconnect_request_threshold: int = field(default=10 * 1024 * 1024) ray_reconnect_check_interval_seconds: int = field(default=60 * 10) ray_reconnect_wait_timeout_seconds: int = field(default=30) + ray_reconnect_max_attempts: int = field(default=2) + ray_reconnect_retry_backoff_seconds: float = field(default=5.0) def __post_init__(self): if self.temp_dir: diff --git a/tests/unit/admin/core/test_ray_service.py b/tests/unit/admin/core/test_ray_service.py index 0a9602d3b7..8b4da1482c 100644 --- a/tests/unit/admin/core/test_ray_service.py +++ b/tests/unit/admin/core/test_ray_service.py @@ -2,12 +2,129 @@ import pytest +from rock import InternalServerRockError from rock.admin.core.ray_service import RayService +from rock.config import RayConfig from rock.deployments.config import RayDeploymentConfig from rock.deployments.ray import RayDeployment from rock.sandbox.sandbox_actor import SandboxActor +def _make_service(**overrides) -> RayService: + cfg_kwargs = dict( + address=None, + ray_reconnect_enabled=False, + ray_reconnect_wait_timeout_seconds=1, + ray_reconnect_max_attempts=3, + ray_reconnect_retry_backoff_seconds=0, + ) + cfg_kwargs.update(overrides) + return RayService(RayConfig(**cfg_kwargs)) + + +@pytest.mark.asyncio +async def test_reconnect_ray_retries_on_init_failure_and_eventually_succeeds(): + service = _make_service(ray_reconnect_max_attempts=3) + service._ray_request_count = 99 + old_establish_time = service._ray_establish_time + + init_calls = {"n": 0} + + def init_side_effect(**_kwargs): + init_calls["n"] += 1 + if init_calls["n"] < 3: + raise ConnectionError("ray head unreachable") + + with ( + patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, + patch("rock.admin.core.ray_service.ray.init", side_effect=init_side_effect) as mock_init, + patch("rock.admin.core.ray_service.time.time", return_value=old_establish_time + 5), + ): + await service._reconnect_ray() + + assert mock_init.call_count == 3 + assert mock_shutdown.call_count == 3 + assert service._ray_request_count == 0 + assert service._ray_establish_time == old_establish_time + 5 + + +@pytest.mark.asyncio +async def test_reconnect_ray_does_not_reset_counters_when_all_attempts_fail(): + service = _make_service(ray_reconnect_max_attempts=2) + service._ray_request_count = 99 + old_establish_time = service._ray_establish_time + + with ( + patch("rock.admin.core.ray_service.ray.shutdown"), + patch("rock.admin.core.ray_service.ray.init", side_effect=ConnectionError("down")), + ): + await service._reconnect_ray() + + # Counters preserved so the next scheduler tick will retry promptly. + assert service._ray_request_count == 99 + assert service._ray_establish_time == old_establish_time + + +@pytest.mark.asyncio +async def test_reconnect_ray_releases_write_lock_after_init_failure(): + service = _make_service(ray_reconnect_max_attempts=1) + + with ( + patch("rock.admin.core.ray_service.ray.shutdown"), + patch("rock.admin.core.ray_service.ray.init", side_effect=ConnectionError("down")), + ): + await service._reconnect_ray() + + # After the failed reconnect a reader must still be able to acquire the lock. + async with service._ray_rwlock.read_lock(): + pass + + +@pytest.mark.asyncio +async def test_async_ray_get_raises_when_ray_not_initialized(): + service = _make_service() + + fake_ref = MagicMock() + with ( + patch("rock.admin.core.ray_service.ray.is_initialized", return_value=False), + patch("rock.admin.core.ray_service.ray.get") as mock_ray_get, + ): + with pytest.raises(InternalServerRockError): + await service.async_ray_get(fake_ref, timeout=1) + + # Must short-circuit BEFORE touching ray.get; otherwise auto-init could spawn a + # local Ray cluster on the admin host. + mock_ray_get.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_ray_get_actor_raises_when_ray_not_initialized(): + service = _make_service() + + with ( + patch("rock.admin.core.ray_service.ray.is_initialized", return_value=False), + patch("rock.admin.core.ray_service.ray.get_actor") as mock_get_actor, + ): + with pytest.raises(InternalServerRockError): + await service.async_ray_get_actor("any-actor", namespace="ns") + + mock_get_actor.assert_not_called() + + +@pytest.mark.asyncio +async def test_reconnect_ray_logs_critical_when_all_attempts_exhausted(): + service = _make_service(ray_reconnect_max_attempts=2) + + with ( + patch("rock.admin.core.ray_service.ray.shutdown"), + patch("rock.admin.core.ray_service.ray.init", side_effect=ConnectionError("down")), + patch("rock.admin.core.ray_service.logger") as mock_logger, + ): + await service._reconnect_ray() + + assert mock_logger.critical.called, "expected logger.critical when ray reconnect exhausts all retries" + + @pytest.mark.need_ray @pytest.mark.asyncio async def test_reconnect_ray_calls_ray_shutdown_and_init_and_reset_counters(ray_service: RayService): From 3b22e10ca1c4ffe353870f167314c57cd00c8cbf Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Fri, 15 May 2026 17:21:25 +0800 Subject: [PATCH 098/226] =?UTF-8?q?[REFACTOR]=20OSS=20=E4=B8=8A=E4=BC=A0/?= =?UTF-8?q?=E4=B8=8B=E8=BD=BD=E4=BB=8E=E5=AE=A2=E6=88=B7=E7=AB=AF=20env=20?= =?UTF-8?q?vars=20=E8=A7=A3=E8=80=A6=EF=BC=9B3=20=E5=B1=82=E9=85=8D?= =?UTF-8?q?=E7=BD=AE=E8=A7=A3=E6=9E=90=E6=9C=BA=E5=88=B6=20(#943)=20(#949)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(sdk): scaffold OssClient class for OSS operations * feat(sdk): add deterministic OSS object name computation * feat(sdk): add layered OSS config resolution (env > server) * feat(sdk): migrate STS credentials fetch + token expiration check to OssClient * fix(sdk): preserve 5min buffer + fromisoformat in OssClient._is_token_expired * feat(sdk): implement OssClient setup with layered config resolution * feat(sdk): implement OssClient.upload_via_oss with deterministic naming * fix(sdk): unique session name in upload_via_oss verification step * feat(sdk): implement OssClient.download_via_oss * feat(sdk): add OssClient.close placeholder * feat(server): extend /get_token response with Endpoint/Bucket/Region * feat(server): warn at startup when OSS config is partially set * refactor(sdk): make Sandbox compose OssClient for OSS operations * fix(sdk): preserve explicit UploadMode.OSS failure when OSS unavailable * refactor(sdk): delegate download_file to OssClient.download_via_oss * test(sdk): update integration mock to OssClient._get_sts_credentials * feat(sdk): add OssClient.schedule_async_persistence (fire-and-forget) * feat(sdk): OssClient.close awaits pending persistence tasks with 5s timeout * feat(sdk): trigger async OSS persistence for small file uploads * fix(sdk): ensure OSS setup before small-file persistence to avoid permanently disabled state * feat(sdk): Sandbox.close awaits OssClient persistence tasks * test: attach caplog handler directly to logger to bypass propagate=False * test(integration): update assertion to match new 'OSS is not available' message * feat(server): prefer env vars over YAML in OSS config (align with client Layer 1) * refactor(sdk): rename _oss_client.py to oss_client.py OssClient 是 Sandbox 公开组合的成员（self._oss），底层模块没必要 _ 私有前缀。 Co-Authored-By: Claude Opus 4.7 * fix(sdk): restore _oss_bucket and _oss_token_expire_time in Sandbox.__repr__ 重构后 OssClient 持有这两个字段，__repr__ 通过 self._oss 读取，输出格式与重构前保持一致，方便日志/调试对比。 Co-Authored-By: Claude Opus 4.7 * refactor(sdk): unify file-existence check via arun in OssClient upload_via_oss / download_via_oss 的 verify 步骤都改成 arun(test -f, NORMAL) 单调用，跟同函数里 mkdir -p 的风格一致。download 路径加 shlex.quote 保留原 list-form 的 shell 安全性。顺手清掉因此变成 unused 的 6 个 import。 Co-Authored-By: Claude Opus 4.7 * style(sdk): use f-string in OssClient logging; drop redundant sleep(0) - 6 处 logger 调用从 %s/%d 占位符改为 f-string，统一风格 - 删除 schedule_async_persistence 末尾的 await asyncio.sleep(0)：close() 用 _pending_persistence_tasks 等任务，create_task 那一刻 task 已加入 set，让出事件循环并不必要 - 修 test_schedules_task_when_available：把 gather + mock 断言挪到 with patch 块内（原本依赖 sleep(0) 让 task 在 patch 仍生效时跑过 mock，删了就会暴露 patch 已退出后才 gather 的测试 bug） Co-Authored-By: Claude Opus 4.7 * chore(sdk): drop log-content assertions and translate Chinese comments - Remove the _capture_on helper and log-content assertions: - test_oss_client: simplify two log-warning tests to behavior-only (failure does not raise; close timeout does not hang) - test_sandbox_proxy: drop TestStartupOssConfigWarning class entirely (3 tests reduced to no-op without log capture) - Translate Chinese comments/docstrings to English across the OSS refactor scope (sandbox_proxy_service, client, oss_client, the two test files above) Co-Authored-By: Claude Opus 4.7 * fix(sdk): add BC shims for legacy _oss_token_expire_time / _is_token_expired Downstream xrl/intetest test_timestamp_expire still writes sandbox._oss_token_expire_time and calls sandbox._is_token_expired(). Refactor moved both onto OssClient; expose a property (with setter) and a method wrapper that proxy back to self._oss to keep the test surface intact. _oss_bucket was only touched by helpers under @pytest.mark.skip and is deliberately not shimmed. Co-Authored-By: Claude Opus 4.7 * fix(sdk): use execute() not arun() for test -f file-existence checks in OSS arun() raises an exception when the underlying command exits non-zero (rocklet maps non-zero exit_code to status='Failed'). For 'test -f X', a missing file is the very branch we need to inspect, so the exception path short-circuits download_via_oss / upload_via_oss verify and surfaces as "Failed to execute command" instead of a clean "not found" response. Switch the two test -f sites back to execute(Command([...])), which always returns the exit_code without raising. Update unit-test mocks accordingly. Surfaced by integration test test_download_file when checking a non-existent remote path. Co-Authored-By: Claude Opus 4.7 * fix(sdk): use sandbox basename for OSS object names, not local basename OSS objects mirror sandbox-side files; naming them after the sandbox path keeps the OSS-side identity stable when the same sandbox file is downloaded to differently-named local destinations. The previous local-first ordering also broke integration test_download_file Test 3, where remote '/tmp/test_download.txt' downloaded to 'downloaded.txt' produced an OSS object the test mock could not look up by remote basename. Falls back to local basename when sandbox_path is empty (mirror of the old fallback path). Co-Authored-By: Claude Opus 4.7 * refactor(sdk): inline _OSS_CLOSE_TIMEOUT_SECONDS into close() default arg Drop the module-level constant in favor of a default parameter on OssClient.close(timeout=5.0). Production callers stay unchanged; the single-test timeout override is now an explicit close(timeout=0.05) call instead of patching a module symbol. Co-Authored-By: Claude Opus 4.7 * refactor(sdk): drop BC shims for _is_token_expired / _oss_token_expire_time Downstream xrl/intetest .. test_timestamp_expire (the only consumer that hit these shims via a non-skipped path) was deleted. The remaining references in mock_setup_oss are reached only by tests that already carry @pytest.mark.skip, so the proxy property + method are dead weight now. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 --- rock/sandbox/service/sandbox_proxy_service.py | 33 +- rock/sdk/sandbox/client.py | 137 +---- rock/sdk/sandbox/file_system.py | 119 +---- rock/sdk/sandbox/oss_client.py | 338 ++++++++++++ .../sdk/sandbox/test_file_system.py | 8 +- tests/unit/sandbox/test_sandbox_proxy.py | 110 ++++ tests/unit/sdk/sandbox/test_file_system.py | 53 ++ tests/unit/sdk/sandbox/test_oss_client.py | 492 ++++++++++++++++++ tests/unit/sdk/sandbox/test_upload_by_path.py | 104 ++++ 9 files changed, 1170 insertions(+), 224 deletions(-) create mode 100644 rock/sdk/sandbox/oss_client.py create mode 100644 tests/unit/sdk/sandbox/test_oss_client.py create mode 100644 tests/unit/sdk/sandbox/test_upload_by_path.py diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index bcb138c6a2..55aa23fdc2 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -78,6 +78,29 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): ) self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count + self._validate_oss_config_or_warn() + + def _validate_oss_config_or_warn(self) -> None: + # Same resolution order as gen_oss_sts_token: env > YAML + endpoint = env_vars.ROCK_OSS_BUCKET_ENDPOINT or self.oss_config.endpoint + bucket = env_vars.ROCK_OSS_BUCKET_NAME or self.oss_config.bucket + region = env_vars.ROCK_OSS_BUCKET_REGION + fields = { + "endpoint": endpoint, + "bucket": bucket, + "region": region, + "access_key_id": self.oss_config.access_key_id, + "access_key_secret": self.oss_config.access_key_secret, + "role_arn": self.oss_config.role_arn, + } + set_count = sum(1 for v in fields.values() if v) + if 0 < set_count < len(fields): + missing = [k for k, v in fields.items() if not v] + logger.warning( + "OSS configuration is partially set. Missing fields: %s. " + "Server will return null for these in /get_token, clients fall back accordingly.", + ", ".join(missing), + ) @monitor_sandbox_operation() async def create_session(self, request: CreateSessionRequest) -> CreateBashSessionResponse: @@ -658,11 +681,19 @@ def gen_oss_sts_token(self): try: body = self.sts_client.do_action_with_exception(request) token = json.loads(oss2.to_unicode(body)) - return token["Credentials"] + credentials = token["Credentials"] except Exception: logger.error("generate oss sts token failed") return None + return { + **credentials, + # env > YAML, matches client-side Layer 1 priority + "Endpoint": env_vars.ROCK_OSS_BUCKET_ENDPOINT or self.oss_config.endpoint or None, + "Bucket": env_vars.ROCK_OSS_BUCKET_NAME or self.oss_config.bucket or None, + "Region": env_vars.ROCK_OSS_BUCKET_REGION or None, + } + async def get_sandbox_websocket_url( self, sandbox_id: str, target_path: str | None = None, port: int | None = None ) -> str: diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index e27f9b60c1..65527d91b8 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -3,19 +3,15 @@ import math import mimetypes import os -import shlex import time import uuid import warnings -from datetime import datetime, timedelta, timezone from enum import Enum from pathlib import Path -import oss2 from httpx import ReadTimeout from typing_extensions import deprecated -from rock import env_vars from rock.actions import ( AbstractSandbox, Action, @@ -30,7 +26,6 @@ ExecuteBashSessionResponse, IsAliveResponse, Observation, - OssSetupResponse, ReadFileRequest, ReadFileResponse, SandboxResponse, @@ -55,6 +50,7 @@ from rock.sdk.sandbox.file_system import FileSystem, LinuxFileSystem from rock.sdk.sandbox.model_service.base import ModelService from rock.sdk.sandbox.network import Network +from rock.sdk.sandbox.oss_client import OssClient from rock.sdk.sandbox.process import Process from rock.sdk.sandbox.remote_user import LinuxRemoteUser, RemoteUser from rock.sdk.sandbox.runtime_env.base import RuntimeEnv, RuntimeEnvId @@ -75,7 +71,6 @@ class Sandbox(AbstractSandbox): _sandbox_id: str | None = None _host_name: str | None = None _host_ip: str | None = None - _oss_bucket: oss2.Bucket | None = None _cluster: str | None = None _namespace: str | None = None _experiment_id: str | None = None @@ -99,7 +94,6 @@ def __init__(self, config: SandboxConfig): else: self._route_key = self.config.route_key - self._oss_token_expire_time = self._generate_utc_iso_time() self._cluster = self.config.cluster self.remote_user = LinuxRemoteUser(self) self.process = Process(self) @@ -108,6 +102,7 @@ def __init__(self, config: SandboxConfig): self.runtime_envs = {} self.deploy = Deploy(self) self.agent = RockAgent(self) + self._oss = OssClient(self) @property def sandbox_id(self) -> str: @@ -718,11 +713,18 @@ async def upload_by_path( if not file_path.exists(): return UploadResponse(success=False, message=f"File not found: {file_path}") if upload_mode == UploadMode.OSS or ( - upload_mode != UploadMode.DIRECT - and env_vars.ROCK_OSS_ENABLE - and os.path.getsize(file_path) > 1024 * 1024 * 1 + upload_mode != UploadMode.DIRECT and os.path.getsize(file_path) > 1024 * 1024 * 1 ): - return await self._upload_via_oss(path_str, target_path) + await self._oss.ensure_setup() + if self._oss.is_available: + return await self._oss.upload_via_oss(path_str, target_path) + # Explicit OSS requested but unavailable -> fail (BC: preserve legacy behavior) + if upload_mode == UploadMode.OSS: + return UploadResponse( + success=False, + message="Failed to upload file, please setup oss bucket first", + ) + # Otherwise fall through to admin /upload (natural degradation for auto / default large files) url = f"{self._url}/upload" headers = self._build_headers() @@ -748,8 +750,13 @@ async def upload_by_path( logging.debug(f"Upload response: {response}") if "Success" != response.get("status"): return UploadResponse(success=False, message=f"Failed to execute command: upload response: {response}") - else: - return UploadResponse(success=True, message=f"Successfully uploaded file {filename} to {target_path}") + # Admin /upload succeeded; opportunistically persist to OSS in background. + # Skipped silently when OSS is not configured/available. + # ensure_setup is idempotent and short-circuits when OSS is unavailable, + # so small-file-only flows still get a chance to bootstrap OSS persistence. + if await self._oss.ensure_setup() and self._oss.is_available: + await self._oss.schedule_async_persistence(path_str, target_path) + return UploadResponse(success=True, message=f"Successfully uploaded file {filename} to {target_path}") async def read_file(self, request: ReadFileRequest) -> ReadFileResponse: url = f"{self._url}/read_file" @@ -832,84 +839,6 @@ async def _generate_tmp_session_name(self) -> str: timestamp = str(time.time_ns()) return f"bash-{timestamp}" - async def _upload_via_oss(self, file_path: str | Path, target_path: str): - if self._oss_bucket is None or self._is_token_expired(): - setup_response: OssSetupResponse = await self._setup_oss() - if not setup_response.success: - return UploadResponse(success=False, message="Failed to upload file, please setup oss bucket first") - timestamp = str(time.time_ns()) - file_name = file_path.split("/")[-1] - tmp_obj_name = f"{timestamp}-{file_name}" - oss2.resumable_upload(self._oss_bucket, tmp_obj_name, file_path) - url = self._oss_bucket.sign_url("GET", tmp_obj_name, 600, slash_safe=True) - try: - # wget -O does not create missing parent dirs; mkdir -p first to - # match the multipart upload path (rocklet /upload mkdirs server-side). - parent_dir = str(Path(target_path).parent) - await self.arun( - cmd=f"mkdir -p {shlex.quote(parent_dir)}", - wait_timeout=10, - mode=RunMode.NORMAL, - ) - - download_cmd = f"wget -c -O {target_path} '{url}'" - await self.arun(cmd=download_cmd, wait_timeout=600, mode=RunMode.NOHUP) - check_file_session = f"bash-{timestamp}" - await self.create_session(CreateBashSessionRequest(session=check_file_session)) - check_file_cmd = f"test -f {target_path}" - check_response: Observation = await self._run_in_session( - action=BashAction(command=check_file_cmd, session=check_file_session) - ) - if not check_response.exit_code == 0: - return UploadResponse( - success=False, message=f"Failed to upload file {file_name}, sandbox download phase failed" - ) - else: - return UploadResponse(success=True, message=f"Successfully uploaded file {file_name} to {target_path}") - except Exception: - return UploadResponse(success=False, message=f"Failed to upload file {file_name} to {target_path}") - - async def _get_oss_sts_credentials(self) -> dict: - """Get OSS STS credentials from sandbox and update token expiration time. - - Side effects: - Updates self._oss_token_expire_time for token expiration checking - - Returns: - dict: STS credentials with keys: AccessKeyId, AccessKeySecret, SecurityToken, Expiration - - Raises: - Exception: If HTTP request fails or response is invalid - """ - url = f"{self._url}/get_token" - headers = self._build_headers() - response = await HttpUtils.get(url, headers) - if response["status"] != "Success": - raise Exception(f"Failed to get OSS STS token: {response.get('message', 'Unknown error')}") - - credentials = response["result"] - self._oss_token_expire_time = credentials["Expiration"] - return credentials - - async def _setup_oss(self) -> OssSetupResponse: - try: - credentials = await self._get_oss_sts_credentials() - auth = oss2.StsAuth( - credentials["AccessKeyId"], - credentials["AccessKeySecret"], - credentials["SecurityToken"], - ) - - self._oss_bucket = oss2.Bucket( - auth=auth, - endpoint=env_vars.ROCK_OSS_BUCKET_ENDPOINT, - bucket_name=env_vars.ROCK_OSS_BUCKET_NAME, - region=env_vars.ROCK_OSS_BUCKET_REGION, - ) - except Exception as e: - return OssSetupResponse(success=False, message=f"Failed to setup oss bucket: {e}") - return OssSetupResponse(success=True, message="Successfully setup oss bucket") - def _add_user_defined_tag_into_headers(self, headers: dict): if self.config.user_id: headers["X-User-Id"] = self.config.user_id @@ -918,22 +847,6 @@ def _add_user_defined_tag_into_headers(self, headers: dict): if self.config.namespace: headers["X-Namespace"] = self.config.namespace - def _is_token_expired(self) -> bool: - try: - expire_time = datetime.fromisoformat(self._oss_token_expire_time.replace("Z", "+00:00")) - current_time = datetime.now(timezone.utc) - - buffer_time = timedelta(minutes=5) - effective_expire_time = expire_time - buffer_time - - return current_time >= effective_expire_time - - except (ValueError, AttributeError): - return True - - def _generate_utc_iso_time(self): - return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - async def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse: url = f"{self._url}/close_session" headers = self._build_headers() @@ -953,6 +866,12 @@ async def close_session(self, request: CloseSessionRequest) -> CloseSessionRespo return CloseSessionResponse(**result) async def close(self) -> CloseResponse: + # Drain pending async OSS persistence tasks (with timeout) before + # tearing down the sandbox so in-flight uploads have a chance to finish. + try: + await self._oss.close() + except Exception as e: + logging.warning(f"OssClient.close() failed, IGNORE: {e}") await self.stop() def __str__(self): @@ -975,11 +894,11 @@ def __repr__(self): f"_sandbox_id={self._sandbox_id!r}, " f"_host_name={self._host_name!r}, " f"_host_ip={self._host_ip!r}, " - f"_oss_bucket={self._oss_bucket!r}, " + f"_oss_bucket={self._oss._bucket!r}, " f"_cluster={self._cluster!r}, " f"_pod_name={self._pod_name!r}, " f"_ip={self._ip!r}, " - f"_oss_token_expire_time={self._oss_token_expire_time!r}" + f"_oss_token_expire_time={self._oss._token_expire_time!r}" f")" ) diff --git a/rock/sdk/sandbox/file_system.py b/rock/sdk/sandbox/file_system.py index bb56a41f60..58a0720d84 100644 --- a/rock/sdk/sandbox/file_system.py +++ b/rock/sdk/sandbox/file_system.py @@ -5,9 +5,6 @@ from abc import ABC, abstractmethod from pathlib import Path -import oss2 - -from rock import env_vars from rock.actions import CreateBashSessionRequest, Observation from rock.actions.sandbox.base import AbstractSandbox from rock.actions.sandbox.request import ChmodRequest, ChownRequest, Command @@ -186,120 +183,22 @@ async def download_file( remote_path: str, local_path: str | Path, ) -> DownloadFileResponse: - """Download file from sandbox container to local machine using OSS as intermediary. - - Flow: - 1. Check OSS is enabled via ROCK_OSS_ENABLE - 2. Verify source file exists in sandbox container (must be a regular file, not directory) - 3. Ensure ossutil is installed (auto-install if missing, checks wget/curl/unzip) - 4. Verify ossutil is working by running 'ossutil version' - 5. Get STS credentials from sandbox via get_token API - 6. Upload file from sandbox to OSS using ossutil (with STS token) - 7. Download file from OSS to local using oss2 (with same STS token) - 8. Verify downloaded file exists locally - - Note: - - Only supports regular files, NOT directories. To download a directory, first create a - tar archive in the sandbox (e.g., `tar czf /tmp/mydir.tar.gz /path/to/mydir`), then - download the tar file. - - Temporary OSS object is NOT cleaned up automatically (aligns with _upload_via_oss behavior). - Consider using OSS lifecycle policies for automatic cleanup. - - Args: - remote_path: File path inside the sandbox container (absolute path recommended, must be a regular file) - local_path: Target file path on the local machine (supports ~/ expansion) - - Returns: - DownloadFileResponse with success status and message + """Download file from sandbox container to local machine via OSS. - Raises: - AssertionError: If sandbox is not an instance of Sandbox class + OSS availability is determined by OssClient (Layer 1 env > Layer 2 server response). + Returns failure if OSS is unavailable. """ - from rock.sdk.sandbox.client import RunMode, Sandbox + from rock.sdk.sandbox.client import Sandbox - # Assert sandbox is Sandbox instance assert isinstance(self.sandbox, Sandbox), "sandbox must be an instance of Sandbox" - timestamp = str(time.time_ns()) - - try: - # Check OSS enable - if not env_vars.ROCK_OSS_ENABLE: - return DownloadFileResponse( - success=False, message="OSS download is not enabled. Please set ROCK_OSS_ENABLE=true" - ) - - # Check remote file exists (must be a regular file, not directory) - check_response: CommandResponse = await self.sandbox.execute(Command(command=["test", "-f", remote_path])) - if check_response.exit_code != 0: - return DownloadFileResponse( - success=False, - message=f"Source file not found or is not a regular file in sandbox: {remote_path}. " - "Note: Only regular files are supported. For directories, create a tar archive first.", - ) - - if not await self.ensure_ossutil(): - return DownloadFileResponse(success=False, message="Failed to ensure ossutil is installed and working") - - # Get STS credentials from sandbox (for both ossutil upload and oss2 download) - try: - credentials = await self.sandbox._get_oss_sts_credentials() - except Exception as e: - return DownloadFileResponse(success=False, message=f"Failed to get OSS STS token: {e}") - - access_key_id = credentials["AccessKeyId"] - access_key_secret = credentials["AccessKeySecret"] - security_token = credentials["SecurityToken"] - - endpoint = env_vars.ROCK_OSS_BUCKET_ENDPOINT - bucket_name = env_vars.ROCK_OSS_BUCKET_NAME - region = env_vars.ROCK_OSS_BUCKET_REGION - - # Upload from sandbox to OSS via ossutil - file_name = remote_path.split("/")[-1] - tmp_obj_name = f"{timestamp}-{file_name}" - oss_url = f"oss://{bucket_name}/{tmp_obj_name}" - - # Build ossutil command with bash -c wrapper (required for nohup mode) - ossutil_inner_cmd = ( - f"ossutil cp {shlex.quote(remote_path)} {shlex.quote(oss_url)}" - f" --access-key-id {shlex.quote(access_key_id)}" - f" --access-key-secret {shlex.quote(access_key_secret)}" - f" --sts-token {shlex.quote(security_token)}" - f" --endpoint {shlex.quote(endpoint)}" - f" --region {shlex.quote(region)}" - ) - ossutil_cmd = f"bash -c {shlex.quote(ossutil_inner_cmd)}" - logger.debug(f"Uploading {remote_path} to OSS via ossutil") - upload_response = await self.sandbox.arun(cmd=ossutil_cmd, mode=RunMode.NOHUP) - if upload_response.exit_code != 0: - return DownloadFileResponse( - success=False, - message=f"Failed to upload file to OSS (exit_code={upload_response.exit_code}): {upload_response.output}", - ) - logger.debug(f"ossutil upload completed: {upload_response.output}") - - # Download from OSS to local via oss2 - oss_auth = oss2.StsAuth(access_key_id, access_key_secret, security_token) - oss_bucket = oss2.Bucket(oss_auth, endpoint, bucket_name, region=region) - - local = Path(local_path).expanduser().resolve() - local.parent.mkdir(parents=True, exist_ok=True) - try: - oss_bucket.get_object_to_file(tmp_obj_name, str(local)) - except Exception as e: - return DownloadFileResponse(success=False, message=f"Failed to download from OSS: {str(e)}") - - # Verify local file exists - if not local.exists(): - return DownloadFileResponse(success=False, message=f"Downloaded file not found at: {local}") + if not await self.sandbox._oss.ensure_setup(): + return DownloadFileResponse(success=False, message="OSS is not available") - # Note: OSS temporary object is NOT cleaned up here to align with _upload_via_oss behavior - return DownloadFileResponse(success=True, message=f"Successfully downloaded {remote_path} to {local}") + if not await self.ensure_ossutil(): + return DownloadFileResponse(success=False, message="Failed to ensure ossutil is installed and working") - except Exception as e: - logger.exception(f"Unexpected error during download_by_oss: {e}") - return DownloadFileResponse(success=False, message=f"Unexpected error: {str(e)}") + return await self.sandbox._oss.download_via_oss(remote_path, Path(local_path)) async def ensure_ossutil(self) -> bool: """Ensure ossutil is installed in the sandbox. Returns True if ready.""" diff --git a/rock/sdk/sandbox/oss_client.py b/rock/sdk/sandbox/oss_client.py new file mode 100644 index 0000000000..51711ff133 --- /dev/null +++ b/rock/sdk/sandbox/oss_client.py @@ -0,0 +1,338 @@ +"""OssClient — encapsulates all OSS interactions for a Sandbox. + +Holds OSS state (bucket, token expiration, async persistence tasks) and +exposes upload / download / persistence operations. Composed by Sandbox. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import shlex +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import TYPE_CHECKING + +import oss2 + +from rock import env_vars +from rock.actions.sandbox.request import Command +from rock.actions.sandbox.response import DownloadFileResponse, UploadResponse +from rock.logger import init_logger +from rock.utils.http import HttpUtils + +if TYPE_CHECKING: + from rock.sdk.sandbox.client import Sandbox + +logger = init_logger(__name__) + + +@dataclass +class OssClientConfig: + """Resolved OSS configuration (Layer 1 env or Layer 2 server).""" + + endpoint: str + bucket: str + region: str + enabled_via_env: bool # True = Layer 1 (gated by ROCK_OSS_ENABLE); False = Layer 2 + + +class OssClient: + """OSS operations for a single Sandbox instance.""" + + def __init__(self, sandbox: Sandbox): + self._sandbox = sandbox + self._bucket = None + self._token_expire_time: str | None = None + self._client_config: OssClientConfig | None = None + self._pending_persistence_tasks: set[asyncio.Task] = set() + + @staticmethod + def _compute_object_name(sandbox_id: str, local_path: str, sandbox_path: str) -> str: + # Prefer sandbox basename: the OSS object mirrors a sandbox-side file, + # so naming it after the sandbox path keeps OSS-side names meaningful + # even when local destinations differ (e.g. download to a renamed file). + payload = f"{sandbox_id}|{local_path}|{sandbox_path}" + digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() + filename = Path(sandbox_path).name or Path(local_path).name + return f"{digest}-{filename}" + + @staticmethod + def _resolve_config(sts_response: dict) -> OssClientConfig | None: + # Layer 1: env var (highest priority) + env_endpoint = env_vars.ROCK_OSS_BUCKET_ENDPOINT + env_bucket = env_vars.ROCK_OSS_BUCKET_NAME + env_region = env_vars.ROCK_OSS_BUCKET_REGION + if env_endpoint and env_bucket and env_region: + return OssClientConfig( + endpoint=env_endpoint, + bucket=env_bucket, + region=env_region, + enabled_via_env=True, + ) + + # Layer 2: server response (fallback default) + resp_endpoint = sts_response.get("Endpoint") + resp_bucket = sts_response.get("Bucket") + resp_region = sts_response.get("Region") + if resp_endpoint and resp_bucket and resp_region: + return OssClientConfig( + endpoint=resp_endpoint, + bucket=resp_bucket, + region=resp_region, + enabled_via_env=False, + ) + + # Layer 3: OSS unavailable + return None + + async def _get_sts_credentials(self) -> dict: + """Fetch STS credentials and OSS config from /get_token endpoint. + + Returns the entire response result dict, which may include: + - STS creds: AccessKeyId, AccessKeySecret, SecurityToken, Expiration + - OSS config (if server is new + configured): Endpoint, Bucket, Region + + Side effect: caches Expiration in self._token_expire_time. + """ + url = f"{self._sandbox._url}/get_token" + headers = self._sandbox._build_headers() + response = await HttpUtils.get(url, headers) + if response["status"] != "Success": + raise Exception(f"Failed to get OSS STS token: {response.get('message', 'Unknown error')}") + credentials = response["result"] + self._token_expire_time = credentials["Expiration"] + return credentials + + def _is_token_expired(self) -> bool: + """Whether cached token is missing, malformed, or within 5min of expiration.""" + try: + expire_time = datetime.fromisoformat(self._token_expire_time.replace("Z", "+00:00")) + current_time = datetime.now(timezone.utc) + effective_expire_time = expire_time - timedelta(minutes=5) + return current_time >= effective_expire_time + except (ValueError, AttributeError): + return True + + @property + def is_available(self) -> bool: + """Whether OSS is available: bucket has been successfully initialized.""" + return self._bucket is not None + + async def ensure_setup(self) -> bool: + """Ensure OSS bucket is set up and token is fresh. Idempotent. + + Returns True if OSS is available, False otherwise. + """ + if self._bucket is not None and not self._is_token_expired(): + return True + return await self._setup() + + async def _setup(self) -> bool: + try: + sts_response = await self._get_sts_credentials() + except Exception as e: + logger.warning(f"Failed to get STS credentials: {e}") + return False + + config = self._resolve_config(sts_response) + if config is None: + return False + + # Layer 1 also requires ROCK_OSS_ENABLE + if config.enabled_via_env and not env_vars.ROCK_OSS_ENABLE: + return False + + try: + auth = oss2.StsAuth( + sts_response["AccessKeyId"], + sts_response["AccessKeySecret"], + sts_response["SecurityToken"], + ) + self._bucket = oss2.Bucket( + auth=auth, + endpoint=config.endpoint, + bucket_name=config.bucket, + region=config.region, + ) + self._client_config = config + return True + except Exception as e: + logger.warning(f"Failed to initialize OSS bucket: {e}") + self._bucket = None + return False + + async def upload_via_oss(self, file_path: str, target_path: str) -> UploadResponse: + """Upload a local file to sandbox via OSS as intermediary (large file path). + + Steps: + 1. Compute deterministic OSS object name + 2. Resumable upload local file to OSS + 3. Sign a temporary GET URL + 4. Inside sandbox: mkdir -p parent + wget the signed URL + 5. Verify target file exists in sandbox + """ + from rock.sdk.sandbox.client import RunMode # late import to avoid circular + + if self._bucket is None: + return UploadResponse(success=False, message="OSS bucket not set up") + + file_name = Path(file_path).name + oss_object_name = self._compute_object_name( + sandbox_id=self._sandbox.sandbox_id, + local_path=file_path, + sandbox_path=target_path, + ) + + try: + oss2.resumable_upload(self._bucket, oss_object_name, file_path) + url = self._bucket.sign_url("GET", oss_object_name, 600, slash_safe=True) + + # mkdir -p target parent (wget -O does not auto-create dirs in NOHUP mode) + parent_dir = str(Path(target_path).parent) + await self._sandbox.arun( + cmd=f"mkdir -p {shlex.quote(parent_dir)}", + wait_timeout=10, + mode=RunMode.NORMAL, + ) + + # wget the signed URL + download_cmd = f"wget -c -O {target_path} '{url}'" + await self._sandbox.arun(cmd=download_cmd, wait_timeout=600, mode=RunMode.NOHUP) + + # Verify target exists in sandbox. Use execute() instead of arun(), + # because arun() raises on non-zero exit codes; here exit_code=1 + # (file missing) is a normal branch we need to inspect. + check = await self._sandbox.execute(Command(command=["test", "-f", target_path])) + if check.exit_code != 0: + return UploadResponse( + success=False, + message=f"Failed to upload file {file_name}, sandbox download phase failed", + ) + return UploadResponse( + success=True, + message=f"Successfully uploaded file {file_name} to {target_path}", + ) + except Exception as e: + logger.warning(f"upload_via_oss failed: {e}") + return UploadResponse( + success=False, + message=f"Failed to upload file {file_name} to {target_path}: {e}", + ) + + async def download_via_oss(self, remote_path: str, local_path: Path) -> DownloadFileResponse: + """Download file from sandbox to local via OSS as intermediary. + + Note: ensure_setup must succeed before calling. Caller is LinuxFileSystem, + which holds an `ensure_ossutil` helper for installing ossutil in the sandbox. + """ + from rock.sdk.sandbox.client import RunMode # late import + + if self._bucket is None or self._client_config is None: + return DownloadFileResponse(success=False, message="OSS is not available") + + # Verify source file exists in sandbox (must be regular file). + # Use execute() instead of arun(): arun() raises on non-zero exit + # codes, but exit_code=1 (file missing) is the branch we want to act on. + check = await self._sandbox.execute(Command(command=["test", "-f", remote_path])) + if check.exit_code != 0: + return DownloadFileResponse( + success=False, + message=( + f"Source file not found or is not a regular file in sandbox: {remote_path}. " + "Note: Only regular files are supported. For directories, create a tar archive first." + ), + ) + + # Caller (LinuxFileSystem) is responsible for ensure_ossutil before calling here. + + # Refresh STS creds for ossutil (use existing helper, will refresh if expired) + if self._is_token_expired(): + await self._setup() + sts_response = await self._get_sts_credentials() + access_key_id = sts_response["AccessKeyId"] + access_key_secret = sts_response["AccessKeySecret"] + security_token = sts_response["SecurityToken"] + + # Upload sandbox file to OSS via ossutil + oss_object_name = self._compute_object_name( + sandbox_id=self._sandbox.sandbox_id, + local_path=str(local_path), + sandbox_path=remote_path, + ) + oss_url = f"oss://{self._client_config.bucket}/{oss_object_name}" + + ossutil_inner = ( + f"ossutil cp {shlex.quote(remote_path)} {shlex.quote(oss_url)}" + f" --access-key-id {shlex.quote(access_key_id)}" + f" --access-key-secret {shlex.quote(access_key_secret)}" + f" --sts-token {shlex.quote(security_token)}" + f" --endpoint {shlex.quote(self._client_config.endpoint)}" + f" --region {shlex.quote(self._client_config.region)}" + ) + upload_cmd = f"bash -c {shlex.quote(ossutil_inner)}" + upload_resp = await self._sandbox.arun(cmd=upload_cmd, mode=RunMode.NOHUP) + if upload_resp.exit_code != 0: + return DownloadFileResponse( + success=False, + message=f"Failed to upload file to OSS (exit_code={upload_resp.exit_code}): {upload_resp.output}", + ) + + # Download from OSS to local via oss2 + auth = oss2.StsAuth(access_key_id, access_key_secret, security_token) + bucket = oss2.Bucket( + auth, self._client_config.endpoint, self._client_config.bucket, region=self._client_config.region + ) + local = Path(local_path).expanduser().resolve() + local.parent.mkdir(parents=True, exist_ok=True) + try: + bucket.get_object_to_file(oss_object_name, str(local)) + except Exception as e: + return DownloadFileResponse(success=False, message=f"Failed to download from OSS: {e}") + + if not local.exists(): + return DownloadFileResponse(success=False, message=f"Downloaded file not found at: {local}") + + return DownloadFileResponse(success=True, message=f"Successfully downloaded {remote_path} to {local}") + + async def schedule_async_persistence(self, local_path: str, sandbox_path: str) -> str | None: + """Fire-and-forget: upload local file to OSS in the background. + + Returns the OSS object key if scheduled, None if OSS is unavailable. + Failures are logged as warnings; main flow is unaffected. + """ + if self._bucket is None: + return None + + oss_object_name = self._compute_object_name( + sandbox_id=self._sandbox.sandbox_id, + local_path=local_path, + sandbox_path=sandbox_path, + ) + task = asyncio.create_task(self._persist_to_oss(local_path, oss_object_name)) + self._pending_persistence_tasks.add(task) + task.add_done_callback(self._pending_persistence_tasks.discard) + return oss_object_name + + async def _persist_to_oss(self, local_path: str, oss_object_name: str) -> None: + try: + await asyncio.to_thread(oss2.resumable_upload, self._bucket, oss_object_name, local_path) + logger.info(f"OSS persisted: {oss_object_name}") + except Exception as e: + logger.warning(f"OSS persistence failed for {oss_object_name}: {e}") + + async def close(self, timeout: float = 5.0) -> None: + """Wait for pending persistence tasks (with timeout).""" + if not self._pending_persistence_tasks: + return + try: + await asyncio.wait_for( + asyncio.gather(*self._pending_persistence_tasks, return_exceptions=True), + timeout=timeout, + ) + except asyncio.TimeoutError: + logger.warning( + f"OSS persistence tasks did not finish within {timeout}s on close " + f"({len(self._pending_persistence_tasks)} pending)" + ) diff --git a/tests/integration/sdk/sandbox/test_file_system.py b/tests/integration/sdk/sandbox/test_file_system.py index 78a73479d8..0aab40d7c6 100644 --- a/tests/integration/sdk/sandbox/test_file_system.py +++ b/tests/integration/sdk/sandbox/test_file_system.py @@ -104,9 +104,9 @@ async def test_download_file(sandbox_instance: Sandbox, monkeypatch): ) assert not response.success, "Download should fail when OSS is disabled" - assert "not enabled" in response.message.lower(), ( - f"Error message should mention OSS disabled: {response.message}" - ) + assert ( + "not available" in response.message.lower() + ), f"Error message should mention OSS unavailable: {response.message}" logger.info("✓ OSS disabled error handling works correctly") # Setup mocks for remaining tests @@ -123,7 +123,7 @@ async def mock_get_sts_credentials(): "Expiration": "2026-03-15T00:00:00Z", } - monkeypatch.setattr(sandbox_instance, "_get_oss_sts_credentials", mock_get_sts_credentials) + monkeypatch.setattr(sandbox_instance._oss, "_get_sts_credentials", mock_get_sts_credentials) original_arun = sandbox_instance.arun diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index f98802c6b7..d8f44fa40b 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -1,8 +1,10 @@ import uuid +from unittest.mock import MagicMock, patch import pytest from rock.actions.sandbox.response import State +from rock.config import OssConfig from rock.deployments.config import DockerDeploymentConfig from rock.sandbox.sandbox_manager import SandboxManager from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService @@ -84,3 +86,111 @@ async def test_list_sandbox(sandbox_manager: SandboxManager, sandbox_proxy_servi assert len(result.items) == 0 await sandbox_manager.stop(sandbox_id1) await sandbox_manager.stop(sandbox_id2) + + +class TestGenOssStsToken: + @pytest.fixture + def sandbox_proxy_service(self): + # Build a minimal SandboxProxyService without going through __init__ + # (which requires real Redis / metrics / RAM-Acs client setup). + service = SandboxProxyService.__new__(SandboxProxyService) + service.oss_config = OssConfig(role_arn="test_role_arn") + service.sts_client = MagicMock() + return service + + def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): + sandbox_proxy_service.oss_config.endpoint = "ep" + sandbox_proxy_service.oss_config.bucket = "bk" + + fake_token_body = ( + b'{"Credentials": {"AccessKeyId":"ak","AccessKeySecret":"sk",' + b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' + ) + with ( + patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, + ): + mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" + mock_env.ROCK_OSS_BUCKET_NAME = "" + mock_env.ROCK_OSS_BUCKET_REGION = "rg" + result = sandbox_proxy_service.gen_oss_sts_token() + + assert result["AccessKeyId"] == "ak" + assert result["Endpoint"] == "ep" + assert result["Bucket"] == "bk" + assert result["Region"] == "rg" + + def test_sts_failure_returns_none(self, sandbox_proxy_service): + with patch.object( + sandbox_proxy_service.sts_client, "do_action_with_exception", side_effect=Exception("sts fail") + ): + result = sandbox_proxy_service.gen_oss_sts_token() + assert result is None + + def test_partial_oss_config_returns_creds_with_none_extras(self, sandbox_proxy_service): + # endpoint not set, but STS still works + sandbox_proxy_service.oss_config.endpoint = "" + sandbox_proxy_service.oss_config.bucket = "bk" + + fake_token_body = ( + b'{"Credentials": {"AccessKeyId":"ak","AccessKeySecret":"sk",' + b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' + ) + with ( + patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, + ): + mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" + mock_env.ROCK_OSS_BUCKET_NAME = "" + mock_env.ROCK_OSS_BUCKET_REGION = "" + result = sandbox_proxy_service.gen_oss_sts_token() + + assert result["AccessKeyId"] == "ak" # STS still returns normally + assert result["Endpoint"] is None + assert result["Bucket"] == "bk" + assert result["Region"] is None + + def test_env_var_overrides_yaml_for_endpoint_and_bucket(self, sandbox_proxy_service): + """env_vars.ROCK_OSS_BUCKET_* take precedence over oss_config (YAML), aligned with client Layer 1.""" + sandbox_proxy_service.oss_config.endpoint = "yaml.endpoint" + sandbox_proxy_service.oss_config.bucket = "yaml-bucket" + + fake_token_body = ( + b'{"Credentials": {"AccessKeyId":"ak","AccessKeySecret":"sk",' + b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' + ) + with ( + patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, + ): + mock_env.ROCK_OSS_BUCKET_ENDPOINT = "env.endpoint" + mock_env.ROCK_OSS_BUCKET_NAME = "env-bucket" + mock_env.ROCK_OSS_BUCKET_REGION = "env-region" + result = sandbox_proxy_service.gen_oss_sts_token() + + # env wins everywhere + assert result["Endpoint"] == "env.endpoint" + assert result["Bucket"] == "env-bucket" + assert result["Region"] == "env-region" + + def test_yaml_used_when_env_var_empty(self, sandbox_proxy_service): + """Fall back to YAML when env is empty.""" + sandbox_proxy_service.oss_config.endpoint = "yaml.endpoint" + sandbox_proxy_service.oss_config.bucket = "yaml-bucket" + + fake_token_body = ( + b'{"Credentials": {"AccessKeyId":"ak","AccessKeySecret":"sk",' + b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' + ) + with ( + patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, + ): + mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" + mock_env.ROCK_OSS_BUCKET_NAME = "" + mock_env.ROCK_OSS_BUCKET_REGION = "rg" # only region is set + result = sandbox_proxy_service.gen_oss_sts_token() + + assert result["Endpoint"] == "yaml.endpoint" # YAML fallback + assert result["Bucket"] == "yaml-bucket" + assert result["Region"] == "rg" # env diff --git a/tests/unit/sdk/sandbox/test_file_system.py b/tests/unit/sdk/sandbox/test_file_system.py index 402da4dd9a..59421a4f82 100644 --- a/tests/unit/sdk/sandbox/test_file_system.py +++ b/tests/unit/sdk/sandbox/test_file_system.py @@ -2,6 +2,8 @@ from unittest.mock import AsyncMock, MagicMock +from rock.actions.sandbox.response import DownloadFileResponse +from rock.sdk.sandbox.client import Sandbox from rock.sdk.sandbox.file_system import LinuxFileSystem @@ -12,9 +14,60 @@ def _sandbox(exit_code=0): return sb +def _sandbox_with_oss(oss_setup_returns: bool = True, ossutil_ok: bool = True, download_response=None): + """Build a mock Sandbox with _oss + ensure_ossutil dependencies. + + Uses spec=Sandbox so isinstance(self.sandbox, Sandbox) check inside + download_file passes. + """ + sb = AsyncMock(spec=Sandbox) + sb.process = MagicMock() + sb.process.execute_script = AsyncMock(return_value=MagicMock(exit_code=0 if ossutil_ok else 1, output="")) + sb.execute = AsyncMock(return_value=MagicMock(exit_code=0 if ossutil_ok else 1, stdout="v2", stderr="")) + + sb._oss = MagicMock() + sb._oss.ensure_setup = AsyncMock(return_value=oss_setup_returns) + sb._oss.download_via_oss = AsyncMock( + return_value=download_response or DownloadFileResponse(success=True, message="ok") + ) + return sb + + class TestEnsureOssutil: async def test_success(self): assert await LinuxFileSystem(_sandbox()).ensure_ossutil() is True async def test_install_failure(self): assert await LinuxFileSystem(_sandbox(exit_code=1)).ensure_ossutil() is False + + +class TestDownloadFileDelegatesToOssClient: + async def test_delegates_to_oss_client_after_ensure_setup_and_ossutil(self, tmp_path): + sb = _sandbox_with_oss(oss_setup_returns=True, ossutil_ok=True) + + fs = LinuxFileSystem(sb) + resp = await fs.download_file("/sandbox/foo.txt", tmp_path / "foo.txt") + + assert resp.success is True + sb._oss.ensure_setup.assert_awaited_once() + sb._oss.download_via_oss.assert_awaited_once() + + async def test_returns_oss_unavailable_when_setup_fails(self, tmp_path): + sb = _sandbox_with_oss(oss_setup_returns=False) + + fs = LinuxFileSystem(sb) + resp = await fs.download_file("/sandbox/foo.txt", tmp_path / "foo.txt") + + assert resp.success is False + assert "OSS is not available" in resp.message + sb._oss.download_via_oss.assert_not_awaited() + + async def test_returns_failure_when_ossutil_install_fails(self, tmp_path): + sb = _sandbox_with_oss(oss_setup_returns=True, ossutil_ok=False) + + fs = LinuxFileSystem(sb) + resp = await fs.download_file("/sandbox/foo.txt", tmp_path / "foo.txt") + + assert resp.success is False + assert "ossutil" in resp.message + sb._oss.download_via_oss.assert_not_awaited() diff --git a/tests/unit/sdk/sandbox/test_oss_client.py b/tests/unit/sdk/sandbox/test_oss_client.py new file mode 100644 index 0000000000..361ff55efc --- /dev/null +++ b/tests/unit/sdk/sandbox/test_oss_client.py @@ -0,0 +1,492 @@ +"""Tests for OssClient — encapsulates all OSS operations for Sandbox.""" + +import asyncio +import re +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock import env_vars +from rock.actions.sandbox.response import DownloadFileResponse, UploadResponse +from rock.sdk.sandbox.oss_client import OssClient, OssClientConfig + + +def _make_sandbox(base_url="http://admin:8080", headers=None): + sb = MagicMock() + sb._url = base_url + sb._build_headers = MagicMock(return_value=headers or {}) + return sb + + +def test_oss_client_module_imports(): + assert OssClient is not None + assert OssClientConfig is not None + + +class TestComputeObjectName: + def test_format_is_hash_dash_filename(self): + name = OssClient._compute_object_name("sb-1", "/local/file.json", "/sandbox/file.json") + assert re.match(r"^[0-9a-f]{64}-file\.json$", name) + + def test_deterministic_same_inputs_same_output(self): + a = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/x") + b = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/x") + assert a == b + + def test_different_sandbox_id_yields_different_hash(self): + a = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/x") + b = OssClient._compute_object_name("sb-2", "/local/x", "/sandbox/x") + assert a != b + + def test_different_local_path_yields_different_hash(self): + a = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/x") + b = OssClient._compute_object_name("sb-1", "/local/y", "/sandbox/x") + assert a != b + + def test_different_sandbox_path_yields_different_hash(self): + a = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/x") + b = OssClient._compute_object_name("sb-1", "/local/x", "/sandbox/y") + assert a != b + + def test_filename_is_basename_of_sandbox_path(self): + name = OssClient._compute_object_name("sb-1", "/dir1/dir2/foo.txt", "/other/bar.txt") + assert name.endswith("-bar.txt") + + def test_filename_falls_back_to_local_path_basename_when_sandbox_empty(self): + name = OssClient._compute_object_name("sb-1", "/local/baz.txt", "") + assert name.endswith("-baz.txt") + + +class TestResolveConfig: + def test_layer1_env_takes_precedence_over_server(self): + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", "env.endpoint"), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", "env-bucket"), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", "env-region"), + ): + cfg = OssClient._resolve_config( + { + "Endpoint": "srv.endpoint", + "Bucket": "srv-bucket", + "Region": "srv-region", + } + ) + assert cfg.endpoint == "env.endpoint" + assert cfg.bucket == "env-bucket" + assert cfg.region == "env-region" + assert cfg.enabled_via_env is True + + def test_layer2_used_when_env_not_all_set(self): + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + ): + cfg = OssClient._resolve_config( + { + "Endpoint": "srv.endpoint", + "Bucket": "srv-bucket", + "Region": "srv-region", + } + ) + assert cfg.endpoint == "srv.endpoint" + assert cfg.enabled_via_env is False + + def test_partial_env_does_not_promote_to_layer1(self): + # Only endpoint set; bucket / region missing -> not Layer 1, fall back to Layer 2 + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", "env.endpoint"), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + ): + cfg = OssClient._resolve_config( + { + "Endpoint": "srv.endpoint", + "Bucket": "srv-bucket", + "Region": "srv-region", + } + ) + assert cfg.endpoint == "srv.endpoint" + assert cfg.enabled_via_env is False + + def test_layer3_returns_none_when_neither_layer_complete(self): + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + ): + cfg = OssClient._resolve_config({"Endpoint": None, "Bucket": None, "Region": None}) + assert cfg is None + + def test_server_partial_treated_as_unavailable(self): + # Server returns endpoint/bucket but no region -> Layer 2 incomplete -> unavailable + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + ): + cfg = OssClient._resolve_config({"Endpoint": "x", "Bucket": "y", "Region": None}) + assert cfg is None + + def test_empty_dict_returns_none(self): + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + ): + cfg = OssClient._resolve_config({}) + assert cfg is None + + +class TestGetStsCredentials: + async def test_success_returns_credentials_dict(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + + mock_response = { + "status": "Success", + "result": { + "AccessKeyId": "ak", + "AccessKeySecret": "sk", + "SecurityToken": "tok", + "Expiration": "2026-12-31T00:00:00Z", + "Endpoint": "endpoint", + "Bucket": "bucket", + "Region": "region", + }, + } + with patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http: + mock_http.get = AsyncMock(return_value=mock_response) + result = await client._get_sts_credentials() + + assert result["AccessKeyId"] == "ak" + assert result["Endpoint"] == "endpoint" + assert client._token_expire_time == "2026-12-31T00:00:00Z" + + async def test_failure_raises(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + with patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http: + mock_http.get = AsyncMock(return_value={"status": "Fail", "message": "boom"}) + with pytest.raises(Exception, match="boom"): + await client._get_sts_credentials() + + +class TestIsTokenExpired: + def test_no_token_means_expired(self): + client = OssClient(_make_sandbox()) + client._token_expire_time = None + assert client._is_token_expired() is True + + def test_future_expiration_not_expired(self): + client = OssClient(_make_sandbox()) + client._token_expire_time = "2099-01-01T00:00:00Z" + assert client._is_token_expired() is False + + def test_past_expiration_is_expired(self): + client = OssClient(_make_sandbox()) + client._token_expire_time = "2000-01-01T00:00:00Z" + assert client._is_token_expired() is True + + def test_within_5min_buffer_is_expired(self): + client = OssClient(_make_sandbox()) + # 1 minute in the future (within the 5-minute buffer) + near_future = (datetime.now(timezone.utc) + timedelta(minutes=1)).strftime("%Y-%m-%dT%H:%M:%SZ") + client._token_expire_time = near_future + assert client._is_token_expired() is True + + def test_attribute_error_is_treated_as_expired(self): + client = OssClient(_make_sandbox()) + client._token_expire_time = 12345 # int, no .replace method → AttributeError + assert client._is_token_expired() is True + + +class TestSetup: + async def test_layer1_env_with_enable_true(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", "env.endpoint"), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", "env-bucket"), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", "env-region"), + patch.object(env_vars, "ROCK_OSS_ENABLE", True), + patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http, + patch("rock.sdk.sandbox.oss_client.oss2") as mock_oss2, + ): + mock_http.get = AsyncMock( + return_value={ + "status": "Success", + "result": { + "AccessKeyId": "ak", + "AccessKeySecret": "sk", + "SecurityToken": "tok", + "Expiration": "2099-01-01T00:00:00Z", + }, + } + ) + mock_oss2.Bucket = MagicMock(return_value="bucket-instance") + ok = await client.ensure_setup() + + assert ok is True + assert client.is_available is True + mock_oss2.Bucket.assert_called_once() + kwargs = mock_oss2.Bucket.call_args.kwargs + assert kwargs["endpoint"] == "env.endpoint" + assert kwargs["bucket_name"] == "env-bucket" + + async def test_layer1_env_with_enable_false_returns_unavailable(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", "env.endpoint"), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", "env-bucket"), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", "env-region"), + patch.object(env_vars, "ROCK_OSS_ENABLE", False), + patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http, + ): + mock_http.get = AsyncMock( + return_value={ + "status": "Success", + "result": { + "AccessKeyId": "ak", + "AccessKeySecret": "sk", + "SecurityToken": "tok", + "Expiration": "2099-01-01T00:00:00Z", + }, + } + ) + ok = await client.ensure_setup() + + assert ok is False + assert client.is_available is False + + async def test_layer2_server_response_used_when_env_unset(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http, + patch("rock.sdk.sandbox.oss_client.oss2") as mock_oss2, + ): + mock_http.get = AsyncMock( + return_value={ + "status": "Success", + "result": { + "AccessKeyId": "ak", + "AccessKeySecret": "sk", + "SecurityToken": "tok", + "Expiration": "2099-01-01T00:00:00Z", + "Endpoint": "srv.endpoint", + "Bucket": "srv-bucket", + "Region": "srv-region", + }, + } + ) + mock_oss2.Bucket = MagicMock(return_value="bucket-instance") + ok = await client.ensure_setup() + + assert ok is True + assert client.is_available is True + kwargs = mock_oss2.Bucket.call_args.kwargs + assert kwargs["endpoint"] == "srv.endpoint" + + async def test_layer3_unavailable_when_neither_set(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + + with ( + patch.object(env_vars, "ROCK_OSS_BUCKET_ENDPOINT", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_NAME", ""), + patch.object(env_vars, "ROCK_OSS_BUCKET_REGION", ""), + patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http, + ): + mock_http.get = AsyncMock( + return_value={ + "status": "Success", + "result": { + "AccessKeyId": "ak", + "AccessKeySecret": "sk", + "SecurityToken": "tok", + "Expiration": "2099-01-01T00:00:00Z", + }, + } + ) + ok = await client.ensure_setup() + + assert ok is False + assert client.is_available is False + + async def test_ensure_setup_idempotent_when_token_valid(self): + sandbox = _make_sandbox() + client = OssClient(sandbox) + client._bucket = MagicMock() # already set up + client._token_expire_time = "2099-01-01T00:00:00Z" + + with patch("rock.sdk.sandbox.oss_client.HttpUtils") as mock_http: + mock_http.get = AsyncMock() + ok = await client.ensure_setup() + + assert ok is True + mock_http.get.assert_not_called() # /get_token not re-invoked + + +class TestUploadViaOss: + async def test_success(self): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-123" + sandbox.arun = AsyncMock(return_value=MagicMock(exit_code=0)) + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=0)) + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._bucket.sign_url = MagicMock(return_value="https://oss/signed?...") + + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload") as mock_upload: + response = await client.upload_via_oss("/local/foo.json", "/sandbox/dst/foo.json") + + assert isinstance(response, UploadResponse) + assert response.success is True + # Verify OSS object naming follows the new convention: sha256-filename + expected_obj = OssClient._compute_object_name("sb-123", "/local/foo.json", "/sandbox/dst/foo.json") + mock_upload.assert_called_once() + # resumable_upload(bucket, obj_name, file_path) + assert mock_upload.call_args.args[1] == expected_obj + assert mock_upload.call_args.args[2] == "/local/foo.json" + + async def test_sandbox_verification_fail_returns_failure(self): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-123" + # mkdir/wget go through arun and succeed; the final test -f check goes + # through execute() and fails (exit_code=1 = file missing). + sandbox.arun = AsyncMock(return_value=MagicMock(exit_code=0)) + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=1)) + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._bucket.sign_url = MagicMock(return_value="url") + + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload"): + response = await client.upload_via_oss("/local/foo.json", "/sandbox/dst/foo.json") + + assert response.success is False + assert "sandbox download phase failed" in response.message + + async def test_oss_upload_exception_returns_failure(self): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-123" + client = OssClient(sandbox) + client._bucket = MagicMock() + + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload", side_effect=Exception("oss boom")): + response = await client.upload_via_oss("/local/foo.json", "/sandbox/dst/foo.json") + + assert response.success is False + + +class TestDownloadViaOss: + async def test_oss_unavailable_returns_failure(self, tmp_path): + sandbox = _make_sandbox() + client = OssClient(sandbox) + # _bucket is still None + response = await client.download_via_oss("/sandbox/foo.txt", tmp_path / "foo.txt") + assert isinstance(response, DownloadFileResponse) + assert response.success is False + assert "OSS is not available" in response.message + + async def test_remote_file_not_found_returns_failure(self, tmp_path): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-1" + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=1)) # test -f fails + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._client_config = OssClientConfig("ep", "bk", "rg", enabled_via_env=False) + + response = await client.download_via_oss("/sandbox/foo.txt", tmp_path / "foo.txt") + assert response.success is False + assert "not found" in response.message.lower() + + +class TestClose: + async def test_close_with_no_pending_tasks_is_noop(self): + client = OssClient(_make_sandbox()) + await client.close() # should not raise + + +class TestScheduleAsyncPersistence: + async def test_schedules_task_when_available(self): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-1" + client = OssClient(sandbox) + client._bucket = MagicMock() + + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_to_thread: + mock_to_thread.return_value = None + key = await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") + # wait for task to complete (must be within the patch scope) + await asyncio.gather(*client._pending_persistence_tasks, return_exceptions=True) + mock_to_thread.assert_awaited_once() + + assert key.endswith("-foo.json") + + async def test_no_op_when_unavailable(self): + client = OssClient(_make_sandbox()) + # _bucket is still None + key = await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") + assert key is None + assert len(client._pending_persistence_tasks) == 0 + + async def test_failure_does_not_raise(self): + """OSS upload failure is swallowed; main flow is unaffected.""" + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-1" + client = OssClient(sandbox) + client._bucket = MagicMock() + + with patch("asyncio.to_thread", new_callable=AsyncMock) as mock_to_thread: + mock_to_thread.side_effect = Exception("oss boom") + key = await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") + # wait for task to complete (must not raise) + await asyncio.gather(*client._pending_persistence_tasks, return_exceptions=True) + + assert key is not None + + +class TestCloseAwaitsPendingTasks: + async def test_close_awaits_completion(self): + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-1" + client = OssClient(sandbox) + client._bucket = MagicMock() + + completion_marker = asyncio.Event() + + async def slow_upload(*args, **kwargs): + await asyncio.sleep(0.1) + completion_marker.set() + + with patch("asyncio.to_thread", new=slow_upload): + await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") + assert not completion_marker.is_set() + await client.close() + assert completion_marker.is_set() + + async def test_close_timeout_does_not_hang(self): + """close does not hang or raise on timeout (pending tasks beyond timeout are abandoned).""" + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-1" + client = OssClient(sandbox) + client._bucket = MagicMock() + + async def hang(*args, **kwargs): + await asyncio.sleep(100) + + with patch("asyncio.to_thread", new=hang): + await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") + await client.close(timeout=0.05) diff --git a/tests/unit/sdk/sandbox/test_upload_by_path.py b/tests/unit/sdk/sandbox/test_upload_by_path.py new file mode 100644 index 0000000000..20e21c34ab --- /dev/null +++ b/tests/unit/sdk/sandbox/test_upload_by_path.py @@ -0,0 +1,104 @@ +"""Test Sandbox.upload_by_path async OSS persistence integration.""" + +from unittest.mock import AsyncMock, MagicMock, patch + + +async def test_small_file_triggers_async_persistence(tmp_path): + """Small admin /upload success path schedules async OSS persistence.""" + f = tmp_path / "small.txt" + f.write_text("hi") + + from rock.sdk.sandbox.client import Sandbox + from rock.sdk.sandbox.config import SandboxConfig + + sandbox = Sandbox(SandboxConfig(base_url="http://x")) + # Replace the whole _oss with a MagicMock so we can set is_available + # (it's a @property on the real OssClient and not assignable). + sandbox._oss = MagicMock() + sandbox._oss.ensure_setup = AsyncMock(return_value=True) + sandbox._oss.is_available = True + sandbox._oss.schedule_async_persistence = AsyncMock(return_value="hash-small.txt") + + with patch( + "rock.sdk.sandbox.client.HttpUtils.post_multipart", + AsyncMock(return_value={"status": "Success"}), + ): + response = await sandbox.upload_by_path(str(f), "/sandbox/small.txt") + + assert response.success is True + sandbox._oss.ensure_setup.assert_awaited() + sandbox._oss.schedule_async_persistence.assert_awaited_once_with(str(f), "/sandbox/small.txt") + + +async def test_small_file_no_persistence_when_oss_unavailable(tmp_path): + """Admin /upload success but OSS unavailable: no persistence scheduled.""" + f = tmp_path / "small.txt" + f.write_text("hi") + + from rock.sdk.sandbox.client import Sandbox + from rock.sdk.sandbox.config import SandboxConfig + + sandbox = Sandbox(SandboxConfig(base_url="http://x")) + sandbox._oss = MagicMock() + sandbox._oss.ensure_setup = AsyncMock(return_value=False) + sandbox._oss.is_available = False + sandbox._oss.schedule_async_persistence = AsyncMock() + + with patch( + "rock.sdk.sandbox.client.HttpUtils.post_multipart", + AsyncMock(return_value={"status": "Success"}), + ): + response = await sandbox.upload_by_path(str(f), "/sandbox/small.txt") + + assert response.success is True + sandbox._oss.schedule_async_persistence.assert_not_awaited() + + +async def test_failed_upload_no_persistence(tmp_path): + """Admin /upload failure: persistence must NOT be scheduled.""" + f = tmp_path / "small.txt" + f.write_text("hi") + + from rock.sdk.sandbox.client import Sandbox + from rock.sdk.sandbox.config import SandboxConfig + + sandbox = Sandbox(SandboxConfig(base_url="http://x")) + sandbox._oss = MagicMock() + sandbox._oss.ensure_setup = AsyncMock(return_value=True) + sandbox._oss.is_available = True + sandbox._oss.schedule_async_persistence = AsyncMock() + + with patch( + "rock.sdk.sandbox.client.HttpUtils.post_multipart", + AsyncMock(return_value={"status": "Failed", "message": "boom"}), + ): + response = await sandbox.upload_by_path(str(f), "/sandbox/small.txt") + + assert response.success is False + sandbox._oss.schedule_async_persistence.assert_not_awaited() + + +async def test_sandbox_close_awaits_oss_pending_tasks(): + """Sandbox.close() must await OssClient.close() so pending persistence + tasks have a chance to finish before the sandbox lifecycle ends.""" + from rock.sdk.sandbox.client import Sandbox + from rock.sdk.sandbox.config import SandboxConfig + + sandbox = Sandbox(SandboxConfig(base_url="http://x")) + sandbox._oss = MagicMock() + sandbox._oss.close = AsyncMock() + + # stop() would otherwise issue a real HTTP request; replace with a noop. + # Track call order to verify _oss.close runs BEFORE stop (so pending OSS + # tasks aren't aborted by sandbox teardown). + call_order: list[str] = [] + sandbox._oss.close.side_effect = lambda: call_order.append("oss_close") + + async def fake_stop(): + call_order.append("stop") + + with patch.object(sandbox, "stop", AsyncMock(side_effect=fake_stop)): + await sandbox.close() + + sandbox._oss.close.assert_awaited_once() + assert call_order == ["oss_close", "stop"] From aae3e9788edc9b68c7d87cd4ecad95a28c21e02b Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 15 May 2026 14:44:35 +0800 Subject: [PATCH 099/226] chore(sdk): make sandbox cluster default configurable via env var Introduce ROCK_DEFAULT_CLUSTER env var (default: vpc-nt-a) and use it as the default for SandboxConfig.cluster, replacing the hardcoded "zb". Refs #947 --- rock/env_vars.py | 2 ++ rock/sdk/sandbox/config.py | 2 +- tests/unit/sdk/agent/test_job_config_serialization.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/rock/env_vars.py b/rock/env_vars.py index 0d9f169395..7844da9c62 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -14,6 +14,7 @@ ROCK_CONFIG: str | None = None ROCK_CONFIG_DIR_NAME: str | None = None ROCK_BASE_URL: str | None = "http://localhost:8080" + ROCK_DEFAULT_CLUSTER: str = "vpc-nt-a" ROCK_WORKER_ROCKLET_PORT: int | None = None ROCK_SANDBOX_STARTUP_TIMEOUT_SECONDS: int = 180 ROCK_CODE_SANDBOX_BASE_URL: str | None = None @@ -77,6 +78,7 @@ "ROCK_CONFIG": lambda: os.getenv("ROCK_CONFIG"), "ROCK_CONFIG_DIR_NAME": lambda: os.getenv("ROCK_CONFIG_DIR_NAME", "rock-conf"), "ROCK_BASE_URL": lambda: os.getenv("ROCK_BASE_URL", "http://localhost:8080"), + "ROCK_DEFAULT_CLUSTER": lambda: os.getenv("ROCK_DEFAULT_CLUSTER", "vpc-nt-a"), "ROCK_WORKER_ROCKLET_PORT": lambda: int(val) if (val := os.getenv("ROCK_WORKER_ROCKLET_PORT")) else None, "ROCK_SANDBOX_STARTUP_TIMEOUT_SECONDS": lambda: int(os.getenv("ROCK_SANDBOX_STARTUP_TIMEOUT_SECONDS", "180")), "ROCK_CODE_SANDBOX_BASE_URL": lambda: os.getenv("ROCK_CODE_SANDBOX_BASE_URL", ""), diff --git a/rock/sdk/sandbox/config.py b/rock/sdk/sandbox/config.py index 130aaf40a8..4fcf59e030 100644 --- a/rock/sdk/sandbox/config.py +++ b/rock/sdk/sandbox/config.py @@ -38,7 +38,7 @@ class SandboxConfig(BaseConfig): limit_cpus: float | None = None user_id: str | None = None experiment_id: str | None = None - cluster: str = "zb" + cluster: str = env_vars.ROCK_DEFAULT_CLUSTER namespace: str | None = None registry_username: str | None = None registry_password: str | None = None diff --git a/tests/unit/sdk/agent/test_job_config_serialization.py b/tests/unit/sdk/agent/test_job_config_serialization.py index fbccc3d833..71c4690992 100644 --- a/tests/unit/sdk/agent/test_job_config_serialization.py +++ b/tests/unit/sdk/agent/test_job_config_serialization.py @@ -26,7 +26,7 @@ def test_inherits_sandbox_config_fields(self): assert env.image == "python:3.11" assert env.memory == "8g" assert env.cpus == 2.0 - assert env.cluster == "zb" + assert env.cluster == "vpc-nt-a" def test_inherits_harbor_env_fields(self): env = RockEnvironmentConfig() From 3537359e647aaca947a4833b8ce91dbb59de662e Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 15 May 2026 17:08:58 +0800 Subject: [PATCH 100/226] test(sdk): fix test_sandbox_cluster by moving it to its own class Co-Authored-By: Claude Opus 4.7 --- tests/unit/sdk/test_sandbox_config.py | 31 +++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/unit/sdk/test_sandbox_config.py b/tests/unit/sdk/test_sandbox_config.py index 1480965f91..f677c7d35e 100644 --- a/tests/unit/sdk/test_sandbox_config.py +++ b/tests/unit/sdk/test_sandbox_config.py @@ -1,6 +1,7 @@ import pytest from pydantic import ValidationError +from rock.sdk.sandbox.client import Sandbox from rock.sdk.sandbox.config import SandboxConfig @@ -28,3 +29,33 @@ def test_negative_value_raises_error(self): def test_large_negative_value_raises_error(self): with pytest.raises(ValidationError, match="auto_delete_seconds must be >= 0"): SandboxConfig(auto_delete_seconds=-100) + + +class TestSandboxCluster: + def test_sandbox_cluster(self): + fake_route_key = "fake_route_key" + fake_auth_token = "fake_auth_token" + config = SandboxConfig( + image="python:3.11", + route_key=fake_route_key, + xrl_authorization=fake_auth_token, + cluster="sg", + ) + sandbox = Sandbox(config) + assert sandbox.cluster == "sg" + common_headers = sandbox._build_headers() + assert common_headers["X-Cluster"] == "sg" + assert common_headers["ROUTE-KEY"] == fake_route_key + assert common_headers["XRL-Authorization"] == f"Bearer {fake_auth_token}" + + # default cluster is vpc-nt-a + config = SandboxConfig( + image="python:3.11", + route_key=fake_route_key, + xrl_authorization=fake_auth_token, + ) + + sandbox = Sandbox(config) + assert sandbox.cluster == "vpc-nt-a" + common_headers = sandbox._build_headers() + assert common_headers["X-Cluster"] == "vpc-nt-a" From 7c2118fc97c404984c35cfb88c01864bfd1ea58c Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Sun, 17 May 2026 16:16:20 +0800 Subject: [PATCH 101/226] docs(scheduler): add scheduler user guide for v1.7.x #974 (#975) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * add scheduler doc * fix format --- .../version-1.7.x/User Guides/scheduler.md | 283 ++++++++++++++++++ .../version-1.7.x/User Guides/scheduler.md | 283 ++++++++++++++++++ rock-conf/rock-local.yml | 42 +++ rock/env_vars.py | 2 + rock/utils/system.py | 5 + 5 files changed, 615 insertions(+) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/scheduler.md create mode 100644 docs/versioned_docs/version-1.7.x/User Guides/scheduler.md diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/scheduler.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/scheduler.md new file mode 100644 index 0000000000..88662edd51 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.7.x/User Guides/scheduler.md @@ -0,0 +1,283 @@ +--- +sidebar_position: 5 +--- + +# 任务调度器(Scheduler) + +ROCK 调度器是内嵌于 `admin` 服务中的周期性任务框架。它会按可配置的时间间隔,把后台维护任务(镜像清理、文件清理、容器清理、镜像预拉取、自定义任务……)分发到所有存活的 Ray worker 上,从而在无人工干预的情况下保持 worker 节点健康。 + +本文介绍如何启用调度器、配置内置任务、编写自定义任务,以及如何观测运行状态。 + +## 1. 工作原理 + +- 调度器以独立的守护线程(`SchedulerThread`)运行在 `admin` 进程内,使用自己的 `asyncio` 事件循环。 +- 任务通过 [APScheduler](https://apscheduler.readthedocs.io/) 以固定间隔(`interval_seconds`)触发。 +- 每次触发时,调度器会先获取存活 Ray worker 列表(由 `worker_cache_ttl` 秒级缓存),然后并发地把任务下发到每个 worker(默认并发度:50)。 +- **下发动作通过 worker 上的 rocklet 服务以 HTTP 方式完成**:admin 端构造 `RemoteSandboxRuntime(host=worker_ip, port=Port.PROXY)`(参见 `rock.deployments.constants.Port`),并调用 `runtime.execute / read_file / write_file`。**因此每个 worker 都必须运行 `rocklet` 服务并在 `Port.PROXY` 上可达**,否则调度器无法下发命令、也无法在 worker 上读写状态文件。 +- 每个任务都继承自 `rock.admin.scheduler.task_base.BaseTask`,并必须实现 `run_action(runtime: RemoteSandboxRuntime)` —— 单个 worker 上的实际执行逻辑。 +- 每个 worker 的执行状态会被持久化到 `ROCK_SCHEDULER_STATUS_DIR`(默认 `/data/scheduler_status`)目录下的 JSON 文件中;每次执行结束后还会写入聚合报告 `/_run_report.json`。 +- 当配置了 Nacos 配置源时,调度器会订阅配置变更,并按 diff 应用:仅 hash 发生变化的任务被重新安装,被删除的任务会同步从所有 worker 上清理。 + +### 前置条件 + +启用调度器之前,请确认每个 Ray worker 满足以下条件: + +| 条件 | 原因 | +|------|------| +| worker 上正在运行 `rocklet` 进程 | 调度器通过 rocklet HTTP 接口下发每一个任务;若 rocklet 不存在,`runtime.execute` 调用会超时。 | +| admin 可访问 rocklet 监听端口 | 调度器固定使用 `Port.PROXY`(定义于 `rock.deployments.constants.Port`)作为下发目标,请确认防火墙 / 安全组未阻断该端口。 | +| `ROCK_SCHEDULER_STATUS_DIR` 在 worker 内可写 | 任务会在该目录下读写 `_status.json`,用于幂等控制和 PID 跟踪。 | +| 任务依赖的工具在 worker 上可用 | 例如清理 / 拉取类任务需要 `docker`;`ImageCleanupTask` 首次运行会通过 `curl` 联网安装 `docuum`。 | + +rocklet 服务由 worker 标准启动脚本(`docker_run.sh`、`docker_run_with_uv.sh`、`docker_run_with_pip.sh`)自动拉起,通常等价于 `rocklet --port `。如果你使用了自定义 entrypoint 来启动 worker,请确保等价命令被执行。具体的运行时类型与 rocklet 启动方式可参考 [Configuration](./configuration.md)。 + +### 幂等性 + +每个任务都需要声明自己的幂等模式,该模式直接影响重复触发时的行为: + +| 模式 | 行为 | +|------|------| +| `IDEMPOTENT` | 每次 tick 都会执行,可安全重复(例如 `docker pull`、`find -exec rm`)。 | +| `NON_IDEMPOTENT` | 任务会拉起一个后台守护进程(例如 `docuum`)。调度器会读取上一次的状态文件,检查记录的 PID 是否仍存活,若仍在运行则跳过本次启动。当任务从配置中移除时,调度器会通过 `pkill` 杀掉该进程。 | + +## 2. 启用调度器 + +调度器配置位于 ROCK admin YAML 顶层 `scheduler:` 字段下(例如 `rock-conf/rock-local.yml`、`rock-conf/rock-dev.yml`)。 + +```yaml +scheduler: + enabled: true # 总开关 + worker_cache_ttl: 43200 # Worker IP 缓存 TTL(秒) + tasks: + # ... 任务列表,详见下文 +``` + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `enabled` | bool | `false` | 总开关。设为 `false` 时所有任务会被卸载,且不再触发。 | +| `worker_cache_ttl` | int | `3600` | 存活 worker IP 列表缓存时长(秒);超过该时长后会从 `ray.nodes()` 重新获取。 | +| `tasks` | list | `[]` | `TaskConfig` 列表,详见 [第 4 节](#4-任务配置schema)。 | + +### 相关环境变量 + +| 变量 | 默认值 | 用途 | +|------|--------|------| +| `ROCK_SCHEDULER_STATUS_DIR` | `/data/scheduler_status` | worker 上写入任务状态 JSON 与执行报告的目录。 | +| `ROCK_LOGGING_PATH` | (未设置) | 设置后,调度器拉起的守护进程(docuum、container_cleanup、image_pull)会将 stdout/stderr 重定向到 `/.log`。 | +| `ROCK_DOCUUM_INSTALL_URL` | `https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh` | `ImageCleanupTask` 按需拉取 `docuum` 安装脚本的 URL。 | + +## 3. 内置任务 + +ROCK 在 `rock.admin.scheduler.tasks` 下提供了 4 个内置任务,通过把 `task_class` 设置为对应的全限定类路径即可注册。 + +### 3.1 ImageCleanupTask + +在每个 worker 上运行 [`docuum`](https://github.com/stepchowfun/docuum),当磁盘占用超过阈值时按 LRU 策略淘汰镜像。**非幂等** —— `docuum` 是常驻守护进程;调度器会跟踪其 PID,只要进程仍存活就跳过重复拉起。 + +```yaml +- task_class: rock.admin.scheduler.tasks.image_cleanup_task.ImageCleanupTask + enabled: true + interval_seconds: 43200 # 每 12 小时检查一次守护进程 + params: + disk_threshold: "70%" # 磁盘占用超过 70% 时触发淘汰 + image_whitelist: # 匹配 repository:tag 的 glob 模式,白名单内的镜像不会被淘汰 + - "python:3.11" + - "my-registry.example.com/base/*" +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `disk_threshold` | str | `"1T"` | 传给 `docuum --threshold` 的磁盘阈值,支持容量(`100G`、`1T`)或百分比(`70%`)。 | +| `image_whitelist` | list[str] | `[]` | 透传给 `docuum --keep` 的 glob 模式列表。 | + +### 3.2 FileCleanupTask + +遍历配置的目录,删除超过 `max_age_mins` 或大于 `max_file_size` 的文件,然后清理留下的空目录。**幂等**。 + +```yaml +- task_class: rock.admin.scheduler.tasks.file_cleanup_task.FileCleanupTask + enabled: true + interval_seconds: 86400 # 每天执行一次 + params: + target_dirs: + # 字符串形式 —— 不配置排除项 + - "/data/service_status" + # 对象形式 —— 配置该目录独有的排除项 + - path: "/data/logs" + exclude_files: # 支持纯文件名 / 相对路径 / 绝对路径 + - "docuum.log" + - "./rocklet.log" + - "./access.log" + exclude_dirs: + - ".cache" + max_age_mins: 10080 # 7 天,超出此时间的文件会被删除 + max_file_size: "1G" # 大于此大小的文件会被删除(支持 K/M/G/T) +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `target_dirs` | list | `[]` | 每个条目是一个字符串(只填路径)或 `{path, exclude_files, exclude_dirs}`。 | +| `max_age_mins` | int | `10080` | mtime 早于该分钟数的文件会被删除。 | +| `max_file_size` | str | `"1G"` | 大于该阈值的文件会被删除,支持 `K/M/G/T` 单位。 | + +删除条件为 `(-mmin +max_age_mins) OR (-size +max_file_size)`。文件清理后,会再用 `find -depth -type d -empty -delete` 清理留下的空目录(同样遵循 `exclude_dirs` 配置)。 + +### 3.3 ContainerCleanupTask + +删除停止时间超过指定时长的 Docker 容器,避免 worker 上的容器列表无限增长。**幂等**。 + +```yaml +- task_class: rock.admin.scheduler.tasks.container_cleanup_task.ContainerCleanupTask + enabled: true + interval_seconds: 86400 + params: + max_age_hours: 72 # 删除超过 72 小时的 exited 容器 +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `max_age_hours` | int | `24` | 已退出容器的最大保留时间(以 `FinishedAt` 起算的小时数);超出后被 `docker rm`。 | + +每次执行还会顺带清理处于 `created` 状态(从未启动)的容器。 + +### 3.4 ImagePullTask + +在每个 worker 上预拉取一组 Docker 镜像,并可选地先登录私有仓库,以降低沙箱冷启动延迟。**幂等**(若镜像已是最新,`docker pull` 等同于空操作)。 + +```yaml +- task_class: rock.admin.scheduler.tasks.image_pull_task.ImagePullTask + enabled: true + interval_seconds: 21600 # 每 6 小时刷新一次 + params: + images: + # 字符串形式 —— 公开镜像,无需鉴权 + - "python:3.11" + # 对象形式 —— 私有镜像,需要登录 + - image: "my-registry.example.com/chatos/python:313" + registry_username: "myuser" + registry_password: "bXlwYXNzd29yZA==" # base64 编码 +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `images` | list | `[]` | 每个条目是一个镜像字符串或 `{image, registry_username, registry_password}`。 | + +`registry_password` 必须是 base64 编码,worker 端会先解码再通过 `docker login --password-stdin` 登录。仓库地址会从镜像名称中解析,因此每个镜像可以指向不同的仓库。 + +## 4. 任务配置 Schema + +`scheduler.tasks` 下的每一项都会被解析为 `rock.config.TaskConfig`: + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `task_class` | str | `""` | Python 类的全限定路径,**必填**。 | +| `enabled` | bool | `true` | 设为 `false` 的任务在加载阶段会被跳过,在 reload 阶段会被卸载。 | +| `interval_seconds` | int | `3600` | APScheduler `interval` 间隔(秒)。 | +| `params` | dict | `{}` | 任务自定义参数,会在 `from_config()` 中被消费。 | + +只要某个任务条目中任何字段发生变化,调度器就会卸载旧任务(非幂等任务还会清理 worker 上的守护进程与状态文件)再安装新任务,**整个过程不需要重启 admin 进程**。 + +## 5. 编写自定义任务 + +任何位于 Python 路径下、继承自 `BaseTask` 的类都可以注册为调度任务。最小契约示例如下: + +```python +# my_pkg/my_tasks/disk_report_task.py +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + + +class DiskReportTask(BaseTask): + """记录每个 worker 的 `df -h` 输出。""" + + def __init__(self, interval_seconds: int = 3600, mount_point: str = "/"): + super().__init__( + type="disk_report", # 同时作为 APScheduler job id 与状态文件名前缀 + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + self.mount_point = mount_point + + @classmethod + def from_config(cls, task_config) -> "DiskReportTask": + return cls( + interval_seconds=task_config.interval_seconds, + mount_point=task_config.params.get("mount_point", "/"), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + result = await runtime.execute( + Command(command=f"df -h {self.mount_point}", shell=True), + ) + return { + "status": TaskStatusEnum.SUCCESS, + "exit_code": result.exit_code, + "stdout": result.stdout, + } +``` + +随后在 YAML 中注册: + +```yaml +scheduler: + enabled: true + tasks: + - task_class: my_pkg.my_tasks.disk_report_task.DiskReportTask + enabled: true + interval_seconds: 600 + params: + mount_point: "/data" +``` + +### 自定义任务编写要点 + +- **`super().__init__()` 中的 `type` 必须全局唯一**:它会同时被用作 APScheduler job id、状态文件名(`_status.json`)与执行报告文件名(`_run_report.json`)。两个任务不能共用同一个 `type`。 +- **正确选择 `IdempotencyType`**: + - 当 `run_action` 同步执行完毕、可安全重入时,使用 `IDEMPOTENT`。 + - 当任务通过 `nohup` 拉起常驻守护进程并返回 PID 时,使用 `NON_IDEMPOTENT`;调度器会跟踪 PID,在其存活期间跳过重复启动,在卸载时通过 `pkill` 杀掉。 +- **`run_action` 必须返回 dict**,推荐字段: + - `status` —— `TaskStatusEnum` 值,会写入状态文件。 + - `pid` —— `NON_IDEMPOTENT` 守护型任务必须返回(可使用 `rock.utils.system.extract_nohup_pid` 从 `nohup ... & echo PID_PREFIX${{!}}PID_SUFFIX` 输出中提取)。 + - 其他诊断字段会落到状态文件的 `extra` 块里。 +- **重写 `from_config(cls, task_config)`** 用于把 `task_config.params` 翻译成 `__init__` 的入参。 +- **使用 `runtime.execute / read_file / write_file`** 与 worker 通信,**不要**在本地直接执行 shell —— 调度器是把任务下发到远端 worker 的 `RemoteSandboxRuntime`。 + +## 6. 可观测性 + +每个任务会在每个 worker 上产生两类产物: + +| 路径 | 写入方 | 内容 | +|------|--------|------| +| `/_status.json` | `BaseTask.save_task_status` | 单 worker 最新状态:`task_name`、`worker_ip`、`pid`、`status`(`pending`/`running`/`success`/`failed`)、`last_run`、`error`,以及任务自定义的 `extra` 字段。 | +| `/_run_report.json` | `BaseTask.run`(由 admin 端在本轮 tick 结束后写入) | 聚合报告:总数 / 成功数 / 失败数、`success_ips` 列表、`failed_details`(`ip` + 错误堆栈)。 | + +调度器内部日志会输出到 ROCK admin 标准日志路径下,logger 名称包括 `name="scheduler"`、`name="task_base"`、`name="image_clean"` 等。当设置了 `ROCK_LOGGING_PATH` 时,被调度器拉起的守护进程会把自身日志写入 `/.log`(例如 `docuum.log`、`container_cleanup.log`、`image_pull.log`)。 + +## 7. 通过 Nacos 动态热更(可选) + +当 admin 服务启用了 Nacos 配置源时,调度器会注册一个 YAML 监听器,并对配置推送做出响应: + +- 仅检查 `scheduler:` 段,其他段被忽略。 +- 新配置段会被计算 hash 并与上一次的 hash 比对 —— 重复推送会被自动跳过。 +- 通过对新旧任务列表 diff,决定哪些任务需要安装、卸载或重新安装(`params` / `interval_seconds` / `enabled` 任意改动都会触发)。 +- 被删除或重新安装的非幂等任务会先做清理:杀掉守护进程 PID 并删除状态文件。 + +由此,任务间隔调整、参数微调、增删任务等都可以在 admin 进程不重启的前提下生效。 + +## 8. 常见问题排查 + +| 现象 | 可能原因 | 检查项 | +|------|----------|--------| +| admin 日志输出 `Scheduler disabled, all tasks removed` | `scheduler.enabled` 为 `false` | 把 YAML 中的 `enabled` 改为 `true`。 | +| `No alive workers found for task ''` | Ray 集群没有存活的 worker | 确认 `ray.nodes()` 返回的 CPU worker 处于 alive 状态;新加 worker 时可调小 `worker_cache_ttl`。 | +| 任务到点触发,但每个 worker 都进入 `failed_details` 且报连接错误 | worker 上 rocklet 未运行,或 `Port.PROXY` 被防火墙阻断 | 在 worker 上访问 rocklet 存活探针 `GET /is_alive`(例如 `curl http://:/is_alive`);若无响应,使用 `rocklet --port ` 重启或在防火墙放行该端口。 | +| 非幂等任务在某个 worker 上始终不再触发 | 状态文件中记录的 PID 仍存活 | 查看 `/_status.json`,如 `status: running` 且 PID 仍在,则 `should_run` 会返回 `False`,属预期行为。 | +| 日志报 `Failed to create task ''` | `task_class` 导入失败 | 确认对应模块在 admin 进程中可被导入(同一个 venv、`PYTHONPATH` 中可见)。 | +| `ImagePullTask` 中 `docker login` 失败 | `registry_password` 未做 base64 编码,或镜像名称解析出的仓库地址有误 | 用 `echo -n '' \| base64` 重新编码;确认镜像名中的仓库 host 正确。 | + +## 相关文档 + +- [Configuration](./configuration.md) —— 环境变量与运行时部署说明 +- [API Documentation](../References/api.md) —— admin HTTP API +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) —— SDK 编程式使用 diff --git a/docs/versioned_docs/version-1.7.x/User Guides/scheduler.md b/docs/versioned_docs/version-1.7.x/User Guides/scheduler.md new file mode 100644 index 0000000000..68a63cb427 --- /dev/null +++ b/docs/versioned_docs/version-1.7.x/User Guides/scheduler.md @@ -0,0 +1,283 @@ +--- +sidebar_position: 5 +--- + +# Scheduler + +The ROCK scheduler is a periodic task framework embedded in the `admin` service. It dispatches background maintenance tasks (image cleanup, file cleanup, container cleanup, image pre-pull, custom tasks, ...) to every alive Ray worker on a configurable interval, so worker nodes stay healthy without manual intervention. + +This guide covers how to enable the scheduler, configure built-in tasks, write your own task, and inspect runtime status. + +## 1. How It Works + +- The scheduler runs inside the `admin` process as a dedicated daemon thread (`SchedulerThread`) with its own `asyncio` event loop. +- Tasks are scheduled by [APScheduler](https://apscheduler.readthedocs.io/) using fixed intervals (`interval_seconds`). +- For each tick, the scheduler resolves the list of alive Ray workers (cached via `worker_cache_ttl` seconds) and dispatches the task to every worker concurrently (default concurrency: 50). +- Dispatch is done over HTTP through the worker's **rocklet** service: the admin builds a `RemoteSandboxRuntime(host=worker_ip, port=Port.PROXY)` (see `rock.deployments.constants.Port`) and calls `runtime.execute / read_file / write_file` against it. **Every worker must therefore have the `rocklet` server running and reachable on `Port.PROXY`** — otherwise the scheduler cannot push commands or read/write status files on that worker. +- Each task subclasses `rock.admin.scheduler.task_base.BaseTask` and must implement `run_action(runtime: RemoteSandboxRuntime)` — the work performed on a single worker. +- Per-worker execution status is persisted to the worker filesystem under `ROCK_SCHEDULER_STATUS_DIR` (default `/data/scheduler_status`), and an aggregated execution report is written to `/_run_report.json` after every run. +- If a Nacos config provider is enabled, the scheduler subscribes to config changes and applies a diff: only tasks whose hash changed are re-installed; removed tasks are cleaned up from all workers. + +### Prerequisites + +Before enabling the scheduler, make sure each Ray worker meets the following requirements: + +| Requirement | Why | +|-------------|-----| +| `rocklet` process is running on the worker | The scheduler dispatches every task through the rocklet HTTP API; without it, `runtime.execute` calls time out. | +| Rocklet's listening port is reachable from the admin | The scheduler uses `Port.PROXY` (defined in `rock.deployments.constants.Port`) as the dispatch target. Make sure no firewall / security group blocks it. | +| `ROCK_SCHEDULER_STATUS_DIR` is writable inside the worker | Tasks read and write `_status.json` here for idempotency / PID tracking. | +| Tools required by the task are available on the worker | e.g. `docker` for the cleanup / pull tasks, `curl` and outbound network for `ImageCleanupTask` to install `docuum` on first run. | + +The rocklet server is started automatically by the standard worker bootstrap scripts (`docker_run.sh`, `docker_run_with_uv.sh`, `docker_run_with_pip.sh`) — typically `rocklet --port `. If you bring up workers with a custom entrypoint, ensure the equivalent command is invoked. See the [Configuration](./configuration.md) guide for the runtime-environment options that govern how rocklet is started. + +### Idempotency + +Each task declares an idempotency mode that affects how it is re-run: + +| Mode | Behavior | +|------|----------| +| `IDEMPOTENT` | Always run on every tick. Safe to repeat (e.g. `docker pull`, `find -exec rm`). | +| `NON_IDEMPOTENT` | Spawns a background daemon (e.g. `docuum`). The scheduler reads the previous status file, checks whether the recorded PID is still alive, and skips re-launch if the daemon is still running. On task removal the daemon is killed via `pkill`. | + +## 2. Enabling the Scheduler + +The scheduler is configured under the top-level `scheduler:` key of the ROCK admin YAML (e.g. `rock-conf/rock-local.yml`, `rock-conf/rock-dev.yml`). + +```yaml +scheduler: + enabled: true # Master switch + worker_cache_ttl: 43200 # Worker IP cache TTL in seconds + tasks: + # ... task list, see below +``` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `enabled` | bool | `false` | Master switch. When `false`, all tasks are removed and no new ticks fire. | +| `worker_cache_ttl` | int | `3600` | Seconds the alive-worker IP list is cached before refreshing from `ray.nodes()`. | +| `tasks` | list | `[]` | List of `TaskConfig` entries (see [Section 4](#4-task-config-schema)). | + +### Related Environment Variables + +| Variable | Default | Purpose | +|----------|---------|---------| +| `ROCK_SCHEDULER_STATUS_DIR` | `/data/scheduler_status` | Directory on workers where per-task status JSON and run reports are written. | +| `ROCK_LOGGING_PATH` | (unset) | When set, scheduler-spawned daemons (docuum, container_cleanup, image_pull) redirect their stdout/stderr to `/.log`. | +| `ROCK_DOCUUM_INSTALL_URL` | `https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh` | Install script URL for `docuum`, fetched on demand by `ImageCleanupTask`. | + +## 3. Built-in Tasks + +ROCK ships with four built-in tasks under `rock.admin.scheduler.tasks`. Each task is registered by setting `task_class` to its fully qualified class path. + +### 3.1 ImageCleanupTask + +Runs [`docuum`](https://github.com/stepchowfun/docuum) on every worker to evict the least-recently-used Docker images once disk usage crosses a threshold. **Non-idempotent** — `docuum` runs as a long-lived daemon; the scheduler tracks its PID and skips re-launch while the daemon is alive. + +```yaml +- task_class: rock.admin.scheduler.tasks.image_cleanup_task.ImageCleanupTask + enabled: true + interval_seconds: 43200 # Re-check daemon every 12 hours + params: + disk_threshold: "70%" # Trigger eviction when disk usage exceeds 70% + image_whitelist: # Glob patterns matching repository:tag — never evicted + - "python:3.11" + - "my-registry.example.com/base/*" +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `disk_threshold` | str | `"1T"` | Disk usage threshold passed to `docuum --threshold`. Accepts size (`100G`, `1T`) or percentage (`70%`). | +| `image_whitelist` | list[str] | `[]` | Glob patterns forwarded to `docuum --keep`. | + +### 3.2 FileCleanupTask + +Walks each configured directory and removes files that are either older than `max_age_mins` or larger than `max_file_size`, then prunes empty subdirectories. **Idempotent**. + +```yaml +- task_class: rock.admin.scheduler.tasks.file_cleanup_task.FileCleanupTask + enabled: true + interval_seconds: 86400 # Run daily + params: + target_dirs: + # Plain string form — no exclusions + - "/data/service_status" + # Object form — per-directory exclusions + - path: "/data/logs" + exclude_files: # Plain name | relative path | absolute path + - "docuum.log" + - "./rocklet.log" + - "./access.log" + exclude_dirs: + - ".cache" + max_age_mins: 10080 # 7 days; older files are removed + max_file_size: "1G" # Files larger than this are removed (supports K/M/G/T) +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `target_dirs` | list | `[]` | Each entry is either a string (path only) or `{path, exclude_files, exclude_dirs}`. | +| `max_age_mins` | int | `10080` | Files whose mtime is older than this many minutes are deleted. | +| `max_file_size` | str | `"1G"` | Files larger than this are deleted. Suffixes `K/M/G/T` accepted. | + +The deletion condition is `(-mmin +max_age_mins) OR (-size +max_file_size)`. After file removal, a second `find -depth -type d -empty -delete` pass removes empty directories left behind (also honoring `exclude_dirs`). + +### 3.3 ContainerCleanupTask + +Removes stopped Docker containers older than a configurable age. Helps prevent the worker's container list from growing unbounded between sandbox runs. **Idempotent**. + +```yaml +- task_class: rock.admin.scheduler.tasks.container_cleanup_task.ContainerCleanupTask + enabled: true + interval_seconds: 86400 + params: + max_age_hours: 72 # Remove exited containers older than 72 hours +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `max_age_hours` | int | `24` | Maximum age (hours since `FinishedAt`) for kept exited containers. Older ones are `docker rm`'d. | + +The task also removes any container in the `created` state (never started) on every run. + +### 3.4 ImagePullTask + +Pre-pulls a list of Docker images on every worker, optionally logging in to private registries first. Reduces sandbox cold-start latency. **Idempotent** (`docker pull` is a no-op when the image is already up-to-date). + +```yaml +- task_class: rock.admin.scheduler.tasks.image_pull_task.ImagePullTask + enabled: true + interval_seconds: 21600 # Refresh every 6 hours + params: + images: + # Plain string form — public image, no auth + - "python:3.11" + # Object form — private image with registry login + - image: "my-registry.example.com/chatos/python:313" + registry_username: "myuser" + registry_password: "bXlwYXNzd29yZA==" # base64-encoded +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `images` | list | `[]` | Each entry is either an image string or `{image, registry_username, registry_password}`. | + +`registry_password` must be base64-encoded; the worker decodes it and pipes it to `docker login --password-stdin`. The registry host is parsed from the image name, so each image can target a different registry. + +## 4. Task Config Schema + +Every entry under `scheduler.tasks` is loaded as a `rock.config.TaskConfig`: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `task_class` | str | `""` | Fully qualified Python class path. Required. | +| `enabled` | bool | `true` | Disabled tasks are skipped at install time and torn down on reload. | +| `interval_seconds` | int | `3600` | APScheduler `interval` in seconds. | +| `params` | dict | `{}` | Task-specific kwargs forwarded to `from_config()`. | + +A change in any field of an existing task entry causes the scheduler to uninstall the old task (cleaning up its worker-side state when non-idempotent) and install the new one — without restarting the admin process. + +## 5. Writing a Custom Task + +Any class under your Python path that subclasses `BaseTask` can be registered. The minimum contract is: + +```python +# my_pkg/my_tasks/disk_report_task.py +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + + +class DiskReportTask(BaseTask): + """Log `df -h` output from every worker.""" + + def __init__(self, interval_seconds: int = 3600, mount_point: str = "/"): + super().__init__( + type="disk_report", # Used as job id and status filename prefix + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + self.mount_point = mount_point + + @classmethod + def from_config(cls, task_config) -> "DiskReportTask": + return cls( + interval_seconds=task_config.interval_seconds, + mount_point=task_config.params.get("mount_point", "/"), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + result = await runtime.execute( + Command(command=f"df -h {self.mount_point}", shell=True), + ) + return { + "status": TaskStatusEnum.SUCCESS, + "exit_code": result.exit_code, + "stdout": result.stdout, + } +``` + +Then register it from YAML: + +```yaml +scheduler: + enabled: true + tasks: + - task_class: my_pkg.my_tasks.disk_report_task.DiskReportTask + enabled: true + interval_seconds: 600 + params: + mount_point: "/data" +``` + +### Authoring Checklist + +- **Always set a unique `type` string** in `super().__init__()`. It is used as the APScheduler job id, the status filename (`_status.json`), and the run-report filename (`_run_report.json`). Two tasks must not share a `type`. +- **Pick the right `IdempotencyType`**: + - Use `IDEMPOTENT` when `run_action` finishes synchronously and is safe to re-execute. + - Use `NON_IDEMPOTENT` when you `nohup` a long-running daemon and return its PID. The scheduler will then track the PID, skip re-launch while it is alive, and `pkill` it on uninstall. +- **Return a dict from `run_action`**. Recommended keys: + - `status` — a `TaskStatusEnum` value, persisted into the status file. + - `pid` — required for `NON_IDEMPOTENT` daemons (use `rock.utils.system.extract_nohup_pid` on the `nohup ... & echo PID_PREFIX${{!}}PID_SUFFIX` output). + - Any other diagnostic fields are written to the status file's `extra` section. +- **Override `from_config(cls, task_config)`** to translate `task_config.params` into your `__init__` kwargs. +- **Use `runtime.execute / read_file / write_file`** rather than running shell commands locally — the scheduler is dispatching to remote workers via `RemoteSandboxRuntime`. + +## 6. Observability + +For each task, the scheduler writes two artifacts on every worker: + +| Path | Written By | Contents | +|------|------------|----------| +| `/_status.json` | `BaseTask.save_task_status` | Latest per-worker status: `task_name`, `worker_ip`, `pid`, `status` (`pending`/`running`/`success`/`failed`), `last_run`, `error`, plus task-specific `extra` fields. | +| `/_run_report.json` | `BaseTask.run` (admin side, after the tick completes) | Aggregated report: total/success/failed counts, list of `success_ips`, and `failed_details` (`ip` + traceback). | + +Scheduler-internal logs are written under the standard ROCK admin log path, with `name="scheduler"`, `name="task_base"`, `name="image_clean"`, etc. Scheduler-spawned daemons additionally write their own logs to `/.log` when `ROCK_LOGGING_PATH` is set (e.g. `docuum.log`, `container_cleanup.log`, `image_pull.log`). + +## 7. Dynamic Reload via Nacos (Optional) + +When the admin service is configured with a Nacos provider, the scheduler installs a YAML listener and reacts to config pushes: + +- Only the `scheduler:` section is inspected; other sections are ignored. +- The new section is hashed and compared against the previous one — duplicate notifications are skipped. +- A diff between old and new task lists determines which tasks to install, uninstall, or reinstall (changed `params` / `interval_seconds` / `enabled`). +- Non-idempotent tasks that are removed or re-installed are first cleaned up: their daemon PID is killed and the status file is removed. + +This means task interval changes, parameter tweaks, and adding/removing tasks can be applied without restarting the admin process. + +## 8. Troubleshooting + +| Symptom | Likely Cause | What to Check | +|---------|--------------|----------------| +| `Scheduler disabled, all tasks removed` in admin log | `scheduler.enabled` is `false` | Set `enabled: true` in YAML. | +| `No alive workers found for task ''` | Ray cluster has no live worker nodes | Verify `ray.nodes()` reports alive CPU workers; consider lowering `worker_cache_ttl` if workers were just added. | +| Task ticks fire but every worker shows up in `failed_details` with connection errors | `rocklet` is not running on the workers, or `Port.PROXY` is blocked | On the worker host, hit the rocklet liveness endpoint `GET /is_alive` on `Port.PROXY` (e.g. `curl http://:/is_alive`); if it does not respond, restart rocklet (`rocklet --port `) or open the port in the firewall. | +| Task runs but never repeats on a non-idempotent worker | Recorded PID still alive | Inspect `/_status.json`; if `status: running` and the PID is alive, `should_run` returns `False`. | +| `Failed to create task ''` | `task_class` import failed | Ensure the module is importable inside the admin process (installed in the same venv, on `PYTHONPATH`). | +| `docker login` failing in `ImagePullTask` | `registry_password` not base64-encoded, or wrong registry parsed from image | Re-encode the password with `echo -n '' \| base64`; double-check the image's registry host. | + +## Related Documents + +- [Configuration](./configuration.md) — Environment variables and runtime layout +- [API Documentation](../References/api.md) — Admin HTTP API +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) — Programmatic sandbox usage diff --git a/rock-conf/rock-local.yml b/rock-conf/rock-local.yml index a4f364d5f0..cab2958f5b 100644 --- a/rock-conf/rock-local.yml +++ b/rock-conf/rock-local.yml @@ -7,3 +7,45 @@ ray: warmup: images: - "python:3.11" + +# Scheduler configuration +scheduler: + enabled: true # Whether to enable the scheduler + worker_cache_ttl: 43200 # Worker IP cache TTL (seconds) + tasks: # Task list + - task_class: rock.admin.scheduler.tasks.image_cleanup_task.ImageCleanupTask # Task class path + enabled: true # Whether to enable this task + interval_seconds: 43200 # Execution interval (seconds) + params: # Task-specific parameters + disk_threshold: "70%" # docuum cleanup threshold + image_whitelist: + - "python:3.11" + - task_class: rock.admin.scheduler.tasks.file_cleanup_task.FileCleanupTask # File cleanup task class path + enabled: true # Whether to enable this task + interval_seconds: 360 # Execution interval (seconds); runs once per day by default + params: # Task-specific parameters + target_dirs: # List of target directories to clean up; each directory can configure its own exclusions + - path: "/data/logs" # Directory path + exclude_files: # Files to exclude in this directory (supports name, relative path, absolute path) + - "docuum.log" + - "./rocklet.log" + - "./rock_worker.log" + - "./access.log" + - "/data/service_status" + max_age_mins: 180 # Maximum file retention time (minutes); files whose last-modified time exceeds this value will be cleaned up. Default: 7 days (10080 minutes) + max_file_size: "1G" # File size threshold; files larger than this will be cleaned up. Supports K/M/G/T units + - task_class: "rock.admin.scheduler.tasks.container_cleanup_task.ContainerCleanupTask" + interval_seconds: 86400 + enabled: true + params: + max_age_hours: 72 + - task_class: rock.admin.scheduler.tasks.image_pull_task.ImagePullTask + enabled: true + interval_seconds: 21600 + params: + images: + # - image: "python:3.11" + # registry_username: "user" + # registry_password: "ZNDQ0NWVlZDUzZDTVmMQo=" + - "python:3.11" + diff --git a/rock/env_vars.py b/rock/env_vars.py index 7844da9c62..659ae4255c 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -45,6 +45,7 @@ ROCK_PYTHON_ENV_PATH: str | None = None ROCK_ADMIN_ENV: str | None = "dev" ROCK_ADMIN_ROLE: str | None = "write" + ROCK_FORCE_PRIMARY_POD: bool = False ROCK_CLI_LOAD_PATHS: str = str(Path(__file__).parent / "cli" / "command") ROCK_CLI_DEFAULT_CONFIG_PATH: str @@ -102,6 +103,7 @@ "ROCK_PYTHON_ENV_PATH": lambda: os.getenv("ROCK_PYTHON_ENV_PATH", sys.base_prefix), "ROCK_ADMIN_ENV": lambda: os.getenv("ROCK_ADMIN_ENV", "dev"), "ROCK_ADMIN_ROLE": lambda: os.getenv("ROCK_ADMIN_ROLE", "write"), + "ROCK_FORCE_PRIMARY_POD": lambda: os.getenv("ROCK_FORCE_PRIMARY_POD", "false").lower() == "true", "ROCK_CLI_LOAD_PATHS": lambda: os.getenv("ROCK_CLI_LOAD_PATHS", str(Path(__file__).parent / "cli" / "command")), "ROCK_CLI_DEFAULT_CONFIG_PATH": lambda: os.getenv( "ROCK_CLI_DEFAULT_CONFIG_PATH", Path.home() / ".rock" / "config.ini" diff --git a/rock/utils/system.py b/rock/utils/system.py index b0d1e035e0..a79cc55224 100644 --- a/rock/utils/system.py +++ b/rock/utils/system.py @@ -238,7 +238,12 @@ def is_primary_pod() -> bool: Check if the current pod is the primary pod (index 0). Reads /etc/hostname file and parses the pod index from the pod name. Hostname format example: rock-admin-write-nt-gray-0.rock-admin-write-nt-gray-hs.chatos.svc.cluster.local + + If ROCK_FORCE_PRIMARY_POD is set to true, treat the current pod as primary + without inspecting the hostname. """ + if env_vars.ROCK_FORCE_PRIMARY_POD: + return True try: with open("/etc/hostname") as f: hostname = f.read().strip() From 93ea33096d2d290688775454cda91d003694a2cc Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Fri, 15 May 2026 23:47:12 +0800 Subject: [PATCH 102/226] feat(oss): unify dual-account STS via /get_token?account= and push transfer prefix to sdk #952 --- rock/admin/entrypoints/sandbox_proxy_api.py | 19 +- rock/config.py | 39 +++ rock/env_vars.py | 5 + rock/sandbox/service/sandbox_proxy_service.py | 73 ++++-- rock/sdk/sandbox/oss_client.py | 19 +- tests/unit/sandbox/test_sts_dual_account.py | 235 ++++++++++++++++++ tests/unit/sdk/sandbox/test_oss_client.py | 32 +++ tests/unit/test_config.py | 33 +++ 8 files changed, 435 insertions(+), 20 deletions(-) create mode 100644 tests/unit/sandbox/test_sts_dual_account.py diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 808a2b0624..1e67231f21 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -321,8 +321,23 @@ async def portforward(websocket: WebSocket, id: str, port: int): @sandbox_proxy_router.get("/get_token") @handle_exceptions(error_message="get oss sts token failed") -async def get_token(): - result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token) +async def get_token(account: str = "legacy"): + """STS token for OSS upload/download. + + Query param `account` selects the credential pool: + - "legacy" (default): xrl-sandbox bucket, BC for SDK < 1.8 + - "primary": chatos-rock bucket, used by SDK >= 1.8 + """ + result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token, account) + return RockResponse(result=result) + + +@sandbox_proxy_router.get("/get_token_v2") +@handle_exceptions(error_message="get oss sts token v2 failed") +async def get_token_v2(): + """Primary-account STS for SDK >= 1.8. + Response format is identical to /get_token (same Credentials shape).""" + result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token_v2) return RockResponse(result=result) diff --git a/rock/config.py b/rock/config.py index 1495bdb327..4c1b7937ba 100644 --- a/rock/config.py +++ b/rock/config.py @@ -60,6 +60,18 @@ class SandboxConfig: remove_container_enabled: bool = True +@dataclass +class OssAccountConfig: + endpoint: str = "" + bucket: str = "" + access_key_id: str = "" + access_key_secret: str = "" + role_arn: str = "" + region: str = "" + """Region used to construct the STS AcsClient for this account. Falls back + to env_vars.ROCK_OSS_BUCKET_REGION when empty, to preserve legacy behavior.""" + + @dataclass class OssConfig: endpoint: str = "" @@ -67,6 +79,33 @@ class OssConfig: access_key_id: str = "" access_key_secret: str = "" role_arn: str = "" + region: str = "" + """Region for the legacy STS AcsClient. Empty falls back to + env_vars.ROCK_OSS_BUCKET_REGION, YAML-level values always win over env.""" + + primary: OssAccountConfig = field(default_factory=OssAccountConfig) + """Primary account used by SDK >= 1.8 (`/get_token_v2` endpoint) and by + all host-side archival (OssArchiver, landed in PR-1). An empty + `primary.bucket` disables v2 STS and archival, leaving legacy path + fully operational.""" + + transfer_prefix: str = "" + """Prefix under the PRIMARY bucket for ephemeral host↔container file + transfers ({timestamp}-{filename} objects). The legacy bucket keeps + its pre-existing flat layout (no prefix) for backward compatibility — + xrl-sandbox has a 3-day lifecycle rule at bucket root (configured + in the Aliyun OSS console, not in repo) that we do not disturb. + + Note: this field lives in admin-side RockConfig and is NOT what the + SDK reads. The SDK reads ROCK_OSS_TRANSFER_PREFIX directly from the + process env. xrl package is no longer maintained, so internal users + must export this env var themselves when upgrading to SDK >= 1.8.""" + + def __post_init__(self): + # Allow YAML to pass a dict for `primary` (dataclass deserialization + # from yaml.safe_load returns dicts, not nested dataclasses). + if isinstance(self.primary, dict): + self.primary = OssAccountConfig(**self.primary) @dataclass diff --git a/rock/env_vars.py b/rock/env_vars.py index 659ae4255c..11a9830613 100644 --- a/rock/env_vars.py +++ b/rock/env_vars.py @@ -37,6 +37,10 @@ ROCK_OSS_BUCKET_ENDPOINT: str | None = None ROCK_OSS_BUCKET_NAME: str | None = None ROCK_OSS_BUCKET_REGION: str | None = None + ROCK_OSS_TRANSFER_PREFIX: str | None = None + """Optional key prefix for host↔container transfer objects under the + bucket identified by ROCK_OSS_BUCKET_NAME. Empty = flat layout (legacy + bucket). New SDK clusters set this to "rock-transfer/".""" ROCK_PIP_INDEX_URL: str | None = "https://mirrors.aliyun.com/pypi/simple/" ROCK_MONITOR_ENABLE: bool = False @@ -92,6 +96,7 @@ "ROCK_RAY_NAMESPACE": lambda: os.getenv("ROCK_RAY_NAMESPACE", "xrl-sandbox"), "ROCK_SANDBOX_EXPIRE_TIME_KEY": lambda: os.getenv("ROCK_SANDBOX_EXPIRE_TIME_KEY", "expire_time"), "ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY": lambda: os.getenv("ROCK_SANDBOX_AUTO_CLEAR_TIME_KEY", "auto_clear_time"), + "ROCK_OSS_TRANSFER_PREFIX": lambda: os.getenv("ROCK_OSS_TRANSFER_PREFIX"), "ROCK_OSS_ENABLE": lambda: os.getenv("ROCK_OSS_ENABLE", "false").lower() == "true", "ROCK_OSS_BUCKET_ENDPOINT": lambda: os.getenv("ROCK_OSS_BUCKET_ENDPOINT"), "ROCK_OSS_BUCKET_NAME": lambda: os.getenv("ROCK_OSS_BUCKET_NAME"), diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 55aa23fdc2..8c6d8b55e1 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -71,11 +71,22 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): ), ) - self.sts_client = client.AcsClient( - self.oss_config.access_key_id, - self.oss_config.access_key_secret, - env_vars.ROCK_OSS_BUCKET_REGION, - ) + # Replace single self.sts_client with a dict keyed by account name, + # so /get_token?account=legacy|primary maps to the right credentials. + legacy_region = self.oss_config.region or env_vars.ROCK_OSS_BUCKET_REGION + primary_region = self.oss_config.primary.region or env_vars.ROCK_OSS_BUCKET_REGION + self._sts_clients = { + "legacy": client.AcsClient( + self.oss_config.access_key_id, + self.oss_config.access_key_secret, + legacy_region, + ), + "primary": client.AcsClient( + self.oss_config.primary.access_key_id, + self.oss_config.primary.access_key_secret, + primary_region, + ), + } self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count self._validate_oss_config_or_warn() @@ -668,30 +679,60 @@ def _api_url(self, host_ip: str, service_status: ServiceStatus) -> str: port = service_status.get_mapped_port(Port.PROXY) return f"http://{host_ip}:{port}" - def gen_oss_sts_token(self): - role_arn = self.oss_config.role_arn + def gen_oss_sts_token(self, account: str = "legacy") -> dict | None: # CHANGED: account param, default "legacy" preserves BC + """Generate STS credentials and OSS config for the given account. + Args: + account: "legacy" (xrl-sandbox, BC for SDK < 1.8) or + "primary" (chatos-rock, SDK >= 1.8). + Returns: + Dict with STS credentials (AccessKeyId, AccessKeySecret, SecurityToken, Expiration) PLUS account-scoped OSS config: + Endpoint, Bucket, Region, Prefix. None on failure or when the requested account is unconfigured. + """ + if account not in self._sts_clients: + logger.error(f"unknown OSS account: {account!r}") + return None + + if account == "primary": + primary = self.oss_config.primary + role_arn = primary.role_arn + session_name = "rock-sandbox-primary" + endpoint = primary.endpoint or None + bucket = primary.bucket or None + region = primary.region or env_vars.ROCK_OSS_BUCKET_REGION or None + prefix = self.oss_config.transfer_prefix or None + else: # legacy + role_arn = self.oss_config.role_arn + session_name = "rock-sandbox-legacy" + endpoint = env_vars.ROCK_OSS_BUCKET_ENDPOINT or self.oss_config.endpoint or None + bucket = env_vars.ROCK_OSS_BUCKET_NAME or self.oss_config.bucket or None + region = env_vars.ROCK_OSS_BUCKET_REGION or None + prefix = env_vars.ROCK_OSS_TRANSFER_PREFIX or None + + if not role_arn: + logger.warning(f"oss role_arn not configured for account={account!r}") + return None + request = CommonRequest(product="Sts", version="2015-04-01", action_name="AssumeRole") request.set_method("POST") request.set_protocol_type("https") request.add_query_param("RoleArn", role_arn) - request.add_query_param("RoleSessionName", "sessiontest") + request.add_query_param("RoleSessionName", session_name) # at least 900s request.add_query_param("DurationSeconds", "900") request.set_accept_format("JSON") try: - body = self.sts_client.do_action_with_exception(request) - token = json.loads(oss2.to_unicode(body)) - credentials = token["Credentials"] + body = self._sts_clients[account].do_action_with_exception(request) + credentials = json.loads(oss2.to_unicode(body))["Credentials"] except Exception: - logger.error("generate oss sts token failed") + logger.error(f"generate oss sts token failed (account={account})", exc_info=True) return None return { **credentials, - # env > YAML, matches client-side Layer 1 priority - "Endpoint": env_vars.ROCK_OSS_BUCKET_ENDPOINT or self.oss_config.endpoint or None, - "Bucket": env_vars.ROCK_OSS_BUCKET_NAME or self.oss_config.bucket or None, - "Region": env_vars.ROCK_OSS_BUCKET_REGION or None, + "Endpoint": endpoint, + "Bucket": bucket, + "Region": region, + "Prefix": prefix, # transfer-object key prefix, scoped per account } async def get_sandbox_websocket_url( diff --git a/rock/sdk/sandbox/oss_client.py b/rock/sdk/sandbox/oss_client.py index 51711ff133..c16306ccfa 100644 --- a/rock/sdk/sandbox/oss_client.py +++ b/rock/sdk/sandbox/oss_client.py @@ -36,6 +36,7 @@ class OssClientConfig: bucket: str region: str enabled_via_env: bool # True = Layer 1 (gated by ROCK_OSS_ENABLE); False = Layer 2 + prefix: str = "" # Transfer-object key prefix (Layer 1 from env; Layer 2 from server response) class OssClient: @@ -49,13 +50,22 @@ def __init__(self, sandbox: Sandbox): self._pending_persistence_tasks: set[asyncio.Task] = set() @staticmethod - def _compute_object_name(sandbox_id: str, local_path: str, sandbox_path: str) -> str: + def _compute_object_name( + sandbox_id: str, + local_path: str, + sandbox_path: str, + prefix: str | None = None, # NEW: server-pushed transfer prefix (e.g. "rock-transfer/") + ) -> str: # Prefer sandbox basename: the OSS object mirrors a sandbox-side file, # so naming it after the sandbox path keeps OSS-side names meaningful # even when local destinations differ (e.g. download to a renamed file). payload = f"{sandbox_id}|{local_path}|{sandbox_path}" digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() filename = Path(sandbox_path).name or Path(local_path).name + # Honor server-pushed Prefix so chatos-rock's rock-transfer/ lifecycle catches the object. + clean_prefix = (prefix or "").strip("/") + if clean_prefix: + return f"{clean_prefix}/{digest}-{filename}" return f"{digest}-{filename}" @staticmethod @@ -70,6 +80,7 @@ def _resolve_config(sts_response: dict) -> OssClientConfig | None: bucket=env_bucket, region=env_region, enabled_via_env=True, + prefix=env_vars.ROCK_OSS_TRANSFER_PREFIX or "", # NEW: env can carry prefix too ) # Layer 2: server response (fallback default) @@ -82,6 +93,7 @@ def _resolve_config(sts_response: dict) -> OssClientConfig | None: bucket=resp_bucket, region=resp_region, enabled_via_env=False, + prefix=sts_response.get("Prefix") or "", # NEW: pull prefix from server ) # Layer 3: OSS unavailable @@ -96,7 +108,7 @@ async def _get_sts_credentials(self) -> dict: Side effect: caches Expiration in self._token_expire_time. """ - url = f"{self._sandbox._url}/get_token" + url = f"{self._sandbox._url}/get_token?account=primary" headers = self._sandbox._build_headers() response = await HttpUtils.get(url, headers) if response["status"] != "Success": @@ -183,6 +195,7 @@ async def upload_via_oss(self, file_path: str, target_path: str) -> UploadRespon sandbox_id=self._sandbox.sandbox_id, local_path=file_path, sandbox_path=target_path, + prefix=self._client_config.prefix if self._client_config else None, # NEW ) try: @@ -260,6 +273,7 @@ async def download_via_oss(self, remote_path: str, local_path: Path) -> Download sandbox_id=self._sandbox.sandbox_id, local_path=str(local_path), sandbox_path=remote_path, + prefix=self._client_config.prefix if self._client_config else None, ) oss_url = f"oss://{self._client_config.bucket}/{oss_object_name}" @@ -309,6 +323,7 @@ async def schedule_async_persistence(self, local_path: str, sandbox_path: str) - sandbox_id=self._sandbox.sandbox_id, local_path=local_path, sandbox_path=sandbox_path, + prefix=self._client_config.prefix if self._client_config else None, ) task = asyncio.create_task(self._persist_to_oss(local_path, oss_object_name)) self._pending_persistence_tasks.add(task) diff --git a/tests/unit/sandbox/test_sts_dual_account.py b/tests/unit/sandbox/test_sts_dual_account.py new file mode 100644 index 0000000000..8c3b59009f --- /dev/null +++ b/tests/unit/sandbox/test_sts_dual_account.py @@ -0,0 +1,235 @@ +import json +from unittest.mock import MagicMock, patch + +from rock.config import ( + OssAccountConfig, + OssConfig, + ProxyServiceConfig, + RockConfig, + RuntimeConfig, +) +from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService + + +def _make_rock_config( + *, + legacy_role: str, + primary_role: str, + legacy_region: str = "cn-hangzhou", + transfer_prefix: str = "rock-transfer/", +) -> RockConfig: + return RockConfig( + oss=OssConfig( + endpoint="oss-cn-hangzhou.aliyuncs.com", + bucket="xrl-sandbox", + access_key_id="legacy-ak", + access_key_secret="legacy-sk", + role_arn=legacy_role, + region=legacy_region, + transfer_prefix=transfer_prefix, + primary=OssAccountConfig( + endpoint="oss-cn-hangzhou.aliyuncs.com", + bucket="chatos-rock", + access_key_id="primary-ak", + access_key_secret="primary-sk", + role_arn=primary_role, + region="cn-hangzhou", + ), + ), + proxy_service=ProxyServiceConfig(), + runtime=RuntimeConfig( + python_env_path="/usr/bin/python3", + envhub_db_url="sqlite:////tmp/rock_envs.db", + ), + ) + + +def _fake_assume(ak: str, sk: str, tok: str): + def _do(req): + return json.dumps( + { + "Credentials": { + "AccessKeyId": ak, + "AccessKeySecret": sk, + "SecurityToken": tok, + "Expiration": "2099-01-01T00:00:00Z", + } + } + ).encode() + + return _do + + +def _build_service(rock_config: RockConfig) -> SandboxProxyService: + # Each AcsClient() call returns a fresh MagicMock so legacy/primary clients + # are not aliased — otherwise per-account fakes would leak across accounts. + with patch( + "rock.sandbox.service.sandbox_proxy_service.client.AcsClient", + side_effect=lambda *a, **kw: MagicMock(), + ): + return SandboxProxyService(rock_config, meta_store=MagicMock()) + + +def test_legacy_account_uses_legacy_role_and_bucket(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + ) + ) + svc._sts_clients["legacy"].do_action_with_exception = _fake_assume("L-AK", "L-SK", "L-TOK") + svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") + + creds = svc.gen_oss_sts_token() # default account="legacy" + assert creds["AccessKeyId"] == "L-AK" + assert creds["Bucket"] == "xrl-sandbox" + assert creds["Prefix"] is None + + +def test_primary_account_uses_primary_role_and_bucket_with_prefix(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + ) + ) + svc._sts_clients["legacy"].do_action_with_exception = _fake_assume("L-AK", "L-SK", "L-TOK") + svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") + + creds = svc.gen_oss_sts_token(account="primary") + assert creds["AccessKeyId"] == "P-AK" + assert creds["Bucket"] == "chatos-rock" + assert creds["Prefix"] == "rock-transfer/" + + +def test_primary_returns_none_when_role_arn_empty(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="", + ) + ) + assert svc.gen_oss_sts_token(account="primary") is None + + +def test_primary_prefix_is_none_when_yaml_does_not_set_it(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + transfer_prefix="", + ) + ) + svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") + creds = svc.gen_oss_sts_token(account="primary") + assert creds["Prefix"] is None + + +def test_unknown_account_returns_none(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + ) + ) + assert svc.gen_oss_sts_token(account="bogus") is None + + +def test_legacy_prefix_falls_back_to_env(monkeypatch): + monkeypatch.setenv("ROCK_OSS_TRANSFER_PREFIX", "legacy-pref/") + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + ) + ) + svc._sts_clients["legacy"].do_action_with_exception = _fake_assume("L-AK", "L-SK", "L-TOK") + creds = svc.gen_oss_sts_token() + assert creds["Prefix"] == "legacy-pref/" # env wins for legacy + + +def test_assume_role_session_names_are_distinct(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + ) + ) + captured: dict[str, dict] = {} + + def capture(label: str): + def _do(req): + captured[label] = dict(req.get_query_params()) + return json.dumps( + { + "Credentials": { + "AccessKeyId": "x", + "AccessKeySecret": "y", + "SecurityToken": "z", + "Expiration": "2099-01-01T00:00:00Z", + } + } + ).encode() + + return _do + + svc._sts_clients["legacy"].do_action_with_exception = capture("legacy") + svc._sts_clients["primary"].do_action_with_exception = capture("primary") + svc.gen_oss_sts_token(account="legacy") + svc.gen_oss_sts_token(account="primary") + + assert captured["legacy"]["RoleSessionName"] == "rock-sandbox-legacy" + assert captured["primary"]["RoleSessionName"] == "rock-sandbox-primary" + assert captured["legacy"]["RoleArn"] == "acs:ram::1933967579503727:role/legacy-role" + assert captured["primary"]["RoleArn"] == "acs:ram::1771269394322852:role/chatos-rock-sts-role" + + +def test_legacy_region_from_yaml_overrides_env(monkeypatch): + monkeypatch.setenv("ROCK_OSS_BUCKET_REGION", "cn-shanghai") + captured: list[tuple[str, str]] = [] + + def fake_acs_client(ak, sk, region): + captured.append((ak, region)) + return MagicMock() + + with patch( + "rock.sandbox.service.sandbox_proxy_service.client.AcsClient", + side_effect=fake_acs_client, + ): + SandboxProxyService( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + legacy_region="cn-hangzhou", + ), + meta_store=MagicMock(), + ) + + legacy_ak, legacy_region = captured[0] + assert legacy_ak == "legacy-ak" + assert legacy_region == "cn-hangzhou" # yaml won + + +def test_legacy_region_falls_back_to_env_when_yaml_empty(monkeypatch): + monkeypatch.setenv("ROCK_OSS_BUCKET_REGION", "cn-shanghai") + captured: list[tuple[str, str]] = [] + + def fake_acs_client(ak, sk, region): + captured.append((ak, region)) + return MagicMock() + + with patch( + "rock.sandbox.service.sandbox_proxy_service.client.AcsClient", + side_effect=fake_acs_client, + ): + SandboxProxyService( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + legacy_region="", + ), + meta_store=MagicMock(), + ) + + legacy_ak, legacy_region = captured[0] + assert legacy_region == "cn-shanghai" # env fallback diff --git a/tests/unit/sdk/sandbox/test_oss_client.py b/tests/unit/sdk/sandbox/test_oss_client.py index 361ff55efc..66cfde92c1 100644 --- a/tests/unit/sdk/sandbox/test_oss_client.py +++ b/tests/unit/sdk/sandbox/test_oss_client.py @@ -490,3 +490,35 @@ async def hang(*args, **kwargs): with patch("asyncio.to_thread", new=hang): await client.schedule_async_persistence("/local/foo.json", "/sandbox/foo.json") await client.close(timeout=0.05) + + +# prefix propagation +def test_compute_object_name_with_prefix(): + name = OssClient._compute_object_name( + sandbox_id="sb-1", + local_path="/tmp/x", + sandbox_path="/data/x", + prefix="rock-transfer/", + ) + assert name.startswith("rock-transfer/") + assert name.endswith("-x") + + +def test_compute_object_name_strips_slashes_in_prefix(): + name = OssClient._compute_object_name( + sandbox_id="sb-1", + local_path="/tmp/x", + sandbox_path="/data/x", + prefix="/rock-transfer//", + ) + assert name.startswith("rock-transfer/") + assert "//" not in name + + +def test_compute_object_name_no_prefix_keeps_legacy_layout(): + name = OssClient._compute_object_name( + sandbox_id="sb-1", + local_path="/tmp/x", + sandbox_path="/data/x", + ) + assert "/" not in name # flat layout diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 1edf62d1ab..d2d0941f12 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -48,3 +48,36 @@ async def test_runtime_config(): assert runtime_config.max_allowed_spec.cpus == 16 assert runtime_config.standard_spec.memory == "8g" assert runtime_config.standard_spec.cpus == 2 + + +def test_oss_config_defaults(): + from rock.config import OssAccountConfig, OssConfig + + cfg = OssConfig() + assert cfg.bucket == "" + assert cfg.region == "" + assert cfg.transfer_prefix == "" # default is now empty; deployments must opt-in via YAML + assert isinstance(cfg.primary, OssAccountConfig) + assert cfg.primary.bucket == "" + assert cfg.primary.region == "" + + +def test_oss_config_primary_dict_coerced(): + from rock.config import OssAccountConfig, OssConfig + + cfg = OssConfig( + primary={ + "endpoint": "e", + "bucket": "chatos-rock", + "access_key_id": "a", + "access_key_secret": "s", + "role_arn": "r", + "region": "cn-hangzhou", + } + ) + assert isinstance(cfg.primary, OssAccountConfig) + assert cfg.primary.bucket == "chatos-rock" + assert cfg.primary.region == "cn-hangzhou" + # legacy 顶层字段未提供时仍为默认空,确认 primary 不会污染 legacy + assert cfg.bucket == "" + From 0ed625dac6318dbce0d96f0ebc804715b8503529 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 10:01:58 +0800 Subject: [PATCH 103/226] fix(sandbox): update test_sandbox_proxy after gen_oss_sts_token dual-account refactor --- tests/unit/sandbox/test_sandbox_proxy.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index d8f44fa40b..f12081a30c 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -95,7 +95,8 @@ def sandbox_proxy_service(self): # (which requires real Redis / metrics / RAM-Acs client setup). service = SandboxProxyService.__new__(SandboxProxyService) service.oss_config = OssConfig(role_arn="test_role_arn") - service.sts_client = MagicMock() + # gen_oss_sts_token routes by account name; legacy is the default. + service._sts_clients = {"legacy": MagicMock(), "primary": MagicMock()} return service def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): @@ -107,7 +108,7 @@ def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" @@ -122,7 +123,7 @@ def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): def test_sts_failure_returns_none(self, sandbox_proxy_service): with patch.object( - sandbox_proxy_service.sts_client, "do_action_with_exception", side_effect=Exception("sts fail") + sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", side_effect=Exception("sts fail") ): result = sandbox_proxy_service.gen_oss_sts_token() assert result is None @@ -137,7 +138,7 @@ def test_partial_oss_config_returns_creds_with_none_extras(self, sandbox_proxy_s b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" @@ -160,7 +161,7 @@ def test_env_var_overrides_yaml_for_endpoint_and_bucket(self, sandbox_proxy_serv b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "env.endpoint" @@ -183,7 +184,7 @@ def test_yaml_used_when_env_var_empty(self, sandbox_proxy_service): b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service.sts_client, "do_action_with_exception", return_value=fake_token_body), + patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" From e0fd6f4c0a6589f704dd500a31468cc639e3c109 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 11:57:53 +0800 Subject: [PATCH 104/226] fix(api): wire /get_token_v2 to unified gen_oss_sts_token("primary") The old gen_oss_sts_token_v2() was removed during the dual-account refactor but this route was not updated, causing an AttributeError at runtime. --- rock/admin/entrypoints/sandbox_proxy_api.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 1e67231f21..31ad699522 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -336,8 +336,9 @@ async def get_token(account: str = "legacy"): @handle_exceptions(error_message="get oss sts token v2 failed") async def get_token_v2(): """Primary-account STS for SDK >= 1.8. + Delegates to the unified ``gen_oss_sts_token`` with account="primary". Response format is identical to /get_token (same Credentials shape).""" - result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token_v2) + result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token, "primary") return RockResponse(result=result) From 2e88d6fe2b962a6ab269ad3767b95ed5a00b3635 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 12:13:58 +0800 Subject: [PATCH 105/226] refactor(api): remove redundant /get_token_v2 endpoint, unify on /get_token?account= PR0 design unifies dual-account STS through a single /get_token endpoint with an `account` query param. The /get_token_v2 route was a transitional shim that stayed alive after the unification; it now adds nothing beyond /get_token?account=primary and contradicts the design doc. Drop the route entirely and update OssConfig.primary docstring to point at the canonical endpoint. Co-Authored-By: Claude Opus 4.7 (1M context) --- rock/admin/entrypoints/sandbox_proxy_api.py | 10 ---------- rock/config.py | 7 +++---- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 31ad699522..68f2d28689 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -332,16 +332,6 @@ async def get_token(account: str = "legacy"): return RockResponse(result=result) -@sandbox_proxy_router.get("/get_token_v2") -@handle_exceptions(error_message="get oss sts token v2 failed") -async def get_token_v2(): - """Primary-account STS for SDK >= 1.8. - Delegates to the unified ``gen_oss_sts_token`` with account="primary". - Response format is identical to /get_token (same Credentials shape).""" - result = await asyncio.to_thread(sandbox_proxy_service.gen_oss_sts_token, "primary") - return RockResponse(result=result) - - @sandbox_proxy_router.api_route( "/sandboxes/{sandbox_id}/vnc", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"], diff --git a/rock/config.py b/rock/config.py index 4c1b7937ba..6913da454c 100644 --- a/rock/config.py +++ b/rock/config.py @@ -84,10 +84,9 @@ class OssConfig: env_vars.ROCK_OSS_BUCKET_REGION, YAML-level values always win over env.""" primary: OssAccountConfig = field(default_factory=OssAccountConfig) - """Primary account used by SDK >= 1.8 (`/get_token_v2` endpoint) and by - all host-side archival (OssArchiver, landed in PR-1). An empty - `primary.bucket` disables v2 STS and archival, leaving legacy path - fully operational.""" + """Primary account used by SDK >= 1.8 (`/get_token?account=primary`) and by + host-side archival. An empty `primary.bucket` disables v2 STS and archival, + leaving legacy path fully operational.""" transfer_prefix: str = "" """Prefix under the PRIMARY bucket for ephemeral host↔container file From 3cad424d2b099a121c5e92df02848d75e73fbc4d Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Sat, 16 May 2026 23:38:59 +0800 Subject: [PATCH 106/226] fix: conditionally create primary AcsClient only when credentials configured MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Avoid constructing an AcsClient with empty AK/SK when primary OSS account is not configured in YAML. The existing `if account not in self._sts_clients` check in gen_oss_sts_token already handles the unconfigured case gracefully by returning None. 🤖 Generated with [Qoder][https://qoder.com] --- rock/sandbox/service/sandbox_proxy_service.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 8c6d8b55e1..dc5de7860d 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -74,19 +74,22 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): # Replace single self.sts_client with a dict keyed by account name, # so /get_token?account=legacy|primary maps to the right credentials. legacy_region = self.oss_config.region or env_vars.ROCK_OSS_BUCKET_REGION - primary_region = self.oss_config.primary.region or env_vars.ROCK_OSS_BUCKET_REGION self._sts_clients = { "legacy": client.AcsClient( self.oss_config.access_key_id, self.oss_config.access_key_secret, legacy_region, ), - "primary": client.AcsClient( + } + # Only create primary client when credentials are configured, + # avoiding an AcsClient with empty AK/SK that would fail at call time. + if self.oss_config.primary.access_key_id: + primary_region = self.oss_config.primary.region or env_vars.ROCK_OSS_BUCKET_REGION + self._sts_clients["primary"] = client.AcsClient( self.oss_config.primary.access_key_id, self.oss_config.primary.access_key_secret, primary_region, - ), - } + ) self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count self._validate_oss_config_or_warn() From fcb3c0f4849f87f0bc7fa66f1a72dad7c507efe4 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Mon, 18 May 2026 15:42:54 +0800 Subject: [PATCH 107/226] fix(rocklet): use cgroup metrics for container CPU instead of psutil (#945) (#946) * fix(rocklet): use cgroup metrics for container CPU instead of psutil Co-Authored-By: Claude Opus 4.6 * refactor(rocklet): simplify cgroup CPU cache to single scalar and cache cpu_quota Remove per-thread dict for CPU usage/time (no concurrency concern) and cache cpu_quota since it doesn't change at runtime. Co-Authored-By: Claude Opus 4.6 * refactor(rocklet): convert cgroup_stats to OOP with CgroupCpuStats class Replace module-level globals and functions with a CgroupCpuStats class, eliminating shared mutable state and simplifying test isolation. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- rock/rocklet/linux.py | 7 +- rock/utils/cgroup_stats.py | 113 ++++++++++ tests/unit/utils/test_cgroup_stats.py | 288 ++++++++++++++++++++++++++ 3 files changed, 407 insertions(+), 1 deletion(-) create mode 100644 rock/utils/cgroup_stats.py create mode 100644 tests/unit/utils/test_cgroup_stats.py diff --git a/rock/rocklet/linux.py b/rock/rocklet/linux.py index 00bb49499b..6f91eb8ca5 100644 --- a/rock/rocklet/linux.py +++ b/rock/rocklet/linux.py @@ -34,6 +34,7 @@ SessionNotInitializedError, ) from rock.utils import get_executor +from rock.utils.cgroup_stats import CgroupCpuStats from .rocklet import Rocklet, Session @@ -342,12 +343,16 @@ def interact(self) -> None: class LinuxRocklet(Rocklet): """Rocklet implementation for sys.platform in {'linux', 'darwin'}.""" + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._cgroup_cpu = CgroupCpuStats() + def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: return BashSession(request) async def get_statistics(self) -> dict: return { - "cpu": psutil.cpu_percent(), + "cpu": self._cgroup_cpu.cpu_percent(), "mem": psutil.virtual_memory().percent, "disk": psutil.disk_usage("/").percent, "net": psutil.net_io_counters().bytes_recv + psutil.net_io_counters().bytes_sent, diff --git a/rock/utils/cgroup_stats.py b/rock/utils/cgroup_stats.py new file mode 100644 index 0000000000..233a414191 --- /dev/null +++ b/rock/utils/cgroup_stats.py @@ -0,0 +1,113 @@ +"""Container-aware CPU metrics via cgroup v1/v2. + +Falls back to psutil when cgroup files are unavailable (e.g. running +outside a container or on non-Linux platforms). +""" + +import os +import time +from pathlib import Path + +import psutil + + +class CgroupCpuStats: + """Reads container CPU utilization from cgroup v1/v2 pseudo-files.""" + + def __init__(self): + self._last_cpu_usage: int | None = None + self._last_cpu_time: int | None = None + self._cgroup_version: int | None = None + self._cpu_quota: float | None = None + + def _detect_cgroup_version(self) -> int: + if self._cgroup_version is not None: + return self._cgroup_version + + if Path("/sys/fs/cgroup/cgroup.controllers").exists(): + self._cgroup_version = 2 + elif Path("/sys/fs/cgroup/cpu/cpuacct.usage").exists() or Path("/sys/fs/cgroup/cpuacct/cpuacct.usage").exists(): + self._cgroup_version = 1 + else: + self._cgroup_version = 0 + + return self._cgroup_version + + def _read_cpu_usage_ns(self) -> int | None: + try: + ver = self._detect_cgroup_version() + if ver == 2: + text = Path("/sys/fs/cgroup/cpu.stat").read_text() + for line in text.splitlines(): + if line.startswith("usage_usec"): + return int(line.split()[1]) * 1000 + return None + elif ver == 1: + for path in ("/sys/fs/cgroup/cpu/cpuacct.usage", "/sys/fs/cgroup/cpuacct/cpuacct.usage"): + p = Path(path) + if p.exists(): + return int(p.read_text().strip()) + return None + return None + except Exception: + return None + + def _read_cpu_quota(self) -> float: + if self._cpu_quota is not None: + return self._cpu_quota + + try: + ver = self._detect_cgroup_version() + if ver == 2: + text = Path("/sys/fs/cgroup/cpu.max").read_text().strip() + parts = text.split() + if parts[0] == "max": + self._cpu_quota = float(os.cpu_count() or 1) + return self._cpu_quota + quota = int(parts[0]) + period = int(parts[1]) + if quota <= 0 or period <= 0: + self._cpu_quota = float(os.cpu_count() or 1) + return self._cpu_quota + self._cpu_quota = quota / period + return self._cpu_quota + elif ver == 1: + quota = int(Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read_text().strip()) + period = int(Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read_text().strip()) + if quota <= 0 or period <= 0: + self._cpu_quota = float(os.cpu_count() or 1) + return self._cpu_quota + self._cpu_quota = quota / period + return self._cpu_quota + except Exception: + pass + self._cpu_quota = float(os.cpu_count() or 1) + return self._cpu_quota + + def cpu_percent(self) -> float: + """Return container CPU utilization % since last call. + + First call returns 0.0 (no baseline). Falls back to psutil if cgroup + files are unavailable. + """ + usage_ns = self._read_cpu_usage_ns() + now_ns = time.monotonic_ns() + + if usage_ns is None: + return psutil.cpu_percent() + + prev_usage = self._last_cpu_usage + prev_time = self._last_cpu_time + self._last_cpu_usage = usage_ns + self._last_cpu_time = now_ns + + if prev_usage is None or prev_time is None: + return 0.0 + + delta_usage = usage_ns - prev_usage + delta_time = now_ns - prev_time + if delta_time <= 0 or delta_usage < 0: + return 0.0 + + num_cpus = self._read_cpu_quota() + return min(round((delta_usage / delta_time) / num_cpus * 100, 1), 100.0) diff --git a/tests/unit/utils/test_cgroup_stats.py b/tests/unit/utils/test_cgroup_stats.py new file mode 100644 index 0000000000..b996268279 --- /dev/null +++ b/tests/unit/utils/test_cgroup_stats.py @@ -0,0 +1,288 @@ +"""Tests for rock.utils.cgroup_stats — container-aware CPU metrics.""" + +from unittest.mock import patch + +import pytest + +from rock.utils.cgroup_stats import CgroupCpuStats, Path + + +def _mock_path_exists(mapping: dict[str, bool]): + """Return a side_effect for Path.exists() based on a path-string mapping.""" + original_exists = Path.exists + + def _exists(self): + s = str(self) + if s in mapping: + return mapping[s] + return original_exists(self) + + return _exists + + +def _mock_path_read_text(mapping: dict[str, str]): + """Return a side_effect for Path.read_text() based on a path-string mapping.""" + + def _read_text(self, *args, **kwargs): + s = str(self) + if s in mapping: + return mapping[s] + raise FileNotFoundError(s) + + return _read_text + + +# ---------- cgroup version detection ---------- + + +class TestDetectCgroupVersion: + def test_detects_v2(self): + stats = CgroupCpuStats() + exists_map = {"/sys/fs/cgroup/cgroup.controllers": True} + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 2 + + def test_detects_v1(self): + stats = CgroupCpuStats() + exists_map = { + "/sys/fs/cgroup/cgroup.controllers": False, + "/sys/fs/cgroup/cpu/cpuacct.usage": True, + } + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 1 + + def test_detects_v1_alternate_path(self): + stats = CgroupCpuStats() + exists_map = { + "/sys/fs/cgroup/cgroup.controllers": False, + "/sys/fs/cgroup/cpu/cpuacct.usage": False, + "/sys/fs/cgroup/cpuacct/cpuacct.usage": True, + } + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 1 + + def test_detects_no_cgroup(self): + stats = CgroupCpuStats() + exists_map = { + "/sys/fs/cgroup/cgroup.controllers": False, + "/sys/fs/cgroup/cpu/cpuacct.usage": False, + "/sys/fs/cgroup/cpuacct/cpuacct.usage": False, + } + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 0 + + def test_caches_result(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + assert stats._detect_cgroup_version() == 2 + + +# ---------- CPU percent ---------- + + +class TestCpuPercent: + def test_first_call_returns_zero(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map = {"/sys/fs/cgroup/cpu.stat": "usage_usec 1000000\n"} + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + assert stats.cpu_percent() == 0.0 + + def test_v2_cpu_calculation(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map_1 = { + "/sys/fs/cgroup/cpu.stat": "usage_usec 1000000\n", + "/sys/fs/cgroup/cpu.max": "100000 100000\n", + } + read_map_2 = { + "/sys/fs/cgroup/cpu.stat": "usage_usec 1500000\n", + "/sys/fs/cgroup/cpu.max": "100000 100000\n", + } + time_values = [100_000_000_000, 200_000_000_000] + time_idx = {"i": 0} + + def mock_monotonic_ns(): + val = time_values[time_idx["i"]] + time_idx["i"] += 1 + return val + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_1)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + stats.cpu_percent() + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_2)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + result = stats.cpu_percent() + + # delta_usage = (1500000 - 1000000) * 1000 = 500_000_000 ns + # delta_time = 100_000_000_000 ns + # quota = 100000/100000 = 1.0 CPU + # pct = (500_000_000 / 100_000_000_000) / 1.0 * 100 = 0.5% + assert result == 0.5 + + def test_v1_cpu_calculation(self): + stats = CgroupCpuStats() + stats._cgroup_version = 1 + exists_map = {"/sys/fs/cgroup/cpu/cpuacct.usage": True} + read_map_1 = { + "/sys/fs/cgroup/cpu/cpuacct.usage": "1000000000\n", + "/sys/fs/cgroup/cpu/cpu.cfs_quota_us": "200000\n", + "/sys/fs/cgroup/cpu/cpu.cfs_period_us": "100000\n", + } + read_map_2 = { + "/sys/fs/cgroup/cpu/cpuacct.usage": "2000000000\n", + "/sys/fs/cgroup/cpu/cpu.cfs_quota_us": "200000\n", + "/sys/fs/cgroup/cpu/cpu.cfs_period_us": "100000\n", + } + time_values = [100_000_000_000, 200_000_000_000] + time_idx = {"i": 0} + + def mock_monotonic_ns(): + val = time_values[time_idx["i"]] + time_idx["i"] += 1 + return val + + with ( + patch.object(Path, "exists", _mock_path_exists(exists_map)), + patch.object(Path, "read_text", _mock_path_read_text(read_map_1)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + stats.cpu_percent() + + with ( + patch.object(Path, "exists", _mock_path_exists(exists_map)), + patch.object(Path, "read_text", _mock_path_read_text(read_map_2)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + result = stats.cpu_percent() + + # delta_usage = 1_000_000_000 ns, delta_time = 100_000_000_000 ns + # quota = 200000/100000 = 2.0 CPUs + # pct = (1_000_000_000 / 100_000_000_000) / 2.0 * 100 = 0.5% + assert result == 0.5 + + def test_fallback_to_psutil_when_no_cgroup(self): + stats = CgroupCpuStats() + stats._cgroup_version = 0 + with patch("rock.utils.cgroup_stats.psutil.cpu_percent", return_value=42.0): + assert stats.cpu_percent() == 42.0 + + def test_capped_at_100(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map_1 = { + "/sys/fs/cgroup/cpu.stat": "usage_usec 0\n", + "/sys/fs/cgroup/cpu.max": "100000 100000\n", + } + read_map_2 = { + "/sys/fs/cgroup/cpu.stat": "usage_usec 200000000\n", + "/sys/fs/cgroup/cpu.max": "100000 100000\n", + } + time_values = [100_000_000_000, 200_000_000_000] + time_idx = {"i": 0} + + def mock_monotonic_ns(): + val = time_values[time_idx["i"]] + time_idx["i"] += 1 + return val + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_1)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + stats.cpu_percent() + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_2)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + result = stats.cpu_percent() + + assert result == 100.0 + + def test_zero_delta_time_returns_zero(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map = {"/sys/fs/cgroup/cpu.stat": "usage_usec 1000000\n"} + same_time = 100_000_000_000 + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", return_value=same_time), + ): + stats.cpu_percent() + result = stats.cpu_percent() + + assert result == 0.0 + + def test_negative_delta_usage_returns_zero(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map_1 = {"/sys/fs/cgroup/cpu.stat": "usage_usec 2000000\n"} + read_map_2 = { + "/sys/fs/cgroup/cpu.stat": "usage_usec 1000000\n", + "/sys/fs/cgroup/cpu.max": "100000 100000\n", + } + time_values = [100_000_000_000, 200_000_000_000] + time_idx = {"i": 0} + + def mock_monotonic_ns(): + val = time_values[time_idx["i"]] + time_idx["i"] += 1 + return val + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_1)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + stats.cpu_percent() + + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map_2)), + patch("rock.utils.cgroup_stats.time.monotonic_ns", mock_monotonic_ns), + ): + result = stats.cpu_percent() + + assert result == 0.0 + + +# ---------- CPU quota ---------- + + +class TestCpuQuota: + def test_v2_unlimited(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + read_map = {"/sys/fs/cgroup/cpu.max": "max 100000\n"} + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map)), + patch("rock.utils.cgroup_stats.os.cpu_count", return_value=4), + ): + assert stats._read_cpu_quota() == 4.0 + + def test_v1_unlimited(self): + stats = CgroupCpuStats() + stats._cgroup_version = 1 + read_map = { + "/sys/fs/cgroup/cpu/cpu.cfs_quota_us": "-1\n", + "/sys/fs/cgroup/cpu/cpu.cfs_period_us": "100000\n", + } + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map)), + patch("rock.utils.cgroup_stats.os.cpu_count", return_value=8), + ): + assert stats._read_cpu_quota() == 8.0 + + def test_fallback_on_error(self): + stats = CgroupCpuStats() + stats._cgroup_version = 2 + with ( + patch.object(Path, "read_text", side_effect=PermissionError), + patch("rock.utils.cgroup_stats.os.cpu_count", return_value=2), + ): + assert stats._read_cpu_quota() == 2.0 From c223c1cc5a7efe65498cb908b1a4bdddd0eb3394 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 08:41:05 +0800 Subject: [PATCH 108/226] feature(oss): add archive command builder + OSS config fields for sandbox log archival --- rock/config.py | 19 +++++++ rock/utils/archive_command.py | 41 +++++++++++++++ tests/unit/test_config.py | 23 ++++++++- tests/unit/utils/test_archive_command.py | 63 ++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 rock/utils/archive_command.py create mode 100644 tests/unit/utils/test_archive_command.py diff --git a/rock/config.py b/rock/config.py index 6913da454c..c901696189 100644 --- a/rock/config.py +++ b/rock/config.py @@ -100,6 +100,25 @@ class OssConfig: process env. xrl package is no longer maintained, so internal users must export this env var themselves when upgrading to SDK >= 1.8.""" + archive_prefix: str = "" + """OSS object key prefix under the PRIMARY bucket for sandbox-log + archives. Empty (default) means each deployment's YAML must opt-in + explicitly to a value matching its OSS bucket lifecycle rule (e.g. + "rock-archives/").""" + + archive_ttl_days: int = 30 + """OSS-side lifecycle expiration days for archives. Audit-only: + ROCK does not enforce TTL itself; the bucket lifecycle rule on + `archive_prefix` is the source of truth.""" + + keep_days_before_archive: int = 3 + """Days to wait after sandbox stop before archiving. Gives + operators a short investigation window without bloating disk.""" + + archive_max_attempts: int = 3 + """Max retry attempts before giving up archival and degrading + to KEEP (FileCleanupTask is the eventual janitor).""" + def __post_init__(self): # Allow YAML to pass a dict for `primary` (dataclass deserialization # from yaml.safe_load returns dicts, not nested dataclasses). diff --git a/rock/utils/archive_command.py b/rock/utils/archive_command.py new file mode 100644 index 0000000000..7e93e0481d --- /dev/null +++ b/rock/utils/archive_command.py @@ -0,0 +1,41 @@ +"""Pure functions for building sandbox-log archive bash commands and OSS keys. + +Used by SandboxLogArchiveTask to drive archival via `runtime.execute()` — +no rocklet endpoint is added; the worker only needs `tar` and `ossutil`. + +Credentials must be passed via `SandboxCommand.env` (not in the command +string), so they never appear in `ps` argv output. +""" + +import shlex +from pathlib import Path + + +def build_sandbox_log_key(sandbox_id: str, prefix: str = "") -> str: + """Construct the OSS object key for a sandbox-log archive. + + Layout: ``sandbox-logs/.tar.gz``. ``prefix`` may be + empty (flat layout under bucket root) or end with ``/``. + """ + cleaned = (prefix or "").strip("/") + sub = f"sandbox-logs/{sandbox_id}.tar.gz" + return f"{cleaned}/{sub}" if cleaned else sub + + +def build_archive_command(log_dir: str, oss_key: str, bucket: str, endpoint: str) -> str: + """Build the bash one-liner that tar+gzips ``log_dir`` and streams to OSS. + + On non-zero exit (tar / ossutil failure), the trailing ``rm -rf`` is + skipped — caller relies on the exit code to decide retry vs delete. + AK/SK come from ``OSS_ACCESS_KEY_ID`` / ``OSS_ACCESS_KEY_SECRET`` env + vars set by the caller via ``SandboxCommand.env``. + """ + log_path = Path(log_dir) + parent = str(log_path.parent) + name = log_path.name + oss_url = f"oss://{bucket}/{oss_key}" + return ( + f"tar -czf - -C {shlex.quote(parent)} {shlex.quote(name)} " + f"| ossutil cp -f - {shlex.quote(oss_url)} --endpoint {shlex.quote(endpoint)} " + f"&& rm -rf {shlex.quote(log_dir)}" + ) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index d2d0941f12..41cc1171e8 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -56,10 +56,15 @@ def test_oss_config_defaults(): cfg = OssConfig() assert cfg.bucket == "" assert cfg.region == "" - assert cfg.transfer_prefix == "" # default is now empty; deployments must opt-in via YAML + assert cfg.transfer_prefix == "" # default empty; YAML must opt-in assert isinstance(cfg.primary, OssAccountConfig) assert cfg.primary.bucket == "" assert cfg.primary.region == "" + # archive defaults: prefix empty (YAML opt-in), other timing fields preset + assert cfg.archive_prefix == "" + assert cfg.archive_ttl_days == 30 + assert cfg.keep_days_before_archive == 3 + assert cfg.archive_max_attempts == 3 def test_oss_config_primary_dict_coerced(): @@ -81,3 +86,19 @@ def test_oss_config_primary_dict_coerced(): # legacy 顶层字段未提供时仍为默认空,确认 primary 不会污染 legacy assert cfg.bucket == "" + +def test_oss_config_archive_fields_overridable(): + from rock.config import OssConfig + + cfg = OssConfig( + archive_prefix="custom-prefix/", + archive_ttl_days=7, + keep_days_before_archive=1, + archive_max_attempts=5, + ) + assert cfg.archive_prefix == "custom-prefix/" + assert cfg.archive_ttl_days == 7 + assert cfg.keep_days_before_archive == 1 + assert cfg.archive_max_attempts == 5 + + diff --git a/tests/unit/utils/test_archive_command.py b/tests/unit/utils/test_archive_command.py new file mode 100644 index 0000000000..f06f162afa --- /dev/null +++ b/tests/unit/utils/test_archive_command.py @@ -0,0 +1,63 @@ +"""Tests for rock.utils.archive_command.""" + +from rock.utils.archive_command import build_archive_command, build_sandbox_log_key + + +class TestBuildSandboxLogKey: + def test_with_prefix(self): + assert build_sandbox_log_key("sb-123", "rock-archives/") == "rock-archives/sandbox-logs/sb-123.tar.gz" + + def test_strips_leading_and_trailing_slashes(self): + # Avoid double slashes regardless of how the prefix is configured in YAML. + assert build_sandbox_log_key("sb-1", "/rock-archives//") == "rock-archives/sandbox-logs/sb-1.tar.gz" + + def test_empty_prefix_yields_flat_layout(self): + assert build_sandbox_log_key("sb-1") == "sandbox-logs/sb-1.tar.gz" + assert build_sandbox_log_key("sb-1", "") == "sandbox-logs/sb-1.tar.gz" + + def test_no_double_slash_at_join_boundary(self): + # Strip only guards against leading/trailing slash duplication; internal + # slashes are preserved as-is (caller's responsibility to pass a sane prefix). + assert build_sandbox_log_key("sb-1", "/a/b/").startswith("a/b/sandbox-logs/") + + +class TestBuildArchiveCommand: + def test_contains_tar_pipe_ossutil_then_rm(self): + cmd = build_archive_command( + log_dir="/data/logs/sb-1", + oss_key="rock-archives/sandbox-logs/sb-1.tar.gz", + bucket="chatos-rock", + endpoint="oss-cn-hangzhou.aliyuncs.com", + ) + # tar streams to ossutil, only rm on success (&& chains) + assert "tar -czf -" in cmd + assert "| ossutil cp -f -" in cmd + assert "&& rm -rf" in cmd + + def test_uses_parent_dir_for_tar_so_archive_does_not_embed_full_path(self): + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + # `-C ` keeps the tarball flat at /, not + # /data/logs/sb-1/ — important for restore. + assert "-C /data/logs sb-1" in cmd + + def test_paths_are_shell_quoted(self): + cmd = build_archive_command( + log_dir="/data/logs/has space/sb-1", + oss_key="rock-archives/has space/sb-1.tar.gz", + bucket="b", + endpoint="e", + ) + # shlex.quote wraps anything with spaces in single quotes + assert "'/data/logs/has space/sb-1'" in cmd + assert "'/data/logs/has space'" in cmd + + def test_oss_url_is_built_from_bucket_and_key(self): + cmd = build_archive_command("/data/logs/sb-1", "rock-archives/sandbox-logs/sb-1.tar.gz", "chatos-rock", "e") + assert "oss://chatos-rock/rock-archives/sandbox-logs/sb-1.tar.gz" in cmd + + def test_no_credentials_in_command_string(self): + # AK/SK MUST flow via SandboxCommand.env, never the command string. + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + assert "ACCESS_KEY" not in cmd.upper() + assert "SECRET" not in cmd.upper() + assert "--access-key" not in cmd From 1a320b5a9666e797769b488ca013cbb18597ce9b Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Sat, 16 May 2026 23:46:15 +0800 Subject: [PATCH 109/226] fix: remove unused archive_ttl_days field from OssConfig MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This field was audit-only and never referenced by any code logic. Its presence could mislead operators into thinking ROCK enforces OSS-side TTL, when in fact the bucket lifecycle rule is the only source of truth. Removing it avoids confusion. 🤖 Generated with [Qoder][https://qoder.com] --- rock/config.py | 5 ----- tests/unit/test_config.py | 3 --- 2 files changed, 8 deletions(-) diff --git a/rock/config.py b/rock/config.py index c901696189..978135c74d 100644 --- a/rock/config.py +++ b/rock/config.py @@ -106,11 +106,6 @@ class OssConfig: explicitly to a value matching its OSS bucket lifecycle rule (e.g. "rock-archives/").""" - archive_ttl_days: int = 30 - """OSS-side lifecycle expiration days for archives. Audit-only: - ROCK does not enforce TTL itself; the bucket lifecycle rule on - `archive_prefix` is the source of truth.""" - keep_days_before_archive: int = 3 """Days to wait after sandbox stop before archiving. Gives operators a short investigation window without bloating disk.""" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 41cc1171e8..734e499203 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -62,7 +62,6 @@ def test_oss_config_defaults(): assert cfg.primary.region == "" # archive defaults: prefix empty (YAML opt-in), other timing fields preset assert cfg.archive_prefix == "" - assert cfg.archive_ttl_days == 30 assert cfg.keep_days_before_archive == 3 assert cfg.archive_max_attempts == 3 @@ -92,12 +91,10 @@ def test_oss_config_archive_fields_overridable(): cfg = OssConfig( archive_prefix="custom-prefix/", - archive_ttl_days=7, keep_days_before_archive=1, archive_max_attempts=5, ) assert cfg.archive_prefix == "custom-prefix/" - assert cfg.archive_ttl_days == 7 assert cfg.keep_days_before_archive == 1 assert cfg.archive_max_attempts == 5 From 4da53c22d9f7216f133c20c9acb3a8c340c587e6 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 05:40:22 +0800 Subject: [PATCH 110/226] fix oss archive command --- rock/utils/archive_command.py | 32 +++++++++--- tests/unit/utils/test_archive_command.py | 63 ++++++++++++++++++------ 2 files changed, 72 insertions(+), 23 deletions(-) diff --git a/rock/utils/archive_command.py b/rock/utils/archive_command.py index 7e93e0481d..ba356122bf 100644 --- a/rock/utils/archive_command.py +++ b/rock/utils/archive_command.py @@ -23,19 +23,35 @@ def build_sandbox_log_key(sandbox_id: str, prefix: str = "") -> str: def build_archive_command(log_dir: str, oss_key: str, bucket: str, endpoint: str) -> str: - """Build the bash one-liner that tar+gzips ``log_dir`` and streams to OSS. + """Build the bash one-liner that tar+gzips ``log_dir`` and uploads to OSS. - On non-zero exit (tar / ossutil failure), the trailing ``rm -rf`` is - skipped — caller relies on the exit code to decide retry vs delete. - AK/SK come from ``OSS_ACCESS_KEY_ID`` / ``OSS_ACCESS_KEY_SECRET`` env - vars set by the caller via ``SandboxCommand.env``. + Why a temp-file pipeline instead of ``tar | ossutil cp -``: ossutil + 1.7.x neither reads ``OSS_ACCESS_KEY_ID`` env vars nor accepts stdin + (``-``) as source, so we materialize the tarball under ``mktemp -d`` + and write a temporary ossutil config carrying the credentials. + + AK/SK still flow via ``OSS_ACCESS_KEY_ID`` / ``OSS_ACCESS_KEY_SECRET`` + env vars set by the caller via ``SandboxCommand.env`` — they are + referenced by name in the command string, never substituted in, so + ``ps`` / shell history never sees the literal values. + + ``set -e`` aborts the chain before the final ``rm -rf `` if + any step fails, so the caller can rely on the exit code to retry. + The scratch dir is removed via ``trap EXIT`` regardless of outcome. """ log_path = Path(log_dir) parent = str(log_path.parent) name = log_path.name oss_url = f"oss://{bucket}/{oss_key}" return ( - f"tar -czf - -C {shlex.quote(parent)} {shlex.quote(name)} " - f"| ossutil cp -f - {shlex.quote(oss_url)} --endpoint {shlex.quote(endpoint)} " - f"&& rm -rf {shlex.quote(log_dir)}" + "set -e && " + "ARCHIVE_DIR=$(mktemp -d -t sb-archive-XXXXXX) && " + "trap 'rm -rf \"$ARCHIVE_DIR\"' EXIT && " + "umask 077 && " + "printf '[Credentials]\\nlanguage=EN\\nendpoint=%s\\naccessKeyID=%s\\naccessKeySecret=%s\\n' " + f'{shlex.quote(endpoint)} "$OSS_ACCESS_KEY_ID" "$OSS_ACCESS_KEY_SECRET" ' + '> "$ARCHIVE_DIR/ossconfig" && ' + f'tar -czf "$ARCHIVE_DIR/archive.tar.gz" -C {shlex.quote(parent)} {shlex.quote(name)} && ' + f'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz" {shlex.quote(oss_url)} && ' + f"rm -rf {shlex.quote(log_dir)}" ) diff --git a/tests/unit/utils/test_archive_command.py b/tests/unit/utils/test_archive_command.py index f06f162afa..621f00e91d 100644 --- a/tests/unit/utils/test_archive_command.py +++ b/tests/unit/utils/test_archive_command.py @@ -22,17 +22,41 @@ def test_no_double_slash_at_join_boundary(self): class TestBuildArchiveCommand: - def test_contains_tar_pipe_ossutil_then_rm(self): - cmd = build_archive_command( - log_dir="/data/logs/sb-1", - oss_key="rock-archives/sandbox-logs/sb-1.tar.gz", - bucket="chatos-rock", - endpoint="oss-cn-hangzhou.aliyuncs.com", - ) - # tar streams to ossutil, only rm on success (&& chains) - assert "tar -czf -" in cmd - assert "| ossutil cp -f -" in cmd - assert "&& rm -rf" in cmd + def test_starts_with_set_e_so_failures_short_circuit(self): + # `set -e` is what guarantees the trailing `rm -rf ` is + # skipped on tar / ossutil failure, so retry can re-archive cleanly. + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + assert cmd.startswith("set -e &&") + + def test_uses_mktemp_and_traps_exit_for_cleanup(self): + # ossutil 1.7.x cannot read stdin, so we must land a temp tarball; + # mktemp -d isolates concurrent archives, trap EXIT cleans up. + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + assert "mktemp -d" in cmd + assert "trap " in cmd and "EXIT" in cmd + + def test_writes_ossutil_config_with_endpoint_and_env_var_credentials(self): + # ossutil 1.7.x ignores OSS_ACCESS_KEY_* env vars, so we materialize + # an [Credentials] config file. AK/SK are still passed by env-var + # reference (printf "%s" + "$OSS_ACCESS_KEY_*"), never substituted. + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "oss-cn-hangzhou.aliyuncs.com") + assert "[Credentials]" in cmd + assert "language=EN" in cmd + assert "endpoint=%s" in cmd + assert "accessKeyID=%s" in cmd + assert "accessKeySecret=%s" in cmd + assert ( + "oss-cn-hangzhou.aliyuncs.com" in cmd + ) # endpoint passed through (shlex.quote skips quoting for safe chars) + assert '"$OSS_ACCESS_KEY_ID"' in cmd + assert '"$OSS_ACCESS_KEY_SECRET"' in cmd + + def test_tar_then_ossutil_cp_against_temp_files(self): + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + # tar produces a real file under the scratch dir. + assert 'tar -czf "$ARCHIVE_DIR/archive.tar.gz"' in cmd + # ossutil cp uses -c -f (1.7.x compatible). + assert 'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz"' in cmd def test_uses_parent_dir_for_tar_so_archive_does_not_embed_full_path(self): cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") @@ -55,9 +79,18 @@ def test_oss_url_is_built_from_bucket_and_key(self): cmd = build_archive_command("/data/logs/sb-1", "rock-archives/sandbox-logs/sb-1.tar.gz", "chatos-rock", "e") assert "oss://chatos-rock/rock-archives/sandbox-logs/sb-1.tar.gz" in cmd - def test_no_credentials_in_command_string(self): - # AK/SK MUST flow via SandboxCommand.env, never the command string. + def test_log_dir_rm_is_last_and_chained_on_success(self): + # The final rm is gated by the && chain — set -e + && ensures it + # never runs if any prior step failed. + cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + assert cmd.endswith("&& rm -rf /data/logs/sb-1") + + def test_no_literal_credentials_in_command_string(self): + # AK/SK appear ONLY as env-var references; never as literals or + # ossutil --access-key flag (which would expose them in argv). cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") - assert "ACCESS_KEY" not in cmd.upper() - assert "SECRET" not in cmd.upper() + assert '"$OSS_ACCESS_KEY_ID"' in cmd + assert '"$OSS_ACCESS_KEY_SECRET"' in cmd assert "--access-key" not in cmd + # `LTAI` is the prefix of every alibaba live AK — cheap regression check. + assert "LTAI" not in cmd From fedc655c1ab39895c824f04d1f9e516152d592e7 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 07:07:39 +0800 Subject: [PATCH 111/226] include ArchivePrefix in /get_token response --- rock/sandbox/service/sandbox_proxy_service.py | 10 ++++-- tests/unit/sandbox/test_sts_dual_account.py | 35 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index dc5de7860d..9035a0a869 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -682,14 +682,17 @@ def _api_url(self, host_ip: str, service_status: ServiceStatus) -> str: port = service_status.get_mapped_port(Port.PROXY) return f"http://{host_ip}:{port}" - def gen_oss_sts_token(self, account: str = "legacy") -> dict | None: # CHANGED: account param, default "legacy" preserves BC + def gen_oss_sts_token( + self, account: str = "legacy" + ) -> dict | None: # CHANGED: account param, default "legacy" preserves BC """Generate STS credentials and OSS config for the given account. Args: account: "legacy" (xrl-sandbox, BC for SDK < 1.8) or "primary" (chatos-rock, SDK >= 1.8). Returns: Dict with STS credentials (AccessKeyId, AccessKeySecret, SecurityToken, Expiration) PLUS account-scoped OSS config: - Endpoint, Bucket, Region, Prefix. None on failure or when the requested account is unconfigured. + Endpoint, Bucket, Region, Prefix (transfer/upload prefix), ArchivePrefix (recovery prefix used by `rock storage get`). + None on failure or when the requested account is unconfigured. """ if account not in self._sts_clients: logger.error(f"unknown OSS account: {account!r}") @@ -736,6 +739,9 @@ def gen_oss_sts_token(self, account: str = "legacy") -> dict | None: # CHANGED: "Bucket": bucket, "Region": region, "Prefix": prefix, # transfer-object key prefix, scoped per account + # ArchivePrefix is the same across both accounts (it lives on oss_config root, not per-account). + # Letting the client pull it from STS lets `rock storage get` skip --archive-prefix. + "ArchivePrefix": self.oss_config.archive_prefix or None, } async def get_sandbox_websocket_url( diff --git a/tests/unit/sandbox/test_sts_dual_account.py b/tests/unit/sandbox/test_sts_dual_account.py index 8c3b59009f..d19fa6ac7d 100644 --- a/tests/unit/sandbox/test_sts_dual_account.py +++ b/tests/unit/sandbox/test_sts_dual_account.py @@ -17,6 +17,7 @@ def _make_rock_config( primary_role: str, legacy_region: str = "cn-hangzhou", transfer_prefix: str = "rock-transfer/", + archive_prefix: str = "rock-archives/", ) -> RockConfig: return RockConfig( oss=OssConfig( @@ -27,6 +28,7 @@ def _make_rock_config( role_arn=legacy_role, region=legacy_region, transfer_prefix=transfer_prefix, + archive_prefix=archive_prefix, primary=OssAccountConfig( endpoint="oss-cn-hangzhou.aliyuncs.com", bucket="chatos-rock", @@ -125,6 +127,39 @@ def test_primary_prefix_is_none_when_yaml_does_not_set_it(): assert creds["Prefix"] is None +def test_archive_prefix_returned_for_both_accounts(): + # ArchivePrefix lives on oss_config root (not per-account), so legacy and primary + # should both surface the same value — `rock storage get` reads it to skip + # --archive-prefix, regardless of which account it queries. + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + archive_prefix="rock-archives/", + ) + ) + svc._sts_clients["legacy"].do_action_with_exception = _fake_assume("L-AK", "L-SK", "L-TOK") + svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") + + legacy = svc.gen_oss_sts_token() + primary = svc.gen_oss_sts_token(account="primary") + assert legacy["ArchivePrefix"] == "rock-archives/" + assert primary["ArchivePrefix"] == "rock-archives/" + + +def test_archive_prefix_none_when_yaml_does_not_set_it(): + svc = _build_service( + _make_rock_config( + legacy_role="acs:ram::1933967579503727:role/legacy-role", + primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", + archive_prefix="", + ) + ) + svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") + creds = svc.gen_oss_sts_token(account="primary") + assert creds["ArchivePrefix"] is None + + def test_unknown_account_returns_none(): svc = _build_service( _make_rock_config( From 6d32eeed7a3f880baa88c126f8a3858f2813f5ea Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 13:11:38 +0800 Subject: [PATCH 112/226] wrap archive_command in ArchiveCommand class --- rock/utils/archive_command.py | 85 +++++++++++++----------- tests/unit/utils/test_archive_command.py | 36 +++++----- 2 files changed, 66 insertions(+), 55 deletions(-) diff --git a/rock/utils/archive_command.py b/rock/utils/archive_command.py index ba356122bf..73a9f5fb8d 100644 --- a/rock/utils/archive_command.py +++ b/rock/utils/archive_command.py @@ -1,4 +1,4 @@ -"""Pure functions for building sandbox-log archive bash commands and OSS keys. +"""Sandbox-log archive bash command + OSS key builder. Used by SandboxLogArchiveTask to drive archival via `runtime.execute()` — no rocklet endpoint is added; the worker only needs `tar` and `ossutil`. @@ -11,47 +11,56 @@ from pathlib import Path -def build_sandbox_log_key(sandbox_id: str, prefix: str = "") -> str: - """Construct the OSS object key for a sandbox-log archive. +class ArchiveCommand: + """Namespace for building sandbox-log archive commands and OSS keys. - Layout: ``sandbox-logs/.tar.gz``. ``prefix`` may be - empty (flat layout under bucket root) or end with ``/``. + Stateless: all methods are `@staticmethod`. Grouped under a class so + admin / CLI sides go through one explicit entry point (``ArchiveCommand.build_key``, + ``ArchiveCommand.build_command``) and cannot drift on key layout / command shape. """ - cleaned = (prefix or "").strip("/") - sub = f"sandbox-logs/{sandbox_id}.tar.gz" - return f"{cleaned}/{sub}" if cleaned else sub + @staticmethod + def build_key(sandbox_id: str, prefix: str = "") -> str: + """Construct the OSS object key for a sandbox-log archive. -def build_archive_command(log_dir: str, oss_key: str, bucket: str, endpoint: str) -> str: - """Build the bash one-liner that tar+gzips ``log_dir`` and uploads to OSS. + Layout: ``sandbox-logs/.tar.gz``. ``prefix`` may be + empty (flat layout under bucket root) or end with ``/``. + """ + cleaned = (prefix or "").strip("/") + sub = f"sandbox-logs/{sandbox_id}.tar.gz" + return f"{cleaned}/{sub}" if cleaned else sub - Why a temp-file pipeline instead of ``tar | ossutil cp -``: ossutil - 1.7.x neither reads ``OSS_ACCESS_KEY_ID`` env vars nor accepts stdin - (``-``) as source, so we materialize the tarball under ``mktemp -d`` - and write a temporary ossutil config carrying the credentials. + @staticmethod + def build_command(log_dir: str, oss_key: str, bucket: str, endpoint: str) -> str: + """Build the bash one-liner that tar+gzips ``log_dir`` and uploads to OSS. - AK/SK still flow via ``OSS_ACCESS_KEY_ID`` / ``OSS_ACCESS_KEY_SECRET`` - env vars set by the caller via ``SandboxCommand.env`` — they are - referenced by name in the command string, never substituted in, so - ``ps`` / shell history never sees the literal values. + Why a temp-file pipeline instead of ``tar | ossutil cp -``: ossutil + 1.7.x neither reads ``OSS_ACCESS_KEY_ID`` env vars nor accepts stdin + (``-``) as source, so we materialize the tarball under ``mktemp -d`` + and write a temporary ossutil config carrying the credentials. - ``set -e`` aborts the chain before the final ``rm -rf `` if - any step fails, so the caller can rely on the exit code to retry. - The scratch dir is removed via ``trap EXIT`` regardless of outcome. - """ - log_path = Path(log_dir) - parent = str(log_path.parent) - name = log_path.name - oss_url = f"oss://{bucket}/{oss_key}" - return ( - "set -e && " - "ARCHIVE_DIR=$(mktemp -d -t sb-archive-XXXXXX) && " - "trap 'rm -rf \"$ARCHIVE_DIR\"' EXIT && " - "umask 077 && " - "printf '[Credentials]\\nlanguage=EN\\nendpoint=%s\\naccessKeyID=%s\\naccessKeySecret=%s\\n' " - f'{shlex.quote(endpoint)} "$OSS_ACCESS_KEY_ID" "$OSS_ACCESS_KEY_SECRET" ' - '> "$ARCHIVE_DIR/ossconfig" && ' - f'tar -czf "$ARCHIVE_DIR/archive.tar.gz" -C {shlex.quote(parent)} {shlex.quote(name)} && ' - f'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz" {shlex.quote(oss_url)} && ' - f"rm -rf {shlex.quote(log_dir)}" - ) + AK/SK still flow via ``OSS_ACCESS_KEY_ID`` / ``OSS_ACCESS_KEY_SECRET`` + env vars set by the caller via ``SandboxCommand.env`` — they are + referenced by name in the command string, never substituted in, so + ``ps`` / shell history never sees the literal values. + + ``set -e`` aborts the chain before the final ``rm -rf `` if + any step fails, so the caller can rely on the exit code to retry. + The scratch dir is removed via ``trap EXIT`` regardless of outcome. + """ + log_path = Path(log_dir) + parent = str(log_path.parent) + name = log_path.name + oss_url = f"oss://{bucket}/{oss_key}" + return ( + "set -e && " + "ARCHIVE_DIR=$(mktemp -d -t sb-archive-XXXXXX) && " + "trap 'rm -rf \"$ARCHIVE_DIR\"' EXIT && " + "umask 077 && " + "printf '[Credentials]\\nlanguage=EN\\nendpoint=%s\\naccessKeyID=%s\\naccessKeySecret=%s\\n' " + f'{shlex.quote(endpoint)} "$OSS_ACCESS_KEY_ID" "$OSS_ACCESS_KEY_SECRET" ' + '> "$ARCHIVE_DIR/ossconfig" && ' + f'tar -czf "$ARCHIVE_DIR/archive.tar.gz" -C {shlex.quote(parent)} {shlex.quote(name)} && ' + f'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz" {shlex.quote(oss_url)} && ' + f"rm -rf {shlex.quote(log_dir)}" + ) diff --git a/tests/unit/utils/test_archive_command.py b/tests/unit/utils/test_archive_command.py index 621f00e91d..ef82f096d5 100644 --- a/tests/unit/utils/test_archive_command.py +++ b/tests/unit/utils/test_archive_command.py @@ -1,37 +1,37 @@ """Tests for rock.utils.archive_command.""" -from rock.utils.archive_command import build_archive_command, build_sandbox_log_key +from rock.utils.archive_command import ArchiveCommand -class TestBuildSandboxLogKey: +class TestArchiveCommandBuildKey: def test_with_prefix(self): - assert build_sandbox_log_key("sb-123", "rock-archives/") == "rock-archives/sandbox-logs/sb-123.tar.gz" + assert ArchiveCommand.build_key("sb-123", "rock-archives/") == "rock-archives/sandbox-logs/sb-123.tar.gz" def test_strips_leading_and_trailing_slashes(self): # Avoid double slashes regardless of how the prefix is configured in YAML. - assert build_sandbox_log_key("sb-1", "/rock-archives//") == "rock-archives/sandbox-logs/sb-1.tar.gz" + assert ArchiveCommand.build_key("sb-1", "/rock-archives//") == "rock-archives/sandbox-logs/sb-1.tar.gz" def test_empty_prefix_yields_flat_layout(self): - assert build_sandbox_log_key("sb-1") == "sandbox-logs/sb-1.tar.gz" - assert build_sandbox_log_key("sb-1", "") == "sandbox-logs/sb-1.tar.gz" + assert ArchiveCommand.build_key("sb-1") == "sandbox-logs/sb-1.tar.gz" + assert ArchiveCommand.build_key("sb-1", "") == "sandbox-logs/sb-1.tar.gz" def test_no_double_slash_at_join_boundary(self): # Strip only guards against leading/trailing slash duplication; internal # slashes are preserved as-is (caller's responsibility to pass a sane prefix). - assert build_sandbox_log_key("sb-1", "/a/b/").startswith("a/b/sandbox-logs/") + assert ArchiveCommand.build_key("sb-1", "/a/b/").startswith("a/b/sandbox-logs/") -class TestBuildArchiveCommand: +class TestArchiveCommandBuildCommand: def test_starts_with_set_e_so_failures_short_circuit(self): # `set -e` is what guarantees the trailing `rm -rf ` is # skipped on tar / ossutil failure, so retry can re-archive cleanly. - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") assert cmd.startswith("set -e &&") def test_uses_mktemp_and_traps_exit_for_cleanup(self): # ossutil 1.7.x cannot read stdin, so we must land a temp tarball; # mktemp -d isolates concurrent archives, trap EXIT cleans up. - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") assert "mktemp -d" in cmd assert "trap " in cmd and "EXIT" in cmd @@ -39,7 +39,7 @@ def test_writes_ossutil_config_with_endpoint_and_env_var_credentials(self): # ossutil 1.7.x ignores OSS_ACCESS_KEY_* env vars, so we materialize # an [Credentials] config file. AK/SK are still passed by env-var # reference (printf "%s" + "$OSS_ACCESS_KEY_*"), never substituted. - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "oss-cn-hangzhou.aliyuncs.com") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "oss-cn-hangzhou.aliyuncs.com") assert "[Credentials]" in cmd assert "language=EN" in cmd assert "endpoint=%s" in cmd @@ -52,20 +52,20 @@ def test_writes_ossutil_config_with_endpoint_and_env_var_credentials(self): assert '"$OSS_ACCESS_KEY_SECRET"' in cmd def test_tar_then_ossutil_cp_against_temp_files(self): - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") # tar produces a real file under the scratch dir. assert 'tar -czf "$ARCHIVE_DIR/archive.tar.gz"' in cmd # ossutil cp uses -c -f (1.7.x compatible). assert 'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz"' in cmd def test_uses_parent_dir_for_tar_so_archive_does_not_embed_full_path(self): - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") # `-C ` keeps the tarball flat at /, not # /data/logs/sb-1/ — important for restore. assert "-C /data/logs sb-1" in cmd def test_paths_are_shell_quoted(self): - cmd = build_archive_command( + cmd = ArchiveCommand.build_command( log_dir="/data/logs/has space/sb-1", oss_key="rock-archives/has space/sb-1.tar.gz", bucket="b", @@ -76,19 +76,21 @@ def test_paths_are_shell_quoted(self): assert "'/data/logs/has space'" in cmd def test_oss_url_is_built_from_bucket_and_key(self): - cmd = build_archive_command("/data/logs/sb-1", "rock-archives/sandbox-logs/sb-1.tar.gz", "chatos-rock", "e") + cmd = ArchiveCommand.build_command( + "/data/logs/sb-1", "rock-archives/sandbox-logs/sb-1.tar.gz", "chatos-rock", "e" + ) assert "oss://chatos-rock/rock-archives/sandbox-logs/sb-1.tar.gz" in cmd def test_log_dir_rm_is_last_and_chained_on_success(self): # The final rm is gated by the && chain — set -e + && ensures it # never runs if any prior step failed. - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") assert cmd.endswith("&& rm -rf /data/logs/sb-1") def test_no_literal_credentials_in_command_string(self): # AK/SK appear ONLY as env-var references; never as literals or # ossutil --access-key flag (which would expose them in argv). - cmd = build_archive_command("/data/logs/sb-1", "k", "b", "e") + cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") assert '"$OSS_ACCESS_KEY_ID"' in cmd assert '"$OSS_ACCESS_KEY_SECRET"' in cmd assert "--access-key" not in cmd From 3e409f561c72261f934eae33a2e442e74e5edb4c Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 13:17:50 +0800 Subject: [PATCH 113/226] use triple-quoted multiline bash for build_command --- rock/utils/archive_command.py | 34 ++++++++++++++---------- tests/unit/utils/test_archive_command.py | 4 ++- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/rock/utils/archive_command.py b/rock/utils/archive_command.py index 73a9f5fb8d..b21f7abcba 100644 --- a/rock/utils/archive_command.py +++ b/rock/utils/archive_command.py @@ -8,6 +8,7 @@ """ import shlex +import textwrap from pathlib import Path @@ -49,18 +50,23 @@ def build_command(log_dir: str, oss_key: str, bucket: str, endpoint: str) -> str The scratch dir is removed via ``trap EXIT`` regardless of outcome. """ log_path = Path(log_dir) - parent = str(log_path.parent) - name = log_path.name - oss_url = f"oss://{bucket}/{oss_key}" - return ( - "set -e && " - "ARCHIVE_DIR=$(mktemp -d -t sb-archive-XXXXXX) && " - "trap 'rm -rf \"$ARCHIVE_DIR\"' EXIT && " - "umask 077 && " - "printf '[Credentials]\\nlanguage=EN\\nendpoint=%s\\naccessKeyID=%s\\naccessKeySecret=%s\\n' " - f'{shlex.quote(endpoint)} "$OSS_ACCESS_KEY_ID" "$OSS_ACCESS_KEY_SECRET" ' - '> "$ARCHIVE_DIR/ossconfig" && ' - f'tar -czf "$ARCHIVE_DIR/archive.tar.gz" -C {shlex.quote(parent)} {shlex.quote(name)} && ' - f'ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz" {shlex.quote(oss_url)} && ' - f"rm -rf {shlex.quote(log_dir)}" + parent = shlex.quote(str(log_path.parent)) + name = shlex.quote(log_path.name) + endpoint_q = shlex.quote(endpoint) + oss_url_q = shlex.quote(f"oss://{bucket}/{oss_key}") + log_dir_q = shlex.quote(log_dir) + # textwrap.dedent strips the common leading whitespace so the source can be + # indented for readability without polluting the emitted shell string. + return textwrap.dedent( + f"""\ + set -e \\ + && ARCHIVE_DIR=$(mktemp -d -t sb-archive-XXXXXX) \\ + && trap 'rm -rf "$ARCHIVE_DIR"' EXIT \\ + && umask 077 \\ + && printf '[Credentials]\\nlanguage=EN\\nendpoint=%s\\naccessKeyID=%s\\naccessKeySecret=%s\\n' \\ + {endpoint_q} "$OSS_ACCESS_KEY_ID" "$OSS_ACCESS_KEY_SECRET" \\ + > "$ARCHIVE_DIR/ossconfig" \\ + && tar -czf "$ARCHIVE_DIR/archive.tar.gz" -C {parent} {name} \\ + && ossutil cp -c "$ARCHIVE_DIR/ossconfig" -f "$ARCHIVE_DIR/archive.tar.gz" {oss_url_q} \\ + && rm -rf {log_dir_q}""" ) diff --git a/tests/unit/utils/test_archive_command.py b/tests/unit/utils/test_archive_command.py index ef82f096d5..5910928187 100644 --- a/tests/unit/utils/test_archive_command.py +++ b/tests/unit/utils/test_archive_command.py @@ -25,8 +25,10 @@ class TestArchiveCommandBuildCommand: def test_starts_with_set_e_so_failures_short_circuit(self): # `set -e` is what guarantees the trailing `rm -rf ` is # skipped on tar / ossutil failure, so retry can re-archive cleanly. + # Triple-quoted multi-line string emits `set -e \\\n&& ...` (bash treats + # `\` as line continuation, so it's still one chained command). cmd = ArchiveCommand.build_command("/data/logs/sb-1", "k", "b", "e") - assert cmd.startswith("set -e &&") + assert cmd.startswith("set -e \\\n&&") def test_uses_mktemp_and_traps_exit_for_cleanup(self): # ossutil 1.7.x cannot read stdin, so we must land a temp tarball; From bdbacdfd748780fac8bd2746ff5e39719425c3b4 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 14:07:47 +0800 Subject: [PATCH 114/226] extract SandboxLogConfig and SandboxFileTransferConfig from OssConfig --- rock/config.py | 84 +++++++++++++------ rock/sandbox/service/sandbox_proxy_service.py | 9 +- tests/unit/sandbox/test_sts_dual_account.py | 9 +- tests/unit/test_config.py | 62 ++++++++++++-- 4 files changed, 124 insertions(+), 40 deletions(-) diff --git a/rock/config.py b/rock/config.py index 978135c74d..a09ffc2d04 100644 --- a/rock/config.py +++ b/rock/config.py @@ -52,12 +52,70 @@ class RedisConfig: password: str = "" +@dataclass +class SandboxLogConfig: + """Policy for archiving stopped-sandbox log directories to OSS. + + Lives under SandboxConfig.log: the fields are domain knobs of "what to do + with stopped sandbox logs" — when to archive, how many retries, what OSS + key prefix to use — colocated with other sandbox lifecycle / cleanup + policy (image_keep_patterns, remove_container_enabled). OSS endpoint / + bucket / credentials still belong to OssConfig.primary. + """ + + archive_prefix: str = "" + """OSS object key prefix under OssConfig.primary.bucket for sandbox-log + archives. Empty (default) means each deployment's YAML must opt-in + explicitly to a value matching its OSS bucket lifecycle rule (e.g. + "rock-archives/").""" + + keep_days_before_archive: int = 3 + """Days to wait after sandbox stop before archiving. Gives operators + a short investigation window without bloating disk.""" + + archive_max_attempts: int = 3 + """Max retry attempts before giving up archival and degrading to KEEP + (FileCleanupTask is the eventual janitor).""" + + +@dataclass +class SandboxFileTransferConfig: + """Policy for sandbox <-> host file transfer via OSS as intermediary. + + Lives under SandboxConfig.file_transfer: the prefix governs where the + SDK puts ephemeral transfer objects under OssConfig.primary.bucket, + a sandbox-side concern not OSS connectivity. + """ + + prefix: str = "" + """Prefix under OssConfig.primary.bucket for ephemeral host↔container + file transfers ({timestamp}-{filename} objects). The legacy bucket keeps + its pre-existing flat layout (no prefix) for backward compatibility — + xrl-sandbox has a 3-day lifecycle rule at bucket root (configured in + the Aliyun OSS console, not in repo) that we do not disturb. + + Note: this field lives in admin-side RockConfig and is NOT what the SDK + reads. The SDK reads ROCK_OSS_TRANSFER_PREFIX directly from the process + env. xrl package is no longer maintained, so internal users must export + this env var themselves when upgrading to SDK >= 1.8.""" + + @dataclass class SandboxConfig: actor_resource: str = "" actor_resource_num: float = 0.0 gateway_num: int = 1 remove_container_enabled: bool = True + log: SandboxLogConfig = field(default_factory=SandboxLogConfig) + file_transfer: SandboxFileTransferConfig = field(default_factory=SandboxFileTransferConfig) + + def __post_init__(self): + # Allow YAML to pass dicts for nested dataclasses (yaml.safe_load + # returns dicts, not nested dataclass instances). + if isinstance(self.log, dict): + self.log = SandboxLogConfig(**self.log) + if isinstance(self.file_transfer, dict): + self.file_transfer = SandboxFileTransferConfig(**self.file_transfer) @dataclass @@ -88,32 +146,6 @@ class OssConfig: host-side archival. An empty `primary.bucket` disables v2 STS and archival, leaving legacy path fully operational.""" - transfer_prefix: str = "" - """Prefix under the PRIMARY bucket for ephemeral host↔container file - transfers ({timestamp}-{filename} objects). The legacy bucket keeps - its pre-existing flat layout (no prefix) for backward compatibility — - xrl-sandbox has a 3-day lifecycle rule at bucket root (configured - in the Aliyun OSS console, not in repo) that we do not disturb. - - Note: this field lives in admin-side RockConfig and is NOT what the - SDK reads. The SDK reads ROCK_OSS_TRANSFER_PREFIX directly from the - process env. xrl package is no longer maintained, so internal users - must export this env var themselves when upgrading to SDK >= 1.8.""" - - archive_prefix: str = "" - """OSS object key prefix under the PRIMARY bucket for sandbox-log - archives. Empty (default) means each deployment's YAML must opt-in - explicitly to a value matching its OSS bucket lifecycle rule (e.g. - "rock-archives/").""" - - keep_days_before_archive: int = 3 - """Days to wait after sandbox stop before archiving. Gives - operators a short investigation window without bloating disk.""" - - archive_max_attempts: int = 3 - """Max retry attempts before giving up archival and degrading - to KEEP (FileCleanupTask is the eventual janitor).""" - def __post_init__(self): # Allow YAML to pass a dict for `primary` (dataclass deserialization # from yaml.safe_load returns dicts, not nested dataclasses). diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 9035a0a869..a7110a5322 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -705,7 +705,7 @@ def gen_oss_sts_token( endpoint = primary.endpoint or None bucket = primary.bucket or None region = primary.region or env_vars.ROCK_OSS_BUCKET_REGION or None - prefix = self.oss_config.transfer_prefix or None + prefix = self._rock_config.sandbox_config.file_transfer.prefix or None else: # legacy role_arn = self.oss_config.role_arn session_name = "rock-sandbox-legacy" @@ -739,9 +739,10 @@ def gen_oss_sts_token( "Bucket": bucket, "Region": region, "Prefix": prefix, # transfer-object key prefix, scoped per account - # ArchivePrefix is the same across both accounts (it lives on oss_config root, not per-account). - # Letting the client pull it from STS lets `rock storage get` skip --archive-prefix. - "ArchivePrefix": self.oss_config.archive_prefix or None, + # ArchivePrefix is the same across both accounts (it lives on the dedicated + # SandboxLogConfig under SandboxConfig.log, not per-OSS-account). Letting + # the client pull it from STS lets `rock storage get` skip --archive-prefix. + "ArchivePrefix": self._rock_config.sandbox_config.log.archive_prefix or None, } async def get_sandbox_websocket_url( diff --git a/tests/unit/sandbox/test_sts_dual_account.py b/tests/unit/sandbox/test_sts_dual_account.py index d19fa6ac7d..85331fb941 100644 --- a/tests/unit/sandbox/test_sts_dual_account.py +++ b/tests/unit/sandbox/test_sts_dual_account.py @@ -7,6 +7,9 @@ ProxyServiceConfig, RockConfig, RuntimeConfig, + SandboxConfig, + SandboxFileTransferConfig, + SandboxLogConfig, ) from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService @@ -27,8 +30,6 @@ def _make_rock_config( access_key_secret="legacy-sk", role_arn=legacy_role, region=legacy_region, - transfer_prefix=transfer_prefix, - archive_prefix=archive_prefix, primary=OssAccountConfig( endpoint="oss-cn-hangzhou.aliyuncs.com", bucket="chatos-rock", @@ -38,6 +39,10 @@ def _make_rock_config( region="cn-hangzhou", ), ), + sandbox_config=SandboxConfig( + log=SandboxLogConfig(archive_prefix=archive_prefix), + file_transfer=SandboxFileTransferConfig(prefix=transfer_prefix), + ), proxy_service=ProxyServiceConfig(), runtime=RuntimeConfig( python_env_path="/usr/bin/python3", diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 734e499203..5bfdf77b90 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -56,14 +56,16 @@ def test_oss_config_defaults(): cfg = OssConfig() assert cfg.bucket == "" assert cfg.region == "" - assert cfg.transfer_prefix == "" # default empty; YAML must opt-in assert isinstance(cfg.primary, OssAccountConfig) assert cfg.primary.bucket == "" assert cfg.primary.region == "" - # archive defaults: prefix empty (YAML opt-in), other timing fields preset - assert cfg.archive_prefix == "" - assert cfg.keep_days_before_archive == 3 - assert cfg.archive_max_attempts == 3 + # transfer_prefix moved to SandboxConfig.file_transfer.prefix; archive_prefix / + # keep_days_before_archive / archive_max_attempts moved to SandboxConfig.log. + # OssConfig is now purely OSS connectivity (endpoint / bucket / credentials). + assert not hasattr(cfg, "transfer_prefix") + assert not hasattr(cfg, "archive_prefix") + assert not hasattr(cfg, "keep_days_before_archive") + assert not hasattr(cfg, "archive_max_attempts") def test_oss_config_primary_dict_coerced(): @@ -86,10 +88,21 @@ def test_oss_config_primary_dict_coerced(): assert cfg.bucket == "" -def test_oss_config_archive_fields_overridable(): - from rock.config import OssConfig +def test_sandbox_log_config_defaults(): + from rock.config import SandboxLogConfig - cfg = OssConfig( + cfg = SandboxLogConfig() + # prefix defaults empty: each deployment YAML must opt-in to a value + # matching its OSS bucket lifecycle rule (e.g. "rock-archives/"). + assert cfg.archive_prefix == "" + assert cfg.keep_days_before_archive == 3 + assert cfg.archive_max_attempts == 3 + + +def test_sandbox_log_config_overridable(): + from rock.config import SandboxLogConfig + + cfg = SandboxLogConfig( archive_prefix="custom-prefix/", keep_days_before_archive=1, archive_max_attempts=5, @@ -99,3 +112,36 @@ def test_oss_config_archive_fields_overridable(): assert cfg.archive_max_attempts == 5 +def test_sandbox_file_transfer_config_defaults(): + from rock.config import SandboxFileTransferConfig + + cfg = SandboxFileTransferConfig() + # prefix defaults empty; YAML opts in to "rock-transfer/" for the + # primary bucket's lifecycle-managed transfer area. + assert cfg.prefix == "" + + +def test_sandbox_config_nests_log_and_file_transfer(): + # The reorg puts log + file_transfer under SandboxConfig (not OssConfig + # or root) since they're sandbox-domain concerns. + from rock.config import SandboxConfig, SandboxFileTransferConfig, SandboxLogConfig + + cfg = SandboxConfig() + assert isinstance(cfg.log, SandboxLogConfig) + assert isinstance(cfg.file_transfer, SandboxFileTransferConfig) + + +def test_sandbox_config_coerces_nested_dicts_from_yaml(): + # yaml.safe_load returns dicts for nested keys; __post_init__ must coerce + # them into the right dataclass so callers can keep dotted access. + from rock.config import SandboxConfig, SandboxFileTransferConfig, SandboxLogConfig + + cfg = SandboxConfig( + log={"archive_prefix": "rock-archives/", "keep_days_before_archive": 7}, + file_transfer={"prefix": "rock-transfer/"}, + ) + assert isinstance(cfg.log, SandboxLogConfig) + assert cfg.log.archive_prefix == "rock-archives/" + assert cfg.log.keep_days_before_archive == 7 + assert isinstance(cfg.file_transfer, SandboxFileTransferConfig) + assert cfg.file_transfer.prefix == "rock-transfer/" From b0b2b2da3f85cb038d9fdb66d5f4158fbe584db7 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 14:53:09 +0800 Subject: [PATCH 115/226] set _rock_config in test_sandbox_proxy fixture for SandboxLogConfig access --- tests/unit/sandbox/test_sandbox_proxy.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index f12081a30c..356184af1a 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -97,6 +97,12 @@ def sandbox_proxy_service(self): service.oss_config = OssConfig(role_arn="test_role_arn") # gen_oss_sts_token routes by account name; legacy is the default. service._sts_clients = {"legacy": MagicMock(), "primary": MagicMock()} + # gen_oss_sts_token reads sandbox_config.log.archive_prefix and + # sandbox_config.file_transfer.prefix from rock_config after the + # SandboxLogConfig / SandboxFileTransferConfig extraction. + service._rock_config = MagicMock() + service._rock_config.sandbox_config.log.archive_prefix = "" + service._rock_config.sandbox_config.file_transfer.prefix = "" return service def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): @@ -108,7 +114,9 @@ def test_success_returns_dict_with_extra_fields(self, sandbox_proxy_service): b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), + patch.object( + sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body + ), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" @@ -138,7 +146,9 @@ def test_partial_oss_config_returns_creds_with_none_extras(self, sandbox_proxy_s b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), + patch.object( + sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body + ), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" @@ -161,7 +171,9 @@ def test_env_var_overrides_yaml_for_endpoint_and_bucket(self, sandbox_proxy_serv b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), + patch.object( + sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body + ), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "env.endpoint" @@ -184,7 +196,9 @@ def test_yaml_used_when_env_var_empty(self, sandbox_proxy_service): b'"SecurityToken":"tok","Expiration":"2099-01-01T00:00:00Z"}}' ) with ( - patch.object(sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body), + patch.object( + sandbox_proxy_service._sts_clients["legacy"], "do_action_with_exception", return_value=fake_token_body + ), patch("rock.sandbox.service.sandbox_proxy_service.env_vars") as mock_env, ): mock_env.ROCK_OSS_BUCKET_ENDPOINT = "" From 4c1d44f1ed5d69568da2e10aa8e9ae601f8004fc Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 15:39:33 +0800 Subject: [PATCH 116/226] remove ArchivePrefix from /get_token (wrong domain layering) --- rock/sandbox/service/sandbox_proxy_service.py | 11 +++--- tests/unit/sandbox/test_sandbox_proxy.py | 7 ++-- tests/unit/sandbox/test_sts_dual_account.py | 36 ------------------- 3 files changed, 9 insertions(+), 45 deletions(-) diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index a7110a5322..e42564a7bd 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -686,12 +686,17 @@ def gen_oss_sts_token( self, account: str = "legacy" ) -> dict | None: # CHANGED: account param, default "legacy" preserves BC """Generate STS credentials and OSS config for the given account. + + Returns ONLY OSS connectivity / account-scoped config. Sandbox-side + domain policy (e.g. archive prefix, retry counts) does NOT belong here + and must be fetched from a separate endpoint when a client needs it. + Args: account: "legacy" (xrl-sandbox, BC for SDK < 1.8) or "primary" (chatos-rock, SDK >= 1.8). Returns: Dict with STS credentials (AccessKeyId, AccessKeySecret, SecurityToken, Expiration) PLUS account-scoped OSS config: - Endpoint, Bucket, Region, Prefix (transfer/upload prefix), ArchivePrefix (recovery prefix used by `rock storage get`). + Endpoint, Bucket, Region, Prefix (transfer/upload prefix). None on failure or when the requested account is unconfigured. """ if account not in self._sts_clients: @@ -739,10 +744,6 @@ def gen_oss_sts_token( "Bucket": bucket, "Region": region, "Prefix": prefix, # transfer-object key prefix, scoped per account - # ArchivePrefix is the same across both accounts (it lives on the dedicated - # SandboxLogConfig under SandboxConfig.log, not per-OSS-account). Letting - # the client pull it from STS lets `rock storage get` skip --archive-prefix. - "ArchivePrefix": self._rock_config.sandbox_config.log.archive_prefix or None, } async def get_sandbox_websocket_url( diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index 356184af1a..5a117a0642 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -97,11 +97,10 @@ def sandbox_proxy_service(self): service.oss_config = OssConfig(role_arn="test_role_arn") # gen_oss_sts_token routes by account name; legacy is the default. service._sts_clients = {"legacy": MagicMock(), "primary": MagicMock()} - # gen_oss_sts_token reads sandbox_config.log.archive_prefix and - # sandbox_config.file_transfer.prefix from rock_config after the - # SandboxLogConfig / SandboxFileTransferConfig extraction. + # gen_oss_sts_token reads sandbox_config.file_transfer.prefix from + # rock_config (primary-account branch) after the SandboxFileTransferConfig + # extraction; legacy branch reads ROCK_OSS_TRANSFER_PREFIX env directly. service._rock_config = MagicMock() - service._rock_config.sandbox_config.log.archive_prefix = "" service._rock_config.sandbox_config.file_transfer.prefix = "" return service diff --git a/tests/unit/sandbox/test_sts_dual_account.py b/tests/unit/sandbox/test_sts_dual_account.py index 85331fb941..4d4f47f87f 100644 --- a/tests/unit/sandbox/test_sts_dual_account.py +++ b/tests/unit/sandbox/test_sts_dual_account.py @@ -9,7 +9,6 @@ RuntimeConfig, SandboxConfig, SandboxFileTransferConfig, - SandboxLogConfig, ) from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService @@ -20,7 +19,6 @@ def _make_rock_config( primary_role: str, legacy_region: str = "cn-hangzhou", transfer_prefix: str = "rock-transfer/", - archive_prefix: str = "rock-archives/", ) -> RockConfig: return RockConfig( oss=OssConfig( @@ -40,7 +38,6 @@ def _make_rock_config( ), ), sandbox_config=SandboxConfig( - log=SandboxLogConfig(archive_prefix=archive_prefix), file_transfer=SandboxFileTransferConfig(prefix=transfer_prefix), ), proxy_service=ProxyServiceConfig(), @@ -132,39 +129,6 @@ def test_primary_prefix_is_none_when_yaml_does_not_set_it(): assert creds["Prefix"] is None -def test_archive_prefix_returned_for_both_accounts(): - # ArchivePrefix lives on oss_config root (not per-account), so legacy and primary - # should both surface the same value — `rock storage get` reads it to skip - # --archive-prefix, regardless of which account it queries. - svc = _build_service( - _make_rock_config( - legacy_role="acs:ram::1933967579503727:role/legacy-role", - primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", - archive_prefix="rock-archives/", - ) - ) - svc._sts_clients["legacy"].do_action_with_exception = _fake_assume("L-AK", "L-SK", "L-TOK") - svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") - - legacy = svc.gen_oss_sts_token() - primary = svc.gen_oss_sts_token(account="primary") - assert legacy["ArchivePrefix"] == "rock-archives/" - assert primary["ArchivePrefix"] == "rock-archives/" - - -def test_archive_prefix_none_when_yaml_does_not_set_it(): - svc = _build_service( - _make_rock_config( - legacy_role="acs:ram::1933967579503727:role/legacy-role", - primary_role="acs:ram::1771269394322852:role/chatos-rock-sts-role", - archive_prefix="", - ) - ) - svc._sts_clients["primary"].do_action_with_exception = _fake_assume("P-AK", "P-SK", "P-TOK") - creds = svc.gen_oss_sts_token(account="primary") - assert creds["ArchivePrefix"] is None - - def test_unknown_account_returns_none(): svc = _build_service( _make_rock_config( From a0bb5a92a551a13d3bdf2bb13284610178d35965 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 19:20:14 +0800 Subject: [PATCH 117/226] feature(scheduler): add RayLogCleanupTask and disable worker-to-driver log forwarding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related changes touching ray log paths: 1. RayLogCleanupTask: drop /data/tmp/ray/session__ dirs on each worker that are NOT session_latest. Resolves the session_latest symlink and skips its real target; sessions younger than min_age_hours are also kept as a buffer against a stale symlink. Idempotent — `rm -rf` is resumable across task interruptions because mtime still matches on the next tick. Constructor rejects `min_age_hours < 1` to prevent a misconfig that would delete the live session. Defaults: ray_temp_dir="/data/tmp/ray", min_age_hours=24, interval 24h. NOTE: This is the WORKER side. ray-head's /data/tmp/ray is cleaned by a daily cron baked into the head Dockerfile (rock-internal repo); rocklet is not deployed on the head and the worker scheduler does not reach it. 2. ray.init(log_to_driver=False) on both init paths in ray_service.py (first init + reconnect). Stops worker actor stdout/stderr from being forwarded to the admin process via gRPC, reducing log fanout / network traffic. Worker-side log files in /data/tmp/ray/... are unaffected, which is exactly the set RayLogCleanupTask is designed to manage. Trade-off: admin's own stdout will no longer surface actor print or exception traceback. Operators debugging actor issues must SSH the worker and read worker-*.out under the relevant session dir. Enable per-environment via the existing `enabled:` flag in scheduler.tasks. Co-Authored-By: Claude Opus 4.7 (1M context) --- rock/admin/core/ray_service.py | 2 + rock/admin/scheduler/tasks/__init__.py | 9 +- .../scheduler/tasks/ray_log_cleanup_task.py | 91 +++++++++++++ tests/unit/admin/core/test_ray_service.py | 1 + tests/unit/admin/scheduler/__init__.py | 0 .../scheduler/test_ray_log_cleanup_task.py | 120 ++++++++++++++++++ 6 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 rock/admin/scheduler/tasks/ray_log_cleanup_task.py create mode 100644 tests/unit/admin/scheduler/__init__.py create mode 100644 tests/unit/admin/scheduler/test_ray_log_cleanup_task.py diff --git a/rock/admin/core/ray_service.py b/rock/admin/core/ray_service.py index f6f5ec61ec..9d0fd3379f 100644 --- a/rock/admin/core/ray_service.py +++ b/rock/admin/core/ray_service.py @@ -40,6 +40,7 @@ def init(self): namespace=self._config.namespace, resources=self._config.resources, _temp_dir=self._config.temp_dir, + log_to_driver=False, ) if self._config.ray_reconnect_enabled: self._setup_ray_reconnect_scheduler() @@ -94,6 +95,7 @@ async def _reconnect_ray(self): namespace=self._config.namespace, resources=self._config.resources, _temp_dir=self._config.temp_dir, + log_to_driver=False, ) except Exception as e: last_exc = e diff --git a/rock/admin/scheduler/tasks/__init__.py b/rock/admin/scheduler/tasks/__init__.py index d9a663bd99..700b0c73a6 100644 --- a/rock/admin/scheduler/tasks/__init__.py +++ b/rock/admin/scheduler/tasks/__init__.py @@ -3,5 +3,12 @@ from rock.admin.scheduler.tasks.file_cleanup_task import FileCleanupTask from rock.admin.scheduler.tasks.image_cleanup_task import ImageCleanupTask from rock.admin.scheduler.tasks.image_pull_task import ImagePullTask +from rock.admin.scheduler.tasks.ray_log_cleanup_task import RayLogCleanupTask -__all__ = ["ContainerCleanupTask", "FileCleanupTask", "ImageCleanupTask", "ImagePullTask"] +__all__ = [ + "ContainerCleanupTask", + "FileCleanupTask", + "ImageCleanupTask", + "ImagePullTask", + "RayLogCleanupTask", +] diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py new file mode 100644 index 0000000000..840bab9c7c --- /dev/null +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -0,0 +1,91 @@ +"""Drop stale /data/tmp/ray/session_* dirs on each worker.""" + +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.common.constants import SCHEDULER_LOG_NAME +from rock.logger import init_logger +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + +logger = init_logger(name="ray_log_cleanup", file_name=SCHEDULER_LOG_NAME) + + +class RayLogCleanupTask(BaseTask): + """Drop /data/tmp/ray/session_* dirs that are NOT the live session. + + Ray restarts (cluster up/down, head failover) leave dozens of + session__ dirs behind. The currently active one is + symlinked as `session_latest`; we resolve that link and skip its target. + Sessions younger than `min_age_hours` are also kept as a buffer against + a stale symlink. + + NOTE: This is the WORKER side. The ray-head's /data/tmp/ray is cleaned + by a daily cron baked into the head Dockerfile (rock-internal repo); + rocklet is not deployed on the head and the worker scheduler does not + reach it. + """ + + def __init__( + self, + interval_seconds: int = 86400, + ray_temp_dir: str = "/data/tmp/ray", + min_age_hours: int = 24, + ): + """ + Args: + interval_seconds: Execution interval, default 24 hours. + ray_temp_dir: Ray's --temp-dir, default /data/tmp/ray. + min_age_hours: Only delete session dirs whose mtime is older than + this AND that are not session_latest. Default 24h. + """ + super().__init__( + type="ray_log_cleanup", + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + if min_age_hours < 1: + raise ValueError(f"ray_log_cleanup.min_age_hours must be >= 1, got {min_age_hours}") + self.ray_temp_dir = ray_temp_dir.rstrip("/") + self.min_age_hours = min_age_hours + + @classmethod + def from_config(cls, task_config) -> "RayLogCleanupTask": + return cls( + interval_seconds=task_config.interval_seconds, + ray_temp_dir=task_config.params.get("ray_temp_dir", "/data/tmp/ray"), + min_age_hours=task_config.params.get("min_age_hours", 24), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + ray_dir = self.ray_temp_dir + max_age_min = self.min_age_hours * 60 + # Resolve session_latest -> live basename, then list session_* dirs + # older than threshold and rm -rf those that are not the live one. + command = ( + f'if [ -d "{ray_dir}" ]; then ' + f' LIVE=$(readlink "{ray_dir}/session_latest" 2>/dev/null | xargs -I{{}} basename {{}} 2>/dev/null); ' + f' echo "live_session=${{LIVE:-}}"; ' + f' find "{ray_dir}" -maxdepth 1 -type d -name "session_*" ' + f' ! -name "session_latest" -mmin +{max_age_min} ' + f' | while read -r d; do ' + f' bn=$(basename "$d"); ' + f' if [ "$bn" != "$LIVE" ]; then ' + f' rm -rf "$d" && echo "removed=$bn"; ' + f' fi; ' + f' done; ' + f' echo "ray_log_cleanup_done"; ' + f'else echo "ray_temp_dir_not_found"; fi' + ) + result = await runtime.execute(Command(command=command, shell=True, check=False)) + output = (result.stdout or "").strip() + removed = [line.split("=", 1)[1] for line in output.splitlines() if line.startswith("removed=")] + logger.info( + f"[{self.type}] [{runtime._config.host}] ray_log_cleanup done: " + f"removed={len(removed)} sessions, output_head={output[:300]}" + ) + return { + "status": TaskStatusEnum.SUCCESS, + "exit_code": result.exit_code, + "removed_count": len(removed), + "removed_sessions": removed, + "output_head": output[:1000], + } diff --git a/tests/unit/admin/core/test_ray_service.py b/tests/unit/admin/core/test_ray_service.py index 8b4da1482c..6341ea9f1b 100644 --- a/tests/unit/admin/core/test_ray_service.py +++ b/tests/unit/admin/core/test_ray_service.py @@ -160,6 +160,7 @@ async def test_reconnect_ray_calls_ray_shutdown_and_init_and_reset_counters(ray_ namespace=ray_service._config.namespace, resources=ray_service._config.resources, _temp_dir=ray_service._config.temp_dir, + log_to_driver=False, ) assert service._ray_request_count == 0 diff --git a/tests/unit/admin/scheduler/__init__.py b/tests/unit/admin/scheduler/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py new file mode 100644 index 0000000000..8f2991fd8a --- /dev/null +++ b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py @@ -0,0 +1,120 @@ +"""Tests for RayLogCleanupTask.""" + +from unittest.mock import AsyncMock + +import pytest + +from rock.admin.scheduler.task_base import TaskStatusEnum +from rock.admin.scheduler.tasks.ray_log_cleanup_task import RayLogCleanupTask + + +class _FakeTaskConfig: + def __init__(self, params=None, interval_seconds=86400): + self.params = params or {} + self.interval_seconds = interval_seconds + + +class _FakeExecResult: + def __init__(self, exit_code=0, stdout="ray_log_cleanup_done"): + self.exit_code = exit_code + self.stdout = stdout + + +def _runtime(stdout="ray_log_cleanup_done", exit_code=0): + rt = AsyncMock() + rt._config = type("C", (), {"host": "10.0.0.1"})() + rt.execute = AsyncMock(return_value=_FakeExecResult(exit_code=exit_code, stdout=stdout)) + return rt + + +class TestInit: + def test_default(self): + task = RayLogCleanupTask() + assert task.type == "ray_log_cleanup" + assert task.ray_temp_dir == "/data/tmp/ray" + assert task.min_age_hours == 24 + + def test_strips_trailing_slash(self): + task = RayLogCleanupTask(ray_temp_dir="/data/ray/") + assert task.ray_temp_dir == "/data/ray" + + def test_rejects_min_age_below_one(self): + with pytest.raises(ValueError, match="min_age_hours must be >= 1"): + RayLogCleanupTask(min_age_hours=0) + + +class TestFromConfig: + def test_from_config_defaults(self): + task = RayLogCleanupTask.from_config(_FakeTaskConfig()) + assert task.ray_temp_dir == "/data/tmp/ray" + assert task.min_age_hours == 24 + + def test_from_config_custom(self): + cfg = _FakeTaskConfig( + params={"ray_temp_dir": "/data/ray", "min_age_hours": 48}, + interval_seconds=3600, + ) + task = RayLogCleanupTask.from_config(cfg) + assert task.ray_temp_dir == "/data/ray" + assert task.min_age_hours == 48 + assert task.interval_seconds == 3600 + + +class TestRunAction: + @pytest.mark.asyncio + async def test_command_skips_session_latest(self): + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # `! -name "session_latest"` skips the symlink itself; readlink resolves + # the target so we also skip whichever real session it points at. + assert '! -name "session_latest"' in cmd + assert "readlink" in cmd + assert 'name "session_*"' in cmd + + @pytest.mark.asyncio + async def test_command_uses_min_age_in_minutes(self): + task = RayLogCleanupTask(min_age_hours=48) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # 48h * 60 = 2880 + assert "-mmin +2880" in cmd + + @pytest.mark.asyncio + async def test_command_respects_custom_temp_dir(self): + task = RayLogCleanupTask(ray_temp_dir="/data/ray") + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + assert '"/data/ray"' in cmd + + @pytest.mark.asyncio + async def test_extracts_removed_count_from_output(self): + stdout = ( + "live_session=session_2026_03_01_xyz_111\n" + "removed=session_2026_02_15_aaa_222\n" + "removed=session_2026_02_20_bbb_333\n" + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + assert result["removed_count"] == 2 + assert "session_2026_02_15_aaa_222" in result["removed_sessions"] + assert "session_2026_02_20_bbb_333" in result["removed_sessions"] + + @pytest.mark.asyncio + async def test_handles_missing_ray_dir(self): + task = RayLogCleanupTask() + runtime = _runtime(stdout="ray_temp_dir_not_found") + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + assert result["removed_count"] == 0 From 0ef3d98349699d813332b19acff8e0f1bebec034 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 14:55:45 +0800 Subject: [PATCH 118/226] use triple-quoted multiline bash for ray_log_cleanup command --- .../scheduler/tasks/ray_log_cleanup_task.py | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index 840bab9c7c..264d65ee3a 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -1,5 +1,7 @@ """Drop stale /data/tmp/ray/session_* dirs on each worker.""" +import textwrap + from rock.admin.proto.request import SandboxCommand as Command from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum from rock.common.constants import SCHEDULER_LOG_NAME @@ -60,20 +62,25 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: max_age_min = self.min_age_hours * 60 # Resolve session_latest -> live basename, then list session_* dirs # older than threshold and rm -rf those that are not the live one. - command = ( - f'if [ -d "{ray_dir}" ]; then ' - f' LIVE=$(readlink "{ray_dir}/session_latest" 2>/dev/null | xargs -I{{}} basename {{}} 2>/dev/null); ' - f' echo "live_session=${{LIVE:-}}"; ' - f' find "{ray_dir}" -maxdepth 1 -type d -name "session_*" ' - f' ! -name "session_latest" -mmin +{max_age_min} ' - f' | while read -r d; do ' - f' bn=$(basename "$d"); ' - f' if [ "$bn" != "$LIVE" ]; then ' - f' rm -rf "$d" && echo "removed=$bn"; ' - f' fi; ' - f' done; ' - f' echo "ray_log_cleanup_done"; ' - f'else echo "ray_temp_dir_not_found"; fi' + # textwrap.dedent strips common leading whitespace so source can be + # indented for readability without polluting the emitted shell. + command = textwrap.dedent( + f"""\ + if [ -d "{ray_dir}" ]; then + LIVE=$(readlink "{ray_dir}/session_latest" 2>/dev/null | xargs -I{{}} basename {{}} 2>/dev/null) + echo "live_session=${{LIVE:-}}" + find "{ray_dir}" -maxdepth 1 -type d -name "session_*" \\ + ! -name "session_latest" -mmin +{max_age_min} \\ + | while read -r d; do + bn=$(basename "$d") + if [ "$bn" != "$LIVE" ]; then + rm -rf "$d" && echo "removed=$bn" + fi + done + echo "ray_log_cleanup_done" + else + echo "ray_temp_dir_not_found" + fi""" ) result = await runtime.execute(Command(command=command, shell=True, check=False)) output = (result.stdout or "").strip() From 60c3e106eea4ca0159ef6c0fdf889824f985d5c8 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Mon, 18 May 2026 19:29:19 +0800 Subject: [PATCH 119/226] feat(sandbox): add CPU overcommit with grayscale rollout, lifecycle summary, and absolute-cores CPU gauge #978 (#979) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * refactor(rocklet): simplify cgroup CPU cache to single scalar and cache cpu_quota Remove per-thread dict for CPU usage/time (no concurrency concern) and cache cpu_quota since it doesn't change at runtime. Co-Authored-By: Claude Opus 4.6 * support cpu over * add log to check cpu * opt ut * add log * update version --------- Co-authored-by: qianyang Co-authored-by: Claude Opus 4.6 --- pyproject.toml | 2 +- rock/admin/entrypoints/sandbox_api.py | 49 ++++++++++++++++++++------- rock/common/constants.py | 13 ++++++- rock/deployments/config.py | 2 +- rock/sandbox/base_actor.py | 41 ++++++++++++++++++++++ rock/sandbox/operator/abstract.py | 3 +- rock/sandbox/operator/k8s/operator.py | 6 +++- rock/sandbox/operator/ray.py | 6 ++-- rock/sandbox/sandbox_actor.py | 7 ++-- rock/sandbox/sandbox_manager.py | 9 ++--- tests/unit/test_base_actor.py | 33 +++++++++++++++++- uv.lock | 2 +- 12 files changed, 145 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7fbe447e70..fbd220c545 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.7.0" +version = "1.8.0" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index e8a264a659..3c1760736f 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -1,3 +1,4 @@ +import math from typing import Annotated, Any from fastapi import APIRouter, Body, Depends, File, Form, UploadFile @@ -26,7 +27,8 @@ ) from rock.admin.proto.response import SandboxStartResponse from rock.common.constants import ( - CPU_PREEMPT_SWITCH, + CPU_OVERCOMMIT_ALLOWED_KEYS_KEY, + CPU_OVERCOMMIT_HEADROOM_KEY, GET_STATUS_SWITCH, KATA_DIND_DISK_SIZE_KEY, KATA_RUNTIME_SWITCH, @@ -94,17 +96,41 @@ async def _apply_disk_limits(config: DockerDeploymentConfig) -> None: config.disk_limit_log = disk_limit_log -async def _apply_cpu_preempt_switch(config: DockerDeploymentConfig) -> None: - """Check nacos switch and enable CPU preemption on the config if the switch is on. +async def _apply_cpu_overcommit_default(config: DockerDeploymentConfig, rock_authorization: str | None) -> None: + """Derive limit_cpus from cpus + Nacos headroom when SDK did not set it. - When the switch is off, limit_cpus will be - cleared so that --cpus is not passed to docker run. + Formula: limit_cpus = min(2 * cpus, cpus + headroom) + - SDK-supplied limit_cpus always wins (function is a no-op in that case). + - Grayscale gate driven by Nacos list `cpu_overcommit_allowed_keys`: + * key absent from Nacos -> gate is open for every caller (full rollout). + * key present as a list -> only `rock_authorization` values in the list pass. + * key present but not a list (misconfigured) -> gate closed. + - headroom is read from Nacos key `cpu_overcommit_headroom` (default 0). + - headroom <= 0 keeps limit_cpus = None (docker run gets no --cpus flag). """ - if ( - sandbox_manager.rock_config.nacos_provider is not None - and not await sandbox_manager.rock_config.nacos_provider.get_switch_status(CPU_PREEMPT_SWITCH, False) - ): - config.limit_cpus = None + if config.limit_cpus is not None: + return + + nacos = sandbox_manager.rock_config.nacos_provider + if nacos is None: + return + + nacos_config = await nacos.get_config() or {} + allowed_keys = nacos_config.get(CPU_OVERCOMMIT_ALLOWED_KEYS_KEY) + if allowed_keys is not None and (not isinstance(allowed_keys, list) or rock_authorization not in allowed_keys): + return + + raw = nacos_config.get(CPU_OVERCOMMIT_HEADROOM_KEY) + try: + headroom = float(raw) if raw is not None else 0.0 + except (TypeError, ValueError): + headroom = 0.0 + + # Reject NaN / inf so a fat-fingered Nacos edit can't break sandbox startup + if not math.isfinite(headroom) or headroom <= 0: + return + + config.limit_cpus = min(2 * config.cpus, config.cpus + headroom) @sandbox_router.post("/start") @@ -113,7 +139,6 @@ async def start(request: SandboxStartRequest) -> RockResponse[SandboxStartRespon config = DockerDeploymentConfig.from_request(request) await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) - await _apply_cpu_preempt_switch(config) await _apply_disk_limits(config) sandbox_start_response = await sandbox_manager.start(config) return RockResponse(result=sandbox_start_response) @@ -128,7 +153,7 @@ async def start_async( config = DockerDeploymentConfig.from_request(request) await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) - await _apply_cpu_preempt_switch(config) + await _apply_cpu_overcommit_default(config, headers.user_info.get("rock_authorization")) await _apply_disk_limits(config) sandbox_start_response = await sandbox_manager.start_async( config, diff --git a/rock/common/constants.py b/rock/common/constants.py index 68a5cc27bd..bf19e498e4 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -3,7 +3,8 @@ GET_STATUS_SWITCH = "get_status_v2_enabled" KATA_RUNTIME_SWITCH = "use_kata_enabled" SUPPORT_KATA_SWITCH = "support_kata_enabled" -CPU_PREEMPT_SWITCH = "cpu_preempt_enabled" +CPU_OVERCOMMIT_HEADROOM_KEY = "cpu_overcommit_headroom" +CPU_OVERCOMMIT_ALLOWED_KEYS_KEY = "cpu_overcommit_allowed_keys" KATA_DIND_DISK_SIZE_KEY = "kata_dind_disk_size" SANDBOX_DISK_LIMIT_ROOTFS_KEY = "sandbox_disk_limit_rootfs" SANDBOX_DISK_LIMIT_LOG_KEY = "sandbox_disk_limit_log" @@ -20,3 +21,13 @@ class DeploymentHookStep(str, Enum): PULLING_IMAGE = "Pulling docker image" STARTING_RUNTIME = "Starting runtime" + + +class StopReason(str, Enum): + """Why a sandbox was stopped. Propagated through the SandboxManager → Operator → Actor + stop chain so the actor-side lifecycle summary can distinguish user-initiated stops + from auto-cleanup of expired sandboxes. + """ + + MANUAL = "manual" + EXPIRED = "expired" diff --git a/rock/deployments/config.py b/rock/deployments/config.py index b9b9da17fd..2ae440aae4 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -91,7 +91,7 @@ class DockerDeploymentConfig(DeploymentConfig): """Number of CPU cores to allocate for the container. Used as --cpu-shares (cpus * 1024).""" limit_cpus: float | None = None - """Hard limit on the number of CPU cores the container can use. Used as --cpus when CPU preemption is enabled via nacos switch.""" + """Hard limit on the number of CPU cores the container can use. Used as --cpus. When the SDK leaves this None, the admin gateway may derive it from cpus + Nacos `cpu_overcommit_headroom`.""" disk_limit_rootfs: str | None = None """Maximum rootfs disk size for the container (e.g., '20g', '50g'). Maps to --storage-opt size=. Only supported on overlay2 storage driver with xfs backing filesystem. None means no limit.""" diff --git a/rock/sandbox/base_actor.py b/rock/sandbox/base_actor.py index 0111083736..4b8a94bdd7 100644 --- a/rock/sandbox/base_actor.py +++ b/rock/sandbox/base_actor.py @@ -11,6 +11,7 @@ from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader from rock import env_vars +from rock.common.constants import StopReason from rock.deployments.abstract import AbstractDeployment from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig from rock.deployments.docker import DockerDeployment @@ -39,6 +40,7 @@ class BaseActor: _user_defined_tags: dict = {} _created_time: float = None _host_name: str = None + _max_cpus_used: float = 0.0 def __init__( self, @@ -51,6 +53,7 @@ def __init__( if isinstance(config, DockerDeploymentConfig) and config.auto_clear_time: self._auto_clear_time_in_minutes = config.auto_clear_time self._created_time = time.monotonic() + self._max_cpus_used = 0.0 self._stop_time = datetime.datetime.now() + datetime.timedelta(minutes=self._auto_clear_time_in_minutes) # Initialize the user and environment info - can be overridden by subclasses self._role = "test" @@ -110,6 +113,12 @@ def _init_monitor(self): self._gauges["rt"] = self.meter.create_gauge( name="xrl_gateway.system.lifespan_rt", description="Life Span Rt", unit="1" ) + self._gauges["cpus_used"] = self.meter.create_gauge( + name="xrl_gateway.system.cpus_used", + description="Sandbox actual CPU cores used (cpu_percent/100 * cpus_limit). " + "cpus_allocated and cpus_limit are reported as tags on this and every other gauge.", + unit="1", + ) async def _setup_monitor(self): if not env_vars.ROCK_MONITOR_ENABLE: @@ -127,6 +136,19 @@ async def _setup_monitor(self): ) self._metrics_report_scheduler.start() + def log_lifecycle_summary(self, reason: StopReason = StopReason.MANUAL) -> None: + """Emit a single line summarising CPU usage and lifespan for this sandbox.""" + duration_seconds = time.monotonic() - self._created_time if self._created_time is not None else 0.0 + logger.info( + f"[{self._config.container_name}] lifecycle summary: " + f"reason={reason.value}, " + f"image={self._config.image}, " + f"cpus={self._config.cpus}, " + f"limit_cpus={self._config.limit_cpus}, " + f"max_cpus_used={self._max_cpus_used:.4f}, " + f"duration={duration_seconds:.2f}s" + ) + def stop_monitoring(self): if env_vars.ROCK_MONITOR_ENABLE and self._metrics_report_scheduler and self._metrics_report_scheduler.running: logger.info("Stopping APScheduler...") @@ -159,6 +181,11 @@ async def _collect_sandbox_metrics(self, sandbox_id: str): return logger.debug(f"sandbox [{sandbox_id}] metrics = {metrics}") + cpus = float(self._config.cpus) + limit_cpus_raw = self._config.limit_cpus + # When limit_cpus is None, treat it as equal to cpus (no overcommit). + effective_limit = float(limit_cpus_raw) if limit_cpus_raw is not None else cpus + attributes = { "sandbox_id": sandbox_id, "env": self._env, @@ -169,6 +196,11 @@ async def _collect_sandbox_metrics(self, sandbox_id: str): "experiment_id": self._experiment_id, "namespace": self._namespace, "host_name": self._host_name, + # cpus_allocated / cpus_limit are static per-sandbox values; reported + # as tags on every gauge so dashboards can group/filter without + # needing dedicated time-series for them. + "cpus_allocated": f"{cpus:g}", + "cpus_limit": f"{effective_limit:g}", } if self._user_defined_tags is not None: attributes.update(self._user_defined_tags) @@ -179,6 +211,15 @@ async def _collect_sandbox_metrics(self, sandbox_id: str): self._gauges["disk"].set(metrics["disk"], attributes=attributes) self._gauges["net"].set(metrics["net"], attributes=attributes) + # cpus_used inherits the semantic of metrics["cpu"]: it is reported + # as a percentage of the sandbox's allocated CPU (see the legend of + # xrl_gateway.system.cpu in admin/metrics/monitor.py). Multiplying by + # effective_limit converts it back to absolute cores. + cpus_used = (metrics["cpu"] / 100.0) * effective_limit if effective_limit > 0 else 0.0 + if cpus_used > self._max_cpus_used: + self._max_cpus_used = cpus_used + self._gauges["cpus_used"].set(cpus_used, attributes=attributes) + logger.debug(f"Successfully reported metrics for sandbox: {sandbox_id}") else: logger.warning(f"No metrics returned for sandbox: {sandbox_id}") diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index dc326d7504..72efea0551 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -2,6 +2,7 @@ from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.redis_key import alive_sandbox_key +from rock.common.constants import StopReason from rock.config import RuntimeConfig from rock.deployments.config import DeploymentConfig from rock.utils.providers.nacos_provider import NacosConfigProvider @@ -22,7 +23,7 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: ... @abstractmethod - async def stop(self, sandbox_id: str) -> bool: + async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> bool: ... def set_redis_provider(self, redis_provider: RedisProvider): diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 07f3d49aaf..503c236774 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -1,6 +1,7 @@ """K8s Operator implementation for managing sandboxes via Kubernetes.""" from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.common.constants import StopReason from rock.config import K8sConfig from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger @@ -113,13 +114,16 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: raise Exception(f"Sandbox {sandbox_id} not found in Redis") return sandbox_info - async def stop(self, sandbox_id: str) -> bool: + async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> bool: """Stop and delete a sandbox. Args: sandbox_id: Sandbox identifier + reason: Why the stop was triggered. Logged here for traceability; K8s + path has no actor-side lifecycle summary to attach it to. Returns: True if successful, False otherwise """ + logger.info(f"[{sandbox_id}] k8s stop (reason={reason.value})") return await self._provider.stop(sandbox_id) diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index bb255e4ebf..e0efbae4ef 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -6,7 +6,7 @@ from rock.actions.sandbox.response import IsAliveResponse, State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.ray_service import RayService -from rock.common.constants import GET_STATUS_SWITCH +from rock.common.constants import GET_STATUS_SWITCH, StopReason from rock.config import RuntimeConfig from rock.deployments.config import DockerDeploymentConfig from rock.deployments.constants import Port @@ -104,10 +104,10 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: else: return sandbox_info - async def stop(self, sandbox_id: str) -> bool: + async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> bool: async with self._ray_service.get_ray_rwlock().read_lock(): 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()) + await self._ray_service.async_ray_get(actor.stop.remote(reason)) logger.info(f"run time stop over {sandbox_id}") ray.kill(actor) return True diff --git a/rock/sandbox/sandbox_actor.py b/rock/sandbox/sandbox_actor.py index 92a940e346..96c1799d7c 100644 --- a/rock/sandbox/sandbox_actor.py +++ b/rock/sandbox/sandbox_actor.py @@ -25,6 +25,7 @@ from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest +from rock.common.constants import StopReason from rock.deployments.abstract import AbstractDeployment from rock.deployments.config import DeploymentConfig from rock.deployments.constants import Status @@ -133,8 +134,8 @@ async def start(self): self._clean_container_background() await self._setup_monitor() - async def stop(self): - logger.info(f"[{self._config.container_name}] start to stop") + async def stop(self, reason: StopReason = StopReason.MANUAL): + logger.info(f"[{self._config.container_name}] start to stop (reason={reason.value})") try: await self._deployment.stop() logger.info(f"[{self._config.container_name}] deployment stopped") @@ -142,6 +143,8 @@ async def stop(self): logger.info(f"[{self._config.container_name}] actor stopped") except Exception as e: logger.error(f"[{self._config.container_name}] Error occurred while stopping container: {e}", exc_info=True) + finally: + self.log_lifecycle_summary(reason) async def commit(self, image_tag: str, username: str, password: str) -> CommandResponse: logger.info(f"start to commit {self._config.container_name} to {image_tag}") diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 3d85fd0671..dd033df490 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -25,6 +25,7 @@ from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest from rock.admin.proto.response import SandboxStartResponse, SandboxStatusResponse +from rock.common.constants import StopReason from rock.config import RockConfig, RuntimeConfig from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig from rock.logger import init_logger @@ -170,8 +171,8 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: ) @monitor_sandbox_operation() - async def stop(self, sandbox_id): - logger.info(f"stop sandbox {sandbox_id}") + async def stop(self, sandbox_id, reason: StopReason = StopReason.MANUAL): + logger.info(f"stop sandbox {sandbox_id} (reason={reason.value})") sandbox_info: SandboxInfo | None = await self._meta_store.get(sandbox_id) if sandbox_info is None: sandbox_info = {} @@ -180,7 +181,7 @@ async def stop(self, sandbox_id): sandbox_info["stop_time"] = get_iso8601_timestamp() log_billing_info(sandbox_info=sandbox_info) try: - await self._operator.stop(sandbox_id) + await self._operator.stop(sandbox_id, reason) except ValueError as e: logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) await self._meta_store.archive(sandbox_id, sandbox_info) @@ -330,7 +331,7 @@ async def _check_job_background(self): is_expired = await self._is_expired(sandbox_id) if is_expired: logger.info(f"sandbox_id:[{sandbox_id}] is expired, start to stop") - asyncio.create_task(self.stop(sandbox_id)) + asyncio.create_task(self.stop(sandbox_id, reason=StopReason.EXPIRED)) except asyncio.CancelledError as e: logger.error("check_job_background CancelledError", exc_info=e) continue diff --git a/tests/unit/test_base_actor.py b/tests/unit/test_base_actor.py index 61586c3d24..76c98077a9 100644 --- a/tests/unit/test_base_actor.py +++ b/tests/unit/test_base_actor.py @@ -150,6 +150,8 @@ def _make_actor() -> ConcreteBaseActor: config = MagicMock() config.container_name = "test-container" config.auto_clear_time = None # skip DockerDeploymentConfig branch + config.cpus = 2.0 + config.limit_cpus = None deployment = MagicMock() deployment.__class__ = object # make isinstance(deployment, DockerDeployment) return False @@ -157,7 +159,7 @@ def _make_actor() -> ConcreteBaseActor: actor = ConcreteBaseActor(config, deployment) actor.host = "127.0.0.1" # Pre-populate all gauges with mocks so tests can override selectively - for key in ("cpu", "mem", "disk", "net", "rt"): + for key in ("cpu", "mem", "disk", "net", "rt", "cpus_used"): actor._gauges[key] = MagicMock() return actor @@ -304,3 +306,32 @@ async def test_metrics_attributes_host_name_matches_actor_field(): attributes = mock_cpu_gauge.set.call_args[1]["attributes"] assert attributes["host_name"] == custom_hostname + + +@pytest.mark.asyncio +async def test_cpu_metric_no_overcommit(): + """When limit_cpus is None: cpus_used scales by cpus, and cpus_allocated/cpus_limit + are reported as tags equal to cpus.""" + actor = _make_actor() # cpus=2.0, limit_cpus=None; ConcreteBaseActor returns cpu=10.0 + await actor._collect_sandbox_metrics("test-container") + + # cpu_pct=10, effective_limit=2 -> cpus_used = 0.1 * 2 = 0.2 + assert actor._gauges["cpus_used"].set.call_args[0][0] == pytest.approx(0.2) + attrs = actor._gauges["cpus_used"].set.call_args[1]["attributes"] + assert attrs["cpus_allocated"] == "2" + assert attrs["cpus_limit"] == "2" + + +@pytest.mark.asyncio +async def test_cpu_metric_overcommit(): + """When limit_cpus > cpus: cpus_used scales by the limit; cpus_limit tag reflects the limit.""" + actor = _make_actor() + actor._config.limit_cpus = 6.0 # overcommit: 2 allocated, 6 hard cap + + await actor._collect_sandbox_metrics("test-container") + + # cpu_pct=10, effective_limit=6 -> cpus_used = 0.1 * 6 = 0.6 + assert actor._gauges["cpus_used"].set.call_args[0][0] == pytest.approx(0.6) + attrs = actor._gauges["cpus_used"].set.call_args[1]["attributes"] + assert attrs["cpus_allocated"] == "2" + assert attrs["cpus_limit"] == "6" diff --git a/uv.lock b/uv.lock index bfd81a2782..56fba7522c 100644 --- a/uv.lock +++ b/uv.lock @@ -4174,7 +4174,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.7.0" +version = "1.7.1" source = { editable = "." } dependencies = [ { name = "anyio" }, From a1797a591ab104007fa0d5cf7f98e5e3a885fcf0 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 19 May 2026 07:51:29 +0800 Subject: [PATCH 120/226] feat(sandbox): support include_all_states parameter in get_status API (#951) * feat(sandbox): support include_all_states parameter in get_status API Add include_all_states=False parameter to SandboxManager.get_status/get_status_v2, sandbox_api GET /get_status endpoint, and SDK Sandbox.get_status. When enabled and the operator returns None (sandbox not in active scheduling), fall back to meta_store.batch_get so stopped/historical sandbox info can still be retrieved. Also fix RayOperator.get_status (use_rocklet path) to return None instead of raising when the sandbox is absent from Redis. Signed-off-by: Jiachen Zhang * test(sandbox): add unit tests for get_status include_all_states feature Signed-off-by: Jiachen Zhang * refactor(sandbox): flatten lifecycle_info dict into top-level fields on SandboxStatusResponse --------- Signed-off-by: Jiachen Zhang --- rock/actions/sandbox/response.py | 4 + rock/admin/entrypoints/sandbox_api.py | 8 +- rock/admin/proto/response.py | 3 + rock/sandbox/operator/abstract.py | 2 +- rock/sandbox/operator/k8s/operator.py | 4 +- rock/sandbox/operator/ray.py | 4 +- rock/sandbox/sandbox_manager.py | 35 +++-- rock/sandbox/sandbox_meta_store.py | 6 +- rock/sdk/sandbox/client.py | 4 +- .../sandbox/operator/test_k8s_operator.py | 8 +- .../test_get_status_include_all_states.py | 122 ++++++++++++++++++ 11 files changed, 174 insertions(+), 26 deletions(-) create mode 100644 tests/unit/sandbox/test_get_status_include_all_states.py diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index 7f7baaab89..5fdb4f10e5 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -50,6 +50,10 @@ class SandboxStatusResponse(BaseModel): memory: str | None = None disk_limit_rootfs: str | None = None disk_limit_log: str | None = None + state: State | None = None + start_time: str | None = None + stop_time: str | None = None + create_time: str | None = None class CommandResponse(BaseModel): diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index 3c1760736f..cc2940e4fc 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -183,14 +183,16 @@ async def get_sandbox_statistics(sandbox_id: str): @sandbox_router.get("/get_status") @handle_exceptions(error_message="get sandbox status failed") -async def get_status(sandbox_id: str): +async def get_status(sandbox_id: str, include_all_states: bool = False): # TODO: do judgement inside operator if ( sandbox_manager.rock_config.nacos_provider is not None and await sandbox_manager.rock_config.nacos_provider.get_switch_status(GET_STATUS_SWITCH) ): - return RockResponse(result=await sandbox_manager.get_status_v2(sandbox_id)) - return RockResponse(result=await sandbox_manager.get_status(sandbox_id)) + return RockResponse( + result=await sandbox_manager.get_status_v2(sandbox_id, include_all_states=include_all_states) + ) + return RockResponse(result=await sandbox_manager.get_status(sandbox_id, include_all_states=include_all_states)) @sandbox_router.post("/execute") diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index d1edabc7a5..c1e1cd9249 100644 --- a/rock/admin/proto/response.py +++ b/rock/admin/proto/response.py @@ -34,6 +34,9 @@ class SandboxStatusResponse(BaseModel): memory: str | None = None disk_limit_rootfs: str | None = None disk_limit_log: str | None = None + start_time: str | None = None + stop_time: str | None = None + create_time: str | None = None @classmethod def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusResponse": diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index 72efea0551..997c95546b 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -19,7 +19,7 @@ async def submit(self, config: DeploymentConfig, user_info: dict = {}) -> Sandbo ... @abstractmethod - async def get_status(self, sandbox_id: str) -> SandboxInfo: + async def get_status(self, sandbox_id: str) -> SandboxInfo | None: ... @abstractmethod diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 503c236774..ca9e1ac66e 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -90,7 +90,7 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> """ return await self._provider.submit(config, user_info) - async def get_status(self, sandbox_id: str) -> SandboxInfo: + async def get_status(self, sandbox_id: str) -> SandboxInfo | None: """Get sandbox status with user info from Redis. This method first gets status from provider (IP, port_mapping, is_alive), @@ -111,7 +111,7 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo: if redis_info: return _merge_sandbox_info(redis_info, sandbox_info) else: - raise Exception(f"Sandbox {sandbox_id} not found in Redis") + return None return sandbox_info async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> bool: diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index e0efbae4ef..fa8b03150b 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -74,9 +74,11 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> logger.info(f"sandbox {sandbox_id} is submitted") return sandbox_info - async def get_status(self, sandbox_id: str) -> SandboxInfo: + async def get_status(self, sandbox_id: str) -> SandboxInfo | None: if self.use_rocklet(): sandbox_info: SandboxInfo = await build_sandbox_from_redis(self._redis_provider, sandbox_id) + if sandbox_info is None: + return None host_ip = sandbox_info.get("host_ip") remote_status = await self.get_remote_status(sandbox_id, host_ip) is_alive = await self._check_alive_status(sandbox_id, host_ip, remote_status) diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index dd033df490..7a6c1316d9 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -215,16 +215,24 @@ async def commit(self, sandbox_id, image_tag: str, username: str, password: str) return result @monitor_sandbox_operation() - async def get_status(self, sandbox_id) -> SandboxStatusResponse: - sandbox_info: SandboxInfo = await self._operator.get_status(sandbox_id=sandbox_id) - is_alive = sandbox_info.get("state") == State.RUNNING - if sandbox_info.get("state") == State.STOPPED: - raise BadRequestRockError(f"Sandbox {sandbox_id} is already stopped") - self._update_sandbox_alive_info(sandbox_info, is_alive) - current = await self._meta_store.get(sandbox_id) - if current is None or current.get("state") != sandbox_info.get("state"): - await self._meta_store.update(sandbox_id, sandbox_info) - await self._refresh_timeout(sandbox_id) + async def get_status(self, sandbox_id, include_all_states: bool = False) -> SandboxStatusResponse: + is_alive = False + + sandbox_info: SandboxInfo | None = await self._operator.get_status(sandbox_id=sandbox_id) + if sandbox_info is not None: + is_alive = sandbox_info.get("state") == State.RUNNING + self._update_sandbox_alive_info(sandbox_info, is_alive) + if sandbox_info.get("state") in (State.PENDING, State.RUNNING): + current = await self._meta_store.get(sandbox_id) + if current is None or current.get("state") != sandbox_info.get("state"): + await self._meta_store.update(sandbox_id, sandbox_info) + await self._refresh_timeout(sandbox_id) + elif include_all_states: + sandbox_info = await self._meta_store.get(sandbox_id, check_db=True) + + if sandbox_info is None: + raise BadRequestRockError(f"Sandbox {sandbox_id} not found") + return SandboxStatusResponse( sandbox_id=sandbox_id, status=sandbox_info.get("phases"), @@ -243,6 +251,9 @@ async def get_status(self, sandbox_id) -> SandboxStatusResponse: memory=sandbox_info.get("memory"), disk_limit_rootfs=sandbox_info.get("disk_limit_rootfs"), disk_limit_log=sandbox_info.get("disk_limit_log"), + start_time=sandbox_info.get("start_time"), + stop_time=sandbox_info.get("stop_time"), + create_time=sandbox_info.get("create_time"), ) async def build_sandbox_info_from_redis(self, sandbox_id: str, deployment_info: SandboxInfo) -> SandboxInfo | None: @@ -266,12 +277,12 @@ def _update_sandbox_alive_info(self, sandbox_info: SandboxInfo, is_alive: bool) if sandbox_info.get("start_time") is None: sandbox_info["start_time"] = get_iso8601_timestamp() - async def get_status_v2(self, sandbox_id) -> SandboxStatusResponse: + async def get_status_v2(self, sandbox_id, include_all_states: bool = False) -> SandboxStatusResponse: """ Deprecated: Use get_status(sandbox_id, use_rocklet=True) instead. This method is kept for backward compatibility. """ - return await self.get_status(sandbox_id) + return await self.get_status(sandbox_id, include_all_states=include_all_states) async def create_session(self, request: CreateSessionRequest) -> CreateBashSessionResponse: return await self._proxy_service.create_session(request) diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index b1dec0246b..ed6b97f45d 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -115,11 +115,15 @@ async def archive(self, sandbox_id: str, final_info: SandboxInfo) -> None: await self._redis.json_delete(timeout_sandbox_key(sandbox_id)) @monitor_metastore_operation - async def get(self, sandbox_id: str) -> SandboxInfo | None: + async def get(self, sandbox_id: str, check_db: bool = False) -> SandboxInfo | None: """Read sandbox info from the Redis alive key.""" result = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") if result and len(result) > 0: return result[0] + if check_db: + result = await self._db.get(sandbox_id) + if result: + return result return None async def exists(self, sandbox_id: str) -> bool: diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 65527d91b8..3a57eacda5 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -222,8 +222,8 @@ async def is_alive(self) -> IsAliveResponse: logging.warning(f"Failed to get is alive, {str(e)}") raise Exception(f"Failed to get is alive: {str(e)}") - async def get_status(self) -> SandboxStatusResponse: - url = f"{self._url}/get_status?sandbox_id={self.sandbox_id}" + async def get_status(self, include_all_states: bool = False) -> SandboxStatusResponse: + url = f"{self._url}/get_status?sandbox_id={self.sandbox_id}&include_all_states={include_all_states}" headers = self._build_headers() response = await HttpUtils.get(url, headers) logging.debug(f"Get status response: {response}") diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 0bb2be04fd..67da376313 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -179,7 +179,7 @@ async def test_get_sandbox_info_from_redis_no_provider(self, k8s_operator): @pytest.mark.asyncio async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, redis_provider): - """Test get_status raises error when sandbox not found in Redis.""" + """Test get_status returns None when sandbox not found in Redis.""" k8s_operator.set_redis_provider(redis_provider) # Mock provider returns sandbox info @@ -192,9 +192,9 @@ async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, } mock_provider.get_status = AsyncMock(return_value=SandboxInfo(**mock_sandbox_info)) - # Sandbox not in Redis (no data stored) - with pytest.raises(Exception, match="Sandbox test-sandbox not found in Redis"): - await k8s_operator.get_status("test-sandbox") + # Sandbox not in Redis (no data stored) → returns None + result = await k8s_operator.get_status("test-sandbox") + assert result is None class TestMergeSandboxInfo: diff --git a/tests/unit/sandbox/test_get_status_include_all_states.py b/tests/unit/sandbox/test_get_status_include_all_states.py new file mode 100644 index 0000000000..abd67896f0 --- /dev/null +++ b/tests/unit/sandbox/test_get_status_include_all_states.py @@ -0,0 +1,122 @@ +""" +Unit tests for SandboxManager.get_status changes in feat(sandbox): support include_all_states. + +Key behaviour changes covered: + - operator.get_status() may now return None + - include_all_states=False + operator None → raise BadRequestRockError("not found") + - include_all_states=True + operator None → fall back to meta_store.get(check_db=True) + - include_all_states=True + operator data → skip fallback, normal path + - meta_store.update only when state is PENDING or RUNNING + - start_time/stop_time/create_time populated in every response +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from rock.actions.sandbox.response import State +from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.admin.proto.response import SandboxStatusResponse +from rock.sdk.common.exceptions import BadRequestRockError + + +def _make_sandbox_info(sandbox_id: str = "sandbox-1", state: State = State.RUNNING) -> SandboxInfo: + return SandboxInfo( + sandbox_id=sandbox_id, + state=state, + host_ip="10.0.0.1" if state != State.PENDING else None, + host_name="node-1" if state != State.PENDING else None, + image="python:3.11", + phases={}, + port_mapping={}, + ) + + +@pytest.fixture +def mock_operator(): + return AsyncMock() + + +@pytest.fixture +def mock_meta_store(): + store = AsyncMock() + store.get = AsyncMock(return_value=None) + store.update = AsyncMock() + return store + + +@pytest.fixture +def sandbox_manager(mock_operator, mock_meta_store, rock_config): + from rock.sandbox.sandbox_manager import SandboxManager + + with patch("rock.sandbox.sandbox_manager.SandboxProxyService"): + manager = SandboxManager.__new__(SandboxManager) + manager.rock_config = rock_config + manager._operator = mock_operator + manager._meta_store = mock_meta_store + manager._refresh_timeout = AsyncMock() + return manager + + +class TestGetStatusIncludeAllStates: + @pytest.mark.asyncio + async def test_running_sandbox_returns_alive_response(self, sandbox_manager, mock_operator): + """Operator returns RUNNING → is_alive=True, state field populated.""" + mock_operator.get_status = AsyncMock(return_value=_make_sandbox_info(state=State.RUNNING)) + + result = await sandbox_manager.get_status("sandbox-1") + + assert isinstance(result, SandboxStatusResponse) + assert result.state == State.RUNNING + assert result.is_alive is True + + @pytest.mark.asyncio + async def test_running_state_triggers_meta_store_update(self, sandbox_manager, mock_operator, mock_meta_store): + """RUNNING state writes back to meta_store when state changed.""" + mock_operator.get_status = AsyncMock(return_value=_make_sandbox_info(state=State.RUNNING)) + + await sandbox_manager.get_status("sandbox-1") + + mock_meta_store.update.assert_awaited_once() + + @pytest.mark.asyncio + async def test_operator_none_flag_false_raises_not_found(self, sandbox_manager, mock_operator, mock_meta_store): + """operator=None + include_all_states=False → not found, no DB fallback triggered.""" + mock_operator.get_status = AsyncMock(return_value=None) + + with pytest.raises(BadRequestRockError, match="not found"): + await sandbox_manager.get_status("sandbox-1", include_all_states=False) + + for c in mock_meta_store.get.call_args_list: + assert not c.kwargs.get("check_db") + + @pytest.mark.asyncio + async def test_operator_none_flag_true_calls_db_fallback(self, sandbox_manager, mock_operator, mock_meta_store): + """operator=None + include_all_states=True → meta_store.get(check_db=True) called.""" + mock_operator.get_status = AsyncMock(return_value=None) + mock_meta_store.get = AsyncMock(return_value=_make_sandbox_info(state=State.PENDING)) + + result = await sandbox_manager.get_status("sandbox-1", include_all_states=True) + + mock_meta_store.get.assert_awaited_once_with("sandbox-1", check_db=True) + assert result.state == State.PENDING + + @pytest.mark.asyncio + async def test_operator_none_flag_true_db_empty_raises(self, sandbox_manager, mock_operator): + """operator=None + include_all_states=True + DB empty → not found.""" + mock_operator.get_status = AsyncMock(return_value=None) + + with pytest.raises(BadRequestRockError, match="not found"): + await sandbox_manager.get_status("sandbox-1", include_all_states=True) + + @pytest.mark.asyncio + async def test_operator_data_with_flag_true_skips_db_fallback( + self, sandbox_manager, mock_operator, mock_meta_store + ): + """operator returns data + include_all_states=True → check_db=True never triggered.""" + mock_operator.get_status = AsyncMock(return_value=_make_sandbox_info(state=State.RUNNING)) + + await sandbox_manager.get_status("sandbox-1", include_all_states=True) + + for c in mock_meta_store.get.call_args_list: + assert not c.kwargs.get("check_db") From 30287d451194810f76656bfc78dfd9141e07447d Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Tue, 19 May 2026 11:28:44 +0800 Subject: [PATCH 121/226] feat(k8s): GPU support with Jinja2 templates and extensible accelerator types (#981) --- pyproject.toml | 1 + requirements_admin.txt | 4 + rock/admin/entrypoints/sandbox_api.py | 33 ++++- rock/admin/proto/request.py | 4 + rock/common/constants.py | 1 + rock/config.py | 45 +++++++ rock/deployments/config.py | 30 +++++ rock/sandbox/operator/k8s/constants.py | 5 + rock/sandbox/operator/k8s/provider.py | 30 +++-- rock/sandbox/operator/k8s/template_loader.py | 116 ++++++++--------- rock/utils/jinja_render.py | 42 ++++++ tests/unit/conftest.py | 29 ++++- .../sandbox/operator/test_k8s_provider.py | 26 ++++ .../operator/test_k8s_template_loader.py | 82 +++++++++++- tests/unit/test_config.py | 94 ++++++++++++- uv.lock | 123 ++++++++++++++++-- 16 files changed, 572 insertions(+), 93 deletions(-) create mode 100644 rock/utils/jinja_render.py diff --git a/pyproject.toml b/pyproject.toml index fbd220c545..883fd83e70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "rich", "oss2", "pyyaml", + "jinja2", "tzdata", ] diff --git a/requirements_admin.txt b/requirements_admin.txt index 332a15229d..60d0a43140 100644 --- a/requirements_admin.txt +++ b/requirements_admin.txt @@ -242,6 +242,8 @@ importlib-metadata==8.7.0 # opentelemetry-api incremental==24.7.2 # via twisted +jinja2==3.1.6 + # via rl-rock jmespath==0.10.0 # via # aliyun-python-sdk-core @@ -265,6 +267,8 @@ magiccube==0.3.0 # via reasoning-gym markdown-it-py==4.0.0 # via rich +markupsafe==3.0.3 + # via jinja2 math-verify==0.8.0 # via gem-llm matplotlib==3.10.7 diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index cc2940e4fc..ab49081b16 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -29,6 +29,7 @@ from rock.common.constants import ( CPU_OVERCOMMIT_ALLOWED_KEYS_KEY, CPU_OVERCOMMIT_HEADROOM_KEY, + EXTRA_ACCELERATOR_TYPES_KEY, GET_STATUS_SWITCH, KATA_DIND_DISK_SIZE_KEY, KATA_RUNTIME_SWITCH, @@ -37,8 +38,9 @@ SUPPORT_KATA_SWITCH, ) from rock.common.exception import handle_exceptions -from rock.deployments.config import DockerDeploymentConfig +from rock.deployments.config import AcceleratorType, DockerDeploymentConfig from rock.sandbox.sandbox_manager import SandboxManager +from rock.sdk.common.exceptions import BadRequestRockError sandbox_router = APIRouter() sandbox_manager: SandboxManager @@ -96,6 +98,33 @@ async def _apply_disk_limits(config: DockerDeploymentConfig) -> None: config.disk_limit_log = disk_limit_log +async def _apply_accelerator_type_validation(config: DockerDeploymentConfig) -> None: + """Validate ``config.accelerator_type`` against the built-in enum union with + Nacos-provided extras. + + Allowed set = ``AcceleratorType`` enum values ∪ list under Nacos key + ``extra_accelerator_types``. When Nacos is unavailable or the key is missing, + only the built-in enum applies. Raises :class:`BadRequestRockError` on + mismatch. ``None`` is always allowed (caller did not request a specific GPU). + """ + if config.accelerator_type is None: + return + + allowed: set[str] = {item.value for item in AcceleratorType} + + nacos = sandbox_manager.rock_config.nacos_provider + if nacos is not None: + nacos_config = await nacos.get_config() or {} + extras = nacos_config.get(EXTRA_ACCELERATOR_TYPES_KEY) or [] + if isinstance(extras, list): + allowed.update(str(item) for item in extras) + + if config.accelerator_type not in allowed: + raise BadRequestRockError( + f"Invalid accelerator_type {config.accelerator_type!r}. " f"Allowed values: {sorted(allowed)}" + ) + + async def _apply_cpu_overcommit_default(config: DockerDeploymentConfig, rock_authorization: str | None) -> None: """Derive limit_cpus from cpus + Nacos headroom when SDK did not set it. @@ -137,6 +166,7 @@ async def _apply_cpu_overcommit_default(config: DockerDeploymentConfig, rock_aut @handle_exceptions(error_message="start sandbox failed") async def start(request: SandboxStartRequest) -> RockResponse[SandboxStartResponse]: config = DockerDeploymentConfig.from_request(request) + await _apply_accelerator_type_validation(config) await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) await _apply_disk_limits(config) @@ -151,6 +181,7 @@ async def start_async( headers: Annotated[StartHeaders, Depends()], ) -> RockResponse[SandboxStartResponse]: config = DockerDeploymentConfig.from_request(request) + await _apply_accelerator_type_validation(config) await _apply_kata_runtime_switch(config) await _apply_kata_disk_size(config) await _apply_cpu_overcommit_default(config, headers.user_info.get("rock_authorization")) diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index d51de25241..1b49b141f2 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -39,6 +39,10 @@ class SandboxStartRequest(BaseModel): """Whether to use kata container runtime (io.containerd.kata.v2) instead of --privileged mode.""" auto_delete_seconds: int | None = None """The time for automatic container deletion, with the unit being seconds.""" + num_gpus: float | None = None + """Number of GPUs to allocate. Supports fractional values (e.g. 0.5 for GPU sharing).""" + accelerator_type: str | None = None + """GPU accelerator type (e.g. 'A100', 'V100'). If not specified, any available GPU will be used.""" class SandboxCommand(Command): diff --git a/rock/common/constants.py b/rock/common/constants.py index bf19e498e4..0af5006c7d 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -8,6 +8,7 @@ KATA_DIND_DISK_SIZE_KEY = "kata_dind_disk_size" SANDBOX_DISK_LIMIT_ROOTFS_KEY = "sandbox_disk_limit_rootfs" SANDBOX_DISK_LIMIT_LOG_KEY = "sandbox_disk_limit_log" +EXTRA_ACCELERATOR_TYPES_KEY = "extra_accelerator_types" PID_PREFIX = "PIDSTART" PID_SUFFIX = "PIDEND" SCHEDULER_LOG_NAME = "scheduler.log" diff --git a/rock/config.py b/rock/config.py index a09ffc2d04..d912219431 100644 --- a/rock/config.py +++ b/rock/config.py @@ -225,6 +225,13 @@ class K8sConfig: namespace: str = "rock" templates: dict[str, dict] = field(default_factory=dict) + # Paths (relative to env yaml dir, or absolute) of YAML files holding shared + # template definitions. Each file's top-level keys are template names. The + # resolver in RockConfig.from_env loads these in order (later wins on key + # conflict) then lets the inline `templates` block override. After loading + # the field is consumed (cleared) so K8sConfig only carries `templates`. + template_includes: list[str] = field(default_factory=list) + # Template mapping: image_os -> template_name, e.g., {"windows": "windows_template", "linux": "default"} template_map: dict[str, str] = field(default_factory=dict) @@ -281,6 +288,43 @@ def __post_init__(self) -> None: raise Exception("ROCK_ENVHUB_DB_URL is not an absolute path") +def _resolve_k8s_template_includes(k8s_dict: dict, base_dir: Path) -> None: + """Resolve K8sConfig.template_includes in place. + + Reads each path under `template_includes` (relative paths are anchored at + `base_dir`, typically the env yaml's parent directory), parses it as YAML + whose top-level keys are template names, and merges the result into + `templates`. Resolution rules: + + - Multiple includes: processed in declaration order; later wins per key. + - Inline `templates`: takes precedence over any include. + - Each template entry is replaced wholesale (no deep merge per template) so + env-level overrides are explicit and unambiguous. + + The `template_includes` key is removed from the dict after resolution so + that K8sConfig(**dict) sees an already-merged `templates` only. + """ + includes = k8s_dict.pop("template_includes", None) or [] + if not includes: + return + + inline = k8s_dict.get("templates") or {} + merged: dict[str, dict] = {} + for rel in includes: + path = Path(rel) + if not path.is_absolute(): + path = base_dir / path + if not path.exists(): + raise FileNotFoundError(f"k8s.template_includes references missing file: {path}") + with open(path) as f: + loaded = yaml.safe_load(f) or {} + if not isinstance(loaded, dict): + raise ValueError(f"k8s template include {path} must be a mapping at top level") + merged.update(loaded) + merged.update(inline) + k8s_dict["templates"] = merged + + @dataclass class RockConfig: ray: RayConfig = field(default_factory=RayConfig) @@ -318,6 +362,7 @@ def from_env(cls, config_path: str | None = None): if "ray" in config: kwargs["ray"] = RayConfig(**config["ray"]) if "k8s" in config: + _resolve_k8s_template_includes(config["k8s"], config_file.parent) kwargs["k8s"] = K8sConfig(**config["k8s"]) if "warmup" in config: kwargs["warmup"] = WarmupConfig(**config["warmup"]) diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 2ae440aae4..248bffcb57 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -7,6 +7,7 @@ """ from abc import abstractmethod +from enum import Enum from typing import Literal from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -17,6 +18,26 @@ from rock.utils import REQUEST_TIMEOUT_SECONDS +class AcceleratorType(str, Enum): + """GPU accelerator type, following Ray's naming convention. + + See: https://docs.ray.io/en/latest/ray-core/accelerator-types.html + """ + + V100 = "V100" + T4 = "T4" + A10G = "A10G" + L4 = "L4" + L40S = "L40S" + A100 = "A100" + A100_80G = "A100-80G" + H100 = "H100" + H200 = "H200" + B200 = "B200" + A10 = "A10" + GeForce_RTX_2080_Ti = "GeForce-RTX-2080-Ti" + + class DeploymentConfig(BaseModel): """Base configuration class for all deployment types.""" @@ -136,6 +157,15 @@ class DockerDeploymentConfig(DeploymentConfig): runtime_config: RuntimeConfig = Field(default_factory=RuntimeConfig) """Runtime configuration settings.""" + num_gpus: float | None = None + """Number of GPUs to allocate. Supports fractional values (e.g. 0.5 for GPU sharing).""" + + accelerator_type: str | None = None + """GPU accelerator type. Accepts any value present in the built-in + :class:`AcceleratorType` enum, or any extra type advertised through Nacos + (key ``extra_accelerator_types``). Validation is performed at the admin + API entry. If not specified, any available GPU will be used.""" + extended_params: dict[str, str] = Field(default_factory=dict) """Generic extension field for storing custom string key-value pairs.""" diff --git a/rock/sandbox/operator/k8s/constants.py b/rock/sandbox/operator/k8s/constants.py index fc53a8e9bf..b69b9b2def 100644 --- a/rock/sandbox/operator/k8s/constants.py +++ b/rock/sandbox/operator/k8s/constants.py @@ -25,5 +25,10 @@ class K8sConstants: EXT_TEMPLATE_NAME = "template_name" EXT_RESOURCE_VERSION = "k8s_resource_version" + # Built-in template names + TEMPLATE_DEFAULT = "default" + TEMPLATE_GPU_SINGLE = "gpu-single" + TEMPLATE_GPU_MULTI = "gpu-multi" + # Nacos config keys NACOS_POOLS_KEY = "pools" diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index bf4a073793..86eb875f2f 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -406,12 +406,15 @@ async def _get_pool_name(self, config: DockerDeploymentConfig) -> str | None: return ResourceMatchingPoolSelector().select_pool(config, pools) def _get_template_name(self, config: DockerDeploymentConfig) -> str: - """Get template name from extended_params or config template_map. + """Get template name from extended_params, GPU detection, or template_map. Priority: - 1. Check extended_params for template name - 2. Fallback to template_map based on image_os matching - 3. Return 'default' if not found + 1. Check extended_params for explicit template name + 2. If config.num_gpus == 1, auto-select 'gpu-single' (single full card) + 3. If config.num_gpus > 0 and != 1, auto-select 'gpu-multi' + (covers fractional shares <1 and multi-card >1) + 4. Fallback to template_map based on image_os matching + 5. Return 'default' if not found Args: config: Docker deployment configuration @@ -419,19 +422,28 @@ def _get_template_name(self, config: DockerDeploymentConfig) -> str: Returns: Template name (defaults to 'default') """ - # Priority 1: Check extended_params + # Priority 1: Check extended_params (explicit override) template_name = config.extended_params.get(K8sConstants.EXT_TEMPLATE_NAME) if template_name: return template_name - # Priority 2: Check template_map based on image_os + # Priority 2: Single full GPU goes to the single-card template + if config.num_gpus is not None and config.num_gpus == 1: + return K8sConstants.TEMPLATE_GPU_SINGLE + + # Priority 3: Any other positive num_gpus (fractional <1 or multi >1) + # routes to the multi-GPU template + if config.num_gpus is not None and config.num_gpus > 0: + return K8sConstants.TEMPLATE_GPU_MULTI + + # Priority 4: Check template_map based on image_os if config.image_os and self._k8s_config.template_map: mapped_template = self._k8s_config.template_map.get(config.image_os) if mapped_template: return mapped_template - # Priority 3: Return default - return "default" + # Priority 5: Return default + return K8sConstants.TEMPLATE_DEFAULT def _normalize_memory(self, memory: str) -> str: """Normalize memory format to Kubernetes standard. @@ -553,6 +565,8 @@ async def _build_batchsandbox_manifest(self, config: DockerDeploymentConfig) -> image=config.image, cpus=config.cpus, memory=self._normalize_memory(config.memory), + num_gpus=config.num_gpus, + accelerator_type=config.accelerator_type, ) logger.debug( diff --git a/rock/sandbox/operator/k8s/template_loader.py b/rock/sandbox/operator/k8s/template_loader.py index 57d58e658c..5832145089 100644 --- a/rock/sandbox/operator/k8s/template_loader.py +++ b/rock/sandbox/operator/k8s/template_loader.py @@ -4,8 +4,11 @@ import json from typing import Any +import jinja2 + from rock.logger import init_logger from rock.sandbox.operator.k8s.constants import K8sConstants +from rock.utils.jinja_render import render_node logger = init_logger(__name__) @@ -26,6 +29,8 @@ def __init__(self, templates: dict[str, dict[str, Any]], default_namespace: str if not self._templates: raise ValueError("No templates provided. At least one template must be defined in K8sConfig.templates.") + self._jinja_env = jinja2.Environment(undefined=jinja2.StrictUndefined, autoescape=False) + logger.info(f"Loaded {len(self._templates)} K8S templates from config") logger.debug(f"Available templates: {', '.join(self._templates.keys())}") @@ -50,48 +55,45 @@ def get_template(self, template_name: str = "default") -> dict[str, Any]: def build_manifest( self, template_name: str = "default", - sandbox_id: str = None, - image: str = None, - cpus: float = None, - memory: str = None, + sandbox_id: str | None = None, + image: str | None = None, + cpus: float | None = None, + memory: str | None = None, + num_gpus: int | None = None, + accelerator_type: str | None = None, ) -> dict[str, Any]: """Build a complete BatchSandbox manifest from template. - Template structure: - - namespace: K8S namespace for the sandbox (REQUIRED) - - ports: custom port configuration (not part of K8S manifest) - - template: corresponds to spec.template in BatchSandbox CRD - - template.metadata -> spec.template.metadata - - template.spec -> spec.template.spec (Pod spec) + The template is rendered with Jinja2: every string value is treated as + a Jinja2 template against a ``ctx`` built from the call arguments. + ``None`` arguments enter ``ctx`` as ``""`` so that: + + * plain ``{{ var }}`` placeholders collapse to empty strings and the + drop-empty rule removes the surrounding dict key / list element; + * ``{{ var | default('x', true) }}`` placeholders fall back to the + template-supplied default. - Top-level fields are hardcoded: - - apiVersion: sandbox.opensandbox.io/v1alpha1 - - kind: BatchSandbox - - metadata: constructed from parameters - - spec.replicas: always 1 + The CRD wrapper (apiVersion/kind/metadata/spec.replicas) and the + sandbox-id / template / resource-speedup labels and ports annotation + are still assembled in code, since they are structural rather than + configurable. Args: - template_name: Name of the template to use - sandbox_id: Sandbox identifier - image: Container image - cpus: CPU resource limit - memory: Memory resource limit (normalized format like '2Gi') + template_name: Name of the template to use. + sandbox_id: Sandbox identifier (auto-generated if missing). + image: Container image (rendered into the template via {{ image }}). + cpus: CPU resource value (rendered via {{ cpus }}). + memory: Memory resource value (rendered via {{ memory }}). + num_gpus: GPU count (rendered via {{ num_gpus }}). + accelerator_type: GPU model (rendered via {{ accelerator_type }}). Returns: - Complete BatchSandbox manifest + Complete BatchSandbox manifest. """ import uuid - # Get template configuration config = self.get_template(template_name) - # Use default namespace (configured at startup) - namespace = self._default_namespace - - # Get enable_resource_speedup from template (default to True) - enable_resource_speedup = config.get("enable_resource_speedup", True) - - # Get port configuration from template (required) ports_config = config.get("ports") if not ports_config: raise ValueError( @@ -99,22 +101,33 @@ def build_manifest( f"Each template must define ports (proxy, server, ssh)." ) - # Extract template (corresponds to spec.template in BatchSandbox) - pod_template = config.get("template", {}) - template_metadata = copy.deepcopy(pod_template.get("metadata", {})) - pod_spec = copy.deepcopy(pod_template.get("spec", {})) - - # Generate sandbox_id if not provided if not sandbox_id: sandbox_id = f"sandbox-{uuid.uuid4().hex[:8]}" - # Build top-level BatchSandbox manifest (hardcoded structure) + # num_gpus stays numeric so templates can do arithmetic; cpus str-coerced to pin float->"4.0" formatting. + ctx = { + "sandbox_id": sandbox_id, + "template_name": template_name, + "image": image if image is not None else "", + "cpus": str(cpus) if cpus is not None else "", + "memory": memory if memory is not None else "", + "num_gpus": num_gpus if num_gpus is not None else "", + "accelerator_type": accelerator_type if accelerator_type is not None else "", + } + + rendered = render_node(config, self._jinja_env, ctx) + + enable_resource_speedup = rendered.get("enable_resource_speedup", True) + pod_template = rendered.get("template", {}) + template_metadata = pod_template.get("metadata", {}) + pod_spec = pod_template.get("spec", {}) + manifest = { "apiVersion": K8sConstants.CRD_API_VERSION, "kind": K8sConstants.CRD_KIND, "metadata": { "name": sandbox_id, - "namespace": namespace, + "namespace": self._default_namespace, "labels": { K8sConstants.LABEL_SANDBOX_ID: sandbox_id, K8sConstants.LABEL_TEMPLATE: template_name, @@ -124,45 +137,18 @@ def build_manifest( }, }, "spec": { - "replicas": 1, # Always 1 for sandbox + "replicas": 1, "template": {"metadata": template_metadata, "spec": pod_spec}, }, } - # Add resource speedup label if enabled if enable_resource_speedup: manifest["metadata"]["labels"][K8sConstants.LABEL_RESOURCE_SPEEDUP] = "true" - # Add sandbox-id label to template metadata if "labels" not in manifest["spec"]["template"]["metadata"]: manifest["spec"]["template"]["metadata"]["labels"] = {} manifest["spec"]["template"]["metadata"]["labels"][K8sConstants.LABEL_SANDBOX_ID] = sandbox_id - # Set container image - if image: - containers = pod_spec.get("containers", []) - if containers and len(containers) > 0: - containers[0]["image"] = image - - # Set resources if provided - if cpus is not None or memory is not None: - containers = pod_spec.get("containers", []) - if containers and len(containers) > 0: - if "resources" not in containers[0]: - containers[0]["resources"] = {} - - if cpus is not None or memory is not None: - containers[0]["resources"]["requests"] = {} - containers[0]["resources"]["limits"] = {} - - if cpus is not None: - containers[0]["resources"]["requests"]["cpu"] = str(cpus) - containers[0]["resources"]["limits"]["cpu"] = str(cpus) - - if memory is not None: - containers[0]["resources"]["requests"]["memory"] = memory - containers[0]["resources"]["limits"]["memory"] = memory - return manifest @property diff --git a/rock/utils/jinja_render.py b/rock/utils/jinja_render.py new file mode 100644 index 0000000000..905c9631e7 --- /dev/null +++ b/rock/utils/jinja_render.py @@ -0,0 +1,42 @@ +"""Jinja2 helpers for rendering nested template structures.""" + +from collections.abc import Mapping +from typing import Any + +import jinja2 + +_DROP = object() # sentinel: rendered placeholder collapsed to empty + + +def render_node(node: Any, env: jinja2.Environment, ctx: Mapping[str, Any]) -> Any: + """Recursively render a template node with Jinja2. + + Strings containing ``{{`` are rendered against ``ctx``; an empty + rendered result causes the surrounding dict key to be dropped or + the surrounding list element to be skipped. Non-string scalars + pass through unchanged. + """ + if isinstance(node, str): + if "{{" not in node: + return node + rendered = env.from_string(node).render(**ctx).strip() + if rendered == "": + return _DROP + return rendered + if isinstance(node, dict): + result: dict[Any, Any] = {} + for k, v in node.items(): + rendered = render_node(v, env, ctx) + if rendered is _DROP: + continue + result[k] = rendered + return result + if isinstance(node, list): + result_list: list[Any] = [] + for item in node: + rendered = render_node(item, env, ctx) + if rendered is _DROP: + continue + result_list.append(rendered) + return result_list + return node diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 49837907b6..9e4a5172b2 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -104,7 +104,9 @@ async def sandbox_manager( ray_operator, _memory_sandbox_table: SandboxTable, ): - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config) + meta_store = SandboxMetaStore( + redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config + ) sandbox_manager = SandboxManager( rock_config, meta_store=meta_store, @@ -120,7 +122,9 @@ async def sandbox_manager( async def sandbox_proxy_service( rock_config: RockConfig, redis_provider: RedisProvider, _memory_sandbox_table: SandboxTable ): - meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config) + meta_store = SandboxMetaStore( + redis_provider=redis_provider, sandbox_table=_memory_sandbox_table, rock_config=rock_config + ) sandbox_proxy_service = SandboxProxyService(rock_config, meta_store=meta_store) return sandbox_proxy_service @@ -202,7 +206,7 @@ def k8s_config(): @pytest.fixture def basic_templates(): - """Create basic template configuration.""" + """Create basic template configuration with Jinja2 placeholders.""" return { "default": { "ports": { @@ -212,7 +216,24 @@ def basic_templates(): }, "template": { "metadata": {"labels": {"app": "rock-sandbox"}}, - "spec": {"containers": [{"name": "main", "image": "python:3.11"}]}, + "spec": { + "containers": [ + { + "name": "main", + "image": "{{ image | default('python:3.11', true) }}", + "resources": { + "requests": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + }, + "limits": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + }, + }, + } + ] + }, }, } } diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index d403913ba2..0c81064a6d 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -36,6 +36,7 @@ def make_config( memory: str = "4Gi", extended_params: dict = None, image_os: str = "linux", + num_gpus: float | None = None, ) -> DockerDeploymentConfig: return DockerDeploymentConfig( image=image, @@ -44,6 +45,7 @@ def make_config( container_name="test-sandbox", extended_params=extended_params or {}, image_os=image_os, + num_gpus=num_gpus, ) @@ -210,6 +212,30 @@ def test_returns_default_when_no_params_and_no_template_map(self): config = make_config() assert provider._get_template_name(config) == "default" + def test_returns_gpu_single_when_num_gpus_is_one(self): + """num_gpus == 1 (single full card) routes to 'gpu-single'.""" + provider = make_provider() + config = make_config(num_gpus=1) + assert provider._get_template_name(config) == K8sConstants.TEMPLATE_GPU_SINGLE + + def test_returns_gpu_multi_when_fractional_lt_one(self): + """Fractional GPU (0 < num_gpus < 1) routes to 'gpu-multi'.""" + provider = make_provider() + config = make_config(num_gpus=0.5) + assert provider._get_template_name(config) == K8sConstants.TEMPLATE_GPU_MULTI + + def test_returns_gpu_multi_when_num_gpus_gt_one(self): + """Multi-GPU (num_gpus > 1) routes to 'gpu-multi'.""" + provider = make_provider() + config = make_config(num_gpus=2) + assert provider._get_template_name(config) == K8sConstants.TEMPLATE_GPU_MULTI + + def test_extended_params_takes_priority_over_gpu_routing(self): + """extended_params template_name beats the GPU auto-selection.""" + provider = make_provider() + config = make_config(extended_params={"template_name": "custom"}, num_gpus=4) + assert provider._get_template_name(config) == "custom" + # ========== _get_pool_ports ========== diff --git a/tests/unit/sandbox/operator/test_k8s_template_loader.py b/tests/unit/sandbox/operator/test_k8s_template_loader.py index 5d5edfe7fc..3e7fdfa259 100644 --- a/tests/unit/sandbox/operator/test_k8s_template_loader.py +++ b/tests/unit/sandbox/operator/test_k8s_template_loader.py @@ -93,8 +93,13 @@ def test_build_manifest_without_resources(self, template_loader): container = manifest["spec"]["template"]["spec"]["containers"][0] - # Should not have resources section if not specified - assert "resources" not in container or not container.get("resources") + # Should not have any concrete resource values when nothing is specified. + # The Jinja2-based render keeps the template's resources skeleton but + # drops any keys whose placeholder rendered to empty (cpus, memory). + assert "resources" in container + resources = container["resources"] + assert resources.get("requests", {}) == {} + assert resources.get("limits", {}) == {} def test_build_manifest_with_custom_image(self, template_loader): """Test building manifest with custom image.""" @@ -150,3 +155,76 @@ def test_build_manifest_adds_sandbox_id_to_pod_labels(self, template_loader): pod_labels = manifest["spec"]["template"]["metadata"]["labels"] assert pod_labels[K8sConstants.LABEL_SANDBOX_ID] == "test-sandbox" + + def test_build_manifest_gpu_template(self): + """GPU placeholders fill correctly when num_gpus and accelerator_type provided.""" + templates = { + "gpu": { + "ports": {"proxy": 8000, "server": 8080, "ssh": 22}, + "template": { + "spec": { + "containers": [ + { + "name": "main", + "image": "{{ image | default('cuda:12', true) }}", + "resources": { + "limits": { + "nvidia.com/gpu": ("{{ num_gpus if num_gpus else '' }}"), + } + }, + } + ], + "nodeSelector": { + "nvidia.com/gpu.product": "{{ accelerator_type }}", + }, + } + }, + } + } + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") + + manifest = loader.build_manifest( + template_name="gpu", + sandbox_id="test-gpu", + num_gpus=4, + accelerator_type="A100", + ) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert container["resources"]["limits"]["nvidia.com/gpu"] == "4" + + nodeSelector = manifest["spec"]["template"]["spec"]["nodeSelector"] + assert nodeSelector["nvidia.com/gpu.product"] == "A100" + + def test_build_manifest_drops_gpu_when_no_gpu(self): + """When num_gpus omitted, GPU keys collapse out of resources.limits.""" + templates = { + "gpu": { + "ports": {"proxy": 8000, "server": 8080, "ssh": 22}, + "template": { + "spec": { + "containers": [ + { + "name": "main", + "image": "{{ image | default('cuda:12', true) }}", + "resources": { + "limits": { + "cpu": "{{ cpus | default('2', true) }}", + "nvidia.com/gpu": ("{{ num_gpus if num_gpus else '' }}"), + } + }, + } + ], + } + }, + } + } + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") + + manifest = loader.build_manifest(template_name="gpu", sandbox_id="test-cpu") + + limits = manifest["spec"]["template"]["spec"]["containers"][0]["resources"]["limits"] + # cpu has a default → present + assert limits["cpu"] == "2" + # GPU placeholders rendered to empty → keys dropped + assert "nvidia.com/gpu" not in limits diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 5bfdf77b90..f6a993e1c8 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,6 +1,10 @@ +import tempfile +from pathlib import Path + import pytest +import yaml -from rock.config import RockConfig, RuntimeConfig +from rock.config import RockConfig, RuntimeConfig, _resolve_k8s_template_includes @pytest.mark.asyncio @@ -145,3 +149,91 @@ def test_sandbox_config_coerces_nested_dicts_from_yaml(): assert cfg.log.keep_days_before_archive == 7 assert isinstance(cfg.file_transfer, SandboxFileTransferConfig) assert cfg.file_transfer.prefix == "rock-transfer/" + + +# ===== _resolve_k8s_template_includes ===== + + +def _write_yaml(path: Path, data: dict) -> None: + path.write_text(yaml.safe_dump(data)) + + +@pytest.mark.asyncio +async def test_resolve_includes_noop_without_includes(): + """No template_includes key → templates passed through unchanged.""" + k8s = {"templates": {"default": {"a": 1}}} + _resolve_k8s_template_includes(k8s, Path("/unused")) + assert k8s == {"templates": {"default": {"a": 1}}} + + +@pytest.mark.asyncio +async def test_resolve_includes_single_file_relative(): + """Single relative-path include populates templates.""" + with tempfile.TemporaryDirectory() as tmp: + base = Path(tmp) + inc_dir = base / "templates" / "k8s" + inc_dir.mkdir(parents=True) + _write_yaml(inc_dir / "default.yml", {"default": {"ports": {"proxy": 8000}}}) + k8s = {"template_includes": ["templates/k8s/default.yml"]} + _resolve_k8s_template_includes(k8s, base) + assert "template_includes" not in k8s + assert k8s["templates"] == {"default": {"ports": {"proxy": 8000}}} + + +@pytest.mark.asyncio +async def test_resolve_includes_multiple_later_wins(): + """Conflicts across includes resolved by declaration order — later wins.""" + with tempfile.TemporaryDirectory() as tmp: + base = Path(tmp) + _write_yaml(base / "a.yml", {"shared": {"v": 1}, "only_a": {"x": "a"}}) + _write_yaml(base / "b.yml", {"shared": {"v": 2}, "only_b": {"x": "b"}}) + k8s = {"template_includes": ["a.yml", "b.yml"]} + _resolve_k8s_template_includes(k8s, base) + assert k8s["templates"]["shared"] == {"v": 2} + assert k8s["templates"]["only_a"] == {"x": "a"} + assert k8s["templates"]["only_b"] == {"x": "b"} + + +@pytest.mark.asyncio +async def test_resolve_includes_inline_overrides_include(): + """Inline templates entry replaces the include version wholesale.""" + with tempfile.TemporaryDirectory() as tmp: + base = Path(tmp) + _write_yaml(base / "base.yml", {"default": {"from": "include"}}) + k8s = { + "template_includes": ["base.yml"], + "templates": {"default": {"from": "inline"}}, + } + _resolve_k8s_template_includes(k8s, base) + assert k8s["templates"] == {"default": {"from": "inline"}} + + +@pytest.mark.asyncio +async def test_resolve_includes_absolute_path(): + """Absolute paths bypass base_dir resolution.""" + with tempfile.TemporaryDirectory() as tmp: + abs_file = Path(tmp) / "abs.yml" + _write_yaml(abs_file, {"default": {"k": "v"}}) + k8s = {"template_includes": [str(abs_file)]} + _resolve_k8s_template_includes(k8s, Path("/nonexistent")) + assert k8s["templates"] == {"default": {"k": "v"}} + + +@pytest.mark.asyncio +async def test_resolve_includes_missing_file_raises(): + """Missing include path surfaces a clear FileNotFoundError.""" + with tempfile.TemporaryDirectory() as tmp: + k8s = {"template_includes": ["does-not-exist.yml"]} + with pytest.raises(FileNotFoundError, match="does-not-exist.yml"): + _resolve_k8s_template_includes(k8s, Path(tmp)) + + +@pytest.mark.asyncio +async def test_resolve_includes_non_mapping_raises(): + """A list-shaped include yaml is rejected.""" + with tempfile.TemporaryDirectory() as tmp: + base = Path(tmp) + (base / "bad.yml").write_text("- not_a_mapping\n") + k8s = {"template_includes": ["bad.yml"]} + with pytest.raises(ValueError, match="must be a mapping"): + _resolve_k8s_template_includes(k8s, base) diff --git a/uv.lock b/uv.lock index 56fba7522c..3ef30586fa 100644 --- a/uv.lock +++ b/uv.lock @@ -1863,7 +1863,7 @@ name = "hyperlink" version = "21.0.0" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } dependencies = [ - { name = "idna" }, + { name = "idna", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://mirrors.aliyun.com/pypi/packages/3a/51/1947bd81d75af87e3bb9e34593a4cf118115a8feb451ce7a69044ef1412e/hyperlink-21.0.0.tar.gz", hash = "sha256:427af957daa58bc909471c6c40f74c5450fa123dd093fc53efd2e91d2705a56b" } wheels = [ @@ -1919,8 +1919,8 @@ name = "incremental" version = "24.7.2" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } dependencies = [ - { name = "setuptools" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "setuptools", marker = "sys_platform != 'win32'" }, + { name = "tomli", marker = "python_full_version < '3.11' and sys_platform != 'win32'" }, ] sdist = { url = "https://mirrors.aliyun.com/pypi/packages/27/87/156b374ff6578062965afe30cc57627d35234369b3336cf244b240c8d8e6/incremental-24.7.2.tar.gz", hash = "sha256:fb4f1d47ee60efe87d4f6f0ebb5f70b9760db2b2574c59c8e8912be4ebd464c9" } wheels = [ @@ -1936,6 +1936,18 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12" }, ] +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67" }, +] + [[package]] name = "jiter" version = "0.14.0" @@ -2266,6 +2278,91 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147" }, ] +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/e8/4b/3541d44f3937ba468b75da9eebcae497dcf67adb65caa16760b0a6807ebb/markupsafe-3.0.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559" }, + { url = "https://mirrors.aliyun.com/pypi/packages/98/1b/fbd8eed11021cabd9226c37342fa6ca4e8a98d8188a8d9b66740494960e4/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419" }, + { url = "https://mirrors.aliyun.com/pypi/packages/40/01/e560d658dc0bb8ab762670ece35281dec7b6c1b33f5fbc09ebb57a185519/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ba88449deb3de88bd40044603fafffb7bc2b055d626a330323a9ed736661695" }, + { url = "https://mirrors.aliyun.com/pypi/packages/af/cd/ce6e848bbf2c32314c9b237839119c5a564a59725b53157c856e90937b7a/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f42d0984e947b8adf7dd6dde396e720934d12c506ce84eea8476409563607591" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c9/2a/b5c12c809f1c3045c4d580b035a743d12fcde53cf685dbc44660826308da/markupsafe-3.0.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c0c0b3ade1c0b13b936d7970b1d37a57acde9199dc2aecc4c336773e1d86049c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/cf/e3/9427a68c82728d0a88c50f890d0fc072a1484de2f3ac1ad0bfc1a7214fd5/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:0303439a41979d9e74d18ff5e2dd8c43ed6c6001fd40e5bf2e43f7bd9bbc523f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bc/36/23578f29e9e582a4d0278e009b38081dbe363c5e7165113fad546918a232/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:d2ee202e79d8ed691ceebae8e0486bd9a2cd4794cec4824e1c99b6f5009502f6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/56/21/dca11354e756ebd03e036bd8ad58d6d7168c80ce1fe5e75218e4945cbab7/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:177b5253b2834fe3678cb4a5f0059808258584c559193998be2601324fdeafb1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/87/99/faba9369a7ad6e4d10b6a5fbf71fa2a188fe4a593b15f0963b73859a1bbd/markupsafe-3.0.3-cp310-cp310-win32.whl", hash = "sha256:2a15a08b17dd94c53a1da0438822d70ebcd13f8c3a95abe3a9ef9f11a94830aa" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d6/25/55dc3ab959917602c96985cb1253efaa4ff42f71194bddeb61eb7278b8be/markupsafe-3.0.3-cp310-cp310-win_amd64.whl", hash = "sha256:c4ffb7ebf07cfe8931028e3e4c85f0357459a3f9f9490886198848f4fa002ec8" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d0/9e/0a02226640c255d1da0b8d12e24ac2aa6734da68bff14c05dd53b94a0fc3/markupsafe-3.0.3-cp310-cp310-win_arm64.whl", hash = "sha256:e2103a929dfa2fcaf9bb4e7c091983a49c9ac3b19c9061b6d5427dd7d14d81a1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50" }, + { url = "https://mirrors.aliyun.com/pypi/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf" }, + { url = "https://mirrors.aliyun.com/pypi/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115" }, + { url = "https://mirrors.aliyun.com/pypi/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19" }, + { url = "https://mirrors.aliyun.com/pypi/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01" }, + { url = "https://mirrors.aliyun.com/pypi/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219" }, + { url = "https://mirrors.aliyun.com/pypi/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12" }, + { url = "https://mirrors.aliyun.com/pypi/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed" }, + { url = "https://mirrors.aliyun.com/pypi/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73" }, + { url = "https://mirrors.aliyun.com/pypi/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19" }, + { url = "https://mirrors.aliyun.com/pypi/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb" }, + { url = "https://mirrors.aliyun.com/pypi/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009" }, + { url = "https://mirrors.aliyun.com/pypi/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287" }, + { url = "https://mirrors.aliyun.com/pypi/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026" }, + { url = "https://mirrors.aliyun.com/pypi/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737" }, + { url = "https://mirrors.aliyun.com/pypi/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe" }, + { url = "https://mirrors.aliyun.com/pypi/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581" }, + { url = "https://mirrors.aliyun.com/pypi/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab" }, + { url = "https://mirrors.aliyun.com/pypi/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634" }, + { url = "https://mirrors.aliyun.com/pypi/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50" }, + { url = "https://mirrors.aliyun.com/pypi/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5" }, + { url = "https://mirrors.aliyun.com/pypi/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523" }, + { url = "https://mirrors.aliyun.com/pypi/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa" }, +] + [[package]] name = "math-verify" version = "0.8.0" @@ -3043,7 +3140,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://mirrors.aliyun.com/pypi/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f" } wheels = [ @@ -4174,12 +4271,13 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.7.1" +version = "1.8.0" source = { editable = "." } dependencies = [ { name = "anyio" }, { name = "build" }, { name = "httpx" }, + { name = "jinja2" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-exporter-prometheus" }, @@ -4322,6 +4420,7 @@ requires-dist = [ { name = "gem-llm", marker = "extra == 'sandbox-actor'", specifier = ">=0.1.0" }, { name = "httpx" }, { name = "httpx", marker = "extra == 'model-service'" }, + { name = "jinja2" }, { name = "kubernetes", marker = "extra == 'admin'", specifier = ">=35.0.0" }, { name = "nacos-sdk-python", marker = "extra == 'admin'", specifier = ">=0.1.14" }, { name = "nacos-sdk-python", marker = "extra == 'sandbox-actor'", specifier = ">=0.1.14" }, @@ -4888,13 +4987,13 @@ name = "twisted" version = "25.5.0" source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } dependencies = [ - { name = "attrs" }, - { name = "automat" }, - { name = "constantly" }, - { name = "hyperlink" }, - { name = "incremental" }, - { name = "typing-extensions" }, - { name = "zope-interface" }, + { name = "attrs", marker = "sys_platform != 'win32'" }, + { name = "automat", marker = "sys_platform != 'win32'" }, + { name = "constantly", marker = "sys_platform != 'win32'" }, + { name = "hyperlink", marker = "sys_platform != 'win32'" }, + { name = "incremental", marker = "sys_platform != 'win32'" }, + { name = "typing-extensions", marker = "sys_platform != 'win32'" }, + { name = "zope-interface", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://mirrors.aliyun.com/pypi/packages/13/0f/82716ed849bf7ea4984c21385597c949944f0f9b428b5710f79d0afc084d/twisted-25.5.0.tar.gz", hash = "sha256:1deb272358cb6be1e3e8fc6f9c8b36f78eb0fa7c2233d2dbe11ec6fee04ea316" } wheels = [ From cc3a0f59d36237536bd1e1ae5cfcbd73576be6b1 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 19:19:13 +0800 Subject: [PATCH 122/226] feature(scheduler): add BuildCacheCleanupTask for uv/pip cache pruning Periodically prunes uv and pip caches on each worker. Each tool runs in its own self-skipping shell snippet (if/then/else with `command -v` guard), so a worker that lacks one of them does not fail the whole task. Key behavior vs naive `a && b || c`: the if/then/else structure emits "skipped (not installed)" only when `command -v` actually fails; if the tool is present but `prune` itself errors (permission, disk full, etc.), the output is "prune failed". The legacy form would also fire the "skipped" branch on a prune error, misleading operators in logs. tools is yml-configurable (subset of {"uv", "pip"}, default both); unknown tool names raise at construction time so misconfig fails fast. Idempotent (uv/pip caches only drop unreferenced entries). Default interval 24h. Enable per-environment via the existing `enabled:` flag in scheduler.tasks. Co-Authored-By: Claude Opus 4.7 (1M context) --- rock/admin/scheduler/tasks/__init__.py | 2 + .../tasks/build_cache_cleanup_task.py | 86 ++++++++++++ .../test_build_cache_cleanup_task.py | 131 ++++++++++++++++++ 3 files changed, 219 insertions(+) create mode 100644 rock/admin/scheduler/tasks/build_cache_cleanup_task.py create mode 100644 tests/unit/admin/scheduler/test_build_cache_cleanup_task.py diff --git a/rock/admin/scheduler/tasks/__init__.py b/rock/admin/scheduler/tasks/__init__.py index 700b0c73a6..41b251635f 100644 --- a/rock/admin/scheduler/tasks/__init__.py +++ b/rock/admin/scheduler/tasks/__init__.py @@ -1,4 +1,5 @@ # rock/admin/scheduler/tasks/__init__.py +from rock.admin.scheduler.tasks.build_cache_cleanup_task import BuildCacheCleanupTask from rock.admin.scheduler.tasks.container_cleanup_task import ContainerCleanupTask from rock.admin.scheduler.tasks.file_cleanup_task import FileCleanupTask from rock.admin.scheduler.tasks.image_cleanup_task import ImageCleanupTask @@ -6,6 +7,7 @@ from rock.admin.scheduler.tasks.ray_log_cleanup_task import RayLogCleanupTask __all__ = [ + "BuildCacheCleanupTask", "ContainerCleanupTask", "FileCleanupTask", "ImageCleanupTask", diff --git a/rock/admin/scheduler/tasks/build_cache_cleanup_task.py b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py new file mode 100644 index 0000000000..d7d2443e4a --- /dev/null +++ b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py @@ -0,0 +1,86 @@ +"""Prune build/install caches (uv, pip) on each worker.""" + +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.common.constants import SCHEDULER_LOG_NAME +from rock.logger import init_logger +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + +logger = init_logger(name="build_cache_cleanup", file_name=SCHEDULER_LOG_NAME) + +# Map tool name -> shell snippet that prunes its cache. Each snippet uses an +# explicit if/then/else so the "not installed" message is only emitted when +# `command -v` actually fails — otherwise the legacy `a && b || c` form would +# also fire `c` when `b` failed (permission/disk error), misleading operators +# into thinking the tool was missing when in fact the prune itself errored. +_TOOL_COMMANDS: dict[str, str] = { + "uv": ( + "if command -v uv >/dev/null 2>&1; then " + ' uv cache prune 2>&1 || echo "uv: prune failed"; ' + 'else echo "uv: skipped (not installed)"; fi' + ), + "pip": ( + "if command -v pip >/dev/null 2>&1; then " + ' pip cache purge 2>&1 || echo "pip: prune failed"; ' + 'else echo "pip: skipped (not installed)"; fi' + ), +} + + +class BuildCacheCleanupTask(BaseTask): + """Prune build/install caches for the configured tools. + + Each tool runs in its own self-skipping shell snippet, so a worker that + lacks one of them won't fail the whole task. Default `tools` covers both + common tools used in the standard worker image; restrict via yml if you + want to skip one of them. + """ + + def __init__( + self, + interval_seconds: int = 86400, + tools: list[str] | None = None, + ): + """ + Args: + interval_seconds: Execution interval, default 24 hours. + tools: Subset of {"uv", "pip"}; default both. Unknown names raise. + """ + super().__init__( + type="build_cache_cleanup", + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + tools = tools if tools is not None else ["uv", "pip"] + unknown = [t for t in tools if t not in _TOOL_COMMANDS] + if unknown: + raise ValueError( + f"Unsupported build cache tool(s): {unknown}. Supported: {sorted(_TOOL_COMMANDS)}" + ) + self.tools = tools + + @classmethod + def from_config(cls, task_config) -> "BuildCacheCleanupTask": + return cls( + interval_seconds=task_config.interval_seconds, + tools=task_config.params.get("tools"), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + # Sequential `; ` (not `&& `): each tool's snippet already converts + # "missing" to a soft echo, but `;` keeps any future tool that emits a + # non-zero exit from short-circuiting the rest. + snippets = [_TOOL_COMMANDS[t] for t in self.tools] + command = "; ".join(snippets) if snippets else "echo 'no tools configured'" + result = await runtime.execute(Command(command=command, shell=True, check=False)) + output = (result.stdout or "").strip() + logger.info( + f"[{self.type}] [{runtime._config.host}] cache prune done: " + f"tools={self.tools}, exit={result.exit_code}, output_head={output[:200]}" + ) + return { + "status": TaskStatusEnum.SUCCESS, + "tools": self.tools, + "exit_code": result.exit_code, + "output_head": output[:500], + } diff --git a/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py b/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py new file mode 100644 index 0000000000..bb42576bfe --- /dev/null +++ b/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py @@ -0,0 +1,131 @@ +"""Tests for BuildCacheCleanupTask.""" + +from unittest.mock import AsyncMock + +import pytest + +from rock.admin.scheduler.task_base import TaskStatusEnum +from rock.admin.scheduler.tasks.build_cache_cleanup_task import BuildCacheCleanupTask + + +class _FakeTaskConfig: + def __init__(self, params=None, interval_seconds=86400): + self.params = params or {} + self.interval_seconds = interval_seconds + + +class _FakeExecResult: + def __init__(self, exit_code=0, stdout="Pruned 0 entries"): + self.exit_code = exit_code + self.stdout = stdout + + +def _runtime(stdout="Pruned 0 entries", exit_code=0): + rt = AsyncMock() + rt._config = type("C", (), {"host": "10.0.0.1"})() + rt.execute = AsyncMock(return_value=_FakeExecResult(exit_code=exit_code, stdout=stdout)) + return rt + + +class TestInit: + def test_default_tools(self): + task = BuildCacheCleanupTask() + assert task.type == "build_cache_cleanup" + assert task.interval_seconds == 86400 + assert task.tools == ["uv", "pip"] + + def test_custom_tools_subset(self): + task = BuildCacheCleanupTask(tools=["uv"]) + assert task.tools == ["uv"] + + def test_empty_tools_list_is_allowed(self): + # Explicit empty list -> task no-ops (echo line). Different from None, + # which means "use default". + task = BuildCacheCleanupTask(tools=[]) + assert task.tools == [] + + def test_unknown_tool_raises(self): + with pytest.raises(ValueError, match="Unsupported build cache tool"): + BuildCacheCleanupTask(tools=["uv", "npm"]) + + +class TestFromConfig: + def test_from_config_defaults(self): + task = BuildCacheCleanupTask.from_config(_FakeTaskConfig()) + assert task.tools == ["uv", "pip"] + + def test_from_config_custom_tools(self): + cfg = _FakeTaskConfig(params={"tools": ["pip"]}, interval_seconds=3600) + task = BuildCacheCleanupTask.from_config(cfg) + assert task.tools == ["pip"] + assert task.interval_seconds == 3600 + + +class TestRunAction: + @pytest.mark.asyncio + async def test_default_command_invokes_both_tools(self): + task = BuildCacheCleanupTask() + runtime = _runtime() + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + assert result["tools"] == ["uv", "pip"] + + cmd = runtime.execute.await_args.args[0].command + assert "uv cache prune" in cmd + assert "pip cache purge" in cmd + # Each step has its own command -v guard so missing tool != failure. + assert "command -v uv" in cmd + assert "command -v pip" in cmd + + @pytest.mark.asyncio + async def test_command_skips_missing_tool_with_soft_echo(self): + # if/then/else structure emits "skipped (not installed)" only when + # `command -v` fails — distinct from the "prune failed" branch below. + task = BuildCacheCleanupTask() + runtime = _runtime(stdout="No cache entries to prune\npip: skipped (not installed)") + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + assert "skipped" in result["output_head"] + + @pytest.mark.asyncio + async def test_command_distinguishes_prune_failure_from_missing_tool(self): + # Regression: the old `a && b || c` form fired "skipped (not installed)" + # even when the tool was present but prune itself errored (permission, + # disk full, etc.). The if/then/else version must surface "prune failed" + # in that case so operators can tell the two apart in logs. + task = BuildCacheCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + cmd = runtime.execute.await_args.args[0].command + + # Both branches must be present in the generated command for both tools. + assert "uv: prune failed" in cmd + assert "uv: skipped (not installed)" in cmd + assert "pip: prune failed" in cmd + assert "pip: skipped (not installed)" in cmd + # The mutually-exclusive shape: if/then/else replaces `a && b || c`. + assert "if command -v uv" in cmd + assert "if command -v pip" in cmd + assert "else echo" in cmd + + @pytest.mark.asyncio + async def test_subset_of_tools(self): + task = BuildCacheCleanupTask(tools=["uv"]) + runtime = _runtime() + + await task.run_action(runtime) + cmd = runtime.execute.await_args.args[0].command + assert "uv cache prune" in cmd + assert "pip cache purge" not in cmd + + @pytest.mark.asyncio + async def test_empty_tools_list_runs_noop(self): + task = BuildCacheCleanupTask(tools=[]) + runtime = _runtime(stdout="no tools configured") + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + cmd = runtime.execute.await_args.args[0].command + assert "no tools configured" in cmd From 1e64ecbeca30dd023d153c8f1c0d5fc2ebde9d4d Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 18 May 2026 15:02:13 +0800 Subject: [PATCH 123/226] use triple-quoted multiline bash for build_cache snippets --- .../tasks/build_cache_cleanup_task.py | 28 +++++++++++-------- .../test_build_cache_cleanup_task.py | 4 ++- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/rock/admin/scheduler/tasks/build_cache_cleanup_task.py b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py index d7d2443e4a..38c4f8e490 100644 --- a/rock/admin/scheduler/tasks/build_cache_cleanup_task.py +++ b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py @@ -1,5 +1,7 @@ """Prune build/install caches (uv, pip) on each worker.""" +import textwrap + from rock.admin.proto.request import SandboxCommand as Command from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum from rock.common.constants import SCHEDULER_LOG_NAME @@ -14,15 +16,21 @@ # also fire `c` when `b` failed (permission/disk error), misleading operators # into thinking the tool was missing when in fact the prune itself errored. _TOOL_COMMANDS: dict[str, str] = { - "uv": ( - "if command -v uv >/dev/null 2>&1; then " - ' uv cache prune 2>&1 || echo "uv: prune failed"; ' - 'else echo "uv: skipped (not installed)"; fi' + "uv": textwrap.dedent( + """\ + if command -v uv >/dev/null 2>&1; then + uv cache prune 2>&1 || echo "uv: prune failed" + else + echo "uv: skipped (not installed)" + fi""" ), - "pip": ( - "if command -v pip >/dev/null 2>&1; then " - ' pip cache purge 2>&1 || echo "pip: prune failed"; ' - 'else echo "pip: skipped (not installed)"; fi' + "pip": textwrap.dedent( + """\ + if command -v pip >/dev/null 2>&1; then + pip cache purge 2>&1 || echo "pip: prune failed" + else + echo "pip: skipped (not installed)" + fi""" ), } @@ -54,9 +62,7 @@ def __init__( tools = tools if tools is not None else ["uv", "pip"] unknown = [t for t in tools if t not in _TOOL_COMMANDS] if unknown: - raise ValueError( - f"Unsupported build cache tool(s): {unknown}. Supported: {sorted(_TOOL_COMMANDS)}" - ) + raise ValueError(f"Unsupported build cache tool(s): {unknown}. Supported: {sorted(_TOOL_COMMANDS)}") self.tools = tools @classmethod diff --git a/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py b/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py index bb42576bfe..f3d8e36468 100644 --- a/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_build_cache_cleanup_task.py @@ -108,7 +108,9 @@ async def test_command_distinguishes_prune_failure_from_missing_tool(self): # The mutually-exclusive shape: if/then/else replaces `a && b || c`. assert "if command -v uv" in cmd assert "if command -v pip" in cmd - assert "else echo" in cmd + # Triple-quote bash puts `else` and `echo` on separate lines (newline as + # statement separator); just check the keyword is present. + assert "else" in cmd @pytest.mark.asyncio async def test_subset_of_tools(self): From 949eb5a5bb3a3f5e37beccc6cc3f1de8dc2a8a2d Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 19 May 2026 12:16:49 +0800 Subject: [PATCH 124/226] fix(docker): remove dead remove_images branch + add cls to remove_image Per review feedback on #965, scope down to two tightly-coupled fixes: 1. DockerUtil.remove_image was decorated @classmethod but missing the `cls` parameter, so the first positional arg got swallowed as `cls` and `docker rmi ` always failed. The bug stayed hidden because DockerDeployment._stop() catches CalledProcessError silently. Add the missing `cls` and a regression test that asserts the actual cmdline passed to subprocess.check_output. 2. DockerDeploymentConfig.remove_images is never exposed (default False) so the entire image-removal-on-stop branch in DockerDeployment._stop() is dead code. Image cleanup is the responsibility of ImageCleanupTask (docuum-driven background sweep), not the per-sandbox stop path. Delete the field and the dead branch. Also drop a stale `image_keep_patterns` reference from a SandboxLogConfig docstring (the field itself was already removed earlier). Co-Authored-By: Claude Opus 4.7 --- rock/config.py | 2 +- rock/deployments/config.py | 3 --- rock/deployments/docker.py | 7 ----- rock/utils/docker.py | 2 +- tests/unit/utils/test_docker_remove_image.py | 28 ++++++++++++++++++++ 5 files changed, 30 insertions(+), 12 deletions(-) create mode 100644 tests/unit/utils/test_docker_remove_image.py diff --git a/rock/config.py b/rock/config.py index d912219431..cec3899e7a 100644 --- a/rock/config.py +++ b/rock/config.py @@ -59,7 +59,7 @@ class SandboxLogConfig: Lives under SandboxConfig.log: the fields are domain knobs of "what to do with stopped sandbox logs" — when to archive, how many retries, what OSS key prefix to use — colocated with other sandbox lifecycle / cleanup - policy (image_keep_patterns, remove_container_enabled). OSS endpoint / + policy (remove_container_enabled). OSS endpoint / bucket / credentials still belong to OssConfig.primary. """ diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 248bffcb57..8f03a98b15 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -90,9 +90,6 @@ class DockerDeploymentConfig(DeploymentConfig): pull: Literal["never", "always", "missing"] = "missing" """Docker image pull policy: 'never', 'always', or 'missing'.""" - remove_images: bool = False - """Whether to remove the Docker image after the container stops.""" - python_standalone_dir: str | None = None """Directory path for Python standalone installation within the container.""" diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 99b0ade1f5..9c76a0ad82 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -664,13 +664,6 @@ def _stop(self): self._cleanup_log_dir_xfs_quota() self._container_name = None - if self._config and self._config.remove_images and DockerUtil.is_image_available(self._config.image): - logger.info(f"Removing image {self._config.image}") - try: - DockerUtil.remove_image(self._config.image) - except subprocess.CalledProcessError: - logger.error(f"Failed to remove image {self._config.image}", exc_info=True) - if self._check_stop_task is not None: logger.info("Stopping check task") self._check_stop_task.cancel() diff --git a/rock/utils/docker.py b/rock/utils/docker.py index 41e7f99981..f8a67fd87e 100644 --- a/rock/utils/docker.py +++ b/rock/utils/docker.py @@ -206,7 +206,7 @@ def logout(cls, registry: str, timeout: int = 30) -> str: raise @classmethod - def remove_image(image: str) -> bytes: + def remove_image(cls, image: str) -> bytes: """Remove a Docker image""" return subprocess.check_output(["docker", "rmi", image], timeout=30) diff --git a/tests/unit/utils/test_docker_remove_image.py b/tests/unit/utils/test_docker_remove_image.py new file mode 100644 index 0000000000..ac6d4e5ca7 --- /dev/null +++ b/tests/unit/utils/test_docker_remove_image.py @@ -0,0 +1,28 @@ +"""Regression tests for DockerUtil.remove_image cls bug fix.""" + +import subprocess +from unittest.mock import patch + +import pytest + +from rock.utils.docker import DockerUtil + + +class TestRemoveImageClsFix: + """Regression: remove_image previously had `@classmethod` without `cls`, + so the first positional arg was swallowed as `cls` and the actual image + name was lost.""" + + def test_remove_image_passes_image_to_subprocess(self): + with patch("subprocess.check_output", return_value=b"") as mock_run: + DockerUtil.remove_image("nginx:latest") + cmd = mock_run.call_args.args[0] + assert cmd == ["docker", "rmi", "nginx:latest"] + + def test_remove_image_propagates_error(self): + with patch( + "subprocess.check_output", + side_effect=subprocess.CalledProcessError(1, "docker"), + ): + with pytest.raises(subprocess.CalledProcessError): + DockerUtil.remove_image("nginx:latest") From 9611befa3b868d0ddc6a92f064d70a20c7c41d1a Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 19 May 2026 12:31:47 +0800 Subject: [PATCH 125/226] feature(scheduler): merge dangling/BuildKit prune into ImageCleanupTask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback on #970, fold the previously-proposed DockerImagePruneTask back into ImageCleanupTask rather than registering a parallel scheduler task. Single task, single schedule, single from_config entry — co-locating the two cleanups in the docuum task also reflects that they are complementary (docuum handles whole-image LRU; the prune step handles dangling layers + BuildKit cache that docuum cannot see) rather than independent. Behavior: - New `keep_build_storage` parameter (default "20GB"). After docuum is launched (still nohup &), runs `docker image prune --filter dangling=true` and `docker builder prune --keep-storage ` in one fail-soft synchronous shell pipeline. - Volume pruning is intentionally NOT invoked: long-running sandboxes may attach named volumes that we do not want dropped on schedule. - `image_whitelist` is intentionally NOT plumbed through to the prune step: dangling layers and BuildKit cache entries have no image:tag. - Set `keep_build_storage` to None / empty string to disable the prune step (preserves pre-#970 behavior). Tests live in tests/unit/admin/scheduler/test_image_cleanup_task.py. Co-Authored-By: Claude Opus 4.7 --- .../scheduler/tasks/image_cleanup_task.py | 60 ++++-- .../scheduler/test_image_cleanup_task.py | 171 ++++++++++++++++++ 2 files changed, 220 insertions(+), 11 deletions(-) create mode 100644 tests/unit/admin/scheduler/test_image_cleanup_task.py diff --git a/rock/admin/scheduler/tasks/image_cleanup_task.py b/rock/admin/scheduler/tasks/image_cleanup_task.py index 89c7f3599d..26b021d23b 100644 --- a/rock/admin/scheduler/tasks/image_cleanup_task.py +++ b/rock/admin/scheduler/tasks/image_cleanup_task.py @@ -11,21 +11,37 @@ class ImageCleanupTask(BaseTask): - """Docker image cleanup task using docuum.""" + """Docker image cleanup: docuum LRU + dangling/BuildKit prune. + + Two complementary cleanups in one task: + - docuum: long-running daemon, LRU eviction of whole image:tag entries. + Honors ``image_whitelist``. + - ``docker image prune --filter dangling=true`` + ``docker builder prune + --keep-storage ``: one-shot synchronous sweep of dangling layers + (``:``) and BuildKit cache, which docuum never touches. + Whitelist is intentionally NOT plumbed through here — dangling layers + and BuildKit cache entries have no image:tag, so a whitelist would be + misleading on both subcommands. + + Set ``keep_build_storage`` to a falsy value to disable the prune step. + """ def __init__( self, interval_seconds: int = 3600, disk_threshold: str = "1T", image_whitelist: list[str] | None = None, + keep_build_storage: str | None = "20GB", ): """ - Initialize image cleanup task. - Args: interval_seconds: Execution interval, default 1 hour - disk_threshold: Disk threshold to trigger cleanup, default 1T - image_whitelist: List of regex patterns for images to keep (matched against repository:tag) + disk_threshold: Disk threshold to trigger docuum cleanup, default 1T + image_whitelist: Regex patterns of images to keep (matched against + repository:tag). Applies to docuum only. + keep_build_storage: Lower bound for BuildKit cache retention, + passed to ``docker builder prune --keep-storage``. Default + "20GB". Set to None / empty to skip the prune step. """ super().__init__( type="image_cleanup", @@ -34,21 +50,21 @@ def __init__( ) self.disk_threshold = disk_threshold self.image_whitelist = image_whitelist or [] + self.keep_build_storage = keep_build_storage @classmethod def from_config(cls, task_config) -> "ImageCleanupTask": """Create task instance from config.""" - disk_threshold = task_config.params.get("disk_threshold", "1T") - image_whitelist = task_config.params.get("image_whitelist", []) return cls( interval_seconds=task_config.interval_seconds, - disk_threshold=disk_threshold, - image_whitelist=image_whitelist, + disk_threshold=task_config.params.get("disk_threshold", "1T"), + image_whitelist=task_config.params.get("image_whitelist", []), + keep_build_storage=task_config.params.get("keep_build_storage", "20GB"), ) async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: - """Run docuum image cleanup action.""" - # Check if docuum exists, install if not + """Start docuum daemon, then synchronously prune dangling/build cache.""" + # 1) docuum: LRU image eviction (long-running, nohup &) check_and_install_cmd = ( f"command -v docuum > /dev/null 2>&1 || curl {env_vars.ROCK_DOCUUM_INSTALL_URL} -LSfs | sh" ) @@ -67,9 +83,31 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: pid = extract_nohup_pid(result.stdout) logger.info(f"image cleanup task [{pid}] run successfully on worker[{runtime._config.host}]") + # 2) Dangling/BuildKit prune (sync, fail-soft so an old/missing docker + # subcommand on one worker doesn't abort the rest of the pipeline). + prune_exit = None + prune_output = "" + if self.keep_build_storage: + prune_steps = [ + "docker image prune -f --filter dangling=true", + f"docker builder prune -f --keep-storage {self.keep_build_storage}", + ] + prune_cmd = "; ".join(f"({s}) 2>&1 || true" for s in prune_steps) + prune_result = await runtime.execute(Command(command=prune_cmd, shell=True, check=False)) + prune_output = (prune_result.stdout or "").strip()[:1000] + prune_exit = prune_result.exit_code + logger.info( + f"docker prune done on worker[{runtime._config.host}]: " + f"keep_build_storage={self.keep_build_storage}, exit={prune_exit}, " + f"output_head={prune_output[:300]}" + ) + return { "pid": pid, "disk_threshold": self.disk_threshold, "image_whitelist": self.image_whitelist, + "keep_build_storage": self.keep_build_storage, + "prune_exit_code": prune_exit, + "prune_output_head": prune_output, "status": TaskStatusEnum.RUNNING, } diff --git a/tests/unit/admin/scheduler/test_image_cleanup_task.py b/tests/unit/admin/scheduler/test_image_cleanup_task.py new file mode 100644 index 0000000000..0121d1e5bc --- /dev/null +++ b/tests/unit/admin/scheduler/test_image_cleanup_task.py @@ -0,0 +1,171 @@ +"""Tests for ImageCleanupTask (docuum LRU + dangling/BuildKit prune).""" + +from unittest.mock import AsyncMock + +import pytest + +from rock.admin.scheduler.task_base import TaskStatusEnum +from rock.admin.scheduler.tasks.image_cleanup_task import ImageCleanupTask + + +class _FakeTaskConfig: + def __init__(self, params=None, interval_seconds=3600): + self.params = params or {} + self.interval_seconds = interval_seconds + + +class _FakeExecResult: + def __init__(self, exit_code=0, stdout=""): + self.exit_code = exit_code + self.stdout = stdout + + +def _runtime(side_effects): + """Build a runtime whose execute() returns successive results from side_effects.""" + rt = AsyncMock() + rt._config = type("C", (), {"host": "10.0.0.1"})() + rt.execute = AsyncMock(side_effect=side_effects) + return rt + + +# Fixed call sequence for run_action when keep_build_storage is set: +# 1) docuum install check +# 2) docuum start (returns PID-tagged stdout) +# 3) docker image prune + docker builder prune (one combined cmd) +def _default_results(pid=12345, prune_stdout="Total reclaimed space: 1.2GB"): + return [ + _FakeExecResult(), + _FakeExecResult(stdout=f"PIDSTART{pid}PIDEND"), + _FakeExecResult(stdout=prune_stdout), + ] + + +class TestInit: + def test_default(self): + task = ImageCleanupTask() + assert task.type == "image_cleanup" + assert task.interval_seconds == 3600 + assert task.disk_threshold == "1T" + assert task.image_whitelist == [] + assert task.keep_build_storage == "20GB" + + def test_custom(self): + task = ImageCleanupTask( + interval_seconds=600, + disk_threshold="500G", + image_whitelist=[r"^rock-base.*$"], + keep_build_storage="5GB", + ) + assert task.interval_seconds == 600 + assert task.disk_threshold == "500G" + assert task.image_whitelist == [r"^rock-base.*$"] + assert task.keep_build_storage == "5GB" + + def test_disable_prune(self): + task = ImageCleanupTask(keep_build_storage=None) + assert task.keep_build_storage is None + + +class TestFromConfig: + def test_from_config_defaults(self): + task = ImageCleanupTask.from_config(_FakeTaskConfig()) + assert task.disk_threshold == "1T" + assert task.image_whitelist == [] + assert task.keep_build_storage == "20GB" + + def test_from_config_custom(self): + cfg = _FakeTaskConfig( + params={ + "disk_threshold": "200G", + "image_whitelist": [r"^pinned:.*$"], + "keep_build_storage": "10GB", + }, + interval_seconds=7200, + ) + task = ImageCleanupTask.from_config(cfg) + assert task.interval_seconds == 7200 + assert task.disk_threshold == "200G" + assert task.image_whitelist == [r"^pinned:.*$"] + assert task.keep_build_storage == "10GB" + + +class TestRunAction: + @pytest.mark.asyncio + async def test_docuum_command_includes_threshold(self): + task = ImageCleanupTask(disk_threshold="500G") + runtime = _runtime(_default_results()) + + await task.run_action(runtime) + docuum_cmd = runtime.execute.await_args_list[1].args[0].command + assert "docuum --threshold 500G" in docuum_cmd + + @pytest.mark.asyncio + async def test_docuum_command_passes_whitelist(self): + task = ImageCleanupTask(image_whitelist=[r"^rock-base.*$", r"^pinned:.*$"]) + runtime = _runtime(_default_results()) + + await task.run_action(runtime) + docuum_cmd = runtime.execute.await_args_list[1].args[0].command + assert "--keep '^rock-base.*$'" in docuum_cmd + assert "--keep '^pinned:.*$'" in docuum_cmd + + @pytest.mark.asyncio + async def test_prune_command_includes_image_and_builder_prune(self): + task = ImageCleanupTask(keep_build_storage="10GB") + runtime = _runtime(_default_results()) + + await task.run_action(runtime) + prune_cmd = runtime.execute.await_args_list[2].args[0].command + assert "docker image prune -f --filter dangling=true" in prune_cmd + assert "docker builder prune -f --keep-storage 10GB" in prune_cmd + + @pytest.mark.asyncio + async def test_prune_command_does_not_invoke_volume_prune(self): + # Long-running sandboxes may attach named volumes; we don't want to + # drop them on schedule. + task = ImageCleanupTask() + runtime = _runtime(_default_results()) + + await task.run_action(runtime) + prune_cmd = runtime.execute.await_args_list[2].args[0].command + assert "docker volume prune" not in prune_cmd + + @pytest.mark.asyncio + async def test_prune_step_is_fail_soft(self): + task = ImageCleanupTask() + runtime = _runtime(_default_results()) + + await task.run_action(runtime) + prune_cmd = runtime.execute.await_args_list[2].args[0].command + # `(...) 2>&1 || true` makes a missing/old docker subcommand non-fatal. + assert "|| true" in prune_cmd + + @pytest.mark.asyncio + async def test_prune_skipped_when_keep_build_storage_falsy(self): + task = ImageCleanupTask(keep_build_storage=None) + # Only 2 execute calls expected: install check + docuum start. + runtime = _runtime(_default_results()[:2]) + + result = await task.run_action(runtime) + assert runtime.execute.await_count == 2 + assert result["prune_exit_code"] is None + assert result["prune_output_head"] == "" + + @pytest.mark.asyncio + async def test_run_action_returns_pid_and_status(self): + task = ImageCleanupTask() + runtime = _runtime(_default_results(pid=98765)) + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.RUNNING + assert result["pid"] == 98765 + assert result["disk_threshold"] == "1T" + assert result["keep_build_storage"] == "20GB" + + @pytest.mark.asyncio + async def test_run_action_records_prune_output(self): + task = ImageCleanupTask() + runtime = _runtime(_default_results(prune_stdout="Total reclaimed space: 3.5GB\n(more)")) + + result = await task.run_action(runtime) + assert "Total reclaimed space" in result["prune_output_head"] From faa5fec73841d6944a776111aae935150ad1e555 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 21 May 2026 11:16:19 +0800 Subject: [PATCH 126/226] docs: add v1.8.0 release version (CN + EN) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Copy version-1.7.x → version-1.8.x for both English (docs/versioned_docs/) and Chinese (docs/i18n/zh-Hans/...) trees. - Wipe old Release Notes (major.minor bump 1.7 → 1.8), keep a fresh index.md listing only v1.8.0. - Create v1.8.0 release note in both languages with full content (SDK / Sandbox / Deployments / Proxy / Metrics / Scheduler / Rocklet / etc). - Copy version-1.7.x-sidebars.json for the new version. - Copy + update version-1.7.x.json (zh-Hans translation) with version.label message and description bumped to 1.8.x. - Insert "1.8.x" at the start of docs/versions.json. lastVersion in docusaurus.config.js is intentionally NOT changed yet — default-display will stay at 1.7.x until release content is finalized. Co-Authored-By: Claude Opus 4.7 --- .../version-1.8.x.json | 34 ++ .../Getting Started/installation.md | 141 ++++++++ .../Getting Started/quickstart.md | 172 ++++++++++ .../Getting Started/rock-agent.md | 76 +++++ .../version-1.8.x/Getting Started/rockroll.md | 200 +++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../References/Python SDK References/job.md | 150 +++++++++ .../Python SDK References/model-service.md | 298 +++++++++++++++++ .../Python SDK References/python_sdk.md | 265 +++++++++++++++ .../Python SDK References/remote_user.md | 69 ++++ .../Python SDK References/rock-agent.md | 310 ++++++++++++++++++ .../Python SDK References/runtime-env.md | 137 ++++++++ .../Python SDK References/sandbox.md | 113 +++++++ .../swe-bench-evaluation.md | 228 +++++++++++++ .../version-1.8.x/References/api.md | 194 +++++++++++ .../version-1.8.x/Release Notes/index.md | 6 + .../version-1.8.x/Release Notes/v1.8.0.md | 154 +++++++++ .../User Guides/configuration.md | 188 +++++++++++ .../version-1.8.x/overview.md | 40 +++ .../Getting Started/installation.md | 143 ++++++++ .../Getting Started/quickstart.md | 166 ++++++++++ .../Getting Started/rock-agent.md | 76 +++++ .../version-1.8.x/Getting Started/rockroll.md | 194 +++++++++++ .../References/Python SDK References/codes.md | 93 ++++++ .../Python SDK References/deploy.md | 68 ++++ .../Python SDK References/file_system.md | 94 ++++++ .../References/Python SDK References/job.md | 150 +++++++++ .../Python SDK References/model-service.md | 298 +++++++++++++++++ .../Python SDK References/python_sdk.md | 265 +++++++++++++++ .../Python SDK References/remote_user.md | 70 ++++ .../Python SDK References/rock-agent.md | 310 ++++++++++++++++++ .../Python SDK References/runtime-env.md | 136 ++++++++ .../Python SDK References/sandbox.md | 114 +++++++ .../swe-bench-evaluation.md | 229 +++++++++++++ .../version-1.8.x/References/api.md | 195 +++++++++++ .../version-1.8.x/Release Notes/index.md | 6 + .../version-1.8.x/Release Notes/v1.8.0.md | 138 ++++++++ .../User Guides/configuration.md | 189 +++++++++++ docs/versioned_docs/version-1.8.x/overview.md | 33 ++ .../version-1.8.x-sidebars.json | 64 ++++ docs/versions.json | 1 + 43 files changed, 6062 insertions(+) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x.json create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/installation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/quickstart.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rockroll.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/codes.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/deploy.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/file_system.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/job.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/model-service.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/python_sdk.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/remote_user.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/rock-agent.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/runtime-env.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/sandbox.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/api.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/index.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/configuration.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/overview.md create mode 100644 docs/versioned_docs/version-1.8.x/Getting Started/installation.md create mode 100644 docs/versioned_docs/version-1.8.x/Getting Started/quickstart.md create mode 100644 docs/versioned_docs/version-1.8.x/Getting Started/rock-agent.md create mode 100644 docs/versioned_docs/version-1.8.x/Getting Started/rockroll.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/codes.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/deploy.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/file_system.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/job.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/model-service.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/python_sdk.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/remote_user.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/rock-agent.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/runtime-env.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/sandbox.md create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md create mode 100644 docs/versioned_docs/version-1.8.x/References/api.md create mode 100644 docs/versioned_docs/version-1.8.x/Release Notes/index.md create mode 100644 docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md create mode 100644 docs/versioned_docs/version-1.8.x/User Guides/configuration.md create mode 100644 docs/versioned_docs/version-1.8.x/overview.md create mode 100644 docs/versioned_sidebars/version-1.8.x-sidebars.json diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x.json b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x.json new file mode 100644 index 0000000000..35a87b2824 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x.json @@ -0,0 +1,34 @@ +{ + "version.label": { + "message": "1.8.x", + "description": "The label for version 1.8.x" + }, + "sidebar.tutorialSidebar.category.Getting Started": { + "message": "快速上手", + "description": "The label for category 'Getting Started' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.User Guides": { + "message": "用户指南", + "description": "The label for category 'User Guides' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.References": { + "message": "参考", + "description": "The label for category 'References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Release Notes": { + "message": "版本说明", + "description": "The label for category 'Release Notes' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.model-service": { + "message": "Model Service 参考", + "description": "The label for category 'Model Service References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.sandbox-agent": { + "message": "Sandbox Agent参考", + "description": "The label for category 'Sandbox Agent References' in sidebar 'tutorialSidebar'" + }, + "sidebar.tutorialSidebar.category.Python SDK References": { + "message": "Python SDK 参考", + "description": "The label for category 'Python SDK References' in sidebar 'tutorialSidebar'" + } +} diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/installation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/installation.md new file mode 100644 index 0000000000..0ab70e55d1 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/installation.md @@ -0,0 +1,141 @@ +--- +sidebar_position: 3 +--- + +# 安装指南 + +本文档介绍如何使用 `uv` 和 `pip` 安装和设置 ROCK 开发环境。该项目是一个强化学习开放构建工具包，支持多种组件。 + +## 使用 uv（推荐） + +### 快速安装所有依赖 + +```bash +# 安装所有依赖（包括可选依赖） +uv sync --all-extras + +# 安装开发/测试依赖 +uv sync --all-extras --all-groups +``` + +### 安装不同依赖组 + +#### 仅核心依赖 +```bash +uv sync +``` + +#### 管理组件依赖 +```bash +uv sync --extra admin +``` + +#### Rocklet 执行环境依赖 +```bash +uv sync --extra rocklet +``` + +#### 所有依赖 +```bash +uv sync --all-extras +``` + +#### 开发/测试依赖 +```bash +uv sync --all-extras --group test +``` + +## 使用 pip + +### 从 pip 源安装 + +#### 仅核心依赖 +```bash +pip install rl-rock +``` + +#### 管理组件依赖 +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install "rl-rock[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install "rl-rock[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install "rl-rock[all]" +``` + +### 使用 pip 从源码安装 + +#### 仅核心依赖 +```bash +pip install . +``` + +#### 管理组件依赖 +```bash +pip install ".[admin]" +``` + +#### Rocklet 执行环境依赖 +```bash +pip install ".[rocklet]" +``` + +#### 构建器依赖 +```bash +pip install ".[builder]" +``` + +#### 安装所有可选依赖 +```bash +pip install ".[all]" +``` + +## 可用入口点 + +该包提供以下命令行脚本： + +- `rocklet`: ROCK 执行环境服务器 (rock.rocklet.server:main) +- `admin`: 管理服务器 (rock.admin.main:main) +- `envhub`: 环境中心服务器 (rock.envhub.server:main) +- `rock`: 主 ROCK 命令行接口 (rock.cli.main:main) + +## 开发设置 + +### 使用 uv（推荐） + +```bash +# 克隆并设置开发环境 +git clone +cd ROCK +uv sync --all-extras --group test + +# 运行测试 +uv run pytest + + +### 使用 pip + +```bash +# 开发模式安装所有可选依赖 +pip install -e ".[all]" + +# 分别安装 +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" +``` + +## 附加说明 + +- 项目配置为默认使用阿里云 PyPI 镜像: `https://mirrors.aliyun.com/pypi/simple/` +- 对于本地开发，运行测试需要 `test` 依赖组 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/quickstart.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..e0a0891c0e --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/quickstart.md @@ -0,0 +1,172 @@ +--- +sidebar_position: 2 +--- + +# 快速上手 + +本指南将通过完整的示例演示如何使用 ROCK 创建和管理强化学习环境。ROCK (Reinforcement Open Construction Kit) 是一个全面的沙箱环境管理框架，主要用于强化学习和AI开发环境。 + +## 1. 环境准备 + +我们推荐在 Linux 系统下启动 ROCK，能够尽量复用项目依赖，提升环境拉起速度。如果需要在 macOS 上尝试，可以参考 [MacOS 启动](#7-macos-启动) 一节。 + +在开始之前，请确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且示例中依赖python:3.11镜像 +docker pull python:3.11 + +# 验证 uv 安装 +uv --version + + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone +cd ROCK + +# 创建虚拟环境（使用 uv 托管的 Python, 以python 3.11 版本为例） +uv venv --python 3.11 --python-preference only-managed + +# 安装所有依赖组 +uv sync --all-extras +``` + +> **重要提示**: 为确保 ROCK 能正确挂载项目和虚拟环境及其依赖的 base Python 解释器，强烈推荐使用 uv 托管的 Python 环境而非系统 Python。 + +## 2. 激活虚拟环境 + +在运行任何 ROCK 命令之前，需要先激活虚拟环境。确保 sys.base_prefix 是 uv 管理的环境，类似于 `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` 等路径。 + +```bash +# 激活虚拟环境 +source .venv/bin/activate + +# 验证 Python 环境 +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **验证要点**: 确保输出的 base prefix 路径指向 uv 管理的 Python 环境，而非系统 Python。 + +## 3. 验证环境配置 + +激活虚拟环境后，验证依赖安装是否正确： + +```bash +# 检查关键依赖 +python -c "import rock; print(\"Hello ROCK\")" +``` + + +## 4. 启动 ROCK 服务 + +激活虚拟环境后，在项目根目录下，启动 ROCK Admin 服务： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 启动 ROCK Admin 服务（本地环境） +rock admin start +``` + +服务启动后，您将看到类似以下的输出： + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **服务说明**: ROCK Admin 服务默认运行在 `http://127.0.0.1:8080`。 + +## 5. 运行示例环境 + +现在可以运行示例环境来验证安装。确保 ROCK 服务正在运行，然后打开一个新的终端窗口执行以下命令： + +```bash +# 确保虚拟环境已激活 +source .venv/bin/activate + +# 运行沙箱示例 +python examples/sandbox_demo.py + +# 运行 GEM 协议示例 +python examples/sokoban_demo.py +``` + +### 5.1 示例说明 + +- **sandbox_demo.py**: 演示如何使用 ROCK 的沙箱 SDK 创建和管理容器化环境 +- **sokoban_demo.py**: 演示如何使用 ROCK 的 GEM 协议兼容接口创建强化学习环境 + +> **运行要求**: 确保 ROCK Admin 服务正在运行，因为示例需要与服务进行通信。 + +## 6. 分布式环境配置（可选） + +对于分布式多机器环境，请确保以下配置一致： + +1. 所有机器上 ROCK 和 uv 的 Python 配置使用相同的根 Python 解释器 +2. Docker 版本在所有节点上保持一致 +3. 网络配置允许各节点间正常通信 + + + +## 7. MacOS 启动 + +在 macOS 上，如果需要启动 Linux 镜像的环境，需要先设置环境变量： + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +在容器启动时，会安装对应的 uv 环境，细节可以参考 `rock/rocklet/local_files/docker_run_with_uv.sh` 脚本。 + +> **注意**: 相比 Linux 系统，macOS 上的启动速度会较慢，且比较依赖网络环境，可以根据实际情况调整脚本。ROCK_WORKER_ENV_TYPE的细节可以参考 [Configuration Guide](../User%20Guides/configuration.md). + + +## 8. 从Pip源启动 + +如果从Pip源启动Admin Server，在参照[安装指南](./installation.md)安装完成ROCK后, 需要设置额外环境变量: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(这一启动方式在容器环境启动时会从Pypi源上拉取最新的rocklet并安装, 相对启动速度比较慢, 仅推荐测试使用, 生产上依旧推荐其他的启动方式) + + +## 总结 + +恭喜！您已经成功完成了 ROCK 的快速开始指南。现在您应该能够： + +- 正确设置 ROCK 开发环境 +- 使用 uv 管理的 Python 环境 +- 启动和管理 ROCK 服务 +- 运行示例程序验证安装 +- 在分布式环境中配置 ROCK（如果需要） + +如需深入了解 ROCK 的更多功能，请参考以下文档： + +## 下一步学习 + +- [配置指南](../User%20Guides/configuration.md) - 详细了解 ROCK 的配置选项 +- [API 文档](../References/api.md) - 查看完整的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +- [安装指南](./installation.md) - 详细了解 ROCK 安装和配置 +- [概览](../overview.md) - 了解 ROCK 的设计理念 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..1898fbdcd0 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rock-agent.md @@ -0,0 +1,76 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent 快速启动 + +ROCK 提供两种并列的 agent 使用能力,各自适用不同场景: + +- **Job**:通过 BashJob / HarborJob 在 sandbox 里跑一次 agent 评测/任务(典型基准:SWE-bench、Terminal Bench),是入门主要场景。 +- **install-agent**:直接在单个沙箱里安装并运行 agent,适合本地开发、单次调试。 + +下面优先介绍 Job 用法,install-agent 用法见末尾或 [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md)。 + +## 前置条件 + +- 确保有可用的 ROCK 服务,如果需要本地拉起服务端,参考[快速启动](quickstart.md) + +--- + +## 一、用 Job 运行 Agent + +Job 有两种 backend:**Harbor Job** 用于运行 AI agent 基准评测任务(SWE-bench、Terminal Bench 等);**Bash Job** 用于在沙箱里跑自定义 shell 脚本。 + +### 1.1 准备 yaml + +挑一类作为起点,直接复制对应模板: + +- Harbor Job(Terminal Bench):[`examples/job/harbor/tb_job_config.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor/tb_job_config.yaml.template) +- Bash Job(claw-eval):[`examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template) + +按模板填好对应字段即可。两类 Job 的完整字段说明见 [Use Job to Run Agent](../References/Python%20SDK%20References/job.md)。 + +### 1.2 通过 Python SDK 启动 + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +BashJob 的用法、完整字段说明、结果处理详见 [Use Job to Run Agent](../References/Python%20SDK%20References/job.md)。 + +--- + +## 二、install-agent:在沙箱里安装并运行 Agent + +适合本地开发或单次调试 agent 的场景,核心 API: + +```python +await sandbox.agent.install(config="rock_agent_config.yaml") +result = await sandbox.agent.run(prompt="hello") +``` + +`examples/install-agents/` 下提供了多个开箱即用的示例: + +- `examples/install-agents/iflow_cli/` — IFlowCli +- `examples/install-agents/claude_code/` — Claude Code +- `examples/install-agents/cursor_cli/`、`qwen_code/`、`swe_agent/`、`openclaw/` — 其他 + +运行 Claude Code 示例: + +```bash +cd examples/install-agents/claude_code +python claude_code_demo.py +``` + +完整 RockAgentConfig 字段说明、占位符语义、API 参考详见 [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rockroll.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..3b53810ba3 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Getting Started/rockroll.md @@ -0,0 +1,200 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL 快速开始指南 + +本指南将引导您使用 ROLL (训练框架) 和 ROCK (环境管理) 来运行一个基于 Sokoban 游戏（推箱子）的强化学习训练示例。 + +## 1. 单机环境准备 + +在开始之前，请先确保您的系统已安装以下依赖项： + +### 1.1 系统要求 + +- **操作系统**: 推荐使用 Linux (如 Ubuntu 20.04+) +- **硬件**: 建议使用 NVIDIA GPU 并安装对应的驱动程序 +- **Docker**: ROCK 使用 Docker 进行容器化环境管理 +- **uv**: ROCK 使用 uv 进行依赖管理和虚拟环境创建 + +### 1.2 验证依赖安装 + +```bash +# 验证 Docker 安装 +docker --version + +# 验证 Docker 可用, 且可提前拉取 Sokoban 游戏环境镜像，避免训练时等待 +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# 验证 uv 安装 +uv --version + +``` + +### 1.3 项目初始化 + +```bash +# 克隆项目仓库 +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# 确保两个仓库位于同一级目录下，如下所示： +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. 启动训练流程 + +> 说明：下文均以 *torch2.6.0 + vLLM0.8.4* 为例。 + + +### 方式一： 使用虚拟环境启动（推荐） + +#### 为什么推荐这种方式？ +- 隔离性：uv 虚拟环境能确保项目依赖与系统环境隔离，避免冲突。 +- 速度快：ROCK 可以复用此虚拟环境，大大加快了后续环境的启动速度。 +- 稳定性：依赖关系更清晰，环境更易复现。 + + +```bash +# 进入 ROCK 目录 +cd ROCK + +# 使用 uv 创建并激活 Python 3.10 虚拟环境（ROLL推荐使用Python 3.10） +uv venv --python 3.10 --python-preference only-managed + +# 激活虚拟环境 +source .venv/bin/activate + +# 使用uv安装ROCK的依赖 +uv sync --all-extras + +# 若使用Python 3.10, 启动 ray 时会报错：ValueError: is not a valid Sentinel +# 原因是 ray 与 click>=8.3 版本不兼容，需要降级到 click<8.3 +# Python 3.11 不会有这个问题 +uv pip install 'click>=8.2,<8.3' + +# 切换到 ROLL 目录以安装其依赖 +cd ../ROLL + +# 设置国内 PyPI 镜像源以加速下载 +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装核心 PyTorch 组件 +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -i $PYPI_MIRROR + +# 安装transformer-engine，--no-build-isolation 避免因环境隔离导致找不到 torch +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation -i $PYPI_MIRROR + +# 安装预编译的 flash-attention，以匹配特定的 CUDA 和 PyTorch 版本 +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# 安装其余依赖 +uv pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# (可选) 安装Tensorboard，用于查看训练指标 +uv pip install tensorboard -i $PYPI_MIRROR + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### 方式二：使用系统环境启动（备选方案） + +为获得最佳兼容性，推荐使用 ROLL 官方提供的基础 Docker 镜像，因为它们已经预装了匹配的 CUDA、cuDNN 和其他基础库。 + +> [ROLL 官方镜像列表](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Installation/image_address/) + + +#### 注意 +此方式会将所有 Python 包直接安装到您的当前环境（例如，容器的基础环境）中，可能会与系统自带的包或其他项目产生冲突。 + +由于 ROCK 无法复用环境，每次启动任务时都可能需要重新安装部分依赖，启动速度较慢且受网络影响。 + + +```bash +PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/" + +# 安装ROCK的依赖 +cd ROCK +pip install . -i $PYPI_MIRROR +pip install ".[admin]" -i $PYPI_MIRROR + +# 安装ROLL的依赖 +cd ../ROLL +pip install -r requirements_torch260_vllm.txt -i $PYPI_MIRROR + +# 配置ROCK用uv启动的环境变量 +export ROCK_WORKER_ENV_TYPE=uv + +# 启动ROLL脚本（包含ROCK服务的启动） +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +至此，您已成功启动了 Sokoban 强化学习训练流程。祝您 Rock & Roll 愉快！ + + +## 3. 多机部署 + +除了在单机上运行，您也可以将 **ROCK 服务** 和 **ROLL 训练** 部署在不同的机器上，通过网络进行通信。这是一种常见的服务化部署模式。 + +### 3.1 在机器 A 上部署 ROCK 服务 + +在一台独立的机器（或容器）上，参照[ROCK快速指南](./quickstart.md)部署并启动 ROCK 服务。 + +> **重要提示** +> 启动服务后，请记下ROCK服务的IP地址和端口，例如`http://192.168.1.10:8000`，后续步骤将需要这个地址。 + +### 3.2 在机器 B 上准备 ROLL 客户端 + +在另一台将要运行训练任务的机器上，执行以下操作。 + +1. 验证网络连通性 + +首先，使用 curl 命令检查是否能从机器 B 访问到机器 A 上的 ROCK 服务。 +```bash +# 将 : 替换为您的 ROCK 服务实际地址 +# 如果成功，会收到 ROCK 服务的响应 {"message":"hello, ROCK!"} +curl http://: +``` + +2. 准备 ROLL 环境 + +```bash +# 克隆 ROLL 仓库 +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# 安装依赖 +pip install -r requirements_torch260_vllm.txt -i https://mirrors.aliyun.com/pypi/simple/ +``` + +3. 配置 ROLL 连接地址 + +修改 ROLL 的配置文件，使其能够找到并连接到远程的 ROCK 服务。 +- 打开配置文件：examples/agentic_demo/agentic_val_sokoban_sandbox.yaml +- 找到 SokobanSandbox 下的 env_config 部分 +- 将 base_url 的值修改为您的 ROCK 服务地址 +```yaml +custom_envs: + SokobanSandbox: + env_config: + # 将这里的地址修改为您的 ROCK 服务地址 + # 例如: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. 启动训练 +配置完成后，即可在机器 B 上启动 ROLL 训练脚本。 + +```bash +# 此脚本现在会通过网络请求机器 A 上的 ROCK 服务来创建环境 +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### 进阶：分布式 ROLL 训练 + +如果您希望将 ROLL 训练任务本身进行分布式部署，可以参考 ROLL 的官方分布式部署文档。 +> [快速上手：多节点部署指南](https://alibaba.github.io/ROLL/zh-Hans/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/codes.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..47b74166de --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +错误码定义和分类，用于错误处理和重试策略。 + +## 使用示例 + +```python +import rock + +def test_codes_values(): + """测试基本状态码值""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes 分类 + +```python +OK = 2000, "OK" +""" +成功状态码 (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +客户端错误码 (4xxx): + +这些错误表示客户端请求有问题， +SDK 会抛出异常。 +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +服务端错误码 (5xxx): + +这些错误表示服务端出现问题， +SDK 会抛出异常。 +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +命令/执行错误码 (6xxx): + +这些错误与命令执行相关，由模型处理， +SDK 不会抛出异常。 +""" +``` + +## 重试策略建议 + +- **重试触发条件**: 只有当 `INTERNAL_SERVER_ERROR` 时才需要重试 +- **其他情况的处理策略**: + - `BAD_REQUEST`: 需要检查 arun 调用逻辑是否有异常 + - `COMMAND_ERROR`: stdout 输出到 `observation.output`，stderr 输出到 `observation.failure_reason` +- `COMMAND_ERROR` 说明: 由于 bash 执行失败时，stdout/stderr 可能全部非空，建议将 observation 中 output 和 failure_reason 全部 prompt 给模型进行推理 + +## 重试示例 + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/deploy.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..b7bd2da08f --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +沙箱资源部署管理器，用于本地目录部署和模板格式化。 + +## deploy_working_dir - 部署本地目录 + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# 部署本地目录到沙箱（自动生成目标路径） +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"部署到: {target}") # 例如: /tmp/rock_workdir_abc123 + +# 部署到指定目标路径 +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - 模板变量替换 + +`format` 方法支持两种模板语法： + +- **`${variable}`** - 标准 Python 字符串模板语法 +- **`<>`** - 替代语法（内部转换为 `${variable}`） + +```python +# 使用 ${working_dir} 模板变量 +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# 结果: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# 使用 <<>> 替代语法 +cmd = deploy.format("cat <>/file.txt") +# 结果: cat /tmp/rock_workdir_abc123/file.txt + +# 结合自定义变量使用 +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# 结果: cat /tmp/rock_workdir_abc123/settings.json + +# Shell 语法保持不变 +cmd = deploy.format("echo $((3 << 2 >> 1))") +# 结果: echo $((3 << 2 >> 1)) + +# 直接访问 working_dir +if deploy.working_dir: + print(f"当前工作目录: {deploy.working_dir}") +``` + +## 多次部署 + +后续调用会覆盖之前的工作目录路径： + +```python +# 第一次部署 +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# 第二次部署（覆盖之前的路径） +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/file_system.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..741b14f5e4 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +文件系统操作接口，提供沙箱环境中的权限管理和目录上传功能。 + +## chown - 修改所有者 + +```python +from rock.actions.sandbox.request import ChownRequest + +# 创建远程用户后修改所有者 +await sandbox.remote_user.create_remote_user("deploy") + +# 获取当前目录 +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# 修改目录所有者 +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# 递归修改目录及其内容所有者 +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - 修改权限 + +```python +from rock.actions.sandbox.request import ChmodRequest + +# 创建测试目录 +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# 修改目录权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# 递归修改权限（包括子目录和文件） +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# 设置最高权限 +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - 上传目录 + +```python +import os +from pathlib import Path + +# 准备本地目录 +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# 上传到沙箱 +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"上传成功: {result.output}") +else: + print(f"上传失败: {result.failure_reason}") +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/job.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/job.md new file mode 100644 index 0000000000..6dee1f678b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/job.md @@ -0,0 +1,150 @@ +# Use Job to Run Agent + +> 这是 ROCK 两种并列的 agent 使用能力中 **Job** 的参考文档,核心 API 是 `rock.sdk.job.Job` 与 `JobConfig`,用于在沙箱里跑一次 agent 评测/任务。有两种 backend:**Bash Job** 与 **Harbor Bench Job**。 +> +> 另一种能力是在单个沙箱里安装并运行 agent,见 [Install Agent in Sandbox](./rock-agent.md)。两种能力使用各自独立的配置 schema,**不要互相套用**。 + +`rock.sdk.job` 通过同一套 `Job` API 支持两种模式,通过配置类型区分: + +- **Bash Job**:在沙箱中运行自定义 Shell 脚本,适合数据处理、外部评测工具等 +- **Harbor Bench Job**:通过 Harbor 框架运行 AI agent 基准评测任务(SWE-bench、Terminal Bench 等) + +## 端到端示例 + +最小可跑通的 Python 用法: + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") # 含 agents: 与 datasets: + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +完整 yaml 模板参考 `examples/job/harbor/swe_job_config.yaml.template`。 + +--- + +## Bash Job + +Bash Job 适用于在沙箱内执行任意 Shell 脚本的场景,例如运行评测工具、数据处理流程等。 + +完整示例参考:[`examples/job/bash/claw_eval/`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval) + +- `run_claw_eval.py` — 主入口,演示 `JobConfig.from_yaml()` + `Job(config).run()` +- `claw_eval_bashjob.yaml.template` — YAML 配置模板,含 `script_path`、`environment`、`uploads`、`env` 等字段 +- `run_claw_eval.sh` — 沙箱内实际执行的脚本,演示 DinD 启动、日志写入和评分输出 + +### BashJobConfig 配置字段 + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `script` | `str \| None` | `None` | 内联脚本内容,与 `script_path` 二选一 | +| `script_path` | `str \| None` | `None` | 本地脚本文件路径,运行时读取并上传执行 | +| `job_name` | `str` | 当前时间戳 | 任务名称,用于日志和产物路径区分 | +| `environment` | `EnvironmentConfig` | — | 沙箱连接及资源配置,详见下表 | +| `namespace` | `str \| None` | `None` | 命名空间 | +| `experiment_id` | `str \| None` | `None` | 实验 ID | +| `timeout` | `int` | `7200` | 整体超时秒数(2 小时) | + +**`environment` 常用字段:** + +| 字段 | 类型 | 说明 | +|------|------|------| +| `image` | `str` | 沙箱 Docker 镜像 | +| `base_url` | `str` | ROCK 平台地址 | +| `xrl_authorization` | `str` | 鉴权 Token | +| `cluster` | `str` | 目标集群 | +| `memory` | `str` | 内存大小(如 `"64g"`) | +| `cpus` | `int` | CPU 核数 | +| `auto_stop` | `bool` | 任务完成后是否自动停止沙箱 | +| `uploads` | `list` | 本地文件/目录上传列表,格式:`[本地路径, 沙箱目标路径]` | +| `env` | `dict[str, str]` | 注入沙箱会话的环境变量 | + +--- + +## Harbor Bench Job + +Harbor Bench Job 适用于通过 Harbor 框架运行 AI agent 基准评测任务,如 SWE-bench、Terminal Bench 等。 + +> **注意**:`rock.sdk.bench.Job` 已废弃,将在未来移除。请改用 `rock.sdk.job.Job` + `HarborJobConfig`。 + +完整示例参考:[`examples/job/harbor/`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor) + +- `harbor_demo.py` — 主入口,演示 `JobConfig.from_yaml()` + `Job(config).run()` + 结果遍历 +- `swe_job_config.yaml.template` — SWE-bench 任务配置模板 +- `swe_job_config-verifier.yaml.template` — 附带 `verifier.mode: native` 的变体 +- `tb_job_config.yaml.template` — Terminal Bench 任务配置模板 + +### HarborJobConfig 核心配置字段 + +**基础字段:** + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `experiment_id` | `str` | 必填 | 实验 ID,Harbor 中必须提供 | +| `job_name` | `str \| None` | 自动生成 | 格式:`{dataset}_{task}_{uuid[:8]}` | +| `namespace` | `str \| None` | `None` | 命名空间,从沙箱自动反填 | +| `environment` | `RockEnvironmentConfig` | — | 沙箱连接及资源配置 | + +**执行控制字段:** + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `n_attempts` | `int` | `1` | 每个 Trial 的尝试次数 | +| `timeout` | `int` | `7200` | 整体超时秒数(自动从 agent_timeout 推算) | +| `debug` | `bool` | `False` | 调试模式,保留更多中间产物 | + +**组件字段:** + +| 字段 | 类型 | 说明 | +|------|------|------| +| `agents` | `list[AgentConfig]` | Harbor 框架自身的 agent 配置(典型字段:`name`、`model_name`),完整字段见 `examples/job/harbor/swe_job_config.yaml.template` | +| `datasets` | `list[DatasetConfig]` | 数据集配置列表 | +| `verifier` | `VerifierConfig` | Verifier 评测配置 | +| `orchestrator` | `OrchestratorConfig` | 并发调度配置 | + +--- + +## 结果处理 + +两种 Job 模式均返回 `JobResult`: + +```python +result = await Job(config).run() + +print(f"status={result.status}, score={result.score}") +for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + if trial.exception_info: + print(f" {trial.exception_info.exception_type}: {trial.exception_info.exception_message}") +``` + +### JobResult 字段 + +| 字段 / 属性 | 类型 | 说明 | +|------------|------|------| +| `status` | `JobStatus` | 任务整体状态 | +| `trial_results` | `list[TrialResult]` | 所有 Trial 结果列表 | +| `score` | `float`(属性) | 所有 Trial `score` 的平均值 | +| `n_completed` | `int`(属性) | 状态为 `completed` 的 Trial 数 | +| `n_failed` | `int`(属性) | 状态为 `failed` 的 Trial 数 | + +### TrialResult 字段 + +| 字段 / 属性 | 类型 | 说明 | +|------------|------|------| +| `task_name` | `str` | 任务名称 | +| `exit_code` | `int` | 进程退出码 | +| `raw_output` | `str` | 进程原始输出 | +| `exception_info` | `ExceptionInfo \| None` | 若有异常则填充 | +| `status` | `str`(属性) | `"completed"` 或 `"failed"` | +| `duration_sec` | `float`(属性) | 执行耗时(秒) | +| `score` | `float`(属性) | 评分(Bash Job 默认 `0.0`,Harbor 模式来自 verifier) | diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/model-service.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..ba158cf75a --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service（实验性） + +ROCK 提供的 Model Service 负责处理 AI 模型调用的通信，为代理(Agent)和训练框架(如 Roll)或实际的 LLM 推理服务之间提供通信桥梁。 + +## 与 RockAgent 集成 + +ModelService 通常由 **RockAgent** 自动管理，无需手动调用生命周期方法。只需在配置中启用即可： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # 启用 ModelService，RockAgent 会自动管理其生命周期 +) +``` + +RockAgent 会自动： +- 安装 ModelService（安装 Python 运行时环境、安装模型服务包） +- 启动/停止 ModelService +- 监控 Agent 进程 + +## 架构概述（Local 模式） + +Local 模式下，模型服务使用**文件系统**作为通信媒介，实现代理和模型间的请求-响应机制。 + +当 Agent 需要调用模型时，请求首先写入日志文件，然后由负责监听的组件处理响应。当模型生成响应后，结果将写回日志文件，并由等待的 Agent 读取。 + +## anti_call_llm - 核心 API + +`anti_call_llm()` 是 **Local 模式**下最重要的 API，用于手动触发 LLM 反调用，实现模型调用的精细控制： + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='OpenAI type response', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +**使用场景：** +- Agent 捕获到 LLM 响应后，调用此方法通知 Roll 运行时 +- 支持携带响应数据，用于错误处理或重试 +- 超时和检查间隔可配置，适应不同网络环境 + +## CLI 命令 + +如果需要通过 CLI 使用模型服务，ROCK 提供了一个 CLI 命令集，可以在沙箱中安装 ROCK 后，通过 `rock model-service` 访问： + +### start 命令 +开始模型服务进程 +```bash +rock model-service start --type [local|proxy] [选项] +``` + +参数: + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `--type` | str | `local` | 服务类型：`local` 或 `proxy` | +| `--config-file` | str | None | 配置文件路径 | +| `--host` | str | None | 服务器地址（覆盖配置） | +| `--port` | int | None | 服务器端口（覆盖配置） | +| `--proxy-base-url` | str | None | 代理基础 URL | +| `--retryable-status-codes` | str | None | 可重试状态码，逗号分隔 | +| `--request-timeout` | int | None | 请求超时秒数 | + +### watch-agent 命令 +监控代理进程，当进程退出时发送 SESSION_END 消息 +```bash +rock model-service watch-agent --pid <进程ID> +``` + +参数: +- `--pid`: 需要监控的代理进程 ID + +### stop 命令 +停止模型服务 +```bash +rock model-service stop +``` + +### anti-call-llm 命令 +反调用 LLM 接口 +```bash +rock model-service anti-call-llm --index <索引> [--response <响应>] +``` + +参数: +- `--index`: 上一个 LLM 调用的索引，从 0 开始 +- `--response`: 上一次 LLM 调用的响应（可选） + +## 文件通信协议 + +模型服务使用文件进行进程间通信，定义了特定的标记格式用于区分请求和响应： + +### 请求格式 +``` +LLM_REQUEST_START{JSON请求数据}LLM_REQUEST_END{元数据JSON} +``` + +### 响应格式 +``` +LLM_RESPONSE_START{JSON响应数据}LLM_RESPONSE_END{元数据JSON} +``` + +### 会话结束标识 +``` +SESSION_END +``` + +元数据包含时间戳和索引信息，用于保证消息顺序和处理。 + +## SDK 使用 + +### ModelServiceConfig + +模型服务配置类，位于 `rock/sdk/sandbox/model_service/base.py`： + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # 服务类型 + install_cmd="pip install rock-model-service", # 安装命令 + install_timeout=300, # 安装超时（秒） + start_cmd="rock model-service start --type ${type}", # 启动命令 + stop_cmd="rock model-service stop", # 停止命令 + logging_path="/data/logs", # 日志路径 + logging_file_name="model_service.log", # 日志文件名 +) +``` + +| 配置项 | 默认值 | 说明 | +|--------|--------|------| +| `enabled` | `False` | 是否启用模型服务（RockAgent 自动管理） | +| `type` | `"local"` | 服务类型：`local` 或 `proxy` | +| `install_cmd` | - | 模型服务包安装命令 | +| `install_timeout` | `300` | 安装超时时间（秒） | +| `start_cmd` | - | 启动命令模板 | +| `stop_cmd` | - | 停止命令 | +| `logging_path` | `/data/logs` | 日志目录路径 | +| `logging_file_name` | `model_service.log` | 日志文件名 | + +### ModelService + +模型服务管理类，处理沙箱内模型服务的生命周期： + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# 通常由 RockAgent 自动管理，无需手动调用 +# 以下方法仅在需要手动控制时使用 + +# 安装模型服务 +await model_service.install() + +# 启动模型服务 +await model_service.start() + +# 监控代理进程 +await model_service.watch_agent(pid="12345") + +# 执行反调用 LLM（Local 模式核心 API） +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# 停止模型服务 +await model_service.stop() +``` + +## API 参考 + +### install() + +在沙箱中安装模型服务依赖。 + +```python +await model_service.install() +``` + +执行步骤： +1. 创建并初始化 Python 运行时环境 +2. 创建 Rock 配置文件 +3. 安装模型服务包 + +**注意：** 通常由 RockAgent 自动调用。 + +### start() + +启动模型服务。 + +```python +await model_service.start() +``` + +前提条件：必须先调用 `install()`。 + +**注意：** 通常由 RockAgent 自动调用。 + +### stop() + +停止模型服务。 + +```python +await model_service.stop() +``` + +如果服务未运行，会跳过此操作。 + +**注意：** 通常由 RockAgent 自动调用。 + +### watch_agent(pid) + +监控代理进程。 + +```python +await model_service.watch_agent(pid="12345") +``` + +当进程退出时，发送 `SESSION_END` 消息。 + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +执行反调用 LLM 操作。**这是 Local 模式下最重要的 API。** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM 调用索引 + response_payload='{"result": "..."}', # 响应数据（可选） + call_timeout=600, # 操作超时（秒） + check_interval=3, # 状态检查间隔（秒） +) +``` + +## 配置选项 + +### 服务配置 +- `SERVICE_HOST`: 服务主机地址，默认为 `"0.0.0.0"` +- `SERVICE_PORT`: 服务端口，默认为 `8080` + +### 日志配置 +- `LOG_FILE`: 用以通信的日志文件路径，包含请求和响应数据 + +### 轨迹（Traj）日志记录 +模型服务将 LLM 调用轨迹（traj）记录到 JSONL 文件中，用于调试和分析。 + +| 环境变量 | 默认值 | 说明 | +|----------|--------|------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | traj 日志文件目录 | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | 追加模式（true/false） | + +**traj 文件位置**: `{DATA_DIR}/LLMTraj.jsonl` + +**traj 文件格式**（JSONL - 每行一个 JSON 对象）： +```json +{"request": {...}, "response": {...}} +``` + +### 轮询配置 +- `POLLING_INTERVAL_SECONDS`: 轮询间隔，默认为 `0.1` 秒 +- `REQUEST_TIMEOUT`: 请求超时时间，默认为无限 + +### 标记配置 +定义了用于区分日志文件中不同类型消息的标记： +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig（服务端） + +服务端配置类定义了模型服务如何处理请求： + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `host` | str | `"0.0.0.0"` | 服务器地址 | +| `port` | int | `8080` | 服务器端口 | +| `proxy_base_url` | str \| None | `None` | 直接代理 URL | +| `proxy_rules` | dict | 见下方 | 模型名称到 URL 的映射 | +| `retryable_status_codes` | list[int] | `[429, 500]` | 可重试的 HTTP 状态码 | +| `request_timeout` | int | `120` | 请求超时时间（秒） | + +**默认 proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/python_sdk.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..c1083f29b0 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK 参考 + +本指南详细介绍如何使用 ROCK SDK 进行开发，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK SDK为开发者提供了便捷的Python接口来使用ROCK平台的功能，包括沙箱环境管理和GEM环境交互。 + +> **重要提示**: 使用 SDK 之前，请确保 ROCK Admin 服务正在运行。可以通过以下命令启动： +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 基本沙箱操作 + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 沙箱组管理 + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# 创建沙箱组配置 +config = SandboxGroupConfig( + image="python:3.11", + size=4, # 创建4个沙箱 + start_concurrency=2, # 并发启动级别为2 +) + +# 创建并启动沙箱组 +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# 批量操作 +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# 批量停止 +await sandbox_group.stop() +``` + +### 2.3 配置示例 + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 沙箱加速配置 + +ROCK 提供沙箱网络加速功能，支持配置 APT、PIP 和 GitHub 镜像源，提升受限网络环境下的包下载速度。 + +#### 支持的加速类型 + +**APT 镜像配置** + +配置 APT 包管理器镜像源，加速 Debian/Ubuntu 软件包下载。 + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# 配置 APT 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP 镜像配置** + +配置 Python 包索引镜像，加速 pip 安装。 + +```python +# HTTP 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS 镜像 +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub 加速** + +通过添加自定义 DNS 解析条目加速 GitHub 访问。 + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### 完整示例 + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """创建沙箱并配置加速""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # 配置加速（在安装包之前配置） + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # speedup 不会主动安装 PIP，仅配置镜像源进行加速 + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # 可以通过镜像 IP 加速 GitHub 访问 + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### 注意事项 + +1. **配置顺序**: 在安装包之前配置加速 +2. **HTTPS vs HTTP**: HTTPS 镜像不需要为 PIP 配置 trusted-host +3. **GitHub IP**: 不同区域可能需要不同的 IP 以获得最佳性能 +4. **持久性**: 配置在沙箱生命周期内持久有效 +5. **多次调用**: 后续的加速调用会覆盖之前的配置 +6. **PIP 安装**: speedup 功能仅配置镜像源，不会自动安装 PIP + +## 3. GEM SDK + +### 3.1 Python SDK 方式 + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## 相关文档 +- [快速开始指南](../../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK SDK +- [API 文档](../api.md) - 查看 SDK 封装的底层 API 接口 +- [配置指南](../../User%20Guides/configuration.md) - 了解 SDK 相关的配置选项 +- [安装指南](../../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/remote_user.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..791ca85fdd --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/remote_user.md @@ -0,0 +1,69 @@ +# Remote User + +远程用户管理，用于在沙箱中创建和管理用户。 + +## 使用示例 + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +创建远程用户。 + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +检查用户是否存在。 + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/rock-agent.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..cd252750aa --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,310 @@ +# Install Agent in Sandbox (Experimental) + +> 这是 ROCK 两种并列的 agent 使用能力中 **install-agent** 的参考文档,核心 API 是 `sandbox.agent.install()` 与 `sandbox.agent.run(prompt)`,用于在单个沙箱里安装并运行 agent。 +> +> 另一种能力是用 Job 在沙箱里跑一次 agent 评测/任务,见 [Use Job to Run Agent](./job.md)。两种能力使用各自独立的配置 schema。 + +RockAgent 是 ROCK 框架用来在沙箱中 install 自定义 agent 的能力,负责完整的 agent 生命周期管理:环境初始化、ModelService 集成、命令执行等。 + +使用 `sandbox.agent.install()` 与 `sandbox.agent.run(prompt)` 就可以在 Rock 提供的 Sandbox 环境中安装和运行 Agent。 + +## 核心概念 + +RockAgent 的核心工作流程分为两个阶段： + +1. **install(config)**: 初始化 Agent 环境，包括部署工作目录、设置环境变量、初始化运行时环境等 +2. **run(prompt)**: 执行 Agent 任务，替换占位符并启动 Agent 进程 + +## 快速开始 + +### Claude Code 示例 + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli 示例 + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} 必须 + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # 环境变量 + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent 示例 + +```yaml +working_dir: "." # 上传包含 langgraph_agent.py 的本地当前目录到 sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # 运行本地脚本 + +runtime_env_config: + type: python + pip: # 安装 pip 依赖 + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## 配置详解 + +### 基础配置 + +```yaml +agent_type: "default" # Agent 类型标识（默认: "default"） +agent_name: "demo-agent" # Agent 实例名称（默认: 随机 uuid） +version: "1.0.0" # 版本标识（默认: "default"） +instance_id: "instance-001" # 实例 ID（默认: "instance-id-<随机uuid>"） +agent_installed_dir: "/tmp/installed_agent" # Agent 安装目录（默认: "/tmp/installed_agent"） +agent_session: "my-session" # bash 会话标识（默认: "agent-session-<随机uuid>"） +env: # 环境变量（默认: {}） + OPENAI_API_KEY: "xxxxxxx" +``` + +### 工作目录配置 + +```yaml +working_dir: "./my_project" # 本地目录，上传到 sandbox（默认: None 不上传） +project_path: "/testbed" # sandbox 中工作目录，用于 cd（默认: None） +use_deploy_working_dir_as_fallback: true # project_path 为空时是否回退到 deploy.working_dir（默认: true） +``` + +### 执行配置 + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent 执行命令，必须包含 ${prompt}（默认: None） + +skip_wrap_run_cmd: false # 跳过为 run_cmd 添加 PATH 的包装（默认: false） + +# 超时配置 +agent_install_timeout: 600 # 安装超时，单位秒（默认: 600） +agent_run_timeout: 1800 # 运行超时，单位秒（默认: 1800） +agent_run_check_interval: 30 # 检查间隔，单位秒（默认: 30） +``` + +**`skip_wrap_run_cmd`**: +- `false`（默认）：为命令添加 `export PATH=:$PATH &&` 包装，确保使用运行时环境的可执行文件 +- `true`：跳过 PATH 包装，直接使用 `bash -c` 运行命令 + +### 初始化钩子 + +```yaml +pre_init_cmds: # 初始化前执行的命令（默认: 从 env_vars 读取） + - command: "apt update && apt install -y git" + timeout_seconds: 300 # 命令超时，单位秒（默认: 300） + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # 初始化后执行的命令（默认: []） + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**注意事项**： +- `pre_init_cmds` 和 `post_init_cmds` 不继承 Agent 的 `env` 环境变量 +- 通常用于执行安装操作和配置文件移动操作 +- 常用命令示例： + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv 配置 + +```yaml +runtime_env_config: # 具体参考 RuntimeEnv 有关文档 + type: "python" # 运行时类型: python / node（默认: "python"） + version: "3.11" # 版本号 + pip: # Python 依赖包列表 + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node 运行时示例**： + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**自动执行的操作**： +- 根据 `type` 安装对应的运行时（Python 或 Node.js） +- 安装 `pip` 依赖（如果配置了） +- 执行 `custom_install_cmd` 自定义安装命令（如果配置了） +- 支持 `npm_registry` 配置 Node.js 的 npm 镜像源 + +### ModelService 配置 + +```yaml +model_service_config: # 具体参考 ModelService 有关文档 + enabled: true # 启用 ModelService（默认: false） +``` + +**自动执行的操作**： +- 安装阶段：安装 ModelService（仅安装，不启动） +- 运行阶段：启动 ModelService + `watch_agent` 监控进程 + +**注意事项**：需要将模型请求的 URL 设置为 ModelService 的 URL。例如 ModelService 提供的 OpenAI-compatible 的 URL 为 `http://127.0.0.1:8080/v1/chat/completions`，则通常需要将 Agent 向 LLM 请求的 URL 设置为 `http://127.0.0.1:8080/v1/`。 + +## API 参考 + +### install(config) + +初始化 Agent 环境。 + +**执行流程**： +1. 如果配置了 `working_dir`，部署到 sandbox +2. 设置 bash session，以及配置 env 环境变量 +3. 执行 `pre_init_cmds` +4. 并行初始化 RuntimeEnv 和 ModelService（如果启用） +5. 执行 `post_init_cmds` + +**参数**： +- `config`: Agent 配置文件，支持两种传入方式： + - **字符串路径**: YAML 配置文件路径，默认值为 `"rock_agent_config.yaml"` + - **RockAgentConfig 对象**: 直接传入 `RockAgentConfig` 实例 + +### run(prompt) + +执行 Agent 任务。 + +**执行流程**： +1. 替换占位符,准备 Agent 运行命令 +2. 启动 agent 进程 +3. 如果启用 ModelService,启动 `watch_agent` +4. 等待任务完成并返回结果 + +## 高级用法 + +### working_dir 与 project_path 的区别与联动 + +| 配置项 | 作用 | 联动方式 | +|--------|------|----------| +| `working_dir` | 本地目录，上传到 sandbox | 调用 `deploy.deploy_working_dir()` 上传，上传后 `deploy.working_dir` 变为 sandbox 中的路径 | +| `${working_dir}` | 命令中的占位符 | 被 `deploy.format()` 替换为 `deploy.working_dir` 的值，会在配置中的 init_cmds 和 run_cmd 中替换 | +| `project_path` | sandbox 中的工作目录 | 用于运行前 `cd project_path`，不设置时会进入到 `deploy.working_dir` 工作目录 | +| `use_deploy_working_dir_as_fallback` | run 时 project_path 未设置时是否回退到 deploy.working_dir | 默认为 `true`，设为 `false` 时即使未设置 project_path 也不会进入 working_dir | + +**使用建议**： +- 使用 `working_dir` 上传本地项目代码到 sandbox +- 使用 `project_path` 指定 sandbox 中的工作目录（如 `/testbed`） +- 设置 `use_deploy_working_dir_as_fallback: false` 的场景：需要进行本地文件挂载，但希望在镜像默认工作目录下运行 Agent + +### 占位符使用 + +Rock Agent 在支持在配置文件中替换以下占位符： + +- `${prompt}`: 在run_cmd 中必需，会被替换为 `run(prompt)` 传入的提示词 +- `${working_dir}`: 可选，会被替换为 sandbox 中实际的工作目录路径, 同时支持在 init_cmds和 run_cmd 中使用 +- `${bin_dir}`: 可选，会被替换为运行时环境的 bin 目录路径 + +**示例**： +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback 说明 + +当 `project_path` 未设置时： +- `true`（默认）：运行 Agent 前会自动 `cd` 到 `deploy.working_dir` +- `false`：运行 Agent 前不会自动切换目录，保持在当前目录 + +适用场景： +- `true`: 大多数场景，希望 Agent 在上传的代码目录中运行 +- `false`: 需要挂载本地文件，但希望在镜像默认工作目录（如 `/app, /testbed`）下运行 Agent + +## 完整配置示例 + +```yaml +# ========== 基础配置 ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== 工作目录配置 ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== 运行配置 ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# 超时配置 +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== 初始化命令 ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== 运行时环境配置 ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService 集成 ========== +model_service_config: + enabled: true +``` + +## 使用示例 + +### 使用 YAML 配置文件(推荐) + +```python +import asyncio +from rock.sdk.sandbox import Sandbox, SandboxConfig + +async def main(): + sandbox = Sandbox(SandboxConfig()) + await sandbox.start() + try: + # rock_agent_config.yaml 与本文档「快速开始」中的示例一致 + await sandbox.agent.install(config="rock_agent_config.yaml") + result = await sandbox.agent.run(prompt="hello") + print(result) + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +更多开箱即用的示例参见 `examples/install-agents/`(Claude Code、IFlowCli、Cursor CLI、Qwen Code、SWE-agent、OpenClaw 等)。 + +如需通过 Job 跑 agent 评测/基准任务(另一条代码路径,有独立的配置 schema),见 [Use Job to Run Agent](./job.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/runtime-env.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..a5532e900b --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,137 @@ +# RuntimeEnv + +RuntimeEnv 模块用于在沙箱中管理语言运行时环境（目前提供了 Python / Node.js）。 + +## 快速开始(使用示例) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +异步工厂方法，根据配置创建 RuntimeEnv 实例并初始化，自动注册到 `sandbox.runtime_envs`。 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# 自动注册，可通过 sandbox.runtime_envs[env.runtime_env_id] 访问 +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +包装命令，将 `bin_dir` 加入 PATH，确保优先使用运行时环境中的可执行文件。 + +```python +wrapped = env.wrapped_cmd("node script.js") +# 返回: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +在运行时环境中执行命令。内部基于 `wrapped_cmd` 实现 + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | 类型标识 | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python 版本，默认 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | pip 包列表或 requirements.txt 路径 | +| `pip_index_url` | `str \| None` | 环境变量 | pip 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | 要创建符号链接的可执行文件列表 | + +## NodeRuntimeEnvConfig + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | 类型标识 | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node 版本，默认 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm 镜像源 | +| `extra_symlink_dir` | `str \| None` | `None` | 符号链接的目标目录 | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | 要创建符号链接的可执行文件列表 | + +## 自定义 RuntimeEnv 实现约束 + +自定义 RuntimeEnv 需遵循以下规则： + +1. **定义 `runtime_env_type` 类属性**：作为类型标识符，用于自动注册到 RuntimeEnv 工厂 +2. **重写 `_get_install_cmd()`**：返回安装命令 +3. **安装命令最后必须**：将目录重命名为 `runtime-env` + + +## NodeRuntimeEnv 简化版实现示例 + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config 类：定义配置类型，用于 RuntimeEnv.create() 路由到对应实现 +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # 必须与 runtime_env_type 一致 + +# RuntimeEnv 实现类：定义如何安装和运行该运行时环境 +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # 自动注册到 RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # 下载 Node 二进制包并解压，最后重命名为 runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## 加速基础环境安装 + +`PythonRuntimeEnv` 默认从 https://github.com/astral-sh/python-build-standalone/releases/ 下载 Python 安装包。若网络不可达或下载较慢，可通过环境变量 `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` 或 `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` 覆盖默认安装命令（例如切换到内网源/镜像源）。 + +默认值示例： + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +例如，替换为镜像源下载： + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +请确保该命令执行完成后，会在 `runtime_env` 的默认工作目录下生成 `runtime-env` 目录，并且 `${workdir}/runtime-env/bin/` 下包含对应可执行文件，例如： + +- `${workdir}/runtime-env/bin/python` + +Node 环境同理，可通过修改环境变量 `ROCK_RTENV_NODE_V22180_INSTALL_CMD` 来指定更快的下载/安装命令。 \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/sandbox.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..088f1e3110 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/sandbox.md @@ -0,0 +1,113 @@ +# 处理大文件和长命令输出 + +## `arun` +`arun()` 在 `nohup` 模式下提供了两个关键参数，帮助 Agent / 调用方在"执行"与"查看"之间按需解耦： + +1. **`response_limited_bytes_in_nohup`**（int 型） + 限制返回内容的最大字符数（例如 `64 * 1024`），适合仍需立刻查看部分日志、但必须控制带宽的场景。默认值 `None` 表示不加限制。 + +2. **`ignore_output`**（bool，默认 `False`） + 当设为 `True` 时，`arun()` 不再读取 nohup 输出文件，而是在命令执行完毕后立即返回一段提示信息（包含输出文件路径、**文件大小**及查看方式）。日志仍写入 `/tmp/tmp_.out`，后续可通过 `read_file`、下载接口或自定义命令按需读取，实现"执行"与"查看"彻底解耦。返回的文件大小信息可帮助用户决定是直接下载还是分块读取。 + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# 示例 1：限制最多 1024 个字符 +resp_limit = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# 示例 2：完全跳过日志读取，后续再通过 read_file / 下载获取 +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# 可通过 Sandbox.read_file(...) / 下载接口 / cat /tmp/tmp_xxx.out 查看日志 +``` + +## `read_file_by_line_range` + +按行范围异步读取文件内容，支持自动分块读取和会话管理，支持大文件读取。 + +### 重要特性 +- **大文件分块读取**: 自动将大文件分成多个小块进行读取 +- **自动统计行数**: 未指定结束行时，自动计算文件总行数 +- **内置重试机制**: 关键操作支持最多 3 次重试，提高可靠性 +- **参数验证**: 自动验证输入参数的合法性 +- **会话管理**: 支持指定会话或自动创建临时会话 + +### 参数说明 +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `file_path` | str | - | 要读取的文件路径（沙箱中的绝对路径或相对路径） | +| `start_line` | int \| None | 1 | 起始行号（从 1 开始） | +| `end_line` | int \| None | None | 结束行号（包含），默认为文件末尾 | +| `lines_per_request` | int | 1000 | 每次请求读取的行数，范围 1-10000 | + +### 返回值 +- `ReadFileResponse`: 包含文件内容的响应对象 + - `content` (str): 读取的文件内容 + +### 异常说明 +- `Exception`: 当 `start_line < 1` 时抛出 +- `Exception`: 当 `end_line < start_line` 时抛出 +- `Exception`: 当 `lines_per_request` 不在 1-10000 范围内时抛出 +- `Exception`: 当文件读取失败时抛出 + +### 使用示例 + +```python +# 读取整个文件 +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# 读取指定行范围（第 100 到 500 行） +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# 从第 1990 行读取到文件末尾 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# 使用自定义分块大小 +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### 注意事项 +- 行号从 1 开始计数，而非 0 +- 对于大文件建议适当增加 `lines_per_request` 以提高效率 +- 文件路径必须是沙箱内的有效路径 +- 使用 `sed` 命令进行文件读取，确保沙箱镜像支持该命令 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..6a74f397c7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,228 @@ +# SWE-Bench 评测 + +本文档介绍如何使用 ROCK SDK 运行 SWE-Bench Verified 评测，包括沙箱启动、Agent 集成、测试环境准备和结果解析。 + +### 快速开始 +SWE-Bench-Verified 是一个用于评估 AI 编程 Agent 在真实软件工程任务上表现的基准测试。 + +在ROCK上运行一个SWE-Bench任务包含以下步骤： + +1. **load_task_config** — 加载 `task.yaml` 获取任务指令 +2. **start_sandbox** — 使用任务专属的 Docker 镜像启动沙箱 +3. **agent.install / agent.run** — 安装并运行 Agent 来解决任务 +4. **setup_test_env** — 上传测试文件和运行测试脚本到沙箱 +5. **运行测试** — 通过 `sandbox.arun()` 执行测试脚本，支持超时控制 +6. **parse_swebench_result** — 解析测试输出，判断 PASSED / FAILED +7. **sandbox.stop** — 清理沙箱资源 + +**下面是示例代码** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. 加载任务指令 + task_config = await load_task_config(task_dir) # 参见 load_task_config 章节 + instruction = task_config["instruction"] + + # 2. 启动沙箱 + sandbox = await start_sandbox(task_name) # 参见 start_sandbox 章节 + + try: + # 3. 安装并运行 Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. 准备测试环境 + await setup_test_env(sandbox, task_dir) # 参见 setup_test_env 章节 + + # 5. 运行测试 + resp = await run_tests(sandbox) # 参见"运行测试"章节 + + # 6. 解析结果 + is_resolved = parse_swebench_result(resp.output) # 参见 parse_swebench_result 章节 + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +以下章节详细介绍评测流程中使用的各个函数。 + +--- + +## start_sandbox + +使用任务专属的 SWE-Bench Docker 镜像启动沙箱实例。每个任务都有一个预构建的镜像，包含目标仓库和运行环境。 + +`image` 参数格式如下： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +例如，任务 `django__django-14539` 对应的镜像为： + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +从任务目录中加载 `task.yaml` 配置文件。YAML 文件包含 `instruction` 字段，用于描述 Agent 需要完成的编程任务。 + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# 使用示例 +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +使用 `sandbox.agent.install()` 和 `sandbox.agent.run()` 在沙箱中部署和执行 Agent。详细的 Agent 配置请参考 [Rock Agent](./rock-agent.md)。 + +```python +# 使用 YAML 配置文件安装 Agent（以 iflow_config.yaml 为例） +await sandbox.agent.install(config="iflow_config.yaml") + +# 使用任务指令运行 Agent +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +在沙箱中准备测试环境：安装 [uv](https://github.com/astral-sh/uv) 包管理器，并上传测试文件和运行测试脚本。 + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """准备测试环境并返回会话名称。""" + # 1. 创建带有自定义环境变量的会话 + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. 安装 uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. 上传测试文件 + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. 上传运行测试脚本 + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## 运行测试 + +使用 `RunMode.NOHUP` 模式执行测试脚本，支持可配置的超时时间。 + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +解析测试输出以判断 SWE-Bench 任务是否通过。解析器会查找由标记行分隔的结果块，并检查是否包含 `PASSED`。 + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """解析 SWE-Bench 测试输出，判断任务是否通过。 + + 匹配 'SWEBench results starts here' 和 + 'SWEBench results ends here' 之间的内容块， + 然后检查其中是否包含 'PASSED'。 + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# 使用示例 +is_resolved = parse_swebench_result(resp.output) +``` + +## 注意事项 + +- **任务数据集**：任务目录（包含 `task.yaml`、`tests/` 和 `run-tests.sh`）可从 [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) 仓库获取。 +- **任务镜像**：每个 SWE-Bench 任务需要特定的 Docker 镜像（如 `sweb.eval.x86_64.`）。请确保镜像在对应的环境中可用。 +- **Agent 配置**：Agent 配置 YAML 定义了运行时、依赖和执行命令。详情请参考 [Rock Agent](./rock-agent.md)。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/api.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/api.md new file mode 100644 index 0000000000..06f44c326c --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/References/api.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 1 +--- + +# API 参考 + +本指南详细介绍 ROCK 平台提供的核心 API 服务，包括沙箱环境管理和 GEM 环境交互。 + +## 1. 概述 + +ROCK平台提供两种核心API服务： +- Sandbox API：沙箱环境管理 +- GEM API：GEM环境交互 + +所有 API 接口都遵循 RESTful 设计原则，支持 JSON 格式的数据交换。 + +## 2. Sandbox API + +沙箱环境全生命周期管理功能： + +### 沙箱管理接口 + +1. **Start Sandbox** - 启动沙箱环境 + - 创建一个新的沙箱实例 + - 支持指定镜像、资源配置等参数 + +2. **Start Sandbox Async** - 异步启动沙箱环境 + - 异步方式创建沙箱实例 + - 适用于需要快速响应的场景 + +3. **Check Sandbox Alive Status** - 检查沙箱存活状态 + - 验证沙箱是否正常运行 + +4. **Get Sandbox Statistics** - 获取沙箱统计信息 + - 获取沙箱的资源使用统计 + +5. **Get Sandbox Status** - 获取沙箱详细状态 + - 获取沙箱的完整状态信息 + +6. **Stop Sandbox** - 停止沙箱环境 + - 安全关闭沙箱实例 + +7. **Commit Sandbox** - 提交沙箱为镜像 + - 将当前沙箱状态保存为新镜像 + +### 命令执行接口 + +8. **Execute Command** - 在沙箱中执行命令 + - 直接在沙箱中运行指定命令 + +9. **Create Bash Session** - 创建Bash会话 + - 创建持久化的Bash会话环境 + +10. **Run Command in Session** - 在会话中执行命令 + - 在已创建的会话中执行命令 + +11. **Close Session** - 关闭会话 + - 释放会话资源 + +### 文件操作接口 + +12. **Read File** - 读取沙箱文件 + - 从沙箱中读取指定文件内容 + +13. **Write File** - 写入沙箱文件 + - 向沙箱中写入文件 + +14. **Upload File** - 上传文件到沙箱 + - 将本地文件上传到沙箱 + +## 3. GEM API + +GEM环境交互功能： + +1. **Make Environment** - 创建GEM环境 + - 初始化一个新的GEM环境实例 + +2. **Reset Environment** - 重置GEM环境 + - 将GEM环境重置到初始状态 + +3. **Step Environment** - 执行GEM环境步骤 + - 在GEM环境中执行一个动作步骤 + +4. **Close Environment** - 关闭GEM环境 + - 释放GEM环境资源 + +## 4. HTTP API 使用示例 + +### 4.1 Sandbox API 示例 + +#### 启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 异步启动沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### 执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### 创建会话 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### 在会话中执行命令 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### 上传文件 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### 停止沙箱 +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API 示例 + +```bash +# 创建GEM环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# 重置环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# 执行步骤 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# 关闭环境 +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速开始使用 ROCK API +- [Python SDK 文档](./Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 调用 API +- [配置指南](../User%20Guides/configuration.md) - 了解 API 相关的配置选项 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/index.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/index.md new file mode 100644 index 0000000000..02a4d7b6b6 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/index.md @@ -0,0 +1,6 @@ +--- +sidebar_position: 1 +--- +# 版本说明 + +* [release v1.8.0](v1.8.0.md) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md new file mode 100644 index 0000000000..c9a18d7a69 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md @@ -0,0 +1,154 @@ +# v1.8.0 + +## 发布日期 +2026 年 5 月 21 日 + +--- + +## 🐍 SDK 相关改动 + +### 新功能 + +* **SDK**: 使沙箱集群默认值可通过环境变量配置  + + +### Bug 修复 + +* **SDK**: 在 Python 运行时环境中增加配置类型校验 (#652) + +* **SDK**: 在 OSS 上传路径中使用 wget 前先创建目标父目录 (#940) + + +--- + +## 📦 沙箱 (Sandbox) + +### 新功能 + +* 新增 CPU 超分能力，支持灰度发布、生命周期摘要和绝对核心数 CPU 指标 (#978) + +* `get_status` 接口支持 `include_all_states` 参数，可查询所有状态的沙箱 (#951) + +* 上传/下载文件到sandbox的账号、bucket迁移（向前兼容，老bucket仍支持使用）（#953） + +* 优化中转文件直接存储到bucket根目录问题（向前兼容） + + +--- + +## 🚀 部署 (Deployments) + +### 新功能 + +* **Kubernetes**: 新增 GPU 支持，采用 Jinja2 模板和可扩展加速器类型 (#981) + +* 区域集群配置优化：抽取公共配置，降低维护成本 + +* 支持rayhead上cron定时任务清理过期日志和目录 + + +### Bug 修复 + +* **Docker**: 容器停止时清理 XFS 项目配额 + +* **Docker**: 移除无效的删除镜像分支，添加 CLS 日志服务支持到删除镜像功能 (#965) + +* 修复删除镜像传参bug + + +--- + +## 🔀 代理 (Proxy) + +### 新功能 + +* 模型服务代理支持流式传输和回放模式，实现字节透传，提供转发和回放两种后端 (#935) + + +--- + +## 📊 指标监控 (Metrics) + +### 新功能 + +* 记录 OTLP 导出的数据点数量和耗时 + +* 新增沙箱启动各阶段的时序埋点监控 + + +### Bug 修复 + +* 从元数据存储读取沙箱镜像信息，替代内存字典 + +* 将 rock\_config 传递给沙箱表和元数据存储，确保指标监控使用正确的端点 + +* 修复 `_get_user_info` 指标问题 (#911) + + +--- + +## ⚙️ 任务调度 (Scheduler) + +### 新功能 + +* 支持通过 Nacos 动态重载配置 (#888) + +* 新增 Ray 日志清理任务，禁用 worker 到 driver 的日志转发 + +* 新增构建缓存清理任务，用于修剪 uv/pip 缓存 + +* 将悬空镜像和 BuildKit 修剪合并到镜像清理任务中 (#970) + +* 优化文件清理定时任务的性能和配置安全验证 + + +### Bug 修复 + +* 处理 Ray 后台任务重连时 `ray.init` 引发的异常 + + +--- + +## 🎯 Rocklet + +### 新功能 + +* 新增 Windows PowerShell 支持 (#921) + + +### Bug 修复 + +* 将循环设备磁盘挂载到 Docker 数据根目录，替代硬编码路径 + +* 为 Kata 运行时的 Nix 镜像添加 `/bin` 符号链接挂载 (#936) + +* 使用 cgroup 指标获取容器 CPU 使用率，替代 psutil + + +--- + +## ✨ 其他新功能 + +* **命令行工具**: 新增 `-v` 参数控制日志详细程度，统一日志级别管理 + +* **任务系统**: 集成沙箱内模型服务代理，支持录制和回放功能 + +* **对象存储**: 统一双账号 STS 令牌获取接口，将传输前缀推送至 SDK + + +## 🐛 其他 Bug 修复 + +* **BashJob**: 将 OSS 上传改为脚本注入方式，修复仅提交模式下的文件丢失和环境变量凭证不生效问题 + + +## ♻️ 代码重构 + +* **对象存储**: 将 OSS 上传/下载与客户端环境变量解耦，实现三层配置解析机制 (#943)[《oss重构提测-ROCK-V1.8》](https://alidocs.dingtalk.com/i/nodes/o14dA3GK8gQlkoYwcK5yyo2QV9ekBD76?corpId=dingd8e1123006514592&utm_medium=im_card&iframeQuery=utm_medium%3Dportal_main_colum_create%26utm_source%3Dportal&utm_scene=person_space&utm_source=im&cid=72193729861) + + +## 🔧 构建与工具 + +* 移除 `need_database` 标记 (#901) + + +--- \ No newline at end of file diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/configuration.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/configuration.md new file mode 100644 index 0000000000..a212189bc7 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/configuration.md @@ -0,0 +1,188 @@ +--- +sidebar_position: 4 +--- + +# 配置指南 + +本指南详细介绍如何配置 ROCK 环境以满足不同的使用需求，包括本地开发、测试和生产部署。 + +## 1. 环境变量配置 + +ROCK 支持通过环境变量配置关键参数。以下是主要的环境变量： + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK服务基础URL +export ROCK_LOG_LEVEL=INFO # 日志级别 +export ROCK_LOGGING_PATH=/path/to/logs # 日志文件路径，默认 None (输出到控制台) +export ROCK_LOGGING_FILE_NAME=rocklet.log # 日志文件名，默认 "rocklet.log", 启动admin时可以自定义日志文件名, 如admin.log +export ROCK_LOGGING_LEVEL=INFO # 日志输出级别，默认 "INFO" +export ROCK_WORKER_ENV_TYPE=local # 运行时环境类型，可选值: local, docker, uv, pip +``` + +更多环境变量可参考 `rock/env_vars.py` 文件。 + +### 1.1 运行时环境 (Runtime Environments) + +ROCK 提供了多种不同的运行时环境来满足不同场景的需求，选择通过环境变量 `ROCK_WORKER_ENV_TYPE` 进行配置。每种环境有不同的部署要求、性能特征和适用场景。每种环境都有其独特的优势和限制，开发者可以根据部署环境的需要选择最适合的运行时环境。 + +#### 1.1.1 Docker 运行时环境 + +Docker 运行时环境适用于已经预安装了所需依赖的 Docker 镜像环境。这种环境要求部署环境中直接可用 `/tmp/miniforge/bin/rocklet` 可执行文件。 + +**挂载配置:** +- `/tmp/miniforge` - 包含预安装的 Python 环境 +- `/tmp/local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 容器化部署环境 +- 已经构建了包含 `rocklet` 的自定义 Docker 镜像 +- 适合生产环境，启动速度快 + +**要求:** +- 需要使用定制的 Docker 镜像，其中包含 `/tmp/miniforge/bin/rocklet` 可执行文件 +- Docker 环境支持 + +#### 1.1.2 本地运行时环境 + +本地运行时环境直接利用当前部署环境的 Python 环境和项目文件。该环境要求宿主机和容器之间具有相同的操作系统，以便能够直接挂载虚拟环境和 Python 解释器。 + +**挂载配置:** +- `python_env_path` - Python 环境路径 +- `project_root` - 项目根目录 +- `.venv` - 虚拟环境目录（挂载为容器中的 `/tmp/miniforge`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**适用场景:** +- 开发环境 +- 宿主机和目标容器使用相同操作系统的场景 +- 需要快速重新使用现有 Python 环境 + +**要求:** +- 相同的操作系统（主机/容器） +- 可直接访问当前部署的 `.venv` 虚拟环境 +- Python 解释器路径兼容 + +#### 1.1.3 UV 运行时环境 + +UV 运行时环境只依赖于可用的 ROCK 项目，但初始化相对较慢且网络要求较高。这种环境最适合没有预配置环境的场景。它从原始项目重新构建 rocklet 环境。这是推荐在 Mac 操作系统上使用的环境。 + +**挂载配置:** +- `project_root` - 项目根目录（挂载为容器中的 `/tmp + project_root`） +- `local_files` - 执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**适用场景:** +- Mac 操作系统 +- 跨操作系统启动 +- 没有预配置环境的场景 +- 没有使用 uv 管理 Rock + +**优势:** +- 无需预构建镜像 +- 跨平台兼容性好 +- 特别适合开发和测试 + +**限制:** +- 初始化速度较慢 +- 网络要求较高 +- 启动时间较长 + +#### 1.1.4 PIP 运行时环境 + +PIP 运行时环境使用 pip 在容器内安装所需依赖。这种环境适合快速设置并能在容器中完成依赖安装的场景，是默认的运行时环境。它不需要预先构建包含依赖的镜像，通过 pip 直接管理 Python 包。 + +**挂载配置:** +- `local_files` - 包含执行所需的本地文件 + +**启动命令:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**适用场景:** +- 使用PIP源安装的ROCK +- 快速测试ROCK + +**优势:** +- 简单的部署设置 + +**限制:** +- 依赖安装时间较长 +- 需要网络访问以安装依赖包 +- 每次启动时都需要安装依赖 + +#### 1.1.5 配置指南 + +根据不同的使用场景，可以参考以下选择指南： + +| 场景 | 推荐环境 | 原因 | +|------|----------|------| +| 生产环境 | Docker 运行时 | 快速启动，稳定性能 | +| 开发环境，同一 OS | 本地运行时 | 环境重用，开发周期快 | +| Mac 开发 | UV 运行时 | 支持最佳的跨平台兼容性 | +| 跨平台开发 | UV 运行时 | 避免环境兼容性问题 | +| 快速测试 | UV 运行时 | 无需预配置工作 | +| PIP源安装 | PIP 运行时 | 直接使用 pip 安装依赖 | + +这些运行时环境通过 `ROCK_WORKER_ENV_TYPE` 环境变量进行配置，该变量可设置为 "local"、"docker"、"uv" 或 "pip"。 + +### 1.2 日志配置 + +在日志配置方面，ROCK 的日志系统具有以下特性： + +- 日志系统不能同时输出到文件和控制台，只有当设置了 `ROCK_LOGGING_PATH` 时，日志才会输出到指定文件，否则输出到控制台。 +- `ROCK_LOGGING_LEVEL` 用于控制日志输出级别，`ROCK_LOG_LEVEL` 用于通用日志级别设置。 + +## 2. 分布式部署要求 + +由于 ROCK 支持分布式部署，当在 Ray 集群的不同节点上运行时，需要满足以下一致性要求： + +#### 目录结构一致性 +在所有 Ray 节点上，必须保证以下目录结构完全一致： +- ROCK 项目仓库目录 +- `.venv` 虚拟环境目录 +- `.venv` 依赖的 base Python 目录 + + +#### 挂载要求 +ROCK 的启动依赖于挂载 ROCK 项目和对应的 base Python 环境，要求在多机环境中保持一致性： + +#### 验证分布式配置 +可以通过以下方式验证分布式部署配置： + +```bash +# 在所有节点上检查目录一致性 +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# 验证 Python 环境可用性 +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# 检查所有节点上的环境变量设置 +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + + + +## 相关文档 + +- [快速开始指南](../Getting%20Started/quickstart.md) - 了解如何快速搭建 ROCK 环境 +- [API 文档](../References/api.md) - 查看沙箱相关的 API 接口 +- [Python SDK 文档](../References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 SDK 配置沙箱 +- [安装指南](../Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/overview.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/overview.md new file mode 100644 index 0000000000..a02536f50d --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/overview.md @@ -0,0 +1,40 @@ +--- +sidebar_position: 1 +--- + +# 概览 + +ROCK (Reinforcement Open Construction Kit) 是一个开源的强化学习环境开发框架，旨在简化强化学习环境的开发、部署和管理流程。 + +## 什么是 ROCK + +ROCK (Reinforcement Open Construction Kit) 是一个开源强化学习环境开发框架。通过使用 ROCK，开发者可以快速地开发强化学习环境，并结合其他强化学习训练框架，实现高效的强化学习训练。 + +ROCK 提供了完整的沙箱环境管理功能，支持容器化部署，能够实现环境的快速创建、运行和销毁。同时，ROCK 兼容 GEM 协议，为强化学习环境提供了标准化的接口。 + +## ROCK 的核心功能 + +1. **简化开发流程**：简化强化学习环境的开发、构建和管理流程，支持多种开源的强化学习环境 +2. **大规模调度部署**：支持快速强化学习环境的大规模调度部署，通过 GEM 协议可以方便地访问强化学习环境 +3. **框架集成**：与其他强化学习训练框架集成，实现大规模可扩展的强化学习训练 + +## ROCK 的价值 + +ROCK 为不同角色的工程师提供了显著价值： + +- **强化学习算法工程师**：ROCK 可以简化强化学习环境的开发流程，让工程师专注于算法实现 +- **强化学习应用工程师**：ROCK 可以进行快速强化学习环境的大规模部署，提高应用开发效率 + +## 相关文档 + +如果您是第一次使用 ROCK，建议按以下顺序阅读文档： +1. [快速开始指南](./Getting%20Started/quickstart.md) - 快速搭建开发环境 +2. [配置指南](./User%20Guides/configuration.md) - 配置您的 ROCK 环境 +3. [Python SDK 文档](./References/Python%20SDK%20References/python_sdk.md) - 学习如何使用 Python SDK 进行开发 +4. [API 文档](./References/api.md) - 了解完整的 API 接口 +5. [安装指南](./Getting%20Started/installation.md) - 详细了解 ROCK 安装和配置 + + + + + diff --git a/docs/versioned_docs/version-1.8.x/Getting Started/installation.md b/docs/versioned_docs/version-1.8.x/Getting Started/installation.md new file mode 100644 index 0000000000..c45b09a8fe --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Getting Started/installation.md @@ -0,0 +1,143 @@ +--- +sidebar_position: 3 +--- + +# Installation + +This document explains how to install and set up the ROCK development environment using both `uv` and `pip`. The project is a Reinforcement Open Construction Kit that supports various components. + +## Using uv (Recommended) + +### Quick Install All Dependencies + +```bash +# Install all dependencies including optional ones +uv sync --all-extras + +# Install development/testing dependencies +uv sync --all-extras --all-groups +``` + +### Install Different Dependency Groups + +#### Core Dependencies Only +```bash +uv sync +``` + +#### Admin Component Dependencies +```bash +uv sync --extra admin +``` + +#### Rocklet Execution Environment Dependencies +```bash +uv sync --extra rocklet +``` + + +#### All Dependencies at Once +```bash +uv sync --all-extras +``` + +#### Development/Testing Dependencies +```bash +uv sync --all-extras --group test +``` + +## Using pip + +### Install from pip source + +#### Core Dependencies Only +```bash +pip install rl-rock +``` + +#### Admin Component Dependencies +```bash +pip install "rl-rock[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install "rl-rock[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install "rl-rock[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install "rl-rock[all]" +``` + +### Install with pip from source code + +#### Core Dependencies Only +```bash +pip install . +``` + +#### Admin Component Dependencies +```bash +pip install ".[admin]" +``` + +#### Rocklet Execution Environment Dependencies +```bash +pip install ".[rocklet]" +``` + +#### Builder Dependencies +```bash +pip install ".[builder]" +``` + +#### Install All Optional Dependencies +```bash +pip install ".[all]" +``` + +## Available Entry Points + +The package provides the following command line scripts: + +- `rocklet`: ROCK execution environment server (rock.rocklet.server:main) +- `admin`: Admin management server (rock.admin.main:main) +- `envhub`: Environment hub server (rock.envhub.server:main) +- `rock`: Main ROCK command line interface (rock.cli.main:main) + +## Development Setup + +### Using uv (Recommended) + +```bash +# Clone and set up development environment +git clone +cd ROCK +uv sync --all-extras --group test + +# Run tests +uv run pytest + +``` + +### Using pip + +```bash +# For development, install in editable mode with all extras +pip install -e ".[all]" + +# Or separately +pip install -e . +pip install ".[admin]" ".[rocklet]" ".[builder]" # Optional extras +``` + +## Additional Notes + +- The project is configured to use the Alibaba cloud PyPI mirror by default: `https://mirrors.aliyun.com/pypi/simple/` +- For local development, running tests requires the `test` dependency group diff --git a/docs/versioned_docs/version-1.8.x/Getting Started/quickstart.md b/docs/versioned_docs/version-1.8.x/Getting Started/quickstart.md new file mode 100644 index 0000000000..2f14808f5b --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Getting Started/quickstart.md @@ -0,0 +1,166 @@ +--- +sidebar_position: 2 +--- + +# Getting Started + +This guide will demonstrate how to use ROCK to create and manage reinforcement learning environments through complete examples. + +## 1. Environment Preparation + +We recommend starting ROCK on Linux systems to maximize dependency reuse and improve environment startup speed. If you need to try on macOS, please refer to the [MacOS Startup](#7-macos-startup) section. + +Before starting, please ensure your system has the following dependencies installed: + +### 1.1 System Requirements + +- **Docker**: ROCK uses Docker for containerized environment management +- **uv**: ROCK uses uv for dependency management and virtual environment creation + +### 1.2 Verify Dependency Installation + +```bash +# Verify Docker installation +docker --version + +# Verify Docker image, and example depends on python:3.11 image +docker pull python:3.11 + +# Verify uv installation +uv --version +``` + +### 1.3 Project Initialization + +```bash +# Clone repository +git clone +cd ROCK + +# Create virtual environment (using uv-managed Python, use python 3.11 as an example) +uv venv --python 3.11 --python-preference only-managed + +# Install all dependency groups +uv sync --all-extras +``` + +> **Important Note**: To ensure ROCK can correctly mount the project and virtual environment along with its base Python interpreter, it is strongly recommended to use uv-managed Python environments to create virtual environments rather than system Python. + +## 2. Activate Virtual Environment + +Before running any ROCK commands, you need to activate the virtual environment. Ensure sys.base_prefix is a uv-managed environment, such as `/root/.local/share/uv/python/cpython-3.11.8-linux-x86_64-gnu` or similar paths. + +```bash +# Activate virtual environment +source .venv/bin/activate + +# Verify Python environment +python -c "import sys; print('Base prefix:', sys.base_prefix)" +``` + +> **Verification Point**: Ensure the output base prefix path points to a uv-managed Python environment, not system Python. + +## 3. Verify Environment Configuration + +After activating the virtual environment, verify that dependencies are installed correctly: + +```bash +# Check key dependencies +python -c "import rock; print(\"Hello ROCK\")" +``` + +## 4. Start ROCK Service + +After activating the virtual environment, start the ROCK Admin service on project root: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Start ROCK Admin service (local environment) +rock admin start +``` + +After the service starts, you will see output similar to the following: + +``` +INFO: Started server process [12345] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://127.0.0.1:8080 (Press CTRL+C to quit) +``` + +> **Service Information**: The ROCK Admin service runs by default on `http://127.0.0.1:8080`. + +## 5. Run Example Environments + +Now you can run example environments to verify the installation. Ensure the ROCK service is running, then open a new terminal window to execute the following commands: + +```bash +# Ensure virtual environment is activated +source .venv/bin/activate + +# Run sandbox example +python examples/sandbox_demo.py + +# Run GEM protocol example +python examples/sokoban_demo.py +``` + +### 5.1 Example Descriptions + +- **sandbox_demo.py**: Demonstrates how to use ROCK's sandbox SDK to create and manage containerized environments +- **sokoban_demo.py**: Demonstrates how to use ROCK's GEM protocol compatible interface to create reinforcement learning environments + +> **Running Requirements**: Ensure the ROCK Admin service is running, as examples need to communicate with the service. + +## 6. Distributed Environment Configuration (Optional) + +For distributed multi-machine environments, ensure the following configurations are consistent: + +1. All machines use the same root Python interpreter for ROCK and uv Python configurations +2. Docker versions are consistent across all nodes +3. Network configuration allows normal communication between nodes + + +## 7. MacOS Startup + +On macOS, if you need to start Linux image environments, you first need to set the environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=uv +``` + +During container startup, the corresponding uv environment will be installed. For details, please refer to the `rock/rocklet/local_files/docker_run_with_uv.sh` script. + +> **Note**: Compared to Linux systems, the startup speed on macOS will be slower and more dependent on network conditions. You can adjust the script according to actual conditions.You can find detatils for ROCK_WORKER_ENV_TYPE in [Configuration Guide](../User%20Guides/configuration.md). + +## 8. Starting from Pip Source + +If starting the Admin Server from Pip source, after completing the ROCK installation by referring to [installation](./installation.md), you need to set an additional environment variable: + +```bash +export ROCK_WORKER_ENV_TYPE=pip +``` + +(This startup method will pull and install the latest rocklet from the PyPI source when starting the container environment. The startup speed is relatively slow, so it is only recommended for testing purposes. For production environments, other startup methods are still recommended.) + +## Summary + +Congratulations! You have successfully completed the ROCK quick start guide. You should now be able to: + +- Properly set up the ROCK development environment +- Use uv-managed Python environments +- Start and manage ROCK services +- Run example programs to verify installation +- Configure ROCK in distributed environments (if needed) + +For a deeper understanding of ROCK's additional features, please refer to the following documents: + +## Next Steps + +- [Configuration Guide](../User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](../References/api.md) - View complete API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the Python SDK for development +- [Installation Guide](./installation.md) - Detailed information about ROCK installation and setup +- [Overview](../overview.md) - Understand ROCK's design philosophy \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/Getting Started/rock-agent.md b/docs/versioned_docs/version-1.8.x/Getting Started/rock-agent.md new file mode 100644 index 0000000000..19369c1aff --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Getting Started/rock-agent.md @@ -0,0 +1,76 @@ +--- +sidebar_position: 4 +--- + +# Rock Agent Quick Start + +ROCK provides two parallel ways to use agents, each suited to a different scenario: + +- **Job**: Run an agent evaluation/task in a sandbox via BashJob / HarborJob (typical benchmarks: SWE-bench, Terminal Bench) — the primary entry point. +- **install-agent**: Install and run an agent directly inside a single sandbox — for local development and one-off debugging. + +Job is covered first. The install-agent section follows at the end, with full reference at [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md). + +## Prerequisites + +- Make sure you have a working ROCK service. If you need to start the service locally, refer to [Quick Start](quickstart.md). + +--- + +## 1. Use Job to Run Agent + +Job has two backends: **Harbor Job** runs an AI agent benchmark task (SWE-bench, Terminal Bench, etc.); **Bash Job** runs a custom shell script inside a sandbox. + +### 1.1 Prepare a yaml + +Pick a starting point and copy the matching template: + +- Harbor Job (Terminal Bench): [`examples/job/harbor/tb_job_config.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor/tb_job_config.yaml.template) +- Bash Job (claw-eval): [`examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval/claw_eval_bashjob.yaml.template) + +Fill in the fields per the template. See [Use Job to Run Agent](../References/Python%20SDK%20References/job.md) for the full field reference of both backends. + +### 1.2 Launch via Python SDK + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +For BashJob usage, full field references, and result-handling details, see [Use Job to Run Agent](../References/Python%20SDK%20References/job.md). + +--- + +## 2. install-agent: Install and Run an Agent in a Sandbox + +For local development or debugging a single agent run, the core API is: + +```python +await sandbox.agent.install(config="rock_agent_config.yaml") +result = await sandbox.agent.run(prompt="hello") +``` + +The `examples/install-agents/` directory ships ready-to-run examples: + +- `examples/install-agents/iflow_cli/` — IFlowCli +- `examples/install-agents/claude_code/` — Claude Code +- `examples/install-agents/cursor_cli/`, `qwen_code/`, `swe_agent/`, `openclaw/` — others + +Run the Claude Code example: + +```bash +cd examples/install-agents/claude_code +python claude_code_demo.py +``` + +For full RockAgentConfig field details, placeholder semantics, and API reference, see [Install Agent in Sandbox (Experimental)](../References/Python%20SDK%20References/rock-agent.md). diff --git a/docs/versioned_docs/version-1.8.x/Getting Started/rockroll.md b/docs/versioned_docs/version-1.8.x/Getting Started/rockroll.md new file mode 100644 index 0000000000..2465a7733f --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Getting Started/rockroll.md @@ -0,0 +1,194 @@ +--- +sidebar_position: 7 +--- + +# ROCK & ROLL Quick Start Guide + +This guide will walk you through running a reinforcement learning training example based on the Sokoban game, using ROLL (the training framework) and ROCK (the environment management tool). + +## 1. Prerequisites + +Before you begin, please ensure your system has the following dependencies installed. + +### 1.1 System Requirements + +- **OS**: A Linux-based system is recommended (e.g., Ubuntu 20.04+). +- **Hardware**: An NVIDIA GPU with the corresponding drivers is recommended. +- **Docker**: ROCK uses Docker for containerized environment management. +- **uv**: ROCK uses uv for dependency management and virtual environment creation. + +### 1.2 Verify Dependencies & Pre-pull Image + +```bash +# Verify Docker installation +docker --version + +# Verify Docker is running and pre-pull the Sokoban environment image +# This will save time when the training starts. +docker pull rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/sokoban-sandbox:latest + +# Verify uv installation +uv --version + +``` + +### 1.3 Initialize the Project + +```bash +# Clone the project repositories +git clone https://github.com/alibaba/ROCK.git +git clone https://github.com/alibaba/ROLL.git + +# Ensure both repositories are in the same parent directory, like this: +# your-workspace/ +# ├── ROCK/ +# └── ROLL/ +``` + + +## 2. Launch the Training Process + +> Note: The following instructions use torch==2.6.0 and vLLM==0.8.4 as an example. + + +### Option 1: Using a Virtual Environment (Recommended) + +#### Why is this method recommended? +- Isolation: A uv virtual environment ensures that project dependencies are isolated from your system, preventing conflicts. +- Fast Startup: ROCK can reuse this virtual environment, significantly speeding up subsequent task initializations. +- Stability & Reproducibility: Dependency management is cleaner and more reliable. + + +```bash +# Navigate to the ROCK directory +cd ROCK + +# Create and activate a Python 3.10 virtual environment (ROLL recommends Python 3.10) +uv venv --python 3.10 --python-preference only-managed +source .venv/bin/activate + +# Install all of ROCK's dependencies using uv +uv sync --all-extras + +# If using Python 3.10, starting Ray may raise a `ValueError: is not a valid Sentinel`. +# This is due to an incompatibility between `ray` and `click` versions 8.3+. +# To fix this, downgrade `click` to a version below 8.3. This issue does not affect Python 3.11. +uv pip install 'click>=8.2,click<8.3' + +# Navigate to the ROLL directory to install its dependencies +cd ../ROLL + +# Install core PyTorch components +uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 + +# Install transformer-engine. The --no-build-isolation flag prevents errors where torch cannot be found. +uv pip install transformer-engine[pytorch]==2.2.0 --no-build-isolation + +# Install a pre-compiled version of flash-attention matching the specific CUDA and PyTorch versions +uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# Install the remaining dependencies +uv pip install -r requirements_torch260_vllm.txt + +# (Optional) Install Tensorboard to check training metrics +uv pip install tensorboard -i $PYPI_MIRROR + +# All set! Launch the training script. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +### Option 2: Using the System Environment (Alternative) + +For optimal compatibility with this method, we recommend running these commands inside one of ROLL's official base Docker images. These images come pre-installed with matching CUDA, cuDNN, and other foundational libraries. + +> [ROLL's Official Docker Image List](https://alibaba.github.io/ROLL/docs/Getting%20Started/Installation/image_address) + + +#### Warning +This method will install all Python packages directly into your current environment (e.g., the container's base system), which may cause conflicts with system packages or other projects. + +Since ROCK cannot reuse the environment, it may need to reinstall some dependencies each time a task starts, leading to slower startup times that are dependent on network speed. + + +```bash +# Install ROCK's dependencies +cd ROCK +pip install . +pip install ".[admin]" + +# Install ROLL's dependencies +cd ../ROLL +pip install -r requirements_torch260_vllm.txt + +# Crucial: Configure ROCK to use uv as its worker environment manager +export ROCK_WORKER_ENV_TYPE=uv + +# Launch the training script +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh +``` + +You have now successfully launched the Sokoban reinforcement learning training process. Happy Rock & Roll！ + + +## 3. Multi-Node Deployment + +Instead of running everything on a single machine, you can deploy the **ROCK Service** and **ROLL job** on separate machines. This is a common client-server setup where they communicate over the network. + +### 3.1 Deploy the ROCK Service on Machine A + +On a dedicated machine (or container), follow the [ROCK Quick Start Guide](./quickstart.md) to deploy and start the ROCK service. + +> **Important** +> After starting the service, take note of its IP address and port (e.g., `http://192.168.1.10:8000`). You will need this address for the subsequent steps. + +### 3.2 Prepare the ROLL Client on Machine B + +On the other machine where you will run the training task, perform the following steps. + +1. Verify Network Connectivity + +First, use the curl command to check if you can reach the ROCK service on Machine A from Machine B. +```bash +# Replace : with the actual address of your ROCK service +# If successful, you should receive a response like {"message":"hello, ROCK!"} +curl http://: +``` + +2. Prepare the ROLL Environment + +```bash +# Clone the ROLL repository +git clone https://github.com/alibaba/ROLL.git +cd ROLL + +# Install dependencies +pip install -r requirements_torch260_vllm.txt +``` + +3. Configure the ROLL Connection Address + +Modify ROLL's configuration file to point to the remote ROCK service. +- Open the configuration file: examples/agentic_demo/agentic_val_sokoban_sandbox.yaml. +- Find the "SokobanSandbox" section under "env_config". +- Update the base_url value to your ROCK service's address. +```yaml +custom_envs: + SokobanSandbox: + env_config: + # Change the address here to your ROCK service's address + # Example: base_url: 'http://192.168.1.10:8000' + base_url: 'http://:' +``` + +4. Start Training +Once configured, you can start the ROLL training script on Machine B. + +```bash +# This script will now request environments from the ROCK service on Machine A over the network. +bash examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh +``` + +### Advanced: Distributed ROLL Training + +If you wish to deploy the ROLL training task itself in a distributed manner, you can refer to ROLL's official documentation for distributed deployment. +> [Quick Start: Multi-Node Deployment Guide](https://alibaba.github.io/ROLL/docs/Getting%20Started/Quick%20Start/multi_nodes_quick_start) \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/codes.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/codes.md new file mode 100644 index 0000000000..dceb8d3182 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/codes.md @@ -0,0 +1,93 @@ +# Error Codes + +Error code definitions and categories for error handling and retry strategies. + +## Usage Example + +```python +import rock + +def test_codes_values(): + """Test basic status code values""" + assert rock.codes.OK == 2000 + assert rock.codes.BAD_REQUEST == 4000 + assert rock.codes.INTERNAL_SERVER_ERROR == 5000 + assert rock.codes.COMMAND_ERROR == 6000 +``` + +## Codes Categories + +```python +OK = 2000, "OK" +""" +Success codes (2xxx) +""" + +BAD_REQUEST = 4000, "Bad Request" +""" +Client error codes (4xxx): + +These errors indicate issues with the client request, +SDK will raise Exceptions for these errors. +""" + +INTERNAL_SERVER_ERROR = 5000, "Internal Server Error" +""" +Server error codes (5xxx): + +These errors indicate issues on the server side, +SDK will raise Exceptions for these errors. +""" + +COMMAND_ERROR = 6000, "Command Error" +""" +Command/execution error codes (6xxx): + +These errors are related to command execution and should be handled by the model, +SDK will NOT raise Exceptions for these errors. +""" +``` + +## Retry Strategy Recommendations + +- **Retry trigger**: Only retry when `INTERNAL_SERVER_ERROR` occurs +- **Other error handling**: + - `BAD_REQUEST`: Check if there are issues with the arun call logic + - `COMMAND_ERROR`: stdout goes to `observation.output`, stderr goes to `observation.failure_reason` +- `COMMAND_ERROR` note: When bash execution fails, both stdout and stderr may be non-empty. It is recommended to prompt the model with both output and failure_reason from the observation. + +## Retry Example + +```python +# Background execution with nohup +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 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 +``` diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/deploy.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/deploy.md new file mode 100644 index 0000000000..5fd5b70546 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/deploy.md @@ -0,0 +1,68 @@ +# Deploy + +Sandbox resource deployment manager for local directory deployment and template formatting. + +## deploy_working_dir - Deploy Local Directory + +```python +sandbox = Sandbox(config) +deploy = sandbox.deploy + +# Deploy local directory (auto-generated target path) +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", +) +print(f"Deployed to: {target}") # e.g., /tmp/rock_workdir_abc123 + +# Deploy to specific target path +target = await deploy.deploy_working_dir( + local_path="/path/to/local/project", + target_path="/root/workdir", +) +``` + +## format - Template Variable Substitution + +The `format` method supports two template syntaxes: + +- **`${variable}`** - Standard Python string template syntax +- **`<>`** - Alternative syntax (converted to `${variable}` internally) + +```python +# After deploy_working_dir, use ${working_dir} placeholder +cmd = deploy.format("mv ${working_dir}/config.json /root/.app/") +# Result: mv /tmp/rock_workdir_abc123/config.json /root/.app/ + +# Alternative <<>> syntax +cmd = deploy.format("cat <>/file.txt") +# Result: cat /tmp/rock_workdir_abc123/file.txt + +# Combine with custom variables +cmd = deploy.format( + "cat ${working_dir}/${config_file}", + config_file="settings.json" +) +# Result: cat /tmp/rock_workdir_abc123/settings.json + +# Shell syntax is preserved +cmd = deploy.format("echo $((3 << 2 >> 1))") +# Result: echo $((3 << 2 >> 1)) + +# Access working_dir directly +if deploy.working_dir: + print(f"Current working directory: {deploy.working_dir}") +``` + +## Multiple Deployments + +Subsequent calls overwrite previous working directory paths: + +```python +# First deployment +path1 = await deploy.deploy_working_dir(local_path="/project/v1") +print(deploy.working_dir) # /tmp/rock_workdir_xxx1 + +# Second deployment (overwrites previous path) +path2 = await deploy.deploy_working_dir(local_path="/project/v2") +print(deploy.working_dir) # /tmp/rock_workdir_xxx2 +``` diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/file_system.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/file_system.md new file mode 100644 index 0000000000..a64228e8f1 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/file_system.md @@ -0,0 +1,94 @@ +# FileSystem + +File system interface for sandbox environment operations including permission and ownership management. + +## chown - Change Owner + +```python +from rock.actions.sandbox.request import ChownRequest + +# Create remote user before changing ownership +await sandbox.remote_user.create_remote_user("deploy") + +# Get current working directory +pwd_response = await sandbox.execute(Command(command=["pwd"])) +pwd = pwd_response.stdout.strip() + +# Change directory owner +await sandbox.fs.chown( + ChownRequest( + paths=[pwd], + remote_user="deploy", + recursive=False, + ) +) + +# Recursively change owner for directory and contents +await sandbox.fs.chown( + ChownRequest( + paths=["/home/user/project"], + remote_user="deploy", + recursive=True, + ) +) +``` + +## chmod - Change Permissions + +```python +from rock.actions.sandbox.request import ChmodRequest + +# Create test directory +await sandbox.execute(Command(command=["mkdir", "-p", "/tmp/app"])) + +# Change directory permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="755", + recursive=False, + ) +) + +# Recursively change permissions (includes subdirectories and files) +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/app"], + mode="644", + recursive=True, + ) +) + +# Set maximum permissions +await sandbox.fs.chmod( + ChmodRequest( + paths=["/tmp/shared"], + mode="777", + recursive=True, + ) +) +``` + +## upload_dir - Upload Directory + +```python +import os +from pathlib import Path + +# Prepare local directory +local_dir = Path("/Users/foo/my-project") +(local_dir / "config.json").write_text('{"key": "value"}') +(local_dir / "app.py").write_text("print('hello')") + +# Upload to sandbox +result = await sandbox.fs.upload_dir( + source_dir=str(local_dir), + target_dir="/root/project", + extract_timeout=600, +) + +if result.exit_code == 0: + print(f"Upload success: {result.output}") +else: + print(f"Upload failed: {result.failure_reason}") +``` diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/job.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/job.md new file mode 100644 index 0000000000..a7dd7bcfd1 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/job.md @@ -0,0 +1,150 @@ +# Use Job to Run Agent + +> This is the reference for **Job**, one of ROCK's two parallel ways to use agents. Its core API is `rock.sdk.job.Job` with `JobConfig`, used to run an agent evaluation/task in a sandbox. Two backends are supported: **Bash Job** and **Harbor Bench Job**. +> +> The other way is to install and run an agent inside a single sandbox — see [Install Agent in Sandbox](./rock-agent.md). The two ways use distinct config schemas — **do not mix them**. + +`rock.sdk.job` exposes a single `Job` API that supports two modes, distinguished by the config type: + +- **Bash Job**: Runs an arbitrary shell script inside a sandbox — useful for data processing, external evaluation tools, etc. +- **Harbor Bench Job**: Runs an AI agent benchmark task via the Harbor framework (SWE-bench, Terminal Bench, etc.). + +## End-to-End Example + +A minimal runnable Python snippet: + +```python +import asyncio +from rock.sdk.job import Job, JobConfig + +async def main(): + config = JobConfig.from_yaml("swe_job_config.yaml") # contains agents: and datasets: + result = await Job(config).run() + + print(f"status={result.status}, score={result.score}") + for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + +asyncio.run(main()) +``` + +The full yaml template is in `examples/job/harbor/swe_job_config.yaml.template`. + +--- + +## Bash Job + +Bash Job is for running arbitrary shell scripts inside a sandbox — running an external evaluation tool, processing data, etc. + +Full example: [`examples/job/bash/claw_eval/`](https://github.com/alibaba/ROCK/tree/master/examples/job/bash/claw_eval) + +- `run_claw_eval.py` — Entry point demonstrating `JobConfig.from_yaml()` + `Job(config).run()` +- `claw_eval_bashjob.yaml.template` — YAML template with `script_path`, `environment`, `uploads`, `env`, etc. +- `run_claw_eval.sh` — The script that actually runs in the sandbox (DinD startup, log writing, score output) + +### BashJobConfig Fields + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `script` | `str \| None` | `None` | Inline script content (mutually exclusive with `script_path`) | +| `script_path` | `str \| None` | `None` | Local script path; the file is read and uploaded at runtime | +| `job_name` | `str` | current timestamp | Name used for log and artifact paths | +| `environment` | `EnvironmentConfig` | — | Sandbox connection and resource config (see below) | +| `namespace` | `str \| None` | `None` | Namespace | +| `experiment_id` | `str \| None` | `None` | Experiment ID | +| `timeout` | `int` | `7200` | Overall timeout in seconds (2 hours) | + +**Common `environment` fields:** + +| Field | Type | Description | +|------|------|-------------| +| `image` | `str` | Sandbox Docker image | +| `base_url` | `str` | ROCK platform URL | +| `xrl_authorization` | `str` | Auth token | +| `cluster` | `str` | Target cluster | +| `memory` | `str` | Memory size (e.g. `"64g"`) | +| `cpus` | `int` | Number of CPUs | +| `auto_stop` | `bool` | Whether to stop the sandbox after the job | +| `uploads` | `list` | Local-to-sandbox file/dir uploads, format: `[local_path, sandbox_path]` | +| `env` | `dict[str, str]` | Environment variables injected into the sandbox session | + +--- + +## Harbor Bench Job + +Harbor Bench Job runs AI agent benchmark tasks like SWE-bench and Terminal Bench via the Harbor framework. + +> **Note**: `rock.sdk.bench.Job` is deprecated and will be removed in a future release. Use `rock.sdk.job.Job` + `HarborJobConfig` instead. + +Full example: [`examples/job/harbor/`](https://github.com/alibaba/ROCK/tree/master/examples/job/harbor) + +- `harbor_demo.py` — Entry point demonstrating `JobConfig.from_yaml()` + `Job(config).run()` + result iteration +- `swe_job_config.yaml.template` — SWE-bench task config template +- `swe_job_config-verifier.yaml.template` — Variant with `verifier.mode: native` +- `tb_job_config.yaml.template` — Terminal Bench task config template + +### HarborJobConfig Core Fields + +**Basic fields:** + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `experiment_id` | `str` | required | Experiment ID — required by Harbor | +| `job_name` | `str \| None` | auto-generated | Format: `{dataset}_{task}_{uuid[:8]}` | +| `namespace` | `str \| None` | `None` | Namespace, auto-filled from the sandbox | +| `environment` | `RockEnvironmentConfig` | — | Sandbox connection and resource config | + +**Execution control:** + +| Field | Type | Default | Description | +|------|------|---------|-------------| +| `n_attempts` | `int` | `1` | Attempts per Trial | +| `timeout` | `int` | `7200` | Overall timeout (auto-derived from agent_timeout) | +| `debug` | `bool` | `False` | Debug mode — keeps more intermediate artifacts | + +**Components:** + +| Field | Type | Description | +|------|------|-------------| +| `agents` | `list[AgentConfig]` | Harbor's own agent config (typical fields: `name`, `model_name`) — see `examples/job/harbor/swe_job_config.yaml.template` for the canonical shape | +| `datasets` | `list[DatasetConfig]` | Dataset configs | +| `verifier` | `VerifierConfig` | Verifier evaluation config | +| `orchestrator` | `OrchestratorConfig` | Concurrency / scheduling config | + +--- + +## Result Handling + +Both Job modes return a `JobResult`: + +```python +result = await Job(config).run() + +print(f"status={result.status}, score={result.score}") +for trial in result.trial_results: + print(f" {trial.task_name}: score={trial.score} ({trial.status})") + if trial.exception_info: + print(f" {trial.exception_info.exception_type}: {trial.exception_info.exception_message}") +``` + +### JobResult Fields + +| Field / Property | Type | Description | +|------------------|------|-------------| +| `status` | `JobStatus` | Overall task status | +| `trial_results` | `list[TrialResult]` | List of all Trial results | +| `score` | `float` (property) | Average `score` across all Trials | +| `n_completed` | `int` (property) | Number of Trials with status `completed` | +| `n_failed` | `int` (property) | Number of Trials with status `failed` | + +### TrialResult Fields + +| Field / Property | Type | Description | +|------------------|------|-------------| +| `task_name` | `str` | Task name | +| `exit_code` | `int` | Process exit code | +| `raw_output` | `str` | Raw process output | +| `exception_info` | `ExceptionInfo \| None` | Populated if an exception occurred | +| `status` | `str` (property) | `"completed"` or `"failed"` | +| `duration_sec` | `float` (property) | Execution time in seconds | +| `score` | `float` (property) | Score (Bash Job defaults to `0.0`; Harbor mode comes from the verifier) | diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/model-service.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/model-service.md new file mode 100644 index 0000000000..23dbc21bfc --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/model-service.md @@ -0,0 +1,298 @@ +# Model Service (Experimental) + +The Model Service provided by ROCK is responsible for handling AI model call communications, serving as a communication bridge between agents and training frameworks (such as Roll) or actual LLM inference services. + +## RockAgent Integration + +ModelService is typically **automatically managed by RockAgent** - no manual lifecycle management is required. Simply enable it in the configuration: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, # Enable ModelService, RockAgent manages its lifecycle +) +``` + +RockAgent will automatically: +- Install ModelService (install Python runtime, install model service package) +- Start/stop ModelService +- Monitor Agent process + +## Architecture Overview (Local Mode) + +In local mode, the model service uses the **file system** as the communication medium, implementing a request-response mechanism between agents and models. + +When an agent needs to call a model, the request is first written to a log file, then processed by the listening component. When the model generates a response, the result is written back to the log file and read by the waiting agent. + +## anti_call_llm - Core API + +`anti_call_llm()` is the **most important API in Local mode**, used to manually trigger LLM anti-calls for fine-grained control over model calls: + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='OpenAI type response', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +**Use cases:** +- After Agent captures LLM response, call this method to notify Roll runtime +- Supports carrying response data for error handling or retry +- Configurable timeout and check interval for different network environments + +## CLI Commands + +To use the model service via CLI, ROCK provides a set of CLI commands that can be accessed via `rock model-service` after installing ROCK in the sandbox: + +### start command +Start the model service process +```bash +rock model-service start --type [local|proxy] [options] +``` + +Parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `--type` | str | `local` | Service type: `local` or `proxy` | +| `--config-file` | str | None | Path to configuration file | +| `--host` | str | None | Server host address (overrides config) | +| `--port` | int | None | Server port (overrides config) | +| `--proxy-base-url` | str | None | Proxy base URL | +| `--retryable-status-codes` | str | None | Comma-separated list of retryable status codes | +| `--request-timeout` | int | None | Request timeout in seconds | + +### watch-agent command +Monitor the agent process and send a SESSION_END message when the process exits +```bash +rock model-service watch-agent --pid +``` + +Parameters: +- `--pid`: The ID of the agent process to monitor + +### stop command +Stop the model service +```bash +rock model-service stop +``` + +### anti-call-llm command +Anti-call the LLM interface +```bash +rock model-service anti-call-llm --index [--response ] +``` + +Parameters: +- `--index`: Index of the previous LLM call, starting from 0 +- `--response`: Response from the previous LLM call (optional) + +## File Communication Protocol + +The model service uses files for inter-process communication, defining specific marker formats to distinguish requests and responses: + +### Request Format +``` +LLM_REQUEST_START{JSON request data}LLM_REQUEST_END{metadata JSON} +``` + +### Response Format +``` +LLM_RESPONSE_START{JSON response data}LLM_RESPONSE_END{metadata JSON} +``` + +### Session End Marker +``` +SESSION_END +``` + +Metadata contains timestamp and index information to ensure message order and processing. + +## SDK Usage + +### ModelServiceConfig + +Model service configuration class, located in `rock/sdk/sandbox/model_service/base.py`: + +```python +from rock.sdk.sandbox.model_service.base import ModelServiceConfig + +config = ModelServiceConfig( + enabled=True, + type="local", # Service type + install_cmd="pip install rock-model-service", # Install command + install_timeout=300, # Install timeout (seconds) + start_cmd="rock model-service start --type ${type}", # Start command + stop_cmd="rock model-service stop", # Stop command + logging_path="/data/logs", # Log path + logging_file_name="model_service.log", # Log filename +) +``` + +| Config | Default | Description | +|--------|---------|-------------| +| `enabled` | `False` | Whether to enable model service (RockAgent manages) | +| `type` | `"local"` | Service type: `local` or `proxy` | +| `install_cmd` | - | Model service package install command | +| `install_timeout` | `300` | Install timeout in seconds | +| `start_cmd` | - | Start command template | +| `stop_cmd` | - | Stop command | +| `logging_path` | `/data/logs` | Log directory path | +| `logging_file_name` | `model_service.log` | Log filename | + +### ModelService + +Model service management class, handles the lifecycle of model services within the sandbox: + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.model_service.base import ModelServiceConfig, ModelService + +sandbox = Sandbox(config) +model_service = ModelService(sandbox, ModelServiceConfig()) + +# Typically auto-managed by RockAgent, no manual calls needed +# The following methods are only for manual control when needed + +# Install model service +await model_service.install() + +# Start model service +await model_service.start() + +# Monitor agent process +await model_service.watch_agent(pid="12345") + +# Execute anti-call LLM (Core API for Local mode) +result = await model_service.anti_call_llm( + index=0, + response_payload='{"content": "response"}', + call_timeout=600, + check_interval=3, +) + +# Stop model service +await model_service.stop() +``` + +## API Reference + +### install() + +Install model service dependencies in the sandbox. + +```python +await model_service.install() +``` + +Execution steps: +1. Create and initialize Python runtime environment +2. Create Rock config file +3. Install model service package + +**Note:** Typically auto-called by RockAgent. + +### start() + +Start the model service. + +```python +await model_service.start() +``` + +Prerequisite: Must call `install()` first. + +**Note:** Typically auto-called by RockAgent. + +### stop() + +Stop the model service. + +```python +await model_service.stop() +``` + +If the service is not running, this operation will be skipped. + +**Note:** Typically auto-called by RockAgent. + +### watch_agent(pid) + +Monitor the agent process. + +```python +await model_service.watch_agent(pid="12345") +``` + +Sends `SESSION_END` message when the process exits. + +### anti_call_llm(index, response_payload, call_timeout, check_interval) + +Execute anti-call LLM operation. **This is the most important API in Local mode.** + +```python +result = await model_service.anti_call_llm( + index=0, # LLM call index + response_payload='{"result": "..."}', # Response data (optional) + call_timeout=600, # Operation timeout (seconds) + check_interval=3, # Status check interval (seconds) +) +``` + +## Configuration Options + +### Service Configuration +- `SERVICE_HOST`: Service host address, defaults to `"0.0.0.0"` +- `SERVICE_PORT`: Service port, defaults to `8080` + +### Log Configuration +- `LOG_FILE`: Log file path used for communication, containing request and response data + +### Trajectory (Traj) Logging +The model service records LLM call trajectories (traj) to a JSONL file for debugging and analysis. + +| Environment Variable | Default | Description | +|---------------------|---------|-------------| +| `ROCK_MODEL_SERVICE_DATA_DIR` | `/data/logs` | Directory for traj log files | +| `ROCK_MODEL_SERVICE_TRAJ_APPEND_MODE` | `false` | Append mode (true/false) | + +**Traj file location**: `{DATA_DIR}/LLMTraj.jsonl` + +**Traj file format** (JSONL - one JSON object per line): +```json +{"request": {...}, "response": {...}} +``` + +### Polling Configuration +- `POLLING_INTERVAL_SECONDS`: Polling interval, defaults to `0.1` seconds +- `REQUEST_TIMEOUT`: Request timeout, defaults to unlimited + +### Marker Configuration +Defines markers used to distinguish different types of messages in the log file: +- `REQUEST_START_MARKER` / `REQUEST_END_MARKER` +- `RESPONSE_START_MARKER` / `RESPONSE_END_MARKER` +- `SESSION_END_MARKER` + +### ModelServiceConfig (Server-side) + +The server-side configuration class defines how the model service handles requests: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `host` | str | `"0.0.0.0"` | Server host address | +| `port` | int | `8080` | Server port | +| `proxy_base_url` | str \| None | `None` | Direct proxy URL | +| `proxy_rules` | dict | See below | Model name to URL mapping | +| `retryable_status_codes` | list[int] | `[429, 500]` | Retryable HTTP status codes | +| `request_timeout` | int | `120` | Request timeout in seconds | + +**Default proxy_rules**: +```python +{ + "gpt-3.5-turbo": "https://api.openai.com/v1", + "default": "https://api-inference.modelscope.cn/v1", +} +``` diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/python_sdk.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/python_sdk.md new file mode 100644 index 0000000000..5272d7edb6 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/python_sdk.md @@ -0,0 +1,265 @@ +--- +sidebar_position: 2 +--- + +# Python SDK Reference + +This guide provides detailed information on how to use the ROCK SDK for development, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +ROCK SDK provides developers with convenient Python interfaces to use ROCK platform features, including sandbox environment management and GEM environment interaction. + +> **Important Note**: Before using the SDK, ensure that the ROCK Admin service is running. You can start it with the following command: +> ```bash +> rock admin start +> ``` + +## 2. Sandbox SDK + +### 2.1 Basic Sandbox Operations + +```python +import asyncio + +from rock.actions import CreateBashSessionRequest +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def run_sandbox(): + """Run sandbox demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + Default admin server port is 8080. + """ + # Create sandbox configuration + config = SandboxConfig(image="python:3.11", memory="8g", cpus=2.0) + + # Create sandbox instance + sandbox = Sandbox(config) + + # Start sandbox (connects to admin server) + await sandbox.start() + + # Create session in sandbox for command execution + await sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + + # Execute command in sandbox session + result = await sandbox.arun(cmd="echo Hello ROCK", session="bash-1") + print("\n" + "*" * 50 + "\n" + result.output + "\n" + "*" * 50 + "\n") + + # Stop and clean up sandbox resources + await sandbox.stop() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print("IMPORTANT: Make sure the admin server is running before executing this demo!") + print("Start the admin server with: rock admin start") + asyncio.run(run_sandbox()) +``` + +### 2.2 Sandbox Group Management + +```python +from rock.sdk.sandbox.config import SandboxGroupConfig + +# Create sandbox group configuration +config = SandboxGroupConfig( + image="python:3.11", + size=4, # Create 4 sandboxes + start_concurrency=2, # Concurrency level for startup is 2 +) + +# Create and start sandbox group +sandbox_group = SandboxGroup(config) +await sandbox_group.start() + +# Batch operations +for sandbox in sandbox_group.sandbox_list: + await sandbox.run_in_session(Action(session="default", command="echo Hello")) + +# Batch stop +await sandbox_group.stop() +``` + +### 2.3 Configuration Example + +```python +config = SandboxConfig( + image="python:3.11", + auto_clear_seconds=60 * 20, + experiment_id="test", +) +``` + +### 2.4 Sandbox Speedup Configuration + +ROCK provides sandbox network acceleration capabilities, supporting configuration of APT, PIP, and GitHub mirror sources to improve package download speeds in restricted network environments. + +#### Supported Speedup Types + +**APT Mirror Configuration** + +Configure APT package manager mirror sources for faster Debian/Ubuntu package downloads. + +```python +from rock.sdk.sandbox.speedup import SpeedupType + +# Configure APT mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) +``` + +**PIP Mirror Configuration** + +Configure Python package index mirrors for faster pip installations. + +```python +# HTTP mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="http://mirrors.cloud.aliyuncs.com" +) + +# HTTPS mirror +await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" +) +``` + +**GitHub Acceleration** + +Configure GitHub IP acceleration by adding custom DNS resolution entries. + +```python +await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" +) +``` + +#### Complete Example + +```python +from rock.sdk.sandbox.speedup import SpeedupType +from rock.actions import RunMode + +async def setup_sandbox_with_speedup(): + """Create sandbox and configure acceleration""" + config = SandboxConfig(image="python:3.11") + sandbox = Sandbox(config) + + await sandbox.start() + + # Configure acceleration (before installing packages) + await sandbox.network.speedup( + speedup_type=SpeedupType.APT, + speedup_value="http://mirrors.cloud.aliyuncs.com" + ) + + await sandbox.arun(cmd="apt-get update && apt-get install -y git", mode=RunMode.NOHUP) + + await sandbox.network.speedup( + speedup_type=SpeedupType.PIP, + speedup_value="https://mirrors.aliyun.com" + ) + + # Speedup does not automatically install PIP, it only configures mirror sources for acceleration + await sandbox.arun(cmd="pip install numpy", mode=RunMode.NOHUP) + + # GitHub can be accelerated through mirror IP + await sandbox.network.speedup( + speedup_type=SpeedupType.GITHUB, + speedup_value="11.11.11.11" + ) + + return sandbox +``` + +#### Important Notes + +1. **Configuration Order**: Configure speedup before installing packages +2. **HTTPS vs HTTP**: HTTPS mirrors don't require trusted-host configuration for PIP +3. **GitHub IP**: Different regions may require different IPs for optimal performance +4. **Persistence**: Configurations persist within the sandbox lifecycle +5. **Multiple Calls**: Subsequent speedup calls will override previous configurations +6. **PIP Installation**: The speedup feature only configures mirror sources and does not automatically install PIP + +## 3. GEM SDK + +### 3.1 Python SDK Approach + +```python +import random +import rock + +def main(): + """Main function to run the Sokoban demo with admin server requirement. + + NOTE: This demo requires the admin server to be running for proper execution. + Make sure to start the admin server before running this script. + """ + # Create environment using GEM standard interface + # NOTE: This requires the admin server to be running + env_id = "game:Sokoban-v0-easy" + env = rock.make(env_id) + + # Reset environment to initial state + observation, info = env.reset(seed=42) + print( + "\n" + + "=" * 80 + + "\nInitial Observation:\n" + + str(observation) + + "\n\nInitial Info:\n" + + str(info) + + "\n" + + "=" * 80 + + "\n" + ) + + # Run environment loop until termination + step_count = 0 + while True: + # Interactive environment operation with random actions + action = f"\\boxed{{{random.choice(['up', 'left', 'right', 'down'])}}}" + observation, reward, terminated, truncated, info = env.step(action) + + step_count += 1 + print( + "\n" + + "-" * 80 + + f"\nStep {step_count} - Action: {action}\nReward: {reward}\nObservation:\n{observation}\nInfo: {info}\nTerminated: {terminated}, Truncated: {truncated}\n" + + "-" * 80 + + "\n" + ) + + # Check if environment has reached terminal state + if terminated or truncated: + print("\n" + "=" * 80 + "\nEpisode finished!\n" + "=" * 80 + "\n") + break + + # Clean up environment resources + env.close() + +if __name__ == "__main__": + # Ensure admin server is running before executing + print( + "\n" + + "=" * 80 + + "\nIMPORTANT: Make sure the admin server is running before executing this demo!\nStart the admin server with: rock admin start\n" + + "=" * 80 + + "\n" + ) + main() +``` + +## Related Documents +- [Quick Start Guide](../../Getting%20Started/quickstart.md) - Learn how to quickly get started with the ROCK SDK +- [API Documentation](../api.md) - View the underlying API interfaces encapsulated by the SDK +- [Configuration Guide](../../User%20Guides/configuration.md) - Learn about SDK-related configuration options +- [Installation Guide](../../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/remote_user.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/remote_user.md new file mode 100644 index 0000000000..810bbcf028 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/remote_user.md @@ -0,0 +1,70 @@ +# Remote User + +Remote user management for creating and managing users in the sandbox. + +## Usage Examples + +```python +import asyncio +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.client import Sandbox + +from rock.actions import Action, CreateBashSessionRequest, Observation + + +async def test_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + print('test remote user success') + +async def test_create_session_with_remote_user(): + config = SandboxConfig( + image='hub.docker.alibaba-inc.com/chatos/python:3.11', + xrl_authorization='xxx', + cluster='nt-c' + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.remote_user.create_remote_user('rock') + assert await sandbox.remote_user.is_user_exist('rock') + + await sandbox.create_session(CreateBashSessionRequest(remote_user="rock", session="bash")) + + observation: Observation = await sandbox.run_in_session( + action=Action(session="bash", command="whoami") + ) + print(observation) + assert observation.output.strip() == "rock" + print('test create session with remote user success') + +if __name__ == '__main__': + asyncio.run(test_remote_user()) + asyncio.run(test_create_session_with_remote_user()) +``` + +## API + +### create_remote_user(username) + +Create a remote user. + +```python +await sandbox.remote_user.create_remote_user('username') +``` + +### is_user_exist(username) + +Check if a user exists. + +```python +exists = await sandbox.remote_user.is_user_exist('username') +``` diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/rock-agent.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/rock-agent.md new file mode 100644 index 0000000000..9cc184d46c --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/rock-agent.md @@ -0,0 +1,310 @@ +# Install Agent in Sandbox (Experimental) + +> This is the reference for **install-agent**, one of ROCK's two parallel ways to use agents. Its core API is `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, used to install and run an agent inside a single sandbox. +> +> The other way is to run an agent evaluation/task via Job — see [Use Job to Run Agent](./job.md). The two ways use distinct config schemas. + +RockAgent is the ROCK framework's mechanism for installing a custom agent inside a sandbox. It manages the full agent lifecycle — environment initialization, ModelService integration, command execution, and so on. + +Using `sandbox.agent.install()` and `sandbox.agent.run(prompt)`, you can install and run Agents in the Sandbox environment provided by Rock. + +## Core Concepts + +The core workflow of RockAgent is divided into two phases: + +1. **install(config)**: Initialize the Agent environment, including deploying the working directory, setting environment variables, initializing the runtime environment, etc. +2. **run(prompt)**: Execute the Agent task, replace placeholders, and start the Agent process + +## Quick Start + +### Claude Code Example + +```yaml +run_cmd: "claude -p ${prompt}" + +runtime_env_config: + type: node + custom_install_cmd: "npm install -g @anthropic-ai/claude-code" + +env: + ANTHROPIC_BASE_URL: "" + ANTHROPIC_API_KEY: "" +``` + +### IFlowCli Example + +```yaml +run_cmd: "iflow -p ${prompt} --yolo" # ${prompt} is required + +runtime_env_config: + type: node + custom_install_cmd: "npm i -g @iflow-ai/iflow-cli@latest" + +env: # Environment variables + IFLOW_API_KEY: "xxxxxxx" + IFLOW_BASE_URL: "xxxxxxx" + IFLOW_MODEL_NAME: "xxxxxxx" +``` + +### LangGraph Agent Example + +```yaml +working_dir: "." # Upload local current directory containing langgraph_agent.py to sandbox + +run_cmd: "python langgraph_agent.py ${prompt}" # Run local script + +runtime_env_config: + type: python + pip: # Install pip dependencies + - langchain==1.2.3 + - langchain-openai==1.1.7 + - langgraph==1.0.6 + +env: + OPENAI_API_KEY: xxxxxxx +``` + +## Configuration Details + +### Basic Configuration + +```yaml +agent_type: "default" # Agent type identifier (default: "default") +agent_name: "demo-agent" # Agent instance name (default: random uuid) +version: "1.0.0" # Version identifier (default: "default") +instance_id: "instance-001" # Instance ID (default: "instance-id-") +agent_installed_dir: "/tmp/installed_agent" # Agent installation directory (default: "/tmp/installed_agent") +agent_session: "my-session" # Bash session identifier (default: "agent-session-") +env: # Environment variables (default: {}) + OPENAI_API_KEY: "xxxxxxx" +``` + +### Working Directory Configuration + +```yaml +working_dir: "./my_project" # Local directory to upload to sandbox (default: None, no upload) +project_path: "/testbed" # Working directory in sandbox for cd (default: None) +use_deploy_working_dir_as_fallback: true # Whether to fall back to deploy.working_dir when project_path is empty (default: true) +``` + +### Execution Configuration + +```yaml +run_cmd: "python main.py --prompt ${prompt}" # Agent execution command, must contain ${prompt} (default: None) + +skip_wrap_run_cmd: false # Skip wrapping run_cmd with PATH (default: false) + +# Timeout configuration +agent_install_timeout: 600 # Installation timeout in seconds (default: 600) +agent_run_timeout: 1800 # Run timeout in seconds (default: 1800) +agent_run_check_interval: 30 # Check interval in seconds (default: 30) +``` + +**`skip_wrap_run_cmd`**: +- `false` (default): Wraps the command with `export PATH=:$PATH &&` to ensure runtime environment executables are used +- `true`: Skips PATH wrapping, runs the command directly with `bash -c` + +### Initialization Hooks + +```yaml +pre_init_cmds: # Commands executed before initialization (default: read from env_vars) + - command: "apt update && apt install -y git" + timeout_seconds: 300 # Command timeout in seconds (default: 300) + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: # Commands executed after initialization (default: []) + - command: "echo 'Installation complete'" + timeout_seconds: 30 +``` + +**Notes**: +- `pre_init_cmds` and `post_init_cmds` do not inherit the Agent's `env` environment variables +- Typically used for installation operations and configuration file movement +- Common command examples: + - `apt update && apt install -y git wget tar` + - `cp ${working_dir}/config.json /root/.config/config.json` + +### RuntimeEnv Configuration + +```yaml +runtime_env_config: # Refer to RuntimeEnv documentation for details + type: "python" # Runtime type: python / node (default: "python") + version: "3.11" # Version number + pip: # Python dependency package list + - package1==1.0.0 + - package2==2.0.0 + custom_install_cmd: "git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e ." +``` + +**Node Runtime Example**: + +```yaml +runtime_env_config: + type: "node" + version: "22.18.0" + npm_registry: "https://registry.npmmirror.com" + custom_install_cmd: "npm i -g some-package" +``` + +**Automatic Operations**: +- Install corresponding runtime based on `type` (Python or Node.js) +- Install `pip` dependencies (if configured) +- Execute `custom_install_cmd` custom installation command (if configured) +- Support `npm_registry` configuration for Node.js npm mirror source + +### ModelService Configuration + +```yaml +model_service_config: # Refer to ModelService documentation for details + enabled: true # Enable ModelService (default: false) +``` + +**Automatic Operations**: +- Installation phase: Install ModelService (install only, do not start) +- Run phase: Start ModelService + `watch_agent` monitoring process + +**Notes**: You need to set the model request URL to the ModelService URL. For example, if the ModelService provides an OpenAI-compatible URL at `http://127.0.0.1:8080/v1/chat/completions`, you typically need to set the Agent's LLM request URL to `http://127.0.0.1:8080/v1/`. + +## API Reference + +### install(config) + +Initialize the Agent environment. + +**Execution Flow**: +1. If `working_dir` is configured, deploy to sandbox +2. Set up bash session and configure env environment variables +3. Execute `pre_init_cmds` +4. Initialize RuntimeEnv and ModelService in parallel (if enabled) +5. Execute `post_init_cmds` + +**Parameters**: +- `config`: Agent configuration file, supports two input methods: + - **String path**: YAML configuration file path, default value is `"rock_agent_config.yaml"` + - **RockAgentConfig object**: Directly pass a `RockAgentConfig` instance + +### run(prompt) + +Execute the Agent task. + +**Execution Flow**: +1. Replace placeholders and prepare Agent run command +2. Start the agent process +3. If ModelService is enabled, start `watch_agent` +4. Wait for task completion and return results + +## Advanced Usage + +### Difference and Interaction between working_dir and project_path + +| Configuration | Function | Interaction Method | +|--------------|----------|-------------------| +| `working_dir` | Local directory uploaded to sandbox | Calls `deploy.deploy_working_dir()` to upload, after upload `deploy.working_dir` becomes the path in sandbox | +| `${working_dir}` | Placeholder in commands | Replaced by `deploy.format()` with the value of `deploy.working_dir`, replaced in init_cmds and run_cmd in the configuration | +| `project_path` | Working directory in sandbox | Used for `cd project_path` before running, when not set it enters the `deploy.working_dir` working directory | +| `use_deploy_working_dir_as_fallback` | Whether to fall back to deploy.working_dir when project_path is not set at runtime | Default is `true`, when set to `false` it will not enter working_dir even if project_path is not set | + +**Usage Recommendations**: +- Use `working_dir` to upload local project code to sandbox +- Use `project_path` to specify the working directory in sandbox (e.g., `/testbed`) +- Set `use_deploy_working_dir_as_fallback: false` scenario: Need to perform local file mounting, but want to run Agent in the image's default working directory + +### Placeholder Usage + +Rock Agent supports replacing the following placeholders in the configuration file: + +- `${prompt}`: Required in run_cmd, will be replaced with the prompt passed to `run(prompt)` +- `${working_dir}`: Optional, will be replaced with the actual working directory path in sandbox, also supported in init_cmds and run_cmd +- `${bin_dir}`: Optional, will be replaced with the runtime environment's bin directory path + +**Example**: +```yaml +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" +``` + +### use_deploy_working_dir_as_fallback Explanation + +When `project_path` is not set: +- `true` (default): Before running Agent, it will automatically `cd` to `deploy.working_dir` +- `false`: Before running Agent, it will not automatically switch directories, staying in the current directory + +Applicable Scenarios: +- `true`: Most scenarios, where you want Agent to run in the uploaded code directory +- `false`: Need to mount local files, but want to run Agent in the image's default working directory (e.g., `/app`, `/testbed`) + +## Complete Configuration Example + +```yaml +# ========== Basic Configuration ========== +agent_type: "default" +agent_name: "demo-agent" +version: "1.0.0" +instance_id: "instance-001" +agent_installed_dir: "/tmp/installed_agent" +agent_session: "my-session" +env: + OPENAI_API_KEY: "xxxxxxx" + +# ========== Working Directory Configuration ========== +working_dir: "./my_project" +project_path: "/testbed" +use_deploy_working_dir_as_fallback: true + +# ========== Run Configuration ========== +run_cmd: "python ${working_dir}/main.py --prompt ${prompt}" + +# Timeout configuration +agent_install_timeout: 600 +agent_run_timeout: 1800 +agent_run_check_interval: 30 + +# ========== Initialization Commands ========== +pre_init_cmds: + - command: "apt update && apt install -y git" + timeout_seconds: 300 + - command: "cp ${working_dir}/config.json /root/.config/config.json" + timeout_seconds: 60 + +post_init_cmds: + - command: "echo 'Installation complete'" + timeout_seconds: 30 + +# ========== Runtime Environment Configuration ========== +runtime_env_config: + type: "python" + version: "3.11" + pip: + - langchain==1.2.3 + - langchain-openai==1.1.7 + +# ========== ModelService Integration ========== +model_service_config: + enabled: true +``` + +## Usage Examples + +### Using YAML Configuration File (Recommended) + +```python +import asyncio +from rock.sdk.sandbox import Sandbox, SandboxConfig + +async def main(): + sandbox = Sandbox(SandboxConfig()) + await sandbox.start() + try: + # rock_agent_config.yaml matches the examples in "Quick Start" above + await sandbox.agent.install(config="rock_agent_config.yaml") + result = await sandbox.agent.run(prompt="hello") + print(result) + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +More ready-to-run examples are in `examples/install-agents/` (Claude Code, IFlowCli, Cursor CLI, Qwen Code, SWE-agent, OpenClaw, etc.). + +To run an agent evaluation/benchmark task via Job (a different code path with its own config schema), see [Use Job to Run Agent](./job.md). diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/runtime-env.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/runtime-env.md new file mode 100644 index 0000000000..e1996cfcdb --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/runtime-env.md @@ -0,0 +1,136 @@ +# RuntimeEnv + +The RuntimeEnv module is used to manage language runtime environments in the sandbox (currently providing Python / Node.js). + +## Quick Start (Example) + +```python +from rock.sdk.sandbox import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +sandbox_config = SandboxConfig() +sandbox = Sandbox() +await sandbox.start() + +node_runtime_env_config = NodeRuntimeEnvConfig(version="default") +env = await RuntimeEnv.create(sandbox, node_runtime_env_config) + +await env.run("node --version") +``` + +## RuntimeEnv.create + +An async factory method that creates and initializes a RuntimeEnv instance based on the configuration, and automatically registers it to `sandbox.runtime_envs`. + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, NodeRuntimeEnvConfig + +env = await RuntimeEnv.create( + sandbox, + NodeRuntimeEnvConfig(version="22.18.0"), +) + +# Auto-registered; accessible via sandbox.runtime_envs[env.runtime_env_id] +print(env.runtime_env_id in sandbox.runtime_envs) # True +``` + +## wrapped_cmd + +Wraps a command by adding `bin_dir` to PATH to ensure executables from the runtime environment are used with priority. + +```python +wrapped = env.wrapped_cmd("node script.js") +# Returns: bash -c 'export PATH=/tmp/rock-runtime-envs/node/22.18.0/xxx/runtime-env/bin:$PATH && node script.js' +``` + +## run + +Executes a command within the runtime environment. Internally implemented based on `wrapped_cmd`. + +```python +await env.run("node script.js") +await env.run("npm install express") +``` + +## PythonRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["python"]` | `"python"` | Type identifier | +| `version` | `"3.11" \| "3.12" \| "default"` | `"default"` | Python version; default is 3.11 | +| `pip` | `list[str] \| str \| None` | `None` | List of pip packages or a requirements.txt path | +| `pip_index_url` | `str \| None` | Environment variable | pip index mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["python", "python3", "pip", "pip3"]` | List of executables to symlink | + +## NodeRuntimeEnvConfig + +| Field | Type | Default | Description | +|------|------|--------|------| +| `type` | `Literal["node"]` | `"node"` | Type identifier | +| `version` | `"22.18.0" \| "default"` | `"default"` | Node version; default is 22.18.0 | +| `npm_registry` | `str \| None` | `None` | npm registry mirror | +| `extra_symlink_dir` | `str \| None` | `None` | Target directory for executable symlinks | +| `extra_symlink_executables` | `list[str]` | `["node", "npm", "npx"]` | List of executables to symlink | + +## Constraints for Custom RuntimeEnv Implementations + +A custom RuntimeEnv must follow these rules: + +1. **Define the `runtime_env_type` class attribute**: used as a type identifier for automatic registration into the RuntimeEnv factory +2. **Override `_get_install_cmd()`**: return the install command +3. **The install command must end with**: renaming the directory to `runtime-env` + +## Simplified NodeRuntimeEnv Implementation Example + +```python +from rock.sdk.sandbox.runtime_env import RuntimeEnv, RuntimeEnvConfig +from typing import Literal +from pydantic import Field +from typing_extensions import override + +# Config class: defines the config type so RuntimeEnv.create() can route to the corresponding implementation +class NodeRuntimeEnvConfig(RuntimeEnvConfig): + type: Literal["node"] = "node" # Must match runtime_env_type + +# RuntimeEnv implementation class: defines how to install and run this runtime environment +class NodeRuntimeEnv(RuntimeEnv): + runtime_env_type = "node" # Auto-registered to RuntimeEnv._REGISTRY + + @override + def _get_install_cmd(self) -> str: + # Download the Node binary tarball and extract it, then rename to runtime-env + return ( + "wget -q -O node.tar.xz https://npmmirror.com/mirrors/node/v22.18.0/node-v22.18.0-linux-x64.tar.xz && " + "tar -xf node.tar.xz && " + "mv node-v22.18.0-linux-x64 runtime-env" + ) +``` + +## Speeding Up Base Runtime Installation + +`PythonRuntimeEnv` downloads Python packages from https://github.com/astral-sh/python-build-standalone/releases/ by default. If the network is unavailable or slow, you can override the default install command via `ROCK_RTENV_PYTHON_V31114_INSTALL_CMD` or `ROCK_RTENV_PYTHON_V31212_INSTALL_CMD` (e.g., switch to an internal registry or a mirror). + +Default value example: + +```python +"ROCK_RTENV_PYTHON_V31114_INSTALL_CMD": lambda: os.getenv( + "ROCK_RTENV_PYTHON_V31114_INSTALL_CMD", + "[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; " + "wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz " + "&& tar -xzf cpython31114.tar.gz && mv python runtime-env", +), +``` + +For example, override it to download from a mirror: + +```bash +export ROCK_RTENV_PYTHON_V31114_INSTALL_CMD='[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://mirror.nju.edu.cn/github-release/astral-sh/python-build-standalone/20251209/cpython-3.11.14+20251209-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz && mv python runtime-env' +``` + +Make sure the command creates a `runtime-env` directory under the default working directory of `runtime_env`, and that `${workdir}/runtime-env/bin/` contains the expected executables, e.g.: + +- `${workdir}/runtime-env/bin/python` + +The same applies to Node.js: you can override the install command via `ROCK_RTENV_NODE_V22180_INSTALL_CMD` to use a faster download/install method. \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/sandbox.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/sandbox.md new file mode 100644 index 0000000000..e6e1e43124 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/sandbox.md @@ -0,0 +1,114 @@ +# Handling Large Files and Long Command Outputs + +## `arun` + +`arun()` provides two knobs to control how `nohup` output is handled: + +1. **`response_limited_bytes_in_nohup`** *(integer type)* + Caps the number of characters returned from the nohup output file. Useful when you still need to stream some logs back but want an upper bound (default `None` = no cap). + +2. **`ignore_output`** *(bool, default `False`)* + When set to `True`, `arun()` skips reading the nohup output file entirely. The command still runs to completion and writes logs to `/tmp/tmp_.out`, but the SDK immediately returns a lightweight hint telling agents where to fetch the logs later (via `read_file`, download APIs, or custom commands). This fully decouples "execute command" from "inspect logs". The response also includes the **file size** to help users decide whether to download directly or read in chunks. + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig +from rock.sdk.sandbox.request import CreateBashSessionRequest + +config = SandboxConfig( + image=f"{image}", + xrl_authorization=f"{xrl_authorization}", + user_id=f"{user_id}", + cluster=f"{cluster}", +) +sandbox = Sandbox(config) + +session = sandbox.create_session(CreateBashSessionRequest(session="bash-1")) + +# Example 1: limit the returned logs to 1024 characters +resp_limited = asyncio.run( + sandbox.arun( + cmd="cat /tmp/test.txt", + mode="nohup", + session="bash-1", + response_limited_bytes_in_nohup=1024, + ) +) + +# Example 2: skip collecting logs; agent will download/read them later +resp_detached = asyncio.run( + sandbox.arun( + cmd="bash run_long_job.sh", + mode="nohup", + session="bash-1", + ignore_output=True, + ) +) +print(resp_detached.output) +# Command executed in nohup mode without streaming the log content. +# Status: completed +# Output file: /tmp/tmp_xxx.out +# File size: 15.23 MB +# Use Sandbox.read_file(...), download APIs, or run 'cat /tmp/tmp_xxx.out' ... +``` + +## `read_file_by_line_range` + +Asynchronously reads file content by line range, with built-in support for automatic chunking and session management. Supports large file reading. + +### Key Features +- **Chunked reading for large files**: Automatically splits large files into chunks +- **Automatic line count**: Estimates total lines when end_line is not specified +- **Built-in retry mechanism**: Up to 3 retries for critical operations +- **Input validation**: Validates input parameters automatically +- **Session management**: Supports custom session or auto-created temporary session + +### Parameters +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `file_path` | str | - | File path to read (absolute or relative path in sandbox) | +| `start_line` | int \| None | 1 | Starting line number (1-based) | +| `end_line` | int \| None | None | Ending line number (inclusive), defaults to file end | +| `lines_per_request` | int | 1000 | Lines per request, range 1-10000 | + +### Return Value +- `ReadFileResponse`: Response object containing file content + - `content` (str): The file content read + +### Exception Handling +- `Exception`: Raised when `start_line < 1` +- `Exception`: Raised when `end_line < start_line` +- `Exception`: Raised when `lines_per_request` is not in range 1-10000 +- `Exception`: Raised when file reading fails + +### Usage Examples + +```python +# Read the entire file +response = await sandbox.read_file_by_line_range("/path/to/file.txt") + +# Read a specific line range (lines 100 to 500) +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=100, + end_line=500 +) + +# Read from line 1990 to the end of file +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + start_line=1990 +) + +# Use custom chunk size +response = await sandbox.read_file_by_line_range( + "/path/to/file.txt", + lines_per_request=5000 +) +``` + +### Notes +- Line numbers are 1-based, not 0-based +- For large files, consider increasing `lines_per_request` for better efficiency +- File path must be a valid path within the sandbox +- Uses `sed` command for file reading; ensure the sandbox image supports this command diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md new file mode 100644 index 0000000000..85f34cbe7c --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/swe-bench-evaluation.md @@ -0,0 +1,229 @@ +# SWE-Bench Evaluation + +This guide demonstrates how to use the ROCK SDK to run SWE-Bench Verified evaluations, including sandbox setup, Agent integration, test environment preparation, and result parsing. + +### Quick Start + +SWE-Bench is a benchmark for evaluating AI coding agents on real-world software engineering tasks. + +Running a SWE-Bench task on ROCK involves the following steps: + +1. **load_task_config** — Load `task.yaml` to get the task instruction +2. **start_sandbox** — Start a sandbox with a task-specific Docker image +3. **agent.install / agent.run** — Install and run the Agent to solve the task +4. **setup_test_env** — Upload test files and run-test script to the sandbox +5. **Run tests** — Execute the test script via `sandbox.arun()` with timeout +6. **parse_swebench_result** — Parse test output to determine PASSED / FAILED +7. **sandbox.stop** — Clean up sandbox resources + +**Here is an example code** + +```python +import asyncio +from pathlib import Path + +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def main(): + task_name = "django__django-14539" + task_dir = Path("/root/terminal-bench-datasets/datasets/swebench-verified") / task_name + agent_config_path = "/path/to/iflow_config.yaml" + + # 1. Load task instruction + task_config = await load_task_config(task_dir) # see load_task_config section + instruction = task_config["instruction"] + + # 2. Start sandbox + sandbox = await start_sandbox(task_name) # see start_sandbox section + + try: + # 3. Install and run Agent + await sandbox.agent.install(config=agent_config_path) + result = await sandbox.agent.run(instruction) + + # 4. Setup test environment + await setup_test_env(sandbox, task_dir) # see setup_test_env section + + # 5. Run tests + resp = await run_tests(sandbox) # see Running Tests section + + # 6. Parse results + is_resolved = parse_swebench_result(resp.output) # see parse_swebench_result section + print(f"Task {task_name} resolved: {is_resolved}") + finally: + await sandbox.stop() + +asyncio.run(main()) +``` + +The following sections describe each function used in the workflow in detail. + +--- + +## start_sandbox + +Start a sandbox instance with a task-specific SWE-Bench Docker image. Each task has a pre-built image containing the target repository and environment. + +The `image` parameter follows the format: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.{task_name} +``` + +For example, task `django__django-14539` maps to: + +``` +slimshetty/swebench-verified:sweb.eval.x86_64.django__django-14539 +``` + +```python +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +async def start_sandbox(task_name: str) -> Sandbox: + image = f"slimshetty/swebench-verified:sweb.eval.x86_64.{task_name}" + config = SandboxConfig(image=image) + sandbox = Sandbox(config) + await sandbox.start() + return sandbox +``` + +## load_task_config + +Load task configuration from a `task.yaml` file in the task directory. The YAML file contains the `instruction` field that describes the coding task for the Agent. + +```python +import yaml +from pathlib import Path + +async def load_task_config(task_dir: Path) -> dict: + task_yaml_path = task_dir / "task.yaml" + if not task_yaml_path.exists(): + raise FileNotFoundError(f"task.yaml not found in {task_dir}") + + with open(task_yaml_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + return config + +# Usage +task_config = await load_task_config(task_dir) +instruction = task_config["instruction"] +``` + +## agent.install / agent.run + +Use `sandbox.agent.install()` and `sandbox.agent.run()` to deploy and execute an Agent inside the sandbox. Refer to [Rock Agent](./rock-agent.md) for detailed Agent configuration. + +```python +# Install Agent with a YAML configuration file（e.g., iflow_config.yaml） +await sandbox.agent.install(config="iflow_config.yaml") + +# Run Agent with the task instruction +result = await sandbox.agent.run(instruction) +``` + +## setup_test_env + +Prepare the test environment in the sandbox: install the [uv](https://github.com/astral-sh/uv) package manager, and upload test files and the run-test script. + +```python +from pathlib import Path + +from rock.actions.sandbox.request import CreateBashSessionRequest +from rock.sdk.sandbox.client import RunMode, Sandbox + +async def setup_test_env(sandbox: Sandbox, task_dir: Path) -> str: + """Set up the test environment and return the session name.""" + # 1. Create a session with custom environment variables + session_name = "swe-evaluation" + await sandbox.create_session( + CreateBashSessionRequest( + session=session_name, + env_enable=True, + env={ + "UV_PYTHON_INSTALL_MIRROR": "https://registry.npmmirror.com/-/binary/python-build-standalone" + }, + ) + ) + + # 2. Install uv + for cmd in [ + "wget https://github.com/astral-sh/uv/releases/download/0.10.5/uv-x86_64-unknown-linux-gnu.tar.gz", + "tar -xzf uv-x86_64-unknown-linux-gnu.tar.gz --strip-components=1 -C /usr/local/bin", + ]: + await sandbox.arun(cmd, session=session_name, mode=RunMode.NOHUP) + + # 3. Upload test files + sandbox_test_dir = "/tests" + result = await sandbox.fs.upload_dir(task_dir / "tests", sandbox_test_dir) + if result.exit_code != 0: + raise RuntimeError("Failed to upload test files") + + # 4. Upload run-tests script + run_tests_script = task_dir / "run-tests.sh" + result = await sandbox.upload_by_path( + run_tests_script, + f"{sandbox_test_dir}/{run_tests_script.name}", + ) + if not result.success: + raise RuntimeError("Failed to upload run-tests script") + + return session_name +``` + +## Running Tests + +Execute the test script with a configurable timeout using `RunMode.NOHUP`. + +```python +import shlex +from rock.actions.sandbox.response import Observation +from rock.sdk.sandbox.client import RunMode + +test_timeout_sec = 3600 +sandbox_test_dir = "/tests" + +session_name = "swe-evaluation" + +run_tests_command = f"sh -c 'bash {sandbox_test_dir}/run-tests.sh'" +resp: Observation = await sandbox.arun( + run_tests_command, + session=session_name, + mode=RunMode.NOHUP, + wait_timeout=test_timeout_sec, +) +``` + +## parse_swebench_result + +Parse the test output to determine whether the SWE-Bench task is resolved. The parser looks for a result block delimited by marker lines and checks for `PASSED`. + +```python +import re + +def parse_swebench_result(output: str) -> bool: + """Parse SWE-Bench test output to determine if the task is resolved. + + Matches the block between 'SWEBench results starts here' and + 'SWEBench results ends here', then checks whether it contains 'PASSED'. + """ + match = re.search( + r"SWEBench results starts here\s*(.*?)\s*SWEBench results ends here", + output, + re.DOTALL, + ) + if not match: + return False + return match.group(1).strip() == "PASSED" + +# Usage +is_resolved = parse_swebench_result(resp.output) +``` + +## Notes + +- **Task Datasets**: Task directories (containing `task.yaml`, `tests/`, and `run-tests.sh`) can be obtained from the [terminal-bench-datasets](https://github.com/laude-institute/terminal-bench-datasets) repository. +- **Task Images**: Each SWE-Bench task requires a specific Docker image (e.g., `sweb.eval.x86_64.`). Ensure the image is available before running tests. +- **Agent Config**: The Agent configuration YAML defines the runtime, dependencies, and execution command. See [Rock Agent](./rock-agent.md) for details. + diff --git a/docs/versioned_docs/version-1.8.x/References/api.md b/docs/versioned_docs/version-1.8.x/References/api.md new file mode 100644 index 0000000000..d73bf49d6c --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/api.md @@ -0,0 +1,195 @@ +--- +sidebar_position: 1 +--- + +# API Reference + +This guide provides detailed information about the core API services provided by the ROCK platform, including sandbox environment management and GEM environment interaction. + +## 1. Overview + +The ROCK platform provides two core API services: +- Sandbox API: Sandbox environment management +- GEM API: GEM environment interaction + +All API interfaces follow RESTful design principles and support JSON format data exchange. + +## 2. Sandbox API + +Full lifecycle management functions for sandbox environments: + +### Sandbox Management Interfaces + +1. **Start Sandbox** - Start a sandbox environment + - Create a new sandbox instance + - Support specifying image, resource configuration and other parameters + +2. **Start Sandbox Async** - Asynchronously start a sandbox environment + - Asynchronously create a sandbox instance + - Suitable for scenarios requiring quick response + +3. **Check Sandbox Alive Status** - Check sandbox alive status + - Verify if the sandbox is running normally + +4. **Get Sandbox Statistics** - Get sandbox statistics + - Get resource usage statistics of the sandbox + +5. **Get Sandbox Status** - Get detailed sandbox status + - Get complete status information of the sandbox + +6. **Stop Sandbox** - Stop sandbox environment + - Safely shut down the sandbox instance + +7. **Commit Sandbox** - Commit sandbox as image + - Save current sandbox state as a new image + +### Command Execution Interfaces + +8. **Execute Command** - Execute command in sandbox + - Run specified command directly in the sandbox + +9. **Create Bash Session** - Create Bash session + - Create a persistent Bash session environment + +10. **Run Command in Session** - Run command in session + - Execute command in a created session + +11. **Close Session** - Close session + - Release session resources + +### File Operation Interfaces + +12. **Read File** - Read sandbox file + - Read specified file content from the sandbox + +13. **Write File** - Write sandbox file + - Write file to the sandbox + +14. **Upload File** - Upload file to sandbox + - Upload local file to the sandbox + +## 3. GEM API + +GEM environment interaction functions: + +1. **Make Environment** - Create GEM environment + - Initialize a new GEM environment instance + +2. **Reset Environment** - Reset GEM environment + - Reset GEM environment to initial state + +3. **Step Environment** - Execute GEM environment step + - Execute an action step in the GEM environment + +4. **Close Environment** - Close GEM environment + - Release GEM environment resources + + +## 4. HTTP API Usage Examples + +### 4.1 Sandbox API Examples + +#### Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Asynchronously Start Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/start_async' \ +-H 'Content-Type: application/json' \ +-d '{ + "image": "python:3.11", + "resources": { + "cpu": "2", + "memory": "8g" + } +}' +``` + +#### Execute Command +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/execute' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "command": "ls -la" +}' +``` + +#### Create Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/create_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session" +}' +``` + +#### Run Command in Session +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/run_in_session' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345", + "session": "my_session", + "command": "python script.py" +}' +``` + +#### Upload File +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/upload' \ +-F 'file=@./local_file.txt' \ +-F 'target_path=./remote_file.txt' \ +-F 'sandbox_id=sandbox-12345' +``` + +#### Stop Sandbox +```bash +curl -X POST 'http://localhost:8080/apis/envs/sandbox/v1/stop' \ +-H 'Content-Type: application/json' \ +-d '{ + "sandbox_id": "sandbox-12345" +}' +``` + +### 4.2 GEM API Examples + +```bash +# Create GEM environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/make' \ +-H 'Content-Type: application/json' \ +-d '{"env_id": "game:Sokoban-v0-easy"}' + +# Reset environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/reset' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "seed": 42}' + +# Execute step +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/step' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345", "action": "random_action"}' + +# Close environment +curl -X POST 'http://localhost:8080/apis/v1/envs/gem/close' \ +-H 'Content-Type: application/json' \ +-d '{"sandbox_id": "sandbox-12345"}' +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly get started with ROCK API +- [Python SDK Documentation](./Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to call APIs +- [Configuration Guide](../User%20Guides/configuration.md) - Learn about API-related configuration options +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/Release Notes/index.md b/docs/versioned_docs/version-1.8.x/Release Notes/index.md new file mode 100644 index 0000000000..b31bb727b7 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Release Notes/index.md @@ -0,0 +1,6 @@ +--- +sidebar_position: 1 +--- +# Release Notes + +* [release v1.8.0](v1.8.0.md) diff --git a/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md new file mode 100644 index 0000000000..47a00935c8 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md @@ -0,0 +1,138 @@ +# v1.8.0 + +## Release Date +May 21, 2026 + +--- + +## 🐍 SDK Changes + +### New Features + +* **SDK**: Make sandbox cluster default value configurable via environment variable + +### Bug Fixes + +* **SDK**: Add runtime config type validation in `PythonRuntimeEnv` (#652) + +* **SDK**: `mkdir` target parent dir before `wget` in OSS upload path (#940) + +--- + +## 📦 Sandbox + +### New Features + +* CPU overcommit with grayscale rollout, lifecycle summary, and absolute-cores CPU gauge (#978) + +* `get_status` API supports `include_all_states` parameter — query sandboxes in all states (#951) + +* Account / bucket migration for sandbox upload/download (backward compatible — legacy bucket still supported) (#953) + +* Fix transfer file being stored directly under bucket root (backward compatible) + +--- + +## 🚀 Deployments + +### New Features + +* **Kubernetes**: GPU support with Jinja2 templates and extensible accelerator types (#981) + +* Region cluster config refactor: extract common config, reduce maintenance cost + +* Support cron-based scheduled cleanup of expired logs and directories on ray-head + +### Bug Fixes + +* **Docker**: Cleanup XFS project quota on container stop + +* **Docker**: Remove dead `remove_images` branch + add CLS log support to image removal (#965) + +* Fix `remove_image` missing `cls` parameter bug + +--- + +## 🔀 Proxy + +### New Features + +* Model-service proxy supports streaming and replay — byte passthrough, with Forward and Replay backends (#935) + +--- + +## 📊 Metrics + +### New Features + +* Log OTLP export data points count and duration + +* Add startup timing instrumentation across sandbox launch stages + +### Bug Fixes + +* Read sandbox image from meta store instead of in-memory dict + +* Pass `rock_config` to `SandboxTable` / `SandboxMetaStore` so `MetricsMonitor` uses the correct endpoint + +* Fix `_get_user_info` metrics issue (#911) + +--- + +## ⚙️ Scheduler + +### New Features + +* Dynamic config reloading via Nacos (#888) + +* Add `RayLogCleanupTask` and disable worker-to-driver log forwarding + +* Add `BuildCacheCleanupTask` for pruning uv/pip caches + +* Merge dangling-layer and BuildKit prune into `ImageCleanupTask` (#970) + +* Improve `FileCleanupTask` performance and config safety validation + +### Bug Fixes + +* Handle exception raised by `ray.init` during Ray background reconnection task + +--- + +## 🎯 Rocklet + +### New Features + +* Add Windows PowerShell support (#921) + +### Bug Fixes + +* Mount loop disk to Docker data-root instead of hardcoded path + +* Symlink mount into `/bin` for Nix images with Kata runtime (#936) + +* Use cgroup metrics for container CPU instead of psutil + +--- + +## ✨ Other New Features + +* **CLI**: Add `-v` verbosity control and unify log-level management + +* **Job system**: Integrate in-sandbox model-service proxy with record/replay support + +* **OSS**: Unify dual-account STS token endpoint, push transfer prefix down to SDK + +## 🐛 Other Bug Fixes + +* **BashJob**: Switch OSS upload to script-injection mode — fixes file loss in submit-only mode and env credentials not taking effect + +## ♻️ Refactoring + +* **OSS**: Decouple OSS upload/download from client env vars; implement 3-layer config resolution (#943) [《OSS Refactor Test Plan — ROCK V1.8》](https://alidocs.dingtalk.com/i/nodes/o14dA3GK8gQlkoYwcK5yyo2QV9ekBD76?corpId=dingd8e1123006514592&utm_medium=im_card&iframeQuery=utm_medium%3Dportal_main_colum_create%26utm_source%3Dportal&utm_scene=person_space&utm_source=im&cid=72193729861) + +## 🔧 Build & Tooling + +* Remove `need_database` test marker (#901) + +--- diff --git a/docs/versioned_docs/version-1.8.x/User Guides/configuration.md b/docs/versioned_docs/version-1.8.x/User Guides/configuration.md new file mode 100644 index 0000000000..604256878b --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/User Guides/configuration.md @@ -0,0 +1,189 @@ +--- +sidebar_position: 4 +--- + +# Configuration + +This guide provides detailed instructions on how to configure the ROCK environment to meet different usage requirements, including local development, testing, and production deployment. + +## 1. Environment Variable Configuration + +ROCK supports configuring key parameters through environment variables. The main environment variables are as follows: + +```bash +export ROCK_BASE_URL=http://localhost:8080 # ROCK service base URL +export ROCK_LOG_LEVEL=INFO # Log level +export ROCK_LOGGING_PATH=/path/to/logs # Log file path, default None (output to console) +export ROCK_LOGGING_FILE_NAME=rocklet.log # Log file name, default "rocklet.log", can be customized by admin like admin.log +export ROCK_LOGGING_LEVEL=INFO # Log output level, default "INFO" +export ROCK_WORKER_ENV_TYPE=local # Runtime environment type, options: local, docker, uv, pip +``` + +More environment variables can be found in `rock/env_vars.py`. + +### 1.1 Runtime Environments + +ROCK provides multiple different runtime environments to meet the needs of different scenarios, configured through the `ROCK_WORKER_ENV_TYPE` environment variable. Each environment has different deployment requirements, performance characteristics and applicable scenarios. Each environment has its own unique advantages and limitations, and developers can choose the most suitable runtime environment according to their deployment needs. + +#### 1.1.1 Docker Runtime Environment + +The Docker runtime environment is suitable for Docker image environments where dependencies are pre-installed. This environment requires the `/tmp/miniforge/bin/rocklet` executable to be directly available in the deployment environment. + +**Mount Configuration:** +- `/tmp/miniforge` - Contains pre-installed Python environment +- `/tmp/local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Containerized deployment environments +- Already built custom Docker image containing `rocklet` +- Suitable for production, fast startup + +**Requirements:** +- Requires a custom Docker image containing `/tmp/miniforge/bin/rocklet` executable +- Docker environment support + +#### 1.1.2 Local Runtime Environment + +The local runtime environment directly uses the Python environment and project files of the current deployment. This environment requires the same operating system between the host and container to directly mount the virtual environment and Python interpreter. + +**Mount Configuration:** +- `python_env_path` - Python environment path +- `project_root` - Project root directory +- `.venv` - Virtual environment directory (mounted as `/tmp/miniforge` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run.sh && /tmp/local_files/docker_run.sh +``` + +**Use Cases:** +- Development environments +- Scenarios where host and target container use the same operating system +- Need to quickly reuse existing Python environment + +**Requirements:** +- Same operating system (host/container) +- Direct access to the currently deployed `.venv` virtual environment +- Python interpreter path compatibility + +#### 1.1.3 UV Runtime Environment + +The UV runtime environment only depends on the available ROCK project, but initialization is relatively slow and network requirements are higher. This environment is most suitable for scenarios without preconfigured environments. It rebuilds the rocklet environment from the original project. This is the recommended environment for Mac OS. + +**Mount Configuration:** +- `project_root` - Project root directory (mounted as `/tmp + project_root` in container) +- `local_files` - Local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_uv.sh && /tmp/local_files/docker_run_with_uv.sh '' +``` + +**Use Cases:** +- Mac OS +- Cross-OS startup +- Scenarios without preconfigured environment +- No uv management Rock + +**Advantages:** +- No pre-built image required +- Good cross-platform compatibility +- Suitable for development and testing especially + +**Limitations:** +- Initialization is relatively slow +- Higher network requirements +- Longer startup time + +#### 1.1.4 PIP Runtime Environment + +The PIP runtime environment uses pip to install required dependencies in the container. This environment is suitable for quick setup and scenarios where dependencies can be installed in the container. It is the default runtime environment. It does not require pre-built images containing dependencies, and manages Python packages directly through pip. + +**Mount Configuration:** +- `local_files` - Contains local files required for execution + +**Start Command:** +```bash +chmod +x /tmp/local_files/docker_run_with_pip.sh && /tmp/local_files/docker_run_with_pip.sh +``` + +**Use Cases:** +- ROCK installation from PIP source +- Fast testing of ROCK + +**Advantages:** +- Simple deployment setup + +**Limitations:** +- Long dependency installation time +- Requires network access to install dependency packages +- Dependencies need to be installed each time on startup + +#### 1.1.5 Configuration Guide + +Refer to the following selection guide for different use cases: + +| Scenario | Recommended Environment | Reason | +|----------|--------------------------|-------| +| Production environment | Docker Runtime | Fast startup, stable performance | +| Development environment, same OS | Local Runtime | Environment reuse, fast development cycle | +| Mac development | UV Runtime | Best cross-platform compatibility support | +| Cross-platform development | UV Runtime | Avoids environment compatibility issues | +| Fast testing | UV Runtime | Requires no pre-configuration | +| PIP source installation | PIP Runtime | Install dependencies directly with pip | + +These runtime environments are configured through the `ROCK_WORKER_ENV_TYPE` environment variable, which can be set to "local", "docker", "uv" or "pip". + +### 1.2 Logging Configuration + +Regarding logging configuration, ROCK's logging system has the following characteristics: + +- The logging system cannot output to both file and console simultaneously. If `ROCK_LOGGING_PATH` is set, logs will be output to the designated file, otherwise to console. +- `ROCK_LOGGING_LEVEL` is used to control the output log level, while `ROCK_LOG_LEVEL` is used for general log level settings. + +## 2. Distributed Deployment Requirements + +Since ROCK supports distributed deployment, when running on different nodes of a Ray cluster, the following consistency requirements must be met: + +#### Directory Structure Consistency + +On all Ray nodes, the following directory structure must be completely consistent: +- ROCK project repository directory +- `.venv` virtual environment directory +- The base Python directory that `.venv` depends on + + +#### Mounting Requirements + +ROCK's startup depends on mounting the ROCK project and the corresponding base Python environment, requiring consistency in multi-machine environments: + +#### Verifying Distributed Configuration + +Distributed deployment configuration can be verified through the following methods: + +```bash +# Check directory consistency on all nodes +ls -la /path/to/rock +ls -la /path/to/rock/.venv +ls -la $ROCK_PYTHON_ENV_PATH + +# Verify Python environment availability +$ROCK_PYTHON_ENV_PATH/bin/python --version + +# Check environment variable settings on all nodes +echo $ROCK_PYTHON_ENV_PATH +echo $ROCK_PROJECT_ROOT +``` + +## Related Documents + +- [Quick Start Guide](../Getting%20Started/quickstart.md) - Learn how to quickly set up the ROCK environment +- [API Documentation](../References/api.md) - View sandbox-related API interfaces +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) - Learn how to use the SDK to configure sandboxes +- [Installation Guide](../Getting%20Started/installation.md) - Detailed information about ROCK installation and setup \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/overview.md b/docs/versioned_docs/version-1.8.x/overview.md new file mode 100644 index 0000000000..0c377e8b25 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/overview.md @@ -0,0 +1,33 @@ +--- +sidebar_position: 1 +--- + +# Overview + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework designed to simplify the development, deployment, and management of reinforcement learning environments. + +## What is ROCK + +ROCK (Reinforcement Open Construction Kit) is an open-source reinforcement learning environment development framework. By using ROCK, developers can quickly develop reinforcement learning environments and integrate with other reinforcement learning training frameworks to implement efficient reinforcement learning training. + +ROCK provides comprehensive sandbox environment management capabilities, supports containerized deployment, and enables rapid creation, execution, and destruction of environments. Additionally, ROCK is compatible with the GEM protocol, providing standardized interfaces for reinforcement learning environments. + +## Core Capabilities of ROCK + +1. **Simplified Development Process**: Simplifies the development, construction, and management of reinforcement learning environments, supporting various open-source reinforcement learning environments +2. **Large-scale Scheduling and Deployment**: Enables large-scale scheduling and deployment of rapid reinforcement learning environments. By supporting the GEM protocol, reinforcement learning environments can be easily accessed +3. **Framework Integration**: Integrates with other reinforcement learning training frameworks to achieve large-scale and scalable reinforcement learning training + +## Value of ROCK + +ROCK provides significant value to different roles of engineers: + +- **Reinforcement Learning Algorithm Engineers**: ROCK simplifies the development process of reinforcement learning environments, allowing engineers to focus on algorithm implementation +- **Reinforcement Learning Application Engineers**: ROCK enables large-scale deployment of rapid reinforcement learning environments, improving application development efficiency + +## Learn More + +- [Quick Start Guide](./Getting%20Started/quickstart.md) - Get started with ROCK quickly +- [Configuration Guide](./User%20Guides/configuration.md) - Detailed information about ROCK configuration options +- [API Documentation](./References/api.md) - View ROCK's API interfaces +- [Python SDK Documentation](./References/Python%20SDK%20References/python_sdk.md) - Learn how to use ROCK's Python SDK \ No newline at end of file diff --git a/docs/versioned_sidebars/version-1.8.x-sidebars.json b/docs/versioned_sidebars/version-1.8.x-sidebars.json new file mode 100644 index 0000000000..b475b11530 --- /dev/null +++ b/docs/versioned_sidebars/version-1.8.x-sidebars.json @@ -0,0 +1,64 @@ +{ + "tutorialSidebar": [ + "overview", + { + "type": "category", + "label": "Getting Started", + "link": { + "type": "doc", + "id": "Getting Started/quickstart" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Getting Started" + } + ] + }, + { + "type": "category", + "label": "User Guides", + "items": [ + { + "type": "autogenerated", + "dirName": "User Guides" + } + ] + }, + { + "type": "category", + "label": "References", + "items": [ + "References/api", + { + "type": "category", + "label": "Python SDK References", + "link": { + "type": "doc", + "id": "References/Python SDK References/python_sdk" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "References/Python SDK References" + } + ] + } + ] + }, + { + "type": "category", + "label": "Release Notes", + "link": { + "type": "doc", + "id": "Release Notes/index" + }, + "items": [ + { + "type": "autogenerated", + "dirName": "Release Notes" + } + ] + } + ] +} diff --git a/docs/versions.json b/docs/versions.json index fb1f7d2f57..86a66bc1bf 100644 --- a/docs/versions.json +++ b/docs/versions.json @@ -1,4 +1,5 @@ [ + "1.8.x", "1.7.x", "1.6.x", "1.5.x", From 8e2b9b5adfb1fcd17fc8ff32661a8dee111a201a Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 21 May 2026 16:38:28 +0800 Subject: [PATCH 127/226] fix(sdk): drop wget -c so OSS upload overwrites existing target_path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repeat upload_by_path() to the same sandbox target_path was leaving the sandbox file at its old content while the SDK still reported success=True — a silent data correctness bug. ## Root cause `oss_client.py:upload_via_oss` used: download_cmd = f"wget -c -O {target_path} '{url}'" GNU `wget -c` (continue/resume) skips the download when the local file already exists with size matching the remote object, exiting 0 without fetching. Combined with `_compute_object_name` deriving the OSS key from sha256(sandbox_id|local|sandbox_path) — i.e. path-based, not content-based — a repeat upload to the same target re-uses the same OSS key. resumable_upload DID overwrite the OSS object with new content, but wget on the sandbox side compared sizes, saw they match, and skipped. The downstream `test -f target_path` only verifies file existence, not freshness, so the SDK returned success=True with stale content in the sandbox. This bug is older than PR #949: the same `wget -c -O` line lived at client.py:842 in v1.7.3. PR #949 moved it verbatim to oss_client.py:201 and removed the ROCK_OSS_ENABLE env gate, switching AUTO mode to default-on for files >1MB. The combination promoted a latent corner-case into a path most SDK users hit. Affected real-world flows: model checkpoint overwrites, config hot-reload, training data refresh — anywhere users push to a stable sandbox path repeatedly. ## Fix Drop -c, keep -O. -O alone forces wget to truncate-and-rewrite, regardless of any existing local file. Also wrap target_path in shlex.quote to defend against paths containing spaces or shell metacharacters. ## Test coverage Three regression tests against the wget command shape: - test_wget_command_has_no_continue_flag: asserts ` -c ` is absent and ` -O ` is present in the wget invocation. - test_wget_command_quotes_target_path: target_path with spaces is shell-quoted. - test_repeat_upload_to_same_target_does_not_skip_via_wget: simulates the production overwrite scenario; both rounds must invoke wget without -c. The previous test suite (`test_success`, etc.) used unique target_paths per call, so it never exercised the overwrite path — that's the gap this fix also closes. Co-Authored-By: Claude Opus 4.7 --- rock/sdk/sandbox/oss_client.py | 14 +++- tests/unit/sdk/sandbox/test_oss_client.py | 86 +++++++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/rock/sdk/sandbox/oss_client.py b/rock/sdk/sandbox/oss_client.py index c16306ccfa..c0d7b97e46 100644 --- a/rock/sdk/sandbox/oss_client.py +++ b/rock/sdk/sandbox/oss_client.py @@ -210,8 +210,18 @@ async def upload_via_oss(self, file_path: str, target_path: str) -> UploadRespon mode=RunMode.NORMAL, ) - # wget the signed URL - download_cmd = f"wget -c -O {target_path} '{url}'" + # wget the signed URL. + # Note: NO `-c` (continue/resume). With `-c`, wget skips download + # entirely if the local file already exists with size matching the + # remote object — but the OSS object name is derived from + # (sandbox_id|local_path|sandbox_path), not file content, so a + # repeat upload to the same target re-uses the same OSS key. The + # remote object IS overwritten by `resumable_upload` above (new + # content), but `wget -c` would compare sizes, see they match, and + # exit 0 without fetching, leaving the sandbox file at its old + # content while we still report success. `-O` alone forces wget + # to truncate and rewrite, regardless of existing local state. + download_cmd = f"wget -O {shlex.quote(target_path)} '{url}'" await self._sandbox.arun(cmd=download_cmd, wait_timeout=600, mode=RunMode.NOHUP) # Verify target exists in sandbox. Use execute() instead of arun(), diff --git a/tests/unit/sdk/sandbox/test_oss_client.py b/tests/unit/sdk/sandbox/test_oss_client.py index 66cfde92c1..4844aaab22 100644 --- a/tests/unit/sdk/sandbox/test_oss_client.py +++ b/tests/unit/sdk/sandbox/test_oss_client.py @@ -388,6 +388,92 @@ async def test_oss_upload_exception_returns_failure(self): assert response.success is False + async def test_wget_command_has_no_continue_flag(self): + """Regression: `wget -c` would skip download when target exists with + matching size — but OSS object key is path-based (not content-based), + so a repeat upload re-uses the same key and `-c` would leave the + sandbox file at its old content. We must use `-O` alone to force + truncate-and-rewrite.""" + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-123" + sandbox.arun = AsyncMock(return_value=MagicMock(exit_code=0)) + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=0)) + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._bucket.sign_url = MagicMock(return_value="https://oss/signed?token=xxx") + + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload"): + await client.upload_via_oss("/local/foo.json", "/sandbox/dst/foo.json") + + # arun() is invoked twice: first for `mkdir -p`, second for `wget`. + wget_call = next( + c for c in sandbox.arun.await_args_list + if "wget" in (c.kwargs.get("cmd") or (c.args[0] if c.args else "")) + ) + wget_cmd = wget_call.kwargs.get("cmd") or wget_call.args[0] + # The bug was `wget -c -O ...` — `-c` makes wget skip a same-sized local file. + assert " -c " not in f" {wget_cmd} ", f"wget should not use -c: {wget_cmd!r}" + assert " -O " in wget_cmd, f"wget must use -O to force overwrite: {wget_cmd!r}" + + async def test_wget_command_quotes_target_path(self): + """target_path is interpolated into a shell command; if it contains + spaces or shell metachars (`;`, `$`, `&`, etc.) the unquoted form + would either fail or, worse, execute injected commands. shlex.quote + is the standard fix.""" + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-123" + sandbox.arun = AsyncMock(return_value=MagicMock(exit_code=0)) + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=0)) + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._bucket.sign_url = MagicMock(return_value="https://oss/signed") + + unsafe_target = "/sandbox/dir with spaces/foo.json" + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload"): + await client.upload_via_oss("/local/foo.json", unsafe_target) + + wget_call = next( + c for c in sandbox.arun.await_args_list + if "wget" in (c.kwargs.get("cmd") or (c.args[0] if c.args else "")) + ) + wget_cmd = wget_call.kwargs.get("cmd") or wget_call.args[0] + # shlex.quote wraps a path-with-spaces in single quotes. + assert "'/sandbox/dir with spaces/foo.json'" in wget_cmd, \ + f"target_path with spaces must be shell-quoted: {wget_cmd!r}" + + async def test_repeat_upload_to_same_target_does_not_skip_via_wget(self): + """End-to-end intent: repeat upload(local_v2, /tmp/x) after upload(local_v1, /tmp/x) + must issue a wget command that *will* overwrite — i.e. wget must be invoked + with -O alone (not -c -O). We can't observe sandbox file content in unit + tests, but verifying the command shape forces the correct behavior.""" + sandbox = _make_sandbox() + sandbox.sandbox_id = "sb-fixed" + sandbox.arun = AsyncMock(return_value=MagicMock(exit_code=0)) + sandbox.execute = AsyncMock(return_value=MagicMock(exit_code=0)) + + client = OssClient(sandbox) + client._bucket = MagicMock() + client._bucket.sign_url = MagicMock(return_value="https://oss/v1") + + with patch("rock.sdk.sandbox.oss_client.oss2.resumable_upload"): + # Round 1: upload local_a → /sandbox/x + await client.upload_via_oss("/local/a.bin", "/sandbox/x.bin") + # Round 2: upload local_b (different content, hypothetically) → SAME target + client._bucket.sign_url = MagicMock(return_value="https://oss/v2") + await client.upload_via_oss("/local/b.bin", "/sandbox/x.bin") + + # Both rounds must produce a wget cmd without `-c`. + wget_calls = [ + (c.kwargs.get("cmd") or (c.args[0] if c.args else "")) + for c in sandbox.arun.await_args_list + if "wget" in (c.kwargs.get("cmd") or (c.args[0] if c.args else "")) + ] + assert len(wget_calls) == 2 + for wget_cmd in wget_calls: + assert " -c " not in f" {wget_cmd} ", f"`-c` regressed: {wget_cmd!r}" + class TestDownloadViaOss: async def test_oss_unavailable_returns_failure(self, tmp_path): From cb67e7914ebb0c0d738abc0ba5ab981af9b3d2e5 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Fri, 22 May 2026 11:11:11 +0800 Subject: [PATCH 128/226] docs(v1.8.0): remove internal dingtalk links from release note These point at internal Aliyun documents (alidocs.dingtalk.com) that external readers can't open; remove from both EN and ZH release notes. Co-Authored-By: Claude Opus 4.7 --- .../version-1.8.x/Release Notes/v1.8.0.md | 2 +- docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md index c9a18d7a69..060d386c5c 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md @@ -143,7 +143,7 @@ ## ♻️ 代码重构 -* **对象存储**: 将 OSS 上传/下载与客户端环境变量解耦，实现三层配置解析机制 (#943)[《oss重构提测-ROCK-V1.8》](https://alidocs.dingtalk.com/i/nodes/o14dA3GK8gQlkoYwcK5yyo2QV9ekBD76?corpId=dingd8e1123006514592&utm_medium=im_card&iframeQuery=utm_medium%3Dportal_main_colum_create%26utm_source%3Dportal&utm_scene=person_space&utm_source=im&cid=72193729861) +* **对象存储**: 将 OSS 上传/下载与客户端环境变量解耦，实现三层配置解析机制 (#943) ## 🔧 构建与工具 diff --git a/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md index 47a00935c8..4754f420bb 100644 --- a/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md +++ b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md @@ -129,7 +129,7 @@ May 21, 2026 ## ♻️ Refactoring -* **OSS**: Decouple OSS upload/download from client env vars; implement 3-layer config resolution (#943) [《OSS Refactor Test Plan — ROCK V1.8》](https://alidocs.dingtalk.com/i/nodes/o14dA3GK8gQlkoYwcK5yyo2QV9ekBD76?corpId=dingd8e1123006514592&utm_medium=im_card&iframeQuery=utm_medium%3Dportal_main_colum_create%26utm_source%3Dportal&utm_scene=person_space&utm_source=im&cid=72193729861) +* **OSS**: Decouple OSS upload/download from client env vars; implement 3-layer config resolution (#943) ## 🔧 Build & Tooling From 238ae1c32aaf99793f731faf2b9620b1ed58609d Mon Sep 17 00:00:00 2001 From: Generalwin <52099674+Generalwin@users.noreply.github.com> Date: Fri, 22 May 2026 12:01:15 +0800 Subject: [PATCH 129/226] feat(k8s): Support disk quota limits for K8s operator sandbox (#994) * feat(k8s): k8s operator support disk limits * fix format * feat(k8s): fix code review --- rock/config.py | 1 + rock/sandbox/operator/k8s/provider.py | 34 +++++--- rock/sandbox/operator/k8s/template_loader.py | 3 + rock/utils/format.py | 4 + .../sandbox/operator/test_k8s_provider.py | 48 ++++++++++ .../operator/test_k8s_template_loader.py | 87 +++++++++++++++++++ tests/unit/utils/test_format.py | 5 ++ 7 files changed, 172 insertions(+), 10 deletions(-) diff --git a/rock/config.py b/rock/config.py index cec3899e7a..cc5d7ffc39 100644 --- a/rock/config.py +++ b/rock/config.py @@ -207,6 +207,7 @@ class PoolConfig: image: str cpus: float memory: str + disk: str = "" ports: dict[str, int] = field(default_factory=dict) def __post_init__(self): diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index 86eb875f2f..874fc6f794 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -48,15 +48,15 @@ class ResourceMatchingPoolSelector(PoolSelector): 3. Among all matching pools, select the one with smallest resource capacity """ - def _parse_memory_to_mb(self, memory: str) -> float: - """Parse memory string to MB for comparison.""" - memory = memory.lower().strip() + def _parse_size_to_mb(self, size: str) -> float: + """Parse size string (memory/disk) to MB for comparison.""" + size = size.lower().strip() # Extract number and unit - match = re.match(r"^(\d+(\.\d+)?)\s*([a-z]*)$", memory) + match = re.match(r"^(\d+(\.\d+)?)\s*([a-z]*)$", size) if not match: try: - return float(memory) / (1024 * 1024) # Assume bytes + return float(size) / (1024 * 1024) # Assume bytes except (ValueError, TypeError): return 0 @@ -77,18 +77,23 @@ def _parse_memory_to_mb(self, memory: str) -> float: else: return 0 + # Disk weight factor: disk typically has much larger values than memory, + # divide by this constant to reduce its impact on pool scoring + _DISK_SCORE_WEIGHT_DIVISOR = 100 + def _get_pool_resource_score(self, pool_config: PoolConfig) -> float: """Calculate resource score for a pool (lower is better for selection).""" - memory_mb = self._parse_memory_to_mb(pool_config.memory) - # Normalize memory to GB and add to cpus for a unified score (lower = smaller capacity) - return pool_config.cpus + memory_mb / 1024 + memory_mb = self._parse_size_to_mb(pool_config.memory) + disk_mb = self._parse_size_to_mb(pool_config.disk) if pool_config.disk else 0 + return pool_config.cpus + memory_mb / 1024 + disk_mb / 1024 / self._DISK_SCORE_WEIGHT_DIVISOR def select_pool(self, config: DockerDeploymentConfig, pools: dict[str, PoolConfig]) -> str | None: """Select best matching pool based on image and resource requirements.""" if not pools: return None - config_memory_mb = self._parse_memory_to_mb(config.memory) + config_memory_mb = self._parse_size_to_mb(config.memory) + config_disk_mb = self._parse_size_to_mb(config.disk_limit_rootfs) if config.disk_limit_rootfs else 0 matching_pools: list[tuple[str, PoolConfig, float]] = [] for pool_name, pool_config in pools.items(): @@ -97,10 +102,18 @@ def select_pool(self, config: DockerDeploymentConfig, pools: dict[str, PoolConfi continue # Check resource capacity (pool must have >= required resources) - pool_memory_mb = self._parse_memory_to_mb(pool_config.memory) + pool_memory_mb = self._parse_size_to_mb(pool_config.memory) if pool_config.cpus < config.cpus or pool_memory_mb < config_memory_mb: continue + # Check disk capacity + if config_disk_mb > 0: + if not pool_config.disk: + continue + pool_disk_mb = self._parse_size_to_mb(pool_config.disk) + if pool_disk_mb < config_disk_mb: + continue + # Calculate score for this matching pool score = self._get_pool_resource_score(pool_config) matching_pools.append((pool_name, pool_config, score)) @@ -565,6 +578,7 @@ async def _build_batchsandbox_manifest(self, config: DockerDeploymentConfig) -> image=config.image, cpus=config.cpus, memory=self._normalize_memory(config.memory), + disk=self._normalize_memory(config.disk_limit_rootfs) if config.disk_limit_rootfs else None, num_gpus=config.num_gpus, accelerator_type=config.accelerator_type, ) diff --git a/rock/sandbox/operator/k8s/template_loader.py b/rock/sandbox/operator/k8s/template_loader.py index 5832145089..f11831a46c 100644 --- a/rock/sandbox/operator/k8s/template_loader.py +++ b/rock/sandbox/operator/k8s/template_loader.py @@ -59,6 +59,7 @@ def build_manifest( image: str | None = None, cpus: float | None = None, memory: str | None = None, + disk: str | None = None, num_gpus: int | None = None, accelerator_type: str | None = None, ) -> dict[str, Any]: @@ -84,6 +85,7 @@ def build_manifest( image: Container image (rendered into the template via {{ image }}). cpus: CPU resource value (rendered via {{ cpus }}). memory: Memory resource value (rendered via {{ memory }}). + disk: Disk resource value (rendered via {{ disk }}). num_gpus: GPU count (rendered via {{ num_gpus }}). accelerator_type: GPU model (rendered via {{ accelerator_type }}). @@ -111,6 +113,7 @@ def build_manifest( "image": image if image is not None else "", "cpus": str(cpus) if cpus is not None else "", "memory": memory if memory is not None else "", + "disk": disk if disk is not None else "", "num_gpus": num_gpus if num_gpus is not None else "", "accelerator_type": accelerator_type if accelerator_type is not None else "", } diff --git a/rock/utils/format.py b/rock/utils/format.py index 47f9130c1f..12c028b995 100644 --- a/rock/utils/format.py +++ b/rock/utils/format.py @@ -7,12 +7,16 @@ def parse_size_to_bytes(size_str: str) -> int: "b": 1, "k": 1024, "kb": 1024, + "ki": 1024, "m": 1024**2, "mb": 1024**2, + "mi": 1024**2, "g": 1024**3, "gb": 1024**3, + "gi": 1024**3, "t": 1024**4, "tb": 1024**4, + "ti": 1024**4, } match = re.match(r"^(\d+(?:\.\d+)?)\s*([a-z]+)?$", size_str) diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index 0c81064a6d..785c6da797 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -37,6 +37,7 @@ def make_config( extended_params: dict = None, image_os: str = "linux", num_gpus: float | None = None, + disk_limit_rootfs: str | None = None, ) -> DockerDeploymentConfig: return DockerDeploymentConfig( image=image, @@ -46,6 +47,7 @@ def make_config( extended_params=extended_params or {}, image_os=image_os, num_gpus=num_gpus, + disk_limit_rootfs=disk_limit_rootfs, ) @@ -115,6 +117,52 @@ def test_memory_unit_conversion(self): config = make_config(image="python:3.11", cpus=2, memory="4Gi") assert selector.select_pool(config, pools) == "pool_mi" + def test_skip_pool_when_disk_not_enough(self): + """Return None when pool disk is smaller than required.""" + selector = ResourceMatchingPoolSelector() + pools = { + "pool_small_disk": PoolConfig(image="python:3.11", cpus=2, memory="4Gi", disk="20Gi"), + } + config = make_config(image="python:3.11", cpus=2, memory="4Gi", disk_limit_rootfs="50Gi") + assert selector.select_pool(config, pools) is None + + def test_skip_pool_without_disk_when_disk_required(self): + """Return None when config requires disk but pool has no disk.""" + selector = ResourceMatchingPoolSelector() + pools = { + "pool_no_disk": PoolConfig(image="python:3.11", cpus=2, memory="4Gi"), + } + config = make_config(image="python:3.11", cpus=2, memory="4Gi", disk_limit_rootfs="50Gi") + assert selector.select_pool(config, pools) is None + + def test_select_pool_with_sufficient_disk(self): + """Select pool when disk capacity meets requirement.""" + selector = ResourceMatchingPoolSelector() + pools = { + "pool_disk": PoolConfig(image="python:3.11", cpus=2, memory="4Gi", disk="100Gi"), + } + config = make_config(image="python:3.11", cpus=2, memory="4Gi", disk_limit_rootfs="50Gi") + assert selector.select_pool(config, pools) == "pool_disk" + + def test_select_best_fit_pool_with_disk(self): + """Select pool with smallest cpu+mem+disk when multiple pools match.""" + selector = ResourceMatchingPoolSelector() + pools = { + "pool_large": PoolConfig(image="python:3.11", cpus=8, memory="16Gi", disk="200Gi"), + "pool_exact": PoolConfig(image="python:3.11", cpus=2, memory="4Gi", disk="50Gi"), + } + config = make_config(image="python:3.11", cpus=2, memory="4Gi", disk_limit_rootfs="50Gi") + assert selector.select_pool(config, pools) == "pool_exact" + + def test_no_disk_filter_when_config_has_no_disk(self): + """Pools without disk field are still selectable when config has no disk requirement.""" + selector = ResourceMatchingPoolSelector() + pools = { + "pool_no_disk": PoolConfig(image="python:3.11", cpus=2, memory="4Gi"), + } + config = make_config(image="python:3.11", cpus=2, memory="4Gi") + assert selector.select_pool(config, pools) == "pool_no_disk" + # ========== _get_pool_name ========== diff --git a/tests/unit/sandbox/operator/test_k8s_template_loader.py b/tests/unit/sandbox/operator/test_k8s_template_loader.py index 3e7fdfa259..f7d5997bdf 100644 --- a/tests/unit/sandbox/operator/test_k8s_template_loader.py +++ b/tests/unit/sandbox/operator/test_k8s_template_loader.py @@ -228,3 +228,90 @@ def test_build_manifest_drops_gpu_when_no_gpu(self): assert limits["cpu"] == "2" # GPU placeholders rendered to empty → keys dropped assert "nvidia.com/gpu" not in limits + + def test_build_manifest_with_disk(self): + """Disk placeholder fills correctly when disk is provided.""" + templates = { + "with-disk": { + "ports": {"proxy": 8000, "server": 8080, "ssh": 22}, + "template": { + "spec": { + "containers": [ + { + "name": "main", + "image": "{{ image }}", + "resources": { + "requests": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + "ephemeral-storage": "{{ disk }}", + }, + "limits": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + "ephemeral-storage": "{{ disk }}", + }, + }, + } + ], + } + }, + } + } + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") + + manifest = loader.build_manifest( + template_name="with-disk", + sandbox_id="test-disk", + image="python:3.11", + cpus=2.0, + memory="4Gi", + disk="100Gi", + ) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert container["resources"]["requests"]["ephemeral-storage"] == "100Gi" + assert container["resources"]["limits"]["ephemeral-storage"] == "100Gi" + + def test_build_manifest_drops_disk_when_none(self): + """When disk is None, ephemeral-storage keys are dropped.""" + templates = { + "with-disk": { + "ports": {"proxy": 8000, "server": 8080, "ssh": 22}, + "template": { + "spec": { + "containers": [ + { + "name": "main", + "image": "{{ image }}", + "resources": { + "requests": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + "ephemeral-storage": "{{ disk }}", + }, + "limits": { + "cpu": "{{ cpus }}", + "memory": "{{ memory }}", + "ephemeral-storage": "{{ disk }}", + }, + }, + } + ], + } + }, + } + } + loader = K8sTemplateLoader(templates=templates, default_namespace="rock-test") + + manifest = loader.build_manifest( + template_name="with-disk", + sandbox_id="test-no-disk", + image="python:3.11", + cpus=2.0, + memory="4Gi", + ) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert "ephemeral-storage" not in container["resources"]["requests"] + assert "ephemeral-storage" not in container["resources"]["limits"] diff --git a/tests/unit/utils/test_format.py b/tests/unit/utils/test_format.py index 57a49cfad7..98b1068f85 100644 --- a/tests/unit/utils/test_format.py +++ b/tests/unit/utils/test_format.py @@ -20,6 +20,7 @@ def test_kilobytes(): assert parse_size_to_bytes("1K") == 1024 assert parse_size_to_bytes("1kb") == 1024 assert parse_size_to_bytes("1KB") == 1024 + assert parse_size_to_bytes("1Ki") == 1024 assert parse_size_to_bytes("2k") == 2048 @@ -28,6 +29,7 @@ def test_megabytes(): assert parse_size_to_bytes("1M") == 1024**2 assert parse_size_to_bytes("1mb") == 1024**2 assert parse_size_to_bytes("1MB") == 1024**2 + assert parse_size_to_bytes("1Mi") == 1024**2 assert parse_size_to_bytes("2m") == 2 * 1024**2 @@ -36,6 +38,8 @@ def test_gigabytes(): assert parse_size_to_bytes("1G") == 1024**3 assert parse_size_to_bytes("1gb") == 1024**3 assert parse_size_to_bytes("1GB") == 1024**3 + assert parse_size_to_bytes("1Gi") == 1024**3 + assert parse_size_to_bytes("100Gi") == 100 * 1024**3 assert parse_size_to_bytes("2g") == 2 * 1024**3 @@ -44,6 +48,7 @@ def test_terabytes(): assert parse_size_to_bytes("1T") == 1024**4 assert parse_size_to_bytes("1tb") == 1024**4 assert parse_size_to_bytes("1TB") == 1024**4 + assert parse_size_to_bytes("1Ti") == 1024**4 def test_decimal_values(): From f8b456dd5bd464ac54703a5e8336752c3c9d1da5 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 22 May 2026 15:00:57 +0800 Subject: [PATCH 130/226] fix(admin): retry SandboxTable ops once on stale connection after DB restart (#987) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test(admin): add SandboxTable reconnect tests with real PG process restart Covers the scenario where the postgres process is killed and restarted inside a running container (pg_ctl stop/start), leaving the container port stable but invalidating existing connections. pool_pre_ping=False forces the decorator — not the pool — to handle recovery. * fix(admin): retry SandboxTable ops once on stale connection after DB restart Adds _retry_on_disconnect decorator applied to all six SandboxTable methods. Retries once when DBAPIError.connection_invalidated is True, which SQLAlchemy sets when asyncpg detects "connection is closed" — meaning the query never executed and is safe to retry. Addresses stale connections caused by DB process restart or NAT idle timeout dropping the TCP connection. * test(admin): simulate 3s PG outage to enforce back-off requirement A bare single-attempt retry fires immediately after the DB stops and finds it still down. Only a retry strategy with cumulative back-off exceeding the 3-second outage window can bridge the gap. This makes the test RED against the old no-sleep implementation and GREEN once sufficient exponential back-off is in place. * fix(admin): retry SandboxTable ops with exponential back-off across DB outages The retry decorator now spans both failure modes seen during a PG restart: 1. statement-execution path - an already-checked-out connection goes stale and asyncpg raises sqlalchemy.exc.InterfaceError / OperationalError (DBAPIError subclasses, wrapped by SQLAlchemy's _handle_dbapi_exception). 2. connect path - the pool tries to dial a fresh connection while PG is still down; asyncpg raises ConnectionError / ConnectionResetError / OSError directly. SQLAlchemy does NOT wrap connect-path failures into DBAPIError, so the previous "except DBAPIError" missed this path entirely - retries fired only on the first stale-connection error and then crashed on the second attempt's connect failure. Exception set: (OperationalError, InterfaceError, DisconnectionError, ConnectionError, OSError, asyncio.TimeoutError) Excluded on purpose: DatabaseError - it would swallow IntegrityError / ProgrammingError / DataError, all permanent failures that must fast-fail. ATTEMPTS=4 with exponential back-off (1s, 2s, 4s) gives a cumulative 7s window, sufficient to bridge typical PG process-restart outages. --- rock/admin/core/sandbox_table.py | 58 ++++++ .../core/test_sandbox_table_reconnect.py | 170 ++++++++++++++++++ 2 files changed, 228 insertions(+) create mode 100644 tests/unit/admin/core/test_sandbox_table_reconnect.py diff --git a/rock/admin/core/sandbox_table.py b/rock/admin/core/sandbox_table.py index fb1d0c1949..3db3c939eb 100644 --- a/rock/admin/core/sandbox_table.py +++ b/rock/admin/core/sandbox_table.py @@ -2,9 +2,12 @@ from __future__ import annotations +import asyncio +import functools from typing import TYPE_CHECKING, Any from sqlalchemy import select +from sqlalchemy.exc import DisconnectionError, InterfaceError, OperationalError from sqlalchemy.ext.asyncio import AsyncSession from rock.admin.core.db_provider import DatabaseProvider @@ -21,6 +24,55 @@ logger = init_logger(__name__) +_DISCONNECT_RETRY_ATTEMPTS = 4 + +# Exceptions retried with exponential back-off across DB outages. +# - OperationalError / InterfaceError: SQLAlchemy-wrapped runtime connection +# problems on the statement-execution path (stale connection, server gone, +# socket-level failures observed mid-query). +# - DisconnectionError: explicit pool-level "connection is invalid" signal. +# - OSError / ConnectionError / asyncio.TimeoutError: asyncpg's connect path +# raises these directly; SQLAlchemy does NOT wrap them into DBAPIError +# because they fire before a statement is ever issued. Without catching +# them here, retries cannot bridge a multi-second PG restart window. +# Excluded on purpose: DatabaseError (would swallow IntegrityError, +# DataError, ProgrammingError — all permanent failures that must fast-fail). +_RETRY_EXCEPTIONS: tuple[type[BaseException], ...] = ( + OperationalError, + InterfaceError, + DisconnectionError, + ConnectionError, + OSError, + asyncio.TimeoutError, +) + + +def _retry_on_disconnect(func): + """Retry up to _DISCONNECT_RETRY_ATTEMPTS times across DB outages.""" + + @functools.wraps(func) + async def wrapper(*args, **kwargs): + last_exc: BaseException | None = None + for attempt in range(1, _DISCONNECT_RETRY_ATTEMPTS + 1): + try: + return await func(*args, **kwargs) + except _RETRY_EXCEPTIONS as exc: + last_exc = exc + logger.warning( + "DB connection lost on %s (attempt %d/%d): %r", + func.__name__, + attempt, + _DISCONNECT_RETRY_ATTEMPTS, + exc, + ) + if attempt < _DISCONNECT_RETRY_ATTEMPTS: + await asyncio.sleep(1.0 * 2 ** (attempt - 1)) + assert last_exc is not None + raise last_exc + + return wrapper + + class SandboxTable: """Sandbox-specific database access layer backed by DatabaseProvider. @@ -46,6 +98,7 @@ def __init__(self, db_provider: DatabaseProvider, rock_config: RockConfig | None metric_prefix="meta_store.db", ) + @_retry_on_disconnect @monitor_metastore_operation async def create( self, @@ -76,6 +129,7 @@ async def create( session.add(record) await session.commit() + @_retry_on_disconnect @monitor_metastore_operation async def get(self, sandbox_id: str) -> dict | None: """Return a sandbox row as a plain dict, or ``None`` if not found.""" @@ -85,6 +139,7 @@ async def get(self, sandbox_id: str) -> dict | None: return None return record.to_dict() + @_retry_on_disconnect @monitor_metastore_operation async def update(self, sandbox_id: str, info: SandboxInfo) -> None: """Partial update of scalar columns; always overwrites ``status`` with *info*.""" @@ -100,6 +155,7 @@ async def update(self, sandbox_id: str, info: SandboxInfo) -> None: setattr(record, key, value) await session.commit() + @_retry_on_disconnect @monitor_metastore_operation async def delete(self, sandbox_id: str) -> None: """Hard-delete a sandbox record.""" @@ -109,6 +165,7 @@ async def delete(self, sandbox_id: str) -> None: await session.delete(record) await session.commit() + @_retry_on_disconnect @monitor_metastore_operation async def list_by(self, column: str, value: str | int | float | bool) -> list[dict]: """Equality query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" @@ -120,6 +177,7 @@ async def list_by(self, column: str, value: str | int | float | bool) -> list[di result = await session.execute(stmt) return [r.to_dict() for r in result.scalars().all()] + @_retry_on_disconnect @monitor_metastore_operation async def list_by_in(self, column: str, values: list[str | int | float | bool]) -> list[dict]: """IN query on a single column. Only columns in ``SandboxRecord.LIST_BY_ALLOWLIST`` are permitted.""" diff --git a/tests/unit/admin/core/test_sandbox_table_reconnect.py b/tests/unit/admin/core/test_sandbox_table_reconnect.py new file mode 100644 index 0000000000..b4aafe6ab8 --- /dev/null +++ b/tests/unit/admin/core/test_sandbox_table_reconnect.py @@ -0,0 +1,170 @@ +"""SandboxTable reconnect tests — PostgreSQL process restart inside the container. + +Setup +----- +PID 1 of the container is ``sh`` blocked on ``sleep infinity``. +postgres runs as a background child. ``pg_ctl stop / start`` restarts the +postgres process without touching the container, so the host port stays stable +and data is preserved (same PGDATA directory, new process). +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from rock.admin.core.sandbox_table import SandboxTable + +_PGUSER = "test" +_PGPASS = "test" +_PGDB = "testdb" +_PGDATA = "/var/lib/postgresql/data" + + +def _wait_pg_ready_sql(container, user: str, db: str, timeout: int = 30) -> None: + """Two-stage wait: pg_isready (socket up) then SELECT 1 (queries accepted). + + Mirrors the logic in tests/unit/conftest.py::pg_container to close the + startup race window between the socket accepting and WAL replay finishing. + """ + import time + + deadline = time.time() + timeout + while time.time() < deadline: + code, _ = container.exec_run(f"pg_isready -U {user}") + if code == 0: + code, _ = container.exec_run(f'psql -U {user} -d {db} -c "SELECT 1"') + if code == 0: + return + time.sleep(0.5) + raise TimeoutError(f"PostgreSQL did not become ready within {timeout}s") + + +@pytest.mark.need_docker +class TestSandboxTablePgProcessRestart: + """PostgreSQL process restart inside a running container. + + pool_pre_ping=False so the pool does NOT silently reconnect — the + @_retry_on_disconnect decorator must handle recovery. + """ + + @pytest.fixture + def restartable_pg(self): + """Container where postgres runs as a background child of PID 1 (sleep infinity).""" + import socket + import uuid + + import docker + + client = docker.from_env() + name = f"rock-test-pg-proc-{uuid.uuid4().hex[:8]}" + + hostname = socket.gethostname() + try: + current = client.containers.get(hostname) + networks = current.attrs["NetworkSettings"]["Networks"] + network_name = "bridge" if "bridge" in networks else next(iter(networks), None) + except Exception: + network_name = None + + env = {"POSTGRES_USER": _PGUSER, "POSTGRES_PASSWORD": _PGPASS, "POSTGRES_DB": _PGDB} + run_kwargs = { + "image": "postgres:16-alpine", + "name": name, + "detach": True, + "environment": env, + "entrypoint": ["sh", "-c"], + # Single-element list: Docker passes the whole string as argv[1] to sh -c. + # A plain string would be split on spaces, breaking the & operator. + "command": ["docker-entrypoint.sh postgres & sleep infinity"], + } + if network_name: + run_kwargs["network"] = network_name + else: + run_kwargs["ports"] = {"5432/tcp": None} + + container = client.containers.run(**run_kwargs) + try: + _wait_pg_ready_sql(container, _PGUSER, _PGDB) + container.reload() + if network_name: + host = container.attrs["NetworkSettings"]["Networks"][network_name]["IPAddress"] + port = 5432 + else: + host = "127.0.0.1" + port = int(container.ports["5432/tcp"][0]["HostPort"]) + + yield { + "container": container, + "url": f"postgresql://{_PGUSER}:{_PGPASS}@{host}:{port}/{_PGDB}", + } + finally: + try: + container.stop(timeout=5) + container.remove() + except Exception: + pass + + @pytest.fixture + async def table(self, restartable_pg): + """pool_size=1, pool_pre_ping=False — decorator must handle stale connections.""" + from sqlalchemy.ext.asyncio import create_async_engine + + from rock.admin.core.schema import Base + + url = restartable_pg["url"].replace("postgresql://", "postgresql+asyncpg://") + engine = create_async_engine( + url, + pool_size=1, + max_overflow=0, + pool_pre_ping=False, + connect_args={"statement_cache_size": 0}, + ) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + provider = MagicMock() + provider.engine = engine + t = SandboxTable(provider) + yield t + await engine.dispose() + + _OUTAGE_SECONDS = 4 + + def _do_pg_restart_blocking(self, restartable_pg) -> None: + """Stop, hold the outage open for _OUTAGE_SECONDS, then start. + + Runs in an executor so the asyncio loop stays free to drive the + retry/back-off path inside SandboxTable while PG is down. + Cumulative back-off in the production wrapper (1+2+4+8 = 15s) must + exceed _OUTAGE_SECONDS for the test to pass; a no-sleep retry fires + immediately, finds the DB still down, and fails. + """ + import time + + container = restartable_pg["container"] + container.exec_run(f"su postgres -c 'pg_ctl stop -D {_PGDATA} -m fast'") + time.sleep(self._OUTAGE_SECONDS) + container.exec_run(f"su postgres -c 'pg_ctl start -D {_PGDATA} -l /tmp/pg.log'") + _wait_pg_ready_sql(container, _PGUSER, _PGDB) + + async def test_retry_recovers_after_pg_restart(self, table, restartable_pg): + import asyncio + + await table.create("pgr-1", {"user_id": "bob", "create_time": "2025-01-01T00:00:00Z"}) + await table.list_by_in("sandbox_id", ["pgr-1"]) # warm pool + + # Kick off the outage in the background so the query below races against it. + restart_task = asyncio.create_task(asyncio.to_thread(self._do_pg_restart_blocking, restartable_pg)) + + # Let `pg_ctl stop` actually land and the asyncpg reader observe the RST + # before issuing the query; the query must therefore traverse the outage + # window via the retry decorator's back-off. + await asyncio.sleep(0.5) + + result = await table.list_by_in("sandbox_id", ["pgr-1"]) + + await restart_task + assert len(result) == 1 + assert result[0]["sandbox_id"] == "pgr-1" From 745b100d39c02258f7174fb65ec81042e259435c Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Fri, 22 May 2026 22:30:10 +0800 Subject: [PATCH 131/226] deduplicate region scheduler.tasks via base config inheritance Add `_base: ` resolution in RockConfig.from_env() with deep merge support for dicts and identity-keyed lists. Multi-region YAML configs can now factor out a single base file; previously the `_base` key was silently dropped by the kwargs whitelist, leading to dataclass-default fallbacks (e.g. redis port=0) at runtime. fixes #1004 --- rock/config.py | 70 +++++++++ tests/unit/test_config.py | 294 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 364 insertions(+) diff --git a/rock/config.py b/rock/config.py index cc5d7ffc39..ab68b23d30 100644 --- a/rock/config.py +++ b/rock/config.py @@ -358,6 +358,17 @@ def from_env(cls, config_path: str | None = None): with open(config_file) as f: config = yaml.safe_load(f) + # Handle _base config inheritance + if "_base" in config: + base_path = Path(config.pop("_base")) + if not base_path.is_absolute(): + base_path = config_file.parent / base_path + if not base_path.exists(): + raise Exception(f"base config file {base_path} not found") + with open(base_path) as f: + base_config = yaml.safe_load(f) + config = cls._deep_merge(base_config, config) + # Convert nested dictionaries to dataclass objects kwargs = {} if "ray" in config: @@ -386,6 +397,65 @@ def from_env(cls, config_path: str | None = None): return cls(**kwargs) + # ============================================================================ + # Merging Rules: + # 1. Dictionary elements within the list are matched based on their `task_class`. + # 2. Matched elements: Regional configurations are deep-merged to override the base library (applied at the field level, not as a complete replacement). + # 3. Unmatched base tasks: Retained as-is. + # 4. Newly added regional tasks: Appended to the list. + # 5. To "disable" a specific base task within a region: Set `enabled: false`. + # ============================================================================ + + @staticmethod + def _deep_merge(base: dict, override: dict) -> dict: + """Deep merge override into base. Override values take precedence.""" + result = base.copy() + for key, value in override.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = RockConfig._deep_merge(result[key], value) + elif key in result and isinstance(result[key], list) and isinstance(value, list): + result[key] = RockConfig._merge_lists(result[key], value) + else: + result[key] = value + return result + + @staticmethod + def _merge_lists(base_list: list, override_list: list) -> list: + """Merge two lists. For lists of dicts with 'task_class' key, merge by that key.""" + if not base_list or not override_list: + return override_list if override_list else base_list + + # Check if both lists contain dicts with a common identity key + merge_key = None + for candidate in ("task_class", "name", "id"): + if all(isinstance(item, dict) and candidate in item for item in base_list) and all( + isinstance(item, dict) and candidate in item for item in override_list + ): + merge_key = candidate + break + + if not merge_key: + # No merge key found, override replaces base entirely + return override_list + + # Merge by key: base items are kept/overridden, new override items appended + override_map = {item[merge_key]: item for item in override_list} + result = [] + seen = set() + for item in base_list: + key_val = item[merge_key] + if key_val in override_map: + # Deep merge the matching item + result.append(RockConfig._deep_merge(item, override_map[key_val])) + else: + result.append(item) + seen.add(key_val) + # Append new items from override that aren't in base + for item in override_list: + if item[merge_key] not in seen: + result.append(item) + return result + def __post_init__(self) -> None: logger.info(f"init RockConfig: {self}") diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f6a993e1c8..bc69baf1a8 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,4 +1,5 @@ import tempfile +import textwrap from pathlib import Path import pytest @@ -237,3 +238,296 @@ async def test_resolve_includes_non_mapping_raises(): k8s = {"template_includes": ["bad.yml"]} with pytest.raises(ValueError, match="must be a mapping"): _resolve_k8s_template_includes(k8s, base) + + +# --------------------------------------------------------------------------- +# Unit tests for RockConfig._deep_merge +# --------------------------------------------------------------------------- + + +class TestDeepMerge: + """Tests for RockConfig._deep_merge static method.""" + + def test_disjoint_keys(self): + """Non-overlapping keys are all preserved.""" + base = {"a": 1, "b": 2} + override = {"c": 3} + result = RockConfig._deep_merge(base, override) + assert result == {"a": 1, "b": 2, "c": 3} + + def test_override_scalar(self): + """Override value replaces base for scalar keys.""" + base = {"a": 1, "b": 2} + override = {"b": 99} + result = RockConfig._deep_merge(base, override) + assert result == {"a": 1, "b": 99} + + def test_nested_dict_merge(self): + """Nested dicts are recursively merged.""" + base = {"x": {"a": 1, "b": 2}} + override = {"x": {"b": 20, "c": 30}} + result = RockConfig._deep_merge(base, override) + assert result == {"x": {"a": 1, "b": 20, "c": 30}} + + def test_deeply_nested_merge(self): + """Three-level nested dicts merge correctly.""" + base = {"l1": {"l2": {"a": 1, "b": 2}}} + override = {"l1": {"l2": {"b": 99, "c": 3}}} + result = RockConfig._deep_merge(base, override) + assert result == {"l1": {"l2": {"a": 1, "b": 99, "c": 3}}} + + def test_override_dict_with_scalar(self): + """Override a dict value with a scalar replaces entirely.""" + base = {"a": {"nested": 1}} + override = {"a": "flat"} + result = RockConfig._deep_merge(base, override) + assert result == {"a": "flat"} + + def test_override_scalar_with_dict(self): + """Override a scalar with a dict replaces entirely.""" + base = {"a": "flat"} + override = {"a": {"nested": 1}} + result = RockConfig._deep_merge(base, override) + assert result == {"a": {"nested": 1}} + + def test_empty_base(self): + base = {} + override = {"a": 1} + assert RockConfig._deep_merge(base, override) == {"a": 1} + + def test_empty_override(self): + base = {"a": 1} + override = {} + assert RockConfig._deep_merge(base, override) == {"a": 1} + + def test_both_empty(self): + assert RockConfig._deep_merge({}, {}) == {} + + def test_base_not_mutated(self): + """_deep_merge must not mutate the base dict.""" + base = {"a": {"b": 1}} + override = {"a": {"b": 2}} + RockConfig._deep_merge(base, override) + assert base == {"a": {"b": 1}} + + def test_list_values_delegate_to_merge_lists(self): + """When both values are lists, _merge_lists is invoked.""" + base = {"items": [1, 2]} + override = {"items": [3, 4]} + result = RockConfig._deep_merge(base, override) + # No merge key → override replaces base + assert result == {"items": [3, 4]} + + +# --------------------------------------------------------------------------- +# Unit tests for RockConfig._merge_lists +# --------------------------------------------------------------------------- + + +class TestMergeLists: + """Tests for RockConfig._merge_lists static method.""" + + # --- edge cases: empty inputs --- + + def test_empty_base_returns_override(self): + assert RockConfig._merge_lists([], [{"a": 1}]) == [{"a": 1}] + + def test_empty_override_returns_base(self): + assert RockConfig._merge_lists([{"a": 1}], []) == [{"a": 1}] + + def test_both_empty(self): + assert RockConfig._merge_lists([], []) == [] + + # --- no merge key: override replaces --- + + def test_no_merge_key_plain_values(self): + """Lists of non-dicts → override replaces base.""" + assert RockConfig._merge_lists([1, 2], [3, 4]) == [3, 4] + + def test_no_common_key_in_dicts(self): + """Dicts without a shared identity key → override replaces base.""" + base = [{"foo": 1}] + override = [{"bar": 2}] + assert RockConfig._merge_lists(base, override) == [{"bar": 2}] + + # --- merge by task_class --- + + def test_merge_by_task_class_override(self): + """Matched items are deep-merged by task_class.""" + base = [ + {"task_class": "cleanup", "interval": 60, "enabled": True}, + {"task_class": "report", "interval": 300}, + ] + override = [ + {"task_class": "cleanup", "interval": 120}, + ] + result = RockConfig._merge_lists(base, override) + assert len(result) == 2 + assert result[0] == {"task_class": "cleanup", "interval": 120, "enabled": True} + assert result[1] == {"task_class": "report", "interval": 300} + + def test_merge_by_task_class_append_new(self): + """New items in override are appended.""" + base = [{"task_class": "cleanup", "interval": 60}] + override = [ + {"task_class": "cleanup", "interval": 120}, + {"task_class": "audit", "interval": 600}, + ] + result = RockConfig._merge_lists(base, override) + assert len(result) == 2 + assert result[0]["task_class"] == "cleanup" + assert result[0]["interval"] == 120 + assert result[1] == {"task_class": "audit", "interval": 600} + + def test_merge_by_task_class_disable(self): + """Override can disable a base task via enabled: false.""" + base = [{"task_class": "cleanup", "interval": 60, "enabled": True}] + override = [{"task_class": "cleanup", "enabled": False}] + result = RockConfig._merge_lists(base, override) + assert result[0]["enabled"] is False + assert result[0]["interval"] == 60 # preserved from base + + # --- merge by name --- + + def test_merge_by_name(self): + base = [{"name": "svc-a", "port": 80}] + override = [{"name": "svc-a", "port": 8080}] + result = RockConfig._merge_lists(base, override) + assert result == [{"name": "svc-a", "port": 8080}] + + # --- merge by id --- + + def test_merge_by_id(self): + base = [{"id": "x1", "value": 10}] + override = [{"id": "x1", "value": 20}] + result = RockConfig._merge_lists(base, override) + assert result == [{"id": "x1", "value": 20}] + + # --- key priority: task_class > name > id --- + + def test_merge_key_priority_task_class_over_name(self): + """When both task_class and name exist, task_class is used.""" + base = [{"task_class": "A", "name": "na", "v": 1}] + override = [{"task_class": "A", "name": "nb", "v": 2}] + result = RockConfig._merge_lists(base, override) + assert result[0]["v"] == 2 + assert result[0]["name"] == "nb" # overridden + + # --- nested deep merge within list items --- + + def test_nested_dict_merge_within_list_item(self): + """Dict values inside matched list items are recursively merged.""" + base = [{"task_class": "t1", "params": {"a": 1, "b": 2}}] + override = [{"task_class": "t1", "params": {"b": 20, "c": 30}}] + result = RockConfig._merge_lists(base, override) + assert result[0]["params"] == {"a": 1, "b": 20, "c": 30} + + # --- preserving order --- + + def test_order_preserved(self): + """Base order is preserved; new override items appended at end.""" + base = [ + {"task_class": "B", "v": 1}, + {"task_class": "A", "v": 2}, + ] + override = [ + {"task_class": "C", "v": 3}, + {"task_class": "A", "v": 22}, + ] + result = RockConfig._merge_lists(base, override) + assert [item["task_class"] for item in result] == ["B", "A", "C"] + assert result[1]["v"] == 22 + + +# --------------------------------------------------------------------------- +# Integration test: from_env with _base inheritance +# --------------------------------------------------------------------------- + + +class TestFromEnvBaseInheritance: + """Test RockConfig.from_env _base config inheritance using temp files.""" + + def test_base_inheritance_deep_merges(self, tmp_path: Path): + """Child config inherits and overrides base via _deep_merge.""" + base_file = tmp_path / "base.yml" + base_file.write_text( + textwrap.dedent("""\ + ray: + namespace: "base-ns" + runtime_env: + working_dir: ./ + warmup: + images: + - "python:3.11" + """) + ) + + child_file = tmp_path / "child.yml" + child_file.write_text( + textwrap.dedent("""\ + _base: base.yml + ray: + namespace: "child-ns" + """) + ) + + config = RockConfig.from_env(config_path=str(child_file)) + # Overridden + assert config.ray.namespace == "child-ns" + # Inherited from base (runtime_env is deep-merged: child has no runtime_env so base is kept) + assert config.ray.runtime_env == {"working_dir": "./"} + assert config.warmup.images == ["python:3.11"] + + def test_base_inheritance_scheduler_tasks_merge(self, tmp_path: Path): + """Scheduler tasks list is merged by task_class key.""" + base_file = tmp_path / "base.yml" + base_file.write_text( + textwrap.dedent("""\ + scheduler: + tasks: + - task_class: "rock.admin.scheduler.tasks.cleanup.CleanupTask" + enabled: true + interval_seconds: 60 + - task_class: "rock.admin.scheduler.tasks.report.ReportTask" + enabled: true + interval_seconds: 300 + """) + ) + + child_file = tmp_path / "child.yml" + child_file.write_text( + textwrap.dedent("""\ + _base: base.yml + scheduler: + tasks: + - task_class: "rock.admin.scheduler.tasks.cleanup.CleanupTask" + interval_seconds: 120 + - task_class: "rock.admin.scheduler.tasks.audit.AuditTask" + enabled: true + interval_seconds: 600 + """) + ) + + config = RockConfig.from_env(config_path=str(child_file)) + tasks = config.scheduler.tasks + task_map = {t.task_class: t for t in tasks} + + # Overridden interval + assert task_map["rock.admin.scheduler.tasks.cleanup.CleanupTask"].interval_seconds == 120 + assert task_map["rock.admin.scheduler.tasks.cleanup.CleanupTask"].enabled is True # inherited + # Preserved from base + assert task_map["rock.admin.scheduler.tasks.report.ReportTask"].interval_seconds == 300 + # Newly appended + assert task_map["rock.admin.scheduler.tasks.audit.AuditTask"].interval_seconds == 600 + + def test_base_not_found_raises(self, tmp_path: Path): + child_file = tmp_path / "child.yml" + child_file.write_text( + textwrap.dedent("""\ + _base: nonexistent.yml + ray: + namespace: "ns" + """) + ) + with pytest.raises(Exception, match="base config file.*not found"): + RockConfig.from_env(config_path=str(child_file)) From 127d18e74bb3d07c67a40ac9816fdb3ddea63654 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 26 May 2026 11:43:25 +0800 Subject: [PATCH 132/226] refactor(sandbox): introduce SandboxStateMachine for lifecycle management (#988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add SandboxStateMachine (python-statemachine StateChart) as the single source of truth for sandbox state transitions: States: pending (initial) → running → stopped Transitions: - alive: pending → running (on_alive: update meta store + start_time) - stop: pending/running → stopped (on_stop: stop operator, archive + billing) - stop_noop: stopped → stopped (idempotent, logs skip) SandboxManager changes: - Add _get_current_statemachine(): restore SM from meta store (check_db=True) - stop(): route by SM state — None (dangling) handled inline; STOPPED via stop_noop; else via stop - get_status(): use SM for existence check and pending→running alive transition; remove _update_sandbox_alive_info() (logic moved into on_alive callback) Also add scripts/gen_statemachine_diagram.py for SM visualisation, and tests covering transitions and manager behaviour. Dependencies: python-statemachine>=3.1.1; pydot>=4.0.1 (test group) --- pyproject.toml | 1 + rock/sandbox/sandbox_manager.py | 79 +++---- rock/sandbox/sandbox_statemachine.py | 106 ++++++++++ scripts/gen_statemachine_diagram.py | 12 ++ .../test_get_status_include_all_states.py | 46 +---- .../unit/sandbox/test_sandbox_statemachine.py | 193 ++++++++++++++++++ .../unit/sandbox/test_sandbox_transitions.py | 184 +++++++++++++++++ uv.lock | 11 + 8 files changed, 559 insertions(+), 73 deletions(-) create mode 100644 rock/sandbox/sandbox_statemachine.py create mode 100644 scripts/gen_statemachine_diagram.py create mode 100644 tests/unit/sandbox/test_sandbox_statemachine.py create mode 100644 tests/unit/sandbox/test_sandbox_transitions.py diff --git a/pyproject.toml b/pyproject.toml index 883fd83e70..1f58c2792c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "pyyaml", "jinja2", "tzdata", + "python-statemachine>=3.1.1", ] [project.optional-dependencies] diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 7a6c1316d9..789c267640 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -15,7 +15,6 @@ from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.core.ray_service import RayService -from rock.admin.metrics.billing import log_billing_info from rock.admin.metrics.decorator import monitor_sandbox_operation from rock.admin.proto.request import ClusterInfo, UserInfo from rock.admin.proto.request import SandboxAction as Action @@ -35,6 +34,7 @@ from rock.sandbox.operator.abstract import AbstractOperator from rock.sandbox.sandbox_actor import SandboxActor from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sandbox.sandbox_statemachine import SandboxStateMachine from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError @@ -70,6 +70,13 @@ def __init__( self._proxy_service = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) logger.info("sandbox service init success") + async def _get_current_statemachine(self, sandbox_id: str) -> SandboxStateMachine | None: + """Fetch current state from meta store and return a restored SandboxStateMachine, or None if not found.""" + info = await self._meta_store.get(sandbox_id, check_db=True) + if info is None: + return None + return await SandboxStateMachine.from_state_value(info.get("state"), sandbox_info=info) + async def refresh_aes_key(self): try: await self.rock_config.update() @@ -171,23 +178,20 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: ) @monitor_sandbox_operation() - async def stop(self, sandbox_id, reason: StopReason = StopReason.MANUAL): - logger.info(f"stop sandbox {sandbox_id} (reason={reason.value})") - sandbox_info: SandboxInfo | None = await self._meta_store.get(sandbox_id) - if sandbox_info is None: - sandbox_info = {} - sandbox_info["state"] = State.STOPPED - if sandbox_info.get("start_time"): - sandbox_info["stop_time"] = get_iso8601_timestamp() - log_billing_info(sandbox_info=sandbox_info) - try: - await self._operator.stop(sandbox_id, reason) - except ValueError as e: - logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) + async def stop(self, sandbox_id: str): + sm = await self._get_current_statemachine(sandbox_id) + if sm is None: + logger.info(f"stop dangling sandbox {sandbox_id}") + sandbox_info: SandboxInfo = {"state": State.STOPPED} + try: + await self._operator.stop(sandbox_id) + except ValueError as e: + logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) await self._meta_store.archive(sandbox_id, sandbox_info) - return - logger.info(f"sandbox {sandbox_id} stopped") - await self._meta_store.archive(sandbox_id, sandbox_info) + elif sm.current_state.value == State.STOPPED: + await sm.send("stop_noop", sandbox_id=sandbox_id) + else: + await sm.send("stop", sandbox_id=sandbox_id, operator=self._operator, meta_store=self._meta_store) async def get_mount(self, sandbox_id): async with self._ray_service.get_ray_rwlock().read_lock(): @@ -216,23 +220,33 @@ async def commit(self, sandbox_id, image_tag: str, username: str, password: str) @monitor_sandbox_operation() async def get_status(self, sandbox_id, include_all_states: bool = False) -> SandboxStatusResponse: - is_alive = False + # get status from meta_store + sm = await self._get_current_statemachine(sandbox_id) + if sm is None: + raise BadRequestRockError(f"Sandbox {sandbox_id} not found") - sandbox_info: SandboxInfo | None = await self._operator.get_status(sandbox_id=sandbox_id) - if sandbox_info is not None: - is_alive = sandbox_info.get("state") == State.RUNNING - self._update_sandbox_alive_info(sandbox_info, is_alive) - if sandbox_info.get("state") in (State.PENDING, State.RUNNING): - current = await self._meta_store.get(sandbox_id) - if current is None or current.get("state") != sandbox_info.get("state"): - await self._meta_store.update(sandbox_id, sandbox_info) + # update status from operator + is_alive = False + operator_sandbox_info: SandboxInfo | None = await self._operator.get_status(sandbox_id=sandbox_id) + if operator_sandbox_info is not None: + is_alive = operator_sandbox_info.get("state") == State.RUNNING + if sm.current_state.value == State.PENDING and is_alive: + await sm.send( + "alive", sandbox_id=sandbox_id, meta_store=self._meta_store, sandbox_info=operator_sandbox_info + ) + if operator_sandbox_info.get("state") in (State.PENDING, State.RUNNING): await self._refresh_timeout(sandbox_id) - elif include_all_states: - sandbox_info = await self._meta_store.get(sandbox_id, check_db=True) - if sandbox_info is None: + # compat with legacy get_status behavior by default (include_all_states == False), + # raise 'not found' if not on pending or running status. + if not include_all_states and sm.current_state.value not in (State.PENDING, State.RUNNING): raise BadRequestRockError(f"Sandbox {sandbox_id} not found") + if operator_sandbox_info is not None: + sandbox_info = operator_sandbox_info + else: + sandbox_info = sm.sandbox_info + return SandboxStatusResponse( sandbox_id=sandbox_id, status=sandbox_info.get("phases"), @@ -270,13 +284,6 @@ async def build_sandbox_info_from_redis(self, sandbox_id: str, deployment_info: sandbox_info = deployment_info return sandbox_info - def _update_sandbox_alive_info(self, sandbox_info: SandboxInfo, is_alive: bool) -> None: - if is_alive: - sandbox_info["state"] = State.RUNNING - # Set start_time for the first time the sandbox becomes alive - if sandbox_info.get("start_time") is None: - sandbox_info["start_time"] = get_iso8601_timestamp() - async def get_status_v2(self, sandbox_id, include_all_states: bool = False) -> SandboxStatusResponse: """ Deprecated: Use get_status(sandbox_id, use_rocklet=True) instead. diff --git a/rock/sandbox/sandbox_statemachine.py b/rock/sandbox/sandbox_statemachine.py new file mode 100644 index 0000000000..87bfb5069f --- /dev/null +++ b/rock/sandbox/sandbox_statemachine.py @@ -0,0 +1,106 @@ +""" +Sandbox State Machine using python-statemachine library. + +Defines sandbox lifecycle states and transitions. The on_start / on_stop +callbacks contain the actual async business logic so the state machine is +the single place that owns both transition validation and execution. +""" + +from statemachine import State as SMState +from statemachine import StateChart + +from rock.actions.sandbox.response import State as RockState +from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.admin.metrics.billing import log_billing_info +from rock.logger import init_logger +from rock.utils.system import get_iso8601_timestamp + +logger = init_logger(__name__) + + +class SandboxStateMachine(StateChart): + """ + State machine for sandbox lifecycle management. + + States: + - pending: Sandbox is being created / starting + - running: Sandbox is actively running + - stopped: Sandbox has been stopped + + Transitions: + - stop: pending/running → stopped (stops operator, archives meta) + - stop_noop: stopped → stopped (idempotent; logs and returns) + - alive: pending → running (called from get_status on pending→running; also usable by reconciler) + """ + + allow_event_without_transition = False # raise TransitionNotAllowed instead of silently ignoring invalid events + catch_errors_as_events = False + + # States + pending = SMState("Pending", initial=True, value=RockState.PENDING) + running = SMState("Running", value=RockState.RUNNING) + stopped = SMState("Stopped", value=RockState.STOPPED) + + # Transitions + stop = pending.to(stopped) | running.to(stopped) + stop_noop = stopped.to(stopped) + alive = pending.to(running) + + def __init__(self, **kwargs): + """Initialize with optional sandbox_info.""" + super().__init__(**kwargs) + self.sandbox_info: SandboxInfo | None = kwargs.get("sandbox_info") + + # Callbacks + + async def on_stop(self, sandbox_id: str, operator, meta_store) -> None: + logger.info(f"stop sandbox {sandbox_id}") + sandbox_info = self.sandbox_info or {} + + # Initialize sandbox_info with default values if not set + if "sandbox_id" not in sandbox_info: + sandbox_info["sandbox_id"] = sandbox_id + + sandbox_info["state"] = RockState.STOPPED + if sandbox_info.get("start_time"): + sandbox_info["stop_time"] = get_iso8601_timestamp() + log_billing_info(sandbox_info=sandbox_info) + + try: + await operator.stop(sandbox_id) + except ValueError as e: + logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) + + logger.info(f"sandbox {sandbox_id} stopped") + await meta_store.archive(sandbox_id, sandbox_info) + + # Update self.sandbox_info for potential future use + self.sandbox_info = sandbox_info + + async def on_stop_noop(self, sandbox_id: str) -> None: + logger.info(f"Sandbox {sandbox_id} already stopped, skipping") + + async def on_alive(self, sandbox_id: str, meta_store, sandbox_info: SandboxInfo) -> None: + sandbox_info["state"] = RockState.RUNNING + if not sandbox_info.get("start_time"): + sandbox_info["start_time"] = get_iso8601_timestamp() + await meta_store.update(sandbox_id, sandbox_info) + + # Update self.sandbox_info for potential future use + self.sandbox_info = sandbox_info + + @classmethod + async def from_state_value(cls, state_value: str | None, sandbox_info: SandboxInfo) -> "SandboxStateMachine": + """Create a state machine restored to *state_value* (from Redis/DB).""" + state_map = { + RockState.PENDING: "pending", + RockState.RUNNING: "running", + RockState.STOPPED: "stopped", + } + sm = ( + cls(start_value=state_map[state_value], sandbox_info=sandbox_info) + if state_value and state_value in state_map + else cls(sandbox_info=sandbox_info) + ) + await sm.activate_initial_state() + return sm diff --git a/scripts/gen_statemachine_diagram.py b/scripts/gen_statemachine_diagram.py new file mode 100644 index 0000000000..c618c9be52 --- /dev/null +++ b/scripts/gen_statemachine_diagram.py @@ -0,0 +1,12 @@ +"""Generate a PNG diagram for SandboxStateMachine. + +Usage: + uv run python scripts/gen_statemachine_diagram.py +""" + +from statemachine.contrib.diagram import DotGraphMachine + +from rock.sandbox.sandbox_statemachine import SandboxStateMachine + +DotGraphMachine(SandboxStateMachine)().write_png("sandbox_statemachine.png") +print("Written to sandbox_statemachine.png") diff --git a/tests/unit/sandbox/test_get_status_include_all_states.py b/tests/unit/sandbox/test_get_status_include_all_states.py index abd67896f0..1a25340c2b 100644 --- a/tests/unit/sandbox/test_get_status_include_all_states.py +++ b/tests/unit/sandbox/test_get_status_include_all_states.py @@ -1,12 +1,10 @@ """ -Unit tests for SandboxManager.get_status changes in feat(sandbox): support include_all_states. +Unit tests for SandboxManager.get_status with include_all_states support. Key behaviour changes covered: - operator.get_status() may now return None - - include_all_states=False + operator None → raise BadRequestRockError("not found") - - include_all_states=True + operator None → fall back to meta_store.get(check_db=True) - - include_all_states=True + operator data → skip fallback, normal path - - meta_store.update only when state is PENDING or RUNNING + - include_all_states=True + operator data → skip extra DB fallback, normal path + - meta_store.update only when state transitions PENDING → RUNNING - start_time/stop_time/create_time populated in every response """ @@ -17,7 +15,6 @@ from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.proto.response import SandboxStatusResponse -from rock.sdk.common.exceptions import BadRequestRockError def _make_sandbox_info(sandbox_id: str = "sandbox-1", state: State = State.RUNNING) -> SandboxInfo: @@ -46,8 +43,9 @@ def mock_meta_store(): @pytest.fixture -def sandbox_manager(mock_operator, mock_meta_store, rock_config): +async def sandbox_manager(mock_operator, mock_meta_store, rock_config): from rock.sandbox.sandbox_manager import SandboxManager + from rock.sandbox.sandbox_statemachine import SandboxStateMachine with patch("rock.sandbox.sandbox_manager.SandboxProxyService"): manager = SandboxManager.__new__(SandboxManager) @@ -55,6 +53,10 @@ def sandbox_manager(mock_operator, mock_meta_store, rock_config): manager._operator = mock_operator manager._meta_store = mock_meta_store manager._refresh_timeout = AsyncMock() + + mock_sm = await SandboxStateMachine.from_state_value(State.PENDING, sandbox_info={}) + manager._get_current_statemachine = AsyncMock(return_value=mock_sm) + return manager @@ -79,36 +81,6 @@ async def test_running_state_triggers_meta_store_update(self, sandbox_manager, m mock_meta_store.update.assert_awaited_once() - @pytest.mark.asyncio - async def test_operator_none_flag_false_raises_not_found(self, sandbox_manager, mock_operator, mock_meta_store): - """operator=None + include_all_states=False → not found, no DB fallback triggered.""" - mock_operator.get_status = AsyncMock(return_value=None) - - with pytest.raises(BadRequestRockError, match="not found"): - await sandbox_manager.get_status("sandbox-1", include_all_states=False) - - for c in mock_meta_store.get.call_args_list: - assert not c.kwargs.get("check_db") - - @pytest.mark.asyncio - async def test_operator_none_flag_true_calls_db_fallback(self, sandbox_manager, mock_operator, mock_meta_store): - """operator=None + include_all_states=True → meta_store.get(check_db=True) called.""" - mock_operator.get_status = AsyncMock(return_value=None) - mock_meta_store.get = AsyncMock(return_value=_make_sandbox_info(state=State.PENDING)) - - result = await sandbox_manager.get_status("sandbox-1", include_all_states=True) - - mock_meta_store.get.assert_awaited_once_with("sandbox-1", check_db=True) - assert result.state == State.PENDING - - @pytest.mark.asyncio - async def test_operator_none_flag_true_db_empty_raises(self, sandbox_manager, mock_operator): - """operator=None + include_all_states=True + DB empty → not found.""" - mock_operator.get_status = AsyncMock(return_value=None) - - with pytest.raises(BadRequestRockError, match="not found"): - await sandbox_manager.get_status("sandbox-1", include_all_states=True) - @pytest.mark.asyncio async def test_operator_data_with_flag_true_skips_db_fallback( self, sandbox_manager, mock_operator, mock_meta_store diff --git a/tests/unit/sandbox/test_sandbox_statemachine.py b/tests/unit/sandbox/test_sandbox_statemachine.py new file mode 100644 index 0000000000..435f37ff7e --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_statemachine.py @@ -0,0 +1,193 @@ +""" +Unit tests for SandboxStateMachine. + +Covers: +- State transitions (valid and invalid) +- State.active properties for querying state +- State restoration via from_state_value() +- Async action callbacks: on_stop +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from rock.actions.sandbox.response import State +from rock.sandbox.sandbox_statemachine import SandboxStateMachine + +# --------------------------------------------------------------------------- +# Transitions +# --------------------------------------------------------------------------- + + +class TestTransitions: + @pytest.mark.asyncio + async def test_initial_state_is_pending(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + assert sm.pending.is_active + assert not sm.running.is_active + assert not sm.stopped.is_active + + @pytest.mark.asyncio + async def test_alive_goes_to_running(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + await sm.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + assert sm.running.is_active + + @pytest.mark.asyncio + async def test_stop_from_pending(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + await sm.send("stop", sandbox_id="sb", operator=AsyncMock(), meta_store=AsyncMock()) + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_stop_from_running(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + await sm.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + await sm.send("stop", sandbox_id="sb", operator=AsyncMock(), meta_store=AsyncMock()) + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_stop_noop_from_stopped(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + await sm.send("stop", sandbox_id="sb", operator=AsyncMock(), meta_store=AsyncMock()) + await sm.send("stop_noop", sandbox_id="sb") + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_full_lifecycle(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + assert sm.pending.is_active + await sm.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + assert sm.running.is_active + await sm.send("stop", sandbox_id="sb", operator=AsyncMock(), meta_store=AsyncMock()) + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_multiple_instances_are_independent(self): + sm1, sm2 = SandboxStateMachine(), SandboxStateMachine() + await sm1.activate_initial_state() + await sm2.activate_initial_state() + await sm1.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + assert sm1.running.is_active + assert sm2.pending.is_active + + +# --------------------------------------------------------------------------- +# State query helpers +# --------------------------------------------------------------------------- + + +class TestStateHelpers: + @pytest.mark.asyncio + async def test_state_active_properties_track_state(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + assert sm.pending.is_active and not sm.running.is_active + + await sm.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + assert sm.running.is_active + + await sm.send("stop", sandbox_id="sb", operator=AsyncMock(), meta_store=AsyncMock()) + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_repr(self): + sm = SandboxStateMachine() + await sm.activate_initial_state() + assert "pending" in repr(sm) + await sm.send("alive", sandbox_id="sb", meta_store=AsyncMock(), sandbox_info={}) + assert "running" in repr(sm) + + +# --------------------------------------------------------------------------- +# State restoration +# --------------------------------------------------------------------------- + + +class TestFromStateValue: + @pytest.mark.asyncio + async def test_none_starts_in_pending(self): + sm = await SandboxStateMachine.from_state_value(None, sandbox_info={}) + assert sm.pending.is_active + + @pytest.mark.asyncio + async def test_restores_pending(self): + sm = await SandboxStateMachine.from_state_value(State.PENDING, sandbox_info={}) + assert sm.pending.is_active + + @pytest.mark.asyncio + async def test_restores_running(self): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + assert sm.running.is_active + + @pytest.mark.asyncio + async def test_restores_stopped(self): + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info={}) + assert sm.stopped.is_active + + @pytest.mark.asyncio + async def test_unknown_value_defaults_to_pending(self): + sm = await SandboxStateMachine.from_state_value("bogus", sandbox_info={}) + assert sm.pending.is_active + + +# --------------------------------------------------------------------------- +# on_stop callback +# --------------------------------------------------------------------------- + + +class TestOnStop: + @pytest.fixture + def mock_operator(self): + return AsyncMock() + + @pytest.fixture + def mock_meta_store(self): + store = AsyncMock() + store.get = AsyncMock(return_value={"state": State.RUNNING}) + return store + + @pytest.mark.asyncio + async def test_stops_operator_and_archives(self, mock_operator, mock_meta_store): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=mock_meta_store) + mock_operator.stop.assert_awaited_once_with("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_archives_stopped_state(self, mock_operator, mock_meta_store): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=mock_meta_store) + archived_info = mock_meta_store.archive.call_args[0][1] + assert archived_info["state"] == State.STOPPED + + @pytest.mark.asyncio + async def test_actor_not_found_still_archives(self, mock_meta_store): + op = AsyncMock() + op.stop = AsyncMock(side_effect=ValueError("not found")) + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + await sm.send("stop", sandbox_id="sb-1", operator=op, meta_store=mock_meta_store) + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_logs_billing_when_start_time_present(self, mock_operator, mock_meta_store): + sm = await SandboxStateMachine.from_state_value( + State.RUNNING, sandbox_info={"state": State.RUNNING, "start_time": "2024-01-01T00:00:00"} + ) + with patch("rock.sandbox.sandbox_statemachine.log_billing_info") as mock_billing: + await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=mock_meta_store) + mock_billing.assert_called_once() + + @pytest.mark.asyncio + async def test_meta_store_none_still_archives(self, mock_operator): + store = AsyncMock() + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=store) + store.archive.assert_awaited_once() diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py new file mode 100644 index 0000000000..67b6386a34 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -0,0 +1,184 @@ +"""Tests for SandboxManager lifecycle methods. + +Verifies that stop(), get_status(), and start_async() behave correctly for +each sandbox state, using lightweight mocks (no Ray / Docker). +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.actions.sandbox.response import State +from rock.sandbox.sandbox_manager import SandboxManager +from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError + + +@pytest.fixture +def mock_meta_store(): + store = AsyncMock() + store.get = AsyncMock(return_value=None) + store.create = AsyncMock() + store.update = AsyncMock() + store.archive = AsyncMock() + store.get_timeout = AsyncMock(return_value=None) + store.update_timeout = AsyncMock() + return store + + +@pytest.fixture +def mock_operator(): + op = AsyncMock() + op.submit = AsyncMock(return_value={"host_name": "h1", "host_ip": "1.2.3.4", "memory": "1g", "cpus": 1.0}) + op.stop = AsyncMock() + op.get_status = AsyncMock(return_value={"state": State.RUNNING}) + return op + + +@pytest.fixture +async def mgr(mock_meta_store, mock_operator): + """Minimal SandboxManager with real lifecycle logic and mocked infrastructure.""" + from rock.sandbox.sandbox_statemachine import SandboxStateMachine + + m = MagicMock(spec=SandboxManager) + m._meta_store = mock_meta_store + m._operator = mock_operator + + m._aes_encrypter = MagicMock() + m._aes_encrypter.encrypt = MagicMock(return_value="enc") + m.refresh_aes_key = AsyncMock() + + # Function to get state machine based on meta_store data — mirrors _get_current_statemachine + async def get_current_statemachine(sandbox_id: str) -> SandboxStateMachine | None: + info = await mock_meta_store.get(sandbox_id, check_db=True) + if info is None: + return None + state = info.get("state") + if state is None: + raise InternalServerRockError(f"Sandbox {sandbox_id} exists in store but has no state field") + return await SandboxStateMachine.from_state_value(state, sandbox_info=info) + + m._get_current_statemachine = AsyncMock(side_effect=get_current_statemachine) + + m.stop = SandboxManager.stop.__get__(m, SandboxManager) + m.get_status = SandboxManager.get_status.__get__(m, SandboxManager) + m._refresh_timeout = AsyncMock() + return m + + +# --------------------------------------------------------------------------- +# TestManagerStop +# --------------------------------------------------------------------------- + + +class TestManagerStop: + @pytest.mark.asyncio + async def test_stop_not_found_attempts_cleanup(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = None + await mgr.stop("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_not_found_actor_missing_still_archives(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = None + mock_operator.stop.side_effect = ValueError("actor not found") + await mgr.stop("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_already_stopped_is_noop(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.STOPPED} + await mgr.stop("sb-1") + mock_operator.stop.assert_not_awaited() + + @pytest.mark.asyncio + async def test_stop_running_calls_operator_and_archives(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING} + await mgr.stop("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_pending_calls_operator_and_archives(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.PENDING} + await mgr.stop("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_actor_not_found_still_archives(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING} + mock_operator.stop.side_effect = ValueError("actor not found") + await mgr.stop("sb-1") + mock_meta_store.archive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_sets_billing_info_when_start_time_present(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING, "start_time": "2024-01-01T00:00:00"} + with patch("rock.sandbox.sandbox_statemachine.log_billing_info") as mock_billing: + await mgr.stop("sb-1") + mock_billing.assert_called_once() + + @pytest.mark.asyncio + async def test_stop_archived_info_has_stopped_state(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING} + await mgr.stop("sb-1") + archived_info = mock_meta_store.archive.call_args[0][1] + assert archived_info["state"] == State.STOPPED + + @pytest.mark.asyncio + async def test_stop_missing_state_field_raises(self, mgr, mock_meta_store): + mock_meta_store.get.return_value = {} + with pytest.raises(InternalServerRockError, match="no state field"): + await mgr.stop("sb-1") + + +# --------------------------------------------------------------------------- +# TestManagerGetStatus +# --------------------------------------------------------------------------- + + +class TestManagerGetStatus: + @pytest.mark.asyncio + async def test_not_found_raises_when_operator_returns_none(self, mgr, mock_meta_store, mock_operator): + mock_operator.get_status.return_value = None + mock_meta_store.get.return_value = None + with pytest.raises(BadRequestRockError, match="not found"): + await mgr.get_status("sb-1") + + @pytest.mark.asyncio + async def test_include_all_states_falls_back_to_meta_store(self, mgr, mock_meta_store, mock_operator): + mock_operator.get_status.return_value = None + mock_meta_store.get.return_value = {"state": State.STOPPED, "phases": {}, "port_mapping": {}} + result = await mgr.get_status("sb-1", include_all_states=True) + assert result.state == State.STOPPED + assert result.is_alive is False + mock_meta_store.get.assert_awaited_with("sb-1", check_db=True) + + @pytest.mark.asyncio + async def test_running_returns_response(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING} + mock_operator.get_status.return_value = { + "state": State.RUNNING, + "host_name": "h1", + "host_ip": "1.2.3.4", + "phases": {}, + "port_mapping": {}, + } + result = await mgr.get_status("sb-1") + assert result.sandbox_id == "sb-1" + assert result.is_alive is True + + @pytest.mark.asyncio + async def test_updates_meta_on_state_change(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.PENDING} + mock_operator.get_status.return_value = {"state": State.RUNNING, "phases": {}, "port_mapping": {}} + await mgr.get_status("sb-1") + mock_meta_store.update.assert_awaited_once() + + @pytest.mark.asyncio + async def test_no_update_if_state_unchanged(self, mgr, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = {"state": State.RUNNING} + mock_operator.get_status.return_value = {"state": State.RUNNING, "phases": {}, "port_mapping": {}} + await mgr.get_status("sb-1") + mock_meta_store.update.assert_not_awaited() diff --git a/uv.lock b/uv.lock index 3ef30586fa..eb6e373dfc 100644 --- a/uv.lock +++ b/uv.lock @@ -3929,6 +3929,15 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/45/58/38b5afbc1a800eeea951b9285d3912613f2603bdf897a4ab0f4bd7f405fc/python_multipart-0.0.20-py3-none-any.whl", hash = "sha256:8a62d3a8335e06589fe01f2a3e178cdcc632f3fbe0d492ad9ee0ec35aab1f104" }, ] +[[package]] +name = "python-statemachine" +version = "3.1.2" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/af/e2/dcc389e9855f65a3e54c859499dd409341d5ba561da74af22ce824bc4964/python_statemachine-3.1.2.tar.gz", hash = "sha256:3a64c4ae91d628eb256580f328c68b16e472f25794d2c478e91f2f1dfb8b0668" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/f2/f5/fe9072bfc9d245138561c09bf05b841bf6d92b4ce93d44dbfa18d96e498e/python_statemachine-3.1.2-py3-none-any.whl", hash = "sha256:89dcdbcbf7b197adade16138c75bee4f32e0c98b7eeb18d30aeb1c830dbbdd62" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -4285,6 +4294,7 @@ dependencies = [ { name = "oss2" }, { name = "pydantic" }, { name = "python-multipart" }, + { name = "python-statemachine" }, { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, @@ -4437,6 +4447,7 @@ requires-dist = [ { name = "psutil", marker = "extra == 'rocklet'" }, { name = "pydantic" }, { name = "python-multipart" }, + { name = "python-statemachine", specifier = ">=3.1.1" }, { name = "pyyaml" }, { name = "ray", extras = ["default"], marker = "extra == 'admin'", specifier = "==2.43.0" }, { name = "redis", marker = "extra == 'admin'" }, From 94bcc1a8e2c7bccde2f9255579d43bbb9cf01f36 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Tue, 26 May 2026 11:46:13 +0800 Subject: [PATCH 133/226] =?UTF-8?q?fix(rocklet):=20add=20per-disk=20usage?= =?UTF-8?q?=20monitoring=20for=20rootfs,=20log,=20and=20kat=E2=80=A6=20(#9?= =?UTF-8?q?83)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(rocklet): add per-disk usage monitoring for rootfs, log, and kata DinD Split the single disk percentage metric into three independent gauges (rootfs, log directory, kata DinD) to enable fine-grained disk usage alerting in Grafana. The original "disk" field is preserved for backward compatibility. Co-Authored-By: Claude Opus 4.6 * fix(rocklet): remove duplicate disk_rootfs_percent metric, keep disk disk and disk_rootfs_percent were identical; remove the duplicate and annotate disk as the rootfs usage percent. Also remove the redundant disk_rootfs gauge from base_actor and guard disk_log/disk_dind reporting independently with truthy checks. Co-Authored-By: Claude Opus 4.6 * fix(rocklet): broaden exception handling in _get_docker_data_root() Catch all exceptions instead of only FileNotFoundError and ValueError, so PermissionError, OSError, TypeError etc. also fall back safely. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- rock/admin/metrics/constants.py | 2 + rock/admin/metrics/monitor.py | 2 + rock/rocklet/linux.py | 34 +++- rock/sandbox/base_actor.py | 11 ++ tests/unit/rocklet/test_disk_statistics.py | 176 +++++++++++++++++++++ 5 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 tests/unit/rocklet/test_disk_statistics.py diff --git a/rock/admin/metrics/constants.py b/rock/admin/metrics/constants.py index 1ef95f14ab..48c69b8c04 100644 --- a/rock/admin/metrics/constants.py +++ b/rock/admin/metrics/constants.py @@ -13,6 +13,8 @@ class MetricsConstants: SANDBOX_CPU = "system.cpu" SANDBOX_MEM = "system.memory" SANDBOX_DISK = "system.disk" + SANDBOX_DISK_LOG = "system.disk.log" + SANDBOX_DISK_DIND = "system.disk.dind" SANDBOX_NET = "system.network" TOTAL_CPU_RESOURCE = "resource.cpu.total" diff --git a/rock/admin/metrics/monitor.py b/rock/admin/metrics/monitor.py index 34ceddb8a1..44c8d077bb 100644 --- a/rock/admin/metrics/monitor.py +++ b/rock/admin/metrics/monitor.py @@ -90,6 +90,8 @@ def _register_metrics(self): self._register_gauge( MetricsConstants.SANDBOX_DISK, "Single sandbox disk usage percentage of allocated resources" ) + self._register_gauge(MetricsConstants.SANDBOX_DISK_LOG, "Sandbox log dir disk usage percent") + self._register_gauge(MetricsConstants.SANDBOX_DISK_DIND, "Sandbox kata DinD disk usage percent") self._register_gauge( MetricsConstants.SANDBOX_NET, "Single sandbox network usage percentage of allocated resources" ) diff --git a/rock/rocklet/linux.py b/rock/rocklet/linux.py index 6f91eb8ca5..2e14faee09 100644 --- a/rock/rocklet/linux.py +++ b/rock/rocklet/linux.py @@ -346,14 +346,46 @@ class LinuxRocklet(Rocklet): def __init__(self, **kwargs): super().__init__(**kwargs) self._cgroup_cpu = CgroupCpuStats() + self._docker_data_root: str | None = None def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: return BashSession(request) async def get_statistics(self) -> dict: + disk_root = psutil.disk_usage("/") + + log_path = os.environ.get("ROCK_LOGGING_PATH", "/data/logs") + try: + disk_log_percent = psutil.disk_usage(log_path).percent if os.path.exists(log_path) else 0.0 + except OSError: + disk_log_percent = 0.0 + + disk_dind_percent = 0.0 + if os.environ.get("ROCK_KATA_RUNTIME") == "true": + if self._docker_data_root is None: + self._docker_data_root = self._get_docker_data_root() + dind_path = self._docker_data_root + try: + disk_dind_percent = psutil.disk_usage(dind_path).percent if os.path.exists(dind_path) else 0.0 + except OSError: + disk_dind_percent = 0.0 + return { "cpu": self._cgroup_cpu.cpu_percent(), "mem": psutil.virtual_memory().percent, - "disk": psutil.disk_usage("/").percent, + "disk": disk_root.percent, # legacy metric name, actually rootfs usage percent + "disk_log_percent": disk_log_percent, + "disk_dind_percent": disk_dind_percent, "net": psutil.net_io_counters().bytes_recv + psutil.net_io_counters().bytes_sent, } + + @staticmethod + def _get_docker_data_root() -> str: + import json as _json + + try: + with open("/etc/docker/daemon.json") as f: + cfg = _json.load(f) + return cfg.get("data-root", "/var/lib/docker") + except Exception: + return "/var/lib/docker" diff --git a/rock/sandbox/base_actor.py b/rock/sandbox/base_actor.py index 4b8a94bdd7..a55963e3be 100644 --- a/rock/sandbox/base_actor.py +++ b/rock/sandbox/base_actor.py @@ -107,6 +107,12 @@ def _init_monitor(self): self._gauges["disk"] = self.meter.create_gauge( name="xrl_gateway.system.disk", description="Disk Usage", unit="1" ) + self._gauges["disk_log"] = self.meter.create_gauge( + name="xrl_gateway.system.disk.log", description="Sandbox log dir disk usage percent", unit="1" + ) + self._gauges["disk_dind"] = self.meter.create_gauge( + name="xrl_gateway.system.disk.dind", description="Sandbox kata DinD disk usage percent", unit="1" + ) self._gauges["net"] = self.meter.create_gauge( name="xrl_gateway.system.network", description="Network Usage", unit="1" ) @@ -211,6 +217,11 @@ async def _collect_sandbox_metrics(self, sandbox_id: str): self._gauges["disk"].set(metrics["disk"], attributes=attributes) self._gauges["net"].set(metrics["net"], attributes=attributes) + if metrics.get("disk_log_percent"): + self._gauges["disk_log"].set(metrics["disk_log_percent"], attributes=attributes) + if metrics.get("disk_dind_percent"): + self._gauges["disk_dind"].set(metrics["disk_dind_percent"], attributes=attributes) + # cpus_used inherits the semantic of metrics["cpu"]: it is reported # as a percentage of the sandbox's allocated CPU (see the legend of # xrl_gateway.system.cpu in admin/metrics/monitor.py). Multiplying by diff --git a/tests/unit/rocklet/test_disk_statistics.py b/tests/unit/rocklet/test_disk_statistics.py new file mode 100644 index 0000000000..b2e8b66636 --- /dev/null +++ b/tests/unit/rocklet/test_disk_statistics.py @@ -0,0 +1,176 @@ +import json +from collections import namedtuple +from unittest.mock import mock_open, patch + +import pytest + +from rock.rocklet.linux import LinuxRocklet + +DiskUsage = namedtuple("DiskUsage", ["total", "used", "free", "percent"]) + + +@pytest.fixture +def runtime(): + return LinuxRocklet() + + +@pytest.mark.asyncio +async def test_get_statistics_returns_all_disk_fields(runtime): + with ( + patch.object(runtime._cgroup_cpu, "cpu_percent", return_value=10.0), + patch("psutil.virtual_memory") as mock_vmem, + patch("psutil.disk_usage") as mock_disk, + patch("psutil.net_io_counters") as mock_net, + patch("os.path.exists", return_value=True), + patch.dict("os.environ", {"ROCK_LOGGING_PATH": "/data/logs", "ROCK_KATA_RUNTIME": "true"}), + patch("builtins.open", mock_open(read_data=json.dumps({"data-root": "/var/lib/docker"}))), + ): + mock_vmem.return_value = type("obj", (object,), {"percent": 50.0})() + mock_net.return_value = type("obj", (object,), {"bytes_recv": 100, "bytes_sent": 200})() + mock_disk.side_effect = lambda path: { + "/": DiskUsage(total=100, used=75, free=25, percent=75.0), + "/data/logs": DiskUsage(total=100, used=60, free=40, percent=60.0), + "/var/lib/docker": DiskUsage(total=100, used=80, free=20, percent=80.0), + }[path] + + stats = await runtime.get_statistics() + + assert stats["cpu"] == 10.0 + assert stats["mem"] == 50.0 + assert stats["disk"] == 75.0 + assert stats["disk_log_percent"] == 60.0 + assert stats["disk_dind_percent"] == 80.0 + assert stats["net"] == 300 + + +@pytest.mark.asyncio +async def test_get_statistics_no_kata_runtime(runtime): + with ( + patch.object(runtime._cgroup_cpu, "cpu_percent", return_value=5.0), + patch("psutil.virtual_memory") as mock_vmem, + patch("psutil.disk_usage") as mock_disk, + patch("psutil.net_io_counters") as mock_net, + patch("os.path.exists", return_value=True), + patch.dict("os.environ", {"ROCK_LOGGING_PATH": "/data/logs"}, clear=False), + ): + mock_vmem.return_value = type("obj", (object,), {"percent": 40.0})() + mock_net.return_value = type("obj", (object,), {"bytes_recv": 50, "bytes_sent": 50})() + mock_disk.side_effect = lambda path: { + "/": DiskUsage(total=100, used=30, free=70, percent=30.0), + "/data/logs": DiskUsage(total=100, used=20, free=80, percent=20.0), + }[path] + + # Ensure ROCK_KATA_RUNTIME is not set + import os + + os.environ.pop("ROCK_KATA_RUNTIME", None) + + stats = await runtime.get_statistics() + + assert stats["disk_dind_percent"] == 0.0 + assert stats["disk"] == 30.0 + assert stats["disk_log_percent"] == 20.0 + + +@pytest.mark.asyncio +async def test_get_statistics_log_path_not_exists(runtime): + with ( + patch.object(runtime._cgroup_cpu, "cpu_percent", return_value=5.0), + patch("psutil.virtual_memory") as mock_vmem, + patch("psutil.disk_usage") as mock_disk, + patch("psutil.net_io_counters") as mock_net, + patch("os.path.exists", return_value=False), + patch.dict("os.environ", {}, clear=False), + ): + mock_vmem.return_value = type("obj", (object,), {"percent": 40.0})() + mock_net.return_value = type("obj", (object,), {"bytes_recv": 10, "bytes_sent": 10})() + mock_disk.return_value = DiskUsage(total=100, used=50, free=50, percent=50.0) + + import os + + os.environ.pop("ROCK_KATA_RUNTIME", None) + + stats = await runtime.get_statistics() + + assert stats["disk_log_percent"] == 0.0 + assert stats["disk_dind_percent"] == 0.0 + + +@pytest.mark.asyncio +async def test_get_statistics_log_disk_oserror(runtime): + def disk_usage_side_effect(path): + if path == "/": + return DiskUsage(total=100, used=50, free=50, percent=50.0) + raise OSError("Permission denied") + + with ( + patch.object(runtime._cgroup_cpu, "cpu_percent", return_value=5.0), + patch("psutil.virtual_memory") as mock_vmem, + patch("psutil.disk_usage", side_effect=disk_usage_side_effect), + patch("psutil.net_io_counters") as mock_net, + patch("os.path.exists", return_value=True), + patch.dict("os.environ", {"ROCK_LOGGING_PATH": "/data/logs"}, clear=False), + ): + mock_vmem.return_value = type("obj", (object,), {"percent": 40.0})() + mock_net.return_value = type("obj", (object,), {"bytes_recv": 10, "bytes_sent": 10})() + + import os + + os.environ.pop("ROCK_KATA_RUNTIME", None) + + stats = await runtime.get_statistics() + + assert stats["disk_log_percent"] == 0.0 + + +@pytest.mark.asyncio +async def test_get_statistics_dind_disk_oserror(runtime): + def disk_usage_side_effect(path): + if path == "/": + return DiskUsage(total=100, used=50, free=50, percent=50.0) + if path == "/data/logs": + return DiskUsage(total=100, used=30, free=70, percent=30.0) + raise OSError("Permission denied") + + with ( + patch.object(runtime._cgroup_cpu, "cpu_percent", return_value=5.0), + patch("psutil.virtual_memory") as mock_vmem, + patch("psutil.disk_usage", side_effect=disk_usage_side_effect), + patch("psutil.net_io_counters") as mock_net, + patch("os.path.exists", return_value=True), + patch.dict( + "os.environ", {"ROCK_LOGGING_PATH": "/data/logs", "ROCK_KATA_RUNTIME": "true"}, clear=False + ), + patch("builtins.open", mock_open(read_data=json.dumps({"data-root": "/var/lib/docker"}))), + ): + mock_vmem.return_value = type("obj", (object,), {"percent": 40.0})() + mock_net.return_value = type("obj", (object,), {"bytes_recv": 10, "bytes_sent": 10})() + + stats = await runtime.get_statistics() + + assert stats["disk_dind_percent"] == 0.0 + assert stats["disk_log_percent"] == 30.0 + + +class TestGetDockerDataRoot: + def test_reads_data_root_from_daemon_json(self): + daemon_config = json.dumps({"data-root": "/mnt/docker-data"}) + with patch("builtins.open", mock_open(read_data=daemon_config)): + result = LinuxRocklet._get_docker_data_root() + assert result == "/mnt/docker-data" + + def test_returns_default_when_no_data_root_key(self): + daemon_config = json.dumps({"storage-driver": "overlay2"}) + with patch("builtins.open", mock_open(read_data=daemon_config)): + result = LinuxRocklet._get_docker_data_root() + assert result == "/var/lib/docker" + + def test_returns_default_when_file_not_found(self): + with patch("builtins.open", side_effect=FileNotFoundError): + result = LinuxRocklet._get_docker_data_root() + assert result == "/var/lib/docker" + + def test_returns_default_when_invalid_json(self): + with patch("builtins.open", mock_open(read_data="not valid json")): + result = LinuxRocklet._get_docker_data_root() + assert result == "/var/lib/docker" From cefdaa60afb1deda8d3d8abb3a1eae91631014d5 Mon Sep 17 00:00:00 2001 From: lkc Date: Wed, 27 May 2026 13:41:48 +0800 Subject: [PATCH 134/226] feat(datasets): make 'rock datasets list' fast on cross-region OSS (#1010) (#1011) * feat(datasets): add list_organizations to registry * feat(datasets): add list_org_datasets to registry * feat(datasets): add list_dataset_splits to registry * feat(datasets): add list_all_datasets with bounded concurrency * feat(datasets): add fast-path methods to DatasetClient * feat(datasets): rewrite list with --depth and fast paths AI-Model: claude-opus-4-7 AI-Contributed/Feature: 48/48 AI-Contributed/UT: 107/107 * feat(datasets): add splits subcommand AI-Model: claude-opus-4-7 AI-Contributed/Feature: 24/24 AI-Contributed/UT: 63/63 * chore(datasets): apply lint and format AI-Model: claude-opus-4-7 AI-Contributed/Feature: 36/53 AI-Contributed/UT: 44/65 * fix(datasets): preserve list parser exclusivity Defer the default list depth until runtime so argparse still treats an explicit --depth value as present in the mutually exclusive group on Python 3.10. Update envhub CLI docs for the rewritten list output and splits command. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 47/47 AI-Contributed/UT: 8/8 --- docs/dev/envhub/README.md | 42 +++- rock/cli/command/datasets.py | 109 ++++++++--- rock/sdk/envhub/datasets/client.py | 13 +- rock/sdk/envhub/datasets/registry/base.py | 21 +- rock/sdk/envhub/datasets/registry/oss.py | 53 +++-- tests/unit/datasets/test_client.py | 39 ++++ tests/unit/datasets/test_datasets_command.py | 175 ++++++++++++++++- tests/unit/datasets/test_oss_registry.py | 194 +++++++++++++++++-- 8 files changed, 593 insertions(+), 53 deletions(-) diff --git a/docs/dev/envhub/README.md b/docs/dev/envhub/README.md index 6255efe1f6..7f5b4080e4 100644 --- a/docs/dev/envhub/README.md +++ b/docs/dev/envhub/README.md @@ -248,20 +248,52 @@ class DatasetClient: rock datasets list [OPTIONS] Options: - --org TEXT 只列出指定 organization 的 datasets + --depth {1,2} 1: 只列出 organizations；2: 列出 organizations 和 datasets（默认） + --org TEXT 只列出指定 organization 的 datasets（与 --depth 互斥） --bucket TEXT OSS bucket 名称（覆盖 config.ini） --endpoint TEXT OSS endpoint（覆盖 config.ini） --access-key-id TEXT OSS access key ID（覆盖 config.ini） --access-key-secret TEXT OSS access key secret（覆盖 config.ini） + --region TEXT OSS region（覆盖 config.ini） ``` 输出示例： ``` -Dataset Split Tasks -qwen/my-bench train 42 -qwen/my-bench test 10 -alibaba/code-eval train 100 +Organization Dataset +-------------------------- +alibaba code-eval +qwen my-bench + +2 datasets in 2 organizations. +``` + +#### rock datasets splits + +``` +rock datasets splits [OPTIONS] + +Required: + --org TEXT Organization 名称 + --dataset TEXT Dataset 名称 + +Options: + --bucket TEXT OSS bucket 名称（覆盖 config.ini） + --endpoint TEXT OSS endpoint（覆盖 config.ini） + --access-key-id TEXT OSS access key ID（覆盖 config.ini） + --access-key-secret TEXT OSS access key secret（覆盖 config.ini） + --region TEXT OSS region（覆盖 config.ini） +``` + +输出示例： + +``` +Split +----- +test +train + +2 splits. ``` #### rock datasets upload diff --git a/rock/cli/command/datasets.py b/rock/cli/command/datasets.py index b53636cb79..d35b8b87ea 100644 --- a/rock/cli/command/datasets.py +++ b/rock/cli/command/datasets.py @@ -35,6 +35,8 @@ async def arun(self, args: argparse.Namespace) -> None: await self._list(args) elif args.datasets_command == "tasks": await self._tasks(args) + elif args.datasets_command == "splits": + await self._splits(args) elif args.datasets_command == "upload": await self._upload(args) else: @@ -59,20 +61,48 @@ def _build_oss_registry_info(self, args: argparse.Namespace) -> OssRegistryInfo: async def _list(self, args: argparse.Namespace) -> None: registry_info = self._build_oss_registry_info(args) client = DatasetClient(registry_info) - datasets = client.list_datasets(org=getattr(args, "org", None)) - if not datasets: - print("No datasets found.") + if getattr(args, "org", None): + datasets = client.list_org_datasets(args.org) + pairs = [(args.org, d) for d in datasets] + self._render_org_dataset_pairs(pairs) + return + + depth = getattr(args, "depth", None) or 2 + if depth == 1: + orgs = client.list_organizations() + self._render_orgs(orgs) return - col_id = max(len("Dataset"), max(len(d.id) for d in datasets)) - col_split = max(len("Split"), max(len(d.split) for d in datasets)) + pairs = client.list_all_datasets() + self._render_org_dataset_pairs(pairs) - header = f"{'Dataset':<{col_id}} {'Split':<{col_split}} {'Tasks':>6}" + @staticmethod + def _render_org_dataset_pairs(pairs: list[tuple[str, str]]) -> None: + if not pairs: + print("No datasets found.") + return + col_org = max(len("Organization"), max(len(o) for o, _ in pairs)) + col_ds = max(len("Dataset"), max(len(d) for _, d in pairs)) + header = f"{'Organization':<{col_org}} {'Dataset':<{col_ds}}" print(header) print("-" * len(header)) - for ds in sorted(datasets, key=lambda d: (d.id, d.split)): - print(f"{ds.id:<{col_id}} {ds.split:<{col_split}} {len(ds.task_ids):>6}") + for o, d in pairs: + print(f"{o:<{col_org}} {d:<{col_ds}}") + n_orgs = len({o for o, _ in pairs}) + print(f"\n{len(pairs)} datasets in {n_orgs} organizations.") + + @staticmethod + def _render_orgs(orgs: list[str]) -> None: + if not orgs: + print("No organizations found.") + return + width = max(len("Organization"), max(len(o) for o in orgs)) + print(f"{'Organization':<{width}}") + print("-" * width) + for o in orgs: + print(o) + print(f"\n{len(orgs)} organizations.") async def _tasks(self, args: argparse.Namespace) -> None: registry_info = self._build_oss_registry_info(args) @@ -92,8 +122,6 @@ async def _tasks(self, args: argparse.Namespace) -> None: print("No tasks found after applying offset/limit.") return - limit_text = str(args.limit) if args.limit is not None else "all" - print() print("=" * 80) print(f"Dataset: {spec.id} Split: {spec.split} Total: {total} Shown: {len(shown_task_ids)}") @@ -103,6 +131,23 @@ async def _tasks(self, args: argparse.Namespace) -> None: for task_id in shown_task_ids: print(task_id) + async def _splits(self, args: argparse.Namespace) -> None: + registry_info = self._build_oss_registry_info(args) + client = DatasetClient(registry_info) + splits = client.list_dataset_splits(args.org, args.dataset) + + if not splits: + print(f"No splits found for dataset '{args.org}/{args.dataset}'.") + return + + width = max(len("Split"), max(len(s) for s in splits)) + print(f"{'Split':<{width}}") + print("-" * width) + for s in splits: + print(s) + word = "split" if len(splits) == 1 else "splits" + print(f"\n{len(splits)} {word}.") + async def _upload(self, args: argparse.Namespace) -> None: local_dir = Path(args.dir) if not local_dir.is_dir(): @@ -135,14 +180,24 @@ async def add_parser_to(subparsers: argparse._SubParsersAction) -> None: def add_oss_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--bucket", help="OSS bucket name (overrides config.ini)") parser.add_argument("--endpoint", help="OSS endpoint URL (overrides config.ini)") - parser.add_argument("--access-key-id", dest="access_key_id", - help="OSS access key ID (overrides config.ini)") - parser.add_argument("--access-key-secret", dest="access_key_secret", - help="OSS access key secret (overrides config.ini)") + parser.add_argument( + "--access-key-id", dest="access_key_id", help="OSS access key ID (overrides config.ini)" + ) + parser.add_argument( + "--access-key-secret", dest="access_key_secret", help="OSS access key secret (overrides config.ini)" + ) parser.add_argument("--region", help="OSS region (overrides config.ini)") list_parser = datasets_subparsers.add_parser("list", help="List datasets in OSS registry") - list_parser.add_argument("--org", help="Filter by organization") + list_group = list_parser.add_mutually_exclusive_group() + list_group.add_argument( + "--depth", + type=int, + choices=[1, 2], + default=None, + help="1: list orgs only. 2 (default): list orgs and datasets.", + ) + list_group.add_argument("--org", help="List datasets under the given organization only") add_oss_args(list_parser) tasks_parser = datasets_subparsers.add_parser("tasks", help="List task IDs under one dataset split") tasks_parser.add_argument("--org", required=True, help="Organization name") @@ -152,15 +207,25 @@ def add_oss_args(parser: argparse.ArgumentParser) -> None: tasks_parser.add_argument("--limit", type=_positive_int, default=None, help="Maximum number of tasks to show") add_oss_args(tasks_parser) + splits_parser = datasets_subparsers.add_parser("splits", help="List splits under one dataset") + splits_parser.add_argument("--org", required=True, help="Organization name") + splits_parser.add_argument("--dataset", required=True, help="Dataset name") + add_oss_args(splits_parser) + upload_parser = datasets_subparsers.add_parser("upload", help="Upload local task dirs to OSS") upload_parser.add_argument("--org", required=True, help="Organization name") upload_parser.add_argument("--dataset", required=True, help="Dataset name") upload_parser.add_argument("--split", required=True, help="Split name (e.g. train, test, v1.0)") - upload_parser.add_argument("--dir", required=True, - help="Local directory containing {task_id}/ subdirectories") - upload_parser.add_argument("--concurrency", type=int, default=4, - choices=range(1, 17), metavar="[1-16]", - help="Upload concurrency (default: 4)") - upload_parser.add_argument("--overwrite", action="store_true", - help="Overwrite existing tasks in OSS (default: skip)") + upload_parser.add_argument("--dir", required=True, help="Local directory containing {task_id}/ subdirectories") + upload_parser.add_argument( + "--concurrency", + type=int, + default=4, + choices=range(1, 17), + metavar="[1-16]", + help="Upload concurrency (default: 4)", + ) + upload_parser.add_argument( + "--overwrite", action="store_true", help="Overwrite existing tasks in OSS (default: skip)" + ) add_oss_args(upload_parser) diff --git a/rock/sdk/envhub/datasets/client.py b/rock/sdk/envhub/datasets/client.py index 2ea4d70af0..16189defec 100644 --- a/rock/sdk/envhub/datasets/client.py +++ b/rock/sdk/envhub/datasets/client.py @@ -4,7 +4,6 @@ class DatasetClient: - def __init__(self, registry: OssRegistryInfo) -> None: self._registry = OssDatasetRegistry(registry) @@ -14,6 +13,18 @@ def list_datasets(self, org: str | None = None) -> list[DatasetSpec]: def list_dataset_tasks(self, organization: str, dataset: str, split: str = "test") -> DatasetSpec | None: return self._registry.list_dataset_tasks(organization, dataset, split) + def list_organizations(self) -> list[str]: + return self._registry.list_organizations() + + def list_org_datasets(self, organization: str) -> list[str]: + return self._registry.list_org_datasets(organization) + + def list_all_datasets(self, concurrency: int = 10) -> list[tuple[str, str]]: + return self._registry.list_all_datasets(concurrency) + + def list_dataset_splits(self, organization: str, dataset: str) -> list[str]: + return self._registry.list_dataset_splits(organization, dataset) + def upload_dataset( self, source: LocalDatasetConfig, diff --git a/rock/sdk/envhub/datasets/registry/base.py b/rock/sdk/envhub/datasets/registry/base.py index 5c9192d2d6..0532671ffe 100644 --- a/rock/sdk/envhub/datasets/registry/base.py +++ b/rock/sdk/envhub/datasets/registry/base.py @@ -5,7 +5,6 @@ class BaseDatasetRegistry(ABC): - @abstractmethod def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: """List all datasets. Filtered to `organization` if provided.""" @@ -16,6 +15,26 @@ def list_dataset_tasks(self, organization: str, dataset: str, split: str = "test """List task ids for one dataset split. Returns None if dataset/split has no tasks.""" ... + @abstractmethod + def list_organizations(self) -> list[str]: + """List organization names under the dataset registry. Single backend call.""" + ... + + @abstractmethod + def list_org_datasets(self, organization: str) -> list[str]: + """List dataset names under one organization. Single backend call.""" + ... + + @abstractmethod + def list_dataset_splits(self, organization: str, dataset: str) -> list[str]: + """List split names under one dataset. Single backend call.""" + ... + + @abstractmethod + def list_all_datasets(self, concurrency: int = 10) -> list[tuple[str, str]]: + """List all (org, dataset) pairs. 1 + N_org backend calls with bounded concurrency.""" + ... + @abstractmethod def upload_dataset( self, diff --git a/rock/sdk/envhub/datasets/registry/oss.py b/rock/sdk/envhub/datasets/registry/oss.py index 4e90613571..af36481ce9 100644 --- a/rock/sdk/envhub/datasets/registry/oss.py +++ b/rock/sdk/envhub/datasets/registry/oss.py @@ -1,6 +1,6 @@ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import oss2 @@ -14,7 +14,6 @@ class OssDatasetRegistry(BaseDatasetRegistry): - def __init__(self, registry: OssRegistryInfo) -> None: self._registry = registry @@ -58,7 +57,7 @@ def _extract_tasks_from_split(self, bucket: oss2.Bucket, split_prefix: str) -> l if key.endswith("/"): continue # Get the relative path from split_prefix - relative = key[len(split_prefix):] + relative = key[len(split_prefix) :] # Only direct files (no nested paths with "/") if "/" in relative: continue @@ -70,6 +69,37 @@ def _extract_tasks_from_split(self, bucket: oss2.Bucket, split_prefix: str) -> l all_tasks = sorted(set(dir_tasks + file_tasks)) return all_tasks + def list_organizations(self) -> list[str]: + bucket = self._build_bucket() + base = self._registry.oss_dataset_path or "datasets" + result = bucket.list_objects_v2(prefix=f"{base}/", delimiter="/", max_keys=1000) + return sorted(self._last_segment(p) for p in result.prefix_list) + + def list_org_datasets(self, organization: str) -> list[str]: + bucket = self._build_bucket() + base = self._registry.oss_dataset_path or "datasets" + result = bucket.list_objects_v2(prefix=f"{base}/{organization}/", delimiter="/", max_keys=1000) + return sorted(self._last_segment(p) for p in result.prefix_list) + + def list_dataset_splits(self, organization: str, dataset: str) -> list[str]: + bucket = self._build_bucket() + base = self._registry.oss_dataset_path or "datasets" + result = bucket.list_objects_v2(prefix=f"{base}/{organization}/{dataset}/", delimiter="/", max_keys=1000) + return sorted(self._last_segment(p) for p in result.prefix_list) + + def list_all_datasets(self, concurrency: int = 10) -> list[tuple[str, str]]: + orgs = self.list_organizations() + if not orgs: + return [] + pairs: list[tuple[str, str]] = [] + with ThreadPoolExecutor(max_workers=concurrency) as ex: + future_to_org = {ex.submit(self.list_org_datasets, o): o for o in orgs} + for fut in as_completed(future_to_org): + org = future_to_org[fut] + for ds in fut.result(): + pairs.append((org, ds)) + return sorted(pairs) + def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: bucket = self._build_bucket() base = self._registry.oss_dataset_path or "datasets" @@ -93,11 +123,13 @@ def list_datasets(self, organization: str | None = None) -> list[DatasetSpec]: split = self._last_segment(split_prefix) task_ids = self._extract_tasks_from_split(bucket, split_prefix) - datasets.append(DatasetSpec( - id=f"{org}/{name}", - split=split, - task_ids=task_ids, - )) + datasets.append( + DatasetSpec( + id=f"{org}/{name}", + split=split, + task_ids=task_ids, + ) + ) return datasets @@ -157,10 +189,7 @@ def upload_dataset( raw: dict[str, int | None | Exception] = {} with ThreadPoolExecutor(max_workers=concurrency) as executor: - futures = { - executor.submit(self._upload_task, bucket, org, name, split, d, overwrite): d - for d in task_dirs - } + futures = {executor.submit(self._upload_task, bucket, org, name, split, d, overwrite): d for d in task_dirs} for future, task_dir in futures.items(): try: raw[task_dir.name] = future.result() diff --git a/tests/unit/datasets/test_client.py b/tests/unit/datasets/test_client.py index baf8a8f753..f811763120 100644 --- a/tests/unit/datasets/test_client.py +++ b/tests/unit/datasets/test_client.py @@ -42,3 +42,42 @@ def test_dataset_client_list_tasks_delegates_to_registry_with_default_split(): mock_list_tasks.assert_called_once_with("qwen", "bench", "test") assert result == expected + + +def test_dataset_client_list_organizations_delegates(): + client = DatasetClient(make_registry_info()) + with patch.object(client._registry, "list_organizations", return_value=["a", "b"]) as m: + result = client.list_organizations() + m.assert_called_once_with() + assert result == ["a", "b"] + + +def test_dataset_client_list_org_datasets_delegates(): + client = DatasetClient(make_registry_info()) + with patch.object(client._registry, "list_org_datasets", return_value=["d1"]) as m: + result = client.list_org_datasets("qwen") + m.assert_called_once_with("qwen") + assert result == ["d1"] + + +def test_dataset_client_list_all_datasets_delegates_with_default_concurrency(): + client = DatasetClient(make_registry_info()) + with patch.object(client._registry, "list_all_datasets", return_value=[("a", "x")]) as m: + result = client.list_all_datasets() + m.assert_called_once_with(10) + assert result == [("a", "x")] + + +def test_dataset_client_list_all_datasets_passes_custom_concurrency(): + client = DatasetClient(make_registry_info()) + with patch.object(client._registry, "list_all_datasets", return_value=[]) as m: + client.list_all_datasets(concurrency=5) + m.assert_called_once_with(5) + + +def test_dataset_client_list_dataset_splits_delegates(): + client = DatasetClient(make_registry_info()) + with patch.object(client._registry, "list_dataset_splits", return_value=["test", "train"]) as m: + result = client.list_dataset_splits("qwen", "bench") + m.assert_called_once_with("qwen", "bench") + assert result == ["test", "train"] diff --git a/tests/unit/datasets/test_datasets_command.py b/tests/unit/datasets/test_datasets_command.py index 7c0b0d068e..7fc17561f9 100644 --- a/tests/unit/datasets/test_datasets_command.py +++ b/tests/unit/datasets/test_datasets_command.py @@ -21,6 +21,7 @@ def make_base_args(**kwargs): org=None, dataset=None, split=None, + depth=2, offset=0, limit=None, ) @@ -28,6 +29,7 @@ def make_base_args(**kwargs): setattr(args, k, v) return args + def make_registry_info(): return OssRegistryInfo(oss_bucket="b", oss_access_key_id="k", oss_access_key_secret="s") @@ -45,7 +47,9 @@ def test_command_name(): def test_build_oss_registry_info_from_cli_args(): cmd = DatasetsCommand() - args = make_base_args(bucket="cli-bucket", endpoint="https://oss.example.com", access_key_id="kid", access_key_secret="ksec") + args = make_base_args( + bucket="cli-bucket", endpoint="https://oss.example.com", access_key_id="kid", access_key_secret="ksec" + ) with patch("rock.cli.command.datasets.ConfigManager") as mock_mgr: ds_cfg = mock_mgr.return_value.get_config.return_value.dataset_config @@ -204,3 +208,172 @@ def test_tasks_prints_no_tasks_message_when_not_found(capsys): out = capsys.readouterr().out assert "No tasks found" in out + + +# --------------------------------------------------------------------------- +# list subcommand tests (depth + --org rewrite) +# --------------------------------------------------------------------------- + + +def test_list_default_depth_calls_list_all_datasets_and_renders_two_columns(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="list", depth=None, org=None) + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_all_datasets.return_value = [ + ("alibaba", "pinch"), + ("qwen", "bench-1"), + ] + asyncio.run(cmd._list(args)) + + MockClient.return_value.list_all_datasets.assert_called_once_with() + out = capsys.readouterr().out + assert "Organization" in out + assert "Dataset" in out + assert "alibaba" in out and "pinch" in out + assert "qwen" in out and "bench-1" in out + assert "2 datasets in 2 organizations." in out + + +def test_list_depth_1_calls_list_organizations_and_renders_one_column(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="list", depth=1, org=None) + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_organizations.return_value = ["alibaba", "qwen"] + asyncio.run(cmd._list(args)) + + MockClient.return_value.list_organizations.assert_called_once_with() + MockClient.return_value.list_all_datasets.assert_not_called() + out = capsys.readouterr().out + assert "Organization" in out + assert "alibaba" in out + assert "qwen" in out + assert "2 organizations." in out + + +def test_list_with_org_calls_list_org_datasets(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="list", depth=2, org="alibaba") + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_org_datasets.return_value = ["pinch", "webdev"] + asyncio.run(cmd._list(args)) + + MockClient.return_value.list_org_datasets.assert_called_once_with("alibaba") + MockClient.return_value.list_all_datasets.assert_not_called() + MockClient.return_value.list_organizations.assert_not_called() + out = capsys.readouterr().out + assert "alibaba" in out and "pinch" in out and "webdev" in out + assert "2 datasets in 1 organizations." in out + + +def test_list_empty_prints_no_datasets_message(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="list", depth=2, org=None) + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_all_datasets.return_value = [] + asyncio.run(cmd._list(args)) + + out = capsys.readouterr().out + assert "No datasets found." in out + + +def test_list_depth_1_empty_prints_no_organizations_message(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="list", depth=1, org=None) + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_organizations.return_value = [] + asyncio.run(cmd._list(args)) + + out = capsys.readouterr().out + assert "No organizations found." in out + + +def test_list_parser_depth_and_org_mutually_exclusive(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["datasets", "list", "--depth", "2", "--org", "alibaba"]) + + +def test_list_parser_depth_default_is_deferred_to_runtime(): + parser = _build_parser() + parsed = parser.parse_args(["datasets", "list"]) + assert parsed.depth is None + assert parsed.org is None + + +def test_list_parser_rejects_invalid_depth(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["datasets", "list", "--depth", "3"]) + + +# --------------------------------------------------------------------------- +# splits subcommand tests +# --------------------------------------------------------------------------- + + +def test_splits_lists_split_names(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="splits", org="alibaba", dataset="pinch") + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_dataset_splits.return_value = ["test", "train"] + asyncio.run(cmd._splits(args)) + + MockClient.return_value.list_dataset_splits.assert_called_once_with("alibaba", "pinch") + out = capsys.readouterr().out + assert "Split" in out + assert "test" in out + assert "train" in out + assert "2 splits." in out + + +def test_splits_empty_prints_no_splits_message(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="splits", org="alibaba", dataset="missing") + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_dataset_splits.return_value = [] + asyncio.run(cmd._splits(args)) + + out = capsys.readouterr().out + assert "No splits found for dataset 'alibaba/missing'." in out + + +def test_splits_singular_footer_for_one_split(capsys): + cmd = DatasetsCommand() + args = make_base_args(datasets_command="splits", org="alibaba", dataset="pinch") + + with patch.object(DatasetsCommand, "_build_oss_registry_info", return_value=make_registry_info()): + with patch("rock.cli.command.datasets.DatasetClient") as MockClient: + MockClient.return_value.list_dataset_splits.return_value = ["test"] + asyncio.run(cmd._splits(args)) + + out = capsys.readouterr().out + assert "1 split." in out + + +def test_splits_parser_requires_org_and_dataset(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["datasets", "splits"]) + with pytest.raises(SystemExit): + parser.parse_args(["datasets", "splits", "--org", "alibaba"]) + + +def test_splits_parser_accepts_org_and_dataset(): + parser = _build_parser() + parsed = parser.parse_args(["datasets", "splits", "--org", "alibaba", "--dataset", "pinch"]) + assert parsed.org == "alibaba" + assert parsed.dataset == "pinch" diff --git a/tests/unit/datasets/test_oss_registry.py b/tests/unit/datasets/test_oss_registry.py index 9237ac1d6f..fc3dd9be70 100644 --- a/tests/unit/datasets/test_oss_registry.py +++ b/tests/unit/datasets/test_oss_registry.py @@ -27,10 +27,12 @@ def test_list_datasets_returns_all(): make_list_result(prefixes=["datasets/qwen/"]), make_list_result(prefixes=["datasets/qwen/my-bench/"]), make_list_result(prefixes=["datasets/qwen/my-bench/train/"]), - make_list_result(prefixes=[ - "datasets/qwen/my-bench/train/task-001/", - "datasets/qwen/my-bench/train/task-002/", - ]), + make_list_result( + prefixes=[ + "datasets/qwen/my-bench/train/task-001/", + "datasets/qwen/my-bench/train/task-002/", + ] + ), ] with patch.object(registry, "_build_bucket", return_value=mock_bucket): @@ -58,6 +60,7 @@ def test_list_datasets_filter_by_org(): assert first_call_kwargs["prefix"] == "datasets/qwen/" assert len(datasets) == 1 + def test_list_datasets_counts_directory_and_file_tasks(): registry = OssDatasetRegistry(make_registry_info()) mock_bucket = MagicMock() @@ -104,6 +107,7 @@ def test_build_prefix_with_split(): registry = OssDatasetRegistry(make_registry_info()) assert registry._build_prefix("qwen", "my-bench", "train") == "datasets/qwen/my-bench/train" + # --------------------------------------------------------------------------- # list_dataset_tasks tests # --------------------------------------------------------------------------- @@ -112,10 +116,12 @@ def test_build_prefix_with_split(): def test_list_dataset_tasks_uses_default_test_split_and_sorts_task_ids(): registry = OssDatasetRegistry(make_registry_info()) mock_bucket = MagicMock() - mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[ - "datasets/qwen/my-bench/test/task-002/", - "datasets/qwen/my-bench/test/task-001/", - ]) + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[ + "datasets/qwen/my-bench/test/task-002/", + "datasets/qwen/my-bench/test/task-001/", + ] + ) with patch.object(registry, "_build_bucket", return_value=mock_bucket): spec = registry.list_dataset_tasks("qwen", "my-bench") @@ -132,9 +138,11 @@ def test_list_dataset_tasks_uses_default_test_split_and_sorts_task_ids(): def test_list_dataset_tasks_supports_custom_split(): registry = OssDatasetRegistry(make_registry_info()) mock_bucket = MagicMock() - mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[ - "datasets/qwen/my-bench/train/task-001/", - ]) + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[ + "datasets/qwen/my-bench/train/task-001/", + ] + ) with patch.object(registry, "_build_bucket", return_value=mock_bucket): spec = registry.list_dataset_tasks("qwen", "my-bench", "train") @@ -146,6 +154,7 @@ def test_list_dataset_tasks_supports_custom_split(): first_call_kwargs = mock_bucket.list_objects_v2.call_args_list[0][1] assert first_call_kwargs["prefix"] == "datasets/qwen/my-bench/train/" + def test_list_dataset_tasks_includes_directory_and_file_tasks_with_suffix_stripped(): registry = OssDatasetRegistry(make_registry_info()) mock_bucket = MagicMock() @@ -281,3 +290,166 @@ def test_upload_dataset_oss_key_format(tmp_path): key = mock_bucket.put_object.call_args[0][0] assert key == "datasets/qwen/my-bench/train/task-001/task.toml" + + +# --------------------------------------------------------------------------- +# list_organizations tests +# --------------------------------------------------------------------------- + + +def test_list_organizations_returns_sorted_org_names(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[ + "datasets/qwen/", + "datasets/alibaba/", + "datasets/AoneBenchDev/", + ] + ) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + orgs = registry.list_organizations() + + call_kwargs = mock_bucket.list_objects_v2.call_args[1] + assert call_kwargs["prefix"] == "datasets/" + assert call_kwargs["delimiter"] == "/" + assert call_kwargs["max_keys"] == 1000 + assert orgs == ["AoneBenchDev", "alibaba", "qwen"] + + +def test_list_organizations_returns_empty_when_no_orgs(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + orgs = registry.list_organizations() + + assert orgs == [] + + +def test_list_org_datasets_returns_sorted_dataset_names(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[ + "datasets/qwen/bench-2/", + "datasets/qwen/bench-1/", + ] + ) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + datasets = registry.list_org_datasets("qwen") + + call_kwargs = mock_bucket.list_objects_v2.call_args[1] + assert call_kwargs["prefix"] == "datasets/qwen/" + assert call_kwargs["delimiter"] == "/" + assert call_kwargs["max_keys"] == 1000 + assert datasets == ["bench-1", "bench-2"] + + +def test_list_org_datasets_returns_empty_when_org_missing(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + assert registry.list_org_datasets("nonexistent") == [] + + +def test_list_dataset_splits_returns_sorted_split_names(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result( + prefixes=[ + "datasets/qwen/bench/train/", + "datasets/qwen/bench/test/", + ] + ) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + splits = registry.list_dataset_splits("qwen", "bench") + + call_kwargs = mock_bucket.list_objects_v2.call_args[1] + assert call_kwargs["prefix"] == "datasets/qwen/bench/" + assert call_kwargs["delimiter"] == "/" + assert splits == ["test", "train"] + + +def test_list_dataset_splits_returns_empty_when_dataset_missing(): + registry = OssDatasetRegistry(make_registry_info()) + mock_bucket = MagicMock() + mock_bucket.list_objects_v2.return_value = make_list_result(prefixes=[]) + + with patch.object(registry, "_build_bucket", return_value=mock_bucket): + assert registry.list_dataset_splits("qwen", "nope") == [] + + +def test_list_all_datasets_returns_sorted_pairs(): + registry = OssDatasetRegistry(make_registry_info()) + + def fake_list_org_datasets(org): + return {"qwen": ["bench-2", "bench-1"], "alibaba": ["pinch"]}[org] + + with patch.object(registry, "list_organizations", return_value=["qwen", "alibaba"]): + with patch.object(registry, "list_org_datasets", side_effect=fake_list_org_datasets): + pairs = registry.list_all_datasets() + + assert pairs == [("alibaba", "pinch"), ("qwen", "bench-1"), ("qwen", "bench-2")] + + +def test_list_all_datasets_uses_bounded_concurrency(): + registry = OssDatasetRegistry(make_registry_info()) + + with patch.object(registry, "list_organizations", return_value=["o1", "o2"]): + with patch.object(registry, "list_org_datasets", return_value=["d"]): + with patch("rock.sdk.envhub.datasets.registry.oss.ThreadPoolExecutor") as mock_pool: + with patch("rock.sdk.envhub.datasets.registry.oss.as_completed", side_effect=lambda d: list(d)): + mock_executor = MagicMock() + mock_pool.return_value.__enter__.return_value = mock_executor + future = MagicMock() + future.result.return_value = ["d"] + mock_executor.submit.return_value = future + registry.list_all_datasets(concurrency=7) + + mock_pool.assert_called_once_with(max_workers=7) + + +def test_list_all_datasets_default_concurrency_is_10(): + registry = OssDatasetRegistry(make_registry_info()) + + with patch.object(registry, "list_organizations", return_value=["o1"]): + with patch.object(registry, "list_org_datasets", return_value=["d"]): + with patch("rock.sdk.envhub.datasets.registry.oss.ThreadPoolExecutor") as mock_pool: + with patch("rock.sdk.envhub.datasets.registry.oss.as_completed", side_effect=lambda d: list(d)): + mock_executor = MagicMock() + mock_pool.return_value.__enter__.return_value = mock_executor + future = MagicMock() + future.result.return_value = ["d"] + mock_executor.submit.return_value = future + registry.list_all_datasets() + + mock_pool.assert_called_once_with(max_workers=10) + + +def test_list_all_datasets_propagates_exception_from_worker(): + import pytest as _pytest + + registry = OssDatasetRegistry(make_registry_info()) + + def fake_list_org_datasets(org): + if org == "bad": + raise RuntimeError("oss boom") + return ["d"] + + with patch.object(registry, "list_organizations", return_value=["good", "bad"]): + with patch.object(registry, "list_org_datasets", side_effect=fake_list_org_datasets): + with _pytest.raises(RuntimeError, match="oss boom"): + registry.list_all_datasets() + + +def test_list_all_datasets_empty_when_no_orgs(): + registry = OssDatasetRegistry(make_registry_info()) + with patch.object(registry, "list_organizations", return_value=[]): + assert registry.list_all_datasets() == [] From bc2c6537038c54f0f086c04cb5ebbf9c69324c56 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 27 May 2026 14:54:18 +0800 Subject: [PATCH 135/226] fix(sandbox): preserve stop reason in stop() lost in #988 FSM refactor (#1021) The SandboxStateMachine refactor (#988) dropped the `reason` parameter from SandboxManager.stop() and from the on_stop callback chain, breaking the EXPIRED-vs-MANUAL distinction that auto-cleanup relies on (the _check_stop watchdog passes reason=StopReason.EXPIRED to differentiate timed-out sandboxes from user-initiated stops). Restore the parameter and forward it through: - SandboxManager.stop(sandbox_id, reason=StopReason.MANUAL) - dangling path: operator.stop(sandbox_id, reason=reason) - normal path: sm.send("stop", ..., reason=reason) - SandboxStateMachine.on_stop(... , reason=StopReason.MANUAL) - log message: "stop sandbox {id} (reason={reason.value})" - operator.stop(sandbox_id, reason=reason) inside on_stop Behavior matches pre-#988 master. --- rock/sandbox/sandbox_manager.py | 12 +++++++++--- rock/sandbox/sandbox_statemachine.py | 7 ++++--- tests/unit/sandbox/test_sandbox_statemachine.py | 3 ++- tests/unit/sandbox/test_sandbox_transitions.py | 7 ++++--- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 789c267640..0ba97e7a9b 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -178,20 +178,26 @@ async def start(self, config: DeploymentConfig) -> SandboxStartResponse: ) @monitor_sandbox_operation() - async def stop(self, sandbox_id: str): + async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL): sm = await self._get_current_statemachine(sandbox_id) if sm is None: logger.info(f"stop dangling sandbox {sandbox_id}") sandbox_info: SandboxInfo = {"state": State.STOPPED} try: - await self._operator.stop(sandbox_id) + await self._operator.stop(sandbox_id, reason=reason) except ValueError as e: logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) await self._meta_store.archive(sandbox_id, sandbox_info) elif sm.current_state.value == State.STOPPED: await sm.send("stop_noop", sandbox_id=sandbox_id) else: - await sm.send("stop", sandbox_id=sandbox_id, operator=self._operator, meta_store=self._meta_store) + await sm.send( + "stop", + sandbox_id=sandbox_id, + operator=self._operator, + meta_store=self._meta_store, + reason=reason, + ) async def get_mount(self, sandbox_id): async with self._ray_service.get_ray_rwlock().read_lock(): diff --git a/rock/sandbox/sandbox_statemachine.py b/rock/sandbox/sandbox_statemachine.py index 87bfb5069f..889c4cfe77 100644 --- a/rock/sandbox/sandbox_statemachine.py +++ b/rock/sandbox/sandbox_statemachine.py @@ -12,6 +12,7 @@ from rock.actions.sandbox.response import State as RockState from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.metrics.billing import log_billing_info +from rock.common.constants import StopReason from rock.logger import init_logger from rock.utils.system import get_iso8601_timestamp @@ -53,8 +54,8 @@ def __init__(self, **kwargs): # Callbacks - async def on_stop(self, sandbox_id: str, operator, meta_store) -> None: - logger.info(f"stop sandbox {sandbox_id}") + async def on_stop(self, sandbox_id: str, operator, meta_store, reason: StopReason = StopReason.MANUAL) -> None: + logger.info(f"stop sandbox {sandbox_id} (reason={reason.value})") sandbox_info = self.sandbox_info or {} # Initialize sandbox_info with default values if not set @@ -67,7 +68,7 @@ async def on_stop(self, sandbox_id: str, operator, meta_store) -> None: log_billing_info(sandbox_info=sandbox_info) try: - await operator.stop(sandbox_id) + await operator.stop(sandbox_id, reason=reason) except ValueError as e: logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) diff --git a/tests/unit/sandbox/test_sandbox_statemachine.py b/tests/unit/sandbox/test_sandbox_statemachine.py index 435f37ff7e..c9c12613b1 100644 --- a/tests/unit/sandbox/test_sandbox_statemachine.py +++ b/tests/unit/sandbox/test_sandbox_statemachine.py @@ -13,6 +13,7 @@ import pytest from rock.actions.sandbox.response import State +from rock.common.constants import StopReason from rock.sandbox.sandbox_statemachine import SandboxStateMachine # --------------------------------------------------------------------------- @@ -158,7 +159,7 @@ def mock_meta_store(self): async def test_stops_operator_and_archives(self, mock_operator, mock_meta_store): sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=mock_meta_store) - mock_operator.stop.assert_awaited_once_with("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) mock_meta_store.archive.assert_awaited_once() @pytest.mark.asyncio diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py index 67b6386a34..320660f855 100644 --- a/tests/unit/sandbox/test_sandbox_transitions.py +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -9,6 +9,7 @@ import pytest from rock.actions.sandbox.response import State +from rock.common.constants import StopReason from rock.sandbox.sandbox_manager import SandboxManager from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError @@ -75,7 +76,7 @@ class TestManagerStop: async def test_stop_not_found_attempts_cleanup(self, mgr, mock_meta_store, mock_operator): mock_meta_store.get.return_value = None await mgr.stop("sb-1") - mock_operator.stop.assert_awaited_once_with("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) mock_meta_store.archive.assert_awaited_once() @pytest.mark.asyncio @@ -95,14 +96,14 @@ async def test_stop_already_stopped_is_noop(self, mgr, mock_meta_store, mock_ope async def test_stop_running_calls_operator_and_archives(self, mgr, mock_meta_store, mock_operator): mock_meta_store.get.return_value = {"state": State.RUNNING} await mgr.stop("sb-1") - mock_operator.stop.assert_awaited_once_with("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) mock_meta_store.archive.assert_awaited_once() @pytest.mark.asyncio async def test_stop_pending_calls_operator_and_archives(self, mgr, mock_meta_store, mock_operator): mock_meta_store.get.return_value = {"state": State.PENDING} await mgr.stop("sb-1") - mock_operator.stop.assert_awaited_once_with("sb-1") + mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) mock_meta_store.archive.assert_awaited_once() @pytest.mark.asyncio From dca3cfe4b353d3d0e74cb46f5b37cff582f6e178 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 27 May 2026 14:58:00 +0800 Subject: [PATCH 136/226] refactor(deployments): split `docker run` into `docker create` + `docker start -a` (#1012) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the one-shot `docker run` in start() with `docker create` followed by `docker start -a`. Equivalent at the docker engine level (same args, same stdout/stderr behavior with -a) but introduces a window between the two steps where per-container host state can be set up before any process in the container can run — used by the next commit to attach the docker-allocated XFS prjid to the host log dir. Helpers in rock/deployments/docker.py: - _docker_create: synchronous `docker create`; raises on failure. Container is left in `created` state on success. - _docker_start: replaces _docker_run. `docker start -a ` via Popen; the returned handle behaves like the old _docker_run handle for _wait_until_alive (.poll() / .stdout / .stderr) and for _stop (`docker kill ` first, then kill the Popen — same as before). Cleanup contract around the create→start gap: - DockerUtil.remove_container_force (new): generic best-effort `docker rm -f`; swallows errors with a warning. - start() wraps _docker_start in try/except: if create succeeded but start failed before _container_process is set, the orphan `created` container is removed before re-raising. _wait_until_alive's own failure path is unchanged — it still calls self.stop(), which manages the container via the live _container_process handle. - _docker_start itself does NOT clean up on failure; the caller knows whether the container is freshly created (remove) or pre-existing (keep — leaves room for a future restart path that reuses _docker_start). --- rock/deployments/docker.py | 34 +++++++++++++++++++++++++++++----- rock/utils/docker.py | 16 ++++++++++++++++ 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 9c76a0ad82..15b58500b7 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -539,7 +539,7 @@ async def start(self): runtime_args = self._build_runtime_args() cmds = [ "docker", - "run", + "create", "--entrypoint", "", *env_arg, @@ -568,8 +568,16 @@ async def start(self): ) logger.info(f"Command: {cmd_str!r}") # shell=True required for && etc. - with StageTimer("startup_timing", f"[{self._container_name}] Docker run", logger): - self._container_process = await loop.run_in_executor(executor, self._docker_run, cmds) + with StageTimer("startup_timing", f"[{self._container_name}] Docker start", logger): + await loop.run_in_executor(executor, self._docker_create, cmds) + # After docker create succeeds, the container exists in `created` state. If anything below + # fails before _wait_until_alive sets _container_process up for _stop to manage, we own the + # orphan and must remove it. _wait_until_alive's own failure path already calls self.stop(). + try: + self._container_process = await loop.run_in_executor(executor, self._docker_start) + except Exception: + DockerUtil.remove_container_force(self._container_name) + raise await loop.run_in_executor(executor, self._hooks.on_custom_step, DeploymentHookStep.STARTING_RUNTIME) logger.info(f"Starting runtime at {self._config.port}") self._runtime = RemoteSandboxRuntime.from_config( @@ -603,9 +611,25 @@ def _prepare_volume_mounts(self) -> list[str]: logger.info(f"volume_args: {volume_args}") return volume_args - def _docker_run(self, cmd: list[str]): + def _docker_create(self, cmd: list[str]) -> None: + """Create the container without starting it.""" + try: + subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, timeout=60) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + logger.error(f"Failed to create container {self._container_name}") + self._service_status.update_status( + phase_name="docker_run", status=Status.FAILED, message="docker run failed" + ) + raise + + def _docker_start(self) -> subprocess.Popen: + """Start a previously-created container with stdout/stderr attached.""" try: - exec_rlt = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + exec_rlt = subprocess.Popen( + ["docker", "start", "-a", self._container_name], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) self._service_status.update_status( phase_name="docker_run", status=Status.RUNNING, message="docker run running" ) diff --git a/rock/utils/docker.py b/rock/utils/docker.py index f8a67fd87e..96444880b4 100644 --- a/rock/utils/docker.py +++ b/rock/utils/docker.py @@ -210,6 +210,22 @@ def remove_image(cls, image: str) -> bytes: """Remove a Docker image""" return subprocess.check_output(["docker", "rmi", image], timeout=30) + @classmethod + def remove_container_force(cls, name: str, timeout: int = 10) -> None: + """Swallows errors (missing container, daemon hiccups, timeout) and only + warns — intended for cleanup paths where the caller's primary error + must still propagate. + """ + try: + subprocess.run( + ["docker", "rm", "-f", name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=timeout, + ) + except Exception as e: + logger.warning(f"Remove of container {name} failed: {e}") + class ImageUtil: """Docker image name utilities""" From 98c4a9146c8364046b64050b6c2b428a6dcb4903 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Wed, 27 May 2026 17:40:40 +0800 Subject: [PATCH 137/226] fix(rocklet): use cgroup metrics for container memory instead of psutil (#1017) * fix(rocklet): use cgroup metrics for container memory instead of psutil Same approach as CPU: read memory.current/memory.max (cgroup v2) or memory.usage_in_bytes/memory.limit_in_bytes (cgroup v1), falling back to psutil when running outside a container or with unlimited memory. Co-Authored-By: Claude Opus 4.6 * fix(rocklet): exclude reclaimable page cache from container memory usage Subtract inactive_file (cgroup v2) / total_inactive_file (v1) from raw cgroup memory counters to match the working-set definition used by docker stats and kubelet, so reported memory reflects application pressure rather than reclaimable file cache. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.6 --- rock/rocklet/linux.py | 5 +- rock/utils/cgroup_stats.py | 94 +++++++++++- tests/unit/utils/test_cgroup_stats.py | 204 +++++++++++++++++++++++++- 3 files changed, 296 insertions(+), 7 deletions(-) diff --git a/rock/rocklet/linux.py b/rock/rocklet/linux.py index 2e14faee09..b2f1149e24 100644 --- a/rock/rocklet/linux.py +++ b/rock/rocklet/linux.py @@ -34,7 +34,7 @@ SessionNotInitializedError, ) from rock.utils import get_executor -from rock.utils.cgroup_stats import CgroupCpuStats +from rock.utils.cgroup_stats import CgroupCpuStats, CgroupMemStats from .rocklet import Rocklet, Session @@ -346,6 +346,7 @@ class LinuxRocklet(Rocklet): def __init__(self, **kwargs): super().__init__(**kwargs) self._cgroup_cpu = CgroupCpuStats() + self._cgroup_mem = CgroupMemStats() self._docker_data_root: str | None = None def _build_bash_session(self, request: CreateBashSessionRequest) -> Session: @@ -372,7 +373,7 @@ async def get_statistics(self) -> dict: return { "cpu": self._cgroup_cpu.cpu_percent(), - "mem": psutil.virtual_memory().percent, + "mem": self._cgroup_mem.mem_percent(), "disk": disk_root.percent, # legacy metric name, actually rootfs usage percent "disk_log_percent": disk_log_percent, "disk_dind_percent": disk_dind_percent, diff --git a/rock/utils/cgroup_stats.py b/rock/utils/cgroup_stats.py index 233a414191..e5cf1bc834 100644 --- a/rock/utils/cgroup_stats.py +++ b/rock/utils/cgroup_stats.py @@ -1,4 +1,4 @@ -"""Container-aware CPU metrics via cgroup v1/v2. +"""Container-aware CPU/memory metrics via cgroup v1/v2. Falls back to psutil when cgroup files are unavailable (e.g. running outside a container or on non-Linux platforms). @@ -111,3 +111,95 @@ def cpu_percent(self) -> float: num_cpus = self._read_cpu_quota() return min(round((delta_usage / delta_time) / num_cpus * 100, 1), 100.0) + + +class CgroupMemStats: + """Reads container memory utilization from cgroup v1/v2 pseudo-files.""" + + def __init__(self): + self._cgroup_version: int | None = None + self._mem_limit: int | None = None + + def _detect_cgroup_version(self) -> int: + if self._cgroup_version is not None: + return self._cgroup_version + + if Path("/sys/fs/cgroup/cgroup.controllers").exists(): + self._cgroup_version = 2 + elif Path("/sys/fs/cgroup/memory/memory.usage_in_bytes").exists(): + self._cgroup_version = 1 + else: + self._cgroup_version = 0 + + return self._cgroup_version + + def _read_mem_stat(self, stat_path: str, key: str) -> int: + """Read a single counter from a cgroup memory.stat file. Returns 0 on any error.""" + try: + for line in Path(stat_path).read_text().splitlines(): + parts = line.split() + if len(parts) >= 2 and parts[0] == key: + return int(parts[1]) + except Exception: + pass + return 0 + + def _read_mem_usage_bytes(self) -> int | None: + """Return container memory usage minus reclaimable page cache (inactive_file). + + Raw cgroup usage counters include page cache, which is reclaimable and not + a meaningful signal of application memory pressure. Subtracting inactive_file + matches the working-set definition used by docker stats and kubelet. + """ + try: + ver = self._detect_cgroup_version() + if ver == 2: + usage = int(Path("/sys/fs/cgroup/memory.current").read_text().strip()) + inactive_file = self._read_mem_stat("/sys/fs/cgroup/memory.stat", "inactive_file") + return max(usage - inactive_file, 0) + elif ver == 1: + usage = int(Path("/sys/fs/cgroup/memory/memory.usage_in_bytes").read_text().strip()) + inactive_file = self._read_mem_stat("/sys/fs/cgroup/memory/memory.stat", "total_inactive_file") + return max(usage - inactive_file, 0) + return None + except Exception: + return None + + def _read_mem_limit_bytes(self) -> int | None: + if self._mem_limit is not None: + return self._mem_limit + + try: + ver = self._detect_cgroup_version() + if ver == 2: + text = Path("/sys/fs/cgroup/memory.max").read_text().strip() + if text == "max": + return None + limit = int(text) + if limit <= 0: + return None + self._mem_limit = limit + return self._mem_limit + elif ver == 1: + limit = int(Path("/sys/fs/cgroup/memory/memory.limit_in_bytes").read_text().strip()) + # cgroup v1 uses a very large number (like PAGE_COUNTER_MAX) for unlimited + if limit <= 0 or limit >= (1 << 62): + return None + self._mem_limit = limit + return self._mem_limit + except Exception: + pass + return None + + def mem_percent(self) -> float: + """Return container memory utilization %. + + Falls back to psutil if cgroup files are unavailable or memory is unlimited. + """ + usage = self._read_mem_usage_bytes() + limit = self._read_mem_limit_bytes() + + if usage is None or limit is None: + return psutil.virtual_memory().percent + + return min(round(usage / limit * 100, 1), 100.0) diff --git a/tests/unit/utils/test_cgroup_stats.py b/tests/unit/utils/test_cgroup_stats.py index b996268279..157e5ceb41 100644 --- a/tests/unit/utils/test_cgroup_stats.py +++ b/tests/unit/utils/test_cgroup_stats.py @@ -1,10 +1,8 @@ -"""Tests for rock.utils.cgroup_stats — container-aware CPU metrics.""" +"""Tests for rock.utils.cgroup_stats — container-aware CPU/memory metrics.""" from unittest.mock import patch -import pytest - -from rock.utils.cgroup_stats import CgroupCpuStats, Path +from rock.utils.cgroup_stats import CgroupCpuStats, CgroupMemStats, Path def _mock_path_exists(mapping: dict[str, bool]): @@ -286,3 +284,201 @@ def test_fallback_on_error(self): patch("rock.utils.cgroup_stats.os.cpu_count", return_value=2), ): assert stats._read_cpu_quota() == 2.0 + + +# ---------- Memory cgroup version detection ---------- + + +class TestMemDetectCgroupVersion: + def test_detects_v2(self): + stats = CgroupMemStats() + exists_map = {"/sys/fs/cgroup/cgroup.controllers": True} + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 2 + + def test_detects_v1(self): + stats = CgroupMemStats() + exists_map = { + "/sys/fs/cgroup/cgroup.controllers": False, + "/sys/fs/cgroup/memory/memory.usage_in_bytes": True, + } + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 1 + + def test_detects_no_cgroup(self): + stats = CgroupMemStats() + exists_map = { + "/sys/fs/cgroup/cgroup.controllers": False, + "/sys/fs/cgroup/memory/memory.usage_in_bytes": False, + } + with patch.object(Path, "exists", _mock_path_exists(exists_map)): + assert stats._detect_cgroup_version() == 0 + + def test_caches_result(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + assert stats._detect_cgroup_version() == 2 + + +# ---------- Memory percent ---------- + + +class TestMemPercent: + def test_v2_mem_calculation(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "524288000\n", + "/sys/fs/cgroup/memory.max": "1073741824\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + # 524288000 / 1073741824 * 100 = 48.8% + assert result == 48.8 + + def test_v1_mem_calculation(self): + stats = CgroupMemStats() + stats._cgroup_version = 1 + read_map = { + "/sys/fs/cgroup/memory/memory.usage_in_bytes": "734003200\n", + "/sys/fs/cgroup/memory/memory.limit_in_bytes": "1073741824\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + # 734003200 / 1073741824 * 100 = 68.4% + assert result == 68.4 + + def test_fallback_to_psutil_when_no_cgroup(self): + stats = CgroupMemStats() + stats._cgroup_version = 0 + with patch("rock.utils.cgroup_stats.psutil.virtual_memory") as mock_vmem: + mock_vmem.return_value.percent = 55.3 + assert stats.mem_percent() == 55.3 + + def test_fallback_to_psutil_when_unlimited_v2(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "524288000\n", + "/sys/fs/cgroup/memory.max": "max\n", + } + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map)), + patch("rock.utils.cgroup_stats.psutil.virtual_memory") as mock_vmem, + ): + mock_vmem.return_value.percent = 33.0 + assert stats.mem_percent() == 33.0 + + def test_fallback_to_psutil_when_unlimited_v1(self): + stats = CgroupMemStats() + stats._cgroup_version = 1 + read_map = { + "/sys/fs/cgroup/memory/memory.usage_in_bytes": "524288000\n", + "/sys/fs/cgroup/memory/memory.limit_in_bytes": "9223372036854771712\n", + } + with ( + patch.object(Path, "read_text", _mock_path_read_text(read_map)), + patch("rock.utils.cgroup_stats.psutil.virtual_memory") as mock_vmem, + ): + mock_vmem.return_value.percent = 45.0 + assert stats.mem_percent() == 45.0 + + def test_capped_at_100(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + stats._mem_limit = 1073741824 + read_map = { + "/sys/fs/cgroup/memory.current": "2147483648\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + assert result == 100.0 + + def test_caches_mem_limit(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "524288000\n", + "/sys/fs/cgroup/memory.max": "1073741824\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + stats.mem_percent() + assert stats._mem_limit == 1073741824 + + def test_fallback_to_psutil_on_read_error(self): + stats = CgroupMemStats() + stats._cgroup_version = 2 + with ( + patch.object(Path, "read_text", side_effect=PermissionError), + patch("rock.utils.cgroup_stats.psutil.virtual_memory") as mock_vmem, + ): + mock_vmem.return_value.percent = 60.0 + assert stats.mem_percent() == 60.0 + + def test_v2_subtracts_inactive_file_cache(self): + """memory.current includes page cache; inactive_file should be subtracted.""" + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "1073741824\n", + "/sys/fs/cgroup/memory.stat": "anon 524288000\nfile 400000000\ninactive_file 262144000\nactive_file 137856000\n", + "/sys/fs/cgroup/memory.max": "2147483648\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + # (1073741824 - 262144000) / 2147483648 * 100 = 37.8% + assert result == 37.8 + + def test_v1_subtracts_total_inactive_file_cache(self): + """memory.usage_in_bytes includes page cache; total_inactive_file should be subtracted.""" + stats = CgroupMemStats() + stats._cgroup_version = 1 + read_map = { + "/sys/fs/cgroup/memory/memory.usage_in_bytes": "1073741824\n", + "/sys/fs/cgroup/memory/memory.stat": "cache 400000000\ntotal_cache 400000000\ntotal_inactive_file 262144000\n", + "/sys/fs/cgroup/memory/memory.limit_in_bytes": "2147483648\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + # (1073741824 - 262144000) / 2147483648 * 100 = 37.8% + assert result == 37.8 + + def test_v2_no_subtraction_when_memory_stat_missing(self): + """If memory.stat is unreadable, fall back to raw usage rather than crashing.""" + stats = CgroupMemStats() + stats._cgroup_version = 2 + # No memory.stat entry — _mock_path_read_text raises FileNotFoundError. + read_map = { + "/sys/fs/cgroup/memory.current": "524288000\n", + "/sys/fs/cgroup/memory.max": "1073741824\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + assert result == 48.8 + + def test_v2_inactive_file_larger_than_usage_clamped_to_zero(self): + """Sanity guard: subtraction must not produce a negative usage.""" + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "100\n", + "/sys/fs/cgroup/memory.stat": "inactive_file 9999\n", + "/sys/fs/cgroup/memory.max": "1000\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + assert result == 0.0 + + def test_v2_ignores_malformed_memory_stat_lines(self): + """memory.stat parser should skip blank/malformed lines without raising.""" + stats = CgroupMemStats() + stats._cgroup_version = 2 + read_map = { + "/sys/fs/cgroup/memory.current": "1000\n", + "/sys/fs/cgroup/memory.stat": "\nmalformed\ninactive_file 200\nanon\n", + "/sys/fs/cgroup/memory.max": "10000\n", + } + with patch.object(Path, "read_text", _mock_path_read_text(read_map)): + result = stats.mem_percent() + # (1000 - 200) / 10000 * 100 = 8.0% + assert result == 8.0 From dd5b6e4d3d521d6e2f924488d56afb87ff884137 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Wed, 27 May 2026 20:08:31 +0800 Subject: [PATCH 138/226] docs(v1.8.x): add sandbox concurrent-creation benchmark report & scheduler guide (#1035) (#1036) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * add benchmark doc & copy scheduler guide --- .../User Guides/sandbox_create_benchmark.md | 43 +++ .../version-1.8.x/User Guides/scheduler.md | 283 ++++++++++++++++++ .../sandbox_create_latency_distribution.png | Bin 0 -> 100042 bytes .../User Guides/sandbox_create_benchmark.md | 43 +++ .../version-1.8.x/User Guides/scheduler.md | 283 ++++++++++++++++++ 5 files changed, 652 insertions(+) create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/sandbox_create_benchmark.md create mode 100644 docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/scheduler.md create mode 100644 docs/static/img/sandbox_create_latency_distribution.png create mode 100644 docs/versioned_docs/version-1.8.x/User Guides/sandbox_create_benchmark.md create mode 100644 docs/versioned_docs/version-1.8.x/User Guides/scheduler.md diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/sandbox_create_benchmark.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/sandbox_create_benchmark.md new file mode 100644 index 0000000000..a5bbb0f361 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/sandbox_create_benchmark.md @@ -0,0 +1,43 @@ +# ROCK Sandbox 大规模并发创建压测报告 + +## 1. 测试背景 + +在 Agentic 强化学习训练以及大规模 Agent Rollout 场景中，**单个训练 step 往往需要同时拉起成千上万个 Sandbox**，因此 Sandbox 的并发创建吞吐与延迟直接决定了整体训练效率。 + +本报告分别在 1000 / 2000 / 4000 / 8000 / 16000 并发规模下完成 Sandbox 批量创建压测，验证 ROCK 在大规模并发下的稳定性与延迟表现。 + +## 2. 测试范围与口径 + +- **被测对象**：ROCK Sandbox 的创建链路。 +- **统计指标**：单个 Sandbox 从发起创建请求到 Sandbox 处于存活（可用）状态的端到端耗时（秒）。 +- **并发规模**：1000、2000、4000、8000、16000 个 Sandbox 同时发起创建。 +- **并发模型**：采用 **多机分布式 + 纯多进程** 的方式驱动并发——由多台机器同时发压，每台机器内部以独立 OS 进程承载每个 Sandbox 创建任务，避免单进程内 GIL、事件循环、TLS 握手等成为客户端侧瓶颈，使数据真实反映服务端的处理能力。 + +## 3. 测试结果 + +### 3.1 Sandbox 创建耗时统计 + +| 并发规模 | 样本数 | 成功率 | 最小值 | 最大值 | 平均值 | P50 | P95 | P99 | +|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| **1000** | 1000 | 100% | 0.47s | 9.84s | 4.72s | 4.93s | 6.89s | 8.69s | +| **2000** | 2000 | 100% | 0.64s | 11.20s | 4.90s | 4.92s | 9.76s | 10.53s | +| **4000** | 4000 | 100% | 0.93s | 15.51s | 7.16s | 7.05s | 11.26s | 13.34s | +| **8000** | 8000 | 100% | 3.85s | 33.99s | 18.06s | 17.86s | 30.09s | 32.03s | +| **16000** | 16000 | 100% | 2.66s | 63.84s | 37.17s | 39.98s | 56.68s | 60.11s | + +> 时间单位均为秒（s）。所有规模下成功率均为 **100%**（0 失败）。 +> +> 说明：随着并发规模的增大，为保护服务端稳定性，ROCK 会在控制面侧进行限流，因此耗时随并发上升而有所增加。 + +![Sandbox 创建耗时分布](../../../../../static/img/sandbox_create_latency_distribution.png) + +## 4. 结论 + +ROCK 在 1000 至 16000 并发规模下完成的 Sandbox 批量创建压测取得了以下结果： + +1. **100% 成功率**，验证了 ROCK 在大规模并发下的可靠性。 +2. **小规模并发（≤2000）几乎无排队**，P50 稳定在 5s 以内，可以满足绝大多数实时性敏感的训练 / 评估场景。 +3. **大规模并发（≥4000）延迟随并发数近似线性增长**，行为可预测，便于按 Sandbox 池规模做容量规划。 +4. **16000 并发场景下 P99 仍可控制在 60s 量级**，足以支撑超大规模 Agent Rollout 与并行 RL Step 等任务。 + +ROCK 已经具备承载 **万级并发 Sandbox 创建** 的能力，为大规模 Agentic RL 训练提供了坚实的环境基础设施。 diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/scheduler.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/scheduler.md new file mode 100644 index 0000000000..88662edd51 --- /dev/null +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/User Guides/scheduler.md @@ -0,0 +1,283 @@ +--- +sidebar_position: 5 +--- + +# 任务调度器(Scheduler) + +ROCK 调度器是内嵌于 `admin` 服务中的周期性任务框架。它会按可配置的时间间隔,把后台维护任务(镜像清理、文件清理、容器清理、镜像预拉取、自定义任务……)分发到所有存活的 Ray worker 上,从而在无人工干预的情况下保持 worker 节点健康。 + +本文介绍如何启用调度器、配置内置任务、编写自定义任务,以及如何观测运行状态。 + +## 1. 工作原理 + +- 调度器以独立的守护线程(`SchedulerThread`)运行在 `admin` 进程内,使用自己的 `asyncio` 事件循环。 +- 任务通过 [APScheduler](https://apscheduler.readthedocs.io/) 以固定间隔(`interval_seconds`)触发。 +- 每次触发时,调度器会先获取存活 Ray worker 列表(由 `worker_cache_ttl` 秒级缓存),然后并发地把任务下发到每个 worker(默认并发度:50)。 +- **下发动作通过 worker 上的 rocklet 服务以 HTTP 方式完成**:admin 端构造 `RemoteSandboxRuntime(host=worker_ip, port=Port.PROXY)`(参见 `rock.deployments.constants.Port`),并调用 `runtime.execute / read_file / write_file`。**因此每个 worker 都必须运行 `rocklet` 服务并在 `Port.PROXY` 上可达**,否则调度器无法下发命令、也无法在 worker 上读写状态文件。 +- 每个任务都继承自 `rock.admin.scheduler.task_base.BaseTask`,并必须实现 `run_action(runtime: RemoteSandboxRuntime)` —— 单个 worker 上的实际执行逻辑。 +- 每个 worker 的执行状态会被持久化到 `ROCK_SCHEDULER_STATUS_DIR`(默认 `/data/scheduler_status`)目录下的 JSON 文件中;每次执行结束后还会写入聚合报告 `/_run_report.json`。 +- 当配置了 Nacos 配置源时,调度器会订阅配置变更,并按 diff 应用:仅 hash 发生变化的任务被重新安装,被删除的任务会同步从所有 worker 上清理。 + +### 前置条件 + +启用调度器之前,请确认每个 Ray worker 满足以下条件: + +| 条件 | 原因 | +|------|------| +| worker 上正在运行 `rocklet` 进程 | 调度器通过 rocklet HTTP 接口下发每一个任务;若 rocklet 不存在,`runtime.execute` 调用会超时。 | +| admin 可访问 rocklet 监听端口 | 调度器固定使用 `Port.PROXY`(定义于 `rock.deployments.constants.Port`)作为下发目标,请确认防火墙 / 安全组未阻断该端口。 | +| `ROCK_SCHEDULER_STATUS_DIR` 在 worker 内可写 | 任务会在该目录下读写 `_status.json`,用于幂等控制和 PID 跟踪。 | +| 任务依赖的工具在 worker 上可用 | 例如清理 / 拉取类任务需要 `docker`;`ImageCleanupTask` 首次运行会通过 `curl` 联网安装 `docuum`。 | + +rocklet 服务由 worker 标准启动脚本(`docker_run.sh`、`docker_run_with_uv.sh`、`docker_run_with_pip.sh`)自动拉起,通常等价于 `rocklet --port `。如果你使用了自定义 entrypoint 来启动 worker,请确保等价命令被执行。具体的运行时类型与 rocklet 启动方式可参考 [Configuration](./configuration.md)。 + +### 幂等性 + +每个任务都需要声明自己的幂等模式,该模式直接影响重复触发时的行为: + +| 模式 | 行为 | +|------|------| +| `IDEMPOTENT` | 每次 tick 都会执行,可安全重复(例如 `docker pull`、`find -exec rm`)。 | +| `NON_IDEMPOTENT` | 任务会拉起一个后台守护进程(例如 `docuum`)。调度器会读取上一次的状态文件,检查记录的 PID 是否仍存活,若仍在运行则跳过本次启动。当任务从配置中移除时,调度器会通过 `pkill` 杀掉该进程。 | + +## 2. 启用调度器 + +调度器配置位于 ROCK admin YAML 顶层 `scheduler:` 字段下(例如 `rock-conf/rock-local.yml`、`rock-conf/rock-dev.yml`)。 + +```yaml +scheduler: + enabled: true # 总开关 + worker_cache_ttl: 43200 # Worker IP 缓存 TTL(秒) + tasks: + # ... 任务列表,详见下文 +``` + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `enabled` | bool | `false` | 总开关。设为 `false` 时所有任务会被卸载,且不再触发。 | +| `worker_cache_ttl` | int | `3600` | 存活 worker IP 列表缓存时长(秒);超过该时长后会从 `ray.nodes()` 重新获取。 | +| `tasks` | list | `[]` | `TaskConfig` 列表,详见 [第 4 节](#4-任务配置schema)。 | + +### 相关环境变量 + +| 变量 | 默认值 | 用途 | +|------|--------|------| +| `ROCK_SCHEDULER_STATUS_DIR` | `/data/scheduler_status` | worker 上写入任务状态 JSON 与执行报告的目录。 | +| `ROCK_LOGGING_PATH` | (未设置) | 设置后,调度器拉起的守护进程(docuum、container_cleanup、image_pull)会将 stdout/stderr 重定向到 `/.log`。 | +| `ROCK_DOCUUM_INSTALL_URL` | `https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh` | `ImageCleanupTask` 按需拉取 `docuum` 安装脚本的 URL。 | + +## 3. 内置任务 + +ROCK 在 `rock.admin.scheduler.tasks` 下提供了 4 个内置任务,通过把 `task_class` 设置为对应的全限定类路径即可注册。 + +### 3.1 ImageCleanupTask + +在每个 worker 上运行 [`docuum`](https://github.com/stepchowfun/docuum),当磁盘占用超过阈值时按 LRU 策略淘汰镜像。**非幂等** —— `docuum` 是常驻守护进程;调度器会跟踪其 PID,只要进程仍存活就跳过重复拉起。 + +```yaml +- task_class: rock.admin.scheduler.tasks.image_cleanup_task.ImageCleanupTask + enabled: true + interval_seconds: 43200 # 每 12 小时检查一次守护进程 + params: + disk_threshold: "70%" # 磁盘占用超过 70% 时触发淘汰 + image_whitelist: # 匹配 repository:tag 的 glob 模式,白名单内的镜像不会被淘汰 + - "python:3.11" + - "my-registry.example.com/base/*" +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `disk_threshold` | str | `"1T"` | 传给 `docuum --threshold` 的磁盘阈值,支持容量(`100G`、`1T`)或百分比(`70%`)。 | +| `image_whitelist` | list[str] | `[]` | 透传给 `docuum --keep` 的 glob 模式列表。 | + +### 3.2 FileCleanupTask + +遍历配置的目录,删除超过 `max_age_mins` 或大于 `max_file_size` 的文件,然后清理留下的空目录。**幂等**。 + +```yaml +- task_class: rock.admin.scheduler.tasks.file_cleanup_task.FileCleanupTask + enabled: true + interval_seconds: 86400 # 每天执行一次 + params: + target_dirs: + # 字符串形式 —— 不配置排除项 + - "/data/service_status" + # 对象形式 —— 配置该目录独有的排除项 + - path: "/data/logs" + exclude_files: # 支持纯文件名 / 相对路径 / 绝对路径 + - "docuum.log" + - "./rocklet.log" + - "./access.log" + exclude_dirs: + - ".cache" + max_age_mins: 10080 # 7 天,超出此时间的文件会被删除 + max_file_size: "1G" # 大于此大小的文件会被删除(支持 K/M/G/T) +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `target_dirs` | list | `[]` | 每个条目是一个字符串(只填路径)或 `{path, exclude_files, exclude_dirs}`。 | +| `max_age_mins` | int | `10080` | mtime 早于该分钟数的文件会被删除。 | +| `max_file_size` | str | `"1G"` | 大于该阈值的文件会被删除,支持 `K/M/G/T` 单位。 | + +删除条件为 `(-mmin +max_age_mins) OR (-size +max_file_size)`。文件清理后,会再用 `find -depth -type d -empty -delete` 清理留下的空目录(同样遵循 `exclude_dirs` 配置)。 + +### 3.3 ContainerCleanupTask + +删除停止时间超过指定时长的 Docker 容器,避免 worker 上的容器列表无限增长。**幂等**。 + +```yaml +- task_class: rock.admin.scheduler.tasks.container_cleanup_task.ContainerCleanupTask + enabled: true + interval_seconds: 86400 + params: + max_age_hours: 72 # 删除超过 72 小时的 exited 容器 +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `max_age_hours` | int | `24` | 已退出容器的最大保留时间(以 `FinishedAt` 起算的小时数);超出后被 `docker rm`。 | + +每次执行还会顺带清理处于 `created` 状态(从未启动)的容器。 + +### 3.4 ImagePullTask + +在每个 worker 上预拉取一组 Docker 镜像,并可选地先登录私有仓库,以降低沙箱冷启动延迟。**幂等**(若镜像已是最新,`docker pull` 等同于空操作)。 + +```yaml +- task_class: rock.admin.scheduler.tasks.image_pull_task.ImagePullTask + enabled: true + interval_seconds: 21600 # 每 6 小时刷新一次 + params: + images: + # 字符串形式 —— 公开镜像,无需鉴权 + - "python:3.11" + # 对象形式 —— 私有镜像,需要登录 + - image: "my-registry.example.com/chatos/python:313" + registry_username: "myuser" + registry_password: "bXlwYXNzd29yZA==" # base64 编码 +``` + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `images` | list | `[]` | 每个条目是一个镜像字符串或 `{image, registry_username, registry_password}`。 | + +`registry_password` 必须是 base64 编码,worker 端会先解码再通过 `docker login --password-stdin` 登录。仓库地址会从镜像名称中解析,因此每个镜像可以指向不同的仓库。 + +## 4. 任务配置 Schema + +`scheduler.tasks` 下的每一项都会被解析为 `rock.config.TaskConfig`: + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `task_class` | str | `""` | Python 类的全限定路径,**必填**。 | +| `enabled` | bool | `true` | 设为 `false` 的任务在加载阶段会被跳过,在 reload 阶段会被卸载。 | +| `interval_seconds` | int | `3600` | APScheduler `interval` 间隔(秒)。 | +| `params` | dict | `{}` | 任务自定义参数,会在 `from_config()` 中被消费。 | + +只要某个任务条目中任何字段发生变化,调度器就会卸载旧任务(非幂等任务还会清理 worker 上的守护进程与状态文件)再安装新任务,**整个过程不需要重启 admin 进程**。 + +## 5. 编写自定义任务 + +任何位于 Python 路径下、继承自 `BaseTask` 的类都可以注册为调度任务。最小契约示例如下: + +```python +# my_pkg/my_tasks/disk_report_task.py +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + + +class DiskReportTask(BaseTask): + """记录每个 worker 的 `df -h` 输出。""" + + def __init__(self, interval_seconds: int = 3600, mount_point: str = "/"): + super().__init__( + type="disk_report", # 同时作为 APScheduler job id 与状态文件名前缀 + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + self.mount_point = mount_point + + @classmethod + def from_config(cls, task_config) -> "DiskReportTask": + return cls( + interval_seconds=task_config.interval_seconds, + mount_point=task_config.params.get("mount_point", "/"), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + result = await runtime.execute( + Command(command=f"df -h {self.mount_point}", shell=True), + ) + return { + "status": TaskStatusEnum.SUCCESS, + "exit_code": result.exit_code, + "stdout": result.stdout, + } +``` + +随后在 YAML 中注册: + +```yaml +scheduler: + enabled: true + tasks: + - task_class: my_pkg.my_tasks.disk_report_task.DiskReportTask + enabled: true + interval_seconds: 600 + params: + mount_point: "/data" +``` + +### 自定义任务编写要点 + +- **`super().__init__()` 中的 `type` 必须全局唯一**:它会同时被用作 APScheduler job id、状态文件名(`_status.json`)与执行报告文件名(`_run_report.json`)。两个任务不能共用同一个 `type`。 +- **正确选择 `IdempotencyType`**: + - 当 `run_action` 同步执行完毕、可安全重入时,使用 `IDEMPOTENT`。 + - 当任务通过 `nohup` 拉起常驻守护进程并返回 PID 时,使用 `NON_IDEMPOTENT`;调度器会跟踪 PID,在其存活期间跳过重复启动,在卸载时通过 `pkill` 杀掉。 +- **`run_action` 必须返回 dict**,推荐字段: + - `status` —— `TaskStatusEnum` 值,会写入状态文件。 + - `pid` —— `NON_IDEMPOTENT` 守护型任务必须返回(可使用 `rock.utils.system.extract_nohup_pid` 从 `nohup ... & echo PID_PREFIX${{!}}PID_SUFFIX` 输出中提取)。 + - 其他诊断字段会落到状态文件的 `extra` 块里。 +- **重写 `from_config(cls, task_config)`** 用于把 `task_config.params` 翻译成 `__init__` 的入参。 +- **使用 `runtime.execute / read_file / write_file`** 与 worker 通信,**不要**在本地直接执行 shell —— 调度器是把任务下发到远端 worker 的 `RemoteSandboxRuntime`。 + +## 6. 可观测性 + +每个任务会在每个 worker 上产生两类产物: + +| 路径 | 写入方 | 内容 | +|------|--------|------| +| `/_status.json` | `BaseTask.save_task_status` | 单 worker 最新状态:`task_name`、`worker_ip`、`pid`、`status`(`pending`/`running`/`success`/`failed`)、`last_run`、`error`,以及任务自定义的 `extra` 字段。 | +| `/_run_report.json` | `BaseTask.run`(由 admin 端在本轮 tick 结束后写入) | 聚合报告:总数 / 成功数 / 失败数、`success_ips` 列表、`failed_details`(`ip` + 错误堆栈)。 | + +调度器内部日志会输出到 ROCK admin 标准日志路径下,logger 名称包括 `name="scheduler"`、`name="task_base"`、`name="image_clean"` 等。当设置了 `ROCK_LOGGING_PATH` 时,被调度器拉起的守护进程会把自身日志写入 `/.log`(例如 `docuum.log`、`container_cleanup.log`、`image_pull.log`)。 + +## 7. 通过 Nacos 动态热更(可选) + +当 admin 服务启用了 Nacos 配置源时,调度器会注册一个 YAML 监听器,并对配置推送做出响应: + +- 仅检查 `scheduler:` 段,其他段被忽略。 +- 新配置段会被计算 hash 并与上一次的 hash 比对 —— 重复推送会被自动跳过。 +- 通过对新旧任务列表 diff,决定哪些任务需要安装、卸载或重新安装(`params` / `interval_seconds` / `enabled` 任意改动都会触发)。 +- 被删除或重新安装的非幂等任务会先做清理:杀掉守护进程 PID 并删除状态文件。 + +由此,任务间隔调整、参数微调、增删任务等都可以在 admin 进程不重启的前提下生效。 + +## 8. 常见问题排查 + +| 现象 | 可能原因 | 检查项 | +|------|----------|--------| +| admin 日志输出 `Scheduler disabled, all tasks removed` | `scheduler.enabled` 为 `false` | 把 YAML 中的 `enabled` 改为 `true`。 | +| `No alive workers found for task ''` | Ray 集群没有存活的 worker | 确认 `ray.nodes()` 返回的 CPU worker 处于 alive 状态;新加 worker 时可调小 `worker_cache_ttl`。 | +| 任务到点触发,但每个 worker 都进入 `failed_details` 且报连接错误 | worker 上 rocklet 未运行,或 `Port.PROXY` 被防火墙阻断 | 在 worker 上访问 rocklet 存活探针 `GET /is_alive`(例如 `curl http://:/is_alive`);若无响应,使用 `rocklet --port ` 重启或在防火墙放行该端口。 | +| 非幂等任务在某个 worker 上始终不再触发 | 状态文件中记录的 PID 仍存活 | 查看 `/_status.json`,如 `status: running` 且 PID 仍在,则 `should_run` 会返回 `False`,属预期行为。 | +| 日志报 `Failed to create task ''` | `task_class` 导入失败 | 确认对应模块在 admin 进程中可被导入(同一个 venv、`PYTHONPATH` 中可见)。 | +| `ImagePullTask` 中 `docker login` 失败 | `registry_password` 未做 base64 编码,或镜像名称解析出的仓库地址有误 | 用 `echo -n '' \| base64` 重新编码;确认镜像名中的仓库 host 正确。 | + +## 相关文档 + +- [Configuration](./configuration.md) —— 环境变量与运行时部署说明 +- [API Documentation](../References/api.md) —— admin HTTP API +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) —— SDK 编程式使用 diff --git a/docs/static/img/sandbox_create_latency_distribution.png b/docs/static/img/sandbox_create_latency_distribution.png new file mode 100644 index 0000000000000000000000000000000000000000..2f2307dd3c5bdc91a700985ddd542c16d6dd99c6 GIT binary patch literal 100042 zcmeFZ_dnNR`v&}J(2%Agil%HSD_IpPqcSqH3E5eh6-|nURgt0($^Mv`Nfe@tjD#X2 zdvBiO>VCf8AMQWkc|FVP?(XhRAMf{do!5CD=W!h8b@$wvlZx4iJo^ib?9=Ey1@ETqFMITxO-UHPYpdU56;Uu1n{4y6 zf31Bk?Hrwk#MpeX)0CTFUUsqkhkV1HA5KPYuC7@_{XVV=`)_RAwdcRSDpv80=2GPT z`>V_(m0{C=eRZ*|bKU)4|2*2uuUO!}|8cPW!KRyc{_AU$kv?zZe|@5AVX+bOe|_Sf z3Ne=c&V<--TL`H4nD771rA9deh=m(bBKq<&T>U*mtj-9BL0b{h0HmYK+40)YR?t^uvisNx%O7 z_4nV(+Yu8a{l23#U}Sp4v{aH+%#KfQ)qOUxh9h6X!&x*^4D=maRidP&S}LQfJ%7Ew zY1v*VH9tT9^Y`zSyu7^P5)#qcX(q4FCv2OY90=g>nO-(GHFQ!#LnE;JEs`(^UbPZo^>BS@GndC}d&!$0z<8>>iAMuxDE=cLWAU%!l< z7hV3=Z#o&oE@Gv#INK+vnq{LC(_?IGy!+O=?RzTYuwb-|`>Y~r5f*ZCH(qPRKQdWN zxA-C&Kd-Ye(-ZZ+w6yfqKrPL@a!1u+nJdFTuXoURRTZax7W-xE{d9|X8~AK0q;|$= zX0*#)Khdkau~92aU~Zx>tg&?S;{#`X4pdGUhDOe9R#H-`dFy21i0`xMiqdPOS5MaI zcKGvkcXPg%xA$27h31gf%aeVLXGN>9rA|ffhQb4P9{nqM@xp}*awaAw2FwqYnV6Xm zlm+hO-LZ4$eQBB8w{P$6JnH(}rRHwi>Dc-O?#H!K3olRH9?jLp(u}jGFDGmBMC$6| z4t(2(yIT&Ni&0>V?2V3&9&}AJ-4t`?aKD6OHzf?9Sj9lh6Nrs@hK() z-!nO@_U_%g!s5$w5i3@eFHtfl(=S#zjCHFt=gf@tnAX2mZLan#_20^J&`35b?!cja zDS>QzqJKPi{yb3SnW*u{z@5`0A2~+K4@CV@K5Ft|b@M?Bi=H2RZF-#p1N?f0$8WCf zQneaAbMD+myin$?Z;3jGgA^Z&?b*9G&8g>1WT#=F*Ijc92ze(!HG>=hE)6vQD@_SUIq zxVwtYRMYRw+2@kSuv%v3n(-Ri`BxkRpWX_(f4?+b$keQafr;zd>_lJ5$B(PFZQFJS zFEY@Qy=n7i-}?G`n=e6&o}+(rcJHZ3W)->g;$qd)PoF-?+Su?7wiN^*&i6`5>GSGe zzI;gG9?NwcU$2FUMt^ty3pvBKOh?YHq^B=zth7jMEAWhQ8)|dczbBb_v5M7i>Nc-t zqEV`nwBx{ck(G3GyQQTKM9)ywUS961;L^+Di-?HW$i^1X*O;<*-@cW{??`g-|69_B zGyh5F%AziRHKu1~PA|!L&oXg%4KI5xb2T|+5HWP>?AhB?DwQv2%@$5B`dneVA1u6D z$v1+6f*OWxTe4mn7WrJ68KpK1rkij}?-^Ep-0HEGjL60B5J4j)8)SCv z-Fw3NhgiFpAJ6$WFZA%3r~knn5Cw!eMz=H!5qhld2o=jFAxZ{6bYDd_L* zy;~P|HbyVoj>bGmF|xJ2KiHbrSs$-H-5VEW);aOkZHS#!)P{?}IV~gOL)6tptI*`z z$p!^HY@!Z7WlLq5DGm+}vN(fj*#G&lFEajCwY9Yl!#|IGjy%dNCx2GbZ~XO6&1ztJ%T zf!BK30^0-)?jhCsRbK6_-9J9m&LJ%=O)_^}+yULpOE+Em8n#C?F!ePgok-A3`q|sN zMo&*~dV2Z=Vj}0tNVIAl*2ie#@3%kS5}2{Bbj-}m+Io5=t*tvlc(uw<9n!8Y`9#*? zY#C&&PdJ|tJH2&sa-hZBO8>$IWxF!=PX(Q2U%s5xEGxOa;l%3eBGwHruuH?EqrM>_ zTah@KDz|OgbX_O?;!7l?;gON+ets+Z2L^t1cKWT~DPh!GOQmg#rLARPD0wPkJ&bH0 zgYV@sER@6suJ!cvY-(

G13QYl8v_6boOh+9!htE8e00Q_aOM|Gr~0e{I5&f#T~n4VTdeXEUp&AI zoIdhYj~{O##k#xu7Rt=OVvQ>3Lci=yb}KgAyi+sz0(VQ<(`7H7-q_fbZn5DOBik)} z2U#akQBf|9*Bjqx#QWC7oa`L!s`Lmu?G#lhQ6y&HW>g(<6y+?k2K}GA4Xk!)mz_q`YP7XE2P3N@7@=n*TQM~3yKIXU$ zJ+JP?vYL#)Rc#-&3jA_^BCZnkcX%MXJ=%S&dw6nE;renqBYXP|_V)I|)(z{_Hlw(>c-fgVXON;;3zrr>)3YO7yDGNgb>ywJi+$;8UY=N$WW0U* zcJ0hdAvH5^zPNE*ofmp@K7oG2h7AOpefzeRsuZU1Ow{(E+u(&8H*S2275V7xdn9Yq z1NSrYh_9Mt{iEavm6xv}XU@s#1VZKSw}jifPuwEmnB&qHTJ7aJ_~R~8TlA%GZ}2UC z_wO_A*s&(mjr-R58lOF+JFN;AD}g3NNo-ThZz zN>N#Bo?FeU(_12EihLGWmBIw%$%aSEtr0(Zw4|(zo}HcjwQ3A~*}d(isSK5F0}`

8=zv>IwWk>-w%``g7zs+?D^I-e@RD>b((zoL;kKcTOJ;#sd z2^~X4W0CQZ#LKKANVU1S7)bT{_3M9HvIUW}FQ&?K=r}yUpV4C_+vX_^?l6AAd1-_$6s^JB|d(B z>WJdv;+(zCj9!MNvBpf>t?B72d>3*yMoPH&78S{$p1pEu&Ubf0veTM$3_umAGphth zijJ+NDpkfFaqSN``3~f6GWxsR@_Pn%enEjTVq?{+Rnv2GF`iR{+UXbXtGc+FwiaFG zX!japL^0kMnQd!hbMDotO~Jfc>yI5f#&Y=LI`R|^chR4m5<3Ri^{NC%kdD$kRP48k zmJz@#sHx1hb?bhmqvjP%r5UT%?>ML(i{i|CW)3GM$vH;hzM672@}y15hTS&+>L2ob zQMk`~_rZgWzW|chEX>XKH>4Ow#kgN$J!E#rhP}XZs#^6q5~Ah>ty~{hS67p^e0Pa4 z6ul^$r{#|i1Pm3DGEy4d?Uj-9tU#P=)q+WpAc zjb({(V#nW^G$iVHH1?pxY@MQ`8sdO`046bafA#9s9c1$I+N%@av~+t>+)x{D^R>z> z&BtWl%l+Np*xA>|Tj_*4Lrxj(7aha4`lq7&{3^=NC3SElST969fBxKLsI5SvQpob_ zuE&!T-!rW%tKDv|-+73810@=zfJU=vNY*} zv)g${Nk{6+ux#2oL4*8_BA4pFdKO{cFlM08D1kQUW?N?}is2qkP6eSa{fM!9CPzf=esHlF#42x(8I{Cw;nGTuP;H4e z{TN8=b-Em7RBh`_Qs<0fd&;?&$7zxVgBR&g`)eu zy}kXXBh${E{JE|Jm+C!G6m$Qj-HA8CRpeA6 z+~j~dUmc2YMr{b6j@n4^J)yC$Cws4J$;w_E9v=^8%uY*F^?De(1yv}sV^qxUhc=%M zP3}U0$N5*Wu}_2J;)@0TwhI=`M~6q%mKaedV#c3IxV$lWt`y2AbZHhfn|HmINyWI% zRRqiV&rW~~oCWhMcae0({P5{udy(MUZG5l)MPr1tx2|Gcn(0X$A5@vN>-@=jc00us zH%NutTJ~Kq4e638Ott;X`}fD^7pB_%yAFyp(Ou&@|0ej4$Rj*Sap-wU&u%*HL3&9P>c*q4{SUgHaQNw~sj0~!pDRDt!l`63 zl;lZTVLB8%1)3PoL(L^03t?448(uTMX&IrYr7rXnfWiN000eD$4ou6p{hW~}FnJofc;-+SZ=)M_q)1T_!S>hpjq9e725_r>wLvz#Gk z4>xrovC>|wMp-Le*dr>c+0@kZDNx;;L*tKWwNcony$a!8%2iJepA0<=j=d#|G9Xz{q+f28>*g(okYitlg-#|YixWS z<*KBvZi}6ror8;uJR02R(%$~SM!&i$!*qL#=0}yr|9)dY@}wWr6gzqWttYB#Kj>|5lP7%Xhq)goefdEiLAJ5*DeZ(uMYq-W>url{^9_nrl!gQ z4gVS%3PP3r_4VZm+QLBk7wHo?d#V|hs!3M2LCYvvTU!U<6qi(1GVI>Hn|3i~%jV7a z-c>&ez1b!wCu(@@U|xbfGHI40kwZ6_`rLtpkF5ol(UhA z<+b<_`KJIu1dcs>#@658zmb{wHk+8;r`)x9PCe_l@$23||4X3&M^L-p1AFuJH6)N` zLOu3gLn=z1($%4s>Ww-at$I3m>W5OGpQ>EETy2o=CXxh}ty02zSGytG+x=r>(22*vD%v#-pF> zd~BqnG&47xP;}r0)+4nhWdNs!RGEUTp%1Hxz_JkD?MJ@5y7e@kAw8rWv0z@}<>mFt zbQ&?a8a)_V%#A27N^DF_D^H$0xo_XTpI>B_g70hnee|xSMbG|8eUAnjNMYOWOjw%| zfGcO$^AC1vC11D>h^(7oxeO=tLY}KigNwX^0zLf(_In61EB(69++qz;k5UR}*W%A1 zjHt}wPlc_PBSuf$USFPDhZUh9T5o{AVU4&wyZ9PCA2>+5+Pb6PJ=kK(u_#4YXRjX3a5_)UA0d$Ix4WTq0cT zpFfsJ%Jd+IgmM$hejEsQVjE~mJK>tYjq9B3UjxFnp%AUcKee=cZK8EBK0baLINq!T zz3(k-4C$oMLat_HEXA5R{gD3_uO7#AyCZ3Ftri*89Gs zq=u@*jGtQev%u8A@M9b2L?&v-W=1wqIayf>IxuHD8I`}CFU-mxk(U12XFFq&Mtfv> z^9fa~5MX@_{_0LZKnWmi=HJUgLPCG)W*%}<|BB?%)z|7l2i*0Asas4o@k?A%)` z&pQA@@4=%-(ISax08`V_gmUh^wQ8oBYiC$Kw8;UqC)p5Da3RNG8S27URvF{>E9(<= zw%~<=Z!vB^B-i(R>!*B$4EuHiYi$pY+{6S!B$IJ4tJY0m6t?E=NkBVa5hfS}s&B_p zSLMDYWkto~I1Y|;leRgzxvbhAN2u1>w{Nckzj6mhX*cLSJWMI#)2wGD;d-9r-_L{` z67BXRIaD94I4{V!-RR0WL7yZ{Pk110Wayqar*!JnxzA4hO=${YO7TZ*s{|`j42zP* zH+bt0u_*QSB`3-pJ$kf2KmK>H+T+1lpPLNK+nl;TZvs~HpV_}}-yQ66pY0wQ0f94V zdON8ppPaoNTjSLaJPZ!rM@8clX_5;<(7J1?E6lK>7U!{|&==VAMRl@WFc&pFD{I@T z4ea%b*`S+B1w+$ax$Dt=>ae0$ZPt2>9#w5V*J<#_TmP;dqBh}tw%K;AW`-pa;^KGk z+1%a6_kKv)IeUkEFYt70cHe`CSp9G&orOR_j^wK-Ou|G?A`|? z(TOaLn0ga=EC5p-PvnMBLBnJ>eT^jD^P*KBIj-`}bvRc&6`?JB`10k;@6czX3{IK- zm$B@mn@w_3(FE$efFcu273oT_7oiF`p%>l{W;nA}zsN^F=k43_$&tjK zb^1wb-sa^6AQ7MHeJfrk~t5x>4#P~G-x)Q537}+>_Xy?aala#*m1+8alKxajZluwEwAZc zxBVM;WugtMEX&KuIW>1q6PY-asuizc(*5}fty2;#<_>WO*|cvD`;kgK*{Z^YPJ}AN zcS-HOamU|(Gf+M=FnJv-+I)e^6X`eCssKcHKwU5a{X(zzrN;b6j-!7cKR6`4nfUXG z+Nz-z;2_LfnGkqq9T(X|ZN4f_B^!Kf)2tMr$U(m`cZmSq7te4)zvu9-^!Q1r6#IIoQ||w8JVrlD8e7k$4SaSNDvvipu$Zf!9b9w0X32oJtm^rh5?@ z0(!~#xHPLe&a^8dCP}|6sv{nDZDL~D_bg0JY6eA<=3I8@WTn9A)2H88Rw{*lxwrkW z9FBCb0ABrb#1ZvSWFZgT88FBG;F&C`Y1!FflWQ_jzRE`{zOQXZK$=h7Urx974yZdz zs^?hu{f>#f*52ByJ9S;lb0eXRadvceZf$HwU;P0s2&3zMWpwdN`u~wTBc18;{>(Lh z|L%WQW4&S8^xPzPI4f5-aR*FF7*faaWhg~nn z4>pyRDNNR$T)i7KB#+Xw%)5@%RfD4SO?(vf$8gZ~glrMzug|^oEgpC#vigI@TlXJaQsd)eV|6p5%B#DWU*|qe@tOD= z!ndout+Il~$FE4Kb*G#`^E1>UekyE~2*c zD_tG`qZ&EWW6{0VRkiZdr;XTsCKSJo>aoh3I?L{FiSn6^&xi~$kT7F(-m=HDSU|By zL|;&sbb(4oZNCdA3FjJY=Arw-)7 zXqdz-h9e$aWsSh)tUAC;bum^>RYK!$R;Qz<#};z6f$}J^u>=fWx$l#E#??lQ3=9nD z^;HV$x`dnLwY2RF&6S42+&@RC1u2rr5RS z(wvVlO5eYCPa|LJ2uR*wPBiPN3l+7A`gxk6hR2WlCgyFVnEd|qkTZ<`=+v6-HU;4I ze>g=c>c)+g%>kXDcr5Z5ww?e6Nz6IhweQ6+C@+6^yYCsNg}i|?&8*O8W!B~=G^S;5 z@28>PDKM{bk*>@YNdH?CbKB-AcoeSK)&kFLK3kw=^WNZ!m-P!j@v-K;!g<@b>^b~S zo4W{F?6INC+~pny8dx`wH?`B)gin80fCv|Mmwx(rQ$Nk(3%d?G1e@#C?pMetUOa3%0WoY*r&VW?6dmY^|8wQH3i9dM6hZMMCk&g1%6 zkCu;+!@Z3xEaIH<)#?TjdcwI%ROEWJKbxQ?KfG1oHM1FUUJR;0SvMK4`k|~W8mLj> zY>b@M6NCIK#u{R#mp%q@ob;n-RDL=HB39_?qSxct^1yfRWbd&{Ni_Dv1Z`Z-u#N9s ze{=dKc6PhkjI!rn#@8?~JdUiRyf7*CSLkt^pLUF?ofv4*5>@V%OmII#7o0a350EB3 z{d2`$rSNjVJRr*wBrYl%9OI@;Yi1x22HCd^F_5302y-~y4G4J5Sm)Aq_?3LXHu0;A z`QuAS?xfKo941vU442z??-TssmQ+uF;4Z2x2PyV%?(Qf&*1zn#W?iITH z`&7)RjAO^j8yo3|QmWo2+1S{ud(IYu0;IqzvwNk_Ocygi!79oyI+Rbj8lEk!oII-E z74Sv}!M{Ub}>xUxxk7tWDHOr$<9cc`##RUKz@-+v)A?T?2mF&wLK_`7lID zmh;KbIy=#J`S<91qnmm8jZsHVZsqdj%g=%M3l`94-LDud`{a=MIRF!lHyWokWjKLF zYd`$+(Hu$u0R9O@lr%GMG?RGPhTz`-P=-ig&4^t6P<7_OzI|MUskL5f^ub}5T7i_c z26+Hx;`Xao`~C?Vgk4ZmS7(%d^?5zOe~|R-Y7X!5O#ohw$o3HuE^FZiC?(xypO!7c zgk}_m*D?O|V7CI@XESq~2ak%1iaH#qVJuKdYTJ>+b(a!mrl+~OwZI(2RKx^u&s;P! z+r`OA#0L;EH*NAmLPDy;SO~Fe48rM7V4&u@*C04|bXP|PoRyxDOS4h*e)sGzns~*Y zRVs>o^%}4uge{&qbxIb&xYaN2p0T3rbAfIHPScyWZ$C5V5PAhqhdXt&(6JHI_xoSNf^Ao*)-=~H3+@iP*mQ7r^%=lv;0cI z`KV~fzsw^ntRbr0=>Fw_VBz#~pyO{6QQ-SJ(8A_i?(*0Cjss|lHDD35i|D9ybP{qt zBeF_8`C<2J_DUtikF?~jRlzBJiXU|~n+s$r0wHgFy5#~4IQd>gTwHrwKka`tXFb-~ z3@@VT>CnsqE8<}sxe;L#1H!Wtxbg#{FQZn{WBion*|hX@^sYM(LRQdu(2O9whra1j z!T8xzr{1BzHWyqZEhCZKAqz=VCIuCGW<$a-<2_5A&;{QaDflcsZ>WeD_2(sUS@Uzz z?4dC5uOp*fN~m>zRpOkJ)xK-$^gUK=ser`Vb6aSZC^6V3wCay1l*CUI zkfd=iM6_K6`48BpZb9@R_N2faNA@s+B;QCFC1^b+hF&RrN+;umyFd z6fM3dEEUn!a@@x_?b`C=LmD7$lFSxn=(9#lEPF>-7&zQEnRNmC5%G`FeJT94gg3tNR%TWa~aJp%+6+=K8`M^p`v)`?mv|Ub^g~Q!vn3_Lcc#tk=Nn zfbM`;*chZ+t_xfnnah<%$kt|Ie{5FeQf92#n0Yi#IZn4o$A$8oUu*k4qlPuhx=}#V zXPyNiCl55&kdr_Aeil{&Fa!@78*OIFnIOcAdrtCQtd6(~et;!nintoU0DLMG7mMr9 z>Ua?5OlL2rpitQLh~;E^=Fy{#EWPVa=N2xG0}D6D=85kNNaLn*hF= zHyqmjDAHRZuElb8yjQJHW7fA0)~i&X1rJgw8~kgCcL4nb@c`)-A5%n`y(8^C`x8Mb ztS!19`ika`Lh@o=aY}amb3+v_dU=wMDbsUQ*#CxujPL}Kp; zs|I%c7~0B-tRI`VZoNqyYEabVj~_pN#$OSrWr#cLpR2GrS}p$Pz`%XS&a#(8Cb@k8 zH52XPEt8F>pbMEBib`XxajK33a!pjbq2TzJ`1$%~EE3@?@#wRoZi7ZGS+;+^zVvHz zH~WZZH8wT%!{NJNowIRyl-KAQFe~@$&Pc?8XTJt;Rr3A2?ud03kV9{AFYuXp>v0`j zT@%P|a0%SX&K9$=wOvMe%&GDo#l)##{^^e&Ka8NDcn;71so4uwLp^j6+kNC|J7SVB zxG|%dsi_YN1=MG0qF^OJ7Wo?{UGjPO`l6Lp{8;f47Ri62uaP$pi*y3?6S14^DC+lm z_VND9+X;%rx&#^aDw$Uep70vNc+DJ|9BfVaX8swgK>EbU)|Rf|BL>>rN3lMoKiR3)W(~|Sih?v6 z{fT%fujzxHbi#Y7opQb3lk6jEXLhBS4fw0tg4RbYur4-RMEwHwCsFhYSuAHf^HD6aOEJocbrYU1&8&R$D=O8Bkwesm1rFE5rr%kd zyH}!a_?F#g>r)BKoJU*BB62C%<_DPS2ce;)ps_|qJzAKXf)B3swNJZ?lx2h0!#DNb z8c5enMF-FStvi~Ho-hWy9W7MJwZG}*bnQ6Yja4=CkVbB}4t(DpVh{S**gVAjsD6^J z2r{>ILDR}>=xlg>jWpL6 zfHT$vHWUSwsL)VOTMk)wL&N+lLTKfaMfj|DJ(pHA%=l~l_89lM{nFCx;5bG9UccI4 zen=C|aaHZ5^1L%vY!LF<(S+8s5}z5U>#En1D*r@oc$nVfpS&VfkvNLbMy#toyZT{Y z%Nv9Tjkz7V^M}#^=1kqwhczLjol7&3PjW%qUQ}UFIElldCzy5l60#a?m-Mdm^Q=yE zBOuCR5SYYdAl#qb?o$FiC9^hC(p`^SJc)rVlD93aW?r znpH>12zrjhD%mB=P_4hB1a9Ks*nsvg>DS1}h&W_$6ZRorFgXCZoF8W(rTrZ5t*gE4 zJzTsB4DQ3$FVJ);klVhD%&#jSur)){gz?VO<=7a|1~>!dKDxx0bM!JYGCERJ&}h0Y zPtj5lrW~wx{WC31I1~^DOtEJo--yfu-moBG7U###7VvXYp6K6_hW@yAQd9Q;z+3H% zs&B27t=-UEcWz);nR~Bof_`2WUlSEzfM#dyQa)HfO$*jNr*!&sF@)MjNgj^G*#oGU zH2V%&#=P$ZWR%xhDbLGidg!U0yG2B{LMTr!?%eXV;68+(uL)X$Ny&;S7vB2*>soUo zlI=nWA!036|MBC2mWibP%gA#D9upZ^&4u1`U!&z2lAL4f>wsYX6#2EX5Ld2_*+4t{ zCz}qv`28C>I5?UacI_Jbr3_A%$xz#_C5sWk7?wF~j0$OAw4fO(eemG9hR;F|n&X@S zbBIl*l-VERfOe6O)TPE`+k@|9?@DOP_VVU`k)B&eLhZ;|mBtjqlQL{idfh6_AOE&Kb%o7^U_uv7OF@;w0rGbXRE;O zOz0cUl@AcFTv=SNwlfj20q42AmUvKeJg1TyQqi`RKMkdSF`c;tYK+*7%L@U!EV*oCh!N_&l#AJFX#o}5VJ zKTDD;n9Reu_Y?(Bk9+G}z~kst9$T|!4N)S{ZNC&(fg^VP4)JS5jIDoDb@ic8i#Sdl z1mTC7mehdw7M63;p>yNCeqhYFVwsQsMEVOCwPjTd<%_?#gKU*UV^#zyji5PuIS zMTLjo*{>A(r-0r}YH;tKJ!ZvM?eiF7?%Y{@4`&t`_PnG$G~WX)%QhXp`1u{`w_{%e zf20)*$9vd9BO($fijq@$ooZ{(*ulIW<*+d8l4K0l1r{LjT}sbH95UqYZl($Me2Myq zsG}}*f`(ji%cZ<(q7kGC=puv%*5D1Q`2kC_SXkhwcJGc%$ z!@>)^e^k|Vb)R;K$XH94VKN}XB!JPnA@LaK6^EYcou!0Tzl9jw22d{!Z9zC6DhEKW!E1j=3mH=uQ>thhJv7UUg# zE^M_K!Ly$?X5(`Fu=m6?arg^zOQo+I&A_QWL{p#gf_U?T&K*B-LI`E$Luo1Z8Q9FM z=K8!s`XZaRpqymN>0A7ps3_QkOka4u_4nt#R1-hx-YLf7rdVT<+E&!qyehL;#+ z04~sqf5kTvLIRYuxZ7Yv^{PzJNgQ!{;x@TVocm>D+vxoiGZ2m z4)}my@E;Q$mqLu6v5Cvep+_0R*$caUg|8Dh;cHM-yx_EIuFyk?v7YV-5FAGFeL?0g zF1$UomRVpOfE^q_On?QN>sJ41&C70BQWX4xRz5(Dsz}uBjmLi?7$Ff>sk%FLbOyun93-I;DH-~|zu-Nnuz^7`i`t@I=mltkr@?EI&B7m65ciYEpDALi z^;wwp*B9E(n-cgeF+V^5;Kk1@I`4kqIZ4}HpQNXCX;4N&f;HQ|op=_hnqmMgaU%Ng zp_vJq;avl z5L*Kc{{ob@v~OY%)nNfqvL@7@IsXf4TwDye#6L}Gd%^0ZfnR>?d+4&$g1fuBF?K!t zuti|QFhFDn+9S)_S8J|)p3|hhGo^+yjwv=YoTa=#jiA>nQNI*v*_!J@-0ZK6>R~tU z1W5aq%e?9M&j7)~cZgYX)L>5|ToGuscZgOMy8=08BsF4fCXz7f3JUT$7z^O3JvFC# z<%;-v4rxBmWf2oO(jMbn$S-iMPEDNzX-4J*^jo03$FtOji*T480eFLEZ~67bGVsco zWjJpPd2T}rAm<6`h@?tL1t`q#fbg9Q&7L1|`ls%S^aGp1X2nGO1|YSu3u}R@SEp=|JKd%3tySzlE;b4CHgst|@QAmy+@IY+K~hkP_V-m4DwSl7Y_wMh;Un+AVh z=?GvlY_=)QVJ8OAWh%5r!CcOSH}Q>(+zDRj7T~@0eZ#kxiI+48(T(W_Hu#V4U?K(O zf8*xOdj`UuCY70=e|y=S%!7?+YxK1j6;(HAL9V+6u9jHLY7g#S32*An?T0Utu@o-M zzo6aSxPALw&=bFqGitLE^|F(TCtm!=K7GQ9B5xxjq>fd`Za7TRjA`osgihzF|%Pbmaf0$dVe;$G{ z$lO--&cIn(JCdE-U?&Mzb-&nJ*zP^0pB0J`vc^BQ;q6=ld&?H< z`=;Qn#L!=}kq7WGG2S#gP@TRcXr&1bAiOd^iVv`w6`_U|7&Alse!`JB$ZW=_S%BPV3s+}wan zcRy=?am;VsnP(IEmnvK|AE(XLg2QI2H2y473#%bI`V)?eWqb)9`rGro-o9NVZa#Dw zDs2rQV75SCrp;p^`IHizIm{KeGv0*`4LXySPtdwwodPmi+s9A|?=_#KS#ghvErXJ=mlo6VQv zHY0IyaO@HH$G4pNJ^H3_?=oP8()){!v$*XP&iss~^zrE~Rw<8+HgjoWi?BjpZ&UAa zSiv>wJlB9d_|qy)>08ws(BXbNlB7u8j2!yysJE9_{`6134Yc#3RMU6iV{dC6P5z+S zIRItQ{DUT-VDu8Jt0@WFi*-4UroxtAWku&vHvUP9B~p@-4sgDBO|`l?xVedWNP}@9 z{u*FuP0gmO3lm``QGNE6fD%GZ-O2?EPvNR% zhIId_#*6p?aoYQK*u&_D>`_tpSsmwE2J;2nnH|QbP+lDt=iN2p)#R;n)Z8R^Le>Hv z2NxcsKlE$_M95psTo}<~3PtM=2nMmKAY%}R7a5^zE=&l@%R7pmiYk(b+B`j&H^h|3 zqHf)me*`DOdfe-8V~U_Pe>n>wx6nWNL7!)lx(~XT_>xES7nS^(xP2j6{OSL`@>4&} z2pkOtSstjm!+`P=wA0tTb}lQF4+6lDgjl8SD^ic6sruR9Z&RC3D|!EZIXH@sfDq>T z$IKcJM0rIwg#;%hB_)OgWmi>HtVLogL;OI>`Lv@aZbLrMaGVEFBWWYDuBEy87f9%bnOrZ=qWuNoNYE!C`m0kqM09j? zjBzBI-Z~vXEp(zP$C3`*j6W2X0oHH1{@|YLCH!L(`AZ3R4t`Ut*QfXb0$o`uIVmym zAXIQnP05^`K@d^!I$~C=`+=I?LGB_D2?PVHls={`nDm66!9Lz!s2&B2C>exAanGq4 zt6uXzm1ZlyE)$mU@Y}i95Ex5squ3OJ*zd!&tCd!n)Q>uRx4L0pJ5}JG@WN=NpmvVK z%hOZtNMe>vsW01vr98&30l0PK;7}5`7k!W8?@ybEy_(Zy=(U}`=uu^Ch3H1WikLEFU}Iy0 zfd@3u21Zt4S-3|g==c7qMO!`I-@FC=(L2ZjjvXJ?Kd7-E$LF?P;oTN_hyaxzB5R+Gd>jyt+J`3o0z zNJ>h6#Y7I#5rG`*g4D1MSERm&hlRBGawktJl|cRu1ir=dvLDfAb9eh;Nu$x;gMlN0W}D!dTzldq&(%XAHHKNK+EhkH*x&1c6--$Q8g>1M2s zI~!GoSi1Lj@C}-toGT-H*0W3AL1*)ebUq9WC$P|<#>$$xem~OiE#+z z02unqPO20$fuPJx5i*tqTPMUV@>b;J62iPez&uny-=7E~(<7ba@zBAJ0J#CUVDxka zq%ty&hTW}+RbgAXYE|F5CO~i&0mH(3M7$t22o&96P=b-I0l)Xd07FnFA!6Yl0$TIM zTD`>kST@u41-bP<0M0$%hUdItR@)1JpK{o$FZ3GtVWhoS$6w zT)>3`WHgXP%0n0~tU$~*{{$0{LF8o*+}nxREPSbcHRAXhoJBg_)&e>~CwCp5u4@`I zf`Fcik<27x?C^;zS8_Cs>=B?4R2ZRhuYz^0n1IyZ3i$lgwFdwoU35(>P5Jl3R zp=Ef1ea14w0^s}xz~`EX>9BS7PbaiDxLl)Tjh?% zG$Lrx0W~EW7O|su(+1U!6~2p!-o5A?z}YxBJKsP{ifZl$rDPWu*C!kOAl3x*KRH*I zq_b?AeNm%2uyC~nDe!_LHT+w9>xR#-$36f*A2fPLNis&59PSe`x7(|5PZmIR7)Zky zQYHyn$nqlyk|ZbPZ*R_%H71k%aL0R$IRk5>lTQO)ZO%b3v7HxKwQ3cgLH@sb42%9< z%ntoRs=%fnQ&ZcFKJga}9fU$a(?ka12(5Uz;Pe0*+~a`C4v_ze`-ODq@W6qz3<5>V z#C48>@h+m)lY|JxE1jf6If?iWdDAQZ=N^#J5a9OphTan)4XNnj($8ZbP&x21WFB0% zq;PT8+G}ng+g*eF#hwa=|GjG@ddJ;+_vG+$FvtinGOS+x4tYF%p-T`*RaFyQmStn| zPITh?8PhQ#4iDeYo}Mr_I#QXUvjhK8t@zGO56ej)zKa-q-op+cWw}R@SXYBMA?LfEo zZ#Brb74`yJVN^F>yc#xbxYCaq8U54~*h-$DOC%s)fI|Kkz6xXr_Ceow=qp*G%P>G} zg(s$1mfVLDy~Opum!+r%V$x*S!GjxdfQylV(CIek3v4DIaT_y`@ZsD7+3J6iLs?bx;W8(b=Omp$4xacS->??HSVhd1bd(V#6t*|12J zBt|>5LFdqBGx2F}1jk7(-XLRG(AZ+aDp~LizPJJkA9p4~+!!DC<(%+U@Fj#tgv;~b zm0=_F3kvqq5Uw$;&j@ynT(^%nT0Q)|I58r*`>(PR4ca`=O7y=XopY5LUj*O2K zcU%rC!58%*0ue*t zTIh07y?dvP!}rdaNHuido zivDOmd}~^k;P_<1^;GWq{X49n$X5?UG1t$T8t6Z*AA^K=`T6H1@0hltdozUggdiCJ z@}RHkuvnI=U1T2_o#>jwHAu}@&>Jk1h0pqn&DiJlC%Wp#lOqpb=D{678p#)QC>w#? z)Ray8rPs+43Ll=u!Y=B;Ahc2}1rcC>Hj@#Xu7@o>=S+j1K)_dpu)0hZHg{323P7Ks z^$ow4HuMcAmK510mMppHW6rs68tV2o%rRB>nLKF)!9quYgvAf7wexF;Orj|APrc0q=3p{;);LLr=X6beD|(EPd` z%kQ>%gXks$*giEiHDpW=^67AuSo;fbPq--QBgl9-&*(6=vnwS)d1({rH2~~zw}rt@y2>b2Na+KV)7(p%AwWI?dmX z+%D-E8Ns??PsAOgO#-9iHjj*jJ; z?1TY*V|erC&0mo6WRZ22QI_V%Hjv)ZX%-u0tr9hnjkFM1$GLCc3-Eoo2W?s6e+Voh z&Xegl93jN>b&U3rdyQahscslGry5v*jJbVNe$qf{;Ky6_gu$Jxto~rMOF(Z>4#Gz@ z4`Exkj(EgyB0j+1L!l7+652@2g5C%HNX!MPxZ(&b08v%_%s{?l2mTT56oHtsWpCN4ndU6QUq-QaO zEdhHaudKWd>d9@=lwsBDp%}Pd5fc+5j~f?WIpd^#WsSnYRu6oF149h|%Vgnn-ltV@ zLQc*-$mI%7bYGc-*B}WJZV+dXq*ha+L6nejc*;hwDOG*ZPxO&#S5nd(ejR@tTPMkZ zd~rdj3XRqc{3OaDy+HTH`L_*MoOTT1Fojl(5dTpb2ItPBaS0!2mP5Wngl(dHC|X*y zuFb$9ch)aC6D7?rb@Z~6Q{5#U7fAHA<0wVvBYY5)YEQw0@y4n~HLjDYJQOxxSdDGE zi3W+X3(5Sz=FOYUtgzSuneQiYtz3AN=2mV%ra?Od`Km8h`T9C=D=KGDNLFz?~w;yUTh9gpiBL50#&$aQp> z<$L9b)I|S4hfEH~c8srBwLVyvZvk}RXM@82PeXk+AVV-?vw}j_1GSqbOeW~TPf#et z0jSb|!Ye3JO%N}*edAn$CO=_E043<- zWI)>jAg#e360im?>9rblEkPz^x~lf}jT`T9PZY!w43!;&rf5DEH!k5 zy}sZ43Gcn@yh$VOTHP4=NyZi7P@uPvl!cBhUf=+VPl}7w#Oys@JF-(|1gy5wB zB`McX=8PcbV+>&*K}aZ?;EF&Q?Ns1jLhKN`93-gs5TFQdL!3#^51#ZP^OKc3DO4t?yyZ#QJ~AD(GGjm)J5u{)x1gIQ%%|LN6>ci0|bTf`_|C zaIcbIWf5!%xPf3jbwZ_e`)6Azik085UbyN6AJZm-jm(Jd`7t^@x8_QW7jI}HkxyA*_q z`X_@>Il|#>!u>>hB2SPhpLC0}y%CArAbxQx2ey(}c{?1JQI4-B83HQ~#x4iey1=m0 zr}5gaia4Y^%tprn1Bj|dB9o{kI6q`y9t?DrUF#e1`yNH-WW*uqQ*C6W6 zd80ao`PJZ})b;i(eZ$|iZ(HnSH`>{@152}P*VvC1%wY8zTv$!$?21~HrtSFgFb{Y% z3O#t6|2KC6tfN`y;X{W~ZJG~}+d{y6g2j}(xppg=Gy#8oYT`$t%S=DHyp8h0W8&|{ z_`^KqRtt(z$cc*jWrWODj$;plRTtWm*vT;sw~uwXF4j1H-CTnD&PnNj8p|k%=@qb* zMAjvI2?5VWoUWIt_hOM(lp|MBl6tIjM^7m#9-4V_1{nocjQvEy2o7c{u)@Wj{E;hY z^gV!O zL_tD^yM;R=2oDFo(-=j_V<_^oyCKhwcc>?ucI;ROyh_+jY`7eX1ldH`yHqD51P^@S z`9_mgzmE$MFiNco#{_9#a5b=mtFpUh!X(IV*vvP9wDg0c9tM9^u<%$HLPASadZ9y1 zvQ|cl4&~GFMJH5JT1p44LzXNtR-62gV;4* zWsu$lYACP`Bl{{C>K|9kkXtufQw^Kr6NW$^*<;`#|JW!d#V83OESOSu2`~|$mW?;XDx6;8s3OuzciU24$ zRTAW!TKM&IHNe+j;w!GkFYtk!#)#i5G!P?$g&qnkF`__+~?DRN;AO^+Tfyqb{naK5ao0W9#3e?Xq|)cU5iwBXdNrM9)xvKVKk+Omfcu!Pt8M zRh4bs!pF7^ts>f{76YP3Q4vsrNH(H?fPzR8Fpv>ZqGZepMxua#5=D{{C2OM~lCwxs zNs^T)`HgwJ_xtPp^Nny=Nv z`z);rfhu!VXJDbXvc;>v3Sn%3aRmYGW{r^lM7FY|ZYBK8#_%ktcGkQ`KTCOnzIZd8FkqP=I-OQ`%+vlozj7&6hM=SduN?3iG*;!*DCOpq;}rvXq+sh-YXn zO5|5Zx_X-!D$ZhRibG@ zuU3F7sk9ddWeEWP8o$HHx|atxN0EV=Tlmb*oFF6PUtmat+BV+j5SC*sL3wZ^oW86vtt7lO zH`eVPw3gn3o_swzI-0Qo!QAqR#GdMCU#ruS-&g!Zqi&vkRMp{@ZY3gXcuXASc_3r(GKt*!g&z$iHg!}%Wl+n_0OzoVx;al-4h z^He&x?olUXA@t-qd~Fgn#zzkyZYXxZ)=3y9SR)(|lkav*gqozHL1G0PA@B?Gkvr1j zJD`LuWFRa$P6oImA^r7h5|Ft=2N{WtI|2|Sr(D!}l@*qZjfirZPa9FgphvXPI|HP& zvj7A)N=TeS$(`hw-4Q?@qFO(KjQ#$7@HiDhm=vJccnrK@!5g#TSCANI1|d96Bny0m0Y;MYAmcPi`#9o0oz3>8}lgXQoh9qz!_^c?Ghc ztP@<;Nc0d=l=~gF+5QlShTt-5SLfR5+`4xBz-tFN-UyxZzD$kUnj$XrIi^_2)mv_abbBoSh(*}2O zb9mu+nOgG@^0kU#*c_ds!BW~l6b_@{OO??(*t*>kBdAFq#5RIN?^4D=z>#_rQE|7X z0Z*)89K^~h=x_sNxo{3PbqD)@ZR*Sou}slx+CsCp3U!dy^pV`e^`rae8#h<0>ME-Q zcYF0U&s+4*hT>2p*kn3cyXWRjHy~*X=gpf(a=hCV6t8MSM~iQcD;VXut)2JSY{meS zgj;D`TUXV=4Uz(s$~Sl;qf|f$DMxX)`Nb;+Wg&6widCc;((74g-CneM)?PF>xnno@ zOJu#evIay}p9)}~`#Itk?*-B8Nz<57+#vD7@}eJDL++U5Ry##+CMG5lVTcr$npS;^ zxXfAOXYeB^?#n{1>*`P)@`^*=l*jvJ#`n=9E=l%ZXkNc3R|~%rdkJ~`ioFKMI>Tf= z-%%22`VFuC9UV59@OPj|9Nn&kKGva_eEbPB_C^jZJ8Dn*qj=D7h=w+biD{B-XSakp z;@ zMOz-{nYKGMrY!k4$uOCyZ3Aoz7W@r&?iu$M1z(@OGqY|}gUw04VoL{T0TPT0FlZU> z_HQtczQN;?gqpHHRRi1#<5vIZhBbAb9y0`gVI#D&4yvm^hM;o*8j>CZR0^!I>T+1& zlz$BIhShn~x#53j^+R`ZXvPczr}mG+rb(;_0j4`3NI4>s8Qi%(2}%1u8RcLSzW?d3 z`y8z7Q@iXPQ3O!-r?bCXXDZnZcyx_Ij3yLXDB$(F^=zI#;~mT*pS&mD0n&f2a8XFN z)cgX5aikj*wg|UE3$UNIxg{gKP56mUnT=>6N>3v`=KlG_$uiD{>Vk)Y;i6*5O0h4x z8=lsS294;!wtCbx!sTAX@jHV-66W-t`WS1))Bnkc?=%VjYUV4#&HizbEMhPSF2LOj z=)0Ze7|tBd`o3cSZP+i^hOcpPtJ=r}jnq?HO`a*ADNf6ZoEMqU=&Z@jh(RvNeGdR* z>9CKeXM*{9nE08!vePs+UIv@C{wV_3bbAZueI4%(lGQzRH*-(mCX>+j$Q`#O|iG z`8Re|>E>x_XsmWT5+3BM>`F@5M5#Fle)yjZ$Vo>I%9jBJS-=QEIpiCBSi(=}nN_Uf zLU{l4N*7#(er?RBPwLxFe+09a=W*fkygxJE<>$)}{;`3Y350c;hq`>sp%nWm)jasT zBjTs;b%~K)`D9{^?>06IM{%Y{pV;0xeFg38oH4nk=hih@Q8^`kM5^{{@#-FBlzV3F zMTg1$?%dme8;;ESC1+g=phQ8_CVc$ilKp^NmE=5==O`9iQ#SXWzmI+Go~PCdA7o#; zr}uIE^Uveh&U;wD$zQLA%CBSN%x7n>h}8>$jhW_4k$-!Ccj{8z|r)H zZCmkeMzg9c^+DIUqhIo{(qUgb;jc}!L#t|!KuG5&qS3&?^7D-cI`;@g=$C_q1?|M- z^SqG}qh|u0Kr60z4e!=nbU^TQ`!#l{oo4&q&uGs;E7_{1S6Iun?&X$I_cG`rEvVU zeh*%int;@C9jF{*12Xw%g3`mgk=T2#Tjk$jgO-^0H60M^3qQFiD>GBC1st{QjfZ`= zsmMgn35iLFt4Gu%m)`?g-110`44c>!d1$xjEDHGUlnd<)HnImHhThJ|!_PuiNsNC> z9t`s>hEamm?Cp*+&=~$P*WEQ)M)Rja0Kf(Fjhjy z_=32T$T9-j)<4{;T2_q?ovm%Yf$>du8v%xHb?@(ius?lP@%W9K?Fn7=qgCu4(4N}9CUV*%8Lc2&MCzv@sVlGV08e>-5(}|N zYOOF+{4GgqRoEY2kWij{Pl)pgBqRM#l=RHaAf+R7O_i(KLz2wnE>5PPKRfIl=9`nj zDl7TvCVZbrKiXLQ2o4x!jam*va@P=oeDu-e_On~3x>}yYNQ?0_lg2Q&0lpV^@~Ia*WrwGZU&W#5nbQK*gYCR7u-_#DHr+MF-^a70x_owG28#Q6 zg?Ro?n-LS!*aXemIB2tf7@&r;U+n?r@6e)l2)}hiQE??f%?Q^twGAPrqnMI?V+o01 zD{~xZue6e8b3;iL7z>*A6`d?rD6n&9XwH<{*bm(Oc&>K76Up`EVXF~ zM7763QG9<1rC80p^pJn)`fr~(ao|L*q;nf^;}>cle=x3Iz4`-5HX6L-eyZVHuQM`cGr&}je(WncB7C4oqp88m%j-H!X$ay` zwB7%M2DIdwTlSpk=+{Zj@^5i8 z4K;_yq|ko_7+>wEXJgC1>bNvRptQREAwC^>9ZBt?@ZfzSckX;T;QMRqWFVBBcM8;L>dF1W4Hp}0#`uAv;i%+vls-~f$<7^(AQA+ zT)!OY(JZ7Oe=?$qM0XJvbLum01WT5Stb)OS5y}H#v?JgGq8q9YW6%`?&D?}wL%s$c zk{T(0;_aYF0aXc@hECpV&pLP_BwUVLAXi6@xDfa&dfbbS>oABGZtO%wA(We;Hn0R* z7#Ke{Z}B^X9mx@^B`&|Tn&j2ZUJE!Kt*2|LI}shaRDZ$k+|0_#61sQe#$@#;^#1NN zELn>wKh)2OL=_ddeb`SF^^r6K!xgsqE3^g+Mg+3?yJNtQrY*(zg55*yaNtwkKzNOm z7D&zXUPPiKII2D2U6a9u_D_!>I}%k&_=)zax2V9(V&H}nChp=8&0>)F1>iC+Uf$wX zfz}jQi%r-*uG5Cx)S~t|nTPQOz&gzXI-B0pHAEq&oxOnYUqZ_0y=*pAh6FT``7Ve$0~nrjw-V zjOH808}u+8{TQ${2XifsuSDO95#+=Oa0Qe@ZI(b%a?yzw=DP(^K5i@8>Z=;WA&EkA z4Xtxu5L(L^8Ew_cagMdrNxvY&cnZdjeW3=1VHgMbS9xz{fb#mJNn+7G5#eBznoU*J ztM$uAH?CsfB=f7rz;{1?Yyg44ieeXZ2QW?pz>$&eQ3)9fQE5)usfrS^)qRJlabZwB z{$LPO88xV(=_+rJ82-+pwD`_*aRcosWo5@sJU=!IMZ;Y-{Z01(9$;+1wus__l*Q2h znMNaBVv{vUCT!bc3YU5f{77Zj1q`ZWnp5&`XuP>>{O@tXF zow_X@!P+RvR5bhZN3s1J9nHspgSiaSk0ZA7e=2QAIi-J~c%JA}iTws zHl_Xd%9UGGUoHyl2XB4rRBX3nFd2ugmU1sA(4n1|`0;L(%OSzM9Ordt1lg~r$IPA@ z(ZZfA=wlN~4H%!5&et^@YTEc~e4M-?W4d4S9|(Us80kZa{DKkMh^BwFup=~3w%>Mk zQE{7y%?weMAqh9KAbQ`RSiFuda_{Lg>zalTKO>BW^RM|j3{{6me>ANTHu;mpB>dqD z<%j>YK7DtK^^y``>S4Z4B+8eMEHb&mSUEVrD(khXpit94NFny4fd%8{hUdQ3_XW8A zXyYM}cIII+#Vmk#u7QDp1$IRNBJ4m57_LZ>3M)`Nm2)x78qp75A1``u8w&i3HVo#@ zxArS4Kywn@eEwlhr|a;0XYhatu6H1FLk`9}ZWs3Q0k|3Bq5v>Dhh0aKn4&vmDWswH0sB7=8LvegLi52iAyI{)wY=ip z>we`}<|5inD(Pjs#CazV#S<6-%%2LdFy2%&UcP4p?VK&Zk~BL`s5~CbW{Q_W`amq$ z&R{fKC=#J}rR{@=g0e)-VW~u1!kq*M<%yf5r%smx%Hhe%_gY z2R35h7#f>Fn%TR5{~r)+;R>%^bttC4@Unu&rah2-56fXwRlXbh?w<)OU+L${gN)7Gm$c}iB z_Awv;i2_aHBi;%u1L=xjo7{q--vIDIle5eBt(@+X)rAe+n=)1B(xv)OBUcTJP@YA)+RPfjB@zbn>`{bh;O{yzNFC9*qYTiGFhsleFf0R?~6wg!8 z|7fS#JhqG9*r%{;Ovf@auC1W&(wQB@?tVRr@dl9K!yzRKS}r<}X6ZSzNfK;s)zy}i zm{>V+dhM^A;ZSoeHK_h|7rr3A&o6li1)MC8k^ls_H0cfR7 zxqtnqC*1(Xn!xo)Tx`C*=Bvi7Qa4aP5K$elvJ!GFY8eNRme7oH;ISNMCpgcL7*121 zQd$0aAh1UXD>f=$iyOQ?E}U6>fr1-2d>NjYDZ&ljZ$@Y(OTZGj5m}CclRKO{ss60x&vXB=`@!L=KzAo zoX1ixXj3_jO0FO_P!Ah?2(_dsk0Ww6jCSnWTLOvI*f$K7Or`D^r0Jisoy4gn5~v4h z3cgH9xry;(Qsdt@J&|(afNS)Tx>(p}KLeQMw5jRBy`!RGB=CZ!%k!TR>$G42FhoEc z5IITu!x@u`*+G=pXy6h64Zq1JAD}*OU*cIg;bGfW=Pw$k)~?qhsD7X`!BYBZf78~; zZjp@c6i~C2UH!+;?O|y0=Aa;7a6T3t{;SZ-AaDg*;@J~@;+-zV1`5rYcAK=1(fyr= z&)w>!chx7*21waMf2r?cwVlj(yobwc<)#KpJcr zO)jdK0nG9aar&ZFIuF2fkaOWgS^TiwYI(~sP|dcOZrfX{G#LFjH7(vRG&8LiqD!xB zdN~-NG|Zt{c8Ft%l?p)7*hw#S#e4<-ZYopC#IsvlX19mf9eC1M!;VGbP8N7(XBjhH z50V9lm8JW(DMViISJ#$OfjH?Z>d(Oqk-edrHkx@Ltvmw_rvH@M-iRlWi9wwAoVNPZ z7bp7v{q(e4n`G+q9IF!|?52xxo_w?OCa@8b^P+qOw{0uGk)Es?v{mj}^M>|Il124QdjV>on<|Y zYIu&NgVXbz1#<>$oQ!r)#imV0C^xS)$7C{iK*)yrtXE`PyQeIM{mi0*4<8<`ufWPC z8I^CWx5Hf65FkF3eHn zH4b!4!1~C$@C5cZ(}*`lJDUyoQ;*VRKa9@_-BznVJYJrBfds9y!UBX%SCZ-<>Tsai z@*LL9Y~R8PTTn?7EL?T4D_&_9HTF#TzFHjbfQ04dtHtKdF}b=WHd1!SZ@FM2A;>cu zuma;cvd7(K*!?@=r)_WaVXKwjxW=uj?ixH=6DOY_Ak_2_RtZGX?9lrI_G_y4 zZFkKy{l{8OoI!nZfp$cV#xNwsb1`__{qwINrqkz9Z#Pq!Cury3AYZ#orBkav?-zfZ zp!IqNO;J7Xd&uejWelN#iT`s$;+puocNby$(vex#NE*TO$vK@+K5p)k0c zu^sdCD0&_mMZRVNdSAj|&q45gNCs>0ghyp}b;dx90MGE$UY+JOQP^F5bH+#e+0p-t zK=uqGAjUNRk7Vc|q2bKsfYRW?wc7ko2-uUlU@4E9f7mTm zAlx9?z$VFvBtK4$GE_iNL*yt1LjJj$EmgYOpe)Xz@vVOOC@6&IAQ)t>k+%ANOaiK* zaQNyOHoQnvc*EBKZ!Bk@^m-ofe#QRN5OtIpEgwTxALXa60?_Pl#t-aIgYM{Mmc6E3 z1}oXJHWVA%SqALtZh`mOV?Xq~$+4=Ky-SOem~SXPFpTjNBG zh^9tNQKKoB>X6ieUIGLjiQ(m+_+c(&rnMLRO?r|VrkLjjsRDLO zcaQ{AEaPEJ6*~h;)k)(I=P^K=VT^Nrhv>r+^aG3K2b>#&>^5}(BO+PxeM0PjQv>)y zt9!2PNqyb{#VeI-PQGh+(8Y)#i_xlJXzF|gBoq5UJG8m2f#d$;sUt5+U)%uY=P`a^ zF9_QB?X9LF(b``dA8B~^E`42+ef-+|C%my(0N zv!-+Fbk=nL01nrbseY?e*ZLO%GuiITsKd?8xjqBu&Y)R9UOF<-2sacLJU?c8Af(SZ zjIaom$opqkS2b?ZetmJDd9+Vg+?fRVb$hCaIP~dL>(=-NpU024&RGJMtW0#Y7A6mD zu>PV|Q*&4tdwcQT_McTLW?G(w&vXpiO`b7DZd|{9r29Npwa(=IiTPZ;hBzmh+LQlu z=cJUNGEB1upc?nGwbf1+lCT|e`x^hK#|;DL>O9|e{Po2%os+M?vjWTFEDkxt6&v`~8uM`_n!hd( zUL4 zAAhf85K#@KZ_D%Vq0>BF^tj8Jkli=~V7txd`T@tVJ)4n*NPzS_8Ggn87~ESzgW1B@ zWJ0tKC8W_R6r+!Kv^+|b=d1x*_1nLUgZg-F*HL6~y%G+-~t=?4zUc`<6q#+i%!njB@T-~OG4Pu_DZma z6h5{9CbbWkR7Co>T|i@EM|DsB=dzSiUcAJU4Km<0JsRmmtE9$~Lx5#w>1F>yS9Zwx zagtaA*lOQl7S~3RlWUQImSNH3XO9h#vQ@0oKX7xAT{Q1JV-vLfAXxr}fGkBeh_}a?<<*I$6NR@G z7>_=b18y8{0qp^;1t%&tcMpGDTe{sc_x93F;EFy#1OlOz+S5CKfCFP$RsE0;vb}ez zv`O_QjJ?C_dxnULvHL>vi@)E0J+6B4HMyTq2kB9dnRG6vBK!-c3lsI46Ax>NSN@rK36p;CQPL>9O&tTran z^TGVAn>D|{mX^1J%|6vXUcrq4K*TxvSO~xrN|g!_gJ;oRs1gq+w zSYFjD&gX()=k_)g=v~1*E}&Qz$T|-ZN)R(|TFG~6t3T$aATA}AMa@wAr&{O9pNEYn ziE2j0brX{p0GC$)`ZwT>-&NH%cLN?v|A#FG260`d@j%Oo9LQ>eKsJ+2R1}ow&Vtwz zf*(v6einvI9FR;XT4hlH{y(TxbIKuKjgy?Z1@?}5{C(er?7`^<7ZZR>$! zGgy(Z)KhuKIN8w{M2wj6LwG|F`gcbm!42Xm3xOCLXTx{6O}y|IYXm5hd=o4h4YXrm z%gWrwu5v^Aa0lTy9S)iu9^0Z1_(!ucB0QYNYRuY8HW-2a0Yp`@Ao66M>#>1c5*|l> z(DBEYE9W!g>_Vm>~)D88S zMH>Yu055L<+GF-WM0iIq7wg!sBt?fJOHtzP6>1Q;(7<1ZO zyto8-uVxGKkAvWMByob^p92puI;hzn* zvfJ1L(DKrCL3vJuKC+s`VN)n$uVlP~qfJZ(+j^+)>^hyGc~pW3NX*PTy-sIT%+9y# zM_0O`O_vhUOxkxruq8B7CBC0V5uwN#dkMjqa^n#vMU|)%8hNyaSNDkwoh0Y!jma zZZ3S1NA(4hFzm-lAjMMC3yfDx%saG`zK2Pnoc$eq{6RpSXlMcgK^x4oFbx7#v-VBrbE-K%>8NQNQqg-F^8|?Z z7j+k5!H0GalML|MYo_jxYld_~g7nK3bb~Digyr403Y-9Ro?jXV{v*(l1HHpVgi5+X3Xsa!g@q=1n0dVsQ#}rK@%0PU+*D$oUsmchn531Tn-VxRmecU;e?PJ zbOFu+Y0FT}I zdhzWhH6TuZ`?33gd9Lq$NIM7U>m3$}7`5G{u=F*e+Q1Fp0f)1k+UK`qKRHT!^@yPW zFnO#5Pk-N6N33t;B7v--KTb3&g)85|p(@-i93ioF*~Y;w09HPKnZz#C6X}7yhl0rs z%vdNfGbtJ1I57D3#xCjN1#b~Fg|9IVZw3th9tJUd#+z8KFBjxTRGAH)*Dz57!%v9n%{@6m%0yYPdNacRDsI^dsUX#Sl#p||< zPYdau_%7yC3UT4U;6#>|1Zpz80Q;=s5&^^$RMa^G&*zsEvuw*z92ulR3jT_BM6YV!ty2pSGCsTW z3(31y6k3MsVoN7#>yEQnYP(Pl0Xw2#?Ugk+fYlk5_`30;l|B?R92#l z8iF?Ae=!$ukErM!)JfezP2Mi`=NcGZt8>q+S|Xtmj}12!135(3pGC|GJBly(EJ%#EH5C&~@fh;33QL`LPO zK-Y8!A7egRIl#i>MltjtK=DIp4SENphs`-WQsWusC1f@wLubF)g+^Opy&TlisO~OmA1Ri@J{U)ZI9K+xUT?A^KPe#nGO zBt7r@#8ri6(IbPIG!`#(q=m@tc=fbU#Kt~jYRcobik}Zz6?J&(cOvE5-ur-b?*c~I zr6)Idn;`|gsdn{d>e8N#2itR}ki@$aj}b*NdPPuNK8Sc!Ussa{3Z@A__-|S1h6Y7e zWJ_LX*n@Pt^9Y*#xof^Azgk=>mYG&)I=}lzdTl|*iCd)?Udj_=x6KI|v4a0@M6e5* zqHzSyPr^+kjj%cK))|3Kw2D*^2*)IFIm>;!k?}K+z6@{q!EbD=i zXBtGImCy^MM2-$NjKpnVPBgUg`}t>RC|P+2TY?|GrI^Y9fl->D{tnFkp{9Jl!E7O@ zJ%FmD2=QUr%MK0{8&i~y&BKj;naxPYb^!HS5m{USIBJp2(v68XUO|vc5v|V0mP3#M zS@~Zee`6vyKLAs1OhAg>qiZ=gWBwz%2fKj<9!LI0?e6sCv<;~qM5YY2&S+av%6=Os zmh);>J)dw08TFQ%&$SmZ*^EZc%ydQ-XEzF5n4tYQfAhwVPRK~|ktzLbv~50vh6+3F zwF03)nV0kHtR)2j)B$&*Kkm1})zU|GFBiw3$6034Mo^pO05L#DVHJ9V(Y3K3SQCbg zqrWpvqiFPB?2~e!g?q%sHwg%k;;fx^nC<%7)l2Pdl}bRUrQlTcRM>i_6* zwZ+0!u49yeel5ieow*BVHk(0Ak=E8$2|hMIG%t3O=rXh>o}w`bx6D!kN<_$C2gGz1 z12frnij~bF5r{t+IR9H_Yu^piF+|#CPyq0ao)F~-{ffBdGpGyW5+M@rS&CxQ zxg;rpmo)2x<36;@jY{w>Ew5KWB^KBpYc?Q?M*g+$xdh>C4g)WsC=k?4aus&k%-?r1 z7GLufHO1K=>H??VcE1@T6(MK`;j~4LV{N|Z$I8><;{uFS5220w*frmPAVPXu7-6Je z?V*hyRalCi-=()oO0J_)9@5`g3{q{unOyJd_PrNL?RvB}le#AiSDW>v4~tI5>)eB? z2O))vtT2gSW5{TpOV(nnkP~$EPNOur$8LabkVG7iu$@CxO&lCh$|1yk1^3iyed!Lg z&(K1`#N?Pes3Ol|ppwGcC=UO*qdZ82Qj5k6(oAF?B8XrHUXs=t{qb7S{o@U2$E3TN zWeoAWm!Q@`pM8+m+YOU8dcjOzz#14rgcP`eLru$Xe_!|S?}ZOtuyV(38ea%>dM4_Y z9_xw+4~`YPq%GDBIrtZAwEXzYr;Nd!{Dt-|U;0-YxSp|O>|3fQyVSd(tVml3(9n*= z3OQL`$-o9>9^V+}l*h&)>lBpS*426%@I78&68F%@)Ho!>^~)dg=P#OFvb z|2%W^>%#b2>tgGm{;7t-j=7m}LH%~6Q)@3BO!;)y9i>2FeSQ5-yd#@>v((9=O7I)TGB*>wa$%QEX*-3t7U zYt2is=6Pau;d4~+(-9F7$%XjatqX*Og`-})FzFSj5f9Ryg}*Xliym}(>Zz?+meB0x z0J^eP+TZx~!INR-pSt5R z_{94&Edir@A(?4=dTsGXP)Ia?KVjK;X5M^UuEzDANYkGDwU$RcN4vVZEOOaIn6e`c zmZft?6(xtO?yeOTSg~S79lDYe!GA(q@NVQR1!k#@KJXK%R-JT9FHdgUwtT^j8osmI z$vsobO%99kcW1i?1*@_~n*y#a7TX$dA>9~Sf>b4B7R+XscsMmu&yVP(6xgdEbf1|8 z^aY40_=n3p^QX*q9v(Hj-r=#*b-2E3O`=${yxqN}_&O;;KQeu-(@^>FK3`^9tB1S? zvm?mif>YUMS(`@d+9#L4-9cxLVq99gzp&O0y1OI+))N~Y!uqETmWX>YbtgX8cGaV; zZm{eC`)1kud%oS4Y4J&JbQ*0meJV%4qwe69Y(p>QrBnuPRYg_Rd!5ZEE-mvcYo`h9 zs|4`chLH7RPKvS61&lD1^;o5I`OLlQ@%{xj_Ghk3+XEh}&>mr0;ir%P%D{j2@@S4M z2xSlM@?c{0fUL)sYkaa5$O=!KJZS*v*-yk=N6fZw@=^B=hj{}z0edc;MLOpT3bAl{ zEf5DUU*E`xh%K4a{I#Z9!O7u1Zf#t9$$gaBDtc&8wLe}5S#7;m}^ z(|Dj+yJ9FO%cT|)|4ezMeF%G)b$8G-aXSPEjY!Diw zD3tCV`>2*(@udl2rzaCxw|d_%L@=1ZC(5(J#(43K{agT0<*lu)J-Khh6cupsbXBgk zj4^Yn@zM0T;RUQC7`eB2C!NEsnw#fzu%$lM&k-I4HkM+s8GmV-GyHRPmxnxGfo;x2 zqvJ?oa51)D{g);Mkw!0&7R103a##?n1jw;*J! zj57pL<@nzg#rZ+RKUigYL-Z&rTbE-x<_x}9a{4+eEHOSNB;^L*$ObK znFWZ~eZ;e0yx1;V^C5~)Rsgz5Oj*RA3FR_tVjl?`uSXtOEVIVU3wb?4G)Y7>*(`H& zEJ;p_mX4v$In}K@^D(hhDuB>iB&zrNz@ZY66#z+QOTL<_sy>#k^K)`H$HI&q4#^dB z*PMSTQxiJIJCR}9mDNPoc)y=zwlv?OMitB)o9vLL2?s*y(%`G2t%47Y+#j$89E~pR zeuf^*`U@Fmg)7Hhu!T)bO+A^wv2{0@R!sU#HCETuGSt#vJXtr)`tf)pMSh5wf62--eK$gueNX77L~J-b*|jJ z%_DJ@(`*hsukr=#23m&@ib$1Jc%d;(Qr6<_oK>yWYPzxJ$X*=d(z{;Zl(fHYxo`T@ zUc=UL59G>ogvSTGL?KXR&A*zEX*k`h0!M%^-i(JKXZ3?%eLjaJ9;f}&xd&!`SMeRj(+_*< z=tqopC(CzfW?c8KP_~vDu2Mk>8Z}2*=K;^)qz|k+=@#c^z4qdWm?VYlv3J6L*%vRE z`X-!`thX(w&fu)wXo<97c(_)i>GC<}U%8wqAH?Vrx&>L*?Wu->?j8|ywNhg^l-D@K z4HkvGvL0mH<1+{3f-n}!alA74LK~#7@YuXM=l;Zh%N9W?BO8MZ1?g-|w@6QKlFjwh zp&v0?&srHq@t2!xE*$Z29FEp4R?8W$W){p$_s_tMZIK-)5-dfDzIBt5ecI$ep$ZUb zyuGY_6KG?!s5*W)9BDK8Lz~aWk#sx%Ia~s;@ji9S@fmM4kFfUVY2mDy>4`KrtO|6@ zuC;4@w0h_1z2wK~)-Vx3*Fz6Bwhq@AtWr_ZN2Z`qV%c&}WQ(km`W9Iohu-9#%)yeh zvGxEJ4Gq5o?2Wjq1S?&BVBhJs z+;{;oU^Sjhkwz<(P*IAgAdN4ebs=kD!i0s{HV#=zOH9edRE9#C^CXmFn1Kt$P=dBQ2Aj& zI^*57{JP^6z7}0kk?lTO7YYuo()PkElr2iubdjgVzqefUcT`nNa;AslnFXftSw1+q zu?5cX54815q$YZhX-dvpdvPxc9DxuH--CBxdm*jO2R>57p7%J_tC(nsH}D+fi>LVE zwDuNyD>0rhY#4J@_RZ23VaDdon|Tuys(;Y6;+^G!~@)RTFE*6g|_^GC2r*DIaE0eV@Q))O7 zo`E)>x0z=RpIi@9KAG<_dDwX{pgD~ETLL$3!u9Ni--|DXfrluSr+#R_W7HXi-t%B0 zoxb26m|cOg|4P4An##Q%>8G+LF2!o+EvV*;lGiW}sP=EcXb7c&^2qcZgWmCun|)7N z$X>uGjma*x968L9xDQ6?Wa-SLl1D2Vs+Hh}JkLQ_f^~X3o>(f<9g6q8-b5$ad1j*7 z5S@x+#pceH--i#=8|r37ho=KiC@Uw(=*+5mb?$IH$5nQXMn-{u#ainot4f8m&WM1M zFZvb5d%gBQD|>2&b8KFfB7x&zWD%=7`E8Y~M=vz$*2H!9JJlR4fMZ3lv` zClgjTz5>dTC*zP0Fisc$$Bwz1y?i@Y6x&t+$RP-%dWEL%g}*%DnZT zK5r&x9drEm$ufO~sP2?i`Jq^CGKp_ALz9yKC&zL1{_WZE4aVN`CBGj6Uq3Zs#d7r5 zjAKH-O$U9wmE-qC>FfWeyou|#X>>z?#WCsS<3rZWp)bwQc{HOr{Mc5Z?=yIVA{>KX zgSBny^wUfg0>y+Oa`9nHzW1#TDId^56O3X^$`H%!)NI#}r5J$mvhTUJhu4a? z@ijuWft1X6gdr4N#;sViOiAJqO6(HsQol~qzLA*{sXw0xFTGwNMWijcR;Ms7qt72H zLp!o2JxY97HcttzSpwqtWTGpfv(R3@1RwK8^XTsbwOc>DV6;!p;AFDR=EBr7`K|vR z2(T@I1B-jx0UOXxxfX9~g9nA3!4nM~91YFq9!PTB8{3+{-r+oAUO&)8L?J(M)C19F zJ4}92@bpmRBqFrZ2hKyG)eRIi?7j5+E%;X#z0ez9p8ueW#GN1pg1K$ezu`o3*=ECmg}czayb)mnGy+8wn0 zjyc+VI&+IzfR?5-%|ZN#q^!q*#iC*5(Qdmx8-@0ybp-iII@nOiT&+0|zwlK3_fAN1 z1!lfkHq=$lgLdiITXyG%<~a4D5RVz`ke(5x+`}RWEy+@{WMj7U^P9SSy|Kza6qGMy zv_XkzOk^Ml0vcufroiu`&^Ru(d-y!P$d&4+?R0c!NyNC(POSl$UXqETrI#34&MS zsv5Gy)1x`11_5R#n^RyB^=_@RF)A~MK`aJuEXZl3j{AY+hUZL&-|qtqIQNV%Rj_5r zWwm~Sdtu7w3gcpo?%AxE&$ZeL8E@>!TA;Q^Ju(sOvV8_8BLF=~yTN5KvuH>d0plRn zhK+vd{nP)P^$qk=J;G4OV*8wmke^u~SSnH3`dYu2;{Syw=}o^J`MvYh7qa}0$D-8& zVf!ZD&cmvt!FePi{hd47grG(OZXP*^DR~58m96;Iy*$fUrbGPzU)E4XCSo$S&+xkx z#7~ab=74G@t0>kPHh0t(jYcthD#~^wJ zV`F|Bsv6!kHm5EH*HrlrYuk5*v&#

jVOdL1TEW($erE0$|suuXh-Ci1jbfXRadvYTc#Alp3W>- zsdfI5Wr-+va&q!&Q42lyTGm;>Mc(AA44v^s1lDjdmyNO%$RoT`KGz4M$1GTScy0Oj zri2t^MHypmlqu1lry9$rOz%$P_#2`uSqFoz3@n8b0Gy>_IQzu+E43(mFM`KXZgxR> zj(Uuo2g^hVH6wiVXA^1I30g-9h{KgAUI?fB4~!?JS$<2KwiQzrc#R_t@w#9HMVAObw|=|LG(#Ow z3OR}wX<}j`0IUkgM_;3OK}Dr0z@O^x&iD?vsD4a8aoY<(f6_Y$21SCP_ z!OeAQ+imkr{(SmHMNe%=k6C^fBG{0 z{%e*C11&Y?M?~>k{dX4_Np8M@aP0G;nEMo2i6YqZU=j#dBGw%T z=b;uz6%9c2s)NC*1d;^;Y;u9#^p29MfL^FoA7)uQ`7_dXS+|Yxu>tY zle+*+G>%5yni1>?L^og0<i;yn(4h@-I|l?acj zoak&I%d$S%(?%N4+nRaBf9lsY{Qe!pbL-rsdJ^D$n6h|DWQ}D;4FTu1UR$OO zKgf&cL#HVTqF@7F0K1KkS9LdiXFXa0|CM;F>G)xdd66R=zBWx6xChXo#l~Mh;(eZ8mlU>T8?SjUrd|7W zOOT7)!Il6u$X!YjKMzv;L7^C>_uUiK)zyTBt(LU6B4Ar`2A zdXMhNkr$D>weA(ig(xMIwiFAf!b^GxiM=KCiO?%$Ev*0;BXNo)vCDNGFVRVXc8@f> z#E-qZ5z#UW@aCSBA}zFodlNvrR8M+$99uiaEV#fzG_=a@;>^zGI$n z7HW2aQl#dY>CrkImqfgz@bq*sip5ZfwfOfh?y7m(9(pLirF{=^w}68%#A+3RU+fF4 z22N)eK;IJh+LX`x$e5{{t12n1JbIybyJj{0P`p#G>1tt<6Me{K$6prGK46Y@D!FKK zTFH7KL>7=9FMlwAgiMg^|6ru_RFrBLB_~_wb{EC>gKV-Lx`o#Ch!T@Oa_s;eP)K>o zw$VBbvFDm$>E5JDvL!4J51Z^WkJE9dcyE$QAB+p-LB=--17LJ}=*OR+S)QIufUVAi z@Gw`4Ss8^+!zm+s#AYbP6W1JlL_dAVC~!~RW)1uEKmOu=@^G2ro-^V3Yb`!-u`tb- z+wl{poiKeZlPOCoN;{n9d=`LVHmk}3G9&8&w~{&Vkq6@*sHkK!nUEZKYmsB^66^!3S=hMgxKGzO4ufcofj+?-_P{Q2`faJlH zMW8qsrR20-)&eo4Fb9r#ZH+F!_zD9Se15lZ^5H)9 zzwx`g-=(wYDBys{wV=^CIMF;K6=0U4;Z#2gAy#nCC1|(qhL>eQh^SQyAfTBXNzG9} z_D;@0^Yf2UPri(vgt@_>1wH}Qbq^)$*@!0E9hxU}C*u@bf z6kvpM7gTs^os8{JOX~t9PQT%W}) z>0)Oo){`ZemLqT@ayKkEizbSofN zi)#)pNQnpm>4C*#`>cEU>kUUFwTegY?rk>O6Jrh-$MEeM;|jf|fX;~W`2At@m?*gqY-d?o+KyPSa! zrCc1P%*52lVl+UL8V=^cF#^E zCloFZ0-lOQ{XPw;C585L&rvo$`72+VG?){Bn}#E)w8JolL4wOsfx!PQCjWBsGu1Pf z$9Hk}m45bp zh8Q-m0#6SaX*aEE-M$w&7sQPwiuUxc(HMFVnGj~8^HHQ;Y>i`EeW%r%N3tF_TwI0@ zIHk%PXN<_A0PK-eZ_q(EFRMH?6pGs4nc?5NvDi>n0p2fG`@I`)l&oB#nuGuTXFpzz z+v1D^b!^M$NCd({&UZ@v96lU*z{q$AGb+8`WlfVqXD8i3l;416deB81lV<-DF7XPZ zL2TnGo!Q5ytb40OG$YhGwwm>lVRW@VBKYT#n%ifiC=oN-CEuRY{Kc{(_pn~*BNsXR zw<@(bZRFpoqXAx@BilcC=4m1L&o(im|#-fE`l zgO`_2#J9@ef~)y`3XMg!9_i^Pq=%2m@rc!o@>cO(*r{~p=e%I)h6g?^|5Ly}ou!g; zshK@aii|7Xo#m}4?ZOUM4AN$dfAT?X_xu(=_{OR{6Cvhem z#)D;;N4(UN^$5fNm^3>jPV3a=1w>3KI#j#pNPNAqLqlD=+4ie+P)v^6tCo3FMxwt< zKXi#0TQN?@6YZxa3YWiVkGPOmZ`bb|8J^p8OYiL_+erL>m{`8XiZiuS$`#oMK8`zX zOjgk?^%(0n%uCCU7(A07*2QldK2m%~e&RwI?>gy;@bNK2-Kol=vBpwQX5eBeG0Bpg z7t>AOWS1!mc6u{qi(A#Y>P=xkRN0#YncMnv4t*ZstV(d1e@O0svRU>1q|O>a$&#t~ zqtdnqmDK|5QguD?@8G^$11WVsZmC5israIn8hpyrHrvu@f$CtwmvRUCmt&~+(sZ?P zWq7w-thR7_%J_z&TfC0rsk3|h-Pv5Z`Z&0g@=-48wVJDX{eB_qX zEkn=a?2)p$y}p-xl(SJw}0@?(66uNYA?)d&4WDS3Pt_e zLi=8wRs4U2SYO^Oe!OcF<;kGlcu{>Hu>qK~1#oS3dNULlQV9f9R1ehg_0YkIpx&Y8 zVm0dtZCGB&RhU&(fz9l?CMX8LXgpINqi4`AoR4jz3k;DEdQXFbMgQh?3;+ z38~Sv(l(KtF`Rn3Btuh8!<3b|cDrM$eQRRy)Oq9~mK}k%gbkQooQaAIGWlT#@)7|w z6)=m4?OX~7X=U2Ad8<*n)s|GLSZF9cQFk7wtm%0c%v1LH_X?V_j1qyeA?|!8%MmFC z+BT)V?SLS61gsW!(B-p4_#ooX6J+;Mv9W~hXl3?22tB4%`FV4U7r+#RByW}Amyr-9 zVmaD@-)IAf3f+a6i0YNBItMcMI-rD99d1K!=*)-YX@FGiFLM_cqVD+^uygH<_;c`d ze6r}lm%DrHS7}Z!+c6iB3X^h07D?x+DAKxcU~aM=sBlE$qOqsS)6ZzT+a70tG14&-L$?D8od7&o==CV7i4jQmK&ID4jFI%wI<_Zx0RU%$iB8#+_6_eI+uGv%-CAkZ5Hq*@Si$ZeJY zJ2};ygH=Q2x)f+Dl>>WlMGa_yqF=+8LOgs53wD<2Zn)fX4;MP8j8qWtQ5b*(@^XT> z-}A_J>6QUp`mU&%xHP?BoJD2;q@7cIRUZ8ohUn_5EUTBn{2K`1?}hmcVyjz z>8oRJOqD=Ha2+Uld_d6c!?Xj14+)*)7D5CqQC1b&Q(y@3G?GVB!LymhrEO{cl_%oC z1nIp5)qnsvJyF9TI|BkLg%dUKJgz5yN9Lw)P(}7)61skqoihlWs>HWb;AM!{p0ZS+ z(j8OQsVp49XQ8Cz^Zy-`Y2Q?`OnwkibdQm#i4i2k&_DFKPKBkL|RR&Lc5vxkp0%vK9 z>~!<@&*2JrxJn4h8zXOfF zB`AUfQ1t=ix%vg8Nj)>!s3;F}R}$suf3f!-U{RiF+c4Q|G;4{QjlGdFf&zjfC5V8= zL>&;Yf`EWV5Rn$C(n*ZEF+p@_O4UfQP?RclEKwY~^gb$0x(-bkX68Sy$FkXXci(rv z@A$s=_>cGh=P+tyW}bPTxu5&GuJbz2^MdS>&A~Tv7>HM3*Ht{=Gn!(3UjQrDf?v@O4PD#caMLSd0#Or7k%+(gRX3Kh~{lhQQwiq;o z(vAhq-fh~J#ha+#qIYzDeSKcaloo3Jc<_ew&zgVCBw_pBhXFw}iy;btP{ec~3tv4@ zIt<*2u*YDcZ_w&xh#P2BSNx4@6jk!o8hziKO?}M=PcNeq1R#(qo3a}ld^c7U!HNqnZBJ>%sv$Xp@8jIvkZafCS#dkW?Pa3aP zs*XyvvM>&r8#XSlo0R1y@*T!)0UFSCU2yU_EHFAP2d^#h+0Zw+uBO6oES+{QHfN?% z{u-NTH=`56t;;>F91BM4mb%sK%mI$Qj_P>Su|43*`R;qhAHYgwvNGVCR#cJ1=NgSw z_+s(AZWW?(Y~333M9_86EVVXutz2KNdSu(guJqE(eyfm#LJN~I55GsBo$kO+TnsYX zUda-Dh{di-W?!?$r8So{8!P3%_)27}1$4uo?p*nl+n>uuyY9N;MrBwF;8_i@wq3CC z`CAo8cMnktTjypc^jKJ;((SRxF1U6=yHcGPEjOnr%-pu(?;$WQ40JlnO@OOe^}I|^ zqV1UOTvj8{-e3i-lh%!o^8t@Ks&s*waFr{5Awr?7;8A+IX!VdbbFp6+fOP@bLRXkT z*=GIp;@q!Ok(a9}q2X3ID%tWOgR2)iOoKU}6O zxcieL9R5eY4@dtcYo`Q|^}2O7<=asO(2f}<7a2--A|4^3q1JSt;$nt_(#B9pn10d_ zj^@Yd0Qg>>>Nk1=b+#sfIRf18DwK}{M-5}mwor0K9uAaMo6zNtv_>kYhGQkA*mjix z$M`&*X?et%+Ix3$OS3=4pEsvLxvh?sL6=)}oTVww?ufYek_P`t25C6aO8puU=8v-MSEwo0Byhz!$37(4hxu!aHUb}@Ye20eynv`1X+So!pPk*e7md8qIesrrOmm?6nAnN< zNg|An10NDGmU_o$fp|_EK!T>n!OWA4_+n}l@4LZNP|=*4w-vTjNM4dl_N&aIxr8T zKt>CEs5Rc=GSMNcb1^M#PGRvxJy#s6gvEb5Gbl=BQqze2)Qp%+wq2b% zT1ocU9%1VF9bqb9R&Bs`MAbE+(=Et?ZP81;ax5{%>nJr-2$hBp^42z-F-D?jgY!d| zCJp<1#yv;+Momw@K8pQ}Ox4B?Zc&!|BKbX6li6&OiN%)L=youSFL1BA(EZ$w-2|q? z=b+)$p>!O}yr$ze-K#m^AX2gh9)|`U5PTY?PC}oMRB{iiQ>Ii|{HC>7z4$%O-pBjh z=Fd18RLjL}kprwyXKX!zYiFUY3DQFQL8|WRbvCcyeeK+en>NA4152{>C3?y))p_3S z!QZS4&I)Cp_G6E0aR&Sgp1E00?G&U+r{3mU__Mt}GC!vCns-F!C$+Kf_|wvkP)k*N zGnGR<75&`Q{SDUEq4SHBdTxIEj=`~UxL4=%9!cQCf#_*|--#X$X@8g-;%1;GF{m1- zX4^Jf-R-b?_=S9SU4Rp7$!xn8?a{i9%#tt*kLP)=j&*c(J8h? zeU0Z}N8IGVoI&n*z134Qy;S*0R<-oD!6g|Pl31Z^MQlt>srl$Ae^GF$sJp5!t(+I| ziIJ~j-TL2s6jPNiu0_Wxn1q!2ul5`ebbLP$?3tr?&{HbayvAmv8=|WWg zWM9Wqt+4W2aux+0C;HwilI=7%@{~{VWxg`0x=+w*%GAS$JP>Wf*)w=^V)yQf5&C(jkXvM%*qn^<)tZBgB~7Tf6VP1*spU*v|hrK=|rZ(nPQ%v>~jT4I1zbIHWDKYW&V0^<4m z!BttjA7hn7Q-ck0jt8p}eDC`;TSVDUHOAKuTk6T|zdbS;?H!z4Aw2GuKY2M|`;d}C z^3QH(oD%5b&cmI`VAU+JtnE6F?%P`;9z@$!mP*z;5!Jrgb5dqz@}8NnPy5fj8@Iu>S*!88p!n2Ft0k?9yk6muSR0tB)+#uZ?oLb zZF&}5!3lBX^3l7aKkMj+OP6O);wXw zieVX+i{#=jOWJNq5-i!1%L@c{hB<3ZMg}FxnI>tx7eW07FaN?o14gExgb){8a0p1ThJYGjTW*D7d_H}y3O|R>TGBahe-RH?C{%UAv zeCMcOF*R$M$G=qhyk@J>bN(Ao@c%H^?tkmk2fBUUPXrMxGllVPBolVdl=p`Ng3$m* zOjE_2iF%BNy4`JJlraw@BYY>7QH3u~D#ZZ^32Y*Y^3?&*wF#Qp4Wa}jt2$JJf@=ZY z!G#v&{EIN{eeur}gQt)PIPV^T-(2<_fhlDO5 z%l_`0H<0pC==;s?7wr^1#T*vf8JS}p^)9Jf`;#o$wsdtw*rnJe3*J6|-`P>ce$vvM zUhaiz`nQj_s^oH$29W1iTwy55^h79`YV3X?K@ri%h*?CW5~!aNQE$wiq@D@(0bTU& ztB~(*aRFC~#TML~sg!h{;eOG3?DOo~)6n_DOvNKFCVp90{s>!VII{cHLK;)buCfA{ zMGbfK<7rzV9SnMD#Jnd_p5p`BN6ye4IhcJ7ZRTxBi@rY z>n6RgCnxKhG_D`LFoRAmBf#yM>K3Y;QlMt0wmlh;P%=%03mB;k7Yf*2-FWoh=p3m) z$=FpL3v%f#=#de_`RekHBLknC7DIgL49q4D1W{rZ%gM<(SL>^HA`7D&GzN)MKJOUz zU`mmd(ImBpisweYIPDWl>SFDU!S=gPNa8$#FB00d4z#NbZI4l%(j}~7i)-XK+Lidt#F$K0$#J8bVb`@Zh$uU^) zcHXcQ9Yx0O{QUB1AUE2h<5aIuLe0n~Q83kWdyl~Ljnd~`&@7sJkQB zkUjxhUteD|y2vK#JZb}@q{FpDC8nCdE}U|3sTa5A`Zjc8UAhtpal=GM@I*6T90q8Q zP}@c@7KR@v`K=|$1FLRvVBScTUOJeS-GqMA!HM($Zw8H_{9!_|Gw&H7zY&e}a@(9h zu+PfPZJO-+>60#8*Js`ti*MIXOEzb|wmWjG3Z1MZWdy(jRQWdaA@SukK%b1HD*eeJ zpClV6%q4G<`XOL{*y4Z(QcT@X!hOz-w#s<9NwPT z2MhLBJ{hKa2q8)dO&ts0W(bftalA`<9I%*3H^EdAh~ICKBWdxA)3oFG^3lct zE5r5q&6WuY%q)Gs9A^LMyZ$Cf)$J|uvm3-rDf>=cT3JNGBVfaG)N-i%#@3iffdIF1 zmDlBj!^a$7-Tz<_wgSR>s2NtLc&%-6YL+Upprzm4dGA_MQYpaLP&D@5(bKS;yTNvF za^8>7OCgRbtw_WaSGD%rejM(Aj7 zr2df4wak(u?l*=DLW0flic6sbpp`e;CPtGUl{P<+Lf$aUZ?+smbYiF7Gp5Ug>g8*u z=&Jj}J90Uj15Snv8{Pm-RuY5?xMs@&BA>^UDhCQaXCj*2_iyNT7*^* z1qV**fC9nh0HSs;oA)OiM9S_x2G>7gO+r-SZm||ATW#Y|>j)RXiPHgO%|WFL-6}l} zv9Ii|;#woNYrcWI)&L%FIaag`3IMxGZQH82rGBrTLRDXU`SkJ#IC9qTH31I=vli_C zZWE!vlVk9S>jvWl5>RVm0I8-k-|9ZScXIFJXsHlG${sXw`Y`Aol5~w=CKFiCQs0jR zkU$J63Ur3-RtDqP9t26k_0r%&L=t&K#Me4g>ppvufJ>SF#O%UJ?@^z{rHU=v3Cq(` z;-FK$M`A_bM6tvQ(ddZ_2CewMQej3J~Ytp{;Am|lsf`#OrGb?1eRpI z*I!gXa#oWZl!y=H0s#>oA|a}p(gLc}`yWma!6_aD;9<8dxOP~q?&ELj=3lt!q^qw^6s={ne=L5dU((I)XW_!kgcU^jxg z(665t^ZQ>$)`o~ro`OGdsCbme^P>8=$TvvK2zUybd@_% zsFYt2#ypB!B~~AfP(zQ2w3{Tq9@y3E368k91UzUxyxBPMcdJ8E&~DuL?Tr-1#s;BZZdLEi5->$bnhyEotxZB`TY_0+n1OUV zm_>SE%CV@e_}Y!4I|(rtd23asoiRo9eeXAgWU)_wWWGzWKW$4*je>!vNpnc(ywUbg zX$9U~E2I!Pfw6l0f0h)i<%Tzz~26BjK z0ku0<&W)Jzgt&vlVkpOEfQ^D#WLM+C_r>RI^Rm05;e^f}i}%^-rD=!NXkpOfj;=gT z@%sAzJUOV7V@OGOen96wrTdB$d{-q#aex9Tp{wIdWDufZ5b^;cGYPc_j&i@q90R{i zMUEGcz&q*Q)J+kv58(Z8Gt`ZtQZm39o}k)9tgFK;})lO!7>?lzZ2GdZ zTErP1GaPIFwx>HHH{PYxGUM53gp1;lzjcys&d(D~bXX3_Xsur{-Lcj+mdm}EujlLOOS9za^A8=V9Gl!e?53di4P?67h*%Qa&*?P+^Yf@_+I+?9~yDiDP zD$Bmdjm=iRYEvx=OVsH8A&gO|GauOQazt(ihs(;;#W>#)=eYWHPl8@k_m#R%aQ<5Po@UjF$wckkp?eW?QcA?D|McQ%#Z0?@e7f@P zZkE4~pY83Ak;`n)$e=-f;naz~rWLb?YU_NjS~i}kcP!7jv?Ak?P62hbNiGJSIk_*9 z?*a8dfSiWI$0g}Fg^GReZ{X6Ob~hSla#oy-*GUv#>(QQ5kUP>XznK%bUVUlvfs|hG zO9#`QWk%8^F8hl+k-vy}0Uw7|6#7WYDI?6pBqHh-19!bZV%kQCWg-tz8M<#(j3w<# zZnasFEDw=O&uHJ!il2`LyIbqsSC)8XALOcPEXeq=c19irC2I7#JeX%pdm;L{Yz``$ z5e!47-5&xoG9IFPh?0Ji4x(8rM;ST^*P}4Tr4bo@r1d%+wsLMoD*o7btEYB_hF|!@ zfr{A$94-%wUv zCw;_=|2@xkbGQYFrA4tb>*Dr)l(!#NDpYt66Y(x7Nr4n(E+BO<$PmeL(a%y_JPe9I zlc$t-mo+31K?v>e^D^5F9-!8!YaWGt8 ze#zLxpf=Cb-{86(M^;uzWRj5NYI`Amciz0z)`X0St{=iGUo9yJCEkSfh&jTe(jr|3$8M_!_gTpDB|&O*Lx3-f^hW)DhJ8ALL?nFb)7ZPe-7vVZ}g?Jla$cPAA%iTzPYtAYc$e>eFv-6=|wiLsOh(}f{wfT{4}{NKSf z=;X&DZTx~vW|qTM41lr>c@7|$nZ_}o?mjbBy9R6MoEZBg(#mfVV)u`x#FKXZLQq?6 zAdi}WK}SoJ9!}KAI{nVAzub z7AZUje4q5suLKWae$F9^qqBfo3@IwPRoXCcA6pT-XLpu@Fie$~$qv z55#A54fh3srbL3~2?$&OJ_tI&xkxDon>U1|@$9dR4$&QFscz;#N7*qeLbYFj?hNSy z2@&XRrpy&5riF}geo>{n8@UJ>Fa#OLH#VHzk|uU*kb+Cc7s9HK+Nqc@UD zpA+xz|6e4huM^cX)`sAh(3_B184Go~M?_LKZrf3BFWe3h{pn!z$+F3M2(2iUTuX43 zND$0chL8;13e{dxz2KB(OeSY;W|E?MU{~;&yTZ$~UDQ|XZwTr#HJF1B8ONF2SEN)H zg096G7&@ScLr;|2Xz%75fn^4ITqnvKr?>0l;cy>Rn=wQpBQYLj!n;&ukfALLYbM|u zwN<(9N#_20wgazUuBC3^**IP4sLA|&CNWq|^TurXg`yckr0kw2`D+zs7HqVW!`^VHI`>b;4grm!JD=vofF=nQG0xGru=%yE#TH#sF5q`ID*R z!?j;j=zDrL*D2Z-?AV?DrTWymj@fAYq zT+So?DK58Ba7Cl3NNnKQP2mN6nlVnIM@~l= z?3W%caf*6}>2*Z?$L`b}uXrpnGFl=x3q~x5dxzFGnY8gI>#`S*On%|eML68B?YQAt&ES?H{_Dd8>$4#P4pBf=a;v3@b}aD&#MEYYmAkE zsou}8X|>mrK*Mx_KD!%gY>QfFyzmm+Zg|)ZB@;{yjT)Gp7<|BH1v`qHllH|^7bir&6%7lb zvXxyP=a=3vxsVla-FWre6a5VZ{)M@XgN<>U>sdU*7~$^0i<2!gx~HlO6YGLZ#JNM( zrDl>xF6lBhM{AQxE%t|^V;3kdmu^<^{lsLKcbr+CLD{Wu_9-PHhX7Yynk*BE3Ecy0 zTykU~J*%gHCrfvVOc4>aNsmemN&DQM1R{|V^>QgPfQh>7!p!1GoJ~OIOD6wP7=Adl zXpyKZ-P55h?_&D#O}#Vpc9~~?8QfLzNLp61+S0#2`=XI*-=Q)SqpHIveZ_)uy@?@X zF^5~`qXqvw-o4xi1=C%TyuIsOhL^Mtk8bsk^tj%g^ep9oD7{!VZ|Sk$KGqni+&pY4 z8MP!C$?%=`=X4M^kVClx0C9Z~B&u**5wZ&?OR0#X9wG&IcU=fr`p z^Q((ztp{9o#15ow&KCDH7V3q?os(_zU6HNtC&yM!J$!J`YxrVRM(53?c@q{I8AIo` zxJ*3_lpTK3tR%zE>dCeBGIBM!mMPzqQ6`S`QH~wF<+(m2v+U520dvl5TC^~=%7}tsF$Y#$_UuixNR4;Ac`Ok6eW&)U`Jcl5qtN2UJhIJ>;FqYimvOE2`y?C!edA(PK*B;tJVH7By&6Pkk<_e@%GZf!NDXt;B_ zEY;+~%XJndA=k%yo85bEO*jiBp1s4gC^%NZADORnZC`DCG&)({TfMO@yM6k>h>ska zQfl^ZTzAr3Yz|%pydtV4x9E0gx;HG3w@!*J_w1_M@S+Y<*wG{P8bz~O)_L2jT;;6= z;-8)WvVO9{a`{a?T|HgnoKp|nGCfBm25xMQmb#;t`K4QXB}#7rg}Tq+2N~w=H1$ip zvaPtv_%jbhr`XMEtp`KvubJ$M9F8ei?Ih;^q-^He_~lmLZ|nEQ%&Ajf?flNPI4QX> z=Yh7yul(|hX6c`{CvwAd^e4OzD@26rS?20E=#>agiZ#_2yA1cW_okj6Z(X~rxGA}+ zY1nRlS1FetWz!ZJ!fMPP-m++-CpEsHe&Jo+*QPx<#{&gbsnw@?>TU0{cx_9ZLVI%` z=e{A53!L@l??GUbsmtVLDH{(ckJ zloDM$A6$3*pib{yyCG@8{JPpW-$=ohffXDZde?(uzg0r6e*^6r8jZe7_#Rh8bNjS_l z7q1s(PVP46bOf6VRTd7n{W`H|-GKCHELZT8dAO#8-M*#sqDe6Pf*sU1#5TLa`D$i* zLPtiTW9d{={HBOyd2=3h3w`C3$}bw}WSQHVWL@tsX!8E?4G$SdmFheBjaz1pB-y`1 zoGpIbA!=Ke+hu0f-+XBxi*0EAlGDWEof{OTcTS7+>355bT6;_TIfrG)9p4hOT(aea zc7dFnOOcI%u)*d2-sTw{p!R1ufYU6Zuw2 zaqAY9M!$QPe4X*1k=ODkS3=dgt+)`)2!Ok*bcf&*RQA6{5H`L%T6eRg(Iz$;yB8yY z1E?Je&QMnrrT!5onY~JCrA66dcOs8)Sb*6)KGK1~VG?XzXT3aLw9hGuqjz*s;JIb=8`lEjU1 zfzv@30$mBoo?tJjMxJ{lEGMT5qHgPj{XuHuI=Q(sR=YpDd%AlNm?s2ErdH9aYjt*P zczdaij0KxVkpB;9*FB$nKV5)&oEo)kPClXned2WZ26iE?d%8#Ha3yayhNF0d2XE&_JzYEU%<p$3Kir?M`vuPJz)i= zF2@Q^V~>cvPmM?9_r9DCv>?i0r4vj~BTo?DgC9!W48w%g&175oxeM72W5?A^8}XM# zPx96`b?mUg_wT&k9fifZ%^i`3`-8ESXSuJE�Xi+0m`(i!O9PrfE8w4h7XyKypjq zj*Kcc4V-z1wa)VB>SKtf>f&^`BX^SEWH~&dHXz+CGslizmINv1H=9?MfO~K(FzNlI zSmsN(Gig0}FUK>_%K)yLU%hs?Jdj+UbTAz6Mb3VDmv?MK^dMuG4q6~Dp5BM-uJGPV zFtW-pA1N}Og{yVsH|E1bh<|u?9@*IexYwr+Fy$9Le6FNuu>}5Lp8^#)QqI8(-hM+n z@=-!NljapEm4U5ok6fS*Qj>^uvZ^@3`Q7@bn#2`q+>i=g2h&l?gKMucLg>OJp|5S+!XRCV5c7Wq0vKVLN!>-|KtrDfj1GdiOh(L z?i^Xb>;}(;#CZSs&Lw~;!r=Ka#g&Xp1cVZif9MGz{Z3+}Db!qSn2JzN^G}XNq7kyB zWAUi74q>7VWJx8cPSb$*#e)SzZt}fDPWab1NW^ROQHuGZIV`H()&bz~!8cS&+FX=R z9Y8-ymN?~~Cw(tTh^Q81Oi*2vs63D_d$YI_)50!{AX6+QY_%Q{PpM#J7qMR;*zO?b zNep=A|K;zuKSNX&VWjIkZ|L*vlhru{?h!8QNTl^b#@Li_;R^iZPeMahE zh#ryBhcyq;-lG>vPmu3;vpDo{7m8s{(S?QMNUE67-l3AO@C?%vZ@y;(?CxYka~1|# zsP754rDg{?LWYESnNfyiH|O;p#atKx)ZYqhMkWtP+)Rs=f>ZQH1tVnS0 zYBo_G62TWDJS|WdCpznGed3BC+jB_F<}_U#vJwN9xo0D7SQclzOWk|R4Aq}^*QI;N zX=i$d=S@Z`5_erq(T!Xl8FxQsdV=$B7j3IEx6bTecxgff zmzCKldT=V2zF{_sOPQa=}Zttk6&u*6fLvtYW_({(}C zV@~CEDq9CK)Rmo@mR!^s3+i#}EU?Xn7jWiSOLHpEdqGTu&+~B2qlpFQbF%lJzs_w5 zwJ;81d!`oNNY@IOPn!`SCe?0N9J=yz=I*a-($#chC!RQ~EBwqmPCrPWbC08N{8?dX z-e@n2$MJMg9xCT2Hv3N<=?O>^xVDVnbZuzAl+ly>v?|y{+zrE4#go@PFRBXqM{)%f z^82$77u$IbJzEEE+>)iktr^0OM1goT{P5$M{`#$XyVngT@(oh?Ev>uErtfpWnk5hc zj{^_n+?zD#X%4hrO1|$Xw0WG~p7tu!Uxp2%2o8&UVkTS0_7|&9JPB8q(=KqrKP!*V z;p%e~CcF7YM_Y6Jr_6VyAMaV3H%l+uN%AGT%CTgbpzE@k47{VYN5d6jUVYo-sCIl* z?N60onJV-PfbVuv(S(L8(aEM@{%n6Oo<2#aC)Gy>9Bx{Pg(hqK?2jhooAQlrmE2^p zn+@%2J1z+P-tSA*TkpF6n%jJB(Gihbo`a6Lzj@Zqp%90B+oSv#(0awfEB6%dnBD+r zLAF(e=zcTuRFaD~>r7ccB~3lSwuU+)SH#x5E=9T$Mvs_{$L>T=_tN0&xBTrZ>tIo` zXSw&5_5GetoQvkoIJrmbJ8x04x4z2IyoQgeeX3u-|^_YCmf&1zswq7G~@CLJ4Th>Y#>MEMol`@Nb6SW^yHxX zR1G&A#zUJ-O;O`^6D5$4OJsm3oM1C{7*016*fcI;a;-JOd^v;*=|C7K{sPaK1Iy}7 znN9G;kQ~?3z{}jt-%Bj$a@C$%!sg^IaUxJY+5I9O_m`-HqedEV3Z1w! zm6XVIHVk`Hyh&XFGHS-Z0;$aGrc}UlgaejSA_XiZA<4eSUHUaCRUk?C=zi&G1e(qpc(f8-y+p^XA&+D(AIP>DK z?_XJS=GhNn-^^9``O;jom74MK=fYA+9MTTS3~GsFU2EtI==bRi+NJGj8e3t! z)@?jvgGTlZuCZm~5zFRj&7Hf)8?xSeZ`Kcf(T$xz!-~mN9dtQ|pB+3fhakzKAlI^- zhE9xa1M!6X@*0urf|j0Yg<){5qvYk!=|Gg|H^9#s|7KA3FdedZKe|1lzbwHg-fVbz zklLXKqR(dP7u1D^B>Df?ypx-j*D_Z8tHMMFyY(lUnsS|ZHYprvFd84(WFp9=Gt9Pk{( zeT4r^hSP^U>Vy1-8rh!p>uNIRw1jpVOI#KYPMGwHdrAk@==}B!x;@!Ad)a}!(LUuZ zK8Yp6P_CD^t<^_)MzP>`CEptOoR>Z z+opP7^R_Ih_Af7q^)wDTG5%gDt`m8GYW5{*zv@ zMeptSP!!QG0HsXRQY;8}>4BpyUp&B9YaEo8yp*nPY`h4-Ia$XrmKfgD99q11vEWsE z5Ve~~a_NbNaPzyo5au^_~dETrMQ< zIP}3Yva(=dlTzOfLx_ktJGBT5cN=f8L?HZOAPMtbMhrfL5vw5B7$1oNP*s~nFs--0=D=TZc2%V?6U%wUg2QFK-Osnaag*!vj z(o)P~@G6V;`iM@raXCNmFK-)>wF0n7jXgshJrv^Wi?tkcTx* zbyen$M~@a`+?-d&s+w^NObIP65ss?41<&whZh1{ns!R2k-Dfkj*Xaz%hsUSKYs6o1 z*L^T3jPJVbad^bL!oYUmPN~ZJ$zKfPznv^J8MREWf;W2Hm229apSZKjjtxD!kS*SF+(Z?0{;hw`YEjs~<^tZ*lais1X2wW`>z?)c2D5TJWo&g$BvJb!7WK_4X1NLS*I4!RC-~`SNF9z-qoK^lE-i z=Qk@I%gU(6cM;UI-w<%MzdJadBr_|it}e_+ayr~PR^(9(x0YO+s6{iT36~#T*~0zk zgF%{5wouW4J8WlCm629?(EjmfhdfHM&uVGf@Iz`GJgu+)(CqIl2;Ch#_3e1v%adQ7 zu2ydk(r%U7U zEm{BVQ};v!)=J$Roo2Om1a`<>fM%Sjd6anT)>;e)E`_Ktul}gC0@9fYn6wka?@Mx1 zp-Eumjl|}FsWqcMk)I*qWPjr`DMf1*`Y#WS4dqNnpAD1;VUF&!kt}XZVbsJi%P$CuWg%Bzahkb@zC?!G0AJ+M4WOJA`&Hxni(*FnMaZ zA|iZ;MV$0$8_4K)#-ZQ69}g0t#zu`#9n=}qWf03_FeyE-7x(5cZF^*3KuSompJ)}9 z>qw?pJaSM6Ip=zw-y4bowOyh>MQC2egu(g){MVmXbj31!ZPQ{TDh8OXVsXI+x|py_UO(mK z<>d5=e=T({Ldm%CWnb!>Z@_<+SDGGXO{1X_LmX3z>F(p>pv-7c5uPX9j{}nT7fC5d zb4jokC@y(L$B;u=<4n#IoOk;R3I|4_-ZN2eW-w<=8X$=vw!5>J%l8%3A z>$S?({9=gF3e~(7!EKy=?@hae5yHICAQ{6cyR?yPcHfz}qq~$R;;bY;du6poh5JS1 zyV|8ao5Sm!VB2^bi^P`OKFYrxJMF_e>zhOS1xgZu5`hkC*J0r%++}1`m)K5zt^*pr zudz97OkBt6TVzW`0LQ>tEW&`x4!9+eC$G6gf~2H58Q(Fvv4Jo<+y)0Pjf}JL{4L>sY5?own+SWV_$mTcy`M67;e^$mb48M z_lxRebNl_=ol6aa%?(WWV|Rm+VlP|kUH8yWDA3V(Uc&n=d!+uaQVl;2L^O87)}}9J zQG9y(IvV~#w;a_>{h7YN>fj=`4lVNmjnz25{X8ZCxPK=u+qFO~qePQq&DAq18)$Mg zZyg<)$sf`^&gSeC%AWSuM5!)QvE4B;DkGU)r?4$-O8KH-aCNM(QnJs|Ld!p0U$fbC z_b!XaMVsbX+sp#Ax^4HgQ9db{Ir{N|>(4AH0J@U*g`Ty?x;+Q{xp% z{I<}HwDOK@H&Ncm;HX=BP%q_lX}wmZ2YG4Ub=iG-X(dK0SEWAEi`}ql!0l&YCja{# zYPLrce7xN+-Vzmv-hJryMFJ~TI5qiO=qa1} zp&JoX5m}n4O6nWdefn8e{c3||&dK+y)bD(^!}K{`%si&;-&!*1Y<7Rr*|ziHfLe>= z4e5ppQVlYq_30(+A%L7{tndpXc#y1rtC5}9ajMp)whMMm)N`eneuW1Vj**l(g#ntzgn{_?+6FvubpRwmeR9yNxERd$f z!J+>9G0Dqp+RA-WSzEa>dDLjOwX-XsId0!*B^+|xA0YkJyR-2bC zFPH15p`xnlP&0_Eg-34RwDXEQ3I2Qj$?H{q$ycJ-*xAZ18)Ta}?AjHx$6%Yx((t7t zp;le7oZf4F&4wDQHC}#i)w_sY8<+B1crWN$>w>ejIlp*s{(jmY?u3n?wxF~Zqn5|V zP=W`u$V31K%ih_Hh)P!K|KW%QI?ba^ix17wKU<6jn)okG>NnvuQ}LSUh-iiCH`vQ_@9%;yQ4* zNP?z;c^1`GRUC9?uKknupdLG{0|qQgrhtd|Qhddg>m=vQk1&}&f55@1gw%(mo+qAw zaCih4u~L@Vfu$}T4;{C$0Y_`Qon7A+H?p~Nfl#}PSb{Vr0g1;_xa3a=xe_4~z_<6s zuB*a&GH?Uj;Y6Vqg@ZLvnMN?vhAB?^k9+usgTa2azMYi&j7J)Q$6(i8IlEqte7?w4 z((flmAaRC`Yhho!RaA5b&JQ5tdZhRGPG_))VD)9H`E8)DQW_@#l(_a-y zT*=YV(HvwuGv_fo*G~U7t{6B!k6hFt*0J9sA!~C&b%uX3^!c2e7(6}e2PR3_L#9(O zJ+A;Ci5#C;Y@#?42tl{E!Iw{eMAddU?-21%1HRp9di4`N+%IVa+DtcQ2?vYw{Ygvb z&B7!R8}yO-(N!YW|=d)JS2(rZ+(XUDtl$U8TG!5E~y7z2y8#yin@Scf+-*eyU_2fy|ZWCX`01E0L8LZE(=bvv{QB zh^zlPoJE~P*es-d```?+eHrkKn))@Mj=E#f7YqI`k+2}@AH3xRx&;mVz$A7sTD`RV zV4}&+J^ZUlp^s;f2njTe?ci9A_SQn#s~@EH#siIGg|XW>+1|&8;aJJo+S5d?+G z-?PFxl-KH{8!><0`jFL2gr_>5o6T7E!N$k(J^lUt3cvkCY*H%DBi_4xW#9O<9O#ti zefbxufM<^@E6>5rLZ-*764|ralfN!p`5!APIXRy{*sVda-67P;R}l7dcN6Iy+tpX6 zYaFh&ZJymB>cgPIOYKGd>5Mx}7a^OaflQ@40#gJnM_E<2aCV&=?cW%EOa(#+gNys58F zKX&?8rCaTS^%1LImu|Z^)E|6iy7cp`Zr^KF;*Y3}|2*SkoBEp92R-O{pV=f1a9j0g=!6MP^ zF4#$e#e5j#`HN8WZyAur08thXd(j34C6l_dBh8$rNTwk_Ql5~PnVCIz{g25ghc03W z9gA&k?Y6UZ!Mx`BT<+ZaqCOB3f>{I#G6s)~{TsFR;pgZAbRa*liO-ByC6~b+?MQi~ z0={Zau?i=@JkeZSv^7zpafwfS z5txr1WQO`mTmz|^G~w`??u3|I#Wynu3Kj`1?)^O8cMtM@L5ke_Q3cusJR#$McG_r{^&-&lDjAnf+I%Y4?|s z(u{iMP^R32nra_6J1X)~qnob7L% zIC=I!UuJfyl8)qSl%luyW!1(^pSzx-E|8CdY=owUyXIo!Aw6r)m#klC1qdWV1#C{l z^=GCtuEr;8E$|ugpRtU#4WKZUtUUbwM6RS{*Ip~V_X4jbDY!D%Tk6iA_x`Pvk-C&h zPkjB@q4NL#ovy_H-0N^gyD=UZ94Rwt=OR}h+H;_ubRk<}ghIBjuSh1S07;lj$N8;? zfk=?0Y8$}ZE-DW{-SU<~iBK?U6;4gbA!d#HN-;( zQSBU=)*3JuA(k#M?d^i0Lxq#S5VF@$1`0#CpQ%hdI_9C4j3a1| z-3sc2BN-1a)ZCpS%)sFgn~cqo7CqmJyS=B)11dTeTQUE|Jth|%C0TTjFQ;!w(p zr+-F?D)Y$N@>n*{z8peV4vYFKG!gSZq?=3AQwD)L&u<+C( zL-ORVw^u0%R~ov-MbaEBqp!M#mm?%4 zB!AnkBLu*hS5aNn#HkC5xfK!`Zu6U^_>^gLOVk)TAr5kK^1DoSUX1JOtAG6X@fH2r zzg`ascZ=p;zI?g<_|`p8PiLqMEL&#MTvD&tGpV_n-PF`vm7kSUB)+^Rp{p({BP%mL z{_ zP*L}xy#G3nTw0X7$P&Vl>@r^@lRtNLby4Kyr4O&`YnR@(_nU7nZ(ckp4{sU~%W$nuD7{FNy8N6)Mv@2mG-xB-4}|B&0DW$Safa;JUct3$!rEs+Zh zH9vLGn?E0UrX1foyMHox;^aqH_)&h{^=Wao)(F+VZhklz=CsV)y}WSMwHr#VDrQ-d zjt=?6gySjqf6M&xYx@e}F3pxwMY8ILTi009zFVJT(3^P6Pm1hZ{<*gAnqJnMs?YPa z3|sR$Tp`PQ(Z7SIh;>u}qm?LPMp4>B(~g!qyiGiR_8Pc?)(ABOfGezc7lvl7)!d$@ z=j-d+_VTO;b%hCUiKJx?pZ=c5G!0J=IR+)5NXN$;B+$CB*i^>97evm*^aw~!-^|KW zQCIJ#F&8waJG)@OfhI{4%xl+-T8czDYz_$ui^LdswSwVL&Dwv5c|Cr!5B}BA-PXye z0+!TI8RvD(P!n8{7Og}n6sgsP@w|*ndEiM#riYdF z=0n%zX}%OJ}fJ_Re5lh=ILZ4Z6DtGb4^3T z1wxz1n9mt%z%F9MfviYl#PTu2y##W-SX7l`Ucg>0V@ZBZvoF!WaH-k3wG)H{+}=yq zJ-R|i3b>d43pwx4`r-e4{9Cz(@%UZy;s3Sxw{mU?Jf`~Z9PTf)^Lsf_k#POjUn?)q znpu{dJaKx6ACXq@(Ms2qcm6zbE67(tTG&wJRWk&W_X!MUi6I3FO$|PO ze;$*1@67~OB6O8@LHO+^dq1kTX&Xf&*kAL@6U8h*S zvo}PxQ(QVN>+eG+;N=y7(rMNNheclBL5HX=A@34&3L{@Hzf<$*Or$n`uZ02obLV|*2VSbSqW%qW3Sauq$^nvmuY>mXeGsubstdR&zWNRc%Muxr=l%?EUpcP|B zXy60}u*$Mi+t(Igo^2WS5%wa_slG7?d*n^jR>krd?r>E3zsPJU#&?}(#AnFe2UFs=JN5)F^{ z&+0!8DhL^VvCRMUCnPx^oIO&sX6ex77;NQauWxee+iz$CB&}qt&=D45KA1&X2DL#X zYT0a%KJ+8Q2*M$(0vv0#2Wysc2a)ivbSe<(U2`&Pz9L?-HTvRx+5THQTlqS!y9Yfx zf6^bztau*n(i*9K@l4ff|18lNPj2a)m^|)QQC>pXBb|!l6*56W-XUIlD$8)Ses6Dm zz5h#d-liojKeW91&bj&vvsqV)jzw10ou9I|aIUE4{Tj;NYNPDA$GLPl&wGx+QLmDh z6AuR2?K8~=^D7@pN9K37XX@ySxQUELvAS?(Fw%i`rZ? z6tgv4sWdw_Uz07U=?x#QOzUBPu+WPmbfPl|mCLbyvROG)X^^LSQ=I8p2g-6RB6s;ZnZ2c>X$ynJj)lC zego@jR*cpb{o+^G&wFO(Xx}hhGMxN@hg(H|n&M;YZ)>mSKULL}9?e#Y%`KeL33AMB zjXQ6iT0hw}O+GEg`2hd0yXS!u^L7|`Lu}M$-ud_o4p(SvJds_&lGc%1w7Z~FEY+}f z^rN*v_zTO^X)@HG)RCErh!M{1vM5?uf4RE|lILT|*$73}Ns zoa}-Wry_^j_Emp9(^fJN&6}!MaQWk%7oRf=blN7gmb9my&grJ9IHV||@(X>)mj1Og zPMFhy;8OxCidu6W*e3oSdDCC^5gCy0?EE6lY~WQ$T=({Kb5?sR`umKYv^cwXy`P|A zQos|gP>tk@dej_BCk6c6@+|XSXGP{7}W0EYDXD;*!7c z^fGZC^B-?ZiT4$@|MAjs!h!H%Z=!E7D9DWqZ`%*Y5i)PFNly@f#GoCam*Nx~RhnHdl+P?YyAHVPBc_>+h41 zwLI?aTs9Q^jeLE*^lcY`Aflu8RLa2l%^n}fN=^HoiN*{6_^*=5&Gl!(9u7KhHXGg| z>{V~@8z{VdVDD~0WMaKjtY2l__}W891!d#1N~s1$>aw4pS#G91BB{%6^lcd{nz8Za z>pKgJ_QVr$ykTNafhs#~ObKn;EAsOb3LP9c9f6;W7h7DB@fh-cRog8nlO23?uzmbd z^$pAN)cF+yZR5Go39jMXf+v_qe2vdyx2Za85qm!L&K=jXXsUl8z2 z_AN>|?Jl>hV6|R^(b!e);k)Loozk(#EzVxd@HiWF(`NbpF?&wOWTJm~U#z%wOHn_! zIwb#4(_L&x{SW*H4SsB_yq~Hae%jJkkuCVi*Wi)txOd4Gp|nYeO14{Qy(p43Y-!MI zmRo$&AX1V* zLdbV7(0R))XaByl|2yjyUk5_6*0Y{<*XzEnYQB}huXbHWk=$!9nKii1bwjVRf#(JS*p|#8E@Jq{O`JDgWK8U6o0(!`F*|qcSr-Q zMnFGP-A=eyn`ft@R!COmb5(#|H7Nd(hb7%y@9)8hKwISIwT2yRK`#K{Bq1fW_s$3; z_=1A{L6_VZHQ+`(Kv`%sdAiI5q9>@8?fPD1Ws^I+9)GUCT3h1d z`l}&mohFgCr;;HWp_3m_J_Lb z+n<3*`RvJNR##7tvF&=PcH8gX_C|^xBc+)pxXSiq$4{yIQx#wlT+9% z3)W2va$ZtJImeOw|5p>KVbxi-I_ys$kZ&W$$)i6$?*Em46CZ$YL5{$`%BR!w?PgQl zIu@3BJIlr%%+cd(SDxtc`mbJgGF+YDH)3RrWO#fkp1M{lAJ5G#$SHR`lEH)3$w-VD z7#-=(aPusauv-Z*eKEcgmO7g@(OKhj2_RorrTks9_MVeZQTjF&a&D>5pZf$`!RYh# zI07D)-MOn&j!vE0EG;KP);B;l5=iE%x&>^P-YCahEoCJe&Bu$STdj0x zlyT1ojp}-|bLyUPYKsWvitph_Xlf$+3C=c?u&C@cQkEh=4zzXC=o;iJ6wVRsV?kU- zYL={sR%^>+cv<{Y3R*?QU}(t$p2@eiKQ`3iAOwQt<=wr_x*hD{N}Yb&bSUBV)N!au&B28!+4}Q6<$yuYDV$LxZ|8)4SLX*ZkPQ?RN1ZE%ul% znGzps(r|h7*X7Q-7TgYTpgl4uF-*(eLyxT(0zWP%^}{MS>WseZ_a+YZ3XAFK#0(tJ zYFn;l&~m38C@jAd`Ti?aPp4J09K;ufcw73kSc{&Qce2e){-fS;DJMn2-i4G{GeE%6 z+1dJbi+`_A63z&-qMAAusLI3)#a#6zbB&Q0fL z4o*p?JWZ1)X?YW`gV43CNxQTHp|0ji`RRX zY;R9@-C@N~nG#;0rk0y)!NF>|brhxx-A;q0w(Z_ym_A=)kw$&rsar@g7tB2p6E(rp zmPNMc#;BuPKIqL<1kk7*0V^Y!td#(Mbc5elmhDW&njm^%+OSOvAat2BbMcuk%D^(k zL~RXd$|7TO8GQ?xzNwCbsdqWgv1u>;xYN6p^ZlpR7m|k}lwVPTlb`MJ2IRK!xYEn0 zm3_2f9V2gn)w)Q*+`OR$w47`absx=omUnK-!wc=?+Je;p4b}XlyEBNaKJ+2d}gf@dXq_CHqK4?3B(40+}a2xwkj0n@L@XMVWs%xW+}>z%w0XjTEMk@m>t z4%70fGW~F@nda;*9KF!`Qwza=_|XZ0(J;-+bMbmJ5jjTBKC+@yS0eM6TIMx6MQUj8 zOXl=3w0*uw)$9Zp$|gMBwyE5XLv}Zt<&J^lSTi0#%_EF72p=-S2_$SjNZD>|XE$LR zhJ$9?B7&kq_{3n{LuIF5BUBBDlIk`-$MySSsWKlmk}D>*yf99fW35n#-W)C?-eD`B z8VKQ<>I1z3Zf)cNXOp4PXVl7iC>;F3Xx+u8eO~TjQ8O}aod~npCLSM=790Y~0)&GM z5WW=yP98a@b-qXh$hg$0#rwi?p4&sC;$F@>a zi*ws`*w~hE+`skPI3Nf=*tx3Mm$w^^C{PpOZ={d}gC{?{1iNLohDFt9^Rz_1K>y{7z5t(@H0_%oB3VV7m49j%7dBugnLbKeHuS z?7>+h!p{IDkU=KK0Tt;pqXc^3SQKdba?iA{uTIc|$dR=hveSl?OsB=w;KhZJmz6pU zD$&FK&8Uj~>TH|%l#q1V_>1uEr}R$ljkLo48e!V+r4y(jM9wNhv8$az0EG0hm%}$!b75c`Xs_(N%NB z8E4&1me;aiHc@Y;Oe>43Bq3C@Snn=u0nqE#va&8ScTcZZYAy5R*4LeDr}oss6rBsh zT-GGEcpawT({orT8<&xc`g7lJs(oEz24_Mo+)=^%G6D!_ zEx(kyuXa|{aL((BYW8Im8u20??a%Gami4fyQ(s1$^^4R9MbX9I;(aqP^+ge(JmFdF zVIMSk>FQZwerl|}Fnsx3i~p$@+%vCF^aZ&DwI(ku^HnkhFT`o=4JLLL~4&(ro3XTVz)Q-g= zV$nS3iMFVxiB&a2uO`a!Mth=4#C&D!UDO&M+aSR>5~ty_;%aIXzP$TPSmDOc@oHVJ zRm|jyj)>s5DaB8+V-=L*Kl-6!)@pZh^Ru0=<@Dr@1N7YxsGlDanf3_ymx`ukIO6I% zIqA6f2t0O>v+B$> z^}i98;6L873FVsCE~(lP8hB-=20tsGwWw8IgSBtV_Rc^8i zy093kYIDgbxo$EXnU3-w)`4JrKk-uOui9!3zM^jy*4E)*+I~Gf{JUyi3a=P-EMK^4 z9225=x_|Tle$I}cAG3PwVx1fQUOvjhi(f-bfIAmWd}IAF#5zn%ptE(7)~-_ZWjSY! zoPPpseCVS_<3M^u=xa?o6_p-E>Ul5RbY1!CUk}TG85ODsS_x%3nC4wt2nHJ6|7Scy z5D+PrNi#pgpa`(`BGa{=F6FB>^d^9mXW`~vJrQlf4O8+^c+ihoXTw80HUP0*K@}wp z?k+Nr1y2eR@geYR?M4{-%E~#Zrf+i_g+t*VW^sS*Rf0vv=&;-u;zq+u)p~k8Bu+Gw z$v)jfFLEd{xO-x=qMU!#u=-it9fVrOl`f6q8o30w>{UJ=w za!;)-(W@o`bz~(k%=Rls9vWLB>UT;k(aq=9h^z6fP>e}agO(=gm0j^P%j^@z5&tl@ zi_6V5u>nr~*VeuBB@VzXi||JIbh%tBhq1t?R5O5=zD6!ez2KtVPBi#V*_!*WKy7{s7+*b6W5dMiYf4HrHub05x?eG$C z;|TB7_StjDi3T3e&#iG7kMgiJjSP`{l|8t012oONA`@#|oAIK#LcP%$nA80fWm0)B zKhC|x(dB*h(Rn)GSxOM2XaQ$|?!51nhFR%AEbe~s(3TfrEKC+v;f)qkZ#LkXfa&s6 ziD55aB~c9@>R{}t`U|3*SYl?S?VWVuRP4?*5|uBGm=XZe&xWW{{!Y&qx}u=$U2 za&E68!&C<}Zgzk0-j))8{P^&%{5E)!V&cc@BnQR{nX2w;l!2LzVr2J=WF{2NV=sdU zoUDcQ8VtAh84x)-m4b9 zWCC%>Bq8#Z%@p00Fijy+$Md$t4paa7qod^H229C_QgT5BI#OcDQh|@#X0rlOopJMS zBo8@n(KR$Ax!BZ2LzViX94ycYwNgg(xICs!H-3}FwUIRrTz-IM!0He_^+07)%0AT| z;mK31m3Jt9NEOHWUHE4;NQ|SHw{OFM;^Z)OBUwC;?Wu3GgQ`^XfYPdzQ@*gOor#00 z*U`8v%w7l(+tg=GMvChdzfVrT>6aqnW}?gj&eqnx3u+-bYMED5lPJb%d{{o08IyZW zn=GzO+d}GzP3H&?UdytJV`oVU@>Vd($3SVU669^5D3V4*2%91bH{5c^@oCI?11j9KK#xP%wU?*Xu-Oghp)a?FrMU4jrqaX&`zS2o8ut8F_RDIG*XZ0S58J z)~Hg2>1cC$uf6Yl3~H)#_nCH=TeK^~%N7_HR_lcHNbmKL!MSH}R%Ou}U7~e44uoaR zvHJ|nlb9pucEsT$x-7tLhCs|$0cp~T3{%7WA;$9FgX`%j{`y`O)Kr?y7@kZXHjkdc z_VkEQbv@o*SD22SQ0j=E#6E5cS$msNE_Mw*>+tycjedLVe|YGNOu= z__&q`2Jj~BqbMoFuQ9kRldLbj)<{>>l;s0zgy9;?FC$PhN-zj@7T}%8Y$w(%DG?%ysG^_bGOSvfo=(eg8f1LQ zGtqVXI+gk2EGC<7%B*9l4}i2OU+B5x`V5{fi!T>G-O4w4xjym&)>!-f_Zn(@C$J16 z+!jCpU<|VK|;Q3zV~1-TQQDE48X5H?7p^Bx&1$uyI?ui+O8bo9RoG{o10k zH}T%S)Z&zsHjJl)q;*SfuZIhH>605>4k?KojzZckzy;I(?txk29kK|38Xy@%_`eZ& zB3uyuwN8o@VdURouL;d)1e5b0VXx}{4eV8RwBe~57C=4n(B6Rws{ft}q|>%>X=G$3 z&n2MOA}ufVlu_ZSlwn3@#PARVkQU&G&RN;_eOq~Go?}a8_+%*P;tHQL! zSnd{WCSkI)Y)Sl8b+rPd3)MV0e%tkK@o@x{fWE4$Ta1?i)C>%M#c{svnKCV1(}28&``>)a6SkytXN zap_XME$xI^BGM=6g07IQdJw@=c9esUo@KSsBn2D?ZO2YnpIwFKQ#%dYZPS#XmW|h0n_HNqVa}%M^U*~NkgYejVO@SMfl(@c^dS-BF47ug!8xWC zXMxKP{yIrFiLTWH5up}>_{jo zE4vW9neG0Ke*?t6n5Frci?0x%9 z8)zy(>Eq0L#)NQI1ET=pCQGE;T|rnL*e`be%+02(+W88hAyFx`>F(Ug4X(W4{z(y} zpv`uk{C{A#c(8MQ#$e~Wk4PGR`U=9*J^}2rv%}8z_HbDqT{=S^2%r_x7%zVG{}4GV z3&sgzQVFO;7jWCo$SKyondZ*S5*BhlDOWWfMP04S*?v{Z`*O1Kra-4MNC*J!$F49f zonunzK6RP#f5O#nAdiENuD^X#v#PkQUZi#O@9l#f!xRru+QcrldeV1jWcc;f#e$GghG_x+gl5-D6k@wqytEsAl?BwBdcwP?%^@L_0so3r${ zd81xRLV|+a%eUEPlD^Lsx)wWmX7YrajL<7u%2r0R+a?mC6F&?d+@;%FA}eIVsjFk0 zFywGzY%DXs?Ffz+fNZk?sM}Nxv>9`ppnf#Ql<`G`*8M4Fu<+Ol<(~DI$9c1va{MP| z;;R^MQ{UOeL6|MPQN!OIlFknk?OFycLvLzsU+>Mn^))KoY&ucX-_5Vc6Hl)U9k*U# zn&nr@h4ZoCG^be0o33CZD91CuA?OE2v|_c6nS+Y&c)sz=mFz>ub4uV zmu6g{+5Fn*$JDfHJElGjrDot2IZ`qa#i&@s!w2&qVaE9Zt+~ROP`4u6!CC zyC3Xq_y+0{kdz}ge5p1%;JGlv%+R}Q4R11zB1)2IS{QN^<7PHFw!$XK8EgVn)o!kC@q9@RZ6cUaB%QD#QVCnO;N`ksVii`w7#}y3NY9Z+VJSM1QH_k?#Zo z5IiWd2aQ$Azg7acAcp#k5v{8OPK}smruMTFUW-*#vwBxKyu8|(Nv$0LCXcppemkxX zTFebFZSrduk(N>tOJog-X?=Cr?MktOWo`?jv16@#;SZ#pD#Cgs(Xu@*(P1;yRnvvg zl@amm?eD21y6rt)e@Lug^v!&~%n&g=y-M-GWk+v;JZa5AhZGoiJ_4|@_-wD~;XKjo znc8F0e~_$5&5h*9knDbU0xHu*GlB5Aw_z>ZS1jLA)!ejS?D%H3_O=Ztf1U+ofjU^> z)%>Tl+&d?|zwx1&;Sm4#x%2V~@JEhfqw@N^2q@4*+bdmHgJ{}oi zIF_cB5CLrwlQYlWos5%^L`t+UeplBnq>=49f%P6D@D4A&aGGJ{<{DaEWw$5K#OjkC{^V;r@=4A`SII*o+wCu9 zX%CtQ-8tB<;ms8D(1ff~Tr?Dq-H6Ww@jP*DzUtT9lPylW_X1LOTZVywVL zfIkgMNchY79imb3vv{H6a0R8(DSs-Hke3?WD5hrIlPjYWtL__6G`XEb&=2K)cAaha z2@NNa3Xd;!u-g3%MrH{{Ox0H3Zw!F8$f6g@1h*5(`+U4=-@i#N&Bf1Z<%}pYKi6{lDep@taL3yxO#$!P7ohzQD~ly=Pa;l&>MR6jsq` zkCulEuT*)2)AON(hE5O2ahOUj>&6d6W0m}+VqAR|I&#ej=&LfK{VTI(#Px~#c74UQ zIqQy}%!rk^@U0f;^qLcyLwrUI)rHkke8-6UtXQ_@-sN{mdv;l>P;iEataU~gQ24&4 zfhx=`vcu-;c{@_T_}cRO#j|PY4o52D5`B-c6v@$3>UP+%`@UR}y-#DXN+JQtIOdzC z=lxD&%l?mOcR8^9X4xnkb&rK84fP9dP367OwTsD=;>wqNU9oG_9bz==-UE_u>ZUANJ|gya}#41_l3mLGLNxJLV`N>rAJdu=1E#8 zR8H8f-=~i4=@r=LcW%m2ErHpV*DDKMQO-QIQ)LB;7wvo(nq~p*dy4xxm$%hfdf?hq zZp*`!>&8~xQ|UK-8xv`(9y7IZ@D3=vMq3N=)&IKX2SS~?mCm8$?TA^4XPk&MSt%LK zWe9``tNPO%i^@bQ=w6&;r907Mm|({M5V1x2^>8`=gh#vi=EWVqWWGMx++1`e*6i$; zzSMVmWO!wgy&XytcVbsf1UR&gqTe!OpQ+z*Tc~8kg-q4yzJKn^kMl!pj|@2Hnc0JV>h-9G;87H77I-`Zn6qFq3nuKdC8ecZ^hCwPT6qg0 zLxbrwnD|yUM|iZdvUYMe7yxMpNfapi4-%mkh#ZH9A1|VX3zQw~?eV~eri-+A;6+|5 z>t^@4=st0W9^s$9&%wTZ1`*HUqP6AeT@UtN&n|H}ZPk#{%#oh6!_HI2eRT=Bto5V7 zsGf$3^8fT{I4YeUX`g%Ax9+@ygVUT(X(VuKyz2ZkoDdVhOpZ8&-z^Ap>@xGZybzIv zy*`zd7AR|MvPjBy8NB(3bK}hdldlvtpDvx8QX(d5j#R)5l~8Ulx54Dkk6W=@+lb9B zjEl*2hYeKy5|Q*94rSB-6xt$J03b&Mxbj%U9Aex=n~boWRB$}CC*oj#fRDG$rKTssw#X&Y z)17Chh3$pI1oaMfq__+wX2HOgU*>Ju}AE| zYks9Pui&NUz?Obh#7}5W5(%? z{3Z8u@9`PSTTR(l&0|;Vw8CxnJP5UIrp7#59YsF~-OOgW47+}}4AP%~x)IWf+p}j+ z1+Y7V)~$uVH-b>*`cwu8`t_0j^Sr`F@qVb8P^+{gxP{B-(F>1PDe)iOEK3a)=8&o-8zjEjp;G4pc^9H57 z%&b#2{AJL)uWW`0D90I7Uj9x`)RbqgDhPnyw;O@;2A~ARR|Z_LrA}xY zqmGUmkH&z-8p^&t3u(cXx5;RXTdy)BQUmXjh_+SN&>YC$7cVKIdeowC;P~#EY5NDE zfr{6<3gnB&Sw3;0n8GjJeso%|jnQK4TxC>f*{V1374jutGE-sdI?1@&q1Df)JS_`d z#KX2Qpk_7e1P^+08`6<2DkII+x?a?%trk3U`0*5JH#onSPHoE6nv6XTLO^H7vyb;CFoH-@taHWQ(ZMc#VKz$bfSNEs>=K(97VVzyW z*sP8P39k~Tm83m>#KJ@{;9L@(y)P?h^T-^5Q5Z|MIgYeQisotFe}MWi6UW^nHeb0H zQtF*5u(>_gV}%`U{IA9|6Lf$U6U-iECTh;;7=FRNM>FIB5L)a}p zJHpbU`a!d)TY71^sUA3;52o3yyq$9yoI7Vi2p5{ugvYRRpX%P&BFesgucOZaA>Je5 z<;1$<>`MJ))|@doKh^+rpbFM@f5=zUBmGEtoGuo-x_F$7=gU{D)o+w8D^2Ai6{}=5 zpTaduaz-SYFRB`pO4w#J#X5y{djJMo`wrGDCGtdVFJW7{K^T` zK^wwwukT*q1;z5N-e!LMC5<0n=!>c~sr^O$GX3MW`=1F~YXVxi{xZIyu@SYyF5^-Z z9fPZ};xt9kVjh+UcC3tp)yIdoiI1Y@nxs1_ub*1-??t;#=e&5tmK4y{)s1x)oe3|>7&PeU7Y>b?v{9YqR^o)3hb!d)r)X?ug*+O&iKUBMu!Q5 z6un+jT`)hd3^ZTkJ+&9+*l|3=<(??r%U&%8g>%P!^}UOlH3{5Ojy=05%?nJsIW=#z zk!ZsFuwIlCgAmC(MCyx(iWslV@Q>5aH4D%JTA~x0GN-+&mL}#rQ%+@-uW$TBU4tZ2#trH*N<7leXxUNzR)P065obT2w=Zz%^nh`6+JVfeUm*((e zI{2n^YwEqRLVh!bP;;eiTHp;mxQz-b9%fr0a;R9#9PuRkcBW!<%`0ae_EjKOC}!?*j)L=<)8R)QS#ooK~cyv%H8Q3mvGy8^$&mLNE;-y z!#y3!`y*K^y*sqodwW&23X`dp*a}~oSgwN3u9xzZq$>ZJ5)>ZX0y0=5fUddS6(i&K z4yv0j+a#6B%NPSQT>&)VlSJpBK{qgi~Za^r}D+uc`+vvOa~<6E}E#Etr3uMH(p4%kF(7W{OfduPln%}wiMs4 z-txTJn!U8cl7+oe`G^UZf%AP?%T=BXHr&*wk5OsocVy@K9U^G1w_{HWQzn$D9?pKA znqDg*Lt$6+jq0h}E4T8>RFS&l9FzuqnXM@}+oyrt_>zI!-h_zzXI+S2k-;$&I~9S@ z!;w%w5#k+M>{HL?mbRSmT)!OAjXDzZ^L(TjQLEYFN{8f=_l$I2qt6eM?qFPN%g2zJ znqCazLaTeUGLOWm?Ul1SR)^f8ha2svg1b)M9}-ZEuo&$2VWt@S!9z%(6QJ{7QUeKF zfm57$whlw?J9i8tEg`XFJm<1LHzRnYU=8WkvORu|H2Expk!gDvrr*2lEufVkZ;9#c z?FF(8&$o2&Nf1Y0ULUzb{t1IP>%E1(SQB{)DetpTvM8&R2O@t2&^y!H8VFS-A5LB>NZfUm^e!gY55*RI|RSBaW?a3=; zj)_0vl@jzOPoCUH1pm>6=|zG$q&tQTx?1?uZ8t+al=92<_M~O7o?|gO2S$!EIbS@= zq<1ghZ9ZYWU9a5JBT6Dy7!nK#`}dQm4P&rzhYOTXrT3$-6meDLdJ4QZwi>rV;YWez zIf*(*j2hS?v{Cy;UcuIGxc-v$fB)aj)xeBOAh*hYbf5yrYWRGgUtzN(T=7gEfiWk9 z&4=AT0Uv&1-%S=k7=aUvyWMU5Jm4=v!`aI}{{8D7acIUvr|_zUh0-yM%bSZG=f%YJ zballoF-Uz7IS9a91)T%1gwd=cA^cdjpnH1`tkO*Y%$85Qd%7L}XybLewcv6sAKE)Q zaz$*Z{rkh&W4Ay{9*(!)^!{bj{o5GTp6Dw_CvAe@T(;%7Y`kNT80`Pk?-6juRpSUd)_8+U-d&vbGFSqrb7tS`{;G{Sj0 zAq3v0K2vF#`T1>#z!k7VN+Jp5UWKK0L`nF;-tS;#w!q%QWq480;xt0IaJfJR77_oxc#gKEVYa^wwc_vo(5j)Z352{ce0Kt?o>(gL(b2(vk=I<)a$*(QN}(%!cb=G~$e7Os28 zA&o{h#|PJCiV^Y`gry+8m$}ie72 z2>|Dtn!{89($+$#OlufJB9dKWNb3qg zf&w{XaXVDzH;!*d>pX3N%M=uffd3fS!zxnSlfVCrt$jQE@E?pZVe9sF{fGYx>@kPG z0y=APnRc$m}!5#D=|SO1N%PtX8|l)KpTh( zAROFP|MpvYz?+Z22T^#0-xvs*ihM^9Vi`c+oI{uvktvW1&@=d_lhdIMd080b8b<>B zLnxTX=Vllap*w+y8JNO0CIONxzc-vIlPO^W=)H0?1;|sJDCaeoRKI@T9~c0eX1~$o zF;6{Ul2wJ05+YErULZL=rT}{k2q+P46d=xdb{kQlcCcU=E7B_3h_I0+3t5PO7J>oo zyB_RxTQCJKv!Iub&~Foo#C}i<^_>wC5>kW07UKSaW`r~kKR_4ZrNGeuX@DA(jNMaH zQml%d1_*>sg2_Q>g#dlgkYk+3yZ|x|)YRO(cy+X?#EZ(q(AI_M4Befw*c6*uJrQmH zAyZgOpi991j_8!5ko!3k1F7FKVMJT3`dKTr=YmS{C4QTOf{AApoSaCtWe=(ne2OKN z_#P9D-MRM`{8qg;JlR9RshK=@gC0LA%Vl}VhdU4F8T5|QWfE7j8ggY-k?Uk$-Q8}& zNE+wSCfc9X_0zIkhc9_?3taLQ$7OBRZ5PnSW>Q#gwXHJ?|s@v9$ z!5A1Xb)f=^p7Sh$U~4973JyC?SDt!LDG|g} ztUZ}Ut}^EaZsRIRWYgX+Z~x%l8#-JwvCCVb+Q1@@_1@#0wrTKvtCMD0a2>rK=T1MS z$f`dk8=KmlG}VGGCV3D}#K~VtyQrDX2dK7palhMZGnDc3L`$6E=PQ<-eu-Ci|I?SF zUw+5PD%YRy+_*1%so0zR%`{m1PM-44>x-F+uy4=*waz;RIEp^fnn&@Y(4+b&Udnou zPy=Q0Kn9B81^1~M(^wensW6dKW}UqUP>LKS=S3=%+| zPhJn_0*_AvBWeo>Nk$+`0|2Q+pP5=>lnE9EeDF=jRMsEw2Yef-l-F>1t1Tm|8rU&} z&l@OWDrsi7>kr1rc{vXQE9?Zw$cX7ZgtuoqJ^7uo4S8zXaitEunB;L;r3g z5kEfz?)kyq1ITZ0d%6K(Tem%FVi-Y03dS7x7Z65egr5vSTO(Mf2Yc5M4BP&hG3)Yw zS|x=`{+FS@|7sZUU%!8$)(2EQ7F}3WmvwG5APR)~C<-1%6ZttrNDUFCgZSGP^m@G~ zXHm8rx+0)4LN+~Hzarh_C_&l z;43$g)O>6RohavON)Mtwf{^|q0>No3P(fs>%?g$7R5E^#m z^@!|~3soM`hC%M69pIu6E^CB4d=>PFMrO_c*>*DnE)V$)WUOs6={~}Qkq*WP*YL5u zI5*%vu>6g~4F3AnLS$kG-W_oo>lzBkkOSp~(@KpoPzSYwcSDG`p(Ggt5=02OBw%`6 z!y6q*rJDTl-6P$>$dDT@XH~g9jJ>BuL~YKcH`%Br+znTa?bu$`~9Be3l1E_Cp|vGm;WN3 z*LmE2%WQycVn>>&-bJHJn{*$egg5PsI(*^PmJb&e)|g+!bQ@y+{P}`FxIFK~SYlf> z*;mQUyX}t&4@Gf4w+d4!*Rd;{XRRMTc<{jJ=FQi+#-&Xz8pg)44}yYX>%J5KQQ0S8 z@wIX;C!qRg`iq%8GlcTC8&*~tSdXcW`i3xGm)o~*3*bQq)dk)nn*MG7mMCGPuT#zM z-(TdEc7D0%$l3R?Qckz9IWVFY6(2vVtCILzJ817%?+va)8TtA7Jz#vweyu^Q=EM6? zXY7I6xEzFljjmmLfi6Pk56u}61rGoyPaRt^b@?BkBa{!?wfdBgSt0RzwM!_b< z&P-?Lmkkb@CdN)H8nS(6+zH#eg-2$Kw)FF2w$d=JpI$3i+uS^PZRXxl3T(7t;6Ap3 zd6ZBl(4k6ZbQ(=PnSdzX5-Thq&~R2wCv%-9d5XxsOabk{qVRr2OS+90X%T^ zeL14!zxrmP{x*|NQJk0_sf9snn6bFi0I+9yu~1mpT*x!Nb2SGKCd?$d=76PYa z9Et&!Phph!d%%SY_qdh(2gMVjqHM9q5*S>)`qd+++LvYwdqpN*lThVdY$wmhS9jya z4Pjy7kBB0pzOJdMBe$&kX|z!Qj)cOnHqeM&$4n%MAcwaw!KPob^Kk4`T z^7)1%WE}hUz1K7~Lwm1zFs4Dv5q{BU>1O_-{8GEJ03No=n67U8#7?&5I$!r=-##1< ztPFm+^)+Vf%K6#fgp`&4y1`*;z~M8`CF}kgT?FHi%Z@Ok7=$i@@T|wd4=OzZcp&1K z(uIjT)o|;&pgoI6M?mftT?8jY&$)-J?StT86L)tBu|M|=!U||$;nIXoomvR#%-E3% zGpX17{Hn}=x;rW+#>jq$a{enRCW!Q&MfYUto~R`>`y!NN35BrZ!^dcyE)Wc#MnmQ{ zF6qx=(&A8YGH&TwEr7|9XLc+%7_X4$^S? zD=)gIfCZM`bm77U;op9HM|%1u#_0O>$bUsQcvz8 zU8Hj%_DfxG^OUQmH4^H(qNtpx}E-R8$k^ zMeYH8=&Yr~+APqLx`>6FyHCF&T^2)`!;!Oca@66&Z*^xy8b0lu>H^JgePiP+aI`)R z*$mT7foksr5k+@$+Bzv8ms_{UPZxZ2Al4@mp0b ztzKZ#%z(WAq>M~O2WCpYtL)92bM=jlLVSG6E>-hEy=7vxt5@@D*D7D7q`)GcM6*bU z67=$(UzrpQfB%|LE}7rZ)1Oj;wKNE2P6s(Tvq9~h1TwFbkj|9XFn|9@%6ag0LBVOt zRHw?+)KsjvT_^|zm@}#6wUnhcntn#UXxjLV1R*x<18)+y+~Ng$g8N_;`1ZKU=4sKQ zvs=Uu;f`50xyY_o8$C;f+tfVk2@0cHAWDcuVL)4jJ|5?CDp!bC$#40O$TnjBwmogT z#Hx;=Rnh4A9El9A{U-t88!PLPHRxQP0CPXP=Ctx(m4U6F1~_>*d=tFuUJ!*)1J*Iz zwtYM~&O?=%U{mOz?KAb$Y<8zKQ4*+iZwu(2qkx za{fqM0Ky^Z-9*XAZcJ_Zsfh){`$35QiE*WS=DWH&3LyS{8m4r&G(lRHI!h=AhMew2 zN_=Fp;r<;=2CX+}59hBO92wklo|EvP_0TD#gTMUp3J6<`z*8k);2OH8NYe?s_Nah# zHA+xg+7KCJ_ZY#@rrzFVm%aVcT!bI31+u}%GN3pxea3Z1k|+2a4|Z5i=^xA4nTkGu?f{9lQ-c#wX@IJTMx0N# zo&k*+oAt8UV$abau;zmN{Fl7Fy&=Ji0yOMYX}sd4WAdPD2hVE{P|anVf~6~i71jbi z=c|s6IUo>_f}A~ADb*`i+JWRO2Sj%58$-E~@Qq(T-zBTFw_AbZXy~E^1LWzy7FWx! zDtgYmppr+(x?U9)VpA4uci+Jtko)l@^ZfH3a04=K;Ctp#nD$}-*>q%LV*cdh&9+@& zf+FA~4eZTU?%+IUv;0H@Bww>M6XY|2fcH}HKHeUfBVGdD0Q!t_;%`4Ar}^Q-6{tsG zj-LA?X@U%#-qA}#8a69^^P`RY5N%*wswq+6JR`8!uauzIgTOq%kLC?nEli!JXn1SD zzL?NxuJ1xN_$n}z%n}%mVId)U%~2x!Xo~=5som;TpBKH?*!~J+pg%A01+#6w|8naN zc~?}@9C)PDR)xDZr-N3Yt89D|9{&k+2>r^)eTB-{u3Qu8{P~4^PmvM-8Z7pg+2Lvv z2st{ySDUSyC3<#J(s||Jz1?6nCBO-afTe^u$xey>>s>4?#~x0pJ51O-8d$w+*Cvk$ z2@6MD%zs6VtbKHH-$G5kU)8F#MKi;eHaB!iP;kC~#A4*lkaJsH43zfhT@d-QKEYT; zSDu+1VD-EkrgZeuVT^yeJ^6kri!!kSn=T?TQ4G%r7XcTfmRa!v{7l!Dxf6GrpKxv7 zA?$kSSiy(qft62kcK;PpFS_F}JS5-_)Uaz{zg8Tz)Ic$D53k~&!zyC`2mEV*MI5Z1 zw)BvD_)*DQW4gsYlv{}-eji&~dq7F(72I43d4%-~&;EOtKuHyxbSCU;J>{U?pO6Qo zy2b?*ZZPxy_I+7UWA#_c3=Fys0@2viFh%+YI|4*R`lhDBa2=ZQ zvfd5ZorVxX?)>SePjGr)DR_=49OHR+wxN)OgdULJXWLa(wjE@9qc^CvIeMSWmY=VI z=l%SgZ<7b7(UE&mo1&k0PL(tb%s|-&;?V;KtYTtf66Ratq~kiLB+3a+@T{3jLcDqC z;6WR#BXB4r%+Jg?lKKi$&_xtdFXLr)0BcPEk35S<00LbISqDT9kdAur3f%ip;Wpt2 zlJPu6|X66*-O+6nXb?rkRZ<|m>-gCwf^z0G-Uk7VC=z>Pe4_^<#m zG!%UZ=O}YSkTXDz{8?XjsHPIg_SLaZkBB#9*V=u2aUJUj3Pp;N-QIa@=k>qwa6=UH z^ZB7mt(&`eTv4#inzSvzElA#KiIwarcd=+otow2rZEghvI`dgDUYm|Cf=8TF#x->h zw`|wDckhHHB)Xs+Mx@CD1OK64ez^gU&9*&9t-;Jm!qQ_M!7BJWG`ZNKjC%9ToPetF z5C?}Hb_UYka`^Sw8m`y(69U$}+cYSUYal1*&w{wWYI;()Yb(dj>mMieyYdj7w1nO+ z$eEGU2KMd<#AxbRB9caStt;eLRH_;*bN+9oV<4tl2!fo+iysEAZ6FeB?oCl zn!Mwd=dfSR3sm@x3=E>QL@>e7j}g25Z*$Z{6GW1Zb81ar`WCZGATF$9Dpg zAhr4NAISf2*>ec@&;NbwWc@{mRjx1vVoO+Q*wlzz5L61LK;*pxS_8jK3T|`;yU=6E6zc0;%TWD4Uu|Gu zkPbp|laNqEH4{5P6*(e2JS8M#H{$&mG;%MvatgE8pbFpN(Q9V-r3MVK4@ApI%-8*OePEc4?H;=fa(AY8)Pzmho7&m`oYJ1AT2|$ zuLEvJca^srd^E(#&2Cj-pB0!(ozp|Oa)_G*z`eD`%aIIg>5&jL9VsA1BqSi80qRom zU`pRvwY-$k-5ab7k9+nK;CSJX&>Is*{s)aA%`>nnnc)7KQAmY|I5uYL&df8z<$&=Q zz(Yiys}jb8WWwQ46hX|XOU+6{MxGzz)uks${t9PcTkS(x;8Me;rKi&sCO}d}1Zp>K zt?(w`#)opC4!M!%(@c@-YaY4?Y=_U&e2dy)6!cnWXVw^GX|yNQ=Zwk2)w$5cgA$1y zu<$30w#7-ep^ShICC7_A+*9e{K&Wubh>1a**4NQd4J=aD`kuTW;^A?|=3p39s68Ak zAay}eDBK#PCTPAnkUl>Aih zmoL+|oP+B*o6%{Hbp%DesJJ*%DF&)KNZVya!SIskukZeZ!do-32w1o8l)=3=vbWDP z+;L$4{_Bv6lz@*B7YEx`3TF(4Hiut8AQlyO?)EDiJ3HUd@bGY}_2NTMZ2&9KEQ4Py z3Vg|F%fJzbC1@;lqrkLYebBq!FAfX}efi*JeaK5NAYI&undmZs9?)4d6{IRAUpau} zI$SG=+bg;(w0wjDf)iF4g?c~HDM9>lK~?qfh!rgO5UM{*Kf9`-nf=L=7;;+B=}1i0 z#OlxEP?=xCwZZ*Y$L`@e*M6Y=Wqwl9TPQqPVG*C=KjB}I--*5``#Ektl{jP^Yy8nK}`CRgfYM0fjUIK0pXQs;HDm45AT4iwGnIY=TOV zgsyX$PW#!<^1+aBC->%_bM{_q?Q{3?hYqg*&1KCC)1lfI{fh0KFuJ#ybHv4xTSn2hrxSTP|8%b5`-`9<{j>_7*U=ID*pY@ z>5kzkBkYJJH%t1CfH6=*zh}C+CLe%QCeW$GCKk~*+Nva>xLk|_yG?2=mDj!<&FF#F zQ)W`pw`aS}o~`*;bOJ7JIIC5g>NI>9T#KS~4;50#|k zGI~%8$VT>pGH?oCC;948Q2jU7d0XkTzu=Mx+eJ_CJvZjPg|o-&D~e&ivXFxBY*7WG z(dyS42y!6f&J@^94g<47%-U@#)YPIl9fm0c0GfmjM*@@JeK4$Lo3Y^~$9P+|e>c(% z=g}0WlJ)}0~#_JP$ z+tuyj6huoRu_3kqbi^{+TD8K?MTlt*EP2Q>gSgt0kWi`ok zAG*_Fz6By{rRcXgfw`+6u_*hDR5a)22<+or?{ClKT~`MuUMv3;Q__ zVaPfOS~?H1YAB??W@wOz%_V@-Ea`v(8dMRqNY#z!dME%jP=pdM3D#atR;{g!Ki$_> zeFpRkd4l9|plO!!I4JBH)F6iNFp31y${7qc=A zoCsFw)_6O)yG5ayXNvOKTVbx*OPIeiSor`hNL^Op%2~iY!AC~(A z2!tU+S8ljsU%Fl`ME!!vf=uWyh9DtD( All values are in seconds. **Success rate is 100% (zero failures)** at every concurrency level. +> +> Note: as concurrency grows, ROCK applies rate limiting on the control plane to protect server-side stability, so the observed latency increases accordingly with higher concurrency. + +![Sandbox Create Latency Distribution](../../../static/img/sandbox_create_latency_distribution.png) + +## 4. Conclusion + +The benchmark from 1000 up to 16000 concurrent sandbox creations on ROCK shows: + +1. **100% success rate**, validating ROCK's reliability under heavy concurrent load. +2. **Small-scale concurrency (≤ 2000) is essentially queue-free**, with P50 below 5s — comfortably suitable for latency-sensitive training and evaluation scenarios. +3. **At large scale (≥ 4000), latency grows roughly linearly with concurrency**, giving predictable, easy-to-capacity-plan behaviour for sandbox pools. +4. **Even at 16000 concurrency, P99 stays under ~60s**, sufficient to support very large agent rollouts and parallel RL steps. + +ROCK is therefore proven capable of supporting **tens of thousands of concurrent sandbox creations**, providing solid environment infrastructure for large-scale agentic RL training. diff --git a/docs/versioned_docs/version-1.8.x/User Guides/scheduler.md b/docs/versioned_docs/version-1.8.x/User Guides/scheduler.md new file mode 100644 index 0000000000..68a63cb427 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/User Guides/scheduler.md @@ -0,0 +1,283 @@ +--- +sidebar_position: 5 +--- + +# Scheduler + +The ROCK scheduler is a periodic task framework embedded in the `admin` service. It dispatches background maintenance tasks (image cleanup, file cleanup, container cleanup, image pre-pull, custom tasks, ...) to every alive Ray worker on a configurable interval, so worker nodes stay healthy without manual intervention. + +This guide covers how to enable the scheduler, configure built-in tasks, write your own task, and inspect runtime status. + +## 1. How It Works + +- The scheduler runs inside the `admin` process as a dedicated daemon thread (`SchedulerThread`) with its own `asyncio` event loop. +- Tasks are scheduled by [APScheduler](https://apscheduler.readthedocs.io/) using fixed intervals (`interval_seconds`). +- For each tick, the scheduler resolves the list of alive Ray workers (cached via `worker_cache_ttl` seconds) and dispatches the task to every worker concurrently (default concurrency: 50). +- Dispatch is done over HTTP through the worker's **rocklet** service: the admin builds a `RemoteSandboxRuntime(host=worker_ip, port=Port.PROXY)` (see `rock.deployments.constants.Port`) and calls `runtime.execute / read_file / write_file` against it. **Every worker must therefore have the `rocklet` server running and reachable on `Port.PROXY`** — otherwise the scheduler cannot push commands or read/write status files on that worker. +- Each task subclasses `rock.admin.scheduler.task_base.BaseTask` and must implement `run_action(runtime: RemoteSandboxRuntime)` — the work performed on a single worker. +- Per-worker execution status is persisted to the worker filesystem under `ROCK_SCHEDULER_STATUS_DIR` (default `/data/scheduler_status`), and an aggregated execution report is written to `/_run_report.json` after every run. +- If a Nacos config provider is enabled, the scheduler subscribes to config changes and applies a diff: only tasks whose hash changed are re-installed; removed tasks are cleaned up from all workers. + +### Prerequisites + +Before enabling the scheduler, make sure each Ray worker meets the following requirements: + +| Requirement | Why | +|-------------|-----| +| `rocklet` process is running on the worker | The scheduler dispatches every task through the rocklet HTTP API; without it, `runtime.execute` calls time out. | +| Rocklet's listening port is reachable from the admin | The scheduler uses `Port.PROXY` (defined in `rock.deployments.constants.Port`) as the dispatch target. Make sure no firewall / security group blocks it. | +| `ROCK_SCHEDULER_STATUS_DIR` is writable inside the worker | Tasks read and write `_status.json` here for idempotency / PID tracking. | +| Tools required by the task are available on the worker | e.g. `docker` for the cleanup / pull tasks, `curl` and outbound network for `ImageCleanupTask` to install `docuum` on first run. | + +The rocklet server is started automatically by the standard worker bootstrap scripts (`docker_run.sh`, `docker_run_with_uv.sh`, `docker_run_with_pip.sh`) — typically `rocklet --port `. If you bring up workers with a custom entrypoint, ensure the equivalent command is invoked. See the [Configuration](./configuration.md) guide for the runtime-environment options that govern how rocklet is started. + +### Idempotency + +Each task declares an idempotency mode that affects how it is re-run: + +| Mode | Behavior | +|------|----------| +| `IDEMPOTENT` | Always run on every tick. Safe to repeat (e.g. `docker pull`, `find -exec rm`). | +| `NON_IDEMPOTENT` | Spawns a background daemon (e.g. `docuum`). The scheduler reads the previous status file, checks whether the recorded PID is still alive, and skips re-launch if the daemon is still running. On task removal the daemon is killed via `pkill`. | + +## 2. Enabling the Scheduler + +The scheduler is configured under the top-level `scheduler:` key of the ROCK admin YAML (e.g. `rock-conf/rock-local.yml`, `rock-conf/rock-dev.yml`). + +```yaml +scheduler: + enabled: true # Master switch + worker_cache_ttl: 43200 # Worker IP cache TTL in seconds + tasks: + # ... task list, see below +``` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `enabled` | bool | `false` | Master switch. When `false`, all tasks are removed and no new ticks fire. | +| `worker_cache_ttl` | int | `3600` | Seconds the alive-worker IP list is cached before refreshing from `ray.nodes()`. | +| `tasks` | list | `[]` | List of `TaskConfig` entries (see [Section 4](#4-task-config-schema)). | + +### Related Environment Variables + +| Variable | Default | Purpose | +|----------|---------|---------| +| `ROCK_SCHEDULER_STATUS_DIR` | `/data/scheduler_status` | Directory on workers where per-task status JSON and run reports are written. | +| `ROCK_LOGGING_PATH` | (unset) | When set, scheduler-spawned daemons (docuum, container_cleanup, image_pull) redirect their stdout/stderr to `/.log`. | +| `ROCK_DOCUUM_INSTALL_URL` | `https://raw.githubusercontent.com/stepchowfun/docuum/main/install.sh` | Install script URL for `docuum`, fetched on demand by `ImageCleanupTask`. | + +## 3. Built-in Tasks + +ROCK ships with four built-in tasks under `rock.admin.scheduler.tasks`. Each task is registered by setting `task_class` to its fully qualified class path. + +### 3.1 ImageCleanupTask + +Runs [`docuum`](https://github.com/stepchowfun/docuum) on every worker to evict the least-recently-used Docker images once disk usage crosses a threshold. **Non-idempotent** — `docuum` runs as a long-lived daemon; the scheduler tracks its PID and skips re-launch while the daemon is alive. + +```yaml +- task_class: rock.admin.scheduler.tasks.image_cleanup_task.ImageCleanupTask + enabled: true + interval_seconds: 43200 # Re-check daemon every 12 hours + params: + disk_threshold: "70%" # Trigger eviction when disk usage exceeds 70% + image_whitelist: # Glob patterns matching repository:tag — never evicted + - "python:3.11" + - "my-registry.example.com/base/*" +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `disk_threshold` | str | `"1T"` | Disk usage threshold passed to `docuum --threshold`. Accepts size (`100G`, `1T`) or percentage (`70%`). | +| `image_whitelist` | list[str] | `[]` | Glob patterns forwarded to `docuum --keep`. | + +### 3.2 FileCleanupTask + +Walks each configured directory and removes files that are either older than `max_age_mins` or larger than `max_file_size`, then prunes empty subdirectories. **Idempotent**. + +```yaml +- task_class: rock.admin.scheduler.tasks.file_cleanup_task.FileCleanupTask + enabled: true + interval_seconds: 86400 # Run daily + params: + target_dirs: + # Plain string form — no exclusions + - "/data/service_status" + # Object form — per-directory exclusions + - path: "/data/logs" + exclude_files: # Plain name | relative path | absolute path + - "docuum.log" + - "./rocklet.log" + - "./access.log" + exclude_dirs: + - ".cache" + max_age_mins: 10080 # 7 days; older files are removed + max_file_size: "1G" # Files larger than this are removed (supports K/M/G/T) +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `target_dirs` | list | `[]` | Each entry is either a string (path only) or `{path, exclude_files, exclude_dirs}`. | +| `max_age_mins` | int | `10080` | Files whose mtime is older than this many minutes are deleted. | +| `max_file_size` | str | `"1G"` | Files larger than this are deleted. Suffixes `K/M/G/T` accepted. | + +The deletion condition is `(-mmin +max_age_mins) OR (-size +max_file_size)`. After file removal, a second `find -depth -type d -empty -delete` pass removes empty directories left behind (also honoring `exclude_dirs`). + +### 3.3 ContainerCleanupTask + +Removes stopped Docker containers older than a configurable age. Helps prevent the worker's container list from growing unbounded between sandbox runs. **Idempotent**. + +```yaml +- task_class: rock.admin.scheduler.tasks.container_cleanup_task.ContainerCleanupTask + enabled: true + interval_seconds: 86400 + params: + max_age_hours: 72 # Remove exited containers older than 72 hours +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `max_age_hours` | int | `24` | Maximum age (hours since `FinishedAt`) for kept exited containers. Older ones are `docker rm`'d. | + +The task also removes any container in the `created` state (never started) on every run. + +### 3.4 ImagePullTask + +Pre-pulls a list of Docker images on every worker, optionally logging in to private registries first. Reduces sandbox cold-start latency. **Idempotent** (`docker pull` is a no-op when the image is already up-to-date). + +```yaml +- task_class: rock.admin.scheduler.tasks.image_pull_task.ImagePullTask + enabled: true + interval_seconds: 21600 # Refresh every 6 hours + params: + images: + # Plain string form — public image, no auth + - "python:3.11" + # Object form — private image with registry login + - image: "my-registry.example.com/chatos/python:313" + registry_username: "myuser" + registry_password: "bXlwYXNzd29yZA==" # base64-encoded +``` + +| Param | Type | Default | Description | +|-------|------|---------|-------------| +| `images` | list | `[]` | Each entry is either an image string or `{image, registry_username, registry_password}`. | + +`registry_password` must be base64-encoded; the worker decodes it and pipes it to `docker login --password-stdin`. The registry host is parsed from the image name, so each image can target a different registry. + +## 4. Task Config Schema + +Every entry under `scheduler.tasks` is loaded as a `rock.config.TaskConfig`: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `task_class` | str | `""` | Fully qualified Python class path. Required. | +| `enabled` | bool | `true` | Disabled tasks are skipped at install time and torn down on reload. | +| `interval_seconds` | int | `3600` | APScheduler `interval` in seconds. | +| `params` | dict | `{}` | Task-specific kwargs forwarded to `from_config()`. | + +A change in any field of an existing task entry causes the scheduler to uninstall the old task (cleaning up its worker-side state when non-idempotent) and install the new one — without restarting the admin process. + +## 5. Writing a Custom Task + +Any class under your Python path that subclasses `BaseTask` can be registered. The minimum contract is: + +```python +# my_pkg/my_tasks/disk_report_task.py +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + + +class DiskReportTask(BaseTask): + """Log `df -h` output from every worker.""" + + def __init__(self, interval_seconds: int = 3600, mount_point: str = "/"): + super().__init__( + type="disk_report", # Used as job id and status filename prefix + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + self.mount_point = mount_point + + @classmethod + def from_config(cls, task_config) -> "DiskReportTask": + return cls( + interval_seconds=task_config.interval_seconds, + mount_point=task_config.params.get("mount_point", "/"), + ) + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + result = await runtime.execute( + Command(command=f"df -h {self.mount_point}", shell=True), + ) + return { + "status": TaskStatusEnum.SUCCESS, + "exit_code": result.exit_code, + "stdout": result.stdout, + } +``` + +Then register it from YAML: + +```yaml +scheduler: + enabled: true + tasks: + - task_class: my_pkg.my_tasks.disk_report_task.DiskReportTask + enabled: true + interval_seconds: 600 + params: + mount_point: "/data" +``` + +### Authoring Checklist + +- **Always set a unique `type` string** in `super().__init__()`. It is used as the APScheduler job id, the status filename (`_status.json`), and the run-report filename (`_run_report.json`). Two tasks must not share a `type`. +- **Pick the right `IdempotencyType`**: + - Use `IDEMPOTENT` when `run_action` finishes synchronously and is safe to re-execute. + - Use `NON_IDEMPOTENT` when you `nohup` a long-running daemon and return its PID. The scheduler will then track the PID, skip re-launch while it is alive, and `pkill` it on uninstall. +- **Return a dict from `run_action`**. Recommended keys: + - `status` — a `TaskStatusEnum` value, persisted into the status file. + - `pid` — required for `NON_IDEMPOTENT` daemons (use `rock.utils.system.extract_nohup_pid` on the `nohup ... & echo PID_PREFIX${{!}}PID_SUFFIX` output). + - Any other diagnostic fields are written to the status file's `extra` section. +- **Override `from_config(cls, task_config)`** to translate `task_config.params` into your `__init__` kwargs. +- **Use `runtime.execute / read_file / write_file`** rather than running shell commands locally — the scheduler is dispatching to remote workers via `RemoteSandboxRuntime`. + +## 6. Observability + +For each task, the scheduler writes two artifacts on every worker: + +| Path | Written By | Contents | +|------|------------|----------| +| `/_status.json` | `BaseTask.save_task_status` | Latest per-worker status: `task_name`, `worker_ip`, `pid`, `status` (`pending`/`running`/`success`/`failed`), `last_run`, `error`, plus task-specific `extra` fields. | +| `/_run_report.json` | `BaseTask.run` (admin side, after the tick completes) | Aggregated report: total/success/failed counts, list of `success_ips`, and `failed_details` (`ip` + traceback). | + +Scheduler-internal logs are written under the standard ROCK admin log path, with `name="scheduler"`, `name="task_base"`, `name="image_clean"`, etc. Scheduler-spawned daemons additionally write their own logs to `/.log` when `ROCK_LOGGING_PATH` is set (e.g. `docuum.log`, `container_cleanup.log`, `image_pull.log`). + +## 7. Dynamic Reload via Nacos (Optional) + +When the admin service is configured with a Nacos provider, the scheduler installs a YAML listener and reacts to config pushes: + +- Only the `scheduler:` section is inspected; other sections are ignored. +- The new section is hashed and compared against the previous one — duplicate notifications are skipped. +- A diff between old and new task lists determines which tasks to install, uninstall, or reinstall (changed `params` / `interval_seconds` / `enabled`). +- Non-idempotent tasks that are removed or re-installed are first cleaned up: their daemon PID is killed and the status file is removed. + +This means task interval changes, parameter tweaks, and adding/removing tasks can be applied without restarting the admin process. + +## 8. Troubleshooting + +| Symptom | Likely Cause | What to Check | +|---------|--------------|----------------| +| `Scheduler disabled, all tasks removed` in admin log | `scheduler.enabled` is `false` | Set `enabled: true` in YAML. | +| `No alive workers found for task ''` | Ray cluster has no live worker nodes | Verify `ray.nodes()` reports alive CPU workers; consider lowering `worker_cache_ttl` if workers were just added. | +| Task ticks fire but every worker shows up in `failed_details` with connection errors | `rocklet` is not running on the workers, or `Port.PROXY` is blocked | On the worker host, hit the rocklet liveness endpoint `GET /is_alive` on `Port.PROXY` (e.g. `curl http://:/is_alive`); if it does not respond, restart rocklet (`rocklet --port `) or open the port in the firewall. | +| Task runs but never repeats on a non-idempotent worker | Recorded PID still alive | Inspect `/_status.json`; if `status: running` and the PID is alive, `should_run` returns `False`. | +| `Failed to create task ''` | `task_class` import failed | Ensure the module is importable inside the admin process (installed in the same venv, on `PYTHONPATH`). | +| `docker login` failing in `ImagePullTask` | `registry_password` not base64-encoded, or wrong registry parsed from image | Re-encode the password with `echo -n '' \| base64`; double-check the image's registry host. | + +## Related Documents + +- [Configuration](./configuration.md) — Environment variables and runtime layout +- [API Documentation](../References/api.md) — Admin HTTP API +- [Python SDK Documentation](../References/Python%20SDK%20References/python_sdk.md) — Programmatic sandbox usage From cc1854fc38de8ae7ebdafc7acfebfe4bff947245 Mon Sep 17 00:00:00 2001 From: ShixinPeng <58160712+BCeZn@users.noreply.github.com> Date: Thu, 28 May 2026 11:42:32 +0800 Subject: [PATCH 139/226] chore(ci): run admin+network tests only on push, skip on PRs (#1040) * chore(ci): skip need_admin_and_network test step Temporarily remove the unstable admin+network test step from CI to unblock the pipeline. Tracked in #1039. Co-Authored-By: Claude Opus 4.7 (1M context) * chore(ci): comment out admin+network step instead of deleting Keep the step in the file as a commented-out block so it's easy to re-enable once the env is stable. Refs #1039. Co-Authored-By: Claude Opus 4.7 (1M context) * chore(ci): run admin+network tests only on push, skip on PRs Re-enable the admin+network test step but gate it on the push event so it only runs after merge into master / release/v1.0, not on every PR. Refs #1039. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .github/workflows/python-ci.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/python-ci.yml b/.github/workflows/python-ci.yml index 2dc4c2d89e..08561fdc3f 100644 --- a/.github/workflows/python-ci.yml +++ b/.github/workflows/python-ci.yml @@ -55,8 +55,9 @@ jobs: echo "🔐 Running admin tests..." uv run pytest -n auto -m "need_admin" --reruns 1 + # Only run after merge into master / release/v1.0 (push event), not on PRs. See #1039 - name: Run tests need admin and network - if: success() + if: success() && github.event_name == 'push' run: | echo "🔐 Running admin and network tests..." uv run pytest -n auto -m "need_admin_and_network" --reruns 1 \ No newline at end of file From cbbfd0b1f450e2ce63bc2baf6187daf500aad938 Mon Sep 17 00:00:00 2001 From: daifangwen Date: Thu, 23 Apr 2026 06:57:00 +0000 Subject: [PATCH 140/226] add tracking config into job config --- .../proposals/job-metrics-reporting-config.md | 287 ++++++++++++++++++ rock/sdk/envhub/__init__.py | 4 +- rock/sdk/envhub/config.py | 25 ++ 3 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 docs/proposals/job-metrics-reporting-config.md diff --git a/docs/proposals/job-metrics-reporting-config.md b/docs/proposals/job-metrics-reporting-config.md new file mode 100644 index 0000000000..0c5387eda1 --- /dev/null +++ b/docs/proposals/job-metrics-reporting-config.md @@ -0,0 +1,287 @@ +## Job 级指标汇报配置方案 — 集成 Harbor ml_tracker + +### 背景 + +ROCK 的 Job 系统需要在 **Bench 评测** 和 **RL 训练** 场景中汇报运行指标。Harbor 框架已内置 `ml_tracker` 模块，可汇报以下关键指标： + +| 类别 | 指标 | +|------|------| +| **Reward** | `reward/*`（verifier 输出的各 reward key） | +| **Duration** | `total_duration_sec`、`agent_duration_sec` | +| **Token** | `input_tokens`、`output_tokens`、`cache_tokens`、`cost_usd` | +| **RL 训练** | `logprobs_mean`、`entropy`、`loss`、`kl_divergence`、`advantage`、`grad_norm`、`clip_fraction`、`value_loss`、`explained_variance` | +| **Running** | `pass_rate`、`avg_reward`、`error_rate` | +| **Summary** | `final_pass_rate`、`final_avg_reward`、`final_error_rate`、`total_trials`、`total_errors`、`total_duration_sec` | + +但当前 ml_tracker 的启用方式依赖**环境变量** `ROCK_API_KEY` 的存在性（硬编码判断），用户无法通过 Job 配置声明式地控制是否启用、传入超参数等。 + +**改动前**（Harbor `job.py`）： + +```python +# 硬编码检查环境变量，无配置入口 +if os.environ.get("ROCK_API_KEY"): + self._tracker = MLTrackerFactory.create(...) +``` + +--- + +### 目标 + +在 `EnvironmentConfig` 上新增 **`tracking`** 字段，让用户在 YAML 的 `environment` 段中声明式地启用 Harbor 内置的 ml_tracker，汇报 Bench/RL 训练指标。 + +**设计原则**： +- **字段名不绑定具体 SDK**：用 `tracking`（而非 `ml_tracker`），避免配置字段与具体包名耦合 +- **复用 Harbor 已有能力**：不另起炉灶，底层仍调用 Harbor `ml_tracker` 模块 +- 所有字段可选，零配置向后兼容（默认不启用，保持现有行为） +- 不侵入 `HarborJobConfig.metrics: list[MetricConfig]`（那是评测结果的聚合策略，语义不同） + +--- + +### 方案（已实现） + +#### 模型定义 + +**ROCK 侧** — `rock/sdk/envhub/config.py`： + +`TrackingConfig` 定义在 `EnvironmentConfig` 同级，作为 `EnvironmentConfig` 的二级字段： + +```python +class TrackingConfig(BaseModel): + """Experiment tracking configuration. + + When present and enabled, activates Harbor's built-in ml_tracker to report + per-trial metrics (reward, duration, token usage, RL training signals) + and a final job-level summary. + """ + + enabled: bool = Field( + default=True, + description="Whether to enable experiment tracking for this job.", + ) + params: dict[str, Any] = Field( + default_factory=dict, + description=( + "User-defined hyperparameters merged into ml_tracker.init(config=...). " + "Combined with auto-collected job metadata (agents, datasets, etc.)." + ), + ) + +class EnvironmentConfig(SandboxConfig): + uploads: list[tuple[str, str]] = Field(default_factory=list) + env: dict[str, str] = Field(default_factory=dict) + oss_mirror: OssMirrorConfig | None = None + tracking: TrackingConfig | None = Field( + default=None, + description="Experiment tracking configuration. None = disabled (default).", + ) +``` + +**Harbor 侧** — `harbor/ml_tracker/config.py`（内部模块名保持 `ml_tracker` 不变）： + +```python +class MLTrackerConfig(BaseModel): + enabled: bool = Field(default=True) + params: dict[str, Any] = Field( + default_factory=dict, + description="User-defined hyperparameters merged into ml_tracker.init(config=...).", + ) +``` + +> **命名决策**： +> - 用户配置字段名 = `tracking`（不绑定具体 SDK，未来可扩展到其他 tracker） +> - 子字段 = `params`（而非 `config`，避免 `tracking.config` 语义重复） +> - Harbor 内部模块目录仍叫 `ml_tracker/`（内部实现，不暴露给用户） + +#### 在配置层次中的位置 + +`tracking` 放在 `EnvironmentConfig` 下作为二级字段，而非 `JobConfig` 的一级字段。原因： + +- **与 `oss_mirror` 同层**：`tracking` 和 `oss_mirror` 都是环境级别的能力配置，放在 environment 下更内聚 +- **Harbor 的 `EnvironmentConfig` 天然包含这类配置**：Harbor YAML 中 environment 段是 tracking 信息的自然归属 +- **简化序列化**：`to_harbor_yaml()` 通过 `to_harbor_environment()` 序列化 environment 时自然携带 tracking + +```python +# JobConfig 不直接暴露 tracking，通过 environment 间接访问 +class JobConfig(BaseModel): + environment: EnvironmentConfig = Field(default_factory=EnvironmentConfig) + job_name: str | None = None + namespace: str | None = None + experiment_id: str | None = None + labels: dict[str, str] = Field(default_factory=dict) + timeout: int = 7200 +``` + +> **默认 `None`**：不写 `tracking` 时行为等价于改动前（不启用）。用户显式写 `environment.tracking: {}` 即可启用。 + +--- + +#### YAML 配置示例 + +**最简启用**（所有默认值，自动采集 agent/dataset 信息）： + +```yaml +experiment_id: exp-rl-001 +job_name: qwen-72b-swe-bench +environment: + tracking: {} +``` + +**记录额外超参数**（RL 训练场景）： + +```yaml +experiment_id: exp-rl-002 +job_name: rl-grpo-run-3 +environment: + tracking: + params: + model: qwen-72b-instruct + algorithm: GRPO + learning_rate: 1.0e-5 + batch_size: 64 + kl_coeff: 0.05 + num_rollouts: 4 +``` + +**显式禁用**（覆盖团队默认配置）： + +```yaml +environment: + tracking: + enabled: false +``` + +**不写 `tracking`**（默认行为，等同于禁用）： + +```yaml +experiment_id: exp-001 +job_name: my-job +environment: {} +# tracking 不出现 → None → 不启用 +``` + +--- + +### 字段说明 + +| 字段路径 | 类型 | 默认值 | 说明 | +|----------|------|--------|------| +| **`environment.tracking`** | `TrackingConfig \| None` | `None` | 开关。`None` = 不启用（向后兼容）；写 `{}` = 启用。 | +| **`environment.tracking.enabled`** | `bool` | `True` | 细粒度开关。配合 `tracking: { enabled: false }` 可显式禁用。 | +| **`environment.tracking.params`** | `dict[str, Any]` | `{}` | 用户自定义超参数，与自动采集的 job metadata 合并后传给 `ml_tracker.init(config=...)`。 | + +两层开关的设计意图： +- `tracking` 不写 / `null` → 不启用（向后兼容，默认路径） +- `tracking: {}` → 启用（`enabled` 默认 `True`） +- `tracking: { enabled: false }` → 显式禁用（团队配置模板中可以预留 `tracking` 段落但暂时关闭） + +--- + +### 汇报的指标详情 + +启用 tracking 后，Harbor 框架会在以下时机自动汇报： + +**每个 Trial 结束时**（`TrialEvent.END` hook）： + +``` +reward/* — verifier 输出的 reward 值（每个 key 单独上报） +total_duration_sec — Trial 总耗时 +agent_duration_sec — Agent 执行耗时 +input_tokens — 输入 token 数 +output_tokens — 输出 token 数 +cache_tokens — 缓存 token 数 +cost_usd — 推理花费（USD） +logprobs_mean — rollout log probabilities 均值（RL） +entropy — 策略熵 = -logprobs_mean（RL） +loss — 训练 loss（RL，来自 agent metadata） +kl_divergence — KL 散度（RL） +advantage — 优势值（RL） +grad_norm — 梯度范数（RL） +clip_fraction — PPO clip fraction（RL） +value_loss — 值函数 loss（RL） +explained_variance — 解释方差（RL） +pass_rate — 截至当前的通过率（running） +avg_reward — 截至当前的平均 reward（running） +error_rate — 截至当前的错误率（running） +``` + +**Job 结束时**（`report_job_summary`）： + +``` +final_pass_rate — 最终通过率 +final_avg_reward — 最终平均 reward +final_error_rate — 最终错误率 +total_trials — 总 trial 数 +total_errors — 总错误数 +total_duration_sec — Job 总耗时 +``` + +--- + +### 与现有体系的关系 + +``` +JobConfig +├── environment: EnvironmentConfig +│ ├── uploads, env, ... ← 已有: 环境级配置 +│ ├── oss_mirror: OssMirrorConfig ← 已有: OSS 镜像配置 +│ └── tracking: TrackingConfig | None ← NEW: 实验追踪配置 +├── labels: dict[str, str] ← 已有: Job 级标签 +└── ... + +HarborJobConfig(JobConfig) +├── environment.tracking (inherited) ← NEW: 通过 environment 继承 +├── metrics: list[MetricConfig] ← 已有: 评测结果聚合方式（sum/mean/max） +└── ... + +BashJobConfig(JobConfig) +├── environment.tracking (inherited) ← NEW: 通过 environment 继承 +└── ... +``` + +**关键区分**： +- **`environment.tracking`**（新增）= "实验追踪：每个 Trial 的 **业务指标怎么记录**"（reward/token/RL signals → ml_tracker SDK） +- **`metrics`**（HarborJobConfig 已有）= "评测聚合：多个 Trial 的结果 **怎么聚合成最终分数**"（mean/sum/max） + +两者语义正交，互不冲突。`tracking` 与 `oss_mirror` 同层，都属于环境级别的能力配置。 + +--- + +### 改动文件清单 + +#### ROCK 侧 + +| 文件 | 改动 | +|------|------| +| `rock/sdk/envhub/config.py` | 新增 `TrackingConfig` 类 + `EnvironmentConfig.tracking` 字段 | + +#### Harbor 侧 + +| 文件 | 改动 | +|------|------| +| `harbor/ml_tracker/config.py` | `config` 字段重命名为 `params` | +| `harbor/ml_tracker/factory.py` | 新增 `tracker_config` 参数，合并用户 `params` 到自动采集的 config | +| `harbor/models/job/config.py` | 新增 `tracking: MLTrackerConfig \| None` 字段 | +| `harbor/job.py` | 从 `self.config.tracking` 读取配置替代 env var 硬编码；传递 `tracker_config` 给 factory | +| `tests/unit/ml_tracker/test_config.py` | `config` → `params` 适配 | +| `tests/unit/ml_tracker/test_factory.py` | 新增 `test_create_merges_user_params` 测试 | +| `tests/unit/ml_tracker/test_job_integration.py` | 重写为测试 `tracking` 字段的配置集成 | + +#### 配置传递链路 + +``` +用户 YAML + → rock HarborJobConfig.environment.tracking (解析 + 校验) + → to_harbor_yaml() → environment 段携带 tracking + → harbor JobConfig.tracking (反序列化) + → harbor Job.__init__ 读取 → MLTrackerFactory.create(tracker_config=...) +``` + +--- + +### 向后兼容性 + +- `EnvironmentConfig.tracking` 默认为 `None`，不写等价于改动前行为（不启用）。 +- `SandboxConfig` 基类不受影响（`tracking` 只加在 `EnvironmentConfig` 层）。 +- `BashJobConfig` / `HarborJobConfig`：通过 `environment` 间接访问，不涉及 `extra="forbid"` 问题。 +- `_HarborJobFields`：environment 中 `tracking` 为 `None` 时被序列化过滤，不出现在 Harbor YAML 中。 +- `ROCK_API_KEY` 环境变量：Harbor `job.py` 中同时检查 `tracking is not None and tracking.enabled` **和** `ROCK_API_KEY`，两个条件都满足才启用。这保证了即使配置启用了 tracking，没有 API key 也不会报错。 diff --git a/rock/sdk/envhub/__init__.py b/rock/sdk/envhub/__init__.py index 115ee11588..5a528d5c34 100644 --- a/rock/sdk/envhub/__init__.py +++ b/rock/sdk/envhub/__init__.py @@ -1,3 +1,3 @@ -from rock.sdk.envhub.config import EnvironmentConfig, OssMirrorConfig +from rock.sdk.envhub.config import EnvironmentConfig, OssMirrorConfig, TrackingConfig -__all__ = ["EnvironmentConfig", "OssMirrorConfig"] +__all__ = ["EnvironmentConfig", "OssMirrorConfig", "TrackingConfig"] diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py index 6aad1688aa..5ed3a8cc10 100644 --- a/rock/sdk/envhub/config.py +++ b/rock/sdk/envhub/config.py @@ -6,6 +6,8 @@ from __future__ import annotations +from typing import Any + from pydantic import BaseModel, Field, model_validator from rock.sdk.sandbox.config import SandboxConfig @@ -70,6 +72,25 @@ def _record_replay_mutually_exclusive(self): "set one (recording mode) or the other (replay mode), not both." ) return self +class TrackingConfig(BaseModel): + """Experiment tracking configuration. + + When present and enabled, activates Harbor's built-in ml_tracker to report + per-trial metrics (reward, duration, token usage, RL training signals) + and a final job-level summary. + """ + + enabled: bool = Field( + default=True, + description="Whether to enable experiment tracking for this job.", + ) + params: dict[str, Any] = Field( + default_factory=dict, + description=( + "User-defined hyperparameters merged into ml_tracker.init(config=...). " + "Combined with auto-collected job metadata (agents, datasets, etc.)." + ), + ) class EnvironmentConfig(SandboxConfig): @@ -85,3 +106,7 @@ class EnvironmentConfig(SandboxConfig): proxy: ProxyConfig | None = None """In-sandbox model-service proxy for OpenAI request record/replay. None (default) means no proxy is started.""" + tracking: TrackingConfig | None = Field( + default=None, + description="Experiment tracking configuration. None = disabled (default).", + ) From 5f92314d9132bae84e67a2d7bd6592a5a7223b49 Mon Sep 17 00:00:00 2001 From: daifangwen Date: Thu, 23 Apr 2026 09:02:52 +0000 Subject: [PATCH 141/226] update docs --- .../proposals/job-metrics-reporting-config.md | 41 ++++++++++++------- 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/docs/proposals/job-metrics-reporting-config.md b/docs/proposals/job-metrics-reporting-config.md index 0c5387eda1..01fd693037 100644 --- a/docs/proposals/job-metrics-reporting-config.md +++ b/docs/proposals/job-metrics-reporting-config.md @@ -76,21 +76,33 @@ class EnvironmentConfig(SandboxConfig): ) ``` -**Harbor 侧** — `harbor/ml_tracker/config.py`（内部模块名保持 `ml_tracker` 不变）： +**Harbor 侧** — `harbor/tracker/config.py`（模块目录已从 `ml_tracker/` 重命名为 `tracker/`）： ```python -class MLTrackerConfig(BaseModel): +class TrackingConfig(BaseModel): enabled: bool = Field(default=True) params: dict[str, Any] = Field( default_factory=dict, - description="User-defined hyperparameters merged into ml_tracker.init(config=...).", + description="User-defined hyperparameters merged into tracker init config.", + ) +``` + +Harbor 侧 `EnvironmentConfig`（`harbor/models/trial/config.py`）中同样作为二级字段： + +```python +class EnvironmentConfig(BaseModel): + ... + tracking: TrackingConfig | None = Field( + default=None, + description="Experiment tracking configuration. None = disabled (default).", ) ``` > **命名决策**： > - 用户配置字段名 = `tracking`（不绑定具体 SDK，未来可扩展到其他 tracker） > - 子字段 = `params`（而非 `config`，避免 `tracking.config` 语义重复） -> - Harbor 内部模块目录仍叫 `ml_tracker/`（内部实现，不暴露给用户） +> - ROCK 和 Harbor 两侧配置类统一命名为 `TrackingConfig` +> - Harbor 内部模块目录从 `ml_tracker/` 重命名为 `tracker/`，实现类 `MLTrackerImpl` 保持不变（因为底层仍使用 `ml_tracker` SDK） #### 在配置层次中的位置 @@ -258,13 +270,13 @@ BashJobConfig(JobConfig) | 文件 | 改动 | |------|------| -| `harbor/ml_tracker/config.py` | `config` 字段重命名为 `params` | -| `harbor/ml_tracker/factory.py` | 新增 `tracker_config` 参数，合并用户 `params` 到自动采集的 config | -| `harbor/models/job/config.py` | 新增 `tracking: MLTrackerConfig \| None` 字段 | -| `harbor/job.py` | 从 `self.config.tracking` 读取配置替代 env var 硬编码；传递 `tracker_config` 给 factory | -| `tests/unit/ml_tracker/test_config.py` | `config` → `params` 适配 | -| `tests/unit/ml_tracker/test_factory.py` | 新增 `test_create_merges_user_params` 测试 | -| `tests/unit/ml_tracker/test_job_integration.py` | 重写为测试 `tracking` 字段的配置集成 | +| `harbor/tracker/config.py` | 新模块，`TrackingConfig`（`enabled` + `params`） | +| `harbor/tracker/base.py` | `BaseMLTracker` → `BaseTracker` | +| `harbor/tracker/tracker.py` | `MLTrackerImpl` 改为继承 `BaseTracker`，逻辑不变 | +| `harbor/tracker/factory.py` | `MLTrackerFactory`，新增 `tracker_config` 参数，合并用户 `params` | +| `harbor/models/trial/config.py` | `EnvironmentConfig` 新增 `tracking: TrackingConfig \| None` 字段 | +| `harbor/job.py` | 从 `self.config.environment.tracking` 读取配置；`tracking.enabled` 控制启用（不再硬编码检查 env var） | +| ~~`harbor/ml_tracker/`~~ | 整个目录重命名为 `harbor/tracker/` | #### 配置传递链路 @@ -272,8 +284,9 @@ BashJobConfig(JobConfig) 用户 YAML → rock HarborJobConfig.environment.tracking (解析 + 校验) → to_harbor_yaml() → environment 段携带 tracking - → harbor JobConfig.tracking (反序列化) - → harbor Job.__init__ 读取 → MLTrackerFactory.create(tracker_config=...) + → harbor EnvironmentConfig.tracking (反序列化) + → harbor Job.__init__ 从 self.config.environment.tracking 读取 + → MLTrackerFactory.create(tracker_config=...) ``` --- @@ -284,4 +297,4 @@ BashJobConfig(JobConfig) - `SandboxConfig` 基类不受影响（`tracking` 只加在 `EnvironmentConfig` 层）。 - `BashJobConfig` / `HarborJobConfig`：通过 `environment` 间接访问，不涉及 `extra="forbid"` 问题。 - `_HarborJobFields`：environment 中 `tracking` 为 `None` 时被序列化过滤，不出现在 Harbor YAML 中。 -- `ROCK_API_KEY` 环境变量：Harbor `job.py` 中同时检查 `tracking is not None and tracking.enabled` **和** `ROCK_API_KEY`，两个条件都满足才启用。这保证了即使配置启用了 tracking，没有 API key 也不会报错。 +- `ROCK_API_KEY` 环境变量：不再用于控制启用逻辑（改由 `tracking.enabled` 控制）。`ROCK_API_KEY` 仅在 `MLTrackerImpl.__init__` 中作为 `ml_tracker.login(key=...)` 的凭证使用，未设置时传 `None`（由 SDK 自行处理鉴权 fallback）。 From ec17769678608949ff67a4d82e042d6d49df2ef4 Mon Sep 17 00:00:00 2001 From: daifangwen Date: Thu, 23 Apr 2026 09:43:14 +0000 Subject: [PATCH 142/226] add test cases --- tests/unit/sdk/job/test_config.py | 110 ++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index 8c678ab337..e8b41e9081 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -20,6 +20,7 @@ VerifierConfig, ) from rock.sdk.envhub import EnvironmentConfig +from rock.sdk.envhub.config import TrackingConfig from rock.sdk.job.config import BashJobConfig, JobConfig # --------------------------------------------------------------------------- @@ -684,3 +685,112 @@ def test_defaults_to_datetime_string(self): def test_explicit_name_preserved(self): assert BashJobConfig(job_name="x").job_name == "x" + + +# --------------------------------------------------------------------------- +# TrackingConfig +# --------------------------------------------------------------------------- + + +class TestTrackingConfig: + def test_default_values(self): + config = TrackingConfig() + assert config.enabled is True + assert config.params == {} + + def test_disabled(self): + config = TrackingConfig(enabled=False) + assert config.enabled is False + assert config.params == {} + + def test_custom_params(self): + config = TrackingConfig(params={"learning_rate": 0.01, "epochs": 10, "model": "qwen-72b"}) + assert config.params["learning_rate"] == 0.01 + assert config.params["epochs"] == 10 + assert config.params["model"] == "qwen-72b" + + def test_from_dict(self): + data = {"enabled": True, "params": {"batch_size": 32}} + config = TrackingConfig.model_validate(data) + assert config.enabled is True + assert config.params["batch_size"] == 32 + + def test_from_dict_minimal(self): + config = TrackingConfig.model_validate({}) + assert config.enabled is True + assert config.params == {} + + def test_serialization_roundtrip(self): + config = TrackingConfig(enabled=True, params={"lr": 0.001, "algo": "GRPO"}) + json_str = config.model_dump_json() + restored = TrackingConfig.model_validate_json(json_str) + assert restored == config + + +# --------------------------------------------------------------------------- +# TrackingConfig on base EnvironmentConfig +# --------------------------------------------------------------------------- + + +class TestTrackingConfigOnBaseEnvironment: + def test_tracking_default_none(self): + env = EnvironmentConfig() + assert env.tracking is None + + def test_tracking_enabled(self): + env = EnvironmentConfig(tracking=TrackingConfig()) + assert env.tracking is not None + assert env.tracking.enabled is True + assert env.tracking.params == {} + + def test_tracking_disabled(self): + env = EnvironmentConfig(tracking=TrackingConfig(enabled=False)) + assert env.tracking is not None + assert env.tracking.enabled is False + + def test_tracking_with_params(self): + env = EnvironmentConfig(tracking=TrackingConfig(params={"model": "qwen-72b", "lr": 1e-5})) + assert env.tracking.params["model"] == "qwen-72b" + assert env.tracking.params["lr"] == 1e-5 + + def test_tracking_from_dict(self): + data = {"tracking": {"enabled": True, "params": {"batch_size": 64}}} + env = EnvironmentConfig.model_validate(data) + assert env.tracking is not None + assert env.tracking.params["batch_size"] == 64 + + def test_tracking_none_from_dict(self): + data = {"tracking": None} + env = EnvironmentConfig.model_validate(data) + assert env.tracking is None + + def test_tracking_empty_dict_from_yaml(self): + """Simulates YAML `tracking: {}` — should enable with defaults.""" + data = {"tracking": {}} + env = EnvironmentConfig.model_validate(data) + assert env.tracking is not None + assert env.tracking.enabled is True + assert env.tracking.params == {} + + def test_tracking_coexists_with_other_fields(self): + env = EnvironmentConfig( + image="python:3.11", + env={"MY_VAR": "hello"}, + tracking=TrackingConfig(params={"model": "test"}), + ) + assert env.image == "python:3.11" + assert env.env == {"MY_VAR": "hello"} + assert env.tracking.params["model"] == "test" + + def test_serialization_roundtrip_with_tracking(self): + env = EnvironmentConfig(tracking=TrackingConfig(params={"lr": 0.01})) + json_str = env.model_dump_json() + restored = EnvironmentConfig.model_validate_json(json_str) + assert restored.tracking is not None + assert restored.tracking.params["lr"] == 0.01 + + def test_serialization_roundtrip_without_tracking(self): + env = EnvironmentConfig() + json_str = env.model_dump_json() + restored = EnvironmentConfig.model_validate_json(json_str) + assert restored.tracking is None From a24430d215933b63cbfb98d494df366c03baeb64 Mon Sep 17 00:00:00 2001 From: daifangwen Date: Thu, 23 Apr 2026 10:07:11 +0000 Subject: [PATCH 143/226] add api_key --- rock/sdk/envhub/config.py | 4 +++ tests/unit/sdk/job/test_config.py | 52 +++++++++++++++++++++++++++++-- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/rock/sdk/envhub/config.py b/rock/sdk/envhub/config.py index 5ed3a8cc10..ad80ff56cd 100644 --- a/rock/sdk/envhub/config.py +++ b/rock/sdk/envhub/config.py @@ -84,6 +84,10 @@ class TrackingConfig(BaseModel): default=True, description="Whether to enable experiment tracking for this job.", ) + api_key: str | None = Field( + default=None, + description="API key for the tracking platform. Falls back to ROCK_API_KEY env var if not set.", + ) params: dict[str, Any] = Field( default_factory=dict, description=( diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index e8b41e9081..b9896133fa 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -696,11 +696,13 @@ class TestTrackingConfig: def test_default_values(self): config = TrackingConfig() assert config.enabled is True + assert config.api_key is None assert config.params == {} def test_disabled(self): config = TrackingConfig(enabled=False) assert config.enabled is False + assert config.api_key is None assert config.params == {} def test_custom_params(self): @@ -709,22 +711,56 @@ def test_custom_params(self): assert config.params["epochs"] == 10 assert config.params["model"] == "qwen-72b" + def test_api_key(self): + config = TrackingConfig(api_key="sk-test-key-123") + assert config.api_key == "sk-test-key-123" + assert config.enabled is True + + def test_api_key_with_disabled(self): + config = TrackingConfig(enabled=False, api_key="sk-key") + assert config.enabled is False + assert config.api_key == "sk-key" + + def test_api_key_none_by_default(self): + config = TrackingConfig() + assert config.api_key is None + def test_from_dict(self): data = {"enabled": True, "params": {"batch_size": 32}} config = TrackingConfig.model_validate(data) assert config.enabled is True + assert config.api_key is None assert config.params["batch_size"] == 32 + def test_from_dict_with_api_key(self): + data = {"api_key": "sk-from-dict", "params": {"lr": 0.01}} + config = TrackingConfig.model_validate(data) + assert config.api_key == "sk-from-dict" + assert config.params["lr"] == 0.01 + def test_from_dict_minimal(self): config = TrackingConfig.model_validate({}) assert config.enabled is True + assert config.api_key is None assert config.params == {} def test_serialization_roundtrip(self): - config = TrackingConfig(enabled=True, params={"lr": 0.001, "algo": "GRPO"}) + config = TrackingConfig(enabled=True, api_key="sk-round", params={"lr": 0.001, "algo": "GRPO"}) json_str = config.model_dump_json() restored = TrackingConfig.model_validate_json(json_str) assert restored == config + assert restored.api_key == "sk-round" + + def test_exclude_none_omits_api_key_when_not_set(self): + config = TrackingConfig(params={"lr": 0.01}) + data = config.model_dump(mode="json", exclude_none=True) + assert "api_key" not in data + assert data["params"] == {"lr": 0.01} + + def test_exclude_none_includes_api_key_when_set(self): + config = TrackingConfig(api_key="sk-present") + data = config.model_dump(mode="json", exclude_none=True) + assert data["api_key"] == "sk-present" # --------------------------------------------------------------------------- @@ -753,12 +789,23 @@ def test_tracking_with_params(self): assert env.tracking.params["model"] == "qwen-72b" assert env.tracking.params["lr"] == 1e-5 + def test_tracking_with_api_key(self): + env = EnvironmentConfig(tracking=TrackingConfig(api_key="sk-env-key", params={"model": "test"})) + assert env.tracking.api_key == "sk-env-key" + assert env.tracking.params["model"] == "test" + def test_tracking_from_dict(self): data = {"tracking": {"enabled": True, "params": {"batch_size": 64}}} env = EnvironmentConfig.model_validate(data) assert env.tracking is not None assert env.tracking.params["batch_size"] == 64 + def test_tracking_from_dict_with_api_key(self): + data = {"tracking": {"api_key": "sk-yaml-key", "params": {"lr": 0.01}}} + env = EnvironmentConfig.model_validate(data) + assert env.tracking.api_key == "sk-yaml-key" + assert env.tracking.params["lr"] == 0.01 + def test_tracking_none_from_dict(self): data = {"tracking": None} env = EnvironmentConfig.model_validate(data) @@ -783,10 +830,11 @@ def test_tracking_coexists_with_other_fields(self): assert env.tracking.params["model"] == "test" def test_serialization_roundtrip_with_tracking(self): - env = EnvironmentConfig(tracking=TrackingConfig(params={"lr": 0.01})) + env = EnvironmentConfig(tracking=TrackingConfig(api_key="sk-rt", params={"lr": 0.01})) json_str = env.model_dump_json() restored = EnvironmentConfig.model_validate_json(json_str) assert restored.tracking is not None + assert restored.tracking.api_key == "sk-rt" assert restored.tracking.params["lr"] == 0.01 def test_serialization_roundtrip_without_tracking(self): From 7a69d17826b95133ed043f3d68e066b4ae3ff9f6 Mon Sep 17 00:00:00 2001 From: daifangwen Date: Fri, 24 Apr 2026 03:36:30 +0000 Subject: [PATCH 144/226] add tracking to harbor environment config --- rock/sdk/bench/models/trial/config.py | 3 ++- tests/unit/sdk/job/test_config.py | 23 +++++++++++++++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/rock/sdk/bench/models/trial/config.py b/rock/sdk/bench/models/trial/config.py index 05f26ce20d..d4a6982859 100644 --- a/rock/sdk/bench/models/trial/config.py +++ b/rock/sdk/bench/models/trial/config.py @@ -7,7 +7,7 @@ from rock.sdk.bench.models.environment_type import EnvironmentType from rock.sdk.envhub import EnvironmentConfig as _EnvConfig -from rock.sdk.envhub.config import OssMirrorConfig +from rock.sdk.envhub.config import OssMirrorConfig, TrackingConfig class AgentConfig(BaseModel): @@ -33,6 +33,7 @@ class EnvironmentConfig(BaseModel): suppress_override_warnings: bool = False mounts_json: list[dict[str, Any]] | None = None oss_mirror: OssMirrorConfig | None = None + tracking: TrackingConfig | None = None oss_deps: dict[str, str] = Field(default_factory=dict) env: dict[str, str] = Field(default_factory=dict) kwargs: dict[str, Any] = Field(default_factory=dict) diff --git a/tests/unit/sdk/job/test_config.py b/tests/unit/sdk/job/test_config.py index b9896133fa..631d576eef 100644 --- a/tests/unit/sdk/job/test_config.py +++ b/tests/unit/sdk/job/test_config.py @@ -293,12 +293,31 @@ def test_returns_valid_yaml_string(self): parsed = yaml.safe_load(yaml_str) assert isinstance(parsed, dict) + def test_tracking_config_preserved_in_harbor_yaml(self): + """tracking config on environment must survive to_harbor_yaml() serialization.""" + tracking = TrackingConfig(enabled=True, api_key="sk-test-123", params={"lr": 0.01}) + env = RockEnvironmentConfig(tracking=tracking) + cfg = HarborJobConfig(experiment_id="test-exp", environment=env) + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + assert "environment" in data + assert "tracking" in data["environment"], "tracking must not be stripped by to_harbor_yaml()" + assert data["environment"]["tracking"]["enabled"] is True + assert data["environment"]["tracking"]["api_key"] == "sk-test-123" + assert data["environment"]["tracking"]["params"] == {"lr": 0.01} + + def test_tracking_config_none_omitted_in_harbor_yaml(self): + """When tracking is None (default), it should not appear in harbor YAML.""" + cfg = HarborJobConfig(experiment_id="test-exp") + yaml_str = cfg.to_harbor_yaml() + data = yaml.safe_load(yaml_str) + env_data = data.get("environment", {}) + assert "tracking" not in env_data + # --------------------------------------------------------------------------- # HarborJobConfig.from_yaml # --------------------------------------------------------------------------- - - class TestHarborJobConfigFromYaml: def test_round_trip(self, tmp_path): """Write a YAML config, read it back, verify fields.""" From 908f20d4bf2fcbc362dba6a068c1cd399a201f82 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Thu, 28 May 2026 16:14:48 +0800 Subject: [PATCH 145/226] feat(deployments): share docker rootfs XFS prjid with sandbox log dir (#1013) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(deployments): share docker rootfs XFS prjid with sandbox log dir When `--storage-opt size=` is in effect, docker allocates an XFS project id on the container's overlay2 upper dir and binds a bhard limit to it. Previously the sandbox log dir got an independently-allocated prjid + its own bhard (disk_limit_log), so one sandbox had two separate quotas — log writes could exhaust the log quota while rootfs sat idle (or vice versa) even though there was no operational reason to split them. Reuse docker's prjid for the log dir instead. Both paths now share a single bhard (the rootfs limit) — the user-facing knob becomes "this sandbox's total disk write budget", which matches how operators reason about it. Implementation (rock/deployments/docker.py): - _get_docker_rootfs_prjid_and_mountpoint (new): reads the prjid docker assigned to the container's upper dir via `docker inspect` + `xfs_io -r -c lsproj`, plus the mountpoint via findmnt. Only meaningful after `docker create` — the upper dir doesn't exist before that. - _setup_log_dir_quota_shared (new): runs in the create→start gap added by the previous commit. If the log dir lives on the same XFS mountpoint as the rootfs upper dir, attaches the same prjid to it with `xfs_quota project -s` (binds prjid + sets inherit flag). Does NOT call `limit -p` — the bhard belongs to docker; a second limit would clobber it. Cleanup is also docker's responsibility: `docker rm` resets the prjid's limit when the upper dir is torn down. - start() invokes the shared path unconditionally. On any failure (no rootfs prjid, log dir on a different XFS, subprocess error) the log dir simply has no quota in effect — there is no independent fallback. Deployments that need shared-quota enforcement must keep the log dir and docker root on the same filesystem. - _effective_disk_limit_log mirrors the rootfs limit on success, stays None otherwise. Drops the legacy independent-prjid code path entirely: - _try_set_log_dir_quota / _cleanup_log_dir_xfs_quota methods removed - XFS_PRJID_MIN / XFS_PRJID_RANGE constants removed - log_dir_xfs_prjid / log_dir_xfs_mountpoint fields removed (no longer need to track the prjid for manual cleanup) * refactor(deployments): remove disk_limit_log field entirely Now that log dir shares the rootfs XFS prjid (previous commit), the disk_limit_log field is redundant — the log quota always equals the rootfs quota. Remove it from config, models, responses, and tests. The _setup_log_dir_quota_shared() logic that physically binds the log dir to the rootfs prjid via xfs_quota is preserved unchanged. Add backward-compat tests proving old servers sending disk_limit_log are safely ignored by new SDK models (Pydantic extra='ignore'). --- rock/actions/sandbox/response.py | 1 - rock/actions/sandbox/sandbox_info.py | 1 - rock/admin/entrypoints/sandbox_api.py | 7 +- rock/admin/proto/response.py | 3 - rock/common/constants.py | 1 - rock/config.py | 2 - rock/deployments/config.py | 3 - rock/deployments/docker.py | 184 ++++++++++-------- rock/sandbox/sandbox_actor.py | 2 - rock/sandbox/sandbox_manager.py | 1 - .../unit/admin/proto/test_sandbox_response.py | 65 ++----- tests/unit/compat/__init__.py | 0 .../unit/compat/test_disk_limit_log_compat.py | 53 +++++ .../test_docker_deployment_disk_limit.py | 100 +--------- 14 files changed, 169 insertions(+), 254 deletions(-) create mode 100644 tests/unit/compat/__init__.py create mode 100644 tests/unit/compat/test_disk_limit_log_compat.py diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index 5fdb4f10e5..96cf8c172c 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -49,7 +49,6 @@ class SandboxStatusResponse(BaseModel): cpus: float | None = None memory: str | None = None disk_limit_rootfs: str | None = None - disk_limit_log: str | None = None state: State | None = None start_time: str | None = None stop_time: str | None = None diff --git a/rock/actions/sandbox/sandbox_info.py b/rock/actions/sandbox/sandbox_info.py index ae2aac16f7..b0c6282019 100644 --- a/rock/actions/sandbox/sandbox_info.py +++ b/rock/actions/sandbox/sandbox_info.py @@ -22,7 +22,6 @@ class SandboxInfo(TypedDict, total=False): cpus: float memory: str disk_limit_rootfs: str - disk_limit_log: str create_time: str start_time: str stop_time: str diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index ab49081b16..dc17d52687 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -33,7 +33,6 @@ GET_STATUS_SWITCH, KATA_DIND_DISK_SIZE_KEY, KATA_RUNTIME_SWITCH, - SANDBOX_DISK_LIMIT_LOG_KEY, SANDBOX_DISK_LIMIT_ROOTFS_KEY, SUPPORT_KATA_SWITCH, ) @@ -79,23 +78,19 @@ async def _apply_disk_limits(config: DockerDeploymentConfig) -> None: """Apply disk limits from RuntimeConfig (rock-xxx.yml), overridable by Nacos at runtime. Priority: Nacos > RuntimeConfig (rock-xxx.yml). None in both means no limit. + The log dir shares the rootfs prjid + bhard at runtime, so only rootfs is configurable. """ runtime = sandbox_manager.rock_config.runtime nacos = sandbox_manager.rock_config.nacos_provider disk_limit_rootfs = runtime.sandbox_disk_limit_rootfs - disk_limit_log = runtime.sandbox_disk_limit_log if nacos is not None: nacos_rootfs = await nacos.get_config_value(SANDBOX_DISK_LIMIT_ROOTFS_KEY) if nacos_rootfs: disk_limit_rootfs = nacos_rootfs - nacos_log = await nacos.get_config_value(SANDBOX_DISK_LIMIT_LOG_KEY) - if nacos_log: - disk_limit_log = nacos_log config.disk_limit_rootfs = disk_limit_rootfs - config.disk_limit_log = disk_limit_log async def _apply_accelerator_type_validation(config: DockerDeploymentConfig) -> None: diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index c1e1cd9249..eeb9911329 100644 --- a/rock/admin/proto/response.py +++ b/rock/admin/proto/response.py @@ -12,7 +12,6 @@ class SandboxStartResponse(SandboxResponse): cpus: float | None = None memory: str | None = None disk_limit_rootfs: str | None = None - disk_limit_log: str | None = None # TODO: inherit from SandboxStartResponse @@ -33,7 +32,6 @@ class SandboxStatusResponse(BaseModel): cpus: float | None = None memory: str | None = None disk_limit_rootfs: str | None = None - disk_limit_log: str | None = None start_time: str | None = None stop_time: str | None = None create_time: str | None = None @@ -54,7 +52,6 @@ def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusRespons cpus=sandbox_info.get("cpus"), memory=sandbox_info.get("memory"), disk_limit_rootfs=sandbox_info.get("disk_limit_rootfs"), - disk_limit_log=sandbox_info.get("disk_limit_log"), ) diff --git a/rock/common/constants.py b/rock/common/constants.py index 0af5006c7d..c7d867fb22 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -7,7 +7,6 @@ CPU_OVERCOMMIT_ALLOWED_KEYS_KEY = "cpu_overcommit_allowed_keys" KATA_DIND_DISK_SIZE_KEY = "kata_dind_disk_size" SANDBOX_DISK_LIMIT_ROOTFS_KEY = "sandbox_disk_limit_rootfs" -SANDBOX_DISK_LIMIT_LOG_KEY = "sandbox_disk_limit_log" EXTRA_ACCELERATOR_TYPES_KEY = "extra_accelerator_types" PID_PREFIX = "PIDSTART" PID_SUFFIX = "PIDEND" diff --git a/rock/config.py b/rock/config.py index ab68b23d30..e2b2b7d8f7 100644 --- a/rock/config.py +++ b/rock/config.py @@ -264,8 +264,6 @@ class RuntimeConfig: user_defined_tags: dict = field(default_factory=dict) sandbox_disk_limit_rootfs: str | None = None """Default rootfs quota per container. None means no limit. Can be overridden by nacos key 'default_disk_limit'.""" - sandbox_disk_limit_log: str | None = None - """Default log-dir quota per container. None means no limit. Can be overridden by nacos key 'default_log_dir_quota'.""" def __post_init__(self) -> None: # Convert dict to StandardSpec if needed diff --git a/rock/deployments/config.py b/rock/deployments/config.py index 8f03a98b15..a54e944b06 100644 --- a/rock/deployments/config.py +++ b/rock/deployments/config.py @@ -114,9 +114,6 @@ class DockerDeploymentConfig(DeploymentConfig): disk_limit_rootfs: str | None = None """Maximum rootfs disk size for the container (e.g., '20g', '50g'). Maps to --storage-opt size=. Only supported on overlay2 storage driver with xfs backing filesystem. None means no limit.""" - disk_limit_log: str | None = None - """XFS project quota for the sandbox log directory. Server-side only, applied via xfs_quota. None means no limit.""" - container_name: str | None = None """Custom name for the container. If None, a random name will be generated.""" diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 15b58500b7..64a6c0b949 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -1,8 +1,8 @@ import asyncio import datetime -import hashlib import os import random +import re import shlex import subprocess import time @@ -43,8 +43,6 @@ __all__ = ["DockerDeployment", "DockerDeploymentConfig"] CHECK_CLEAR_INTERVAL_SECONDS = 300 -XFS_PRJID_MIN = (1 << 31) # low project IDs are reserved for Docker -XFS_PRJID_RANGE = (1 << 32) - XFS_PRJID_MIN # use remaining 32-bit space: [XFS_PRJID_MIN, 2^32) logger = init_logger(__name__) @@ -65,7 +63,6 @@ def __init__( if registry_password: self._config.registry_password = registry_password self._effective_disk_limit_rootfs: str | None = self._config.disk_limit_rootfs - self._effective_disk_limit_log: str | None = self._config.disk_limit_log self._runtime: RemoteSandboxRuntime | None = None self._container_process = None self._runtime_timeout = 0.15 @@ -89,8 +86,6 @@ def __init__( raise Exception(f"Invalid ROCK_WORKER_ENV_TYPE: {env_vars.ROCK_WORKER_ENV_TYPE}") self.sandbox_validator: DockerSandboxValidator | None = DockerSandboxValidator() - self.log_dir_xfs_prjid: int | None = None - self.log_dir_xfs_mountpoint: str | None = None def add_hook(self, hook: DeploymentHook): self._hooks.add_hook(hook) @@ -212,37 +207,6 @@ def _prepare_kata_disk(self) -> None: os.remove(disk_path) raise - def _cleanup_log_dir_xfs_quota(self) -> None: - """Remove XFS project quota for the sandbox log directory on exit.""" - if self.log_dir_xfs_prjid is None or self.log_dir_xfs_mountpoint is None: - return - if not self._container_name: - return - - log_file_path = f"{env_vars.ROCK_LOGGING_PATH}/{self._container_name}" - project_id = self.log_dir_xfs_prjid - mount_point = self.log_dir_xfs_mountpoint - try: - clear_limit_cmd = f"limit -p bhard=0 bsoft=0 {project_id}" - clear_project_cmd = f"project -C -p {shlex.quote(log_file_path)} {project_id}" - for cmd in (clear_limit_cmd, clear_project_cmd): - result = subprocess.run( - ["xfs_quota", "-x", "-c", cmd, mount_point], - capture_output=True, - text=True, - timeout=5, - ) - if result.returncode != 0: - logger.warning( - f"xfs_quota cleanup failed for {log_file_path!r} cmd={cmd!r}: {result.stderr.strip() or result.stdout.strip()}" - ) - logger.info(f"Cleaned up XFS project quota (prjid={project_id}) for {log_file_path!r}") - except Exception as e: - logger.warning(f"Failed to cleanup XFS project quota for {log_file_path!r}: {e}") - finally: - self.log_dir_xfs_prjid = None - self.log_dir_xfs_mountpoint = None - def _cleanup_kata_disk(self) -> None: """Remove the kata disk image file from the host. @@ -409,60 +373,113 @@ def _storage_opts(self): return ["--storage-opt", f"size={self._effective_disk_limit_rootfs}"] return [] - def _try_set_log_dir_quota(self, log_file_path: str) -> None: - """Best-effort: set XFS project quota for sandbox log directory. + def _get_docker_rootfs_prjid_and_upper_dir(self) -> tuple[int | None, str | None]: + """Read the XFS project id docker assigned to the container's overlay2 + upper dir, together with the upper dir path itself. - Requires the log path to be on an XFS mount with prjquota/pquota enabled. - This check is independent of Docker's storage driver (no overlay2 requirement). + Only meaningful after `docker create` (the upper dir exists by then). + Returns (None, None) on any failure — the caller should fall back to + an independent prjid path. """ - if self._effective_disk_limit_log is None: + try: + inspect_result = subprocess.run( + ["docker", "inspect", "--format={{.GraphDriver.Data.UpperDir}}", self._container_name], + capture_output=True, + text=True, + timeout=5, + ) + if inspect_result.returncode != 0 or not inspect_result.stdout.strip(): + return None, None + upper_dir = inspect_result.stdout.strip() + + lsproj_result = subprocess.run( + ["xfs_io", "-r", "-c", "lsproj", upper_dir], + capture_output=True, + text=True, + timeout=5, + ) + if lsproj_result.returncode != 0: + return None, None + # output looks like: `projid = 12345` (or similar). Grab the trailing integer. + match = re.search(r"(\d+)\s*$", lsproj_result.stdout.strip()) + if not match: + return None, None + prjid = int(match.group(1)) + if prjid <= 0: + return None, None + return prjid, upper_dir + except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError) as e: + logger.warning(f"Failed to read docker rootfs prjid for {self._container_name}: {e}") + return None, None + + def _setup_log_dir_quota_shared(self, log_file_path: str) -> None: + """Attach the docker-allocated rootfs prjid to the host log dir so that + rootfs and log share a single XFS quota (the bhard already set by + `--storage-opt size=`). + + Must be called between `docker create` and `docker start`. On any + failure (no rootfs prjid, log dir on a different filesystem, quota + command error) the log dir is simply left unbound. + + We only `project -s` here (binds prjid + sets inherit flag); we do + NOT call `limit -p`. The bhard belongs to docker; setting a separate + limit would clobber it. Cleanup is also docker's responsibility: + `docker rm` resets the prjid's limit when the upper dir is torn down. + """ + if self._effective_disk_limit_rootfs is None: return - if not DockerUtil.is_xfs_prjquota_path(log_file_path): - logger.info(f"Log path {log_file_path!r} is not on XFS+prjquota, skipping quota setup") - self._effective_disk_limit_log = None + project_id, upper_dir = self._get_docker_rootfs_prjid_and_upper_dir() + logger.info(f"setup_log_dir_quota_shared: log={log_file_path!r}, prjid={project_id}, upper_dir={upper_dir!r}") + if project_id is None or upper_dir is None: + logger.info(f"docker rootfs prjid unavailable for {log_file_path!r}; cannot share, fall back") + return + + try: + if os.stat(log_file_path).st_dev != os.stat(upper_dir).st_dev: + logger.info( + f"log dir {log_file_path!r} on different filesystem from rootfs {upper_dir!r}; " + f"cannot share prjid, fall back" + ) + return + except OSError as e: + logger.warning(f"stat failed while checking prjid sharing for {log_file_path!r}: {e}") return - # Derive a deterministic project id from container name; reserve low ids. - project_id = (int(hashlib.sha1(self.container_name.encode("utf-8")).hexdigest()[:8], 16) % XFS_PRJID_RANGE) + XFS_PRJID_MIN try: findmnt_result = subprocess.run( - ["findmnt", "-T", log_file_path, "-o", "TARGET", "--noheadings"], + ["findmnt", "-T", upper_dir, "-o", "TARGET", "--noheadings"], capture_output=True, text=True, timeout=5, ) - if findmnt_result.returncode != 0: - logger.warning(f"Failed to find mountpoint for log path {log_file_path!r}, skip quota setup") - self._effective_disk_limit_log = None - return - mount_point = findmnt_result.stdout.strip() - if not mount_point: - logger.warning(f"Empty mountpoint for log path {log_file_path!r}, skip quota setup") - self._effective_disk_limit_log = None + if findmnt_result.returncode != 0 or not findmnt_result.stdout.strip(): + logger.warning( + f"findmnt failed for upper_dir {upper_dir!r}: " f"{findmnt_result.stderr.strip() or 'empty output'}" + ) return + xfs_mountpoint = findmnt_result.stdout.strip() set_project_cmd = f"project -s -p {shlex.quote(log_file_path)} {project_id}" - set_limit_cmd = f"limit -p bhard={self._effective_disk_limit_log} {project_id}" - for cmd in (set_project_cmd, set_limit_cmd): - result = subprocess.run( - ["xfs_quota", "-x", "-c", cmd, mount_point], - capture_output=True, - text=True, - timeout=5, + result = subprocess.run( + ["xfs_quota", "-x", "-c", set_project_cmd, xfs_mountpoint], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode != 0: + logger.warning( + f"xfs_quota project -s failed for {log_file_path!r} prjid={project_id}: " + f"{result.stderr.strip() or result.stdout.strip()}" ) - if result.returncode != 0: - logger.warning( - f"xfs_quota failed for {log_file_path!r} with cmd={cmd!r}: {result.stderr.strip() or result.stdout.strip()}" - ) - self._effective_disk_limit_log = None - return - self.log_dir_xfs_prjid = project_id - self.log_dir_xfs_mountpoint = mount_point - logger.info(f"Set XFS project quota {self._effective_disk_limit_log} for log path {log_file_path!r}") - except Exception as e: - logger.warning(f"Failed to set XFS project quota for {log_file_path!r}: {e}") - self._effective_disk_limit_log = None + return + + logger.info( + f"Attached log dir {log_file_path!r} to docker rootfs prjid={project_id} " + f"(shared rootfs+log quota, limit managed by docker --storage-opt)" + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError) as e: + logger.warning(f"Failed to share rootfs prjid with log dir {log_file_path!r}: {e}") async def start(self): """Starts the runtime.""" @@ -480,8 +497,6 @@ async def start(self): self._effective_disk_limit_rootfs = None else: self._effective_disk_limit_rootfs = self._config.disk_limit_rootfs - # Resolve effective log quota; _try_set_log_dir_quota will downgrade to None if XFS+prjquota is unavailable. - self._effective_disk_limit_log = self._config.disk_limit_log if self._container_name is None: self.set_container_name(self._get_container_name()) @@ -508,13 +523,13 @@ async def start(self): env_arg = [] - # Conditionally set up logging path mount based on ROCK_LOGGING_PATH + # Conditionally set up logging path mount based on ROCK_LOGGING_PATH. + log_file_path: str | None = None volume_args = self._prepare_volume_mounts() if env_vars.ROCK_LOGGING_PATH: # Only mount if ROCK_LOGGING_PATH is set (not None or empty) log_file_path = f"{env_vars.ROCK_LOGGING_PATH}/{self.container_name}" os.makedirs(log_file_path, exist_ok=True) os.chmod(log_file_path, 0o777) - self._try_set_log_dir_quota(log_file_path) volume_args.extend(["-v", f"{log_file_path}:{env_vars.ROCK_LOGGING_PATH}"]) env_arg = [ "-e", @@ -574,6 +589,11 @@ async def start(self): # fails before _wait_until_alive sets _container_process up for _stop to manage, we own the # orphan and must remove it. _wait_until_alive's own failure path already calls self.stop(). try: + # Bind log dir to the docker-allocated rootfs prjid in the create→start gap, + # before any container process can write. The bhard is set by `--storage-opt + # size=` on the rootfs prjid, so log and rootfs share one quota. + if log_file_path is not None: + self._setup_log_dir_quota_shared(log_file_path) self._container_process = await loop.run_in_executor(executor, self._docker_start) except Exception: DockerUtil.remove_container_force(self._container_name) @@ -685,7 +705,6 @@ def _stop(self): self._container_process = None self._cleanup_kata_disk() - self._cleanup_log_dir_xfs_quota() self._container_name = None if self._check_stop_task is not None: @@ -720,11 +739,6 @@ def effective_disk_limit_rootfs(self) -> str | None: """Returns the actual rootfs quota in effect after runtime capability checks (may differ from config.disk_limit_rootfs).""" return self._effective_disk_limit_rootfs - @property - def effective_disk_limit_log(self) -> str | None: - """Returns the actual log-dir quota in effect after runtime capability checks (may differ from config.disk_limit_log).""" - return self._effective_disk_limit_log - async def _check_stop(self): logger.info(f"Start check container to stop: {self._container_name}") try: diff --git a/rock/sandbox/sandbox_actor.py b/rock/sandbox/sandbox_actor.py index 96c1799d7c..123402dfca 100644 --- a/rock/sandbox/sandbox_actor.py +++ b/rock/sandbox/sandbox_actor.py @@ -130,7 +130,6 @@ async def start(self): raise ex if isinstance(self._deployment, DockerDeployment): self._config.disk_limit_rootfs = self._deployment.effective_disk_limit_rootfs - self._config.disk_limit_log = self._deployment.effective_disk_limit_log self._clean_container_background() await self._setup_monitor() @@ -280,6 +279,5 @@ async def sandbox_info(self) -> SandboxInfo: "cpus": self._config.cpus, "memory": self._config.memory, "disk_limit_rootfs": self._config.disk_limit_rootfs, - "disk_limit_log": self._config.disk_limit_log, } return {} diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 0ba97e7a9b..81946285b5 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -270,7 +270,6 @@ async def get_status(self, sandbox_id, include_all_states: bool = False) -> Sand cpus=sandbox_info.get("cpus"), memory=sandbox_info.get("memory"), disk_limit_rootfs=sandbox_info.get("disk_limit_rootfs"), - disk_limit_log=sandbox_info.get("disk_limit_log"), start_time=sandbox_info.get("start_time"), stop_time=sandbox_info.get("stop_time"), create_time=sandbox_info.get("create_time"), diff --git a/tests/unit/admin/proto/test_sandbox_response.py b/tests/unit/admin/proto/test_sandbox_response.py index d4868ab10d..64ba8ff10f 100644 --- a/tests/unit/admin/proto/test_sandbox_response.py +++ b/tests/unit/admin/proto/test_sandbox_response.py @@ -1,10 +1,10 @@ """ -Unit tests for admin proto response models — disk_limit_rootfs and disk_limit_log fields. +Unit tests for admin proto response models — disk_limit_rootfs fields. Tests cover: -- SandboxStartResponse.disk_limit_rootfs / disk_limit_log fields -- SandboxStatusResponse.disk_limit_rootfs / disk_limit_log fields -- SandboxStatusResponse.from_sandbox_info() extraction of both fields +- SandboxStartResponse.disk_limit_rootfs field +- SandboxStatusResponse.disk_limit_rootfs field +- SandboxStatusResponse.from_sandbox_info() extraction """ from rock.admin.proto.response import SandboxStartResponse, SandboxStatusResponse @@ -17,30 +17,20 @@ def test_disk_limit_rootfs_default_is_none(self): response = SandboxStartResponse() assert response.disk_limit_rootfs is None - def test_disk_limit_log_default_is_none(self): - response = SandboxStartResponse() - assert response.disk_limit_log is None - def test_disk_limit_rootfs_set_value(self): response = SandboxStartResponse(disk_limit_rootfs="20g") assert response.disk_limit_rootfs == "20g" - def test_disk_limit_log_set_value(self): - response = SandboxStartResponse(disk_limit_log="5g") - assert response.disk_limit_log == "5g" - - def test_all_fields_with_both_limits(self): + def test_all_fields_with_rootfs_limit(self): response = SandboxStartResponse( sandbox_id="test-sandbox", host_ip="10.0.0.1", cpus=4.0, memory="16g", disk_limit_rootfs="50g", - disk_limit_log="5g", ) assert response.sandbox_id == "test-sandbox" assert response.disk_limit_rootfs == "50g" - assert response.disk_limit_log == "5g" assert response.cpus == 4.0 assert response.memory == "16g" @@ -53,20 +43,12 @@ def test_disk_limit_rootfs_default_is_none(self): response = SandboxStatusResponse() assert response.disk_limit_rootfs is None - def test_disk_limit_log_default_is_none(self): - response = SandboxStatusResponse() - assert response.disk_limit_log is None - def test_disk_limit_rootfs_set_value(self): response = SandboxStatusResponse(disk_limit_rootfs="20g") assert response.disk_limit_rootfs == "20g" - def test_disk_limit_log_set_value(self): - response = SandboxStatusResponse(disk_limit_log="5g") - assert response.disk_limit_log == "5g" - - def test_from_sandbox_info_with_both_limits(self): - """from_sandbox_info() should extract both limit fields from SandboxInfo dict.""" + def test_from_sandbox_info_with_rootfs_limit(self): + """from_sandbox_info() should extract disk_limit_rootfs from SandboxInfo dict.""" sandbox_info = { "sandbox_id": "test-sandbox", "phases": {}, @@ -75,16 +57,14 @@ def test_from_sandbox_info_with_both_limits(self): "cpus": 2.0, "memory": "8g", "disk_limit_rootfs": "30g", - "disk_limit_log": "5g", } response = SandboxStatusResponse.from_sandbox_info(sandbox_info) assert response.disk_limit_rootfs == "30g" - assert response.disk_limit_log == "5g" assert response.cpus == 2.0 assert response.memory == "8g" - def test_from_sandbox_info_without_limits(self): - """from_sandbox_info() should yield None for both when absent.""" + def test_from_sandbox_info_without_limit(self): + """from_sandbox_info() should yield None when disk_limit_rootfs is absent.""" sandbox_info = { "sandbox_id": "test-sandbox", "phases": {}, @@ -94,49 +74,32 @@ def test_from_sandbox_info_without_limits(self): } response = SandboxStatusResponse.from_sandbox_info(sandbox_info) assert response.disk_limit_rootfs is None - assert response.disk_limit_log is None - def test_from_sandbox_info_with_none_limits(self): - """from_sandbox_info() should surface None when fields are explicitly None.""" + def test_from_sandbox_info_with_none_limit(self): + """from_sandbox_info() should surface None when field is explicitly None.""" sandbox_info = { "sandbox_id": "test-sandbox", "phases": {}, "port_mapping": {}, "disk_limit_rootfs": None, - "disk_limit_log": None, } response = SandboxStatusResponse.from_sandbox_info(sandbox_info) assert response.disk_limit_rootfs is None - assert response.disk_limit_log is None - - def test_from_sandbox_info_partial_limits(self): - """from_sandbox_info() handles one field set, one absent.""" - sandbox_info = { - "sandbox_id": "test-sandbox", - "phases": {}, - "port_mapping": {}, - "disk_limit_rootfs": "50g", - } - response = SandboxStatusResponse.from_sandbox_info(sandbox_info) - assert response.disk_limit_rootfs == "50g" - assert response.disk_limit_log is None # ---- actions/sandbox/response.SandboxStatusResponse tests ---- class TestActionsSandboxStatusResponseDiskLimit: - def test_actions_status_response_both_limits(self): - """rock.actions.sandbox.response.SandboxStatusResponse should have both limit fields.""" + def test_actions_status_response_rootfs_limit(self): + """rock.actions.sandbox.response.SandboxStatusResponse should have disk_limit_rootfs.""" from rock.actions.sandbox.response import SandboxStatusResponse as ActionStatusResponse - response = ActionStatusResponse(disk_limit_rootfs="20g", disk_limit_log="5g") + response = ActionStatusResponse(disk_limit_rootfs="20g") assert response.disk_limit_rootfs == "20g" - assert response.disk_limit_log == "5g" def test_actions_status_response_defaults_none(self): from rock.actions.sandbox.response import SandboxStatusResponse as ActionStatusResponse response = ActionStatusResponse() assert response.disk_limit_rootfs is None - assert response.disk_limit_log is None diff --git a/tests/unit/compat/__init__.py b/tests/unit/compat/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/compat/test_disk_limit_log_compat.py b/tests/unit/compat/test_disk_limit_log_compat.py new file mode 100644 index 0000000000..64026fe287 --- /dev/null +++ b/tests/unit/compat/test_disk_limit_log_compat.py @@ -0,0 +1,53 @@ +""" +Backward compatibility test: verify that old server responses containing +disk_limit_log are safely ignored by the new SDK models (which no longer +define that field). + +Pydantic v2 default behavior (extra='ignore') discards unknown fields silently. +""" + +from rock.actions.sandbox.response import SandboxStatusResponse +from rock.admin.proto.response import SandboxStartResponse +from rock.admin.proto.response import SandboxStatusResponse as AdminSandboxStatusResponse + + +class TestNewSdkOldServerCompat: + """New SDK (disk_limit_log removed) + Old server (still sends disk_limit_log).""" + + def test_actions_status_response_ignores_extra_disk_limit_log(self): + server_response = { + "sandbox_id": "test-123", + "status": {}, + "port_mapping": {}, + "host_ip": "10.0.0.1", + "cpus": 2.0, + "memory": "8g", + "disk_limit_rootfs": "50g", + "disk_limit_log": "50g", + } + response = SandboxStatusResponse(**server_response) + assert response.disk_limit_rootfs == "50g" + assert not hasattr(response, "disk_limit_log") + + def test_admin_start_response_ignores_extra_disk_limit_log(self): + server_response = { + "sandbox_id": "test-123", + "host_ip": "10.0.0.1", + "disk_limit_rootfs": "50g", + "disk_limit_log": "10g", + } + response = SandboxStartResponse(**server_response) + assert response.disk_limit_rootfs == "50g" + assert not hasattr(response, "disk_limit_log") + + def test_admin_status_response_ignores_extra_disk_limit_log(self): + server_response = { + "sandbox_id": "test-123", + "status": {}, + "port_mapping": {}, + "disk_limit_rootfs": "50g", + "disk_limit_log": "50g", + } + response = AdminSandboxStatusResponse(**server_response) + assert response.disk_limit_rootfs == "50g" + assert not hasattr(response, "disk_limit_log") diff --git a/tests/unit/deployments/test_docker_deployment_disk_limit.py b/tests/unit/deployments/test_docker_deployment_disk_limit.py index 28ed165db0..a678a383d2 100644 --- a/tests/unit/deployments/test_docker_deployment_disk_limit.py +++ b/tests/unit/deployments/test_docker_deployment_disk_limit.py @@ -2,7 +2,7 @@ Unit tests for disk_limit support in DockerDeployment and DockerDeploymentConfig. Tests cover: -- DockerDeploymentConfig default and custom disk_limit_rootfs / disk_limit_log values +- DockerDeploymentConfig default and custom disk_limit_rootfs values - DockerDeployment._storage_opts() argument generation - DockerDeployment.start() graceful degradation when storage-opt is unsupported """ @@ -22,36 +22,19 @@ def test_default_disk_limit_rootfs_is_none(self): config = DockerDeploymentConfig() assert config.disk_limit_rootfs is None - def test_default_disk_limit_log_is_none(self): - config = DockerDeploymentConfig() - assert config.disk_limit_log is None - def test_custom_disk_limit_rootfs(self): config = DockerDeploymentConfig(disk_limit_rootfs="50g") assert config.disk_limit_rootfs == "50g" - def test_custom_disk_limit_log(self): - config = DockerDeploymentConfig(disk_limit_log="5g") - assert config.disk_limit_log == "5g" - def test_disk_limit_rootfs_none(self): config = DockerDeploymentConfig(disk_limit_rootfs=None) assert config.disk_limit_rootfs is None - def test_disk_limit_log_none(self): - config = DockerDeploymentConfig(disk_limit_log=None) - assert config.disk_limit_log is None - def test_disk_limit_rootfs_preserved_in_model_dump(self): config = DockerDeploymentConfig(disk_limit_rootfs="50g") dump = config.model_dump() assert dump["disk_limit_rootfs"] == "50g" - def test_disk_limit_log_preserved_in_model_dump(self): - config = DockerDeploymentConfig(disk_limit_log="5g") - dump = config.model_dump() - assert dump["disk_limit_log"] == "5g" - def test_disk_limit_rootfs_none_preserved_in_model_dump(self): config = DockerDeploymentConfig(disk_limit_rootfs=None) dump = config.model_dump() @@ -125,7 +108,7 @@ async def _run_start(deployment): class TestDockerDeploymentStartDiskLimit: - """Tests that start() applies correct effective values for rootfs and log quotas.""" + """Tests that start() applies correct effective values for rootfs quota.""" @pytest.mark.asyncio @patch("rock.deployments.docker.DockerSandboxValidator") @@ -165,82 +148,3 @@ async def test_no_error_when_rootfs_already_none(self, _mock_detect, _mock_valid assert deployment.config.disk_limit_rootfs is None assert deployment.effective_disk_limit_rootfs is None - - @pytest.mark.asyncio - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=True) - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - async def test_log_downgraded_when_not_xfs_prjquota(self, _mock_prjquota, _mock_detect, _mock_validator): - """When log path is not XFS+prjquota: effective_disk_limit_log=None; config unchanged. - - Note: log quota has NO dependency on docker being overlay2 — - is_xfs_prjquota_path() is the only gate. - """ - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - _make_start_mocks(deployment) - - with ( - patch("rock.deployments.docker.get_executor"), - patch("rock.deployments.docker.asyncio.get_running_loop") as mock_loop, - patch("rock.deployments.docker.wait_until_alive", new_callable=AsyncMock), - patch("rock.deployments.docker.env_vars") as mock_env, - patch("rock.deployments.docker.subprocess"), - ): - mock_env.ROCK_LOGGING_PATH = "/var/log/rock" - mock_env.ROCK_TIME_ZONE = "UTC" - mock_loop.return_value.run_in_executor = AsyncMock() - try: - await deployment.start() - except Exception: - pass - - assert deployment.config.disk_limit_log == "5g" - assert deployment.effective_disk_limit_log is None - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - def test_log_not_downgraded_when_no_log_path(self, _mock_prjquota, _mock_validator): - """When ROCK_LOGGING_PATH is empty, _try_set_log_dir_quota is never called, - so effective_disk_limit_log remains equal to config.disk_limit_log.""" - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - # effective starts equal to config before start() is called - assert deployment.effective_disk_limit_log == "5g" - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - def test_try_set_log_dir_quota_downgrades_when_not_xfs_prjquota(self, _mock_prjquota, _mock_validator): - """_try_set_log_dir_quota: is_xfs_prjquota_path=False → effective_disk_limit_log=None.""" - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - deployment._effective_disk_limit_log = "5g" - deployment._container_name = "test-container" - - deployment._try_set_log_dir_quota("/var/log/rock/test-container") - - assert deployment.effective_disk_limit_log is None - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=True) - def test_try_set_log_dir_quota_independent_of_docker_driver(self, _mock_prjquota, _mock_validator): - """_try_set_log_dir_quota passes the XFS gate regardless of Docker storage driver. - - Log quota only requires is_xfs_prjquota_path(); overlay2 is irrelevant. - The subprocess calls inside (findmnt, xfs_quota) are mocked to succeed. - """ - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - deployment._effective_disk_limit_log = "5g" - deployment._container_name = "test-container" - - with patch("rock.deployments.docker.subprocess") as mock_sub: - ok = MagicMock() - ok.returncode = 0 - ok.stdout = "/var/log/rock" - mock_sub.run.return_value = ok - deployment._try_set_log_dir_quota("/var/log/rock/test-container") - - # xfs_quota succeeded → effective value preserved and prjid recorded - assert deployment.effective_disk_limit_log == "5g" - assert deployment.log_dir_xfs_prjid is not None From 10285c376532b529834296700541e5f24237e6ec Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 29 May 2026 10:26:35 +0800 Subject: [PATCH 146/226] refactor(meta-store): add Redis-merge semantics to archive and filter alive-key fields (#1037) archive() now reads the current Redis alive-key, merges final_info on top (e.g. stop_time, state), and persists the complete snapshot to DB before evicting Redis keys. This ensures the DB always receives full sandbox state regardless of what the caller passes. Also filters SandboxInfo fields when writing the Redis alive key and removes the no-op archive call from stop()'s dangling actor path. --- rock/actions/sandbox/sandbox_info.py | 15 ++++++- rock/sandbox/sandbox_manager.py | 2 - rock/sandbox/sandbox_meta_store.py | 40 ++++++++++++------- tests/unit/sandbox/test_sandbox_meta_store.py | 9 ++--- .../unit/sandbox/test_sandbox_transitions.py | 8 ++-- 5 files changed, 48 insertions(+), 26 deletions(-) diff --git a/rock/actions/sandbox/sandbox_info.py b/rock/actions/sandbox/sandbox_info.py index b0c6282019..3fc96c6f31 100644 --- a/rock/actions/sandbox/sandbox_info.py +++ b/rock/actions/sandbox/sandbox_info.py @@ -1,4 +1,4 @@ -from typing import TypedDict +from typing import Any, TypedDict from rock.actions.sandbox.response import State from rock.deployments.status import PhaseStatus @@ -26,3 +26,16 @@ class SandboxInfo(TypedDict, total=False): start_time: str stop_time: str extended_params: dict[str, str] + + +_SANDBOX_INFO_KEYS = frozenset(SandboxInfo.__annotations__.keys()) + + +def pick_sandbox_info_fields(data: dict[str, Any]) -> "SandboxInfo": + """Return a dict containing only keys declared in :class:`SandboxInfo`. + + Used by ``SandboxMetaStore`` to keep DB-only columns (e.g. ``spec`` / + ``status``, surfaced by ``SandboxRecord.to_dict()`` on the DB-fallback + read path) out of the Redis alive key. + """ + return {k: v for k, v in data.items() if k in _SANDBOX_INFO_KEYS} # type: ignore[return-value] diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 81946285b5..eb19cd270e 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -182,12 +182,10 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL): sm = await self._get_current_statemachine(sandbox_id) if sm is None: logger.info(f"stop dangling sandbox {sandbox_id}") - sandbox_info: SandboxInfo = {"state": State.STOPPED} try: await self._operator.stop(sandbox_id, reason=reason) except ValueError as e: logger.error(f"ray get actor, actor {sandbox_id} not exist", exc_info=e) - await self._meta_store.archive(sandbox_id, sandbox_info) elif sm.current_state.value == State.STOPPED: await sm.send("stop_noop", sandbox_id=sandbox_id) else: diff --git a/rock/sandbox/sandbox_meta_store.py b/rock/sandbox/sandbox_meta_store.py index ed6b97f45d..9b04f66a3f 100644 --- a/rock/sandbox/sandbox_meta_store.py +++ b/rock/sandbox/sandbox_meta_store.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any from rock.actions.sandbox.response import State -from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.actions.sandbox.sandbox_info import SandboxInfo, pick_sandbox_info_fields from rock.admin.core.redis_key import alive_sandbox_key, timeout_sandbox_key from rock.admin.core.sandbox_table import SandboxTable from rock.admin.metrics.decorator import monitor_metastore_operation @@ -71,8 +71,13 @@ async def create( deployment_config: ``DockerDeploymentConfig`` snapshot written once to the ``spec`` DB column. Redis does not store this. + + The Redis payload is filtered to keys declared in ``SandboxInfo`` so any + DB-only fields the caller may carry (e.g. ``spec`` / ``status`` from a + prior DB-fallback read) cannot leak into the alive key. """ - await self._redis.json_set(alive_sandbox_key(sandbox_id), "$", sandbox_info) + redis_payload = pick_sandbox_info_fields(sandbox_info) + await self._redis.json_set(alive_sandbox_key(sandbox_id), "$", redis_payload) if timeout_info is not None: await self._redis.json_set(timeout_sandbox_key(sandbox_id), "$", timeout_info) @@ -80,9 +85,16 @@ async def create( @monitor_metastore_operation async def update(self, sandbox_id: str, sandbox_info: SandboxInfo) -> None: - """Merge *sandbox_info* into the existing Redis alive key and await DB update.""" + """Merge *sandbox_info* into the existing Redis alive key and await DB update. + + The Redis-side merge uses only keys declared in ``SandboxInfo``; DB-only + fields like ``spec`` / ``status`` are dropped so they don't pollute the + alive key. The DB write keeps the full dict — the DB layer has its own + column-based filtering. + """ + redis_payload = pick_sandbox_info_fields(sandbox_info) current = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") - merged: dict[str, Any] = {**(current[0] if current else {}), **sandbox_info} + merged: dict[str, Any] = {**(current[0] if current else {}), **redis_payload} await self._redis.json_set(alive_sandbox_key(sandbox_id), "$", merged) await self._db.update(sandbox_id, sandbox_info) @@ -96,20 +108,20 @@ async def delete(self, sandbox_id: str) -> None: await self._db.delete(sandbox_id) @monitor_metastore_operation - async def archive(self, sandbox_id: str, final_info: SandboxInfo) -> None: - """Persist final state to DB, then remove sandbox from Redis. + async def archive(self, sandbox_id: str, final_info: SandboxInfo | None = None) -> None: + """Snapshot Redis state to DB, then evict Redis keys. - Unlike ``delete``, the DB record is preserved and updated with - ``final_info`` (e.g. ``stop_time``, ``state``). Use this when a - sandbox has finished its lifecycle and the final state should be - queryable from the DB. + Reads the current alive-key from Redis, merges *final_info* on top + (e.g. ``stop_time``, ``state``), persists the result to the DB, and + then deletes the Redis alive + timeout keys. The DB write is awaited before the Redis keys are deleted so that - the final state is always durably stored before the alive key - disappears. If the DB write fails the exception propagates and - Redis cleanup is skipped. + the snapshot is always durably stored before the cache disappears. """ - await self._db.update(sandbox_id, final_info) + current = await self._redis.json_get(alive_sandbox_key(sandbox_id), "$") + merged: dict[str, Any] = {**(current[0] if current else {}), **(final_info or {})} + if merged: + await self._db.update(sandbox_id, merged) await self._redis.json_delete(alive_sandbox_key(sandbox_id)) await self._redis.json_delete(timeout_sandbox_key(sandbox_id)) diff --git a/tests/unit/sandbox/test_sandbox_meta_store.py b/tests/unit/sandbox/test_sandbox_meta_store.py index b90a17a966..c59df7836d 100644 --- a/tests/unit/sandbox/test_sandbox_meta_store.py +++ b/tests/unit/sandbox/test_sandbox_meta_store.py @@ -147,8 +147,8 @@ async def test_remove_deletes_redis_and_db(self, repo, redis, db): class TestArchive: - async def test_archive_removes_redis_and_updates_db(self, repo, redis, db): - """archive() should update DB first, then remove Redis keys.""" + async def test_archive_merges_redis_and_updates_db(self, repo, redis, db): + """archive() should read Redis, merge final_info, write DB, then evict Redis keys.""" await repo.create(SANDBOX_ID, SANDBOX_INFO) await redis.json_set(timeout_sandbox_key(SANDBOX_ID), "$", {"auto_clear_time": "30", "expire_time": "9999"}) await asyncio.sleep(0.1) # let the create fire-and-forget DB insert settle @@ -169,14 +169,13 @@ async def test_archive_removes_redis_and_updates_db(self, repo, redis, db): assert db_record["user_id"] == "user-1" # original fields preserved async def test_archive_db_written_before_redis_deleted(self, repo, redis, db): - """DB must be durably updated before the Redis alive key is removed.""" + """DB must be durably updated before the Redis alive key is evicted.""" await repo.create(SANDBOX_ID, SANDBOX_INFO) await asyncio.sleep(0.1) - # Intercept: check DB state immediately after archive returns (no extra sleep). await repo.archive(SANDBOX_ID, {"state": "stopped"}) - # At this point archive() has already awaited the DB write and deleted Redis. + # At this point archive() has already awaited the DB write and evicted Redis. assert await redis.json_get(alive_sandbox_key(SANDBOX_ID), "$") is None db_record = await db.get(SANDBOX_ID) assert db_record is not None diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py index 320660f855..9f5b503101 100644 --- a/tests/unit/sandbox/test_sandbox_transitions.py +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -73,18 +73,18 @@ async def get_current_statemachine(sandbox_id: str) -> SandboxStateMachine | Non class TestManagerStop: @pytest.mark.asyncio - async def test_stop_not_found_attempts_cleanup(self, mgr, mock_meta_store, mock_operator): + async def test_stop_not_found_attempts_operator_stop(self, mgr, mock_meta_store, mock_operator): mock_meta_store.get.return_value = None await mgr.stop("sb-1") mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) - mock_meta_store.archive.assert_awaited_once() + mock_meta_store.archive.assert_not_awaited() @pytest.mark.asyncio - async def test_stop_not_found_actor_missing_still_archives(self, mgr, mock_meta_store, mock_operator): + async def test_stop_not_found_actor_missing_is_silent(self, mgr, mock_meta_store, mock_operator): mock_meta_store.get.return_value = None mock_operator.stop.side_effect = ValueError("actor not found") await mgr.stop("sb-1") - mock_meta_store.archive.assert_awaited_once() + mock_meta_store.archive.assert_not_awaited() @pytest.mark.asyncio async def test_stop_already_stopped_is_noop(self, mgr, mock_meta_store, mock_operator): From 251f9ba4f30ca7fd2f3c52dae40dffb6e53f652a Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 29 May 2026 18:14:40 +0800 Subject: [PATCH 147/226] feature(sandbox): support sandbox restart (#1001) * fix(sandbox): make container cleanup watchdog controllable from Python The previous shell script spawned the watchdog via `nohup ... &`, so the process saved in `self._clean_container_background_process` was only the short-lived setup script (which exited within milliseconds after detaching the real watchdog). The subsequent `.kill()` call in `actor.stop()` was a no-op and the real watchdog was unreachable from Python. Run the watchdog as the foreground process of the Popen call and add `start_new_session=True` to preserve the SIGHUP isolation that `nohup` provided. With this, the Popen handle points at the actual watchdog and `.kill()` works, which is the prerequisite for the upcoming restart fix (restart must terminate the old watchdog before docker start, otherwise the old watchdog races and `docker stop`s the freshly started container). * feat(sandbox): add /restart endpoint that reuses the existing container restart brings a stopped sandbox back up by running `docker start` on the original container, preserving its filesystem state across the stop/restart cycle. start() remains the path for fresh containers via `docker run`. - POST /restart admin route + SDK Sandbox.restart() - AbstractOperator / RayOperator / SandboxActor / DockerDeployment each expose a restart() method; DockerDeployment.restart() runs `docker start` and validates the container is running. - SandboxStateMachine adds a stopped -> pending transition. on_restart rebuilds DockerDeploymentConfig from the spec snapshot in sandbox_record.spec (DockerDeploymentConfig.model_dump written once by sandbox_table.create), so the new actor wraps the existing container with the same image / memory / cpus / auto_clear. Sandboxes without a spec snapshot fall back to flat sandbox_info fields plus pydantic field defaults. - SandboxManager.restart_async validates the transition and dispatches the SM event; symmetric with stop(). perf(ports): find_free_port resolves docker-published host ports from a module-level cache. do_port_mapping refreshes the cache once per call so its three find_free_port lookups share a single docker scan; standalone find_free_port callers lazy-refresh when the cache is empty. Tests: integration suites for find_free_port and DockerDeployment diagnostics; unit tests for the SM restart transition and the manager restart path. * fix(sandbox): reject restart for kata runtime containers with clear error _stop() calls _cleanup_kata_disk() which deletes the host .img file bound to the container via -v. On restart, docker start cannot mount the missing volume and kata-agent's createContainer fails. Until a dedicated delete API moves disk cleanup out of _stop(), restart is blocked for kata containers. Raise NotImplementedError early in DockerDeployment.restart() so callers get an actionable error instead of a cryptic kata-agent gRPC failure. Signed-off-by: Jiachen Zhang --------- Signed-off-by: Jiachen Zhang --- rock/admin/entrypoints/sandbox_api.py | 7 + .../scripts/clean_container_background.sh | 24 ++- rock/deployments/docker.py | 74 ++++++++ rock/sandbox/operator/abstract.py | 10 ++ rock/sandbox/operator/k8s/operator.py | 3 + rock/sandbox/operator/ray.py | 39 +++- rock/sandbox/sandbox_actor.py | 28 ++- rock/sandbox/sandbox_manager.py | 24 +++ rock/sandbox/sandbox_statemachine.py | 44 ++++- rock/sdk/sandbox/client.py | 20 +++ rock/utils/__init__.py | 2 + rock/utils/system.py | 118 +++++++++++-- .../deployments/test_docker_diagnostics.py | 166 ++++++++++++++++++ tests/integration/utils/__init__.py | 0 .../integration/utils/test_find_free_port.py | 157 +++++++++++++++++ .../unit/sandbox/test_sandbox_statemachine.py | 100 +++++++++++ .../unit/sandbox/test_sandbox_transitions.py | 141 +++++++++++++-- 17 files changed, 915 insertions(+), 42 deletions(-) create mode 100644 tests/integration/deployments/test_docker_diagnostics.py create mode 100644 tests/integration/utils/__init__.py create mode 100644 tests/integration/utils/test_find_free_port.py diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index dc17d52687..9a6a469d34 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -277,6 +277,13 @@ async def close(sandbox_id: str = Body(..., embed=True)) -> RockResponse[str]: return RockResponse(result=f"{sandbox_id} stopped") +@sandbox_router.post("/restart") +@handle_exceptions(error_message="restart sandbox failed") +async def restart(sandbox_id: str = Body(..., embed=True)) -> RockResponse[SandboxStartResponse]: + result = await sandbox_manager.restart_async(sandbox_id) + return RockResponse(result=result) + + @sandbox_router.post("/commit") @handle_exceptions(error_message="commit sandbox failed") async def commit( diff --git a/rock/admin/scripts/clean_container_background.sh b/rock/admin/scripts/clean_container_background.sh index 691f7b7d2b..227e8bcc0e 100644 --- a/rock/admin/scripts/clean_container_background.sh +++ b/rock/admin/scripts/clean_container_background.sh @@ -1,18 +1,14 @@ #!/bin/bash -# Container cleanup function -setup_container_cleanup() { - PID=$1 - CONTAINER_NAME=$2 - echo "ray actor pid is $PID" +# Watchdog: poll until the actor process disappears, then stop the container. +# Runs as the foreground process of the Popen call in SandboxActor so the +# Python side can terminate it via process.kill() during stop/restart. +PID=$1 +CONTAINER_NAME=$2 +echo "ray actor pid is $PID" - nohup bash -c " - while [ -e /proc/$PID ]; do - sleep 1 - done - docker stop $CONTAINER_NAME - " > /dev/null 2>&1 & -} +while [ -e /proc/$PID ]; do + sleep 1 +done -setup_container_cleanup "$@" -echo "start container cleanup success" \ No newline at end of file +docker stop $CONTAINER_NAME diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 64a6c0b949..adcfb598ff 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -35,6 +35,7 @@ StageTimer, find_free_port, get_executor, + refresh_docker_used_ports, release_port, sandbox_id_ctx_var, timeout, @@ -667,6 +668,78 @@ async def stop(self): if self._runtime: await loop.run_in_executor(stop_executor, self._stop) + async def restart(self): + """Restart an existing stopped container using docker start. + + Precondition: caller (SandboxStateMachine) guarantees the container + is in a stopped/exited state. A nonexistent container surfaces via + `docker start` failing — see is_alive()'s poll-based detection. + """ + # TODO: once a sandbox delete API exists, move _cleanup_kata_disk() there; + # until then kata restart is blocked because _stop() deletes the .img file. + if self._config.use_kata_runtime: + raise NotImplementedError( + f"Restart is not supported for kata runtime containers (container={self._container_name}). " + ) + + executor = get_executor() + loop = asyncio.get_running_loop() + + logger.info(f"Restarting container {self._container_name} with docker start") + + # Reuse the same Popen-based attached start used by start(), so the + # restart path also produces a valid self._container_process. Without + # this, _stop() would skip its `if self._container_process is not None` + # branch and never call docker kill / cleanup. + self._container_process = await loop.run_in_executor(executor, self._docker_start) + + # Recover the rocklet port from the container's port bindings if not set in config. + # When a new actor is created for restart, config.port may be None. + if self._config.port is None: + self._config.port = await loop.run_in_executor(executor, self._get_rocklet_port_from_inspect) + if self._config.port is None: + raise Exception(f"Cannot determine rocklet port for container {self._container_name}") + + # Re-establish runtime connection + logger.info(f"Starting runtime at {self._config.port}") + self._runtime = RemoteSandboxRuntime.from_config( + RemoteSandboxRuntimeConfig(port=self._config.port, timeout=self._runtime_timeout) + ) + self._runtime.set_executor(executor) + + # Wait until container is alive + with StageTimer("startup_timing", f"[{self._container_name}] Wait until alive", logger): + await self._wait_until_alive(timeout=self._config.startup_timeout) + + # Re-enable auto-clear if configured + if self._config.enable_auto_clear: + self._check_stop_task = asyncio.create_task(self._check_stop()) + + logger.info(f"Container {self._container_name} restarted successfully") + + def _get_rocklet_port_from_inspect(self) -> int | None: + """Read the host-side port mapped to the rocklet (container port 22555) from docker inspect.""" + try: + result = subprocess.run( + [ + "docker", + "inspect", + "--format", + '{{(index (index .HostConfig.PortBindings "22555/tcp") 0).HostPort}}', + self._container_name, + ], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode == 0: + port_str = result.stdout.strip() + if port_str.isdigit(): + return int(port_str) + except Exception as e: + logger.warning(f"Failed to get rocklet port from inspect for {self._container_name}: {e}") + return None + def _stop(self): """Stops the runtime.""" if self._container_name in ENV_POOL: @@ -772,6 +845,7 @@ def get_status(self) -> ServiceStatus: return self._service_status async def do_port_mapping(self): + refresh_docker_used_ports() proxy_port = await find_free_port() self._service_status.add_port_mapping(Port.PROXY, proxy_port) ssh_port = await find_free_port() diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index 997c95546b..8f6347095f 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -18,6 +18,16 @@ class AbstractOperator(ABC): async def submit(self, config: DeploymentConfig, user_info: dict = {}) -> SandboxInfo: ... + @abstractmethod + async def restart(self, config: DeploymentConfig, host_ip: str | None = None) -> SandboxInfo: + """Restart an existing stopped container using docker start. + + The actor for this sandbox has already been killed by stop(). + Implementations must create a new actor and invoke docker start + on the existing (stopped) container — not docker run. + """ + ... + @abstractmethod async def get_status(self, sandbox_id: str) -> SandboxInfo | None: ... diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index ca9e1ac66e..aeaf3c756a 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -90,6 +90,9 @@ async def submit(self, config: DockerDeploymentConfig, user_info: dict = {}) -> """ return await self._provider.submit(config, user_info) + async def restart(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> SandboxInfo: + raise NotImplementedError("K8sOperator does not support container-reuse restart") + async def get_status(self, sandbox_id: str) -> SandboxInfo | None: """Get sandbox status with user info from Redis. diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index fa8b03150b..1ce9acb826 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -32,19 +32,25 @@ def __init__(self, ray_service: RayService, runtime_config: RuntimeConfig): def _get_actor_name(self, sandbox_id: str) -> str: return f"sandbox-{sandbox_id}" - async def create_actor(self, config: DockerDeploymentConfig): - actor_options = self._generate_actor_options(config) + async def create_actor(self, config: DockerDeploymentConfig, pin_to_host_ip: str | None = None): + actor_options = self._generate_actor_options(config, pin_to_host_ip=pin_to_host_ip) deployment: DockerDeployment = config.get_deployment() sandbox_actor = SandboxActor.options(**actor_options).remote(config, deployment) return sandbox_actor - def _generate_actor_options(self, config: DockerDeploymentConfig) -> dict: + def _generate_actor_options(self, config: DockerDeploymentConfig, pin_to_host_ip: str | None = None) -> dict: actor_name = self._get_actor_name(config.container_name) actor_options = {"name": actor_name, "lifetime": "detached"} try: memory = parse_size_to_bytes(config.memory) actor_options["num_cpus"] = config.cpus actor_options["memory"] = memory + # Pin to a specific node via Ray's implicit `node:` resource + # (registered automatically per node with value 1.0; we consume a + # negligible 0.001 so we don't block other actors). Used by restart + # to land on the host that owns the existing container. + if pin_to_host_ip: + actor_options["resources"] = {f"node:{pin_to_host_ip}": 0.001} return actor_options except ValueError as e: logger.warning(f"Invalid memory size: {config.memory}", exc_info=e) @@ -114,6 +120,33 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> ray.kill(actor) return True + async def restart(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> SandboxInfo: + """Restart an existing sandbox using docker start (container is preserved). + + Flow: + 1. Create a fresh detached actor pinned to ``host_ip`` (the node that + owns the existing container) — without this Ray would schedule the + new actor on any free node and `docker inspect ` would + fail with "container does not exist" on that wrong node. + 2. actor.restart() → deployment.restart() — DockerDeployment.restart() + handles a still-running container by issuing `docker kill` first. + """ + async with self._ray_service.get_ray_rwlock().read_lock(): + sandbox_id = config.container_name + + if not host_ip: + logger.warning( + f"restart for {sandbox_id} called without host_ip; new actor " + f"may be scheduled on a node that does not own the container" + ) + sandbox_actor: SandboxActor = await self.create_actor(config, pin_to_host_ip=host_ip) + await self._ray_service.async_ray_get(sandbox_actor.restart.remote()) + + sandbox_info: SandboxInfo = await self._ray_service.async_ray_get(sandbox_actor.sandbox_info.remote()) + sandbox_info["state"] = State.PENDING + logger.info(f"sandbox {sandbox_id} restarted") + return sandbox_info + async def _check_alive_status(self, sandbox_id: str, host_ip: str, remote_status: ServiceStatus) -> bool: """Check if sandbox is alive""" try: diff --git a/rock/sandbox/sandbox_actor.py b/rock/sandbox/sandbox_actor.py index 123402dfca..b29379995f 100644 --- a/rock/sandbox/sandbox_actor.py +++ b/rock/sandbox/sandbox_actor.py @@ -74,9 +74,14 @@ def _clean_container_background(self): logger.info("start to run background script") actor_pid = str(os.getpid()) logger.info(f"actor_pid is {actor_pid}") - # Execute script + # start_new_session: detach from actor's session so SIGHUP doesn't + # take the watchdog down (replaces the script's prior `nohup ... &` + # form, which made the watchdog uncontrollable from Python). process = subprocess.Popen( ["bash", self._clean_container_background_script, actor_pid, self._config.container_name], + start_new_session=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, ) logger.info(f"Background script started successfully, pid is {process.pid}") self._clean_container_background_process = process @@ -145,6 +150,27 @@ async def stop(self, reason: StopReason = StopReason.MANUAL): finally: self.log_lifecycle_summary(reason) + async def restart(self): + """Restart an existing stopped container using docker start.""" + logger.info(f"[{self._config.container_name}] start to restart") + try: + await self._deployment.restart() + logger.info(f"[{self._config.container_name}] deployment restarted") + + # Re-arm the cleanup watchdog so a future actor death still triggers + # docker stop — start() does this too; restart must stay symmetric. + if isinstance(self._deployment, DockerDeployment): + self._clean_container_background() + + # Re-establish monitoring after restart + await self._setup_monitor() + logger.info(f"[{self._config.container_name}] actor restarted") + except Exception as e: + logger.error( + f"[{self._config.container_name}] Error occurred while restarting container: {e}", exc_info=True + ) + raise + async def commit(self, image_tag: str, username: str, password: str) -> CommandResponse: logger.info(f"start to commit {self._config.container_name} to {image_tag}") with tempfile.TemporaryDirectory() as docker_config_dir: diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index eb19cd270e..f0e6d58a06 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -150,6 +150,30 @@ async def start_async( host_ip=sandbox_info.get("host_ip"), ) + @monitor_sandbox_operation() + async def restart_async(self, sandbox_id: str) -> SandboxStartResponse: + sm = await self._get_current_statemachine(sandbox_id) + if sm is None: + raise BadRequestRockError(f"Sandbox {sandbox_id} not found") + + state = sm.current_state.value + if state != State.STOPPED: + raise BadRequestRockError(f"Sandbox {sandbox_id} cannot be restarted: current state is '{state.value}'") + + await sm.send( + "restart", + sandbox_id=sandbox_id, + operator=self._operator, + meta_store=self._meta_store, + ) + + info: SandboxInfo = sm.sandbox_info or {} + return SandboxStartResponse( + sandbox_id=sandbox_id, + host_name=info.get("host_name"), + host_ip=info.get("host_ip"), + ) + @monitor_sandbox_operation() async def start(self, config: DeploymentConfig) -> SandboxStartResponse: docker_deployment_config: DockerDeploymentConfig = await self.deployment_manager.init_config(config) diff --git a/rock/sandbox/sandbox_statemachine.py b/rock/sandbox/sandbox_statemachine.py index 889c4cfe77..bf215c05b3 100644 --- a/rock/sandbox/sandbox_statemachine.py +++ b/rock/sandbox/sandbox_statemachine.py @@ -13,7 +13,10 @@ from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.metrics.billing import log_billing_info from rock.common.constants import StopReason +from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger +from rock.sandbox.utils.timeout import SandboxTimeoutHelper +from rock.sdk.common.exceptions import BadRequestRockError from rock.utils.system import get_iso8601_timestamp logger = init_logger(__name__) @@ -46,6 +49,7 @@ class SandboxStateMachine(StateChart): stop = pending.to(stopped) | running.to(stopped) stop_noop = stopped.to(stopped) alive = pending.to(running) + restart = stopped.to(pending) def __init__(self, **kwargs): """Initialize with optional sandbox_info.""" @@ -87,8 +91,44 @@ async def on_alive(self, sandbox_id: str, meta_store, sandbox_info: SandboxInfo) sandbox_info["start_time"] = get_iso8601_timestamp() await meta_store.update(sandbox_id, sandbox_info) - # Update self.sandbox_info for potential future use - self.sandbox_info = sandbox_info + async def on_restart(self, sandbox_id: str, operator, meta_store) -> None: + info = self.sandbox_info or {} + + host_ip = info.get("host_ip") + if not host_ip: + raise BadRequestRockError(f"Sandbox {sandbox_id} has no host_ip; cannot pin restart to original node") + + # Prefer the spec snapshot (DockerDeploymentConfig.model_dump persisted to + # the DB at start time) so the new actor wraps the existing container with + # the exact same config. + spec = info.get("spec") or {} + if spec: + restart_config = DockerDeploymentConfig(**spec) + else: + logger.warning( + f"sandbox {sandbox_id} has no spec snapshot; rebuilding config from flat fields with model defaults" + ) + restart_config = DockerDeploymentConfig( + container_name=sandbox_id, + image=info.get("image") or DockerDeploymentConfig.model_fields["image"].default, + memory=info.get("memory") or DockerDeploymentConfig.model_fields["memory"].default, + cpus=float(info.get("cpus") or DockerDeploymentConfig.model_fields["cpus"].default), + ) + timeout_info = SandboxTimeoutHelper.make_timeout_info(restart_config.auto_clear_time) + + logger.info(f"restart sandbox {sandbox_id} (pin host_ip={host_ip})") + await operator.restart(restart_config, host_ip=host_ip) + + # The previous stop() called meta_store.archive() which removed the Redis + # alive key, so a partial update would lose every field except `state`. + # Re-seed Redis with the full sandbox_info restored from the DB. + # meta_store.update filters to SandboxInfo-declared keys, so DB-only + # fields (spec/status) won't pollute the alive key. + new_info = dict(info) + new_info["state"] = RockState.PENDING + new_info.pop("stop_time", None) + await meta_store.update(sandbox_id, new_info) + await meta_store.update_timeout(sandbox_id, timeout_info) @classmethod async def from_state_value(cls, state_value: str | None, sandbox_info: SandboxInfo) -> "SandboxStateMachine": diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 3a57eacda5..a6029f9415 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -267,6 +267,26 @@ async def stop(self): except Exception as e: logging.warning(f"Failed to stop sandbox, IGNORE: {e}") + async def restart(self): + """Restart a stopped sandbox using 'docker start' (reuses existing container). + + The sandbox must be in STOPPED state before calling this method. + After restart, the sandbox will be back in RUNNING state with the same container. + """ + if not self.sandbox_id: + raise Exception("sandbox_id is not set, cannot restart") + url = f"{self._url}/restart" + headers = self._build_headers() + data = {"sandbox_id": self.sandbox_id} + response = await HttpUtils.post(url, headers, data) + logging.debug(f"Restart sandbox response: {response}") + if "Success" != response.get("status"): + result = response.get("result", None) + if result is not None: + rock_response = SandboxResponse(**result) + raise_for_code(rock_response.code, f"Failed to restart sandbox: {response}") + raise Exception(f"Failed to restart sandbox: {response}") + async def commit(self, image_tag: str, username: str, password: str): if not self.sandbox_id: return diff --git a/rock/utils/__init__.py b/rock/utils/__init__.py index c33716fb46..04f48f8f0d 100644 --- a/rock/utils/__init__.py +++ b/rock/utils/__init__.py @@ -35,6 +35,7 @@ get_host_ip, get_instance_id, get_uniagent_endpoint, + refresh_docker_used_ports, release_port, run_command_with_output, run_shell_command, @@ -57,6 +58,7 @@ "run_shell_command", "extract_nohup_pid", "find_free_port", + "refresh_docker_used_ports", "release_port", "get_instance_id", "get_host_ip", diff --git a/rock/utils/system.py b/rock/utils/system.py index a79cc55224..a4a3d8a329 100644 --- a/rock/utils/system.py +++ b/rock/utils/system.py @@ -1,5 +1,6 @@ import asyncio import datetime +import json import logging import os import re @@ -8,7 +9,6 @@ import time import zoneinfo from pathlib import Path -from threading import Lock from rock import env_vars from rock.common.constants import PID_PREFIX @@ -17,7 +17,13 @@ _REGISTERED_PORTS = set() -_REGISTERED_PORTS_LOCK = Lock() + +# Scratchpad cache for docker-published host ports. Populated by +# `refresh_docker_used_ports()` (typically at the top of a batch port +# allocation like `do_port_mapping`) so we don't re-shell out for every +# `find_free_port` call inside the batch. Stale across batches is fine — +# each batch refreshes. +_DOCKER_USED_PORTS: set[int] = set() def run_command_with_output(cmd, wait=False): @@ -86,8 +92,96 @@ def extract_nohup_pid(nohup_output: str) -> int: return None +def _get_docker_used_host_ports() -> set[int]: + """Return host ports reserved by any docker container (running OR stopped). + + If the module-level ``_DOCKER_USED_PORTS`` cache has been populated by a + recent ``refresh_docker_used_ports()`` call (e.g. at the top of + ``do_port_mapping``), return the cached snapshot to avoid re-shelling + out for each ``find_free_port`` in the same batch. + + Stopped containers keep their ``HostConfig.PortBindings`` metadata; on + ``docker start`` the daemon re-binds the same host port. If we handed that + port out to another sandbox in the meantime, the later restart would fail + with ``address already in use``. ``docker ps`` ``Ports`` column hides this + for stopped containers, so we inspect ``HostConfig.PortBindings`` directly. + """ + if _DOCKER_USED_PORTS: + return set(_DOCKER_USED_PORTS) + used: set[int] = set() + try: + ls = subprocess.run( + ["docker", "ps", "-aq"], + capture_output=True, + text=True, + timeout=10, + ) + if ls.returncode != 0: + logger.debug(f"docker ps failed ({ls.returncode}): {ls.stderr.strip()}") + return used + ids = ls.stdout.split() + if not ids: + return used + inspect = subprocess.run( + ["docker", "inspect", "--format", "{{json .HostConfig.PortBindings}}", *ids], + capture_output=True, + text=True, + timeout=15, + ) + if inspect.returncode != 0: + logger.debug(f"docker inspect failed ({inspect.returncode}): {inspect.stderr.strip()}") + return used + for line in inspect.stdout.splitlines(): + line = line.strip() + if not line or line == "null": + continue + try: + bindings = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(bindings, dict): + continue + # bindings == {"22555/tcp": [{"HostIp": "", "HostPort": "55555"}, ...], ...} + for host_specs in bindings.values(): + for spec in host_specs or []: + p = (spec or {}).get("HostPort") + if p and p.isdigit(): + used.add(int(p)) + except Exception as e: + logger.warning(f"failed to enumerate docker-used ports: {e}") + return used + + +def refresh_docker_used_ports() -> set[int]: + """Refresh the module-level ``_DOCKER_USED_PORTS`` cache from docker. + + Call this once at the start of a batch port allocation (e.g. + ``do_port_mapping``) so subsequent ``find_free_port`` calls in the same + batch reuse a single docker scan. Returns a snapshot of the cache. + + No locking: each Ray sandbox actor is a single-threaded asyncio process + and ``do_port_mapping`` calls this once per sandbox creation; there are no + concurrent writers in practice. + """ + _DOCKER_USED_PORTS.clear() + _DOCKER_USED_PORTS.update(_get_docker_used_host_ports()) + return set(_DOCKER_USED_PORTS) + + async def find_free_port(max_attempts: int = 10, sleep_between_attempts: float = 0.1) -> int: - """Find a free port that is not yet registered + """Find a free port avoiding three classes of collisions: + + 1. **Same-process concurrent callers** — via in-memory ``_REGISTERED_PORTS``. + 2. **Restart-safety across processes** — via ``_DOCKER_USED_PORTS``. + Callers that need protection against stopped docker containers + reclaiming a port on restart MUST call ``refresh_docker_used_ports()`` + before their batch; this function reads the cache as-is and does not + refresh on its own. Callers that only need an OS-listenable port + (e.g. binding a test server) may skip the refresh. + 3. **Kernel-level races on currently-listening ports** — OS + ``bind(("", 0))`` only hands out a port the kernel considers free, + so we don't need to enumerate listeners separately; on collision + we just retry. Args: max_attempts: Maximum number of attempts to find a port @@ -99,16 +193,18 @@ async def find_free_port(max_attempts: int = 10, sleep_between_attempts: float = Raises: RuntimeError: If unable to find a free port after max_attempts """ + cached = set(_DOCKER_USED_PORTS) for _ in range(max_attempts): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - with _REGISTERED_PORTS_LOCK: - s.bind(("", 0)) - port = s.getsockname()[1] - if port not in _REGISTERED_PORTS: - _REGISTERED_PORTS.add(port) - logger.debug(f"Found free port {port}") - return port - logger.debug(f"Port {port} already registered, trying again after {sleep_between_attempts}s") + s.bind(("", 0)) + port = s.getsockname()[1] + if port not in _REGISTERED_PORTS and port not in cached: + _REGISTERED_PORTS.add(port) + logger.debug(f"Found free port {port}") + return port + logger.debug( + f"Port {port} already registered or held by docker, " f"trying again after {sleep_between_attempts}s" + ) time.sleep(sleep_between_attempts) msg = f"Failed to find a unique free port after {max_attempts} attempts" raise RuntimeError(msg) diff --git a/tests/integration/deployments/test_docker_diagnostics.py b/tests/integration/deployments/test_docker_diagnostics.py new file mode 100644 index 0000000000..f8005c0048 --- /dev/null +++ b/tests/integration/deployments/test_docker_diagnostics.py @@ -0,0 +1,166 @@ +""" +Integration tests for the docker-start side helpers used by ``restart()``: + +- ``_docker_start()`` — attached ``docker start -a`` (also used by start()) +- ``_get_rocklet_port_from_inspect()`` — recovers the rocklet host port from a live container + +We don't drive ``restart()`` end-to-end here (that needs an HTTP probe against +rocklet, covered by ``test_restart_e2e.py``); instead we verify the building +blocks behave correctly against real docker so the test catches docker CLI +format regressions as well as our parsing logic. +""" + +from __future__ import annotations + +import socket +import subprocess +import time +import uuid + +import pytest + +from rock.deployments.config import DockerDeploymentConfig +from rock.deployments.docker import DockerDeployment +from tests.integration.conftest import SKIP_IF_NO_DOCKER + + +def _free_port() -> int: + """Synchronous helper: bind to port 0 to get an OS-allocated free port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _docker_state(name: str) -> str: + """Read container's State.Status via docker inspect; "unknown" on failure.""" + result = subprocess.run( + ["docker", "inspect", "--format={{.State.Status}}", name], + capture_output=True, + text=True, + timeout=10, + ) + return result.stdout.strip() if result.returncode == 0 else "unknown" + + +pytestmark = [pytest.mark.need_docker, SKIP_IF_NO_DOCKER] + + +@pytest.fixture +def container_name() -> str: + """Unique container name so parallel test runs don't clash.""" + return f"diag-test-{uuid.uuid4().hex[:12]}" + + +@pytest.fixture +def deployment(container_name: str) -> DockerDeployment: + cfg = DockerDeploymentConfig( + container_name=container_name, + image="python:3.11", + memory="512m", + cpus=0.5, + ) + return cfg.get_deployment() + + +class TestDockerStartHelpers: + """Tests for the docker-start side helpers used by restart() — without + requiring a real rocklet image.""" + + def test_docker_start_revives_exited_container_with_same_id_and_ports( + self, container_name: str, deployment: DockerDeployment + ): + # Create container with **explicit host:container port mapping** — + # this mirrors DockerDeployment._docker_run() which always passes + # `-p {host_port}:{container_port}`. Without an explicit host port, + # docker re-allocates a fresh host port on each start, so we must + # match production behavior to test the "ports preserved" invariant. + host_port = _free_port() + subprocess.run( + [ + "docker", + "run", + "-d", + "--name", + container_name, + "-p", + f"{host_port}:22555", + "--entrypoint", + "", + "python:3.11", + "sh", + "-c", + "sleep 30", + ], + check=True, + capture_output=True, + timeout=60, + ) + try: + cid_before = subprocess.run( + ["docker", "inspect", "--format={{.Id}}", container_name], + capture_output=True, + text=True, + check=True, + timeout=10, + ).stdout.strip() + host_port_before = deployment._get_rocklet_port_from_inspect() + assert host_port_before is not None + + # Stop, then docker start via the helper under test + subprocess.run(["docker", "kill", container_name], check=True, capture_output=True, timeout=15) + for _ in range(20): + if _docker_state(container_name) == "exited": + break + time.sleep(0.2) + assert _docker_state(container_name) == "exited" + + deployment._docker_start() # the function actually used by restart() + + # Container should be running again + for _ in range(20): + if _docker_state(container_name) == "running": + break + time.sleep(0.2) + assert _docker_state(container_name) == "running" + + cid_after = subprocess.run( + ["docker", "inspect", "--format={{.Id}}", container_name], + capture_output=True, + text=True, + check=True, + timeout=10, + ).stdout.strip() + host_port_after = deployment._get_rocklet_port_from_inspect() + + # Critical invariants for restart semantics + assert cid_after == cid_before, "container id must be preserved (same container, not new one)" + assert host_port_after == host_port_before, "published port mapping must be preserved" + finally: + subprocess.run(["docker", "rm", "-f", container_name], capture_output=True, timeout=15) + + def test_get_rocklet_port_returns_none_when_no_22555_mapping( + self, container_name: str, deployment: DockerDeployment + ): + # Container without `-p 22555` shouldn't expose port; helper returns None + subprocess.run( + [ + "docker", + "run", + "-d", + "--name", + container_name, + "--entrypoint", + "", + "python:3.11", + "sh", + "-c", + "sleep 30", + ], + check=True, + capture_output=True, + timeout=60, + ) + try: + assert deployment._get_rocklet_port_from_inspect() is None + finally: + subprocess.run(["docker", "rm", "-f", container_name], capture_output=True, timeout=15) diff --git a/tests/integration/utils/__init__.py b/tests/integration/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integration/utils/test_find_free_port.py b/tests/integration/utils/test_find_free_port.py new file mode 100644 index 0000000000..4d79e484de --- /dev/null +++ b/tests/integration/utils/test_find_free_port.py @@ -0,0 +1,157 @@ +""" +Integration tests for find_free_port — covers the cross-process protection +added by querying `docker ps` for host ports already published by any container +(running OR stopped). + +These tests use real docker (no mocks) so they validate both our parsing logic +and docker CLI output format. +""" + +from __future__ import annotations + +import asyncio +import socket +import subprocess +import uuid + +import pytest + +from rock.utils import system as _system +from rock.utils.system import ( + _REGISTERED_PORTS, + _get_docker_used_host_ports, + find_free_port, + refresh_docker_used_ports, + release_port, +) +from tests.integration.conftest import SKIP_IF_NO_DOCKER + +pytestmark = [pytest.mark.need_docker, SKIP_IF_NO_DOCKER] + + +@pytest.fixture(autouse=True) +def _clear_docker_used_ports_cache(): + # ``_get_docker_used_host_ports`` short-circuits to the module-level + # ``_DOCKER_USED_PORTS`` snapshot when it is non-empty. Other tests in this + # worker may have populated it via ``refresh_docker_used_ports`` / + # ``do_port_mapping``; clear before and after each test so each one starts + # from a clean slate. + _system._DOCKER_USED_PORTS.clear() + yield + _system._DOCKER_USED_PORTS.clear() + + +def _free_port_via_os() -> int: + """Get an OS-allocated free port (TOCTOU: port may be re-allocated soon).""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def container_holding_port(): + """Spin up a container that publishes a chosen host port; yield (name, port).""" + name = f"fft-{uuid.uuid4().hex[:10]}" + host_port = _free_port_via_os() + subprocess.run( + [ + "docker", + "run", + "-d", + "--name", + name, + "-p", + f"{host_port}:22555", + "--entrypoint", + "", + "python:3.11", + "sh", + "-c", + "sleep 30", + ], + check=True, + capture_output=True, + timeout=60, + ) + try: + yield name, host_port + finally: + subprocess.run(["docker", "rm", "-f", name], capture_output=True, timeout=15) + + +class TestGetDockerUsedHostPorts: + def test_lists_running_container_ports(self, container_holding_port): + _name, port = container_holding_port + used = _get_docker_used_host_ports() + assert port in used, f"running container's host port {port} should be in used set" + + def test_lists_stopped_container_ports(self, container_holding_port): + # Stopped containers still hold the host_port:container_port reservation + # in docker metadata — must also be avoided by find_free_port. + name, port = container_holding_port + subprocess.run(["docker", "kill", name], check=True, capture_output=True, timeout=15) + used = _get_docker_used_host_ports() + assert port in used, f"stopped container's host port {port} should remain in used set" + + def test_returns_set_of_ints(self): + # Sanity: format parsing shouldn't break even when no containers exist with our naming + result = _get_docker_used_host_ports() + assert isinstance(result, set) + for p in result: + assert isinstance(p, int) + assert 0 < p < 65536 + + +class TestFindFreePortAvoidsDockerPorts: + # find_free_port reads _DOCKER_USED_PORTS as-is and does not refresh on + # its own. Production callers (do_port_mapping) refresh upfront — these + # tests mirror that contract by calling refresh_docker_used_ports() first. + + def test_does_not_return_port_held_by_running_container(self, container_holding_port): + _name, blocked_port = container_holding_port + refresh_docker_used_ports() + try: + ports = [] + for _ in range(5): + p = asyncio.run(find_free_port()) + ports.append(p) + assert p != blocked_port, f"find_free_port returned blocked port {p}" + finally: + for p in ports: + release_port(p) + + def test_does_not_return_port_held_by_stopped_container(self, container_holding_port): + # Same guarantee must hold for stopped containers (they will reclaim + # the port on docker start). + name, blocked_port = container_holding_port + subprocess.run(["docker", "kill", name], check=True, capture_output=True, timeout=15) + refresh_docker_used_ports() + try: + ports = [] + for _ in range(5): + p = asyncio.run(find_free_port()) + ports.append(p) + assert p != blocked_port, f"find_free_port returned blocked port {p}" + finally: + for p in ports: + release_port(p) + + +class TestFindFreePortStillRespectsInProcessRegistry: + def test_consecutive_calls_return_distinct_ports(self): + """Same-process concurrent protection (the original behaviour) must still work.""" + ports = [] + try: + for _ in range(5): + p = asyncio.run(find_free_port()) + ports.append(p) + assert len(ports) == len(set(ports)), f"got duplicates: {ports}" + finally: + for p in ports: + release_port(p) + + def test_release_port_makes_it_reusable(self): + p = asyncio.run(find_free_port()) + assert p in _REGISTERED_PORTS + release_port(p) + assert p not in _REGISTERED_PORTS diff --git a/tests/unit/sandbox/test_sandbox_statemachine.py b/tests/unit/sandbox/test_sandbox_statemachine.py index c9c12613b1..9e16d1a0e5 100644 --- a/tests/unit/sandbox/test_sandbox_statemachine.py +++ b/tests/unit/sandbox/test_sandbox_statemachine.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, patch import pytest +from statemachine.exceptions import TransitionNotAllowed from rock.actions.sandbox.response import State from rock.common.constants import StopReason @@ -162,6 +163,18 @@ async def test_stops_operator_and_archives(self, mock_operator, mock_meta_store) mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.MANUAL) mock_meta_store.archive.assert_awaited_once() + @pytest.mark.asyncio + async def test_propagates_reason_to_operator(self, mock_operator, mock_meta_store): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + await sm.send( + "stop", + sandbox_id="sb-1", + operator=mock_operator, + meta_store=mock_meta_store, + reason=StopReason.EXPIRED, + ) + mock_operator.stop.assert_awaited_once_with("sb-1", reason=StopReason.EXPIRED) + @pytest.mark.asyncio async def test_archives_stopped_state(self, mock_operator, mock_meta_store): sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) @@ -192,3 +205,90 @@ async def test_meta_store_none_still_archives(self, mock_operator): sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=store) store.archive.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# restart transitions +# --------------------------------------------------------------------------- + + +_VALID_RESTART_INFO = { + "host_ip": "1.2.3.4", + "spec": { + "container_name": "sb", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, +} + + +class TestRestartTransitions: + def _restart_kwargs(self, meta_store=None): + return dict( + sandbox_id="sb", + operator=AsyncMock(), + meta_store=meta_store or AsyncMock(), + ) + + @pytest.mark.asyncio + async def test_restart_from_stopped_transitions_to_pending(self): + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info=dict(_VALID_RESTART_INFO)) + await sm.send("restart", **self._restart_kwargs()) + assert sm.pending.is_active + + @pytest.mark.asyncio + async def test_restart_from_running_raises(self): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + with pytest.raises(TransitionNotAllowed): + await sm.send("restart", **self._restart_kwargs()) + + @pytest.mark.asyncio + async def test_restart_from_pending_raises(self): + sm = await SandboxStateMachine.from_state_value(State.PENDING, sandbox_info={}) + with pytest.raises(TransitionNotAllowed): + await sm.send("restart", **self._restart_kwargs()) + + +# --------------------------------------------------------------------------- +# on_restart callback +# --------------------------------------------------------------------------- + + +class TestOnRestart: + @pytest.fixture + def mock_meta_store(self): + return AsyncMock() + + async def _send_restart(self, mock_meta_store, sandbox_info=None): + info = sandbox_info if sandbox_info is not None else dict(_VALID_RESTART_INFO) + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info=info) + await sm.send( + "restart", + sandbox_id="sb-1", + operator=AsyncMock(), + meta_store=mock_meta_store, + ) + + @pytest.mark.asyncio + async def test_calls_meta_store_update_not_create(self, mock_meta_store): + await self._send_restart(mock_meta_store) + mock_meta_store.update.assert_awaited_once() + mock_meta_store.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_updates_state_to_pending(self, mock_meta_store): + await self._send_restart(mock_meta_store) + updated_info = mock_meta_store.update.call_args[0][1] + assert updated_info["state"] == State.PENDING + + @pytest.mark.asyncio + async def test_writes_timeout_built_from_spec(self, mock_meta_store): + # auto_clear_time_minutes=30 in spec → make_timeout_info uses 30 + await self._send_restart(mock_meta_store) + mock_meta_store.update_timeout.assert_awaited_once() + sandbox_id, timeout_info = mock_meta_store.update_timeout.call_args[0] + assert sandbox_id == "sb-1" + # SandboxTimeoutHelper.make_timeout_info stores auto_clear_time as the env-var key + assert any("30" == str(v) for v in timeout_info.values()) diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py index 9f5b503101..1a25ca7624 100644 --- a/tests/unit/sandbox/test_sandbox_transitions.py +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -9,9 +9,10 @@ import pytest from rock.actions.sandbox.response import State +from rock.admin.proto.response import SandboxStartResponse from rock.common.constants import StopReason from rock.sandbox.sandbox_manager import SandboxManager -from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError +from rock.sdk.common.exceptions import BadRequestRockError @pytest.fixture @@ -50,13 +51,12 @@ async def mgr(mock_meta_store, mock_operator): # Function to get state machine based on meta_store data — mirrors _get_current_statemachine async def get_current_statemachine(sandbox_id: str) -> SandboxStateMachine | None: + from rock.sandbox.sandbox_statemachine import SandboxStateMachine + info = await mock_meta_store.get(sandbox_id, check_db=True) if info is None: return None - state = info.get("state") - if state is None: - raise InternalServerRockError(f"Sandbox {sandbox_id} exists in store but has no state field") - return await SandboxStateMachine.from_state_value(state, sandbox_info=info) + return await SandboxStateMachine.from_state_value(info.get("state"), sandbox_info=info) m._get_current_statemachine = AsyncMock(side_effect=get_current_statemachine) @@ -127,12 +127,6 @@ async def test_stop_archived_info_has_stopped_state(self, mgr, mock_meta_store, archived_info = mock_meta_store.archive.call_args[0][1] assert archived_info["state"] == State.STOPPED - @pytest.mark.asyncio - async def test_stop_missing_state_field_raises(self, mgr, mock_meta_store): - mock_meta_store.get.return_value = {} - with pytest.raises(InternalServerRockError, match="no state field"): - await mgr.stop("sb-1") - # --------------------------------------------------------------------------- # TestManagerGetStatus @@ -183,3 +177,128 @@ async def test_no_update_if_state_unchanged(self, mgr, mock_meta_store, mock_ope mock_operator.get_status.return_value = {"state": State.RUNNING, "phases": {}, "port_mapping": {}} await mgr.get_status("sb-1") mock_meta_store.update.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# TestManagerRestart +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_docker_config(): + cfg = MagicMock() + cfg.container_name = "sb-1" + cfg.auto_clear_time = 30 + return cfg + + +@pytest.fixture +def mgr_restart(mgr, mock_docker_config): + mgr.deployment_manager = MagicMock() + mgr.deployment_manager.init_config = AsyncMock(return_value=mock_docker_config) + mgr.restart_async = SandboxManager.restart_async.__wrapped__.__get__(mgr) + return mgr + + +class TestManagerRestart: + @pytest.mark.asyncio + async def test_sandbox_not_found_raises(self, mgr_restart, mock_meta_store): + mock_meta_store.get.return_value = None + with pytest.raises(BadRequestRockError, match="not found"): + await mgr_restart.restart_async(MagicMock()) + + @pytest.mark.asyncio + async def test_non_stopped_state_raises(self, mgr_restart, mock_meta_store): + mock_meta_store.get.return_value = {"state": State.RUNNING, "host_ip": "1.2.3.4", "host_name": "w1"} + with pytest.raises(BadRequestRockError): + await mgr_restart.restart_async(MagicMock()) + + @pytest.mark.asyncio + async def test_stopped_success_returns_response(self, mgr_restart, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = { + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "host_name": "worker-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, + } + mock_operator.restart = AsyncMock(return_value={"host_name": "worker-1", "host_ip": "1.2.3.4"}) + result = await mgr_restart.restart_async("sb-1") + assert isinstance(result, SandboxStartResponse) + assert result.sandbox_id == "sb-1" + assert result.host_ip == "1.2.3.4" + + @pytest.mark.asyncio + async def test_stopped_success_calls_operator_restart(self, mgr_restart, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = { + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "host_name": "worker-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, + } + mock_operator.restart = AsyncMock(return_value={"host_name": "worker-1", "host_ip": "1.2.3.4"}) + await mgr_restart.restart_async("sb-1") + mock_operator.restart.assert_awaited_once() + called_config = mock_operator.restart.await_args.args[0] + assert called_config.container_name == "sb-1" + assert called_config.image == "python:3.11" + + @pytest.mark.asyncio + async def test_stopped_success_calls_meta_update(self, mgr_restart, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = { + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "host_name": "worker-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, + } + mock_operator.restart = AsyncMock(return_value={"host_name": "worker-1", "host_ip": "1.2.3.4"}) + await mgr_restart.restart_async("sb-1") + mock_meta_store.update.assert_awaited() + mock_meta_store.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_operator_failure_propagates(self, mgr_restart, mock_meta_store, mock_operator): + mock_meta_store.get.return_value = { + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "host_name": "worker-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, + } + mock_operator.restart = AsyncMock(side_effect=BadRequestRockError("docker start failed")) + with pytest.raises(BadRequestRockError, match="docker start failed"): + await mgr_restart.restart_async("sb-1") From a4679dee897dc7b242a8c318a035449134597d18 Mon Sep 17 00:00:00 2001 From: jiaoliao <38124819+zhongwen666@users.noreply.github.com> Date: Mon, 1 Jun 2026 14:28:45 +0800 Subject: [PATCH 148/226] Feat/sdk support gpu params (#1047) * add release note 120 * Revert "add release note 120" This reverts commit 65a11fd929d9e743c0320664c9599111c6425392. * sdk support gpu * update version --- pyproject.toml | 2 +- rock/sdk/sandbox/client.py | 2 ++ rock/sdk/sandbox/config.py | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1f58c2792c..c7f88a6fca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.8.0" +version = "1.8.3" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index a6029f9415..fc9fd55459 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -171,6 +171,8 @@ async def start(self): "startup_timeout": self.config.startup_timeout, "memory": self.config.memory, "cpus": self.config.cpus, + "num_gpus": self.config.num_gpus, + "accelerator_type": self.config.accelerator_type, "registry_username": self.config.registry_username, "registry_password": self.config.registry_password, "use_kata_runtime": self.config.use_kata_runtime, diff --git a/rock/sdk/sandbox/config.py b/rock/sdk/sandbox/config.py index 4fcf59e030..6ca3fa1c01 100644 --- a/rock/sdk/sandbox/config.py +++ b/rock/sdk/sandbox/config.py @@ -36,6 +36,8 @@ class SandboxConfig(BaseConfig): memory: str = "8g" cpus: float = 2 limit_cpus: float | None = None + num_gpus: float | None = None + accelerator_type: str | None = None user_id: str | None = None experiment_id: str | None = None cluster: str = env_vars.ROCK_DEFAULT_CLUSTER From 33b173494a58e05c2871012f200757df607d4d6b Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 1 Jun 2026 10:29:36 +0800 Subject: [PATCH 149/226] fix(sandbox): always write stop_time on stop() even when start_time absent (start-failed sandboxes) --- rock/sandbox/sandbox_statemachine.py | 7 ++++++- tests/unit/sandbox/test_sandbox_statemachine.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/rock/sandbox/sandbox_statemachine.py b/rock/sandbox/sandbox_statemachine.py index bf215c05b3..e0285f5862 100644 --- a/rock/sandbox/sandbox_statemachine.py +++ b/rock/sandbox/sandbox_statemachine.py @@ -67,8 +67,13 @@ async def on_stop(self, sandbox_id: str, operator, meta_store, reason: StopReaso sandbox_info["sandbox_id"] = sandbox_id sandbox_info["state"] = RockState.STOPPED + # Always record stop_time — sandboxes that never started (e.g. image + # pull / docker run failed before sandbox_actor wrote start_time) also + # need this for downstream consumers like SandboxLogArchiveTask. The + # billing call below stays gated on start_time because billing is + # only meaningful for actually-started sandboxes. + sandbox_info["stop_time"] = get_iso8601_timestamp() if sandbox_info.get("start_time"): - sandbox_info["stop_time"] = get_iso8601_timestamp() log_billing_info(sandbox_info=sandbox_info) try: diff --git a/tests/unit/sandbox/test_sandbox_statemachine.py b/tests/unit/sandbox/test_sandbox_statemachine.py index 9e16d1a0e5..fc46f1f106 100644 --- a/tests/unit/sandbox/test_sandbox_statemachine.py +++ b/tests/unit/sandbox/test_sandbox_statemachine.py @@ -206,6 +206,21 @@ async def test_meta_store_none_still_archives(self, mock_operator): await sm.send("stop", sandbox_id="sb-1", operator=mock_operator, meta_store=store) store.archive.assert_awaited_once() + @pytest.mark.asyncio + async def test_stop_time_always_written_even_when_start_failed(self, mock_operator, mock_meta_store): + """REGRESSION: sandboxes that fail before sandbox_actor writes start_time + (image pull / docker run errors) still get stop_time. Without this, + SandboxLogArchiveTask can't age them and their log dirs leak forever. + Billing stays gated on start_time (billing only meaningful for started).""" + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) # no start_time + with patch("rock.sandbox.sandbox_statemachine.log_billing_info") as mock_billing: + await sm.send("stop", sandbox_id="sb-failed", operator=mock_operator, meta_store=mock_meta_store) + + archived = mock_meta_store.archive.call_args[0][1] + assert archived["state"] == State.STOPPED + assert archived.get("stop_time"), "stop_time must be set even when start_time absent" + mock_billing.assert_not_called() # no billing when start_time absent + # --------------------------------------------------------------------------- # restart transitions From 31839ee9a477e2069ecb3d5133ba6406722894a5 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 26 May 2026 15:00:34 +0800 Subject: [PATCH 150/226] feat(scheduler): clean session_latest/logs PID-aware + logs/old in ray task --- .../scheduler/tasks/ray_log_cleanup_task.py | 159 +++++++++++--- .../scheduler/test_ray_log_cleanup_task.py | 204 +++++++++++++++++- 2 files changed, 329 insertions(+), 34 deletions(-) diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index 264d65ee3a..9b1626ae0f 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -1,4 +1,17 @@ -"""Drop stale /data/tmp/ray/session_* dirs on each worker.""" +"""Clean up Ray temp dir on each worker. + +Three layers of cleanup, run together in one shell pipeline: + 1. Stale full ``session__`` dirs (excluding the live session). + 2. Inside ``session_latest/logs/``: per-file PID-aware cleanup + (delete files whose PID is no longer alive) plus an mtime fallback + for non-PID, non-daemon files. + 3. Inside ``session_latest/logs/old/``: time-based cleanup of Ray's + own rotation backups. + +Daemon-written files (raylet*, gcs_server*, runtime_env_agent*, etc.) +are NEVER deleted while the session is alive — Ray holds open fds, so +removing the inode wouldn't free disk and breaks log following. +""" import textwrap @@ -12,18 +25,23 @@ class RayLogCleanupTask(BaseTask): - """Drop /data/tmp/ray/session_* dirs that are NOT the live session. + """Clean Ray temp dir: stale session dirs + active session per-file cleanup. Ray restarts (cluster up/down, head failover) leave dozens of - session__ dirs behind. The currently active one is - symlinked as `session_latest`; we resolve that link and skip its target. - Sessions younger than `min_age_hours` are also kept as a buffer against - a stale symlink. + ``session__`` dirs behind. The currently active one is + symlinked as ``session_latest``; we resolve that link and skip its target + when removing stale session dirs. + + Inside the live ``session_latest/logs/``, Ray accumulates per-worker + logs (named after the worker PID) that are never reaped after the + worker exits, plus its own rotation backups under ``logs/old/`` that + Ray never cleans. Without intervention, file count grows to tens or + hundreds of thousands on long-running clusters. NOTE: This is the WORKER side. The ray-head's /data/tmp/ray is cleaned by a daily cron baked into the head Dockerfile (rock-internal repo); rocklet is not deployed on the head and the worker scheduler does not - reach it. + reach it. The cron script mirrors this task's shell pipeline. """ def __init__( @@ -31,6 +49,8 @@ def __init__( interval_seconds: int = 86400, ray_temp_dir: str = "/data/tmp/ray", min_age_hours: int = 24, + live_log_keep_days: int = 7, + old_logs_keep_hours: int = 24, ): """ Args: @@ -38,6 +58,11 @@ def __init__( ray_temp_dir: Ray's --temp-dir, default /data/tmp/ray. min_age_hours: Only delete session dirs whose mtime is older than this AND that are not session_latest. Default 24h. + live_log_keep_days: Mtime threshold for non-PID, non-daemon files + in session_latest/logs/. Default 7 days. + old_logs_keep_hours: Mtime threshold for files under + session_latest/logs/old/ (Ray's own rotation backups). + Default 24 hours. """ super().__init__( type="ray_log_cleanup", @@ -46,8 +71,14 @@ def __init__( ) if min_age_hours < 1: raise ValueError(f"ray_log_cleanup.min_age_hours must be >= 1, got {min_age_hours}") + if live_log_keep_days < 1: + raise ValueError(f"ray_log_cleanup.live_log_keep_days must be >= 1, got {live_log_keep_days}") + if old_logs_keep_hours < 1: + raise ValueError(f"ray_log_cleanup.old_logs_keep_hours must be >= 1, got {old_logs_keep_hours}") self.ray_temp_dir = ray_temp_dir.rstrip("/") self.min_age_hours = min_age_hours + self.live_log_keep_days = live_log_keep_days + self.old_logs_keep_hours = old_logs_keep_hours @classmethod def from_config(cls, task_config) -> "RayLogCleanupTask": @@ -55,44 +86,116 @@ def from_config(cls, task_config) -> "RayLogCleanupTask": interval_seconds=task_config.interval_seconds, ray_temp_dir=task_config.params.get("ray_temp_dir", "/data/tmp/ray"), min_age_hours=task_config.params.get("min_age_hours", 24), + live_log_keep_days=task_config.params.get("live_log_keep_days", 7), + old_logs_keep_hours=task_config.params.get("old_logs_keep_hours", 24), ) async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: ray_dir = self.ray_temp_dir - max_age_min = self.min_age_hours * 60 - # Resolve session_latest -> live basename, then list session_* dirs - # older than threshold and rm -rf those that are not the live one. + session_age_min = self.min_age_hours * 60 + live_age_min = self.live_log_keep_days * 24 * 60 + old_age_min = self.old_logs_keep_hours * 60 + + # 3-stage shell pipeline: + # PART 1 — drop stale full session__ dirs + # PART 2 — session_latest/logs/ per-file: PID-aware + mtime fallback + # PART 3 — session_latest/logs/old/ time-based + # # textwrap.dedent strips common leading whitespace so source can be # indented for readability without polluting the emitted shell. command = textwrap.dedent( f"""\ - if [ -d "{ray_dir}" ]; then - LIVE=$(readlink "{ray_dir}/session_latest" 2>/dev/null | xargs -I{{}} basename {{}} 2>/dev/null) - echo "live_session=${{LIVE:-}}" - find "{ray_dir}" -maxdepth 1 -type d -name "session_*" \\ - ! -name "session_latest" -mmin +{max_age_min} \\ - | while read -r d; do - bn=$(basename "$d") - if [ "$bn" != "$LIVE" ]; then - rm -rf "$d" && echo "removed=$bn" + set +e + if [ ! -d "{ray_dir}" ]; then + echo "ray_temp_dir_not_found" + exit 0 + fi + + # PART 1: stale full session__ dirs (not session_latest) + LIVE=$(readlink "{ray_dir}/session_latest" 2>/dev/null | xargs -I{{}} basename {{}} 2>/dev/null) + echo "live_session=${{LIVE:-}}" + find "{ray_dir}" -maxdepth 1 -type d -name "session_*" \\ + ! -name "session_latest" -mmin +{session_age_min} \\ + | while read -r d; do + bn=$(basename "$d") + if [ "$bn" != "$LIVE" ]; then + rm -rf "$d" && echo "removed=$bn" + fi + done + + LOGS="{ray_dir}/session_latest/logs" + if [ -d "$LOGS" ]; then + # PART 2a: PID-aware — files matching *[_-].{{log,err,out}} + # Probe `kill -0 `; if PID is dead, remove. Only files older + # than 60 minutes are considered, to avoid racing with new + # worker startups still writing their first log line. + find "$LOGS" -maxdepth 1 -type f -mmin +60 \\ + -regextype posix-extended \\ + -regex '.*[_-][0-9]+\\.(log|err|out)$' \\ + | while read -r f; do + bn=$(basename "$f") + pid=$(echo "$bn" | grep -oE '[_-][0-9]+\\.(log|err|out)$' | grep -oE '[0-9]+' | head -1) + [ -z "$pid" ] && continue + if ! kill -0 "$pid" 2>/dev/null; then + rm -f "$f" && echo "removed_dead_pid_log=$bn" fi done - echo "ray_log_cleanup_done" - else - echo "ray_temp_dir_not_found" - fi""" + + # PART 2b: non-PID, non-daemon stale files older than + # {self.live_log_keep_days} days. Daemon files (raylet*, + # gcs_server*, runtime_env_agent*, dashboard*, monitor*, + # log_monitor*) are NEVER deleted while session is alive — + # Ray holds open fds, removal wouldn't free disk. + find "$LOGS" -maxdepth 1 -type f -mmin +{live_age_min} \\ + -regextype posix-extended \\ + ! -regex '.*[_-][0-9]+\\.(log|err|out)$' \\ + ! -name 'raylet*' \\ + ! -name 'gcs_server*' \\ + ! -name 'runtime_env_agent*' \\ + ! -name 'dashboard*' \\ + ! -name 'monitor*' \\ + ! -name 'log_monitor*' \\ + | while read -r f; do + rm -f "$f" && echo "removed_stale_file=$(basename "$f")" + done + fi + + # PART 3: session_latest/logs/old/ — Ray's own rotation backups + OLD="$LOGS/old" + if [ -d "$OLD" ]; then + find "$OLD" -type f -mmin +{old_age_min} \\ + | while read -r f; do + rm -f "$f" && echo "removed_old=$(basename "$f")" + done + fi + + echo "ray_log_cleanup_done" + """ ) result = await runtime.execute(Command(command=command, shell=True, check=False)) output = (result.stdout or "").strip() - removed = [line.split("=", 1)[1] for line in output.splitlines() if line.startswith("removed=")] + + # Parse per-category removal counts from output. + # `removed=` retained for backward compat (PART 1 session dirs). + removed_sessions = [line.split("=", 1)[1] for line in output.splitlines() if line.startswith("removed=")] + removed_dead_pid = sum(1 for line in output.splitlines() if line.startswith("removed_dead_pid_log=")) + removed_stale = sum(1 for line in output.splitlines() if line.startswith("removed_stale_file=")) + removed_old = sum(1 for line in output.splitlines() if line.startswith("removed_old=")) + logger.info( f"[{self.type}] [{runtime._config.host}] ray_log_cleanup done: " - f"removed={len(removed)} sessions, output_head={output[:300]}" + f"sessions={len(removed_sessions)}, dead_pid={removed_dead_pid}, " + f"stale={removed_stale}, old={removed_old}, output_head={output[:300]}" ) return { "status": TaskStatusEnum.SUCCESS, "exit_code": result.exit_code, - "removed_count": len(removed), - "removed_sessions": removed, - "output_head": output[:1000], + # Backward-compatible fields (session-level removal): + "removed_count": len(removed_sessions), + "removed_sessions": removed_sessions, + # New per-category counters: + "removed_dead_pid_count": removed_dead_pid, + "removed_stale_count": removed_stale, + "removed_old_count": removed_old, + "output_head": output[:1500], } diff --git a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py index 8f2991fd8a..c37823f553 100644 --- a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py @@ -1,4 +1,4 @@ -"""Tests for RayLogCleanupTask.""" +"""Tests for RayLogCleanupTask (3-stage: stale sessions + logs/ PID-aware + logs/old/).""" from unittest.mock import AsyncMock @@ -27,12 +27,19 @@ def _runtime(stdout="ray_log_cleanup_done", exit_code=0): return rt +# --------------------------------------------------------------------------- +# Init / from_config +# --------------------------------------------------------------------------- + + class TestInit: def test_default(self): task = RayLogCleanupTask() assert task.type == "ray_log_cleanup" assert task.ray_temp_dir == "/data/tmp/ray" assert task.min_age_hours == 24 + assert task.live_log_keep_days == 7 + assert task.old_logs_keep_hours == 24 def test_strips_trailing_slash(self): task = RayLogCleanupTask(ray_temp_dir="/data/ray/") @@ -42,27 +49,54 @@ def test_rejects_min_age_below_one(self): with pytest.raises(ValueError, match="min_age_hours must be >= 1"): RayLogCleanupTask(min_age_hours=0) + def test_rejects_live_log_keep_days_below_one(self): + with pytest.raises(ValueError, match="live_log_keep_days must be >= 1"): + RayLogCleanupTask(live_log_keep_days=0) + + def test_rejects_old_logs_keep_hours_below_one(self): + with pytest.raises(ValueError, match="old_logs_keep_hours must be >= 1"): + RayLogCleanupTask(old_logs_keep_hours=0) + + def test_custom_thresholds(self): + task = RayLogCleanupTask(live_log_keep_days=3, old_logs_keep_hours=12) + assert task.live_log_keep_days == 3 + assert task.old_logs_keep_hours == 12 + class TestFromConfig: def test_from_config_defaults(self): task = RayLogCleanupTask.from_config(_FakeTaskConfig()) assert task.ray_temp_dir == "/data/tmp/ray" assert task.min_age_hours == 24 + assert task.live_log_keep_days == 7 + assert task.old_logs_keep_hours == 24 def test_from_config_custom(self): cfg = _FakeTaskConfig( - params={"ray_temp_dir": "/data/ray", "min_age_hours": 48}, + params={ + "ray_temp_dir": "/data/ray", + "min_age_hours": 48, + "live_log_keep_days": 14, + "old_logs_keep_hours": 6, + }, interval_seconds=3600, ) task = RayLogCleanupTask.from_config(cfg) assert task.ray_temp_dir == "/data/ray" assert task.min_age_hours == 48 + assert task.live_log_keep_days == 14 + assert task.old_logs_keep_hours == 6 assert task.interval_seconds == 3600 -class TestRunAction: +# --------------------------------------------------------------------------- +# Shell command shape — verify the 3 stages and their parameters +# --------------------------------------------------------------------------- + + +class TestCommandShape: @pytest.mark.asyncio - async def test_command_skips_session_latest(self): + async def test_part1_skips_session_latest(self): task = RayLogCleanupTask() runtime = _runtime() await task.run_action(runtime) @@ -75,7 +109,7 @@ async def test_command_skips_session_latest(self): assert 'name "session_*"' in cmd @pytest.mark.asyncio - async def test_command_uses_min_age_in_minutes(self): + async def test_part1_uses_min_age_in_minutes(self): task = RayLogCleanupTask(min_age_hours=48) runtime = _runtime() await task.run_action(runtime) @@ -92,9 +126,78 @@ async def test_command_respects_custom_temp_dir(self): cmd = runtime.execute.await_args.args[0].command assert '"/data/ray"' in cmd + assert "/data/ray/session_latest/logs" in cmd + + @pytest.mark.asyncio + async def test_part2a_uses_kill_zero_pid_probe(self): + """PART 2a must probe PID with `kill -0` (no-signal liveness check).""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + assert "kill -0" in cmd + # PID regex must match Ray's worker file naming + assert "[_-][0-9]+" in cmd + # Only files older than 60 min are candidates (race-window guard) + assert "-mmin +60" in cmd + + @pytest.mark.asyncio + async def test_part2b_uses_live_log_keep_days(self): + """PART 2b stale-file mtime threshold = live_log_keep_days * 24 * 60 minutes.""" + task = RayLogCleanupTask(live_log_keep_days=7) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # 7d * 24h * 60min = 10080 + assert "-mmin +10080" in cmd @pytest.mark.asyncio - async def test_extracts_removed_count_from_output(self): + async def test_part2b_daemon_whitelist_protected(self): + """raylet*, gcs_server*, runtime_env_agent*, dashboard*, monitor*, + log_monitor* must all be excluded by name (regression guard — Ray + holds fds; deletion wouldn't free disk).""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + for daemon in ("raylet", "gcs_server", "runtime_env_agent", "dashboard", "monitor", "log_monitor"): + assert f"! -name '{daemon}*'" in cmd, f"daemon whitelist missing: {daemon}" + + @pytest.mark.asyncio + async def test_part3_uses_old_logs_keep_hours(self): + """PART 3 old-dir mtime threshold = old_logs_keep_hours * 60 minutes.""" + task = RayLogCleanupTask(old_logs_keep_hours=24) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # 24h * 60 = 1440 + assert "-mmin +1440" in cmd + assert "session_latest/logs/old" in cmd + + @pytest.mark.asyncio + async def test_part3_uses_old_logs_keep_hours_custom(self): + task = RayLogCleanupTask(old_logs_keep_hours=6) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # 6h * 60 = 360 + assert "-mmin +360" in cmd + + +# --------------------------------------------------------------------------- +# Output parsing — per-category counters +# --------------------------------------------------------------------------- + + +class TestOutputParsing: + @pytest.mark.asyncio + async def test_extracts_part1_removed_sessions(self): + """PART 1 session removals reported via `removed=` (backward compat).""" stdout = ( "live_session=session_2026_03_01_xyz_111\n" "removed=session_2026_02_15_aaa_222\n" @@ -110,6 +213,62 @@ async def test_extracts_removed_count_from_output(self): assert "session_2026_02_15_aaa_222" in result["removed_sessions"] assert "session_2026_02_20_bbb_333" in result["removed_sessions"] + @pytest.mark.asyncio + async def test_extracts_part2a_dead_pid_count(self): + stdout = ( + "live_session=session_xxx\n" + "removed_dead_pid_log=python-core-worker-aaaa_12345.log\n" + "removed_dead_pid_log=worker-bbbb-c205-67890.err\n" + "removed_dead_pid_log=worker-bbbb-c205-67890.out\n" + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_dead_pid_count"] == 3 + + @pytest.mark.asyncio + async def test_extracts_part2b_stale_count(self): + stdout = ( + "live_session=session_xxx\n" + "removed_stale_file=runtime_env_setup-60010000.log\n" + "removed_stale_file=runtime_env_setup-5c010000.log\n" + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_stale_count"] == 2 + + @pytest.mark.asyncio + async def test_extracts_part3_old_count(self): + stdout = ( + "live_session=session_xxx\n" + "removed_old=python-core-worker-aaa.log.1\n" + "removed_old=python-core-worker-aaa.log.2\n" + "removed_old=raylet.out.1\n" + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_old_count"] == 3 + + @pytest.mark.asyncio + async def test_all_counters_zero_when_nothing_removed(self): + stdout = "live_session=session_xxx\nray_log_cleanup_done" + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_count"] == 0 + assert result["removed_dead_pid_count"] == 0 + assert result["removed_stale_count"] == 0 + assert result["removed_old_count"] == 0 + @pytest.mark.asyncio async def test_handles_missing_ray_dir(self): task = RayLogCleanupTask() @@ -118,3 +277,36 @@ async def test_handles_missing_ray_dir(self): result = await task.run_action(runtime) assert result["status"] == TaskStatusEnum.SUCCESS assert result["removed_count"] == 0 + assert result["removed_dead_pid_count"] == 0 + assert result["removed_stale_count"] == 0 + assert result["removed_old_count"] == 0 + + +# --------------------------------------------------------------------------- +# Regression guards — must-have command properties +# --------------------------------------------------------------------------- + + +class TestRegressionGuards: + @pytest.mark.asyncio + async def test_command_does_not_recurse_into_session_dirs_at_top_level(self): + """Top-level walk MUST keep `-maxdepth 1`; we explicitly DO recurse + into `session_latest/logs/old/` in PART 3 but never into other + `session__` internals (only whole-dir rm -rf for those).""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # PART 1 stale-session walk: must have maxdepth 1 + assert '-maxdepth 1 -type d -name "session_*"' in cmd + + @pytest.mark.asyncio + async def test_command_uses_session_latest_logs_path(self): + """PART 2/3 must scope to session_latest/logs explicitly, not arbitrary session dirs.""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + assert "session_latest/logs" in cmd From db75d6cc73d56e5b583a33df5d016d76e2d46b5d Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 26 May 2026 20:04:09 +0800 Subject: [PATCH 151/226] fix(scheduler): protect agent-* and other daemon files from PART 2a PID probe --- .../scheduler/tasks/ray_log_cleanup_task.py | 19 +++++++++- .../scheduler/test_ray_log_cleanup_task.py | 37 ++++++++++++++++++- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index 9b1626ae0f..a8cc037ef2 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -129,9 +129,23 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: # Probe `kill -0 `; if PID is dead, remove. Only files older # than 60 minutes are considered, to avoid racing with new # worker startups still writing their first log line. + # + # Daemon files are excluded by name FIRST (same whitelist as + # PART 2b). Without this guard, names like `agent-.err` — + # where is a Ray-generated agent identifier, NOT a PID — + # match the PID regex; kill -0 fails because exceeds + # the Linux PID range, so the file gets wrongly removed even + # though Ray's runtime env agent is still writing to it. find "$LOGS" -maxdepth 1 -type f -mmin +60 \\ -regextype posix-extended \\ -regex '.*[_-][0-9]+\\.(log|err|out)$' \\ + ! -name 'raylet*' \\ + ! -name 'gcs_server*' \\ + ! -name 'runtime_env_agent*' \\ + ! -name 'dashboard*' \\ + ! -name 'monitor*' \\ + ! -name 'log_monitor*' \\ + ! -name 'agent-*' \\ | while read -r f; do bn=$(basename "$f") pid=$(echo "$bn" | grep -oE '[_-][0-9]+\\.(log|err|out)$' | grep -oE '[0-9]+' | head -1) @@ -144,8 +158,8 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: # PART 2b: non-PID, non-daemon stale files older than # {self.live_log_keep_days} days. Daemon files (raylet*, # gcs_server*, runtime_env_agent*, dashboard*, monitor*, - # log_monitor*) are NEVER deleted while session is alive — - # Ray holds open fds, removal wouldn't free disk. + # log_monitor*, agent-*) are NEVER deleted while session is + # alive — Ray holds open fds, removal wouldn't free disk. find "$LOGS" -maxdepth 1 -type f -mmin +{live_age_min} \\ -regextype posix-extended \\ ! -regex '.*[_-][0-9]+\\.(log|err|out)$' \\ @@ -155,6 +169,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: ! -name 'dashboard*' \\ ! -name 'monitor*' \\ ! -name 'log_monitor*' \\ + ! -name 'agent-*' \\ | while read -r f; do rm -f "$f" && echo "removed_stale_file=$(basename "$f")" done diff --git a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py index c37823f553..4b2b148828 100644 --- a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py @@ -156,8 +156,8 @@ async def test_part2b_uses_live_log_keep_days(self): @pytest.mark.asyncio async def test_part2b_daemon_whitelist_protected(self): """raylet*, gcs_server*, runtime_env_agent*, dashboard*, monitor*, - log_monitor* must all be excluded by name (regression guard — Ray - holds fds; deletion wouldn't free disk).""" + log_monitor*, agent-* must all be excluded by name (regression guard + — Ray holds fds; deletion wouldn't free disk).""" task = RayLogCleanupTask() runtime = _runtime() await task.run_action(runtime) @@ -165,6 +165,39 @@ async def test_part2b_daemon_whitelist_protected(self): cmd = runtime.execute.await_args.args[0].command for daemon in ("raylet", "gcs_server", "runtime_env_agent", "dashboard", "monitor", "log_monitor"): assert f"! -name '{daemon}*'" in cmd, f"daemon whitelist missing: {daemon}" + assert "! -name 'agent-*'" in cmd, "daemon whitelist missing: agent-*" + + @pytest.mark.asyncio + async def test_part2a_daemon_whitelist_protects_agent_files(self): + """CORE REGRESSION (task #25): PART 2a (PID-aware) MUST also exclude + daemon-named files BEFORE the kill -0 probe. Without this guard, + `agent-.err` matches the PID regex; kill -0 fails because + is a Ray-generated agent identifier (NOT a Linux PID) that + exceeds the PID range, so the file gets wrongly removed even though + Ray's runtime env agent is still writing to it.""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # PART 2a section starts with `-mmin +60` (race-window guard). + # Find the segment from "+60" up to the next pipe (end of find expr). + idx = cmd.find("-mmin +60") + assert idx >= 0, "PART 2a -mmin +60 marker not found" + part2a = cmd[idx : cmd.find("| while read -r f", idx)] + # All daemon whitelist names must appear in PART 2a's find filter, + # not just PART 2b. Otherwise agent-* / runtime_env_agent* etc. get + # wrongly probed and removed. + for daemon in ( + "raylet", + "gcs_server", + "runtime_env_agent", + "dashboard", + "monitor", + "log_monitor", + "agent-", + ): + assert f"! -name '{daemon}*'" in part2a, f"PART 2a missing daemon whitelist: {daemon}" @pytest.mark.asyncio async def test_part3_uses_old_logs_keep_hours(self): From 3c05f1d8319d54dac5ab66f10e87972cc5489bd2 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 11:28:56 +0800 Subject: [PATCH 152/226] feat(scheduler): clean runtime_env_setup-* in PART 2c (covers hex suffix that PART 2b waited 7d on) --- .../scheduler/tasks/ray_log_cleanup_task.py | 43 +++++++++++++-- .../scheduler/test_ray_log_cleanup_task.py | 55 +++++++++++++++++++ 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index a8cc037ef2..ca8452eda0 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -51,6 +51,7 @@ def __init__( min_age_hours: int = 24, live_log_keep_days: int = 7, old_logs_keep_hours: int = 24, + setup_log_keep_minutes: int = 60, ): """ Args: @@ -63,6 +64,11 @@ def __init__( old_logs_keep_hours: Mtime threshold for files under session_latest/logs/old/ (Ray's own rotation backups). Default 24 hours. + setup_log_keep_minutes: Mtime threshold for runtime_env_setup-* + files. These are one-shot Ray runtime-env materialisation + scripts whose trailing token is a Ray job id (digit OR hex), + NOT a Linux PID — PART 2a's PID probe is meaningless for them. + Default 60min (matches PART 2a race-window guard). """ super().__init__( type="ray_log_cleanup", @@ -75,10 +81,13 @@ def __init__( raise ValueError(f"ray_log_cleanup.live_log_keep_days must be >= 1, got {live_log_keep_days}") if old_logs_keep_hours < 1: raise ValueError(f"ray_log_cleanup.old_logs_keep_hours must be >= 1, got {old_logs_keep_hours}") + if setup_log_keep_minutes < 1: + raise ValueError(f"ray_log_cleanup.setup_log_keep_minutes must be >= 1, got {setup_log_keep_minutes}") self.ray_temp_dir = ray_temp_dir.rstrip("/") self.min_age_hours = min_age_hours self.live_log_keep_days = live_log_keep_days self.old_logs_keep_hours = old_logs_keep_hours + self.setup_log_keep_minutes = setup_log_keep_minutes @classmethod def from_config(cls, task_config) -> "RayLogCleanupTask": @@ -88,6 +97,7 @@ def from_config(cls, task_config) -> "RayLogCleanupTask": min_age_hours=task_config.params.get("min_age_hours", 24), live_log_keep_days=task_config.params.get("live_log_keep_days", 7), old_logs_keep_hours=task_config.params.get("old_logs_keep_hours", 24), + setup_log_keep_minutes=task_config.params.get("setup_log_keep_minutes", 60), ) async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: @@ -95,11 +105,14 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: session_age_min = self.min_age_hours * 60 live_age_min = self.live_log_keep_days * 24 * 60 old_age_min = self.old_logs_keep_hours * 60 - - # 3-stage shell pipeline: - # PART 1 — drop stale full session__ dirs - # PART 2 — session_latest/logs/ per-file: PID-aware + mtime fallback - # PART 3 — session_latest/logs/old/ time-based + setup_age_min = self.setup_log_keep_minutes + + # 4-stage shell pipeline: + # PART 1 — drop stale full session__ dirs + # PART 2a — session_latest/logs/ PID-aware (worker/runtime_env_setup with digit pid) + # PART 2c — session_latest/logs/ runtime_env_setup-* mtime-only (covers hex variant) + # PART 2b — session_latest/logs/ non-PID, non-daemon stale (long tail) + # PART 3 — session_latest/logs/old/ time-based # # textwrap.dedent strips common leading whitespace so source can be # indented for readability without polluting the emitted shell. @@ -155,6 +168,21 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: fi done + # PART 2c: runtime_env_setup-* — one-shot Ray runtime-env + # materialisation scripts. Trailing token is a Ray job id, + # NOT a PID; comes in two flavours: + # - pure digits: `runtime_env_setup-31050000.log` — accidentally + # removed by PART 2a (digit > pid_max → kill -0 fails) + # - hex: `runtime_env_setup-f4060000.log` — skipped by PART 2a + # (regex needs digits) and waited 7d for PART 2b + # Same retention window as PART 2a ({setup_age_min}min default) + # to avoid races with in-flight setups while clearing the bulk. + find "$LOGS" -maxdepth 1 -type f -mmin +{setup_age_min} \\ + -name 'runtime_env_setup-*' \\ + | while read -r f; do + rm -f "$f" && echo "removed_setup=$(basename "$f")" + done + # PART 2b: non-PID, non-daemon stale files older than # {self.live_log_keep_days} days. Daemon files (raylet*, # gcs_server*, runtime_env_agent*, dashboard*, monitor*, @@ -194,13 +222,15 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: # `removed=` retained for backward compat (PART 1 session dirs). removed_sessions = [line.split("=", 1)[1] for line in output.splitlines() if line.startswith("removed=")] removed_dead_pid = sum(1 for line in output.splitlines() if line.startswith("removed_dead_pid_log=")) + removed_setup = sum(1 for line in output.splitlines() if line.startswith("removed_setup=")) removed_stale = sum(1 for line in output.splitlines() if line.startswith("removed_stale_file=")) removed_old = sum(1 for line in output.splitlines() if line.startswith("removed_old=")) logger.info( f"[{self.type}] [{runtime._config.host}] ray_log_cleanup done: " f"sessions={len(removed_sessions)}, dead_pid={removed_dead_pid}, " - f"stale={removed_stale}, old={removed_old}, output_head={output[:300]}" + f"setup={removed_setup}, stale={removed_stale}, old={removed_old}, " + f"output_head={output[:300]}" ) return { "status": TaskStatusEnum.SUCCESS, @@ -210,6 +240,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: "removed_sessions": removed_sessions, # New per-category counters: "removed_dead_pid_count": removed_dead_pid, + "removed_setup_count": removed_setup, "removed_stale_count": removed_stale, "removed_old_count": removed_old, "output_head": output[:1500], diff --git a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py index 4b2b148828..13d2de537d 100644 --- a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py @@ -40,6 +40,7 @@ def test_default(self): assert task.min_age_hours == 24 assert task.live_log_keep_days == 7 assert task.old_logs_keep_hours == 24 + assert task.setup_log_keep_minutes == 60 def test_strips_trailing_slash(self): task = RayLogCleanupTask(ray_temp_dir="/data/ray/") @@ -57,6 +58,10 @@ def test_rejects_old_logs_keep_hours_below_one(self): with pytest.raises(ValueError, match="old_logs_keep_hours must be >= 1"): RayLogCleanupTask(old_logs_keep_hours=0) + def test_rejects_setup_log_keep_minutes_below_one(self): + with pytest.raises(ValueError, match="setup_log_keep_minutes must be >= 1"): + RayLogCleanupTask(setup_log_keep_minutes=0) + def test_custom_thresholds(self): task = RayLogCleanupTask(live_log_keep_days=3, old_logs_keep_hours=12) assert task.live_log_keep_days == 3 @@ -70,6 +75,7 @@ def test_from_config_defaults(self): assert task.min_age_hours == 24 assert task.live_log_keep_days == 7 assert task.old_logs_keep_hours == 24 + assert task.setup_log_keep_minutes == 60 def test_from_config_custom(self): cfg = _FakeTaskConfig( @@ -78,6 +84,7 @@ def test_from_config_custom(self): "min_age_hours": 48, "live_log_keep_days": 14, "old_logs_keep_hours": 6, + "setup_log_keep_minutes": 30, }, interval_seconds=3600, ) @@ -86,6 +93,7 @@ def test_from_config_custom(self): assert task.min_age_hours == 48 assert task.live_log_keep_days == 14 assert task.old_logs_keep_hours == 6 + assert task.setup_log_keep_minutes == 30 assert task.interval_seconds == 3600 @@ -199,6 +207,33 @@ async def test_part2a_daemon_whitelist_protects_agent_files(self): ): assert f"! -name '{daemon}*'" in part2a, f"PART 2a missing daemon whitelist: {daemon}" + @pytest.mark.asyncio + async def test_part2c_uses_setup_log_keep_minutes(self): + """PART 2c (runtime_env_setup-*) mtime threshold = setup_log_keep_minutes.""" + task = RayLogCleanupTask(setup_log_keep_minutes=30) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # The first textual occurrence is the comment block; search for the + # actual filter line "-name 'runtime_env_setup-*'" to skip past it. + idx = cmd.find("-name 'runtime_env_setup-*'") + assert idx >= 0, "PART 2c -name 'runtime_env_setup-*' filter not found" + window = cmd[max(0, idx - 150):idx] + assert "-mmin +30" in window, f"PART 2c missing -mmin +30 in nearby window: {window!r}" + + @pytest.mark.asyncio + async def test_part2c_targets_runtime_env_setup_glob(self): + """PART 2c must use -name 'runtime_env_setup-*' (covers both digit and hex suffix).""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + assert "-name 'runtime_env_setup-*'" in cmd + # Emits removed_setup= marker so output parser can count it + assert "removed_setup=" in cmd + @pytest.mark.asyncio async def test_part3_uses_old_logs_keep_hours(self): """PART 3 old-dir mtime threshold = old_logs_keep_hours * 60 minutes.""" @@ -275,6 +310,25 @@ async def test_extracts_part2b_stale_count(self): result = await task.run_action(runtime) assert result["removed_stale_count"] == 2 + @pytest.mark.asyncio + async def test_extracts_part2c_setup_count_digit_and_hex(self): + """REGRESSION: PART 2c must count both digit-suffix AND hex-suffix + runtime_env_setup files. Pre-fix, only the digit variant was cleaned + (accidentally by PART 2a's pid_max overflow); hex variant waited 7d + for PART 2b.""" + stdout = ( + "live_session=session_xxx\n" + "removed_setup=runtime_env_setup-31050000.log\n" # pure digits + "removed_setup=runtime_env_setup-f4060000.log\n" # hex — the main fix + "removed_setup=runtime_env_setup-b6000000.log\n" # hex + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_setup_count"] == 3 + @pytest.mark.asyncio async def test_extracts_part3_old_count(self): stdout = ( @@ -299,6 +353,7 @@ async def test_all_counters_zero_when_nothing_removed(self): result = await task.run_action(runtime) assert result["removed_count"] == 0 assert result["removed_dead_pid_count"] == 0 + assert result["removed_setup_count"] == 0 assert result["removed_stale_count"] == 0 assert result["removed_old_count"] == 0 From 011e9a27be7481dddff13dd92c2851dd922d608d Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 1 Jun 2026 11:39:59 +0800 Subject: [PATCH 153/226] feat(scheduler): add PART 2d rotated daemon log cleanup (raylet.N.out, gcs_server.N.err, etc.) --- .../scheduler/tasks/ray_log_cleanup_task.py | 34 ++++++++- .../scheduler/test_ray_log_cleanup_task.py | 73 +++++++++++++++++++ 2 files changed, 105 insertions(+), 2 deletions(-) diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index ca8452eda0..1c151a701a 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -52,6 +52,7 @@ def __init__( live_log_keep_days: int = 7, old_logs_keep_hours: int = 24, setup_log_keep_minutes: int = 60, + rotated_daemon_keep_hours: int = 24, ): """ Args: @@ -69,6 +70,11 @@ def __init__( scripts whose trailing token is a Ray job id (digit OR hex), NOT a Linux PID — PART 2a's PID probe is meaningless for them. Default 60min (matches PART 2a race-window guard). + rotated_daemon_keep_hours: Mtime threshold for rotated daemon log + files (raylet.N.out, gcs_server.N.err, etc.). Ray performs + in-place rotation: active file (raylet.out) stays open with + an fd held; rotated copies (raylet.1.out, raylet.2.out, ...) + have NO open fd and are safe to delete. Default 24h. """ super().__init__( type="ray_log_cleanup", @@ -83,11 +89,14 @@ def __init__( raise ValueError(f"ray_log_cleanup.old_logs_keep_hours must be >= 1, got {old_logs_keep_hours}") if setup_log_keep_minutes < 1: raise ValueError(f"ray_log_cleanup.setup_log_keep_minutes must be >= 1, got {setup_log_keep_minutes}") + if rotated_daemon_keep_hours < 1: + raise ValueError(f"ray_log_cleanup.rotated_daemon_keep_hours must be >= 1, got {rotated_daemon_keep_hours}") self.ray_temp_dir = ray_temp_dir.rstrip("/") self.min_age_hours = min_age_hours self.live_log_keep_days = live_log_keep_days self.old_logs_keep_hours = old_logs_keep_hours self.setup_log_keep_minutes = setup_log_keep_minutes + self.rotated_daemon_keep_hours = rotated_daemon_keep_hours @classmethod def from_config(cls, task_config) -> "RayLogCleanupTask": @@ -98,6 +107,7 @@ def from_config(cls, task_config) -> "RayLogCleanupTask": live_log_keep_days=task_config.params.get("live_log_keep_days", 7), old_logs_keep_hours=task_config.params.get("old_logs_keep_hours", 24), setup_log_keep_minutes=task_config.params.get("setup_log_keep_minutes", 60), + rotated_daemon_keep_hours=task_config.params.get("rotated_daemon_keep_hours", 24), ) async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: @@ -106,11 +116,13 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: live_age_min = self.live_log_keep_days * 24 * 60 old_age_min = self.old_logs_keep_hours * 60 setup_age_min = self.setup_log_keep_minutes + rotated_age_min = self.rotated_daemon_keep_hours * 60 - # 4-stage shell pipeline: + # 5-stage shell pipeline: # PART 1 — drop stale full session__ dirs # PART 2a — session_latest/logs/ PID-aware (worker/runtime_env_setup with digit pid) # PART 2c — session_latest/logs/ runtime_env_setup-* mtime-only (covers hex variant) + # PART 2d — session_latest/logs/ rotated daemon logs (raylet.N.out, gcs_server.N.err, etc.) # PART 2b — session_latest/logs/ non-PID, non-daemon stale (long tail) # PART 3 — session_latest/logs/old/ time-based # @@ -183,6 +195,21 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: rm -f "$f" && echo "removed_setup=$(basename "$f")" done + # PART 2d: rotated daemon log backups — Ray performs in-place + # rotation for long-running daemon processes: the active file + # (e.g. raylet.out) keeps its fd open; rotated copies + # (raylet.1.out, raylet.2.out, ...) have NO open fd and are + # safe to delete. Without this stage the daemon whitelist + # (! -name 'raylet*') protects rotated copies indefinitely, + # causing multi-GB accumulation on long-running clusters. + # Pattern: .. where N is a positive integer. + find "$LOGS" -maxdepth 1 -type f -mmin +{rotated_age_min} \\ + -regextype posix-extended \\ + -regex '.*/((raylet|gcs_server|runtime_env_agent|dashboard|monitor|log_monitor)\\.[0-9]+\\.(out|err|log))$' \\ + | while read -r f; do + rm -f "$f" && echo "removed_rotated_daemon=$(basename "$f")" + done + # PART 2b: non-PID, non-daemon stale files older than # {self.live_log_keep_days} days. Daemon files (raylet*, # gcs_server*, runtime_env_agent*, dashboard*, monitor*, @@ -223,13 +250,15 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: removed_sessions = [line.split("=", 1)[1] for line in output.splitlines() if line.startswith("removed=")] removed_dead_pid = sum(1 for line in output.splitlines() if line.startswith("removed_dead_pid_log=")) removed_setup = sum(1 for line in output.splitlines() if line.startswith("removed_setup=")) + removed_rotated_daemon = sum(1 for line in output.splitlines() if line.startswith("removed_rotated_daemon=")) removed_stale = sum(1 for line in output.splitlines() if line.startswith("removed_stale_file=")) removed_old = sum(1 for line in output.splitlines() if line.startswith("removed_old=")) logger.info( f"[{self.type}] [{runtime._config.host}] ray_log_cleanup done: " f"sessions={len(removed_sessions)}, dead_pid={removed_dead_pid}, " - f"setup={removed_setup}, stale={removed_stale}, old={removed_old}, " + f"setup={removed_setup}, rotated_daemon={removed_rotated_daemon}, " + f"stale={removed_stale}, old={removed_old}, " f"output_head={output[:300]}" ) return { @@ -241,6 +270,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: # New per-category counters: "removed_dead_pid_count": removed_dead_pid, "removed_setup_count": removed_setup, + "removed_rotated_daemon_count": removed_rotated_daemon, "removed_stale_count": removed_stale, "removed_old_count": removed_old, "output_head": output[:1500], diff --git a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py index 13d2de537d..1a161b6419 100644 --- a/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_ray_log_cleanup_task.py @@ -41,6 +41,7 @@ def test_default(self): assert task.live_log_keep_days == 7 assert task.old_logs_keep_hours == 24 assert task.setup_log_keep_minutes == 60 + assert task.rotated_daemon_keep_hours == 24 def test_strips_trailing_slash(self): task = RayLogCleanupTask(ray_temp_dir="/data/ray/") @@ -62,6 +63,10 @@ def test_rejects_setup_log_keep_minutes_below_one(self): with pytest.raises(ValueError, match="setup_log_keep_minutes must be >= 1"): RayLogCleanupTask(setup_log_keep_minutes=0) + def test_rejects_rotated_daemon_keep_hours_below_one(self): + with pytest.raises(ValueError, match="rotated_daemon_keep_hours must be >= 1"): + RayLogCleanupTask(rotated_daemon_keep_hours=0) + def test_custom_thresholds(self): task = RayLogCleanupTask(live_log_keep_days=3, old_logs_keep_hours=12) assert task.live_log_keep_days == 3 @@ -76,6 +81,7 @@ def test_from_config_defaults(self): assert task.live_log_keep_days == 7 assert task.old_logs_keep_hours == 24 assert task.setup_log_keep_minutes == 60 + assert task.rotated_daemon_keep_hours == 24 def test_from_config_custom(self): cfg = _FakeTaskConfig( @@ -85,6 +91,7 @@ def test_from_config_custom(self): "live_log_keep_days": 14, "old_logs_keep_hours": 6, "setup_log_keep_minutes": 30, + "rotated_daemon_keep_hours": 12, }, interval_seconds=3600, ) @@ -94,6 +101,7 @@ def test_from_config_custom(self): assert task.live_log_keep_days == 14 assert task.old_logs_keep_hours == 6 assert task.setup_log_keep_minutes == 30 + assert task.rotated_daemon_keep_hours == 12 assert task.interval_seconds == 3600 @@ -234,6 +242,51 @@ async def test_part2c_targets_runtime_env_setup_glob(self): # Emits removed_setup= marker so output parser can count it assert "removed_setup=" in cmd + @pytest.mark.asyncio + async def test_part2d_uses_rotated_daemon_keep_hours(self): + """PART 2d mtime threshold = rotated_daemon_keep_hours * 60 minutes.""" + task = RayLogCleanupTask(rotated_daemon_keep_hours=48) + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # 48h * 60 = 2880 + assert "-mmin +2880" in cmd + + @pytest.mark.asyncio + async def test_part2d_matches_rotated_daemon_pattern(self): + """PART 2d must match .. but NOT the active file ..""" + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + cmd = runtime.execute.await_args.args[0].command + # Must use regex matching rotated copies (daemon.N.ext) + assert "raylet|gcs_server|runtime_env_agent|dashboard|monitor|log_monitor" in cmd + assert r"\.[0-9]+\." in cmd + # Emits removed_rotated_daemon= marker + assert "removed_rotated_daemon=" in cmd + + @pytest.mark.asyncio + async def test_part2d_does_not_match_active_daemon_files(self): + """Active daemon files (raylet.out, raylet.err) must NOT match PART 2d regex. + + The regex requires a numeric segment between daemon name and extension: + raylet.1.out matches, raylet.out does NOT.""" + import re + + task = RayLogCleanupTask() + runtime = _runtime() + await task.run_action(runtime) + + # The regex uses \.[0-9]+ between daemon name and extension, + # so `raylet.out` (no number) can never match. + pattern = r"(raylet|gcs_server|runtime_env_agent|dashboard|monitor|log_monitor)\.[0-9]+\.(out|err|log)" + assert not re.match(pattern, "raylet.out") + assert not re.match(pattern, "raylet.err") + assert re.match(pattern, "raylet.1.out") + assert re.match(pattern, "gcs_server.2.err") + @pytest.mark.asyncio async def test_part3_uses_old_logs_keep_hours(self): """PART 3 old-dir mtime threshold = old_logs_keep_hours * 60 minutes.""" @@ -329,6 +382,24 @@ async def test_extracts_part2c_setup_count_digit_and_hex(self): result = await task.run_action(runtime) assert result["removed_setup_count"] == 3 + @pytest.mark.asyncio + async def test_extracts_part2d_rotated_daemon_count(self): + """PART 2d must count rotated daemon log files (raylet.N.out etc.).""" + stdout = ( + "live_session=session_xxx\n" + "removed_rotated_daemon=raylet.1.out\n" + "removed_rotated_daemon=raylet.2.out\n" + "removed_rotated_daemon=raylet.3.out\n" + "removed_rotated_daemon=gcs_server.1.err\n" + "removed_rotated_daemon=dashboard.1.log\n" + "ray_log_cleanup_done" + ) + task = RayLogCleanupTask() + runtime = _runtime(stdout=stdout) + + result = await task.run_action(runtime) + assert result["removed_rotated_daemon_count"] == 5 + @pytest.mark.asyncio async def test_extracts_part3_old_count(self): stdout = ( @@ -354,6 +425,7 @@ async def test_all_counters_zero_when_nothing_removed(self): assert result["removed_count"] == 0 assert result["removed_dead_pid_count"] == 0 assert result["removed_setup_count"] == 0 + assert result["removed_rotated_daemon_count"] == 0 assert result["removed_stale_count"] == 0 assert result["removed_old_count"] == 0 @@ -366,6 +438,7 @@ async def test_handles_missing_ray_dir(self): assert result["status"] == TaskStatusEnum.SUCCESS assert result["removed_count"] == 0 assert result["removed_dead_pid_count"] == 0 + assert result["removed_rotated_daemon_count"] == 0 assert result["removed_stale_count"] == 0 assert result["removed_old_count"] == 0 From 60eb249b331dbc4ceffb5d57dd8eb362843c7391 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 11:29:01 +0800 Subject: [PATCH 154/226] feature(cli): add 'rock storage get' to download archived sandbox logs from OSS --- rock/cli/command/storage.py | 109 ++++++++ rock/sdk/sandbox/storage_client.py | 159 +++++++++++ tests/unit/cli/command/test_storage.py | 248 ++++++++++++++++++ tests/unit/sdk/sandbox/test_storage_client.py | 81 ++++++ 4 files changed, 597 insertions(+) create mode 100644 rock/cli/command/storage.py create mode 100644 rock/sdk/sandbox/storage_client.py create mode 100644 tests/unit/cli/command/test_storage.py create mode 100644 tests/unit/sdk/sandbox/test_storage_client.py diff --git a/rock/cli/command/storage.py b/rock/cli/command/storage.py new file mode 100644 index 0000000000..ac99bc66e1 --- /dev/null +++ b/rock/cli/command/storage.py @@ -0,0 +1,109 @@ +"""`rock storage get ` — download an archived sandbox log tarball. + +Thin CLI wrapper around :class:`rock.sdk.sandbox.storage_client.StorageClient`. +This file only handles argparse + user-facing output; all OSS/STS logic lives +in the SDK so it can be reused programmatically. +""" + +import argparse +import os + +from rock.cli.command.command import Command +from rock.logger import init_logger +from rock.sdk.sandbox.storage_client import ArchiveNotFoundError, StorageClient + +logger = init_logger("rock.cli.storage") + + +class StorageCommand(Command): + """rock storage get [-o PATH]""" + + name = "storage" + + async def arun(self, args: argparse.Namespace): + action = getattr(args, "storage_action", None) + if action == "get": + await self._get(args) + return + # argparse with required=True surfaces missing-action errors itself, + # but guard explicitly so subclasses/tests get a clear message. + raise ValueError(f"Unknown storage action: {action!r}") + + async def _get(self, args: argparse.Namespace): + client = StorageClient( + base_url=args.base_url, + auth_token=getattr(args, "auth_token", None), + extra_headers=getattr(args, "extra_headers", None), + ) + out_path = self._resolve_output_path(args.output, args.sandbox_id) + + try: + written = await client.download_archived_log( + sandbox_id=args.sandbox_id, + output_path=out_path, + archive_prefix=args.archive_prefix or "", + bucket=args.bucket, + endpoint=args.endpoint, + ) + except ArchiveNotFoundError as e: + print(f"NOT FOUND: {e}") + logger.error(f"Archive not found for sandbox {args.sandbox_id}: {e}") + return + except RuntimeError: + # Pre-flight errors (STS fetch, OSS config missing) — let them + # propagate so the user sees the configuration explanation. + raise + except Exception as e: + # OSS connectivity / unexpected oss2 errors — print + swallow. + print(f"FAILED: {e}") + logger.exception(f"Failed to download archive for {args.sandbox_id}: {e}") + return + + print(f"OK: {written}") + print(f"To extract: tar -xzf {written}") + + @staticmethod + def _resolve_output_path(output: str | None, sandbox_id: str) -> str: + if not output: + return f"./{sandbox_id}.tar.gz" + # If caller passed a directory (existing or with trailing /), drop the file inside. + if output.endswith("/") or os.path.isdir(output): + return os.path.join(output.rstrip("/"), f"{sandbox_id}.tar.gz") + return output + + @staticmethod + async def add_parser_to(subparsers: argparse._SubParsersAction): + storage = subparsers.add_parser( + "storage", + help="Manage sandbox archive storage on OSS", + description="Download archived sandbox log tarballs from OSS.", + ) + storage_sub = storage.add_subparsers(dest="storage_action", required=True) + + get_p = storage_sub.add_parser("get", help="Download an archived sandbox log tarball") + get_p.add_argument("sandbox_id", help="Sandbox id (matches the directory name under ROCK_LOGGING_PATH)") + get_p.add_argument( + "-o", + "--output", + default=None, + help="Output file path or directory (default: ./.tar.gz)", + ) + get_p.add_argument( + "--archive-prefix", + dest="archive_prefix", + default="", + help=( + "OSS key prefix used at archive time (must match admin's " + "sandbox_config.log.archive_prefix, e.g. 'rock-archives/')." + ), + ) + get_p.add_argument( + "--bucket", + default=None, + help="Override the OSS bucket returned by admin /get_token", + ) + get_p.add_argument( + "--endpoint", + default=None, + help="Override the OSS endpoint returned by admin /get_token", + ) diff --git a/rock/sdk/sandbox/storage_client.py b/rock/sdk/sandbox/storage_client.py new file mode 100644 index 0000000000..dbb09c9b15 --- /dev/null +++ b/rock/sdk/sandbox/storage_client.py @@ -0,0 +1,159 @@ +"""StorageClient — download archived sandbox log tarballs from OSS. + +Separate from :class:`rock.sdk.sandbox.oss_client.OssClient` because they serve +different OSS prefixes and lifecycles: + - OssClient — sandbox <-> host file transfer (rock-transfer/), alive sandbox + - StorageClient — stopped sandbox log archives (rock-archives/), historical + +Recovery flow: + 1. Call admin ``/get_token?account=primary`` to obtain a short-lived STS + for the primary OSS account. + 2. Use oss2 + STS auth to download + ``oss:///sandbox-logs/.tar.gz``. + +The CLI (``rock storage get``) is a thin wrapper around this client. +""" + +from __future__ import annotations + +import asyncio +import os +from typing import Any + +import oss2 + +from rock.logger import init_logger +from rock.utils.archive_command import ArchiveCommand +from rock.utils.http import HttpUtils + +logger = init_logger(__name__) + + +class ArchiveNotFoundError(FileNotFoundError): + """Raised when the requested sandbox's archive does not exist in OSS.""" + + +class StorageClient: + """SDK client for downloading archived sandbox log tarballs from OSS. + + Args: + base_url: admin base URL. May be the bare host (``https://admin/``) or + include the ``/apis/envs/sandbox/v1`` prefix; both are accepted. + auth_token: optional ``xrl-authorization`` token for admin requests. + extra_headers: optional extra HTTP headers attached to every admin call. + """ + + _API_PREFIX = "/apis/envs/sandbox/v1" + + def __init__( + self, + base_url: str, + auth_token: str | None = None, + extra_headers: dict[str, str] | None = None, + ): + self._base_url = base_url + self._auth_token = auth_token + self._extra_headers = dict(extra_headers or {}) + + async def download_archived_log( + self, + sandbox_id: str, + output_path: str, + archive_prefix: str = "", + bucket: str | None = None, + endpoint: str | None = None, + ) -> str: + """Download ``sandbox_id``'s archived log tarball to ``output_path``. + + Args: + sandbox_id: target sandbox id (matches the log directory name). + output_path: local file path to write the downloaded tarball. + archive_prefix: OSS key prefix used at archive time. Must match + admin's ``sandbox_config.log.archive_prefix``. + bucket: optional OSS bucket override (defaults to value returned + by admin ``/get_token``). + endpoint: optional OSS endpoint override (same default rule). + + Returns: + The local path the tarball was written to (== ``output_path``). + + Raises: + ArchiveNotFoundError: archive object does not exist in OSS. + RuntimeError: STS fetch failed or admin response missing fields. + """ + sts = await self._fetch_primary_sts() + bucket_name, endpoint_url, region = self._extract_oss_target(sts, bucket, endpoint) + oss_key = ArchiveCommand.build_key(sandbox_id, archive_prefix) + + oss_bucket = oss2.Bucket( + auth=oss2.StsAuth(sts["AccessKeyId"], sts["AccessKeySecret"], sts["SecurityToken"]), + endpoint=endpoint_url, + bucket_name=bucket_name, + region=region, + ) + + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + try: + await asyncio.to_thread(oss_bucket.get_object_to_file, oss_key, output_path) + except oss2.exceptions.NoSuchKey as e: + raise ArchiveNotFoundError(f"oss://{bucket_name}/{oss_key}") from e + + return output_path + + async def _fetch_primary_sts(self) -> dict[str, Any]: + url = self._build_get_token_url() + headers = self._build_headers() + try: + response = await HttpUtils.get(url, headers) + except Exception as e: + # Most common 404 cause: caller passed admin-write URL but /get_token only + # lives on the proxy/read role. Augment the error so the user knows what to flip. + if "404" in str(e): + raise RuntimeError( + f"admin /get_token returned 404 at {url}. " + "/get_token is only mounted on the proxy/read admin role, not the write admin. " + "If you used the write URL, switch base_url to the proxy/read URL." + ) from e + raise + if response.get("status") != "Success": + raise RuntimeError(f"admin /get_token returned: {response.get('message') or response}") + result = response.get("result") + if not result: + raise RuntimeError("admin /get_token returned an empty result; check OssConfig.primary on admin") + return result + + def _build_get_token_url(self) -> str: + """Normalize ``base_url`` and append ``/get_token?account=primary``. + + Accepts either the bare admin host or a base that already contains + the API prefix; in either case the result hits the right route. + """ + clean = self._base_url.rstrip("/") + if not clean.endswith(self._API_PREFIX): + clean = f"{clean}{self._API_PREFIX}" + return f"{clean}/get_token?account=primary" + + def _build_headers(self) -> dict[str, str]: + headers = dict(self._extra_headers) + if self._auth_token: + headers["xrl-authorization"] = self._auth_token + return headers + + @staticmethod + def _extract_oss_target( + sts: dict[str, Any], + bucket_override: str | None, + endpoint_override: str | None, + ) -> tuple[str, str, str | None]: + # The /get_token response from a recent admin includes Endpoint/Bucket/Region + # for the primary account. Caller overrides win (useful when testing against + # a non-default bucket without redeploying admin). + bucket = bucket_override or sts.get("Bucket") + endpoint = endpoint_override or sts.get("Endpoint") + region = sts.get("Region") + if not bucket or not endpoint: + raise RuntimeError( + "OSS bucket/endpoint missing — pass bucket/endpoint explicitly or configure " + "OssConfig.primary.bucket/endpoint on admin" + ) + return bucket, endpoint, region diff --git a/tests/unit/cli/command/test_storage.py b/tests/unit/cli/command/test_storage.py new file mode 100644 index 0000000000..d66b3db305 --- /dev/null +++ b/tests/unit/cli/command/test_storage.py @@ -0,0 +1,248 @@ +"""Tests for `rock storage get` — download archived sandbox log tarball.""" + +import argparse +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.cli.command.storage import StorageCommand + + +def _args(**overrides): + base = dict( + storage_action="get", + sandbox_id="sb-abc", + output=None, + archive_prefix="rock-archives/", + bucket=None, + endpoint=None, + base_url="http://admin.local:8080", + auth_token="tok-1", + extra_headers={"cluster": "default"}, + ) + base.update(overrides) + return argparse.Namespace(**base) + + +def _sts_payload(**extra): + payload = { + "AccessKeyId": "AK", + "AccessKeySecret": "SK", + "SecurityToken": "ST", + "Endpoint": "oss-cn-hangzhou.aliyuncs.com", + "Bucket": "chatos-rock", + "Region": "cn-hangzhou", + } + payload.update(extra) + return payload + + +class TestStorageGet: + @pytest.mark.asyncio + async def test_success_downloads_to_default_path_and_uses_correct_key(self, tmp_path, capsys): + cmd = StorageCommand() + args = _args(output=str(tmp_path / "out.tar.gz")) + + bucket_mock = MagicMock() + bucket_mock.get_object_to_file = MagicMock() + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.StsAuth.return_value = "stsauth" + # NoSuchKey lookup path is exercised separately; here the call succeeds. + oss2_mod.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + + await cmd.arun(args) + + # Asserts on what we sent to oss2: + oss2_mod.StsAuth.assert_called_once_with("AK", "SK", "ST") + oss2_mod.Bucket.assert_called_once() + kwargs = oss2_mod.Bucket.call_args.kwargs + assert kwargs["bucket_name"] == "chatos-rock" + assert kwargs["endpoint"] == "oss-cn-hangzhou.aliyuncs.com" + assert kwargs["region"] == "cn-hangzhou" + + # Key is built via shared helper, must match archive-side layout exactly. + bucket_mock.get_object_to_file.assert_called_once() + oss_key, local_path = bucket_mock.get_object_to_file.call_args.args + assert oss_key == "rock-archives/sandbox-logs/sb-abc.tar.gz" + assert local_path == str(tmp_path / "out.tar.gz") + + out = capsys.readouterr().out + assert "OK:" in out + assert "tar -xzf" in out + + @pytest.mark.asyncio + async def test_default_output_when_no_flag(self, capsys): + cmd = StorageCommand() + args = _args(output=None) + + bucket_mock = MagicMock() + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + patch("rock.sdk.sandbox.storage_client.os.makedirs"), + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + + await cmd.arun(args) + + _, local_path = bucket_mock.get_object_to_file.call_args.args + assert local_path == "./sb-abc.tar.gz" + + @pytest.mark.asyncio + async def test_directory_output_appends_filename(self, capsys): + cmd = StorageCommand() + args = _args(output="/tmp/recover/") # trailing slash → directory semantics + + bucket_mock = MagicMock() + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + patch("rock.sdk.sandbox.storage_client.os.makedirs"), + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + + await cmd.arun(args) + + _, local_path = bucket_mock.get_object_to_file.call_args.args + assert local_path == "/tmp/recover/sb-abc.tar.gz" + + @pytest.mark.asyncio + async def test_no_such_key_prints_not_found(self, capsys): + cmd = StorageCommand() + args = _args() + + # Build the NoSuchKey class first so storage.py can raise it via the patched module. + no_such_key = type("NoSuchKey", (Exception,), {}) + bucket_mock = MagicMock() + bucket_mock.get_object_to_file = MagicMock(side_effect=no_such_key("missing")) + + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + patch("rock.sdk.sandbox.storage_client.os.makedirs"), + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.exceptions.NoSuchKey = no_such_key + + await cmd.arun(args) + + out = capsys.readouterr().out + assert "NOT FOUND" in out + assert "rock-archives/sandbox-logs/sb-abc.tar.gz" in out + + @pytest.mark.asyncio + async def test_generic_oss_exception_prints_failed(self, capsys): + cmd = StorageCommand() + args = _args() + + bucket_mock = MagicMock() + # Use IOError (subclass of OSError, not RuntimeError) — real oss2 errors + # surface as oss2.exceptions.OssError/ServerError, which inherit from + # Exception (not RuntimeError). CLI swallows these and prints FAILED; + # RuntimeError is reserved for pre-flight (STS/config) errors and bubbles. + bucket_mock.get_object_to_file = MagicMock(side_effect=IOError("network exploded")) + + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + patch("rock.sdk.sandbox.storage_client.os.makedirs"), + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + + await cmd.arun(args) + + out = capsys.readouterr().out + assert "FAILED" in out + assert "network exploded" in out + + @pytest.mark.asyncio + async def test_admin_get_token_failure_raises(self): + cmd = StorageCommand() + args = _args() + with patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Failed", "message": "no primary configured"}), + ): + with pytest.raises(RuntimeError, match="no primary configured"): + await cmd.arun(args) + + @pytest.mark.asyncio + async def test_missing_oss_target_raises(self): + """If admin returns STS without Endpoint/Bucket and no CLI override, fail loud.""" + cmd = StorageCommand() + args = _args(bucket=None, endpoint=None) + + bare = _sts_payload() + bare.pop("Endpoint") + bare.pop("Bucket") + with patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": bare}), + ): + with pytest.raises(RuntimeError, match="bucket/endpoint missing"): + await cmd.arun(args) + + @pytest.mark.asyncio + async def test_cli_overrides_take_precedence_over_get_token_response(self): + cmd = StorageCommand() + args = _args(bucket="my-bucket", endpoint="my-endpoint") + + bucket_mock = MagicMock() + with ( + patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(return_value={"status": "Success", "result": _sts_payload()}), + ), + patch("rock.sdk.sandbox.storage_client.oss2") as oss2_mod, + patch("rock.sdk.sandbox.storage_client.os.makedirs"), + ): + oss2_mod.Bucket.return_value = bucket_mock + oss2_mod.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + + await cmd.arun(args) + + kwargs = oss2_mod.Bucket.call_args.kwargs + assert kwargs["bucket_name"] == "my-bucket" + assert kwargs["endpoint"] == "my-endpoint" + + @pytest.mark.asyncio + async def test_unknown_action_raises(self): + cmd = StorageCommand() + with pytest.raises(ValueError, match="Unknown storage action"): + await cmd.arun(argparse.Namespace(storage_action="delete")) + + @pytest.mark.asyncio + async def test_404_error_message_explains_proxy_role(self): + # A 404 on /get_token nearly always means the user pointed at the write-admin URL + # by mistake; surface that explicitly instead of bubbling a raw HTTPStatusError. + # Integration-level: CLI → SDK → patched HttpUtils → RuntimeError surfaces. + cmd = StorageCommand() + args = _args() + with patch( + "rock.sdk.sandbox.storage_client.HttpUtils.get", + AsyncMock(side_effect=Exception("Client error '404 Not Found' for url ...")), + ): + with pytest.raises(RuntimeError, match="proxy/read admin role"): + await cmd.arun(args) diff --git a/tests/unit/sdk/sandbox/test_storage_client.py b/tests/unit/sdk/sandbox/test_storage_client.py new file mode 100644 index 0000000000..951780560c --- /dev/null +++ b/tests/unit/sdk/sandbox/test_storage_client.py @@ -0,0 +1,81 @@ +"""SDK-level tests for StorageClient (URL/headers/STS extraction pure logic). + +Integration-level oss2/HttpUtils mocking is exercised end-to-end via +tests/unit/cli/command/test_storage.py (CLI -> SDK -> mocked oss2 path). +This file focuses on the pure helpers + auth header logic that don't need +the full download path. +""" + +import pytest + +from rock.sdk.sandbox.storage_client import StorageClient + + +class TestBuildGetTokenUrl: + def test_appends_api_prefix_when_base_url_is_bare_host(self): + client = StorageClient(base_url="https://admin.local") + assert client._build_get_token_url() == "https://admin.local/apis/envs/sandbox/v1/get_token?account=primary" + + def test_does_not_double_append_when_base_url_already_has_prefix(self): + client = StorageClient(base_url="https://admin.local/apis/envs/sandbox/v1") + assert client._build_get_token_url() == "https://admin.local/apis/envs/sandbox/v1/get_token?account=primary" + + def test_strips_trailing_slash(self): + client = StorageClient(base_url="https://admin.local/") + assert client._build_get_token_url() == "https://admin.local/apis/envs/sandbox/v1/get_token?account=primary" + + +class TestBuildHeaders: + def test_auth_token_lifts_to_xrl_authorization_header(self): + client = StorageClient(base_url="https://x", auth_token="tok-2", extra_headers={"cluster": "c1"}) + headers = client._build_headers() + assert headers["xrl-authorization"] == "tok-2" + assert headers["cluster"] == "c1" + + def test_no_auth_token_keeps_extra_headers_only(self): + client = StorageClient(base_url="https://x", auth_token=None, extra_headers={"cluster": "c1"}) + headers = client._build_headers() + assert "xrl-authorization" not in headers + assert headers["cluster"] == "c1" + + def test_extra_headers_isolated_from_caller_mutation(self): + # Defensive copy: caller mutating their dict after construction must not + # leak into subsequent client requests. + original = {"cluster": "c1"} + client = StorageClient(base_url="https://x", extra_headers=original) + original["leaked"] = "yes" + assert "leaked" not in client._build_headers() + + +class TestExtractOssTarget: + _STS_FULL = { + "AccessKeyId": "AK", + "AccessKeySecret": "SK", + "SecurityToken": "ST", + "Endpoint": "oss-cn-hangzhou.aliyuncs.com", + "Bucket": "chatos-rock", + "Region": "cn-hangzhou", + } + + def test_caller_bucket_override_wins(self): + bucket, endpoint, region = StorageClient._extract_oss_target(self._STS_FULL, "override-bucket", None) + assert bucket == "override-bucket" + assert endpoint == self._STS_FULL["Endpoint"] + assert region == "cn-hangzhou" + + def test_caller_endpoint_override_wins(self): + bucket, endpoint, region = StorageClient._extract_oss_target(self._STS_FULL, None, "override-endpoint") + assert bucket == self._STS_FULL["Bucket"] + assert endpoint == "override-endpoint" + + def test_missing_bucket_raises(self): + sts = dict(self._STS_FULL) + del sts["Bucket"] + with pytest.raises(RuntimeError, match="bucket/endpoint missing"): + StorageClient._extract_oss_target(sts, None, None) + + def test_missing_endpoint_raises(self): + sts = dict(self._STS_FULL) + del sts["Endpoint"] + with pytest.raises(RuntimeError, match="bucket/endpoint missing"): + StorageClient._extract_oss_target(sts, None, None) From e1d110170d08a89dcd7a1b3cb79ae8b8adb778cb Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 15:59:33 +0800 Subject: [PATCH 155/226] docs(v1.8.0): wrap PR numbers in release notes with GitHub links (EN + zh-Hans) --- .../version-1.8.x/Release Notes/v1.8.0.md | 28 ++++++++--------- .../version-1.8.x/Release Notes/v1.8.0.md | 30 +++++++++---------- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md index 060d386c5c..2978342039 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md +++ b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.8.x/Release Notes/v1.8.0.md @@ -14,9 +14,9 @@ ### Bug 修复 -* **SDK**: 在 Python 运行时环境中增加配置类型校验 (#652) +* **SDK**: 在 Python 运行时环境中增加配置类型校验 ([#652](https://github.com/alibaba/ROCK/pull/652)) -* **SDK**: 在 OSS 上传路径中使用 wget 前先创建目标父目录 (#940) +* **SDK**: 在 OSS 上传路径中使用 wget 前先创建目标父目录 ([#940](https://github.com/alibaba/ROCK/pull/940)) --- @@ -25,9 +25,9 @@ ### 新功能 -* 新增 CPU 超分能力，支持灰度发布、生命周期摘要和绝对核心数 CPU 指标 (#978) +* 新增 CPU 超分能力，支持灰度发布、生命周期摘要和绝对核心数 CPU 指标 ([#978](https://github.com/alibaba/ROCK/pull/978)) -* `get_status` 接口支持 `include_all_states` 参数，可查询所有状态的沙箱 (#951) +* `get_status` 接口支持 `include_all_states` 参数，可查询所有状态的沙箱 ([#951](https://github.com/alibaba/ROCK/pull/951)) * 上传/下载文件到sandbox的账号、bucket迁移（向前兼容，老bucket仍支持使用）（#953） @@ -40,7 +40,7 @@ ### 新功能 -* **Kubernetes**: 新增 GPU 支持，采用 Jinja2 模板和可扩展加速器类型 (#981) +* **Kubernetes**: 新增 GPU 支持，采用 Jinja2 模板和可扩展加速器类型 ([#981](https://github.com/alibaba/ROCK/pull/981)) * 区域集群配置优化：抽取公共配置，降低维护成本 @@ -51,7 +51,7 @@ * **Docker**: 容器停止时清理 XFS 项目配额 -* **Docker**: 移除无效的删除镜像分支，添加 CLS 日志服务支持到删除镜像功能 (#965) +* **Docker**: 移除无效的删除镜像分支，添加 CLS 日志服务支持到删除镜像功能 ([#965](https://github.com/alibaba/ROCK/pull/965)) * 修复删除镜像传参bug @@ -62,7 +62,7 @@ ### 新功能 -* 模型服务代理支持流式传输和回放模式，实现字节透传，提供转发和回放两种后端 (#935) +* 模型服务代理支持流式传输和回放模式，实现字节透传，提供转发和回放两种后端 ([#935](https://github.com/alibaba/ROCK/pull/935)) --- @@ -82,7 +82,7 @@ * 将 rock\_config 传递给沙箱表和元数据存储，确保指标监控使用正确的端点 -* 修复 `_get_user_info` 指标问题 (#911) +* 修复 `_get_user_info` 指标问题 ([#911](https://github.com/alibaba/ROCK/pull/911)) --- @@ -91,13 +91,13 @@ ### 新功能 -* 支持通过 Nacos 动态重载配置 (#888) +* 支持通过 Nacos 动态重载配置 ([#888](https://github.com/alibaba/ROCK/pull/888)) * 新增 Ray 日志清理任务，禁用 worker 到 driver 的日志转发 * 新增构建缓存清理任务，用于修剪 uv/pip 缓存 -* 将悬空镜像和 BuildKit 修剪合并到镜像清理任务中 (#970) +* 将悬空镜像和 BuildKit 修剪合并到镜像清理任务中 ([#970](https://github.com/alibaba/ROCK/pull/970)) * 优化文件清理定时任务的性能和配置安全验证 @@ -113,14 +113,14 @@ ### 新功能 -* 新增 Windows PowerShell 支持 (#921) +* 新增 Windows PowerShell 支持 ([#921](https://github.com/alibaba/ROCK/pull/921)) ### Bug 修复 * 将循环设备磁盘挂载到 Docker 数据根目录，替代硬编码路径 -* 为 Kata 运行时的 Nix 镜像添加 `/bin` 符号链接挂载 (#936) +* 为 Kata 运行时的 Nix 镜像添加 `/bin` 符号链接挂载 ([#936](https://github.com/alibaba/ROCK/pull/936)) * 使用 cgroup 指标获取容器 CPU 使用率，替代 psutil @@ -143,12 +143,12 @@ ## ♻️ 代码重构 -* **对象存储**: 将 OSS 上传/下载与客户端环境变量解耦，实现三层配置解析机制 (#943) +* **对象存储**: 将 OSS 上传/下载与客户端环境变量解耦，实现三层配置解析机制 ([#943](https://github.com/alibaba/ROCK/pull/943)) ## 🔧 构建与工具 -* 移除 `need_database` 标记 (#901) +* 移除 `need_database` 标记 ([#901](https://github.com/alibaba/ROCK/pull/901)) --- \ No newline at end of file diff --git a/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md index 4754f420bb..b87ef807b8 100644 --- a/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md +++ b/docs/versioned_docs/version-1.8.x/Release Notes/v1.8.0.md @@ -13,9 +13,9 @@ May 21, 2026 ### Bug Fixes -* **SDK**: Add runtime config type validation in `PythonRuntimeEnv` (#652) +* **SDK**: Add runtime config type validation in `PythonRuntimeEnv` ([#652](https://github.com/alibaba/ROCK/pull/652)) -* **SDK**: `mkdir` target parent dir before `wget` in OSS upload path (#940) +* **SDK**: `mkdir` target parent dir before `wget` in OSS upload path ([#940](https://github.com/alibaba/ROCK/pull/940)) --- @@ -23,11 +23,11 @@ May 21, 2026 ### New Features -* CPU overcommit with grayscale rollout, lifecycle summary, and absolute-cores CPU gauge (#978) +* CPU overcommit with grayscale rollout, lifecycle summary, and absolute-cores CPU gauge ([#978](https://github.com/alibaba/ROCK/pull/978)) -* `get_status` API supports `include_all_states` parameter — query sandboxes in all states (#951) +* `get_status` API supports `include_all_states` parameter — query sandboxes in all states ([#951](https://github.com/alibaba/ROCK/pull/951)) -* Account / bucket migration for sandbox upload/download (backward compatible — legacy bucket still supported) (#953) +* Account / bucket migration for sandbox upload/download (backward compatible — legacy bucket still supported) ([#953](https://github.com/alibaba/ROCK/pull/953)) * Fix transfer file being stored directly under bucket root (backward compatible) @@ -37,7 +37,7 @@ May 21, 2026 ### New Features -* **Kubernetes**: GPU support with Jinja2 templates and extensible accelerator types (#981) +* **Kubernetes**: GPU support with Jinja2 templates and extensible accelerator types ([#981](https://github.com/alibaba/ROCK/pull/981)) * Region cluster config refactor: extract common config, reduce maintenance cost @@ -47,7 +47,7 @@ May 21, 2026 * **Docker**: Cleanup XFS project quota on container stop -* **Docker**: Remove dead `remove_images` branch + add CLS log support to image removal (#965) +* **Docker**: Remove dead `remove_images` branch + add CLS log support to image removal ([#965](https://github.com/alibaba/ROCK/pull/965)) * Fix `remove_image` missing `cls` parameter bug @@ -57,7 +57,7 @@ May 21, 2026 ### New Features -* Model-service proxy supports streaming and replay — byte passthrough, with Forward and Replay backends (#935) +* Model-service proxy supports streaming and replay — byte passthrough, with Forward and Replay backends ([#935](https://github.com/alibaba/ROCK/pull/935)) --- @@ -75,7 +75,7 @@ May 21, 2026 * Pass `rock_config` to `SandboxTable` / `SandboxMetaStore` so `MetricsMonitor` uses the correct endpoint -* Fix `_get_user_info` metrics issue (#911) +* Fix `_get_user_info` metrics issue ([#911](https://github.com/alibaba/ROCK/pull/911)) --- @@ -83,13 +83,13 @@ May 21, 2026 ### New Features -* Dynamic config reloading via Nacos (#888) +* Dynamic config reloading via Nacos ([#888](https://github.com/alibaba/ROCK/pull/888)) * Add `RayLogCleanupTask` and disable worker-to-driver log forwarding * Add `BuildCacheCleanupTask` for pruning uv/pip caches -* Merge dangling-layer and BuildKit prune into `ImageCleanupTask` (#970) +* Merge dangling-layer and BuildKit prune into `ImageCleanupTask` ([#970](https://github.com/alibaba/ROCK/pull/970)) * Improve `FileCleanupTask` performance and config safety validation @@ -103,13 +103,13 @@ May 21, 2026 ### New Features -* Add Windows PowerShell support (#921) +* Add Windows PowerShell support ([#921](https://github.com/alibaba/ROCK/pull/921)) ### Bug Fixes * Mount loop disk to Docker data-root instead of hardcoded path -* Symlink mount into `/bin` for Nix images with Kata runtime (#936) +* Symlink mount into `/bin` for Nix images with Kata runtime ([#936](https://github.com/alibaba/ROCK/pull/936)) * Use cgroup metrics for container CPU instead of psutil @@ -129,10 +129,10 @@ May 21, 2026 ## ♻️ Refactoring -* **OSS**: Decouple OSS upload/download from client env vars; implement 3-layer config resolution (#943) +* **OSS**: Decouple OSS upload/download from client env vars; implement 3-layer config resolution ([#943](https://github.com/alibaba/ROCK/pull/943)) ## 🔧 Build & Tooling -* Remove `need_database` test marker (#901) +* Remove `need_database` test marker ([#901](https://github.com/alibaba/ROCK/pull/901)) --- From d5f1ef73184fbb032b68156ff9cfed404e9d95c6 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 11:23:17 +0800 Subject: [PATCH 156/226] fix(scheduler): split ImageCleanupTask prune (idempotent) from docuum launch --- .../scheduler/tasks/image_cleanup_task.py | 71 ++++-- .../scheduler/test_image_cleanup_task.py | 217 +++++++++++++++--- 2 files changed, 231 insertions(+), 57 deletions(-) diff --git a/rock/admin/scheduler/tasks/image_cleanup_task.py b/rock/admin/scheduler/tasks/image_cleanup_task.py index 26b021d23b..1c099dbd7e 100644 --- a/rock/admin/scheduler/tasks/image_cleanup_task.py +++ b/rock/admin/scheduler/tasks/image_cleanup_task.py @@ -63,8 +63,11 @@ def from_config(cls, task_config) -> "ImageCleanupTask": ) async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: - """Start docuum daemon, then synchronously prune dangling/build cache.""" - # 1) docuum: LRU image eviction (long-running, nohup &) + """NON-idempotent path: launch docuum daemon (gated by should_run via base run_on_worker).""" + return await self._launch_docuum(runtime) + + async def _launch_docuum(self, runtime: RemoteSandboxRuntime) -> dict: + """docuum: LRU image eviction (long-running, nohup &). NON-idempotent.""" check_and_install_cmd = ( f"command -v docuum > /dev/null 2>&1 || curl {env_vars.ROCK_DOCUUM_INSTALL_URL} -LSfs | sh" ) @@ -81,33 +84,53 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: result = await runtime.execute(Command(command=command, shell=True)) pid = extract_nohup_pid(result.stdout) - logger.info(f"image cleanup task [{pid}] run successfully on worker[{runtime._config.host}]") - - # 2) Dangling/BuildKit prune (sync, fail-soft so an old/missing docker - # subcommand on one worker doesn't abort the rest of the pipeline). - prune_exit = None - prune_output = "" - if self.keep_build_storage: - prune_steps = [ - "docker image prune -f --filter dangling=true", - f"docker builder prune -f --keep-storage {self.keep_build_storage}", - ] - prune_cmd = "; ".join(f"({s}) 2>&1 || true" for s in prune_steps) - prune_result = await runtime.execute(Command(command=prune_cmd, shell=True, check=False)) - prune_output = (prune_result.stdout or "").strip()[:1000] - prune_exit = prune_result.exit_code - logger.info( - f"docker prune done on worker[{runtime._config.host}]: " - f"keep_build_storage={self.keep_build_storage}, exit={prune_exit}, " - f"output_head={prune_output[:300]}" - ) - + logger.info(f"docuum launched with PID [{pid}] on worker[{runtime._config.host}]") return { "pid": pid, "disk_threshold": self.disk_threshold, "image_whitelist": self.image_whitelist, + "status": TaskStatusEnum.RUNNING, + } + + async def _run_prune(self, runtime: RemoteSandboxRuntime) -> dict: + """Dangling/BuildKit prune (sync, fail-soft). IDEMPOTENT — runs every cycle.""" + if not self.keep_build_storage: + return {"prune_exit_code": None, "prune_output_head": ""} + prune_steps = [ + "docker image prune -f --filter dangling=true", + f"docker builder prune -f --keep-storage {self.keep_build_storage}", + ] + prune_cmd = "; ".join(f"({s}) 2>&1 || true" for s in prune_steps) + prune_result = await runtime.execute(Command(command=prune_cmd, shell=True, check=False)) + prune_output = (prune_result.stdout or "").strip()[:1000] + prune_exit = prune_result.exit_code + logger.info( + f"docker prune done on worker[{runtime._config.host}]: " + f"keep_build_storage={self.keep_build_storage}, exit={prune_exit}, " + f"output_head={prune_output[:300]}" + ) + return { "keep_build_storage": self.keep_build_storage, "prune_exit_code": prune_exit, "prune_output_head": prune_output, - "status": TaskStatusEnum.RUNNING, } + + async def run_on_worker(self, ip): + """Override base: prune unconditionally (idempotent), then gate docuum on should_run. + + The base run_on_worker skips entire task when should_run returns False, which + correctly prevents docuum re-launch but wrongly blocks the idempotent prune + step (dangling layers / BuildKit cache pile up forever once docuum is alive). + """ + runtime = self._get_runtime(ip) + # prune always runs — idempotent, fail-soft + try: + await self._run_prune(runtime) + except Exception as e: + logger.warning(f"[{self.type}] prune failed on worker[{ip}]: {e}") + # docuum gated by should_run — non-idempotent + if not await self.should_run(runtime): + logger.info(f"[{self.type}] docuum already running on worker[{ip}], skip launch") + return + logger.info(f"[{self.type}] launch docuum on worker[{ip}]") + await self.single_run(runtime, ip) diff --git a/tests/unit/admin/scheduler/test_image_cleanup_task.py b/tests/unit/admin/scheduler/test_image_cleanup_task.py index 0121d1e5bc..b001d9fd79 100644 --- a/tests/unit/admin/scheduler/test_image_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_image_cleanup_task.py @@ -1,4 +1,13 @@ -"""Tests for ImageCleanupTask (docuum LRU + dangling/BuildKit prune).""" +"""Tests for ImageCleanupTask (docuum LRU + dangling/BuildKit prune). + +Architecture (post split): +- ``run_action`` only launches docuum daemon (NON-idempotent). +- ``_run_prune`` performs dangling/BuildKit prune (IDEMPOTENT). +- ``run_on_worker`` overrides base: prune unconditionally, then gate + docuum on ``should_run``. This decoupling fixes the regression where + the whole task was skipped once docuum was alive, causing dangling + layers and BuildKit cache to accumulate forever. +""" from unittest.mock import AsyncMock @@ -28,18 +37,22 @@ def _runtime(side_effects): return rt -# Fixed call sequence for run_action when keep_build_storage is set: +# Call sequence for _launch_docuum (run_action): # 1) docuum install check # 2) docuum start (returns PID-tagged stdout) -# 3) docker image prune + docker builder prune (one combined cmd) -def _default_results(pid=12345, prune_stdout="Total reclaimed space: 1.2GB"): +def _docuum_results(pid=12345): return [ _FakeExecResult(), _FakeExecResult(stdout=f"PIDSTART{pid}PIDEND"), - _FakeExecResult(stdout=prune_stdout), ] +# Call sequence for _run_prune (single combined cmd): +# 1) docker image prune + docker builder prune (one combined cmd) +def _prune_results(prune_stdout="Total reclaimed space: 1.2GB"): + return [_FakeExecResult(stdout=prune_stdout)] + + class TestInit: def test_default(self): task = ImageCleanupTask() @@ -90,32 +103,69 @@ def test_from_config_custom(self): class TestRunAction: + """run_action now only launches docuum — prune moved to _run_prune.""" + @pytest.mark.asyncio async def test_docuum_command_includes_threshold(self): task = ImageCleanupTask(disk_threshold="500G") - runtime = _runtime(_default_results()) + runtime = _runtime(_docuum_results()) await task.run_action(runtime) + + # call 0: install check; call 1: docuum start docuum_cmd = runtime.execute.await_args_list[1].args[0].command assert "docuum --threshold 500G" in docuum_cmd @pytest.mark.asyncio async def test_docuum_command_passes_whitelist(self): task = ImageCleanupTask(image_whitelist=[r"^rock-base.*$", r"^pinned:.*$"]) - runtime = _runtime(_default_results()) + runtime = _runtime(_docuum_results()) await task.run_action(runtime) + docuum_cmd = runtime.execute.await_args_list[1].args[0].command assert "--keep '^rock-base.*$'" in docuum_cmd assert "--keep '^pinned:.*$'" in docuum_cmd + @pytest.mark.asyncio + async def test_run_action_returns_pid_and_status(self): + task = ImageCleanupTask() + runtime = _runtime(_docuum_results(pid=98765)) + + result = await task.run_action(runtime) + + assert result["status"] == TaskStatusEnum.RUNNING + assert result["pid"] == 98765 + assert result["disk_threshold"] == "1T" + assert result["image_whitelist"] == [] + + @pytest.mark.asyncio + async def test_run_action_does_not_invoke_prune(self): + """Regression guard: run_action MUST NOT call prune (decoupled to _run_prune).""" + task = ImageCleanupTask() + runtime = _runtime(_docuum_results()) + + await task.run_action(runtime) + + # exactly 2 execute calls: install check + docuum start (NO prune call) + assert runtime.execute.await_count == 2 + for call in runtime.execute.await_args_list: + cmd = call.args[0].command + assert "docker image prune" not in cmd + assert "docker builder prune" not in cmd + + +class TestRunPrune: + """_run_prune is idempotent and runs every cycle (via run_on_worker).""" + @pytest.mark.asyncio async def test_prune_command_includes_image_and_builder_prune(self): task = ImageCleanupTask(keep_build_storage="10GB") - runtime = _runtime(_default_results()) + runtime = _runtime(_prune_results()) - await task.run_action(runtime) - prune_cmd = runtime.execute.await_args_list[2].args[0].command + await task._run_prune(runtime) + + prune_cmd = runtime.execute.await_args_list[0].args[0].command assert "docker image prune -f --filter dangling=true" in prune_cmd assert "docker builder prune -f --keep-storage 10GB" in prune_cmd @@ -124,48 +174,149 @@ async def test_prune_command_does_not_invoke_volume_prune(self): # Long-running sandboxes may attach named volumes; we don't want to # drop them on schedule. task = ImageCleanupTask() - runtime = _runtime(_default_results()) + runtime = _runtime(_prune_results()) - await task.run_action(runtime) - prune_cmd = runtime.execute.await_args_list[2].args[0].command + await task._run_prune(runtime) + + prune_cmd = runtime.execute.await_args_list[0].args[0].command assert "docker volume prune" not in prune_cmd @pytest.mark.asyncio async def test_prune_step_is_fail_soft(self): task = ImageCleanupTask() - runtime = _runtime(_default_results()) + runtime = _runtime(_prune_results()) - await task.run_action(runtime) - prune_cmd = runtime.execute.await_args_list[2].args[0].command + await task._run_prune(runtime) + + prune_cmd = runtime.execute.await_args_list[0].args[0].command # `(...) 2>&1 || true` makes a missing/old docker subcommand non-fatal. assert "|| true" in prune_cmd + @pytest.mark.asyncio + async def test_prune_records_output(self): + task = ImageCleanupTask() + runtime = _runtime(_prune_results(prune_stdout="Total reclaimed space: 3.5GB\n(more)")) + + result = await task._run_prune(runtime) + + assert "Total reclaimed space" in result["prune_output_head"] + assert result["prune_exit_code"] == 0 + assert result["keep_build_storage"] == "20GB" + @pytest.mark.asyncio async def test_prune_skipped_when_keep_build_storage_falsy(self): + """keep_build_storage=None → fast return, no docker exec.""" task = ImageCleanupTask(keep_build_storage=None) - # Only 2 execute calls expected: install check + docuum start. - runtime = _runtime(_default_results()[:2]) + runtime = AsyncMock() + runtime.execute = AsyncMock() + + result = await task._run_prune(runtime) + + assert result == {"prune_exit_code": None, "prune_output_head": ""} + runtime.execute.assert_not_awaited() - result = await task.run_action(runtime) - assert runtime.execute.await_count == 2 - assert result["prune_exit_code"] is None - assert result["prune_output_head"] == "" + +class TestRunOnWorker: + """run_on_worker overrides base: prune unconditionally, then gate docuum on should_run. + + Verifies the fix for the regression where the entire task was skipped once + docuum daemon was alive (NON_IDEMPOTENT idempotency type → should_run False + → run_action skipped → prune never ran → dangling layers accumulated forever). + """ @pytest.mark.asyncio - async def test_run_action_returns_pid_and_status(self): + async def test_first_run_launches_both(self, monkeypatch): + """status=None / docuum dead → prune runs AND single_run (docuum launch) runs.""" task = ImageCleanupTask() - runtime = _runtime(_default_results(pid=98765)) + runtime = AsyncMock() + monkeypatch.setattr(task, "_get_runtime", lambda ip: runtime) + monkeypatch.setattr(task, "should_run", AsyncMock(return_value=True)) + monkeypatch.setattr(task, "single_run", AsyncMock()) + prune_spy = AsyncMock(return_value={}) + monkeypatch.setattr(task, "_run_prune", prune_spy) - result = await task.run_action(runtime) - assert result["status"] == TaskStatusEnum.RUNNING - assert result["pid"] == 98765 - assert result["disk_threshold"] == "1T" - assert result["keep_build_storage"] == "20GB" + await task.run_on_worker("10.0.0.1") + + prune_spy.assert_awaited_once_with(runtime) + task.single_run.assert_awaited_once_with(runtime, "10.0.0.1") @pytest.mark.asyncio - async def test_run_action_records_prune_output(self): + async def test_docuum_alive_still_prunes(self, monkeypatch): + """CORE REGRESSION: docuum pid alive (should_run=False) → prune STILL runs. + + Before the fix, base run_on_worker would skip the entire task when + should_run returned False, blocking prune indefinitely once docuum + was alive (which is forever — docuum is a long-running daemon). + """ task = ImageCleanupTask() - runtime = _runtime(_default_results(prune_stdout="Total reclaimed space: 3.5GB\n(more)")) + runtime = AsyncMock() + monkeypatch.setattr(task, "_get_runtime", lambda ip: runtime) + monkeypatch.setattr(task, "should_run", AsyncMock(return_value=False)) + single_run = AsyncMock() + monkeypatch.setattr(task, "single_run", single_run) + prune_spy = AsyncMock(return_value={}) + monkeypatch.setattr(task, "_run_prune", prune_spy) - result = await task.run_action(runtime) - assert "Total reclaimed space" in result["prune_output_head"] + await task.run_on_worker("10.0.0.1") + + prune_spy.assert_awaited_once_with(runtime) + single_run.assert_not_awaited() # docuum NOT relaunched (correctly) + + @pytest.mark.asyncio + async def test_prune_exception_does_not_block_docuum(self, monkeypatch): + """prune raises → docuum launch still proceeds per should_run; no crash propagates.""" + task = ImageCleanupTask() + runtime = AsyncMock() + monkeypatch.setattr(task, "_get_runtime", lambda ip: runtime) + monkeypatch.setattr(task, "_run_prune", AsyncMock(side_effect=RuntimeError("docker down"))) + monkeypatch.setattr(task, "should_run", AsyncMock(return_value=True)) + single_run = AsyncMock() + monkeypatch.setattr(task, "single_run", single_run) + + # Must not raise (run_on_worker catches prune exception with try/except) + await task.run_on_worker("10.0.0.1") + + single_run.assert_awaited_once_with(runtime, "10.0.0.1") + + @pytest.mark.asyncio + async def test_docuum_dead_relaunches(self, monkeypatch): + """docuum pid not alive (should_run=True) → single_run called → docuum relaunched.""" + task = ImageCleanupTask() + runtime = AsyncMock() + monkeypatch.setattr(task, "_get_runtime", lambda ip: runtime) + monkeypatch.setattr(task, "should_run", AsyncMock(return_value=True)) + single_run = AsyncMock() + monkeypatch.setattr(task, "single_run", single_run) + monkeypatch.setattr(task, "_run_prune", AsyncMock(return_value={})) + + await task.run_on_worker("10.0.0.1") + + single_run.assert_awaited_once_with(runtime, "10.0.0.1") + + @pytest.mark.asyncio + async def test_prune_runs_before_docuum_check(self, monkeypatch): + """Order matters: prune must complete (or fail-soft) BEFORE should_run check. + + Guarantees prune always gets its chance even if should_run later + decides to skip docuum launch. + """ + task = ImageCleanupTask() + runtime = AsyncMock() + call_order = [] + monkeypatch.setattr(task, "_get_runtime", lambda ip: runtime) + + async def fake_prune(rt): + call_order.append("prune") + return {} + + async def fake_should_run(rt): + call_order.append("should_run") + return False + + monkeypatch.setattr(task, "_run_prune", fake_prune) + monkeypatch.setattr(task, "should_run", fake_should_run) + monkeypatch.setattr(task, "single_run", AsyncMock()) + + await task.run_on_worker("10.0.0.1") + + assert call_order == ["prune", "should_run"] From cf1ac33d9c848c86304e93d4d8fd5aaadc40d50c Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 25 May 2026 19:42:06 +0800 Subject: [PATCH 157/226] feat(scheduler): db-driven SandboxLogArchiveTask, drop sentinel file design --- rock/admin/main.py | 11 + rock/admin/scheduler/tasks/__init__.py | 2 + .../tasks/sandbox_log_archive_task.py | 234 ++++++++++++ .../test_sandbox_log_archive_task.py | 350 ++++++++++++++++++ 4 files changed, 597 insertions(+) create mode 100644 rock/admin/scheduler/tasks/sandbox_log_archive_task.py create mode 100644 tests/unit/admin/scheduler/test_sandbox_log_archive_task.py diff --git a/rock/admin/main.py b/rock/admin/main.py index b1eeef12c0..511da0959a 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -21,6 +21,12 @@ from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service from rock.admin.scheduler.scheduler import SchedulerThread +from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( + set_rock_config_provider as set_archive_rock_config_provider, +) +from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( + set_sandbox_table_provider as set_archive_sandbox_table_provider, +) from rock.config import DatabaseConfig, RockConfig, SchedulerConfig from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager @@ -89,6 +95,11 @@ async def lifespan(app: FastAPI): sandbox_table = SandboxTable(db_provider, rock_config=rock_config) meta_store = SandboxMetaStore(redis_provider=redis_provider, sandbox_table=sandbox_table, rock_config=rock_config) + # Wire SandboxLogArchiveTask deps. Providers (vs static set) so Nacos + # config reload propagates to the next task run without re-injection. + set_archive_sandbox_table_provider(lambda: sandbox_table) + set_archive_rock_config_provider(lambda: rock_config) + # init scheduler thread scheduler_thread = None diff --git a/rock/admin/scheduler/tasks/__init__.py b/rock/admin/scheduler/tasks/__init__.py index 41b251635f..9069adba6d 100644 --- a/rock/admin/scheduler/tasks/__init__.py +++ b/rock/admin/scheduler/tasks/__init__.py @@ -5,6 +5,7 @@ from rock.admin.scheduler.tasks.image_cleanup_task import ImageCleanupTask from rock.admin.scheduler.tasks.image_pull_task import ImagePullTask from rock.admin.scheduler.tasks.ray_log_cleanup_task import RayLogCleanupTask +from rock.admin.scheduler.tasks.sandbox_log_archive_task import SandboxLogArchiveTask __all__ = [ "BuildCacheCleanupTask", @@ -13,4 +14,5 @@ "ImageCleanupTask", "ImagePullTask", "RayLogCleanupTask", + "SandboxLogArchiveTask", ] diff --git a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py new file mode 100644 index 0000000000..e8a94028ec --- /dev/null +++ b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py @@ -0,0 +1,234 @@ +"""Deferred archival of stopped sandbox log directories — DB-driven. + +Per-worker daily task. Each run: + 1. /execute on worker: ``ls ${log_root}/`` returns candidate sandbox_ids + (directories under the log root are named after the sandbox they belong to). + 2. SandboxTable.list_by_in("sandbox_id", candidate_ids): batch query + state + stop_time from sandbox_record. + 3. For each candidate: + - state != "stopped" → skip (sandbox still alive) + - stop_time missing / unparseable → log warning, skip + - age_days < keep_days_before_archive → skip (too young) + - else → tar | ossutil cp && rm -rf + +No sentinel files; single source of truth is sandbox_record. +Credentials are passed via SandboxCommand.env, never argv. +""" + +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from rock import env_vars +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.scheduler.task_base import BaseTask, IdempotencyType, TaskStatusEnum +from rock.common.constants import SCHEDULER_LOG_NAME +from rock.logger import init_logger +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime +from rock.utils.archive_command import ArchiveCommand + +if TYPE_CHECKING: + pass + +logger = init_logger(name="sandbox_log_archive", file_name=SCHEDULER_LOG_NAME) + + +# Module-level providers, injected by main.py lifespan. +# Lazy callable so we always read the *current* config / table reference, +# not a stale snapshot — config can be hot-reloaded via Nacos. +_sandbox_table_provider = None # callable[[], SandboxTable | None] +_rock_config_provider = None # callable[[], RockConfig | None] + + +def set_sandbox_table_provider(provider) -> None: + global _sandbox_table_provider + _sandbox_table_provider = provider + + +def set_rock_config_provider(provider) -> None: + global _rock_config_provider + _rock_config_provider = provider + + +class SandboxLogArchiveTask(BaseTask): + def __init__( + self, + interval_seconds: int = 86400, + log_root: str | None = None, + ): + super().__init__( + type="sandbox_log_archive", + interval_seconds=interval_seconds, + idempotency=IdempotencyType.IDEMPOTENT, + ) + # Resolved at run time, not init time, so YAML override / env var + # still has effect when ROCK_LOGGING_PATH is exported late. + self._log_root_override = log_root + + @classmethod + def from_config(cls, task_config) -> "SandboxLogArchiveTask": + return cls( + interval_seconds=task_config.interval_seconds, + log_root=task_config.params.get("log_root"), + ) + + @property + def log_root(self) -> str: + root = self._log_root_override or env_vars.ROCK_LOGGING_PATH or "" + return root.rstrip("/") + + async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: + if not self.log_root: + logger.warning(f"[{self.type}] log_root unconfigured (set ROCK_LOGGING_PATH); skip") + return {"status": TaskStatusEnum.SUCCESS, "message": "no log root configured"} + + sandbox_table = _sandbox_table_provider() if _sandbox_table_provider else None + if sandbox_table is None: + logger.warning(f"[{self.type}] sandbox_table provider not set; skip") + return {"status": TaskStatusEnum.SUCCESS, "message": "sandbox_table not available"} + + rock_config = _rock_config_provider() if _rock_config_provider else None + if rock_config is None: + logger.warning(f"[{self.type}] rock_config provider not set; skip") + return {"status": TaskStatusEnum.SUCCESS, "message": "rock_config not available"} + + primary = rock_config.oss.primary + bucket = primary.bucket + endpoint = primary.endpoint + access_key_id = primary.access_key_id + access_key_secret = primary.access_key_secret + if not (bucket and endpoint and access_key_id and access_key_secret): + logger.warning(f"[{self.type}] OSS primary account incomplete; skip archival") + return {"status": TaskStatusEnum.SUCCESS, "message": "oss primary account not configured"} + + log_cfg = rock_config.sandbox_config.log + keep_days = int(log_cfg.keep_days_before_archive or 3) + archive_prefix = log_cfg.archive_prefix or "" + + # Step 1: discover candidate sandbox_ids on this worker + candidate_ids = await self._discover_candidates(runtime) + if not candidate_ids: + return { + "status": TaskStatusEnum.SUCCESS, + "scanned": 0, + "archived": 0, + } + + # Step 2: batch query DB for state + stop_time + rows = await sandbox_table.list_by_in("sandbox_id", candidate_ids) + rows_by_id = {r["sandbox_id"]: r for r in rows} + + now = datetime.now(timezone.utc) + archived = 0 + skipped_alive = 0 + skipped_too_young = 0 + skipped_orphan = 0 + failed = 0 + + for sandbox_id in candidate_ids: + row = rows_by_id.get(sandbox_id) + if row is None: + logger.warning(f"[{self.type}] orphan log dir for {sandbox_id} (no DB row); skip") + skipped_orphan += 1 + continue + if row.get("state") != "stopped": + skipped_alive += 1 + continue + + stop_time = self._parse_stop_time(row.get("stop_time")) + if stop_time is None: + logger.warning(f"[{self.type}] {sandbox_id} state=stopped but stop_time missing/unparseable; skip") + skipped_orphan += 1 + continue + + age_days = (now - stop_time).days + if age_days < keep_days: + skipped_too_young += 1 + continue + + try: + await self._archive_one( + runtime, + sandbox_id, + archive_prefix, + bucket, + endpoint, + access_key_id, + access_key_secret, + ) + archived += 1 + except Exception as e: + logger.exception(f"[{self.type}] archive {sandbox_id} failed: {e}") + failed += 1 + + return { + "status": TaskStatusEnum.SUCCESS, + "scanned": len(candidate_ids), + "archived": archived, + "skipped_alive": skipped_alive, + "skipped_too_young": skipped_too_young, + "skipped_orphan": skipped_orphan, + "failed": failed, + } + + async def _discover_candidates(self, runtime: RemoteSandboxRuntime) -> list[str]: + """List sandbox_ids that still have a log directory on this worker. + + Top-level entries under ``log_root`` are treated as candidate + sandbox_ids. Non-sandbox entries (e.g. logrotate's own files) are + filtered out by the DB lookup step (Step 2) — orphans are logged. + """ + cmd = f"ls -1 {self.log_root} 2>/dev/null || true" + result = await runtime.execute(Command(command=cmd, shell=True, check=False)) + if result.exit_code != 0: + return [] + names = (result.stdout or "").strip().split("\n") + return [n.strip() for n in names if n.strip()] + + @staticmethod + def _parse_stop_time(raw) -> datetime | None: + """Parse stop_time from DB row. Returns None on missing/malformed. + + ``sandbox_record.stop_time`` is ``String(64)``; accept ISO 8601 with + or without timezone. Naive datetimes are assumed UTC. + """ + if not raw: + return None + try: + s = str(raw).replace("Z", "+00:00") + dt = datetime.fromisoformat(s) + return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc) + except (ValueError, TypeError): + return None + + async def _archive_one( + self, + runtime: RemoteSandboxRuntime, + sandbox_id: str, + archive_prefix: str, + bucket: str, + endpoint: str, + access_key_id: str, + access_key_secret: str, + ) -> None: + """tar+gzip the sandbox's log dir, upload via ossutil, rm -rf on success. + + Credentials go through ``Command.env`` so they never appear in argv / + ``ps`` output. The command itself is built by the pure-function + ``build_archive_command`` (PR #957) — single source of truth for + archive key naming. + """ + log_dir = f"{self.log_root}/{sandbox_id}" + oss_key = ArchiveCommand.build_key(sandbox_id, archive_prefix) + cmd = ArchiveCommand.build_command(log_dir, oss_key, bucket, endpoint) + await runtime.execute( + Command( + command=cmd, + shell=True, + check=True, + env={ + "OSS_ACCESS_KEY_ID": access_key_id, + "OSS_ACCESS_KEY_SECRET": access_key_secret, + }, + ) + ) + logger.info(f"[{self.type}] archived {sandbox_id} -> oss://{bucket}/{oss_key}") diff --git a/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py b/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py new file mode 100644 index 0000000000..7e09992de8 --- /dev/null +++ b/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py @@ -0,0 +1,350 @@ +"""Tests for SandboxLogArchiveTask (DB-driven, no sentinel files).""" + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from rock.admin.scheduler.task_base import TaskStatusEnum +from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( + SandboxLogArchiveTask, + set_rock_config_provider, + set_sandbox_table_provider, +) + +# --------------------------------------------------------------------------- +# Test helpers +# --------------------------------------------------------------------------- + + +class _FakeTaskConfig: + def __init__(self, params=None, interval_seconds=86400): + self.params = params or {} + self.interval_seconds = interval_seconds + + +class _FakeExecResult: + def __init__(self, exit_code=0, stdout=""): + self.exit_code = exit_code + self.stdout = stdout + + +def _runtime(side_effects): + rt = AsyncMock() + rt._config = SimpleNamespace(host="10.0.0.1") + rt.execute = AsyncMock(side_effect=side_effects) + return rt + + +def _row(sandbox_id, state="stopped", stop_time=None): + return {"sandbox_id": sandbox_id, "state": state, "stop_time": stop_time} + + +def _iso_days_ago(days, tz=timezone.utc) -> str: + return (datetime.now(tz) - timedelta(days=days)).isoformat() + + +def _fake_rock_config( + bucket="b", + endpoint="oss-cn-hangzhou.aliyuncs.com", + access_key_id="AKID", + access_key_secret="AKSEC", + keep_days=3, + archive_prefix="rock-archives/", +): + return SimpleNamespace( + oss=SimpleNamespace( + primary=SimpleNamespace( + bucket=bucket, + endpoint=endpoint, + access_key_id=access_key_id, + access_key_secret=access_key_secret, + ) + ), + sandbox_config=SimpleNamespace( + log=SimpleNamespace( + keep_days_before_archive=keep_days, + archive_prefix=archive_prefix, + ) + ), + ) + + +@pytest.fixture(autouse=True) +def reset_providers(): + yield + set_sandbox_table_provider(None) + set_rock_config_provider(None) + + +@pytest.fixture +def fake_table(): + table = MagicMock() + table.list_by_in = AsyncMock(return_value=[]) + return table + + +# --------------------------------------------------------------------------- +# Init / from_config +# --------------------------------------------------------------------------- + + +class TestInit: + def test_default(self): + task = SandboxLogArchiveTask() + assert task.type == "sandbox_log_archive" + assert task.interval_seconds == 86400 + + def test_log_root_override(self): + task = SandboxLogArchiveTask(log_root="/custom/path/") + assert task.log_root == "/custom/path" # trailing slash stripped + + def test_log_root_falls_back_to_env(self, monkeypatch): + import rock.env_vars + + monkeypatch.setattr(rock.env_vars, "ROCK_LOGGING_PATH", "/data/logs/") + task = SandboxLogArchiveTask() + assert task.log_root == "/data/logs" + + def test_from_config(self): + task = SandboxLogArchiveTask.from_config( + _FakeTaskConfig(params={"log_root": "/var/log/rock"}, interval_seconds=3600) + ) + assert task.interval_seconds == 3600 + assert task.log_root == "/var/log/rock" + + +# --------------------------------------------------------------------------- +# Early-exit guards +# --------------------------------------------------------------------------- + + +class TestEarlyExit: + @pytest.mark.asyncio + async def test_skip_when_log_root_empty(self, monkeypatch, fake_table): + import rock.env_vars + + monkeypatch.setattr(rock.env_vars, "ROCK_LOGGING_PATH", None) + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + + task = SandboxLogArchiveTask() + runtime = _runtime([]) + result = await task.run_action(runtime) + + assert result["status"] == TaskStatusEnum.SUCCESS + assert "no log root" in result["message"] + runtime.execute.assert_not_awaited() + + @pytest.mark.asyncio + async def test_skip_when_sandbox_table_unset(self, fake_table): + set_sandbox_table_provider(None) + set_rock_config_provider(lambda: _fake_rock_config()) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([]) + result = await task.run_action(runtime) + + assert result["status"] == TaskStatusEnum.SUCCESS + assert "sandbox_table not available" in result["message"] + + @pytest.mark.asyncio + async def test_skip_when_rock_config_unset(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(None) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([]) + result = await task.run_action(runtime) + + assert result["status"] == TaskStatusEnum.SUCCESS + assert "rock_config not available" in result["message"] + + @pytest.mark.asyncio + async def test_skip_when_oss_primary_incomplete(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(bucket="")) # bucket empty → skip + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([]) + result = await task.run_action(runtime) + + assert result["status"] == TaskStatusEnum.SUCCESS + assert "oss primary account not configured" in result["message"] + runtime.execute.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + + +class TestDiscovery: + @pytest.mark.asyncio + async def test_no_candidates_returns_scanned_zero(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="")]) + + result = await task.run_action(runtime) + + assert result["scanned"] == 0 + assert result["archived"] == 0 + fake_table.list_by_in.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Classification (orphan / alive / too_young / archived) +# --------------------------------------------------------------------------- + + +class TestClassification: + @pytest.mark.asyncio + async def test_orphan_log_dir_skipped(self, fake_table): + # ls returns 'sb-orphan', but DB has no row for it + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + fake_table.list_by_in = AsyncMock(return_value=[]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-orphan")]) + + result = await task.run_action(runtime) + + assert result["skipped_orphan"] == 1 + assert result["archived"] == 0 + + @pytest.mark.asyncio + async def test_alive_sandbox_skipped(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-alive", state="alive")]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-alive")]) + + result = await task.run_action(runtime) + + assert result["skipped_alive"] == 1 + assert result["archived"] == 0 + + @pytest.mark.asyncio + async def test_stopped_too_young_skipped(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(keep_days=3)) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-young", state="stopped", stop_time=_iso_days_ago(1))]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-young")]) + + result = await task.run_action(runtime) + + assert result["skipped_too_young"] == 1 + assert result["archived"] == 0 + + @pytest.mark.asyncio + async def test_stopped_old_enough_archived(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(keep_days=3)) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-old", state="stopped", stop_time=_iso_days_ago(5))]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime( + [ + _FakeExecResult(stdout="sb-old"), # discover + _FakeExecResult(stdout=""), # archive cmd + ] + ) + + result = await task.run_action(runtime) + + assert result["archived"] == 1 + assert result["failed"] == 0 + assert runtime.execute.await_count == 2 + + @pytest.mark.asyncio + async def test_stop_time_malformed_skipped(self, fake_table): + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-bad", state="stopped", stop_time="not-a-date")]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-bad")]) + + result = await task.run_action(runtime) + + assert result["skipped_orphan"] == 1 + assert result["archived"] == 0 + + +# --------------------------------------------------------------------------- +# Archive command details (credentials in env, key format, isolation) +# --------------------------------------------------------------------------- + + +class TestArchiveCommand: + @pytest.mark.asyncio + async def test_credentials_passed_via_env_not_argv(self, fake_table): + """OSS_ACCESS_KEY_ID/SECRET MUST go through Command.env (so they + don't leak into ps argv / audit logs / shell history).""" + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(access_key_id="SECRET_AK", access_key_secret="SECRET_SK")) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-1", state="stopped", stop_time=_iso_days_ago(10))]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-1"), _FakeExecResult()]) + + await task.run_action(runtime) + + archive_call = runtime.execute.await_args_list[1] + cmd_obj = archive_call.args[0] + assert "SECRET_AK" not in cmd_obj.command + assert "SECRET_SK" not in cmd_obj.command + assert cmd_obj.env["OSS_ACCESS_KEY_ID"] == "SECRET_AK" + assert cmd_obj.env["OSS_ACCESS_KEY_SECRET"] == "SECRET_SK" + + @pytest.mark.asyncio + async def test_archive_command_uses_build_sandbox_log_key(self, fake_table): + """Archive key follows the format from rock/utils/archive_command.py + — single source of truth shared with rock storage get (PR #962).""" + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(bucket="my-bucket", archive_prefix="archives/")) + fake_table.list_by_in = AsyncMock(return_value=[_row("sb-x", state="stopped", stop_time=_iso_days_ago(5))]) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="sb-x"), _FakeExecResult()]) + + await task.run_action(runtime) + + archive_cmd = runtime.execute.await_args_list[1].args[0].command + # build_sandbox_log_key("sb-x", "archives/") => "archives/sandbox-logs/sb-x.tar.gz" + assert "oss://my-bucket/archives/sandbox-logs/sb-x.tar.gz" in archive_cmd + + @pytest.mark.asyncio + async def test_one_failure_does_not_abort_loop(self, fake_table): + """If one sandbox's archive raises, the loop continues for others.""" + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config(keep_days=3)) + fake_table.list_by_in = AsyncMock( + return_value=[ + _row("sb-fail", state="stopped", stop_time=_iso_days_ago(5)), + _row("sb-ok", state="stopped", stop_time=_iso_days_ago(5)), + ] + ) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime( + [ + _FakeExecResult(stdout="sb-fail\nsb-ok"), # discover + RuntimeError("ossutil down"), # archive sb-fail raises + _FakeExecResult(), # archive sb-ok succeeds + ] + ) + + result = await task.run_action(runtime) + + assert result["archived"] == 1 + assert result["failed"] == 1 From 5bac2135ef733b7f4ac2891ac18733f98513a7ba Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 26 May 2026 21:02:39 +0800 Subject: [PATCH 158/226] fix(scheduler): dispatch SandboxLogArchiveTask DB calls to main loop (fix cross-loop asyncpg pool) --- rock/admin/main.py | 10 +++++ .../tasks/sandbox_log_archive_task.py | 38 ++++++++++++++++++- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/rock/admin/main.py b/rock/admin/main.py index 511da0959a..8355d20b79 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -1,4 +1,5 @@ import argparse +import asyncio import json import logging import time @@ -21,6 +22,9 @@ from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service from rock.admin.scheduler.scheduler import SchedulerThread +from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( + set_main_loop_provider as set_archive_main_loop_provider, +) from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( set_rock_config_provider as set_archive_rock_config_provider, ) @@ -99,6 +103,12 @@ async def lifespan(app: FastAPI): # config reload propagates to the next task run without re-injection. set_archive_sandbox_table_provider(lambda: sandbox_table) set_archive_rock_config_provider(lambda: rock_config) + # Capture lifespan loop (uvicorn main loop). SandboxLogArchiveTask runs + # inside SchedulerThread's child loop; it must dispatch DB calls back to + # this main loop so asyncpg pool stays bound here and HTTP handlers don't + # break with "Future attached to a different loop". + _main_loop = asyncio.get_running_loop() + set_archive_main_loop_provider(lambda: _main_loop) # init scheduler thread scheduler_thread = None diff --git a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py index e8a94028ec..03b88e06a7 100644 --- a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py +++ b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py @@ -15,6 +15,7 @@ Credentials are passed via SandboxCommand.env, never argv. """ +import asyncio from datetime import datetime, timezone from typing import TYPE_CHECKING @@ -37,6 +38,12 @@ # not a stale snapshot — config can be hot-reloaded via Nacos. _sandbox_table_provider = None # callable[[], SandboxTable | None] _rock_config_provider = None # callable[[], RockConfig | None] +# Main event loop (uvicorn's loop where lifespan + HTTP handlers run). +# Required because this task runs inside SchedulerThread's child loop, but +# `sandbox_table` (asyncpg/SQLAlchemy) has a pool bound to the main loop — +# awaiting it directly from the child loop pollutes pool affinity and breaks +# subsequent HTTP handlers with "Future attached to a different loop". +_main_loop_provider = None # callable[[], asyncio.AbstractEventLoop | None] def set_sandbox_table_provider(provider) -> None: @@ -49,6 +56,34 @@ def set_rock_config_provider(provider) -> None: _rock_config_provider = provider +def set_main_loop_provider(provider) -> None: + global _main_loop_provider + _main_loop_provider = provider + + +async def _run_on_main_loop(coro): + """Dispatch ``coro`` to the main event loop if we're on a different one. + + Why: SchedulerThread runs tasks inside its own asyncio loop. asyncpg / + SQLAlchemy pool is bound to whichever loop first uses it (the main loop, + via lifespan ``create_tables`` + HTTP handlers). Calling ``await + sandbox_table.xxx()`` directly from the scheduler's child loop binds the + pool to *that* loop instead, breaking subsequent HTTP requests on the + main loop with ``Future attached to a different loop``. + """ + main_loop = _main_loop_provider() if _main_loop_provider else None + try: + current = asyncio.get_running_loop() + except RuntimeError: + current = None + if main_loop is None or current is main_loop: + # No main loop wired, or we're already on it — direct await is safe. + return await coro + # We're on a child loop. Dispatch to main loop and await via wrap_future. + future = asyncio.run_coroutine_threadsafe(coro, main_loop) + return await asyncio.wrap_future(future) + + class SandboxLogArchiveTask(BaseTask): def __init__( self, @@ -114,7 +149,8 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: } # Step 2: batch query DB for state + stop_time - rows = await sandbox_table.list_by_in("sandbox_id", candidate_ids) + # Dispatch DB query to main loop (see _run_on_main_loop docstring). + rows = await _run_on_main_loop(sandbox_table.list_by_in("sandbox_id", candidate_ids)) rows_by_id = {r["sandbox_id"]: r for r in rows} now = datetime.now(timezone.utc) From 9f613603809b55251727b831584e83f9f55e5e46 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 26 May 2026 21:08:54 +0800 Subject: [PATCH 159/226] fix(scheduler): cap cross-loop dispatch with 60s timeout (prevent hang if main loop dies) --- .../scheduler/tasks/sandbox_log_archive_task.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py index 03b88e06a7..786fb4c32f 100644 --- a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py +++ b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py @@ -61,6 +61,13 @@ def set_main_loop_provider(provider) -> None: _main_loop_provider = provider +# Safety cap on cross-loop dispatch: if the main loop is gone (e.g. admin +# was SIGKILLed before SchedulerThread had a chance to stop), the child +# loop's await on the dispatched future would hang forever. 60s is generous +# for any reasonable sandbox_table query; raises TimeoutError if exceeded. +_CROSS_LOOP_DISPATCH_TIMEOUT = 60.0 + + async def _run_on_main_loop(coro): """Dispatch ``coro`` to the main event loop if we're on a different one. @@ -70,6 +77,9 @@ async def _run_on_main_loop(coro): sandbox_table.xxx()`` directly from the scheduler's child loop binds the pool to *that* loop instead, breaking subsequent HTTP requests on the main loop with ``Future attached to a different loop``. + + The dispatched future is bounded by ``_CROSS_LOOP_DISPATCH_TIMEOUT`` + to prevent the child loop from hanging if the main loop has stopped. """ main_loop = _main_loop_provider() if _main_loop_provider else None try: @@ -79,9 +89,10 @@ async def _run_on_main_loop(coro): if main_loop is None or current is main_loop: # No main loop wired, or we're already on it — direct await is safe. return await coro - # We're on a child loop. Dispatch to main loop and await via wrap_future. + # We're on a child loop. Dispatch to main loop and await via wrap_future, + # bounded by a timeout so a dead main loop doesn't block us forever. future = asyncio.run_coroutine_threadsafe(coro, main_loop) - return await asyncio.wrap_future(future) + return await asyncio.wait_for(asyncio.wrap_future(future), timeout=_CROSS_LOOP_DISPATCH_TIMEOUT) class SandboxLogArchiveTask(BaseTask): From 67a7b3d8e6d0f8fce057b13b70f5e4ade4a4a3c3 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 11:28:47 +0800 Subject: [PATCH 160/226] fix(scheduler): _discover_candidates use find -type d to skip daemon log files --- .../tasks/sandbox_log_archive_task.py | 14 ++++++++---- .../test_sandbox_log_archive_task.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py index 786fb4c32f..6e05cffc94 100644 --- a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py +++ b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py @@ -220,11 +220,17 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: async def _discover_candidates(self, runtime: RemoteSandboxRuntime) -> list[str]: """List sandbox_ids that still have a log directory on this worker. - Top-level entries under ``log_root`` are treated as candidate - sandbox_ids. Non-sandbox entries (e.g. logrotate's own files) are - filtered out by the DB lookup step (Step 2) — orphans are logged. + Only **directories** directly under ``log_root`` are returned. The + daemon-written files in the same root (docuum.log, rocklet.log, + rock_worker.log, access.log, command.log, rsync_logs_to_host.log, + worker_metrics_monitor.log, image_pull.log, ...) must NOT be + treated as sandbox_ids — they would otherwise generate spurious + "orphan log dir" warnings and waste DB lookups. + + ``find -maxdepth 1 -mindepth 1 -type d -printf '%f\\n'`` is GNU-find + portable (worker images are Linux-based; macOS not supported here). """ - cmd = f"ls -1 {self.log_root} 2>/dev/null || true" + cmd = f"find {self.log_root} -maxdepth 1 -mindepth 1 -type d -printf '%f\\n' 2>/dev/null || true" result = await runtime.execute(Command(command=cmd, shell=True, check=False)) if result.exit_code != 0: return [] diff --git a/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py b/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py index 7e09992de8..63ddf050e8 100644 --- a/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py +++ b/tests/unit/admin/scheduler/test_sandbox_log_archive_task.py @@ -195,6 +195,28 @@ async def test_no_candidates_returns_scanned_zero(self, fake_table): assert result["archived"] == 0 fake_table.list_by_in.assert_not_awaited() + @pytest.mark.asyncio + async def test_discover_uses_find_type_d_not_ls(self, fake_table): + """REGRESSION: _discover_candidates must use `find -type d`, not `ls`. + + `ls` lists daemon-written files (docuum.log / rocklet.log / etc.) + directly under log_root alongside real sandbox subdirs; each file + then triggers a useless DB lookup and emits a spurious "orphan log + dir" warning. Verified in pre against real /data/logs that mixes + files + dirs. + """ + set_sandbox_table_provider(lambda: fake_table) + set_rock_config_provider(lambda: _fake_rock_config()) + + task = SandboxLogArchiveTask(log_root="/data/logs") + runtime = _runtime([_FakeExecResult(stdout="")]) + await task.run_action(runtime) + + cmd = runtime.execute.await_args_list[0].args[0].command + assert "-type d" in cmd, "discovery must restrict to directories" + assert "-maxdepth 1" in cmd, "must not recurse into sandbox log subdirs" + assert cmd.lstrip().startswith("find "), f"expected find, got: {cmd[:50]}" + # --------------------------------------------------------------------------- # Classification (orphan / alive / too_young / archived) From 115d75e4c3313fd9c630d5aed1e1db31e6dd42c0 Mon Sep 17 00:00:00 2001 From: lkc Date: Mon, 1 Jun 2026 16:28:55 +0800 Subject: [PATCH 161/226] fix(sdk): sanitize generated Harbor job names (#1031) Reject explicit Harbor job names containing path separators and strip dataset/task prefixes when auto-generating names. Co-Authored-By: Codex AI-Model: GPT-5 Codex AI-Contributed/Feature: 6/6 AI-Contributed/UT: 0/0 --- rock/sdk/bench/models/job/config.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rock/sdk/bench/models/job/config.py b/rock/sdk/bench/models/job/config.py index fcd836c02a..3305572c49 100644 --- a/rock/sdk/bench/models/job/config.py +++ b/rock/sdk/bench/models/job/config.py @@ -230,16 +230,18 @@ def _auto_job_name(self): import uuid as _uuid if self.job_name is not None: + if "/" in self.job_name: + raise ValueError("job_name must not contain '/'") return self parts: list[str] = [] if self.datasets: ds = self.datasets[0] if getattr(ds, "name", None): - parts.append(ds.name) + parts.append(ds.name.rsplit("/", 1)[-1]) task_names = getattr(ds, "task_names", None) or [] if len(task_names) == 1: - parts.append(task_names[0]) + parts.append(task_names[0].rsplit("/", 1)[-1]) parts.append(_uuid.uuid4().hex[:8]) self.job_name = "_".join(parts) From 999f081eedb85ce5348a509630a4954e1c6018d6 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Tue, 2 Jun 2026 18:00:46 +0800 Subject: [PATCH 162/226] feat(admin): add parameter validation for API endpoints (#985) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(admin): parameter validation with InvalidParameterRockError Introduce `validate_required_str` helper that rejects empty / whitespace-only string parameters at the API boundary, plus a dedicated exception type for unambiguous client-side parameter errors. New exception - `InvalidParameterRockError` (subclass of `BadRequestRockError`) marks request shapes that cannot be caused by server state. Backward compatible: existing `except BadRequestRockError` handlers keep catching it, HTTP stays 400. Metric routing - New `*.invalid_param` bucket: only `InvalidParameterRockError`. Clean signal of "the caller sent something wrong." - All other exceptions (including `BadRequestRockError`) go to `*.failure`, so failure rate stays a server-health signal undiluted by client mistakes. Validated parameters - `sandbox_id` across all sandbox / gem endpoints - `image` in start / start_async - `image_tag` in commit - `rock-host-ip` header in host proxy - CPU / memory / disk_limit_rootfs spec limits in SandboxManager - batch sandbox_ids + pagination in SandboxProxyService Kept as `BadRequestRockError` (server-state or HTTP-semantics ambiguity): accelerator_type (Nacos availability risk), sandbox already exists (409), sandbox not found (404 / GC masking), Ray/Deployment defensive memory parse, target_port multi-source conflicts, SDK-side output_file. Co-Authored-By: Claude Opus 4.7 * fix(admin): correct param validation contract for RockResponse[T] endpoints Review fixes on top of 7a3a44a4f: 1. Missing sandbox_id checks added to body-param endpoints: /execute, /create_session, /run_in_session, /close_session, /read_file, /write_file (both sandbox_api.py and sandbox_proxy_api.py). 2. Validation moved from service back to API layer: removed validate_required_str from SandboxProxyService.get_service_status and _update_expire_time. New API-layer checks cover every path into them. 3. Fixed ResponseValidationError on RockResponse[T] endpoints (/stop, /commit, /start, ...). Root cause: handle_exceptions wrapped raised RockExceptions as RockResponse(result=SandboxResponse(...)), but `result` must satisfy the endpoint's declared T -- FastAPI's response_model check then 500'd with a confusing traceback after the handler had already returned. Fix: validate_required_str now returns a failed RockResponse instead of raising; callers early-return with if err := validate_required_str(...): return err handle_exceptions is intentionally untouched -- the existing path for all other (non-validation) exceptions still behaves as before. gem endpoints don't return RockResponse, so they use a local _require_sandbox_id helper that raises HTTPException(400) -- cleaner than the previous InvalidParameterRockError that fell through to a 500. Tests: - test_validation.py rewritten for the return-value contract. - test_param_validation.py adds coverage for /stop, /commit image_tag, and the 12 newly-validated body-param endpoints. - test_sandbox_id_validation.py deleted (tested removed service-layer checks). Co-Authored-By: Claude Opus 4.7 * refactor(admin): replace validate_required_str with NonBlankStr Pydantic type Move required-string validation from imperative helper calls inside each endpoint to a declarative Pydantic Annotated type, and register a global RequestValidationError handler that maps validation failures back to the RockResponse(status=Failed, error=...) envelope so the contract stays consistent regardless of where validation happens. - common/validation.py: drop validate_required_str, expose NonBlankStr - common/exception.py + admin/main.py: register request_validation_exception_handler so 422 becomes the RockResponse envelope on HTTP 200 - admin/entrypoints/{sandbox_api,sandbox_proxy_api}.py, admin/gem/api.py, admin/proto/request.py, actions/envs/request.py: switch to NonBlankStr - admin/scheduler/{task_base,tasks/*}.py: SandboxCommand/Read/Write requests now require sandbox_id, so scheduler-side calls pass a "scheduler-task" placeholder - tests: cover NonBlankStr, the global handler, and the gem endpoints Co-Authored-By: Claude Opus 4.7 * fix(tests): use base action types in SandboxActor test SandboxActor.execute/create_session/run_in_session accept the base rock.actions types (Command/BashAction/CreateBashSessionRequest), not the API-layer SandboxCommand/SandboxBashAction/SandboxCreateBashSessionRequest wrappers which now require sandbox_id via NonBlankStr. Co-Authored-By: Claude Opus 4.7 * fix(tests): pass sandbox_id to SandboxCommand/Action in SandboxActor test 277f32a6a switched the test to base rock.actions types to avoid the new sandbox_id NonBlankStr requirement, but rocklet.execute and SandboxActor.execute are typed against the admin proto Sandbox* variants and read fields (shell/check/error_msg) that only exist on those subclasses. The base Command therefore raised AttributeError inside the Ray actor. Restore the SandboxCommand/SandboxBashAction/SandboxCreateBashSessionRequest imports and pass a placeholder sandbox_id so construction satisfies NonBlankStr without breaking the runtime contract. Co-Authored-By: Claude Opus 4.7 * refactor(admin): drop InvalidParameterRockError, fold into BadRequestRockError The dedicated subclass only differed from BadRequestRockError by metric bucket — same HTTP code, same response shape. Per review, the extra class and the request.invalid_param counter weren't pulling their weight, so collapse them. - Remove InvalidParameterRockError and the public re-export - Replace all raise sites with BadRequestRockError - Drop the request.invalid_param counter; all 4xx now lands in request.failure alongside other exceptions - Update metrics decorator test accordingly Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 --- rock/actions/envs/request.py | 10 +- rock/admin/entrypoints/sandbox_api.py | 27 +- rock/admin/entrypoints/sandbox_proxy_api.py | 7 +- rock/admin/main.py | 8 + rock/admin/proto/request.py | 15 +- rock/admin/scheduler/task_base.py | 16 +- .../tasks/build_cache_cleanup_task.py | 2 +- .../scheduler/tasks/container_cleanup_task.py | 2 +- .../scheduler/tasks/file_cleanup_task.py | 4 +- .../scheduler/tasks/image_cleanup_task.py | 8 +- rock/admin/scheduler/tasks/image_pull_task.py | 2 +- .../scheduler/tasks/ray_log_cleanup_task.py | 2 +- .../tasks/sandbox_log_archive_task.py | 3 +- rock/common/exception.py | 27 ++ rock/common/validation.py | 18 ++ rock/sandbox/sandbox_manager.py | 4 +- tests/unit/admin/entrypoints/__init__.py | 0 .../entrypoints/test_param_validation.py | 257 ++++++++++++++++++ tests/unit/admin/metrics/test_decorator.py | 19 ++ .../unit/admin/proto/test_sandbox_request.py | 10 +- tests/unit/common/test_exception_handlers.py | 84 ++++++ tests/unit/common/test_validation.py | 29 ++ .../unit/deployments/test_local_deployment.py | 13 +- tests/unit/deployments/test_sandbox.py | 14 +- tests/unit/rocklet/test_docker_deployment.py | 29 +- .../rocklet/test_local_sandbox_runtime.py | 28 +- tests/unit/utils/test_shell_util.py | 22 +- 27 files changed, 578 insertions(+), 82 deletions(-) create mode 100644 rock/common/validation.py create mode 100644 tests/unit/admin/entrypoints/__init__.py create mode 100644 tests/unit/admin/entrypoints/test_param_validation.py create mode 100644 tests/unit/common/test_exception_handlers.py create mode 100644 tests/unit/common/test_validation.py diff --git a/rock/actions/envs/request.py b/rock/actions/envs/request.py index 3f921d1d76..273b3c4e2b 100644 --- a/rock/actions/envs/request.py +++ b/rock/actions/envs/request.py @@ -1,20 +1,22 @@ from pydantic import BaseModel +from rock.common.validation import NonBlankStr + class EnvMakeRequest(BaseModel): env_id: str - sandbox_id: str + sandbox_id: NonBlankStr class EnvResetRequest(BaseModel): - sandbox_id: str + sandbox_id: NonBlankStr seed: int | None = None class EnvStepRequest(BaseModel): - sandbox_id: str + sandbox_id: NonBlankStr action: str class EnvCloseRequest(BaseModel): - sandbox_id: str + sandbox_id: NonBlankStr diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index 9a6a469d34..b5e757e3bc 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -37,6 +37,7 @@ SUPPORT_KATA_SWITCH, ) from rock.common.exception import handle_exceptions +from rock.common.validation import NonBlankStr from rock.deployments.config import AcceleratorType, DockerDeploymentConfig from rock.sandbox.sandbox_manager import SandboxManager from rock.sdk.common.exceptions import BadRequestRockError @@ -191,7 +192,7 @@ async def start_async( @sandbox_router.get("/is_alive") @handle_exceptions(error_message="get sandbox is alive failed") -async def is_alive(sandbox_id: str): +async def is_alive(sandbox_id: NonBlankStr): try: status_response = await sandbox_manager.get_status(sandbox_id) alive_response = IsAliveResponse(is_alive=status_response.is_alive, message=status_response.host_name) @@ -203,13 +204,13 @@ async def is_alive(sandbox_id: str): @sandbox_router.get("/get_sandbox_statistics") @handle_exceptions(error_message="get sandbox statistics failed") -async def get_sandbox_statistics(sandbox_id: str): +async def get_sandbox_statistics(sandbox_id: NonBlankStr): return RockResponse(result=await sandbox_manager.get_sandbox_statistics(sandbox_id)) @sandbox_router.get("/get_status") @handle_exceptions(error_message="get sandbox status failed") -async def get_status(sandbox_id: str, include_all_states: bool = False): +async def get_status(sandbox_id: NonBlankStr, include_all_states: bool = False): # TODO: do judgement inside operator if ( sandbox_manager.rock_config.nacos_provider is not None @@ -265,14 +266,14 @@ async def write_file(request: SandboxWriteFileRequest) -> RockResponse[WriteFile async def upload( file: UploadFile = File(...), target_path: str = Form(...), - sandbox_id: str | None = Form(None), + sandbox_id: Annotated[NonBlankStr, Form()] = ..., ) -> RockResponse[UploadResponse]: return RockResponse(result=await sandbox_manager.upload(file, target_path, sandbox_id)) @sandbox_router.post("/stop") @handle_exceptions(error_message="stop sandbox failed") -async def close(sandbox_id: str = Body(..., embed=True)) -> RockResponse[str]: +async def close(sandbox_id: Annotated[NonBlankStr, Body(embed=True)]) -> RockResponse[str]: await sandbox_manager.stop(sandbox_id) return RockResponse(result=f"{sandbox_id} stopped") @@ -287,13 +288,15 @@ async def restart(sandbox_id: str = Body(..., embed=True)) -> RockResponse[Sandb @sandbox_router.post("/commit") @handle_exceptions(error_message="commit sandbox failed") async def commit( - sandbox_id: str = Body(..., embed=True), - image_tag: str = Body( - ..., - embed=True, - example="docker.io/library/nginx:1.25", - description="commited image tag: /:", - ), + sandbox_id: Annotated[NonBlankStr, Body(embed=True)], + image_tag: Annotated[ + NonBlankStr, + Body( + embed=True, + example="docker.io/library/nginx:1.25", + description="commited image tag: /:", + ), + ], username: str = Body(..., embed=True), password: str = Body(..., embed=True), ) -> RockResponse[str]: diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 68f2d28689..510e4b3e02 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -1,5 +1,5 @@ import asyncio -from typing import Any +from typing import Annotated, Any from fastapi import APIRouter, Body, File, Form, Query, Request, UploadFile, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse as _JSONResponse @@ -28,6 +28,7 @@ from rock.admin.proto.response import BatchSandboxStatusResponse, SandboxListResponse from rock.common.exception import handle_exceptions from rock.common.port_validation import validate_port_forward_port +from rock.common.validation import NonBlankStr from rock.logger import init_logger from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sdk.common.exceptions import BadRequestRockError @@ -147,7 +148,7 @@ async def close_session(request: SandboxCloseBashSessionRequest) -> RockResponse @sandbox_proxy_router.get("/is_alive") @handle_exceptions(error_message="get sandbox is alive failed") -async def is_alive(sandbox_id: str): +async def is_alive(sandbox_id: NonBlankStr): return RockResponse(result=await sandbox_proxy_service.is_alive(sandbox_id)) @@ -168,7 +169,7 @@ async def write_file(request: SandboxWriteFileRequest) -> RockResponse[WriteFile async def upload( file: UploadFile = File(...), target_path: str = Form(...), - sandbox_id: str | None = Form(None), + sandbox_id: Annotated[NonBlankStr, Form()] = ..., ) -> RockResponse[UploadResponse]: return RockResponse(result=await sandbox_proxy_service.upload(file, target_path, sandbox_id)) diff --git a/rock/admin/main.py b/rock/admin/main.py index 8355d20b79..61e322b000 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -10,6 +10,7 @@ import uvicorn from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError from starlette.middleware.cors import CORSMiddleware from starlette.responses import JSONResponse @@ -31,6 +32,7 @@ from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( set_sandbox_table_provider as set_archive_sandbox_table_provider, ) +from rock.common.exception import request_validation_exception_handler from rock.config import DatabaseConfig, RockConfig, SchedulerConfig from rock.logger import init_logger from rock.sandbox.gem_manager import GemManager @@ -204,6 +206,12 @@ async def lifespan(app: FastAPI): # --- CORS configuration end --- +# Pydantic validation errors are matched by FastAPI before the catch-all Exception +# handler below — register an explicit override so they come out as RockResponse +# envelopes instead of the default 422 ``{"detail": [...]}``. +app.add_exception_handler(RequestValidationError, request_validation_exception_handler) + + @app.exception_handler(Exception) async def base_exception_handler(request: Request, exc: Exception): exc_content = {"detail": str(exc), "traceback": traceback.format_exc().split("\n")} diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index 1b49b141f2..7954b1af56 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -12,10 +12,11 @@ ReadFileRequest, WriteFileRequest, ) +from rock.common.validation import NonBlankStr class SandboxStartRequest(BaseModel): - image: str = "" + image: NonBlankStr """image""" image_os: str = "linux" """The operating system of the image (e.g., 'linux', 'windows').""" @@ -62,14 +63,14 @@ class SandboxCommand(Command): """Environment variables to pass to the command.""" cwd: str | None = None """The current working directory to run the command in.""" - sandbox_id: str | None = None + sandbox_id: NonBlankStr """The id of the sandbox.""" class SandboxCreateBashSessionRequest(CreateBashSessionRequest): startup_timeout: float = 1.0 max_read_size: int = 2000 - sandbox_id: str | None = None + sandbox_id: NonBlankStr remote_user: str | None = Field(default=None) @@ -77,7 +78,7 @@ class SandboxCreateBashSessionRequest(CreateBashSessionRequest): class SandboxBashAction(BashAction): - sandbox_id: str | None = None + sandbox_id: NonBlankStr """The id of the sandbox.""" is_interactive_command: bool = False """For a non-exiting command to an interactive program @@ -99,18 +100,18 @@ class SandboxBashAction(BashAction): class SandboxCloseBashSessionRequest(CloseBashSessionRequest): - sandbox_id: str | None = None + sandbox_id: NonBlankStr SandboxCloseSessionRequest = Annotated[SandboxCloseBashSessionRequest, Field(discriminator="session_type")] class SandboxReadFileRequest(ReadFileRequest): - sandbox_id: str | None = None + sandbox_id: NonBlankStr class SandboxWriteFileRequest(WriteFileRequest): - sandbox_id: str | None = None + sandbox_id: NonBlankStr class WarmupRequest(BaseModel): diff --git a/rock/admin/scheduler/task_base.py b/rock/admin/scheduler/task_base.py index 0c6ef29d18..77a5d92724 100644 --- a/rock/admin/scheduler/task_base.py +++ b/rock/admin/scheduler/task_base.py @@ -154,12 +154,14 @@ async def single_run(self, runtime: RemoteSandboxRuntime, ip: str) -> dict: async def get_task_status(self, runtime: RemoteSandboxRuntime) -> TaskStatus | None: """Get task status from worker.""" - check_file_resp = await runtime.execute(Command(command=f"ls {self.status_file_path}", shell=True)) + check_file_resp = await runtime.execute( + Command(command=f"ls {self.status_file_path}", shell=True, sandbox_id="scheduler-task") + ) if check_file_resp.exit_code == 2: logger.info(f"task status file not exist: {self.status_file_path}") return None - response = await runtime.read_file(ReadFileRequest(path=self.status_file_path)) + response = await runtime.read_file(ReadFileRequest(path=self.status_file_path, sandbox_id="scheduler-task")) if response.content: try: return TaskStatus.from_json(response.content) @@ -169,11 +171,15 @@ async def get_task_status(self, runtime: RemoteSandboxRuntime) -> TaskStatus | N async def save_task_status(self, runtime: RemoteSandboxRuntime, status: TaskStatus): """Save task status to worker file.""" - await runtime.write_file(WriteFileRequest(path=self.status_file_path, content=status.to_json())) + await runtime.write_file( + WriteFileRequest(path=self.status_file_path, content=status.to_json(), sandbox_id="scheduler-task") + ) async def _clear_task_status(self, runtime: RemoteSandboxRuntime) -> None: """Remove the status file from worker.""" - await runtime.execute(Command(command=f"rm -f {self.status_file_path}", shell=True)) + await runtime.execute( + Command(command=f"rm -f {self.status_file_path}", shell=True, sandbox_id="scheduler-task") + ) async def cleanup_on_worker(self, ip: str) -> None: """Stop any long-running process spawned by this task on a single worker. @@ -188,7 +194,7 @@ async def cleanup_on_worker(self, ip: str) -> None: return if await runtime.check_pid_exists(status.pid): kill_cmd = f"pkill -9 -P {status.pid}; kill -9 {status.pid}" - await runtime.execute(Command(command=kill_cmd, shell=True)) + await runtime.execute(Command(command=kill_cmd, shell=True, sandbox_id="scheduler-task")) logger.info(f"[{self.type}] killed pid {status.pid} on worker[{ip}]") await self._clear_task_status(runtime) diff --git a/rock/admin/scheduler/tasks/build_cache_cleanup_task.py b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py index 38c4f8e490..1c945da9aa 100644 --- a/rock/admin/scheduler/tasks/build_cache_cleanup_task.py +++ b/rock/admin/scheduler/tasks/build_cache_cleanup_task.py @@ -78,7 +78,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: # non-zero exit from short-circuiting the rest. snippets = [_TOOL_COMMANDS[t] for t in self.tools] command = "; ".join(snippets) if snippets else "echo 'no tools configured'" - result = await runtime.execute(Command(command=command, shell=True, check=False)) + result = await runtime.execute(Command(command=command, shell=True, check=False, sandbox_id="scheduler-task")) output = (result.stdout or "").strip() logger.info( f"[{self.type}] [{runtime._config.host}] cache prune done: " diff --git a/rock/admin/scheduler/tasks/container_cleanup_task.py b/rock/admin/scheduler/tasks/container_cleanup_task.py index 6f69d2e0dd..c79809b418 100644 --- a/rock/admin/scheduler/tasks/container_cleanup_task.py +++ b/rock/admin/scheduler/tasks/container_cleanup_task.py @@ -60,7 +60,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: f'\' > "$CLEANUP_LOG" 2>&1 & echo {PID_PREFIX}${{!}}{PID_SUFFIX}' ) - result = await runtime.execute(Command(command=command, shell=True)) + result = await runtime.execute(Command(command=command, shell=True, sandbox_id="scheduler-task")) pid = extract_nohup_pid(result.stdout) logger.info( diff --git a/rock/admin/scheduler/tasks/file_cleanup_task.py b/rock/admin/scheduler/tasks/file_cleanup_task.py index 15ef90d81a..464005f00f 100644 --- a/rock/admin/scheduler/tasks/file_cleanup_task.py +++ b/rock/admin/scheduler/tasks/file_cleanup_task.py @@ -231,7 +231,9 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: target_dir = dir_config.path try: command = self._build_cleanup_command(dir_config) - result = await runtime.execute(Command(command=command, shell=True, check=True)) + result = await runtime.execute( + Command(command=command, shell=True, check=True, sandbox_id="scheduler-task") + ) output = result.stdout.strip() if result.stdout else "" results[target_dir] = { "exit_code": result.exit_code, diff --git a/rock/admin/scheduler/tasks/image_cleanup_task.py b/rock/admin/scheduler/tasks/image_cleanup_task.py index 1c099dbd7e..84ef3f8b48 100644 --- a/rock/admin/scheduler/tasks/image_cleanup_task.py +++ b/rock/admin/scheduler/tasks/image_cleanup_task.py @@ -71,7 +71,7 @@ async def _launch_docuum(self, runtime: RemoteSandboxRuntime) -> dict: check_and_install_cmd = ( f"command -v docuum > /dev/null 2>&1 || curl {env_vars.ROCK_DOCUUM_INSTALL_URL} -LSfs | sh" ) - await runtime.execute(Command(command=check_and_install_cmd, shell=True)) + await runtime.execute(Command(command=check_and_install_cmd, shell=True, sandbox_id="scheduler-task")) log_redirect = ( '[ -n "$ROCK_LOGGING_PATH" ] && DOCUUM_LOG="$ROCK_LOGGING_PATH/docuum.log" || DOCUUM_LOG="/dev/null"' @@ -81,7 +81,7 @@ async def _launch_docuum(self, runtime: RemoteSandboxRuntime) -> dict: if keep_args: docuum_cmd = f"{docuum_cmd} {keep_args}" command = f'{log_redirect}; nohup {docuum_cmd} > "$DOCUUM_LOG" 2>&1 & echo {PID_PREFIX}${{!}}{PID_SUFFIX}' - result = await runtime.execute(Command(command=command, shell=True)) + result = await runtime.execute(Command(command=command, shell=True, sandbox_id="scheduler-task")) pid = extract_nohup_pid(result.stdout) logger.info(f"docuum launched with PID [{pid}] on worker[{runtime._config.host}]") @@ -101,7 +101,9 @@ async def _run_prune(self, runtime: RemoteSandboxRuntime) -> dict: f"docker builder prune -f --keep-storage {self.keep_build_storage}", ] prune_cmd = "; ".join(f"({s}) 2>&1 || true" for s in prune_steps) - prune_result = await runtime.execute(Command(command=prune_cmd, shell=True, check=False)) + prune_result = await runtime.execute( + Command(command=prune_cmd, shell=True, check=False, sandbox_id="scheduler-task") + ) prune_output = (prune_result.stdout or "").strip()[:1000] prune_exit = prune_result.exit_code logger.info( diff --git a/rock/admin/scheduler/tasks/image_pull_task.py b/rock/admin/scheduler/tasks/image_pull_task.py index a461accfd7..08500b0718 100644 --- a/rock/admin/scheduler/tasks/image_pull_task.py +++ b/rock/admin/scheduler/tasks/image_pull_task.py @@ -183,7 +183,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: f"echo {PID_PREFIX}${{!}}{PID_SUFFIX}" ) - result = await runtime.execute(Command(command=command, shell=True)) + result = await runtime.execute(Command(command=command, shell=True, sandbox_id="scheduler-task")) if result.exit_code != 0: error_msg = f"Failed to start image pull task: exit_code={result.exit_code}, stderr={result.stderr}" diff --git a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py index 1c151a701a..607baa20bd 100644 --- a/rock/admin/scheduler/tasks/ray_log_cleanup_task.py +++ b/rock/admin/scheduler/tasks/ray_log_cleanup_task.py @@ -242,7 +242,7 @@ async def run_action(self, runtime: RemoteSandboxRuntime) -> dict: echo "ray_log_cleanup_done" """ ) - result = await runtime.execute(Command(command=command, shell=True, check=False)) + result = await runtime.execute(Command(command=command, shell=True, check=False, sandbox_id="scheduler-task")) output = (result.stdout or "").strip() # Parse per-category removal counts from output. diff --git a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py index 6e05cffc94..1fbb8e91f0 100644 --- a/rock/admin/scheduler/tasks/sandbox_log_archive_task.py +++ b/rock/admin/scheduler/tasks/sandbox_log_archive_task.py @@ -231,7 +231,7 @@ async def _discover_candidates(self, runtime: RemoteSandboxRuntime) -> list[str] portable (worker images are Linux-based; macOS not supported here). """ cmd = f"find {self.log_root} -maxdepth 1 -mindepth 1 -type d -printf '%f\\n' 2>/dev/null || true" - result = await runtime.execute(Command(command=cmd, shell=True, check=False)) + result = await runtime.execute(Command(command=cmd, shell=True, check=False, sandbox_id="scheduler-task")) if result.exit_code != 0: return [] names = (result.stdout or "").strip().split("\n") @@ -282,6 +282,7 @@ async def _archive_one( "OSS_ACCESS_KEY_ID": access_key_id, "OSS_ACCESS_KEY_SECRET": access_key_secret, }, + sandbox_id=sandbox_id, ) ) logger.info(f"[{self.type}] archived {sandbox_id} -> oss://{bucket}/{oss_key}") diff --git a/rock/common/exception.py b/rock/common/exception.py index 8f5b6e037d..c506049a27 100644 --- a/rock/common/exception.py +++ b/rock/common/exception.py @@ -1,6 +1,10 @@ import functools import logging +from fastapi import Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse + from rock.actions import ResponseStatus, RockResponse from rock.logger import init_logger from rock.sdk.common.exceptions import RockException, from_rock_exception @@ -8,6 +12,29 @@ logger = init_logger(__name__) +async def request_validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse: + """Map FastAPI's RequestValidationError to the project's RockResponse envelope. + + FastAPI registers a default handler for RequestValidationError that returns + 422 ``{"detail": [...]}``. That shape clashes with the rest of the API, where + business failures come back as ``RockResponse(status=Failed, error=...)`` over + HTTP 200. Registering this handler on the FastAPI app aligns Pydantic-driven + validation errors with the same contract used by ``validate_required_str`` — + callers see one shape regardless of where validation happened. + """ + msg = "; ".join(f"{'.'.join(str(p) for p in e['loc'])}: {e['msg']}" for e in exc.errors()) + logger.warning("request validation failed on %s: %s", request.url.path, msg) + return JSONResponse( + status_code=200, + content=RockResponse( + status=ResponseStatus.FAILED, + message="invalid parameter", + error=msg, + result=None, + ).model_dump(), + ) + + def handle_exceptions(error_message: str = "error occurred"): """Exception handling decorator diff --git a/rock/common/validation.py b/rock/common/validation.py new file mode 100644 index 0000000000..f04ee53c4d --- /dev/null +++ b/rock/common/validation.py @@ -0,0 +1,18 @@ +"""Reusable Pydantic types for API request validation.""" + +from typing import Annotated + +from pydantic import StringConstraints + +NonBlankStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] +"""A string that cannot be empty or whitespace-only. + +Used on request models / Path / Query / Body / Form parameters to reject +``""`` and ``" "`` at deserialization time. ``strip_whitespace=True`` causes +Pydantic to trim surrounding whitespace before applying ``min_length=1``, +catching whitespace-only inputs. + +Validation failures surface as ``RequestValidationError``, which the global +handler in ``rock.common.exception.request_validation_exception_handler`` maps +back to the project's ``RockResponse(status=Failed, error=...)`` envelope. +""" diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index f0e6d58a06..95cdd146bc 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -412,4 +412,6 @@ def validate_sandbox_spec(self, runtime_config: RuntimeConfig, deployment_config parse_size_to_bytes(deployment_config.disk_limit_rootfs) except ValueError as e: logger.warning(f"Invalid disk_limit_rootfs size: {deployment_config.disk_limit_rootfs}", exc_info=e) - raise BadRequestRockError(f"Invalid disk_limit_rootfs size: {deployment_config.disk_limit_rootfs}") + raise BadRequestRockError( + f"Invalid disk_limit_rootfs size: {deployment_config.disk_limit_rootfs}" + ) diff --git a/tests/unit/admin/entrypoints/__init__.py b/tests/unit/admin/entrypoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/admin/entrypoints/test_param_validation.py b/tests/unit/admin/entrypoints/test_param_validation.py new file mode 100644 index 0000000000..5f80600278 --- /dev/null +++ b/tests/unit/admin/entrypoints/test_param_validation.py @@ -0,0 +1,257 @@ +"""Test parameter validation at the API endpoint level.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.exceptions import RequestValidationError +from httpx import ASGITransport, AsyncClient + +from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager +from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service +from rock.admin.gem.api import gem_router, set_env_service +from rock.admin.proto.response import SandboxStartResponse +from rock.common.exception import request_validation_exception_handler + + +def _build_app(router): + app = FastAPI() + # Match the production app: pydantic validation errors are mapped back to + # the RockResponse envelope instead of FastAPI's default 422 {"detail": [...]}. + app.add_exception_handler(RequestValidationError, request_validation_exception_handler) + app.include_router(router) + return app + + +@pytest.fixture +def sandbox_app(): + mock_manager = MagicMock() + mock_manager.rock_config = MagicMock() + mock_manager.rock_config.nacos_provider = None + set_sandbox_manager(mock_manager) + return _build_app(sandbox_router), mock_manager + + +@pytest.fixture +def proxy_app(): + mock_service = MagicMock() + set_sandbox_proxy_service(mock_service) + return _build_app(sandbox_proxy_router), mock_service + + +@pytest.fixture +def gem_app(): + mock_service = MagicMock() + set_env_service(mock_service) + return _build_app(gem_router), mock_service + + +def _assert_failed(resp, *, field: str = "sandbox_id"): + """Assert a RockResponse-shaped failure mentioning ``field`` in the error. + + Validation now happens at deserialization via NonBlankStr. The global + RequestValidationError handler maps the failure back to RockResponse with + HTTP 200, ``status=Failed``, ``result=None``, and a ``error`` string that + embeds the offending field's location (e.g. ``body.sandbox_id: ...``). + """ + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "Failed" + assert body["result"] is None + assert field in body["error"] + + +# --- sandbox_api.py tests --- + + +@pytest.mark.asyncio +async def test_sandbox_is_alive_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.get("/is_alive", params={"sandbox_id": ""})) + _assert_failed(await client.get("/is_alive", params={"sandbox_id": " "})) + + +@pytest.mark.asyncio +async def test_sandbox_get_status_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.get("/get_status", params={"sandbox_id": ""})) + + +@pytest.mark.asyncio +async def test_sandbox_get_statistics_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.get("/get_sandbox_statistics", params={"sandbox_id": ""})) + + +@pytest.mark.asyncio +async def test_sandbox_start_empty_image(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/start", json={"image": ""}), field="image") + _assert_failed(await client.post("/start", json={"image": " "}), field="image") + _assert_failed(await client.post("/start", json={}), field="image") + + +@pytest.mark.asyncio +async def test_sandbox_start_valid_image_passes_validation(sandbox_app): + """Verify that a valid image value does not trigger the validation error.""" + app, mock_manager = sandbox_app + mock_manager.start = AsyncMock( + return_value=SandboxStartResponse(sandbox_id="sb-1", host_name="h", host_ip="1.2.3.4") + ) + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.post("/start", json={"image": "python:3.11"}) + assert resp.status_code == 200 + assert resp.json()["status"] == "Success" + + +@pytest.mark.asyncio +async def test_sandbox_stop_empty_sandbox_id_returns_clean_failure(sandbox_app): + """Regression: /stop is typed RockResponse[str]; the old raise-via-handle_exceptions + path produced a ResponseValidationError because the wrapped result was a + SandboxResponse, not a str. The early-return pattern keeps the response shape + consistent with the declared response_model.""" + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/stop", json={"sandbox_id": ""})) + _assert_failed(await client.post("/stop", json={"sandbox_id": " "})) + + +@pytest.mark.asyncio +async def test_sandbox_commit_empty_image_tag(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed( + await client.post( + "/commit", + json={"sandbox_id": "sb-1", "image_tag": "", "username": "u", "password": "p"}, + ), + field="image_tag", + ) + + +@pytest.mark.asyncio +async def test_sandbox_run_in_session_empty_sandbox_id(sandbox_app): + """Regression for issue 1: /run_in_session takes sandbox_id from the request body + and was previously unvalidated.""" + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/run_in_session", json={"sandbox_id": "", "command": "ls"})) + + +@pytest.mark.asyncio +async def test_sandbox_execute_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/execute", json={"sandbox_id": "", "command": "ls"})) + + +@pytest.mark.asyncio +async def test_sandbox_create_session_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/create_session", json={"sandbox_id": ""})) + + +@pytest.mark.asyncio +async def test_sandbox_close_session_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/close_session", json={"sandbox_id": "", "session": "s"})) + + +@pytest.mark.asyncio +async def test_sandbox_read_file_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/read_file", json={"sandbox_id": "", "path": "/tmp/x"})) + + +@pytest.mark.asyncio +async def test_sandbox_write_file_empty_sandbox_id(sandbox_app): + app, _ = sandbox_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/write_file", json={"sandbox_id": "", "path": "/tmp/x", "content": "c"})) + + +# --- sandbox_proxy_api.py tests --- + + +@pytest.mark.asyncio +async def test_proxy_is_alive_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.get("/is_alive", params={"sandbox_id": ""})) + _assert_failed(await client.get("/is_alive", params={"sandbox_id": " "})) + + +@pytest.mark.asyncio +async def test_proxy_run_in_session_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/run_in_session", json={"sandbox_id": "", "command": "ls"})) + + +@pytest.mark.asyncio +async def test_proxy_execute_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/execute", json={"sandbox_id": "", "command": "ls"})) + + +@pytest.mark.asyncio +async def test_proxy_create_session_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/create_session", json={"sandbox_id": ""})) + + +@pytest.mark.asyncio +async def test_proxy_close_session_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/close_session", json={"sandbox_id": "", "session": "s"})) + + +@pytest.mark.asyncio +async def test_proxy_read_file_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/read_file", json={"sandbox_id": "", "path": "/tmp/x"})) + + +@pytest.mark.asyncio +async def test_proxy_write_file_empty_sandbox_id(proxy_app): + app, _ = proxy_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/write_file", json={"sandbox_id": "", "path": "/tmp/x", "content": "c"})) + + +# --- gem/api.py tests --- + + +@pytest.mark.asyncio +async def test_gem_step_empty_sandbox_id(gem_app): + app, _ = gem_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/step", json={"sandbox_id": "", "action": "noop"})) + _assert_failed(await client.post("/step", json={"sandbox_id": " ", "action": "noop"})) + + +@pytest.mark.asyncio +async def test_gem_reset_empty_sandbox_id(gem_app): + app, _ = gem_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/reset", json={"sandbox_id": ""})) + _assert_failed(await client.post("/reset", json={"sandbox_id": " "})) + + +@pytest.mark.asyncio +async def test_gem_close_empty_sandbox_id(gem_app): + app, _ = gem_app + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + _assert_failed(await client.post("/close", json={"sandbox_id": ""})) + _assert_failed(await client.post("/close", json={"sandbox_id": " "})) diff --git a/tests/unit/admin/metrics/test_decorator.py b/tests/unit/admin/metrics/test_decorator.py index 247499bcf6..08ff686353 100644 --- a/tests/unit/admin/metrics/test_decorator.py +++ b/tests/unit/admin/metrics/test_decorator.py @@ -13,6 +13,7 @@ ) from rock.admin.metrics.monitor import MetricsMonitor from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sdk.common.exceptions import BadRequestRockError class SampleObject: @@ -186,3 +187,21 @@ async def do_something(self, sandbox_id): assert attrs["user_id"] == "u1" assert attrs["experiment_id"] == "e1" assert attrs["namespace"] == "n1" + + +def test_record_metrics_bad_request_goes_to_failure(): + """BadRequestRockError lands in the generic `.failure` bucket along with + every other exception type — there is no separate client-fault bucket.""" + mock_metrics_monitor = Mock(spec=MetricsMonitor) + attributes = {"operation": "test_op"} + start_time = 0 + exception = BadRequestRockError("sandbox_id is required") + + with patch("rock.admin.metrics.decorator.time.perf_counter", return_value=1.0): + with pytest.raises(BadRequestRockError): + _record_metrics(mock_metrics_monitor, exception, attributes, start_time, "test") + + error_attrs = {**attributes, "error_type": "BadRequestRockError"} + mock_metrics_monitor.record_counter_by_name.assert_any_call("test.failure", 1, error_attrs) + mock_metrics_monitor.record_gauge_by_name.assert_called_once_with("test.rt", 1000.0, error_attrs) + mock_metrics_monitor.record_counter_by_name.assert_any_call("test.total", 1, error_attrs) diff --git a/tests/unit/admin/proto/test_sandbox_request.py b/tests/unit/admin/proto/test_sandbox_request.py index f98d55bc7b..c199b0207e 100644 --- a/tests/unit/admin/proto/test_sandbox_request.py +++ b/tests/unit/admin/proto/test_sandbox_request.py @@ -13,13 +13,13 @@ def test_image_os_default_value(): """SandboxStartRequest should default image_os to 'linux'.""" - request = SandboxStartRequest() + request = SandboxStartRequest(image="ubuntu:22.04") assert request.image_os == "linux" def test_image_os_custom_value(): """SandboxStartRequest should accept a custom image_os value.""" - request = SandboxStartRequest(image_os="windows") + request = SandboxStartRequest(image="ubuntu:22.04", image_os="windows") assert request.image_os == "windows" @@ -67,14 +67,14 @@ def test_all_fields_propagate_from_request_to_docker_config(): def test_sandbox_id_becomes_container_name(): """sandbox_id from SandboxStartRequest should map to container_name in DockerDeploymentConfig.""" - request = SandboxStartRequest(sandbox_id="my-sandbox") + request = SandboxStartRequest(image="ubuntu:22.04", sandbox_id="my-sandbox") config = DockerDeploymentConfig.from_request(request) assert config.container_name == "my-sandbox" def test_none_sandbox_id_yields_none_container_name(): """When sandbox_id is None, container_name in DockerDeploymentConfig should also be None.""" - request = SandboxStartRequest(sandbox_id=None) + request = SandboxStartRequest(image="ubuntu:22.04", sandbox_id=None) config = DockerDeploymentConfig.from_request(request) assert config.container_name is None @@ -82,7 +82,7 @@ def test_none_sandbox_id_yields_none_container_name(): def test_image_os_default_matches_between_sdk_config_and_start_request(): """SandboxConfig and SandboxStartRequest must share the same default value for image_os.""" sdk_config = SandboxConfig() - start_request = SandboxStartRequest() + start_request = SandboxStartRequest(image="ubuntu:22.04") assert sdk_config.image_os == start_request.image_os diff --git a/tests/unit/common/test_exception_handlers.py b/tests/unit/common/test_exception_handlers.py new file mode 100644 index 0000000000..c6ca34aee4 --- /dev/null +++ b/tests/unit/common/test_exception_handlers.py @@ -0,0 +1,84 @@ +"""Tests for global exception handlers shared across FastAPI apps.""" + +from typing import Annotated + +import pytest +from fastapi import FastAPI +from fastapi.exceptions import RequestValidationError +from httpx import ASGITransport, AsyncClient +from pydantic import BaseModel, StringConstraints + +from rock.common.exception import request_validation_exception_handler + + +class _Body(BaseModel): + sandbox_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] + + +@pytest.fixture +def app(): + app = FastAPI() + app.add_exception_handler(RequestValidationError, request_validation_exception_handler) + + @app.post("/echo") + async def _echo(body: _Body): + return {"ok": True} + + @app.get("/echo_query") + async def _echo_query(sandbox_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]): + return {"ok": True} + + return app + + +@pytest.mark.asyncio +async def test_empty_string_body_returns_rock_envelope(app): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.post("/echo", json={"sandbox_id": ""}) + # Envelope contract aligns with validate_required_str: HTTP 200, business failure inside body. + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "Failed" + assert body["result"] is None + assert "sandbox_id" in body["error"] + + +@pytest.mark.asyncio +async def test_whitespace_only_body_returns_rock_envelope(app): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.post("/echo", json={"sandbox_id": " "}) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "Failed" + assert body["result"] is None + assert "sandbox_id" in body["error"] + + +@pytest.mark.asyncio +async def test_missing_field_returns_rock_envelope(app): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.post("/echo", json={}) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "Failed" + assert body["result"] is None + assert "sandbox_id" in body["error"] + + +@pytest.mark.asyncio +async def test_invalid_query_param_returns_rock_envelope(app): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.get("/echo_query", params={"sandbox_id": ""}) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "Failed" + assert body["result"] is None + assert "sandbox_id" in body["error"] + + +@pytest.mark.asyncio +async def test_valid_request_passes_through(app): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + resp = await client.post("/echo", json={"sandbox_id": "abc"}) + assert resp.status_code == 200 + assert resp.json() == {"ok": True} diff --git a/tests/unit/common/test_validation.py b/tests/unit/common/test_validation.py new file mode 100644 index 0000000000..d4750b26f9 --- /dev/null +++ b/tests/unit/common/test_validation.py @@ -0,0 +1,29 @@ +"""Tests for the NonBlankStr Pydantic type.""" + +import pytest +from pydantic import BaseModel, ValidationError + +from rock.common.validation import NonBlankStr + + +class _Model(BaseModel): + value: NonBlankStr + + +def test_accepts_normal_string(): + assert _Model(value="abc").value == "abc" + + +def test_strips_surrounding_whitespace(): + assert _Model(value=" abc ").value == "abc" + + +@pytest.mark.parametrize("bad", ["", " ", "\t\n"]) +def test_rejects_empty_or_whitespace(bad): + with pytest.raises(ValidationError): + _Model(value=bad) + + +def test_rejects_missing_field(): + with pytest.raises(ValidationError): + _Model() diff --git a/tests/unit/deployments/test_local_deployment.py b/tests/unit/deployments/test_local_deployment.py index cc75df7187..2735134cd8 100644 --- a/tests/unit/deployments/test_local_deployment.py +++ b/tests/unit/deployments/test_local_deployment.py @@ -21,18 +21,19 @@ async def test_local_deployment(): async def test_nohup_output_command(): d = LocalDeployment() await d.start() - await d.runtime.create_session(CreateBashSessionRequest(session_type="bash")) + sid = "local-test" + await d.runtime.create_session(CreateBashSessionRequest(session_type="bash", sandbox_id=sid)) cmd_with_nohup = 'nohup echo "hello, rock" > /tmp/nohup_test.out 2>&1 &' - await d.runtime.run_in_session(BashAction(command=cmd_with_nohup)) - nohup_resp = await d.runtime.run_in_session(BashAction(command="cat /tmp/nohup_test.out")) + await d.runtime.run_in_session(BashAction(command=cmd_with_nohup, sandbox_id=sid)) + nohup_resp = await d.runtime.run_in_session(BashAction(command="cat /tmp/nohup_test.out", sandbox_id=sid)) assert "nohup: ignoring input" in nohup_resp.output cmd_without_nohup = 'nohup echo "hello, rock" < /dev/null > /tmp/nohup_test.out 2>&1 &' - await d.runtime.run_in_session(BashAction(command=cmd_without_nohup)) - nohup_resp = await d.runtime.run_in_session(BashAction(command="cat /tmp/nohup_test.out")) + await d.runtime.run_in_session(BashAction(command=cmd_without_nohup, sandbox_id=sid)) + nohup_resp = await d.runtime.run_in_session(BashAction(command="cat /tmp/nohup_test.out", sandbox_id=sid)) assert "nohup: ignoring input" not in nohup_resp.output - await d.runtime.run_in_session(BashAction(command="rm -rf /tmp/nohup_test.out")) + await d.runtime.run_in_session(BashAction(command="rm -rf /tmp/nohup_test.out", sandbox_id=sid)) await d.stop() assert not await d.is_alive() diff --git a/tests/unit/deployments/test_sandbox.py b/tests/unit/deployments/test_sandbox.py index f144aafa7d..24a8206150 100644 --- a/tests/unit/deployments/test_sandbox.py +++ b/tests/unit/deployments/test_sandbox.py @@ -14,6 +14,8 @@ logger = init_logger(__name__) +SANDBOX_ID = "test-sandbox" + @pytest.mark.need_ray @pytest.mark.asyncio @@ -21,21 +23,23 @@ async def test_execute(ray_init_shutdown): sandbox_config = LocalDeploymentConfig() sandbox_actor = SandboxActor.remote(sandbox_config, sandbox_config.get_deployment()) sandbox_actor.start.remote() - command_response_ref = sandbox_actor.execute.remote(Command(command="ls")) + command_response_ref = sandbox_actor.execute.remote(Command(command="ls", sandbox_id=SANDBOX_ID)) command_response: CommandResponse = ray.get(command_response_ref) assert "rock" in command_response.stdout logger.info(f"observation: {command_response.stdout}") - sandbox_actor.create_session.remote(CreateBashSessionRequest()) + sandbox_actor.create_session.remote(CreateBashSessionRequest(sandbox_id=SANDBOX_ID)) observation: BashObservation = ray.get( - sandbox_actor.run_in_session.remote(BashAction(command="ls", action_type="bash")) + sandbox_actor.run_in_session.remote(BashAction(command="ls", action_type="bash", sandbox_id=SANDBOX_ID)) ) logger.info(f"observation: {observation}") assert "rock" in observation.output - ray.get(sandbox_actor.run_in_session.remote(BashAction(command="cd rock", action_type="bash"))) + ray.get( + sandbox_actor.run_in_session.remote(BashAction(command="cd rock", action_type="bash", sandbox_id=SANDBOX_ID)) + ) observation: BashObservation = ray.get( - sandbox_actor.run_in_session.remote(BashAction(command="ls", action_type="bash")) + sandbox_actor.run_in_session.remote(BashAction(command="ls", action_type="bash", sandbox_id=SANDBOX_ID)) ) logger.info(f"observation: {observation}") assert "sdk" in observation.output diff --git a/tests/unit/rocklet/test_docker_deployment.py b/tests/unit/rocklet/test_docker_deployment.py index c144e5d0bc..e69f2a1f83 100644 --- a/tests/unit/rocklet/test_docker_deployment.py +++ b/tests/unit/rocklet/test_docker_deployment.py @@ -3,7 +3,10 @@ import pytest from rock import env_vars -from rock.actions import BashAction, CloseBashSessionRequest, Command, CreateBashSessionRequest +from rock.admin.proto.request import SandboxBashAction as BashAction +from rock.admin.proto.request import SandboxCloseBashSessionRequest as CloseBashSessionRequest +from rock.admin.proto.request import SandboxCommand as Command +from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest from rock.deployments.config import DockerDeploymentConfig, get_deployment @@ -18,31 +21,32 @@ async def test_docker_deployment(container_name): await d.is_alive() await d.start() assert await d.is_alive() - command = Command(command=["echo", "hello"]) + sid = container_name + command = Command(command=["echo", "hello"], sandbox_id=sid) await d.runtime.execute(command) # test bash session with default env - create_session_request = CreateBashSessionRequest(session_type="bash") + create_session_request = CreateBashSessionRequest(session_type="bash", sandbox_id=sid) await d.runtime.create_session(create_session_request) - action = BashAction(command="echo $PATH") + action = BashAction(command="echo $PATH", sandbox_id=sid) path_result_with_env = await d.runtime.run_in_session(action) print(path_result_with_env.output) - action = BashAction(command="echo $HOME") + action = BashAction(command="echo $HOME", sandbox_id=sid) home_result_with_env = await d.runtime.run_in_session(action) print(home_result_with_env.output) - close_session_request = CloseBashSessionRequest(session_type="bash") + close_session_request = CloseBashSessionRequest(session_type="bash", sandbox_id=sid) await d.runtime.close_session(close_session_request) # test bash session without default env - create_session_request = CreateBashSessionRequest(session_type="bash", env_enable=False) + create_session_request = CreateBashSessionRequest(session_type="bash", env_enable=False, sandbox_id=sid) await d.runtime.create_session(create_session_request) - action = BashAction(command="echo $PATH") + action = BashAction(command="echo $PATH", sandbox_id=sid) path_result = await d.runtime.run_in_session(action) print(path_result.output) - action = BashAction(command="echo $HOME") + action = BashAction(command="echo $HOME", sandbox_id=sid) home_result = await d.runtime.run_in_session(action) print(home_result.output) - close_session_request = CloseBashSessionRequest(session_type="bash") + close_session_request = CloseBashSessionRequest(session_type="bash", sandbox_id=sid) await d.runtime.close_session(close_session_request) await d.stop() @@ -58,14 +62,15 @@ async def test_docker_deployment_mounts_localtime_in_container(container_name): try: await d.start() + sid = container_name if host_has_zoneinfo: - result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%z"])) + result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%z"], sandbox_id=sid)) import subprocess host_offset = subprocess.check_output(["date", "+%z"], env={**os.environ, "TZ": tz}).decode().strip() assert result.stdout.strip() == host_offset else: - result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%Z"])) + result = await d.runtime.execute(Command(command=["/bin/sh", "-c", "date +%Z"], sandbox_id=sid)) assert result.stdout.strip() == "UTC" finally: await d.stop() diff --git a/tests/unit/rocklet/test_local_sandbox_runtime.py b/tests/unit/rocklet/test_local_sandbox_runtime.py index ca1ac5f166..d554c85eaa 100644 --- a/tests/unit/rocklet/test_local_sandbox_runtime.py +++ b/tests/unit/rocklet/test_local_sandbox_runtime.py @@ -23,7 +23,9 @@ async def test_upload_file(local_runtime: Rocklet, tmp_path: Path): file_path.write_text("test") tmp_target = tmp_path / "target.txt" await local_runtime.upload(UploadRequest(source_path=str(file_path), target_path=str(tmp_target))) - assert (await local_runtime.read_file(ReadFileRequest(path=str(tmp_target)))).content == "test" + assert ( + await local_runtime.read_file(ReadFileRequest(path=str(tmp_target), sandbox_id="local-test")) + ).content == "test" @pytest.mark.asyncio @@ -34,8 +36,13 @@ async def test_upload_directory(local_runtime: Rocklet, tmp_path: Path): (dir_path / "file2.txt").write_text("test2") tmp_target = tmp_path / "target_dir" await local_runtime.upload(UploadRequest(source_path=str(dir_path), target_path=str(tmp_target))) - assert (await local_runtime.read_file(ReadFileRequest(path=str(tmp_target / "file1.txt")))).content == "test1" - assert (await local_runtime.read_file(ReadFileRequest(path=str(tmp_target / "file2.txt")))).content == "test2" + sid = "local-test" + assert ( + await local_runtime.read_file(ReadFileRequest(path=str(tmp_target / "file1.txt"), sandbox_id=sid)) + ).content == "test1" + assert ( + await local_runtime.read_file(ReadFileRequest(path=str(tmp_target / "file2.txt"), sandbox_id=sid)) + ).content == "test2" @pytest.mark.asyncio @@ -68,14 +75,19 @@ async def test_gem(local_runtime: Rocklet): @pytest.mark.asyncio async def test_prompt_command(local_runtime: Rocklet): prompt_command = "echo ROCK" + sid = "local-test" await local_runtime.create_session( - CreateBashSessionRequest(env={"PROMPT_COMMAND": prompt_command}, session_type="bash") + CreateBashSessionRequest(env={"PROMPT_COMMAND": prompt_command}, session_type="bash", sandbox_id=sid) + ) + without_prompt_command = await local_runtime.run_in_session( + BashAction(command="echo hello", action_type="bash", sandbox_id=sid) ) - without_prompt_command = await local_runtime.run_in_session(BashAction(command="echo hello", action_type="bash")) assert without_prompt_command.output == "hello" await local_runtime.run_in_session( - BashAction(command=f'export PROMPT_COMMAND="{prompt_command}"', action_type="bash") + BashAction(command=f'export PROMPT_COMMAND="{prompt_command}"', action_type="bash", sandbox_id=sid) + ) + with_prompt_command = await local_runtime.run_in_session( + BashAction(command="echo hello", action_type="bash", sandbox_id=sid) ) - with_prompt_command = await local_runtime.run_in_session(BashAction(command="echo hello", action_type="bash")) assert with_prompt_command.output.__contains__("ROCK") - await local_runtime.close_session(CloseBashSessionRequest(session_type="bash")) + await local_runtime.close_session(CloseBashSessionRequest(session_type="bash", sandbox_id=sid)) diff --git a/tests/unit/utils/test_shell_util.py b/tests/unit/utils/test_shell_util.py index a2eb635c92..4d0065d2ca 100644 --- a/tests/unit/utils/test_shell_util.py +++ b/tests/unit/utils/test_shell_util.py @@ -20,10 +20,14 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): session_name = f"bash-{temp_id}" d = LocalDeployment() await d.start() - await d.runtime.create_session(SandboxCreateBashSessionRequest(session=session_name)) + await d.runtime.create_session( + SandboxCreateBashSessionRequest(session=session_name, sandbox_id="local-test"), + ) cmd = f"/bin/bash -c '{cmd}'" nohup_command = f"nohup {cmd} < /dev/null > {out_file} 2>&1 & echo {PID_PREFIX}$!{PID_SUFFIX};disown" - resp = await d.runtime.run_in_session(BashAction(command=nohup_command, session=session_name)) + resp = await d.runtime.run_in_session( + BashAction(command=nohup_command, session=session_name, sandbox_id="local-test") + ) logger.info(f"nohup_command response: {resp.output}") pid = extract_nohup_pid(resp.output) start_time = time.perf_counter() @@ -31,7 +35,9 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): while time.perf_counter() < end_time: try: await asyncio.wait_for( - d.runtime.run_in_session(BashAction(command=f"kill -0 {pid}", session=session_name)), + d.runtime.run_in_session( + BashAction(command=f"kill -0 {pid}", session=session_name, sandbox_id="local-test") + ), timeout=30, ) await asyncio.sleep(1) @@ -39,10 +45,16 @@ async def mock_arun(cmd: str, response_limited_bytes: int = 1024 * 64): print(str(e)) break nohup_resp = await d.runtime.run_in_session( - BashAction(command=f"head -c {response_limited_bytes} {out_file}", session=session_name) + BashAction( + command=f"head -c {response_limited_bytes} {out_file}", + session=session_name, + sandbox_id="local-test", + ) ) yield pid, nohup_resp.output - await d.runtime.run_in_session(BashAction(command=f"rm -rf {out_file}", session=session_name)) + await d.runtime.run_in_session( + BashAction(command=f"rm -rf {out_file}", session=session_name, sandbox_id="local-test") + ) await d.stop() From 6bedee619d7ec773890f74578acb375a6f140c62 Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Tue, 2 Jun 2026 19:10:51 +0800 Subject: [PATCH 163/226] refactor(sandbox): make start() delegate to start_async() to fix missing meta store write (#1051) start() previously managed actor creation independently, skipping meta store persistence. Now it delegates to start_async() and adds a wait-for-alive loop with REQUEST_TIMEOUT_SECONDS (85s) timeout. --- rock/sandbox/sandbox_manager.py | 37 +++++------- .../unit/sandbox/test_sandbox_transitions.py | 58 +++++++++++++++++++ 2 files changed, 71 insertions(+), 24 deletions(-) diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 95cdd146bc..cc78fec538 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -1,4 +1,5 @@ import asyncio +import time from fastapi import UploadFile @@ -38,7 +39,7 @@ from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError, InternalServerRockError -from rock.utils import StageTimer +from rock.utils import REQUEST_TIMEOUT_SECONDS, StageTimer from rock.utils.crypto_utils import AESEncryption from rock.utils.format import convert_to_gb, parse_size_to_bytes from rock.utils.system import get_iso8601_timestamp @@ -176,30 +177,18 @@ async def restart_async(self, sandbox_id: str) -> SandboxStartResponse: @monitor_sandbox_operation() async def start(self, config: DeploymentConfig) -> SandboxStartResponse: - docker_deployment_config: DockerDeploymentConfig = await self.deployment_manager.init_config(config) - - sandbox_id = docker_deployment_config.container_name - actor_name = self.deployment_manager.get_actor_name(sandbox_id) - deployment = docker_deployment_config.get_deployment() - - sandbox_actor: SandboxActor = await deployment.creator_actor(actor_name) - - with StageTimer("startup_timing", f"[{sandbox_id}] Actor start", logger): - await self._ray_service.async_ray_get(sandbox_actor.start.remote()) - logger.info(f"sandbox {sandbox_id} is started") - - with StageTimer("startup_timing", f"[{sandbox_id}] Wait actor alive", logger): - while not await self._is_actor_alive(sandbox_id): - logger.debug(f"wait actor for sandbox alive, sandbox_id: {sandbox_id}") - # TODO: timeout check + response = await self.start_async(config) + sandbox_id = response.sandbox_id + deadline = time.time() + REQUEST_TIMEOUT_SECONDS + with StageTimer("startup_timing", f"[{sandbox_id}] Wait sandbox running", logger): + while True: + status = await self.get_status(sandbox_id) + if status.is_alive: + break + if time.time() >= deadline: + raise TimeoutError(f"sandbox {sandbox_id} not running after {REQUEST_TIMEOUT_SECONDS}s") await asyncio.sleep(1) - await self.get_status(sandbox_id) - - return SandboxStartResponse( - sandbox_id=sandbox_id, - host_name=await self._ray_service.async_ray_get(sandbox_actor.host_name.remote()), - host_ip=await self._ray_service.async_ray_get(sandbox_actor.host_ip.remote()), - ) + return response @monitor_sandbox_operation() async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL): diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py index 1a25ca7624..ce2756d412 100644 --- a/tests/unit/sandbox/test_sandbox_transitions.py +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -302,3 +302,61 @@ async def test_operator_failure_propagates(self, mgr_restart, mock_meta_store, m mock_operator.restart = AsyncMock(side_effect=BadRequestRockError("docker start failed")) with pytest.raises(BadRequestRockError, match="docker start failed"): await mgr_restart.restart_async("sb-1") + + +# --------------------------------------------------------------------------- +# TestManagerStart +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mgr_start(mgr, mock_meta_store, mock_operator, mock_docker_config): + mgr.deployment_manager = MagicMock() + mgr.deployment_manager.init_config = AsyncMock(return_value=mock_docker_config) + mgr._check_sandbox_exists_in_redis = AsyncMock() + mgr.validate_sandbox_spec = MagicMock() + mgr.rock_config = MagicMock() + mgr.rock_config.runtime.use_standard_spec_only = False + + mgr.start_async = SandboxManager.start_async.__wrapped__.__get__(mgr) + mgr._build_sandbox_info_metadata = AsyncMock() + mgr.start = SandboxManager.start.__wrapped__.__get__(mgr) + mgr.get_status = AsyncMock(return_value=MagicMock(is_alive=True, state=State.RUNNING)) + return mgr + + +class TestManagerStart: + @pytest.mark.asyncio + async def test_start_writes_meta_store_and_waits_running(self, mgr_start, mock_meta_store, mock_operator): + config = MagicMock() + config.image = "python:3.11" + result = await mgr_start.start(config) + assert isinstance(result, SandboxStartResponse) + assert result.sandbox_id == "sb-1" + mock_meta_store.create.assert_awaited_once() + mgr_start.get_status.assert_awaited_once_with("sb-1") + + @pytest.mark.asyncio + async def test_start_retries_until_sandbox_running(self, mgr_start): + call_count = 0 + + async def running_after_retries(sandbox_id): + nonlocal call_count + call_count += 1 + return MagicMock(is_alive=call_count >= 3, state=State.RUNNING if call_count >= 3 else State.PENDING) + + mgr_start.get_status = running_after_retries + config = MagicMock() + config.image = "python:3.11" + result = await mgr_start.start(config) + assert isinstance(result, SandboxStartResponse) + assert call_count >= 3 + + @pytest.mark.asyncio + async def test_start_timeout_raises(self, mgr_start): + mgr_start.get_status = AsyncMock(return_value=MagicMock(is_alive=False, state=State.PENDING)) + config = MagicMock() + config.image = "python:3.11" + with patch("rock.sandbox.sandbox_manager.REQUEST_TIMEOUT_SECONDS", 0): + with pytest.raises(TimeoutError, match="not running after"): + await mgr_start.start(config) From 7eecce94d7d9db18b1bea01046babbb6718c0231 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Sat, 16 May 2026 16:46:03 +0800 Subject: [PATCH 164/226] perf(scheduler): switch FileCleanupTask to find -delete and add minimal path safety guards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Performance: - Replace `find ... -exec rm -f {} +` with `find ... -delete` - Replace `find ... -exec rmdir {} +` with `find ... -depth -type d -empty -delete` - `-delete` calls unlink(2) directly without forking per-batch rm/rmdir; on dirs with tens of thousands of files this is ~10x faster. Safety guards: - Add module-level `_DANGEROUS_PATHS` blacklist with 2 entries: "/" and "/tmp/miniforge". The list is intentionally minimal — both are observed incident sources (root wipe, uv runtime deletion). OS dirs like /etc /var are NOT included: they're hypothetical mistakes, and a longer blacklist raises the cost of changing the policy later. - `FileCleanupTask.__init__` rejects target_dirs entries that match a blacklist entry (exact or subtree) -> ValueError at config-load time. Note: "/" is exact-match only — a naive subtree check would reject every absolute path. - `TargetDirConfig.from_raw` rejects empty / relative / `..` paths. The ".." check happens BEFORE os.path.normpath, since normalize would collapse "/data/../etc" to "/etc" and silently accept the traversal. Compat: - yaml format unchanged (str + dict still both accepted). - Behaviour for legitimate paths (/data/logs, /data/sandbox_logs etc.) unchanged — only the find action verb swapped. Tests: - New tests/unit/admin/scheduler/test_file_cleanup_task.py covers from_raw validation, blacklist guard (root exact-match, miniforge exact + subtree, /data/logs explicitly allowed), command generation (-delete present, -exec rm absent), from_config backward-compat, run_action happy + error paths. 31 tests, all passing. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../scheduler/tasks/file_cleanup_task.py | 90 +++++- .../admin/scheduler/test_file_cleanup_task.py | 261 ++++++++++++++++++ 2 files changed, 340 insertions(+), 11 deletions(-) create mode 100644 tests/unit/admin/scheduler/test_file_cleanup_task.py diff --git a/rock/admin/scheduler/tasks/file_cleanup_task.py b/rock/admin/scheduler/tasks/file_cleanup_task.py index 464005f00f..8b4638349c 100644 --- a/rock/admin/scheduler/tasks/file_cleanup_task.py +++ b/rock/admin/scheduler/tasks/file_cleanup_task.py @@ -1,4 +1,5 @@ # rock/admin/scheduler/tasks/file_cleanup_task.py +import os from dataclasses import dataclass, field from rock.admin.proto.request import SandboxCommand as Command @@ -10,6 +11,18 @@ logger = init_logger(name="file_cleanup", file_name=SCHEDULER_LOG_NAME) +# Paths that must NEVER appear in target_dirs because deleting them takes the +# whole worker (or worse, the host) down. Both entries are real incident sources: +# - "/" : `find / -delete` is a full-OS wipe +# - "/tmp/miniforge" : uv-shared Python runtime; if removed, no sandbox can start +# Other "obviously dangerous" OS dirs (/etc, /var, ...) are intentionally NOT here: +# they're hypothetical mistakes, not observed ones, and a longer list raises the +# cost of changing the policy later. +_DANGEROUS_PATHS: tuple[str, ...] = ( + "/", + "/tmp/miniforge", +) + @dataclass class TargetDirConfig: @@ -41,16 +54,38 @@ def from_raw(cls, raw: str | dict) -> "TargetDirConfig": Returns: A TargetDirConfig instance + + Raises: + ValueError: path is empty, relative, or contains ".." traversal. """ if isinstance(raw, str): - return cls(path=raw) - if isinstance(raw, dict): - return cls( + instance = cls(path=raw) + elif isinstance(raw, dict): + instance = cls( path=raw["path"], exclude_dirs=raw.get("exclude_dirs", []), exclude_files=raw.get("exclude_files", []), ) - raise ValueError(f"Unsupported target_dirs entry type: {type(raw)}") + else: + raise ValueError(f"Unsupported target_dirs entry type: {type(raw)}") + + cls._validate_path(instance.path) + return instance + + @staticmethod + def _validate_path(path: str) -> None: + """Reject empty / non-absolute / traversal paths. + + ".." is checked against the *raw* path components — os.path.normpath + would already collapse "/data/../etc" to "/etc", so a post-normalize + check would silently accept the very thing we want to reject. + """ + if not path or not isinstance(path, str): + raise ValueError(f"target_dirs path must be a non-empty string, got: {path!r}") + if not os.path.isabs(path): + raise ValueError(f"target_dirs path must be an absolute path, got: {path!r}") + if ".." in path.split(os.sep): + raise ValueError(f"target_dirs path must not contain '..', got: {path!r}") class FileCleanupTask(BaseTask): @@ -73,16 +108,42 @@ def __init__( and its own exclude_dirs/exclude_files. max_age_mins: Max file age in minutes since last modification, default 7 days (10080 mins) max_file_size: Max file size threshold (e.g. "500M", "1G"), files exceeding this will be removed + + Raises: + ValueError: a target_dirs entry hits _DANGEROUS_PATHS (config-time fail-fast). """ super().__init__( type="file_cleanup", interval_seconds=interval_seconds, idempotency=IdempotencyType.IDEMPOTENT, ) - self.target_dirs = target_dirs or [] + target_dirs = target_dirs or [] + for dc in target_dirs: + self._assert_not_dangerous(dc.path) + self.target_dirs = target_dirs self.max_age_mins = max_age_mins self.max_file_size = max_file_size + @staticmethod + def _assert_not_dangerous(path: str) -> None: + """Reject `path` if it equals a blacklist entry or sits in its subtree. + + "/" is exact-match only: subtree match for "/" would reject every + absolute path (since `path.startswith("/")` is true for any abs path), + which obviously is not the intent — only `target_dirs: ["/"]` itself + should fail. + """ + normalized = os.path.normpath(path) + for dangerous in _DANGEROUS_PATHS: + in_subtree = dangerous != "/" and normalized.startswith(dangerous + "/") + if normalized == dangerous or in_subtree: + raise ValueError( + f"FileCleanupTask refuses dangerous path {path!r} " + f"(matched blacklist entry {dangerous!r}); " + f"these paths are managed by other components, " + f"not by file_cleanup." + ) + @classmethod def from_config(cls, task_config) -> "FileCleanupTask": """Create task instance from config. @@ -168,10 +229,17 @@ def _build_exclude_expr(dir_config: "TargetDirConfig") -> str: def _build_cleanup_command(self, dir_config: TargetDirConfig) -> str: """Build the shell command for cleaning up files in a single directory. - The command performs two steps: - 1. Delete files matching either condition: older than max_age_mins OR larger than max_file_size - (excluding configured directories and files) - 2. Remove empty directories left behind (excluding configured directories) + Performance: + ``find ... -delete`` (vs ``-exec rm -f {} +``) calls unlink(2) + directly without forking a per-batch rm; on dirs with tens of + thousands of files this is roughly an order of magnitude faster. + + Empty-dir cleanup uses the same idiom: ``-depth -type d -empty -delete``. + + Steps: + 1. Delete files (with exclusions) older than max_age_mins OR + exceeding max_file_size. + 2. Remove empty directories left behind (with exclusions). Args: dir_config: The target directory configuration with its exclusions @@ -200,9 +268,9 @@ def _build_cleanup_command(self, dir_config: TargetDirConfig) -> str: f'find "{target_dir}" {exclude_expr}' f"-type f " f"\\( -mmin +{self.max_age_mins} -o {size_find_expr} \\) " - f"-exec rm -f {{}} +; " + f"-delete; " f'find "{target_dir}" -depth {dir_exclude_expr}' - f"-type d -empty -exec rmdir {{}} +; " + f"-type d -empty -delete; " f'echo "cleanup_done"; ' f'else echo "dir_not_found"; fi' ) diff --git a/tests/unit/admin/scheduler/test_file_cleanup_task.py b/tests/unit/admin/scheduler/test_file_cleanup_task.py new file mode 100644 index 0000000000..ddc4085016 --- /dev/null +++ b/tests/unit/admin/scheduler/test_file_cleanup_task.py @@ -0,0 +1,261 @@ +"""Tests for FileCleanupTask: -delete perf swap, blacklist guard, path validation.""" + +from unittest.mock import AsyncMock + +import pytest + +from rock.admin.scheduler.task_base import TaskStatusEnum +from rock.admin.scheduler.tasks.file_cleanup_task import ( + _DANGEROUS_PATHS, + FileCleanupTask, + TargetDirConfig, +) + +# --------------------------------------------------------------------------- # +# Section A: TargetDirConfig.from_raw — backward-compat + path validation +# --------------------------------------------------------------------------- # + + +class TestTargetDirConfigFromRaw: + def test_from_raw_with_string(self): + cfg = TargetDirConfig.from_raw("/data/cache") + assert cfg.path == "/data/cache" + assert cfg.exclude_dirs == [] + assert cfg.exclude_files == [] + + def test_from_raw_with_dict_full(self): + cfg = TargetDirConfig.from_raw( + { + "path": "/data/cache", + "exclude_dirs": [".git", "important"], + "exclude_files": [".gitkeep"], + } + ) + assert cfg.path == "/data/cache" + assert cfg.exclude_dirs == [".git", "important"] + assert cfg.exclude_files == [".gitkeep"] + + def test_from_raw_with_dict_minimal(self): + cfg = TargetDirConfig.from_raw({"path": "/data/cache"}) + assert cfg.exclude_dirs == [] + assert cfg.exclude_files == [] + + @pytest.mark.parametrize("bad", [123, None, ["/data"], ()]) + def test_from_raw_rejects_unsupported_type(self, bad): + with pytest.raises(ValueError, match="Unsupported target_dirs entry type"): + TargetDirConfig.from_raw(bad) + + @pytest.mark.parametrize("bad_path", ["", "relative/path", "./logs"]) + def test_from_raw_rejects_relative_or_empty(self, bad_path): + with pytest.raises(ValueError): + TargetDirConfig.from_raw({"path": bad_path}) + + def test_from_raw_rejects_dotdot(self): + # The check must happen pre-normalize: os.path.normpath collapses + # "/data/../etc" to "/etc" and would otherwise hide the traversal. + with pytest.raises(ValueError, match="must not contain '..'"): + TargetDirConfig.from_raw("/data/../etc") + + +# --------------------------------------------------------------------------- # +# Section B: FileCleanupTask blacklist guard (only "/" and "/tmp/miniforge") +# --------------------------------------------------------------------------- # + + +class TestFileCleanupTaskBlacklist: + def test_init_accepts_safe_paths(self): + task = FileCleanupTask( + target_dirs=[ + TargetDirConfig(path="/data/cache"), + TargetDirConfig(path="/data/workspace"), + ], + ) + assert len(task.target_dirs) == 2 + + def test_init_rejects_root(self): + with pytest.raises(ValueError, match="dangerous path"): + FileCleanupTask(target_dirs=[TargetDirConfig(path="/")]) + + def test_init_rejects_miniforge_exact(self): + with pytest.raises(ValueError, match="dangerous path"): + FileCleanupTask(target_dirs=[TargetDirConfig(path="/tmp/miniforge")]) + + def test_init_rejects_miniforge_subtree(self): + with pytest.raises(ValueError, match="dangerous path"): + FileCleanupTask(target_dirs=[TargetDirConfig(path="/tmp/miniforge/python311")]) + + def test_init_allows_data_logs(self): + # /data/logs is intentionally NOT blacklisted — internal yml relies on + # configuring it with exclude_files. This regression guards against + # future "tighten the blacklist" PRs that would break deployments. + task = FileCleanupTask(target_dirs=[TargetDirConfig(path="/data/logs")]) + assert task.target_dirs[0].path == "/data/logs" + + def test_init_allows_os_dirs_not_in_minimal_blacklist(self): + # /etc, /var, /usr are intentionally NOT in our minimal blacklist. + # "/" is exact-match only; would otherwise reject every absolute path. + # OS-dir misconfig is a hypothetical mistake, not an observed one; + # keeping the blacklist minimal makes future policy changes easier. + task = FileCleanupTask(target_dirs=[TargetDirConfig(path="/etc")]) + assert task.target_dirs[0].path == "/etc" + + def test_dangerous_blacklist_is_minimal(self): + # Exactly two entries — keep this list short on purpose. + assert _DANGEROUS_PATHS == ("/", "/tmp/miniforge") + + +# --------------------------------------------------------------------------- # +# Section C: _build_cleanup_command — uses -delete (the perf change) +# --------------------------------------------------------------------------- # + + +class TestBuildCleanupCommand: + def _new_task(self, **kwargs): + return FileCleanupTask( + target_dirs=kwargs.pop("target_dirs", [TargetDirConfig(path="/data/cache")]), + max_age_mins=kwargs.pop("max_age_mins", 10080), + max_file_size=kwargs.pop("max_file_size", "1G"), + **kwargs, + ) + + def test_command_uses_delete_for_files(self): + task = self._new_task() + cmd = task._build_cleanup_command(TargetDirConfig(path="/data/cache")) + assert "-delete;" in cmd + # Old forms must be gone — guards against accidental revert. + assert "-exec rm -f" not in cmd + assert "-exec rm " not in cmd + + def test_command_uses_delete_for_empty_dirs(self): + task = self._new_task() + cmd = task._build_cleanup_command(TargetDirConfig(path="/data/cache")) + assert "-type d -empty -delete" in cmd + assert "-exec rmdir" not in cmd + + def test_command_includes_target_dir_existence_check(self): + task = self._new_task() + cmd = task._build_cleanup_command(TargetDirConfig(path="/data/cache")) + assert 'if [ -d "/data/cache" ]; then' in cmd + assert 'else echo "dir_not_found"; fi' in cmd + assert 'echo "cleanup_done"' in cmd + + def test_command_includes_age_and_size_predicates(self): + task = self._new_task(max_age_mins=4320, max_file_size="500M") + cmd = task._build_cleanup_command(TargetDirConfig(path="/data/cache")) + assert "-mmin +4320" in cmd + # 500M = 500 * 1024 * 1024 = 524288000 + assert "-size +524288000c" in cmd + + def test_command_with_exclude_dirs(self): + dir_cfg = TargetDirConfig(path="/data/cache", exclude_dirs=["keep_me"]) + task = self._new_task(target_dirs=[dir_cfg]) + cmd = task._build_cleanup_command(dir_cfg) + assert '-name "keep_me" -prune' in cmd + + def test_command_with_exclude_files(self): + dir_cfg = TargetDirConfig(path="/data/cache", exclude_files=["important.log"]) + task = self._new_task(target_dirs=[dir_cfg]) + cmd = task._build_cleanup_command(dir_cfg) + assert '-name "important.log" -prune' in cmd + + +# --------------------------------------------------------------------------- # +# Section D: from_config — yaml -> task instance, backward compat, fail-fast +# --------------------------------------------------------------------------- # + + +class _FakeTaskConfig: + """Lightweight stand-in for rock.config.TaskConfig in unit tests.""" + + def __init__(self, params, interval_seconds=86400): + self.params = params + self.interval_seconds = interval_seconds + + +class TestFromConfig: + def test_from_config_legacy_string_format(self): + task_config = _FakeTaskConfig( + params={ + "target_dirs": ["/data/cache", "/data/scratch"], + "max_age_mins": 1440, + "max_file_size": "500M", + } + ) + task = FileCleanupTask.from_config(task_config) + assert [dc.path for dc in task.target_dirs] == ["/data/cache", "/data/scratch"] + assert task.max_age_mins == 1440 + assert task.max_file_size == "500M" + + def test_from_config_new_dict_format(self): + task_config = _FakeTaskConfig( + params={ + "target_dirs": [ + {"path": "/data/cache", "exclude_dirs": ["keep"], "exclude_files": ["KEEP.txt"]}, + "/data/scratch", + ], + } + ) + task = FileCleanupTask.from_config(task_config) + assert task.target_dirs[0].path == "/data/cache" + assert task.target_dirs[0].exclude_dirs == ["keep"] + assert task.target_dirs[1].path == "/data/scratch" + + def test_from_config_defaults_when_missing(self): + task_config = _FakeTaskConfig(params={"target_dirs": ["/data/cache"]}) + task = FileCleanupTask.from_config(task_config) + assert task.max_age_mins == 10080 + assert task.max_file_size == "1G" + + def test_from_config_rejects_dangerous_yaml(self): + # Regression: yaml that contains /tmp/miniforge MUST fail at load time. + task_config = _FakeTaskConfig(params={"target_dirs": ["/tmp/miniforge"]}) + with pytest.raises(ValueError, match="dangerous path"): + FileCleanupTask.from_config(task_config) + + +# --------------------------------------------------------------------------- # +# Section E: run_action — happy path uses -delete, errors propagate +# --------------------------------------------------------------------------- # + + +class _FakeExecResult: + def __init__(self, exit_code=0, stdout="cleanup_done"): + self.exit_code = exit_code + self.stdout = stdout + + +class TestRunAction: + @pytest.mark.asyncio + async def test_run_action_no_target_dirs(self): + task = FileCleanupTask(target_dirs=[]) + result = await task.run_action(runtime=AsyncMock()) + assert result["status"] == TaskStatusEnum.SUCCESS + assert "no target directories" in result["message"] + + @pytest.mark.asyncio + async def test_run_action_executes_delete_command(self): + task = FileCleanupTask(target_dirs=[TargetDirConfig(path="/data/cache")]) + + runtime = AsyncMock() + runtime._config = type("C", (), {"host": "10.0.0.1"})() + runtime.execute = AsyncMock(return_value=_FakeExecResult()) + + result = await task.run_action(runtime) + assert result["status"] == TaskStatusEnum.SUCCESS + assert result["target_dirs"] == ["/data/cache"] + + # Verify the executed shell command actually used -delete (not -exec rm). + executed_cmd = runtime.execute.await_args.args[0].command + assert "-delete" in executed_cmd + assert "-exec rm" not in executed_cmd + + @pytest.mark.asyncio + async def test_run_action_re_raises_on_error(self): + task = FileCleanupTask(target_dirs=[TargetDirConfig(path="/data/cache")]) + + runtime = AsyncMock() + runtime._config = type("C", (), {"host": "10.0.0.1"})() + runtime.execute = AsyncMock(side_effect=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await task.run_action(runtime) From 300ae25c4bd93b6a45c0d8fc1b411378adde0514 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 11:29:05 +0800 Subject: [PATCH 165/226] refactor(scheduler): rename _DANGEROUS_PATHS to _PATH_BLACKLIST for consistency with blacklist terminology (PR #967 review) --- rock/admin/scheduler/tasks/file_cleanup_task.py | 17 +++++++++-------- .../admin/scheduler/test_file_cleanup_task.py | 6 +++--- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/rock/admin/scheduler/tasks/file_cleanup_task.py b/rock/admin/scheduler/tasks/file_cleanup_task.py index 8b4638349c..b31946b22a 100644 --- a/rock/admin/scheduler/tasks/file_cleanup_task.py +++ b/rock/admin/scheduler/tasks/file_cleanup_task.py @@ -11,14 +11,15 @@ logger = init_logger(name="file_cleanup", file_name=SCHEDULER_LOG_NAME) -# Paths that must NEVER appear in target_dirs because deleting them takes the -# whole worker (or worse, the host) down. Both entries are real incident sources: +# Path blacklist: entries that must NEVER appear in target_dirs because deleting +# them takes the whole worker (or worse, the host) down. Both entries are real +# incident sources: # - "/" : `find / -delete` is a full-OS wipe # - "/tmp/miniforge" : uv-shared Python runtime; if removed, no sandbox can start # Other "obviously dangerous" OS dirs (/etc, /var, ...) are intentionally NOT here: # they're hypothetical mistakes, not observed ones, and a longer list raises the # cost of changing the policy later. -_DANGEROUS_PATHS: tuple[str, ...] = ( +_PATH_BLACKLIST: tuple[str, ...] = ( "/", "/tmp/miniforge", ) @@ -110,7 +111,7 @@ def __init__( max_file_size: Max file size threshold (e.g. "500M", "1G"), files exceeding this will be removed Raises: - ValueError: a target_dirs entry hits _DANGEROUS_PATHS (config-time fail-fast). + ValueError: a target_dirs entry hits _PATH_BLACKLIST (config-time fail-fast). """ super().__init__( type="file_cleanup", @@ -134,12 +135,12 @@ def _assert_not_dangerous(path: str) -> None: should fail. """ normalized = os.path.normpath(path) - for dangerous in _DANGEROUS_PATHS: - in_subtree = dangerous != "/" and normalized.startswith(dangerous + "/") - if normalized == dangerous or in_subtree: + for entry in _PATH_BLACKLIST: + in_subtree = entry != "/" and normalized.startswith(entry + "/") + if normalized == entry or in_subtree: raise ValueError( f"FileCleanupTask refuses dangerous path {path!r} " - f"(matched blacklist entry {dangerous!r}); " + f"(matched blacklist entry {entry!r}); " f"these paths are managed by other components, " f"not by file_cleanup." ) diff --git a/tests/unit/admin/scheduler/test_file_cleanup_task.py b/tests/unit/admin/scheduler/test_file_cleanup_task.py index ddc4085016..ce87259227 100644 --- a/tests/unit/admin/scheduler/test_file_cleanup_task.py +++ b/tests/unit/admin/scheduler/test_file_cleanup_task.py @@ -6,7 +6,7 @@ from rock.admin.scheduler.task_base import TaskStatusEnum from rock.admin.scheduler.tasks.file_cleanup_task import ( - _DANGEROUS_PATHS, + _PATH_BLACKLIST, FileCleanupTask, TargetDirConfig, ) @@ -99,9 +99,9 @@ def test_init_allows_os_dirs_not_in_minimal_blacklist(self): task = FileCleanupTask(target_dirs=[TargetDirConfig(path="/etc")]) assert task.target_dirs[0].path == "/etc" - def test_dangerous_blacklist_is_minimal(self): + def test_path_blacklist_is_minimal(self): # Exactly two entries — keep this list short on purpose. - assert _DANGEROUS_PATHS == ("/", "/tmp/miniforge") + assert _PATH_BLACKLIST == ("/", "/tmp/miniforge") # --------------------------------------------------------------------------- # From 3576823b1d9f911ac41dff95c7832a6f46a044c5 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Thu, 28 May 2026 14:33:22 +0800 Subject: [PATCH 166/226] fix: replace os.path ops with string checks for cross-platform CI --- rock/admin/scheduler/tasks/file_cleanup_task.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rock/admin/scheduler/tasks/file_cleanup_task.py b/rock/admin/scheduler/tasks/file_cleanup_task.py index b31946b22a..ef22d521f6 100644 --- a/rock/admin/scheduler/tasks/file_cleanup_task.py +++ b/rock/admin/scheduler/tasks/file_cleanup_task.py @@ -83,9 +83,9 @@ def _validate_path(path: str) -> None: """ if not path or not isinstance(path, str): raise ValueError(f"target_dirs path must be a non-empty string, got: {path!r}") - if not os.path.isabs(path): + if not path.startswith("/"): raise ValueError(f"target_dirs path must be an absolute path, got: {path!r}") - if ".." in path.split(os.sep): + if ".." in path.split("/"): raise ValueError(f"target_dirs path must not contain '..', got: {path!r}") @@ -134,9 +134,9 @@ def _assert_not_dangerous(path: str) -> None: which obviously is not the intent — only `target_dirs: ["/"]` itself should fail. """ - normalized = os.path.normpath(path) + normalized = path.rstrip("/") or "/" for entry in _PATH_BLACKLIST: - in_subtree = entry != "/" and normalized.startswith(entry + "/") + in_subtree = entry != "/" and (normalized + "/").startswith(entry + "/") if normalized == entry or in_subtree: raise ValueError( f"FileCleanupTask refuses dangerous path {path!r} " From 09c1a14eda84418007085bef49c4377e60156e16 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Mon, 1 Jun 2026 15:44:35 +0800 Subject: [PATCH 167/226] fix(test): add os.makedirs/chmod mocks to disk_limit test so _try_set_log_dir_quota is reached --- .../test_docker_deployment_disk_limit.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/tests/unit/deployments/test_docker_deployment_disk_limit.py b/tests/unit/deployments/test_docker_deployment_disk_limit.py index a678a383d2..2a58d74830 100644 --- a/tests/unit/deployments/test_docker_deployment_disk_limit.py +++ b/tests/unit/deployments/test_docker_deployment_disk_limit.py @@ -148,3 +148,84 @@ async def test_no_error_when_rootfs_already_none(self, _mock_detect, _mock_valid assert deployment.config.disk_limit_rootfs is None assert deployment.effective_disk_limit_rootfs is None + + @pytest.mark.asyncio + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=True) + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + async def test_log_downgraded_when_not_xfs_prjquota(self, _mock_prjquota, _mock_detect, _mock_validator): + """When log path is not XFS+prjquota: effective_disk_limit_log=None; config unchanged. + + Note: log quota has NO dependency on docker being overlay2 — + is_xfs_prjquota_path() is the only gate. + """ + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + _make_start_mocks(deployment) + + with ( + patch("rock.deployments.docker.get_executor"), + patch("rock.deployments.docker.asyncio.get_running_loop") as mock_loop, + patch("rock.deployments.docker.wait_until_alive", new_callable=AsyncMock), + patch("rock.deployments.docker.env_vars") as mock_env, + patch("rock.deployments.docker.subprocess"), + patch("rock.deployments.docker.os.makedirs"), + patch("rock.deployments.docker.os.chmod"), + ): + mock_env.ROCK_LOGGING_PATH = "/var/log/rock" + mock_env.ROCK_TIME_ZONE = "UTC" + mock_loop.return_value.run_in_executor = AsyncMock() + try: + await deployment.start() + except Exception: + pass + + assert deployment.config.disk_limit_log == "5g" + assert deployment.effective_disk_limit_log is None + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + def test_log_not_downgraded_when_no_log_path(self, _mock_prjquota, _mock_validator): + """When ROCK_LOGGING_PATH is empty, _try_set_log_dir_quota is never called, + so effective_disk_limit_log remains equal to config.disk_limit_log.""" + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + # effective starts equal to config before start() is called + assert deployment.effective_disk_limit_log == "5g" + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) + def test_try_set_log_dir_quota_downgrades_when_not_xfs_prjquota(self, _mock_prjquota, _mock_validator): + """_try_set_log_dir_quota: is_xfs_prjquota_path=False → effective_disk_limit_log=None.""" + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + deployment._effective_disk_limit_log = "5g" + deployment._container_name = "test-container" + + deployment._try_set_log_dir_quota("/var/log/rock/test-container") + + assert deployment.effective_disk_limit_log is None + + @patch("rock.deployments.docker.DockerSandboxValidator") + @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=True) + def test_try_set_log_dir_quota_independent_of_docker_driver(self, _mock_prjquota, _mock_validator): + """_try_set_log_dir_quota passes the XFS gate regardless of Docker storage driver. + + Log quota only requires is_xfs_prjquota_path(); overlay2 is irrelevant. + The subprocess calls inside (findmnt, xfs_quota) are mocked to succeed. + """ + config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") + deployment = DockerDeployment.from_config(config) + deployment._effective_disk_limit_log = "5g" + deployment._container_name = "test-container" + + with patch("rock.deployments.docker.subprocess") as mock_sub: + ok = MagicMock() + ok.returncode = 0 + ok.stdout = "/var/log/rock" + mock_sub.run.return_value = ok + deployment._try_set_log_dir_quota("/var/log/rock/test-container") + + # xfs_quota succeeded → effective value preserved and prjid recorded + assert deployment.effective_disk_limit_log == "5g" + assert deployment.log_dir_xfs_prjid is not None From a78e512b73edec0fd09630c89f8d3db5261122ea Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Tue, 2 Jun 2026 21:42:36 +0800 Subject: [PATCH 168/226] fix(test): restore docker disk-limit tests to master, drop stale log-dir-quota cases from bad rebase --- .../test_docker_deployment_disk_limit.py | 81 ------------------- 1 file changed, 81 deletions(-) diff --git a/tests/unit/deployments/test_docker_deployment_disk_limit.py b/tests/unit/deployments/test_docker_deployment_disk_limit.py index 2a58d74830..a678a383d2 100644 --- a/tests/unit/deployments/test_docker_deployment_disk_limit.py +++ b/tests/unit/deployments/test_docker_deployment_disk_limit.py @@ -148,84 +148,3 @@ async def test_no_error_when_rootfs_already_none(self, _mock_detect, _mock_valid assert deployment.config.disk_limit_rootfs is None assert deployment.effective_disk_limit_rootfs is None - - @pytest.mark.asyncio - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.detect_storage_opt_support", return_value=True) - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - async def test_log_downgraded_when_not_xfs_prjquota(self, _mock_prjquota, _mock_detect, _mock_validator): - """When log path is not XFS+prjquota: effective_disk_limit_log=None; config unchanged. - - Note: log quota has NO dependency on docker being overlay2 — - is_xfs_prjquota_path() is the only gate. - """ - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - _make_start_mocks(deployment) - - with ( - patch("rock.deployments.docker.get_executor"), - patch("rock.deployments.docker.asyncio.get_running_loop") as mock_loop, - patch("rock.deployments.docker.wait_until_alive", new_callable=AsyncMock), - patch("rock.deployments.docker.env_vars") as mock_env, - patch("rock.deployments.docker.subprocess"), - patch("rock.deployments.docker.os.makedirs"), - patch("rock.deployments.docker.os.chmod"), - ): - mock_env.ROCK_LOGGING_PATH = "/var/log/rock" - mock_env.ROCK_TIME_ZONE = "UTC" - mock_loop.return_value.run_in_executor = AsyncMock() - try: - await deployment.start() - except Exception: - pass - - assert deployment.config.disk_limit_log == "5g" - assert deployment.effective_disk_limit_log is None - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - def test_log_not_downgraded_when_no_log_path(self, _mock_prjquota, _mock_validator): - """When ROCK_LOGGING_PATH is empty, _try_set_log_dir_quota is never called, - so effective_disk_limit_log remains equal to config.disk_limit_log.""" - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - # effective starts equal to config before start() is called - assert deployment.effective_disk_limit_log == "5g" - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=False) - def test_try_set_log_dir_quota_downgrades_when_not_xfs_prjquota(self, _mock_prjquota, _mock_validator): - """_try_set_log_dir_quota: is_xfs_prjquota_path=False → effective_disk_limit_log=None.""" - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - deployment._effective_disk_limit_log = "5g" - deployment._container_name = "test-container" - - deployment._try_set_log_dir_quota("/var/log/rock/test-container") - - assert deployment.effective_disk_limit_log is None - - @patch("rock.deployments.docker.DockerSandboxValidator") - @patch("rock.deployments.docker.DockerUtil.is_xfs_prjquota_path", return_value=True) - def test_try_set_log_dir_quota_independent_of_docker_driver(self, _mock_prjquota, _mock_validator): - """_try_set_log_dir_quota passes the XFS gate regardless of Docker storage driver. - - Log quota only requires is_xfs_prjquota_path(); overlay2 is irrelevant. - The subprocess calls inside (findmnt, xfs_quota) are mocked to succeed. - """ - config = DockerDeploymentConfig(disk_limit_log="5g", image="python:3.11") - deployment = DockerDeployment.from_config(config) - deployment._effective_disk_limit_log = "5g" - deployment._container_name = "test-container" - - with patch("rock.deployments.docker.subprocess") as mock_sub: - ok = MagicMock() - ok.returncode = 0 - ok.stdout = "/var/log/rock" - mock_sub.run.return_value = ok - deployment._try_set_log_dir_quota("/var/log/rock/test-container") - - # xfs_quota succeeded → effective value preserved and prjid recorded - assert deployment.effective_disk_limit_log == "5g" - assert deployment.log_dir_xfs_prjid is not None From 8ade404d453330bc12c5151a20a60b2911516aa4 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Wed, 27 May 2026 17:20:27 +0800 Subject: [PATCH 169/226] =?UTF-8?q?docs:=20refresh=20README=20Updates=20wi?= =?UTF-8?q?th=20v1.4.0=20=E2=80=93=20v1.8.0=20minor-release=20lineup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 3e4c86e15a..53f31c4909 100644 --- a/README.md +++ b/README.md @@ -155,11 +155,11 @@ if __name__ == "__main__": | Date | Release | |:-------------|:---| -| **[Latest]** | 🎉 ROCK v1.7.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.7.0) | -| **[2026-04-23]** | 🎉 ROCK v1.5.1 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.5.1) | -| **[2026-04-10]** | 🎉 ROCK v1.4.7 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.7) | -| **[2026-03-27]** | 🎉 ROCK v1.4.4 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.4) | -| **[2026-03-24]** | 🎉 ROCK v1.4.3 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/Release%20Notes/v1.4.3) | +| **[Latest]** | 🎉 ROCK v1.8.0 Released (2026-05-21) — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | +| **[2026-04-24]** | 🎉 ROCK v1.7.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | +| **[2026-04-17]** | 🎉 ROCK v1.6.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | +| **[2026-04-10]** | 🎉 ROCK v1.5.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | +| **[2026-03-14]** | 🎉 ROCK v1.4.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | --- From fe10f954737da5ce78d7f588d0484c0e4ecf6910 Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Wed, 3 Jun 2026 10:31:06 +0800 Subject: [PATCH 170/226] fix: correct release notes links in README to version-specific pages --- README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 53f31c4909..c04d07e767 100644 --- a/README.md +++ b/README.md @@ -155,11 +155,11 @@ if __name__ == "__main__": | Date | Release | |:-------------|:---| -| **[Latest]** | 🎉 ROCK v1.8.0 Released (2026-05-21) — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | -| **[2026-04-24]** | 🎉 ROCK v1.7.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | -| **[2026-04-17]** | 🎉 ROCK v1.6.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | -| **[2026-04-10]** | 🎉 ROCK v1.5.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | -| **[2026-03-14]** | 🎉 ROCK v1.4.0 Released — [Release Notes](https://alibaba.github.io/ROCK/docs/overview) | +| **[Latest]** | 🎉 ROCK v1.8.0 Released (2026-05-21) — [Release Notes](https://rock.io.alibaba-inc.com/docs/1.8.x/Release%20Notes/v1.8.0) | +| **[2026-04-24]** | 🎉 ROCK v1.7.0 Released — [Release Notes](https://rock.io.alibaba-inc.com/docs/1.7.x/Release%20Notes/v1.7.0) | +| **[2026-04-17]** | 🎉 ROCK v1.6.0 Released — [Release Notes](https://rock.io.alibaba-inc.com/docs/1.6.x/Release%20Notes/v1.6.0) | +| **[2026-04-10]** | 🎉 ROCK v1.5.0 Released — [Release Notes](https://rock.io.alibaba-inc.com/docs/1.5.x/Release%20Notes/v1.5.0) | +| **[2026-03-14]** | 🎉 ROCK v1.4.0 Released — [Release Notes](https://rock.io.alibaba-inc.com/docs/1.4.x/Release%20Notes/v1.4.4) | --- From f4618cde81a70a44f4d5965efbd9202be85c5ec1 Mon Sep 17 00:00:00 2001 From: jinbai340997 <15652831212@163.com> Date: Wed, 3 Jun 2026 17:44:04 +0800 Subject: [PATCH 171/226] feat(admin): ops-jobs API with DB-persisted state, multi-pod safe (#1027) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * added operation and maintenance API repository capability, supporting operation and maintenance API to initiate disk cleanup tasks. * fix(scheduler): key get_task_registry by task.type, not full class path * fix(admin): build ops_jobs task registry per pod, not via scheduler thread (any pod can POST) * fix(admin): always run create_tables() so new ORM models auto-create on real DB * refactor rock ops database design * add rock ops test * refactor(admin-ops): split into api/service/table layers, Phase enum, typed CRUD, result→status * fix(scheduler-task-table): add expire_on_commit=False to prevent detached instance error * fix(ops): return rejected/rateLimited conditions in TaskSetResponse instead of discarding --- rock/admin/core/scheduler_task_table.py | 71 +++++ rock/admin/core/schema.py | 18 ++ rock/admin/entrypoints/admin_ops_api.py | 54 ++++ rock/admin/main.py | 41 ++- rock/admin/proto/request.py | 9 + rock/admin/proto/response.py | 44 +++ rock/admin/scheduler/scheduler.py | 22 ++ rock/admin/service/__init__.py | 0 rock/admin/service/ops_service.py | 237 ++++++++++++++++ .../admin/core/test_scheduler_task_table.py | 134 +++++++++ .../admin/entrypoints/test_admin_ops_api.py | 256 ++++++++++++++++++ 11 files changed, 881 insertions(+), 5 deletions(-) create mode 100644 rock/admin/core/scheduler_task_table.py create mode 100644 rock/admin/entrypoints/admin_ops_api.py create mode 100644 rock/admin/service/__init__.py create mode 100644 rock/admin/service/ops_service.py create mode 100644 tests/unit/admin/core/test_scheduler_task_table.py create mode 100644 tests/unit/admin/entrypoints/test_admin_ops_api.py diff --git a/rock/admin/core/scheduler_task_table.py b/rock/admin/core/scheduler_task_table.py new file mode 100644 index 0000000000..713d66764a --- /dev/null +++ b/rock/admin/core/scheduler_task_table.py @@ -0,0 +1,71 @@ +"""SchedulerTaskTable: single-table CRUD for scheduler task executions. + +Tasks are grouped by taskset_id (one group per API call). +""" + +from __future__ import annotations + +from enum import Enum + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from rock.admin.core.db_provider import DatabaseProvider +from rock.admin.core.sandbox_table import _retry_on_disconnect +from rock.admin.core.schema import SchedulerTaskRecord +from rock.logger import init_logger + +logger = init_logger(__name__) + + +class Phase(str, Enum): + PENDING = "Pending" + RUNNING = "Running" + SUCCEEDED = "Succeeded" + FAILED = "Failed" + REJECTED = "Rejected" + RATE_LIMITED = "RateLimited" + NOT_FOUND = "NotFound" + + +class SchedulerTaskTable: + def __init__(self, db_provider: DatabaseProvider) -> None: + self._db = db_provider + + @_retry_on_disconnect + async def insert_tasks(self, records: list[SchedulerTaskRecord]) -> None: + async with AsyncSession(self._db.engine, expire_on_commit=False) as session: + for r in records: + session.add(r) + await session.commit() + + @_retry_on_disconnect + async def get_tasks_by_group(self, taskset_id: str) -> list[SchedulerTaskRecord]: + async with AsyncSession(self._db.engine, expire_on_commit=False) as session: + stmt = select(SchedulerTaskRecord).where(SchedulerTaskRecord.taskset_id == taskset_id) + rows = (await session.execute(stmt)).scalars().all() + return list(rows) + + @_retry_on_disconnect + async def update_task(self, task_id: str, **fields) -> bool: + async with AsyncSession(self._db.engine) as session: + row = await session.get(SchedulerTaskRecord, task_id) + if row is None: + return False + for k, v in fields.items(): + setattr(row, k, v) + await session.commit() + return True + + @_retry_on_disconnect + async def has_recent_task(self, task_type: str, since_epoch: float) -> bool: + async with AsyncSession(self._db.engine) as session: + stmt = ( + select(SchedulerTaskRecord.task_id) + .where( + SchedulerTaskRecord.task_type == task_type, SchedulerTaskRecord.creation_timestamp >= since_epoch + ) + .limit(1) + ) + row = (await session.execute(stmt)).first() + return row is not None diff --git a/rock/admin/core/schema.py b/rock/admin/core/schema.py index 7e5b78b760..d9e22c5579 100644 --- a/rock/admin/core/schema.py +++ b/rock/admin/core/schema.py @@ -102,3 +102,21 @@ def column_names(cls) -> set[str]: def to_dict(self) -> dict[str, Any]: """Return all non-``None`` column values as a plain dict.""" return {c.key: getattr(self, c.key) for c in self.__table__.columns if getattr(self, c.key) is not None} + + +class SchedulerTaskRecord(Base): + """One row per scheduled task execution. Grouped by taskset_id.""" + + __tablename__ = "scheduler_task" + + task_id = Column(String(32), primary_key=True) + taskset_id = Column(String(32), nullable=True, index=True) + task_type = Column(String(64), nullable=False, index=True) + target_workers = Column(_JSONB_VARIANT, nullable=True) + creation_timestamp = Column(Float, nullable=False) + phase = Column(String(32), nullable=False, default="Pending") + assigned_pod = Column(String(128), nullable=False, default="") + start_time = Column(Float, nullable=True) + completion_time = Column(Float, nullable=True) + conditions = Column(_JSONB_VARIANT, nullable=True) + status = Column(_JSONB_VARIANT, nullable=True) diff --git a/rock/admin/entrypoints/admin_ops_api.py b/rock/admin/entrypoints/admin_ops_api.py new file mode 100644 index 0000000000..a0a13f1ab7 --- /dev/null +++ b/rock/admin/entrypoints/admin_ops_api.py @@ -0,0 +1,54 @@ +"""Admin ops API: thin routing layer for TaskSet operations. + +Endpoints (relative to prefix /apis/envs/sandbox/v1/ops): +- POST /tasksets Create a TaskSet (triggers tasks on workers) +- GET /tasksets/{taskset_id} Query TaskSet with child tasks +""" + +from fastapi import APIRouter, Body, Request + +from rock.actions.response import ResponseStatus, RockResponse +from rock.admin.proto.request import CreateTaskSetRequest +from rock.admin.service.ops_service import OpsService +from rock.common.exception import handle_exceptions + +admin_ops_router = APIRouter() + +_ops_service: OpsService | None = None + + +def set_ops_service(service: OpsService) -> None: + global _ops_service + _ops_service = service + + +@admin_ops_router.post("/tasksets") +@handle_exceptions(error_message="create taskset failed") +async def create_taskset( + request: Request, + payload: CreateTaskSetRequest = Body(default_factory=CreateTaskSetRequest), +) -> RockResponse[dict]: + if _ops_service is None: + return RockResponse( + status=ResponseStatus.FAILED, + message="ops service not initialised", + error="server misconfigured", + ) + + caller = request.client.host if request.client else "unknown" + resp = await _ops_service.create_taskset(payload.spec, caller) + return RockResponse(status=ResponseStatus.SUCCESS, message="ok", result=resp.model_dump()) + + +@admin_ops_router.get("/tasksets/{taskset_id}") +@handle_exceptions(error_message="get taskset failed") +async def get_taskset(taskset_id: str) -> RockResponse[dict]: + if _ops_service is None: + return RockResponse( + status=ResponseStatus.FAILED, + message="ops service not initialised", + error="server misconfigured", + ) + + resp = await _ops_service.get_taskset(taskset_id) + return RockResponse(status=ResponseStatus.SUCCESS, message="ok", result=resp.model_dump()) diff --git a/rock/admin/main.py b/rock/admin/main.py index 61e322b000..a3c7537f40 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -18,11 +18,15 @@ from rock.admin.core.db_provider import DatabaseProvider from rock.admin.core.ray_service import RayService from rock.admin.core.sandbox_table import SandboxTable +from rock.admin.core.scheduler_task_table import SchedulerTaskTable +from rock.admin.entrypoints.admin_ops_api import admin_ops_router, set_ops_service from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service -from rock.admin.scheduler.scheduler import SchedulerThread +from rock.admin.scheduler.scheduler import SchedulerThread, WorkerIPCache +from rock.admin.scheduler.task_base import BaseTask +from rock.admin.scheduler.task_factory import TaskFactory from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( set_main_loop_provider as set_archive_main_loop_provider, ) @@ -32,6 +36,7 @@ from rock.admin.scheduler.tasks.sandbox_log_archive_task import ( set_sandbox_table_provider as set_archive_sandbox_table_provider, ) +from rock.admin.service.ops_service import OpsService from rock.common.exception import request_validation_exception_handler from rock.config import DatabaseConfig, RockConfig, SchedulerConfig from rock.logger import init_logger @@ -55,6 +60,30 @@ logging.getLogger("urllib3").setLevel(logging.WARNING) +def _init_ops_service( + rock_config: RockConfig, + scheduler_task_table: "SchedulerTaskTable", +) -> OpsService: + """Build OpsService with task registry from scheduler config.""" + ops_task_registry: dict[str, BaseTask] = {} + if rock_config.scheduler.enabled: + for task_config in rock_config.scheduler.tasks: + if not getattr(task_config, "enabled", True): + continue + try: + task = TaskFactory.create_task(task_config) + ops_task_registry[task.type] = task + except Exception as e: + logger.warning(f"ops_taskset: failed to instantiate '{task_config.task_class}': {e}") + + ops_worker_cache = WorkerIPCache(cache_ttl=rock_config.scheduler.worker_cache_ttl) + return OpsService( + task_table=scheduler_task_table, + task_registry=ops_task_registry, + alive_workers_provider=ops_worker_cache.get_alive_workers, + ) + + @asynccontextmanager async def lifespan(app: FastAPI): config_file_path = ( @@ -105,13 +134,12 @@ async def lifespan(app: FastAPI): # config reload propagates to the next task run without re-injection. set_archive_sandbox_table_provider(lambda: sandbox_table) set_archive_rock_config_provider(lambda: rock_config) - # Capture lifespan loop (uvicorn main loop). SandboxLogArchiveTask runs - # inside SchedulerThread's child loop; it must dispatch DB calls back to - # this main loop so asyncpg pool stays bound here and HTTP handlers don't - # break with "Future attached to a different loop". _main_loop = asyncio.get_running_loop() set_archive_main_loop_provider(lambda: _main_loop) + # init scheduler task table (DB-backed, multi-pod safe) + scheduler_task_table = SchedulerTaskTable(db_provider) + # init scheduler thread scheduler_thread = None @@ -166,6 +194,8 @@ async def lifespan(app: FastAPI): elif rock_config.scheduler.enabled: logger.info("Scheduler thread skipped on non-primary pod") + set_ops_service(_init_ops_service(rock_config, scheduler_task_table)) + else: sandbox_manager = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) set_sandbox_proxy_service(sandbox_manager) @@ -271,6 +301,7 @@ def main(): # config router if args.role == "admin": app.include_router(sandbox_router, prefix="/apis/envs/sandbox/v1", tags=["sandbox"]) + app.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops", tags=["admin-ops"]) else: app.include_router(sandbox_proxy_router, prefix="/apis/envs/sandbox/v1", tags=["sandbox"]) app.include_router(warmup_router, prefix="/apis/envs/sandbox/v1", tags=["warmup"]) diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index 7954b1af56..99398a27a2 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -141,6 +141,15 @@ class UserInfo(TypedDict, total=False): rock_authorization: str +class TaskSetSpec(BaseModel): + taskTypes: list[str] | None = Field(default=None) + targetWorkers: list[str] | None = Field(default=None) + + +class CreateTaskSetRequest(BaseModel): + spec: TaskSetSpec = Field(default_factory=TaskSetSpec) + + class ClusterInfo(TypedDict, total=False): cluster_name: str diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index eeb9911329..41aca0e7fe 100644 --- a/rock/admin/proto/response.py +++ b/rock/admin/proto/response.py @@ -3,6 +3,7 @@ from rock.actions import SandboxResponse from rock.actions.sandbox.response import State from rock.actions.sandbox.sandbox_info import SandboxInfo +from rock.admin.proto.request import TaskSetSpec class SandboxStartResponse(SandboxResponse): @@ -74,3 +75,46 @@ class SandboxListResponse(SandboxResponse): items: list[SandboxListStatusResponse] = [] total: int = 0 has_more: bool = False + + +class TaskSetMetadata(BaseModel): + tasksetId: str + creationTimestamp: float + + +class TaskSetStatusModel(BaseModel): + phase: str + assignedPod: str = "" + active: int = 0 + succeeded: int = 0 + failed: int = 0 + startTime: float | None = None + completionTime: float | None = None + conditions: list[dict] | None = None + + +class TaskMetadata(BaseModel): + taskId: str + tasksetId: str + creationTimestamp: float + + +class TaskStatusModel(BaseModel): + phase: str + startTime: float | None = None + completionTime: float | None = None + conditions: list[dict] | None = None + status: list[dict] | None = None + + +class TaskResponse(BaseModel): + metadata: TaskMetadata + spec: dict + status: TaskStatusModel + + +class TaskSetResponse(BaseModel): + metadata: TaskSetMetadata + spec: "TaskSetSpec" + status: TaskSetStatusModel + tasks: list[TaskResponse] | None = None diff --git a/rock/admin/scheduler/scheduler.py b/rock/admin/scheduler/scheduler.py index cf0b8122e3..39d88f3a74 100644 --- a/rock/admin/scheduler/scheduler.py +++ b/rock/admin/scheduler/scheduler.py @@ -297,3 +297,25 @@ def stop(self) -> None: def is_alive(self) -> bool: """Check if the scheduler thread is alive.""" return self._thread is not None and self._thread.is_alive() + + def get_task_registry(self) -> dict[str, "BaseTask"]: + """Return a snapshot {task.type: instance} of currently scheduled tasks. + + Keyed by ``task.type`` (e.g. ``"image_cleanup"``), not by full class + path, so admin ops API can match user-friendly task names supplied + from clients. The internal ``_tasks_by_class`` dict is keyed by class + path (used by Nacos config-driven install/uninstall); this method + builds a fresh ``{type: task}`` dict for external consumers. + + Safe to call from other threads — returns a fresh dict so concurrent + mutation by Nacos config reload (in scheduler thread) doesn't fault. + """ + if self._task_scheduler is None: + return {} + return {t.type: t for t in self._task_scheduler._tasks_by_class.values()} + + def get_alive_workers(self) -> list[str]: + """Return currently alive worker IPs (TTL-cached via WorkerIPCache).""" + if self._task_scheduler is None or self._task_scheduler._worker_cache is None: + return [] + return self._task_scheduler._worker_cache.get_alive_workers() diff --git a/rock/admin/service/__init__.py b/rock/admin/service/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/rock/admin/service/ops_service.py b/rock/admin/service/ops_service.py new file mode 100644 index 0000000000..a866cda884 --- /dev/null +++ b/rock/admin/service/ops_service.py @@ -0,0 +1,237 @@ +"""OpsService: business logic for admin ops TaskSet lifecycle.""" + +from __future__ import annotations + +import asyncio +import os +import time +import uuid +from collections.abc import Callable + +from rock.admin.core.scheduler_task_table import Phase, SchedulerTaskTable +from rock.admin.core.schema import SchedulerTaskRecord +from rock.admin.proto.request import TaskSetSpec +from rock.admin.proto.response import ( + TaskMetadata, + TaskResponse, + TaskSetMetadata, + TaskSetResponse, + TaskSetStatusModel, + TaskStatusModel, +) +from rock.admin.scheduler.task_base import BaseTask +from rock.logger import init_logger + +logger = init_logger(__name__) +audit_logger = init_logger("admin_ops_audit") + +_RATE_LIMIT_SECONDS = 60 +_WHITELIST_SUFFIXES = ("_cleanup", "_prune", "_archive") + + +class OpsService: + def __init__( + self, + task_table: SchedulerTaskTable, + task_registry: dict[str, BaseTask], + alive_workers_provider: Callable[[], list[str]], + ) -> None: + self._task_table = task_table + self._task_registry = task_registry + self._alive_workers_provider = alive_workers_provider + + async def create_taskset(self, spec: TaskSetSpec, caller: str) -> TaskSetResponse: + audit_logger.info(f"create_taskset: caller={caller}, spec={spec.model_dump()}") + + worker_ips = self._resolve_workers(spec) + allowed, rejected = self._resolve_tasks(spec.taskTypes) + + if not allowed: + return TaskSetResponse( + metadata=TaskSetMetadata(tasksetId="", creationTimestamp=time.time()), + spec=spec, + status=TaskSetStatusModel( + phase=Phase.REJECTED, + conditions=[{"type": "Rejected", "rejectedTaskTypes": rejected}], + ), + ) + + since = time.time() - _RATE_LIMIT_SECONDS + in_cooldown: list[str] = [] + for t in allowed: + if await self._task_table.has_recent_task(t.type, since): + in_cooldown.append(t.type) + + runnable = [t for t in allowed if t.type not in in_cooldown] + + if not runnable: + return TaskSetResponse( + metadata=TaskSetMetadata(tasksetId="", creationTimestamp=time.time()), + spec=spec, + status=TaskSetStatusModel( + phase=Phase.RATE_LIMITED, + conditions=[ + { + "type": "RateLimited", + "rateLimitedTaskTypes": sorted(in_cooldown), + "cooldownSeconds": _RATE_LIMIT_SECONDS, + "rejectedTaskTypes": rejected, + } + ], + ), + ) + + now = time.time() + taskset_id = uuid.uuid4().hex + pod_id = _pod_id() + + records = [ + SchedulerTaskRecord( + task_id=uuid.uuid4().hex, + taskset_id=taskset_id, + task_type=t.type, + target_workers=worker_ips, + creation_timestamp=now, + phase=Phase.PENDING, + assigned_pod=pod_id, + ) + for t in runnable + ] + await self._task_table.insert_tasks(records) + + asyncio.create_task(self._run_tasks_async(taskset_id, runnable, worker_ips, records)) + + audit_logger.info( + f"create_taskset: taskset_id={taskset_id}, caller={caller}, " + f"tasks={[t.type for t in runnable]}, workers={len(worker_ips)}, pod={pod_id}" + ) + + resp = _aggregate_taskset(taskset_id, records) + if rejected or in_cooldown: + resp.status.conditions = [ + {"type": "Partial", "rejectedTaskTypes": rejected, "rateLimitedTaskTypes": in_cooldown} + ] + return resp + + async def get_taskset(self, taskset_id: str) -> TaskSetResponse: + tasks = await self._task_table.get_tasks_by_group(taskset_id) + if not tasks: + return TaskSetResponse( + metadata=TaskSetMetadata(tasksetId=taskset_id, creationTimestamp=0), + spec=TaskSetSpec(), + status=TaskSetStatusModel(phase=Phase.NOT_FOUND), + ) + return _aggregate_taskset(taskset_id, tasks) + + def _resolve_workers(self, spec: TaskSetSpec) -> list[str]: + if spec.targetWorkers is not None: + return list(spec.targetWorkers) + try: + return list(self._alive_workers_provider()) + except Exception as e: + logger.warning(f"alive workers provider failed: {e}") + return [] + + def _resolve_tasks(self, requested: list[str] | None) -> tuple[list[BaseTask], list[str]]: + if requested is None: + return [t for name, t in self._task_registry.items() if _is_whitelisted(name)], [] + allowed: list[BaseTask] = [] + rejected: list[str] = [] + for name in requested: + if not _is_whitelisted(name): + rejected.append(name) + continue + if name not in self._task_registry: + rejected.append(name) + continue + allowed.append(self._task_registry[name]) + return allowed, rejected + + async def _run_tasks_async( + self, + taskset_id: str, + tasks: list[BaseTask], + worker_ips: list[str], + records: list[SchedulerTaskRecord], + ) -> None: + for task, record in zip(tasks, records): + tid = record.task_id + await self._task_table.update_task(tid, phase=Phase.RUNNING, start_time=time.time()) + try: + await task.run(worker_ips) + status = [{"worker": ip, "success": True} for ip in worker_ips] + await self._task_table.update_task( + tid, phase=Phase.SUCCEEDED, completion_time=time.time(), status=status + ) + except Exception as e: + logger.exception(f"taskset '{taskset_id}' task '{task.type}' failed") + status = [{"worker": ip, "success": False, "message": str(e)} for ip in worker_ips] + conditions = [{"type": "Failed", "reason": "ExecutionError", "message": str(e)[:2048]}] + await self._task_table.update_task( + tid, phase=Phase.FAILED, completion_time=time.time(), status=status, conditions=conditions + ) + + audit_logger.info(f"taskset '{taskset_id}' done") + + +def _is_whitelisted(task_type: str) -> bool: + return any(task_type.endswith(s) for s in _WHITELIST_SUFFIXES) + + +def _pod_id() -> str: + return os.environ.get("HOSTNAME") or "unknown" + + +def _aggregate_taskset(taskset_id: str, tasks: list[SchedulerTaskRecord]) -> TaskSetResponse: + task_responses = [ + TaskResponse( + metadata=TaskMetadata( + taskId=t.task_id, + tasksetId=t.taskset_id, + creationTimestamp=t.creation_timestamp, + ), + spec={"taskType": t.task_type, "targetWorkers": t.target_workers}, + status=TaskStatusModel( + phase=t.phase, + startTime=t.start_time, + completionTime=t.completion_time, + conditions=t.conditions, + status=t.status, + ), + ) + for t in tasks + ] + + succeeded = sum(1 for t in tasks if t.phase == Phase.SUCCEEDED) + failed = sum(1 for t in tasks if t.phase == Phase.FAILED) + active = len(tasks) - succeeded - failed + + if active > 0: + phase = Phase.RUNNING + elif failed > 0: + phase = Phase.FAILED + else: + phase = Phase.SUCCEEDED + + start_times = [t.start_time for t in tasks if t.start_time] + completion_times = [t.completion_time for t in tasks if t.completion_time] + + return TaskSetResponse( + metadata=TaskSetMetadata( + tasksetId=taskset_id, + creationTimestamp=min(t.creation_timestamp for t in tasks), + ), + spec=TaskSetSpec( + targetWorkers=tasks[0].target_workers if tasks else None, + ), + status=TaskSetStatusModel( + phase=phase, + assignedPod=tasks[0].assigned_pod if tasks else "", + active=active, + succeeded=succeeded, + failed=failed, + startTime=min(start_times) if start_times else None, + completionTime=max(completion_times) if completion_times and active == 0 else None, + ), + tasks=task_responses, + ) diff --git a/tests/unit/admin/core/test_scheduler_task_table.py b/tests/unit/admin/core/test_scheduler_task_table.py new file mode 100644 index 0000000000..799842d180 --- /dev/null +++ b/tests/unit/admin/core/test_scheduler_task_table.py @@ -0,0 +1,134 @@ +"""Tests for SchedulerTaskTable CRUD (single-table: scheduler_task).""" + +from __future__ import annotations + +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.admin.core.schema import SchedulerTaskRecord +from rock.admin.core.scheduler_task_table import Phase, SchedulerTaskTable + + +def _make_table_with_mock_session(): + table = SchedulerTaskTable.__new__(SchedulerTaskTable) + table._db = MagicMock() + + mock_session = AsyncMock() + cm = AsyncMock() + cm.__aenter__ = AsyncMock(return_value=mock_session) + cm.__aexit__ = AsyncMock(return_value=None) + + return table, cm, mock_session + + +@pytest.mark.asyncio +async def test_insert_tasks(): + table, cm, session = _make_table_with_mock_session() + + record = SchedulerTaskRecord( + task_id="b" * 32, + taskset_id="a" * 32, + task_type="image_cleanup", + target_workers=["10.0.0.1"], + creation_timestamp=time.time(), + phase=Phase.PENDING, + assigned_pod="pod-x", + ) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + await table.insert_tasks([record]) + + session.add.assert_called_once_with(record) + session.commit.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_tasks_by_group(): + table, cm, session = _make_table_with_mock_session() + + mock_record = SchedulerTaskRecord( + task_id="b" * 32, + taskset_id="a" * 32, + task_type="image_cleanup", + target_workers=["10.0.0.1"], + creation_timestamp=time.time(), + phase=Phase.RUNNING, + assigned_pod="pod-x", + ) + + scalars_mock = MagicMock() + scalars_mock.all.return_value = [mock_record] + result_mock = MagicMock() + result_mock.scalars.return_value = scalars_mock + session.execute = AsyncMock(return_value=result_mock) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + results = await table.get_tasks_by_group("a" * 32) + + assert len(results) == 1 + assert isinstance(results[0], SchedulerTaskRecord) + assert results[0].taskset_id == "a" * 32 + + +@pytest.mark.asyncio +async def test_update_task(): + table, cm, session = _make_table_with_mock_session() + + row = SchedulerTaskRecord( + task_id="b" * 32, + taskset_id="a" * 32, + task_type="image_cleanup", + target_workers=["10.0.0.1"], + creation_timestamp=time.time(), + phase=Phase.PENDING, + assigned_pod="pod-x", + ) + session.get = AsyncMock(return_value=row) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + ok = await table.update_task("b" * 32, phase=Phase.RUNNING, start_time=123.0) + + assert ok is True + assert row.phase == Phase.RUNNING + assert row.start_time == 123.0 + + +@pytest.mark.asyncio +async def test_update_task_not_found(): + table, cm, session = _make_table_with_mock_session() + session.get = AsyncMock(return_value=None) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + ok = await table.update_task("nonexistent", phase=Phase.RUNNING) + + assert ok is False + + +@pytest.mark.asyncio +async def test_has_recent_task_true(): + table, cm, session = _make_table_with_mock_session() + + result_mock = MagicMock() + result_mock.first.return_value = ("some_id",) + session.execute = AsyncMock(return_value=result_mock) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + found = await table.has_recent_task("image_cleanup", time.time() - 60) + + assert found is True + + +@pytest.mark.asyncio +async def test_has_recent_task_false(): + table, cm, session = _make_table_with_mock_session() + + result_mock = MagicMock() + result_mock.first.return_value = None + session.execute = AsyncMock(return_value=result_mock) + + with patch("rock.admin.core.scheduler_task_table.AsyncSession", return_value=cm): + found = await table.has_recent_task("image_cleanup", time.time() - 60) + + assert found is False diff --git a/tests/unit/admin/entrypoints/test_admin_ops_api.py b/tests/unit/admin/entrypoints/test_admin_ops_api.py new file mode 100644 index 0000000000..e946a8c0d1 --- /dev/null +++ b/tests/unit/admin/entrypoints/test_admin_ops_api.py @@ -0,0 +1,256 @@ +"""Tests for admin ops API (layered: api → OpsService → SchedulerTaskTable).""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient + +from rock.admin.core.scheduler_task_table import Phase +from rock.admin.core.schema import SchedulerTaskRecord +from rock.admin.entrypoints.admin_ops_api import admin_ops_router, set_ops_service +from rock.admin.service.ops_service import OpsService + + +def _fake_task(type_: str): + t = MagicMock() + t.type = type_ + t.run = AsyncMock(return_value=None) + return t + + +class FakeTable: + """In-memory fake SchedulerTaskTable using SchedulerTaskRecord objects.""" + + def __init__(self): + self._tasks: dict[str, SchedulerTaskRecord] = {} + + async def insert_tasks(self, records: list[SchedulerTaskRecord]): + for r in records: + self._tasks[r.task_id] = r + + async def get_tasks_by_group(self, taskset_id: str) -> list[SchedulerTaskRecord]: + return [t for t in self._tasks.values() if t.taskset_id == taskset_id] + + async def update_task(self, task_id: str, **fields) -> bool: + if task_id not in self._tasks: + return False + record = self._tasks[task_id] + for k, v in fields.items(): + setattr(record, k, v) + return True + + async def has_recent_task(self, task_type: str, since_epoch: float) -> bool: + return any(t.task_type == task_type and t.creation_timestamp >= since_epoch for t in self._tasks.values()) + + +@pytest.fixture +def app_with_router(): + app = FastAPI() + app.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops") + return app + + +@pytest.fixture +def fake_table(): + return FakeTable() + + +@pytest.fixture(autouse=True) +def setup_module(fake_table): + registry = { + "image_cleanup": _fake_task("image_cleanup"), + "build_cache_cleanup": _fake_task("build_cache_cleanup"), + "ray_log_cleanup": _fake_task("ray_log_cleanup"), + } + service = OpsService( + task_table=fake_table, + task_registry=registry, + alive_workers_provider=lambda: ["10.0.0.1", "10.0.0.2"], + ) + set_ops_service(service) + yield + set_ops_service(None) + + +@pytest.fixture +async def client(app_with_router): + transport = ASGITransport(app=app_with_router) + async with AsyncClient(transport=transport, base_url="http://test") as c: + yield c + + +class TestCreateTaskSet: + @pytest.mark.asyncio + async def test_accepted_default_tasks_default_workers(self, client, fake_table): + r = await client.post("/apis/envs/sandbox/v1/ops/tasksets", json={"spec": {}}) + assert r.status_code == 200 + body = r.json() + assert body["status"] == "Success" + result = body["result"] + assert result["status"]["phase"] == Phase.RUNNING + assert result["metadata"]["tasksetId"] != "" + assert len(result["metadata"]["tasksetId"]) == 32 + assert result["status"]["active"] == 3 + assert len(result["tasks"]) == 3 + task_types = {t["spec"]["taskType"] for t in result["tasks"]} + assert task_types == {"image_cleanup", "build_cache_cleanup", "ray_log_cleanup"} + + @pytest.mark.asyncio + async def test_accepted_specific_tasks(self, client): + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.RUNNING + assert len(body["result"]["tasks"]) == 1 + assert body["result"]["tasks"][0]["spec"]["taskType"] == "image_cleanup" + + @pytest.mark.asyncio + async def test_rejected_non_whitelisted_task(self, client): + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_pull"]}}, + ) + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.REJECTED + + @pytest.mark.asyncio + async def test_rejected_unknown_whitelisted_task(self, client): + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["nonexistent_cleanup"]}}, + ) + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.REJECTED + + @pytest.mark.asyncio + async def test_rate_limited(self, client): + await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.RATE_LIMITED + + @pytest.mark.asyncio + async def test_partial_rate_limit_runs_remainder(self, client): + await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup", "build_cache_cleanup"]}}, + ) + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.RUNNING + assert len(body["result"]["tasks"]) == 1 + assert body["result"]["tasks"][0]["spec"]["taskType"] == "build_cache_cleanup" + + @pytest.mark.asyncio + async def test_tasks_persisted(self, client, fake_table): + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + taskset_id = r.json()["result"]["metadata"]["tasksetId"] + child_tasks = [t for t in fake_table._tasks.values() if t.taskset_id == taskset_id] + assert len(child_tasks) == 1 + assert child_tasks[0].task_type == "image_cleanup" + + @pytest.mark.asyncio + async def test_taskset_id_is_128_bit_uuid(self, client): + r = await client.post("/apis/envs/sandbox/v1/ops/tasksets", json={"spec": {}}) + taskset_id = r.json()["result"]["metadata"]["tasksetId"] + assert len(taskset_id) == 32 + int(taskset_id, 16) + + @pytest.mark.asyncio + async def test_child_task_has_taskset_reference(self, client): + r = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + result = r.json()["result"] + parent_id = result["metadata"]["tasksetId"] + child = result["tasks"][0] + assert child["metadata"]["tasksetId"] == parent_id + + +class TestGetTaskSet: + @pytest.mark.asyncio + async def test_get_existing_taskset_with_tasks(self, client): + post = await client.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + taskset_id = post.json()["result"]["metadata"]["tasksetId"] + + get = await client.get(f"/apis/envs/sandbox/v1/ops/tasksets/{taskset_id}") + body = get.json() + assert body["status"] == "Success" + assert body["result"]["metadata"]["tasksetId"] == taskset_id + assert body["result"]["status"]["phase"] in (Phase.RUNNING, Phase.SUCCEEDED) + assert len(body["result"]["tasks"]) == 1 + + @pytest.mark.asyncio + async def test_get_nonexistent_taskset(self, client): + r = await client.get("/apis/envs/sandbox/v1/ops/tasksets/doesnotexist") + body = r.json() + assert body["status"] == "Success" + assert body["result"]["status"]["phase"] == Phase.NOT_FOUND + + +class TestMultiPod: + @pytest.mark.asyncio + async def test_post_pod_a_get_pod_b_shares_state(self, fake_table): + app_a = FastAPI() + app_a.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops") + app_b = FastAPI() + app_b.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops") + + async with AsyncClient(transport=ASGITransport(app=app_a), base_url="http://a") as ca: + post = await ca.post( + "/apis/envs/sandbox/v1/ops/tasksets", + json={"spec": {"taskTypes": ["image_cleanup"]}}, + ) + taskset_id = post.json()["result"]["metadata"]["tasksetId"] + + async with AsyncClient(transport=ASGITransport(app=app_b), base_url="http://b") as cb: + get = await cb.get(f"/apis/envs/sandbox/v1/ops/tasksets/{taskset_id}") + body = get.json() + + assert body["status"] == "Success" + assert body["result"]["metadata"]["tasksetId"] == taskset_id + assert body["result"]["status"]["phase"] != Phase.NOT_FOUND + + +class TestMisconfiguration: + @pytest.mark.asyncio + async def test_post_returns_failed_when_service_unset(self, app_with_router): + set_ops_service(None) + transport = ASGITransport(app=app_with_router) + async with AsyncClient(transport=transport, base_url="http://test") as c: + r = await c.post("/apis/envs/sandbox/v1/ops/tasksets", json={"spec": {}}) + assert r.json()["status"] == "Failed" + + @pytest.mark.asyncio + async def test_get_returns_failed_when_service_unset(self, app_with_router): + set_ops_service(None) + transport = ASGITransport(app=app_with_router) + async with AsyncClient(transport=transport, base_url="http://test") as c: + r = await c.get("/apis/envs/sandbox/v1/ops/tasksets/x") + assert r.json()["status"] == "Failed" From a5e1a9bf6f44309f9df0eaaa7c891e02bc7c0c6f Mon Sep 17 00:00:00 2001 From: jinbai <15652831212@163.com> Date: Wed, 3 Jun 2026 15:30:51 +0800 Subject: [PATCH 172/226] fix(rocklet): set success=true and file_name in UploadResponse after successful upload --- rock/rocklet/local_api.py | 2 +- rock/rocklet/rocklet.py | 2 +- tests/unit/rocklet/test_local_sandbox_runtime.py | 8 ++++++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/rock/rocklet/local_api.py b/rock/rocklet/local_api.py index 362688a74d..fc3ee49d6d 100644 --- a/rock/rocklet/local_api.py +++ b/rock/rocklet/local_api.py @@ -107,7 +107,7 @@ async def upload( file_path.unlink() else: shutil.move(file_path, target_path) - return UploadResponse() + return UploadResponse(success=True, file_name=target_path.name) @local_router.post("/close") diff --git a/rock/rocklet/rocklet.py b/rock/rocklet/rocklet.py index 0a50d5186d..3881f68744 100644 --- a/rock/rocklet/rocklet.py +++ b/rock/rocklet/rocklet.py @@ -250,7 +250,7 @@ async def upload(self, request: UploadRequest) -> UploadResponse: else: shutil.copy(request.source_path, request.target_path) self.command_logger.info("[upload output]: upload success!") - return UploadResponse() + return UploadResponse(success=True, file_name=Path(request.target_path).name) async def close(self) -> CloseResponse: """Closes the runtime.""" diff --git a/tests/unit/rocklet/test_local_sandbox_runtime.py b/tests/unit/rocklet/test_local_sandbox_runtime.py index d554c85eaa..11f4be7d1d 100644 --- a/tests/unit/rocklet/test_local_sandbox_runtime.py +++ b/tests/unit/rocklet/test_local_sandbox_runtime.py @@ -22,7 +22,9 @@ async def test_upload_file(local_runtime: Rocklet, tmp_path: Path): file_path = tmp_path / "source.txt" file_path.write_text("test") tmp_target = tmp_path / "target.txt" - await local_runtime.upload(UploadRequest(source_path=str(file_path), target_path=str(tmp_target))) + resp = await local_runtime.upload(UploadRequest(source_path=str(file_path), target_path=str(tmp_target))) + assert resp.success is True + assert resp.file_name == "target.txt" assert ( await local_runtime.read_file(ReadFileRequest(path=str(tmp_target), sandbox_id="local-test")) ).content == "test" @@ -35,7 +37,9 @@ async def test_upload_directory(local_runtime: Rocklet, tmp_path: Path): (dir_path / "file1.txt").write_text("test1") (dir_path / "file2.txt").write_text("test2") tmp_target = tmp_path / "target_dir" - await local_runtime.upload(UploadRequest(source_path=str(dir_path), target_path=str(tmp_target))) + resp = await local_runtime.upload(UploadRequest(source_path=str(dir_path), target_path=str(tmp_target))) + assert resp.success is True + assert resp.file_name == "target_dir" sid = "local-test" assert ( await local_runtime.read_file(ReadFileRequest(path=str(tmp_target / "file1.txt"), sandbox_id=sid)) From 215b118488364e7b9da051931f13140f2236620a Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Wed, 3 Jun 2026 17:57:46 +0800 Subject: [PATCH 173/226] =?UTF-8?q?feat(sandbox):=20add=20/delete=20endpoi?= =?UTF-8?q?nt=20+=20cascade=20STOPPED=20=E2=86=92=20DELETED=20for=20--rm?= =?UTF-8?q?=20containers=20(#1038)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a DELETE /sandbox/{id} endpoint that transitions STOPPED sandboxes to DELETED state, cleaning up containers and archiving metadata. When `remove_container=True`, the stop flow cascades directly to DELETED via DeleteReason.IMMEDIATE — the container is already gone after operator.stop, so we skip the intermediate STOPPED state entirely. --- rock/actions/sandbox/response.py | 1 + rock/actions/sandbox/sandbox_info.py | 1 + rock/admin/entrypoints/sandbox_api.py | 18 ++ rock/common/constants.py | 12 + rock/deployments/docker.py | 18 ++ rock/sandbox/operator/abstract.py | 4 + rock/sandbox/operator/k8s/operator.py | 3 + rock/sandbox/operator/ray.py | 26 +++ rock/sandbox/sandbox_actor.py | 9 + rock/sandbox/sandbox_manager.py | 37 ++- rock/sandbox/sandbox_statemachine.py | 41 +++- rock/sdk/sandbox/client.py | 15 ++ tests/integration/sdk/sandbox/test_basic.py | 56 +++++ .../sandbox/test_sandbox_manager_delete.py | 211 ++++++++++++++++++ .../unit/sandbox/test_sandbox_statemachine.py | 104 +++++++++ 15 files changed, 554 insertions(+), 2 deletions(-) create mode 100644 tests/unit/sandbox/test_sandbox_manager_delete.py diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index 96cf8c172c..1e44d26017 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -16,6 +16,7 @@ class State(str, Enum): PENDING = "pending" RUNNING = "running" STOPPED = "stopped" + DELETED = "deleted" class IsAliveResponse(BaseModel): diff --git a/rock/actions/sandbox/sandbox_info.py b/rock/actions/sandbox/sandbox_info.py index 3fc96c6f31..33ee9ab1d5 100644 --- a/rock/actions/sandbox/sandbox_info.py +++ b/rock/actions/sandbox/sandbox_info.py @@ -25,6 +25,7 @@ class SandboxInfo(TypedDict, total=False): create_time: str start_time: str stop_time: str + delete_time: str extended_params: dict[str, str] diff --git a/rock/admin/entrypoints/sandbox_api.py b/rock/admin/entrypoints/sandbox_api.py index b5e757e3bc..f83fc0ec7f 100644 --- a/rock/admin/entrypoints/sandbox_api.py +++ b/rock/admin/entrypoints/sandbox_api.py @@ -278,6 +278,24 @@ async def close(sandbox_id: Annotated[NonBlankStr, Body(embed=True)]) -> RockRes return RockResponse(result=f"{sandbox_id} stopped") +@sandbox_router.post("/delete") +@handle_exceptions(error_message="delete sandbox failed") +async def delete(sandbox_id: str = Body(..., embed=True)) -> RockResponse: + """Soft-delete a stopped sandbox. + + Returns 400-equivalent (status=Failed) when the sandbox is not in ``stopped`` + state. Unknown sandbox is idempotent (Success). After this call the DB + record holds ``state='deleted'`` and the worker container has been removed + via ``operator.delete``. + + Return type is the bare ``RockResponse`` (not ``RockResponse[str]``) so the + ``handle_exceptions`` decorator can fill ``result`` with a ``SandboxResponse`` + on the failure path without tripping FastAPI's response_model validation. + """ + await sandbox_manager.delete(sandbox_id) + return RockResponse(result=f"{sandbox_id} deleted") + + @sandbox_router.post("/restart") @handle_exceptions(error_message="restart sandbox failed") async def restart(sandbox_id: str = Body(..., embed=True)) -> RockResponse[SandboxStartResponse]: diff --git a/rock/common/constants.py b/rock/common/constants.py index c7d867fb22..bf333f6dd8 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -31,3 +31,15 @@ class StopReason(str, Enum): MANUAL = "manual" EXPIRED = "expired" + + +class DeleteReason(str, Enum): + """Why a sandbox was deleted. Distinguishes operator-initiated /delete calls from + background scanner cleanups driven by ``auto_delete_seconds``. + """ + + MANUAL = "manual" + # TODO: implement background auto-delete scan driven by auto_delete_seconds + EXPIRED = "expired" + # `--rm` containers: cascade STOPPED → DELETED on stop since the container is already gone. + IMMEDIATE = "immediate" diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index adcfb598ff..9b77844eb8 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -717,6 +717,24 @@ async def restart(self): logger.info(f"Container {self._container_name} restarted successfully") + async def delete(self) -> None: + """Remove the container via ``docker rm -f``. + + Idempotent — a container that doesn't exist counts as success because + nothing remains to clean up. The actor was previously stopped (or + freshly created without start), so there is no + ``self._container_process`` / ``self._runtime`` to unwind here. + Quota / log cleanup already ran during ``_stop`` and is not repeated. + """ + container_name = self._container_name or (self._config.container_name if self._config else None) + if not container_name: + logger.warning("delete: no container_name available, skipping docker rm") + return + + executor = get_executor() + loop = asyncio.get_running_loop() + await loop.run_in_executor(executor, DockerUtil.remove_container_force, container_name) + def _get_rocklet_port_from_inspect(self) -> int | None: """Read the host-side port mapped to the rocklet (container port 22555) from docker inspect.""" try: diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index 8f6347095f..b54afd77cb 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -36,6 +36,10 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo | None: async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> bool: ... + @abstractmethod + async def delete(self, config: DeploymentConfig, host_ip: str | None = None) -> bool: + ... + def set_redis_provider(self, redis_provider: RedisProvider): self._redis_provider = redis_provider diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index aeaf3c756a..4ba327e3ea 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -130,3 +130,6 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> """ logger.info(f"[{sandbox_id}] k8s stop (reason={reason.value})") return await self._provider.stop(sandbox_id) + + async def delete(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> bool: + raise NotImplementedError("delete is not yet implemented for K8sOperator") diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index 1ce9acb826..740146b56e 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -120,6 +120,32 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> ray.kill(actor) return True + async def delete(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> bool: + async with self._ray_service.get_ray_rwlock().read_lock(): + sandbox_id = config.container_name + actor_name = self._get_actor_name(sandbox_id) + + try: + existing_actor = await self._ray_service.async_ray_get_actor(actor_name) + ray.kill(existing_actor) + except Exception: + logger.info(f"Actor {actor_name} already gone, proceeding with delete") + + if not host_ip: + logger.warning( + f"delete for {sandbox_id} called without host_ip; new actor " + f"may be scheduled on a node that does not own the container" + ) + config.cpus = 0.01 + config.memory = "128m" + sandbox_actor: SandboxActor = await self.create_actor(config, pin_to_host_ip=host_ip) + try: + await self._ray_service.async_ray_get(sandbox_actor.delete.remote()) + logger.info(f"sandbox {sandbox_id} deleted on host_ip={host_ip}") + return True + finally: + ray.kill(sandbox_actor) + async def restart(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> SandboxInfo: """Restart an existing sandbox using docker start (container is preserved). diff --git a/rock/sandbox/sandbox_actor.py b/rock/sandbox/sandbox_actor.py index b29379995f..f34aea31ed 100644 --- a/rock/sandbox/sandbox_actor.py +++ b/rock/sandbox/sandbox_actor.py @@ -171,6 +171,15 @@ async def restart(self): ) raise + async def delete(self): + container_name = self._config.container_name if self._config else None + logger.info(f"[{container_name}] start to delete") + try: + await self._deployment.delete() + logger.info(f"[{container_name}] deployment deleted") + except Exception as e: + logger.error(f"[{container_name}] Error occurred while deleting container: {e}", exc_info=True) + async def commit(self, image_tag: str, username: str, password: str) -> CommandResponse: logger.info(f"start to commit {self._config.container_name} to {image_tag}") with tempfile.TemporaryDirectory() as docker_config_dir: diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index cc78fec538..1c1ff3070a 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -25,7 +25,7 @@ from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest from rock.admin.proto.response import SandboxStartResponse, SandboxStatusResponse -from rock.common.constants import StopReason +from rock.common.constants import DeleteReason, StopReason from rock.config import RockConfig, RuntimeConfig from rock.deployments.config import DeploymentConfig, DockerDeploymentConfig from rock.logger import init_logger @@ -209,6 +209,41 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL): meta_store=self._meta_store, reason=reason, ) + # `--rm` containers are already gone after stop; cascade to DELETED + # so the metadata row doesn't linger in STOPPED. + # Redis keys are gone after archive; re-read from DB to get spec. + sm = await self._get_current_statemachine(sandbox_id) + spec = ((sm.sandbox_info or {}).get("spec") or {}) if sm else {} + if spec.get("remove_container"): + await sm.send( + "delete", + sandbox_id=sandbox_id, + operator=self._operator, + meta_store=self._meta_store, + reason=DeleteReason.IMMEDIATE, + ) + + @monitor_sandbox_operation() + async def delete(self, sandbox_id: str, reason: DeleteReason = DeleteReason.MANUAL) -> None: + sm = await self._get_current_statemachine(sandbox_id) + if sm is None: + logger.info(f"delete: sandbox {sandbox_id} not found, noop") + return + state = sm.current_state.value + if state == State.DELETED: + logger.info(f"delete: sandbox {sandbox_id} already deleted, noop") + return + if state != State.STOPPED: + raise BadRequestRockError( + f"Sandbox {sandbox_id} cannot be deleted: current state is '{state.value}', must be stopped first" + ) + await sm.send( + "delete", + sandbox_id=sandbox_id, + operator=self._operator, + meta_store=self._meta_store, + reason=reason, + ) async def get_mount(self, sandbox_id): async with self._ray_service.get_ray_rwlock().read_lock(): diff --git a/rock/sandbox/sandbox_statemachine.py b/rock/sandbox/sandbox_statemachine.py index e0285f5862..105b2423d9 100644 --- a/rock/sandbox/sandbox_statemachine.py +++ b/rock/sandbox/sandbox_statemachine.py @@ -12,7 +12,7 @@ from rock.actions.sandbox.response import State as RockState from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.admin.metrics.billing import log_billing_info -from rock.common.constants import StopReason +from rock.common.constants import DeleteReason, StopReason from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger from rock.sandbox.utils.timeout import SandboxTimeoutHelper @@ -30,11 +30,13 @@ class SandboxStateMachine(StateChart): - pending: Sandbox is being created / starting - running: Sandbox is actively running - stopped: Sandbox has been stopped + - deleted: Sandbox has been soft-deleted (DB record kept, state=deleted) Transitions: - stop: pending/running → stopped (stops operator, archives meta) - stop_noop: stopped → stopped (idempotent; logs and returns) - alive: pending → running (called from get_status on pending→running; also usable by reconciler) + - delete: stopped → deleted (operator removes docker container, soft-deletes meta) """ allow_event_without_transition = False # raise TransitionNotAllowed instead of silently ignoring invalid events @@ -44,12 +46,14 @@ class SandboxStateMachine(StateChart): pending = SMState("Pending", initial=True, value=RockState.PENDING) running = SMState("Running", value=RockState.RUNNING) stopped = SMState("Stopped", value=RockState.STOPPED) + deleted = SMState("Deleted", final=True, value=RockState.DELETED) # Transitions stop = pending.to(stopped) | running.to(stopped) stop_noop = stopped.to(stopped) alive = pending.to(running) restart = stopped.to(pending) + delete = stopped.to(deleted) def __init__(self, **kwargs): """Initialize with optional sandbox_info.""" @@ -135,6 +139,40 @@ async def on_restart(self, sandbox_id: str, operator, meta_store) -> None: await meta_store.update(sandbox_id, new_info) await meta_store.update_timeout(sandbox_id, timeout_info) + async def on_delete( + self, + sandbox_id: str, + operator, + meta_store, + reason: DeleteReason = DeleteReason.MANUAL, + ) -> None: + logger.info(f"delete sandbox {sandbox_id} (reason={reason.value})") + sandbox_info = self.sandbox_info or {} + if "sandbox_id" not in sandbox_info: + sandbox_info["sandbox_id"] = sandbox_id + + host_ip = sandbox_info.get("host_ip") + if reason == DeleteReason.IMMEDIATE: + logger.info(f"sandbox {sandbox_id}: skip operator.delete (container already removed by --rm)") + else: + spec = sandbox_info.get("spec") or {} + if spec: + try: + delete_config = DockerDeploymentConfig(**spec) + await operator.delete(delete_config, host_ip=host_ip) + except Exception as e: + logger.warning(f"operator.delete({sandbox_id}, {host_ip}) failed: {e}", exc_info=True) + else: + logger.warning( + f"sandbox {sandbox_id} has no spec snapshot; skip operator.delete, " + "rely on ContainerCleanupTask to reap docker container" + ) + + sandbox_info["state"] = RockState.DELETED + sandbox_info["delete_time"] = get_iso8601_timestamp() + await meta_store.archive(sandbox_id, sandbox_info) + self.sandbox_info = sandbox_info + @classmethod async def from_state_value(cls, state_value: str | None, sandbox_info: SandboxInfo) -> "SandboxStateMachine": """Create a state machine restored to *state_value* (from Redis/DB).""" @@ -142,6 +180,7 @@ async def from_state_value(cls, state_value: str | None, sandbox_info: SandboxIn RockState.PENDING: "pending", RockState.RUNNING: "running", RockState.STOPPED: "stopped", + RockState.DELETED: "deleted", } sm = ( cls(start_value=state_map[state_value], sandbox_info=sandbox_info) diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index fc9fd55459..259314915e 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -269,6 +269,21 @@ async def stop(self): except Exception as e: logging.warning(f"Failed to stop sandbox, IGNORE: {e}") + async def delete(self): + if not self.sandbox_id: + raise Exception("sandbox_id is not set, cannot delete") + url = f"{self._url}/delete" + headers = self._build_headers() + data = {"sandbox_id": self.sandbox_id} + response = await HttpUtils.post(url, headers, data) + logging.debug(f"Delete sandbox response: {response}") + if "Success" != response.get("status"): + result = response.get("result", None) + if result is not None: + rock_response = SandboxResponse(**result) + raise_for_code(rock_response.code, f"Failed to delete sandbox: {response}") + raise Exception(f"Failed to delete sandbox: {response}") + async def restart(self): """Restart a stopped sandbox using 'docker start' (reuses existing container). diff --git a/tests/integration/sdk/sandbox/test_basic.py b/tests/integration/sdk/sandbox/test_basic.py index fda538a9a8..a1226ccb95 100644 --- a/tests/integration/sdk/sandbox/test_basic.py +++ b/tests/integration/sdk/sandbox/test_basic.py @@ -103,6 +103,62 @@ async def test_sandbox_file_operations(admin_remote_server: RemoteServer): await sandbox.stop() +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_sandbox_restart(admin_remote_server: RemoteServer): + """Test stop → restart → verify running → stop lifecycle.""" + config = SandboxConfig( + image="python:3.11", + startup_timeout=60, + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + ) + sandbox = Sandbox(config) + try: + await sandbox.start() + await sandbox.create_session(CreateBashSessionRequest(session="default")) + + result = await sandbox.arun(cmd="echo before_restart", session="default") + assert "before_restart" in result.output + + await sandbox.stop() + status = await sandbox.get_status(include_all_states=True) + assert status.state == "stopped" + + await sandbox.restart() + status = await sandbox.get_status(include_all_states=True) + assert status.state in ("pending", "running") + + await sandbox.create_session(CreateBashSessionRequest(session="default")) + result = await sandbox.arun(cmd="echo after_restart", session="default") + assert "after_restart" in result.output + finally: + await sandbox.stop() + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_sandbox_delete(admin_remote_server: RemoteServer): + """Test stop → delete lifecycle. Deleted sandbox should not be found.""" + config = SandboxConfig( + image="python:3.11", + startup_timeout=60, + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + ) + sandbox = Sandbox(config) + await sandbox.start() + + await sandbox.stop() + status = await sandbox.get_status(include_all_states=True) + assert status.state == "stopped" + + await sandbox.delete() + + with pytest.raises(Exception): + await sandbox.get_status(include_all_states=True) + + @pytest.mark.need_admin @SKIP_IF_NO_DOCKER @pytest.mark.asyncio diff --git a/tests/unit/sandbox/test_sandbox_manager_delete.py b/tests/unit/sandbox/test_sandbox_manager_delete.py new file mode 100644 index 0000000000..c411196baa --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_manager_delete.py @@ -0,0 +1,211 @@ +"""Unit tests for SandboxManager.delete. + +Avoids ray / docker dependencies by patching out the BaseManager scheduler +setup and stubbing the meta_store / operator. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.actions.sandbox.response import State +from rock.common.constants import DeleteReason +from rock.config import RockConfig, SandboxConfig +from rock.sandbox.sandbox_manager import SandboxManager +from rock.sdk.common.exceptions import BadRequestRockError + + +@pytest.fixture +def rock_config_min(): + cfg = RockConfig() + cfg.sandbox_config = SandboxConfig() + return cfg + + +@pytest.fixture +def manager(rock_config_min): + operator = AsyncMock() + meta_store = AsyncMock() + meta_store.get = AsyncMock(return_value=None) + # Patch BaseManager scheduler setup so tests don't spawn APScheduler. + with patch("rock.sandbox.base_manager.BaseManager._setup_scheduler"): + m = SandboxManager( + rock_config=rock_config_min, + meta_store=meta_store, + ray_namespace="test", + ray_service=MagicMock(), + enable_runtime_auto_clear=False, + operator=operator, + ) + return m + + +class TestDelete: + @pytest.mark.asyncio + async def test_delete_unknown_sandbox_is_noop(self, manager): + manager._meta_store.get = AsyncMock(return_value=None) + await manager.delete("sb-unknown") + manager._meta_store.archive.assert_not_called() + manager._operator.delete.assert_not_called() + + @pytest.mark.asyncio + async def test_delete_from_pending_raises_400(self, manager): + manager._meta_store.get = AsyncMock( + return_value={"sandbox_id": "sb-1", "state": State.PENDING, "host_ip": "1.2.3.4"} + ) + with pytest.raises(BadRequestRockError): + await manager.delete("sb-1") + manager._operator.delete.assert_not_called() + manager._meta_store.archive.assert_not_called() + + @pytest.mark.asyncio + async def test_delete_from_running_raises_400(self, manager): + manager._meta_store.get = AsyncMock( + return_value={"sandbox_id": "sb-1", "state": State.RUNNING, "host_ip": "1.2.3.4"} + ) + with pytest.raises(BadRequestRockError): + await manager.delete("sb-1") + manager._operator.delete.assert_not_called() + + @pytest.mark.asyncio + async def test_delete_from_stopped_archives_with_deleted_state(self, manager): + manager._meta_store.get = AsyncMock( + return_value={ + "sandbox_id": "sb-1", + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "spec": {"container_name": "sb-1", "image": "python:3.11", "memory": "2g", "cpus": 1}, + } + ) + await manager.delete("sb-1") + manager._operator.delete.assert_awaited_once() + args, kwargs = manager._operator.delete.call_args + assert args[0].container_name == "sb-1" + assert kwargs.get("host_ip") == "1.2.3.4" + manager._meta_store.archive.assert_awaited_once() + info = manager._meta_store.archive.call_args[0][1] + assert info["state"] == State.DELETED + assert info["delete_time"] + + @pytest.mark.asyncio + async def test_delete_already_deleted_is_noop(self, manager): + manager._meta_store.get = AsyncMock( + return_value={"sandbox_id": "sb-1", "state": State.DELETED, "host_ip": "1.2.3.4"} + ) + await manager.delete("sb-1") + manager._operator.delete.assert_not_called() + manager._meta_store.archive.assert_not_called() + + @pytest.mark.asyncio + async def test_operator_delete_failure_still_archives(self, manager): + manager._meta_store.get = AsyncMock( + return_value={ + "sandbox_id": "sb-1", + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "spec": {"container_name": "sb-1", "image": "python:3.11", "memory": "2g", "cpus": 1}, + } + ) + manager._operator.delete = AsyncMock(side_effect=RuntimeError("worker unreachable")) + await manager.delete("sb-1") + manager._meta_store.archive.assert_awaited_once() + info = manager._meta_store.archive.call_args[0][1] + assert info["state"] == State.DELETED + + @pytest.mark.asyncio + async def test_delete_propagates_reason(self, manager): + manager._meta_store.get = AsyncMock( + return_value={ + "sandbox_id": "sb-1", + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "spec": {"container_name": "sb-1", "image": "python:3.11", "memory": "2g", "cpus": 1}, + } + ) + await manager.delete("sb-1", reason=DeleteReason.EXPIRED) + # No public assertion target — reason is logged. Just ensure it doesn't raise + # and archive happened. + manager._meta_store.archive.assert_awaited_once() + + +class TestStopCascadeDelete: + """`docker run --rm` sandboxes collapse STOPPED → DELETED in one stop call, + so users don't observe a STOPPED row that auto-delete would reap later.""" + + @pytest.mark.asyncio + async def test_stop_running_with_remove_container_cascades_to_deleted(self, manager): + redis_info = { + "sandbox_id": "sb-1", + "state": State.RUNNING, + "host_ip": "1.2.3.4", + "start_time": "2026-05-28T00:00:00+00:00", + } + db_info = { + **redis_info, + "state": State.STOPPED, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "remove_container": True, + }, + } + # Call order: (1) @monitor decorator reads user_info, (2) _get_current_statemachine + # reads Redis (no spec), (3) cascade reads DB fallback (has spec). + manager._meta_store.get = AsyncMock(side_effect=[redis_info, redis_info, db_info]) + await manager.stop("sb-1") + manager._operator.stop.assert_awaited_once() + manager._operator.delete.assert_not_called() + # on_stop archives with STOPPED, on_delete archives with DELETED (but + # skips operator.delete because IMMEDIATE means --rm already removed it). + assert manager._meta_store.archive.await_count == 2 + last_info = manager._meta_store.archive.call_args[0][1] + assert last_info["state"] == State.DELETED + assert last_info["delete_time"] + + @pytest.mark.asyncio + async def test_stop_running_without_remove_container_stays_stopped(self, manager): + redis_info = { + "sandbox_id": "sb-1", + "state": State.RUNNING, + "host_ip": "1.2.3.4", + "start_time": "2026-05-28T00:00:00+00:00", + } + db_info = { + **redis_info, + "state": State.STOPPED, + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "remove_container": False, + }, + } + manager._meta_store.get = AsyncMock(side_effect=[redis_info, redis_info, db_info]) + await manager.stop("sb-1") + manager._operator.stop.assert_awaited_once() + manager._operator.delete.assert_not_called() + manager._meta_store.archive.assert_awaited_once() + info = manager._meta_store.archive.call_args[0][1] + assert info["state"] == State.STOPPED + + @pytest.mark.asyncio + async def test_stop_noop_on_already_stopped_does_not_cascade(self, manager): + # A redundant stop on an already-STOPPED sandbox must stay idempotent + # even if remove_container=True — auto-delete still owns this row's + # eventual STOPPED → DELETED transition. + manager._meta_store.get = AsyncMock( + return_value={ + "sandbox_id": "sb-1", + "state": State.STOPPED, + "host_ip": "1.2.3.4", + "spec": {"container_name": "sb-1", "remove_container": True}, + } + ) + await manager.stop("sb-1") + manager._operator.stop.assert_not_called() + manager._operator.delete.assert_not_called() diff --git a/tests/unit/sandbox/test_sandbox_statemachine.py b/tests/unit/sandbox/test_sandbox_statemachine.py index fc46f1f106..54f93b4952 100644 --- a/tests/unit/sandbox/test_sandbox_statemachine.py +++ b/tests/unit/sandbox/test_sandbox_statemachine.py @@ -307,3 +307,107 @@ async def test_writes_timeout_built_from_spec(self, mock_meta_store): assert sandbox_id == "sb-1" # SandboxTimeoutHelper.make_timeout_info stores auto_clear_time as the env-var key assert any("30" == str(v) for v in timeout_info.values()) + + +# --------------------------------------------------------------------------- +# delete transitions +# --------------------------------------------------------------------------- + + +class TestDeleteTransitions: + def _kwargs(self, operator=None, meta_store=None): + return dict( + sandbox_id="sb", + operator=operator or AsyncMock(), + meta_store=meta_store or AsyncMock(), + ) + + @pytest.mark.asyncio + async def test_delete_from_stopped_transitions_to_deleted(self): + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info={"host_ip": "1.2.3.4"}) + await sm.send("delete", **self._kwargs()) + assert sm.deleted.is_active + + @pytest.mark.asyncio + async def test_delete_from_pending_raises(self): + sm = await SandboxStateMachine.from_state_value(State.PENDING, sandbox_info={}) + with pytest.raises(TransitionNotAllowed): + await sm.send("delete", **self._kwargs()) + + @pytest.mark.asyncio + async def test_delete_from_running_raises(self): + sm = await SandboxStateMachine.from_state_value(State.RUNNING, sandbox_info={}) + with pytest.raises(TransitionNotAllowed): + await sm.send("delete", **self._kwargs()) + + @pytest.mark.asyncio + async def test_deleted_is_final_no_transitions_allowed(self): + sm = await SandboxStateMachine.from_state_value(State.DELETED, sandbox_info={}) + for event in ("stop", "stop_noop", "alive", "restart", "delete"): + with pytest.raises(TransitionNotAllowed): + await sm.send(event, **self._kwargs(), sandbox_info={}) + + +_VALID_DELETE_INFO = { + "host_ip": "10.0.0.1", + "spec": { + "container_name": "sb-1", + "image": "python:3.11", + "memory": "2g", + "cpus": 1, + "auto_clear_time_minutes": 30, + }, +} + + +class TestOnDelete: + @pytest.fixture + def mock_meta_store(self): + return AsyncMock() + + @pytest.mark.asyncio + async def test_calls_operator_delete_with_config_and_host_ip(self, mock_meta_store): + op = AsyncMock() + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info=dict(_VALID_DELETE_INFO)) + await sm.send("delete", sandbox_id="sb-1", operator=op, meta_store=mock_meta_store) + op.delete.assert_awaited_once() + args, kwargs = op.delete.call_args + config = args[0] + assert config.container_name == "sb-1" + assert kwargs.get("host_ip") == "10.0.0.1" + + @pytest.mark.asyncio + async def test_archives_with_state_deleted_and_delete_time(self, mock_meta_store): + op = AsyncMock() + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info=dict(_VALID_DELETE_INFO)) + await sm.send("delete", sandbox_id="sb-1", operator=op, meta_store=mock_meta_store) + archived_info = mock_meta_store.archive.call_args[0][1] + assert archived_info["state"] == State.DELETED + assert archived_info["delete_time"] + + @pytest.mark.asyncio + async def test_operator_delete_failure_still_archives(self, mock_meta_store): + op = AsyncMock() + op.delete = AsyncMock(side_effect=RuntimeError("worker unreachable")) + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info=dict(_VALID_DELETE_INFO)) + await sm.send("delete", sandbox_id="sb-1", operator=op, meta_store=mock_meta_store) + mock_meta_store.archive.assert_awaited_once() + archived_info = mock_meta_store.archive.call_args[0][1] + assert archived_info["state"] == State.DELETED + + @pytest.mark.asyncio + async def test_missing_spec_skips_operator_but_still_archives(self, mock_meta_store): + """No ``spec`` snapshot → cannot rebuild config → skip operator.delete, + still soft-delete the record (ContainerCleanupTask becomes fallback).""" + op = AsyncMock() + sm = await SandboxStateMachine.from_state_value(State.STOPPED, sandbox_info={"host_ip": "10.0.0.1"}) + await sm.send("delete", sandbox_id="sb-1", operator=op, meta_store=mock_meta_store) + op.delete.assert_not_called() + mock_meta_store.archive.assert_awaited_once() + archived_info = mock_meta_store.archive.call_args[0][1] + assert archived_info["state"] == State.DELETED + + @pytest.mark.asyncio + async def test_restores_deleted_from_state_value(self): + sm = await SandboxStateMachine.from_state_value(State.DELETED, sandbox_info={}) + assert sm.deleted.is_active From c1293136ad997d8656e143e9822a4e0a1239557a Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Thu, 4 Jun 2026 18:36:18 +0800 Subject: [PATCH 174/226] fix(sandbox): handle actor not found in RayOperator.get_status() (#1062) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(sandbox): handle actor not found in RayOperator.get_status() After sandbox.stop() kills the Ray actor, subsequent get_status() calls crash with "Failed to look up actor" because the non-rocklet path in RayOperator.get_status() has no exception handling for missing actors. Wrap the async_ray_get_actor() call in a try-except and return None when the actor is not found. This is consistent with the rocklet path behavior and allows SandboxManager.get_status() to fall back to the state machine's stored sandbox_info (which correctly reflects the STOPPED state). Fixes test_sandbox_restart and test_sandbox_delete CI failures. * fix(sandbox): restore port_mapping from persisted status file after restart After stop(), the Ray actor is killed. Restart creates a new actor whose _service_status is empty — get_status() returns port_mapping={}, which causes on_alive() to overwrite meta_store with empty port_mapping, breaking create_session() with KeyError on Port.PROXY. Fix: in DockerDeployment.restart(), read the PersistedServiceStatus JSON file (written during start()) to recover all port mappings into the new actor's _service_status. This is the same file that the use_rocklet path reads via rocklet — hence rocklet environments were never affected. * fix(test): adapt restart/delete integration tests for --rm cascade 215b118 introduced STOPPED→DELETED cascade when remove_container=True (the default). The new tests didn't account for this: - Set auto_delete_seconds=300 so remove_container=False, keeping the sandbox in STOPPED state for restart/delete assertions. - Use include_all_states=False in test_sandbox_delete post-delete assertion (correctly raises for deleted sandboxes). * feat(sdk): add polling loop to Sandbox.restart() Match start() behavior: poll get_status() until is_alive before returning, so the caller can immediately use the sandbox after restart() resolves. --- rock/deployments/docker.py | 16 +++++++++++++--- rock/sandbox/operator/ray.py | 6 +++++- rock/sdk/sandbox/client.py | 14 ++++++++++++++ tests/integration/sdk/sandbox/test_basic.py | 6 ++++-- 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/rock/deployments/docker.py b/rock/deployments/docker.py index 9b77844eb8..f10cabf971 100644 --- a/rock/deployments/docker.py +++ b/rock/deployments/docker.py @@ -1,5 +1,6 @@ import asyncio import datetime +import json import os import random import re @@ -693,10 +694,19 @@ async def restart(self): # branch and never call docker kill / cleanup. self._container_process = await loop.run_in_executor(executor, self._docker_start) - # Recover the rocklet port from the container's port bindings if not set in config. - # When a new actor is created for restart, config.port may be None. + # Recover port mappings from the persisted service status file. + # When a new actor is created for restart, _service_status is empty, + # but the file written by the original actor during start() is still on disk. + self._service_status.set_sandbox_id(self._container_name) + status_path = PersistedServiceStatus.gen_service_status_path(self._container_name) + if os.path.exists(status_path): + with open(status_path) as f: + data = json.load(f) + for port_value, mapping in data.get("port_mapping", {}).items(): + self._service_status.add_port_mapping(int(port_value), mapping) + if self._config.port is None: - self._config.port = await loop.run_in_executor(executor, self._get_rocklet_port_from_inspect) + self._config.port = self._service_status.port_mapping.get(Port.PROXY) if self._config.port is None: raise Exception(f"Cannot determine rocklet port for container {self._container_name}") diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index 740146b56e..e8cd4ae36b 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -94,7 +94,11 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo | None: sandbox_info.update(remote_status.to_dict()) return sandbox_info async with self._ray_service.get_ray_rwlock().read_lock(): - actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id)) + try: + actor: SandboxActor = await self._ray_service.async_ray_get_actor(self._get_actor_name(sandbox_id)) + except (ValueError, Exception): + logger.debug(f"Actor for sandbox {sandbox_id} not found, returning None") + return None 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"] = {name: phase.to_dict() for name, phase in remote_status.phases.items()} diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 259314915e..23290da9cd 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -304,6 +304,20 @@ async def restart(self): raise_for_code(rock_response.code, f"Failed to restart sandbox: {response}") raise Exception(f"Failed to restart sandbox: {response}") + start_time = time.time() + while time.time() - start_time < self.config.startup_timeout: + sandbox_info = await self.get_status(include_all_states=True) + logging.debug(f"Restart get status response: {sandbox_info}") + if sandbox_info.is_alive: + return + error_msg = await self._parse_error_message_from_status(sandbox_info.status) + if error_msg: + raise InternalServerRockError(f"Failed to restart sandbox because {error_msg}, sandbox: {str(self)}") + await asyncio.sleep(3) + raise InternalServerRockError( + f"Failed to restart sandbox within {self.config.startup_timeout}s, sandbox: {str(self)}" + ) + async def commit(self, image_tag: str, username: str, password: str): if not self.sandbox_id: return diff --git a/tests/integration/sdk/sandbox/test_basic.py b/tests/integration/sdk/sandbox/test_basic.py index a1226ccb95..c67808f340 100644 --- a/tests/integration/sdk/sandbox/test_basic.py +++ b/tests/integration/sdk/sandbox/test_basic.py @@ -112,6 +112,7 @@ async def test_sandbox_restart(admin_remote_server: RemoteServer): image="python:3.11", startup_timeout=60, base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + auto_delete_seconds=300, ) sandbox = Sandbox(config) try: @@ -127,7 +128,7 @@ async def test_sandbox_restart(admin_remote_server: RemoteServer): await sandbox.restart() status = await sandbox.get_status(include_all_states=True) - assert status.state in ("pending", "running") + assert status.state == "running" await sandbox.create_session(CreateBashSessionRequest(session="default")) result = await sandbox.arun(cmd="echo after_restart", session="default") @@ -145,6 +146,7 @@ async def test_sandbox_delete(admin_remote_server: RemoteServer): image="python:3.11", startup_timeout=60, base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + auto_delete_seconds=300, ) sandbox = Sandbox(config) await sandbox.start() @@ -156,7 +158,7 @@ async def test_sandbox_delete(admin_remote_server: RemoteServer): await sandbox.delete() with pytest.raises(Exception): - await sandbox.get_status(include_all_states=True) + await sandbox.get_status(include_all_states=False) @pytest.mark.need_admin From 7291a09151102b82ab179507ffd5d634f85a9808 Mon Sep 17 00:00:00 2001 From: "Qianyang(Ji Kai)" <111677149+jake11-oho@users.noreply.github.com> Date: Fri, 5 Jun 2026 01:08:11 +0800 Subject: [PATCH 175/226] fix(rocklet): pass sandbox_id in body to /execute and /read_file (#1065) (#1066) * fix(rocklet): pass sandbox_id in body to /execute and /read_file (#1065) PR #985 added NonBlankStr to SandboxCommand.sandbox_id / SandboxReadFileRequest.sandbox_id, but two internal call sites still passed sandbox_id only in the HTTP header (rocklet) or omitted it entirely (Python construction). After the upgrade rocklet rejects the request with 422 and Pydantic raises ValidationError at Command(...) construction: - rock/sandbox/operator/ray.py:170-186 (get_remote_status): rocklet /execute and /read_file payloads carried only headers["sandbox_id"]; the rocklet deserialises into Sandbox* models from the body and now rejects the missing field with 422. Surfaces as "Failed to get status: ... '422 Unprocessable Entity' for url 'http://:22555/execute'" whenever Nacos GET_STATUS_SWITCH is on. - rock/sandbox/remote_sandbox.py:252-257 (check_pid_exists): builds a SandboxCommand without sandbox_id, raising pydantic.ValidationError at construction time and breaking task_base.cleanup_on_worker for non-idempotent scheduler tasks (docuum, ImageCleanup, ...). Both sites fixed to supply sandbox_id in the body. check_pid_exists uses "scheduler-task" to match the placeholder convention PR #985 already established in task_base.py. Audited the rest of the repo (HTTP callers, model construction, scheduler tasks, SDK, tests) - no other regressions. Co-Authored-By: Claude Opus 4.7 * refactor(remote_sandbox): make check_pid_exists take sandbox_id from caller Previous commit hard-coded sandbox_id="scheduler-task" inside RemoteSandboxRuntime.check_pid_exists. That leaks a scheduler-layer placeholder into a generic remote-runtime method and breaks the convention used elsewhere in task_base.py, where every Command/Read/Write request spells out sandbox_id="scheduler-task" at the call site. Move the value to the caller: check_pid_exists now takes sandbox_id as a required positional arg, and the two task_base.py call sites pass sandbox_id="scheduler-task" explicitly -- matching the surrounding code. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 --- rock/admin/scheduler/task_base.py | 4 +- rock/sandbox/operator/ray.py | 4 +- rock/sandbox/remote_sandbox.py | 15 ++++- .../test_ray_operator_get_remote_status.py | 57 +++++++++++++++++++ .../test_remote_sandbox_check_pid_exists.py | 26 +++++++++ 5 files changed, 99 insertions(+), 7 deletions(-) create mode 100644 tests/unit/sandbox/operator/test_ray_operator_get_remote_status.py create mode 100644 tests/unit/sandbox/test_remote_sandbox_check_pid_exists.py diff --git a/rock/admin/scheduler/task_base.py b/rock/admin/scheduler/task_base.py index 77a5d92724..13b310598d 100644 --- a/rock/admin/scheduler/task_base.py +++ b/rock/admin/scheduler/task_base.py @@ -192,7 +192,7 @@ async def cleanup_on_worker(self, ip: str) -> None: status = await self.get_task_status(runtime) if status is None or not status.pid: return - if await runtime.check_pid_exists(status.pid): + if await runtime.check_pid_exists(status.pid, sandbox_id="scheduler-task"): kill_cmd = f"pkill -9 -P {status.pid}; kill -9 {status.pid}" await runtime.execute(Command(command=kill_cmd, shell=True, sandbox_id="scheduler-task")) logger.info(f"[{self.type}] killed pid {status.pid} on worker[{ip}]") @@ -230,7 +230,7 @@ async def should_run(self, runtime: RemoteSandboxRuntime) -> bool: # Check if process is still running if status.pid and status.status == TaskStatusEnum.RUNNING: - pid_exists = await runtime.check_pid_exists(status.pid) + pid_exists = await runtime.check_pid_exists(status.pid, sandbox_id="scheduler-task") if pid_exists: return False # Process still running, skip if status.pid is None and status.status == TaskStatusEnum.FAILED: diff --git a/rock/sandbox/operator/ray.py b/rock/sandbox/operator/ray.py index e8cd4ae36b..0f03f31d72 100644 --- a/rock/sandbox/operator/ray.py +++ b/rock/sandbox/operator/ray.py @@ -200,7 +200,7 @@ async def get_remote_status(self, sandbox_id: str, host_ip: str) -> ServiceStatu find_file_rsp = await HttpUtils.post( url=execute_url, headers=headers, - data={"command": ["ls", service_status_path]}, + data={"command": ["ls", service_status_path], "sandbox_id": sandbox_id}, read_timeout=60, ) @@ -211,7 +211,7 @@ async def get_remote_status(self, sandbox_id: str, host_ip: str) -> ServiceStatu response: dict = await HttpUtils.post( url=read_file_url, headers=headers, - data={"path": service_status_path}, + data={"path": service_status_path, "sandbox_id": sandbox_id}, read_timeout=60, ) if response.get("content"): diff --git a/rock/sandbox/remote_sandbox.py b/rock/sandbox/remote_sandbox.py index 8ad1fa395f..f073528221 100644 --- a/rock/sandbox/remote_sandbox.py +++ b/rock/sandbox/remote_sandbox.py @@ -249,10 +249,19 @@ def _get_statistics(self): msg += traceback.format_exc() return {} - async def check_pid_exists(self, pid: int) -> bool: - """Check if a process exists on the remote host.""" + async def check_pid_exists(self, pid: int, sandbox_id: str) -> bool: + """Check if a process exists on the remote host. + + ``sandbox_id`` satisfies the rocklet's NonBlankStr contract on + ``SandboxCommand`` (PR #985) and shows up in the rocklet access log + for tracing -- pass the caller's own context value. + """ result = await self.execute( - Command(command=f"kill -0 {pid} 2>/dev/null && echo 'exists' || echo 'not_exists'", shell=True) + Command( + command=f"kill -0 {pid} 2>/dev/null && echo 'exists' || echo 'not_exists'", + shell=True, + sandbox_id=sandbox_id, + ) ) return result.stdout.strip() == "exists" diff --git a/tests/unit/sandbox/operator/test_ray_operator_get_remote_status.py b/tests/unit/sandbox/operator/test_ray_operator_get_remote_status.py new file mode 100644 index 0000000000..2243484aec --- /dev/null +++ b/tests/unit/sandbox/operator/test_ray_operator_get_remote_status.py @@ -0,0 +1,57 @@ +"""Regression tests for RayOperator.get_remote_status request shape. + +PR #985 added ``NonBlankStr sandbox_id`` to ``SandboxCommand`` and +``SandboxReadFileRequest`` in ``rock/admin/proto/request.py``. The rocklet's +``/execute`` and ``/read_file`` endpoints deserialise into those models, so +the request body must carry ``sandbox_id`` -- passing it only in the HTTP +header (the pre-985 behaviour) now triggers 422 Unprocessable Entity from +the rocklet, which bubbles up through ``get_status_v2`` as +``Failed to get status``. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rock.sandbox.operator.ray import RayOperator + + +@pytest.fixture +def operator() -> RayOperator: + return RayOperator(ray_service=MagicMock(), runtime_config=MagicMock()) + + +@pytest.mark.asyncio +async def test_get_remote_status_passes_sandbox_id_to_execute_body(operator: RayOperator): + sandbox_id = "sb-abc" + host_ip = "10.0.0.1" + with patch("rock.sandbox.operator.ray.HttpUtils.post", new=AsyncMock()) as mock_post: + # exit_code == 2 short-circuits before /read_file is called. + mock_post.return_value = {"exit_code": 2} + await operator.get_remote_status(sandbox_id, host_ip) + + assert mock_post.call_count == 1 + kwargs = mock_post.call_args.kwargs + assert "/execute" in kwargs["url"] + assert kwargs["data"].get("sandbox_id") == sandbox_id, ( + "rocklet /execute requires sandbox_id in body (NonBlankStr); header-only sandbox_id triggers 422" + ) + + +@pytest.mark.asyncio +async def test_get_remote_status_passes_sandbox_id_to_read_file_body(operator: RayOperator): + sandbox_id = "sb-abc" + host_ip = "10.0.0.1" + with patch("rock.sandbox.operator.ray.HttpUtils.post", new=AsyncMock()) as mock_post: + mock_post.side_effect = [ + {"exit_code": 0}, # ls succeeded -> proceed to read_file + {"content": ""}, # empty content path + ] + await operator.get_remote_status(sandbox_id, host_ip) + + assert mock_post.call_count == 2 + read_file_kwargs = mock_post.call_args_list[1].kwargs + assert "/read_file" in read_file_kwargs["url"] + assert read_file_kwargs["data"].get("sandbox_id") == sandbox_id, ( + "rocklet /read_file requires sandbox_id in body (NonBlankStr); header-only sandbox_id triggers 422" + ) diff --git a/tests/unit/sandbox/test_remote_sandbox_check_pid_exists.py b/tests/unit/sandbox/test_remote_sandbox_check_pid_exists.py new file mode 100644 index 0000000000..69b3b9f5b2 --- /dev/null +++ b/tests/unit/sandbox/test_remote_sandbox_check_pid_exists.py @@ -0,0 +1,26 @@ +"""Regression test for RemoteSandboxRuntime.check_pid_exists. + +PR #985 added ``NonBlankStr sandbox_id`` to ``SandboxCommand``, so the +caller-supplied ``sandbox_id`` must reach the wire request -- otherwise +construction raises ``pydantic.ValidationError`` (breaking the scheduler's +non-idempotent task cleanup path ``task_base.cleanup_on_worker`` -> +``runtime.check_pid_exists``). +""" + +from unittest.mock import AsyncMock + +import pytest + +from rock.actions import CommandResponse +from rock.sandbox.remote_sandbox import RemoteSandboxRuntime + + +@pytest.mark.asyncio +async def test_check_pid_exists_forwards_sandbox_id_to_command(): + runtime = RemoteSandboxRuntime(host="http://127.0.0.1", port=22555) + runtime.execute = AsyncMock(return_value=CommandResponse(exit_code=0, stdout="exists\n", stderr="")) + + assert await runtime.check_pid_exists(1234, sandbox_id="scheduler-task") is True + + cmd_arg = runtime.execute.call_args.args[0] + assert cmd_arg.sandbox_id == "scheduler-task" From c622091fe5e7acc1077cc4f832da3aec077d447b Mon Sep 17 00:00:00 2001 From: Jiachen Zhang Date: Fri, 5 Jun 2026 17:24:20 +0800 Subject: [PATCH 176/226] fix(sandbox): handle K8s CRD not found in K8sOperator.get_status() (#1068) When a sandbox exists in Redis/DB but its K8s BatchSandbox CRD has been deleted, get_status() leaked the raw ApiException(404) to the client. Catch exceptions in K8sOperator.get_status() and return None, consistent with RayOperator behavior. This allows SandboxManager to fall back to the stored state from meta_store when include_all_states=True. --- rock/sandbox/operator/k8s/operator.py | 15 +++- .../sandbox/operator/test_k8s_operator.py | 87 +++++++++++++++++-- 2 files changed, 92 insertions(+), 10 deletions(-) diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 4ba327e3ea..d28ea74591 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -103,10 +103,19 @@ async def get_status(self, sandbox_id: str) -> SandboxInfo | None: sandbox_id: Sandbox identifier Returns: - SandboxInfo with current status and user info + SandboxInfo with current status and user info, or None if K8s resource not found """ - # Get sandbox info from provider (includes is_alive check) - sandbox_info = await self._provider.get_status(sandbox_id) + try: + sandbox_info = await self._provider.get_status(sandbox_id) + except Exception as e: + if hasattr(e, "status") and e.status == 404: + logger.debug(f"K8s resource for sandbox {sandbox_id} not found, returning None") + return None + if "is being deleted" in str(e): + logger.debug(f"K8s resource for sandbox {sandbox_id} is being deleted, returning None") + return None + logger.warning(f"Failed to get status from K8s for sandbox {sandbox_id}: {e}") + return None # Get user info from redis if available if self._redis_provider: diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 67da376313..d0dcb784fa 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -113,22 +113,31 @@ async def test_get_status_not_alive(self, k8s_operator, mock_provider): @pytest.mark.asyncio async def test_get_status_not_found(self, k8s_operator, mock_provider): - """Test status retrieval when sandbox not found in cache.""" + """Test status retrieval returns None when sandbox not found in K8s.""" mock_provider.get_status = AsyncMock(side_effect=Exception("Sandbox test-sandbox not found")) - with pytest.raises(Exception, match="not found"): - await k8s_operator.get_status("test-sandbox") + result = await k8s_operator.get_status("test-sandbox") + assert result is None + + @pytest.mark.asyncio + async def test_get_status_k8s_api_404(self, k8s_operator, mock_provider): + """Test status retrieval returns None when K8s API returns 404.""" + from kubernetes.client.exceptions import ApiException + + mock_provider.get_status = AsyncMock(side_effect=ApiException(status=404, reason="Not Found")) + + result = await k8s_operator.get_status("test-sandbox") + assert result is None @pytest.mark.asyncio async def test_get_status_missing_ports_annotation(self, k8s_operator, mock_provider): - """Test that missing ports annotation raises error.""" - # Mock get_status to raise ValueError for missing ports + """Test that missing ports annotation returns None.""" mock_provider.get_status = AsyncMock( side_effect=ValueError("Sandbox 'test-sandbox' is missing required 'rock.sandbox/ports' annotation") ) - with pytest.raises(Exception, match="missing required.*annotation"): - await k8s_operator.get_status("test-sandbox") + result = await k8s_operator.get_status("test-sandbox") + assert result is None @pytest.mark.asyncio async def test_stop_success(self, k8s_operator, mock_provider): @@ -197,6 +206,70 @@ async def test_get_status_not_found_in_redis(self, k8s_operator, mock_provider, assert result is None +class TestK8sGetStatusWithManager: + """End-to-end test: K8sOperator + SandboxManager.get_status(include_all_states=True). + + Verifies that when a sandbox exists in Redis/DB but the K8s CRD is deleted, + get_status(include_all_states=True) returns the stored stopped state instead of crashing. + """ + + @pytest.fixture + def mock_meta_store(self): + from unittest.mock import AsyncMock + + store = AsyncMock() + store.get = AsyncMock(return_value=None) + store.create = AsyncMock() + store.update = AsyncMock() + store.archive = AsyncMock() + store.get_timeout = AsyncMock(return_value=None) + store.update_timeout = AsyncMock() + return store + + @pytest.fixture + async def mgr_with_k8s_operator(self, k8s_operator, mock_meta_store): + from unittest.mock import AsyncMock, MagicMock + + from rock.sandbox.sandbox_manager import SandboxManager + from rock.sandbox.sandbox_statemachine import SandboxStateMachine + + m = MagicMock(spec=SandboxManager) + m._meta_store = mock_meta_store + m._operator = k8s_operator + + async def get_current_statemachine(sandbox_id: str): + info = await mock_meta_store.get(sandbox_id, check_db=True) + if info is None: + return None + return await SandboxStateMachine.from_state_value(info.get("state"), sandbox_info=info) + + m._get_current_statemachine = AsyncMock(side_effect=get_current_statemachine) + m.get_status = SandboxManager.get_status.__get__(m, SandboxManager) + m._refresh_timeout = AsyncMock() + return m + + @pytest.mark.asyncio + async def test_get_status_include_all_states_with_k8s_404( + self, mgr_with_k8s_operator, mock_meta_store, mock_provider + ): + """K8s CRD deleted + include_all_states=True → returns stopped state from meta_store.""" + from kubernetes.client.exceptions import ApiException + + mock_meta_store.get.return_value = { + "state": State.STOPPED, + "phases": {}, + "port_mapping": {22: 41125, 8080: 51275, 22555: 43823}, + "host_ip": "10.5.170.134", + } + mock_provider.get_status = AsyncMock(side_effect=ApiException(status=404, reason="Not Found")) + + result = await mgr_with_k8s_operator.get_status("test-sandbox", include_all_states=True) + + assert result.state == State.STOPPED + assert result.is_alive is False + assert result.port_mapping == {22: 41125, 8080: 51275, 22555: 43823} + + class TestMergeSandboxInfo: """Test cases for _merge_sandbox_info function.""" From f4220355c9d4dc789311b87b5852ce5ce9420a63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:04:18 +0800 Subject: [PATCH 177/226] docs: add rock sdk mcp migration design --- ...06-rock-sdk-mcp-mcpenv-migration-design.md | 279 ++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-06-rock-sdk-mcp-mcpenv-migration-design.md diff --git a/docs/superpowers/specs/2026-06-06-rock-sdk-mcp-mcpenv-migration-design.md b/docs/superpowers/specs/2026-06-06-rock-sdk-mcp-mcpenv-migration-design.md new file mode 100644 index 0000000000..e84fd7b419 --- /dev/null +++ b/docs/superpowers/specs/2026-06-06-rock-sdk-mcp-mcpenv-migration-design.md @@ -0,0 +1,279 @@ +# ROCK SDK MCP McpEnv Migration Design + +## Context + +ScaffoldHub currently owns `McpEnv` and `RockRuntime` under `scaffoldhub.sdk`. +`McpEnv` starts MCP servers inside a ROCK sandbox, resolves server auth +placeholders through ScaffoldHub auth/tool resources, and delegates data +lifecycle operations to ScaffoldHub lifecycle implementations. + +The desired ownership is: + +- ROCK owns the MCP sandbox SDK facade and ROCK runtime integration. +- ScaffoldHub owns tool resources such as auth providers, data lifecycle + factories, and concrete Slack/Notion lifecycle implementations. +- ScaffoldHub no longer depends on the ROCK SDK after the migration. + +## Goals + +- Add a new Python SDK package at `rock.sdk.mcp`. +- Move the public `McpEnv` API from ScaffoldHub into `rock.sdk.mcp`. +- Move `RockRuntime`, `RockRuntimeConfig`, runtime errors, and + `BeforeLaunchHook` into `rock.sdk.mcp`. +- Preserve the current `McpEnv` behavior rather than redesigning the API. +- Make ScaffoldHub an optional dependency for users of `rock.sdk.mcp`. +- Allow ScaffoldHub to remove its `rl-rock` dependency after it stops exporting + `scaffoldhub.sdk.McpEnv`. + +## Non-Goals + +- Do not add a top-level `rock.mcp` package. +- Do not redesign `McpEnv` method names, return values, or lifecycle semantics. +- Do not move ScaffoldHub tool implementations into ROCK. +- Do not introduce a compatibility re-export from `scaffoldhub.sdk` back to + `rock.sdk.mcp`, because that would keep ScaffoldHub dependent on ROCK. +- Do not implement generic data upload APIs in `McpEnv`. + +## Proposed Package Layout + +Add: + +```text +rock/sdk/mcp/ +├── __init__.py +├── mcp_env.py +└── rock_runtime.py +``` + +Public imports: + +```python +from rock.sdk.mcp import McpEnv +from rock.sdk.mcp.rock_runtime import RockRuntime, RockRuntimeConfig +``` + +`rock.sdk.mcp.__init__` exports only the MCP SDK surface needed by callers: + +```python +__all__ = ["McpEnv"] +``` + +## Dependency Direction + +The target dependency graph is: + +```text +rock.sdk.mcp + -> rock.sdk.sandbox.client.Sandbox + -> rock.sdk.sandbox.config.SandboxConfig + -> rock.actions.Command + -> scaffoldhub.tools.base.DataLifecycleFactory +``` + +ScaffoldHub keeps: + +```text +scaffoldhub.auth +scaffoldhub.tools.base +scaffoldhub.tools.slack +scaffoldhub.tools.notion +``` + +ROCK should add an optional extra: + +```toml +[project.optional-dependencies] +mcp = [ + "scaffoldhub>=0.1.0", +] +``` + +The base `rl-rock` install should not force ScaffoldHub onto every user. +Callers that import or instantiate `McpEnv` without ScaffoldHub installed should +receive a clear import error telling them to install `rl-rock[mcp]` or +`scaffoldhub`. + +## McpEnv Behavior + +`McpEnv` remains a thin facade with the existing ScaffoldHub semantics. + +### Construction + +- `McpEnv(servers=None)` uses `{}`. +- `servers` must be a dict or `None`; otherwise raise + `TypeError("servers must be a dict")`. +- Store a defensive copy in `self.servers`. +- Initialize: + - `self.running = False` + - `self.urls = {}` + - `self.resolved_servers = {}` + - `self.data_lifecycle_factory = DataLifecycleFactory()` + - `self.data_lifecycles` for server keys supported by the factory + - `self._rock_runtime = RockRuntime()` + +### Sandbox Access + +Expose `sandbox` as a property returning the raw ROCK `Sandbox | None` from the +runtime. This preserves the temporary low-level escape hatch used by +`before_launch` integrations. + +### Start + +`await start(before_launch=None)` should: + +- Set `running` to `False`. +- Clear `urls`. +- Resolve each server config into `resolved_servers`. +- Replace full-string env placeholders such as `${SLACK_MCP_XOXP_TOKEN}` with + auth values from the ScaffoldHub auth provider when available. +- Leave partial templates, unknown placeholders, non-string values, and + non-dict server configs unchanged. +- Delegate sandbox startup to `RockRuntime.start`. +- Set `urls` and `running=True` only after successful launch and health checks. + +### Data Lifecycle + +`init(data)`: + +- Requires `data` to be a dict. +- Only dispatches keys present in both `data_lifecycles` and `data`. +- Requires intersection values to be dicts. +- Does not store `self.data`. + +`reset()`: + +- Calls `reset()` on every configured lifecycle. +- Does not stop the ROCK runtime. +- Does not clear `urls`, `resolved_servers`, or `running`. + +`dump()`: + +- Calls `dump()` on every configured lifecycle. +- Returns only non-empty lifecycle dumps. +- Does not require a prior `init()`. + +`release()`: + +- Stops the ROCK runtime only when `running` is true. +- Logs runtime stop failures. +- Always sets `running=False`, clears `urls`, and clears `resolved_servers`. +- Preserves lifecycle instances and leaves data cleanup to explicit `reset()`. + +## RockRuntime Behavior + +`RockRuntime` also preserves the current ScaffoldHub implementation. + +`RockRuntimeConfig.from_env()` reads: + +- `ROCK_API_KEY` and `ROCK_USER_ID` as required values. +- `ROCK_BASE_URL`, defaulting to `https://xrl.alibaba-inc.com`. +- `ROCK_SANDBOX_IMAGE`, defaulting to + `rock-registry.cn-hangzhou.cr.aliyuncs.com/envs/mcp-atlas-local:v0.4.0`. +- `ROCK_EXPERIMENT_ID`, defaulting to `mcpenv`. +- `ROCK_CLUSTER`, defaulting to `nt-a`. +- `ROCK_SANDBOX_CPUS`, defaulting to `4`. +- `ROCK_SANDBOX_MEMORY`, defaulting to `8g`. +- `ROCK_AUTO_CLEAR_SECONDS`, defaulting to `3600`. + +`start(servers, before_launch=None)` should: + +- Reject repeated start attempts. +- Create `Sandbox(SandboxConfig(...))` using ROCK SDK. +- Start the sandbox. +- Create `/app/workspace` and `/data`. +- Write `/app/mcp-servers.json` with `{"mcpServers": servers}`. +- Run the optional sync or async `before_launch(sandbox)` callback. +- Execute `bash /app/launch.sh > /tmp/launch.log 2>&1 &`. +- Health-check each server's SSE proxy URL. +- Return `{server_name: sse_url}`. +- Stop the sandbox on startup failure and preserve the original startup error + as the cause. + +`stop()` clears runtime state and stops the sandbox when present. + +Diagnostic helpers such as `dump_sandbox_logs`, `upload_file`, and `read_file` +should move with `RockRuntime` because they operate on the owned sandbox. + +## ScaffoldHub Changes + +After ROCK gains `rock.sdk.mcp`, ScaffoldHub should: + +- Remove `rl-rock>=...` from its project dependencies. +- Remove or stop exporting `scaffoldhub.sdk.McpEnv`. +- Remove or move `scaffoldhub.sdk.rock_runtime`. +- Keep `DataLifecycle`, `DataLifecycleFactory`, `AuthProvider`, and concrete + lifecycle implementations. +- Update examples and docs from: + +```python +from scaffoldhub.sdk import McpEnv +``` + +to: + +```python +from rock.sdk.mcp import McpEnv +``` + +No compatibility re-export should be added in ScaffoldHub, because it would +reintroduce a package-level cycle. + +## Testing + +Move ScaffoldHub MCP SDK tests into ROCK: + +```text +tests/unit/sdk/mcp/test_mcp_env.py +tests/unit/sdk/mcp/test_rock_runtime.py +tests/integration/sdk/mcp/test_mcp_env_rock_integration.py +``` + +Unit tests should avoid real ROCK services by using fake runtime or fake +sandbox objects. They should cover: + +- `McpEnv` constructor validation. +- Server env placeholder resolution. +- Unavailable auth keeps placeholders unchanged. +- `init`, `reset`, `dump`, and `release` behavior. +- `before_launch` signature compatibility. +- Raw sandbox property exposure. +- `RockRuntimeConfig.from_env()` required and default values. +- MCP server JSON rendering. +- Server URL construction. +- Repeated runtime start rejection. +- Startup cleanup and original error preservation. + +Integration tests that require real ROCK services should remain marked so they +do not run in the fast default test set unless the environment is prepared. + +## Verification + +Expected verification commands after implementation: + +```bash +uv run pytest tests/unit/sdk/mcp -v +uv run ruff check rock/sdk/mcp tests/unit/sdk/mcp +uv run ruff format rock/sdk/mcp tests/unit/sdk/mcp +``` + +If ScaffoldHub changes are made in a separate repository, run its lifecycle +tests there after removing `rl-rock` from dependencies. + +## Migration Risks + +- If `scaffoldhub` is installed with a version that lacks + `scaffoldhub.tools.base.DataLifecycleFactory`, `McpEnv` cannot create + lifecycles. The import error should identify the missing optional dependency. +- If ScaffoldHub keeps `scaffoldhub.sdk.McpEnv` as a re-export, package + dependency cycles may persist. The migration should avoid that compatibility + path. +- `RockRuntimeConfig` environment variable names are currently inherited from + ScaffoldHub. This migration keeps them unchanged to preserve behavior. + +## Release Coordination + +- ROCK adds `scaffoldhub>=0.1.0` to the `mcp` optional extra. +- ROCK releases `rock.sdk.mcp` before ScaffoldHub removes its SDK facade. +- ScaffoldHub repository edits are a separate implementation unit from this + ROCK repository change. They remove the `rl-rock` dependency and update + ScaffoldHub docs/tests after the ROCK API is available. From 9f7fceeec28218ca897239338bf55bf42500a013 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:21:39 +0800 Subject: [PATCH 178/226] test(sdk): add mcp rock runtime tests --- tests/unit/sdk/mcp/test_rock_runtime.py | 154 ++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 tests/unit/sdk/mcp/test_rock_runtime.py diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py new file mode 100644 index 0000000000..d412909c14 --- /dev/null +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -0,0 +1,154 @@ +import asyncio +import inspect +import json +from types import SimpleNamespace + +import pytest + +from rock.sdk.mcp import rock_runtime +from rock.sdk.mcp.rock_runtime import RockRuntime, RockRuntimeConfig, RockRuntimeConfigError, RockRuntimeError + + +def test_rock_runtime_config_requires_api_key(monkeypatch): + monkeypatch.delenv("ROCK_API_KEY", raising=False) + monkeypatch.setenv("ROCK_USER_ID", "user-001") + + with pytest.raises(RockRuntimeConfigError, match="ROCK_API_KEY is required"): + RockRuntimeConfig.from_env() + + +def test_rock_runtime_config_requires_user_id(monkeypatch): + monkeypatch.setenv("ROCK_API_KEY", "rock-key") + monkeypatch.delenv("ROCK_USER_ID", raising=False) + + with pytest.raises(RockRuntimeConfigError, match="ROCK_USER_ID is required"): + RockRuntimeConfig.from_env() + + +def test_rock_runtime_config_reads_defaults_and_numeric_values(monkeypatch): + monkeypatch.setenv("ROCK_API_KEY", "rock-key") + monkeypatch.setenv("ROCK_USER_ID", "user-001") + monkeypatch.delenv("ROCK_BASE_URL", raising=False) + monkeypatch.delenv("ROCK_SANDBOX_IMAGE", raising=False) + monkeypatch.delenv("ROCK_EXPERIMENT_ID", raising=False) + monkeypatch.delenv("ROCK_CLUSTER", raising=False) + monkeypatch.delenv("ROCK_SANDBOX_CPUS", raising=False) + monkeypatch.delenv("ROCK_SANDBOX_MEMORY", raising=False) + monkeypatch.delenv("ROCK_AUTO_CLEAR_SECONDS", raising=False) + + config = RockRuntimeConfig.from_env() + + assert config.api_key == "rock-key" + assert config.user_id == "user-001" + assert config.base_url == "https://xrl.alibaba-inc.com" + assert config.image == "rock-registry.cn-hangzhou.cr.aliyuncs.com/envs/mcp-atlas-local:v0.4.0" + assert config.experiment_id == "mcpenv" + assert config.cluster == "nt-a" + assert config.cpus == 4.0 + assert config.memory == "8g" + assert config.auto_clear_seconds == 3600 + + +def test_rock_runtime_builds_server_urls_from_sandbox_id(monkeypatch): + monkeypatch.setenv("ROCK_API_KEY", "rock-key") + monkeypatch.setenv("ROCK_USER_ID", "user-001") + monkeypatch.delenv("ROCK_BASE_URL", raising=False) + runtime = RockRuntime(config=RockRuntimeConfig.from_env()) + runtime._sandbox_id = "sandbox-123" + + assert runtime.get_all_server_urls(["calculator", "slack"]) == { + "calculator": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/calculator/sse", + "slack": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/slack/sse", + } + + +def test_rock_runtime_serializes_mcp_server_config(monkeypatch): + monkeypatch.setenv("ROCK_API_KEY", "rock-key") + monkeypatch.setenv("ROCK_USER_ID", "user-001") + monkeypatch.delenv("ROCK_BASE_URL", raising=False) + runtime = RockRuntime(config=RockRuntimeConfig.from_env()) + + rendered = runtime.build_mcp_servers_json( + { + "calculator": { + "command": "uvx", + "args": ["mcp-server-calculator==0.2.0"], + } + } + ) + + assert json.loads(rendered) == { + "mcpServers": { + "calculator": { + "command": "uvx", + "args": ["mcp-server-calculator==0.2.0"], + } + } + } + + +def test_rock_runtime_rejects_start_when_sandbox_is_already_started(monkeypatch): + monkeypatch.setenv("ROCK_API_KEY", "rock-key") + monkeypatch.setenv("ROCK_USER_ID", "user-001") + monkeypatch.delenv("ROCK_BASE_URL", raising=False) + runtime = RockRuntime(config=RockRuntimeConfig.from_env()) + runtime._sandbox_id = "sandbox-123" + + with pytest.raises(RockRuntimeError, match="ROCK runtime has already been started"): + asyncio.run(runtime.start({})) + + +def test_rock_runtime_start_accepts_before_launch_hook(): + signature = inspect.signature(RockRuntime.start) + + assert "before_launch" in signature.parameters + assert signature.parameters["before_launch"].default is None + + +def test_rock_runtime_exposes_raw_sandbox_property(): + assert isinstance(RockRuntime.sandbox, property) + + +def test_rock_runtime_start_preserves_before_launch_error_when_cleanup_fails(monkeypatch): + class StopFailingSandbox: + sandbox_id = "sandbox-123" + + def __init__(self, config): + self.config = config + + async def start(self): + pass + + async def execute(self, command): + return SimpleNamespace(exit_code=0) + + async def write_file_by_path(self, content, path): + pass + + async def stop(self): + raise RuntimeError("stop failed") + + monkeypatch.setattr(rock_runtime, "Sandbox", StopFailingSandbox) + runtime = RockRuntime( + config=RockRuntimeConfig( + base_url="https://xrl.alibaba-inc.com", + api_key="rock-key", + image="image", + user_id="user-001", + experiment_id="experiment", + cluster="cluster", + cpus=1.0, + memory="1g", + auto_clear_seconds=60, + ) + ) + + async def before_launch(_sandbox): + raise RuntimeError("hook failed") + + with pytest.raises(RockRuntimeError, match="hook failed") as exc_info: + asyncio.run(runtime.start({"calculator": {}}, before_launch=before_launch)) + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "hook failed" + assert runtime.sandbox is None From 2952710c45b6d7c7f32e372711ad97d6e2d2589c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:24:08 +0800 Subject: [PATCH 179/226] feat(sdk): add mcp rock runtime --- rock/sdk/mcp/__init__.py | 1 + rock/sdk/mcp/rock_runtime.py | 311 +++++++++++++++++++++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 rock/sdk/mcp/__init__.py create mode 100644 rock/sdk/mcp/rock_runtime.py diff --git a/rock/sdk/mcp/__init__.py b/rock/sdk/mcp/__init__.py new file mode 100644 index 0000000000..c9c2ef67bd --- /dev/null +++ b/rock/sdk/mcp/__init__.py @@ -0,0 +1 @@ +__all__: list[str] = [] diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py new file mode 100644 index 0000000000..860afadd1c --- /dev/null +++ b/rock/sdk/mcp/rock_runtime.py @@ -0,0 +1,311 @@ +from __future__ import annotations + +import asyncio +import inspect +import json +import logging +import os +from collections.abc import Awaitable, Callable, Iterable +from dataclasses import dataclass +from typing import Any + +import httpx + +from rock.actions import Command +from rock.sdk.sandbox.client import Sandbox +from rock.sdk.sandbox.config import SandboxConfig + +BeforeLaunchHook = Callable[[Sandbox], Awaitable[None] | None] + +logger = logging.getLogger(__name__) + + +class RockRuntimeError(RuntimeError): + """Raised when the ROCK runtime cannot start, health check, or stop cleanly.""" + + +class RockRuntimeConfigError(RockRuntimeError): + """Raised when required ROCK runtime configuration is missing or invalid.""" + + +@dataclass(frozen=True) +class RockRuntimeConfig: + base_url: str + api_key: str + image: str + user_id: str + experiment_id: str + cluster: str + cpus: float + memory: str + auto_clear_seconds: int + + @classmethod + def from_env(cls) -> RockRuntimeConfig: + api_key = os.getenv("ROCK_API_KEY", "").strip() + if not api_key: + raise RockRuntimeConfigError("ROCK_API_KEY is required") + + user_id = os.getenv("ROCK_USER_ID", "").strip() + if not user_id: + raise RockRuntimeConfigError("ROCK_USER_ID is required") + + try: + cpus = float(os.getenv("ROCK_SANDBOX_CPUS", "4")) + except ValueError as error: + raise RockRuntimeConfigError("ROCK_SANDBOX_CPUS must be a number") from error + + try: + auto_clear_seconds = int(os.getenv("ROCK_AUTO_CLEAR_SECONDS", "3600")) + except ValueError as error: + raise RockRuntimeConfigError("ROCK_AUTO_CLEAR_SECONDS must be an integer") from error + + return cls( + base_url=os.getenv("ROCK_BASE_URL", "https://xrl.alibaba-inc.com").rstrip("/"), + api_key=api_key, + image=os.getenv( + "ROCK_SANDBOX_IMAGE", + "rock-registry.cn-hangzhou.cr.aliyuncs.com/envs/mcp-atlas-local:v0.4.0", + ), + user_id=user_id, + experiment_id=os.getenv("ROCK_EXPERIMENT_ID", "mcpenv"), + cluster=os.getenv("ROCK_CLUSTER", "nt-a"), + cpus=cpus, + memory=os.getenv("ROCK_SANDBOX_MEMORY", "8g"), + auto_clear_seconds=auto_clear_seconds, + ) + + +class RockRuntime: + def __init__( + self, + config: RockRuntimeConfig | None = None, + *, + health_check_retries: int = 10, + health_check_interval_seconds: float = 10.0, + http_timeout_seconds: float = 10.0, + ): + self.config = config + self.health_check_retries = health_check_retries + self.health_check_interval_seconds = health_check_interval_seconds + self.http_timeout_seconds = http_timeout_seconds + self._sandbox: Sandbox | None = None + self._sandbox_id: str | None = None + self._started = False + + @property + def sandbox_id(self) -> str | None: + if self._sandbox is not None: + return self._sandbox.sandbox_id + return self._sandbox_id + + @property + def sandbox(self) -> Sandbox | None: + return self._sandbox + + @property + def sse_headers(self) -> dict[str, str]: + config = self._require_config() + return {"XRL-Authorization": f"Bearer {config.api_key}"} + + def build_mcp_servers_json(self, servers: dict[str, Any]) -> str: + return json.dumps({"mcpServers": servers}, indent=2) + + def get_server_url(self, server_name: str) -> str: + sandbox_id = self.sandbox_id + if not sandbox_id: + raise RockRuntimeError("ROCK sandbox has not been started") + + config = self._require_config() + return f"{config.base_url}/apis/envs/sandbox/v1/sandboxes/{sandbox_id}/proxy/{server_name}/sse" + + def get_all_server_urls(self, server_names: Iterable[str]) -> dict[str, str]: + return {name: self.get_server_url(name) for name in server_names} + + async def start( + self, + servers: dict[str, Any], + before_launch: BeforeLaunchHook | None = None, + ) -> dict[str, str]: + if self._started or self._sandbox is not None or self._sandbox_id is not None: + raise RockRuntimeError("ROCK runtime has already been started") + + config = self._require_config() + self._started = False + self._sandbox = Sandbox( + SandboxConfig( + base_url=config.base_url, + extra_headers=self.sse_headers, + image=config.image, + user_id=config.user_id, + experiment_id=config.experiment_id, + cluster=config.cluster, + auto_clear_seconds=config.auto_clear_seconds, + cpus=config.cpus, + memory=config.memory, + ) + ) + + try: + await self._sandbox.start() + self._sandbox_id = self._sandbox.sandbox_id + await self._prepare_directories() + await self._write_mcp_config(servers) + await self._run_before_launch_hook(before_launch) + await self._launch_servers() + await self._health_check(sorted(servers.keys())) + self._started = True + return self.get_all_server_urls(sorted(servers.keys())) + except Exception as error: + try: + await self.stop() + except Exception as cleanup_error: + logger.warning("Failed to stop ROCK runtime after startup failure: %s", cleanup_error) + if isinstance(error, RockRuntimeError): + raise + raise RockRuntimeError(f"Failed to start ROCK runtime: {error}") from error + + async def stop(self) -> None: + sandbox = self._sandbox + self._started = False + self._sandbox = None + self._sandbox_id = None + if sandbox is not None: + await sandbox.stop() + + async def _prepare_directories(self) -> None: + sandbox = self._require_sandbox() + result = await sandbox.execute(Command(command=["bash", "-c", "mkdir -p /app/workspace /data"])) + if result.exit_code != 0: + raise RockRuntimeError(f"Failed to prepare sandbox directories: {result}") + + async def _write_mcp_config(self, servers: dict[str, Any]) -> None: + sandbox = self._require_sandbox() + await sandbox.write_file_by_path( + content=self.build_mcp_servers_json(servers), + path="/app/mcp-servers.json", + ) + + async def _run_before_launch_hook( + self, + before_launch: BeforeLaunchHook | None, + ) -> None: + if before_launch is None: + return + + result = before_launch(self._require_sandbox()) + if inspect.isawaitable(result): + await result + + async def _launch_servers(self) -> None: + sandbox = self._require_sandbox() + result = await sandbox.execute( + Command(command=["bash", "-c", "bash /app/launch.sh > /tmp/launch.log 2>&1 &"]) + ) + if result.exit_code != 0: + raise RockRuntimeError(f"Failed to launch MCP servers: {result}") + + async def _health_check(self, server_names: list[str]) -> None: + pending = set(server_names) + if not pending: + return + + async def check_one(client: httpx.AsyncClient, server_name: str) -> tuple[str, bool, str]: + url = self.get_server_url(server_name) + try: + async with client.stream("GET", url, headers=self.sse_headers) as response: + content_type = response.headers.get("content-type", "") + if "text/event-stream" in content_type: + return server_name, True, "SSE stream ready" + return server_name, False, f"status={response.status_code}, content-type={content_type}" + except Exception as error: + return server_name, False, str(error) + + last_details: dict[str, str] = {} + for attempt in range(self.health_check_retries): + async with httpx.AsyncClient(timeout=self.http_timeout_seconds) as client: + results = await asyncio.gather(*(check_one(client, name) for name in sorted(pending))) + + for server_name, ok, detail in results: + if ok: + pending.discard(server_name) + else: + last_details[server_name] = detail + + if not pending: + return + + if attempt < self.health_check_retries - 1: + await asyncio.sleep(self.health_check_interval_seconds) + + detail = ", ".join(f"{name}: {last_details.get(name, 'not ready')}" for name in sorted(pending)) + raise RockRuntimeError(f"ROCK MCP server health check failed: {detail}") + + def _require_config(self) -> RockRuntimeConfig: + if self.config is None: + self.config = RockRuntimeConfig.from_env() + return self.config + + def _require_sandbox(self) -> Sandbox: + if self._sandbox is None: + raise RockRuntimeError("ROCK sandbox has not been started") + return self._sandbox + + async def dump_sandbox_logs(self) -> dict[str, str]: + """Dump sandbox logs for debugging when startup fails.""" + sandbox = self._sandbox + if sandbox is None: + return {} + + results: dict[str, str] = {} + for label, path in [ + ("launch", "/tmp/launch.log"), + ("github", "/app/config/dynamic/logs/github.log"), + ("dynamic_config", "/app/config/dynamic/github.json"), + ]: + try: + result = await sandbox.execute( + Command(command=["bash", "-c", f"cat {path} 2>/dev/null || echo '(no {label}.log)'"]), + ) + results[label] = result.stdout + except Exception as error: + results[label] = f"(failed to read: {error})" + + try: + result = await sandbox.execute( + Command( + command=[ + "bash", + "-c", + ( + "ls -la /usr/local/bin/github-mcp-server 2>/dev/null; " + "echo '---'; ps aux 2>/dev/null || ps -ef 2>/dev/null || echo '(no ps)'" + ), + ] + ), + ) + results["process_info"] = result.stdout + except Exception as error: + results["process_info"] = f"(failed to read: {error})" + + return results + + async def upload_file(self, local_path: str, remote_path: str) -> None: + """Upload a file to the sandbox.""" + from rock.actions import UploadRequest + + sandbox = self._require_sandbox() + request = UploadRequest(source_path=local_path, target_path=remote_path) + result = await sandbox.upload(request) + if not result.success: + raise RockRuntimeError(f"Upload failed: {result.message}") + logger.info("Uploaded %s -> %s", local_path, remote_path) + + async def read_file(self, remote_path: str) -> str: + """Read a file from the sandbox.""" + from rock.actions import ReadFileRequest + + sandbox = self._require_sandbox() + request = ReadFileRequest(path=remote_path) + result = await sandbox.read_file(request) + return result.content From 54e12fea4ddaf2664e5f8f1ec8fd1d194fbe6c74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:26:32 +0800 Subject: [PATCH 180/226] test(sdk): add mcpenv tests --- tests/unit/sdk/mcp/test_mcp_env.py | 383 +++++++++++++++++++++++++++++ 1 file changed, 383 insertions(+) create mode 100644 tests/unit/sdk/mcp/test_mcp_env.py diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py new file mode 100644 index 0000000000..989330c7f4 --- /dev/null +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -0,0 +1,383 @@ +import asyncio +import importlib +import inspect +import sys +from copy import deepcopy +from types import ModuleType + +import pytest + + +class RecordingRuntime: + def __init__(self): + self.stopped = False + + async def stop(self): + self.stopped = True + + +class FailingStopRuntime: + async def stop(self): + raise RuntimeError("stop failed") + + +class RecordingDataLifecycle: + def __init__(self): + self.initialized_data = {} + self.reset_calls = 0 + + def init(self, data: dict) -> None: + self.initialized_data = deepcopy(data) + + def dump(self) -> dict: + return deepcopy(self.initialized_data) + + def reset(self) -> None: + self.reset_calls += 1 + self.initialized_data = {} + + +class FakeAuthProvider: + def __init__(self): + self.auth = { + "slack": { + "SLACK_MCP_XOXP_TOKEN": "xoxp-test-token", + "SLACK_MCP_XOXB_TOKEN": "xoxb-test-token", + } + } + + def provide(self, platform: str) -> dict: + if platform not in self.auth: + raise ValueError(f"Unsupported platform: {platform}") + return self.auth[platform] + + +class FakeDataLifecycleFactory: + def __init__(self): + self.auth_provider = FakeAuthProvider() + self.created = {} + + def supports(self, lifecycle_type: str) -> bool: + return lifecycle_type == "slack" + + def create(self, lifecycle_type: str): + if lifecycle_type != "slack": + raise ValueError(f"Unsupported data lifecycle type: {lifecycle_type}") + lifecycle = RecordingDataLifecycle() + self.created[lifecycle_type] = lifecycle + return lifecycle + + +def install_fake_scaffoldhub(monkeypatch): + scaffoldhub = ModuleType("scaffoldhub") + tools = ModuleType("scaffoldhub.tools") + base = ModuleType("scaffoldhub.tools.base") + base.DataLifecycleFactory = FakeDataLifecycleFactory + + monkeypatch.setitem(sys.modules, "scaffoldhub", scaffoldhub) + monkeypatch.setitem(sys.modules, "scaffoldhub.tools", tools) + monkeypatch.setitem(sys.modules, "scaffoldhub.tools.base", base) + + +def reload_mcp_env(monkeypatch): + install_fake_scaffoldhub(monkeypatch) + sys.modules.pop("rock.sdk.mcp.mcp_env", None) + module = importlib.import_module("rock.sdk.mcp.mcp_env") + return importlib.reload(module) + + +def slack_server_config() -> dict: + return { + "command": "npx", + "args": [ + "slack-mcp-server@1.1.23", + "--transport", + "stdio", + ], + "env": { + "SLACK_MCP_XOXP_TOKEN": "${SLACK_MCP_XOXP_TOKEN}", + "SLACK_MCP_XOXB_TOKEN": "${SLACK_MCP_XOXB_TOKEN}", + "STATIC_VALUE": "unchanged", + "PARTIAL_TEMPLATE": "token-${SLACK_MCP_XOXP_TOKEN}", + "UNKNOWN_PLACEHOLDER": "${missing_token}", + }, + } + + +def test_mcp_env_init_dump_and_release_with_declared_server(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingDataLifecycle() + env.data_lifecycles["slack"] = lifecycle + data = { + "slack": { + "seed_fixture_path": "src/scaffoldhub/tools/slack/slack_mcp_eval_fixture.json", + "channels": ["general", "random"], + "users": { + "u001": "alice", + "u002": "bob", + }, + }, + } + expected_data = deepcopy(data) + + assert "slack" in env.data_lifecycles + assert env.dump() == {} + + env.init(data) + + dumped_data = env.dump() + + assert dumped_data == expected_data + assert not hasattr(env, "data") + + dumped_data["slack"]["channels"].append("alerts") + + assert env.dump() == expected_data + assert data == expected_data + + asyncio.run(env.release()) + + assert env.is_alive() is False + assert env.urls == {} + assert env.data_lifecycles["slack"] is lifecycle + assert env.resolved_servers == {} + + +def test_mcp_env_constructor_requires_servers_dict(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + + with pytest.raises(TypeError, match="servers must be a dict"): + mcp_env.McpEnv(servers=[]) + + +def test_mcp_env_constructor_reports_missing_scaffoldhub(monkeypatch): + sys.modules.pop("scaffoldhub.tools.base", None) + sys.modules.pop("rock.sdk.mcp.mcp_env", None) + module = importlib.import_module("rock.sdk.mcp.mcp_env") + module = importlib.reload(module) + + def raise_missing_dependency(): + raise ImportError("rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`.") + + monkeypatch.setattr(module, "_load_data_lifecycle_factory", raise_missing_dependency) + + with pytest.raises(ImportError, match=r"rl-rock\[mcp\]"): + module.McpEnv() + + +def test_mcp_env_resolves_server_env_placeholders_without_starting_runtime(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + resolved = env._resolve_server_config("slack", slack_server_config()) + + assert env.is_alive() is False + assert resolved == { + "command": "npx", + "args": [ + "slack-mcp-server@1.1.23", + "--transport", + "stdio", + ], + "env": { + "SLACK_MCP_XOXP_TOKEN": "xoxp-test-token", + "SLACK_MCP_XOXB_TOKEN": "xoxb-test-token", + "STATIC_VALUE": "unchanged", + "PARTIAL_TEMPLATE": "token-${SLACK_MCP_XOXP_TOKEN}", + "UNKNOWN_PLACEHOLDER": "${missing_token}", + }, + } + assert env.servers["slack"]["env"]["SLACK_MCP_XOXP_TOKEN"] == "${SLACK_MCP_XOXP_TOKEN}" + + +def test_mcp_env_resolution_keeps_placeholders_when_auth_is_unavailable(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv( + servers={ + "github": { + "command": "github-mcp-server", + "env": { + "GITHUB_TOKEN": "${github_token}", + }, + } + } + ) + + resolved = env._resolve_server_config("github", env.servers["github"]) + + assert env.is_alive() is False + assert env.data_lifecycles == {} + assert resolved == { + "command": "github-mcp-server", + "env": { + "GITHUB_TOKEN": "${github_token}", + }, + } + + +def test_mcp_env_init_with_no_servers_allows_empty_data_but_not_urls_before_start(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv() + + env.init({}) + + assert env.dump() == {} + with pytest.raises(RuntimeError, match="McpEnv has not been started"): + env.get_urls() + + asyncio.run(env.release()) + + assert env.is_alive() is False + + +def test_mcp_env_init_requires_dict(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv() + + with pytest.raises(TypeError, match="data must be a dict"): + env.init([]) + + +def test_mcp_env_init_ignores_data_keys_missing_from_servers(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + env.init({"github": {"repositories": ["example"]}}) + + assert env.dump() == {} + + +def test_mcp_env_init_ignores_server_keys_missing_from_data(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + env.init({}) + + assert env.dump() == {} + + +def test_mcp_env_init_rejects_non_dict_values_only_for_intersection_keys(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + with pytest.raises(TypeError, match="data for slack must be a dict"): + env.init({"slack": []}) + + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + env.init({"github": []}) + + assert env.dump() == {} + + +def test_mcp_env_reset_delegates_to_configured_lifecycles_without_stopping_runtime(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingDataLifecycle() + env.data_lifecycles["slack"] = lifecycle + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + env.init({"slack": {"channels": ["general"]}}) + + env.reset() + + assert lifecycle.reset_calls == 1 + assert env.dump() == {} + assert env.is_alive() is True + assert env.urls == {"slack": "https://example.test/slack/sse"} + assert env.resolved_servers == {"slack": slack_server_config()} + + +def test_mcp_env_reset_does_not_require_prior_init(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingDataLifecycle() + env.data_lifecycles["slack"] = lifecycle + + env.reset() + + assert lifecycle.reset_calls == 1 + assert env.dump() == {} + + +def test_mcp_env_dump_before_init_returns_empty_data(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + assert env.dump() == {} + + +def test_mcp_env_get_urls_requires_started_runtime_not_data_init(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + + assert env.get_urls() == {"slack": "https://example.test/slack/sse"} + + +def test_mcp_env_release_without_start_or_init_is_noop(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + asyncio.run(env.release()) + + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + + +def test_mcp_env_release_after_start_before_init_stops_runtime_and_preserves_lifecycles(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + runtime = RecordingRuntime() + lifecycle = RecordingDataLifecycle() + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = runtime + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + + asyncio.run(env.release()) + + assert runtime.stopped is True + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + env.reset() + assert lifecycle.reset_calls == 1 + + +def test_mcp_env_release_preserves_lifecycles_when_runtime_stop_fails(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingDataLifecycle() + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = FailingStopRuntime() + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + + asyncio.run(env.release()) + + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + env.reset() + assert lifecycle.reset_calls == 1 + + +def test_mcp_env_start_accepts_before_launch_hook(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + signature = inspect.signature(mcp_env.McpEnv.start) + + assert "before_launch" in signature.parameters + assert signature.parameters["before_launch"].default is None + + +def test_mcp_env_exposes_raw_sandbox_property(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + + assert isinstance(mcp_env.McpEnv.sandbox, property) From d199ef9c139f0cb948896c2986cf0b9dd2e768e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:28:36 +0800 Subject: [PATCH 181/226] feat(sdk): add mcpenv --- rock/sdk/mcp/__init__.py | 4 +- rock/sdk/mcp/mcp_env.py | 212 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 215 insertions(+), 1 deletion(-) create mode 100644 rock/sdk/mcp/mcp_env.py diff --git a/rock/sdk/mcp/__init__.py b/rock/sdk/mcp/__init__.py index c9c2ef67bd..9c70ec3c7b 100644 --- a/rock/sdk/mcp/__init__.py +++ b/rock/sdk/mcp/__init__.py @@ -1 +1,3 @@ -__all__: list[str] = [] +from rock.sdk.mcp.mcp_env import McpEnv + +__all__ = ["McpEnv"] diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py new file mode 100644 index 0000000000..1822c30b9e --- /dev/null +++ b/rock/sdk/mcp/mcp_env.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import logging +from copy import deepcopy +from typing import Any + +from rock.sdk.mcp.rock_runtime import BeforeLaunchHook, RockRuntime +from rock.sdk.sandbox.client import Sandbox + +logger = logging.getLogger(__name__) + + +def _load_data_lifecycle_factory(): + try: + from scaffoldhub.tools.base import DataLifecycleFactory + except ImportError as error: + raise ImportError( + "rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`." + ) from error + return DataLifecycleFactory + + +class McpEnv: + """ + MCP environment manager. + + It resolves MCP server configs, starts real ROCK sandboxes, delegates data + lifecycle operations to ScaffoldHub resources, exposes server URLs, and + releases runtime resources. + """ + + def __init__(self, servers: dict | None = None): + """ + Create an uninitialized MCP environment. + + Args: + servers: MCP server config. Top-level keys are server or lifecycle + types, such as ``slack``. + + Raises: + TypeError: Raised when servers is neither dict nor None. + """ + if servers is None: + servers = {} + if not isinstance(servers, dict): + raise TypeError("servers must be a dict") + + self.running = False + self.urls = {} + self.servers = deepcopy(servers) + self.resolved_servers = {} + data_lifecycle_factory = _load_data_lifecycle_factory() + self.data_lifecycle_factory = data_lifecycle_factory() + self.data_lifecycles: dict[str, Any] = {} + self._rock_runtime = RockRuntime() + + for lifecycle_type in self.servers: + if self.data_lifecycle_factory.supports(lifecycle_type): + self.data_lifecycles[lifecycle_type] = self.data_lifecycle_factory.create(lifecycle_type) + + @property + def sandbox(self) -> Sandbox | None: + """ + Return the raw ROCK Sandbox object for advanced pre-launch integrations. + + This is intentionally a low-level escape hatch for callers that need + direct ROCK SDK access while the higher-level lifecycle API evolves. + """ + return self._rock_runtime.sandbox + + async def start(self, before_launch: BeforeLaunchHook | None = None) -> None: + """ + Start a real ROCK MCP sandbox. + + The environment is marked running only after sandbox creation, optional + pre-launch callback execution, MCP server launch, and SSE health checks + all succeed. + """ + self.running = False + self.urls = {} + self.resolved_servers = { + server_name: self._resolve_server_config(server_name, server_config) + for server_name, server_config in self.servers.items() + } + urls = await self._rock_runtime.start( + self.resolved_servers, + before_launch=before_launch, + ) + self.urls = urls + self.running = True + + def is_alive(self) -> bool: + """ + Return the runtime state recorded by this facade. + + Returns: + True after successful start, false after release. + """ + return self.running + + def init(self, data: dict): + """ + Initialize MCP environment data. + + Args: + data: Layered environment data, such as ``{"slack": {...}}``. + + Raises: + TypeError: Raised when data is not a dict or an intersecting + lifecycle value is not a dict. + """ + if not isinstance(data, dict): + raise TypeError("data must be a dict") + + for lifecycle_type in self.data_lifecycles.keys() & data.keys(): + lifecycle_data = data[lifecycle_type] + if not isinstance(lifecycle_data, dict): + raise TypeError(f"data for {lifecycle_type} must be a dict") + self.data_lifecycles[lifecycle_type].init(lifecycle_data) + + def reset(self) -> None: + """ + Reset configured MCP environment data. + + This only delegates to configured data lifecycles. It does not stop the + ROCK runtime, clear URLs, or change the recorded running state. + """ + for lifecycle in self.data_lifecycles.values(): + lifecycle.reset() + + def get_urls(self) -> dict: + """ + Get MCP server URLs. + + Returns: + A defensive copy of the server URL mapping. + + Raises: + RuntimeError: Raised when the runtime has not started successfully. + """ + if not self.running: + raise RuntimeError("McpEnv has not been started") + + return deepcopy(self.urls) + + def dump(self) -> dict: + """ + Export current MCP environment data. + + Returns: + Layered data returned by configured data lifecycles. + """ + dumped_data = {} + for lifecycle_type, lifecycle in self.data_lifecycles.items(): + lifecycle_data = lifecycle.dump() + if lifecycle_data != {}: + dumped_data[lifecycle_type] = lifecycle_data + return dumped_data + + async def release(self): + """ + Release the physical MCP runtime. + + Data cleanup is intentionally handled by explicit reset calls. + """ + try: + if self.running: + try: + await self._rock_runtime.stop() + except Exception as error: + logger.warning("Failed to stop ROCK runtime during release: %s", error) + finally: + self.running = False + self.urls = {} + self.resolved_servers = {} + + def _resolve_server_config(self, server_name: str, server_config: Any) -> Any: + if not isinstance(server_config, dict): + return deepcopy(server_config) + + resolved_config = deepcopy(server_config) + env = resolved_config.get("env") + if not isinstance(env, dict): + return resolved_config + + auth = self._server_auth(server_name) + resolved_config["env"] = { + env_key: self._resolve_env_value(env_value, auth) + for env_key, env_value in env.items() + } + return resolved_config + + def _server_auth(self, server_name: str) -> dict: + auth_provider = getattr(self.data_lifecycle_factory, "auth_provider", None) + if auth_provider is None: + return {} + + try: + return auth_provider.provide(server_name) + except ValueError: + return {} + + def _resolve_env_value(self, value: Any, auth: dict) -> Any: + if not isinstance(value, str): + return value + if not value.startswith("${") or not value.endswith("}"): + return value + + auth_key = value[2:-1] + if not auth_key: + return value + return auth.get(auth_key, value) From caf31d90ad9963ee856d1fcec4f044e5263eb61d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:32:26 +0800 Subject: [PATCH 182/226] feat(sdk): add mcp optional dependency --- pyproject.toml | 5 +++++ tests/unit/sdk/mcp/test_packaging.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 tests/unit/sdk/mcp/test_packaging.py diff --git a/pyproject.toml b/pyproject.toml index c7f88a6fca..72a0d20f3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,10 @@ model-service = [ "httpx", ] +mcp = [ + "scaffoldhub>=0.1.0.dev1", +] + all = [ "rl-rock[admin]", @@ -120,6 +124,7 @@ test = [ "pytest-trio", "pytest-twisted", "pytest-env", + "tomli>=2.0.0; python_version < '3.11'", ] [tool.setuptools.packages.find] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py new file mode 100644 index 0000000000..f21f4d8fa3 --- /dev/null +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -0,0 +1,14 @@ +from pathlib import Path + +try: + import tomllib +except ModuleNotFoundError: + import tomli as tomllib + + +def test_mcp_extra_declares_scaffoldhub_dependency(): + pyproject = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8")) + + mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] + + assert "scaffoldhub>=0.1.0.dev1" in mcp_dependencies From 485dcbd5bfc02d25845b5f47d3a1e9fcff17721c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:33:16 +0800 Subject: [PATCH 183/226] docs(sdk): add mcp sdk reference --- .../References/Python SDK References/mcp.md | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md new file mode 100644 index 0000000000..81e220fd56 --- /dev/null +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -0,0 +1,71 @@ +# MCP SDK + +`rock.sdk.mcp` provides `McpEnv`, a small SDK facade for running MCP servers +inside ROCK sandboxes. + +## Installation + +Install the MCP extra when using ScaffoldHub-backed tool lifecycles: + +```bash +pip install "rl-rock[mcp]" +``` + +## Basic Usage + +```python +import asyncio + +from rock.sdk.mcp import McpEnv + + +async def main(): + env = McpEnv( + servers={ + "calculator": { + "command": "uvx", + "args": ["mcp-server-calculator==0.2.0"], + } + } + ) + + try: + await env.start() + urls = env.get_urls() + print(urls["calculator"]) + finally: + if env.is_alive(): + await env.release() + + +asyncio.run(main()) +``` + +## Lifecycle Data + +`McpEnv.init(data)`, `McpEnv.reset()`, and `McpEnv.dump()` delegate tool data +lifecycle work to ScaffoldHub resources. The top-level keys in `data` match the +server or lifecycle type: + +```python +env.init({"slack": {"seed_fixture_path": "fixture.json"}}) +snapshot = env.dump() +env.reset() +``` + +`reset()` only resets configured tool data. It does not stop the ROCK sandbox. +Use `await env.release()` to stop the sandbox runtime. + +## Launch Hook + +`start()` accepts an optional sync or async `before_launch` callback. The +callback receives the raw ROCK `Sandbox` after `/app/mcp-servers.json` is +written and before `/app/launch.sh` starts MCP servers. + +```python +async def before_launch(sandbox): + await sandbox.write_file_by_path(content="ready", path="/data/marker.txt") + + +await env.start(before_launch=before_launch) +``` From cca98690b312e60a5b6e5b4ddbed5bbc78a47ce8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:34:43 +0800 Subject: [PATCH 184/226] test(sdk): add mcp rock integration tests --- .../sdk/mcp/test_mcp_env_rock_integration.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/integration/sdk/mcp/test_mcp_env_rock_integration.py diff --git a/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py b/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py new file mode 100644 index 0000000000..19c37f3712 --- /dev/null +++ b/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py @@ -0,0 +1,112 @@ +import asyncio +import os + +import httpx +import pytest + +from rock.actions import Command +from rock.sdk.mcp import McpEnv + +pytestmark = pytest.mark.integration + + +def calculator_server_config() -> dict: + return { + "command": "uvx", + "args": [ + "mcp-server-calculator==0.2.0", + ], + } + + +def require_rock_credentials(): + missing = [ + name + for name in ("ROCK_API_KEY", "ROCK_USER_ID") + if not os.getenv(name, "").strip() + ] + if missing: + pytest.skip(f"Missing ROCK credentials: {', '.join(missing)}") + + +async def run_real_rock_calculator_server_case(): + require_rock_credentials() + env = McpEnv(servers={"calculator": calculator_server_config()}) + + try: + await env.start() + assert env.is_alive() is True + + env.init({}) + urls = env.get_urls() + + assert set(urls) == {"calculator"} + assert urls["calculator"].endswith("/calculator/sse") + + headers = {"XRL-Authorization": f"Bearer {os.environ['ROCK_API_KEY']}"} + with httpx.stream("GET", urls["calculator"], headers=headers, timeout=10.0) as response: + assert response.status_code == 200 + assert "text/event-stream" in response.headers.get("content-type", "") + finally: + if env.is_alive(): + await env.release() + + assert env.is_alive() is False + + +def test_mcp_env_starts_real_rock_calculator_server_and_returns_sse_url(): + asyncio.run(run_real_rock_calculator_server_case()) + + +async def run_real_rock_async_before_launch_case(): + require_rock_credentials() + env = McpEnv(servers={"calculator": calculator_server_config()}) + marker_path = "/data/scaffoldhub-before-launch-async.txt" + seen_sandbox_ids: list[str] = [] + + async def before_launch(sandbox): + seen_sandbox_ids.append(sandbox.sandbox_id) + await sandbox.write_file_by_path(content="async hook was here", path=marker_path) + + try: + await env.start(before_launch=before_launch) + env.init({}) + + assert env.is_alive() is True + assert env.sandbox is not None + assert seen_sandbox_ids == [env.sandbox.sandbox_id] + + result = await env.sandbox.execute(Command(command=["bash", "-c", f"test -f {marker_path}"])) + assert result.exit_code == 0 + assert env.get_urls()["calculator"].endswith("/calculator/sse") + finally: + if env.is_alive(): + await env.release() + + +def test_mcp_env_async_before_launch_receives_real_sandbox_and_runs_before_health_check(): + asyncio.run(run_real_rock_async_before_launch_case()) + + +async def run_real_rock_sync_before_launch_case(): + require_rock_credentials() + env = McpEnv(servers={"calculator": calculator_server_config()}) + seen_sandbox_ids: list[str] = [] + + def before_launch(sandbox): + seen_sandbox_ids.append(sandbox.sandbox_id) + + try: + await env.start(before_launch=before_launch) + env.init({}) + + assert env.is_alive() is True + assert env.sandbox is not None + assert seen_sandbox_ids == [env.sandbox.sandbox_id] + finally: + if env.is_alive(): + await env.release() + + +def test_mcp_env_sync_before_launch_receives_real_sandbox(): + asyncio.run(run_real_rock_sync_before_launch_case()) From acea5fea8c3578ac4b6c65a7fb743a6e7902fb95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 07:35:53 +0800 Subject: [PATCH 185/226] style(sdk): format mcp sdk files --- rock/sdk/mcp/mcp_env.py | 7 ++----- rock/sdk/mcp/rock_runtime.py | 4 +--- tests/integration/sdk/mcp/test_mcp_env_rock_integration.py | 6 +----- 3 files changed, 4 insertions(+), 13 deletions(-) diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index 1822c30b9e..a03764be7d 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -14,9 +14,7 @@ def _load_data_lifecycle_factory(): try: from scaffoldhub.tools.base import DataLifecycleFactory except ImportError as error: - raise ImportError( - "rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`." - ) from error + raise ImportError("rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`.") from error return DataLifecycleFactory @@ -185,8 +183,7 @@ def _resolve_server_config(self, server_name: str, server_config: Any) -> Any: auth = self._server_auth(server_name) resolved_config["env"] = { - env_key: self._resolve_env_value(env_value, auth) - for env_key, env_value in env.items() + env_key: self._resolve_env_value(env_value, auth) for env_key, env_value in env.items() } return resolved_config diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index 860afadd1c..f55a55fa2d 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -199,9 +199,7 @@ async def _run_before_launch_hook( async def _launch_servers(self) -> None: sandbox = self._require_sandbox() - result = await sandbox.execute( - Command(command=["bash", "-c", "bash /app/launch.sh > /tmp/launch.log 2>&1 &"]) - ) + result = await sandbox.execute(Command(command=["bash", "-c", "bash /app/launch.sh > /tmp/launch.log 2>&1 &"])) if result.exit_code != 0: raise RockRuntimeError(f"Failed to launch MCP servers: {result}") diff --git a/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py b/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py index 19c37f3712..0137ad8772 100644 --- a/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py +++ b/tests/integration/sdk/mcp/test_mcp_env_rock_integration.py @@ -20,11 +20,7 @@ def calculator_server_config() -> dict: def require_rock_credentials(): - missing = [ - name - for name in ("ROCK_API_KEY", "ROCK_USER_ID") - if not os.getenv(name, "").strip() - ] + missing = [name for name in ("ROCK_API_KEY", "ROCK_USER_ID") if not os.getenv(name, "").strip()] if missing: pytest.skip(f"Missing ROCK credentials: {', '.join(missing)}") From 58428e75ab6a34ec2871a6ca4eb6afda2f2bc413 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 08:20:37 +0800 Subject: [PATCH 186/226] fix(sdk): configure scaffoldhub mcp source --- .../References/Python SDK References/mcp.md | 6 +++++- pyproject.toml | 10 +++++++++- tests/unit/sdk/mcp/test_packaging.py | 13 ++++++++++++- 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index 81e220fd56..efd49f7fd0 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -8,9 +8,13 @@ inside ROCK sandboxes. Install the MCP extra when using ScaffoldHub-backed tool lifecycles: ```bash -pip install "rl-rock[mcp]" +pip install --extra-index-url https://artlab.alibaba-inc.com/1/pypi/simple "rl-rock[mcp]" ``` +The MCP extra currently supports Python 3.11 and 3.12 because the published +ScaffoldHub package is Python 3.11+ and ROCK officially supports Python 3.10 to +3.12. + ## Basic Usage ```python diff --git a/pyproject.toml b/pyproject.toml index 72a0d20f3a..dea85c6e40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub>=0.1.0.dev1", + "scaffoldhub==0.1.0.dev1; python_version >= '3.11' and python_version < '3.13'", ] @@ -133,6 +133,14 @@ where = ["."] [tool.setuptools.package-data] rock = ["rocklet/local_files/*"] +[tool.uv.sources] +scaffoldhub = { index = "artlab" } + +[[tool.uv.index]] +name = "artlab" +url = "https://artlab.alibaba-inc.com/1/pypi/simple" +explicit = true + [[tool.uv.index]] url = "https://mirrors.aliyun.com/pypi/simple/" default = true diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index f21f4d8fa3..0fefae3432 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,4 +11,15 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub>=0.1.0.dev1" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev1; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + + +def test_scaffoldhub_resolves_from_artlab_index(): + pyproject = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8")) + + assert pyproject["tool"]["uv"]["sources"]["scaffoldhub"] == {"index": "artlab"} + assert { + "name": "artlab", + "url": "https://artlab.alibaba-inc.com/1/pypi/simple", + "explicit": True, + } in pyproject["tool"]["uv"]["index"] From bdeb9284263e318c207802908d317a54b0b17970 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Sat, 6 Jun 2026 23:49:18 +0800 Subject: [PATCH 187/226] chore: bump version to 1.9.0.dev1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index dea85c6e40..6402bbba46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.8.3" +version = "1.9.0.dev1" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From f59119a540a41c501e49b4dd614d27d67424078a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Tue, 9 Jun 2026 18:24:22 +0800 Subject: [PATCH 188/226] chore(sdk): upgrade scaffoldhub mcp extra --- pyproject.toml | 2 +- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6402bbba46..e8a8063fa3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev1; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev2; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index 0fefae3432..fcc6ee5944 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev1; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev2; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From 8426af141e7bfa0362e968875feda41d74126c6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 12 Jun 2026 09:11:42 +0800 Subject: [PATCH 189/226] chore: bump version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e8a8063fa3..0a5b228390 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.9.0.dev1" +version = "1.10.0.dev1" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 5318f40ee52263cc3a1246b9bdf86d2567534678 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 12 Jun 2026 09:21:01 +0800 Subject: [PATCH 190/226] chore: bump mcp release versions --- pyproject.toml | 4 ++-- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0a5b228390..1057bf286d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev1" +version = "1.10.0.dev2" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev2; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev3; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index fcc6ee5944..6bb81b597f 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev2; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev3; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From 861a97efa58cda0efb7dd26db5bd02edf137d2a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 15:25:36 +0800 Subject: [PATCH 191/226] docs(mcp): design auth lease release --- ...06-15-mcp-env-auth-lease-release-design.md | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-15-mcp-env-auth-lease-release-design.md diff --git a/docs/superpowers/specs/2026-06-15-mcp-env-auth-lease-release-design.md b/docs/superpowers/specs/2026-06-15-mcp-env-auth-lease-release-design.md new file mode 100644 index 0000000000..311a917f0e --- /dev/null +++ b/docs/superpowers/specs/2026-06-15-mcp-env-auth-lease-release-design.md @@ -0,0 +1,148 @@ +# McpEnv Auth Lease Release Design + +## Context + +ScaffoldHub 0.1.0.dev4 adds a PostgreSQL-backed, lease-aware `AuthProvider`. +When ROCK `McpEnv` resolves MCP server environment placeholders, auth values may +come from a database lease. ROCK must release those active leases when the MCP +environment is released. + +This design covers only the ROCK repository changes. ScaffoldHub owns the +database schema, auth borrowing, active lease tracking, and +`AuthProvider.release_active_leases()` implementation. + +## Goals + +- Upgrade `rl-rock[mcp]` to depend on `scaffoldhub==0.1.0.dev4`. +- Make `McpEnv` explicitly own the `AuthProvider` instance used for both server + environment resolution and lifecycle construction. +- Call `AuthProvider.release_active_leases()` from `McpEnv.release()`. +- Surface auth lease release failures to callers. +- Preserve the caller's ability to call `release()` again after an auth lease + release failure. + +## Non-Goals + +- Do not implement auth borrow or release SQL in ROCK. +- Do not add `McpEnv.__init__` parameters for custom auth providers or data + lifecycle factories. +- Do not change `reset()` semantics. `reset()` remains data lifecycle cleanup + only and does not stop the runtime or release auth leases. +- Do not silently support auth providers that lack `release_active_leases()`. + The new method is a required ScaffoldHub 0.1.0.dev4 contract. + +## Architecture + +`McpEnv` will load both `DataLifecycleFactory` and `AuthProvider` from +ScaffoldHub. During construction it will create one `AuthProvider` and pass the +same instance into `DataLifecycleFactory`: + +```python +self.auth_provider = AuthProvider() +self.data_lifecycle_factory = DataLifecycleFactory(auth_provider=self.auth_provider) +``` + +This makes auth resource ownership explicit at the ROCK boundary. The factory +still creates tool data lifecycles, but `McpEnv` owns the auth provider whose +leases must be released. + +`_server_auth()` will use `self.auth_provider.provide(server_name)` directly. +That keeps server environment placeholder resolution and +`DataLifecycleFactory.create()` on the same `AuthProvider` instance. A lease +borrowed while resolving a server env placeholder will therefore be recorded in +the same provider that `release()` later asks to release active leases. + +## Release Semantics + +`McpEnv.release()` will perform two cleanup actions: + +1. Stop the ROCK runtime if `self.running` is true. +2. Call `self.auth_provider.release_active_leases()`. + +The auth release call must not be gated by `self.running`. If the first +`release()` call stops the runtime but fails while releasing auth leases, the +caller must be able to call `release()` again and retry auth lease release even +though the runtime is no longer running. + +`release()` will clear `running`, `urls`, and `resolved_servers` in a `finally` +block. This mirrors the existing behavior that a release attempt makes the +runtime facade unusable for server URL access, even if part of cleanup reports +an error. + +## Error Handling + +Runtime stop failures keep the existing behavior: log a warning and continue. +This preserves current tests and avoids masking auth lease cleanup behind a +sandbox stop failure. + +Auth lease release failures are caller-visible. If +`release_active_leases()` raises, `McpEnv.release()` raises: + +```python +RuntimeError("Failed to release MCP auth leases") +``` + +with the original exception chained via `raise ... from error`. + +The state cleanup still happens before the exception is raised. The remaining +active lease state is owned by ScaffoldHub `AuthProvider`, so a later +`await env.release()` will call `release_active_leases()` again and allow the +provider to retry any leases it still tracks. + +## Data Flow + +1. `McpEnv.__init__()` creates `self.auth_provider`. +2. `McpEnv.__init__()` creates `self.data_lifecycle_factory` with that provider. +3. During construction, supported lifecycles are created through the factory. +4. `start()` resolves each server config. +5. `_server_auth()` calls `self.auth_provider.provide(server_name)`. +6. ScaffoldHub may borrow and track a database auth lease. +7. `release()` calls `self.auth_provider.release_active_leases()` on every + release attempt. + +## Tests + +Unit tests should cover: + +- The fake ScaffoldHub `DataLifecycleFactory` receives the same auth provider + instance that `McpEnv` stores on `env.auth_provider`. +- `release()` calls `auth_provider.release_active_leases()` on success. +- A second `release()` call still calls `release_active_leases()` when + `running` is false. +- If `release_active_leases()` raises, `release()` raises `RuntimeError` and + clears `running`, `urls`, and `resolved_servers`. +- If runtime stop raises, `release()` still calls `release_active_leases()`. +- Existing server env placeholder resolution still uses the provider returned + auth data. + +The fake auth provider in `tests/unit/sdk/mcp/test_mcp_env.py` must implement +the required `release_active_leases()` method to match the ScaffoldHub 0.1.0.dev4 +contract. + +## Documentation + +Update MCP SDK documentation and the client migration guide to state: + +- `rl-rock[mcp]` depends on `scaffoldhub==0.1.0.dev4`. +- `await env.release()` stops the ROCK sandbox and releases ScaffoldHub auth + leases. +- If auth lease release fails, `release()` raises and callers may call it again + to retry release. + +## Verification + +Run focused tests: + +```bash +uv run pytest tests/unit/sdk/mcp/test_mcp_env.py -v +``` + +Run formatting and linting for touched files: + +```bash +uv run ruff format rock/sdk/mcp/mcp_env.py tests/unit/sdk/mcp/test_mcp_env.py +uv run ruff check rock/sdk/mcp/mcp_env.py tests/unit/sdk/mcp/test_mcp_env.py +``` + +If dependency metadata changes, run the relevant packaging or lockfile checks +according to the implementation plan. From f4650b774f73c728dc251103f7c5ffa89575b40b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:05:18 +0800 Subject: [PATCH 192/226] test(mcp): cover auth provider ownership --- tests/unit/sdk/mcp/test_mcp_env.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index 989330c7f4..a3d120ae00 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -45,16 +45,23 @@ def __init__(self): "SLACK_MCP_XOXB_TOKEN": "xoxb-test-token", } } + self.release_active_leases_calls = 0 def provide(self, platform: str) -> dict: if platform not in self.auth: raise ValueError(f"Unsupported platform: {platform}") return self.auth[platform] + def release_active_leases(self) -> None: + self.release_active_leases_calls += 1 + class FakeDataLifecycleFactory: - def __init__(self): - self.auth_provider = FakeAuthProvider() + last_auth_provider = None + + def __init__(self, auth_provider=None): + self.auth_provider = auth_provider or FakeAuthProvider() + FakeDataLifecycleFactory.last_auth_provider = self.auth_provider self.created = {} def supports(self, lifecycle_type: str) -> bool: @@ -70,11 +77,14 @@ def create(self, lifecycle_type: str): def install_fake_scaffoldhub(monkeypatch): scaffoldhub = ModuleType("scaffoldhub") + auth = ModuleType("scaffoldhub.auth") tools = ModuleType("scaffoldhub.tools") base = ModuleType("scaffoldhub.tools.base") + auth.AuthProvider = FakeAuthProvider base.DataLifecycleFactory = FakeDataLifecycleFactory monkeypatch.setitem(sys.modules, "scaffoldhub", scaffoldhub) + monkeypatch.setitem(sys.modules, "scaffoldhub.auth", auth) monkeypatch.setitem(sys.modules, "scaffoldhub.tools", tools) monkeypatch.setitem(sys.modules, "scaffoldhub.tools.base", base) @@ -151,6 +161,16 @@ def test_mcp_env_constructor_requires_servers_dict(monkeypatch): mcp_env.McpEnv(servers=[]) +def test_mcp_env_owns_auth_provider_and_passes_it_to_lifecycle_factory(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + assert isinstance(env.auth_provider, FakeAuthProvider) + assert FakeDataLifecycleFactory.last_auth_provider is env.auth_provider + assert env.data_lifecycle_factory.auth_provider is env.auth_provider + + def test_mcp_env_constructor_reports_missing_scaffoldhub(monkeypatch): sys.modules.pop("scaffoldhub.tools.base", None) sys.modules.pop("rock.sdk.mcp.mcp_env", None) From 1538965d8a45ea51192f609d40c432a2ce86fe55 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:06:04 +0800 Subject: [PATCH 193/226] feat(mcp): own scaffoldhub auth provider --- rock/sdk/mcp/mcp_env.py | 16 +++++++--------- tests/unit/sdk/mcp/test_mcp_env.py | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index a03764be7d..7f6816b016 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -10,12 +10,13 @@ logger = logging.getLogger(__name__) -def _load_data_lifecycle_factory(): +def _load_scaffoldhub_components(): try: + from scaffoldhub.auth import AuthProvider from scaffoldhub.tools.base import DataLifecycleFactory except ImportError as error: raise ImportError("rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`.") from error - return DataLifecycleFactory + return AuthProvider, DataLifecycleFactory class McpEnv: @@ -47,8 +48,9 @@ def __init__(self, servers: dict | None = None): self.urls = {} self.servers = deepcopy(servers) self.resolved_servers = {} - data_lifecycle_factory = _load_data_lifecycle_factory() - self.data_lifecycle_factory = data_lifecycle_factory() + auth_provider_class, data_lifecycle_factory_class = _load_scaffoldhub_components() + self.auth_provider = auth_provider_class() + self.data_lifecycle_factory = data_lifecycle_factory_class(auth_provider=self.auth_provider) self.data_lifecycles: dict[str, Any] = {} self._rock_runtime = RockRuntime() @@ -188,12 +190,8 @@ def _resolve_server_config(self, server_name: str, server_config: Any) -> Any: return resolved_config def _server_auth(self, server_name: str) -> dict: - auth_provider = getattr(self.data_lifecycle_factory, "auth_provider", None) - if auth_provider is None: - return {} - try: - return auth_provider.provide(server_name) + return self.auth_provider.provide(server_name) except ValueError: return {} diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index a3d120ae00..c0a4640692 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -180,7 +180,7 @@ def test_mcp_env_constructor_reports_missing_scaffoldhub(monkeypatch): def raise_missing_dependency(): raise ImportError("rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`.") - monkeypatch.setattr(module, "_load_data_lifecycle_factory", raise_missing_dependency) + monkeypatch.setattr(module, "_load_scaffoldhub_components", raise_missing_dependency) with pytest.raises(ImportError, match=r"rl-rock\[mcp\]"): module.McpEnv() From 4e3098b79fed3d34cf983d043010d1012a466f63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:06:55 +0800 Subject: [PATCH 194/226] test(mcp): cover auth lease release --- tests/unit/sdk/mcp/test_mcp_env.py | 73 ++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index c0a4640692..1dbc17104e 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -75,6 +75,12 @@ def create(self, lifecycle_type: str): return lifecycle +class FailingReleaseAuthProvider(FakeAuthProvider): + def release_active_leases(self) -> None: + self.release_active_leases_calls += 1 + raise RuntimeError("database release failed") + + def install_fake_scaffoldhub(monkeypatch): scaffoldhub = ModuleType("scaffoldhub") auth = ModuleType("scaffoldhub.auth") @@ -389,6 +395,73 @@ def test_mcp_env_release_preserves_lifecycles_when_runtime_stop_fails(monkeypatc assert lifecycle.reset_calls == 1 +def test_mcp_env_release_releases_auth_leases(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + runtime = RecordingRuntime() + env._rock_runtime = runtime + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + + asyncio.run(env.release()) + + assert runtime.stopped is True + assert env.auth_provider.release_active_leases_calls == 1 + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + + +def test_mcp_env_release_retries_auth_release_when_runtime_is_not_running(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + env.running = False + + asyncio.run(env.release()) + asyncio.run(env.release()) + + assert env.auth_provider.release_active_leases_calls == 2 + + +def test_mcp_env_release_raises_when_auth_release_fails_and_clears_state(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + failing_provider = FailingReleaseAuthProvider() + env.auth_provider = failing_provider + env.data_lifecycle_factory.auth_provider = failing_provider + env._rock_runtime = RecordingRuntime() + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + + with pytest.raises(RuntimeError, match="Failed to release MCP auth leases") as exc_info: + asyncio.run(env.release()) + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "database release failed" + assert failing_provider.release_active_leases_calls == 1 + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + + +def test_mcp_env_release_still_releases_auth_when_runtime_stop_fails(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + env._rock_runtime = FailingStopRuntime() + env.running = True + env.urls = {"slack": "https://example.test/slack/sse"} + env.resolved_servers = {"slack": slack_server_config()} + + asyncio.run(env.release()) + + assert env.auth_provider.release_active_leases_calls == 1 + assert env.is_alive() is False + assert env.urls == {} + assert env.resolved_servers == {} + + def test_mcp_env_start_accepts_before_launch_hook(monkeypatch): mcp_env = reload_mcp_env(monkeypatch) signature = inspect.signature(mcp_env.McpEnv.start) From 46b5b32149710b0a57fde64114f1b51a855f1efb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:07:41 +0800 Subject: [PATCH 195/226] feat(mcp): release scaffoldhub auth leases --- rock/sdk/mcp/mcp_env.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index 7f6816b016..cacb673823 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -159,21 +159,30 @@ def dump(self) -> dict: async def release(self): """ - Release the physical MCP runtime. + Release the physical MCP runtime and any active ScaffoldHub auth leases. Data cleanup is intentionally handled by explicit reset calls. """ + auth_release_error: Exception | None = None try: if self.running: try: await self._rock_runtime.stop() except Exception as error: logger.warning("Failed to stop ROCK runtime during release: %s", error) + + try: + self.auth_provider.release_active_leases() + except Exception as error: + auth_release_error = error finally: self.running = False self.urls = {} self.resolved_servers = {} + if auth_release_error is not None: + raise RuntimeError("Failed to release MCP auth leases") from auth_release_error + def _resolve_server_config(self, server_name: str, server_config: Any) -> Any: if not isinstance(server_config, dict): return deepcopy(server_config) From 0fee8d15331ac54ecfb810e4bea01425d7e87fe2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:08:28 +0800 Subject: [PATCH 196/226] build(mcp): upgrade scaffoldhub dependency --- pyproject.toml | 2 +- uv.lock | 200 ++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 198 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1057bf286d..0dceb11e7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev3; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev4; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/uv.lock b/uv.lock index eb6e373dfc..f08689bfe8 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 1 requires-python = ">=3.10, <4.0" resolution-markers = [ "python_full_version >= '3.13' and sys_platform != 'win32'", @@ -1828,6 +1828,15 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad" }, ] +[[package]] +name = "httpx-sse" +version = "0.4.3" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc" }, +] + [[package]] name = "huggingface-hub" version = "1.0.1" @@ -2454,6 +2463,37 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/9a/cc/3fe688ff1355010937713164caacf9ed443675ac48a997bab6ed23b3f7c0/matplotlib-3.10.7-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3886e47f64611046bc1db523a09dd0a0a6bed6081e6f90e13806dd1d1d1b5e91" }, ] +[[package]] +name = "mcp" +version = "1.27.2" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "anyio", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "httpx", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "httpx-sse", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "jsonschema", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "pydantic", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "pydantic-settings", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "pyjwt", extra = ["crypto"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "python-multipart", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "pywin32", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and sys_platform == 'win32'" }, + { name = "sse-starlette", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "starlette", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typing-inspection", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "uvicorn", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and sys_platform != 'emscripten'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/27/3c/347cf965d313f5d41764e7d46bea6ffe7d9ef13b983cc429b0340962a082/mcp-1.27.2.tar.gz", hash = "sha256:8e02db104096d1c25b28e64bde29a5c32b31bc241710213e12fd4d84985bdfef" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/c9/11/252c6f971dc4f16af1d98a1c469d8ba523aab00d1bb76b4d3bc1ff32eacc/mcp-1.27.2-py3-none-any.whl", hash = "sha256:d6ff5160c6ca65d93013626efb3fc249de683c30b2d8570755ceddd490344de5" }, +] + +[package.optional-dependencies] +cli = [ + { name = "python-dotenv", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typer", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -3473,6 +3513,86 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/c9/ad/33b2ccec09bf96c2b2ef3f9a6f66baac8253d7565d8839e024a6b905d45d/psutil-7.1.3-cp37-abi3-win_arm64.whl", hash = "sha256:bd0d69cee829226a761e92f28140bec9a5ee9d5b4fb4b0cc589068dbfff559b1" }, ] +[[package]] +name = "psycopg" +version = "3.3.4" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "tzdata", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and sys_platform == 'win32'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/db/2f/cb91e5502ec9de1de6f1b76cfbf69531932725361168bb06963620c77e2e/psycopg-3.3.4.tar.gz", hash = "sha256:e21207764952cff81b6b8bdacad9a3939f2793367fdac2987b3aac36a651b5bc" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/5c/e0/7b3dee031daae7743609ce3c746565d4a3ed7c2c186479eb48e34e838c64/psycopg-3.3.4-py3-none-any.whl", hash = "sha256:b6bbc25ccf05c8fad3b061d9db2ef0909a555171b84b07f29458a447253d679a" }, +] + +[package.optional-dependencies] +binary = [ + { name = "psycopg-binary", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and implementation_name != 'pypy'" }, +] + +[[package]] +name = "psycopg-binary" +version = "3.3.4" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/b7/bf/70d8a60488f9955cbbcd538beae44d56bb2f1d19e673b72788f2d343ff55/psycopg_binary-3.3.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b7bfff1ca23732b488cbca3076fc11bc98d520ee122514fdb17a8e20d3338f5a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/db/b0/29e98ba210c9dbc75a6dc91e3f99b9e06ea901a62ca95804e02a1ae13e6b/psycopg_binary-3.3.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:32a6fbf8481e3a370d0d72b860d35948a693cb01281da217f7b2f307636e591a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/8e/ab/3df087b3c12bf74e47c08204172b2fabb5a144679110d5c7ad12d9201323/psycopg_binary-3.3.4-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:bdef84570ebbce1d42b4e7ea952d21c414c5f118ad02fee00c5625f35e134429" }, + { url = "https://mirrors.aliyun.com/pypi/packages/87/9a/f088207b4cd6772f9e0d8a91807e79fa2458d4eb9eb1ae406c68415f2bec/psycopg_binary-3.3.4-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa1cbc10768a796c96d3243656016bf4e337c81c71097270bb7b0ad6210d9765" }, + { url = "https://mirrors.aliyun.com/pypi/packages/48/45/4523a857f253871d75c22e1c2e79fd47e599e736bcba1bad58d83e24be02/psycopg_binary-3.3.4-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cf7f73a4a792bc5db58a4b385d8a1467e8d468f7548702fb0ed1e9b7501b1c13" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7c/d1/925bf776503345bef428e6c45fb017d0139ddbe0e211814b585c4253dca8/psycopg_binary-3.3.4-cp310-cp310-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d7b4d40c153fa352ab3cca530f3a0baedf7621b2ebcbd7f084009522c21788fc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/6f/aa/99727337206fbba357ca084bf4ea8b29dc986f61842a2685859af61416db/psycopg_binary-3.3.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f9b1c2533af01cd7648378599f82b0b8ae32f293296e6eec5753a625bc97ef28" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0b/a4/567ba2c37d19d8c2f63d836385dfd2495aa5897bbee6cfab104d9ee58624/psycopg_binary-3.3.4-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:ad3bc94054876155549fdaedf4a46d1ec69d39a5bcee377148afe498e84c4b8e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b7/23/86457f5a82731685d7701de7bfaa5eb783dd1fecbf875321897d9d9ce33a/psycopg_binary-3.3.4-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:eb4eed2079c01a4850bf467deacfab56d356d4225040170af03dc9958321242d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a7/d8/249456df16d47de082abd9b73bce8ccdeb0293eb12e590f9150c7cbdb788/psycopg_binary-3.3.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f80e3f2b5331dbbf0901bcb658056c03eeb2c1ef31d774afb0d61598b242e744" }, + { url = "https://mirrors.aliyun.com/pypi/packages/15/6b/c4abe228acafd8a385c1fb615d4f1e3c9b8ad7a4e4f0e84118ba3ffeed9c/psycopg_binary-3.3.4-cp310-cp310-win_amd64.whl", hash = "sha256:574ea21a9651958f1535c5a1c649c7409e9168bcbffa29a3f2f961f58b322949" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b6/82/df3312c0ca083d5b43b352f27d4dd8b1e614bd334473074715d9e0000da4/psycopg_binary-3.3.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:612a627d733f695b1de1f9b4bd511c15f999a5d8b915d444bbd7dd71cf3370da" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1f/b5/d74d542458d3e8ac0571d8a88f57ca369999b9a82f4fa528052d0d7d3e4c/psycopg_binary-3.3.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:13a7f380824c35896dcac7fe0f61440f7ca49d6dc73f3c13a9a4471e6a3b302e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/09/67/06bab9c60671999f4c6ceff1b334f3ac1f9fc5789eb467c714623ea21de9/psycopg_binary-3.3.4-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:276904e3452d6a23d474ef9a21eee19f20eed3d53ddd2576af033827e0ba0992" }, + { url = "https://mirrors.aliyun.com/pypi/packages/72/9b/023433e2b20f970de1e22d29132a95281277646da0b2e2879dd4ee94b8c1/psycopg_binary-3.3.4-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ab8cca8ef8fb1ccf5b048ae5bd78ba55b9e4b5d472e3ce5ca39ff4d2a9c249e4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/08/cd/ae16da8fde228a38b2fe9269bbc13cf89e0186173f2265600f02d6a71e64/psycopg_binary-3.3.4-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7465bfe6087d2d5b42d4c53b9b11ca9f218e477317a4a162a10e3c19e984ba8e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4f/81/0ba09fa5f5f88779093a2541a8e02489825721f258ab88058b11d68b3eb5/psycopg_binary-3.3.4-cp311-cp311-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:22cdbf5f91ef7bb91fe0c5757e1962d3127a8010256eefd9c61fcaf441802097" }, + { url = "https://mirrors.aliyun.com/pypi/packages/73/6a/629136040cc3497adb442a305710b5913f2a754d4630fc3d3717c4c0df65/psycopg_binary-3.3.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e2631da29253a98bd496e6c4813b24e09a4fe3fb2a9e88513305d6f8747cce95" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7c/32/1027f843c6dc2d5d51960ee62cc0c2cf755a4c39455aff1371173edbef7d/psycopg_binary-3.3.4-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:7f7668f30b9dd5163197e5cbf4e0efd54e00f0a859cc566ce56cfc31f4054839" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0b/e1/380a724d9093c74adb14d4fce920ea8327838abb61f760b1448586b14a8e/psycopg_binary-3.3.4-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:cffc3408d77a27973f33e5d909b624cce683db5fc25964b02fe0aae7886c1007" }, + { url = "https://mirrors.aliyun.com/pypi/packages/db/cd/895893ae575a09c97ccfd5def070d88993d955ef34df45a881fd5ff506d6/psycopg_binary-3.3.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0579252a1202cd73e4da137a1426e2dae993ae44e757605344282af3a082848c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/dd/c6/2330a20794e37a3ec609ef2fd8522919ec7a4395a1abf979a8e2d1775cd5/psycopg_binary-3.3.4-cp311-cp311-win_amd64.whl", hash = "sha256:41f2ec0fea529832982bcb6c9415de3c86264ebe562b77a467c0fbcd7efbba8d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/95/7d/03818e13ba7f36de93573c93ee3482006d3dfa8b0f8d28df511bad0a1a92/psycopg_binary-3.3.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5ab28a2a7649df3b72e6b674b4c190e448e8e77cf496a65bd846472048de2089" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a5/b9/11b341edf8d54e2694726b273fe9652b254d989f4f63e3ac6816ad6b55f4/psycopg_binary-3.3.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6402a9d8146cf4b3974ded3fd28a971e83dc6a0333eb7822524a3aa20b546578" }, + { url = "https://mirrors.aliyun.com/pypi/packages/8b/18/4665bacd65e7865b4372fcd8abb8b9186ada4b0025f8c2ca691b364a556c/psycopg_binary-3.3.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:580ae30a5f95ccd90008ec697d3ed6a4a2047a516407ad904283fa42086936e9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7c/b1/b83136c6e510593d9b0c759ba5384337bc4ad82d19fda675adc4b2703c84/psycopg_binary-3.3.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e7510c37550f91a187e3660a8cc50d4b760f8c3b8b2f89ebc5698cd2c7f2c85d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/67/8d/a9821e2a648afe6091989929982a3b0f00b2631a859cb81379728f08fb75/psycopg_binary-3.3.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77df19583501ea288eaf15ac0fe7ad01e6d8091a91d5c41df5c718f307d8e31b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7e/58/2e349e8d23905dc2317b80ac65f48fb6f821a4777a4e994a60da91c4850f/psycopg_binary-3.3.4-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:018fbed325936da502feb546642c982dcc4b9ffdea32dfef78dbf3b7f7ad4070" }, + { url = "https://mirrors.aliyun.com/pypi/packages/45/48/57b00d03b4721878326122a1f1e6b0a90b85bcaec56b5b2f8ea6cfa45235/psycopg_binary-3.3.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:17a21953a9e5ff3a16dab692625a3676e2f101db5e40072f39dbee2250194d68" }, + { url = "https://mirrors.aliyun.com/pypi/packages/25/37/33b47d8c007df69aec500df5889767c4d313748e8e9e27a2fef8a6dabcee/psycopg_binary-3.3.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:eb05ee1c2b817d27c537333224c9e83c7afb86fe7296ba970990068baf819b16" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ca/c6/32b0835dbc2122617902b649d76a91c1e75406e76bf3d595b0c3bb5ffad6/psycopg_binary-3.3.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:773d573e11f437ce0bdb95b7c18dc58390494f96d43f8b45b9760436114f7652" }, + { url = "https://mirrors.aliyun.com/pypi/packages/cd/68/d190ef0c0c5b16ded07831dabc8ddd412f4cdab07ec6e30ed38d9bda0e1f/psycopg_binary-3.3.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:71e55ccbdfae79a2ed9c6369c3008a3025817ff9d7e27b32a2d84e2a4267e66e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/25/8f/81dcbc2e8454b74d14881275ea45f00791052dac531a9fa8be1730d1685b/psycopg_binary-3.3.4-cp312-cp312-win_amd64.whl", hash = "sha256:494ca54901be8cf9eb7e02c25b731f2317c378efa44f43e8f9bd0e1184ae7be4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/09/43/13e9c406fbbf354580476e248a16b64802a376873ebe6339e30bb655572d/psycopg_binary-3.3.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fbd1d4ed566895ad2d3bf4ddfd8bae90026930ddf29df3b9d91d32c8c47866a7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/22/be/2923cd7c3683e7afdecf4f10796a18de02f5c5ddc0969aa2ad0a8cdd3bbd/psycopg_binary-3.3.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:75a9067e236f9b9ae3535b66fe99bddb33d39c0de10112e49b9ab11eee53dc31" }, + { url = "https://mirrors.aliyun.com/pypi/packages/96/a0/2c913d6fe13d6a8bd13597d36739bf47af063ad9399e402cfecab16f3c1e/psycopg_binary-3.3.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:b56b603ebcea8aa10b46228b8410ba7f13e7c2ee54389d4d9be0927fd8ce2a70" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e7/38/205d10bc1ad0df4a21c5c51659126bd3ea0ef98fcad1e852f78c249bb9c3/psycopg_binary-3.3.4-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c677c4ad433cb7150c8cd304a0769ae3bcfbe5ea0676eb53faa7b1443b16d0d3" }, + { url = "https://mirrors.aliyun.com/pypi/packages/36/fc/f0381ddcd45eff3bb70dbca6823a996048d7f507b2ec3fc92c6fabc0fe87/psycopg_binary-3.3.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:26df2717e59c0473e4465a97dfb1b7afebaa479277870fd5784d1436470db47c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/95/40/fa545ae152c24327651e5624e4902121e808270be36c10b12e9939be09bc/psycopg_binary-3.3.4-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1dc1f79fd16bb1f3f4421417a514607539f17804d95c7ed617265369d1981cae" }, + { url = "https://mirrors.aliyun.com/pypi/packages/86/e4/2f8a47ee97f90cd2b933d0463081d35631ff419de2b8c984a5f369857de0/psycopg_binary-3.3.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:136f199a407b5348b9b857c504aff60c77622a28482e7195839ce1b51238c4cc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0e/0e/94e842ff4a7f98ed162580ca2e8b8864b28c1e0350f2443f8ee47f821167/psycopg_binary-3.3.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:b6f5a29e9c775b9f12a1a717aa7a2c80f9e1db6f27ba44a5b59c80ac61d2ffcf" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d0/83/fc6c174b672e29b7de996ea77b6cbddf46c891751c3355f6974292baa6b4/psycopg_binary-3.3.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:ee17a2cf4943cde261adfad1bbc5bf38d6b3776d7afff74c7cabcbeaeb08c260" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e9/65/768364d4a97a15b1a7f47ba52688c1686f22941d8332a8398cefc468e25f/psycopg_binary-3.3.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5c4ab71be17bdca30cb34c34c4e1496e2f5d6f20c199c12bad226070b22ef9bf" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bd/3b/218efbc9e645becd80cdf651acda05f85cfe546b7a9c0458c7cbc8fe1f74/psycopg_binary-3.3.4-cp313-cp313-win_amd64.whl", hash = "sha256:dbfdb9b6cc79f31104a7b162a2b921b765fcc62af6c00540a167a8de47e4ed38" }, + { url = "https://mirrors.aliyun.com/pypi/packages/48/a6/828c9185701dab71b234c2a76c38a08b098ebfec5020716b4e93807492b5/psycopg_binary-3.3.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:28b7398fdd19db3232c884fb24550bdfe951221f510e195e233299e4c9b78f97" }, + { url = "https://mirrors.aliyun.com/pypi/packages/92/58/5b40dbc9d839045c9dae956960e4fb6d20bcabe6c59a2aa34fc3a371913f/psycopg_binary-3.3.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1fbaa292a3c8bb61b45df1ad3da1908ccee7cb889db9425e3557d9e34e2a4829" }, + { url = "https://mirrors.aliyun.com/pypi/packages/85/a9/793f0ac107a9003b48441d0d1f9f616d96e0f37458dd8dc12528ceff55fb/psycopg_binary-3.3.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:94596f9e7633ee3f6440711d43bb70aa31cc0a46a900ab8b4201a366ace5c9e7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/8f/26/42e8533497e2592334f68ec529cf5f840f7fa4e99575a4bb61aa184dbfbf/psycopg_binary-3.3.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8c0056529e68dbe9184cd4019a1f3d8f3a4ead2f6fc7a5afcf27d3314edd1277" }, + { url = "https://mirrors.aliyun.com/pypi/packages/15/af/b7151776cc08d5935d45c833ec818a9beb417cf7c08239af1aafbdae78ee/psycopg_binary-3.3.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c09aad7051326e7603c14e50636db9c01f78272dc54b3accff03d46370461e6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d0/ed/c92533b9124712d592cbf1cd6c76da933a2e0acea81dfe1fbe7e735f0cff/psycopg_binary-3.3.4-cp314-cp314-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:514404ed543efd620c85602b747df2a23cf1241b4067199e1a66f2d2757aaa41" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a2/23/ccadfd0de416aa188356daa199453af24087b042e296088706d190ae0295/psycopg_binary-3.3.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:46893c26858be12cc49ca4226ed6a60b4bfccadd946b3bebb783a60b38788228" }, + { url = "https://mirrors.aliyun.com/pypi/packages/fd/a0/c8f43cee36386f7bc891ab41a9d31ea07cf9826038e732da79f26b1e5f34/psycopg_binary-3.3.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:df1d567fc430f6df15c9fcf67d87685fc49bdb325adc0db5af1adfb2f44eb5c9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4e/2c/c1547871be3790676e8868b38655496422f94f0978dfb66b74bdba2f1676/psycopg_binary-3.3.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:6b9016b1714da4dd5ecaaa75b82098aa5a0b87854ce9b092e21c27c4ae23e014" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c4/b1/f6670f00fa7ea601584623f6c11602ab92117d83eaff885e0210f6de7418/psycopg_binary-3.3.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:47c656a8a7ba6eb0cff1801a4caaa9c8bdc12d03080e273aff1c8ac39971a77e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/eb/e6/5fff07a70d1f945ed90ae131c3bd76cab32beff7c58c6db15ad5820b6d1f/psycopg_binary-3.3.4-cp314-cp314-win_amd64.whl", hash = "sha256:c37e024c07308cd06cf3ec51bfd0e7f6157585a4d84d1bce4a7f5f7913719bf8" }, +] + [[package]] name = "ptyprocess" version = "0.7.0" @@ -3727,6 +3847,20 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/32/56/8a7ca5d2cd2cda1d245d34b1c9a942920a718082ae8e54e5f3e5a58b7add/pydantic_core-2.33.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:329467cecfb529c925cf2bbd4d60d2c509bc2fb52a20c1045bf09bb70971a9c1" }, ] +[[package]] +name = "pydantic-settings" +version = "2.14.1" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "pydantic", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "python-dotenv", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typing-inspection", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/07/60/1d1e59c9c90d54591469ada7d268251f71c24bdb765f1a8a832cee8c6653/pydantic_settings-2.14.1.tar.gz", hash = "sha256:e874d3bec7e787b0c9958277956ed9b4dd5de6a80e162188fdaff7c5e26fd5fa" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/ae/8d/f1af3832f5e6eb13ba94ee809e72b8ecb5eef226d27ee0bef7d963d943c7/pydantic_settings-2.14.1-py3-none-any.whl", hash = "sha256:6e3c7edfd8277687cdc598f56e5cff0e9bfff0910a3749deaa8d4401c3a2b9de" }, +] + [[package]] name = "pyfiglet" version = "1.0.2" @@ -3745,6 +3879,20 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b" }, ] +[[package]] +name = "pyjwt" +version = "2.13.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728" }, +] + +[package.optional-dependencies] +crypto = [ + { name = "cryptography", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] + [[package]] name = "pyparsing" version = "3.2.5" @@ -4280,7 +4428,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.8.0" +version = "1.10.0.dev2" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -4362,6 +4510,9 @@ builder = [ { name = "gem-llm" }, { name = "swebench" }, ] +mcp = [ + { name = "scaffoldhub", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] model-service = [ { name = "alibabacloud-cr20181201" }, { name = "fastapi" }, @@ -4404,6 +4555,7 @@ test = [ { name = "pytest-twisted" }, { name = "pytest-xdist" }, { name = "ruff" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, ] [package.metadata] @@ -4458,6 +4610,7 @@ requires-dist = [ { name = "rl-rock", extras = ["builder"], marker = "extra == 'all'" }, { name = "rl-rock", extras = ["rocklet"], marker = "extra == 'admin'" }, { name = "rl-rock", extras = ["rocklet"], marker = "extra == 'all'" }, + { name = "scaffoldhub", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'mcp'", specifier = "==0.1.0.dev4", index = "https://artlab.alibaba-inc.com/1/pypi/simple" }, { name = "sqlmodel", marker = "extra == 'admin'" }, { name = "sqlmodel", marker = "extra == 'sandbox-actor'" }, { name = "swebench", marker = "extra == 'builder'" }, @@ -4469,7 +4622,7 @@ requires-dist = [ { name = "uvicorn", marker = "extra == 'rocklet'" }, { name = "websockets", marker = "extra == 'admin'", specifier = ">=15.0.1" }, ] -provides-extras = ["admin", "rocklet", "sandbox-actor", "builder", "model-service", "all"] +provides-extras = ["admin", "rocklet", "sandbox-actor", "builder", "model-service", "mcp", "all"] [package.metadata.requires-dev] test = [ @@ -4486,6 +4639,7 @@ test = [ { name = "pytest-twisted" }, { name = "pytest-xdist", specifier = ">=3.8.0" }, { name = "ruff", specifier = ">=0.14.0" }, + { name = "tomli", marker = "python_full_version < '3.11'", specifier = ">=2.0.0" }, ] [[package]] @@ -4660,6 +4814,24 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/48/f0/ae7ca09223a81a1d890b2557186ea015f6e0502e9b8cb8e1813f1d8cfa4e/s3transfer-0.14.0-py3-none-any.whl", hash = "sha256:ea3b790c7077558ed1f02a3072fb3cb992bbbd253392f4b6e9e8976941c7d456" }, ] +[[package]] +name = "scaffoldhub" +version = "0.1.0.dev4" +source = { registry = "https://artlab.alibaba-inc.com/1/pypi/simple" } +dependencies = [ + { name = "httpx", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "mcp", extra = ["cli"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "psycopg", extra = ["binary"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "python-dotenv", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "requests", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "rl-rock", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "slack-sdk", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +sdist = { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev4.tar.gz", hash = "sha256:1b9d9b712b7f369b73d4650a37a1a972d454d2cde24c38b8c16acbcc5287f92a" } +wheels = [ + { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev4-py3-none-any.whl", hash = "sha256:36d0e8d4026d6b12850d7c4d8c7c689af1fff644f395e58625cfea54fe515207" }, +] + [[package]] name = "setuptools" version = "80.9.0" @@ -4699,6 +4871,15 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274" }, ] +[[package]] +name = "slack-sdk" +version = "3.42.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/0e/00/16258bfa547559b2c936b50c882b4f0a36ebf6b69639eb763d8fa5e8d6cb/slack_sdk-3.42.0.tar.gz", hash = "sha256:873db9e1f632ac650ffdbf9d8ba825f3e9e7e576a1e4f9604ccb2a15b3727e3d" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/ce/ef/8a1556bd4843443993fc116783790a7cc553601a37f7d965ec26eef95e76/slack_sdk-3.42.0-py2.py3-none-any.whl", hash = "sha256:eb39aff97e476e10cc5a8ac29bd2e79a9959e880d9fe0c03b4e8f05b2ac996ff" }, +] + [[package]] name = "smart-open" version = "7.4.4" @@ -4805,6 +4986,19 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/8c/92/c35e036151fe53822893979f8a13e6f235ae8191f4164a79ae60a95d66aa/sqlmodel-0.0.27-py3-none-any.whl", hash = "sha256:667fe10aa8ff5438134668228dc7d7a08306f4c5c4c7e6ad3ad68defa0e7aa49" }, ] +[[package]] +name = "sse-starlette" +version = "3.4.4" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "anyio", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "starlette", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/f7/2b/58abc2d1fd397e7dde08e947e05c884d8ef2f78d5e2588c17a12d42d6994/sse_starlette-3.4.4.tar.gz", hash = "sha256:07e0fa0460138baf25cdd5fb28683472c3995dc1642225191b3832d62526bcb0" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/dc/67/805710444ea8cc75fbf70b920ed431a560c4bf9c57f7d5a3117213189399/sse_starlette-3.4.4-py3-none-any.whl", hash = "sha256:3f4dd50d8aed2771a091f3a83000323fc3844541c16b4fe585ae2420cc6df973" }, +] + [[package]] name = "starlette" version = "0.49.3" From 9e49cd50d93c9248f16c45540365ff7e451fc766 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:09:09 +0800 Subject: [PATCH 197/226] docs(mcp): describe auth lease release --- .../References/Python SDK References/mcp.md | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index efd49f7fd0..e7c0a3e015 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -13,7 +13,7 @@ pip install --extra-index-url https://artlab.alibaba-inc.com/1/pypi/simple "rl-r The MCP extra currently supports Python 3.11 and 3.12 because the published ScaffoldHub package is Python 3.11+ and ROCK officially supports Python 3.10 to -3.12. +3.12. The MCP extra depends on `scaffoldhub==0.1.0.dev4`. ## Basic Usage @@ -45,6 +45,11 @@ async def main(): asyncio.run(main()) ``` +`release()` is the cleanup boundary for both ROCK runtime resources and +ScaffoldHub auth leases. Keep it in a `finally` block. If it raises because auth +lease release failed, call it again after handling or logging the error to retry +lease release. + ## Lifecycle Data `McpEnv.init(data)`, `McpEnv.reset()`, and `McpEnv.dump()` delegate tool data @@ -58,7 +63,10 @@ env.reset() ``` `reset()` only resets configured tool data. It does not stop the ROCK sandbox. -Use `await env.release()` to stop the sandbox runtime. +Use `await env.release()` to stop the sandbox runtime and release any +ScaffoldHub auth leases borrowed while resolving server credentials. +If auth lease release fails, `release()` raises an exception and can be called +again to retry lease release. ## Launch Hook From ed363d5140edba3f2c1fc005e3b8e6a7c22341cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:22:23 +0800 Subject: [PATCH 198/226] docs(mcp): clarify release retry guidance --- .../References/Python SDK References/mcp.md | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index e7c0a3e015..52ea21e244 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -38,17 +38,17 @@ async def main(): urls = env.get_urls() print(urls["calculator"]) finally: - if env.is_alive(): - await env.release() + await env.release() asyncio.run(main()) ``` `release()` is the cleanup boundary for both ROCK runtime resources and -ScaffoldHub auth leases. Keep it in a `finally` block. If it raises because auth -lease release failed, call it again after handling or logging the error to retry -lease release. +ScaffoldHub auth leases. Keep it in a `finally` block and call it even if +`start()` fails before `env.is_alive()` becomes true. If auth lease release +fails, `release()` raises `RuntimeError("Failed to release MCP auth leases")`; +call it again after handling or logging the error to retry lease release. ## Lifecycle Data @@ -65,8 +65,10 @@ env.reset() `reset()` only resets configured tool data. It does not stop the ROCK sandbox. Use `await env.release()` to stop the sandbox runtime and release any ScaffoldHub auth leases borrowed while resolving server credentials. -If auth lease release fails, `release()` raises an exception and can be called -again to retry lease release. +Call it even when `env.is_alive()` is false, because server credential +resolution can borrow auth before the runtime is marked alive. If auth lease +release fails, `release()` raises `RuntimeError("Failed to release MCP auth leases")` +and can be called again to retry lease release. ## Launch Hook From c71a8136d0eb5d84f06f7d5dac2df8acd2f0fa1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 15 Jun 2026 16:35:47 +0800 Subject: [PATCH 199/226] chore: bump version to 1.10.0.dev3 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0dceb11e7a..f4c9f7f5e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev2" +version = "1.10.0.dev3" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/uv.lock b/uv.lock index f08689bfe8..7e2d59c231 100644 --- a/uv.lock +++ b/uv.lock @@ -4428,7 +4428,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.10.0.dev2" +version = "1.10.0.dev3" source = { editable = "." } dependencies = [ { name = "anyio" }, From feec0a820fbf5a8476bb1cae9aa2b9beca4d426e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Tue, 16 Jun 2026 18:13:13 +0800 Subject: [PATCH 200/226] chore: bump mcp dev versions --- pyproject.toml | 4 ++-- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f4c9f7f5e5..218ac87726 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev3" +version = "1.10.0.dev4" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev4; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev5; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index 6bb81b597f..9b8246028e 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev3; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev5; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From 6eb70f2dbc391a4e447d77f971cac144c4e886fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Wed, 17 Jun 2026 18:05:14 +0800 Subject: [PATCH 201/226] chore: bump mcp package versions --- pyproject.toml | 4 ++-- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 218ac87726..2780e86e5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev4" +version = "1.10.0.dev5" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev5; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev6; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index 9b8246028e..1a07d70e0d 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev5; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev6; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From d946c440b5cea532b63d93466322a9af6c37fc65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Thu, 18 Jun 2026 17:02:16 +0800 Subject: [PATCH 202/226] fix: update MCP runtime defaults --- rock/sdk/mcp/rock_runtime.py | 4 ++-- tests/unit/sdk/mcp/test_rock_runtime.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index f55a55fa2d..378b2539f9 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -65,11 +65,11 @@ def from_env(cls) -> RockRuntimeConfig: api_key=api_key, image=os.getenv( "ROCK_SANDBOX_IMAGE", - "rock-registry.cn-hangzhou.cr.aliyuncs.com/envs/mcp-atlas-local:v0.4.0", + "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.10.1", ), user_id=user_id, experiment_id=os.getenv("ROCK_EXPERIMENT_ID", "mcpenv"), - cluster=os.getenv("ROCK_CLUSTER", "nt-a"), + cluster=os.getenv("ROCK_CLUSTER", "vpc-nt-a"), cpus=cpus, memory=os.getenv("ROCK_SANDBOX_MEMORY", "8g"), auto_clear_seconds=auto_clear_seconds, diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index d412909c14..3a24f7df54 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -41,9 +41,9 @@ def test_rock_runtime_config_reads_defaults_and_numeric_values(monkeypatch): assert config.api_key == "rock-key" assert config.user_id == "user-001" assert config.base_url == "https://xrl.alibaba-inc.com" - assert config.image == "rock-registry.cn-hangzhou.cr.aliyuncs.com/envs/mcp-atlas-local:v0.4.0" + assert config.image == "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.10.1" assert config.experiment_id == "mcpenv" - assert config.cluster == "nt-a" + assert config.cluster == "vpc-nt-a" assert config.cpus == 4.0 assert config.memory == "8g" assert config.auto_clear_seconds == 3600 From 27d80dc10e5e9a9a56837b750506ec0a7af01188 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Thu, 2 Jul 2026 15:10:34 +0800 Subject: [PATCH 203/226] chore: bump mcp dev package versions --- pyproject.toml | 4 ++-- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2780e86e5e..75c08dfd24 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev5" +version = "1.10.0.dev9" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev6; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.0.dev11; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index 1a07d70e0d..5504d65bc9 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev6; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.0.dev11; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From 05dcf08f95e7ce5e954a5d4e6a37e208668df972 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Thu, 2 Jul 2026 20:45:07 +0800 Subject: [PATCH 204/226] chore: bump mcp base image version --- pyproject.toml | 2 +- rock/sdk/mcp/rock_runtime.py | 2 +- tests/unit/sdk/mcp/test_rock_runtime.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 75c08dfd24..5e6aef887b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev9" +version = "1.10.0.dev10" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index 378b2539f9..bbb4845209 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -65,7 +65,7 @@ def from_env(cls) -> RockRuntimeConfig: api_key=api_key, image=os.getenv( "ROCK_SANDBOX_IMAGE", - "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.10.1", + "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.12.0", ), user_id=user_id, experiment_id=os.getenv("ROCK_EXPERIMENT_ID", "mcpenv"), diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index 3a24f7df54..25300ec35d 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -41,7 +41,7 @@ def test_rock_runtime_config_reads_defaults_and_numeric_values(monkeypatch): assert config.api_key == "rock-key" assert config.user_id == "user-001" assert config.base_url == "https://xrl.alibaba-inc.com" - assert config.image == "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.10.1" + assert config.image == "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.12.0" assert config.experiment_id == "mcpenv" assert config.cluster == "vpc-nt-a" assert config.cpus == 4.0 From ad773267964c0aa4b7d71fda02b6b40a8d2405a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 14:31:22 +0800 Subject: [PATCH 205/226] chore: bump mcp base image to v0.13.0 --- pyproject.toml | 2 +- rock/sdk/mcp/rock_runtime.py | 2 +- tests/unit/sdk/mcp/test_rock_runtime.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5e6aef887b..78c225f2b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev10" +version = "1.10.0.dev11" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index bbb4845209..27363e0c8d 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -65,7 +65,7 @@ def from_env(cls) -> RockRuntimeConfig: api_key=api_key, image=os.getenv( "ROCK_SANDBOX_IMAGE", - "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.12.0", + "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.13.0", ), user_id=user_id, experiment_id=os.getenv("ROCK_EXPERIMENT_ID", "mcpenv"), diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index 25300ec35d..264d33dc27 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -41,7 +41,7 @@ def test_rock_runtime_config_reads_defaults_and_numeric_values(monkeypatch): assert config.api_key == "rock-key" assert config.user_id == "user-001" assert config.base_url == "https://xrl.alibaba-inc.com" - assert config.image == "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.12.0" + assert config.image == "rock-instances-registry-vpc.cn-shanghai.cr.aliyuncs.com/instance/rock-mcp-base:v0.13.0" assert config.experiment_id == "mcpenv" assert config.cluster == "vpc-nt-a" assert config.cpus == 4.0 From 0badc7b8b20abb12a5f4748a1988d3690d000e13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:12:08 +0800 Subject: [PATCH 206/226] docs: add McpEnv SandboxAware design --- ...2026-07-03-mcp-env-sandbox-aware-design.md | 276 ++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md diff --git a/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md b/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md new file mode 100644 index 0000000000..e656808973 --- /dev/null +++ b/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md @@ -0,0 +1,276 @@ +# McpEnv SandboxAware Lifecycle Injection Design + +## Context + +ROCK owns `rock.sdk.mcp.McpEnv` and `RockRuntime`. `McpEnv` creates +ScaffoldHub data lifecycles through `DataLifecycleFactory`, while +`RockRuntime` creates the ROCK sandbox, writes `/app/mcp-servers.json`, starts +MCP server processes, health-checks SSE endpoints, and stops the sandbox. + +ScaffoldHub now exposes `SandboxAware` from `scaffoldhub.tools.base`. +`SandboxAware` is a narrow lifecycle marker whose public contract is only: + +```python +def set_sandbox(self, sandbox) -> None: + ... +``` + +ScaffoldHub's design states that external launchers, including ROCK +`McpEnv`, should inject the externally created ROCK sandbox into lifecycles +that implement `SandboxAware`. The marker does not ask the launcher to call +lifecycle `before_launch()` methods. Those methods remain lifecycle-internal +compatibility helpers and are not part of the public integration contract. + +## Goals + +- Automatically inject the ROCK sandbox into ScaffoldHub lifecycles that + implement `SandboxAware`. +- Run injection after the sandbox is started and `/app/mcp-servers.json` has + been written. +- Run injection before the user-provided `before_launch(sandbox)` callback. +- Preserve existing `McpEnv.start(before_launch=None)` API shape. +- Keep `RockRuntime` independent from ScaffoldHub data lifecycle concepts. +- Preserve current startup cleanup behavior when injection fails. + +## Non-Goals + +- Do not call lifecycle `before_launch()` automatically. +- Do not add `prepare`, `before_launch`, or sandbox setup methods to + `SandboxAware`. +- Do not add lifecycle parameters to `RockRuntime.start()`. +- Do not move ScaffoldHub lifecycle handling into `RockRuntime`. +- Do not change `init`, `dump`, `reset`, `release`, URL generation, auth + placeholder resolution, or auth lease release semantics. + +## Current Flow + +Today `McpEnv.start()` resolves server configs and passes the caller's +`before_launch` callback directly into `RockRuntime.start()`. + +```text +McpEnv.start(before_launch=user_hook) + resolve server configs + RockRuntime.start(resolved_servers, before_launch=user_hook) + Sandbox.start() + prepare /app/workspace and /data + write /app/mcp-servers.json + user_hook(sandbox) + launch /app/launch.sh + health-check SSE endpoints +``` + +This gives callers one pre-launch hook but requires each caller to hand-write +SandboxAware injection if a lifecycle needs the sandbox. + +## Proposed Flow + +`McpEnv` will compose an internal pre-launch hook and pass that hook to +`RockRuntime.start()`. `RockRuntime` keeps the same public API and the same +internal sequencing. + +```text +McpEnv.start(before_launch=user_hook) + resolve server configs + compose McpEnv pre-launch hook + RockRuntime.start(resolved_servers, before_launch=mcp_env_hook) + Sandbox.start() + prepare /app/workspace and /data + write /app/mcp-servers.json + mcp_env_hook(sandbox) + inject sandbox into SandboxAware lifecycles + user_hook(sandbox) + launch /app/launch.sh + health-check SSE endpoints +``` + +The externally visible order becomes: + +1. Sandbox is created and started. +2. Runtime directories are prepared. +3. `/app/mcp-servers.json` is written. +4. `McpEnv` calls `set_sandbox(sandbox)` on every `SandboxAware` lifecycle. +5. The caller's `before_launch(sandbox)` callback runs, if provided. +6. MCP servers are launched and health-checked. + +## Component Responsibilities + +### `RockRuntime` + +`RockRuntime` remains the MCP runtime orchestrator. It abstracts the fixed +runtime sequence for running MCP server configs in a ROCK sandbox: + +- create `Sandbox(SandboxConfig(...))`; +- prepare runtime directories; +- render and write MCP server config; +- expose one generic pre-launch hook slot; +- launch `/app/launch.sh`; +- health-check SSE endpoints; +- stop the sandbox during release or failed startup cleanup. + +It should not import or reference `DataLifecycle`, `DataLifecycleFactory`, or +`SandboxAware`. + +### `McpEnv` + +`McpEnv` remains the integration facade between ROCK MCP runtime and +ScaffoldHub resources. It already owns ScaffoldHub auth providers, data +lifecycle creation, placeholder resolution, lifecycle `init/dump/reset`, and +auth lease release. SandboxAware injection belongs here because it is a +ScaffoldHub lifecycle integration concern. + +## Detailed Design + +### Loading ScaffoldHub Components + +`_load_scaffoldhub_components()` currently loads `AuthProvider` and +`DataLifecycleFactory`. It will also attempt to load `SandboxAware` from +`scaffoldhub.tools.base`. + +The preferred import source is: + +```python +from scaffoldhub.auth import AuthProvider +from scaffoldhub.tools.base import DataLifecycleFactory +``` + +`SandboxAware` should be imported separately so compatibility handling is +precise: + +```python +try: + from scaffoldhub.tools.base import SandboxAware +except ImportError: + SandboxAware = None +``` + +For compatibility with older ScaffoldHub versions, missing `SandboxAware` +should not break `McpEnv` construction. If `AuthProvider` or +`DataLifecycleFactory` cannot be imported, `McpEnv` should keep the existing +clear optional dependency error. If only `SandboxAware` is missing, +`McpEnv` stores `None` and skips automatic injection. + +### Construction State + +`McpEnv.__init__()` will store the loaded marker class: + +```python +self.sandbox_aware_class = sandbox_aware_class +``` + +No public constructor arguments are added. + +### Start Hook Composition + +`McpEnv.start(before_launch=None)` will compose a hook before calling +`RockRuntime.start()`: + +```python +runtime_before_launch = self._compose_before_launch(before_launch) +urls = await self._rock_runtime.start( + self.resolved_servers, + before_launch=runtime_before_launch, +) +``` + +The composed hook always runs SandboxAware injection first. It then invokes the +caller-provided hook, preserving support for both synchronous and asynchronous +callbacks. + +### SandboxAware Injection + +`McpEnv` will add a private helper: + +```python +def _inject_sandbox_into_lifecycles(self, sandbox: Sandbox) -> None: + sandbox_aware_class = self.sandbox_aware_class + if sandbox_aware_class is None: + return + + for lifecycle in self.data_lifecycles.values(): + if isinstance(lifecycle, sandbox_aware_class): + lifecycle.set_sandbox(sandbox) +``` + +The helper does not inspect lifecycle names, does not special-case tools, and +does not call lifecycle `before_launch()`. + +### User Hook Invocation + +`McpEnv` will mirror the current `RockRuntime` sync/async hook behavior: + +```python +result = before_launch(sandbox) +if inspect.isawaitable(result): + await result +``` + +This keeps existing callers compatible while guaranteeing that any caller hook +sees lifecycles after sandbox injection. + +## Error Handling + +If `set_sandbox()` raises, the composed hook raises. `RockRuntime.start()` +already treats hook failures as startup failures: + +- it calls `stop()` to clean up the sandbox; +- it logs cleanup failures without masking the original startup error; +- it raises `RockRuntimeError` with the original error chained. + +`McpEnv` should not add a second cleanup path. + +If the caller's `before_launch()` raises after successful injection, behavior +remains the same as today: startup fails and `RockRuntime` cleans up. + +## Compatibility + +Existing callers that do not use ScaffoldHub `SandboxAware` continue to work. +Existing callers with a `before_launch(sandbox)` callback continue to receive +the raw ROCK sandbox at the same runtime point, but after automatic lifecycle +injection. + +Older ScaffoldHub versions that lack `SandboxAware` keep the previous behavior: +no automatic lifecycle injection is performed, and the caller can still inject +manually from their own `before_launch` callback if needed. + +## Testing + +Unit tests should cover: + +- `McpEnv.start()` injects the sandbox into lifecycles that implement + `SandboxAware`. +- User `before_launch(sandbox)` runs after SandboxAware injection. +- Non-SandboxAware lifecycles are ignored. +- If `SandboxAware` is unavailable from ScaffoldHub, `McpEnv` still starts + using the previous behavior. +- If `set_sandbox()` raises, startup fails and the runtime cleanup path runs. +- Existing placeholder resolution, lifecycle `init/dump/reset`, auth lease + release, and raw sandbox property tests continue to pass. + +Focused verification after implementation: + +```bash +uv run pytest tests/unit/sdk/mcp/test_mcp_env.py -v +uv run pytest tests/unit/sdk/mcp/test_rock_runtime.py -v +uv run ruff check rock/sdk/mcp tests/unit/sdk/mcp +uv run ruff format rock/sdk/mcp tests/unit/sdk/mcp +``` + +## Documentation + +Update MCP SDK documentation to state: + +- `McpEnv` automatically injects the started ROCK sandbox into ScaffoldHub + lifecycles that implement `SandboxAware`. +- Injection happens after `/app/mcp-servers.json` is written and before the + caller's `before_launch` callback. +- `McpEnv` does not call lifecycle `before_launch()` methods automatically. + +## Risks + +- A lifecycle may implement `SandboxAware` and rely on the caller to run an + additional preparation method. This design intentionally does only dependency + injection; full lifecycle preparation requires a separate public contract. +- Older ScaffoldHub packages do not export `SandboxAware`. The compatibility + behavior avoids breaking construction but does not provide injection. +- If multiple lifecycles share mutable sandbox state, they all receive the same + sandbox object. This matches the external launcher ownership model. From 7b67c2a2d3ede2fe26159fcc92349188c2e81124 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:14:13 +0800 Subject: [PATCH 207/226] docs: translate McpEnv SandboxAware design --- ...2026-07-03-mcp-env-sandbox-aware-design.md | 281 ++++++++---------- 1 file changed, 132 insertions(+), 149 deletions(-) diff --git a/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md b/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md index e656808973..f16a3bc2eb 100644 --- a/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md +++ b/docs/superpowers/specs/2026-07-03-mcp-env-sandbox-aware-design.md @@ -1,140 +1,131 @@ -# McpEnv SandboxAware Lifecycle Injection Design +# McpEnv SandboxAware 生命周期注入设计 -## Context +## 背景 -ROCK owns `rock.sdk.mcp.McpEnv` and `RockRuntime`. `McpEnv` creates -ScaffoldHub data lifecycles through `DataLifecycleFactory`, while -`RockRuntime` creates the ROCK sandbox, writes `/app/mcp-servers.json`, starts -MCP server processes, health-checks SSE endpoints, and stops the sandbox. +ROCK 拥有 `rock.sdk.mcp.McpEnv` 和 `RockRuntime`。`McpEnv` 通过 +ScaffoldHub 的 `DataLifecycleFactory` 创建工具数据生命周期；`RockRuntime` +负责创建 ROCK sandbox、写入 `/app/mcp-servers.json`、启动 MCP server 进程、 +检查 SSE endpoint，并在释放时停止 sandbox。 -ScaffoldHub now exposes `SandboxAware` from `scaffoldhub.tools.base`. -`SandboxAware` is a narrow lifecycle marker whose public contract is only: +ScaffoldHub 现在从 `scaffoldhub.tools.base` 导出 `SandboxAware`。 +`SandboxAware` 是一个很窄的 lifecycle 标记接口，公共契约只有： ```python def set_sandbox(self, sandbox) -> None: ... ``` -ScaffoldHub's design states that external launchers, including ROCK -`McpEnv`, should inject the externally created ROCK sandbox into lifecycles -that implement `SandboxAware`. The marker does not ask the launcher to call -lifecycle `before_launch()` methods. Those methods remain lifecycle-internal -compatibility helpers and are not part of the public integration contract. +ScaffoldHub 的设计要求外部 launcher（包括 ROCK `McpEnv`）把外部创建好的 +ROCK sandbox 注入到实现了 `SandboxAware` 的 lifecycle 中。这个标记接口不要求 +launcher 调用 lifecycle 的 `before_launch()` 方法。那些方法仍然是具体 lifecycle +内部的兼容 helper，不属于公开集成契约。 -## Goals +## 目标 -- Automatically inject the ROCK sandbox into ScaffoldHub lifecycles that - implement `SandboxAware`. -- Run injection after the sandbox is started and `/app/mcp-servers.json` has - been written. -- Run injection before the user-provided `before_launch(sandbox)` callback. -- Preserve existing `McpEnv.start(before_launch=None)` API shape. -- Keep `RockRuntime` independent from ScaffoldHub data lifecycle concepts. -- Preserve current startup cleanup behavior when injection fails. +- 自动把 ROCK sandbox 注入到实现了 `SandboxAware` 的 ScaffoldHub lifecycle。 +- 注入发生在 sandbox 启动并写入 `/app/mcp-servers.json` 之后。 +- 注入发生在用户传入的 `before_launch(sandbox)` 回调之前。 +- 保持现有 `McpEnv.start(before_launch=None)` API 形状不变。 +- 保持 `RockRuntime` 不感知 ScaffoldHub 数据生命周期概念。 +- 注入失败时复用当前启动失败清理逻辑。 -## Non-Goals +## 非目标 -- Do not call lifecycle `before_launch()` automatically. -- Do not add `prepare`, `before_launch`, or sandbox setup methods to - `SandboxAware`. -- Do not add lifecycle parameters to `RockRuntime.start()`. -- Do not move ScaffoldHub lifecycle handling into `RockRuntime`. -- Do not change `init`, `dump`, `reset`, `release`, URL generation, auth - placeholder resolution, or auth lease release semantics. +- 不自动调用 lifecycle 的 `before_launch()`。 +- 不向 `SandboxAware` 增加 `prepare`、`before_launch` 或 sandbox setup 方法。 +- 不向 `RockRuntime.start()` 增加 lifecycle 参数。 +- 不把 ScaffoldHub lifecycle 处理移动到 `RockRuntime`。 +- 不改变 `init`、`dump`、`reset`、`release`、URL 生成、auth 占位符解析或 + auth lease release 语义。 -## Current Flow +## 当前流程 -Today `McpEnv.start()` resolves server configs and passes the caller's -`before_launch` callback directly into `RockRuntime.start()`. +现在 `McpEnv.start()` 解析 server 配置后，会把调用方传入的 `before_launch` +回调原样传给 `RockRuntime.start()`。 ```text McpEnv.start(before_launch=user_hook) - resolve server configs + 解析 server 配置 RockRuntime.start(resolved_servers, before_launch=user_hook) Sandbox.start() - prepare /app/workspace and /data - write /app/mcp-servers.json + 准备 /app/workspace 和 /data + 写入 /app/mcp-servers.json user_hook(sandbox) - launch /app/launch.sh - health-check SSE endpoints + 启动 /app/launch.sh + 检查 SSE endpoints ``` -This gives callers one pre-launch hook but requires each caller to hand-write -SandboxAware injection if a lifecycle needs the sandbox. +这给调用方提供了一个 pre-launch hook，但如果某个 lifecycle 需要 sandbox， +调用方必须自己手写 SandboxAware 注入逻辑。 -## Proposed Flow +## 设计流程 -`McpEnv` will compose an internal pre-launch hook and pass that hook to -`RockRuntime.start()`. `RockRuntime` keeps the same public API and the same -internal sequencing. +`McpEnv` 会组合一个内部 pre-launch hook，再把这个 hook 传给 +`RockRuntime.start()`。`RockRuntime` 保持相同的公共 API 和内部启动顺序。 ```text McpEnv.start(before_launch=user_hook) - resolve server configs - compose McpEnv pre-launch hook + 解析 server 配置 + 组合 McpEnv pre-launch hook RockRuntime.start(resolved_servers, before_launch=mcp_env_hook) Sandbox.start() - prepare /app/workspace and /data - write /app/mcp-servers.json + 准备 /app/workspace 和 /data + 写入 /app/mcp-servers.json mcp_env_hook(sandbox) - inject sandbox into SandboxAware lifecycles + 向 SandboxAware lifecycles 注入 sandbox user_hook(sandbox) - launch /app/launch.sh - health-check SSE endpoints + 启动 /app/launch.sh + 检查 SSE endpoints ``` -The externally visible order becomes: +外部可见顺序为： -1. Sandbox is created and started. -2. Runtime directories are prepared. -3. `/app/mcp-servers.json` is written. -4. `McpEnv` calls `set_sandbox(sandbox)` on every `SandboxAware` lifecycle. -5. The caller's `before_launch(sandbox)` callback runs, if provided. -6. MCP servers are launched and health-checked. +1. 创建并启动 sandbox。 +2. 准备 runtime 目录。 +3. 写入 `/app/mcp-servers.json`。 +4. `McpEnv` 对每个 `SandboxAware` lifecycle 调用 `set_sandbox(sandbox)`。 +5. 如果调用方提供了 `before_launch(sandbox)`，再执行该回调。 +6. 启动 MCP servers 并做 health check。 -## Component Responsibilities +## 组件职责 ### `RockRuntime` -`RockRuntime` remains the MCP runtime orchestrator. It abstracts the fixed -runtime sequence for running MCP server configs in a ROCK sandbox: +`RockRuntime` 仍然是 MCP runtime 编排器。它抽象的是“把 MCP server 配置跑进 +ROCK sandbox，并产出 SSE URL”的固定流程： -- create `Sandbox(SandboxConfig(...))`; -- prepare runtime directories; -- render and write MCP server config; -- expose one generic pre-launch hook slot; -- launch `/app/launch.sh`; -- health-check SSE endpoints; -- stop the sandbox during release or failed startup cleanup. +- 创建 `Sandbox(SandboxConfig(...))`； +- 准备 runtime 目录； +- 渲染并写入 MCP server 配置； +- 暴露一个通用 pre-launch hook 插槽； +- 启动 `/app/launch.sh`； +- 检查 SSE endpoints； +- release 或启动失败时停止 sandbox。 -It should not import or reference `DataLifecycle`, `DataLifecycleFactory`, or -`SandboxAware`. +它不应该 import 或引用 `DataLifecycle`、`DataLifecycleFactory`、`SandboxAware`。 ### `McpEnv` -`McpEnv` remains the integration facade between ROCK MCP runtime and -ScaffoldHub resources. It already owns ScaffoldHub auth providers, data -lifecycle creation, placeholder resolution, lifecycle `init/dump/reset`, and -auth lease release. SandboxAware injection belongs here because it is a -ScaffoldHub lifecycle integration concern. +`McpEnv` 仍然是 ROCK MCP runtime 和 ScaffoldHub 资源之间的集成 facade。它已经 +拥有 ScaffoldHub auth provider、data lifecycle 创建、占位符解析、lifecycle +`init/dump/reset` 以及 auth lease release。`SandboxAware` 注入属于 +ScaffoldHub lifecycle 集成职责，因此应放在 `McpEnv` 中。 -## Detailed Design +## 详细设计 -### Loading ScaffoldHub Components +### 加载 ScaffoldHub 组件 -`_load_scaffoldhub_components()` currently loads `AuthProvider` and -`DataLifecycleFactory`. It will also attempt to load `SandboxAware` from -`scaffoldhub.tools.base`. +`_load_scaffoldhub_components()` 当前加载 `AuthProvider` 和 +`DataLifecycleFactory`。它将额外尝试加载 `SandboxAware`。 -The preferred import source is: +首选导入来源为： ```python from scaffoldhub.auth import AuthProvider from scaffoldhub.tools.base import DataLifecycleFactory ``` -`SandboxAware` should be imported separately so compatibility handling is -precise: +`SandboxAware` 应单独导入，这样兼容处理更精确： ```python try: @@ -143,26 +134,24 @@ except ImportError: SandboxAware = None ``` -For compatibility with older ScaffoldHub versions, missing `SandboxAware` -should not break `McpEnv` construction. If `AuthProvider` or -`DataLifecycleFactory` cannot be imported, `McpEnv` should keep the existing -clear optional dependency error. If only `SandboxAware` is missing, -`McpEnv` stores `None` and skips automatic injection. +为了兼容旧版 ScaffoldHub，缺少 `SandboxAware` 不应导致 `McpEnv` 构造失败。 +如果 `AuthProvider` 或 `DataLifecycleFactory` 无法导入，`McpEnv` 仍保留当前 +清晰的 optional dependency 错误。如果只有 `SandboxAware` 缺失，`McpEnv` +保存 `None` 并跳过自动注入。 -### Construction State +### 构造状态 -`McpEnv.__init__()` will store the loaded marker class: +`McpEnv.__init__()` 保存加载到的标记类： ```python self.sandbox_aware_class = sandbox_aware_class ``` -No public constructor arguments are added. +不新增公开构造参数。 -### Start Hook Composition +### Start Hook 组合 -`McpEnv.start(before_launch=None)` will compose a hook before calling -`RockRuntime.start()`: +`McpEnv.start(before_launch=None)` 在调用 `RockRuntime.start()` 前组合 hook： ```python runtime_before_launch = self._compose_before_launch(before_launch) @@ -172,13 +161,12 @@ urls = await self._rock_runtime.start( ) ``` -The composed hook always runs SandboxAware injection first. It then invokes the -caller-provided hook, preserving support for both synchronous and asynchronous -callbacks. +组合后的 hook 总是先执行 SandboxAware 注入，然后执行调用方传入的 hook。 +调用方 hook 继续支持同步和异步两种形式。 -### SandboxAware Injection +### SandboxAware 注入 -`McpEnv` will add a private helper: +`McpEnv` 新增私有 helper： ```python def _inject_sandbox_into_lifecycles(self, sandbox: Sandbox) -> None: @@ -191,12 +179,12 @@ def _inject_sandbox_into_lifecycles(self, sandbox: Sandbox) -> None: lifecycle.set_sandbox(sandbox) ``` -The helper does not inspect lifecycle names, does not special-case tools, and -does not call lifecycle `before_launch()`. +该 helper 不检查 lifecycle 名称，不为具体工具写白名单，也不调用 lifecycle 的 +`before_launch()`。 -### User Hook Invocation +### 用户 Hook 调用 -`McpEnv` will mirror the current `RockRuntime` sync/async hook behavior: +`McpEnv` 复用当前 `RockRuntime` 的同步/异步 hook 语义： ```python result = before_launch(sandbox) @@ -204,49 +192,45 @@ if inspect.isawaitable(result): await result ``` -This keeps existing callers compatible while guaranteeing that any caller hook -sees lifecycles after sandbox injection. +这样保持现有调用方兼容，同时保证调用方 hook 执行时，相关 lifecycles 已经完成 +sandbox 注入。 -## Error Handling +## 错误处理 -If `set_sandbox()` raises, the composed hook raises. `RockRuntime.start()` -already treats hook failures as startup failures: +如果 `set_sandbox()` 抛异常，组合 hook 直接抛出。`RockRuntime.start()` 已经把 +hook 失败视为启动失败： -- it calls `stop()` to clean up the sandbox; -- it logs cleanup failures without masking the original startup error; -- it raises `RockRuntimeError` with the original error chained. +- 调用 `stop()` 清理 sandbox； +- 记录 cleanup 失败，但不遮蔽原始启动错误； +- 抛出 `RockRuntimeError`，并把原始错误挂在异常链上。 -`McpEnv` should not add a second cleanup path. +`McpEnv` 不应增加第二套清理路径。 -If the caller's `before_launch()` raises after successful injection, behavior -remains the same as today: startup fails and `RockRuntime` cleans up. +如果调用方的 `before_launch()` 在注入成功后抛异常，行为与今天一致：启动失败， +由 `RockRuntime` 清理 sandbox。 -## Compatibility +## 兼容性 -Existing callers that do not use ScaffoldHub `SandboxAware` continue to work. -Existing callers with a `before_launch(sandbox)` callback continue to receive -the raw ROCK sandbox at the same runtime point, but after automatic lifecycle -injection. +不使用 ScaffoldHub `SandboxAware` 的现有调用方继续正常工作。已有 +`before_launch(sandbox)` 回调的调用方仍会在同一个 runtime 时机拿到原始 ROCK +sandbox，只是此时自动 lifecycle 注入已经完成。 -Older ScaffoldHub versions that lack `SandboxAware` keep the previous behavior: -no automatic lifecycle injection is performed, and the caller can still inject -manually from their own `before_launch` callback if needed. +缺少 `SandboxAware` 的旧版 ScaffoldHub 保持旧行为：不执行自动 lifecycle 注入。 +如果调用方确实需要注入 sandbox，仍可在自己的 `before_launch` 回调中手动完成。 -## Testing +## 测试 -Unit tests should cover: +单测应覆盖： -- `McpEnv.start()` injects the sandbox into lifecycles that implement - `SandboxAware`. -- User `before_launch(sandbox)` runs after SandboxAware injection. -- Non-SandboxAware lifecycles are ignored. -- If `SandboxAware` is unavailable from ScaffoldHub, `McpEnv` still starts - using the previous behavior. -- If `set_sandbox()` raises, startup fails and the runtime cleanup path runs. -- Existing placeholder resolution, lifecycle `init/dump/reset`, auth lease - release, and raw sandbox property tests continue to pass. +- `McpEnv.start()` 会向实现 `SandboxAware` 的 lifecycle 注入 sandbox。 +- 用户 `before_launch(sandbox)` 在 SandboxAware 注入之后执行。 +- 非 SandboxAware lifecycle 会被忽略。 +- 当 ScaffoldHub 不导出 `SandboxAware` 时，`McpEnv` 仍按旧行为启动。 +- `set_sandbox()` 抛异常时，启动失败并走 runtime cleanup 路径。 +- 现有占位符解析、lifecycle `init/dump/reset`、auth lease release、raw sandbox + property 测试继续通过。 -Focused verification after implementation: +实现后的聚焦验证命令： ```bash uv run pytest tests/unit/sdk/mcp/test_mcp_env.py -v @@ -255,22 +239,21 @@ uv run ruff check rock/sdk/mcp tests/unit/sdk/mcp uv run ruff format rock/sdk/mcp tests/unit/sdk/mcp ``` -## Documentation +## 文档 -Update MCP SDK documentation to state: +更新 MCP SDK 文档，说明： -- `McpEnv` automatically injects the started ROCK sandbox into ScaffoldHub - lifecycles that implement `SandboxAware`. -- Injection happens after `/app/mcp-servers.json` is written and before the - caller's `before_launch` callback. -- `McpEnv` does not call lifecycle `before_launch()` methods automatically. +- `McpEnv` 会自动把已启动的 ROCK sandbox 注入到实现 `SandboxAware` 的 + ScaffoldHub lifecycle。 +- 注入发生在 `/app/mcp-servers.json` 写入之后、调用方 `before_launch` 回调之前。 +- `McpEnv` 不会自动调用 lifecycle 的 `before_launch()` 方法。 -## Risks +## 风险 -- A lifecycle may implement `SandboxAware` and rely on the caller to run an - additional preparation method. This design intentionally does only dependency - injection; full lifecycle preparation requires a separate public contract. -- Older ScaffoldHub packages do not export `SandboxAware`. The compatibility - behavior avoids breaking construction but does not provide injection. -- If multiple lifecycles share mutable sandbox state, they all receive the same - sandbox object. This matches the external launcher ownership model. +- 某些 lifecycle 可能实现了 `SandboxAware`，但仍需要调用方执行额外准备方法。 + 本设计有意只做依赖注入；完整 lifecycle 准备需要另一个公开契约。 +- 旧版 ScaffoldHub package 不导出 `SandboxAware`。兼容行为可以避免构造失败， + 但不会提供自动注入。 +- 多个 lifecycles 会收到同一个 sandbox 对象。如果它们共享可变 sandbox 状态， + 需要由具体 lifecycle 自身保证使用方式正确。这符合外部 launcher 拥有 sandbox + 的模型。 From 2a91ab8684bb4d56562af8d95538fd1b6aed046f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:35:42 +0800 Subject: [PATCH 208/226] fix: update mcp proxy sse urls --- rock/sdk/mcp/rock_runtime.py | 2 +- tests/unit/sdk/mcp/test_rock_runtime.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index 27363e0c8d..f860f9a8b4 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -117,7 +117,7 @@ def get_server_url(self, server_name: str) -> str: raise RockRuntimeError("ROCK sandbox has not been started") config = self._require_config() - return f"{config.base_url}/apis/envs/sandbox/v1/sandboxes/{sandbox_id}/proxy/{server_name}/sse" + return f"{config.base_url}/apis/envs/sandbox/v1/sandboxes/{sandbox_id}/proxy/mcp/{server_name}/sse" def get_all_server_urls(self, server_names: Iterable[str]) -> dict[str, str]: return {name: self.get_server_url(name) for name in server_names} diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index 264d33dc27..00bebaa227 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -57,8 +57,8 @@ def test_rock_runtime_builds_server_urls_from_sandbox_id(monkeypatch): runtime._sandbox_id = "sandbox-123" assert runtime.get_all_server_urls(["calculator", "slack"]) == { - "calculator": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/calculator/sse", - "slack": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/slack/sse", + "calculator": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/mcp/calculator/sse", + "slack": "https://xrl.alibaba-inc.com/apis/envs/sandbox/v1/sandboxes/sandbox-123/proxy/mcp/slack/sse", } From 215e203e865396250de1900e511a0bb10f40180e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:39:31 +0800 Subject: [PATCH 209/226] test: add McpEnv SandboxAware fakes --- tests/unit/sdk/mcp/test_mcp_env.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index 1dbc17104e..52f3366f6d 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -37,6 +37,27 @@ def reset(self) -> None: self.initialized_data = {} +class FakeSandboxAware: + def set_sandbox(self, sandbox) -> None: + self.sandbox = sandbox + + +class RecordingSandboxAwareLifecycle(RecordingDataLifecycle, FakeSandboxAware): + def __init__(self): + super().__init__() + self.sandbox = None + self.events: list[str] = [] + + def set_sandbox(self, sandbox) -> None: + self.events.append("set_sandbox") + self.sandbox = sandbox + + +class FailingSandboxAwareLifecycle(RecordingDataLifecycle, FakeSandboxAware): + def set_sandbox(self, sandbox) -> None: + raise RuntimeError("sandbox injection failed") + + class FakeAuthProvider: def __init__(self): self.auth = { @@ -81,13 +102,15 @@ def release_active_leases(self) -> None: raise RuntimeError("database release failed") -def install_fake_scaffoldhub(monkeypatch): +def install_fake_scaffoldhub(monkeypatch, *, include_sandbox_aware: bool = True): scaffoldhub = ModuleType("scaffoldhub") auth = ModuleType("scaffoldhub.auth") tools = ModuleType("scaffoldhub.tools") base = ModuleType("scaffoldhub.tools.base") auth.AuthProvider = FakeAuthProvider base.DataLifecycleFactory = FakeDataLifecycleFactory + if include_sandbox_aware: + base.SandboxAware = FakeSandboxAware monkeypatch.setitem(sys.modules, "scaffoldhub", scaffoldhub) monkeypatch.setitem(sys.modules, "scaffoldhub.auth", auth) @@ -95,8 +118,8 @@ def install_fake_scaffoldhub(monkeypatch): monkeypatch.setitem(sys.modules, "scaffoldhub.tools.base", base) -def reload_mcp_env(monkeypatch): - install_fake_scaffoldhub(monkeypatch) +def reload_mcp_env(monkeypatch, *, include_sandbox_aware: bool = True): + install_fake_scaffoldhub(monkeypatch, include_sandbox_aware=include_sandbox_aware) sys.modules.pop("rock.sdk.mcp.mcp_env", None) module = importlib.import_module("rock.sdk.mcp.mcp_env") return importlib.reload(module) From 15beb12ad7db249e4da746d477cf0df8a6bc1289 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:40:25 +0800 Subject: [PATCH 210/226] test: cover McpEnv SandboxAware injection --- tests/unit/sdk/mcp/test_mcp_env.py | 113 +++++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index 52f3366f6d..8455a9f513 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -21,6 +21,37 @@ async def stop(self): raise RuntimeError("stop failed") +class RecordingStartRuntime: + def __init__(self, sandbox=None): + self.sandbox = sandbox or object() + self.started_servers = None + self.received_before_launch = None + self.stop_calls = 0 + + async def start(self, servers, before_launch=None): + self.started_servers = deepcopy(servers) + self.received_before_launch = before_launch + if before_launch is not None: + await before_launch(self.sandbox) + return {name: f"https://example.test/{name}/sse" for name in sorted(servers)} + + async def stop(self): + self.stop_calls += 1 + + +class CleanupRecordingStartRuntime(RecordingStartRuntime): + async def start(self, servers, before_launch=None): + self.started_servers = deepcopy(servers) + self.received_before_launch = before_launch + try: + if before_launch is not None: + await before_launch(self.sandbox) + except Exception: + await self.stop() + raise + return {name: f"https://example.test/{name}/sse" for name in sorted(servers)} + + class RecordingDataLifecycle: def __init__(self): self.initialized_data = {} @@ -493,6 +524,88 @@ def test_mcp_env_start_accepts_before_launch_hook(monkeypatch): assert signature.parameters["before_launch"].default is None +def test_mcp_env_start_injects_sandbox_into_sandbox_aware_lifecycle(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingSandboxAwareLifecycle() + sandbox = object() + runtime = RecordingStartRuntime(sandbox=sandbox) + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = runtime + + asyncio.run(env.start()) + + assert lifecycle.sandbox is sandbox + assert lifecycle.events == ["set_sandbox"] + assert env.is_alive() is True + assert env.get_urls() == {"slack": "https://example.test/slack/sse"} + + +def test_mcp_env_start_runs_user_before_launch_after_sandbox_injection(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingSandboxAwareLifecycle() + sandbox = object() + events: list[str] = [] + runtime = RecordingStartRuntime(sandbox=sandbox) + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = runtime + + async def before_launch(received_sandbox): + events.extend(lifecycle.events) + events.append("user_before_launch") + assert received_sandbox is sandbox + assert lifecycle.sandbox is sandbox + + asyncio.run(env.start(before_launch=before_launch)) + + assert events == ["set_sandbox", "user_before_launch"] + + +def test_mcp_env_start_ignores_non_sandbox_aware_lifecycle(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingDataLifecycle() + runtime = RecordingStartRuntime() + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = runtime + + asyncio.run(env.start()) + + assert not hasattr(lifecycle, "sandbox") + assert env.is_alive() is True + + +def test_mcp_env_start_skips_injection_when_sandbox_aware_is_unavailable(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch, include_sandbox_aware=False) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + lifecycle = RecordingSandboxAwareLifecycle() + runtime = RecordingStartRuntime() + env.data_lifecycles["slack"] = lifecycle + env._rock_runtime = runtime + + asyncio.run(env.start()) + + assert lifecycle.sandbox is None + assert lifecycle.events == [] + assert env.is_alive() is True + + +def test_mcp_env_start_propagates_sandbox_injection_failure_and_runtime_cleans_up(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + runtime = CleanupRecordingStartRuntime() + env.data_lifecycles["slack"] = FailingSandboxAwareLifecycle() + env._rock_runtime = runtime + + with pytest.raises(RuntimeError, match="sandbox injection failed"): + asyncio.run(env.start()) + + assert runtime.stop_calls == 1 + assert env.is_alive() is False + assert env.urls == {} + + def test_mcp_env_exposes_raw_sandbox_property(monkeypatch): mcp_env = reload_mcp_env(monkeypatch) From bcf3b732cf40774b516bd7ae9aa3125914e9af03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:41:23 +0800 Subject: [PATCH 211/226] feat: inject sandbox into SandboxAware lifecycles --- rock/sdk/mcp/mcp_env.py | 35 ++++++++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index cacb673823..e9faecb877 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -1,5 +1,6 @@ from __future__ import annotations +import inspect import logging from copy import deepcopy from typing import Any @@ -16,7 +17,13 @@ def _load_scaffoldhub_components(): from scaffoldhub.tools.base import DataLifecycleFactory except ImportError as error: raise ImportError("rock.sdk.mcp requires scaffoldhub. Install it with `pip install 'rl-rock[mcp]'`.") from error - return AuthProvider, DataLifecycleFactory + + try: + from scaffoldhub.tools.base import SandboxAware + except ImportError: + SandboxAware = None + + return AuthProvider, DataLifecycleFactory, SandboxAware class McpEnv: @@ -48,9 +55,10 @@ def __init__(self, servers: dict | None = None): self.urls = {} self.servers = deepcopy(servers) self.resolved_servers = {} - auth_provider_class, data_lifecycle_factory_class = _load_scaffoldhub_components() + auth_provider_class, data_lifecycle_factory_class, sandbox_aware_class = _load_scaffoldhub_components() self.auth_provider = auth_provider_class() self.data_lifecycle_factory = data_lifecycle_factory_class(auth_provider=self.auth_provider) + self.sandbox_aware_class = sandbox_aware_class self.data_lifecycles: dict[str, Any] = {} self._rock_runtime = RockRuntime() @@ -84,7 +92,7 @@ async def start(self, before_launch: BeforeLaunchHook | None = None) -> None: } urls = await self._rock_runtime.start( self.resolved_servers, - before_launch=before_launch, + before_launch=self._compose_before_launch(before_launch), ) self.urls = urls self.running = True @@ -183,6 +191,27 @@ async def release(self): if auth_release_error is not None: raise RuntimeError("Failed to release MCP auth leases") from auth_release_error + def _compose_before_launch(self, before_launch: BeforeLaunchHook | None) -> BeforeLaunchHook: + async def composed_before_launch(sandbox: Sandbox) -> None: + self._inject_sandbox_into_lifecycles(sandbox) + if before_launch is None: + return + + result = before_launch(sandbox) + if inspect.isawaitable(result): + await result + + return composed_before_launch + + def _inject_sandbox_into_lifecycles(self, sandbox: Sandbox) -> None: + sandbox_aware_class = self.sandbox_aware_class + if sandbox_aware_class is None: + return + + for lifecycle in self.data_lifecycles.values(): + if isinstance(lifecycle, sandbox_aware_class): + lifecycle.set_sandbox(sandbox) + def _resolve_server_config(self, server_name: str, server_config: Any) -> Any: if not isinstance(server_config, dict): return deepcopy(server_config) From 1183f2d5b08ecb21dd0cad74bbfe011df129ad84 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:42:02 +0800 Subject: [PATCH 212/226] docs: document McpEnv SandboxAware injection --- .../References/Python SDK References/mcp.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index 52ea21e244..66d9b3fbef 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -70,11 +70,23 @@ resolution can borrow auth before the runtime is marked alive. If auth lease release fails, `release()` raises `RuntimeError("Failed to release MCP auth leases")` and can be called again to retry lease release. +## SandboxAware Lifecycles + +When ScaffoldHub exports `SandboxAware`, `McpEnv` automatically injects the +started ROCK sandbox into each configured lifecycle that implements that +interface. + +Injection happens after `/app/mcp-servers.json` is written and before the +caller-provided `before_launch(sandbox)` callback runs. `McpEnv` only calls +`set_sandbox(sandbox)`; it does not call lifecycle `before_launch()` methods +automatically. + ## Launch Hook `start()` accepts an optional sync or async `before_launch` callback. The callback receives the raw ROCK `Sandbox` after `/app/mcp-servers.json` is -written and before `/app/launch.sh` starts MCP servers. +written, after `SandboxAware` lifecycle injection, and before `/app/launch.sh` +starts MCP servers. ```python async def before_launch(sandbox): From 4ce8cc2817127b7b1dacad4f8d4049d07e5c7f00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 15:42:48 +0800 Subject: [PATCH 213/226] chore: update MCP dependency lockfile --- uv.lock | 101 +++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 96 insertions(+), 5 deletions(-) diff --git a/uv.lock b/uv.lock index 7e2d59c231..502088c9b8 100644 --- a/uv.lock +++ b/uv.lock @@ -897,6 +897,20 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/0a/4c/925909008ed5a988ccbb72dcc897407e5d6d3bd72410d69e051fc0c14647/charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f" }, ] +[[package]] +name = "circus" +version = "0.19.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "psutil", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "pyzmq", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "tornado", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/94/97/824bfce6949716ea93adcd5ff8aa4c277f40a735d7f644669674ec132ae4/circus-0.19.0.tar.gz", hash = "sha256:fbe6a5029998ac1239b17ebdd38251ac8b22627d30e4ec6f68cb10233911b0f4" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/f4/2c/1b09e40d512b7b9f9e58f2ee6c4648461e3fb40de2201856adaa1d22e96f/circus-0.19.0-py3-none-any.whl", hash = "sha256:15cac59d2bac8d8793f801a3a57e54acb261590c93e29fbfe639eaef8a680d39" }, +] + [[package]] name = "click" version = "8.3.0" @@ -4181,6 +4195,79 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b" }, ] +[[package]] +name = "pyzmq" +version = "27.1.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "cffi", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and implementation_name == 'pypy'" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/04/0b/3c9baedbdf613ecaa7aa07027780b8867f57b6293b6ee50de316c9f3222b/pyzmq-27.1.0.tar.gz", hash = "sha256:ac0765e3d44455adb6ddbf4417dcce460fc40a05978c08efdf2948072f6db540" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/67/b9/52aa9ec2867528b54f1e60846728d8b4d84726630874fee3a91e66c7df81/pyzmq-27.1.0-cp310-cp310-macosx_10_15_universal2.whl", hash = "sha256:508e23ec9bc44c0005c4946ea013d9317ae00ac67778bd47519fdf5a0e930ff4" }, + { url = "https://mirrors.aliyun.com/pypi/packages/99/64/5653e7b7425b169f994835a2b2abf9486264401fdef18df91ddae47ce2cc/pyzmq-27.1.0-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:507b6f430bdcf0ee48c0d30e734ea89ce5567fd7b8a0f0044a369c176aa44556" }, + { url = "https://mirrors.aliyun.com/pypi/packages/73/78/7d713284dbe022f6440e391bd1f3c48d9185673878034cfb3939cdf333b2/pyzmq-27.1.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bf7b38f9fd7b81cb6d9391b2946382c8237fd814075c6aa9c3b746d53076023b" }, + { url = "https://mirrors.aliyun.com/pypi/packages/30/76/8f099f9d6482450428b17c4d6b241281af7ce6a9de8149ca8c1c649f6792/pyzmq-27.1.0-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:03ff0b279b40d687691a6217c12242ee71f0fba28bf8626ff50e3ef0f4410e1e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/59/f0/37fbfff06c68016019043897e4c969ceab18bde46cd2aca89821fcf4fb2e/pyzmq-27.1.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:677e744fee605753eac48198b15a2124016c009a11056f93807000ab11ce6526" }, + { url = "https://mirrors.aliyun.com/pypi/packages/47/14/7254be73f7a8edc3587609554fcaa7bfd30649bf89cd260e4487ca70fdaa/pyzmq-27.1.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:dd2fec2b13137416a1c5648b7009499bcc8fea78154cd888855fa32514f3dad1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/22/dc/49f2be26c6f86f347e796a4d99b19167fc94503f0af3fd010ad262158822/pyzmq-27.1.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:08e90bb4b57603b84eab1d0ca05b3bbb10f60c1839dc471fc1c9e1507bef3386" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a3/3e/154fb963ae25be70c0064ce97776c937ecc7d8b0259f22858154a9999769/pyzmq-27.1.0-cp310-cp310-win32.whl", hash = "sha256:a5b42d7a0658b515319148875fcb782bbf118dd41c671b62dae33666c2213bda" }, + { url = "https://mirrors.aliyun.com/pypi/packages/62/b2/f4ab56c8c595abcb26b2be5fd9fa9e6899c1e5ad54964e93ae8bb35482be/pyzmq-27.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:c0bb87227430ee3aefcc0ade2088100e528d5d3298a0a715a64f3d04c60ba02f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3b/e3/be2cc7ab8332bdac0522fdb64c17b1b6241a795bee02e0196636ec5beb79/pyzmq-27.1.0-cp310-cp310-win_arm64.whl", hash = "sha256:9a916f76c2ab8d045b19f2286851a38e9ac94ea91faf65bd64735924522a8b32" }, + { url = "https://mirrors.aliyun.com/pypi/packages/06/5d/305323ba86b284e6fcb0d842d6adaa2999035f70f8c38a9b6d21ad28c3d4/pyzmq-27.1.0-cp311-cp311-macosx_10_15_universal2.whl", hash = "sha256:226b091818d461a3bef763805e75685e478ac17e9008f49fce2d3e52b3d58b86" }, + { url = "https://mirrors.aliyun.com/pypi/packages/bd/a0/fc7e78a23748ad5443ac3275943457e8452da67fda347e05260261108cbc/pyzmq-27.1.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:0790a0161c281ca9723f804871b4027f2e8b5a528d357c8952d08cd1a9c15581" }, + { url = "https://mirrors.aliyun.com/pypi/packages/7e/22/37d15eb05f3bdfa4abea6f6d96eb3bb58585fbd3e4e0ded4e743bc650c97/pyzmq-27.1.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c895a6f35476b0c3a54e3eb6ccf41bf3018de937016e6e18748317f25d4e925f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b1/c4/2a6fe5111a01005fc7af3878259ce17684fabb8852815eda6225620f3c59/pyzmq-27.1.0-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bbf8d3630bf96550b3be8e1fc0fea5cbdc8d5466c1192887bd94869da17a63e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/cb/eb/bfdcb41d0db9cd233d6fb22dc131583774135505ada800ebf14dfb0a7c40/pyzmq-27.1.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:15c8bd0fe0dabf808e2d7a681398c4e5ded70a551ab47482067a572c054c8e2e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ab/21/e3180ca269ed4a0de5c34417dfe71a8ae80421198be83ee619a8a485b0c7/pyzmq-27.1.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:bafcb3dd171b4ae9f19ee6380dfc71ce0390fefaf26b504c0e5f628d7c8c54f2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3b/b1/5e21d0b517434b7f33588ff76c177c5a167858cc38ef740608898cd329f2/pyzmq-27.1.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e829529fcaa09937189178115c49c504e69289abd39967cd8a4c215761373394" }, + { url = "https://mirrors.aliyun.com/pypi/packages/03/f2/44913a6ff6941905efc24a1acf3d3cb6146b636c546c7406c38c49c403d4/pyzmq-27.1.0-cp311-cp311-win32.whl", hash = "sha256:6df079c47d5902af6db298ec92151db82ecb557af663098b92f2508c398bb54f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/23/6d/d8d92a0eb270a925c9b4dd039c0b4dc10abc2fcbc48331788824ef113935/pyzmq-27.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:190cbf120fbc0fc4957b56866830def56628934a9d112aec0e2507aa6a032b97" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ae/14/01afebc96c5abbbd713ecfc7469cfb1bc801c819a74ed5c9fad9a48801cb/pyzmq-27.1.0-cp311-cp311-win_arm64.whl", hash = "sha256:eca6b47df11a132d1745eb3b5b5e557a7dae2c303277aa0e69c6ba91b8736e07" }, + { url = "https://mirrors.aliyun.com/pypi/packages/92/e7/038aab64a946d535901103da16b953c8c9cc9c961dadcbf3609ed6428d23/pyzmq-27.1.0-cp312-abi3-macosx_10_15_universal2.whl", hash = "sha256:452631b640340c928fa343801b0d07eb0c3789a5ffa843f6e1a9cee0ba4eb4fc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e8/5e/c3c49fdd0f535ef45eefcc16934648e9e59dace4a37ee88fc53f6cd8e641/pyzmq-27.1.0-cp312-abi3-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1c179799b118e554b66da67d88ed66cd37a169f1f23b5d9f0a231b4e8d44a113" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f8/e5/b0b2504cb4e903a74dcf1ebae157f9e20ebb6ea76095f6cfffea28c42ecd/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3837439b7f99e60312f0c926a6ad437b067356dc2bc2ec96eb395fd0fe804233" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f8/9b/c108cdb55560eaf253f0cbdb61b29971e9fb34d9c3499b0e96e4e60ed8a5/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43ad9a73e3da1fab5b0e7e13402f0b2fb934ae1c876c51d0afff0e7c052eca31" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c2/bb/b79798ca177b9eb0825b4c9998c6af8cd2a7f15a6a1a4272c1d1a21d382f/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0de3028d69d4cdc475bfe47a6128eb38d8bc0e8f4d69646adfbcd840facbac28" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9c/80/2df2e7977c4ede24c79ae39dcef3899bfc5f34d1ca7a5b24f182c9b7a9ca/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_i686.whl", hash = "sha256:cf44a7763aea9298c0aa7dbf859f87ed7012de8bda0f3977b6fb1d96745df856" }, + { url = "https://mirrors.aliyun.com/pypi/packages/46/bd/2d45ad24f5f5ae7e8d01525eb76786fa7557136555cac7d929880519e33a/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f30f395a9e6fbca195400ce833c731e7b64c3919aa481af4d88c3759e0cb7496" }, + { url = "https://mirrors.aliyun.com/pypi/packages/e6/2f/104c0a3c778d7c2ab8190e9db4f62f0b6957b53c9d87db77c284b69f33ea/pyzmq-27.1.0-cp312-abi3-win32.whl", hash = "sha256:250e5436a4ba13885494412b3da5d518cd0d3a278a1ae640e113c073a5f88edd" }, + { url = "https://mirrors.aliyun.com/pypi/packages/fc/7f/a21b20d577e4100c6a41795842028235998a643b1ad406a6d4163ea8f53e/pyzmq-27.1.0-cp312-abi3-win_amd64.whl", hash = "sha256:9ce490cf1d2ca2ad84733aa1d69ce6855372cb5ce9223802450c9b2a7cba0ccf" }, + { url = "https://mirrors.aliyun.com/pypi/packages/78/c2/c012beae5f76b72f007a9e91ee9401cb88c51d0f83c6257a03e785c81cc2/pyzmq-27.1.0-cp312-abi3-win_arm64.whl", hash = "sha256:75a2f36223f0d535a0c919e23615fc85a1e23b71f40c7eb43d7b1dedb4d8f15f" }, + { url = "https://mirrors.aliyun.com/pypi/packages/60/cb/84a13459c51da6cec1b7b1dc1a47e6db6da50b77ad7fd9c145842750a011/pyzmq-27.1.0-cp313-cp313-android_24_arm64_v8a.whl", hash = "sha256:93ad4b0855a664229559e45c8d23797ceac03183c7b6f5b4428152a6b06684a5" }, + { url = "https://mirrors.aliyun.com/pypi/packages/dc/b6/94414759a69a26c3dd674570a81813c46a078767d931a6c70ad29fc585cb/pyzmq-27.1.0-cp313-cp313-android_24_x86_64.whl", hash = "sha256:fbb4f2400bfda24f12f009cba62ad5734148569ff4949b1b6ec3b519444342e6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a5/ad/15906493fd40c316377fd8a8f6b1f93104f97a752667763c9b9c1b71d42d/pyzmq-27.1.0-cp313-cp313t-macosx_10_15_universal2.whl", hash = "sha256:e343d067f7b151cfe4eb3bb796a7752c9d369eed007b91231e817071d2c2fec7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/14/1d/d343f3ce13db53a54cb8946594e567410b2125394dafcc0268d8dda027e0/pyzmq-27.1.0-cp313-cp313t-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:08363b2011dec81c354d694bdecaef4770e0ae96b9afea70b3f47b973655cc05" }, + { url = "https://mirrors.aliyun.com/pypi/packages/69/2d/d83dd6d7ca929a2fc67d2c3005415cdf322af7751d773524809f9e585129/pyzmq-27.1.0-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d54530c8c8b5b8ddb3318f481297441af102517602b569146185fa10b63f4fa9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3e/cd/9822a7af117f4bc0f1952dbe9ef8358eb50a24928efd5edf54210b850259/pyzmq-27.1.0-cp313-cp313t-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6f3afa12c392f0a44a2414056d730eebc33ec0926aae92b5ad5cf26ebb6cc128" }, + { url = "https://mirrors.aliyun.com/pypi/packages/9a/12/f003e824a19ed73be15542f172fd0ec4ad0b60cf37436652c93b9df7c585/pyzmq-27.1.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c65047adafe573ff023b3187bb93faa583151627bc9c51fc4fb2c561ed689d39" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d5/4a/e82d788ed58e9a23995cee70dbc20c9aded3d13a92d30d57ec2291f1e8a3/pyzmq-27.1.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:90e6e9441c946a8b0a667356f7078d96411391a3b8f80980315455574177ec97" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d9/94/2da0a60841f757481e402b34bf4c8bf57fa54a5466b965de791b1e6f747d/pyzmq-27.1.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:add071b2d25f84e8189aaf0882d39a285b42fa3853016ebab234a5e78c7a43db" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4f/6f/55c10e2e49ad52d080dc24e37adb215e5b0d64990b57598abc2e3f01725b/pyzmq-27.1.0-cp313-cp313t-win32.whl", hash = "sha256:7ccc0700cfdf7bd487bea8d850ec38f204478681ea02a582a8da8171b7f90a1c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/87/4d/2534970ba63dd7c522d8ca80fb92777f362c0f321900667c615e2067cb29/pyzmq-27.1.0-cp313-cp313t-win_amd64.whl", hash = "sha256:8085a9fba668216b9b4323be338ee5437a235fe275b9d1610e422ccc279733e2" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f6/fa/f8aea7a28b0641f31d40dea42d7ef003fded31e184ef47db696bc74cd610/pyzmq-27.1.0-cp313-cp313t-win_arm64.whl", hash = "sha256:6bb54ca21bcfe361e445256c15eedf083f153811c37be87e0514934d6913061e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/87/45/19efbb3000956e82d0331bafca5d9ac19ea2857722fa2caacefb6042f39d/pyzmq-27.1.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:ce980af330231615756acd5154f29813d553ea555485ae712c491cd483df6b7a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/48/43/d72ccdbf0d73d1343936296665826350cb1e825f92f2db9db3e61c2162a2/pyzmq-27.1.0-cp314-cp314t-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1779be8c549e54a1c38f805e56d2a2e5c009d26de10921d7d51cfd1c8d4632ea" }, + { url = "https://mirrors.aliyun.com/pypi/packages/2f/2e/a483f73a10b65a9ef0161e817321d39a770b2acf8bcf3004a28d90d14a94/pyzmq-27.1.0-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7200bb0f03345515df50d99d3db206a0a6bee1955fbb8c453c76f5bf0e08fb96" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f5/d2/5f36552c2d3e5685abe60dfa56f91169f7a2d99bbaf67c5271022ab40863/pyzmq-27.1.0-cp314-cp314t-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01c0e07d558b06a60773744ea6251f769cd79a41a97d11b8bf4ab8f034b0424d" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c4/2a/404b331f2b7bf3198e9945f75c4c521f0c6a3a23b51f7a4a401b94a13833/pyzmq-27.1.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:80d834abee71f65253c91540445d37c4c561e293ba6e741b992f20a105d69146" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1c/0b/f4107e33f62a5acf60e3ded67ed33d79b4ce18de432625ce2fc5093d6388/pyzmq-27.1.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:544b4e3b7198dde4a62b8ff6685e9802a9a1ebf47e77478a5eb88eca2a82f2fd" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0d/01/add31fe76512642fd6e40e3a3bd21f4b47e242c8ba33efb6809e37076d9b/pyzmq-27.1.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:cedc4c68178e59a4046f97eca31b148ddcf51e88677de1ef4e78cf06c5376c9a" }, + { url = "https://mirrors.aliyun.com/pypi/packages/c4/59/a5f38970f9bf07cee96128de79590bb354917914a9be11272cfc7ff26af0/pyzmq-27.1.0-cp314-cp314t-win32.whl", hash = "sha256:1f0b2a577fd770aa6f053211a55d1c47901f4d537389a034c690291485e5fe92" }, + { url = "https://mirrors.aliyun.com/pypi/packages/70/d8/78b1bad170f93fcf5e3536e70e8fadac55030002275c9a29e8f5719185de/pyzmq-27.1.0-cp314-cp314t-win_amd64.whl", hash = "sha256:19c9468ae0437f8074af379e986c5d3d7d7bfe033506af442e8c879732bedbe0" }, + { url = "https://mirrors.aliyun.com/pypi/packages/81/d6/4bfbb40c9a0b42fc53c7cf442f6385db70b40f74a783130c5d0a5aa62228/pyzmq-27.1.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dc5dbf68a7857b59473f7df42650c621d7e8923fb03fa74a526890f4d33cc4d7" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f3/81/a65e71c1552f74dec9dff91d95bafb6e0d33338a8dfefbc88aa562a20c92/pyzmq-27.1.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:c17e03cbc9312bee223864f1a2b13a99522e0dc9f7c5df0177cd45210ac286e6" }, + { url = "https://mirrors.aliyun.com/pypi/packages/58/ed/0202ca350f4f2b69faa95c6d931e3c05c3a397c184cacb84cb4f8f42f287/pyzmq-27.1.0-pp310-pypy310_pp73-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:f328d01128373cb6763823b2b4e7f73bdf767834268c565151eacb3b7a392f90" }, + { url = "https://mirrors.aliyun.com/pypi/packages/47/42/1ff831fa87fe8f0a840ddb399054ca0009605d820e2b44ea43114f5459f4/pyzmq-27.1.0-pp310-pypy310_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c1790386614232e1b3a40a958454bdd42c6d1811837b15ddbb052a032a43f62" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d1/db/5c4d6807434751e3f21231bee98109aa57b9b9b55e058e450d0aef59b70f/pyzmq-27.1.0-pp310-pypy310_pp73-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:448f9cb54eb0cee4732b46584f2710c8bc178b0e5371d9e4fc8125201e413a74" }, + { url = "https://mirrors.aliyun.com/pypi/packages/26/af/78ce193dbf03567eb8c0dc30e3df2b9e56f12a670bf7eb20f9fb532c7e8a/pyzmq-27.1.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:05b12f2d32112bf8c95ef2e74ec4f1d4beb01f8b5e703b38537f8849f92cb9ba" }, + { url = "https://mirrors.aliyun.com/pypi/packages/4c/c6/c4dcdecdbaa70969ee1fdced6d7b8f60cfabe64d25361f27ac4665a70620/pyzmq-27.1.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:18770c8d3563715387139060d37859c02ce40718d1faf299abddcdcc6a649066" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3e/79/f38c92eeaeb03a2ccc2ba9866f0439593bb08c5e3b714ac1d553e5c96e25/pyzmq-27.1.0-pp311-pypy311_pp73-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:ac25465d42f92e990f8d8b0546b01c391ad431c3bf447683fdc40565941d0604" }, + { url = "https://mirrors.aliyun.com/pypi/packages/49/0e/3f0d0d335c6b3abb9b7b723776d0b21fa7f3a6c819a0db6097059aada160/pyzmq-27.1.0-pp311-pypy311_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:53b40f8ae006f2734ee7608d59ed661419f087521edbfc2149c3932e9c14808c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/a1/cf/f2b3784d536250ffd4be70e049f3b60981235d70c6e8ce7e3ef21e1adb25/pyzmq-27.1.0-pp311-pypy311_pp73-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f605d884e7c8be8fe1aa94e0a783bf3f591b84c24e4bc4f3e7564c82ac25e271" }, + { url = "https://mirrors.aliyun.com/pypi/packages/01/1b/5dbe84eefc86f48473947e2f41711aded97eecef1231f4558f1f02713c12/pyzmq-27.1.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c9f7f6e13dff2e44a6afeaf2cf54cee5929ad64afaf4d40b50f93c58fc687355" }, +] + [[package]] name = "ray" version = "2.43.0" @@ -4428,7 +4515,7 @@ wheels = [ [[package]] name = "rl-rock" -version = "1.10.0.dev3" +version = "1.10.0.dev11" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -4610,7 +4697,7 @@ requires-dist = [ { name = "rl-rock", extras = ["builder"], marker = "extra == 'all'" }, { name = "rl-rock", extras = ["rocklet"], marker = "extra == 'admin'" }, { name = "rl-rock", extras = ["rocklet"], marker = "extra == 'all'" }, - { name = "scaffoldhub", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'mcp'", specifier = "==0.1.0.dev4", index = "https://artlab.alibaba-inc.com/1/pypi/simple" }, + { name = "scaffoldhub", marker = "python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'mcp'", specifier = "==0.1.0.dev11", index = "https://artlab.alibaba-inc.com/1/pypi/simple" }, { name = "sqlmodel", marker = "extra == 'admin'" }, { name = "sqlmodel", marker = "extra == 'sandbox-actor'" }, { name = "swebench", marker = "extra == 'builder'" }, @@ -4816,20 +4903,24 @@ wheels = [ [[package]] name = "scaffoldhub" -version = "0.1.0.dev4" +version = "0.1.0.dev11" source = { registry = "https://artlab.alibaba-inc.com/1/pypi/simple" } dependencies = [ + { name = "circus", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "fastapi", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "httpx", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "jsonschema", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "mcp", extra = ["cli"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "psycopg", extra = ["binary"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "python-dotenv", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "requests", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "rl-rock", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, { name = "slack-sdk", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "uvicorn", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, ] -sdist = { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev4.tar.gz", hash = "sha256:1b9d9b712b7f369b73d4650a37a1a972d454d2cde24c38b8c16acbcc5287f92a" } +sdist = { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev11.tar.gz", hash = "sha256:6a8bb9ffec9be6317104484e1d6b87c804482f3b3aa4ef4dcf0893f2924191d2" } wheels = [ - { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev4-py3-none-any.whl", hash = "sha256:36d0e8d4026d6b12850d7c4d8c7c689af1fff644f395e58625cfea54fe515207" }, + { url = "https://artlab.alibaba-inc.com/1/pypi/simple/scaffoldhub/scaffoldhub-0.1.0.dev11-py3-none-any.whl", hash = "sha256:0cacc40a3d1119972544162639e7c1b19b024d791cd1551a6095f83edeae65be" }, ] [[package]] From 5d9535a45767ffc3ef7b5574823ff6629f7b548a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 16:02:47 +0800 Subject: [PATCH 214/226] chore: bump package version to 1.10.0.dev12 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 78c225f2b2..b96e9271c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev11" +version = "1.10.0.dev12" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 0785748a6eea9c030a97478e7c4b40a0ae4cf360 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:16:08 +0800 Subject: [PATCH 215/226] docs: design McpEnv servers validation --- ...26-07-03-mcp-env-require-servers-design.md | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-03-mcp-env-require-servers-design.md diff --git a/docs/superpowers/specs/2026-07-03-mcp-env-require-servers-design.md b/docs/superpowers/specs/2026-07-03-mcp-env-require-servers-design.md new file mode 100644 index 0000000000..70027e18af --- /dev/null +++ b/docs/superpowers/specs/2026-07-03-mcp-env-require-servers-design.md @@ -0,0 +1,100 @@ +# McpEnv 非空 servers 构造校验设计 + +## 背景 + +`rock.sdk.mcp.McpEnv` 是 ROCK MCP SDK 的 facade，负责解析 MCP server 配置、 +启动 ROCK sandbox、写入 `/app/mcp-servers.json`、暴露 SSE URL，并代理 +ScaffoldHub lifecycle 的 `init/dump/reset/release`。 + +当前 `McpEnv.__init__(servers=None)` 会把 `None` 转成 `{}`，因此 `McpEnv()` +和 `McpEnv(servers={})` 都可以构造成功。单测中也有用例覆盖“无 servers 时允许 +空数据初始化，但未启动前不能获取 URL”的行为。 + +这个行为容易掩盖调用方误用。`McpEnv` 的主要目的就是运行一组 MCP servers。 +当没有任何 server 配置时，继续构造一个空环境没有实际运行价值，后续甚至可能 +启动一个不包含 MCP server 的 sandbox。更清晰的契约是在构造期直接失败。 + +## 目标 + +- `McpEnv()`、`McpEnv(servers=None)` 和 `McpEnv(servers={})` 都在构造期报错。 +- 非 dict 类型的 `servers` 仍然保持类型错误语义。 +- 合法的非空 dict 配置保持现有行为不变。 +- 错误信息清晰指向 `servers` 必须是非空 dict。 +- 单测覆盖无参、`None`、空 dict、非 dict 和合法非空 dict 的构造行为。 +- 文档明确 `servers` 是必需的非空配置。 + +## 非目标 + +- 不改变 `McpEnv.start()` 的运行顺序。 +- 不改变 auth 占位符解析、lifecycle 创建、`init/dump/reset/release` 语义。 +- 不改变 `RockRuntime` 的接口或空 server health check 行为。 +- 不新增兼容期 warning。 + +## 方案选择 + +采用构造期严格校验。 + +备选方案包括在 `start()` 时校验，或先 warning 后续版本再报错。启动期校验会让 +错误离根因更远，并继续保留空环境的半有效状态。warning 适合必须兼容旧调用方的 +公开 API，但这里空 `servers` 对 `McpEnv` 来说是无实际运行意义的配置，直接失败 +更符合 SDK facade 的职责边界。 + +## 详细设计 + +`McpEnv.__init__` 的签名可以暂时保持： + +```python +def __init__(self, servers: dict | None = None): + ... +``` + +保留 `None` 类型是为了让调用方得到清晰的 `ValueError`，而不是 Python 缺参时的 +默认 `TypeError`。构造逻辑调整为： + +1. 如果 `servers` 不是 dict： + - 当 `servers is None` 时抛 `ValueError("servers must be a non-empty dict")`。 + - 其他非 dict 类型抛 `TypeError("servers must be a dict")`。 +2. 如果 `servers` 是空 dict，抛 `ValueError("servers must be a non-empty dict")`。 +3. 如果 `servers` 是非空 dict，继续按现有逻辑深拷贝、创建 auth provider、 + 创建 lifecycle、初始化 `RockRuntime`。 + +错误类型区分的目的： + +- `servers=[]`、`servers="slack"` 这类输入是类型错误，继续使用 `TypeError`。 +- `McpEnv()`、`servers=None`、`servers={}` 代表缺少必需配置，使用 `ValueError` + 更准确。 + +## 测试设计 + +更新 `tests/unit/sdk/mcp/test_mcp_env.py`： + +- 保留 `test_mcp_env_constructor_requires_servers_dict` 覆盖非 dict 类型抛 + `TypeError("servers must be a dict")`。 +- 新增或调整构造校验测试，覆盖： + - `McpEnv()` 抛 `ValueError("servers must be a non-empty dict")`； + - `McpEnv(servers=None)` 抛同样的 `ValueError`； + - `McpEnv(servers={})` 抛同样的 `ValueError`。 +- 删除或改写当前允许无 servers 初始化的测试。 +- 现有带 `servers={"slack": ...}` 的 `init/dump/reset/release/start` 测试保持通过。 +- 当前部分测试使用 `McpEnv()` 只为了测试 `init` 输入类型；这些测试应改为传入 + 最小合法 server 配置，避免依赖空环境。 + +推荐验证命令： + +```bash +uv run pytest tests/unit/sdk/mcp/test_mcp_env.py -q +``` + +## 文档设计 + +在 MCP SDK 文档中补充一句：`servers` 必须是非空 dict，顶层 key 是 MCP server +名称或 lifecycle 类型。 + +已有示例已经使用非空 `servers`，无需改动示例结构。 + +## 兼容性影响 + +这是一个有意的行为收紧。依赖 `McpEnv()` 或 `McpEnv(servers={})` 创建空环境的代码 +会在构造期失败。调用方应传入至少一个 MCP server 配置。 + +该变更不会影响正常启动真实 MCP servers 的调用方。 From 7ac02c20479e48376d7509b2726599609fd9a71b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:24:35 +0800 Subject: [PATCH 216/226] fix: require McpEnv servers config --- .../References/Python SDK References/mcp.md | 3 ++ rock/sdk/mcp/mcp_env.py | 11 ++++-- tests/unit/sdk/mcp/test_mcp_env.py | 37 +++++++++---------- 3 files changed, 28 insertions(+), 23 deletions(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index 66d9b3fbef..62bf1d9607 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -3,6 +3,9 @@ `rock.sdk.mcp` provides `McpEnv`, a small SDK facade for running MCP servers inside ROCK sandboxes. +`McpEnv` requires `servers` to be a non-empty `dict`; each top-level key is an +MCP server name or lifecycle type. + ## Installation Install the MCP extra when using ScaffoldHub-backed tool lifecycles: diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index e9faecb877..611a78413c 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -40,16 +40,19 @@ def __init__(self, servers: dict | None = None): Create an uninitialized MCP environment. Args: - servers: MCP server config. Top-level keys are server or lifecycle - types, such as ``slack``. + servers: Non-empty MCP server config. Top-level keys are server or + lifecycle types, such as ``slack``. Raises: - TypeError: Raised when servers is neither dict nor None. + TypeError: Raised when servers is not a dict. + ValueError: Raised when servers is None or empty. """ if servers is None: - servers = {} + raise ValueError("servers must be a non-empty dict") if not isinstance(servers, dict): raise TypeError("servers must be a dict") + if not servers: + raise ValueError("servers must be a non-empty dict") self.running = False self.urls = {} diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index 8455a9f513..6bbf106026 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -214,11 +214,25 @@ def test_mcp_env_init_dump_and_release_with_declared_server(monkeypatch): assert env.resolved_servers == {} -def test_mcp_env_constructor_requires_servers_dict(monkeypatch): +def test_mcp_env_constructor_rejects_missing_or_empty_servers(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + + with pytest.raises(ValueError, match="servers must be a non-empty dict"): + mcp_env.McpEnv() + + with pytest.raises(ValueError, match="servers must be a non-empty dict"): + mcp_env.McpEnv(servers=None) + + with pytest.raises(ValueError, match="servers must be a non-empty dict"): + mcp_env.McpEnv(servers={}) + + +@pytest.mark.parametrize("servers", [[], "slack"]) +def test_mcp_env_constructor_requires_servers_dict(monkeypatch, servers): mcp_env = reload_mcp_env(monkeypatch) with pytest.raises(TypeError, match="servers must be a dict"): - mcp_env.McpEnv(servers=[]) + mcp_env.McpEnv(servers=servers) def test_mcp_env_owns_auth_provider_and_passes_it_to_lifecycle_factory(monkeypatch): @@ -243,7 +257,7 @@ def raise_missing_dependency(): monkeypatch.setattr(module, "_load_scaffoldhub_components", raise_missing_dependency) with pytest.raises(ImportError, match=r"rl-rock\[mcp\]"): - module.McpEnv() + module.McpEnv(servers={"slack": slack_server_config()}) def test_mcp_env_resolves_server_env_placeholders_without_starting_runtime(monkeypatch): @@ -296,24 +310,9 @@ def test_mcp_env_resolution_keeps_placeholders_when_auth_is_unavailable(monkeypa } -def test_mcp_env_init_with_no_servers_allows_empty_data_but_not_urls_before_start(monkeypatch): - mcp_env = reload_mcp_env(monkeypatch) - env = mcp_env.McpEnv() - - env.init({}) - - assert env.dump() == {} - with pytest.raises(RuntimeError, match="McpEnv has not been started"): - env.get_urls() - - asyncio.run(env.release()) - - assert env.is_alive() is False - - def test_mcp_env_init_requires_dict(monkeypatch): mcp_env = reload_mcp_env(monkeypatch) - env = mcp_env.McpEnv() + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) with pytest.raises(TypeError, match="data must be a dict"): env.init([]) From 2d79d4b2b37016abc3185ecca48af8d47b4c09ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:29:13 +0800 Subject: [PATCH 217/226] docs: add McpEnv runtime options design --- ...26-07-03-mcp-env-runtime-options-design.md | 196 ++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md diff --git a/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md b/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md new file mode 100644 index 0000000000..5bd60f309f --- /dev/null +++ b/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md @@ -0,0 +1,196 @@ +# McpEnv Runtime Options Design + +## Purpose + +`McpEnv` currently constructs `RockRuntime()` with fixed runtime behavior. +`RockRuntime` already has timeout-related constructor parameters for MCP server +health checks, but `McpEnv` does not expose them. Users need a stable way to +adjust these runtime health-check settings when creating an MCP environment +without expanding `McpEnv.__init__` into a long list of low-level parameters. + +This change introduces a small options object for `RockRuntime` settings and +lets `McpEnv` accept it during construction. + +## Goals + +- Expose all existing `RockRuntime` timeout-related health-check settings + through one object. +- Keep current behavior unchanged when users do not pass options. +- Let users create `RockRuntimeOptions()`, mutate only the fields they need, + and pass the object to `McpEnv`. +- Lock the options at `McpEnv` construction time so later mutations to the + caller-owned object do not change an existing environment. +- Preserve direct `RockRuntime(...)` construction compatibility. + +## Non-Goals + +- Do not introduce a new total health-check deadline parameter. +- Do not change the existing retry-loop semantics in `_health_check()`. +- Do not include sandbox startup, MCP config writing, `before_launch`, or server + launch time in any health-check timeout setting. +- Do not expose unrelated ROCK sandbox configuration through this object. + +## Public API + +Add a non-frozen dataclass: + +```python +@dataclass +class RockRuntimeOptions: + health_check_retries: int = 10 + health_check_interval_seconds: float = 10.0 + http_timeout_seconds: float = 10.0 +``` + +Export it from `rock.sdk.mcp`: + +```python +from rock.sdk.mcp import McpEnv, RockRuntimeOptions +``` + +Usage: + +```python +options = RockRuntimeOptions() +options.health_check_retries = 12 +options.health_check_interval_seconds = 5.0 + +env = McpEnv( + servers={"calculator": {"command": "uvx", "args": ["mcp-server-calculator==0.2.0"]}}, + runtime_options=options, +) + +options.health_check_retries = 1 # Does not affect env. +``` + +`McpEnv.__init__` accepts: + +```python +def __init__( + self, + servers: dict | None = None, + runtime_options: RockRuntimeOptions | None = None, +): + ... +``` + +When `runtime_options` is `None`, `McpEnv` uses `RockRuntimeOptions()` and +therefore keeps current default behavior. + +## Option Snapshot Semantics + +`RockRuntimeOptions` remains mutable for ergonomic caller-side setup. Once +passed to `McpEnv`, the environment snapshots it immediately and does not hold a +reference to the caller-owned object. + +This means: + +- Mutating options before `McpEnv(...)` affects that environment. +- Mutating the same options object after `McpEnv(...)` does not affect that + environment. +- Each `McpEnv` instance owns independent runtime options. + +Implementation snapshots by constructing a new `RockRuntimeOptions` from the +input fields. The dataclass only contains scalar values, and an explicit +field-by-field copy makes the API boundary clear. + +## RockRuntime Construction + +`RockRuntime` accepts an `options` object while keeping the existing flat +parameters for compatibility: + +```python +class RockRuntime: + def __init__( + self, + config: RockRuntimeConfig | None = None, + *, + options: RockRuntimeOptions | None = None, + health_check_retries: int | None = None, + health_check_interval_seconds: float | None = None, + http_timeout_seconds: float | None = None, + ): + ... +``` + +Normalization rules: + +- Start from a snapshot of `options` if provided, otherwise + `RockRuntimeOptions()`. +- If any flat parameter is provided, it overrides the corresponding option + field. This preserves existing direct `RockRuntime(...)` use cases and lets + tests continue to customize one value directly. +- Store the final snapshot as `self.options`. +- Keep existing runtime attributes `self.health_check_retries`, + `self.health_check_interval_seconds`, and `self.http_timeout_seconds` as + assigned aliases copied from `self.options`. This preserves direct attribute + access for tests or internal callers while keeping option normalization in one + place. + +`McpEnv` should construct: + +```python +self._rock_runtime = RockRuntime(options=runtime_options_snapshot) +``` + +## Validation + +Validate normalized runtime options during `RockRuntime` construction: + +- `health_check_retries` must be an integer greater than or equal to `1`. +- `health_check_interval_seconds` must be a number greater than `0`. +- `http_timeout_seconds` must be a number greater than `0`. + +Invalid values raise `RockRuntimeConfigError` with a field-specific message, +for example: + +- `health_check_retries must be >= 1` +- `health_check_interval_seconds must be > 0` +- `http_timeout_seconds must be > 0` + +Validation happens before sandbox startup. `McpEnv(...)` will fail fast because +it constructs `RockRuntime` during initialization. + +## Data Flow + +1. Caller creates and optionally mutates `RockRuntimeOptions`. +2. Caller passes it to `McpEnv`. +3. `McpEnv` snapshots the options and constructs `RockRuntime(options=...)`. +4. `McpEnv.start()` resolves server configuration as it does today. +5. `RockRuntime.start()` launches the sandbox and eventually calls + `_health_check()`. +6. `_health_check()` reads the normalized option values for retry count, + interval, and per-request HTTP timeout. + +No lifecycle, auth, server-config resolution, sandbox injection, or release +behavior changes. + +## Testing + +Add unit coverage for: + +- `RockRuntimeOptions()` exposes current defaults. +- `McpEnv(runtime_options=...)` passes option values into its `RockRuntime`. +- Mutating the original options after `McpEnv(...)` does not affect + `env._rock_runtime.options`. +- `RockRuntime(options=..., health_check_retries=...)` applies flat parameter + overrides for direct-construction compatibility. +- Invalid option values raise `RockRuntimeConfigError`. +- Existing `McpEnv(servers=...)` and `RockRuntime(...)` tests still pass without + API changes. + +Integration tests do not need to exercise non-default timing because this is +configuration plumbing and runtime validation. Existing real ROCK MCP +integration coverage is sufficient for default behavior. + +## Documentation + +Update MCP SDK documentation with a small example showing: + +```python +options = RockRuntimeOptions() +options.health_check_retries = 12 +env = McpEnv(servers=servers, runtime_options=options) +``` + +Mention that options are snapshotted when passed to `McpEnv`. From 154a5bf21e90c065ed87df692b355f44d8c9669a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:33:39 +0800 Subject: [PATCH 218/226] docs: translate McpEnv runtime options design --- ...26-07-03-mcp-env-runtime-options-design.md | 178 ++++++++---------- 1 file changed, 81 insertions(+), 97 deletions(-) diff --git a/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md b/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md index 5bd60f309f..c4cb355592 100644 --- a/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md +++ b/docs/superpowers/specs/2026-07-03-mcp-env-runtime-options-design.md @@ -1,38 +1,36 @@ -# McpEnv Runtime Options Design +# McpEnv Runtime Options 设计 -## Purpose +## 目的 -`McpEnv` currently constructs `RockRuntime()` with fixed runtime behavior. -`RockRuntime` already has timeout-related constructor parameters for MCP server -health checks, but `McpEnv` does not expose them. Users need a stable way to -adjust these runtime health-check settings when creating an MCP environment -without expanding `McpEnv.__init__` into a long list of low-level parameters. +`McpEnv` 当前会固定构造 `RockRuntime()`，用户无法通过 `McpEnv` +调整 runtime 行为。`RockRuntime` 已经有 MCP server 健康检查相关的超时构造参数， +但这些参数没有暴露到 `McpEnv`。 -This change introduces a small options object for `RockRuntime` settings and -lets `McpEnv` accept it during construction. +本设计新增一个小型 options 对象，用来承载 `RockRuntime` 的健康检查配置， +并让 `McpEnv` 在构造时接收该对象。这样既能暴露配置能力，也避免把 +`McpEnv.__init__` 扩展成一长串低层参数。 -## Goals +## 目标 -- Expose all existing `RockRuntime` timeout-related health-check settings - through one object. -- Keep current behavior unchanged when users do not pass options. -- Let users create `RockRuntimeOptions()`, mutate only the fields they need, - and pass the object to `McpEnv`. -- Lock the options at `McpEnv` construction time so later mutations to the - caller-owned object do not change an existing environment. -- Preserve direct `RockRuntime(...)` construction compatibility. +- 通过一个对象暴露现有 `RockRuntime` 的全部健康检查超时相关参数。 +- 用户不传 options 时，保持当前默认行为不变。 +- 用户可以先创建 `RockRuntimeOptions()`，只修改自己关心的字段，再传给 + `McpEnv`。 +- `McpEnv` 构造时锁定 options；调用方后续修改原始 options 对象，不影响已创建的 + environment。 +- 保留直接构造 `RockRuntime(...)` 的兼容性。 -## Non-Goals +## 非目标 -- Do not introduce a new total health-check deadline parameter. -- Do not change the existing retry-loop semantics in `_health_check()`. -- Do not include sandbox startup, MCP config writing, `before_launch`, or server - launch time in any health-check timeout setting. -- Do not expose unrelated ROCK sandbox configuration through this object. +- 不新增“整体健康检查截止时间”参数。 +- 不改变 `_health_check()` 当前基于 retry loop 的语义。 +- 不把 sandbox startup、MCP config 写入、`before_launch` 或 server launch 时间计入任何 + 健康检查 timeout 参数。 +- 不通过该对象暴露无关的 ROCK sandbox 配置。 -## Public API +## 公共 API -Add a non-frozen dataclass: +新增一个非 frozen dataclass： ```python @dataclass @@ -42,13 +40,13 @@ class RockRuntimeOptions: http_timeout_seconds: float = 10.0 ``` -Export it from `rock.sdk.mcp`: +从 `rock.sdk.mcp` 导出： ```python from rock.sdk.mcp import McpEnv, RockRuntimeOptions ``` -Usage: +用法示例： ```python options = RockRuntimeOptions() @@ -60,10 +58,10 @@ env = McpEnv( runtime_options=options, ) -options.health_check_retries = 1 # Does not affect env. +options.health_check_retries = 1 # 不影响 env。 ``` -`McpEnv.__init__` accepts: +`McpEnv.__init__` 接收： ```python def __init__( @@ -74,30 +72,26 @@ def __init__( ... ``` -When `runtime_options` is `None`, `McpEnv` uses `RockRuntimeOptions()` and -therefore keeps current default behavior. +当 `runtime_options` 为 `None` 时，`McpEnv` 使用 `RockRuntimeOptions()`， +因此默认行为与当前实现一致。 -## Option Snapshot Semantics +## Options 快照语义 -`RockRuntimeOptions` remains mutable for ergonomic caller-side setup. Once -passed to `McpEnv`, the environment snapshots it immediately and does not hold a -reference to the caller-owned object. +`RockRuntimeOptions` 保持可变，方便调用方先构造对象，再按需修改字段。 +对象一旦传入 `McpEnv`，`McpEnv` 会立即创建一份快照，不持有调用方原始对象引用。 -This means: +具体语义： -- Mutating options before `McpEnv(...)` affects that environment. -- Mutating the same options object after `McpEnv(...)` does not affect that - environment. -- Each `McpEnv` instance owns independent runtime options. +- 在调用 `McpEnv(...)` 之前修改 options，会影响该 environment。 +- 在调用 `McpEnv(...)` 之后继续修改同一个 options 对象，不影响该 environment。 +- 每个 `McpEnv` 实例拥有独立的 runtime options。 -Implementation snapshots by constructing a new `RockRuntimeOptions` from the -input fields. The dataclass only contains scalar values, and an explicit -field-by-field copy makes the API boundary clear. +实现时通过读取输入对象字段并构造新的 `RockRuntimeOptions` 来创建快照。 +该 dataclass 只包含标量字段，显式逐字段复制可以让 API 边界更清楚。 -## RockRuntime Construction +## RockRuntime 构造 -`RockRuntime` accepts an `options` object while keeping the existing flat -parameters for compatibility: +`RockRuntime` 接收一个 `options` 对象，同时保留现有平铺参数以兼容直接构造用法： ```python class RockRuntime: @@ -113,79 +107,69 @@ class RockRuntime: ... ``` -Normalization rules: +归一化规则： -- Start from a snapshot of `options` if provided, otherwise - `RockRuntimeOptions()`. -- If any flat parameter is provided, it overrides the corresponding option - field. This preserves existing direct `RockRuntime(...)` use cases and lets - tests continue to customize one value directly. -- Store the final snapshot as `self.options`. -- Keep existing runtime attributes `self.health_check_retries`, - `self.health_check_interval_seconds`, and `self.http_timeout_seconds` as - assigned aliases copied from `self.options`. This preserves direct attribute - access for tests or internal callers while keeping option normalization in one - place. +- 如果传入 `options`，先基于它创建一份快照；否则使用 `RockRuntimeOptions()`。 +- 如果同时传入任意平铺参数，则平铺参数覆盖对应的 option 字段。 + 这样可以保留既有 `RockRuntime(...)` 直接使用方式，也方便测试继续单独覆盖某个值。 +- 将最终快照保存为 `self.options`。 +- 保留现有 runtime 属性 `self.health_check_retries`、 + `self.health_check_interval_seconds` 和 `self.http_timeout_seconds`，这些属性从 + `self.options` 复制赋值。这样既保留测试或内部调用方可能依赖的直接属性访问， + 又把 options 归一化集中在一个地方。 -`McpEnv` should construct: +`McpEnv` 构造 runtime 时使用： ```python self._rock_runtime = RockRuntime(options=runtime_options_snapshot) ``` -## Validation +## 校验 -Validate normalized runtime options during `RockRuntime` construction: +在 `RockRuntime` 构造期间校验归一化后的 runtime options： -- `health_check_retries` must be an integer greater than or equal to `1`. -- `health_check_interval_seconds` must be a number greater than `0`. -- `http_timeout_seconds` must be a number greater than `0`. +- `health_check_retries` 必须是大于等于 `1` 的整数。 +- `health_check_interval_seconds` 必须是大于 `0` 的数字。 +- `http_timeout_seconds` 必须是大于 `0` 的数字。 -Invalid values raise `RockRuntimeConfigError` with a field-specific message, -for example: +非法值抛出 `RockRuntimeConfigError`，错误信息应包含具体字段，例如： - `health_check_retries must be >= 1` - `health_check_interval_seconds must be > 0` - `http_timeout_seconds must be > 0` -Validation happens before sandbox startup. `McpEnv(...)` will fail fast because -it constructs `RockRuntime` during initialization. +校验发生在 sandbox 启动之前。因为 `McpEnv(...)` 会在初始化阶段构造 +`RockRuntime`，所以非法 options 会快速失败。 -## Data Flow +## 数据流 -1. Caller creates and optionally mutates `RockRuntimeOptions`. -2. Caller passes it to `McpEnv`. -3. `McpEnv` snapshots the options and constructs `RockRuntime(options=...)`. -4. `McpEnv.start()` resolves server configuration as it does today. -5. `RockRuntime.start()` launches the sandbox and eventually calls - `_health_check()`. -6. `_health_check()` reads the normalized option values for retry count, - interval, and per-request HTTP timeout. +1. 调用方创建并按需修改 `RockRuntimeOptions`。 +2. 调用方将 options 传给 `McpEnv`。 +3. `McpEnv` 创建 options 快照，并构造 `RockRuntime(options=...)`。 +4. `McpEnv.start()` 按现有逻辑解析 server 配置。 +5. `RockRuntime.start()` 启动 sandbox，并最终调用 `_health_check()`。 +6. `_health_check()` 读取归一化后的 retry 次数、间隔和单次 HTTP 请求 timeout。 -No lifecycle, auth, server-config resolution, sandbox injection, or release -behavior changes. +生命周期、auth、server config 解析、sandbox 注入和 release 行为均不改变。 -## Testing +## 测试 -Add unit coverage for: +新增单元测试覆盖： -- `RockRuntimeOptions()` exposes current defaults. -- `McpEnv(runtime_options=...)` passes option values into its `RockRuntime`. -- Mutating the original options after `McpEnv(...)` does not affect - `env._rock_runtime.options`. -- `RockRuntime(options=..., health_check_retries=...)` applies flat parameter - overrides for direct-construction compatibility. -- Invalid option values raise `RockRuntimeConfigError`. -- Existing `McpEnv(servers=...)` and `RockRuntime(...)` tests still pass without - API changes. +- `RockRuntimeOptions()` 暴露当前默认值。 +- `McpEnv(runtime_options=...)` 将 option 值传入其 `RockRuntime`。 +- `McpEnv(...)` 之后修改原始 options，不影响 `env._rock_runtime.options`。 +- `RockRuntime(options=..., health_check_retries=...)` 会应用平铺参数覆盖， + 保持直接构造兼容性。 +- 非法 option 值抛出 `RockRuntimeConfigError`。 +- 现有 `McpEnv(servers=...)` 和 `RockRuntime(...)` 测试无需 API 调整即可继续通过。 -Integration tests do not need to exercise non-default timing because this is -configuration plumbing and runtime validation. Existing real ROCK MCP -integration coverage is sufficient for default behavior. +集成测试不需要覆盖非默认 timing。该改动主要是配置传递和 runtime 校验； +现有真实 ROCK MCP 集成测试足以覆盖默认行为。 -## Documentation +## 文档 -Update MCP SDK documentation with a small example showing: +更新 MCP SDK 文档，增加一个简短示例： ```python options = RockRuntimeOptions() @@ -193,4 +177,4 @@ options.health_check_retries = 12 env = McpEnv(servers=servers, runtime_options=options) ``` -Mention that options are snapshotted when passed to `McpEnv`. +文档中需要说明：options 传入 `McpEnv` 时会被快照锁定。 From 1bfeb5743b0598c8b3e0915cf79f18476775c621 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:49:39 +0800 Subject: [PATCH 219/226] feat: add MCP runtime options --- rock/sdk/mcp/rock_runtime.py | 52 ++++++++++++++++--- tests/unit/sdk/mcp/test_rock_runtime.py | 67 +++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 6 deletions(-) diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index f860f9a8b4..9059e80192 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -76,19 +76,59 @@ def from_env(cls) -> RockRuntimeConfig: ) +@dataclass +class RockRuntimeOptions: + health_check_retries: int = 10 + health_check_interval_seconds: float = 10.0 + http_timeout_seconds: float = 10.0 + + +def _snapshot_runtime_options(options: RockRuntimeOptions | None) -> RockRuntimeOptions: + if options is None: + return RockRuntimeOptions() + return RockRuntimeOptions( + health_check_retries=options.health_check_retries, + health_check_interval_seconds=options.health_check_interval_seconds, + http_timeout_seconds=options.http_timeout_seconds, + ) + + +def _validate_runtime_options(options: RockRuntimeOptions) -> None: + if not isinstance(options.health_check_retries, int) or options.health_check_retries < 1: + raise RockRuntimeConfigError("health_check_retries must be >= 1") + if ( + not isinstance(options.health_check_interval_seconds, int | float) + or options.health_check_interval_seconds <= 0 + ): + raise RockRuntimeConfigError("health_check_interval_seconds must be > 0") + if not isinstance(options.http_timeout_seconds, int | float) or options.http_timeout_seconds <= 0: + raise RockRuntimeConfigError("http_timeout_seconds must be > 0") + + class RockRuntime: def __init__( self, config: RockRuntimeConfig | None = None, *, - health_check_retries: int = 10, - health_check_interval_seconds: float = 10.0, - http_timeout_seconds: float = 10.0, + options: RockRuntimeOptions | None = None, + health_check_retries: int | None = None, + health_check_interval_seconds: float | None = None, + http_timeout_seconds: float | None = None, ): + runtime_options = _snapshot_runtime_options(options) + if health_check_retries is not None: + runtime_options.health_check_retries = health_check_retries + if health_check_interval_seconds is not None: + runtime_options.health_check_interval_seconds = health_check_interval_seconds + if http_timeout_seconds is not None: + runtime_options.http_timeout_seconds = http_timeout_seconds + _validate_runtime_options(runtime_options) + self.config = config - self.health_check_retries = health_check_retries - self.health_check_interval_seconds = health_check_interval_seconds - self.http_timeout_seconds = http_timeout_seconds + self.options = runtime_options + self.health_check_retries = runtime_options.health_check_retries + self.health_check_interval_seconds = runtime_options.health_check_interval_seconds + self.http_timeout_seconds = runtime_options.http_timeout_seconds self._sandbox: Sandbox | None = None self._sandbox_id: str | None = None self._started = False diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index 00bebaa227..3e0071468b 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -49,6 +49,73 @@ def test_rock_runtime_config_reads_defaults_and_numeric_values(monkeypatch): assert config.auto_clear_seconds == 3600 +def test_rock_runtime_options_exposes_current_defaults(): + options = rock_runtime.RockRuntimeOptions() + + assert options.health_check_retries == 10 + assert options.health_check_interval_seconds == 10.0 + assert options.http_timeout_seconds == 10.0 + + +def test_rock_runtime_snapshots_options_at_construction(): + options = rock_runtime.RockRuntimeOptions( + health_check_retries=12, + health_check_interval_seconds=5.0, + http_timeout_seconds=2.5, + ) + + runtime = RockRuntime(options=options) + options.health_check_retries = 1 + options.health_check_interval_seconds = 1.0 + options.http_timeout_seconds = 1.0 + + assert runtime.options == rock_runtime.RockRuntimeOptions( + health_check_retries=12, + health_check_interval_seconds=5.0, + http_timeout_seconds=2.5, + ) + assert runtime.health_check_retries == 12 + assert runtime.health_check_interval_seconds == 5.0 + assert runtime.http_timeout_seconds == 2.5 + + +def test_rock_runtime_flat_parameters_override_options(): + runtime = RockRuntime( + options=rock_runtime.RockRuntimeOptions( + health_check_retries=12, + health_check_interval_seconds=5.0, + http_timeout_seconds=2.5, + ), + health_check_retries=3, + http_timeout_seconds=1.5, + ) + + assert runtime.options == rock_runtime.RockRuntimeOptions( + health_check_retries=3, + health_check_interval_seconds=5.0, + http_timeout_seconds=1.5, + ) + assert runtime.health_check_retries == 3 + assert runtime.health_check_interval_seconds == 5.0 + assert runtime.http_timeout_seconds == 1.5 + + +@pytest.mark.parametrize( + ("options_kwargs", "message"), + [ + ({"health_check_retries": 0}, "health_check_retries must be >= 1"), + ({"health_check_retries": 1.5}, "health_check_retries must be >= 1"), + ({"health_check_interval_seconds": 0}, "health_check_interval_seconds must be > 0"), + ({"http_timeout_seconds": 0}, "http_timeout_seconds must be > 0"), + ], +) +def test_rock_runtime_rejects_invalid_options(options_kwargs, message): + options = rock_runtime.RockRuntimeOptions(**options_kwargs) + + with pytest.raises(RockRuntimeConfigError, match=message): + RockRuntime(options=options) + + def test_rock_runtime_builds_server_urls_from_sandbox_id(monkeypatch): monkeypatch.setenv("ROCK_API_KEY", "rock-key") monkeypatch.setenv("ROCK_USER_ID", "user-001") From f9511add7df24b7b322bc62470c3237e8a5e5a77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:50:55 +0800 Subject: [PATCH 220/226] feat: expose MCP runtime options --- rock/sdk/mcp/__init__.py | 3 ++- rock/sdk/mcp/mcp_env.py | 18 ++++++++++++++--- tests/unit/sdk/mcp/test_mcp_env.py | 31 ++++++++++++++++++++++++++++++ 3 files changed, 48 insertions(+), 4 deletions(-) diff --git a/rock/sdk/mcp/__init__.py b/rock/sdk/mcp/__init__.py index 9c70ec3c7b..a418a8283a 100644 --- a/rock/sdk/mcp/__init__.py +++ b/rock/sdk/mcp/__init__.py @@ -1,3 +1,4 @@ from rock.sdk.mcp.mcp_env import McpEnv +from rock.sdk.mcp.rock_runtime import RockRuntimeOptions -__all__ = ["McpEnv"] +__all__ = ["McpEnv", "RockRuntimeOptions"] diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index 611a78413c..ecee319e75 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -5,7 +5,7 @@ from copy import deepcopy from typing import Any -from rock.sdk.mcp.rock_runtime import BeforeLaunchHook, RockRuntime +from rock.sdk.mcp.rock_runtime import BeforeLaunchHook, RockRuntime, RockRuntimeOptions from rock.sdk.sandbox.client import Sandbox logger = logging.getLogger(__name__) @@ -26,6 +26,16 @@ def _load_scaffoldhub_components(): return AuthProvider, DataLifecycleFactory, SandboxAware +def _snapshot_runtime_options(options: RockRuntimeOptions | None) -> RockRuntimeOptions: + if options is None: + return RockRuntimeOptions() + return RockRuntimeOptions( + health_check_retries=options.health_check_retries, + health_check_interval_seconds=options.health_check_interval_seconds, + http_timeout_seconds=options.http_timeout_seconds, + ) + + class McpEnv: """ MCP environment manager. @@ -35,13 +45,15 @@ class McpEnv: releases runtime resources. """ - def __init__(self, servers: dict | None = None): + def __init__(self, servers: dict | None = None, runtime_options: RockRuntimeOptions | None = None): """ Create an uninitialized MCP environment. Args: servers: Non-empty MCP server config. Top-level keys are server or lifecycle types, such as ``slack``. + runtime_options: Optional ROCK runtime health-check options. The + values are snapshotted during construction. Raises: TypeError: Raised when servers is not a dict. @@ -63,7 +75,7 @@ def __init__(self, servers: dict | None = None): self.data_lifecycle_factory = data_lifecycle_factory_class(auth_provider=self.auth_provider) self.sandbox_aware_class = sandbox_aware_class self.data_lifecycles: dict[str, Any] = {} - self._rock_runtime = RockRuntime() + self._rock_runtime = RockRuntime(options=_snapshot_runtime_options(runtime_options)) for lifecycle_type in self.servers: if self.data_lifecycle_factory.supports(lifecycle_type): diff --git a/tests/unit/sdk/mcp/test_mcp_env.py b/tests/unit/sdk/mcp/test_mcp_env.py index 6bbf106026..f666829f0f 100644 --- a/tests/unit/sdk/mcp/test_mcp_env.py +++ b/tests/unit/sdk/mcp/test_mcp_env.py @@ -609,3 +609,34 @@ def test_mcp_env_exposes_raw_sandbox_property(monkeypatch): mcp_env = reload_mcp_env(monkeypatch) assert isinstance(mcp_env.McpEnv.sandbox, property) + + +def test_mcp_env_accepts_runtime_options_and_passes_snapshot_to_rock_runtime(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + options = mcp_env.RockRuntimeOptions( + health_check_retries=12, + health_check_interval_seconds=5.0, + http_timeout_seconds=2.5, + ) + + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}, runtime_options=options) + options.health_check_retries = 1 + options.health_check_interval_seconds = 1.0 + options.http_timeout_seconds = 1.0 + + assert env._rock_runtime.options == mcp_env.RockRuntimeOptions( + health_check_retries=12, + health_check_interval_seconds=5.0, + http_timeout_seconds=2.5, + ) + assert env._rock_runtime.health_check_retries == 12 + assert env._rock_runtime.health_check_interval_seconds == 5.0 + assert env._rock_runtime.http_timeout_seconds == 2.5 + + +def test_mcp_env_uses_default_runtime_options_when_not_provided(monkeypatch): + mcp_env = reload_mcp_env(monkeypatch) + + env = mcp_env.McpEnv(servers={"slack": slack_server_config()}) + + assert env._rock_runtime.options == mcp_env.RockRuntimeOptions() From 156fa52ec86eddce5bd8b62c4f519f12d859db50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:51:30 +0800 Subject: [PATCH 221/226] docs: document MCP runtime options --- .../References/Python SDK References/mcp.md | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md index 62bf1d9607..f8e0fbbb0b 100644 --- a/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md +++ b/docs/versioned_docs/version-1.8.x/References/Python SDK References/mcp.md @@ -23,7 +23,7 @@ ScaffoldHub package is Python 3.11+ and ROCK officially supports Python 3.10 to ```python import asyncio -from rock.sdk.mcp import McpEnv +from rock.sdk.mcp import McpEnv, RockRuntimeOptions async def main(): @@ -47,6 +47,32 @@ async def main(): asyncio.run(main()) ``` +## Runtime Options + +`McpEnv` accepts `RockRuntimeOptions` for ROCK runtime health-check settings. +Create the object, adjust only the fields you need, and pass it to `McpEnv`: + +```python +options = RockRuntimeOptions() +options.health_check_retries = 12 +options.health_check_interval_seconds = 5.0 +options.http_timeout_seconds = 2.5 + +env = McpEnv( + servers={ + "calculator": { + "command": "uvx", + "args": ["mcp-server-calculator==0.2.0"], + } + }, + runtime_options=options, +) +``` + +`McpEnv` snapshots these values during construction. Mutating the same +`RockRuntimeOptions` object after creating `McpEnv` does not change that +environment. + `release()` is the cleanup boundary for both ROCK runtime resources and ScaffoldHub auth leases. Keep it in a `finally` block and call it even if `start()` fails before `env.is_alive()` becomes true. If auth lease release From bc6f472e292891556203ca00f0913dc334feb48e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:52:51 +0800 Subject: [PATCH 222/226] style: format MCP runtime options --- rock/sdk/mcp/rock_runtime.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index 9059e80192..8cb4b862d4 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -96,10 +96,7 @@ def _snapshot_runtime_options(options: RockRuntimeOptions | None) -> RockRuntime def _validate_runtime_options(options: RockRuntimeOptions) -> None: if not isinstance(options.health_check_retries, int) or options.health_check_retries < 1: raise RockRuntimeConfigError("health_check_retries must be >= 1") - if ( - not isinstance(options.health_check_interval_seconds, int | float) - or options.health_check_interval_seconds <= 0 - ): + if not isinstance(options.health_check_interval_seconds, int | float) or options.health_check_interval_seconds <= 0: raise RockRuntimeConfigError("health_check_interval_seconds must be > 0") if not isinstance(options.http_timeout_seconds, int | float) or options.http_timeout_seconds <= 0: raise RockRuntimeConfigError("http_timeout_seconds must be > 0") From 69935484f8cf96dfbf08b9bb488d31766dbaf649 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 17:59:09 +0800 Subject: [PATCH 223/226] fix: reject invalid MCP runtime option values --- rock/sdk/mcp/rock_runtime.py | 17 ++++++++++++++--- tests/unit/sdk/mcp/test_rock_runtime.py | 20 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/rock/sdk/mcp/rock_runtime.py b/rock/sdk/mcp/rock_runtime.py index 8cb4b862d4..b88b89c9dc 100644 --- a/rock/sdk/mcp/rock_runtime.py +++ b/rock/sdk/mcp/rock_runtime.py @@ -4,6 +4,7 @@ import inspect import json import logging +import math import os from collections.abc import Awaitable, Callable, Iterable from dataclasses import dataclass @@ -94,11 +95,21 @@ def _snapshot_runtime_options(options: RockRuntimeOptions | None) -> RockRuntime def _validate_runtime_options(options: RockRuntimeOptions) -> None: - if not isinstance(options.health_check_retries, int) or options.health_check_retries < 1: + if type(options.health_check_retries) is not int or options.health_check_retries < 1: raise RockRuntimeConfigError("health_check_retries must be >= 1") - if not isinstance(options.health_check_interval_seconds, int | float) or options.health_check_interval_seconds <= 0: + if ( + not isinstance(options.health_check_interval_seconds, int | float) + or isinstance(options.health_check_interval_seconds, bool) + or not math.isfinite(options.health_check_interval_seconds) + or options.health_check_interval_seconds <= 0 + ): raise RockRuntimeConfigError("health_check_interval_seconds must be > 0") - if not isinstance(options.http_timeout_seconds, int | float) or options.http_timeout_seconds <= 0: + if ( + not isinstance(options.http_timeout_seconds, int | float) + or isinstance(options.http_timeout_seconds, bool) + or not math.isfinite(options.http_timeout_seconds) + or options.http_timeout_seconds <= 0 + ): raise RockRuntimeConfigError("http_timeout_seconds must be > 0") diff --git a/tests/unit/sdk/mcp/test_rock_runtime.py b/tests/unit/sdk/mcp/test_rock_runtime.py index 3e0071468b..7d888440b8 100644 --- a/tests/unit/sdk/mcp/test_rock_runtime.py +++ b/tests/unit/sdk/mcp/test_rock_runtime.py @@ -105,8 +105,13 @@ def test_rock_runtime_flat_parameters_override_options(): [ ({"health_check_retries": 0}, "health_check_retries must be >= 1"), ({"health_check_retries": 1.5}, "health_check_retries must be >= 1"), + ({"health_check_retries": True}, "health_check_retries must be >= 1"), ({"health_check_interval_seconds": 0}, "health_check_interval_seconds must be > 0"), + ({"health_check_interval_seconds": True}, "health_check_interval_seconds must be > 0"), + ({"health_check_interval_seconds": float("nan")}, "health_check_interval_seconds must be > 0"), ({"http_timeout_seconds": 0}, "http_timeout_seconds must be > 0"), + ({"http_timeout_seconds": True}, "http_timeout_seconds must be > 0"), + ({"http_timeout_seconds": float("nan")}, "http_timeout_seconds must be > 0"), ], ) def test_rock_runtime_rejects_invalid_options(options_kwargs, message): @@ -116,6 +121,21 @@ def test_rock_runtime_rejects_invalid_options(options_kwargs, message): RockRuntime(options=options) +@pytest.mark.parametrize( + ("runtime_kwargs", "message"), + [ + ({"health_check_retries": True}, "health_check_retries must be >= 1"), + ({"health_check_interval_seconds": True}, "health_check_interval_seconds must be > 0"), + ({"health_check_interval_seconds": float("nan")}, "health_check_interval_seconds must be > 0"), + ({"http_timeout_seconds": True}, "http_timeout_seconds must be > 0"), + ({"http_timeout_seconds": float("nan")}, "http_timeout_seconds must be > 0"), + ], +) +def test_rock_runtime_rejects_invalid_flat_option_overrides(runtime_kwargs, message): + with pytest.raises(RockRuntimeConfigError, match=message): + RockRuntime(options=rock_runtime.RockRuntimeOptions(), **runtime_kwargs) + + def test_rock_runtime_builds_server_urls_from_sandbox_id(monkeypatch): monkeypatch.setenv("ROCK_API_KEY", "rock-key") monkeypatch.setenv("ROCK_USER_ID", "user-001") From c266e320e33261cf0a10a3b0ad9928ab6f5f0ea3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Fri, 3 Jul 2026 18:26:57 +0800 Subject: [PATCH 224/226] chore: bump package version to 1.10.0.dev13 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b96e9271c5..5d46f49c2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev12" +version = "1.10.0.dev13" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ From 450c56800af0121ff189703df27034631083ea9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9A=86=E5=AE=87?= Date: Mon, 6 Jul 2026 14:41:46 +0800 Subject: [PATCH 225/226] chore: bump scaffoldhub to 0.1.2 --- pyproject.toml | 4 ++-- tests/unit/sdk/mcp/test_packaging.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5d46f49c2c..8e0249c04b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" authors = [{ name = "chatos@alibaba" }] requires-python = "<4.0,>=3.10" name = "rl-rock" -version = "1.10.0.dev13" +version = "1.10.0.dev14" description = "ROCK-Reinforcement Open Construction Kit" readme = "README.md" dependencies = [ @@ -93,7 +93,7 @@ model-service = [ ] mcp = [ - "scaffoldhub==0.1.0.dev11; python_version >= '3.11' and python_version < '3.13'", + "scaffoldhub==0.1.2; python_version >= '3.11' and python_version < '3.13'", ] diff --git a/tests/unit/sdk/mcp/test_packaging.py b/tests/unit/sdk/mcp/test_packaging.py index 5504d65bc9..9f7c2ab3e0 100644 --- a/tests/unit/sdk/mcp/test_packaging.py +++ b/tests/unit/sdk/mcp/test_packaging.py @@ -11,7 +11,7 @@ def test_mcp_extra_declares_scaffoldhub_dependency(): mcp_dependencies = pyproject["project"]["optional-dependencies"]["mcp"] - assert "scaffoldhub==0.1.0.dev11; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies + assert "scaffoldhub==0.1.2; python_version >= '3.11' and python_version < '3.13'" in mcp_dependencies def test_scaffoldhub_resolves_from_artlab_index(): From eb6522058ccc58fbdbc1b6b65b44a52cb6556082 Mon Sep 17 00:00:00 2001 From: pengda Date: Mon, 6 Jul 2026 15:32:10 +0800 Subject: [PATCH 226/226] =?UTF-8?q?refactor(mcp=5Fenv):=20=E6=94=B9?= =?UTF-8?q?=E8=BF=9B=20before=5Flaunch=20=E9=92=A9=E5=AD=90=E5=A4=84?= =?UTF-8?q?=E7=90=86=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此提交改进了 `_compose_before_launch` 方法中的钩子处理逻辑，确保所有生命周期对象的 `before_launch` 方法被正确调用，并且用户自定义的 `before_launch` 钩子也能按预期执行。 Co-developed-by: Aone Copilot --- rock/sdk/mcp/mcp_env.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/rock/sdk/mcp/mcp_env.py b/rock/sdk/mcp/mcp_env.py index ecee319e75..9dff442799 100644 --- a/rock/sdk/mcp/mcp_env.py +++ b/rock/sdk/mcp/mcp_env.py @@ -209,12 +209,20 @@ async def release(self): def _compose_before_launch(self, before_launch: BeforeLaunchHook | None) -> BeforeLaunchHook: async def composed_before_launch(sandbox: Sandbox) -> None: self._inject_sandbox_into_lifecycles(sandbox) - if before_launch is None: - return - result = before_launch(sandbox) - if inspect.isawaitable(result): - await result + # 自动调用所有 lifecycle 的 before_launch 方法 + for lifecycle in self.data_lifecycles.values(): + hook = getattr(lifecycle, "before_launch", None) + if callable(hook): + result = hook(sandbox) + if inspect.isawaitable(result): + await result + + # 再执行用户传入的 before_launch hook + if before_launch is not None: + result = before_launch(sandbox) + if inspect.isawaitable(result): + await result return composed_before_launch