File size: 6,969 Bytes
6a1771b 4ea7c69 5eb23eb 6a1771b 4ea7c69 6a1771b | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | #!/usr/bin/env python3
"""Upload scratch/gandalf staging files to Hugging Face.
Runs on the login node (python3 + huggingface_hub). The feanor sidecar only
needs to land files in --stage-dir; this process does the HTTP/LFS upload.
"""
import argparse
import json
import os
import sys
import time
import traceback
def _load_json(path):
try:
with open(path) as f:
return json.load(f)
except Exception:
return {}
def _write_json(path, payload):
tmp = path + ".tmp"
with open(tmp, "w") as f:
json.dump(payload, f, indent=2, sort_keys=True)
f.write("\n")
os.replace(tmp, path)
_KEEP_HB = (
"event", "step", "stage", "loss",
"val_acc", "val_mass", "val_excess", "val_spread",
"val_mass_multi", "val_excess_multi", "val_out_multi", "val_kl_multi",
"val_hbound", "val_hgap", "val_hbound_puz",
"val_ceq", "val_hq", "val_ceq_ok", "val_ceq_gap",
"loc_wave", "loc_coverage", "loc_lcs", "loc_dup",
"puzzle_acc", "last_ckpt_step",
)
def _slim_heartbeat(src, dest):
raw = _load_json(src)
slim = {k: raw[k] for k in _KEEP_HB if k in raw}
_write_json(dest, slim)
return dest
def _slim_log(src, dest):
"""Keep only eval / promote lines: step, stage, loss, accuracies."""
keep = []
try:
with open(src, errors="replace") as f:
for line in f:
s = line.strip()
if not s:
continue
if s.startswith("[curriculum]"):
keep.append(s + "\n")
continue
# "2000 stage 1 loss 1.47 loc_acc 0.03 val_acc 0.11 inset 0.66"
if " loss " in s and "stage" in s and not s.startswith("I"):
keep.append(s + "\n")
except OSError:
return None
with open(dest, "w") as f:
f.writelines(keep)
return dest
def _upload(api, token, local, repo, dest):
api.upload_file(
path_or_fileobj=local,
path_in_repo=dest,
repo_id=repo,
repo_type="model",
token=token,
commit_message=f"sync {dest}",
)
print(f"[hf_push] uploaded {dest} ({os.path.getsize(local)} bytes)", flush=True)
def sync_once(args, state, api, token):
prefix = args.prefix.strip("/")
repo = args.repo
for local, dest, key in (
(args.heartbeat, f"{prefix}/heartbeat.json", "heartbeat"),
(args.log, f"{prefix}/train.log", "log"),
):
if not local or not os.path.isfile(local):
continue
mtime = os.path.getmtime(local)
size = os.path.getsize(local)
last = state.get(key, {})
min_age = 0 if key == "heartbeat" else 300
due = (time.time() - float(last.get("uploaded_at", 0))) >= min_age
if due and (mtime > float(last.get("mtime", 0)) or size != last.get("size", -1)):
upload_path = local
if key == "heartbeat":
upload_path = _slim_heartbeat(
local, os.path.join(os.path.dirname(local) or ".",
"_hf_heartbeat_slim.json"))
elif key == "log":
upload_path = _slim_log(
local, os.path.join(os.path.dirname(local) or ".",
"_hf_train_slim.log"))
if upload_path:
_upload(api, token, upload_path, repo, dest)
state[key] = {"mtime": mtime, "size": size, "uploaded_at": time.time()}
if args.stage_dir and os.path.isdir(args.stage_dir):
for name in sorted(os.listdir(args.stage_dir)):
if name in ("hf_push_state.json", "heartbeat.json"):
continue
if not name.endswith(".tar") and name not in ("latest.json",):
continue
local = os.path.join(args.stage_dir, name)
if not os.path.isfile(local):
continue
# step_/stage_ names are unique per checkpoint, so once uploaded
# they never need revisiting. latest.tar/latest.json keep the same
# name and are rewritten every save, so skipping them by name would
# freeze the rolling checkpoint at whichever step happened to land
# first; track those by mtime instead.
rolling = name in ("latest.tar", "latest.json")
mtime = os.path.getmtime(local)
if rolling:
if mtime <= float(state.get("rolling", {}).get(name, 0)):
continue
elif name in state.get("staged", []):
continue
dest = f"{prefix}/{name}"
if name.startswith("stage"):
dest = f"{prefix}/stages/{name}"
elif name.startswith("step_"):
dest = f"{prefix}/steps/{name}"
_upload(api, token, local, repo, dest)
if rolling:
state.setdefault("rolling", {})[name] = mtime
continue # keep the file; the sidecar overwrites it in place
state.setdefault("staged", []).append(name)
# Drop the local copy after a successful upload to save quota.
try:
os.remove(local)
except OSError:
pass
return state
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="Avra98/Sudoku_superposition")
ap.add_argument("--prefix", default="runs/w12_inst_latent")
ap.add_argument("--heartbeat", default="")
ap.add_argument("--log", default="")
ap.add_argument("--stage-dir", default="")
ap.add_argument("--token-file", default="/scratch/users/gatmiry/.hf_token")
ap.add_argument("--state", default="")
ap.add_argument("--interval", type=int, default=60)
ap.add_argument("--once", action="store_true")
args = ap.parse_args()
token = ""
if args.token_file and os.path.isfile(args.token_file):
with open(args.token_file) as f:
token = f.read().strip()
token = token or os.environ.get("HF_TOKEN", "")
if not token:
print("[hf_push] no token", file=sys.stderr)
return 1
from huggingface_hub import HfApi
api = HfApi(token=token)
state_path = args.state or os.path.join(
args.stage_dir or "/tmp/sudoku_hf_uploads", "hf_push_state.json")
os.makedirs(os.path.dirname(state_path) or ".", exist_ok=True)
state = _load_json(state_path)
print(f"[hf_push] watching heartbeat={args.heartbeat} log={args.log} "
f"stage={args.stage_dir}", flush=True)
while True:
try:
state = sync_once(args, state, api, token)
_write_json(state_path, state)
except Exception:
traceback.print_exc()
print("[hf_push] cycle failed; will retry", flush=True)
if args.once:
break
time.sleep(max(5, args.interval))
return 0
if __name__ == "__main__":
sys.exit(main())
|