scripts/diagnostics/hf_stack_exchange_diagnosis.py

#!/usr/bin/env python3
from __future__ import annotations

import importlib
import json
import os
import platform
import socket
import ssl
import sys
import time
import traceback
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any

DATASET_ID = "HuggingFaceH4/stack-exchange-preferences"
API_URL = f"https://huggingface.co/api/datasets/{DATASET_ID}"
BASE_URL = f"https://huggingface.co/datasets/{DATASET_ID}"
TIMEOUT = 30


def safe_version(package: str) -> str:
    try:
        module = importlib.import_module(package)
        return str(getattr(module, "__version__", "unknown"))
    except Exception as exc:
        return f"IMPORT_ERROR: {type(exc).__name__}: {exc}"


def record(
    results: list[dict[str, Any]],
    test: str,
    status: str,
    detail: Any,
) -> None:
    row = {
        "test": test,
        "status": status,
        "detail": detail,
    }
    results.append(row)
    print(
        f"[{status}] {test}: "
        f"{json.dumps(detail, ensure_ascii=False, default=str)}",
        flush=True,
    )


def masked_environment() -> dict[str, Any]:
    names = [
        "HF_HUB_OFFLINE",
        "HF_DATASETS_OFFLINE",
        "TRANSFORMERS_OFFLINE",
        "HF_ENDPOINT",
        "HF_HOME",
        "HF_HUB_CACHE",
        "HF_DATASETS_CACHE",
        "HTTP_PROXY",
        "HTTPS_PROXY",
        "ALL_PROXY",
        "NO_PROXY",
        "REQUESTS_CA_BUNDLE",
        "CURL_CA_BUNDLE",
        "SSL_CERT_FILE",
    ]
    output: dict[str, Any] = {}
    for name in names:
        value = os.environ.get(name)
        if value is None:
            output[name] = "UNSET"
        elif "PROXY" in name and "@" in value:
            prefix, suffix = value.rsplit("@", 1)
            scheme = prefix.split("://", 1)[0]
            output[name] = f"{scheme}://***@{suffix}"
        else:
            output[name] = value
    output["HF_TOKEN"] = "SET" if os.environ.get("HF_TOKEN") else "UNSET"
    output["HUGGING_FACE_HUB_TOKEN"] = (
        "SET" if os.environ.get("HUGGING_FACE_HUB_TOKEN") else "UNSET"
    )
    return output


def url_request(
    url: str,
    method: str = "GET",
) -> dict[str, Any]:
    request = urllib.request.Request(
        url,
        method=method,
        headers={
            "User-Agent": "DGX-AI-Factory-HF-Diagnosis/1.0",
            "Accept": "application/json,text/plain,*/*",
        },
    )
    started = time.time()
    with urllib.request.urlopen(request, timeout=TIMEOUT) as response:
        body = response.read(4096)
        return {
            "url": url,
            "status_code": response.status,
            "elapsed_sec": round(time.time() - started, 3),
            "content_type": response.headers.get("content-type"),
            "content_length": response.headers.get("content-length"),
            "sample": body.decode("utf-8", errors="replace")[:600],
        }


