[LukGudDev]

Summary

On Windows with ACESTEP_OFFLOAD_DIT_TO_CPU=true, the server's private memory grows by 4–12 GB per generation until the process dies. This is very likely what #1338 reported (an access violation in c10.dll after ~8–9 offloaded generations on Windows, closed as stale).

Cause. _run_generate_music_service_with_progress starts a new threading.Thread for every generation. PyTorch's Windows builds allocate CPU tensors with mimalloc (compiled into c10.dll), which keeps one heap per thread and abandons it when the thread exits. Each generation's thread allocates the offloaded DiT's new CPU copy. The next generation frees that copy from a different thread, after the owning thread has exited, and the memory is not reused. Every generation leaves one DiT-sized block behind.

#947 handled the Linux side of this with mallopt/malloc_trim. Windows has nothing equivalent to call, because mimalloc lives inside c10.dll and exports nothing.

Fix. A new service_thread.run_on_service_thread(fn, timeout) runs each call on an idle long-lived daemon thread. It starts a new thread only when every existing one is busy. Sequential generations therefore always run on the same thread, so the DiT's CPU copy is allocated and freed in a live heap and gets reused. generate_music_execute.py calls it instead of creating a thread. The timeout check and the error it raises are unchanged.

Evidence

Setup: RTX 5070 Ti Laptop (12 GB), 64 GB RAM, Windows 11, Python 3.12.10, torch 2.7.1+cu128, REST API server, acestep-v15-xl-turbo with acestep-5Hz-lm-1.7B (pt backend), ACESTEP_OFFLOAD_DIT_TO_CPU=true.

Server, back-to-back 60 s generations:

Build	Server private memory
main	+4 GB per generation (50 → 54 → 58.6 GB); one silent crash
this fix	flat at 42.6 GB over 4 generations

Causes ruled out before this one:

• Python references. A census of live CPU tensors stayed at 11.04 GB while private memory went 18.7 → 31.7 → 44.2 → 56.6 GB.
• Pinned memory. The pinned host cache was 0.
• The Windows CRT heap. HeapCompact/_heapmin changed nothing.
• mimalloc purge settings. MIMALLOC_PURGE_DELAY=0 and MIMALLOC_ABANDONED_PAGE_PURGE changed nothing.

Plain-PyTorch repro, no ACE-Step code. A 6 GB bf16 module is moved CPU↔CUDA; private memory is shown per cycle.

Where the moves run	Private memory
Main thread	flat at 9.0 GB
New short-lived thread per move	16 → 23 → 30 → 37 GB
Two long-lived threads	flat at 16.1 GB

The same repro against this PR's module. A 3 GB module is moved CPU→CUDA→CPU once per "take", like one offloaded generation.

