scripts/tail_training_job.py

#!/usr/bin/env python
"""Re-attach to an in-flight HF Jobs run and stream its logs.

Usage::

    $env:HF_TOKEN = "hf_..."
    python scripts/tail_training_job.py 69ec88dfd70108f37acde39d
"""
from __future__ import annotations

import os
import sys

from huggingface_hub import HfApi

from submit_training_job import tail_logs  # type: ignore[import-not-found]


def main() -> int:
    if len(sys.argv) < 2:
        print("usage: python scripts/tail_training_job.py <job_id> [namespace]", file=sys.stderr)
        return 2
    job_id = sys.argv[1]
    namespace = sys.argv[2] if len(sys.argv) > 2 else "akhiilll"
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("ERROR: set HF_TOKEN in the environment first.", file=sys.stderr)
        return 2
    api = HfApi()
    return tail_logs(api, token, job_id, namespace=namespace)


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