prepare_deeph_resume_config.py

#!/usr/bin/env python3
"""Prepare a DeepH training config that resumes from the latest output.

This helper is used by slurm_train_deeph_soc_243_gpu.sh. It prints one line:
  CONFIG <path>       when training should run with the generated config
  COMPLETE <path>     when the latest output already reached DeepH test stage
"""

from __future__ import annotations

import argparse
import re
import sys
import tomllib
from pathlib import Path


SECTION_RE = re.compile(r"^\s*\[([^\]]+)\]\s*$")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-config", required=True)
    parser.add_argument("--generated-config", required=True)
    return parser.parse_args()


def latest_output_dir(outputs_dir: Path) -> Path | None:
    if not outputs_dir.is_dir():
        return None
    candidates = [
        path
        for path in outputs_dir.iterdir()
        if path.is_dir() and (path / "deepx.log").is_file() and (path / "model").is_dir()
    ]
    if not candidates:
        return None
    return max(candidates, key=lambda path: path.name)


def completed(output_dir: Path) -> bool:
    log_file = output_dir / "deepx.log"
    if not log_file.is_file():
        return False
    text = log_file.read_text(errors="replace")
    return "[test] Use-Best-Epoch" in text


def has_latest_checkpoint(output_dir: Path) -> bool:
    params = output_dir / "model" / "params" / "latest.pytree"
    states = output_dir / "model" / "states" / "latest.pytree"
    if not params.is_dir() or not states.is_dir():
        return False
    return any(params.glob("epoch_*")) and any(states.glob("epoch_*"))


def toml_value(value) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, str):
        return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"'
    return str(value)


def set_key(text: str, section: str, key: str, value) -> str:
    lines = text.splitlines()
    current_section = None
    replaced = False
    rendered = f"{key} = {toml_value(value)}"

    for idx, line in enumerate(lines):
        match = SECTION_RE.match(line)
        if match:
            current_section = match.group(1)
            continue
        if current_section == section and re.match(rf"^\s*{re.escape(key)}\s*=", line):
            lines[idx] = rendered
            replaced = True
            break

    if not replaced:
        raise KeyError(f"Missing key [{section}].{key} in base config")
    return "\n".join(lines) + "\n"


def main() -> int:
    args = parse_args()
    base_config = Path(args.base_config).resolve()
    generated_config = Path(args.generated_config).resolve()

    base_text = base_config.read_text()
    config = tomllib.loads(base_text)
    outputs_dir = Path(config["data"]["outputs_dir"]).resolve()

    previous = latest_output_dir(outputs_dir)
    if previous is not None and completed(previous):
        print(f"COMPLETE {previous}")
        return 0

    generated_text = base_text
    if previous is not None and has_latest_checkpoint(previous):
        generated_text = set_key(generated_text, "process.train.continued", "enable", True)
        generated_text = set_key(generated_text, "process.train.continued", "new_training_data", False)
        generated_text = set_key(generated_text, "process.train.continued", "new_optimizer", False)
        generated_text = set_key(
            generated_text,
            "process.train.continued",
            "previous_output_dir",
            str(previous),
        )
        generated_text = set_key(generated_text, "process.train.continued", "load_model_type", "latest")
        generated_text = set_key(generated_text, "process.train.continued", "load_model_epoch", -1)

        split_file = previous / "dataset_split.json"
        if split_file.is_file():
            generated_text = set_key(
                generated_text,
                "process.train.dataloader",
                "dataset_split_json",
                str(split_file),
            )
        print(f"[resume] previous output: {previous}", file=sys.stderr)
    else:
        generated_text = set_key(generated_text, "process.train.continued", "enable", False)
        generated_text = set_key(generated_text, "process.train.continued", "previous_output_dir", "")
        generated_text = set_key(generated_text, "process.train.dataloader", "dataset_split_json", "")
        if previous is None:
            print("[resume] no previous output; starting fresh", file=sys.stderr)
        else:
            print(
                f"[resume] previous output has no latest checkpoint; starting fresh: {previous}",
                file=sys.stderr,
            )

    generated_config.write_text(generated_text)
    print(f"CONFIG {generated_config}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())