def main(report_path: Path) -> int:
    results: list[dict[str, Any]] = []

    print("==== STAGE 1: RUNTIME AND ENVIRONMENT ====", flush=True)
    record(
        results,
        "runtime",
        "PASS",
        {
            "python": sys.version,
            "executable": sys.executable,
            "platform": platform.platform(),
            "machine": platform.machine(),
            "openssl": ssl.OPENSSL_VERSION,
        },
    )
    record(
        results,
        "package_versions",
        "PASS",
        {
            "datasets": safe_version("datasets"),
            "huggingface_hub": safe_version("huggingface_hub"),
            "fsspec": safe_version("fsspec"),
            "requests": safe_version("requests"),
            "pyarrow": safe_version("pyarrow"),
        },
    )

    env = masked_environment()
    offline_names = [
        "HF_HUB_OFFLINE",
        "HF_DATASETS_OFFLINE",
        "TRANSFORMERS_OFFLINE",
    ]
    offline_active = any(
        str(env.get(name, "")).strip().lower()
        in {"1", "true", "yes", "on"}
        for name in offline_names
    )
    record(
        results,
        "environment",
        "FAIL" if offline_active else "PASS",
        {
            **env,
            "offline_mode_detected": offline_active,
        },
    )

    print("==== STAGE 2: DNS AND HTTPS ====", flush=True)
    try:
        addresses = sorted({
            item[4][0]
            for item in socket.getaddrinfo(
                "huggingface.co",
                443,
                type=socket.SOCK_STREAM,
            )
        })
        record(
            results,
            "dns_huggingface_co",
            "PASS",
            {"addresses": addresses},
        )
    except Exception as exc:
        record(
            results,
            "dns_huggingface_co",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
            },
        )

    for test_name, url in (
        ("https_huggingface_root", "https://huggingface.co"),
        ("https_dataset_page", BASE_URL),
        ("https_dataset_api", API_URL),
    ):
        try:
            record(
                results,
                test_name,
                "PASS",
                url_request(url),
            )
        except urllib.error.HTTPError as exc:
            body = exc.read(1000).decode("utf-8", errors="replace")
            record(
                results,
                test_name,
                "FAIL",
                {
                    "status_code": exc.code,
                    "reason": str(exc.reason),
                    "body": body,
                    "headers": dict(exc.headers),
                },
            )
        except Exception as exc:
            record(
                results,
                test_name,
                "FAIL",
                {
                    "error_type": type(exc).__name__,
                    "error": str(exc),
                },
            )

    print("==== STAGE 3: HUGGING FACE HUB CLIENT ====", flush=True)
    api_info: Any = None
    try:
        from huggingface_hub import HfApi

        api = HfApi(
            endpoint=os.environ.get(
                "HF_ENDPOINT",
                "https://huggingface.co",
            )
        )
        info = api.dataset_info(
            DATASET_ID,
            timeout=TIMEOUT,
        )
        api_info = {
            "id": info.id,
            "private": info.private,
            "gated": getattr(info, "gated", None),
            "sha": info.sha,
            "last_modified": str(info.last_modified),
            "siblings": len(info.siblings or []),
        }
        record(
            results,
            "hf_api_dataset_info",
            "PASS",
            api_info,
        )
    except Exception as exc:
        record(
            results,
            "hf_api_dataset_info",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
                "traceback": traceback.format_exc(limit=5),
            },
        )

    print("==== STAGE 4: DATASETS LIBRARY RESOLUTION ====", flush=True)
    try:
        from datasets import get_dataset_config_names

        configs = get_dataset_config_names(
            DATASET_ID,
            trust_remote_code=False,
        )
        record(
            results,
            "datasets_config_names",
            "PASS",
            {
                "count": len(configs),
                "sample": configs[:20],
            },
        )
    except Exception as exc:
        record(
            results,
            "datasets_config_names",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
                "traceback": traceback.format_exc(limit=6),
            },
        )

    try:
        from datasets import load_dataset_builder

        builder = load_dataset_builder(
            DATASET_ID,
            trust_remote_code=False,
        )
        record(
            results,
            "load_dataset_builder",
            "PASS",
            {
                "builder_name": builder.info.builder_name,
                "config_name": builder.config.name,
                "features": str(builder.info.features),
                "splits": str(builder.info.splits),
            },
        )
    except Exception as exc:
        record(
            results,
            "load_dataset_builder",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
                "traceback": traceback.format_exc(limit=6),
            },
        )

    print("==== STAGE 5: NORMAL STREAMING LOAD ====", flush=True)
    try:
        from datasets import load_dataset

        stream = load_dataset(
            DATASET_ID,
            split="train",
            streaming=True,
            trust_remote_code=False,
        )
        first = next(iter(stream))
        record(
            results,
            "load_dataset_streaming",
            "PASS",
            {
                "keys": sorted(first.keys()),
                "qid": first.get("qid"),
                "question_chars": len(str(first.get("question", ""))),
                "answer_count": len(first.get("answers") or []),
            },
        )
    except Exception as exc:
        record(
            results,
            "load_dataset_streaming",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
                "traceback": traceback.format_exc(limit=8),
            },
        )

    print("==== STAGE 6: DIRECT PARQUET FALLBACK TEST ====", flush=True)
    parquet_paths: list[str] = []
    try:
        if api_info is None:
            from huggingface_hub import HfApi

            info = HfApi().dataset_info(
                DATASET_ID,
                timeout=TIMEOUT,
            )
        else:
            from huggingface_hub import HfApi

            info = HfApi().dataset_info(
                DATASET_ID,
                timeout=TIMEOUT,
            )

        parquet_paths = [
            sibling.rfilename
            for sibling in (info.siblings or [])
            if sibling.rfilename.endswith(".parquet")
        ]
        preferred = [
            name
            for name in parquet_paths
            if "Stackoverflow.com" in name
            or "askubuntu" in name.lower()
            or "unix" in name.lower()
        ]
        selected = (preferred or parquet_paths)[:1]
        record(
            results,
            "parquet_file_discovery",
            "PASS" if selected else "FAIL",
            {
                "parquet_file_count": len(parquet_paths),
                "selected": selected,
            },
        )

        if selected:
            direct_url = (
                f"https://huggingface.co/datasets/{DATASET_ID}"
                f"/resolve/main/{selected[0]}?download=true"
            )
            try:
                direct_http = url_request(direct_url)
                direct_http["sample"] = "<binary parquet response omitted>"
                record(
                    results,
                    "direct_parquet_https",
                    "PASS",
                    direct_http,
                )
            except Exception as exc:
                record(
                    results,
                    "direct_parquet_https",
                    "FAIL",
                    {
                        "url": direct_url,
                        "error_type": type(exc).__name__,
                        "error": str(exc),
                    },
                )

            try:
                from datasets import load_dataset

                parquet_stream = load_dataset(
                    "parquet",
                    data_files={"train": direct_url},
                    split="train",
                    streaming=True,
                )
                first = next(iter(parquet_stream))
                record(
                    results,
                    "direct_parquet_streaming",
                    "PASS",
                    {
                        "url": direct_url,
                        "keys": sorted(first.keys()),
                        "qid": first.get("qid"),
                    },
                )
            except Exception as exc:
                record(
                    results,
                    "direct_parquet_streaming",
                    "FAIL",
                    {
                        "url": direct_url,
                        "error_type": type(exc).__name__,
                        "error": str(exc),
                        "traceback": traceback.format_exc(limit=8),
                    },
                )
    except Exception as exc:
        record(
            results,
            "parquet_file_discovery",
            "FAIL",
            {
                "error_type": type(exc).__name__,
                "error": str(exc),
                "traceback": traceback.format_exc(limit=6),
            },
        )

    statuses = {
        row["test"]: row["status"]
        for row in results
    }

    if offline_active:
        cause = "OFFLINE_ENVIRONMENT_VARIABLE_ACTIVE"
        next_action = "UNSET_OFFLINE_VARIABLES_AND_RETRY"
    elif statuses.get("dns_huggingface_co") == "FAIL":
        cause = "DNS_RESOLUTION_FAILURE"
        next_action = "REPAIR_DNS_OR_NETWORK"
    elif statuses.get("https_huggingface_root") == "FAIL":
        cause = "HTTPS_NETWORK_OR_TLS_FAILURE"
        next_action = "REPAIR_NETWORK_PROXY_OR_CA"
    elif statuses.get("https_dataset_api") == "FAIL":
        cause = "HUGGING_FACE_DATASET_API_BLOCKED"
        next_action = "CHECK_PROXY_FIREWALL_OR_HF_ENDPOINT"
    elif statuses.get("hf_api_dataset_info") == "FAIL":
        cause = "HUGGINGFACE_HUB_CLIENT_FAILURE"
        next_action = "REVIEW_HUB_CLIENT_ERROR_AND_VERSION"
    elif statuses.get("load_dataset_streaming") == "PASS":
        cause = "TRANSIENT_OR_ALREADY_RESOLVED"
        next_action = "RETRY_COLLECTION_WITH_CLEAN_ENV"
    elif statuses.get("direct_parquet_streaming") == "PASS":
        cause = "DATASETS_REPOSITORY_RESOLUTION_FAILURE"
        next_action = "USE_DIRECT_PARQUET_FALLBACK"
    elif statuses.get("load_dataset_builder") == "FAIL":
        cause = "DATASETS_LIBRARY_OR_REPOSITORY_LAYOUT_INCOMPATIBILITY"
        next_action = "PATCH_LOADER_OR_UPDATE_COMPATIBLE_PACKAGES"
    else:
        cause = "UNRESOLVED_HUB_ACCESS_FAILURE"
        next_action = "REVIEW_DIAGNOSTIC_REPORT"

    report = {
        "task": "HF_STACK_EXCHANGE_DIAGNOSIS",
        "dataset_id": DATASET_ID,
        "cause": cause,
        "next_action": next_action,
        "tests": results,
        "environment_modified": False,
        "packages_modified": False,
        "collection_started": False,
        "training_started": False,
        "stable_model_modified": False,
    }
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    print()
    print("==== FINAL RESULT ====")
    print("TASK=HF_STACK_EXCHANGE_DIAGNOSIS")
    print("EXECUTION_STATUS=PASS")
    print(f"SYSTEM_STATUS={cause}")
    print(f"DATASET_ID={DATASET_ID}")
    print(
        "OFFLINE_MODE_DETECTED="
        + str(offline_active)
    )
    print(
        "DNS_PASS="
        + str(statuses.get("dns_huggingface_co") == "PASS")
    )
    print(
        "HTTPS_ROOT_PASS="
        + str(statuses.get("https_huggingface_root") == "PASS")
    )
    print(
        "DATASET_API_PASS="
        + str(statuses.get("https_dataset_api") == "PASS")
    )
    print(
        "HF_API_PASS="
        + str(statuses.get("hf_api_dataset_info") == "PASS")
    )
    print(
        "LOAD_DATASET_STREAMING_PASS="
        + str(statuses.get("load_dataset_streaming") == "PASS")
    )
    print(
        "DIRECT_PARQUET_STREAMING_PASS="
        + str(statuses.get("direct_parquet_streaming") == "PASS")
    )
    print("ENVIRONMENT_MODIFIED=False")
    print("PACKAGES_MODIFIED=False")
    print("COLLECTION_STARTED=False")
    print("TRAINING_STARTED=False")
    print("STABLE_MODEL_MODIFIED=False")
    print(f"REPORT_FILE={report_path}")
    print(f"NEXT_ACTION={next_action}")
    print("==== END FINAL RESULT ====")
    return 0


if __name__ == "__main__":
    output = Path(sys.argv[1]).expanduser().resolve()
    node-7.example.invalid SystemExit(main(output))