wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-16 14:47:33 +08:00
Code Issues Actions 22 Packages Projects Releases Wiki Activity

feat(isolation): wheel support for isolated custom nodes

Browse Source 
Extends pyisolate process isolation with wheel-based dependency
management, sandbox mode policy, and compatibility fixes validated
against DA3 as the first complex isolated custom node.

- Add sandbox_mode policy (required/disabled) with COMFY_HOST_POLICY_PATH
  env override for host security configuration
- Plumb cuda_wheels config and standardize child environment detection
- Add PLY, NPZ, File3D, VIDEO serializers for core save nodes
- Register isolated extension web directories on host side
  (sandbox_mode=disabled) for frontend JS widget serving
- Capture ACCEPT_ALL_INPUTS from child node class to prevent
  @classproperty trigger on proxy class
- Serialize NodeOutput with ui/expand/block_execution through JSON-RPC
- Cherry-pick V3 python_module metadata fix (e4592190)
- Remove improperly committed cache artifacts
This commit is contained in:
John Pollock 2026-03-15 01:25:40 -05:00
parent 54461f9ecc
commit 6aa0b838a0
11 changed files with 162 additions and 93117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.pyisolate_venvs/ComfyUI-DepthAnythingV2/cache/cache_key vendored
View File

@ -1 +0,0 @@
f03e4c88e21504c3

81
.pyisolate_venvs/ComfyUI-DepthAnythingV2/cache/node_info.json vendored
View File

@ -1,81 +0,0 @@
{
"DepthAnything_V2": {
"input_types": {
"required": {
"da_model": {
"__pyisolate_tuple__": [
"DAMODEL"
]
},
"images": {
"__pyisolate_tuple__": [
"IMAGE"
]
}
}
},
"return_types": [
"IMAGE"
],
"return_names": [
"image"
],
"function": "process",
"category": "DepthAnythingV2",
"output_node": false,
"output_is_list": null,
"is_v3": false,
"display_name": "Depth Anything V2"
},
"DownloadAndLoadDepthAnythingV2Model": {
"input_types": {
"required": {
"model": {
"__pyisolate_tuple__": [
[
"depth_anything_v2_vits_fp16.safetensors",
"depth_anything_v2_vits_fp32.safetensors",
"depth_anything_v2_vitb_fp16.safetensors",
"depth_anything_v2_vitb_fp32.safetensors",
"depth_anything_v2_vitl_fp16.safetensors",
"depth_anything_v2_vitl_fp32.safetensors",
"depth_anything_v2_vitg_fp32.safetensors",
"depth_anything_v2_metric_hypersim_vitl_fp32.safetensors",
"depth_anything_v2_metric_vkitti_vitl_fp32.safetensors"
],
{
"default": "depth_anything_v2_vitl_fp32.safetensors"
}
]
}
},
"optional": {
"precision": {
"__pyisolate_tuple__": [
[
"auto",
"bf16",
"fp16",
"fp32"
],
{
"default": "auto"
}
]
}
}
},
"return_types": [
"DAMODEL"
],
"return_names": [
"da_v2_model"
],
"function": "loadmodel",
"category": "DepthAnythingV2",
"output_node": false,
"output_is_list": null,
"is_v3": false,
"display_name": "DownloadAndLoadDepthAnythingV2Model"
}
}

1
.pyisolate_venvs/ComfyUI-IsolationToolkit/cache/cache_key vendored
View File

@ -1 +0,0 @@
4b90e6876f4c0b8c

93028
.pyisolate_venvs/ComfyUI-IsolationToolkit/cache/node_info.json vendored
View File

File diff suppressed because it is too large Load Diff

47
comfy/isolation/extension_loader.py
View File

