"""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, } ]