nodes/huchukato-ComfyUI-HuggingFace
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix canceling download causing complete lock

Browse Source
This commit is contained in:
MoonGoblin
2025-04-04 21:42:58 +07:00
parent 0089081ee2
commit 636c3bcfa8
3 changed files with 103 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
downloader/chunk_downloader.py
+20 -8
View File
@@ -81,10 +81,12 @@ class ChunkDownloader:
def cancel(self):
"""Signal the download to cancel."""
if not self.is_cancelled:
print(f"[Downloader {self.download_id or 'N/A'}] Cancellation requested.")
print(f"[Downloader {self.download_id or 'N/A'}] Cancellation requested by user.")
self.cancel_event.set()
print("masok 1")
self.error = "Download cancelled by user"
if self.manager and self.download_id:
print("masok 2")
self.manager._update_download_status(self.download_id, status="cancelled", error=self.error)
def _cleanup_temp(self, success: bool):
@@ -178,7 +180,8 @@ class ChunkDownloader:
self.manager._update_download_status(
self.download_id,
progress=progress,
speed=self._speed
speed=self._speed,
status="downloading"
)
def download_segment(self, segment_index: int, start_byte: int, end_byte: int):
@@ -189,7 +192,10 @@ class ChunkDownloader:
for current_try in range(retries):
if self.is_cancelled:
return
print(f"[Downloader {self.download_id}] Segment {segment_index} cancelled before request (Try {current_try+1}).")
# Ensure error is set if not already
if not self.error: self.error = "Cancelled during segment download"
return
response = None
try:
@@ -200,8 +206,12 @@ class ChunkDownloader:
with open(part_file_path, 'wb') as f:
for chunk in response.iter_content(self.chunk_size):
if self.is_cancelled:
return
print(f"[Downloader {self.download_id}] Segment {segment_index} cancelled mid-stream.")
# Ensure error is set if not already
if not self.error: self.error = "Cancelled during segment download"
# No need to return immediately, finally block will close response
raise OperationCancelled("Download cancelled by user signal.") # Raise specific exception
if chunk:
bytes_written = f.write(chunk)
bytes_written_this_segment += bytes_written
@@ -216,6 +226,7 @@ class ChunkDownloader:
return # Success
except (requests.exceptions.RequestException, ValueError) as e:
# Handle HTTP status codes
status_code = None
@@ -312,7 +323,7 @@ class ChunkDownloader:
"""Fallback to standard single-connection download."""
self.connection_type = "Single"
if self.manager and self.download_id:
self.manager._update_download_status(self.download_id, connection_type=self.connection_type)
self.manager._update_download_status(self.download_id, connection_type=self.connection_type, status="downloading")
print(f"[Downloader {self.download_id}] Using standard single-connection download for {self.output_path.name}...")
@@ -352,11 +363,12 @@ class ChunkDownloader:
if self.is_cancelled:
print(f"[Downloader {self.download_id}] Fallback download cancelled.")
return False
if chunk:
bytes_written = f.write(chunk)
self._update_progress(bytes_written)
# Verify download size if known
if self.total_size > 0 and self.downloaded != self.total_size and not self.error:
print(f"[Downloader {self.download_id}] Warning: Size mismatch. Expected {self.total_size}, got {self.downloaded}.")
@@ -474,7 +486,7 @@ class ChunkDownloader:
"""Handle multi-connection download process."""
self.connection_type = f"Multi ({self.num_connections})"
if self.manager and self.download_id:
self.manager._update_download_status(self.download_id, connection_type=self.connection_type)
self.manager._update_download_status(self.download_id, connection_type=self.connection_type, status="downloading")
print(f"[Downloader {self.download_id}] Starting multi-connection download for {self.output_path.name} "
f"({self.total_size / (1024 * 1024):.2f} MB) using {self.num_connections} connections.")
downloader/manager.py
+75 -46
View File
@@ -77,46 +77,83 @@ class DownloadManager:
return download_id
def cancel_download(self, download_id: str) -> bool:
"""Requests cancellation of a queued or active download."""
with self.lock:
# Check queue first
for i, item in enumerate(self.queue):
if item["id"] == download_id:
cancelled_info = self.queue.pop(i)
cancelled_info["status"] = "cancelled"
cancelled_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
cancelled_info["error"] = "Cancelled from queue"
self._add_to_history(cancelled_info)
print(f"[Manager] Cancelled queued download: {download_id}")
return True
"""Requests cancellation of a queued or active download."""
print(f"[Manager] Received cancellation request for: {download_id}") # Moved print earlier
downloader_to_cancel: Optional['ChunkDownloader'] = None
found_in_active = False
# Check active downloads
if download_id in self.active_downloads:
active_info = self.active_downloads[download_id]
downloader: Optional['ChunkDownloader'] = active_info.get("downloader_instance")
with self.lock:
# 1. Check queue first (can be fully handled under lock)
for i, item in enumerate(self.queue):
if item["id"] == download_id:
cancelled_info = self.queue.pop(i)
cancelled_info["status"] = "cancelled"
cancelled_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
cancelled_info["error"] = "Cancelled from queue"
self._add_to_history(cancelled_info) # _add_to_history is safe under lock
print(f"[Manager] Cancelled queued download: {download_id}")
return True # Cancelled from queue, we are done
if downloader:
# Check if already cancelled (to avoid duplicate logs/actions)
if not downloader.is_cancelled:
downloader.cancel() # Signal the downloader thread to stop
# Status ("cancelled") will be updated by the _download_file_wrapper or downloader.cancel()
print(f"[Manager] Cancellation requested for active download: {download_id}")
# Don't move to history yet, let the thread finish and update status
# 2. Check active downloads - Find the instance *under lock*
if download_id in self.active_downloads:
found_in_active = True
active_info = self.active_downloads[download_id]
downloader_to_cancel = active_info.get("downloader_instance")
current_status = active_info.get("status")
# If downloader instance doesn't exist yet (status 'starting')
# or if already terminal, handle it here under lock
if not downloader_to_cancel and current_status == "starting":
print(f"[Manager] Marking 'starting' download as cancelled: {download_id}")
# Mark as cancelled, it won't start or will be caught by wrapper
active_info["status"] = "cancelled"
if not active_info.get("end_time"):
active_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
active_info["error"] = "Cancelled before download thread fully started"
# Don't assign downloader_to_cancel, let process_queue clean it up
downloader_to_cancel = None # Explicitly clear
# Indicate we handled it (or attempted to)
return True # Exit, processed within lock
elif current_status in ["completed", "failed", "cancelled"]:
print(f"[Manager] Download {download_id} is already in terminal state '{current_status}'. Cannot cancel.")
# Indicate we don't need to proceed further outside the lock
downloader_to_cancel = None # Explicitly clear
return False # Already finished/cancelled
# If we found an instance and it's potentially running,
# store it to call cancel *after* releasing the lock.
print(f"[Manager] Found active downloader instance for {download_id}. Will signal outside lock.")
# Lock is released automatically here when exiting 'with'
# 3. Signal the downloader instance *outside* the lock
if downloader_to_cancel:
try:
# Check if already cancelled (to avoid duplicate logs/actions) - This check is thread-safe
if not downloader_to_cancel.is_cancelled:
print(f"[Manager] Calling downloader.cancel() for {download_id}")
downloader_to_cancel.cancel() # This can now call _update_download_status safely
print(f"[Manager] Signalled downloader.cancel() for {download_id}")
# Let the download thread's final status update handle moving to history
return True # Signal sent
else:
print(f"[Manager] Active download {download_id} was already cancelling.")
return True
else:
# Downloader instance not yet created (e.g., status 'starting')
# Mark as cancelled, it won't start
active_info["status"] = "cancelled"
active_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
active_info["error"] = "Cancelled before download started"
# The _process_queue will move it later
print(f"[Manager] Cancelled download (status: starting): {download_id}")
return True
return True # Already in cancelling state is considered a success here
except Exception as e:
print(f"[Manager] Error calling downloader.cancel() for {download_id}: {e}")
# Update status to failed maybe? Or just log.
# Use _update_download_status directly here as we are outside the lock
self._update_download_status(download_id, status="failed", error=f"Error during cancel signaling: {e}")
return False # Failed to signal
print(f"[Manager] Could not cancel - ID not found in queue or active: {download_id}")
return False
# 4. Handle cases where it wasn't in queue and wasn't running/starting
if not found_in_active:
print(f"[Manager] Could not cancel - ID not found in queue or active: {download_id}")
return False # Not found
# It was found in active but was already terminal or couldn't be signalled
# Return value determined above
return False # Should have returned True earlier if successful
def get_status(self) -> Dict[str, List[Dict[str, Any]]]:
"""Returns the current state of the queue, active downloads, and history."""
@@ -237,7 +274,8 @@ class DownloadManager:
if download_id in self.active_downloads:
item = self.active_downloads[download_id]
updated = False # Track if any field was actually updated
if status == "cancelled":
print("masok 5 cancel lur")
# Only update if value is provided
if status is not None and item.get("status") != status:
item["status"] = status
@@ -266,17 +304,8 @@ class DownloadManager:
updated = True
# Optional: Log when an update actually happens
# if updated: print(f"DEBUG - Updated status for {download_id}: Status={status}, Progress={progress}, Speed={speed}, Conn={connection_type}, Error={error}")
print(f"DEBUG - Updated status for {download_id}: Status={status}, Progress={progress}, Speed={speed}, Conn={connection_type}, Error={error}")
# else:
# # This can happen normally if status update arrives after item moved to history
# # It can also happen if cancel is processed very quickly
# # If status indicates a terminal state, log a warning as it might indicate a race condition
# if status in ["completed", "failed", "cancelled"]:
# print(f"[Manager] Warning: Attempted to update status for finished/unknown ID: {download_id}. Status was: {status}. This might indicate a race condition if unexpected.")
# # else: # Non-terminal updates for unknown IDs are less concerning
# # print(f"Debug: Attempted to update progress/speed for finished/unknown ID: {download_id}.")
# pass
def _save_civitai_metadata(self, download_info: Dict[str, Any]):
"""Saves the .cminfo.json file."""
server/routes.py
+8 -2
View File
@@ -455,14 +455,20 @@ async def route_get_status(request):
async def route_cancel_download(request):
"""API Endpoint to cancel a download."""
try:
print("masok 6")
print(request)
data = await _get_request_json(request)
print("masok 7 ")
print(data)
download_id = data.get("download_id")
print("masok 8 " + download_id)
if not download_id:
print("not download id " + download_id)
raise web.HTTPBadRequest(reason="Missing 'download_id'")
success = download_manager.cancel_download(download_id)
print("masok 9")
print(success)
if success:
return web.json_response({
"status": "cancelled", # Or "cancellation_requested" ?