Where each take runs	Private memory after takes 1–5
New thread per take (main's pattern)	10.1 → 13.6 → 17.1 → 20.6 → 24.1 GB
run_on_service_thread	10.1 → 10.0 → 10.0 → 10.0 → 10.0 GB

The server numbers come from a first version of this change that used a single worker thread. For sequential generations, the final version takes the same path: every call runs on one reused thread, which test_sequential_calls_reuse_one_thread covers. The last table was measured with the final module.

Repro script (Windows, CUDA)

"""Move a ~3 GB module CPU->CUDA->CPU once per 'take' and print private memory.

python repro.py fresh    # new thread per take, like main
python repro.py reused   # run_on_service_thread, like this PR
"""
import ctypes
import gc
import sys
import threading
from ctypes import wintypes

import torch

from acestep.core.generation.handler.service_thread import run_on_service_thread


class PMC(ctypes.Structure):
    _fields_ = [("cb", wintypes.DWORD), ("PageFaultCount", wintypes.DWORD)] + [
        (n, ctypes.c_size_t) for n in ("PeakWorkingSetSize", "WorkingSetSize", "a", "b", "c",
                                       "d", "PagefileUsage", "PeakPagefileUsage", "PrivateUsage")]


def private_gb():
    k32, psapi = ctypes.WinDLL("kernel32"), ctypes.WinDLL("psapi")
    k32.GetCurrentProcess.restype = wintypes.HANDLE
    psapi.GetProcessMemoryInfo.argtypes = [wintypes.HANDLE, ctypes.c_void_p, wintypes.DWORD]
    c = PMC(); c.cb = ctypes.sizeof(PMC)
    psapi.GetProcessMemoryInfo(k32.GetCurrentProcess(), ctypes.byref(c), c.cb)
    return round(c.PrivateUsage / 2**30, 1)


m = torch.nn.Sequential(*[torch.nn.Linear(4096, 4096, bias=False) for _ in range(90)])
m = m.to(torch.bfloat16)


def take():
    m.to("cuda"); torch.cuda.synchronize()
    m.to("cpu")


for i in range(5):
    if sys.argv[1] == "fresh":
        t = threading.Thread(target=take); t.start(); t.join()
    else:
        run_on_service_thread(take, timeout=120)
    gc.collect(); torch.cuda.empty_cache()
    print(f"take {i + 1}: private {private_gb()} GB")

Scope

• acestep/core/generation/handler/service_thread.py (new, 73 lines): the long-lived threads.
• acestep/core/generation/handler/service_thread_test.py (new): 4 tests. They cover reusing one thread for sequential calls, the timeout result with the next call not blocked, logging an exception and keeping the thread, and not keeping a finished call's closure alive.
• acestep/core/generation/handler/generate_music_execute.py: calls run_on_service_thread instead of starting a threading.Thread per call.
• acestep/core/generation/handler/generate_music_execute_test.py: 1 new test for the timeout branch (raises TimeoutError and stops the progress estimator).

Out of scope:

• The Windows RSS log line, which always reads "RSS: 0 MB". That fix is in a separate PR (fix(memory): Windows RSS readout always reported 0 MB #1369).
• The timeout default and the abandoned-work problem in SFT (non-turbo) DiT exceeds the 600 s default ACESTEP_GENERATION_TIMEOUT, and the abandoned service-generate thread keeps running #1350. They are unchanged here.

Risk and compatibility

• Cross-platform on purpose. The per-call thread lives in platform-independent code, so every platform now reuses threads. The leak was measured on Windows only. Concurrency, the timeout, the daemon flag and the thread name (service-generate) are unchanged. The only difference is that a call may run on a thread that ran an earlier call. Gating this to Windows would keep two code paths for no functional gain, but I'm happy to gate it if you prefer.
• Concurrent calls (e.g. ACESTEP_API_WORKERS > 1) each get their own thread immediately, as before. Nothing queues behind another call, so timeouts are measured exactly as before.
• After a timeout, the stuck call keeps its thread busy and the next call gets a new thread, as before. When the stuck call finishes, its thread returns to the pool.
• Thread count is bounded by the peak number of simultaneous calls, including timed-out calls that are still running. With the default single API worker that is 1 thread, or 2 after a timeout. Threads never exit.
• Daemon threads, as before, so a hung CUDA call can't block interpreter exit. ThreadPoolExecutor workers would block it.
• No state carried between generations. An idle thread drops its last call before going idle, so a finished generation's inputs, outputs and exception frames aren't kept alive (tested). The service_generate path sets no thread-local state that would persist on a reused thread: inference_mode is restored on exit, and the path has no set_device, set_stream, autocast, threading.local or ContextVar.

Regression checks

• uv run python -m unittest acestep.core.generation.handler.service_thread_test acestep.core.generation.handler.generate_music_execute_test acestep.core.generation.handler.generate_music_test acestep.core.generation.handler.service_generate_test: 20 tests, OK (Windows 11, Python 3.12.10).
• service_thread_test run 200 times in one process: 0 failures, and only 2 service-generate threads alive at the end.
• An independent review ran a concurrency check against the module: 8 caller threads × 300 jobs of random length. Every job ran exactly once, no thread ran two jobs at once, and the thread count equalled peak concurrency.
• The rest of acestep/core/generation/handler has the same 5 pre-existing test errors on main and on this branch. The affected files are init_service_test, io_audio_test, reference_audio_validation_test, service_generate_flow_edit_test and vae_decode_chunks_test.
• The server and repro measurements above.
• Not run: Linux, macOS/MPS, ROCm or XPU hardware. The new tests are platform-independent.

Reviewer notes

• Only the service_generate call moves. Offloads outside it, such as the VAE context in the decode step, run on the caller's thread.
  • REST API server: that thread is the API's long-lived executor thread, and the flat 42.6 GB above includes those offloads.
  • Gradio UI: callers run on anyio worker threads, which exit after 10 s idle, so the VAE's smaller CPU copy may still follow the old pattern there. I haven't measured it.
• feat(gradio): surface live generation progress across LM, diffusion, and decode #954 (open) also edits this gen_thread block, to poll it for progress. If it lands first, I'll rebase and let it poll the call's completion event instead.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Generation requests now report a timeout error when the service call does not finish within the configured limit, rather than waiting indefinitely.
  • Errors raised during generation continue to be surfaced to the caller.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 5d977b45-0047-4daa-ab4b-e6a3d9b6d5ac

Commits

Reviewing files that changed from the base of the PR and between ca1e85f and 2422ec4.

Files selected for processing: 4

• acestep/core/generation/handler/generate_music_execute.py
• acestep/core/generation/handler/generate_music_execute_test.py
• acestep/core/generation/handler/service_thread.py
• acestep/core/generation/handler/service_thread_test.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

Walkthrough

Music generation now runs its service callable through a reusable daemon worker pool. The caller waits up to the existing generation timeout. A timeout returns an error to the caller but does not stop the running callable.

Changes: Service Thread Execution

Layer / File(s)	Summary
Worker reuse and lifecycle acestep/core/generation/handler/service_thread.py, acestep/core/generation/handler/service_thread_test.py	The new runner reuses idle daemon workers and starts a worker when none is idle. Tests cover sequential reuse, timeout behavior, callable cleanup, and exception logging.
Music generation integration acestep/core/generation/handler/generate_music_execute.py, acestep/core/generation/handler/generate_music_execute_test.py	Music generation uses the worker runner with the existing timeout. If the runner returns False, the helper raises TimeoutError; the test also checks that progress stops.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Generation as generate_music_execute
  participant Runner as run_on_service_thread
  participant Worker as service worker
  Generation->>Runner: Submit service callable and generation timeout
  Runner->>Worker: Enqueue callable
  Worker->>Worker: Execute callable
  alt Callable completes before timeout
    Worker-->>Runner: Signal completion
    Runner-->>Generation: Return true
  else Caller reaches timeout
    Runner-->>Generation: Return false
    Worker->>Worker: Continue running callable
  end

Loading

Possibly related issues

• ace-step/ACE-Step-1.5, issue 1350: The issue concerns work that continues after a service_generate timeout. This change retains that timeout behavior, so it does not address the issue.

Suggested reviewers: uni404x64

Merge Risk: Minimal · up to 2422e

No actionable issue was found in worker reuse or its music-generation integration. The change is ready to merge after normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 78.57% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 14 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title names the Windows RAM leak fix and the mechanism (reusing service_generate threads). The changeset adds run_on_service_thread, which reuses idle daemon worker threads for service_generate ca…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title names the Windows RAM leak fix and the mechanism (reusing service_generate threads). The changeset adds run_on_service_thread, which reuses idle daemon worker threads for service_generate calls, and updates generate_music_execute.py to use it. The title is specific and matches the main change. It does not name the DiT offload trigger, but the title still covers the primary change. The title also follows the conventional commit prefix style.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{A rabbit sends a task to run,
The worker hums beneath the sun.
If time runs out, the task may stay,
The caller leaves without delay.
Then idle threads await the day.}

_{Comment @coderabbitai help to get the list of available commands.}