nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Queue SSH work behind occupied capacity

Browse Source
This commit is contained in:
Ken Simpson
2026-08-25 14:27:07 -07:00
parent 017af31c84
commit eaca5e1ea0
5 changed files with 244 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -206,7 +206,7 @@ For offline development and CI, a local API simulator stands in for the live ser
### Cost- And Capability-Aware Scheduling
Before comparing costs, the planner resolves the model files referenced by each remote component and derives conservative memory floors: the largest resident model plus LoRA/adapter/ControlNet weights, a weight-overhead margin, and a fixed reserve for activations and runtime state. Candidates are checked against live free VRAM and RAM where a probe reports them, so an undersized host is rejected before any model upload. Enabled SSH hosts are re-probed immediately before every placement decision, and an unreachable host is simply excluded, letting another provider serve as the fallback under automatic policy.
Before comparing costs, the planner resolves the model files referenced by each remote component and derives conservative memory floors: the largest resident model plus LoRA/adapter/ControlNet weights, a weight-overhead margin, and a fixed reserve for activations and runtime state. Candidates are checked against live free VRAM and RAM where a probe reports them, so an undersized host is rejected before any model upload. When an earlier ComfyUI prompt already owns a configured SSH host, a later queued prompt is admitted against that host's real total capacity rather than its transient free memory; prompt-level execution serialization ensures the earlier workflow releases the capacity first. Unrelated GPU use receives no such exemption. Enabled SSH hosts are re-probed immediately before every placement decision, and an unreachable host is simply excluded, letting another provider serve as the fallback under automatic policy.
Provider boundaries are retained when values can be serialized between workers. When a `MODEL`, `CLIP`, `VAE`, `NOISE`, or `SAMPLER` producer is safely reproducible from loader or literal inputs, the planner replicates that dependency on the downstream worker instead of transferring the in-memory object or forcing unrelated phases onto the same provider. Non-reproducible runtime objects remain hard co-location constraints, and the strictest required provider applies to their complete component.
api_intercept.py
+87 -2
View File
@@ -1320,6 +1320,7 @@ def _plan_component_execution(
status_callback: SetupStatusCallback | None = None,
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> ComponentExecutionPlan:
"""Plan through connected configurations or synthesize the legacy behavior."""
try:
@@ -1345,6 +1346,7 @@ def _plan_component_execution(
status_callback=status_callback,
environment_status_callback=environment_status_callback,
plan_callback=plan_callback,
occupied_environment_ids=occupied_environment_ids,
)
@@ -1358,6 +1360,7 @@ def _plan_configured_component_execution(
status_callback: SetupStatusCallback | None = None,
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> ComponentExecutionPlan:
"""Resolve and prepare a capacity-aware plan from connected configurations."""
configurations_by_id = {
@@ -1387,6 +1390,7 @@ def _plan_configured_component_execution(
ssh_hosts_by_id=ssh_hosts_by_id,
vast_service=vast_service,
vast_unavailable_reason=vast_unavailable_reason,
occupied_environment_ids=occupied_environment_ids,
)
_apply_historical_execution_estimates(
components=components,
@@ -1411,6 +1415,15 @@ def _plan_configured_component_execution(
except NoCompatibleExecutionEnvironmentError as exc:
raise ModalPromptValidationError(str(exc)) from exc
for component_id, assignment in tuple(assignments.items()):
if assignment.environment_id not in occupied_environment_ids:
continue
assignments[component_id] = replace(
assignment,
reasons=assignment.reasons
+ ("capacity available after the earlier workflow completes",),
)
safe_configurations = _safe_remote_configuration_payload(configuration_set)
if plan_callback is not None:
plan_callback(
@@ -1656,6 +1669,7 @@ def _configured_candidate_environments(
ssh_hosts_by_id: Mapping[str, SshHostConfig],
vast_service: VastService | None,
vast_unavailable_reason: str | None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> tuple[
dict[str, list[EnvironmentSchedulingState]],
dict[tuple[str, str], VastProfileQuote],
@@ -1674,6 +1688,7 @@ def _configured_candidate_environments(
ssh_hosts_by_id=ssh_hosts_by_id,
vast_service=vast_service,
vast_unavailable_reason=vast_unavailable_reason,
occupied_environment_ids=occupied_environment_ids,
)
if state is not None:
candidates[component_id].append(state)
@@ -1725,17 +1740,26 @@ def _configured_candidate_environment(
ssh_hosts_by_id: Mapping[str, SshHostConfig],
vast_service: VastService | None,
vast_unavailable_reason: str | None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> tuple[EnvironmentSchedulingState | None, VastProfileQuote | None]:
"""Resolve one configuration as a candidate for one component."""
if isinstance(configuration, ModalRemoteConfiguration):
return _modal_configuration_environment_state(configuration, settings), None
if isinstance(configuration, SshRemoteConfiguration):
host = ssh_hosts_by_id[configuration.configuration_id]
return replace(
state = replace(
host.scheduling_state(),
configuration_id=configuration.configuration_id,
display_name=configuration.display_name,
), None
)
if state.environment_id in occupied_environment_ids:
state = _maximum_capacity_state(state)
logger.info(
"Planning queued work against reclaimable SSH capacity "
"environment=%s because an earlier workflow owns that host.",
state.environment_id,
)
return state, None
if not isinstance(configuration, VastRemoteConfiguration):
raise TypeError(
f"Unsupported remote configuration type {type(configuration).__name__}."
@@ -7434,6 +7458,7 @@ def rewrite_prompt_for_modal(
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
cancellation_check: Callable[[], bool] | None = None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> tuple[dict[str, Any], RewriteSummary]:
"""Rewrite connected remote components into Modal proxy nodes."""
if cancellation_check is not None and cancellation_check():
@@ -7496,6 +7521,7 @@ def rewrite_prompt_for_modal(
status_callback=status_callback,
environment_status_callback=environment_status_callback,
plan_callback=plan_callback,
occupied_environment_ids=occupied_environment_ids,
)
assignments_by_component_id = execution_plan.assignments
if execution_plan.configuration_set is not None:
@@ -7972,6 +7998,7 @@ async def rewrite_prompt_for_modal_async(
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
cancellation_check: Callable[[], bool] | None = None,
occupied_environment_ids: frozenset[str] = frozenset(),
) -> tuple[dict[str, Any], RewriteSummary]:
"""Prepare one Modal prompt without blocking ComfyUI's event loop."""
return await asyncio.to_thread(
@@ -7986,6 +8013,7 @@ async def rewrite_prompt_for_modal_async(
environment_status_callback=environment_status_callback,
plan_callback=plan_callback,
cancellation_check=cancellation_check,
occupied_environment_ids=occupied_environment_ids,
)
@@ -8478,6 +8506,52 @@ def _queue_item_prompt_ids(items: Iterable[Any]) -> set[str]:
}
def _queued_ssh_environment_ids(
prompt_server: Any,
*,
excluding_prompt_id: str | None = None,
) -> frozenset[str]:
"""Return SSH environments reserved by prompts already in ComfyUI's queue."""
prompt_queue = getattr(prompt_server, "prompt_queue", None)
get_current_queue = getattr(prompt_queue, "get_current_queue", None)
if not callable(get_current_queue):
return frozenset()
queue_state = get_current_queue()
if not isinstance(queue_state, (list, tuple)) or len(queue_state) != 2:
return frozenset()
running, queued = queue_state
queue_items = [
item
for collection in (running, queued)
if isinstance(collection, (list, tuple))
for item in collection
]
environment_ids: set[str] = set()
for item in queue_items:
if not isinstance(item, (list, tuple)) or len(item) <= 3:
continue
prompt_id = str(item[1]) if len(item) > 1 else ""
if excluding_prompt_id is not None and prompt_id == excluding_prompt_id:
continue
extra_data = item[3]
if not isinstance(extra_data, Mapping):
continue
remote_execution = extra_data.get("remote_execution")
if not isinstance(remote_execution, Mapping):
continue
assignments = remote_execution.get("assignments")
if not isinstance(assignments, Mapping):
continue
for assignment in assignments.values():
if not isinstance(assignment, Mapping):
continue
provider = str(assignment.get("provider") or "").strip().lower()
environment_id = str(assignment.get("environment_id") or "").strip()
if provider == ExecutionProvider.SSH_DOCKER.value and environment_id:
environment_ids.add(environment_id)
return frozenset(environment_ids)
def _set_remote_preparation(
prompt_server: Any,
*,
@@ -9237,6 +9311,16 @@ def setup_modal_queue_route(
if "prompt" in json_data:
emit_setup_status("Preparing remote workflow")
rewrite_started_at = time.perf_counter()
occupied_environment_ids = _queued_ssh_environment_ids(
prompt_server,
excluding_prompt_id=prompt_id,
)
if occupied_environment_ids:
logger.info(
"Queued workflow may reuse SSH capacity after earlier "
"prompts finish environments=%s.",
sorted(occupied_environment_ids),
)
rewritten_prompt, summary = await rewrite_prompt_for_modal_async(
prompt=json_data["prompt"],
workflow=workflow,
@@ -9247,6 +9331,7 @@ def setup_modal_queue_route(
environment_status_callback=emit_environment_setup_status,
plan_callback=emit_execution_plan,
cancellation_check=preparation_cancellation.is_set,
occupied_environment_ids=occupied_environment_ids,
)
selected_modal_gpus = _selected_modal_gpus(
summary,
docs/LEARNINGS.md
+4
View File
@@ -574,3 +574,7 @@
- A Vast contract can disappear after returning a usable SSH endpoint. Revalidate the exact provider contract immediately before each worker-readiness connection; otherwise a deleted direct endpoint looks like a 15-minute SSH authentication failure. Remove the stale registry record, exclude its offer, and cold-start one bounded replacement so transient host loss does not fail the workflow or return it to the same vanished capacity.
- Asset synchronization is not worker readiness. An SSH environment can finish custom-node and prompt-asset preparation while its fingerprinted image still needs to build and its container still needs to start. End queue-time setup with a provider-specific intermediate state, then thread the execution-time runtime lifecycle through the same environment-scoped status stream; reserve **Ready** for a worker that has passed its actual readiness probe.
- Importing Triton successfully during image construction does not prove the runtime can execute Triton-backed nodes. Triton lazily compiles native launchers on first use, so slim worker images must retain a C/C++ compiler toolchain, expose deterministic `CC`/`CXX` paths, and validate the compiler binaries while building the image instead of discovering their absence inside a billable workflow.
## 2026-08-25: Queue admission must distinguish occupied capacity from undersized capacity
- Live free VRAM is the correct admission signal for an otherwise idle SSH host, but it is misleading when the same ComfyUI queue already has a workflow assigned to that host. Preserve each queued prompt's provider assignment in queue metadata, identify the exact SSH environments owned by earlier prompts, and admit later serialized work against those hosts' reserve-adjusted total capacity. Do not generalize this to every busy GPU: unrelated workloads still need the conservative live-free-memory check, and physically undersized hardware must remain incompatible.
tests/test_api_intercept.py
+66
View File
@@ -6488,6 +6488,72 @@ def test_remote_preparation_bridge_exposes_work_to_all_queue_views(
assert queue_update_counts == [1, 1]
def test_queued_ssh_environment_ids_reads_earlier_prompt_assignments(
api_intercept_module: Any,
) -> None:
"""Queue-time planning should recognize SSH hosts owned by earlier prompts."""
def assignment(provider: str, environment_id: str) -> dict[str, str]:
"""Return minimal serialized placement metadata."""
return {
"provider": provider,
"environment_id": environment_id,
}
prompt_server = SimpleNamespace(
prompt_queue=SimpleNamespace(
get_current_queue=lambda: (
[
(
1,
"prompt-running",
{},
{
"remote_execution": {
"assignments": {
"257": assignment("ssh_docker", "lambda")
}
}
},
)
],
[
(
2,
"prompt-current",
{},
{
"remote_execution": {
"assignments": {
"300": assignment("ssh_docker", "ignored")
}
}
},
),
(
3,
"prompt-modal",
{},
{
"remote_execution": {
"assignments": {
"400": assignment("modal", "modal:H100")
}
}
},
),
],
)
)
)
environment_ids = api_intercept_module._queued_ssh_environment_ids(
prompt_server,
excluding_prompt_id="prompt-current",
)
assert environment_ids == frozenset({"lambda"})
def test_remote_preparation_bridge_clears_failed_submission(
api_intercept_module: Any,
) -> None:
tests/test_configured_remote_planning.py
+86
View File
@@ -5,6 +5,8 @@ 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."""
@@ -338,6 +340,90 @@ def test_workflow_ssh_configuration_is_probed_and_carried_into_execution_metadat
assert metadata["ssh_host_config"]["health"] == "unknown"
def test_queued_workflow_uses_nameplate_capacity_of_occupied_ssh_host(
api_intercept_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 = api_intercept_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)
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
)
def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload(
api_intercept_module: Any,
remote_configuration_nodes_module: Any,