mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-03 13:19:23 +08:00
Add delete and clear all downloads funcitonalities.
This commit is contained in:
parent
fe4d0c9722
commit
4c82c708a7
@ -6,11 +6,13 @@ envelope used by ``app/assets/api/routes.py``:
|
||||
POST /api/download/enqueue
|
||||
GET /api/download
|
||||
POST /api/download/availability
|
||||
POST /api/download/clear
|
||||
POST /api/download/credentials
|
||||
GET /api/download/credentials
|
||||
GET /api/download/credentials/{id}
|
||||
DELETE /api/download/credentials/{id}
|
||||
GET /api/download/{id}
|
||||
DELETE /api/download/{id}
|
||||
POST /api/download/{id}/pause
|
||||
POST /api/download/{id}/resume
|
||||
POST /api/download/{id}/cancel
|
||||
@ -107,6 +109,12 @@ async def availability(request: web.Request) -> web.Response:
|
||||
return _ok({"models": await DOWNLOAD_MANAGER.availability(parsed.models)})
|
||||
|
||||
|
||||
@ROUTES.post("/api/download/clear")
|
||||
async def clear(request: web.Request) -> web.Response:
|
||||
deleted = await DOWNLOAD_MANAGER.clear()
|
||||
return _ok({"deleted": deleted})
|
||||
|
||||
|
||||
# ----- credentials (secrets are write-only) — must precede /{id} -----
|
||||
|
||||
|
||||
@ -164,6 +172,15 @@ async def get_download(request: web.Request) -> web.Response:
|
||||
return _ok(view)
|
||||
|
||||
|
||||
@ROUTES.delete("/api/download/{id}")
|
||||
async def delete_download(request: web.Request) -> web.Response:
|
||||
try:
|
||||
await DOWNLOAD_MANAGER.delete(request.match_info["id"])
|
||||
except DownloadError as e:
|
||||
return _from_download_error(e)
|
||||
return _ok({"deleted": True})
|
||||
|
||||
|
||||
@ROUTES.post("/api/download/{id}/pause")
|
||||
async def pause(request: web.Request) -> web.Response:
|
||||
try:
|
||||
|
||||
@ -10,7 +10,7 @@ from __future__ import annotations
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import delete, select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from app.database.db import create_session
|
||||
@ -84,6 +84,24 @@ def delete_download(download_id: str) -> None:
|
||||
session.commit()
|
||||
|
||||
|
||||
def delete_downloads(download_ids: list[str]) -> int:
|
||||
"""Delete many downloads in one transaction; returns the number removed.
|
||||
|
||||
Uses a bulk ``DELETE ... WHERE id IN (...)``. Segment rows are removed by
|
||||
the ``ON DELETE CASCADE`` foreign key (SQLite ``PRAGMA foreign_keys=ON`` is
|
||||
set in ``app/database/db.py``), so this stays consistent without loading the
|
||||
ORM relationship.
|
||||
"""
|
||||
if not download_ids:
|
||||
return 0
|
||||
with create_session() as session:
|
||||
result = session.execute(
|
||||
delete(Download).where(Download.id.in_(download_ids))
|
||||
)
|
||||
session.commit()
|
||||
return result.rowcount or 0
|
||||
|
||||
|
||||
def replace_segments(download_id: str, segments: list[dict]) -> None:
|
||||
"""Atomically replace the segment plan for a download."""
|
||||
with create_session() as session:
|
||||
|
||||
@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from typing import Callable, Optional
|
||||
|
||||
@ -210,6 +211,64 @@ class DownloadManager:
|
||||
# picked up the next time a slot frees. Pump in case a slot is free now.
|
||||
await self._scheduler.pump()
|
||||
|
||||
async def delete(self, download_id: str) -> None:
|
||||
"""Delete a terminal download so it stays gone from history.
|
||||
|
||||
Refuses to delete a live download so a record is never removed out from
|
||||
under a running worker; cancel it first. Any leftover ``.part`` temp
|
||||
file (e.g. from a failed transfer) is removed, but the finished model
|
||||
file on disk is never touched.
|
||||
"""
|
||||
if self._scheduler.get_job(download_id) is not None:
|
||||
raise DownloadError(
|
||||
"DOWNLOAD_ACTIVE",
|
||||
"Cannot delete a download that is still in progress.",
|
||||
status=409,
|
||||
)
|
||||
row = await asyncio.to_thread(queries.get_download, download_id)
|
||||
if row is None:
|
||||
raise DownloadError("NOT_FOUND", "No such download.", status=404)
|
||||
if row.status in _LIVE_STATUSES:
|
||||
raise DownloadError(
|
||||
"DOWNLOAD_ACTIVE",
|
||||
"Cannot delete a download that is still in progress.",
|
||||
status=409,
|
||||
)
|
||||
|
||||
try:
|
||||
os.remove(row.temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
await asyncio.to_thread(queries.delete_download, download_id)
|
||||
|
||||
async def clear(self) -> int:
|
||||
"""Delete all terminal downloads from history in one transaction.
|
||||
|
||||
Skips anything still live (queued/active/paused/verifying, or a running
|
||||
job) so an in-flight download is never removed out from under a worker.
|
||||
Finished model files on disk are never touched; only leftover ``.part``
|
||||
temp files from failed/cancelled transfers are removed. Returns the
|
||||
number of history rows deleted.
|
||||
"""
|
||||
|
||||
rows = await asyncio.to_thread(queries.list_downloads)
|
||||
deletable = [
|
||||
r
|
||||
for r in rows
|
||||
if r.status not in _LIVE_STATUSES
|
||||
and self._scheduler.get_job(r.id) is None
|
||||
]
|
||||
if not deletable:
|
||||
return 0
|
||||
for r in deletable:
|
||||
try:
|
||||
os.remove(r.temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
return await asyncio.to_thread(
|
||||
queries.delete_downloads, [r.id for r in deletable]
|
||||
)
|
||||
|
||||
# ----- read models -----
|
||||
|
||||
def _view(self, row) -> dict:
|
||||
|
||||
61
openapi.yaml
61
openapi.yaml
@ -2560,6 +2560,28 @@ paths:
|
||||
summary: Bulk model availability + status
|
||||
tags:
|
||||
- download
|
||||
/api/download/clear:
|
||||
post:
|
||||
description: |
|
||||
Delete all terminal downloads (completed, failed, cancelled) from history
|
||||
in one transaction, so the cleared history persists across reloads. Live
|
||||
downloads (queued, active, paused, verifying) are skipped. Finished model
|
||||
files on disk are never removed; only leftover .part temp files are cleaned up.
|
||||
operationId: clearDownloads
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
properties:
|
||||
deleted:
|
||||
description: Number of history rows removed.
|
||||
type: integer
|
||||
type: object
|
||||
description: History cleared
|
||||
summary: Clear terminal downloads from history
|
||||
tags:
|
||||
- download
|
||||
/api/download/credentials:
|
||||
get:
|
||||
description: List stored per-host credentials. Secrets are never returned; only masked metadata (last 4 chars, scheme, label).
|
||||
@ -2701,6 +2723,45 @@ paths:
|
||||
tags:
|
||||
- download
|
||||
/api/download/{id}:
|
||||
delete:
|
||||
description: |
|
||||
Delete a single terminal download from history so it stays gone across
|
||||
reloads. Refuses (409) to delete a live download (queued, active, paused,
|
||||
verifying) — cancel it first. The finished model file on disk is never
|
||||
removed; only a leftover .part temp file is cleaned up.
|
||||
operationId: deleteDownload
|
||||
parameters:
|
||||
- in: path
|
||||
name: id
|
||||
required: true
|
||||
schema:
|
||||
format: uuid
|
||||
type: string
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
properties:
|
||||
deleted:
|
||||
type: boolean
|
||||
type: object
|
||||
description: Deleted
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: No such download
|
||||
"409":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Download is still in progress
|
||||
summary: Delete a download from history
|
||||
tags:
|
||||
- download
|
||||
get:
|
||||
description: Get the current status + progress of a single download.
|
||||
operationId: getDownloadStatus
|
||||
|
||||
136
tests-unit/model_downloader_test/test_delete.py
Normal file
136
tests-unit/model_downloader_test/test_delete.py
Normal file
@ -0,0 +1,136 @@
|
||||
"""Unit tests for ``DownloadManager.delete`` and ``DownloadManager.clear``.
|
||||
|
||||
Deleting a terminal row must remove it from history for good (so it does not
|
||||
reappear on the next ``list``), leave live rows untouched, and clean up any
|
||||
leftover ``.part`` temp file without touching the finished model file.
|
||||
|
||||
``clear()`` is the bulk variant: it removes all terminal rows atomically, skips
|
||||
live ones, and returns the count of rows deleted.
|
||||
|
||||
Async methods are driven via ``asyncio.run`` so no pytest-asyncio plugin is
|
||||
required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from app.model_downloader.constants import DownloadStatus
|
||||
from app.model_downloader.database import queries
|
||||
from app.model_downloader.manager import DOWNLOAD_MANAGER, DownloadError
|
||||
|
||||
|
||||
def _insert(download_id: str, status: str, *, temp_path: str = "/tmp/none.part") -> None:
|
||||
queries.insert_download(
|
||||
{
|
||||
"id": download_id,
|
||||
"url": "https://huggingface.co/org/model.safetensors",
|
||||
"model_id": "loras/model.safetensors",
|
||||
"dest_path": "/tmp/model.safetensors",
|
||||
"temp_path": temp_path,
|
||||
"status": status,
|
||||
"priority": 0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_delete_removes_terminal_row_from_history():
|
||||
_insert("done", DownloadStatus.COMPLETED)
|
||||
|
||||
asyncio.run(DOWNLOAD_MANAGER.delete("done"))
|
||||
|
||||
assert queries.get_download("done") is None
|
||||
|
||||
|
||||
def test_delete_refuses_live_row():
|
||||
_insert("live", DownloadStatus.QUEUED)
|
||||
|
||||
with pytest.raises(DownloadError) as excinfo:
|
||||
asyncio.run(DOWNLOAD_MANAGER.delete("live"))
|
||||
|
||||
assert excinfo.value.code == "DOWNLOAD_ACTIVE"
|
||||
assert queries.get_download("live") is not None
|
||||
|
||||
|
||||
def test_delete_missing_row_raises_not_found():
|
||||
with pytest.raises(DownloadError) as excinfo:
|
||||
asyncio.run(DOWNLOAD_MANAGER.delete("nope"))
|
||||
|
||||
assert excinfo.value.code == "NOT_FOUND"
|
||||
|
||||
|
||||
def test_delete_removes_leftover_temp_file(tmp_path):
|
||||
partial = tmp_path / "model.safetensors.part"
|
||||
partial.write_bytes(b"partial")
|
||||
_insert("failed", DownloadStatus.FAILED, temp_path=str(partial))
|
||||
|
||||
asyncio.run(DOWNLOAD_MANAGER.delete("failed"))
|
||||
|
||||
assert not os.path.exists(partial)
|
||||
assert queries.get_download("failed") is None
|
||||
|
||||
|
||||
# ----- clear -----
|
||||
|
||||
|
||||
def test_clear_removes_all_terminal_rows():
|
||||
_insert("c-done", DownloadStatus.COMPLETED)
|
||||
_insert("c-fail", DownloadStatus.FAILED)
|
||||
_insert("c-canc", DownloadStatus.CANCELLED)
|
||||
|
||||
deleted = asyncio.run(DOWNLOAD_MANAGER.clear())
|
||||
|
||||
assert deleted == 3
|
||||
assert queries.get_download("c-done") is None
|
||||
assert queries.get_download("c-fail") is None
|
||||
assert queries.get_download("c-canc") is None
|
||||
|
||||
|
||||
def test_clear_skips_live_rows():
|
||||
_insert("cl-queued", DownloadStatus.QUEUED)
|
||||
_insert("cl-paused", DownloadStatus.PAUSED)
|
||||
_insert("cl-done", DownloadStatus.COMPLETED)
|
||||
|
||||
deleted = asyncio.run(DOWNLOAD_MANAGER.clear())
|
||||
|
||||
assert deleted == 1
|
||||
assert queries.get_download("cl-queued") is not None
|
||||
assert queries.get_download("cl-paused") is not None
|
||||
assert queries.get_download("cl-done") is None
|
||||
|
||||
|
||||
def test_clear_returns_zero_when_nothing_to_delete():
|
||||
_insert("cl-only-live", DownloadStatus.QUEUED)
|
||||
|
||||
deleted = asyncio.run(DOWNLOAD_MANAGER.clear())
|
||||
|
||||
assert deleted == 0
|
||||
assert queries.get_download("cl-only-live") is not None
|
||||
|
||||
|
||||
def test_clear_removes_leftover_temp_files(tmp_path):
|
||||
partial = tmp_path / "clear_partial.part"
|
||||
partial.write_bytes(b"partial data")
|
||||
finished = tmp_path / "finished.safetensors"
|
||||
finished.write_bytes(b"real model weights")
|
||||
|
||||
_insert("cl-part", DownloadStatus.FAILED, temp_path=str(partial))
|
||||
# The finished file is not the temp_path; temp_path for a completed download
|
||||
# no longer exists (already renamed), so use a non-existent path here to
|
||||
# verify clear() tolerates a missing temp file without raising.
|
||||
_insert("cl-comp", DownloadStatus.COMPLETED, temp_path=str(tmp_path / "gone.part"))
|
||||
|
||||
asyncio.run(DOWNLOAD_MANAGER.clear())
|
||||
|
||||
# Leftover .part from the failed download is cleaned up.
|
||||
assert not partial.exists()
|
||||
# Finished model file is never touched.
|
||||
assert finished.exists()
|
||||
|
||||
|
||||
def test_clear_empty_db_returns_zero():
|
||||
deleted = asyncio.run(DOWNLOAD_MANAGER.clear())
|
||||
assert deleted == 0
|
||||
Loading…
Reference in New Issue
Block a user