diff --git a/backends/native/README.md b/backends/native/README.md index 9ce1b69d277..27a35c6930f 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -45,6 +45,15 @@ A `.ptn` is an uncompressed zip with fixed entry names: | `program.ptg` | Methods as graphs; FlatBuffer defined by `serialization/native_graph.fbs`. | | `program.safetensors` | Constants. Byte-identical immutable tensors are stored once. | | `aliases.json` | Deduplicated key to owner key; present only when needed. | +| `metadata.json` | Producer-defined values by name; present only when given. | + +`to_native(ep, metadata={...})` writes `metadata.json`, the counterpart to a +PTE's constant methods: a JSON object whose values are each a bool, number, +string, or a flat list of those. Names mean nothing to the format; a runner +reads the ones it knows with `Package::metadata_int` and `metadata_ints`, and +`exir.native.read_metadata` reads them back in Python. LLM exports reuse the +constant-method names in `extension/llm/export/model_metadata.py`, such as +`get_windows` for the per-layer sliding windows of an off-graph KV cache. ## Samples (optional) diff --git a/backends/native/runtime/deserialize/Package.cpp b/backends/native/runtime/deserialize/Package.cpp index 4b5f33e3cd5..407811eb28f 100644 --- a/backends/native/runtime/deserialize/Package.cpp +++ b/backends/native/runtime/deserialize/Package.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -72,6 +73,42 @@ std::unordered_map parse_aliases( return aliases; } +Json parse_metadata(std::string_view text) { + Json doc; + try { + doc = parse_json(text); + } catch (const Json::exception& error) { + throw std::runtime_error( + "package: invalid metadata.json: " + std::string(error.what())); + } + if (!doc.is_object()) { + throw std::runtime_error("package: metadata.json is not a JSON object"); + } + for (auto entry = doc.begin(); entry != doc.end(); ++entry) { + const Json& value = entry.value(); + const auto scalar = [](const Json& v) { + return v.is_boolean() || v.is_number() || v.is_string(); + }; + if (entry.key().empty() || + !(scalar(value) || + (value.is_array() && + std::all_of(value.begin(), value.end(), scalar)))) { + throw std::runtime_error( + "package: metadata '" + entry.key() + + "' must be a JSON scalar or a list of scalars"); + } + } + return doc; +} + +// nlohmann parses an integer past int64_t's range as unsigned. +bool is_metadata_int(const Json& value) { + return value.is_number_integer() && + (!value.is_number_unsigned() || + value.get() <= + static_cast(std::numeric_limits::max())); +} + } // namespace bool Package::looks_like_package(ByteSpan bytes) { @@ -124,6 +161,7 @@ Package& Package::operator=(Package&& other) noexcept { tensors_ = std::move(other.tensors_); tensor_data_offset_ = other.tensor_data_offset_; aliases_ = std::move(other.aliases_); + metadata_ = std::move(other.metadata_); } return *this; } @@ -179,6 +217,50 @@ void Package::load_metadata() { const std::vector aliases = zip_->read(kAliasesEntry); aliases_ = parse_aliases(ByteSpan(aliases), *tensors_); } + + const std::optional metadata_size = zip_->member_size(kMetadataEntry); + if (metadata_size) { + if (*metadata_size > detail::kMaxJsonBytes) { + throw ResourceLimitError("package: metadata exceeds size limit"); + } + const std::vector metadata = zip_->read(kMetadataEntry); + metadata_.assign(metadata.begin(), metadata.end()); + parse_metadata(metadata_); + } +} + +std::optional Package::metadata_int(const std::string& name) const { + if (metadata_.empty()) { + return std::nullopt; + } + const Json doc = parse_metadata(metadata_); + const auto found = doc.find(name); + if (found == doc.end()) { + return std::nullopt; + } + if (!is_metadata_int(*found)) { + throw std::runtime_error( + "package: metadata '" + name + "' is not an integer"); + } + return found->get(); +} + +std::optional> Package::metadata_ints( + const std::string& name) const { + if (metadata_.empty()) { + return std::nullopt; + } + const Json doc = parse_metadata(metadata_); + const auto found = doc.find(name); + if (found == doc.end()) { + return std::nullopt; + } + if (!found->is_array() || + !std::all_of(found->begin(), found->end(), is_metadata_int)) { + throw std::runtime_error( + "package: metadata '" + name + "' is not a list of integers"); + } + return found->get>(); } std::optional Package::constant_info( @@ -254,8 +336,10 @@ void Package::verify() const { if (tensors_) { zip_->verify(kSafeTensorsEntry); } - if (zip_->member_size(kAliasesEntry)) { - zip_->verify(kAliasesEntry); + for (const char* member : {kAliasesEntry, kMetadataEntry}) { + if (zip_->member_size(member)) { + zip_->verify(member); + } } } diff --git a/backends/native/runtime/deserialize/Package.h b/backends/native/runtime/deserialize/Package.h index a81654216a6..918ada644d0 100644 --- a/backends/native/runtime/deserialize/Package.h +++ b/backends/native/runtime/deserialize/Package.h @@ -27,6 +27,7 @@ namespace ptn { constexpr const char* kProgramEntry = "program.ptg"; constexpr const char* kSafeTensorsEntry = "program.safetensors"; constexpr const char* kAliasesEntry = "aliases.json"; +constexpr const char* kMetadataEntry = "metadata.json"; // Metadata for one constant resolved out of a package. `sizes` is borrowed from // the Package and must not outlive it. @@ -58,6 +59,8 @@ class Package { std::optional tensors_; size_t tensor_data_offset_ = 0; std::unordered_map aliases_; + // The metadata.json object, validated at load; empty when absent. + std::string metadata_; Package(); @@ -101,6 +104,12 @@ class Package { // Metadata for `key`, resolving an alias to its owner. nullopt when absent. std::optional constant_info(const std::string& key) const; + // Producer-defined metadata (the metadata.json member) by name. nullopt when + // absent; throws std::runtime_error when the value has another type. + std::optional metadata_int(const std::string& name) const; + std::optional> metadata_ints( + const std::string& name) const; + // Load one constant into a new owning buffer. nullopt when absent. std::optional acquire_constant(const std::string& key) const; diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp index 0afa787051a..f9d47fae185 100644 --- a/backends/native/test/runtime/deserialize/test_package.cpp +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -159,6 +159,46 @@ TEST(PackageTest, IgnoresSampleMembers) { EXPECT_NO_THROW(package.verify()); } +TEST(PackageTest, ReadsMetadataByName) { + const TempPackage file( + {}, + {{kMetadataEntry, + R"({"get_n_caches":2,"get_windows":[0,4096],"arch":"llama","tied":true})"}}); + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.metadata_int("get_n_caches"), 2); + EXPECT_EQ( + package.metadata_ints("get_windows"), (std::vector{0, 4096})); + EXPECT_EQ(package.metadata_int("missing"), std::nullopt); + EXPECT_EQ(package.metadata_ints("missing"), std::nullopt); + EXPECT_THROW(package.metadata_int("get_windows"), std::runtime_error); + EXPECT_THROW(package.metadata_ints("get_n_caches"), std::runtime_error); + EXPECT_THROW(package.metadata_int("arch"), std::runtime_error); + EXPECT_THROW(package.metadata_int("tied"), std::runtime_error); + EXPECT_NO_THROW(package.verify()); + + const TempPackage none; + EXPECT_EQ( + Package::load(none.path()).metadata_int("get_n_caches"), std::nullopt); +} + +TEST(PackageTest, RejectsMalformedMetadata) { + for (const std::string metadata : + {R"([1])", + R"({"nested":[[1]]})", + R"({"object":{"a":1}})", + R"({"":1})", + R"({"null":null})", + "{"}) { + const TempPackage file({}, {{kMetadataEntry, metadata}}); + EXPECT_THROW(Package::load(file.path()), std::runtime_error) << metadata; + } + const TempPackage huge( + {}, {{kMetadataEntry, R"({"n":18446744073709551615})"}}); + EXPECT_THROW( + Package::load(huge.path()).metadata_int("n"), std::runtime_error); +} + TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) { const TempPackage file; const Package package = Package::load(file.path()); diff --git a/exir/_serialize/_ptn.py b/exir/_serialize/_ptn.py index d06e0942c63..82b9f34cbc2 100644 --- a/exir/_serialize/_ptn.py +++ b/exir/_serialize/_ptn.py @@ -16,6 +16,7 @@ program.ptg native Program flatbuffer (file identifier "NPTG") program.safetensors referenced constants, content-deduped (owners only) aliases.json duplicate data_key -> owner key, only when duplicates exist + metadata.json producer-defined name -> JSON scalar or list, when given Byte-identical *immutable* tensors (same dtype, shape, and content) are stored once: the first key owns the safetensors entry and the rest alias to it. An owner @@ -50,6 +51,7 @@ import hashlib import json +import math import os import time import zipfile @@ -59,6 +61,7 @@ PTG_ENTRY = "program.ptg" SAFETENSORS_ENTRY = "program.safetensors" ALIASES_ENTRY = "aliases.json" +METADATA_ENTRY = "metadata.json" _SAFETENSORS_METADATA_KEY = "__metadata__" # Tensor data starts on this boundary in the file, and each tensor is naturally @@ -276,11 +279,39 @@ def _write_safetensors( stream.write(_raw(owners[key])) +MetadataValue = bool | int | float | str | list[bool | int | float | str] + + +def _check_metadata(metadata: dict[str, MetadataValue], where: str) -> None: + """Names are non-empty strings; values are JSON scalars or flat lists of them. + + Floats must be finite: JSON has no infinity or NaN. + """ + if not isinstance(metadata, dict): + raise TypeError(f"{where}: metadata must be a dict.") + + def check_scalar(name: str, value: object) -> None: + if not isinstance(value, (bool, int, float, str)): + raise TypeError( + f"{where}: metadata {name!r} holds a {type(value).__name__}; " + "values are bool, int, float, str, or a list of those." + ) + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{where}: metadata {name!r} holds a non-finite float.") + + for name, value in metadata.items(): + if not isinstance(name, str) or not name: + raise TypeError(f"{where}: metadata names must be non-empty strings.") + for item in value if isinstance(value, list) else [value]: + check_scalar(name, item) + + def write_ptn( path: str, ptg: bytes, constants: dict[str, torch.Tensor], mutable_keys: frozenset[str] = frozenset(), + metadata: dict[str, MetadataValue] | None = None, ) -> None: """Write a .ptn package, streaming the constants a tensor at a time. @@ -293,15 +324,21 @@ def write_ptn( dedup, since equal bytes do not make two mutable buffers the same state. Passed in rather than derived here, so this stays independent of the graph format. Every mutable key must occur in ``constants``. + metadata: Producer-defined values a runtime reads by name, such as an + LLM's cache geometry; written as metadata.json when non-empty. Raises: KeyError: If a constant has a dtype safetensors cannot represent. - TypeError: If a data key is not a string or a value is not a tensor. - ValueError: If mutable-key policy is inconsistent or a data key is reserved. + TypeError: If a data key is not a string, a value is not a tensor, or + metadata is not a name -> scalar-or-list dict. + ValueError: If mutable-key policy is inconsistent, a data key is reserved, + or a metadata float is not finite. """ # TODO: save atomically. This opens the destination directly, so a failure # part way through truncates an existing valid package. Write to a temporary # file beside the destination and os.replace it once complete. + metadata = metadata or {} + _check_metadata(metadata, "write_ptn") owners, aliases = _dedup(constants, mutable_keys) parent = os.path.dirname(path) if parent: @@ -317,6 +354,26 @@ def write_ptn( pkg.writestr( ALIASES_ENTRY, json.dumps(aliases, indent=2, sort_keys=True) ) + if metadata: + pkg.writestr( + METADATA_ENTRY, + json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False), + ) + + +def read_metadata(path: str) -> dict[str, MetadataValue]: + """The package's metadata.json, or an empty dict when it has none. + + Raises: + TypeError: If metadata.json is not a name -> scalar-or-list object. + ValueError: If metadata.json is not valid JSON. + """ + with zipfile.ZipFile(path) as pkg: + if METADATA_ENTRY not in pkg.namelist(): + return {} + metadata = json.loads(pkg.read(METADATA_ENTRY)) + _check_metadata(metadata, "read_metadata") + return metadata def read_ptn(path: str) -> tuple[bytes, dict[str, torch.Tensor]]: @@ -371,4 +428,12 @@ def read_ptn(path: str) -> tuple[bytes, dict[str, torch.Tensor]]: return ptg, constants -__all__ = ["ALIASES_ENTRY", "PTG_ENTRY", "SAFETENSORS_ENTRY", "read_ptn", "write_ptn"] +__all__ = [ + "ALIASES_ENTRY", + "METADATA_ENTRY", + "PTG_ENTRY", + "SAFETENSORS_ENTRY", + "read_metadata", + "read_ptn", + "write_ptn", +] diff --git a/exir/native/__init__.py b/exir/native/__init__.py index f7e3b0dbacd..23fe5e06465 100644 --- a/exir/native/__init__.py +++ b/exir/native/__init__.py @@ -13,6 +13,7 @@ SampleSet, write_samples, ) +from executorch.exir._serialize._ptn import read_metadata from executorch.exir.native.native import NativeProgramManager, to_native __all__ = [ @@ -20,6 +21,7 @@ "NativeProgramManager", "Sample", "SampleSet", + "read_metadata", "read_samples", "to_native", "write_samples", diff --git a/exir/native/native.py b/exir/native/native.py index 4df444e5340..2c484c827ee 100644 --- a/exir/native/native.py +++ b/exir/native/native.py @@ -47,7 +47,7 @@ validate_program, ) from executorch.backends.native.serialization.schema import OutputKind, Program -from executorch.exir._serialize._ptn import write_ptn +from executorch.exir._serialize._ptn import _check_metadata, MetadataValue, write_ptn from torch.utils import _pytree as pytree if TYPE_CHECKING: @@ -76,10 +76,12 @@ def __init__( self, ptg: bytes, constants: dict[str, torch.Tensor], + metadata: dict[str, MetadataValue] | None = None, ) -> None: # TODO add args help self._ptg = ptg self._constants = constants + self._metadata = dict(metadata or {}) @property def methods(self) -> set[str]: @@ -89,7 +91,13 @@ def methods(self) -> set[str]: def save(self, path: str) -> None: """Write the program and its constants to ``path`` as a .ptn package.""" mutable_keys = frozenset(_mutable_data_keys(deserialize_program(self._ptg))) - write_ptn(path, self._ptg, self._constants, mutable_keys=mutable_keys) + write_ptn( + path, + self._ptg, + self._constants, + mutable_keys=mutable_keys, + metadata=self._metadata, + ) def _check_delegate_outputs( @@ -362,6 +370,7 @@ def _normalize_programs( def to_native( programs: "ExportedProgram | dict[str, ExportedProgram]", + metadata: dict[str, MetadataValue] | None = None, ) -> NativeProgramManager: """Lower ExportedProgram(s) through the native backend. @@ -376,17 +385,21 @@ def to_native( Nothing about lowering is configurable yet. The passes, the partitioner, and the edge compile config are all owned here, because correctness of the package depends on them. Exposing the compile config is worth reconsidering. - ``constant_methods`` and ETRecord are likewise deferred. + ETRecord is likewise deferred. Args: programs: The exported program(s) to lower. + metadata: Values a runtime reads by name, the package's counterpart to a + PTE's ``constant_methods``: each a bool, int, float, str, or a list + of those. Saved as the package's metadata.json. Returns: A ``NativeProgramManager``; call ``save`` to write a .ptn package. Raises: TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name - dictionary containing only ``ExportedProgram`` values. + dictionary containing only ``ExportedProgram`` values, or + ``metadata`` holds a name or value outside the types above. ValueError: If the method dictionary is empty, a method name is empty or in the reserved ``state`` namespace, a method does not fully delegate to the native backend, a mutable constant differs between methods, @@ -394,6 +407,8 @@ def to_native( backing tensor. """ method_programs = _normalize_programs(programs) + if metadata is not None: + _check_metadata(metadata, "to_native") from executorch.backends.native import get_default_compile_config from executorch.backends.native.partitioner import NativePartitioner @@ -493,7 +508,9 @@ def to_native( constants = _merge_constants(method_constants) _validate_merged_program(program, constants) - return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants) + return NativeProgramManager( + ptg=_compile_to_bytes(program), constants=constants, metadata=metadata + ) __all__ = ["NativeProgramManager", "to_native"] diff --git a/exir/native/test/test_native.py b/exir/native/test/test_native.py index 135e92d34d8..704e3b92df9 100644 --- a/exir/native/test/test_native.py +++ b/exir/native/test/test_native.py @@ -29,6 +29,7 @@ from executorch.exir import to_edge_transform_and_lower from executorch.exir._serialize._ptn import ( ALIASES_ENTRY, + METADATA_ENTRY, PTG_ENTRY, read_ptn, SAFETENSORS_ENTRY, @@ -39,6 +40,7 @@ ) from executorch.exir.native import ( Call, + read_metadata, read_samples, Sample, SampleSet, @@ -744,5 +746,45 @@ def test_samples_round_trip_on_saved_package(self): ) +class MetadataTest(unittest.TestCase): + def _save(self, td: str, metadata) -> str: + path = os.path.join(td, "m.ptn") + model = nn.Linear(8, 4) + program = torch.export.export(model, (torch.randn(2, 8),)) + to_native(program, metadata=metadata).save(path) + return path + + def test_metadata_round_trips(self): + metadata = { + "get_n_caches": 2, + "get_windows": [0, 4096], + "rope_theta": 5e5, + "tied": True, + "arch": "llama", + } + with tempfile.TemporaryDirectory() as td: + self.assertEqual(read_metadata(self._save(td, metadata)), metadata) + + def test_no_metadata_writes_no_member(self): + with tempfile.TemporaryDirectory() as td: + path = self._save(td, None) + with zipfile.ZipFile(path) as pkg: + self.assertNotIn(METADATA_ENTRY, pkg.namelist()) + self.assertEqual(read_metadata(path), {}) + + def test_invalid_metadata_is_rejected(self): + program = torch.export.export(nn.Linear(8, 4), (torch.randn(2, 8),)) + for metadata in ( + {"": 1}, + {"nested": [[1]]}, + {"tensor": torch.tensor([1])}, + {"none": None}, + ): + with self.subTest(metadata=metadata), self.assertRaises(TypeError): + to_native(program, metadata=metadata) + with self.assertRaisesRegex(ValueError, "non-finite"): + to_native(program, metadata={"nan": float("nan")}) + + if __name__ == "__main__": unittest.main()