From 9a526e254660fa6c9b48ea20deba16cdf631024e Mon Sep 17 00:00:00 2001 From: John Pollock Date: Tue, 30 Sep 2025 09:10:37 -0500 Subject: [PATCH] prepare for final release candidate --- .clinerules | 333 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 200 insertions(+), 133 deletions(-) diff --git a/.clinerules b/.clinerules index 9c8d208..f04aadf 100644 --- a/.clinerules +++ b/.clinerules @@ -3,17 +3,22 @@ ## Project Context This is ComfyUI-MultiGPU: a production-grade multi-device AI inference platform that transforms ComfyUI from single-GPU to universal multi-device support. The project enables previously impossible AI workflows across diverse hardware configurations. +**Current Version**: v2.5.0 Release Candidate +**Status**: PRODUCTION READY +**Stability**: 9/10 - Verified working in production +**Community**: 300+ commits, 90+ resolved issues, active ecosystem + ## Memory Bank System **CRITICAL**: Always read ALL files in the `memory-bank/` folder at the start of every session. The Memory Bank contains complete project context: ### Core Documentation (Read These First) 1. `memory-bank/projectbrief.md` - Project identity, mission, evolution timeline 2. `memory-bank/productContext.md` - Problem space, user goals, success metrics -3. `memory-bank/activeContext.md` - Current work focus and priorities -4. `memory-bank/progress.md` - Production status, roadmap, lessons learned +3. `memory-bank/activeContext.md` - Current work focus and priorities (UPDATED 2025-09-30) +4. `memory-bank/progress.md` - Production status, roadmap, lessons learned (UPDATED 2025-09-30) ### Technical Deep Dive -5. `memory-bank/systemPatterns.md` - Architecture patterns and design decisions +5. `memory-bank/systemPatterns.md` - Architecture patterns and design decisions (UPDATED 2025-09-30) 6. `memory-bank/techContext.md` - Technology stack and development environment 7. `memory-bank/performance-benchmarks.md` - Quantified performance across hardware configurations 8. `memory-bank/comfyui-lineage.md` - Integration analysis with ComfyUI core @@ -24,150 +29,91 @@ This is ComfyUI-MultiGPU: a production-grade multi-device AI inference platform - **User Agency**: Let users specify device placement explicitly - **Production Quality**: Stability and reliability over experimental features - **Community First**: Solutions should benefit the entire ComfyUI ecosystem +- **Clean Code**: Remove debug artifacts, comprehensive production logging only ## Key Technical Patterns -- **City96's Dynamic Class Override**: Elegant inheritance pattern for node creation +- **City96's Dynamic Class Override**: Elegant inheritance pattern for node creation (50 lines vs 400+) - **Load-Patch-Distribute Pipeline**: Quality-preserving LoRA application workflow - **ComfyCore Alignment**: Work WITH existing ComfyUI patterns, not against them - **Multi-Device Native**: Treat all devices as equal citizens +- **Selective Unload**: Per-model granular control over memory management -## Current Status (v2.4.7) -- Production Grade: 300+ commits, 90 resolved issues, active community -- Performance Validated: NVLink near-native (5-7% slowdown), PCIe 4.0 excellent (40-50%) -- Ecosystem Integration: 10+ custom nodes with dynamic detection -- Hardware Support: Universal compatibility across CUDA, CPU, MPS, XPU, NPU, DirectML +## Production Status (v2.5.0) -When working on this project, always reference the Memory Bank for context and maintain the established patterns and philosophy. +### Core Features ✅ +- **DisTorch2 Distributed Loading**: Universal SafeTensor support with CLIP head preservation +- **Selective Unload System**: Verified working - keeps models with `keep_loaded=True`, ejects others +- **Multi-Device VRAM Management**: Clears allocator caches across all devices +- **Manager Parity**: Mirrors ComfyUI-Manager "Free model and node cache" behavior +- **Universal Device Support**: CUDA, CPU, MPS, XPU, NPU, MLU, DirectML, CoreX -## CRITICAL ACTIVE ISSUE: CPU Memory Leak Investigation +### Recent Achievements (2025-09-30) +- **Code Refactoring** (-219 lines total): + - DisTorch2 allocation consolidation (-179 lines): Unified UNET and CLIP allocation functions + - Production cleanup (-40 lines): Removed diagnostic instrumentation wrapper +- **Verified Working**: Selective unload tested in production with comprehensive logging +- **Clean Architecture**: Single responsibility modules, clear dependency direction -### Problem Description -**SEVERE**: Monotonic CPU DRAM increase across successive model swaps. VRAM management stable, but CPU memory allocated and never reclaimed, leading to eventual OOM failure of ComfyUI process. +### Performance Validation +- **NVLink**: 5-7% slowdown (near-native) +- **PCIe 4.0 x16**: 40-50% slowdown (excellent) +- **PCIe 3.0 x16**: 70-80% slowdown (good) +- **PCIe 4.0 x8**: 80-100% slowdown (acceptable) +- **PCIe 3.0 x8**: 150-200% slowdown (workable) +- **PCIe 3.0 x4**: 300-400% slowdown (last resort) -**Pattern**: Stepwise increases (4GB → 13GB → 33GB → 57GB → 73GB) corresponding to model loading events. Large objects (tensors, model structures, intermediates) retained in Python/CPU memory across executions. - -### The Benchmark (Known Working Solution) -**ONLY** successful CPU memory reclamation: ComfyUI "Free model and node cache" button -- Sets `"free_memory": True` → PromptExecutor.reset() → drops execution caches -- Sets `"unload_models": True` → comfy.model_management.unload_all_models() - -### Investigation History (All Failed) - -**Phase 1: Missing Executor Reset Hypothesis** ❌ -- Theory: MultiGPU cleared VRAM but failed PromptExecutor.reset() -- Action: Implemented trigger_executor_cache_reset utilities -- Result: FAILED - e.reset() occurred but CPU memory not reclaimed - -**Phase 2: Implementation Fixes** ❌ -- Theory: Flawed implementation or insufficient visibility -- Action: Fixed bugs, added extensive memory logging -- Result: FAILED - Logs confirmed resets occurring, CPU usage still rising - -**Phase 3: Aggressive Reclamation** ❌ -- Theory: References outside execution cache or allocator fragmentation -- Actions: malloc_trim(0), store pruning, reference cycle patching, deep diagnostics -- Result: FAILED - OOM persisted, diagnostics showed "Tracked ModelPatchers=0" - -**CRITICAL DIAGNOSTIC FAILURE**: Deep diagnostics completely failed due to patching `partially_load` (conditional) instead of `__init__` (universal) - zero visibility into actual leak. - -### Root Cause Analysis -- Execution cache and allocator fragmentation ruled out -- Live Python references held to large data structures -- References NOT cleared by gc.collect() or e.reset() -- Benchmark button works = "unload_models": True is the critical difference -- unload_all_models() successfully breaks reference chains holding CPU memory - -### Mandated Plan Forward (FINALIZED SOLUTION) -**Resolution**: CPU memory leaks eliminated via transient 3-flag selective ejection system - -#### Core Principle: `keep_loaded` Boolean Drives 3 Execution Behaviors -The `keep_loaded` boolean serves triple duty when set to "False": -1. **Load-Time Preservation**: Returns MAX_VRAM in `model_memory_required()` → forces Comfy to evict other models pre-loading -2. **Ejection Trigger**: Workflow detects `keep_loaded=False` → sets transient flags for selective unloading -3. **Surgical Destruction**: End-of-workflow unload applies wrecking ball ONLY to flagged DisTorch models - -#### 3-Transient-Flags Architecture - - -**Per-Model Flag**: `_distorch2_unload_model = TRUE/FALSE` (object-scoped) -- Marks specific DisTorch models for distributed device ejection -- Applied during load phase to models with `keep_loaded=False` -- Cleared after ejection (transient marker) - -**Comfy Core Flag**: `PromptExecutor.unload_all_models = TRUE` (standard) -- Triggered by DisTorch logic at end-of-workflow -- Calls our patched `unload_all_models()` method -- Generates the selective ejection signal - -#### Implementation Plan: Code Changes Required - -**Phase 1: Flag Setting (distorch_2.py)** -```python -# In DistTorch load override - detect keep_loaded=False during execution -if hasattr(out[0], 'model') and hasattr(out[0].model, '_mgpu_keep_loaded'): - is_distorch2_keep_false = (out[0].model._mgpu_keep_loaded == False) - if is_distorch2_keep_false: - # Set transient flags for selective ejection - globals()['DISTORCH2_UNLOAD_MODEL'] = True - out[0].model._distorch2_unload_model = True - set_prompt_executor_unload_flag() -``` - -**Phase 2: Surgical Unload Logic (model_management_mgpu.py)** -```python -# Check: Are we in DisTorch ejection mode? -distorch_ejection_mode = any( - getattr(getattr(lm.model, 'model', None), '_distorch2_unload_model', False) - for lm in mm.current_loaded_models -) - -if not distorch_ejection_mode: - # Normal Comfy unload - delegate to original - return _mgpu_original_unload_all_models() - -# SURGICAL MODE: Only process flagged models -for lm in mm.current_loaded_models: - if hasattr(getattr(lm.model, 'model', None), '_distorch2_unload_model'): - # WRECKING BALL: Eject from all distributed device locations - apply_distributed_device_cleanup(lm.model) - # else: SKIP ENTIRELY - no processing of any kind - -# Reset transient flags after surgical operation -globals()['DISTORCH2_UNLOAD_MODEL'] = False -for lm in mm.current_loaded_models: - if hasattr(lm.model, 'model') and hasattr(lm.model.model, '_distorch2_unload_model'): - delattr(lm.model.model, '_distorch2_unload_model') -``` - -#### Behavioral Guarantee -- **Same workflow re-run**: Deterministic - flags reset per execution -- **External unload calls**: No flags set → normal Comfy behavior -- **Normal Comfy models**: Never flag-munged → standard unload behavior -- **DisTorch models with `keep_loaded=True`**: Handle via standard Comfy unload -- **DisTorch models with `keep_loaded=False`**: Surgical ejection from distributed devices - -#### Key Advantages -- **No persistent state**: Flags reset after each operation -- **Surgical precision**: Only tagged models processed -- **Comfy compatibility**: External calls unaffected -- **Execution isolation**: Each workflow manages its own ejection -- **Memory safety**: Designed for CPU leaks elimination through proper distributed cleanup - -**Implementation Status**: ✅ IMPLEMENTED - Phase 3 selective ejection fully deployed in code. Per-model `_mgpu_unload_distorch_model` flags, patched `unload_all_models`, and Manager parity via `force_full_system_cleanup` are all active. 🐛 REMAINING BUG: Selective retention not working correctly - retained models (keep_loaded=True) are still being ejected in some flows despite the selective unload logic being present. +### Ecosystem Integration +- 10+ custom node integrations with automatic detection +- Dynamic node creation for compatible loaders +- Fail-loudly compatibility with ComfyCore API ## Module Architecture Rules ### Module Boundary Principles - **Single Responsibility**: Each module should have ONE clear purpose -- **Dependency Direction**: Dependencies should flow in ONE direction only -- **Import Hierarchy**: Lower-level modules (device_utils) should NOT import from higher-level modules (distorch_2) +- **Dependency Direction**: Dependencies flow UPWARD only - violations create circular imports +- **Import Hierarchy**: Base modules NEVER import from Feature/UI modules ### Module Hierarchy (Dependency Order) -1. `device_utils.py` - **BASE**: Device detection, VRAM management only -2. `model_management_mgpu.py` - **CORE**: Model lifecycle, memory logging, cleanup functions -3. `distorch_2.py`, `distorch.py` - **FEATURES**: DisTorch distribution logic -4. `nodes.py`, `checkpoint_multigpu.py` - **UI**: Node implementations -5. `__init__.py` - **ASSEMBLY**: Final integration and registration +1. **`device_utils.py`** - BASE LAYER + - Device detection and enumeration + - VRAM cache management (`soft_empty_cache_multigpu`) + - Pure hardware abstraction - NO model tracking + +2. **`model_management_mgpu.py`** - CORE LAYER + - Model lifecycle tracking + - Memory logging infrastructure + - Cleanup orchestration (`force_full_system_cleanup`, `trigger_executor_cache_reset`) + - Patched `mm.unload_all_models` (selective ejection) + +3. **`distorch_2.py`, `distorch.py`** - FEATURE LAYER + - DisTorch distribution algorithms + - Allocation analysis and device assignment + - Per-model flag setting (`_mgpu_unload_distorch_model`) + - Imports from CORE/BASE only + +4. **`nodes.py`, `checkpoint_multigpu.py`** - UI LAYER + - Device-aware user interfaces + - Node implementations and definitions + - Imports from any lower level + +5. **`__init__.py`** - ASSEMBLY LAYER + - Final integration and patch registration + - Node mapping and registration + - Imports from all lower levels + +### Import Flow Architecture +``` + __init__.py ← Assembly + ↑ + UI Layer ← nodes.py, checkpoint_multigpu.py + ↑ + Feature Layer ← distorch_2.py, distorch.py + ↑ + Core Layer ← model_management_mgpu.py + ↑ + Base Layer ← device_utils.py +``` ### Mandatory Architecture Checks **BEFORE adding ANY import statement:** @@ -176,14 +122,135 @@ for lm in mm.current_loaded_models: 3. **Check Cycles**: Run `python -c "import sys; sys.path.append('.'); import "` to detect circular imports ### Function Placement Rules -- **device_utils.py**: ONLY device detection, VRAM cache management -- **model_management_mgpu.py**: Model tracking, memory logging, cleanup utilities +- **device_utils.py**: ONLY device detection, VRAM cache management +- **model_management_mgpu.py**: Model tracking, memory logging, cleanup utilities - **Feature modules**: Import from CORE/BASE only, never each other - **UI modules**: Import from any lower level, implement user interfaces only ### Violation Detection If import fails with "circular import" or "cannot import name": 1. STOP immediately - do not work around -2. Identify which module boundary was violated +2. Identify which module boundary was violated 3. Move misplaced function to correct architectural layer 4. Update ALL imports consistently + +## Memory Management System (Verified Working) + +### Selective Unload Pipeline +**Load Phase**: +```python +# DisTorch2 wrapper sets per-model flag based on keep_loaded parameter +if hasattr(out[0], 'model') and hasattr(out[0].model, '_mgpu_keep_loaded'): + keep_loaded = out[0].model._mgpu_keep_loaded + out[0].model._mgpu_unload_distorch_model = (not keep_loaded) +``` + +**Unload Phase** (patched `mm.unload_all_models`): +```python +# Categorize models by flag +models_to_unload = [flagged models] +kept_models = [unflagged models] + +if kept_models: + # Selective: eject flagged, retain others with GC anchors + for lm in models_to_unload: + lm.model_unload(unpatch_weights=True) + mm.current_loaded_models = kept_models +else: + # Standard cleanup when no models to keep + _mgpu_original_unload_all_models() +``` + +**Verified Working** (Production Logs 2025-09-30): +``` +[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 +[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 +[REMAINING_MODEL] 0: AutoencodingEngine +[REMAINING_MODEL] 1: FluxClipModel_ +``` + +### Manager Parity +`force_full_system_cleanup()` mirrors ComfyUI-Manager "Free model and node cache": +- Sets `unload_models=True`, `free_memory=True` on PromptQueue +- Triggers patched `mm.unload_all_models` for selective ejection +- Triggers `PromptExecutor.reset()` for CPU memory management + +## Code Quality Standards + +### Production Requirements +- **No Debug Cruft**: Remove all diagnostic-only code before release +- **Comprehensive Logging**: Production-grade telemetry at major operations +- **Clean Modules**: Single responsibility, clear boundaries +- **Fail Loudly**: Surface API changes immediately, no defensive masking + +### Logging Conventions +```python +# Model Management logs +logger.mgpu_mm_log("[OPERATION] Description with context") + +# Memory state logging +multigpu_memory_log("identifier", "tag") + +# Debug logging (use sparingly) +logger.debug("[Component] Detailed diagnostic information") +``` + +### Code Style +- Self-documenting code over excessive comments +- Clear function/variable names conveying intent +- Minimal comments for non-obvious constraints only +- Structured logging for production debugging + +## Development Workflow + +### Before Making Changes +1. Read relevant Memory Bank files +2. Understand module architecture and dependencies +3. Check if change violates architectural boundaries +4. Consider impact on existing patterns + +### When Adding Features +1. Determine correct module placement (BASE/CORE/FEATURE/UI) +2. Verify no circular dependencies created +3. Add comprehensive logging at key operations +4. Test with production workflows +5. Update Memory Bank documentation + +### When Refactoring +1. Eliminate code duplication (DRY principle) +2. Remove debug artifacts and diagnostic code +3. Maintain or improve architectural clarity +4. Verify no functionality regressions +5. Document pattern changes in systemPatterns.md + +## Testing Philosophy + +### Manual Validation +- Test across hardware configurations (NVLink, PCIe variants, CPU) +- Verify selective unload with keep_loaded combinations +- Check memory usage patterns (VRAM + CPU) +- Validate quality parity with single-GPU baselines + +### Community Testing +- Active users provide hardware configuration validation +- Integration testing with custom node ecosystem +- Performance feedback across diverse setups + +## Next Steps (v2.5.0 Release) + +### Immediate +- [ ] Final testing pass across hardware configurations +- [ ] GitHub release notes and changelog +- [ ] Community announcement + +### Short-term +- [ ] Issue triage and community feedback integration +- [ ] New model format support (Mochi, community requests) +- [ ] Documentation refresh and tutorials + +### Long-term +- [ ] Model parallelism research +- [ ] Streaming inference for video +- [ ] Multi-node orchestration + +When working on this project, always reference the Memory Bank for context and maintain the established patterns and philosophy. The codebase is production-ready - focus on stability, community needs, and quality over experimental features.