From dd3fbfadbe94548b096612a5fa500016489b1dd4 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 2 Mar 2026 06:57:41 +0000 Subject: [PATCH] Fix branch join convergence and remote log visibility gating --- api/orchestration/prompt_transform.py | 5 +++- api/queue_orchestration.py | 1 + tests/test_prompt_transform.py | 18 +++++++++++++ web/constants.js | 1 + web/tests/workerLifecycle.test.js | 37 ++++++++++++++++++++++++++- web/ui/entityCard.js | 2 +- web/workerLifecycle.js | 25 +++++++++++++++--- 7 files changed, 82 insertions(+), 7 deletions(-) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 75d2b68..993bbb6 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -603,7 +603,10 @@ def _override_join_nodes( assigned_branch = -1 inputs = node.setdefault("inputs", {}) - inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + # Group all joins under the same DistributedBranch into one queue so + # worker-branch senders converge into the master's active join node. + job_source_id = upstream_branch_id or node_id + inputs["multi_job_id"] = job_id_map.get(job_source_id, job_source_id) inputs["is_worker"] = not is_master inputs["enabled_worker_ids"] = enabled_json inputs["assigned_branch"] = assigned_branch diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 1a32c62..5879d9d 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -327,6 +327,7 @@ async def orchestrate_distributed_execution( prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class("DistributedListCollector") + prompt_index.nodes_for_class("DistributedJoin") + + prompt_index.nodes_for_class("DistributedBranch") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") ) for node_id in queue_node_ids: diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 55c7e52..5407ebf 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -717,5 +717,23 @@ class ApplyOverridesJoinTests(unittest.TestCase): result = _apply(prompt, "master", enabled_worker_ids=["worker-a"]) self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1") + def test_join_uses_upstream_branch_job_id_for_grouped_convergence(self): + master_prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 0], "num_branches": 2}}, + } + worker_prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 1], "num_branches": 2}}, + } + + master_result = _apply(master_prompt, "master", enabled_worker_ids=["worker-a"]) + worker_result = _apply(worker_prompt, "worker-a", enabled_worker_ids=["worker-a"]) + + self.assertEqual(master_result["3"]["inputs"]["multi_job_id"], "run_2") + self.assertEqual(worker_result["3"]["inputs"]["multi_job_id"], "run_2") + if __name__ == "__main__": unittest.main() diff --git a/web/constants.js b/web/constants.js index 1ec83d7..f9885e8 100644 --- a/web/constants.js +++ b/web/constants.js @@ -175,6 +175,7 @@ export const TIMEOUTS = { LAUNCH: 90000, // ms for worker launch (longer for model loading) RETRY_DELAY: 1000, // initial delay for exponential backoff MAX_RETRIES: 3, // max retry attempts + REMOTE_LOG_AVAILABILITY_WINDOW: 30000, // show remote log button for this long after last online probe // UI feedback delays BUTTON_RESET: 3000, // button text/state reset after actions diff --git a/web/tests/workerLifecycle.test.js b/web/tests/workerLifecycle.test.js index b402b82..f5c8cfc 100644 --- a/web/tests/workerLifecycle.test.js +++ b/web/tests/workerLifecycle.test.js @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest"; -import { getWorkerUrl } from "../workerLifecycle.js"; +import { TIMEOUTS } from "../constants.js"; +import { getWorkerUrl, shouldShowRemoteLogButton } from "../workerLifecycle.js"; describe("workerLifecycle URL construction", () => { @@ -50,3 +51,37 @@ describe("workerLifecycle URL construction", () => { expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt"); }); }); + +describe("workerLifecycle remote log visibility", () => { + it("shows remote log when worker is online", () => { + const worker = { enabled: true }; + const status = { online: true }; + expect(shouldShowRemoteLogButton(worker, status, 100_000)).toBe(true); + }); + + it("shows remote log shortly after last online detection", () => { + const worker = { enabled: true }; + const now = 100_000; + const status = { + online: false, + lastSeenOnlineAt: now - (TIMEOUTS.REMOTE_LOG_AVAILABILITY_WINDOW - 1), + }; + expect(shouldShowRemoteLogButton(worker, status, now)).toBe(true); + }); + + it("hides remote log when worker has been unavailable past the freshness window", () => { + const worker = { enabled: true }; + const now = 100_000; + const status = { + online: false, + lastSeenOnlineAt: now - (TIMEOUTS.REMOTE_LOG_AVAILABILITY_WINDOW + 1), + }; + expect(shouldShowRemoteLogButton(worker, status, now)).toBe(false); + }); + + it("hides remote log when worker is disabled", () => { + const worker = { enabled: false }; + const status = { online: true, lastSeenOnlineAt: 99_000 }; + expect(shouldShowRemoteLogButton(worker, status, 100_000)).toBe(false); + }); +}); diff --git a/web/ui/entityCard.js b/web/ui/entityCard.js index 69bcd14..a499d39 100644 --- a/web/ui/entityCard.js +++ b/web/ui/entityCard.js @@ -120,7 +120,7 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { ui.addPlaceholderHover(card, leftColumn, entityType); } - if (isWorker && !isRemote) { + if (isWorker) { updateWorkerControls(extension, data.id); } diff --git a/web/workerLifecycle.js b/web/workerLifecycle.js index 8acdc2e..d3b5510 100644 --- a/web/workerLifecycle.js +++ b/web/workerLifecycle.js @@ -41,6 +41,20 @@ function setButtonVisibility(button, visible) { button.style.display = visible ? "" : "none"; } +export function shouldShowRemoteLogButton(worker, status = {}, now = Date.now()) { + if (!worker?.enabled) { + return false; + } + if (status?.online) { + return true; + } + const lastSeenOnlineAt = Number(status?.lastSeenOnlineAt || 0); + if (!Number.isFinite(lastSeenOnlineAt) || lastSeenOnlineAt <= 0) { + return false; + } + return (now - lastSeenOnlineAt) <= TIMEOUTS.REMOTE_LOG_AVAILABILITY_WINDOW; +} + export async function checkAllWorkerStatuses(extension) { if (_statusCheckRunning || !extension.panelElement) { return; @@ -142,6 +156,7 @@ export function getWorkerUrl(extension, worker, endpoint = "") { export async function checkWorkerStatus(extension, worker) { // Assume caller ensured enabled; proceed with check const workerUrl = getWorkerUrl(extension, worker); + const previousStatus = extension.state.getWorkerStatus(worker.id); try { const signal = extension.statusCheckAbortController?.signal || null; @@ -162,6 +177,7 @@ export async function checkWorkerStatus(extension, worker) { online: true, processing: isProcessing, queueCount: queueRemaining, + lastSeenOnlineAt: Date.now(), }); // Update status dot based on probe result @@ -183,6 +199,7 @@ export async function checkWorkerStatus(extension, worker) { online: false, processing: false, queueCount: 0, + lastSeenOnlineAt: previousStatus?.lastSeenOnlineAt || 0, }); // Check if worker is launching @@ -385,13 +402,15 @@ export function updateWorkerControls(extension, workerId) { const launchBtn = document.getElementById(`launch-${workerId}`); const stopBtn = document.getElementById(`stop-${workerId}`); const logBtn = document.getElementById(`log-${workerId}`); + const status = extension.state.getWorkerStatus(workerId); if (isRemoteWorker(extension, worker)) { setButtonVisibility(launchBtn, false); setButtonVisibility(stopBtn, false); if (logBtn) { - setButtonVisibility(logBtn, true); - logBtn.disabled = false; + const showRemoteLog = shouldShowRemoteLogButton(worker, status); + setButtonVisibility(logBtn, showRemoteLog); + logBtn.disabled = !showRemoteLog; logBtn.textContent = "View Log"; setButtonClass(logBtn, "btn--log"); } @@ -400,8 +419,6 @@ export function updateWorkerControls(extension, workerId) { // Ensure we check for string ID const managedInfo = extension.state.getWorker(workerId).managed; - const status = extension.state.getWorkerStatus(workerId); - // Show log button immediately if we have log file info (even if worker is still starting) if (logBtn) { const showLog = Boolean(managedInfo?.log_file);