Skip to content

Commit 035aef5

Browse files
committed
prover: the derive gate's two checks never panic
A thread that holds a derive permit and asks again now passes straight through with no bytes of its own instead of asserting: it cannot park the thread, at worst it over-admits, and the gate counts it as a re-entry. A build's fork (map_maybe_parallel) counts a permit held across it on that permit's gate instead of asserting. Both counts are in the gate's summary line and the tests read them, in debug and release alike; the prover's release path no longer has a panic for either state.
1 parent 81ea9bd commit 035aef5

3 files changed

Lines changed: 143 additions & 97 deletions

File tree

‎prover/src/lfm/block_tree_tests.rs‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4152,8 +4152,8 @@ fn the_streaming_derivation_derives_the_kept_trees_top() {
41524152
/// under a fixed budget of half the open peak, never under the largest set.
41534153
/// The gated derivation's running total stays within its budget and waits,
41544154
/// and its top is the open one's; each derivation's pool high-water is
4155-
/// printed. Then the mutation: a permit held across a build's forks is
4156-
/// refused at the first fork.
4155+
/// printed. Then the mutation: a permit held across each build's forks is
4156+
/// counted there, and the tree still derives.
41574157
#[cfg(feature = "cuda")]
41584158
#[test]
41594159
#[ignore = "box tier: --features cuda, three derivations of a six-leaf tree"]
@@ -4196,30 +4196,31 @@ fn the_derive_gate_bounds_the_derivation_on_the_card() {
41964196
assert_eq!(gated.dispatches, open.dispatches);
41974197
assert_eq!(gated_top, open_top, "the gate is scheduling only");
41984198

4199-
// ⛔ The mutation: each build holds a permit across its walk's forks.
4199+
assert_eq!((gated.reentries, gated.forks_held), (0, 0), "{gated}");
4200+
4201+
// ⛔ The mutation: each build holds a permit across its walk's forks. Every
4202+
// fork is counted, and a job stolen there that asks again passes through:
4203+
// the tree still derives, to the same top.
42004204
pin_setting(Some(Setting::Fixed(u64::MAX)));
42014205
let armed = super::derive_gate::arm().expect("the gate arms");
4202-
let refused = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4203-
plan.derive_tree(&wrap_opts, &|p| {
4206+
let (tree, _) = plan
4207+
.derive_tree(&wrap_opts, &|p| {
42044208
let _held = super::derive_gate::admit(0).expect("armed");
42054209
super::block_plan::artifacts_of(p, &wrap_opts)
42064210
})
4207-
}));
4211+
.expect("the tree derives under the mutation");
42084212
drop(armed);
42094213
pin_setting(None);
4210-
let msg = match refused {
4211-
Ok(_) => panic!("a permit held across a build's forks went unnoticed"),
4212-
Err(e) => e
4213-
.downcast_ref::<String>()
4214-
.cloned()
4215-
.or_else(|| e.downcast_ref::<&str>().map(|s| s.to_string()))
4216-
.unwrap_or_default(),
4217-
};
4218-
// A job stolen at an earlier fork that asks again is the same hazard,
4219-
// refused at the gate.
4214+
let held = take_summary().expect("the mutation armed the gate");
4215+
println!("DERIVE GATE mutation (a permit around each build): {held}");
42204216
assert!(
4221-
msg.contains("held across a rayon fork") || msg.contains("already holds a derive permit"),
4222-
"refused for: {msg}"
4217+
held.forks_held > 0,
4218+
"no fork under a held permit was counted: {held}"
4219+
);
4220+
assert_eq!(
4221+
tree.last().expect("a top level")[0].1.program_id,
4222+
open_top,
4223+
"the mutation changes scheduling only"
42234224
);
42244225
}
42254226

‎prover/src/lfm/derive_gate.rs‎

