process replay without global list [pr] (tinygrad#8143)

tinygrad/engine/schedule.py

@@ -1,4 +1,4 @@
-import sys, atexit, functools
+import sys, functools
 from collections import defaultdict, deque
 from dataclasses import dataclass, field
 from typing import FrozenSet, Set, Tuple, List, Dict, Optional, DefaultDict
@@ -211,15 +211,9 @@ def full_ast_rewrite(pre:UOp, ctx:ScheduleContext) -> Tuple[UOp, ScheduleItemCon
         and ShapeTracker.from_shape(s.shape).shrink(m) == s.shrink(m)) for x in ops):
       raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n"
                          +colored("   - a += a.T\n", "red")+colored("   + a += a.T.contiguous()", "green"))
-  if getenv("RUN_PROCESS_REPLAY"): PROCESS_REPLAY_CAPTURE.append(((pre, ctx.assigns), sink))
+  if getenv("RUN_PROCESS_REPLAY"): diskcache_put("schedule_process_replay", str(pre.key), (pre, si_ctx.assigns, {}, sink))
   return sink, si_ctx
 
-PROCESS_REPLAY_CAPTURE: List[Tuple[Tuple[UOp, Set[UOp]], UOp]] = []
-if getenv("RUN_PROCESS_REPLAY"):
-  @atexit.register
-  def save_process_replay() -> None:
-    for x,ret in PROCESS_REPLAY_CAPTURE: diskcache_put("schedule_process_replay", str(x[0].key), (*x, {}, ret))
-
 # **** Schedule grouping
 
 def uval(u:UOp) -> UOp: