ins-pricing 0.1.6__py3-none-any.whl

user_packages legacy/notebook_utils.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, cast
+
+
+def _find_user_packages_dir(cwd: Optional[Path] = None) -> Path:
+    cwd = (cwd or Path().resolve()).resolve()
+    candidates = [cwd / "user_packages", cwd, cwd.parent / "user_packages"]
+    for cand in candidates:
+        if (cand / "BayesOpt_entry.py").exists() and (cand / "watchdog_run.py").exists():
+            return cand
+    raise FileNotFoundError(
+        "Cannot locate user_packages directory (expected BayesOpt_entry.py and watchdog_run.py). "
+        f"cwd={cwd}"
+    )
+
+
+def _stringify_cmd(cmd: Sequence[object]) -> List[str]:
+    return [str(x) for x in cmd]
+
+
+def build_bayesopt_entry_cmd(
+    config_json: str | Path,
+    model_keys: Sequence[str],
+    *,
+    nproc_per_node: int = 1,
+    standalone: bool = True,
+    entry_script: str | Path = "BayesOpt_entry.py",
+    extra_args: Optional[Sequence[str]] = None,
+) -> List[str]:
+    """构造运行 BayesOpt_entry.py 的命令（可选 torchrun/DDP）。"""
+    pkg_dir = _find_user_packages_dir()
+    entry_path = (pkg_dir / entry_script).resolve() if not Path(entry_script).is_absolute() else Path(entry_script).resolve()
+    config_path = Path(config_json)
+    if not config_path.is_absolute():
+        config_path = (pkg_dir / config_path).resolve() if (pkg_dir / config_path).exists() else config_path.resolve()
+
+    cmd: List[object]
+    if int(nproc_per_node) > 1:
+        cmd = [
+            sys.executable,
+            "-m",
+            "torch.distributed.run",
+            *(["--standalone"] if standalone else []),
+            f"--nproc_per_node={int(nproc_per_node)}",
+            str(entry_path),
+        ]
+    else:
+        cmd = [sys.executable, str(entry_path)]
+
+    cmd += ["--config-json", str(config_path), "--model-keys", *list(model_keys)]
+    if extra_args:
+        cmd += list(extra_args)
+    return _stringify_cmd(cmd)
+
+
+def build_incremental_cmd(
+    config_json: str | Path,
+    *,
+    entry_script: str | Path = "BayesOpt_incremental.py",
+    extra_args: Optional[Sequence[str]] = None,
+) -> List[str]:
+    """构造运行 BayesOpt_incremental.py 的命令。"""
+    pkg_dir = _find_user_packages_dir()
+    entry_path = (pkg_dir / entry_script).resolve() if not Path(entry_script).is_absolute() else Path(entry_script).resolve()
+    config_path = Path(config_json)
+    if not config_path.is_absolute():
+        config_path = (pkg_dir / config_path).resolve() if (pkg_dir / config_path).exists() else config_path.resolve()
+
+    cmd: List[object] = [sys.executable, str(entry_path), "--config-json", str(config_path)]
+    if extra_args:
+        cmd += list(extra_args)
+    return _stringify_cmd(cmd)
+
+
+def wrap_with_watchdog(
+    cmd: Sequence[str],
+    *,
+    idle_seconds: int = 7200,
+    max_restarts: int = 50,
+    restart_delay_seconds: int = 10,
+    stop_on_nonzero_exit: bool = True,
+    watchdog_script: str | Path = "watchdog_run.py",
+) -> List[str]:
+    """用 watchdog 包一层命令：超过 idle_seconds 无输出则自动杀进程树并重启。"""
+    pkg_dir = _find_user_packages_dir()
+    watchdog_path = (pkg_dir / watchdog_script).resolve() if not Path(watchdog_script).is_absolute() else Path(watchdog_script).resolve()
+    wd_cmd: List[object] = [
+        sys.executable,
+        str(watchdog_path),
+        "--idle-seconds",
+        str(int(idle_seconds)),
+        "--max-restarts",
+        str(int(max_restarts)),
+        "--restart-delay-seconds",
+        str(int(restart_delay_seconds)),
+    ]
+    if stop_on_nonzero_exit:
+        wd_cmd.append("--stop-on-nonzero-exit")
+    wd_cmd.append("--")
+    wd_cmd.extend(list(cmd))
+    return _stringify_cmd(wd_cmd)
+
+
+def run(cmd: Sequence[str], *, check: bool = True) -> subprocess.CompletedProcess:
+    """在 notebook 中运行外部命令（同步等待）。"""
+    return subprocess.run(list(cmd), check=check)
+
+
+def run_bayesopt_entry(
+    *,
+    config_json: str | Path,
+    model_keys: Sequence[str],
+    max_evals: int = 50,
+    plot_curves: bool = True,
+    ft_role: Optional[str] = None,
+    nproc_per_node: int = 1,
+    use_watchdog: bool = False,
+    idle_seconds: int = 7200,
+    max_restarts: int = 50,
+    restart_delay_seconds: int = 10,
+    extra_args: Optional[Sequence[str]] = None,
+) -> subprocess.CompletedProcess:
+    """便捷封装：构造并运行 BayesOpt_entry（可选 torchrun + watchdog）。"""
+    args: List[str] = ["--max-evals", str(int(max_evals))]
+    if plot_curves:
+        args.append("--plot-curves")
+    if ft_role:
+        args += ["--ft-role", str(ft_role)]
+    if extra_args:
+        args += list(extra_args)
+
+    cmd = build_bayesopt_entry_cmd(
+        config_json=config_json,
+        model_keys=model_keys,
+        nproc_per_node=nproc_per_node,
+        extra_args=args,
+    )
+    if use_watchdog:
+        cmd = wrap_with_watchdog(
+            cmd,
+            idle_seconds=idle_seconds,
+            max_restarts=max_restarts,
+            restart_delay_seconds=restart_delay_seconds,
+        )
+    return run(cmd, check=True)
+
+
+def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
+    """Notebook 统一入口：只修改 config 文件即可切换不同运行方式。
+
+    约定：config.json 可选增加一个 `runner` 字段，用于 notebook 执行控制：
+    - runner.mode: "entry"（默认）或 "incremental"
+    - runner.nproc_per_node: >1 则使用 torchrun/DDP（仅 entry）
+    - runner.model_keys: 需要运行的模型列表（仅 entry）
+    - runner.max_evals / runner.plot_curves / runner.ft_role（仅 entry，可覆盖 config 内同名字段）
+    - runner.use_watchdog / runner.idle_seconds / runner.max_restarts / runner.restart_delay_seconds
+    - runner.incremental_args: List[str]（仅 incremental，等价于直接传给 BayesOpt_incremental.py 的额外参数）
+    """
+    pkg_dir = _find_user_packages_dir()
+    config_path = Path(config_json)
+    if not config_path.is_absolute():
+        config_path = (pkg_dir / config_path).resolve() if (pkg_dir / config_path).exists() else config_path.resolve()
+    raw = json.loads(config_path.read_text(encoding="utf-8", errors="replace"))
+    runner = cast(dict, raw.get("runner") or {})
+
+    mode = str(runner.get("mode") or "entry").strip().lower()
+    use_watchdog = bool(runner.get("use_watchdog", False))
+    idle_seconds = int(runner.get("idle_seconds", 7200))
+    max_restarts = int(runner.get("max_restarts", 50))
+    restart_delay_seconds = int(runner.get("restart_delay_seconds", 10))
+
+    if mode == "incremental":
+        inc_args = runner.get("incremental_args") or []
+        if not isinstance(inc_args, list):
+            raise ValueError("config.runner.incremental_args must be a list of strings.")
+        cmd = build_incremental_cmd(config_path, extra_args=[str(x) for x in inc_args])
+        if use_watchdog:
+            cmd = wrap_with_watchdog(
+                cmd,
+                idle_seconds=idle_seconds,
+                max_restarts=max_restarts,
+                restart_delay_seconds=restart_delay_seconds,
+            )
+        return run(cmd, check=True)
+
+    if mode != "entry":
+        raise ValueError(f"Unsupported runner.mode={mode!r}, expected 'entry' or 'incremental'.")
+
+    model_keys = runner.get("model_keys")
+    if not model_keys:
+        model_keys = raw.get("model_keys")
+    if not model_keys:
+        model_keys = ["ft"]
+    if not isinstance(model_keys, list):
+        raise ValueError("runner.model_keys must be a list of strings.")
+
+    nproc_per_node = int(runner.get("nproc_per_node", 1))
+    max_evals = int(runner.get("max_evals", raw.get("max_evals", 50)))
+    plot_curves = bool(runner.get("plot_curves", raw.get("plot_curves", True)))
+    ft_role = runner.get("ft_role", None)
+    if ft_role is None:
+        ft_role = raw.get("ft_role")
+
+    cmd = build_bayesopt_entry_cmd(
+        config_path,
+        model_keys=[str(x) for x in model_keys],
+        nproc_per_node=nproc_per_node,
+        extra_args=[
+            "--max-evals",
+            str(max_evals),
+            *(["--plot-curves"] if plot_curves else []),
+            *(["--ft-role", str(ft_role)] if ft_role else []),
+        ],
+    )
+
+    if use_watchdog:
+        cmd = wrap_with_watchdog(
+            cmd,
+            idle_seconds=idle_seconds,
+            max_restarts=max_restarts,
+            restart_delay_seconds=restart_delay_seconds,
+        )
+    return run(cmd, check=True)

