瀏覽代碼

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>
Sihyeon Jang 2 周之前
父節點
當前提交
7042318c34
共有 1 個文件被更改,包括 21 次插入19 次删除
  1. 21 19
      backend/open_webui/routers/ollama.py

+ 21 - 19
backend/open_webui/routers/ollama.py

@@ -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)
-
-                    url = f"{ollama_url}/api/blobs/sha256:{hashed}"
-                    response = requests.post(url, data=file)
-
-                    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."
+                    file.close()
+
+                    with open(file_path, "rb") as file:
+                        chunk_size = 1024 * 1024 * 2
+                        hashed = calculate_sha256(file, chunk_size)
+
+                        url = f"{ollama_url}/api/blobs/sha256:{hashed}"
+                        with requests.Session() as session:
+                            response = session.post(url, data=file, timeout=30)
+
+                            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"