From 890c0479b1c24deed4a473399d84d664f8705ab8 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Sat, 31 Oct 2026 17:23:50 -0400 Subject: [PATCH 1/3] fix(coglet): delete returned output paths after consumption Returning or yielding a Path now transfers ownership to Cog. The worker snapshots the file into a unique temp file under output_dir and unlinks the source entry before user code resumes, so a later write to the same path cannot change an already-streamed output. Symlink targets are copied but never unlinked. The parent reads the snapshot into memory and deletes it before uploading, so failed or canceled uploads no longer leave files behind. Staged files use tempfile names instead of a per-sender counter, which fixes collisions between senders sharing an output directory. Nested Paths are encoded inline and collected in a HashSet so repeated paths in one output still work, then deleted once encoding finishes. Upload filenames keep their original basename and are appended as a single escaped URL path segment, preserving signed query parameters. Fixes #1434. --- architecture/04-container-runtime.md | 4 +- crates/coglet-python/src/output.rs | 94 +++++++++++++- crates/coglet-python/src/predictor.rs | 4 +- crates/coglet-python/tests/test_coglet.py | 45 +++++++ crates/coglet/Cargo.toml | 2 +- crates/coglet/src/bridge/protocol.rs | 6 +- crates/coglet/src/orchestrator.rs | 135 ++++++++++++++++--- crates/coglet/src/worker.rs | 151 +++++++++++++++++++--- docs/llms.txt | 2 + docs/python.md | 2 + 10 files changed, 402 insertions(+), 43 deletions(-) diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 17c94ed55c..3bfce4f291 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -179,7 +179,7 @@ Per-prediction data. Using separate sockets per slot avoids head-of-line blockin | ---------------------------------------------- | ---------------------------------------------------------------- | | `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 | +| `FileOutput { filename, kind, mime_type, upload_filename }` | Owned output snapshot referenced by path; optional original basename for uploads | | `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 | @@ -373,6 +373,8 @@ When a prediction input exceeds 6MiB, it's too large to send inline through the 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. +Returning or yielding a path transfers ownership to Cog. The worker snapshots its contents into a unique temporary file and unlinks the source entry before resuming user code. This prevents a later overwrite from changing an earlier streamed output. The parent deletes the snapshot after reading it; uploads use the bytes in memory. Nested paths are encoded inline and unlinked after the containing output is encoded. Cleanup removes files only, never prediction directories derived from public IDs. + `FileOutputKind` distinguishes between normal file outputs (`FileType`) and oversized outputs (`Oversized`) that exceeded an inline size limit. ## Custom Metrics diff --git a/crates/coglet-python/src/output.rs b/crates/coglet-python/src/output.rs index c4300d6bb1..a1128686a9 100644 --- a/crates/coglet-python/src/output.rs +++ b/crates/coglet-python/src/output.rs @@ -24,7 +24,17 @@ pub fn process_output<'py>( output: &Bound<'py, PyAny>, ) -> PyResult> { let encodeable = make_encodeable(py, output)?; - encode_files(py, &encodeable) + let mut paths = std::collections::HashSet::new(); + let encoded = encode_files(py, &encodeable, &mut paths)?; + // Wait until every occurrence is encoded so repeated paths in one output work. + for path in paths { + if let Err(error) = std::fs::remove_file(path) + && error.kind() != std::io::ErrorKind::NotFound + { + return Err(error.into()); + } + } + Ok(encoded) } /// Process a single output item (for generator outputs). @@ -142,7 +152,11 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult> { +fn encode_files<'py>( + py: Python<'py>, + obj: &Bound<'py, PyAny>, + paths: &mut std::collections::HashSet, +) -> PyResult> { // str -- return as-is (don't recurse into characters) if obj.is_instance_of::() { return Ok(obj.clone()); @@ -152,7 +166,7 @@ fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult() { let new_dict = PyDict::new(py); for (key, value) in dict.iter() { - new_dict.set_item(&key, encode_files(py, &value)?)?; + new_dict.set_item(&key, encode_files(py, &value, paths)?)?; } return Ok(new_dict.into_any()); } @@ -161,7 +175,7 @@ fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult() { let items: Vec> = list .iter() - .map(|item| encode_files(py, &item)) + .map(|item| encode_files(py, &item, paths)) .collect::>()?; let new_list = PyList::new(py, &items)?; return Ok(new_list.into_any()); @@ -175,7 +189,9 @@ fn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult()?); + return Ok(encoded); } // io.IOBase -> base64 encode @@ -442,6 +458,74 @@ mod tests { // ── encode_files / file_to_base64 ──────────────────────────────── + #[test] + fn nested_repeated_paths_are_encoded_then_consumed() { + let root = tempfile::tempdir().unwrap(); + let file = root.path().join("result.txt"); + std::fs::write(&file, b"nested").unwrap(); + pyo3::Python::initialize(); + Python::attach(|py| { + let path = py + .import("pathlib") + .unwrap() + .getattr("Path") + .unwrap() + .call1((&file,)) + .unwrap(); + let obj = PyDict::new(py); + obj.set_item("first", &path).unwrap(); + obj.set_item("second", PyList::new(py, [&path]).unwrap()) + .unwrap(); + let encoded = process_output(py, obj.as_any()).unwrap(); + assert_eq!( + encoded + .get_item("first") + .unwrap() + .extract::() + .unwrap(), + "data:text/plain;base64,bmVzdGVk" + ); + assert_eq!( + encoded + .get_item("second") + .unwrap() + .get_item(0) + .unwrap() + .extract::() + .unwrap(), + "data:text/plain;base64,bmVzdGVk" + ); + assert!(!file.exists()); + assert!( + process_output(py, &path).is_err(), + "a stale path must not be reused" + ); + }); + } + + #[cfg(unix)] + #[test] + fn nested_symlink_preserves_target() { + let root = tempfile::tempdir().unwrap(); + let target = root.path().join("asset.txt"); + let link = root.path().join("result.txt"); + std::fs::write(&target, b"asset").unwrap(); + std::os::unix::fs::symlink("asset.txt", &link).unwrap(); + pyo3::Python::initialize(); + Python::attach(|py| { + let path = py + .import("pathlib") + .unwrap() + .getattr("Path") + .unwrap() + .call1((&link,)) + .unwrap(); + process_output(py, &path).unwrap(); + }); + assert!(!link.exists()); + assert_eq!(std::fs::read(target).unwrap(), b"asset"); + } + #[test] fn encode_pathlike_to_base64() { let result = processed( diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index de6f188114..ae09c60725 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -212,7 +212,7 @@ fn submit_async_coroutine( /// 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 ownership via send_file_output. /// 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. fn send_output_item( @@ -235,7 +235,7 @@ 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 + // Snapshot and consume the returned path before the predictor resumes. let path_str: String = item .call_method0("__fspath__") .and_then(|p| p.extract()) diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 57dcdd40a8..214d290162 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -520,6 +520,51 @@ def test_includes_predict_time(self, sync_predictor: Path): assert result["metrics"]["predict_time"] >= 0 +@pytest.mark.parametrize("shape", ["scalar", "nested", "generator"]) +def test_returned_paths_are_consumed(tmp_path: Path, shape: str) -> None: + predictor = tmp_path / "predict.py" + predictor.write_text(""" +from pathlib import Path +from cog import BasePredictor + +class Predictor(BasePredictor): + def predict(self, write: bool = True) -> Path: + path = Path(__file__).parent / "result.txt" + if write: + path.write_text("hello") + return path +""") + if shape == "nested": + predictor.write_text( + predictor.read_text() + .replace("-> Path:", "-> dict[str, Path]:") + .replace("return path", 'return {"file": path}') + ) + elif shape == "generator": + predictor.write_text( + "from typing import Iterator\n" + + predictor.read_text() + .replace("-> Path:", "-> Iterator[Path]:") + .replace( + "return path", + 'yield path\n assert not path.exists(), "yielded path was not consumed"', + ) + ) + (tmp_path / "cog.yaml").write_text('predict: "predict.py:Predictor"\n') + with CogletServer(predictor) as server: + result = server.predict({}) + assert result["status"] == "succeeded", result + output = result["output"] + if shape == "nested": + output = output["file"] + elif shape == "generator": + output = output[0] + assert output == "data:text/plain;base64,aGVsbG8=" + assert not (tmp_path / "result.txt").exists() + result = server.predict({"write": False}) + assert result["status"] == "failed", result + + class TestSecretInput: """Tests for cog.Secret input coercion.""" diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index 32502e9ec3..a935eb3fb0 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -46,6 +46,7 @@ mime_guess = "2.0.5" # Identifiers uuid.workspace = true +tempfile = "3" # HTTP server axum.workspace = true @@ -84,7 +85,6 @@ nix = { version = "0.31", features = ["signal", "fs"] } [dev-dependencies] insta.workspace = true -tempfile = "3" wiremock = "0.6" tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 1d8e9bf741..61b2c29663 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -352,13 +352,17 @@ pub enum SlotResponse { }, /// Output for a file/path-like output return type or an output that exceeds the size threshold - /// for bridge codec serialization. + /// for bridge codec serialization. The parent owns the file and unlinks it + /// after reading its bytes, including when the prediction has been canceled. FileOutput { filename: String, kind: FileOutputKind, /// Explicit MIME type from the predictor. Falls back to mime_guess when None. #[serde(skip_serializing_if = "Option::is_none")] mime_type: Option, + /// Original basename when the file was staged under a temporary name. + #[serde(default, skip_serializing_if = "Option::is_none")] + upload_filename: Option, }, /// Streaming output chunk for generator and iterator output. diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index c51bd47801..0655fa354b 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -33,6 +33,17 @@ use crate::prediction::Prediction; const MAX_PENDING_CANCELLATIONS: usize = 1000; +/// Read an owned output snapshot and release its disk space before uploading. +fn read_output_file(filename: &str) -> std::io::Result> { + let result = std::fs::read(filename); + if let Err(error) = std::fs::remove_file(filename) + && error.kind() != std::io::ErrorKind::NotFound + { + tracing::warn!(%filename, %error, "Failed to remove output file"); + } + result +} + /// Upload a file to a signed endpoint, returning the final URL. /// /// Matches the behavior of Python cog's `put_file_to_signed_endpoint`: @@ -46,10 +57,14 @@ async fn upload_file( content_type: &str, prediction_id: &str, ) -> Result { - let url = format!("{endpoint}{filename}"); + let mut url = reqwest::Url::parse(endpoint).map_err(|_| "invalid upload URL".to_string())?; + url.path_segments_mut() + .map_err(|_| "upload URL cannot contain path segments".to_string())? + .pop_if_empty() + .push(filename); let client = reqwest::Client::new(); let resp = client - .put(&url) + .put(url) .header("Content-Type", content_type) .header("X-Prediction-Id", prediction_id) .body(data.to_vec()) @@ -86,14 +101,6 @@ async fn upload_file( } } -fn ensure_trailing_slash(s: &str) -> String { - if s.ends_with('/') { - s.to_string() - } else { - format!("{s}/") - } -} - /// Try to lock a prediction mutex. /// On poison: logs error, recovers to fail the prediction, returns None. /// Caller should remove the prediction from tracking if None is returned. @@ -1119,9 +1126,9 @@ async fn run_event_loop( predictions.remove(&slot_id); } } - Ok(SlotResponse::FileOutput { filename, kind, mime_type }) => { + Ok(SlotResponse::FileOutput { filename, kind, mime_type, upload_filename }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); - let bytes = match std::fs::read(&filename) { + let bytes = match read_output_file(&filename) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); @@ -1153,7 +1160,7 @@ async fn run_event_loop( } FileOutputKind::FileType => { let mime = mime_type.unwrap_or_else(|| { - mime_guess::from_path(&filename) + mime_guess::from_path(upload_filename.as_deref().unwrap_or(&filename)) .first_or_octet_stream() .to_string() }); @@ -1168,12 +1175,12 @@ async fn run_event_loop( }; #[cfg(feature = "tracing")] let trace = try_lock_prediction(&pred).and_then(|p| p.trace_carrier()); - let endpoint = ensure_trailing_slash(url); - let basename = std::path::Path::new(&filename) + let endpoint = url.clone(); + let basename = upload_filename.unwrap_or_else(|| std::path::Path::new(&filename) .file_name() .and_then(|n| n.to_str()) .unwrap_or("output") - .to_string(); + .to_string()); let upload_span = crate::cog_span!( info_span, "cog.prediction.upload_output", @@ -1431,6 +1438,7 @@ mod tests { filename: file.to_str().unwrap().to_string(), kind: FileOutputKind::FileType, mime_type: Some("application/octet-stream".into()), + upload_filename: None, }) .await .unwrap(); @@ -1464,6 +1472,96 @@ mod tests { } } + #[test] + fn reading_owned_output_deletes_only_that_file() { + let directory = tempfile::tempdir().unwrap(); + let file = directory.path().join("output.json"); + let sibling = directory.path().join("other.json"); + std::fs::write(&file, b"{\"value\":1}").unwrap(); + std::fs::write(&sibling, b"keep").unwrap(); + assert_eq!( + read_output_file(file.to_str().unwrap()).unwrap(), + b"{\"value\":1}" + ); + assert!(!file.exists()); + assert_eq!(std::fs::read(sibling).unwrap(), b"keep"); + } + + #[tokio::test] + async fn upload_preserves_reserved_filename_and_signed_query() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(201)) + .expect(1) + .mount(&server) + .await; + let filename = "result?draft#1%20 final.txt"; + let result = upload_file( + &format!("{}/outputs?token=secret", server.uri()), + filename, + b"data", + "text/plain", + "test", + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests[0].url.path(), + "/outputs/result%3Fdraft%231%2520%20final.txt" + ); + assert_eq!(requests[0].url.query(), Some("token=secret")); + assert_eq!( + result, + format!( + "{}/outputs/result%3Fdraft%231%2520%20final.txt", + server.uri() + ) + ); + } + + #[tokio::test] + async fn file_upload_streams_and_cleans_snapshot_before_done() { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .and(path("/result%3Fdraft.txt")) + .and(header("Content-Type", "text/plain")) + .and(wiremock::matchers::body_bytes(b"snapshot")) + .respond_with(ResponseTemplate::new(201)) + .expect(1) + .mount(&server) + .await; + let root = tempfile::tempdir().unwrap(); + let outputs = tempfile::tempdir().unwrap(); + let file = root.path().join("result?draft.txt"); + std::fs::write(&file, b"snapshot").unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel(); + let sender = crate::worker::SlotSender::new(tx, outputs.path().to_owned()); + sender.send_file_output(file.clone(), None).unwrap(); + assert!(!file.exists()); + let mut prediction = TestPrediction::new(server.uri(), true).await; + let mut events = prediction.prediction.lock().unwrap().subscribe_stream(); + prediction + .writer + .send(rx.try_recv().unwrap()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + events.recv().await.unwrap(); + if !prediction.prediction.lock().unwrap().outputs().is_empty() { + break; + } + } + }) + .await + .expect("file output should stream without Done"); + assert_eq!(std::fs::read_dir(outputs.path()).unwrap().count(), 0); + assert!(!prediction.prediction.lock().unwrap().is_terminal()); + assert_eq!(prediction.finish(true).await["status"], "succeeded"); + } + #[tokio::test] async fn rejected_output_upload_fails_scalar_prediction() { let server = MockServer::start().await; @@ -1484,6 +1582,10 @@ mod tests { "Failed to upload output: upload returned status 503 Service Unavailable" ); assert!(response["output"].is_null()); + assert!( + !file.path().exists(), + "failed uploads must not retain snapshots" + ); } #[tokio::test] @@ -1556,6 +1658,7 @@ mod tests { assert_eq!(response["status"], "succeeded", "{response}"); assert_eq!(response["output"], "https://example.com/result.bin"); + assert!(!file.exists()); assert!(response["error"].is_null()); assert_eq!( server.received_requests().await.unwrap()[0].body, diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 613c058849..30adf06f18 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -162,7 +162,6 @@ type SlotWriter = pub struct SlotSender { tx: mpsc::UnboundedSender, output_dir: PathBuf, - file_counter: Arc, output_counter: Arc, } @@ -171,7 +170,6 @@ impl SlotSender { Self { tx, output_dir, - file_counter: Arc::new(AtomicUsize::new(0)), output_counter: Arc::new(AtomicU64::new(0)), } } @@ -180,12 +178,6 @@ impl SlotSender { self.output_counter.fetch_add(1, Ordering::Relaxed) } - /// Generate a unique 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}")) - } - pub fn send_log(&self, source: LogSource, data: &str) -> io::Result<()> { if data.is_empty() { return Ok(()); @@ -211,28 +203,64 @@ 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) + use std::io::Write; + let mut file = tempfile::Builder::new() + .suffix(&format!(".{extension}")) + .tempfile_in(&self.output_dir)?; + file.write_all(data)?; + self.send_staged_file(file, mime_type, None) } /// 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. + /// Snapshot the file before the predictor resumes, then unlink the source + /// entry. Symlink targets are copied but never unlinked. /// `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 + let upload_filename = path + .file_name() + .and_then(|s| s.to_str()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 filename"))? + .to_owned(); + let mut source = std::fs::File::open(&path)?; + if !source.metadata()?.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "output path must be a regular file", + )); + } + let mut file = tempfile::NamedTempFile::new_in(&self.output_dir)?; + std::io::copy(&mut source, &mut file)?; + std::fs::remove_file(&path)?; + self.send_staged_file(file, mime_type, Some(upload_filename)) + } + + fn send_staged_file( + &self, + file: tempfile::NamedTempFile, + mime_type: Option, + upload_filename: Option, + ) -> io::Result<()> { + let filename = file + .path() .to_str() .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path"))? - .to_string(); + .to_owned(); let msg = SlotResponse::FileOutput { filename, kind: FileOutputKind::FileType, mime_type, + upload_filename, }; - self.tx - .send(msg) - .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) + let (_, path) = file.keep().map_err(|e| e.error)?; + if self.tx.send(msg).is_err() { + std::fs::remove_file(path)?; + return Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "slot channel closed", + )); + } + Ok(()) } /// Send a user metric to the parent process. @@ -277,6 +305,7 @@ fn build_output_message( filename, kind: FileOutputKind::Oversized, mime_type: None, + upload_filename: None, }) } else { Ok(SlotResponse::OutputChunk { output, index }) @@ -1019,6 +1048,94 @@ async fn run_prediction( mod tests { use super::*; + #[test] + fn returned_path_is_consumed_and_each_yield_has_its_own_snapshot() { + let root = tempfile::tempdir().unwrap(); + let outputs = tempfile::tempdir().unwrap(); + let source = root.path().join("result?draft.txt"); + let (tx, mut rx) = mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, outputs.path().to_owned()); + let mut snapshots = Vec::new(); + for contents in [b"first".as_slice(), b"second".as_slice()] { + std::fs::write(&source, contents).unwrap(); + sender.send_file_output(source.clone(), None).unwrap(); + assert!(!source.exists()); + let SlotResponse::FileOutput { + filename, + upload_filename, + .. + } = rx.try_recv().unwrap() + else { + panic!("expected FileOutput"); + }; + assert_eq!(upload_filename.as_deref(), Some("result?draft.txt")); + snapshots.push(filename); + } + assert_ne!(snapshots[0], snapshots[1]); + assert_eq!(std::fs::read(&snapshots[0]).unwrap(), b"first"); + assert_eq!(std::fs::read(&snapshots[1]).unwrap(), b"second"); + assert_eq!( + sender.send_file_output(source, None).unwrap_err().kind(), + io::ErrorKind::NotFound + ); + } + + #[test] + fn output_staging_failure_preserves_source() { + let root = tempfile::tempdir().unwrap(); + let source = root.path().join("result.txt"); + std::fs::write(&source, b"keep me").unwrap(); + let (tx, _rx) = mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, root.path().join("missing")); + assert!(sender.send_file_output(source.clone(), None).is_err()); + assert_eq!(std::fs::read(source).unwrap(), b"keep me"); + } + + #[test] + fn failed_send_does_not_leave_a_snapshot() { + let outputs = tempfile::tempdir().unwrap(); + let (tx, rx) = mpsc::unbounded_channel(); + drop(rx); + let sender = SlotSender::new(tx, outputs.path().to_owned()); + assert!(sender.write_file_output(b"data", "txt", None).is_err()); + assert_eq!(std::fs::read_dir(outputs.path()).unwrap().count(), 0); + } + + #[cfg(unix)] + #[test] + fn returned_symlink_consumes_link_but_preserves_target() { + let root = tempfile::tempdir().unwrap(); + let target = root.path().join("asset.txt"); + let link = root.path().join("result.txt"); + std::fs::write(&target, b"asset").unwrap(); + std::os::unix::fs::symlink("asset.txt", &link).unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, root.path().to_owned()); + sender.send_file_output(link.clone(), None).unwrap(); + assert!(!link.exists()); + assert_eq!(std::fs::read(target).unwrap(), b"asset"); + let SlotResponse::FileOutput { filename, .. } = rx.try_recv().unwrap() else { + panic!("expected FileOutput"); + }; + assert_eq!(std::fs::read(filename).unwrap(), b"asset"); + } + + #[test] + fn independent_senders_sharing_output_directory_do_not_overwrite() { + let outputs = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel(); + for contents in [b"first".as_slice(), b"second".as_slice()] { + let sender = SlotSender::new(tx.clone(), outputs.path().to_owned()); + sender.write_file_output(contents, "txt", None).unwrap(); + } + for contents in [b"first".as_slice(), b"second".as_slice()] { + let SlotResponse::FileOutput { filename, .. } = rx.try_recv().unwrap() else { + panic!("expected FileOutput"); + }; + assert_eq!(std::fs::read(filename).unwrap(), contents); + } + } + #[test] fn predict_result_success() { let r = PredictResult::success(serde_json::json!("hello"), 0.5, false); diff --git a/docs/llms.txt b/docs/llms.txt index 90700b8a6c..7b7b799957 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3358,6 +3358,8 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. +Returning or yielding a path transfers ownership of that file to Cog. Cog copies its contents and deletes the source entry before continuing, including files outside temporary directories. Return a copy if you need to keep the original. A symlink output deletes the link, not its target. To return the same path again, write a new file there first. Cog does not delete the containing directory. + This example takes an input file, resizes it, and returns the resized image: ```python diff --git a/docs/python.md b/docs/python.md index 069a17abe8..6e49be3ffb 100644 --- a/docs/python.md +++ b/docs/python.md @@ -548,6 +548,8 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. +Returning or yielding a path transfers ownership of that file to Cog. Cog copies its contents and deletes the source entry before continuing, including files outside temporary directories. Return a copy if you need to keep the original. A symlink output deletes the link, not its target. To return the same path again, write a new file there first. Cog does not delete the containing directory. + This example takes an input file, resizes it, and returns the resized image: ```python From 341112927b4ba30c85b5fe95e9b7f2a5c6b998f5 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Sat, 31 Oct 2026 17:49:22 -0400 Subject: [PATCH 2/3] fix(coglet): only delete returned paths in the temp dir Returned paths are now copied and then deleted only when their directory resolves under the system temp dir. Baked-in assets, weights, and mounted files are left alone, and a failed delete logs a warning instead of failing the prediction. - Release list, tuple, and async-generator paths after the last item so `[p, p]` works. Generators still release before resuming. - Check for a regular file before opening, on top-level and nested paths. - Copy outside the GIL. - Replace control characters and dot-only names in the upload segment so a basename can't move the PUT to the parent path. - Make tempfile a workspace dependency. - Add integration tests for path ownership and upload names, and escape the mock upload server's Location path. --- architecture/04-container-runtime.md | 2 +- crates/Cargo.toml | 3 + crates/coglet-python/Cargo.toml | 2 +- crates/coglet-python/src/output.rs | 94 +++++++++++++++--- crates/coglet-python/src/predictor.rs | 42 +++++--- crates/coglet-python/tests/test_coglet.py | 38 ++++++- crates/coglet/Cargo.toml | 4 +- crates/coglet/src/orchestrator.rs | 62 +++++++++++- crates/coglet/src/worker.rs | 98 +++++++++++++++++-- docs/llms.txt | 4 +- docs/python.md | 4 +- integration-tests/harness/harness.go | 2 +- .../tests/coglet_path_output_ownership.txtar | 39 ++++++++ .../coglet_path_output_upload_name.txtar | 31 ++++++ 14 files changed, 379 insertions(+), 46 deletions(-) create mode 100644 integration-tests/tests/coglet_path_output_ownership.txtar create mode 100644 integration-tests/tests/coglet_path_output_upload_name.txtar diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 3bfce4f291..ed7bd8a9bb 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -373,7 +373,7 @@ When a prediction input exceeds 6MiB, it's too large to send inline through the 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. -Returning or yielding a path transfers ownership to Cog. The worker snapshots its contents into a unique temporary file and unlinks the source entry before resuming user code. This prevents a later overwrite from changing an earlier streamed output. The parent deletes the snapshot after reading it; uploads use the bytes in memory. Nested paths are encoded inline and unlinked after the containing output is encoded. Cleanup removes files only, never prediction directories derived from public IDs. +The worker snapshots a returned or yielded path into a unique temporary file in `output_dir` before resuming user code, so a later overwrite can't change an earlier streamed output. After the snapshot is sent, the worker deletes the source only if its directory resolves under the system temp dir; files elsewhere are left alone. Paths in a list, tuple, or nested output are deleted after the whole output is sent, so the same path can appear more than once. The parent deletes the snapshot right after reading it, before the upload starts; uploads use the bytes in memory. Cleanup removes files only, never directories. `FileOutputKind` distinguishes between normal file outputs (`FileType`) and oversized outputs (`Oversized`) that exceeded an inline size limit. diff --git a/crates/Cargo.toml b/crates/Cargo.toml index f6376844ce..5e5df71d5b 100644 --- a/crates/Cargo.toml +++ b/crates/Cargo.toml @@ -34,6 +34,9 @@ serde_json = "1" # Identifiers uuid = { version = "1", features = ["v4", "serde"] } +# Filesystem +tempfile = "3" + # Observability tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } diff --git a/crates/coglet-python/Cargo.toml b/crates/coglet-python/Cargo.toml index 779361b96d..1db0cee4fb 100644 --- a/crates/coglet-python/Cargo.toml +++ b/crates/coglet-python/Cargo.toml @@ -31,7 +31,7 @@ libc = "0.2" [dev-dependencies] pyo3 = { workspace = true, features = ["auto-initialize"] } -tempfile = "3" +tempfile.workspace = true [features] default = ["tracing", "tracing-grpc"] diff --git a/crates/coglet-python/src/output.rs b/crates/coglet-python/src/output.rs index a1128686a9..a0d9a86a11 100644 --- a/crates/coglet-python/src/output.rs +++ b/crates/coglet-python/src/output.rs @@ -12,37 +12,45 @@ //! //! This replaces the Python modules cog.json and cog.files. +use std::collections::HashSet; +use std::path::PathBuf; + +use coglet_core::worker::release_output_path; 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. +/// remaining Path/IOBase objects to base64 data URLs. Encoded paths are +/// released once the whole output is encoded, so repeated paths work. pub fn process_output<'py>( py: Python<'py>, output: &Bound<'py, PyAny>, ) -> PyResult> { - let encodeable = make_encodeable(py, output)?; - let mut paths = std::collections::HashSet::new(); - let encoded = encode_files(py, &encodeable, &mut paths)?; - // Wait until every occurrence is encoded so repeated paths in one output work. - for path in paths { - if let Err(error) = std::fs::remove_file(path) - && error.kind() != std::io::ErrorKind::NotFound - { - return Err(error.into()); - } - } + let mut paths = HashSet::new(); + let encoded = process_output_item(py, output, &mut paths)?; + release_output_paths(paths); Ok(encoded) } -/// Process a single output item (for generator outputs). +/// Process a single output item, collecting encoded paths instead of +/// releasing them. Callers release `paths` after the last item that might +/// reference them. pub fn process_output_item<'py>( py: Python<'py>, item: &Bound<'py, PyAny>, + paths: &mut HashSet, ) -> PyResult> { - process_output(py, item) + let encodeable = make_encodeable(py, item)?; + encode_files(py, &encodeable, paths) +} + +/// Release every collected output path. See [`release_output_path`]. +pub fn release_output_paths(paths: HashSet) { + for path in paths { + release_output_path(&path); + } } /// Normalize a Python object into a JSON-friendly form. @@ -155,7 +163,7 @@ fn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult( py: Python<'py>, obj: &Bound<'py, PyAny>, - paths: &mut std::collections::HashSet, + paths: &mut HashSet, ) -> PyResult> { // str -- return as-is (don't recurse into characters) if obj.is_instance_of::() { @@ -185,12 +193,20 @@ fn encode_files<'py>( let os_mod = py.import("os")?; let pathlike_cls = os_mod.getattr("PathLike")?; if obj.is_instance(&pathlike_cls)? { + let path = obj.extract::()?; + // Check before opening so a FIFO or device can't block the open. + if !std::fs::metadata(&path)?.is_file() { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "output path must be a regular file: {}", + path.display() + ))); + } let builtins = py.import("builtins")?; let fh = builtins.getattr("open")?.call1((obj, "rb"))?; let result = file_to_base64(py, &fh); fh.call_method0("close")?; let encoded = result?; - paths.insert(obj.extract::()?); + paths.insert(path); return Ok(encoded); } @@ -526,6 +542,52 @@ mod tests { assert_eq!(std::fs::read(target).unwrap(), b"asset"); } + #[test] + fn nested_path_outside_temp_dir_is_kept() { + let root = tempfile::Builder::new() + .tempdir_in(env!("CARGO_MANIFEST_DIR")) + .unwrap(); + let file = root.path().join("asset.txt"); + std::fs::write(&file, b"asset").unwrap(); + pyo3::Python::initialize(); + Python::attach(|py| { + let path = py + .import("pathlib") + .unwrap() + .getattr("Path") + .unwrap() + .call1((&file,)) + .unwrap(); + let obj = PyDict::new(py); + obj.set_item("file", &path).unwrap(); + process_output(py, obj.as_any()).unwrap(); + }); + assert_eq!(std::fs::read(file).unwrap(), b"asset"); + } + + #[cfg(unix)] + #[test] + fn nested_fifo_is_rejected_without_blocking() { + let root = tempfile::tempdir().unwrap(); + let fifo = root.path().join("pipe"); + let c_path = std::ffi::CString::new(fifo.as_os_str().as_encoded_bytes()).unwrap(); + assert_eq!(unsafe { libc::mkfifo(c_path.as_ptr(), 0o600) }, 0); + pyo3::Python::initialize(); + Python::attach(|py| { + let path = py + .import("pathlib") + .unwrap() + .getattr("Path") + .unwrap() + .call1((&fifo,)) + .unwrap(); + let obj = PyDict::new(py); + obj.set_item("file", &path).unwrap(); + assert!(process_output(py, obj.as_any()).is_err()); + }); + assert!(fifo.exists()); + } + #[test] fn encode_pathlike_to_base64() { let result = processed( diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index ae09c60725..089b5b06c3 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -1,12 +1,14 @@ //! Python predictor loading and invocation. +use std::collections::HashSet; +use std::path::PathBuf; use std::sync::{Arc, OnceLock}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; -use coglet_core::worker::SlotSender; +use coglet_core::worker::{SlotSender, release_output_path}; use coglet_core::{PredictionError, PredictionOutput, PredictionResult}; use crate::cancel; @@ -212,14 +214,18 @@ fn submit_async_coroutine( /// Send a single output item over IPC, routing file outputs to disk. /// -/// For Path outputs (os.PathLike): transfers ownership via send_file_output. +/// For Path outputs (os.PathLike): snapshots the file via send_file_output. /// 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. +/// +/// Sent paths are added to `paths`; the caller releases them once no later +/// item can reference them. fn send_output_item( py: Python<'_>, item: &Bound<'_, PyAny>, json_module: &Bound<'_, PyAny>, slot_sender: &SlotSender, + paths: &mut HashSet, ) -> Result<(), PredictionError> { let os = py .import("os") @@ -235,14 +241,15 @@ fn send_output_item( .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if item.is_instance(&pathlike).unwrap_or(false) { - // Snapshot and consume the returned path before the predictor resumes. 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) + let path = PathBuf::from(path_str); + // Copying a large file shouldn't block other slots on the GIL. + py.detach(|| slot_sender.send_file_output(path.clone(), None)) .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; + paths.insert(path); return Ok(()); } @@ -281,7 +288,7 @@ fn send_output_item( } // Non-file output - process normally - let processed = output::process_output_item(py, item) + let processed = output::process_output_item(py, item, paths) .map_err(|e| PredictionError::Failed(format!("Failed to process output item: {}", e)))?; let item_str: String = json_module @@ -981,7 +988,10 @@ impl PythonPredictor { PredictionError::Failed(format!("Generator iteration error: {}", e)) })?; - send_output_item(py, &item, json_module, slot_sender)?; + // Release before the generator resumes so it can reuse the path. + let mut paths = HashSet::new(); + send_output_item(py, &item, json_module, slot_sender, &mut paths)?; + output::release_output_paths(paths); } // Outputs already streamed over IPC — return empty stream @@ -1018,11 +1028,12 @@ impl PythonPredictor { .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) + let path = PathBuf::from(path_str); + py.detach(|| slot_sender.send_file_output(path.clone(), None)) .map_err(|e| { PredictionError::Failed(format!("Failed to send file output: {}", e)) })?; + release_output_path(&path); return Ok(PredictionOutput::Single(serde_json::Value::Null)); } @@ -1060,16 +1071,21 @@ impl PythonPredictor { // 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. + // Paths are released after the last item so `[p, p]` works. if let Ok(list) = result.cast::() { + let mut paths = HashSet::new(); for item in list.iter() { - send_output_item(py, &item, json_module, slot_sender)?; + send_output_item(py, &item, json_module, slot_sender, &mut paths)?; } + output::release_output_paths(paths); return Ok(PredictionOutput::Stream(vec![])); } if let Ok(tuple) = result.cast::() { + let mut paths = HashSet::new(); for item in tuple.iter() { - send_output_item(py, &item, json_module, slot_sender)?; + send_output_item(py, &item, json_module, slot_sender, &mut paths)?; } + output::release_output_paths(paths); return Ok(PredictionOutput::Stream(vec![])); } @@ -1198,9 +1214,11 @@ impl PythonPredictor { let output = if is_async_gen { // Result is a pre-collected list — stream each item over IPC if let Ok(list) = result.extract::>>() { + let mut paths = HashSet::new(); for item in list { - send_output_item(py, &item, &json_module, slot_sender)?; + send_output_item(py, &item, &json_module, slot_sender, &mut paths)?; } + output::release_output_paths(paths); } PredictionOutput::Stream(vec![]) } else { diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 214d290162..c42c32d4b5 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -5,6 +5,7 @@ import socket import subprocess import sys +import tempfile import threading import time from pathlib import Path @@ -520,7 +521,7 @@ def test_includes_predict_time(self, sync_predictor: Path): assert result["metrics"]["predict_time"] >= 0 -@pytest.mark.parametrize("shape", ["scalar", "nested", "generator"]) +@pytest.mark.parametrize("shape", ["scalar", "nested", "generator", "list"]) def test_returned_paths_are_consumed(tmp_path: Path, shape: str) -> None: predictor = tmp_path / "predict.py" predictor.write_text(""" @@ -540,6 +541,12 @@ def predict(self, write: bool = True) -> Path: .replace("-> Path:", "-> dict[str, Path]:") .replace("return path", 'return {"file": path}') ) + elif shape == "list": + predictor.write_text( + predictor.read_text() + .replace("-> Path:", "-> list[Path]:") + .replace("return path", "return [path, path]") + ) elif shape == "generator": predictor.write_text( "from typing import Iterator\n" @@ -559,12 +566,41 @@ def predict(self, write: bool = True) -> Path: output = output["file"] elif shape == "generator": output = output[0] + elif shape == "list": + assert output[0] == output[1], output + output = output[0] assert output == "data:text/plain;base64,aGVsbG8=" assert not (tmp_path / "result.txt").exists() result = server.predict({"write": False}) assert result["status"] == "failed", result +@pytest.mark.parametrize("nested", [False, True]) +def test_returned_paths_outside_temp_dir_are_kept(tmp_path: Path, nested: bool) -> None: + # pytest's tmp_path lives in the system temp dir, so the asset goes next to + # this file to stand in for something baked into the image. + with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as asset_dir: + asset = Path(asset_dir) / "asset.txt" + asset.write_text("hello") + output = '{"file": asset}' if nested else "asset" + predictor = tmp_path / "predict.py" + predictor.write_text(f""" +from pathlib import Path +from cog import BasePredictor + +class Predictor(BasePredictor): + def predict(self) -> {"dict[str, Path]" if nested else "Path"}: + asset = Path({str(asset)!r}) + return {output} +""") + (tmp_path / "cog.yaml").write_text('predict: "predict.py:Predictor"\n') + with CogletServer(predictor) as server: + for _ in range(2): + result = server.predict({}) + assert result["status"] == "succeeded", result + assert asset.read_text() == "hello" + + class TestSecretInput: """Tests for cog.Secret input coercion.""" diff --git a/crates/coglet/Cargo.toml b/crates/coglet/Cargo.toml index a935eb3fb0..e32e5f2e30 100644 --- a/crates/coglet/Cargo.toml +++ b/crates/coglet/Cargo.toml @@ -46,7 +46,9 @@ mime_guess = "2.0.5" # Identifiers uuid.workspace = true -tempfile = "3" + +# Filesystem +tempfile.workspace = true # HTTP server axum.workspace = true diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 0655fa354b..e94a61e918 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -34,6 +34,9 @@ use crate::prediction::Prediction; const MAX_PENDING_CANCELLATIONS: usize = 1000; /// Read an owned output snapshot and release its disk space before uploading. +/// +/// The snapshot is deleted as soon as it is read, before the upload starts, so +/// failed and canceled uploads have nothing left to clean up. fn read_output_file(filename: &str) -> std::io::Result> { let result = std::fs::read(filename); if let Err(error) = std::fs::remove_file(filename) @@ -44,12 +47,28 @@ fn read_output_file(filename: &str) -> std::io::Result> { result } +/// Make a filename safe to push as a single URL path segment. +/// +/// URL parsing drops tabs and newlines and treats `.`/`..` as dot segments, so +/// a name like `".\n."` would otherwise move the upload to the parent path. +fn upload_segment(filename: &str) -> String { + let segment: String = filename + .chars() + .map(|c| if c.is_control() { '_' } else { c }) + .collect(); + match segment.as_str() { + "" | "." | ".." => "output".to_string(), + _ => segment, + } +} + /// Upload a file to a signed endpoint, returning the final URL. /// /// Matches the behavior of Python cog's `put_file_to_signed_endpoint`: -/// PUT to `{endpoint}{filename}` with Content-Type header, then extract -/// the final URL from the Location header (falling back to response URL), -/// stripping query parameters. Follows redirects automatically. +/// PUT to the endpoint with `filename` appended as one escaped path segment +/// (the endpoint's query string is kept), then extract the final URL from the +/// Location header (falling back to response URL), stripping query parameters. +/// Follows redirects automatically. async fn upload_file( endpoint: &str, filename: &str, @@ -59,9 +78,9 @@ async fn upload_file( ) -> Result { let mut url = reqwest::Url::parse(endpoint).map_err(|_| "invalid upload URL".to_string())?; url.path_segments_mut() - .map_err(|_| "upload URL cannot contain path segments".to_string())? + .map_err(|_| "upload URL does not support path segments".to_string())? .pop_if_empty() - .push(filename); + .push(&upload_segment(filename)); let client = reqwest::Client::new(); let resp = client .put(url) @@ -1521,6 +1540,38 @@ mod tests { ); } + #[test] + fn upload_segment_cannot_escape_its_segment() { + assert_eq!(upload_segment(".\n."), "._."); + assert_eq!(upload_segment("a\tb.png"), "a_b.png"); + assert_eq!(upload_segment(""), "output"); + assert_eq!(upload_segment(".."), "output"); + assert_eq!(upload_segment("result.png"), "result.png"); + } + + #[tokio::test] + async fn upload_with_control_characters_stays_under_endpoint() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(201)) + .expect(1) + .mount(&server) + .await; + upload_file( + &format!("{}/bucket/pred/?token=secret", server.uri()), + ".\n.", + b"data", + "text/plain", + "test", + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests[0].url.path(), "/bucket/pred/._."); + assert_eq!(requests[0].url.query(), Some("token=secret")); + } + #[tokio::test] async fn file_upload_streams_and_cleans_snapshot_before_done() { let server = MockServer::start().await; @@ -1539,6 +1590,7 @@ mod tests { let (tx, mut rx) = mpsc::unbounded_channel(); let sender = crate::worker::SlotSender::new(tx, outputs.path().to_owned()); sender.send_file_output(file.clone(), None).unwrap(); + crate::worker::release_output_path(&file); assert!(!file.exists()); let mut prediction = TestPrediction::new(server.uri(), true).await; let mut events = prediction.prediction.lock().unwrap().subscribe_stream(); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 30adf06f18..e85fad60fd 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -213,8 +213,9 @@ impl SlotSender { /// Send a file-typed output (e.g. Path, File return types). /// - /// Snapshot the file before the predictor resumes, then unlink the source - /// entry. Symlink targets are copied but never unlinked. + /// Snapshots the file into the output dir. The source is left in place; + /// callers pass it to [`release_output_path`] once every reference to it + /// has been sent. /// `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 upload_filename = path @@ -222,16 +223,16 @@ impl SlotSender { .and_then(|s| s.to_str()) .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 filename"))? .to_owned(); - let mut source = std::fs::File::open(&path)?; - if !source.metadata()?.is_file() { + // Check before opening so a FIFO or device can't block the open. + if !std::fs::metadata(&path)?.is_file() { return Err(io::Error::new( io::ErrorKind::InvalidInput, "output path must be a regular file", )); } + let mut source = std::fs::File::open(&path)?; let mut file = tempfile::NamedTempFile::new_in(&self.output_dir)?; std::io::copy(&mut source, &mut file)?; - std::fs::remove_file(&path)?; self.send_staged_file(file, mime_type, Some(upload_filename)) } @@ -285,6 +286,38 @@ impl SlotSender { } } +/// Delete a returned output path once its contents have been sent. +/// +/// Only files whose directory resolves under the system temp dir are deleted, +/// so baked-in assets, weights, and mounted files are never touched. The +/// directory is resolved rather than the file, so a returned symlink removes +/// the link and keeps its target. Failures are logged, not returned, because +/// the output has already been sent. +pub fn release_output_path(path: &std::path::Path) { + if !is_in_temp_dir(path) { + return; + } + if let Err(error) = std::fs::remove_file(path) + && error.kind() != io::ErrorKind::NotFound + { + tracing::warn!(path = %path.display(), %error, "Failed to remove output file"); + } +} + +fn is_in_temp_dir(path: &std::path::Path) -> bool { + let parent = match path.parent() { + Some(parent) if !parent.as_os_str().is_empty() => parent, + _ => std::path::Path::new("."), + }; + match ( + std::fs::canonicalize(parent), + std::fs::canonicalize(std::env::temp_dir()), + ) { + (Ok(parent), Ok(temp_dir)) => parent.starts_with(temp_dir), + _ => false, + } +} + /// Build an output message, spilling to disk if larger than the IPC frame limit. fn build_output_message( output_dir: &std::path::Path, @@ -1059,6 +1092,7 @@ mod tests { for contents in [b"first".as_slice(), b"second".as_slice()] { std::fs::write(&source, contents).unwrap(); sender.send_file_output(source.clone(), None).unwrap(); + release_output_path(&source); assert!(!source.exists()); let SlotResponse::FileOutput { filename, @@ -1080,6 +1114,48 @@ mod tests { ); } + #[test] + fn release_keeps_files_outside_temp_dir() { + let root = tempfile::Builder::new() + .tempdir_in(env!("CARGO_MANIFEST_DIR")) + .unwrap(); + let asset = root.path().join("asset.txt"); + std::fs::write(&asset, b"asset").unwrap(); + release_output_path(&asset); + assert_eq!(std::fs::read(asset).unwrap(), b"asset"); + } + + #[cfg(unix)] + #[test] + fn release_keeps_files_behind_symlinked_directory() { + let outside = tempfile::Builder::new() + .tempdir_in(env!("CARGO_MANIFEST_DIR")) + .unwrap(); + let temp = tempfile::tempdir().unwrap(); + let link = temp.path().join("outputs"); + std::os::unix::fs::symlink(outside.path(), &link).unwrap(); + std::fs::write(outside.path().join("asset.txt"), b"asset").unwrap(); + release_output_path(&link.join("asset.txt")); + assert_eq!( + std::fs::read(outside.path().join("asset.txt")).unwrap(), + b"asset" + ); + } + + #[cfg(unix)] + #[test] + fn fifo_output_is_rejected_without_blocking() { + let root = tempfile::tempdir().unwrap(); + let fifo = root.path().join("pipe"); + nix::unistd::mkfifo(&fifo, nix::sys::stat::Mode::S_IRWXU).unwrap(); + let (tx, _rx) = mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, root.path().to_owned()); + assert_eq!( + sender.send_file_output(fifo, None).unwrap_err().kind(), + io::ErrorKind::InvalidInput + ); + } + #[test] fn output_staging_failure_preserves_source() { let root = tempfile::tempdir().unwrap(); @@ -1112,6 +1188,7 @@ mod tests { let (tx, mut rx) = mpsc::unbounded_channel(); let sender = SlotSender::new(tx, root.path().to_owned()); sender.send_file_output(link.clone(), None).unwrap(); + release_output_path(&link); assert!(!link.exists()); assert_eq!(std::fs::read(target).unwrap(), b"asset"); let SlotResponse::FileOutput { filename, .. } = rx.try_recv().unwrap() else { @@ -1124,11 +1201,20 @@ mod tests { fn independent_senders_sharing_output_directory_do_not_overwrite() { let outputs = tempfile::tempdir().unwrap(); let (tx, mut rx) = mpsc::unbounded_channel(); + let root = tempfile::tempdir().unwrap(); + let source = root.path().join("result.txt"); for contents in [b"first".as_slice(), b"second".as_slice()] { let sender = SlotSender::new(tx.clone(), outputs.path().to_owned()); sender.write_file_output(contents, "txt", None).unwrap(); + std::fs::write(&source, contents).unwrap(); + sender.send_file_output(source.clone(), None).unwrap(); } - for contents in [b"first".as_slice(), b"second".as_slice()] { + for contents in [ + b"first".as_slice(), + b"first".as_slice(), + b"second".as_slice(), + b"second".as_slice(), + ] { let SlotResponse::FileOutput { filename, .. } = rx.try_recv().unwrap() else { panic!("expected FileOutput"); }; diff --git a/docs/llms.txt b/docs/llms.txt index 7b7b799957..b74dfc8b59 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3358,7 +3358,9 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. -Returning or yielding a path transfers ownership of that file to Cog. Cog copies its contents and deletes the source entry before continuing, including files outside temporary directories. Return a copy if you need to keep the original. A symlink output deletes the link, not its target. To return the same path again, write a new file there first. Cog does not delete the containing directory. +Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. + +Don't build a returned path from an input value. A caller could point it at any file in the container. This example takes an input file, resizes it, and returns the resized image: diff --git a/docs/python.md b/docs/python.md index 6e49be3ffb..431576249e 100644 --- a/docs/python.md +++ b/docs/python.md @@ -548,7 +548,9 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. -Returning or yielding a path transfers ownership of that file to Cog. Cog copies its contents and deletes the source entry before continuing, including files outside temporary directories. Return a copy if you need to keep the original. A symlink output deletes the link, not its target. To return the same path again, write a new file there first. Cog does not delete the containing directory. +Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. + +Don't build a returned path from an input value. A caller could point it at any file in the container. This example takes an input file, resizes it, and returns the resized image: diff --git a/integration-tests/harness/harness.go b/integration-tests/harness/harness.go index 6512897783..352c7a2dd5 100644 --- a/integration-tests/harness/harness.go +++ b/integration-tests/harness/harness.go @@ -1130,7 +1130,7 @@ func (h *Harness) cmdUploadServerStart(ts *testscript.TestScript, neg bool, args mus.mu.Unlock() // Return a clean URL without query params (simulates a signed URL redirect) - location := fmt.Sprintf("http://host.docker.internal:%d%s", mus.port, r.URL.Path) + location := fmt.Sprintf("http://host.docker.internal:%d%s", mus.port, r.URL.EscapedPath()) w.Header().Set("Location", location) w.WriteHeader(http.StatusOK) }) diff --git a/integration-tests/tests/coglet_path_output_ownership.txtar b/integration-tests/tests/coglet_path_output_ownership.txtar new file mode 100644 index 0000000000..28bf022e0d --- /dev/null +++ b/integration-tests/tests/coglet_path_output_ownership.txtar @@ -0,0 +1,39 @@ +# Test which returned paths Cog deletes after sending them. +# A temp file is deleted before the generator resumes, so the same path can be +# rewritten for the next yield. A file baked into the image is never deleted, +# so the second prediction can return it again. + +cog build -t $TEST_IMAGE +cog serve + +curl POST /predictions '{"input":{}}' +stdout '"status":"succeeded"' +# "first", "second", "asset\n" as base64 data URLs +stdout '"output":\["data:text/plain;base64,Zmlyc3Q=","data:text/plain;base64,c2Vjb25k","data:text/plain;base64,YXNzZXQK"\]' + +curl POST /predictions '{"input":{}}' +stdout '"status":"succeeded"' +stdout 'data:text/plain;base64,YXNzZXQK' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- asset.txt -- +asset +-- predict.py -- +import tempfile +from typing import Iterator + +from cog import BasePredictor, Path + + +class Predictor(BasePredictor): + def predict(self) -> Iterator[Path]: + path = Path(tempfile.mkdtemp()) / "out.txt" + for text in ["first", "second"]: + path.write_text(text) + yield path + assert not path.exists(), "temp output was not deleted" + yield Path("/src/asset.txt") diff --git a/integration-tests/tests/coglet_path_output_upload_name.txtar b/integration-tests/tests/coglet_path_output_upload_name.txtar new file mode 100644 index 0000000000..87cc36ba38 --- /dev/null +++ b/integration-tests/tests/coglet_path_output_upload_name.txtar @@ -0,0 +1,31 @@ +# Test that an uploaded output keeps its basename as one escaped path segment +# and that a query string on --upload-url (as with signed URLs) is preserved. + +cog build -t $TEST_IMAGE + +upload-server-start + +cog serve --upload-url ${UPLOAD_SERVER_URL}?token=abc + +curl POST /predictions '{"input":{}}' +stdout '"status":"succeeded"' +stdout '"output":"http://host.docker.internal:[0-9]+/my%20result%20%231\.txt"' + +upload-server-count 1 + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +import tempfile + +from cog import BasePredictor, Path + + +class Predictor(BasePredictor): + def predict(self) -> Path: + path = Path(tempfile.mkdtemp()) / "my result #1.txt" + path.write_text("hello") + return path From 50abc7ca98abd4e9b1c6ad17a9023c46cb2b8b05 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Sat, 31 Oct 2026 17:56:54 -0400 Subject: [PATCH 3/3] docs(coglet): note async generators send outputs after finishing Async generators are collected before any item is sent, so the per-yield cleanup only applies to sync generators. Behavior is unchanged from main. --- architecture/04-container-runtime.md | 2 +- docs/llms.txt | 2 +- docs/python.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index ed7bd8a9bb..a89bd2106c 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -373,7 +373,7 @@ When a prediction input exceeds 6MiB, it's too large to send inline through the 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. -The worker snapshots a returned or yielded path into a unique temporary file in `output_dir` before resuming user code, so a later overwrite can't change an earlier streamed output. After the snapshot is sent, the worker deletes the source only if its directory resolves under the system temp dir; files elsewhere are left alone. Paths in a list, tuple, or nested output are deleted after the whole output is sent, so the same path can appear more than once. The parent deletes the snapshot right after reading it, before the upload starts; uploads use the bytes in memory. Cleanup removes files only, never directories. +The worker snapshots a returned or yielded path into a unique temporary file in `output_dir` before resuming user code, so a later overwrite can't change an earlier streamed output. Async generators are collected before any item is sent, so their paths are snapshotted and deleted after the generator finishes. After the snapshot is sent, the worker deletes the source only if its directory resolves under the system temp dir; files elsewhere are left alone. Paths in a list, tuple, or nested output are deleted after the whole output is sent, so the same path can appear more than once. The parent deletes the snapshot right after reading it, before the upload starts; uploads use the bytes in memory. Cleanup removes files only, never directories. `FileOutputKind` distinguishes between normal file outputs (`FileType`) and oversized outputs (`Oversized`) that exceeded an inline size limit. diff --git a/docs/llms.txt b/docs/llms.txt index b74dfc8b59..15d2499fef 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -3358,7 +3358,7 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. -Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. +Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a sync generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. An async generator runs to completion before any output is sent, so give each `yield` its own path. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. Don't build a returned path from an input value. A caller could point it at any file in the container. diff --git a/docs/python.md b/docs/python.md index 431576249e..2ed97f6e16 100644 --- a/docs/python.md +++ b/docs/python.md @@ -548,7 +548,7 @@ These types can be used directly as input parameter types and output return type For models that return a `cog.Path` object, the output returned by Cog's built-in HTTP server will be a URL. -Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. +Cog copies the contents of a returned or yielded path. If the file is in the system temp directory (for example, one made with `tempfile`, or a downloaded input file), Cog then deletes it. For a sync generator, this happens before the generator resumes, so you can write to the same path again for the next `yield`. An async generator runs to completion before any output is sent, so give each `yield` its own path. Files anywhere else, such as assets baked into the image or model weights, are never deleted. A symlink in the temp directory is removed but its target is kept. Cog does not delete directories. Don't build a returned path from an input value. A caller could point it at any file in the container.