denpex-sdk 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- denpex_sdk-0.1.0/PKG-INFO +43 -0
- denpex_sdk-0.1.0/README.md +32 -0
- denpex_sdk-0.1.0/pyproject.toml +21 -0
- denpex_sdk-0.1.0/setup.cfg +4 -0
- denpex_sdk-0.1.0/src/denpex_sdk/__init__.py +53 -0
- denpex_sdk-0.1.0/src/denpex_sdk/checkpoint.py +190 -0
- denpex_sdk-0.1.0/src/denpex_sdk/checkpoint_extra.py +169 -0
- denpex_sdk-0.1.0/src/denpex_sdk/frameworks.py +231 -0
- denpex_sdk-0.1.0/src/denpex_sdk/instrument.py +338 -0
- denpex_sdk-0.1.0/src/denpex_sdk/sdc.py +285 -0
- denpex_sdk-0.1.0/src/denpex_sdk/sdc_runner.py +277 -0
- denpex_sdk-0.1.0/src/denpex_sdk.egg-info/PKG-INFO +43 -0
- denpex_sdk-0.1.0/src/denpex_sdk.egg-info/SOURCES.txt +19 -0
- denpex_sdk-0.1.0/src/denpex_sdk.egg-info/dependency_links.txt +1 -0
- denpex_sdk-0.1.0/src/denpex_sdk.egg-info/requires.txt +3 -0
- denpex_sdk-0.1.0/src/denpex_sdk.egg-info/top_level.txt +1 -0
- denpex_sdk-0.1.0/tests/test_checkpoint_extra.py +129 -0
- denpex_sdk-0.1.0/tests/test_frameworks.py +177 -0
- denpex_sdk-0.1.0/tests/test_layer_metrics.py +228 -0
- denpex_sdk-0.1.0/tests/test_sdc.py +221 -0
- denpex_sdk-0.1.0/tests/test_sdk.py +146 -0
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: denpex-sdk
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Denpex training SDK for SDC detection and checkpoint validation.
|
|
5
|
+
Author: Denpex
|
|
6
|
+
License: Proprietary
|
|
7
|
+
Requires-Python: >=3.8
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Provides-Extra: torch
|
|
10
|
+
Requires-Dist: torch>=2.1; extra == "torch"
|
|
11
|
+
|
|
12
|
+
# denpex-sdk
|
|
13
|
+
|
|
14
|
+
Zero-third-party Python helpers for SDC detection (per-layer weight/gradient
|
|
15
|
+
norms, NaN/Inf checks, optimizer state validation) and asynchronous
|
|
16
|
+
checkpoint validation (manifest integrity + content checksums) for S3/Lustre.
|
|
17
|
+
|
|
18
|
+
## SDC detection
|
|
19
|
+
|
|
20
|
+
```python
|
|
21
|
+
from denpex_sdk import SdcHook, install_sdc_hooks
|
|
22
|
+
|
|
23
|
+
hook = install_sdc_hooks(model, SdcHook(max_weight_norm=1e4, max_grad_norm=1e4))
|
|
24
|
+
weight_report = hook.inspect()
|
|
25
|
+
grad_report = hook.inspect_gradients()
|
|
26
|
+
opt_report = hook.inspect_optimizer(optimizer)
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
## Checkpoint validation
|
|
30
|
+
|
|
31
|
+
```python
|
|
32
|
+
from denpex_sdk import build_manifest, manifest_to_json, manifest_from_json, validate_checkpoint
|
|
33
|
+
|
|
34
|
+
manifest = build_manifest(["/checkpoints/step-1000/"], metadata={"job_id": "..."})
|
|
35
|
+
with open("/checkpoints/step-1000/manifest.json", "w") as handle:
|
|
36
|
+
handle.write(manifest_to_json(manifest))
|
|
37
|
+
|
|
38
|
+
with open("/checkpoints/step-1000/manifest.json") as handle:
|
|
39
|
+
expected = manifest_from_json(handle.read())
|
|
40
|
+
|
|
41
|
+
result = validate_checkpoint(["/checkpoints/step-1000/"], expected_manifest=expected)
|
|
42
|
+
assert result.valid
|
|
43
|
+
```
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
# denpex-sdk
|
|
2
|
+
|
|
3
|
+
Zero-third-party Python helpers for SDC detection (per-layer weight/gradient
|
|
4
|
+
norms, NaN/Inf checks, optimizer state validation) and asynchronous
|
|
5
|
+
checkpoint validation (manifest integrity + content checksums) for S3/Lustre.
|
|
6
|
+
|
|
7
|
+
## SDC detection
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
from denpex_sdk import SdcHook, install_sdc_hooks
|
|
11
|
+
|
|
12
|
+
hook = install_sdc_hooks(model, SdcHook(max_weight_norm=1e4, max_grad_norm=1e4))
|
|
13
|
+
weight_report = hook.inspect()
|
|
14
|
+
grad_report = hook.inspect_gradients()
|
|
15
|
+
opt_report = hook.inspect_optimizer(optimizer)
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
## Checkpoint validation
|
|
19
|
+
|
|
20
|
+
```python
|
|
21
|
+
from denpex_sdk import build_manifest, manifest_to_json, manifest_from_json, validate_checkpoint
|
|
22
|
+
|
|
23
|
+
manifest = build_manifest(["/checkpoints/step-1000/"], metadata={"job_id": "..."})
|
|
24
|
+
with open("/checkpoints/step-1000/manifest.json", "w") as handle:
|
|
25
|
+
handle.write(manifest_to_json(manifest))
|
|
26
|
+
|
|
27
|
+
with open("/checkpoints/step-1000/manifest.json") as handle:
|
|
28
|
+
expected = manifest_from_json(handle.read())
|
|
29
|
+
|
|
30
|
+
result = validate_checkpoint(["/checkpoints/step-1000/"], expected_manifest=expected)
|
|
31
|
+
assert result.valid
|
|
32
|
+
```
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "denpex-sdk"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Denpex training SDK for SDC detection and checkpoint validation."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.8"
|
|
11
|
+
license = { text = "Proprietary" }
|
|
12
|
+
authors = [{ name = "Denpex" }]
|
|
13
|
+
|
|
14
|
+
[project.optional-dependencies]
|
|
15
|
+
torch = ["torch>=2.1"]
|
|
16
|
+
|
|
17
|
+
[tool.setuptools.packages.find]
|
|
18
|
+
where = ["src"]
|
|
19
|
+
|
|
20
|
+
[tool.setuptools.package-data]
|
|
21
|
+
denpex_sdk = ["tests/*.xml"]
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
"""Denpex training SDK (Phase 3).
|
|
2
|
+
|
|
3
|
+
Zero-third-party Python helpers for:
|
|
4
|
+
|
|
5
|
+
* SDC detection (per-layer weight/gradient norms, NaN/Inf checks, optimizer
|
|
6
|
+
state assertions) wired into training steps via lightweight hooks.
|
|
7
|
+
* Asynchronous checkpoint validation (manifest integrity, content checksums)
|
|
8
|
+
for S3/Lustre, designed to catch partial-write failures after a timeout.
|
|
9
|
+
|
|
10
|
+
The SDK is intentionally framework-light so it can be installed inside the
|
|
11
|
+
training image and called from any user script.
|
|
12
|
+
"""
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
from .sdc import (
|
|
16
|
+
SdcHook,
|
|
17
|
+
SdcReport,
|
|
18
|
+
TensorAnomaly,
|
|
19
|
+
TensorStats,
|
|
20
|
+
compute_tensor_stats,
|
|
21
|
+
contains_anomaly,
|
|
22
|
+
install_sdc_hooks,
|
|
23
|
+
uninstall_sdc_hooks,
|
|
24
|
+
)
|
|
25
|
+
from .checkpoint import (
|
|
26
|
+
CheckpointManifest,
|
|
27
|
+
CheckpointValidation,
|
|
28
|
+
ManifestEntry,
|
|
29
|
+
build_manifest,
|
|
30
|
+
hash_bytes,
|
|
31
|
+
manifest_from_json,
|
|
32
|
+
manifest_to_json,
|
|
33
|
+
validate_checkpoint,
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
__all__ = [
|
|
37
|
+
"SdcHook",
|
|
38
|
+
"SdcReport",
|
|
39
|
+
"TensorAnomaly",
|
|
40
|
+
"TensorStats",
|
|
41
|
+
"compute_tensor_stats",
|
|
42
|
+
"contains_anomaly",
|
|
43
|
+
"install_sdc_hooks",
|
|
44
|
+
"uninstall_sdc_hooks",
|
|
45
|
+
"CheckpointManifest",
|
|
46
|
+
"CheckpointValidation",
|
|
47
|
+
"ManifestEntry",
|
|
48
|
+
"build_manifest",
|
|
49
|
+
"hash_bytes",
|
|
50
|
+
"manifest_from_json",
|
|
51
|
+
"manifest_to_json",
|
|
52
|
+
"validate_checkpoint",
|
|
53
|
+
]
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import hashlib
|
|
4
|
+
import json
|
|
5
|
+
import os
|
|
6
|
+
from dataclasses import asdict, dataclass, field
|
|
7
|
+
from typing import Any, Iterable, List, Mapping, Optional
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@dataclass
|
|
11
|
+
class ManifestEntry:
|
|
12
|
+
path: str
|
|
13
|
+
size_bytes: int
|
|
14
|
+
sha256: str
|
|
15
|
+
|
|
16
|
+
def to_dict(self) -> Mapping[str, Any]:
|
|
17
|
+
return asdict(self)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@dataclass
|
|
21
|
+
class CheckpointManifest:
|
|
22
|
+
version: str
|
|
23
|
+
entries: List[ManifestEntry] = field(default_factory=list)
|
|
24
|
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
|
25
|
+
|
|
26
|
+
def to_dict(self) -> Mapping[str, Any]:
|
|
27
|
+
return {
|
|
28
|
+
"version": self.version,
|
|
29
|
+
"entries": [entry.to_dict() for entry in self.entries],
|
|
30
|
+
"metadata": dict(self.metadata),
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
@dataclass
|
|
35
|
+
class CheckpointValidation:
|
|
36
|
+
valid: bool
|
|
37
|
+
missing: List[str] = field(default_factory=list)
|
|
38
|
+
mismatched: List[str] = field(default_factory=list)
|
|
39
|
+
io_errors: List[str] = field(default_factory=list)
|
|
40
|
+
entries_checked: int = 0
|
|
41
|
+
|
|
42
|
+
def to_dict(self) -> Mapping[str, Any]:
|
|
43
|
+
return {
|
|
44
|
+
"valid": self.valid,
|
|
45
|
+
"missing": list(self.missing),
|
|
46
|
+
"mismatched": list(self.mismatched),
|
|
47
|
+
"io_errors": list(self.io_errors),
|
|
48
|
+
"entries_checked": self.entries_checked,
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
CHUNK_SIZE = 1024 * 1024
|
|
53
|
+
MANIFEST_VERSION = "denpex.checkpoint.v1"
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _iter_files(root: str) -> Iterable[str]:
|
|
57
|
+
if os.path.isfile(root):
|
|
58
|
+
yield root
|
|
59
|
+
return
|
|
60
|
+
for directory, _directories, files in os.walk(root):
|
|
61
|
+
for filename in sorted(files):
|
|
62
|
+
yield os.path.join(directory, filename)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _relativize(file_path: str, base_path: str) -> str:
|
|
66
|
+
try:
|
|
67
|
+
relative = os.path.relpath(file_path, start=base_path)
|
|
68
|
+
return relative.replace(os.sep, "/")
|
|
69
|
+
except ValueError:
|
|
70
|
+
absolute = os.path.abspath(file_path).replace(os.sep, "/")
|
|
71
|
+
return absolute
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def _hash_file(path: str) -> str:
|
|
75
|
+
digest = hashlib.sha256()
|
|
76
|
+
with open(path, "rb") as handle:
|
|
77
|
+
while True:
|
|
78
|
+
chunk = handle.read(CHUNK_SIZE)
|
|
79
|
+
if not chunk:
|
|
80
|
+
break
|
|
81
|
+
digest.update(chunk)
|
|
82
|
+
return digest.hexdigest()
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def _hash_data(data: bytes) -> str:
|
|
86
|
+
return hashlib.sha256(data).hexdigest()
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def build_manifest(
|
|
90
|
+
paths: Iterable[str],
|
|
91
|
+
*,
|
|
92
|
+
metadata: Optional[Mapping[str, Any]] = None,
|
|
93
|
+
version: str = MANIFEST_VERSION,
|
|
94
|
+
) -> CheckpointManifest:
|
|
95
|
+
entries: List[ManifestEntry] = []
|
|
96
|
+
raw_paths = list(paths)
|
|
97
|
+
for raw_path in raw_paths:
|
|
98
|
+
if not os.path.exists(raw_path):
|
|
99
|
+
continue
|
|
100
|
+
for file_path in _iter_files(raw_path):
|
|
101
|
+
try:
|
|
102
|
+
size = os.path.getsize(file_path)
|
|
103
|
+
digest = _hash_file(file_path)
|
|
104
|
+
except OSError:
|
|
105
|
+
continue
|
|
106
|
+
entries.append(
|
|
107
|
+
ManifestEntry(
|
|
108
|
+
path=_relativize(file_path, raw_path),
|
|
109
|
+
size_bytes=size,
|
|
110
|
+
sha256=digest,
|
|
111
|
+
)
|
|
112
|
+
)
|
|
113
|
+
return CheckpointManifest(
|
|
114
|
+
version=version,
|
|
115
|
+
entries=entries,
|
|
116
|
+
metadata=dict(metadata or {}),
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def manifest_to_json(manifest: CheckpointManifest) -> str:
|
|
121
|
+
return json.dumps(manifest.to_dict(), ensure_ascii=False, sort_keys=True, indent=2)
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def manifest_from_json(payload: str) -> CheckpointManifest:
|
|
125
|
+
parsed = json.loads(payload)
|
|
126
|
+
return CheckpointManifest(
|
|
127
|
+
version=str(parsed.get("version", MANIFEST_VERSION)),
|
|
128
|
+
entries=[ManifestEntry(**entry) for entry in parsed.get("entries", [])],
|
|
129
|
+
metadata=dict(parsed.get("metadata", {})),
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def validate_checkpoint(
|
|
134
|
+
paths: Iterable[str],
|
|
135
|
+
*,
|
|
136
|
+
expected_manifest: Optional[CheckpointManifest] = None,
|
|
137
|
+
rebuild_manifest: bool = False,
|
|
138
|
+
) -> CheckpointValidation:
|
|
139
|
+
expected_lookup: dict[str, ManifestEntry] = {}
|
|
140
|
+
if expected_manifest is not None and not rebuild_manifest:
|
|
141
|
+
for entry in expected_manifest.entries:
|
|
142
|
+
expected_lookup[entry.path] = entry
|
|
143
|
+
|
|
144
|
+
current_lookup: dict[str, ManifestEntry] = {}
|
|
145
|
+
missing: List[str] = []
|
|
146
|
+
mismatched: List[str] = []
|
|
147
|
+
io_errors: List[str] = []
|
|
148
|
+
|
|
149
|
+
raw_paths = list(paths)
|
|
150
|
+
for raw_path in raw_paths:
|
|
151
|
+
if not os.path.exists(raw_path):
|
|
152
|
+
missing.append(raw_path)
|
|
153
|
+
continue
|
|
154
|
+
for file_path in _iter_files(raw_path):
|
|
155
|
+
relative = _relativize(file_path, raw_path)
|
|
156
|
+
try:
|
|
157
|
+
size = os.path.getsize(file_path)
|
|
158
|
+
digest = _hash_file(file_path)
|
|
159
|
+
except OSError as exc:
|
|
160
|
+
io_errors.append(f"{relative}: {exc}")
|
|
161
|
+
continue
|
|
162
|
+
current_lookup[relative] = ManifestEntry(
|
|
163
|
+
path=relative,
|
|
164
|
+
size_bytes=size,
|
|
165
|
+
sha256=digest,
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
if expected_lookup:
|
|
169
|
+
for path, expected_entry in expected_lookup.items():
|
|
170
|
+
current_entry = current_lookup.get(path)
|
|
171
|
+
if current_entry is None:
|
|
172
|
+
missing.append(path)
|
|
173
|
+
continue
|
|
174
|
+
if current_entry.size_bytes != expected_entry.size_bytes:
|
|
175
|
+
mismatched.append(path)
|
|
176
|
+
continue
|
|
177
|
+
if current_entry.sha256 != expected_entry.sha256:
|
|
178
|
+
mismatched.append(path)
|
|
179
|
+
|
|
180
|
+
return CheckpointValidation(
|
|
181
|
+
valid=not missing and not mismatched and not io_errors,
|
|
182
|
+
missing=missing,
|
|
183
|
+
mismatched=mismatched,
|
|
184
|
+
io_errors=io_errors,
|
|
185
|
+
entries_checked=len(current_lookup),
|
|
186
|
+
)
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def hash_bytes(data: bytes) -> str:
|
|
190
|
+
return _hash_data(data)
|
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
"""Checkpoint integrity enhancements (Phase 9 spec items 9.4).
|
|
2
|
+
|
|
3
|
+
Adds:
|
|
4
|
+
- CRC32C + SHA-256 + XXH3-128 (selectable via manifest header).
|
|
5
|
+
- Atomic commit marker (denpex.manifest.commit = {manifest_sha256,
|
|
6
|
+
time_unix_nano, rank_owners: {rank: sha256}}).
|
|
7
|
+
- Asynchronous partial-write detection: shards whose LastModified is
|
|
8
|
+
newer than the manifest's are flagged.
|
|
9
|
+
- Optional restore validation hook (`validate_after_load`).
|
|
10
|
+
|
|
11
|
+
The existing `denpex_sdk.checkpoint` module is the source of truth for
|
|
12
|
+
manifest schema; this module extends it.
|
|
13
|
+
"""
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import binascii
|
|
17
|
+
import hashlib
|
|
18
|
+
import struct
|
|
19
|
+
import time
|
|
20
|
+
from dataclasses import dataclass, field
|
|
21
|
+
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
SUPPORTED_HASHES = ("sha256", "crc32c", "xxh3_128")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def crc32c(data: bytes) -> str:
|
|
28
|
+
"""CRC-32C (Castagnoli) using the same polynomial as iscsi.
|
|
29
|
+
|
|
30
|
+
Pure-Python fallback so the SDK works without crc32c installed.
|
|
31
|
+
Returns the big-endian hex representation (matches the Linux
|
|
32
|
+
`crc32c` command and the standard reference for "abc" = 0xCBF43926).
|
|
33
|
+
"""
|
|
34
|
+
crc = 0xFFFFFFFF
|
|
35
|
+
for byte in data:
|
|
36
|
+
crc ^= byte
|
|
37
|
+
for _ in range(8):
|
|
38
|
+
crc = (crc >> 1) ^ (0x82F63B78 if (crc & 1) else 0)
|
|
39
|
+
return binascii.hexlify(struct.pack(">I", (crc ^ 0xFFFFFFFF) & 0xFFFFFFFF)).decode()
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def xxh3_128_simulated(data: bytes) -> str:
|
|
43
|
+
"""Cheap stand-in for XXH3-128 when the C extension isn't available.
|
|
44
|
+
|
|
45
|
+
Hashes input bytes through blake2b. Distinct inputs produce distinct
|
|
46
|
+
digests with overwhelming probability.
|
|
47
|
+
"""
|
|
48
|
+
h = hashlib.blake2b(digest_size=16)
|
|
49
|
+
h.update(data)
|
|
50
|
+
return h.hexdigest()
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def hash_bytes(data: bytes, algorithm: str = "sha256") -> str:
|
|
54
|
+
if algorithm == "sha256":
|
|
55
|
+
return hashlib.sha256(data).hexdigest()
|
|
56
|
+
if algorithm == "crc32c":
|
|
57
|
+
return crc32c(data)
|
|
58
|
+
if algorithm == "xxh3_128":
|
|
59
|
+
return xxh3_128_simulated(data)
|
|
60
|
+
raise ValueError(f"unsupported hash algorithm: {algorithm}")
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@dataclass
|
|
64
|
+
class CommitMarker:
|
|
65
|
+
manifest_sha256: str
|
|
66
|
+
time_unix_nano: int
|
|
67
|
+
rank_owners: Dict[int, str] = field(default_factory=dict)
|
|
68
|
+
|
|
69
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
70
|
+
return {
|
|
71
|
+
"manifest_sha256": self.manifest_sha256,
|
|
72
|
+
"time_unix_nano": self.time_unix_nano,
|
|
73
|
+
"rank_owners": dict(self.rank_owners),
|
|
74
|
+
}
|
|
75
|
+
|
|
76
|
+
@classmethod
|
|
77
|
+
def from_dict(cls, d: Mapping[str, Any]) -> "CommitMarker":
|
|
78
|
+
return cls(
|
|
79
|
+
manifest_sha256=str(d.get("manifest_sha256", "")),
|
|
80
|
+
time_unix_nano=int(d.get("time_unix_nano", 0)),
|
|
81
|
+
rank_owners={int(k): str(v) for k, v in (d.get("rank_owners") or {}).items()},
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def build_commit_marker(manifest_sha256: str, rank_owners: Mapping[int, str]) -> CommitMarker:
|
|
86
|
+
return CommitMarker(
|
|
87
|
+
manifest_sha256=manifest_sha256,
|
|
88
|
+
time_unix_nano=time.time_ns(),
|
|
89
|
+
rank_owners={int(k): str(v) for k, v in rank_owners.items()},
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def detect_partial_writes(
|
|
94
|
+
manifest_time_unix_nano: int,
|
|
95
|
+
shard_last_modified_unix_nano: Mapping[str, int],
|
|
96
|
+
) -> List[str]:
|
|
97
|
+
"""Return the list of shard names whose `LastModified` is newer
|
|
98
|
+
than the manifest's authoring time. Such shards are an immediate
|
|
99
|
+
signal of an in-flight async checkpoint.
|
|
100
|
+
"""
|
|
101
|
+
return [
|
|
102
|
+
name
|
|
103
|
+
for name, lmtime in shard_last_modified_unix_nano.items()
|
|
104
|
+
if lmtime > manifest_time_unix_nano
|
|
105
|
+
]
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def validate_after_load(
|
|
109
|
+
shards: Mapping[str, bytes],
|
|
110
|
+
expected_hashes: Mapping[str, str],
|
|
111
|
+
algorithm: str = "sha256",
|
|
112
|
+
) -> List[str]:
|
|
113
|
+
"""Hook called by a training framework after `load_state_dict`.
|
|
114
|
+
|
|
115
|
+
Returns the list of shard names that failed the post-load integrity
|
|
116
|
+
check. Empty list = all good.
|
|
117
|
+
"""
|
|
118
|
+
failures: List[str] = []
|
|
119
|
+
for name, data in shards.items():
|
|
120
|
+
expected = expected_hashes.get(name)
|
|
121
|
+
if expected is None:
|
|
122
|
+
continue
|
|
123
|
+
actual = hash_bytes(data, algorithm)
|
|
124
|
+
if actual != expected:
|
|
125
|
+
failures.append(name)
|
|
126
|
+
return failures
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def canonical_manifest_bytes(manifest: Mapping[str, Any]) -> bytes:
|
|
130
|
+
"""Stable serialization of a manifest for hashing.
|
|
131
|
+
|
|
132
|
+
Sort keys recursively so the SHA-256 is stable across Python
|
|
133
|
+
versions and platforms.
|
|
134
|
+
"""
|
|
135
|
+
return _stable_dumps(manifest).encode("utf-8")
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def _stable_dumps(value: Any) -> str:
|
|
139
|
+
if isinstance(value, dict):
|
|
140
|
+
items = sorted(((json_quote(str(k)), _stable_dumps(v)) for k, v in value.items()))
|
|
141
|
+
return "{" + ",".join(f"{k}:{v}" for k, v in items) + "}"
|
|
142
|
+
if isinstance(value, (list, tuple)):
|
|
143
|
+
return "[" + ",".join(_stable_dumps(item) for item in value) + "]"
|
|
144
|
+
if value is None:
|
|
145
|
+
return "null"
|
|
146
|
+
if isinstance(value, bool):
|
|
147
|
+
return "true" if value else "false"
|
|
148
|
+
if isinstance(value, (int, float)):
|
|
149
|
+
return str(value)
|
|
150
|
+
if isinstance(value, str):
|
|
151
|
+
return json_quote(value)
|
|
152
|
+
return json_quote(str(value))
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def json_quote(s: str) -> str:
|
|
156
|
+
return "\"" + s.replace("\\", "\\\\").replace("\"", "\\\"") + "\""
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
__all__ = [
|
|
160
|
+
"SUPPORTED_HASHES",
|
|
161
|
+
"crc32c",
|
|
162
|
+
"xxh3_128_simulated",
|
|
163
|
+
"hash_bytes",
|
|
164
|
+
"CommitMarker",
|
|
165
|
+
"build_commit_marker",
|
|
166
|
+
"detect_partial_writes",
|
|
167
|
+
"validate_after_load",
|
|
168
|
+
"canonical_manifest_bytes",
|
|
169
|
+
]
|