Lines changed: 121 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,15 @@
1717
//! around its device dispatch, which forks nothing. A holder that waited inside
1818
//! rayon — a `join` whose other half was stolen, a parallel iterator — would
1919
//! run the pool's queued jobs on its own thread meanwhile, and one of them that
20-
//! asks this gate again parks the thread with the bytes it holds (the count
21-
//! gate's hazard, and #1014's W3 hang on BIG 569). Every fork of a build
22-
//! asserts it ([`assert_none_held`], at `registry::map_maybe_parallel`), and
23-
//! so does [`ByteGate::admit`]. A waiter is a rayon worker parked on a condvar:
24-
//! it steals nothing, and the holders it waits for need no rayon work to
25-
//! finish.
20+
//! asks this gate again would park the thread with the bytes it holds (the
21+
//! count gate's hazard, and #1014's W3 hang on BIG 569). That state is made
22+
//! unreachable rather than fatal: a thread that holds a permit and asks again
23+
//! passes straight through without taking more bytes ([`ByteGate::admit`]: at
24+
//! worst it over-admits, never deadlocks), and every fork of a build counts a
25+
//! permit held across it ([`note_fork`], at `registry::map_maybe_parallel`).
26+
//! Both counts are in the gate's [`Summary`], and the tests read them. A
27+
//! waiter is a rayon worker parked on a condvar: it steals nothing, and the
28+
//! holders it waits for need no rayon work to finish.
2629
//!
2730
//! `LAMBDA_VM_BLOCK_DERIVE_GATE`: unset or `auto` (the default) calibrates the
2831
//! budget at the arming, `off` keeps no total (each commit checked alone, as
@@ -34,7 +37,7 @@
3437
// and the fork's assert reach it.
3538
#![cfg_attr(not(feature = "cuda"), allow(dead_code))]
3639

37-
use std::cell::Cell;
40+
use std::cell::RefCell;
3841
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
3942
use std::sync::{Arc, Condvar, Mutex};
4043
use std::time::Instant;
@@ -94,21 +97,20 @@ pub(crate) fn pin_setting(s: Option<Setting>) {
9497
}
9598

9699
thread_local! {
97-
/// [`BytePermit`]s held by this thread.
98-
static HELD: Cell<usize> = const { Cell::new(0) };
100+
/// The gate whose permit this thread holds, and how many permits of it
101+
/// (re-entries included).
102+
static HELD: RefCell<Option<(Arc<ByteGate>, usize)>> = const { RefCell::new(None) };
99103
}
100104

101-
/// Panic if this thread holds a derive permit at `site`, a rayon fork: a job
102-
/// the thread runs while it waits there could ask the gate again and park the
103-
/// thread with the bytes it holds.
104-
pub(crate) fn assert_none_held(site: &str) {
105-
assert_eq!(
106-
HELD.with(Cell::get),
107-
0,
108-
"a derive permit is held across a rayon fork ({site}): a job this thread runs while it \
109-
waits there can ask the gate again and park it; take the permit around the device \
110-
dispatch alone"
111-
);
105+
/// A rayon fork of a build: counted on the gate whose permit this thread
106+
/// holds, if any — a job the thread runs while it waits there could ask the
107+
/// gate again ([`ByteGate::admit`] lets it through).
108+
pub(crate) fn note_fork() {
109+
HELD.with(|h| {
110+
if let Some((gate, _)) = h.borrow().as_ref() {
111+
gate.forks_held.fetch_add(1, Ordering::Relaxed);
112+
}
113+
});
112114
}
113115

114116
/// Whether a set of `bytes` enters beside `used` held under `budget`: alone,
@@ -126,12 +128,18 @@ pub(crate) struct ByteGate {
126128
dispatches: AtomicUsize,
127129
waits: AtomicUsize,
128130
waited_nanos: AtomicU64,
131+
/// Admissions by a thread that already held a permit, let through.
132+
reentries: AtomicUsize,
133+
/// Rayon forks of a build made while holding a permit ([`note_fork`]).
134+
forks_held: AtomicUsize,
129135
}
130136

131137
/// Bytes of a [`ByteGate`], given back on drop.
132138
pub(crate) struct BytePermit {
133139
gate: Arc<ByteGate>,
134140
bytes: u64,
141+
/// Taken by a thread that already held a permit: no bytes of its own.
142+
reentry: bool,
135143
}
136144

137145
impl ByteGate {
@@ -144,17 +152,30 @@ impl ByteGate {
144152
dispatches: AtomicUsize::new(0),
145153
waits: AtomicUsize::new(0),
146154
waited_nanos: AtomicU64::new(0),
155+
reentries: AtomicUsize::new(0),
156+
forks_held: AtomicUsize::new(0),
147157
}
148158
}
149159

150-
/// Take `bytes`, waiting until they [`fits`] beside what is held.
160+
/// Take `bytes`, waiting until they [`fits`] beside what is held. A
161+
/// thread that already holds a permit passes straight through with no
162+
/// bytes: waiting would park it with the bytes it holds.
151163
pub(crate) fn admit(self: &Arc<Self>, bytes: u64) -> BytePermit {
152-
assert_eq!(
153-
HELD.with(Cell::get),
154-
0,
155-
"this thread already holds a derive permit; asking again could park it forever \
156-
with the bytes it holds"
157-
);
164+
let reentry = HELD.with(|h| match h.borrow_mut().as_mut() {
165+
Some((_, n)) => {
166+
*n += 1;
167+
true
168+
}
169+
None => false,
170+
});
171+
if reentry {
172+
self.reentries.fetch_add(1, Ordering::Relaxed);
173+
return BytePermit {
174+
gate: Arc::clone(self),
175+
bytes: 0,
176+
reentry: true,
177+
};
178+
}
158179
let mut used = self.used.lock().unwrap_or_else(|e| e.into_inner());
159180
let mut since = None;
160181
while !fits(*used, bytes, self.budget) {
@@ -170,10 +191,11 @@ impl ByteGate {
170191
self.waited_nanos
171192
.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
172193
}
173-
HELD.with(|h| h.set(h.get() + 1));
194+
HELD.with(|h| *h.borrow_mut() = Some((Arc::clone(self), 1)));
174195
BytePermit {
175196
gate: Arc::clone(self),
176197
bytes,
198+
reentry: false,
177199
}
178200
}
179201

@@ -189,19 +211,31 @@ impl ByteGate {
189211
dispatches: self.dispatches.load(Ordering::Relaxed),
190212
waits: self.waits.load(Ordering::Relaxed),
191213
waited_secs: self.waited_nanos.load(Ordering::Relaxed) as f64 * 1e-9,
214+
reentries: self.reentries.load(Ordering::Relaxed),
215+
forks_held: self.forks_held.load(Ordering::Relaxed),
192216
}
193217
}
194218
}
195219

