diff --git a/Dockerfile b/Dockerfile index bb0c04d1b..942cacb57 100644 --- a/Dockerfile +++ b/Dockerfile @@ -33,7 +33,7 @@ RUN pip install uv && uv --version && \ # install sageattention ADD pkg/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl /workspace/pkg/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl -RUN uv pip install -U --no-deps --no-build-isolation spandrel timm tensorboard poetry flash-attn "xformers==0.0.31.post1" "file:./pkg/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl" +RUN uv pip install -U --no-deps --no-build-isolation spandrel timm tensorboard poetry "flash-attn<=2.8.0" "xformers==0.0.31.post1" "file:./pkg/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl" # this exotic command will determine the correct torchaudio to install for the image RUN <<-EOF python -c 'import torch, re, subprocess @@ -66,7 +66,7 @@ WORKDIR /workspace # addresses https://github.com/pytorch/pytorch/issues/104801 # and issues reported by importing nodes_canny # smoke test -RUN python -c "import torch; import xformers; import sageattention; import cv2" && comfyui --quick-test-for-ci --cpu --cwd /workspace +RUN python -c "import torch; import xformers; import sageattention; import cv2; import diffusers.hooks" && comfyui --quick-test-for-ci --cpu --cwd /workspace EXPOSE 8188 CMD ["python", "-m", "comfy.cmd.main", "--listen", "--use-sage-attention", "--reserve-vram=0", "--logging-level=INFO", "--enable-cors"] diff --git a/comfy/cli_args.py b/comfy/cli_args.py index fe964ba45..63157dd24 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -156,6 +156,7 @@ def _create_parser() -> EnhancedConfigArgParser: parser.add_argument("--whitelist-custom-nodes", type=str, action=FlattenAndAppendAction, nargs='+', default=[], help="Specify custom node folders to load even when --disable-all-custom-nodes is enabled.") parser.add_argument("--blacklist-custom-nodes", type=str, action=FlattenAndAppendAction, nargs='+', default=[], help="Specify custom node folders to never load. Accepts shell-style globs.") parser.add_argument("--disable-api-nodes", action="store_true", help="Disable loading all api nodes.") + parser.add_argument("--enable-eval", action="store_true", help="Enable nodes that can evaluate Python code in workflows.") parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.") parser.add_argument("--create-directories", action="store_true", diff --git a/comfy/cli_args_types.py b/comfy/cli_args_types.py index 953747f95..903f46f1a 100644 --- a/comfy/cli_args_types.py +++ b/comfy/cli_args_types.py @@ -169,6 +169,7 @@ class Configuration(dict): whitelist_custom_nodes (list[str]): Specify custom node folders to load even when --disable-all-custom-nodes is enabled. default_device (Optional[int]): Set the id of the default device, all other devices will stay visible. block_runtime_package_installation (Optional[bool]): When set, custom nodes like ComfyUI Manager, Easy Use, Nunchaku and others will not be able to use pip or uv to install packages at runtime (experimental). + enable_eval (Optional[bool]): Enable nodes that can evaluate Python code in workflows. """ def __init__(self, **kwargs): @@ -288,6 +289,7 @@ class Configuration(dict): self.database_url: str = db_config() self.default_device: Optional[int] = None self.block_runtime_package_installation = None + self.enable_eval: Optional[bool] = False for key, value in kwargs.items(): self[key] = value @@ -420,6 +422,7 @@ class FlattenAndAppendAction(argparse.Action): Custom action to handle comma-separated values and multiple invocations of the same argument, flattening them into a single list. """ + def __call__(self, parser, namespace, values, option_string=None): items = getattr(namespace, self.dest, None) if items is None: diff --git a/comfy_extras/eval_web/__init__.py b/comfy_extras/eval_web/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/comfy_extras/eval_web/ace_utils.js b/comfy_extras/eval_web/ace_utils.js new file mode 100644 index 000000000..78c00d809 --- /dev/null +++ b/comfy_extras/eval_web/ace_utils.js @@ -0,0 +1,769 @@ +/** + * Uses code adapted from https://github.com/yorkane/ComfyUI-KYNode + * + * MIT License + * + * Copyright (c) 2024 Kevin Yuan + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +// Make modal window +function makeModal({ title = "Message", text = "No text", type = "info", parent = null, stylePos = "fixed", classes = [] } = {}) { + const overlay = document.createElement("div"); + Object.assign(overlay.style, { + display: "none", + position: stylePos, + background: "rgba(0 0 0 / 0.8)", + opacity: 0, + top: "0", + left: "0", + right: "0", + bottom: "0", + zIndex: "500", + transition: "all .8s", + cursor: "pointer", + }); + + const boxModal = document.createElement("div"); + Object.assign(boxModal.style, { + transition: "all 0.5s", + opacity: 0, + display: "none", + position: stylePos, + top: "50%", + left: "50%", + transform: "translate(-50%,-50%)", + background: "#525252", + minWidth: "300px", + fontFamily: "sans-serif", + zIndex: "501", + border: "1px solid rgb(255 255 255 / 45%)", + }); + + boxModal.className = "alekpet_modal_window"; + boxModal.classList.add(...classes); + + const boxModalBody = document.createElement("div"); + Object.assign(boxModalBody.style, { + display: "flex", + flexDirection: "column", + textAlign: "center", + }); + + boxModalBody.className = "alekpet_modal_body"; + + const boxModalHtml = ` +
+
${title}
+
✕
+
+
${text}
`; + boxModalBody.innerHTML = boxModalHtml; + + const alekpet_modal_header = boxModalBody.querySelector(".alekpet_modal_header"); + Object.assign(alekpet_modal_header.style, { + display: "flex", + alignItems: "center", + }); + + const close = boxModalBody.querySelector(".alekpet_modal_close"); + Object.assign(close.style, { + cursor: "pointer", + }); + + let parentElement = document.body; + if (parent && parent.nodeType === 1) { + parentElement = parent; + } + + boxModal.append(boxModalBody); + parentElement.append(overlay, boxModal); + + const removeEvent = new Event("removeElements"); + const remove = () => { + animateTransitionProps(boxModal, { opacity: 0 }).then(() => + animateTransitionProps(overlay, { opacity: 0 }).then(() => { + parentElement.removeChild(boxModal); + parentElement.removeChild(overlay); + }), + ); + }; + + boxModal.addEventListener("removeElements", remove); + overlay.addEventListener("removeElements", remove); + + animateTransitionProps(overlay) + .then(() => { + overlay.addEventListener("click", () => { + overlay.dispatchEvent(removeEvent); + }); + animateTransitionProps(boxModal); + }) + .then(() => boxModal.querySelector(".alekpet_modal_close").addEventListener("click", () => boxModal.dispatchEvent(removeEvent))); +} + +function findWidget(node, value, attr = "name", func = "find") { + return node?.widgets ? node.widgets[func]((w) => (Array.isArray(value) ? value.includes(w[attr]) : w[attr] === value)) : null; +} + +function animateTransitionProps(el, props = { opacity: 1 }, preStyles = { display: "block" }) { + Object.assign(el.style, preStyles); + + el.style.transition = !el.style.transition || !window.getComputedStyle(el).getPropertyValue("transition") ? "all .8s" : el.style.transition; + + return new Promise((res) => { + setTimeout(() => { + Object.assign(el.style, props); + + const transstart = () => (el.isAnimating = true); + const transchancel = () => (el.isAnimating = false); + el.addEventListener("transitionstart", transstart); + el.addEventListener("transitioncancel", transchancel); + + el.addEventListener("transitionend", function transend() { + el.isAnimating = false; + el.removeEventListener("transitionend", transend); + el.removeEventListener("transitionend", transchancel); + el.removeEventListener("transitionend", transstart); + res(el); + }); + }, 100); + }); +} + +function animateClick(target, params = {}) { + const { opacityVal = 0.9, callback = () => {} } = params; + if (target?.isAnimating) return; + + const hide = +target.style.opacity === 0; + return animateTransitionProps(target, { + opacity: hide ? opacityVal : 0, + }).then((el) => { + const isHide = hide || el.style.display === "none"; + showHide({ elements: [target], hide: !hide }); + callback(); + return isHide; + }); +} + +function showHide({ elements = [], hide = null, displayProp = "block" } = {}) { + Array.from(elements).forEach((el) => { + if (hide !== null) { + el.style.display = !hide ? displayProp : "none"; + } else { + el.style.display = !el.style.display || el.style.display === "none" ? displayProp : "none"; + } + }); +} + +function isEmptyObject(obj) { + if (!obj) return true; + return Object.keys(obj).length === 0 && obj.constructor === Object; +} + +function makeElement(tag, attrs = {}) { + if (!tag) tag = "div"; + const element = document.createElement(tag); + Object.keys(attrs).forEach((key) => { + const currValue = attrs[key]; + if (key === "class") { + if (Array.isArray(currValue)) { + element.classList.add(...currValue); + } else if (currValue instanceof String || typeof currValue === "string") { + element.className = currValue; + } + } else if (key === "dataset") { + try { + if (Array.isArray(currValue)) { + currValue.forEach((datasetArr) => { + const [prop, propval] = Object.entries(datasetArr)[0]; + element.dataset[prop] = propval; + }); + } else { + Object.entries(currValue).forEach((datasetArr) => { + const [prop, propval] = datasetArr; + element.dataset[prop] = propval; + }); + } + } catch (err) { + console.log(err); + } + } else if (key === "style") { + if (typeof currValue === "object" && !Array.isArray(currValue) && Object.keys(currValue).length) { + Object.assign(element[key], currValue); + } else if (typeof currValue === "object" && Array.isArray(currValue) && currValue.length) { + element[key] = [...currValue]; + } else if (currValue instanceof String || typeof currValue === "string") { + element[key] = currValue; + } + } else if (["for"].includes(key)) { + element.setAttribute(key, currValue); + } else if (key === "children") { + element.append(...(currValue instanceof Array ? currValue : [currValue])); + } else if (key === "parent") { + currValue.append(element); + } else { + element[key] = currValue; + } + }); + return element; +} + +function isValidStyle(opt, strColor) { + let op = new Option().style; + if (!op.hasOwnProperty(opt)) return { result: false, color: "", color_hex: "" }; + + op[opt] = strColor; + + return { + result: op[opt] !== "", + color_rgb: op[opt], + color_hex: rgbToHex(op[opt]), + }; +} + +function rgbToHex(rgb) { + const regEx = new RegExp(/\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)/); + if (regEx.test(rgb)) { + let [, r, g, b] = regEx.exec(rgb); + r = parseInt(r).toString(16); + g = parseInt(g).toString(16); + b = parseInt(b).toString(16); + + r = r.length === 1 ? r + "0" : r; + g = g.length === 1 ? g + "0" : g; + b = b.length === 1 ? b + "0" : b; + + return `#${r}${g}${b}`; + } +} + +async function getDataJSON(url) { + try { + const response = await fetch(url); + const jsonData = await response.json(); + return jsonData; + } catch (err) { + return new Error(err); + } +} + +function deepMerge(target, source) { + if (source?.nodeType) return; + for (let key in source) { + if (source[key] instanceof Object && key in target) { + Object.assign(source[key], deepMerge(target[key], source[key])); + } + } + + Object.assign(target || {}, source); + return target; +} + +const THEME_MODAL_WINDOW_BASE = { + stylesTitle: { + background: "auto", + padding: "5px", + borderRadius: "6px", + marginBottom: "5px", + alignSelf: "stretch", + }, + stylesWrapper: { + display: "none", + opacity: 0, + minWidth: "220px", + position: "absolute", + left: "50%", + top: "50%", + transform: "translate(-50%, -50%)", + transition: "all .8s", + fontFamily: "monospace", + zIndex: 99999, + }, + stylesBox: { + display: "flex", + flexDirection: "column", + background: "#0e0e0e", + padding: "6px", + justifyContent: "center", + alignItems: "center", + gap: "3px", + textAlign: "center", + borderRadius: "6px", + color: "white", + border: "2px solid silver", + boxShadow: "2px 2px 4px silver", + maxWidth: "300px", + }, + stylesClose: { + position: "absolute", + top: "-10px", + right: "-10px", + background: "silver", + borderRadius: "50%", + width: "20px", + height: "20px", + cursor: "pointer", + display: "flex", + justifyContent: "center", + alignItems: "center", + fontSize: "0.8rem", + }, +}; + +const THEMES_MODAL_WINDOW = { + error: { + stylesTitle: { + ...THEME_MODAL_WINDOW_BASE.stylesTitle, + background: "#8f210f", + }, + stylesBox: { + ...THEME_MODAL_WINDOW_BASE.stylesBox, + background: "#3b2222", + boxShadow: "3px 3px 6px #141414", + border: "1px solid #f91b1b", + }, + stylesWrapper: { ...THEME_MODAL_WINDOW_BASE.stylesWrapper }, + stylesClose: { + ...THEME_MODAL_WINDOW_BASE.stylesClose, + background: "#3b2222", + }, + }, + warning: { + stylesTitle: { + ...THEME_MODAL_WINDOW_BASE.stylesTitle, + background: "#e99818", + }, + stylesBox: { + ...THEME_MODAL_WINDOW_BASE.stylesBox, + background: "#594e32", + boxShadow: "3px 3px 6px #141414", + border: "1px solid #e99818", + }, + stylesWrapper: { ...THEME_MODAL_WINDOW_BASE.stylesWrapper }, + stylesClose: { + ...THEME_MODAL_WINDOW_BASE.stylesClose, + background: "#594e32", + }, + }, + normal: { + stylesTitle: { + ...THEME_MODAL_WINDOW_BASE.stylesTitle, + background: "#108f0f", + }, + stylesBox: { + ...THEME_MODAL_WINDOW_BASE.stylesBox, + background: "#223b2a", + boxShadow: "3px 3px 6px #141414", + border: "1px solid #108f0f", + }, + stylesWrapper: { ...THEME_MODAL_WINDOW_BASE.stylesWrapper }, + stylesClose: { + ...THEME_MODAL_WINDOW_BASE.stylesClose, + background: "#223b2a", + }, + }, +}; + +const defaultOptions = { + auto: { + autohide: false, + autoshow: false, + autoremove: false, + propStyles: { opacity: 0 }, + propPreStyles: {}, + timewait: 2000, + }, + overlay: { + overlay_enabled: false, + overlayClasses: [], + overlayStyles: {}, + }, + close: { closeRemove: false, showClose: true }, + parent: null, +}; + +function createWindowModal({ textTitle = "Message", textBody = "Hello world!", textFooter = null, classesWrapper = [], stylesWrapper = {}, classesBox = [], stylesBox = {}, classesTitle = [], stylesTitle = {}, classesBody = [], stylesBody = {}, classesClose = [], stylesClose = {}, classesFooter = [], stylesFooter = {}, options = defaultOptions } = {}) { + // Check all options exist + const _options = deepMerge(JSON.parse(JSON.stringify(defaultOptions)), options); + + const { + parent, + overlay: { overlay_enabled, overlayClasses, overlayStyles }, + close: { closeRemove, showClose }, + auto: { autohide, autoshow, autoremove, timewait, propStyles, propPreStyles }, + } = _options; + + // Function past text(html) + function addText(text, parent) { + if (!parent) return; + + switch (typeof text) { + case "string": + if (/^\<.*\/?\>$/.test(text)) { + parent.innerHTML = text; + } else { + parent.textContent = text; + } + break; + case "object": + default: + if (Array.isArray(text)) { + text.forEach((element) => (element.nodeType === 1 || element.nodeType === 3) && parent.append(element)); + } else if (text.nodeType === 1 || text.nodeType === 3) parent.append(text); + } + } + + // Overlay + let overlayElement = null; + if (overlay_enabled) { + overlayElement = makeElement("div", { + class: [...overlayClasses], + style: { + display: "none", + position: "fixed", + background: "rgba(0 0 0 / 0.8)", + opacity: 0, + top: 0, + left: 0, + right: 0, + bottom: 0, + zIndex: 99999, + transition: "all .8s", + cursor: "pointer", + ...overlayStyles, + }, + }); + } + + // Wrapper + const wrapper_settings = makeElement("div", { + class: ["alekpet__wrapper__window", ...classesWrapper], + }); + + Object.assign(wrapper_settings.style, { + ...THEME_MODAL_WINDOW_BASE.stylesWrapper, + ...stylesWrapper, + }); + + // Box + const box__settings = makeElement("div", { + class: ["alekpet__window__box", ...classesBox], + }); + Object.assign(box__settings.style, { + ...THEME_MODAL_WINDOW_BASE.stylesBox, + ...stylesBox, + }); + + // Title + let box_settings_title = ""; + if (textTitle) { + box_settings_title = makeElement("div", { + class: ["alekpet__window__title", ...classesTitle], + }); + + Object.assign(box_settings_title.style, { + ...THEME_MODAL_WINDOW_BASE.stylesTitle, + ...stylesTitle, + }); + + // Add text (html) to title + addText(textTitle, box_settings_title); + } + // Body + let box_settings_body = ""; + if (textBody) { + box_settings_body = makeElement("div", { + class: ["alekpet__window__body", ...classesBody], + }); + + Object.assign(box_settings_body.style, { + display: "flex", + flexDirection: "column", + alignItems: "flex-end", + gap: "5px", + textWrap: "wrap", + ...stylesBody, + }); + + // Add text (html) to body + addText(textBody, box_settings_body); + } + + // Close button + const close__box__button = makeElement("div", { + class: ["close__box__button", ...classesClose], + textContent: "✖", + }); + + Object.assign(close__box__button.style, { + ...THEME_MODAL_WINDOW_BASE.stylesClose, + ...stylesClose, + }); + + if (!showClose) close__box__button.style.display = "none"; + + const closeEvent = new Event("closeModal"); + const closeModalWindow = function () { + overlay_enabled + ? animateTransitionProps(overlayElement, { + opacity: 0, + }) + .then(() => + animateTransitionProps(wrapper_settings, { + opacity: 0, + }), + ) + .then(() => { + if (closeRemove) { + parent.removeChild(wrapper_settings); + parent.removeChild(overlayElement); + } else { + showHide({ elements: [wrapper_settings, overlayElement] }); + } + }) + : animateTransitionProps(wrapper_settings, { + opacity: 0, + }).then(() => { + showHide({ elements: [wrapper_settings] }); + }); + }; + + close__box__button.addEventListener("closeModal", closeModalWindow); + + close__box__button.addEventListener("click", () => close__box__button.dispatchEvent(closeEvent)); + + close__box__button.onmouseenter = () => { + close__box__button.style.opacity = 0.8; + }; + + close__box__button.onmouseleave = () => { + close__box__button.style.opacity = 1; + }; + + box__settings.append(box_settings_title, box_settings_body); + + // Footer + if (textFooter) { + const box_settings_footer = makeElement("div", { + class: [...classesFooter], + }); + Object.assign(box_settings_footer.style, { + ...stylesFooter, + }); + + // Add text (html) to body + addText(textFooter, box_settings_footer); + + box__settings.append(box_settings_footer); + } + + wrapper_settings.append(close__box__button, box__settings); + + if (parent && parent.nodeType === 1) { + if (overlay_enabled) parent.append(overlayElement); + parent.append(wrapper_settings); + + if (autoshow) { + overlay_enabled + ? animateClick(overlayElement).then(() => + animateClick(wrapper_settings).then( + () => + autohide && + setTimeout( + () => + animateTransitionProps(wrapper_settings, { ...propStyles }, { ...propPreStyles }) + .then(() => animateTransitionProps(overlayElement, { ...propStyles }, { ...propPreStyles })) + .then(() => { + if (autoremove) { + parent.removeChild(wrapper_settings); + parent.removeChild(overlayElement); + } + }), + timewait, + ), + ), + ) + : animateClick(wrapper_settings).then(() => autohide && setTimeout(() => animateTransitionProps(wrapper_settings, { ...propStyles }, { ...propPreStyles }).then(() => autoremove && parent.removeChild(wrapper_settings)), timewait)); + } + } + + return wrapper_settings; +} + +// Prompt +async function comfyuiDesktopPrompt(title, message, defaultValue) { + try { + return await app.extensionManager.dialog.prompt({ + title, + message, + defaultValue, + }); + } catch (err) { + return prompt(title, message); + } +} + +// Alert +function comfyuiDesktopAlert(message) { + try { + app.extensionManager.toast.addAlert(message); + } catch (err) { + alert(message); + } +} + +// Confirm +function confirmModal({ title, message }) { + return new Promise((res) => { + const overlay = makeElement("div", { + class: ["alekpet_confOverlay"], + style: { + background: "rgba(0, 0, 0, 0.7)", + position: "fixed", + top: 0, + left: 0, + right: 0, + bottom: 0, + zIndex: 9999, + userSelect: "none", + }, + }); + + const modal = makeElement("div", { + class: ["alekpet_confModal"], + style: { + ...THEME_MODAL_WINDOW_BASE.stylesBox, + position: "fixed", + top: "50%", + left: "50%", + fontFamily: "monospace", + background: "rgb(92 186 255 / 20%)", + transform: "translate(-50%, -50%)", + borderColor: "rgba(92, 186, 255, 0.63)", + boxShadow: "rgba(92, 186, 255, 0.63) 2px 2px 4px", + }, + }); + + const titleEl = makeElement("div", { + class: ["alekpet_confTitle"], + style: { + ...THEME_MODAL_WINDOW_BASE.stylesTitle, + background: "rgba(92, 186, 255, 0.63)", + }, + textContent: title, + }); + + const messageEl = makeElement("div", { + class: ["alekpet_confMessage"], + style: { + display: "flex", + flexDirection: "column", + alignItems: "flex-end", + gap: "5px", + textWrap: "wrap", + }, + textContent: message, + }); + + const action_box = makeElement("div", { + class: ["alekpet_confActions"], + style: { + display: "flex", + gap: "5px", + width: "100%", + padding: "4px", + justifyContent: "flex-end", + }, + }); + + const remove = () => { + modal.remove(); + overlay.remove(); + }; + + const ok = makeElement("div", { + class: ["alekpet_confButtons", "alekpet_confButtonOk"], + style: { + background: "linear-gradient(45deg, green, limegreen) rgb(21, 100, 6)", + }, + textContent: "Ok", + onclick: (e) => { + res(true); + remove(); + }, + }); + + const Cancel = makeElement("div", { + class: ["alekpet_confButtons", "alekpet_confButtonCancel"], + style: { + background: "linear-gradient(45deg, #b64396, #a52a8b) rgb(135 3 161)", + }, + textContent: "Cancel", + onclick: (e) => { + res(false); + remove(); + }, + }); + + action_box.append(ok, Cancel); + modal.append(titleEl, messageEl, action_box); + overlay.append(modal); + document.body.append(overlay); + }); +} + +async function comfyuiDesktopConfirm(message) { + try { + const result = await confirmModal({ + title: "Confirm", + message: message, + }); + + // Wait update comfyui frontend! Confirm Cancel not return value! Fixed in ComfyUI_frontend ver. v1.10.8 + // https://github.com/Comfy-Org/ComfyUI_frontend/issues/2649 + // const result = await app.extensionManager.dialog.confirm({ + // title: "Confirm", + // message: message, + // }); + return result; + } catch (err) { + return confirm(message); + } +} + +export { + makeModal, + createWindowModal, + animateTransitionProps, + animateClick, + showHide, + makeElement, + getDataJSON, + isEmptyObject, + isValidStyle, + rgbToHex, + findWidget, + THEMES_MODAL_WINDOW, + // + comfyuiDesktopConfirm, + comfyuiDesktopPrompt, + comfyuiDesktopAlert, +}; diff --git a/comfy_extras/eval_web/ky_eval_python.js b/comfy_extras/eval_web/ky_eval_python.js new file mode 100644 index 000000000..3d65aa5c0 --- /dev/null +++ b/comfy_extras/eval_web/ky_eval_python.js @@ -0,0 +1,377 @@ +/** + * Uses code adapted from https://github.com/yorkane/ComfyUI-KYNode + * + * MIT License + * + * Copyright (c) 2024 Kevin Yuan + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +import { app } from "../../scripts/app.js"; + +import * as ace from "https://cdn.jsdelivr.net/npm/ace-code@1.43.4/+esm"; +import { makeElement, findWidget } from "./ace_utils.js"; + +// Constants +const varTypes = ["int", "boolean", "string", "float", "json", "list", "dict"]; +const typeMap = { + int: "int", + boolean: "bool", + string: "str", + float: "float", + json: "json", + list: "list", + dict: "dict", +}; + +ace.config.setModuleLoader('ace/mode/python', () => + import('https://cdn.jsdelivr.net/npm/ace-builds@1.43.4/src/mode-python.js') +); + +ace.config.setModuleLoader('ace/theme/monokai', () => + import('https://cdn.jsdelivr.net/npm/ace-builds@1.43.4/src/theme-monokai.js') +); + +function getPostition(node, ctx, w_width, y, n_height) { + const margin = 5; + + const rect = ctx.canvas.getBoundingClientRect(); + const transform = new DOMMatrix() + .scaleSelf(rect.width / ctx.canvas.width, rect.height / ctx.canvas.height) + .multiplySelf(ctx.getTransform()) + .translateSelf(margin, margin + y); + const scale = new DOMMatrix().scaleSelf(transform.a, transform.d); + + return { + transformOrigin: "0 0", + transform: scale, + left: `${transform.a + transform.e + rect.left}px`, + top: `${transform.d + transform.f + rect.top}px`, + maxWidth: `${w_width - margin * 2}px`, + maxHeight: `${n_height - margin * 2 - y - 15}px`, + width: `${w_width - margin * 2}px`, + height: "90%", + position: "absolute", + scrollbarColor: "var(--descrip-text) var(--bg-color)", + scrollbarWidth: "thin", + zIndex: app.graph._nodes.indexOf(node), + }; +} + +// Create editor code +function codeEditor(node, inputName, inputData) { + const widget = { + type: "pycode", + name: inputName, + options: { hideOnZoom: true }, + value: + inputData[1]?.default || + `def my(a, b=1): + return a * b
+ +r0 = str(my(23, 9))`, + draw(ctx, node, widget_width, y, widget_height) { + const hidden = node.flags?.collapsed || (!!widget.options.hideOnZoom && app.canvas.ds.scale < 0.5) || widget.type === "converted-widget" || widget.type === "hidden"; + + widget.codeElement.hidden = hidden; + + if (hidden) { + widget.options.onHide?.(widget); + return; + } + + Object.assign(this.codeElement.style, getPostition(node, ctx, widget_width, y, node.size[1])); + }, + computeSize(...args) { + return [500, 250]; + }, + }; + + widget.codeElement = makeElement("pre", { + innerHTML: widget.value, + }); + + widget.editor = ace.edit(widget.codeElement); + widget.editor.setTheme("ace/theme/monokai"); + widget.editor.session.setMode("ace/mode/python"); + widget.editor.setOptions({ + enableAutoIndent: true, + enableLiveAutocompletion: true, + enableBasicAutocompletion: true, + fontFamily: "monospace", + }); + widget.codeElement.hidden = true; + + document.body.appendChild(widget.codeElement); + + const collapse = node.collapse; + node.collapse = function () { + collapse.apply(this, arguments); + if (this.flags?.collapsed) { + widget.codeElement.hidden = true; + } else { + if (this.flags?.collapsed === false) { + widget.codeElement.hidden = false; + } + } + }; + + return widget; +} + +// Save data to workflow forced! +function saveValue() { + app?.extensionManager?.workflow?.activeWorkflow?.changeTracker?.checkState(); +} + +// Register extensions +app.registerExtension({ + name: "KYNode.KY_Eval_Python", + getCustomWidgets(app) { + return { + PYCODE: (node, inputName, inputData, app) => { + const widget = codeEditor(node, inputName, inputData); + + widget.editor.getSession().on("change", function (e) { + widget.value = widget.editor.getValue(); + saveValue(); + }); + + const varTypeList = node.addWidget( + "combo", + "select_type", + "string", + (v) => { + // widget.editor.setTheme(`ace/theme/${varTypeList.value}`); + }, + { + values: varTypes, + serialize: false, + }, + ); + + // 6. 使用 addDOMWidget 将容器添加到节点上 + // - 第一个参数是 widget 的名称，在节点内部需要是唯一的。 + // - 第二个参数是 widget 的类型，对于自定义 DOM 元素，通常是 "div"。 + // - 第三个参数是您创建的 DOM 元素。 + // - 第四个参数是一个选项对象，可以用来配置 widget。 + // node.addDOMWidget("rowOfButtons", "div", container, { + // }); + node.addWidget("button", "Add Input variable", "add_input_variable", async () => { + // Input name variable and check + let nameInput = node?.inputs?.length ? `p${node.inputs.length - 1}` : "p0"; + + const currentWidth = node.size[0]; + let tp = varTypeList.value; + nameInput = nameInput + "_" + typeMap[tp]; + node.addInput(nameInput, "*"); + node.setSize([currentWidth, node.size[1]]); + let cv = widget.editor.getValue(); + if (tp === "json") { + cv = cv + "\n" + nameInput + " = json.loads(" + nameInput + ")"; + } else if (tp === "list") { + cv = cv + "\n" + nameInput + " = []"; + } else if (tp === "dict") { + cv = cv + "\n" + nameInput + " = {}"; + } else { + cv = cv + "\n" + nameInput + " = " + typeMap[tp] + "(" + nameInput + ")"; + } + widget.editor.setValue(cv); + saveValue(); + }); + + node.addWidget("button", "Add Output variable", "add_output_variable", async () => { + const currentWidth = node.size[0]; + // Output name variable + let nameOutput = node?.outputs?.length ? `r${node.outputs.length}` : "r0"; + let tp = varTypeList.value; + nameOutput = nameOutput + "_" + typeMap[tp]; + node.addOutput(nameOutput, tp); + node.setSize([currentWidth, node.size[1]]); + let cv = widget.editor.getValue(); + if (tp === "json") { + cv = cv + "\n" + nameOutput + " = json.dumps(" + nameOutput + ")"; + } else if (tp === "list") { + cv = cv + "\n" + nameOutput + " = []"; + } else if (tp === "dict") { + cv = cv + "\n" + nameOutput + " = {}"; + } else { + cv = cv + "\n" + nameOutput + " = " + typeMap[tp] + "(" + nameOutput + ")"; + } + widget.editor.setValue(cv); + saveValue(); + }); + + node.onRemoved = function () { + for (const w of node?.widgets) { + if (w?.codeElement) w.codeElement.remove(); + } + }; + + node.addCustomWidget(widget); + + return widget; + }, + }; + }, + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // --- IDENode + if (nodeData.name === "KY_Eval_Python") { + // Node Created + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = async function () { + const ret = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + + const node_title = await this.getTitle(); + const nodeName = `${nodeData.name}_${this.id}`; + + this.name = nodeName; + + // Create default inputs, when first create node + if (this?.inputs?.length < 2) { + ["p0_str"].forEach((inputName) => { + const currentWidth = this.size[0]; + this.addInput(inputName, "*"); + this.setSize([currentWidth, this.size[1]]); + }); + } + + const widgetEditor = findWidget(this, "pycode", "type"); + + this.setSize([530, this.size[1]]); + + return ret; + }; + + const onDrawForeground = nodeType.prototype.onDrawForeground; + nodeType.prototype.onDrawForeground = function (ctx) { + const r = onDrawForeground?.apply?.(this, arguments); + + // if (this.flags?.collapsed) return r; + + if (this?.outputs?.length) { + for (let o = 0; o < this.outputs.length; o++) { + const { name, type } = this.outputs[o]; + const colorType = LGraphCanvas.link_type_colors[type.toUpperCase()]; + const nameSize = ctx.measureText(name); + const typeSize = ctx.measureText(`[${type === "*" ? "any" : type.toLowerCase()}]`); + + ctx.fillStyle = colorType === "" ? "#AAA" : colorType; + ctx.font = "12px Arial, sans-serif"; + ctx.textAlign = "right"; + ctx.fillText(`[${type === "*" ? "any" : type.toLowerCase()}]`, this.size[0] - nameSize.width - typeSize.width, o * 20 + 19); + } + } + + if (this?.inputs?.length) { + const not_showing = ["select_type", "pycode"]; + for (let i = 1; i < this.inputs.length; i++) { + const { name, type } = this.inputs[i]; + if (not_showing.includes(name)) continue; + const colorType = LGraphCanvas.link_type_colors[type.toUpperCase()]; + const nameSize = ctx.measureText(name); + + ctx.fillStyle = !colorType || colorType === "" ? "#AAA" : colorType; + ctx.font = "12px Arial, sans-serif"; + ctx.textAlign = "left"; + ctx.fillText(`[${type === "*" ? "any" : type.toLowerCase()}]`, nameSize.width + 25, i * 20); + } + } + return r; + }; + + // Node Configure + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function (node) { + onConfigure?.apply(this, arguments); + if (node?.widgets_values?.length) { + const widget_code_id = findWidget(this, "pycode", "type", "findIndex"); + const widget_theme_id = findWidget(this, "varTypeList", "name", "findIndex"); + const widget_language_id = findWidget(this, "language", "name", "findIndex"); + + const editor = this.widgets[widget_code_id]?.editor; + + if (editor) { + // editor.setTheme( + // `ace/theme/${this.widgets_values[widget_theme_id]}` + // ); + // editor.session.setMode( + // `ace/mode/${this.widgets_values[widget_language_id]}` + // ); + editor.setValue(this.widgets_values[widget_code_id]); + editor.clearSelection(); + } + } + }; + + // ExtraMenuOptions + const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (_, options) { + getExtraMenuOptions?.apply(this, arguments); + + const past_index = options.length - 1; + const past = options[past_index]; + + if (!!past) { + // Inputs remove + for (const input_idx in this.inputs) { + const input = this.inputs[input_idx]; + + if (["language", "select_type"].includes(input.name)) continue; + + options.splice(past_index + 1, 0, { + content: `Remove Input ${input.name}`, + callback: (e) => { + const currentWidth = this.size[0]; + if (input.link) { + app.graph.removeLink(input.link); + } + this.removeInput(input_idx); + this.setSize([80, this.size[1]]); + saveValue(); + }, + }); + } + + // Output remove + for (const output_idx in this.outputs) { + const output = this.outputs[output_idx]; + + if (output.name === "r0") continue; + + options.splice(past_index + 1, 0, { + content: `Remove Output ${output.name}`, + callback: (e) => { + const currentWidth = this.size[0]; + if (output.link) { + app.graph.removeLink(output.link); + } + this.removeOutput(output_idx); + this.setSize([currentWidth, this.size[1]]); + saveValue(); + }, + }); + } + } + }; + // end - ExtraMenuOptions + } + }, +}); diff --git a/comfy_extras/nodes/nodes_eval.py b/comfy_extras/nodes/nodes_eval.py new file mode 100644 index 000000000..a09739c21 --- /dev/null +++ b/comfy_extras/nodes/nodes_eval.py @@ -0,0 +1,110 @@ +import re +import traceback +import types + +from comfy.execution_context import current_execution_context +from comfy.node_helpers import export_package_as_web_directory, export_custom_nodes +from comfy.nodes.package_typing import CustomNode + +remove_type_name = re.compile(r"(\{.*\})", re.I | re.M) + + +# Hack: string type that is always equal in not equal comparisons, thanks pythongosssss +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + + +PY_CODE = AnyType("*") +IDEs_DICT = {} + + +# - Thank you very much for the class -> Trung0246 - +# - https://github.com/Trung0246/ComfyUI-0246/blob/main/utils.py#L51 +class TautologyStr(str): + def __ne__(self, other): + return False + + +class ByPassTypeTuple(tuple): + def __getitem__(self, index): + if index > 0: + index = 0 + item = super().__getitem__(index) + if isinstance(item, str): + return TautologyStr(item) + return item + + +# --------------------------- + + +class KY_Eval_Python(CustomNode): + @classmethod + def INPUT_TYPES(s): + + return { + "required": { + "pycode": ( + "PYCODE", + { + "default": """import re, json, os, traceback +from time import strftime + +def runCode(): + nowDataTime = strftime("%Y-%m-%d %H:%M:%S") + return f"Hello ComfyUI with us today {nowDataTime}!" +r0_str = runCode() + unique_id +""" + }, + ), + }, + "hidden": {"unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = ByPassTypeTuple((PY_CODE,)) + RETURN_NAMES = ("r0_str",) + FUNCTION = "exec_py" + DESCRIPTION = "IDE Node is an node that allows you to run code written in Python or Javascript directly in the node." + CATEGORY = "KYNode/Code" + + def exec_py(self, pycode, unique_id, extra_pnginfo, **kwargs): + ctx = current_execution_context() + if ctx.configuration.enable_eval is not True: + raise ValueError("Python eval is disabled") + + if unique_id not in IDEs_DICT: + IDEs_DICT[unique_id] = self + + outputs = {unique_id: unique_id} + if extra_pnginfo and 'workflow' in extra_pnginfo and extra_pnginfo['workflow']: + for node in extra_pnginfo['workflow']['nodes']: + if node['id'] == int(unique_id): + outputs_valid = [ouput for ouput in node.get('outputs', []) if ouput.get('name', '') != '' and ouput.get('type', '') != ''] + outputs = {ouput['name']: None for ouput in outputs_valid} + self.RETURN_TYPES = ByPassTypeTuple(out["type"] for out in outputs_valid) + self.RETURN_NAMES = tuple(name for name in outputs.keys()) + my_namespace = types.SimpleNamespace() + # 从 prompt 对象中提取 prompt_id + # if extra_data and 'extra_data' in extra_data and 'prompt_id' in extra_data['extra_data']: + # prompt_id = prompt['extra_data']['prompt_id'] + # outputs['p0_str'] = p0_str + + my_namespace.__dict__.update(outputs) + my_namespace.__dict__.update({prop: kwargs[prop] for prop in kwargs}) + # my_namespace.__dict__.setdefault("r0_str", "The r0 variable is not assigned") + + try: + exec(pycode, my_namespace.__dict__) + except Exception as e: + err = traceback.format_exc() + mc = re.search(r'line (\d+), in ([\w\W]+)$', err, re.MULTILINE) + msg = mc[1] + ':' + mc[2] + my_namespace.r0 = f"Error Line{msg}" + + new_dict = {key: my_namespace.__dict__[key] for key in my_namespace.__dict__ if key not in ['__builtins__', *kwargs.keys()] and not callable(my_namespace.__dict__[key])} + return (*new_dict.values(),) + + +export_custom_nodes() +export_package_as_web_directory("comfy_extras.eval_web") diff --git a/tests/conftest.py b/tests/conftest.py index 931ec8deb..1c5b3df20 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -100,7 +100,7 @@ def frontend_backend_worker_with_rabbitmq(request, tmp_path_factory, num_workers frontend_command = [ "comfyui", - "--listen=127.0.0.1", + "--listen=0.0.0.0", "--port=19001", "--cpu", "--distributed-queue-frontend", diff --git a/tests/distributed/test_tracing_integration.py b/tests/distributed/test_tracing_integration.py index 412b14c02..36a9fabd6 100644 --- a/tests/distributed/test_tracing_integration.py +++ b/tests/distributed/test_tracing_integration.py @@ -8,6 +8,7 @@ full distributed trace. import asyncio import logging import os +import tempfile import time import uuid @@ -21,6 +22,7 @@ from opentelemetry.sdk.trace.export import BatchSpanProcessor from opentelemetry.semconv.attributes import service_attributes from testcontainers.core.container import DockerContainer from testcontainers.core.waiting_utils import wait_for_logs +from testcontainers.nginx import NginxContainer from comfy.client.sdxl_with_refiner_workflow import sdxl_workflow_with_refiner @@ -54,6 +56,102 @@ class JaegerContainer(DockerContainer): return self +@pytest.fixture(scope="function") +def nginx_proxy(frontend_backend_worker_with_rabbitmq): + """ + Provide an nginx proxy in front of the ComfyUI frontend. + This tests if nginx is blocking W3C trace context propagation. + """ + import socket + import subprocess + + # Extract host and port from frontend address + frontend_url = frontend_backend_worker_with_rabbitmq + # frontend_url is like "http://127.0.0.1:19001" + import re + match = re.match(r'http://([^:]+):(\d+)', frontend_url) + if not match: + raise ValueError(f"Could not parse frontend URL: {frontend_url}") + + frontend_host = match.group(1) + frontend_port = match.group(2) + nginx_port = 8085 + + # Get the Docker bridge gateway IP (this is how containers reach the host on Linux) + # Try to get the default Docker bridge gateway + try: + result = subprocess.run( + ["docker", "network", "inspect", "bridge", "-f", "{{range .IPAM.Config}}{{.Gateway}}{{end}}"], + capture_output=True, + text=True, + check=True + ) + docker_gateway = result.stdout.strip() + logger.info(f"Using Docker gateway IP: {docker_gateway}") + except Exception as e: + # Fallback: try common gateway IPs + docker_gateway = "172.17.0.1" # Default Docker bridge gateway on Linux + logger.warning(f"Could not detect Docker gateway, using default: {docker_gateway}") + + # Create nginx config that proxies to the frontend and passes trace headers + nginx_conf = f""" +events {{ + worker_connections 1024; +}} + +http {{ + upstream backend {{ + server {docker_gateway}:{frontend_port}; + }} + + server {{ + listen {nginx_port}; + + location / {{ + proxy_pass http://backend; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + }} + }} +}} +""" + + # Write config to a temporary file + with tempfile.NamedTemporaryFile(mode='w', suffix='.conf', delete=False) as f: + f.write(nginx_conf) + nginx_conf_path = f.name + + try: + # Start nginx container with the config + nginx = NginxContainer(port=nginx_port) + nginx.with_volume_mapping(nginx_conf_path, "/etc/nginx/nginx.conf") + nginx.start() + + # Get the nginx URL + host = nginx.get_container_host_ip() + port = nginx.get_exposed_port(nginx_port) + nginx_url = f"http://{host}:{port}" + + logger.info(f"Nginx proxy started at {nginx_url} -> {frontend_url}") + + # Wait for nginx to be ready + for _ in range(30): + try: + response = requests.get(nginx_url, timeout=1) + if response.status_code: + break + except Exception: + pass + time.sleep(0.5) + + yield nginx_url + finally: + nginx.stop() + os.unlink(nginx_conf_path) + + @pytest.fixture(scope="module") def jaeger_container(): """ @@ -201,21 +299,21 @@ def verify_trace_continuity(trace: dict, expected_services: list[str]) -> bool: # order matters, execute jaeger_container first @pytest.mark.asyncio -async def test_tracing_integration(jaeger_container, frontend_backend_worker_with_rabbitmq): +async def test_tracing_integration(jaeger_container, nginx_proxy): """ - Integration test for distributed tracing across services. + Integration test for distributed tracing across services with nginx proxy. This test: 1. Starts ComfyUI frontend and worker with RabbitMQ - 2. Configures OTLP export to Jaeger testcontainer - 3. Submits a workflow through the frontend - 4. Queries Jaeger to verify trace propagation - 5. Validates that the trace spans multiple services with proper relationships + 2. Starts nginx proxy in front of the frontend to test trace context propagation through nginx + 3. Configures OTLP export to Jaeger testcontainer + 4. Submits a workflow through the nginx proxy + 5. Queries Jaeger to verify trace propagation + 6. Validates that the trace spans multiple services with proper relationships - Note: The frontend_backend_worker_with_rabbitmq fixture is parameterized, - so this test will run with both ThreadPoolExecutor and ProcessPoolExecutor. + This specifically tests if nginx is blocking W3C trace context (traceparent/tracestate headers). """ - server_address = frontend_backend_worker_with_rabbitmq + server_address = nginx_proxy jaeger_url = jaeger_container.get_query_url() otlp_endpoint = jaeger_container.get_otlp_endpoint() @@ -410,31 +508,27 @@ async def test_multiple_requests_different_traces(frontend_backend_worker_with_r # Query Jaeger and verify we have multiple distinct traces jaeger_url = jaeger_container.get_query_url() - try: - traces_response = query_jaeger_traces(jaeger_url, "comfyui", lookback="5m", limit=10) - traces = traces_response.get("data", []) + traces_response = query_jaeger_traces(jaeger_url, "comfyui", lookback="5m", limit=10) + traces = traces_response.get("data", []) - if len(traces) >= 2: - # Get trace IDs - trace_ids = [trace.get("traceID") for trace in traces] - unique_trace_ids = set(trace_ids) + assert len(traces) >= 2 + # Get trace IDs + trace_ids = [trace.get("traceID") for trace in traces] + unique_trace_ids = set(trace_ids) - logger.info(f"Found {len(unique_trace_ids)} unique traces") + logger.info(f"Found {len(unique_trace_ids)} unique traces") - # Verify we have multiple distinct traces - assert len(unique_trace_ids) >= 2, ( - f"Expected at least 2 distinct traces, found {len(unique_trace_ids)}. " - "Each request should create its own trace." - ) + # Verify we have multiple distinct traces + assert len(unique_trace_ids) >= 2, ( + f"Expected at least 2 distinct traces, found {len(unique_trace_ids)}. " + "Each request should create its own trace." + ) - logger.info("✓ Multiple requests created distinct traces") - else: - pytest.skip("Not enough traces to validate") - except Exception as e: - pytest.skip(f"Could not query Jaeger: {e}") + logger.info("✓ Multiple requests created distinct traces") @pytest.mark.asyncio +@pytest.mark.skip(reason="rabbitmq has to be configured for observability?") async def test_trace_contains_rabbitmq_operations(frontend_backend_worker_with_rabbitmq, jaeger_container): """ Test that traces include RabbitMQ publish/consume operations. @@ -455,43 +549,21 @@ async def test_trace_contains_rabbitmq_operations(frontend_backend_worker_with_r await asyncio.sleep(5) - try: - traces_response = query_jaeger_traces(jaeger_url, "comfyui", lookback="5m") - traces = traces_response.get("data", []) + traces_response = query_jaeger_traces(jaeger_url, "comfyui", lookback="5m") + traces = traces_response.get("data", []) - if traces: - # Look for RabbitMQ-related operations in any trace - rabbitmq_operations = [ - "publish", "consume", "amq_queue_publish", "amq_queue_consume", - "amq.basic.publish", "amq.basic.consume", "send", "receive" - ] + # Look for RabbitMQ-related operations in any trace + rabbitmq_operations = [ + "publish", "consume", "amq_queue_publish", "amq_queue_consume", + "amq.basic.publish", "amq.basic.consume", "send", "receive" + ] - found_rabbitmq_ops = [] - for trace in traces: - for span in trace.get("spans", []): - op_name = span.get("operationName", "").lower() - for rmq_op in rabbitmq_operations: - if rmq_op in op_name: - found_rabbitmq_ops.append(op_name) + found_rabbitmq_ops = [] + for trace in traces: + for span in trace.get("spans", []): + op_name = span.get("operationName", "").lower() + for rmq_op in rabbitmq_operations: + if rmq_op in op_name: + found_rabbitmq_ops.append(op_name) - if found_rabbitmq_ops: - logger.info(f"✓ Found RabbitMQ operations in traces: {set(found_rabbitmq_ops)}") - else: - logger.warning( - "No RabbitMQ operations found in traces. " - "This suggests that either:\n" - "1. AioPikaInstrumentor is not creating spans, or\n" - "2. The spans are being filtered out by the collector, or\n" - "3. The spans exist but use different operation names" - ) - - # Log all operation names to help debug - all_ops = set() - for trace in traces[:3]: # First 3 traces - for span in trace.get("spans", []): - all_ops.add(span.get("operationName")) - logger.info(f"Sample operation names: {all_ops}") - else: - pytest.skip("No traces found") - except Exception as e: - pytest.skip(f"Could not query Jaeger: {e}") + assert found_rabbitmq_ops, "No RabbitMQ-related operations found in traces" \ No newline at end of file