diff --git a/ollama/_client.py b/ollama/_client.py index 13dd6441..4ff5ad8c 100644 --- a/ollama/_client.py +++ b/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 diff --git a/tests/test_client.py b/tests/test_client.py index c212a11d..a291a616 100644 --- a/tests/test_client.py +++ b/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'))