diff --git a/README.md b/README.md index fffee6b..27680b7 100644 --- a/README.md +++ b/README.md @@ -197,7 +197,7 @@ mlx-dspark resolution is text-only because its Qwen path intentionally omits the Modal-Sync explicitly installs `hf-xet` beside `huggingface-hub` for Apple-local, Modal, and SSH-host staging. Hugging Face automatically uses the Rust Xet transfer backend for Xet-backed repositories; no alternate download command is required. When `HF_TOKEN` or `HUGGING_FACE_HUB_TOKEN` is present, Modal-Sync passes it to repository inspection, config retrieval, and the snapshot download without logging its value. Xet's adaptive concurrency remains enabled by default. Advanced users may set `HF_XET_HIGH_PERFORMANCE=1` in the matching local, Modal-secret, or SSH Docker environment only on high-bandwidth machines with at least 64 GiB of RAM; the default 16 GiB Modal staging worker is intentionally not configured for that mode. -For Modal execution, the local dispatcher finds fixed LLM profiles in the prompt and calls the deployed CPU-only `ModelStager`. It writes the immutable snapshot and completion marker to the shared Modal Volume before the selected GPU worker starts. For SSH execution, the controller runs the same provider-neutral resolver and stager inside the selected persistent worker, writes the profile and weights to that host's named Docker volume, streams progress back to ComfyUI, and rewrites the payload to the immutable generated profile before GPU execution. A linked/dynamic remote `model_profile` is rejected because it cannot be staged before GPU allocation. For local execution, resolution and staging happen directly in the ComfyUI process before MLX allocation. +Before automatic placement, the planner performs metadata-only resolution for raw Hugging Face IDs and applies the generated profile's conservative VRAM and RAM floors, so an inexpensive small GPU cannot win cost ranking for a checkpoint it cannot load. For Modal execution, the local dispatcher then calls the deployed CPU-only `ModelStager`, which writes the immutable snapshot and completion marker to the shared Modal Volume before the selected GPU worker starts. For SSH execution, the controller runs the same provider-neutral resolver and stager inside the selected persistent worker, writes the profile and weights to that host's named Docker volume, streams progress back to ComfyUI, and rewrites the payload to the immutable generated profile before GPU execution. Cancelling during SSH staging stops and later recreates that managed worker while preserving resumable partial files in its named volume. A linked/dynamic remote `model_profile` is rejected because it cannot be staged before GPU allocation. For local execution, resolution and staging happen directly in the ComfyUI process before MLX allocation. Each target and local engine holds a process-global serialized LRU. A cache hit reuses the processor and weights across node and workflow executions. Switching local engines unloads the other engine's residents so separate caches cannot silently double unified-memory use. Before a cold load, the manager asks ComfyUI to release idle managed models, evicts older resident LLMs if necessary, and enforces the target's reserve: `local_reserve_free_memory_gb` defaults to 4 GiB for macOS unified memory, while `reserve_free_vram_gb` remains the Modal GPU control. Local MLX defaults to one resident model and can be changed with `COMFY_MODAL_LOCAL_LLM_MAX_RESIDENT_MODELS`; `COMFY_MODAL_LOCAL_LLM_RESERVE_FREE_GB` supplies the programmatic default when the runner is called without a node value. All backends pass generated token IDs through the same model-aware reasoning abstraction. diff --git a/api_intercept.py b/api_intercept.py index b516978..e4498a8 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -486,17 +486,66 @@ def _component_minimum_vram_bytes( ).minimum_vram_bytes -def _component_profile_minimum_vram_bytes( +def _prompt_llm_model_references(prompt: Mapping[str, Any]) -> tuple[str, ...]: + """Return every fixed Modal LLM model reference in a prompt.""" + references: set[str] = set() + for prompt_node in prompt.values(): + if not isinstance(prompt_node, Mapping): + continue + if str(prompt_node.get("class_type") or "") != "ModalLLM": + continue + inputs = prompt_node.get("inputs") + if not isinstance(inputs, Mapping): + continue + model_reference = inputs.get("model_profile") + if isinstance(model_reference, str) and model_reference.strip(): + references.add(model_reference.strip()) + return tuple(sorted(references)) + + +def _resolve_prompt_llm_profiles( + prompt: Mapping[str, Any], + settings: ModalSyncSettings, +) -> dict[str, Any]: + """Resolve LLM metadata before environment admission and cost ranking.""" + from .llm_profiles import get_llm_profile + from .llm_resolver import resolve_model_profile + + storage_root = Path( + getattr(settings, "local_storage_root", "/tmp/comfyui-modal-sync-storage") + ) + profiles: dict[str, Any] = {} + for model_reference in _prompt_llm_model_references(prompt): + try: + profile = get_llm_profile(model_reference, storage_root=storage_root) + except ValueError as profile_error: + if model_reference.startswith("hf-"): + raise ModalPromptValidationError(str(profile_error)) from profile_error + try: + profile = resolve_model_profile(model_reference, storage_root).profile + except ValueError as resolution_error: + raise ModalPromptValidationError(str(resolution_error)) from resolution_error + profiles[model_reference] = profile + logger.info( + "Resolved planner LLM profile model=%s profile=%s weights_gib=%.2f " + "estimated_vram_gib=%.2f.", + model_reference, + profile.profile_id, + profile.artifact_bytes / 1024**3, + profile.estimated_vram_gb, + ) + return profiles + + +def _component_profile_memory_estimate( component: RemoteComponentPlan, prompt: Mapping[str, Any], preferences: WorkflowExecutionPreferences, -) -> int: - """Estimate a component VRAM floor from workflow and LLM profile metadata.""" + resolved_profiles: Mapping[str, Any], +) -> ComponentMemoryEstimate: + """Estimate a component's RAM and VRAM floors from resolved LLM profiles.""" minimum_vram_bytes = preferences.minimum_vram_bytes - try: - from .llm_profiles import get_llm_profile - except ImportError: - return minimum_vram_bytes + profiles: dict[str, Any] = {} for node_id in component.node_ids: prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): @@ -509,15 +558,23 @@ def _component_profile_minimum_vram_bytes( model_profile = inputs.get("model_profile") if not isinstance(model_profile, str) or not model_profile.strip(): continue - try: - profile = get_llm_profile(model_profile.strip()) - except (KeyError, ValueError): + profile = resolved_profiles.get(model_profile.strip()) + if profile is None: continue + profiles[profile.profile_id] = profile minimum_vram_bytes = max( minimum_vram_bytes, int(max(0.0, profile.estimated_vram_gb) * 1024**3), ) - return minimum_vram_bytes + artifact_sizes = [int(profile.artifact_bytes) for profile in profiles.values()] + return ComponentMemoryEstimate( + minimum_vram_bytes=minimum_vram_bytes, + minimum_ram_bytes=(sum(artifact_sizes) + _MODEL_RAM_HEADROOM_BYTES) + if artifact_sizes + else 0, + model_asset_count=len(artifact_sizes), + largest_model_bytes=max(artifact_sizes, default=0), + ) def _iter_prompt_string_values(value: object) -> Iterator[str]: @@ -591,12 +648,16 @@ def _component_memory_estimate( prompt: Mapping[str, Any], preferences: WorkflowExecutionPreferences, settings: ModalSyncSettings, + resolved_llm_profiles: Mapping[str, Any] | None = None, ) -> ComponentMemoryEstimate: """Infer conservative RAM and VRAM floors from resident model weight sizes.""" - profile_vram_bytes = _component_profile_minimum_vram_bytes( + profile_estimate = _component_profile_memory_estimate( component, prompt, preferences, + resolved_llm_profiles + if resolved_llm_profiles is not None + else _resolve_prompt_llm_profiles(prompt, settings), ) primary_sizes, additive_sizes = _component_model_asset_sizes( component, @@ -605,10 +666,7 @@ def _component_memory_estimate( ) all_sizes = primary_sizes + additive_sizes if not all_sizes: - return ComponentMemoryEstimate( - minimum_vram_bytes=profile_vram_bytes, - minimum_ram_bytes=0, - ) + return profile_estimate primary_peak_bytes = max(primary_sizes or all_sizes) additive_bytes = sum(additive_sizes) if primary_sizes else 0 @@ -618,10 +676,19 @@ def _component_memory_estimate( + _MODEL_VRAM_HEADROOM_BYTES ) return ComponentMemoryEstimate( - minimum_vram_bytes=max(profile_vram_bytes, model_vram_bytes), - minimum_ram_bytes=resident_weight_bytes + _MODEL_RAM_HEADROOM_BYTES, - model_asset_count=len(all_sizes), - largest_model_bytes=max(all_sizes), + minimum_vram_bytes=max( + profile_estimate.minimum_vram_bytes, + model_vram_bytes, + ), + minimum_ram_bytes=max( + profile_estimate.minimum_ram_bytes, + resident_weight_bytes + _MODEL_RAM_HEADROOM_BYTES, + ), + model_asset_count=len(all_sizes) + profile_estimate.model_asset_count, + largest_model_bytes=max( + max(all_sizes), + profile_estimate.largest_model_bytes, + ), ) @@ -708,6 +775,7 @@ def _plan_component_execution_assignments( scheduler = CostAwareEnvironmentScheduler() history = _execution_history(settings) environment_ids = [environment.environment_id for environment in environments] + resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings) assignments: dict[str, ExecutionAssignment] = {} for component in components: component_signature = _component_execution_signature(component, prompt) @@ -721,6 +789,7 @@ def _plan_component_execution_assignments( prompt, preferences, settings, + resolved_llm_profiles, ) if memory_estimate.model_asset_count: logger.info( diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index edd6d4d..bed20e2 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -470,3 +470,5 @@ - The SSH runtime fingerprint hashed top-level ComfyUI Python modules while its Docker context omitted them, producing an image with package directories but no `nodes.py`. Keep the image selection aligned with the fingerprint selection: include every top-level `.py` file plus the selected package trees, and exclude unrelated top-level assets explicitly. - Cost optimization is safe only after component resource admission. Ordinary loader nodes often have no explicit VRAM declaration, so resolve their local model files before placement, budget the largest primary weight set plus additive LoRA/adapter/ControlNet weights, apply runtime headroom, and compare against probed free VRAM and available RAM. Otherwise an explicit zero-cost 24 GiB host can beat a large Modal GPU for a 66 GiB BF16 model and begin a doomed upload before failing at execution. - A raw Hugging Face model ID cannot be passed directly to an SSH worker just because the Modal dispatcher knows how to resolve it. Run the provider-neutral resolver and snapshot stager inside the selected SSH container, persist both generated manifest and weights in its named volume, stream CPU-stage progress to ComfyUI, and rewrite the payload to the immutable profile ID before invoking the GPU execution kernel. +- SSH LLM staging happens too late to influence an already chosen environment. Resolve Hugging Face metadata before scheduler admission and apply the generated profile's artifact size and estimated VRAM; otherwise a zero-cost 24 GiB host begins downloading a 52 GiB checkpoint before discovering the placement was impossible. +- Cancelling an asyncio task that is awaiting `to_thread` does not stop the blocking Docker staging process, and terminating the local `docker exec` client does not reliably kill its container child. Track active staging invocations and stop the managed worker on cancellation; its named volume preserves partial Hugging Face files for a later compatible run. diff --git a/ssh_executor.py b/ssh_executor.py index ad02b75..dcc5ff8 100644 --- a/ssh_executor.py +++ b/ssh_executor.py @@ -53,6 +53,11 @@ else: # pragma: no cover - top-level remote imports. logger = logging.getLogger(__name__) _STAGED_SSH_PROFILES_LOCK = threading.Lock() _STAGED_SSH_PROFILE_RESULTS: dict[tuple[str, str], dict[str, Any]] = {} +_ACTIVE_SSH_STAGERS_LOCK = threading.Lock() +_ACTIVE_SSH_STAGERS: dict[ + str, + tuple[SshRuntimeManager, SshRuntimeSpec], +] = {} class SshRemoteInvocationError(RuntimeError): @@ -111,6 +116,20 @@ class SshDockerExecutorClient: invocation_id: str, ) -> bool: """Request cancellation of one active invocation on its assigned worker.""" + with _ACTIVE_SSH_STAGERS_LOCK: + active_stager = _ACTIVE_SSH_STAGERS.get(invocation_id) + if active_stager is not None: + active_manager, active_spec = active_stager + stopped = active_manager.stop_worker(active_spec.worker_index) + logger.info( + "Stopped SSH worker to cancel active model staging invocation=%s " + "environment=%s worker_index=%d stopped=%s.", + invocation_id, + active_manager.controller.host.environment_id, + active_spec.worker_index, + stopped, + ) + return stopped manager, spec = self._runtime(payload) request = encode_json_frame( RemoteFrameKind.CANCEL, @@ -280,11 +299,18 @@ class SshDockerExecutorClient: for model_reference in model_references: arguments.extend(("--model-reference", model_reference)) process = manager.controller.docker_popen(tuple(arguments)) - results = self._consume_stager_output( - process, - payload, - _emit_local_llm_staging_progress, - ) + invocation_id = str(payload["invocation_id"]) + with _ACTIVE_SSH_STAGERS_LOCK: + _ACTIVE_SSH_STAGERS[invocation_id] = (manager, spec) + try: + results = self._consume_stager_output( + process, + payload, + _emit_local_llm_staging_progress, + ) + finally: + with _ACTIVE_SSH_STAGERS_LOCK: + _ACTIVE_SSH_STAGERS.pop(invocation_id, None) downloaded_gib = sum( float(result.get("artifact_bytes") or 0) / 1024**3 for result in results diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 6ecebee..ee50131 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -567,6 +567,103 @@ def test_automatic_policy_rejects_zero_cost_host_for_oversized_model( assert assignments["6"].environment_id == "modal:H200" +def test_planner_resolves_hugging_face_metadata_before_cost_ranking( + api_intercept_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(llm_resolver_module, "resolve_model_profile", resolve) + monkeypatch.setattr( + api_intercept_module, + "_schedulable_ssh_hosts", + lambda _settings: (host,), + ) + monkeypatch.setattr( + api_intercept_module, + "_execution_history", + lambda _settings: None, + ) + + assignments = api_intercept_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, ) -> None: diff --git a/tests/test_ssh_executor.py b/tests/test_ssh_executor.py index e191492..0f83b45 100644 --- a/tests/test_ssh_executor.py +++ b/tests/test_ssh_executor.py @@ -182,3 +182,32 @@ def test_ssh_dispatch_stages_and_rewrites_hugging_face_model_reference( assert result == b"serialized" assert stage_calls == [[requested_model]] + + +def test_ssh_cancel_stops_worker_during_model_staging( + ssh_executor_module: Any, + tmp_path: Path, +) -> None: + """Cancellation must stop a worker whose CPU stager is still downloading.""" + stopped: list[int] = [] + manager = SimpleNamespace( + controller=SimpleNamespace(host=SimpleNamespace(environment_id="lambda")), + stop_worker=lambda worker_index: stopped.append(worker_index) or True, + ) + spec = SimpleNamespace(worker_index=3) + invocation_id = "RIV_staging" + client = ssh_executor_module.SshDockerExecutorClient( + registry=SimpleNamespace(), + repo_root=tmp_path, + settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"), + ) + with ssh_executor_module._ACTIVE_SSH_STAGERS_LOCK: + ssh_executor_module._ACTIVE_SSH_STAGERS[invocation_id] = (manager, spec) + try: + cancelled = client.cancel({}, invocation_id) + finally: + with ssh_executor_module._ACTIVE_SSH_STAGERS_LOCK: + ssh_executor_module._ACTIVE_SSH_STAGERS.pop(invocation_id, None) + + assert cancelled is True + assert stopped == [3]