code/hfapi.py

"""IPv4-forced Hugging Face API helpers (the local network stalls on IPv6).

Used for bucket creation/upload/download and Jobs submission/polling.
"""
import socket as _socket

_orig_gai = _socket.getaddrinfo


def _ipv4_gai(host, port, family=0, type=0, proto=0, flags=0):
    return _orig_gai(host, port, _socket.AF_INET, type, proto, flags)


_socket.getaddrinfo = _ipv4_gai

import json
import sys
import time

from huggingface_hub import HfApi, get_token

api = HfApi()
USER = "visv-Bro"
BUCKET = f"{USER}/bot-entucb-reproduction"


def _http(method, url, **kw):
    import httpx
    tok = get_token()
    headers = kw.pop("headers", {})
    headers["Authorization"] = f"Bearer {tok}"
    r = httpx.request(method, url, headers=headers, timeout=60, **kw)
    return r


def create_bucket():
    r = _http("POST", "https://huggingface.co/api/repos/create",
              json={"name": BUCKET.split("/")[1], "type": "bucket", "private": False})
    print(r.status_code, r.text[:200])


def upload(local_path, path_in_repo):
    api.upload_file(path_or_fileobj=local_path, path_in_repo=path_in_repo,
                    repo_id=BUCKET, repo_type="bucket")
    print(f"uploaded {path_in_repo}")


def download(path_in_repo, local_path):
    from huggingface_hub import hf_hub_download
    import shutil
    p = hf_hub_download(repo_id=BUCKET, repo_type="bucket", filename=path_in_repo)
    shutil.copy(p, local_path)
    print(f"downloaded {path_in_repo} -> {local_path}")


def submit_job(command, flavor="cpu-upgrade", timeout_s=3600, image="ghcr.io/astral-sh/uv:python3.12-bookworm",
               bucket_mount=True):
    payload = {
        "dockerImage": image,
        "command": command,
        "flavor": flavor,
        "timeout": timeout_s,
        "arguments": [],
        "environment": {},
    }
    if bucket_mount:
        payload["volumes"] = [{"type": "bucket", "source": BUCKET,
                               "mountPath": "/bucket", "readOnly": False}]
    r = _http("POST", f"https://huggingface.co/api/jobs/{USER}", json=payload)
    d = r.json()
    print(json.dumps({k: d.get(k) for k in ("id", "flavor", "status")}, default=str))
    return d.get("id")


def job_status(job_id):
    r = _http("GET", f"https://huggingface.co/api/jobs/{USER}/{job_id}")
    d = r.json()
    return d.get("status", {})


def job_logs(job_id, tail=40):
    r = _http("GET", f"https://huggingface.co/api/jobs/{USER}/{job_id}/logs-stream")
    if r.status_code != 200:
        r = _http("GET", f"https://huggingface.co/api/jobs/{USER}/{job_id}/logs")
    lines = r.text.splitlines()[-tail:]
    return "\n".join(lines)


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "create-bucket":
        create_bucket()
    elif cmd == "upload":
        upload(sys.argv[2], sys.argv[3])
    elif cmd == "download":
        download(sys.argv[2], sys.argv[3])
    elif cmd == "submit":
        submit_job(json.loads(sys.argv[2]), *sys.argv[3:4])
    elif cmd == "status":
        print(job_status(sys.argv[2]))
    elif cmd == "logs":
        print(job_logs(sys.argv[2]))