opal-optimizer 0.3.5__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.
Files changed (49) hide show
  1. opal/__init__.py +7 -0
  2. opal/_version.py +3 -0
  3. opal/agent.py +394 -0
  4. opal/backends/__init__.py +31 -0
  5. opal/backends/_codex_client.py +450 -0
  6. opal/backends/_privacy.py +116 -0
  7. opal/backends/base.py +248 -0
  8. opal/backends/claude.py +665 -0
  9. opal/backends/codex.py +356 -0
  10. opal/backends/mailbox.py +1034 -0
  11. opal/backends/unavailable.py +161 -0
  12. opal/cli.py +396 -0
  13. opal/context.py +111 -0
  14. opal/defaults.json +11 -0
  15. opal/defaults.py +54 -0
  16. opal/estimate.py +120 -0
  17. opal/hf.py +109 -0
  18. opal/misc/__init__.py +2 -0
  19. opal/misc/aggressive.py +157 -0
  20. opal/misc/audit.py +84 -0
  21. opal/misc/boundary.py +226 -0
  22. opal/misc/checkpoint.py +399 -0
  23. opal/misc/ddp.py +147 -0
  24. opal/misc/integrity.py +263 -0
  25. opal/misc/nested.py +67 -0
  26. opal/misc/neutral_analysis.py +166 -0
  27. opal/misc/precision.py +210 -0
  28. opal/misc/recovery.py +231 -0
  29. opal/misc/schedulers.py +214 -0
  30. opal/misc/setup.py +253 -0
  31. opal/misc/skip.py +165 -0
  32. opal/misc/taskfacts.py +293 -0
  33. opal/misc/waiting.py +58 -0
  34. opal/programs.py +482 -0
  35. opal/prompts/example_updater.py +252 -0
  36. opal/prompts/neutral.md +206 -0
  37. opal/report.py +1001 -0
  38. opal/report_page.py +175 -0
  39. opal/report_privacy.py +159 -0
  40. opal/results.py +80 -0
  41. opal/schedule.py +134 -0
  42. opal/summaries.py +224 -0
  43. opal/wrapper.py +1674 -0
  44. opal_optimizer-0.3.5.dist-info/METADATA +232 -0
  45. opal_optimizer-0.3.5.dist-info/RECORD +49 -0
  46. opal_optimizer-0.3.5.dist-info/WHEEL +5 -0
  47. opal_optimizer-0.3.5.dist-info/entry_points.txt +2 -0
  48. opal_optimizer-0.3.5.dist-info/licenses/LICENSE +21 -0
  49. opal_optimizer-0.3.5.dist-info/top_level.txt +1 -0
