diff --git a/orchestrators/imajin-pipeline/src/image_pipeline/stages/generate.py b/orchestrators/imajin-pipeline/src/image_pipeline/stages/generate.py index 656affb3..afec0b34 100644 --- a/orchestrators/imajin-pipeline/src/image_pipeline/stages/generate.py +++ b/orchestrators/imajin-pipeline/src/image_pipeline/stages/generate.py @@ -343,6 +343,8 @@ async def init_gpu_boss(): Must be called at service startup before loading any models. Initializes both the GPUBoss for lease coordination and ManagedModelLoader for coordinated model loading. + + Automatically detects GPUs and registers them with the coordinator. """ global _boss, _diffusers_loader @@ -351,12 +353,30 @@ async def init_gpu_boss(): return from model_boss import GPUBoss, Priority + from model_boss.gpu.utils import get_gpu_info_safe from model_boss_loaders import ManagedModelLoader, DiffusersLoader # noqa: F401 - import for side-effect (registry registration) _boss = GPUBoss() await _boss.connect() logger.info("GPU coordination initialized (model-boss v3)") + # Detect and initialize GPUs + gpus = await get_gpu_info_safe(timeout_s=10.0) + if not gpus: + raise RuntimeError( + "No GPUs detected. GPU coordination requires at least one CUDA-capable GPU." + ) + + for gpu in gpus: + await _boss.initialize_gpu( + gpu_index=gpu["index"], + vram_total_mb=gpu["total_memory_mb"], + gpu_name=gpu["name"], + ) + logger.info( + f"Registered GPU {gpu['index']}: {gpu['name']} ({gpu['total_memory_mb']} MB)" + ) + # Initialize managed loader with GPU boss for coordinated loading _diffusers_loader = ManagedModelLoader( boss=_boss,