probes/p0_cpu_probe.py

# Phase 0 CPU probe. Runs on a Kaggle CPU instance (zero GPU quota).
# Purpose: establish empirically what the main run will depend on --
# in-job internet, pip behavior, image contents, disk/memory limits,
# Hub file download, and the artifact-retrieval path. Writes one artifact
# to /kaggle/working so retrieval can be verified by hash.
# This file contains NO credentials and must never gain any.

import hashlib
import json
import os
import platform
import shutil
import socket
import subprocess
import sys
import time

R = {}
T_START = time.time()
MONO0 = time.monotonic()


def step(name, fn):
    a, b = time.monotonic(), time.time()
    try:
        R[name] = fn()
    except Exception as e:
        R[name] = {"error": f"{type(e).__name__}: {e}"}
    R.setdefault("_timing", {})[name] = round(time.monotonic() - a, 2)
    R["_timing"][name + "_epoch"] = round(b, 0)


def run(cmd, timeout=90):
    p = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
    return {"rc": p.returncode, "out": p.stdout[-3000:], "err": p.stderr[-1500:]}


def get_http(url, timeout=25, nbytes=400):
    import urllib.request

    t = time.monotonic()
    req = urllib.request.Request(url, headers={"User-Agent": "ounce100m-probe"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        body = r.read(nbytes)
        return {
            "status": r.status,
            "ms": round((time.monotonic() - t) * 1000),
            "bytes_head": len(body),
            "final_url": r.url[:160],
        }


def pkg_versions():
    from importlib.metadata import PackageNotFoundError, version

    want = [
        "torch", "torchvision", "transformers", "tokenizers", "datasets", "evaluate",
        "accelerate", "trl", "numpy", "pandas", "pyarrow", "sentencepiece", "tiktoken",
        "protobuf", "einops", "flash-attn", "xformers", "deepspeed", "bitsandbytes",
        "huggingface-hub", "safetensors", "timm", "scipy", "sentence-transformers",
        "wandb", "liger-kernel", "causal-conv1d", "mamba-ssm", "unsloth", "peft",
    ]
    out = {}
    for n in want:
        try:
            out[n] = version(n)
        except PackageNotFoundError:
            pass
    return out


def cuda_info():
    import torch

    d = {
        "torch": torch.__version__,
        "cuda_built": torch.version.cuda,
        "cudnn": torch.backends.cudnn.version(),
        "mkl": torch.cuda.get_mkl_version() if hasattr(torch.cuda, "get_mkl_version") else None,
        "allow_tf32_default": torch.backends.cuda.matmul.allow_tf32,
        "arch_list": list(torch.cuda.get_arch_list()),
        "device_count": torch.cuda.device_count(),
        "cpu_threads": torch.get_num_threads(),
    }
    if torch.cuda.is_available():
        props = []
        for i in range(torch.cuda.device_count()):
            p = torch.cuda.get_device_properties(i)
            props.append({
                "i": i, "name": p.name, "cc": f"{p.major}.{p.minor}",
                "sm": p.multi_processor_count, "mem_GB": round(p.total_memory / 1024**3, 2),
                "bf16_capable_cc": p.major >= 8,
            })
        d["devices"] = props
    return d


def precision_probe():
    # Turing T4 = cc 7.5: fp16 tensor cores yes, bf16 no. Prove it rather than recall it.
    import torch

    out = {}
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    for dt in (torch.float16, torch.bfloat16, torch.float32):
        try:
            x = torch.randn(1024, 1024, device=dev, dtype=dt)
            y = (x @ x).float()
            torch.cuda.synchronize() if dev == "cuda" else None
            out[str(dt)] = {"ok": True, "checksum": round(float(y.abs().mean()), 4)}
        except Exception as e:
            out[str(dt)] = {"ok": False, "err": f"{type(e).__name__}: {e}"[:160]}
    return out


def net_checks():
    out = {}
    for host in ["huggingface.co", "cdn-lfs.huggingface.co", "pypi.org", "files.pythonhosted.org"]:
        try:
            infos = socket.getaddrinfo(host, 443)
            out[host] = sorted({i[4][0] for i in infos})[:3]
        except Exception as e:
            out[host] = f"DNSFAIL {type(e).__name__}: {e}"
    return out


def hub_download():
    url = "https://huggingface.co/datasets/wikimedia/wikipedia/resolve/main/README.md"
    d = get_http(url)
    try:
        import datasets

        t = time.monotonic()
        ds = datasets.load_dataset("wikimedia/wikipedia", "20231101.en",
                                  split="train", streaming=True)
        row = next(iter(ds))
        d["streaming"] = {
            "ok": True, "ms_to_first_row": round((time.monotonic() - t) * 1000),
            "keys": list(row.keys()), "title": str(row.get("title"))[:60],
        }
    except Exception as e:
        d["streaming"] = {"ok": False, "err": f"{type(e).__name__}: {e}"[:300]}
    return d


def pip_probe():
    # xxhash: small, not in the Kaggle image, useful later for shard hashing.
    t = time.monotonic()
    r = run("python -m pip install --quiet --disable-pip-version-check --no-input xxhash", timeout=240)
    r["install_seconds"] = round(time.monotonic() - t, 1)
    try:
        import xxhash

        r["import_ok"] = xxhash.xxh64(b"ounce100m").hexdigest()
    except Exception as e:
        r["import_ok"] = f"FAIL {e}"
    return r


def fs_probe():
    out = {}
    for p in ["/kaggle/working", "/kaggle/input", "/tmp", "/", "/dev/shm"]:
        try:
            u = shutil.disk_usage(p)
            out[p] = {"total_GB": round(u.total / 1024**3, 1), "free_GB": round(u.free / 1024**3, 1)}
        except Exception as e:
            out[p] = f"ERR {e}"
    with open("/proc/meminfo") as f:
        mm = {l.split(":")[0]: l.split()[0] for l in f if ":" in l}
    out["mem_MemTotal_MB"] = int(mm.get("MemTotal", 0)) // 1024
    out["mem_MemAvailable_MB"] = int(mm.get("MemAvailable", 0)) // 1024
    out["cgroup_v2_max"] = None
    for lp in ["/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"]:
        if os.path.exists(lp):
            out["cgroup_v2_max"] = open(lp).read().strip()
            break
    out["cpu_count_os"] = os.cpu_count()
    try:
        out["cpu_affinity"] = len(os.sched_getaffinity(0))
    except Exception:
        pass
    out["nproc_smi"] = run("nvidia-smi -L", timeout=30)
    out["df_h"] = run("df -h | head -20", timeout=30).get("out", "")
    return out


def artifact():
    payload = json.dumps({"probe": "p0", "t": T_START, "tokens": 12345}).encode()
    path = "ounce100m_p0_artifact.json"
    with open(path, "wb") as f:
        f.write(payload)
    return {
        "path": path, "abs": os.path.abspath(path), "size": len(payload),
        "sha256": hashlib.sha256(payload).hexdigest(),
        "cwd": os.getcwd(),
        "working_listing": sorted(os.listdir("/kaggle/working")),
    }


step("env", lambda: {
    "python": sys.version.split()[0], "platform": platform.platform(),
    "hostname": socket.gethostname(), "user": os.environ.get("USER"),
    "cwd": os.getcwd(), "start_epoch": T_START,
})
step("env_var_names", lambda: sorted(k for k in os.environ if any(
    s in k.upper() for s in ["HF_", "KAGGLE_", "HUGGING", "TOKEN", "SECRET", "CUDA", "TORCH"])))
step("packages", pkg_versions)
step("net_dns", net_checks)
step("http", lambda: {
    "hub_api": get_http("https://huggingface.co/api/models?limit=1"),
    "pypi_simple": get_http("https://pypi.org/simple/xxhash/"),
    "google": get_http("https://www.google.com/generate_204"),
})
step("hub_download", hub_download)
step("pip_install", pip_probe)
step("filesystem", fs_probe)
step("artifact", artifact)
if os.environ.get("CUDA_VISIBLE_DEVICES") != "None":
    step("cuda", cuda_info)
    step("precision", precision_probe)
else:
    step("cuda", lambda: {"skipped": "CUDA_VISIBLE_DEVICES=None -> CPU shape"})

R["_summary"] = {
    "total_seconds": round(time.monotonic() - MONO0, 1),
    "end_epoch": time.time(),
    "errors_seen": sorted(k for k, v in R.items() if isinstance(v, dict) and "error" in v),
}
print("PROBE_JSON_BEGIN")
print(json.dumps(R, indent=1, default=str))
print("PROBE_JSON_END")