nodes/JasonHoku-ComfyUI-Ultimate-Auto-Sampler-Config-Grid-Testing-Suite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(distribution): delete distribution_manager.py (~650 lines)

Browse Source 
DistributionManager + DistributedJob + JobState moved to
ComfyUI-USCG-Distributed companion. Manager now accepts
job_completed_check/build_prompt_with_triggers/get_model_cache_key as
callbacks to avoid bidirectional imports.
This commit is contained in:
Jason Hoku
2026-05-19 15:26:06 -07:00
parent cc27f63040
commit 12ff7bf39b
1 changed files with 0 additions and 650 deletions
Download Patch File Download Diff File Expand all files Collapse all files
distribution_manager.py
-650
View File
@@ -1,650 +0,0 @@
"""
Distributed Task Manager
Thread-safe job queue for distributing generation work across multiple ComfyUI instances.
Uses a pull-based model: workers claim jobs from the master via API.
"""
import threading
import time
import json
import os
import uuid
class JobState:
"""Job lifecycle states."""
PENDING = "pending"
CLAIMED = "claimed"
COMPLETED = "completed"
FAILED = "failed"
class DistributedJob:
"""Represents a single job in the distributed queue."""
__slots__ = [
'job_id', 'config', 'input_job', 'state', 'worker_id',
'claimed_at', 'completed_at', 'error', 'gen_index', 'fail_count'
]
def __init__(self, job_id, config, input_job, gen_index):
self.job_id = job_id
self.config = config # Expanded config dict (sampler, model, lora, etc.)
self.input_job = input_job # {label, width, height, batch_idx}
self.state = JobState.PENDING
self.worker_id = None
self.claimed_at = None
self.completed_at = None
self.error = None
self.gen_index = gen_index # Sequential index for deterministic sort ordering
self.fail_count = 0
class DistributionManager:
"""
Thread-safe job distribution coordinator.
Manages an in-memory job list with atomic claim/complete/fail operations.
The master populates jobs at startup, then workers claim them one at a time.
Jobs that time out (worker crash) are automatically re-queued.
"""
# Max times a job can fail before being permanently abandoned
MAX_FAIL_COUNT = 3
def __init__(self, session_name, paths, unique_id, existing_data,
claim_timeout_seconds=600):
self._lock = threading.Lock()
self._jobs = {} # job_id -> DistributedJob
self._pending_queue = [] # List of job_ids in FIFO order
self._workers = {} # worker_id -> {url, last_heartbeat, jobs_completed, ...}
self.session_name = session_name
self.paths = paths
self.unique_id = unique_id
self.existing_data = existing_data # Shared manifest reference
self.claim_timeout = claim_timeout_seconds
self._persistence_path = os.path.join(paths["base"], "distribution_state.json")
self._active = True
self._total_completed = 0
# Pre-encoded conditionings from master, keyed by "(model_key|lora_key)"
# Structure: {"model|lora": {"positive": {prompt: serialized}, "negative": {...}}}
# Populated by master pre-encoding phase; attached to job claim responses
self._encoded_conditionings = {}
# Model sync: when enabled, job claim responses include upcoming_models for pre-fetching
self.sync_models_to_workers = False
# Session-level settings (upscaling, cooldown) passed through to workers
self._session_settings = None
def populate_jobs(self, expanded_configs, input_jobs, existing_data,
overwrite_existing, has_optional_inputs, lora_triggerwords_mode):
"""
Build the full job list from expanded configs × input jobs.
Filters out already-completed jobs using check_if_job_completed().
Args:
expanded_configs: List of expanded config dicts (from expand_configs())
input_jobs: List of input job dicts (from prepare_input_jobs())
existing_data: Loaded manifest data for skip detection
overwrite_existing: If True, don't skip completed jobs
has_optional_inputs: If True, relaxed skip matching
lora_triggerwords_mode: Trigger word mode string
"""
from .generation_orchestrator import check_if_job_completed
from .trigger_words import build_prompt_with_triggers
existing_items = existing_data.get("items", [])
gen_index_offset = len(existing_items)
skipped = 0
total = 0
with self._lock:
job_index = 0
for job_idx, job in enumerate(input_jobs):
w = job["width"]
h = job["height"]
batch_idx = job.get("batch_idx", 0)
for conf_idx, conf in enumerate(expanded_configs):
total += 1
# Per-config resolution override
if conf.get("resolution") is not None:
if job_idx > 0:
continue
w = conf["resolution"][0]
h = conf["resolution"][1]
batch_idx = 0
# Compute actual prompts (with trigger words applied)
# build_prompt_with_triggers returns (final_positive, trigger_string)
try:
actual_positive, _ = build_prompt_with_triggers(
conf, lora_triggerwords_mode
)
except Exception:
actual_positive = conf.get("positive", "")
actual_negative = conf.get("negative", "")
# Skip detection (same logic as main generation loop)
if not overwrite_existing:
match_index = check_if_job_completed(
existing_items, conf,
conf.get("seed", 0), w, h, batch_idx,
actual_positive, actual_negative,
has_optional_inputs=has_optional_inputs
)
if match_index != -1:
skipped += 1
continue
# Create job with serializable input_job (no latent tensors)
serializable_input_job = {
"label": job.get("label", f"{w}x{h}"),
"width": w,
"height": h,
"batch_idx": batch_idx
}
job_id = str(uuid.uuid4())[:12]
gen_index = gen_index_offset + job_index
distributed_job = DistributedJob(
job_id=job_id,
config=conf,
input_job=serializable_input_job,
gen_index=gen_index
)
self._jobs[job_id] = distributed_job
self._pending_queue.append(job_id)
job_index += 1
self._persist_state()
print(f"[Distribution] 📋 Populated {len(self._jobs)} jobs "
f"({skipped} skipped, {total} total before filtering)")
def claim_job(self, worker_id):
"""
Atomically claim the next pending job for a worker.
Returns:
dict with job details, or None if no pending jobs.
"""
with self._lock:
# Release timed-out claimed jobs first
self._release_timed_out_jobs()
if not self._pending_queue:
return None
job_id = self._pending_queue.pop(0)
job = self._jobs[job_id]
job.state = JobState.CLAIMED
job.worker_id = worker_id
job.claimed_at = time.time()
self._persist_state()
return self._job_to_dict(job)
def complete_job(self, job_id, metadata=None):
"""
Mark a job as completed.
Guards against double-completion (can happen if a job times out and is
reclaimed by another worker, then the original worker finishes late).
Args:
job_id: The job ID string
metadata: Optional metadata dict (not stored in job, just for logging)
Returns:
True if job was found and completed, False otherwise.
"""
with self._lock:
job = self._jobs.get(job_id)
if not job:
print(f"[Distribution] ⚠️ complete_job: unknown job_id {job_id}")
return False
# Prevent double-completion (job may have been reclaimed after timeout)
if job.state == JobState.COMPLETED:
print(f"[Distribution] ⚠️ complete_job: job {job_id} already completed, ignoring duplicate")
return False
job.state = JobState.COMPLETED
job.completed_at = time.time()
self._total_completed += 1
# Update worker stats
if job.worker_id and job.worker_id in self._workers:
self._workers[job.worker_id]["jobs_completed"] += 1
self._persist_state()
return True
def fail_job(self, job_id, error_msg=""):
"""
Mark a job as failed and re-queue to pending (up to MAX_FAIL_COUNT times).
Args:
job_id: The job ID string
error_msg: Error description
Returns:
True if job was found and re-queued, False otherwise.
"""
with self._lock:
job = self._jobs.get(job_id)
if not job:
print(f"[Distribution] ⚠️ fail_job: unknown job_id {job_id}")
return False
job.fail_count += 1
# Update worker stats
if job.worker_id and job.worker_id in self._workers:
self._workers[job.worker_id]["jobs_failed"] += 1
if job.fail_count >= self.MAX_FAIL_COUNT:
job.state = JobState.FAILED
job.error = error_msg
print(f"[Distribution] ❌ Job {job_id} permanently failed after "
f"{job.fail_count} attempts: {error_msg}")
else:
job.state = JobState.PENDING
job.worker_id = None
job.claimed_at = None
job.error = error_msg
self._pending_queue.append(job_id)
print(f"[Distribution] 🔄 Job {job_id} re-queued "
f"(attempt {job.fail_count}/{self.MAX_FAIL_COUNT}): {error_msg}")
self._persist_state()
return True
def claim_batch_for_local(self, count):
"""
Claim a batch of jobs for the master's local processing.
Uses worker_id "master" to distinguish from remote workers.
Args:
count: Maximum number of jobs to claim
Returns:
List of job dicts (may be empty).
"""
with self._lock:
self._release_timed_out_jobs()
batch = []
for _ in range(min(count, len(self._pending_queue))):
job_id = self._pending_queue.pop(0)
job = self._jobs[job_id]
job.state = JobState.CLAIMED
job.worker_id = "master"
job.claimed_at = time.time()
batch.append(self._job_to_dict(job))
if batch:
self._persist_state()
return batch
def heartbeat(self, worker_id):
"""Update worker heartbeat timestamp."""
with self._lock:
if worker_id in self._workers:
self._workers[worker_id]["last_heartbeat"] = time.time()
def register_worker(self, worker_id, worker_url=""):
"""
Register a new worker.
Args:
worker_id: Unique worker identifier
worker_url: Worker's base URL (for display/monitoring only)
"""
with self._lock:
self._workers[worker_id] = {
"url": worker_url,
"last_heartbeat": time.time(),
"jobs_completed": 0,
"jobs_failed": 0,
"jobs_duplicated": 0,
"status": "idle"
}
print(f"[Distribution] 🤝 Worker registered: {worker_id} ({worker_url})")
def record_late_submission(self, worker_id):
"""
Record a duplicate/late submission from a worker.
Called when a worker submits a result for a job that was already completed
(e.g. because the job timed out, was reclaimed by master, and completed
before the worker finished). The worker did real work — this tracks it.
Args:
worker_id: The worker that submitted late
"""
with self._lock:
if worker_id in self._workers:
self._workers[worker_id]["jobs_duplicated"] += 1
def get_status(self):
"""Return overall progress status."""
with self._lock:
total = len(self._jobs)
pending = len(self._pending_queue)
claimed = sum(1 for j in self._jobs.values()
if j.state == JobState.CLAIMED)
completed = sum(1 for j in self._jobs.values()
if j.state == JobState.COMPLETED)
failed = sum(1 for j in self._jobs.values()
if j.state == JobState.FAILED)
return {
"total": total,
"pending": pending,
"claimed": claimed,
"completed": completed,
"failed": failed,
"workers": {
wid: info.copy()
for wid, info in self._workers.items()
}
}
def deactivate(self):
"""Mark the distribution manager as inactive (generation complete)."""
with self._lock:
self._active = False
@property
def is_active(self):
"""Check if distribution is still active."""
with self._lock:
return self._active
@property
def has_pending_or_claimed(self):
"""True if any jobs are still pending or claimed (excluding failed)."""
with self._lock:
return any(
j.state in (JobState.PENDING, JobState.CLAIMED)
for j in self._jobs.values()
)
@property
def total_completed(self):
"""Number of completed jobs."""
with self._lock:
return self._total_completed
def set_encoded_conditionings(self, encoded_dict):
"""
Store pre-encoded conditionings from the master's pre-encoding phase.
Args:
encoded_dict: Dict keyed by "(model_cache_key|lora_key)" with structure:
{"positive": {prompt_text: serialized_list}, "negative": {prompt_text: serialized_list}}
"""
with self._lock:
self._encoded_conditionings = encoded_dict
total_groups = len(encoded_dict)
total_prompts = sum(
len(v.get("positive", {})) + len(v.get("negative", {}))
for v in encoded_dict.values()
)
print(f"[Distribution] 🧠 Stored pre-encoded conditionings: "
f"{total_groups} model/LoRA groups, {total_prompts} total prompt encodings")
def get_encoded_conditionings_for_job(self, job):
"""
Look up pre-encoded conditionings for a specific job.
Uses the job's config to compute the group key and prompt text,
then returns the matching serialized positive/negative conditionings.
Args:
job: DistributedJob instance
Returns:
dict with "encoded_positive" and "encoded_negative" keys, or None if not available.
"""
if not self._encoded_conditionings:
return None
try:
from .generation_orchestrator import get_model_cache_key
from .trigger_words import build_prompt_with_triggers
config = job.config
# Build the group key: "model_cache_key|lora_key"
model_key = get_model_cache_key(config)
lora_key = config.get("lora_expanded", config.get("lora", "None"))
group_key = f"{model_key}|{lora_key}"
group = self._encoded_conditionings.get(group_key)
if not group:
return None
# Build the actual prompt text (with trigger words applied)
lora_triggerwords_mode = config.get("_lora_triggerwords_mode", "None")
try:
actual_positive, _ = build_prompt_with_triggers(config, lora_triggerwords_mode)
except Exception:
actual_positive = config.get("positive", "")
actual_negative = config.get("negative", "")
encoded_positive = group.get("positive", {}).get(actual_positive)
encoded_negative = group.get("negative", {}).get(actual_negative)
if encoded_positive is not None and encoded_negative is not None:
return {
"encoded_positive": encoded_positive,
"encoded_negative": encoded_negative,
# Include final prompt text so workers can use it for metadata
# without needing to re-run trigger word lookups
"actual_positive_text": actual_positive,
"actual_negative_text": actual_negative,
}
return None
except Exception as e:
print(f"[Distribution] ⚠️ Error looking up conditionings for job {job.job_id}: {e}")
return None
def _release_timed_out_jobs(self):
"""
Release jobs that have been claimed but not completed within timeout.
Must be called with lock held.
Uses the worker's last heartbeat timestamp (not just job claimed_at) to
determine if the worker is still alive. A slow-but-alive worker that
sends heartbeats will NOT have its jobs reclaimed. Only workers that
have stopped heartbeating (crash/disconnect) will have jobs reclaimed.
"""
now = time.time()
released = 0
for job in self._jobs.values():
if (job.state == JobState.CLAIMED
and job.worker_id != "master"
and job.claimed_at
and now - job.claimed_at > self.claim_timeout):
# Check if the worker is still heartbeating — if so, it's alive
# and just slow (e.g. slow GPU, first-time model loading).
# Only reclaim if the worker has gone silent.
worker = self._workers.get(job.worker_id)
if worker:
last_hb = worker.get("last_heartbeat", 0)
if now - last_hb < self.claim_timeout:
# Worker is still alive — don't reclaim
continue
print(f"[Distribution] ⏰ Job {job.job_id} timed out "
f"from worker {job.worker_id} (no heartbeat), re-queuing")
job.state = JobState.PENDING
job.worker_id = None
job.claimed_at = None
self._pending_queue.append(job.job_id)
released += 1
if released > 0:
print(f"[Distribution] Released {released} timed-out jobs")
def release_timed_out_jobs(self):
"""Public wrapper — call from the wait loop to reclaim dead worker jobs."""
with self._lock:
self._release_timed_out_jobs()
def _persist_state(self):
"""
Persist job states to disk for crash recovery.
Must be called with lock held.
"""
state = {
"session_name": self.session_name,
"timestamp": time.time(),
"jobs": {
jid: {
"state": j.state,
"worker_id": j.worker_id,
"fail_count": j.fail_count
}
for jid, j in self._jobs.items()
}
}
try:
with open(self._persistence_path, "w") as f:
json.dump(state, f)
except Exception as e:
print(f"[Distribution] ⚠️ Persist error: {e}")
def _get_upcoming_models(self, exclude_job_id=None, limit=10):
"""
Scan pending jobs and collect unique model references for pre-fetching.
Must be called with lock held.
Args:
exclude_job_id: Job to skip (the current job being claimed)
limit: Maximum number of unique model references to return
Returns:
List of {category, filename} dicts
"""
seen = set()
upcoming = []
for job_id in self._pending_queue:
if job_id == exclude_job_id:
continue
job = self._jobs.get(job_id)
if not job:
continue
config = job.config
# Main model/checkpoint
model_name = config.get("model", "None")
if model_name and model_name != "None":
model_type = config.get("model_type", "checkpoint")
if model_type == "checkpoint":
cat = "checkpoints"
elif model_type == "gguf":
cat = "unet_gguf"
elif model_type == "diffusion_model":
cat = "diffusion_models"
else:
cat = "checkpoints"
key = f"{cat}:{model_name}"
if key not in seen:
seen.add(key)
upcoming.append({"category": cat, "filename": model_name})
# LoRA — combo string: "name:str:str + name2:str:str", may include folders
lora = config.get("lora_expanded", config.get("lora", "None"))
if lora and lora != "None":
for lora_part in lora.split(" + "):
lora_filename = lora_part.strip().split(":")[0].strip()
if not lora_filename or lora_filename == "None":
continue
# Skip folder references — workers expand those locally
if lora_filename.endswith("/") or lora_filename.endswith("/*") or ("[" in lora_filename and "]" in lora_filename):
continue
key = f"loras:{lora_filename}"
if key not in seen:
seen.add(key)
upcoming.append({"category": "loras", "filename": lora_filename})
# VAE (skip remote URLs and "Default")
vae = config.get("vae", "Default")
if vae and vae != "Default" and not vae.startswith("remote:"):
key = f"vae:{vae}"
if key not in seen:
seen.add(key)
upcoming.append({"category": "vae", "filename": vae})
# Text encoders
text_encoders = config.get("text_encoders", [])
for te in text_encoders:
if te and te != "None":
key = f"text_encoders:{te}"
if key not in seen:
seen.add(key)
upcoming.append({"category": "text_encoders", "filename": te})
if len(upcoming) >= limit:
break
return upcoming
def _job_to_dict(self, job):
"""
Convert a DistributedJob to a serializable dict for API responses.
Note: input_job latent tensors are NOT included (workers create their own).
Config is round-tripped through JSON to ensure all values are serializable
(e.g. numpy int64 → Python int, sets → lists, etc.).
If pre-encoded conditionings are available, they are attached to the response
so workers can skip CLIP encoding entirely.
"""
try:
safe_config = json.loads(json.dumps(job.config, default=str))
except Exception:
safe_config = job.config
result = {
"job_id": job.job_id,
"config": safe_config,
"input_job": job.input_job,
"gen_index": job.gen_index
}
# Attach session-level cooldown settings to worker config
if hasattr(self, '_session_settings') and self._session_settings and 'cooldown' in self._session_settings:
safe_config["_session_cooldown"] = self._session_settings["cooldown"]
# Attach image_format to worker config so distributed workers honour the user's choice
if hasattr(self, '_session_settings') and self._session_settings:
safe_config["_session_image_format"] = self._session_settings.get("image_format", "webp")
# Attach pre-encoded conditionings if available (master encoding feature)
encoded = self.get_encoded_conditionings_for_job(job)
if encoded:
result["encoded_positive"] = encoded["encoded_positive"]
result["encoded_negative"] = encoded["encoded_negative"]
# Include final prompt text (with triggers applied) so workers can
# use it for metadata without re-running CivitAI trigger lookups
result["actual_positive_text"] = encoded.get("actual_positive_text")
result["actual_negative_text"] = encoded.get("actual_negative_text")
# Attach upcoming models list if sync is enabled (for pre-fetching)
if self.sync_models_to_workers:
result["upcoming_models"] = self._get_upcoming_models(
exclude_job_id=job.job_id, limit=10
)
return result