Coverage for /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equine/upgrade.py: 94%
105 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 00:14 +0000
1# Copyright 2024, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
2# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
3# SPDX-License-Identifier: MIT
4"""
5Rewrite model files from earlier EQUINE releases in the current format.
7 python -m equine.upgrade MODEL.eq [-o OUT.eq]
9Earlier releases saved pickled Python objects (an ``io.BytesIO`` holding the
10TorchScript archive, ``CovType`` members, scipy KDE objects, numpy strings),
11which ``torch.load(weights_only=True)`` refuses, so EQUINE no longer opens
12those files. The upgrade reads the file with unrestricted unpickling, which
13runs any code embedded in it, so only upgrade files you trust. It then writes
14the same model in the current layout and checks that the safe loader reads it.
16Reading the old layout is deprecated: this module is removed in a later
17release, and files not upgraded by then cannot be opened.
18"""
20import argparse
21import io
22import os
23import shutil
24import sys
25import warnings
26from enum import Enum
27from typing import Any, Optional
29import numpy as np
30import torch
31from scipy.stats import gaussian_kde
33from .equine_protonet import _kde_to_state
34from .utils import (
35 EQUINE_FORMAT_VERSION,
36 LegacyFormatError,
37 jit_archive_to_tensor,
38 load_checkpoint,
39)
41_DEPRECATION = (
42 "Reading model files from earlier EQUINE releases is deprecated and will be "
43 "removed in a later release; upgrade your files now."
44)
47def _plain(value: Any) -> Any:
48 """``value`` with numpy scalars and strings, at any depth, as Python values."""
49 if isinstance(value, dict):
50 return {_plain(k): _plain(v) for k, v in value.items()}
51 if isinstance(value, (list, tuple)):
52 return type(value)(_plain(v) for v in value)
53 if isinstance(value, np.generic):
54 return value.item()
55 return value
58def _current_layout(checkpoint: dict[str, Any]) -> dict[str, Any]:
59 """The legacy checkpoint dictionary rewritten with plain values only."""
60 out = dict(checkpoint)
61 out["equine_format_version"] = EQUINE_FORMAT_VERSION
62 archive = out.get("embed_jit_save")
63 if isinstance(archive, (bytes, bytearray)):
64 archive = io.BytesIO(archive)
65 if isinstance(archive, io.BytesIO): 65 ↛ 67line 65 didn't jump to line 67 because the condition on line 65 was always true
66 out["embed_jit_save"] = jit_archive_to_tensor(archive)
67 settings = dict(out.get("settings", {}))
68 if isinstance(settings.get("cov_type"), Enum):
69 settings["cov_type"] = settings["cov_type"].value
70 out["settings"] = _plain(settings)
71 if isinstance(out.get("support"), dict): 71 ↛ 73line 71 didn't jump to line 73 because the condition on line 71 was always true
72 out["support"] = {int(label): x for label, x in out["support"].items()}
73 if isinstance(out.get("outlier_kde"), dict):
74 out["outlier_kde"] = {
75 int(label): _kde_to_state(kde) if isinstance(kde, gaussian_kde) else kde
76 for label, kde in out["outlier_kde"].items()
77 }
78 for key in ("feature_names", "label_names"):
79 names = out.get(key)
80 if isinstance(names, (list, tuple)) or (
81 isinstance(names, np.ndarray) and names.ndim == 1
82 ): # anything else is left for the loader to report
83 out[key] = [str(name) for name in names]
84 if "train_summary" in out: 84 ↛ 86line 84 didn't jump to line 86 because the condition on line 84 was always true
85 out["train_summary"] = _plain(out["train_summary"])
86 return out
89def upgrade_model_file(path: str, out_path: Optional[str] = None) -> Optional[str]:
90 """
91 Rewrite a trusted model file from an earlier EQUINE release in the current
92 format.
94 Parameters
95 ----------
96 path : str
97 The file to upgrade. It is read with unrestricted unpickling.
98 out_path : Optional[str]
99 Where to write the result; defaults to replacing ``path``. The result
100 is written next to it and moved into place only after the safe loader
101 has read it, so a failed upgrade leaves the original as it was. A
102 symbolic link is followed, so its target is replaced, and the
103 replaced file keeps its permission bits where the filesystem allows
104 it (a ``UserWarning`` is issued otherwise).
106 Returns
107 -------
108 Optional[str]
109 The path of the written file, or None if ``path`` is already in the
110 current format and nothing was written.
112 Raises
113 ------
114 ValueError
115 If the file is not an EQUINE model file, if the safe loader refuses it
116 for a reason other than its age (for example a format version newer
117 than this EQUINE knows), or if the upgraded file is still refused.
118 OSError
119 If the result cannot be written next to ``out_path``.
120 """
121 warnings.warn(_DEPRECATION, FutureWarning, stacklevel=2)
122 try:
123 load_checkpoint(path)
124 except LegacyFormatError:
125 pass
126 else:
127 return None
128 try:
129 checkpoint = torch.load(path, map_location="cpu", weights_only=False)
130 except Exception as err:
131 raise ValueError(
132 f"'{path}' could not be read as a PyTorch file "
133 f"({type(err).__name__}: {err})."
134 ) from err
135 if not isinstance(checkpoint, dict) or "train_summary" not in checkpoint:
136 raise ValueError(f"'{path}' is not an EQUINE model file.")
137 current = _current_layout(checkpoint)
138 out_path = out_path or path
139 if os.path.islink(out_path):
140 out_path = os.path.realpath(out_path) # replace the link's target, not the link
141 directory = os.path.dirname(os.path.abspath(out_path))
142 partial = os.path.join(directory, f".{os.path.basename(out_path)}.upgrading")
143 try:
144 try:
145 torch.save(current, partial)
146 except Exception as err:
147 raise OSError(
148 f"Could not write '{partial}' ({type(err).__name__}: {err})."
149 ) from err
150 try:
151 load_checkpoint(partial)
152 except ValueError as err:
153 reason = str(err).replace(partial, path).split(". Rewrite it once with")[0]
154 raise ValueError(
155 f"The safe loader refuses the rewritten file: {reason}. '{path}' "
156 "holds a value the upgrade does not convert; it was left unchanged."
157 ) from err
158 if os.path.exists(out_path):
159 try:
160 shutil.copymode(out_path, partial)
161 except OSError as err: # a filesystem that rejects mode changes
162 warnings.warn(
163 f"Could not copy the permission bits of '{out_path}' onto the "
164 f"upgraded file ({err}).",
165 UserWarning,
166 stacklevel=2,
167 )
168 os.replace(partial, out_path)
169 finally:
170 if os.path.exists(partial):
171 os.remove(partial)
172 return out_path
175def main(argv: Optional[list[str]] = None) -> int:
176 parser = argparse.ArgumentParser(
177 prog="python -m equine.upgrade",
178 description="Rewrite a model file from an earlier EQUINE release in the "
179 "current format. Only upgrade files you trust: the file is unpickled "
180 "without restrictions, which can run code embedded in it.",
181 )
182 parser.add_argument("path", help="the model file to upgrade")
183 parser.add_argument(
184 "-o", "--output", help="write the result here instead of replacing PATH"
185 )
186 args = parser.parse_args(argv)
187 try:
188 with warnings.catch_warnings(record=True) as record:
189 warnings.simplefilter("always")
190 written = upgrade_model_file(args.path, args.output)
191 except (ValueError, OSError) as err:
192 print(f"error: {err}", file=sys.stderr)
193 return 1
194 for w in record:
195 if issubclass(w.category, UserWarning): 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 print(f"note: {w.message}", file=sys.stderr)
197 if written is None:
198 print(f"{args.path} is already in the current format; nothing written.")
199 return 0
200 if written == args.path:
201 print(f"Upgraded {args.path} in place.")
202 else:
203 print(f"Upgraded {args.path} -> {written}")
204 print(_DEPRECATION, file=sys.stderr)
205 return 0
208if __name__ == "__main__": 208 ↛ 209line 208 didn't jump to line 209 because the condition on line 208 was never true
209 sys.exit(main())