full model: the compile check no longer consumes a step; run() fails instead of spinning when no progress is made

Jeremy Magland <jmagland@flatironinstitute.org> committed 1 month ago commit ff3df5d71a14 parent 5a24e5c Browse files

1 changed file+4−3

src/cymsim/full/simulation.pymodified+4−3View file

@@ -179,9 +179,7 @@ class Simulation:

179179 if self.device.type == "cuda" and self._compile_mode in ("reduce-overhead", "max-autotune"):

180180 self._mark_step = torch.compiler.cudagraph_mark_step_begin

181181 self._mark_step()

182- out = compiled(self.static, *self._state, self._stick)

183- self._state, self._stick = tuple(x.clone() for x in out[:7]), out[7].clone()

184- self.total_steps += 1

182+ compiled(self.static, *self._state, self._stick) # compiles; the result is discarded and the state is unchanged

185183 self._step = compiled

186184 except Exception as e: # noqa: BLE001

187185 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3)

@@ -209,7 +207,10 @@ class Simulation:

209207 target = self.target_samples if seconds is None else min(self.target_samples, self.samples + round(seconds * self.fs))

210208 next_report = self.samples + self.fs

211209 while self.samples < target and not self.done:

210+ before = self.total_steps

212211 self.run_batch()

212+ if self.total_steps == before:

213+ raise RuntimeError("the run made no progress; the step budget is exhausted before the sample target")

213214 if progress and self.samples >= next_report:

214215 print(f"{self.t:.2f} s simulated in {self.seconds:.1f} s wall; {self.status()}", flush=True)

215216 next_report += self.fs