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

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

Related Issues

• Fixes #307
• Follow-up to #306
• Original issue #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.