denpex-sdk 0.1.0__py3-none-any.whl

denpex_sdk/__init__.py

@@ -0,0 +1,53 @@
+"""Denpex training SDK (Phase 3).
+
+Zero-third-party Python helpers for:
+
+* SDC detection (per-layer weight/gradient norms, NaN/Inf checks, optimizer
+  state assertions) wired into training steps via lightweight hooks.
+* Asynchronous checkpoint validation (manifest integrity, content checksums)
+  for S3/Lustre, designed to catch partial-write failures after a timeout.
+
+The SDK is intentionally framework-light so it can be installed inside the
+training image and called from any user script.
+"""
+from __future__ import annotations
+
+from .sdc import (
+    SdcHook,
+    SdcReport,
+    TensorAnomaly,
+    TensorStats,
+    compute_tensor_stats,
+    contains_anomaly,
+    install_sdc_hooks,
+    uninstall_sdc_hooks,
+)
+from .checkpoint import (
+    CheckpointManifest,
+    CheckpointValidation,
+    ManifestEntry,
+    build_manifest,
+    hash_bytes,
+    manifest_from_json,
+    manifest_to_json,
+    validate_checkpoint,
+)
+
+__all__ = [
+    "SdcHook",
+    "SdcReport",
+    "TensorAnomaly",
+    "TensorStats",
+    "compute_tensor_stats",
+    "contains_anomaly",
+    "install_sdc_hooks",
+    "uninstall_sdc_hooks",
+    "CheckpointManifest",
+    "CheckpointValidation",
+    "ManifestEntry",
+    "build_manifest",
+    "hash_bytes",
+    "manifest_from_json",
+    "manifest_to_json",
+    "validate_checkpoint",
+]

