mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
[feature] Download Queue (#5225)
* add base definition of download manager * basic functionality working * add unit tests for download queue * add documentation and FastAPI route * fix docs * add missing test dependency; fix import ordering * fix file path length checking on windows * fix ruff check error * move release() into the __del__ method * disable testing of stderr messages due to issues with pytest capsys fixture * fix unsorted imports * harmonized implementation of start() and stop() calls in download and & install modules * Update invokeai/app/services/download/download_base.py Co-authored-by: Ryan Dick <ryanjdick3@gmail.com> * replace test datadir fixture with tmp_path * replace DownloadJobBase->DownloadJob in download manager documentation * make source and dest arguments to download_queue.download() an AnyHttpURL and Path respectively * fix pydantic typecheck errors in the download unit test * ruff formatting * add "job cancelled" as an event rather than an exception * fix ruff errors * Update invokeai/app/services/download/download_default.py Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com> * use threading.Event to stop service worker threads; handle unfinished job edge cases * remove dangling STOP job definition * fix ruff complaint * fix ruff check again * avoid race condition when start() and stop() are called simultaneously from different threads * avoid race condition in stop() when a job becomes active while shutting down --------- Co-authored-by: Lincoln Stein <lstein@gmail.com> Co-authored-by: Ryan Dick <ryanjdick3@gmail.com> Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Co-authored-by: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com>
This commit is contained in:
111
invokeai/app/api/routers/download_queue.py
Normal file
111
invokeai/app/api/routers/download_queue.py
Normal file
@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein
|
||||
"""FastAPI route for the download queue."""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import Body, Path, Response
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
from starlette.exceptions import HTTPException
|
||||
|
||||
from invokeai.app.services.download import (
|
||||
DownloadJob,
|
||||
UnknownJobIDException,
|
||||
)
|
||||
|
||||
from ..dependencies import ApiDependencies
|
||||
|
||||
download_queue_router = APIRouter(prefix="/v1/download_queue", tags=["download_queue"])
|
||||
|
||||
|
||||
@download_queue_router.get(
|
||||
"/",
|
||||
operation_id="list_downloads",
|
||||
)
|
||||
async def list_downloads() -> List[DownloadJob]:
|
||||
"""Get a list of active and inactive jobs."""
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
return queue.list_jobs()
|
||||
|
||||
|
||||
@download_queue_router.patch(
|
||||
"/",
|
||||
operation_id="prune_downloads",
|
||||
responses={
|
||||
204: {"description": "All completed jobs have been pruned"},
|
||||
400: {"description": "Bad request"},
|
||||
},
|
||||
)
|
||||
async def prune_downloads():
|
||||
"""Prune completed and errored jobs."""
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
queue.prune_jobs()
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@download_queue_router.post(
|
||||
"/i/",
|
||||
operation_id="download",
|
||||
)
|
||||
async def download(
|
||||
source: AnyHttpUrl = Body(description="download source"),
|
||||
dest: str = Body(description="download destination"),
|
||||
priority: int = Body(default=10, description="queue priority"),
|
||||
access_token: Optional[str] = Body(default=None, description="token for authorization to download"),
|
||||
) -> DownloadJob:
|
||||
"""Download the source URL to the file or directory indicted in dest."""
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
return queue.download(source, dest, priority, access_token)
|
||||
|
||||
|
||||
@download_queue_router.get(
|
||||
"/i/{id}",
|
||||
operation_id="get_download_job",
|
||||
responses={
|
||||
200: {"description": "Success"},
|
||||
404: {"description": "The requested download JobID could not be found"},
|
||||
},
|
||||
)
|
||||
async def get_download_job(
|
||||
id: int = Path(description="ID of the download job to fetch."),
|
||||
) -> DownloadJob:
|
||||
"""Get a download job using its ID."""
|
||||
try:
|
||||
job = ApiDependencies.invoker.services.download_queue.id_to_job(id)
|
||||
return job
|
||||
except UnknownJobIDException as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@download_queue_router.delete(
|
||||
"/i/{id}",
|
||||
operation_id="cancel_download_job",
|
||||
responses={
|
||||
204: {"description": "Job has been cancelled"},
|
||||
404: {"description": "The requested download JobID could not be found"},
|
||||
},
|
||||
)
|
||||
async def cancel_download_job(
|
||||
id: int = Path(description="ID of the download job to cancel."),
|
||||
):
|
||||
"""Cancel a download job using its ID."""
|
||||
try:
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
job = queue.id_to_job(id)
|
||||
queue.cancel_job(job)
|
||||
return Response(status_code=204)
|
||||
except UnknownJobIDException as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@download_queue_router.delete(
|
||||
"/i",
|
||||
operation_id="cancel_all_download_jobs",
|
||||
responses={
|
||||
204: {"description": "Download jobs have been cancelled"},
|
||||
},
|
||||
)
|
||||
async def cancel_all_download_jobs():
|
||||
"""Cancel all download jobs."""
|
||||
ApiDependencies.invoker.services.download_queue.cancel_all_jobs()
|
||||
return Response(status_code=204)
|
Reference in New Issue
Block a user