From 243086391dfcb96dfadf30f696181d89989e9f86 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 28 Feb 2026 02:20:28 +0000 Subject: [PATCH] Add remote worker log viewing and local log proxy endpoints --- api/worker_routes.py | 124 +++++++++++++++++++++++++++++++- tests/api/test_worker_routes.py | 102 ++++++++++++++++++++++++++ web/apiClient.js | 4 ++ web/constants.js | 2 + web/ui.js | 16 ++--- web/ui/logModal.js | 14 +++- web/workerLifecycle.js | 39 ++++++---- 7 files changed, 276 insertions(+), 25 deletions(-) diff --git a/api/worker_routes.py b/api/worker_routes.py index e3d042f..336eec7 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -13,7 +13,13 @@ import server from ..utils.config import load_config from ..utils.logging import debug_log, log -from ..utils.network import build_worker_url, handle_api_error, normalize_host, probe_worker +from ..utils.network import ( + build_worker_url, + get_client_session, + handle_api_error, + normalize_host, + probe_worker, +) from ..utils.constants import CHUNK_SIZE from ..workers import get_worker_manager from .schemas import require_fields, validate_worker_id @@ -319,6 +325,71 @@ def _read_worker_log_sync(log_file, lines_to_read): } +def _parse_positive_int_query(value, default, minimum=1, maximum=10000): + """Parse bounded positive integer query params with sane fallback.""" + try: + parsed = int(value) + except (TypeError, ValueError): + return default + parsed = max(minimum, parsed) + if maximum is not None: + parsed = min(maximum, parsed) + return parsed + + +def _find_worker_by_id(config, worker_id): + worker_id_str = str(worker_id).strip() + for worker in config.get("workers", []): + if str(worker.get("id")).strip() == worker_id_str: + return worker + return None + + +@server.PromptServer.instance.routes.get("/distributed/local_log") +async def get_local_log_endpoint(request): + """Return this instance's in-memory ComfyUI log buffer.""" + try: + from app.logger import get_logs + except Exception as e: + return await handle_api_error(request, f"Failed to import app.logger: {e}", 500) + + try: + lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000) + logs = get_logs() + if logs is None: + return web.json_response( + { + "status": "success", + "content": "", + "entries": 0, + "source": "memory", + "truncated": False, + "lines_shown": 0, + } + ) + + entries = list(logs) + selected_entries = entries[-lines_to_read:] + content = "".join( + entry.get("m", "") if isinstance(entry, dict) else str(entry) + for entry in selected_entries + ) + lines_shown = content.count("\n") + (1 if content else 0) + + return web.json_response( + { + "status": "success", + "content": content, + "entries": len(selected_entries), + "source": "memory", + "truncated": len(entries) > len(selected_entries), + "lines_shown": lines_shown, + } + ) + except Exception as e: + return await handle_api_error(request, e, 500) + + @server.PromptServer.instance.routes.get("/distributed/network_info") async def get_network_info_endpoint(request): """Get network interfaces and recommend best IP for master.""" @@ -553,7 +624,7 @@ async def get_worker_log_endpoint(request): return await handle_api_error(request, "Log file not found", 404) # Read last N lines (or full file if small) - lines_to_read = int(request.query.get('lines', 1000)) + lines_to_read = _parse_positive_int_query(request.query.get('lines'), default=1000) try: loop = asyncio.get_running_loop() @@ -573,3 +644,52 @@ async def get_worker_log_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}") +async def get_remote_worker_log_endpoint(request): + """Proxy a remote worker log request to the worker's local in-memory log endpoint.""" + try: + worker_id = str(request.match_info["worker_id"]).strip() + config = load_config() + worker = _find_worker_by_id(config, worker_id) + if not worker: + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + + # Remote log proxy is only meaningful for remote/cloud workers. + host = normalize_host(worker.get("host")) or "" + if not host: + return await handle_api_error( + request, + f"Worker {worker_id} is local; use /distributed/worker_log/{worker_id} instead.", + 400, + ) + + lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000) + worker_url = build_worker_url(worker, "/distributed/local_log") + session = await get_client_session() + async with session.get( + worker_url, + params={"lines": str(lines_to_read)}, + timeout=aiohttp.ClientTimeout(total=5), + ) as resp: + if resp.status >= 400: + body = await resp.text() + return await handle_api_error( + request, + f"Remote worker {worker_id} returned HTTP {resp.status}: {body[:400]}", + resp.status, + ) + + try: + data = await resp.json() + except Exception as e: + return await handle_api_error( + request, + f"Remote worker {worker_id} returned invalid JSON: {e}", + 502, + ) + + return web.json_response(data) + except Exception as e: + return await handle_api_error(request, e, 500) diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index 34b49f2..9a4c562 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -4,6 +4,7 @@ import sys import tempfile import types import unittest +from collections import deque from pathlib import Path from unittest.mock import patch @@ -24,6 +25,35 @@ class _FakeRequest: return self._payload +class _FakeHTTPClientResponse: + def __init__(self, payload, status=200): + self._payload = payload + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, _exc_type, _exc, _tb): + return False + + async def json(self): + return self._payload + + async def text(self): + return str(self._payload) + + +class _FakeHTTPClientSession: + def __init__(self, payload, status=200): + self._payload = payload + self._status = status + self.calls = [] + + def get(self, url, params=None, timeout=None): + self.calls.append({"url": url, "params": params, "timeout": timeout}) + return _FakeHTTPClientResponse(self._payload, status=self._status) + + class _DummyWorkerManager: def __init__(self): self.processes = {} @@ -295,6 +325,78 @@ class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase): self.assertIn("content", response.payload) self.assertIn("line-3", response.payload["content"]) + async def test_local_log_reads_memory_buffer(self): + request = _FakeRequest(query={"lines": "2"}) + fake_logs = deque( + [ + {"m": "line-1\n"}, + {"m": "line-2\n"}, + {"m": "line-3\n"}, + ], + maxlen=300, + ) + app_module = types.ModuleType("app") + app_module.__path__ = [] + logger_module = types.ModuleType("app.logger") + logger_module.get_logs = lambda: fake_logs + app_module.logger = logger_module + + with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}): + response = await worker_routes.get_local_log_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("status"), "success") + self.assertEqual(response.payload.get("source"), "memory") + self.assertEqual(response.payload.get("entries"), 2) + self.assertIn("line-3", response.payload.get("content", "")) + + async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self): + config = { + "workers": [ + { + "id": "worker-remote", + "name": "Remote Worker", + "host": "worker.example.com", + "port": 8188, + "type": "remote", + } + ] + } + request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"}) + proxied_payload = { + "status": "success", + "content": "remote-log-content\n", + "entries": 1, + "source": "memory", + "truncated": False, + "lines_shown": 1, + } + fake_session = _FakeHTTPClientSession(proxied_payload) + + async def _fake_get_client_session(): + return fake_session + + with patch.object(worker_routes, "load_config", return_value=config), patch.object( + worker_routes, "get_client_session", side_effect=_fake_get_client_session + ): + response = await worker_routes.get_remote_worker_log_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("content"), "remote-log-content\n") + self.assertEqual(len(fake_session.calls), 1) + self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"}) + self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log")) + + async def test_remote_worker_log_rejects_local_workers(self): + config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]} + request = _FakeRequest(match_info={"worker_id": "worker-local"}) + + with patch.object(worker_routes, "load_config", return_value=config): + response = await worker_routes.get_remote_worker_log_endpoint(request) + + self.assertEqual(response.status, 400) + self.assertIn("local", response.payload.get("message", "").lower()) + if __name__ == "__main__": unittest.main() diff --git a/web/apiClient.js b/web/apiClient.js index 7c1292e..a731ff7 100644 --- a/web/apiClient.js +++ b/web/apiClient.js @@ -135,6 +135,10 @@ export function createApiClient(baseUrl) { async getWorkerLog(workerId, lines = 1000) { return request(`/distributed/worker_log/${workerId}?lines=${lines}`); }, + + async getRemoteWorkerLog(workerId, lines = 300) { + return request(`/distributed/remote_worker_log/${workerId}?lines=${lines}`); + }, async clearLaunchingFlag(workerId) { return request('/distributed/worker/clear_launching', { diff --git a/web/constants.js b/web/constants.js index d89376c..3ce53c3 100644 --- a/web/constants.js +++ b/web/constants.js @@ -205,6 +205,8 @@ export const ENDPOINTS = { STOP_WORKER: '/distributed/stop_worker', MANAGED_WORKERS: '/distributed/managed_workers', WORKER_LOG: '/distributed/worker_log', + REMOTE_WORKER_LOG: '/distributed/remote_worker_log', + LOCAL_LOG: '/distributed/local_log', CLEAR_LAUNCHING: '/distributed/worker/clear_launching', PREPARE_JOB: '/distributed/prepare_job', LOAD_IMAGE: '/distributed/load_image', diff --git a/web/ui.js b/web/ui.js index c16a6ac..f5af377 100644 --- a/web/ui.js +++ b/web/ui.js @@ -390,7 +390,7 @@ export class DistributedUI { statusDot.classList.toggle('status-pulsing', pulsing); } - showLogModal(extension, workerId, logData) { + showLogModal(extension, workerId, logData, fetchLog = null) { if (this._logModal) { this._logModal.unmount(); this._logModal = null; @@ -404,7 +404,7 @@ export class DistributedUI { modal.mount(document.body, { workerName, logData, - fetchLog: async () => extension.api.getWorkerLog(workerId, 1000), + fetchLog: fetchLog || (async () => extension.api.getWorkerLog(workerId, 1000)), onClose: () => { if (this._logModal === modal) { this._logModal = null; @@ -584,12 +584,12 @@ export class DistributedUI { controlsWrapper.appendChild(badge); } else if (config.dynamic && data) { if (isRemote) { - const isCloud = data.type === 'cloud'; - const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; - const remoteInfo = this.createButton(workerTypeText, null, BUTTON_STYLES.info); - remoteInfo.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.info + " color: #999; cursor: default;"; - remoteInfo.disabled = true; - controlsWrapper.appendChild(remoteInfo); + const logBtn = this.createButton('View Log', () => viewWorkerLog(extension, data.id, true)); + logBtn.id = `log-${data.id}`; + logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; + logBtn.classList.add("btn--log"); + logBtn.title = "View remote worker log"; + controlsWrapper.appendChild(logBtn); } else { const controls = this.createWorkerControls(data.id, { launch: () => launchWorker(extension, data.id), diff --git a/web/ui/logModal.js b/web/ui/logModal.js index 80b7532..12ad3da 100644 --- a/web/ui/logModal.js +++ b/web/ui/logModal.js @@ -29,9 +29,17 @@ export function createLogModal() { _logContentEl.scrollTop = _logContentEl.scrollHeight; } - let statusText = `Log file: ${logData.log_file || 'unknown'}`; - if (logData.truncated) { - statusText += ` (showing last ${logData.lines_shown} lines of ${formatFileSize(logData.file_size || 0)})`; + let statusText; + if (logData.source === "memory") { + statusText = "Remote worker log (in-memory buffer)"; + if (logData.truncated) { + statusText += ` (showing last ${logData.lines_shown || 0} lines)`; + } + } else { + statusText = `Log file: ${logData.log_file || 'unknown'}`; + if (logData.truncated) { + statusText += ` (showing last ${logData.lines_shown} lines of ${formatFileSize(logData.file_size || 0)})`; + } } _statusBarEl.textContent = statusText; }; diff --git a/web/workerLifecycle.js b/web/workerLifecycle.js index 4893a8f..8acdc2e 100644 --- a/web/workerLifecycle.js +++ b/web/workerLifecycle.js @@ -381,8 +381,20 @@ export function updateWorkerControls(extension, workerId) { return; } - // Skip button updates for remote workers + // Update button states - buttons are now inside a wrapper div + const launchBtn = document.getElementById(`launch-${workerId}`); + const stopBtn = document.getElementById(`stop-${workerId}`); + const logBtn = document.getElementById(`log-${workerId}`); + if (isRemoteWorker(extension, worker)) { + setButtonVisibility(launchBtn, false); + setButtonVisibility(stopBtn, false); + if (logBtn) { + setButtonVisibility(logBtn, true); + logBtn.disabled = false; + logBtn.textContent = "View Log"; + setButtonClass(logBtn, "btn--log"); + } return; } @@ -390,11 +402,6 @@ export function updateWorkerControls(extension, workerId) { const managedInfo = extension.state.getWorker(workerId).managed; const status = extension.state.getWorkerStatus(workerId); - // Update button states - buttons are now inside a wrapper div - const launchBtn = document.getElementById(`launch-${workerId}`); - const stopBtn = document.getElementById(`stop-${workerId}`); - const logBtn = document.getElementById(`log-${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); @@ -429,9 +436,11 @@ export function updateWorkerControls(extension, workerId) { } } -export async function viewWorkerLog(extension, workerId) { +export async function viewWorkerLog(extension, workerId, isRemote = false) { + const worker = extension.config.workers.find((w) => w.id === workerId); + const isRemoteLog = isRemote || (worker ? isRemoteWorker(extension, worker) : false); const managedInfo = extension.state.getWorker(workerId).managed; - if (!managedInfo?.log_file) { + if (!isRemoteLog && !managedInfo?.log_file) { return; } @@ -445,11 +454,13 @@ export async function viewWorkerLog(extension, workerId) { } try { - // Fetch log content - const data = await extension.api.getWorkerLog(workerId, 1000); + const fetchLog = isRemoteLog + ? async () => extension.api.getRemoteWorkerLog(workerId, 300) + : async () => extension.api.getWorkerLog(workerId, 1000); + const data = await fetchLog(); // Create modal dialog - extension.ui.showLogModal(extension, workerId, data); + extension.ui.showLogModal(extension, workerId, data, fetchLog); // Restore button if (logBtn) { @@ -486,7 +497,11 @@ export async function refreshLog(extension, workerId, silent = false) { } try { - const data = await extension.api.getWorkerLog(workerId, 1000); + const worker = extension.config.workers.find((w) => w.id === workerId); + const isRemoteLog = worker ? isRemoteWorker(extension, worker) : false; + const data = isRemoteLog + ? await extension.api.getRemoteWorkerLog(workerId, 300) + : await extension.api.getWorkerLog(workerId, 1000); // Update content const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50;