API Reference

The public Python API is intentionally small. Use the CLI for humans and the API when another Python tool needs structured findings.

Public API

mlenvdoctor.api

Public Python API for ML Environment Doctor.

diagnose

diagnose(*, full: bool = False, parallel: bool = True) -> List[DiagnosticIssue]

Run ML environment diagnostics from Python code.

Parameters:

Name	Type	Description	Default
full	bool	Include slower checks such as disk, Docker GPU, and network checks.	False
parallel	bool	Run independent checks concurrently when possible.	True

Returns:

Type	Description
List[DiagnosticIssue]	A list of normalized diagnostic findings.

Source code in src/mlenvdoctor/api.py

def diagnose(*, full: bool = False, parallel: bool = True) -> List[DiagnosticIssue]:
    """Run ML environment diagnostics from Python code.

    Args:
        full: Include slower checks such as disk, Docker GPU, and network checks.
        parallel: Run independent checks concurrently when possible.

    Returns:
        A list of normalized diagnostic findings.
    """
    return diagnose_env(full=full, parallel=parallel, show_header=False)

Diagnostic Models

mlenvdoctor.diagnose.DiagnosticIssue dataclass

Represents a diagnostic outcome.

Source code in src/mlenvdoctor/diagnose.py

@dataclass
class DiagnosticIssue:
    """Represents a diagnostic outcome."""

    name: str
    status: str
    severity: str
    fix: str
    details: Optional[str] = None
    check_id: str = ""
    category: str = "general"
    recommendation: str = ""
    likely_cause: str = ""
    verify_steps: List[str] = field(default_factory=list)
    confidence: str = "medium"
    evidence: List[str] = field(default_factory=list)
    mismatch_code: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_row(self) -> Tuple[str, str, str, str]:
        """Convert to table row."""
        status_icon_map = {
            "PASS": icon_check(),
            "FAIL": icon_cross(),
            "WARN": icon_warning(),
            "INFO": icon_info(),
        }
        status_icon = status_icon_map.get(self.status.split()[0], "?")
        return (
            self.name,
            f"{status_icon} {self.status}",
            self.severity.upper(),
            self.fix,
        )

to_row

to_row() -> Tuple[str, str, str, str]

Convert to table row.

Source code in src/mlenvdoctor/diagnose.py

def to_row(self) -> Tuple[str, str, str, str]:
    """Convert to table row."""
    status_icon_map = {
        "PASS": icon_check(),
        "FAIL": icon_cross(),
        "WARN": icon_warning(),
        "INFO": icon_info(),
    }
    status_icon = status_icon_map.get(self.status.split()[0], "?")
    return (
        self.name,
        f"{status_icon} {self.status}",
        self.severity.upper(),
        self.fix,
    )

mlenvdoctor.diagnose.DoctorFinding dataclass

Prioritized summary item for the doctor command.

Source code in src/mlenvdoctor/diagnose.py

@dataclass
class DoctorFinding:
    """Prioritized summary item for the `doctor` command."""

    problem: str
    severity: str
    confidence: str
    likely_cause: str
    best_fix: str
    verify_steps: List[str]
    evidence: List[str]
    check_id: str
    category: str
    linked_checks: List[str] = field(default_factory=list)

Core Diagnostic Engine

mlenvdoctor.diagnose.diagnose_env

diagnose_env(full: bool = False, parallel: bool = True, show_header: bool = True) -> List[DiagnosticIssue]

Run diagnostic checks and return normalized issues.

Source code in src/mlenvdoctor/diagnose.py

def diagnose_env(
    full: bool = False,
    parallel: bool = True,
    show_header: bool = True,
) -> List[DiagnosticIssue]:
    """Run diagnostic checks and return normalized issues."""
    if show_header:
        console.print(
            f"[bold blue]{icon_search()} Running ML Environment Diagnostics...[/bold blue]\n"
        )

    core_checks = [
        check_python_runtime,
        check_accelerator_backend,
        check_cuda_driver,
        check_pytorch_cuda,
        check_ml_libraries,
        check_tensorflow_keras,
        check_jax_flax,
    ]
    issues = _run_check_group(
        core_checks, parallel=parallel, timeout=60.0, failure_severity="critical"
    )

    if full:
        extended_checks = [
            check_gpu_memory,
            check_disk_space,
            check_docker_gpu,
            check_internet_connectivity,
        ]
        issues.extend(
            _run_check_group(
                extended_checks, parallel=parallel, timeout=120.0, failure_severity="warning"
            )
        )

    issues.extend(_compatibility_matrix_issues(issues))

    return issues

mlenvdoctor.diagnose.summarize_for_doctor

summarize_for_doctor(issues: List['DiagnosticIssue']) -> List[DoctorFinding]

Convert detailed issues into prioritized doctor findings.

Source code in src/mlenvdoctor/diagnose.py

