mirror of
https://github.com/dograh-hq/dograh.git
synced 2026-06-07 07:55:16 +02:00
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
|
"""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
|