ins-pricing 0.4.5__py3-none-any.whl → 0.5.1__py3-none-any.whl

ins_pricing/cli/utils/notebook_utils.py

@@ -1,184 +1,184 @@
-from __future__ import annotations
-
-import argparse
-import json
-import subprocess
-import sys
-from pathlib import Path
-from typing import Iterable, List, Optional, Sequence, cast
-
-try:
-    from .cli_config import add_config_json_arg, set_env  # type: ignore
-except Exception:  # pragma: no cover
-    from cli_config import add_config_json_arg, set_env  # type: ignore
-
-
-def _find_ins_pricing_dir(cwd: Optional[Path] = None) -> Path:
-    cwd = (cwd or Path().resolve()).resolve()
-    pkg_root = Path(__file__).resolve().parents[2]
-    candidates = [pkg_root, cwd / "ins_pricing", cwd, cwd.parent / "ins_pricing"]
-    for cand in candidates:
-        cli_entry = cand / "cli" / "BayesOpt_entry.py"
-        cli_watchdog = cand / "cli" / "watchdog_run.py"
-        if cli_entry.exists() and cli_watchdog.exists():
-            return cand
-    raise FileNotFoundError(
-        "Cannot locate ins_pricing directory (expected cli/BayesOpt_entry.py and "
-        "cli/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 = "cli/BayesOpt_entry.py",
-    extra_args: Optional[Sequence[str]] = None,
-) -> List[str]:
-    """Build a command to run cli/BayesOpt_entry.py (optional torchrun/DDP)."""
-    pkg_dir = _find_ins_pricing_dir()
-    entry_script_path = Path(entry_script)
-    if entry_script_path.is_absolute():
-        entry_path = entry_script_path.resolve()
-    else:
-        candidate = pkg_dir / entry_script_path
-        legacy = pkg_dir / "modelling" / entry_script_path
-        entry_path = (
-            candidate.resolve()
-            if candidate.exists()
-            else legacy.resolve()
-            if legacy.exists()
-            else candidate.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 = "cli/BayesOpt_incremental.py",
-    extra_args: Optional[Sequence[str]] = None,
-) -> List[str]:
-    """Build a command to run cli/BayesOpt_incremental.py."""
-    pkg_dir = _find_ins_pricing_dir()
-    entry_script_path = Path(entry_script)
-    if entry_script_path.is_absolute():
-        entry_path = entry_script_path.resolve()
-    else:
-        candidate = pkg_dir / entry_script_path
-        legacy = pkg_dir / "modelling" / entry_script_path
-        entry_path = (
-            candidate.resolve()
-            if candidate.exists()
-            else legacy.resolve()
-            if legacy.exists()
-            else candidate.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 build_explain_cmd(
-    config_json: str | Path,
-    *,
-    entry_script: str | Path = "cli/Explain_entry.py",
-    extra_args: Optional[Sequence[str]] = None,
-) -> List[str]:
-    """Build a command to run cli/Explain_entry.py."""
-    pkg_dir = _find_ins_pricing_dir()
-    entry_script_path = Path(entry_script)
-    if entry_script_path.is_absolute():
-        entry_path = entry_script_path.resolve()
-    else:
-        candidate = pkg_dir / entry_script_path
-        legacy = pkg_dir / "modelling" / entry_script_path
-        entry_path = (
-            candidate.resolve()
-            if candidate.exists()
-            else legacy.resolve()
-            if legacy.exists()
-            else candidate.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)
-
-
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, cast
+
+try:
+    from ins_pricing.cli.utils.cli_config import add_config_json_arg, set_env  # type: ignore
+except Exception:  # pragma: no cover
+    from cli_config import add_config_json_arg, set_env  # type: ignore
+
+
+def _find_ins_pricing_dir(cwd: Optional[Path] = None) -> Path:
+    cwd = (cwd or Path().resolve()).resolve()
+    pkg_root = Path(__file__).resolve().parents[2]
+    candidates = [pkg_root, cwd / "ins_pricing", cwd, cwd.parent / "ins_pricing"]
+    for cand in candidates:
+        cli_entry = cand / "cli" / "BayesOpt_entry.py"
+        cli_watchdog = cand / "cli" / "watchdog_run.py"
+        if cli_entry.exists() and cli_watchdog.exists():
+            return cand
+    raise FileNotFoundError(
+        "Cannot locate ins_pricing directory (expected cli/BayesOpt_entry.py and "
+        "cli/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 = "cli/BayesOpt_entry.py",
+    extra_args: Optional[Sequence[str]] = None,
+) -> List[str]:
+    """Build a command to run cli/BayesOpt_entry.py (optional torchrun/DDP)."""
+    pkg_dir = _find_ins_pricing_dir()
+    entry_script_path = Path(entry_script)
+    if entry_script_path.is_absolute():
+        entry_path = entry_script_path.resolve()
+    else:
+        candidate = pkg_dir / entry_script_path
+        legacy = pkg_dir / "modelling" / entry_script_path
+        entry_path = (
+            candidate.resolve()
+            if candidate.exists()
+            else legacy.resolve()
+            if legacy.exists()
+            else candidate.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 = "cli/BayesOpt_incremental.py",
+    extra_args: Optional[Sequence[str]] = None,
+) -> List[str]:
+    """Build a command to run cli/BayesOpt_incremental.py."""
+    pkg_dir = _find_ins_pricing_dir()
+    entry_script_path = Path(entry_script)
+    if entry_script_path.is_absolute():
+        entry_path = entry_script_path.resolve()
+    else:
+        candidate = pkg_dir / entry_script_path
+        legacy = pkg_dir / "modelling" / entry_script_path
+        entry_path = (
+            candidate.resolve()
+            if candidate.exists()
+            else legacy.resolve()
+            if legacy.exists()
+            else candidate.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 build_explain_cmd(
+    config_json: str | Path,
+    *,
+    entry_script: str | Path = "cli/Explain_entry.py",
+    extra_args: Optional[Sequence[str]] = None,
+) -> List[str]:
+    """Build a command to run cli/Explain_entry.py."""
+    pkg_dir = _find_ins_pricing_dir()
+    entry_script_path = Path(entry_script)
+    if entry_script_path.is_absolute():
+        entry_path = entry_script_path.resolve()
+    else:
+        candidate = pkg_dir / entry_script_path
+        legacy = pkg_dir / "modelling" / entry_script_path
+        entry_path = (
+            candidate.resolve()
+            if candidate.exists()
+            else legacy.resolve()
+            if legacy.exists()
+            else candidate.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 = "cli/watchdog_run.py",
-) -> List[str]:
-    """Wrap a command with watchdog: restart when idle_seconds elapses with no output."""
-    pkg_dir = _find_ins_pricing_dir()
-    watchdog_script_path = Path(watchdog_script)
-    if watchdog_script_path.is_absolute():
-        watchdog_path = watchdog_script_path.resolve()
-    else:
-        candidate = pkg_dir / watchdog_script_path
-        legacy = pkg_dir / "modelling" / watchdog_script_path
-        watchdog_path = (
-            candidate.resolve()
-            if candidate.exists()
-            else legacy.resolve()
-            if legacy.exists()
-            else candidate.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("--")
+    *,
+    idle_seconds: int = 7200,
+    max_restarts: int = 50,
+    restart_delay_seconds: int = 10,
+    stop_on_nonzero_exit: bool = True,
+    watchdog_script: str | Path = "cli/watchdog_run.py",
+) -> List[str]:
+    """Wrap a command with watchdog: restart when idle_seconds elapses with no output."""
+    pkg_dir = _find_ins_pricing_dir()
+    watchdog_script_path = Path(watchdog_script)
+    if watchdog_script_path.is_absolute():
+        watchdog_path = watchdog_script_path.resolve()
+    else:
+        candidate = pkg_dir / watchdog_script_path
+        legacy = pkg_dir / "modelling" / watchdog_script_path
+        watchdog_path = (
+            candidate.resolve()
+            if candidate.exists()
+            else legacy.resolve()
+            if legacy.exists()
+            else candidate.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)
 
@@ -189,72 +189,19 @@ def run(cmd: Sequence[str], *, check: bool = True) -> subprocess.CompletedProces
 
 
 def _build_config_parser(description: str) -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description=description)
-    add_config_json_arg(
-        parser,
-        help_text="Path to config.json (relative paths are resolved from ins_pricing/ when possible).",
+    parser = argparse.ArgumentParser(description=description)
+    add_config_json_arg(
+        parser,
+        help_text="Path to config.json (relative paths are resolved from ins_pricing/ when possible).",
     )
     return parser
 
 
-def run_from_config_cli(
-    description: str,
-    argv: Optional[Sequence[str]] = None,
-) -> subprocess.CompletedProcess:
-    parser = _build_config_parser(description)
-    args = parser.parse_args(argv)
-    return run_from_config(args.config_json)
-
+def build_cmd_from_config(config_json: str | Path) -> tuple[List[str], str]:
+    """Build a command list from config.json runner settings.
 
-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:
-    """Convenience wrapper: build and run BayesOpt_entry (optional 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 entry point: switch execution modes by editing config.json.
-
-    Convention: config.json may include a `runner` section for notebook control:
-    - runner.mode: "entry" (default), "incremental", or "explain"
-    - runner.nproc_per_node: >1 enables torchrun/DDP (entry only)
-    - runner.model_keys: list of models to run (entry only)
-    - runner.max_evals / runner.plot_curves / runner.ft_role (entry only; override config fields)
-    - runner.use_watchdog / runner.idle_seconds / runner.max_restarts / runner.restart_delay_seconds
-    - runner.incremental_args: List[str] (incremental only; extra args for cli/BayesOpt_incremental.py)
+    Returns:
+        (cmd, mode) where mode is one of: entry, incremental, explain.
     """
     pkg_dir = _find_ins_pricing_dir()
     config_path = Path(config_json)
@@ -266,6 +213,10 @@ def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
 
     mode = str(runner.get("mode") or "entry").strip().lower()
     use_watchdog = bool(runner.get("use_watchdog", False))
+    if mode == "watchdog":
+        use_watchdog = True
+        mode = "entry"
+
     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))
@@ -282,7 +233,7 @@ def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
                 max_restarts=max_restarts,
                 restart_delay_seconds=restart_delay_seconds,
             )
-        return run(cmd, check=True)
+        return cmd, "incremental"
 
     if mode == "explain":
         exp_args = runner.get("explain_args") or []
@@ -296,7 +247,7 @@ def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
                 max_restarts=max_restarts,
                 restart_delay_seconds=restart_delay_seconds,
             )
-        return run(cmd, check=True)
+        return cmd, "explain"
 
     if mode != "entry":
         raise ValueError(
@@ -337,4 +288,67 @@ def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
             max_restarts=max_restarts,
             restart_delay_seconds=restart_delay_seconds,
         )
+    return cmd, "entry"
+
+
+def run_from_config_cli(
+    description: str,
+    argv: Optional[Sequence[str]] = None,
+) -> subprocess.CompletedProcess:
+    parser = _build_config_parser(description)
+    args = parser.parse_args(argv)
+    return run_from_config(args.config_json)
+
+
+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:
+    """Convenience wrapper: build and run BayesOpt_entry (optional 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 entry point: switch execution modes by editing config.json.
+
+    Convention: config.json may include a `runner` section for notebook control:
+    - runner.mode: "entry" (default), "incremental", or "explain"
+    - runner.nproc_per_node: >1 enables torchrun/DDP (entry only)
+    - runner.model_keys: list of models to run (entry only)
+    - runner.max_evals / runner.plot_curves / runner.ft_role (entry only; override config fields)
+    - runner.use_watchdog / runner.idle_seconds / runner.max_restarts / runner.restart_delay_seconds
+    - runner.incremental_args: List[str] (incremental only; extra args for cli/BayesOpt_incremental.py)
+    """
+    cmd, _mode = build_cmd_from_config(config_json)
     return run(cmd, check=True)