Improve aria2 and Python download robustness and retry handling
- Drain aria2 daemon stderr in a background thread to prevent Windows pipe buffer saturation and deadlock - Promote critical disk write and permission errors from aria2 stderr to warnings - Automatically refresh expired signed CDN URLs and resume downloads via aria2 continue=true - Shield addUri against mid-registration cancellations to prevent orphaned downloads - Retry Range requests directly on target URL when redirects drop the Range header in Python downloads - Add auto-reconnect stream loop on transient socket disruptions with seamless partial-hash resumption - Increase default download chunk size to 4MB for improved I/O throughput - Parse HTTP 429 Retry-After headers for informative rate limit reporting - Add unit test suite covering new aria2 and Python downloader resilience behaviors
This commit is contained in:
Dariusz L
+18
-1
@@ -56,12 +56,15 @@ except ImportError:
|
||||
log = create_module_logger("core.downloader")
|
||||
|
||||
SPEED_HISTORY_SIZE = 5
|
||||
CHUNK_SIZE = 1024 * 1024
|
||||
CHUNK_SIZE = 4 * 1024 * 1024
|
||||
CLI_LOG_INTERVAL = 5
|
||||
ARIA2_RPC_TIMEOUT = (2, 5)
|
||||
ARIA2_STATUS_RPC_RETRIES = 4
|
||||
ARIA2_STATUS_RPC_RETRY_DELAY = 0.15
|
||||
ARIA2_IDLE_STOP_SECONDS = 5 * 60
|
||||
MAX_ARIA2_URL_RECOVERY_ATTEMPTS = 2
|
||||
PYTHON_DOWNLOAD_MAX_RETRIES = 2
|
||||
PYTHON_DOWNLOAD_RETRY_DELAY = 0.5
|
||||
# Keep validation aligned with the install location used by the aria2 installer.
|
||||
# ARIA2_INSTALL_ROOT resolves the custom node package first, so paths reached
|
||||
# through a shared ComfyUI/custom_nodes symlink remain valid while external
|
||||
@@ -95,10 +98,14 @@ context.ARIA2_RPC_TIMEOUT = ARIA2_RPC_TIMEOUT
|
||||
context.ARIA2_STATUS_RPC_RETRIES = ARIA2_STATUS_RPC_RETRIES
|
||||
context.ARIA2_STATUS_RPC_RETRY_DELAY = ARIA2_STATUS_RPC_RETRY_DELAY
|
||||
context.ARIA2_IDLE_STOP_SECONDS = ARIA2_IDLE_STOP_SECONDS
|
||||
context.MAX_ARIA2_URL_RECOVERY_ATTEMPTS = MAX_ARIA2_URL_RECOVERY_ATTEMPTS
|
||||
context.PYTHON_DOWNLOAD_MAX_RETRIES = PYTHON_DOWNLOAD_MAX_RETRIES
|
||||
context.PYTHON_DOWNLOAD_RETRY_DELAY = PYTHON_DOWNLOAD_RETRY_DELAY
|
||||
context.MANAGED_ARIA2_ROOT = MANAGED_ARIA2_ROOT
|
||||
context.HF_XET_ARIA2_AUTH_HOSTS = HF_XET_ARIA2_AUTH_HOSTS
|
||||
|
||||
context.aria2_process = None
|
||||
context.aria2_recent_stderr = deque(maxlen=50)
|
||||
context.aria2_rpc_url = ""
|
||||
context.aria2_rpc_secret = ""
|
||||
context.aria2_rpc_lock = threading.Lock()
|
||||
@@ -247,6 +254,16 @@ context._aria2_action_error_is_ok = aria2_backend.aria2_action_error_is_ok
|
||||
context._is_aria2_missing_control_file_error = (
|
||||
aria2_backend.is_aria2_missing_control_file_error
|
||||
)
|
||||
context._is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr
|
||||
context.is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr
|
||||
context._is_aria2_recoverable_url_error = (
|
||||
aria2_backend.is_aria2_recoverable_url_error
|
||||
)
|
||||
context.is_aria2_recoverable_url_error = (
|
||||
aria2_backend.is_aria2_recoverable_url_error
|
||||
)
|
||||
context._drain_aria2_stderr = _bind(aria2_backend.drain_aria2_stderr)
|
||||
context.drain_aria2_stderr = context._drain_aria2_stderr
|
||||
context.Aria2Error = aria2_backend.Aria2Error
|
||||
context._find_free_port = aria2_backend.find_free_port
|
||||
context._parse_aria2_int = aria2_backend.parse_aria2_int
|
||||
|
||||
+241
-53
@@ -87,6 +87,84 @@ def parse_aria2_int(value: Any) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
ARIA2_IMPORTANT_STDERR_MARKERS = (
|
||||
"write disk cache flush failure",
|
||||
"error when trying to flush write cache",
|
||||
"failed to write into the file",
|
||||
"failed to open the file",
|
||||
"failed to seek the file",
|
||||
"no space left on device",
|
||||
"not enough space on the disk",
|
||||
"input/output error",
|
||||
"permission denied",
|
||||
"access is denied",
|
||||
"disk quota exceeded",
|
||||
)
|
||||
|
||||
MAX_ARIA2_URL_RECOVERY_ATTEMPTS = 2
|
||||
|
||||
|
||||
def is_aria2_important_stderr(line: str) -> bool:
|
||||
"""Identify critical disk or permission error messages in aria2's stderr."""
|
||||
lowered = line.lower()
|
||||
return any(marker in lowered for marker in ARIA2_IMPORTANT_STDERR_MARKERS)
|
||||
|
||||
|
||||
def is_aria2_recoverable_url_error(error: Any) -> bool:
|
||||
"""Identify aria2 transfer errors that may be resolved by refreshing a signed URL."""
|
||||
message = str(error or "").lower()
|
||||
return (
|
||||
"no uri available" in message
|
||||
or "403 forbidden" in message
|
||||
or "401 unauthorized" in message
|
||||
or "410 gone" in message
|
||||
or "status code 403" in message
|
||||
or "status code 401" in message
|
||||
or "bad or unexpected: 403" in message
|
||||
or "bad or unexpected: 401" in message
|
||||
)
|
||||
|
||||
|
||||
def drain_aria2_stderr(
|
||||
process: Any,
|
||||
recent_lines: Any = None,
|
||||
*,
|
||||
dependencies: Any = None,
|
||||
) -> None:
|
||||
"""Continuously drain aria2's stderr pipe so it never saturates Windows pipe buffers."""
|
||||
if not process or not getattr(process, "stderr", None):
|
||||
return
|
||||
stderr = process.stderr
|
||||
try:
|
||||
while True:
|
||||
raw_line = stderr.readline()
|
||||
if not raw_line:
|
||||
break
|
||||
if isinstance(raw_line, bytes):
|
||||
line = raw_line.decode("utf-8", errors="replace").rstrip()
|
||||
else:
|
||||
line = str(raw_line).rstrip()
|
||||
if not line:
|
||||
continue
|
||||
if recent_lines is not None:
|
||||
try:
|
||||
recent_lines.append(line)
|
||||
except Exception:
|
||||
pass
|
||||
if is_aria2_important_stderr(line):
|
||||
log.warning(f"aria2: {line}")
|
||||
else:
|
||||
log.debug(f"aria2 stderr: {line}")
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
if stderr and not getattr(stderr, "closed", False):
|
||||
stderr.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def is_aria2_missing_control_file_error(error: Any) -> bool:
|
||||
"""Identify aria2's stale-output error caused by a missing .aria2 file."""
|
||||
message = str(error or "").lower()
|
||||
@@ -440,6 +518,12 @@ def stop_aria2_daemon(
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if getattr(process, "stderr", None) and not getattr(process.stderr, "closed", False):
|
||||
process.stderr.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
facade.aria2_process = None
|
||||
facade.aria2_rpc_url = ""
|
||||
facade.aria2_rpc_secret = ""
|
||||
@@ -557,6 +641,15 @@ def ensure_aria2_daemon(
|
||||
process.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
if (
|
||||
process is not None
|
||||
and getattr(process, "stderr", None)
|
||||
and not getattr(process.stderr, "closed", False)
|
||||
):
|
||||
try:
|
||||
process.stderr.close()
|
||||
except Exception:
|
||||
pass
|
||||
facade.aria2_process = None
|
||||
facade.aria2_process_started_by_resolver = False
|
||||
|
||||
@@ -592,6 +685,7 @@ def ensure_aria2_daemon(
|
||||
creationflags = getattr(facade.subprocess, "CREATE_NO_WINDOW", 0)
|
||||
|
||||
log.info(f"Starting aria2 RPC daemon from {executable}")
|
||||
facade.aria2_recent_stderr = facade.deque(maxlen=50)
|
||||
facade.aria2_process = facade.subprocess.Popen(
|
||||
command,
|
||||
stdout=facade.subprocess.DEVNULL,
|
||||
@@ -600,18 +694,22 @@ def ensure_aria2_daemon(
|
||||
)
|
||||
facade.aria2_process_started_by_resolver = True
|
||||
|
||||
drain_thread = facade.threading.Thread(
|
||||
target=facade._drain_aria2_stderr,
|
||||
args=(facade.aria2_process, facade.aria2_recent_stderr),
|
||||
daemon=True,
|
||||
name="aria2-stderr-drain",
|
||||
)
|
||||
drain_thread.start()
|
||||
|
||||
start_time = facade.time.time()
|
||||
last_error = ""
|
||||
while facade.time.time() - start_time < 10:
|
||||
process = facade.aria2_process
|
||||
if process.poll() is not None:
|
||||
stderr = ""
|
||||
try:
|
||||
stderr = (
|
||||
process.stderr.read() if process.stderr else b""
|
||||
).decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
stderr = ""
|
||||
facade.time.sleep(0.05)
|
||||
recent = getattr(facade, "aria2_recent_stderr", None)
|
||||
stderr = "\n".join(recent) if recent else ""
|
||||
raise Aria2Error(
|
||||
"aria2 RPC process exited early with code "
|
||||
f"{process.returncode}: {stderr.strip()}"
|
||||
@@ -736,6 +834,12 @@ def download_file_with_aria2(
|
||||
start_time = facade.time.time()
|
||||
filename = facade.get_filename_from_path(dest_path)
|
||||
expected_sha256 = facade._extract_expected_sha256(metadata)
|
||||
max_url_recovery = getattr(
|
||||
facade,
|
||||
"MAX_ARIA2_URL_RECOVERY_ATTEMPTS",
|
||||
MAX_ARIA2_URL_RECOVERY_ATTEMPTS,
|
||||
)
|
||||
url_recovery_attempts = 0
|
||||
|
||||
with facade.download_lock:
|
||||
facade.download_progress[download_id] = create_initial_progress(
|
||||
@@ -757,45 +861,60 @@ def download_file_with_aria2(
|
||||
headers,
|
||||
)
|
||||
|
||||
options: Dict[str, Any] = {
|
||||
"dir": facade.os.path.dirname(dest_path),
|
||||
"out": filename,
|
||||
"continue": "true",
|
||||
"max-connection-per-server": "4",
|
||||
"split": "4",
|
||||
"min-split-size": "1M",
|
||||
"allow-overwrite": "true",
|
||||
"auto-file-renaming": "false",
|
||||
"file-allocation": "none",
|
||||
"no-want-digest-header": "true",
|
||||
# Redirects were already resolved and validated above. Keeping them
|
||||
# disabled prevents sensitive headers from reaching another host.
|
||||
"max-redirect": "0",
|
||||
}
|
||||
# Checksum and check-integrity are deliberately omitted from aria2 options.
|
||||
# On Windows, aria2c's single-threaded, unaccelerated software checksum
|
||||
# verification blocks download completion and freezes the process at 100%
|
||||
# for many minutes on 20+ GB models. Instead, verification is handled
|
||||
# asynchronously in the background via hardware-accelerated Python hashlib.
|
||||
user_agent = facade._get_header_value(request_headers, "User-Agent")
|
||||
referer = facade._get_header_value(request_headers, "Referer")
|
||||
if user_agent:
|
||||
options["user-agent"] = user_agent
|
||||
if referer:
|
||||
options["referer"] = referer
|
||||
def build_transfer_options(target_headers: Dict[str, str]) -> Dict[str, Any]:
|
||||
opts: Dict[str, Any] = {
|
||||
"dir": facade.os.path.dirname(dest_path),
|
||||
"out": filename,
|
||||
"continue": "true",
|
||||
"max-connection-per-server": "4",
|
||||
"split": "4",
|
||||
"min-split-size": "1M",
|
||||
"allow-overwrite": "true",
|
||||
"auto-file-renaming": "false",
|
||||
"file-allocation": "none",
|
||||
"no-want-digest-header": "true",
|
||||
# Redirects were already resolved and validated above. Keeping them
|
||||
# disabled prevents sensitive headers from reaching another host.
|
||||
"max-redirect": "0",
|
||||
}
|
||||
# Checksum and check-integrity are deliberately omitted from aria2 options.
|
||||
# On Windows, aria2c's single-threaded, unaccelerated software checksum
|
||||
# verification blocks download completion and freezes the process at 100%
|
||||
# for many minutes on 20+ GB models. Instead, verification is handled
|
||||
# asynchronously in the background via hardware-accelerated Python hashlib.
|
||||
user_agent = facade._get_header_value(target_headers, "User-Agent")
|
||||
referer = facade._get_header_value(target_headers, "Referer")
|
||||
if user_agent:
|
||||
opts["user-agent"] = user_agent
|
||||
if referer:
|
||||
opts["referer"] = referer
|
||||
|
||||
header_values = [
|
||||
f"{key}: {value}"
|
||||
for key, value in request_headers.items()
|
||||
if str(key).lower() not in {"user-agent", "referer"}
|
||||
]
|
||||
if header_values:
|
||||
options["header"] = header_values
|
||||
header_values = [
|
||||
f"{key}: {value}"
|
||||
for key, value in target_headers.items()
|
||||
if str(key).lower() not in {"user-agent", "referer"}
|
||||
]
|
||||
if header_values:
|
||||
opts["header"] = header_values
|
||||
return opts
|
||||
|
||||
options = build_transfer_options(request_headers)
|
||||
|
||||
def add_aria2_transfer(
|
||||
target_url: str = "",
|
||||
target_options: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
preserve_progress: bool = False,
|
||||
) -> str:
|
||||
eff_url = target_url or aria2_url
|
||||
eff_options = target_options if target_options is not None else options
|
||||
|
||||
if download_id in facade.cancelled_downloads:
|
||||
raise Aria2Error("Download was cancelled")
|
||||
|
||||
def add_aria2_transfer() -> str:
|
||||
new_gid = facade._aria2_rpc(
|
||||
"aria2.addUri",
|
||||
[[aria2_url], options],
|
||||
[[eff_url], eff_options],
|
||||
)
|
||||
if not isinstance(new_gid, str) or not new_gid:
|
||||
raise Aria2Error("aria2 did not return a download gid")
|
||||
@@ -805,18 +924,41 @@ def download_file_with_aria2(
|
||||
"gid": new_gid,
|
||||
"path": dest_path,
|
||||
}
|
||||
|
||||
if download_id in facade.cancelled_downloads:
|
||||
try:
|
||||
facade._aria2_rpc("aria2.forceRemove", [new_gid])
|
||||
except Exception:
|
||||
pass
|
||||
raise Aria2Error("Download was cancelled")
|
||||
|
||||
with facade.download_lock:
|
||||
facade.download_progress[download_id].update(
|
||||
{
|
||||
"aria2_gid": new_gid,
|
||||
"status": "downloading",
|
||||
"error": None,
|
||||
"progress": 0,
|
||||
"downloaded": 0,
|
||||
"total_size": 0,
|
||||
"speed": 0,
|
||||
}
|
||||
progress_dict = facade.download_progress.setdefault(
|
||||
download_id,
|
||||
create_initial_progress(
|
||||
url=url,
|
||||
path=dest_path,
|
||||
filename=filename,
|
||||
directory=facade.os.path.dirname(dest_path),
|
||||
download_backend="aria2",
|
||||
start_time=start_time,
|
||||
),
|
||||
)
|
||||
updates: Dict[str, Any] = {
|
||||
"aria2_gid": new_gid,
|
||||
"status": "downloading",
|
||||
"error": None,
|
||||
}
|
||||
if not preserve_progress:
|
||||
updates.update(
|
||||
{
|
||||
"progress": 0,
|
||||
"downloaded": 0,
|
||||
"total_size": 0,
|
||||
"speed": 0,
|
||||
}
|
||||
)
|
||||
progress_dict.update(updates)
|
||||
return new_gid
|
||||
|
||||
recovery_attempted = False
|
||||
@@ -1000,8 +1142,49 @@ def download_file_with_aria2(
|
||||
settings,
|
||||
)
|
||||
if recovered:
|
||||
gid = add_aria2_transfer()
|
||||
gid = add_aria2_transfer(
|
||||
aria2_url,
|
||||
options,
|
||||
preserve_progress=True,
|
||||
)
|
||||
continue
|
||||
|
||||
if (
|
||||
url_recovery_attempts < max_url_recovery
|
||||
and download_id not in facade.cancelled_downloads
|
||||
and facade._is_aria2_recoverable_url_error(error_msg)
|
||||
):
|
||||
url_recovery_attempts += 1
|
||||
log.warning(
|
||||
f"aria2 transfer for {filename} failed with '{error_msg}'; "
|
||||
f"refreshing signed URL and resuming partial download "
|
||||
f"(attempt {url_recovery_attempts}/{max_url_recovery})"
|
||||
)
|
||||
try:
|
||||
facade._aria2_rpc("aria2.removeDownloadResult", [gid])
|
||||
except Exception:
|
||||
pass
|
||||
facade.time.sleep(1.0)
|
||||
try:
|
||||
facade._ensure_aria2_daemon(settings)
|
||||
aria2_url, request_headers = (
|
||||
facade._resolve_download_url_for_aria2(
|
||||
url,
|
||||
headers,
|
||||
)
|
||||
)
|
||||
options = build_transfer_options(request_headers)
|
||||
gid = add_aria2_transfer(
|
||||
aria2_url,
|
||||
options,
|
||||
preserve_progress=True,
|
||||
)
|
||||
continue
|
||||
except Exception as refresh_exc:
|
||||
log.warning(
|
||||
f"Could not refresh signed URL for {filename}: {refresh_exc}"
|
||||
)
|
||||
|
||||
with facade.download_lock:
|
||||
facade.download_progress[download_id]["status"] = "error"
|
||||
facade.download_progress[download_id]["error"] = error_msg
|
||||
@@ -1022,6 +1205,11 @@ def download_file_with_aria2(
|
||||
facade.time.sleep(0.5)
|
||||
|
||||
except Exception as exc:
|
||||
if download_id in facade.cancelled_downloads:
|
||||
facade._delete_partial_download_files(dest_path)
|
||||
facade.cancelled_downloads.discard(download_id)
|
||||
result["error"] = "Download cancelled"
|
||||
return result
|
||||
error_msg = facade._sanitize_download_error(exc)
|
||||
with facade.download_lock:
|
||||
if download_id in facade.download_progress:
|
||||
|
||||
+182
-36
@@ -10,6 +10,29 @@ class DownloadCancelled(Exception):
|
||||
"""The caller cancelled before the download response became available."""
|
||||
|
||||
|
||||
def parse_retry_after(value: Optional[str]) -> Optional[int]:
|
||||
"""Parse HTTP Retry-After header into integer seconds."""
|
||||
if not value or not str(value).strip():
|
||||
return None
|
||||
val = str(value).strip()
|
||||
try:
|
||||
return max(1, int(val))
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
from datetime import datetime, timezone
|
||||
from email.utils import parsedate_to_datetime
|
||||
|
||||
parsed = parsedate_to_datetime(val)
|
||||
now = datetime.now(timezone.utc)
|
||||
if parsed.tzinfo is None:
|
||||
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
||||
delta = (parsed - now).total_seconds()
|
||||
return max(1, int(delta))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def request_download_response(facade: Any, download_id: str, url: str, **kwargs: Any):
|
||||
"""Allow cancellation while waiting for HTTP headers without opening a file."""
|
||||
finished = threading.Event()
|
||||
@@ -240,6 +263,37 @@ def download_file(
|
||||
is_resumed = bool(
|
||||
existing_size > 0 and getattr(response, "status_code", None) == 206
|
||||
)
|
||||
if not is_resumed and existing_size > 0 and final_url and final_url != url:
|
||||
log.info(
|
||||
"Range request was not honored across redirect (HTTP 200). "
|
||||
f"Retrying Range request directly on target URL for {filename}..."
|
||||
)
|
||||
try:
|
||||
response.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
retry_resp, retry_final_url, _ = request_download_response(
|
||||
facade,
|
||||
download_id,
|
||||
final_url,
|
||||
headers=request_headers,
|
||||
stream=True,
|
||||
timeout=(20, 60),
|
||||
)
|
||||
retry_resp.raise_for_status()
|
||||
if getattr(retry_resp, "status_code", None) == 206:
|
||||
response = retry_resp
|
||||
final_url = retry_final_url
|
||||
is_resumed = True
|
||||
log.info(
|
||||
f"Successfully resumed download from target URL at byte {existing_size}"
|
||||
)
|
||||
else:
|
||||
response = retry_resp
|
||||
except Exception as retry_exc:
|
||||
log.debug(f"Direct Range retry on target URL failed: {retry_exc}")
|
||||
|
||||
if not is_resumed and existing_size > 0:
|
||||
log.info(
|
||||
f"Server returned HTTP {getattr(response, 'status_code', None)} without partial content support. "
|
||||
@@ -286,51 +340,135 @@ def download_file(
|
||||
break
|
||||
sha256_hasher.update(existing_chunk)
|
||||
|
||||
max_stream_retries = getattr(facade, "PYTHON_DOWNLOAD_MAX_RETRIES", 2)
|
||||
stream_retry_delay = getattr(facade, "PYTHON_DOWNLOAD_RETRY_DELAY", 0.5)
|
||||
stream_attempts = 0
|
||||
|
||||
cancelled = False
|
||||
file_mode = "ab" if is_resumed else "wb"
|
||||
with open(partial_path, file_mode) as file_handle:
|
||||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||||
if download_id in cancelled_downloads:
|
||||
cancelled = True
|
||||
while True:
|
||||
try:
|
||||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||||
if download_id in cancelled_downloads:
|
||||
cancelled = True
|
||||
break
|
||||
|
||||
if chunk:
|
||||
file_handle.write(chunk)
|
||||
if sha256_hasher is not None:
|
||||
sha256_hasher.update(chunk)
|
||||
downloaded += len(chunk)
|
||||
|
||||
current_time = time.time()
|
||||
time_delta = current_time - last_speed_update
|
||||
if time_delta >= 0.5:
|
||||
bytes_delta = downloaded - last_downloaded
|
||||
instant_speed = bytes_delta / time_delta if time_delta > 0 else 0
|
||||
speed_history.append(instant_speed)
|
||||
smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0
|
||||
last_speed_update = current_time
|
||||
last_downloaded = downloaded
|
||||
|
||||
with download_lock:
|
||||
download_progress[download_id]["downloaded"] = downloaded
|
||||
download_progress[download_id]["speed"] = int(smoothed_speed)
|
||||
if total_size > 0:
|
||||
download_progress[download_id]["progress"] = int((downloaded / total_size) * 100)
|
||||
|
||||
if current_time - last_cli_log >= cli_log_interval:
|
||||
last_cli_log = current_time
|
||||
progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0
|
||||
downloaded_str = facade.format_bytes(downloaded)
|
||||
total_str = facade.format_bytes(total_size) if total_size > 0 else "?"
|
||||
speed_str = facade.format_bytes(int(smoothed_speed)) + "/s"
|
||||
log.info(f"Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}")
|
||||
else:
|
||||
with download_lock:
|
||||
download_progress[download_id]["downloaded"] = downloaded
|
||||
if total_size > 0:
|
||||
download_progress[download_id]["progress"] = int((downloaded / total_size) * 100)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(downloaded, total_size)
|
||||
|
||||
if cancelled or download_id in cancelled_downloads:
|
||||
break
|
||||
|
||||
# Check if stream closed prematurely
|
||||
if total_size > 0 and downloaded < total_size:
|
||||
raise requests.exceptions.RequestException(
|
||||
f"Connection closed prematurely ({facade.format_bytes(downloaded)} / {total_size_str})"
|
||||
)
|
||||
|
||||
# Finished transferring all data
|
||||
break
|
||||
|
||||
if chunk:
|
||||
file_handle.write(chunk)
|
||||
if sha256_hasher is not None:
|
||||
sha256_hasher.update(chunk)
|
||||
downloaded += len(chunk)
|
||||
except (requests.exceptions.RequestException, OSError, TimeoutError) as stream_exc:
|
||||
if cancelled or download_id in cancelled_downloads:
|
||||
cancelled = True
|
||||
break
|
||||
|
||||
current_time = time.time()
|
||||
time_delta = current_time - last_speed_update
|
||||
if time_delta >= 0.5:
|
||||
bytes_delta = downloaded - last_downloaded
|
||||
instant_speed = bytes_delta / time_delta if time_delta > 0 else 0
|
||||
speed_history.append(instant_speed)
|
||||
smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0
|
||||
last_speed_update = current_time
|
||||
last_downloaded = downloaded
|
||||
if stream_attempts < max_stream_retries:
|
||||
stream_attempts += 1
|
||||
file_handle.flush()
|
||||
try:
|
||||
response.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with download_lock:
|
||||
download_progress[download_id]["downloaded"] = downloaded
|
||||
download_progress[download_id]["speed"] = int(smoothed_speed)
|
||||
if total_size > 0:
|
||||
download_progress[download_id]["progress"] = int((downloaded / total_size) * 100)
|
||||
log.warning(
|
||||
f"Download stream interrupted for {filename} ({facade._sanitize_download_error(stream_exc)}); "
|
||||
f"auto-reconnecting and resuming from byte {downloaded} "
|
||||
f"(attempt {stream_attempts}/{max_stream_retries})..."
|
||||
)
|
||||
time.sleep(stream_retry_delay)
|
||||
|
||||
if current_time - last_cli_log >= cli_log_interval:
|
||||
last_cli_log = current_time
|
||||
progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0
|
||||
downloaded_str = facade.format_bytes(downloaded)
|
||||
total_str = facade.format_bytes(total_size) if total_size > 0 else "?"
|
||||
speed_str = facade.format_bytes(int(smoothed_speed)) + "/s"
|
||||
log.info(f"Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}")
|
||||
else:
|
||||
with download_lock:
|
||||
download_progress[download_id]["downloaded"] = downloaded
|
||||
if total_size > 0:
|
||||
download_progress[download_id]["progress"] = int((downloaded / total_size) * 100)
|
||||
resume_headers = facade.build_download_headers(url, headers)
|
||||
if downloaded > 0:
|
||||
resume_headers["Range"] = f"bytes={downloaded}-"
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(downloaded, total_size)
|
||||
reconnected = False
|
||||
if final_url and final_url != url:
|
||||
try:
|
||||
resume_resp, resume_final_url, _ = request_download_response(
|
||||
facade,
|
||||
download_id,
|
||||
final_url,
|
||||
headers=resume_headers,
|
||||
stream=True,
|
||||
timeout=(20, 60),
|
||||
)
|
||||
resume_resp.raise_for_status()
|
||||
if getattr(resume_resp, "status_code", None) == 206:
|
||||
response = resume_resp
|
||||
final_url = resume_final_url
|
||||
reconnected = True
|
||||
except Exception as rec_exc:
|
||||
log.debug(f"Direct reconnect on target URL failed: {rec_exc}")
|
||||
|
||||
if not reconnected:
|
||||
try:
|
||||
resume_resp, resume_final_url, _ = request_download_response(
|
||||
facade,
|
||||
download_id,
|
||||
url,
|
||||
headers=resume_headers,
|
||||
stream=True,
|
||||
timeout=(20, 60),
|
||||
)
|
||||
resume_resp.raise_for_status()
|
||||
if getattr(resume_resp, "status_code", None) == 206 or downloaded == 0:
|
||||
response = resume_resp
|
||||
final_url = resume_final_url
|
||||
reconnected = True
|
||||
except Exception as rec_exc:
|
||||
log.debug(f"Origin reconnect failed: {rec_exc}")
|
||||
|
||||
if reconnected:
|
||||
continue
|
||||
|
||||
raise
|
||||
|
||||
if cancelled or download_id in cancelled_downloads:
|
||||
with download_lock:
|
||||
@@ -457,6 +595,14 @@ def download_file(
|
||||
error_msg = f"Unauthorized (HTTP {status_code}): Authentication required."
|
||||
elif status_code == 404:
|
||||
error_msg = "Model not found (HTTP 404): The file may have been moved or deleted."
|
||||
elif status_code == 429:
|
||||
resp = getattr(exc, "response", None)
|
||||
retry_hdr = resp.headers.get("Retry-After") if resp is not None else None
|
||||
retry_sec = parse_retry_after(retry_hdr) if retry_hdr else None
|
||||
if retry_sec:
|
||||
error_msg = f"Rate limit exceeded (HTTP 429): please retry after {retry_sec}s."
|
||||
else:
|
||||
error_msg = "Rate limit exceeded (HTTP 429): too many requests, please try again later."
|
||||
|
||||
with download_lock:
|
||||
download_progress[download_id]["status"] = "error"
|
||||
|
||||
@@ -1,19 +1,28 @@
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from collections import deque
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
from core.download.api import context as downloader
|
||||
from core.download.aria2_backend import (
|
||||
download_file_with_aria2,
|
||||
drain_aria2_stderr,
|
||||
is_aria2_important_stderr,
|
||||
is_aria2_recoverable_url_error,
|
||||
pause_download,
|
||||
recover_aria2_missing_control_file,
|
||||
resume_download,
|
||||
start_background_sha256_verification,
|
||||
)
|
||||
from core.download.orchestrator import download_file, find_active_download_for_path
|
||||
from core.download.orchestrator import (
|
||||
download_file,
|
||||
find_active_download_for_path,
|
||||
parse_retry_after,
|
||||
)
|
||||
from core.path_utils import normalize_metadata_file_path
|
||||
|
||||
|
||||
@@ -483,6 +492,240 @@ class DownloaderBoundaryTests(unittest.TestCase):
|
||||
sidecar_meta = downloader.read_json_safe(sidecar_path, {})
|
||||
self.assertFalse(sidecar_meta.get("sha256_verified"))
|
||||
|
||||
def test_is_aria2_important_stderr(self):
|
||||
self.assertTrue(is_aria2_important_stderr("Error: Write disk cache flush failure"))
|
||||
self.assertTrue(is_aria2_important_stderr("cause: No space left on device"))
|
||||
self.assertTrue(is_aria2_important_stderr("failed to write into the file: Access is denied."))
|
||||
self.assertFalse(is_aria2_important_stderr("aria2 progress: 50% - 20 MB/s"))
|
||||
self.assertFalse(is_aria2_important_stderr("Connecting to server 127.0.0.1:8000"))
|
||||
|
||||
def test_drain_aria2_stderr_reads_lines_without_blocking(self):
|
||||
mock_process = MagicMock()
|
||||
mock_process.stderr = io.BytesIO(
|
||||
b"connecting to host\n"
|
||||
b"no space left on device\n"
|
||||
b"closing connection\n"
|
||||
)
|
||||
recent_lines = deque(maxlen=10)
|
||||
drain_aria2_stderr(mock_process, recent_lines)
|
||||
self.assertEqual(
|
||||
list(recent_lines),
|
||||
[
|
||||
"connecting to host",
|
||||
"no space left on device",
|
||||
"closing connection",
|
||||
],
|
||||
)
|
||||
|
||||
def test_is_aria2_recoverable_url_error(self):
|
||||
self.assertTrue(is_aria2_recoverable_url_error("No URI available. See --max-tries option."))
|
||||
self.assertTrue(is_aria2_recoverable_url_error("HTTP response header was bad or unexpected: 403 Forbidden"))
|
||||
self.assertTrue(is_aria2_recoverable_url_error("Failed with status code 401"))
|
||||
self.assertTrue(is_aria2_recoverable_url_error("Remote server returned 410 Gone"))
|
||||
self.assertFalse(is_aria2_recoverable_url_error("Checksum mismatch error"))
|
||||
self.assertFalse(is_aria2_recoverable_url_error("File already exists and cannot be overwritten"))
|
||||
|
||||
def test_aria2_download_recovers_from_expired_signed_url(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dest_path = os.path.join(temp_dir, "model_resumed.safetensors")
|
||||
with open(dest_path, "wb") as f:
|
||||
f.write(b"partially downloaded bytes")
|
||||
|
||||
mock_resolve = MagicMock(side_effect=[
|
||||
("https://cdn.example.com/expired-url", {"User-Agent": "test-ua"}),
|
||||
("https://cdn.example.com/fresh-url", {"User-Agent": "test-ua"}),
|
||||
])
|
||||
mock_rpc = MagicMock(side_effect=[
|
||||
"gid-1", # 1st addUri
|
||||
"OK", # removeDownloadResult for gid-1
|
||||
"gid-2", # 2nd addUri
|
||||
])
|
||||
mock_status = MagicMock(side_effect=[
|
||||
{"status": "error", "errorMessage": "No URI available. See --max-tries option."},
|
||||
{"status": "complete", "completedLength": "500", "totalLength": "500"},
|
||||
])
|
||||
|
||||
with patch.object(downloader, "_ensure_aria2_daemon"), \
|
||||
patch.object(downloader, "_resolve_download_url_for_aria2", mock_resolve), \
|
||||
patch.object(downloader, "_aria2_rpc", mock_rpc), \
|
||||
patch.object(downloader, "_aria2_tell_status", mock_status), \
|
||||
patch.object(downloader, "write_model_resolver_metadata", return_value=""), \
|
||||
patch.object(downloader, "_schedule_aria2_idle_stop"), \
|
||||
patch.object(downloader, "invalidate_model_caches"):
|
||||
|
||||
result = download_file_with_aria2(
|
||||
"https://civitai.com/api/download/models/999",
|
||||
dest_path,
|
||||
"url-refresh-test",
|
||||
dependencies=downloader,
|
||||
)
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(mock_resolve.call_count, 2)
|
||||
self.assertEqual(mock_rpc.call_args_list[0].args[0], "aria2.addUri")
|
||||
self.assertEqual(mock_rpc.call_args_list[0].args[1][0], ["https://cdn.example.com/expired-url"])
|
||||
self.assertEqual(mock_rpc.call_args_list[1].args, ("aria2.removeDownloadResult", ["gid-1"]))
|
||||
self.assertEqual(mock_rpc.call_args_list[2].args[0], "aria2.addUri")
|
||||
self.assertEqual(mock_rpc.call_args_list[2].args[1][0], ["https://cdn.example.com/fresh-url"])
|
||||
|
||||
def test_aria2_cancellation_shielding_prevents_orphaned_transfer(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dest_path = os.path.join(temp_dir, "cancelled.safetensors")
|
||||
download_id = "test-shielded-cancel"
|
||||
downloader.cancelled_downloads.add(download_id)
|
||||
|
||||
mock_rpc = MagicMock()
|
||||
mock_resolve = MagicMock(return_value=("https://example.com/file.safetensors", {}))
|
||||
|
||||
with patch.object(downloader, "_ensure_aria2_daemon"), \
|
||||
patch.object(downloader, "_resolve_download_url_for_aria2", mock_resolve), \
|
||||
patch.object(downloader, "_aria2_rpc", mock_rpc), \
|
||||
patch.object(downloader, "_schedule_aria2_idle_stop"):
|
||||
|
||||
result = download_file_with_aria2(
|
||||
"https://example.com/file.safetensors",
|
||||
dest_path,
|
||||
download_id,
|
||||
dependencies=downloader,
|
||||
)
|
||||
|
||||
self.assertFalse(result["success"])
|
||||
self.assertEqual(result["error"], "Download cancelled")
|
||||
self.assertEqual(mock_rpc.call_count, 0)
|
||||
|
||||
def test_parse_retry_after(self):
|
||||
self.assertEqual(parse_retry_after("120"), 120)
|
||||
self.assertEqual(parse_retry_after("0"), 1)
|
||||
self.assertIsNone(parse_retry_after(""))
|
||||
self.assertIsNone(parse_retry_after(None))
|
||||
self.assertIsNone(parse_retry_after("invalid-text"))
|
||||
|
||||
def test_python_download_handles_rate_limit_429_with_retry_after(self):
|
||||
import requests
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dest_path = os.path.join(temp_dir, "model.safetensors")
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 429
|
||||
mock_resp.headers = {"Retry-After": "45"}
|
||||
http_err = requests.exceptions.HTTPError(response=mock_resp)
|
||||
|
||||
with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \
|
||||
patch.object(downloader, "request_public_url", side_effect=http_err):
|
||||
|
||||
result = download_file(
|
||||
"https://civitai.com/api/download/models/12345",
|
||||
dest_path,
|
||||
"rate-limit-test",
|
||||
dependencies=downloader,
|
||||
)
|
||||
|
||||
self.assertFalse(result["success"])
|
||||
self.assertIn("Rate limit exceeded (HTTP 429)", result["error"])
|
||||
self.assertIn("45s", result["error"])
|
||||
|
||||
def test_python_download_resumes_via_direct_target_retry_when_redirect_drops_range(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dest_path = os.path.join(temp_dir, "model.safetensors")
|
||||
partial_path = f"{dest_path}.part"
|
||||
initial_data = b"0123456789"
|
||||
with open(partial_path, "wb") as f:
|
||||
f.write(initial_data)
|
||||
|
||||
# Response 1: from origin redirect, drops range -> returns 200
|
||||
resp_redirected = MagicMock()
|
||||
resp_redirected.status_code = 200
|
||||
resp_redirected.headers = {"Content-Length": "20"}
|
||||
|
||||
# Response 2: from direct target URL with range -> returns 206
|
||||
remaining_data = b"ABCDEFGHIJ"
|
||||
total_data = initial_data + remaining_data
|
||||
resp_direct = MagicMock()
|
||||
resp_direct.status_code = 206
|
||||
resp_direct.headers = {
|
||||
"Content-Range": f"bytes 10-19/{len(total_data)}",
|
||||
"Content-Length": "10",
|
||||
}
|
||||
resp_direct.iter_content.return_value = [remaining_data]
|
||||
|
||||
mock_request = MagicMock(side_effect=[
|
||||
(resp_redirected, "https://cdn.example.com/direct-file", {}),
|
||||
(resp_direct, "https://cdn.example.com/direct-file", {}),
|
||||
])
|
||||
|
||||
with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \
|
||||
patch.object(downloader, "request_public_url", mock_request), \
|
||||
patch.object(downloader, "write_model_resolver_metadata", return_value=""):
|
||||
|
||||
result = download_file(
|
||||
"https://civitai.com/api/download/models/100",
|
||||
dest_path,
|
||||
"range-redirect-test",
|
||||
dependencies=downloader,
|
||||
)
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(len(total_data), result["size"])
|
||||
with open(dest_path, "rb") as f:
|
||||
self.assertEqual(total_data, f.read())
|
||||
self.assertEqual(mock_request.call_count, 2)
|
||||
# Second call was direct to the CDN target
|
||||
self.assertEqual(mock_request.call_args_list[1].args[1], "https://cdn.example.com/direct-file")
|
||||
|
||||
def test_python_download_auto_reconnects_on_stream_interruption(self):
|
||||
import requests
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dest_path = os.path.join(temp_dir, "model_reconnect.safetensors")
|
||||
chunk1 = b"AAAAA"
|
||||
chunk2 = b"BBBBB"
|
||||
total_data = chunk1 + chunk2
|
||||
|
||||
# Stream 1: yields chunk1 then drops connection
|
||||
resp1 = MagicMock()
|
||||
resp1.status_code = 200
|
||||
resp1.headers = {"Content-Length": str(len(total_data))}
|
||||
|
||||
def failing_stream(*args, **kwargs):
|
||||
yield chunk1
|
||||
raise requests.exceptions.ConnectionError("Connection reset by peer")
|
||||
|
||||
resp1.iter_content.side_effect = failing_stream
|
||||
|
||||
# Stream 2: resumes with Range 206, yields chunk2
|
||||
resp2 = MagicMock()
|
||||
resp2.status_code = 206
|
||||
resp2.headers = {
|
||||
"Content-Range": f"bytes 5-9/{len(total_data)}",
|
||||
"Content-Length": "5",
|
||||
}
|
||||
resp2.iter_content.return_value = [chunk2]
|
||||
|
||||
mock_request = MagicMock(side_effect=[
|
||||
(resp1, "https://example.com/model.safetensors", {}),
|
||||
(resp2, "https://example.com/model.safetensors", {}),
|
||||
])
|
||||
|
||||
with patch.object(downloader, "_download_backend_from_settings", return_value="python"), \
|
||||
patch.object(downloader, "request_public_url", mock_request), \
|
||||
patch.object(downloader, "write_model_resolver_metadata", return_value=""):
|
||||
|
||||
result = download_file(
|
||||
"https://example.com/model.safetensors",
|
||||
dest_path,
|
||||
"stream-reconnect-test",
|
||||
dependencies=downloader,
|
||||
)
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(len(total_data), result["size"])
|
||||
with open(dest_path, "rb") as f:
|
||||
self.assertEqual(total_data, f.read())
|
||||
self.assertEqual(mock_request.call_count, 2)
|
||||
# Reconnected call should have requested Range from byte 5
|
||||
second_call_headers = mock_request.call_args_list[1].kwargs.get("headers", {})
|
||||
self.assertEqual(second_call_headers.get("Range"), "bytes=5-")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user