|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2026 Flare community contributors |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +"""Locate a checkpoint's assets when a TrellisMX release overlays a carrier. |
| 5 | +
|
| 6 | +A TrellisMX release is a routed-expert overlay: its root holds the manifest, the |
| 7 | +per-layer sidecars and the design record, while the safetensors index, tokenizer, |
| 8 | +chat template and generation config stay in ``carrier/`` with every tensor the |
| 9 | +sidecars do not replace. Shared code that resolves or inspects a checkpoint asks |
| 10 | +this module where those files are, without importing a recipe. |
| 11 | +
|
| 12 | +Only the manifest is read, and no payload. This module does not validate a release; |
| 13 | +the recipe that serves one owns every gate. A root without this manifest is treated |
| 14 | +exactly as before so no other checkpoint's behaviour moves, and a manifest that |
| 15 | +cannot be read raises instead of letting its root pass for a plain checkpoint. |
| 16 | +""" |
| 17 | + |
| 18 | +from __future__ import annotations |
| 19 | + |
| 20 | +import json |
| 21 | +from pathlib import Path |
| 22 | + |
| 23 | +MANIFEST_NAME = "trellismx-manifest.json" |
| 24 | +MANIFEST_SCHEMA = "trellismx.hf-overlay-release.v1" |
| 25 | +CARRIER_DIRECTORY = "carrier" |
| 26 | +SIDECAR_DIRECTORY = "sidecars" |
| 27 | +# The encoder transform the sidecars pin, which preparation reads beside them. |
| 28 | +TRANSFORM_PATH = "design/transform.json" |
| 29 | + |
| 30 | + |
| 31 | +def sidecar_path(layer: int, rank: int, world_size: int) -> str: |
| 32 | + """The release-relative path of one layer and tensor-parallel rank's sidecar.""" |
| 33 | + return f"{SIDECAR_DIRECTORY}/p8-layer-{layer:03d}-tp{world_size}-rank-{rank}.safetensors" |
| 34 | + |
| 35 | + |
| 36 | +def _manifest(root: Path) -> dict | None: |
| 37 | + """``root``'s overlay manifest when it declares this schema, else None. |
| 38 | +
|
| 39 | + Raises ``ValueError`` when the manifest exists but is not a readable JSON object. |
| 40 | + """ |
| 41 | + manifest_path = root / MANIFEST_NAME |
| 42 | + if not manifest_path.is_file(): |
| 43 | + return None |
| 44 | + try: |
| 45 | + manifest = json.loads(manifest_path.read_text()) |
| 46 | + except (OSError, ValueError) as error: |
| 47 | + raise ValueError(f"{manifest_path} is unreadable: {error}") from error |
| 48 | + if not isinstance(manifest, dict): |
| 49 | + raise ValueError(f"{manifest_path} is not a JSON object") |
| 50 | + return manifest if manifest.get("schema") == MANIFEST_SCHEMA else None |
| 51 | + |
| 52 | + |
| 53 | +def _carrier_index(root: Path) -> Path: |
| 54 | + return root / CARRIER_DIRECTORY / "model.safetensors.index.json" |
| 55 | + |
| 56 | + |
| 57 | +def trellismx_assets_root(root: str | Path) -> Path: |
| 58 | + """The directory holding ``root``'s index and frontend assets. |
| 59 | +
|
| 60 | + The carrier directory when ``root`` declares this overlay schema and carries a |
| 61 | + carrier index, and ``root`` itself otherwise. |
| 62 | + """ |
| 63 | + root = Path(root).expanduser() |
| 64 | + return root / CARRIER_DIRECTORY if _manifest(root) is not None and _carrier_index(root).is_file() else root |
| 65 | + |
| 66 | + |
| 67 | +def overlay_schema(root: str | Path) -> str | None: |
| 68 | + """The overlay manifest schema ``root`` is packaged with, or None for any other checkpoint.""" |
| 69 | + return MANIFEST_SCHEMA if _manifest(Path(root).expanduser()) is not None else None |
| 70 | + |
| 71 | + |
| 72 | +def _allocation(manifest: dict) -> dict[int, int]: |
| 73 | + """The manifest's allocation: each layer index and the tensor-parallel size it is sliced for.""" |
| 74 | + allocation = manifest.get("allocation") |
| 75 | + if not isinstance(allocation, dict) or not allocation: |
| 76 | + raise ValueError("TrellisMX manifest allocation must be a nonempty object of layer indices") |
| 77 | + layers = {} |
| 78 | + for key, world_size in allocation.items(): |
| 79 | + try: |
| 80 | + layer = int(key) |
| 81 | + except ValueError: |
| 82 | + layer = None |
| 83 | + if layer is None or key != str(layer) or layer < 0 or type(world_size) is not int or world_size <= 0: |
| 84 | + raise ValueError(f"TrellisMX manifest allocation entry {key!r}: {world_size!r} is invalid") |
| 85 | + layers[layer] = world_size |
| 86 | + return layers |
| 87 | + |
| 88 | + |
| 89 | +def _manifest_files(manifest: dict) -> list[tuple[str, int]]: |
| 90 | + """Every file the manifest lists, which must be exactly the release's canonical inventory. |
| 91 | +
|
| 92 | + That inventory is one rank-named sidecar for every layer and rank the allocation |
| 93 | + declares, so a list that omits, repeats, renames or adds an entry is rejected |
| 94 | + instead of making an incomplete download look complete. |
| 95 | + """ |
| 96 | + files = manifest.get("files") |
| 97 | + if not isinstance(files, list) or not files: |
| 98 | + raise ValueError("TrellisMX manifest lists no files") |
| 99 | + entries = [] |
| 100 | + for entry in files: |
| 101 | + path = entry.get("path") if isinstance(entry, dict) else None |
| 102 | + size = entry.get("bytes") if isinstance(entry, dict) else None |
| 103 | + if not isinstance(path, str) or not path or Path(path).is_absolute() or ".." in Path(path).parts \ |
| 104 | + or type(size) is not int or size <= 0: |
| 105 | + raise ValueError(f"TrellisMX manifest file entry {entry!r} is invalid") |
| 106 | + entries.append((path, size)) |
| 107 | + inventory = {sidecar_path(layer, rank, world_size) |
| 108 | + for layer, world_size in _allocation(manifest).items() for rank in range(world_size)} |
| 109 | + listed = [path for path, _ in entries] |
| 110 | + if len(listed) != len(set(listed)) or set(listed) != inventory: |
| 111 | + missing, unexpected = sorted(inventory - set(listed)), sorted(set(listed) - inventory) |
| 112 | + repeated = sorted({path for path in listed if listed.count(path) > 1}) |
| 113 | + raise ValueError("TrellisMX manifest files must be one sidecar per allocated layer and rank:" |
| 114 | + f" missing {missing}, unexpected {unexpected}, repeated {repeated}") |
| 115 | + return entries |
| 116 | + |
| 117 | + |
| 118 | +def trellismx_files_present(root: str | Path) -> bool: |
| 119 | + """Whether every file ``root``'s release needs beside its carrier is present at its declared size. |
| 120 | +
|
| 121 | + Those files are the canonical inventory (one sidecar per allocated layer and rank) |
| 122 | + and the encoder transform. True for a root that is not a release, which needs |
| 123 | + nothing more; a manifest whose file list is not that inventory raises |
| 124 | + ``ValueError``. Only sizes are compared; preparation verifies each sidecar's |
| 125 | + digest when it loads it. |
| 126 | + """ |
| 127 | + root = Path(root).expanduser() |
| 128 | + manifest = _manifest(root) |
| 129 | + if manifest is None: |
| 130 | + return True |
| 131 | + return _carrier_index(root).is_file() and (root / TRANSFORM_PATH).is_file() and all( |
| 132 | + (root / path).is_file() and (root / path).stat().st_size == size for path, size in _manifest_files(manifest)) |
| 133 | + |
| 134 | + |
| 135 | +def trellismx_stored_bytes(root: str | Path) -> int | None: |
| 136 | + """Stored tensor bytes of ``root``'s release: the carrier's payload plus every manifest file. |
| 137 | +
|
| 138 | + None for a root that is not a release. |
| 139 | + """ |
| 140 | + root = Path(root).expanduser() |
| 141 | + manifest = _manifest(root) |
| 142 | + if manifest is None: |
| 143 | + return None |
| 144 | + try: |
| 145 | + index = json.loads(_carrier_index(root).read_text()) |
| 146 | + except OSError as error: |
| 147 | + raise ValueError(f"{_carrier_index(root)} is unreadable: {error}") from error |
| 148 | + carrier = index.get("metadata", {}).get("total_size") if isinstance(index, dict) else None |
| 149 | + if type(carrier) is not int or carrier <= 0: |
| 150 | + raise ValueError(f"{root / CARRIER_DIRECTORY} has an invalid safetensors total_size: {carrier!r}") |
| 151 | + return carrier + sum(size for _, size in _manifest_files(manifest)) |
| 152 | + |
| 153 | + |
| 154 | +__all__ = [ |
| 155 | + "CARRIER_DIRECTORY", "MANIFEST_NAME", "MANIFEST_SCHEMA", "SIDECAR_DIRECTORY", "TRANSFORM_PATH", "overlay_schema", |
| 156 | + "sidecar_path", "trellismx_assets_root", "trellismx_files_present", "trellismx_stored_bytes", |
| 157 | +] |
0 commit comments