fix: close async blob files after interrupted uploads

ollama/_client.py

@@ -1301,15 +1301,21 @@ async def create_blob(self, path: Union[str, Path]) -> str:
 
     digest = f'sha256:{sha256sum.hexdigest()}'
 
-    async def upload_bytes():
-      async with await anyio.open_file(path, 'rb') as r:
+    r = await anyio.open_file(path, 'rb')
+    try:
+
+      async def upload_bytes():
         while True:
           chunk = await r.read(32 * 1024)
           if not chunk:
             break
           yield chunk
 
-    await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes())
+      await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes())
+    finally:
+      # Closing the local file must complete even when the upload is cancelled.
+      with anyio.CancelScope(shield=True):
+        await r.aclose()
 
     return digest
 

tests/test_client.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import inspect
 import json
@@ -7,7 +8,10 @@
 from pathlib import Path
 from typing import Any
 
+import anyio
+import httpx
 import pytest
+from anyio.lowlevel import checkpoint
 from httpx import Response as httpxResponse
 from pydantic import BaseModel
 from pytest_httpserver import HTTPServer, URIPattern
@@ -1243,6 +1247,51 @@ async def test_async_client_create_blob_exists(httpserver: HTTPServer):
     assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'
 
 
+@pytest.mark.parametrize('interruption', ('write_error', 'cancelled', 'cancel_scope'))
+async def test_async_client_create_blob_closes_file_on_interrupted_upload(tmp_path, monkeypatch, interruption):
+  path = tmp_path / 'weights.gguf'
+  path.write_bytes(b'x' * 65537)
+  opened_files = []
+  open_file = anyio.open_file
+
+  async def track_open_file(*args, **kwargs):
+    file = await open_file(*args, **kwargs)
+    opened_files.append(file)
+    return file
+
+  monkeypatch.setattr(anyio, 'open_file', track_open_file)
+
+  class InterruptedUpload(httpx.AsyncBaseTransport):
+    async def handle_async_request(self, request):
+      async for chunk in request.stream:
+        assert chunk
+        if interruption == 'cancelled':
+          asyncio.current_task().cancel()
+          await asyncio.sleep(0)
+        elif interruption == 'cancel_scope':
+          cancel_scope.cancel()
+          await checkpoint()
+        raise httpx.WriteError('upload interrupted', request=request)
+
+  async with AsyncClient(transport=InterruptedUpload()) as client:
+    try:
+      if interruption == 'cancel_scope':
+        with anyio.CancelScope() as cancel_scope:
+          await client.create_blob(path)
+        assert cancel_scope.cancelled_caught
+      else:
+        error = asyncio.CancelledError if interruption == 'cancelled' else httpx.WriteError
+        with pytest.raises(error):
+          await client.create_blob(path)
+
+      assert opened_files
+      assert all(file.wrapped.closed for file in opened_files)
+    finally:
+      # Also close real files when running the regression against an unfixed client.
+      for file in opened_files:
+        await file.aclose()
+
+
 async def test_async_client_delete(httpserver: HTTPServer):
   httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200))
   client = AsyncClient(httpserver.url_for('/api/delete'))