opal/misc/boundary.py ADDED
@@ -0,0 +1,226 @@
1
+ """Best-effort static diagnostics for program boundary crossings, never an execution policy or a sandbox.
2
+
3
+ Only source text is inspected. Common import aliases are resolved, but runtime aliases, computed attributes and
4
+ other indirect access can evade these heuristics; a flag can also be benign. Nothing here executes agent code,
5
+ touches training objects or RNGs, or changes whether a program is accepted. All diagnostic failures are contained.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import ast
11
+ import sys
12
+
13
+
14
+ _IO_MODULES = {"glob", "shutil", "subprocess", "socket", "urllib", "requests", "http", "httpx", "aiohttp"}
15
+ _DYNAMIC = {"__import__", "__builtins__", "eval", "exec", "compile"}
16
+ _PATH_METHODS = {
17
+ "open",
18
+ "read_text",
19
+ "read_bytes",
20
+ "write_text",
21
+ "write_bytes",
22
+ "iterdir",
23
+ "glob",
24
+ "rglob",
25
+ "mkdir",
26
+ "touch",
27
+ "unlink",
28
+ "rename",
29
+ "rmdir",
30
+ "tofile",
31
+ }
32
+ _SERIALIZE = {
33
+ "load",
34
+ "loads",
35
+ "loadtxt",
36
+ "genfromtxt",
37
+ "save",
38
+ "saves",
39
+ "savetxt",
40
+ "savez",
41
+ "savez_compressed",
42
+ "dump",
43
+ "dumps",
44
+ "memmap",
45
+ "fromfile",
46
+ "tofile",
47
+ }
48
+ _FRAMES = {
49
+ "__globals__",
50
+ "__code__",
51
+ "__closure__",
52
+ "__subclasses__",
53
+ "__mro__",
54
+ "f_back",
55
+ "f_code",
56
+ "f_globals",
57
+ "f_locals",
58
+ "gi_frame",
59
+ "cr_frame",
60
+ "ag_frame",
61
+ "tb_frame",
62
+ }
63
+ _MODEL_NAMES = {"model", "module", "loss", "loss_fn", "loss_func", "loss_function", "criterion"}
64
+
65
+
66
+ def _name(node, aliases):
67
+ if isinstance(node, ast.Name):
68
+ return aliases.get(node.id, node.id)
69
+ if isinstance(node, ast.Attribute):
70
+ return _name(node.value, aliases) + "." + node.attr
71
+ return ""
72
+
73
+
74
+ def _category(name):
75
+ root, _, tail = name.partition(".")
76
+ if root in _IO_MODULES or name in {"open", "io.open"}:
77
+ return "io"
78
+ if root == "os" and (
79
+ tail.startswith("path.")
80
+ or tail
81
+ in {
82
+ "open",
83
+ "read",
84
+ "write",
85
+ "listdir",
86
+ "scandir",
87
+ "walk",
88
+ "fwalk",
89
+ "system",
90
+ "popen",
91
+ "remove",
92
+ "unlink",
93
+ "rename",
94
+ "replace",
95
+ "mkdir",
96
+ "makedirs",
97
+ "rmdir",
98
+ "removedirs",
99
+ }
100
+ ):
101
+ return "io"
102
+ if root in {"numpy", "torch", "pickle"} and name.rsplit(".", 1)[-1] in _SERIALIZE:
103
+ return "io"
104
+ if root in {"importlib", "builtins"} or name in _DYNAMIC:
105
+ return "dynamic"
106
+ if root in {"inspect", "__main__"} or name in {
107
+ "globals",
108
+ "sys.modules",
109
+ "sys._getframe",
110
+ "sys._current_frames",
111
+ "gc.get_objects",
112
+ "gc.get_referrers",
113
+ "traceback.walk_stack",
114
+ "traceback.walk_tb",
115
+ }:
116
+ return "introspection"
117
+ return None
118
+
119
+
120
+ def _imported(node, aliases):
121
+ return _name(node, aliases).split(".")[0] in {name.split(".")[0] for name in aliases.values()}
122
+
123
+
124
+ def _attribute(attr, obj, aliases):
125
+ if attr in {"__builtins__", "__import__"}:
126
+ return "dynamic"
127
+ if attr in _FRAMES or (attr == "__dict__" and _imported(obj, aliases)):
128
+ return "introspection"
129
+ if attr.startswith("register_") and attr.endswith("_hook"):
130
+ return "model_mutation"
131
+ if attr in _PATH_METHODS:
132
+ return "io"
133
+ return None
134
+
135
+
136
+ def _model_object(node, aliases):
137
+ name = _name(node, aliases)
138
+ if isinstance(node, ast.Call) and _name(node.func, aliases) == "type" and node.args:
139
+ return _model_object(node.args[0], aliases)
140
+ return bool(_MODEL_NAMES.intersection(name.split("."))) or name.startswith("torch.nn.")
141
+
142
+
143
+ def _scan(files):
144
+ flags = {}
145
+ for file, source in sorted(files.items()):
146
+ if not file.endswith(".py"):
147
+ continue
148
+ tree = ast.parse(source, filename=file)
149
+ lines, aliases = source.splitlines(), {}
150
+ for node in ast.walk(tree):
151
+ if isinstance(node, ast.Import):
152
+ for a in node.names:
153
+ aliases[a.asname or a.name.split(".")[0]] = a.name if a.asname else a.name.split(".")[0]
154
+ elif isinstance(node, ast.ImportFrom) and node.module and not node.level:
155
+ for a in node.names:
156
+ aliases[a.asname or a.name] = node.module + "." + a.name
157
+
158
+ def flag(node, category, file=file, lines=lines):
159
+ if category:
160
+ line = node.lineno
161
+ flags[file, line, category] = dict(
162
+ file=file, line=line, category=category, snippet=lines[line - 1].strip()[:200]
163
+ )
164
+
165
+ for node in ast.walk(tree):
166
+ if isinstance(node, ast.Import):
167
+ for a in node.names:
168
+ flag(node, _category(a.name))
169
+ elif isinstance(node, ast.ImportFrom) and node.module and not node.level:
170
+ flag(node, _category(node.module))
171
+ for a in node.names:
172
+ flag(node, _category(node.module + "." + a.name))
173
+ elif isinstance(node, (ast.Name, ast.Attribute)):
174
+ flag(node, _category(_name(node, aliases)))
175
+ if isinstance(node, ast.Attribute):
176
+ flag(node, _attribute(node.attr, node.value, aliases))
177
+ if isinstance(node.ctx, ast.Store) and _model_object(node.value, aliases):
178
+ flag(node, "model_mutation")
179
+ elif isinstance(node, ast.Call):
180
+ func = _name(node.func, aliases)
181
+ if func in {"getattr", "setattr", "delattr"} and len(node.args) >= 2:
182
+ obj, attr = node.args[:2]
183
+ if isinstance(attr, ast.Constant) and isinstance(attr.value, str):
184
+ flag(node, _attribute(attr.value, obj, aliases))
185
+ if func != "getattr" and _model_object(obj, aliases):
186
+ flag(node, "model_mutation")
187
+ if func == "vars" and node.args and _imported(node.args[0], aliases):
188
+ flag(node, "introspection")
189
+ return [flags[k] for k in sorted(flags)]
190
+
191
+
192
+ def _failed(exc):
193
+ try:
194
+ print(f"[opal] boundary check failed ({type(exc).__name__}); training continues", file=sys.stderr, flush=True)
195
+ except BaseException: # even a broken stderr must not turn a diagnostic into a training failure
196
+ pass
197
+
198
+
199
+ def check(files):
200
+ """Return file/line/category/snippet records, or [] on failure (including malformed source). Never raises."""
201
+ try:
202
+ return _scan(files)
203
+ except BaseException as exc:
204
+ _failed(exc)
205
+ return []
206
+
207
+
208
+ def record(files, version, run):
209
+ """Called on rank 0 before loading a candidate, including starting, resumed and recovery programs.
210
+
211
+ A candidate that subsequently fails to build can still have flags; these are observations, not rejections.
212
+ """
213
+ try:
214
+ flags = check(files)
215
+ for flag in flags:
216
+ run.event("boundary", version=version, **flag)
217
+ if flags:
218
+ categories = ", ".join(sorted({f["category"] for f in flags}))
219
+ print(
220
+ f"[opal] boundary flags: v{version:03d}, {len(flags)} ({categories}); "
221
+ "review events.jsonl and report.html; execution is unchanged (not sandboxed)",
222
+ file=sys.stderr,
223
+ flush=True,
224
+ )
225
+ except BaseException as exc:
226
+ _failed(exc)
@@ -0,0 +1,399 @@
1
+ """Checkpoints: state_dict / load_state_dict, resuming into the same run folder, and the fallback when a program
2
+ crashed after its trial window in an earlier attempt."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import bisect
7
+ import copy
8
+ import io
9
+ import json
10
+ import shutil
11
+ import traceback
12
+ from pathlib import Path
13
+
14
+ import torch
15
+ import torch.distributed as dist
16
+
17
+ from .. import programs
18
+ from ..agent import Run, now_stamp
19
+ from .nested import devices, restore_devices
20
+ from .schedulers import follows
21
+
22
+
23
+ _OBSERVATION_LISTS = (
24
+ "_float_losses", "_pending_losses", "_analysis_losses", "_analysis_metrics", "_metrics_new",
25
+ "_shown_new", "metrics", "reports", "errors", "_reasons", "_skipped",
26
+ )
27
+
28
+
29
+ def _empty_observations(update):
30
+ return dict(
31
+ **{name: [] for name in _OBSERVATION_LISTS},
32
+ _losses_from=update + 1, _pending_from=update + 1, _pending_omitted=0, _reports_made=0,
33
+ _nonfinite_episode=False, _skips_in_a_row=0, _skip_alarm=False,
34
+ _last_applied=None, _last_answer=None, _observed_loss=None, activations={},
35
+ )
36
+
37
+
38
+ class CheckpointMixin:
39
+ def state_dict(self):
40
+ """Everything needed to resume, including the programs and the agent's session. Never raises because of the
41
+ agent's state: if the updater's state cannot be saved (its state_dict() raises, or holds what a new process
42
+ could not load), the checkpoint leaves it out (the program then restarts with fresh state on resume) and the
43
+ agent is told."""
44
+ up, bad = None, []
45
+ try:
46
+ with programs.keep_param_group_identity():
47
+ up = self.updater.state_dict()
48
+ except programs.AGENT_ERRORS:
49
+ if self.digests[0] == self._starting_digest:
50
+ raise # the training script's own optimizers, as in the plain run
51
+ tb = traceback.format_exc(limit=6)
52
+ self.errors.append(
53
+ dict(
54
+ update=self.update,
55
+ where="checkpointing your updater state",
56
+ error="state_dict() raised; checkpoints leave it out until you fix it:\n" + tb[-3000:],
57
+ )
58
+ )
59
+ last = tb.strip().splitlines()[-1]
60
+ self._warn("state-raised", f"the updater's state_dict() raised ({last}); checkpoints leave it out")
61
+ self._state_dropped("updater", last)
62
+ else:
63
+ if self.digests[0] != self._starting_digest:
64
+ up = programs.normalize_state(up)
65
+ bad = programs.unpicklable(up) if self.digests[0] != self._starting_digest else []
66
+ if bad:
67
+ msg = "the updater's state_dict() holds objects a new process could not load (" + "; ".join(bad) + ")"
68
+ self.errors.append(
69
+ dict(
70
+ update=self.update,
71
+ where="checkpointing your updater state",
72
+ error=msg + "; checkpoints leave it out until you fix it",
73
+ )
74
+ )
75
+ self._warn("state-nonplain", msg + "; checkpoints leave it out")
76
+ self._state_dropped("updater", msg)
77
+ up = None
78
+ mem = io.BytesIO()
79
+ try:
80
+ torch.save(self.memory, mem)
81
+ except Exception as exc:
82
+ mem = io.BytesIO()
83
+ torch.save({}, mem)
84
+ self.errors.append(
85
+ dict(
86
+ update=self.update,
87
+ where="checkpointing analysis memory",
88
+ error=f"memory could not be saved and is left out of checkpoints: {exc}",
89
+ )
90
+ )
91
+ self._warn("memory-pickle", f"analysis memory could not be pickled ({exc}); checkpoints leave it out")
92
+ self._state_dropped("memory", str(exc))
93
+ observations = self._observation_state()
94
+ try:
95
+ event_offset = (self.run.dir / "events.jsonl").stat().st_size if self.run is not None else None
96
+ except OSError:
97
+ event_offset = None
98
+ return dict(
99
+ opal=1,
100
+ observations=observations,
101
+ update=self.update,
102
+ version=self.version,
103
+ files=self.files,
104
+ prev_files=self._prev_files,
105
+ updater=up,
106
+ devices=devices(up),
107
+ memory=mem.getvalue(),
108
+ decisions=self.worker.k if self.worker else 0,
109
+ cost=self.worker.cost if self.worker else 0.0,
110
+ unknown_cost_decisions=sorted(self.worker.unknown_cost_decisions) if self.worker else [],
111
+ **(dict(usage=dict(self.worker.usage)) if self.worker and self.worker.usage is not None else {}),
112
+ backend=self.worker.backend.state() if self.worker else {},
113
+ notes_seen=self.run._notes_seen if self.run is not None else 0,
114
+ dev=list(getattr(self, "_dev", [])), # validation=N: the development measurements so far
115
+ run_dir=str(self.run.dir) if self.run is not None else None,
116
+ event_offset=event_offset,
117
+ **(dict(ranks_repaired_at=list(self._ranks_repaired_at)) if self._ranks_repaired_at else {}),
118
+ # the training script's own optimizers and schedulers, whatever the program's state_dict() holds, and the
119
+ # user's schedule while a program owns it (misc/schedulers.py)
120
+ user_optimizers=[o.state_dict() for o in self.user_optimizers],
121
+ user_schedulers=[s.state_dict() for s in self.user_schedulers],
122
+ ledger=self._ledger,
123
+ follows=follows(self.updater),
124
+ next_analysis=self._next_analysis,
125
+ # the decisions the first check scheduled when decisions was not set (opal/estimate.py), and whether it ran
126
+ decision_schedule=self._schedule_segments if getattr(self, "_schedule_changed", False) else None,
127
+ estimated=getattr(self, "_estimated", False),
128
+ )
129
+
130
+ def _observation_state(self):
131
+ values = {name: getattr(self, name) for name in _empty_observations(self.update)}
132
+ values["reports"], values["errors"] = list(self.reports), list(self.errors)
133
+ if torch.is_tensor(self._observed_loss):
134
+ values["_observed_loss"] = self._observed_loss.detach().clone()
135
+ values["loss_buffer"] = self._loss_buf[:self._loss_n].detach().clone()
136
+ defaults = dict(_empty_observations(self.update), loss_buffer=None)
137
+ for name, value in values.items():
138
+ value = programs.normalize_state(value)
139
+ bad = programs.unpicklable(value)
140
+ if bad:
141
+ self._state_dropped("observations." + name, "; ".join(bad))
142
+ value = defaults[name]
143
+ values[name] = value
144
+ return copy.deepcopy(values)
145
+
146
+ def _state_dropped(self, field, reason):
147
+ if self.run is not None:
148
+ self.run.event("state_dropped", update=self.update, version=self.version, field=field, reason=reason)
149
+ self._live.refresh()
150
+
151
+ def _continue_run_dir(self, old):
152
+ """On resume, keep writing into the checkpoint's run folder when this run's folder was auto-named and that folder
153
+ is the same run interrupted: it lies in this run's opal_runs/ (under the current working directory) and its run
154
+ did not finish. Otherwise (a folder named with run_dir=, a checkpoint copied from another place, or one of a
155
+ run that finished) this run writes to its own folder and never touches the checkpoint's."""
156
+ if not self.rank0 or not old or not Path(old).is_dir() or Path(old).resolve() == self.run.dir:
157
+ return
158
+ why = (
159
+ "as you asked (run_dir=)"
160
+ if not self._auto_run_dir
161
+ else _not_the_same_run(Path(old), self.run.dir.parent)
162
+ )
163
+ if why:
164
+ self._log(f"the checkpoint's run folder was {old}; this run writes to {self.run.dir}, {why}")
165
+ # the agent's inbox holds every input so far: bring the earlier ones (the lost attempt's are archived later)
166
+ for f in sorted((Path(old) / "workspace" / "inbox").glob("decision-[0-9][0-9][0-9]*.json")):
167
+ dest = self.run.workspace / "inbox" / f.name
168
+ if not dest.exists() and not f.is_symlink():
169
+ shutil.copy2(f, dest)
170
+ return
171
+ new = self.run.dir
172
+ self.run = Run(old)
173
+ self._prog_root = self.run.programs
174
+ self._run_path = self.run.dir # crash markers go to the folder the run writes to
175
+ self._live.dir = self.run.dir
176
+ if self.worker is not None:
177
+ self.worker.run = self.run
178
+ shutil.rmtree(new, ignore_errors=True)
179
+
180
+ def _archive_lost_attempt(self, version, decisions, update, event_offset=None):
181
+ """Records written after the checkpoint by the attempt that was interrupted move to attempts/<time>/, so the
182
+ resumed run neither overwrites them nor mixes them into its curves."""
183
+ dest = self.run.dir / "attempts" / now_stamp()
184
+ moved = 0
185
+ for d in list(self.run.decisions.glob("[0-9][0-9][0-9]*")):
186
+ if d.name.isdigit() and int(d.name) > decisions:
187
+ (dest / "decisions").mkdir(parents=True, exist_ok=True)
188
+ shutil.move(str(d), str(dest / "decisions" / d.name))
189
+ moved += 1
190
+ for d in list(self.run.programs.glob("v[0-9][0-9][0-9]*")):
191
+ if d.name[1:].isdigit() and int(d.name[1:]) > version:
192
+ (dest / "programs").mkdir(parents=True, exist_ok=True)
193
+ shutil.move(str(d), str(dest / "programs" / d.name))
194
+ moved += 1
195
+ for f in list((self.run.workspace / "inbox").glob("decision-[0-9][0-9][0-9]*.json")):
196
+ try:
197
+ k = int(f.stem.split("-")[1])
198
+ except ValueError:
199
+ continue
200
+ if k > decisions: # the agent reads inbox/ as every input so far: not the lost attempt's
201
+ (dest / "inbox").mkdir(parents=True, exist_ok=True)
202
+ shutil.move(str(f), str(dest / "inbox" / f.name))
203
+ moved += 1
204
+ for name, key in (("losses.jsonl", "to_update"), ("metrics.jsonl", "update"), ("reports.jsonl", "update")):
205
+ f = self.run.dir / name
206
+ if not f.exists():
207
+ continue
208
+ keep, later = [], []
209
+ for line in f.read_text().splitlines():
210
+ try:
211
+ (later if json.loads(line).get(key, 0) > update else keep).append(line)
212
+ except json.JSONDecodeError:
213
+ keep.append(line)
214
+ if later:
215
+ dest.mkdir(parents=True, exist_ok=True)
216
+ (dest / name).write_text("\n".join(later) + "\n")
217
+ f.write_text("".join(x + "\n" for x in keep))
218
+ moved += 1
219
+ f = self.run.dir / "events.jsonl"
220
+ if f.exists():
221
+ data = f.read_bytes()
222
+ if event_offset is not None and 0 <= event_offset <= len(data):
223
+ keep, later = data[:event_offset], data[event_offset:]
224
+ else: # checkpoints from before the exact event offset was recorded
225
+ keep, later = [], []
226
+ for line in data.splitlines(keepends=True):
227
+ try:
228
+ e = json.loads(line)
229
+ lost = (e.get("update", e.get("applied_at", 0)) > update or e.get("k", 0) > decisions
230
+ or e.get("version", 0) > version or e.get("kind") in ("end", "interrupted"))
231
+ except (ValueError, TypeError):
232
+ lost = False
233
+ (later if lost else keep).append(line)
234
+ keep, later = b"".join(keep), b"".join(later)
235
+ if later:
236
+ dest.mkdir(parents=True, exist_ok=True)
237
+ (dest / "events.jsonl").write_bytes(later)
238
+ f.write_bytes(keep)
239
+ moved += 1
240
+ return dest if moved else None
241
+
242
+ def _finish_resume(self):
243
+ """At the first step after load_state_dict, when the training script has also loaded its own scheduler states:
244
+ a crashed program that was replaced hands the user's schedule over as at a swap, and the trial window that a
245
+ resume opens gets its record of the schedule."""
246
+ pending, self._resume_pending = self._resume_pending, None
247
+ if pending.get("crashed"):
248
+ self._schedule_handover(None, self.updater, self._schedule_state(), None)
249
+ self._new_install_record()
250
+
251
+ def _crashed_digests(self, *folders):
252
+ out = set()
253
+ for folder in folders:
254
+ if folder and Path(folder).is_dir():
255
+ for f in Path(folder).glob("crashed-*-r*.json"):
256
+ out.add(f.name.split("-")[1])
257
+ return out
258
+
259
+ def load_state_dict(self, state):
260
+ """Resume: call after wrap(...) with the same arguments, on every rank."""
261
+ if state.get("opal") != 1:
262
+ raise ValueError("opal: not an OPAL optimizer state")
263
+ self._continue_run_dir(state.get("run_dir"))
264
+ if self.dist: # every rank writes its crash marker to the folder rank 0 writes to
265
+ box = [str(self.run.dir) if self.rank0 else None]
266
+ dist.broadcast_object_list(box, src=0, group=self.ctrl)
267
+ self._run_path = Path(box[0])
268
+ files, prev = state["files"], state.get("prev_files")
269
+ # programs that crashed after their trial window (markers left by any rank of any earlier attempt, in this
270
+ # run folder or the checkpoint's) are not resumed: fall back to the predecessor, or to the starting program
271
+ choice = None
272
+ if self.rank0:
273
+ crashed = self._crashed_digests(self.run.dir, state.get("run_dir"))
274
+ digest = lambda fs: programs.side_digests(fs)[0] if fs else None # noqa: E731
275
+ if digest(files) in crashed:
276
+ choice = prev if prev and digest(prev) not in crashed else self._starting_files
277
+ if self.dist:
278
+ box = [choice]
279
+ dist.broadcast_object_list(box, src=0, group=self.ctrl)
280
+ choice = box[0]
281
+ crashed = choice is not None
282
+ if crashed:
283
+ self.errors.append(
284
+ dict(
285
+ update=state["update"],
286
+ where=f"program v{state['version']:03d} (crashed after "
287
+ "this checkpoint in an earlier attempt; the run resumed with an earlier program)",
288
+ error="see crashed-*.json in the run folder",
289
+ )
290
+ )
291
+ if self.rank0:
292
+ same_folder = state.get("run_dir") and Path(state["run_dir"]).resolve() == self.run.dir
293
+ lost = self._archive_lost_attempt(state["version"], state.get("decisions", 0), state["update"],
294
+ state.get("event_offset") if same_folder else None)
295
+ if lost:
296
+ self._log(f"records of the interrupted attempt after the checkpoint moved to {lost}")
297
+ self.update = state["update"]
298
+ self._ranks_repaired_at = list(state.get("ranks_repaired_at", []))
299
+ if getattr(self, "_integrity", None) is not None: # misc/integrity.py: the sampled history is not this run's
300
+ self._integrity.reset()
301
+ self.version = state["version"] - (0 if crashed else 1)
302
+ # Existing checkpoints may contain values outside today's new-program allowlist. The checkpoint
303
+ # loader already reconstructed them; do not reject a legacy program's initial bookkeeping state.
304
+ err = self._install(choice if crashed else files, previous=None, handover=False, restoring=True)
305
+ if err:
306
+ raise RuntimeError(f"opal: the checkpointed programs failed to load:\n{err}")
307
+ with programs.keep_param_group_identity(): # before the program's own state, which may hold them too
308
+ for o, s in zip(self.user_optimizers, state.get("user_optimizers") or []):
309
+ o.load_state_dict(s)
310
+ for s, st in zip(self.user_schedulers, state.get("user_schedulers") or []):
311
+ s.load_state_dict(st)
312
+ self._ledger = state.get("ledger")
313
+ if state.get("updater") is not None and not crashed:
314
+ with programs.keep_param_group_identity():
315
+ self.updater.load_state_dict(restore_devices(state["updater"], state.get("devices"), self.device))
316
+ elif not crashed:
317
+ self._warn(
318
+ "state-missing",
319
+ "the checkpoint has no updater state (it could not be saved); the program restarts its state",
320
+ )
321
+ self.memory = torch.load(
322
+ io.BytesIO(state["memory"]), weights_only=False,
323
+ map_location=lambda storage, location: storage if location == "cpu" else storage.to(device=self.device),
324
+ )
325
+ self._dev = [d for d in state.get("dev", []) if d.get("update", 0) <= self.update]
326
+ if state.get("estimated"): # the first check of the checkpointed run measured it (and may have capped it)
327
+ self._estimate_due, self._estimated = False, True
328
+ if state.get("decision_schedule") is not None and self.rank0:
329
+ self._set_schedule([list(seg) for seg in state["decision_schedule"]])
330
+ self._sched_i = bisect.bisect_right(self.schedule, self.update)
331
+ observations = _empty_observations(self.update)
332
+ observations.update(copy.deepcopy(state.get("observations", {})))
333
+ loss_buffer = observations.pop("loss_buffer", None)
334
+ self._loss_n = 0 if loss_buffer is None else loss_buffer.numel()
335
+ self._loss_buf = torch.zeros(max(self.default_every, self._loss_n), dtype=torch.float32,
336
+ device=self.device)
337
+ if self._loss_n:
338
+ self._loss_buf[:self._loss_n].copy_(loss_buffer)
339
+ # Restore only known fields, keeping the report/error bounds and clearing legacy buffers.
340
+ for name in _empty_observations(self.update):
341
+ if name in ("reports", "errors"):
342
+ target = getattr(self, name)
343
+ target.clear()
344
+ target.extend(observations[name])
345
+ else:
346
+ setattr(self, name, observations[name])
347
+ self._next_analysis = state.get("next_analysis") or self.update + self.every # the analysis keeps its grid
348
+ self._audit_at = self.update + 2
349
+ # A checkpoint has no pre-swap rollback copy or evidence that ranks agreed before its updater was built.
350
+ self._rank_snapshot, self._rank_trial_equal, self._rank_checked_at = None, False, None
351
+ # a new trial window (decided by rank 0): if the restored program raises soon, every rank goes back together
352
+ if self._from_rank0(self.files != self._starting_files):
353
+ fallback = self._starting_files if crashed else (prev or self._starting_files)
354
+ self._prev_files, self._prev_updater = fallback, None
355
+ self._trial_until, self._snapshot, self._first_pending = self.update + self.trial_steps, None, False
356
+ self._forget_retired()
357
+ # finished at the first step, once the training script has loaded its own scheduler states too
358
+ self._resume_pending = dict(crashed=crashed)
359
+ if self.worker:
360
+ self.worker.k, self.worker.cost = state.get("decisions", 0), state.get("cost", 0.0)
361
+ self.worker.unknown_cost_decisions = set(state.get("unknown_cost_decisions", []))
362
+ self.worker.usage = dict(state["usage"]) if state.get("usage") is not None else None
363
+ self.worker.backend.restore(state.get("backend", {}))
364
+ self.worker.backend.setup(self.run.dir, self.run.workspace)
365
+ if self.run is not None:
366
+ self.run._notes_seen = state.get("notes_seen", 0)
367
+ self._resume_note = (
368
+ f"The run was resumed from a checkpoint at update {self.update:,}; anything you decided after that "
369
+ "point in the interrupted attempt was lost, and the programs are the ones running at the checkpoint"
370
+ + (" (except that a program which crashed after it was replaced by an earlier one)" if crashed else "")
371
+ + "."
372
+ )
373
+ self.run.event(
374
+ "resumed", update=self.update, version=self.version, crashed_program_replaced=crashed,
375
+ schedule=self.schedule,
376
+ **(dict(ranks_repaired_at=self._ranks_repaired_at) if self._ranks_repaired_at else {}),
377
+ )
378
+
379
+
380
+ def _not_the_same_run(old: Path, runs: Path):
381
+ """Why a checkpoint's run folder is not this run's to continue (None if it is): it lies outside this run's
382
+ opal_runs/ (a checkpoint copied from another place), or its run finished (its last start, resume or end is an
383
+ end), so a resume from it is a new run."""
384
+ try:
385
+ if old.resolve().parent != Path(runs).resolve():
386
+ return f"since that folder is not in {runs} (a checkpoint copied from another place)"
387
+ last = None
388
+ for line in (old / "events.jsonl").read_text().splitlines():
389
+ try:
390
+ kind = json.loads(line).get("kind")
391
+ except json.JSONDecodeError:
392
+ continue
393
+ if kind in ("start", "resumed", "end", "interrupted"):
394
+ last = kind
395
+ if last == "end":
396
+ return "since the run of that folder finished (a resume from its checkpoint is a new run)"
397
+ except OSError:
398
+ return "since that folder could not be read"
399
+ return None