196220
impl Drop for BytePermit {
197221
fn drop(&mut self) {
198-
let mut used = self.gate.used.lock().unwrap_or_else(|e| e.into_inner());
199-
*used = used.saturating_sub(self.bytes);
200-
drop(used);
201-
// All: the waiters ask for different sizes, and the one woken alone
202-
// might not fit while another would.
203-
self.gate.freed.notify_all();
204-
HELD.with(|h| h.set(h.get().saturating_sub(1)));
222+
if !self.reentry {
223+
let mut used = self.gate.used.lock().unwrap_or_else(|e| e.into_inner());
224+
*used = used.saturating_sub(self.bytes);
225+
drop(used);
226+
// All: the waiters ask for different sizes, and the one woken alone
227+
// might not fit while another would.
228+
self.gate.freed.notify_all();
229+
}
230+
HELD.with(|h| {
231+
let mut h = h.borrow_mut();
232+
if let Some((_, n)) = h.as_mut() {
233+
*n -= 1;
234+
if *n == 0 {
235+
*h = None;
236+
}
237+
}
238+
});
205239
}
206240
}
207241

@@ -213,18 +247,23 @@ pub(crate) struct Summary {
213247
pub dispatches: usize,
214248
pub waits: usize,
215249
pub waited_secs: f64,
250+
pub reentries: usize,
251+
pub forks_held: usize,
216252
}
217253

218254
impl std::fmt::Display for Summary {
219255
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
220256
write!(
221257
f,
222-
"peak {:.2} GiB of a {:.2} GiB budget · {} device sets · {} waited (Σ {:.2}s)",
258+
"peak {:.2} GiB of a {:.2} GiB budget · {} device sets · {} waited (Σ {:.2}s) · {} \
259+
re-entries · {} forks under a permit",
223260
self.peak as f64 / GIB,
224261
self.budget as f64 / GIB,
225262
self.dispatches,
226263
self.waits,
227-
self.waited_secs
264+
self.waited_secs,
265+
self.reentries,
266+
self.forks_held
228267
)
229268
}
230269
}
@@ -305,7 +344,7 @@ impl Drop for DeriveGate {
305344
}
306345

307346
/// Take `bytes` of the armed gate for one device dispatch; `None` while it is
308-
/// not armed. Hold the permit around the dispatch alone ([`assert_none_held`]).
347+
/// not armed. Hold the permit around the dispatch alone ([`note_fork`]).
309348
pub(crate) fn admit(bytes: u64) -> Option<BytePermit> {
310349
let gate = ARMED
311350
.lock()
@@ -353,15 +392,34 @@ mod tests {
353392
assert!(fits(u64::MAX - 1, 10, u64::MAX));
354393
}
355394

395+
/// A thread that holds a permit and asks again passes straight through,
396+
/// with no bytes of its own, and the gate counts it.
356397
#[test]
357-
fn a_thread_asking_twice_is_refused() {
398+
fn a_thread_asking_twice_passes_through() {
358399
let gate = Arc::new(ByteGate::new(10));
359400
let held = gate.admit(1);
360-
let again = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| gate.admit(1)));
361-
assert!(again.is_err(), "a second permit on one thread is refused");
401+
let again = gate.admit(1);
402+
assert_eq!(
403+
*gate.used.lock().unwrap(),
404+
1,
405+
"the second permit took no bytes"
406+
);
407+
drop(again);
408+
assert_eq!(
409+
*gate.used.lock().unwrap(),
410+
1,
411+
"the first still holds its byte"
412+
);
362413
drop(held);
414+
assert_eq!(*gate.used.lock().unwrap(), 0);
415+
assert_eq!(gate.summary().reentries, 1);
363416
drop(gate.admit(1));
364-
assert_eq!(HELD.with(Cell::get), 0);
417+
assert_eq!(
418+
gate.summary().reentries,
419+
1,
420+
"once released, a permit is a first one"
421+
);
422+
assert!(HELD.with(|h| h.borrow().is_none()));
365423
}
366424