def summarize_for_doctor(issues: List["DiagnosticIssue"]) -> List[DoctorFinding]:
    """Convert detailed issues into prioritized doctor findings."""
    findings = _root_cause_findings(issues)
    if findings:
        findings.sort(
            key=lambda finding: (
                0 if finding.severity == "critical" else 1,
                finding.category,
                finding.problem,
            )
        )
        return findings

    actionable_issues = [
        issue
        for issue in issues
        if issue.status.startswith(("FAIL", "WARN")) and issue.severity in {"critical", "warning"}
    ]
    actionable_issues.sort(
        key=lambda issue: (
            0 if issue.severity == "critical" else 1,
            issue.category,
            issue.name,
        )
    )

    fallback_findings: List[DoctorFinding] = []
    for issue in actionable_issues:
        fallback_findings.append(
            DoctorFinding(
                problem=issue.name,
                severity=issue.severity,
                confidence=issue.confidence,
                likely_cause=issue.likely_cause or _default_likely_cause(issue),
                best_fix=issue.recommendation
                or issue.fix
                or "Run `mlenvdoctor diagnose` for more detail.",
                verify_steps=issue.verify_steps or _default_verify_steps(issue),
                evidence=issue.evidence or ([issue.details] if issue.details else []),
                check_id=issue.check_id,
                category=issue.category,
                linked_checks=[issue.check_id] if issue.check_id else [],
            )
        )
    return fallback_findings

Fix Planning

mlenvdoctor.fix.FixAction dataclass

A validated action the fix engine can perform.

Source code in src/mlenvdoctor/fix.py

@dataclass
class FixAction:
    """A validated action the fix engine can perform."""

    kind: str
    description: str
    risk: str = "low"
    command: Optional[List[str]] = None
    output_path: Optional[Path] = None
    details: str = ""

mlenvdoctor.fix.FixResult dataclass

Structured result for a fix run.

Source code in src/mlenvdoctor/fix.py

@dataclass
class FixResult:
    """Structured result for a fix run."""

    success: bool
    actions: List[FixAction] = field(default_factory=list)
    reasons: List[str] = field(default_factory=list)
    selected_stack: str = "trl-peft"
    executed_actions: List[str] = field(default_factory=list)
    verification_issues: List[DiagnosticIssue] = field(default_factory=list)
    verification_summary: Optional["FixVerificationSummary"] = None
    created_paths: List[Path] = field(default_factory=list)
    message: str = ""

mlenvdoctor.fix.plan_fixes

plan_fixes(issues: List[DiagnosticIssue], *, use_conda: bool, create_venv: bool, stack: str, doctor_findings: Optional[List[DoctorFinding]] = None, requirements_output: str = DEFAULT_REQUIREMENTS_FILE, conda_output: str = 'environment-mlenvdoctor.yml', venv_path: str = '.venv') -> List[FixAction]

Turn diagnostics into a predictable fix plan.

Source code in src/mlenvdoctor/fix.py

def plan_fixes(
    issues: List[DiagnosticIssue],
    *,
    use_conda: bool,
    create_venv: bool,
    stack: str,
    doctor_findings: Optional[List[DoctorFinding]] = None,
    requirements_output: str = DEFAULT_REQUIREMENTS_FILE,
    conda_output: str = "environment-mlenvdoctor.yml",
    venv_path: str = ".venv",
) -> List[FixAction]:
    """Turn diagnostics into a predictable fix plan."""
    findings = doctor_findings if doctor_findings is not None else summarize_for_doctor(issues)
    critical_issues = [
        issue for issue in issues if issue.severity == "critical" and "FAIL" in issue.status
    ]
    actionable_library_issue = any(
        issue.category in {"dependencies", "pytorch"} for issue in critical_issues
    )
    fixable_root_causes = {
        "root_pytorch_missing",
        "root_pytorch_cpu_only_build",
        "root_pytorch_cuda_mismatch",
        "root_ml_stack_dependencies",
        "root_tensorflow_gpu_path",
        "root_jax_backend",
    }
    root_cause_requires_requirements = any(
        finding.check_id in fixable_root_causes
        or finding.category in {"dependencies", "pytorch", "tensorflow", "jax"}
        for finding in findings
    )

    actions: List[FixAction] = []
    if create_venv:
        actions.append(
            FixAction(
                kind="create_venv",
                description=f"Create virtual environment at {venv_path}",
                risk="low",
                output_path=Path(venv_path),
            )
        )

    if use_conda:
        actions.append(
            FixAction(
                kind="write_conda",
                description=f"Generate Conda environment file for stack '{stack}'",
                risk="low",
                output_path=Path(conda_output),
            )
        )
        return actions

    if root_cause_requires_requirements or critical_issues or actionable_library_issue:
        actions.append(
            FixAction(
                kind="write_requirements",
                description=f"Generate requirements file for stack '{stack}'",
                risk="low",
                output_path=Path(requirements_output),
                details=(
                    "Generated from grouped root-cause findings"
                    if root_cause_requires_requirements
                    else "Generated from actionable raw issues"
                ),
            )
        )

    return actions