Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions architecture/04-container-runtime.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion crates/coglet-python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
235 changes: 226 additions & 9 deletions crates/coglet-python/src/output.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,20 +5,23 @@
//! - 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};

/// Process prediction output for JSON serialization.
///
/// 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>,
Expand All @@ -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<Bound<'py, PyAny>> {
process_output(py, item)
output: &Bound<'py, PyAny>,
slot_sender: &SlotSender,
) -> PyResult<(Bound<'py, PyAny>, Vec<StructuredOutputFile>)> {
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<Bound<'py, PyAny>> {
// Pydantic v2: model_dump()
if let Ok(method) = obj.getattr("model_dump")
Expand Down Expand Up @@ -108,7 +119,7 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult<Bo
return obj.call_method0("isoformat");
}

// os.PathLike -> 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)? {
Expand Down Expand Up @@ -142,6 +153,7 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult<Bo
}

/// Recursively walk the output and encode any Path/IOBase objects to base64 data URLs.
#[cfg(test)]
fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
// str -- return as-is (don't recurse into characters)
if obj.is_instance_of::<PyString>() {
Expand Down Expand Up @@ -189,10 +201,121 @@ fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult<Bound
Ok(obj.clone())
}

fn stage_files<'py>(
py: Python<'py>,
obj: &Bound<'py, PyAny>,
slot_sender: &SlotSender,
pointer: &str,
files: &mut Vec<StructuredOutputFile>,
) -> PyResult<Bound<'py, PyAny>> {
if obj.is_instance_of::<PyString>() {
return Ok(obj.clone());
}

if let Ok(dict) = obj.cast_exact::<PyDict>() {
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::<PyList>() {
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::<PyResult<Vec<_>>>()?;
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::<PyString>() {
content.extract::<String>()?.into_bytes()
} else {
content.extract::<Vec<u8>>()?
};
let upload_filename = obj
.getattr("name")
.and_then(|name| name.extract::<String>())
.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<String> {
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::<serde_json::Map<String, serde_json::Value>>(&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<Bound<'py, PyAny>> {
// Seek to start if possible
if let Ok(seekable) = fh.call_method0("seekable")
Expand Down Expand Up @@ -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::<serde_json::Value>(&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();
Expand Down
Loading
Loading