From a40b531ab7fb8a7ec37104c275396082a59fd764 Mon Sep 17 00:00:00 2001 From: Esteban Safranchik Date: Wed, 5 Aug 2026 17:32:36 -0700 Subject: [PATCH 1/2] Add directory-native model state --- AGENTS.md | 10 ++ CHANGELOG.md | 7 + GLOSSARY.md | 4 +- README.md | 37 ++++- SPEC.md | 34 ++++ hytorch/__init__.py | 4 + hytorch/_git.py | 11 ++ hytorch/graph.py | 18 +++ hytorch/state_dir.py | 334 ++++++++++++++++++++++++++++++++++++++++ tests/test_linear.py | 1 + tests/test_state_dir.py | 221 ++++++++++++++++++++++++++ 11 files changed, 678 insertions(+), 3 deletions(-) create mode 100644 hytorch/state_dir.py create mode 100644 tests/test_state_dir.py diff --git a/AGENTS.md b/AGENTS.md index 937d3e9..80e3643 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -39,6 +39,10 @@ real text-agent equivalent exists. | `torch.optim` | `hytorch.optim` | | learning rate | mutation temperature `temp` | | applied parameter update | committed Git mutation | +| `model.state_dict()` | `model.state_dir()` | +| `torch.save(...)` | `hytorch.save(...)` | +| `torch.load(...)` | `hytorch.load(...)` | +| `model.load_state_dict(...)` | `model.load_state_dir(...)` | `Module.__setattr__` registers Parameters and child Modules. Registration owns state; calls in `forward()` dynamically define topology. `model.parameters()` @@ -75,6 +79,12 @@ candidate model branch. `mn.init.DEFAULT_PRIORS`. It can later contain arbitrary code, tools, examples, and data. +Model checkpoint syntax follows PyTorch with a directory-native representation: +`hytorch.save(model.state_dir(), path)` and +`model.load_state_dir(hytorch.load(path))`. A StateDir fixes one canonical model +commit and preserves the complete model Git history. It excludes feedback, +sessions, temporary node trees, and unpromoted optimizer candidates. + Forward returns the complete committed statespace, never a special answer file. Do not inject Space contents into the agent prompt. Mount complete directory trees. `zero_feed()` clears accumulated feedback and discards an unpromoted diff --git a/CHANGELOG.md b/CHANGELOG.md index 7464549..3ef896e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ All notable changes to HyTorch will be recorded in this file. The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). HyTorch uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- PyTorch-shaped, directory-native model checkpoints with `state_dir()`, + `hytorch.save()`, `hytorch.load()`, and `load_state_dir()`. + ## [0.1.0] - 2026-08-05 ### Added diff --git a/GLOSSARY.md b/GLOSSARY.md index ba47987..5c945e1 100644 --- a/GLOSSARY.md +++ b/GLOSSARY.md @@ -34,7 +34,9 @@ | parameter delta | candidate workspace mutation | Proposed change made during backward | | updated parameter storage | promoted global Git commit | Canonical model generation after `step()` | | saved forward activations | statespace commit and harness session | Context resumed during backward | -| `state_dict()` | model workspace Git repository | Serialization target; API not yet implemented | +| `state_dict()` | `state_dir()` | Immutable handle to the canonical model-state revision | +| `torch.save(model.state_dict(), path)` | `hytorch.save(model.state_dir(), path)` | Save complete parameter state and canonical history | +| `load_state_dict(torch.load(path))` | `load_state_dir(hytorch.load(path))` | Validate and restore registered workspaces | The canonical training form is: diff --git a/README.md b/README.md index b6e36f9..779a7d6 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,39 @@ previous iteration. `backward()` resumes the agents and creates candidate workspace changes. `step()` promotes all completed changes as one new model generation. +### Save and load model state + +Save the complete canonical model state as a self-contained Git directory: + +```python +hytorch.save(model.state_dir(), "research-network") +``` + +Create the same model structure, then load the saved workspaces: + +```python +restored = ResearchNetwork().to("pi") +result = restored.load_state_dir(hytorch.load("research-network")) +``` + +This follows the PyTorch `state_dict()` pattern with a directory-native state: + +```python +# PyTorch +torch.save(model.state_dict(), "model.pth") +model.load_state_dict(torch.load("model.pth", weights_only=True)) + +# HyTorch +hytorch.save(model.state_dir(), "model-state") +model.load_state_dir(hytorch.load("model-state")) +``` + +A `StateDir` identifies one immutable model commit. The saved directory +contains `MODEL.json`, every registered workspace, and the canonical model Git +history. Loading is strict by default. Pass `strict=False` to permit missing or +unexpected workspace keys with compatible shapes. The save destination must +not already exist. + ## PyTorch-shaped composition HyTorch follows PyTorch syntax and ownership where a direct agent equivalent @@ -215,8 +248,8 @@ environments and review agent-created changes before production use. Version 0.1.0 includes Spaces, Parameters, dynamic Module graphs, dense Linear layers, directional backward feedback, atomic DFM optimizer generations, and -the Dockerized Pi harness. It does not yet implement `state_dict()`. The -`codex` and `claude-code` harnesses are reserved but unavailable. +the Dockerized Pi harness. The `codex` and `claude-code` harnesses are reserved +but unavailable. ## Resources diff --git a/SPEC.md b/SPEC.md index ea40f94..e5e0758 100644 --- a/SPEC.md +++ b/SPEC.md @@ -249,6 +249,37 @@ Feedback is transient text and does not use a third Git repository. Git gives Spaces and model generations stable identity, ancestry, diffs, audit history, and atomic promotion. +## Model state directories + +HyTorch serializes model state as a directory because each Parameter element is +already a complete directory. The public form follows PyTorch checkpoint +syntax: + +```python +hytorch.save(model.state_dir(), path) +model.load_state_dir(hytorch.load(path)) +``` + +`model.state_dir()` returns a `StateDir` fixed to the canonical model commit at +the time of the call. It does not include an unpromoted DFM candidate. +`hytorch.save()` creates a self-contained Git directory at that exact commit. +The saved state contains `MODEL.json`, all registered workspace directories, +and the canonical model history. It excludes feedback, active harness sessions, +temporary node trees, and optimizer candidates. + +`hytorch.load()` validates the repository root, committed `MODEL.json`, format, +and workspace paths. It returns a `StateDir`; it does not modify a model. +`model.load_state_dir()` copies matching workspaces into an initialized model +and records one canonical load commit. The default `strict=True` requires the +saved and destination workspace keys to match exactly. `strict=False` permits +missing and unexpected keys, but shape and module-type mismatches remain +errors. The return value reports missing and unexpected keys in the same style +as PyTorch's `load_state_dict()`. + +State capture and load require a clean canonical model worktree. Loading while +an optimizer candidate is pending is an error. Validation must finish before +HyTorch changes any destination workspace. + ## Required invariants 1. One output feature executes one agent. @@ -264,3 +295,6 @@ and atomic promotion. 11. Backward closes each resumed forward session. 12. Only `optimizer.step()` promotes candidate workspace commits. 13. `zero_feed()` never changes committed canonical workspace history. +14. A StateDir identifies one committed, immutable model generation. +15. Saved model state never includes transient feedback, sessions, or candidates. +16. A failed state load leaves every destination workspace unchanged. diff --git a/hytorch/__init__.py b/hytorch/__init__.py index 3dd9795..0bf4831 100644 --- a/hytorch/__init__.py +++ b/hytorch/__init__.py @@ -6,6 +6,7 @@ from ._random import manual_seed from .backward import Loss, Report, WorkspaceRevision from .space import Space, SpaceBatch, space +from .state_dir import StateDir, load, save __version__ = "0.1.0" @@ -17,13 +18,16 @@ "Report", "Space", "SpaceBatch", + "StateDir", "WorkspaceRevision", "__version__", "harness", "inference_mode", "is_inference_mode_enabled", + "load", "manual_seed", "mn", "optim", + "save", "space", ] diff --git a/hytorch/_git.py b/hytorch/_git.py index 50ed595..a30abd2 100644 --- a/hytorch/_git.py +++ b/hytorch/_git.py @@ -93,6 +93,17 @@ def resolve(self, ref: str) -> str: except GitError as exc: raise GitError(f"resolve git ref {ref!r}: {exc}") from exc + def read_file(self, commit: str, path: str) -> bytes: + """Read one file from a committed tree.""" + result = subprocess.run( + ["git", "-C", self.root, "show", f"{commit}:{path}"], + capture_output=True, + check=False, + ) + if result.returncode != 0: + raise GitError(result.stderr.decode(errors="replace").strip()) + return result.stdout + def current_branch(self) -> str: branch = _run(self.root, ["symbolic-ref", "--quiet", "--short", "HEAD"]) if not branch: diff --git a/hytorch/graph.py b/hytorch/graph.py index 908e3f3..f3ae51d 100644 --- a/hytorch/graph.py +++ b/hytorch/graph.py @@ -110,6 +110,23 @@ def parameters(self, recurse: bool = True): self._ensure_parameter_store() return (parameter for _, parameter in self.named_parameters(recurse=recurse)) + def state_dir(self): + """Return an immutable handle to the canonical model-state revision.""" + from .state_dir import StateDir + + store = self._ensure_parameter_store() + if not store.repo.is_clean(): + raise RuntimeError( + "hytorch.mn.Module.state_dir: model state has uncommitted changes" + ) + return StateDir(store.root, store.repo.resolve("HEAD")) + + def load_state_dir(self, state_dir, strict: bool = True): + """Copy a StateDir into this Module and its descendants.""" + from .state_dir import load_module_state + + return load_module_state(self, state_dir, strict) + def apply(self, fn): for child in self.children(): child.apply(fn) @@ -173,6 +190,7 @@ def _ensure_parameter_store(self) -> ParameterStore: "parameters": { parameter_name: { "shape": list(parameter.shape), + "input_features": parameter.input_features, "workspaces": [ view.relative_path for view in parameter.views() ], diff --git a/hytorch/state_dir.py b/hytorch/state_dir.py new file mode 100644 index 0000000..fd1ba6c --- /dev/null +++ b/hytorch/state_dir.py @@ -0,0 +1,334 @@ +"""Directory-native model state serialization.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import tempfile +import uuid +from collections import namedtuple +from dataclasses import dataclass, field + +from ._git import GitError, Repo +from .parameter import copy_tree + + +class _IncompatibleKeys( + namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]) +): + __slots__ = () + + def __repr__(self) -> str: + if not self.missing_keys and not self.unexpected_keys: + return "" + return super().__repr__() + + +@dataclass(frozen=True, init=False) +class StateDir: + """An immutable model-state revision stored in a Git directory.""" + + repo: Repo = field(repr=False, compare=False) + commit: str + dir: str + _manifest: dict = field(repr=False, compare=False) + + def __init__(self, directory: str | os.PathLike[str], commit: str = "HEAD"): + resolved_dir = os.path.realpath(os.fspath(directory)) + try: + repo = Repo.discover(resolved_dir) + except GitError as exc: + raise ValueError( + "hytorch.StateDir: directory must be a Git repository" + ) from exc + if repo.root != resolved_dir: + raise ValueError( + "hytorch.StateDir: directory must be a Git repository root" + ) + resolved_commit = repo.resolve(commit) + object.__setattr__(self, "repo", repo) + object.__setattr__(self, "commit", resolved_commit) + object.__setattr__(self, "dir", resolved_dir) + object.__setattr__(self, "_manifest", _read_manifest(repo, resolved_commit)) + + +def save(state_dir: StateDir, path: str | os.PathLike[str]) -> None: + """Save one exact model-state revision as a self-contained Git directory.""" + if not isinstance(state_dir, StateDir): + raise TypeError("hytorch.save: value must be a hytorch.StateDir") + destination = os.path.abspath(os.fspath(path)) + if os.path.lexists(destination): + raise FileExistsError(destination) + if os.path.commonpath((state_dir.dir, destination)) == state_dir.dir: + raise ValueError("hytorch.save: destination cannot be inside the model state") + + os.makedirs(os.path.dirname(destination), exist_ok=True) + try: + _git( + state_dir.dir, + "clone", + "--quiet", + "--no-local", + "--no-checkout", + "--no-tags", + state_dir.dir, + destination, + ) + _git(destination, "checkout", "--quiet", "-B", "main", state_dir.commit) + _git(destination, "remote", "remove", "origin") + StateDir(destination) + except Exception: + if os.path.lexists(destination): + shutil.rmtree(destination) + raise + + +def load(path: str | os.PathLike[str]) -> StateDir: + """Load a model-state directory for use with ``Module.load_state_dir``.""" + return StateDir(path) + + +def load_module_state(module, state_dir: StateDir, strict: bool) -> _IncompatibleKeys: + if not isinstance(state_dir, StateDir): + raise TypeError( + "hytorch.mn.Module.load_state_dir: state_dir must be a hytorch.StateDir" + ) + if not isinstance(strict, bool): + raise TypeError("hytorch.mn.Module.load_state_dir: strict must be a bool") + + store = module._ensure_parameter_store() + if not store.repo.is_clean(): + raise RuntimeError( + "hytorch.mn.Module.load_state_dir: model state has uncommitted changes" + ) + optimizers = { + parameter._optimizer + for parameter in module.parameters() + if parameter._optimizer is not None + } + if any( + getattr(optimizer, "_pending", None) is not None for optimizer in optimizers + ): + raise RuntimeError( + "hytorch.mn.Module.load_state_dir: discard or promote the pending optimizer update first" + ) + + destination_manifest = _read_manifest(store.repo, store.repo.resolve("HEAD")) + source_entries = _workspace_entries(state_dir._manifest) + destination_entries = _workspace_entries(destination_manifest) + source_keys = set(source_entries) + destination_keys = set(destination_entries) + missing = sorted(destination_keys - source_keys) + unexpected = sorted(source_keys - destination_keys) + errors = _compatibility_errors( + state_dir._manifest, + destination_manifest, + source_entries, + destination_entries, + ) + if strict: + if missing: + errors.append("Missing key(s): " + ", ".join(repr(key) for key in missing)) + if unexpected: + errors.append( + "Unexpected key(s): " + ", ".join(repr(key) for key in unexpected) + ) + if errors: + detail = "\n\t".join(errors) + raise RuntimeError( + f"Error(s) in loading state_dir for {type(module).__name__}:\n\t{detail}" + ) + + with tempfile.TemporaryDirectory(prefix="hytorch-state-load-") as snapshot: + state_dir.repo.export_tree(state_dir.commit, snapshot) + matched = sorted(source_keys & destination_keys) + for key in matched: + source = os.path.join(snapshot, source_entries[key]["path"]) + if not os.path.isdir(source): + raise RuntimeError( + f"hytorch.load: state directory is missing workspace {key!r}" + ) + _promote_loaded_workspaces( + store, + snapshot, + matched, + source_entries, + destination_entries, + ) + + for parameter in module.parameters(): + parameter.zero_feed() + return _IncompatibleKeys(missing, unexpected) + + +def _promote_loaded_workspaces( + store, + snapshot: str, + matched: list[str], + source_entries: dict[str, dict], + destination_entries: dict[str, dict], +) -> None: + base = store.repo.resolve("HEAD") + branch = f"hytorch/load/{uuid.uuid4().hex}" + candidate = tempfile.mkdtemp(prefix="hytorch-state-candidate-") + branched = False + added = False + try: + store.repo.branch(branch, base) + branched = True + store.repo.add_worktree(candidate, branch) + added = True + for key in matched: + source = os.path.join(snapshot, source_entries[key]["path"]) + destination = os.path.join(candidate, destination_entries[key]["path"]) + copy_tree(source, destination) + commit, _ = store.repo.commit_all_workdir( + candidate, "hytorch: load model state" + ) + store.repo.fast_forward(store.root, commit) + except Exception: + _discard_load_candidate(store.repo, candidate, branch, branched, added) + raise + _discard_load_candidate(store.repo, candidate, branch, branched, added) + + +def _discard_load_candidate( + repo: Repo, candidate: str, branch: str, branched: bool, added: bool +) -> None: + if added: + try: + repo.remove_worktree(candidate) + except GitError: + return + elif os.path.isdir(candidate): + shutil.rmtree(candidate) + if branched: + try: + repo.delete_branch(branch) + except GitError: + pass + + +def _read_manifest(repo: Repo, commit: str) -> dict: + try: + value = json.loads(repo.read_file(commit, "MODEL.json")) + except (GitError, json.JSONDecodeError, UnicodeDecodeError) as exc: + raise ValueError( + "hytorch.StateDir: committed MODEL.json is missing or invalid" + ) from exc + if not isinstance(value, dict) or value.get("format") != "hytorch-model-v1": + raise ValueError("hytorch.StateDir: unsupported model-state format") + if not isinstance(value.get("modules"), dict): + raise ValueError("hytorch.StateDir: MODEL.json modules must be an object") + _workspace_entries(value) + return value + + +def _workspace_entries(manifest: dict) -> dict[str, dict]: + entries: dict[str, dict] = {} + for module_name, module in manifest["modules"].items(): + if not isinstance(module_name, str) or not isinstance(module, dict): + raise ValueError("hytorch.StateDir: invalid module entry") + parameters = module.get("parameters") + if not isinstance(parameters, dict): + raise ValueError("hytorch.StateDir: invalid parameter entries") + for parameter_name, parameter in parameters.items(): + if not isinstance(parameter_name, str) or not isinstance(parameter, dict): + raise ValueError("hytorch.StateDir: invalid parameter entry") + shape = parameter.get("shape") + input_features = parameter.get("input_features") + workspaces = parameter.get("workspaces") + if ( + not isinstance(shape, list) + or not shape + or any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in shape + ) + or ( + input_features is not None + and ( + not isinstance(input_features, int) + or isinstance(input_features, bool) + or input_features <= 0 + ) + ) + or not isinstance(workspaces, list) + or len(workspaces) != shape[0] + ): + raise ValueError("hytorch.StateDir: invalid workspace metadata") + for index, relative in enumerate(workspaces): + if not isinstance(relative, str) or not _safe_relative_path(relative): + raise ValueError("hytorch.StateDir: invalid workspace path") + key = f"{module_name}.{parameter_name}.{index}" + if key in entries or any( + item["path"] == relative for item in entries.values() + ): + raise ValueError("hytorch.StateDir: duplicate workspace entry") + entries[key] = { + "path": relative, + "shape": shape, + "input_features": input_features, + "module_type": module.get("type"), + } + return entries + + +def _compatibility_errors( + source_manifest: dict, + destination_manifest: dict, + source_entries: dict[str, dict], + destination_entries: dict[str, dict], +) -> list[str]: + del source_manifest, destination_manifest + errors = [] + for key in sorted(set(source_entries) & set(destination_entries)): + source = source_entries[key] + destination = destination_entries[key] + if source["shape"] != destination["shape"]: + errors.append( + f"size mismatch for {key}: source shape {source['shape']} " + f"does not match model shape {destination['shape']}" + ) + if ( + source["input_features"] is not None + and destination["input_features"] is not None + and source["input_features"] != destination["input_features"] + ): + errors.append( + f"input feature mismatch for {key}: source " + f"{source['input_features']} does not match model " + f"{destination['input_features']}" + ) + if source["module_type"] != destination["module_type"]: + errors.append( + f"module type mismatch for {key}: source {source['module_type']!r} " + f"does not match model {destination['module_type']!r}" + ) + return errors + + +def _safe_relative_path(path: str) -> bool: + normalized = os.path.normpath(path) + return ( + bool(path) + and not os.path.isabs(path) + and normalized not in {".", ".."} + and not normalized.startswith(".." + os.sep) + ) + + +def _git(directory: str, *args: str) -> None: + result = subprocess.run( + ["git", "-C", directory, *args], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise GitError(result.stderr.strip() or f"git {' '.join(args)} failed") + + +__all__ = ["StateDir", "load", "save"] diff --git a/tests/test_linear.py b/tests/test_linear.py index 7fee220..ea83ba9 100644 --- a/tests/test_linear.py +++ b/tests/test_linear.py @@ -88,6 +88,7 @@ def test_parameter_tree_is_one_directory_per_agent(): with open(os.path.join(store.root, "MODEL.json"), encoding="utf-8") as file: manifest = json.load(file) assert manifest["modules"]["layer"]["parameters"]["weight"]["shape"] == [3] + assert manifest["modules"]["layer"]["parameters"]["weight"]["input_features"] == 2 def test_init_functions_can_reset_materialized_parameters(): diff --git a/tests/test_state_dir.py b/tests/test_state_dir.py new file mode 100644 index 0000000..9d4dd06 --- /dev/null +++ b/tests/test_state_dir.py @@ -0,0 +1,221 @@ +import json +import os +from dataclasses import FrozenInstanceError + +import pytest +from conftest import run_git + +import hytorch +import hytorch.state_dir as state_dir_module +from hytorch import mn +from hytorch._git import Repo + + +class Model(mn.Module): + def __init__(self, width: int = 2, bias: str = "source"): + super().__init__() + self.layer = mn.Linear(1, width, bias=bias) + + def forward(self, value): + return self.layer(value) + + +def test_state_dir_save_and_load_round_trip(tmp_path): + source = Model(bias="trained policy") + source.layer.weight[0]._set_text("print('check')", "tools/check.py") + source._ensure_parameter_store().commit("test: add workspace tool") + source_state = source.state_dir() + checkpoint = tmp_path / "model-state" + + hytorch.save(source_state, checkpoint) + loaded = hytorch.load(checkpoint) + + assert isinstance(source_state, hytorch.StateDir) + assert loaded.commit == source_state.commit + assert os.path.isdir(checkpoint / ".git") + assert json.loads((checkpoint / "MODEL.json").read_text())["format"] == ( + "hytorch-model-v1" + ) + + restored = Model(bias="different policy") + result = restored.load_state_dir(loaded) + assert repr(result) == "" + assert result.missing_keys == [] + assert result.unexpected_keys == [] + assert restored.layer.weight[0].text() == source.layer.weight[0].text() + assert restored.layer.weight[1].text() == source.layer.weight[1].text() + assert restored.layer.weight[0].text("tools/check.py") == "print('check')\n" + assert restored.state_dir().commit != source_state.commit + + with pytest.raises(FrozenInstanceError): + source_state.commit = "changed" + + +def test_save_uses_committed_revision_not_worktree(tmp_path): + model = Model(width=1) + state = model.state_dir() + agents = os.path.join(state.dir, "layers", "layer", "0", "AGENTS.md") + with open(agents, "w", encoding="utf-8") as file: + file.write("uncommitted\n") + + checkpoint = tmp_path / "model-state" + hytorch.save(state, checkpoint) + + assert (checkpoint / "layers/layer/0/AGENTS.md").read_text() != "uncommitted\n" + with pytest.raises(RuntimeError, match="uncommitted changes"): + model.state_dir() + + +def test_save_rejects_existing_or_nested_destination(tmp_path): + state = Model(width=1).state_dir() + existing = tmp_path / "existing" + existing.mkdir() + + with pytest.raises(FileExistsError): + hytorch.save(state, existing) + with pytest.raises(ValueError, match="inside the model state"): + hytorch.save(state, os.path.join(state.dir, "nested")) + + +def test_load_rejects_non_model_repository(new_repo): + with pytest.raises(ValueError, match="MODEL.json"): + hytorch.load(new_repo.root) + + +def test_strict_load_reports_missing_and_unexpected_keys(tmp_path): + class ExpandedModel(Model): + def __init__(self): + super().__init__(width=1) + self.extra = mn.Linear(1, 1, bias="extra") + + one = Model(width=1) + checkpoint = tmp_path / "one" + hytorch.save(one.state_dir(), checkpoint) + state = hytorch.load(checkpoint) + two = ExpandedModel() + + with pytest.raises(RuntimeError, match="Missing key.*extra.weight.0"): + two.load_state_dir(state) + + result = two.load_state_dir(state, strict=False) + assert result.missing_keys == ["extra.weight.0"] + assert result.unexpected_keys == [] + assert two.layer.weight[0].text() == one.layer.weight[0].text() + + expanded_checkpoint = tmp_path / "expanded" + hytorch.save(ExpandedModel().state_dir(), expanded_checkpoint) + simple = Model(width=1) + result = simple.load_state_dir(hytorch.load(expanded_checkpoint), strict=False) + assert result.missing_keys == [] + assert result.unexpected_keys == ["extra.weight.0"] + + +def test_load_rejects_shape_mismatch_even_when_not_strict(tmp_path): + source = Model(width=1) + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + destination = Model(width=2) + + with pytest.raises(RuntimeError, match="size mismatch"): + destination.load_state_dir(hytorch.load(checkpoint), strict=False) + + +def test_load_rejects_logical_input_width_mismatch(tmp_path): + class WideInputModel(mn.Module): + def __init__(self): + super().__init__() + self.layer = mn.Linear(2, 1) + + def forward(self, *values): + return self.layer(*values) + + source = Model(width=1) + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + + with pytest.raises(RuntimeError, match="input feature mismatch"): + WideInputModel().load_state_dir(hytorch.load(checkpoint)) + + +def test_load_accepts_v1_manifest_without_logical_width(tmp_path): + source = Model(width=1) + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + manifest_path = checkpoint / "MODEL.json" + manifest = json.loads(manifest_path.read_text()) + del manifest["modules"]["layer"]["parameters"]["weight"]["input_features"] + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + run_git(str(checkpoint), "add", "MODEL.json") + run_git(str(checkpoint), "commit", "-m", "test: emulate old manifest") + + restored = Model(width=1) + restored.load_state_dir(hytorch.load(checkpoint)) + assert restored.layer.weight[0].text() == source.layer.weight[0].text() + + +def test_state_dir_preserves_model_generation_history(tmp_path): + model = Model(width=1) + initial = model.state_dir() + model.layer.weight[0]._set_text("new policy") + model._parameter_store.commit("test: update policy") + updated = model.state_dir() + checkpoint = tmp_path / "history" + hytorch.save(updated, checkpoint) + saved_repo = Repo.discover(str(checkpoint)) + + assert saved_repo.is_ancestor(initial.commit, updated.commit) + assert saved_repo.resolve("HEAD") == updated.commit + + +def test_load_rejects_pending_optimizer_update(tmp_path): + source = Model(width=1) + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + destination = Model(width=1) + optimizer = hytorch.optim.DFM(destination.parameters()) + optimizer._pending = object() + + with pytest.raises(RuntimeError, match="pending optimizer update"): + destination.load_state_dir(hytorch.load(checkpoint)) + + +def test_missing_saved_workspace_does_not_partially_load(tmp_path): + source = Model(width=2, bias="source") + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + os.unlink(checkpoint / "layers/layer/1/AGENTS.md") + run_git(str(checkpoint), "add", "-A") + run_git(str(checkpoint), "commit", "-m", "test: remove workspace") + state = hytorch.load(checkpoint) + destination = Model(width=2, bias="destination") + before = [view.text() for view in destination.layer.weight.views()] + + with pytest.raises(RuntimeError, match="missing workspace"): + destination.load_state_dir(state) + + assert [view.text() for view in destination.layer.weight.views()] == before + + +def test_copy_failure_does_not_change_canonical_model(tmp_path, monkeypatch): + source = Model(width=2, bias="source") + checkpoint = tmp_path / "source" + hytorch.save(source.state_dir(), checkpoint) + destination = Model(width=2, bias="destination") + before = [view.text() for view in destination.layer.weight.views()] + before_commit = destination.state_dir().commit + original = state_dir_module.copy_tree + copies = 0 + + def fail_second_copy(source, destination): + nonlocal copies + copies += 1 + if copies == 2: + raise OSError("simulated copy failure") + original(source, destination) + + monkeypatch.setattr(state_dir_module, "copy_tree", fail_second_copy) + with pytest.raises(OSError, match="simulated copy failure"): + destination.load_state_dir(hytorch.load(checkpoint)) + + assert destination.state_dir().commit == before_commit + assert [view.text() for view in destination.layer.weight.views()] == before From 84924fb371634f28f4c540ed9551af9e5d7ca762 Mon Sep 17 00:00:00 2001 From: Esteban Safranchik Date: Wed, 5 Aug 2026 23:17:45 -0700 Subject: [PATCH 2/2] Fix Git identity in state directory tests --- tests/conftest.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 01ac9da..b3105bd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,8 +10,19 @@ def run_git(root: str, *args: str) -> str: + env = dict(os.environ) + env.update( + GIT_AUTHOR_NAME="HyTorch Test", + GIT_AUTHOR_EMAIL="hytorch-test@localhost", + GIT_COMMITTER_NAME="HyTorch Test", + GIT_COMMITTER_EMAIL="hytorch-test@localhost", + ) result = subprocess.run( - ["git", "-C", root, *args], capture_output=True, text=True, check=False + ["git", "-C", root, *args], + env=env, + capture_output=True, + text=True, + check=False, ) if result.returncode != 0: raise RuntimeError(f"git {args}: {result.stderr}")