1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/docs/runtime/async_handler.md"))]
2
3use crate::core::address::Address;
4use crate::core::distribution::Distribution;
5use crate::core::model::Model;
6use crate::runtime::handler::Handler;
7use crate::runtime::trace::Trace;
8
9#[allow(async_fn_in_trait)]
59pub trait AsyncHandler {
60 async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64;
62
63 async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool;
65
66 async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64;
68
69 async fn on_sample_usize(&mut self, addr: &Address, dist: &dyn Distribution<usize>) -> usize;
71
72 async fn on_sample_i64(&mut self, addr: &Address, _dist: &dyn Distribution<i64>) -> i64 {
81 panic!(
82 "handler does not implement on_sample_i64 (i64 sample site at {})",
83 addr
84 )
85 }
86
87 async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, value: f64);
89
90 async fn on_observe_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>, value: bool);
92
93 async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, value: u64);
95
96 async fn on_observe_usize(
98 &mut self,
99 addr: &Address,
100 dist: &dyn Distribution<usize>,
101 value: usize,
102 );
103
104 async fn on_observe_i64(&mut self, addr: &Address, _dist: &dyn Distribution<i64>, _value: i64) {
111 panic!(
112 "handler does not implement on_observe_i64 (i64 observe site at {})",
113 addr
114 )
115 }
116
117 async fn on_factor(&mut self, logw: f64);
126
127 async fn finish(self) -> Trace
132 where
133 Self: Sized;
134}
135
136#[derive(Clone, Copy, Debug, Default)]
166pub struct FromSync<H>(pub H);
167
168impl<H: Handler> AsyncHandler for FromSync<H> {
169 async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64 {
170 self.0.on_sample_f64(addr, dist)
171 }
172
173 async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool {
174 self.0.on_sample_bool(addr, dist)
175 }
176
177 async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64 {
178 self.0.on_sample_u64(addr, dist)
179 }
180
181 async fn on_sample_usize(&mut self, addr: &Address, dist: &dyn Distribution<usize>) -> usize {
182 self.0.on_sample_usize(addr, dist)
183 }
184
185 async fn on_sample_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>) -> i64 {
186 self.0.on_sample_i64(addr, dist)
187 }
188
189 async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, value: f64) {
190 self.0.on_observe_f64(addr, dist, value)
191 }
192
193 async fn on_observe_bool(
194 &mut self,
195 addr: &Address,
196 dist: &dyn Distribution<bool>,
197 value: bool,
198 ) {
199 self.0.on_observe_bool(addr, dist, value)
200 }
201
202 async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, value: u64) {
203 self.0.on_observe_u64(addr, dist, value)
204 }
205
206 async fn on_observe_usize(
207 &mut self,
208 addr: &Address,
209 dist: &dyn Distribution<usize>,
210 value: usize,
211 ) {
212 self.0.on_observe_usize(addr, dist, value)
213 }
214
215 async fn on_observe_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>, value: i64) {
216 self.0.on_observe_i64(addr, dist, value)
217 }
218
219 async fn on_factor(&mut self, logw: f64) {
220 self.0.on_factor(logw)
221 }
222
223 async fn finish(self) -> Trace {
224 self.0.finish()
225 }
226}
227
228pub async fn run_async<A, H: AsyncHandler>(mut h: H, m: Model<A>) -> (A, Trace) {
279 let mut m = m;
287 let a = loop {
288 m = match m {
289 Model::Pure(a) => break a,
290 Model::SampleF64 { addr, dist, k } => {
291 let x = h.on_sample_f64(&addr, &*dist).await;
292 k(x)
293 }
294 Model::SampleBool { addr, dist, k } => {
295 let x = h.on_sample_bool(&addr, &*dist).await;
296 k(x)
297 }
298 Model::SampleU64 { addr, dist, k } => {
299 let x = h.on_sample_u64(&addr, &*dist).await;
300 k(x)
301 }
302 Model::SampleUsize { addr, dist, k } => {
303 let x = h.on_sample_usize(&addr, &*dist).await;
304 k(x)
305 }
306 Model::SampleI64 { addr, dist, k } => {
307 let x = h.on_sample_i64(&addr, &*dist).await;
308 k(x)
309 }
310 Model::ObserveF64 {
311 addr,
312 dist,
313 value,
314 k,
315 } => {
316 h.on_observe_f64(&addr, &*dist, value).await;
317 k(())
318 }
319 Model::ObserveBool {
320 addr,
321 dist,
322 value,
323 k,
324 } => {
325 h.on_observe_bool(&addr, &*dist, value).await;
326 k(())
327 }
328 Model::ObserveU64 {
329 addr,
330 dist,
331 value,
332 k,
333 } => {
334 h.on_observe_u64(&addr, &*dist, value).await;
335 k(())
336 }
337 Model::ObserveUsize {
338 addr,
339 dist,
340 value,
341 k,
342 } => {
343 h.on_observe_usize(&addr, &*dist, value).await;
344 k(())
345 }
346 Model::ObserveI64 {
347 addr,
348 dist,
349 value,
350 k,
351 } => {
352 h.on_observe_i64(&addr, &*dist, value).await;
353 k(())
354 }
355 Model::Factor { logw, k } => {
356 h.on_factor(logw).await;
357 k(())
358 }
359 };
360 };
361 let t = h.finish().await;
362 (a, t)
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368 use crate::addr;
369 use crate::core::distribution::*;
370 use crate::core::model::{factor, observe, pure, sample, ModelExt};
371 use crate::runtime::handler::run;
372 use crate::runtime::interpreters::{PriorHandler, ReplayHandler, ScoreGivenTrace};
373 use rand::rngs::StdRng;
374 use rand::{RngCore, SeedableRng};
375 use std::future::Future;
376 use std::panic::{catch_unwind, AssertUnwindSafe};
377 use std::pin::Pin;
378 use std::sync::Arc;
379 use std::task::{Context, Poll, Wake, Waker};
380
381 fn block_on<F: Future>(fut: F) -> F::Output {
384 struct ThreadWaker(std::thread::Thread);
385 impl Wake for ThreadWaker {
386 fn wake(self: Arc<Self>) {
387 self.0.unpark();
388 }
389 }
390 let waker = Waker::from(Arc::new(ThreadWaker(std::thread::current())));
391 let mut cx = Context::from_waker(&waker);
392 let mut fut = std::pin::pin!(fut);
393 loop {
394 match fut.as_mut().poll(&mut cx) {
395 Poll::Ready(out) => return out,
396 Poll::Pending => std::thread::park(),
397 }
398 }
399 }
400
401 struct YieldNow(bool);
404
405 impl Future for YieldNow {
406 type Output = ();
407 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
408 if self.0 {
409 Poll::Ready(())
410 } else {
411 self.0 = true;
412 cx.waker().wake_by_ref();
413 Poll::Pending
414 }
415 }
416 }
417
418 struct Yielding<H>(H);
421
422 impl<H: AsyncHandler> AsyncHandler for Yielding<H> {
423 async fn on_sample_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>) -> f64 {
424 YieldNow(false).await;
425 self.0.on_sample_f64(addr, dist).await
426 }
427 async fn on_sample_bool(&mut self, addr: &Address, dist: &dyn Distribution<bool>) -> bool {
428 YieldNow(false).await;
429 self.0.on_sample_bool(addr, dist).await
430 }
431 async fn on_sample_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>) -> u64 {
432 YieldNow(false).await;
433 self.0.on_sample_u64(addr, dist).await
434 }
435 async fn on_sample_usize(
436 &mut self,
437 addr: &Address,
438 dist: &dyn Distribution<usize>,
439 ) -> usize {
440 YieldNow(false).await;
441 self.0.on_sample_usize(addr, dist).await
442 }
443 async fn on_sample_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>) -> i64 {
444 YieldNow(false).await;
445 self.0.on_sample_i64(addr, dist).await
446 }
447 async fn on_observe_f64(&mut self, addr: &Address, dist: &dyn Distribution<f64>, v: f64) {
448 YieldNow(false).await;
449 self.0.on_observe_f64(addr, dist, v).await
450 }
451 async fn on_observe_bool(
452 &mut self,
453 addr: &Address,
454 dist: &dyn Distribution<bool>,
455 v: bool,
456 ) {
457 YieldNow(false).await;
458 self.0.on_observe_bool(addr, dist, v).await
459 }
460 async fn on_observe_u64(&mut self, addr: &Address, dist: &dyn Distribution<u64>, v: u64) {
461 YieldNow(false).await;
462 self.0.on_observe_u64(addr, dist, v).await
463 }
464 async fn on_observe_usize(
465 &mut self,
466 addr: &Address,
467 dist: &dyn Distribution<usize>,
468 v: usize,
469 ) {
470 YieldNow(false).await;
471 self.0.on_observe_usize(addr, dist, v).await
472 }
473 async fn on_observe_i64(&mut self, addr: &Address, dist: &dyn Distribution<i64>, v: i64) {
474 YieldNow(false).await;
475 self.0.on_observe_i64(addr, dist, v).await
476 }
477 async fn on_factor(&mut self, logw: f64) {
478 YieldNow(false).await;
479 self.0.on_factor(logw).await
480 }
481 async fn finish(self) -> Trace {
482 YieldNow(false).await;
483 self.0.finish().await
484 }
485 }
486
487 type AllTypes = (f64, bool, u64, usize, i64);
488
489 fn all_effects_model() -> Model<AllTypes> {
493 crate::prob! {
494 let x <- sample(addr!("x"), Normal::new(0.0, 1.0).unwrap());
495 let b <- sample(addr!("b"), Bernoulli::new(0.3).unwrap());
496 let n <- sample(addr!("n"), Poisson::new(4.0).unwrap());
497 let c <- sample(addr!("c"), Categorical::new(vec![0.2, 0.5, 0.3]).unwrap());
498 let z <- sample(addr!("z"), DiscreteUniform::new(-3, 3).unwrap());
499 observe(addr!("obs_f64"), Normal::new(x, 0.5).unwrap(), 0.4);
500 observe(addr!("obs_bool"), Bernoulli::new(if b { 0.9 } else { 0.2 }).unwrap(), true);
501 observe(addr!("obs_u64"), Poisson::new(n as f64 + 1.0).unwrap(), 3u64);
502 observe(addr!("obs_usize"), Categorical::uniform(3).unwrap(), c);
503 observe(addr!("obs_i64"), DiscreteUniform::new(z - 1, z + 1).unwrap(), z);
504 factor(-0.5 * x * x);
505 pure((x, b, n, c, z))
506 }
507 }
508
509 fn assert_same_run(expected: &(AllTypes, Trace), actual: &(AllTypes, Trace)) {
512 let ((x, b, n, c, z), t) = expected;
513 let ((ax, ab, an, ac, az), at) = actual;
514 assert_eq!(x.to_bits(), ax.to_bits(), "value x");
515 assert_eq!((b, n, c, z), (ab, an, ac, az), "value");
516 assert_eq!(t.choices.len(), at.choices.len(), "number of choices");
517 for ((addr, choice), (a_addr, a_choice)) in t.choices.iter().zip(&at.choices) {
518 assert_eq!(addr, a_addr);
519 assert_eq!(choice.addr, a_choice.addr);
520 assert_eq!(choice.value, a_choice.value, "value at {addr}");
521 assert_eq!(
522 choice.logp.to_bits(),
523 a_choice.logp.to_bits(),
524 "logp at {addr}"
525 );
526 }
527 assert_eq!(t.log_prior.to_bits(), at.log_prior.to_bits(), "log_prior");
528 assert_eq!(
529 t.log_likelihood.to_bits(),
530 at.log_likelihood.to_bits(),
531 "log_likelihood"
532 );
533 assert_eq!(
534 t.log_factors.to_bits(),
535 at.log_factors.to_bits(),
536 "log_factors"
537 );
538 }
539
540 #[test]
545 fn from_sync_matches_run_for_builtin_handlers() {
546 let mut rngs = [0, 1, 2].map(|_| StdRng::seed_from_u64(62));
548 let [r_sync, r_async, r_yield] = &mut rngs;
549 let prior = |rng| PriorHandler {
550 rng,
551 trace: Trace::default(),
552 };
553 let expected = run(prior(r_sync), all_effects_model());
554 let from_sync = block_on(run_async(FromSync(prior(r_async)), all_effects_model()));
555 let yielding = block_on(run_async(
556 Yielding(FromSync(prior(r_yield))),
557 all_effects_model(),
558 ));
559 assert_same_run(&expected, &from_sync);
560 assert_same_run(&expected, &yielding);
561 let next = rngs.map(|mut rng| rng.next_u64());
562 assert!(next[0] == next[1] && next[0] == next[2], "RNG consumption");
563
564 let prior_trace = expected.1;
566 assert_eq!(prior_trace.choices.len(), 5);
567 assert!(prior_trace.log_prior.is_finite());
568 assert!(prior_trace.log_likelihood.is_finite() && prior_trace.log_likelihood < 0.0);
569 assert!(prior_trace.log_factors <= 0.0);
570
571 let mut base = prior_trace.clone();
574 base.choices.remove(&addr!("n"));
575 let mut rngs = [0, 1, 2].map(|_| StdRng::seed_from_u64(7));
576 let [r_sync, r_async, r_yield] = &mut rngs;
577 let replay = |rng| ReplayHandler {
578 rng,
579 base: base.clone(),
580 trace: Trace::default(),
581 };
582 let expected = run(replay(r_sync), all_effects_model());
583 let from_sync = block_on(run_async(FromSync(replay(r_async)), all_effects_model()));
584 let yielding = block_on(run_async(
585 Yielding(FromSync(replay(r_yield))),
586 all_effects_model(),
587 ));
588 assert_same_run(&expected, &from_sync);
589 assert_same_run(&expected, &yielding);
590 let next = rngs.map(|mut rng| rng.next_u64());
591 assert!(next[0] == next[1] && next[0] == next[2], "RNG consumption");
592
593 let score = || ScoreGivenTrace {
595 base: prior_trace.clone(),
596 trace: Trace::default(),
597 };
598 let expected = run(score(), all_effects_model());
599 let from_sync = block_on(run_async(FromSync(score()), all_effects_model()));
600 let yielding = block_on(run_async(Yielding(FromSync(score())), all_effects_model()));
601 assert_same_run(&expected, &from_sync);
602 assert_same_run(&expected, &yielding);
603 assert_eq!(
604 expected.1.total_log_weight().to_bits(),
605 prior_trace.total_log_weight().to_bits()
606 );
607 }
608
609 #[test]
616 fn run_async_is_stack_safe_for_deep_models() {
617 fn build(i: usize, n: usize, acc: f64) -> Model<f64> {
618 if i >= n {
619 pure(acc)
620 } else {
621 sample(addr!("x", i), Normal::new(0.0, 1.0).unwrap())
622 .bind(move |x| build(i + 1, n, acc + x))
623 }
624 }
625
626 let handle = std::thread::Builder::new()
627 .stack_size(512 * 1024)
628 .spawn(|| {
629 let n = 100_000;
630 let mut rng = StdRng::seed_from_u64(2024);
631 let (sum, trace) = block_on(run_async(
632 FromSync(PriorHandler {
633 rng: &mut rng,
634 trace: Trace::default(),
635 }),
636 build(0, n, 0.0),
637 ));
638 assert!(sum.is_finite());
639 assert_eq!(trace.choices.len(), n);
640 assert!(trace.log_prior.is_finite());
641
642 let mut rng = StdRng::seed_from_u64(2024);
643 let (sum_yielding, trace_yielding) = block_on(run_async(
644 Yielding(FromSync(PriorHandler {
645 rng: &mut rng,
646 trace: Trace::default(),
647 })),
648 build(0, n, 0.0),
649 ));
650 assert_eq!(sum_yielding.to_bits(), sum.to_bits());
651 assert_eq!(trace_yielding.choices.len(), n);
652 assert_eq!(
653 trace_yielding.log_prior.to_bits(),
654 trace.log_prior.to_bits()
655 );
656 })
657 .expect("spawn thread");
658 handle
659 .join()
660 .expect("deep model async interpretation overflowed the stack");
661 }
662
663 struct NoI64;
667
668 impl Handler for NoI64 {
669 fn on_sample_f64(&mut self, _: &Address, _: &dyn Distribution<f64>) -> f64 {
670 unreachable!()
671 }
672 fn on_sample_bool(&mut self, _: &Address, _: &dyn Distribution<bool>) -> bool {
673 unreachable!()
674 }
675 fn on_sample_u64(&mut self, _: &Address, _: &dyn Distribution<u64>) -> u64 {
676 unreachable!()
677 }
678 fn on_sample_usize(&mut self, _: &Address, _: &dyn Distribution<usize>) -> usize {
679 unreachable!()
680 }
681 fn on_observe_f64(&mut self, _: &Address, _: &dyn Distribution<f64>, _: f64) {
682 unreachable!()
683 }
684 fn on_observe_bool(&mut self, _: &Address, _: &dyn Distribution<bool>, _: bool) {
685 unreachable!()
686 }
687 fn on_observe_u64(&mut self, _: &Address, _: &dyn Distribution<u64>, _: u64) {
688 unreachable!()
689 }
690 fn on_observe_usize(&mut self, _: &Address, _: &dyn Distribution<usize>, _: usize) {
691 unreachable!()
692 }
693 fn on_factor(&mut self, _: f64) {
694 unreachable!()
695 }
696 fn finish(self) -> Trace {
697 unreachable!()
698 }
699 }
700
701 impl AsyncHandler for NoI64 {
702 async fn on_sample_f64(&mut self, _: &Address, _: &dyn Distribution<f64>) -> f64 {
703 unreachable!()
704 }
705 async fn on_sample_bool(&mut self, _: &Address, _: &dyn Distribution<bool>) -> bool {
706 unreachable!()
707 }
708 async fn on_sample_u64(&mut self, _: &Address, _: &dyn Distribution<u64>) -> u64 {
709 unreachable!()
710 }
711 async fn on_sample_usize(&mut self, _: &Address, _: &dyn Distribution<usize>) -> usize {
712 unreachable!()
713 }
714 async fn on_observe_f64(&mut self, _: &Address, _: &dyn Distribution<f64>, _: f64) {
715 unreachable!()
716 }
717 async fn on_observe_bool(&mut self, _: &Address, _: &dyn Distribution<bool>, _: bool) {
718 unreachable!()
719 }
720 async fn on_observe_u64(&mut self, _: &Address, _: &dyn Distribution<u64>, _: u64) {
721 unreachable!()
722 }
723 async fn on_observe_usize(&mut self, _: &Address, _: &dyn Distribution<usize>, _: usize) {
724 unreachable!()
725 }
726 async fn on_factor(&mut self, _: f64) {
727 unreachable!()
728 }
729 async fn finish(self) -> Trace {
730 unreachable!()
731 }
732 }
733
734 fn panic_message(f: impl FnOnce()) -> String {
735 let payload = catch_unwind(AssertUnwindSafe(f)).expect_err("expected a panic");
736 match payload.downcast::<String>() {
737 Ok(message) => *message,
738 Err(payload) => payload
739 .downcast_ref::<&str>()
740 .expect("panic payload is a string")
741 .to_string(),
742 }
743 }
744
745 #[test]
750 fn i64_defaults_panic_with_the_handler_message() {
751 let sample_i64 = || sample(addr!("z"), DiscreteUniform::new(-1, 1).unwrap());
752 let documented = "handler does not implement on_sample_i64 (i64 sample site at z)";
753 assert_eq!(
754 panic_message(|| drop(block_on(run_async(NoI64, sample_i64())))),
755 documented
756 );
757 assert_eq!(panic_message(|| drop(run(NoI64, sample_i64()))), documented);
758 assert_eq!(
759 panic_message(|| drop(block_on(run_async(FromSync(NoI64), sample_i64())))),
760 documented
761 );
762
763 let observe_i64 = || observe(addr!("w"), DiscreteUniform::new(-1, 1).unwrap(), 0);
764 let documented = "handler does not implement on_observe_i64 (i64 observe site at w)";
765 assert_eq!(
766 panic_message(|| drop(block_on(run_async(NoI64, observe_i64())))),
767 documented
768 );
769 assert_eq!(
770 panic_message(|| drop(run(NoI64, observe_i64()))),
771 documented
772 );
773 assert_eq!(
774 panic_message(|| drop(block_on(run_async(FromSync(NoI64), observe_i64())))),
775 documented
776 );
777 }
778}