1175 lines
40 KiB
Python
1175 lines
40 KiB
Python
"""Tests for the execution scheduling boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from api_intercept_test_support import * # noqa: F401,F403
|
|
|
|
def test_auto_placement_selects_every_eligible_prompt_node(
|
|
api_intercept_module: Any,
|
|
remote_graph_analysis_module: Any,
|
|
) -> None:
|
|
"""Workflow auto placement should not require per-node remote toggles."""
|
|
prompt = {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
|
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
|
|
"3": {"class_type": "ModalEndpointChat", "inputs": {}},
|
|
"4": {"class_type": "VastAILeaseConfiguration", "inputs": {}},
|
|
"5": {"class_type": "ModalRemoteConfiguration", "inputs": {}},
|
|
"6": {"class_type": "VastRemoteConfiguration", "inputs": {}},
|
|
"7": {"class_type": "SshRemoteConfiguration", "inputs": {}},
|
|
"8": {"class_type": "RemoteExecutionConfigurator", "inputs": {}},
|
|
}
|
|
workflow = {
|
|
"extra": {
|
|
"remote_execution": {
|
|
"policy": "automatic",
|
|
"auto_place": True,
|
|
}
|
|
},
|
|
"nodes": [],
|
|
}
|
|
|
|
selected = remote_graph_analysis_module.requested_remote_node_ids(
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=SimpleNamespace(marker_property="is_modal_remote"),
|
|
)
|
|
|
|
assert selected == {"1", "2"}
|
|
|
|
def test_disabled_auto_placement_preserves_explicit_markers(
|
|
api_intercept_module: Any,
|
|
remote_graph_analysis_module: Any,
|
|
) -> None:
|
|
"""Manual workflows should continue honoring the existing node property."""
|
|
prompt = {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
|
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
|
|
}
|
|
workflow = {
|
|
"extra": {"remote_execution": {"policy": "self_hosted", "auto_place": False}},
|
|
"nodes": [
|
|
{"id": 1, "properties": {"is_modal_remote": False}},
|
|
{"id": 2, "properties": {"is_modal_remote": True}},
|
|
],
|
|
}
|
|
|
|
selected = remote_graph_analysis_module.requested_remote_node_ids(
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=SimpleNamespace(marker_property="is_modal_remote"),
|
|
)
|
|
|
|
assert selected == {"2"}
|
|
|
|
def test_scheduler_refreshes_recent_ssh_capabilities_before_placement(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
remote_hosts_module: Any,
|
|
execution_environments_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Automatic placement must not trust even recently persisted free VRAM."""
|
|
registry = remote_hosts_module.RemoteHostRegistry.for_user_directory(tmp_path)
|
|
previous_capabilities = execution_environments_module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=64 * 1024**3,
|
|
available_ram_bytes=60 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
probed_at_epoch=2_000_000_000.0,
|
|
)
|
|
registry.replace_hosts(
|
|
[
|
|
remote_hosts_module.SshHostConfig(
|
|
environment_id="freshened",
|
|
display_name="Freshened",
|
|
ssh_target="freshened",
|
|
capabilities=previous_capabilities,
|
|
health=execution_environments_module.EnvironmentHealth.READY,
|
|
)
|
|
]
|
|
)
|
|
capabilities = execution_environments_module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=64 * 1024**3,
|
|
available_ram_bytes=60 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
probed_at_epoch=1_700_000_000.0,
|
|
)
|
|
|
|
class FakeController:
|
|
"""Return current capability data for one configured host."""
|
|
|
|
def __init__(self, host: Any) -> None:
|
|
"""Retain the probed host."""
|
|
self.host = host
|
|
|
|
def probe_capabilities(self) -> Any:
|
|
"""Return current capabilities."""
|
|
return capabilities
|
|
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_ssh_host_registry",
|
|
lambda _settings: registry,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"SshDockerController",
|
|
FakeController,
|
|
)
|
|
|
|
hosts = execution_scheduling_module._schedulable_ssh_hosts(SimpleNamespace())
|
|
|
|
assert hosts[0].health.value == "ready"
|
|
assert hosts[0].capabilities == capabilities
|
|
|
|
def test_modal_only_policy_rejects_ssh_only_llm_backend(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A provider-specific backend must not be dispatched to Modal by policy."""
|
|
component = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["257"],
|
|
representative_node_id="257",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["257"],
|
|
contains_output_node=False,
|
|
)
|
|
|
|
with pytest.raises(
|
|
api_intercept_module.ModalPromptValidationError,
|
|
match="Modal-only execution cannot run SSH-only component",
|
|
):
|
|
execution_scheduling_module._plan_component_execution_assignments(
|
|
components=[component],
|
|
prompt={
|
|
"257": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {
|
|
"model_profile": (
|
|
"huihui-qwen3.8-27b-abliterated-q2-k-gguf"
|
|
)
|
|
},
|
|
}
|
|
},
|
|
workflow={"extra": {"remote_execution": {"policy": "modal"}}},
|
|
settings=SimpleNamespace(
|
|
modal_gpu="H200",
|
|
max_containers=1,
|
|
local_storage_root=tmp_path,
|
|
),
|
|
)
|
|
|
|
def test_automatic_policy_assigns_component_to_lower_cost_ready_host(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
remote_hosts_module: Any,
|
|
execution_environments_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Automatic policy should compare a compatible SSH host with Modal."""
|
|
module = execution_environments_module
|
|
capabilities = module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=64 * 1024**3,
|
|
available_ram_bytes=60 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(module.GpuCapability("GPU-1", "GPU", 80 * 1024**3),),
|
|
probed_at_epoch=2_000_000_000.0,
|
|
)
|
|
host = remote_hosts_module.SshHostConfig(
|
|
environment_id="cheap-host",
|
|
display_name="Cheap host",
|
|
ssh_target="cheap-host",
|
|
cost_usd_per_second=0.0001,
|
|
capabilities=capabilities,
|
|
health=module.EnvironmentHealth.READY,
|
|
)
|
|
component = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["1"],
|
|
representative_node_id="1",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["1"],
|
|
contains_output_node=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_schedulable_ssh_hosts",
|
|
lambda _settings: (host,),
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_execution_history",
|
|
lambda _settings: None,
|
|
)
|
|
|
|
assignments = execution_scheduling_module._plan_component_execution_assignments(
|
|
components=[component],
|
|
prompt={"1": {"class_type": "KSampler", "inputs": {"steps": 20}}},
|
|
workflow={
|
|
"extra": {
|
|
"remote_execution": {
|
|
"policy": "automatic",
|
|
"auto_place": True,
|
|
}
|
|
}
|
|
},
|
|
settings=SimpleNamespace(modal_gpu="RTX-PRO-6000", max_containers=1),
|
|
)
|
|
|
|
assert assignments["1"].provider.value == "ssh_docker"
|
|
assert assignments["1"].environment_id == "cheap-host"
|
|
|
|
def test_planner_recycles_idle_ssh_worker_before_cost_ranking(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
remote_hosts_module: Any,
|
|
execution_environments_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Resident managed-worker VRAM should not hide a cheaper host's capacity."""
|
|
module = execution_environments_module
|
|
gib = 1024**3
|
|
occupied_capabilities = module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=64 * gib,
|
|
available_ram_bytes=48 * gib,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(
|
|
module.GpuCapability(
|
|
"GPU-4090",
|
|
"RTX 4090",
|
|
24 * gib,
|
|
free_vram_bytes=11 * gib,
|
|
),
|
|
),
|
|
)
|
|
reclaimed_capabilities = replace(
|
|
occupied_capabilities,
|
|
gpus=(
|
|
replace(
|
|
occupied_capabilities.gpus[0],
|
|
free_vram_bytes=23 * gib,
|
|
),
|
|
),
|
|
)
|
|
host = remote_hosts_module.SshHostConfig(
|
|
environment_id="lambda",
|
|
display_name="Lambda 4090",
|
|
ssh_target="lambda",
|
|
cost_usd_per_second=0.0,
|
|
capabilities=occupied_capabilities,
|
|
health=module.EnvironmentHealth.READY,
|
|
)
|
|
component = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["1"],
|
|
representative_node_id="1",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["1"],
|
|
contains_output_node=False,
|
|
)
|
|
lifecycle_calls: list[str] = []
|
|
|
|
class FakeController:
|
|
"""Expose one idle worker and its post-reclaim capability probe."""
|
|
|
|
def __init__(self, configured_host: Any) -> None:
|
|
"""Retain the selected host."""
|
|
assert configured_host.environment_id == "lambda"
|
|
|
|
def remove_idle_managed_workers(self) -> tuple[str, ...]:
|
|
"""Report one safely recycled warm worker."""
|
|
lifecycle_calls.append("remove")
|
|
return ("comfy-remote-lambda-fingerprint-w0",)
|
|
|
|
def probe_capabilities(self) -> Any:
|
|
"""Return free VRAM after the managed container stopped."""
|
|
lifecycle_calls.append("probe")
|
|
return reclaimed_capabilities
|
|
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_schedulable_ssh_hosts",
|
|
lambda _settings: (host,),
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"SshDockerController",
|
|
FakeController,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_ssh_host_registry",
|
|
lambda _settings: None,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_execution_history",
|
|
lambda _settings: None,
|
|
)
|
|
|
|
assignments = execution_scheduling_module._plan_component_execution_assignments(
|
|
components=[component],
|
|
prompt={"1": {"class_type": "KSampler", "inputs": {"steps": 20}}},
|
|
workflow={
|
|
"extra": {
|
|
"remote_execution": {
|
|
"policy": "automatic",
|
|
"auto_place": True,
|
|
"minimum_vram_gb": 16,
|
|
}
|
|
}
|
|
},
|
|
settings=SimpleNamespace(modal_gpu="RTX-PRO-6000", max_containers=1),
|
|
)
|
|
|
|
assert lifecycle_calls == ["remove", "probe"]
|
|
assert assignments["1"].provider is module.ExecutionProvider.SSH_DOCKER
|
|
assert assignments["1"].environment_id == "lambda"
|
|
|
|
def test_planner_does_not_recycle_for_equal_cost_tie_break(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_environments_module: Any,
|
|
) -> None:
|
|
"""A lexical tie alone must not discard a compatible worker's warm cache."""
|
|
module = execution_environments_module
|
|
actual = module.ExecutionAssignment(
|
|
environment_id="ready-host",
|
|
provider=module.ExecutionProvider.SSH_DOCKER,
|
|
predicted_cost_usd=0.0,
|
|
predicted_completion_seconds=60.0,
|
|
)
|
|
optimistic = replace(actual, environment_id="idle-host")
|
|
|
|
assert not execution_scheduling_module._reclaim_improves_assignment(
|
|
optimistic,
|
|
actual,
|
|
module.ComponentResourceRequirements(),
|
|
)
|
|
|
|
def test_automatic_policy_rejects_zero_cost_host_for_oversized_model(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
remote_hosts_module: Any,
|
|
execution_environments_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Model weights plus headroom must exclude a cheap host before cost ranking."""
|
|
module = execution_environments_module
|
|
model_path = tmp_path / "minimax_h3_bf16.safetensors"
|
|
with model_path.open("wb") as model_file:
|
|
model_file.truncate(66 * 1024**3)
|
|
|
|
capabilities = module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=64 * 1024**3,
|
|
available_ram_bytes=60 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(module.GpuCapability("GPU-4090", "RTX 4090", 24 * 1024**3),),
|
|
probed_at_epoch=2_000_000_000.0,
|
|
)
|
|
host = remote_hosts_module.SshHostConfig(
|
|
environment_id="lambda",
|
|
display_name="Lambda 4090",
|
|
ssh_target="lambda",
|
|
cost_usd_per_second=0.0,
|
|
capabilities=capabilities,
|
|
health=module.EnvironmentHealth.READY,
|
|
)
|
|
component = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["6"],
|
|
representative_node_id="6",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["6"],
|
|
contains_output_node=False,
|
|
)
|
|
prompt = {
|
|
"6": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": str(model_path), "weight_dtype": "default"},
|
|
}
|
|
}
|
|
workflow = {
|
|
"extra": {
|
|
"remote_execution": {
|
|
"policy": "automatic",
|
|
"auto_place": True,
|
|
}
|
|
}
|
|
}
|
|
settings = SimpleNamespace(
|
|
modal_gpu="H200",
|
|
max_containers=1,
|
|
comfyui_root=None,
|
|
)
|
|
preferences = module.WorkflowExecutionPreferences.from_workflow(workflow)
|
|
estimate = execution_scheduling_module._component_memory_estimate(
|
|
component,
|
|
prompt,
|
|
preferences,
|
|
settings,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_schedulable_ssh_hosts",
|
|
lambda _settings: (host,),
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_execution_history",
|
|
lambda _settings: None,
|
|
)
|
|
|
|
assignments = execution_scheduling_module._plan_component_execution_assignments(
|
|
components=[component],
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=settings,
|
|
)
|
|
|
|
assert estimate.model_asset_count == 1
|
|
assert estimate.largest_model_bytes == 66 * 1024**3
|
|
assert 24 * 1024**3 < estimate.minimum_vram_bytes < 96 * 1024**3
|
|
assert estimate.minimum_ram_bytes == 70 * 1024**3
|
|
assert assignments["6"].provider is module.ExecutionProvider.MODAL
|
|
assert assignments["6"].environment_id == "modal:H200"
|
|
|
|
def test_planner_resolves_hugging_face_metadata_before_cost_ranking(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_resource_estimates_module: Any,
|
|
execution_environments_module: Any,
|
|
remote_hosts_module: Any,
|
|
llm_resolver_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Raw Hugging Face IDs must expose their VRAM floor before SSH placement."""
|
|
module = execution_environments_module
|
|
capabilities = module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=128 * 1024**3,
|
|
available_ram_bytes=120 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="28.0.0",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(module.GpuCapability("GPU-4090", "RTX 4090", 24 * 1024**3),),
|
|
probed_at_epoch=2_000_000_000.0,
|
|
)
|
|
host = remote_hosts_module.SshHostConfig(
|
|
environment_id="lambda",
|
|
display_name="Lambda 4090",
|
|
ssh_target="lambda",
|
|
cost_usd_per_second=0.0,
|
|
capabilities=capabilities,
|
|
health=module.EnvironmentHealth.READY,
|
|
)
|
|
component = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["257"],
|
|
representative_node_id="257",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["257"],
|
|
contains_output_node=False,
|
|
)
|
|
prompt = {
|
|
"257": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {"model_profile": "owner/large-model"},
|
|
}
|
|
}
|
|
workflow = {
|
|
"extra": {
|
|
"remote_execution": {
|
|
"policy": "automatic",
|
|
"auto_place": True,
|
|
}
|
|
}
|
|
}
|
|
profile = SimpleNamespace(
|
|
profile_id="hf-" + "d" * 64,
|
|
artifact_bytes=55_563_006_216,
|
|
estimated_vram_gb=67.9,
|
|
)
|
|
resolved_references: list[str] = []
|
|
|
|
def resolve(model_reference: str, storage_root: Path) -> Any:
|
|
"""Return deterministic metadata without downloading model weights."""
|
|
assert storage_root == tmp_path
|
|
resolved_references.append(model_reference)
|
|
return SimpleNamespace(profile=profile)
|
|
|
|
settings = SimpleNamespace(
|
|
modal_gpu="H200",
|
|
max_containers=1,
|
|
comfyui_root=None,
|
|
local_storage_root=tmp_path,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_resource_estimates_module,
|
|
"resolve_model_profile",
|
|
resolve,
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_schedulable_ssh_hosts",
|
|
lambda _settings: (host,),
|
|
)
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_execution_history",
|
|
lambda _settings: None,
|
|
)
|
|
|
|
assignments = execution_scheduling_module._plan_component_execution_assignments(
|
|
components=[component],
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=settings,
|
|
)
|
|
|
|
assert resolved_references == ["owner/large-model"]
|
|
assert assignments["257"].provider is module.ExecutionProvider.MODAL
|
|
assert assignments["257"].environment_id == "modal:H200"
|
|
assert "requires at least 67.90 GiB GPU VRAM" in assignments["257"].reasons
|
|
|
|
def test_planner_keeps_unmarked_llm_local_between_remote_text_nodes(
|
|
api_intercept_module: Any,
|
|
remote_graph_analysis_module: Any,
|
|
) -> None:
|
|
"""Transportable text boundaries must not absorb a local LLM into Modal."""
|
|
prompt = {
|
|
"1": {"class_type": "RemoteTextSource", "inputs": {}},
|
|
"2": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {"prompt": ["1", 0], "model_profile": "owner/model"},
|
|
},
|
|
"3": {
|
|
"class_type": "RemoteTextConsumer",
|
|
"inputs": {"prompt": ["2", 0]},
|
|
},
|
|
}
|
|
workflow = {
|
|
"nodes": [
|
|
{"id": 1, "properties": {"is_modal_remote": True}},
|
|
{"id": 2, "properties": {"is_modal_remote": False}},
|
|
{"id": 3, "properties": {"is_modal_remote": True}},
|
|
]
|
|
}
|
|
fake_nodes_module = SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"RemoteTextSource": _FakeTextNode,
|
|
"ModalLLM": _FakeTextNode,
|
|
"RemoteTextConsumer": _FakeTextNode,
|
|
},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
|
|
analysis = remote_graph_analysis_module.analyze_remote_node_selection(
|
|
prompt,
|
|
workflow,
|
|
[],
|
|
settings=SimpleNamespace(marker_property="is_modal_remote"),
|
|
nodes_module=fake_nodes_module,
|
|
)
|
|
|
|
assert analysis.resolved_remote_node_ids == ["1", "3"]
|
|
assert analysis.sandwiched_local_node_ids == ["2"]
|
|
assert analysis.added_node_ids == []
|
|
|
|
def test_configured_non_modal_plan_omits_modal_status_gpu(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""Vast-only and SSH-only plans must not activate Modal status polling."""
|
|
summary = api_intercept_module.RewriteSummary(
|
|
execution_assignments_by_representative={
|
|
"vast-component": api_intercept_module.ExecutionAssignment(
|
|
environment_id="vast:instance-1",
|
|
provider=api_intercept_module.ExecutionProvider.VAST,
|
|
predicted_cost_usd=0.01,
|
|
predicted_completion_seconds=20.0,
|
|
configuration_id="vast-config",
|
|
),
|
|
"ssh-component": api_intercept_module.ExecutionAssignment(
|
|
environment_id="ssh-host",
|
|
provider=api_intercept_module.ExecutionProvider.SSH_DOCKER,
|
|
predicted_cost_usd=0.0,
|
|
predicted_completion_seconds=15.0,
|
|
configuration_id="ssh-config",
|
|
),
|
|
},
|
|
remote_configurations=[
|
|
{
|
|
"configuration_id": "vast-config",
|
|
"provider": "vast",
|
|
"display_name": "Vast pool",
|
|
},
|
|
{
|
|
"configuration_id": "ssh-config",
|
|
"provider": "ssh_docker",
|
|
"display_name": "SSH host",
|
|
},
|
|
],
|
|
)
|
|
|
|
assert api_intercept_module._selected_modal_gpus(summary, "B300") == []
|
|
assert api_intercept_module._prompt_uses_remote_execution_configurator(
|
|
{
|
|
"99": {
|
|
"class_type": "RemoteExecutionConfigurator",
|
|
"inputs": {},
|
|
}
|
|
}
|
|
)
|
|
|
|
def test_configured_modal_plan_reports_only_selected_modal_gpus(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""Status metadata should come from selected configurations, not legacy GPU state."""
|
|
summary = api_intercept_module.RewriteSummary(
|
|
execution_assignments_by_representative={
|
|
"modal-component": api_intercept_module.ExecutionAssignment(
|
|
environment_id="modal:modal-config:H200",
|
|
provider=api_intercept_module.ExecutionProvider.MODAL,
|
|
predicted_cost_usd=0.01,
|
|
predicted_completion_seconds=10.0,
|
|
configuration_id="modal-config",
|
|
)
|
|
},
|
|
remote_configurations=[
|
|
{
|
|
"configuration_id": "modal-config",
|
|
"provider": "modal",
|
|
"display_name": "Modal H200",
|
|
"gpu_type": "H200",
|
|
}
|
|
],
|
|
)
|
|
|
|
assert api_intercept_module._selected_modal_gpus(summary, "B300") == ["H200"]
|
|
|
|
def test_planner_resolved_llm_profile_is_attached_to_matching_payload(
|
|
api_intercept_module: Any,
|
|
prompt_payload_metadata_module: Any,
|
|
llm_profiles_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""The execution payload should carry metadata already resolved by planning."""
|
|
profile = llm_profiles_module.get_llm_profile("smolvlm2-2.2b-instruct")
|
|
payload = {
|
|
"component_id": "llm",
|
|
"subgraph_prompt": {
|
|
"llm": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {"model_profile": profile.profile_id},
|
|
}
|
|
},
|
|
}
|
|
|
|
prompt_payload_metadata_module._attach_resolved_llm_profiles(
|
|
payload,
|
|
{profile.profile_id: profile},
|
|
SimpleNamespace(local_storage_root=tmp_path),
|
|
)
|
|
|
|
entry = payload["resolved_llm_profiles"][profile.profile_id]
|
|
assert entry["profile"] == profile.to_mapping()
|
|
assert entry["security_scan_complete"] is True
|
|
|
|
def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target(
|
|
api_intercept_module: Any,
|
|
prompt_affinity_planning_module: Any,
|
|
) -> None:
|
|
"""Each proxy should prepare only its nearest reachable future worker group."""
|
|
rewritten_prompt = {
|
|
"spec-a": {
|
|
"class_type": "ModalProxy",
|
|
"inputs": {
|
|
"original_node_data": {
|
|
"component_id": "spec-a",
|
|
"prompt_id": "prompt-spec",
|
|
"modal_gpu": "RTX-PRO-6000",
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:RTX-PRO-6000",
|
|
"remote_worker_affinity_group": "llm",
|
|
"subgraph_prompt": {"1": {"class_type": "ModalLLM", "inputs": {}}},
|
|
}
|
|
},
|
|
},
|
|
"spec-b": {
|
|
"class_type": "ModalProxy",
|
|
"inputs": {
|
|
"upstream": ["spec-local", 0],
|
|
"original_node_data": {
|
|
"component_id": "spec-b",
|
|
"prompt_id": "prompt-spec",
|
|
"modal_gpu": "RTX-PRO-6000",
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:RTX-PRO-6000",
|
|
"remote_worker_affinity_group": "comfy",
|
|
"remote_local_gap_pool": True,
|
|
"snapshot_profile_key": "loader-profile:abc",
|
|
"subgraph_prompt": {
|
|
"2": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": "video-model.safetensors"},
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
"spec-local": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["spec-a", 0]},
|
|
},
|
|
"spec-c": {
|
|
"class_type": "ModalProxy",
|
|
"inputs": {
|
|
"upstream": ["spec-b", 0],
|
|
"original_node_data": {
|
|
"component_id": "spec-c",
|
|
"prompt_id": "prompt-spec",
|
|
"modal_gpu": "RTX-PRO-6000",
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:RTX-PRO-6000",
|
|
"remote_worker_affinity_group": "llm",
|
|
"subgraph_prompt": {"3": {"class_type": "ModalLLM", "inputs": {}}},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
prompt_affinity_planning_module._configure_speculative_affinity_prewarm_payloads(
|
|
rewritten_prompt=rewritten_prompt,
|
|
execution_stages=[["spec-a", "spec-b"], ["spec-c"]],
|
|
)
|
|
|
|
first_payload = api_intercept_module.registered_proxy_execution_payload(
|
|
"spec-a", rewritten_prompt["spec-a"]["inputs"]["original_node_data"]
|
|
)
|
|
second_payload = api_intercept_module.registered_proxy_execution_payload(
|
|
"spec-b", rewritten_prompt["spec-b"]["inputs"]["original_node_data"]
|
|
)
|
|
third_payload = api_intercept_module.registered_proxy_execution_payload(
|
|
"spec-c", rewritten_prompt["spec-c"]["inputs"]["original_node_data"]
|
|
)
|
|
|
|
first_target = first_payload["speculative_remote_prewarm_target"]
|
|
second_target = second_payload["speculative_remote_prewarm_target"]
|
|
assert first_target["component_id"] == "spec-b"
|
|
assert first_target["remote_worker_affinity_group"] == "comfy"
|
|
assert first_target["snapshot_profile_key"] == "loader-profile:abc"
|
|
assert second_target["component_id"] == "spec-c"
|
|
assert second_target["remote_worker_affinity_group"] == "llm"
|
|
assert "speculative_remote_prewarm_target" not in third_payload
|
|
|
|
def test_planner_does_not_bridge_local_gap_keepalive_across_providers(
|
|
api_intercept_module: Any,
|
|
prompt_affinity_planning_module: Any,
|
|
) -> None:
|
|
"""A Modal producer must not retain a slot for an SSH continuation."""
|
|
rewritten_prompt = {
|
|
"modal-producer": {
|
|
"class_type": "ModalProxy",
|
|
"inputs": {
|
|
"original_node_data": {
|
|
"component_id": "modal-producer",
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:H200",
|
|
}
|
|
},
|
|
},
|
|
"local-gap": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["modal-producer", 0]},
|
|
},
|
|
"ssh-consumer": {
|
|
"class_type": "ModalProxy",
|
|
"inputs": {
|
|
"source": ["local-gap", 0],
|
|
"original_node_data": {
|
|
"component_id": "ssh-consumer",
|
|
"execution_provider": "ssh_docker",
|
|
"execution_environment_id": "lambda",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
prompt_affinity_planning_module._configure_local_gap_keepalive_payloads(
|
|
rewritten_prompt=rewritten_prompt,
|
|
remote_component_ids=["modal-producer", "ssh-consumer"],
|
|
sandwiched_local_node_ids={"local-gap"},
|
|
)
|
|
|
|
for component_id in ("modal-producer", "ssh-consumer"):
|
|
payload = api_intercept_module.registered_proxy_execution_payload(
|
|
component_id,
|
|
rewritten_prompt[component_id]["inputs"]["original_node_data"],
|
|
)
|
|
assert "remote_local_gap_pool" not in payload
|
|
assert "keepalive_after_remote_component" not in payload
|
|
assert "stop_local_gap_keepalive_before_remote_component" not in payload
|
|
|
|
def test_delete_modal_cache_dicts_deletes_configured_dicts(
|
|
api_intercept_module: Any,
|
|
modal_admin_ops_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Deleting Modal caches should clear and delete every configured cache Dict."""
|
|
|
|
class FakeNotFoundError(Exception):
|
|
"""Stand-in for Modal object misses."""
|
|
|
|
class FakeDictObject:
|
|
"""Minimal Modal Dict object used only for existence checks."""
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Store the configured Dict name."""
|
|
self.name = name
|
|
|
|
def delete(self, name: str) -> None:
|
|
"""Fail if the deprecated instance deletion path is used."""
|
|
raise AssertionError(f"deprecated Dict.delete path used for {name}")
|
|
|
|
class FakeDictObjects:
|
|
"""Minimal Modal Dict manager namespace."""
|
|
|
|
@staticmethod
|
|
def delete(name: str, allow_missing: bool = False) -> None:
|
|
"""Record a manager delete call."""
|
|
assert allow_missing is True
|
|
deleted.append(name)
|
|
|
|
class FakeDict:
|
|
"""Minimal Modal Dict namespace."""
|
|
|
|
objects = FakeDictObjects()
|
|
|
|
@staticmethod
|
|
def from_name(name: str, create_if_missing: bool = False) -> FakeDictObject:
|
|
"""Return fake Dict objects, except for one missing cache."""
|
|
assert create_if_missing is False
|
|
if name == "app-interrupts":
|
|
raise FakeNotFoundError(name)
|
|
return FakeDictObject(name)
|
|
|
|
class FakeModal:
|
|
"""Minimal Modal SDK double."""
|
|
|
|
exception = SimpleNamespace(NotFoundError=FakeNotFoundError)
|
|
Dict = FakeDict
|
|
|
|
settings = settings_module.ModalSyncSettings(
|
|
app_name="app",
|
|
auto_deploy=True,
|
|
allow_ephemeral_fallback=False,
|
|
enable_memory_snapshot=True,
|
|
enable_gpu_memory_snapshot=False,
|
|
execution_mode="remote",
|
|
sync_custom_nodes=False,
|
|
volume_name="volume",
|
|
route_path="/modal/queue_prompt",
|
|
marker_property="is_modal_remote",
|
|
local_storage_root=tmp_path / "storage",
|
|
remote_storage_root="/storage",
|
|
custom_nodes_archive_name="custom_nodes_bundle.zip",
|
|
comfyui_root=None,
|
|
custom_nodes_dir=None,
|
|
interrupt_dict_name="app-interrupts",
|
|
node_output_cache_dict_name="app-node-cache",
|
|
session_bridge_dict_name="app-session-bridges",
|
|
sync_index_dict_name="app-sync-index",
|
|
snapshot_profile_dict_name="app-snapshot-profiles",
|
|
)
|
|
deleted: list[str] = []
|
|
monkeypatch.setattr(modal_admin_ops_module, "modal", FakeModal)
|
|
|
|
result = asyncio.run(api_intercept_module.delete_modal_cache_dicts(settings))
|
|
|
|
assert result == {
|
|
"deleted": [
|
|
"app-node-cache",
|
|
"app-session-bridges",
|
|
"app-sync-index",
|
|
"app-snapshot-profiles",
|
|
],
|
|
"skipped": ["app-interrupts"],
|
|
}
|
|
assert deleted == result["deleted"]
|
|
|
|
def test_delete_modal_volume_deletes_configured_volume(
|
|
api_intercept_module: Any,
|
|
modal_admin_ops_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Deleting the Modal volume should target only the configured volume name."""
|
|
|
|
class FakeVolumeObject:
|
|
"""Minimal Modal Volume object used only for existence checks."""
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Store the configured Volume name."""
|
|
self.name = name
|
|
|
|
def delete(self, name: str) -> None:
|
|
"""Fail if the deprecated instance deletion path is used."""
|
|
raise AssertionError(f"deprecated Volume.delete path used for {name}")
|
|
|
|
class FakeVolumeObjects:
|
|
"""Minimal Modal Volume manager namespace."""
|
|
|
|
@staticmethod
|
|
def delete(name: str, allow_missing: bool = False) -> None:
|
|
"""Record a manager delete call."""
|
|
assert allow_missing is True
|
|
deleted.append(name)
|
|
|
|
class FakeVolume:
|
|
"""Minimal Modal Volume namespace."""
|
|
|
|
objects = FakeVolumeObjects()
|
|
|
|
@staticmethod
|
|
def from_name(name: str, create_if_missing: bool = False) -> FakeVolumeObject:
|
|
"""Return a fake Volume object."""
|
|
assert create_if_missing is False
|
|
return FakeVolumeObject(name)
|
|
|
|
class FakeModal:
|
|
"""Minimal Modal SDK double."""
|
|
|
|
exception = SimpleNamespace()
|
|
Volume = FakeVolume
|
|
|
|
settings = settings_module.ModalSyncSettings(
|
|
app_name="app",
|
|
auto_deploy=True,
|
|
allow_ephemeral_fallback=False,
|
|
enable_memory_snapshot=True,
|
|
enable_gpu_memory_snapshot=False,
|
|
execution_mode="remote",
|
|
sync_custom_nodes=False,
|
|
volume_name="configured-volume",
|
|
route_path="/modal/queue_prompt",
|
|
marker_property="is_modal_remote",
|
|
local_storage_root=tmp_path / "storage",
|
|
remote_storage_root="/storage",
|
|
custom_nodes_archive_name="custom_nodes_bundle.zip",
|
|
comfyui_root=None,
|
|
custom_nodes_dir=None,
|
|
)
|
|
deleted: list[str] = []
|
|
monkeypatch.setattr(modal_admin_ops_module, "modal", FakeModal)
|
|
|
|
result = asyncio.run(api_intercept_module.delete_modal_volume(settings))
|
|
|
|
assert result == {"deleted": ["configured-volume"], "skipped": []}
|
|
assert deleted == ["configured-volume"]
|
|
|
|
def test_selected_vast_capacity_streams_setup_status(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
vast_models_module: Any,
|
|
) -> None:
|
|
"""Queue-time Vast acquisition should stream provider progress into the UI."""
|
|
profile = vast_models_module.VastResourceProfile(
|
|
profile_id="vast-config",
|
|
profile_name="Vast-pool",
|
|
maximum_instances=1,
|
|
)
|
|
configuration = api_intercept_module.VastRemoteConfiguration(
|
|
configuration_id="vast-config",
|
|
display_name="Vast-pool",
|
|
profile=profile,
|
|
)
|
|
configuration_set = api_intercept_module.RemoteConfigurationSet(
|
|
configurations=(configuration,)
|
|
)
|
|
assignments = {
|
|
"component-1": api_intercept_module.ExecutionAssignment(
|
|
environment_id=profile.environment_id,
|
|
provider=api_intercept_module.ExecutionProvider.VAST,
|
|
predicted_cost_usd=0.02,
|
|
predicted_completion_seconds=30.0,
|
|
configuration_id="vast-config",
|
|
capacity_slot_index=0,
|
|
)
|
|
}
|
|
requirements = {
|
|
"component-1": api_intercept_module.ComponentResourceRequirements(
|
|
estimated_execution_seconds=30.0
|
|
)
|
|
}
|
|
quote = SimpleNamespace(
|
|
profile=profile,
|
|
predicted_incremental_cost_usd=0.02,
|
|
)
|
|
status_events: list[tuple[str, int | None, int | None]] = []
|
|
environment_status_events: list[
|
|
tuple[str, str, int | None, int | None]
|
|
] = []
|
|
|
|
class FakeVastService:
|
|
"""Emit representative readiness phases without renting an instance."""
|
|
|
|
def acquire_sync(
|
|
self,
|
|
selected_quote: Any,
|
|
*,
|
|
slot: int,
|
|
status_callback: Any,
|
|
) -> Any:
|
|
"""Emit image and runtime phases before returning a fake lease."""
|
|
assert selected_quote is quote
|
|
assert slot == 0
|
|
status_callback("Vast.ai instance 42 is downloading the worker image")
|
|
status_callback("Initializing Vast.ai worker")
|
|
return SimpleNamespace(
|
|
environment_id="vast:vast-config:42",
|
|
idle_retention_seconds=3600.0,
|
|
)
|
|
|
|
leases = execution_scheduling_module._prepare_selected_vast_capacity(
|
|
assignments=assignments,
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements,
|
|
vast_quotes={("component-1", "vast-config"): quote},
|
|
vast_service=FakeVastService(),
|
|
status_callback=lambda message, current, total: status_events.append(
|
|
(message, current, total)
|
|
),
|
|
environment_status_callback=(
|
|
lambda environment_id, message, current, total: (
|
|
environment_status_events.append(
|
|
(environment_id, message, current, total)
|
|
)
|
|
)
|
|
),
|
|
)
|
|
|
|
assert list(leases) == ["vast:vast-config:42"]
|
|
assert status_events == [
|
|
("Acquiring Vast.ai capacity 1 of 1", 0, 1),
|
|
("Vast.ai instance 42 is downloading the worker image", 0, 1),
|
|
("Initializing Vast.ai worker", 0, 1),
|
|
("Vast.ai capacity 1 of 1 is ready", 1, 1),
|
|
]
|
|
assert environment_status_events == [
|
|
("vast:vast-config", "Acquiring Vast.ai capacity 1 of 1", None, None),
|
|
(
|
|
"vast:vast-config",
|
|
"Vast.ai instance 42 is downloading the worker image",
|
|
None,
|
|
None,
|
|
),
|
|
("vast:vast-config", "Initializing Vast.ai worker", None, None),
|
|
(
|
|
"vast:vast-config:42",
|
|
"Vast.ai worker ready; preparing remote assets next",
|
|
None,
|
|
None,
|
|
),
|
|
]
|
|
assert assignments["component-1"].environment_id == "vast:vast-config:42"
|
|
|
|
def test_selected_vast_capacity_preserves_intentional_cancellation(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
vast_models_module: Any,
|
|
) -> None:
|
|
"""Cancelling capacity acquisition should not become a provider failure."""
|
|
profile = vast_models_module.VastResourceProfile(
|
|
profile_id="vast-config",
|
|
profile_name="Vast-pool",
|
|
maximum_instances=1,
|
|
)
|
|
configuration = api_intercept_module.VastRemoteConfiguration(
|
|
configuration_id="vast-config",
|
|
display_name="Vast-pool",
|
|
profile=profile,
|
|
)
|
|
assignment = api_intercept_module.ExecutionAssignment(
|
|
environment_id=profile.environment_id,
|
|
provider=api_intercept_module.ExecutionProvider.VAST,
|
|
predicted_cost_usd=0.02,
|
|
predicted_completion_seconds=30.0,
|
|
configuration_id="vast-config",
|
|
capacity_slot_index=0,
|
|
)
|
|
quote = SimpleNamespace(
|
|
profile=profile,
|
|
predicted_incremental_cost_usd=0.02,
|
|
)
|
|
|
|
class CancelledVastService:
|
|
"""Stop acquisition as though the user cancelled queue preparation."""
|
|
|
|
def acquire_sync(self, selected_quote: Any, *, slot: int) -> Any:
|
|
"""Raise the prompt-scoped cancellation without a provider failure."""
|
|
assert selected_quote is quote
|
|
assert slot == 0
|
|
raise api_intercept_module.SyncCancelledError(
|
|
"Remote workflow preparation was cancelled."
|
|
)
|
|
|
|
with pytest.raises(
|
|
api_intercept_module.SyncCancelledError,
|
|
match="Remote workflow preparation was cancelled",
|
|
):
|
|
execution_scheduling_module._prepare_selected_vast_capacity(
|
|
assignments={"component-1": assignment},
|
|
configuration_set=api_intercept_module.RemoteConfigurationSet(
|
|
configurations=(configuration,)
|
|
),
|
|
requirements_by_component={
|
|
"component-1": api_intercept_module.ComponentResourceRequirements(
|
|
estimated_execution_seconds=30.0
|
|
)
|
|
},
|
|
vast_quotes={("component-1", "vast-config"): quote},
|
|
vast_service=CancelledVastService(),
|
|
)
|
|
|