increase download chunksize for better speed

This commit is contained in:
Lincoln Stein
2023-09-09 22:15:34 -04:00
parent b7a6a536e6
commit 8636015d92
2 changed files with 16 additions and 12 deletions

View File

@ -280,11 +280,6 @@ class DownloadQueue(DownloadQueueBase):
if self._in_terminal_state(job):
del self._jobs[job.id]
if job.status == "error":
self._logger.warning(f"{job.source}: Download finished with error: {job.error}")
else:
self._logger.info(f"{job.source}: Download finished with status {job.status}")
self._queue.task_done()
def _fetch_metadata(self, job: DownloadJobBase) -> Tuple[AnyHttpUrl, ModelSourceMetadata]:
@ -380,7 +375,7 @@ class DownloadQueue(DownloadQueueBase):
if resp.status_code == 206 or exist_size > 0:
self._logger.warning(f"{dest}: partial file found. Resuming")
elif resp.status_code != 200:
raise HTTPError(resp.reason)
raise HTTPError(f"status code {resp.status_code}: {resp.reason}")
else:
self._logger.info(f"{job.source}: Downloading {job.destination}")
@ -389,7 +384,7 @@ class DownloadQueue(DownloadQueueBase):
self._update_job_status(job, DownloadJobStatus.RUNNING)
with open(dest, open_mode) as file:
for data in resp.iter_content(chunk_size=16384):
for data in resp.iter_content(chunk_size=100000):
if job.status != DownloadJobStatus.RUNNING: # cancelled, paused or errored
return
job.bytes += file.write(data)
@ -428,12 +423,12 @@ class DownloadQueue(DownloadQueueBase):
elif new_status in [DownloadJobStatus.COMPLETED, DownloadJobStatus.ERROR]:
job.job_ended = time.time()
if job.event_handlers:
try:
for handler in job.event_handlers:
for handler in job.event_handlers:
try:
handler(job)
except Exception as excp:
job.status = DownloadJobStatus.ERROR
job.error = excp
except Exception as excp:
job.error = excp
self._update_job_status(job, DownloadJobStatus.ERROR)
def _download_repoid(self, job: DownloadJobBase):
"""Download a job that holds a huggingface repoid."""

View File

@ -78,6 +78,7 @@ class ModelInstallBase(ABC):
store: Optional[ModelConfigStore] = None,
config: Optional[InvokeAIAppConfig] = None,
logger: Optional[InvokeAILogger] = None,
download: Optional[DownloadQueueBase] = None,
):
"""
Create ModelInstall object.
@ -88,6 +89,8 @@ class ModelInstallBase(ABC):
uses the system-wide default app config.
:param logger: Optional InvokeAILogger. If None passed,
uses the system-wide default logger.
:param download: Optional DownloadQueueBase object. If None passed,
a default queue object will be created.
"""
pass
@ -323,6 +326,7 @@ class ModelInstall(ModelInstallBase):
def complete_installation(job: DownloadJobBase):
if job.status == "completed":
self._logger.info(f"{job.source}: Download finished with status {job.status}. Installing.")
model_id = self.install(job.destination)
info = self._store.get_model(model_id)
info.source = str(job.source)
@ -334,8 +338,13 @@ class ModelInstall(ModelInstallBase):
info.thumbnail_url = metadata.thumbnail_url
self._store.update_model(model_id, info)
self._async_installs[job.source] = model_id
elif job.status == "error":
self._logger.warning(f"{job.source}: Download finished with error: {job.error}")
elif job.status == "cancelled":
self._logger.warning(f"{job.source}: Download cancelled at caller's request.")
jobs = queue.list_jobs()
if len(jobs) <= 1 and job.status in ["completed", "error", "cancelled"]:
self._tmpdir.cleanup()
self._tmpdir = None
# note - this is probably not going to work. The tmpdir