From 414679c676be554f23db9b283cf6f2fbb995a053 Mon Sep 17 00:00:00 2001 From: Sebastian Monroy Date: Thu, 25 Sep 2025 11:30:25 +0100 Subject: [PATCH] simplified the logging when worker_consumer.py fails to initialize SQS, improved logging everywhere --- media_stream.py | 68 ++++++++++++++++++---------- nilornodes.py | 23 +++++++--- worker_consumer.py | 109 ++++++++++++++++++++++++++++++--------------- 3 files changed, 136 insertions(+), 64 deletions(-) diff --git a/media_stream.py b/media_stream.py index dc22279..e994a5f 100644 --- a/media_stream.py +++ b/media_stream.py @@ -69,7 +69,7 @@ class MediaStreamInput: input_name: str = "default_input", ): logging.info( - f"MediaStreamInput: Downloading from {presigned_download_url} for input '{input_name}' with format '{format}'" + f"ℹī¸ Nilor-Nodes: MediaStreamInput: Downloading from {presigned_download_url} for input '{input_name}' with format '{format}'" ) try: # Two-phase download for batches: manifest first, then assets @@ -79,7 +79,7 @@ class MediaStreamInput: manifest = manifest_response.json() logging.info( - f"Processing manifest for '{manifest.get('input_name')}' with {len(manifest.get('files', []))} assets." + f"ℹī¸ Nilor-Nodes: Processing manifest for '{manifest.get('input_name')}' with {len(manifest.get('files', []))} assets." ) # Sort files by sequence number to ensure correct order @@ -96,7 +96,7 @@ class MediaStreamInput: asset_responses.append(resp.content) except requests.RequestException as e: logging.error( - f"Failed to download asset {file_info.get('filename')}: {e}" + f"🛑 Nilor-Nodes: Failed to download asset {file_info.get('filename')}: {e}" ) raise # Re-raise to fail the entire process @@ -118,14 +118,20 @@ class MediaStreamInput: ) except requests.RequestException as e: - logging.error(f"MediaStreamInput: Failed to download file: {e}") + logging.error( + f"🛑 Nilor-Nodes (MediaStreamInput): Failed to download file: {e}" + ) return (None,) except Exception as e: - logging.error(f"MediaStreamInput: Failed to process media: {e}") + logging.error( + f"🛑 Nilor-Nodes (MediaStreamInput): Failed to process media: {e}" + ) return (None,) def _process_image_batch(self, image_bytes_list): - logging.info(f"Processing image batch with {len(image_bytes_list)} images...") + logging.info( + f"ℹī¸ Nilor-Nodes (MediaStreamInput): Processing image batch with {len(image_bytes_list)} images..." + ) output_images = [] for image_bytes in image_bytes_list: @@ -142,12 +148,12 @@ class MediaStreamInput: images_tensor = torch.cat(output_images, dim=0) logging.info( - f"Image batch processing successful. Batch shape: {images_tensor.shape}" + f"✅ Nilor-Nodes (MediaStreamInput): Image batch processing successful. Batch shape: {images_tensor.shape}" ) return (images_tensor,) def _process_image(self, image_bytes): - logging.info("Processing as image...") + logging.info("ℹī¸ Nilor-Nodes (MediaStreamInput): Processing as image...") image_pil = Image.open(io.BytesIO(image_bytes)) # Ensure image is in RGB @@ -156,11 +162,11 @@ class MediaStreamInput: np.array(rgb_image_pil).astype(np.float32) / 255.0 ).unsqueeze(0) - logging.info("Image processing successful.") + logging.info("✅ Nilor-Nodes (MediaStreamInput): Image processing successful.") return (image_tensor,) def _process_video(self, video_bytes): - logging.info("Processing as video...") + logging.info("ℹī¸ Nilor-Nodes (MediaStreamInput): Processing as video...") frames = [] with imageio.get_reader(io.BytesIO(video_bytes), format="mp4") as reader: for frame in reader: @@ -171,12 +177,16 @@ class MediaStreamInput: frames.append(tensor_frame) if not frames: - raise ValueError("No frames could be read from the video.") + raise ValueError( + "🛑 Nilor-Nodes (MediaStreamInput): No frames could be read from the video." + ) # Stack frames into a single tensor (batch of images) video_tensor = torch.stack(frames) - logging.info(f"Video processing successful. Image Shape: {video_tensor.shape}") + logging.info( + f"✅ Nilor-Nodes (MediaStreamInput): Video processing successful. Image Shape: {video_tensor.shape}" + ) return (video_tensor,) @@ -245,7 +255,7 @@ class MediaStreamOutput: final_outputs_dict = json.loads(output_object_keys.replace("'", '"')) except Exception as e: logging.error( - f"FATAL: Could not parse output_object_keys from string: {output_object_keys}. Error: {e}" + f"🛑 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not parse output_object_keys from string: {output_object_keys}. Error: {e}" ) final_outputs_dict = {} # Send empty dict on failure. @@ -260,7 +270,7 @@ class MediaStreamOutput: output_key_for_this_node = final_outputs_dict.get(output_name) if not output_key_for_this_node: logging.error( - f"FATAL: Could not find object key for output name '{output_name}' in output_object_keys." + f"🛑 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not find object key for output name '{output_name}' in output_object_keys." ) # Send an empty dictionary to signal failure. final_outputs_for_sqs = {} @@ -284,21 +294,25 @@ class MediaStreamOutput: region_name=os.getenv("AWS_DEFAULT_REGION", "us-east-1"), ) logging.info( - f"Sending completion message for job {job_id} to queue: {job_completions_queue_url}" + f"ℹī¸ Nilor-Nodes (MediaStreamOutput): Sending completion message for job {job_id} to queue: {job_completions_queue_url}" ) sqs_client.send_message( QueueUrl=job_completions_queue_url, MessageBody=json.dumps(completion_message), ) - logging.info("Completion message sent successfully.") + logging.info( + "✅ Nilor-Nodes (MediaStreamOutput): Completion message sent successfully." + ) except Exception as e: - logging.error(f"Failed to send completion message to SQS: {e}") + logging.error( + f"🛑 Nilor-Nodes (MediaStreamOutput): Failed to send completion message to SQS: {e}" + ) raise # Re-raise to fail the ComfyUI job return {"ui": {"images": []}} def _upload_image(self, image_tensor, url): - logging.info("Uploading as PNG image...") + logging.info("ℹī¸ Nilor-Nodes (MediaStreamOutput): Uploading as PNG image...") i = 255.0 * image_tensor.cpu().numpy() img_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) @@ -309,7 +323,9 @@ class MediaStreamOutput: self._perform_upload(buffer, url, "image/png") def _upload_video(self, image_batch_tensor, url, framerate): - logging.info(f"Uploading as MP4 video. Frame count: {len(image_batch_tensor)}") + logging.info( + f"ℹī¸ Nilor-Nodes (MediaStreamOutput): Uploading as MP4 video. Frame count: {len(image_batch_tensor)}" + ) frames = [] for image_tensor in image_batch_tensor: i = 255.0 * image_tensor.cpu().numpy() @@ -324,18 +340,24 @@ class MediaStreamOutput: def _perform_upload(self, buffer, url, content_type): try: - logging.info(f"Uploading to {url} with Content-Type: {content_type}") + logging.info( + f"ℹī¸ Nilor-Nodes (MediaStreamOutput): Uploading to {url} with Content-Type: {content_type}" + ) headers = {"Content-Type": content_type} response = requests.put( url, data=buffer.read(), headers=headers, timeout=300 ) response.raise_for_status() - logging.info("Upload successful.") + logging.info("✅ Nilor-Nodes (MediaStreamOutput): Upload successful.") except requests.RequestException as e: - logging.error(f"MediaStreamOutput: Failed to upload media: {e}") + logging.error( + f"🛑 Nilor-Nodes (MediaStreamOutput): Failed to upload media: {e}" + ) raise except Exception as e: - logging.error(f"MediaStreamOutput: Failed to process and upload media: {e}") + logging.error( + f"🛑 Nilor-Nodes (MediaStreamOutput): Failed to process and upload media: {e}" + ) raise diff --git a/nilornodes.py b/nilornodes.py index d051c9c..da2f6ac 100644 --- a/nilornodes.py +++ b/nilornodes.py @@ -18,6 +18,7 @@ from pathlib import Path import cv2 import warnings from .utils import pil2tensor, tensor2pil +import logging BIGMIN = -(2**53 - 1) BIGMAX = 2**53 - 1 @@ -404,7 +405,7 @@ class NilorSaveEXRArbitrary: self, channels=None, filename_prefix="output", prompt=None, extra_pnginfo=None ): - print("Running save_exr_arbitrary") + logging.info("ℹī¸ Nilor-Nodes (SaveEXRArbitrary): Running save_exr_arbitrary") # print(f"channels: {channels}") # print(f"filename_prefix: {filename_prefix}") @@ -416,7 +417,9 @@ class NilorSaveEXRArbitrary: try: actual_channels[0] except TypeError: - print("actual_channels is not subscriptable") + logging.error( + "🛑 Nilor-Nodes (SaveEXRArbitrary): actual_channels is not subscriptable" + ) return # File path handling @@ -503,9 +506,13 @@ class NilorSaveEXRArbitrary: exr_file.writePixels(channel_data) exr_file.close() - print(f"EXR file saved successfully to {writepath}") + logging.info( + f"✅ Nilor-Nodes (SaveEXRArbitrary): EXR file saved successfully to {writepath}" + ) except Exception as e: - print(f"Failed to write EXR file: {e}") + logging.error( + f"🛑 Nilor-Nodes (SaveEXRArbitrary): Failed to write EXR file: {e}" + ) class NilorSaveVideoToHFDataset: @@ -782,12 +789,16 @@ class NilorOutputFilenameString: if unique_id is not None and extra_pnginfo is not None: if not isinstance(extra_pnginfo, list): - print("Error: extra_pnginfo is not a list") + logging.error( + "🛑 Nilor-Nodes (OutputFilenameString): extra_pnginfo is not a list" + ) elif ( not isinstance(extra_pnginfo[0], dict) or "workflow" not in extra_pnginfo[0] ): - print("Error: extra_pnginfo[0] is not a dict or missing 'workflow' key") + logging.error( + "🛑 Nilor-Nodes (OutputFilenameString): extra_pnginfo[0] is not a dict or missing 'workflow' key" + ) else: workflow = extra_pnginfo[0]["workflow"] node = next( diff --git a/worker_consumer.py b/worker_consumer.py index 18dbe73..4e9a495 100644 --- a/worker_consumer.py +++ b/worker_consumer.py @@ -16,6 +16,7 @@ import aiohttp import websockets from aiobotocore.session import get_session from dotenv import load_dotenv +from botocore.exceptions import EndpointConnectionError # --- Load Environment Variables --- # Load from the .env file in the same directory @@ -23,9 +24,11 @@ current_dir = os.path.dirname(os.path.abspath(__file__)) dotenv_path = os.path.join(current_dir, ".env") if os.path.exists(dotenv_path): load_dotenv(dotenv_path=dotenv_path) - logging.info(f"Loaded environment variables from {dotenv_path}") + logging.info(f"ℹī¸ Nilor-Nodes: Loaded environment variables from {dotenv_path}") else: - logging.info("No .env file found, relying on shell environment variables.") + logging.info( + "ℹī¸ Nilor-Nodes: No .env file found, relying on shell environment variables." + ) # --- Configuration --- SQS_ENDPOINT_URL = os.getenv("SQS_ENDPOINT_URL", "http://localhost:9324") @@ -59,7 +62,7 @@ class WorkerConsumer: self.http_session = None async def _initialize_sqs(self): - """Initializes SQS queue URLs.""" + """Initializes SQS queue URLs. Returns True on success, False on failure.""" async with self.session.create_client( "sqs", region_name=AWS_DEFAULT_REGION, @@ -74,9 +77,17 @@ class WorkerConsumer: self.status_updates_queue_url = await self._get_queue_url( client, SQS_JOB_STATUS_UPDATES_QUEUE_NAME ) + return True + except EndpointConnectionError as e: + # Quiet the noisy traceback by logging a concise warning instead + logging.warning( + f"⚠ī¸ Nilor-Nodes: SQS endpoint is unreachable at {SQS_ENDPOINT_URL}: {e}. " + "Disabling SQS worker consumer functionality." + ) + return False except Exception as e: - logging.error(f"Failed to initialize SQS queues: {e}") - raise + logging.error(f"⚠ī¸ Nilor-Nodes: Failed to initialize SQS queues: {e}") + return False async def _get_queue_url(self, client, queue_name): """Retrieves the SQS queue URL.""" @@ -84,14 +95,16 @@ class WorkerConsumer: response = await client.get_queue_url(QueueName=queue_name) return response["QueueUrl"] except client.exceptions.QueueDoesNotExist: - logging.error(f"SQS queue '{queue_name}' does not exist.") + logging.error(f"⚠ī¸ Nilor-Nodes: SQS queue '{queue_name}' does not exist.") raise async def listen_for_comfy_events(self): while True: try: async with websockets.connect(COMFYUI_WS_URL) as websocket: - logging.info(f"Connected to ComfyUI websocket at {COMFYUI_WS_URL}") + logging.info( + f"ℹī¸ Nilor-Nodes: Connected to ComfyUI websocket at {COMFYUI_WS_URL}" + ) while True: message = await websocket.recv() if isinstance(message, str): @@ -116,7 +129,7 @@ class WorkerConsumer: ): job_id = self.prompt_id_to_job_id_map[prompt_id] logging.info( - f"Execution started for prompt_id {prompt_id} (job_id: {job_id}) via '{event_type}' event. Sending 'running' status." + f"ℹī¸ Nilor-Nodes: Execution started for prompt_id {prompt_id} (job_id: {job_id}) via '{event_type}' event. Sending 'running' status." ) await self._send_status_update(job_id, "running") self.sent_running_status_prompts.add(prompt_id) @@ -124,7 +137,7 @@ class WorkerConsumer: # Handle execution errors elif event_type == "execution_error": logging.error( - f"Received execution error for prompt_id {prompt_id}: {data}" + f"🛑 Nilor-Nodes: Received execution error for prompt_id {prompt_id}: {data}" ) if prompt_id in self.prompt_id_to_job_id_map: self.prompt_id_to_job_id_map.pop(prompt_id) @@ -133,7 +146,7 @@ class WorkerConsumer: # Log successful execution elif event_type == "executed": logging.info( - f"Prompt {prompt_id} executed successfully according to websocket event. Final node is responsible for sending completion message." + f"✅ Nilor-Nodes: Prompt {prompt_id} executed successfully according to websocket event. Final node is responsible for sending completion message." ) if prompt_id in self.prompt_id_to_job_id_map: self.prompt_id_to_job_id_map.pop(prompt_id) @@ -141,43 +154,56 @@ class WorkerConsumer: elif event_type not in ["progress", "progress_state"]: logging.info( - f"Received ComfyUI websocket event of type '{event_type}': {data}" + f"ℹī¸ Nilor-Nodes: Received ComfyUI websocket event of type '{event_type}': {data}" ) except json.JSONDecodeError: logging.debug( - "Received non-JSON text message from websocket, ignoring." + "⚠ī¸ Nilor-Nodes: Received non-JSON text message from websocket, ignoring." ) else: logging.debug( - "Received binary message from websocket, ignoring." + "⚠ī¸ Nilor-Nodes: Received binary message from websocket, ignoring." ) except ( websockets.exceptions.ConnectionClosedError, ConnectionRefusedError, ) as e: logging.warning( - f"ComfyUI websocket connection failed: {e}. Retrying in 5 seconds..." + f"🛑 Nilor-Nodes: ComfyUI websocket connection failed: {e}. Retrying in 5 seconds..." ) await asyncio.sleep(5) except Exception as e: logging.error( - f"An unexpected error occurred in the websocket listener: {e}", + f"🛑 Nilor-Nodes: An unexpected error occurred in the websocket listener: {e}", exc_info=True, ) await asyncio.sleep(10) async def consume_loop(self): - """The main loop to continuously poll for and process messages.""" - await self._initialize_sqs() - logging.info(f"Starting worker consumer. Polling queue: {self.jobs_queue_url}") + """The main loop to continuously poll for and process messages. + Keeps retrying SQS initialization and polling if the endpoint is down. + """ - # Start the websocket listener in the background + # Start the websocket listener in the background immediately listener_task = asyncio.create_task(self.listen_for_comfy_events()) try: while True: - logging.debug("Polling for messages...") + # Ensure SQS is initialized; if not, keep attempting to initialize + if self.jobs_queue_url is None or self.status_updates_queue_url is None: + initialized = await self._initialize_sqs() + if not initialized: + logging.warning( + "⚠ī¸ Nilor-Nodes: SQS initialization failed. Retrying in 10 seconds..." + ) + await asyncio.sleep(10) + continue + logging.info( + f"ℹī¸ Nilor-Nodes: Starting worker consumer. Polling queue: {self.jobs_queue_url}" + ) + + logging.debug("ℹī¸ Nilor-Nodes: Polling for messages...") try: async with self.session.create_client( "sqs", @@ -194,7 +220,7 @@ class WorkerConsumer: messages = response.get("Messages", []) if not messages: - logging.debug("No messages received.") + logging.debug("ℹī¸ Nilor-Nodes: No messages received.") continue for message in messages: @@ -214,30 +240,40 @@ class WorkerConsumer: ReceiptHandle=message["ReceiptHandle"], ) logging.info( - f"Deleted message {message['MessageId']} from queue." + f"ℹī¸ Nilor-Nodes: Deleted message {message['MessageId']} from queue." ) except json.JSONDecodeError: # This is a poison pill message, log it but don't retry. # It will be moved to the DLQ after enough failed receives. logging.error( - f"Message {message['MessageId']} is a poison pill (JSON decode failed) and will be ignored." + f"🛑 Nilor-Nodes: Message {message['MessageId']} is a poison pill (JSON decode failed) and will be ignored." ) except Exception as e: logging.error( - f"Processing failed for message {message['MessageId']}: {e}. It will be returned to the queue for retry." + f"🛑 Nilor-Nodes: Processing failed for message {message['MessageId']}: {e}. It will be returned to the queue for retry." ) + except EndpointConnectionError as e: + # Lost connection to SQS; reset and re-initialize on next loop + logging.warning( + f"⚠ī¸ Nilor-Nodes: Lost connection to SQS at {SQS_ENDPOINT_URL}: {e}. Will retry initialization in 10 seconds." + ) + self.jobs_queue_url = None + self.status_updates_queue_url = None + await asyncio.sleep(10) except Exception as e: - logging.error(f"An error occurred in the consume loop: {e}") + logging.error( + f"🛑 Nilor-Nodes: An error occurred in the consume loop: {e}" + ) await asyncio.sleep(10) # Wait before retrying finally: listener_task.cancel() await asyncio.gather(listener_task, return_exceptions=True) - logging.info("Websocket listener stopped.") + logging.info("⚠ī¸ Nilor-Nodes: Websocket listener stopped.") async def process_message(self, message): """Processes a single SQS message.""" - logging.info(f"Processing message: {message['MessageId']}") + logging.info(f"ℹī¸ Nilor-Nodes: Processing message: {message['MessageId']}") try: body = json.loads(message["Body"]) @@ -254,7 +290,7 @@ class WorkerConsumer: # Validate that the payload has the required keys before submitting. if not job_id or "prompt" not in job_payload: logging.error( - f"Invalid message format: missing 'client_id' or 'prompt'. Payload: {job_payload}" + f"🛑 Nilor-Nodes: Invalid message format: missing 'client_id' or 'prompt'. Payload: {job_payload}" ) return @@ -263,7 +299,7 @@ class WorkerConsumer: except Exception as e: logging.error( - f"An unexpected error occurred while processing message: {e}. It will be retried." + f"🛑 Nilor-Nodes: An unexpected error occurred while processing message: {e}. It will be retried." ) # Re-raise to prevent deletion from queue if we want SQS to handle retry raise @@ -279,22 +315,22 @@ class WorkerConsumer: response_json = await response.json() prompt_id = response_json.get("prompt_id") logging.info( - f"Successfully submitted job to ComfyUI. Prompt ID: {prompt_id}" + f"✅ Nilor-Nodes: Successfully submitted job to ComfyUI. Prompt ID: {prompt_id}" ) self.prompt_id_to_job_id_map[prompt_id] = job_id # No need to delete here, the consume_loop handles message deletion except aiohttp.ClientError as e: logging.error( - f"Failed to submit job to ComfyUI: {e}. Message will be retried." + f"🛑 Nilor-Nodes: Failed to submit job to ComfyUI: {e}. Message will be retried." ) except (json.JSONDecodeError, KeyError) as e: logging.error( - f"Failed to parse ComfyUI response: {e}. Discarding malformed response." + f"🛑 Nilor-Nodes: Failed to parse ComfyUI response: {e}. Discarding malformed response." ) except Exception as e: logging.error( - f"An unexpected error occurred while submitting job to ComfyUI: {e}", + f"🛑 Nilor-Nodes: An unexpected error occurred while submitting job to ComfyUI: {e}", exc_info=True, ) @@ -311,10 +347,13 @@ class WorkerConsumer: await client.send_message( QueueUrl=self.status_updates_queue_url, MessageBody=message_body ) - logging.info(f"Sent status update for job {job_id}: {status}") + logging.info( + f"✅ Nilor-Nodes: Sent status update for job {job_id}: {status}" + ) except Exception as e: logging.error( - f"Failed to send status update for job {job_id}: {e}", exc_info=True + f"🛑 Nilor-Nodes: Failed to send status update for job {job_id}: {e}", + exc_info=True, )