diff --git a/ProjectStructure.md b/ProjectStructure.md index 4af201f..41c4020 100644 --- a/ProjectStructure.md +++ b/ProjectStructure.md @@ -582,6 +582,7 @@ The Dashboard is a standalone single-page app rendered inside an iframe. All cod "sync_models_to_workers": false } ``` +`wait_timeout` (seconds, optional, not sent by the Builder UI): how long the master waits on remote workers before taking their claimed jobs over; default `max(600, claim_timeout)` (`_dist_wait_timeout`). **Top-level state fields** (outside `config_arrays`): ```json @@ -802,7 +803,7 @@ __init__.py |---|---|---|---| | `ultimate_grid.update` | `generation_orchestrator.py` / `image_generation.py` / `remote_vae.py` / `workflow_sweep_collect.py` | `session_name, node, manifest, meta, new_items` | New images generated | | `ultimate_grid.update_data` | `generation_orchestrator.py`, `upscale_runner.py` | `session_name, new_items` (+ `node` from the orchestrator) | Per-item replacement: the listed items replace the cards with the same id (deferred upscales, dashboard upscales) | -| `ultimate_grid.progress` | `generation_orchestrator.py`, `upscale_runner.py`, `workflow_sweep_runner.py` | `session_name, node, progress_pct, eta_str, current_job, total_jobs, avg_duration, last_duration, finish_time` | ETA/progress bar. Optional `phase`: `upscaling` (deferred and dashboard upscales) or `sweep` (Workflow Sweep, label "Sweep job"). The sweep payload is `{node, session_name, phase: "sweep", sweep_id, state, errors, current_job, total_jobs, progress_pct, eta_str, finish_time, avg_duration, last_duration}` plus, on completion, `complete, total_elapsed, total_generated` | +| `ultimate_grid.progress` | `generation_orchestrator.py`, `upscale_runner.py`, `workflow_sweep_runner.py` | `session_name, node, progress_pct, eta_str, current_job, total_jobs, avg_duration, last_duration, finish_time` | ETA/progress bar. Optional `phase`: `upscaling` (deferred and dashboard upscales), `distributed` (distributed run: adds `workers`, the machines currently working, and the label gets ` | N machine(s)`; sent on every worker result via the manager's `on_status` hook, jobs skipped by Start At Job / claim filter / done on the master included in `current_job`) or `sweep` (Workflow Sweep, label "Sweep job"). The sweep payload is `{node, session_name, phase: "sweep", sweep_id, state, errors, current_job, total_jobs, progress_pct, eta_str, finish_time, avg_duration, last_duration}` plus, on completion, `complete, total_elapsed, total_generated` | | `ultimate_grid.next_run_seeds` | `generation_orchestrator.py` | `node, seeds` | Per-config seeds the NEXT run should use after a `random_after` / `increment_after` / `decrement_after` full-run seed behavior, as `{config_index: seed}` (no listener in this package) | | `ultimate_grid.metrics_progress` | `metrics_runner.py` | `session_name, job_id, status, completed, total` | Progress of a "compute metrics for this session" job (no listener in this package; the SPA polls `/config_tester/metrics_status`) | | `ultimate_grid.distribution_status` | `generation_orchestrator.py` | `session_name, workers, stats` | Distribution worker status | diff --git a/README.md b/README.md index 2d12eae..b748030 100644 --- a/README.md +++ b/README.md @@ -867,7 +867,7 @@ Limitations: no per-region LoRAs, equal splits only (no ratios), hard region edg ### Job Order & Preview Pass (Run Settings) Both live in the **Run Settings** section of the Config Builder and default to the old behaviour. -* **Job Order:** `As written (default)` runs one resolution pass after another, exactly as before. `Minimize model/LoRA reloads` re-groups the job list by model, LoRA stack, VAE and resolution. `As written` is already grouped by model (and Builder per-config resolutions stay inside their model group), so Minimize only changes anything when LoRA folders or random LoRA picks expand into different stacks (e.g. a model with 2 LoRA folders alternating A1 B1 A2 B2 becomes A1 A2 B1 B2), or when a batched latent is connected (one pass per latent). LoRA stacks are grouped before VAEs, so a model with several VAEs and several LoRA stacks can switch VAE more often than as written. The log prints the model/LoRA/VAE switch counts before and after and says when nothing changed. **Start At Job #** numbers follow the chosen order. In distributed mode the master computes the order and queues jobs to workers in it (see **Distributed parity** below); Start At Job is not applied there. +* **Job Order:** `As written (default)` runs one resolution pass after another, exactly as before. `Minimize model/LoRA reloads` re-groups the job list by model, LoRA stack, VAE and resolution. `As written` is already grouped by model (and Builder per-config resolutions stay inside their model group), so Minimize only changes anything when LoRA folders or random LoRA picks expand into different stacks (e.g. a model with 2 LoRA folders alternating A1 B1 A2 B2 becomes A1 A2 B1 B2), or when a batched latent is connected (one pass per latent). LoRA stacks are grouped before VAEs, so a model with several VAEs and several LoRA stacks can switch VAE more often than as written. The log prints the model/LoRA/VAE switch counts before and after and says when nothing changed. **Start At Job #** numbers follow the chosen order. In distributed mode the master computes the order and queues jobs to workers in it (see **Distributed parity** below); Start At Job skips the earlier plan positions there too, and they count as done in the progress bar. * **Preview Pass:** renders every cell cheaply first (fewer **Preview Steps**, smaller image via **Preview Scale** x0.25/0.33/0.5/0.75 snapped to the model's latent grid, optional **Preview CFG Override**, upscaling skipped) so you see the whole grid in a fraction of the time. Preview images carry a `PREVIEW` badge on the Dashboard. Reject the losers, then **untick Preview Pass and queue again**: only the surviving cells render at full quality, and each full image replaces its preview in the default Dashboard view. Interrupted preview passes resume (previews already done are skipped), and a preview pass skips cells that already have a full image. Not used when a latent input is connected; LTX video configs always render normally. Works in distributed mode too (the master decides which cells to render and stamps every preview, whichever machine rendered it). * **Skip Rejected Previews** (default on) and **Only Favorited Previews**: apply when Preview Pass is OFF. Cells whose preview you rejected are skipped, including after **Delete Rejected** (the rejection is remembered in `meta.rejected_preview_fps`). With *Only Favorited*, a cell that has a preview is rendered only if that preview is favorited; cells with no preview are still rendered. `Overwrite Existing` does not override a rejection. * A cell is identified by a **config fingerprint** (model, LoRA stack, VAE, text encoders, sampler, scheduler, CFG, denoise, seed, raw prompts, full resolution, plus advanced-sampling/SLG/deep-shrink settings; *not* steps). Cells whose seed is re-randomised on every run (`random_before`, extra random seeds) cannot be matched to a preview and simply render again. @@ -1706,7 +1706,12 @@ The master decides everything stateful and workers stay plain renderers: * Every item records which machine rendered it in `dist_worker` (`master` or the worker id). It is bookkeeping, not a setting: the pairwise review ignores it when deciding whether two images differ. * Rejected-preview tombstones, job order and the run-API session settings now survive a distributed run (they used to be overwritten). * Needs **ComfyUI-USCG-Distributed with `populate_from_plan`** on the master and an updated USCG on every worker (the worker needs `config_utils.preview_dims` for previews). With an older companion the master logs a warning and falls back to the old behaviour (job order and previews ignored). A preview job answered by a worker that ignored the render override is rejected and re-queued, never saved as a mislabelled full-size image. -* Not covered yet: re-checking rejects made while a distributed run is in progress, claim affinity (a worker preferring jobs for the model it already loaded), Start At Job #, LTX video on workers. +* **Rejects made during the run:** a full pass re-reads the dashboard's reject state when a job is claimed (at most every 2 s, only when the manifest changed). A preview you reject while the run is going is dropped at claim time (log: `dropped at claim: rejected_preview`) and counts as skipped in the progress bar. Not applied to the preview pass itself, or when optional model/conditioning inputs are connected. Needs the updated companion (`set_claim_filter`). +* **Claim affinity:** a worker (and the master) is preferably handed the next job for the model and LoRA stack it already has loaded, so a mixed run reloads far fewer checkpoints. Jobs stay in the master's plan order otherwise. +* **Start At Job #:** the first N-1 plan positions are skipped before any other check and counted as skipped. +* **LTX video** always renders on the master (workers cannot load it): the master renders the videos while the workers render the images, and an LTX job is never skipped as "already exists". With an older companion that has no `populate_from_plan`, a run containing LTX configs stops with an error asking you to update the companion. +* **Live progress and ETA:** the dashboard bar shows `Job x/y (p%) | ETA ... | Ns/job | N machine(s)` and updates whenever a worker returns a result. The ETA divides the remaining jobs by the combined throughput of every machine that is currently working (rolling average of the last 10 jobs each). +* Needs the updated **ComfyUI-USCG-Distributed** for the features above; with the old companion the master simply skips them. ### Model Sync @@ -1714,6 +1719,10 @@ When enabled, the master automatically transfers checkpoint and LoRA files to wo **Enable:** Toggle "Sync Models to Workers" in Distribution Settings. +**What is synced:** the checkpoint, the LoRAs, the VAE and the text encoders / diffusion models a job needs, into the matching folder on the worker (sub-folders are kept; model names with backslashes work between Windows and Linux). A file is downloaded once per worker even when several jobs need it. A VAE of `None` means the model's own VAE and is never requested. + +**Slow downloads:** the master waits for workers for the Distribution **wait timeout** (default: the larger of 600 s and the Claim Timeout; set `wait_timeout` in the distribution config to change it). When it is over, the master reclaims the jobs still claimed by workers and renders them itself, so a long model download no longer loses images; a late result from the worker is discarded as a duplicate. + ### Distribution Troubleshooting * **Workers not connecting:** Verify the worker URL is correct (include port). Check firewalls. Both master and worker must have this node suite installed. @@ -1721,6 +1730,8 @@ When enabled, the master automatically transfers checkpoint and LoRA files to wo * **Duplicate images:** If a worker times out but actually completes the job, duplicates may appear. Increase claim timeout to prevent this. * **Workers not using pre-encoded prompts:** Ensure "Use Master Text Encoding" is enabled. Check worker console for `"Using master pre-encoded conditionings"` log message. * **Model sync failing:** Ensure both master and worker have sufficient disk space. Check worker console for download errors. +* **Master says "rendering N reclaimed job(s) on the master":** the workers did not finish within the wait timeout (usually a big model download). The master took those jobs over; raise `wait_timeout` or sync the model once beforehand. +* **"LTX video configs ... need an updated ComfyUI-USCG-Distributed":** update the companion on the master, or turn Distribution off for the run. --- @@ -1820,7 +1831,7 @@ __init__.py |---|---|---|---| | `ultimate_grid.update` | `generation_orchestrator.py` | `{session_name, node, manifest, meta, new_items}` | New images generated | | `ultimate_grid.update_data` | `generation_orchestrator.py`, `upscale_runner.py` | per-item replacement payload | An existing item changed (for example a preview replaced by its full image) | -| `ultimate_grid.progress` | `generation_orchestrator.py` | `{session_name, progress_pct, eta_str, ...}` | ETA/progress updates. Optional `phase`: `upscaling` (deferred upscales) or `sweep` (Workflow Sweep, which adds `sweep_id, state, errors`) | +| `ultimate_grid.progress` | `generation_orchestrator.py` | `{session_name, progress_pct, eta_str, ...}` | ETA/progress updates. Optional `phase`: `upscaling` (deferred upscales), `distributed` (distributed run, adds `workers` = machines currently working) or `sweep` (Workflow Sweep, which adds `sweep_id, state, errors`) | | `ultimate_grid.next_run_seeds` | `generation_orchestrator.py` | `{node, seeds}` | Seeds the next run will use | | `ultimate_grid.metrics_progress` | `metrics_runner.py` | `{session_name, job_id, status, completed, total}` | Progress of a "compute metrics for this session" job | | `ultimate_grid.distribution_status` | `generation_orchestrator.py` | `{session_name, workers, stats}` | Distribution status | diff --git a/Roadmap.md b/Roadmap.md index f533808..08ba281 100644 --- a/Roadmap.md +++ b/Roadmap.md @@ -114,7 +114,7 @@ A button in the builder UI to pack short sha256 into config with an explanation #### **11. Model-Specific Prompts** (Needs improvement) -#### **13. Real-Time ETA** (Improved — now uses rolling 10-job window for more responsive estimates. Distribution times still TODO) +#### **13. Real-Time ETA** (Improved — now uses rolling 10-job window for more responsive estimates. Distributed runs: DONE, combined-throughput ETA over all active machines, pushed to the dashboard as `phase: "distributed"` progress events) #### **14. Cache Trigger Word Placement** (low priority) * **Problem:** `trigger_words.py` logic runs every loop iteration. @@ -184,7 +184,10 @@ README is 2000+ words; most users won't read it. Efficiency Nodes grew via commu Verify `te_batch_size` on the 5060 Ti with the real MOSS-Image, Krea 2 and Qwen-Image encoders, then decide whether the default of 1 (off) should change. #### **30. Distributed companion parity** -**Done (wave 3, `feat/dist-parity`):** `job_order`, the preview pass and image metrics now work in distributed mode (master-owned plan, `populate_from_plan`, `dist_worker` on every item). Open: re-check rejects made mid-run at claim time, claim affinity, Start At Job in distributed mode, LTX workers. +**Done (wave 3, `feat/dist-parity`):** `job_order`, the preview pass and image metrics now work in distributed mode (master-owned plan, `populate_from_plan`, `dist_worker` on every item). DONE (feat/dist-open): rejects made mid-run are re-checked at claim time, claim affinity, Start At Job in distributed mode, LTX rendered on the master (never on workers). Open: live two-machine verification of the whole set. + +#### **73. Distributed Model Sync hardening** +DONE (feat/dist-open, both repos): the master now takes over jobs still claimed after the wait timeout instead of discarding the workers' images (`wait_timeout`, default `max(600, claim_timeout)`); the companion de-duplicates concurrent downloads of one file, resolves model folders with fallbacks, normalises backslash names, refuses `..` paths and treats a VAE of `None` as the default. Open: confirm with a real slow-download run which defect the original failure was. #### **31. Scorer-pack adapters** Verify the node ids of the optional aesthetic / reward scorer packs against the real packs, then enable the shipped adapter entries. diff --git a/generation_orchestrator.py b/generation_orchestrator.py index eeb5ae0..e05d4b4 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -28,7 +28,8 @@ from .model_loader import ( from .lora_utils import expand_lora_folder from .image_generation import ( generate_image, generate_image_attn_safe, flush_batch_with_vae, flush_batch_with_remote_vae, - create_image_metadata, decode_latent_with_vae, pil_to_image_array, calculate_eta, print_generation_progress + create_image_metadata, decode_latent_with_vae, pil_to_image_array, calculate_eta, print_generation_progress, + format_eta_str ) from .config_utils import sanitize_session_name from .html_generator import get_html_template @@ -1665,64 +1666,18 @@ def run_generation_loop( except: pass + # ==== START AT JOB # (skip earlier jobs; LTX jobs too) ==== + if start_at_job > 0 and current_job < start_at_job: + skipped_count += 1 + continue + # LTX Video dispatch — must run BEFORE any image-gen-specific code that # accesses conf['sampler']/conf['steps'] (LTX configs use sampler_stage1/2). if conf.get("model_type") == "ltx_video": - from .ltx_video_generation import preflight_ltx, load_ltx_models, ltx_video_generate - - if has_regions(conf.get("positive", "")) or has_regions(conf.get("negative", "")): - print("[GridTester] ⚠️ COUPLE regional prompts are not supported for LTX video; text is encoded literally") - - # Pre-flight on first LTX config of this run (idempotent — cached check internally) - if not _ltx_preflight_done[0]: - try: - preflight_ltx(conf) - _ltx_preflight_done[0] = True - # Free image-gen caches to make room for LTX VRAM - try: - from . import model_cache as mc - if hasattr(mc, "clear_image_caches"): - mc.clear_image_caches() - except Exception: - pass - except RuntimeError as e: - print(f"[GridTester] LTX preflight FAILED: {e}") - raise - - # Override conf width/height with the current loop's resolution. - # The loop's `w, h` comes from the sampler node's resolutions widget - # (or a per-config override). LTX configs don't have width/height in - # their builder UI yet, so the cartesian-expansion defaults (446x576) - # would otherwise stick. - conf["width"] = w - conf["height"] = h - - # Filename uses gen index + timestamp suffix to guarantee uniqueness - # per run. Skip-on-exists is intentionally NOT enabled for LTX (v1) — - # the resume/match logic image-gen uses doesn't yet understand LTX - # config fields (sigmas, strengths, audio_mode, etc.) so it produced - # false-positive skips on real config changes (e.g., LoRA strength, - # seed). Until we have a proper LTX-aware manifest match, every LTX - # gen runs fresh. Users can manually delete unwanted mp4s. - _ts_suffix = str(int(time.time() * 1000) % 10_000_000) - ltx_output_filename = f"{conf_idx:06d}_seed{conf.get('seed', 0)}_{w}x{h}_dur{conf.get('duration_seconds', 5)}s_{conf.get('frame_rate', 25)}fps_{_ts_suffix}" - ltx_output_path = os.path.join(paths["images"], ltx_output_filename + ".mp4") - progress_pct = int((current_job / total_jobs) * 100) print(f"[GridTester] 📊 {current_job}/{total_jobs} ({progress_pct}%) | " - f"LTX {conf.get('duration_seconds', 5)}s @ {conf.get('frame_rate', 25)}fps | {conf.get('width')}x{conf.get('height')}") - - try: - ltx_models = load_ltx_models(conf) - gen_result = ltx_video_generate(conf, ltx_models, ltx_output_path) - except Exception as e: - print(f"[GridTester] LTX gen FAILED for config {conf_idx}: {e}") - raise - - # Build manifest entry and persist - item = _build_ltx_manifest_entry(conf, gen_result, ltx_output_filename, gen_index=conf_idx, session_name=session_name) - existing_data["items"].insert(0, item) - save_manifest(paths["manifest"], existing_data) + f"LTX {conf.get('duration_seconds', 5)}s @ {conf.get('frame_rate', 25)}fps | {w}x{h}") + _render_ltx_job(conf, w, h, conf_idx, paths, existing_data, session_name, _ltx_preflight_done) # Skip the standard image-gen path continue @@ -1742,11 +1697,6 @@ def run_generation_loop( ) actual_negative_prompt = conf["negative"] - # ==== START AT JOB # (skip earlier jobs) ==== - if start_at_job > 0 and current_job < start_at_job: - skipped_count += 1 - continue - # ==== PREVIEW PROMOTION SKIP ==== # Full run: skip cells whose preview was rejected (or not favorited). Preview run: skip cells # that already have a full image. A rejection is an explicit user signal, so overwrite_existing @@ -3171,8 +3121,219 @@ def _preencode_all_conditionings( return encoded_all +def _render_ltx_job(conf, w, h, conf_idx, paths, existing_data, session_name, preflight_state, + dist_worker=None): + """Render ONE LTX video job on this machine and record it in the manifest. + + Shared by the normal loop and the distributed master (LTX never runs on workers). ``preflight_state`` + is a one-element list used as a mutable flag so the preflight + image-cache clear run once per run. + Returns the manifest item (already inserted at index 0 and saved). Raises on failure. + """ + from .ltx_video_generation import preflight_ltx, load_ltx_models, ltx_video_generate + + if has_regions(conf.get("positive", "")) or has_regions(conf.get("negative", "")): + print("[GridTester] ⚠️ COUPLE regional prompts are not supported for LTX video; text is encoded literally") + + # Pre-flight on first LTX config of this run (idempotent — cached check internally) + if not preflight_state[0]: + try: + preflight_ltx(conf) + preflight_state[0] = True + # Free image-gen caches to make room for LTX VRAM + try: + from . import model_cache as mc + if hasattr(mc, "clear_image_caches"): + mc.clear_image_caches() + except Exception: + pass + except RuntimeError as e: + print(f"[GridTester] LTX preflight FAILED: {e}") + raise + + # Override conf width/height with the current loop's resolution. + # The loop's `w, h` comes from the sampler node's resolutions widget + # (or a per-config override). LTX configs don't have width/height in + # their builder UI yet, so the cartesian-expansion defaults (446x576) + # would otherwise stick. + conf["width"] = w + conf["height"] = h + + # Filename uses gen index + timestamp suffix to guarantee uniqueness + # per run. Skip-on-exists is intentionally NOT enabled for LTX (v1) — + # the resume/match logic image-gen uses doesn't yet understand LTX + # config fields (sigmas, strengths, audio_mode, etc.) so it produced + # false-positive skips on real config changes (e.g., LoRA strength, + # seed). Until we have a proper LTX-aware manifest match, every LTX + # gen runs fresh. Users can manually delete unwanted mp4s. + _ts_suffix = str(int(time.time() * 1000) % 10_000_000) + ltx_output_filename = f"{conf_idx:06d}_seed{conf.get('seed', 0)}_{w}x{h}_dur{conf.get('duration_seconds', 5)}s_{conf.get('frame_rate', 25)}fps_{_ts_suffix}" + ltx_output_path = os.path.join(paths["images"], ltx_output_filename + ".mp4") + + try: + ltx_models = load_ltx_models(conf) + gen_result = ltx_video_generate(conf, ltx_models, ltx_output_path) + except Exception as e: + print(f"[GridTester] LTX gen FAILED for config {conf_idx}: {e}") + raise + + # Build manifest entry and persist + item = _build_ltx_manifest_entry(conf, gen_result, ltx_output_filename, gen_index=conf_idx, session_name=session_name) + if dist_worker: + item["dist_worker"] = dist_worker + existing_data["items"].insert(0, item) + save_manifest(paths["manifest"], existing_data) + return item + + +def _dist_wait_timeout(distribution_config): + """Seconds the master waits on remote workers before taking their claimed jobs over.""" + claim_timeout = distribution_config.get("claim_timeout", 600) + return int(distribution_config.get("wait_timeout", max(600, claim_timeout))) + + +def _force_reclaim_remote(manager, JobState): + """Move every job claimed by a remote worker back to pending (the master renders them itself). + Returns how many jobs were moved.""" + n = 0 + with manager._lock: + for job in manager._jobs.values(): + if job.state == JobState.CLAIMED and job.worker_id != "master": + job.state = JobState.PENDING + job.worker_id = None + job.claimed_at = None + manager._pending_queue.append(job.job_id) + n += 1 + return n + + +def _make_reject_recheck(manifest_path, existing_data, *, skip_rejected, only_favorited, + min_interval=2.0, clock=time.monotonic): + """Build ``check(extra, config) -> reason | None`` for the manager's claim filter. + + A preview rejected in the dashboard while a distributed full pass is running must still stop its full + render. The manifest on disk is the source of truth (save_changes writes it); it is re-read at most + once per ``min_interval`` seconds and only when (mtime_ns, size) changed. Preview jobs and jobs without + a stamped fingerprint are never dropped. A read/parse failure keeps the previous index and retries on + the next call. + + Lock order: this runs with the manager lock held and takes manifest_utils._manifest_lock, so nothing + may take the manager lock while holding the manifest lock. + """ + from .config_utils import build_preview_index, should_skip_for_previews + from . import manifest_utils + + def _tombstones(meta): + return list((meta or {}).get("rejected_preview_fps", []) or []) + + state = { + "key": None, "next": None, + "pindex": build_preview_index(existing_data.get("items", []), + _tombstones(existing_data.get("meta"))), + } + + def _refresh(): + now = clock() + if state["next"] is not None and now < state["next"]: + return + state["next"] = now + min_interval + try: + st = os.stat(manifest_path) + except OSError: + return + key = (st.st_mtime_ns, st.st_size) + if key == state["key"]: + return + try: + with manifest_utils._manifest_lock: + with open(manifest_path, "r", encoding="utf-8") as f: + disk = json.load(f) + if not isinstance(disk, dict): + return + tombs = set(_tombstones(disk.get("meta"))) | set(_tombstones(existing_data.get("meta"))) + state["pindex"] = build_preview_index(disk.get("items", []) or [], tombs) + state["key"] = key + except Exception: + return # keep the old index; retry on the next call + + def check(extra, config): + extra = extra or {} + if extra.get("render"): + return None + fp = (extra.get("stamp") or {}).get("config_fp") + if not fp: + return None + _refresh() + return should_skip_for_previews(fp, state["pindex"], skip_rejected, only_favorited) + + return check + + +def _install_claim_filter(manager, paths, existing_data, *, preview_on, has_optional_inputs, + skip_rejected, only_favorited): + """Install the reject re-check on the manager when it applies (full pass, no optional inputs, a + skip rule is on) and the companion supports it. Returns True when installed.""" + if preview_on or has_optional_inputs or not (skip_rejected or only_favorited): + return False + if not hasattr(manager, "set_claim_filter"): + print("[Distribution] ℹ️ Rejects made during the run are not re-checked at claim time: update " + "ComfyUI-USCG-Distributed to enable it") + return False + manager.set_claim_filter(_make_reject_recheck( + paths["manifest"], existing_data, skip_rejected=skip_rejected, only_favorited=only_favorited)) + return True + + +def _distributed_progress_payload(status, eta, base, unique_id, session_name): + """The ``ultimate_grid.progress`` payload for a distributed run (pure).""" + done = (base["pre_skipped"] + base["ltx_done"] + status.get("completed", 0) + + status.get("skipped", 0) + status.get("failed", 0)) + total = base["plan_total"] + pct = min(100, int(done * 100 / total)) if total else 0 + secs = (eta or {}).get("eta_seconds") + if secs is None: + eta_str, finish = "—", "" + else: + eta_str = format_eta_str(secs) + finish = time.strftime("%H:%M:%S", time.localtime(time.time() + secs)) + avg = (eta or {}).get("avg_duration") + return { + "node": unique_id, + "session_name": session_name, + "current_job": done, + "total_jobs": total, + "progress_pct": pct, + "eta_str": eta_str, + "finish_time": finish, + "avg_duration": round(avg, 1) if avg else 0, + "last_duration": 0, + "phase": "distributed", + "workers": (eta or {}).get("active_workers") or 0, + } + + +def _emit_distributed_progress(manager, base, unique_id, session_name, status=None): + """Send the distributed progress event to the dashboard; never raises.""" + try: + if PromptServer is None: + return + if status is None: + status = manager.get_status() + eta = None + if hasattr(manager, "get_eta"): + try: + eta = manager.get_eta() + except Exception: + eta = None + PromptServer.instance.send_sync( + "ultimate_grid.progress", + _distributed_progress_payload(status, eta, base, unique_id, session_name)) + except Exception: + pass + + def _send_distribution_status(manager): """Send distribution status update to the dashboard.""" + status = None try: status = manager.get_status() status["active"] = manager.is_active @@ -3184,10 +3345,17 @@ def _send_distribution_status(manager): }) except Exception: pass + cb = getattr(manager, "on_status", None) + if cb is not None and status is not None: + try: + cb(status) + except Exception: + pass def _distributed_plan_entries(job_plan, existing_data, *, overwrite_existing, has_optional_inputs, - lora_triggerwords_mode, preview=None, skip_rejected=True, only_favorited=False): + lora_triggerwords_mode, preview=None, skip_rejected=True, only_favorited=False, + start_at_job=0): """Turn the (already ordered) job plan into the entries the distributed manager queues. The master decides everything stateful here, with the same rules and order as the normal loop @@ -3197,7 +3365,9 @@ def _distributed_plan_entries(job_plan, existing_data, *, overwrite_existing, ha ``preview`` is None (full pass) or ``{"steps", "scale", "cfg"}`` (cfg may be None = keep). Returns ``(entries, counts)``; an entry is ``{"config", "input_job", "extra"}`` where ``extra`` holds ``stamp`` (manifest fields the master adds on collection), optional ``render`` (``preview_scale``) and - optional ``seed`` (survivor seed to re-render). ``counts`` = {"exists", "promotion", "full_exists"}. + optional ``seed`` (survivor seed to re-render); every entry also carries ``conf_idx``. ``counts`` = + {"exists", "promotion", "full_exists"} plus "start_at" when ``start_at_job`` > 1. ``start_at_job`` N > 1 skips plan positions 1..N-1 + (1-based, like the normal loop) before any other check; LTX entries are never skipped as "existing". """ from .config_utils import ( config_fingerprint, build_preview_index, build_full_fp_set, should_skip_for_previews @@ -3208,8 +3378,13 @@ def _distributed_plan_entries(job_plan, existing_data, *, overwrite_existing, ha full_fps = build_full_fp_set(items) if preview else set() entries = [] counts = {"exists": 0, "promotion": 0, "full_exists": 0} + if start_at_job > 1: + counts["start_at"] = 0 # only present when Start At Job is in effect - for pj in job_plan: + for pos, pj in enumerate(job_plan, 1): + if start_at_job > 1 and pos < start_at_job: + counts["start_at"] += 1 + continue conf, w, h, batch_idx = pj["conf"], pj["w"], pj["h"], pj["batch_idx"] is_ltx = conf.get("model_type") == "ltx_video" is_pv = bool(preview) and not is_ltx @@ -3241,7 +3416,7 @@ def _distributed_plan_entries(job_plan, existing_data, *, overwrite_existing, ha # A randomized seed never matches an old item (the normal loop rolls a fresh one before the check), # except a survivor's approved seed. - if not overwrite_existing and (not randomize or survivor_seed is not None): + if not is_ltx and not overwrite_existing and (not randomize or survivor_seed is not None): pos_, _ = build_prompt_with_triggers(conf, lora_triggerwords_mode) if check_if_job_completed( items, eff_conf, survivor_seed if survivor_seed is not None else conf["seed"], @@ -3270,6 +3445,7 @@ def _distributed_plan_entries(job_plan, existing_data, *, overwrite_existing, ha "input_job": {"label": pj["job"].get("label", f"{w}x{h}"), "width": w, "height": h, "batch_idx": batch_idx}, "extra": extra, + "conf_idx": pj["conf_idx"], }) return entries, counts @@ -3319,39 +3495,62 @@ def _distributed_metrics_pass(existing_data, pre_ids, images_dir, image_path_for return out +LTX_LEGACY_COMPANION_ERROR = ( + "LTX video configs in a distributed run are rendered on the master and need an updated " + "ComfyUI-USCG-Distributed (populate_from_plan). Update the companion, or turn Distribution off for this run." +) + + def _populate_distributed_manager(manager, job_plan, expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode, *, job_order="as_written", preview_on=False, pv_steps=8, pv_scale=0.5, pv_cfg=None, - skip_rejected=True, only_favorited=False): + skip_rejected=True, only_favorited=False, start_at_job=0): """Queue the run on the distributed manager: the master's own plan when the companion supports it - (populate_from_plan), else the companion's legacy nested loop (job order / preview ignored).""" + (populate_from_plan), else the companion's legacy nested loop (job order / preview ignored). + + Returns ``{"counts", "ltx_entries", "queued", "plan_total"}``. LTX entries are NOT queued: the master + renders them itself. With the legacy companion ``counts`` is {} and ``plan_total``/``queued`` are None. + """ if job_plan is not None and hasattr(manager, "populate_from_plan"): # The master decides order, skips, preview promotion and fingerprints; workers just render. _pv = {"steps": pv_steps, "scale": pv_scale, "cfg": pv_cfg} if preview_on else None entries, _counts = _distributed_plan_entries( job_plan, existing_data, overwrite_existing=overwrite_existing, has_optional_inputs=has_optional_inputs, lora_triggerwords_mode=lora_triggerwords_mode, - preview=_pv, skip_rejected=skip_rejected, only_favorited=only_favorited + preview=_pv, skip_rejected=skip_rejected, only_favorited=only_favorited, + start_at_job=start_at_job ) - print(f"[Distribution] 📋 Plan: {len(entries)} job(s) to render | skipped: {_counts['exists']} already exist, " + ltx_entries = [e for e in entries if e["config"].get("model_type") == "ltx_video"] + queued = [e for e in entries if e["config"].get("model_type") != "ltx_video"] + print(f"[Distribution] 📋 Plan: {len(queued)} job(s) to render | skipped: {_counts['exists']} already exist, " f"{_counts['promotion']} rejected/unfavorited previews, {_counts['full_exists']} already full") + if start_at_job > 1: + print(f"[Distribution] ⏭️ Start At Job #{start_at_job}: {start_at_job - 1} job(s) skipped") + if ltx_entries: + print(f"[Distribution] 🎬 {len(ltx_entries)} LTX job(s) render on the master") if job_order != "as_written": print(f"[Distribution] 🧭 Job order {job_order}: jobs are claimed in the master's plan order") if preview_on: print(f"[GridTester] 🔎 Distributed preview pass: {pv_steps} steps, x{pv_scale} scale") _spec = image_metrics.current_spec() manager.populate_from_plan( - entries, run_extras={"metrics_builtins": list(_spec.get("builtins") or [])} if _spec else {}) + queued, run_extras={"metrics_builtins": list(_spec.get("builtins") or [])} if _spec else {}) + return {"counts": _counts, "ltx_entries": ltx_entries, "queued": queued, "plan_total": len(job_plan)} else: + if any(c.get("model_type") == "ltx_video" for c in (expanded or [])): + raise RuntimeError(LTX_LEGACY_COMPANION_ERROR) if (job_order != "as_written" or preview_on or any(it.get("preview") is True for it in existing_data.get("items", []))): print(f"[GridTester] ⚠️ job_order / preview pass / preview promotion need an updated ComfyUI-USCG-Distributed " f"companion - ignored in distributed mode") + if start_at_job > 1: + print(f"[GridTester] ⚠️ Start At Job needs an updated ComfyUI-USCG-Distributed companion - " + f"ignored in distributed mode") manager.populate_jobs( expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode ) - + return {"counts": {}, "ltx_entries": [], "queued": None, "plan_total": None} def _standard_run_meta(existing_data, positive_text, negative_text, ckpt_name, seed, vae_batch_size, @@ -3449,6 +3648,7 @@ def _run_distributed_generation( worker_urls = distribution_config.get("worker_urls", []) claim_timeout = distribution_config.get("claim_timeout", 600) + dist_wait_timeout = _dist_wait_timeout(distribution_config) # Inject trigger word mode into configs so workers can use it for conf in expanded: @@ -3487,12 +3687,32 @@ def _run_distributed_generation( pv_scale, pv_cfg, meta_run) _pre_ids = {it.get("id") for it in existing_data.get("items", [])} - _populate_distributed_manager( + pop = _populate_distributed_manager( manager, job_plan, expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode, job_order=job_order, preview_on=preview_on, pv_steps=pv_steps, pv_scale=pv_scale, pv_cfg=pv_cfg, skip_rejected=skip_rejected_previews, - only_favorited=only_favorited_previews + only_favorited=only_favorited_previews, + start_at_job=int((session_settings or {}).get("start_at_job", 0) or 0) ) + if not isinstance(pop, dict): + pop = {"counts": {}, "ltx_entries": [], "queued": None, "plan_total": None} + _ltx_entries = pop.get("ltx_entries") or [] + _pop_total = manager.get_status()["total"] + # Progress accounting for the dashboard: everything the plan resolved without a worker render + # (skipped / start-at / LTX done on the master) plus what the manager reports. + prog_base = { + "pre_skipped": sum((pop.get("counts") or {}).values()), + "ltx_done": 0, + "plan_total": pop["plan_total"] if pop.get("plan_total") is not None else _pop_total, + } + try: + manager.on_status = lambda st: _emit_distributed_progress( + manager, prog_base, unique_id, session_name, st) + except Exception: + pass + _install_claim_filter( + manager, paths, existing_data, preview_on=preview_on, has_optional_inputs=has_optional_inputs, + skip_rejected=skip_rejected_previews, only_favorited=only_favorited_previews) # A legacy companion (no populate_from_plan) ignores the preview pass, so its workers render FULL # images: the master must still upscale them, with the settings from before the preview override. @@ -3515,7 +3735,8 @@ def _run_distributed_generation( print(f"[Distribution] 🧠 MASTER PRE-ENCODING: Encoding all unique prompts for workers...") encoded_conditionings = _preencode_all_conditionings( - expanded, lora_triggerwords_mode, ckpt_name, use_remote_vae, + [c for c in expanded if c.get("model_type") != "ltx_video"], + lora_triggerwords_mode, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, model_cache ) @@ -3530,7 +3751,7 @@ def _run_distributed_generation( # Early exit if all jobs were already completed (resume with nothing to do) initial_status = manager.get_status() - if initial_status["total"] == 0: + if initial_status["total"] == 0 and not _ltx_entries: print(f"[Distribution] ⏭️ All jobs already completed, nothing to distribute") manager.deactivate() distribution.clear_active_manager() @@ -3541,19 +3762,25 @@ def _run_distributed_generation( html = get_html_template(session_name, existing_data, unique_id) return (html,) + # First progress event: the plan's skipped jobs count as done, ETA unknown yet + _emit_distributed_progress(manager, prog_base, unique_id, session_name, initial_status) + # === Notify remote workers === - master_url = distribution.get_master_url() - print(f"[Distribution] 🌐 Master URL: {master_url}") - print(f"[Distribution] 🌐 Notifying {len(worker_urls)} worker(s)...") + if initial_status["total"] == 0: + print(f"[Distribution] 🎬 Only LTX job(s) in this run: they render on the master, workers are not needed") + else: + master_url = distribution.get_master_url() + print(f"[Distribution] 🌐 Master URL: {master_url}") + print(f"[Distribution] 🌐 Notifying {len(worker_urls)} worker(s)...") - sync_models = distribution_config.get("sync_models_to_workers", False) - if sync_models: - manager.sync_models_to_workers = True - print(f"[Distribution] ☁️ Model sync enabled — workers will download missing models from master") + sync_models = distribution_config.get("sync_models_to_workers", False) + if sync_models: + manager.sync_models_to_workers = True + print(f"[Distribution] ☁️ Model sync enabled — workers will download missing models from master") - worker_results = distribution.notify_workers_to_start(worker_urls, master_url, session_name, sync_models_to_workers=sync_models) - successful_workers = sum(1 for _, ok, _ in worker_results if ok) - print(f"[Distribution] ✅ {successful_workers}/{len(worker_urls)} workers started") + worker_results = distribution.notify_workers_to_start(worker_urls, master_url, session_name, sync_models_to_workers=sync_models) + successful_workers = sum(1 for _, ok, _ in worker_results if ok) + print(f"[Distribution] ✅ {successful_workers}/{len(worker_urls)} workers started") # === Master local processing state === loaded_model, loaded_clip, loaded_vae = None, None, None @@ -3608,7 +3835,42 @@ def _run_distributed_generation( # Master claims jobs from its own distribution manager and processes them # using the same model loading, encoding, and generation code paths as the # normal (non-distributed) generation loop. + ltx_failed = 0 + _legacy_ltx_pf = [False] + _idle_since = None + _idle_passes = 0 try: + # --- LTX phase: video jobs always render on the master (workers cannot load them). Workers + # were already notified, so they render the image jobs in parallel. + if _ltx_entries: + _ltx_pf = [False] + for _le in _ltx_entries: + try: + import comfy.model_management as mm + if mm.processing_interrupted(): + print(f"\n[Distribution/Master] 🛑 INTERRUPTED") + raise InterruptProcessingException() + except InterruptProcessingException: + raise + except: + pass + _lc, _lj = _le["config"], _le["input_job"] + print(f"[Distribution/Master] 🎬 LTX job {prog_base['ltx_done'] + 1}/{len(_ltx_entries)} " + f"| {_lj['width']}x{_lj['height']}") + try: + _render_ltx_job(_lc, _lj["width"], _lj["height"], _le.get("conf_idx", 0), paths, + existing_data, session_name, _ltx_pf, dist_worker="master") + except InterruptProcessingException: + raise + except Exception as e: + ltx_failed += 1 + print(f"[Distribution/Master] ❌ LTX job failed: {e}") + prog_base["ltx_done"] += 1 + _emit_distributed_progress(manager, prog_base, unique_id, session_name) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + while True: # Check for interrupt try: @@ -3624,7 +3886,44 @@ def _run_distributed_generation( # Claim next job from manager batch = manager.claim_batch_for_local(1) if not batch: - break # No more pending jobs for master + # Nothing for the master right now: finish its own batch first, then stay and wait for + # the workers (their results are collected by the routes). After the wait timeout the + # master takes the still-claimed jobs over instead of throwing the workers' images away. + if pending_batch: + _flush_pending_batch( + pending_batch, current_vae_is_remote, current_remote_vae_url, + per_config_remote_workers, use_remote_vae, remote_vae_worker, + loaded_vae, paths, existing_data, session_name, + paths["manifest"], unique_id, + config_overrides_vae=config_overrides_vae, image_format=image_format + ) + for jid in pending_job_ids: + manager.complete_job(jid) + pending_batch = [] + pending_job_ids = [] + _send_distribution_status(manager) + if not manager.has_pending_or_claimed: + break + if _idle_since is None: + _idle_since = time.time() + _idle_passes = 0 + manager.release_timed_out_jobs() + if time.time() - _idle_since > dist_wait_timeout: + n_rec = _force_reclaim_remote(manager, JobState) + print(f"[Distribution] ⏰ Wait timeout ({dist_wait_timeout}s): rendering {n_rec} reclaimed " + f"job(s) on the master") + _idle_since = None + continue + time.sleep(2) + _send_distribution_status(manager) + _idle_passes += 1 + if _idle_passes % 5 == 0: + _st = manager.get_status() + print(f"[Distribution] ⏳ Waiting for workers - Pending: {_st['pending']}, " + f"Claimed: {_st['claimed']}, Completed: {_st['completed']}/{_st['total']}") + continue + _idle_since = None + _job_t0 = time.time() claimed_job = batch[0] conf = claimed_job["config"] @@ -3643,6 +3942,19 @@ def _run_distributed_generation( import random current_seed = random.randint(0, 2**63 - 1) + if conf.get("model_type") == "ltx_video": + # Only reachable with a queue built outside the plan split: render it here. + try: + _render_ltx_job(conf, w, h, claimed_job.get("gen_index", master_jobs_processed), paths, + existing_data, session_name, _legacy_ltx_pf, dist_worker="master") + manager.complete_job(job_id) + except InterruptProcessingException: + raise + except Exception as e: + print(f"[Distribution/Master] ❌ LTX job failed: {e}") + manager.fail_job(job_id, str(e)) + continue + master_jobs_processed += 1 completed_total = manager.total_completed print(f"[Distribution/Master] 📊 Processing job {job_id} | " @@ -3818,7 +4130,25 @@ def _run_distributed_generation( ) job_durations.append(duration) + if hasattr(manager, "record_duration"): + try: + manager.record_duration("master", time.time() - _job_t0) + except Exception: + pass eta_info = calculate_eta(job_durations, completed_total + 1, dist_total) + if hasattr(manager, "get_eta"): + try: + _ge = manager.get_eta() + except Exception: + _ge = None + if _ge and _ge.get("eta_seconds") is not None: + _es = int(_ge["eta_seconds"]) + eta_info = { + "hours": _es // 3600, "minutes": (_es % 3600) // 60, "seconds": _es % 60, + "finish_time": time.time() + _es, + "finish_formatted": time.strftime("%H:%M:%S", time.localtime(time.time() + _es)), + "avg_duration": _ge.get("avg_duration") or duration, + } if eta_info: print_generation_progress( completed_total + 1, dist_total, conf, w, h, duration, eta_info @@ -3926,7 +4256,6 @@ def _run_distributed_generation( # === Wait for remote workers to finish === # Max wait timeout prevents infinite loops if workers crash without cleanup. # After timeout, remaining claimed jobs are reclaimed and processed on master. - dist_wait_timeout = distribution_config.get("wait_timeout", 600) # Default 10 minutes if manager.has_pending_or_claimed: print(f"[Distribution] ⏳ Waiting for remote workers to finish (timeout: {dist_wait_timeout}s)...") wait_start = time.time() @@ -3941,13 +4270,7 @@ def _run_distributed_generation( remaining = status['pending'] + status['claimed'] print(f"[Distribution] ⏰ Wait timeout ({dist_wait_timeout}s) reached with {remaining} jobs remaining — reclaiming for master") # Force-reclaim all remaining claimed jobs - with manager._lock: - for job in manager._jobs.values(): - if job.state == JobState.CLAIMED and job.worker_id != "master": - job.state = JobState.PENDING - job.worker_id = None - job.claimed_at = None - manager._pending_queue.append(job.job_id) + _force_reclaim_remote(manager, JobState) break # Check for interrupt while waiting @@ -4102,6 +4425,13 @@ def _run_distributed_generation( print(f"[Distribution] ⏱️ {total_hours}h {total_minutes}m {total_seconds}s total") if final_status.get('failed', 0) > 0: print(f"[Distribution] ❌ {final_status['failed']} jobs failed permanently") + if final_status.get('skipped', 0) > 0: + print(f"[Distribution] ⏭️ {final_status['skipped']} dropped at claim (preview rejected during the run)") + if pop.get("counts", {}).get("start_at", 0) > 0: + print(f"[Distribution] ⏭️ {pop['counts']['start_at']} skipped by Start At Job") + if _ltx_entries: + print(f"[Distribution] 🎬 LTX: {len(_ltx_entries) - ltx_failed}/{len(_ltx_entries)} rendered on the master" + + (f", {ltx_failed} failed" if ltx_failed else "")) # Per-worker breakdown workers = final_status.get("workers", {}) @@ -4126,7 +4456,7 @@ def _run_distributed_generation( if PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.progress", _completion_payload( - unique_id, session_name, dist_total, avg_per_job, + unique_id, session_name, (prog_base.get("plan_total") or dist_total), avg_per_job, f"{total_hours}h {total_minutes}m {total_seconds}s", final_status.get("completed", total_generated), next_phase)) except Exception: diff --git a/image_generation.py b/image_generation.py index 5f9d193..5ee0374 100644 --- a/image_generation.py +++ b/image_generation.py @@ -1055,6 +1055,17 @@ def calculate_eta(job_durations, current_job, total_jobs): } +def format_eta_str(seconds): + """Human ETA text: "1h 2m", "2m 5s" or "9s" (same rules as the console progress lines).""" + seconds = max(0, int(seconds)) + hours, minutes, secs = seconds // 3600, (seconds % 3600) // 60, seconds % 60 + if hours > 0: + return f"{hours}h {minutes}m" + if minutes > 0: + return f"{minutes}m {secs}s" + return f"{secs}s" + + def print_generation_progress(current_job, total_jobs, config, width, height, duration, eta_info): """ Print progress information for current generation. diff --git a/resources/logic_events.js b/resources/logic_events.js index d321bcc..ca2d978 100644 --- a/resources/logic_events.js +++ b/resources/logic_events.js @@ -56,7 +56,8 @@ function formatProgressText(p) { } const phaseLabel = p.phase === 'upscaling' ? 'Upscaling' : (p.phase === 'sweep' ? 'Sweep job' : 'Job'); const finishPart = p.finish_time ? ` | ~${p.finish_time}` : ''; - return `${phaseLabel} ${p.current_job}/${p.total_jobs} (${p.progress_pct}%) | ETA: ${p.eta_str}${finishPart} | ${p.avg_duration}s/job${p.errors ? ` | ${p.errors} error${p.errors === 1 ? '' : 's'}` : ''}${(p.state === 'paused' || p.state === 'cancelled') ? ` | ${p.state}` : ''}`; + const machinesPart = (p.phase === 'distributed' && p.workers) ? ` | ${p.workers} machine(s)` : ''; + return `${phaseLabel} ${p.current_job}/${p.total_jobs} (${p.progress_pct}%) | ETA: ${p.eta_str}${finishPart} | ${p.avg_duration}s/job${p.errors ? ` | ${p.errors} error${p.errors === 1 ? '' : 's'}` : ''}${(p.state === 'paused' || p.state === 'cancelled') ? ` | ${p.state}` : ''}${machinesPart}`; } // WebSocket Listener - OPTIMIZED diff --git a/tests/conftest.py b/tests/conftest.py index f3706dd..6f86d63 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ We stub those out here so pytest can collect and run unit tests without ComfyUI. import sys import types import os +import threading # Try to import real torch first — some tests (crop/paste, integration) need it. # If torch isn't installed in the active Python, fall through to the stub loop below. @@ -106,6 +107,7 @@ _GEN_ORCH_DEPS = { "encode_pair_cached": lambda *a, **kw: (None, None), }, "manifest_utils": { + "_manifest_lock": threading.Lock(), "load_existing_manifest": lambda *a, **kw: {}, "save_manifest": lambda *a, **kw: None, "preserve_disk_comparisons": lambda *a, **kw: None, @@ -138,6 +140,7 @@ _GEN_ORCH_DEPS = { "pil_to_image_array": lambda *a, **kw: None, "calculate_eta": lambda *a, **kw: 0.0, "print_generation_progress": lambda *a, **kw: None, + "format_eta_str": lambda seconds: f"{int(seconds)}s", }, "html_generator": { "get_html_template": lambda *a, **kw: "", diff --git a/tests/js/progress_text.test.js b/tests/js/progress_text.test.js new file mode 100644 index 0000000..fc61992 --- /dev/null +++ b/tests/js/progress_text.test.js @@ -0,0 +1,47 @@ +// Run: node --test tests/js/progress_text.test.js (ETA bar label for progress events, no browser needed) +const fs = require('fs'); +const path = require('path'); +const test = require('node:test'); +const assert = require('assert'); + +const src = fs.readFileSync(path.join(__dirname, '..', '..', 'resources', 'logic_events.js'), 'utf8'); +function grab(name) { + const i = src.indexOf('function ' + name + '('); + assert.ok(i >= 0, name + ' missing'); + let depth = 0, j = src.indexOf('{', i); + for (let k = j; k < src.length; k++) { + if (src[k] === '{') depth++; + if (src[k] === '}') { depth--; if (depth === 0) return src.slice(i, k + 1); } + } +} +const formatProgressText = new Function(grab('formatProgressText') + '\nreturn formatProgressText;')(); + +const base = { current_job: 3, total_jobs: 6, progress_pct: 50, eta_str: '2m 5s', finish_time: '12:00:00', avg_duration: 10.5 }; + +test('distributed label shows the machine count', () => { + const t = formatProgressText({ ...base, phase: 'distributed', workers: 2 }); + assert.strictEqual(t, 'Job 3/6 (50%) | ETA: 2m 5s | ~12:00:00 | 10.5s/job | 2 machine(s)'); +}); + +test('distributed without workers omits the machine part', () => { + const t = formatProgressText({ ...base, phase: 'distributed', workers: 0 }); + assert.strictEqual(t, 'Job 3/6 (50%) | ETA: 2m 5s | ~12:00:00 | 10.5s/job'); +}); + +test('upscaling label is unchanged', () => { + const t = formatProgressText({ ...base, phase: 'upscaling', workers: 3 }); + assert.strictEqual(t, 'Upscaling 3/6 (50%) | ETA: 2m 5s | ~12:00:00 | 10.5s/job'); +}); + +test('sweep label keeps errors and paused state', () => { + assert.strictEqual( + formatProgressText({ ...base, phase: 'sweep', errors: 1, state: 'paused', finish_time: '' }), + 'Sweep job 3/6 (50%) | ETA: 2m 5s | 10.5s/job | 1 error | paused'); + assert.strictEqual( + formatProgressText({ ...base, phase: 'sweep', errors: 2, state: 'cancelled' }), + 'Sweep job 3/6 (50%) | ETA: 2m 5s | ~12:00:00 | 10.5s/job | 2 errors | cancelled'); +}); + +test('plain run (no phase) uses the Job label', () => { + assert.strictEqual(formatProgressText(base), 'Job 3/6 (50%) | ETA: 2m 5s | ~12:00:00 | 10.5s/job'); +}); diff --git a/tests/test_dist_idle_wait.py b/tests/test_dist_idle_wait.py new file mode 100644 index 0000000..ea1f025 --- /dev/null +++ b/tests/test_dist_idle_wait.py @@ -0,0 +1,76 @@ +"""Master waits for remote workers inside the claim loop; reclaim + timeout helpers.""" +import os +import re +import sys +import threading +from collections import OrderedDict + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import generation_orchestrator as go # preloaded by tests/conftest.py + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +class _JobState: + PENDING = "pending" + CLAIMED = "claimed" + COMPLETED = "completed" + + +class _Job: + def __init__(self, job_id, state, worker_id=None, claimed_at=None): + self.job_id, self.state, self.worker_id, self.claimed_at = job_id, state, worker_id, claimed_at + + +class _Mgr: + def __init__(self, jobs): + self._lock = threading.Lock() + self._jobs = OrderedDict((j.job_id, j) for j in jobs) + self._pending_queue = [j.job_id for j in jobs if j.state == _JobState.PENDING] + + +def test_force_reclaim_moves_only_remote_claimed_jobs(): + jobs = [ + _Job("a", _JobState.CLAIMED, "worker-1", 1.0), + _Job("b", _JobState.CLAIMED, "master", 2.0), + _Job("c", _JobState.PENDING), + _Job("d", _JobState.COMPLETED, "worker-1"), + _Job("e", _JobState.CLAIMED, "worker-2", 3.0), + ] + mgr = _Mgr(jobs) + assert go._force_reclaim_remote(mgr, _JobState) == 2 + assert mgr._jobs["a"].state == _JobState.PENDING and mgr._jobs["a"].worker_id is None + assert mgr._jobs["a"].claimed_at is None + assert mgr._jobs["e"].state == _JobState.PENDING + assert mgr._jobs["b"].state == _JobState.CLAIMED and mgr._jobs["b"].worker_id == "master" + assert mgr._jobs["d"].state == _JobState.COMPLETED + assert mgr._pending_queue == ["c", "a", "e"] + + +def test_force_reclaim_with_nothing_to_do(): + assert go._force_reclaim_remote(_Mgr([_Job("a", _JobState.PENDING)]), _JobState) == 0 + + +def test_wait_timeout_default_is_max_of_600_and_claim_timeout(): + assert go._dist_wait_timeout({}) == 600 + assert go._dist_wait_timeout({"claim_timeout": 300}) == 600 + assert go._dist_wait_timeout({"claim_timeout": 1200}) == 1200 + assert go._dist_wait_timeout({"wait_timeout": 30}) == 30 + assert go._dist_wait_timeout({"wait_timeout": "45", "claim_timeout": 1200}) == 45 + + +def test_empty_batch_branch_flushes_waits_and_reclaims(): + with open(os.path.join(ROOT, "generation_orchestrator.py"), encoding="utf-8") as f: + src = f.read() + body = re.search(r"\ndef _run_distributed_generation\(.*?(?=\ndef |\Z)", src, re.S).group(0) + start = body.index(" if not batch:") + end = body.index(" _idle_since = None\n _job_t0", start) + branch = body[start:end] + assert "_flush_pending_batch(" in branch + assert branch.index("_flush_pending_batch(") < branch.index("has_pending_or_claimed") + assert branch.index("has_pending_or_claimed") < branch.index("break") + assert branch.index("break") < branch.index("_force_reclaim_remote(") + assert "release_timed_out_jobs()" in branch + assert "time.sleep(2)" in branch + assert "continue" in branch diff --git a/tests/test_dist_ltx_master.py b/tests/test_dist_ltx_master.py new file mode 100644 index 0000000..79dba34 --- /dev/null +++ b/tests/test_dist_ltx_master.py @@ -0,0 +1,118 @@ +"""LTX jobs in distributed mode: the master renders them, they never reach workers.""" +import os +import sys +import types + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import generation_orchestrator as go # preloaded by tests/conftest.py +from config_utils import build_job_plan + + +@pytest.fixture(autouse=True) +def _prompts(monkeypatch): + monkeypatch.setattr(go, "build_prompt_with_triggers", lambda conf, mode: (conf["positive"], "")) + + +def _conf(**kw): + c = {"model": "m", "lora": "None", "lora_expanded": "None", "sampler": "euler", "scheduler": "normal", + "steps": 30, "cfg": 7.0, "denoise": 1.0, "seed": 5, "positive": "p", "negative": "n"} + c.update(kw) + return c + + +def _ltx(**kw): + c = {"model_type": "ltx_video", "model": "ltx", "seed": 9, "positive": "video", "negative": "", + "duration_seconds": 5, "frame_rate": 25} + c.update(kw) + return c + + +def _job(w=1024, h=1024): + return {"label": f"{w}x{h}", "width": w, "height": h, "batch_idx": 0} + + +class _OldManager: + def populate_jobs(self, *a): + raise AssertionError("must not queue anything") + + +class _NewManager: + def __init__(self): + self.entries = None + + def populate_from_plan(self, entries, run_extras=None): + self.entries = entries + + +def _populate(mgr, confs, items=()): + plan = build_job_plan(confs, [_job()]) + return go._populate_distributed_manager( + mgr, plan, confs, [_job()], {"items": list(items), "meta": {}}, False, False, "off") + + +def test_ltx_split_out_and_never_queued(): + mgr = _NewManager() + confs = [_conf(seed=1), _ltx(), _conf(seed=2)] + ret = _populate(mgr, confs) + assert [e["config"]["seed"] for e in mgr.entries] == [1, 2] + assert all(e["config"].get("model_type") != "ltx_video" for e in mgr.entries) + assert len(ret["ltx_entries"]) == 1 + assert ret["ltx_entries"][0]["conf_idx"] == 1 + assert ret["plan_total"] == 3 + assert all("conf_idx" in e for e in mgr.entries) + + +def test_ltx_not_skipped_as_existing(): + ltx = _ltx() + item = {"media_type": "video", "model_type": "ltx_video", "seed": 9, "width": 1024, "height": 1024, + "batch_idx": 0, "positive": "video", "negative": "", "model": "ltx"} + plan = build_job_plan([ltx], [_job()]) + entries, counts = go._distributed_plan_entries( + plan, {"items": [item], "meta": {}}, overwrite_existing=False, has_optional_inputs=False, + lora_triggerwords_mode="off") + assert len(entries) == 1 + assert counts["exists"] == 0 + + +def test_legacy_manager_with_ltx_raises(): + with pytest.raises(RuntimeError, match="rendered on the master"): + go._populate_distributed_manager( + _OldManager(), None, [_conf(), _ltx()], [_job()], {"items": []}, False, False, "off") + + +def test_render_ltx_job_inserts_item_and_preflights_once(monkeypatch, tmp_path): + pkg = go.__package__ + calls = {"preflight": 0, "load": 0, "gen": 0, "saved": 0} + stub = types.ModuleType(f"{pkg}.ltx_video_generation") + + def preflight(conf): + calls["preflight"] += 1 + + def load(conf): + calls["load"] += 1 + return object() + + def gen(conf, models, out_path): + calls["gen"] += 1 + return {"path": out_path} + + stub.preflight_ltx, stub.load_ltx_models, stub.ltx_video_generate = preflight, load, gen + monkeypatch.setitem(sys.modules, f"{pkg}.ltx_video_generation", stub) + monkeypatch.setattr(go, "_build_ltx_manifest_entry", + lambda conf, res, name, gen_index=None, session_name="": {"id": name, "media_type": "video"}) + monkeypatch.setattr(go, "save_manifest", lambda path, data: calls.__setitem__("saved", calls["saved"] + 1)) + paths = {"images": str(tmp_path), "manifest": str(tmp_path / "manifest.json")} + data = {"items": [{"id": "old"}]} + flag = [False] + + a = go._render_ltx_job(_ltx(), 640, 480, 3, paths, data, "sess", flag, dist_worker="master") + b = go._render_ltx_job(_ltx(), 640, 480, 4, paths, data, "sess", flag) + + assert calls["preflight"] == 1 and calls["load"] == 2 and calls["gen"] == 2 and calls["saved"] == 2 + assert a["dist_worker"] == "master" + assert "dist_worker" not in b + assert data["items"][0] is b and data["items"][1] is a + assert data["items"][2] == {"id": "old"} diff --git a/tests/test_dist_progress.py b/tests/test_dist_progress.py new file mode 100644 index 0000000..2dca67c --- /dev/null +++ b/tests/test_dist_progress.py @@ -0,0 +1,135 @@ +"""Distributed progress events: payload maths, emit helper, on_status hook, format_eta_str.""" +import ast +import os +import sys +import types + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + +import generation_orchestrator as go # preloaded by tests/conftest.py + + +def _real_format_eta_str(): + """image_generation is stubbed under test (needs torch): rebuild the real function from source.""" + with open(os.path.join(ROOT, "image_generation.py"), encoding="utf-8") as f: + tree = ast.parse(f.read()) + node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "format_eta_str") + code = compile(ast.Module(body=[node], type_ignores=[]), "format_eta_str", "exec") + fn_code = next(c for c in code.co_consts if isinstance(c, types.CodeType) and c.co_name == "format_eta_str") + return types.FunctionType(fn_code, {}) + + +@pytest.fixture +def fmt(monkeypatch): + f = _real_format_eta_str() + monkeypatch.setattr(go, "format_eta_str", f) + return f + + +def test_format_eta_str_cases(fmt): + assert fmt(9) == "9s" + assert fmt(0) == "0s" + assert fmt(125) == "2m 5s" + assert fmt(60) == "1m 0s" + assert fmt(3725) == "1h 2m" + assert fmt(7200) == "2h 0m" + assert fmt(-4) == "0s" + assert fmt(9.9) == "9s" + + +BASE = {"pre_skipped": 2, "ltx_done": 1, "plan_total": 10} +STATUS = {"completed": 3, "skipped": 1, "failed": 0, "pending": 3, "claimed": 0, "total": 7} + + +def test_payload_maths(fmt): + eta = {"eta_seconds": 125, "avg_duration": 10.26, "active_workers": 2} + p = go._distributed_progress_payload(STATUS, eta, BASE, "n1", "sess") + assert p["current_job"] == 2 + 1 + 3 + 1 + 0 == 7 + assert p["total_jobs"] == 10 + assert p["progress_pct"] == 70 + assert p["phase"] == "distributed" + assert p["workers"] == 2 + assert p["eta_str"] == "2m 5s" + assert p["avg_duration"] == 10.3 + assert p["node"] == "n1" and p["session_name"] == "sess" + assert p["last_duration"] == 0 + assert len(p["finish_time"]) == 8 + + +def test_payload_counts_failed_and_caps_at_100(fmt): + p = go._distributed_progress_payload({"completed": 20, "failed": 2}, None, BASE, "n", "s") + assert p["progress_pct"] == 100 + + +def test_payload_without_eta(fmt): + p = go._distributed_progress_payload(STATUS, None, BASE, "n", "s") + assert p["eta_str"] == "—" + assert p["finish_time"] == "" + assert p["workers"] == 0 + p2 = go._distributed_progress_payload( + STATUS, {"eta_seconds": None, "avg_duration": None, "active_workers": 0}, BASE, "n", "s") + assert p2["eta_str"] == "—" + + +class _Server: + def __init__(self): + self.events = [] + self.instance = self + + def send_sync(self, name, payload): + self.events.append((name, payload)) + + +class _Mgr: + unique_id = "n1" + session_name = "sess" + is_active = True + + def __init__(self, with_eta=True): + self.on_status = None + if with_eta: + self.get_eta = lambda: {"eta_seconds": 30, "avg_duration": 5.0, "active_workers": 3} + + def get_status(self): + return dict(STATUS) + + +def test_emit_sends_progress_event(monkeypatch, fmt): + srv = _Server() + monkeypatch.setattr(go, "PromptServer", srv) + go._emit_distributed_progress(_Mgr(), BASE, "n1", "sess") + assert len(srv.events) == 1 + name, p = srv.events[0] + assert name == "ultimate_grid.progress" + assert p["workers"] == 3 and p["current_job"] == 7 and p["eta_str"] == "30s" + + +def test_emit_without_get_eta_and_never_raises(monkeypatch, fmt): + srv = _Server() + monkeypatch.setattr(go, "PromptServer", srv) + go._emit_distributed_progress(_Mgr(with_eta=False), BASE, "n1", "sess") + assert srv.events[0][1]["eta_str"] == "—" + monkeypatch.setattr(go, "PromptServer", None) + go._emit_distributed_progress(_Mgr(), BASE, "n1", "sess") # no server: no error + go._emit_distributed_progress(object(), BASE, "n1", "sess") # broken manager: no error + + +def test_send_distribution_status_calls_on_status(): + mgr = _Mgr() + seen = [] + mgr.on_status = seen.append + go._send_distribution_status(mgr) # no `server` module under test: the callback still fires + assert len(seen) == 1 and seen[0]["completed"] == 3 and seen[0]["active"] is True + + +def test_send_distribution_status_survives_failing_callback(): + mgr = _Mgr() + + def boom(st): + raise RuntimeError("x") + + mgr.on_status = boom + go._send_distribution_status(mgr) diff --git a/tests/test_dist_reject_recheck.py b/tests/test_dist_reject_recheck.py new file mode 100644 index 0000000..5bf520a --- /dev/null +++ b/tests/test_dist_reject_recheck.py @@ -0,0 +1,134 @@ +"""Claim-time re-check of previews rejected while a distributed full pass runs.""" +import json +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import generation_orchestrator as go # preloaded by tests/conftest.py + +FP = "fp:abc" +EXTRA = {"stamp": {"config_fp": FP}} + + +def _preview(**kw): + p = {"id": "p1", "preview": True, "preview_of": FP, "seed": 1} + p.update(kw) + return p + + +def _write(path, items, meta=None, bump=0): + with open(path, "w", encoding="utf-8") as f: + json.dump({"items": items, "meta": meta or {}}, f) + st = os.stat(path) + os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns + bump)) + + +class _Clock: + def __init__(self): + self.t = 0.0 + + def __call__(self): + return self.t + + +def _make(tmp_path, items=(), meta=None, disk_items=None, disk_meta=None, **kw): + mp = str(tmp_path / "manifest.json") + _write(mp, list(items if disk_items is None else disk_items), disk_meta if disk_meta is not None else meta) + mem = {"items": list(items), "meta": meta or {}} + kw.setdefault("skip_rejected", True) + kw.setdefault("only_favorited", False) + clock = _Clock() + chk = go._make_reject_recheck(mp, mem, min_interval=2.0, clock=clock, **kw) + return mp, chk, clock + + +def test_not_rejected_returns_none(tmp_path): + _, chk, _ = _make(tmp_path, [_preview()]) + assert chk(EXTRA, {}) is None + + +def test_rewritten_manifest_with_rejection_is_seen(tmp_path): + mp, chk, clock = _make(tmp_path, [_preview()]) + assert chk(EXTRA, {}) is None + _write(mp, [_preview(rejected=True)], bump=5_000_000_000) + clock.t = 10.0 + assert chk(EXTRA, {}) == "rejected_preview" + + +def test_disk_tombstone(tmp_path): + _, chk, _ = _make(tmp_path, [_preview()], disk_meta={"rejected_preview_fps": [FP]}) + assert chk(EXTRA, {}) == "rejected_preview" + + +def test_in_memory_tombstone(tmp_path): + _, chk, _ = _make(tmp_path, [_preview()], meta={"rejected_preview_fps": [FP]}, disk_meta={}) + assert chk(EXTRA, {}) == "rejected_preview" + + +def test_only_favorited_with_unfavorited_preview(tmp_path): + _, chk, _ = _make(tmp_path, [_preview()], skip_rejected=False, only_favorited=True) + assert chk(EXTRA, {}) == "no_favorited_preview" + + +def test_preview_jobs_and_unstamped_jobs_never_dropped(tmp_path): + _, chk, _ = _make(tmp_path, [_preview(rejected=True)]) + assert chk({"stamp": {"config_fp": FP}, "render": {"preview_scale": 0.5}}, {}) is None + assert chk({}, {}) is None + assert chk(None, {}) is None + assert chk({"stamp": {}}, {}) is None + + +def test_corrupt_file_keeps_old_index(tmp_path): + mp, chk, clock = _make(tmp_path, [_preview(rejected=True)]) + assert chk(EXTRA, {}) == "rejected_preview" + with open(mp, "w", encoding="utf-8") as f: + f.write("{not json") + clock.t = 10.0 + assert chk(EXTRA, {}) == "rejected_preview" + + +def test_no_reread_until_min_interval(tmp_path): + mp, chk, clock = _make(tmp_path, [_preview()]) + assert chk(EXTRA, {}) is None + _write(mp, [_preview(rejected=True)], bump=5_000_000_000) + clock.t = 1.0 # < min_interval: still the cached index + assert chk(EXTRA, {}) is None + clock.t = 2.5 + assert chk(EXTRA, {}) == "rejected_preview" + + +# -- _install_claim_filter ---------------------------------------------------------------------------------- + +class _Mgr: + def __init__(self): + self.filter = None + + def set_claim_filter(self, fn): + self.filter = fn + + +def _install(mgr, tmp_path, **kw): + base = dict(preview_on=False, has_optional_inputs=False, skip_rejected=True, only_favorited=False) + base.update(kw) + return go._install_claim_filter( + mgr, {"manifest": str(tmp_path / "m.json")}, {"items": [], "meta": {}}, **base) + + +def test_installed_for_a_full_pass(tmp_path): + m = _Mgr() + assert _install(m, tmp_path) is True + assert callable(m.filter) + + +def test_not_installed_for_preview_pass_or_optional_inputs_or_no_rules(tmp_path): + for kw in ({"preview_on": True}, {"has_optional_inputs": True}, + {"skip_rejected": False, "only_favorited": False}): + m = _Mgr() + assert _install(m, tmp_path, **kw) is False + assert m.filter is None + + +def test_manager_without_set_claim_filter_is_fine(tmp_path, capsys): + assert _install(object(), tmp_path) is False + assert "update" in capsys.readouterr().out.lower() diff --git a/tests/test_dist_start_at.py b/tests/test_dist_start_at.py new file mode 100644 index 0000000..7dce6e4 --- /dev/null +++ b/tests/test_dist_start_at.py @@ -0,0 +1,120 @@ +"""Distributed mode honours Start At Job: plan positions before N are skipped by the master.""" +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import generation_orchestrator as go # preloaded by tests/conftest.py +from config_utils import build_job_plan, order_job_plan + + +@pytest.fixture(autouse=True) +def _prompts(monkeypatch): + monkeypatch.setattr(go, "build_prompt_with_triggers", lambda conf, mode: (conf["positive"], "")) + + +def _conf(**kw): + c = {"model": "m", "lora": "None", "lora_expanded": "None", "sampler": "euler", "scheduler": "normal", + "steps": 30, "cfg": 7.0, "denoise": 1.0, "seed": 5, "positive": "p", "negative": "n"} + c.update(kw) + return c + + +def _job(w=1024, h=1024): + return {"label": f"{w}x{h}", "width": w, "height": h, "batch_idx": 0} + + +def _plan(confs): + return build_job_plan(confs, [_job()]) + + +def _full_item(conf): + return {"seed": conf["seed"], "width": 1024, "height": 1024, "batch_idx": 0, "sampler": conf["sampler"], + "scheduler": conf["scheduler"], "steps": conf["steps"], "cfg": conf["cfg"], "denoise": 1.0, + "model": conf["model"], "positive": conf["positive"], "negative": conf["negative"], "lora": "None"} + + +def _run(plan, items=(), **kw): + kw.setdefault("overwrite_existing", False) + kw.setdefault("has_optional_inputs", False) + kw.setdefault("lora_triggerwords_mode", "off") + return go._distributed_plan_entries(plan, {"items": list(items), "meta": {}}, **kw) + + +def _five(): + return [_conf(seed=i, positive=f"p{i}") for i in range(1, 6)] + + +def test_start_at_three_skips_first_two(): + entries, counts = _run(_plan(_five()), start_at_job=3) + assert [e["config"]["seed"] for e in entries] == [3, 4, 5] + assert counts["start_at"] == 2 + + +@pytest.mark.parametrize("n", [0, 1]) +def test_zero_and_one_skip_nothing(n): + entries, counts = _run(_plan(_five()), start_at_job=n) + assert [e["config"]["seed"] for e in entries] == [1, 2, 3, 4, 5] + assert counts.get("start_at", 0) == 0 + + +def test_larger_than_plan_skips_everything(): + entries, counts = _run(_plan(_five()), start_at_job=99) + assert entries == [] + assert counts["start_at"] == 5 + + +def test_job_before_n_that_exists_counts_as_start_at(): + confs = _five() + entries, counts = _run(_plan(confs), [_full_item(confs[0])], start_at_job=3) + assert counts["start_at"] == 2 + assert counts["exists"] == 0 + assert [e["config"]["seed"] for e in entries] == [3, 4, 5] + + +def test_positions_follow_the_ordered_plan(): + confs = [_conf(seed=1, model="A"), _conf(seed=2, model="B"), _conf(seed=3, model="A"), + _conf(seed=4, model="B")] + plan = order_job_plan(_plan(confs), "minimize_reloads", go.get_model_cache_key) + seeds = [pj["conf"]["seed"] for pj in plan] + entries, counts = _run(plan, start_at_job=3) + assert [e["config"]["seed"] for e in entries] == seeds[2:] + assert counts["start_at"] == 2 + + +class _OldManager: + def __init__(self): + self.called = None + + def populate_jobs(self, *a): + self.called = ("populate_jobs", a) + + +class _NewManager(_OldManager): + def populate_from_plan(self, entries, run_extras=None): + self.called = ("populate_from_plan", entries, run_extras) + + +def _populate(mgr, plan, **kw): + confs = [pj["conf"] for pj in plan] + return go._populate_distributed_manager(mgr, plan, confs, [_job()], {"items": []}, False, False, "off", **kw) + + +def test_legacy_manager_warns_that_start_at_is_ignored(capsys): + mgr = _OldManager() + ret = _populate(mgr, _plan(_five()), start_at_job=3) + assert mgr.called[0] == "populate_jobs" + assert "Start At Job" in capsys.readouterr().out + assert ret["queued"] is None and ret["plan_total"] is None + + +def test_populate_returns_plan_total_and_queued(): + mgr = _NewManager() + ret = _populate(mgr, _plan(_five()), start_at_job=3) + assert ret["plan_total"] == 5 + assert [e["config"]["seed"] for e in ret["queued"]] == [3, 4, 5] + assert ret["counts"]["start_at"] == 2 + assert ret["ltx_entries"] == [] + assert [e["config"]["seed"] for e in mgr.called[1]] == [3, 4, 5]