diff --git a/.gitignore b/.gitignore index 0bdb9efb..19c655a0 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,11 @@ tmp/ # Local Harbor dataset tooling /env/harbor/generate_vulhub_dataset.py + +# Local cybergym env overrides (tracked files kept via skip-worktree) +/env/cybergym/ + +/.safactory-locks/ + +# Local credentials / private config (do not commit) +config.yaml diff --git a/args.py b/args.py index 4d831459..8718f7db 100644 --- a/args.py +++ b/args.py @@ -138,7 +138,14 @@ def parse_simulation_args(argv: Sequence[str] | None = None) -> argparse.Namespa "--gateway-close-timeout-s", type=float, default=120.0, - help="Total timeout for polling gateway session close completion.", + help=( + "HTTP timeout for gateway session close requests. Must exceed the" + " gateway's drain_timeout_s (default 30s): the close endpoint blocks" + " up to drain_timeout_s waiting for in-flight LLM requests to finish," + " so a runner timeout shorter than drain_timeout_s abandons the close" + " before the gateway responds, leaving the session unsealed" + " (is_terminal=0) and orphaning the rollout group." + ), ) parser.add_argument( "--gateway-close-retries", diff --git a/clusters/rjob_cluster.py b/clusters/rjob_cluster.py index 558e3d39..0261a5e1 100644 --- a/clusters/rjob_cluster.py +++ b/clusters/rjob_cluster.py @@ -24,6 +24,9 @@ _DEFAULT_RUNNER_CONTAINER_PATH = "/tmp/safactory-openclaw-runner.mjs" _DEFAULT_RUN_COMMAND = f"node {_DEFAULT_RUNNER_CONTAINER_PATH}" +# RJob task ids are "-" and must match the cluster regex +# ^[a-zA-Z0-9][-a-zA-Z0-9]{1,61}[a-zA-Z0-9]$ — dots are NOT allowed, so strip them +# (replace with "-") rather than only allowing alnum + "." + "-". _INVALID_NAME_CHARS = re.compile(r"[^a-z0-9-]+") _MAX_RJOB_NAME_LEN = 49 _MAX_RJOB_AGENT_NAME_LEN = 12 @@ -279,6 +282,10 @@ async def wait_terminal( trace.update_context(rjob_submit_to_starting_ms=submit_to_starting_ms) if status == "Running" and submit_to_running_ms is None: submit_to_running_ms = elapsed_ms + # Absolute epoch seconds at the moment the RJob entered Running. + # Joined with rjob_submit_ts to derive cluster queue time + # (rjob_running_ts - rjob_submit_ts) in the episode record. + rjob_running_ts = time.time() if trace is not None: trace.mark( "rjob_running", @@ -287,7 +294,10 @@ async def wait_terminal( job_name=job_name, submit_to_running_ms=submit_to_running_ms, ) - trace.update_context(rjob_submit_to_running_ms=submit_to_running_ms) + trace.update_context( + rjob_submit_to_running_ms=submit_to_running_ms, + rjob_running_ts=rjob_running_ts, + ) if status != last_status: if trace is not None: trace.mark( diff --git a/config.yaml b/config.yaml deleted file mode 100644 index 682ea758..00000000 --- a/config.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Global RJob connection/auth settings shared by all RJob environments. -# Override with: python launcher.py --mode rjob --rjob-config path/to/config.yaml - -rjob: - cluster_entry: "https://h.pjlab.org.cn" - namespace: "ailab-evobox" - access_key: "" - secret_key: "" - verifyssl: true - retries: 3 - - # Optional but commonly shared across RJob submissions. - charged_group: "evobox_proxy" - name_prefix: safactory - poll_interval_s: 5 - cleanup_on_finish: true - keep_failed_jobs: false - no_packaging: true - auto_delete_duration: 12h - submit_concurrency: 1 diff --git a/config.yaml.example b/config.yaml.example new file mode 100644 index 00000000..f3ddc538 --- /dev/null +++ b/config.yaml.example @@ -0,0 +1,33 @@ +# Global RJob connection/auth settings shared by all RJob environments. +# Override with: python launcher.py --mode rjob --rjob-config path/to/config.yaml +# +# Copy this file to config.yaml and fill in your own credentials. +# Do NOT commit config.yaml (it is gitignored). + +rjob: + cluster_entry: "https://h.pjlab.org.cn" + namespace: "ailab-evobox" + access_key: "" + secret_key: "" + verifyssl: true + retries: 3 + + # Optional but commonly shared across RJob submissions. + charged_group: "evobox_proxy" + # NOTE: intentionally NOT set. rjob_cluster.resolve_gateway_base_url does + # `cfg.get("gateway_base_url") or request.gateway_base_url`; if we hardcode an + # IP here it goes stale the moment the training pod restarts (the runner then + # gets a dead gateway URL and LLM calls hang). Leave this unset so it falls + # back to request.gateway_base_url, which the buffer server fills from + # AIEVOBOX_GATEWAY_BASE_URL (env.rjob.sh: hostname -I of THIS pod). + # gateway_base_url: "http://:8000/v1/sessions" + name_prefix: safactory + poll_interval_s: 5 + # Set to false temporarily so succeeded-but-empty pods are kept for log inspection. + cleanup_on_finish: false + keep_failed_jobs: true + no_packaging: true + auto_delete_duration: 12h + # Raised from 1 so multiple rollout episodes can be submitted in parallel + # (must be >= AIEVOBOX_POOL_SIZE to actually parallelize). + submit_concurrency: 8 diff --git a/core/data_manager/manager.py b/core/data_manager/manager.py index 3eafb20c..8e8095fc 100644 --- a/core/data_manager/manager.py +++ b/core/data_manager/manager.py @@ -314,6 +314,21 @@ async def list_session_steps( checkout_latest=checkout_latest, )) + async def list_terminal_steps_for_sessions( + self, + session_ids: List[str], + *, + job_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Return terminal step rows for a batch of sessions (used by the RL buffer).""" + if not session_ids: + return [] + return await self._strategy.list_session_step_rows(SessionStepQuery( + job_id=job_id or self.job_id or None, + session_ids=tuple(session_ids), + is_terminal=True, + )) + async def update_session_step_rows( self, *, @@ -418,24 +433,7 @@ async def mark_latest_session_completed( async def close(self) -> None: """Close the storage strategy""" await self._strategy.close() - - async def fetch_done_steps_with_context( - self, - after_id: int = 0, - limit: int = 100 - ) -> List[Dict]: - """Fetch completed steps for training data collection""" - if hasattr(self._strategy, 'fetch_done_steps_with_context'): - return await self._strategy.fetch_done_steps_with_context(self.job_id, after_id, limit) - return [] - - async def get_max_step_id(self) -> int: - """Get maximum primary key for pagination""" - if hasattr(self._strategy, 'get_max_step_id'): - return await self._strategy.get_max_step_id(self.job_id) - return 0 - @property def buffer_stats(self) -> Optional[dict]: """Get buffer statistics (SQLite only)""" if hasattr(self._strategy, 'buffer_stats'): diff --git a/core/data_manager/strategy/cloud_strategy_impl.py b/core/data_manager/strategy/cloud_strategy_impl.py index 7d7703e5..00867f4b 100644 --- a/core/data_manager/strategy/cloud_strategy_impl.py +++ b/core/data_manager/strategy/cloud_strategy_impl.py @@ -1133,78 +1133,6 @@ async def _flush_records(self) -> int: self._stats["flush_count"] += 1 log.debug("Flushed %d cloud records", len(records)) return len(records) - - async def fetch_done_steps_with_context( - self, - job_id: str, - after_id: int = 0, - limit: int = 100 - ) -> List[Dict]: - """ - Fetch completed steps for training data collection. - Uses cursor-based pagination. - """ - await self.init() - - results = self.client.pull_data( - dataset_type=CLOUD_DATASET_TYPE, - cursor=after_id, - checkout_latest=True, - where_sql="job_id = '{}' AND is_terminal = True".format(_escape_sql_literal(job_id)), - limit=limit, - deserialize_json=True, - ) - - if results is None or len(results) == 0: - log.debug("No completed cloud steps to fetch: result_count=%s", 0 if results is None else len(results)) - return [] - - cursor = self.client.extract_cursor(results) - - rows = [] - for _, row in results.iterrows(): - meta = _meta_json_object(row.get("meta_json")) - messages = _json_value(row.get("messages"), []) - if not isinstance(messages, (dict, list)): - messages = [] - response = _json_value( - row.get("response"), - row.get("response"), - ) - rows.append( - { - "step_pk": cursor, - "step_id": row["step_id"], - "env_name": row["env_name"], - "env_id": row["session_id"], - "meta_json": json.dumps(meta, ensure_ascii=False, default=str), - "prompt": self.normalize_messages(messages), - "request": meta.get("request"), - "response": _response_text(response), - "reward": row["reward"], - "step_reward": row["step_reward"], - "total_reward": row["reward"], - "session_id": row["session_id"], - "session_end_time": row["created_at"] if row["created_at"] else None, - "group_id": meta.get("group_id"), - "truncated": row["is_truncated"], - "is_session_completed": row["is_session_completed"], - } - ) - return rows - - async def get_max_step_id(self, job_id: str) -> int: - """Get maximum primary key for pagination""" - await self.init() - - last_cursor = self.client.get_max_created_at( - where_sql=( - "dataset_type = '{}' AND job_id = '{}' AND is_terminal = True" - .format(CLOUD_DATASET_TYPE, _escape_sql_literal(job_id)) - ), - ) - - return last_cursor # --- Helpers --- def extract_image_path(self, item: dict) -> str | None: diff --git a/core/data_manager/strategy/sqlite_strategy_impl.py b/core/data_manager/strategy/sqlite_strategy_impl.py index 4bcdb942..3186136c 100644 --- a/core/data_manager/strategy/sqlite_strategy_impl.py +++ b/core/data_manager/strategy/sqlite_strategy_impl.py @@ -536,82 +536,3 @@ def buffer_stats(self) -> Optional[dict]: """Get buffer statistics""" return self._write_buffer.stats if self._write_buffer else None - async def fetch_done_steps_with_context( - self, - job_id: str, - after_id: int = 0, - limit: int = 100 - ) -> List[Dict]: - """ - Fetch completed steps for training data collection. - Uses cursor-based pagination. - """ - await self.init() - - trace = PerfTrace( - "sqlite_strategy.fetch_done_steps_with_context", - logger=log, - context={ - "operation": "db_read", - "table": "session_steps", - "job_id": job_id, - "after_id": after_id, - "limit": limit, - }, - ) - try: - with trace.span("db_read.fetch_done_steps", limit=limit): - steps = await SessionStep.filter( - job_id=job_id, - is_trainable=True, - id__gt=after_id - ).order_by("id").limit(limit) - - rows = [ - { - "step_pk": s.id, - "step_id": s.step_id, - "env_name": s.env_name, - "env_id": s.session_id, - "meta_json": s.meta_json, - "prompt": s.messages, - "request": s.request, - "response": s.response, - "reward": s.step_reward, - "step_reward": s.step_reward, - "total_reward": s.reward, - "session_id": s.session_id, - "session_end_time": s.created_at.isoformat() if s.created_at else None, - "group_id": s.group_id, - "truncated": s.is_truncated, - "is_session_completed": s.is_session_completed, - } - for s in steps - ] - trace.emit_summary(status="success", row_count=len(rows)) - return rows - except Exception as exc: - trace.emit_summary(status="failed", error_type=type(exc).__name__, error=str(exc)) - raise - - async def get_max_step_id(self, job_id: str) -> int: - """Get maximum primary key for pagination""" - await self.init() - trace = PerfTrace( - "sqlite_strategy.get_max_step_id", - logger=log, - context={ - "operation": "db_read", - "table": "session_steps", - "job_id": job_id, - }, - ) - try: - with trace.span("db_read.max_terminal_step_id"): - latest = await SessionStep.filter(job_id=job_id, is_terminal=True).order_by("-id").first() - max_id = latest.id if latest else 0 - trace.emit_summary(status="success", row_count=1 if latest else 0, max_step_id=max_id) - return max_id - except Exception as exc: - trace.emit_summary(status="failed", error_type=type(exc).__name__, error=str(exc)) - raise diff --git a/docs/guides/patcheval-rl-guide_CN.md b/docs/guides/patcheval-rl-guide_CN.md new file mode 100644 index 00000000..bf0ff47d --- /dev/null +++ b/docs/guides/patcheval-rl-guide_CN.md @@ -0,0 +1,123 @@ +# PatchEval RL 调通指南 + +汇总 patcheval RL(RJob 模式,Qwen3.5-9B / Qwen3.8-27B)从跑不起来到能产出训练数据的关键改动。按问题分类,每条给出现象、根因、修复。 + +--- + +## 一、生成 / 模板相关 Bug + +### B1. `TemplateError: System message must be at the beginning` / `No user query found` + +- **现象**:RolloutManager `apply_chat_template` 崩溃。Qwen3.x 模板要求 system 在最前、必须有 user query。 +- **根因**:`TrajectoryMaskBuilder._render_message_delta_str` 单独渲染 system message 时,模板注入合成 system → guard 失败。 +- **修复**(`rl/mask/trajectory_mask_builder.py`):新增 `_USER_ONLY_BASE` + `_get_user_suffix_str()`,system message 改渲染 `[system_msg] + _USER_ONLY_BASE` 再剥离,同时满足两条 guard。 + +### B2. `IndexError` in `_init_suffix_tokens` + +- **根因**:`apply_chat_template(tokenize=True)` 返回 `BatchEncoding` 而非 list,按下标迭代错位。 +- **修复**:解包成纯 token list 再处理。 + +### B3. `prepare_generate_input() takes 3 positional arguments but 4 were given` + +- **根因**:`tools` 支持是未提交改动,文件从 HEAD 恢复后丢签名。 +- **修复**:给 `prepare_generate_input` / `_ensure_path` / `_add_prompt_message` 加 `tools` 参数;首条 system 带 `` 块时用 `_render_first_system_delta_str` 渲染(F1)。 + +### B4. `TypeError: Can only get item pairs from a mapping` + +- **根因**:OpenHands 发 OpenAI 格式 `tool_calls`,`arguments` 是 JSON 字符串,Qwen 模板对它用 `|items` 要求 dict。 +- **修复**(`rl/llm_proxy.py`):`_normalize_messages_for_qwen_template` 把 `arguments` JSON string→dict、`content` None→"",在 `prepare_generate_input` 前调用。 + +--- + +## 二、Buffer / 取数相关 Bug + +### B5. Buffer 凑不齐 group → `rollout data is not ready` 死锁 + +- **现象**:buffer 持续 `new_items=0, ready_groups=0`,pending 组永远凑不齐 `group_size`,训练拿不到数据。DB 里其实已有满组,buffer 看不到。 +- **根因**:terminal step 是 `reward_committer` **UPDATE 现有行**翻转 `is_terminal`(不是 INSERT),行 `id` 在创建时就定了。用 `id` 自增游标增量捞 terminal step,eval 晚翻转的行 id 已被游标越过 → 永远漏捞 → 组凑不齐 → 死锁。 +- **现行修复(finished-env 两阶段 fetch,已取代早期 lookback 方案)**: + - 不再用 `id` 游标捞单步 terminal。改为: + - Phase 1:`list_environment_rows(finished=True)` —— 先捞被 `mark_environment_finished` 标记完成的 env。 + - Phase 2:`list_terminal_steps_for_sessions(env_ids)` —— 再批量取这些 env 的 terminal step。 + - 不变量:`mark_environment_finished` 只在 env 全部 step 都 `is_terminal=True` 后才调用 → `finished=True` 保证训练 step 全到齐 → 无 late-flip → 不需要 lookback / 去重。 + - 早期 lookback + `served_pks` 方案(`fetch_done_steps_with_context`)已废弃删除(提交 `2ba46ef`)。 +- **改动文件**:`rl/buffer_server.py`(`fetch_new_items_from_db` 两阶段逻辑)、`core/data_manager/manager.py` + `sqlite/cloud_strategy_impl.py`(`list_environment_rows` / `list_session_step_rows`)。 + +### B11. `get_training_info matched=0` → 0 trainable groups → weight_version 卡 1 + +- **现象**:slime.log 大量 `matched=0, expected=N`,`Trainable groups added this round: 0`,weight_version 永远 1。 +- **根因**:生成时 `llm_proxy` 先 `_normalize_messages_for_qwen_template`(arguments→dict),训练取数时直接读 DB(arguments 是 JSON string)→ `_message_matches` 对不上 → matched=0 → 0 trainable groups → 永不更新权重。**比 reward=0 更根本**。 +- **修复**(`rl/slime_generator.py::_get_record_training_info`):调 `get_training_info` 前对 DB 消息做同样的 `_normalize_messages_for_qwen_template` 归一化。 + +--- + +## 三、Episode / 封盘相关 Bug + +### B6. 每步生成被截断 → agent 1 步结束 → patch 空 → eval 全 failed → 熔断 → pool 停 → 死锁 + +- **根因**:OpenHands 请求不带 `max_tokens`(自定义 gateway 路由的 model 自动检测失败)→ sglang 用极小默认值 → 生成在工具调用中途被截断(`finish_reason=length`)→ agent 拿不到完整 tool call → 1 步结束 → 没改文件 → `patch=""` → rule_evaluator 返回 failed → 连续失败触发熔断 → pool 停。 +- **修复**(gateway 兜底,保证生效): + - `gateway/app.py`:`_ensure_default_max_tokens`,请求缺 `max_tokens` 时注入默认(env `GATEWAY_DEFAULT_MAX_TOKENS`,默认 8192,设 0 关闭)。 + - `env/patcheval/openhands_runner.py`:`_run_openhands` 显式设 `LLM_MAX_OUTPUT_TOKENS`(best-effort,当前镜像版本未翻译进请求,靠 gateway 兜底)。 + +### B7. 封盘超时不匹配 → 孤儿 session(`is_terminal=0`)→ group 凑不满 → 死锁 + +- **根因**:runner `close_session` 超时 15s < gateway `drain_timeout=30s` → runner 15s 抛超时放弃,gateway 仍强封写 `is_terminal=1`,两边对不上 → 孤儿。 +- **修复**(治本 + 治源头): + - `args.py` / `manager/types.py`:`gateway_close_timeout_s` 默认 15→45(>gateway 30s drain)。**治本。** + - `rl/buffer_server.py`:注入 `--gateway-close-timeout-s`(env `AIEVOBOX_GATEWAY_CLOSE_TIMEOUT_S`,默认 45)。 + - `rl/examples/patcheval/env.rjob.sh`:`AIEVOBOX_GATEWAY_MAX_STEPS` 收紧,缩短 episode。 + - `gateway/app.py`:`GATEWAY_DEFAULT_MAX_TOKENS` 16384→6144,收紧单步生成,降低 drain 压力(有意权衡:宁可截断但封盘,不要完整但孤儿)。 + +--- + +## 四、Megatron GDN 不支持 Packed Sequence + +### GDN packed-seq monkey-patch + +- **现象**:Qwen3.8-27B(48 层 GDN + 16 层 Full Attention)第一步 `compute_log_prob` 崩溃:`NotImplementedError: GDN does not support packed sequence for now.`(`megatron/core/ssm/gated_delta_net.py:302`)。 +- **根因**:slime 默认 `thd`(packed)布局,多条 trajectory concat 成长序列用 `cu_seqlens` 标边界。Megatron 原生 GDN forward 入口直接 raise,没把 `cu_seqlens` 传给底层 `chunk_gated_delta_rule`(fla 本身已支持 cu_seqlens)。bridge 模式忽略 `--spec`(slime 的 `qwen3_5.py` 有支持 cu_seqlens 的 GDN),用不上。 +- **修复**(运行时 monkey-patch,不重打镜像、不改 Megatron/slime 核心): + - `rl/patches/gdn_packed_seq.py`:patch `GatedDeltaNet.forward`,删 raise,从 `packed_seq_params.cu_seqlens_q` 提取 cu_seqlens 传给 `chunk_gated_delta_rule` 和 `causal_conv1d_fn`。 + - `rl/patches/sitecustomize.py`:启动时自动加载。 + - `env.rjob.sh`:`export PYTHONPATH="${REPO_ROOT}/rl/patches${PYTHONPATH:+:${PYTHONPATH}}"`。 +- **为什么不用 bshd(padding)**:27B TP=4 在 140GB 卡 OOM(差 822 MiB),thd packing 内存更省。 +- **回退**:删 `PYTHONPATH` 那行即禁用。 + +--- + +## 五、配置 / Feature + +- **F2 DAPO filter 默认关闭**:`env.rjob.sh` `DAPO_filter=false`,避免早期 reward 全 0 时 pipeline 卡死。 +- **F3 env.rjob.sh 自包含**:移除 `source geo3k_vl/env.sh`,内联默认值,避免 VL 任务配置串味。 +- **F4 RL_EPOCH 默认 100**(从 2)。 +- **reward pre-gate 放宽**(`env/patcheval/rule_evaluator.py`):缺 cve_id/patch/language 或容器起不来时返回 `SUCCEEDED+0.0`(对齐 bench 的 validation_fail=0 语义),不再返回 FAILED(FAILED 会让 RewardCommitter 拒写 → reward NULL → session 不 seal)。同步覆盖 162 个 per-CVE 副本。 +- **LOSS_MASK_TYPE 注入**(`rl/run_slime_generator.sh`):`RUNTIME_ENV_JSON.env_vars` 加 `LOSS_MASK_TYPE=qwen3_5`,修复 RolloutManager Ray actor 拿不到 → 用 base adapter → B1 报错的 bug。 +- **gateway host**:env 脚本 `hostname -I` → `hostname -i`。 + +--- + +## 六、改动文件清单 + +| 文件 | 说明 | +|---|---| +| `rl/mask/trajectory_mask_builder.py` | B1/B2/B3 + F1:模板渲染、BatchEncoding 解包、tools | +| `rl/llm_proxy.py` | B4:`_normalize_messages_for_qwen_template` | +| `rl/slime_generator.py` | B11:训练取数归一化;加 `rl/mask` 到 sys.path | +| `rl/buffer_server.py` | B5:finished-env 两阶段 fetch;B7:`--gateway-close-timeout-s` | +| `core/data_manager/manager.py` / `*_strategy_impl.py` | B5:`list_environment_rows` / `list_session_step_rows` | +| `env/patcheval/rule_evaluator.py` | reward pre-gate 放宽(+ 162 per-CVE 副本) | +| `env/patcheval/openhands_runner.py` | B6:`LLM_MAX_OUTPUT_TOKENS` | +| `gateway/app.py` | B6:`_ensure_default_max_tokens`;B7:默认 6144 | +| `args.py` / `manager/types.py` | B7:`gateway_close_timeout_s` 15→45 | +| `rl/patches/gdn_packed_seq.py` + `sitecustomize.py` | GDN packed-seq monkey-patch | +| `rl/examples/patcheval/env.rjob*.sh` | F2/F3/F4 + B7 + gateway host + max_steps | +| `rl/run_slime_generator.sh` | LOSS_MASK_TYPE 注入 | + +--- + +## 七、已知遗留 / Follow-up + +- **cloud 后端 late-flip**:`cloud_strategy_impl.py` 的游标是 created_at 时间戳,理论上同样有 late-flip 风险,finished-env 方案需在云后端验证。 +- **无 dockerd 时的学习信号**:launcher 侧无 dockerd 时 CVE 容器起不来,所有 reward=0.0 → 组内无方差 → GRPO 梯度≈0。管道能跑(`ready_groups>0`)但无真实学习。要 1.0 正样本需让 `openhands_runner` 在 pod 内判分并上报 `strict_success/poc_passed`,使 `rule_evaluator` 的 runner fallback 命中。 +- **github 屏蔽**:`openhands_runner._block_github_cdn` 是 patcheval 故意的防作弊 + 快速失败,**不是 bug**。 diff --git a/env/patcheval/.gitignore b/env/patcheval/.gitignore new file mode 100644 index 00000000..2293470e --- /dev/null +++ b/env/patcheval/.gitignore @@ -0,0 +1,3 @@ +generated_openhands_exp1*/ +push_patcheval_done.txt +push_patcheval_images.sh diff --git a/env/patcheval/generate_full_config.py b/env/patcheval/generate_full_config.py index 732cd418..3ce51fdb 100644 --- a/env/patcheval/generate_full_config.py +++ b/env/patcheval/generate_full_config.py @@ -73,6 +73,79 @@ def parse_args() -> argparse.Namespace: "--no-proxy", default="host.docker.internal,localhost,127.0.0.1,::1", ) + # --- RJob mode --- + parser.add_argument( + "--mode", + choices=["docker", "rjob"], + default="docker", + help="docker: CVE containers on one Docker host. rjob: each episode " + "submitted as a cluster RJob pod pulling images from the registry.", + ) + parser.add_argument( + "--rjob-name-prefix", + default="patcheval", + help="RJob name_prefix (RJob pod name component).", + ) + parser.add_argument( + "--rjob-registry", + default="registry.h.pjlab.org.cn", + help="RJob image registry host. In rjob mode, env_image is rewritten from " + "the dataset's ghcr.io path to //:-latest (the " + "tag pushed by push_patcheval_images.sh).", + ) + parser.add_argument( + "--rjob-registry-ns", + default="ailab-evobox-evobox_proxy", + help="RJob image registry namespace.", + ) + parser.add_argument( + "--rjob-repo", + default="patcheval", + help="RJob image repository name.", + ) + parser.add_argument( + "--rjob-openhands-bin", + default="/mnt/shared-storage-user/evobox-share/leishanzhe/openhands-install/openhands", + help="Pre-downloaded openhands binary on shared storage; embedded into " + "the RJob pod so the curl installer (flaky from pods) is skipped.", + ) + parser.add_argument( + "--rjob-results-root", + default="", + help="safactory_results_root written into rjob env_params. Defaults to " + "/results when empty.", + ) + parser.add_argument( + "--rjob-no-proxy", + default="localhost,127.0.0.1,::1,10.0.0.0/8,100.96.0.0/12,.pjlab.org.cn", + help="no_proxy for RJob pods (drops docker-only host.docker.internal " + "and the docker-mode pod IP; 100.x gateway IPs are covered by " + "100.96.0.0/12).", + ) + parser.add_argument( + "--rjob-mount-config", + nargs="+", + default=[ + "gpfs://gpfs1/leishanzhe:/mnt/shared-storage-user/leishanzhe", + "gpfs://gpfs1/evobox-share:/mnt/shared-storage-user/evobox-share", + ], + help="Cluster gpfs mounts for the RJob pod (gpfs:///:).", + ) + parser.add_argument("--rjob-cpu", type=int, default=2) + parser.add_argument("--rjob-memory-mb", type=int, default=5120) + parser.add_argument( + "--rjob-custom-resources", + nargs="+", + default=["brainpp.cn/fuse=1"], + help="RJob resources.custom_resources (resource=value).", + ) + parser.add_argument( + "--rjob-gateway-base-url", + default="", + help="If set, written uncommented as PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL. " + "If empty (default), left commented out — the runner uses the dynamic " + "SAFACTORY_GATEWAY_BASE_URL injected by the launcher.", + ) return parser.parse_args() @@ -152,6 +225,19 @@ def write_configs( claude_model: str, claude_max_thinking_tokens: int, openhands_install_timeout_s: float, + mode: str = "docker", + rjob_name_prefix: str = "patcheval", + rjob_registry: str = "registry.h.pjlab.org.cn", + rjob_registry_ns: str = "ailab-evobox-evobox_proxy", + rjob_repo: str = "patcheval", + rjob_openhands_bin: str = "", + rjob_results_root: str = "", + rjob_no_proxy: str = "localhost,127.0.0.1,::1,10.0.0.0/8,100.96.0.0/12,.pjlab.org.cn", + rjob_mount_config: list[str] | None = None, + rjob_cpu: int = 2, + rjob_memory_mb: int = 5120, + rjob_custom_resources: list[str] | None = None, + rjob_gateway_base_url: str = "", ) -> None: output_dir.mkdir(parents=True, exist_ok=True) dataset_dir = output_dir / "datasets" @@ -165,7 +251,7 @@ def write_configs( runner_name = { "claudecode": "claudecode_runner.py", "openhands": "openhands_runner.py", - }.get(baseline, "strict_runner.py") + }.get(baseline, "runner.py") runner_path = Path(__file__).resolve().with_name(runner_name) container_env = { "PYTHONDONTWRITEBYTECODE": "1", @@ -213,27 +299,113 @@ def write_configs( } ) - common_agent = { - "container": { - "workdir": "/workspace", - "runner_entrypoint": { - "source": str(runner_path), - "target": "/tmp/safactory-patcheval-runner.py", - "command": "python /tmp/safactory-patcheval-runner.py", + common_agent: dict[str, Any] + if mode == "rjob": + # RJob pods cannot reach the k8s-internal docker proxy, and the + # docker-only host.docker.internal / pod IP are not meaningful from a + # cluster pod. Proxy clearing + rjob no_proxy apply to ALL baselines. + # The runner is embedded via base64 for every baseline (rjob mode does + # not add install_runner_script outputs to volumes). + rjob_container_env = dict(container_env) + # RJob does not apply Docker bind-mount volumes (only mount_config/gpfs), + # so /opt/patcheval would be empty. Point the runner at the official + # source via its gpfs-accessible path instead. + rjob_container_env["PATCHEVAL_OFFICIAL_ROOT"] = str(official_runtime_dir) + rjob_container_env.update( + { + "NO_PROXY": rjob_no_proxy, + "no_proxy": rjob_no_proxy, + "HTTP_PROXY": "", + "HTTPS_PROXY": "", + "http_proxy": "", + "https_proxy": "", + } + ) + if baseline == "openhands": + # OpenHands runs DinD inside the pod to load CVE images and uses a + # pre-downloaded openhands binary on shared storage (the curl + # installer is flaky from pods). The LLM baseline needs neither. + rjob_container_env.update( + { + "DOCKER_HOST": "unix:///var/run/docker.sock", + "DOCKER_TLS_CERTDIR": "", + "PATCHEVAL_OPENHANDS_BIN": rjob_openhands_bin, + } + ) + # PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL: left commented out by default + # (the runner prefers the dynamic SAFACTORY_GATEWAY_BASE_URL injected + # by the launcher). Write it uncommented only if explicitly provided. + if rjob_gateway_base_url: + rjob_container_env["PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL"] = rjob_gateway_base_url + + rjob_block: dict[str, Any] = { + "name_prefix": rjob_name_prefix, + "private_machine": "group", + "image_pull_policy": "IfNotPresent", + "no_packaging": True, + "cleanup_on_finish": False, + "keep_failed_jobs": True, + # NOTE: do NOT set preemptible: false — the rjob SDK stringifies the + # Python bool to "False", which the cluster admission webhook + # rejects. Omitting the key uses the SDK default (preemptible). + "privileged": True, + "resources": { + "cpu": rjob_cpu, + "gpu": 0, + "memory_in_mb": rjob_memory_mb, + "custom_resources": list(rjob_custom_resources or []), }, - "install_runner_script": True, - "env": container_env, - "extra_args": ["--add-host=host.docker.internal:host-gateway"], - "volumes": [ + "mount_config": list(rjob_mount_config or []), + "embedded_files": [ { - "source": str(official_runtime_dir), - "target": "/opt/patcheval", - "read_only": True, + "source": str(runner_path), + "target": "/tmp/safactory-patcheval-runner.py", } ], - "idle_command": "tail -f /dev/null", } - } + common_agent = { + "container": { + "workdir": "/workspace", + "runner_entrypoint": { + "source": str(runner_path), + "target": "/tmp/safactory-patcheval-runner.py", + "command": "python /tmp/safactory-patcheval-runner.py", + }, + "install_runner_script": True, + "env": rjob_container_env, + "volumes": [ + { + "source": str(official_runtime_dir), + "target": "/opt/patcheval", + "read_only": True, + } + ], + "idle_command": "tail -f /dev/null", + }, + "rjob": rjob_block, + } + else: + common_agent = { + "container": { + "workdir": "/workspace", + "runner_entrypoint": { + "source": str(runner_path), + "target": "/tmp/safactory-patcheval-runner.py", + "command": "python /tmp/safactory-patcheval-runner.py", + }, + "install_runner_script": True, + "env": container_env, + "extra_args": ["--add-host=host.docker.internal:host-gateway"], + "volumes": [ + { + "source": str(official_runtime_dir), + "target": "/opt/patcheval", + "read_only": True, + } + ], + "idle_command": "tail -f /dev/null", + } + } missing_archives: list[str] = [] for record in records: @@ -262,23 +434,44 @@ def write_configs( task.update({"setting": setting, "prompt_template": prompt_template}) dataset_path.write_text(json.dumps(task, ensure_ascii=False) + "\n", encoding="utf-8") + env_params: dict[str, Any] = { + "task_family": "patcheval", + "rule_evaluator_timeout_s": evaluation_timeout_s, + "patcheval_official_root": str(official_root), + "patcheval_docker_adapter": str(docker_adapter_path), + "patcheval_image_archive_dir": str(archive_dir or ""), + "patcheval_http_proxy": http_proxy, + "patcheval_no_proxy": no_proxy, + "patcheval_shared_tmp": shared_tmp, + } + if mode == "rjob": + # RJob pods cannot reach the k8s-internal docker proxy, and the + # docker-only host.docker.internal / pod IP are not meaningful + # from a cluster pod. Mirror the hand-written rjob sample. + env_params["patcheval_http_proxy"] = "" + env_params["patcheval_no_proxy"] = rjob_no_proxy + if rjob_results_root: + env_params["safactory_results_root"] = rjob_results_root + else: + # Default to the SAfactory repo's results dir (this file lives at + # /env/patcheval/generate_full_config.py). + repo_root = Path(__file__).resolve().parents[2] + env_params["safactory_results_root"] = str(repo_root / "results") environments.append( { "env_name": name, - "env_image": image, + "env_image": ( + # RJob pods pull from the internal registry (the tag pushed by + # push_patcheval_images.sh), not the ghcr.io path the docker + # mode loads from the local tar archive. + f"{rjob_registry}/{rjob_registry_ns}/{rjob_repo}:{cve_id.lower()}-latest" + if mode == "rjob" + else image + ), "env_num": 1, "dataset": f"./datasets/{dataset_filename}", "dataset_load_mode": "eager", - "env_params": { - "task_family": "patcheval", - "rule_evaluator_timeout_s": evaluation_timeout_s, - "patcheval_official_root": str(official_root), - "patcheval_docker_adapter": str(docker_adapter_path), - "patcheval_image_archive_dir": str(archive_dir or ""), - "patcheval_http_proxy": http_proxy, - "patcheval_no_proxy": no_proxy, - "patcheval_shared_tmp": shared_tmp, - }, + "env_params": env_params, } ) agents[name] = common_agent @@ -297,9 +490,11 @@ def write_configs( suffix = "" if len(missing_archives) <= 10 else f"\n ... and {len(missing_archives) - 10} more" raise FileNotFoundError(f"Missing {len(missing_archives)} image archive(s):\n{sample}{suffix}") - with (output_dir / "patcheval_config.yaml").open("w", encoding="utf-8") as handle: + config_name = "patcheval_config.rjob.yaml" if mode == "rjob" else "patcheval_config.yaml" + start_name = "patcheval_start.rjob.yaml" if mode == "rjob" else "patcheval_start.yaml" + with (output_dir / config_name).open("w", encoding="utf-8") as handle: yaml.safe_dump({"environments": environments}, handle, sort_keys=False, allow_unicode=True) - with (output_dir / "patcheval_start.yaml").open("w", encoding="utf-8") as handle: + with (output_dir / start_name).open("w", encoding="utf-8") as handle: yaml.safe_dump({"agents": agents}, handle, sort_keys=False, allow_unicode=True) def main() -> None: args = parse_args() @@ -373,9 +568,22 @@ def main() -> None: str(args.claude_model).strip(), int(args.claude_max_thinking_tokens), float(args.openhands_install_timeout_s), + mode=str(args.mode), + rjob_name_prefix=str(args.rjob_name_prefix), + rjob_registry=str(args.rjob_registry), + rjob_registry_ns=str(args.rjob_registry_ns), + rjob_repo=str(args.rjob_repo), + rjob_openhands_bin=str(args.rjob_openhands_bin), + rjob_results_root=str(args.rjob_results_root).strip(), + rjob_no_proxy=str(args.rjob_no_proxy).strip(), + rjob_mount_config=list(args.rjob_mount_config or []), + rjob_cpu=int(args.rjob_cpu), + rjob_memory_mb=int(args.rjob_memory_mb), + rjob_custom_resources=list(args.rjob_custom_resources or []), + rjob_gateway_base_url=str(args.rjob_gateway_base_url).strip(), ) print( - f"Generated PatchEval {args.baseline} configuration " + f"Generated PatchEval {args.baseline} ({args.mode}) configuration " f"for {len(records)} task(s) in {output_dir}" ) diff --git a/env/patcheval/openhands_runner.py b/env/patcheval/openhands_runner.py index fcf617c9..2cbf73b1 100644 --- a/env/patcheval/openhands_runner.py +++ b/env/patcheval/openhands_runner.py @@ -14,13 +14,79 @@ DEFAULT_TIMEOUT_S = 2700.0 DEFAULT_INSTALL_TIMEOUT_S = 900.0 +DEFAULT_MAX_OUTPUT_TOKENS = 8192 MAX_LOG_CHARS = 32_000 +# Live log file on shared storage so the training node can `tail -f` and see +# exactly where OpenHands is stuck (subprocess.run(capture_output=True) buffers +# everything in memory until exit, so a hang shows nothing). Opened in main(). +_LOG_FILE = None + + +def _log(msg: str) -> None: + if _LOG_FILE is None: + return + try: + _LOG_FILE.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}\n") + _LOG_FILE.flush() + except Exception: + pass + + +def _open_log(session_id: str): + global _LOG_FILE + candidates = [] + result_path = os.environ.get("SAFACTORY_RESULT_PATH", "") + if result_path: + candidates.append(os.path.join(os.path.dirname(result_path), "openhands.log")) + subdir = os.environ.get("SAFACTORY_OUTPUT_SUBDIR", "") + if subdir: + candidates.append(os.path.join(subdir, "openhands.log")) + candidates.append(f"/tmp/openhands-{session_id}.log") + for path in candidates: + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + f = open(path, "a", buffering=1) + _LOG_FILE = f + _log(f"===== openhands_runner start session={session_id} pid={os.getpid()} =====") + _log(f"log file: {path}") + return + except Exception: + continue + + +def _block_github_cdn() -> None: + """Point github.com + raw.githubusercontent.com at 127.0.0.1 in /etc/hosts + so the openhands binary's startup calls (version check via + raw.githubusercontent.com, public-skills git clone via github.com) fail fast + (ECONNREFUSED) instead of hanging on the slow/flaky GitHub CDN from RJob + pods. The binary catches these errors and proceeds to the LLM calls. Only + done when using the pre-installed local binary (PATCHEVAL_OPENHANDS_BIN), + since the curl installer path needs real github.com access. + """ + if not os.environ.get("PATCHEVAL_OPENHANDS_BIN", "").strip(): + return + hosts = ["raw.githubusercontent.com", "github.com", "objects.githubusercontent.com"] + try: + with open("/etc/hosts", "r") as f: + current = f.read() + additions = [f"127.0.0.1 {h}" for h in hosts if h not in current] + if not additions: + _log("block_github: already blocked in /etc/hosts") + return + with open("/etc/hosts", "a") as f: + f.write("\n# safactory: fast-fail openhands startup github calls\n") + f.write("\n".join(additions) + "\n") + _log(f"block_github: added to /etc/hosts: {additions}") + except Exception as exc: + _log(f"block_github: failed to edit /etc/hosts: {exc!r}") + def main() -> int: started_at = time.perf_counter() request = _read_request() session_id = _required_text(request.get("session_id"), "session_id") + _open_log(session_id) cve_id = "" try: @@ -99,16 +165,42 @@ def main() -> int: def _ensure_openhands(timeout_s: float) -> str: + # 1) Explicit override: a pre-installed openhands binary on shared storage. + _log("ensure_openhands: step 1) PATCHEVAL_OPENHANDS_BIN override") + override = os.environ.get("PATCHEVAL_OPENHANDS_BIN", "").strip() + if override and Path(override).is_file() and os.access(override, os.X_OK): + _log(f"ensure_openhands: using override {override}") + return override + # 2) Already on PATH (e.g. baked into the env image). + _log("ensure_openhands: step 2) which openhands") existing = shutil.which("openhands") if existing: + _log(f"ensure_openhands: using PATH binary {existing}") return existing + # 3) pip install from the cluster's internal PyPI mirror (no external egress + # needed). RJob env pods can reach mirrors.h.pjlab.org.cn (same domain as + # the image registry they already pull from). PIP_INDEX_URL is set in the + # start yaml env so pip uses the internal mirror instead of pypi.org. + import sys + _log("ensure_openhands: step 3) pip install openhands-ai") + try: + _run_streamed([sys.executable, "-m", "pip", "install", "openhands-ai"], timeout_s) + except Exception as exc: + _log(f"ensure_openhands: pip install raised {exc!r}") + found = shutil.which("openhands") + if found: + _log(f"ensure_openhands: using pip-installed binary {found}") + return found + # 4) Last resort: download + run the official installer (needs external egress). + _log("ensure_openhands: step 4) curl install.openhands.dev installer") install_dir = Path("/opt/openhands") install_dir.mkdir(parents=True, exist_ok=True) script = Path("/tmp/install-openhands.sh") - _run(["curl", "-fsSL", "https://install.openhands.dev/install.sh", "-o", str(script)], timeout_s) + _run_streamed(["curl", "-fsSL", "https://install.openhands.dev/install.sh", "-o", str(script)], timeout_s) env = os.environ.copy() env["OPENHANDS_INSTALL_DIR"] = str(install_dir) - install = _run(["bash", str(script)], timeout_s, env=env, check=False) + install_rc, install_out = _run_streamed(["bash", str(script)], timeout_s, env=env) + _log(f"ensure_openhands: installer exit={install_rc}") candidates = ( install_dir / "openhands", Path("/usr/local/bin/openhands"), @@ -118,10 +210,12 @@ def _ensure_openhands(timeout_s: float) -> str: ) for candidate in candidates: if candidate.is_file() and os.access(candidate, os.X_OK): + _log(f"ensure_openhands: using installer binary {candidate}") return str(candidate) + _log("ensure_openhands: NO binary found after all steps") raise RuntimeError( "OpenHands installation completed but no executable was found. " - f"installer exit={install.returncode}; output={_trim_log(install.stdout + install.stderr)}" + f"installer exit={install_rc}; output={_trim_log(install_out)}" ) @@ -134,9 +228,14 @@ def _run_openhands( problem_statement: str, timeout_s: float, ) -> dict[str, Any]: + # Prefer the dynamic gateway URL injected by the launcher + # (SAFACTORY_GATEWAY_BASE_URL, derived from AIEVOBOX_GATEWAY_HOST) over the + # static PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL baked into the start yaml, + # which can go stale when the training pod IP changes between runs. gateway_base = _required_text( - os.environ.get("PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL"), - "PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL", + os.environ.get("SAFACTORY_GATEWAY_BASE_URL") + or os.environ.get("PATCHEVAL_OPENHANDS_GATEWAY_BASE_URL"), + "SAFACTORY_GATEWAY_BASE_URL", ).rstrip("/") route_model = _required_text( os.environ.get("PATCHEVAL_OPENHANDS_MODEL"), @@ -148,6 +247,15 @@ def _run_openhands( "PatchEval evaluator scripts, test.patch, fix.patch, or other benchmark artifacts. " "Leave the final code changes in the git working tree." ) + # OpenHands defaults max_output_tokens to 0 ("auto-detect from model"), but + # auto-detection fails for the custom gateway-routed model, so no max_tokens + # is sent in the LLM request and sglang falls back to a tiny default (~128), + # truncating generation mid-tool-call (finish_reason=length). Set an + # explicit cap so tool calls + reasoning render fully. + max_output_tokens = _positive_int( + os.environ.get("PATCHEVAL_OPENHANDS_MAX_OUTPUT_TOKENS"), + DEFAULT_MAX_OUTPUT_TOKENS, + ) env = os.environ.copy() env.update( { @@ -155,6 +263,7 @@ def _run_openhands( "LLM_MODEL": f"openai/{route_model}", "LLM_API_KEY": "safactory", "LLM_BASE_URL": f"{gateway_base}/{session_id}", + "LLM_MAX_OUTPUT_TOKENS": str(max_output_tokens), } ) command = [ @@ -167,16 +276,55 @@ def _run_openhands( "--task", task, ] + _log(f"run_openhands: cwd={work_dir} gateway={gateway_base} model={route_model}") + _log(f"run_openhands: cmd={' '.join(command)}") + _log(f"run_openhands: LLM_BASE_URL={env['LLM_BASE_URL']}") + _block_github_cdn() try: - completed = _run(command, timeout_s, cwd=work_dir, check=False, env=env) - except subprocess.TimeoutExpired as exc: - output = (exc.stdout or "") + (exc.stderr or "") - return {"exit_code": None, "output": output, "timed_out": True} - return { - "exit_code": completed.returncode, - "output": completed.stdout + completed.stderr, - "timed_out": False, - } + proc = subprocess.Popen( + command, + cwd=str(work_dir), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + buffer_parts: list[str] = [] + deadline = time.perf_counter() + timeout_s + timed_out = False + assert proc.stdout is not None + while True: + line = proc.stdout.readline() + if line == "": + if proc.poll() is not None: + break + if time.perf_counter() > deadline: + timed_out = True + break + continue + buffer_parts.append(line) + _log(f"OH| {line.rstrip()}") + if timed_out and proc.poll() is None: + _log(f"run_openhands: TIMEOUT after {timeout_s}s, killing pid={proc.pid}") + proc.kill() + try: + remaining = proc.communicate(timeout=10)[0] or "" + except Exception: + remaining = "" + if remaining: + buffer_parts.append(remaining) + _log(f"OH| (tail) {remaining.rstrip()}") + exit_code = proc.returncode + _log(f"run_openhands: exit_code={exit_code} timed_out={timed_out}") + return { + "exit_code": exit_code, + "output": "".join(buffer_parts), + "timed_out": timed_out, + } + except Exception as exc: + _log(f"run_openhands: exception {exc!r}") + return {"exit_code": None, "output": str(exc), "timed_out": False} def _hide_evaluation_artifacts() -> None: @@ -212,7 +360,20 @@ def _read_request() -> dict[str, Any]: def _write_result(value: dict[str, Any]) -> None: - print(json.dumps(value, ensure_ascii=False)) + text = json.dumps(value, ensure_ascii=False) + print(text, flush=True) + # RJob mode: logs_rjob sometimes returns empty stdout, so also persist the + # result to the shared artifact path (SAFACTORY_RESULT_PATH) on gpfs. The + # launcher's parse_result_artifact fallback reads this same file. + result_path = os.environ.get("SAFACTORY_RESULT_PATH") + if result_path: + try: + Path(result_path).parent.mkdir(parents=True, exist_ok=True) + Path(result_path).write_text(text) + except Exception: + # stdout print above is the primary path; never let a file-write + # failure change the process exit status. + pass def _run( @@ -234,6 +395,61 @@ def _run( ) +def _run_streamed( + args: list[str], + timeout_s: float, + *, + cwd: Path = Path("/workspace"), + env: dict[str, str] | None = None, +) -> tuple[int, str]: + """Run a command, streaming merged stdout/stderr to the live log file. + + Used for the OpenHands install steps (pip / curl / bash installer) so a + hang during download is visible in the shared log instead of silent. + Returns (returncode, combined_output). Never raises on timeout/non-zero. + """ + try: + proc = subprocess.Popen( + args, + cwd=str(cwd), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + except Exception as exc: + _log(f"run_streamed: spawn failed {args[0]}: {exc!r}") + return -1, str(exc) + parts: list[str] = [] + deadline = time.perf_counter() + timeout_s + timed_out = False + assert proc.stdout is not None + while True: + line = proc.stdout.readline() + if line == "": + if proc.poll() is not None: + break + if time.perf_counter() > deadline: + timed_out = True + break + continue + parts.append(line) + _log(f"INSTALL| {line.rstrip()}") + if timed_out and proc.poll() is None: + _log(f"run_streamed: TIMEOUT after {timeout_s}s, killing {args[0]}") + proc.kill() + try: + rem = proc.communicate(timeout=10)[0] or "" + except Exception: + rem = "" + if rem: + parts.append(rem) + _log(f"INSTALL| (tail) {rem.rstrip()}") + _log(f"run_streamed: {args[0]} exit={proc.returncode} timed_out={timed_out}") + return proc.returncode, "".join(parts) + + def _required_text(value: Any, name: str) -> str: text = str(value or "").strip() if not text: @@ -249,6 +465,14 @@ def _positive_float(value: Any, default: float) -> float: return parsed if parsed > 0 else default +def _positive_int(value: Any, default: int) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + def _trim_log(value: str) -> str: return str(value or "")[-MAX_LOG_CHARS:] diff --git a/env/patcheval/rule_evaluator.py b/env/patcheval/rule_evaluator.py index df539f7e..52a9a438 100644 --- a/env/patcheval/rule_evaluator.py +++ b/env/patcheval/rule_evaluator.py @@ -35,17 +35,40 @@ async def evaluate_rule( official_record = official_record if isinstance(official_record, dict) else {} language = str(official_record.get("programming_language") or "").strip() + # Pre-gate (relaxed): align with PatchEval's native scoring semantics. + # The official bench has NO upfront rejection of empty/missing patches -- an + # empty patch flows through the Docker oracle and scores 0 (validation_fail), + # it is a legitimate negative sample, NOT "no data". Returning FAILED here + # made RewardCommitter refuse to write a reward (reward stays NULL), which + # silently dropped these trajectories and starved RL of negative samples. + # So instead of FAILED/null, return SUCCEEDED with score 0 so a reward of 0.0 + # is committed and the trajectory can join a GRPO group. if not cve_id or not patch or not language: - return EvalResult.failed( + missing = [] + if not cve_id: + missing.append("cve_id") + if not patch: + missing.append("patch") + if not language: + missing.append("language") + return EvalResult( session_id=request.session_id, eval_id=spec.eval_id, method=spec.method.value, - reason="PatchEval runner did not provide cve_id, patch, and programming language", + status=EvalStatus.SUCCEEDED.value, + raw_score=0.0, + normalized_score_10=0.0, + reason=( + "PatchEval pre-gate: missing " + + ", ".join(missing) + + " -- scored 0 (no valid patch produced)" + ), artifacts={ "bench": "patcheval", "cve_id": cve_id or None, "patch_generated": bool(patch), "language": language or None, + "validation_type": "validation_fail", "metrics": metrics, }, ) @@ -65,28 +88,56 @@ async def evaluate_rule( [], ) except Exception as exc: - return EvalResult.failed( + # In rjob mode the runner already runs the official PatchEval evaluation + # (PoC + unit tests) inside the RJob pod using the real CVE image, and + # records strict_success/poc_passed/unit_tests_passed in its metrics. + # The launcher-side re-evaluation here needs a local Docker daemon, + # which rjob coordinator hosts typically lack (k8s worker nodes run + # containerd, not dockerd). When Docker is unavailable, fall back to the + # runner's authoritative in-pod result instead of failing the episode. + if _is_docker_unavailable(exc): + fallback = _fallback_from_runner_metrics(request, spec, metrics, exc) + if fallback is not None: + return fallback + # Align with PatchEval native semantics: an exception during official + # evaluation (typically Docker unavailable on the rjob coordinator) is + # a validation failure, NOT "no data". Score 0 so a reward of 0.0 is + # committed and the trajectory joins a GRPO group instead of being + # silently dropped (FAILED -> RewardCommitter refuses -> reward NULL). + return EvalResult( session_id=request.session_id, eval_id=spec.eval_id, method=spec.method.value, - reason="official PatchEval evaluation raised an exception", + status=EvalStatus.SUCCEEDED.value, + raw_score=0.0, + normalized_score_10=0.0, + reason="official PatchEval evaluation raised an exception (Docker unavailable?) -- scored 0", error_text=str(exc), - artifacts={"bench": "patcheval", "cve_id": cve_id, "patch": patch}, + artifacts={"bench": "patcheval", "cve_id": cve_id, "patch": patch, "validation_type": "validation_fail"}, ) strict_success = validation_type == "Repair Success" if poc_passed is None: - return EvalResult.failed( + # Align with PatchEval native semantics: the bench swallows the + # container-start exception (run_evaluation.py) and treats an + # unstartable CVE container as validation_fail -> 0, a legitimate + # negative sample. Returning FAILED here made RewardCommitter refuse + # to write a reward, silently dropping these patch-producing + # trajectories and starving RL of negative samples. + return EvalResult( session_id=request.session_id, eval_id=spec.eval_id, method=spec.method.value, - reason="official PatchEval evaluator could not start the CVE container", + status=EvalStatus.SUCCEEDED.value, + raw_score=0.0, + normalized_score_10=0.0, + reason="official PatchEval evaluator could not start the CVE container -- scored 0 (validation_fail)", error_text=_trim_log(poc_log or "unknown Docker evaluation error"), artifacts={ "bench": "patcheval", "cve_id": cve_id, "patch": patch, - "validation_type": validation_type, + "validation_type": "validation_fail", "poc_log": _trim_log(poc_log), "unit_test_log": _trim_log(unit_test_log), }, @@ -131,6 +182,74 @@ def _trim_log(value: Any) -> str: return str(value or "")[-_MAX_LOG_CHARS:] +def _is_docker_unavailable(exc: BaseException) -> bool: + """True when the exception indicates the Docker daemon is not reachable.""" + text = str(exc).lower() + if "fetching server api version" in text: + return True + if "no such file or directory" in text and "docker" in text: + return True + walked = exc + while walked is not None: + cls = type(walked) + if cls.__module__ == "docker.errors" or cls.__name__ == "DockerException": + return True + walked = walked.__cause__ or walked.__context__ + return False + + +def _fallback_from_runner_metrics( + request: EvalRequest, + spec: EvalSpec, + metrics: dict[str, Any], + exc: BaseException, +) -> EvalResult | None: + """Build an EvalResult from the runner's in-pod official evaluation. + + The runner runs the official CVE evaluation (PoC + unit tests) inside the + RJob pod with the real CVE image; its metrics carry strict_success, + poc_passed, unit_tests_passed, etc. When the launcher cannot re-run that + evaluation (no local Docker daemon), those metrics are the authoritative + result. + """ + if not isinstance(metrics, dict): + return None + if "strict_success" not in metrics or "poc_passed" not in metrics: + return None + strict_success = bool(metrics.get("strict_success")) + score = 1.0 if strict_success else 0.0 + failure_stage = metrics.get("failure_stage") + return EvalResult( + session_id=request.session_id, + eval_id=spec.eval_id, + method=spec.method.value, + status=EvalStatus.SUCCEEDED.value, + raw_score=score, + normalized_score_10=10.0 if strict_success else 0.0, + reason=( + "runner in-pod evaluation passed (launcher Docker unavailable)" + if strict_success + else f"runner in-pod evaluation did not pass: {failure_stage or 'patch not fixed'}" + ), + error_text=str(exc), + artifacts={ + "bench": "patcheval", + "cve_id": metrics.get("cve_id") or None, + "setting": metrics.get("setting"), + "eval_source": "runner_in_pod_fallback", + "strict_success": strict_success, + "poc_passed": metrics.get("poc_passed") is True, + "unit_test_present": metrics.get("unit_test_present") is True, + "unit_tests_passed": metrics.get("unit_tests_passed") is True, + "failure_stage": failure_stage, + "poc_log": _trim_log(metrics.get("poc_log")), + "unit_test_log": _trim_log(metrics.get("unit_test_log")), + "patch": metrics.get("patch"), + "launcher_docker_error": str(exc), + }, + ) + + def _load_official_evaluation(env_params: dict[str, Any]) -> type[Any]: global _ADAPTER_INSTALLED, _OFFICIAL_EVALUATION with _OFFICIAL_LOCK: diff --git a/env/patcheval/strict_runner.py b/env/patcheval/runner.py similarity index 78% rename from env/patcheval/strict_runner.py rename to env/patcheval/runner.py index fee8cf43..566d29e9 100644 --- a/env/patcheval/strict_runner.py +++ b/env/patcheval/runner.py @@ -13,11 +13,51 @@ from typing import Any from urllib.parse import urlsplit, urlunsplit from urllib.request import Request, urlopen +from urllib.error import HTTPError, URLError + + +_RETRY_STATUS_CODES = {429, 500, 502, 503, 504} +_GATEWAY_MAX_RETRIES = 3 +_GATEWAY_BACKOFF_S = (5.0, 10.0, 20.0) + + +def _post_json(url: str, payload: bytes, timeout_s: float) -> dict[str, Any]: + """POST JSON with bounded retry on transient upstream errors. + + The opus-5 upstream proxy intermittently returns 503 ("upstream load + saturated") under load; without retry a single transient failure kills the + whole episode. Retry idempotent chat requests on 429/5xx and connection + errors so the LLM baseline survives flaky upstreams. + """ + last_exc: Exception | None = None + for attempt in range(_GATEWAY_MAX_RETRIES + 1): + try: + request = Request( + url, + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urlopen(request, timeout=timeout_s) as response: + return json.loads(response.read().decode("utf-8")) + except HTTPError as exc: + last_exc = exc + if exc.code not in _RETRY_STATUS_CODES or attempt >= _GATEWAY_MAX_RETRIES: + raise + time.sleep(_GATEWAY_BACKOFF_S[attempt]) + except URLError as exc: + last_exc = exc + if attempt >= _GATEWAY_MAX_RETRIES: + raise + time.sleep(_GATEWAY_BACKOFF_S[attempt]) + raise last_exc # type: ignore[misc] DEFAULT_TIMEOUT_S = 900.0 MAX_LOG_CHARS = 16_000 SETTING_EPOCHS = {"s1.1": 1, "s1.2": 1, "s1.3": 1, "s1.4": 5} +RESULT_JSON_PREFIX = "SAFACTORY_RESULT_JSON " +RESULT_PATH_ENV = "SAFACTORY_RESULT_PATH" LANGUAGE_COMMENT_MAP = { "py": ("Python", "#"), "js": ("JavaScript", "//"), @@ -366,28 +406,59 @@ def _evaluate_official_scripts(patch: str, timeout_s: float) -> dict[str, Any]: return result +def _is_claude_model(model: str) -> bool: + return model.lower().startswith("claude") + + def _call_gateway(*, base_url: str, model: str, prompt: str, timeout_s: float) -> str: - payload = json.dumps( - { - "model": model, - "messages": [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": prompt}, - ], - "temperature": 0, - "max_tokens": 16384, - "stream": False, - } - ).encode("utf-8") - request = Request( - f"{base_url}/chat/completions", - data=payload, - headers={"Content-Type": "application/json"}, - method="POST", + if _is_claude_model(model): + return _call_gateway_anthropic( + base_url=base_url, model=model, prompt=prompt, timeout_s=timeout_s + ) + return _call_gateway_openai( + base_url=base_url, model=model, prompt=prompt, timeout_s=timeout_s ) - with urlopen(request, timeout=timeout_s) as response: - body = json.loads(response.read().decode("utf-8")) - return str(body["choices"][0]["message"]["content"]) + + +def _call_gateway_anthropic(*, base_url: str, model: str, prompt: str, timeout_s: float) -> str: + # Native Anthropic Messages API. The gateway injects the upstream api key + # and anthropic-version, so the runner only sends the body. `temperature` + # is deprecated for claude-opus-5 and rejected by the upstream proxy, so it + # is omitted entirely. + body: dict[str, Any] = { + "model": model, + "max_tokens": 16384, + "system": "You are a helpful assistant", + "messages": [{"role": "user", "content": prompt}], + "stream": False, + } + payload = json.dumps(body).encode("utf-8") + data = _post_json(f"{base_url}/v1/messages", payload, timeout_s) + # content is a list of blocks (e.g. thinking + text); concatenate text parts. + parts = [ + str(block.get("text") or "") + for block in (data.get("content") or []) + if isinstance(block, dict) and block.get("type") == "text" + ] + if not parts: + raise RuntimeError(f"anthropic response had no text content: {data!r}") + return "".join(parts) + + +def _call_gateway_openai(*, base_url: str, model: str, prompt: str, timeout_s: float) -> str: + body: dict[str, Any] = { + "model": model, + "messages": [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": prompt}, + ], + "temperature": 0, + "max_tokens": 16384, + "stream": False, + } + payload = json.dumps(body).encode("utf-8") + data = _post_json(f"{base_url}/chat/completions", payload, timeout_s) + return str(data["choices"][0]["message"]["content"]) def _read_request() -> dict[str, Any]: @@ -401,7 +472,18 @@ def _read_request() -> dict[str, Any]: def _write_result(value: dict[str, Any]) -> None: - print(json.dumps(value, ensure_ascii=False)) + payload = json.dumps(value, ensure_ascii=False) + result_path = os.environ.get(RESULT_PATH_ENV, "").strip() + if result_path: + try: + path = Path(result_path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(payload + "\n", encoding="utf-8") + temporary.replace(path) + except OSError: + pass + print(RESULT_JSON_PREFIX + payload, flush=True) def _required_text(value: Any, name: str) -> str: diff --git a/evaluator/service.py b/evaluator/service.py index cb1af44c..6f601f61 100644 --- a/evaluator/service.py +++ b/evaluator/service.py @@ -76,11 +76,13 @@ async def evaluate(self, request: EvalRequest) -> EvalResult: final_result.session_id = request.session_id log.info( - "EVAL complete: session=%s status=%s score=%.4f elapsed=%.2fs", + "EVAL complete: session=%s status=%s score=%.4f elapsed=%.2fs reason=%s error_text=%s", request.session_id, final_result.status, final_result.normalized_score_10, time.perf_counter() - started_at, + final_result.reason, + final_result.error_text, ) trace.emit_summary( status=final_result.status, diff --git a/gateway/admission_control.py b/gateway/admission_control.py index c4c9237b..d578794e 100644 --- a/gateway/admission_control.py +++ b/gateway/admission_control.py @@ -10,10 +10,20 @@ from gateway.models import GatewayRequestContext, GatewaySessionBinding -@dataclass(frozen=True) class AdmissionRejected(Exception): - reason: str - status_code: int + """Admission rejected by the gateway (draining, queue full, concurrency, etc). + + NB: must NOT be a @dataclass(frozen=True). Python's `raise` statement sets + `__traceback__` on the raised instance; frozen dataclasses forbid attribute + assignment, so `raise AdmissionRejected(...)` would itself raise + `AttributeError: cannot assign to field '__traceback__'` and surface as a + 500 instead of the intended status_code. + """ + + def __init__(self, reason: str, status_code: int) -> None: + super().__init__(reason) + self.reason = reason + self.status_code = status_code @dataclass(frozen=True) @@ -34,6 +44,13 @@ def __init__(self, cfg: GatewayConfig): self._per_route_inflight: dict[str, int] = {} self._request_acquired: set[str] = set() self._route_acquired: set[tuple[str, str]] = set() + # Per-route semaphores: instead of hard-rejecting when a route is at + # max_concurrency, requests WAIT for a slot. Hard 503-rejects fail RL + # episodes on transient concurrency spikes (slow sglang, retry bursts, + # launcher over-subscription). A bounded wait lets episodes queue and + # proceed as soon as a slot frees, which is the correct behavior for an + # RL gateway. Capacity is fixed at first use from target.max_concurrency. + self._route_semaphores: dict[str, asyncio.Semaphore] = {} self.accepted_total = 0 self.rejected_total = 0 @@ -43,54 +60,78 @@ async def acquire_request( binding: GatewaySessionBinding, target: LLMRouteTarget | None = None, ) -> AdmissionDecision: + # Acquire the route concurrency slot BEFORE taking the admission lock: + # the semaphore WAITS when the route is at max_concurrency instead of + # hard-rejecting. Waiting outside the lock means a queued request does + # not block other admissions (draining checks, other routes, etc.). + route_sem: asyncio.Semaphore | None = None + if target is not None and target.max_concurrency > 0: + route_sem = self._route_semaphore(target.route_model, target.max_concurrency) + await route_sem.acquire() + async with self._lock: - if binding.status != "active": - self.rejected_total += 1 - raise AdmissionRejected(f"session is {binding.status}", 409) - if self.draining: - self.rejected_total += 1 - raise AdmissionRejected("gateway is draining", 503) - - if self.cfg.max_steps >= 0 and binding.step_count_for(ctx.requested_model) >= self.cfg.max_steps: - binding.mark_model_truncated( - ctx.requested_model, - "max_steps_reached", - datetime.now(timezone.utc), - ) - self.accepted_total += 1 - return AdmissionDecision(action="stop", stop_reason="max_steps_reached") + try: + if binding.status != "active": + self.rejected_total += 1 + raise AdmissionRejected(f"session is {binding.status}", 409) + if self.draining: + self.rejected_total += 1 + raise AdmissionRejected("gateway is draining", 503) + + if self.cfg.max_steps >= 0 and binding.step_count_for(ctx.requested_model) >= self.cfg.max_steps: + binding.mark_model_truncated( + ctx.requested_model, + "max_steps_reached", + datetime.now(timezone.utc), + ) + self.accepted_total += 1 + if route_sem is not None: + route_sem.release() + return AdmissionDecision(action="stop", stop_reason="max_steps_reached") + + if self._inflight_requests >= self.cfg.max_inflight_requests: + self.rejected_total += 1 + raise AdmissionRejected("gateway inflight limit reached", 503) - if self._inflight_requests >= self.cfg.max_inflight_requests: - self.rejected_total += 1 - raise AdmissionRejected("gateway inflight limit reached", 503) + if ctx.is_stream and self._active_streams >= self.cfg.max_active_streams: + self.rejected_total += 1 + raise AdmissionRejected("gateway active stream limit reached", 503) - if ctx.is_stream and self._active_streams >= self.cfg.max_active_streams: - self.rejected_total += 1 - raise AdmissionRejected("gateway active stream limit reached", 503) + session_inflight = self._per_session_inflight.get(ctx.session_id, 0) + if session_inflight >= self.cfg.per_session_max_inflight: + self.rejected_total += 1 + raise AdmissionRejected("per-session inflight limit reached", 429) - session_inflight = self._per_session_inflight.get(ctx.session_id, 0) - if session_inflight >= self.cfg.per_session_max_inflight: - self.rejected_total += 1 - raise AdmissionRejected("per-session inflight limit reached", 429) + if target is not None: + # The semaphore already guarantees a concurrency slot; just + # track the counter for snapshot/reporting. + route_inflight = self._per_route_inflight.get(target.route_model, 0) + self._per_route_inflight[target.route_model] = route_inflight + 1 + self._route_acquired.add((ctx.request_id, target.route_model)) - if target is not None: - route_inflight = self._per_route_inflight.get(target.route_model, 0) - if route_inflight >= target.max_concurrency: - self.rejected_total += 1 - raise AdmissionRejected("LLM route concurrency limit reached", 503) - self._per_route_inflight[target.route_model] = route_inflight + 1 - self._route_acquired.add((ctx.request_id, target.route_model)) - - self._inflight_requests += 1 - if ctx.is_stream: - self._active_streams += 1 - binding.active_stream_count += 1 - binding.active_request_count += 1 - self._per_session_inflight[ctx.session_id] = session_inflight + 1 - self._request_acquired.add(ctx.request_id) - self.accepted_total += 1 - llm_step_index = binding.increment_step_count(ctx.requested_model) - return AdmissionDecision(action="forward", llm_step_index=llm_step_index) + self._inflight_requests += 1 + if ctx.is_stream: + self._active_streams += 1 + binding.active_stream_count += 1 + binding.active_request_count += 1 + self._per_session_inflight[ctx.session_id] = session_inflight + 1 + self._request_acquired.add(ctx.request_id) + self.accepted_total += 1 + llm_step_index = binding.increment_step_count(ctx.requested_model) + return AdmissionDecision(action="forward", llm_step_index=llm_step_index) + except BaseException: + # If we got a route slot but a later admission check rejected us, + # release the slot so a queued request can proceed. + if route_sem is not None: + route_sem.release() + raise + + def _route_semaphore(self, route_model: str, capacity: int) -> asyncio.Semaphore: + sem = self._route_semaphores.get(route_model) + if sem is None: + sem = asyncio.Semaphore(max(1, capacity)) + self._route_semaphores[route_model] = sem + return sem async def release( self, @@ -124,6 +165,10 @@ async def release( self._per_route_inflight[target.route_model] = route_inflight else: self._per_route_inflight.pop(target.route_model, None) + # Release the concurrency slot so a queued request can proceed. + sem = self._route_semaphores.get(target.route_model) + if sem is not None: + sem.release() async def snapshot(self) -> dict[str, int | bool]: async with self._lock: diff --git a/gateway/app.py b/gateway/app.py index 6b8f579f..7e7738e6 100644 --- a/gateway/app.py +++ b/gateway/app.py @@ -3,6 +3,7 @@ import asyncio import json import logging +import os import time from collections.abc import AsyncIterator, Coroutine from contextlib import asynccontextmanager @@ -31,6 +32,25 @@ log = logging.getLogger("gateway.app") +def _ensure_default_max_tokens(payload: dict[str, Any], default_max_tokens: int) -> None: + """Inject a default max_tokens when the upstream client omits it. + + OpenHands defaults max_output_tokens to 0 ("auto-detect"), which fails for + custom gateway-routed models, so no max_tokens is sent and sglang falls back + to a tiny default (~128), truncating generation mid-tool-call + (finish_reason=length). The gateway is the choke point for every LLM call, + so filling a sane default here fixes it for all agents regardless of their + own config/env-var support. Set default_max_tokens<=0 (e.g. + GATEWAY_DEFAULT_MAX_TOKENS=0) to disable. + Default 32768: large enough that the model's "overthinking" monologue + (~7-9k tokens) is not truncated, while still capping runaway generation. + """ + if "max_tokens" in payload or "max_completion_tokens" in payload: + return + if default_max_tokens > 0: + payload["max_tokens"] = default_max_tokens + + def _without_beta_query(query: str) -> str | None: filtered = [ (name, value) @@ -138,6 +158,7 @@ async def handle_inference_request( payload = await request.json() if not isinstance(payload, dict): raise ValueError("request body must be a JSON object") + _ensure_default_max_tokens(payload, request.app.state.gateway_config.default_max_tokens) with trace.span("resolve_request"): ctx = await resolver.resolve( @@ -560,6 +581,7 @@ async def handle_standard_inference_request( payload = await request.json() if not isinstance(payload, dict): raise ValueError("request body must be a JSON object") + _ensure_default_max_tokens(payload, request.app.state.gateway_config.default_max_tokens) requested_model = payload.get("model") if not isinstance(requested_model, str) or not requested_model: diff --git a/gateway/config.py b/gateway/config.py index 066ee0fc..4fc643a2 100644 --- a/gateway/config.py +++ b/gateway/config.py @@ -1,7 +1,8 @@ from __future__ import annotations import json -from dataclasses import dataclass, fields +import os +from dataclasses import dataclass, fields, replace from typing import Any import yaml @@ -9,6 +10,14 @@ DEFAULT_SQLITE_DB_URL = "sqlite://env_trajs.db" +def _safe_int(value: Any, default: int) -> int: + try: + return int(value) + except (TypeError, ValueError): + return int(default) + + + @dataclass(frozen=True) class LLMRouteConfig: base_url: str @@ -56,6 +65,7 @@ class GatewayConfig: drain_timeout_s: int = 30 session_close_timeout_s: float = 90.0 session_close_retry_after_s: int = 10 + default_max_tokens: int = 32768 storage_type: str = "sqlite" storage_config: dict[str, Any] | None = None llm_routes: dict[str, LLMRouteConfig] | None = None @@ -105,6 +115,13 @@ def load_gateway_config(path: str | None = None) -> GatewayConfig: file_data = _load_file(path) if path else {} cfg = _dict_to_config(file_data) + # Env-var overrides (centralized here so app.py never reads os.environ for + # gateway config). GATEWAY_DEFAULT_MAX_TOKENS overrides the file value; <=0 + # disables default-token injection. + env_max_tokens = os.environ.get("GATEWAY_DEFAULT_MAX_TOKENS") + if env_max_tokens is not None and env_max_tokens.strip() != "": + cfg = replace(cfg, default_max_tokens=_safe_int(env_max_tokens, cfg.default_max_tokens)) + storage_config = _storage_config_for(cfg.storage_type, cfg.storage_config) llm_routes = cfg.llm_routes or _default_routes() diff --git a/gateway/telemetry.py b/gateway/telemetry.py index e0068545..16a8bacf 100644 --- a/gateway/telemetry.py +++ b/gateway/telemetry.py @@ -17,6 +17,25 @@ from gateway.models import GatewayRequestContext, GatewaySessionBinding, GatewayTelemetryRecord from gateway.storage import GatewayStorage +# Structured timing log (rl/timing_log.py). The gateway process does not have +# rl/ on PYTHONPATH (only AIEVOBOX_ROOT), so add it defensively. On import +# failure fall back to a noop so call sites can emit unconditionally; the +# on/off switch lives inside timing_log itself (set_enabled / env +# SAFACTORY_TIMING_LOG_ENABLED), not behind a per-call guard here. +try: + from timing_log import emit as _timing_emit # type: ignore +except Exception: # pragma: no cover - import path fixup + import os as _os + import sys as _sys + _rl_dir = _os.path.join(_os.environ.get("AIEVOBOX_ROOT", ""), "rl") + if _rl_dir and _rl_dir not in _sys.path: + _sys.path.insert(0, _rl_dir) + try: + from timing_log import emit as _timing_emit # type: ignore + except Exception: + def _timing_emit(*_args: Any, **_kwargs: Any) -> None: # type: ignore + return None + log = logging.getLogger("gateway.telemetry") SENSITIVE_KEY_PARTS = ( @@ -160,6 +179,17 @@ async def enqueue_success( self._latest_success_step.get(key, 0), ) + # Per-LLM-step timing for offline analysis. See _emit_llm_step for the + # field semantics; emission is gated inside timing_log. + self._emit_llm_step( + binding, + latency_ms, + upstream_latency_ms, + stream_stats, + 200, + response_body=response_body, + ) + async def enqueue_failure( self, ctx: GatewayRequestContext, @@ -193,6 +223,15 @@ async def enqueue_failure( ) await self._enqueue(binding, record) + self._emit_llm_step( + binding, + latency_ms, + upstream_latency_ms, + stream_stats, + status_code, + error_text=error_text, + ) + async def wait_for_session_flush(self, binding: GatewaySessionBinding) -> None: if self._writer_tasks: future: asyncio.Future[None] = asyncio.get_running_loop().create_future() @@ -400,6 +439,45 @@ async def _record_binding( binding.upstream_base_url = target.base_url binding.last_seen_at = datetime.now(timezone.utc) + def _emit_llm_step( + self, + binding: GatewaySessionBinding, + latency_ms: float, + upstream_latency_ms: float | None, + stream_stats: StreamTelemetryStats | None, + status_code: int, + response_body: dict[str, Any] | None = None, + error_text: str | None = None, + ) -> None: + """Emit one structured ``llm_step`` timing record for offline analysis. + + ``upstream_latency_ms`` is the actual LLM inference time (gateway -> + llm_proxy -> sglang); ``latency_ms`` is the end-to-end step time. Joined + with the worker's episode record (same session_id) to split env-startup + vs rollout vs llm-inference. Emission is gated inside ``timing_log`` + (``set_enabled`` / ``SAFACTORY_TIMING_LOG_ENABLED``), so this call is + unconditional; the import-time noop fallback covers a missing ``rl/`` + on PYTHONPATH. + """ + usage = response_body.get("usage") if isinstance(response_body, dict) else None + usage = usage if isinstance(usage, dict) else {} + _timing_emit( + "llm_step", + session_id=binding.session_id, + group_id=getattr(binding, "group_id", None), + env_name=getattr(binding, "env_name", None), + model=binding.model, + step_index=getattr(binding, "llm_step_count", None), + latency_ms=latency_ms, + upstream_latency_ms=upstream_latency_ms, + ttft_ms=getattr(stream_stats, "ttft_ms", None) if stream_stats else None, + status_code=status_code, + error=error_text, + prompt_tokens=_usage_int(usage, "prompt_tokens", "input_tokens"), + completion_tokens=_usage_int(usage, "completion_tokens", "output_tokens"), + total_tokens=_usage_int(usage, "total_tokens"), + ) + async def _next_seq(self, session_id: str, model: str) -> int: async with self._lock: key = (session_id, model) diff --git a/manager/rjob_episode_runner.py b/manager/rjob_episode_runner.py index 1c2bd7cf..9f6b56e2 100644 --- a/manager/rjob_episode_runner.py +++ b/manager/rjob_episode_runner.py @@ -70,10 +70,12 @@ async def start( logs_error = "" result: SimulationStartResult | None = None timings_ms: Dict[str, float] = {} + timings_abs: Dict[str, float] = {} status_poll_count = 0 summary_status = "failed" summary_extra: Dict[str, Any] = {} episode_started = time.perf_counter() + episode_started_ts = time.time() try: if not lease.image: raise RuntimeError(f"RJob lease missing image: {lease.agent_name}/{lease.agent_id}") @@ -135,6 +137,7 @@ async def start( with trace.span("submit_job", requested_rjob_name=rjob_name): submitted = await self._cluster.submit_job(client, job, submit_kwargs) timings_ms["rjob_submit_ms"] = _elapsed_ms(started) + timings_abs["rjob_submit_ts"] = time.time() trace.update_context(rjob_submit_ms=timings_ms["rjob_submit_ms"]) submitted_name = str(submitted or rjob_name).strip() trace.update_context(submitted_rjob_name=submitted_name) @@ -168,6 +171,13 @@ async def start( trace=trace, ) timings_ms["rjob_wait_terminal_ms"] = _elapsed_ms(started) + # rjob_running_ts is set inside wait_terminal via trace.update_context + # when the RJob first enters Running state. Surface it as an absolute + # epoch timestamp so the episode record can derive cluster queue time + # (rjob_running_ts - rjob_submit_ts). + _running_ts = trace.context.get("rjob_running_ts") + if _running_ts is not None: + timings_abs["rjob_running_ts"] = _running_ts trace.update_context(rjob_status=terminal_status, status_poll_count=status_poll_count) try: started = time.perf_counter() @@ -264,6 +274,7 @@ async def start( result, trace=trace, timings_ms=timings_ms, + timings_abs=timings_abs, submitted_name=submitted_name, terminal_status=terminal_status, status_poll_count=status_poll_count, @@ -389,6 +400,7 @@ def _attach_timing_metrics( *, trace: PerfTrace, timings_ms: Dict[str, float], + timings_abs: Dict[str, float], submitted_name: str, terminal_status: str, status_poll_count: int, @@ -403,6 +415,7 @@ def _attach_timing_metrics( "rjob_status": terminal_status, "rjob_status_poll_count": status_poll_count, **timings_ms, + **timings_abs, } ) diff --git a/manager/simulation_config.py b/manager/simulation_config.py index d9a7ba7f..cee6470e 100644 --- a/manager/simulation_config.py +++ b/manager/simulation_config.py @@ -905,10 +905,18 @@ def _normalize_embedded_file(item: Any, cfg_path: Path) -> Dict[str, str]: def expand_rl_group_size(yaml_config_list: List[Dict[str, Any]], group_size: int) -> List[Dict[str, Any]]: if int(group_size) <= 0: return yaml_config_list + # Oversample: launch more envs per prompt than group_size so the first + # group_size episodes to finish form a group, and long-tail episodes + # don't block the group. buffer_server still pops group_size at a time; + # the surplus stays in the bucket for the next group (or gets discarded + # at rollout end). Set RL_OVERSAMPLE=0 to disable. + oversample = int(os.environ.get("RL_OVERSAMPLE", "0")) + env_num = int(group_size) + oversample expanded = [dict(item) for item in yaml_config_list] for item in expanded: - item["env_num"] = int(group_size) - log.debug("Override agent parallelism env_num=%d for %d config(s)", int(group_size), len(expanded)) + item["env_num"] = env_num + log.debug("Override agent parallelism env_num=%d (group_size=%d + oversample=%d) for %d config(s)", + env_num, int(group_size), oversample, len(expanded)) return expanded diff --git a/manager/simulation_worker.py b/manager/simulation_worker.py index 16cd0231..1337f095 100644 --- a/manager/simulation_worker.py +++ b/manager/simulation_worker.py @@ -3,6 +3,7 @@ import asyncio import inspect import logging +import os import time from collections import deque from typing import Any, Dict, Optional @@ -20,6 +21,25 @@ from evaluator.rule_evaluator import discover_rule_eval_spec from evaluator.service import EvaluationService +# Structured timing log (rl/timing_log.py). The launcher subprocess does not +# have rl/ on PYTHONPATH (only AIEVOBOX_ROOT), so add it defensively. On +# import failure fall back to a noop so call sites emit unconditionally; the +# on/off switch lives inside timing_log (set_enabled / env +# SAFACTORY_TIMING_LOG_ENABLED), not behind a per-call guard here. +try: + from timing_log import emit as _timing_emit # type: ignore +except Exception: # pragma: no cover - import path fixup + import os as _os + import sys as _sys + _rl_dir = _os.path.join(_os.environ.get("AIEVOBOX_ROOT", ""), "rl") + if _rl_dir and _rl_dir not in _sys.path: + _sys.path.insert(0, _rl_dir) + try: + from timing_log import emit as _timing_emit # type: ignore + except Exception: + def _timing_emit(*_args: Any, **_kwargs: Any) -> None: # type: ignore + return None + from .agent_start_client import AgentStartClient from .session_lifecycle import complete_latest_session_step from .simulation_lease_pool import SimulationLeasePool @@ -34,6 +54,18 @@ log = logging.getLogger("manager.simulation_worker") +def _iso_to_epoch(s: str | None) -> float | None: + """Parse an ISO-8601 timestamp (from gateway session status) to epoch seconds.""" + if not s: + return None + try: + from datetime import datetime + + return datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp() + except Exception: + return None + + class _SimulationCircuitBreaker: _TIMEOUT_MARKERS = ( "timed out", @@ -148,6 +180,10 @@ def __init__( self._results: Dict[str, SimulationStartResult] = {} self._results_lock = asyncio.Lock() self._circuit_breaker = _SimulationCircuitBreaker(cfg) + # Number of episodes currently running (acquired a lease, not yet + # released). Used for the periodic active_envs timing snapshot. + self._active_episodes = 0 + self._active_lock = asyncio.Lock() async def run_all(self) -> SimulationRunSummary: log.info( @@ -159,6 +195,10 @@ async def run_all(self) -> SimulationRunSummary: asyncio.create_task(self._worker_loop(worker_id), name=f"simulation-worker-{worker_id}") for worker_id in range(self.worker_count) ] + # Periodic snapshot of concurrent active episodes for capacity / + # GPU-ratio analysis (env pods vs inference GPUs). + snapshot_task = asyncio.create_task(self._active_snapshot_loop(), name="active-envs-snapshot") + tasks.append(snapshot_task) cancelled = False try: await asyncio.gather(*tasks) @@ -211,6 +251,23 @@ async def run_all(self) -> SimulationRunSummary: results={key: result.total_reward for key, result in results.items()}, ) + async def _active_snapshot_loop(self) -> None: + """Emit periodic active-env snapshots for capacity / GPU-ratio analysis.""" + interval = float(os.environ.get("SAFACTORY_ACTIVE_SNAPSHOT_INTERVAL_S", "10") or 10) + try: + while True: + await asyncio.sleep(interval) + async with self._active_lock: + active = self._active_episodes + _timing_emit( + "active_envs", + active_episodes=active, + pool_size=self.lease_pool.pool_size, + worker_count=self.worker_count, + ) + except asyncio.CancelledError: + return + async def _worker_loop(self, worker_id: int) -> None: while True: trace = PerfTrace( @@ -233,6 +290,8 @@ async def _worker_loop(self, worker_id: int) -> None: return agent_key = f"{lease.agent_name}_{lease.agent_id}" + async with self._active_lock: + self._active_episodes += 1 trace.update_context( agent_key=agent_key, agent_name=lease.agent_name, @@ -326,11 +385,17 @@ async def _worker_loop(self, worker_id: int) -> None: release_reusable = False elif self.evaluation_service is not None and self.reward_committer is not None: with trace.span("eval_discover_rule"): - public_env_params = strip_internal_env_params(lease.env_params) + public_env_params = materialize_dataset_env_params( + strip_internal_env_params(lease.env_params) + ) eval_spec = discover_rule_eval_spec( agent_name=lease.agent_name, env_root=self.cfg.agent_root, ) + if eval_spec is not None: + _eval_timeout = float(lease.env_params.get("rule_evaluator_timeout_s") or 0) + if _eval_timeout > 0: + eval_spec.timeout_s = _eval_timeout trace.mark("eval_spec_resolved", eval_spec_found=eval_spec is not None) if eval_spec is not None: log.debug( @@ -351,12 +416,40 @@ async def _worker_loop(self, worker_id: int) -> None: env_params=public_env_params, eval_spec=eval_spec, ) + _eval_started = time.perf_counter() with trace.span("evaluation_service"): eval_result = await self.evaluation_service.evaluate(eval_request) + _eval_elapsed = time.perf_counter() - _eval_started + result.metrics = dict(result.metrics or {}) + result.metrics["eval_elapsed_s"] = round(_eval_elapsed, 3) trace.update_context( eval_status=eval_result.status, eval_score=eval_result.normalized_score_10, ) + # Rule-evaluator timing: for PatchEval this is the + # apply-patch + PoC + unit-test cost, which for + # compiled CVE projects can dominate the episode. + _eval_artifacts = eval_result.artifacts if isinstance(eval_result.artifacts, dict) else {} + _patch = _eval_artifacts.get("patch") or "" + _patch_lines = _patch.count("\n") + 1 if _patch else 0 + _timing_emit( + "eval", + session_id=result.session_id, + env_name=lease.agent_name, + group_id=lease.group_id, + cve_id=_eval_artifacts.get("cve_id"), + eval_status=eval_result.status, + score=eval_result.normalized_score_10, + raw_score=eval_result.raw_score, + reason=eval_result.reason, + eval_elapsed_s=round(_eval_elapsed, 3), + strict_success=_eval_artifacts.get("strict_success"), + poc_passed=_eval_artifacts.get("poc_passed"), + unit_tests_passed=_eval_artifacts.get("unit_tests_passed"), + validation_type=_eval_artifacts.get("validation_type"), + patch_produced=bool(_patch), + patch_lines=_patch_lines or None, + ) if eval_result.status == "succeeded": with trace.span("reward_commit"): await self.reward_committer.commit( @@ -431,6 +524,9 @@ async def _worker_loop(self, worker_id: int) -> None: except Exception as exc: trace.update_context(release_error_type=type(exc).__name__, release_error=str(exc)) log.exception("worker=%d agent=%s critical error in lease_pool.done()", worker_id, agent_key) + finally: + async with self._active_lock: + self._active_episodes = max(0, self._active_episodes - 1) if cancelled: trace.update_context(final_status="cancelled") trace.emit_summary(status="cancelled") @@ -452,6 +548,70 @@ async def _worker_loop(self, worker_id: int) -> None: elapsed, ) + # Structured per-episode timing record for offline analysis. + # rjob_*_ms come from RJobEpisodeRunner._attach_timing_metrics. + # gw_first_seen_ts / gw_closed_ts come from the gateway session + # status captured during _finalize_gateway_session. Together with + # rjob_submit_ts they yield: + # env_startup_s = first LLM call - rjob submit (pod boot + agent init) + # env_active_s = session close - first LLM call (active rollout) + # env_lifecycle_s = session close - rjob submit (full env lifetime) + if result is not None: + ep_metrics = result.metrics if isinstance(result.metrics, dict) else {} + _rjob_submit_ts = ep_metrics.get("rjob_submit_ts") + _gw_first_ts = ep_metrics.get("gw_first_seen_ts") + _gw_closed_ts = ep_metrics.get("gw_closed_ts") + _env_startup = ( + round(_gw_first_ts - _rjob_submit_ts, 3) + if _gw_first_ts is not None and _rjob_submit_ts is not None + else None + ) + _env_active = ( + round(_gw_closed_ts - _gw_first_ts, 3) + if _gw_closed_ts is not None and _gw_first_ts is not None + else None + ) + _env_lifecycle = ( + round(_gw_closed_ts - _rjob_submit_ts, 3) + if _gw_closed_ts is not None and _rjob_submit_ts is not None + else None + ) + _timing_emit( + "episode", + worker_id=worker_id, + env_name=lease.agent_name, + agent_id=lease.agent_id, + group_id=lease.group_id, + session_id=session.session_id if session is not None else lease.agent_id, + runtime=lease.runtime, + status=result.status, + reward=result.total_reward, + step_count=result.step_count, + truncated=bool(result.truncated), + episode_elapsed_s=round(elapsed, 3), + rjob_submit_ms=ep_metrics.get("rjob_submit_ms"), + rjob_wait_terminal_ms=ep_metrics.get("rjob_wait_terminal_ms"), + rjob_fetch_logs_ms=ep_metrics.get("rjob_fetch_logs_ms"), + rjob_parse_result_ms=ep_metrics.get("rjob_parse_result_ms"), + rjob_cleanup_ms=ep_metrics.get("rjob_cleanup_ms"), + rjob_total_ms=ep_metrics.get("rjob_total_ms"), + rjob_status=ep_metrics.get("rjob_status"), + rjob_name=ep_metrics.get("rjob_name"), + # Absolute epoch seconds at RJob submit; join with the + # gateway's first llm_step ts (same session_id) to derive + # env-startup time = first_llm_call - rjob_submit. + rjob_submit_ts=ep_metrics.get("rjob_submit_ts"), + rjob_running_ts=ep_metrics.get("rjob_running_ts"), + # Per-env startup / active / lifecycle durations (seconds), + # derived from gateway session timing + rjob_submit_ts. + env_startup_s=_env_startup, + env_active_s=_env_active, + env_lifecycle_s=_env_lifecycle, + gw_first_seen_ts=_gw_first_ts, + gw_closed_ts=_gw_closed_ts, + eval_elapsed_s=ep_metrics.get("eval_elapsed_s"), + ) + async def _acquire_lease_or_stop(self, worker_id: int) -> SimulationAgentLease | None: del worker_id if self._circuit_breaker.is_open(): @@ -507,14 +667,20 @@ async def _run_one_episode( ) if result.status != "succeeded": + logs_tail = "" + metrics = result.metrics if isinstance(result.metrics, dict) else {} + if metrics.get("logs_tail"): + logs_tail = self._tail(str(metrics.get("logs_tail"))) log.warning( - "worker=%d runtime=%s env=%s agent_id=%s returned status=%s error=%s", + "worker=%d runtime=%s env=%s agent_id=%s returned status=%s error=%s | rjob=%s | LOGS_TAIL:\n%s\n[/LOGS_TAIL]", worker_id, lease.runtime, lease.agent_name, lease.agent_id, result.status, self._tail(result.error_text or ""), + metrics.get("rjob_name", ""), + logs_tail, ) return result @@ -549,6 +715,28 @@ async def _finalize_gateway_session( reason=reason, completion_mode=completion_mode, ) + with trace.span("gateway_wait_telemetry_flush"): + await self.gateway_client.wait_telemetry_flush(result.session_id) + # Capture gateway-side session timing for per-env startup / + # lifecycle analysis. first_seen_at = wall-clock of the FIRST LLM + # call reaching the gateway (= env ready / agent started making + # requests); closed_at = when we closed the session. Joined with + # the worker's rjob_submit_ts (in result.metrics) to derive + # env_startup_s / env_active_s / env_lifecycle_s in the episode + # record. Best-effort: a failure here must not fail finalization. + try: + status = await self.gateway_client.get_session_status(result.session_id) + if isinstance(status, dict): + result.metrics = dict(result.metrics or {}) + result.metrics["gw_first_seen_ts"] = _iso_to_epoch(status.get("first_seen_at")) + result.metrics["gw_closed_ts"] = _iso_to_epoch(status.get("closed_at")) + except Exception as exc: + log.debug( + "worker=%d agent=%s gateway session status fetch failed: %s", + worker_id, + agent_key, + exc, + ) return True except httpx.HTTPError as exc: log.warning( diff --git a/rl/.gitignore b/rl/.gitignore new file mode 100644 index 00000000..c6b313a2 --- /dev/null +++ b/rl/.gitignore @@ -0,0 +1,7 @@ +# Local-only utility / one-off experiment scripts (kept on disk, not tracked). +# - cleanup_rl.sh: one-shot "kill all RL residual processes" helper. +# - restart_pool_test.sh: one-off POOL_SIZE sweep experiment (hardcoded 27B). +# - collect_pool_metrics.sh: companion metrics collector for the sweep above. +cleanup_rl.sh +restart_pool_test.sh +collect_pool_metrics.sh diff --git a/rl/buffer_server.py b/rl/buffer_server.py index e7ba685c..9037d6a0 100644 --- a/rl/buffer_server.py +++ b/rl/buffer_server.py @@ -22,6 +22,7 @@ from utils import get_env import gateway_autostart +from timing_log import emit as timing_emit import uvicorn from fastapi import FastAPI, HTTPException, Request @@ -34,10 +35,21 @@ from core.data_manager.manager import DataManager -# Setup logging -LOG_DIR = os.path.join(AIEVOBOX_ROOT, "logs") +# Setup logging — write into the per-run directory (same place as slime.log, +# e.g. logs/patcheval_qwen3_8_27b/20260825-221716/) so each run's logs are +# co-located instead of appending to a flat cross-run file. run_buffer_server.sh +# exports AIEVOBOX_RUN_DIR (from .current_run or a fresh timestamp) before +# launching us; fall back to the flat logs/ dir if it isn't set. +_RUN_DIR = os.environ.get("AIEVOBOX_RUN_DIR", "").strip() +LOG_DIR = _RUN_DIR or os.path.join(AIEVOBOX_ROOT, "logs") os.makedirs(LOG_DIR, exist_ok=True) LOG_FILE = os.path.join(LOG_DIR, "buffer_server.log") +# Co-locate the gateway's own log in the run dir too. The gateway subprocess +# (launched below via gateway_autostart) inherits this env and reads +# SAFACTORY_GATEWAY_LOG_PATH at startup (gateway/__main__.py). Only override +# when running per-run (don't clobber an explicit user override). +if _RUN_DIR and not os.environ.get("SAFACTORY_GATEWAY_LOG_PATH"): + os.environ["SAFACTORY_GATEWAY_LOG_PATH"] = os.path.join(LOG_DIR, "gateway.log") logger = logging.getLogger("buffer_server") logger.setLevel(logging.DEBUG) @@ -73,8 +85,10 @@ # DataManager for querying the database data_manager: Optional[DataManager] = None -# Track last served step ID for cursor-based pagination -last_served_id: int = 0 +# Env-id cursor for finished-environment fetch. Advances forward only: +# finished=True implies all steps are already is_terminal=True, so there is no +# late-flip window to re-scan (see list_environment_rows + list_terminal_steps_for_sessions). +last_env_cursor: int = 0 # Pending items by instance_id (for grouping) pending_items_by_instance: Dict[str, List[Dict[str, Any]]] = {} @@ -162,9 +176,14 @@ def _assistant_message_from_stored_response(response: Any) -> Optional[Dict[str, def _build_item_from_row(row: Dict[str, Any]) -> Dict[str, Any]: - """Convert a database row to the expected item format.""" + """Convert a database row to the expected item format. + + Accepts the raw format from list_terminal_steps_for_sessions (keys + ``messages``/``prompt``, ``meta_json``/``env_state``, ``session_id``/``env_id``, + ``created_at``/``session_end_time``, ``is_truncated``/``truncated``). + """ # Parse stored prompt (JSON serialized messages list) - prompt_str = row.get("prompt", "") + prompt_str = row.get("prompt") or row.get("messages", "") if isinstance(prompt_str, str): base_messages = json.loads(prompt_str) if prompt_str else [] else: @@ -177,12 +196,13 @@ def _build_item_from_row(row: Dict[str, Any]) -> Dict[str, Any]: messages.append(assistant) session_id = row.get("session_id", "") - env_id = row.get("env_id", "") + env_id = row.get("env_id") or session_id group_id = row.get("group_id", "") - # 从 env_state 中解析 weight_version + # 从 env_state / meta_json 中解析 weight_version weight_version = 0 - if env_state_raw := row.get("env_state"): + env_state_raw = row.get("env_state") or row.get("meta_json") + if env_state_raw: try: env_state = json.loads(env_state_raw) if isinstance(env_state_raw, str) else env_state_raw raw_weight_version = env_state.get("weight_version") if isinstance(env_state, dict) else None @@ -191,16 +211,18 @@ def _build_item_from_row(row: Dict[str, Any]) -> Dict[str, Any]: except (TypeError, ValueError, json.JSONDecodeError): weight_version = 0 + truncated = row.get("truncated", row.get("is_truncated", False)) + end_time = row.get("session_end_time") or row.get("created_at") extra_info = { - "timestamp": _parse_timestamp(row.get("session_end_time")) or _parse_timestamp(row.get("timestamp")) or time.time(), + "timestamp": _parse_timestamp(end_time) or _parse_timestamp(row.get("timestamp")) or time.time(), "steps": row.get("step_id", 0), # 注意:finish_reason 与 truncated 不完全等价,finish_reason 仅用于训练侧标记截断状态 - "finish_reason": "length" if row.get("truncated", False) else "stop", + "finish_reason": "length" if truncated else "stop", "session_id": session_id, "env_id": env_id, "group_id": group_id, "weight_version": weight_version, - "truncated": row.get("truncated", False), + "truncated": truncated, } return { @@ -213,34 +235,57 @@ def _build_item_from_row(row: Dict[str, Any]) -> Dict[str, Any]: async def fetch_new_items_from_db(limit: Optional[int] = None) -> List[Dict[str, Any]]: - """Fetch new completed steps from the database using cursor-based pagination.""" - global data_manager, last_served_id + """Fetch new completed steps using an env-id cursor (two-phase fetch). + + Phase 1: discover finished environments with id > last_env_cursor. + Phase 2: fetch their terminal steps via list_terminal_steps_for_sessions. + Because mark_environment_finished is only called after all steps are + is_terminal=True, finished=True guarantees all training-ready steps are + terminal — no late-flip window, no dedup needed. + """ + global data_manager, last_env_cursor if data_manager is None: return [] - items = [] try: - rows = await data_manager.fetch_done_steps_with_context( - after_id=last_served_id, - limit=limit or 100 + envs = await data_manager.list_environment_rows( + after_id=last_env_cursor, + finished=True, + limit=limit or 100, ) except Exception as e: - logger.error(f"fetch_done_steps_with_context error: {e}") + logger.error(f"list_environment_rows error: {e}") + return [] + + if not envs: + return [] + + env_ids = [str(e.get("env_id") or "") for e in envs if e.get("env_id")] + next_cursor = max(int(e.get("id") or 0) for e in envs) + if not env_ids: + if next_cursor > last_env_cursor: + last_env_cursor = next_cursor return [] + try: + rows = await data_manager.list_terminal_steps_for_sessions(env_ids) + except Exception as e: + logger.error(f"list_terminal_steps_for_sessions error: {e}") + return [] + + items = [] for row in rows: - step_pk = row.get("step_pk") try: item = _build_item_from_row(row) items.append(item) - # Update cursor to the latest processed id - if not last_served_id or step_pk > last_served_id: - last_served_id = step_pk except Exception as e: logger.error(f"Error building item from row: {e}") continue + if next_cursor > last_env_cursor: + last_env_cursor = next_cursor + return items @@ -338,15 +383,16 @@ async def get_rollout_data(request: Request): async def init_data_manager(job_session: str, storage_type: str, db_url: str, restart_training: bool = False): """Initialize the DataManager for querying the database.""" - global data_manager, last_served_id + global data_manager, last_env_cursor data_manager = DataManager(job_id=job_session, storage_type=storage_type, db_url=db_url) await data_manager.init() logger.info(f"DataManager initialized with {storage_type} DB: {db_url}, job_session: {job_session}") # Initialize cursor based on restart_training flag if restart_training: - last_served_id = await data_manager.get_max_step_id() - logger.info(f"restart_training=True, initialized last_served_id={last_served_id}") + envs = await data_manager.list_environment_rows(finished=True, limit=100000) + last_env_cursor = max((int(e.get("id") or 0) for e in envs), default=0) + logger.info(f"restart_training=True, initialized last_env_cursor={last_env_cursor}") def start_aievobox_process(data: dict): @@ -355,7 +401,7 @@ def start_aievobox_process(data: dict): NOTE: LLM Proxy is now hosted in-process by slime_generator. It must already be running before this function is called. """ - global aievobox_process, group_size, last_served_id, pending_items_by_instance, data_manager + global aievobox_process, group_size, last_env_cursor, pending_items_by_instance, data_manager # Set group size (num_repeat_per_sample) group_size = int(data.get("num_repeat_per_sample", 16)) @@ -395,6 +441,7 @@ def start_aievobox_process(data: dict): agent_root = get_env("AIEVOBOX_AGENT_ROOT") or "env" agent_config = os.environ.get("AIEVOBOX_AGENT_CONFIG") agent_start_config = os.environ.get("AIEVOBOX_AGENT_START_CONFIG") + rjob_config = str(get_env("AIEVOBOX_RJOB_CONFIG") or "").strip() # v2 docker/rjob runs need the container startup definition (env_types). When # not set explicitly, derive it from the agent config path: # env//_config.yaml -> env//_start.yaml. @@ -408,15 +455,46 @@ def start_aievobox_process(data: dict): llm_temperature = float(get_env("LLM_TEMPERATURE") or 1.0) pool_size = int(get_env("AIEVOBOX_POOL_SIZE") or 16) rl_epoch = int(get_env("RL_EPOCH") or 1) + docker_image_archive_dir = str( + get_env("AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR") or "" + ).strip() + docker_pull_policy = str( + get_env("AIEVOBOX_DOCKER_PULL_POLICY") or "never" + ).strip() + agent_start_timeout_s = str( + get_env("AIEVOBOX_AGENT_START_TIMEOUT_S") or "" + ).strip() evaluation_flag = str(os.environ.get("AIEVOBOX_ENABLE_EVALUATION") or "").strip().lower() evaluation_enabled = evaluation_flag in {"1", "true", "yes", "on"} + # Circuit breaker: tolerate a large batch of failures before tripping. The + # breaker's min_samples is capped at window (simulation_worker.py), so both + # must be raised together. Default 240 lets the first 240 episodes fail + # without stopping scheduling (useful while the model is weak / during + # bring-up). The consecutive-timeout limit must also be raised: the default + # of 5 trips as soon as 5 episodes eval-timeout in a row, which is normal + # during bring-up (e.g. npm/network flakiness in rule evaluators). Override + # via AIEVOBOX_CIRCUIT_BREAKER_WINDOW / AIEVOBOX_CIRCUIT_BREAKER_MIN_SAMPLES / + # AIEVOBOX_CIRCUIT_BREAKER_CONSECUTIVE_TIMEOUTS; set 0 to fall back to + # launcher defaults (window=50, min_samples=20, consecutive=5). + cb_window = int(get_env("AIEVOBOX_CIRCUIT_BREAKER_WINDOW") or 240) + cb_min_samples = int(get_env("AIEVOBOX_CIRCUIT_BREAKER_MIN_SAMPLES") or 240) + cb_consecutive_timeouts = int(get_env("AIEVOBOX_CIRCUIT_BREAKER_CONSECUTIVE_TIMEOUTS") or 240) + + # Gateway close timeout: must exceed the gateway's drain_timeout_s (30s) or + # the runner abandons the close before the gateway finishes draining in-flight + # LLM requests, leaving sessions unsealed (is_terminal=0) and orphaning rollout + # groups. Default 45s > 30s drain. Override via AIEVOBOX_GATEWAY_CLOSE_TIMEOUT_S. + gateway_close_timeout_s = float(get_env("AIEVOBOX_GATEWAY_CLOSE_TIMEOUT_S") or 45.0) + cmd = [ "python3", launcher_script, "--mode", mode, + *(["--rjob-config", rjob_config] if rjob_config else []), "--db-path", db_url, "--storage-type", storage_type, - *(["--agent-config", agent_config] if agent_config else ["--agent-root", agent_root]), + "--agent-root", agent_root, + *(["--agent-config", agent_config] if agent_config else []), *(["--agent-start-config", agent_start_config] if agent_start_config else []), *(["--enable-evaluation"] if evaluation_enabled else []), "--gateway-base-url", gateway_base_url, @@ -428,6 +506,33 @@ def start_aievobox_process(data: dict): "--no-rebuild-table", "--rl-group-size", str(group_size), "--rl-epoch", str(rl_epoch), + "--docker-pull-policy", docker_pull_policy, + *( + ["--agent-start-timeout-s", agent_start_timeout_s] + if agent_start_timeout_s + else [] + ), + *( + ["--docker-image-archive-dir", docker_image_archive_dir] + if docker_image_archive_dir + else [] + ), + *( + ["--circuit-breaker-window", str(cb_window)] + if cb_window > 0 + else [] + ), + *( + ["--circuit-breaker-min-samples", str(cb_min_samples)] + if cb_min_samples > 0 + else [] + ), + *( + ["--circuit-breaker-consecutive-timeouts", str(cb_consecutive_timeouts)] + if cb_consecutive_timeouts > 0 + else [] + ), + "--gateway-close-timeout-s", str(gateway_close_timeout_s), ] logger.info(f"Starting launcher.py: {' '.join(cmd)}") @@ -442,6 +547,31 @@ def start_aievobox_process(data: dict): stderr=None, # Inherit stderr ) logger.info(f"launcher.py started with PID: {aievobox_process.pid}") + + # Record the RL capacity / GPU-ratio config once per rollout start so + # the timing log is self-describing for offline capacity analysis. + timing_emit( + "rl_config", + mode=get_env("AIEVOBOX_MODE") or "docker", + pool_size=pool_size, + group_size=group_size, + max_steps=max_steps, + rollout_num_gpus=int(get_env("ROLLOUT_NUM_GPUS") or 0), + rollout_num_gpus_per_engine=int(get_env("ROLLOUT_NUM_GPUS_PER_ENGINE") or 0), + num_gpus=int(get_env("NUM_GPUS") or 0), + actor_num_gpus_per_node=int(get_env("ACTOR_NUM_GPUS_PER_NODE") or 0), + global_batch_size=int(get_env("RL_GLOBAL_BATCH_SIZE") or 0), + rollout_batch_size=int(get_env("SLIME_ROLLOUT_BATCH_SIZE") or 0), + num_rollout=int(get_env("NUM_ROLLOUT") or 0), + rl_epoch=int(get_env("RL_EPOCH") or 0), + gateway_max_steps=int(get_env("AIEVOBOX_GATEWAY_MAX_STEPS") or -1), + # Ratio of concurrent env pods to inference GPUs (envs per GPU). + envs_per_inference_gpu=( + pool_size / int(get_env("ROLLOUT_NUM_GPUS") or 1) + if int(get_env("ROLLOUT_NUM_GPUS") or 0) > 0 + else None + ), + ) except Exception as e: logger.error(f"Failed to start launcher.py: {e}") raise @@ -469,6 +599,76 @@ async def start_rollout(request: Request): return {"message": "Rollout started"} +@app.post("/stop_rollout") +async def stop_rollout(): + """Kill the AIEvoBox launcher process to stop all envs. + + Called by slime_generator after collecting enough rollout data, + before the training step begins. This ensures no envs are still + sending LLM requests to SGLang, so flush_cache can succeed immediately. + """ + global aievobox_process + + if aievobox_process is None or aievobox_process.poll() is not None: + logger.info("[stop_rollout] AIEvoBox not running, nothing to stop") + return {"message": "AIEvoBox not running"} + + pid = aievobox_process.pid + logger.info(f"[stop_rollout] Killing AIEvoBox process tree (pid={pid})") + + try: + import signal + import os + import subprocess + + # Kill only the launcher.py process and its children, NOT the entire + # process group (which would also kill the buffer server itself). + # Use ps to find all descendant PIDs, then kill them individually. + try: + # Find all child/descendant processes of the launcher + result = subprocess.run( + ["ps", "--ppid", str(pid), "-o", "pid=", "--no-header"], + capture_output=True, text=True, timeout=5, + ) + child_pids = [int(p.strip()) for p in result.stdout.split() if p.strip()] + + # Recursively find grandchildren + all_pids = list(child_pids) + for child_pid in child_pids: + try: + result2 = subprocess.run( + ["ps", "--ppid", str(child_pid), "-o", "pid=", "--no-header"], + capture_output=True, text=True, timeout=5, + ) + all_pids.extend(int(p.strip()) for p in result2.stdout.split() if p.strip()) + except Exception: + pass + + # Kill children first (bottom-up), then the launcher itself + for kill_pid in all_pids: + try: + os.kill(kill_pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + + # Finally kill the launcher process itself + aievobox_process.kill() + except (ProcessLookupError, PermissionError): + aievobox_process.kill() + + aievobox_process.wait(timeout=10) + logger.info(f"[stop_rollout] AIEvoBox process {pid} killed successfully") + except Exception as e: + logger.warning(f"[stop_rollout] Error killing AIEvoBox: {e}") + try: + aievobox_process.kill() + except Exception: + pass + + aievobox_process = None + return {"message": "AIEvoBox stopped", "pid": pid} + + @app.get("/health") async def health_check(): """Health check endpoint.""" diff --git a/rl/examples/harbor/.gitignore b/rl/examples/harbor/.gitignore new file mode 100644 index 00000000..98eb6713 --- /dev/null +++ b/rl/examples/harbor/.gitignore @@ -0,0 +1 @@ +wandb_logs/ \ No newline at end of file diff --git a/rl/examples/harbor/README.md b/rl/examples/harbor/README.md new file mode 100644 index 00000000..4db83aca --- /dev/null +++ b/rl/examples/harbor/README.md @@ -0,0 +1,84 @@ +# Harbor RL + +一个 Harbor (vulhub-exploit) 环境的 RL 训练样例,setting 对齐 `rl/examples/patcheval`。 + +## 与 PatchEval 的区别 + +| 维度 | PatchEval | Harbor | +|---|---|---| +| 环境 | 77 个 CVE,每个一个 env_name,env_num=300 | 单个 env_name `harbor`,dataset 474 个 vulhub 任务 | +| 配置生成 | 需先跑 `env/patcheval/generate_full_config.py` 生成 generated dir | 无生成步骤,rjob 配置直接提交在 `env/harbor/` | +| Agent | openhands | claude-code(`env_params.agent`,runner.py 把 `ANTHROPIC_BASE_URL` 指向 RL gateway) | +| 镜像 | 每个 CVE 一个 tarball,从 archive dir 加载 | 单一 runtime 镜像 `safactory-harbor-runtime-v0.21.0`,RJob pod 从 registry 拉 | +| Reward | patcheval rule 校验 | Harbor verifier 0/1,`rule_evaluator.py` 归一化到 0-10 | +| 单 episode 时长 | ~40 LLM 步 | timeout_s=9000(2.5h),更长 | + +## 硬件要求 + +- 2 台 8 卡 H200 机器(训练机 + 推理机,共 16 卡) +- 训练机:Megatron TP=4,跑 8 卡 +- 推理机:SGLang 8 引擎,每引擎 1 卡 + +## 训练配置 + +| 参数 | 值 | 说明 | +|---|---|---| +| 模型 | Qwen3.8-27B | GQA,TP 必须整除 num_query_groups=4 | +| TP_SIZE | 4 | 张量并行 | +| POOL_SIZE | 16 | 并发环境数 | +| RL_GROUP_SIZE | 8 | 每个 vulhub 任务采样 8 条轨迹 | +| RL_ROLLOUT_GROUP_BATCH_SIZE | 8 | 每批 8 个任务 | +| RL_GLOBAL_BATCH_SIZE | 64 | 每步训练 64 条轨迹 | +| RL_EPOCH | 100 | 训练轮数(474 任务 / 8 ≈ 59 batch/epoch) | +| MAX_TOKENS_PER_GPU | 2048 | 微批 token 上限 | +| TRAJ_TRUNCATION_MAX_SEQ_LEN | 8192 | 训练时截断长轨迹,防 OOM | +| OPTIMIZER_CPU_OFFLOAD | true | 优化器卸到 CPU,省 ~40GB 显存 | +| SGLANG_MEM_FRACTION_STATIC | 0.7 | KV cache 池占比 | +| GATEWAY_MAX_STEPS | 60 | 每条轨迹最大 LLM 步数(vulhub 比 CVE 补丁长) | +| AGENT_START_TIMEOUT_S | 3600 | RJob pod 启动超时(嵌套 dockerd + 镜像拉取) | +| Vulhub 任务 | 474 个 | `env/harbor/datasets/harbor_vulhub_all.jsonl` | + +## 启动 + +```bash +# 推理机 +ray start --address="<训练机IP>:6379" --num-gpus=8 --disable-usage-stats + +# 训练机 +ray start --head --node-ip-address="<训练机IP>" --port=6379 --num-gpus=8 --disable-usage-stats + +# 训练机 - 窗口1:buffer +export HARBOR_VARIANT=vulhub_all +export RL_ENV_SH=$PWD/rl/examples/harbor/env.rjob.sh +export CLEANUP_BEFORE_RUN=false +bash rl/run_buffer_server.sh + +# 训练机 - 窗口2:训练(等 gateway 起来后) +export SKIP_RAY_START=true +export MASTER_ADDR="<训练机IP>" +bash rl/run_slime_generator.sh +``` + +### 变体 + +`HARBOR_VARIANT` 选择不同的 harbor rjob 配置(都已在 `env/harbor/` 提交): + +| 变体 | config / start | 任务数 | 用途 | +|---|---|---|---| +| `vulhub_all`(默认) | `harbor_vulhub_all_*` | 474 | 正式训练 | +| `cvebench` | `harbor_cvebench_*` | 1 | cvebench smoke,端到端联调 | +| `smoke` | `harbor_*`(无后缀) | 1 | oracle smoke,验证 rjob 管线 | + +## 注意事项 + +- TP 不能设 8(GQA 约束:num_query_groups=4 必须被 TP 整除) +- 启动前确认 8000 端口空闲,否则 gateway 起不来导致 0 轨迹 +- Harbor episode 比 PatchEval 长得多(timeout_s=9000 vs 900),单步训练时间显著更长,先小 POOL_SIZE / 小 NUM_ROLLOUT 跑通再放量 +- reward 全 0 是正常的(基座模型难解 vulhub),有解出才有学习信号 +- `AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR` 留空:harbor runtime 镜像由 RJob pod 从 registry 拉,不需要本地 tarball +- RJob pod 是 privileged(嵌套 Docker 跑 vulnerable 服务),已在 `*_start.rjob.yaml` 里设 `privileged: true` + +## TODO + +- GATEWAY_MAX_STEPS=60 是初值,需根据实际 vulhub 解出率 / 吞吐调优 +- LLM_MAX_LENGTH=131072:vulhub 多轮 exploit 轨迹长,需确认 SGLang 显存够用,必要时降 POOL_SIZE diff --git a/rl/examples/harbor/check_rjob.py b/rl/examples/harbor/check_rjob.py new file mode 100644 index 00000000..4ff59e04 --- /dev/null +++ b/rl/examples/harbor/check_rjob.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +"""Query RJob cluster for harbor jobs: status + logs. Run on the training machine.""" +import sys, yaml, traceback + +CFG = "/mnt/shared-storage-user/leishanzhe/repo/SAfactory/config.yaml" + +def main(): + with open(CFG) as f: + cfg = yaml.safe_load(f) or {} + rjob = cfg.get("rjob", {}) + try: + from brainpp.rjob import RJobClient + except ImportError: + print("ERROR: brainpp.rjob not installed in this env"); sys.exit(1) + + client = RJobClient( + cluster_entry=rjob.get("cluster_entry"), + namespace=rjob.get("namespace"), + access_key=rjob.get("access_key"), + secret_key=rjob.get("secret_key"), + verifyssl=bool(rjob.get("verifyssl", True)), + retries=int(rjob.get("retries", 3) or 0), + ) + print(f"RJobClient -> {rjob.get('cluster_entry')} ns={rjob.get('namespace')}") + + # list all jobs, filter harbor + try: + jobs = client.list([]) + except Exception: + # some versions need a name filter; try with prefix + try: + jobs = client.list("harbor") + except Exception: + jobs = client.list("safactory") + print(f"\n=== total jobs returned: {len(jobs) if hasattr(jobs,'__len__') else '?'} ===") + + harbor_jobs = [] + for j in (jobs or []): + name = getattr(j, "name", None) or (j.get("name") if isinstance(j, dict) else str(j)) + if "harbor" in str(name).lower() or "safactory" in str(name).lower(): + harbor_jobs.append(j) + + print(f"harbor/safactory jobs: {len(harbor_jobs)}") + if not harbor_jobs: + print("No harbor jobs found in cluster. Launcher may not be submitting, or jobs already cleaned.") + # show first few of any jobs + print("\n=== first 10 of ALL jobs ===") + for j in (jobs or [])[:10]: + print(" ", getattr(j, "name", None) or (j.get("name") if isinstance(j, dict) else j)) + return + + print("\n=== harbor job statuses ===") + for j in harbor_jobs[:30]: + name = getattr(j, "name", None) or (j.get("name") if isinstance(j, dict) else str(j)) + status = getattr(j, "status", None) or (j.get("status") if isinstance(j, dict) else "?") + print(f" {name} status={status}") + + # fetch logs for first 2 non-succeeded + print("\n=== logs for first 2 harbor jobs ===") + for j in harbor_jobs[:2]: + name = getattr(j, "name", None) or (j.get("name") if isinstance(j, dict) else str(j)) + print(f"\n----- LOGS: {name} -----") + try: + raw = client.logs_rjob(name) + txt = raw if isinstance(raw, str) else getattr(raw, "text", None) or str(raw) + print(txt[-4000:] if len(txt) > 4000 else txt) + except Exception as e: + print(f"logs_rjob failed: {e}") + traceback.print_exc() + +if __name__ == "__main__": + main() diff --git a/rl/examples/harbor/env.rjob.sh b/rl/examples/harbor/env.rjob.sh new file mode 100755 index 00000000..c6bc3cb1 --- /dev/null +++ b/rl/examples/harbor/env.rjob.sh @@ -0,0 +1,292 @@ +#!/usr/bin/env bash +# +# ============================================================================= +# [RJOB MODE] Harbor RL environment +# ============================================================================= +# Harbor (vulhub-exploit) RL training settings for rl/run_buffer_server.sh and +# rl/run_slime_generator.sh. Mirrors rl/examples/patcheval/env.rjob.sh but +# targets the Harbor environment (env/harbor/), whose runner.py drives the +# `harbor` CLI to spin up vulnerable services inside a nested-Docker RJob pod +# and scores the agent's exploit via Harbor's verifier. +# +# Key differences vs PatchEval: +# 1. No "generated dir" step. PatchEval needs generate_full_config.py to +# materialize one env per CVE (77 envs x 300 copies). Harbor ships a +# single env_name "harbor" with a JSONL dataset of 474 vulhub tasks +# (env/harbor/datasets/harbor_vulhub_all.jsonl); the buffer server fans +# the dataset out to per-task rollout groups. So AIEVOBOX_AGENT_CONFIG / +# AIEVOBOX_AGENT_START_CONFIG point directly at the committed yamls in +# env/harbor/. +# 2. Agent is `claude-code` (set in env_params of harbor_vulhub_all_config). +# runner.py rewrites ANTHROPIC_BASE_URL -> the RL gateway URL, so the +# agent's LLM calls hit the model under training. No per-task image +# archive is needed; the single harbor runtime image +# (safactory-harbor-runtime-v0.21.0) is pulled by the RJob pod. +# 3. Reward: Harbor verifier emits a 0/1 reward; rule_evaluator.py +# (env/harbor/rule_evaluator.py) normalizes it onto SAfactory's 0-10 +# scale. Same binary-reward dynamic as PatchEval -> base model rarely +# solves a vulhub task, so most groups are all-0 (DAPO_filter stays off +# by default to avoid stalling the buffer). +# 4. Tasks are LONG: timeout_s=9000 (2.5h) per episode in the vulhub config. +# AIEVOBOX_GATEWAY_MAX_STEPS and AIEVOBOX_AGENT_START_TIMEOUT_S are raised +# accordingly vs PatchEval. +# +# Usage: +# export HARBOR_VARIANT=vulhub_all # vulhub_all (default) | cvebench | smoke +# RL_ENV_SH=$this rl/run_buffer_server.sh +# ============================================================================= +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# Self-contained: same infrastructure defaults that patcheval/env.rjob.sh +# inlines from geo3k_vl/env.sh. Only the vars actually consumed by +# run_slime_generator.sh / buffer_server.py / llm_proxy.py / slime_generator.py +# are set here; no geo3k-specific values are carried over. + +: "${HARBOR_VARIANT:?Set HARBOR_VARIANT to one of: vulhub_all | cvebench | smoke}" + +# --- Harbor agent config selection ----------------------------------------- +# Harbor commits its rjob configs directly in env/harbor/ (no generation step). +# _config.rjob.yaml -> AIEVOBOX_AGENT_CONFIG (env_types/datasets) +# _start.rjob.yaml -> AIEVOBOX_AGENT_START_CONFIG (container/rjob spec) +# rule_evaluator.py lives next to them, so AIEVOBOX_AGENT_ROOT = env/harbor. +HARBOR_ENV_DIR="${REPO_ROOT}/env/harbor" +case "${HARBOR_VARIANT}" in + vulhub_all) + HARBOR_CONFIG="${HARBOR_ENV_DIR}/harbor_vulhub_all_config.rjob.yaml" + HARBOR_START_CONFIG="${HARBOR_ENV_DIR}/harbor_vulhub_start.rjob.yaml" + HARBOR_DATASET_N=474 + ;; + cvebench) + # cvebench smoke (oracle, 1 task). Useful for end-to-end bring-up only. + HARBOR_CONFIG="${HARBOR_ENV_DIR}/harbor_cvebench_config.rjob.yaml" + HARBOR_START_CONFIG="${HARBOR_ENV_DIR}/harbor_cvebench_start.rjob.yaml" + HARBOR_DATASET_N=1 + ;; + smoke) + # harbor oracle smoke (1 task). Useful for verifying the rjob pipeline. + HARBOR_CONFIG="${HARBOR_ENV_DIR}/harbor_config.rjob.yaml" + HARBOR_START_CONFIG="${HARBOR_ENV_DIR}/harbor_start.rjob.yaml" + HARBOR_DATASET_N=1 + ;; + *) + echo "Unknown HARBOR_VARIANT='${HARBOR_VARIANT}'. Use vulhub_all | cvebench | smoke." >&2 + exit 1 + ;; +esac + +if [[ ! -f "${HARBOR_CONFIG}" || ! -f "${HARBOR_START_CONFIG}" ]]; then + echo "Missing Harbor rjob config for variant '${HARBOR_VARIANT}':" >&2 + echo " config: ${HARBOR_CONFIG}" >&2 + echo " start : ${HARBOR_START_CONFIG}" >&2 + exit 1 +fi + +# --- AIEVOBOX / env wiring ------------------------------------------------- +export AIEVOBOX_EXAMPLE_NAME="harbor_qwen3_8_27b" +export AIEVOBOX_ROOT="${AIEVOBOX_ROOT:-${REPO_ROOT}}" +export AIEVOBOX_MODE="rjob" +export AIEVOBOX_RJOB_CONFIG="${AIEVOBOX_RJOB_CONFIG:-${REPO_ROOT}/config.yaml}" +export STORAGE_TYPE="${STORAGE_TYPE:-sqlite}" +export AIEVOBOX_DB_URL="${HARBOR_DB_URL:-sqlite:///${AIEVOBOX_ROOT}/rl/examples/harbor/harbor_qwen3_8_27b.db}" +# Harbor runtime image is pulled by the RJob pod from the registry; no local +# archive dir is needed (unlike PatchEval's per-CVE tarballs). Keep empty so +# the launcher skips archive mounting. +export AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR="${HARBOR_IMAGE_ARCHIVE_DIR:-}" +export AIEVOBOX_DOCKER_PULL_POLICY="${AIEVOBOX_DOCKER_PULL_POLICY:-always}" +export AIEVOBOX_AGENT_CONFIG="${HARBOR_CONFIG}" +export AIEVOBOX_AGENT_START_CONFIG="${HARBOR_START_CONFIG}" +export AIEVOBOX_AGENT_ROOT="${HARBOR_ENV_DIR}" +# Per-task rollout rounds (NOT LLM steps per episode). Each harbor task is +# rolled out this many times per rollout step. +export AIEVOBOX_MAX_STEPS="${HARBOR_MAX_STEPS:-1}" +export AIEVOBOX_ENABLE_EVALUATION="${AIEVOBOX_ENABLE_EVALUATION:-1}" +# RJob scales across the cluster; 16 concurrent episodes keeps SGLang decode +# batches full (same rationale as PatchEval). Override via HARBOR_POOL_SIZE. +export AIEVOBOX_POOL_SIZE="${HARBOR_POOL_SIZE:-16}" +# Harbor episodes are long (timeout_s up to 9000s). Give the RJob pod ample +# startup headroom (nested dockerd + image pull + harbor init). +export AIEVOBOX_AGENT_START_TIMEOUT_S="${HARBOR_AGENT_START_TIMEOUT_S:-3600}" +# Hard cap on LLM steps per episode, enforced by the RL gateway. Vulhub +# exploit tasks need more explore/exploit steps than PatchEval CVE patches; +# 60 is a first guess, tune per throughput (each step is a gateway round-trip). +export AIEVOBOX_GATEWAY_MAX_STEPS="${HARBOR_GATEWAY_MAX_STEPS:-60}" + +# --- RL / GRPO batch sizing ------------------------------------------------ +# Same shape as PatchEval: 8 trajectories per task, 8 tasks per rollout batch, +# 64 trajectories per training step. With 474 vulhub tasks one epoch covers +# 474/8 ~ 59 rollout batches. Override via HARBOR_* if needed. +export RL_GROUP_SIZE="${HARBOR_GROUP_SIZE:-8}" +export RL_GLOBAL_BATCH_SIZE="${HARBOR_GLOBAL_BATCH_SIZE:-64}" +export RL_ROLLOUT_GROUP_BATCH_SIZE="${HARBOR_ROLLOUT_GROUP_BATCH_SIZE:-8}" +export SLIME_ROLLOUT_BATCH_SIZE="${HARBOR_SLIME_ROLLOUT_BATCH_SIZE:-${RL_ROLLOUT_GROUP_BATCH_SIZE}}" +export SLIME_GLOBAL_BATCH_SIZE="${HARBOR_SLIME_GLOBAL_BATCH_SIZE:-${RL_GLOBAL_BATCH_SIZE}}" +export RL_EPOCH="${HARBOR_EPOCH:-100}" +export RL_MODEL="${RL_MODEL:-model}" +export RL_API_KEY="${RL_API_KEY:-openai_api_key}" + +# --- Networking: buffer / proxy / gateway --------------------------------- +export BUFFER_SERVER_HOST="${BUFFER_SERVER_HOST:-127.0.0.1}" +export BUFFER_SERVER_PORT="${BUFFER_SERVER_PORT:-18889}" +export LLM_PROXY_HOST="${LLM_PROXY_HOST:-127.0.0.1}" +export LLM_PROXY_PORT="${LLM_PROXY_PORT:-18890}" +# Vulhub trajectories are long (multi-turn exploit + verification). Raise the +# rollout max response length vs PatchEval so the gateway doesn't truncate +# agent reasoning mid-exploit. 131072 = SGLang default max batch token budget. +export LLM_MAX_LENGTH="${LLM_MAX_LENGTH:-131072}" +export LLM_TEMPERATURE="${LLM_TEMPERATURE:-1.0}" +# Gateway runs on THIS training pod (started by the buffer server via +# gateway_autostart). Default to this pod's IP so RJob pods can reach it. +export AIEVOBOX_GATEWAY_HOST="${HARBOR_GATEWAY_HOST:-$(hostname -i | awk '{print $1}')}" +export AIEVOBOX_GATEWAY_PORT="${HARBOR_GATEWAY_PORT:-8000}" +export AIEVOBOX_GATEWAY_BASE_URL="http://${AIEVOBOX_GATEWAY_HOST}:${AIEVOBOX_GATEWAY_PORT}/v1/sessions" + +# --- Slime / Megatron / model ---------------------------------------------- +export SLIME_HOME="${SLIME_HOME:-/root/slime}" +export MEGATRON_HOME="${MEGATRON_HOME:-/root/Megatron-LM}" +# Model: Qwen3.8-27B (GQA, num_query_groups=4 -> TP must divide 4, so TP=4). +export HF_CKPT_DIR="${QWEN3_8_27B_CKPT_DIR:-/mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/hub/models--Qwen--Qwen3.8-27B/snapshots/1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0}" +export LOAD_DIR="${QWEN3_8_27B_LOAD_DIR:-${HF_CKPT_DIR}}" +export SAVE_DIR="${HARBOR_SAVE_DIR:-${AIEVOBOX_ROOT}/rl/examples/harbor/checkpoints/Qwen3.8-27B_megatron}" +export WANDB_DIR="${HARBOR_WANDB_DIR:-${AIEVOBOX_ROOT}/rl/examples/harbor/wandb_logs}" +export LOG_ROOT="${HARBOR_LOG_ROOT:-${AIEVOBOX_ROOT}/logs/harbor_qwen3_8_27b}" +export MODEL_SCRIPT="${QWEN3_8_27B_MODEL_SCRIPT:-/root/slime/scripts/models/qwen3.5-27B.sh}" +export MODEL_ARGS_ROTARY_BASE=10000000 + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export NUM_GPUS="${HARBOR_NUM_GPUS:-8}" + +# Multi-node NCCL: disable IB (ibv_modify_qp fails across some node pairs) and +# force TCP socket transport. Override via NCCL_* env vars if needed. +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" +export NCCL_NET="${NCCL_NET:-Socket}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-bond0}" +export ACTOR_NUM_NODES=1 +export ACTOR_NUM_GPUS_PER_NODE="${HARBOR_ACTOR_NUM_GPUS_PER_NODE:-8}" +export ROLLOUT_NUM_GPUS="${HARBOR_ROLLOUT_NUM_GPUS:-8}" +export ROLLOUT_NUM_GPUS_PER_ENGINE=1 +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-${SLIME_HOME}/train.py}" +export ROLLOUT_FUNCTION_PATH="${ROLLOUT_FUNCTION_PATH:-rl.slime_generator.generate_rollout}" +export NUM_ROLLOUT="${NUM_ROLLOUT:-10}" +export LOSS_MASK_TYPE="qwen3_5" +export TRAIN_BACKEND="${TRAIN_BACKEND:-megatron}" +export MEGATRON_TO_HF_MODE="${MEGATRON_TO_HF_MODE:-bridge}" +export TP_SIZE="${HARBOR_TP_SIZE:-4}" PP_SIZE="${HARBOR_PP_SIZE:-1}" CP_SIZE=1 EP_SIZE=1 ETP_SIZE=1 +export RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}" +export RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}" +export RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-64}" +export ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}" +export MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-2048}" +# Trajectory truncation for training: long agent trajectories (60-step vulhub +# exploit) can exceed 50k tokens and OOM the training GPU. Truncates each +# trajectory to the last N tokens for training only; the full trajectory is +# still used for reward/advantage computation during rollout. 0 = disable. +export TRAJ_TRUNCATION_MAX_SEQ_LEN="${TRAJ_TRUNCATION_MAX_SEQ_LEN:-8192}" +# GDN packed-seq monkey-patch (see rl/patches/gdn_packed_seq.py). +export PYTHONPATH="${REPO_ROOT}/rl/patches${PYTHONPATH:+:${PYTHONPATH}}" +export USE_DYNAMIC_BATCH_SIZE="${USE_DYNAMIC_BATCH_SIZE:-true}" +export CALCULATE_PER_TOKEN_LOSS="${CALCULATE_PER_TOKEN_LOSS:-true}" +export ADVANTAGE_ESTIMATOR="${ADVANTAGE_ESTIMATOR:-grpo}" +# DAPO group filter: drop groups where all samples share the same reward. +# Harbor reward is binary (0/1) and the base model rarely solves a vulhub +# task, so most groups are all-0 -> filtering would stall the buffer. Off. +export DAPO_filter="${HARBOR_DAPO_FILTER:-false}" +export LR="${LR:-1e-6}" +export OPTIMIZER="${OPTIMIZER:-adam}" +export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +export ADAM_BETA1="${ADAM_BETA1:-0.9}" +export ADAM_BETA2="${ADAM_BETA2:-0.98}" +# Colocate: training (Megatron TP=8) and inference (sglang) share all 8 GPUs. +export SLIME_COLOCATE="${SLIME_COLOCATE:-false}" +# CPU offload optimizer: moves fp32 master weights + Adam states (~41GB at +# TP=4) to CPU. Critical for 27B on 140GB GPUs. +export OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-true}" +export USE_WANDB="${USE_WANDB:-true}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_PROJECT="${WANDB_PROJECT:-slime}" +export WANDB_GROUP="${WANDB_GROUP:-harbor_qwen3_8_27b}" +# KV cache pool fraction. 0.7 gives enough KV capacity to hold POOL_SIZE=16 +# concurrent multi-turn episodes without eviction. Safe on H200 141GB. +export SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.7}" +export SGLANG_ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-fa3}" +export SGLANG_LOG_LEVEL="${SGLANG_LOG_LEVEL:-info}" +export SGLANG_LOG_LEVEL_HTTP="${SGLANG_LOG_LEVEL_HTTP:-error}" +export CLEANUP_BEFORE_RUN="${CLEANUP_BEFORE_RUN:-true}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export PYTHONUNBUFFERED=1 + +# ============================================================================= +# Infrastructure defaults — inlined from rl/examples/geo3k_vl/env.sh. +# Only vars actually consumed by run_slime_generator.sh / buffer_server.py / +# llm_proxy.py / slime_generator.py. Placed at the end so ${VAR:-default} can +# reference harbor values set earlier (POOL_SIZE, RL_GROUP_SIZE, ports). +# ============================================================================= + +# --- Ray / Python launcher --- +export PYTHON_BIN="${PYTHON_BIN:-python3}" +export RAY_BIN="${RAY_BIN:-ray}" +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +export RAY_ADDRESS="${RAY_ADDRESS:-http://127.0.0.1:8265}" +export RAY_PORT="${RAY_PORT:-}" +export KILL_PYTHON_BEFORE_RUN="${KILL_PYTHON_BEFORE_RUN:-false}" + +# --- Slime train / checkpoint args --- +export SAVE_INTERVAL="${SAVE_INTERVAL:-20}" +export MODEL_ARGS_EXTRA="${MODEL_ARGS_EXTRA:-}" +export REF_LOAD_DIR="${REF_LOAD_DIR:-}" +export CUSTOM_REWARD_POST_PROCESS_PATH="${CUSTOM_REWARD_POST_PROCESS_PATH:-}" +export SGLANG_LOGGING_CONFIG_PATH="${SGLANG_LOGGING_CONFIG_PATH:-}" + +# --- Optimizer / GRPO extras --- +export LR_DECAY_STYLE="${LR_DECAY_STYLE:-constant}" +export ENTROPY_COEF="${ENTROPY_COEF:-0.00}" +export EPS_CLIP="${EPS_CLIP:-0.2}" +export EPS_CLIP_HIGH="${EPS_CLIP_HIGH:-0.2}" +export USE_DYNAMIC_GLOBAL_BATCH_SIZE="${USE_DYNAMIC_GLOBAL_BATCH_SIZE:-false}" + +# --- W&B extras --- +export WANDB_TEAM="${WANDB_TEAM:-}" +export WANDB_ALWAYS_USE_TRAIN_STEP="${WANDB_ALWAYS_USE_TRAIN_STEP:-false}" + +# --- SGLang extras --- +# Tuned for high-concurrency rollout (POOL_SIZE=16), same as PatchEval. +export SGLANG_CUDA_GRAPH_BS="${SGLANG_CUDA_GRAPH_BS:-1 2 4 8 16 32}" +export SGLANG_MAX_RUNNING_REQUESTS="${SGLANG_MAX_RUNNING_REQUESTS:-64}" +export SGLANG_SCHEDULE_CONSERVATIVENESS="${SGLANG_SCHEDULE_CONSERVATIVENESS:-}" +export SGLANG_CHUNKED_PREFILL_SIZE="${SGLANG_CHUNKED_PREFILL_SIZE:-8192}" +export SGLANG_ENABLE_MIXED_CHUNK="${SGLANG_ENABLE_MIXED_CHUNK:-false}" + +# --- LLM proxy / buffer server workers & perf --- +export LLM_TOP_P="${LLM_TOP_P:-1.0}" +export LLM_PROXY_ENABLE_CONSOLE_LOG="${LLM_PROXY_ENABLE_CONSOLE_LOG:-0}" +export AIEVOBOX_LLM_MAX_CONCURRENCY="${AIEVOBOX_LLM_MAX_CONCURRENCY:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_PROXY_WORKERS="${AIEVOBOX_LLM_PROXY_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_STARTUP_JITTER_S="${AIEVOBOX_LLM_STARTUP_JITTER_S:-0}" +export AIEVOBOX_TRAININFO_WORKERS="${AIEVOBOX_TRAININFO_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE="${AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE:-256}" +export AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S="${AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S:-0.01}" +export AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS="${AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS:-1800}" +export ROLLBUF_HOST="${ROLLBUF_HOST:-${BUFFER_SERVER_HOST}}" +export ROLLBUF_PORT="${ROLLBUF_PORT:-${BUFFER_SERVER_PORT}}" + +# --- Slime rollout-buffer / GRPO filter --- +export SLIME_ROLLBUF_RESTART_TRAINING="${SLIME_ROLLBUF_RESTART_TRAINING:-True}" +export SLIME_N_SAMPLES_PER_PROMPT="${SLIME_N_SAMPLES_PER_PROMPT:-${RL_GROUP_SIZE}}" +export RL_OFF_BY_N="${RL_OFF_BY_N:-0}" + +# --- AIEVOBOX env extras --- +export AIEVOBOX_MESSAGE_CUT="${AIEVOBOX_MESSAGE_CUT:-0}" +export AIEVOBOC_MULTIPLIER="${AIEVOBOC_MULTIPLIER:-1.2}" + +# --- Runtime --- +# NOTE: expandable_segments:True is incompatible with torch_memory_saver +# (used in colocate mode). Disable it when colocate is on. +if [[ "${SLIME_COLOCATE:-false}" == "true" || "${SLIME_COLOCATE:-false}" == "1" ]]; then + export PYTORCH_CUDA_ALLOC_CONF="" + export PYTORCH_ALLOC_CONF="" +else + # PyTorch >= 2.5 renamed PYTORCH_CUDA_ALLOC_CONF -> PYTORCH_ALLOC_CONF. + # Set both so old and new versions pick up expandable_segments. + export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +fi diff --git a/rl/examples/patcheval/.gitignore b/rl/examples/patcheval/.gitignore index c1f36519..750108ed 100644 --- a/rl/examples/patcheval/.gitignore +++ b/rl/examples/patcheval/.gitignore @@ -2,4 +2,10 @@ official-results/* export_core_trajectories.py PATCHEVAL_DB_FIELDS.md *.jsonl -*.json \ No newline at end of file +*.json + +wandb_logs/ + +generated_openhands_exp1_js77/ +patcheval_eval_gateway.yaml +patcheval_eval_gateway.db diff --git a/rl/examples/patcheval/README.md b/rl/examples/patcheval/README.md index e8fcaf91..58af20b0 100644 --- a/rl/examples/patcheval/README.md +++ b/rl/examples/patcheval/README.md @@ -1,232 +1,356 @@ -# PatchEval Environment +# PatchEval RL + +一个 cyber env 的 RL 训练样例。 + +## 硬件要求 + +- 4 台 8 卡 H200 机器(1 训练 head + 3 推理 worker,共 32 卡) +- 训练 head:`10.102.242.51`,Megatron TP=4 × PP=2 × CP=4,跑 8 卡 +- 推理 worker:`10.102.217.14`、`10.102.217.27`、`10.102.217.42`,SGLang colocate +- colocate 模式:训练 + 推理共享全部 32 卡 + +## 训练配置 + +| 参数 | 值 | 说明 | +|---|---|---| +| 模型 | Qwen3.8-27B | GQA,TP 必须整除 num_query_groups=4 | +| TP_SIZE | 4 | 张量并行 | +| PP_SIZE | **2** | 流水线并行,64 层拆 2×32 层 | +| CP_SIZE | **4** | 上下文并行,处理长序列 | +| MEGATRON_TO_HF_MODE | **raw** | 对齐官方,CP 需要 raw 模式 | +| SLIME_COLOCATE | **true** | 训练+推理共享 GPU | +| POOL_SIZE | 16 | 并发环境数 | +| RL_GROUP_SIZE | 8 | 每个 CVE 采样 8 条轨迹 | +| RL_ROLLOUT_GROUP_BATCH_SIZE | 8 | 每批 8 个 CVE | +| RL_GLOBAL_BATCH_SIZE | 64 | = 8×8 | +| RL_OFF_BY_N | 1 | 允许 1 版本偏差(rollout_id 从 1 开始) | +| MAX_TOKENS_PER_GPU | 8192 | 对齐官方 | +| RECOMPUTE_NUM_LAYERS | 1 | full recompute,对齐官方 | +| TRAJ_TRUNCATION_MAX_SEQ_LEN | **0**(禁用) | 不截断,完整保留长轨迹 | +| OPTIMIZER_CPU_OFFLOAD | true | 优化器卸到 CPU,省 ~40GB 显存 | +| SGLANG_MEM_FRACTION_STATIC | 0.75 | KV cache 池占比 | +| SGLANG_MAMBA_SCHEDULER_STRATEGY | extra_buffer | 修复 mamba pool 问题 | +| SGLANG_SPECULATIVE_ALGORITHM | **EAGLE** | 投机解码,对齐官方 | +| ROLLOUT_MAX_RESPONSE_LEN | 32768 | 单轮生成上限,与 LLM_MAX_LENGTH 分离 | +| LLM_MAX_LENGTH | 65536 | 轨迹上限(实际~40K) | +| LOAD_DIR | Megatron 格式 | raw 模式加载 Megatron checkpoint | +| CVE 任务 | 77 个 JS | 每个 300 副本 | +| max_steps | 40 | 每条轨迹最大 LLM 步数 | + +## 启动 + +### 第 1 步:启动 Ray 集群 -PatchEval 现在统一为 SAfactory 调用链: +```bash +# 训练机(10.102.242.51)先起 head +ray stop --force +ray start --head --port=6379 --num-gpus=8 --num-cpus=100 --disable-usage-stats -- **正式评测**:SAfactory Launcher、Runner 和 Gateway 负责生成与轨迹记录; - Launcher 按约定自动发现 `rule_evaluator.py`,每个 CVE 都由官方 - `evaluation/run_evaluation.py:Evaluation` 在运行资源释放前评分。 -原版 PatchEval 提供两类 baseline: +# 3 台推理机分别执行(不需要 --node-ip-address) +ray stop --force +ray start --address=10.102.242.51:6379 --num-gpus=8 --num-cpus=100 --disable-usage-stats +``` -- **LLM baseline**:将漏洞知识和相关代码包装进 prompt,由 LLM 直接生成补丁; - LLM 不能使用仓库浏览或编辑工具。 -- **Agent baseline**:SWE-agent、OpenHands 或 Claude Code 等 agent 在容器仓库 - 中运行,可以在有限工具调用次数内搜索、读取和修改完整 codebase。 +### 第 2 步:验证集群 -所有配置均由 `generate_full_config.py` 动态生成。LLM baseline 使用 -`strict_runner.py`;当前已接入的 Claude Code Exp1 使用 -`claudecode_runner.py`。 -`run_eval.sh` 先将官方 helper 同步到远程 Docker 可见的 -`patcheval-runtime` 共享目录,再只读挂载到容器的 `/opt/patcheval`。Runner -不再复制实现 `LLMClient.build_prompt`、`PatchParser`、`FuncReplacer`、 -`CodeApplier` 和 `FeedbackHelper` 的逻辑。 +```bash +# 训练机上跑,必须看到 4 个节点各 8 GPU +python3 -c " +import ray +ray.init(address='auto') +for n in ray.nodes(): + r = n['Resources'] + print(f\"Node: {n['NodeManagerAddress']} Alive: {n['Alive']} GPU: {r.get('GPU', 0)}\") +ray.shutdown() +" +``` -## 1. Environment Data +### 第 3 步:训练机窗口 1 — 启动 buffer server -运行时生成的 `patcheval_config.yaml` 为每个任务指定 Docker 镜像和 JSONL: +```bash +cd /mnt/shared-storage-user/leishanzhe/repo/SAfactory +export PATCH_EVAL_GENERATED_DIR=$PWD/rl/examples/patcheval/generated_openhands_exp1_js77 +export RL_ENV_SH=$PWD/rl/examples/patcheval/env.rjob.sh +export CLEANUP_BEFORE_RUN=false +bash rl/run_buffer_server.sh +``` -```yaml -- env_name: patcheval_ - env_image: - env_num: 1 - dataset: ./datasets/.jsonl +### 第 4 步:训练机窗口 2 — 启动训练(等 gateway 起来后) + +```bash +cd /mnt/shared-storage-user/leishanzhe/repo/SAfactory +export PATCH_EVAL_GENERATED_DIR=$PWD/rl/examples/patcheval/generated_openhands_exp1_js77 +export RL_ENV_SH=$PWD/rl/examples/patcheval/env.rjob.sh +export CLEANUP_BEFORE_RUN=false +export SKIP_RAY_START=true +bash rl/run_slime_generator.sh ``` -严格模式 JSONL 每行包含: +> **关键**:窗口 2 必须设 `SKIP_RAY_START=true`,否则脚本会 `ray stop` 把已建好的集群杀掉。 + +## 清理 -```json -{ - "cve_id": "CVE-YYYY-NNNN", - "work_dir": "/workspace/", - "setting": "s1.1", - "prompt_template": "", - "official_record": {"cve_id": "...", "vul_func": []} -} +```bash +# 训练机和推理机都跑一遍 +bash rl/cleanup_rl.sh ``` -任务元数据由官方 `datasets/input.json` 提供;SWE-Agent `dataset.jsonl` 只 -用于补充镜像名和容器内仓库路径。 +## 注意事项 -### 能否直接使用原版 PatchEval 数据 +- TP 不能设 8(GQA 约束:num_query_groups=4 必须被 TP 整除) +- 启动前确认 8000 端口空闲,否则 gateway 起不来导致 0 轨迹 +- reward 全 0 是正常的(基座模型难解 CVE),有解出才有学习信号 +- **启动前必须验证 Ray 集群**:两个节点不同 IP + 各 8 GPU,否则 SPREAD 会把所有 bundle 放到一台机器导致 Duplicate GPU -`generate_full_config.py` 按 CVE ID 合并官方 `input.json` 与 Agent -`dataset.jsonl`,不会使用 `problem_statement` 重新构造 prompt。 +## PP vs CP 对比 -## 2. Environment–LLM Interaction +| | PP=2 | CP=2 | +|---|------|------| +| 切分方式 | 按层切(Stage 0: 0-31层, Stage 1: 32-63层) | 按序列长度切(各处理一半 token) | +| 权重显存 | **减半** ✅ | 不变 ❌ | +| 激活显存 | **减半** ✅ | **减半** ✅ | +| 长序列支持 | ✅ | ✅ | +| 效率 | 有流水线 bubble ⚠️ | 无 bubble ✅ | +| GDN 兼容 | ✅ 原生支持 | ❌ 需要 raw 模式 + hack | +| 实现难度 | 简单 | 复杂 | + +**选择 PP=2 的原因**:Qwen3.8-27B 有 64 层 GDN,单卡 TP=4 装不下全部权重 + 长序列激活。PP=2 权重和激活都减半,GDN 原生兼容,不需要 hack。 -### Environment 输入 +## Monkey-Patch 补丁说明 -Runner 接收 `SimulationStartRequest` JSON,包含 session、任务数据、Gateway -地址、模型名、temperature 和 timeout。 +所有补丁在 `rl/patches/` 目录,通过 `sitecustomize.py` 在 Python 启动时自动加载,无需修改 slime/Megatron 源码。 -### LLM 输入 +| 补丁文件 | 作用 | 加载条件 | +|---|---|---| +| `gdn_packed_seq.py` | 让 GDN 支持 packed sequence(thd 格式) | 总是加载 | +| `traj_truncation.py` | 轨迹截断(TRAJ_TRUNCATION_MAX_SEQ_LEN>0 时生效) | 总是加载 | +| `raw_hf_checkpoint.py` | 允许 raw 模式直接加载 HF checkpoint | 仅 `MEGATRON_TO_HF_MODE=raw` | +| `spread_placement.py` | Ray placement group 策略 PACK→SPREAD(多机分配) | 总是加载 | +| `flush_cache_fix.py` | flush_cache 用 SGLang 内置 `?timeout=60`(v5) | 总是加载 | -Runner 用官方 `LLMClient.build_prompt` 和对应 S1.x 模板生成 user message, -并发送官方 system message `You are a helpful assistant`,固定 -`temperature=0`、`max_tokens=16384`。请求 URL 为当前 SAfactory session -的 `/v1/sessions//chat/completions`。 +> **注意**:`attention_mask_fix.py` 已删除(2026-09-05),仅 bridge 模式需要,raw 模式不需要。 + +## 2026-09-03 更新日志 -### LLM 输出 +### 配置变更 -LLM 输出官方函数级 JSON:`[{"id": "vul_*", "patch": "..."}]`。Runner -直接调用官方 `PatchParser`、`FuncReplacer` 和 `CodeApplier` 解析输出、 -按行范围替换函数并生成 unified diff。 +| 参数 | 之前 | 现在 | 原因 | +|------|------|------|------| +| PP_SIZE | 1 | **2** | 流水线并行,权重+激活减半 | +| CP_SIZE | 2 | **1** | PP 替代 CP,GDN 原生兼容 | +| MEGATRON_TO_HF_MODE | raw | **bridge** | PP 需要 bridge 模式 | +| MAX_TOKENS_PER_GPU | 1024 | **2048** | PP 省显存,可增大 | +| TRAJ_TRUNCATION_MAX_SEQ_LEN | 8192 | **0** | 不截断,完整保留长轨迹 | -### Environment 输出 +### 新建文件 -Runner 返回补丁和生成阶段 metrics。启用 `--enable-evaluation` 后, -`rule_evaluator.py` 调用官方 evaluator:PoC 与单测均通过时 -`raw_score=1`、SAfactory 标准化 reward 为 `10`,否则均为 `0`。该结果直接 -写入当前 trajectory,不再导出补丁或执行批量回写。 +1. `rl/patches/attention_mask_fix.py` — 修复 bridge 模式 `preprocess_packed_seqs` 收到 `attention_mask=None` 的 bug(import hook) +2. `rl/examples/patcheval/generated_openhands_exp1_js77/` — 重新生成 77 个 JS CVE 的 rjob 配置(env_num=300) -## 3. How the Environment Validates a Patch +### 修改文件 -Runner 将官方函数级输出转换为 unified diff,并写入: +1. `rl/patches/sitecustomize.py` — `raw_hf_checkpoint` 改为仅 raw 模式加载;新增 `attention_mask_fix` 注册 +2. `rl/examples/patcheval/env.rjob.sh` — PP=2 / CP=1 / bridge / MAX_TOKENS=2048 / TRAJ_TRUNCATION=0 -```text -/workspace/fix.patch -``` +### 解决的问题 -随后环境执行以下步骤。 +| 问题 | 根因 | 修复 | +|------|------|------| +| Duplicate GPU (PP=2) | 推理机没加入 Ray 集群,16 bundle 全挤训练机 | 在推理机启动 `ray start --address` | +| KeyError: 0 (PP=2) | `raw_hf_checkpoint` 补丁在 bridge 模式也生效,层映射不兼容 PP | 改为仅 raw 模式加载 | +| AttributeError: 'NoneType' (PP=2) | slime 传 `attention_mask=None`,bridge 模型需要有效 tensor | `attention_mask_fix.py` 补丁(import hook) | -### 3.1 验证漏洞是否修复 +## 2026-09-04 更新日志 -```bash -bash /workspace/fix-run.sh -``` +### 配置变更(对齐官方 Qwen3.5-27B 脚本) -`fix-run.sh` 由每个 PatchEval Docker 镜像提供。它会应用候选补丁和安全测试, -然后运行该 CVE 对应的 PoC 回归测试。例如 Gogs 的验证逻辑是: +| 参数 | 之前 | 现在 | 原因 | +|------|------|------|------| +| TP_SIZE | 4 | 4 | 不变(GQA 约束) | +| PP_SIZE | 2 | **2** | 不变 | +| CP_SIZE | 1 | **4** | 对齐官方,CP 处理长序列 | +| MEGATRON_TO_HF_MODE | bridge | **raw** | CP 需要 raw 模式 | +| RECOMPUTE_NUM_LAYERS | 32 | **1** | 对齐官方,full recompute | +| MAX_TOKENS_PER_GPU | 2048 | **8192** | 对齐官方 | +| SLIME_COLOCATE | false | **true** | 训练+推理共享 GPU | +| SGLANG_MEM_FRACTION_STATIC | 0.7 | **0.75** | 对齐官方 | +| ROLLOUT_NUM_GPUS_PER_ENGINE | 1 | **2** | 对齐官方 | +| DECODER_LAST_PIPELINE_NUM_LAYERS | - | **30** | 对齐官方,PP 不均匀切分 | +| SGLANG_MAMBA_SCHEDULER_STRATEGY | auto | **extra_buffer** | 修复 mamba pool illegal memory access | +| SGLANG_SPECULATIVE_ALGORITHM | - | **关闭** | raw 模式下 EAGLE 导致 mamba_pool 崩溃(详见下方) | +| ROLLOUT_NUM_GPUS | 16 | **不传**(colocate 自动=32) | 对齐官方,用满 32 卡 | +| RL_GROUP_SIZE | 8 | **2** | 减小长尾影响 | +| RL_ROLLOUT_GROUP_BATCH_SIZE | 8 | **1** | 减小长尾影响 | +| RL_GLOBAL_BATCH_SIZE | 64 | **2** | = group × batch | +| ROLLOUT_MAX_RESPONSE_LEN | =LLM_MAX_LENGTH(131072) | **32768** | 限制单轮生成,避免长尾卡死 offload | +| ROLLOUT_NUM_PROCESS | 100(默认) | **2**(=global_batch_size) | 避免 offload 时大量 pending 请求卡死 flush | +| LOAD_DIR | HF 格式 | **Megatron 格式** | 对齐官方,raw 模式加载 Megatron checkpoint | -```bash -cd /workspace/gogs -git apply /workspace/test.patch /workspace/fix.patch -go test -run Test_isRepositoryGitPath -``` +### 新增参数 -- 返回码非 0:安全测试失败。 -- 返回码为 0:漏洞攻击已被阻止,继续运行普通单元测试。 +- `ROLLOUT_MAX_RESPONSE_LEN=32768` — 单轮生成 token 上限,与 `LLM_MAX_LENGTH`(轨迹上限)分离 +- `ROLLOUT_NUM_PROCESS` — 并发 env 进程数,默认 = `RL_GLOBAL_BATCH_SIZE`,避免 flush_cache 超时 -因此 `poc_passed=true` 表示安全验证通过,不表示攻击成功。 +### 新增补丁 -### 3.2 检查原有功能 +- `rl/patches/flush_cache_fix.py` — flush_cache 前先 abort 所有 pending 请求,避免长尾 env 卡死 offload(colocate 模式必需) -如果镜像存在 `/workspace/unit_test.sh`,环境继续执行: +### 修改文件 -```bash -bash /workspace/unit_test.sh -``` +1. `rl/examples/patcheval/env.rjob.sh` — CP=4 / raw / colocate / recompute=1 / MAX_TOKENS=8192 / EAGLE 关闭 / mamba extra_buffer / group=2 / ROLLOUT_MAX_RESPONSE_LEN=32768 / LOAD_DIR=Megatron格式 +2. `rl/run_slime_generator.sh` — colocate 时不传 `--rollout-num-gpus`(自动=actor GPUs);`--rollout-max-response-len` 改用 `ROLLOUT_MAX_RESPONSE_LEN`;新增 `--rollout-num-process`;EAGLE 条件传递;mamba/decoder-last-pipeline 参数传递;GRPO args 对齐官方 +3. `rl/patches/sitecustomize.py` — 新增 `flush_cache_fix` 注册 +4. `rl/patches/flush_cache_fix.py` — flush_cache 前 abort pending 请求 +5. HF checkpoint 已转换为 Megatron 格式(`/mnt/shared-storage-user/evobox-share-gpfs2/leishanzhe/model/qwen3_8_27b_megatron`),raw 模式直接加载 -- 安全测试通过、单元测试失败:strict success 为 `false`。 -- 安全测试和单元测试都通过:strict success 为 `true`。 -- 没有 `unit_test.sh`:安全测试通过后 strict success 为 `true`。 +### 解决的问题 -Gateway 保存 prompt、模型回答、token 和延迟;Evaluator 只提交官方二值 -strict-success reward,不再提交 1/7/10 阶段 reward。 +| 问题 | 根因 | 修复 | +|------|------|------| +| `TimeoutError: Timeout while flushing cache`(第一次) | 单轮 `max_tokens=131072`,长尾轨迹跑满 128K token(~9分钟),offload 时 flush_cache 60秒超时 | 新增 `ROLLOUT_MAX_RESPONSE_LEN=32768`,单轮最多 32K token(~2分钟) | +| `TimeoutError: Timeout while flushing cache`(第二次) | `num_process=100`(默认),收齐 2 条轨迹后还有 98 个 env 在跑,SGLang 有 pending 请求,flush_cache 60秒等不完 | 新增 `--rollout-num-process`,默认 = `RL_GLOBAL_BATCH_SIZE`(当前=2);新增 `flush_cache_fix.py` 补丁,flush 前 abort 所有 pending 请求 | +| `CUDA error: illegal memory access` (mamba_pool, rollout 阶段) | SGLang mamba cache pool 默认 `auto` 策略不兼容 | 设 `SGLANG_MAMBA_SCHEDULER_STRATEGY=extra_buffer` | +| `CUDA error: illegal memory access` (mamba_pool, log_probs 阶段) | EAGLE 投机解码 + mamba + raw 模式 + colocate flush/wake_up 循环不兼容:flush_cache 清了主模型 mamba state,但 EAGLE draft model 的 mamba state 不一致 → log_probs 时非法内存访问 | **关闭 EAGLE**(`SGLANG_SPECULATIVE_ALGORITHM` 留空)。官方脚本用 bridge 模式 EAGLE 正常,但 raw 模式下不兼容 | +| `RuntimeError: TorchMemorySaver is disabled` | colocate 时 `expandable_segments:True` 与 `torch_memory_saver` 冲突 | `run_slime_generator.sh` 的 `RUNTIME_ENV_JSON` 改为不强制传 `expandable_segments` | +| `IndexError: index 32 is out of range` (recompute) | PP=2 每 stage 32 层,`RECOMPUTE_NUM_LAYERS=64` 越界 | 改为 `RECOMPUTE_NUM_LAYERS=1`(对齐官方) | +| `KeyError: 'model.language_model.layers.0...'` (转换) | mbridge 转换缺 model spec | 转换命令加 `--spec slime_plugins.models.qwen3_5 get_qwen3_5_spec` | +| `ModuleNotFoundError: megatron.training` (转换) | PYTHONPATH 缺 Megatron | `PYTHONPATH=/root/Megatron-LM:$PYTHONPATH` | -## 4. Running +### EAGLE + raw 模式不兼容问题详解 -### Claude Code Agent baseline(Exp1) +**现象**:rollout 阶段正常,但训练的 log_probs 计算阶段 `mamba_pool.alloc` 报 `CUDA error: illegal memory access`。 -Claude Code Exp1 使用官方 `exp_agent/claudecode/dataset.jsonl` 和 -`templates/default.md`,包含漏洞知识和位置,不向 Agent 提供 PoC 或单元测试 -反馈。Agent 最多执行 100 次工具调用,并在任务容器内浏览和修改完整仓库。 +**根因**: +1. EAGLE 投机解码有独立的 draft model,也有自己的 mamba state +2. colocate 模式下,offload 时 `flush_cache` 清了主模型的 KV cache + mamba state +3. `wake_up` 重新加载 Megatron 模型到 GPU +4. log_probs 阶段重新调用 SGLang,但 EAGLE draft model 的 mamba state 在 flush 后**状态不一致** → 非法内存访问 -`run_eval.sh` 使用 Gateway 原生 Anthropic Messages/SSE 接口。任务容器中的 -Claude Code 直接请求 -`/v1/sessions//v1/messages`,不再启动 Claude Adapter,也不再经过 -Anthropic → OpenAI → Anthropic 转换。Gateway 透明转发原生流式事件,把实际发往 -Provider 的 JSON body 写入 `session_steps.request`,并把 Provider 返回的原始 -Anthropic SSE 文本写入 `session_steps.response`。SSE 聚合由导出脚本完成; -Gateway 不再改写请求或构造 Provider Artifact。 +**官方 vs 我们**: +- 官方用 `bridge` 模式 + EAGLE,正常工作 +- 我们用 `raw` 模式(因 GDN+CP 在 bridge 模式下不兼容)+ EAGLE,崩溃 +- **结论**:raw 模式下 EAGLE 不兼容,必须关闭 -先运行一个样本: +**代价**:推理速度变慢(无投机解码加速),但训练能跑通。待 SGLang 修复后可重新启用。 -```bash -export PATCH_EVAL_API_KEY="" -export DOCKER_HOST="tcp://:2376" -export PATCH_EVAL_MODEL="claude-opus-4-8" -export PATCH_EVAL_BASELINE="claudecode" -export PATCH_EVAL_AGENT_EXPERIMENT="exp1" -export PATCH_EVAL_TASK_LIMIT=1 -export PATCH_EVAL_POOL_SIZE=1 - -./rl/examples/patcheval/run_eval.sh -``` +### `max_tokens` vs `LLM_MAX_LENGTH` 说明 -`PATCH_EVAL_MODEL` 是底层模型路由,不是 Agent 名称。Claude Code baseline -要求显式设置它,避免误用 LLM baseline 默认的 DeepSeek 模型。 +| 参数 | 限制对象 | 值 | 代码位置 | +|------|----------|-----|----------| +| `LLM_MAX_LENGTH` | **整个轨迹**(多轮对话 input+output 累计) | 131072 | `llm_proxy.py` `state.max_length` | +| `ROLLOUT_MAX_RESPONSE_LEN` | **单轮生成**(每次 LLM 调用) | 32768 | `slime_generator.py` → SGLang `sampling_params.max_tokens` | -Claude Code 使用单一原生 Anthropic 路径。Gateway 保留客户端请求 body、除 -`beta` 外的 query string 和支持的 Anthropic headers,并以路由配置的上游凭据 -转发请求;不会改写 thinking、token budget 或 context management。为兼容 -Bedrock,可通过路由的 `anthropic_interleaved_thinking` 仅添加 -`interleaved-thinking-2025-05-14` Beta Header;Claude Code 的内部 Beta 标志 -不会被转发。 +`llm_proxy.py` 实际生效逻辑:`max_new_tokens = min(ROLLOUT_MAX_RESPONSE_LEN, LLM_MAX_LENGTH - 当前input_ids长度)` -首次启动每个 CVE 容器时会安装 Node.js 和 `@anthropic-ai/claude-code`,因此 -Agent baseline 的启动时间和网络开销明显高于 LLM baseline。确认单样本运行 -正常后,将 `PATCH_EVAL_TASK_LIMIT` 改为 `0` 再运行全量。 +**之前的问题**:两个值都是 131072,单轮就能跑满 128K,长尾轨迹 ~9 分钟不结束 → flush_cache 超时 → job 崩溃。 -### LLM baseline(S1.x) +**修复后**:单轮最多 32K(~2 分钟),多轮累计可达 131K,长尾可控。 -一键启动 Gateway、生成配置并运行标准 Launcher: +## 2026-09-05 更新日志 -```bash -export PATCH_EVAL_API_KEY="" -export DOCKER_HOST="tcp://:2376" -export PATCH_EVAL_MODEL="bailian/deepseek-v4-flash" -export PATCH_EVAL_BASELINE="llm" -export PATCH_EVAL_SETTING="s1.1" -export PATCH_EVAL_TASK_LIMIT=1 # smoke test;全量改为 0 - -./rl/examples/patcheval/run_eval.sh -``` +### 发现:补丁可能是 CUDA 崩溃的根因 -脚本只负责进程编排,最终执行的仍是标准 -`launcher.py --enable-evaluation`;没有额外的批量评测或结果回写阶段。每次运行 -使用带时间戳的新数据库,路径会在启动时打印。 +**现象**:`20260904-115405` run 中,rollout 成功收集 64 条样本(`RL_OFF_BY_N=1` 修复生效),但训练 wake_up 后 SGLang 引擎 `CUDA error: illegal memory access`,所有引擎崩溃,job 失败。 -与 OpenRT 相同的标准 SAfactory 启动方式(Gateway 需已单独启动): +**根因分析**: +- 官方 Qwen3.5-27B 脚本就是 `raw 模式 + EAGLE + colocate + CP=4`,**不需要任何补丁**就能跑通 +- 我们之前加了两个补丁,反而可能破坏了稳定性: + 1. `flush_cache_fix.py` — flush_cache 前发 `/abort_request` 强制终止所有 pending 请求,可能破坏 mamba state pool 的一致性 → resume_memory_occupation 时非法内存访问 + 2. `attention_mask_fix.py` — bridge 模式的补丁,raw 模式不需要,但仍在加载,可能干扰 raw 模式的 preprocess 流程 -```bash -GENERATED_DIR=/tmp/safactory-patcheval-s1.1 -python env/patcheval/generate_full_config.py \ - --output-dir "${GENERATED_DIR}" \ - --archive-dir /mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images \ - --official-runtime-dir /mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-runtime \ - --setting s1.1 \ - --evaluation-timeout-s 3600 \ - --shared-tmp /mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-tmp - -python launcher.py \ - --mode docker \ - --docker-pull-policy never \ - --docker-image-archive-dir /mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images \ - --cleanup-docker-image \ - --agent-config "${GENERATED_DIR}/patcheval_config.yaml" \ - --agent-start-config "${GENERATED_DIR}/patcheval_start.yaml" \ - --gateway-base-url http://127.0.0.1:8000/v1/sessions \ - --llm-model YOUR_ROUTE_KEY \ - --enable-evaluation \ - --db-path sqlite://env_trajs.db \ - --pool-size 5 \ - --max-workers 5 -``` +**官方 vs 我们对比**: + +| | 官方 Qwen3.5-27B | 我们 | +|---|---|---| +| megatron-to-hf-mode | raw(默认) | raw | +| EAGLE | ✅ 开 | ❌ 关了(误判) | +| colocate | ✅ | ✅ | +| CP/TP/PP | 4/4/2 | 4/4/2 | +| 补丁 | 无 | flush_cache_fix + attention_mask_fix | +| 数据 | dapo-math-17k(单轮) | PatchEval(多轮 agent) | + +**结论**:应该去掉补丁、开 EAGLE,完全对齐官方脚本。flush_cache 超时问题用其他方式解决(见下方)。 + +### 补丁处理 -`generate_full_config.py` 会将 `rule_evaluator.py` 放到生成配置目录,因此 -Launcher 按标准约定自动发现它,不需要 evaluation YAML。 +| 补丁 | 处置 | 原因 | +|---|---|---| +| `attention_mask_fix.py` | ❌ **删除** | bridge 模式补丁,raw 模式不需要 | +| `flush_cache_fix.py` | ✅ **保留(v5)** | flush_cache 用 SGLang 内置 `?timeout=60`;abort 移到 slime_generator rollout 阶段 | +| `gdn_packed_seq.py` | ✅ 保留 | GDN packed sequence 支持,官方也需要 | +| `traj_truncation.py` | ✅ 保留 | 轨迹截断(当前禁用,TRAJ_TRUNCATION_MAX_SEQ_LEN=0) | +| `raw_hf_checkpoint.py` | ✅ 保留 | raw 模式加载 checkpoint(已转 Megatron 格式,可能不需要) | +| `spread_placement.py` | ✅ 保留 | 多机 placement group SPREAD 策略 | -默认运行 Docker 支持的全部 230 个 CVE。镜像归档默认从以下目录按需加载: +### `flush_cache_fix.py` v5 — 用 SGLang 内置 timeout,abort 移到 rollout 阶段 -```text -/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images +**v1-v4 的问题**:在 flush_cache 阶段发 `/abort_request`,但 abort 在 scheduler stream 操作 mamba state,forward stream 可能还在写 → **race condition** → `CUDA error: illegal memory access`(SGLang issue #24221, #24954)。固定 sleep(2s/10s/60s)不能可靠避免这个 race。 + +**v5 的改进**(基于 SGLang 源码调研): +1. **abort 移到 rollout 阶段**(`slime_generator.py`):收齐数据后先 `/stop_rollout` kill envs,再向所有 SGLang worker 发 `/abort_request`(和 slime 原版 `abort()` 一致)。abort 在 rollout 阶段发生,forward stream 有充足时间在 offload/flush 前稳定。 +2. **flush_cache 用 SGLang 内置 `?timeout=60`**(SGLang PR #21413):scheduler 在 event loop 里轮询 `is_fully_idle()`,继续跑 forward stream 排空 in-flight batch,idle 后自动 flush。不再在 flush 阶段 abort,避免 scheduler/forward stream race。 +3. **不需要固定 sleep**:让 SGLang 自己的 idleness check 决定时机。 + +**三步流程**: +``` +1. /stop_rollout → kill envs ← 断新请求来源 +2. /abort_request → 清残余请求(rollout 阶段) ← 强制终止 SGLang 队列长生成 +3. /flush_cache?timeout=60 → scheduler 轮询 idle ← SGLang 自己确认 idle 后 flush ``` -每个任务启动前,Launcher 会在 Docker 中检查对应的 -`ghcr.io/anonymous2578-data/cve-*:latest`。若镜像不存在,则加载匹配的 -`cve-*-latest.tar`;任务容器结束后再删除本次加载的镜像。这样无需把约 -503 GB 的镜像同时放进 Docker 数据目录。 +### 配置变更 + +| 参数 | 之前 | 现在 | 原因 | +|------|------|------|------| +| SGLANG_SPECULATIVE_ALGORITHM | 空(关闭) | **EAGLE** | 对齐官方,raw 模式 + EAGLE 官方支持 | +| RL_OFF_BY_N | 0 | **1** | 修复首条 rollout 版本过滤问题(rollout_id 从1开始,current_version=rollout_id+1=2,但 weight_version=1) | +| LLM_MAX_LENGTH | 131072 | **65536** | 轨迹实际~40K,64K够用,降低长尾风险 | +| RL_GROUP_SIZE | 2 | **8** | 恢复大 batch 实验 | +| RL_ROLLOUT_GROUP_BATCH_SIZE | 1 | **8** | 恢复大 batch 实验 | +| RL_GLOBAL_BATCH_SIZE | 2 | **64** | = 8×8 | + +### 修改文件 + +1. `rl/patches/sitecustomize.py` — 移除 `attention_mask_fix`;`flush_cache_fix` 注释更新为 v5 +2. `rl/examples/patcheval/env.rjob.sh` — EAGLE 重新开启;LLM_MAX_LENGTH=65536;RL_OFF_BY_N=1;group=8/group_size=8/global_batch=64 +3. `rl/patches/flush_cache_fix.py` — v5:去掉 abort,改用 `/flush_cache?timeout=60` +4. `rl/slime_generator.py` — 收齐数据后先 `/stop_rollout`(kill envs),再向所有 SGLang worker 发 `/abort_request`(清残余),然后 sleep(3) +5. `rl/buffer_server.py` — 新增 `/stop_rollout` 端点(kill `aievobox_process` 子进程树,不杀进程组) + +### 解决的问题 + +| 问题 | 根因 | 修复 | +|------|------|------| +| 所有轨迹被版本过滤丢弃 | `current_version = rollout_id + 1 = 2`,但 `weight_version = 1`,`off_by_n=0` 导致全过滤 | `RL_OFF_BY_N=1`,允许 1 版本偏差 | +| CUDA illegal memory access (wake_up 后) | `flush_cache_fix.py` v1-v4 在 flush 阶段 abort,scheduler stream 清 mamba state 与 forward stream 写竞争(SGLang issue #24221,PR #24954 修复但 0.5.9 未包含) | v5:abort 移到 rollout 阶段,flush 用 `?timeout=60` 不 abort | +| EAGLE 误判为不兼容 | 之前在 raw 模式崩溃,误判为 EAGLE+raw 不兼容;实际是 flush_cache_fix 的 abort 破坏了 mamba state | v5 修复后重新开启 EAGLE | +| flush_cache 超时 | 长尾 env 持续发新请求,SGLang 永不 idle | `/stop_rollout` kill envs + abort 残余 + `?timeout=60` | +| step 2 连不上 buffer server | `/stop_rollout` 用 `os.killpg(os.getpgid(pid))` 杀整个进程组,buffer server 和 launcher.py 同组 → 一起被杀 | 改用 `ps --ppid` 逐层找子进程单独 kill,不杀进程组 | + +### step 1 验证结果(20260905-164831) + +step 1 完整跑通,v5 方案验证成功: + +| 阶段 | 耗时 | 状态 | +|------|------|------| +| rollout(64 samples) | 1932s (~32min) | ✅ | +| flush_cache | <1s | ✅ 无超时 | +| wake_up | 7.4s | ✅ 无 CUDA error | +| log_probs | 1249.9s (~21min) | ✅ | +| actor_train | 2370.6s (~39min) | ✅ loss=0.0(reward 全 0) | +| update_weights | 37.5s | ✅ | +| step 2 rollout | — | ❌ buffer server 被 `/stop_rollout` 误杀(已修复) | + +### SGLang 版本与 mamba race 说明 + +SGLang 0.5.9(当前训练机版本)**不包含** PR #24954(把 mamba state 操作移到 forward stream,消除 scheduler/forward stream race)。该 PR 合并于 2026-05-19,首个包含的正式版本是 **v0.5.13**(2026-06-13)。 + +v5 方案通过将 abort 提前到 rollout 阶段(flush 阶段不 abort),规避了这个 race。如未来升级 SGLang 到 v0.5.13+,abort 将天然安全,可考虑简化补丁。 + +## TODO -Smoke test 时给 `generate_full_config.py` 增加 `--limit 1`;全量评测省略 -`--limit` 或设置为 `--limit 0`。并发由 Launcher 的 `--pool-size` 和 -`--max-workers` 控制。 +- **验证**:`/stop_rollout` 修复后(不杀进程组),step 2+ 能否连续跑通多个 step +- 如 flush_cache 仍超时:可升级 SGLang 到 v0.5.13+(包含 PR #24954,abort 天然安全),或加 `SGLANG_DISABLE_OVERLAP_SCHEDULE=true` 彻底消除 race +- MAX_TOKENS_PER_GPU:当前 8192,可根据显存余量继续调优 +- reward 全 0 问题:考虑 partial-credit reward(参考 SWE-RL / SecureCodeRL)或混合简单任务 +- step 1 耗时分析:rollout 32min + log_probs 21min + actor_train 39min = ~93min/step,100 epoch 需要 ~155 小时(6.5 天) diff --git a/rl/examples/patcheval/env.rjob.qwen3_5_9b.sh b/rl/examples/patcheval/env.rjob.qwen3_5_9b.sh new file mode 100755 index 00000000..23f1879e --- /dev/null +++ b/rl/examples/patcheval/env.rjob.qwen3_5_9b.sh @@ -0,0 +1,324 @@ +#!/usr/bin/env bash +# +# ============================================================================= +# [RJOB MODE] PatchEval RL environment — Qwen3.5-9B variant +# ============================================================================= +# This is the RJOB mode variant for Qwen3.5-9B (dense, ~18GB BF16 weights). +# It mirrors env.rjob.sh (the Qwen3.8-27B variant) but scales down the +# parallelism and resource footprint to match the smaller model: +# - 1 node x 8 GPUs for training (vs 4 nodes x 8 GPUs for 27B) +# - TP=2 / PP=1 / CP=2 (vs TP=4 / PP=2 / CP=4 for 27B) +# - OPTIMIZER_CPU_OFFLOAD=false (9B weights fit in GPU memory easily) +# All other RJob infrastructure (gateway, pool size, EAGLE, mamba scheduler) +# is kept identical to the 27B variant so behavior is comparable. +# +# Prerequisites (vs docker env.sh): +# 1. AIEVOBOX_RJOB_CONFIG must point to a cluster config with valid +# access_key/secret_key (see ${REPO_ROOT}/config.yaml). The cybergym +# example reuses that file; fill in the credentials before running. +# 2. Agent configs are the rjob variants: +# AIEVOBOX_AGENT_CONFIG -> patcheval_config.rjob.yaml +# AIEVOBOX_AGENT_START_CONFIG -> patcheval_start.rjob.yaml +# Both live next to the docker-generated configs in PATCH_EVAL_GENERATED_DIR +# so they reuse the same datasets/ and per-env rule_evaluator.py. +# 3. The RL gateway (running on this training pod) must be reachable from the +# RJob pods. PATCHEVAL_GATEWAY_HOST defaults to this pod's IP; confirm it +# is routable from the cluster namespace (100.x pod IPs usually are). +# 4. RJob pods run DinD (privileged) to load CVE images from the gpfs-mounted +# archive dir, so privileged=true is set in patcheval_start.rjob.yaml. +# +# Usage: +# export PATCH_EVAL_GENERATED_DIR= +# RL_ENV_SH=$this rl/run_buffer_server.sh +# ============================================================================= +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# Self-contained: previously this sourced rl/examples/geo3k_vl/env.sh for +# infrastructure defaults, but that leaked geo3k-specific values (DAPO_filter=true, +# geo3k db path, qwen3-vl-2b model, RL_GLOBAL_BATCH_SIZE=512, ...) into patcheval. +# Only the infrastructure vars actually consumed by run_slime_generator.sh / +# buffer_server.py / llm_proxy.py / slime_generator.py are inlined at the bottom. + +: "${PATCH_EVAL_GENERATED_DIR:?Set PATCH_EVAL_GENERATED_DIR to a generated PatchEval config directory}" + +export AIEVOBOX_EXAMPLE_NAME="patcheval_qwen3_5_9b" +export AIEVOBOX_ROOT="${AIEVOBOX_ROOT:-${REPO_ROOT}}" +export AIEVOBOX_MODE="rjob" +export AIEVOBOX_RJOB_CONFIG="${AIEVOBOX_RJOB_CONFIG:-${REPO_ROOT}/config.yaml}" +export PATCH_EVAL_BASELINE="openhands" +export PATCH_EVAL_AGENT_EXPERIMENT="exp1" +export STORAGE_TYPE="${STORAGE_TYPE:-sqlite}" +export AIEVOBOX_DB_URL="${PATCHEVAL_DB_URL:-sqlite:///${AIEVOBOX_ROOT}/rl/examples/patcheval/patcheval_qwen3_5_9b.db}" +export DOCKER_HOST="${DOCKER_HOST:-tcp://100.99.17.62:2376}" +export AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR="${PATCH_EVAL_IMAGE_ARCHIVE_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images}" +export AIEVOBOX_DOCKER_PULL_POLICY="${AIEVOBOX_DOCKER_PULL_POLICY:-never}" +# RJob variants of the agent configs (live alongside the docker-generated ones +# so datasets/ and rule_evaluator.py are reused). +export AIEVOBOX_AGENT_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_config.rjob.yaml" +export AIEVOBOX_AGENT_START_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_start.rjob.yaml" +export AIEVOBOX_AGENT_ROOT="${PATCH_EVAL_GENERATED_DIR}" +# Per-task rollout rounds (NOT LLM steps per episode). Each CVE task is rolled +# out this many times per rollout step. +export AIEVOBOX_MAX_STEPS="${PATCHEVAL_MAX_STEPS:-60}" +export AIEVOBOX_ENABLE_EVALUATION="${AIEVOBOX_ENABLE_EVALUATION:-1}" +# RJob can scale across the cluster; default higher than docker's 1. +# Raised from 4 to 16 to fix the rollout throughput bottleneck (SGLang was +# only seeing #running-req: 1, ~56 tok/s, because only a few agent episodes +# were in flight). 16 concurrent episodes gives the LLM proxy enough +# in-flight requests to keep SGLang's decode batches full. The derived +# concurrency vars (AIEVOBOX_LLM_MAX_CONCURRENCY, AIEVOBOX_LLM_PROXY_WORKERS, +# AIEVOBOX_TRAININFO_WORKERS) auto-track this via ${AIEVOBOX_POOL_SIZE}. +# Override via PATCHEVAL_POOL_SIZE if cluster capacity is tight. +# 8 concurrent episodes on a single 8-GPU machine (colocate mode: training + +# rollout share the same 8 GPUs). Override via PATCHEVAL_POOL_SIZE if needed. +export AIEVOBOX_POOL_SIZE="${PATCHEVAL_POOL_SIZE:-8}" +export AIEVOBOX_AGENT_START_TIMEOUT_S="${PATCHEVAL_AGENT_START_TIMEOUT_S:-1200}" +# Hard cap on LLM steps per episode, enforced by the RL gateway +# (see rl/gateway_autostart.py). -1 = unlimited. Set >=0 to stop runaway +# agent rollouts (e.g. OpenHands looping 200+ steps without finishing). +# 12 was too few for CVE-fix tasks (binary reward → 0 solve → 0 RL signal). +# 40 gives the model a real shot at explore+edit+test while keeping +# throughput workable (~1.5hr/train step, ~6 days/100 epoch). max_tokens stays +# at 6144 (gateway default) — not lowered, per user choice. +export AIEVOBOX_GATEWAY_MAX_STEPS="${PATCHEVAL_GATEWAY_MAX_STEPS:-60}" + +# NOTE: geo3k_vl/env.sh (sourced above) already sets these to its own defaults +# (e.g. RL_GLOBAL_BATCH_SIZE=512, RL_ROLLOUT_GROUP_BATCH_SIZE=64). Using +# ${VAR:-default} here would keep geo3k's values, so we override +# unconditionally. Override via PATCHEVAL_* if needed. +# group_size=2 + rollout_batch=2 + num_rollout=10 -> 40 episodes total. +# global_batch_size must equal rollout_batch_size * group_size = 2*2 = 4 so +# each rollout's samples exactly fill one global batch (1 train step/rollout, +# no waste). This is the 27B invariant (gbs=64=8*8). A larger gbs (e.g. 8) +# breaks it: each rollout yields 4 samples < 8 -> rollout.py:608 raises +# "Not enough samples 4 for global_batch_size 8". Also keeps train_iters +# = 40//4 = 10 > 0 so Megatron OptimizerParamScheduler's lr_decay_steps>0 +# assert passes. +export RL_GROUP_SIZE="${PATCHEVAL_GROUP_SIZE:-2}" +export RL_GLOBAL_BATCH_SIZE="${PATCHEVAL_GLOBAL_BATCH_SIZE:-4}" +export RL_ROLLOUT_GROUP_BATCH_SIZE="${PATCHEVAL_ROLLOUT_GROUP_BATCH_SIZE:-2}" +export SLIME_ROLLOUT_BATCH_SIZE="${PATCHEVAL_SLIME_ROLLOUT_BATCH_SIZE:-${RL_ROLLOUT_GROUP_BATCH_SIZE}}" +export SLIME_GLOBAL_BATCH_SIZE="${PATCHEVAL_SLIME_GLOBAL_BATCH_SIZE:-${RL_GLOBAL_BATCH_SIZE}}" +# RL_EPOCH=100 training rounds. Override via PATCHEVAL_EPOCH / NUM_ROLLOUT +# env vars if needed. +export RL_EPOCH="${PATCHEVAL_EPOCH:-100}" +export RL_MODEL="${RL_MODEL:-model}" +export RL_API_KEY="${RL_API_KEY:-openai_api_key}" + +export BUFFER_SERVER_HOST="${BUFFER_SERVER_HOST:-127.0.0.1}" +export BUFFER_SERVER_PORT="${BUFFER_SERVER_PORT:-18889}" +export LLM_PROXY_HOST="${LLM_PROXY_HOST:-127.0.0.1}" +export LLM_PROXY_PORT="${LLM_PROXY_PORT:-18890}" +export LLM_MAX_LENGTH="${LLM_MAX_LENGTH:-65536}" +# Single-turn generation limit (per LLM call). Keep separate from LLM_MAX_LENGTH +# (trajectory cap) to avoid long-tail requests blocking offload. Official +# Qwen3.5-27B uses 32768. +export ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-32768}" +export LLM_TEMPERATURE="${LLM_TEMPERATURE:-1.0}" +# Gateway runs on THIS training pod (started by the buffer server via +# gateway_autostart). Default to this pod's IP so it always points at the live +# gateway, not a stale hardcoded IP. Override via PATCHEVAL_GATEWAY_HOST. +export AIEVOBOX_GATEWAY_HOST="${PATCHEVAL_GATEWAY_HOST:-$(hostname -i | awk '{print $1}')}" +export AIEVOBOX_GATEWAY_PORT="${PATCHEVAL_GATEWAY_PORT:-8000}" +export AIEVOBOX_GATEWAY_BASE_URL="http://${AIEVOBOX_GATEWAY_HOST}:${AIEVOBOX_GATEWAY_PORT}/v1/sessions" + +export SLIME_HOME="${SLIME_HOME:-/root/slime}" +export MEGATRON_HOME="${MEGATRON_HOME:-/root/Megatron-LM}" +# Model: Qwen3.5-9B (dense, qwen3_5 architecture, uses qwen3.5-9B.sh spec). +# Override via QWEN3_5_9B_CKPT_DIR / PATCHEVAL_* if needed. +export HF_CKPT_DIR="${QWEN3_5_9B_CKPT_DIR:-/mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/hub/models--Qwen--Qwen3.5-9B/snapshots/c202236235762e1c871ad0ccb60c8ee5ba337b9a}" +# --load must point to Megatron format checkpoint (not HF), matching official script. +# Convert HF->Megatron first with slime/tools/convert_hf_to_torch_dist.py + +# scripts/models/qwen3.5-9B.sh, saving to the path below. +export LOAD_DIR="${QWEN3_5_9B_LOAD_DIR:-/mnt/shared-storage-user/evobox-share-gpfs2/leishanzhe/model/Qwen3.5-9B_megatron}" +# Override geo3k defaults unconditionally (geo3k sets these to its own paths). +export SAVE_DIR="${PATCHEVAL_SAVE_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/checkpoints/Qwen3.5-9B_megatron}" +export WANDB_DIR="${PATCHEVAL_WANDB_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/wandb_logs}" +export LOG_ROOT="${PATCHEVAL_LOG_ROOT:-${AIEVOBOX_ROOT}/logs/patcheval_qwen3_5_9b}" +export MODEL_SCRIPT="${QWEN3_5_9B_MODEL_SCRIPT:-${REPO_ROOT}/../slime/scripts/models/qwen3.5-9B.sh}" +export MODEL_ARGS_ROTARY_BASE=10000000 + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export NUM_GPUS="${PATCHEVAL_NUM_GPUS:-8}" # total on machine; split 4 train + 4 rollout + +# Multi-node NCCL: disable InfiniBand (ibv_modify_qp fails across some node +# pairs) and force TCP socket transport. Also set the network interface to +# ensure NCCL binds to the correct IP. Override via NCCL_* env vars if needed. +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" +export NCCL_NET="${NCCL_NET:-Socket}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-bond0}" +# Single 8-GPU machine, non-colocate: 4 GPUs for training + 4 for rollout. +export ACTOR_NUM_NODES="${PATCHEVAL_ACTOR_NUM_NODES:-1}" +export ACTOR_NUM_GPUS_PER_NODE="${PATCHEVAL_ACTOR_NUM_GPUS_PER_NODE:-4}" +# Inference GPUs for sglang (dedicated, separate from the 4 training GPUs). +# 4 engines x 1 GPU each (9B fits on a single H200 for inference). +export ROLLOUT_NUM_GPUS="${PATCHEVAL_ROLLOUT_NUM_GPUS:-4}" +export ROLLOUT_NUM_GPUS_PER_ENGINE="${PATCHEVAL_ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-${SLIME_HOME}/train.py}" +export ROLLOUT_FUNCTION_PATH="${ROLLOUT_FUNCTION_PATH:-rl.slime_generator.generate_rollout}" +# Debug-friendly: 300 rollout iterations is too many for a debug run. +export NUM_ROLLOUT="${NUM_ROLLOUT:-10}" +export LOSS_MASK_TYPE="qwen3_5" +export TRAIN_BACKEND="${TRAIN_BACKEND:-megatron}" +export MEGATRON_TO_HF_MODE="${MEGATRON_TO_HF_MODE:-raw}" +# 9B dense on 4 training GPUs: TP=2 / PP=1 / CP=1 gives DP=2 (4/2=2). +# vs 27B's TP=4 / PP=2 / CP=4. PP=1 means no DECODER_LAST_PIPELINE_NUM_LAYERS. +export TP_SIZE="${PATCHEVAL_TP_SIZE:-2}" PP_SIZE="${PATCHEVAL_PP_SIZE:-1}" CP_SIZE="${PATCHEVAL_CP_SIZE:-1}" EP_SIZE=1 ETP_SIZE=1 +export RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}" +export RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}" +export RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}" +# PP=1: no last-pipeline layer override needed (27B used PP=2 -> 30). +export ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}" +export MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" +# PP=2 halves both weights and activations per GPU (32 of 64 layers each). +# This provides enough memory headroom to disable truncation entirely — +# full long trajectories (50k+ tokens) can be trained without losing context. +export TRAJ_TRUNCATION_MAX_SEQ_LEN="${TRAJ_TRUNCATION_MAX_SEQ_LEN:-0}" +# GDN packed-seq monkey-patch: patches Megatron GDN forward to pass cu_seqlens +# to chunk_gated_delta_rule, enabling thd (packing) mode without NotImplementedError. +# See rl/patches/gdn_packed_seq.py for details. +export PYTHONPATH="${REPO_ROOT}/rl/patches${PYTHONPATH:+:${PYTHONPATH}}" +export USE_DYNAMIC_BATCH_SIZE="${USE_DYNAMIC_BATCH_SIZE:-true}" +export CALCULATE_PER_TOKEN_LOSS="${CALCULATE_PER_TOKEN_LOSS:-true}" +export ADVANTAGE_ESTIMATOR="${ADVANTAGE_ESTIMATOR:-grpo}" +# DAPO group filter: drop groups where all samples share the same reward +# (zero advantage, no learning signal). For patcheval, base model rarely +# solves CVEs, so most groups are all-0 and get filtered → buffer never fills +# → pipeline stalls. Default off; flip with PATCHEVAL_DAPO_FILTER=true. +export DAPO_filter="${PATCHEVAL_DAPO_FILTER:-false}" +export LR="${LR:-1e-6}" +export OPTIMIZER="${OPTIMIZER:-adam}" +export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +export ADAM_BETA1="${ADAM_BETA1:-0.9}" +export ADAM_BETA2="${ADAM_BETA2:-0.98}" +# Non-colocate mode: training (Megatron) and inference (SGLang) use SEPARATE +# GPUs on this single 8-GPU machine. Split: 4 training + 4 rollout. +export SLIME_COLOCATE="${SLIME_COLOCATE:-false}" +# CPU offload optimizer: 9B weights (~18GB BF16) + optimizer fit comfortably on +# GPU, so CPU offload is OFF by default (27B needed it to avoid OOM on 140GB +# GPUs). Flip with OPTIMIZER_CPU_OFFLOAD=true if memory is tight. +export OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" +export USE_WANDB="${USE_WANDB:-true}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_PROJECT="${WANDB_PROJECT:-slime}" +export WANDB_GROUP="${WANDB_GROUP:-patcheval_qwen3_5_9b}" +# KV cache pool as a fraction of GPU memory (after weights). This is the +# primary throughput lever for multi-turn agent rollouts: each concurrently +# decoding request must keep its prompt+history KV resident, so KV capacity +# directly caps parallel decode (#running-req). At 0.45 the KV hits ~97% with +# 4 sessions/engine, forcing eviction (cached=0 thrashing) and queueing +# (waiting for decode to free KV, median ~26s/request). 0.6 gives +33% KV +# capacity, enough to hold 4 sessions/engine without eviction so all 4 decode +# in parallel — no queue, ~100% prefix reuse. Safe on H200 141GB: 27B weights +# ~54GB + KV 0.6*~87GB-free ≈ 52GB ≈ 106GB < 141GB. Override via env var. +# Non-colocate mode: SGLang has dedicated rollout GPUs (no sharing with +# training), so can use a high fraction like the 27B variant. +export SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.9}" +export SGLANG_ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-fa3}" +export SGLANG_MAMBA_SCHEDULER_STRATEGY="${SGLANG_MAMBA_SCHEDULER_STRATEGY:-extra_buffer}" +# EAGLE speculative decoding (official Qwen3.5-27B config). +# NOTE: Previously disabled due to mamba_pool CUDA illegal memory access, but that +# was actually caused by flush_cache_fix.py's abort_request corrupting mamba state. +# With flush_cache_fix removed, EAGLE should work in raw mode (as in official script). +export SGLANG_SPECULATIVE_ALGORITHM="${SGLANG_SPECULATIVE_ALGORITHM:-EAGLE}" +export SGLANG_SPECULATIVE_NUM_STEPS="${SGLANG_SPECULATIVE_NUM_STEPS:-3}" +export SGLANG_SPECULATIVE_EAGLE_TOPK="${SGLANG_SPECULATIVE_EAGLE_TOPK:-1}" +export SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS="${SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS:-4}" +export SGLANG_LOG_LEVEL="${SGLANG_LOG_LEVEL:-info}" +export SGLANG_LOG_LEVEL_HTTP="${SGLANG_LOG_LEVEL_HTTP:-error}" +export CLEANUP_BEFORE_RUN="${CLEANUP_BEFORE_RUN:-true}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export PYTHONUNBUFFERED=1 + +# ============================================================================= +# Infrastructure defaults — inlined from rl/examples/geo3k_vl/env.sh. +# Only vars actually consumed by run_slime_generator.sh / buffer_server.py / +# llm_proxy.py / slime_generator.py. geo3k-specific paths (geo3k db, geo3k +# config, qwen3-vl-2b model) are intentionally NOT carried over; patcheval +# overrides above already set those. Placed at the end so ${VAR:-default} can +# reference patcheval values set earlier (POOL_SIZE, RL_GROUP_SIZE, ports). +# ============================================================================= + +# --- Ray / Python launcher --- +export PYTHON_BIN="${PYTHON_BIN:-python3}" +export RAY_BIN="${RAY_BIN:-ray}" +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +export RAY_ADDRESS="${RAY_ADDRESS:-http://127.0.0.1:8265}" +export RAY_PORT="${RAY_PORT:-}" +export KILL_PYTHON_BEFORE_RUN="${KILL_PYTHON_BEFORE_RUN:-false}" + +# --- Slime train / checkpoint args --- +export SAVE_INTERVAL="${SAVE_INTERVAL:-20}" +export MODEL_ARGS_EXTRA="${MODEL_ARGS_EXTRA:-}" +export REF_LOAD_DIR="${REF_LOAD_DIR:-/mnt/shared-storage-user/evobox-share-gpfs2/leishanzhe/model/Qwen3.5-9B_megatron}" +export CUSTOM_REWARD_POST_PROCESS_PATH="${CUSTOM_REWARD_POST_PROCESS_PATH:-}" +export SGLANG_LOGGING_CONFIG_PATH="${SGLANG_LOGGING_CONFIG_PATH:-}" + +# --- Optimizer / GRPO extras --- +export LR_DECAY_STYLE="${LR_DECAY_STYLE:-constant}" +export ENTROPY_COEF="${ENTROPY_COEF:-0.00}" +export EPS_CLIP="${EPS_CLIP:-0.2}" +export EPS_CLIP_HIGH="${EPS_CLIP_HIGH:-0.2}" +export USE_DYNAMIC_GLOBAL_BATCH_SIZE="${USE_DYNAMIC_GLOBAL_BATCH_SIZE:-false}" + +# --- W&B extras --- +export WANDB_TEAM="${WANDB_TEAM:-}" +export WANDB_ALWAYS_USE_TRAIN_STEP="${WANDB_ALWAYS_USE_TRAIN_STEP:-false}" + +# --- SGLang extras --- +# Tuned for high-concurrency rollout (POOL_SIZE=16). These three together let +# SGLang actually batch dozens of decode requests instead of running 1 at a +# time (the #running-req:1 symptom seen before): +# - MAX_RUNNING_REQUESTS: hard cap on the running decode batch. 64 gives +# headroom over POOL_SIZE=16 (multi-turn episodes overlap). +# - CUDA_GRAPH_BS: capture graphs for the batch sizes we expect to hit, so +# decode doesn't fall back to the slow non-graph path on shape change +# (fixes the first-step drop to 6.5 tok/s seen in the logs). +# - CHUNKED_PREFILL_SIZE: split long prefills into 8192-token chunks so a +# single long prompt can't starve the running decode batch. +export SGLANG_CUDA_GRAPH_BS="${SGLANG_CUDA_GRAPH_BS:-1 2 4 8 16 32}" +export SGLANG_MAX_RUNNING_REQUESTS="${SGLANG_MAX_RUNNING_REQUESTS:-64}" +export SGLANG_SCHEDULE_CONSERVATIVENESS="${SGLANG_SCHEDULE_CONSERVATIVENESS:-}" +export SGLANG_CHUNKED_PREFILL_SIZE="${SGLANG_CHUNKED_PREFILL_SIZE:-8192}" +export SGLANG_ENABLE_MIXED_CHUNK="${SGLANG_ENABLE_MIXED_CHUNK:-false}" + +# --- LLM proxy / buffer server workers & perf --- +export LLM_TOP_P="${LLM_TOP_P:-1.0}" +export LLM_PROXY_ENABLE_CONSOLE_LOG="${LLM_PROXY_ENABLE_CONSOLE_LOG:-0}" +export AIEVOBOX_LLM_MAX_CONCURRENCY="${AIEVOBOX_LLM_MAX_CONCURRENCY:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_PROXY_WORKERS="${AIEVOBOX_LLM_PROXY_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_STARTUP_JITTER_S="${AIEVOBOX_LLM_STARTUP_JITTER_S:-0}" +export AIEVOBOX_TRAININFO_WORKERS="${AIEVOBOX_TRAININFO_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE="${AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE:-256}" +export AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S="${AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S:-0.01}" +export AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS="${AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS:-1800}" +export ROLLBUF_HOST="${ROLLBUF_HOST:-${BUFFER_SERVER_HOST}}" +export ROLLBUF_PORT="${ROLLBUF_PORT:-${BUFFER_SERVER_PORT}}" + +# --- Slime rollout-buffer / GRPO filter --- +export SLIME_ROLLBUF_RESTART_TRAINING="${SLIME_ROLLBUF_RESTART_TRAINING:-True}" +export SLIME_N_SAMPLES_PER_PROMPT="${SLIME_N_SAMPLES_PER_PROMPT:-${RL_GROUP_SIZE}}" +export RL_OFF_BY_N="${RL_OFF_BY_N:-3}" +# Oversample: launch extra envs per prompt beyond group_size so the first +# group_size episodes to finish form a group; long-tail episodes don't block. +# buffer_server still pops group_size at a time; surplus stays in the bucket. +# E.g. group_size=8 + oversample=4 = 12 envs per prompt, first 8 done = 1 group. +export RL_OVERSAMPLE="${PATCHEVAL_OVERSAMPLE:-4}" + +# --- AIEVOBOX env extras --- +export AIEVOBOX_MESSAGE_CUT="${AIEVOBOX_MESSAGE_CUT:-0}" +export AIEVOBOC_MULTIPLIER="${AIEVOBOC_MULTIPLIER:-1.2}" + +# --- Runtime --- +# NOTE: expandable_segments:True is incompatible with torch_memory_saver +# (used in colocate mode). Disable it when colocate is on. +if [[ "${SLIME_COLOCATE:-false}" == "true" || "${SLIME_COLOCATE:-false}" == "1" ]]; then + export PYTORCH_CUDA_ALLOC_CONF="" + export PYTORCH_ALLOC_CONF="" +else + # PyTorch >= 2.5 renamed PYTORCH_CUDA_ALLOC_CONF → PYTORCH_ALLOC_CONF. + # Set both so old and new versions pick up expandable_segments. + export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +fi diff --git a/rl/examples/patcheval/env.rjob.sh b/rl/examples/patcheval/env.rjob.sh new file mode 100755 index 00000000..232f00cb --- /dev/null +++ b/rl/examples/patcheval/env.rjob.sh @@ -0,0 +1,313 @@ +#!/usr/bin/env bash +# +# ============================================================================= +# [RJOB MODE] PatchEval RL environment +# ============================================================================= +# This is the RJOB mode variant: AIEVOBOX_MODE=rjob, each rollout episode is +# submitted as a cluster job (RJob) via h.pjlab.org.cn, so many episodes can +# run in parallel across the cluster (not limited to a single Docker host). +# +# Prerequisites (vs docker env.sh): +# 1. AIEVOBOX_RJOB_CONFIG must point to a cluster config with valid +# access_key/secret_key (see ${REPO_ROOT}/config.yaml). The cybergym +# example reuses that file; fill in the credentials before running. +# 2. Agent configs are the rjob variants: +# AIEVOBOX_AGENT_CONFIG -> patcheval_config.rjob.yaml +# AIEVOBOX_AGENT_START_CONFIG -> patcheval_start.rjob.yaml +# Both live next to the docker-generated configs in PATCH_EVAL_GENERATED_DIR +# so they reuse the same datasets/ and per-env rule_evaluator.py. +# 3. The RL gateway (running on this training pod) must be reachable from the +# RJob pods. PATCHEVAL_GATEWAY_HOST defaults to this pod's IP; confirm it +# is routable from the cluster namespace (100.x pod IPs usually are). +# 4. RJob pods run DinD (privileged) to load CVE images from the gpfs-mounted +# archive dir, so privileged=true is set in patcheval_start.rjob.yaml. +# +# Usage: +# export PATCH_EVAL_GENERATED_DIR= +# RL_ENV_SH=$this rl/run_buffer_server.sh +# ============================================================================= +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# Self-contained: previously this sourced rl/examples/geo3k_vl/env.sh for +# infrastructure defaults, but that leaked geo3k-specific values (DAPO_filter=true, +# geo3k db path, qwen3-vl-2b model, RL_GLOBAL_BATCH_SIZE=512, ...) into patcheval. +# Only the infrastructure vars actually consumed by run_slime_generator.sh / +# buffer_server.py / llm_proxy.py / slime_generator.py are inlined at the bottom. + +: "${PATCH_EVAL_GENERATED_DIR:?Set PATCH_EVAL_GENERATED_DIR to a generated PatchEval config directory}" + +export AIEVOBOX_EXAMPLE_NAME="patcheval_qwen3_8_27b" +export AIEVOBOX_ROOT="${AIEVOBOX_ROOT:-${REPO_ROOT}}" +export AIEVOBOX_MODE="rjob" +export AIEVOBOX_RJOB_CONFIG="${AIEVOBOX_RJOB_CONFIG:-${REPO_ROOT}/config.yaml}" +export PATCH_EVAL_BASELINE="openhands" +export PATCH_EVAL_AGENT_EXPERIMENT="exp1" +export STORAGE_TYPE="${STORAGE_TYPE:-sqlite}" +export AIEVOBOX_DB_URL="${PATCHEVAL_DB_URL:-sqlite:///${AIEVOBOX_ROOT}/rl/examples/patcheval/patcheval_qwen3_8_27b.db}" +export DOCKER_HOST="${DOCKER_HOST:-tcp://100.99.17.62:2376}" +export AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR="${PATCH_EVAL_IMAGE_ARCHIVE_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images}" +export AIEVOBOX_DOCKER_PULL_POLICY="${AIEVOBOX_DOCKER_PULL_POLICY:-never}" +# RJob variants of the agent configs (live alongside the docker-generated ones +# so datasets/ and rule_evaluator.py are reused). +export AIEVOBOX_AGENT_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_config.rjob.yaml" +export AIEVOBOX_AGENT_START_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_start.rjob.yaml" +export AIEVOBOX_AGENT_ROOT="${PATCH_EVAL_GENERATED_DIR}" +# Per-task rollout rounds (NOT LLM steps per episode). Each CVE task is rolled +# out this many times per rollout step. +export AIEVOBOX_MAX_STEPS="${PATCHEVAL_MAX_STEPS:-1}" +export AIEVOBOX_ENABLE_EVALUATION="${AIEVOBOX_ENABLE_EVALUATION:-1}" +# RJob can scale across the cluster; default higher than docker's 1. +# Raised from 4 to 16 to fix the rollout throughput bottleneck (SGLang was +# only seeing #running-req: 1, ~56 tok/s, because only a few agent episodes +# were in flight). 16 concurrent episodes gives the LLM proxy enough +# in-flight requests to keep SGLang's decode batches full. The derived +# concurrency vars (AIEVOBOX_LLM_MAX_CONCURRENCY, AIEVOBOX_LLM_PROXY_WORKERS, +# AIEVOBOX_TRAININFO_WORKERS) auto-track this via ${AIEVOBOX_POOL_SIZE}. +# Override via PATCHEVAL_POOL_SIZE if cluster capacity is tight. +# NOTE: 实验扫 POOL_SIZE=8/16/24/32 找效率甜点。当前测试值=16。 +export AIEVOBOX_POOL_SIZE="${PATCHEVAL_POOL_SIZE:-16}" +export AIEVOBOX_AGENT_START_TIMEOUT_S="${PATCHEVAL_AGENT_START_TIMEOUT_S:-1200}" +# Hard cap on LLM steps per episode, enforced by the RL gateway +# (see rl/gateway_autostart.py). -1 = unlimited. Set >=0 to stop runaway +# agent rollouts (e.g. OpenHands looping 200+ steps without finishing). +# 12 was too few for CVE-fix tasks (binary reward → 0 solve → 0 RL signal). +# 40 gives the model a real shot at explore+edit+test while keeping +# throughput workable (~1.5hr/train step, ~6 days/100 epoch). max_tokens stays +# at 6144 (gateway default) — not lowered, per user choice. +export AIEVOBOX_GATEWAY_MAX_STEPS="${PATCHEVAL_GATEWAY_MAX_STEPS:-40}" + +# NOTE: geo3k_vl/env.sh (sourced above) already sets these to its own defaults +# (e.g. RL_GLOBAL_BATCH_SIZE=512, RL_ROLLOUT_GROUP_BATCH_SIZE=64). Using +# ${VAR:-default} here would keep geo3k's values, so we override +# unconditionally. Override via PATCHEVAL_* if needed. +export RL_GROUP_SIZE="${PATCHEVAL_GROUP_SIZE:-8}" +export RL_GLOBAL_BATCH_SIZE="${PATCHEVAL_GLOBAL_BATCH_SIZE:-64}" +export RL_ROLLOUT_GROUP_BATCH_SIZE="${PATCHEVAL_ROLLOUT_GROUP_BATCH_SIZE:-8}" +export SLIME_ROLLOUT_BATCH_SIZE="${PATCHEVAL_SLIME_ROLLOUT_BATCH_SIZE:-${RL_ROLLOUT_GROUP_BATCH_SIZE}}" +export SLIME_GLOBAL_BATCH_SIZE="${PATCHEVAL_SLIME_GLOBAL_BATCH_SIZE:-${RL_GLOBAL_BATCH_SIZE}}" +# RL_EPOCH=100 training rounds. Override via PATCHEVAL_EPOCH / NUM_ROLLOUT +# env vars if needed. +export RL_EPOCH="${PATCHEVAL_EPOCH:-100}" +export RL_MODEL="${RL_MODEL:-model}" +export RL_API_KEY="${RL_API_KEY:-openai_api_key}" + +export BUFFER_SERVER_HOST="${BUFFER_SERVER_HOST:-127.0.0.1}" +export BUFFER_SERVER_PORT="${BUFFER_SERVER_PORT:-18889}" +export LLM_PROXY_HOST="${LLM_PROXY_HOST:-127.0.0.1}" +export LLM_PROXY_PORT="${LLM_PROXY_PORT:-18890}" +export LLM_MAX_LENGTH="${LLM_MAX_LENGTH:-65536}" +# Single-turn generation limit (per LLM call). Keep separate from LLM_MAX_LENGTH +# (trajectory cap) to avoid long-tail requests blocking offload. Official +# Qwen3.5-27B uses 32768. +export ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-32768}" +export LLM_TEMPERATURE="${LLM_TEMPERATURE:-1.0}" +# Gateway runs on THIS training pod (started by the buffer server via +# gateway_autostart). Default to this pod's IP so it always points at the live +# gateway, not a stale hardcoded IP. Override via PATCHEVAL_GATEWAY_HOST. +export AIEVOBOX_GATEWAY_HOST="${PATCHEVAL_GATEWAY_HOST:-$(hostname -i | awk '{print $1}')}" +export AIEVOBOX_GATEWAY_PORT="${PATCHEVAL_GATEWAY_PORT:-8000}" +export AIEVOBOX_GATEWAY_BASE_URL="http://${AIEVOBOX_GATEWAY_HOST}:${AIEVOBOX_GATEWAY_PORT}/v1/sessions" + +export SLIME_HOME="${SLIME_HOME:-/root/slime}" +export MEGATRON_HOME="${MEGATRON_HOME:-/root/Megatron-LM}" +# Model: Qwen3.8-27B (same architecture as Qwen3.5-27B, uses qwen3.5-27B.sh spec). +# Override via QWEN3_8_27B_CKPT_DIR / PATCHEVAL_* if needed. +export HF_CKPT_DIR="${QWEN3_8_27B_CKPT_DIR:-/mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/hub/models--Qwen--Qwen3.8-27B/snapshots/1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0}" +# --load must point to Megatron format checkpoint (not HF), matching official script. +# We converted HF→Megatron to qwen3_8_27b_megatron/. +export LOAD_DIR="${QWEN3_8_27B_LOAD_DIR:-/mnt/shared-storage-user/evobox-share-gpfs2/leishanzhe/model/qwen3_8_27b_megatron}" +# Override geo3k defaults unconditionally (geo3k sets these to its own paths). +export SAVE_DIR="${PATCHEVAL_SAVE_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/checkpoints/Qwen3.8-27B_megatron}" +export WANDB_DIR="${PATCHEVAL_WANDB_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/wandb_logs}" +export LOG_ROOT="${PATCHEVAL_LOG_ROOT:-${AIEVOBOX_ROOT}/logs/patcheval_qwen3_8_27b}" +export MODEL_SCRIPT="${QWEN3_8_27B_MODEL_SCRIPT:-/root/slime/scripts/models/qwen3.5-27B.sh}" +export MODEL_ARGS_ROTARY_BASE=10000000 + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export NUM_GPUS="${PATCHEVAL_NUM_GPUS:-8}" + +# Multi-node NCCL: disable InfiniBand (ibv_modify_qp fails across some node +# pairs) and force TCP socket transport. Also set the network interface to +# ensure NCCL binds to the correct IP. Override via NCCL_* env vars if needed. +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" +export NCCL_NET="${NCCL_NET:-Socket}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-bond0}" +export ACTOR_NUM_NODES="${PATCHEVAL_ACTOR_NUM_NODES:-4}" +export ACTOR_NUM_GPUS_PER_NODE="${PATCHEVAL_ACTOR_NUM_GPUS_PER_NODE:-8}" +# Inference GPUs for sglang. In non-colocate mode, this is the dedicated +# rollout GPU count (separate from training GPUs). +# Total: 32 training + 8 rollout = 40 GPUs (5 machines). +export ROLLOUT_NUM_GPUS="${PATCHEVAL_ROLLOUT_NUM_GPUS:-8}" +export ROLLOUT_NUM_GPUS_PER_ENGINE="${PATCHEVAL_ROLLOUT_NUM_GPUS_PER_ENGINE:-2}" +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-${SLIME_HOME}/train.py}" +export ROLLOUT_FUNCTION_PATH="${ROLLOUT_FUNCTION_PATH:-rl.slime_generator.generate_rollout}" +# Debug-friendly: 300 rollout iterations is too many for a debug run. +export NUM_ROLLOUT="${NUM_ROLLOUT:-10}" +export LOSS_MASK_TYPE="qwen3_5" +export TRAIN_BACKEND="${TRAIN_BACKEND:-megatron}" +export MEGATRON_TO_HF_MODE="${MEGATRON_TO_HF_MODE:-raw}" +export TP_SIZE="${PATCHEVAL_TP_SIZE:-4}" PP_SIZE="${PATCHEVAL_PP_SIZE:-2}" CP_SIZE=4 EP_SIZE=1 ETP_SIZE=1 +export RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}" +export RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}" +export RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}" +export DECODER_LAST_PIPELINE_NUM_LAYERS="${DECODER_LAST_PIPELINE_NUM_LAYERS:-30}" +export ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}" +export MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" +# PP=2 halves both weights and activations per GPU (32 of 64 layers each). +# This provides enough memory headroom to disable truncation entirely — +# full long trajectories (50k+ tokens) can be trained without losing context. +export TRAJ_TRUNCATION_MAX_SEQ_LEN="${TRAJ_TRUNCATION_MAX_SEQ_LEN:-0}" +# GDN packed-seq monkey-patch: patches Megatron GDN forward to pass cu_seqlens +# to chunk_gated_delta_rule, enabling thd (packing) mode without NotImplementedError. +# See rl/patches/gdn_packed_seq.py for details. +export PYTHONPATH="${REPO_ROOT}/rl/patches${PYTHONPATH:+:${PYTHONPATH}}" +export USE_DYNAMIC_BATCH_SIZE="${USE_DYNAMIC_BATCH_SIZE:-true}" +export CALCULATE_PER_TOKEN_LOSS="${CALCULATE_PER_TOKEN_LOSS:-true}" +export ADVANTAGE_ESTIMATOR="${ADVANTAGE_ESTIMATOR:-grpo}" +# DAPO group filter: drop groups where all samples share the same reward +# (zero advantage, no learning signal). For patcheval, base model rarely +# solves CVEs, so most groups are all-0 and get filtered → buffer never fills +# → pipeline stalls. Default off; flip with PATCHEVAL_DAPO_FILTER=true. +export DAPO_filter="${PATCHEVAL_DAPO_FILTER:-false}" +export LR="${LR:-1e-6}" +export OPTIMIZER="${OPTIMIZER:-adam}" +export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +export ADAM_BETA1="${ADAM_BETA1:-0.9}" +export ADAM_BETA2="${ADAM_BETA2:-0.98}" +# Non-colocate mode: training (Megatron) and inference (SGLang) use SEPARATE GPUs. +# Training: 24 GPUs (3 machines, TP=4 × PP=2 × CP=3) +# Rollout: 8 GPUs (1 machine, 4 engines × 2 GPUs each) +# No release/resume memory cycle → no mamba state corruption → no flush_cache_fix needed. +export SLIME_COLOCATE="${SLIME_COLOCATE:-false}" +# CPU offload optimizer: moves fp32 master weights + Adam states (~41GB at +# TP=4/DP=2) to CPU, leaving only bf16 weights + bf16 grad on GPU. Critical +# for 27B model on 140GB GPUs where weights+optimizer would otherwise OOM. +# The env.rjob.sh default was false because of an earlier AssertionError, but +# run_slime_generator.sh now passes both --optimizer-cpu-offload and +# --use-precision-aware-optimizer, which resolves the assertion. +export OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-true}" +export USE_WANDB="${USE_WANDB:-true}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_PROJECT="${WANDB_PROJECT:-slime}" +export WANDB_GROUP="${WANDB_GROUP:-patcheval_qwen3_5_9b}" +# KV cache pool as a fraction of GPU memory (after weights). This is the +# primary throughput lever for multi-turn agent rollouts: each concurrently +# decoding request must keep its prompt+history KV resident, so KV capacity +# directly caps parallel decode (#running-req). At 0.45 the KV hits ~97% with +# 4 sessions/engine, forcing eviction (cached=0 thrashing) and queueing +# (waiting for decode to free KV, median ~26s/request). 0.6 gives +33% KV +# capacity, enough to hold 4 sessions/engine without eviction so all 4 decode +# in parallel — no queue, ~100% prefix reuse. Safe on H200 141GB: 27B weights +# ~54GB + KV 0.6*~87GB-free ≈ 52GB ≈ 106GB < 141GB. Override via env var. +# Non-colocate mode: SGLang has dedicated GPUs (no sharing with training). +# Can use much higher mem_fraction_static since no need to reserve memory +# for Megatron training weights/activations. +export SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.9}" +export SGLANG_ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-fa3}" +export SGLANG_MAMBA_SCHEDULER_STRATEGY="${SGLANG_MAMBA_SCHEDULER_STRATEGY:-extra_buffer}" +# EAGLE speculative decoding (official Qwen3.5-27B config). +# NOTE: Previously disabled due to mamba_pool CUDA illegal memory access, but that +# was actually caused by flush_cache_fix.py's abort_request corrupting mamba state. +# With flush_cache_fix removed, EAGLE should work in raw mode (as in official script). +export SGLANG_SPECULATIVE_ALGORITHM="${SGLANG_SPECULATIVE_ALGORITHM:-EAGLE}" +export SGLANG_SPECULATIVE_NUM_STEPS="${SGLANG_SPECULATIVE_NUM_STEPS:-3}" +export SGLANG_SPECULATIVE_EAGLE_TOPK="${SGLANG_SPECULATIVE_EAGLE_TOPK:-1}" +export SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS="${SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS:-4}" +export SGLANG_LOG_LEVEL="${SGLANG_LOG_LEVEL:-info}" +export SGLANG_LOG_LEVEL_HTTP="${SGLANG_LOG_LEVEL_HTTP:-error}" +export CLEANUP_BEFORE_RUN="${CLEANUP_BEFORE_RUN:-true}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export PYTHONUNBUFFERED=1 + +# ============================================================================= +# Infrastructure defaults — inlined from rl/examples/geo3k_vl/env.sh. +# Only vars actually consumed by run_slime_generator.sh / buffer_server.py / +# llm_proxy.py / slime_generator.py. geo3k-specific paths (geo3k db, geo3k +# config, qwen3-vl-2b model) are intentionally NOT carried over; patcheval +# overrides above already set those. Placed at the end so ${VAR:-default} can +# reference patcheval values set earlier (POOL_SIZE, RL_GROUP_SIZE, ports). +# ============================================================================= + +# --- Ray / Python launcher --- +export PYTHON_BIN="${PYTHON_BIN:-python3}" +export RAY_BIN="${RAY_BIN:-ray}" +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +export RAY_ADDRESS="${RAY_ADDRESS:-http://127.0.0.1:8265}" +export RAY_PORT="${RAY_PORT:-}" +export KILL_PYTHON_BEFORE_RUN="${KILL_PYTHON_BEFORE_RUN:-false}" + +# --- Slime train / checkpoint args --- +export SAVE_INTERVAL="${SAVE_INTERVAL:-20}" +export MODEL_ARGS_EXTRA="${MODEL_ARGS_EXTRA:-}" +export REF_LOAD_DIR="${REF_LOAD_DIR:-/mnt/shared-storage-user/evobox-share-gpfs2/leishanzhe/model/qwen3_8_27b_megatron}" +export CUSTOM_REWARD_POST_PROCESS_PATH="${CUSTOM_REWARD_POST_PROCESS_PATH:-}" +export SGLANG_LOGGING_CONFIG_PATH="${SGLANG_LOGGING_CONFIG_PATH:-}" + +# --- Optimizer / GRPO extras --- +export LR_DECAY_STYLE="${LR_DECAY_STYLE:-constant}" +export ENTROPY_COEF="${ENTROPY_COEF:-0.00}" +export EPS_CLIP="${EPS_CLIP:-0.2}" +export EPS_CLIP_HIGH="${EPS_CLIP_HIGH:-0.2}" +export USE_DYNAMIC_GLOBAL_BATCH_SIZE="${USE_DYNAMIC_GLOBAL_BATCH_SIZE:-false}" + +# --- W&B extras --- +export WANDB_TEAM="${WANDB_TEAM:-}" +export WANDB_ALWAYS_USE_TRAIN_STEP="${WANDB_ALWAYS_USE_TRAIN_STEP:-false}" + +# --- SGLang extras --- +# Tuned for high-concurrency rollout (POOL_SIZE=16). These three together let +# SGLang actually batch dozens of decode requests instead of running 1 at a +# time (the #running-req:1 symptom seen before): +# - MAX_RUNNING_REQUESTS: hard cap on the running decode batch. 64 gives +# headroom over POOL_SIZE=16 (multi-turn episodes overlap). +# - CUDA_GRAPH_BS: capture graphs for the batch sizes we expect to hit, so +# decode doesn't fall back to the slow non-graph path on shape change +# (fixes the first-step drop to 6.5 tok/s seen in the logs). +# - CHUNKED_PREFILL_SIZE: split long prefills into 8192-token chunks so a +# single long prompt can't starve the running decode batch. +export SGLANG_CUDA_GRAPH_BS="${SGLANG_CUDA_GRAPH_BS:-1 2 4 8 16 32}" +export SGLANG_MAX_RUNNING_REQUESTS="${SGLANG_MAX_RUNNING_REQUESTS:-64}" +export SGLANG_SCHEDULE_CONSERVATIVENESS="${SGLANG_SCHEDULE_CONSERVATIVENESS:-}" +export SGLANG_CHUNKED_PREFILL_SIZE="${SGLANG_CHUNKED_PREFILL_SIZE:-8192}" +export SGLANG_ENABLE_MIXED_CHUNK="${SGLANG_ENABLE_MIXED_CHUNK:-false}" + +# --- LLM proxy / buffer server workers & perf --- +export LLM_TOP_P="${LLM_TOP_P:-1.0}" +export LLM_PROXY_ENABLE_CONSOLE_LOG="${LLM_PROXY_ENABLE_CONSOLE_LOG:-0}" +export AIEVOBOX_LLM_MAX_CONCURRENCY="${AIEVOBOX_LLM_MAX_CONCURRENCY:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_PROXY_WORKERS="${AIEVOBOX_LLM_PROXY_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_LLM_STARTUP_JITTER_S="${AIEVOBOX_LLM_STARTUP_JITTER_S:-0}" +export AIEVOBOX_TRAININFO_WORKERS="${AIEVOBOX_TRAININFO_WORKERS:-${AIEVOBOX_POOL_SIZE}}" +export AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE="${AIEVOBOX_SQLITE_BULK_INSERT_BATCH_SIZE:-256}" +export AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S="${AIEVOBOX_SQLITE_BULK_INSERT_PAUSE_S:-0.01}" +export AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS="${AIEVOBOX_BUFFER_INCOMPLETE_GROUP_TTL_SECONDS:-1800}" +export ROLLBUF_HOST="${ROLLBUF_HOST:-${BUFFER_SERVER_HOST}}" +export ROLLBUF_PORT="${ROLLBUF_PORT:-${BUFFER_SERVER_PORT}}" + +# --- Slime rollout-buffer / GRPO filter --- +export SLIME_ROLLBUF_RESTART_TRAINING="${SLIME_ROLLBUF_RESTART_TRAINING:-True}" +export SLIME_N_SAMPLES_PER_PROMPT="${SLIME_N_SAMPLES_PER_PROMPT:-${RL_GROUP_SIZE}}" +export RL_OFF_BY_N="${RL_OFF_BY_N:-3}" +# Oversample: launch extra envs per prompt beyond group_size so the first +# group_size episodes to finish form a group; long-tail episodes don't block. +# buffer_server still pops group_size at a time; surplus stays in the bucket. +# E.g. group_size=8 + oversample=4 = 12 envs per prompt, first 8 done = 1 group. +export RL_OVERSAMPLE="${PATCHEVAL_OVERSAMPLE:-4}" + +# --- AIEVOBOX env extras --- +export AIEVOBOX_MESSAGE_CUT="${AIEVOBOX_MESSAGE_CUT:-0}" +export AIEVOBOC_MULTIPLIER="${AIEVOBOC_MULTIPLIER:-1.2}" + +# --- Runtime --- +# NOTE: expandable_segments:True is incompatible with torch_memory_saver +# (used in colocate mode). Disable it when colocate is on. +if [[ "${SLIME_COLOCATE:-false}" == "true" || "${SLIME_COLOCATE:-false}" == "1" ]]; then + export PYTORCH_CUDA_ALLOC_CONF="" + export PYTORCH_ALLOC_CONF="" +else + # PyTorch >= 2.5 renamed PYTORCH_CUDA_ALLOC_CONF → PYTORCH_ALLOC_CONF. + # Set both so old and new versions pick up expandable_segments. + export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +fi diff --git a/rl/examples/patcheval/env.sh b/rl/examples/patcheval/env.sh new file mode 100755 index 00000000..667455eb --- /dev/null +++ b/rl/examples/patcheval/env.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# +# ============================================================================= +# [DOCKER MODE] PatchEval RL environment +# ============================================================================= +# This is the DOCKER mode variant: AIEVOBOX_MODE=docker, agent containers run +# on a single remote Docker daemon (DOCKER_HOST). For the RJob variant (agent +# containers submitted as cluster jobs), see env.rjob.sh in this directory. +# ============================================================================= +# +# PatchEval RL settings for rl/run_buffer_server.sh and +# rl/run_slime_generator.sh. Generate PATCH_EVAL_GENERATED_DIR first with +# env/patcheval/generate_full_config.py. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# Keep this environment aligned with the common Slime launcher contract +# (PYTHON_BIN, RAY_BIN, optimizer, and SGLang defaults), then override all +# Geo3K-specific values below. +source "${REPO_ROOT}/rl/examples/geo3k_vl/env.sh" + +: "${PATCH_EVAL_GENERATED_DIR:?Set PATCH_EVAL_GENERATED_DIR to a generated PatchEval config directory}" + +export AIEVOBOX_EXAMPLE_NAME="patcheval_qwen3_5_9b" +export AIEVOBOX_ROOT="${AIEVOBOX_ROOT:-${REPO_ROOT}}" +export AIEVOBOX_MODE="docker" +export PATCH_EVAL_BASELINE="openhands" +export PATCH_EVAL_AGENT_EXPERIMENT="exp1" +export STORAGE_TYPE="${STORAGE_TYPE:-sqlite}" +export AIEVOBOX_DB_URL="${PATCHEVAL_DB_URL:-sqlite:///${AIEVOBOX_ROOT}/rl/examples/patcheval/patcheval_qwen3_5_9b.db}" +export DOCKER_HOST="${DOCKER_HOST:-tcp://100.99.17.62:2376}" +export AIEVOBOX_DOCKER_IMAGE_ARCHIVE_DIR="${PATCH_EVAL_IMAGE_ARCHIVE_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images}" +export AIEVOBOX_DOCKER_PULL_POLICY="${AIEVOBOX_DOCKER_PULL_POLICY:-never}" +export AIEVOBOX_AGENT_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_config.yaml" +export AIEVOBOX_AGENT_START_CONFIG="${PATCH_EVAL_GENERATED_DIR}/patcheval_start.yaml" +# Rule evaluators live in //rule_evaluator.py; the +# launcher discovers them via --agent_root, so it must point at the generated dir. +export AIEVOBOX_AGENT_ROOT="${PATCH_EVAL_GENERATED_DIR}" +export AIEVOBOX_MAX_STEPS="${PATCHEVAL_MAX_STEPS:-1}" +export AIEVOBOX_ENABLE_EVALUATION="${AIEVOBOX_ENABLE_EVALUATION:-1}" +export AIEVOBOX_POOL_SIZE="${PATCHEVAL_POOL_SIZE:-1}" +export AIEVOBOX_AGENT_START_TIMEOUT_S="${PATCHEVAL_AGENT_START_TIMEOUT_S:-1800}" + +export RL_GROUP_SIZE="${RL_GROUP_SIZE:-8}" +export RL_GLOBAL_BATCH_SIZE="${RL_GLOBAL_BATCH_SIZE:-64}" +export RL_ROLLOUT_GROUP_BATCH_SIZE="${RL_ROLLOUT_GROUP_BATCH_SIZE:-8}" +export SLIME_ROLLOUT_BATCH_SIZE="${SLIME_ROLLOUT_BATCH_SIZE:-${RL_ROLLOUT_GROUP_BATCH_SIZE}}" +export SLIME_GLOBAL_BATCH_SIZE="${SLIME_GLOBAL_BATCH_SIZE:-${RL_GLOBAL_BATCH_SIZE}}" +export RL_EPOCH="${RL_EPOCH:-1000}" +export RL_MODEL="${RL_MODEL:-model}" +export RL_API_KEY="${RL_API_KEY:-openai_api_key}" + +export BUFFER_SERVER_HOST="${BUFFER_SERVER_HOST:-127.0.0.1}" +export BUFFER_SERVER_PORT="${BUFFER_SERVER_PORT:-18889}" +export LLM_PROXY_HOST="${LLM_PROXY_HOST:-127.0.0.1}" +export LLM_PROXY_PORT="${LLM_PROXY_PORT:-18890}" +export LLM_MAX_LENGTH="${LLM_MAX_LENGTH:-8192}" +export LLM_TEMPERATURE="${LLM_TEMPERATURE:-1.0}" +export AIEVOBOX_GATEWAY_HOST="${PATCHEVAL_GATEWAY_HOST:-$(hostname -i | awk '{print $1}')}" +export AIEVOBOX_GATEWAY_PORT="${PATCHEVAL_GATEWAY_PORT:-8000}" +export AIEVOBOX_GATEWAY_BASE_URL="http://${AIEVOBOX_GATEWAY_HOST}:${AIEVOBOX_GATEWAY_PORT}/v1/sessions" + +export SLIME_HOME="${SLIME_HOME:-/root/slime}" +export MEGATRON_HOME="${MEGATRON_HOME:-/root/Megatron-LM}" +export HF_CKPT_DIR="${QWEN3_5_9B_CKPT_DIR:-/mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/hub/models--Qwen--Qwen3.5-9B/snapshots/c202236235762e1c871ad0ccb60c8ee5ba337b9a}" +export LOAD_DIR="${QWEN3_5_9B_LOAD_DIR:-${HF_CKPT_DIR}}" +export SAVE_DIR="${SAVE_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/checkpoints/Qwen3.5-9B_megatron}" +export WANDB_DIR="${WANDB_DIR:-${AIEVOBOX_ROOT}/rl/examples/patcheval/wandb_logs}" +export LOG_ROOT="${LOG_ROOT:-${AIEVOBOX_ROOT}/logs/patcheval_qwen3_5_9b}" +export MODEL_SCRIPT="${QWEN3_5_9B_MODEL_SCRIPT:-${AIEVOBOX_ROOT}/rl/examples/geo3k_vl/qwen3_5_9b.sh}" +export MODEL_ARGS_ROTARY_BASE=10000000 + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export NUM_GPUS=4 +export ACTOR_NUM_NODES=1 +export ACTOR_NUM_GPUS_PER_NODE=1 +export ROLLOUT_NUM_GPUS=3 +export ROLLOUT_NUM_GPUS_PER_ENGINE=1 +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-${SLIME_HOME}/train.py}" +export ROLLOUT_FUNCTION_PATH="${ROLLOUT_FUNCTION_PATH:-rl.slime_generator.generate_rollout}" +export NUM_ROLLOUT="${NUM_ROLLOUT:-300}" +export LOSS_MASK_TYPE="qwen3_5" +export TRAIN_BACKEND="${TRAIN_BACKEND:-megatron}" +export MEGATRON_TO_HF_MODE="${MEGATRON_TO_HF_MODE:-bridge}" +export TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 EP_SIZE=1 ETP_SIZE=1 +export RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}" +export RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}" +export RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}" +export ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}" +export MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-5000}" +export USE_DYNAMIC_BATCH_SIZE="${USE_DYNAMIC_BATCH_SIZE:-true}" +export CALCULATE_PER_TOKEN_LOSS="${CALCULATE_PER_TOKEN_LOSS:-true}" +export ADVANTAGE_ESTIMATOR="${ADVANTAGE_ESTIMATOR:-grpo}" +export LR="${LR:-1e-6}" +export OPTIMIZER="${OPTIMIZER:-adam}" +export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +export ADAM_BETA1="${ADAM_BETA1:-0.9}" +export ADAM_BETA2="${ADAM_BETA2:-0.98}" +export USE_WANDB="${USE_WANDB:-true}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_PROJECT="${WANDB_PROJECT:-slime}" +export WANDB_GROUP="${WANDB_GROUP:-patcheval_qwen3_5_9b}" +export SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.45}" +export SGLANG_ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-fa3}" +export SGLANG_LOG_LEVEL="${SGLANG_LOG_LEVEL:-info}" +export SGLANG_LOG_LEVEL_HTTP="${SGLANG_LOG_LEVEL_HTTP:-error}" +export CLEANUP_BEFORE_RUN="${CLEANUP_BEFORE_RUN:-true}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export PYTHONUNBUFFERED=1 diff --git a/rl/examples/patcheval/run_eval.sh b/rl/examples/patcheval/run_eval.sh index 59c35164..1d597537 100755 --- a/rl/examples/patcheval/run_eval.sh +++ b/rl/examples/patcheval/run_eval.sh @@ -194,6 +194,11 @@ if [[ "${PATCH_EVAL_STORAGE_TYPE}" == "sqlite" ]]; then launcher_storage_args+=(--db-path "sqlite:///${PATCH_EVAL_DB}") fi +launcher_job_id_args=() +if [[ -n "${PATCH_EVAL_JOB_ID:-}" ]]; then + launcher_job_id_args=(--job-id "${PATCH_EVAL_JOB_ID}") +fi + "${PYTHON_BIN}" launcher.py \ --mode docker \ --docker-pull-policy never \ @@ -212,6 +217,7 @@ fi --llm-model "${PATCH_EVAL_MODEL}" \ --llm-temperature 0 \ "${launcher_storage_args[@]}" \ + "${launcher_job_id_args[@]}" \ --pool-size "${PATCH_EVAL_POOL_SIZE}" \ --max-workers "${PATCH_EVAL_POOL_SIZE}" \ --max-steps 1 \ diff --git a/rl/examples/patcheval/run_eval_rjob.qwen3_5_9b.sh b/rl/examples/patcheval/run_eval_rjob.qwen3_5_9b.sh new file mode 100755 index 00000000..aee4d8f8 --- /dev/null +++ b/rl/examples/patcheval/run_eval_rjob.qwen3_5_9b.sh @@ -0,0 +1,290 @@ +#!/usr/bin/env bash +# ============================================================================= +# PatchEval — RJob mode evaluation runner (LLM baseline, s1.1~s1.4) +# ============================================================================= +# Qwen3.5-9B variant of run_eval_rjob.sh. Mirrors the 27B runner but points +# the LLM endpoint at the Qwen3.5-9B SGLang service. +# Architecture: +# SGLang (Qwen3.5-9B) -> :30000 (GPU inference host) +# Gateway (this host) -> 100.99.17.62:8000 (fronts SGLang, reachable +# from RJob pods in cluster) +# launcher.py -> submits each CVE episode as a cluster RJob pod +# that pulls CVE images from the private registry +# and calls back the gateway for LLM completions. +# +# Prereqs (must already be done once): +# 1. SGLang running for Qwen3.5-9B on :30000. +# 2. CVE images pushed to registry.h.pjlab.org.cn (push_patcheval_images.sh). +# 3. config.yaml at repo root has valid rjob access_key/secret_key. +# +# Usage: +# PATCH_EVAL_API_KEY=sk-qwen35-9b-local ./run_eval_rjob.qwen3_5_9b.sh s1.1 # one setting +# PATCH_EVAL_API_KEY=sk-qwen35-9b-local ./run_eval_rjob.qwen3_5_9b.sh s1.1 s1.2 s1.3 s1.4 +# PATCH_EVAL_API_KEY=sk-qwen35-9b-local PATCH_EVAL_TASK_LIMIT=2 ./run_eval_rjob.qwen3_5_9b.sh s1.1 # smoke test +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# --- Required --- +: "${PATCH_EVAL_API_KEY:?Set PATCH_EVAL_API_KEY (the SGLang api-key, e.g. sk-qwen35-9b-local)}" + +# --- Tunables (env-overridable) --- +PYTHON_BIN=${PYTHON_BIN:-/mnt/shared-storage-user/evobox-share/leishanzhe/env/slime-env-0.3.1/bin/python} + +# The conda env bundles libicui18n which needs CXXABI_1.3.15 from a newer +# libstdc++ than the system one. Prepend the env's own lib/ so its bundled +# libstdc++.so.6 is loaded (otherwise `import sqlite3` fails and the gateway +# can't init its sqlite storage). +ENV_LIB_DIR=$(dirname "$(dirname "${PYTHON_BIN}")")/lib +if [[ -d "${ENV_LIB_DIR}" ]]; then + export LD_LIBRARY_PATH="${ENV_LIB_DIR}:${LD_LIBRARY_PATH:-}" +fi + +# SGLang endpoint (inference host running Qwen3.5-9B) +SGLANG_HOST=${SGLANG_HOST:-100.104.113.76} +SGLANG_PORT=${SGLANG_PORT:-30000} +SGLANG_MODEL=${SGLANG_MODEL:-qwen3.5-9b} + +# Gateway (this host, must be reachable from RJob pods) +GATEWAY_HOST=${GATEWAY_HOST:-$(hostname -I | awk '{print $1}')} +GATEWAY_PORT=${GATEWAY_PORT:-18000} # 8000 is taken by nginx on this host + +# Baseline / settings +PATCH_EVAL_BASELINE=${PATCH_EVAL_BASELINE:-llm} +SETTINGS=() +if [[ $# -gt 0 ]]; then SETTINGS=("$@"); else SETTINGS=(s1.1 s1.2 s1.3 s1.4); fi + +# Eval knobs +PATCH_EVAL_TASK_LIMIT=${PATCH_EVAL_TASK_LIMIT:-0} # 0 = all 230 CVEs +PATCH_EVAL_POOL_SIZE=${PATCH_EVAL_POOL_SIZE:-8} # parallel RJob pods (8-GPU machine) +PATCH_EVAL_AGENT_TIMEOUT_S=${PATCH_EVAL_AGENT_TIMEOUT_S:-900} +PATCH_EVAL_EVALUATION_TIMEOUT_S=${PATCH_EVAL_EVALUATION_TIMEOUT_S:-3600} +PATCH_EVAL_SHUTDOWN_TIMEOUT_S=${PATCH_EVAL_SHUTDOWN_TIMEOUT_S:-600} +PATCH_EVAL_STORAGE_TYPE=${PATCH_EVAL_STORAGE_TYPE:-sqlite} + +# Paths +PATCH_EVAL_IMAGE_ARCHIVE_DIR=${PATCH_EVAL_IMAGE_ARCHIVE_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images} +PATCH_EVAL_OFFICIAL_RUNTIME_DIR=${PATCH_EVAL_OFFICIAL_RUNTIME_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-runtime} +PATCH_EVAL_SHARED_TMP=${PATCH_EVAL_SHARED_TMP:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-tmp} +PATCH_EVAL_HTTP_PROXY=${PATCH_EVAL_HTTP_PROXY:-http://httpproxy-headless.kubebrain.svc.pjlab.local:3128} +PATCH_EVAL_NO_PROXY=${PATCH_EVAL_NO_PROXY:-localhost,127.0.0.1,::1,${GATEWAY_HOST},10.0.0.0/8,100.96.0.0/12,.pjlab.org.cn} + +# RJob registry (must match push_patcheval_images.sh) +RJOB_REGISTRY=${RJOB_REGISTRY:-registry.h.pjlab.org.cn} +RJOB_REGISTRY_NS=${RJOB_REGISTRY_NS:-ailab-evobox-evobox_proxy} +RJOB_REPO=${RJOB_REPO:-patcheval} +RJOB_CONFIG=${RJOB_CONFIG:-${ROOT}/config.yaml} + +run_id=${PATCH_EVAL_RUN_ID:-$(date +%Y%m%d-%H%M%S)} + +# Results root for safactory_result.json. MUST be on shared storage (RJob pods +# in the cluster write here via gpfs mount), but isolated from RL runs by an +# eval-specific subdir + run_id so eval and RL results never mix. +RJOB_RESULTS_ROOT=${RJOB_RESULTS_ROOT:-${ROOT}/results/patcheval_eval/${run_id}} + +mkdir -p "${PATCH_EVAL_SHARED_TMP}" "${SCRIPT_DIR}/logs" +export TMPDIR="${PATCH_EVAL_SHARED_TMP}" + +# SQLite DBs live under the SAfactory repo (so they're persistent & collectible +# from shared storage), in an eval-specific subdir + run_id so they never clash +# with RL's patcheval_qwen3_8_27b.db. The repo is on GPFS; if sqlite hits a +# transient "disk I/O error" under concurrent load, override +# PATCH_EVAL_LOCAL_DB_DIR to a local /tmp path for that run. +LOCAL_DB_DIR="${PATCH_EVAL_LOCAL_DB_DIR:-${SCRIPT_DIR}/eval_runs/${run_id}}" +mkdir -p "${LOCAL_DB_DIR}" + +# --------------------------------------------------------------------------- +# Sanity checks +# --------------------------------------------------------------------------- +echo "=== Pre-flight ===" +echo "LLM : http://${SGLANG_HOST}:${SGLANG_PORT} (model=${SGLANG_MODEL})" +echo "Gateway : http://${GATEWAY_HOST}:${GATEWAY_PORT} (this host, reachable from RJob pods)" +echo "RJob cfg : ${RJOB_CONFIG}" +echo "Registry : ${RJOB_REGISTRY}/${RJOB_REGISTRY_NS}/${RJOB_REPO}" +echo "Settings : ${SETTINGS[*]} (limit=${PATCH_EVAL_TASK_LIMIT}, pool=${PATCH_EVAL_POOL_SIZE})" + +# LLM endpoint reachability — try /health (SGLang), then /v1/models (OpenAI proxies +# like the claude-opus-5 gateway that lack /health). Either passing is enough. +LLM_BASE="http://${SGLANG_HOST}:${SGLANG_PORT}" +llm_ok=0 +if curl -fsS --max-time 8 "${LLM_BASE}/health" >/dev/null 2>&1; then + llm_ok=1 +elif curl -fsS --max-time 10 "${LLM_BASE}/v1/models" -H "Authorization: Bearer ${PATCH_EVAL_API_KEY}" >/dev/null 2>&1; then + llm_ok=1 +fi +if [[ "${llm_ok}" -ne 1 ]]; then + echo "ERROR: LLM endpoint not reachable at ${LLM_BASE} (tried /health and /v1/models)." >&2 + echo " For SGLang: start it first (start_sglang_qwen35_9b.sh)." >&2 + echo " For opus-5: check the proxy at ${LLM_BASE} and PATCH_EVAL_API_KEY." >&2 + exit 1 +fi +[[ -f "${RJOB_CONFIG}" ]] || { echo "ERROR: RJob config not found: ${RJOB_CONFIG}" >&2; exit 1; } + +# --------------------------------------------------------------------------- +# Stage official runtime helpers (same as docker run_eval.sh) +# --------------------------------------------------------------------------- +OFFICIAL_SOURCE="${ROOT}/env/patcheval/PatchEval/patcheval" +mkdir -p "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/helper" +mkdir -p "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_agent/claudecode/templates" +cp \ + "${OFFICIAL_SOURCE}/exp_llm/helper/llm_suite.py" \ + "${OFFICIAL_SOURCE}/exp_llm/helper/func_replacer.py" \ + "${OFFICIAL_SOURCE}/exp_llm/helper/__init__.py" \ + "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/helper/" 2>/dev/null || true +touch "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/__init__.py" +cp \ + "${OFFICIAL_SOURCE}/exp_agent/claudecode/templates/default.md" \ + "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_agent/claudecode/templates/default.md" 2>/dev/null || true + +# --------------------------------------------------------------------------- +# Start the gateway (one shared instance across all settings) +# --------------------------------------------------------------------------- +GATEWAY_CONFIG="${PATCH_EVAL_SHARED_TMP}/safactory-patcheval-gateway-rjob-${run_id}.yaml" +GATEWAY_LOG="${SCRIPT_DIR}/logs/gateway-rjob-${run_id}.log" + +PATCH_EVAL_DB_BASE="${PATCH_EVAL_DB_BASE:-${LOCAL_DB_DIR}/patcheval_${SGLANG_MODEL}_rjob_${run_id}.db}" + +model_slug="${SGLANG_MODEL//\//_}" +model_slug="${model_slug//:/_}" + +PATCH_EVAL_DB="${PATCH_EVAL_DB_BASE}" \ +PATCH_EVAL_API_BASE="http://${SGLANG_HOST}:${SGLANG_PORT}/v1" \ +PATCH_EVAL_API_KEY="${PATCH_EVAL_API_KEY}" \ +PATCH_EVAL_MODEL="${SGLANG_MODEL}" \ +PATCH_EVAL_GATEWAY_PORT="${GATEWAY_PORT}" \ +PATCH_EVAL_STORAGE_TYPE="${PATCH_EVAL_STORAGE_TYPE}" \ +GATEWAY_CONFIG="${GATEWAY_CONFIG}" \ +"${PYTHON_BIN}" - <<'PY' +import os +from pathlib import Path +import yaml + +db = Path(os.environ["PATCH_EVAL_DB"]).expanduser().resolve() +storage_type = os.environ["PATCH_EVAL_STORAGE_TYPE"] +config = { + "listen_host": "0.0.0.0", + "listen_port": int(os.environ["PATCH_EVAL_GATEWAY_PORT"]), + "base_session_path": "/v1/sessions", + "max_steps": -1, + "storage_type": storage_type, + "storage_config": ( + {"db_url": f"sqlite:///{db}"} if storage_type == "sqlite" else {} + ), + "llm_routes": { + os.environ["PATCH_EVAL_MODEL"]: { + "base_url": os.environ["PATCH_EVAL_API_BASE"].rstrip("/") + "/", + "api_key": os.environ["PATCH_EVAL_API_KEY"], + "supports_stream": True, + "max_concurrency": 64, + } + }, +} +path = Path(os.environ["GATEWAY_CONFIG"]) +path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") +path.chmod(0o600) +PY + +cleanup() { + if [[ -n "${GATEWAY_PID:-}" ]] && kill -0 "${GATEWAY_PID}" 2>/dev/null; then + kill "${GATEWAY_PID}" 2>/dev/null || true + wait "${GATEWAY_PID}" 2>/dev/null || true + fi + rm -f -- "${GATEWAY_CONFIG}" +} +trap cleanup EXIT INT TERM + +cd "${ROOT}" +"${PYTHON_BIN}" -u -m gateway --config "${GATEWAY_CONFIG}" >"${GATEWAY_LOG}" 2>&1 & +GATEWAY_PID=$! + +for _ in $(seq 1 60); do + if ! kill -0 "${GATEWAY_PID}" 2>/dev/null; then + echo "Gateway exited early; inspect ${GATEWAY_LOG}" >&2 + exit 1 + fi + if curl -fsS --max-time 2 "http://127.0.0.1:${GATEWAY_PORT}/readyz" >/dev/null; then + break + fi + sleep 1 +done +if ! curl -fsS --max-time 2 "http://127.0.0.1:${GATEWAY_PORT}/readyz" >/dev/null; then + echo "Gateway did not become ready; inspect ${GATEWAY_LOG}" >&2 + exit 1 +fi +echo "Gateway ready: http://${GATEWAY_HOST}:${GATEWAY_PORT}/v1/sessions (log: ${GATEWAY_LOG})" + +# --------------------------------------------------------------------------- +# Per-setting loop +# --------------------------------------------------------------------------- +GATEWAY_BASE_URL="http://${GATEWAY_HOST}:${GATEWAY_PORT}/v1/sessions" + +for setting in "${SETTINGS[@]}"; do + echo "" + echo "############################################################" + echo "# Setting ${setting} (baseline=${PATCH_EVAL_BASELINE})" + echo "############################################################" + + GENERATED_DIR="${PATCH_EVAL_SHARED_TMP}/safactory-patcheval-rjob-${setting}-${run_id}" + mkdir -p "${GENERATED_DIR}" + + # Per-setting DB so results don't collide across settings. + # The launcher --db-path MUST be the same file the gateway writes to (the + # evaluator reads the gateway's trajectory DB). One shared DB per run, + # episodes are distinguished by session_id/job_id (matches run_eval.sh). + setting_db="${PATCH_EVAL_DB_BASE}" + + "${PYTHON_BIN}" env/patcheval/generate_full_config.py \ + --output-dir "${GENERATED_DIR}" \ + --archive-dir "${PATCH_EVAL_IMAGE_ARCHIVE_DIR}" \ + --official-runtime-dir "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}" \ + --baseline "${PATCH_EVAL_BASELINE}" \ + --setting "${setting}" \ + --claude-gateway-base-url "${GATEWAY_BASE_URL}" \ + --claude-model "${SGLANG_MODEL}" \ + --limit "${PATCH_EVAL_TASK_LIMIT}" \ + --evaluation-timeout-s "${PATCH_EVAL_EVALUATION_TIMEOUT_S}" \ + --shared-tmp "${PATCH_EVAL_SHARED_TMP}" \ + --http-proxy "${PATCH_EVAL_HTTP_PROXY}" \ + --no-proxy "${PATCH_EVAL_NO_PROXY}" \ + --mode rjob \ + --rjob-registry "${RJOB_REGISTRY}" \ + --rjob-registry-ns "${RJOB_REGISTRY_NS}" \ + --rjob-repo "${RJOB_REPO}" \ + --rjob-results-root "${RJOB_RESULTS_ROOT}" + + echo "Generated rjob config: ${GENERATED_DIR}" + echo " config: ${GENERATED_DIR}/patcheval_config.rjob.yaml" + echo " start : ${GENERATED_DIR}/patcheval_start.rjob.yaml" + + launcher_storage_args=(--storage-type "${PATCH_EVAL_STORAGE_TYPE}") + if [[ "${PATCH_EVAL_STORAGE_TYPE}" == "sqlite" ]]; then + launcher_storage_args+=(--db-path "sqlite:///${setting_db}") + fi + + "${PYTHON_BIN}" launcher.py \ + --mode rjob \ + --rjob-config "${RJOB_CONFIG}" \ + --agent-root "${GENERATED_DIR}" \ + --agent-config "${GENERATED_DIR}/patcheval_config.rjob.yaml" \ + --agent-start-config "${GENERATED_DIR}/patcheval_start.rjob.yaml" \ + --gateway-base-url "${GATEWAY_BASE_URL}" \ + --llm-model "${SGLANG_MODEL}" \ + --llm-temperature 0 \ + --agent-start-timeout-s "${PATCH_EVAL_AGENT_TIMEOUT_S}" \ + --shutdown-timeout-s "${PATCH_EVAL_SHUTDOWN_TIMEOUT_S}" \ + "${launcher_storage_args[@]}" \ + --pool-size "${PATCH_EVAL_POOL_SIZE}" \ + --max-workers "${PATCH_EVAL_POOL_SIZE}" \ + --max-steps 1 \ + --enable-evaluation \ + --no-circuit-breaker \ + ${PATCHEVAL_RESUME:+--resume} + + echo "Setting ${setting} done. Results DB: ${setting_db}" +done + +echo "" +echo "=== All settings complete ===" +echo "Gateway log: ${GATEWAY_LOG}" diff --git a/rl/examples/patcheval/run_eval_rjob.sh b/rl/examples/patcheval/run_eval_rjob.sh new file mode 100755 index 00000000..f69bb46c --- /dev/null +++ b/rl/examples/patcheval/run_eval_rjob.sh @@ -0,0 +1,288 @@ +#!/usr/bin/env bash +# ============================================================================= +# PatchEval — RJob mode evaluation runner (LLM baseline, s1.1~s1.4) +# ============================================================================= +# Architecture: +# SGLang (Qwen3.8-27B) -> 100.104.113.76:30000 (GPU inference host) +# Gateway (this host) -> 100.99.17.62:8000 (fronts SGLang, reachable +# from RJob pods in cluster) +# launcher.py -> submits each CVE episode as a cluster RJob pod +# that pulls CVE images from the private registry +# and calls back the gateway for LLM completions. +# +# Prereqs (must already be done once): +# 1. SGLang running on 100.104.113.76:30000 (see start_sglang_qwen38_27b.sh). +# 2. CVE images pushed to registry.h.pjlab.org.cn (push_patcheval_images.sh). +# 3. config.yaml at repo root has valid rjob access_key/secret_key. +# +# Usage: +# PATCH_EVAL_API_KEY=sk-qwen38-27b-local ./run_eval_rjob.sh s1.1 # one setting +# PATCH_EVAL_API_KEY=sk-qwen38-27b-local ./run_eval_rjob.sh s1.1 s1.2 s1.3 s1.4 +# PATCH_EVAL_API_KEY=sk-qwen38-27b-local PATCH_EVAL_TASK_LIMIT=2 ./run_eval_rjob.sh s1.1 # smoke test +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +ROOT="$(cd -- "${SCRIPT_DIR}/../../.." &>/dev/null && pwd)" + +# --- Required --- +: "${PATCH_EVAL_API_KEY:?Set PATCH_EVAL_API_KEY (the SGLang api-key, e.g. sk-qwen38-27b-local)}" + +# --- Tunables (env-overridable) --- +PYTHON_BIN=${PYTHON_BIN:-/mnt/shared-storage-user/evobox-share/leishanzhe/env/slime-env-0.3.1/bin/python} + +# The conda env bundles libicui18n which needs CXXABI_1.3.15 from a newer +# libstdc++ than the system one. Prepend the env's own lib/ so its bundled +# libstdc++.so.6 is loaded (otherwise `import sqlite3` fails and the gateway +# can't init its sqlite storage). +ENV_LIB_DIR=$(dirname "$(dirname "${PYTHON_BIN}")")/lib +if [[ -d "${ENV_LIB_DIR}" ]]; then + export LD_LIBRARY_PATH="${ENV_LIB_DIR}:${LD_LIBRARY_PATH:-}" +fi + +# SGLang endpoint (inference host) +SGLANG_HOST=${SGLANG_HOST:-100.104.113.76} +SGLANG_PORT=${SGLANG_PORT:-30000} +SGLANG_MODEL=${SGLANG_MODEL:-qwen3.8-27b} + +# Gateway (this host, must be reachable from RJob pods) +GATEWAY_HOST=${GATEWAY_HOST:-$(hostname -I | awk '{print $1}')} +GATEWAY_PORT=${GATEWAY_PORT:-18000} # 8000 is taken by nginx on this host + +# Baseline / settings +PATCH_EVAL_BASELINE=${PATCH_EVAL_BASELINE:-llm} +SETTINGS=() +if [[ $# -gt 0 ]]; then SETTINGS=("$@"); else SETTINGS=(s1.1 s1.2 s1.3 s1.4); fi + +# Eval knobs +PATCH_EVAL_TASK_LIMIT=${PATCH_EVAL_TASK_LIMIT:-0} # 0 = all 230 CVEs +PATCH_EVAL_POOL_SIZE=${PATCH_EVAL_POOL_SIZE:-16} # parallel RJob pods +PATCH_EVAL_AGENT_TIMEOUT_S=${PATCH_EVAL_AGENT_TIMEOUT_S:-900} +PATCH_EVAL_EVALUATION_TIMEOUT_S=${PATCH_EVAL_EVALUATION_TIMEOUT_S:-3600} +PATCH_EVAL_SHUTDOWN_TIMEOUT_S=${PATCH_EVAL_SHUTDOWN_TIMEOUT_S:-600} +PATCH_EVAL_STORAGE_TYPE=${PATCH_EVAL_STORAGE_TYPE:-sqlite} + +# Paths +PATCH_EVAL_IMAGE_ARCHIVE_DIR=${PATCH_EVAL_IMAGE_ARCHIVE_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-images} +PATCH_EVAL_OFFICIAL_RUNTIME_DIR=${PATCH_EVAL_OFFICIAL_RUNTIME_DIR:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-runtime} +PATCH_EVAL_SHARED_TMP=${PATCH_EVAL_SHARED_TMP:-/mnt/shared-storage-user/evobox-share/leishanzhe/dataset/patcheval-tmp} +PATCH_EVAL_HTTP_PROXY=${PATCH_EVAL_HTTP_PROXY:-http://httpproxy-headless.kubebrain.svc.pjlab.local:3128} +PATCH_EVAL_NO_PROXY=${PATCH_EVAL_NO_PROXY:-localhost,127.0.0.1,::1,${GATEWAY_HOST},10.0.0.0/8,100.96.0.0/12,.pjlab.org.cn} + +# RJob registry (must match push_patcheval_images.sh) +RJOB_REGISTRY=${RJOB_REGISTRY:-registry.h.pjlab.org.cn} +RJOB_REGISTRY_NS=${RJOB_REGISTRY_NS:-ailab-evobox-evobox_proxy} +RJOB_REPO=${RJOB_REPO:-patcheval} +RJOB_CONFIG=${RJOB_CONFIG:-${ROOT}/config.yaml} + +run_id=${PATCH_EVAL_RUN_ID:-$(date +%Y%m%d-%H%M%S)} + +# Results root for safactory_result.json. MUST be on shared storage (RJob pods +# in the cluster write here via gpfs mount), but isolated from RL runs by an +# eval-specific subdir + run_id so eval and RL results never mix. +RJOB_RESULTS_ROOT=${RJOB_RESULTS_ROOT:-${ROOT}/results/patcheval_eval/${run_id}} + +mkdir -p "${PATCH_EVAL_SHARED_TMP}" "${SCRIPT_DIR}/logs" +export TMPDIR="${PATCH_EVAL_SHARED_TMP}" + +# SQLite DBs live under the SAfactory repo (so they're persistent & collectible +# from shared storage), in an eval-specific subdir + run_id so they never clash +# with RL's patcheval_qwen3_8_27b.db. The repo is on GPFS; if sqlite hits a +# transient "disk I/O error" under concurrent load, override +# PATCH_EVAL_LOCAL_DB_DIR to a local /tmp path for that run. +LOCAL_DB_DIR="${PATCH_EVAL_LOCAL_DB_DIR:-${SCRIPT_DIR}/eval_runs/${run_id}}" +mkdir -p "${LOCAL_DB_DIR}" + +# --------------------------------------------------------------------------- +# Sanity checks +# --------------------------------------------------------------------------- +echo "=== Pre-flight ===" +echo "LLM : http://${SGLANG_HOST}:${SGLANG_PORT} (model=${SGLANG_MODEL})" +echo "Gateway : http://${GATEWAY_HOST}:${GATEWAY_PORT} (this host, reachable from RJob pods)" +echo "RJob cfg : ${RJOB_CONFIG}" +echo "Registry : ${RJOB_REGISTRY}/${RJOB_REGISTRY_NS}/${RJOB_REPO}" +echo "Settings : ${SETTINGS[*]} (limit=${PATCH_EVAL_TASK_LIMIT}, pool=${PATCH_EVAL_POOL_SIZE})" + +# LLM endpoint reachability — try /health (SGLang), then /v1/models (OpenAI proxies +# like the claude-opus-5 gateway that lack /health). Either passing is enough. +LLM_BASE="http://${SGLANG_HOST}:${SGLANG_PORT}" +llm_ok=0 +if curl -fsS --max-time 8 "${LLM_BASE}/health" >/dev/null 2>&1; then + llm_ok=1 +elif curl -fsS --max-time 10 "${LLM_BASE}/v1/models" -H "Authorization: Bearer ${PATCH_EVAL_API_KEY}" >/dev/null 2>&1; then + llm_ok=1 +fi +if [[ "${llm_ok}" -ne 1 ]]; then + echo "ERROR: LLM endpoint not reachable at ${LLM_BASE} (tried /health and /v1/models)." >&2 + echo " For SGLang: start it first (start_sglang_qwen38_27b.sh)." >&2 + echo " For opus-5: check the proxy at ${LLM_BASE} and PATCH_EVAL_API_KEY." >&2 + exit 1 +fi +[[ -f "${RJOB_CONFIG}" ]] || { echo "ERROR: RJob config not found: ${RJOB_CONFIG}" >&2; exit 1; } + +# --------------------------------------------------------------------------- +# Stage official runtime helpers (same as docker run_eval.sh) +# --------------------------------------------------------------------------- +OFFICIAL_SOURCE="${ROOT}/env/patcheval/PatchEval/patcheval" +mkdir -p "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/helper" +mkdir -p "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_agent/claudecode/templates" +cp \ + "${OFFICIAL_SOURCE}/exp_llm/helper/llm_suite.py" \ + "${OFFICIAL_SOURCE}/exp_llm/helper/func_replacer.py" \ + "${OFFICIAL_SOURCE}/exp_llm/helper/__init__.py" \ + "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/helper/" 2>/dev/null || true +touch "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_llm/__init__.py" +cp \ + "${OFFICIAL_SOURCE}/exp_agent/claudecode/templates/default.md" \ + "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}/exp_agent/claudecode/templates/default.md" 2>/dev/null || true + +# --------------------------------------------------------------------------- +# Start the gateway (one shared instance across all settings) +# --------------------------------------------------------------------------- +GATEWAY_CONFIG="${PATCH_EVAL_SHARED_TMP}/safactory-patcheval-gateway-rjob-${run_id}.yaml" +GATEWAY_LOG="${SCRIPT_DIR}/logs/gateway-rjob-${run_id}.log" + +PATCH_EVAL_DB_BASE="${PATCH_EVAL_DB_BASE:-${LOCAL_DB_DIR}/patcheval_${SGLANG_MODEL}_rjob_${run_id}.db}" + +model_slug="${SGLANG_MODEL//\//_}" +model_slug="${model_slug//:/_}" + +PATCH_EVAL_DB="${PATCH_EVAL_DB_BASE}" \ +PATCH_EVAL_API_BASE="http://${SGLANG_HOST}:${SGLANG_PORT}/v1" \ +PATCH_EVAL_API_KEY="${PATCH_EVAL_API_KEY}" \ +PATCH_EVAL_MODEL="${SGLANG_MODEL}" \ +PATCH_EVAL_GATEWAY_PORT="${GATEWAY_PORT}" \ +PATCH_EVAL_STORAGE_TYPE="${PATCH_EVAL_STORAGE_TYPE}" \ +GATEWAY_CONFIG="${GATEWAY_CONFIG}" \ +"${PYTHON_BIN}" - <<'PY' +import os +from pathlib import Path +import yaml + +db = Path(os.environ["PATCH_EVAL_DB"]).expanduser().resolve() +storage_type = os.environ["PATCH_EVAL_STORAGE_TYPE"] +config = { + "listen_host": "0.0.0.0", + "listen_port": int(os.environ["PATCH_EVAL_GATEWAY_PORT"]), + "base_session_path": "/v1/sessions", + "max_steps": -1, + "storage_type": storage_type, + "storage_config": ( + {"db_url": f"sqlite:///{db}"} if storage_type == "sqlite" else {} + ), + "llm_routes": { + os.environ["PATCH_EVAL_MODEL"]: { + "base_url": os.environ["PATCH_EVAL_API_BASE"].rstrip("/") + "/", + "api_key": os.environ["PATCH_EVAL_API_KEY"], + "supports_stream": True, + "max_concurrency": 64, + } + }, +} +path = Path(os.environ["GATEWAY_CONFIG"]) +path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") +path.chmod(0o600) +PY + +cleanup() { + if [[ -n "${GATEWAY_PID:-}" ]] && kill -0 "${GATEWAY_PID}" 2>/dev/null; then + kill "${GATEWAY_PID}" 2>/dev/null || true + wait "${GATEWAY_PID}" 2>/dev/null || true + fi + rm -f -- "${GATEWAY_CONFIG}" +} +trap cleanup EXIT INT TERM + +cd "${ROOT}" +"${PYTHON_BIN}" -u -m gateway --config "${GATEWAY_CONFIG}" >"${GATEWAY_LOG}" 2>&1 & +GATEWAY_PID=$! + +for _ in $(seq 1 60); do + if ! kill -0 "${GATEWAY_PID}" 2>/dev/null; then + echo "Gateway exited early; inspect ${GATEWAY_LOG}" >&2 + exit 1 + fi + if curl -fsS --max-time 2 "http://127.0.0.1:${GATEWAY_PORT}/readyz" >/dev/null; then + break + fi + sleep 1 +done +if ! curl -fsS --max-time 2 "http://127.0.0.1:${GATEWAY_PORT}/readyz" >/dev/null; then + echo "Gateway did not become ready; inspect ${GATEWAY_LOG}" >&2 + exit 1 +fi +echo "Gateway ready: http://${GATEWAY_HOST}:${GATEWAY_PORT}/v1/sessions (log: ${GATEWAY_LOG})" + +# --------------------------------------------------------------------------- +# Per-setting loop +# --------------------------------------------------------------------------- +GATEWAY_BASE_URL="http://${GATEWAY_HOST}:${GATEWAY_PORT}/v1/sessions" + +for setting in "${SETTINGS[@]}"; do + echo "" + echo "############################################################" + echo "# Setting ${setting} (baseline=${PATCH_EVAL_BASELINE})" + echo "############################################################" + + GENERATED_DIR="${PATCH_EVAL_SHARED_TMP}/safactory-patcheval-rjob-${setting}-${run_id}" + mkdir -p "${GENERATED_DIR}" + + # Per-setting DB so results don't collide across settings. + # The launcher --db-path MUST be the same file the gateway writes to (the + # evaluator reads the gateway's trajectory DB). One shared DB per run, + # episodes are distinguished by session_id/job_id (matches run_eval.sh). + setting_db="${PATCH_EVAL_DB_BASE}" + + "${PYTHON_BIN}" env/patcheval/generate_full_config.py \ + --output-dir "${GENERATED_DIR}" \ + --archive-dir "${PATCH_EVAL_IMAGE_ARCHIVE_DIR}" \ + --official-runtime-dir "${PATCH_EVAL_OFFICIAL_RUNTIME_DIR}" \ + --baseline "${PATCH_EVAL_BASELINE}" \ + --setting "${setting}" \ + --claude-gateway-base-url "${GATEWAY_BASE_URL}" \ + --claude-model "${SGLANG_MODEL}" \ + --limit "${PATCH_EVAL_TASK_LIMIT}" \ + --evaluation-timeout-s "${PATCH_EVAL_EVALUATION_TIMEOUT_S}" \ + --shared-tmp "${PATCH_EVAL_SHARED_TMP}" \ + --http-proxy "${PATCH_EVAL_HTTP_PROXY}" \ + --no-proxy "${PATCH_EVAL_NO_PROXY}" \ + --mode rjob \ + --rjob-registry "${RJOB_REGISTRY}" \ + --rjob-registry-ns "${RJOB_REGISTRY_NS}" \ + --rjob-repo "${RJOB_REPO}" \ + --rjob-results-root "${RJOB_RESULTS_ROOT}" + + echo "Generated rjob config: ${GENERATED_DIR}" + echo " config: ${GENERATED_DIR}/patcheval_config.rjob.yaml" + echo " start : ${GENERATED_DIR}/patcheval_start.rjob.yaml" + + launcher_storage_args=(--storage-type "${PATCH_EVAL_STORAGE_TYPE}") + if [[ "${PATCH_EVAL_STORAGE_TYPE}" == "sqlite" ]]; then + launcher_storage_args+=(--db-path "sqlite:///${setting_db}") + fi + + "${PYTHON_BIN}" launcher.py \ + --mode rjob \ + --rjob-config "${RJOB_CONFIG}" \ + --agent-root "${GENERATED_DIR}" \ + --agent-config "${GENERATED_DIR}/patcheval_config.rjob.yaml" \ + --agent-start-config "${GENERATED_DIR}/patcheval_start.rjob.yaml" \ + --gateway-base-url "${GATEWAY_BASE_URL}" \ + --llm-model "${SGLANG_MODEL}" \ + --llm-temperature 0 \ + --agent-start-timeout-s "${PATCH_EVAL_AGENT_TIMEOUT_S}" \ + --shutdown-timeout-s "${PATCH_EVAL_SHUTDOWN_TIMEOUT_S}" \ + "${launcher_storage_args[@]}" \ + --pool-size "${PATCH_EVAL_POOL_SIZE}" \ + --max-workers "${PATCH_EVAL_POOL_SIZE}" \ + --max-steps 1 \ + --enable-evaluation \ + --no-circuit-breaker \ + ${PATCHEVAL_RESUME:+--resume} + + echo "Setting ${setting} done. Results DB: ${setting_db}" +done + +echo "" +echo "=== All settings complete ===" +echo "Gateway log: ${GATEWAY_LOG}" diff --git a/rl/gateway_autostart.py b/rl/gateway_autostart.py index 5a262e07..2b8c67c9 100644 --- a/rl/gateway_autostart.py +++ b/rl/gateway_autostart.py @@ -47,13 +47,23 @@ def build_gateway_config(*, aievobox_root: str) -> Dict[str, Any]: or os.environ.get("AIEVOBOX_POOL_SIZE") or 256 ) + # Per-session LLM step budget enforced by the gateway. -1 = unlimited + # (default: trust the agent / llm_proxy to bound rollout length). Set + # AIEVOBOX_GATEWAY_MAX_STEPS to a non-negative integer to hard-cap runaway + # rollouts; the gateway will inject a synthetic `max_steps_reached` stop + # once a session reaches that many LLM calls. + max_steps = int(os.environ.get("AIEVOBOX_GATEWAY_MAX_STEPS", "-1")) + if max_steps < -1: + raise ValueError("AIEVOBOX_GATEWAY_MAX_STEPS must be -1 or a non-negative integer") return { "listen_host": "0.0.0.0", "listen_port": port, "base_session_path": "/v1/sessions", # -1: never enforce a per-session step budget / inject synthetic stops, # so gateway does not truncate RL generations. llm_proxy owns rollout length. - "max_steps": -1, + # Set AIEVOBOX_GATEWAY_MAX_STEPS>=0 to hard-cap rollout steps as a fallback + # when the agent / llm_proxy fails to terminate on its own. + "max_steps": max_steps, "storage_type": storage_type, # Must match launcher --db-path (= AIEVOBOX_DB_URL) or launcher /readyz fails. "storage_config": {"db_url": db_url}, diff --git a/rl/llm_proxy.py b/rl/llm_proxy.py index aa6a19db..a1b0183a 100644 --- a/rl/llm_proxy.py +++ b/rl/llm_proxy.py @@ -13,6 +13,8 @@ import asyncio from concurrent.futures import ThreadPoolExecutor +import copy +import json import logging import os import sys @@ -71,9 +73,12 @@ sys.path.insert(0, MASK_DIR) from trajectory_mask_builder import PreparedPrompt, TrajectoryMaskBuilder +from chat_template_adapter import create_adapter app = FastAPI(title="LLM Proxy Server", debug=True) + + def _resolve_proxy_workers() -> int: default_workers = min(32, max(8, os.cpu_count() or 8)) raw = os.getenv("AIEVOBOX_LLM_PROXY_WORKERS") @@ -108,6 +113,7 @@ def __init__(self): self.tokenizer = None self.processor: Optional[Any] = None self.trajectory_mask_builder: Optional[TrajectoryMaskBuilder] = None + self.chat_template_adapter = None # type: ignore[type-arg] self.remote_engine_url: Optional[str] = None # Base URL without /v1 self._http_client: Optional[httpx.AsyncClient] = None self._builder_executor: Optional[ThreadPoolExecutor] = None @@ -183,6 +189,17 @@ async def proxy_chat_completions(request: Request): raise HTTPException(status_code=400, detail=f"Invalid JSON body: {e}") messages = payload.get("messages", []) + # OpenHands sends OpenAI-style `tools` so the model sees the real tool + # definitions (terminal/file_editor/...). Without passing them through, + # the chat template never renders the system block and the model + # hallucinates tool names (e.g. `bash` instead of `terminal`), so OpenHands + # rejects every call ("Tool 'Bash' not found") and the agent never produces + # a patch. Forward them to the mask builder so the prompt (and thus the + # recorded trajectory) includes the tools system block. + tools = payload.get("tools") + # Normalize messages via the chat template adapter (e.g. Qwen needs + # tool_calls.arguments as dict and content as string). + messages = STATE.chat_template_adapter.normalize_messages(messages) # Get sampling params from payload or use defaults temperature = payload.get("temperature", STATE.temperature) @@ -197,7 +214,8 @@ async def proxy_chat_completions(request: Request): builder_executor, STATE.trajectory_mask_builder.prepare_generate_input, session_id, - messages + messages, + tools, ) input_ids = prep.input_ids image_data = prep.image_data @@ -250,12 +268,23 @@ async def proxy_chat_completions(request: Request): http_client = STATE.get_http_client() url = f"{STATE.remote_engine_url}/generate" + # Session affinity: when the SGLang router runs the `consistent_hashing` + # policy, it pins all turns of a session (keyed by this header) to one + # worker so the worker's RadixAttention prefix tree keeps reusing the + # session's growing history. Without it the default cache_aware policy + # scatters turns across workers and ~78% of prefills recompute the full + # prompt from scratch. Harmless when the router uses a non-hashing policy + # (unknown header is ignored). + gen_headers: dict[str, str] = {"Content-Type": "application/json"} + if session_id: + gen_headers["X-SMG-Routing-Key"] = session_id + try: logger.debug(f"Calling /generate: input_ids length={len(input_ids)}, max_new_tokens={max_new_tokens}") resp = await http_client.post( url, json=generate_payload, - headers={"Content-Type": "application/json"} + headers=gen_headers ) resp.raise_for_status() resp_json = resp.json() @@ -279,9 +308,25 @@ async def proxy_chat_completions(request: Request): # Get assistant_text from generate API response (already decoded) assistant_text = resp_json.get("text", "") - # Save trajectory + # Convert model text-format tool calls into OpenAI `tool_calls` so + # OpenHands executes them (e.g. Qwen emits as text). + msg_content, tool_calls, tool_finish = STATE.chat_template_adapter.parse_tool_calls(assistant_text) + if tool_calls: + message_obj = {"role": "assistant", "content": msg_content, "tool_calls": tool_calls} + resp_finish_reason = tool_finish + else: + message_obj = {"role": "assistant", "content": assistant_text} + resp_finish_reason = finish_reason + + # Save trajectory. Pass a NORMALIZED copy of message_obj so the trie + # stores the same message format as the DB after normalization. The raw + # assistant_text is still used for token/mask computation. We normalize a + # COPY so the response sent to OpenHands keeps the original format. if STATE.trajectory_mask_builder is not None: try: + trie_msg = STATE.chat_template_adapter.normalize_messages( + [copy.deepcopy(message_obj)] + )[0] await loop.run_in_executor( builder_executor, STATE.trajectory_mask_builder.record_generation, @@ -290,6 +335,7 @@ async def proxy_chat_completions(request: Request): output_logprobs, assistant_text, finish_reason, + trie_msg, ) except Exception as e: import traceback @@ -303,11 +349,8 @@ async def proxy_chat_completions(request: Request): "model": "proxy", "choices": [{ "index": 0, - "message": { - "role": "assistant", - "content": assistant_text - }, - "finish_reason": finish_reason + "message": message_obj, + "finish_reason": resp_finish_reason }], "usage": { "prompt_tokens": len(input_ids), diff --git a/rl/mask/chat_template_adapter.py b/rl/mask/chat_template_adapter.py new file mode 100644 index 00000000..0ae29691 --- /dev/null +++ b/rl/mask/chat_template_adapter.py @@ -0,0 +1,182 @@ +"""Chat template adapter abstraction. + +Different models (Qwen, Llama, DeepSeek, ...) have different chat template +quirks that the RL pipeline must accommodate: + - Message normalization (OpenAI format → model template format) + - Tool call parsing (model text output → OpenAI structured tool_calls) + - Message delta rendering (for training mask alignment) + +This module provides the base ``ChatTemplateAdapter`` interface and a factory +``create_adapter`` that selects the right adapter by name. New models only need +to add a subclass and register it in the factory. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Shared base chat history used for delta rendering. Most adapters use this +# to compute a stable prefix that can be stripped when rendering a single +# message's template fragment. +BASE_CHAT_HISTORY = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "I am a user."}, +] + + +class ChatTemplateAdapter: + """Base chat template adapter. + + Subclasses override the methods that need model-specific logic. + The defaults work for models whose ``apply_chat_template`` accepts + standard OpenAI-format messages without extra guards. + """ + + def __init__(self, tokenizer: Any, processor: Any = None) -> None: + self.tokenizer = tokenizer + self.processor = processor + self.base_messages_str: str = self.tokenizer.apply_chat_template( + BASE_CHAT_HISTORY, + add_generation_prompt=False, + tokenize=False, + ) + + # -- message normalization ------------------------------------------------ + + def normalize_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Normalize OpenAI-format messages so the model's chat template can + render them. Default: no-op (standard models accept OpenAI format as-is). + """ + return messages + + # -- tool call parsing ----------------------------------------------------- + + def parse_tool_calls( + self, assistant_text: str + ) -> Tuple[Optional[str], Optional[List[Dict[str, Any]]], Optional[str]]: + """Parse the model's raw text output into OpenAI ``tool_calls``. + + Returns ``(content, tool_calls, finish_reason)``: + - content: reasoning text with tool-call markup removed (None if empty) + - tool_calls: list of OpenAI tool_call dicts, or None if none found + - finish_reason: "tool_calls" if any, else None (caller decides) + + Default: no parsing — the model already returns structured tool_calls + via the API, so the raw text is the content. + """ + return assistant_text, None, None + + # -- message delta rendering (for training mask) -------------------------- + + def render_message_delta(self, message: Dict[str, Any]) -> str: + """Render a single message's template fragment for training mask. + + Default: render ``[BASE_CHAT_HISTORY, message]`` and strip the + ``BASE_CHAT_HISTORY`` prefix. This works for models whose template + has no restrictions on system message position or user message presence. + """ + full = self.tokenizer.apply_chat_template( + BASE_CHAT_HISTORY + [message], + add_generation_prompt=False, + tokenize=False, + ) + if not full.startswith(self.base_messages_str): + raise ValueError("failed to extract single-message template fragment") + return full[len(self.base_messages_str):] + + def render_first_system_delta( + self, + message: Dict[str, Any], + tools: Optional[List[Dict[str, Any]]] = None, + ) -> str: + """Render the session's first system message, optionally with tools. + + Default: same as ``render_message_delta`` (most models don't need + special first-system handling). Models that only inject ```` + into the first system message (e.g. Qwen) override this. + """ + return self.render_message_delta(message) + + def needs_tools_on_first_system_only(self) -> bool: + """Whether tools should only be rendered on the first system message. + + Default: False. Qwen overrides to True because its template only + injects the ```` block into the first system message. + """ + return False + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + +_REGISTRY: Dict[str, type[ChatTemplateAdapter]] = {} + + +def register_adapter(name: str, cls: type[ChatTemplateAdapter]) -> None: + """Register a chat template adapter class under ``name``.""" + _REGISTRY[name] = cls + + +def create_adapter( + name_or_type: Optional[str], + tokenizer: Any, + processor: Any = None, +) -> ChatTemplateAdapter: + """Create a chat template adapter by name. + + ``name_or_type`` is matched case-insensitively against registered adapters. + Falls back to the base ``ChatTemplateAdapter`` (no-op) if not found, so + unknown models still work for standard cases. + """ + if name_or_type is None: + name_or_type = "" + key = name_or_type.strip().lower() + + # Map common aliases + alias_map = { + "qwen": "qwen", + "qwen3": "qwen", + "qwen3_5": "qwen", + "qwen3.5": "qwen", + "qwen3_8": "qwen", + "qwen3.8": "qwen", + "qwen3_6": "qwen", + "qwen3.6": "qwen", + } + key = alias_map.get(key, key) + + cls = _REGISTRY.get(key) + if cls is not None: + logger.info("Chat template adapter: %s -> %s", name_or_type, cls.__name__) + return cls(tokenizer, processor) + + logger.warning( + "Unknown chat template adapter %r, falling back to base (no-op). " + "Register it via register_adapter() if the model needs special handling.", + name_or_type, + ) + return ChatTemplateAdapter(tokenizer, processor) + + +# --------------------------------------------------------------------------- +# Auto-register built-in adapters +# --------------------------------------------------------------------------- + +def _autoregister() -> None: + """Import and register built-in adapters. Called once at module load.""" + try: + # Try relative import first (when used as a package, e.g. rl.mask.chat_template_adapter) + from .qwen_chat_template_adapter import QwenChatTemplateAdapter + except ImportError: + # Fall back to absolute import (when rl/mask/ is on sys.path and this + # module is imported as a top-level module, e.g. `from chat_template_adapter import ...`) + from qwen_chat_template_adapter import QwenChatTemplateAdapter + register_adapter("qwen", QwenChatTemplateAdapter) + + +_autoregister() diff --git a/rl/mask/diag_template.py b/rl/mask/diag_template.py new file mode 100644 index 00000000..4b7b47e8 --- /dev/null +++ b/rl/mask/diag_template.py @@ -0,0 +1,48 @@ +"""Diagnose the system-message template fragment extraction failure.""" +import sys +from transformers import AutoTokenizer + +MODEL = sys.argv[1] if len(sys.argv) > 1 else \ + "/mnt/shared-storage-gpfs2/gpfs2-shared-public/huggingface/hub/models--Qwen--Qwen3.8-27B/snapshots/1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + +print(f"Loading tokenizer from: {MODEL}") +tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + +_USER_ONLY_BASE = [{"role": "user", "content": "I am a user."}] +sys_msg = {"role": "system", "content": "You are OpenHands agent, a helpful AI assistant that can interact with a computer to solve tasks."} + +base_str = tok.apply_chat_template(_USER_ONLY_BASE, add_generation_prompt=False, tokenize=False) +with_msg = tok.apply_chat_template([sys_msg] + _USER_ONLY_BASE, add_generation_prompt=False, tokenize=False) + +print("\n===== base_str (render of [user_base]) =====") +print(repr(base_str)) +print("\n===== with_msg (render of [sys_msg, user_base]) =====") +print(repr(with_msg)) +print("\n===== with_msg.endswith(base_str) ? =====") +print(with_msg.endswith(base_str)) + +if not with_msg.endswith(base_str): + i = 0 + while i < len(base_str) and i < len(with_msg) and with_msg[-(i+1)] == base_str[-(i+1)]: + i += 1 + print(f"\nSuffix match length from end: {i}") + print(f"with_msg tail (last 150): {with_msg[-150:]!r}") + print(f"base_str tail (last 150): {base_str[-150:]!r}") + +print("\n===== with enable_thinking=False =====") +try: + b2 = tok.apply_chat_template(_USER_ONLY_BASE, add_generation_prompt=False, tokenize=False, enable_thinking=False) + w2 = tok.apply_chat_template([sys_msg]+_USER_ONLY_BASE, add_generation_prompt=False, tokenize=False, enable_thinking=False) + print(f"endswith base? {w2.endswith(b2)}") + print(f"with_msg2: {w2!r}") +except Exception as e: + print(f"enable_thinking=False error: {e!r}") + +print("\n===== with system content as LIST =====") +sys_msg_list = {"role": "system", "content": [{"type": "text", "text": "You are OpenHands agent."}]} +try: + w3 = tok.apply_chat_template([sys_msg_list]+_USER_ONLY_BASE, add_generation_prompt=False, tokenize=False) + print(f"endswith base? {w3.endswith(base_str)}") + print(f"with_msg3: {w3!r}") +except Exception as e: + print(f"list-content error: {e!r}") diff --git a/rl/mask/qwen_chat_template_adapter.py b/rl/mask/qwen_chat_template_adapter.py new file mode 100644 index 00000000..87405369 --- /dev/null +++ b/rl/mask/qwen_chat_template_adapter.py @@ -0,0 +1,208 @@ +"""Qwen3.5/3.6/3.8 chat template adapter. + +Handles three Qwen-specific quirks that the base adapter cannot: + +1. **Message normalization**: Qwen's chat template iterates ``tool_call.arguments`` + via the Jinja ``items`` filter, so arguments must be a dict (OpenAI sends a + JSON string). The template also accesses ``content`` directly (not ``.get``), + so None content raises ``KeyError`` and non-string content raises + ``AttributeError`` on ``.startswith``. + +2. **Tool call parsing**: Qwen emits tool calls as *text* using + ``...VALUE...`` + wrapped in delimiter tokens. OpenHands only executes structured + ``tool_calls``, so we parse the text format back into OpenAI dicts. + +3. **System message rendering**: Qwen's template has two hard checks: + - system message must be at index 0 + - a user message must exist + Rendering a standalone system message triggers the second check. We work + around it by rendering ``[system, user_base]`` and stripping ``user_base``. + +4. **Tools block injection**: Qwen only injects ``...`` into + the FIRST system message. We render the first system message WITH tools + so the training mask aligns with what sglang rendered. +""" + +from __future__ import annotations + +import json +import re +from typing import Any, Dict, List, Optional, Tuple + +try: + from .chat_template_adapter import BASE_CHAT_HISTORY, ChatTemplateAdapter +except ImportError: + from chat_template_adapter import BASE_CHAT_HISTORY, ChatTemplateAdapter + +# User-only base used for the system-message rendering trick. +_USER_ONLY_BASE = [{"role": "user", "content": "I am a user."}] + +# Regexes for parsing Qwen text-format tool calls. +_FUNCTION_BLOCK_RE = re.compile(r"(.*?)", re.DOTALL) +_PARAM_BLOCK_RE = re.compile(r"(.*?)", re.DOTALL) +# Trailing tool-call delimiter token (a `<...>` tag) right before the first +# `]*>\s*$") + + +class QwenChatTemplateAdapter(ChatTemplateAdapter): + """Chat template adapter for Qwen3.5 / 3.6 / 3.8 models.""" + + def __init__(self, tokenizer: Any, processor: Any = None) -> None: + super().__init__(tokenizer, processor) + self._user_suffix_str: Optional[str] = None + + # -- message normalization ------------------------------------------------ + + def normalize_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Normalize OpenAI-format messages so the Qwen chat template can + render them. + + - ``tool_calls.arguments``: JSON string → dict (template uses ``items``) + - ``content``: None → "" (template accesses content directly) + - ``content``: non-string → string (template calls ``.startswith``) + """ + for msg in messages or []: + if not isinstance(msg, dict): + continue + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list): + for tc in tool_calls: + if not isinstance(tc, dict): + continue + fn = tc.get("function") + if not isinstance(fn, dict): + continue + args = fn.get("arguments") + if isinstance(args, str): + try: + fn["arguments"] = json.loads(args) if args else {} + except (json.JSONDecodeError, ValueError): + fn["arguments"] = {"_raw": args} + content = msg.get("content") + if content is None: + msg["content"] = "" + elif not isinstance(content, str): + try: + msg["content"] = json.dumps(content, ensure_ascii=False) + except Exception: + msg["content"] = str(content) + return messages + + # -- tool call parsing ----------------------------------------------------- + + def parse_tool_calls( + self, assistant_text: str + ) -> Tuple[Optional[str], Optional[List[Dict[str, Any]]], Optional[str]]: + """Convert Qwen text-format tool calls into OpenAI ``tool_calls``. + + Qwen emits ``...VALUE... + `` as text. OpenHands only executes structured + ``tool_calls``, so we parse the text and convert. The raw + ``assistant_text`` is still what gets recorded into the training + trajectory — this conversion only shapes the response handed back + to the agent. + """ + blocks = list(_FUNCTION_BLOCK_RE.finditer(assistant_text)) + if not blocks: + return assistant_text, None, None + + tool_calls: List[Dict[str, Any]] = [] + for idx, blk in enumerate(blocks): + name = blk.group(1) + args: Dict[str, Any] = {} + for p in _PARAM_BLOCK_RE.finditer(blk.group(2)): + args[p.group(1)] = p.group(2).strip("\n") + tool_calls.append({ + "id": f"call_{idx}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + }) + + # Content = text before the first tool-call block, with the trailing + # tool-call delimiter tag stripped. Text between/after blocks is just + # delimiter noise, discard it. + content = assistant_text[: blocks[0].start()] + content = _TRAILING_TAG_RE.sub("", content).strip() + return (content or None), tool_calls, "tool_calls" + + # -- message delta rendering (for training mask) -------------------------- + + def _get_user_suffix_str(self) -> str: + """The rendered form of a single user message as it appears AFTER a + system message. Used to strip the trailing user message when rendering + a standalone system message (Qwen's template guards require a user + message to be present, so we render [system, user] and strip the user + suffix). + + NB: render([user]) alone injects a synthetic default system block + (with reasoning instructions) before the user, so it is NOT a clean + suffix — compute it from BASE_CHAT_HISTORY + [user] instead. + """ + if self._user_suffix_str is None: + full = self.tokenizer.apply_chat_template( + BASE_CHAT_HISTORY + [{"role": "user", "content": "I am a user."}], + add_generation_prompt=False, + tokenize=False, + ) + if not full.startswith(self.base_messages_str): + raise ValueError("failed to extract user-suffix template fragment") + self._user_suffix_str = full[len(self.base_messages_str):] + return self._user_suffix_str + + def render_message_delta(self, message: Dict[str, Any]) -> str: + """Render a single message's template fragment. + + Qwen3.5/3.6 chat template has two hard checks: + 1) system message must be at index 0 + 2) a user message must exist + + For system messages, rendering [system] alone triggers check (2). + We render [system, user_base] and strip user_base to get the clean + system fragment. For non-system messages, the default approach + (render [BASE, msg] and strip BASE prefix) works fine. + """ + if message.get("role") == "system": + user_suffix = self._get_user_suffix_str() + with_msg = self.tokenizer.apply_chat_template( + [message] + _USER_ONLY_BASE, + add_generation_prompt=False, + tokenize=False, + ) + if not with_msg.endswith(user_suffix): + raise ValueError("failed to extract system-message template fragment") + return with_msg[: len(with_msg) - len(user_suffix)] + + # Non-system: use default base-prefix-strip approach. + return super().render_message_delta(message) + + def render_first_system_delta( + self, + message: Dict[str, Any], + tools: Optional[List[Dict[str, Any]]] = None, + ) -> str: + """Render the session's first system message WITH tools so the + template's ``...`` system block lands in the recorded + input_ids, matching what sglang renders for the rollout prompt. + + Qwen's template guards require a user message, so we render + [system_msg, user_base] with tools and strip the clean user suffix. + """ + user_suffix = self._get_user_suffix_str() + with_msg = self.tokenizer.apply_chat_template( + [message] + _USER_ONLY_BASE, + tools=tools, + add_generation_prompt=False, + tokenize=False, + ) + if not with_msg.endswith(user_suffix): + raise ValueError("failed to extract first-system-message template fragment") + return with_msg[: len(with_msg) - len(user_suffix)] + + def needs_tools_on_first_system_only(self) -> bool: + """Qwen only injects ```` into the first system message.""" + return True diff --git a/rl/mask/trajectory_mask_builder.py b/rl/mask/trajectory_mask_builder.py index 02ecff69..f770b22b 100644 --- a/rl/mask/trajectory_mask_builder.py +++ b/rl/mask/trajectory_mask_builder.py @@ -1,4 +1,5 @@ import logging +import os import re from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple @@ -6,15 +7,12 @@ from qwen_vl_utils import process_vision_info from slime.utils.processing_utils import encode_image_for_rollout_engine +from chat_template_adapter import BASE_CHAT_HISTORY, ChatTemplateAdapter, create_adapter + logger = logging.getLogger(__name__) THINK_BLOCK_RE = re.compile(r"\s*.*?\s*", re.DOTALL) -BASE_CHAT_HISTORY = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "I am a user."}, -] - @dataclass class MessageNode: @@ -39,15 +37,14 @@ class PreparedPrompt: class TrajectoryMaskBuilder: - def __init__(self, tokenizer, processor: Any = None) -> None: + def __init__(self, tokenizer, processor: Any = None, adapter: Optional[ChatTemplateAdapter] = None) -> None: self.tokenizer = tokenizer self.processor = processor self.session_roots: Dict[str, MessageNode] = {} - self.base_messages_str = self.tokenizer.apply_chat_template( - BASE_CHAT_HISTORY, - add_generation_prompt=False, - tokenize=False, + self.adapter = adapter or create_adapter( + os.environ.get("LOSS_MASK_TYPE", ""), tokenizer, processor ) + self.base_messages_str = self.adapter.base_messages_str self.generation_tokens = self._init_generation_tokens() self.suffix = self._init_suffix_tokens() @@ -75,6 +72,16 @@ def _init_suffix_tokens(self) -> List[int]: add_generation_prompt=False, tokenize=True, ) + # Some tokenizer versions return BatchEncoding (or a batched tensor) + # here instead of a flat list of token IDs. + if hasattr(test_tokens, "input_ids"): + test_tokens = test_tokens.input_ids + elif isinstance(test_tokens, dict): + test_tokens = test_tokens["input_ids"] + if hasattr(test_tokens, "tolist"): + test_tokens = test_tokens.tolist() + if test_tokens and isinstance(test_tokens[0], (list, tuple)): + test_tokens = test_tokens[0] for idx in range(len(test_tokens) - 1, -1, -1): if test_tokens[idx] == eos_id: return list(test_tokens[idx + 1 :]) @@ -214,14 +221,16 @@ def _build_mm_inputs( return list(input_ids), mm_train_inputs def _render_message_delta_str(self, model_input_message: Dict[str, Any]) -> str: - single_message_chat_template_str = self.tokenizer.apply_chat_template( - BASE_CHAT_HISTORY + [model_input_message], - add_generation_prompt=False, - tokenize=False, - ) - if not single_message_chat_template_str.startswith(self.base_messages_str): - raise ValueError("failed to extract single-message template fragment") - return single_message_chat_template_str[len(self.base_messages_str) :] + """Render a single message's template fragment via the adapter.""" + return self.adapter.render_message_delta(model_input_message) + + def _render_first_system_delta_str( + self, + model_input_message: Dict[str, Any], + tools: List[Dict[str, Any]], + ) -> str: + """Render the first system message WITH tools via the adapter.""" + return self.adapter.render_first_system_delta(model_input_message, tools) def _build_mm_train_inputs_for_images(self, images: List[Any]) -> Optional[Dict[str, Any]]: if self.processor is None or not images: @@ -352,6 +361,7 @@ def _add_prompt_message( tokens: List[int], images: List[Any], image_data: List[str], + tools: Optional[List[Dict[str, Any]]] = None, ) -> Tuple[MessageNode, List[Dict[str, Any]], str, List[int], List[Any], List[str]]: model_input_message = self._message_for_model_input(raw_message) next_model_input_messages = list(model_input_messages) @@ -366,7 +376,10 @@ def _add_prompt_message( next_image_data.extend(new_image_data) delta_mm_train_inputs = self._build_mm_train_inputs_for_images(new_images) - delta_message_str = self._render_message_delta_str(model_input_message) + if tools is not None: + delta_message_str = self._render_first_system_delta_str(model_input_message, tools) + else: + delta_message_str = self._render_message_delta_str(model_input_message) next_messages_str = messages_str + delta_message_str delta_tokens, _ = self._build_mm_inputs(delta_message_str, new_images) delta_tokens = list(delta_tokens) @@ -400,16 +413,24 @@ def _append_assistant_message( output_ids: List[int], assistant_text: str, finish_reason: Optional[str], + assistant_message: Optional[Dict[str, Any]] = None, ) -> MessageNode: del finish_reason - assistant_message = {"role": "assistant", "content": assistant_text} - model_input_message = self._message_for_model_input(assistant_message) + # If caller provides a pre-parsed assistant_message (OpenAI format with + # tool_calls, list content), use it as raw_message so _message_matches + # can compare it against DB messages during get_training_info. The + # raw assistant_text is still used for token/mask computation. + if assistant_message is not None: + raw_message = assistant_message + else: + raw_message = {"role": "assistant", "content": assistant_text} + model_input_message = self._message_for_model_input(raw_message) delta_message_str = self._render_message_delta_str(model_input_message) delta_tokens = list(self.generation_tokens) + list(output_ids) + list(self.suffix) delta_response_mask = [0] * len(self.generation_tokens) + [1] * len(output_ids) + [0] * len(self.suffix) node = MessageNode( - raw_message=assistant_message, + raw_message=raw_message, model_input_message=model_input_message, delta_message_str=delta_message_str, delta_tokens=delta_tokens, @@ -425,12 +446,20 @@ def _ensure_path( self, session_id: str, messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, ) -> Tuple[MessageNode, List[Dict[str, Any]], str, List[int], List[Any], List[str]]: node, matched, model_input_messages, messages_str, tokens, _response_mask, images, image_data, _mm_train_inputs = self._match_prefix( session_id, messages, ) - for message in messages[matched:]: + # Some models (e.g. Qwen) inject a `...` system block + # only into the FIRST system message of the rendered prompt. To get it + # into the recorded input_ids (so the training mask aligns with what + # sglang actually rendered), render the session's first system message + # standalone WITH tools; every other message uses the normal delta. + first_tools = tools if (matched == 0 and not node.children and self.adapter.needs_tools_on_first_system_only()) else None + for idx, message in enumerate(messages[matched:]): + msg_tools = first_tools if (idx == 0 and first_tools is not None and message.get("role") == "system") else None node, model_input_messages, messages_str, tokens, images, image_data = self._add_prompt_message( node, message, @@ -439,6 +468,7 @@ def _ensure_path( tokens, images, image_data, + tools=msg_tools, ) return node, model_input_messages, messages_str, tokens, images, image_data @@ -446,10 +476,12 @@ def prepare_generate_input( self, session_id: str, messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, ) -> PreparedPrompt: node, model_input_messages, messages_str, tokens, _images, image_data = self._ensure_path( session_id, messages, + tools=tools, ) input_ids = list(tokens) input_ids.extend(self.generation_tokens) @@ -468,6 +500,7 @@ def record_generation( output_logprobs: List[List[Any]], assistant_text: str, finish_reason: Optional[str] = None, + assistant_message: Optional[Dict[str, Any]] = None, ) -> MessageNode: del output_logprobs return self._append_assistant_message( @@ -475,6 +508,7 @@ def record_generation( output_ids=list(output_ids), assistant_text=assistant_text, finish_reason=finish_reason, + assistant_message=assistant_message, ) def get_training_info( diff --git a/rl/patches/gdn_packed_seq.py b/rl/patches/gdn_packed_seq.py new file mode 100644 index 00000000..e2385855 --- /dev/null +++ b/rl/patches/gdn_packed_seq.py @@ -0,0 +1,168 @@ +"""Runtime monkey-patch: make Megatron GDN support packed sequences (thd). + +Megatron's `GatedDeltaNet.forward` raises `NotImplementedError` when +`packed_seq_params` is provided. However, the core recurrence op +`chunk_gated_delta_rule` (from fla) already accepts a `cu_seqlens` argument +for variable-length packed sequences — slime's own `qwen3_5.py` uses exactly +this. This patch removes the raise and forwards `cu_seqlens` through, so +slime can use `--qkv-format thd` (packing) without OOM and without hitting +the GDN limitation. + +Loaded automatically when this directory is on PYTHONPATH (set in +env.rjob.sh). No Megatron source modification, no image rebuild. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from megatron.core.ssm.gated_delta_net import GatedDeltaNet +from megatron.core.ssm.gated_delta_net import ( + chunk_gated_delta_rule, + torch_chunk_gated_delta_rule, + causal_conv1d_fn, + l2norm, +) +from megatron.core.utils import deprecate_inference_params, nvtx_range_push, nvtx_range_pop + +_original_forward = GatedDeltaNet.forward + + +def _patched_forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + sequence_len_offset=None, + *, + inference_params=None, +): + inference_context = deprecate_inference_params(inference_context, inference_params) + + seq_len, batch, _ = hidden_states.shape + seq_len = seq_len * self.sp_size + + if inference_context is not None: + raise NotImplementedError("GDN does not support inference for now.") + + # --- packed sequence support: extract cu_seqlens --- + cu_seqlens = None + if packed_seq_params is not None: + cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None) + + # Input projection + nvtx_range_push(suffix="in_proj") + qkvzba, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + + qkvzba = qkvzba.transpose(0, 1) # sbhd -> bshd + + qkv, gate, beta, alpha = torch.split( + qkvzba, + [ + (self.qk_dim * 2 + self.v_dim) // self.tp_size, + self.v_dim // self.tp_size, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + dim=-1, + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + beta = beta.reshape(batch, seq_len, -1) + alpha = alpha.reshape(batch, seq_len, -1) + + # Convolution on qkv + qkv = qkv.transpose(1, 2).contiguous() # b,s,d -> b,d,s + nvtx_range_push(suffix="conv1d") + if (causal_conv1d_fn is None) or self.config.deterministic_mode: + qkv = self.act_fn(self.conv1d(qkv)[..., :seq_len]) + else: + # causal_conv1d_fn supports seq_idx for varlen; convert cu_seqlens + # to seq_idx [batch, seq_len] when available. + seq_idx = None + if cu_seqlens is not None: + seq_idx = torch.zeros(batch, seq_len, dtype=torch.int32, device=qkv.device) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item() + if end > start: + seq_idx[:, start:end] = i + qkv = causal_conv1d_fn( + x=qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + seq_idx=seq_idx, + ) + nvtx_range_pop(suffix="conv1d") + + qkv = qkv.transpose(1, 2) # b,d,s -> b,s,d + query, key, value = torch.split( + qkv, + [self.qk_dim // self.tp_size, self.qk_dim // self.tp_size, self.v_dim // self.tp_size], + dim=-1, + ) + query = query.reshape(batch, seq_len, -1, self.key_head_dim) + key = key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + + if self.use_qk_l2norm: + query = l2norm(query.contiguous()) + key = l2norm(key.contiguous()) + if self.num_value_heads // self.num_key_heads > 1: + query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + nvtx_range_push(suffix="g_and_beta") + g = -self.A_log.exp() * F.softplus(alpha.float() + self.dt_bias) + beta = beta.sigmoid() + nvtx_range_pop(suffix="g_and_beta") + + nvtx_range_push(suffix="gated_delta_rule") + if self.config.deterministic_mode: + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( + query, key, value, + g=g, beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + ) + else: + core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + query, key, value, + g=g, beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + return out, out_bias + + +GatedDeltaNet.forward = _patched_forward diff --git a/rl/patches/sitecustomize.py b/rl/patches/sitecustomize.py new file mode 100644 index 00000000..9e3ad25f --- /dev/null +++ b/rl/patches/sitecustomize.py @@ -0,0 +1,30 @@ +"""Auto-load GDN packed-seq monkey-patch at Python startup. + +Python imports `sitecustomize` automatically during interpreter startup +from any directory on sys.path. This file lives in rl/patches/ which +is added to PYTHONPATH by env.rjob.sh, so the patch is applied before +slime/train.py runs — no code change to slime or Megatron needed. +""" +# Import any pre-existing sitecustomize first (chained sitecustomize). +try: + import _orig_sitecustomize # noqa: F401 +except Exception: + pass + +try: + import gdn_packed_seq # noqa: F401 — applies the monkey-patch +except Exception as _e: + import sys + print(f"[sitecustomize] WARNING: gdn_packed_seq failed to load: {_e}", file=sys.stderr) + +try: + import spread_placement # noqa: F401 — SPREAD strategy for multi-node placement +except Exception as _e: + import sys + print(f"[sitecustomize] WARNING: spread_placement failed to load: {_e}", file=sys.stderr) + +# REMOVED patches (2026-09-08, non-colocate + raw + Megatron-ckpt config): +# - traj_truncation: disabled via TRAJ_TRUNCATION_MAX_SEQ_LEN=0 (PP=2 显存充裕) +# - raw_hf_checkpoint: now loading Megatron-format checkpoints, not HF +# - flush_cache_fix: only needed in colocate mode; non-colocate has no flush issue +# - attention_mask_fix: only needed in bridge mode; using raw mode now diff --git a/rl/patches/spread_placement.py b/rl/patches/spread_placement.py new file mode 100644 index 00000000..a2ccc813 --- /dev/null +++ b/rl/patches/spread_placement.py @@ -0,0 +1,90 @@ +"""Monkey-patch: change placement group strategy from PACK to SPREAD. + +slime uses PACK strategy by default, which tries to pack all bundles on +one node. With 16 bundles (8 actor + 8 rollout) and 2 nodes (8 GPUs each), +PACK can result in all 16 bundles on one node with 2 bundles per GPU, +causing "Duplicate GPU detected" NCCL errors. + +SPREAD strategy distributes bundles evenly across all available nodes, +ensuring each rank gets a unique GPU. This is required for PP>1 where +all 8 actor ranks must be on distinct GPUs. + +NOTE: SPREAD is only needed in COLOCATE mode, where training and rollout +share the same placement group and need to be spread across nodes to avoid +GPU conflicts. In NON-COLOCATE mode, PACK is better: it naturally separates +training bundles (0-31) and rollout bundles (32-39) onto different nodes, +preventing the SGLang engine and Megatron actor from landing on the same GPU. +""" +import logging + +logger = logging.getLogger(__name__) + +_original_create_placement_group = None + + +def _patched_create_placement_group(num_gpus): + """Replacement that uses SPREAD in colocate, PACK in non-colocate.""" + import os + import ray + from ray.util.placement_group import placement_group, PlacementGroupSchedulingStrategy + + # In non-colocate mode, use PACK so training and rollout bundles + # are separated onto different nodes (training fills nodes 1-4, + # rollout fills node 5). SPREAD would mix them on the same nodes, + # causing SGLang and Megatron to share GPUs → OOM. + colocate = os.environ.get("SLIME_COLOCATE", "false").lower() in ("true", "1") + strategy = "SPREAD" if colocate else "PACK" + logger.info(f"Placement group strategy: {strategy} (colocate={colocate})") + + bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] + pg = placement_group(bundles, strategy=strategy) + num_bundles = len(bundles) + + ray.get(pg.ready()) + + # use info actor to get the GPU id + from slime.ray.placement_group import InfoActor + + info_actors = [] + for i in range(num_bundles): + info_actors.append( + InfoActor.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ) + ).remote() + ) + gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors]) + for actor in info_actors: + ray.kill(actor) + + bundle_infos = [(i, gpu_ids[i][0], gpu_ids[i][1]) for i in range(num_bundles)] + + def sort_key(info): + return (info[1], info[2]) + + sorted_bundle_infos = sorted(bundle_infos, key=sort_key) + pg_reordered_bundle_indices = [info[0] for info in sorted_bundle_infos] + pg_reordered_gpu_ids = [gpu_ids[info[0]][1] for info in sorted_bundle_infos] + + for i in range(num_bundles): + actual_bundle_index = pg_reordered_bundle_indices[i] + logger.info( + f" bundle {i:4}, actual_bundle_index: {actual_bundle_index:4}, " + f"node: {gpu_ids[actual_bundle_index][0]}, gpu: {gpu_ids[actual_bundle_index][1]}" + ) + + return pg, pg_reordered_bundle_indices, pg_reordered_gpu_ids + + +def apply_patch(): + global _original_create_placement_group + from slime.ray import placement_group as slime_pg_module + + _original_create_placement_group = slime_pg_module._create_placement_group + slime_pg_module._create_placement_group = _patched_create_placement_group + logger.info("Patched _create_placement_group: strategy now depends on SLIME_COLOCATE") + + +apply_patch() diff --git a/rl/run_buffer_server.sh b/rl/run_buffer_server.sh index 1e18e536..7af85ae4 100755 --- a/rl/run_buffer_server.sh +++ b/rl/run_buffer_server.sh @@ -36,6 +36,27 @@ elif [[ -z "${AIEVOBOX_ROOT:-}" ]]; then exit 1 fi +is_true() { + case "${1:-}" in + 1|true|TRUE|yes|YES|on|ON) return 0 ;; + *) return 1 ;; + esac +} + +if is_true "${AIEVOBOX_RESET_SQLITE_DB:-false}"; then + case "${AIEVOBOX_DB_URL:-}" in + sqlite:///*) + db_path="${AIEVOBOX_DB_URL#sqlite:///}" + rm -f -- "${db_path}" "${db_path}-wal" "${db_path}-shm" + echo "Removed SQLite DB for fresh rollout: ${db_path}" + ;; + *) + echo "AIEVOBOX_RESET_SQLITE_DB requires a sqlite:/// DB URL" >&2 + exit 1 + ;; + esac +fi + require_dir() { local path="$1" local label="$2" @@ -57,14 +78,20 @@ require_file() { export PYTHONPATH="${AIEVOBOX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" mkdir -p "${LOG_ROOT}" -if [[ -z "${AIEVOBOX_RUN_DIR:-}" && -f "${LOG_ROOT}/.current_run" ]]; then +# Always follow .current_run (the single source of truth, refreshed by +# run_slime_generator.sh on each launch). A stale AIEVOBOX_RUN_DIR exported +# into the shell from a previous run must NOT override the latest run dir, +# otherwise restarting only buffer_server (in the same terminal) writes logs +# into the previous run's directory. The env var is kept only as a fallback +# for the very first launch when no .current_run exists yet. +if [[ -f "${LOG_ROOT}/.current_run" ]]; then export AIEVOBOX_RUN_DIR="$(cat "${LOG_ROOT}/.current_run")" -fi -if [[ -z "${AIEVOBOX_RUN_DIR:-}" ]]; then +elif [[ -z "${AIEVOBOX_RUN_DIR:-}" ]]; then export AIEVOBOX_RUN_DIR="${LOG_ROOT}/$(date +%Y%m%d-%H%M%S)" printf '%s\n' "${AIEVOBOX_RUN_DIR}" > "${LOG_ROOT}/.current_run" fi mkdir -p "${AIEVOBOX_RUN_DIR}" +export SAFACTORY_TIMING_LOG="${AIEVOBOX_RUN_DIR}/timing.jsonl" require_dir "${AIEVOBOX_ROOT}" "AIEVOBOX_ROOT" require_file "${AIEVOBOX_ROOT}/rl/buffer_server.py" "buffer server entrypoint" diff --git a/rl/run_slime_generator.sh b/rl/run_slime_generator.sh index b850a0a3..9d2e4983 100755 --- a/rl/run_slime_generator.sh +++ b/rl/run_slime_generator.sh @@ -100,10 +100,18 @@ if (( RL_GROUP_SIZE <= 0 || RL_ROLLOUT_GROUP_BATCH_SIZE <= 0 || RL_GLOBAL_BATCH_ fi if [[ -z "${AIEVOBOX_RUN_DIR:-}" ]]; then - export AIEVOBOX_RUN_DIR="${LOG_ROOT}/$(date +%Y%m%d-%H%M%S)" + # Reuse the buffer_server's run dir if it already created one (written to + # .current_run by run_buffer_server.sh). This keeps all logs (buffer_server, + # gateway, slime, timing) in the same directory. + if [[ -f "${LOG_ROOT}/.current_run" ]]; then + export AIEVOBOX_RUN_DIR="$(cat "${LOG_ROOT}/.current_run")" + else + export AIEVOBOX_RUN_DIR="${LOG_ROOT}/$(date +%Y%m%d-%H%M%S)" + fi fi mkdir -p "${AIEVOBOX_RUN_DIR}" printf '%s\n' "${AIEVOBOX_RUN_DIR}" > "${LOG_ROOT}/.current_run" +export SAFACTORY_TIMING_LOG="${AIEVOBOX_RUN_DIR}/timing.jsonl" ROLLOUT_BUFFER_URL="http://${BUFFER_SERVER_HOST}:${BUFFER_SERVER_PORT}" LLM_PROXY_URL="http://${LLM_PROXY_HOST}:${LLM_PROXY_PORT}" @@ -112,6 +120,8 @@ export ROLLOUT_BUFFER_URL LLM_PROXY_URL export WANDB_MODE export PYTHONUNBUFFERED export PYTORCH_CUDA_ALLOC_CONF +export PYTORCH_ALLOC_CONF +export TRAJ_TRUNCATION_MAX_SEQ_LEN export MODEL_ARGS_ROTARY_BASE source "${MODEL_SCRIPT}" @@ -151,7 +161,8 @@ ROLLOUT_ARGS=( --num-rollout "${NUM_ROLLOUT}" --rollout-batch-size "${RL_ROLLOUT_GROUP_BATCH_SIZE}" --n-samples-per-prompt "${RL_GROUP_SIZE}" - --rollout-max-response-len "${LLM_MAX_LENGTH}" + --rollout-num-process "${ROLLOUT_NUM_PROCESS:-${RL_GLOBAL_BATCH_SIZE}}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-32768}" --rollout-temperature "${LLM_TEMPERATURE}" --global-batch-size "${RL_GLOBAL_BATCH_SIZE}" --loss-mask-type "${LOSS_MASK_TYPE}" @@ -166,6 +177,7 @@ MEGATRON_ARGS=( --tensor-model-parallel-size "${TP_SIZE}" --pipeline-model-parallel-size "${PP_SIZE}" --context-parallel-size "${CP_SIZE}" + --sequence-parallel --expert-model-parallel-size "${EP_SIZE}" --expert-tensor-parallel-size "${ETP_SIZE}" --recompute-granularity "${RECOMPUTE_GRANULARITY}" @@ -177,6 +189,16 @@ MEGATRON_ARGS=( --attention-softmax-in-fp32 --attention-backend "${ATTENTION_BACKEND}" ) +if [[ -n "${DECODER_LAST_PIPELINE_NUM_LAYERS:-}" ]]; then + MEGATRON_ARGS+=(--decoder-last-pipeline-num-layers "${DECODER_LAST_PIPELINE_NUM_LAYERS}") +fi +# CPU offload optimizer: moves fp32 master weights + Adam states (~81GB at TP=4) +# to CPU, leaving only bf16 weights + bf16 grad (~27GB) on GPU. Critical for +# 27B model on 140GB GPUs where weights+optimizer would otherwise OOM. +# Toggle via OPTIMIZER_CPU_OFFLOAD (default: true for 27B on 8-card TP=4). +if is_true "${OPTIMIZER_CPU_OFFLOAD:-true}"; then + MEGATRON_ARGS+=(--optimizer-cpu-offload --use-precision-aware-optimizer) +fi TRAIN_ARGS=( --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}" @@ -195,8 +217,13 @@ GRPO_ARGS=( --advantage-estimator "${ADVANTAGE_ESTIMATOR}" --entropy-coef "${ENTROPY_COEF}" --eps-clip "${EPS_CLIP}" - --eps-clip-high "${EPS_CLIP_HIGH}" + --kl-loss-coef "${KL_LOSS_COEF:-0.00}" + --kl-loss-type "${KL_LOSS_TYPE:-low_var_kl}" + --kl-coef "${KL_COEF:-0.00}" ) +if [[ -n "${EPS_CLIP_HIGH:-}" ]]; then + GRPO_ARGS+=(--eps-clip-high "${EPS_CLIP_HIGH}") +fi if is_true "${USE_OPD:-false}"; then GRPO_ARGS+=(--use-opd --opd-type "${OPD_TYPE:-sglang}" --opd-kl-coef "${OPD_KL_COEF:-1.0}") fi @@ -254,12 +281,88 @@ fi if is_true "${SGLANG_ENABLE_MIXED_CHUNK:-false}"; then SGLANG_ARGS+=(--sglang-enable-mixed-chunk) fi +# Prefix (RadixAttention) caching: reuse KV cache for shared prompt prefixes +# (system prompt, task template, conversation history across multi-turn agent +# steps). Big win for PatchEval where every episode shares the same system +# prompt and the same task description across GRPO samples. Without this the +# SGLang log shows #cached-token: 0 on every prefill. Default on; disable via +# SGLANG_ENABLE_PREFIX_CACHING=false. +if is_true "${SGLANG_ENABLE_PREFIX_CACHING:-true}"; then + SGLANG_ARGS+=(--sglang-enable-prefix-caching) +fi + +# Mamba/GDN scheduler strategy: official Qwen3.5-27B uses extra_buffer to +# avoid illegal memory access in mamba_pool allocation. +SGLANG_ARGS+=(--sglang-mamba-scheduler-strategy "${SGLANG_MAMBA_SCHEDULER_STRATEGY:-extra_buffer}") + +# EAGLE speculative decoding: official Qwen3.5-27B uses EAGLE for faster decode. +if [[ -n "${SGLANG_SPECULATIVE_ALGORITHM:-}" ]]; then + SGLANG_ARGS+=( + --sglang-speculative-algorithm "${SGLANG_SPECULATIVE_ALGORITHM}" + --sglang-speculative-num-steps "${SGLANG_SPECULATIVE_NUM_STEPS:-3}" + --sglang-speculative-eagle-topk "${SGLANG_SPECULATIVE_EAGLE_TOPK:-1}" + --sglang-speculative-num-draft-tokens "${SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS:-4}" + ) +fi + +# Router policy: how the SGLang router distributes requests across engines. +# cache_aware (sglang default) — greedy per-request prefix match; under high +# concurrency it scatters one session's turns across +# engines, so ~78% of prefills recompute the full prompt +# (cached-token=0). Wastes the multi-engine capacity. +# manual (chosen here) — sticky-session routing via the X-SMG-Routing-Key +# header that llm_proxy now sends. Each session_id is pinned +# to one worker and stays there (only remaps if that worker +# dies). Stronger stickiness than consistent_hashing, ideal +# for fixed-engine RL rollouts. Supported since SGLang Model +# Gateway v0.3.1 (PR #15907, 2025-12-27); the installed +# sglang_router 0.3.2 has it. `consistent_hashing` is a +# newer CLI choice (PR #17972, 2026-02-15) NOT in 0.3.2, so do +# NOT set SGLANG_ROUTER_POLICY=consistent_hashing on this +# build (argparse will reject it and crash startup). +# Override via SGLANG_ROUTER_POLICY if needed. +SGLANG_ARGS+=(--router-policy "${SGLANG_ROUTER_POLICY:-manual}") + +# Colocate mode: training (Megatron) and inference (SGLang) share the SAME GPUs. +# Required for big models on few GPUs (e.g. 27B on a single 8-card node): the +# dedicated-pool split (actor + rollout = NUM_GPUS) would need ~16 cards for 27B, +# but colocate time-shares 8 cards via CPU offload between rollout/train phases. +# When on, --rollout-num-gpus is ignored (auto = actor GPUs) and --offload is +# forced by the trainer. Set SLIME_COLOCATE=1 to enable. +COLOCATE_ARGS=() +ROLLOUT_NUM_GPUS_ARG="" +if is_true "${SLIME_COLOCATE:-false}"; then + COLOCATE_ARGS=(--colocate) + # In colocate mode, --rollout-num-gpus is auto-set to actor GPUs. + # Don't pass it (matches official Qwen3.5-27B script). + echo " Colocate: ON (train+rollout share all actor GPUs)" +else + ROLLOUT_NUM_GPUS_ARG="--rollout-num-gpus ${ROLLOUT_NUM_GPUS}" +fi RAY_RUNTIME_PYTHONPATH="${SLIME_HOME}:${AIEVOBOX_ROOT}/rl:${AIEVOBOX_ROOT}:${MEGATRON_HOME}" if [[ -n "${PYTHONPATH:-}" ]]; then RAY_RUNTIME_PYTHONPATH="${RAY_RUNTIME_PYTHONPATH}:${PYTHONPATH}" fi +# Colocate mode requires torch_memory_saver for both training and rollout engines. +# The training actor sets LD_PRELOAD in actor_group.py, but the sglang rollout +# engine does NOT — without it, torch_memory_saver fails to initialize +# (_TorchMemorySaverImpl crashes). Compute the .so path and inject it into the +# Ray runtime env so sglang engines also get the hook. +TMS_HOOK_SO="" +if is_true "${SLIME_COLOCATE:-false}"; then + TMS_HOOK_SO=$(python3 -c " +import torch_memory_saver, os +p = os.path.join(os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), + 'torch_memory_saver_hook_mode_preload.abi3.so') +print(p if os.path.exists(p) else '') +" 2>/dev/null || echo "") + if [[ -z "${TMS_HOOK_SO}" ]]; then + echo "WARNING: torch_memory_saver_hook_mode_preload.abi3.so not found; colocate may fail" + fi +fi + RUNTIME_ENV_JSON="{\ \"env_vars\": {\ \"AIEVOBOX_ROOT\": \"${AIEVOBOX_ROOT}\",\ @@ -284,7 +387,14 @@ RUNTIME_ENV_JSON="{\ \"OPD_TEACHER_MAX_CONCURRENCY\": \"${OPD_TEACHER_MAX_CONCURRENCY:-}\",\ \"OPD_TEACHER_TIMEOUT_SECONDS\": \"${OPD_TEACHER_TIMEOUT_SECONDS:-}\",\ \"WANDB_MODE\": \"${WANDB_MODE}\",\ - \"WANDB_DIR\": \"${WANDB_DIR}\"\ + \"WANDB_DIR\": \"${WANDB_DIR}\",\ + \"NCCL_IB_DISABLE\": \"${NCCL_IB_DISABLE:-1}\",\ + \"NCCL_NET\": \"${NCCL_NET:-Socket}\",\ + \"NCCL_SOCKET_IFNAME\": \"${NCCL_SOCKET_IFNAME:-bond0}\",\ + \"PYTORCH_CUDA_ALLOC_CONF\": \"${PYTORCH_CUDA_ALLOC_CONF}\",\ + \"PYTORCH_ALLOC_CONF\": \"${PYTORCH_ALLOC_CONF}\",\ + \"TRAJ_TRUNCATION_MAX_SEQ_LEN\": \"${TRAJ_TRUNCATION_MAX_SEQ_LEN:-8192}\",\ + \"LOSS_MASK_TYPE\": \"${LOSS_MASK_TYPE:-qwen3_5}\"\ }\ }" @@ -310,14 +420,22 @@ RAY_START_ARGS=(start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NU if [[ -n "${RAY_PORT:-}" ]]; then RAY_START_ARGS+=(--port "${RAY_PORT}") fi -"${RAY_BIN}" "${RAY_START_ARGS[@]}" +# Multi-node: pre-build the Ray cluster manually (head + `ray start --address` +# on workers), then run with SKIP_RAY_START=1 so this script reuses the existing +# cluster instead of `ray start --head` (which would restart Ray and drop the +# workers). Also set CLEANUP_BEFORE_RUN=false so the pre-started cluster survives. +if is_true "${SKIP_RAY_START:-false}"; then + echo "SKIP_RAY_START=1: reusing existing Ray cluster at ${RAY_ADDRESS} (multi-node)" +else + "${RAY_BIN}" "${RAY_START_ARGS[@]}" +fi "${RAY_BIN}" job submit --address="${RAY_ADDRESS}" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ -- "${PYTHON_BIN}" "${TRAIN_ENTRYPOINT}" \ --actor-num-nodes "${ACTOR_NUM_NODES}" \ --actor-num-gpus-per-node "${ACTOR_NUM_GPUS_PER_NODE}" \ - --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" \ + ${ROLLOUT_NUM_GPUS_ARG} \ "${MODEL_ARGS[@]}" \ "${MEGATRON_ARGS[@]}" \ "${CKPT_ARGS[@]}" \ @@ -327,5 +445,6 @@ fi "${WANDB_ARGS[@]}" \ "${TRAIN_ARGS[@]}" \ "${SGLANG_ARGS[@]}" \ + "${COLOCATE_ARGS[@]}" \ "${TEACHER_ARGS[@]}" \ 2>&1 | tee "${AIEVOBOX_RUN_DIR}/slime.log" diff --git a/rl/slime_generator.py b/rl/slime_generator.py index 4128dee2..805c689e 100644 --- a/rl/slime_generator.py +++ b/rl/slime_generator.py @@ -32,12 +32,28 @@ import llm_proxy as _llm_proxy_module from trajectory_mask_builder import TrajectoryMaskBuilder +from chat_template_adapter import create_adapter from opd.teacher_log_probs import attach_teacher_log_probs +from timing_log import emit as _timing_emit, now_s as _timing_now + __all__ = ["generate_rollout"] logger = logging.getLogger(__name__) +# Timestamp (perf_counter) at the end of the previous rollout step, used to +# derive the inter-step "train time" (time slime spends on the GRPO update + +# checkpointing between two rollout calls). None before the first step. +_prev_rollout_end: Optional[float] = None + +# Wall-clock epoch seconds (time.time()) of the previous rollout_step emission. +# The weight update of step N completes during the train phase right after the +# rollout_step N event, so the interval between two consecutive weight updates +# equals the gap between two consecutive rollout_step emissions. We record +# that gap as `weight_update_interval_s` so the real update-weight cadence is +# directly readable from the log without post-processing timestamps. +_prev_rollout_step_ts: Optional[float] = None + # Global variables TOKENIZER = None TRAJECTORY_MASK_BUILDER = None @@ -197,14 +213,19 @@ def _init_llm_proxy_server(args): except Exception: processor = None - # 3. TrajectoryMaskBuilder - TRAJECTORY_MASK_BUILDER = TrajectoryMaskBuilder(TOKENIZER, processor) + # 3. Chat template adapter (model-specific normalization / tool-call parsing) + adapter_type = os.environ.get("LOSS_MASK_TYPE", "") + _chat_adapter = create_adapter(adapter_type, TOKENIZER, processor) + + # 4. TrajectoryMaskBuilder (uses the adapter for message rendering) + TRAJECTORY_MASK_BUILDER = TrajectoryMaskBuilder(TOKENIZER, processor, adapter=_chat_adapter) - # 4. Wire into llm_proxy module STATE (shared in-process) + # 5. Wire into llm_proxy module STATE (shared in-process) state = _llm_proxy_module.STATE state.tokenizer = TOKENIZER state.processor = processor state.trajectory_mask_builder = TRAJECTORY_MASK_BUILDER + state.chat_template_adapter = _chat_adapter remote_engine_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" state.remote_engine_url = remote_engine_url @@ -280,6 +301,16 @@ def build_loss_mask_from_response_mask( def _get_record_training_info(record: Dict[str, Any]) -> Dict[str, Any]: oai_messages = record["messages"] session_id = record["extra_info"].get("session_id", "") + # Re-apply the SAME normalization llm_proxy applied at generation time + # (e.g. Qwen: tool_call.arguments JSON string -> dict; content None -> ""). + # The mask builder's in-memory session tree stores NORMALIZED messages + # (prepare_generate_input is called after adapter.normalize_messages + # in llm_proxy.proxy_chat_completions). The DB, however, stores the raw + # OpenAI format (arguments as JSON string, content possibly null). Without + # re-normalizing here, _message_matches compares dict-arguments vs + # JSON-string-arguments (and "" vs None content) -> matched=0 for every + # session -> 0 trainable groups -> no training -> weight_version stuck at 1. + oai_messages = _llm_proxy_module.STATE.chat_template_adapter.normalize_messages(oai_messages) tokens, response_mask, _image_data, messages_str, mm_train_inputs = TRAJECTORY_MASK_BUILDER.get_training_info( session_id, oai_messages, @@ -669,12 +700,96 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: metrics.record("used/count", float(sum(len(g) for g in final_return_results)), AggType.SUM) metrics.push(step=rollout_id) + # In colocate mode, training and inference share the same GPUs, so before + # the training step begins we must kill all env processes (so no env keeps + # sending LLM requests to SGLang) and abort residual SGLang requests (so + # flush_cache during release/resume succeeds immediately). This whole + # block is colocate-only. + # + # In non-colocate mode, training and inference use SEPARATE GPUs — there is + # no release/resume memory cycle and no flush_cache. Killing the launcher + # here is actively harmful: START_ROLLOUT is only True on the first rollout, + # so once the launcher is killed it is never restarted, and subsequent + # rollout rounds have no envs producing data → buffer never fills → the + # pipeline stalls forever (the run gets stuck on step 2). Therefore in + # non-colocate mode we keep the launcher alive so it keeps producing + # trajectories across rollout rounds. + colocate = os.environ.get("SLIME_COLOCATE", "false").lower() in ("true", "1") + if colocate: + try: + stop_url = f"{base_url}/stop_rollout" + resp = requests.post(stop_url, timeout=15) + if resp.status_code == 200: + logger.info(f"[generate_rollout] Stopped all envs (kill launcher.py)") + print(f"[generate_rollout] Stopped all envs before returning data") + else: + logger.warning(f"[generate_rollout] stop_rollout returned HTTP {resp.status_code}") + except Exception as e: + logger.warning(f"[generate_rollout] Failed to stop envs: {e}") + + # Abort residual pending requests on all SGLang workers. + # Even though envs are killed, there may be requests that envs sent + # just before being killed, still sitting in SGLang's queue. These + # residual requests (especially long generations with max_tokens=32768) + # would keep the scheduler busy for minutes, causing flush_cache to + # time out. Aborting them ensures the scheduler becomes idle quickly. + # This mirrors slime's original abort() in sglang_rollout.py. + try: + router_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + # Try /workers first (sglang_router > 0.2.1), fall back to /list_workers + worker_urls = None + try: + resp = requests.get(f"{router_url}/workers", timeout=10) + if resp.status_code == 200: + worker_urls = [w["url"] for w in resp.json().get("workers", [])] + except Exception: + pass + if not worker_urls: + resp = requests.get(f"{router_url}/list_workers", timeout=10) + if resp.status_code == 200: + worker_urls = resp.json().get("urls", []) + + if worker_urls: + for url in worker_urls: + try: + requests.post( + f"{url}/abort_request", + json={"abort_all": True}, + timeout=10, + ) + except Exception as e: + logger.warning(f"[generate_rollout] abort failed for {url}: {e}") + logger.info( + f"[generate_rollout] Aborted residual requests on " + f"{len(worker_urls)} SGLang workers" + ) + print(f"[generate_rollout] Aborted residual requests on {len(worker_urls)} workers") + else: + logger.warning("[generate_rollout] No SGLang worker URLs found to abort") + except Exception as e: + logger.warning(f"[generate_rollout] Failed to abort SGLang workers: {e}") + + # Give SGLang a moment to process the abort and clean up + time.sleep(3) + else: + # Non-colocate: stop launcher so the next rollout_id can restart it + # with current weights. The launcher starts a fixed batch (pool_size) + # and does NOT replenish, so keeping it alive means no new data after + # the batch finishes. The next generate_rollout call will start_rollout + # a fresh launcher with the updated weight version. + try: + requests.post(f"{base_url}/stop_rollout", timeout=15) + logger.info(f"[generate_rollout] Non-colocate: stopped launcher (will restart next rollout)") + print(f"[generate_rollout] Non-colocate: launcher stopped") + except Exception as e: + logger.warning(f"[generate_rollout] stop_rollout failed: {e}") + return final_return_results def generate_rollout(args, rollout_id, data_buffer, evaluation=False): """Generate rollout for both training and evaluation.""" - global START_ROLLOUT + global START_ROLLOUT, _prev_rollout_end, _prev_rollout_step_ts # Initialize tokenizer + processor + llm_proxy HTTP server (once). # Must happen BEFORE start_rollout, because buffer_server will launch @@ -687,8 +802,60 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False): print(f"start rollout with payload: {start_inform}") print(f"start rollout id: {rollout_id}") START_ROLLOUT = False + elif not evaluation: + # rollout_id > 1: restart launcher so new envs use current weights. + # The old launcher's envs were started with a stale weight_version and + # their data would be filtered out by weight_version check; worse, the + # launcher starts a fixed batch (pool_size) and does NOT replenish, so + # once those envs finish there is no new data at all. Stop the old + # launcher and start a fresh one each rollout round. + try: + requests.post(f"{args.rollout_buffer_url}/stop_rollout", timeout=15) + logger.info(f"[generate_rollout] Stopped old launcher before rollout_id={rollout_id}") + except Exception as e: + logger.warning(f"[generate_rollout] stop_rollout before restart failed: {e}") + metadata = data_buffer.get_metadata() + start_inform = start_rollout(args.rollout_buffer_url, args, metadata) + print(f"restart rollout for rollout_id={rollout_id}: {start_inform}") + + rollout_start = _timing_now() + # train_time = time slime spent on the GRPO update + ckpt between the end + # of the previous rollout call and the start of this one (None for step 0). + train_time_s = (rollout_start - _prev_rollout_end) if _prev_rollout_end is not None else None sample_groups = run(generate_rollout_async(args, rollout_id, data_buffer, evaluation)) if evaluation: + rollout_end = _timing_now() + _prev_rollout_end = rollout_end + _step_ts = time.time() + _wu_interval = round(_step_ts - _prev_rollout_step_ts, 3) if _prev_rollout_step_ts is not None else None + _timing_emit( + "rollout_step", + rollout_id=rollout_id, + evaluation=True, + rollout_time_s=round(rollout_end - rollout_start, 3), + train_time_s=round(train_time_s, 3) if train_time_s is not None else None, + weight_update_interval_s=_wu_interval, + global_batch_size=int(os.environ.get("RL_GLOBAL_BATCH_SIZE") or 0), + rollout_batch_size=int(os.environ.get("SLIME_ROLLOUT_BATCH_SIZE") or 0), + num_groups=len(sample_groups) if sample_groups is not None else None, + ) + _prev_rollout_step_ts = _step_ts return sample_groups - return run(attach_teacher_log_probs(args, sample_groups)) + sample_groups = run(attach_teacher_log_probs(args, sample_groups)) + rollout_end = _timing_now() + _step_ts = time.time() + _wu_interval = round(_step_ts - _prev_rollout_step_ts, 3) if _prev_rollout_step_ts is not None else None + _timing_emit( + "rollout_step", + rollout_id=rollout_id, + rollout_time_s=round(rollout_end - rollout_start, 3), + train_time_s=round(train_time_s, 3) if train_time_s is not None else None, + weight_update_interval_s=_wu_interval, + global_batch_size=int(os.environ.get("RL_GLOBAL_BATCH_SIZE") or 0), + rollout_batch_size=int(os.environ.get("SLIME_ROLLOUT_BATCH_SIZE") or 0), + num_groups=len(sample_groups) if sample_groups is not None else None, + ) + _prev_rollout_step_ts = _step_ts + _prev_rollout_end = rollout_end + return sample_groups diff --git a/rl/timing_log.py b/rl/timing_log.py new file mode 100644 index 00000000..5386439c --- /dev/null +++ b/rl/timing_log.py @@ -0,0 +1,105 @@ +"""Structured timing log for offline analysis. + +Writes one JSON line per event to a single append-only file so the buffer +server (rollout side) and the slime generator (training side) can both record +into the same log even when they run as separate processes. + +Default path is ``${LOG_ROOT}/timing.jsonl`` (shared by both sides via the +env.sh), overridable with ``SAFACTORY_TIMING_LOG``. Each line carries an +``event`` field and a monotonic ``ts`` (epoch seconds) plus whatever fields +the caller passes. Lines are flushed immediately so a crash never loses +already-recorded events. +""" + +from __future__ import annotations + +import json +import os +import time +from typing import Any, Dict, Optional + +_DEFAULT_LOG_NAME = "timing.jsonl" +_file_handle = None +_file_path: Optional[str] = None + +# Module-level enable switch. RL runs rely on timing data being on by +# default; non-RL callers (e.g. a manually-started gateway) can opt out with +# SAFACTORY_TIMING_LOG_ENABLED=0. set_enabled() lets a process flip it at +# runtime (e.g. from loaded config) without touching env vars. +_enabled = str(os.environ.get("SAFACTORY_TIMING_LOG_ENABLED", "1")).strip().lower() not in ( + "0", "false", "no", "off", +) + + +def set_enabled(value: bool) -> None: + """Enable/disable timing emission for this process. + + Emission is on by default (RL depends on it). Call ``set_enabled(False)`` + to silence this process, e.g. a non-RL gateway that pulled in the module + but does not want per-step timing records. + """ + global _enabled + _enabled = bool(value) + + +def _resolve_path() -> str: + override = os.environ.get("SAFACTORY_TIMING_LOG", "").strip() + if override: + return override + log_root = os.environ.get("LOG_ROOT", "").strip() or "/tmp" + # Prefer the current run directory (written by run_slime_generator.sh / + # run_buffer_server.sh) so timing events are per-run instead of + # accumulating in a single root-level file. + current_run_file = os.path.join(log_root, ".current_run") + try: + with open(current_run_file, "r") as f: + run_dir = f.read().strip() + if run_dir and os.path.isdir(run_dir): + return os.path.join(run_dir, _DEFAULT_LOG_NAME) + except Exception: + pass + return os.path.join(log_root, _DEFAULT_LOG_NAME) + + +def _ensure_handle(): + global _file_handle, _file_path + path = _resolve_path() + if _file_handle is not None and _file_path == path: + return _file_handle + if _file_handle is not None: + try: + _file_handle.flush() + _file_handle.close() + except Exception: + pass + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + # Line-buffered append: each write is a full line, flushed right away. + _file_handle = open(path, "a", buffering=1, encoding="utf-8") + _file_path = path + return _file_handle + + +def emit(event: str, **fields: Any) -> None: + """Append one timing event as a JSON line. + + Never raises: logging must not affect the training/rollout process. + Silently drops the event when the module is disabled (see set_enabled() / + SAFACTORY_TIMING_LOG_ENABLED). + """ + if not _enabled: + return + try: + record: Dict[str, Any] = {"event": event, "ts": time.time()} + record.update(fields) + line = json.dumps(record, ensure_ascii=False, default=str) + handle = _ensure_handle() + handle.write(line + "\n") + handle.flush() + except Exception: + # Best-effort: drop on the floor rather than killing the run. + pass + + +def now_s() -> float: + """Monotonic seconds, for callers that measure spans themselves.""" + return time.perf_counter() diff --git a/tests/test_patcheval_strict_protocol.py b/tests/test_patcheval_strict_protocol.py index 9ee54c23..bcfff727 100644 --- a/tests/test_patcheval_strict_protocol.py +++ b/tests/test_patcheval_strict_protocol.py @@ -5,7 +5,7 @@ import unittest from pathlib import Path -from env.patcheval import strict_runner +from env.patcheval import runner ROOT = Path(__file__).resolve().parents[1] @@ -26,7 +26,7 @@ def test_each_s1_prompt_matches_official_builder(self) -> None: functions = [ { "id": item["id"], - "original_code": strict_runner._process_original_code(item["snippet"]), + "original_code": runner._process_original_code(item["snippet"]), } for item in self.record["vul_func"] ] @@ -51,7 +51,7 @@ def test_each_s1_prompt_matches_official_builder(self) -> None: self.record["cve_id"], [], ) - actual = strict_runner._build_prompt( + actual = runner._build_prompt( record=self.record, vul_functions=self.record["vul_func"], feedbacks={}, @@ -63,7 +63,7 @@ def test_each_s1_prompt_matches_official_builder(self) -> None: def test_official_function_json_is_parsed(self) -> None: response = '```json\n[{"id":"vul_py_1","patch":"def fixed():\\n return True"}]\n```' self.assertEqual( - strict_runner._parse_response(response), + runner._parse_response(response), {"vul_py_1": "def fixed():\n return True"}, )