nodes/pollockjj-ComfyUI-MultiGPU
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

prepare for final release candidate

Browse Source
This commit is contained in:
John Pollock
2025-09-30 09:10:37 -05:00
parent 23ed34df1b
commit 9a526e2546
1 changed files with 200 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.clinerules
+200 -133
View File
@@ -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 <module>"` 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.