Skip to content

Run synchronous evaluation callbacks on the managed blocking pool #144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Stephen Belanger (Qard) wants to merge 1 commit into from
Closed

Run synchronous evaluation callbacks on the managed blocking pool #144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions src/eval/blocking.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
//! Managed blocking invocation on the run's fixed execution binding.
use crate::internal::execution::ExecutionBinding;

pub(crate) async fn invoke<T: Send + 'static>(
binding: Option<&ExecutionBinding>,
work: impl FnOnce() -> anyhow::Result<T> + Send + 'static,
) -> anyhow::Result<T> {
let binding = binding.ok_or_else(|| {
anyhow::anyhow!("blocking callbacks require a managed evaluation executor")
})?;
let task = binding.spawn_blocking(work);
// Abort queued work on run cancellation. Tokio cannot interrupt a closure
// already running; its owned data remains valid until that closure returns.
struct AbortOnDrop(tokio::task::AbortHandle);
impl Drop for AbortOnDrop {
fn drop(&mut self) {
self.0.abort();
}
}
let _abort = AbortOnDrop(task.abort_handle());
task.await?
}

#[cfg(test)]
mod tests {
use super::*;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};

#[test]
fn dropping_an_invocation_cancels_queued_work_before_it_can_start() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.max_blocking_threads(1)
.build()
.unwrap();
let binding = ExecutionBinding::external(runtime.handle().clone());
let (release, waiting) = std::sync::mpsc::channel();
let (entered, started) = std::sync::mpsc::channel();
let occupied = binding.spawn_blocking(move || {
entered.send(()).unwrap();
waiting
.recv_timeout(std::time::Duration::from_secs(3))
.unwrap()
});
started
.recv_timeout(std::time::Duration::from_secs(3))
.unwrap();
let called = Arc::new(AtomicBool::new(false));
runtime.block_on(async {
{
let called = called.clone();
let callback = invoke(Some(&binding), move || {
called.store(true, Ordering::SeqCst);
Ok(())
});
futures::pin_mut!(callback);
assert!(futures::poll!(callback).is_pending());
}
release.send(()).unwrap();
occupied.await.unwrap();
binding.spawn_blocking(|| ()).await.unwrap();
});
assert!(!called.load(Ordering::SeqCst));
}
}
107 changes: 73 additions & 34 deletions src/eval/case_runner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,54 +1,75 @@
//! One case's computation, independent of scheduling, recording and reporting.
use super::{Case, CaseError, EvalResult, Scorer, ScorerArgs, Task, TaskHooks};
use crate::SpanHandle;
//! One case's computation, independent of recording and reporting.
use super::{Case, CaseError, EvalDefinition, EvalResult, ScorerArgs, TaskHooks};
use crate::{internal::execution::ExecutionBinding, SpanHandle};
use futures::{stream::FuturesOrdered, StreamExt};
use std::sync::Arc;

