denpex-mlflow 0.1.0__tar.gz

denpex_mlflow-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+*.egg-info/
+build/
+dist/
+.pytest_cache/

denpex_mlflow-0.1.0/PKG-INFO

@@ -0,0 +1,105 @@
+Metadata-Version: 2.5
+Name: denpex-mlflow
+Version: 0.1.0
+Summary: Attach Denpex GPU and training-failure diagnoses to your MLflow runs in one line.
+Project-URL: Homepage, https://denpex.com
+Project-URL: Documentation, https://denpex.com/docs/integrations
+Project-URL: API Reference, https://denpex.com/docs/api
+Project-URL: Python Agent & CLI, https://denpex.com/docs/agent
+Project-URL: VS Code Extension, https://denpex.com/vscode
+Project-URL: Open VSX Extension, https://denpex.com/openvsx
+Project-URL: GPU Failure Encyclopedia, https://denpex.com/failures
+Project-URL: Free GPU Debugging Tools, https://denpex.com/tools
+Project-URL: Changelog, https://denpex.com/changelog
+Project-URL: Repository, https://github.com/Denpex/gpu-failure-cheatsheet
+Project-URL: Issue Tracker, https://denpex.com/contact
+Author: Denpex
+License: MIT
+Keywords: cuda,denpex,diagnosis,gpu,mlflow,nccl,pytorch,training
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Requires-Dist: mlflow>=2.0
+Description-Content-Type: text/markdown
+
+# denpex-mlflow
+
+Attach **[Denpex](https://denpex.com)** GPU and training-failure diagnoses to your **MLflow** runs in one line.
+
+When a run crashes, Denpex reads the traceback (and optionally your log file), diagnoses
+the root cause (NCCL, Xid, CUDA, OOM, checkpoint, etc.), and writes the **root cause + fix** back
+onto the active MLflow run as tags and a `denpex_diagnosis.md` artifact, then re-raises
+your original exception so nothing is hidden.
+
+## Install
+```bash
+pip install denpex-mlflow
+```
+
+## Use
+```python
+import mlflow
+import denpex_mlflow as denpex
+
+with mlflow.start_run():
+    with denpex.watch(job_name="llama-7b-sft"):
+        train()          # if this raises, Denpex diagnoses it and tags the run
+```
+
+Decorator form:
+```python
+@denpex.watch_run("llama-7b-sft")
+def main():
+    train()
+```
+
+One-shot (no MLflow needed):
+```python
+diag = denpex.diagnose(open("train.err").read())
+print(diag.failure_type, diag.confidence, diag.action)
+```
+
+## Auth
+Set an API key (create one at **[denpex.com/dashboard](https://denpex.com/dashboard)** under API keys):
+```bash
+export DENPEX_API_KEY=dpx_...
+# optional, for self-hosted / staging:
+export DENPEX_API_URL=https://api.denpex.com
+```
+
+## What gets written to the run
+| MLflow tag | Value |
+|---|---|
+| `denpex.failure_type` | e.g. `NCCL_TIMEOUT`, `CUDA_OOM`, `GPU_XID_ECC` |
+| `denpex.confidence` | 0 to 100 |
+| `denpex.fix_summary` | the root cause (truncated) |
+| `denpex.diagnosis_id` | link back to [denpex.com/dashboard](https://denpex.com/dashboard) |
+| artifact `denpex_diagnosis.md` | full root cause + fix |
+
+## Notes
+- **Best-effort & non-masking:** if Denpex is unreachable, unauthenticated, or over the
+  free-plan quota, it logs one line to stderr and still re-raises your real error.
+- `KeyboardInterrupt` / `SystemExit` are passed through undiagnosed.
+- The diagnosis runs through the same `/api/diagnose` endpoint as the rest of Denpex, so
+  it also lands in your dashboard history and fires your configured alert channels.
+- Pass `logfile="train.err"` to `watch()` to include the tail of a log file alongside the
+  traceback.
+
+---
+
+## Official Denpex Ecosystem & Documentation
+
+| Resource | Description | Link |
+| :--- | :--- | :--- |
+| **Denpex Platform** | Deterministic GPU and ML training failure diagnostics | [denpex.com](https://denpex.com) |
+| **Documentation Hub** | Quickstart guides, architecture, and operational reference | [denpex.com/docs](https://denpex.com/docs) |
+| **Python Agent & CLI** | On-host monitoring, evidence capture, and zero-egress CLI | [denpex.com/docs/agent](https://denpex.com/docs/agent) |
+| **REST API Reference** | Programmatic `/api/diagnose` and preflight endpoints | [denpex.com/docs/api](https://denpex.com/docs/api) |
+| **VS Code Extension** | Offline GPU crash diagnostics in VS Code | [denpex.com/vscode](https://denpex.com/vscode) |
+| **Open VSX Extension** | Extension for Cursor, Windsurf, VSCodium, and Gitpod | [denpex.com/openvsx](https://denpex.com/openvsx) |
+| **Failure Encyclopedia** | 500+ documented CUDA, NCCL, Xid, Slurm, and vLLM failures | [denpex.com/failures](https://denpex.com/failures) |
+| **Free GPU Tools** | VRAM, NCCL, DeepSpeed ZeRO, and Slurm calculators | [denpex.com/tools](https://denpex.com/tools) |
+| **NVIDIA Xid Reference** | Hardware vs software Xid fault routing table | [denpex.com/xid](https://denpex.com/xid) |
+| **Release Changelog** | Platform, agent, SDK, and extension release notes | [denpex.com/changelog](https://denpex.com/changelog) |
+

denpex_mlflow-0.1.0/README.md

@@ -0,0 +1,80 @@
+# denpex-mlflow
+
+Attach **[Denpex](https://denpex.com)** GPU and training-failure diagnoses to your **MLflow** runs in one line.
+
+When a run crashes, Denpex reads the traceback (and optionally your log file), diagnoses
+the root cause (NCCL, Xid, CUDA, OOM, checkpoint, etc.), and writes the **root cause + fix** back
+onto the active MLflow run as tags and a `denpex_diagnosis.md` artifact, then re-raises
+your original exception so nothing is hidden.
+
+## Install
+```bash
+pip install denpex-mlflow
+```
+
+## Use
+```python
+import mlflow
+import denpex_mlflow as denpex
+
+with mlflow.start_run():
+    with denpex.watch(job_name="llama-7b-sft"):
+        train()          # if this raises, Denpex diagnoses it and tags the run
+```
+
+Decorator form:
+```python
+@denpex.watch_run("llama-7b-sft")
+def main():
+    train()
+```
+
+One-shot (no MLflow needed):
+```python
+diag = denpex.diagnose(open("train.err").read())
+print(diag.failure_type, diag.confidence, diag.action)
+```
+
+## Auth
+Set an API key (create one at **[denpex.com/dashboard](https://denpex.com/dashboard)** under API keys):
+```bash
+export DENPEX_API_KEY=dpx_...
+# optional, for self-hosted / staging:
+export DENPEX_API_URL=https://api.denpex.com
+```
+
+## What gets written to the run
+| MLflow tag | Value |
+|---|---|
+| `denpex.failure_type` | e.g. `NCCL_TIMEOUT`, `CUDA_OOM`, `GPU_XID_ECC` |
+| `denpex.confidence` | 0 to 100 |
+| `denpex.fix_summary` | the root cause (truncated) |
+| `denpex.diagnosis_id` | link back to [denpex.com/dashboard](https://denpex.com/dashboard) |
+| artifact `denpex_diagnosis.md` | full root cause + fix |
+
+## Notes
+- **Best-effort & non-masking:** if Denpex is unreachable, unauthenticated, or over the
+  free-plan quota, it logs one line to stderr and still re-raises your real error.
+- `KeyboardInterrupt` / `SystemExit` are passed through undiagnosed.
+- The diagnosis runs through the same `/api/diagnose` endpoint as the rest of Denpex, so
+  it also lands in your dashboard history and fires your configured alert channels.
+- Pass `logfile="train.err"` to `watch()` to include the tail of a log file alongside the
+  traceback.
+
+---
+
+## Official Denpex Ecosystem & Documentation
+
+| Resource | Description | Link |
+| :--- | :--- | :--- |
+| **Denpex Platform** | Deterministic GPU and ML training failure diagnostics | [denpex.com](https://denpex.com) |
+| **Documentation Hub** | Quickstart guides, architecture, and operational reference | [denpex.com/docs](https://denpex.com/docs) |
+| **Python Agent & CLI** | On-host monitoring, evidence capture, and zero-egress CLI | [denpex.com/docs/agent](https://denpex.com/docs/agent) |
+| **REST API Reference** | Programmatic `/api/diagnose` and preflight endpoints | [denpex.com/docs/api](https://denpex.com/docs/api) |
+| **VS Code Extension** | Offline GPU crash diagnostics in VS Code | [denpex.com/vscode](https://denpex.com/vscode) |
+| **Open VSX Extension** | Extension for Cursor, Windsurf, VSCodium, and Gitpod | [denpex.com/openvsx](https://denpex.com/openvsx) |
+| **Failure Encyclopedia** | 500+ documented CUDA, NCCL, Xid, Slurm, and vLLM failures | [denpex.com/failures](https://denpex.com/failures) |
+| **Free GPU Tools** | VRAM, NCCL, DeepSpeed ZeRO, and Slurm calculators | [denpex.com/tools](https://denpex.com/tools) |
+| **NVIDIA Xid Reference** | Hardware vs software Xid fault routing table | [denpex.com/xid](https://denpex.com/xid) |
+| **Release Changelog** | Platform, agent, SDK, and extension release notes | [denpex.com/changelog](https://denpex.com/changelog) |
+

denpex_mlflow-0.1.0/denpex_mlflow/__init__.py

@@ -0,0 +1,31 @@
+"""denpex-mlflow — attach Denpex GPU/training-failure diagnoses to your MLflow runs.
+
+    import denpex_mlflow as denpex
+    with denpex.watch(job_name="llama-7b-sft"):
+        train()
+
+On failure, Denpex diagnoses the crash and writes the root cause + fix to the active
+MLflow run as tags (`denpex.failure_type`, `denpex.fix_summary`, …) and a
+`denpex_diagnosis.md` artifact, then re-raises your original exception.
+"""
+from ._client import (
+    DenpexAuthError,
+    DenpexClient,
+    DenpexError,
+    DenpexQuotaError,
+    Diagnosis,
+)
+from .watch import diagnose, watch, watch_run
+
+__version__ = "0.1.0"
+__all__ = [
+    "watch",
+    "watch_run",
+    "diagnose",
+    "DenpexClient",
+    "Diagnosis",
+    "DenpexError",
+    "DenpexAuthError",
+    "DenpexQuotaError",
+    "__version__",
+]

denpex_mlflow-0.1.0/denpex_mlflow/_client.py

@@ -0,0 +1,233 @@
+"""Zero-dependency client for the Denpex diagnosis API.
+
+Targets the live `POST /api/diagnose` endpoint (API-key auth via the `X-Denpex-Key`
+header). This is the same path the website console and the cluster agent use, so a diagnosis
+from here is saved to your history and fans out to your alert channels too.
+"""
+from __future__ import annotations
+
+import json
+import os
+import time
+import urllib.error
+import urllib.request
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+from .evidence import select_evidence
+
+DEFAULT_BASE_URL = os.environ.get("DENPEX_API_URL", "https://api.denpex.com")
+MAX_LOG_CHARS = 50_000
+
+
+class DenpexError(Exception):
+    """Base error for any diagnosis failure."""
+
+
+class DenpexAuthError(DenpexError):
+    """Missing/invalid API key (HTTP 401)."""
+
+
+class DenpexQuotaError(DenpexError):
+    """Monthly free-plan quota exhausted (HTTP 429, UPGRADE_REQUIRED)."""
+
+
+@dataclass
+class Diagnosis:
+    """The canonical diagnosis contract, as every other Denpex surface returns it.
+
+    This carried six fields: type, confidence, summary, action, source, and id. It
+    silently discarded the rest of the response. What it dropped is the half an
+    operator acts on: the fallback when the primary fix does not work, the
+    emergency stopgap, the command that PROVES recovery, the evidence the verdict
+    rests on, which rank/GPU initiated the failure, whether a checkpoint is safe
+    to resume from, and the questions the engine asks when it cannot decide.
+
+    An MLflow user therefore saw a confident one-line answer with no escalation
+    path and no way to check whether it worked, from the same API call that
+    returned all of it.
+    """
+
+    failure_type: str
+    # CALIBRATED display value, not the raw score. `meta.confidence` includes
+    # bands measured to be unreliable; `fix.confidenceDisplay` is null when the
+    # band is not accurate enough to show, and None here means "withheld", never
+    # "zero". Printing the raw number re-adds a band with known-bad precision.
+    confidence: Optional[int]
+    summary: str
+    action: str
+    source: Optional[str] = None
+    diagnosis_id: Optional[str] = None
+
+    # ── The escalation ladder ──
+    fallback_action: Optional[str] = None
+    emergency_action: Optional[str] = None
+    verification_command: Optional[str] = None
+    resume_from: Optional[str] = None
+
+    # ── The reasoning trail ──
+    evidence: List[str] = field(default_factory=list)
+    initiator: Optional[str] = None
+    competing_hypotheses: List[str] = field(default_factory=list)
+    clarifying_questions: List[str] = field(default_factory=list)
+
+    def as_markdown(self) -> str:
+        confidence = f"{self.confidence}%" if self.confidence is not None else "withheld (calibration band too wide to report)"
+        parts = [
+            f"# Denpex diagnosis: {self.failure_type}\n",
+            f"**Confidence:** {confidence}\n",
+        ]
+        if self.initiator:
+            parts.append(f"**Initiating failure:** {self.initiator}\n")
+        parts.append(f"## Root cause\n{self.summary or '_n/a_'}\n")
+        parts.append(f"## Fix\n{self.action or '_n/a_'}\n")
+        if self.fallback_action:
+            parts.append(f"## If that does not work\n{self.fallback_action}\n")
+        if self.emergency_action:
+            parts.append(f"## Emergency stopgap\n{self.emergency_action}\n")
+        if self.verification_command:
+            # An exit code is not proof of recovery; this is what proves it.
+            parts.append(f"## Verify the fix worked\n```\n{self.verification_command}\n```\n")
+        if self.resume_from:
+            parts.append(f"## Safe resume point\n`{self.resume_from}`\n")
+        if self.evidence:
+            parts.append("## Evidence\n" + "\n".join(f"- {e}" for e in self.evidence) + "\n")
+        if self.competing_hypotheses:
+            parts.append("## Also considered\n" + "\n".join(f"- {h}" for h in self.competing_hypotheses) + "\n")
+        if self.clarifying_questions:
+            parts.append(
+                "## To narrow this down\n"
+                + "\n".join(f"- {q}" for q in self.clarifying_questions)
+                + "\n"
+            )
+        if self.diagnosis_id:
+            parts.append(f"\n_Diagnosis `{self.diagnosis_id}` · https://denpex.com/dashboard_\n")
+        return "\n".join(parts)
+
+
+class DenpexClient:
+    """Thin client over POST /api/diagnose with bounded retry/backoff."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: str = DEFAULT_BASE_URL,
+        # 300s, not 120s. The canonical deep tier has a 290s deadline, so a
+        # 120s client timeout abandoned answers the server was still producing
+        # and surfaced them to the user as a failure.
+        timeout: float = 300.0,
+        retries: int = 2,
+        backoff: float = 0.5,
+    ):
+        self.api_key = api_key if api_key is not None else os.environ.get("DENPEX_API_KEY", "")
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.retries = retries
+        self.backoff = backoff
+
+    def diagnose(self, logs: str, job_name: Optional[str] = None) -> Diagnosis:
+        if not self.api_key:
+            raise DenpexAuthError("DENPEX_API_KEY is not set (get a key at denpex.com/dashboard).")
+        if not logs or len(logs.strip()) < 6:
+            raise DenpexError("Provide the original machine event or failure log to diagnose.")
+
+        selected_logs, _truncated = select_evidence(logs, MAX_LOG_CHARS)
+        payload = json.dumps({"logs": selected_logs, "jobName": job_name}).encode("utf-8")
+        last_err: Optional[DenpexError] = None
+
+        for attempt in range(self.retries + 1):
+            req = urllib.request.Request(
+                self.base_url + "/api/diagnose",
+                data=payload,
+                method="POST",
+                headers={
+                    "Content-Type": "application/json",
+                    "X-Denpex-Key": self.api_key,
+                    "User-Agent": "Denpex-MLflow/0.1",
+                },
+            )
+            try:
+                with urllib.request.urlopen(req, timeout=self.timeout) as r:
+                    data = json.loads(r.read().decode("utf-8") or "{}")
+                return self._parse(data)
+            except urllib.error.HTTPError as e:
+                body = self._read_err(e)
+                if e.code == 401:
+                    raise DenpexAuthError("Invalid DENPEX_API_KEY.") from None
+                if e.code == 429 and body.get("code") == "UPGRADE_REQUIRED":
+                    raise DenpexQuotaError(body.get("error", "Monthly quota exceeded.")) from None
+                if e.code == 400:
+                    raise DenpexError(body.get("guidance") or body.get("error", "Bad request.")) from None
+                # 429 rate-limit or 5xx -> retryable
+                last_err = DenpexError(body.get("error", f"HTTP {e.code}"))
+                if not (e.code == 429 or 500 <= e.code < 600):
+                    raise last_err from None
+            except urllib.error.URLError as e:
+                last_err = DenpexError(f"Network error: {e.reason}")
+            except (TimeoutError, OSError) as e:
+                last_err = DenpexError(f"Connection error: {e}")
+
+            if attempt < self.retries:
+                time.sleep(self.backoff * (2 ** attempt))
+
+        raise last_err or DenpexError("Diagnosis failed.")
+
+    @staticmethod
+    def _read_err(e: urllib.error.HTTPError) -> dict:
+        try:
+            return json.loads(e.read().decode("utf-8"))
+        except Exception:
+            return {}
+
+    @staticmethod
+    def _parse(data: dict) -> Diagnosis:
+        if not data.get("success"):
+            raise DenpexError(data.get("error", "Diagnosis failed."))
+        meta = data.get("meta") or {}
+        fix = data.get("fix") or {}
+
+        def _strs(value: Any, limit: int = 12) -> List[str]:
+            """Response fields are untrusted input even from our own API."""
+            if not isinstance(value, list):
+                return []
+            out = []
+            for item in value[:limit]:
+                if isinstance(item, str) and item.strip():
+                    out.append(item.strip())
+                elif isinstance(item, dict):
+                    # clarifyingChoices / rankedHypotheses arrive as objects; take
+                    # the human-readable field rather than dropping them entirely.
+                    text = item.get("question") or item.get("theory") or item.get("label")
+                    if isinstance(text, str) and text.strip():
+                        out.append(text.strip())
+            return out
+
+        # `confidenceDisplay` is the calibrated value and may legitimately be
+        # null. That means "withheld", which is not the same as zero, so it is
+        # NOT coerced with `or 0`. Falling back to meta.confidence only when the
+        # key is absent entirely keeps older responses working.
+        display = fix.get("confidenceDisplay", "__missing__")
+        if display == "__missing__":
+            raw = meta.get("confidence")
+            confidence = int(raw) if isinstance(raw, (int, float)) else None
+        elif display is None:
+            confidence = None
+        else:
+            confidence = int(display)
+
+        return Diagnosis(
+            failure_type=meta.get("failureType", "UNKNOWN"),
+            confidence=confidence,
+            summary=fix.get("summary", ""),
+            action=fix.get("action", ""),
+            source=meta.get("source"),
+            diagnosis_id=data.get("diagnosisId"),
+            fallback_action=fix.get("fallbackAction") or None,
+            emergency_action=fix.get("emergencyAction") or None,
+            verification_command=fix.get("verificationCommand") or None,
+            resume_from=fix.get("resumeFrom") or None,
+            evidence=_strs(fix.get("evidence") or meta.get("evidence")),
+            initiator=(fix.get("initiator") or meta.get("initiator") or None),
+            competing_hypotheses=_strs(fix.get("competingHypotheses") or fix.get("rankedHypotheses")),
+            clarifying_questions=_strs(fix.get("clarifyingChoices") or fix.get("clarifyingQuestions")),
+        )

denpex_mlflow-0.1.0/denpex_mlflow/evidence.py

@@ -0,0 +1,72 @@
+"""Bounded causal evidence selection for the MLflow integration."""
+
+from __future__ import annotations
+
+import re
+from typing import Dict, Tuple
+
+_EVIDENCE = re.compile(
+    r"error|fail|fatal|traceback|exception|assert|panic|segfault|xid|nvrm|ecc|"
+    r"fallen off the bus|nccl|cuda|oom|out of memory|timeout|nan\b|\binf\b|"
+    r"node_fail|killed|terminated|version|driver|kernel|container|slurm|kubernetes|"
+    r"rank|node|host|gpu|pci|topology|nvlink|nvswitch|infiniband|rdma|roce",
+    re.I,
+)
+_CAUSAL = re.compile(
+    r"fatal|critical|traceback|exception|assert|panic|segfault|xid|uncorrectable|"
+    r"fallen off the bus|out of memory|watchdog timeout|crc mismatch|loss\s*=\s*(?:nan|inf)|"
+    r"insufficient|unhealthy|ibv_modify_qp|no space left|compil(?:e|ation) (?:error|failed)|"
+    r"no available memory|checksum mismatch|childfailederror|backendcompilerfailed|"
+    r"failedscheduling|input/output error|nonfinite|divergence|first failure|initiator|caused by",
+    re.I,
+)
+
+
+def _clip_line(line: str, maximum: int = 2000) -> str:
+    if len(line) <= maximum:
+        return line
+    marker = " [line clipped by Denpex] "
+    side = max(1, (maximum - len(marker)) // 2)
+    return (line[:side] + marker + line[-side:])[:maximum]
+
+
+def select_evidence(logs: str, maximum: int) -> Tuple[str, bool]:
+    if len(logs) <= maximum:
+        return logs, False
+    lines = logs.splitlines()
+    candidates: Dict[int, int] = {}
+
+    def add(index: int, radius: int, priority: int) -> None:
+        for current in range(max(0, index - radius), min(len(lines), index + radius + 1)):
+            candidates[current] = max(priority, candidates.get(current, 0))
+
+    for index in range(min(48, len(lines))):
+        candidates[index] = 300
+    for index in range(max(0, len(lines) - 96), len(lines)):
+        candidates[index] = 800
+
+    found_first = False
+    for index, line in enumerate(lines):
+        if not _EVIDENCE.search(line):
+            continue
+        causal = bool(_CAUSAL.search(line))
+        add(index, 3 if causal else 1, 1000 if causal else 700)
+        if not found_first:
+            add(index, 6, 1100)
+            found_first = True
+
+    warning = (
+        f"[Denpex condensed input: input exceeded {maximum:,} characters. "
+        "Environment, causal, rank, topology, middle, and final-tail evidence were selected.]"
+    )
+    used = len(warning)
+    selected = set()
+    for index, _priority in sorted(candidates.items(), key=lambda item: (-item[1], item[0])):
+        line = _clip_line(lines[index])
+        if used + len(line) + 1 > maximum:
+            continue
+        selected.add(index)
+        used += len(line) + 1
+    output = [warning]
+    output.extend(_clip_line(lines[index]) for index in sorted(selected))
+    return "\n".join(output), True

denpex_mlflow-0.1.0/denpex_mlflow/watch.py

@@ -0,0 +1,116 @@
+"""`denpex.watch()` — wrap a training run so that, on failure, Denpex diagnoses the
+crash and attaches the root cause + fix straight to the active MLflow run.
+
+    import denpex_mlflow as denpex
+    with denpex.watch(job_name="llama-7b-sft"):
+        train()                         # raises -> diagnosed, tagged, re-raised
+"""
+from __future__ import annotations
+
+import functools
+import os
+import sys
+import traceback
+from contextlib import contextmanager
+from typing import Optional
+
+from ._client import DenpexClient, Diagnosis, DenpexError
+
+
+def _collect_logs(exc: BaseException, logfile: Optional[str], tail_lines: int = 200) -> str:
+    parts = []
+    if logfile:
+        try:
+            with open(logfile, "r", errors="replace") as f:
+                parts.append(f.read())
+        except OSError:
+            pass
+    parts.append("".join(traceback.format_exception(type(exc), exc, exc.__traceback__)))
+    return "\n".join(p for p in parts if p and p.strip())
+
+
+def _attach_to_mlflow(diag: Diagnosis) -> bool:
+    """Write the diagnosis to the active MLflow run as tags + an artifact. No-op if
+    mlflow isn't installed or there's no active run. Never raises."""
+    try:
+        import mlflow
+    except Exception:
+        return False
+    try:
+        if mlflow.active_run() is None:
+            return False
+        mlflow.set_tag("denpex.failure_type", diag.failure_type)
+        mlflow.set_tag("denpex.confidence", diag.confidence)
+        mlflow.set_tag("denpex.fix_summary", (diag.summary or "")[:500])
+        if diag.diagnosis_id:
+            mlflow.set_tag("denpex.diagnosis_id", diag.diagnosis_id)
+        mlflow.log_text(diag.as_markdown(), "denpex_diagnosis.md")
+        return True
+    except Exception:
+        return False
+
+
+def _print_summary(diag: Diagnosis, attached: bool) -> None:
+    where = " (attached to MLflow run)" if attached else ""
+    # ASCII only — this writes to whatever stderr the training box has (may be cp1252/ascii).
+    sys.stderr.write(
+        f"\n[denpex] {diag.failure_type} - {diag.confidence}% confidence{where}\n"
+        f"[denpex] root cause: {diag.summary}\n"
+        f"[denpex] fix: {diag.action}\n"
+    )
+
+
+@contextmanager
+def watch(
+    job_name: Optional[str] = None,
+    *,
+    api_key: Optional[str] = None,
+    base_url: Optional[str] = None,
+    logfile: Optional[str] = None,
+    reraise: bool = True,
+    client: Optional[DenpexClient] = None,
+):
+    """Context manager that diagnoses any exception raised inside it.
+
+    The diagnosis is best-effort and never masks your original error: if Denpex is
+    unreachable or unauthenticated we log a line to stderr and still re-raise.
+    KeyboardInterrupt / SystemExit pass straight through, undiagnosed.
+    """
+    try:
+        yield
+    except (KeyboardInterrupt, SystemExit):
+        raise
+    except BaseException as exc:  # noqa: BLE001 - we re-raise below
+        try:
+            c = client or DenpexClient(
+                api_key=api_key if api_key is not None else os.environ.get("DENPEX_API_KEY", ""),
+                base_url=base_url or os.environ.get("DENPEX_API_URL", "https://api.denpex.com"),
+            )
+            diag = c.diagnose(_collect_logs(exc, logfile), job_name=job_name)
+            _print_summary(diag, _attach_to_mlflow(diag))
+        except DenpexError as inner:
+            sys.stderr.write(f"[denpex] diagnosis skipped: {inner}\n")
+        except Exception as inner:  # never let the SDK mask the real failure
+            sys.stderr.write(f"[denpex] diagnosis error: {inner}\n")
+        if reraise:
+            raise
+
+
+def watch_run(job_name: Optional[str] = None, **watch_kwargs):
+    """Decorator form: `@denpex.watch_run("my-job")` around a training entrypoint."""
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            with watch(job_name=job_name or fn.__name__, **watch_kwargs):
+                return fn(*args, **kwargs)
+        return wrapper
+    return decorator
+
+
+def diagnose(logs: str, job_name: Optional[str] = None, *, api_key: Optional[str] = None, base_url: Optional[str] = None) -> Diagnosis:
+    """One-shot convenience: diagnose a log string directly."""
+    client = DenpexClient(
+        api_key=api_key if api_key is not None else os.environ.get("DENPEX_API_KEY", ""),
+        base_url=base_url or os.environ.get("DENPEX_API_URL", "https://api.denpex.com"),
+    )
+    return client.diagnose(logs, job_name=job_name)

denpex_mlflow-0.1.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "denpex-mlflow"
+version = "0.1.0"
+description = "Attach Denpex GPU and training-failure diagnoses to your MLflow runs in one line."
+readme = "README.md"
+requires-python = ">=3.8"
+license = { text = "MIT" }
+authors = [{ name = "Denpex" }]
+keywords = ["mlflow", "gpu", "training", "diagnosis", "denpex", "nccl", "cuda", "pytorch"]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Intended Audience :: Developers",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+# Runtime is intentionally light: the client uses only the stdlib. mlflow is the host
+# tool you're already running; it's imported lazily so the client also works standalone.
+dependencies = ["mlflow>=2.0"]
+
+[project.urls]
+"Homepage" = "https://denpex.com"
+"Documentation" = "https://denpex.com/docs/integrations"
+"API Reference" = "https://denpex.com/docs/api"
+"Python Agent & CLI" = "https://denpex.com/docs/agent"
+"VS Code Extension" = "https://denpex.com/vscode"
+"Open VSX Extension" = "https://denpex.com/openvsx"
+"GPU Failure Encyclopedia" = "https://denpex.com/failures"
+"Free GPU Debugging Tools" = "https://denpex.com/tools"
+"Changelog" = "https://denpex.com/changelog"
+"Repository" = "https://github.com/Denpex/gpu-failure-cheatsheet"
+"Issue Tracker" = "https://denpex.com/contact"
+
+[tool.hatch.build.targets.wheel]
+packages = ["denpex_mlflow"]

denpex_mlflow-0.1.0/tests/test_denpex_mlflow.py

@@ -0,0 +1,288 @@
+"""Unit tests for denpex-mlflow. Stdlib unittest and mock only. No network or real
+mlflow (a fake module is injected). Run from the package root:
+
+    cd packages/denpex-mlflow && python -m unittest discover -s tests -v
+"""
+import io
+import os
+import sys
+import tempfile
+import types
+import unittest
+import urllib.error
+from unittest import mock
+
+# Make the package importable when run from the package root.
+sys.path.insert(0, ".")
+
+import denpex_mlflow as denpex  # noqa: E402
+from denpex_mlflow._client import DenpexClient  # noqa: E402
+
+
+def _ok_response(payload: dict):
+    body = io.BytesIO()
+    import json
+    body.write(json.dumps(payload).encode())
+    body.seek(0)
+    cm = mock.MagicMock()
+    cm.__enter__.return_value.read.return_value = body.getvalue()
+    cm.__exit__.return_value = False
+    return cm
+
+
+def _http_error(code: int, payload: dict):
+    import json
+    return urllib.error.HTTPError(
+        url="https://api.denpex.com/api/diagnose", code=code, msg="err",
+        hdrs=None, fp=io.BytesIO(json.dumps(payload).encode()),
+    )
+
+
+SUCCESS = {
+    "success": True,
+    "diagnosisId": "diag123",
+    "fix": {"summary": "Rank 3 stalled the all-reduce.", "action": "Restart with NCCL_TIMEOUT=6000."},
+    "meta": {"failureType": "NCCL_TIMEOUT", "confidence": 91, "source": "pattern", "tier": "fast"},
+}
+
+
+class ClientTests(unittest.TestCase):
+    def test_diagnose_success_parses_and_sends_key(self):
+        c = DenpexClient(api_key="dpx_test")
+        with mock.patch("urllib.request.urlopen", return_value=_ok_response(SUCCESS)) as uo:
+            diag = c.diagnose("a" * 50, job_name="job-1")
+        self.assertEqual(diag.failure_type, "NCCL_TIMEOUT")
+        self.assertEqual(diag.confidence, 91)
+        self.assertEqual(diag.diagnosis_id, "diag123")
+        self.assertIn("NCCL_TIMEOUT", diag.as_markdown())
+        # request carried the API key + correct path
+        req = uo.call_args.args[0]
+        self.assertEqual(req.get_header("X-denpex-key"), "dpx_test")
+        self.assertEqual(req.get_header("User-agent"), "Denpex-MLflow/0.1")
+        self.assertTrue(req.full_url.endswith("/api/diagnose"))
+
+    def test_short_logs_rejected_without_network(self):
+        c = DenpexClient(api_key="dpx_test")
+        with mock.patch("urllib.request.urlopen", side_effect=AssertionError("should not call")):
+            with self.assertRaises(denpex.DenpexError):
+                c.diagnose("oom")
+
+    def test_missing_api_key(self):
+        with self.assertRaises(denpex.DenpexAuthError):
+            DenpexClient(api_key="").diagnose("a" * 50)
+
+    def test_auth_error_401(self):
+        c = DenpexClient(api_key="dpx_bad")
+        with mock.patch("urllib.request.urlopen", side_effect=_http_error(401, {"error": "Unauthorized"})):
+            with self.assertRaises(denpex.DenpexAuthError):
+                c.diagnose("a" * 50)
+
+    def test_quota_error_429_upgrade_required(self):
+        c = DenpexClient(api_key="dpx_test")
+        with mock.patch("urllib.request.urlopen", side_effect=_http_error(429, {"error": "Quota exceeded", "code": "UPGRADE_REQUIRED"})):
+            with self.assertRaises(denpex.DenpexQuotaError):
+                c.diagnose("a" * 50)
+
+    def test_input_quality_error_preserves_guidance(self):
+        c = DenpexClient(api_key="dpx_test")
+        response = _http_error(400, {
+            "error": "Input rejected.",
+            "code": "INSUFFICIENT_DIAGNOSTIC_SIGNAL",
+            "guidance": "Paste the original CUDA failure line.",
+        })
+        with mock.patch("urllib.request.urlopen", side_effect=response):
+            with self.assertRaisesRegex(denpex.DenpexError, "original CUDA failure"):
+                c.diagnose("hello hello")
+
+    def test_retry_on_5xx_then_success(self):
+        c = DenpexClient(api_key="dpx_test", retries=1, backoff=0)
+        seq = [_http_error(503, {"error": "boom"}), _ok_response(SUCCESS)]
+        with mock.patch("urllib.request.urlopen", side_effect=seq), mock.patch("time.sleep"):
+            diag = c.diagnose("a" * 50)
+        self.assertEqual(diag.failure_type, "NCCL_TIMEOUT")
+
+
+class FakeClient:
+    def __init__(self):
+        self.calls = []
+
+    def diagnose(self, logs, job_name=None):
+        self.calls.append((logs, job_name))
+        return denpex.Diagnosis(failure_type="CUDA_OOM", confidence=88, summary="OOM at step 42", action="Lower batch size.", diagnosis_id="d1")
+
+
+class WatchTests(unittest.TestCase):
+    def setUp(self):
+        # Inject a fake mlflow module that records tag/artifact writes.
+        self.tags = {}
+        self.artifacts = {}
+        fake = types.ModuleType("mlflow")
+        fake.active_run = lambda: object()
+        fake.set_tag = lambda k, v: self.tags.__setitem__(k, v)
+        fake.log_text = lambda text, path: self.artifacts.__setitem__(path, text)
+        self._saved = sys.modules.get("mlflow")
+        sys.modules["mlflow"] = fake
+
+    def tearDown(self):
+        if self._saved is not None:
+            sys.modules["mlflow"] = self._saved
+        else:
+            sys.modules.pop("mlflow", None)
+
+    def test_watch_diagnoses_attaches_and_reraises(self):
+        fc = FakeClient()
+        with self.assertRaises(ValueError):
+            with denpex.watch(job_name="jx", client=fc):
+                raise ValueError("boom at step 42")
+        self.assertEqual(len(fc.calls), 1)
+        self.assertIn("boom at step 42", fc.calls[0][0])      # traceback captured
+        self.assertEqual(fc.calls[0][1], "jx")                 # job_name passed through
+        self.assertEqual(self.tags["denpex.failure_type"], "CUDA_OOM")
+        self.assertIn("denpex_diagnosis.md", self.artifacts)
+
+    def test_watch_passes_complete_log_to_the_bounded_selector(self):
+        fc = FakeClient()
+        handle, path = tempfile.mkstemp(prefix="denpex-mlflow-", suffix=".log")
+        os.close(handle)
+        try:
+            with open(path, "w", encoding="utf-8") as stream:
+                stream.write("PyTorch version: 2.7.1\n")
+                stream.write("ordinary progress\n" * 400)
+                stream.write("[rank 7][node gpu-03] first failure: NCCL timeout\n")
+            with self.assertRaises(RuntimeError):
+                with denpex.watch(job_name="complete-log", client=fc, logfile=path):
+                    raise RuntimeError("final fatal tail")
+            sent = fc.calls[0][0]
+            self.assertIn("PyTorch version: 2.7.1", sent)
+            self.assertIn("first failure: NCCL timeout", sent)
+            self.assertIn("final fatal tail", sent)
+        finally:
+            os.unlink(path)
+
+    def test_watch_keyboardinterrupt_passes_through_undiagnosed(self):
+        fc = FakeClient()
+        with self.assertRaises(KeyboardInterrupt):
+            with denpex.watch(client=fc):
+                raise KeyboardInterrupt()
+        self.assertEqual(fc.calls, [])
+
+    def test_watch_swallows_sdk_errors_but_reraises_original(self):
+        class Boom:
+            def diagnose(self, *a, **k):
+                raise denpex.DenpexError("api down")
+        with self.assertRaises(RuntimeError):
+            with denpex.watch(client=Boom()):
+                raise RuntimeError("real failure")
+
+
+if __name__ == "__main__":
+    unittest.main()
+
+
+# ── The canonical diagnosis contract (2026-08-21) ────────────────────────────
+#
+# The client parsed six fields and discarded the rest of the response: the
+# fallback fix, the emergency stopgap, the command that PROVES recovery, the
+# evidence, the initiating rank/GPU, the safe resume point, and the engine's
+# clarifying questions. An MLflow user got a confident one-liner with no
+# escalation path and no way to check whether it worked, from the same API
+# call that returned all of it.
+#
+# It also printed `meta.confidence`, the RAW score. The calibrated value is
+# `fix.confidenceDisplay`, which is deliberately null when the band is not
+# accurate enough to report; rendering the raw number re-adds a band measured
+# to be unreliable.
+
+from denpex_mlflow._client import DenpexClient
+
+
+def _full_response():
+    return {
+        "success": True,
+        "diagnosisId": "dx_1",
+        "meta": {"failureType": "GPU_XID_ECC_CASCADE", "confidence": 95, "source": "cascade"},
+        "fix": {
+            "summary": "Rank 3 failed first; the other timeouts are collateral.",
+            "action": "Drain the node and reset the GPU.",
+            "confidenceDisplay": 90,
+            "fallbackAction": "If the Xid recurs after reset, open an RMA.",
+            "emergencyAction": "Cordon the node now to stop new work landing on it.",
+            "verificationCommand": "nvidia-smi -q -d ECC",
+            "resumeFrom": "step_9000",
+            "evidence": ["NVRM: Xid (PCI:0000:1a:00): 48"],
+            "initiator": "rank 3 / GPU-aaaabbbb",
+            "competingHypotheses": [{"theory": "NCCL fabric flap"}],
+            "clarifyingChoices": [{"question": "Does it recur after a reset?"}],
+        },
+    }
+
+
+def test_uses_the_calibrated_confidence_not_the_raw_score():
+    d = DenpexClient._parse(_full_response())
+    assert d.confidence == 90, "must use fix.confidenceDisplay, not meta.confidence (95)"
+
+
+def test_withheld_confidence_is_none_not_zero():
+    # None means "we decline to give a number"; 0 means "we are certain it is
+    # zero percent". Coercing with `or 0` prints the second and means the first.
+    r = _full_response()
+    r["fix"]["confidenceDisplay"] = None
+    d = DenpexClient._parse(r)
+    assert d.confidence is None
+    assert "withheld" in d.as_markdown()
+    assert "0%" not in d.as_markdown()
+
+
+def test_falls_back_to_meta_confidence_for_older_responses():
+    r = _full_response()
+    del r["fix"]["confidenceDisplay"]
+    assert DenpexClient._parse(r).confidence == 95
+
+
+def test_carries_the_whole_escalation_ladder():
+    d = DenpexClient._parse(_full_response())
+    assert d.fallback_action and "RMA" in d.fallback_action
+    assert d.emergency_action and "Cordon" in d.emergency_action
+    assert d.verification_command == "nvidia-smi -q -d ECC"
+    assert d.resume_from == "step_9000"
+
+
+def test_carries_the_reasoning_trail():
+    d = DenpexClient._parse(_full_response())
+    assert d.evidence == ["NVRM: Xid (PCI:0000:1a:00): 48"]
+    assert d.initiator == "rank 3 / GPU-aaaabbbb"
+    # Objects, not strings. The client must read the human-readable field
+    # rather than dropping structured entries on the floor.
+    assert d.competing_hypotheses == ["NCCL fabric flap"]
+    assert d.clarifying_questions == ["Does it recur after a reset?"]
+
+
+def test_markdown_shows_the_operator_everything_it_received():
+    md = DenpexClient._parse(_full_response()).as_markdown()
+    for expected in ("If that does not work", "Emergency stopgap", "Verify the fix worked",
+                     "Safe resume point", "Evidence", "Also considered", "To narrow this down"):
+        assert expected in md, f"markdown dropped: {expected}"
+
+
+def test_a_sparse_response_does_not_render_empty_sections():
+    d = DenpexClient._parse({"success": True, "meta": {"failureType": "X", "confidence": 50},
+                             "fix": {"summary": "s", "action": "a"}})
+    md = d.as_markdown()
+    assert "If that does not work" not in md
+    assert "Verify the fix worked" not in md
+
+
+def test_malformed_list_fields_do_not_raise():
+    # The response is untrusted input even though we serve it.
+    r = _full_response()
+    r["fix"]["evidence"] = "not a list"
+    r["fix"]["clarifyingChoices"] = [None, 42, {"nothing": "useful"}]
+    d = DenpexClient._parse(r)
+    assert d.evidence == []
+    assert d.clarifying_questions == []
+
+
+def test_client_timeout_outlives_the_server_deadline():
+    # The canonical deep tier runs to 290s. A 120s client timeout abandoned
+    # answers the server was still producing and reported them as failures.
+    assert DenpexClient().timeout >= 290.0