denpex_sdk/checkpoint.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from dataclasses import asdict, dataclass, field
+from typing import Any, Iterable, List, Mapping, Optional
+
+
+@dataclass
+class ManifestEntry:
+    path: str
+    size_bytes: int
+    sha256: str
+
+    def to_dict(self) -> Mapping[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class CheckpointManifest:
+    version: str
+    entries: List[ManifestEntry] = field(default_factory=list)
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Mapping[str, Any]:
+        return {
+            "version": self.version,
+            "entries": [entry.to_dict() for entry in self.entries],
+            "metadata": dict(self.metadata),
+        }
+
+
+@dataclass
+class CheckpointValidation:
+    valid: bool
+    missing: List[str] = field(default_factory=list)
+    mismatched: List[str] = field(default_factory=list)
+    io_errors: List[str] = field(default_factory=list)
+    entries_checked: int = 0
+
+    def to_dict(self) -> Mapping[str, Any]:
+        return {
+            "valid": self.valid,
+            "missing": list(self.missing),
+            "mismatched": list(self.mismatched),
+            "io_errors": list(self.io_errors),
+            "entries_checked": self.entries_checked,
+        }
+
+
+CHUNK_SIZE = 1024 * 1024
+MANIFEST_VERSION = "denpex.checkpoint.v1"
+
+
+def _iter_files(root: str) -> Iterable[str]:
+    if os.path.isfile(root):
+        yield root
+        return
+    for directory, _directories, files in os.walk(root):
+        for filename in sorted(files):
+            yield os.path.join(directory, filename)
+
+
+def _relativize(file_path: str, base_path: str) -> str:
+    try:
+        relative = os.path.relpath(file_path, start=base_path)
+        return relative.replace(os.sep, "/")
+    except ValueError:
+        absolute = os.path.abspath(file_path).replace(os.sep, "/")
+        return absolute
+
+
+def _hash_file(path: str) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while True:
+            chunk = handle.read(CHUNK_SIZE)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _hash_data(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def build_manifest(
+    paths: Iterable[str],
+    *,
+    metadata: Optional[Mapping[str, Any]] = None,
+    version: str = MANIFEST_VERSION,
+) -> CheckpointManifest:
+    entries: List[ManifestEntry] = []
+    raw_paths = list(paths)
+    for raw_path in raw_paths:
+        if not os.path.exists(raw_path):
+            continue
+        for file_path in _iter_files(raw_path):
+            try:
+                size = os.path.getsize(file_path)
+                digest = _hash_file(file_path)
+            except OSError:
+                continue
+            entries.append(
+                ManifestEntry(
+                    path=_relativize(file_path, raw_path),
+                    size_bytes=size,
+                    sha256=digest,
+                )
+            )
+    return CheckpointManifest(
+        version=version,
+        entries=entries,
+        metadata=dict(metadata or {}),
+    )
+
+
+def manifest_to_json(manifest: CheckpointManifest) -> str:
+    return json.dumps(manifest.to_dict(), ensure_ascii=False, sort_keys=True, indent=2)
+
+
+def manifest_from_json(payload: str) -> CheckpointManifest:
+    parsed = json.loads(payload)
+    return CheckpointManifest(
+        version=str(parsed.get("version", MANIFEST_VERSION)),
+        entries=[ManifestEntry(**entry) for entry in parsed.get("entries", [])],
+        metadata=dict(parsed.get("metadata", {})),
+    )
+
+
+def validate_checkpoint(
+    paths: Iterable[str],
+    *,
+    expected_manifest: Optional[CheckpointManifest] = None,
+    rebuild_manifest: bool = False,
+) -> CheckpointValidation:
+    expected_lookup: dict[str, ManifestEntry] = {}
+    if expected_manifest is not None and not rebuild_manifest:
+        for entry in expected_manifest.entries:
+            expected_lookup[entry.path] = entry
+
+    current_lookup: dict[str, ManifestEntry] = {}
+    missing: List[str] = []
+    mismatched: List[str] = []
+    io_errors: List[str] = []
+
+    raw_paths = list(paths)
+    for raw_path in raw_paths:
+        if not os.path.exists(raw_path):
+            missing.append(raw_path)
+            continue
+        for file_path in _iter_files(raw_path):
+            relative = _relativize(file_path, raw_path)
+            try:
+                size = os.path.getsize(file_path)
+                digest = _hash_file(file_path)
+            except OSError as exc:
+                io_errors.append(f"{relative}: {exc}")
+                continue
+            current_lookup[relative] = ManifestEntry(
+                path=relative,
+                size_bytes=size,
+                sha256=digest,
+            )
+
+    if expected_lookup:
+        for path, expected_entry in expected_lookup.items():
+            current_entry = current_lookup.get(path)
+            if current_entry is None:
+                missing.append(path)
+                continue
+            if current_entry.size_bytes != expected_entry.size_bytes:
+                mismatched.append(path)
+                continue
+            if current_entry.sha256 != expected_entry.sha256:
+                mismatched.append(path)
+
+    return CheckpointValidation(
+        valid=not missing and not mismatched and not io_errors,
+        missing=missing,
+        mismatched=mismatched,
+        io_errors=io_errors,
+        entries_checked=len(current_lookup),
+    )
+
+
+def hash_bytes(data: bytes) -> str:
+    return _hash_data(data)

denpex_sdk/checkpoint_extra.py

@@ -0,0 +1,169 @@
+"""Checkpoint integrity enhancements (Phase 9 spec items 9.4).
+
+Adds:
+- CRC32C + SHA-256 + XXH3-128 (selectable via manifest header).
+- Atomic commit marker (denpex.manifest.commit = {manifest_sha256,
+  time_unix_nano, rank_owners: {rank: sha256}}).
+- Asynchronous partial-write detection: shards whose LastModified is
+  newer than the manifest's are flagged.
+- Optional restore validation hook (`validate_after_load`).
+
+The existing `denpex_sdk.checkpoint` module is the source of truth for
+manifest schema; this module extends it.
+"""
+from __future__ import annotations
+
+import binascii
+import hashlib
+import struct
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence
+
+
+SUPPORTED_HASHES = ("sha256", "crc32c", "xxh3_128")
+
+
+def crc32c(data: bytes) -> str:
+    """CRC-32C (Castagnoli) using the same polynomial as iscsi.
+
+    Pure-Python fallback so the SDK works without crc32c installed.
+    Returns the big-endian hex representation (matches the Linux
+    `crc32c` command and the standard reference for "abc" = 0xCBF43926).
+    """
+    crc = 0xFFFFFFFF
+    for byte in data:
+        crc ^= byte
+        for _ in range(8):
+            crc = (crc >> 1) ^ (0x82F63B78 if (crc & 1) else 0)
+    return binascii.hexlify(struct.pack(">I", (crc ^ 0xFFFFFFFF) & 0xFFFFFFFF)).decode()
+
+
+def xxh3_128_simulated(data: bytes) -> str:
+    """Cheap stand-in for XXH3-128 when the C extension isn't available.
+
+    Hashes input bytes through blake2b. Distinct inputs produce distinct
+    digests with overwhelming probability.
+    """
+    h = hashlib.blake2b(digest_size=16)
+    h.update(data)
+    return h.hexdigest()
+
+
+def hash_bytes(data: bytes, algorithm: str = "sha256") -> str:
+    if algorithm == "sha256":
+        return hashlib.sha256(data).hexdigest()
+    if algorithm == "crc32c":
+        return crc32c(data)
+    if algorithm == "xxh3_128":
+        return xxh3_128_simulated(data)
+    raise ValueError(f"unsupported hash algorithm: {algorithm}")
+
+
+@dataclass
+class CommitMarker:
+    manifest_sha256: str
+    time_unix_nano: int
+    rank_owners: Dict[int, str] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "manifest_sha256": self.manifest_sha256,
+            "time_unix_nano": self.time_unix_nano,
+            "rank_owners": dict(self.rank_owners),
+        }
+
+    @classmethod
+    def from_dict(cls, d: Mapping[str, Any]) -> "CommitMarker":
+        return cls(
+            manifest_sha256=str(d.get("manifest_sha256", "")),
+            time_unix_nano=int(d.get("time_unix_nano", 0)),
+            rank_owners={int(k): str(v) for k, v in (d.get("rank_owners") or {}).items()},
+        )
+
+
+def build_commit_marker(manifest_sha256: str, rank_owners: Mapping[int, str]) -> CommitMarker:
+    return CommitMarker(
+        manifest_sha256=manifest_sha256,
+        time_unix_nano=time.time_ns(),
+        rank_owners={int(k): str(v) for k, v in rank_owners.items()},
+    )
+
+
+def detect_partial_writes(
+    manifest_time_unix_nano: int,
+    shard_last_modified_unix_nano: Mapping[str, int],
+) -> List[str]:
+    """Return the list of shard names whose `LastModified` is newer
+    than the manifest's authoring time. Such shards are an immediate
+    signal of an in-flight async checkpoint.
+    """
+    return [
+        name
+        for name, lmtime in shard_last_modified_unix_nano.items()
+        if lmtime > manifest_time_unix_nano
+    ]
+
+
+def validate_after_load(
+    shards: Mapping[str, bytes],
+    expected_hashes: Mapping[str, str],
+    algorithm: str = "sha256",
+) -> List[str]:
+    """Hook called by a training framework after `load_state_dict`.
+
+    Returns the list of shard names that failed the post-load integrity
+    check. Empty list = all good.
+    """
+    failures: List[str] = []
+    for name, data in shards.items():
+        expected = expected_hashes.get(name)
+        if expected is None:
+            continue
+        actual = hash_bytes(data, algorithm)
+        if actual != expected:
+            failures.append(name)
+    return failures
+
+
+def canonical_manifest_bytes(manifest: Mapping[str, Any]) -> bytes:
+    """Stable serialization of a manifest for hashing.
+
+    Sort keys recursively so the SHA-256 is stable across Python
+    versions and platforms.
+    """
+    return _stable_dumps(manifest).encode("utf-8")
+
+
+def _stable_dumps(value: Any) -> str:
+    if isinstance(value, dict):
+        items = sorted(((json_quote(str(k)), _stable_dumps(v)) for k, v in value.items()))
+        return "{" + ",".join(f"{k}:{v}" for k, v in items) + "}"
+    if isinstance(value, (list, tuple)):
+        return "[" + ",".join(_stable_dumps(item) for item in value) + "]"
+    if value is None:
+        return "null"
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, (int, float)):
+        return str(value)
+    if isinstance(value, str):
+        return json_quote(value)
+    return json_quote(str(value))
+
+
+def json_quote(s: str) -> str:
+    return "\"" + s.replace("\\", "\\\\").replace("\"", "\\\"") + "\""
+
+
+__all__ = [
+    "SUPPORTED_HASHES",
+    "crc32c",
+    "xxh3_128_simulated",
+    "hash_bytes",
+    "CommitMarker",
+    "build_commit_marker",
+    "detect_partial_writes",
+    "validate_after_load",
+    "canonical_manifest_bytes",
+]