pub(crate) struct CaseRunner;
impl CaseRunner {
pub async fn run<I, O>(
pub async fn run<I: Send + Sync + 'static, O: Send + Sync + 'static>(
case: Case<I, O>,
task: &dyn Task<I, O>,
scorers: &[Box<dyn Scorer<I, O>>],
definition: &EvalDefinition<I, O>,
span: Option<&SpanHandle>,
) -> EvalResult<I, O>
where
I: Send + Sync,
O: Send + Sync,
{
binding: Option<&ExecutionBinding>,
) -> EvalResult<I, O> {
let case = Arc::new(case);
let hooks = TaskHooks::for_case(
span,
case.expected.as_ref(),
case.metadata.clone(),
case.tags.clone(),
);
let task = definition.task();
let computed = if task.is_blocking() {
let case = case.clone();
let state = hooks.rebind(None, None);
let span = span.cloned();
super::blocking::invoke(binding, move || {
let hooks = state.rebind(span.as_ref(), case.expected.as_ref());
task.run_blocking(&case.input, &hooks)
})
.await
} else {
task.run(&case.input, &hooks).await
};
let mut errors = Vec::new();
let output = match task.run(&case.input, &hooks).await {
Ok(output) => Some(output),
let output = match computed {
Ok(output) => Some(Arc::new(output)),
Err(source) => {
errors.push(CaseError::Task { source });
None
}
};
let mut scores = Vec::new();
if let Some(output) = &output {
// Drive scorers concurrently while keeping definition order in the
// outcome, regardless of their relative completion times.
let mut pending = FuturesOrdered::new();
for scorer in scorers {
let args = ScorerArgs {
input: &case.input,
output,
expected: case.expected.as_ref(),
metadata: case.metadata.as_ref(),
};
for scorer in definition.scorers() {
let case = case.clone();
let output = output.clone();
let scorer = scorer.clone();
pending.push_back(async move {
scorer
.score(args)
.await
.map_err(|source| CaseError::Scorer {
name: scorer.name().to_owned(),
source,
let name = scorer.name().to_owned();
let result = if scorer.is_blocking() {
super::blocking::invoke(binding, move || {
scorer.score_blocking(ScorerArgs {
input: &case.input,
output: &output,
expected: case.expected.as_ref(),
metadata: case.metadata.as_ref(),
})
})
.await
} else {
scorer
.score(ScorerArgs {
input: &case.input,
output: &output,
expected: case.expected.as_ref(),
metadata: case.metadata.as_ref(),
})
.await
};
result.map_err(|source| CaseError::Scorer { name, source })
});
}
while let Some(result) = pending.next().await {
Expand All @@ -61,6 +82,13 @@ impl CaseRunner {
let metadata = hooks.get_metadata();
let tags = hooks.get_tags();
drop(hooks);
// All owned invocations have joined. Borrowed callback arguments cannot
// escape, so no input/output clones are required to return ownership.
let case =
Arc::try_unwrap(case).unwrap_or_else(|_| unreachable!("case invocations joined"));
let output = output.map(|output| {
Arc::try_unwrap(output).unwrap_or_else(|_| unreachable!("scorers joined"))
});
EvalResult {
input: case.input,
output,
Expand All @@ -80,7 +108,7 @@ impl CaseRunner {
#[cfg(test)]
mod tests {
use super::*;
use crate::eval::{Score, Scores, SimpleFnTask};
use crate::eval::{Score, Scorer, Scores, Task};
use async_trait::async_trait;
use serde_json::json;
use std::sync::{
Expand Down Expand Up @@ -123,7 +151,12 @@ mod tests {
xact_id: Some("version".into()),
created: Some("timestamp".into()),
};
let result = futures::executor::block_on(CaseRunner::run(case, &LocalTask, &[], None));
let result = futures::executor::block_on(CaseRunner::run(
case,
&EvalDefinition::new(Box::new(LocalTask), vec![]).unwrap(),
None,
None,
));
assert_eq!(result.output.unwrap().0, 6);
assert_eq!(result.case_id.as_deref(), Some("case"));
assert_eq!(result.xact_id.as_deref(), Some("version"));
Expand Down Expand Up @@ -174,8 +207,8 @@ mod tests {
input: 1,
..Default::default()
},
&FailTask,
&scorers,
&EvalDefinition::new(Box::new(FailTask), scorers).unwrap(),
None,
None,
)
.await;
Expand Down Expand Up @@ -233,8 +266,14 @@ mod tests {
input: 1,
..Default::default()
},
&SimpleFnTask::new(|input: &usize| Ok(*input)),
&scorers,
&EvalDefinition::new(
Box::new(crate::eval::task(|input: &usize, _hooks| {
Box::pin(async move { Ok(*input) })
})),
scorers,
)
.unwrap(),
None,
None,
),
)
Expand Down
42 changes: 26 additions & 16 deletions src/eval/definition.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ pub struct EvalDefinition<I, O> {
}

struct DefinitionInner<I, O> {
task: Box<dyn Task<I, O>>,
scorers: Vec<Box<dyn Scorer<I, O>>>,
task: Arc<dyn Task<I, O>>,
scorers: Vec<Arc<dyn Scorer<I, O>>>,
}

impl<I, O> Clone for EvalDefinition<I, O> {
Expand All @@ -37,7 +37,7 @@ impl<I: Send + Sync + 'static, O: Send + Sync + 'static> EvalDefinition<I, O> {
scorers: Vec<Box<dyn Scorer<I, O>>>,
) -> Result<Self, EvalConfigError> {
let mut names = std::collections::HashSet::new();
let mut validated: Vec<Box<dyn Scorer<I, O>>> = Vec::with_capacity(scorers.len());
let mut validated: Vec<Arc<dyn Scorer<I, O>>> = Vec::with_capacity(scorers.len());
for (index, scorer) in scorers.into_iter().enumerate() {
let name = scorer.name().to_owned();
if name.trim().is_empty() {
Expand All @@ -46,11 +46,11 @@ impl<I: Send + Sync + 'static, O: Send + Sync + 'static> EvalDefinition<I, O> {
if !names.insert(name.clone()) {
return Err(EvalConfigError::DuplicateScorerName { name });
}
validated.push(Box::new(NamedScorer { name, scorer }));
validated.push(Arc::new(NamedScorer { name, scorer }));
}
Ok(Self {
inner: Arc::new(DefinitionInner {
task,
task: Arc::from(task),
scorers: validated,
}),
})
Expand All @@ -63,11 +63,11 @@ impl<I, O> EvalDefinition<I, O> {
self.inner.scorers.iter().map(|scorer| scorer.name())
}

pub(crate) fn task(&self) -> &dyn Task<I, O> {
self.inner.task.as_ref()
pub(crate) fn task(&self) -> Arc<dyn Task<I, O>> {
self.inner.task.clone()
}

pub(crate) fn scorers(&self) -> &[Box<dyn Scorer<I, O>>] {
pub(crate) fn scorers(&self) -> &[Arc<dyn Scorer<I, O>>] {
&self.inner.scorers
}
}
Expand All @@ -79,10 +79,18 @@ struct NamedScorer<I, O> {

#[async_trait]
impl<I: Send + Sync, O: Send + Sync> Scorer<I, O> for NamedScorer<I, O> {
fn is_blocking(&self) -> bool {
self.scorer.is_blocking()
}

fn name(&self) -> &str {
&self.name
}

fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result<Scores> {
self.scorer.score_blocking(args)
}

async fn score(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result<Scores> {
self.scorer.score(args).await
}
Expand All @@ -91,7 +99,7 @@ impl<I: Send + Sync, O: Send + Sync> Scorer<I, O> for NamedScorer<I, O> {
#[cfg(test)]
mod tests {
use super::*;
use crate::eval::{case_runner::CaseRunner, task, Case, CaseError, SimpleFnTask};
use crate::eval::{case_runner::CaseRunner, task, Case, CaseError};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

// No Clone, Serialize, equality or runtime is needed for computation.
Expand Down Expand Up @@ -133,8 +141,8 @@ mod tests {
};
let (one, two) = futures::executor::block_on(async {
futures::join!(
CaseRunner::run(case(1, 10), first.task(), first.scorers(), None),
CaseRunner::run(case(2, 20), other.task(), other.scorers(), None),
CaseRunner::run(case(1, 10), &first, None, None),
CaseRunner::run(case(2, 20), &other, None, None),
)
});
assert_eq!(calls.load(Ordering::SeqCst), 2);
Expand All @@ -145,8 +153,8 @@ mod tests {
// Reuse the same callbacks after both prior invocations have completed.
let (one, two) = futures::executor::block_on(async {
futures::join!(
CaseRunner::run(case(3, 30), first.task(), first.scorers(), None),
CaseRunner::run(case(4, 40), other.task(), other.scorers(), None),
CaseRunner::run(case(3, 30), &first, None, None),
CaseRunner::run(case(4, 40), &other, None, None),
)
});
assert_eq!(one.output.unwrap().0, 33);
Expand All @@ -173,7 +181,9 @@ mod tests {
fn scorer_identity_is_captured_once_for_every_run() {
let changed = Arc::new(AtomicBool::new(false));
let definition = EvalDefinition::new(
Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))),
Box::new(task(|input: &usize, _hooks| {
Box::pin(async move { Ok(*input) })
})),
vec![Box::new(ChangingName(changed.clone()))],
)
.unwrap();
Expand All @@ -185,8 +195,8 @@ mod tests {
input: 1,
..Default::default()
},
definition.task(),
definition.scorers(),
&definition,
None,
None,
));
assert!(
Expand Down
Loading
Loading