wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-06 22:51:18 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

Switch to asyncio.to_thread for db calls in job.py

Browse Source
This commit is contained in:
Talmaj Marinc 2026-06-30 10:50:54 +02:00
parent b51d18c641
commit 6cfb7e5b7f
1 changed files with 43 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
app/model_downloader/engine/job.py
View File

@ -149,38 +149,38 @@ class DownloadJob:
async def run(self) -> str: async def run(self) -> str:
"""Run to a terminal/paused state; returns the final status string.""" """Run to a terminal/paused state; returns the final status string."""
self._set_status(DownloadStatus.ACTIVE, error=None) await self._set_status(DownloadStatus.ACTIVE, error=None)
try: try:
pr = await self._probe_and_plan() pr = await self._probe_and_plan()
await self._transfer(pr) await self._transfer(pr)
await self._finalize() await self._finalize()
self._set_status(DownloadStatus.COMPLETED) await self._set_status(DownloadStatus.COMPLETED)
except Paused: except Paused:
await self._persist_progress(force=True) await self._persist_progress(force=True)
self._set_status(DownloadStatus.PAUSED) await self._set_status(DownloadStatus.PAUSED)
except Cancelled: except Cancelled:
await self._close_writer() await self._close_writer()
self._remove_temp() self._remove_temp()
self._set_status(DownloadStatus.CANCELLED) await self._set_status(DownloadStatus.CANCELLED)
except RemoteChanged: except RemoteChanged:
await self._reset_for_restart() await self._reset_for_restart()
self._set_status( await self._set_status(
DownloadStatus.QUEUED, error="remote file changed; restarting" DownloadStatus.QUEUED, error="remote file changed; restarting"
) )
except RetryableError as e: except RetryableError as e:
await self._persist_progress(force=True) await self._persist_progress(force=True)
self._set_status(DownloadStatus.QUEUED, error=str(e)) await self._set_status(DownloadStatus.QUEUED, error=str(e))
except FatalError as e: except FatalError as e:
await self._close_writer() await self._close_writer()
self._remove_temp() self._remove_temp()
self._set_status(DownloadStatus.FAILED, error=str(e)) await self._set_status(DownloadStatus.FAILED, error=str(e))
except Exception as e: # unexpected -> treat as retryable except Exception as e: # unexpected -> treat as retryable
logging.warning( logging.warning(
"[model_downloader] %s unexpected error: %s", "[model_downloader] %s unexpected error: %s",
self.spec.model_id, e, exc_info=True, self.spec.model_id, e, exc_info=True,
) )
await self._persist_progress(force=True) await self._persist_progress(force=True)
self._set_status(DownloadStatus.QUEUED, error=f"{type(e).__name__}: {e}") await self._set_status(DownloadStatus.QUEUED, error=f"{type(e).__name__}: {e}")
finally: finally:
await self._close_writer() await self._close_writer()
return self.state.status return self.state.status
@ -208,7 +208,8 @@ class DownloadJob:
self._etag = pr.etag or self._etag self._etag = pr.etag or self._etag
self.state.total_bytes = pr.total_bytes self.state.total_bytes = pr.total_bytes
queries.update_download( await asyncio.to_thread(
queries.update_download,
self.spec.download_id, self.spec.download_id,
final_url=pr.final_url, final_url=pr.final_url,
total_bytes=pr.total_bytes, total_bytes=pr.total_bytes,
@ -220,7 +221,7 @@ class DownloadJob:
seg_count = effective_segment_count( seg_count = effective_segment_count(
pr.total_bytes, pr.accept_ranges, max(1, args.download_segments) pr.total_bytes, pr.accept_ranges, max(1, args.download_segments)
) )
existing = queries.list_segments(self.spec.download_id) existing = await asyncio.to_thread(queries.list_segments, self.spec.download_id)
if ( if (
seg_count > 1 seg_count > 1
and existing and existing
@ -234,7 +235,8 @@ class DownloadJob:
] ]
elif seg_count > 1 and pr.total_bytes is not None: elif seg_count > 1 and pr.total_bytes is not None:
plans = plan_segments(pr.total_bytes, seg_count) plans = plan_segments(pr.total_bytes, seg_count)
queries.replace_segments( await asyncio.to_thread(
queries.replace_segments,
self.spec.download_id, self.spec.download_id,
[ [
{"idx": p.idx, "start_offset": p.start, "end_offset": p.end, "bytes_done": 0} {"idx": p.idx, "start_offset": p.start, "end_offset": p.end, "bytes_done": 0}
@ -244,7 +246,7 @@ class DownloadJob:
self.state.segments = [SegmentRuntime(p.idx, p.start, p.end, 0) for p in plans] self.state.segments = [SegmentRuntime(p.idx, p.start, p.end, 0) for p in plans]
else: else:
# Single-stream: one logical segment; bytes_done tracked on the row. # Single-stream: one logical segment; bytes_done tracked on the row.
row = queries.get_download(self.spec.download_id) row = await asyncio.to_thread(queries.get_download, self.spec.download_id)
resume_from = row.bytes_done if row else 0 resume_from = row.bytes_done if row else 0
end = (pr.total_bytes - 1) if pr.total_bytes else -1 end = (pr.total_bytes - 1) if pr.total_bytes else -1
# ``row.bytes_done`` may be the SUM of per-segment offsets from a # ``row.bytes_done`` may be the SUM of per-segment offsets from a
@ -263,9 +265,13 @@ class DownloadJob:
) )
resume_from = 0 resume_from = 0
self._remove_temp() self._remove_temp()
if queries.list_segments(self.spec.download_id): if await asyncio.to_thread(queries.list_segments, self.spec.download_id):
queries.replace_segments(self.spec.download_id, []) await asyncio.to_thread(
queries.update_download(self.spec.download_id, bytes_done=0) queries.replace_segments, self.spec.download_id, []
)
await asyncio.to_thread(
queries.update_download, self.spec.download_id, bytes_done=0
)
self.state.segments = [SegmentRuntime(0, 0, end, resume_from)] self.state.segments = [SegmentRuntime(0, 0, end, resume_from)]
self._recompute_bytes_done() self._recompute_bytes_done()
return pr return pr
@ -419,7 +425,7 @@ class DownloadJob:
async def _finalize(self) -> None: async def _finalize(self) -> None:
self._check_control() self._check_control()
await self._close_writer() await self._close_writer()
self._set_status(DownloadStatus.VERIFYING) await self._set_status(DownloadStatus.VERIFYING)
total = self.state.total_bytes total = self.state.total_bytes
segmented = len(self.state.segments) > 1 segmented = len(self.state.segments) > 1
@ -486,20 +492,28 @@ class DownloadJob:
self.state._last_time = now self.state._last_time = now
async def _persist_progress(self, force: bool = False) -> None: async def _persist_progress(self, force: bool = False) -> None:
# Both the DB write and the websocket notify are gated by the same
# throttle: persisting hits SQLite, and notifying broadcasts to every
# client, so doing either per-chunk (small --download-chunk-size or
# many concurrent segments) would overwhelm both. Skip entirely inside
# the window; the next persist (or a forced one) ships the latest bytes.
now = time.monotonic() now = time.monotonic()
if not force and now - self._last_persist < _PERSIST_INTERVAL: if not force and now - self._last_persist < _PERSIST_INTERVAL:
if self._notify:
self._notify(self.spec.download_id)
return return
self._last_persist = now self._last_persist = now
# SQLite is blocking; run it off the event loop per the queries module
# contract so progress persists don't stall the web server.
await asyncio.to_thread(self._write_progress)
if self._notify:
self._notify(self.spec.download_id)
def _write_progress(self) -> None:
queries.update_download(self.spec.download_id, bytes_done=self.state.bytes_done) queries.update_download(self.spec.download_id, bytes_done=self.state.bytes_done)
for seg in self.state.segments: for seg in self.state.segments:
if seg.end >= seg.start: # skip unknown-size sentinel if seg.end >= seg.start: # skip unknown-size sentinel
queries.update_segment_progress( queries.update_segment_progress(
self.spec.download_id, seg.idx, seg.bytes_done self.spec.download_id, seg.idx, seg.bytes_done
) )
if self._notify:
self._notify(self.spec.download_id)
async def _reset_for_restart(self) -> None: async def _reset_for_restart(self) -> None:
await self._close_writer() await self._close_writer()
@ -507,9 +521,13 @@ class DownloadJob:
for seg in self.state.segments: for seg in self.state.segments:
seg.bytes_done = 0 seg.bytes_done = 0
self.state.bytes_done = 0 self.state.bytes_done = 0
queries.update_download(self.spec.download_id, bytes_done=0) await asyncio.to_thread(
if queries.list_segments(self.spec.download_id): queries.update_download, self.spec.download_id, bytes_done=0
queries.replace_segments(self.spec.download_id, []) )
if await asyncio.to_thread(queries.list_segments, self.spec.download_id):
await asyncio.to_thread(
queries.replace_segments, self.spec.download_id, []
)
async def _close_writer(self) -> None: async def _close_writer(self) -> None:
if self._writer is not None: if self._writer is not None:
@ -543,7 +561,7 @@ class DownloadJob:
"[model_downloader] could not remove %s: %s", self.spec.temp_path, e "[model_downloader] could not remove %s: %s", self.spec.temp_path, e
) )
def _set_status(self, status: str, error: Optional[str] = None) -> None: async def _set_status(self, status: str, error: Optional[str] = None) -> None:
# ``error`` is authoritative: passing None clears any prior failure # ``error`` is authoritative: passing None clears any prior failure
# text so transitions out of a failure state (retry/success) don't # text so transitions out of a failure state (retry/success) don't
# leave stale messages on RuntimeState or in the persisted row. # leave stale messages on RuntimeState or in the persisted row.
@ -553,6 +571,6 @@ class DownloadJob:
if status == DownloadStatus.QUEUED: if status == DownloadStatus.QUEUED:
fields["attempts"] = self.spec.attempts + 1 fields["attempts"] = self.spec.attempts + 1
self.spec.attempts += 1 self.spec.attempts += 1
queries.update_download(self.spec.download_id, **fields) await asyncio.to_thread(queries.update_download, self.spec.download_id, **fields)
if self._notify: if self._notify:
self._notify(self.spec.download_id) self._notify(self.spec.download_id)