cortex-training 0.0.3__py3-none-any.whl
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.
- cortex_training/__init__.py +49 -0
- cortex_training/__main__.py +20 -0
- cortex_training/_cli.py +1050 -0
- cortex_training/cli.py +103 -0
- cortex_training/client.py +2866 -0
- cortex_training/peft.py +112 -0
- cortex_training/telemetry.py +526 -0
- cortex_training/tui/__init__.py +23 -0
- cortex_training/tui/__main__.py +109 -0
- cortex_training/tui/app.py +646 -0
- cortex_training/tui/format.py +389 -0
- cortex_training/tui/log_cache.py +415 -0
- cortex_training/wire.py +414 -0
- cortex_training-0.0.3.dist-info/METADATA +143 -0
- cortex_training-0.0.3.dist-info/RECORD +19 -0
- cortex_training-0.0.3.dist-info/WHEEL +5 -0
- cortex_training-0.0.3.dist-info/entry_points.txt +3 -0
- cortex_training-0.0.3.dist-info/licenses/LICENSE +177 -0
- cortex_training-0.0.3.dist-info/top_level.txt +1 -0
cortex_training/wire.py
ADDED
|
@@ -0,0 +1,414 @@
|
|
|
1
|
+
# Copyright 2025 Snowflake Inc.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
"""Cortex Training binary wire protocol (DSSST1).
|
|
17
|
+
|
|
18
|
+
This is the single source of truth for how Cortex Training serializes binary
|
|
19
|
+
training payloads. ``ArcticRL`` imports
|
|
20
|
+
this module, so all three sides share one implementation.
|
|
21
|
+
The wire version and metadata keys below are server-defined protocol literals.
|
|
22
|
+
|
|
23
|
+
It owns three concerns behind one small surface:
|
|
24
|
+
|
|
25
|
+
1. Codec -- ``dumps`` / ``loads`` serialize a nested dict/list/tuple of tensors
|
|
26
|
+
and JSON-able data as a single **safetensors** blob. Pickle is never used, so
|
|
27
|
+
nothing can execute during decode (unlike ``torch.load``).
|
|
28
|
+
|
|
29
|
+
2. Operation metadata -- ``dumps(obj, metadata=...)`` stores small control-plane
|
|
30
|
+
metadata (router replay, chunk descriptors) inside the safetensors
|
|
31
|
+
``__metadata__`` header. ``read_metadata`` reads it back **without**
|
|
32
|
+
deserializing tensors. This replaces the old hand-rolled ``DSSMETA1`` byte
|
|
33
|
+
envelope -- there is now exactly one framing.
|
|
34
|
+
|
|
35
|
+
3. Byte chunking -- ``encode_byte_chunks`` splits an oversized DSSST1 frame into
|
|
36
|
+
self-describing byte-range frames. ``decode_byte_chunks`` reassembles them.
|
|
37
|
+
The split/merge contract lives in one place so client and server cannot
|
|
38
|
+
drift.
|
|
39
|
+
|
|
40
|
+
Wire layout of one frame is just a safetensors blob::
|
|
41
|
+
|
|
42
|
+
[ u64 header_len | JSON header (__metadata__: {dss: <structure>, op: <metadata>}) | tensor bytes ]
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
from __future__ import annotations
|
|
46
|
+
|
|
47
|
+
import hashlib
|
|
48
|
+
import json
|
|
49
|
+
from typing import Any
|
|
50
|
+
|
|
51
|
+
import torch
|
|
52
|
+
from safetensors.torch import load as st_load
|
|
53
|
+
from safetensors.torch import save as st_save
|
|
54
|
+
|
|
55
|
+
WIRE_FORMAT_VERSION = "DSSST1"
|
|
56
|
+
|
|
57
|
+
_STRUCT_KEY = "dss" # structure skeleton (internal)
|
|
58
|
+
_OP_KEY = "op" # operation metadata (router_replay, request/result chunks, ...)
|
|
59
|
+
_REQUEST_CHUNK_KEY = "request_chunk"
|
|
60
|
+
_RESULT_CHUNK_KEY = "result_chunk"
|
|
61
|
+
_PLACEHOLDER_KEY = "__dss_empty__"
|
|
62
|
+
_ROOT_KEY = "__root__"
|
|
63
|
+
_MAX_HEADER_BYTES = 64 * 1024 * 1024
|
|
64
|
+
_CHUNK_PAYLOAD_KEY = "payload"
|
|
65
|
+
_DEFAULT_CHUNK_OVERHEAD_BYTES = 8192
|
|
66
|
+
|
|
67
|
+
_PICKLE_SIGNATURES = (b"\x80\x02", b"\x80\x03", b"\x80\x04", b"\x80\x05")
|
|
68
|
+
_ZIP_SIGNATURE = b"PK\x03\x04"
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
class WireError(ValueError):
|
|
72
|
+
"""Raised for malformed or incompatible wire payloads."""
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def _is_tensor(obj: Any) -> bool:
|
|
76
|
+
return torch.is_tensor(obj)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
# ---------------------------------------------------------------------------
|
|
80
|
+
# Codec
|
|
81
|
+
# ---------------------------------------------------------------------------
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _subtree_has_tensor(obj: Any) -> bool:
|
|
85
|
+
if _is_tensor(obj):
|
|
86
|
+
return True
|
|
87
|
+
if isinstance(obj, dict):
|
|
88
|
+
return any(_subtree_has_tensor(v) for v in obj.values())
|
|
89
|
+
if isinstance(obj, (list, tuple)):
|
|
90
|
+
return any(_subtree_has_tensor(v) for v in obj)
|
|
91
|
+
return False
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def _encode(obj: Any, tensors: dict, seen: set, path: str) -> list:
|
|
95
|
+
if _is_tensor(obj):
|
|
96
|
+
key = path or _ROOT_KEY
|
|
97
|
+
tensor = obj.detach().cpu().contiguous()
|
|
98
|
+
node: list = ["t", key]
|
|
99
|
+
if tensor.dim() == 0:
|
|
100
|
+
tensor = tensor.reshape(1)
|
|
101
|
+
node = ["t", key, {"scalar": True}]
|
|
102
|
+
# safetensors rejects tensors that share storage; clone on collision.
|
|
103
|
+
if tensor.data_ptr() in seen and tensor.numel() > 0:
|
|
104
|
+
tensor = tensor.clone()
|
|
105
|
+
seen.add(tensor.data_ptr())
|
|
106
|
+
tensors[key] = tensor
|
|
107
|
+
return node
|
|
108
|
+
if isinstance(obj, dict):
|
|
109
|
+
if not _subtree_has_tensor(obj):
|
|
110
|
+
return ["j", obj]
|
|
111
|
+
items = []
|
|
112
|
+
for raw_key, value in obj.items():
|
|
113
|
+
if not isinstance(raw_key, str):
|
|
114
|
+
raise WireError(f"DSSST1 dict keys must be strings, got {type(raw_key)!r}")
|
|
115
|
+
child_path = f"{path}.{raw_key}" if path else raw_key
|
|
116
|
+
items.append([raw_key, _encode(value, tensors, seen, child_path)])
|
|
117
|
+
return ["d", items]
|
|
118
|
+
if isinstance(obj, tuple):
|
|
119
|
+
return ["u", [_encode(v, tensors, seen, f"{path}.{i}" if path else str(i)) for i, v in enumerate(obj)]]
|
|
120
|
+
if isinstance(obj, list):
|
|
121
|
+
if not _subtree_has_tensor(obj):
|
|
122
|
+
return ["j", obj]
|
|
123
|
+
return ["l", [_encode(v, tensors, seen, f"{path}.{i}" if path else str(i)) for i, v in enumerate(obj)]]
|
|
124
|
+
return ["j", obj]
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def dumps(obj: Any, *, metadata: dict | None = None) -> bytes:
|
|
128
|
+
"""Serialize ``obj`` to one safetensors frame, with optional op ``metadata``."""
|
|
129
|
+
tensors: dict = {}
|
|
130
|
+
tree = _encode(obj, tensors, set(), "")
|
|
131
|
+
if not tensors:
|
|
132
|
+
tensors[_PLACEHOLDER_KEY] = torch.zeros(1, dtype=torch.uint8)
|
|
133
|
+
header = {_STRUCT_KEY: json.dumps({"v": WIRE_FORMAT_VERSION, "tree": tree})}
|
|
134
|
+
if metadata:
|
|
135
|
+
header[_OP_KEY] = json.dumps(metadata)
|
|
136
|
+
return st_save(tensors, metadata=header)
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def _decode(node: list, tensors: dict) -> Any:
|
|
140
|
+
tag = node[0]
|
|
141
|
+
if tag == "t":
|
|
142
|
+
tensor = tensors[node[1]]
|
|
143
|
+
if len(node) > 2 and isinstance(node[2], dict) and node[2].get("scalar"):
|
|
144
|
+
return tensor.reshape(())
|
|
145
|
+
return tensor
|
|
146
|
+
if tag == "j":
|
|
147
|
+
return node[1]
|
|
148
|
+
if tag == "d":
|
|
149
|
+
return {key: _decode(child, tensors) for key, child in node[1]}
|
|
150
|
+
if tag == "l":
|
|
151
|
+
return [_decode(child, tensors) for child in node[1]]
|
|
152
|
+
if tag == "u":
|
|
153
|
+
return tuple(_decode(child, tensors) for child in node[1])
|
|
154
|
+
raise WireError(f"unknown DSSST1 node tag {tag!r}")
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def _read_header(data: bytes) -> dict:
|
|
158
|
+
if len(data) < 8:
|
|
159
|
+
raise WireError("payload too short to be a DSSST1 safetensors blob")
|
|
160
|
+
header_len = int.from_bytes(data[:8], "little")
|
|
161
|
+
if header_len <= 0 or 8 + header_len > len(data) or header_len > _MAX_HEADER_BYTES:
|
|
162
|
+
raise WireError("invalid DSSST1 safetensors header length")
|
|
163
|
+
header = json.loads(data[8 : 8 + header_len].decode("utf-8"))
|
|
164
|
+
return header.get("__metadata__", {}) or {}
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def loads(data: Any) -> Any:
|
|
168
|
+
"""Deserialize a frame produced by :func:`dumps`.
|
|
169
|
+
|
|
170
|
+
Never invokes pickle/``torch.load``; a legacy pickle or ``torch.save``
|
|
171
|
+
payload is rejected with a clear error instead of being executed.
|
|
172
|
+
"""
|
|
173
|
+
if not isinstance(data, (bytes, bytearray, memoryview)):
|
|
174
|
+
raise WireError("DSSST1 loads expects bytes")
|
|
175
|
+
data = bytes(data)
|
|
176
|
+
try:
|
|
177
|
+
tensors = st_load(data)
|
|
178
|
+
header = _read_header(data)
|
|
179
|
+
except Exception as exc: # not safetensors -> never feed it to pickle
|
|
180
|
+
if data[:2] in _PICKLE_SIGNATURES or data[:4] == _ZIP_SIGNATURE:
|
|
181
|
+
raise WireError(
|
|
182
|
+
"refusing to deserialize a legacy pickle/torch.save payload; "
|
|
183
|
+
"the peer must use the DSSST1 safetensors wire protocol"
|
|
184
|
+
) from exc
|
|
185
|
+
raise WireError(f"not a valid DSSST1 safetensors payload: {exc}") from exc
|
|
186
|
+
|
|
187
|
+
raw = header.get(_STRUCT_KEY)
|
|
188
|
+
if raw is None:
|
|
189
|
+
raise WireError("DSSST1 structure missing from safetensors header")
|
|
190
|
+
info = json.loads(raw)
|
|
191
|
+
if info.get("v") != WIRE_FORMAT_VERSION:
|
|
192
|
+
raise WireError(f"unsupported DSSST1 wire version {info.get('v')!r}")
|
|
193
|
+
tensors.pop(_PLACEHOLDER_KEY, None)
|
|
194
|
+
return _decode(info["tree"], tensors)
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
def read_metadata(data: Any) -> dict:
|
|
198
|
+
"""Return the operation metadata from a frame header without loading tensors.
|
|
199
|
+
|
|
200
|
+
This is a lenient peek: bytes that aren't a valid DSSST1 frame simply have no
|
|
201
|
+
operation metadata, so ``{}`` is returned. The authoritative (pickle-safe)
|
|
202
|
+
decode happens in :func:`loads`.
|
|
203
|
+
"""
|
|
204
|
+
if not isinstance(data, (bytes, bytearray, memoryview)):
|
|
205
|
+
raise WireError("DSSST1 read_metadata expects bytes")
|
|
206
|
+
try:
|
|
207
|
+
raw = _read_header(bytes(data)).get(_OP_KEY)
|
|
208
|
+
except (WireError, ValueError):
|
|
209
|
+
return {}
|
|
210
|
+
return json.loads(raw) if raw else {}
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
# ---------------------------------------------------------------------------
|
|
214
|
+
# Byte chunking -- split and reassemble DSSST1 frames
|
|
215
|
+
# ---------------------------------------------------------------------------
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
def _chunk_key(kind: str) -> str:
|
|
219
|
+
if kind == "request":
|
|
220
|
+
return _REQUEST_CHUNK_KEY
|
|
221
|
+
if kind == "result":
|
|
222
|
+
return _RESULT_CHUNK_KEY
|
|
223
|
+
raise WireError(f"unknown byte chunk kind {kind!r}")
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
def _bytes_to_tensor(data: bytes) -> torch.Tensor:
|
|
227
|
+
# bytearray makes the buffer writable; clone detaches from the temporary.
|
|
228
|
+
return torch.frombuffer(bytearray(data), dtype=torch.uint8).clone()
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def _tensor_to_bytes(tensor: Any) -> bytes:
|
|
232
|
+
if not torch.is_tensor(tensor):
|
|
233
|
+
raise WireError("byte chunk payload must be a tensor")
|
|
234
|
+
if tensor.dtype != torch.uint8:
|
|
235
|
+
raise WireError(f"byte chunk payload tensor must be uint8, got {tensor.dtype}")
|
|
236
|
+
if tensor.ndim != 1:
|
|
237
|
+
raise WireError(f"byte chunk payload tensor must be 1-D, got shape {tuple(tensor.shape)}")
|
|
238
|
+
return tensor.cpu().contiguous().numpy().tobytes()
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
def _make_byte_chunk_frame(
|
|
242
|
+
frame: bytes,
|
|
243
|
+
*,
|
|
244
|
+
kind: str,
|
|
245
|
+
operation: str | None,
|
|
246
|
+
group_id: str | None,
|
|
247
|
+
chunk_idx: int,
|
|
248
|
+
total_chunks: int,
|
|
249
|
+
start: int,
|
|
250
|
+
end: int,
|
|
251
|
+
) -> bytes:
|
|
252
|
+
key = _chunk_key(kind)
|
|
253
|
+
desc: dict[str, Any] = {
|
|
254
|
+
"chunk_idx": chunk_idx,
|
|
255
|
+
"total_chunks": total_chunks,
|
|
256
|
+
"frame_sha256": hashlib.sha256(frame).hexdigest(),
|
|
257
|
+
"frame_size_bytes": len(frame),
|
|
258
|
+
}
|
|
259
|
+
if group_id is not None:
|
|
260
|
+
desc["chunk_group_id"] = group_id
|
|
261
|
+
if operation is not None:
|
|
262
|
+
desc["operation"] = operation
|
|
263
|
+
return dumps({_CHUNK_PAYLOAD_KEY: _bytes_to_tensor(frame[start:end])}, metadata={key: desc})
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
def encode_byte_chunks(
|
|
267
|
+
frame: bytes,
|
|
268
|
+
*,
|
|
269
|
+
kind: str,
|
|
270
|
+
operation: str | None = None,
|
|
271
|
+
max_bytes: int = 0,
|
|
272
|
+
chunk_group_id: str | None = None,
|
|
273
|
+
force_chunk: bool = False,
|
|
274
|
+
) -> list[bytes]:
|
|
275
|
+
"""Split a DSSST1 frame into byte-range DSSST1 chunk frames.
|
|
276
|
+
|
|
277
|
+
Returns ``[frame]`` when no chunking is needed unless ``force_chunk`` is
|
|
278
|
+
true. Chunk frames contain a 1-D uint8 ``payload`` tensor plus metadata
|
|
279
|
+
under ``request_chunk`` or ``result_chunk``.
|
|
280
|
+
"""
|
|
281
|
+
if not isinstance(frame, (bytes, bytearray, memoryview)):
|
|
282
|
+
raise WireError("encode_byte_chunks expects frame bytes")
|
|
283
|
+
frame = bytes(frame)
|
|
284
|
+
if not force_chunk and (max_bytes <= 0 or len(frame) <= max_bytes):
|
|
285
|
+
return [frame]
|
|
286
|
+
|
|
287
|
+
group_id = chunk_group_id if kind == "request" else None
|
|
288
|
+
if kind == "request" and not group_id:
|
|
289
|
+
import uuid
|
|
290
|
+
|
|
291
|
+
group_id = str(uuid.uuid4())
|
|
292
|
+
|
|
293
|
+
if max_bytes <= 0:
|
|
294
|
+
return [
|
|
295
|
+
_make_byte_chunk_frame(
|
|
296
|
+
frame,
|
|
297
|
+
kind=kind,
|
|
298
|
+
operation=operation,
|
|
299
|
+
group_id=group_id,
|
|
300
|
+
chunk_idx=0,
|
|
301
|
+
total_chunks=1,
|
|
302
|
+
start=0,
|
|
303
|
+
end=len(frame),
|
|
304
|
+
)
|
|
305
|
+
]
|
|
306
|
+
|
|
307
|
+
payload_size = max(1, max_bytes - _DEFAULT_CHUNK_OVERHEAD_BYTES)
|
|
308
|
+
while True:
|
|
309
|
+
ranges = [(start, min(start + payload_size, len(frame))) for start in range(0, len(frame), payload_size)] or [
|
|
310
|
+
(0, 0)
|
|
311
|
+
]
|
|
312
|
+
total = len(ranges)
|
|
313
|
+
chunks = [
|
|
314
|
+
_make_byte_chunk_frame(
|
|
315
|
+
frame,
|
|
316
|
+
kind=kind,
|
|
317
|
+
operation=operation,
|
|
318
|
+
group_id=group_id,
|
|
319
|
+
chunk_idx=idx,
|
|
320
|
+
total_chunks=total,
|
|
321
|
+
start=start,
|
|
322
|
+
end=end,
|
|
323
|
+
)
|
|
324
|
+
for idx, (start, end) in enumerate(ranges)
|
|
325
|
+
]
|
|
326
|
+
if all(len(chunk) <= max_bytes for chunk in chunks):
|
|
327
|
+
return chunks
|
|
328
|
+
if payload_size == 1:
|
|
329
|
+
raise WireError("max_bytes is too small to hold one DSSST1 byte chunk frame")
|
|
330
|
+
payload_size = max(1, payload_size // 2)
|
|
331
|
+
|
|
332
|
+
|
|
333
|
+
def read_byte_chunk_metadata(frame: Any) -> dict | None:
|
|
334
|
+
"""Return byte chunk metadata with ``kind`` included, or ``None``."""
|
|
335
|
+
metadata = read_metadata(frame)
|
|
336
|
+
if _REQUEST_CHUNK_KEY in metadata:
|
|
337
|
+
desc = dict(metadata[_REQUEST_CHUNK_KEY])
|
|
338
|
+
desc["kind"] = "request"
|
|
339
|
+
return desc
|
|
340
|
+
if _RESULT_CHUNK_KEY in metadata:
|
|
341
|
+
desc = dict(metadata[_RESULT_CHUNK_KEY])
|
|
342
|
+
desc["kind"] = "result"
|
|
343
|
+
return desc
|
|
344
|
+
return None
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
def decode_byte_chunks(chunks: list[bytes], *, kind: str | None = None) -> bytes:
|
|
348
|
+
"""Reassemble byte-range DSSST1 chunks into the original DSSST1 frame."""
|
|
349
|
+
if not chunks:
|
|
350
|
+
raise WireError("decode_byte_chunks requires at least one frame")
|
|
351
|
+
if len(chunks) == 1 and read_byte_chunk_metadata(chunks[0]) is None:
|
|
352
|
+
return bytes(chunks[0])
|
|
353
|
+
|
|
354
|
+
descriptors = [read_byte_chunk_metadata(chunk) for chunk in chunks]
|
|
355
|
+
if any(desc is None for desc in descriptors):
|
|
356
|
+
raise WireError("decode_byte_chunks received a frame without byte chunk metadata")
|
|
357
|
+
descriptors = [desc for desc in descriptors if desc is not None]
|
|
358
|
+
actual_kind = descriptors[0]["kind"]
|
|
359
|
+
if kind is not None and actual_kind != kind:
|
|
360
|
+
raise WireError(f"expected {kind} chunks, got {actual_kind} chunks")
|
|
361
|
+
|
|
362
|
+
total = int(descriptors[0].get("total_chunks"))
|
|
363
|
+
frame_sha256 = descriptors[0].get("frame_sha256")
|
|
364
|
+
frame_size_bytes = int(descriptors[0].get("frame_size_bytes"))
|
|
365
|
+
group_id = descriptors[0].get("chunk_group_id")
|
|
366
|
+
operation = descriptors[0].get("operation")
|
|
367
|
+
|
|
368
|
+
if len(chunks) != total:
|
|
369
|
+
raise WireError(f"expected {total} byte chunks, got {len(chunks)}")
|
|
370
|
+
|
|
371
|
+
seen: set[int] = set()
|
|
372
|
+
ordered: list[tuple[int, bytes]] = []
|
|
373
|
+
for chunk, desc in zip(chunks, descriptors):
|
|
374
|
+
if desc["kind"] != actual_kind:
|
|
375
|
+
raise WireError("mixed byte chunk kinds")
|
|
376
|
+
if int(desc.get("total_chunks")) != total:
|
|
377
|
+
raise WireError("byte chunk total_chunks differs across frames")
|
|
378
|
+
if desc.get("frame_sha256") != frame_sha256:
|
|
379
|
+
raise WireError("byte chunk frame_sha256 differs across frames")
|
|
380
|
+
if int(desc.get("frame_size_bytes")) != frame_size_bytes:
|
|
381
|
+
raise WireError("byte chunk frame_size_bytes differs across frames")
|
|
382
|
+
if desc.get("chunk_group_id") != group_id:
|
|
383
|
+
raise WireError("byte chunk chunk_group_id differs across frames")
|
|
384
|
+
if desc.get("operation") != operation:
|
|
385
|
+
raise WireError("byte chunk operation differs across frames")
|
|
386
|
+
idx = int(desc.get("chunk_idx"))
|
|
387
|
+
if idx < 0 or idx >= total:
|
|
388
|
+
raise WireError("byte chunk_idx out of range")
|
|
389
|
+
if idx in seen:
|
|
390
|
+
raise WireError("duplicate byte chunk_idx")
|
|
391
|
+
seen.add(idx)
|
|
392
|
+
payload = loads(chunk).get(_CHUNK_PAYLOAD_KEY)
|
|
393
|
+
ordered.append((idx, _tensor_to_bytes(payload)))
|
|
394
|
+
|
|
395
|
+
if seen != set(range(total)):
|
|
396
|
+
missing = sorted(set(range(total)) - seen)
|
|
397
|
+
raise WireError(f"missing byte chunks: {missing}")
|
|
398
|
+
|
|
399
|
+
frame = b"".join(payload for _, payload in sorted(ordered))
|
|
400
|
+
if len(frame) != frame_size_bytes:
|
|
401
|
+
raise WireError("reassembled frame size mismatch")
|
|
402
|
+
if hashlib.sha256(frame).hexdigest() != frame_sha256:
|
|
403
|
+
raise WireError("reassembled frame sha256 mismatch")
|
|
404
|
+
return frame
|
|
405
|
+
|
|
406
|
+
|
|
407
|
+
def encode_result_chunks(result: Any, *, max_bytes: int = 0) -> list[bytes]:
|
|
408
|
+
"""Serialize and optionally byte-chunk a result object."""
|
|
409
|
+
return encode_byte_chunks(dumps(result), kind="result", max_bytes=max_bytes)
|
|
410
|
+
|
|
411
|
+
|
|
412
|
+
def decode_result_chunks(chunks: list[bytes]) -> Any:
|
|
413
|
+
"""Decode a result object returned by :func:`encode_result_chunks`."""
|
|
414
|
+
return loads(decode_byte_chunks(chunks, kind="result"))
|
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cortex-training
|
|
3
|
+
Version: 0.0.3
|
|
4
|
+
Summary: Python SDK and command-line tools for Cortex Training
|
|
5
|
+
License: Apache-2.0
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: requests>=2.28.0
|
|
10
|
+
Requires-Dist: tenacity>=8.2.0
|
|
11
|
+
Requires-Dist: snowflake-connector-python>=3.15.0
|
|
12
|
+
Requires-Dist: torch
|
|
13
|
+
Requires-Dist: safetensors
|
|
14
|
+
Requires-Dist: textual>=0.50
|
|
15
|
+
Provides-Extra: dev
|
|
16
|
+
Requires-Dist: pytest>=7.0.0; extra == "dev"
|
|
17
|
+
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
|
|
18
|
+
Requires-Dist: black>=22.0.0; extra == "dev"
|
|
19
|
+
Requires-Dist: flake8>=5.0.0; extra == "dev"
|
|
20
|
+
Requires-Dist: mypy>=0.990; extra == "dev"
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# Cortex Training Client
|
|
24
|
+
|
|
25
|
+
Python SDK, command-line tools, runnable recipes, and documentation for the
|
|
26
|
+
Cortex Training REST API.
|
|
27
|
+
|
|
28
|
+
## Install
|
|
29
|
+
|
|
30
|
+
Requires Python 3.10 or later and [uv](https://docs.astral.sh/uv/getting-started/installation/).
|
|
31
|
+
Create an environment and install the client:
|
|
32
|
+
|
|
33
|
+
```bash
|
|
34
|
+
uv venv
|
|
35
|
+
source .venv/bin/activate
|
|
36
|
+
uv pip install git+https://github.com/snowflakedb/cortex-training.git
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
The package includes the CLI, log TUI, and Python SDK. `pip install` also works
|
|
40
|
+
in place of `uv pip install` in an active Python environment.
|
|
41
|
+
|
|
42
|
+
## Log In
|
|
43
|
+
|
|
44
|
+
Create `~/cortex-training-config.json` using the
|
|
45
|
+
[connection template](examples/config/connection.json.template). Set `host` to
|
|
46
|
+
your Snowflake account hostname, `pat` to your programmatic access token, and
|
|
47
|
+
`database` and `schema` to the location of your Cortex Training endpoint.
|
|
48
|
+
Keep this file outside the repository and do not commit credentials.
|
|
49
|
+
|
|
50
|
+
Login validates the config and remembers its path for future CLI commands:
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
cortex-training login ~/cortex-training-config.json
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
The equivalent `cortex-training login --config ~/cortex-training-config.json`
|
|
57
|
+
form is also supported.
|
|
58
|
+
|
|
59
|
+
`ct` is an alias for `cortex-training`: every CLI example also works with `ct`.
|
|
60
|
+
See [connection setup](docs/getting-started/setup.md) for more detail.
|
|
61
|
+
|
|
62
|
+
## Try the CLI
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
cortex-training capacity # Check available GPU capacity
|
|
66
|
+
cortex-training list # List jobs
|
|
67
|
+
cortex-training get JOB_ID # Inspect a job from the list
|
|
68
|
+
cortex-training tui # Pick a job and view its logs
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
See the [CLI quick reference](docs/reference/cli.md#quick-reference) for
|
|
72
|
+
submission, training, generation, checkpoints, and log downloads.
|
|
73
|
+
|
|
74
|
+
## Run a Recipe
|
|
75
|
+
|
|
76
|
+
Start with [Run a Quick SFT Job](docs/getting-started/first-sft-run.md), a short
|
|
77
|
+
end-to-end supervised fine-tuning walkthrough.
|
|
78
|
+
|
|
79
|
+
Recipes require a repository checkout; they are not included in the installed
|
|
80
|
+
package:
|
|
81
|
+
|
|
82
|
+
```bash
|
|
83
|
+
git clone https://github.com/snowflakedb/cortex-training.git
|
|
84
|
+
cd cortex-training
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
Install the [recipe dependencies](recipes/README.md#prerequisites), then run
|
|
88
|
+
the chosen recipe's commands from the repository root. Recipes take an explicit
|
|
89
|
+
`config=/path/to/config.json` argument; use the same connection file you logged
|
|
90
|
+
in with.
|
|
91
|
+
|
|
92
|
+
| Task | Recipe |
|
|
93
|
+
|---|---|
|
|
94
|
+
| Fine-tune a chat model with LoRA or full-parameter training | [Conversational SFT](recipes/sft/conversational/README.md) |
|
|
95
|
+
| Train math reasoning with reinforcement learning | [Math GRPO](recipes/rl/math_grpo/README.md) |
|
|
96
|
+
| Serve a model or checkpoint, generate responses, and evaluate | [Inference endpoint](recipes/inference/README.md) |
|
|
97
|
+
|
|
98
|
+
Check `cortex-training capacity` and the recipe's GPU requirements before
|
|
99
|
+
starting a run. Browse the [recipe index](recipes/README.md) for all workflows.
|
|
100
|
+
|
|
101
|
+
## More Documentation
|
|
102
|
+
|
|
103
|
+
- [Getting started and prerequisites](docs/getting-started/README.md)
|
|
104
|
+
- [CLI commands and configuration](docs/reference/cli.md)
|
|
105
|
+
- [Python SDK](docs/reference/python-sdk.md)
|
|
106
|
+
- [REST API](docs/reference/rest-api.md)
|
|
107
|
+
- [Model and training-method compatibility](docs/reference/model-compatibility.md)
|
|
108
|
+
- [Job management and cluster status](docs/guides/operations/manage-jobs.md)
|
|
109
|
+
|
|
110
|
+
## Development
|
|
111
|
+
|
|
112
|
+
### Editable Install
|
|
113
|
+
|
|
114
|
+
From a repository checkout, with your Python environment active:
|
|
115
|
+
|
|
116
|
+
```bash
|
|
117
|
+
uv pip install -e ".[dev]"
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
### Build a Wheel
|
|
121
|
+
|
|
122
|
+
With `uv` installed, run:
|
|
123
|
+
|
|
124
|
+
```bash
|
|
125
|
+
./scripts/build_wheel.sh
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
The script builds the package using `pyproject.toml` in an isolated build
|
|
129
|
+
environment and writes the wheel to `dist/`. It can be invoked from any working
|
|
130
|
+
directory and does not install the package's runtime dependencies.
|
|
131
|
+
|
|
132
|
+
### Repository Map
|
|
133
|
+
|
|
134
|
+
| Path | Purpose |
|
|
135
|
+
|---|---|
|
|
136
|
+
| `model-catalog/` | Supported models, context limits, and recommended job profiles |
|
|
137
|
+
| `docs/` | Getting started material, concepts, guides, and reference |
|
|
138
|
+
| `recipes/` | End-to-end training, sampling, and evaluation workflows |
|
|
139
|
+
| `examples/api/` | Small JSON examples for individual API operations |
|
|
140
|
+
| `examples/config/` | Connection configuration templates |
|
|
141
|
+
| `src/cortex_training/` | Installable Python client |
|
|
142
|
+
| `tests/` | Client and CLI tests |
|
|
143
|
+
| `cluster-status.py` | Optional watch view of running jobs and GPU usage |
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
cortex_training/__init__.py,sha256=W3uHGh4yGYvB3Ye-29NzWcmAuq4wvS03dmR0G6O5aB0,1547
|
|
2
|
+
cortex_training/__main__.py,sha256=bADZ9swQotcDkWUIQ-opq-k_PJZh8yR1IcddxDXio6k,737
|
|
3
|
+
cortex_training/_cli.py,sha256=r38HouCpYj0RR2MFSyzlZ82kaNd93g2LtOrGTc-rfy4,33590
|
|
4
|
+
cortex_training/cli.py,sha256=5ecW-hKpRICUQHTDelL3xcZW1aLz-1UoGBl2DV6q5Ro,2842
|
|
5
|
+
cortex_training/client.py,sha256=EF9pIE1LRuTycKwgL2dFX15uqAH7Tyf5ZU6cARCr55k,115150
|
|
6
|
+
cortex_training/peft.py,sha256=dN0cDbckx6wcLil39hoD6JXrvmPQSN4WC0zFRC5npIQ,4555
|
|
7
|
+
cortex_training/telemetry.py,sha256=n8sEay0umQFkhWPIw44YjcQHZJC23Un3qDuwjtVpmF4,19383
|
|
8
|
+
cortex_training/wire.py,sha256=0jOxKzjdL9ym_b6G5UKlfhuwR6ZRXc0q4mSkYNhC6o8,15768
|
|
9
|
+
cortex_training/tui/__init__.py,sha256=EuG2sQ4PJ4yIZnmSZ2j0SxZ2MHYFHNsBgRT7_N6M3Uo,946
|
|
10
|
+
cortex_training/tui/__main__.py,sha256=_8sL3nktS4bpX-hmDw3wioD-9Ol7pC20MHS3PpRhywY,4365
|
|
11
|
+
cortex_training/tui/app.py,sha256=p7qQxwfgCiMIrmqdMoGR2HbkaQ9r3bYBMGzwy7Cddv8,27040
|
|
12
|
+
cortex_training/tui/format.py,sha256=e83sWcL14ybDq_3oMS7HlAD9vFGREKbN8SiiLx7b_Pg,12835
|
|
13
|
+
cortex_training/tui/log_cache.py,sha256=J6tvefYorOPS-vCDGwLN7EA3Shq_UJtXV_k6S_mNv9g,14977
|
|
14
|
+
cortex_training-0.0.3.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
|
|
15
|
+
cortex_training-0.0.3.dist-info/METADATA,sha256=F7IkUp6ERAW02hSc_tW6i2qRXWSlKVorJMIkiKF1pEc,4930
|
|
16
|
+
cortex_training-0.0.3.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
|
|
17
|
+
cortex_training-0.0.3.dist-info/entry_points.txt,sha256=2_-x1gwFPTJpKb7hbNehYQxUc6h4phAz-NgDBXYRLBQ,91
|
|
18
|
+
cortex_training-0.0.3.dist-info/top_level.txt,sha256=3vIEeuPa44fLTC3IRfhim4oUaPKPhxix7m8fgTHhRUs,16
|
|
19
|
+
cortex_training-0.0.3.dist-info/RECORD,,
|