From 34506f9cec624ddc41f7d3f576e80bfde4be2e01 Mon Sep 17 00:00:00 2001 From: asahoo Date: Tue, 7 Jul 2026 16:09:56 -0500 Subject: [PATCH 1/8] fix: delete coglet-managed output files after upload Coglet creates output files (IOBase writes, oversized JSON spills) in /tmp/coglet/predictions/{id}/outputs/ but never deletes them. This can cause stale-output bugs where models return outputs from previous predictions if an overwrite fails (issue #1434). Add a managed bool field to the FileOutput IPC protocol message so the orchestrator knows which files are safe to delete (coglet-created) vs user-authored Path outputs (returned by reference, must not be deleted). The orchestrator deletes managed files immediately after reading their bytes into memory. A backstop remove_dir_all in remove_prediction cleans up the entire prediction directory to catch files from aborted uploads or cancelled predictions. Closes #1434 --- crates/coglet-python/src/predictor.rs | 4 +- crates/coglet/src/bridge/protocol.rs | 37 +++++++++++++++++ ...__slot_file_output_managed_serializes.snap | 13 ++++++ ...slot_file_output_oversized_serializes.snap | 12 ++++++ ...slot_file_output_unmanaged_serializes.snap | 12 ++++++ crates/coglet/src/orchestrator.rs | 7 +++- crates/coglet/src/service.rs | 40 +++++++++++++++++-- crates/coglet/src/worker.rs | 13 +++++- 8 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6142e48e24..69de338bb3 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -195,7 +195,7 @@ fn send_output_item( .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) + .send_file_output(std::path::PathBuf::from(path_str), None, false) .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; return Ok(()); } @@ -954,7 +954,7 @@ impl PythonPredictor { .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) + .send_file_output(std::path::PathBuf::from(path_str), None, false) .map_err(|e| { PredictionError::Failed(format!("Failed to send file output: {}", e)) })?; diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 0ef46a4e68..68d30e828c 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -327,6 +327,10 @@ 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 created this file (IOBase write or oversized spill). + /// False for user-authored Path outputs — must not be deleted. + #[serde(default)] + managed: bool, }, /// Streaming output chunk for generator and iterator output. @@ -592,6 +596,39 @@ 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, + }; + 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, + }; + insta::assert_json_snapshot!(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, + }; + 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..9acbbb9cd5 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap @@ -0,0 +1,13 @@ +--- +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 +} 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..447bfeefab --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap @@ -0,0 +1,12 @@ +--- +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 +} 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..3a8994bf9b 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -1072,7 +1072,7 @@ async fn run_event_loop( predictions.remove(&slot_id); } } - Ok(SlotResponse::FileOutput { filename, kind, mime_type }) => { + Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); let bytes = match std::fs::read(&filename) { Ok(b) => b, @@ -1081,6 +1081,11 @@ async fn run_event_loop( continue; } }; + if managed + && let Err(e) = std::fs::remove_file(&filename) + { + tracing::debug!(%slot_id, %filename, error = %e, "Failed to delete managed output file"); + } match kind { FileOutputKind::Oversized => { let output: serde_json::Value = match serde_json::from_slice(&bytes) { diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 460a09e09a..0e6c4b1224 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -657,8 +657,7 @@ impl PredictionService { .await; // Create per-prediction dirs for file-based inputs/outputs - let prediction_dir = - std::path::PathBuf::from("/tmp/coglet/predictions").join(&prediction_id); + 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) @@ -797,9 +796,20 @@ impl PredictionService { } } - /// Remove a prediction from the DashMap after completion. + /// Remove a prediction from the DashMap after completion and clean up + /// its on-disk prediction directory as a backstop for any output files + /// that were not deleted individually (e.g. aborted uploads, cancelled + /// predictions). pub fn remove_prediction(&self, id: &str) { self.predictions.remove(id); + let dir = prediction_dir_for(id); + match std::fs::remove_dir_all(&dir) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir") + } + } } pub fn trigger_shutdown(&self) { @@ -811,6 +821,10 @@ impl PredictionService { } } +fn prediction_dir_for(id: &str) -> std::path::PathBuf { + std::path::PathBuf::from("/tmp/coglet/predictions").join(id) +} + fn spawn_orchestrator_cancel(orch: Arc, id: String) { let Ok(handle) = tokio::runtime::Handle::try_current() else { tracing::warn!(prediction_id = %id, "No tokio runtime available to cancel prediction"); @@ -1723,6 +1737,26 @@ mod tests { assert!(!svc.prediction_exists("test-remove")); } + #[test] + fn remove_prediction_deletes_prediction_dir() { + let svc = PredictionService::new_no_pool(); + let dir = prediction_dir_for("test-dir-cleanup"); + 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("test-dir-cleanup"); + + 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("nonexistent-prediction"); + } + #[test] fn build_slot_request_small_input_inline() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index c59a26cc7a..bc0645b62a 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -201,14 +201,21 @@ impl SlotSender { ) -> io::Result<()> { let path = self.next_output_path(extension); std::fs::write(&path, data)?; - self.send_file_output(path, mime_type) + 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<()> { + /// `managed` is true when Coglet created the file (safe to delete after consumption); + /// false for user-authored Path outputs. + pub fn send_file_output( + &self, + path: PathBuf, + mime_type: Option, + managed: bool, + ) -> io::Result<()> { let filename = path .to_str() .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path"))? @@ -217,6 +224,7 @@ impl SlotSender { filename, kind: FileOutputKind::FileType, mime_type, + managed, }; self.tx .send(msg) @@ -265,6 +273,7 @@ fn build_output_message( filename, kind: FileOutputKind::Oversized, mime_type: None, + managed: true, }) } else { Ok(SlotResponse::OutputChunk { output, index }) From a6119b204be570b227607f57dc46f265fe3f47ba Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 13 Jul 2026 11:13:26 -0500 Subject: [PATCH 2/8] fix: safe ownership for returned paths and prediction ID hardening Addresses review feedback on PR #3096 by validating prediction IDs, rejecting duplicate active IDs, hardening remove_dir_all, and safely handling returned cog.Path files via the managed output directory. --- crates/coglet-python/src/predictor.rs | 8 +- crates/coglet/src/service.rs | 195 +++++++++++++++++++-- crates/coglet/src/transport/http/routes.rs | 195 +++++++++++++++++---- crates/coglet/src/worker.rs | 88 ++++++++++ docs/llms.txt | 3 +- docs/python.md | 3 +- 6 files changed, 445 insertions(+), 47 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 69de338bb3..6fc6a335e1 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -189,13 +189,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) { - // Path output — file already on disk, send path reference + // Path output — file already on disk. send_user_file_output copies + // external paths into the managed output directory and deletes them + // after upload, so Coglet never deletes arbitrary user-owned files. 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, false) + .send_user_file_output(std::path::PathBuf::from(path_str), None) .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; return Ok(()); } @@ -954,7 +956,7 @@ impl PythonPredictor { .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, false) + .send_user_file_output(std::path::PathBuf::from(path_str), None) .map_err(|e| { PredictionError::Failed(format!("Failed to send file output: {}", e)) })?; diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 0e6c4b1224..0b6a7c723f 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -51,6 +51,10 @@ pub enum CreatePredictionError { NotReady, #[error("At capacity (no slots available)")] AtCapacity, + #[error("A prediction with this ID already exists")] + DuplicateId, + #[error("Invalid prediction ID: {0}")] + InvalidId(&'static str), } const MAX_STREAM_SUBSCRIBERS: usize = STREAM_CHANNEL_CAPACITY; @@ -518,6 +522,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); @@ -537,16 +543,21 @@ 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, - }, - ); + // Register atomically in DashMap — duplicate active IDs are rejected to + // prevent one prediction from clobbering or deleting another's state. + match self.predictions.entry(id.clone()) { + dashmap::Entry::Vacant(entry) => { + entry.insert(PredictionEntry { + prediction: prediction_arc, + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); + } + dashmap::Entry::Occupied(_) => { + return Err(CreatePredictionError::DuplicateId); + } + } let handle = PredictionHandle { id, cancel_token }; @@ -802,7 +813,18 @@ impl PredictionService { /// predictions). pub fn remove_prediction(&self, id: &str) { self.predictions.remove(id); + + // Defense in depth: validate the ID and ensure we only delete a direct + // child of the prediction root, even though the HTTP layer also validates. + if let Err(e) = validate_prediction_id(id) { + tracing::warn!(prediction_id = %id, error = e, "Refusing to remove prediction dir: invalid ID"); + return; + } let dir = prediction_dir_for(id); + if !is_direct_prediction_child(&dir) { + tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); + return; + } match std::fs::remove_dir_all(&dir) { Ok(()) => {} Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} @@ -821,8 +843,48 @@ impl PredictionService { } } +const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; + +/// Validate a client-supplied prediction ID. +/// +/// IDs must be non-empty, must not contain path separators or traversal +/// sequences, and must not start with a dot. This prevents client-controlled +/// IDs from escaping the prediction root directory. +pub fn validate_prediction_id(id: &str) -> Result<(), &'static str> { + if id.is_empty() { + return Err("prediction ID must not be empty"); + } + if id.starts_with('.') { + return Err("prediction ID must not start with a dot"); + } + if id.contains('/') || id.contains('\\') { + return Err("prediction ID must not contain path separators"); + } + if id.contains("..") { + return Err("prediction ID must not contain traversal sequences"); + } + Ok(()) +} + fn prediction_dir_for(id: &str) -> std::path::PathBuf { - std::path::PathBuf::from("/tmp/coglet/predictions").join(id) + std::path::PathBuf::from(PREDICTION_ROOT).join(id) +} + +/// Verify that `dir` is a direct child of the prediction root. +/// +/// Used as a defense-in-depth check before recursive deletion. Returns true if +/// the canonicalized parent equals the canonicalized prediction root. +fn is_direct_prediction_child(dir: &std::path::Path) -> bool { + let Ok(canonical_dir) = std::fs::canonicalize(dir) else { + return false; + }; + let Some(parent) = canonical_dir.parent() else { + return false; + }; + let Ok(root) = std::fs::canonicalize(PREDICTION_ROOT) else { + return false; + }; + parent == root } fn spawn_orchestrator_cancel(orch: Arc, id: String) { @@ -1757,6 +1819,117 @@ mod tests { svc.remove_prediction("nonexistent-prediction"); } + #[test] + fn remove_prediction_does_not_delete_outside_prediction_root() { + let svc = PredictionService::new_no_pool(); + // Create a directory that looks like it could be a prediction dir but + // is outside the prediction root. + let outside_dir = std::path::PathBuf::from("/tmp/coglet_outside_test"); + std::fs::create_dir_all(&outside_dir).unwrap(); + std::fs::write(outside_dir.join("file.txt"), b"keep").unwrap(); + + // Absolute ID that would resolve to /tmp/coglet_outside_test if not defended. + svc.remove_prediction("/tmp/coglet_outside_test"); + + assert!( + outside_dir.exists(), + "directory outside prediction root should not be removed" + ); + + // Cleanup + std::fs::remove_dir_all(&outside_dir).unwrap(); + } + + #[test] + fn remove_prediction_does_not_delete_traversal_targets() { + let svc = PredictionService::new_no_pool(); + let victim = std::path::PathBuf::from("/tmp/coglet_victim_test"); + std::fs::create_dir_all(&victim).unwrap(); + + svc.remove_prediction("../coglet_victim_test"); + + assert!(victim.exists(), "traversal target should not be removed"); + + std::fs::remove_dir_all(&victim).unwrap(); + } + + #[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(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))); + } + + #[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 invalid_ids = vec![ + "", + "foo/bar", + "foo\\bar", + "../escape", + ".hidden", + "foo..bar", + ]; + + 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()); + } + + #[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(".hidden").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..6b42d3e678 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -23,7 +23,7 @@ use crate::prediction::SharedPredictionStreamEvent; use crate::predictor::PredictionError; use crate::service::{ CreatePredictionError, HealthSnapshot, PredictionService, PredictionStreamSubscription, - SubscribePredictionStreamError, + SubscribePredictionStreamError, validate_prediction_id, }; use crate::version::VersionInfo; use crate::webhook::{TraceContext, WebhookConfig, WebhookEventType, WebhookSender}; @@ -288,6 +288,20 @@ 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() +} + async fn create_prediction( State(service): State>, headers: HeaderMap, @@ -300,7 +314,15 @@ 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 request.id { + Some(id) => { + if let Err(msg) = validate_prediction_id(&id) { + return invalid_prediction_id_response(msg); + } + id + } + None => generate_prediction_id(), + }; let response_mode = prediction_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -323,6 +345,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,20 +357,23 @@ async fn create_prediction_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id - && req_id != &prediction_id - { - 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(); + if let Some(ref req_id) = request.id { + if let Err(msg) = validate_prediction_id(req_id) { + return invalid_prediction_id_response(msg); + } + if req_id != &prediction_id { + 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(); + } } let response_mode = prediction_response_mode(&headers); @@ -488,6 +517,19 @@ async fn create_prediction_with_id( ) .into_response(); } + Err(CreatePredictionError::DuplicateId) => { + 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(); @@ -823,7 +865,15 @@ 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 request.id { + Some(id) => { + if let Err(msg) = validate_prediction_id(&id) { + return invalid_prediction_id_response(msg); + } + id + } + None => generate_prediction_id(), + }; let response_mode = json_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -850,6 +900,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,20 +912,23 @@ 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 Some(ref req_id) = request.id { + if let Err(msg) = validate_prediction_id(req_id) { + return invalid_prediction_id_response(msg); + } + if 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(); + } } // Idempotent: return existing state if already submitted @@ -1578,6 +1635,82 @@ 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/.hidden") + .header("content-type", "application/json") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + + #[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); + + // Small delay to let the async task register the prediction + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + + 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(); + + assert_eq!(response.status(), StatusCode::CONFLICT); + } + #[tokio::test] async fn prediction_at_capacity() { let service = Arc::new(PredictionService::new_no_pool()); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index bc0645b62a..1242e3a183 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -164,6 +164,11 @@ impl SlotSender { } } + /// Directory where Coglet-managed output files for this prediction live. + pub fn output_dir(&self) -> &std::path::Path { + &self.output_dir + } + fn next_output_index(&self) -> u64 { self.output_counter.fetch_add(1, Ordering::Relaxed) } @@ -231,6 +236,33 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } + /// Send a user-returned file path, ensuring Coglet only deletes files it owns. + /// + /// If the path is already inside the Cog-managed output directory, send it as + /// `managed: true` so the orchestrator deletes it after upload. If it lives + /// outside the output directory, copy it to a unique path inside the output + /// directory and send the copy as managed. This preserves the documented + /// behavior that returned files are cleaned up while never deleting arbitrary + /// user-owned paths. + pub fn send_user_file_output( + &self, + path: PathBuf, + mime_type: Option, + ) -> io::Result<()> { + if path.starts_with(&self.output_dir) { + return self.send_file_output(path, mime_type, true); + } + + let ext = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("bin") + .to_string(); + let dest = self.next_output_path(&ext); + std::fs::copy(&path, &dest)?; + self.send_file_output(dest, mime_type, true) + } + /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -982,4 +1014,60 @@ mod tests { let config = WorkerConfig::default(); assert_eq!(config.num_slots, 1); } + + #[test] + fn send_user_file_output_copies_external_path() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + // Create a file outside the managed output dir + let external = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); + std::fs::write(external.path(), b"hello").unwrap(); + + sender + .send_user_file_output(external.path().to_path_buf(), None) + .unwrap(); + + let msg = rx.try_recv().unwrap(); + match msg { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => { + assert!( + filename.starts_with(output_dir.path().to_str().unwrap()), + "external path should be copied into output dir" + ); + assert!(std::path::Path::new(&filename).exists()); + assert_eq!(std::fs::read(&filename).unwrap(), b"hello"); + } + _ => panic!("expected managed FileOutput"), + } + } + + #[test] + fn send_user_file_output_passes_through_output_dir_path() { + 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(); + + let msg = rx.try_recv().unwrap(); + match msg { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => { + assert_eq!(filename, inside.to_str().unwrap()); + } + _ => panic!("expected managed FileOutput for path inside output dir"), + } + } } diff --git a/docs/llms.txt b/docs/llms.txt index e62311b0d6..737e1dcaea 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2771,7 +2771,8 @@ 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 copies returned files into a managed output directory, uploads the copy, + # and deletes it afterward. The original file is left untouched. 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..eb7d9e37a3 100644 --- a/docs/python.md +++ b/docs/python.md @@ -537,7 +537,8 @@ 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 copies returned files into a managed output directory, uploads the copy, + # and deletes it afterward. The original file is left untouched. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) From 8d5a4d52456fe160082cef0e8fa8d5074075ad9c Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 16 Jul 2026 16:21:15 -0500 Subject: [PATCH 3/8] fix: complete managed output cleanup --- crates/coglet-python/src/predictor.rs | 70 +++++++--- crates/coglet/src/bridge/protocol.rs | 3 +- crates/coglet/src/orchestrator.rs | 39 +++++- crates/coglet/src/service.rs | 152 +++++++++++++++------ crates/coglet/src/transport/http/routes.rs | 2 +- crates/coglet/src/worker.rs | 50 +++---- docs/llms.txt | 3 +- docs/python.md | 3 +- 8 files changed, 221 insertions(+), 101 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6fc6a335e1..edb137ed23 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -166,9 +166,22 @@ 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 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. +fn send_path_output( + item: &Bound<'_, PyAny>, + slot_sender: &SlotSender, +) -> Result<(), PredictionError> { + let path_str: String = item + .call_method0("__fspath__") + .and_then(|path| path.extract()) + .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; + slot_sender + .send_user_file_output(std::path::PathBuf::from(path_str), None) + .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e))) +} + fn send_output_item( py: Python<'_>, item: &Bound<'_, PyAny>, @@ -189,16 +202,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_user_file_output copies - // external paths into the managed output directory and deletes them - // after upload, so Coglet never deletes arbitrary user-owned files. - 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_user_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(()); } @@ -951,15 +957,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_user_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)); } @@ -1380,6 +1378,7 @@ mod tests { use std::path::PathBuf; + use coglet_core::bridge::protocol::SlotResponse; use pyo3::types::PyList; fn add_python_sdk_path(py: Python<'_>) { @@ -1437,6 +1436,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/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 68d30e828c..0e0ba4aba1 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -327,8 +327,7 @@ 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 created this file (IOBase write or oversized spill). - /// False for user-authored Path outputs — must not be deleted. + /// True if Coglet owns this file and may delete it after consumption. #[serde(default)] managed: bool, }, diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 3a8994bf9b..96921923e6 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -84,6 +84,14 @@ fn ensure_trailing_slash(s: &str) -> String { } } +fn read_file_output(filename: &str, managed: bool) -> std::io::Result> { + let bytes = std::fs::read(filename)?; + if managed && let Err(e) = std::fs::remove_file(filename) { + tracing::debug!(%filename, error = %e, "Failed to delete managed output file"); + } + Ok(bytes) +} + /// 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. @@ -1074,18 +1082,13 @@ async fn run_event_loop( } Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); - let bytes = match std::fs::read(&filename) { + let bytes = match read_file_output(&filename, managed) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); continue; } }; - if managed - && let Err(e) = std::fs::remove_file(&filename) - { - tracing::debug!(%slot_id, %filename, error = %e, "Failed to delete managed output file"); - } match kind { FileOutputKind::Oversized => { let output: serde_json::Value = match serde_json::from_slice(&bytes) { @@ -1314,6 +1317,30 @@ mod tests { assert!(!pending.contains("overflow")); } + #[test] + fn read_file_output_deletes_managed_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("managed.txt"); + std::fs::write(&path, b"managed").unwrap(); + + let bytes = read_file_output(path.to_str().unwrap(), true).unwrap(); + + assert_eq!(bytes, b"managed"); + assert!(!path.exists()); + } + + #[test] + fn read_file_output_preserves_unmanaged_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("unmanaged.txt"); + std::fs::write(&path, b"unmanaged").unwrap(); + + let bytes = read_file_output(path.to_str().unwrap(), false).unwrap(); + + assert_eq!(bytes, b"unmanaged"); + assert!(path.exists()); + } + #[test] fn wrap_outputs_schema_array_single_item() { // List[Path] with num_outputs=1 → ["url"] not "url" diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 0b6a7c723f..5f1a510929 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -533,6 +533,15 @@ 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(_) => { + return Err(CreatePredictionError::DuplicateId); + } + }; + let permit = pool .try_acquire() .ok_or(CreatePredictionError::AtCapacity)?; @@ -543,21 +552,12 @@ impl PredictionService { let slot = PredictionSlot::new(prediction, permit, idle_rx); let prediction_arc = slot.prediction(); - // Register atomically in DashMap — duplicate active IDs are rejected to - // prevent one prediction from clobbering or deleting another's state. - match self.predictions.entry(id.clone()) { - dashmap::Entry::Vacant(entry) => { - entry.insert(PredictionEntry { - prediction: prediction_arc, - cancel_token: cancel_token.clone(), - input, - cancel_on_stream_drop, - }); - } - dashmap::Entry::Occupied(_) => { - return Err(CreatePredictionError::DuplicateId); - } - } + prediction_entry.insert(PredictionEntry { + prediction: prediction_arc, + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); let handle = PredictionHandle { id, cancel_token }; @@ -812,25 +812,47 @@ impl PredictionService { /// that were not deleted individually (e.g. aborted uploads, cancelled /// predictions). pub fn remove_prediction(&self, id: &str) { - self.predictions.remove(id); - // Defense in depth: validate the ID and ensure we only delete a direct // child of the prediction root, even though the HTTP layer also validates. if let Err(e) = validate_prediction_id(id) { tracing::warn!(prediction_id = %id, error = e, "Refusing to remove prediction dir: invalid ID"); return; } + + // Keep the ID reserved until cleanup completes so a new prediction with + // the same ID cannot create files that this call then removes. + let prediction_entry = self.predictions.entry(id.to_string()); let dir = prediction_dir_for(id); - if !is_direct_prediction_child(&dir) { - tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); - return; - } - match std::fs::remove_dir_all(&dir) { - Ok(()) => {} - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + let cleanup_complete = match std::fs::symlink_metadata(&dir) { + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, Err(e) => { - tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir") + tracing::debug!(prediction_id = %id, error = %e, "Failed to inspect prediction dir"); + false } + Ok(metadata) if metadata.file_type().is_symlink() => match std::fs::remove_file(&dir) { + Ok(()) => true, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir symlink"); + false + } + }, + Ok(_) if !is_direct_prediction_child(&dir) => { + tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); + false + } + Ok(_) => match std::fs::remove_dir_all(&dir) { + Ok(()) => true, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir"); + false + } + }, + }; + + if cleanup_complete && let dashmap::Entry::Occupied(entry) = prediction_entry { + entry.remove(); } } @@ -847,21 +869,18 @@ const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; /// Validate a client-supplied prediction ID. /// -/// IDs must be non-empty, must not contain path separators or traversal -/// sequences, and must not start with a dot. This prevents client-controlled -/// IDs from escaping the prediction root directory. +/// IDs must be non-empty direct path components. This prevents client-controlled +/// IDs from escaping the prediction root directory without rejecting otherwise +/// safe IDs containing dots. pub fn validate_prediction_id(id: &str) -> Result<(), &'static str> { if id.is_empty() { return Err("prediction ID must not be empty"); } - if id.starts_with('.') { - return Err("prediction ID must not start with a dot"); - } if id.contains('/') || id.contains('\\') { return Err("prediction ID must not contain path separators"); } - if id.contains("..") { - return Err("prediction ID must not contain traversal sequences"); + if id == "." || id == ".." { + return Err("prediction ID must not be . or .."); } Ok(()) } @@ -1853,13 +1872,54 @@ mod tests { std::fs::remove_dir_all(&victim).unwrap(); } + #[cfg(unix)] + #[test] + fn remove_prediction_unlinks_directory_symlink_without_following_it() { + let svc = PredictionService::new_no_pool(); + let id = format!("symlink-test-{}", uuid::Uuid::new_v4()); + 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!(!dir.exists(), "prediction symlink should be removed"); + assert!(victim.path().join("keep.txt").exists()); + } + + #[tokio::test] + async fn remove_prediction_retains_id_when_cleanup_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 = format!("cleanup-failure-{}", uuid::Uuid::new_v4()); + let (_handle, _slot) = svc + .submit_prediction(id.clone(), serde_json::json!({}), None, false) + .await + .unwrap(); + 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)); + std::fs::remove_file(dir).unwrap(); + svc.remove_prediction(&id); + assert!(!svc.prediction_exists(&id)); + } + #[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(pool, orchestrator).await; + svc.set_orchestrator(Arc::clone(&pool), orchestrator).await; svc.set_health(Health::Ready).await; let (_handle, _slot) = svc @@ -1882,6 +1942,17 @@ mod tests { .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"); } #[tokio::test] @@ -1893,14 +1964,7 @@ mod tests { svc.set_orchestrator(pool, orchestrator).await; svc.set_health(Health::Ready).await; - let invalid_ids = vec![ - "", - "foo/bar", - "foo\\bar", - "../escape", - ".hidden", - "foo..bar", - ]; + let invalid_ids = vec!["", ".", "..", "foo/bar", "foo\\bar", "../escape"]; for id in invalid_ids { let result = svc @@ -1919,6 +1983,7 @@ mod tests { 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()); } #[test] @@ -1927,7 +1992,8 @@ mod tests { 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(".hidden").is_err()); + assert!(validate_prediction_id(".").is_err()); + assert!(validate_prediction_id("..").is_err()); } #[test] diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 6b42d3e678..f64feb2659 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -1660,7 +1660,7 @@ mod tests { let response = app .oneshot( - Request::put("/predictions/.hidden") + Request::put("/predictions/foo%5Cbar") .header("content-type", "application/json") .body(Body::from(r#"{"input":{}}"#)) .unwrap(), diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 1242e3a183..9464e93e7e 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -164,11 +164,6 @@ impl SlotSender { } } - /// Directory where Coglet-managed output files for this prediction live. - pub fn output_dir(&self) -> &std::path::Path { - &self.output_dir - } - fn next_output_index(&self) -> u64 { self.output_counter.fetch_add(1, Ordering::Relaxed) } @@ -213,8 +208,7 @@ impl SlotSender { /// /// 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. - /// `managed` is true when Coglet created the file (safe to delete after consumption); - /// false for user-authored Path outputs. + /// `managed` is true when Coglet owns the file and may delete it after consumption. pub fn send_file_output( &self, path: PathBuf, @@ -236,30 +230,30 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } - /// Send a user-returned file path, ensuring Coglet only deletes files it owns. + /// Transfer a user-returned file into Coglet's managed output directory. /// - /// If the path is already inside the Cog-managed output directory, send it as - /// `managed: true` so the orchestrator deletes it after upload. If it lives - /// outside the output directory, copy it to a unique path inside the output - /// directory and send the copy as managed. This preserves the documented - /// behavior that returned files are cleaned up while never deleting arbitrary - /// user-owned paths. + /// Returning a path hands ownership to Coglet. Copying before unlinking works + /// across filesystems while ensuring the source is only removed after the + /// managed copy has been written successfully. pub fn send_user_file_output( &self, path: PathBuf, mime_type: Option, ) -> io::Result<()> { - if path.starts_with(&self.output_dir) { - return self.send_file_output(path, mime_type, true); - } - let ext = path .extension() .and_then(|e| e.to_str()) .unwrap_or("bin") .to_string(); - let dest = self.next_output_path(&ext); + let mut dest = self.next_output_path(&ext); + while dest == path || dest.exists() { + dest = self.next_output_path(&ext); + } std::fs::copy(&path, &dest)?; + if let Err(e) = std::fs::remove_file(&path) { + let _ = std::fs::remove_file(&dest); + return Err(e); + } self.send_file_output(dest, mime_type, true) } @@ -1016,19 +1010,20 @@ mod tests { } #[test] - fn send_user_file_output_copies_external_path() { + fn send_user_file_output_transfers_external_path() { let output_dir = tempfile::tempdir().unwrap(); let (tx, mut rx) = mpsc::unbounded_channel::(); let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); - // Create a file outside the managed output dir let external = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); std::fs::write(external.path(), b"hello").unwrap(); + let external_path = external.path().to_path_buf(); sender - .send_user_file_output(external.path().to_path_buf(), None) + .send_user_file_output(external_path.clone(), None) .unwrap(); + assert!(!external_path.exists(), "returned path should be removed"); let msg = rx.try_recv().unwrap(); match msg { SlotResponse::FileOutput { @@ -1045,10 +1040,13 @@ mod tests { } _ => panic!("expected managed FileOutput"), } + + assert!(sender.send_user_file_output(external_path, None).is_err()); + assert!(rx.try_recv().is_err(), "stale path must not be sent again"); } #[test] - fn send_user_file_output_passes_through_output_dir_path() { + fn send_user_file_output_transfers_output_dir_path() { let output_dir = tempfile::tempdir().unwrap(); let (tx, mut rx) = mpsc::unbounded_channel::(); let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); @@ -1058,6 +1056,7 @@ mod tests { sender.send_user_file_output(inside.clone(), None).unwrap(); + assert!(!inside.exists(), "returned path should be removed"); let msg = rx.try_recv().unwrap(); match msg { SlotResponse::FileOutput { @@ -1065,9 +1064,10 @@ mod tests { managed: true, .. } => { - assert_eq!(filename, inside.to_str().unwrap()); + assert_ne!(filename, inside.to_str().unwrap()); + assert_eq!(std::fs::read(filename).unwrap(), b"world"); } - _ => panic!("expected managed FileOutput for path inside output dir"), + _ => panic!("expected managed FileOutput"), } } } diff --git a/docs/llms.txt b/docs/llms.txt index 8de057a814..e8419e39a8 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2784,8 +2784,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. - # Cog copies returned files into a managed output directory, uploads the copy, - # and deletes it afterward. The original file is left untouched. + # This file will automatically be deleted by Cog after it has been returned. 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 eb7d9e37a3..d717cf56f2 100644 --- a/docs/python.md +++ b/docs/python.md @@ -537,8 +537,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. - # Cog copies returned files into a managed output directory, uploads the copy, - # and deletes it afterward. The original file is left untouched. + # This file will automatically be deleted by Cog after it has been returned. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) From 45d06415c87dedcd5d1040529efba99ce32c3b1c Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 3 Aug 2026 23:01:35 -0500 Subject: [PATCH 4/8] fix: harden managed output handoff --- crates/coglet-python/src/predictor.rs | 19 +- crates/coglet-python/tests/test_coglet.py | 63 +++ crates/coglet/src/bridge/protocol.rs | 18 + crates/coglet/src/orchestrator.rs | 190 ++++++-- crates/coglet/src/permit/slot.rs | 52 +++ crates/coglet/src/service.rs | 481 ++++++++++++++++----- crates/coglet/src/transport/http/routes.rs | 135 +++--- crates/coglet/src/worker.rs | 370 ++++++++++++++-- docs/http.md | 3 + docs/llms.txt | 6 +- docs/python.md | 3 +- 11 files changed, 1074 insertions(+), 266 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index edb137ed23..770e67e781 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -164,24 +164,27 @@ fn submit_async_coroutine( Ok(future.unbind()) } -/// Send a single output item over IPC, routing file outputs to disk. -/// -/// 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. +/// 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_str: String = item + let path: String = item .call_method0("__fspath__") .and_then(|path| path.extract()) .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; - slot_sender - .send_user_file_output(std::path::PathBuf::from(path_str), None) + // 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): 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. fn send_output_item( py: Python<'_>, item: &Bound<'_, PyAny>, diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 57dcdd40a8..0b0573def7 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -1,5 +1,6 @@ """Tests for coglet Python bindings.""" +import base64 import queue import re import socket @@ -363,6 +364,41 @@ 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 + + class CogletServer: """Context manager for running coglet server.""" @@ -520,6 +556,33 @@ 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" + + class TestSecretInput: """Tests for cog.Secret input coercion.""" diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 0e0ba4aba1..608448ee7e 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -617,6 +617,24 @@ mod tests { 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, .. } => assert!(!managed), + resp => panic!("expected a FileOutput, got {resp:?}"), + } + } + #[test] fn slot_file_output_oversized_serializes() { let resp = SlotResponse::FileOutput { diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 96921923e6..93ed4213fd 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -84,12 +84,36 @@ fn ensure_trailing_slash(s: &str) -> String { } } -fn read_file_output(filename: &str, managed: bool) -> std::io::Result> { - let bytes = std::fs::read(filename)?; - if managed && let Err(e) = std::fs::remove_file(filename) { - tracing::debug!(%filename, error = %e, "Failed to delete managed output file"); +/// 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"); } - Ok(bytes) +} + +/// Report the first upload that did not succeed, if any. +/// +/// Uploads aborted by cancellation are not failures: the prediction is already +/// terminal and its outputs are no longer wanted. +fn first_upload_error( + results: Vec, tokio::task::JoinError>>, +) -> Option { + results.into_iter().find_map(|result| match result { + Ok(Ok(())) => None, + Ok(Err(error)) => Some(error), + Err(e) if e.is_cancelled() => None, + Err(e) => Some(format!("upload task failed: {e}")), + }) } /// Try to lock a prediction mutex. @@ -745,7 +769,10 @@ async fn run_event_loop( HashMap::new(); let mut pending_healthchecks: Vec> = Vec::new(); let mut healthcheck_counter: u64 = 0; - let mut pending_uploads: HashMap>> = HashMap::new(); + // Upload tasks report failure so `Done` can fail the prediction instead of + // reporting success with a missing output. + let mut pending_uploads: HashMap>>> = + HashMap::new(); let mut pending_cancellations: HashSet = HashSet::new(); let (slot_msg_tx, mut slot_msg_rx) = @@ -1082,7 +1109,7 @@ async fn run_event_loop( } Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); - let bytes = match read_file_output(&filename, managed) { + let bytes = match std::fs::read(&filename) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); @@ -1098,6 +1125,7 @@ async fn run_event_loop( continue; } }; + delete_managed_output(&filename, managed); let poisoned = if let Some(pred) = predictions.get(&slot_id) { if let Some(mut p) = try_lock_prediction(pred) { p.append_output(output); @@ -1128,18 +1156,19 @@ async fn run_event_loop( .unwrap_or("output") .to_string(); 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, &basename, &bytes, &mime) + .await + .inspect_err(|e| { + tracing::error!(%filename, error = %e, "Failed to upload file output"); + })?; + if let Some(pred) = pred + && let Some(mut p) = try_lock_prediction(&pred) + { + p.append_output(serde_json::Value::String(url)); } + // Only safe once the endpoint has the bytes. + delete_managed_output(&filename, managed); + Ok(()) }); pending_uploads.entry(slot_id).or_default().push(handle); } else { @@ -1150,6 +1179,7 @@ async fn run_event_loop( let output = serde_json::Value::String(format!( "data:{mime};base64,{encoded}" )); + delete_managed_output(&filename, managed); let poisoned = if let Some(pred) = predictions.get(&slot_id) { if let Some(mut p) = try_lock_prediction(pred) { p.append_output(output); @@ -1168,13 +1198,14 @@ async fn run_event_loop( } } 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(); if let Some(pred) = predictions.remove(&slot_id) { @@ -1201,11 +1232,11 @@ async fn run_event_loop( None => (None, id.clone()), }; tokio::spawn(async move { - if let Some(token) = cancel_token { + let results = if let Some(token) = cancel_token { let upload_fut = futures::future::join_all(uploads); tokio::pin!(upload_fut); tokio::select! { - _ = &mut upload_fut => {} + results = &mut upload_fut => results, _ = token.cancelled() => { tracing::info!( target: "coglet::prediction", @@ -1219,18 +1250,30 @@ async fn run_event_loop( } } } else { - for h in uploads { - let _ = h.await; - } - } - if let Some(mut p) = try_lock_prediction(&pred) { - let pred_output = wrap_outputs( - p.take_outputs(), - output_is_array, - is_stream, + futures::future::join_all(uploads).await + }; + + let Some(mut p) = try_lock_prediction(&pred) else { + return; + }; + // A dropped output would otherwise be reported as + // success with a short output array. + if let Some(error) = first_upload_error(results) { + tracing::error!( + target: "coglet::prediction", + prediction_id = %upload_pred_id, + %error, + "Failing prediction because an output upload failed" ); - p.set_succeeded(pred_output); + p.set_failed(format!("Failed to upload output file: {error}")); + return; } + let pred_output = wrap_outputs( + p.take_outputs(), + output_is_array, + is_stream, + ); + p.set_succeeded(pred_output); }); } } else { @@ -1293,6 +1336,8 @@ async fn run_event_loop( mod tests { use super::*; use serde_json::json; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; // ── wrap_outputs: schema says array (output_is_array = true) ── @@ -1318,27 +1363,94 @@ mod tests { } #[test] - fn read_file_output_deletes_managed_file() { + 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(); - let bytes = read_file_output(path.to_str().unwrap(), true).unwrap(); + delete_managed_output(path.to_str().unwrap(), true); - assert_eq!(bytes, b"managed"); - assert!(!path.exists()); + assert!(!path.exists(), "managed file should be deleted"); } #[test] - fn read_file_output_preserves_unmanaged_file() { + 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(); - let bytes = read_file_output(path.to_str().unwrap(), false).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) -> Option { + crate::install_crypto_provider(); + let filename = path.to_str().unwrap().to_string(); + let endpoint = ensure_trailing_slash(endpoint); + 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(()) + }); + first_upload_error(vec![handle.await]) + } + + /// 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("a 500 response should fail the upload"); - assert_eq!(bytes, b"unmanaged"); - assert!(path.exists()); + 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(); + + let error = upload_then_delete_managed(&server.uri(), &path).await; + + assert_eq!(error, None); + assert!(!path.exists(), "uploaded file should be deleted"); + } + + #[tokio::test] + async fn first_upload_error_reports_failures_but_not_cancellations() { + let ok: tokio::task::JoinHandle> = tokio::spawn(async { Ok(()) }); + let failed: tokio::task::JoinHandle> = + tokio::spawn(async { Err("upload returned status 500".to_string()) }); + let cancelled: tokio::task::JoinHandle> = + tokio::spawn(async { std::future::pending().await }); + cancelled.abort(); + + assert_eq!(first_upload_error(vec![ok.await]), None); + assert_eq!(first_upload_error(vec![cancelled.await]), None); + assert_eq!( + first_upload_error(vec![failed.await]), + Some("upload returned status 500".to_string()) + ); } #[test] diff --git a/crates/coglet/src/permit/slot.rs b/crates/coglet/src/permit/slot.rs index 95b980fb6b..784cebc6db 100644 --- a/crates/coglet/src/permit/slot.rs +++ b/crates/coglet/src/permit/slot.rs @@ -47,6 +47,18 @@ impl UnregisteredPredictionSlot { pub fn prediction(&self) -> Arc> { self.prediction_slot.prediction() } + + pub fn id(&self) -> String { + 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. @@ -135,6 +147,26 @@ 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()) } @@ -213,6 +245,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/service.rs b/crates/coglet/src/service.rs index 5f1a510929..52a59071f9 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -59,6 +59,13 @@ pub enum CreatePredictionError { 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")] @@ -646,8 +653,19 @@ impl PredictionService { let state = state .ok_or_else(|| PredictionError::Failed("No orchestrator configured".to_string()))?; + let prediction_id = unregistered_slot.id(); + + // 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(); { @@ -667,27 +685,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 = 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| 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() @@ -807,52 +804,49 @@ impl PredictionService { } } - /// Remove a prediction from the DashMap after completion and clean up - /// its on-disk prediction directory as a backstop for any output files - /// that were not deleted individually (e.g. aborted uploads, cancelled - /// predictions). + /// 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) { - // Defense in depth: validate the ID and ensure we only delete a direct - // child of the prediction root, even though the HTTP layer also validates. + // 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; } - // Keep the ID reserved until cleanup completes so a new prediction with - // the same ID cannot create files that this call then removes. - let prediction_entry = self.predictions.entry(id.to_string()); - let dir = prediction_dir_for(id); - let cleanup_complete = match std::fs::symlink_metadata(&dir) { - Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, - Err(e) => { - tracing::debug!(prediction_id = %id, error = %e, "Failed to inspect prediction dir"); - false - } - Ok(metadata) if metadata.file_type().is_symlink() => match std::fs::remove_file(&dir) { - Ok(()) => true, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, + 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) => { - tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir symlink"); - false + // 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; } - }, - Ok(_) if !is_direct_prediction_child(&dir) => { - tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); - false + }; + if let dashmap::Entry::Occupied(occupied) = entry { + occupied.remove(); } - Ok(_) => match std::fs::remove_dir_all(&dir) { - Ok(()) => true, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, - Err(e) => { - tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir"); - false - } - }, + detached }; - if cleanup_complete && let dashmap::Entry::Occupied(entry) = prediction_entry { - entry.remove(); + if let Some(dir) = detached { + delete_detached_dir(&dir); } } @@ -865,23 +859,44 @@ impl PredictionService { } } -const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; +/// 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. /// -/// IDs must be non-empty direct path components. This prevents client-controlled -/// IDs from escaping the prediction root directory without rejecting otherwise -/// safe IDs containing dots. -pub fn validate_prediction_id(id: &str) -> Result<(), &'static str> { +/// 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.contains('/') || id.contains('\\') { - return Err("prediction ID must not contain path separators"); - } 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(()) } @@ -889,21 +904,80 @@ fn prediction_dir_for(id: &str) -> std::path::PathBuf { std::path::PathBuf::from(PREDICTION_ROOT).join(id) } -/// Verify that `dir` is a direct child of the prediction root. +/// Move a prediction directory out of the prediction root, returning its new path. /// -/// Used as a defense-in-depth check before recursive deletion. Returns true if -/// the canonicalized parent equals the canonicalized prediction root. -fn is_direct_prediction_child(dir: &std::path::Path) -> bool { - let Ok(canonical_dir) = std::fs::canonicalize(dir) else { - return false; - }; - let Some(parent) = canonical_dir.parent() else { - return false; - }; - let Ok(root) = std::fs::canonicalize(PREDICTION_ROOT) else { - return false; +/// 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; + } }; - parent == root + 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, id: String) { @@ -1818,15 +1892,23 @@ 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 dir = prediction_dir_for("test-dir-cleanup"); + 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("test-dir-cleanup"); + svc.remove_prediction(&id); assert!(!dir.exists(), "prediction dir should be removed"); } @@ -1835,48 +1917,32 @@ mod tests { 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("nonexistent-prediction"); + svc.remove_prediction(&unique_test_id("nonexistent")); } #[test] - fn remove_prediction_does_not_delete_outside_prediction_root() { + fn remove_prediction_rejects_ids_that_are_not_plain_path_components() { let svc = PredictionService::new_no_pool(); - // Create a directory that looks like it could be a prediction dir but - // is outside the prediction root. - let outside_dir = std::path::PathBuf::from("/tmp/coglet_outside_test"); - std::fs::create_dir_all(&outside_dir).unwrap(); - std::fs::write(outside_dir.join("file.txt"), b"keep").unwrap(); + let victim = tempfile::tempdir().unwrap(); + std::fs::write(victim.path().join("keep.txt"), b"keep").unwrap(); + let victim_path = victim.path().to_str().unwrap(); - // Absolute ID that would resolve to /tmp/coglet_outside_test if not defended. - svc.remove_prediction("/tmp/coglet_outside_test"); + // 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!( - outside_dir.exists(), - "directory outside prediction root should not be removed" + victim.path().join("keep.txt").exists(), + "cleanup must not touch anything outside the prediction root" ); - - // Cleanup - std::fs::remove_dir_all(&outside_dir).unwrap(); - } - - #[test] - fn remove_prediction_does_not_delete_traversal_targets() { - let svc = PredictionService::new_no_pool(); - let victim = std::path::PathBuf::from("/tmp/coglet_victim_test"); - std::fs::create_dir_all(&victim).unwrap(); - - svc.remove_prediction("../coglet_victim_test"); - - assert!(victim.exists(), "traversal target should not be removed"); - - std::fs::remove_dir_all(&victim).unwrap(); } #[cfg(unix)] #[test] fn remove_prediction_unlinks_directory_symlink_without_following_it() { let svc = PredictionService::new_no_pool(); - let id = format!("symlink-test-{}", uuid::Uuid::new_v4()); + 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(); @@ -1885,32 +1951,99 @@ mod tests { svc.remove_prediction(&id); - assert!(!dir.exists(), "prediction symlink should be removed"); - assert!(victim.path().join("keep.txt").exists()); + 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_retains_id_when_cleanup_fails() { + 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 = format!("cleanup-failure-{}", uuid::Uuid::new_v4()); - let (_handle, _slot) = svc + 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)); - std::fs::remove_file(dir).unwrap(); - svc.remove_prediction(&id); - assert!(!svc.prediction_exists(&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] @@ -1955,6 +2088,83 @@ mod tests { 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(); @@ -1964,7 +2174,17 @@ mod tests { svc.set_orchestrator(pool, orchestrator).await; svc.set_health(Health::Ready).await; - let invalid_ids = vec!["", ".", "..", "foo/bar", "foo\\bar", "../escape"]; + 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 @@ -1972,8 +2192,7 @@ mod tests { .await; assert!( matches!(result, Err(CreatePredictionError::InvalidId(_))), - "expected invalid ID error for {:?}", - id + "expected invalid ID error for {id:?}" ); } } @@ -1984,6 +2203,7 @@ mod tests { 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] @@ -1996,6 +2216,27 @@ mod tests { 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 f64feb2659..bdd9437b9b 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -302,6 +302,37 @@ fn invalid_prediction_id_response(msg: &str) -> Response { .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, @@ -314,14 +345,9 @@ async fn create_prediction( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = match request.id { - Some(id) => { - if let Err(msg) = validate_prediction_id(&id) { - return invalid_prediction_id_response(msg); - } - id - } - None => 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); @@ -357,23 +383,10 @@ async fn create_prediction_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id { - if let Err(msg) = validate_prediction_id(req_id) { - return invalid_prediction_id_response(msg); - } - if req_id != &prediction_id { - 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(); - } + if let Err(message) = + check_body_id_matches_url(request.id.as_ref(), &prediction_id, "prediction") + { + return invalid_prediction_id_response(&message); } let response_mode = prediction_response_mode(&headers); @@ -540,6 +553,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); } @@ -865,14 +881,9 @@ async fn create_training( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = match request.id { - Some(id) => { - if let Err(msg) = validate_prediction_id(&id) { - return invalid_prediction_id_response(msg); - } - id - } - None => 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); @@ -912,23 +923,8 @@ async fn create_training_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id { - if let Err(msg) = validate_prediction_id(req_id) { - return invalid_prediction_id_response(msg); - } - if 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 @@ -1671,6 +1667,32 @@ mod tests { 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 @@ -1693,9 +1715,9 @@ mod tests { .await .unwrap(); assert_eq!(response.status(), StatusCode::ACCEPTED); - - // Small delay to let the async task register the prediction - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + // 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 @@ -1708,7 +1730,14 @@ mod tests { .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] diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 9464e93e7e..603735a2a9 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; @@ -150,15 +151,27 @@ type SlotWriter = pub struct SlotSender { tx: mpsc::UnboundedSender, output_dir: PathBuf, + /// Directories whose files Coglet may consume. See `owns_scratch_file`. + scratch_roots: Arc<[PathBuf]>, file_counter: Arc, output_counter: Arc, } impl SlotSender { pub fn new(tx: mpsc::UnboundedSender, output_dir: PathBuf) -> Self { + let scratch_roots = [output_dir.clone(), std::env::temp_dir()]; + Self::with_scratch_roots(tx, output_dir, scratch_roots) + } + + fn with_scratch_roots( + tx: mpsc::UnboundedSender, + output_dir: PathBuf, + scratch_roots: impl Into>, + ) -> Self { Self { tx, output_dir, + scratch_roots: scratch_roots.into(), file_counter: Arc::new(AtomicUsize::new(0)), output_counter: Arc::new(AtomicU64::new(0)), } @@ -168,12 +181,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,8 +232,8 @@ impl SlotSender { extension: &str, mime_type: Option, ) -> io::Result<()> { - let path = self.next_output_path(extension); - std::fs::write(&path, data)?; + let (mut file, path) = self.create_output_file(extension)?; + file.write_all(data)?; self.send_file_output(path, mime_type, true) } @@ -230,33 +263,79 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } - /// Transfer a user-returned file into Coglet's managed output directory. + /// Transfer a run-returned file into Coglet's managed output directory. /// - /// Returning a path hands ownership to Coglet. Copying before unlinking works - /// across filesystems while ensuring the source is only removed after the - /// managed copy has been written successfully. + /// Returning a path hands the data to Coglet, which deletes the managed copy + /// once the parent has read it. Scratch files are moved so nothing is left + /// behind; files from anywhere else are copied so a run cannot delete assets + /// baked into the image. See `owns_scratch_file`. pub fn send_user_file_output( &self, path: PathBuf, mime_type: Option, ) -> io::Result<()> { - let ext = path + 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 mut dest = self.next_output_path(&ext); - while dest == path || dest.exists() { - dest = self.next_output_path(&ext); - } - std::fs::copy(&path, &dest)?; - if let Err(e) = std::fs::remove_file(&path) { + // 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); + + let transfer = if self.owns_scratch_file(&path) { + move_file(&path, &dest) + } else { + std::fs::copy(&path, &dest).map(|_| ()) + }; + if let Err(e) = transfer { let _ = std::fs::remove_file(&dest); return Err(e); } + self.send_file_output(dest, mime_type, true) } + /// Whether Coglet may consume `path` outright instead of copying it. + /// + /// True for the places a run writes per-prediction scratch files: the output + /// dir Coglet handed it, and the temp dir. Reclaiming those is the point of + /// the ownership handoff. Anything else (weights, a bundled sample image, a + /// cache shared between predictions) is copied instead, because unlinking it + /// would break every later prediction. + /// + /// Both the directory entry and its resolved target must be under the same + /// scratch root. This keeps a symlink on either side of the boundary from + /// granting ownership of the other side. + fn owns_scratch_file(&self, path: &std::path::Path) -> bool { + let Ok(resolved) = path.canonicalize() else { + return false; + }; + let Some(parent) = path.parent() else { + return false; + }; + let Ok(resolved_parent) = parent.canonicalize() else { + return false; + }; + self.scratch_roots + .iter() + .filter_map(|root| root.canonicalize().ok()) + .any(|root| resolved.starts_with(&root) && resolved_parent.starts_with(&root)) + } + /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -279,6 +358,31 @@ impl SlotSender { } } +/// 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, @@ -1009,44 +1113,50 @@ mod tests { assert_eq!(config.num_slots, 1); } - #[test] - fn send_user_file_output_transfers_external_path() { - 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 external = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); - std::fs::write(external.path(), b"hello").unwrap(); - let external_path = external.path().to_path_buf(); - - sender - .send_user_file_output(external_path.clone(), None) - .unwrap(); - - assert!(!external_path.exists(), "returned path should be removed"); - let msg = rx.try_recv().unwrap(); - match msg { + /// 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!( - filename.starts_with(output_dir.path().to_str().unwrap()), - "external path should be copied into output dir" + path.starts_with(output_dir), + "managed output should live in the output dir, got {}", + path.display() ); - assert!(std::path::Path::new(&filename).exists()); - assert_eq!(std::fs::read(&filename).unwrap(), b"hello"); + path } - _ => panic!("expected managed FileOutput"), + response => panic!("expected a managed FileOutput, got {response:?}"), } + } - assert!(sender.send_user_file_output(external_path, None).is_err()); - assert!(rx.try_recv().is_err(), "stale path must not be sent again"); + #[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_transfers_output_dir_path() { + 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()); @@ -1056,18 +1166,190 @@ mod tests { sender.send_user_file_output(inside.clone(), None).unwrap(); - assert!(!inside.exists(), "returned path should be removed"); - let msg = rx.try_recv().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"); + } + + /// Returning a file from outside the scratch roots must not delete it, or the + /// first prediction to return a bundled asset would break every later one. + #[test] + fn send_user_file_output_copies_a_file_outside_the_scratch_roots() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + // Only the output dir counts as scratch here, so the asset dir below + // stands in for wherever a run's own files live: the model directory, a + // weights cache, anything baked into the image. + let sender = SlotSender::with_scratch_roots( + tx, + output_dir.path().to_path_buf(), + [output_dir.path().to_path_buf()], + ); + + let asset_dir = tempfile::tempdir().unwrap(); + let asset = asset_dir.path().join("placeholder.png"); + std::fs::write(&asset, b"asset").unwrap(); + + sender.send_user_file_output(asset.clone(), None).unwrap(); + + assert!( + asset.exists(), + "a file outside the scratch roots must survive" + ); + let managed = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_eq!(std::fs::read(&managed).unwrap(), b"asset"); + + // Returning it again still works, because the source is untouched. + sender.send_user_file_output(asset.clone(), None).unwrap(); + let second = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); + assert_ne!(second, managed); + assert_eq!(std::fs::read(&second).unwrap(), b"asset"); + } + + #[test] + fn owns_scratch_file_keeps_symlinks_from_crossing_the_boundary() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, _rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::with_scratch_roots( + tx, + output_dir.path().to_path_buf(), + [output_dir.path().to_path_buf()], + ); + + let outside_dir = tempfile::tempdir().unwrap(); + let outside = outside_dir.path().join("outside.txt"); + std::fs::write(&outside, b"outside").unwrap(); + let link = output_dir.path().join("link.txt"); + #[cfg(unix)] + std::os::unix::fs::symlink(&outside, &link).unwrap(); + + let inside = output_dir.path().join("inside.txt"); + std::fs::write(&inside, b"inside").unwrap(); + let outside_link = outside_dir.path().join("outside-link.txt"); + #[cfg(unix)] + std::os::unix::fs::symlink(&inside, &outside_link).unwrap(); + + assert!(sender.owns_scratch_file(&inside)); + #[cfg(unix)] + { + assert!( + !sender.owns_scratch_file(&link), + "a link inside a scratch root must not grant ownership of its target" + ); + assert!( + !sender.owns_scratch_file(&outside_link), + "a target inside a scratch root must not grant ownership of an outside link" + ); + } + } + + #[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"); + } + + #[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"); + } + + /// 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, .. } => { - assert_ne!(filename, inside.to_str().unwrap()); - assert_eq!(std::fs::read(filename).unwrap(), b"world"); + let spilled: serde_json::Value = + serde_json::from_slice(&std::fs::read(filename).unwrap()).unwrap(); + assert_eq!(spilled, oversized); } - _ => panic!("expected managed FileOutput"), + 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 e8419e39a8..a1b014e733 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1791,6 +1791,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 (`==`). @@ -2784,7 +2787,8 @@ 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 deletes this temporary file once the output has been delivered, 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..1ad9d27441 100644 --- a/docs/python.md +++ b/docs/python.md @@ -537,7 +537,8 @@ 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 deletes this temporary file once the output has been delivered, 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) From 06d813258631b29b73d6962367c34053103817d4 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Wed, 5 Aug 2026 16:23:01 -0500 Subject: [PATCH 5/8] fix: harden structured output handoff --- architecture/04-container-runtime.md | 27 +- crates/coglet-python/README.md | 2 +- crates/coglet-python/src/output.rs | 235 +++++++++- crates/coglet-python/src/predictor.rs | 43 +- crates/coglet-python/tests/test_coglet.py | 85 +++- crates/coglet/src/bridge/protocol.rs | 58 ++- ...__slot_file_output_managed_serializes.snap | 3 +- ...slot_file_output_oversized_serializes.snap | 3 +- crates/coglet/src/orchestrator.rs | 426 ++++++++++++------ crates/coglet/src/permit/slot.rs | 30 +- crates/coglet/src/prediction.rs | 33 +- crates/coglet/src/service.rs | 2 +- crates/coglet/src/worker.rs | 221 ++++++++- 13 files changed, 942 insertions(+), 226 deletions(-) diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 8758609780..1abe778001 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`), the worker transfers them into the prediction's managed output directory. Scratch files are consumed; files outside scratch storage are copied so model assets remain available. 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 770e67e781..8116c1adb5 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -184,7 +184,7 @@ fn send_path_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>, @@ -245,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 @@ -258,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(()) } @@ -995,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)?; @@ -1011,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. diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 0b0573def7..69f1883fd8 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -8,6 +8,7 @@ import sys import threading import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path import coglet @@ -399,12 +400,71 @@ def predict(self, text: str = "hello") -> Path: 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.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) + self.send_header( + "Location", + f"http://127.0.0.1:{self.server.server_port}{self.path}", + ) + 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 = [] @@ -412,10 +472,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, @@ -582,6 +643,26 @@ def test_returned_file_is_encoded_and_then_deleted( ) 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.txt") + assert uploads == [("/upload/result.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.""" diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 608448ee7e..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, }, @@ -330,6 +354,13 @@ pub enum SlotResponse { /// 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. @@ -562,7 +593,7 @@ mod tests { serde_json::to_value(resp).unwrap(), json!({ "type": "protocol_version", - "version": 1 + "version": SLOT_RESPONSE_PROTOCOL_VERSION }) ); } @@ -602,6 +633,8 @@ mod tests { 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); } @@ -613,6 +646,8 @@ mod tests { kind: FileOutputKind::FileType, mime_type: None, managed: false, + upload_filename: None, + index: 0, }; insta::assert_json_snapshot!(resp); } @@ -630,7 +665,16 @@ mod tests { let resp: SlotResponse = serde_json::from_value(json).unwrap(); match resp { - SlotResponse::FileOutput { managed, .. } => assert!(!managed), + SlotResponse::FileOutput { + managed, + upload_filename, + index, + .. + } => { + assert!(!managed); + assert_eq!(upload_filename, None); + assert_eq!(index, 0); + } resp => panic!("expected a FileOutput, got {resp:?}"), } } @@ -642,6 +686,8 @@ mod tests { kind: FileOutputKind::Oversized, mime_type: None, managed: true, + upload_filename: None, + index: 3, }; insta::assert_json_snapshot!(resp); } 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 index 9acbbb9cd5..946de46186 100644 --- 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 @@ -9,5 +9,6 @@ expression: resp "type": "file_type" }, "mime_type": "image/png", - "managed": true + "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 index 447bfeefab..f6347e32e9 100644 --- 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 @@ -8,5 +8,6 @@ expression: resp "kind": { "type": "oversized" }, - "managed": true + "managed": true, + "index": 3 } diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 93ed4213fd..c657a5a8ce 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::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`: @@ -101,18 +99,80 @@ fn delete_managed_output(filename: &str, managed: bool) { } } -/// Report the first upload that did not succeed, if any. -/// -/// Uploads aborted by cancellation are not failures: the prediction is already -/// terminal and its outputs are no longer wanted. -fn first_upload_error( - results: Vec, tokio::task::JoinError>>, -) -> Option { - results.into_iter().find_map(|result| match result { - Ok(Ok(())) => None, - Ok(Err(error)) => Some(error), - Err(e) if e.is_cancelled() => None, - Err(e) => Some(format!("upload task failed: {e}")), +struct ResolvedOutput { + index: u64, + output: serde_json::Value, +} + +type PendingOutput = tokio::task::JoinHandle>; + +fn collect_resolved_outputs( + 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() +} + +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, }) } @@ -732,18 +792,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>, @@ -771,9 +819,8 @@ async fn run_event_loop( let mut healthcheck_counter: u64 = 0; // Upload tasks report failure so `Done` can fail the prediction instead of // reporting success with a missing output. - let mut pending_uploads: HashMap>>> = - HashMap::new(); - let mut pending_cancellations: HashSet = HashSet::new(); + 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); @@ -830,6 +877,10 @@ 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); if let Some(pred) = predictions.remove(&slot) && let Some(mut p) = try_lock_prediction(&pred) && !p.is_terminal() @@ -839,6 +890,8 @@ 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(); for (slot, pred) in predictions.drain() { tracing::warn!(%slot, "Failing prediction due to worker fatal error"); pool.poison(slot); @@ -908,6 +961,8 @@ 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(); for (slot, pred) in predictions.drain() { tracing::warn!(%slot, "Failing prediction due to worker crash"); if let Some(mut p) = try_lock_prediction(&pred) { @@ -973,14 +1028,19 @@ 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); + // PredictionService retries cancellation after registration. + // Retaining an ID-only cancellation here could cancel a later + // prediction that reuses the public ID. + tracing::debug!(%prediction_id, "Cancel requested for unknown prediction"); } } } @@ -1007,24 +1067,7 @@ async fn run_event_loop( ); tracing::debug!(%slot_id, %prediction_id, "Registered 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() => { @@ -1107,12 +1150,30 @@ async fn run_event_loop( predictions.remove(&slot_id); } } - Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { + 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; } }; @@ -1122,13 +1183,16 @@ 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; } }; delete_managed_output(&filename, managed); let poisoned = if let Some(pred) = predictions.get(&slot_id) { if let Some(mut p) = try_lock_prediction(pred) { - p.append_output(output); + p.append_output_chunk(output, index); false } else { true @@ -1141,36 +1205,40 @@ async fn run_event_loop( } } 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 handle = tokio::spawn(async move { - let url = upload_file(&endpoint, &basename, &bytes, &mime) - .await - .inspect_err(|e| { + let url = upload_file( + &endpoint, + &upload_filename, + &bytes, + &mime, + ) + .await + .inspect_err(|e| { tracing::error!(%filename, error = %e, "Failed to upload file output"); })?; - if let Some(pred) = pred - && let Some(mut p) = try_lock_prediction(&pred) - { - p.append_output(serde_json::Value::String(url)); - } // Only safe once the endpoint has the bytes. delete_managed_output(&filename, managed); - Ok(()) + Ok(ResolvedOutput { + index, + output: serde_json::Value::String(url), + }) }); - 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; @@ -1182,7 +1250,7 @@ async fn run_event_loop( delete_managed_output(&filename, managed); let poisoned = if let Some(pred) = predictions.get(&slot_id) { if let Some(mut p) = try_lock_prediction(pred) { - p.append_output(output); + p.append_output_chunk(output, index); false } else { true @@ -1195,6 +1263,32 @@ async fn run_event_loop( } } } + 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.as_deref().map(ensure_trailing_slash); + let handle = tokio::spawn(resolve_structured_output( + structured, + endpoint, + index, + )); + pending_outputs.entry(slot_id).or_default().push(handle); + } } } Ok(SlotResponse::Done { id, output: _, predict_time, is_stream }) => { @@ -1207,67 +1301,116 @@ async fn run_event_loop( output_is_array, "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); 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 { - let results = 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! { - results = &mut upload_fut => results, + 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 { - futures::future::join_all(uploads).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; }; - // A dropped output would otherwise be reported as - // success with a short output array. - if let Some(error) = first_upload_error(results) { + if p.cancel_token().is_cancelled() { + p.set_canceled(); + return; + } + let resolved = match collect_resolved_outputs(results) { + Ok(resolved) => resolved, + Err(error) => { tracing::error!( target: "coglet::prediction", - prediction_id = %upload_pred_id, + prediction_id = %output_pred_id, %error, - "Failing prediction because an output upload failed" + "Failing prediction because an output task failed" ); - p.set_failed(format!("Failed to upload output file: {error}")); + p.set_failed(format!("Failed to process output file: {error}")); return; } + }; + for resolved in resolved { + p.append_output_chunk(resolved.output, resolved.index); + } let pred_output = wrap_outputs( p.take_outputs(), output_is_array, @@ -1287,10 +1430,10 @@ 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); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1303,10 +1446,10 @@ 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); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1315,9 +1458,10 @@ 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); if let Some(pred) = predictions.remove(&slot_id) && let Some(mut p) = try_lock_prediction(&pred) { @@ -1329,6 +1473,7 @@ async fn run_event_loop( } } + abort_all_output_tasks(&mut pending_outputs).await; tracing::info!("Event loop exiting"); } @@ -1349,19 +1494,6 @@ mod tests { assert_eq!(result.into_values(), Vec::::new()); } - #[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}")); - } - - record_pending_cancellation(&mut pending, "overflow".to_string()); - - assert_eq!(pending.len(), MAX_PENDING_CANCELLATIONS); - assert!(!pending.contains("overflow")); - } - #[test] fn delete_managed_output_deletes_managed_file() { let dir = tempfile::tempdir().unwrap(); @@ -1386,7 +1518,10 @@ mod tests { } /// Run the upload-then-delete sequence used by the FileOutput handler. - async fn upload_then_delete_managed(endpoint: &str, path: &std::path::Path) -> Option { + 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 = ensure_trailing_slash(endpoint); @@ -1395,7 +1530,9 @@ mod tests { delete_managed_output(&filename, true); Ok(()) }); - first_upload_error(vec![handle.await]) + 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 @@ -1413,7 +1550,7 @@ mod tests { let error = upload_then_delete_managed(&server.uri(), &path) .await - .expect("a 500 response should fail the upload"); + .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"); @@ -1430,27 +1567,34 @@ mod tests { let path = dir.path().join("0.txt"); std::fs::write(&path, b"payload").unwrap(); - let error = upload_then_delete_managed(&server.uri(), &path).await; + upload_then_delete_managed(&server.uri(), &path) + .await + .unwrap(); - assert_eq!(error, None); assert!(!path.exists(), "uploaded file should be deleted"); } #[tokio::test] - async fn first_upload_error_reports_failures_but_not_cancellations() { - let ok: tokio::task::JoinHandle> = tokio::spawn(async { Ok(()) }); - let failed: tokio::task::JoinHandle> = - tokio::spawn(async { Err("upload returned status 500".to_string()) }); - let cancelled: tokio::task::JoinHandle> = - tokio::spawn(async { std::future::pending().await }); - cancelled.abort(); - - assert_eq!(first_upload_error(vec![ok.await]), None); - assert_eq!(first_upload_error(vec![cancelled.await]), None); - assert_eq!( - first_upload_error(vec![failed.await]), - Some("upload returned status 500".to_string()) - ); + 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); + } + } + + 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!(dropped.load(Ordering::SeqCst)); } #[test] diff --git a/crates/coglet/src/permit/slot.rs b/crates/coglet/src/permit/slot.rs index 784cebc6db..cb77b25534 100644 --- a/crates/coglet/src/permit/slot.rs +++ b/crates/coglet/src/permit/slot.rs @@ -48,7 +48,7 @@ impl UnregisteredPredictionSlot { self.prediction_slot.prediction() } - pub fn id(&self) -> String { + pub fn id(&self) -> &str { self.prediction_slot.id() } @@ -65,6 +65,7 @@ impl UnregisteredPredictionSlot { /// /// 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, @@ -78,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)), @@ -171,11 +174,8 @@ impl PredictionSlot { 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 } } @@ -219,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] 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 52a59071f9..e4ac167d0e 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -653,7 +653,7 @@ impl PredictionService { let state = state .ok_or_else(|| PredictionError::Failed("No orchestrator configured".to_string()))?; - let prediction_id = unregistered_slot.id(); + 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. diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 603735a2a9..47865d575d 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -135,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; @@ -143,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 @@ -248,16 +257,29 @@ impl SlotSender { mime_type: Option, managed: bool, ) -> io::Result<()> { - let filename = path - .to_str() - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path"))? + let upload_filename = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("output") .to_string(); - let msg = SlotResponse::FileOutput { - filename, - kind: FileOutputKind::FileType, + 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: 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")) @@ -274,6 +296,15 @@ impl SlotSender { 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, @@ -291,22 +322,79 @@ impl SlotSender { .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); - let transfer = if self.owns_scratch_file(&path) { - move_file(&path, &dest) - } else { - std::fs::copy(&path, &dest).map(|_| ()) - }; + let is_symlink = std::fs::symlink_metadata(&path)?.file_type().is_symlink(); + let transfer = + if self.owns_scratch_file(&path) || (is_symlink && self.owns_scratch_entry(&path)) { + move_file(&path, &dest) + } else { + std::fs::copy(&path, &dest).map(|_| ()) + }; if let Err(e) = transfer { let _ = std::fs::remove_file(&dest); return Err(e); } - self.send_file_output(dest, mime_type, true) + 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) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } /// Whether Coglet may consume `path` outright instead of copying it. @@ -336,6 +424,19 @@ impl SlotSender { .any(|root| resolved.starts_with(&root) && resolved_parent.starts_with(&root)) } + fn owns_scratch_entry(&self, path: &std::path::Path) -> bool { + let Some(parent) = path.parent() else { + return false; + }; + let Ok(resolved_parent) = parent.canonicalize() else { + return false; + }; + self.scratch_roots + .iter() + .filter_map(|root| root.canonicalize().ok()) + .any(|root| resolved_parent.starts_with(root)) + } + /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -358,6 +459,12 @@ 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. /// @@ -404,6 +511,8 @@ fn build_output_message( kind: FileOutputKind::Oversized, mime_type: None, managed: true, + upload_filename: None, + index, }) } else { Ok(SlotResponse::OutputChunk { output, index }) @@ -1207,6 +1316,30 @@ mod tests { assert_eq!(std::fs::read(&second).unwrap(), b"asset"); } + #[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:?}"), + } + } + #[test] fn owns_scratch_file_keeps_symlinks_from_crossing_the_boundary() { let output_dir = tempfile::tempdir().unwrap(); @@ -1265,6 +1398,31 @@ mod tests { assert_eq!(std::fs::read(managed).unwrap(), b"target bytes"); } + #[cfg(unix)] + #[test] + fn send_user_file_output_unlinks_a_scratch_symlink_to_an_external_file() { + 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::with_scratch_roots( + tx, + output_dir.path().to_path_buf(), + [scratch_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(); @@ -1321,6 +1479,44 @@ mod tests { 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() { @@ -1334,6 +1530,7 @@ mod tests { filename, kind: FileOutputKind::Oversized, managed: true, + index: 0, .. } => { let spilled: serde_json::Value = From 6d9bd832628be5ebbcf31fc5dbe0eb21dc2e11d2 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Tue, 11 Aug 2026 12:32:37 -0500 Subject: [PATCH 6/8] fix: address coglet output review feedback --- architecture/04-container-runtime.md | 2 +- crates/coglet-python/tests/test_coglet.py | 100 ++++++- crates/coglet/src/orchestrator.rs | 284 ++++++++++++------ crates/coglet/src/service.rs | 332 +++++++++++++++------ crates/coglet/src/transport/http/routes.rs | 214 +++++++++++-- crates/coglet/src/worker.rs | 149 ++------- docs/llms.txt | 4 +- docs/python.md | 4 +- 8 files changed, 745 insertions(+), 344 deletions(-) diff --git a/architecture/04-container-runtime.md b/architecture/04-container-runtime.md index 1abe778001..128b43400d 100644 --- a/architecture/04-container-runtime.md +++ b/architecture/04-container-runtime.md @@ -372,7 +372,7 @@ 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 transfers them into the prediction's managed output directory. Scratch files are consumed; files outside scratch storage are copied so model assets remain available. 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. +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. 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. diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 69f1883fd8..1e1366a209 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -1,6 +1,7 @@ """Tests for coglet Python bindings.""" import base64 +import json import queue import re import socket @@ -10,6 +11,7 @@ import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from urllib.parse import quote, urlsplit import coglet import pytest @@ -201,6 +203,49 @@ 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.""" @@ -416,7 +461,7 @@ class Output(BaseModel): class Predictor(BasePredictor): def predict(self, text: str = "nested") -> Output: - output = Path(tempfile.mkdtemp()) / "result.txt" + 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"}) @@ -437,7 +482,7 @@ def do_PUT(self) -> None: self.send_response(200) self.send_header( "Location", - f"http://127.0.0.1:{self.server.server_port}{self.path}", + f"http://127.0.0.1:{self.server.server_port}{quote(urlsplit(self.path).path, safe='/%')}", ) self.send_header("Content-Length", "0") self.end_headers() @@ -656,8 +701,8 @@ def test_structured_path_is_uploaded_with_its_original_name( assert result["status"] == "succeeded", result assert result["output"]["metadata"] == {"label": "kept"} - assert result["output"]["artifact"].endswith("/upload/result.txt") - assert uploads == [("/upload/result.txt", b"structured bytes")] + 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() ) @@ -738,6 +783,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/orchestrator.rs b/crates/coglet/src/orchestrator.rs index c657a5a8ce..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; +use std::collections::{BTreeMap, HashMap}; use std::process::Stdio; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -41,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)) @@ -74,14 +82,6 @@ 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 @@ -104,18 +104,19 @@ struct ResolvedOutput { output: serde_json::Value, } -type PendingOutput = tokio::task::JoinHandle>; +type PendingOutput = tokio::task::JoinHandle>; -fn collect_resolved_outputs( - results: Vec, tokio::task::JoinError>>, -) -> Result, String> { +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() + .collect::, _>>() + .map(|_| ()) } async fn abort_output_tasks(tasks: Vec) { @@ -195,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: @@ -339,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; @@ -460,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, } @@ -485,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())) } @@ -804,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. @@ -813,6 +867,7 @@ 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(); @@ -881,6 +936,7 @@ async fn run_event_loop( 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() @@ -892,6 +948,7 @@ async fn run_event_loop( 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); @@ -963,6 +1020,7 @@ async fn run_event_loop( 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) { @@ -1005,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!( @@ -1037,10 +1095,7 @@ async fn run_event_loop( } } None => { - // PredictionService retries cancellation after registration. - // Retaining an ID-only cancellation here could cancel a later - // prediction that reuses the public ID. - tracing::debug!(%prediction_id, "Cancel requested for unknown prediction"); + tracing::debug!(%prediction_id, "Cancel requested for inactive prediction"); } } } @@ -1066,6 +1121,12 @@ 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 _ = registered_ack.send(()); } @@ -1135,19 +1196,10 @@ 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 { @@ -1190,18 +1242,10 @@ async fn run_event_loop( } }; delete_managed_output(&filename, managed); - 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 - }; - 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); } } FileOutputKind::FileType => { @@ -1219,7 +1263,12 @@ async fn run_event_loop( }); if let Some(ref url) = upload_url { // Spawn upload task so we don't block the event loop - let endpoint = ensure_trailing_slash(url); + let endpoint = url.clone(); + let publisher = Arc::clone( + output_publishers + .get(&slot_id) + .expect("active slot checked above"), + ); let handle = tokio::spawn(async move { let url = upload_file( &endpoint, @@ -1230,13 +1279,14 @@ async fn run_event_loop( .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); - Ok(ResolvedOutput { + publish_output( + &publisher, + serde_json::Value::String(url), index, - output: serde_json::Value::String(url), - }) + ) }); pending_outputs.entry(slot_id).or_default().push(handle); } else { @@ -1248,18 +1298,10 @@ async fn run_event_loop( "data:{mime};base64,{encoded}" )); delete_managed_output(&filename, managed); - 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 - }; - 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); } } } @@ -1281,12 +1323,21 @@ async fn run_event_loop( } }; delete_managed_output(&filename, managed); - let endpoint = upload_url.as_deref().map(ensure_trailing_slash); - let handle = tokio::spawn(resolve_structured_output( - structured, - endpoint, - index, - )); + 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); } } @@ -1303,6 +1354,7 @@ async fn run_event_loop( ); 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 pending.is_empty() { // No pending uploads — complete synchronously to avoid @@ -1395,22 +1447,16 @@ async fn run_event_loop( p.set_canceled(); return; } - let resolved = match collect_resolved_outputs(results) { - Ok(resolved) => resolved, - Err(error) => { + if let Err(error) = collect_output_task_errors(results) { tracing::error!( target: "coglet::prediction", - prediction_id = %output_pred_id, + prediction_id = %output_pred_id, %error, - "Failing prediction because an output task failed" + "Failing prediction because an output task failed" ); - p.set_failed(format!("Failed to process output file: {error}")); + p.set_failed(format!("Failed to process output file: {error}")); return; } - }; - for resolved in resolved { - p.append_output_chunk(resolved.output, resolved.index); - } let pred_output = wrap_outputs( p.take_outputs(), output_is_array, @@ -1434,6 +1480,7 @@ async fn run_event_loop( 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) { @@ -1450,6 +1497,7 @@ async fn run_event_loop( 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) { @@ -1462,6 +1510,7 @@ async fn run_event_loop( 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) { @@ -1481,7 +1530,7 @@ async fn run_event_loop( mod tests { use super::*; use serde_json::json; - use wiremock::matchers::method; + use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; // ── wrap_outputs: schema says array (output_is_array = true) ── @@ -1524,7 +1573,7 @@ mod tests { ) -> Result<(), String> { crate::install_crypto_provider(); let filename = path.to_str().unwrap().to_string(); - let endpoint = ensure_trailing_slash(endpoint); + 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); @@ -1574,6 +1623,51 @@ mod tests { 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}; diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index e4ac167d0e..c553103834 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -52,7 +52,7 @@ pub enum CreatePredictionError { #[error("At capacity (no slots available)")] AtCapacity, #[error("A prediction with this ID already exists")] - DuplicateId, + DuplicateId(ExistingPrediction), #[error("Invalid prediction ID: {0}")] InvalidId(&'static str), } @@ -108,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, } @@ -129,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) } } @@ -162,6 +223,7 @@ impl PredictionStreamSubscription { pub struct PredictionStreamGuard { id: String, + prediction: Arc>, service: Arc, cancel_on_stream_drop: bool, } @@ -172,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); } } @@ -191,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, } } @@ -210,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); } } } @@ -255,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 { @@ -275,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), } } @@ -333,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 { @@ -544,8 +623,10 @@ impl PredictionService { // cannot consume slot capacity. let prediction_entry = match self.predictions.entry(id.clone()) { dashmap::Entry::Vacant(entry) => entry, - dashmap::Entry::Occupied(_) => { - return Err(CreatePredictionError::DuplicateId); + dashmap::Entry::Occupied(entry) => { + return Err(CreatePredictionError::DuplicateId( + ExistingPrediction::from_entry(&id, entry.get()), + )); } }; @@ -560,13 +641,17 @@ impl PredictionService { let prediction_arc = slot.prediction(); prediction_entry.insert(PredictionEntry { - prediction: prediction_arc, + 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))) } @@ -581,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. @@ -709,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!( @@ -787,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(_) => { @@ -796,7 +888,7 @@ impl PredictionService { } }; if let Some(orch) = orchestrator { - spawn_orchestrator_cancel(orch, id_owned); + spawn_orchestrator_cancel(orch, prediction); } true } else { @@ -980,13 +1072,16 @@ fn prepare_slot_request( .map_err(|e| format!("Failed to build slot request: {e}")) } -fn spawn_orchestrator_cancel(orch: Arc, id: String) { +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, @@ -1092,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(()) } @@ -1136,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(()) @@ -1157,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), } } @@ -1178,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(()) } @@ -1557,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(); @@ -1839,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()); @@ -2074,7 +2240,7 @@ mod tests { ) .await; - assert!(matches!(result, Err(CreatePredictionError::DuplicateId))); + assert!(matches!(result, Err(CreatePredictionError::DuplicateId(_)))); assert_eq!(pool.available(), 1, "duplicate must not consume a permit"); let unrelated = svc @@ -2116,7 +2282,7 @@ mod tests { while let Some(result) = submits.join_next().await { match result.unwrap() { Ok(()) => admitted += 1, - Err(CreatePredictionError::DuplicateId) => duplicates += 1, + Err(CreatePredictionError::DuplicateId(_)) => duplicates += 1, Err(e) => panic!("unexpected submit error: {e}"), } } diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index bdd9437b9b..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, validate_prediction_id, + CreatePredictionError, ExistingPrediction, HealthSnapshot, PredictionService, + PredictionStreamSubscription, SubscribePredictionStreamError, validate_prediction_id, }; use crate::version::VersionInfo; use crate::webhook::{TraceContext, WebhookConfig, WebhookEventType, WebhookSender}; @@ -361,6 +361,7 @@ async fn create_prediction( response_mode, trace_context, false, + false, ) .await } @@ -391,17 +392,17 @@ async fn create_prediction_idempotent( 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, @@ -413,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, @@ -452,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 @@ -530,7 +552,14 @@ async fn create_prediction_with_id( ) .into_response(); } - Err(CreatePredictionError::DuplicateId) => { + 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!({ @@ -805,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 => ( @@ -897,6 +917,7 @@ async fn create_training( response_mode, trace_context, true, + false, ) .await } @@ -927,11 +948,21 @@ async fn create_training_idempotent( 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( @@ -944,6 +975,7 @@ async fn create_training_idempotent( response_mode, trace_context, true, + true, ) .await } @@ -1136,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(()) } @@ -1188,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!({ @@ -1606,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; @@ -1906,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 47865d575d..716e998a71 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -160,27 +160,15 @@ pub struct StagedFileOutput { pub struct SlotSender { tx: mpsc::UnboundedSender, output_dir: PathBuf, - /// Directories whose files Coglet may consume. See `owns_scratch_file`. - scratch_roots: Arc<[PathBuf]>, file_counter: Arc, output_counter: Arc, } impl SlotSender { pub fn new(tx: mpsc::UnboundedSender, output_dir: PathBuf) -> Self { - let scratch_roots = [output_dir.clone(), std::env::temp_dir()]; - Self::with_scratch_roots(tx, output_dir, scratch_roots) - } - - fn with_scratch_roots( - tx: mpsc::UnboundedSender, - output_dir: PathBuf, - scratch_roots: impl Into>, - ) -> Self { Self { tx, output_dir, - scratch_roots: scratch_roots.into(), file_counter: Arc::new(AtomicUsize::new(0)), output_counter: Arc::new(AtomicU64::new(0)), } @@ -288,9 +276,8 @@ impl SlotSender { /// 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. Scratch files are moved so nothing is left - /// behind; files from anywhere else are copied so a run cannot delete assets - /// baked into the image. See `owns_scratch_file`. + /// 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, @@ -332,14 +319,7 @@ impl SlotSender { let (file, dest) = self.create_output_file(&extension)?; drop(file); - let is_symlink = std::fs::symlink_metadata(&path)?.file_type().is_symlink(); - let transfer = - if self.owns_scratch_file(&path) || (is_symlink && self.owns_scratch_entry(&path)) { - move_file(&path, &dest) - } else { - std::fs::copy(&path, &dest).map(|_| ()) - }; - if let Err(e) = transfer { + if let Err(e) = move_file(&path, &dest) { let _ = std::fs::remove_file(&dest); return Err(e); } @@ -397,46 +377,6 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } - /// Whether Coglet may consume `path` outright instead of copying it. - /// - /// True for the places a run writes per-prediction scratch files: the output - /// dir Coglet handed it, and the temp dir. Reclaiming those is the point of - /// the ownership handoff. Anything else (weights, a bundled sample image, a - /// cache shared between predictions) is copied instead, because unlinking it - /// would break every later prediction. - /// - /// Both the directory entry and its resolved target must be under the same - /// scratch root. This keeps a symlink on either side of the boundary from - /// granting ownership of the other side. - fn owns_scratch_file(&self, path: &std::path::Path) -> bool { - let Ok(resolved) = path.canonicalize() else { - return false; - }; - let Some(parent) = path.parent() else { - return false; - }; - let Ok(resolved_parent) = parent.canonicalize() else { - return false; - }; - self.scratch_roots - .iter() - .filter_map(|root| root.canonicalize().ok()) - .any(|root| resolved.starts_with(&root) && resolved_parent.starts_with(&root)) - } - - fn owns_scratch_entry(&self, path: &std::path::Path) -> bool { - let Some(parent) = path.parent() else { - return false; - }; - let Ok(resolved_parent) = parent.canonicalize() else { - return false; - }; - self.scratch_roots - .iter() - .filter_map(|root| root.canonicalize().ok()) - .any(|root| resolved_parent.starts_with(root)) - } - /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -1281,39 +1221,27 @@ mod tests { assert_eq!(std::fs::read(&managed).unwrap(), b"world"); } - /// Returning a file from outside the scratch roots must not delete it, or the - /// first prediction to return a bundled asset would break every later one. #[test] - fn send_user_file_output_copies_a_file_outside_the_scratch_roots() { + 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::(); - // Only the output dir counts as scratch here, so the asset dir below - // stands in for wherever a run's own files live: the model directory, a - // weights cache, anything baked into the image. - let sender = SlotSender::with_scratch_roots( - tx, - output_dir.path().to_path_buf(), - [output_dir.path().to_path_buf()], - ); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); - let asset_dir = tempfile::tempdir().unwrap(); - let asset = asset_dir.path().join("placeholder.png"); - std::fs::write(&asset, b"asset").unwrap(); + 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(asset.clone(), None).unwrap(); + sender.send_user_file_output(source.clone(), None).unwrap(); - assert!( - asset.exists(), - "a file outside the scratch roots must survive" - ); + 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"asset"); + assert_eq!(std::fs::read(&managed).unwrap(), b"output"); - // Returning it again still works, because the source is untouched. - sender.send_user_file_output(asset.clone(), None).unwrap(); - let second = expect_managed_output(rx.try_recv().unwrap(), output_dir.path()); - assert_ne!(second, managed); - assert_eq!(std::fs::read(&second).unwrap(), b"asset"); + 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] @@ -1340,43 +1268,6 @@ mod tests { } } - #[test] - fn owns_scratch_file_keeps_symlinks_from_crossing_the_boundary() { - let output_dir = tempfile::tempdir().unwrap(); - let (tx, _rx) = mpsc::unbounded_channel::(); - let sender = SlotSender::with_scratch_roots( - tx, - output_dir.path().to_path_buf(), - [output_dir.path().to_path_buf()], - ); - - let outside_dir = tempfile::tempdir().unwrap(); - let outside = outside_dir.path().join("outside.txt"); - std::fs::write(&outside, b"outside").unwrap(); - let link = output_dir.path().join("link.txt"); - #[cfg(unix)] - std::os::unix::fs::symlink(&outside, &link).unwrap(); - - let inside = output_dir.path().join("inside.txt"); - std::fs::write(&inside, b"inside").unwrap(); - let outside_link = outside_dir.path().join("outside-link.txt"); - #[cfg(unix)] - std::os::unix::fs::symlink(&inside, &outside_link).unwrap(); - - assert!(sender.owns_scratch_file(&inside)); - #[cfg(unix)] - { - assert!( - !sender.owns_scratch_file(&link), - "a link inside a scratch root must not grant ownership of its target" - ); - assert!( - !sender.owns_scratch_file(&outside_link), - "a target inside a scratch root must not grant ownership of an outside link" - ); - } - } - #[cfg(unix)] #[test] fn send_user_file_output_copies_a_relative_symlink_before_unlinking_it() { @@ -1400,7 +1291,7 @@ mod tests { #[cfg(unix)] #[test] - fn send_user_file_output_unlinks_a_scratch_symlink_to_an_external_file() { + 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(); @@ -1409,11 +1300,7 @@ mod tests { 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::with_scratch_roots( - tx, - output_dir.path().to_path_buf(), - [scratch_dir.path().to_path_buf()], - ); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); sender.send_user_file_output(link.clone(), None).unwrap(); diff --git a/docs/llms.txt b/docs/llms.txt index a1b014e733..3e170acd09 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2775,6 +2775,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: @@ -2787,8 +2788,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. - # Cog deletes this temporary file once the output has been delivered, so - # write it again before returning the same path twice. + # 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 1ad9d27441..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,8 +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. - # Cog deletes this temporary file once the output has been delivered, so - # write it again before returning the same path twice. + # 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) From 0de50d109e0122fee5f03bbfee8a4b358e56e414 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Tue, 11 Aug 2026 12:34:33 -0500 Subject: [PATCH 7/8] Potential fix for pull request finding 'CodeQL / HTTP Response Splitting' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Anish Sahoo --- crates/coglet-python/tests/test_coglet.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index 1e1366a209..d842a37e8d 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -480,9 +480,10 @@ 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(self.path).path, safe='/%')}", + 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() From 3902d29b4144d4ffec198e45947aa19f54abfdf4 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Tue, 11 Aug 2026 12:37:07 -0500 Subject: [PATCH 8/8] chore: fix formatting --- crates/coglet-python/tests/test_coglet.py | 28 ++++++++++++----------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/crates/coglet-python/tests/test_coglet.py b/crates/coglet-python/tests/test_coglet.py index d842a37e8d..a2f5111d02 100644 --- a/crates/coglet-python/tests/test_coglet.py +++ b/crates/coglet-python/tests/test_coglet.py @@ -229,19 +229,21 @@ def predict(self, release_path: str) -> Iterator[Path]: 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"}}, - } - }, - }) + 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