diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 8758609780..128b43400d 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -175,15 +175,16 @@ Per-prediction data. Using separate sockets per slot avoids head-of-line blockin **Worker → Parent:** -| Message | Purpose | -| ---------------------------------------------- | ---------------------------------------------------------------- | -| `Log { source, data }` | Log line from run() | -| `Output { output }` | Yielded output value (for generators/streaming) | -| `FileOutput { filename, kind, mime_type }` | File produced by run() -- referenced by path, uploaded by parent | -| `Metric { name, value, mode }` | Custom metric (mode: `replace`, `increment`, or `append`) | -| `Done { id, output, predict_time, is_stream }` | Prediction completed successfully | -| `Failed { id, error }` | Prediction failed | -| `Cancelled { id }` | Prediction was cancelled | +| Message | Purpose | +| ---------------------------------------------------------------- | ------------------------------------------------------------------------ | +| `ProtocolVersion { version }` | Slot protocol compatibility marker | +| `LogLine { source, data }` | Log line from run() | +| `OutputChunk { output, index }` | Indexed yielded or returned output value | +| `FileOutput { filename, kind, managed, upload_filename, index }` | Indexed file, JSON spill, or structured envelope processed by the parent | +| `Metric { name, value, mode }` | Custom metric (mode: `replace`, `increment`, or `append`) | +| `Done { id, output, predict_time, is_stream }` | Prediction completed successfully | +| `Failed { id, error }` | Prediction failed | +| `Cancelled { id }` | Prediction was cancelled | ## Health State Machine @@ -305,7 +306,7 @@ Following a single prediction from HTTP request to response: 6. **`predict(**kwargs)`called** on the singleton predictor instance. Inputs arrive as native Python types -- strings, ints,`pathlib.Path` objects -- not as a request object or raw JSON. -7. **Outputs stream back** over the slot socket. For generators, each `yield` sends an `Output` message immediately -- true streaming, not buffered. For single return values, one `Output` or `FileOutput` message is sent. +7. **Outputs stream back** over the slot socket. For generators, each `yield` sends an indexed `OutputChunk` or `FileOutput` message immediately. Single return values use the same messages. 8. **File outputs uploaded** by the parent process. `cog.Path` return values are uploaded to the configured storage (or base64-encoded for inline responses). This is transparent to the predictor. @@ -371,9 +372,11 @@ When a prediction input exceeds 6MiB, it's too large to send inline through the ## File Outputs -When run() produces file outputs (`cog.Path`), the worker sends a `FileOutput` message with the filename and MIME type. The parent handles uploading the file (or base64-encoding it for inline responses). The `output_dir` field in the `Predict` request tells the worker where to write output files. +When run() produces file outputs (`cog.Path`), returning the path transfers ownership of its directory entry to Coglet. The worker moves the file into the prediction's managed output directory, falling back to copy then unlink across filesystems. Symlink entries are unlinked after their target bytes are copied, without deleting the target. The worker sends a `FileOutput` message with the managed path, ownership flag, output index, MIME type, and original basename. The parent uses the original basename for uploads and deletes the managed copy after encoding or a successful upload. -`FileOutputKind` distinguishes between normal file outputs (`FileType`) and oversized outputs (`Oversized`) that exceeded an inline size limit. +Files nested in object outputs are described in a structured output envelope so the parent can replace each location with its URL without flattening the object. The envelope itself is stored in managed output storage, which keeps large structured outputs below the IPC frame limit. + +`FileOutputKind` distinguishes normal file outputs (`FileType`), oversized JSON outputs (`Oversized`), and structured output envelopes (`Structured`). Output processing is part of the prediction lifecycle: cancellation and failure paths stop and join pending output tasks before the prediction becomes terminal and its directory can be reused. ## Custom Metrics diff --git a/crates/coglet-python/README.md b/crates/coglet-python/README.md index 51a369b3b8..36c275866f 100644 --- a/crates/coglet-python/README.md +++ b/crates/coglet-python/README.md @@ -49,7 +49,7 @@ coglet-python/ ├── predictor.rs # PythonPredictor: wraps Python Predictor class ├── worker_bridge.rs # PythonPredictHandler: implements PredictHandler ├── input.rs # Input processing (Pydantic validation, ADT) - ├── output.rs # Output processing (make_encodeable, upload_files) + ├── output.rs # Output normalization and managed file staging ├── log_writer.rs # SlotLogWriter, ContextVar routing, SetupLogSender ├── audit.rs # Audit hook, TeeWriter for stream protection └── cancel.rs # Cancellation: SIGUSR1, CancelableGuard diff --git a/crates/coglet-python/src/output.rs b/crates/coglet-python/src/output.rs index c4300d6bb1..4eeac9b3b3 100644 --- a/crates/coglet-python/src/output.rs +++ b/crates/coglet-python/src/output.rs @@ -5,13 +5,15 @@ //! - Dataclasses -> dict (via dataclasses.asdict()) //! - Enums -> .value //! - datetime -> .isoformat() -//! - PathLike -> base64 data URL -//! - IOBase -> base64 data URL +//! - PathLike/IOBase -> managed file descriptors //! - numpy int/float/ndarray -> Python int/float/list //! - dict/list/set/tuple/generator -> recursive descent //! //! This replaces the Python modules cog.json and cog.files. +use coglet_core::bridge::protocol::StructuredOutputFile; +use coglet_core::worker::{SlotSender, StagedFileOutput}; +use pyo3::exceptions::{PyOSError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple}; @@ -19,6 +21,7 @@ use pyo3::types::{PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple}; /// /// Calls make_encodeable() to normalize, then encode_files() to convert any /// remaining Path/IOBase objects to base64 data URLs. +#[cfg(test)] pub fn process_output<'py>( py: Python<'py>, output: &Bound<'py, PyAny>, @@ -27,19 +30,27 @@ pub fn process_output<'py>( encode_files(py, &encodeable) } -/// Process a single output item (for generator outputs). -pub fn process_output_item<'py>( +/// Normalize an output and transfer nested files into Coglet-managed storage. +/// +/// File locations are replaced with null placeholders. The returned descriptors +/// contain RFC 6901 pointers that let the parent restore uploaded URLs or data +/// URIs without flattening the surrounding output. +pub fn process_output_for_ipc<'py>( py: Python<'py>, - item: &Bound<'py, PyAny>, -) -> PyResult> { - process_output(py, item) + output: &Bound<'py, PyAny>, + slot_sender: &SlotSender, +) -> PyResult<(Bound<'py, PyAny>, Vec)> { + let encodeable = make_encodeable(py, output)?; + let mut files = Vec::new(); + let output = stage_files(py, &encodeable, slot_sender, "", &mut files)?; + Ok((output, files)) } /// Normalize a Python object into a JSON-friendly form. /// /// Handles Pydantic models, dataclasses, enums, datetime, numpy types, /// and collections. PathLike objects are passed through (handled later -/// by encode_files). +/// by the managed-file staging pass). fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult> { // Pydantic v2: model_dump() if let Ok(method) = obj.getattr("model_dump") @@ -108,7 +119,7 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult pathlib.Path (will be encoded to base64 later by encode_files) + // os.PathLike -> pathlib.Path (handled later by the staging pass) let os_mod = py.import("os")?; let pathlike_cls = os_mod.getattr("PathLike")?; if obj.is_instance(&pathlike_cls)? { @@ -142,6 +153,7 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult> { // str -- return as-is (don't recurse into characters) if obj.is_instance_of::() { @@ -189,10 +201,121 @@ fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult( + py: Python<'py>, + obj: &Bound<'py, PyAny>, + slot_sender: &SlotSender, + pointer: &str, + files: &mut Vec, +) -> PyResult> { + if obj.is_instance_of::() { + return Ok(obj.clone()); + } + + if let Ok(dict) = obj.cast_exact::() { + let new_dict = PyDict::new(py); + for (key, value) in dict.iter() { + let key_string = serialized_dict_key(py, &key)?; + let child_pointer = json_pointer_child(pointer, &key_string); + new_dict.set_item( + &key, + stage_files(py, &value, slot_sender, &child_pointer, files)?, + )?; + } + return Ok(new_dict.into_any()); + } + + if let Ok(list) = obj.cast_exact::() { + let items = list + .iter() + .enumerate() + .map(|(index, item)| { + let child_pointer = json_pointer_child(pointer, &index.to_string()); + stage_files(py, &item, slot_sender, &child_pointer, files) + }) + .collect::>>()?; + return Ok(PyList::new(py, &items)?.into_any()); + } + + let os = py.import("os")?; + let pathlike = os.getattr("PathLike")?; + if obj.is_instance(&pathlike)? { + let path: String = obj.call_method0("__fspath__")?.extract()?; + let staged = py + .detach(|| slot_sender.stage_user_file_output(std::path::PathBuf::from(path), None)) + .map_err(|error| PyOSError::new_err(error.to_string()))?; + files.push(structured_file(pointer, staged)); + return Ok(py.None().into_bound(py)); + } + + let io = py.import("io")?; + let iobase = io.getattr("IOBase")?; + if obj.is_instance(&iobase)? { + if obj.call_method0("seekable")?.is_truthy()? { + obj.call_method1("seek", (0,))?; + } + let content = obj.call_method0("read")?; + let data = if content.is_instance_of::() { + content.extract::()?.into_bytes() + } else { + content.extract::>()? + }; + let upload_filename = obj + .getattr("name") + .and_then(|name| name.extract::()) + .ok() + .and_then(|name| { + std::path::Path::new(&name) + .file_name() + .and_then(|filename| filename.to_str()) + .map(ToString::to_string) + }); + let staged = py + .detach(|| slot_sender.stage_file_output(&data, upload_filename, None)) + .map_err(|error| PyOSError::new_err(error.to_string()))?; + files.push(structured_file(pointer, staged)); + return Ok(py.None().into_bound(py)); + } + + Ok(obj.clone()) +} + +fn structured_file(pointer: &str, staged: StagedFileOutput) -> StructuredOutputFile { + StructuredOutputFile { + pointer: pointer.to_string(), + filename: staged.filename, + upload_filename: staged.upload_filename, + mime_type: staged.mime_type, + managed: staged.managed, + } +} + +fn json_pointer_child(pointer: &str, segment: &str) -> String { + let segment = segment.replace('~', "~0").replace('/', "~1"); + format!("{pointer}/{segment}") +} + +fn serialized_dict_key(py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult { + let holder = PyDict::new(py); + holder.set_item(key, py.None())?; + let encoded: String = py + .import("json")? + .call_method1("dumps", (&holder,))? + .extract()?; + let object = serde_json::from_str::>(&encoded) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + object + .into_iter() + .next() + .map(|(key, _)| key) + .ok_or_else(|| PyValueError::new_err("output dictionary key was not serialized")) +} + /// Encode a file handle to a base64 data URL. /// /// Seeks to start if seekable, reads all bytes, guesses MIME type from /// the file name, and returns "data:{mime};base64,{encoded}". +#[cfg(test)] fn file_to_base64<'py>(py: Python<'py>, fh: &Bound<'py, PyAny>) -> PyResult> { // Seek to start if possible if let Ok(seekable) = fh.call_method0("seekable") @@ -248,6 +371,100 @@ mod tests { use base64::Engine as _; use pyo3::types::PyDict; + #[test] + fn nested_path_is_staged_with_an_escaped_json_pointer() { + pyo3::Python::initialize(); + let source_dir = tempfile::tempdir().unwrap(); + let source = source_dir.path().join("result.txt"); + std::fs::write(&source, b"nested path").unwrap(); + let output_dir = tempfile::tempdir().unwrap(); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let (output, files) = Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("source", source.to_str().unwrap()).unwrap(); + py.run( + c"from pathlib import Path\nobj = {'a/b~c': [Path(source)]}", + None, + Some(&locals), + ) + .unwrap(); + let obj = locals.get_item("obj").unwrap().unwrap(); + let (processed, files) = process_output_for_ipc(py, &obj, &sender).unwrap(); + let output: String = py + .import("json") + .unwrap() + .call_method1("dumps", (&processed,)) + .unwrap() + .extract() + .unwrap(); + ( + serde_json::from_str::(&output).unwrap(), + files, + ) + }); + + assert_eq!(output, serde_json::json!({"a/b~c": [null]})); + assert_eq!(files.len(), 1); + assert_eq!(files[0].pointer, "/a~1b~0c/0"); + assert_eq!(files[0].upload_filename, "result.txt"); + assert_eq!(std::fs::read(&files[0].filename).unwrap(), b"nested path"); + assert!(!source.exists()); + } + + #[test] + fn nested_path_pointer_uses_json_key_serialization() { + pyo3::Python::initialize(); + let source_dir = tempfile::tempdir().unwrap(); + let source = source_dir.path().join("boolean-key.txt"); + std::fs::write(&source, b"boolean key").unwrap(); + let output_dir = tempfile::tempdir().unwrap(); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let files = Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("source", source.to_str().unwrap()).unwrap(); + py.run( + c"from pathlib import Path\nobj = {True: Path(source)}", + None, + Some(&locals), + ) + .unwrap(); + let obj = locals.get_item("obj").unwrap().unwrap(); + process_output_for_ipc(py, &obj, &sender).unwrap().1 + }); + + assert_eq!(files.len(), 1); + assert_eq!(files[0].pointer, "/true"); + } + + #[test] + fn nested_text_iobase_is_staged_as_utf8() { + pyo3::Python::initialize(); + let output_dir = tempfile::tempdir().unwrap(); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let files = Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"import io\nobj = {'artifact': io.StringIO('hello')}", + None, + Some(&locals), + ) + .unwrap(); + let obj = locals.get_item("obj").unwrap().unwrap(); + process_output_for_ipc(py, &obj, &sender).unwrap().1 + }); + + assert_eq!(files.len(), 1); + assert_eq!(files[0].pointer, "/artifact"); + assert_eq!(files[0].upload_filename, "output.bin"); + assert_eq!(std::fs::read(&files[0].filename).unwrap(), b"hello"); + } + /// Helper: evaluate a Python expression and run make_encodeable on it. fn encodeable(py_expr: &str) -> String { pyo3::Python::initialize(); diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6142e48e24..8116c1adb5 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -164,11 +164,27 @@ fn submit_async_coroutine( Ok(future.unbind()) } +/// Hand a run-returned path to Coglet, which takes ownership of the file. +fn send_path_output( + item: &Bound<'_, PyAny>, + slot_sender: &SlotSender, +) -> Result<(), PredictionError> { + let path: String = item + .call_method0("__fspath__") + .and_then(|path| path.extract()) + .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; + // The transfer is blocking file I/O on files that can be very large, so let + // other Python threads run while it happens. + item.py() + .detach(|| slot_sender.send_user_file_output(std::path::PathBuf::from(path), None)) + .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e))) +} + /// Send a single output item over IPC, routing file outputs to disk. /// -/// For Path outputs (os.PathLike): sends the existing file path via send_file_output. +/// For Path outputs (os.PathLike): transfers the file to Coglet-managed storage. /// For IOBase outputs: reads bytes, writes to output_dir via write_file_output. -/// For everything else: processes through make_encodeable + upload_files, then send_output. +/// For everything else: normalizes the value and stages any nested files. fn send_output_item( py: Python<'_>, item: &Bound<'_, PyAny>, @@ -189,14 +205,9 @@ fn send_output_item( .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if item.is_instance(&pathlike).unwrap_or(false) { - // Path output — file already on disk, send path reference - let path_str: String = item - .call_method0("__fspath__") - .and_then(|p| p.extract()) - .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; - slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None) - .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; + // Returning a Path transfers ownership to Coglet, which moves the data + // into managed storage and deletes it after consumption. + send_path_output(item, slot_sender)?; return Ok(()); } @@ -234,8 +245,9 @@ fn send_output_item( return Ok(()); } - // Non-file output - process normally - let processed = output::process_output_item(py, item) + // Preserve nested file locations so the parent can upload them without + // flattening dict or model outputs. + let (processed, files) = output::process_output_for_ipc(py, item, slot_sender) .map_err(|e| PredictionError::Failed(format!("Failed to process output item: {}", e)))?; let item_str: String = json_module @@ -247,9 +259,17 @@ fn send_output_item( let item_json: serde_json::Value = serde_json::from_str(&item_str) .map_err(|e| PredictionError::Failed(format!("Failed to parse output JSON: {}", e)))?; - slot_sender - .send_output(item_json) - .map_err(|e| PredictionError::Failed(format!("Failed to send output: {}", e)))?; + if files.is_empty() { + slot_sender + .send_output(item_json) + .map_err(|e| PredictionError::Failed(format!("Failed to send output: {}", e)))?; + } else { + slot_sender + .send_structured_output(item_json, files) + .map_err(|e| { + PredictionError::Failed(format!("Failed to send structured output: {}", e)) + })?; + } Ok(()) } @@ -949,15 +969,7 @@ impl PythonPredictor { .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if result.is_instance(&pathlike).unwrap_or(false) { - let path_str: String = result - .call_method0("__fspath__") - .and_then(|p| p.extract()) - .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; - slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None) - .map_err(|e| { - PredictionError::Failed(format!("Failed to send file output: {}", e)) - })?; + send_path_output(result, slot_sender)?; return Ok(PredictionOutput::Single(serde_json::Value::Null)); } @@ -992,9 +1004,7 @@ impl PythonPredictor { return Ok(PredictionOutput::Single(serde_json::Value::Null)); } - // List/tuple output — iterate items so file outputs (Path, IOBase) - // go through the FileOutput IPC path for upload instead of being - // base64-encoded inline by process_output. + // List/tuple output — send each item with its own output index. if let Ok(list) = result.cast::() { for item in list.iter() { send_output_item(py, &item, json_module, slot_sender)?; @@ -1008,22 +1018,8 @@ impl PythonPredictor { return Ok(PredictionOutput::Stream(vec![])); } - // Non-file output — process normally - let processed = output::process_output(py, result) - .map_err(|e| PredictionError::Failed(format!("Failed to process output: {}", e)))?; - - let result_str: String = json_module - .call_method1("dumps", (&processed,)) - .map_err(|e| PredictionError::Failed(format!("Failed to serialize output: {}", e)))? - .extract() - .map_err(|e| { - PredictionError::Failed(format!("Failed to extract output string: {}", e)) - })?; - - let output_json: serde_json::Value = serde_json::from_str(&result_str) - .map_err(|e| PredictionError::Failed(format!("Failed to parse output JSON: {}", e)))?; - - Ok(PredictionOutput::Single(output_json)) + send_output_item(py, result, json_module, slot_sender)?; + Ok(PredictionOutput::Single(serde_json::Value::Null)) } /// Worker mode async predict - submits to shared event loop. @@ -1378,6 +1374,7 @@ mod tests { use std::path::PathBuf; + use coglet_core::bridge::protocol::SlotResponse; use pyo3::types::PyList; fn add_python_sdk_path(py: Python<'_>) { @@ -1435,6 +1432,37 @@ sys.modules.setdefault('requests', requests) }) } + #[test] + fn returned_path_transfers_ownership_to_slot_sender() { + pyo3::Python::initialize(); + let source_dir = tempfile::tempdir().unwrap(); + let source = source_dir.path().join("output.txt"); + std::fs::write(&source, b"output").unwrap(); + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + Python::attach(|py| { + let pathlib = py.import("pathlib").unwrap(); + let path = pathlib + .getattr("Path") + .unwrap() + .call1((source.to_str().unwrap(),)) + .unwrap(); + send_path_output(&path, &sender).unwrap(); + }); + + assert!(!source.exists()); + match rx.try_recv().unwrap() { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => assert_eq!(std::fs::read(filename).unwrap(), b"output"), + response => panic!("expected managed FileOutput, got {response:?}"), + } + } + #[test] fn class_with_run_loads() { let predictor = load_predictor_source( diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 57dcdd40a8..a2f5111d02 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -1,5 +1,7 @@ """Tests for coglet Python bindings.""" +import base64 +import json import queue import re import socket @@ -7,7 +9,9 @@ import sys import threading import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from urllib.parse import quote, urlsplit import coglet import pytest @@ -199,6 +203,51 @@ def predict(self, count: int = 3) -> Iterator[str]: return predictor +@pytest.fixture +def blocking_path_generator_predictor(tmp_path: Path) -> Path: + """Create a streaming predictor that blocks after yielding a Path.""" + predictor = tmp_path / "predict.py" + predictor.write_text(""" +import tempfile +import time +from pathlib import Path +from typing import Iterator + +import cog +from cog import BasePredictor + +class Predictor(BasePredictor): + @cog.streaming + def predict(self, release_path: str) -> Iterator[Path]: + output = Path(tempfile.mkdtemp()) / "result.txt" + output.write_text("streamed file") + yield output + while not Path(release_path).exists(): + time.sleep(0.01) +""") + (tmp_path / "cog.yaml").write_text('predict: "predict.py:Predictor"\n') + schema_dir = tmp_path / ".cog" + schema_dir.mkdir() + (schema_dir / "openapi_schema.json").write_text( + json.dumps( + { + "paths": {"/predictions": {"post": {"x-cog-streaming": True}}}, + "components": { + "schemas": { + "Input": { + "type": "object", + "properties": {"release_path": {"type": "string"}}, + "required": ["release_path"], + }, + "Output": {"type": "array", "items": {"type": "string"}}, + } + }, + } + ) + ) + return predictor + + @pytest.fixture def async_predictor(tmp_path: Path) -> Path: """Create an async predictor.""" @@ -363,12 +412,107 @@ async def predict(self, name: str = "test") -> str: return predictor +@pytest.fixture +def path_output_predictor(tmp_path: Path) -> Path: + """Create a predictor that returns a Path it wrote to its own temp directory. + + This is the shape from issue #1434: the run makes a directory, writes a file + into it, hands the path back, and never touches it again. predict() records + the path it returned in ``returned_path.txt`` so tests can check afterwards + whether the file survived. + """ + predictor = tmp_path / "predict.py" + predictor.write_text(""" +import tempfile +from pathlib import Path + +from cog import BasePredictor + +class Predictor(BasePredictor): + def setup(self): + pass + + def predict(self, text: str = "hello") -> Path: + output = Path(tempfile.mkdtemp()) / "output.txt" + output.write_text(text) + (Path(__file__).parent / "returned_path.txt").write_text(str(output)) + return output +""") + + cog_yaml = tmp_path / "cog.yaml" + cog_yaml.write_text(""" +predict: "predict.py:Predictor" +""") + + return predictor + + +@pytest.fixture +def structured_path_output_predictor(tmp_path: Path) -> Path: + """Create a predictor with a Path nested in a structured output.""" + predictor = tmp_path / "predict.py" + predictor.write_text(""" +import tempfile +from pathlib import Path + +from cog import BaseModel, BasePredictor + +class Output(BaseModel): + artifact: Path + metadata: dict[str, str] + +class Predictor(BasePredictor): + def predict(self, text: str = "nested") -> Output: + output = Path(tempfile.mkdtemp()) / "result?draft.txt" + output.write_text(text) + (Path(__file__).parent / "returned_path.txt").write_text(str(output)) + return Output(artifact=output, metadata={"label": "kept"}) +""") + (tmp_path / "cog.yaml").write_text('predict: "predict.py:Predictor"\n') + return predictor + + +@pytest.fixture +def upload_server(): + """Run a local PUT endpoint and record uploaded paths and bodies.""" + uploads = [] + + class UploadHandler(BaseHTTPRequestHandler): + def do_PUT(self) -> None: + length = int(self.headers.get("Content-Length", "0")) + uploads.append((self.path, self.rfile.read(length))) + self.send_response(200) + safe_request_path = self.path.replace("\r", "").replace("\n", "") + self.send_header( + "Location", + f"http://127.0.0.1:{self.server.server_port}{quote(urlsplit(safe_request_path).path, safe='/%')}", + ) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), UploadHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/upload/", uploads + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + class CogletServer: """Context manager for running coglet server.""" - def __init__(self, predictor_path: Path, port: int = 0): + def __init__( + self, predictor_path: Path, port: int = 0, upload_url: str | None = None + ): self.predictor_path = predictor_path self.requested_port = port + self.upload_url = upload_url self.port = None self.process = None self.stderr_lines = [] @@ -376,10 +520,11 @@ def __init__(self, predictor_path: Path, port: int = 0): self.stderr_thread = None def __enter__(self): + predictor_ref = f"{self.predictor_path}:Predictor" cmd = [ sys.executable, "-c", - f"import coglet; coglet.server.serve('{self.predictor_path}:Predictor', port={self.requested_port})", + f"import coglet; coglet.server.serve({predictor_ref!r}, port={self.requested_port}, upload_url={self.upload_url!r})", ] self.process = subprocess.Popen( cmd, @@ -520,6 +665,53 @@ def test_includes_predict_time(self, sync_predictor: Path): assert result["metrics"]["predict_time"] >= 0 +class TestPathOutput: + """Tests for handing a returned Path over to Coglet.""" + + def test_returned_file_is_encoded_and_then_deleted( + self, path_output_predictor: Path + ): + """Returning a Path transfers ownership: Coglet reads the file, then unlinks it. + + The unit tests cover each half of this on its own. This one is the whole + round trip through a real server, so a break in the Python-to-Rust handoff + cannot slip through with the Rust side still passing. + """ + with CogletServer(path_output_predictor) as server: + result = server.predict({"text": "reclaim me"}) + + assert result["status"] == "succeeded" + output = result["output"] + assert output.startswith("data:"), f"expected a data URL, got {output!r}" + payload = base64.b64decode(output.split("base64,", 1)[1]) + assert payload == b"reclaim me" + + returned = Path( + (path_output_predictor.parent / "returned_path.txt").read_text() + ) + assert not returned.exists(), f"{returned} should have been reclaimed" + + def test_structured_path_is_uploaded_with_its_original_name( + self, + structured_path_output_predictor: Path, + upload_server, + ): + upload_url, uploads = upload_server + with CogletServer( + structured_path_output_predictor, upload_url=upload_url + ) as server: + result = server.predict({"text": "structured bytes"}) + + assert result["status"] == "succeeded", result + assert result["output"]["metadata"] == {"label": "kept"} + assert result["output"]["artifact"].endswith("/upload/result%3Fdraft.txt") + assert uploads == [("/upload/result%3Fdraft.txt", b"structured bytes")] + returned = Path( + (structured_path_output_predictor.parent / "returned_path.txt").read_text() + ) + assert not returned.exists(), f"{returned} should have been reclaimed" + + class TestSecretInput: """Tests for cog.Secret input coercion.""" @@ -594,6 +786,53 @@ def test_custom_count(self, generator_predictor: Path): result = server.predict({"count": 5}) assert len(result["output"]) == 5 + def test_path_output_streams_before_prediction_completes( + self, + blocking_path_generator_predictor: Path, + upload_server, + ): + upload_url, _uploads = upload_server + release_path = blocking_path_generator_predictor.parent / "release" + + with CogletServer( + blocking_path_generator_predictor, upload_url=upload_url + ) as server: + response = requests.post( + f"{server.base_url}/predictions", + headers={"Accept": "text/event-stream"}, + json={"input": {"release_path": str(release_path)}}, + stream=True, + timeout=(3, 5), + ) + response.raise_for_status() + events = response.iter_lines(decode_unicode=True) + event_name = None + output = None + try: + for line in events: + if line.startswith("event: "): + event_name = line.removeprefix("event: ") + elif line.startswith("data: ") and event_name == "output": + output = json.loads(line.removeprefix("data: ")) + break + + assert output is not None, "SSE stream ended before the file output" + assert output["index"] == 0 + assert output["chunk"].endswith("/upload/result.txt") + assert not release_path.exists() + finally: + release_path.touch() + + completed = None + for line in events: + if line.startswith("event: "): + event_name = line.removeprefix("event: ") + elif line.startswith("data: ") and event_name == "completed": + completed = json.loads(line.removeprefix("data: ")) + break + assert completed is not None + assert completed["status"] == "succeeded" + class TestAsyncPredictor: """Tests for async predictor.""" diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 0ef46a4e68..a1eab04b18 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -280,6 +280,31 @@ pub enum FileOutputKind { FileType, /// Output exceeds size threshold for bridge codec serialization but is not a file-like return type Oversized, + /// Output contains nested files and is stored as a structured output envelope. + Structured, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct StructuredOutput { + pub output: serde_json::Value, + pub files: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct StructuredOutputFile { + /// RFC 6901 JSON pointer identifying the file's location in `output`. + pub pointer: String, + pub filename: String, + pub upload_filename: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub mime_type: Option, + /// True if Coglet owns this file and may delete it after consumption. + #[serde(default)] + pub managed: bool, +} + +fn is_zero(value: &u64) -> bool { + *value == 0 } /// Accumulation mode for user metrics. @@ -299,7 +324,7 @@ pub enum MetricMode { /// The response enum is already serde-tagged with `type`, so this constant and /// the optional `ProtocolVersion` message give future protocol changes an /// explicit marker without adding a second envelope around every message. -pub const SLOT_RESPONSE_PROTOCOL_VERSION: u32 = 1; +pub const SLOT_RESPONSE_PROTOCOL_VERSION: u32 = 2; /// Messages from worker to parent on slot socket. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -307,9 +332,8 @@ pub const SLOT_RESPONSE_PROTOCOL_VERSION: u32 = 1; pub enum SlotResponse { /// Protocol version handshake message. /// - /// Intended to be sent by the worker when the slot connection opens so the - /// orchestrator can detect version mismatches and adjust behavior. Currently - /// nothing sends this; it is scaffolding for future protocol evolution. + /// Sent by the worker when the slot connection opens so the orchestrator can + /// detect version mismatches before processing prediction output. ProtocolVersion { version: u32, }, @@ -327,6 +351,16 @@ pub enum SlotResponse { /// Explicit MIME type from the predictor. Falls back to mime_guess when None. #[serde(skip_serializing_if = "Option::is_none")] mime_type: Option, + /// True if Coglet owns this file and may delete it after consumption. + #[serde(default)] + managed: bool, + /// Original basename to use at the upload endpoint. The storage filename + /// may be an opaque managed name. + #[serde(default, skip_serializing_if = "Option::is_none")] + upload_filename: Option, + /// Position of this value in the prediction's output sequence. + #[serde(default, skip_serializing_if = "is_zero")] + index: u64, }, /// Streaming output chunk for generator and iterator output. @@ -559,7 +593,7 @@ mod tests { serde_json::to_value(resp).unwrap(), json!({ "type": "protocol_version", - "version": 1 + "version": SLOT_RESPONSE_PROTOCOL_VERSION }) ); } @@ -592,6 +626,72 @@ mod tests { insta::assert_json_snapshot!(resp); } + #[test] + fn slot_file_output_managed_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/tmp/coglet/predictions/pred_123/outputs/0.png".to_string(), + kind: FileOutputKind::FileType, + mime_type: Some("image/png".to_string()), + managed: true, + upload_filename: Some("image.png".to_string()), + index: 0, + }; + insta::assert_json_snapshot!(resp); + } + + #[test] + fn slot_file_output_unmanaged_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/home/user/model/output.wav".to_string(), + kind: FileOutputKind::FileType, + mime_type: None, + managed: false, + upload_filename: None, + index: 0, + }; + insta::assert_json_snapshot!(resp); + } + + /// A worker built before `managed` existed omits the field. Defaulting to + /// false is what keeps the parent from deleting a file it does not own. + #[test] + fn slot_file_output_without_managed_deserializes_as_unmanaged() { + let json = json!({ + "type": "file_output", + "filename": "/home/user/model/output.wav", + "kind": { "type": "file_type" } + }); + + let resp: SlotResponse = serde_json::from_value(json).unwrap(); + + match resp { + SlotResponse::FileOutput { + managed, + upload_filename, + index, + .. + } => { + assert!(!managed); + assert_eq!(upload_filename, None); + assert_eq!(index, 0); + } + resp => panic!("expected a FileOutput, got {resp:?}"), + } + } + + #[test] + fn slot_file_output_oversized_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/tmp/coglet/predictions/pred_123/outputs/spill_abc.json".to_string(), + kind: FileOutputKind::Oversized, + mime_type: None, + managed: true, + upload_filename: None, + index: 3, + }; + insta::assert_json_snapshot!(resp); + } + #[test] fn slot_metric_replace_serializes() { let resp = SlotResponse::Metric { diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap new file mode 100644 index 0000000000..946de46186 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap @@ -0,0 +1,14 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/tmp/coglet/predictions/pred_123/outputs/0.png", + "kind": { + "type": "file_type" + }, + "mime_type": "image/png", + "managed": true, + "upload_filename": "image.png" +} diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap new file mode 100644 index 0000000000..f6347e32e9 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap @@ -0,0 +1,13 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/tmp/coglet/predictions/pred_123/outputs/spill_abc.json", + "kind": { + "type": "oversized" + }, + "managed": true, + "index": 3 +} diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap new file mode 100644 index 0000000000..248622d623 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap @@ -0,0 +1,12 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/home/user/model/output.wav", + "kind": { + "type": "file_type" + }, + "managed": false +} diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index b49ba1fc5e..f655027cfa 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -7,7 +7,7 @@ //! 4. Run event loop routing responses to predictions //! 5. On worker crash: fail all predictions, shut down -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap}; use std::process::Stdio; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -23,14 +23,12 @@ use crate::PredictionOutput; use crate::bridge::codec::JsonCodec; use crate::bridge::protocol::{ ControlRequest, ControlResponse, FileOutputKind, HealthcheckStatus, SlotId, SlotRequest, - SlotResponse, + SlotResponse, StructuredOutput, }; use crate::bridge::transport::create_transport; use crate::permit::{InactiveSlotIdleToken, PermitPool, SlotIdleToken}; use crate::prediction::Prediction; -const MAX_PENDING_CANCELLATIONS: usize = 1000; - /// Upload a file to a signed endpoint, returning the final URL. /// /// Matches the behavior of Python cog's `put_file_to_signed_endpoint`: @@ -43,10 +41,18 @@ async fn upload_file( data: &[u8], content_type: &str, ) -> Result { - let url = format!("{endpoint}{filename}"); + let mut url = reqwest::Url::parse(endpoint) + .map_err(|error| format!("invalid upload endpoint {endpoint}: {error}"))?; + { + let mut segments = url + .path_segments_mut() + .map_err(|_| format!("invalid upload endpoint {endpoint}"))?; + segments.pop_if_empty(); + segments.push(filename); + } let client = reqwest::Client::new(); let resp = client - .put(&url) + .put(url) .header("Content-Type", content_type) .body(data.to_vec()) .timeout(std::time::Duration::from_secs(25)) @@ -76,12 +82,99 @@ async fn upload_file( } } -fn ensure_trailing_slash(s: &str) -> String { - if s.ends_with('/') { - s.to_string() - } else { - format!("{s}/") +/// Delete an output file Coglet owns, now that its bytes have been consumed. +/// +/// Called only once the bytes are safely somewhere else: parsed, encoded into +/// the response, or accepted by the upload endpoint. Deleting any earlier would +/// discard the only copy while it could still be needed. Files left behind by a +/// failure are removed with the prediction directory. +fn delete_managed_output(filename: &str, managed: bool) { + if !managed { + return; + } + if let Err(e) = std::fs::remove_file(filename) + && e.kind() != std::io::ErrorKind::NotFound + { + tracing::warn!(%filename, error = %e, "Failed to delete managed output file"); + } +} + +struct ResolvedOutput { + index: u64, + output: serde_json::Value, +} + +type PendingOutput = tokio::task::JoinHandle>; + +fn collect_output_task_errors( + results: Vec, tokio::task::JoinError>>, +) -> Result<(), String> { + results + .into_iter() + .map(|result| match result { + Ok(result) => result, + Err(error) => Err(format!("output task failed: {error}")), + }) + .collect::, _>>() + .map(|_| ()) +} + +async fn abort_output_tasks(tasks: Vec) { + for task in &tasks { + task.abort(); + } + for task in tasks { + let _ = task.await; + } +} + +async fn abort_all_output_tasks(pending: &mut HashMap>) { + for (_, tasks) in std::mem::take(pending) { + abort_output_tasks(tasks).await; + } +} + +async fn resolve_structured_output( + mut structured: StructuredOutput, + upload_endpoint: Option, + index: u64, +) -> Result { + for file in structured.files { + let bytes = std::fs::read(&file.filename) + .map_err(|error| format!("failed to read {}: {error}", file.filename))?; + let mime = file.mime_type.unwrap_or_else(|| { + mime_guess::from_path(&file.upload_filename) + .first_or_octet_stream() + .to_string() + }); + let output = if let Some(endpoint) = &upload_endpoint { + serde_json::Value::String( + upload_file(endpoint, &file.upload_filename, &bytes, &mime).await?, + ) + } else { + use base64::Engine; + let encoded = base64::engine::general_purpose::STANDARD.encode(bytes); + serde_json::Value::String(format!("data:{mime};base64,{encoded}")) + }; + + if file.pointer.is_empty() { + structured.output = output; + } else { + let target = structured + .output + .pointer_mut(&file.pointer) + .ok_or_else(|| { + format!("structured output pointer does not exist: {}", file.pointer) + })?; + *target = output; + } + delete_managed_output(&file.filename, file.managed); } + + Ok(ResolvedOutput { + index, + output: structured.output, + }) } /// Try to lock a prediction mutex. @@ -103,6 +196,53 @@ fn try_lock_prediction( } } +struct OrderedOutputPublisher { + prediction: Arc>, + next_index: u64, + pending: BTreeMap, +} + +impl OrderedOutputPublisher { + fn new(prediction: Arc>) -> Self { + Self { + prediction, + next_index: 0, + pending: BTreeMap::new(), + } + } + + fn publish(&mut self, output: serde_json::Value, index: u64) -> Result<(), String> { + if index < self.next_index || self.pending.insert(index, output).is_some() { + return Err(format!("duplicate output index {index}")); + } + if !self.pending.contains_key(&self.next_index) { + return Ok(()); + } + + let Some(mut prediction) = try_lock_prediction(&self.prediction) else { + return Err("prediction mutex poisoned".to_string()); + }; + while let Some(output) = self.pending.remove(&self.next_index) { + prediction.append_output_chunk(output, self.next_index); + self.next_index += 1; + } + Ok(()) + } +} + +type OutputPublisher = Arc>; + +fn publish_output( + publisher: &OutputPublisher, + output: serde_json::Value, + index: u64, +) -> Result<(), String> { + publisher + .lock() + .map_err(|_| "output publisher mutex poisoned".to_string())? + .publish(output, index) +} + /// Wrap collected output items into the correct `PredictionOutput` variant. /// /// Priority: @@ -247,11 +387,14 @@ pub trait Orchestrator: Send + Sync { idle_sender: tokio::sync::oneshot::Sender, ); - /// Cancel a prediction by its prediction ID. + /// Cancel the exact registered prediction instance. /// - /// The orchestrator resolves the prediction ID to a slot ID and sends - /// a cancel request to the worker over the control socket. - async fn cancel_by_prediction_id(&self, prediction_id: &str) -> Result<(), OrchestratorError>; + /// Matching by identity prevents a delayed cancellation from affecting a + /// later prediction that reused the same public ID. + async fn cancel_prediction( + &self, + prediction: Arc>, + ) -> Result<(), OrchestratorError>; /// Run user-defined healthcheck if available. async fn healthcheck(&self) -> Result; @@ -368,7 +511,7 @@ pub struct OrchestratorHandle { Arc>>>, register_tx: mpsc::Sender, healthcheck_tx: mpsc::Sender>, - cancel_tx: mpsc::Sender, + cancel_tx: mpsc::Sender>>, slot_ids: Vec, } @@ -393,9 +536,12 @@ impl Orchestrator for OrchestratorHandle { let _ = ack_rx.await; } - async fn cancel_by_prediction_id(&self, prediction_id: &str) -> Result<(), OrchestratorError> { + async fn cancel_prediction( + &self, + prediction: Arc>, + ) -> Result<(), OrchestratorError> { self.cancel_tx - .send(prediction_id.to_string()) + .send(prediction) .await .map_err(|_| OrchestratorError::Protocol("cancel channel closed".to_string())) } @@ -700,18 +846,6 @@ pub async fn spawn_worker( }) } -fn record_pending_cancellation(pending_cancellations: &mut HashSet, prediction_id: String) { - if pending_cancellations.len() >= MAX_PENDING_CANCELLATIONS { - tracing::warn!( - prediction_id = %prediction_id, - cap = MAX_PENDING_CANCELLATIONS, - "Dropping pending cancellation because the pending cancellation buffer is full" - ); - return; - } - pending_cancellations.insert(prediction_id); -} - #[allow(clippy::too_many_arguments)] async fn run_event_loop( mut ctrl_reader: FramedRead>, @@ -724,7 +858,7 @@ async fn run_event_loop( )>, mut register_rx: mpsc::Receiver, mut healthcheck_rx: mpsc::Receiver>, - mut cancel_rx: mpsc::Receiver, + mut cancel_rx: mpsc::Receiver>>, pool: Arc, upload_url: Option, // Schema says Output is "type": "array" — always wrap as Stream. @@ -733,12 +867,15 @@ async fn run_event_loop( output_is_array: bool, ) { let mut predictions: HashMap>> = HashMap::new(); + let mut output_publishers: HashMap = HashMap::new(); let mut idle_senders: HashMap> = HashMap::new(); let mut pending_healthchecks: Vec> = Vec::new(); let mut healthcheck_counter: u64 = 0; - let mut pending_uploads: HashMap>> = HashMap::new(); - let mut pending_cancellations: HashSet = HashSet::new(); + // Upload tasks report failure so `Done` can fail the prediction instead of + // reporting success with a missing output. + let mut pending_outputs: HashMap> = HashMap::new(); + let mut pending_output_errors: HashMap = HashMap::new(); let (slot_msg_tx, mut slot_msg_rx) = mpsc::channel::<(SlotId, Result)>(100); @@ -795,6 +932,11 @@ async fn run_event_loop( Some(Ok(ControlResponse::Failed { slot, error })) => { tracing::warn!(%slot, %error, "Slot poisoned"); pool.poison(slot); + if let Some(tasks) = pending_outputs.remove(&slot) { + abort_output_tasks(tasks).await; + } + pending_output_errors.remove(&slot); + output_publishers.remove(&slot); if let Some(pred) = predictions.remove(&slot) && let Some(mut p) = try_lock_prediction(&pred) && !p.is_terminal() @@ -804,6 +946,9 @@ async fn run_event_loop( } Some(Ok(ControlResponse::Fatal { reason })) => { tracing::error!(%reason, "Worker fatal"); + abort_all_output_tasks(&mut pending_outputs).await; + pending_output_errors.clear(); + output_publishers.clear(); for (slot, pred) in predictions.drain() { tracing::warn!(%slot, "Failing prediction due to worker fatal error"); pool.poison(slot); @@ -873,6 +1018,9 @@ async fn run_event_loop( } None => { tracing::warn!("Control channel closed (worker crashed?)"); + abort_all_output_tasks(&mut pending_outputs).await; + pending_output_errors.clear(); + output_publishers.clear(); for (slot, pred) in predictions.drain() { tracing::warn!(%slot, "Failing prediction due to worker crash"); if let Some(mut p) = try_lock_prediction(&pred) { @@ -915,13 +1063,13 @@ async fn run_event_loop( } } - Some(prediction_id) = cancel_rx.recv() => { - // Resolve prediction_id → slot_id by iterating (fine for small concurrency) - let slot = predictions.iter().find_map(|(sid, pred)| { - try_lock_prediction(pred) - .filter(|p| p.id() == prediction_id) - .map(|_| *sid) - }); + Some(prediction) = cancel_rx.recv() => { + let prediction_id = try_lock_prediction(&prediction) + .map(|prediction| prediction.id().to_string()) + .unwrap_or_else(|| "unknown".to_string()); + let slot = predictions + .iter() + .find_map(|(slot, active)| Arc::ptr_eq(active, &prediction).then_some(*slot)); match slot { Some(slot_id) => { tracing::info!( @@ -938,14 +1086,16 @@ async fn run_event_loop( "Failed to send cancel request to worker" ); } - // Also abort any pending upload tasks for this slot - if let Some(handles) = pending_uploads.remove(&slot_id) { - for h in handles { h.abort(); } + // Keep the handles so the terminal response can join them + // before the prediction ID and directory are released. + if let Some(tasks) = pending_outputs.get(&slot_id) { + for task in tasks { + task.abort(); + } } } None => { - tracing::debug!(%prediction_id, "Cancel requested for unknown prediction; storing pending cancellation"); - record_pending_cancellation(&mut pending_cancellations, prediction_id); + tracing::debug!(%prediction_id, "Cancel requested for inactive prediction"); } } } @@ -971,25 +1121,14 @@ async fn run_event_loop( "Starting prediction" ); tracing::debug!(%slot_id, %prediction_id, "Registered prediction"); + output_publishers.insert( + slot_id, + Arc::new(StdMutex::new(OrderedOutputPublisher::new(Arc::clone( + &prediction, + )))), + ); predictions.insert(slot_id, prediction); - let pending_cancel = pending_cancellations.remove(&prediction_id); let _ = registered_ack.send(()); - if pending_cancel { - tracing::info!( - target: "coglet::prediction", - %prediction_id, - %slot_id, - "Applying pending cancellation" - ); - let mut writer = ctrl_writer.lock().await; - if let Err(e) = writer.send(ControlRequest::Cancel { slot: slot_id }).await { - tracing::error!( - %slot_id, - error = %e, - "Failed to send pending cancel request to worker" - ); - } - } } Some((slot_id, result)) = slot_msg_rx.recv() => { @@ -1057,27 +1196,36 @@ async fn run_event_loop( } } Ok(SlotResponse::OutputChunk { output, index }) => { - let poisoned = if let Some(pred) = predictions.get(&slot_id) { - if let Some(mut p) = try_lock_prediction(pred) { - p.append_output_chunk(output, index); - false - } else { - true - } - } else { - false - }; - // Remove poisoned predictions outside the borrow - if poisoned { - predictions.remove(&slot_id); + if let Some(publisher) = output_publishers.get(&slot_id) + && let Err(error) = publish_output(publisher, output, index) + { + pending_output_errors.entry(slot_id).or_insert(error); } } - Ok(SlotResponse::FileOutput { filename, kind, mime_type }) => { + Ok(SlotResponse::FileOutput { + filename, + kind, + mime_type, + managed, + upload_filename, + index, + }) => { + if !predictions.contains_key(&slot_id) { + tracing::warn!( + %slot_id, + %filename, + "Ignoring FileOutput for an inactive slot" + ); + continue; + } tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); let bytes = match std::fs::read(&filename) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); + pending_output_errors + .entry(slot_id) + .or_insert_with(|| format!("Failed to read output file: {e}")); continue; } }; @@ -1087,53 +1235,60 @@ async fn run_event_loop( Ok(val) => val, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to parse oversized JSON"); + pending_output_errors.entry(slot_id).or_insert_with(|| { + format!("Failed to parse oversized output: {e}") + }); continue; } }; - let poisoned = if let Some(pred) = predictions.get(&slot_id) { - if let Some(mut p) = try_lock_prediction(pred) { - p.append_output(output); - false - } else { - true - } - } else { - false - }; - if poisoned { - predictions.remove(&slot_id); + delete_managed_output(&filename, managed); + if let Some(publisher) = output_publishers.get(&slot_id) + && let Err(error) = publish_output(publisher, output, index) + { + pending_output_errors.entry(slot_id).or_insert(error); } } FileOutputKind::FileType => { + let upload_filename = upload_filename.unwrap_or_else(|| { + std::path::Path::new(&filename) + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("output") + .to_string() + }); let mime = mime_type.unwrap_or_else(|| { - mime_guess::from_path(&filename) + mime_guess::from_path(&upload_filename) .first_or_octet_stream() .to_string() }); if let Some(ref url) = upload_url { // Spawn upload task so we don't block the event loop - let pred = predictions.get(&slot_id).cloned(); - let endpoint = ensure_trailing_slash(url); - let basename = std::path::Path::new(&filename) - .file_name() - .and_then(|n| n.to_str()) - .unwrap_or("output") - .to_string(); + let endpoint = url.clone(); + let publisher = Arc::clone( + output_publishers + .get(&slot_id) + .expect("active slot checked above"), + ); let handle = tokio::spawn(async move { - match upload_file(&endpoint, &basename, &bytes, &mime).await { - Ok(url) => { - if let Some(pred) = pred - && let Some(mut p) = try_lock_prediction(&pred) - { - p.append_output(serde_json::Value::String(url)); - } - } - Err(e) => { - tracing::error!(error = %e, "Failed to upload file output"); - } - } + let url = upload_file( + &endpoint, + &upload_filename, + &bytes, + &mime, + ) + .await + .inspect_err(|e| { + tracing::error!(%filename, error = %e, "Failed to upload file output"); + })?; + // Only safe once the endpoint has the bytes. + delete_managed_output(&filename, managed); + publish_output( + &publisher, + serde_json::Value::String(url), + index, + ) }); - pending_uploads.entry(slot_id).or_default().push(handle); + pending_outputs.entry(slot_id).or_default().push(handle); } else { // No upload URL — base64-encode as data URI use base64::Engine; @@ -1142,87 +1297,172 @@ async fn run_event_loop( let output = serde_json::Value::String(format!( "data:{mime};base64,{encoded}" )); - let poisoned = if let Some(pred) = predictions.get(&slot_id) { - if let Some(mut p) = try_lock_prediction(pred) { - p.append_output(output); - false - } else { - true - } - } else { - false - }; - if poisoned { - predictions.remove(&slot_id); + delete_managed_output(&filename, managed); + if let Some(publisher) = output_publishers.get(&slot_id) + && let Err(error) = publish_output(publisher, output, index) + { + pending_output_errors.entry(slot_id).or_insert(error); } } } + FileOutputKind::Structured => { + let structured: StructuredOutput = + match serde_json::from_slice(&bytes) { + Ok(structured) => structured, + Err(error) => { + tracing::error!( + %slot_id, + %filename, + %error, + "Failed to parse structured output" + ); + pending_output_errors.entry(slot_id).or_insert_with( + || format!("Failed to parse structured output: {error}"), + ); + continue; + } + }; + delete_managed_output(&filename, managed); + let endpoint = upload_url.clone(); + let publisher = Arc::clone( + output_publishers + .get(&slot_id) + .expect("active slot checked above"), + ); + let handle = tokio::spawn(async move { + let resolved = + resolve_structured_output(structured, endpoint, index).await?; + publish_output( + &publisher, + resolved.output, + resolved.index, + ) + }); + pending_outputs.entry(slot_id).or_default().push(handle); + } } } Ok(SlotResponse::Done { id, output: _, predict_time, is_stream }) => { + // Not "succeeded" yet: pending uploads can still fail it. tracing::info!( target: "coglet::prediction", prediction_id = %id, predict_time, is_stream, output_is_array, - "Prediction succeeded" + "Run completed" ); - let uploads = pending_uploads.remove(&slot_id).unwrap_or_default(); + let pending = pending_outputs.remove(&slot_id).unwrap_or_default(); + let output_error = pending_output_errors.remove(&slot_id); + output_publishers.remove(&slot_id); if let Some(pred) = predictions.remove(&slot_id) { - if uploads.is_empty() { + if pending.is_empty() { // No pending uploads — complete synchronously to avoid // a race between tokio::spawn and Notify::notified() in // service.rs. notify_waiters() only wakes already- // registered waiters; spawning a task can fire the // notification before the service registers its waiter. if let Some(mut p) = try_lock_prediction(&pred) { - let pred_output = wrap_outputs( - p.take_outputs(), - output_is_array, - is_stream, - ); - p.set_succeeded(pred_output); + if p.cancel_token().is_cancelled() { + p.set_canceled(); + } else if let Some(error) = output_error { + p.set_failed(error); + } else { + let pred_output = wrap_outputs( + p.take_outputs(), + output_is_array, + is_stream, + ); + p.set_succeeded(pred_output); + } } } else { - // Has pending uploads — must spawn to await them. - // Clone the cancel token so we can abort uploads if - // the prediction is cancelled while uploads are in flight. - let (cancel_token, upload_pred_id) = match try_lock_prediction(&pred) { + // Pending output tasks must finish or be joined after + // abortion before this prediction becomes terminal. + let (cancel_token, output_pred_id) = match try_lock_prediction(&pred) { Some(p) => (Some(p.cancel_token()), p.id().to_string()), None => (None, id.clone()), }; tokio::spawn(async move { - if let Some(token) = cancel_token { - let upload_fut = futures::future::join_all(uploads); - tokio::pin!(upload_fut); + if let Some(error) = output_error { + abort_output_tasks(pending).await; + if let Some(mut p) = try_lock_prediction(&pred) { + if p.cancel_token().is_cancelled() { + p.set_canceled(); + } else { + p.set_failed(error); + } + } + return; + } + + if cancel_token + .as_ref() + .is_some_and(|token| token.is_cancelled()) + { + abort_output_tasks(pending).await; + if let Some(mut p) = try_lock_prediction(&pred) { + p.set_canceled(); + } + return; + } + + let abort_handles = pending + .iter() + .map(tokio::task::JoinHandle::abort_handle) + .collect::>(); + let output_fut = futures::future::join_all(pending); + tokio::pin!(output_fut); + let results = if let Some(token) = cancel_token.as_ref() { tokio::select! { - _ = &mut upload_fut => {} + results = &mut output_fut => Some(results), _ = token.cancelled() => { tracing::info!( target: "coglet::prediction", - prediction_id = %upload_pred_id, - "Aborting in-flight uploads due to cancellation" + prediction_id = %output_pred_id, + "Aborting in-flight output tasks due to cancellation" ); - if let Some(mut p) = try_lock_prediction(&pred) { - p.set_canceled(); + for handle in abort_handles { + handle.abort(); } - return; + let _ = (&mut output_fut).await; + None } } } else { - for h in uploads { - let _ = h.await; + Some((&mut output_fut).await) + }; + + let Some(results) = results else { + if let Some(mut p) = try_lock_prediction(&pred) { + p.set_canceled(); } + return; + }; + + let Some(mut p) = try_lock_prediction(&pred) else { + return; + }; + if p.cancel_token().is_cancelled() { + p.set_canceled(); + return; } - if let Some(mut p) = try_lock_prediction(&pred) { - let pred_output = wrap_outputs( - p.take_outputs(), - output_is_array, - is_stream, + if let Err(error) = collect_output_task_errors(results) { + tracing::error!( + target: "coglet::prediction", + prediction_id = %output_pred_id, + %error, + "Failing prediction because an output task failed" ); - p.set_succeeded(pred_output); + p.set_failed(format!("Failed to process output file: {error}")); + return; } + let pred_output = wrap_outputs( + p.take_outputs(), + output_is_array, + is_stream, + ); + p.set_succeeded(pred_output); }); } } else { @@ -1236,10 +1476,11 @@ async fn run_event_loop( %error, "Prediction failed" ); - // Abort any pending uploads — prediction is terminal - if let Some(handles) = pending_uploads.remove(&slot_id) { - for h in handles { h.abort(); } + if let Some(tasks) = pending_outputs.remove(&slot_id) { + abort_output_tasks(tasks).await; } + pending_output_errors.remove(&slot_id); + output_publishers.remove(&slot_id); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1252,10 +1493,11 @@ async fn run_event_loop( prediction_id = %id, "Prediction cancelled" ); - // Abort any pending uploads — prediction is terminal - if let Some(handles) = pending_uploads.remove(&slot_id) { - for h in handles { h.abort(); } + if let Some(tasks) = pending_outputs.remove(&slot_id) { + abort_output_tasks(tasks).await; } + pending_output_errors.remove(&slot_id); + output_publishers.remove(&slot_id); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1264,9 +1506,11 @@ async fn run_event_loop( } Err(e) => { tracing::error!(%slot_id, error = %e, "Slot socket error"); - if let Some(handles) = pending_uploads.remove(&slot_id) { - for h in handles { h.abort(); } + if let Some(tasks) = pending_outputs.remove(&slot_id) { + abort_output_tasks(tasks).await; } + pending_output_errors.remove(&slot_id); + output_publishers.remove(&slot_id); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1278,6 +1522,7 @@ async fn run_event_loop( } } + abort_all_output_tasks(&mut pending_outputs).await; tracing::info!("Event loop exiting"); } @@ -1285,6 +1530,8 @@ async fn run_event_loop( mod tests { use super::*; use serde_json::json; + use wiremock::matchers::{method, path, query_param}; + use wiremock::{Mock, MockServer, ResponseTemplate}; // ── wrap_outputs: schema says array (output_is_array = true) ── @@ -1297,16 +1544,151 @@ mod tests { } #[test] - fn record_pending_cancellation_caps_stored_ids() { - let mut pending = HashSet::new(); - for index in 0..MAX_PENDING_CANCELLATIONS { - record_pending_cancellation(&mut pending, format!("pred-{index}")); + fn delete_managed_output_deletes_managed_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("managed.txt"); + std::fs::write(&path, b"managed").unwrap(); + + delete_managed_output(path.to_str().unwrap(), true); + + assert!(!path.exists(), "managed file should be deleted"); + } + + #[test] + fn delete_managed_output_preserves_unmanaged_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("unmanaged.txt"); + std::fs::write(&path, b"unmanaged").unwrap(); + + delete_managed_output(path.to_str().unwrap(), false); + + assert!(path.exists(), "unmanaged file must be left alone"); + assert_eq!(std::fs::read(&path).unwrap(), b"unmanaged"); + } + + /// Run the upload-then-delete sequence used by the FileOutput handler. + async fn upload_then_delete_managed( + endpoint: &str, + path: &std::path::Path, + ) -> Result<(), String> { + crate::install_crypto_provider(); + let filename = path.to_str().unwrap().to_string(); + let endpoint = endpoint.to_string(); + let handle: tokio::task::JoinHandle> = tokio::spawn(async move { + upload_file(&endpoint, "0.txt", b"payload", "text/plain").await?; + delete_managed_output(&filename, true); + Ok(()) + }); + handle + .await + .map_err(|error| format!("upload task failed: {error}"))? + } + + /// Deleting has to wait for the upload: on failure the bytes on disk are the + /// only copy left, and the prediction directory cleanup is what reclaims them. + #[tokio::test] + async fn managed_output_survives_a_failed_upload() { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("0.txt"); + std::fs::write(&path, b"payload").unwrap(); + + let error = upload_then_delete_managed(&server.uri(), &path) + .await + .expect_err("a 500 response should fail the upload"); + + assert!(error.contains("status 500"), "error: {error}"); + assert!(path.exists(), "file must survive for directory cleanup"); + } + + #[tokio::test] + async fn managed_output_is_deleted_after_a_successful_upload() { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("0.txt"); + std::fs::write(&path, b"payload").unwrap(); + + upload_then_delete_managed(&server.uri(), &path) + .await + .unwrap(); + + assert!(!path.exists(), "uploaded file should be deleted"); + } + + #[tokio::test] + async fn upload_file_escapes_filename_as_a_path_segment() { + crate::install_crypto_provider(); + let server = MockServer::start().await; + Mock::given(method("PUT")) + .and(path("/upload/result%3Fdraft.png")) + .and(query_param("sig", "abc")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + + let endpoint = format!("{}/upload?sig=abc", server.uri()); + let url = upload_file(&endpoint, "result?draft.png", b"payload", "image/png") + .await + .unwrap(); + + assert!(url.ends_with("/result%3Fdraft.png"), "upload URL: {url}"); + } + + #[test] + fn output_publisher_emits_completed_uploads_in_order() { + let prediction = Arc::new(StdMutex::new(Prediction::new( + "streaming-output".to_string(), + None, + ))); + let mut stream = prediction.lock().unwrap().subscribe_stream(); + let publisher = Arc::new(StdMutex::new(OrderedOutputPublisher::new(prediction))); + + publish_output(&publisher, serde_json::json!("second"), 1).unwrap(); + assert!(stream.try_recv().is_err()); + + publish_output(&publisher, serde_json::json!("first"), 0).unwrap(); + let first = stream.try_recv().unwrap(); + let second = stream.try_recv().unwrap(); + assert_eq!( + first.json_data(), + serde_json::json!({"chunk": "first", "index": 0}) + ); + assert_eq!( + second.json_data(), + serde_json::json!({"chunk": "second", "index": 1}) + ); + } + + #[tokio::test] + async fn abort_output_tasks_joins_aborted_tasks() { + use std::sync::atomic::{AtomicBool, Ordering}; + + struct DropMarker(Arc); + impl Drop for DropMarker { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } } - record_pending_cancellation(&mut pending, "overflow".to_string()); + let dropped = Arc::new(AtomicBool::new(false)); + let marker = DropMarker(Arc::clone(&dropped)); + let task: PendingOutput = tokio::spawn(async move { + let _marker = marker; + std::future::pending().await + }); + + abort_output_tasks(vec![task]).await; - assert_eq!(pending.len(), MAX_PENDING_CANCELLATIONS); - assert!(!pending.contains("overflow")); + assert!(dropped.load(Ordering::SeqCst)); } #[test] diff --git a/crates/coglet/src/permit/slot.rs b/crates/coglet/src/permit/slot.rs index 95b980fb6b..cb77b25534 100644 --- a/crates/coglet/src/permit/slot.rs +++ b/crates/coglet/src/permit/slot.rs @@ -47,12 +47,25 @@ impl UnregisteredPredictionSlot { pub fn prediction(&self) -> Arc> { self.prediction_slot.prediction() } + + pub fn id(&self) -> &str { + self.prediction_slot.id() + } + + /// Fail the prediction and hand the slot back to the pool. + /// + /// For giving up before the request reaches the worker. See + /// [`PredictionSlot::release_unused`]. + pub fn release_unused(self, error: String) { + self.prediction_slot.release_unused(error); + } } /// Holds a prediction and its permit side-by-side. /// /// On drop: Permit returns to pool (if idle and not poisoned at pool level). pub struct PredictionSlot { + id: String, prediction: Arc>, slot_id: SlotId, permit: Option, @@ -66,7 +79,9 @@ impl PredictionSlot { idle_rx: tokio::sync::oneshot::Receiver, ) -> Self { let slot_id = permit.slot_id(); + let id = prediction.id().to_string(); Self { + id, prediction: Arc::new(Mutex::new(prediction)), slot_id, permit: Some(AnyPermit::InUse(permit)), @@ -135,15 +150,32 @@ impl PredictionSlot { } } + /// Fail the prediction and return the permit to the pool without waiting for + /// the worker to confirm the slot is idle. + /// + /// For failures before the request is sent. The worker never saw the slot as + /// busy, so no idle notification is coming and `into_idle` would wait + /// forever. Dropping the slot instead would strand the permit and cost a + /// slot of capacity for the process's lifetime. + pub fn release_unused(mut self, error: String) { + if let Ok(mut prediction) = self.prediction.lock() + && !prediction.is_terminal() + { + prediction.set_failed(error); + } + + self.idle_rx = None; + if let Some(AnyPermit::InUse(permit)) = self.permit.take() { + self.permit = Some(AnyPermit::Idle(permit.into_idle())); + } + } + pub fn is_idle(&self) -> bool { self.permit.as_ref().is_some_and(|p| p.is_idle()) } - pub fn id(&self) -> String { - self.prediction - .try_lock() - .map(|p| p.id().to_string()) - .unwrap_or_default() + pub fn id(&self) -> &str { + &self.id } } @@ -187,6 +219,24 @@ mod tests { let (_idle_tx, idle_rx) = tokio::sync::oneshot::channel(); let slot = PredictionSlot::new(prediction, permit, idle_rx); assert_eq!(slot.slot_id(), slot_id); + assert_eq!(slot.id(), "test_123"); + } + + #[tokio::test] + async fn slot_id_does_not_depend_on_prediction_lock() { + let pool = PermitPool::new(1); + let (a, _b) = UnixStream::pair().unwrap(); + let (_, write) = a.into_split(); + pool.add_permit(SlotId::new(), FramedWrite::new(write, JsonCodec::new())); + + let permit = pool.try_acquire().unwrap(); + let prediction = Prediction::new("locked-prediction".to_string(), None); + let (_idle_tx, idle_rx) = tokio::sync::oneshot::channel(); + let slot = PredictionSlot::new(prediction, permit, idle_rx); + let prediction = slot.prediction(); + let _guard = prediction.lock().unwrap(); + + assert_eq!(slot.id(), "locked-prediction"); } #[tokio::test] @@ -213,6 +263,26 @@ mod tests { assert!(pool.try_acquire().is_some()); } + #[tokio::test] + async fn release_unused_returns_permit() { + let pool = PermitPool::new(1); + + let (a, _b) = UnixStream::pair().unwrap(); + let (_, write) = a.into_split(); + let slot_id = SlotId::new(); + + pool.add_permit(slot_id, FramedWrite::new(write, JsonCodec::new())); + + let permit = pool.try_acquire().unwrap(); + let prediction = Prediction::new("test_123".to_string(), None); + let (_idle_tx, idle_rx) = tokio::sync::oneshot::channel(); + let slot = PredictionSlot::new(prediction, permit, idle_rx); + + slot.release_unused("request was not sent".to_string()); + + assert!(pool.try_acquire().is_some()); + } + #[tokio::test] async fn slot_not_idle_orphans_permit() { let pool = PermitPool::new(1); diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index 2b52e91129..61cb7e05a0 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -1,6 +1,6 @@ //! Prediction state tracking. -use std::collections::{HashMap, VecDeque}; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::Arc; use std::time::Instant; @@ -142,7 +142,7 @@ pub struct Prediction { started_at: Instant, status: PredictionStatus, logs: String, - outputs: Vec, + outputs: BTreeMap, output: Option, error: Option, webhook: Option, @@ -166,7 +166,7 @@ impl Prediction { started_at: Instant::now(), status: PredictionStatus::Starting, logs: String::new(), - outputs: Vec::new(), + outputs: BTreeMap::new(), output: None, error: None, webhook, @@ -466,12 +466,15 @@ impl Prediction { } pub fn append_output(&mut self, output: serde_json::Value) { - let index = self.outputs.len() as u64; + let index = self + .outputs + .last_key_value() + .map_or(0, |(index, _)| index + 1); self.append_output_chunk(output, index); } pub fn append_output_chunk(&mut self, output: serde_json::Value, index: u64) { - self.outputs.push(output.clone()); + self.outputs.insert(index, output.clone()); self.emit_stream_event(PredictionStreamEvent::Output { chunk: output, index, @@ -479,12 +482,12 @@ impl Prediction { self.fire_webhook(WebhookEventType::Output); } - pub fn outputs(&self) -> &[serde_json::Value] { - &self.outputs + pub fn outputs(&self) -> Vec<&serde_json::Value> { + self.outputs.values().collect() } pub fn take_outputs(&mut self) -> Vec { - std::mem::take(&mut self.outputs) + std::mem::take(&mut self.outputs).into_values().collect() } pub fn output(&self) -> Option<&PredictionOutput> { @@ -554,7 +557,7 @@ impl Prediction { if let Some(ref output) = self.output { payload["output"] = serde_json::json!(output); } else if !self.outputs.is_empty() { - payload["output"] = serde_json::json!(self.outputs); + payload["output"] = serde_json::json!(self.outputs.values().collect::>()); } if let Some(ref error) = self.error { @@ -692,6 +695,18 @@ mod tests { assert_eq!(pred.outputs().len(), 2); } + #[test] + fn output_values_are_returned_in_protocol_index_order() { + let mut pred = Prediction::new("test".to_string(), None); + pred.append_output_chunk(serde_json::json!("second"), 1); + pred.append_output_chunk(serde_json::json!("first"), 0); + + assert_eq!( + pred.take_outputs(), + vec![serde_json::json!("first"), serde_json::json!("second")] + ); + } + #[tokio::test] async fn prediction_stream_emits_start_output_log_and_completed() { let mut prediction = Prediction::new("pred_stream".to_string(), None); diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 460a09e09a..c553103834 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -51,10 +51,21 @@ pub enum CreatePredictionError { NotReady, #[error("At capacity (no slots available)")] AtCapacity, + #[error("A prediction with this ID already exists")] + DuplicateId(ExistingPrediction), + #[error("Invalid prediction ID: {0}")] + InvalidId(&'static str), } const MAX_STREAM_SUBSCRIBERS: usize = STREAM_CHANNEL_CAPACITY; +/// Per-prediction input and output directories live here, one per prediction ID. +const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; + +/// Prediction directories are moved here before deletion. See +/// [`detach_prediction_dir`]. +const CLEANUP_ROOT: &str = "/tmp/coglet/cleanup"; + #[derive(Debug, thiserror::Error)] pub enum SubscribePredictionStreamError { #[error("Prediction not found")] @@ -97,9 +108,70 @@ struct PredictionEntry { cancel_on_stream_drop: bool, } +pub struct ExistingPrediction { + id: String, + prediction: Arc>, + input: serde_json::Value, + cancel_on_stream_drop: bool, +} + +impl std::fmt::Debug for ExistingPrediction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExistingPrediction") + .field("id", &self.id) + .finish_non_exhaustive() + } +} + +impl ExistingPrediction { + fn from_entry(id: &str, entry: &PredictionEntry) -> Self { + Self { + id: id.to_string(), + prediction: Arc::clone(&entry.prediction), + input: entry.input.clone(), + cancel_on_stream_drop: entry.cancel_on_stream_drop, + } + } + + pub fn response(&self) -> Option { + let pred = self.prediction.lock().ok()?; + let mut response = pred.build_state_snapshot(); + response["input"] = self.input.clone(); + Some(response) + } + + pub fn subscribe( + &self, + service: &Arc, + ) -> Result { + let stream = { + let Some(prediction) = try_lock_prediction(&self.prediction) else { + return Err(SubscribePredictionStreamError::Unavailable); + }; + if prediction.stream_receiver_count() >= MAX_STREAM_SUBSCRIBERS { + return Err(SubscribePredictionStreamError::TooManySubscribers); + } + prediction.subscribe_stream_replay() + }; + Ok(PredictionStreamSubscription { + id: self.id.clone(), + replay: stream.replay, + skipped: stream.skipped, + receiver: stream.receiver, + guard: PredictionStreamGuard { + id: self.id.clone(), + prediction: Arc::clone(&self.prediction), + service: Arc::clone(service), + cancel_on_stream_drop: self.cancel_on_stream_drop, + }, + }) + } +} + /// Handle to a submitted prediction for cancellation on disconnect. pub struct PredictionHandle { id: String, + prediction: Arc>, cancel_token: CancellationToken, } @@ -118,7 +190,7 @@ impl PredictionHandle { /// `service.cancel(id)` which fires the CancellationToken AND /// delegates to the orchestrator to cancel the worker subprocess. pub fn sync_guard(&self, service: Arc) -> SyncPredictionGuard { - SyncPredictionGuard::new(self.id.clone(), service) + SyncPredictionGuard::new(self.id.clone(), Arc::clone(&self.prediction), service) } } @@ -151,6 +223,7 @@ impl PredictionStreamSubscription { pub struct PredictionStreamGuard { id: String, + prediction: Arc>, service: Arc, cancel_on_stream_drop: bool, } @@ -161,14 +234,8 @@ impl Drop for PredictionStreamGuard { return; } - // Prediction cleanup may remove the service entry before the SSE response - // finishes draining. Missing entries deliberately report zero receivers and - // terminal state so this guard cannot cancel an already-cleaned prediction. - if self.service.stream_receiver_count(&self.id) == 0 - && !self.service.prediction_is_terminal(&self.id) - { - self.service.cancel(&self.id); - } + self.service + .cancel_abandoned_stream(&self.id, &self.prediction); } } @@ -180,13 +247,19 @@ impl Drop for PredictionStreamGuard { /// (Rust-side observers) and the orchestrator (worker subprocess cancel). pub struct SyncPredictionGuard { prediction_id: Option, + prediction: Arc>, service: Arc, } impl SyncPredictionGuard { - pub fn new(prediction_id: String, service: Arc) -> Self { + pub fn new( + prediction_id: String, + prediction: Arc>, + service: Arc, + ) -> Self { Self { prediction_id: Some(prediction_id), + prediction, service, } } @@ -199,7 +272,7 @@ impl SyncPredictionGuard { impl Drop for SyncPredictionGuard { fn drop(&mut self) { if let Some(ref id) = self.prediction_id { - self.service.cancel(id); + self.service.cancel_if_current(id, &self.prediction); } } } @@ -244,6 +317,8 @@ pub struct PredictionService { input_validator: RwLock>, train_validator: RwLock>, supports_prediction_streaming: RwLock, + #[cfg(test)] + idempotent_pre_submit_barrier: RwLock>>, } impl PredictionService { @@ -264,6 +339,8 @@ impl PredictionService { input_validator: RwLock::new(None), train_validator: RwLock::new(None), supports_prediction_streaming: RwLock::new(false), + #[cfg(test)] + idempotent_pre_submit_barrier: RwLock::new(None), } } @@ -322,6 +399,19 @@ impl PredictionService { *self.supports_prediction_streaming.read().await } + #[cfg(test)] + pub async fn set_idempotent_pre_submit_barrier(&self, barrier: Arc) { + *self.idempotent_pre_submit_barrier.write().await = Some(barrier); + } + + #[cfg(test)] + pub async fn wait_idempotent_pre_submit_barrier(&self) { + let barrier = self.idempotent_pre_submit_barrier.read().await.clone(); + if let Some(barrier) = barrier { + barrier.wait().await; + } + } + /// Get the permit pool from orchestrator. pub async fn pool(&self) -> Option> { if let Some(ref state) = *self.orchestrator.read().await { @@ -518,6 +608,8 @@ impl PredictionService { webhook: Option, cancel_on_stream_drop: bool, ) -> Result<(PredictionHandle, UnregisteredPredictionSlot), CreatePredictionError> { + validate_prediction_id(&id).map_err(CreatePredictionError::InvalidId)?; + let health = *self.health.read().await; if health != Health::Ready { return Err(CreatePredictionError::NotReady); @@ -527,6 +619,17 @@ impl PredictionService { let pool = self.pool().await; let pool = pool.as_ref().ok_or(CreatePredictionError::NotReady)?; + // Reserve the ID before acquiring a permit so a rejected duplicate + // cannot consume slot capacity. + let prediction_entry = match self.predictions.entry(id.clone()) { + dashmap::Entry::Vacant(entry) => entry, + dashmap::Entry::Occupied(entry) => { + return Err(CreatePredictionError::DuplicateId( + ExistingPrediction::from_entry(&id, entry.get()), + )); + } + }; + let permit = pool .try_acquire() .ok_or(CreatePredictionError::AtCapacity)?; @@ -537,18 +640,18 @@ impl PredictionService { let slot = PredictionSlot::new(prediction, permit, idle_rx); let prediction_arc = slot.prediction(); - // Register in DashMap — this is the single source of truth - self.predictions.insert( - id.clone(), - PredictionEntry { - prediction: prediction_arc, - cancel_token: cancel_token.clone(), - input, - cancel_on_stream_drop, - }, - ); + prediction_entry.insert(PredictionEntry { + prediction: Arc::clone(&prediction_arc), + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); - let handle = PredictionHandle { id, cancel_token }; + let handle = PredictionHandle { + id, + prediction: prediction_arc, + cancel_token, + }; Ok((handle, UnregisteredPredictionSlot::new(slot, idle_tx))) } @@ -563,65 +666,72 @@ impl PredictionService { /// Locks the real Prediction to read current state — no stale copies. /// Adds `input` from the PredictionEntry on top of the shared snapshot. pub fn get_prediction_response(&self, id: &str) -> Option { - let entry = self.predictions.get(id)?; - let pred = entry.prediction.lock().ok()?; - - let mut response = pred.build_state_snapshot(); - response["input"] = entry.input.clone(); + self.get_existing_prediction(id)?.response() + } - Some(response) + pub fn get_existing_prediction(&self, id: &str) -> Option { + let entry = self.predictions.get(id)?; + Some(ExistingPrediction::from_entry(id, &entry)) } pub fn subscribe_prediction_stream( self: &Arc, id: &str, ) -> Result { - let entry = self - .predictions - .get(id) - .ok_or(SubscribePredictionStreamError::NotFound)?; - let stream = { - let Some(prediction) = try_lock_prediction(&entry.prediction) else { - return Err(SubscribePredictionStreamError::Unavailable); - }; - if prediction.stream_receiver_count() >= MAX_STREAM_SUBSCRIBERS { - return Err(SubscribePredictionStreamError::TooManySubscribers); + self.get_existing_prediction(id) + .ok_or(SubscribePredictionStreamError::NotFound)? + .subscribe(self) + } + + fn cancel_abandoned_stream(&self, id: &str, expected: &Arc>) { + let Some(entry) = self.predictions.get(id) else { + return; + }; + if !Arc::ptr_eq(&entry.prediction, expected) { + return; + } + let should_cancel = entry + .prediction + .lock() + .map(|prediction| prediction.stream_receiver_count() == 0 && !prediction.is_terminal()) + .unwrap_or(false); + if !should_cancel { + return; + } + + entry.cancel_token.cancel(); + let orchestrator = match self.orchestrator.try_read() { + Ok(guard) => guard.as_ref().map(|state| Arc::clone(&state.orchestrator)), + Err(_) => { + tracing::warn!(prediction_id = %id, "Skipped worker cancel: orchestrator lock unavailable"); + None } - prediction.subscribe_stream_replay() }; - let cancel_on_stream_drop = entry.cancel_on_stream_drop; - let id = id.to_string(); - Ok(PredictionStreamSubscription { - id: id.clone(), - replay: stream.replay, - skipped: stream.skipped, - receiver: stream.receiver, - guard: PredictionStreamGuard { - id, - service: Arc::clone(self), - cancel_on_stream_drop, - }, - }) + if let Some(orchestrator) = orchestrator { + spawn_orchestrator_cancel(orchestrator, Arc::clone(expected)); + } } - fn stream_receiver_count(&self, id: &str) -> usize { - self.predictions - .get(id) - .and_then(|entry| { - entry - .prediction - .lock() - .ok() - .map(|p| p.stream_receiver_count()) - }) - .unwrap_or(0) - } + fn cancel_if_current(&self, id: &str, expected: &Arc>) -> bool { + let Some(entry) = self.predictions.get(id) else { + return false; + }; + if !Arc::ptr_eq(&entry.prediction, expected) { + return false; + } - fn prediction_is_terminal(&self, id: &str) -> bool { - self.predictions - .get(id) - .and_then(|entry| entry.prediction.lock().ok().map(|p| p.is_terminal())) - .unwrap_or(true) + entry.cancel_token.cancel(); + let orchestrator = match self.orchestrator.try_read() { + Ok(guard) => guard.as_ref().map(|state| Arc::clone(&state.orchestrator)), + Err(_) => { + tracing::warn!(prediction_id = %id, "Skipped worker cancel: orchestrator lock unavailable"); + None + } + }; + if let Some(orchestrator) = orchestrator { + spawn_orchestrator_cancel(orchestrator, Arc::clone(expected)); + } + true } /// Run a prediction to completion via orchestrator. @@ -635,8 +745,19 @@ impl PredictionService { let state = state .ok_or_else(|| PredictionError::Failed("No orchestrator configured".to_string()))?; + let prediction_id = unregistered_slot.id().to_string(); + + // Prepare everything that can fail before the slot is registered, so a + // failure here can hand the permit straight back to the pool. + let request = match prepare_slot_request(&prediction_id, input, context) { + Ok(request) => request, + Err(error) => { + unregistered_slot.release_unused(error.clone()); + return Err(PredictionError::Failed(error)); + } + }; + let (idle_tx, mut slot) = unregistered_slot.into_parts(); - let prediction_id = slot.id(); let slot_id = slot.slot_id(); { @@ -656,28 +777,6 @@ impl PredictionService { .register_prediction(slot_id, Arc::clone(&prediction_arc), idle_tx) .await; - // Create per-prediction dirs for file-based inputs/outputs - let prediction_dir = - std::path::PathBuf::from("/tmp/coglet/predictions").join(&prediction_id); - let output_dir = prediction_dir.join("outputs"); - let input_dir = prediction_dir.join("inputs"); - std::fs::create_dir_all(&output_dir) - .map_err(|e| PredictionError::Failed(format!("Failed to create output dir: {}", e)))?; - std::fs::create_dir_all(&input_dir) - .map_err(|e| PredictionError::Failed(format!("Failed to create input dir: {}", e)))?; - - let request = build_slot_request( - prediction_id.clone(), - input, - output_dir - .to_str() - .expect("output dir path is valid UTF-8") - .to_string(), - &input_dir, - context, - ) - .map_err(|e| PredictionError::Failed(format!("Failed to build slot request: {}", e)))?; - // permit_mut returns None if permit isn't InUse (shouldn't happen here) let permit = slot .permit_mut() @@ -702,7 +801,7 @@ impl PredictionService { if was_cancelled_before_send && let Err(e) = state .orchestrator - .cancel_by_prediction_id(&prediction_id) + .cancel_prediction(Arc::clone(&prediction_arc)) .await { tracing::error!( @@ -780,7 +879,7 @@ impl PredictionService { // Delegate to orchestrator to actually cancel the worker-side prediction. // This must be non-blocking since cancel() is sync, so we spawn a task. - let id_owned = id.to_string(); + let prediction = Arc::clone(&entry.prediction); let orchestrator = match self.orchestrator.try_read() { Ok(guard) => guard.as_ref().map(|s| Arc::clone(&s.orchestrator)), Err(_) => { @@ -789,7 +888,7 @@ impl PredictionService { } }; if let Some(orch) = orchestrator { - spawn_orchestrator_cancel(orch, id_owned); + spawn_orchestrator_cancel(orch, prediction); } true } else { @@ -797,9 +896,50 @@ impl PredictionService { } } - /// Remove a prediction from the DashMap after completion. + /// Release a completed prediction's ID and delete its on-disk directory. + /// + /// The directory is a backstop for output files that were not deleted + /// individually, such as those from aborted uploads or cancelled predictions. pub fn remove_prediction(&self, id: &str) { - self.predictions.remove(id); + // The HTTP layer validates too, but this decides a path to delete, so it + // does not take the caller's word for it. + if let Err(e) = validate_prediction_id(id) { + tracing::warn!(prediction_id = %id, error = e, "Refusing to remove prediction dir: invalid ID"); + return; + } + + self.remove_prediction_with(id, || detach_prediction_dir(id)); + } + + fn remove_prediction_with( + &self, + id: &str, + detach: impl FnOnce() -> std::io::Result>, + ) { + // Detaching while the ID is still reserved is what stops this cleanup + // from deleting the files of a later prediction that reuses the ID. + // Renaming is a single cheap operation, so the map shard stays locked + // only briefly; the recursive delete happens afterwards. + let detached = { + let entry = self.predictions.entry(id.to_string()); + let detached = match detach() { + Ok(detached) => detached, + Err(e) => { + // Releasing the ID here would let a later prediction reuse + // the stale directory that failed to detach. + tracing::warn!(prediction_id = %id, error = %e, "Failed to detach prediction dir for cleanup"); + return; + } + }; + if let dashmap::Entry::Occupied(occupied) = entry { + occupied.remove(); + } + detached + }; + + if let Some(dir) = detached { + delete_detached_dir(&dir); + } } pub fn trigger_shutdown(&self) { @@ -811,13 +951,137 @@ impl PredictionService { } } -fn spawn_orchestrator_cancel(orch: Arc, id: String) { +/// Longest accepted prediction ID. +/// +/// An ID becomes a directory name, and filesystems generally cap a single name at +/// 255 bytes. 128 stays well inside that while leaving room for the UUIDs and +/// slugs callers actually send, and rejecting the rest up front turns a +/// mid-prediction `ENAMETOOLONG` into a 422 the client can act on. +/// +/// Keep the error message in `validate_prediction_id` in sync with this. +const MAX_PREDICTION_ID_LEN: usize = 128; + +/// Validate a client-supplied prediction ID. +/// +/// An ID is used verbatim as a directory name under the prediction root, so it +/// has to be a single, ordinary path component. The allowlist covers the IDs Cog +/// generates and the usual UUID and slug shapes, and rules out separators, +/// traversal, control characters, and names that differ only by Unicode +/// normalization (which some filesystems fold together, letting two IDs collide +/// on one directory). +pub(crate) fn validate_prediction_id(id: &str) -> Result<(), &'static str> { + if id.is_empty() { + return Err("prediction ID must not be empty"); + } + if id == "." || id == ".." { + return Err("prediction ID must not be . or .."); + } + if !id + .chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.')) + { + return Err( + "prediction ID must contain only letters, numbers, hyphens, underscores, and dots", + ); + } + // Checked after the charset, where one character is one byte, so the limit + // the client is told about is the one being applied. + if id.len() > MAX_PREDICTION_ID_LEN { + return Err("prediction ID must be at most 128 characters"); + } + Ok(()) +} + +fn prediction_dir_for(id: &str) -> std::path::PathBuf { + std::path::PathBuf::from(PREDICTION_ROOT).join(id) +} + +/// Move a prediction directory out of the prediction root, returning its new path. +/// +/// Renaming detaches the directory in one step, so the ID is free to be reused +/// immediately and the slower recursive delete cannot touch whatever the next +/// prediction puts there. Rename acts on the final path component itself, so a +/// symlinked prediction directory is moved as a link rather than followed. +/// +/// Returns `Ok(None)` when there is nothing to clean up. +fn detach_prediction_dir(id: &str) -> std::io::Result> { + let dir = prediction_dir_for(id); + if let Err(e) = std::fs::symlink_metadata(&dir) { + return if e.kind() == std::io::ErrorKind::NotFound { + Ok(None) + } else { + Err(e) + }; + } + + std::fs::create_dir_all(CLEANUP_ROOT)?; + let detached = std::path::PathBuf::from(CLEANUP_ROOT).join(uuid::Uuid::new_v4().to_string()); + std::fs::rename(&dir, &detached)?; + Ok(Some(detached)) +} + +/// Delete a path already detached by [`detach_prediction_dir`]. +/// +/// A symlink is unlinked rather than followed, so a run that replaced its own +/// prediction directory with a link cannot get the target deleted. +fn delete_detached_dir(dir: &std::path::Path) { + let removed = match std::fs::symlink_metadata(dir) { + Ok(metadata) if metadata.is_dir() => std::fs::remove_dir_all(dir), + Ok(_) => std::fs::remove_file(dir), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return, + Err(e) => { + tracing::warn!(dir = %dir.display(), error = %e, "Failed to inspect detached prediction dir"); + return; + } + }; + if let Err(e) = removed + && e.kind() != std::io::ErrorKind::NotFound + { + tracing::warn!(dir = %dir.display(), error = %e, "Failed to delete detached prediction dir"); + } +} + +/// Create a prediction's input and output directories and build its slot request. +/// +/// Kept separate from `predict` so every failure that can happen before the +/// worker is involved is handled in one place, where the permit can still be +/// returned to the pool. +fn prepare_slot_request( + prediction_id: &str, + input: serde_json::Value, + context: std::collections::HashMap, +) -> Result { + let prediction_dir = prediction_dir_for(prediction_id); + let output_dir = prediction_dir.join("outputs"); + let input_dir = prediction_dir.join("inputs"); + std::fs::create_dir_all(&output_dir) + .map_err(|e| format!("Failed to create output dir: {e}"))?; + std::fs::create_dir_all(&input_dir).map_err(|e| format!("Failed to create input dir: {e}"))?; + + let output_dir = output_dir + .to_str() + .ok_or_else(|| "Output dir path is not valid UTF-8".to_string())?; + + build_slot_request( + prediction_id.to_string(), + input, + output_dir.to_string(), + &input_dir, + context, + ) + .map_err(|e| format!("Failed to build slot request: {e}")) +} + +fn spawn_orchestrator_cancel(orch: Arc, prediction: Arc>) { + let id = try_lock_prediction(&prediction) + .map(|prediction| prediction.id().to_string()) + .unwrap_or_else(|| "unknown".to_string()); let Ok(handle) = tokio::runtime::Handle::try_current() else { tracing::warn!(prediction_id = %id, "No tokio runtime available to cancel prediction"); return; }; handle.spawn(async move { - if let Err(e) = orch.cancel_by_prediction_id(&id).await { + if let Err(e) = orch.cancel_prediction(prediction).await { tracing::error!( prediction_id = %id, error = %e, @@ -923,9 +1187,9 @@ mod tests { } } - async fn cancel_by_prediction_id( + async fn cancel_prediction( &self, - _prediction_id: &str, + _prediction: Arc>, ) -> Result<(), crate::orchestrator::OrchestratorError> { Ok(()) } @@ -967,9 +1231,9 @@ mod tests { ) { } - async fn cancel_by_prediction_id( + async fn cancel_prediction( &self, - _prediction_id: &str, + _prediction: Arc>, ) -> Result<(), crate::orchestrator::OrchestratorError> { self.cancel_count.fetch_add(1, Ordering::SeqCst); Ok(()) @@ -988,14 +1252,12 @@ mod tests { struct CancelRecordingOrchestrator { cancel_count: AtomicUsize, - prediction: std::sync::Mutex>>>, } impl CancelRecordingOrchestrator { fn new() -> Self { Self { cancel_count: AtomicUsize::new(0), - prediction: std::sync::Mutex::new(None), } } @@ -1009,21 +1271,18 @@ mod tests { async fn register_prediction( &self, slot_id: SlotId, - prediction: Arc>, + _prediction: Arc>, idle_sender: tokio::sync::oneshot::Sender, ) { - *self.prediction.lock().unwrap() = Some(prediction); let _ = idle_sender.send(InactiveSlotIdleToken::new(slot_id).activate()); } - async fn cancel_by_prediction_id( + async fn cancel_prediction( &self, - _prediction_id: &str, + prediction: Arc>, ) -> Result<(), crate::orchestrator::OrchestratorError> { self.cancel_count.fetch_add(1, Ordering::SeqCst); - if let Some(prediction) = self.prediction.lock().unwrap().as_ref() { - prediction.lock().unwrap().set_canceled(); - } + prediction.lock().unwrap().set_canceled(); Ok(()) } @@ -1388,6 +1647,44 @@ mod tests { assert_eq!(orchestrator_ref.cancel_count(), 0); } + #[tokio::test] + async fn dropping_stale_stream_subscription_does_not_cancel_reused_id() { + let svc = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(CountingCancelOrchestrator::new()); + let orchestrator_ref = Arc::clone(&orchestrator); + + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + + let (_old_handle, _old_slot) = svc + .submit_prediction( + "reused-stream".to_string(), + serde_json::json!({}), + None, + true, + ) + .await + .unwrap(); + let old_subscription = svc.subscribe_prediction_stream("reused-stream").unwrap(); + svc.remove_prediction("reused-stream"); + + let (new_handle, _new_slot) = svc + .submit_prediction( + "reused-stream".to_string(), + serde_json::json!({}), + None, + true, + ) + .await + .unwrap(); + drop(old_subscription); + tokio::time::sleep(Duration::from_millis(25)).await; + + assert!(!new_handle.cancel_token().is_cancelled()); + assert_eq!(orchestrator_ref.cancel_count(), 0); + } + #[tokio::test] async fn submit_returns_at_capacity_when_no_slots() { let svc = PredictionService::new_no_pool(); @@ -1670,6 +1967,44 @@ mod tests { assert!(cancel_token.is_cancelled()); } + #[tokio::test] + async fn stale_sync_guard_does_not_cancel_reused_id() { + let svc = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(CountingCancelOrchestrator::new()); + let orchestrator_ref = Arc::clone(&orchestrator); + + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + + let (old_handle, _old_slot) = svc + .submit_prediction( + "reused-sync".to_string(), + serde_json::json!({}), + None, + false, + ) + .await + .unwrap(); + let old_guard = old_handle.sync_guard(Arc::clone(&svc)); + svc.remove_prediction("reused-sync"); + + let (new_handle, _new_slot) = svc + .submit_prediction( + "reused-sync".to_string(), + serde_json::json!({}), + None, + false, + ) + .await + .unwrap(); + drop(old_guard); + tokio::time::sleep(Duration::from_millis(25)).await; + + assert!(!new_handle.cancel_token().is_cancelled()); + assert_eq!(orchestrator_ref.cancel_count(), 0); + } + #[tokio::test] async fn sync_guard_disarm_prevents_cancel() { let svc = Arc::new(PredictionService::new_no_pool()); @@ -1723,6 +2058,351 @@ mod tests { assert!(!svc.prediction_exists("test-remove")); } + /// Unique per test so concurrent runs of the suite cannot delete each other's + /// fixtures. The prediction root is a fixed global path, so a shared name is + /// a real collision rather than a theoretical one. + fn unique_test_id(prefix: &str) -> String { + format!("{prefix}-{}", uuid::Uuid::new_v4()) + } + + #[test] + fn remove_prediction_deletes_prediction_dir() { + let svc = PredictionService::new_no_pool(); + let id = unique_test_id("dir-cleanup"); + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(dir.join("outputs")).unwrap(); + std::fs::create_dir_all(dir.join("inputs")).unwrap(); + std::fs::write(dir.join("outputs").join("0.png"), b"fake").unwrap(); + + svc.remove_prediction(&id); + + assert!(!dir.exists(), "prediction dir should be removed"); + } + + #[test] + fn remove_prediction_missing_dir_is_ok() { + let svc = PredictionService::new_no_pool(); + // Should not panic or error when the prediction dir doesn't exist + svc.remove_prediction(&unique_test_id("nonexistent")); + } + + #[test] + fn remove_prediction_rejects_ids_that_are_not_plain_path_components() { + let svc = PredictionService::new_no_pool(); + let victim = tempfile::tempdir().unwrap(); + std::fs::write(victim.path().join("keep.txt"), b"keep").unwrap(); + let victim_path = victim.path().to_str().unwrap(); + + // An absolute ID and a traversal ID both name a real directory outside the + // prediction root. Validation has to reject them before any deletion. + svc.remove_prediction(victim_path); + svc.remove_prediction(&format!("..{victim_path}")); + + assert!( + victim.path().join("keep.txt").exists(), + "cleanup must not touch anything outside the prediction root" + ); + } + + #[cfg(unix)] + #[test] + fn remove_prediction_unlinks_directory_symlink_without_following_it() { + let svc = PredictionService::new_no_pool(); + let id = unique_test_id("symlink"); + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(PREDICTION_ROOT).unwrap(); + let victim = tempfile::tempdir().unwrap(); + std::fs::write(victim.path().join("keep.txt"), b"keep").unwrap(); + std::os::unix::fs::symlink(victim.path(), &dir).unwrap(); + + svc.remove_prediction(&id); + + assert!( + std::fs::symlink_metadata(&dir).is_err(), + "prediction symlink should be removed" + ); + assert!( + victim.path().join("keep.txt").exists(), + "symlink target must survive" + ); + } + + #[tokio::test] + async fn remove_prediction_deletes_a_regular_file_and_releases_its_id() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + let id = unique_test_id("regular-file-cleanup"); + let (_handle, slot) = svc + .submit_prediction(id.clone(), serde_json::json!({}), None, false) + .await + .unwrap(); + // Cleanup still has to be safe if a run replaces its directory with a + // regular file. + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(PREDICTION_ROOT).unwrap(); + std::fs::write(&dir, b"not a directory").unwrap(); + + svc.remove_prediction(&id); + + assert!( + !svc.prediction_exists(&id), + "ID should be reusable after cleanup" + ); + assert!(!dir.exists(), "regular file should be removed"); + slot.release_unused("test complete".to_string()); + } + + #[tokio::test] + async fn remove_prediction_retains_id_when_detach_fails() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + let id = unique_test_id("detach-failure"); + let (_handle, slot) = svc + .submit_prediction(id.clone(), serde_json::json!({}), None, false) + .await + .unwrap(); + + svc.remove_prediction_with(&id, || { + Err(std::io::Error::other("injected detach failure")) + }); + + assert!( + svc.prediction_exists(&id), + "a stale directory must not be exposed to a reused ID" + ); + + slot.release_unused("test complete".to_string()); + svc.predictions.remove(&id); + } + + /// Detaching the directory before releasing the ID is what keeps a late + /// cleanup from deleting the files of a prediction that reused the ID. + #[test] + fn detach_prediction_dir_moves_the_directory_out_of_the_prediction_root() { + let id = unique_test_id("detach"); + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("output.txt"), b"payload").unwrap(); + + let detached = detach_prediction_dir(&id).unwrap().expect("dir existed"); + + assert!(!dir.exists(), "original path should be free for reuse"); + assert!(detached.starts_with(CLEANUP_ROOT)); + assert_eq!( + std::fs::read(detached.join("output.txt")).unwrap(), + b"payload" + ); + + delete_detached_dir(&detached); + assert!(!detached.exists()); + } + + #[test] + fn detach_prediction_dir_reports_nothing_to_clean() { + assert!( + detach_prediction_dir(&unique_test_id("absent")) + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn submit_prediction_rejects_duplicate_id() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + + svc.set_orchestrator(Arc::clone(&pool), orchestrator).await; + svc.set_health(Health::Ready).await; + + let (_handle, _slot) = svc + .submit_prediction( + "duplicate-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await + .unwrap(); + + let result = svc + .submit_prediction( + "duplicate-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await; + + assert!(matches!(result, Err(CreatePredictionError::DuplicateId(_)))); + assert_eq!(pool.available(), 1, "duplicate must not consume a permit"); + + let unrelated = svc + .submit_prediction( + "unrelated-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await; + assert!(unrelated.is_ok(), "remaining permit should still be usable"); + } + + /// The ID is reserved and the permit acquired under one map entry guard, so + /// concurrent submits of the same ID cannot both get through. A sequential + /// test passes even with a check-then-act implementation. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn submit_prediction_admits_one_of_many_concurrent_duplicates() { + const ATTEMPTS: usize = 16; + + let svc = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(ATTEMPTS).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(Arc::clone(&pool), orchestrator).await; + svc.set_health(Health::Ready).await; + + let mut submits = tokio::task::JoinSet::new(); + for _ in 0..ATTEMPTS { + let svc = Arc::clone(&svc); + submits.spawn(async move { + svc.submit_prediction("racing-id".to_string(), serde_json::json!({}), None, false) + .await + .map(drop) + }); + } + + let mut admitted = 0; + let mut duplicates = 0; + while let Some(result) = submits.join_next().await { + match result.unwrap() { + Ok(()) => admitted += 1, + Err(CreatePredictionError::DuplicateId(_)) => duplicates += 1, + Err(e) => panic!("unexpected submit error: {e}"), + } + } + + assert_eq!(admitted, 1, "exactly one submit may win the ID"); + assert_eq!(duplicates, ATTEMPTS - 1); + } + + /// A prediction that fails before its request reaches the worker has to give + /// the permit back. The worker never saw the slot, so no idle notification is + /// coming and a dropped permit would cost a slot permanently. + #[tokio::test] + async fn predict_returns_permit_when_request_preparation_fails() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(Arc::clone(&pool), orchestrator).await; + svc.set_health(Health::Ready).await; + + // A regular file where the prediction dir belongs makes create_dir_all fail. + let id = unique_test_id("prep-failure"); + std::fs::create_dir_all(PREDICTION_ROOT).unwrap(); + std::fs::write(prediction_dir_for(&id), b"not a directory").unwrap(); + + let (_handle, slot) = svc + .submit_prediction(id.clone(), serde_json::json!({}), None, false) + .await + .unwrap(); + assert_eq!(pool.available(), 0, "submit takes the only permit"); + let prediction = slot.prediction(); + + let result = svc + .predict(slot, serde_json::json!({}), Default::default()) + .await; + + assert!(matches!(result, Err(PredictionError::Failed(_)))); + assert_eq!(pool.available(), 1, "permit must return to the pool"); + let prediction = prediction.lock().unwrap(); + assert_eq!(prediction.status(), PredictionStatus::Failed); + assert!( + prediction.error().unwrap().contains("Failed to create"), + "error should name the failure: {:?}", + prediction.error() + ); + + let _ = std::fs::remove_file(prediction_dir_for(&id)); + } + + #[tokio::test] + async fn submit_prediction_rejects_invalid_id() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + + let long_id = "a".repeat(MAX_PREDICTION_ID_LEN + 1); + let invalid_ids = [ + "", + ".", + "..", + "foo/bar", + "foo\\bar", + "../escape", + &long_id, + "nul\0byte", + ]; + + for id in invalid_ids { + let result = svc + .submit_prediction(id.to_string(), serde_json::json!({}), None, false) + .await; + assert!( + matches!(result, Err(CreatePredictionError::InvalidId(_))), + "expected invalid ID error for {id:?}" + ); + } + } + + #[test] + fn validate_prediction_id_accepts_safe_ids() { + assert!(validate_prediction_id("pred_123").is_ok()); + assert!(validate_prediction_id("uuid-style-id").is_ok()); + assert!(validate_prediction_id("with.dots").is_ok()); + assert!(validate_prediction_id(".hidden").is_ok()); + assert!(validate_prediction_id(&"a".repeat(MAX_PREDICTION_ID_LEN)).is_ok()); + } + + #[test] + fn validate_prediction_id_rejects_dangerous_ids() { + assert!(validate_prediction_id("").is_err()); + assert!(validate_prediction_id("foo/bar").is_err()); + assert!(validate_prediction_id("foo\\bar").is_err()); + assert!(validate_prediction_id("../up").is_err()); + assert!(validate_prediction_id(".").is_err()); + assert!(validate_prediction_id("..").is_err()); + } + + /// An over-long ID used to pass validation and then fail `mkdir` with + /// ENAMETOOLONG mid-prediction, stranding both the permit and the ID. + #[test] + fn validate_prediction_id_rejects_ids_too_long_for_a_filename() { + let too_long = "a".repeat(MAX_PREDICTION_ID_LEN + 1); + assert!(validate_prediction_id(&too_long).is_err()); + assert!(validate_prediction_id(&"a".repeat(300)).is_err()); + } + + /// Control characters and non-ASCII cannot be trusted as a filename: a NUL + /// makes path syscalls fail outright, and some filesystems fold Unicode + /// spellings together so two distinct IDs would share one directory. + #[test] + fn validate_prediction_id_rejects_characters_unsafe_in_a_filename() { + assert!(validate_prediction_id("nul\0byte").is_err()); + assert!(validate_prediction_id("new\nline").is_err()); + assert!(validate_prediction_id("with space").is_err()); + assert!(validate_prediction_id("héllo").is_err()); + assert!(validate_prediction_id("colon:id").is_err()); + } + #[test] fn build_slot_request_small_input_inline() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 5904737391..d20f400aae 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -22,8 +22,8 @@ use crate::health::{HealthResponse, SetupResult}; use crate::prediction::SharedPredictionStreamEvent; use crate::predictor::PredictionError; use crate::service::{ - CreatePredictionError, HealthSnapshot, PredictionService, PredictionStreamSubscription, - SubscribePredictionStreamError, + CreatePredictionError, ExistingPrediction, HealthSnapshot, PredictionService, + PredictionStreamSubscription, SubscribePredictionStreamError, validate_prediction_id, }; use crate::version::VersionInfo; use crate::webhook::{TraceContext, WebhookConfig, WebhookEventType, WebhookSender}; @@ -288,6 +288,51 @@ fn extract_trace_context(headers: &HeaderMap) -> TraceContext { } } +fn invalid_prediction_id_response(msg: &str) -> Response { + ( + StatusCode::UNPROCESSABLE_ENTITY, + Json(serde_json::json!({ + "detail": [{ + "loc": ["body", "id"], + "msg": msg, + "type": "value_error" + }] + })), + ) + .into_response() +} + +/// Validate a caller-supplied ID, generating one when the request omits it. +/// +/// Rejecting here keeps a bad ID from reaching input validation, so the response +/// names the ID rather than whatever else is wrong with the request. Returns the +/// reason for the rejection, which callers pass to +/// [`invalid_prediction_id_response`]. +fn resolve_prediction_id(id: Option) -> Result { + let Some(id) = id else { + return Ok(generate_prediction_id()); + }; + validate_prediction_id(&id)?; + Ok(id) +} + +/// Check the ID in the body against the one in the URL for idempotent requests. +/// +/// The URL ID is already validated, so an invalid body ID cannot match it and is +/// reported as a mismatch. `subject` names the resource in the error message. +fn check_body_id_matches_url( + body_id: Option<&String>, + url_id: &str, + subject: &str, +) -> Result<(), String> { + match body_id { + Some(body_id) if body_id != url_id => Err(format!( + "{subject} ID must match the ID supplied in the URL" + )), + _ => Ok(()), + } +} + async fn create_prediction( State(service): State>, headers: HeaderMap, @@ -300,7 +345,10 @@ async fn create_prediction( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = request.id.unwrap_or_else(generate_prediction_id); + let prediction_id = match resolve_prediction_id(request.id) { + Ok(id) => id, + Err(message) => return invalid_prediction_id_response(message), + }; let response_mode = prediction_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -313,6 +361,7 @@ async fn create_prediction( response_mode, trace_context, false, + false, ) .await } @@ -323,6 +372,10 @@ async fn create_prediction_idempotent( headers: HeaderMap, body: Option>, ) -> Response { + if let Err(msg) = validate_prediction_id(&prediction_id) { + return invalid_prediction_id_response(msg); + } + let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest { id: None, input: serde_json::json!({}), @@ -331,35 +384,25 @@ async fn create_prediction_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id - && req_id != &prediction_id + if let Err(message) = + check_body_id_matches_url(request.id.as_ref(), &prediction_id, "prediction") { - return ( - StatusCode::UNPROCESSABLE_ENTITY, - Json(serde_json::json!({ - "detail": [{ - "loc": ["body", "id"], - "msg": "prediction ID must match the ID supplied in the URL", - "type": "value_error" - }] - })), - ) - .into_response(); + return invalid_prediction_id_response(&message); } let response_mode = prediction_response_mode(&headers); - // Check if prediction with this ID is already in-flight - if let Some(response) = service.get_prediction_response(&prediction_id) { - if response_mode == PredictionResponseMode::AsyncSse { - if !service.supports_prediction_streaming().await { - return streaming_not_supported_response(); - } - return stream_prediction_response(service, &prediction_id); - } - return (StatusCode::ACCEPTED, Json(response)).into_response(); + // Check if prediction with this ID is already in-flight. + if let Some(existing) = service.get_existing_prediction(&prediction_id) + && let Some(response) = + existing_idempotent_response(Arc::clone(&service), existing, response_mode).await + { + return response; } + #[cfg(test)] + service.wait_idempotent_pre_submit_barrier().await; + let trace_context = extract_trace_context(&headers); create_prediction_with_id( service, @@ -371,10 +414,30 @@ async fn create_prediction_idempotent( response_mode, trace_context, false, + true, ) .await } +async fn existing_idempotent_response( + service: Arc, + existing: ExistingPrediction, + response_mode: PredictionResponseMode, +) -> Option { + if response_mode == PredictionResponseMode::AsyncSse { + if !service.supports_prediction_streaming().await { + return Some(streaming_not_supported_response()); + } + let subscription = match existing.subscribe(&service) { + Ok(subscription) => subscription, + Err(error) => return Some(stream_subscription_error_response(error)), + }; + return Some(stream_prediction_subscription_response(subscription)); + } + let response = existing.response()?; + Some((StatusCode::ACCEPTED, Json(response)).into_response()) +} + fn build_webhook_sender( webhook: Option, events_filter: Vec, @@ -410,6 +473,7 @@ async fn create_prediction_with_id( response_mode: PredictionResponseMode, trace_context: TraceContext, is_training: bool, + is_idempotent: bool, ) -> Response { if !is_training && response_mode == PredictionResponseMode::AsyncSse @@ -488,6 +552,26 @@ async fn create_prediction_with_id( ) .into_response(); } + Err(CreatePredictionError::DuplicateId(existing)) => { + if is_idempotent + && let Some(response) = + existing_idempotent_response(Arc::clone(&service), existing, response_mode) + .await + { + return response; + } + return ( + StatusCode::CONFLICT, + Json(serde_json::json!({ + "error": "A prediction with this ID already exists", + "status": "failed" + })), + ) + .into_response(); + } + Err(CreatePredictionError::InvalidId(msg)) => { + return invalid_prediction_id_response(msg); + } }; let prediction = unregistered_slot.prediction(); @@ -498,6 +582,9 @@ async fn create_prediction_with_id( match service.subscribe_prediction_stream(&prediction_id) { Ok(subscription) => Some(subscription), Err(error) => { + unregistered_slot.release_unused(format!( + "Failed to subscribe to prediction stream: {error}" + )); service.remove_prediction(&prediction_id); return stream_subscription_error_response(error); } @@ -747,15 +834,6 @@ fn prediction_sse_stream( ) } -fn stream_prediction_response(service: Arc, prediction_id: &str) -> Response { - let subscription = match service.subscribe_prediction_stream(prediction_id) { - Ok(subscription) => subscription, - Err(error) => return stream_subscription_error_response(error), - }; - - stream_prediction_subscription_response(subscription) -} - fn stream_subscription_error_response(error: SubscribePredictionStreamError) -> Response { match error { SubscribePredictionStreamError::NotFound => ( @@ -823,7 +901,10 @@ async fn create_training( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = request.id.unwrap_or_else(generate_prediction_id); + let prediction_id = match resolve_prediction_id(request.id) { + Ok(id) => id, + Err(message) => return invalid_prediction_id_response(message), + }; let response_mode = json_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -836,6 +917,7 @@ async fn create_training( response_mode, trace_context, true, + false, ) .await } @@ -850,6 +932,10 @@ async fn create_training_idempotent( return training_streaming_not_supported_response(); } + if let Err(msg) = validate_prediction_id(&training_id) { + return invalid_prediction_id_response(msg); + } + let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest { id: None, input: serde_json::json!({}), @@ -858,27 +944,25 @@ async fn create_training_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id - && req_id != &training_id - { - return ( - StatusCode::UNPROCESSABLE_ENTITY, - Json(serde_json::json!({ - "detail": [{ - "loc": ["body", "id"], - "msg": "training ID must match the ID supplied in the URL", - "type": "value_error" - }] - })), - ) - .into_response(); + if let Err(message) = check_body_id_matches_url(request.id.as_ref(), &training_id, "training") { + return invalid_prediction_id_response(&message); } - // Idempotent: return existing state if already submitted - if let Some(response) = service.get_prediction_response(&training_id) { - return (StatusCode::ACCEPTED, Json(response)).into_response(); + // Idempotent: return existing state if already submitted. + if let Some(existing) = service.get_existing_prediction(&training_id) + && let Some(response) = existing_idempotent_response( + Arc::clone(&service), + existing, + PredictionResponseMode::AsyncJson, + ) + .await + { + return response; } + #[cfg(test)] + service.wait_idempotent_pre_submit_barrier().await; + let response_mode = json_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -891,6 +975,7 @@ async fn create_training_idempotent( response_mode, trace_context, true, + true, ) .await } @@ -1083,9 +1168,9 @@ mod tests { } } - async fn cancel_by_prediction_id( + async fn cancel_prediction( &self, - _prediction_id: &str, + _prediction: Arc>, ) -> Result<(), crate::orchestrator::OrchestratorError> { Ok(()) } @@ -1135,6 +1220,15 @@ mod tests { service } + async fn create_never_completing_service() -> (Arc, Arc) { + let service = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(MockOrchestrator::never_complete()); + service.set_orchestrator(pool, orchestrator.clone()).await; + service.set_health(Health::Ready).await; + (service, orchestrator) + } + async fn enable_prediction_streaming(service: &PredictionService) { service .set_schema(serde_json::json!({ @@ -1553,6 +1647,82 @@ mod tests { assert_eq!(json["status"], "succeeded"); } + #[tokio::test] + async fn concurrent_prediction_idempotent_puts_return_the_existing_prediction() { + let (service, orchestrator) = create_never_completing_service().await; + service + .set_idempotent_pre_submit_barrier(Arc::new(tokio::sync::Barrier::new(2))) + .await; + let app = routes(service); + let first = app.clone().oneshot( + Request::put("/predictions/idempotent-race") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ); + let second = app.oneshot( + Request::put("/predictions/idempotent-race") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ); + + let (first, second) = tokio::join!(first, second); + + let first = first.unwrap(); + let second = second.unwrap(); + assert_eq!(first.status(), StatusCode::ACCEPTED); + assert_eq!(second.status(), StatusCode::ACCEPTED); + let first = response_json(first).await; + let second = response_json(second).await; + assert_eq!(first["id"], "idempotent-race"); + assert_eq!(second["id"], "idempotent-race"); + assert_ne!(first["status"], "failed"); + assert_ne!(second["status"], "failed"); + + tokio::time::timeout(Duration::from_millis(100), async { + while orchestrator.register_count.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!(orchestrator.register_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn idempotent_duplicate_reservation_returns_the_existing_response() { + let (service, _orchestrator) = create_never_completing_service().await; + + for (id, is_training) in [ + ("prediction-duplicate", false), + ("training-duplicate", true), + ] { + let (_handle, _slot) = service + .submit_prediction(id.to_string(), serde_json::json!({}), None, false) + .await + .unwrap(); + + let response = create_prediction_with_id( + Arc::clone(&service), + id.to_string(), + serde_json::json!({}), + Default::default(), + None, + default_webhook_events_filter(), + PredictionResponseMode::AsyncJson, + TraceContext::default(), + is_training, + true, + ) + .await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + } + } + #[tokio::test] async fn prediction_idempotent_id_mismatch() { let service = create_ready_service().await; @@ -1578,6 +1748,115 @@ mod tests { ); } + #[tokio::test] + async fn prediction_rejects_invalid_id_in_body() { + let service = create_ready_service().await; + let app = routes(service); + + let response = app + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .body(Body::from(r#"{"id":"foo/bar","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + + #[tokio::test] + async fn prediction_rejects_invalid_id_in_url() { + let service = create_ready_service().await; + let app = routes(service); + + let response = app + .oneshot( + Request::put("/predictions/foo%5Cbar") + .header("content-type", "application/json") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + + /// Cleanup is wired into the request path, not just available as a helper: + /// without this, deleting the `remove_prediction` calls leaves every test green + /// while prediction directories pile up and completed IDs stay unusable. + #[tokio::test] + async fn completed_prediction_releases_its_id_and_directory() { + let service = create_ready_service().await; + let id = format!("cleanup-wiring-{}", uuid::Uuid::new_v4()); + let dir = std::path::PathBuf::from("/tmp/coglet/predictions").join(&id); + let app = routes(Arc::clone(&service)); + + let response = app + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .body(Body::from(format!(r#"{{"id":"{id}","input":{{}}}}"#))) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response_json(response).await["status"], "succeeded"); + assert!(!service.prediction_exists(&id), "ID should be reusable"); + assert!(!dir.exists(), "prediction dir should be deleted"); + } + + #[tokio::test] + async fn prediction_rejects_duplicate_id() { + // Use an orchestrator that never completes so the first prediction stays + // registered long enough for the duplicate request to be rejected. + let service = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(MockOrchestrator::never_complete()); + service.set_orchestrator(pool, orchestrator).await; + service.set_health(Health::Ready).await; + + let app = routes(Arc::clone(&service)); + let response = app + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"id":"dup-1","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + // submit_prediction registers the ID before the 202 is returned, so the + // duplicate below cannot race it. + assert!(service.prediction_exists("dup-1")); + + let app2 = routes(service); + let response = app2 + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .body(Body::from(r#"{"id":"dup-1","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + // AtCapacity is also 409, so check the body: with the permit taken before + // the duplicate check, this test would pass for the wrong reason. + assert_eq!(response.status(), StatusCode::CONFLICT); + let json = response_json(response).await; + assert!( + json["error"].as_str().unwrap().contains("already exists"), + "expected a duplicate-ID error, got {json}" + ); + } + #[tokio::test] async fn prediction_at_capacity() { let service = Arc::new(PredictionService::new_no_pool()); @@ -1744,6 +2023,51 @@ mod tests { assert_eq!(json["status"], "succeeded"); } + #[tokio::test] + async fn concurrent_training_idempotent_puts_return_the_existing_training() { + let (service, orchestrator) = create_never_completing_service().await; + service + .set_idempotent_pre_submit_barrier(Arc::new(tokio::sync::Barrier::new(2))) + .await; + let app = routes(service); + let first = app.clone().oneshot( + Request::put("/trainings/idempotent-race") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ); + let second = app.oneshot( + Request::put("/trainings/idempotent-race") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ); + + let (first, second) = tokio::join!(first, second); + + let first = first.unwrap(); + let second = second.unwrap(); + assert_eq!(first.status(), StatusCode::ACCEPTED); + assert_eq!(second.status(), StatusCode::ACCEPTED); + let first = response_json(first).await; + let second = response_json(second).await; + assert_eq!(first["id"], "idempotent-race"); + assert_eq!(second["id"], "idempotent-race"); + assert_ne!(first["status"], "failed"); + assert_ne!(second["status"], "failed"); + + tokio::time::timeout(Duration::from_millis(100), async { + while orchestrator.register_count.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!(orchestrator.register_count.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn training_idempotent_id_mismatch() { let service = create_ready_service().await; diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index c59a26cc7a..716e998a71 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -11,6 +11,7 @@ use std::collections::HashMap; use std::io; +use std::io::Write; use std::path::PathBuf; use std::sync::Arc; use std::sync::OnceLock; @@ -134,6 +135,7 @@ use crate::bridge::codec::JsonCodec; use crate::bridge::protocol::{ ControlRequest, ControlResponse, FileOutputKind, LogSource, MAX_INLINE_IPC_SIZE, MetricMode, SLOT_RESPONSE_PROTOCOL_VERSION, SlotId, SlotOutcome, SlotRequest, SlotResponse, + StructuredOutput, StructuredOutputFile, }; use crate::bridge::transport::{ChildTransportInfo, connect_transport}; use crate::orchestrator::HealthcheckResult; @@ -142,6 +144,14 @@ use crate::worker_tracing_layer::WorkerTracingLayer; type SlotWriter = Arc>>>; +#[derive(Debug, Clone)] +pub struct StagedFileOutput { + pub filename: String, + pub upload_filename: String, + pub mime_type: Option, + pub managed: bool, +} + /// Handle for sending messages on a slot socket. /// /// Used by log writers to stream logs during prediction. Thread-safe via @@ -168,12 +178,32 @@ impl SlotSender { self.output_counter.fetch_add(1, Ordering::Relaxed) } - /// Generate a unique filename in the output dir. + /// Generate the next candidate filename in the output dir. fn next_output_path(&self, extension: &str) -> PathBuf { let n = self.file_counter.fetch_add(1, Ordering::Relaxed); self.output_dir.join(format!("{n}.{extension}")) } + /// Create a new file in the output dir, returning it and its path. + /// + /// Creating exclusively rather than testing with `exists()` skips names that + /// are already taken and refuses to write through a pre-existing symlink, + /// which `exists()` cannot detect when the link dangles. + fn create_output_file(&self, extension: &str) -> io::Result<(std::fs::File, PathBuf)> { + loop { + let path = self.next_output_path(extension); + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + { + Ok(file) => return Ok((file, path)), + Err(e) if e.kind() == io::ErrorKind::AlreadyExists => continue, + Err(e) => return Err(e), + } + } + } + pub fn send_log(&self, source: LogSource, data: &str) -> io::Result<()> { if data.is_empty() { return Ok(()); @@ -199,24 +229,148 @@ impl SlotSender { extension: &str, mime_type: Option, ) -> io::Result<()> { - let path = self.next_output_path(extension); - std::fs::write(&path, data)?; - self.send_file_output(path, mime_type) + let (mut file, path) = self.create_output_file(extension)?; + file.write_all(data)?; + self.send_file_output(path, mime_type, true) } /// Send a file-typed output (e.g. Path, File return types). /// /// The file is already on disk at `path` — we just send the path reference. /// `mime_type` is an explicit MIME type; when None the parent guesses from extension. - pub fn send_file_output(&self, path: PathBuf, mime_type: Option) -> io::Result<()> { - let filename = path - .to_str() - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path"))? + /// `managed` is true when Coglet owns the file and may delete it after consumption. + pub fn send_file_output( + &self, + path: PathBuf, + mime_type: Option, + managed: bool, + ) -> io::Result<()> { + let upload_filename = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("output") .to_string(); + let staged = StagedFileOutput { + filename: path_to_string(&path)?, + upload_filename, + mime_type, + managed, + }; + self.send_staged_file_output(staged, self.next_output_index()) + } + + fn send_staged_file_output(&self, file: StagedFileOutput, index: u64) -> io::Result<()> { let msg = SlotResponse::FileOutput { - filename, + filename: file.filename, kind: FileOutputKind::FileType, + mime_type: file.mime_type, + managed: file.managed, + upload_filename: Some(file.upload_filename), + index, + }; + self.tx + .send(msg) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) + } + + /// Transfer a run-returned file into Coglet's managed output directory. + /// + /// Returning a path hands the data to Coglet, which deletes the managed copy + /// once the parent has read it. The source path is consumed during the transfer + /// so it cannot be returned again unless the predictor rewrites it. + pub fn send_user_file_output( + &self, + path: PathBuf, + mime_type: Option, + ) -> io::Result<()> { + let staged = self.stage_user_file_output(path, mime_type)?; + self.send_staged_file_output(staged, self.next_output_index()) + } + + pub fn stage_user_file_output( + &self, + path: PathBuf, + mime_type: Option, + ) -> io::Result { + if !path.exists() { + return Err(io::Error::new( + io::ErrorKind::NotFound, + format!( + "output file {} does not exist. Cog takes ownership of \ + returned files, so a path returned more than once must be \ + rewritten before each output.", + path.display() + ), + )); + } + + let extension = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("bin") + .to_string(); + let upload_filename = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("output.bin") + .to_string(); + // Close our handle before transferring: rename replaces the entry and + // copy reopens it, so holding it open buys nothing. + let (file, dest) = self.create_output_file(&extension)?; + drop(file); + + if let Err(e) = move_file(&path, &dest) { + let _ = std::fs::remove_file(&dest); + return Err(e); + } + + Ok(StagedFileOutput { + filename: path_to_string(&dest)?, + upload_filename, + mime_type, + managed: true, + }) + } + + pub fn stage_file_output( + &self, + data: &[u8], + upload_filename: Option, + mime_type: Option, + ) -> io::Result { + let upload_filename = upload_filename + .filter(|name| !name.is_empty()) + .unwrap_or_else(|| "output.bin".to_string()); + let extension = std::path::Path::new(&upload_filename) + .extension() + .and_then(|extension| extension.to_str()) + .unwrap_or("bin"); + let (mut file, path) = self.create_output_file(extension)?; + file.write_all(data)?; + Ok(StagedFileOutput { + filename: path_to_string(&path)?, + upload_filename, mime_type, + managed: true, + }) + } + + pub fn send_structured_output( + &self, + output: serde_json::Value, + files: Vec, + ) -> io::Result<()> { + let payload = serde_json::to_vec(&StructuredOutput { output, files }) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + let (mut file, path) = self.create_output_file("json")?; + file.write_all(&payload)?; + let msg = SlotResponse::FileOutput { + filename: path_to_string(&path)?, + kind: FileOutputKind::Structured, + mime_type: None, + managed: true, + upload_filename: None, + index: self.next_output_index(), }; self.tx .send(msg) @@ -245,6 +399,37 @@ impl SlotSender { } } +fn path_to_string(path: &std::path::Path) -> io::Result { + path.to_str() + .map(ToString::to_string) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path")) +} + +/// Move a file, falling back to copy-then-unlink when the two paths live on +/// different filesystems and `rename` cannot be used. +/// +/// The fallback copies first so a failure part-way through leaves the source +/// intact. +fn move_file(source: &std::path::Path, dest: &std::path::Path) -> io::Result<()> { + // Renaming a relative symlink into another directory changes what its target + // means. Copy the target bytes into a regular file, then unlink the source + // entry instead. + if std::fs::symlink_metadata(source)?.file_type().is_symlink() { + return copy_then_unlink(source, dest); + } + + match std::fs::rename(source, dest) { + Ok(()) => Ok(()), + Err(e) if e.kind() == io::ErrorKind::CrossesDevices => copy_then_unlink(source, dest), + Err(e) => Err(e), + } +} + +fn copy_then_unlink(source: &std::path::Path, dest: &std::path::Path) -> io::Result<()> { + std::fs::copy(source, dest)?; + std::fs::remove_file(source) +} + /// Build an output message, spilling to disk if larger than the IPC frame limit. fn build_output_message( output_dir: &std::path::Path, @@ -265,6 +450,9 @@ fn build_output_message( filename, kind: FileOutputKind::Oversized, mime_type: None, + managed: true, + upload_filename: None, + index, }) } else { Ok(SlotResponse::OutputChunk { output, index }) @@ -973,4 +1161,279 @@ mod tests { let config = WorkerConfig::default(); assert_eq!(config.num_slots, 1); } + + /// Unwrap a managed FileOutput message, checking it points into `output_dir`. + fn expect_managed_output(response: SlotResponse, output_dir: &std::path::Path) -> PathBuf { + match response { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => { + let path = PathBuf::from(filename); + assert!( + path.starts_with(output_dir), + "managed output should live in the output dir, got {}", + path.display() + ); + path + } + response => panic!("expected a managed FileOutput, got {response:?}"), + } + } + + #[test] + fn send_user_file_output_moves_a_temp_file() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let scratch = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); + std::fs::write(scratch.path(), b"hello").unwrap(); + let scratch_path = scratch.path().to_path_buf(); + + sender + .send_user_file_output(scratch_path.clone(), None) + .unwrap(); + + assert!( + !scratch_path.exists(), + "temp file should be moved, not copied" + ); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(&managed).unwrap(), b"hello"); + } + + #[test] + fn send_user_file_output_moves_a_file_already_in_the_output_dir() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let inside = output_dir.path().join("inside.txt"); + std::fs::write(&inside, b"world").unwrap(); + + sender.send_user_file_output(inside.clone(), None).unwrap(); + + assert!(!inside.exists(), "returned path should be moved"); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_ne!(managed, inside); + assert_eq!(std::fs::read(&managed).unwrap(), b"world"); + } + + #[test] + fn send_user_file_output_consumes_a_file_outside_the_output_dir() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let source_dir = tempfile::tempdir_in(std::env::current_dir().unwrap()).unwrap(); + let source = source_dir.path().join("fixed-output.png"); + std::fs::write(&source, b"output").unwrap(); + + sender.send_user_file_output(source.clone(), None).unwrap(); + + assert!(!source.exists(), "returned source path should be consumed"); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(&managed).unwrap(), b"output"); + + let error = sender + .send_user_file_output(source, None) + .expect_err("a consumed path cannot return stale output"); + assert_eq!(error.kind(), io::ErrorKind::NotFound); + assert!(rx.try_recv().is_err(), "stale path must not be sent again"); + } + + #[test] + fn send_user_file_output_preserves_the_logical_filename() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + let source_dir = tempfile::tempdir().unwrap(); + let source = source_dir.path().join("result.png"); + std::fs::write(&source, b"png").unwrap(); + + sender.send_user_file_output(source, None).unwrap(); + + match rx.try_recv().unwrap() { + SlotResponse::FileOutput { + upload_filename, + index, + .. + } => { + assert_eq!(upload_filename.as_deref(), Some("result.png")); + assert_eq!(index, 0); + } + response => panic!("expected FileOutput, got {response:?}"), + } + } + + #[cfg(unix)] + #[test] + fn send_user_file_output_copies_a_relative_symlink_before_unlinking_it() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let scratch_dir = tempfile::tempdir().unwrap(); + let target = scratch_dir.path().join("target.txt"); + std::fs::write(&target, b"target bytes").unwrap(); + let link = scratch_dir.path().join("output.txt"); + std::os::unix::fs::symlink("target.txt", &link).unwrap(); + + sender.send_user_file_output(link.clone(), None).unwrap(); + + assert!(!link.exists(), "returned symlink should be unlinked"); + assert!(target.exists(), "the symlink target should survive"); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(managed).unwrap(), b"target bytes"); + } + + #[cfg(unix)] + #[test] + fn send_user_file_output_unlinks_a_symlink_without_deleting_its_target() { + let output_dir = tempfile::tempdir().unwrap(); + let scratch_dir = tempfile::tempdir().unwrap(); + let external_dir = tempfile::tempdir().unwrap(); + let target = external_dir.path().join("asset.txt"); + std::fs::write(&target, b"asset bytes").unwrap(); + let link = scratch_dir.path().join("result.txt"); + std::os::unix::fs::symlink(&target, &link).unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + sender.send_user_file_output(link.clone(), None).unwrap(); + + assert!(std::fs::symlink_metadata(&link).is_err()); + assert_eq!(std::fs::read(&target).unwrap(), b"asset bytes"); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(managed).unwrap(), b"asset bytes"); + } + + #[test] + fn send_user_file_output_explains_a_path_that_no_longer_exists() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let scratch = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); + let scratch_path = scratch.path().to_path_buf(); + sender + .send_user_file_output(scratch_path.clone(), None) + .unwrap(); + rx.try_recv().unwrap(); + + let error = sender + .send_user_file_output(scratch_path, None) + .expect_err("the file was already handed over"); + + assert_eq!(error.kind(), io::ErrorKind::NotFound); + assert!( + error.to_string().contains("rewritten before each output"), + "error should explain the ownership handoff: {error}" + ); + assert!(rx.try_recv().is_err(), "stale path must not be sent again"); + } + + /// `exists()` reports false for a dangling symlink, so testing with it would + /// let `copy` create the link's target somewhere outside the output dir. + #[cfg(unix)] + #[test] + fn create_output_file_refuses_to_write_through_a_dangling_symlink() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, _rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let victim_dir = tempfile::tempdir().unwrap(); + let victim = victim_dir.path().join("victim.txt"); + std::os::unix::fs::symlink(&victim, output_dir.path().join("0.txt")).unwrap(); + + let (_file, created) = sender.create_output_file("txt").unwrap(); + + assert_ne!(created, output_dir.path().join("0.txt")); + assert!(!victim.exists(), "symlink target must not be created"); + } + + #[test] + fn write_file_output_marks_the_file_as_managed() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + sender.write_file_output(b"bytes", "bin", None).unwrap(); + + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(&managed).unwrap(), b"bytes"); + } + + #[test] + fn structured_output_is_spilled_with_its_file_descriptors() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + let files = vec![StructuredOutputFile { + pointer: "/artifact".to_string(), + filename: "/tmp/managed.png".to_string(), + upload_filename: "result.png".to_string(), + mime_type: Some("image/png".to_string()), + managed: true, + }]; + let expected = StructuredOutput { + output: serde_json::json!({"artifact": null, "label": "result"}), + files: files.clone(), + }; + + sender + .send_structured_output(expected.output.clone(), files) + .unwrap(); + + match rx.try_recv().unwrap() { + SlotResponse::FileOutput { + filename, + kind: FileOutputKind::Structured, + managed: true, + index: 0, + .. + } => { + let actual = + serde_json::from_slice::(&std::fs::read(filename).unwrap()) + .unwrap(); + assert_eq!(actual, expected); + } + response => panic!("expected structured FileOutput, got {response:?}"), + } + } + + /// Spilled outputs are Coglet's own files, so the parent has to delete them. + #[test] + fn oversized_output_spills_to_a_managed_file() { + let output_dir = tempfile::tempdir().unwrap(); + let oversized = serde_json::json!("x".repeat(MAX_INLINE_IPC_SIZE + 1)); + + let msg = build_output_message(output_dir.path(), oversized.clone(), 0).unwrap(); + + match msg { + SlotResponse::FileOutput { + filename, + kind: FileOutputKind::Oversized, + managed: true, + index: 0, + .. + } => { + let spilled: serde_json::Value = + serde_json::from_slice(&std::fs::read(filename).unwrap()).unwrap(); + assert_eq!(spilled, oversized); + } + msg => panic!("expected an oversized managed FileOutput, got {msg:?}"), + } + } + + #[test] + fn small_output_stays_inline() { + let output_dir = tempfile::tempdir().unwrap(); + + let msg = build_output_message(output_dir.path(), serde_json::json!("small"), 7).unwrap(); + + assert!(matches!(msg, SlotResponse::OutputChunk { index: 7, .. })); + } } diff --git a/docs/http.md b/docs/http.md index de2a36dac4..897bd4c64b 100644 --- a/docs/http.md +++ b/docs/http.md @@ -261,6 +261,9 @@ The client should ensure prediction slots are available before creating a new prediction with a different ID. Clients are responsible for providing unique prediction IDs. +IDs must be 1 to 128 ASCII characters and may contain only letters, numbers, +hyphens (`-`), underscores (`_`), and periods (`.`). The values `.` and `..` +are not valid IDs. We recommend generating a UUIDv4 or [UUIDv7](https://uuid7.com), base32-encoding that value, and removing padding characters (`==`). diff --git a/docs/llms.txt b/docs/llms.txt index c9e6a8dd7c..79af3fcdd2 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1916,6 +1916,9 @@ The client should ensure prediction slots are available before creating a new prediction with a different ID. Clients are responsible for providing unique prediction IDs. +IDs must be 1 to 128 ASCII characters and may contain only letters, numbers, +hyphens (`-`), underscores (`_`), and periods (`.`). The values `.` and `..` +are not valid IDs. We recommend generating a UUIDv4 or [UUIDv7](https://uuid7.com), base32-encoding that value, and removing padding characters (`==`). @@ -2897,6 +2900,7 @@ These types can be used directly as input parameter types and output return type `cog.Path` is a subclass of Python's [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html#basic-use) and can be used as a drop-in replacement. Any `os.PathLike` subclass is also accepted as an input type and treated as `cog.Path`. For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. +Returning a path transfers ownership of that directory entry to Cog. Cog removes the entry while staging the output and does not restore it if later staging or upload work fails. When the path is a symlink, Cog removes the link but leaves its target intact. Models that return assets or shared cache files must copy them to a separate output path first. This example takes an input file, resizes it, and returns the resized image: @@ -2909,7 +2913,7 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # This file will automatically be deleted by Cog after it has been returned. + # Cog consumes the returned path, so write it again before returning the same path twice. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) diff --git a/docs/python.md b/docs/python.md index d717cf56f2..2d2fa0a8a3 100644 --- a/docs/python.md +++ b/docs/python.md @@ -525,6 +525,7 @@ These types can be used directly as input parameter types and output return type `cog.Path` is a subclass of Python's [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html#basic-use) and can be used as a drop-in replacement. Any `os.PathLike` subclass is also accepted as an input type and treated as `cog.Path`. For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. +Returning a path transfers ownership of that directory entry to Cog. Cog removes the entry while staging the output and does not restore it if later staging or upload work fails. When the path is a symlink, Cog removes the link but leaves its target intact. Models that return assets or shared cache files must copy them to a separate output path first. This example takes an input file, resizes it, and returns the resized image: @@ -537,7 +538,7 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # This file will automatically be deleted by Cog after it has been returned. + # Cog consumes the returned path, so write it again before returning the same path twice. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path)