user_packages legacy/watchdog_run.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+import threading
+import time
+from typing import List, Optional
+
+
+def _split_argv(argv: List[str]) -> tuple[List[str], List[str]]:
+    if "--" not in argv:
+        raise ValueError("Missing '--' separator before the command to run.")
+    idx = argv.index("--")
+    return argv[:idx], argv[idx + 1 :]
+
+
+def _kill_process_tree(pid: int) -> None:
+    if pid <= 0:
+        return
+    if os.name == "nt":
+        subprocess.run(
+            ["taskkill", "/PID", str(pid), "/T", "/F"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        )
+        return
+    try:
+        os.killpg(pid, 15)
+        time.sleep(2)
+        os.killpg(pid, 9)
+    except Exception:
+        try:
+            os.kill(pid, 9)
+        except Exception:
+            pass
+
+
+def _reader_thread(
+    proc: subprocess.Popen, last_output_ts: dict, prefix: str = ""
+) -> None:
+    assert proc.stdout is not None
+    for line in proc.stdout:
+        last_output_ts["ts"] = time.time()
+        if prefix:
+            sys.stdout.write(prefix)
+        sys.stdout.write(line)
+        sys.stdout.flush()
+
+
+def _parse_args(before_cmd: List[str], cmd: List[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Run a command under a simple watchdog: if there is no stdout/stderr "
+            "output for N seconds, kill the whole process tree and restart. "
+            "Designed to pair with optuna_storage so BayesOpt can resume."
+        )
+    )
+    parser.add_argument(
+        "--idle-seconds",
+        type=int,
+        default=7200,
+        help="Restart if there is no output for this many seconds (default: 7200).",
+    )
+    parser.add_argument(
+        "--max-restarts",
+        type=int,
+        default=50,
+        help="Maximum restart attempts (default: 50).",
+    )
+    parser.add_argument(
+        "--restart-delay-seconds",
+        type=int,
+        default=10,
+        help="Delay between restarts (default: 10).",
+    )
+    parser.add_argument(
+        "--stop-on-nonzero-exit",
+        action="store_true",
+        help="If the command exits non-zero, stop instead of restarting.",
+    )
+    args = parser.parse_args(before_cmd)
+    if not cmd:
+        parser.error("Empty command after '--'.")
+    return args
+
+
+def run_with_watchdog(
+    cmd: List[str],
+    idle_seconds: int,
+    max_restarts: int,
+    restart_delay_seconds: int,
+    stop_on_nonzero_exit: bool,
+) -> int:
+    idle_seconds = max(1, int(idle_seconds))
+    max_restarts = max(0, int(max_restarts))
+    restart_delay_seconds = max(0, int(restart_delay_seconds))
+
+    attempt = 0
+    while True:
+        attempt += 1
+        print(
+            f"[watchdog] start attempt={attempt} idle_seconds={idle_seconds} cmd={cmd}",
+            flush=True,
+        )
+
+        creationflags = 0
+        start_new_session = False
+        if os.name == "nt":
+            creationflags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
+        else:
+            start_new_session = True
+
+        proc = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,
+            universal_newlines=True,
+            creationflags=creationflags,
+            start_new_session=start_new_session,
+        )
+
+        last_output_ts: dict = {"ts": time.time()}
+        reader = threading.Thread(
+            target=_reader_thread,
+            args=(proc, last_output_ts),
+            kwargs={"prefix": ""},
+            daemon=True,
+        )
+        reader.start()
+
+        killed_for_idle = False
+        exit_code: Optional[int] = None
+        while True:
+            exit_code = proc.poll()
+            if exit_code is not None:
+                break
+            idle_for = time.time() - float(last_output_ts["ts"])
+            if idle_for > idle_seconds:
+                killed_for_idle = True
+                print(
+                    f"[watchdog] idle>{idle_seconds}s (idle_for={int(idle_for)}s), killing pid={proc.pid}",
+                    flush=True,
+                )
+                _kill_process_tree(proc.pid)
+                break
+            time.sleep(5)
+
+        try:
+            proc.wait(timeout=30)
+        except Exception:
+            _kill_process_tree(proc.pid)
+
+        if exit_code is None:
+            exit_code = proc.poll() or 1
+
+        if exit_code == 0:
+            print("[watchdog] finished with exit_code=0", flush=True)
+            return 0
+
+        if stop_on_nonzero_exit and not killed_for_idle:
+            print(
+                f"[watchdog] command exited non-zero (exit_code={exit_code}); stop.",
+                flush=True,
+            )
+            return int(exit_code)
+
+        if attempt > max_restarts + 1:
+            print(
+                f"[watchdog] exceeded max_restarts={max_restarts}; last exit_code={exit_code}",
+                flush=True,
+            )
+            return int(exit_code)
+
+        print(
+            f"[watchdog] restart in {restart_delay_seconds}s (exit_code={exit_code}, killed_for_idle={killed_for_idle})",
+            flush=True,
+        )
+        if restart_delay_seconds:
+            time.sleep(restart_delay_seconds)
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    before_cmd, cmd = _split_argv(argv)
+    args = _parse_args(before_cmd, cmd)
+    return run_with_watchdog(
+        cmd=cmd,
+        idle_seconds=args.idle_seconds,
+        max_restarts=args.max_restarts,
+        restart_delay_seconds=args.restart_delay_seconds,
+        stop_on_nonzero_exit=bool(args.stop_on_nonzero_exit),
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+