Merge pull request #17608 from sihyeonn/fix/sh-prevent-memory-leaks-on-ollama

fix: prevent memory leaks in file handling and HTTP connections
This commit is contained in:
Tim Jaeryang Baek 2025-09-19 22:22:08 -05:00 committed by GitHub
commit 820321ad55
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

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"