fix: delete coglet-managed output files after upload - #3096
Conversation
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
|
LGTM |
michaeldwan
left a comment
michaeldwan
left a comment
There was a problem hiding this comment.
Thanks for taking this on. Cleaning up Coglet-created spill files is worthwhile, and making ownership explicit in the protocol is a useful direction.
I’m requesting changes because the recursive cleanup currently trusts a client-controlled prediction ID, duplicate IDs can delete a live prediction’s state, and predictor-returned paths -- the files described in #1434 -- remain intentionally undeleted. I’ve left details inline.
Please also add behavioral coverage around cleanup. The snapshots prove that managed serializes, but they don’t prove that managed files are deleted, that returned paths follow the intended contract, or that one prediction can’t remove another prediction’s directory.
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.
|
LGTM |
|
LGTM |
|
Follow-up summary for Michael’s review:
Validation completed:
|
![ask-bonk[bot]](https://avatars.githubusercontent.com/in/2454580?s=60&v=4){kind=small}
There was a problem hiding this comment.
I'm Bonk, and I've done a quick review of your PR.
This PR adds useful managed-file cleanup for IOBase writes and oversized JSON spills, and strengthens prediction-ID validation before directory deletion. However, two claims in the PR description and review responses don't match the actual code.
-
Stale-file bug (#1434) not fixed for Path outputs —
send_user_file_outputcopies external paths into the managed output directory but never deletes the original file. The PR description explicitly states "removes the source only after a successful copy," and this was the core concern raised by Michael in the prior review. Without deleting the source, a predictor that returns a fixed temp path and fails to overwrite it on a subsequent prediction will still upload stale data. -
ID reservation not held through cleanup —
remove_predictionremoves the DashMap entry before callingremove_dir_all, so a reused ID could be submitted and have its new files deleted by the old cleanup task. The PR description claims the reservation is held through deletion, but the code does the opposite.
Suggested fixes are inline.
|
Posted review comments on PR #3096: #3096 (review) |
|
LGTM |
|
@anish-sahoo Bonk workflow failed. Check the logs for details. View workflow run · To retry, trigger Bonk again. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
I reviewed PR #3096 by examining the changed files, focusing on the send_user_file_output function in worker.rs and the remove_prediction function in service.rs to verify fixes for the stale file bug and ID reservation race condition mentioned in previous reviews. I checked the code against the review comments and recent commits to confirm whether the reported issues had been properly addressed. |
|
I reviewed PR #3096 by examining the changed files, focusing on the send_user_file_output function in worker.rs and the remove_prediction function in service.rs to verify fixes for the stale file bug and ID reservation race condition mentioned in previous reviews. I checked the code against the review comments and recent commits to confirm whether the reported issues had been properly addressed. Specifically, I looked at:
Based on my examination of the current code in these files, I can confirm that both issues mentioned in the previous review have been fixed:
The code shows these fixes have been implemented as suggested in the previous review comments. Since I'm acting as a reviewer and the working tree is read-only, I cannot make any changes, but I can confirm the issues have been addressed in the current PR state. LGTM |
| .and_then(|e| e.to_str()) | ||
| .unwrap_or("bin") | ||
| .to_string(); | ||
| let upload_filename = path |
There was a problem hiding this comment.
Preserving the original basename makes this user-controlled now. upload_file() later concatenates this directly onto the endpoint, so otherwise-valid names like result?draft.png, result#1.png, or report%20final.txt are interpreted as URL syntax rather than an object-key segment. That can upload to the wrong key (or fail) and return the wrong URL.
Can we append it as one escaped URL path segment instead? A regression test with reserved characters in a returned filename would catch this.
There was a problem hiding this comment.
Fixed in 6d9bd83. Upload filenames are appended as one escaped URL path segment, and signed query parameters are preserved. The regression test uses result?draft.png with a signed query.
| tracing::error!(error = %e, "Failed to upload file output"); | ||
| } | ||
| } | ||
| let url = upload_file( |
There was a problem hiding this comment.
This changes streaming behavior for file outputs. The task now only returns a ResolvedOutput; we keep it in pending_outputs and do not call append_output_chunk until the worker sends Done and every output task has completed.
For a generator that yields a Path and then keeps working, the SSE client will not see that yielded value until the prediction finishes. Before this change, a completed file upload was appended as soon as it completed. The architecture doc still says yielded FileOutputs stream immediately.
Could we publish each successful resolved output as its task completes, while still retaining and joining the task handles before terminal state? Please add an SSE regression test that yields a path, blocks before Done, and asserts the output arrives during the block.
There was a problem hiding this comment.
Fixed in 6d9bd83. Completed uploads now publish immediately through an ordered output publisher, so SSE keeps protocol order without waiting for Done. A subprocess test yields a path, blocks before completion, and confirms the output arrives first.
| ) | ||
| .into_response(); | ||
| } | ||
| Err(CreatePredictionError::DuplicateId) => { |
There was a problem hiding this comment.
There is still a check-then-act race for idempotent PUTs: two requests can both miss the entry at get_prediction_response, then one reserves the ID and the other reaches this branch. The second request gets 409, even though an idempotent retry should return the prediction that just won the race.
Could the idempotent handlers re-read and return the existing prediction on DuplicateId (including the SSE path), or move this into an atomic get-or-create service operation? This applies to the training idempotent endpoint too. A concurrent PUT /predictions/{id} regression test would be useful.
There was a problem hiding this comment.
Fixed in 6d9bd83. Duplicate reservations now carry a stable handle to the existing prediction, so cleanup cannot make the retry fall back to 409. Barrier-based tests cover concurrent prediction and training PUTs.
|
Pushed |
…ing' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Anish Sahoo <anishsahoo2005@gmail.com>
|
I'm Bonk, and I've done a quick review of your PR. This PR implements file output ownership transfer and cleanup for Coglet, addressing issue #1434. It introduces proper handling of returned file paths by moving them into managed storage and cleaning them up after use. It also fixes prediction ID validation and reservation race conditions. The changes look good and address the previous review concerns:
LGTM! |
|
self signed certificate |
|
@anish-sahoo Bonk workflow failed. Check the logs for details. View workflow run · To retry, trigger Bonk again. |
|
LGTM |
|
LGTM |
Fixes #1434.
Pathas an ownership handoff. Coglet consumes the source entry during staging, including paths outside temporary storage, while preserving symlink targets.