"""Dependency-free training client; can also be downloaded as checkpoint_client.py. Configuration: CHECKPOINT_URL, CHECKPOINT_TOKEN or CHECKPOINT_TOKEN_FILE. Checkpoints are uploaded as opaque files. Save training state atomically before uploading. """ import argparse import hashlib import json import os from pathlib import Path import sys import time import urllib.error import urllib.parse import urllib.request class ApiError(RuntimeError): def __init__(self, status, detail): self.status = status super().__init__(f"HTTP {status}: {detail}") def file_hash(path): digest = hashlib.sha256() with open(path, "rb") as file: for chunk in iter(lambda: file.read(4 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() class CheckpointClient: def __init__(self, url=None, token=None, token_file=None): self.url = (url or os.getenv("CHECKPOINT_URL", "https://checkpoint.azuresky.icu")).rstrip("/") self.token = token or os.getenv("CHECKPOINT_TOKEN", "") token_file = token_file or os.getenv("CHECKPOINT_TOKEN_FILE") if not self.token and token_file: text = Path(token_file).read_text(encoding="utf-8-sig").strip() self.token = json.loads(text)["token"] if text.startswith("{") else text if not self.token: raise ValueError("Set CHECKPOINT_TOKEN or CHECKPOINT_TOKEN_FILE") if not self.url.startswith("https://") and not self.url.startswith(("http://127.0.0.1:", "http://localhost:")): raise ValueError("Remote API requires HTTPS") def request(self, method, path, data=None, raw=None): if not path.startswith("/") or path.startswith("//"): raise ValueError("API path must be relative to this server") headers = {"Authorization": "Bearer " + self.token} body = raw if data is not None: body = json.dumps(data, allow_nan=False).encode() headers["Content-Type"] = "application/json" elif raw is not None: headers["Content-Type"] = "application/octet-stream" req = urllib.request.Request(self.url + path, data=body, method=method, headers=headers) try: with urllib.request.urlopen(req, timeout=180) as response: return json.load(response) except urllib.error.HTTPError as error: try: detail = json.load(error).get("detail", "Request failed") except (ValueError, AttributeError): detail = "Request failed" raise ApiError(error.code, detail) from None def create_run(self, project, name, external_id, **metadata): return self.request("POST", "/api/runs", {"project": project, "name": name, "external_id": external_id, **metadata}) def report(self, run_id, step, metrics=None, **progress): return self.request("POST", f"/api/runs/{run_id}/progress", {"step": step, "metrics": metrics or {}, **progress}) def upload(self, run_id, path, step, epoch=None, metrics=None, note="", best=False, pinned=False, on_progress=None): path = Path(path) before = path.stat() digest = file_hash(path) if (before.st_size, before.st_mtime_ns) != (path.stat().st_size, path.stat().st_mtime_ns): raise ValueError("Checkpoint changed while hashing; save it atomically before upload") name_key = hashlib.sha256(path.name.encode()).hexdigest()[:24] meta = {"run_id": run_id, "upload_key": f"{digest}:{step}:{name_key}", "filename": path.name, "size": before.st_size, "sha256": digest, "step": step, "epoch": epoch, "metrics": metrics or {}, "note": note} upload = self.request("POST", "/api/uploads", meta) identity, retries = upload["id"], 0 with path.open("rb") as stream: while upload["offset"] < before.st_size and upload["state"] != "ready": offset = upload["offset"] stream.seek(offset) chunk = stream.read(8 * 1024 * 1024) try: result = self.request("PUT", f"/api/uploads/{identity}?offset={offset}", raw=chunk) upload["offset"] = result["offset"] retries = 0 if on_progress: on_progress(result["offset"], before.st_size) except (OSError, ApiError) as error: if isinstance(error, ApiError) and error.status not in (409, 500, 502, 503, 504): raise retries += 1 if retries > 4: raise time.sleep(min(2 ** (retries - 1), 4)) upload = self.request("GET", f"/api/uploads/{identity}") if (before.st_size, before.st_mtime_ns) != (path.stat().st_size, path.stat().st_mtime_ns): raise ValueError("Checkpoint changed during upload; upload was not finalized") result = self.request("POST", f"/api/uploads/{identity}/complete") if best or pinned: result = self.request("PATCH", f"/api/checkpoints/{identity}", {"best": best, "pinned": pinned}) return result def download(self, checkpoint_id, destination): destination = Path(destination) if destination.exists(): raise FileExistsError("Destination already exists; choose a new filename") a = self.request("GET", f"/api/uploads/{checkpoint_id}") if a["state"] != "ready": raise ValueError("Checkpoint is incomplete") part = destination.with_name(destination.name + ".part") offset = part.stat().st_size if part.exists() else 0 if offset > a["size"]: raise ValueError("Existing partial file is too large") if offset < a["size"]: headers = {"Authorization": "Bearer " + self.token} if offset: headers["Range"] = f"bytes={offset}-" req = urllib.request.Request(self.url + f"/api/checkpoints/{checkpoint_id}/download", headers=headers) with urllib.request.urlopen(req, timeout=180) as response: if offset and (response.status != 206 or not response.headers.get("Content-Range", "").startswith(f"bytes {offset}-")): raise ValueError("Server did not honor resume range; partial file was retained") with part.open("ab") as out: for chunk in iter(lambda: response.read(4 * 1024 * 1024), b""): out.write(chunk) out.flush() os.fsync(out.fileno()) if part.stat().st_size != a["size"] or file_hash(part) != a["sha256"]: raise ValueError("Downloaded checksum mismatch; partial file retained for inspection") os.replace(part, destination) return {"id": checkpoint_id, "path": str(destination), "size": a["size"], "sha256": a["sha256"]} def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--url") parser.add_argument("--token-file", help="Private text/JSON token file; prefer this or an environment variable over shell history") commands = parser.add_subparsers(dest="command", required=True) run = commands.add_parser("run", help="Create or resume run metadata") run.add_argument("--project", required=True) run.add_argument("--name", required=True) run.add_argument("--key", required=True) run.add_argument("--hardware", default="") run.add_argument("--config", default="{}", help="JSON object") run.add_argument("--provenance", default="{}", help="JSON object of version/hash references") progress = commands.add_parser("progress") progress.add_argument("run_id") progress.add_argument("--step", type=int, required=True) progress.add_argument("--total-steps", type=int) progress.add_argument("--epoch", type=float) progress.add_argument("--metrics", default="{}") progress.add_argument("--status", choices=["running", "paused", "completed", "failed"], default="running") upload = commands.add_parser("upload") upload.add_argument("run_id") upload.add_argument("file") upload.add_argument("--step", required=True, type=int) upload.add_argument("--epoch", type=float) upload.add_argument("--metrics", default="{}") upload.add_argument("--best", action="store_true") upload.add_argument("--pin", action="store_true") download = commands.add_parser("download") download.add_argument("checkpoint_id") download.add_argument("destination") request = commands.add_parser("request", help="Any documented API action; prints JSON for agents") request.add_argument("method", choices=["GET", "POST", "PATCH", "DELETE"]) request.add_argument("path") request.add_argument("--json", default=None) args = parser.parse_args() try: client = CheckpointClient(args.url, token_file=args.token_file) if args.command == "run": result = client.create_run(args.project, args.name, args.key, hardware=args.hardware, config=json.loads(args.config), provenance=json.loads(args.provenance)) elif args.command == "progress": result = client.report(args.run_id, args.step, json.loads(args.metrics), total_steps=args.total_steps, epoch=args.epoch, status=args.status) elif args.command == "upload": result = client.upload(args.run_id, args.file, args.step, args.epoch, json.loads(args.metrics), best=args.best, pinned=args.pin) elif args.command == "download": result = client.download(args.checkpoint_id, args.destination) else: result = client.request(args.method, args.path, json.loads(args.json) if args.json is not None else None) print(json.dumps(result, ensure_ascii=False)) except (ValueError, OSError, ApiError) as error: print(json.dumps({"error": str(error)}, ensure_ascii=False), file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()