denpex_sdk/frameworks.py

@@ -0,0 +1,231 @@
+"""Framework-native instrumentation helpers.
+
+These wrappers extract the per-framework telemetry that the spec
+requires without taking a hard dependency on each framework. Each
+helper detects whether the framework is importable and silently
+returns a no-op context manager when it isn't, so the SDK can be
+shipped to environments that don't have DeepSpeed, Megatron, etc.
+
+In production each helper would `import deepspeed`/`import megatron`/
+etc. lazily and emit `TrainingEvent::DeepSpeed` / `MegatronLM` / `FSDP`
+records into the same OTLP envelope.
+"""
+from __future__ import annotations
+
+import contextlib
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterator, Mapping, Optional, Sequence
+
+
+@dataclass
+class FrameworkEvent:
+    framework: str
+    event: str
+    step: int
+    ts_unix_nano: int
+    attributes: Mapping[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "framework": self.framework,
+            "event": self.event,
+            "step": self.step,
+            "ts_unix_nano": self.ts_unix_nano,
+            "attributes": dict(self.attributes),
+        }
+
+
+@dataclass
+class DeepSpeedSnapshot:
+    zero_stage: Optional[int] = None
+    offload_optimizer: Optional[bool] = None
+    offload_param: Optional[bool] = None
+    communication_data_type: Optional[str] = None
+    allgather_partitions: Optional[bool] = None
+    reduce_scatter: Optional[bool] = None
+    overlap_comm: Optional[bool] = None
+    partition_count: Optional[int] = None
+
+
+def parse_deepspeed_config(config: Mapping[str, Any]) -> DeepSpeedSnapshot:
+    """Best-effort parser of a DeepSpeed JSON config.
+
+    Walks the common paths in the spec: `zero_optimization.stage`,
+    `zero_optimization.offload_optimizer`, `communication_data_type`,
+    etc. Returns a snapshot with whatever fields it can extract.
+    """
+    snapshot = DeepSpeedSnapshot()
+    if "zero_optimization" in config:
+        z = config["zero_optimization"]
+        if isinstance(z, dict):
+            if "stage" in z:
+                try:
+                    snapshot.zero_stage = int(z["stage"])
+                except (TypeError, ValueError):
+                    pass
+            if "offload_optimizer" in z and isinstance(z["offload_optimizer"], dict):
+                snapshot.offload_optimizer = True
+            if "offload_param" in z and isinstance(z["offload_param"], dict):
+                snapshot.offload_param = True
+            if "allgather_partitions" in z:
+                snapshot.allgather_partitions = bool(z["allgather_partitions"])
+            if "reduce_scatter" in z:
+                snapshot.reduce_scatter = bool(z["reduce_scatter"])
+            if "overlap_comm" in z:
+                snapshot.overlap_comm = bool(z["overlap_comm"])
+    if "communication_data_type" in config:
+        snapshot.communication_data_type = str(config["communication_data_type"])
+    if "zero_optimization" in config and isinstance(config["zero_optimization"], dict):
+        z = config["zero_optimization"]
+        if "stage" in z:
+            try:
+                snapshot.partition_count = int(z["stage"]) + 1
+            except (TypeError, ValueError):
+                pass
+    return snapshot
+
+
+@dataclass
+class MegatronSnapshot:
+    tensor_parallel_size: Optional[int] = None
+    pipeline_parallel_size: Optional[int] = None
+    context_parallel_size: Optional[int] = None
+    sequence_parallel_size: Optional[int] = None
+    virtual_pipeline_model_parallel_size: Optional[int] = None
+    pipeline_bubble_seconds: Optional[float] = None
+
+
+def parse_megatron_env(environ: Mapping[str, str]) -> MegatronSnapshot:
+    """Best-effort parser of Megatron's distributed-optimizer env vars.
+
+    Production callers would inspect `mpu` directly; this helper works
+    in environments where Megatron isn't installed.
+    """
+    s = MegatronSnapshot()
+    def _int(name: str) -> Optional[int]:
+        v = environ.get(name)
+        if v is None:
+            return None
+        try:
+            return int(v)
+        except ValueError:
+            return None
+    s.tensor_parallel_size = _int("MEGATRON_TENSOR_PARALLEL_SIZE")
+    s.pipeline_parallel_size = _int("MEGATRON_PIPELINE_PARALLEL_SIZE")
+    s.context_parallel_size = _int("MEGATRON_CONTEXT_PARALLEL_SIZE")
+    s.sequence_parallel_size = _int("MEGATRON_SEQUENCE_PARALLEL_SIZE")
+    s.virtual_pipeline_model_parallel_size = _int("MEGATRON_VPP_SIZE")
+    return s
+
+
+@dataclass
+class FsdpSnapshot:
+    sharding_strategy: Optional[str] = None
+    mixed_precision: Optional[bool] = None
+    cpu_offload: Optional[bool] = None
+    reshard_after_forward: Optional[bool] = None
+    state_dict_type: Optional[str] = None
+    use_orig_params: Optional[bool] = None
+    flat_param_count: Optional[int] = None
+
+
+def parse_fsdp_kwargs(kwargs: Mapping[str, Any]) -> FsdpSnapshot:
+    snap = FsdpSnapshot()
+    if "sharding_strategy" in kwargs:
+        snap.sharding_strategy = str(kwargs["sharding_strategy"])
+    if "mixed_precision" in kwargs:
+        snap.mixed_precision = bool(kwargs["mixed_precision"])
+    if "cpu_offload" in kwargs and isinstance(kwargs["cpu_offload"], (bool, dict)):
+        snap.cpu_offload = bool(kwargs["cpu_offload"])
+    if "reshard_after_forward" in kwargs:
+        snap.reshard_after_forward = bool(kwargs["reshard_after_forward"])
+    if "state_dict_type" in kwargs:
+        snap.state_dict_type = str(kwargs["state_dict_type"])
+    if "use_orig_params" in kwargs:
+        snap.use_orig_params = bool(kwargs["use_orig_params"])
+    return snap
+
+
+@dataclass
+class JaxSnapshot:
+    pjit_compile_time_seconds: Optional[float] = None
+    device_count: Optional[int] = None
+    sharding_spec: Optional[str] = None
+    backend: Optional[str] = None
+    has_pjit: bool = False
+    has_pmap: bool = False
+
+
+def parse_jax_config(environ: Mapping[str, str], pjit_seen: bool = False, pmap_seen: bool = False) -> JaxSnapshot:
+    snap = JaxSnapshot()
+    snap.backend = environ.get("JAX_PLATFORMS") or environ.get("JAX_PLATFORM_NAME")
+    snap.has_pjit = pjit_seen
+    snap.has_pmap = pmap_seen
+    if "XLA_FLAGS" in environ:
+        snap.sharding_spec = environ["XLA_FLAGS"]
+    return snap
+
+
+@dataclass
+class TensorFlowSnapshot:
+    strategy: Optional[str] = None
+    cross_device_ops: Optional[str] = None
+    experimental_numa_aware: Optional[bool] = None
+    jit_compile: Optional[bool] = None
+
+
+def parse_tf_config(environ: Mapping[str, str]) -> TensorFlowSnapshot:
+    snap = TensorFlowSnapshot()
+    snap.strategy = environ.get("TF_DISTRIBUTE_STRATEGY")
+    snap.cross_device_ops = environ.get("TF_CROSS_DEVICE_OPS")
+    return snap
+
+
+@contextlib.contextmanager
+def framework_span(framework: str, event: str, *, step: int = 0) -> Iterator[FrameworkEvent]:
+    """Lightweight timing wrapper. Emits a `FrameworkEvent` after the
+    block exits. No-op for unknown frameworks.
+    """
+    started = time.time_ns()
+    fe = FrameworkEvent(framework=framework, event=event, step=step, ts_unix_nano=started)
+    try:
+        yield fe
+    finally:
+        fe.ts_unix_nano = time.time_ns() - started
+
+
+def environment_snapshot() -> Dict[str, Optional[str]]:
+    """Capture a process's relevant env vars for the OTLP envelope.
+
+    Doesn't read `/proc/<pid>/environ` directly — production callers
+    pass the env into this function. Used by tests and by the
+    framework integrations.
+    """
+    keys = [
+        "RANK", "LOCAL_RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT",
+        "NCCL_DEBUG", "NCCL_IB_HCA", "NCCL_P2P_LEVEL", "NCCL_NET_GDR_LEVEL",
+        "CUDA_VISIBLE_DEVICES", "TORCH_DISTRIBUTED_DEBUG",
+        "MEGATRON_TENSOR_PARALLEL_SIZE", "MEGATRON_PIPELINE_PARALLEL_SIZE",
+        "JAX_PLATFORMS", "JAX_PLATFORM_NAME", "XLA_FLAGS",
+        "TF_DISTRIBUTE_STRATEGY", "TF_CROSS_DEVICE_OPS",
+        "RAY_ADDRESS", "RAY_NAMESPACE",
+    ]
+    return {k: os.environ.get(k) for k in keys}
+
+
+def detect_loaded_frameworks() -> Dict[str, bool]:
+    """Best-effort detection of which framework modules are importable.
+
+    Used by `LayerInstrument` to decide which `framework_span` to emit.
+    Never raises.
+    """
+    out: Dict[str, bool] = {}
+    for name in ("torch", "torch.distributed", "deepspeed", "megatron", "flax", "tensorflow", "jax", "ray"):
+        try:
+            __import__(name)
+            out[name] = True
+        except Exception:  # noqa: BLE001
+            out[name] = False
+    return out