@ -27,6 +27,47 @@ except ImportError:
logger = logging.getLogger(__name__)
def _register_web_directory(extension_name: str, node_dir: Path) -> None:
"""Register an isolated extension's web directory on the host side."""
import nodes
# Method 1: pyproject.toml [tool.comfy] web field
pyproject = node_dir / "pyproject.toml"
if pyproject.exists():
try:
with pyproject.open("rb") as f:
data = tomllib.load(f)
web_dir_name = data.get("tool", {}).get("comfy", {}).get("web")
if web_dir_name:
web_dir_path = str(node_dir / web_dir_name)
if os.path.isdir(web_dir_path):
nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
logger.debug("][ Registered web dir for isolated %s: %s", extension_name, web_dir_path)
return
except Exception:
pass
# Method 2: __init__.py WEB_DIRECTORY constant (parse without importing)
init_file = node_dir / "__init__.py"
if init_file.exists():
try:
source = init_file.read_text()
for line in source.splitlines():
stripped = line.strip()
if stripped.startswith("WEB_DIRECTORY"):
# Parse: WEB_DIRECTORY = "./web" or WEB_DIRECTORY = "web"
_, _, value = stripped.partition("=")
value = value.strip().strip("\"'")
if value:
web_dir_path = str((node_dir / value).resolve())
if os.path.isdir(web_dir_path):
nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
logger.debug("][ Registered web dir for isolated %s: %s", extension_name, web_dir_path)
return
except Exception:
pass
async def _stop_extension_safe(
extension: ComfyNodeExtension, extension_name: str
) -> None:
@ -255,6 +296,7 @@ async def load_isolated_node(
"dependencies": dependencies,
"share_torch": share_torch,
"share_cuda_ipc": share_cuda_ipc,
"sandbox_mode": host_policy["sandbox_mode"],
"sandbox": sandbox_config,
}
if cuda_wheels is not None:
@ -263,6 +305,11 @@ async def load_isolated_node(
extension = manager.load_extension(extension_config)
register_dummy_module(extension_name, node_dir)
# Register web directory on the host — only when sandbox is disabled.
# In sandbox mode, serving untrusted JS to the browser is not safe.
if host_policy["sandbox_mode"] == "disabled":
_register_web_directory(extension_name, node_dir)
# Try cache first (lazy spawn)
if is_cache_valid(node_dir, manifest_path, venv_root):
cached_data = load_from_cache(node_dir, venv_root)

25
comfy/isolation/extension_wrapper.py
View File

@ -237,6 +237,9 @@ class ComfyNodeExtension(ExtensionBase):
"not_idempotent": bool(
getattr(node_cls, "NOT_IDEMPOTENT", False)
),
"accept_all_inputs": bool(
getattr(node_cls, "ACCEPT_ALL_INPUTS", False)
),
}
)
except Exception as exc:
@ -369,14 +372,20 @@ class ComfyNodeExtension(ExtensionBase):
handler = getattr(instance, function_name)
try:
import torch
if asyncio.iscoroutinefunction(handler):
result = await handler(**resolved_inputs)
with torch.inference_mode():
result = await handler(**resolved_inputs)
else:
import functools
def _run_with_inference_mode(**kwargs):
with torch.inference_mode():
return handler(**kwargs)
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(
None, functools.partial(handler, **resolved_inputs)
None, functools.partial(_run_with_inference_mode, **resolved_inputs)
)
except Exception:
logger.exception(
@ -388,7 +397,17 @@ class ComfyNodeExtension(ExtensionBase):
raise
if type(result).__name__ == "NodeOutput":
result = result.args
node_output_dict = {
"__node_output__": True,
"args": self._wrap_unpicklable_objects(result.args),
}
if result.ui is not None:
node_output_dict["ui"] = result.ui
if getattr(result, "expand", None) is not None:
node_output_dict["expand"] = result.expand
if getattr(result, "block_execution", None) is not None:
node_output_dict["block_execution"] = result.block_execution
return node_output_dict
if self._is_comfy_protocol_return(result):
wrapped = self._wrap_unpicklable_objects(result)
return wrapped

28
comfy/isolation/host_policy.py
View File

@ -2,6 +2,7 @@
from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import Dict, List, TypedDict
@ -12,8 +13,12 @@ except ImportError:
logger = logging.getLogger(__name__)
HOST_POLICY_PATH_ENV = "COMFY_HOST_POLICY_PATH"
VALID_SANDBOX_MODES = frozenset({"required", "disabled"})
class HostSecurityPolicy(TypedDict):
sandbox_mode: str
allow_network: bool
writable_paths: List[str]
readonly_paths: List[str]
@ -21,6 +26,7 @@ class HostSecurityPolicy(TypedDict):
DEFAULT_POLICY: HostSecurityPolicy = {
"sandbox_mode": "required",
"allow_network": False,
"writable_paths": ["/dev/shm", "/tmp"],
"readonly_paths": [],
@ -30,6 +36,7 @@ DEFAULT_POLICY: HostSecurityPolicy = {
def _default_policy() -> HostSecurityPolicy:
return {
"sandbox_mode": DEFAULT_POLICY["sandbox_mode"],
"allow_network": DEFAULT_POLICY["allow_network"],
"writable_paths": list(DEFAULT_POLICY["writable_paths"]),
"readonly_paths": list(DEFAULT_POLICY["readonly_paths"]),
@ -38,7 +45,8 @@ def _default_policy() -> HostSecurityPolicy:
def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
config_path = comfy_root / "pyproject.toml"
config_override = os.environ.get(HOST_POLICY_PATH_ENV)
config_path = Path(config_override) if config_override else comfy_root / "pyproject.toml"
policy = _default_policy()
if not config_path.exists():
@ -61,6 +69,19 @@ def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
logger.debug("No [tool.comfy.host] section found, using defaults.")
return policy
sandbox_mode = tool_config.get("sandbox_mode")
if isinstance(sandbox_mode, str):
normalized_sandbox_mode = sandbox_mode.strip().lower()
if normalized_sandbox_mode in VALID_SANDBOX_MODES:
policy["sandbox_mode"] = normalized_sandbox_mode
else:
logger.warning(
"Invalid host sandbox_mode %r in %s, using default %r.",
sandbox_mode,
config_path,
DEFAULT_POLICY["sandbox_mode"],
)
if "allow_network" in tool_config:
policy["allow_network"] = bool(tool_config["allow_network"])
@ -75,7 +96,10 @@ def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
policy["whitelist"] = {str(k): str(v) for k, v in whitelist_raw.items()}
logger.debug(
f"Loaded Host Policy: {len(policy['whitelist'])} whitelisted nodes, Network={policy['allow_network']}"
"Loaded Host Policy: %d whitelisted nodes, Sandbox=%s, Network=%s",
len(policy["whitelist"]),
policy["sandbox_mode"],
policy["allow_network"],
)
return policy

22
comfy/isolation/runtime_helpers.py
View File

@ -215,6 +215,19 @@ def build_stub_class(
node_name,
node_unique_id or "-",
)
# Reconstruct NodeOutput if the child serialized one
if isinstance(result, dict) and result.get("__node_output__"):
from comfy_api.latest import io as latest_io
args_raw = result.get("args", ())
deserialized_args = await deserialize_from_isolation(args_raw, extension)
deserialized_args = _detach_shared_cpu_tensors(deserialized_args)
scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
return latest_io.NodeOutput(
*deserialized_args,
ui=result.get("ui"),
expand=result.get("expand"),
block_execution=result.get("block_execution"),
)
deserialized = await deserialize_from_isolation(result, extension)
scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
return _detach_shared_cpu_tensors(deserialized)
@ -278,7 +291,12 @@ def build_stub_class(
return True
def _get_node_info_v1(cls):
return info.get("schema_v1", {})
node_info = copy.deepcopy(info.get("schema_v1", {}))
relative_python_module = node_info.get("python_module")
if not isinstance(relative_python_module, str) or not relative_python_module:
relative_python_module = f"custom_nodes.{extension.name}"
node_info["python_module"] = relative_python_module
return node_info
def _get_base_class(cls):
return latest_io.ComfyNode
@ -308,6 +326,8 @@ def build_stub_class(
attributes["DEPRECATED"] = info.get("deprecated", False)
attributes["API_NODE"] = info.get("api_node", False)
attributes["NOT_IDEMPOTENT"] = info.get("not_idempotent", False)
attributes["ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
attributes["_ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
attributes["INPUT_IS_LIST"] = info.get("input_is_list", False)
class_name = f"PyIsolate_{node_name}".replace(" ", "_")

1
pyproject.toml
View File

@ -11,6 +11,7 @@ repository = "https://github.com/comfyanonymous/ComfyUI"
documentation = "https://docs.comfy.org/"
[tool.comfy.host]
sandbox_mode = "required"
allow_network = false
writable_paths = ["/dev/shm", "/tmp"]

4
tests/isolation/test_cuda_wheels_and_env_flags.py
View File

@ -69,6 +69,7 @@ flash_attn = "flash-attn-special"
monkeypatch.setattr(
"comfy.isolation.extension_loader.load_host_policy",
lambda base_path: {
"sandbox_mode": "required",
"allow_network": False,
"writable_paths": [],
"readonly_paths": [],
@ -95,6 +96,7 @@ flash_attn = "flash-attn-special"
)
assert len(specs) == 1
assert captured["sandbox_mode"] == "required"
assert captured["cuda_wheels"] == {
"index_url": "https://example.invalid/cuda-wheels/",
"packages": ["flash-attn", "sageattention"],
@ -171,6 +173,7 @@ can_isolate = true
monkeypatch.setattr(
"comfy.isolation.extension_loader.load_host_policy",
lambda base_path: {
"sandbox_mode": "disabled",
"allow_network": False,
"writable_paths": [],
"readonly_paths": [],
@ -196,6 +199,7 @@ can_isolate = true
)
)
assert captured["sandbox_mode"] == "disabled"
assert "cuda_wheels" not in captured

41
tests/isolation/test_host_policy.py
View File

@ -10,6 +10,7 @@ def test_load_host_policy_defaults_when_pyproject_missing(tmp_path):
policy = load_host_policy(tmp_path)
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
assert policy["writable_paths"] == DEFAULT_POLICY["writable_paths"]
assert policy["readonly_paths"] == DEFAULT_POLICY["readonly_paths"]
@ -28,6 +29,7 @@ name = "ComfyUI"
)
policy = load_host_policy(tmp_path)
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
assert policy["whitelist"] == {}
@ -39,6 +41,7 @@ def test_load_host_policy_reads_values(tmp_path):
tmp_path / "pyproject.toml",
"""
[tool.comfy.host]
sandbox_mode = "disabled"
allow_network = true
writable_paths = ["/tmp/a", "/tmp/b"]
readonly_paths = ["/opt/readonly"]
@ -49,6 +52,7 @@ ExampleNode = "*"
)
policy = load_host_policy(tmp_path)
assert policy["sandbox_mode"] == "disabled"
assert policy["allow_network"] is True
assert policy["writable_paths"] == ["/tmp/a", "/tmp/b"]
assert policy["readonly_paths"] == ["/opt/readonly"]
@ -70,3 +74,40 @@ whitelist = ["bad"]
policy = load_host_policy(tmp_path)
assert policy["allow_network"] is True
assert policy["whitelist"] == DEFAULT_POLICY["whitelist"]
def test_load_host_policy_ignores_invalid_sandbox_mode(tmp_path):
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
_write_pyproject(
tmp_path / "pyproject.toml",
"""
[tool.comfy.host]
sandbox_mode = "surprise"
""".strip(),
)
policy = load_host_policy(tmp_path)
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
def test_load_host_policy_uses_env_override_path(tmp_path, monkeypatch):
from comfy.isolation.host_policy import load_host_policy
override_path = tmp_path / "host_policy_override.toml"
_write_pyproject(
override_path,
"""
[tool.comfy.host]
sandbox_mode = "disabled"
allow_network = true
""".strip(),
)
monkeypatch.setenv("COMFY_HOST_POLICY_PATH", str(override_path))
policy = load_host_policy(tmp_path / "missing-root")
assert policy["sandbox_mode"] == "disabled"
assert policy["allow_network"] is True