fix: prevent memory leaks in file handling and HTTP connections

- Fix file handle memory leak in download_file_stream by properly closing and reopening files
- Add requests.Session context manager for proper HTTP connection cleanup
- Remove unnecessary file.seek(0) after file reopening
- Add timeout to prevent hanging connections

This prevents memory accumulation during large file downloads and ensures
proper resource cleanup in all scenarios.

Signed-off-by: Sihyeon Jang <sihyeon.jang@navercorp.com>
This commit is contained in:
Sihyeon Jang 2025-09-20 06:37:10 +09:00
parent dbd7b90123
commit 7042318c34

View file

@ -1694,25 +1694,27 @@ async def download_file_stream(
yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n'
if done:
file.seek(0)
chunk_size = 1024 * 1024 * 2
hashed = calculate_sha256(file, chunk_size)
file.seek(0)
file.close()
url = f"{ollama_url}/api/blobs/sha256:{hashed}"
response = requests.post(url, data=file)
with open(file_path, "rb") as file:
chunk_size = 1024 * 1024 * 2
hashed = calculate_sha256(file, chunk_size)
if response.ok:
res = {
"done": done,
"blob": f"sha256:{hashed}",
"name": file_name,
}
os.remove(file_path)
url = f"{ollama_url}/api/blobs/sha256:{hashed}"
with requests.Session() as session:
response = session.post(url, data=file, timeout=30)
yield f"data: {json.dumps(res)}\n\n"
else:
raise "Ollama: Could not create blob, Please try again."
if response.ok:
res = {
"done": done,
"blob": f"sha256:{hashed}",
"name": file_name,
}
os.remove(file_path)
yield f"data: {json.dumps(res)}\n\n"
else:
raise "Ollama: Could not create blob, Please try again."
# url = "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q2_K.gguf"