Files
ttulttul-ComfyUI-Modal/tests/test_runtime_environment.py

481 lines
17 KiB
Python

"""Tests for pinned remote environment and deployment fingerprints."""
from __future__ import annotations
import shlex
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import pytest
def _runtime_settings(**overrides: Any) -> SimpleNamespace:
"""Return minimal settings accepted by the runtime identity builder."""
values: dict[str, Any] = {
"app_name": "comfy-modal-sync",
"buffer_containers": None,
"enable_gpu_memory_snapshot": True,
"enable_memory_snapshot": True,
"execution_timeout_seconds": 3600,
"invocation_dict_name": "comfy-modal-sync-invocations",
"invocation_result_inline_max_bytes": 4 * 1024 * 1024,
"interrupt_dict_name": "comfy-modal-sync-interrupts",
"llm_vllm_execution_mode": "auto",
"max_containers": None,
"min_containers": 0,
"modal_gpu": "A100",
"modal_secret_name": "comfy",
"node_output_cache_dict_name": "comfy-modal-sync-node-cache",
"remote_storage_root": "/storage",
"scaledown_window_seconds": 600,
"session_bridge_dict_name": "comfy-modal-sync-session-bridges",
"bridge_inline_max_bytes": 4 * 1024 * 1024,
"snapshot_profile_dict_name": "comfy-modal-sync-snapshot-profiles",
"startup_timeout_seconds": 900,
"stream_event_queue_maxsize": 256,
"sync_index_dict_name": "comfy-modal-sync-sync-index",
"volume_name": "comfy-universal-storage",
}
values.update(overrides)
return SimpleNamespace(**values)
def _runtime_identity(
runtime_environment_module: Any,
repo_root: Path,
comfyui_root: Path,
custom_nodes_dir: Path | None = None,
**settings_overrides: Any,
) -> Any:
"""Build one test identity from small temporary source trees."""
return runtime_environment_module.build_remote_runtime_identity(
repo_root=repo_root,
comfyui_root=comfyui_root,
custom_nodes_dir=custom_nodes_dir,
settings=_runtime_settings(**settings_overrides),
)
def test_remote_environment_is_fully_pinned(runtime_environment_module: Any) -> None:
"""The supported Python, ComfyUI, and CUDA environments should not float."""
apt_packages = runtime_environment_module.remote_apt_packages()
runtime_packages = runtime_environment_module.remote_runtime_packages()
torch_packages = runtime_environment_module.remote_torch_packages()
assert runtime_environment_module.DEFAULT_MODAL_GPU == "RTX-PRO-6000"
assert runtime_environment_module.REMOTE_APP_PROTOCOL_VERSION == 10
assert runtime_environment_module.REMOTE_PYTHON_VERSION == "3.13"
assert apt_packages == ("build-essential", "libgl1", "libglib2.0-0")
assert all("==" in requirement for requirement in runtime_packages)
assert all("==" in requirement for requirement in torch_packages)
assert len(runtime_packages) == len(set(runtime_packages))
assert "blake3==1.0.9" in runtime_packages
assert "comfy-angle==0.1.0" in runtime_packages
assert "comfy-aimdo==0.4.13" in runtime_packages
assert "comfy-kitchen==0.2.31" in runtime_packages
assert "pyopengl==3.1.10" in runtime_packages
assert "pypdf==6.16.1" in runtime_packages
assert "numpy==2.3.5" in runtime_packages
assert "av==18.1.0" in runtime_packages
assert "safetensors==0.8.0" in runtime_packages
assert "huggingface-hub==1.28.0" in runtime_packages
assert "hf-xet==1.6.0" in runtime_packages
assert "transformers==5.15.0" in runtime_packages
assert "simpleeval==1.0.7" in runtime_packages
def test_remote_video_stack_is_pinned_and_validated(
runtime_environment_module: Any,
) -> None:
"""Remote images must expose the PyAV enum imported by current ComfyUI."""
command = runtime_environment_module.remote_runtime_validation_command()
assert runtime_environment_module.REMOTE_AV_SPEC == "av==18.1.0"
assert command.startswith("python -c ")
assert "ColorPrimaries" in command
assert "Expected PyAV" in command
assert "18.1.0" in command
def test_remote_compiler_toolchain_is_validated(
runtime_environment_module: Any,
) -> None:
"""Triton workers must retain an explicit compiler after image construction."""
command = runtime_environment_module.remote_compiler_validation_command()
assert "/usr/bin/gcc" in command
assert "/usr/bin/g++" in command
assert "gcc --version" in command
def test_remote_huggingface_stack_is_pinned_and_validated(
runtime_environment_module: Any,
) -> None:
"""Model staging should install and import the explicit Xet accelerator."""
assert runtime_environment_module.remote_huggingface_packages() == (
"huggingface-hub==1.28.0",
"hf-xet==1.6.0",
)
command = runtime_environment_module.remote_huggingface_validation_command()
assert "import hf_xet, huggingface_hub" in command
assert "actual_xet=metadata.version" in command
assert "hf-xet" in command
@pytest.mark.parametrize(
"modal_gpu",
(
"T4",
"L4",
"A10",
"L40S",
"A100",
"A100-40GB",
"A100-80GB",
"RTX-PRO-6000",
"H100",
"H100!",
"H200",
"B200",
"B200+",
"B300",
"b300:2",
),
)
def test_all_supported_modal_gpus_use_vllm_compatible_cuda_130_build(
runtime_environment_module: Any,
modal_gpu: str,
) -> None:
"""Every selectable GPU should use one vLLM-compatible CUDA 13 stack."""
build = runtime_environment_module.select_remote_torch_build(modal_gpu)
assert build.cuda_version == "13.0"
assert build.install_layers == (
runtime_environment_module.RemoteTorchInstallLayer(
index_url="https://download.pytorch.org/whl/cu130",
packages=("torch==2.13.0", "torchvision==0.28.0"),
),
runtime_environment_module.RemoteTorchInstallLayer(
index_url="https://download.pytorch.org/whl/cpu",
packages=("torchaudio==2.11.0+cpu",),
extra_options="--no-deps",
),
)
assert runtime_environment_module.remote_torch_packages(modal_gpu) == (
"torch==2.13.0",
"torchvision==0.28.0",
"torchaudio==2.11.0+cpu",
)
assert runtime_environment_module.remote_accelerator_packages(modal_gpu) == (
runtime_environment_module.REMOTE_VLLM_WHEEL_URL,
)
def test_remote_build_validation_imports_complete_torch_and_vllm_stacks(
runtime_environment_module: Any,
) -> None:
"""The image build should fail before deployment when native packages disagree."""
build = runtime_environment_module.select_remote_torch_build("B300")
validation_command = build.validation_command()
validation_script = shlex.split(validation_command)[2]
accelerator_command = runtime_environment_module.remote_accelerator_validation_command(
"RTX-PRO-6000"
)
accelerator_script = shlex.split(accelerator_command)[2]
assert validation_command.startswith("python -c ")
assert "import torch, torchaudio, torchvision" in validation_script
assert "expected_cuda='13.0'" in validation_script
assert "import cv2, numpy, torch, vllm" in accelerator_script
assert "expected_vllm='0.27.1'" in accelerator_script
assert "expected_numpy='2.3.5'" in accelerator_script
assert "expected_opencv='4.13.0.92'" in accelerator_script
def test_empty_modal_gpu_cannot_select_torch_build(runtime_environment_module: Any) -> None:
"""An empty GPU specification should fail before constructing a Modal image."""
with pytest.raises(ValueError, match="cannot be empty"):
runtime_environment_module.select_remote_torch_build(" ")
def test_runtime_identity_changes_with_source_and_runtime_options(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Code or deployment-option changes should produce a different fingerprint."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
repo_root.mkdir()
comfyui_root.mkdir()
repo_source = repo_root / "worker.py"
comfyui_source = comfyui_root / "execution.py"
repo_source.write_text("VALUE = 1\n", encoding="utf-8")
comfyui_source.write_text("VALUE = 1\n", encoding="utf-8")
baseline = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
assert baseline == _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
repo_source.write_text("VALUE = 2\n", encoding="utf-8")
source_changed = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
option_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
modal_gpu="L40S",
)
secret_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
modal_secret_name="workflow-credentials",
)
vllm_mode_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
llm_vllm_execution_mode="throughput",
)
assert source_changed.fingerprint != baseline.fingerprint
assert (
runtime_environment_module.remote_runtime_dependency_fingerprint(
source_changed
)
== runtime_environment_module.remote_runtime_dependency_fingerprint(
baseline
)
)
assert option_changed.fingerprint != source_changed.fingerprint
assert secret_changed.fingerprint != source_changed.fingerprint
assert vllm_mode_changed.fingerprint != baseline.fingerprint
assert baseline.manifest["runtime_options"]["llm_vllm_execution_mode"] == "auto"
assert secret_changed.manifest["runtime_options"]["modal_secret_name"] == (
"workflow-credentials"
)
def test_runtime_identity_tracks_curated_llm_profile_registry(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Changing a pinned model profile should force a remote deployment rebuild."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
repo_root.mkdir()
comfyui_root.mkdir()
(repo_root / "worker.py").write_text("VALUE = 1\n", encoding="utf-8")
profile_path = repo_root / "llm_profiles.json"
profile_path.write_text('{"profiles":[{"id":"one"}]}\n', encoding="utf-8")
baseline = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
profile_path.write_text('{"profiles":[{"id":"two"}]}\n', encoding="utf-8")
changed = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
assert changed.fingerprint != baseline.fingerprint
def test_runtime_identity_records_system_packages(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""The deployment manifest should record system libraries installed in the image."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
repo_root.mkdir()
comfyui_root.mkdir()
identity = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
assert identity.manifest["apt_packages"] == [
"build-essential",
"libgl1",
"libglib2.0-0",
]
def test_runtime_identity_records_universal_torch_and_vllm_build(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""The deployment manifest should describe the universal native stack."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
repo_root.mkdir()
comfyui_root.mkdir()
identity = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
modal_gpu="B300",
)
assert identity.manifest["torch_build"] == {
"cuda_version": "13.0",
"install_layers": [
{
"index_url": "https://download.pytorch.org/whl/cu130",
"packages": ["torch==2.13.0", "torchvision==0.28.0"],
"extra_options": "",
},
{
"index_url": "https://download.pytorch.org/whl/cpu",
"packages": ["torchaudio==2.11.0+cpu"],
"extra_options": "--no-deps",
},
],
}
assert identity.manifest["torch_packages"] == [
"torch==2.13.0",
"torchvision==0.28.0",
"torchaudio==2.11.0+cpu",
]
assert identity.manifest["accelerator_packages"] == [
runtime_environment_module.REMOTE_VLLM_WHEEL_URL
]
def test_runtime_identity_tracks_custom_node_requirements_but_ignores_payload_source(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Image dependencies should affect identity while request-bundled source should not."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
custom_nodes_dir = comfyui_root / "custom_nodes"
custom_node_dir = custom_nodes_dir / "example"
repo_root.mkdir()
comfyui_root.mkdir()
custom_node_dir.mkdir(parents=True)
(repo_root / "worker.py").write_text("VALUE = 1\n", encoding="utf-8")
(comfyui_root / "execution.py").write_text("VALUE = 1\n", encoding="utf-8")
requirements_path = custom_node_dir / "requirements.txt"
source_path = custom_node_dir / "node.py"
requirements_path.write_text("diffusers==0.37.0\n", encoding="utf-8")
source_path.write_text("VALUE = 1\n", encoding="utf-8")
baseline = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
custom_nodes_dir,
)
source_path.write_text("VALUE = 2\n", encoding="utf-8")
source_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
custom_nodes_dir,
)
requirements_path.write_text("diffusers==0.38.0\n", encoding="utf-8")
dependency_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
custom_nodes_dir,
)
assert source_changed.fingerprint == baseline.fingerprint
assert dependency_changed.fingerprint != baseline.fingerprint
assert (
runtime_environment_module.remote_runtime_dependency_fingerprint(
source_changed
)
== runtime_environment_module.remote_runtime_dependency_fingerprint(
baseline
)
)
assert (
runtime_environment_module.remote_runtime_dependency_fingerprint(
dependency_changed
)
!= runtime_environment_module.remote_runtime_dependency_fingerprint(
baseline
)
)
def test_custom_node_packages_exclude_deployment_owned_runtime_conflicts(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Custom-node pins must not override the immutable Python and CUDA stack."""
custom_nodes_dir = tmp_path / "custom_nodes"
custom_node_dir = custom_nodes_dir / "example"
custom_node_dir.mkdir(parents=True)
(custom_node_dir / "requirements.txt").write_text(
"\n".join(
(
"numpy<2",
"opencv-python>=4.7",
"torch>=2.1",
"Pillow>=10",
"triton>=3.3",
"vllm>=0.9",
"diffusers==0.38.0",
)
),
encoding="utf-8",
)
packages = runtime_environment_module.custom_node_runtime_packages(
custom_nodes_dir
)
assert packages == ("diffusers==0.38.0",)
def test_runtime_identity_ignores_comfyui_directories_outside_image_context(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Unshipped ComfyUI scratch source should not trigger a full app rebuild."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
runtime_package = comfyui_root / "comfy"
scratch_package = comfyui_root / "False"
repo_root.mkdir()
runtime_package.mkdir(parents=True)
scratch_package.mkdir(parents=True)
(repo_root / "worker.py").write_text("VALUE = 1\n", encoding="utf-8")
runtime_source = runtime_package / "model_management.py"
scratch_source = scratch_package / "developer_copy.py"
runtime_source.write_text("VALUE = 1\n", encoding="utf-8")
scratch_source.write_text("VALUE = 1\n", encoding="utf-8")
baseline = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
scratch_source.write_text("VALUE = 2\n", encoding="utf-8")
scratch_changed = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
runtime_source.write_text("VALUE = 2\n", encoding="utf-8")
runtime_changed = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
assert scratch_changed.fingerprint == baseline.fingerprint
assert runtime_changed.fingerprint != scratch_changed.fingerprint
def test_runtime_identity_tracks_comfyui_package_data(
runtime_environment_module: Any,
tmp_path: Path,
) -> None:
"""Tokenizer and configuration assets shipped in the image must affect identity."""
repo_root = tmp_path / "repo"
comfyui_root = tmp_path / "ComfyUI"
tokenizer_directory = comfyui_root / "comfy" / "text_encoders" / "qwen25_tokenizer"
repo_root.mkdir()
tokenizer_directory.mkdir(parents=True)
(repo_root / "worker.py").write_text("VALUE = 1\n", encoding="utf-8")
tokenizer_path = tokenizer_directory / "vocab.json"
tokenizer_path.write_text('{"one": 1}\n', encoding="utf-8")
baseline = _runtime_identity(runtime_environment_module, repo_root, comfyui_root)
tokenizer_path.write_text('{"two": 2}\n', encoding="utf-8")
tokenizer_changed = _runtime_identity(
runtime_environment_module,
repo_root,
comfyui_root,
)
assert tokenizer_changed.fingerprint != baseline.fingerprint