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 })