Initial Commit 🚀 🚀

api/utils/worker.py

@@ -0,0 +1,63 @@
+"""Utilities for worker process identification."""
+
+import multiprocessing
+import os
+
+from loguru import logger
+
+
+def get_worker_id() -> int:
+    """Get the current worker ID from environment or process name.
+
+    Returns:
+        Worker ID (0-based index), or 0 if not in a worker process.
+    """
+    # Check for custom ASGI_WORKER_ID (for future compatibility)
+    worker_id = os.getenv("ASGI_WORKER_ID")
+    if worker_id:
+        return int(worker_id)
+
+    # Debug log the process name to understand worker identification
+    process_name = multiprocessing.current_process().name
+
+    # Try to extract worker number from process name
+    # Uvicorn with --workers creates processes like "SpawnProcess-1", "SpawnProcess-2", etc.
+    # TODO FIXME: If a worker process crashes and uvicorn creates a new process,
+    # it assigns ID which may be beyond NUM_FASTAPI_WORKERS. Example: if we have
+    # 2 fastapi workers configured, and one of them dies, we can get a process name with
+    # SpawnProcess-3 which is bad
+    if "SpawnProcess" in process_name:
+        try:
+            # Extract the number after "SpawnProcess-"
+            worker_num = int(process_name.split("-")[-1])
+            # Convert to 0-based index
+            return worker_num - 1
+        except (ValueError, IndexError):
+            logger.warning(
+                f"Could not extract worker ID from process name: {process_name}"
+            )
+
+    # Gunicorn creates workers with names like "Worker-1", "Worker-2", etc.
+    if "Worker" in process_name:
+        try:
+            # Extract the number after "Worker-"
+            worker_num = int(process_name.split("-")[-1])
+            # Convert to 0-based index
+            return worker_num - 1
+        except (ValueError, IndexError):
+            logger.warning(
+                f"Could not extract worker ID from process name: {process_name}"
+            )
+
+    # Not in a worker process (main process or single-process mode)
+    return 0
+
+
+def is_worker_process() -> bool:
+    """Check if we're running in a worker process (not the main process).
+
+    Returns:
+        True if in a worker process, False if in main process or single-process mode.
+    """
+    process_name = multiprocessing.current_process().name
+    return "SpawnProcess" in process_name or "Worker" in process_name