flowmesh-sdk-stack 0.1.0__py3-none-any.whl

flowmesh_stack/doctor.py

@@ -0,0 +1,227 @@
+"""Pure doctor checks shared by FlowMesh tooling."""
+
+import shutil
+import subprocess
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+from flowmesh.config import DEFAULT_CONFIG_PATH, FlowMeshConfig
+from flowmesh.exceptions import ConfigInvalidError, ConfigNotFoundError
+
+from .docker import DockerError, ensure_docker_available
+from .env import validate_env_file
+from .env_schema import EnvSchema, schema_keys, validate_env_values
+
+type FindingLevel = Literal["note", "warning", "error"]
+
+_DEFAULT_CUDA_PROBE_IMAGE = "nvidia/cuda:12.9.1-base-ubuntu24.04"
+_DEFAULT_DOCKER_GPU_RUNTIME = "nvidia"
+
+
+@dataclass(frozen=True)
+class DoctorFinding:
+    level: FindingLevel
+    message: str
+
+
+@dataclass
+class DoctorReport:
+    findings: list[DoctorFinding] = field(default_factory=list)
+    callback: Callable[[DoctorFinding], Any] | None = None
+
+    @property
+    def errors(self) -> list[str]:
+        return [
+            finding.message for finding in self.findings if finding.level == "error"
+        ]
+
+    @property
+    def warnings(self) -> list[str]:
+        return [
+            finding.message for finding in self.findings if finding.level == "warning"
+        ]
+
+    @property
+    def notes(self) -> list[str]:
+        return [finding.message for finding in self.findings if finding.level == "note"]
+
+    def error(self, message: str) -> None:
+        self._add_finding("error", message)
+
+    def warning(self, message: str) -> None:
+        self._add_finding("warning", message)
+
+    def note(self, message: str) -> None:
+        self._add_finding("note", message)
+
+    def extend_errors(self, messages: Iterable[str]) -> None:
+        for message in messages:
+            self.error(message)
+
+    def extend_warnings(self, messages: Iterable[str]) -> None:
+        for message in messages:
+            self.warning(message)
+
+    def extend_notes(self, messages: Iterable[str]) -> None:
+        for message in messages:
+            self.note(message)
+
+    def _add_finding(self, level: FindingLevel, message: str) -> None:
+        finding = DoctorFinding(level, message)
+        self.findings.append(finding)
+        if self.callback:
+            self.callback(finding)
+
+
+def run_doctor_checks(
+    env_file: Path,
+    schema: EnvSchema,
+    callback: Callable[[DoctorFinding], Any] | None = None,
+) -> DoctorReport:
+    """Run shared doctor checks and return structured findings."""
+    report = DoctorReport(callback=callback)
+    env_values, env_errors = validate_env_file(
+        env_file, expected_keys=schema_keys(schema)
+    )
+    report.extend_errors(env_errors)
+    if env_values is not None:
+        errors, warnings = validate_env_values(schema, env_values)
+        report.extend_errors(errors)
+        report.extend_warnings(warnings)
+    validate_config_file(report)
+    validate_docker_availability(report)
+    validate_gpu_visibility(report, env_values or {})
+    return report
+
+
+def validate_config_file(report: DoctorReport) -> None:
+    """Validate the presence and basic correctness of the config file."""
+    try:
+        FlowMeshConfig.from_file(DEFAULT_CONFIG_PATH)
+    except ConfigNotFoundError as exc:
+        report.warning(str(exc))
+    except ConfigInvalidError as exc:
+        report.error(str(exc))
+    else:
+        report.note(f"Config file found at {DEFAULT_CONFIG_PATH}")
+
+
+def validate_docker_availability(report: DoctorReport) -> None:
+    """Validate docker CLI and daemon reachability."""
+    try:
+        ensure_docker_available()
+    except DockerError as exc:
+        report.error(str(exc))
+        return
+    report.note("Docker is available")
+    docker_bin = _require_bin("docker")
+
+    try:
+        version = subprocess.run(
+            [
+                docker_bin,
+                "--version",
+            ],  # nosec B603: argv list, no shell, absolute path.
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if version.stdout:
+            report.note(version.stdout.strip())
+        elif version.stderr:
+            report.note(version.stderr.strip())
+    except FileNotFoundError:
+        report.error("Docker CLI not found")
+        return
+
+    docker_info = subprocess.run(
+        [docker_bin, "info"],  # nosec B603: argv list, no shell, absolute path.
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if docker_info.returncode == 0:
+        report.note("Docker daemon: reachable")
+    else:
+        report.error("Docker daemon: NOT reachable")
+
+    if shutil.which("docker-compose"):
+        report.warning("docker-compose detected (legacy).")
+    else:
+        report.note("Using docker compose plugin.")
+
+
+def validate_gpu_visibility(report: DoctorReport, env_values: dict[str, str]) -> None:
+    """Validate whether GPUs are visible to the host and Docker runtime."""
+    nvidia_smi_bin = shutil.which("nvidia-smi")
+    if nvidia_smi_bin:
+        smi = subprocess.run(
+            [nvidia_smi_bin],  # nosec B603: argv list, no shell, absolute path.
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if smi.stdout:
+            report.note("nvidia-smi output:")
+            report.note(smi.stdout)
+        if smi.returncode != 0:
+            detail = (smi.stderr or smi.stdout).strip()
+            report.warning(f"nvidia-smi failed on host: {detail or 'unknown error'}")
+            return
+        validate_docker_gpu_runtime(report, env_values)
+        return
+    report.warning("nvidia-smi not found; GPU visibility not verified.")
+
+
+def validate_docker_gpu_runtime(
+    report: DoctorReport, env_values: dict[str, str]
+) -> None:
+    """Validate that the configured Docker GPU runtime works with the probe image."""
+    docker_bin = shutil.which("docker")
+    if docker_bin is None:
+        return
+
+    probe_image = env_values.get("SERVER_CUDA_PROBE_IMAGE", _DEFAULT_CUDA_PROBE_IMAGE)
+    runtime = env_values.get("DOCKER_GPU_RUNTIME", _DEFAULT_DOCKER_GPU_RUNTIME).strip()
+    command = [docker_bin, "run", "--rm"]
+    if runtime:
+        command += ["--runtime", runtime]
+    command += [
+        "--gpus",
+        "all",
+        probe_image,
+        "nvidia-smi",
+        "--query-gpu=index,name",
+        "--format=csv,noheader",
+    ]
+    result = subprocess.run(
+        command,  # nosec B603: argv list, no shell, absolute path.
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if result.returncode == 0:
+        report.note("Docker GPU probe succeeded.")
+        return
+
+    stderr = (result.stderr or "").strip()
+    stdout = (result.stdout or "").strip()
+    detail = stderr or stdout or f"exit code {result.returncode}"
+    lowered = detail.lower()
+    if runtime and "unknown or invalid runtime name" in lowered:
+        report.warning(
+            f"DOCKER_GPU_RUNTIME={runtime!r} is not available to Docker on this host. "
+            "If `docker run --rm --gpus all ...` works without `--runtime`, set "
+            "`DOCKER_GPU_RUNTIME=` in the stack env. This is common on DGX Spark."
+        )
+        return
+    report.warning(f"Docker GPU probe failed: {detail}")
+
+
+def _require_bin(name: str) -> str:
+    path = shutil.which(name)
+    if path is None:
+        raise FileNotFoundError(name)
+    return path

flowmesh_stack/env.py

@@ -0,0 +1,145 @@
+"""Environment file helpers shared by FlowMesh tooling."""
+
+import os
+from pathlib import Path
+from urllib.parse import urlparse
+
+
+def parse_env_file(env_file: Path) -> dict[str, str]:
+    """Parse a .env file into key/value pairs."""
+    values: dict[str, str] = {}
+    if not env_file.exists():
+        return values
+    for line in env_file.read_text().splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#") or "=" not in stripped:
+            continue
+        stripped = stripped.removeprefix("export ").strip()
+        key, value = stripped.split("=", 1)
+        values[key.strip()] = _normalize_env_value(value)
+    return values
+
+
+def parse_bool(value: str) -> bool | None:
+    """Parse a string into a boolean value."""
+    lowered = value.strip().lower()
+    if lowered in {"1", "true", "yes", "on"}:
+        return True
+    if lowered in {"0", "false", "no", "off"}:
+        return False
+    return None
+
+
+def parse_int(value: str) -> int | None:
+    """Parse a string into an integer value."""
+    stripped = value.strip()
+    if not stripped:
+        return None
+    try:
+        return int(stripped)
+    except ValueError:
+        return None
+
+
+def parse_float(value: str) -> float | None:
+    """Parse a string into a float value."""
+    stripped = value.strip()
+    if not stripped:
+        return None
+    try:
+        return float(stripped)
+    except ValueError:
+        return None
+
+
+def is_url(value: str, schemes: set[str] | None = None) -> bool:
+    """Check if a string is a valid URL with optional scheme restrictions."""
+    parsed = urlparse(value.strip())
+    if not (parsed.scheme and parsed.netloc):
+        return False
+    if schemes and parsed.scheme not in schemes:
+        return False
+    return True
+
+
+def validate_env_file(
+    env_file: Path,
+    example: Path | None = None,
+    expected_keys: set[str] | None = None,
+) -> tuple[dict[str, str] | None, list[str]]:
+    """Validate an env file against an example template or key set."""
+    errors: list[str] = []
+    if not env_file.exists():
+        return None, [f"env file not found: {env_file}"]
+    if expected_keys is None:
+        if example is None or not example.exists():
+            return parse_env_file(env_file), errors
+        expected_keys = _parse_env_keys(example)
+
+    actual_keys = _parse_env_keys(env_file)
+    missing = sorted(expected_keys - actual_keys)
+    unexpected = sorted(actual_keys - expected_keys)
+    if missing:
+        errors.append(f"Missing required env vars in {env_file}: {', '.join(missing)}")
+    if unexpected:
+        errors.append(f"Unexpected env vars in {env_file}: {', '.join(unexpected)}")
+    return parse_env_file(env_file), errors
+
+
+def ensure_env_file(env_file: Path, example: Path) -> bool:
+    """Create an env file from an example if it does not exist."""
+    if env_file.exists() or not example.exists():
+        return False
+    env_file.write_text(example.read_text())
+    return True
+
+
+def load_env(
+    env_file: Path,
+    base_dir: Path | None = None,
+    path_keys: set[str] | None = None,
+) -> None:
+    """Load env vars from a file into ``os.environ``."""
+    env_key = (env_file, base_dir, path_keys)
+    if getattr(load_env, "_loaded", None) == env_key:
+        return
+    if not env_file.exists():
+        return
+    for line in env_file.read_text().splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#") or "=" not in stripped:
+            continue
+        key, value = stripped.split("=", 1)
+        if path_keys and key in path_keys and value:
+            expanded = Path(value).expanduser()
+            if expanded.is_absolute():
+                os.environ[key] = str(expanded)
+            elif base_dir is not None:
+                os.environ[key] = str((base_dir / expanded).resolve())
+            else:
+                os.environ[key] = value
+        else:
+            os.environ[key] = value
+    load_env._loaded = env_key  # type: ignore[attr-defined]
+
+
+def _parse_env_keys(path: Path) -> set[str]:
+    keys: set[str] = set()
+    if not path.exists():
+        return keys
+    for line in path.read_text().splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#") or "=" not in stripped:
+            continue
+        stripped = stripped.removeprefix("export ").strip()
+        key = stripped.split("=", 1)[0].strip()
+        if key:
+            keys.add(key)
+    return keys
+
+
+def _normalize_env_value(value: str) -> str:
+    stripped = value.strip()
+    if len(stripped) >= 2 and stripped[0] == stripped[-1] and stripped[0] in "\"'":
+        stripped = stripped[1:-1]
+    return stripped.strip()

flowmesh_stack/env_schema.py

@@ -0,0 +1,238 @@
+"""Environment schema definitions and pure validation helpers."""
+
+import enum
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass, field
+from logging import _nameToLevel as LOG_LEVELS
+from pathlib import Path
+from typing import Literal
+
+from .env import is_url, parse_bool, parse_float, parse_int
+
+
+class EnvVarType(enum.StrEnum):
+    STRING = "string"
+    INT = "int"
+    FLOAT = "float"
+    BOOL = "bool"
+    FILE_PATH = "file_path"
+    DIR_PATH = "dir_path"
+    URL = "url"
+    LOG_LEVEL = "log_level"
+    ENUM = "enum"
+    CSV = "csv"
+    CSV_INTS_OR_ALL = "csv_ints_or_all"
+
+
+@dataclass(frozen=True)
+class EnvVar:
+    key: str
+    default: str = ""
+    description: str | list[str] | None = None
+    var_type: EnvVarType = EnvVarType.STRING
+    required: bool = False
+    use_default: bool = False
+    choices: Iterable[str] | None = None
+    min_value: float | None = None
+    max_value: float | None = None
+    min_length: int | None = None
+    ensure_path: Literal["error", "warn", "create"] | None = None
+    url_schemes: set[str] | None = None
+    warn_if_empty: bool = False
+    validator: Callable[[str, list[str], list[str]], None] | None = None
+
+
+@dataclass(frozen=True)
+class EnvSection:
+    title: str
+    description: list[str] = field(default_factory=list)
+    vars: list[EnvVar] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class EnvSchema:
+    name: str
+    header: list[str]
+    sections: list[EnvSection]
+    validators: list[Callable[[dict[str, str], list[str], list[str]], None]] = field(
+        default_factory=list
+    )
+
+
+def schema_keys(schema: EnvSchema) -> set[str]:
+    """Return the set of keys defined by a schema."""
+    keys: set[str] = set()
+    for section in schema.sections:
+        for var in section.vars:
+            keys.add(var.key)
+    return keys
+
+
+def render_env_example(
+    schema: EnvSchema, overrides: Mapping[str, str] | None = None
+) -> str:
+    """Render an example .env file based on the schema.
+
+    ``overrides`` swaps in a different default value for the listed keys to produce a
+    worker-shaped env without rebuilding the schema). Keys not present in ``overrides``
+    use their schema-declared default.
+    """
+    overrides = overrides or {}
+    lines: list[str] = []
+    lines.extend(schema.header)
+    for section in schema.sections:
+        lines.append("")
+        lines.append(f"# ==== {section.title} ====")
+        for desc in section.description:
+            lines.append(f"# {desc}")
+        for var in section.vars:
+            if description := var.description:
+                if isinstance(description, list):
+                    for desc_line in description:
+                        lines.append(f"# {desc_line}")
+                else:
+                    lines.append(f"# {description}")
+            value = overrides.get(var.key, var.default)
+            lines.append(f"{var.key}={value}")
+    lines.append("")
+    return "\n".join(lines)
+
+
+def validate_env_values(
+    schema: EnvSchema, env: dict[str, str]
+) -> tuple[list[str], list[str]]:
+    """Validate environment variable values against the schema.
+
+    Returns a tuple of (errors, warnings) found during validation.
+    """
+    errors: list[str] = []
+    warnings: list[str] = []
+    for section in schema.sections:
+        for var in section.vars:
+            raw = env.get(var.key, "").strip()
+            if not raw:
+                use_default = var.use_default
+                if var.required:
+                    errors.append(f"{var.key} must be set")
+                    use_default = False
+                elif var.warn_if_empty:
+                    message = f"{var.key} is empty"
+                    if use_default:
+                        message += f"; default value '{var.default}' will be used"
+                    warnings.append(message)
+                if not use_default:
+                    continue
+                raw = var.default
+
+            if var.min_length is not None and len(raw) < var.min_length:
+                errors.append(f"{var.key} must be at least {var.min_length} characters")
+
+            match var.var_type:
+                case EnvVarType.INT:
+                    int_value = parse_int(raw)
+                    if int_value is None:
+                        errors.append(f"{var.key} must be an integer")
+                        continue
+                    if var.min_value is not None and int_value < var.min_value:
+                        errors.append(f"{var.key} must be >= {int(var.min_value)}")
+                    if var.max_value is not None and int_value > var.max_value:
+                        errors.append(f"{var.key} must be <= {int(var.max_value)}")
+                case EnvVarType.FLOAT:
+                    float_value = parse_float(raw)
+                    if float_value is None:
+                        errors.append(f"{var.key} must be a number")
+                        continue
+                    if var.min_value is not None and float_value < var.min_value:
+                        errors.append(f"{var.key} must be >= {var.min_value}")
+                    if var.max_value is not None and float_value > var.max_value:
+                        errors.append(f"{var.key} must be <= {var.max_value}")
+                case EnvVarType.BOOL:
+                    if parse_bool(raw) is None:
+                        errors.append(
+                            f"{var.key} must be a boolean (true/false or 1/0)"
+                        )
+                case EnvVarType.FILE_PATH | EnvVarType.DIR_PATH:
+                    _ensure_path(raw, var, errors, warnings)
+                case EnvVarType.URL:
+                    if not is_url(raw, schemes=var.url_schemes):
+                        errors.append(f"{var.key} must be a valid URL")
+                case EnvVarType.LOG_LEVEL:
+                    if raw.upper() not in LOG_LEVELS:
+                        errors.append(f"{var.key} must be a valid log level")
+                case EnvVarType.ENUM:
+                    if var.choices and raw not in var.choices:
+                        allowed = ", ".join(sorted(var.choices))
+                        errors.append(f"{var.key} must be one of: {allowed}")
+                case EnvVarType.CSV:
+                    parts = [part.strip() for part in raw.split(",")]
+                    if any(not part for part in parts):
+                        errors.append(f"{var.key} must not contain empty entries")
+                case EnvVarType.CSV_INTS_OR_ALL:
+                    if raw.lower() != "all":
+                        parts = [part.strip() for part in raw.split(",")]
+                        if any(not part.isdigit() for part in parts if part):
+                            errors.append(
+                                f"{var.key} must be 'all' or a "
+                                "comma-separated list of integers"
+                            )
+
+            if var.validator:
+                var.validator(raw, errors, warnings)
+    for validator in schema.validators:
+        validator(env, errors, warnings)
+
+    return errors, warnings
+
+
+def require_if_true(
+    env: dict[str, str], flag_key: str, required_keys: list[str], errors: list[str]
+) -> None:
+    """Require keys when a boolean-like flag is true."""
+    if parse_bool(env.get(flag_key, "")):
+        for key in required_keys:
+            if not env.get(key, "").strip():
+                errors.append(f"{key} must be set when {flag_key}=1")
+
+
+def require_pair(
+    env: dict[str, str], key_a: str, key_b: str, errors: list[str]
+) -> None:
+    """Require two keys to be either both set or both empty."""
+    a = env.get(key_a, "").strip()
+    b = env.get(key_b, "").strip()
+    if (a or b) and (not a or not b):
+        errors.append(f"{key_a} and {key_b} must both be set")
+
+
+def require_all_or_none(
+    env: dict[str, str], keys: list[str], errors: list[str]
+) -> None:
+    """Require a key group to be fully set or fully empty."""
+    values = [env.get(key, "").strip() for key in keys]
+    if any(values) and not all(values):
+        errors.append(f"Either all or none of {', '.join(keys)} must be set")
+
+
+def _ensure_path(raw: str, var: EnvVar, errors: list[str], warnings: list[str]) -> None:
+    if not raw:
+        errors.append(f"{var.key} must be a non-empty path")
+        return
+    if var.ensure_path is None:
+        return
+
+    path = Path(raw)
+    if path.exists():
+        if var.var_type == EnvVarType.FILE_PATH and not path.is_file():
+            errors.append(f"{var.key} path should be a file: '{raw}'")
+        elif var.var_type == EnvVarType.DIR_PATH and not path.is_dir():
+            errors.append(f"{var.key} path should be a directory: '{raw}'")
+        return
+
+    message = f"{var.key} path does not exist: '{raw}'"
+    match var.ensure_path:
+        case "error":
+            errors.append(message)
+        case "warn":
+            warnings.append(message)
+        case "create":
+            warnings.append(message + "; it will be created at runtime")