Skip to content

Commit b6b37e2

Browse files
committed
Add batch mode for JSON Lines inputs
`--batch` runs one Function against many inputs in one process. It reads JSON Lines from `--input` or stdin and writes one JSON record per input line to stdout. The Function module, its provider, and the optional schema and query are loaded once, so each input pays only for the run. - Each record has the 1-based input `line`, so records match inputs even when blank lines are skipped. - Records are written with serde, so errors with quotes or newlines stay valid JSON Lines. - By default the batch stops at the first failed input. `--batch-continue-on-error` runs all inputs. The exit code is non-zero if any input failed. - `--batch-full-output` writes the full run result for each input. - A summary goes to stderr. `BluejaySchemaAnalyzer::with_analyzer` parses the schema and query once and computes the scale factor for each input.
1 parent 7962a53 commit b6b37e2

5 files changed

Lines changed: 753 additions & 16 deletions

File tree

‎README.md‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,72 @@ contains binaries that can be run on your computer.
2323

2424
To see the list of possible commands and arguments, run `function-runner --help`.
2525

26+
### Batch mode
27+
28+
Use `--batch` to run a Function against many inputs in one process. The
29+
Function is loaded and compiled once, so each input costs only the run itself.
30+
31+
The input is [JSON Lines](https://jsonlines.org/): one JSON input per line,
32+
from `--input` or stdin. Blank lines are skipped.
33+
34+
```sh
35+
function-runner -f function.wasm --batch -i inputs.jsonl > results.jsonl
36+
```
37+
38+
For each input, batch mode writes one JSON record on one line to stdout.
39+
`line` is the 1-based line number of the input, so you can match each record
40+
to its input.
41+
42+
- An input that ran:
43+
`{"line":1,"success":true,"instructions":5069,"memory_usage":1088,"logs":"","output":{...}}`.
44+
`success` is `false` if the Function failed. If the output is not valid
45+
JSON, `output` is `null` and `output_error` gives the reason.
46+
- An input that could not run, for example invalid JSON:
47+
`{"line":2,"success":false,"error":"Invalid input JSON: ..."}`. If the
48+
Function itself cannot run, for example because it imports both WASI and a
49+
provider that does not allow WASI, each input gets an error record with that
50+
reason.
51+
52+
Each record is one complete JSON line. If a result cannot be serialized, the
53+
record for that input is an error record, never partial output.
54+
55+
A summary goes to stderr, for example
56+
`Batch complete: 3 inputs processed, 2 successful, 1 failed`.
57+
58+
Batch options:
59+
60+
- `--batch-continue-on-error`: run all inputs even if some fail. Without it,
61+
the batch stops after the first failed input.
62+
- `--batch-full-output`: write the full result for each input, the same fields
63+
as `--json` plus `line`.
64+
65+
The exit code is `0` only if every input succeeds. `--schema-path` and
66+
`--query-path` work in batch mode; the schema and query are parsed once.
67+
Profiling is not available in batch mode.
68+
69+
## Library usage
70+
71+
To compute scale factors for many inputs, use
72+
`bluejay_schema_analyzer::BluejaySchemaAnalyzer::with_analyzer`. It parses and
73+
validates the schema and query once, then calls your closure with an `analyze`
74+
function that returns the scale factor for one input:
75+
76+
```rust
77+
use function_runner::bluejay_schema_analyzer::BluejaySchemaAnalyzer;
78+
79+
let scale_factors = BluejaySchemaAnalyzer::with_analyzer(
80+
&schema,
81+
Some("schema.graphql"),
82+
&query,
83+
Some("input.graphql"),
84+
|analyze| inputs.iter().map(|input| analyze(input)).collect::<anyhow::Result<Vec<f64>>>(),
85+
)??;
86+
```
87+
88+
The outer `Result` holds schema and query errors. The inner `Result` holds
89+
analysis errors for an input. The `test_with_analyzer_analyzes_many_inputs`
90+
test in `src/bluejay_schema_analyzer.rs` runs this pattern.
91+
2692
## Development
2793

2894
Building requires a rust toolchain of `1.66.0` to `1.67.0`. `cargo install --path . --locked` will build

‎src/batch.rs‎

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
use std::io::{BufRead, Write};
2+
use std::path::Path;
3+
4+
use anyhow::{anyhow, bail, Result};
5+
use function_runner::{
6+
bluejay_schema_analyzer::BluejaySchemaAnalyzer,
7+
engine::{run, FunctionRunParams},
8+
function_run_result::FunctionRunResult,
9+
BytesContainer, BytesContainerType, Codec,
10+
};
11+
use serde::Serialize;
12+
use wasmtime::Module;
13+
14+
type ScaleFactorFn<'a> = dyn Fn(&serde_json::Value) -> Result<f64> + 'a;
15+
16+
pub(crate) struct ScaleLimitsSource<'a> {
17+
pub schema: &'a str,
18+
pub schema_path: Option<&'a str>,
19+
pub query: &'a str,
20+
pub query_path: Option<&'a str>,
21+
}
22+
23+
pub(crate) struct BatchOptions<'a> {
24+
pub function_path: &'a Path,
25+
pub export: &'a str,
26+
pub codec: Codec,
27+
pub default_scale_factor: f64,
28+
pub continue_on_error: bool,
29+
pub full_output: bool,
30+
}
31+
32+
#[derive(Serialize)]
33+
struct MinimalRecord<'a> {
34+
line: usize,
35+
success: bool,
36+
instructions: u64,
37+
memory_usage: u64,
38+
logs: &'a str,
39+
output: Option<&'a serde_json::Value>,
40+
#[serde(skip_serializing_if = "Option::is_none")]
41+
output_error: Option<&'a str>,
42+
}
43+
44+
#[derive(Serialize)]
45+
struct FullRecord<'a> {
46+
line: usize,
47+
#[serde(flatten)]
48+
result: &'a FunctionRunResult,
49+
#[serde(skip_serializing_if = "Option::is_none")]
50+
output_error: Option<&'a str>,
51+
}
52+
53+
#[derive(Serialize)]
54+
struct ErrorRecord<'a> {
55+
line: usize,
56+
success: bool,
57+
error: &'a str,
58+
}
59+
60+
#[derive(Default)]
61+
struct Summary {
62+
processed: usize,
63+
successful: usize,
64+
failed: usize,
65+
}
66+
67+
pub(crate) fn run_batch(
68+
input: impl BufRead,
69+
output: impl Write,
70+
module: &Module,
71+
scale_limits: Option<ScaleLimitsSource>,
72+
options: &BatchOptions,
73+
) -> Result<()> {
74+
match scale_limits {
75+
Some(source) => BluejaySchemaAnalyzer::with_analyzer(
76+
source.schema,
77+
source.schema_path,
78+
source.query,
79+
source.query_path,
80+
|analyze| run_lines(input, output, module, Some(analyze), options),
81+
)?,
82+
None => run_lines(input, output, module, None, options),
83+
}
84+
}
85+
86+
fn run_lines(
87+
mut input: impl BufRead,
88+
mut output: impl Write,
89+
module: &Module,
90+
analyze: Option<&ScaleFactorFn<'_>>,
91+
options: &BatchOptions,
92+
) -> Result<()> {
93+
let mut summary = Summary::default();
94+
let mut line_bytes = Vec::new();
95+
let mut record = Vec::new();
96+
let mut line = 0;
97+
98+
let outcome = loop {
99+
line_bytes.clear();
100+
match input.read_until(b'\n', &mut line_bytes) {
101+
Ok(0) => break Ok(()),
102+
Ok(_) => {}
103+
Err(e) => break Err(anyhow!("Couldn't read input line {}: {}", line + 1, e)),
104+
}
105+
line += 1;
106+
107+
if line_bytes.iter().all(u8::is_ascii_whitespace) {
108+
continue;
109+
}
110+
summary.processed += 1;
111+
112+
let failure = match run_line(std::mem::take(&mut line_bytes), module, analyze, options) {
113+
Ok(result) => match write_result_record(&mut record, line, &result, options.full_output) {
114+
Err(error) => Some(format!("Line {line}: {error}")),
115+
Ok(()) if result.success => None,
116+
Ok(()) => Some(format!(
117+
"The Function execution failed on line {line}. Review the logs for more information."
118+
)),
119+
},
120+
Err(error) => {
121+
let error = format!("{error:#}");
122+
write_error_record(&mut record, line, &error);
123+
Some(format!("Line {line}: {error}"))
124+
}
125+
};
126+
127+
if let Err(e) = output.write_all(&record) {
128+
break Err(e.into());
129+
}
130+
131+
match failure {
132+
None => summary.successful += 1,
133+
Some(reason) => {
134+
summary.failed += 1;
135+
if !options.continue_on_error {
136+
break Err(anyhow!(reason));
137+
}
138+
}
139+
}
140+
};
141+
142+
output.flush()?;
143+
144+
let status = if outcome.is_ok() {
145+
"complete"
146+
} else {
147+
"stopped"
148+
};
149+
eprintln!(
150+
"Batch {status}: {} inputs processed, {} successful, {} failed",
151+
summary.processed, summary.successful, summary.failed
152+
);
153+
154+
outcome?;
155+
156+
if summary.failed > 0 {
157+
bail!("{} of {} inputs failed", summary.failed, summary.processed);
158+
}
159+
160+
Ok(())
161+
}
162+
163+
fn run_line(
164+
line_bytes: Vec<u8>,
165+
module: &Module,
166+
analyze: Option<&ScaleFactorFn<'_>>,
167+
options: &BatchOptions,
168+
) -> Result<FunctionRunResult> {
169+
let input = BytesContainer::new(BytesContainerType::Input, options.codec, line_bytes)?;
170+
171+
let scale_factor = match (analyze, input.json_value.as_ref()) {
172+
(Some(analyze), Some(json_value)) => analyze(json_value)?,
173+
_ => options.default_scale_factor,
174+
};
175+
176+
run(FunctionRunParams {
177+
function_path: options.function_path.to_path_buf(),
178+
input,
179+
export: options.export,
180+
profile_opts: None,
181+
scale_factor,
182+
module: module.clone(),
183+
engine: module.engine().clone(),
184+
})
185+
}
186+
187+
fn write_result_record(
188+
record: &mut Vec<u8>,
189+
line: usize,
190+
result: &FunctionRunResult,
191+
full_output: bool,
192+
) -> Result<(), String> {
193+
record.clear();
194+
let output_error = result.output.encoding_error.as_deref();
195+
let serialized = if full_output {
196+
serde_json::to_writer(
197+
&mut *record,
198+
&FullRecord {
199+
line,
200+
result,
201+
output_error,
202+
},
203+
)
204+
} else {
205+
serde_json::to_writer(
206+
&mut *record,
207+
&MinimalRecord {
208+
line,
209+
success: result.success,
210+
instructions: result.instructions,
211+
memory_usage: result.memory_usage,
212+
logs: &result.logs,
213+
output: result.output.json_value.as_ref(),
214+
output_error,
215+
},
216+
)
217+
};
218+
219+
match serialized {
220+
Ok(()) => {
221+
record.push(b'\n');
222+
Ok(())
223+
}
224+
Err(e) => {
225+
let error = format!("Couldn't serialize result: {e}");
226+
write_error_record(record, line, &error);
227+
Err(error)
228+
}
229+
}
230+
}
231+
232+
fn write_error_record(record: &mut Vec<u8>, line: usize, error: &str) {
233+
record.clear();
234+
serde_json::to_writer(
235+
&mut *record,
236+
&ErrorRecord {
237+
line,
238+
success: false,
239+
error,
240+
},
241+
)
242+
.expect("An error record contains only a number, a bool, and a string");
243+
record.push(b'\n');
244+
}

0 commit comments

Comments
 (0)