672 lines
23 KiB
Python
672 lines
23 KiB
Python
"""Tests for capacity planning from connected remote configuration nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
|
|
class _RemoteImageNode:
|
|
"""Minimal remote image producer used by rewrite integration coverage."""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
|
|
class _LocalImageSinkNode:
|
|
"""Minimal local image consumer used by rewrite integration coverage."""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
|
|
def _component(api_module: Any, node_id: str) -> Any:
|
|
"""Return one independent remote component plan."""
|
|
return api_module.RemoteComponentPlan(
|
|
node_ids=[node_id],
|
|
representative_node_id=node_id,
|
|
boundary_inputs=[],
|
|
boundary_outputs=[],
|
|
execute_node_ids=[node_id],
|
|
contains_output_node=False,
|
|
)
|
|
|
|
|
|
def _vast_inputs(name: str, maximum_instances: int) -> dict[str, Any]:
|
|
"""Return minimal valid Vast configuration inputs."""
|
|
return {
|
|
"profile_name": name,
|
|
"allocated_disk_gb": 200.0,
|
|
"idle_retention_hours": 24.0,
|
|
"maximum_instances": maximum_instances,
|
|
}
|
|
|
|
|
|
def _capabilities(module: Any, name: str = "Configured GPU") -> Any:
|
|
"""Return deterministic scheduler capabilities."""
|
|
return module.EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=16,
|
|
total_ram_bytes=128 * 1024**3,
|
|
available_ram_bytes=128 * 1024**3,
|
|
available_disk_bytes=1024**4,
|
|
docker_version="managed",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(
|
|
module.GpuCapability(
|
|
uuid=f"gpu-{name}",
|
|
name=name,
|
|
total_vram_bytes=96 * 1024**3,
|
|
),
|
|
),
|
|
)
|
|
|
|
|
|
def test_connected_modal_configurations_replace_global_single_gpu_choice(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Two Modal pools should both participate even under the legacy default policy."""
|
|
api = execution_scheduling_module
|
|
nodes = remote_configuration_nodes_module
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
monkeypatch.setitem(api._MODAL_GPU_COST_USD_PER_SECOND, "T4", 0.0)
|
|
monkeypatch.setitem(api._MODAL_GPU_COST_USD_PER_SECOND, "H200", 0.0)
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {}},
|
|
"2": {"class_type": "KSampler", "inputs": {}},
|
|
"10": {
|
|
"class_type": nodes.MODAL_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"configuration_name": "modal-t4",
|
|
"gpu_type": "T4",
|
|
"instance_count": 1,
|
|
},
|
|
},
|
|
"11": {
|
|
"class_type": nodes.MODAL_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"configuration_name": "modal-h200",
|
|
"gpu_type": "H200",
|
|
"instance_count": 1,
|
|
},
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {
|
|
"configuration_0": ["10", 0],
|
|
"configuration_1": ["11", 0],
|
|
},
|
|
},
|
|
}
|
|
|
|
plan_events: list[dict[str, dict[str, Any]]] = []
|
|
plan = api._plan_component_execution(
|
|
components=[_component(api, "1"), _component(api, "2")],
|
|
prompt=prompt,
|
|
workflow={"extra": {"remote_execution": {"policy": "modal"}}},
|
|
settings=settings_module.get_settings(),
|
|
plan_callback=lambda assignments, _configurations: plan_events.append(
|
|
assignments
|
|
),
|
|
)
|
|
|
|
assert plan.configuration_set is not None
|
|
assert {assignment.configuration_id for assignment in plan.assignments.values()} == {
|
|
"10",
|
|
"11",
|
|
}
|
|
assert {assignment.environment_id for assignment in plan.assignments.values()} == {
|
|
"modal:10:T4",
|
|
"modal:11:H200",
|
|
}
|
|
hardware_by_machine = {
|
|
assignment["hardware"]["machine_type"]: assignment["hardware"]
|
|
for assignment in plan_events[-1].values()
|
|
}
|
|
assert hardware_by_machine["T4"] == {
|
|
"machine_type": "T4",
|
|
"gpu_count": 1,
|
|
"gpu_memory_kind": "VRAM",
|
|
"gpu_memory_bytes_per_device": 16 * 1024**3,
|
|
"gpu_memory_bytes_total": 16 * 1024**3,
|
|
"ram_capacity_label": "Provider managed",
|
|
}
|
|
assert hardware_by_machine["H200"] == {
|
|
"machine_type": "H200",
|
|
"gpu_count": 1,
|
|
"gpu_memory_kind": "HBM",
|
|
"gpu_memory_bytes_per_device": 141 * 1024**3,
|
|
"gpu_memory_bytes_total": 141 * 1024**3,
|
|
"ram_capacity_label": "Provider managed",
|
|
}
|
|
|
|
|
|
def test_vast_capacity_is_quoted_globally_and_acquired_only_after_assignment(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_environments_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Independent Vast pools and their capacity slots should survive into leases."""
|
|
api = execution_scheduling_module
|
|
environment_module = execution_environments_module
|
|
nodes = remote_configuration_nodes_module
|
|
events: list[tuple[Any, ...]] = []
|
|
|
|
class FakeVastService:
|
|
"""Expose deterministic quotes and record post-plan acquisitions."""
|
|
|
|
@classmethod
|
|
def from_environment(cls, settings: Any, *, repo_root: Any) -> "FakeVastService":
|
|
"""Construct the fake service without credentials."""
|
|
events.append(("service", settings, repo_root))
|
|
return cls()
|
|
|
|
def quote_best_profile_sync(
|
|
self,
|
|
profiles: Any,
|
|
**requirements: Any,
|
|
) -> Any:
|
|
"""Return a lower incremental price for the broad pool."""
|
|
profile = profiles[0]
|
|
events.append(("quote", profile.profile_name, requirements))
|
|
cost = 0.01 if profile.profile_name == "vast-broad" else 0.02
|
|
return SimpleNamespace(
|
|
profile=profile,
|
|
offer=SimpleNamespace(
|
|
gpu_name=f"{profile.profile_name} GPU",
|
|
num_gpus=1,
|
|
gpu_ram_mb=96 * 1024,
|
|
gpu_total_ram_mb=96 * 1024,
|
|
cpu_ram_mb=192 * 1024,
|
|
),
|
|
existing_lease=None,
|
|
predicted_incremental_cost_usd=cost,
|
|
predicted_execution_seconds=requirements[
|
|
"predicted_execution_seconds"
|
|
],
|
|
)
|
|
|
|
def scheduling_state(self, quote: Any) -> Any:
|
|
"""Expose one virtual pool whose limit is applied by the caller."""
|
|
return environment_module.EnvironmentSchedulingState(
|
|
environment_id=f"vast:{quote.profile.profile_id}",
|
|
provider=environment_module.ExecutionProvider.VAST,
|
|
enabled=True,
|
|
health=environment_module.EnvironmentHealth.READY,
|
|
cost_usd_per_second=(
|
|
quote.predicted_incremental_cost_usd
|
|
/ quote.predicted_execution_seconds
|
|
),
|
|
capabilities=_capabilities(
|
|
environment_module,
|
|
quote.profile.profile_name,
|
|
),
|
|
)
|
|
|
|
def acquire_sync(self, quote: Any, *, slot: int = 0) -> Any:
|
|
"""Create one concrete lease after all assignments have succeeded."""
|
|
profile = quote.profile
|
|
events.append(("acquire", profile.profile_name, slot))
|
|
instance_id = len(
|
|
[event for event in events if event[0] == "acquire"]
|
|
)
|
|
profile_id = (
|
|
profile.profile_id
|
|
if slot == 0
|
|
else f"{profile.profile_id}-slot-{slot}"
|
|
)
|
|
return SimpleNamespace(
|
|
environment_id=f"vast:{profile_id}:{instance_id}",
|
|
instance_id=instance_id,
|
|
profile_id=profile_id,
|
|
profile_name=profile.profile_name,
|
|
gpu_name="Configured GPU",
|
|
gpu_count=1,
|
|
gpu_ram_mb=96 * 1024,
|
|
cpu_ram_mb=192 * 1024,
|
|
hourly_cost_usd=1.0,
|
|
idle_retention_seconds=24 * 3600,
|
|
)
|
|
|
|
monkeypatch.setattr(api, "VastService", FakeVastService)
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
prompt = {
|
|
str(index): {"class_type": "KSampler", "inputs": {}}
|
|
for index in range(1, 4)
|
|
}
|
|
prompt.update(
|
|
{
|
|
"20": {
|
|
"class_type": nodes.VAST_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": _vast_inputs("vast-broad", 2),
|
|
},
|
|
"21": {
|
|
"class_type": nodes.VAST_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": _vast_inputs("vast-premium", 1),
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {
|
|
"configuration_0": ["20", 0],
|
|
"configuration_1": ["21", 0],
|
|
},
|
|
},
|
|
}
|
|
)
|
|
|
|
plan = api._plan_component_execution(
|
|
components=[_component(api, str(index)) for index in range(1, 4)],
|
|
prompt=prompt,
|
|
workflow={"extra": {"remote_execution": {"policy": "automatic"}}},
|
|
settings=settings_module.get_settings(),
|
|
plan_callback=lambda assignments, configurations: events.append(
|
|
("plan", assignments, configurations)
|
|
),
|
|
)
|
|
|
|
acquisitions = [event for event in events if event[0] == "acquire"]
|
|
assert acquisitions == [
|
|
("acquire", "vast-broad", 0),
|
|
("acquire", "vast-broad", 1),
|
|
("acquire", "vast-premium", 0),
|
|
]
|
|
plan_event_index = next(
|
|
index for index, event in enumerate(events) if event[0] == "plan"
|
|
)
|
|
first_acquisition_index = next(
|
|
index for index, event in enumerate(events) if event[0] == "acquire"
|
|
)
|
|
assert plan_event_index < first_acquisition_index
|
|
early_assignments = events[plan_event_index][1]
|
|
assert set(early_assignments) == {"1", "2", "3"}
|
|
assert all(assignment["node_ids"] for assignment in early_assignments.values())
|
|
assert all(
|
|
assignment["hardware"]["gpu_memory_bytes_total"] == 96 * 1024**3
|
|
for assignment in early_assignments.values()
|
|
)
|
|
assert all(
|
|
assignment["hardware"]["ram_bytes"] == 192 * 1024**3
|
|
for assignment in early_assignments.values()
|
|
)
|
|
assert len(events[plan_event_index][2]) == 2
|
|
final_assignments = [event for event in events if event[0] == "plan"][-1][1]
|
|
assert all(
|
|
assignment["hardware"]
|
|
== {
|
|
"machine_type": "Configured GPU",
|
|
"gpu_count": 1,
|
|
"gpu_memory_kind": "VRAM",
|
|
"gpu_memory_bytes_per_device": 96 * 1024**3,
|
|
"gpu_memory_bytes_total": 96 * 1024**3,
|
|
"ram_bytes": 192 * 1024**3,
|
|
}
|
|
for assignment in final_assignments.values()
|
|
)
|
|
assert len(plan.vast_leases_by_environment) == 3
|
|
assert len({assignment.environment_id for assignment in plan.assignments.values()}) == 3
|
|
assert {assignment.configuration_id for assignment in plan.assignments.values()} == {
|
|
"20",
|
|
"21",
|
|
}
|
|
|
|
|
|
def test_workflow_ssh_configuration_is_probed_and_carried_into_execution_metadata(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_environments_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Workflow SSH hosts should not depend on the installation-wide host registry."""
|
|
api = execution_scheduling_module
|
|
nodes = remote_configuration_nodes_module
|
|
|
|
def probe(configuration: Any) -> Any:
|
|
"""Attach deterministic capabilities to the workflow host declaration."""
|
|
return api.replace(
|
|
configuration.host,
|
|
capabilities=_capabilities(execution_environments_module, "SSH GPU"),
|
|
health=execution_environments_module.EnvironmentHealth.READY,
|
|
)
|
|
|
|
monkeypatch.setattr(api, "_probe_workflow_ssh_configuration", probe)
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {}},
|
|
"30": {
|
|
"class_type": nodes.SSH_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"environment_id": "lambda",
|
|
"display_name": "Lambda GPU",
|
|
"ssh_target": "lambda",
|
|
"cost_usd_per_hour": 1.0,
|
|
"maximum_workers": 2,
|
|
},
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {"configuration_0": ["30", 0]},
|
|
},
|
|
}
|
|
|
|
plan_events: list[dict[str, dict[str, Any]]] = []
|
|
plan = api._plan_component_execution(
|
|
components=[_component(api, "1")],
|
|
prompt=prompt,
|
|
workflow={"extra": {"remote_execution": {"policy": "modal"}}},
|
|
settings=settings_module.get_settings(),
|
|
plan_callback=lambda assignments, _configurations: plan_events.append(
|
|
assignments
|
|
),
|
|
)
|
|
assignment = plan.assignments["1"]
|
|
metadata = api._configured_provider_metadata(
|
|
execution_plan=plan,
|
|
assignment=assignment,
|
|
vast_leases_by_environment={},
|
|
)
|
|
|
|
assert assignment.provider is execution_environments_module.ExecutionProvider.SSH_DOCKER
|
|
assert assignment.environment_id == "lambda"
|
|
assert assignment.configuration_id == "lambda"
|
|
assert metadata is not None
|
|
assert metadata["ssh_host_config"]["ssh_target"] == "lambda"
|
|
assert metadata["ssh_host_config"]["capabilities"]["gpus"] == [
|
|
{
|
|
"uuid": "gpu-SSH GPU",
|
|
"name": "SSH GPU",
|
|
"total_vram_bytes": 96 * 1024**3,
|
|
"free_vram_bytes": None,
|
|
"compute_capability": None,
|
|
"driver_version": None,
|
|
}
|
|
]
|
|
assert metadata["ssh_host_config"]["capabilities"][
|
|
"nvidia_container_runtime"
|
|
] is True
|
|
assert metadata["ssh_host_config"]["health"] == "unknown"
|
|
assert plan_events[-1]["1"]["hardware"] == {
|
|
"machine_type": "SSH GPU",
|
|
"gpu_count": 1,
|
|
"gpu_memory_kind": "VRAM",
|
|
"gpu_memory_bytes_per_device": 96 * 1024**3,
|
|
"gpu_memory_bytes_total": 96 * 1024**3,
|
|
"ram_bytes": 128 * 1024**3,
|
|
"ram_available_bytes": 128 * 1024**3,
|
|
}
|
|
|
|
|
|
def test_queued_workflow_uses_nameplate_capacity_of_occupied_ssh_host(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_environments_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""A later prompt should admit capacity that its earlier prompt will release."""
|
|
api = execution_scheduling_module
|
|
environment_module = execution_environments_module
|
|
nodes = remote_configuration_nodes_module
|
|
busy_capabilities = _capabilities(environment_module, "Busy SSH GPU")
|
|
busy_capabilities = api.replace(
|
|
busy_capabilities,
|
|
gpus=tuple(
|
|
api.replace(gpu, free_vram_bytes=11 * 1024**3)
|
|
for gpu in busy_capabilities.gpus
|
|
),
|
|
)
|
|
|
|
def probe(configuration: Any) -> Any:
|
|
"""Report low live VRAM on the workflow-owned SSH host."""
|
|
return api.replace(
|
|
configuration.host,
|
|
capabilities=busy_capabilities,
|
|
health=environment_module.EnvironmentHealth.READY,
|
|
)
|
|
|
|
requirements = {
|
|
"1": environment_module.ComponentResourceRequirements(
|
|
minimum_vram_bytes=16 * 1024**3,
|
|
gpu_required=True,
|
|
architecture="x86_64",
|
|
required_provider=environment_module.ExecutionProvider.SSH_DOCKER,
|
|
)
|
|
}
|
|
monkeypatch.setattr(api, "_probe_workflow_ssh_configuration", probe)
|
|
monkeypatch.setattr(
|
|
api,
|
|
"_configured_component_requirements",
|
|
lambda **_kwargs: requirements,
|
|
)
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
reclaim_attempts: list[str] = []
|
|
monkeypatch.setattr(
|
|
api,
|
|
"_remove_idle_ssh_workers_for_reclaim",
|
|
lambda host: reclaim_attempts.append(host.environment_id) or (),
|
|
)
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {}},
|
|
"30": {
|
|
"class_type": nodes.SSH_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"environment_id": "lambda",
|
|
"display_name": "Lambda GPU",
|
|
"ssh_target": "lambda",
|
|
"maximum_workers": 1,
|
|
},
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {"configuration_0": ["30", 0]},
|
|
},
|
|
}
|
|
|
|
with pytest.raises(api.ModalPromptValidationError, match="insufficient GPU VRAM"):
|
|
api._plan_component_execution(
|
|
components=[_component(api, "1")],
|
|
prompt=prompt,
|
|
workflow=None,
|
|
settings=settings_module.get_settings(),
|
|
)
|
|
|
|
plan = api._plan_component_execution(
|
|
components=[_component(api, "1")],
|
|
prompt=prompt,
|
|
workflow=None,
|
|
settings=settings_module.get_settings(),
|
|
occupied_environment_ids=frozenset({"lambda"}),
|
|
)
|
|
|
|
assignment = plan.assignments["1"]
|
|
assert assignment.environment_id == "lambda"
|
|
assert (
|
|
"capacity available after the earlier workflow completes"
|
|
in assignment.reasons
|
|
)
|
|
assert reclaim_attempts == ["lambda"]
|
|
|
|
|
|
def test_configured_planner_reclaims_idle_worker_memory_before_admission(
|
|
api_intercept_module: Any,
|
|
execution_scheduling_module: Any,
|
|
execution_environments_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Cached models in an idle configured worker should not block a new prompt."""
|
|
api = execution_scheduling_module
|
|
environment_module = execution_environments_module
|
|
nodes = remote_configuration_nodes_module
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {}},
|
|
"30": {
|
|
"class_type": nodes.SSH_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"environment_id": "lambda",
|
|
"display_name": "Lambda GPU",
|
|
"ssh_target": "lambda",
|
|
"maximum_workers": 1,
|
|
},
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {"configuration_0": ["30", 0]},
|
|
},
|
|
}
|
|
full_capabilities = _capabilities(environment_module, "Configured SSH GPU")
|
|
busy_capabilities = api.replace(
|
|
full_capabilities,
|
|
gpus=tuple(
|
|
api.replace(gpu, free_vram_bytes=11 * 1024**3)
|
|
for gpu in full_capabilities.gpus
|
|
),
|
|
)
|
|
probe_count = 0
|
|
|
|
def probe(configuration: Any) -> Any:
|
|
"""Report busy memory first and full memory after worker removal."""
|
|
nonlocal probe_count
|
|
probe_count += 1
|
|
return api.replace(
|
|
configuration.host,
|
|
capabilities=(busy_capabilities if probe_count == 1 else full_capabilities),
|
|
health=environment_module.EnvironmentHealth.READY,
|
|
)
|
|
|
|
removed_hosts: list[str] = []
|
|
monkeypatch.setattr(api, "_probe_workflow_ssh_configuration", probe)
|
|
monkeypatch.setattr(
|
|
api,
|
|
"_remove_idle_ssh_workers_for_reclaim",
|
|
lambda host: removed_hosts.append(host.environment_id) or ("worker-0",),
|
|
)
|
|
monkeypatch.setattr(
|
|
api,
|
|
"_configured_component_requirements",
|
|
lambda **_kwargs: {
|
|
"1": environment_module.ComponentResourceRequirements(
|
|
minimum_vram_bytes=16 * 1024**3,
|
|
gpu_required=True,
|
|
architecture="x86_64",
|
|
required_provider=environment_module.ExecutionProvider.SSH_DOCKER,
|
|
)
|
|
},
|
|
)
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
|
|
plan = api._plan_component_execution(
|
|
components=[_component(api, "1")],
|
|
prompt=prompt,
|
|
workflow=None,
|
|
settings=settings_module.get_settings(),
|
|
)
|
|
|
|
assert plan.assignments["1"].environment_id == "lambda"
|
|
assert removed_hosts == ["lambda"]
|
|
assert probe_count == 2
|
|
|
|
|
|
def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload(
|
|
api_intercept_module: Any,
|
|
remote_configuration_nodes_module: Any,
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Any,
|
|
) -> None:
|
|
"""The configured GPU and count must reach the existing remote executor path."""
|
|
api = api_intercept_module
|
|
nodes = remote_configuration_nodes_module
|
|
custom_nodes_dir = tmp_path / "custom_nodes"
|
|
custom_nodes_dir.mkdir()
|
|
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="local",
|
|
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=custom_nodes_dir,
|
|
)
|
|
sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
|
|
fake_nodes_module = SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"RemoteImage": _RemoteImageNode,
|
|
"LocalSink": _LocalImageSinkNode,
|
|
},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
prompt = {
|
|
"1": {"class_type": "RemoteImage", "inputs": {}},
|
|
"2": {"class_type": "LocalSink", "inputs": {"image": ["1", 0]}},
|
|
"10": {
|
|
"class_type": nodes.MODAL_REMOTE_CONFIGURATION_NODE_ID,
|
|
"inputs": {
|
|
"configuration_name": "modal-t4",
|
|
"gpu_type": "T4",
|
|
"instance_count": 2,
|
|
},
|
|
},
|
|
"99": {
|
|
"class_type": nodes.REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
|
|
"inputs": {"configuration_0": ["10", 0]},
|
|
},
|
|
}
|
|
workflow = {
|
|
"nodes": [
|
|
{"id": 1, "properties": {"is_modal_remote": True}},
|
|
{"id": 2, "properties": {"is_modal_remote": False}},
|
|
]
|
|
}
|
|
monkeypatch.setattr(api, "_execution_history", lambda _settings: None)
|
|
|
|
rewritten, summary = api.rewrite_prompt_for_modal(
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
sync_engine=sync_engine,
|
|
settings=settings,
|
|
nodes_module=fake_nodes_module,
|
|
)
|
|
|
|
payload = rewritten["1"]["inputs"]["original_node_data"]
|
|
assert payload["execution_provider"] == "modal"
|
|
assert payload["execution_environment_id"] == "modal:10:T4"
|
|
assert payload["remote_configuration_id"] == "10"
|
|
assert payload["remote_configuration_name"] == "modal-t4"
|
|
assert payload["modal_gpu"] == "T4"
|
|
assert payload["modal_max_containers"] == 2
|
|
assert summary.remote_configurations == [
|
|
{
|
|
"configuration_id": "10",
|
|
"display_name": "modal-t4",
|
|
"provider": "modal",
|
|
"gpu_type": "T4",
|
|
"instance_count": 2,
|
|
}
|
|
]
|