Multi-GPU PR3: CLI/API flags, --list-gpus, status fields, docs - #1264
Open
greenstephen wants to merge 18 commits into
Open
Multi-GPU PR3: CLI/API flags, --list-gpus, status fields, docs#1264greenstephen wants to merge 18 commits into
greenstephen wants to merge 18 commits into
Conversation
Introduce per-component device mapping so DiT, VAE, text encoder, and LM can be assigned to distinct CUDA devices while preserving single-GPU behavior when no mapping is provided. Co-authored-by: Cursor <cursoragent@cursor.com>
Compute VRAM-aware device placement when gpu_mapping=auto on multi-GPU systems, route conditioning tensors to the DiT GPU at inference boundaries, and initialize the LM on its mapped GPU with correct memory budgeting. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR adds VRAM-aware multi-GPU mapping for DiT, VAE, text encoder, and LM components. Mapping is integrated into initialization, generation, CLI, API, UI, and LM device selection, with GPU runtime reporting, tests, and documentation. ChangesMulti-GPU device mapping support
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant CLI
participant device_map
participant dit_handler
participant llm_handler
CLI->>device_map: resolve gpu_mapping from CLI or environment
CLI->>dit_handler: initialize_service(gpu_mapping)
dit_handler-->>CLI: expose component device_map
CLI->>llm_handler: initialize(device=device_map.lm)
llm_handler-->>CLI: initialize LM
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
acestep/core/generation/handler/generate_music_decode.py (1)
149-164: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winBug:
vae_deviceisNonehere, so the CUDA index is always 0.Line 149 resets
vae_device = None(it is later reused at Line 177 for the CPU-restore path), overwriting the component device computed at Line 133. Consequently at Lines 159–162is_cuda_device(vae_device)is evaluated onNone, always taking theelsebranch and yieldingvae_cuda_index = 0. The whole point of this change — querying effective free VRAM on the VAE's actual device — is lost, and on a multi-GPU layout where the VAE is oncuda:1the budget/auto-CPU-offload decision is made againstcuda:0.Compute the CUDA index from the component device (before it is reset) or from a distinct variable.
🔧 Proposed fix
with self._load_model_context("vae"): pred_latents_cpu = pred_latents.detach().cpu() vae_device = self._get_component_device("vae") + vae_component_device = vae_device pred_latents_for_decode = (else: vae_cuda_index = ( - cuda_device_index(vae_device) - if is_cuda_device(vae_device) + cuda_device_index(vae_component_device) + if is_cuda_device(vae_component_device) else 0 )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/generate_music_decode.py` around lines 149 - 164, The VAE VRAM check is using a reset `vae_device` value, so `generate_music_decode` always falls back to CUDA index 0 instead of the VAE’s actual device. In `generate_music_decode`, preserve the component device computed earlier or use a separate variable before `vae_device = None` is set, and pass that real device into `is_cuda_device`, `cuda_device_index`, and `get_effective_free_vram_gb` so multi-GPU setups use the correct index.acestep/llm_inference.py (1)
547-558: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winFix CUDA checks to accept indexed devices
acestep/llm_inference.py:547-558rewritescudatocuda:0, but the laterdevice == "cuda"branches still gate vLLM, Jetson handling, and VRAM preflight. Explicit CUDA requests will fall through to PyTorch and skip those CUDA-specific paths. Useis_cuda_device(device)in those checks, or keep the raw request separate.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/llm_inference.py` around lines 547 - 558, The CUDA device handling in llm_inference.py is too string-specific: after normalize_component_device() rewrites an explicit CUDA request like cuda to cuda:0, later branches that check for "cuda" will miss vLLM setup, Jetson handling, and VRAM preflight. Update the CUDA-specific conditionals in initialize() and any nearby device-routing logic to use is_cuda_device(device) (or preserve the original request separately) so indexed CUDA devices still take the intended CUDA path.
🧹 Nitpick comments (5)
acestep/api/http/model_init_service.py (1)
141-151: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the device-map/LM-device resolution + deprecation-log boilerplate into a shared helper.
This exact block (read
handler.device_map, overridelm_device, calllog_lm_device_deprecation) is duplicated almost verbatim inacestep/api/startup_llm_init.py(lines 77-87), with a subtly different variant inacestep/ui/gradio/events/generation/service_init.py. Centralizing this inacestep/device_map.py(e.g.resolve_lm_device(explicit_env, device_map, fallback_device) -> tuple[str, bool]) would remove duplication and prevent the entry points from drifting out of sync (see the ordering bug flagged inservice_init.py).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/api/http/model_init_service.py` around lines 141 - 151, The LM-device resolution and deprecation logging logic is duplicated across `model_init_service.py` and the other startup/init entry points. Move the `handler.device_map` lookup, `lm_device` override, and `log_lm_device_deprecation` call into a shared helper in `acestep.device_map` (for example, a `resolve_lm_device(...)` function). Update `model_init_service`, `startup_llm_init`, and the Gradio `service_init` path to call that helper so all entry points use the same resolution order and stay in sync.acestep/api/startup_llm_init.py (1)
26-26: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring not updated for new
dit_handlerparameter.The function signature gained
dit_handler: Any = None, but the docstring (line 28) still just says "Initialize LLM model according to GPU config and environment overrides." with no mention of the new parameter's purpose.As per coding guidelines: "Docstrings: Mandatory for all modules, classes, and public functions. Use concise format with Args, Returns, and exception documentation."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/api/startup_llm_init.py` at line 26, The public function in startup_llm_init.py now accepts dit_handler, but its docstring is stale. Update the docstring for the LLM initialization function to include an Args section that documents dit_handler: Any = None and briefly states its purpose, keeping the existing description aligned with the current signature and docstring style used in the module.Source: Coding guidelines
acestep/core/generation/handler/init_service_orchestrator.py (2)
99-116: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
self.dtypeused for text_encoder is derived only from the DiT device's bfloat16 support, unlike VAE which resolves dtype per-component.
_load_vae_modeluses_get_vae_dtype(device)to pick a device-appropriate dtype, but_load_text_encoder_and_tokenizer(init_service_loader_components.py) appliesself.dtype, which is computed here solely fromcuda_device_index(self.device)i.e. the DiT device. With an explicit heterogeneousgpu_mapping(e.g.text_encoderpinned to a pre-Ampere GPU whileditsits on an Ampere+ GPU), the text encoder could be cast tobfloat16on a GPU that only has software-emulated/degraded bfloat16 support, or trigger aRuntimeErrorin code paths usingtorch.autocast/mixed-precision that explicitly check per-device bf16 support.Consider adding a text-encoder-specific dtype resolution analogous to
_get_vae_dtype(target_device).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/init_service_orchestrator.py` around lines 99 - 116, The text encoder dtype is being chosen from the DiT device only, which can be wrong for heterogeneous gpu_mapping setups. Update the service initialization path in init_service_orchestrator.py so text-encoder loading does not rely solely on self.dtype derived from cuda_device_index(self.device); instead, resolve a text_encoder-specific dtype based on its actual target device, similar to _get_vae_dtype(device). Then wire _load_text_encoder_and_tokenizer to use that per-component dtype so it matches the device-specific bf16 support.
77-92: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
_resolve_component_device_mapcall omitsuse_lm/batch_size, weakening VRAM-aware auto-layout.
_resolve_component_device_map(and the underlyingresolve_component_device_map/compute_auto_device_map) acceptsuse_lmandbatch_sizeto compute a VRAM-awaregpu_mapping="auto"layout, but this call site never forwards them, so the defaults (use_lm=True,batch_size=1) always apply regardless of the caller's actual intent (several callers already knowinit_llmat this point, e.g.service_init.py/model_init_service.py). This can make "auto" layout reserve VRAM for an LM that won't be loaded, or size for batch=1 when actual generation batch sizes differ.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/init_service_orchestrator.py` around lines 77 - 92, The `_resolve_component_device_map` call in `init_service_orchestrator.py` is missing the caller’s actual `use_lm` and `batch_size`, so auto device-map sizing always falls back to default VRAM assumptions. Update the `init_service_orchestrator` flow to thread the real LM intent and generation batch size into `_resolve_component_device_map`, and ensure that method forwards those values into `resolve_component_device_map`/`compute_auto_device_map` so `gpu_mapping="auto"` reflects the actual load.acestep/core/generation/handler/init_service_test.py (1)
264-308: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLGTM!
Solid coverage of the new device-map resolution/routing helpers and the
gpu_mappingfields oninitialize_service's status/params.One optional gap: the multi-GPU integration test (
test_initialize_service_records_multi_gpu_mapping) mapsdit/vae/text_encoderto the same device index and only varieslm, so it doesn't exercise a scenario whereditandtext_encoder/vaesit on genuinely different GPU indices — which is the scenario most relevant to the dtype-selection concern raised ininit_service_orchestrator.py. Consider adding a case with heterogeneous DiT/VAE/text_encoder placement if you want to close that coverage gap.Also applies to: 512-512, 525-525, 968-1005
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/init_service_test.py` around lines 264 - 308, The new multi-GPU coverage in init_service_test.py only verifies a mapping where dit, vae, and text_encoder share the same GPU index, so it does not exercise the dtype-sensitive case where those components are placed on different devices. Update test_initialize_service_records_multi_gpu_mapping or add a companion test to use a heterogeneous gpu_mapping with distinct indices for dit, vae, and text_encoder, and assert the resolved service/device-map behavior through initialize_service and the related helpers in _Host.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@acestep/device_map.py`:
- Around line 1-522: The module is over the LOC hard cap and should be split by
responsibility before merge. Move the dataclasses/types, mapping parsing,
auto-layout computation, and status/serialization helpers out of device_map.py
into smaller modules, then re-export the public API from device_map.py so
callers like resolve_component_device_map, compute_auto_device_map, and
ComponentDeviceMap keep working unchanged. If you cannot split immediately, add
a short PR justification plus a concrete follow-up split plan, but preserve the
stable facade imports either way.
- Around line 256-280: The broad exception handling in discover_gpus is masking
unexpected failures from torch.cuda.get_device_capability. Update the try/except
around the capability lookup in discover_gpus to catch a narrower, expected
exception type such as RuntimeError, and leave other errors unhandled so genuine
CUDA issues surface instead of being silently converted to None.
- Around line 159-204: compute_auto_device_map currently allows the LM to fall
back onto the same GPU as the DiT without accounting for the DiT reservation, so
a layout can pass per-model checks but still overcommit VRAM. Update
compute_auto_device_map to reserve the DiT footprint when evaluating the LM
candidate in the same-gpu fallback path, using dit_gpu and the LM selection
logic so the chosen gpu must have enough free VRAM for both workloads. If no GPU
satisfies the combined requirement, return a LayoutError with an appropriate
LM-specific suggestion.
In `@docs/en/MULTI_GPU.md`:
- Around line 85-87: The fenced mapping example in the MULTI_GPU docs is missing
a language tag and triggers markdownlint MD040. Update the example fence around
the text_encoder/lm mapping to use the existing fenced block in the docs with a
text language label, keeping the content unchanged so the markdown stays
lint-clean.
---
Outside diff comments:
In `@acestep/core/generation/handler/generate_music_decode.py`:
- Around line 149-164: The VAE VRAM check is using a reset `vae_device` value,
so `generate_music_decode` always falls back to CUDA index 0 instead of the
VAE’s actual device. In `generate_music_decode`, preserve the component device
computed earlier or use a separate variable before `vae_device = None` is set,
and pass that real device into `is_cuda_device`, `cuda_device_index`, and
`get_effective_free_vram_gb` so multi-GPU setups use the correct index.
In `@acestep/llm_inference.py`:
- Around line 547-558: The CUDA device handling in llm_inference.py is too
string-specific: after normalize_component_device() rewrites an explicit CUDA
request like cuda to cuda:0, later branches that check for "cuda" will miss vLLM
setup, Jetson handling, and VRAM preflight. Update the CUDA-specific
conditionals in initialize() and any nearby device-routing logic to use
is_cuda_device(device) (or preserve the original request separately) so indexed
CUDA devices still take the intended CUDA path.
---
Nitpick comments:
In `@acestep/api/http/model_init_service.py`:
- Around line 141-151: The LM-device resolution and deprecation logging logic is
duplicated across `model_init_service.py` and the other startup/init entry
points. Move the `handler.device_map` lookup, `lm_device` override, and
`log_lm_device_deprecation` call into a shared helper in `acestep.device_map`
(for example, a `resolve_lm_device(...)` function). Update `model_init_service`,
`startup_llm_init`, and the Gradio `service_init` path to call that helper so
all entry points use the same resolution order and stay in sync.
In `@acestep/api/startup_llm_init.py`:
- Line 26: The public function in startup_llm_init.py now accepts dit_handler,
but its docstring is stale. Update the docstring for the LLM initialization
function to include an Args section that documents dit_handler: Any = None and
briefly states its purpose, keeping the existing description aligned with the
current signature and docstring style used in the module.
In `@acestep/core/generation/handler/init_service_orchestrator.py`:
- Around line 99-116: The text encoder dtype is being chosen from the DiT device
only, which can be wrong for heterogeneous gpu_mapping setups. Update the
service initialization path in init_service_orchestrator.py so text-encoder
loading does not rely solely on self.dtype derived from
cuda_device_index(self.device); instead, resolve a text_encoder-specific dtype
based on its actual target device, similar to _get_vae_dtype(device). Then wire
_load_text_encoder_and_tokenizer to use that per-component dtype so it matches
the device-specific bf16 support.
- Around line 77-92: The `_resolve_component_device_map` call in
`init_service_orchestrator.py` is missing the caller’s actual `use_lm` and
`batch_size`, so auto device-map sizing always falls back to default VRAM
assumptions. Update the `init_service_orchestrator` flow to thread the real LM
intent and generation batch size into `_resolve_component_device_map`, and
ensure that method forwards those values into
`resolve_component_device_map`/`compute_auto_device_map` so `gpu_mapping="auto"`
reflects the actual load.
In `@acestep/core/generation/handler/init_service_test.py`:
- Around line 264-308: The new multi-GPU coverage in init_service_test.py only
verifies a mapping where dit, vae, and text_encoder share the same GPU index, so
it does not exercise the dtype-sensitive case where those components are placed
on different devices. Update test_initialize_service_records_multi_gpu_mapping
or add a companion test to use a heterogeneous gpu_mapping with distinct indices
for dit, vae, and text_encoder, and assert the resolved service/device-map
behavior through initialize_service and the related helpers in _Host.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8f2d8980-8910-47fc-8788-fa4fd843a27a
📒 Files selected for processing (25)
acestep/acestep_v15_pipeline.pyacestep/acestep_v15_pipeline_test.pyacestep/api/http/model_init_service.pyacestep/api/http/model_service_routes.pyacestep/api/http/model_service_routes_test.pyacestep/api/startup_llm_init.pyacestep/api/startup_model_init.pyacestep/core/generation/handler/audio_codes.pyacestep/core/generation/handler/conditioning_embed.pyacestep/core/generation/handler/generate_music_decode.pyacestep/core/generation/handler/init_service_catalog.pyacestep/core/generation/handler/init_service_loader.pyacestep/core/generation/handler/init_service_offload_context.pyacestep/core/generation/handler/init_service_orchestrator.pyacestep/core/generation/handler/init_service_setup.pyacestep/core/generation/handler/init_service_test.pyacestep/core/generation/handler/service_generate_execute.pyacestep/device_map.pyacestep/gpu_config.pyacestep/llm_inference.pyacestep/test_device_map.pyacestep/ui/gradio/events/generation/service_init.pycli.pydocs/en/GPU_COMPATIBILITY.mddocs/en/MULTI_GPU.md
Move GPU discovery, layout, parsing, resolution, and status helpers into focused submodules under acestep/device_map/ with a stable public facade. Co-authored-by: Cursor <cursoragent@cursor.com>
…ndex. Route generate_kwargs tensors to the DiT device, move latents during preprocess, and preserve CUDA device indices in APG project(). Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
acestep/models/common/apg_guidance_test.py (1)
12-59: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding an MPS device-preservation test.
Existing tests cover CPU and non-default CUDA indices but nothing for the
mpsbranch inproject(), which is where a device-preservation regression currently exists (see companion comment onapg_guidance.py). A@unittest.skipUnless(torch.backends.mps.is_available(), ...)test mirroring the CUDA cases would catch this class of bug.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/models/common/apg_guidance_test.py` around lines 12 - 59, Add an MPS device-preservation test alongside the existing `ApgGuidanceDeviceTests` cases, since `project()` and `apg_forward()` are already validated for CPU and CUDA but not the `mps` path. Mirror the structure of `test_project_preserves_non_default_cuda_index` and `test_apg_forward_preserves_non_default_cuda_index`, guarded with `torch.backends.mps.is_available()`, and assert the returned tensors stay on the same MPS device.acestep/core/generation/handler/service_generate_execute.py (1)
135-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRedundant fallback duplicates
_get_component_device.
_get_component_device("dit")already returnsself.devicewhenself.device_mapisNone(seeinit_service_setup.py), so this ternary is dead logic that just duplicates that fallback. Same pattern is repeated ininit_service_memory_basic.py's_ensure_silence_latent_on_device.♻️ Simplify
- dit_device = ( - self._get_component_device("dit") - if getattr(self, "device_map", None) is not None - else self.device - ) + dit_device = self._get_component_device("dit") kwargs["timesteps"] = torch.tensor( timesteps, dtype=torch.float32, device=dit_device, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/service_generate_execute.py` around lines 135 - 144, The device selection in the timesteps setup duplicates the fallback already handled by _get_component_device("dit"), so remove the extra getattr(self, "device_map", None) ternary and rely on _get_component_device directly when assigning dit_device in service_generate_execute.py. Apply the same simplification pattern in _ensure_silence_latent_on_device in init_service_memory_basic.py so both call sites use the centralized device fallback logic consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@acestep/device_map/parsing.py`:
- Around line 66-68: The fallback in the device parsing logic is a no-op because
`backend` is already set by `device_type(default_device)`, so the
`is_cuda_device(default_device)` check cannot change anything. Update the
`parsing.py` logic around `device_type`, `is_cuda_device`, and `backend` so
unknown values actually map to the intended fallback backend instead of
reassigning the same value; keep the explicit `"cuda"`, `"mps"`, `"xpu"`, and
`"cpu"` cases and make the non-matching branch produce a real fallback result.
In `@acestep/models/common/apg_guidance.py`:
- Around line 21-34: The MPS fallback in the APG guidance projection helper is
overwriting the original target device, so the returned tensors end up on CPU
instead of the caller’s MPS device. In the function that captures `dtype` and
`device` before the `mps` check, keep the original `device` unchanged, move `v0`
and `v1` to CPU only for the computation, and then use the preserved original
device in the final `.to(...)` calls. Verify the `v0_parallel` and
`v0_orthogonal` returns still restore the original device for the MPS path.
---
Nitpick comments:
In `@acestep/core/generation/handler/service_generate_execute.py`:
- Around line 135-144: The device selection in the timesteps setup duplicates
the fallback already handled by _get_component_device("dit"), so remove the
extra getattr(self, "device_map", None) ternary and rely on
_get_component_device directly when assigning dit_device in
service_generate_execute.py. Apply the same simplification pattern in
_ensure_silence_latent_on_device in init_service_memory_basic.py so both call
sites use the centralized device fallback logic consistently.
In `@acestep/models/common/apg_guidance_test.py`:
- Around line 12-59: Add an MPS device-preservation test alongside the existing
`ApgGuidanceDeviceTests` cases, since `project()` and `apg_forward()` are
already validated for CPU and CUDA but not the `mps` path. Mirror the structure
of `test_project_preserves_non_default_cuda_index` and
`test_apg_forward_preserves_non_default_cuda_index`, guarded with
`torch.backends.mps.is_available()`, and assert the returned tensors stay on the
same MPS device.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0d3d8584-982e-4fba-be3e-8a65853a71ac
📒 Files selected for processing (20)
acestep/acestep_v15_pipeline.pyacestep/core/generation/handler/conditioning_embed.pyacestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_setup.pyacestep/core/generation/handler/init_service_test.pyacestep/core/generation/handler/service_generate_execute.pyacestep/device_map/__init__.pyacestep/device_map/constants.pyacestep/device_map/devices.pyacestep/device_map/discovery.pyacestep/device_map/errors.pyacestep/device_map/layout.pyacestep/device_map/parsing.pyacestep/device_map/resolve.pyacestep/device_map/status.pyacestep/device_map/types.pyacestep/models/common/apg_guidance.pyacestep/models/common/apg_guidance_test.pyacestep/test_device_map.pydocs/en/MULTI_GPU.md
✅ Files skipped from review due to trivial changes (3)
- acestep/device_map/errors.py
- acestep/device_map/constants.py
- docs/en/MULTI_GPU.md
🚧 Files skipped from review as they are similar to previous changes (4)
- acestep/core/generation/handler/conditioning_embed.py
- acestep/test_device_map.py
- acestep/core/generation/handler/init_service_test.py
- acestep/acestep_v15_pipeline.py
…ack. Also narrow get_device_capability exception handling to RuntimeError. Co-authored-by: Cursor <cursoragent@cursor.com>
Reject unsupported default_device backends explicitly in GPU mapping parsing, and preserve the original MPS device when project() uses CPU math fallback. Co-authored-by: Cursor <cursoragent@cursor.com>
_is_on_target_device only compares backend types, so cuda:0 and cuda:3 look equivalent. Use exact device matching when placing silence_latent on the DiT device from device_map. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
acestep/core/generation/handler/init_service_memory_basic.py (1)
118-124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHoist
normalize_component_deviceand narrow the exception
acestep.device_mapdoesn’t import back intoacestep.core.generation.handler, so this can be a module-level import. CatchDeviceMapErrorandRuntimeErrorhere instead of swallowing every exception.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/init_service_memory_basic.py` around lines 118 - 124, Move the normalize_component_device import in the tensor/device check out of the local try block in the handler module and make it a module-level import, since acestep.device_map has no back-reference into acestep.core.generation.handler. In the device comparison logic that uses normalize_component_device and torch.device, replace the broad except Exception with a narrow catch for DeviceMapError and RuntimeError only, while keeping the existing False fallback.Sources: Coding guidelines, Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@acestep/core/generation/handler/init_service_memory_basic.py`:
- Around line 118-124: Move the normalize_component_device import in the
tensor/device check out of the local try block in the handler module and make it
a module-level import, since acestep.device_map has no back-reference into
acestep.core.generation.handler. In the device comparison logic that uses
normalize_component_device and torch.device, replace the broad except Exception
with a narrow catch for DeviceMapError and RuntimeError only, while keeping the
existing False fallback.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9ffc259d-fe07-4739-9411-dfbb1a884308
📒 Files selected for processing (2)
acestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_test.py
Module-level import avoids per-call overhead; catch DeviceMapError and RuntimeError instead of a broad Exception in _tensor_on_exact_device. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace bare Exception handlers with DeviceMapError/RuntimeError/TypeError for exact device checks and RuntimeError/TypeError/ValueError for backend alias parsing in _is_on_target_device. Co-authored-by: Cursor <cursoragent@cursor.com>
vae_component_device was overwritten with None before the CUDA free-VRAM check, so multi-GPU layouts always queried cuda:0. Keep a separate vae_restore_device for the CPU offload restore path. Co-authored-by: Cursor <cursoragent@cursor.com>
Wrap cuda_device_index int() parsing so strings like cuda:x raise the module's domain error instead of a raw ValueError. Co-authored-by: Cursor <cursoragent@cursor.com>
Author
|
@coderabbitai resume |
Contributor
✅ Action performedReviews resumed. |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
acestep/core/generation/handler/generate_music_decode_test.py (1)
1-358: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftTest module is well past the 200 LOC hard cap.
The visible content of this file already exceeds 350+ lines, well beyond the guideline's hard cap. Consider splitting decode-mixin tests by scenario (basic decode, CPU-offload/error-restoration, VRAM/device-mapping) into separate test modules.
Based on learnings and coding guidelines: Enforce a module size guideline for Python files: only raise module-size concerns when a file exceeds 200 lines of code (LOC)... Apply this across all Python files (pattern **/*.py) to maintain consistency in reviews; when a file exceeds 200 LOC, suggest splitting into smaller modules to improve readability and maintainability. Also per coding guidelines, "If a module would exceed 200 LOC, split by responsibility before merging, or add a short justification in PR notes with a concrete follow-up split plan."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/generate_music_decode_test.py` around lines 1 - 358, The test module exceeds the 200 LOC guideline and should be split by responsibility. Move the scenario groups in GenerateMusicDecodeMixinTests into smaller test modules, such as basic decode, CPU-offload/error restoration, and VRAM/device-mapping, while keeping shared helpers like _load_generate_music_decode_module, _Host, and the fake VAE doubles in a common test utility if needed. Ensure the resulting Python test files each stay under the module-size cap and preserve the same assertions and coverage.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@acestep/core/generation/handler/generate_music_decode_test.py`:
- Around line 1-358: The test module exceeds the 200 LOC guideline and should be
split by responsibility. Move the scenario groups in
GenerateMusicDecodeMixinTests into smaller test modules, such as basic decode,
CPU-offload/error restoration, and VRAM/device-mapping, while keeping shared
helpers like _load_generate_music_decode_module, _Host, and the fake VAE doubles
in a common test utility if needed. Ensure the resulting Python test files each
stay under the module-size cap and preserve the same assertions and coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c9f53c38-bff8-498a-97b8-c2239e434883
📒 Files selected for processing (6)
acestep/core/generation/handler/generate_music_decode.pyacestep/core/generation/handler/generate_music_decode_test.pyacestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_test.pyacestep/device_map/devices.pyacestep/test_device_map.py
🚧 Files skipped from review as they are similar to previous changes (5)
- acestep/core/generation/handler/generate_music_decode.py
- acestep/core/generation/handler/init_service_memory_basic.py
- acestep/device_map/devices.py
- acestep/core/generation/handler/init_service_test.py
- acestep/test_device_map.py
Move shared fixtures into generate_music_decode_test_support.py and split scenario tests into prepare, basic, CPU-offload, and VRAM modules. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py (1)
86-97: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winStrengthen the "not restored to GPU" assertion.
The test only checks
vae_to_callscount andpred_wavsshape; it discards_cpu_latentswithout asserting its device. The test name promises verification that latents aren't restored to GPU, but a regression that keeps latents on CPU via a different code path than an extravae.to()call wouldn't be caught.♻️ Suggested addition
- pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents( + pred_wavs, cpu_latents, _costs = host._decode_generate_music_pred_latents( pred_latents=torch.ones(1, 4, 3), progress=None, use_tiled_decode=False, time_costs={"total_time_cost": 1.0}, ) self.assertEqual(len(host.vae.vae_to_calls), 1) self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8)) + self.assertEqual(str(cpu_latents.device), "cpu")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py` around lines 86 - 97, The test in generate_music_decode_cpu_offload should assert the returned latents stay on CPU, not just that SuccessHost.vae_to_calls stays at 1 and pred_wavs has the expected shape. Update the _decode_generate_music_pred_latents assertion block to also validate _cpu_latents.device (or equivalent tensor device check) so the test covers the “not restored to GPU” behavior promised by the test name.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@acestep/core/generation/handler/generate_music_decode_test_support.py`:
- Around line 12-41: The test helper is stubbing the live
acestep.core.generation.handler package in sys.modules, which can break later
imports that rely on the real package initialization and re-exports. Update
load_generate_music_decode_module to avoid writing placeholder modules under the
acestep.* namespace, ideally by loading generate_music_decode.py under a private
test-only module name or by saving and restoring any existing sys.modules
entries after import. Keep the fix localized to the
load_generate_music_decode_module path and the GENERATE_MUSIC_DECODE_MODULE
setup.
---
Nitpick comments:
In `@acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py`:
- Around line 86-97: The test in generate_music_decode_cpu_offload should assert
the returned latents stay on CPU, not just that SuccessHost.vae_to_calls stays
at 1 and pred_wavs has the expected shape. Update the
_decode_generate_music_pred_latents assertion block to also validate
_cpu_latents.device (or equivalent tensor device check) so the test covers the
“not restored to GPU” behavior promised by the test name.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b7ffe23c-ba38-42da-998e-47d95dcd6e9a
📒 Files selected for processing (6)
acestep/core/generation/handler/generate_music_decode_basic_test.pyacestep/core/generation/handler/generate_music_decode_cpu_offload_test.pyacestep/core/generation/handler/generate_music_decode_prepare_test.pyacestep/core/generation/handler/generate_music_decode_test.pyacestep/core/generation/handler/generate_music_decode_test_support.pyacestep/core/generation/handler/generate_music_decode_vram_test.py
💤 Files with no reviewable changes (1)
- acestep/core/generation/handler/generate_music_decode_test.py
✅ Files skipped from review due to trivial changes (1)
- acestep/core/generation/handler/generate_music_decode_basic_test.py
Drop the custom sys.modules package stubbing loader; import the mixin through the public package path like other handler tests. Co-authored-by: Cursor <cursoragent@cursor.com>
6 tasks
Resolve Gradio lm_device after DiT initialize_service so device_map.lm is applied, normalize/preserve indexed CUDA devices in llm_inference, and move text-encoder token ids onto the text-encoder component device. Co-authored-by: Cursor <cursoragent@cursor.com>
Expose --gpu-mapping and --list-gpus on the Gradio and generation CLIs, surface GPU inventory and device_map in API health/model endpoints, log ACESTEP_LM_DEVICE deprecation, and document multi-GPU usage. Co-authored-by: Cursor <cursoragent@cursor.com>
Reserve DiT VRAM before co-located LM fallback in auto-layout, narrow get_device_capability exception handling to RuntimeError, and fix MULTI_GPU.md fenced-code lint. Co-authored-by: Cursor <cursoragent@cursor.com>
greenstephen
force-pushed
the
feat/multi-gpu-pr3-cli-api
branch
from
July 11, 2026 14:35
f1eae34 to
1eba94f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
acestep/llm_inference_cuda_index_test.py (1)
53-59: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the production vLLM selection path.
This reimplements the predicate locally rather than testing
LLMHandlerbehavior, so it cannot catch a regression in the real gate. Assert the selected backend through the actual initialization/selection boundary instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/llm_inference_cuda_index_test.py` around lines 53 - 59, The test_vllm_gate_accepts_indexed_cuda test duplicates the CUDA predicate instead of exercising production backend selection. Update it to initialize or invoke the actual LLMHandler vLLM selection boundary with device “cuda:1”, then assert that the selected backend remains “vllm” rather than “pt”.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@acestep/acestep_v15_pipeline.py`:
- Around line 296-311: Extract the CLI argument definitions and GPU-mapping
startup logic currently surrounding the gpu_mapping and list_gpus options into
focused modules, keeping each new module under 200 LOC. Preserve main() in
acestep_v15_pipeline.py as the stable entrypoint facade, wiring it to the
extracted parsing and startup helpers without changing existing CLI behavior or
GPU-mapping semantics.
In `@acestep/core/generation/handler/conditioning_embed_test.py`:
- Around line 27-41: Add concise docstrings to the modified test helpers: the
fixture class __init__, _get_component_device, and
_DeviceCheckingEncoder.__call__. Describe each helper’s purpose and key behavior
without changing its implementation.
In `@acestep/llm_inference.py`:
- Around line 744-760: Update the CUDA VRAM calculation in the visible inference
block so total_gb is read from the selected device_index rather than implicitly
device 0, while preserving the MAX_CUDA_VRAM override behavior. Keep free_gb and
the subsequent VRAM warning comparison aligned to the same target GPU.
In `@acestep/ui/gradio/events/generation/service_init_test.py`:
- Around line 279-281: Add a concise docstring to the new `_init_service`
callback describing its initialization behavior and return value, while leaving
its existing mock setup and return unchanged.
- Around line 255-310: Move the InitServiceWrapperDeviceResolutionTests test
class, including test_init_llm_uses_device_map_lm_after_dit_init and its related
setup/imports, into a dedicated device-resolution test module. Update imports or
test discovery as needed, while leaving unrelated service-initialization tests
in service_init_test.py.
---
Nitpick comments:
In `@acestep/llm_inference_cuda_index_test.py`:
- Around line 53-59: The test_vllm_gate_accepts_indexed_cuda test duplicates the
CUDA predicate instead of exercising production backend selection. Update it to
initialize or invoke the actual LLMHandler vLLM selection boundary with device
“cuda:1”, then assert that the selected backend remains “vllm” rather than “pt”.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9da33045-5fab-4d38-8e97-c3d3d7b21489
📒 Files selected for processing (42)
acestep/acestep_v15_pipeline.pyacestep/acestep_v15_pipeline_test.pyacestep/api/http/model_init_service.pyacestep/api/http/model_service_routes.pyacestep/api/http/model_service_routes_test.pyacestep/api/startup_llm_init.pyacestep/core/generation/handler/conditioning_embed.pyacestep/core/generation/handler/conditioning_embed_test.pyacestep/core/generation/handler/generate_music_decode.pyacestep/core/generation/handler/generate_music_decode_basic_test.pyacestep/core/generation/handler/generate_music_decode_cpu_offload_test.pyacestep/core/generation/handler/generate_music_decode_prepare_test.pyacestep/core/generation/handler/generate_music_decode_test.pyacestep/core/generation/handler/generate_music_decode_test_support.pyacestep/core/generation/handler/generate_music_decode_vram_test.pyacestep/core/generation/handler/init_service_memory_basic.pyacestep/core/generation/handler/init_service_setup.pyacestep/core/generation/handler/init_service_test.pyacestep/core/generation/handler/service_generate_execute.pyacestep/core/generation/handler/service_generate_flow_edit_source.pyacestep/device_map/__init__.pyacestep/device_map/constants.pyacestep/device_map/devices.pyacestep/device_map/discovery.pyacestep/device_map/errors.pyacestep/device_map/layout.pyacestep/device_map/parsing.pyacestep/device_map/resolve.pyacestep/device_map/status.pyacestep/device_map/types.pyacestep/llm_backend_compat.pyacestep/llm_backend_compat_test.pyacestep/llm_inference.pyacestep/llm_inference_cuda_index_test.pyacestep/models/common/apg_guidance.pyacestep/models/common/apg_guidance_test.pyacestep/test_device_map.pyacestep/ui/gradio/events/generation/service_init.pyacestep/ui/gradio/events/generation/service_init_test.pycli.pydocs/en/GPU_COMPATIBILITY.mddocs/en/MULTI_GPU.md
💤 Files with no reviewable changes (1)
- acestep/core/generation/handler/generate_music_decode_test.py
✅ Files skipped from review due to trivial changes (4)
- acestep/device_map/errors.py
- acestep/core/generation/handler/service_generate_flow_edit_source.py
- docs/en/MULTI_GPU.md
- docs/en/GPU_COMPATIBILITY.md
🚧 Files skipped from review as they are similar to previous changes (30)
- acestep/core/generation/handler/generate_music_decode_prepare_test.py
- acestep/device_map/init.py
- acestep/device_map/constants.py
- acestep/api/startup_llm_init.py
- acestep/core/generation/handler/generate_music_decode_vram_test.py
- acestep/ui/gradio/events/generation/service_init.py
- acestep/device_map/discovery.py
- acestep/device_map/devices.py
- acestep/device_map/layout.py
- acestep/device_map/status.py
- acestep/api/http/model_service_routes_test.py
- acestep/device_map/types.py
- acestep/models/common/apg_guidance_test.py
- acestep/core/generation/handler/service_generate_execute.py
- acestep/core/generation/handler/generate_music_decode_test_support.py
- acestep/core/generation/handler/generate_music_decode_basic_test.py
- acestep/api/http/model_init_service.py
- acestep/core/generation/handler/generate_music_decode.py
- acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py
- acestep/device_map/parsing.py
- acestep/core/generation/handler/init_service_memory_basic.py
- acestep/core/generation/handler/init_service_setup.py
- acestep/core/generation/handler/conditioning_embed.py
- acestep/device_map/resolve.py
- acestep/models/common/apg_guidance.py
- cli.py
- acestep/core/generation/handler/init_service_test.py
- acestep/acestep_v15_pipeline_test.py
- acestep/test_device_map.py
- acestep/api/http/model_service_routes.py
Split Gradio pipeline CLI/startup into focused modules under the 200 LOC cap, extract device-resolution tests, add missing test helper docstrings, and read LM free/total VRAM from the mapped cuda:N device. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
acestep/ui/gradio/events/generation/service_init_test.py (2)
99-113: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the same keyword-argument pattern here for consistency.
Same concern as the other call — five consecutive positional
Falsevalues are error-prone. Apply keyword arguments matching the device resolution test file.♻️ Proposed refactor for `test_project_root_is_consistent_with_checkpoint_dir`
module.init_service_wrapper( dit_handler, llm_handler, "/any/path/checkpoints", "acestep-v15-turbo", "cpu", - False, - None, - "vllm", - False, - False, - False, - False, - False, + init_llm=False, + lm_model=None, + backend="vllm", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/ui/gradio/events/generation/service_init_test.py` around lines 99 - 113, Update the init_service_wrapper call in test_project_root_is_consistent_with_checkpoint_dir to replace the consecutive positional False values with the corresponding keyword arguments, matching the device resolution test’s established pattern while preserving all existing argument values.
51-65: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse keyword arguments for boolean parameters in
init_service_wrappercalls.Five consecutive positional
Falsevalues make it hard to verify which parameter each corresponds to, and any future parameter addition or reordering could silently break the test. The newservice_init_device_resolution_test.pyalready uses keyword arguments for these same parameters — aligning here would improve consistency and reduce risk.♻️ Proposed refactor for `test_passes_project_root_not_checkpoint_dir`
module.init_service_wrapper( dit_handler, llm_handler, checkpoint_value, "acestep-v15-turbo", "cpu", - False, - None, - "vllm", - False, - False, - False, - False, - False, + init_llm=False, + lm_model=None, + backend="vllm", + use_flash_attention=False, + offload_to_cpu=False, + offload_dit_to_cpu=False, + compile_model=False, + quantization=False, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/ui/gradio/events/generation/service_init_test.py` around lines 51 - 65, Update the init_service_wrapper call in test_passes_project_root_not_checkpoint_dir to pass all boolean parameters as keyword arguments, matching the parameter names and style used in service_init_device_resolution_test.py. Keep the existing boolean values and non-boolean positional arguments unchanged.acestep/gradio_pipeline_mode_defaults.py (1)
48-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
logger.warning()instead ofprint()for warning messages.Per coding guidelines, warnings should be logged with
loguru.logger, notprint(). Thegradio_pipeline_cli.pymodule in this same PR correctly useslogger.warning()for its CUDA probe failure — these new files should be consistent.♻️ Proposed fix
+from loguru import logger + from acestep.gpu_config import VRAM_AUTO_OFFLOAD_THRESHOLD_GB @@ - print( - f"WARNING: 4B LM model is too large for {gpu_memory_gb:.0f}GB GPU. " - f"Downgrading to 1.7B variant: {fallback}" - ) + logger.warning( + "4B LM model is too large for {:.0f}GB GPU. " + "Downgrading to 1.7B variant: {}", + gpu_memory_gb, + fallback, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/gradio_pipeline_mode_defaults.py` around lines 48 - 51, Replace the print call in the 4B-to-1.7B fallback warning with loguru’s logger.warning(), preserving the existing warning message and fallback details. Ensure the module imports logger consistently with gradio_pipeline_cli.py.Source: Coding guidelines
acestep/gradio_pipeline_startup.py (1)
137-140: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
loggerinstead ofprint(..., file=sys.stderr)for warning/error messages.Per coding guidelines, errors and warnings should be logged with
loguru.logger, notprint(). These stderr messages are log output, not CLI output. The sibling modulegradio_pipeline_cli.pyalready useslogger.warning()correctly.♻️ Proposed fix
+from loguru import logger + from acestep.device_map import log_lm_device_deprecation @@ - print( - "Warning: No LM models available, skipping LM initialization", - file=sys.stderr, - ) + logger.warning("No LM models available, skipping LM initialization") @@ - print(f"Warning: LM model download failed: {dl_msg}", file=sys.stderr) + logger.warning("LM model download failed: {}", dl_msg) @@ - print(f"Warning: Failed to download LM model: {exc}", file=sys.stderr) + logger.warning("Failed to download LM model: {}", exc) @@ - print(f"Warning: 5Hz LM initialization failed: {lm_status}", file=sys.stderr) + logger.warning("5Hz LM initialization failed: {}", lm_status)Also applies to: 155-155, 157-157, 180-180
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@acestep/gradio_pipeline_startup.py` around lines 137 - 140, Replace the warning and error-style print calls in the startup flow, including the messages around LM initialization and the additional referenced locations, with the imported loguru logger. Use the appropriate logger.warning or logger.error level while preserving each existing message and keeping genuine CLI output unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@acestep/gradio_pipeline_mode_defaults.py`:
- Line 37: Update the model-size checks and replacement logic in the defaults
configuration to match the “-4B” delimited token rather than the bare “4B”
substring. Apply this consistently to the condition near the offload handling
and the related logic at the other affected locations, preventing “14B” model
names from being misclassified or altered.
---
Nitpick comments:
In `@acestep/gradio_pipeline_mode_defaults.py`:
- Around line 48-51: Replace the print call in the 4B-to-1.7B fallback warning
with loguru’s logger.warning(), preserving the existing warning message and
fallback details. Ensure the module imports logger consistently with
gradio_pipeline_cli.py.
In `@acestep/gradio_pipeline_startup.py`:
- Around line 137-140: Replace the warning and error-style print calls in the
startup flow, including the messages around LM initialization and the additional
referenced locations, with the imported loguru logger. Use the appropriate
logger.warning or logger.error level while preserving each existing message and
keeping genuine CLI output unchanged.
In `@acestep/ui/gradio/events/generation/service_init_test.py`:
- Around line 99-113: Update the init_service_wrapper call in
test_project_root_is_consistent_with_checkpoint_dir to replace the consecutive
positional False values with the corresponding keyword arguments, matching the
device resolution test’s established pattern while preserving all existing
argument values.
- Around line 51-65: Update the init_service_wrapper call in
test_passes_project_root_not_checkpoint_dir to pass all boolean parameters as
keyword arguments, matching the parameter names and style used in
service_init_device_resolution_test.py. Keep the existing boolean values and
non-boolean positional arguments unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2fd849de-0883-4e9a-9240-9f52a5fd713f
📒 Files selected for processing (13)
acestep/acestep_v15_pipeline.pyacestep/acestep_v15_pipeline_gpu_mapping_test.pyacestep/acestep_v15_pipeline_test.pyacestep/core/generation/handler/conditioning_embed_test.pyacestep/gradio_pipeline_banner.pyacestep/gradio_pipeline_cli.pyacestep/gradio_pipeline_cli_service.pyacestep/gradio_pipeline_launch.pyacestep/gradio_pipeline_mode_defaults.pyacestep/gradio_pipeline_startup.pyacestep/llm_inference.pyacestep/ui/gradio/events/generation/service_init_device_resolution_test.pyacestep/ui/gradio/events/generation/service_init_test.py
🚧 Files skipped from review as they are similar to previous changes (2)
- acestep/core/generation/handler/conditioning_embed_test.py
- acestep/llm_inference.py
Avoid false positives on names like 14B when auto-offloading or downgrading large LMs on limited VRAM at Gradio startup. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--gpu-mappingand--list-gpusto the Gradio pipeline (acestep) and generation CLI (cli.py)gpus,gpu_mapping, anddevice_mapin API/health,/v1/models, and/v1/model_inventoryresponsesACESTEP_LM_DEVICEdeprecation in favor ofACESTEP_GPU_MAPPINGdocs/en/MULTI_GPU.mdand cross-link fromGPU_COMPATIBILITY.mdStack
Depends on #1262 (PR1: device map infrastructure) and #1263 (PR2: auto-layout + cross-GPU routing).
Test plan
python -m unittest acestep.test_device_map acestep.acestep_v15_pipeline_test acestep.api.http.model_service_routes_testuv run acestep --list-gpuson a multi-GPU hostACESTEP_GPU_MAPPING=auto uv run acestep --config-path acestep-v15-xl-sft --init-llm truecurl /healthreturnsgpus,gpu_mapping,device_mapafter API startupMade with Cursor
Summary by CodeRabbit
--gpu-mapping(supportsautoand per-component placement) and--list-gpus(prints detected GPUs and exits)./health,/v1/models, and/v1/model_inventorynow include GPU runtime fields:gpus,gpu_mapping, anddevice_map.--list-gpus, GPU mapping wiring, device-map routing, and new GPU fields in responses.