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.
- opal/__init__.py +7 -0
- opal/_version.py +3 -0
- opal/agent.py +394 -0
- opal/backends/__init__.py +31 -0
- opal/backends/_codex_client.py +450 -0
- opal/backends/_privacy.py +116 -0
- opal/backends/base.py +248 -0
- opal/backends/claude.py +665 -0
- opal/backends/codex.py +356 -0
- opal/backends/mailbox.py +1034 -0
- opal/backends/unavailable.py +161 -0
- opal/cli.py +396 -0
- opal/context.py +111 -0
- opal/defaults.json +11 -0
- opal/defaults.py +54 -0
- opal/estimate.py +120 -0
- opal/hf.py +109 -0
- opal/misc/__init__.py +2 -0
- opal/misc/aggressive.py +157 -0
- opal/misc/audit.py +84 -0
- opal/misc/boundary.py +226 -0
- opal/misc/checkpoint.py +399 -0
- opal/misc/ddp.py +147 -0
- opal/misc/integrity.py +263 -0
- opal/misc/nested.py +67 -0
- opal/misc/neutral_analysis.py +166 -0
- opal/misc/precision.py +210 -0
- opal/misc/recovery.py +231 -0
- opal/misc/schedulers.py +214 -0
- opal/misc/setup.py +253 -0
- opal/misc/skip.py +165 -0
- opal/misc/taskfacts.py +293 -0
- opal/misc/waiting.py +58 -0
- opal/programs.py +482 -0
- opal/prompts/example_updater.py +252 -0
- opal/prompts/neutral.md +206 -0
- opal/report.py +1001 -0
- opal/report_page.py +175 -0
- opal/report_privacy.py +159 -0
- opal/results.py +80 -0
- opal/schedule.py +134 -0
- opal/summaries.py +224 -0
- opal/wrapper.py +1674 -0
- opal_optimizer-0.3.5.dist-info/METADATA +232 -0
- opal_optimizer-0.3.5.dist-info/RECORD +49 -0
- opal_optimizer-0.3.5.dist-info/WHEEL +5 -0
- opal_optimizer-0.3.5.dist-info/entry_points.txt +2 -0
- opal_optimizer-0.3.5.dist-info/licenses/LICENSE +21 -0
- 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)
|
opal/misc/checkpoint.py
ADDED
|
@@ -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
|