Run cache: upload the partial run every --upload-every seconds of wall time (default 600), so that a killed job continues from its last upload; Simulator.run(checkpoint=...) - proof-of-concept/cymsim

proof-of-concept / cymsim

Run cache: upload the partial run every --upload-every seconds of wall time (default 600), so that a killed job continues from its last upload; Simulator.run(checkpoint=...)

Jeremy Magland <jmagland@flatironinstitute.org> committed 4 days ago commit ca554abcd0f2 parent a3f35ac Browse files

5 changed files+75−12

README.mdmodified+1−1View file

@@ -26,7 +26,7 @@ cymsim cache-key my_run.yaml --show # the run's cache key (inde

2626 cymsim run my_run.yaml --cache https://artifact-cache.tempory.net/api/apps/cymsim # take the run from the run cache, or continue a shorter cached run

2727 ```

2828

29-With `--cache URL` (or `CYMSIM_CACHE_URL` set), a run is looked up in a run cache under its key (`cymsim cache-key`), which leaves out the duration and the output naming. The cache holds the longest run of each key so far as a `.cymsim` file without the acceleration record: a run it covers is cut from it without building the model, and a longer run continues from it (`--resume`), identical to an uninterrupted run on the same device. Downloads are kept in `~/.cache/cymsim` (`CYMSIM_CACHE_DIR`) and revalidated by ETag. With the cache app's passcode in `CYMSIM_CACHE_KEY`, a run that was computed (from rest or continued) is uploaded, unless the entry meanwhile holds at least as much; without it the cache is read-only. The entry's metadata records the device, jax version, commit and host of each segment of integration. See [cache.py](src/cymsim/cache.py).

29+With `--cache URL` (or `CYMSIM_CACHE_URL` set), a run is looked up in a run cache under its key (`cymsim cache-key`), which leaves out the duration and the output naming. The cache holds the longest run of each key so far as a `.cymsim` file without the acceleration record: a run it covers is cut from it without building the model, and a longer run continues from it (`--resume`), identical to an uninterrupted run on the same device. Downloads are kept in `~/.cache/cymsim` (`CYMSIM_CACHE_DIR`) and revalidated by ETag. With the cache app's passcode in `CYMSIM_CACHE_KEY`, a run that was computed (from rest or continued) is uploaded, at the end and as a partial run every 10 minutes of wall time (`--upload-every`), so that a killed job continues from its last upload; nothing replaces an entry that meanwhile holds at least as much. Without the passcode the cache is read-only. The entry's metadata records the device, jax version, commit and host of each segment of integration. See [cache.py](src/cymsim/cache.py).

3030

3131 A `.cymsim` file is an HDF5 file holding the configuration, the mesh, the listener signals, displacement frames at a low rate, optionally the normal acceleration of every node at the audio rate, and the integrator state; the layout is specified in [docs/format.md](docs/format.md). It is meant for taking the shell's motion into other programs (other listeners, rooms, acoustic solvers) and is the file the web companion cymsimweb saves and loads. A partial file written during a run is a consistent prefix of the complete one.

3232

src/cymsim/cache.pymodified+18−7View file

@@ -16,7 +16,9 @@ it, a run it does not cover continues from it, and without an entry the run star

1616 Downloads are kept in CYMSIM_CACHE_DIR (default ~/.cache/cymsim) and revalidated by ETag.

1717

1818 With the app's passcode (CYMSIM_CACHE_KEY), a run that was computed rather than cut is uploaded,

19-unless the entry meanwhile holds at least as many steps (checked just before the upload; the

19+at the end and as a partial run every ``upload_every`` seconds of wall time (default 10 minutes),

20+so that a killed job continues from its last upload. Nothing replaces an entry that meanwhile holds

21+at least as many steps (checked just before the upload; the

2022 worker has no conditional write). The entry's metadata records, per segment of integration, the

2123 device, jax version, cymsim commit and host; continuing on a different device than the last

2224 segment is allowed but logged, since float32 runs differ bitwise between devices. A failed upload

@@ -182,7 +184,8 @@ def publish(result, url: str, key: str, passcode: str, segments: list[dict], fil

182184 etag_path.write_text(f'"{etag}"')

183185 else:

184186 etag_path.unlink(missing_ok=True)

185- log.info("cache: uploaded %s (%.3f s, %.1f MB)", key, len(result.pressure_rayleigh) / result.fs, len(data) / 1e6)

187+ log.info("cache: uploaded %s%s (%.3f s, %.1f MB)", "partial run " if not result.complete else "", key,

188+ len(result.pressure_rayleigh) / result.fs, len(data) / 1e6)

186189 return True

187190 except (urllib.error.URLError, OSError) as e:

188191 detail = e.read().decode(errors="replace") if isinstance(e, urllib.error.HTTPError) else ""

@@ -229,11 +232,15 @@ def fetch(url: str, key: str, directory: Path | None = None) -> Path | None:

229232 raise

230233

231234

235+UPLOAD_EVERY = 600.0 # default seconds of wall time between partial uploads

236+

237+

232238 def cached_run(config: SimulationConfig, url: str, save_path: str | None = None,

233- save_every: float | None = None, directory: Path | None = None, passcode: str | None = None):

239+ save_every: float | None = None, directory: Path | None = None, passcode: str | None = None,

240+ upload_every: float | None = UPLOAD_EVERY):

234241 """The run of ``config``, taken from or continued from the cache at ``url`` where possible,

235242 and uploaded if it was computed and a passcode is given (``passcode`` or CYMSIM_CACHE_KEY);

236- see the module docstring."""

243+ see the module docstring. ``upload_every`` None or 0 uploads only the finished run."""

237244 passcode = passcode or os.environ.get("CYMSIM_CACHE_KEY") or None

238245 from .cymsim_file import load_cymsim, save_cymsim

239246 from .simulate import Simulator, cut_run

@@ -265,8 +272,12 @@ def cached_run(config: SimulationConfig, url: str, save_path: str | None = None,

265272 if segments and segments[-1].get("to") == prev.n_steps_done and segments[-1].get("device") != device():

266273 log.warning("cache: the cached run was integrated on %s and continues on %s; float32 runs "

267274 "differ bitwise between devices", segments[-1].get("device"), device())

268- result = Simulator(config).run(save_path=save_path, save_every=save_every, resume=prev)

275+ n_from = prev.n_steps_done if prev is not None else 0

276+ upload = lambda res, file=None: publish(res, url, key, passcode, segments + [_segment(n_from, res.n_steps_done)],

277+ file=file, directory=directory)

278+ result = Simulator(config).run(save_path=save_path, save_every=save_every, resume=prev,

279+ checkpoint=upload if passcode and upload_every else None,

280+ checkpoint_every=upload_every)

269281 if passcode:

270- segments.append(_segment(prev.n_steps_done if prev is not None else 0, result.n_steps_done))

271- publish(result, url, key, passcode, segments, file=save_path, directory=directory)

282+ upload(result, save_path)

272283 return result

src/cymsim/cli.pymodified+5−1View file

@@ -41,7 +41,8 @@ def cmd_simulate(args) -> int:

4141 return 2

4242 if cache:

4343 from .cache import cached_run

44- result = cached_run(cfg, cache, save_path=args.output, save_every=args.save_every)

44+ every = {} if args.upload_every is None else {"upload_every": args.upload_every}

45+ result = cached_run(cfg, cache, save_path=args.output, save_every=args.save_every, **every)

4546 else:

4647 sim = Simulator(cfg)

4748 prev = None

@@ -135,6 +136,9 @@ def main(argv=None) -> int:

135136 help="take the run from, or continue it from, the run cache at URL (an artifact-cache app, "

136137 "e.g. https://artifact-cache.tempory.net/api/apps/cymsim); default $CYMSIM_CACHE_URL")

137138 p.add_argument("--no-cache", action="store_true", help="ignore $CYMSIM_CACHE_URL")

139+ p.add_argument("--upload-every", type=float, metavar="SECONDS",

140+ help="with the cache's passcode in $CYMSIM_CACHE_KEY, also upload the partial run every SECONDS "

141+ "of wall time, so that a killed job continues from there (default 600; 0: only at the end)")

138142 p.add_argument("--duration", type=float, help="override numerics.duration [s]")

139143 p.add_argument("--out-prefix", help="override output.prefix")

140144 p.add_argument("--linear", action="store_true", help="drop the quadratic strain terms (reference run)")

src/cymsim/simulate.pymodified+11−3View file

@@ -523,10 +523,13 @@ class Simulator:

523523 self.probe_xy if len(self.probe_xy) else None, None if probe is None else probe[:n_out])

524524

525525 def run(self, save_path: str | None = None, save_every: float | None = None,

526- resume: SimulationResult | None = None) -> SimulationResult:

526+ resume: SimulationResult | None = None, checkpoint=None,

527+ checkpoint_every: float | None = None) -> SimulationResult:

527528 """Integrate the run. With ``save_path`` a .cymsim file is written at the end and, if

528529 ``save_every`` (seconds of wall time) is given, also periodically during the run, each

529- time as a consistent partial file that can be loaded (see docs/format.md).

530+ time as a consistent partial file that can be loaded (see docs/format.md). Likewise

531+ ``checkpoint``, a function of a partial SimulationResult, is called every

532+ ``checkpoint_every`` seconds of wall time (the run cache uploads partial runs with it).

530533

531534 With ``resume``, an earlier run (usually loaded from a .cymsim file) of a configuration

532535 that differs from this one at most in CONTINUATION_FIELDS, the run continues from that

@@ -573,7 +576,7 @@ class Simulator:

573576 from .cymsim_file import CymsimWriter

574577 writer = CymsimWriter(save_path, result_now(n_done, wall0))

575578 t_run = time.time()

576- t_saved = t_run

579+ t_saved = t_checkpoint = t_run

577580 report_every = max(1, self.n_blocks // max(1, num.n_progress))

578581 blk0 = n_done // spb

579582 log.info("running %d blocks of %d steps (nonlinearity meter: max|w|/h < 1 effectively "

@@ -621,12 +624,17 @@ class Simulator:

621624 if not math.isfinite(ke + em + eb):

622625 raise FloatingPointError("energy is not finite: the run has blown up "

623626 "(lower numerics.cfl_safety)")

627+ res = None

624628 if writer is not None and save_every is not None and time.time() - t_saved >= save_every \

625629 and blk < self.n_blocks - 1:

626630 res = result_now(n_done, wall0 + time.time() - t_run)

627631 writer.update(res)

628632 t_saved = time.time()

629633 log.info(" saved partial run (%.3f s) to %s", res.duration_done, save_path)

634+ if checkpoint is not None and checkpoint_every is not None \

635+ and time.time() - t_checkpoint >= checkpoint_every and blk < self.n_blocks - 1:

636+ checkpoint(res or result_now(n_done, wall0 + time.time() - t_run))

637+ t_checkpoint = time.time()

630638

631639 result = result_now(n_done, wall0 + time.time() - t_run)

632640 for h in result.hits:

tests/test_cache_run.pymodified+40−0View file

@@ -266,3 +266,43 @@ def test_a_failed_upload_keeps_the_run(coarse_config, runs, cache, tmp_path, cap

266266 assert res.complete and "upload of" in caplog.text and "Invalid passcode" in caplog.text

267267 assert ("PUT", cache_key(cfg), 401) in cache.requests and not cache.entries

268268 _assert_same_files(out, runs[0.02])

269+

270+

271+def test_partial_runs_are_uploaded(coarse_config, runs, cache, tmp_path):

272+ cfg = _config(coarse_config, 0.05)

273+ key = cache_key(cfg)

274+ cached_run(cfg, cache.url, directory=tmp_path / "c", passcode=PASSCODE, upload_every=1e-9) # after every block

275+ puts = [r for r in cache.requests if r[0] == "PUT"]

276+ assert len(puts) == Simulator(cfg).n_blocks # one per block but the last, and the finished run

277+ _assert_same_files(_load_entry(cache, key, tmp_path / "e.cymsim"), runs[0.05])

278+ (seg,) = cache.metas[key]["segments"] # one segment, extended by each upload

279+ assert seg["from"] == 0 and seg["to"] == cache.metas[key]["n_steps_done"]

280+

281+

282+def test_a_killed_run_continues_from_its_last_upload(coarse_config, runs, cache, tmp_path, monkeypatch):

283+ import cymsim.cache

284+ cfg = _config(coarse_config, 0.05)

285+ key = cache_key(cfg)

286+ real, calls = cymsim.cache.publish, []

287+

288+ def publish_then_die(*args, **kwargs):

289+ real(*args, **kwargs)

290+ calls.append(1)

291+ if len(calls) == 2:

292+ raise KeyboardInterrupt

293+

294+ monkeypatch.setattr(cymsim.cache, "publish", publish_then_die)

295+ with pytest.raises(KeyboardInterrupt):

296+ cached_run(cfg, cache.url, directory=tmp_path / "c", passcode=PASSCODE, upload_every=1e-9)

297+ monkeypatch.undo()

298+ monkeypatch.delenv("CYMSIM_CACHE_KEY", raising=False)

299+ assert not load_cymsim(_load_entry(cache, key, tmp_path / "part.cymsim")).complete

300+ k = cache.metas[key]["n_steps_done"]

301+ cached_run(cfg, cache.url, directory=tmp_path / "c2", passcode=PASSCODE, upload_every=0)

302+ _assert_same_files(_load_entry(cache, key, tmp_path / "e.cymsim"), runs[0.05])

303+ assert [(g["from"], g["to"]) for g in cache.metas[key]["segments"]] == [(0, k), (k, cache.metas[key]["n_steps_done"])]

304+

305+

306+def test_upload_every_zero_uploads_only_the_finished_run(coarse_config, cache, tmp_path):

307+ cached_run(_config(coarse_config, 0.05), cache.url, directory=tmp_path / "c", passcode=PASSCODE, upload_every=0)

308+ assert len([r for r in cache.requests if r[0] == "PUT"]) == 1