Make process_prompt Cancellable Outside Downloads

[rahat2134]

This PR implements protection for downloads at their source in HFShardDownloader, ensuring downloads complete regardless of request cancellation. This is a follow-up fix to #306, moving from blanket request shielding to targeted download protection.

Changes Made

1. Modified HFShardDownloader.ensure_shard to:
   • Protect downloads at their source
   • Handle cancellations gracefully
   • Properly manage concurrent downloads
2. Added comprehensive tests in test_hf_shard_download.py
3. Removed request-level protection from ChatGPTAPI

Implementation Details

async def ensure_shard(self, shard: Shard) -> Path:
    # ... existing download handling ...
    try:
        return await download_task
    except asyncio.CancelledError:
        # Continue download despite cancellation
        await event.wait()
        return result

• Downloads protected where they actually happen
• Clean handling of concurrent requests
• Proper resource cleanup
• Robust error handling

Test Coverage

Added comprehensive tests in test_hf_shard_download.py:

1. test_download_protection: Verifies downloads complete despite cancellation
2. test_multiple_downloads: Ensures concurrent downloads work correctly
3. test_download_error_handling: Validates proper error handling

Testing

test/test_hf_shard_download.py::test_download_protection PASSED
test/test_hf_shard_download.py::test_multiple_downloads PASSED
test/test_hf_shard_download.py::test_download_error_handling PASSED

Breaking Changes

None. This improves download protection while maintaining backward compatibility.

Checklist

• Tests added for all scenarios
• Download protection moved to source
• Concurrent downloads handled properly
• Resource cleanup implemented
• Previous approach removed cleanly

Related Issues

• Fixes [BOUNTY - $200] make it so process_prompt is cancellable outside of a download (follow up to #305) #307
• Follow-up to shield process_prompt so downloads dont get cancelled when chatgpt api request times out fixes #305 #306
• Original issue exo stops downloading when a chatgpt api request times out #305

Note to Reviewer: Per your feedback, I've moved the protection to where downloads actually happen in HFShardDownloader, rather than trying to protect at the request level. This ensures downloads complete regardless of what triggers them (process_prompt or any other request). The solution is focused and thoroughly tested.

[rahat2134]

@AlexCheema PTAL, I am open to any feedback!

[AlexCheema]

I like the idea, the code is well written and the tests are extensive.

However, this doesn't quite solve the problem of preventing downloads from being cancelled -- a download can be started in process_prompt or could be started by another request. I think the idea is correct -- we should prevent downloads from being cancelled but this isn't quite how to do it. Let me know what you think.

[rahat2134]

I like the idea, the code is well written and the tests are extensive.

However, this doesn't quite solve the problem of preventing downloads from being cancelled -- a download can be started in process_prompt or could be started by another request. I think the idea is correct -- we should prevent downloads from being cancelled but this isn't quite how to do it. Let me know what you think.

@AlexCheema Thank you for the feedback! You're absolutely right about the limitation. I see now that checking downloads only at the start of process_prompt doesn't fully protect downloads, since:

1. Downloads could start after the check
2. Downloads could be triggered by other requests
3. We need protection at the download level, not the request level.

Would it be better to:

1. Move the protection closer to the download operation itself in HFShardDownloader?
2. Track active downloads more granularly?

Could you guide me on which approach would be most appropriate? I can then update the PR accordingly.

[AlexCheema]

I like the idea, the code is well written and the tests are extensive.
However, this doesn't quite solve the problem of preventing downloads from being cancelled -- a download can be started in process_prompt or could be started by another request. I think the idea is correct -- we should prevent downloads from being cancelled but this isn't quite how to do it. Let me know what you think.

@AlexCheema Thank you for the feedback! You're absolutely right about the limitation. I see now that checking downloads only at the start of process_prompt doesn't fully protect downloads, since:

1. Downloads could start after the check
2. Downloads could be triggered by other requests
3. We need protection at the download level, not the request level.

Would it be better to:

1. Move the protection closer to the download operation itself in HFShardDownloader?
2. Track active downloads more granularly?

Could you guide me on which approach would be most appropriate? I can then update the PR accordingly.

Yes I think your suggestions are good. I'm not sure I have a full solution in my head right now I'd need to look more deeply. If you want to take a look and suggest something or make more changes in this PR go ahead

[rahat2134]

Hi @AlexCheema , I've updated the PR based on your feedback. Rather than protecting at the request level, I've moved the protection directly into HFShardDownloader where downloads happen. This ensures downloads are protected regardless of what triggers them (process_prompt or any other request).
The approach is simple but effective:

Downloads are protected where they start
Proper handling of cancellations and concurrent requests
Clean resource management
Thorough test coverage

All tests are passing and the changes are focused just on download protection. Let me know if you'd like me to make any adjustments to this approach.

If not then we can merge this PR. Thanks :)

p.s - I have updated the PR description as well.

[rahat2134]

@AlexCheema PTAL!!

[rahat2134]

@AlexCheema PTAL!