File size: 3,585 Bytes
99297d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """vimeml-tiny-zh-v1 arm pretraining entry point; W&B without any machine-wide login.
Credentials come only from the environment (WANDB_API_KEY, optional WANDB_ENTITY),
as cluster_zh13.sh sets them from a private file with isolated W&B directories.
Without a key, or when online initialization fails, the run logs offline under
artifacts/tracking/<model>/ for a later `wandb sync`; training never waits on W&B.
"""
import argparse
import os
from pathlib import Path
import sys
import tomllib
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", type=Path, required=True)
parser.add_argument("--resume", action="store_true")
parser.add_argument("--dry-run", action="store_true", help="Verify configuration, code and data identities, then exit")
args = parser.parse_args()
config = tomllib.loads(args.config.read_text("utf-8"))
output, logs = ROOT / config["output_dir"], ROOT / config["log_dir"]
if args.dry_run:
from vimeml.training.joint_zh13 import build_plan
plan = build_plan(config, ROOT, verify_files=True)
s = plan["settings"]
print({"stage": plan["stage"], "arm": plan["arm"], "parameters": plan["parameters"],
"max_steps": s["max_steps"], "epochs": s["epochs"], "milestone_steps": s["milestone_steps"][:3] + ["..."],
"budget": plan["budget"], "required_files": len(plan["required_files"])})
return
if not args.resume and ((output.exists() and any(output.iterdir())) or (logs.exists() and any(logs.iterdir()))):
parser.error("Fresh output/log directories must be empty")
if args.resume and not (output / "last.pt").exists():
parser.error("No checkpoint to resume")
from vimeml.training.joint_zh13 import run
tracking = config.get("tracking", {})
if not tracking.get("enabled"):
print(run(config, args.resume), flush=True)
return
import wandb
from vimeml.training.data import write_json
directory = ROOT / "artifacts/tracking" / output.name
directory.mkdir(parents=True, exist_ok=True)
def start(mode):
return wandb.init(project=tracking.get("project", "vimeml"), entity=os.environ.get("WANDB_ENTITY"),
group=tracking["group"], name=config["arm"] + ("-resume" if args.resume else ""),
job_type="pretraining", tags=["zh-v1.3", config["arm"]], mode=mode, dir=str(directory),
sync_tensorboard=True, save_code=False, config=config,
settings=wandb.Settings(disable_git=True, init_timeout=90))
mode = tracking.get("mode", "online") if os.environ.get("WANDB_API_KEY") else "offline"
try:
cloud = start(mode)
except Exception as error:
if mode == "offline":
raise
print(f"W&B online initialization failed ({type(error).__name__}: {error}); logging offline", flush=True)
mode = "offline"
cloud = start(mode)
with cloud:
write_json(directory / f"live-{cloud.id}.json", {"id": cloud.id, "url": cloud.url if mode != "offline" else None,
"mode": mode, "resume": args.resume, "directory": str(directory)})
print(f"W&B run ({mode}): {cloud.url if mode != 'offline' else directory}", flush=True)
result = run(config, args.resume)
cloud.summary.update({"training_result": result})
if __name__ == "__main__":
main()
|