367425
/// One synthetic leaf's device sets, in `BuildPlan::walk`'s order: eleven
@@ -459,55 +517,42 @@ mod tests {
459517
assert!(s.peak > largest, "one set at a time: {s}");
460518
assert_eq!(s.dispatches, leaves.len() * 22);
461519
assert!(s.waits > 0, "nothing waited: {s}");
520+
assert_eq!(
521+
(s.reentries, s.forks_held),
522+
(0, 0),
523+
"held only around dispatches: {s}"
524+
);
462525
}
463526

464-
/// The build's fork refuses a thread that holds a permit, before it forks.
527+
/// The build's fork counts a permit its thread holds, on that permit's
528+
/// gate, and nothing when it holds none.
465529
#[cfg(feature = "parallel")]
466530
#[test]
467-
fn the_builds_fork_refuses_a_held_permit() {
531+
fn the_builds_fork_counts_a_held_permit() {
468532
let gate = Arc::new(ByteGate::new(10));
469533
let fork = |held: bool| {
470534
let _p = held.then(|| gate.admit(1));
471-
std::panic::catch_unwind(|| {
472-
super::super::registry::map_maybe_parallel(&[1, 2, 3], |x| x + 1)
473-
})
535+
super::super::registry::map_maybe_parallel(&[1, 2, 3], |x| x + 1)
474536
};
475-
assert_eq!(fork(false).ok(), Some(vec![2, 3, 4]));
476-
let refused = fork(true).expect_err("a fork under a held permit is refused");
477-
let msg = refused
478-
.downcast_ref::<String>()
479-
.cloned()
480-
.unwrap_or_default();
481-
assert!(
482-
msg.contains("held across a rayon fork"),
483-
"refused for: {msg}"
484-
);
537+
assert_eq!(fork(false), vec![2, 3, 4]);
538+
assert_eq!(gate.summary().forks_held, 0);
539+
assert_eq!(fork(true), vec![2, 3, 4]);
540+
assert_eq!(gate.summary().forks_held, 1);
485541
}
486542

487543
/// ⛔ The mutation: a permit held across the walk's forks (the count gate's
488-
/// placement) is refused, at the fork's assert or, when the holder ran a
489-
/// stolen job that asked again first, at the gate's.
544+
/// placement) is counted at every fork, and a job stolen there that asks
545+
/// again passes through rather than parking the thread: the derivation
546+
/// finishes.
490547
#[cfg(feature = "parallel")]
491548
#[test]
492-
fn a_permit_held_across_a_rayon_fork_is_refused() {
549+
fn a_permit_held_across_a_rayon_fork_is_counted() {
493550
let leaves = leaves();
494551
let gate = Arc::new(ByteGate::new(u64::MAX));
495-
let refused = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
496-
pool().install(|| derive(&gate, &leaves, Hold::Walk))
497-
}));
498-
let msg = match refused {
499-
Ok(()) => panic!("a permit held across the walk's forks went unnoticed"),
500-
Err(e) => e
501-
.downcast_ref::<String>()
502-
.cloned()
503-
.or_else(|| e.downcast_ref::<&str>().map(|s| s.to_string()))
504-
.unwrap_or_default(),
505-
};
506-
assert!(
507-
msg.contains("held across a rayon fork")
508-
|| msg.contains("already holds a derive permit"),
509-
"refused for: {msg}"
510-
);
552+
pool().install(|| derive(&gate, &leaves, Hold::Walk));
553+
let s = gate.summary();
554+
// Each leaf's walk forks three row-pair windows and its one-row list.
555+
assert_eq!(s.forks_held, leaves.len() * 4, "{s}");
511556
assert_eq!(*gate.used.lock().unwrap(), 0, "every permit was given back");
512557
}
513558
}

0 commit comments

Comments
 (0)