Mechanical refactor transform: reorder RolloutManager

transform_reorder_rollout_manager.py

#!/usr/bin/env python3
"""Reproducible transform for: reorder RolloutManager methods and privatize some names.

Run from the repo root:  python3 /tmp/transform_reorder_rollout_manager.py
"""
import sys
from pathlib import Path

sys.path.append(".claude/skills/mechanical-refactor-verify")
from mechanical_refactor_verify_utils import verify_mechanical_refactor, git_add_and_commit

BASE_COMMIT = "69d53bdd20c0e187d7960b5a0ed943146f2f2b2e"
TARGET_COMMIT = "9b05c04a9"


def transform(dir_root: Path) -> None:
    source = dir_root / "miles/ray/rollout/rollout_manager.py"
    lines = source.read_text().splitlines(keepends=True)

    # Base file line map (1-indexed):
    # 1-33:    imports + module-level
    # 34-35:   @ray.remote
    # 36-37:   class RolloutManager + docstring
    # 39-86:   __init__
    # 88-108:  _try_ci_fault_injection
    # 110-114: dispose
    # 116-121: server property
    # 123-127: _get_updatable_server
    # 129-132: rollout_engines property
    # 134-141: get_updatable_engines_and_lock
    # 143-145: get_num_rollout_per_epoch
    # 147-163: generate
    # 165-181: eval
    # 183-184: save
    # 186-187: load
    # 189-199: offload
    # 201-203: onload
    # 205-207: health_monitoring_pause
    # 209-211: health_monitoring_resume
    # 213-215: onload_weights
    # 217-219: onload_kv
    # 221-232: recover_updatable_engines
    # 234-238: clear_updatable_num_new_engines
    # 240-241: check_weights
    # 243-261: _get_rollout_data
    # 263-264: set_train_parallel_config

    def L(start: int, end: int) -> str:
        return "".join(lines[start - 1 : end])

    parts = []

    # --- Imports + module-level (lines 1-33, unchanged) ---
    parts.append(L(1, 33))

    # --- @ray.remote + class header + __init__ (lines 34-76, before _metric_checker) ---
    parts.append(L(34, 78))

    # Insert blank line + TODO comment before _health_monitors block
    parts.append("\n")
    parts.append("        # TODO will be replaced by full ft, thus temporarily leave it without modifications\n")

    # Rest of __init__: _health_monitors through end (lines 79-86)
    parts.append(L(79, 86))

    # --- Section: lifecycle ---
    parts.append("\n    # -------------------------- lifecycle -----------------------------\n")
    parts.append("    # TODO: may have a `async def init` here later\n")
    parts.append("\n")
    # dispose (lines 110-114)
    parts.append(L(110, 114))

    # --- Section: data generation ---
    parts.append("\n    # -------------------------- data generation -----------------------------\n")
    parts.append("\n")
    # generate (lines 147-163) — rename health_monitoring_resume
    generate_block = L(147, 163)
    generate_block = generate_block.replace("self.health_monitoring_resume()", "self._health_monitoring_resume()")
    parts.append(generate_block)

    # eval (lines 165-181) — rename health_monitoring_resume
    eval_block = L(165, 181)
    eval_block = eval_block.replace("self.health_monitoring_resume()", "self._health_monitoring_resume()")
    parts.append(eval_block)

    # _get_rollout_data (lines 243-261)
    parts.append(L(243, 261))

    # --- Section: checkpointing ---
    parts.append("\n    # -------------------------- checkpointing -----------------------------\n")
    parts.append("\n")
    parts.append(L(183, 184))
    parts.append("\n")
    parts.append(L(186, 187))
    parts.append("\n")

    # --- Section: offload/onload ---
    parts.append("    # -------------------------- offload/onload -----------------------------\n")
    parts.append("\n")
    offload_block = L(189, 199)
    offload_block = offload_block.replace("self.health_monitoring_pause()", "self._health_monitoring_pause()")
    offload_block = offload_block.replace("self.rollout_engines", "self._rollout_engines")
    parts.append(offload_block)
    parts.append(L(201, 203))
    parts.append("\n")
    parts.append(L(213, 215))
    parts.append("\n")
    parts.append(L(217, 219))
    parts.append("\n")

    # --- Section: engine management ---
    parts.append("    # -------------------------- engine management -----------------------------\n")
    parts.append("\n")
    parts.append(L(134, 141))
    parts.append("\n")
    parts.append(L(234, 238))
    parts.append("\n")
    recover_block = L(221, 232)
    recover_block = recover_block.replace("self.health_monitoring_pause()", "self._health_monitoring_pause()")
    parts.append(recover_block)
    parts.append("\n")
    parts.append(L(123, 127))
    parts.append("\n")

    # --- Section: misc APIs ---
    parts.append("    # -------------------------- misc APIs -----------------------------\n")
    parts.append("\n")
    parts.append(L(143, 145))
    parts.append("\n")
    check_weights_block = L(240, 241)
    check_weights_block = check_weights_block.replace("self.rollout_engines", "self._rollout_engines")
    parts.append(check_weights_block)
    parts.append("\n")
    parts.append(L(263, 264))
    parts.append("\n")

    # --- Section: utils ---
    parts.append("    # -------------------------- utils -----------------------------\n")
    parts.append("\n")
    hm_pause = L(205, 207)
    hm_pause = hm_pause.replace("def health_monitoring_pause(", "def _health_monitoring_pause(")
    parts.append(hm_pause)
    parts.append("\n")
    hm_resume = L(209, 211)
    hm_resume = hm_resume.replace("def health_monitoring_resume(", "def _health_monitoring_resume(")
    parts.append(hm_resume)
    parts.append("\n")
    server_prop = L(116, 121)
    server_prop = server_prop.replace("def server(", "def _server(")
    parts.append(server_prop)
    parts.append("\n")
    engines_prop = L(129, 132)
    engines_prop = engines_prop.replace("def rollout_engines(", "def _rollout_engines(")
    parts.append(engines_prop)
    parts.append("\n")

    # _try_ci_fault_injection (lines 88-108) — self.server -> self._server, updated TODO
    ci_block = L(88, 108)
    ci_block = ci_block.replace(
        "    # TODO will be replaced by full ft\n",
        "    # TODO will be replaced by full ft, thus temporarily leave it without modifications\n",
    )
    ci_block = ci_block.replace("self.server ", "self._server ")
    ci_block = ci_block.replace("self.server.", "self._server.")
    parts.append(ci_block)

    source.write_text("".join(parts))
    git_add_and_commit("mechanical: reorder RolloutManager methods", cwd=str(dir_root))


if __name__ == "__main__":
    verify_mechanical_refactor(
        base_commit=BASE_COMMIT,
        target_commit=TARGET_COMMIT,
        transform=transform,
    )