From 16bccbeb6f328a1af18d311cac3a8b89e5ed45f6 Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 6 Oct 2026 08:32:04 -0400 Subject: [PATCH 1/5] Free tracked contexts at exit before ggml-metal teardown ggml-metal aborts in its static destructor, with GGML_ASSERT([rsets->data count] == 0), when a process exits while a Metal buffer is still allocated. A Flutter macOS quit calls exit() without shutting the Dart isolates down, a hot restart drops them without running finalizers, and an isolate that is killed during new_sd_ctx never sees the context that call returns, so the caller cannot always free its contexts first. src/sd_dart_exit.cpp keeps a mutex-protected registry of live objects, the design of leehack/llamadart-native#99 for this library's own ggml: - sd_dart_new_sd_ctx tracks the context before it returns, and sd_dart_exit_free untracks an object inside the call that frees it, so nothing is left untracked or freed twice whenever an isolate is killed. sd_dart_exit_track tracks any other object for a native caller. - sd_dart_generate_image is a call in flight. Teardown asks every tracked context to cancel, repeats that until the calls in flight have ended, since generate_image clears a cancellation that precedes it, and frees only then, by stage. If a call is still running after wait_ms it frees nothing. - sd_dart_cancel_generation cancels a tracked context under the registry lock, so a cancellation from another thread cannot reach a context that teardown has freed. - Once teardown has begun, a thread that ends its outermost call in flight or reaches the registry outside one never returns to its caller. - On Apple platforms teardown runs during exit(), before the first static of the library is destroyed: the library defines __cxa_atexit for its own image, hidden and not exported, and registers every static destructor behind a call to teardown. A plain atexit handler only precedes the statics that exist when it is registered, and ggml-metal creates its device and buffer types on first use. validate_artifacts.py and apple_xcframework.py validate fail a slice that imports the system function. A context from new_sd_ctx is not tracked, and exiting with one alive behaves as before. With nothing tracked and no load in flight, teardown returns at once. tests/test_exit_teardown.py runs the registry against stand-ins for upstream, also under AddressSanitizer and ThreadSanitizer. Refs #10 --- AGENTS.md | 17 + CMakeLists.txt | 1 + README.md | 103 ++++ src/sd_dart_exit.cpp | 344 ++++++++++++ src/sd_dart_wrapper.h | 120 ++++ tests/native/exit_teardown_test.cpp | 829 ++++++++++++++++++++++++++++ tests/test_exit_teardown.py | 92 +++ tests/test_tools.py | 58 +- tools/apple_xcframework.py | 9 +- tools/validate_artifacts.py | 19 +- 10 files changed, 1588 insertions(+), 4 deletions(-) create mode 100644 src/sd_dart_exit.cpp create mode 100644 tests/native/exit_teardown_test.cpp create mode 100644 tests/test_exit_teardown.py diff --git a/AGENTS.md b/AGENTS.md index 21fe847..176d447 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,6 +17,14 @@ Guidance for coding agents working in `stable-diffusion-native`. would call, and never waits for such a thread: a Dart callback there aborts or deadlocks VM shutdown. Record state natively and let the caller poll. Wrapper state has no destructor; threads still use it while `exit()` runs. + The one wait is exit teardown's, which is bounded and frees nothing when it + runs out. +- Whatever a Dart caller does to a tracked object begins and ends in one + native call: tracking in the call that creates it, untracking in the call + that frees it, and a call in flight around anything that can outlast the + 250 ms settle time. An isolate can be killed between any two native calls. + Give such an upstream function an `sd_dart_` wrapper before `llamadart` + calls it on a tracked context. - Whatever a caller polls must keep its history. Consumers derive state from the order of progress reports, such as which image of a batch is being sampled, so a "latest value" export loses what happened between two polls. @@ -33,11 +41,20 @@ python3 tools/validate_artifacts.py - `tests/test_progress_stress.py` needs a Dart SDK on `PATH` and skips without one; CI runs it with 25 runs per configuration. Run it after any change to `src/`. +- `tests/test_exit_teardown_runtime.py` needs the built macOS runtime and + skips without it. After a change to `src/sd_dart_exit.cpp` or an upstream + bump, run it against the release build and again with + `SD_EXIT_TEARDOWN_SANITIZER=address`, on a Mac where its control reports the + Metal abort: GitHub's runners do not reach it. - Every shipped library must pass `validate_artifacts.py`: it exports exactly the `SD_API` symbols in `stable-diffusion.h` and `src/sd_dart_wrapper.h`, and links only allowlisted system libraries. Never export an upstream-internal or ggml symbol to fix a consumer; a leaked ggml symbol can collide with llama.cpp's ggml in the same process. +- An Apple library must not import `__cxa_atexit`: exit teardown depends on + the hidden definition in `src/sd_dart_exit.cpp` receiving every static + destructor of the image. Keep that definition hidden; libllamadart has its + own. - The Apple XCFramework must pass `apple_xcframework.py validate` and `consumer`. Its Info.plist minimum OS is read from each binary's `LC_BUILD_VERSION`; never hard-code it, since a mismatch fails App Store diff --git a/CMakeLists.txt b/CMakeLists.txt index 3ebbbc5..010f895 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -20,5 +20,6 @@ add_subdirectory("${SD_NATIVE_UPSTREAM_DIR}" sd) # Wrapper exports are compiled into the upstream library target, so they share # its SD_API export macro and reach upstream without patching it. target_sources(stable-diffusion PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/src/sd_dart_exit.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/src/sd_dart_wrapper.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/src/sd_dart_wrapper.h") diff --git a/README.md b/README.md index 82545dd..a472024 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,109 @@ do { dies while a native call is reporting progress, and one that polls a batch of images; see its docstring. +### Exit teardown + +```c +sd_ctx_t* sd_dart_new_sd_ctx(const sd_ctx_params_t* sd_ctx_params); +bool sd_dart_generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, int* num_images_out); +void sd_dart_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); +void sd_dart_exit_free(void* object); + +bool sd_dart_exit_track(void* object, void (*free_fn)(void*), int32_t stage); +bool sd_dart_exit_untrack(void* object); +int32_t sd_dart_exit_tracked_count(void); +void sd_dart_exit_call_begin(void); +void sd_dart_exit_call_end(void); +void sd_dart_exit_set_wait_ms(int32_t wait_ms); +void sd_dart_exit_teardown(void); +``` + +ggml-metal aborts in its static destructor +(`GGML_ASSERT([rsets->data count] == 0)`) when a process exits while a Metal +buffer is still allocated. A caller in a managed runtime cannot always free +its contexts first: a Flutter macOS quit calls `exit()` without shutting the +Dart isolates down, a hot restart drops them without running finalizers, and +an isolate that is killed during `new_sd_ctx` never sees the context that +call returns. So the library keeps a registry of live objects and, on Apple +platforms, frees what is left of it during `exit()`, before the first of its +statics is destroyed. + +| Function | Behavior | +| --- | --- | +| `sd_dart_new_sd_ctx` | `new_sd_ctx` that tracks the context before it returns. | +| `sd_dart_generate_image` | `generate_image` as a call in flight: teardown cancels it and waits for it. | +| `sd_dart_cancel_generation` | `sd_cancel_generation` for a tracked context, from any thread. Does nothing for a context that is not tracked, including one already freed. Never blocks. | +| `sd_dart_exit_free` | Untracks an object and frees it with the function it was tracked with (`free_sd_ctx` for a context). Does nothing when it is not tracked, so it also works as a Dart `NativeFinalizer`. | +| `sd_dart_exit_track`, `_untrack` | For C and C++ callers: track any other object, such as an `upscaler_ctx_t`, with its free function and a stage, or stop tracking it. | +| `sd_dart_exit_tracked_count` | Number of tracked objects. | +| `sd_dart_exit_call_begin`, `_end` | For C and C++ callers: mark a call in flight around an upstream function that has no wrapper. | +| `sd_dart_exit_set_wait_ms` | How long teardown waits for calls in flight; 2000 ms by default. | +| `sd_dart_exit_teardown` | Runs teardown now, for native hosts on platforms where it does not run by itself. Follow it directly with `exit()`. | + +A Dart caller replaces `new_sd_ctx`, `generate_image`, `sd_cancel_generation` +and `free_sd_ctx` with the first four and binds nothing else: the remaining +functions are for native hosts, and a Dart program that calls +`sd_dart_exit_teardown` and returns from `main` waits forever for its blocked +isolates. + +What is guaranteed: + +- **Tracking begins and ends in native code.** A context is tracked before + the call that creates it returns and untracked inside the call that frees + it, so an isolate that is killed in between leaves nothing untracked and + nothing is freed twice. +- **No free under a running call.** Teardown asks every tracked context to + cancel its generation, repeating the request until the calls in flight + have ended, and frees only then. If one is still running after the wait, + it frees nothing at all rather than what happens to be idle. +- **A bounded exit.** The wait ends after `wait_ms`, plus 250 ms for a thread + that has just left a call to finish the short calls that follow it. With + nothing tracked and no load in flight, teardown returns at once. +- **Nothing returns into freed memory.** Once teardown has begun, a thread + that ends its outermost call in flight, or reaches the registry outside + one, never returns to its caller. +- **Order.** Objects are freed by stage, an object before what it uses, and + latest first within a stage. On Apple platforms that happens before any + static of the library is destroyed, whenever the static was created: the + library defines `__cxa_atexit` for its own image, hidden and not exported, + and registers every static destructor behind a call to teardown. + `validate_artifacts.py` and `apple_xcframework.py validate` fail a slice + that imports the system function instead. +- **Callers that do not use it are unaffected.** A context from `new_sd_ctx` + is not tracked, and exiting with one alive behaves as before, the Metal + abort included. Another ggml in the process, such as libllamadart's with + its own copy of this registry, is not touched either. + +What is not: + +- **A load cannot be cancelled.** Teardown waits for one like for any call, + and a load that outlasts `wait_ms` leaves everything allocated, so Metal + aborts as before. +- **A sampling step cannot be interrupted.** stable-diffusion.cpp honors a + cancellation between steps. A step that outlasts `wait_ms` has the same + result as a load that does. +- **Upstream calls on a tracked context are not waited for.** A raw + `generate_image`, `generate_video` or `adetail_image` on a context from + `sd_dart_new_sd_ctx` is a use after free at exit. Native callers bracket + such a call with `sd_dart_exit_call_begin` and `_end`; Dart callers cannot, + because a killed isolate never reaches the end. +- **Platforms other than Apple's.** The abort is ggml-metal's, so elsewhere + teardown runs only when a native host calls it. CI checks on Linux that a + context left alive at exit is harmless on the CPU, and reports what the + Vulkan backend does on Mesa lavapipe; no hardware Vulkan driver was tried. +- **Threads that teardown blocked stay blocked.** A static destructor or + `atexit` handler of another library that joins one hangs the exit. + +`tests/test_exit_teardown.py` runs the registry against stand-ins for +upstream, also under AddressSanitizer and ThreadSanitizer. +`tests/test_exit_teardown_runtime.py` runs it against a built macOS runtime +and a real `sd_ctx_t`, on Metal and on the CPU: the test writes a 12 MB +PixArt model itself, so CI needs no download. Where no Metal residency set is +live, as on GitHub's macOS runners, an untracked context exits cleanly and the +Metal abort is not exercised; there the test shows each free through the +allocator instead. + ## Build ```bash diff --git a/src/sd_dart_exit.cpp b/src/sd_dart_exit.cpp new file mode 100644 index 0000000..7376d07 --- /dev/null +++ b/src/sd_dart_exit.cpp @@ -0,0 +1,344 @@ +#include "sd_dart_wrapper.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__APPLE__) +#include +#endif + +namespace { + +struct Entry { + void (*free_fn)(void*) = nullptr; + // Asks the object to end the call in flight on it early. May be null. + void (*cancel_fn)(void*) = nullptr; + int32_t stage = 0; + uint64_t order = 0; +}; + +// A thread that just left a call in flight usually still makes short calls on +// the same objects, such as reading the model version after a load. Teardown +// gives it this long to finish them. +constexpr std::chrono::milliseconds kSettleTime{250}; +// How often teardown repeats its cancellation while it waits. +constexpr std::chrono::milliseconds kCancelInterval{20}; + +struct Registry { + std::mutex mutex; + std::condition_variable idle; + std::unordered_map objects; + std::chrono::steady_clock::time_point last_call_end{}; + uint64_t next_order = 0; + int32_t calls = 0; + int32_t creating_calls = 0; + int32_t wait_ms = 2000; + std::atomic armed{false}; + std::atomic torn_down{false}; +}; + +Registry& state() { + // Never destroyed: other threads still reach it after static destructors. + static auto* registry = new Registry(); + return *registry; +} + +thread_local bool teardown_thread = false; + +// Calls in flight on this thread. The registry counts the outermost one. +thread_local int32_t call_depth = 0; + +// Called with the registry locked. Returns whether the calling thread may go +// on to use tracked objects. Once teardown has begun, teardown waits for a +// thread that is in a call in flight, and the teardown thread gets false. Any +// other thread may hold objects that teardown frees, so it never returns to +// its caller. +bool admit(Registry& registry, std::unique_lock& lock) { + if (!registry.torn_down) { + return true; + } + if (teardown_thread) { + return false; + } + if (call_depth > 0) { + return true; + } + lock.unlock(); + for (;;) { + std::this_thread::sleep_for(std::chrono::hours(1)); + } +} + +#if defined(__APPLE__) +struct StaticDestructor { + void (*destroy)(void*); + void* object; +}; + +void destroy_static(void* argument) { + const StaticDestructor destructor = *static_cast(argument); + free(argument); + sd_dart_exit_teardown(); + destructor.destroy(destructor.object); +} + +using RegisterDestructor = int (*)(void (*)(void*), void*, void*); + +std::atomic system_cxa_atexit{nullptr}; +#endif + +// Registers teardown to run at exit once something is tracked, so that it +// does not wait for the first static of this library to be destroyed. atexit +// handlers run in reverse order of registration, so this one runs after the +// destructors of statics created later; __cxa_atexit below covers those. +void arm() { +#if defined(__APPLE__) + if (!state().armed.exchange(true)) { + atexit(sd_dart_exit_teardown); + } +#endif +} + +bool insert(void* object, void (*free_fn)(void*), void (*cancel_fn)(void*), int32_t stage) { + if (object == nullptr || free_fn == nullptr || + stage < SD_DART_EXIT_STAGE_CONTEXT_USER || stage > SD_DART_EXIT_STAGE_RESOURCE) { + return false; + } + Registry& registry = state(); + { + std::unique_lock lock(registry.mutex); + if (!admit(registry, lock)) { + return false; + } + registry.objects[object] = {free_fn, cancel_fn, stage, registry.next_order++}; + } + arm(); + return true; +} + +struct Call { + Call() { sd_dart_exit_call_begin(); } + ~Call() { sd_dart_exit_call_end(); } + Call(const Call&) = delete; + Call& operator=(const Call&) = delete; +}; + +// A call in flight that tracks what it creates. Teardown waits for it even +// when nothing is tracked yet. +struct CreatingCall { + CreatingCall() { + sd_dart_exit_call_begin(); + Registry& registry = state(); + std::lock_guard lock(registry.mutex); + ++registry.creating_calls; + } + ~CreatingCall() { + { + Registry& registry = state(); + std::lock_guard lock(registry.mutex); + --registry.creating_calls; + } + sd_dart_exit_call_end(); + } + CreatingCall(const CreatingCall&) = delete; + CreatingCall& operator=(const CreatingCall&) = delete; +}; + +void free_context(void* object) { + free_sd_ctx(static_cast(object)); +} + +void cancel_context(void* object) { + sd_cancel_generation(static_cast(object), SD_CANCEL_ALL); +} + +} // namespace + +#if defined(__APPLE__) +extern "C" { +// The C++ runtime registers the destructor of every static in this library +// through __cxa_atexit, and the linker binds those calls to this definition. +// Tracked objects use such statics, some of which stable-diffusion.cpp and +// ggml create on first use at any time, so teardown has to run before the +// first of them is destroyed: each destructor is registered behind a call to +// teardown. A plain atexit handler cannot do that, as it only precedes the +// statics that exist when it is registered. +__attribute__((visibility("hidden"))) int __cxa_atexit(void (*destroy)(void*), void* object, void* dso_handle) { + RegisterDestructor system_register = system_cxa_atexit.load(); + if (system_register == nullptr) { + system_register = reinterpret_cast(dlsym(RTLD_NEXT, "__cxa_atexit")); + if (system_register == nullptr) { + return -1; + } + system_cxa_atexit.store(system_register); + } + auto* destructor = static_cast(malloc(sizeof(StaticDestructor))); + if (destructor == nullptr) { + return -1; + } + *destructor = {destroy, object}; + const int status = system_register(destroy_static, destructor, dso_handle); + if (status != 0) { + free(destructor); + } + return status; +} +} +#endif + +bool sd_dart_exit_track(void* object, void (*free_fn)(void*), int32_t stage) { + return insert(object, free_fn, nullptr, stage); +} + +bool sd_dart_exit_untrack(void* object) { + Registry& registry = state(); + std::unique_lock lock(registry.mutex); + if (!admit(registry, lock)) { + return false; + } + return registry.objects.erase(object) != 0; +} + +void sd_dart_exit_free(void* object) { + if (object == nullptr) { + return; + } + Registry& registry = state(); + void (*free_fn)(void*) = nullptr; + { + std::unique_lock lock(registry.mutex); + if (!admit(registry, lock)) { + return; + } + const auto found = registry.objects.find(object); + if (found == registry.objects.end()) { + return; + } + free_fn = found->second.free_fn; + registry.objects.erase(found); + if (call_depth++ == 0) { + ++registry.calls; + } + } + free_fn(object); + sd_dart_exit_call_end(); +} + +int32_t sd_dart_exit_tracked_count(void) { + Registry& registry = state(); + std::lock_guard lock(registry.mutex); + return static_cast(registry.objects.size()); +} + +void sd_dart_exit_call_begin(void) { + if (call_depth > 0) { + ++call_depth; + return; + } + Registry& registry = state(); + std::unique_lock lock(registry.mutex); + if (!admit(registry, lock)) { + return; + } + ++registry.calls; + call_depth = 1; +} + +void sd_dart_exit_call_end(void) { + if (call_depth == 0 || --call_depth > 0) { + return; + } + Registry& registry = state(); + std::unique_lock lock(registry.mutex); + --registry.calls; + registry.idle.notify_all(); + if (admit(registry, lock)) { + registry.last_call_end = std::chrono::steady_clock::now(); + } +} + +void sd_dart_exit_set_wait_ms(int32_t wait_ms) { + Registry& registry = state(); + std::lock_guard lock(registry.mutex); + registry.wait_ms = std::max(wait_ms, 0); +} + +void sd_dart_exit_teardown(void) { + Registry& registry = state(); + std::vector> objects; + { + std::unique_lock lock(registry.mutex); + if (registry.torn_down.exchange(true)) { + return; + } + teardown_thread = true; + // With nothing to free and nothing being created, a call in flight is + // no reason to hold up the exit. + if (registry.creating_calls == 0 && registry.objects.empty()) { + return; + } + // A call in flight on this thread cannot end while teardown runs. + const int32_t own_calls = call_depth > 0 ? 1 : 0; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(registry.wait_ms); + while (registry.calls != own_calls) { + // generate_image() clears a cancellation that was requested before + // it started, so the request is repeated until the calls end. + for (const auto& [object, entry] : registry.objects) { + if (entry.cancel_fn != nullptr) { + entry.cancel_fn(object); + } + } + const auto now = std::chrono::steady_clock::now(); + if (now >= deadline) { + return; + } + registry.idle.wait_until(lock, std::min(deadline, now + kCancelInterval)); + } + registry.idle.wait_until(lock, registry.last_call_end + kSettleTime, [] { return false; }); + objects.assign(registry.objects.begin(), registry.objects.end()); + } + std::sort(objects.begin(), objects.end(), [](const auto& a, const auto& b) { + return a.second.stage != b.second.stage ? a.second.stage < b.second.stage : a.second.order > b.second.order; + }); + for (const auto& [object, entry] : objects) { + { + std::lock_guard lock(registry.mutex); + registry.objects.erase(object); + } + entry.free_fn(object); + } +} + +sd_ctx_t* sd_dart_new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { + CreatingCall call; + sd_ctx_t* context = new_sd_ctx(sd_ctx_params); + insert(context, free_context, cancel_context, SD_DART_EXIT_STAGE_CONTEXT); + return context; +} + +bool sd_dart_generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out) { + Call call; + return generate_image(sd_ctx, sd_img_gen_params, images_out, num_images_out); +} + +void sd_dart_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { + Registry& registry = state(); + // Teardown and sd_dart_exit_free() untrack a context under this lock + // before they free it, so a tracked context is alive while it is held. + std::lock_guard lock(registry.mutex); + const auto found = registry.objects.find(sd_ctx); + if (found != registry.objects.end() && found->second.cancel_fn == cancel_context) { + sd_cancel_generation(sd_ctx, mode); + } +} diff --git a/src/sd_dart_wrapper.h b/src/sd_dart_wrapper.h index 9af128f..bc55914 100644 --- a/src/sd_dart_wrapper.h +++ b/src/sd_dart_wrapper.h @@ -64,6 +64,126 @@ SD_API size_t sd_dart_progress_read(uint64_t after, size_t capacity, uint64_t* latest); +// Exit teardown +// +// The library keeps a registry of live native objects and frees what is left +// of it when the process exits without freeing them itself, as a Flutter quit +// or hot restart does. ggml-metal aborts in its static destructor while any +// Metal buffer is still allocated, so teardown has to run before that +// destructor. On Apple platforms it runs during exit(), before the first +// static of this library is destroyed. Elsewhere, where that abort does not +// exist, it runs only when sd_dart_exit_teardown() is called. +// +// Objects are tracked by sd_dart_new_sd_ctx(), before it returns, and by +// sd_dart_exit_track(). +// +// Teardown asks every tracked context to cancel its generation, waits a +// bounded time for the calls in flight and then frees the tracked objects in +// sd_dart_exit_stage order, latest tracked first within a stage. If a call is +// still in flight when the wait ends, it frees nothing. With nothing tracked +// and no sd_dart_new_sd_ctx() in flight, it does not wait. +// +// A call in flight is the time a thread spends inside sd_dart_new_sd_ctx(), +// sd_dart_generate_image() or sd_dart_exit_free(), or between +// sd_dart_exit_call_begin() and sd_dart_exit_call_end(). +// +// Once teardown has begun, the objects a thread holds may be freed as soon as +// it has no call in flight. From then on, the end of a thread's outermost call +// in flight never returns to its caller. Neither does a function documented +// as "blocks after teardown" when it is called outside a call in flight; +// inside one it works as before, since teardown is waiting for that call. The +// thread that runs teardown is exempt: there, such a function returns without +// tracking, untracking or freeing anything. A thread blocked this way stays +// blocked until the process is gone, so a static destructor or atexit handler +// of another library that joins such a thread hangs the exit. +// +// A native call that is not a call in flight is not waited for. Teardown only +// allows a thread 250 ms after its last call in flight to finish what follows +// it, which covers short calls such as sd_get_model_version_name() after a +// load. A longer call on a tracked context that is not a call in flight, such +// as generate_image() or generate_video(), is a use after free at exit, also +// where exiting with the context alive was harmless. Contexts created by +// new_sd_ctx() are not tracked, and exiting with one alive behaves as it did +// before this registry existed. + +// Order in which exit teardown frees tracked objects: every object of a lower +// stage before any object of a higher one, so an object goes before the +// objects it uses. +enum sd_dart_exit_stage { + // State that uses a context. + SD_DART_EXIT_STAGE_CONTEXT_USER = 0, + // sd_ctx_t and the other contexts of stable-diffusion.h. + SD_DART_EXIT_STAGE_CONTEXT = 1, + // What a context uses. + SD_DART_EXIT_STAGE_RESOURCE = 2, +}; + +// Tracks `object` so that exit teardown frees it with `free_fn`, which +// teardown calls on the thread that runs exit(): pass a native function, never +// a callback into a managed runtime. Tracking an address again replaces its +// entry. Returns false for a null argument or an unknown stage. Blocks after +// teardown. +// +// For C and C++ callers. From Dart an object has to be tracked by the native +// call that creates it: an isolate that is killed during that call never +// reaches a later call that would track it. +SD_API bool sd_dart_exit_track(void* object, void (*free_fn)(void*), int32_t stage); + +// Stops tracking `object` without freeing it. Returns whether it was tracked. +// Blocks after teardown. +SD_API bool sd_dart_exit_untrack(void* object); + +// Stops tracking `object` and frees it with the function it was tracked with: +// free_sd_ctx() for a context from sd_dart_new_sd_ctx(). Does nothing when +// `object` is not tracked, so each tracked object is freed once, whether by +// this call or by teardown. Usable as a Dart NativeFinalizer callback. Blocks +// after teardown. +SD_API void sd_dart_exit_free(void* object); + +// Number of tracked objects. +SD_API int32_t sd_dart_exit_tracked_count(void); + +// Mark the start and end of a call in flight on the calling thread, for C and +// C++ callers of native functions that have no sd_dart_ wrapper. They nest. +// Each begin needs one end on the same thread: a thread that never reaches end +// keeps teardown from freeing anything. Do not call them from Dart, where an +// isolate that is killed during the native call in between never reaches end. +// begin blocks after teardown. +SD_API void sd_dart_exit_call_begin(void); +SD_API void sd_dart_exit_call_end(void); + +// Sets how long teardown waits for calls in flight. Negative values are +// treated as zero. The default is 2000 ms. +SD_API void sd_dart_exit_set_wait_ms(int32_t wait_ms); + +// Runs exit teardown now; later runs do nothing. Afterwards tracked objects +// are unusable and other threads that reach the functions above stay blocked, +// so call it only as the last step before the process exits and follow it +// directly with exit() or _exit() on the same thread. It is meant for native +// hosts. Do not bind it from Dart: a Dart program that returns from main +// after it waits forever for its blocked isolates. +SD_API void sd_dart_exit_teardown(void); + +// new_sd_ctx() that tracks the context in the CONTEXT stage before it +// returns. Free the context with sd_dart_exit_free(). A load cannot be +// cancelled: teardown waits for it like for any call in flight. Blocks after +// teardown. +SD_API sd_ctx_t* sd_dart_new_sd_ctx(const sd_ctx_params_t* sd_ctx_params); + +// generate_image() as a call in flight. Takes and returns what generate_image() +// does. Teardown cancels it with SD_CANCEL_ALL, which stable-diffusion.cpp +// honors between sampling steps. Blocks after teardown. +SD_API bool sd_dart_generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out); + +// sd_cancel_generation() for a tracked context, callable from any thread while +// another one generates. Does nothing when `sd_ctx` is not tracked: one that +// sd_dart_exit_free() or teardown already freed, or one from new_sd_ctx(). +// Never blocks, also after teardown. +SD_API void sd_dart_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); + #ifdef __cplusplus } #endif diff --git a/tests/native/exit_teardown_test.cpp b/tests/native/exit_teardown_test.cpp new file mode 100644 index 0000000..25fb872 --- /dev/null +++ b/tests/native/exit_teardown_test.cpp @@ -0,0 +1,829 @@ +// Exercises src/sd_dart_exit.cpp against stand-ins for the upstream functions +// it wraps, so it needs no model and runs under sanitizers. +// +// Each scenario runs in its own process: teardown runs once per process and +// leaves the registry unusable from other threads. + +#include "sd_dart_wrapper.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__has_feature) +#if __has_feature(address_sanitizer) +#define SD_TEST_ADDRESS_SANITIZER 1 +#endif +#endif +#if defined(__SANITIZE_ADDRESS__) && !defined(SD_TEST_ADDRESS_SANITIZER) +#define SD_TEST_ADDRESS_SANITIZER 1 +#endif + +#if defined(SD_TEST_ADDRESS_SANITIZER) +#include +#elif defined(__APPLE__) +#include +#endif + +#define CHECK(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: CHECK(%s) failed\n", __FILE__, __LINE__, #condition); \ + std::_Exit(1); \ + } \ + } while (0) + +// The stand-in for upstream: a context is a heap block that knows whether it +// is being generated on and whether a cancellation is pending. +struct sd_ctx_t { + std::atomic cancel{SD_CANCEL_RESET}; + std::atomic cancel_requests{0}; + std::atomic generating{false}; +}; + +namespace { + +std::mutex log_mutex; +std::vector log_entries; + +void record(const std::string& event) { + std::lock_guard lock(log_mutex); + log_entries.push_back(event); +} + +std::vector recorded() { + std::lock_guard lock(log_mutex); + return log_entries; +} + +void sleep_ms(int ms) { + std::this_thread::sleep_for(std::chrono::milliseconds(ms)); +} + +int64_t elapsed_ms(std::chrono::steady_clock::time_point since) { + return std::chrono::duration_cast(std::chrono::steady_clock::now() - since).count(); +} + +// Whether the heap block that began at `object` has been freed, for a build +// that can tell: AddressSanitizer poisons a freed block, and the macOS +// allocator reports no size for one until it hands the block out again. +#if defined(SD_TEST_ADDRESS_SANITIZER) +const bool kSeesFreedBlocks = true; +bool is_freed(const void* object) { + return __asan_address_is_poisoned(object) != 0; +} +#elif defined(__APPLE__) +const bool kSeesFreedBlocks = true; +bool is_freed(const void* object) { + return malloc_size(object) == 0; +} +#else +const bool kSeesFreedBlocks = false; +bool is_freed(const void*) { + return false; +} +#endif + +#if defined(__APPLE__) +const bool kTeardownRunsAtExit = true; +#else +const bool kTeardownRunsAtExit = false; +#endif + +// What the stand-ins do, set by the scenario before it calls a wrapper. +std::atomic fail_load{false}; +std::atomic hold_load{nullptr}; +std::atomic run_generation{nullptr}; +std::atomic hold_free{nullptr}; +std::atomic contexts_freed{0}; +std::atomic last_freed{nullptr}; +// The arguments the stand-ins were last called with. +std::atomic last_load_params{nullptr}; +std::atomic last_request{nullptr}; +std::atomic last_images_out{nullptr}; + +// Teardown blocks a thread that reaches the registry outside a call in +// flight. This one keeps reaching it, so when its count stops after the +// scenario asked for teardown, teardown has begun. +std::atomic heartbeat{0}; +std::atomic teardown_requested{false}; + +void start_heartbeat() { + std::thread([] { + for (;;) { + sd_dart_exit_untrack(nullptr); + heartbeat.fetch_add(1); + sleep_ms(1); + } + }).detach(); + while (heartbeat.load() == 0) { + sleep_ms(1); + } +} + +bool wait_for_teardown() { + while (!teardown_requested.load()) { + sleep_ms(1); + } + const auto started = std::chrono::steady_clock::now(); + int64_t seen = heartbeat.load(); + auto changed = std::chrono::steady_clock::now(); + while (elapsed_ms(started) < 20000) { + sleep_ms(5); + const int64_t now = heartbeat.load(); + if (now != seen) { + seen = now; + changed = std::chrono::steady_clock::now(); + } else if (elapsed_ms(changed) > 400) { + return true; + } + } + return false; +} + +void free_named(void* object) { + record(static_cast(object)); +} + +char* name(const char* value) { + return const_cast(value); +} + +char kUserLate[] = "user-late"; + +const std::vector kStageOrder = { + "user-late", + "user-early", + "context-late", + "context-early", + "resource-late", + "resource-early", +}; + +void track_out_of_stage_order() { + const auto track = [](const char* value, int32_t stage) { + CHECK(sd_dart_exit_track(name(value), free_named, stage)); + }; + track("resource-early", SD_DART_EXIT_STAGE_RESOURCE); + track("context-early", SD_DART_EXIT_STAGE_CONTEXT); + track("user-early", SD_DART_EXIT_STAGE_CONTEXT_USER); + track("resource-late", SD_DART_EXIT_STAGE_RESOURCE); + track("context-late", SD_DART_EXIT_STAGE_CONTEXT); + track(kUserLate, SD_DART_EXIT_STAGE_CONTEXT_USER); +} + +int test_dispose() { + CHECK(!sd_dart_exit_track(nullptr, free_named, 0)); + CHECK(!sd_dart_exit_track(name("a"), nullptr, 0)); + CHECK(!sd_dart_exit_track(name("a"), free_named, -1)); + CHECK(!sd_dart_exit_track(name("a"), free_named, SD_DART_EXIT_STAGE_RESOURCE + 1)); + CHECK(sd_dart_exit_tracked_count() == 0); + + char* resource = name("resource"); + char* context = name("context"); + char* kept = name("kept"); + CHECK(sd_dart_exit_track(resource, free_named, SD_DART_EXIT_STAGE_RESOURCE)); + CHECK(sd_dart_exit_track(context, free_named, SD_DART_EXIT_STAGE_CONTEXT)); + CHECK(sd_dart_exit_track(kept, free_named, SD_DART_EXIT_STAGE_RESOURCE)); + CHECK(sd_dart_exit_tracked_count() == 3); + + sd_dart_exit_free(context); + sd_dart_exit_free(resource); + sd_dart_exit_free(resource); + sd_dart_exit_free(nullptr); + sd_dart_exit_free(name("never-tracked")); + CHECK(recorded() == std::vector({"context", "resource"})); + + CHECK(sd_dart_exit_untrack(kept)); + CHECK(!sd_dart_exit_untrack(kept)); + sd_dart_exit_free(kept); + CHECK(sd_dart_exit_tracked_count() == 0); + + sd_dart_exit_teardown(); + CHECK(recorded() == std::vector({"context", "resource"})); + return 0; +} + +int test_order() { + track_out_of_stage_order(); + CHECK(sd_dart_exit_tracked_count() == static_cast(kStageOrder.size())); + + sd_dart_exit_teardown(); + CHECK(recorded() == kStageOrder); + CHECK(sd_dart_exit_tracked_count() == 0); + + // The teardown thread is not blocked, and nothing is freed a second time. + sd_dart_exit_free(name("context-early")); + CHECK(!sd_dart_exit_untrack(name("context-early"))); + CHECK(!sd_dart_exit_track(name("late"), free_named, SD_DART_EXIT_STAGE_CONTEXT)); + sd_ctx_params_t params{}; + sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + CHECK(context != nullptr); + CHECK(sd_dart_exit_tracked_count() == 0); + sd_dart_exit_free(context); + CHECK(contexts_freed.load() == 0); + sd_dart_exit_teardown(); + CHECK(recorded() == kStageOrder); + return 0; +} + +int test_repeat() { + char* object = name("object"); + sd_ctx_params_t params{}; + for (int i = 0; i < 2000; ++i) { + CHECK(sd_dart_exit_track(object, free_named, SD_DART_EXIT_STAGE_CONTEXT_USER)); + sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + CHECK(sd_dart_exit_tracked_count() == 2); + sd_dart_exit_free(context); + sd_dart_exit_free(object); + CHECK(sd_dart_exit_tracked_count() == 0); + } + CHECK(recorded().size() == 2000); + CHECK(contexts_freed.load() == 2000); + + // A tracked address that is tracked again holds a new object. + CHECK(sd_dart_exit_track(object, free_named, SD_DART_EXIT_STAGE_CONTEXT)); + CHECK(sd_dart_exit_track(object, free_named, SD_DART_EXIT_STAGE_RESOURCE)); + CHECK(sd_dart_exit_tracked_count() == 1); + sd_dart_exit_teardown(); + CHECK(recorded().size() == 2001); + return 0; +} + +int test_in_flight() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(30000); + + static std::atomic in_call{false}; + static std::atomic returned{false}; + std::thread([] { + sd_dart_exit_call_begin(); + in_call.store(true); + sleep_ms(400); + record("call-finished"); + sd_dart_exit_call_end(); + returned.store(true); + }).detach(); + while (!in_call.load()) { + sleep_ms(1); + } + + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(started) < 20000); + + std::vector expected = {"call-finished"}; + expected.insert(expected.end(), kStageOrder.begin(), kStageOrder.end()); + CHECK(recorded() == expected); + + // The thread that was in the call never gets back to its caller. + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +// Teardown waits for the outermost call of a thread, which stays free to make +// nested calls and to use the registry until that call ends. +int test_nested() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(30000); + start_heartbeat(); + static char inner[] = "inner"; + + static std::atomic in_call{false}; + static std::atomic returned{false}; + std::thread([] { + sd_dart_exit_call_begin(); + in_call.store(true); + CHECK(wait_for_teardown()); + sd_dart_exit_call_begin(); + CHECK(sd_dart_exit_track(inner, free_named, SD_DART_EXIT_STAGE_CONTEXT_USER)); + sd_ctx_params_t params{}; + sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + CHECK(sd_dart_exit_tracked_count() == static_cast(kStageOrder.size()) + 2); + sd_img_gen_params_t request{}; + CHECK(sd_dart_generate_image(context, &request, nullptr, nullptr)); + sd_dart_exit_free(context); + sd_dart_exit_free(kUserLate); + sd_dart_exit_call_end(); + record("call-finished"); + sd_dart_exit_call_end(); + returned.store(true); + }).detach(); + while (!in_call.load()) { + sleep_ms(1); + } + + teardown_requested.store(true); + sd_dart_exit_teardown(); + std::vector expected = {"user-late", "call-finished", "inner"}; + expected.insert(expected.end(), kStageOrder.begin() + 1, kStageOrder.end()); + CHECK(recorded() == expected); + CHECK(contexts_freed.load() == 1); + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +// A thread that left a call in flight gets time to finish what follows it. +int test_settle() { + track_out_of_stage_order(); + + static std::atomic left_call{false}; + std::thread([] { + sd_dart_exit_call_begin(); + sd_dart_exit_call_end(); + left_call.store(true); + sleep_ms(50); + record("tail-finished"); + }).detach(); + while (!left_call.load()) { + sleep_ms(1); + } + + sd_dart_exit_teardown(); + std::vector expected = {"tail-finished"}; + expected.insert(expected.end(), kStageOrder.begin(), kStageOrder.end()); + CHECK(recorded() == expected); + return 0; +} + +// A generation that never ends: teardown gives up and frees nothing, also not +// the objects no call is using. +int test_timeout() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(1000); + sd_ctx_params_t params{}; + static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + static sd_ctx_t* idle = sd_dart_new_sd_ctx(¶ms); + const int32_t tracked = static_cast(kStageOrder.size()) + 2; + CHECK(sd_dart_exit_tracked_count() == tracked); + + run_generation.store([](sd_ctx_t*) { + for (;;) { + sleep_ms(1000); + } + return true; + }); + std::thread([] { + sd_img_gen_params_t request{}; + sd_dart_generate_image(context, &request, nullptr, nullptr); + }).detach(); + while (!context->generating.load()) { + sleep_ms(1); + } + + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + const int64_t waited = elapsed_ms(started); + CHECK(waited >= 900); + CHECK(waited < 20000); + CHECK(recorded().empty()); + CHECK(contexts_freed.load() == 0); + CHECK(sd_dart_exit_tracked_count() == tracked); + CHECK(context->cancel_requests.load() > 0); + CHECK(idle->cancel_requests.load() > 0); + + // Teardown runs once: a second run neither waits nor frees. + const auto restarted = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(restarted) < 500); + CHECK(recorded().empty()); + return 0; +} + +// A call that never ends does not hold up an exit that has nothing to free: +// what a caller that tracks nothing gets. +int test_idle_wait() { + std::thread([] { sd_dart_exit_call_begin(); }).join(); + sd_ctx_params_t params{}; + sd_ctx_t* untracked = new_sd_ctx(¶ms); + sd_dart_exit_set_wait_ms(30000); + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(started) < 5000); + CHECK(contexts_freed.load() == 0); + CHECK(untracked->cancel_requests.load() == 0); + return 0; +} + +// Teardown that runs inside a call in flight, as when a callback of that call +// exits the process, does not wait for the call it is in. +int test_own_call() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_call_begin(); + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(started) < 5000); + CHECK(recorded() == kStageOrder); + sd_dart_exit_call_end(); + return 0; +} + +int test_blocked() { + char* object = name("object"); + CHECK(sd_dart_exit_track(object, free_named, SD_DART_EXIT_STAGE_CONTEXT)); + sd_ctx_params_t params{}; + static sd_ctx_t* raw = new_sd_ctx(¶ms); + sd_dart_exit_teardown(); + CHECK(recorded() == std::vector({"object"})); + + static std::atomic started{0}; + static std::atomic returned{0}; + int launched = 0; + const auto late = [&launched](void (*call)(void*), void* argument) { + ++launched; + std::thread([call, argument] { + started.fetch_add(1); + call(argument); + returned.fetch_add(1); + }).detach(); + }; + late([](void* argument) { sd_dart_exit_free(argument); }, object); + late([](void* argument) { sd_dart_exit_untrack(argument); }, object); + late([](void* argument) { sd_dart_exit_track(argument, free_named, SD_DART_EXIT_STAGE_CONTEXT); }, object); + late([](void*) { sd_dart_exit_call_begin(); }, nullptr); + late([](void*) { + sd_ctx_params_t late_params{}; + sd_dart_new_sd_ctx(&late_params); + }, + nullptr); + late([](void*) { + sd_img_gen_params_t request{}; + sd_dart_generate_image(raw, &request, nullptr, nullptr); + }, + nullptr); + + while (started.load() < launched) { + sleep_ms(1); + } + sleep_ms(300); + CHECK(returned.load() == 0); + CHECK(recorded() == std::vector({"object"})); + CHECK(!raw->generating.load()); + + // Cancelling is the one call that still returns, without touching anything. + std::thread([] { sd_dart_cancel_generation(raw, SD_CANCEL_ALL); }).join(); + CHECK(raw->cancel_requests.load() == 0); + return 0; +} + +void expect_freed_at_exit() { + const std::vector expected = kTeardownRunsAtExit ? kStageOrder : std::vector(); + if (recorded() != expected) { + std::fprintf(stderr, "exit teardown freed %zu objects, expected %zu\n", recorded().size(), expected.size()); + std::_Exit(1); + } +} + +int test_exit() { + // Registered first, so it runs after the teardown registered by tracking. + CHECK(atexit(expect_freed_at_exit) == 0); + track_out_of_stage_order(); + return 0; +} + +struct LateStatic { + std::string text = std::string(256, 's'); + ~LateStatic() { alive.store(false); } + static std::atomic alive; +}; +std::atomic LateStatic::alive{true}; + +const LateStatic& late_static() { + static LateStatic value; + return value; +} + +std::atomic late_static_read{false}; + +void free_and_read_late_static(void*) { + if (!LateStatic::alive.load() || late_static().text != std::string(256, 's')) { + std::fprintf(stderr, "a static was destroyed before teardown\n"); + std::_Exit(1); + } + late_static_read.store(true); +} + +void expect_late_static_read() { + if (late_static_read.load() != kTeardownRunsAtExit) { + std::fprintf(stderr, "exit teardown did not run before the statics were destroyed\n"); + std::_Exit(1); + } +} + +// A static that is created after the first object was tracked is also +// destroyed only after teardown. +int test_late_static() { + static char reader[] = "reader"; + CHECK(atexit(expect_late_static_read) == 0); + CHECK(sd_dart_exit_track(reader, free_and_read_late_static, SD_DART_EXIT_STAGE_RESOURCE)); + CHECK(late_static().text.size() == 256); + return 0; +} + +// The context is tracked by the time the creating call returns, freed once +// whoever frees it, and a failed load tracks nothing. +int test_context() { + sd_ctx_params_t params{}; + fail_load.store(true); + CHECK(sd_dart_new_sd_ctx(¶ms) == nullptr); + CHECK(sd_dart_exit_tracked_count() == 0); + fail_load.store(false); + + sd_ctx_t* first = sd_dart_new_sd_ctx(¶ms); + CHECK(last_load_params.load() == ¶ms); + sd_ctx_t* second = sd_dart_new_sd_ctx(¶ms); + sd_ctx_t* raw = new_sd_ctx(¶ms); + CHECK(first != nullptr && second != nullptr); + CHECK(sd_dart_exit_tracked_count() == 2); + + sd_img_gen_params_t request{}; + sd_image_t* images = nullptr; + int count = 0; + CHECK(sd_dart_generate_image(first, &request, &images, &count)); + CHECK(last_request.load() == &request && last_images_out.load() == &images); + CHECK(count == 7); + CHECK(!first->generating.load()); + run_generation.store([](sd_ctx_t*) { return false; }); + CHECK(!sd_dart_generate_image(first, &request, &images, &count)); + + sd_dart_exit_free(first); + CHECK(contexts_freed.load() == 1); + CHECK(last_freed.load() == first); + CHECK(!kSeesFreedBlocks || is_freed(first)); + sd_dart_exit_free(first); + sd_dart_exit_free(raw); + CHECK(contexts_freed.load() == 1); + CHECK(sd_dart_exit_tracked_count() == 1); + + sd_dart_exit_teardown(); + CHECK(contexts_freed.load() == 2); + CHECK(last_freed.load() == second); + CHECK(!kSeesFreedBlocks || is_freed(second)); + CHECK(!kSeesFreedBlocks || !is_freed(raw)); + CHECK(sd_dart_exit_tracked_count() == 0); + return 0; +} + +// Cancelling reaches a tracked context and nothing else. +int test_cancel() { + sd_ctx_params_t params{}; + sd_ctx_t* tracked = sd_dart_new_sd_ctx(¶ms); + sd_ctx_t* raw = new_sd_ctx(¶ms); + static sd_ctx_t generic; + CHECK(sd_dart_exit_track(&generic, [](void*) {}, SD_DART_EXIT_STAGE_CONTEXT)); + + sd_dart_cancel_generation(tracked, SD_CANCEL_NEW_LATENTS); + CHECK(tracked->cancel_requests.load() == 1); + CHECK(tracked->cancel.load() == SD_CANCEL_NEW_LATENTS); + sd_dart_cancel_generation(raw, SD_CANCEL_ALL); + sd_dart_cancel_generation(&generic, SD_CANCEL_ALL); + sd_dart_cancel_generation(nullptr, SD_CANCEL_ALL); + CHECK(raw->cancel_requests.load() == 0); + CHECK(generic.cancel_requests.load() == 0); + + // Not after the context is freed either: the stand-in would write to a + // freed block, which a sanitizer reports. + sd_dart_exit_free(tracked); + sd_dart_cancel_generation(tracked, SD_CANCEL_ALL); + CHECK(contexts_freed.load() == 1); + return 0; +} + +std::vector tracked_at_free; + +void free_and_count_tracked(void*) { + tracked_at_free.push_back(sd_dart_exit_tracked_count()); +} + +// A context is freed in the CONTEXT stage: what is tracked in the stage before +// still sees it, what is tracked in the stage after does not. +int test_context_stage() { + static char user[] = "user"; + static char resource[] = "resource"; + // Within a stage the latest tracked object goes first, so each of the two + // would be freed after a context that was tracked in its stage. + CHECK(sd_dart_exit_track(user, free_and_count_tracked, SD_DART_EXIT_STAGE_CONTEXT_USER)); + sd_ctx_params_t params{}; + CHECK(sd_dart_new_sd_ctx(¶ms) != nullptr); + CHECK(sd_dart_exit_track(resource, free_and_count_tracked, SD_DART_EXIT_STAGE_RESOURCE)); + sd_dart_exit_teardown(); + CHECK(tracked_at_free == std::vector({2, 0})); + CHECK(contexts_freed.load() == 1); + return 0; +} + +std::atomic free_started{false}; + +// Freeing a context is a call in flight: teardown waits for it, frees nothing +// else meanwhile, and the freeing thread never returns. +int test_free_in_flight() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(30000); + start_heartbeat(); + sd_ctx_params_t params{}; + static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + hold_free.store([] { + free_started.store(true); + CHECK(wait_for_teardown()); + sleep_ms(50); + CHECK(recorded().empty()); + record("free-finished"); + }); + static std::atomic returned{false}; + std::thread([] { + sd_dart_exit_free(context); + returned.store(true); + }).detach(); + while (!free_started.load()) { + sleep_ms(1); + } + + teardown_requested.store(true); + sd_dart_exit_teardown(); + std::vector expected = {"free-finished"}; + expected.insert(expected.end(), kStageOrder.begin(), kStageOrder.end()); + CHECK(recorded() == expected); + CHECK(contexts_freed.load() == 1); + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +std::atomic load_started{false}; + +// A load in flight when teardown begins, with nothing tracked yet: teardown +// waits for it, the context is tracked although teardown has begun, and it is +// freed. The loading thread never returns. +int test_load_in_flight() { + sd_dart_exit_set_wait_ms(30000); + start_heartbeat(); + hold_load.store([] { + load_started.store(true); + CHECK(wait_for_teardown()); + record("load-finished"); + return true; + }); + static std::atomic returned{false}; + std::thread([] { + sd_ctx_params_t params{}; + sd_dart_new_sd_ctx(¶ms); + returned.store(true); + }).detach(); + while (!load_started.load()) { + sleep_ms(1); + } + + teardown_requested.store(true); + sd_dart_exit_teardown(); + CHECK(recorded() == std::vector({"load-finished"})); + CHECK(contexts_freed.load() == 1); + CHECK(sd_dart_exit_tracked_count() == 0); + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +// A generation in flight when teardown begins: teardown cancels it, waits for +// it to end and only then frees its context. The stand-in clears the first +// cancellation the way generate_image() clears one that precedes its start. +int test_generate_in_flight() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(30000); + start_heartbeat(); + sd_ctx_params_t params{}; + static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + CHECK(sd_dart_new_sd_ctx(¶ms) != nullptr); + + run_generation.store([](sd_ctx_t* generating) { + CHECK(wait_for_teardown()); + while (generating->cancel.load() != SD_CANCEL_ALL) { + sleep_ms(1); + } + generating->cancel.store(SD_CANCEL_RESET); + while (generating->cancel.load() != SD_CANCEL_ALL) { + sleep_ms(1); + } + sleep_ms(50); + CHECK(recorded().empty()); + CHECK(contexts_freed.load() == 0); + CHECK(!kSeesFreedBlocks || !is_freed(generating)); + record("generation-finished"); + return false; + }); + static std::atomic returned{false}; + std::thread([] { + sd_img_gen_params_t request{}; + sd_dart_generate_image(context, &request, nullptr, nullptr); + returned.store(true); + }).detach(); + while (!context->generating.load()) { + sleep_ms(1); + } + + teardown_requested.store(true); + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(started) < 20000); + std::vector expected = {"generation-finished"}; + expected.insert(expected.end(), kStageOrder.begin(), kStageOrder.end()); + CHECK(recorded() == expected); + CHECK(contexts_freed.load() == 2); + CHECK(sd_dart_exit_tracked_count() == 0); + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +} // namespace + +sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { + last_load_params.store(sd_ctx_params); + if (bool (*hold)() = hold_load.load()) { + hold(); + } + return fail_load.load() ? nullptr : new sd_ctx_t(); +} + +void free_sd_ctx(sd_ctx_t* sd_ctx) { + CHECK(!sd_ctx->generating.load()); + if (void (*hold)() = hold_free.load()) { + hold(); + } + contexts_freed.fetch_add(1); + last_freed.store(sd_ctx); + delete sd_ctx; +} + +bool generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out) { + last_request.store(sd_img_gen_params); + last_images_out.store(images_out); + sd_ctx->cancel.store(SD_CANCEL_RESET); + sd_ctx->generating.store(true); + bool ok = true; + if (bool (*run)(sd_ctx_t*) = run_generation.load()) { + ok = run(sd_ctx); + } + if (num_images_out != nullptr) { + *num_images_out = 7; + } + sd_ctx->generating.store(false); + return ok; +} + +void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { + sd_ctx->cancel.store(mode); + sd_ctx->cancel_requests.fetch_add(1); +} + +int main(int argc, char** argv) { + const std::string scenario = argc > 1 ? argv[1] : ""; + const struct { + const char* name; + int (*run)(); + } scenarios[] = { + {"dispose", test_dispose}, + {"order", test_order}, + {"repeat", test_repeat}, + {"in-flight", test_in_flight}, + {"nested", test_nested}, + {"settle", test_settle}, + {"timeout", test_timeout}, + {"idle-wait", test_idle_wait}, + {"own-call", test_own_call}, + {"blocked", test_blocked}, + {"exit", test_exit}, + {"late-static", test_late_static}, + {"context", test_context}, + {"context-stage", test_context_stage}, + {"free-in-flight", test_free_in_flight}, + {"cancel", test_cancel}, + {"load-in-flight", test_load_in_flight}, + {"generate-in-flight", test_generate_in_flight}, + }; + for (const auto& candidate : scenarios) { + if (scenario == candidate.name) { + return candidate.run(); + } + } + if (scenario == "list") { + for (const auto& candidate : scenarios) { + std::printf("%s\n", candidate.name); + } + return 0; + } + std::fprintf(stderr, "usage: %s |list\n", argv[0]); + return 2; +} diff --git a/tests/test_exit_teardown.py b/tests/test_exit_teardown.py new file mode 100644 index 0000000..4163124 --- /dev/null +++ b/tests/test_exit_teardown.py @@ -0,0 +1,92 @@ +"""Builds and runs `tests/native/exit_teardown_test.cpp`. + +The test links `src/sd_dart_exit.cpp` against stand-ins for the upstream +functions it wraps, so it needs the submodule's header but no built runtime. +Each scenario runs in its own process, as teardown runs once per process. +""" + +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +SOURCES = [ + REPO_ROOT / "src" / "sd_dart_exit.cpp", + REPO_ROOT / "tests" / "native" / "exit_teardown_test.cpp", +] +INCLUDE_DIRS = [ + REPO_ROOT / "src", + REPO_ROOT / "third_party" / "stable-diffusion.cpp" / "include", +] +COMPILER = os.environ.get("CXX") or shutil.which("c++") +SANITIZERS = ("address", "thread") +SANITIZER_ENV = { + # The scenarios leave contexts and blocked threads behind on purpose. + "ASAN_OPTIONS": "detect_leaks=0", + "TSAN_OPTIONS": "halt_on_error=1", +} + + +def compile_command(output: Path, sources: list[Path], *flags: str) -> list[str]: + includes = [f"-I{directory}" for directory in INCLUDE_DIRS] + return [COMPILER, "-std=c++17", "-O1", "-g", "-pthread", *flags, *includes, + *map(str, sources), "-o", str(output)] + + +@unittest.skipIf(sys.platform == "win32", "uses a GCC or Clang command line") +class ExitTeardownTest(unittest.TestCase): + def require(self, condition: bool, reason: str) -> None: + """Skips locally, but fails in CI, where a skip would hide the test.""" + if condition: + return + if os.environ.get("CI"): + self.fail(reason) + self.skipTest(reason) + + def setUp(self) -> None: + self.require(COMPILER is not None, "no C++ compiler; set CXX") + self.require((INCLUDE_DIRS[1] / "stable-diffusion.h").is_file(), + "third_party/stable-diffusion.cpp is not checked out") + + def build_and_run(self, *flags: str) -> None: + with tempfile.TemporaryDirectory() as directory: + binary = Path(directory) / "exit_teardown_test" + subprocess.run(compile_command(binary, SOURCES, *flags), check=True) + scenarios = subprocess.run([str(binary), "list"], check=True, capture_output=True, + text=True).stdout.split() + self.assertIn("generate-in-flight", scenarios) + for scenario in scenarios: + with self.subTest(scenario=scenario): + result = subprocess.run( + [str(binary), scenario], capture_output=True, text=True, timeout=600, + env={**os.environ, **SANITIZER_ENV}) + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + + def supports(self, flag: str) -> bool: + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "empty.cpp" + source.write_text("int main() { return 0; }\n") + binary = Path(directory) / "empty" + built = subprocess.run(compile_command(binary, [source], flag), + capture_output=True) + return built.returncode == 0 and subprocess.run( + [str(binary)], capture_output=True).returncode == 0 + + def test_tracks_waits_and_frees_in_order(self) -> None: + self.build_and_run() + + def test_scenarios_are_clean_under_sanitizers(self) -> None: + for sanitizer in SANITIZERS: + flag = f"-fsanitize={sanitizer}" + with self.subTest(sanitizer=sanitizer): + self.require(self.supports(flag), + f"{COMPILER} cannot build or run {flag} binaries") + self.build_and_run(flag) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_tools.py b/tests/test_tools.py index ac8dbf9..55f60dc 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1,3 +1,5 @@ +import shutil +import subprocess import sys import tempfile import unittest @@ -9,6 +11,7 @@ from build import HEADERS, TARGETS, WRAPPER_HEADER # noqa: E402 from package_release import TAG_PATTERN # noqa: E402 from sd_api import api_symbols # noqa: E402 +from validate_artifacts import APPLE_STATIC_DESTRUCTOR_IMPORT, imports_symbol # noqa: E402 HEADER = """ #if __GNUC__ >= 4 @@ -47,13 +50,64 @@ def test_merges_the_symbols_of_several_headers(self): ["generate_image", "new_sd_ctx", "sample_method_to_str", "sd_dart_added"], ) - def test_wrapper_header_adds_only_the_progress_exports(self): + def test_wrapper_header_adds_only_the_progress_and_exit_teardown_exports(self): self.assertEqual( api_symbols(WRAPPER_HEADER), - ["sd_dart_progress_enable", "sd_dart_progress_read"], + [ + "sd_dart_cancel_generation", + "sd_dart_exit_call_begin", + "sd_dart_exit_call_end", + "sd_dart_exit_free", + "sd_dart_exit_set_wait_ms", + "sd_dart_exit_teardown", + "sd_dart_exit_track", + "sd_dart_exit_tracked_count", + "sd_dart_exit_untrack", + "sd_dart_generate_image", + "sd_dart_new_sd_ctx", + "sd_dart_progress_enable", + "sd_dart_progress_read", + ], ) +STATIC_WITH_DESTRUCTOR = """ +struct Static { ~Static(); }; +Static::~Static() {} +Static instance; +""" +OWN_REGISTRATION = """ +extern "C" __attribute__((visibility("hidden"))) int __cxa_atexit(void (*)(void*), void*, void*) { + return 0; +} +""" + + +@unittest.skipUnless(sys.platform == "darwin", "reads Mach-O imports with nm") +class StaticDestructorImportTest(unittest.TestCase): + def build(self, directory: str, name: str, source: str, *archs: str) -> Path: + path = Path(directory) / f"{name}.cpp" + path.write_text(source) + library = Path(directory) / f"lib{name}.dylib" + flags = [flag for arch in archs for flag in ("-arch", arch)] + subprocess.run([shutil.which("c++"), "-std=c++17", "-shared", *flags, str(path), + "-o", str(library)], check=True) + return library + + def test_tells_the_system_registration_from_the_library_s_own(self): + with tempfile.TemporaryDirectory() as tmp: + system = self.build(tmp, "system", STATIC_WITH_DESTRUCTOR) + own = self.build(tmp, "own", STATIC_WITH_DESTRUCTOR + OWN_REGISTRATION) + self.assertTrue(imports_symbol(system, APPLE_STATIC_DESTRUCTOR_IMPORT)) + self.assertFalse(imports_symbol(own, APPLE_STATIC_DESTRUCTOR_IMPORT)) + + def test_reads_the_named_architecture_of_a_universal_file(self): + with tempfile.TemporaryDirectory() as tmp: + system = self.build(tmp, "system", STATIC_WITH_DESTRUCTOR, "arm64", "x86_64") + for arch in ("arm64", "x86_64"): + self.assertTrue(imports_symbol(system, APPLE_STATIC_DESTRUCTOR_IMPORT, arch)) + + class ReleaseTagTest(unittest.TestCase): def test_accepts_semver_and_positive_rebuild_counters(self): for tag in ("v0.1.0", "v1.20.3", "v0.1.0-2"): diff --git a/tools/apple_xcframework.py b/tools/apple_xcframework.py index af88cd1..95cd928 100644 --- a/tools/apple_xcframework.py +++ b/tools/apple_xcframework.py @@ -38,7 +38,11 @@ TARGETS, ) from sd_api import api_symbols -from validate_artifacts import APPLE_SYSTEM_DEPENDENCIES +from validate_artifacts import ( + APPLE_STATIC_DESTRUCTOR_IMPORT, + APPLE_SYSTEM_DEPENDENCIES, + imports_symbol, +) FRAMEWORK = "stable_diffusion" XCFRAMEWORK = f"{FRAMEWORK}.xcframework" @@ -305,6 +309,9 @@ def validate_slice(slice_: Slice, root: Path, entry: dict, expected_symbols: set for dep in (line.split(" (")[0].strip() for line in lines): if dep != slice_.install_name and not SYSTEM_DEPENDENCY.match(dep): problems.append(f"{where} {arch}: unexpected dependency {dep}") + if imports_symbol(binary, APPLE_STATIC_DESTRUCTOR_IMPORT, arch): + problems.append(f"{where} {arch}: imports {APPLE_STATIC_DESTRUCTOR_IMPORT}; " + "static destructors bypass exit teardown") names = output(["nm", "-gUj", "-arch", arch, str(binary)]).split() exported = {n[1:] if n.startswith("_") else n for n in names} if exported != expected_symbols: diff --git a/tools/validate_artifacts.py b/tools/validate_artifacts.py index 61c1c0a..47f0b12 100644 --- a/tools/validate_artifacts.py +++ b/tools/validate_artifacts.py @@ -3,7 +3,8 @@ Checks that each library exports exactly the `SD_API` symbols of the shipped headers (so no ggml symbol leaks) and, except on Windows, links only against -allowlisted system libraries. +allowlisted system libraries. An Apple library must also register its static +destructors through its own `__cxa_atexit`. """ from __future__ import annotations @@ -29,6 +30,11 @@ APPLE_DEPENDENCIES = re.compile( rf"^(@rpath/libstable-diffusion\.dylib|{APPLE_SYSTEM_DEPENDENCIES})$" ) +# `src/sd_dart_exit.cpp` defines `__cxa_atexit`, so that every static destructor +# of an Apple library runs exit teardown first. A library that imports the +# system function registers its destructors past that definition, and exit +# teardown then runs after the statics it needs are gone. +APPLE_STATIC_DESTRUCTOR_IMPORT = "___cxa_atexit" ELF_DEPENDENCIES = { "android": {"libc.so", "libm.so", "libdl.so", "liblog.so"}, "linux": { @@ -85,6 +91,12 @@ def exported_symbols(target: Target, library: Path) -> set[str]: return {line.split()[-1] for line in lines.splitlines() if line.strip()} +def imports_symbol(library: Path, symbol: str, arch: str | None = None) -> bool: + """Whether a Mach-O file imports `symbol`; a universal file needs `arch`.""" + select = ["-arch", arch] if arch else [] + return symbol in output(["nm", "-uj", *select, str(library)]).split() + + def dependencies(target: Target, library: Path) -> list[str]: if target.os in ("macos", "ios"): lines = output(["otool", "-L", str(library)]).splitlines()[1:] @@ -129,6 +141,11 @@ def validate(target: Target) -> list[str]: f"{len(leaked)} non-API exports, e.g. {', '.join(leaked[:10])}" ) + if target.os in ("macos", "ios") and imports_symbol(library, APPLE_STATIC_DESTRUCTOR_IMPORT): + problems.append( + f"imports {APPLE_STATIC_DESTRUCTOR_IMPORT}: static destructors bypass exit teardown" + ) + if target.os == "android": small = [a for a in load_alignments(target, library) if a < ANDROID_PAGE_SIZE] if small: From 2299f51f7cb41b80654b8977459139db03b00951 Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 6 Oct 2026 08:32:04 -0400 Subject: [PATCH 2/5] Run exit teardown against a built runtime and a real context The stand-in tests show what the registry does, not that a context is really freed, nor that the shipped library runs teardown before its statics go. tests/test_exit_teardown_runtime.py links a test against the built macOS runtime. The test writes its own model: a PixArt transformer with one small block, a one-block T5 encoder and a TAESD decoder, 12 MB, which the runtime loads and samples on Metal and on the CPU like any other. Only the tensors the runtime reads the dimensions from are listed; the rest come from its own validation errors, so there is no tensor list to maintain and nothing to download. The scenarios leave the process with a tracked context: idle, during a load, during a generation that is held until teardown is waiting for it, and with a second context on another device. Objects tracked around the contexts check that each context is still there before its stage and freed after it, which the allocator shows where Metal does not abort. tests/dart/exit_teardown.dart does the same under a Dart VM: an isolate killed with, or during, a load or a generation, and C exit() while one generates. CI gets an exit-teardown job on macOS, for the runtime as shipped and for one built with AddressSanitizer. On Linux the build job checks that a context left alive at exit is harmless on the CPU and reports what Vulkan does. Refs #10 --- .github/workflows/validate.yml | 43 ++ tests/dart/exit_teardown.dart | 154 +++++ tests/native/exit_teardown_dart_probe.cpp | 47 ++ tests/native/exit_teardown_runtime_test.cpp | 542 ++++++++++++++++++ .../exit_teardown_two_libraries_test.cpp | 115 ++++ tests/test_exit_teardown_runtime.py | 244 ++++++++ 6 files changed, 1145 insertions(+) create mode 100644 tests/dart/exit_teardown.dart create mode 100644 tests/native/exit_teardown_dart_probe.cpp create mode 100644 tests/native/exit_teardown_runtime_test.cpp create mode 100644 tests/native/exit_teardown_two_libraries_test.cpp create mode 100644 tests/test_exit_teardown_runtime.py diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml index 6e3480c..92bfa2d 100644 --- a/.github/workflows/validate.yml +++ b/.github/workflows/validate.yml @@ -45,6 +45,41 @@ jobs: SD_REQUIRE_DART: "1" SD_PROGRESS_STRESS_RUNS: "25" + exit-teardown: + name: exit teardown (${{ matrix.runtime }}) + runs-on: macos-latest + strategy: + fail-fast: false + matrix: + include: + # The runtime as shipped, and one in which a use after free aborts. + - { runtime: release, sanitizer: "" } + - { runtime: address sanitizer, sanitizer: address } + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + - uses: ./.github/actions/setup-toolchain + with: + target: macos-arm64 + - uses: dart-lang/setup-dart@v1 + if: ${{ !matrix.sanitizer }} + with: + sdk: 3.13.5 + - name: Build + if: ${{ !matrix.sanitizer }} + run: python3 tools/build.py build --target macos-arm64 + - name: Free tracked contexts at exit + # With a sanitizer the test builds its own runtime, which the Dart VM + # cannot load. The model is written by the test; nothing is + # downloaded. + run: python3 -m unittest discover -s tests -p test_exit_teardown_runtime.py -v + env: + SD_REQUIRE_RUNTIME: "1" + SD_REQUIRE_DART: ${{ !matrix.sanitizer && '1' || '' }} + SD_EXIT_TEARDOWN_SANITIZER: ${{ matrix.sanitizer }} + SD_EXIT_TEARDOWN_RUNS: "3" + build: name: build (${{ matrix.target }}) runs-on: ${{ matrix.runs_on }} @@ -87,6 +122,14 @@ jobs: env: # ggml skips CPU-type Vulkan devices such as lavapipe unless listed. GGML_VK_VISIBLE_DEVICES: "0" + - name: Exit with a live context + # Exit teardown runs at exit on Apple platforms only; see the + # exit-teardown job. Here a context must be harmless to leave behind. + if: runner.os == 'Linux' && matrix.smoke + run: python3 -m unittest discover -s tests -p test_exit_teardown_runtime.py -v + env: + SD_EXIT_TEARDOWN_TARGET: ${{ matrix.target }} + GGML_VK_VISIBLE_DEVICES: "0" - uses: actions/upload-artifact@v4 if: runner.os == 'macOS' with: diff --git a/tests/dart/exit_teardown.dart b/tests/dart/exit_teardown.dart new file mode 100644 index 0000000..b737833 --- /dev/null +++ b/tests/dart/exit_teardown.dart @@ -0,0 +1,154 @@ +// Exit teardown under a real Dart VM. +// +// Usage: dart exit_teardown.dart +// +// Each scenario leaves the process with a context that no Dart code frees, +// the way a Flutter quit or hot restart does: +// +// idle the main isolate loads a model and returns from main +// leaked a worker isolate loads a model and is killed; then main +// returns: the hot restart case +// killed-in-load a worker isolate is killed while it loads, so no Dart +// code ever sees the context +// killed-in-run a worker isolate is killed while it generates +// exit-in-run main calls C exit() while a worker isolate generates, +// as a Flutter quit does. dart:io's exit() would not do: +// it leaves without running static destructors +// +// With api `tracked` the harness calls sd_dart_new_sd_ctx and +// sd_dart_generate_image, and the process must exit with code 0. With api +// `raw` it calls new_sd_ctx and generate_image: the control, which aborts in +// ggml-metal where Metal residency sets are live. +import 'dart:async'; +import 'dart:ffi'; +import 'dart:io'; +import 'dart:isolate'; + +typedef LoadNative = Pointer Function(Pointer); +typedef GenerateNative = Bool Function( + Pointer, + Pointer, + Pointer>, + Pointer, +); +typedef GenerateDart = bool Function( + Pointer, + Pointer, + Pointer>, + Pointer, +); + +final class Runtime { + final Pointer Function(Pointer) load; + final GenerateDart generate; + final Pointer contextParams; + final Pointer generationParams; + final Pointer> imagesOut; + final Pointer countOut; + + Runtime._( + this.load, + this.generate, + this.contextParams, + this.generationParams, + this.imagesOut, + this.countOut, + ); + + factory Runtime( + String runtimePath, + String probePath, + String model, + bool tracked, + ) { + final runtime = DynamicLibrary.open(runtimePath); + final probe = DynamicLibrary.open(probePath); + final malloc = DynamicLibrary.process() + .lookupFunction< + Pointer Function(IntPtr), + Pointer Function(int) + >('malloc'); + final path = malloc(model.length + 1).cast(); + for (var i = 0; i < model.length; i++) { + path[i] = model.codeUnitAt(i); + } + path[model.length] = 0; + probe.lookupFunction('probe_quiet')(); + return Runtime._( + runtime.lookupFunction Function(Pointer)>( + tracked ? 'sd_dart_new_sd_ctx' : 'new_sd_ctx', + ), + runtime.lookupFunction( + tracked ? 'sd_dart_generate_image' : 'generate_image', + ), + probe.lookupFunction Function(Pointer)>( + 'probe_context_params', + )(path.cast()), + probe.lookupFunction< + Pointer Function(Int), + Pointer Function(int) + >('probe_generation_params')(40), + malloc(sizeOf>()).cast(), + malloc(sizeOf()).cast(), + ); + } + + Pointer loadContext() { + final context = load(contextParams); + if (context == nullptr) { + stderr.writeln('exit teardown: the model did not load'); + exit(3); + } + return context; + } + + void generateOn(Pointer context) => + generate(context, generationParams, imagesOut, countOut); +} + +void worker((SendPort, List) arguments) { + final (events, [runtimePath, probePath, model, scenario, api]) = arguments; + final runtime = Runtime(runtimePath, probePath, model, api == 'tracked'); + events.send('loading'); + final context = runtime.loadContext(); + events.send('loaded'); + if (scenario == 'leaked' || scenario == 'killed-in-load') { + // Keeps the isolate alive until it is killed. + ReceivePort(); + return; + } + for (;;) { + events.send('generating'); + runtime.generateOn(context); + } +} + +Future main(List arguments) async { + final [runtimePath, probePath, model, scenario, api] = arguments; + if (scenario == 'idle') { + final runtime = Runtime(runtimePath, probePath, model, api == 'tracked'); + runtime.generateOn(runtime.loadContext()); + return; + } + + final events = ReceivePort(); + final exited = ReceivePort(); + final isolate = await Isolate.spawn(worker, ( + events.sendPort, + arguments, + ), onExit: exited.sendPort); + final waitFor = switch (scenario) { + 'leaked' => 'loaded', + 'killed-in-load' => 'loading', + _ => 'generating', + }; + await events.firstWhere((event) => event == waitFor); + if (scenario == 'exit-in-run') { + await Future.delayed(const Duration(milliseconds: 20)); + DynamicLibrary.process() + .lookupFunction('exit')(0); + } + isolate.kill(priority: Isolate.immediate); + await exited.first; + exited.close(); +} diff --git a/tests/native/exit_teardown_dart_probe.cpp b/tests/native/exit_teardown_dart_probe.cpp new file mode 100644 index 0000000..b269cf8 --- /dev/null +++ b/tests/native/exit_teardown_dart_probe.cpp @@ -0,0 +1,47 @@ +// Native side of tests/dart/exit_teardown.dart, linked against a built +// runtime: it fills in the parameter structs, so that the harness can call +// the runtime's own functions without mirroring those structs in Dart. + +#include "sd_dart_wrapper.h" + +#include + +#define PROBE_API extern "C" __attribute__((visibility("default"))) + +namespace { + +void quiet_log(enum sd_log_level_t, const char*, void*) {} + +} // namespace + +// Keeps the runtime off stdout and stderr without a Dart callback, and gives +// teardown time for a call in flight on a busy machine. +PROBE_API void probe_quiet(void) { + sd_set_log_callback(quiet_log, nullptr); + sd_dart_progress_enable(); + sd_dart_exit_set_wait_ms(120000); +} + +// Parameters for a model that holds its TAESD decoder itself, as the one +// exit_teardown_runtime_test writes does. Never freed. +PROBE_API sd_ctx_params_t* probe_context_params(const char* model) { + auto* params = static_cast(malloc(sizeof(sd_ctx_params_t))); + sd_ctx_params_init(params); + params->model_path = model; + params->taesd_path = model; + params->eager_load = true; + return params; +} + +PROBE_API sd_img_gen_params_t* probe_generation_params(int steps) { + auto* params = static_cast(malloc(sizeof(sd_img_gen_params_t))); + sd_img_gen_params_init(params); + params->prompt = "a lighthouse"; + params->negative_prompt = ""; + params->width = 64; + params->height = 64; + params->seed = 7; + params->sample_params.sample_steps = steps; + params->sample_params.guidance.txt_cfg = 1.0f; + return params; +} diff --git a/tests/native/exit_teardown_runtime_test.cpp b/tests/native/exit_teardown_runtime_test.cpp new file mode 100644 index 0000000..a6144ac --- /dev/null +++ b/tests/native/exit_teardown_runtime_test.cpp @@ -0,0 +1,542 @@ +// Exit teardown against a built runtime and a real sd_ctx_t. +// +// Usage: exit_teardown_runtime_test make-model +// exit_teardown_runtime_test |default [] +// +// `make-model` writes a PixArt model with one small transformer block, a +// one-block T5 text encoder and a TAESD decoder: 12 MB that the runtime loads +// and samples like any other model, on Metal as well as on the CPU. Nothing +// is downloaded. A model that carries no TAESD decoder, as real ones do, is +// loaded as it is; give the image size it needs. +// +// The other scenarios each run in their own process, as teardown runs once +// per process, and most of them leave by returning from main: on Apple +// platforms teardown then runs during exit(). Elsewhere it does not, and only +// `idle-untracked` applies: a context that is alive at exit must not abort. + +#include "sd_dart_wrapper.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__has_feature) +#if __has_feature(address_sanitizer) +#define SD_TEST_ADDRESS_SANITIZER 1 +#endif +#endif + +#if defined(SD_TEST_ADDRESS_SANITIZER) +#include +#elif defined(__APPLE__) +#include +#endif + +#define CHECK(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: CHECK(%s) failed\n", __FILE__, __LINE__, #condition); \ + std::_Exit(1); \ + } \ + } while (0) + +namespace { + +void sleep_ms(int ms) { + std::this_thread::sleep_for(std::chrono::milliseconds(ms)); +} + +int64_t elapsed_ms(std::chrono::steady_clock::time_point since) { + return std::chrono::duration_cast(std::chrono::steady_clock::now() - since).count(); +} + +// Whether the heap block that began at `object` has been freed. +// AddressSanitizer poisons a freed block, and the macOS allocator reports no +// size for one until it hands the block out again. No scenario that asks +// runs anywhere else. +bool is_freed(const void* object) { +#if defined(SD_TEST_ADDRESS_SANITIZER) + return __asan_address_is_poisoned(object) != 0; +#elif defined(__APPLE__) + return malloc_size(object) == 0; +#else + (void)object; + return false; +#endif +} + +// The object that owns the weights, runners and backends of a context: the +// only member of upstream's sd_ctx_t. Unlike the 8-byte sd_ctx_t itself it is +// large enough for every allocator to tell whether it is freed. +const void* engine_of(const sd_ctx_t* context) { + const void* engine = *reinterpret_cast(context); + CHECK(engine != nullptr && !is_freed(engine)); + return engine; +} + +std::mutex log_mutex; +std::string log_text; +std::atomic on_log{nullptr}; + +void collect_log(enum sd_log_level_t level, const char* text, void*) { + if (void (*hook)() = on_log.load()) { + hook(); + } + std::lock_guard lock(log_mutex); + if (level >= SD_LOG_WARN) { + log_text += text; + } +} + +std::atomic on_progress{nullptr}; + +void collect_progress(int step, int, float, void*) { + if (void (*hook)(int) = on_progress.load()) { + hook(step); + } +} + +int image_size = 64; + +// The model `make-model` writes holds its TAESD decoder itself. +bool holds_taesd(const char* model) { + std::string header(1 << 16, '\0'); + FILE* file = std::fopen(model, "rb"); + if (file == nullptr) { + return false; + } + header.resize(std::fread(&header[0], 1, header.size(), file)); + std::fclose(file); + return header.find("\"tae.decoder.") != std::string::npos; +} + +sd_ctx_params_t context_params(const char* model, const char* backend, bool taesd) { + sd_ctx_params_t params; + sd_ctx_params_init(¶ms); + params.model_path = model; + params.taesd_path = taesd ? model : nullptr; + params.backend = std::strcmp(backend, "default") == 0 ? nullptr : backend; + params.eager_load = true; + return params; +} + +sd_ctx_t* load(const char* model, const char* backend, bool tracked = true) { + const sd_ctx_params_t params = context_params(model, backend, holds_taesd(model)); + sd_ctx_t* context = tracked ? sd_dart_new_sd_ctx(¶ms) : new_sd_ctx(¶ms); + if (context == nullptr) { + std::lock_guard lock(log_mutex); + std::fprintf(stderr, "%s", log_text.c_str()); + } + CHECK(context != nullptr); + return context; +} + +// Returns the number of images, or -1 when the generation failed. +int generate(sd_ctx_t* context, int steps, bool guarded = true) { + sd_img_gen_params_t params; + sd_img_gen_params_init(¶ms); + params.prompt = "a lighthouse"; + params.negative_prompt = ""; + params.width = image_size; + params.height = image_size; + params.seed = 7; + params.sample_params.sample_steps = steps; + params.sample_params.guidance.txt_cfg = 1.0f; + sd_image_t* images = nullptr; + int count = 0; + const bool ok = guarded ? sd_dart_generate_image(context, ¶ms, &images, &count) + : generate_image(context, ¶ms, &images, &count); + if (images != nullptr) { + free_sd_images(images, count); + } + return ok ? count : -1; +} + +using Tensors = std::map>; + +bool write_model(const char* path, const Tensors& tensors) { + std::string header = "{"; + uint64_t offset = 0; + for (const auto& [name, shape] : tensors) { + uint64_t elements = 1; + std::string dims; + for (int64_t dim : shape) { + elements *= static_cast(dim); + dims += (dims.empty() ? "" : ",") + std::to_string(dim); + } + header += (offset == 0 ? "\"" : ",\"") + name + "\":{\"dtype\":\"F32\",\"shape\":[" + dims + + "],\"data_offsets\":[" + std::to_string(offset) + "," + std::to_string(offset + 4 * elements) + "]}"; + offset += 4 * elements; + } + header += "}"; + FILE* file = std::fopen(path, "wb"); + if (file == nullptr) { + return false; + } + const uint64_t header_size = header.size(); + std::fwrite(&header_size, sizeof(header_size), 1, file); + std::fwrite(header.data(), 1, header.size(), file); + const std::vector weights(16384, 0.02f); + for (uint64_t written = 0; written < offset;) { + const uint64_t bytes = std::min(offset - written, weights.size() * sizeof(float)); + std::fwrite(weights.data(), 1, bytes, file); + written += bytes; + } + return std::fclose(file) == 0; +} + +// The tensors the runtime reads the model's dimensions from are given here. +// The rest are whatever its own validation asks for: a load that fails names +// each missing tensor and the shape it wants, so the model follows the +// pinned upstream revision without a tensor list to maintain. +int make_model(const char* path) { + const int64_t hidden = 32, caption = 16, vocabulary = 32128; + const std::string dit = "model.diffusion_model."; + const std::string t5 = "text_encoders.t5xxl.transformer."; + Tensors tensors = { + {dit + "t_block.1.weight", {6 * hidden, hidden}}, + {dit + "x_embedder.proj.weight", {hidden, 4, 2, 2}}, + {dit + "y_embedder.y_proj.fc1.weight", {hidden, caption}}, + {dit + "blocks.0.mlp.fc1.weight", {2 * hidden, hidden}}, + {dit + "blocks.0.cross_attn.kv_linear.weight", {2 * hidden, hidden}}, + {dit + "final_layer.linear.weight", {32, hidden}}, + {t5 + "shared.weight", {vocabulary, caption}}, + {t5 + "encoder.block.0.layer.0.SelfAttention.q.weight", {64, caption}}, + {t5 + "encoder.block.0.layer.1.DenseReluDense.wi_0.weight", {2 * caption, caption}}, + }; + const std::regex missing("tensor '([^']+)' not in model metadata"); + const std::regex wrong("tensor '([^']+)' has wrong shape in model metadata: got \\[[^\\]]*\\], " + "expected \\[(\\d+), (\\d+), (\\d+), (\\d+)\\]"); + for (int attempt = 0; attempt < 8; ++attempt) { + CHECK(write_model(path, tensors)); + { + std::lock_guard lock(log_mutex); + log_text.clear(); + } + // With a TAESD decoder the model needs no VAE, which would be 320 MB. + const sd_ctx_params_t params = context_params(path, "cpu", true); + if (sd_ctx_t* context = new_sd_ctx(¶ms)) { + free_sd_ctx(context); + CHECK(holds_taesd(path)); + return 0; + } + std::lock_guard lock(log_mutex); + bool changed = false; + for (std::sregex_iterator it(log_text.begin(), log_text.end(), missing), end; it != end; ++it) { + changed |= tensors.insert({(*it)[1], {1}}).second; + } + for (std::sregex_iterator it(log_text.begin(), log_text.end(), wrong), end; it != end; ++it) { + // ggml lists dimensions innermost first; safetensors outermost. + std::vector shape; + for (int dim = 5; dim >= 2; --dim) { + const int64_t size = std::stoll((*it)[dim]); + if (size != 1 || !shape.empty() || dim == 2) { + shape.push_back(size); + } + } + changed |= tensors[(*it)[1]] != shape; + tensors[(*it)[1]] = shape; + } + if (!changed) { + std::fprintf(stderr, "the runtime rejected the model without naming a tensor:\n%s", log_text.c_str()); + return 1; + } + } + std::fprintf(stderr, "the model did not settle\n"); + return 1; +} + +// Teardown blocks a thread that reaches the registry outside a call in +// flight. This one keeps reaching it, so when its count stops after the +// scenario asked for teardown, teardown has begun. +std::atomic heartbeat{0}; +std::atomic teardown_requested{false}; + +void start_heartbeat() { + std::thread([] { + for (;;) { + sd_dart_exit_untrack(nullptr); + heartbeat.fetch_add(1); + sleep_ms(1); + } + }).detach(); + while (heartbeat.load() == 0) { + sleep_ms(1); + } +} + +bool wait_for_teardown() { + while (!teardown_requested.load()) { + sleep_ms(1); + } + const auto started = std::chrono::steady_clock::now(); + int64_t seen = heartbeat.load(); + auto changed = std::chrono::steady_clock::now(); + while (elapsed_ms(started) < 60000) { + sleep_ms(5); + const int64_t now = heartbeat.load(); + if (now != seen) { + seen = now; + changed = std::chrono::steady_clock::now(); + } else if (elapsed_ms(changed) > 400) { + return true; + } + } + return false; +} + +// What the two objects tracked around the contexts saw during teardown. +struct Witness { + std::atomic engine{nullptr}; + std::atomic expected_contexts{1}; + std::atomic before_ran{false}; + std::atomic after_ran{false}; + // What else the scenario expects once teardown is over; null if nothing. + std::atomic failure{nullptr}; +}; +Witness witness; +char witness_before[] = "before"; +char witness_after[] = "after"; + +// Freed before any context: every context is still there. +void free_before_contexts(void*) { + CHECK(sd_dart_exit_tracked_count() == witness.expected_contexts.load() + 1); + const void* engine = witness.engine.load(); + CHECK(engine == nullptr || !is_freed(engine)); + witness.before_ran.store(true); +} + +// Freed after every context, with nothing allocated in between. +void free_after_contexts(void*) { + CHECK(sd_dart_exit_tracked_count() == 0); + const void* engine = witness.engine.load(); + CHECK(engine == nullptr || is_freed(engine)); + witness.after_ran.store(true); +} + +void expect_teardown_ran() { + const char* (*failure)() = witness.failure.load(); + const char* message = failure != nullptr ? failure() : nullptr; + if (!witness.before_ran.load() || !witness.after_ran.load()) { + message = "exit teardown did not free the tracked objects"; + } + if (message != nullptr) { + std::fprintf(stderr, "%s\n", message); + std::_Exit(1); + } +} + +// Call before anything is tracked: the check then runs after the teardown +// that tracking registers. +void expect_teardown_at_exit(int contexts) { + witness.expected_contexts.store(contexts); + CHECK(atexit(expect_teardown_ran) == 0); + CHECK(sd_dart_exit_track(witness_before, free_before_contexts, SD_DART_EXIT_STAGE_CONTEXT_USER)); +} + +// Call once the contexts exist, where the scenario can. Within a stage the +// latest tracked object goes first, so this witness would also be freed +// before a context that was tracked in its stage. +void expect_contexts_freed_first() { + CHECK(sd_dart_exit_track(witness_after, free_after_contexts, SD_DART_EXIT_STAGE_RESOURCE)); +} + +// Exits with a loaded context that nothing frees. +int test_idle(const char* model, const char* backend, bool tracked) { + if (tracked) { + expect_teardown_at_exit(1); + } + sd_ctx_t* context = load(model, backend, tracked); + CHECK(generate(context, 2, tracked) == 1); + CHECK(sd_dart_exit_tracked_count() == (tracked ? 2 : 0)); + if (tracked) { + witness.engine.store(engine_of(context)); + expect_contexts_freed_first(); + } + // Tells the caller that whatever follows happened during exit(). + std::printf("generated on %s\n", backend); + std::fflush(stdout); + return 0; +} + +int test_dispose(const char* model, const char* backend) { + for (int i = 0; i < 3; ++i) { + sd_ctx_t* context = load(model, backend); + const void* engine = engine_of(context); + if (i != 1) { + CHECK(generate(context, 1) == 1); + } + CHECK(sd_ctx_supports_image_generation(context)); + CHECK(sd_get_model_version_name(context) != nullptr); + CHECK(sd_dart_exit_tracked_count() == 1); + sd_dart_exit_free(context); + CHECK(is_freed(engine)); + CHECK(sd_dart_exit_tracked_count() == 0); + sd_dart_exit_free(context); + sd_dart_cancel_generation(context, SD_CANCEL_ALL); + } + return 0; +} + +std::atomic steps_seen{0}; +std::atomic cancel_requested{false}; + +// Cancelling through the wrapper ends a generation on a tracked context. +int test_cancel(const char* model, const char* backend) { + static sd_ctx_t* context = load(model, backend); + on_progress.store([](int) { + steps_seen.fetch_add(1); + if (!cancel_requested.exchange(true)) { + std::thread([] { sd_dart_cancel_generation(context, SD_CANCEL_ALL); }).join(); + } + }); + generate(context, 40); + CHECK(steps_seen.load() > 0 && steps_seen.load() < 10); + sd_dart_exit_free(context); + return 0; +} + +std::atomic held{false}; +std::atomic released{false}; +std::atomic steps_after_release{0}; + +// A generation in flight when the process exits: teardown cancels it, waits +// for it to end and then frees the context it ran on. +int test_generate_wait(const char* model, const char* backend) { + witness.failure.store([]() -> const char* { + if (!released.load()) { + return "teardown did not wait for the generation"; + } + return steps_after_release.load() > 3 ? "teardown did not cancel the generation" : nullptr; + }); + expect_teardown_at_exit(1); + start_heartbeat(); + static sd_ctx_t* context = load(model, backend); + witness.engine.store(engine_of(context)); + expect_contexts_freed_first(); + on_progress.store([](int) { + if (released.load()) { + steps_after_release.fetch_add(1); + return; + } + if (held.exchange(true)) { + return; + } + CHECK(wait_for_teardown()); + const void* engine = witness.engine.load(); + CHECK(!is_freed(engine)); + sleep_ms(100); + CHECK(!is_freed(engine)); + CHECK(!witness.before_ran.load()); + released.store(true); + }); + std::thread([] { + generate(context, 40); + std::fprintf(stderr, "a call in flight returned after teardown\n"); + std::_Exit(1); + }).detach(); + while (!held.load()) { + sleep_ms(1); + } + teardown_requested.store(true); + return 0; +} + +// A load in flight when the process exits: teardown waits for it, and the +// context it creates is tracked and freed although teardown has begun. +int test_load_wait(const char* model, const char* backend) { + witness.failure.store([]() -> const char* { + return released.load() ? nullptr : "teardown did not wait for the load"; + }); + expect_teardown_at_exit(1); + expect_contexts_freed_first(); + start_heartbeat(); + static const char* load_model = model; + static const char* load_backend = backend; + on_log.store([] { + if (held.exchange(true)) { + return; + } + CHECK(wait_for_teardown()); + CHECK(sd_dart_exit_tracked_count() == 2); + released.store(true); + }); + std::thread([] { + load(load_model, load_backend); + std::fprintf(stderr, "a call in flight returned after teardown\n"); + std::_Exit(1); + }).detach(); + while (!held.load()) { + sleep_ms(1); + } + teardown_requested.store(true); + return 0; +} + +// A context on another device, loaded after the first one was tracked, makes +// the runtime create statics that did not exist when teardown was registered. +// They are destroyed after teardown all the same. +int test_late_load(const char* model, const char* backend) { + expect_teardown_at_exit(2); + sd_ctx_t* first = load(model, "cpu"); + CHECK(generate(first, 1) == 1); + sd_ctx_t* second = load(model, backend); + CHECK(generate(second, 1) == 1); + // The first context is freed last, so nothing reuses its block before the + // witness looks at it. + witness.engine.store(engine_of(first)); + expect_contexts_freed_first(); + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + const std::string scenario = argc > 1 ? argv[1] : ""; + sd_set_log_callback(collect_log, nullptr); + sd_set_progress_callback(collect_progress, nullptr); + // The scenarios are about what teardown does once the calls in flight + // have ended. A busy machine must not turn that into the timeout. + sd_dart_exit_set_wait_ms(120000); + if (argc == 3 && scenario == "make-model") { + return make_model(argv[2]); + } + if (argc == 4 || argc == 5) { + const char* model = argv[2]; + const char* backend = argv[3]; + image_size = argc == 5 ? std::atoi(argv[4]) : image_size; + if (scenario == "idle") { + return test_idle(model, backend, true); + } + if (scenario == "idle-untracked") { + return test_idle(model, backend, false); + } + if (scenario == "dispose") { + return test_dispose(model, backend); + } + if (scenario == "cancel") { + return test_cancel(model, backend); + } + if (scenario == "generate-wait") { + return test_generate_wait(model, backend); + } + if (scenario == "load-wait") { + return test_load_wait(model, backend); + } + if (scenario == "late-load") { + return test_late_load(model, backend); + } + } + std::fprintf(stderr, "usage: %s make-model | |default []\n", argv[0]); + return 2; +} diff --git a/tests/native/exit_teardown_two_libraries_test.cpp b/tests/native/exit_teardown_two_libraries_test.cpp new file mode 100644 index 0000000..4587a2c --- /dev/null +++ b/tests/native/exit_teardown_two_libraries_test.cpp @@ -0,0 +1,115 @@ +// Exit teardown next to libllamadart in one process. +// +// Usage: exit_teardown_two_libraries_test +// sd|llama exit|sd|llama |default +// +// libllamadart (leehack/llamadart-native) embeds its own ggml and carries the +// same registry and the same hidden __cxa_atexit for it. The third argument +// says which library is loaded first. Each library then tracks an object, and +// this one also a real context. With `exit` the process returns from main and +// both teardowns have to run by themselves; with `sd` or `llama` that +// library's teardown is called first and must leave the other registry alone. +// Every free is printed, so the caller sees each object freed once. + +#include "sd_dart_wrapper.h" + +#include +#include +#include +#include +#include +#include + +#define CHECK(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: CHECK(%s) failed\n", __FILE__, __LINE__, #condition); \ + std::_Exit(1); \ + } \ + } while (0) + +namespace { + +template +Function symbol(void* library, const char* name) { + void* address = dlsym(library, name); + if (address == nullptr) { + std::fprintf(stderr, "missing %s: %s\n", name, dlerror()); + std::_Exit(1); + } + return reinterpret_cast(address); +} + +void report(const char* line) { + CHECK(write(STDOUT_FILENO, line, std::strlen(line)) > 0); +} + +void free_sd_object(void*) { + report("freed by libstable-diffusion\n"); +} + +void free_llama_object(void*) { + report("freed by libllamadart\n"); +} + +void quiet_log(enum sd_log_level_t, const char*, void*) {} +void quiet_progress(int, int, float, void*) {} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 7) { + std::fprintf(stderr, "usage: %s sd|llama exit|sd|llama \n", + argv[0]); + return 2; + } + const bool sd_first = std::strcmp(argv[3], "sd") == 0; + const std::string mode = argv[4]; + void* first = dlopen(argv[sd_first ? 1 : 2], RTLD_NOW | RTLD_LOCAL); + void* second = dlopen(argv[sd_first ? 2 : 1], RTLD_NOW | RTLD_LOCAL); + CHECK(first != nullptr && second != nullptr); + void* sd = sd_first ? first : second; + void* llama = sd_first ? second : first; + + using Track = bool (*)(void*, void (*)(void*), int32_t); + using Count = int32_t (*)(); + using Teardown = void (*)(); + const auto sd_track = symbol(sd, "sd_dart_exit_track"); + const auto sd_count = symbol(sd, "sd_dart_exit_tracked_count"); + const auto sd_teardown = symbol(sd, "sd_dart_exit_teardown"); + const auto llama_track = symbol(llama, "llama_dart_exit_track"); + const auto llama_count = symbol(llama, "llama_dart_exit_tracked_count"); + const auto llama_teardown = symbol(llama, "llama_dart_exit_teardown"); + + symbol(sd, "sd_set_log_callback")(quiet_log, nullptr); + symbol(sd, "sd_set_progress_callback")(quiet_progress, nullptr); + sd_ctx_params_t params; + symbol(sd, "sd_ctx_params_init")(¶ms); + params.model_path = argv[5]; + params.taesd_path = argv[5]; + params.backend = std::strcmp(argv[6], "default") == 0 ? nullptr : argv[6]; + params.eager_load = true; + + static char sd_object[] = "sd"; + static char llama_object[] = "llama"; + // The library that was loaded first tracks last, so that the two orders + // also differ in which teardown is registered first. + if (sd_first) { + CHECK(llama_track(llama_object, free_llama_object, 0)); + } + CHECK(sd_track(sd_object, free_sd_object, SD_DART_EXIT_STAGE_RESOURCE)); + CHECK(symbol(sd, "sd_dart_new_sd_ctx")(¶ms) != nullptr); + if (!sd_first) { + CHECK(llama_track(llama_object, free_llama_object, 0)); + } + CHECK(sd_count() == 2 && llama_count() == 1); + + if (mode == "sd") { + sd_teardown(); + CHECK(sd_count() == 0 && llama_count() == 1); + } else if (mode == "llama") { + llama_teardown(); + CHECK(sd_count() == 2 && llama_count() == 0); + } + return 0; +} diff --git a/tests/test_exit_teardown_runtime.py b/tests/test_exit_teardown_runtime.py new file mode 100644 index 0000000..f29774e --- /dev/null +++ b/tests/test_exit_teardown_runtime.py @@ -0,0 +1,244 @@ +"""Runs `tests/native/exit_teardown_runtime_test.cpp` against a built runtime. + +Needs the macOS runtime of this host under `bin/` (`tools/build.py build +--target macos-arm64`) and skips without it. The test binary writes its own +12 MB model, so nothing is downloaded. + +On Linux, where teardown does not run at exit, it only checks that a context +left alive at exit does not need it: set `SD_EXIT_TEARDOWN_TARGET` to a built +Linux target. + + SD_REQUIRE_RUNTIME=1 fail instead of skipping + SD_REQUIRE_DART=1 fail when no Dart SDK is on PATH; the + Dart VM harness skips without one + SD_EXIT_TEARDOWN_SANITIZER=address + build the runtime with AddressSanitizer + under build/ and test that one instead + SD_EXIT_TEARDOWN_RUNS=5 runs per scenario and backend + SD_EXIT_TEARDOWN_MODEL= also run the scenarios on this model + SD_EXIT_TEARDOWN_MODEL_SIZE=256 the image size that model needs + SD_EXIT_TEARDOWN_LLAMADART= + also load llamadart-native's library into + the same process; it needs a build with + that repository's own exit teardown +""" + +from __future__ import annotations + +import os +import platform +import shutil +import signal +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "tools")) + +import build # noqa: E402 + +TARGET = build.TARGETS["macos-arm64" if platform.machine() == "arm64" else "macos-x64"] +SOURCE = REPO_ROOT / "tests" / "native" / "exit_teardown_runtime_test.cpp" +DART_PROBE = REPO_ROOT / "tests" / "native" / "exit_teardown_dart_probe.cpp" +DART_HARNESS = REPO_ROOT / "tests" / "dart" / "exit_teardown.dart" +DART_SCENARIOS = ("idle", "leaked", "killed-in-load", "killed-in-run", "exit-in-run") +COMPILER = os.environ.get("CXX") or shutil.which("c++") +DART = shutil.which("dart") +TWO_LIBRARIES = REPO_ROOT / "tests" / "native" / "exit_teardown_two_libraries_test.cpp" +LLAMADART = os.environ.get("SD_EXIT_TEARDOWN_LLAMADART", "") +LINUX_TARGET = os.environ.get("SD_EXIT_TEARDOWN_TARGET", "") +SANITIZER = os.environ.get("SD_EXIT_TEARDOWN_SANITIZER", "") +RUNS = int(os.environ.get("SD_EXIT_TEARDOWN_RUNS", "1")) +REAL_MODEL = os.environ.get("SD_EXIT_TEARDOWN_MODEL", "") +REAL_MODEL_SIZE = os.environ.get("SD_EXIT_TEARDOWN_MODEL_SIZE", "256") +SCENARIOS = ("idle", "dispose", "cancel", "generate-wait", "load-wait", "late-load") +# `default` lets the runtime pick its device, which is Metal where there is one. +BACKENDS = ("default", "cpu") +METAL_ABORT = "[rsets->data count] == 0" +SCENARIO_ENV = {"ASAN_OPTIONS": "detect_leaks=0"} + + +def sanitized_library(sanitizer: str) -> Path: + """Builds the runtime as `tools/build.py` does, plus the sanitizer.""" + work_dir = build.BUILD_ROOT / f"{TARGET.name}-{sanitizer}" + work_dir.mkdir(parents=True, exist_ok=True) + flags = f"-fsanitize={sanitizer} -fno-omit-frame-pointer -g" + args = [f"{arg} {flags}" if arg.startswith("-DCMAKE_SHARED_LINKER_FLAGS=") else arg + for arg in build.configure_args(TARGET, work_dir)] + args += [f"-DCMAKE_C_FLAGS={flags}", f"-DCMAKE_CXX_FLAGS={flags}"] + build.run(["cmake", "-S", str(REPO_ROOT), "-B", str(work_dir), *args]) + build.run(["cmake", "--build", str(work_dir), "-j", str(os.cpu_count() or 4)]) + return build.find_library(work_dir, TARGET.library) + + +def build_test(library: Path, output: Path, *flags: str, source: Path = SOURCE) -> None: + subprocess.run( + [COMPILER, "-std=c++17", "-O1", "-g", "-pthread", *flags, f"-I{REPO_ROOT / 'src'}", + f"-I{build.HEADER.parent}", str(source), f"-L{library.parent}", + "-lstable-diffusion", f"-Wl,-rpath,{library.parent}", "-o", str(output)], + check=True) + + +@unittest.skipUnless(sys.platform == "linux" and LINUX_TARGET, + "set SD_EXIT_TEARDOWN_TARGET to a built Linux target") +class LiveContextAtExitTest(unittest.TestCase): + """Nothing frees a context at exit here, so exiting with one must be harmless. + + Checked on the CPU and, for a Vulkan target, on its first Vulkan device. + What happens on Vulkan is reported, not failed: it is a property of the + driver at hand, which CI only has in software (Mesa lavapipe). + """ + + def test_exits_cleanly_with_a_context_alive(self) -> None: + target = build.TARGETS[LINUX_TARGET] + library = build.BIN_ROOT / target.name / "lib" / target.library + self.assertTrue(library.is_file(), f"no runtime at {library}") + with tempfile.TemporaryDirectory() as directory: + binary = Path(directory) / "exit_teardown_runtime_test" + model = Path(directory) / "model.safetensors" + build_test(library, binary) + subprocess.run([str(binary), "make-model", str(model)], check=True, timeout=600) + backends = ["cpu"] + (["Vulkan0"] if "vulkan" in target.accelerators else []) + for backend in backends: + with self.subTest(backend=backend): + result = subprocess.run( + [str(binary), "idle-untracked", str(model), backend], + capture_output=True, text=True, timeout=600, errors="replace") + generated = f"generated on {backend}" in result.stdout + lines = result.stderr.strip().splitlines() + outcome = (f"{target.name} on {backend}: " + + ("exit with a live context" if generated + else "no context to exit with; the run") + + f" returned {result.returncode}" + + (f": {lines[-1]}" if result.returncode != 0 and lines else "")) + level = "notice" if generated and result.returncode == 0 else "warning" + print(f"::{level} title=exit teardown::{outcome}" + if os.environ.get("GITHUB_ACTIONS") else f"{level}: {outcome}", flush=True) + if backend == "cpu": + self.assertTrue(generated, result.stderr) + self.assertEqual(0, result.returncode, result.stderr) + + +@unittest.skipUnless(sys.platform == "darwin", "exit teardown runs at exit on Apple platforms only") +class ExitTeardownRuntimeTest(unittest.TestCase): + binary: Path + model: Path + + @classmethod + def setUpClass(cls) -> None: + library = build.BIN_ROOT / TARGET.name / "lib" / TARGET.library + if SANITIZER: + library = sanitized_library(SANITIZER) + elif not library.is_file(): + reason = f"no runtime at {library}; build {TARGET.name} first" + if os.environ.get("SD_REQUIRE_RUNTIME"): + raise AssertionError(reason) + raise unittest.SkipTest(reason) + cls.library = library + cls.directory = tempfile.TemporaryDirectory() + cls.addClassCleanup(cls.directory.cleanup) + cls.binary = Path(cls.directory.name) / "exit_teardown_runtime_test" + cls.model = Path(cls.directory.name) / "model.safetensors" + cls.compile(SOURCE, cls.binary) + subprocess.run([str(cls.binary), "make-model", str(cls.model)], check=True, timeout=600, + env={**os.environ, **SCENARIO_ENV}) + + @classmethod + def compile(cls, source: Path, output: Path, *flags: str) -> None: + sanitize = [f"-fsanitize={SANITIZER}"] if SANITIZER else [] + build_test(cls.library, output, *sanitize, *flags, source=source) + + def run_scenario(self, scenario: str, backend: str, model: Path | str | None = None, + size: str | None = None) -> subprocess.CompletedProcess: + command = [str(self.binary), scenario, str(model or self.model), backend] + return subprocess.run(command + ([size] if size else []), capture_output=True, text=True, + timeout=600, errors="replace", env={**os.environ, **SCENARIO_ENV}) + + def check_scenarios(self, model: Path | str | None = None, size: str | None = None) -> None: + for scenario in SCENARIOS: + for backend in BACKENDS: + for run in range(RUNS): + with self.subTest(scenario=scenario, backend=backend, run=run): + result = self.run_scenario(scenario, backend, model, size) + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + + def test_frees_tracked_contexts_at_exit(self) -> None: + self.check_scenarios() + + @unittest.skipUnless(REAL_MODEL, "set SD_EXIT_TEARDOWN_MODEL to a model file") + def test_frees_tracked_contexts_of_a_real_model_at_exit(self) -> None: + self.check_scenarios(REAL_MODEL, REAL_MODEL_SIZE) + + def test_untracked_context_shows_whether_metal_aborts(self) -> None: + """The control: a context from new_sd_ctx() is left alone, as before. + + ggml-metal then aborts at exit where Metal residency sets are live. A + clean exit is not a failure of this repository, but it means the + scenarios above showed the frees through the allocator only. + """ + self.assertEqual(0, self.run_scenario("idle-untracked", "cpu").returncode) + result = self.run_scenario("idle-untracked", "default") + if result.returncode == 0: + message = ("an untracked context exited cleanly: no Metal residency set was live, " + "so the Metal abort at exit was not exercised on this machine") + print(f"::warning title=exit teardown::{message}" if os.environ.get("GITHUB_ACTIONS") + else f"note: {message}", flush=True) + return + self.assertEqual(-signal.SIGABRT, result.returncode, result.stdout + result.stderr) + self.assertIn(METAL_ABORT, result.stderr) + print("note: an untracked context aborted in ggml-metal at exit, as before", flush=True) + + def test_dart_vm_exits_cleanly_with_contexts_no_dart_code_frees(self) -> None: + """`tests/dart/exit_teardown.dart`: what a Flutter quit or hot restart leaves.""" + if SANITIZER: + self.skipTest("the Dart VM cannot load a sanitized runtime") + if DART is None: + if os.environ.get("SD_REQUIRE_DART"): + self.fail("needs a Dart SDK on PATH") + self.skipTest("needs a Dart SDK on PATH") + probe = Path(self.directory.name) / "libexit_teardown_dart_probe.dylib" + self.compile(DART_PROBE, probe, "-shared", "-fPIC") + + def run(scenario: str, api: str) -> subprocess.CompletedProcess: + return subprocess.run( + [DART, str(DART_HARNESS), str(self.library), str(probe), str(self.model), + scenario, api], capture_output=True, text=True, timeout=600, errors="replace") + + for scenario in DART_SCENARIOS: + for attempt in range(RUNS): + with self.subTest(scenario=scenario, run=attempt): + result = run(scenario, "tracked") + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + # The control: with upstream's functions the same exits abort where + # Metal residency sets are live, and nowhere else. + aborted = [scenario for scenario in DART_SCENARIOS + if METAL_ABORT in run(scenario, "raw").stderr] + print(f"note: with new_sd_ctx and generate_image, {len(aborted)} of " + f"{len(DART_SCENARIOS)} Dart scenarios aborted in ggml-metal", flush=True) + + @unittest.skipUnless(LLAMADART, "set SD_EXIT_TEARDOWN_LLAMADART to a libllamadart.dylib") + def test_shares_a_process_with_libllamadart(self) -> None: + """Each library tears down its own registry, whichever was loaded first.""" + binary = Path(self.directory.name) / "exit_teardown_two_libraries_test" + subprocess.run([COMPILER, "-std=c++17", "-O1", "-g", f"-I{REPO_ROOT / 'src'}", + f"-I{build.HEADER.parent}", str(TWO_LIBRARIES), "-o", str(binary)], + check=True) + for first in ("sd", "llama"): + for mode in ("exit", "sd", "llama"): + for backend in BACKENDS: + with self.subTest(first=first, mode=mode, backend=backend): + result = subprocess.run( + [str(binary), str(self.library), LLAMADART, first, mode, + str(self.model), backend], + capture_output=True, text=True, timeout=600, errors="replace") + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + self.assertCountEqual( + ["freed by libstable-diffusion", "freed by libllamadart"], + result.stdout.splitlines()) + + +if __name__ == "__main__": + unittest.main() From 2de1e1a2a2c609385817cd53833ade433adcbe0a Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 6 Oct 2026 08:59:48 -0400 Subject: [PATCH 3/5] Only load models on a Metal device that cannot compute On GitHub's macOS runners every scenario that generates on the default device failed, the untracked control included: ggml-metal crashes with a null dereference in ggml_metal_op_bin as soon as the virtual GPU computes. Loading a model there works, and a loaded context already holds Metal buffers. The test now asks first whether each backend can generate at all. Where the default device cannot, the scenarios that can do without a generation skip it there, the two that need one run on the CPU alone, and the Dart VM harness runs on the CPU. Checked locally with GGML_METAL_NO_RESIDENCY=1, which is the other thing those runners differ in: a library whose static destructors bypass teardown still fails idle and late-load on Metal that way. Refs #10 --- README.md | 4 ++- tests/dart/exit_teardown.dart | 39 ++++++++++++--------- tests/native/exit_teardown_dart_probe.cpp | 7 ++-- tests/native/exit_teardown_runtime_test.cpp | 25 ++++++++++--- tests/test_exit_teardown_runtime.py | 30 ++++++++++++++-- 5 files changed, 78 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index a472024..03f02ab 100644 --- a/README.md +++ b/README.md @@ -247,7 +247,9 @@ and a real `sd_ctx_t`, on Metal and on the CPU: the test writes a 12 MB PixArt model itself, so CI needs no download. Where no Metal residency set is live, as on GitHub's macOS runners, an untracked context exits cleanly and the Metal abort is not exercised; there the test shows each free through the -allocator instead. +allocator instead. Those runners also crash in ggml-metal when their virtual +GPU computes, so on Metal the test only loads models there and runs the +scenarios that need a generation on the CPU. ## Build diff --git a/tests/dart/exit_teardown.dart b/tests/dart/exit_teardown.dart index b737833..6f7bdfb 100644 --- a/tests/dart/exit_teardown.dart +++ b/tests/dart/exit_teardown.dart @@ -1,6 +1,7 @@ // Exit teardown under a real Dart VM. // // Usage: dart exit_teardown.dart +// |default // // Each scenario leaves the process with a context that no Dart code frees, // the way a Flutter quit or hot restart does: @@ -25,6 +26,7 @@ import 'dart:io'; import 'dart:isolate'; typedef LoadNative = Pointer Function(Pointer); +typedef ParamsNative = Pointer Function(Pointer, Pointer); typedef GenerateNative = Bool Function( Pointer, Pointer, @@ -55,12 +57,9 @@ final class Runtime { this.countOut, ); - factory Runtime( - String runtimePath, - String probePath, - String model, - bool tracked, - ) { + factory Runtime(List arguments) { + final [runtimePath, probePath, model, _, api, backend] = arguments; + final tracked = api == 'tracked'; final runtime = DynamicLibrary.open(runtimePath); final probe = DynamicLibrary.open(probePath); final malloc = DynamicLibrary.process() @@ -68,11 +67,15 @@ final class Runtime { Pointer Function(IntPtr), Pointer Function(int) >('malloc'); - final path = malloc(model.length + 1).cast(); - for (var i = 0; i < model.length; i++) { - path[i] = model.codeUnitAt(i); + Pointer native(String text) { + final bytes = malloc(text.length + 1).cast(); + for (var i = 0; i < text.length; i++) { + bytes[i] = text.codeUnitAt(i); + } + bytes[text.length] = 0; + return bytes.cast(); } - path[model.length] = 0; + probe.lookupFunction('probe_quiet')(); return Runtime._( runtime.lookupFunction Function(Pointer)>( @@ -81,9 +84,10 @@ final class Runtime { runtime.lookupFunction( tracked ? 'sd_dart_generate_image' : 'generate_image', ), - probe.lookupFunction Function(Pointer)>( - 'probe_context_params', - )(path.cast()), + probe.lookupFunction('probe_context_params')( + native(model), + native(backend), + ), probe.lookupFunction< Pointer Function(Int), Pointer Function(int) @@ -107,8 +111,9 @@ final class Runtime { } void worker((SendPort, List) arguments) { - final (events, [runtimePath, probePath, model, scenario, api]) = arguments; - final runtime = Runtime(runtimePath, probePath, model, api == 'tracked'); + final (events, runtimeArguments) = arguments; + final scenario = runtimeArguments[3]; + final runtime = Runtime(runtimeArguments); events.send('loading'); final context = runtime.loadContext(); events.send('loaded'); @@ -124,9 +129,9 @@ void worker((SendPort, List) arguments) { } Future main(List arguments) async { - final [runtimePath, probePath, model, scenario, api] = arguments; + final scenario = arguments[3]; if (scenario == 'idle') { - final runtime = Runtime(runtimePath, probePath, model, api == 'tracked'); + final runtime = Runtime(arguments); runtime.generateOn(runtime.loadContext()); return; } diff --git a/tests/native/exit_teardown_dart_probe.cpp b/tests/native/exit_teardown_dart_probe.cpp index b269cf8..20fdda4 100644 --- a/tests/native/exit_teardown_dart_probe.cpp +++ b/tests/native/exit_teardown_dart_probe.cpp @@ -5,6 +5,7 @@ #include "sd_dart_wrapper.h" #include +#include #define PROBE_API extern "C" __attribute__((visibility("default"))) @@ -23,12 +24,14 @@ PROBE_API void probe_quiet(void) { } // Parameters for a model that holds its TAESD decoder itself, as the one -// exit_teardown_runtime_test writes does. Never freed. -PROBE_API sd_ctx_params_t* probe_context_params(const char* model) { +// exit_teardown_runtime_test writes does. `default` lets the runtime pick its +// device. Never freed. +PROBE_API sd_ctx_params_t* probe_context_params(const char* model, const char* backend) { auto* params = static_cast(malloc(sizeof(sd_ctx_params_t))); sd_ctx_params_init(params); params->model_path = model; params->taesd_path = model; + params->backend = std::strcmp(backend, "default") == 0 ? nullptr : backend; params->eager_load = true; return params; } diff --git a/tests/native/exit_teardown_runtime_test.cpp b/tests/native/exit_teardown_runtime_test.cpp index a6144ac..eb1e666 100644 --- a/tests/native/exit_teardown_runtime_test.cpp +++ b/tests/native/exit_teardown_runtime_test.cpp @@ -13,6 +13,11 @@ // per process, and most of them leave by returning from main: on Apple // platforms teardown then runs during exit(). Elsewhere it does not, and only // `idle-untracked` applies: a context that is alive at exit must not abort. +// +// `generate` only loads, generates and frees, to tell whether the backend can +// compute on this machine at all. Where it cannot, SD_EXIT_TEARDOWN_LOAD_ONLY +// makes the scenarios that can do without a generation on that backend skip +// it: a loaded context already holds the backend's buffers. #include "sd_dart_wrapper.h" @@ -107,6 +112,7 @@ void collect_progress(int step, int, float, void*) { } int image_size = 64; +bool load_only = false; // The model `make-model` writes holds its TAESD decoder itself. bool holds_taesd(const char* model) { @@ -357,23 +363,30 @@ int test_idle(const char* model, const char* backend, bool tracked) { expect_teardown_at_exit(1); } sd_ctx_t* context = load(model, backend, tracked); - CHECK(generate(context, 2, tracked) == 1); + CHECK(load_only || generate(context, 2, tracked) == 1); CHECK(sd_dart_exit_tracked_count() == (tracked ? 2 : 0)); if (tracked) { witness.engine.store(engine_of(context)); expect_contexts_freed_first(); } // Tells the caller that whatever follows happened during exit(). - std::printf("generated on %s\n", backend); + std::printf("%s on %s\n", load_only ? "loaded" : "generated", backend); std::fflush(stdout); return 0; } +int test_generate(const char* model, const char* backend) { + sd_ctx_t* context = load(model, backend, false); + CHECK(generate(context, 1, false) == 1); + free_sd_ctx(context); + return 0; +} + int test_dispose(const char* model, const char* backend) { for (int i = 0; i < 3; ++i) { sd_ctx_t* context = load(model, backend); const void* engine = engine_of(context); - if (i != 1) { + if (i != 1 && !load_only) { CHECK(generate(context, 1) == 1); } CHECK(sd_ctx_supports_image_generation(context)); @@ -491,7 +504,7 @@ int test_late_load(const char* model, const char* backend) { sd_ctx_t* first = load(model, "cpu"); CHECK(generate(first, 1) == 1); sd_ctx_t* second = load(model, backend); - CHECK(generate(second, 1) == 1); + CHECK(load_only || generate(second, 1) == 1); // The first context is freed last, so nothing reuses its block before the // witness looks at it. witness.engine.store(engine_of(first)); @@ -515,6 +528,10 @@ int main(int argc, char** argv) { const char* model = argv[2]; const char* backend = argv[3]; image_size = argc == 5 ? std::atoi(argv[4]) : image_size; + load_only = std::getenv("SD_EXIT_TEARDOWN_LOAD_ONLY") != nullptr; + if (scenario == "generate") { + return test_generate(model, backend); + } if (scenario == "idle") { return test_idle(model, backend, true); } diff --git a/tests/test_exit_teardown_runtime.py b/tests/test_exit_teardown_runtime.py index f29774e..8a27646 100644 --- a/tests/test_exit_teardown_runtime.py +++ b/tests/test_exit_teardown_runtime.py @@ -55,6 +55,8 @@ REAL_MODEL = os.environ.get("SD_EXIT_TEARDOWN_MODEL", "") REAL_MODEL_SIZE = os.environ.get("SD_EXIT_TEARDOWN_MODEL_SIZE", "256") SCENARIOS = ("idle", "dispose", "cancel", "generate-wait", "load-wait", "late-load") +# What the others do on a device that cannot compute is load, and free. +GENERATING_SCENARIOS = ("cancel", "generate-wait") # `default` lets the runtime pick its device, which is Metal where there is one. BACKENDS = ("default", "cpu") METAL_ABORT = "[rsets->data count] == 0" @@ -145,6 +147,21 @@ def setUpClass(cls) -> None: cls.compile(SOURCE, cls.binary) subprocess.run([str(cls.binary), "make-model", str(cls.model)], check=True, timeout=600, env={**os.environ, **SCENARIO_ENV}) + # GitHub's macOS runners load a model on their virtual GPU but crash in + # ggml-metal as soon as it computes, whoever created the context. + cls.load_only = { + backend for backend in BACKENDS + if subprocess.run([str(cls.binary), "generate", str(cls.model), backend], + capture_output=True, timeout=600, + env={**os.environ, **SCENARIO_ENV}).returncode != 0} + if "cpu" in cls.load_only: + raise AssertionError("the runtime cannot generate an image on the CPU") + if cls.load_only: + message = ("the default device cannot run a generation on this machine: " + "there the scenarios only load a model, and those that need a " + "generation run on the CPU alone") + print(f"::warning title=exit teardown::{message}" if os.environ.get("GITHUB_ACTIONS") + else f"note: {message}", flush=True) @classmethod def compile(cls, source: Path, output: Path, *flags: str) -> None: @@ -154,12 +171,16 @@ def compile(cls, source: Path, output: Path, *flags: str) -> None: def run_scenario(self, scenario: str, backend: str, model: Path | str | None = None, size: str | None = None) -> subprocess.CompletedProcess: command = [str(self.binary), scenario, str(model or self.model), backend] + load_only = {"SD_EXIT_TEARDOWN_LOAD_ONLY": "1"} if backend in self.load_only else {} return subprocess.run(command + ([size] if size else []), capture_output=True, text=True, - timeout=600, errors="replace", env={**os.environ, **SCENARIO_ENV}) + timeout=600, errors="replace", + env={**os.environ, **SCENARIO_ENV, **load_only}) def check_scenarios(self, model: Path | str | None = None, size: str | None = None) -> None: for scenario in SCENARIOS: for backend in BACKENDS: + if scenario in GENERATING_SCENARIOS and backend in self.load_only: + continue for run in range(RUNS): with self.subTest(scenario=scenario, backend=backend, run=run): result = self.run_scenario(scenario, backend, model, size) @@ -202,10 +223,13 @@ def test_dart_vm_exits_cleanly_with_contexts_no_dart_code_frees(self) -> None: probe = Path(self.directory.name) / "libexit_teardown_dart_probe.dylib" self.compile(DART_PROBE, probe, "-shared", "-fPIC") + backend = "cpu" if self.load_only else "default" + def run(scenario: str, api: str) -> subprocess.CompletedProcess: return subprocess.run( [DART, str(DART_HARNESS), str(self.library), str(probe), str(self.model), - scenario, api], capture_output=True, text=True, timeout=600, errors="replace") + scenario, api, backend], + capture_output=True, text=True, timeout=600, errors="replace") for scenario in DART_SCENARIOS: for attempt in range(RUNS): @@ -216,7 +240,7 @@ def run(scenario: str, api: str) -> subprocess.CompletedProcess: # Metal residency sets are live, and nowhere else. aborted = [scenario for scenario in DART_SCENARIOS if METAL_ABORT in run(scenario, "raw").stderr] - print(f"note: with new_sd_ctx and generate_image, {len(aborted)} of " + print(f"note: with new_sd_ctx and generate_image on {backend}, {len(aborted)} of " f"{len(DART_SCENARIOS)} Dart scenarios aborted in ggml-metal", flush=True) @unittest.skipUnless(LLAMADART, "set SD_EXIT_TEARDOWN_LLAMADART to a libllamadart.dylib") From 738dd7158afd77f895f469ee4cd64fc381f7216e Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 6 Oct 2026 14:40:37 -0400 Subject: [PATCH 4/5] Wait for a free in flight, and longer for a load or a generation An audit of 2de1e1a found two ways a quit still reached ggml-metal's abort. A quit while the last tracked context was being freed. sd_dart_exit_free untracks under the lock and frees outside it, and teardown returned at once when nothing was tracked, so the static destructors ran under the free. With SD-Turbo on Metal, an exit within 3 ms of the start of that free failed 35 times out of 36. The registry now counts the calls in flight that create or free a tracked object, in the same critical section as the untracking, and teardown returns early only when nothing is tracked and no such call is in flight. sd_dart_new_sd_ctx had the same gap between the two critical sections it counted itself in: a teardown in between returned at once and the load began after it. Every call scope now enters and leaves in one critical section. A quit during an uninterruptible phase longer than the 2 s wait. stable-diffusion.cpp reads a cancellation only between the phases of a generation, and the VAE decode of one image takes 4.1 s at 768 x 768 and 7.4 s at 1024 x 1024 with SD-Turbo on an M4 Max: 17 of 25 quits during the first aborted. Teardown now waits up to 15 s while sd_dart_new_sd_ctx or sd_dart_generate_image is in flight, and 2 s as before for any other call. It still cancels every 20 ms, so the wait ends when the call does, and it still frees nothing when the bound runs out. With no call in flight it does not wait: an idle quit takes 0.03 s with either bound. sd_dart_exit_set_wait_ms takes both bounds. So quitting during a large generation can delay the exit of the process by up to 15 s. The header and the README say what the bound covers and what still exceeds it. New stand-in scenarios: a free and a load in flight with nothing tracked, a load loop racing teardown (run 400 times), a generation and a load that outlast the shorter bound, each bound running out, an idle teardown, and a generation that exits from its own callback. Refs #10 --- AGENTS.md | 3 +- README.md | 49 +++++-- src/sd_dart_exit.cpp | 151 ++++++++++++++------ src/sd_dart_wrapper.h | 33 ++++- tests/native/exit_teardown_test.cpp | 213 ++++++++++++++++++++++++++-- tests/test_exit_teardown.py | 23 ++- 6 files changed, 392 insertions(+), 80 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 176d447..c435b5e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,7 +45,8 @@ python3 tools/validate_artifacts.py skips without it. After a change to `src/sd_dart_exit.cpp` or an upstream bump, run it against the release build and again with `SD_EXIT_TEARDOWN_SANITIZER=address`, on a Mac where its control reports the - Metal abort: GitHub's runners do not reach it. + Metal abort, with `SD_REQUIRE_METAL_GENERATION=1`: GitHub's runners reach + neither the abort nor a generation on Metal. - Every shipped library must pass `validate_artifacts.py`: it exports exactly the `SD_API` symbols in `stable-diffusion.h` and `src/sd_dart_wrapper.h`, and links only allowlisted system libraries. Never export an upstream-internal diff --git a/README.md b/README.md index 03f02ab..2e37596 100644 --- a/README.md +++ b/README.md @@ -160,7 +160,7 @@ bool sd_dart_exit_untrack(void* object); int32_t sd_dart_exit_tracked_count(void); void sd_dart_exit_call_begin(void); void sd_dart_exit_call_end(void); -void sd_dart_exit_set_wait_ms(int32_t wait_ms); +void sd_dart_exit_set_wait_ms(int32_t wait_ms, int32_t work_wait_ms); void sd_dart_exit_teardown(void); ``` @@ -183,7 +183,7 @@ statics is destroyed. | `sd_dart_exit_track`, `_untrack` | For C and C++ callers: track any other object, such as an `upscaler_ctx_t`, with its free function and a stage, or stop tracking it. | | `sd_dart_exit_tracked_count` | Number of tracked objects. | | `sd_dart_exit_call_begin`, `_end` | For C and C++ callers: mark a call in flight around an upstream function that has no wrapper. | -| `sd_dart_exit_set_wait_ms` | How long teardown waits for calls in flight; 2000 ms by default. | +| `sd_dart_exit_set_wait_ms` | How long teardown waits for calls in flight: `work_wait_ms` (15000 by default) while a load or a generation is among them, `wait_ms` (2000) otherwise. | | `sd_dart_exit_teardown` | Runs teardown now, for native hosts on platforms where it does not run by itself. Follow it directly with `exit()`. | A Dart caller replaces `new_sd_ctx`, `generate_image`, `sd_cancel_generation` @@ -202,9 +202,15 @@ What is guaranteed: cancel its generation, repeating the request until the calls in flight have ended, and frees only then. If one is still running after the wait, it frees nothing at all rather than what happens to be idle. -- **A bounded exit.** The wait ends after `wait_ms`, plus 250 ms for a thread - that has just left a call to finish the short calls that follow it. With - nothing tracked and no load in flight, teardown returns at once. +- **A bounded exit, and no wait without a reason.** With no call in flight + teardown frees at once; with nothing tracked and no call creating or + freeing a tracked object, it returns at once. Otherwise the wait ends when + the calls do, at the latest after 15 s while `sd_dart_new_sd_ctx` or + `sd_dart_generate_image` is in flight and after 2 s for any other call, plus + 250 ms for a thread that has just left a call to finish the short calls + that follow it. +- **A free in flight is waited for** like a load, also when it is freeing the + last tracked object and the registry is already empty. - **Nothing returns into freed memory.** Once teardown has begun, a thread that ends its outermost call in flight, or reaches the registry outside one, never returns to its caller. @@ -222,12 +228,29 @@ What is guaranteed: What is not: -- **A load cannot be cancelled.** Teardown waits for one like for any call, - and a load that outlasts `wait_ms` leaves everything allocated, so Metal - aborts as before. -- **A sampling step cannot be interrupted.** stable-diffusion.cpp honors a - cancellation between steps. A step that outlasts `wait_ms` has the same - result as a load that does. +- **Most of a generation cannot be interrupted.** stable-diffusion.cpp reads + a cancellation only between the phases of a generation, before a sampling + step and before the decode of each image: a load, the text encoder, one + sampling step and the VAE decode of one image run to their end. Teardown waits up to 15 s for + them, so **quitting during a large generation can delay the exit of the + process by up to 15 s**. A phase that outlasts the bound leaves everything + allocated, and Metal aborts as before. Measured with SD-Turbo (q8, 4 steps) + on an M4 Max: + + | Size | One step, Metal | VAE decode, Metal | VAE decode, CPU | + | --- | --- | --- | --- | + | 512 × 512 | 0.3 s | 1.8 s | 6.7 s | + | 640 × 640 | 0.5 s | 2.8 s | | + | 768 × 768 | 0.7 s | 4.1 s | about 12 s | + | 1024 × 1024 | 1.7 s | 7.4 s | | + + The decode scales with the pixel count, so on that machine the bound holds + to about 1024 × 1024 with a factor of two to spare, and at 768 × 768 on a + GPU up to three times slower. It is exceeded by larger images, by slower + devices at those sizes, by a VAE decode on the CPU a little above + 768 × 768, and by a load that takes longer than 15 s. A host that knows its models + sets its own bounds with `sd_dart_exit_set_wait_ms`; passing 2000 for both + restores a 2 s limit for every call. - **Upstream calls on a tracked context are not waited for.** A raw `generate_image`, `generate_video` or `adetail_image` on a context from `sd_dart_new_sd_ctx` is a use after free at exit. Native callers bracket @@ -249,7 +272,9 @@ live, as on GitHub's macOS runners, an untracked context exits cleanly and the Metal abort is not exercised; there the test shows each free through the allocator instead. Those runners also crash in ggml-metal when their virtual GPU computes, so on Metal the test only loads models there and runs the -scenarios that need a generation on the CPU. +scenarios that need a generation on the CPU. It does the same on any machine +whose default device cannot generate and says so in one line; set +`SD_REQUIRE_METAL_GENERATION=1` on a Mac to make that a failure instead. ## Build diff --git a/src/sd_dart_exit.cpp b/src/sd_dart_exit.cpp index 7376d07..b5b3e3f 100644 --- a/src/sd_dart_exit.cpp +++ b/src/sd_dart_exit.cpp @@ -37,10 +37,15 @@ struct Registry { std::condition_variable idle; std::unordered_map objects; std::chrono::steady_clock::time_point last_call_end{}; - uint64_t next_order = 0; - int32_t calls = 0; - int32_t creating_calls = 0; - int32_t wait_ms = 2000; + uint64_t next_order = 0; + int32_t calls = 0; + // The calls in flight that create or free a tracked object. Their object + // is in the registry only before or after the call, never during it. + int32_t object_calls = 0; + // The loads and generations in flight. + int32_t work_calls = 0; + int32_t wait_ms = 2000; + int32_t work_wait_ms = 15000; std::atomic armed{false}; std::atomic torn_down{false}; }; @@ -56,6 +61,18 @@ thread_local bool teardown_thread = false; // Calls in flight on this thread. The registry counts the outermost one. thread_local int32_t call_depth = 0; +// Loads and generations in flight on this thread. +thread_local int32_t work_depth = 0; + +// What a call in flight is to teardown, beyond something to wait for. +enum CallKind : int { + kPlainCall = 0, + // Creates or frees a tracked object: waited for with nothing tracked. + kObjectCall = 1, + // A load or a generation: waited for with the longer bound. + kWorkCall = 2, +}; + // Called with the registry locked. Returns whether the calling thread may go // on to use tracked objects. Once teardown has begun, teardown waits for a // thread that is in a call in flight, and the teardown thread gets false. Any @@ -124,32 +141,67 @@ bool insert(void* object, void (*free_fn)(void*), void (*cancel_fn)(void*), int3 return true; } -struct Call { - Call() { sd_dart_exit_call_begin(); } - ~Call() { sd_dart_exit_call_end(); } - Call(const Call&) = delete; - Call& operator=(const Call&) = delete; -}; +// Begins a call in flight. Called with the registry locked, by a thread that +// admit() let through, so that teardown sees the call and what it is for in +// the same moment as whatever else the caller changes under that lock. +void enter_call(Registry& registry, int kind) { + if (call_depth++ == 0) { + ++registry.calls; + } + if (kind & kObjectCall) { + ++registry.object_calls; + } + if (kind & kWorkCall) { + ++registry.work_calls; + ++work_depth; + } +} + +// Ends a call in flight. Called with the registry locked. Once teardown has +// begun, it does not return from the outermost call of a thread. +void leave_call(Registry& registry, std::unique_lock& lock, int kind) { + if (kind & kObjectCall) { + --registry.object_calls; + } + if (kind & kWorkCall) { + --registry.work_calls; + --work_depth; + } + if (--call_depth > 0) { + return; + } + --registry.calls; + registry.idle.notify_all(); + if (admit(registry, lock)) { + registry.last_call_end = std::chrono::steady_clock::now(); + } +} -// A call in flight that tracks what it creates. Teardown waits for it even -// when nothing is tracked yet. -struct CreatingCall { - CreatingCall() { - sd_dart_exit_call_begin(); +// A call in flight of the given kind, from one critical section to another: +// teardown never sees the call without what it is for. +struct Call { + explicit Call(int kind) + : kind(kind) { Registry& registry = state(); - std::lock_guard lock(registry.mutex); - ++registry.creating_calls; + std::unique_lock lock(registry.mutex); + counted = admit(registry, lock); + if (counted) { + enter_call(registry, kind); + } } - ~CreatingCall() { - { + ~Call() { + if (counted) { Registry& registry = state(); - std::lock_guard lock(registry.mutex); - --registry.creating_calls; + std::unique_lock lock(registry.mutex); + leave_call(registry, lock, kind); } - sd_dart_exit_call_end(); } - CreatingCall(const CreatingCall&) = delete; - CreatingCall& operator=(const CreatingCall&) = delete; + Call(const Call&) = delete; + Call& operator=(const Call&) = delete; + + int kind; + // False on the teardown thread once teardown has begun. + bool counted; }; void free_context(void* object) { @@ -224,12 +276,13 @@ void sd_dart_exit_free(void* object) { } free_fn = found->second.free_fn; registry.objects.erase(found); - if (call_depth++ == 0) { - ++registry.calls; - } + // In the same critical section as the erase: from here on only this + // call tells teardown that the object still needs the library. + enter_call(registry, kObjectCall); } free_fn(object); - sd_dart_exit_call_end(); + std::unique_lock lock(registry.mutex); + leave_call(registry, lock, kObjectCall); } int32_t sd_dart_exit_tracked_count(void) { @@ -245,30 +298,29 @@ void sd_dart_exit_call_begin(void) { } Registry& registry = state(); std::unique_lock lock(registry.mutex); - if (!admit(registry, lock)) { - return; + if (admit(registry, lock)) { + enter_call(registry, kPlainCall); } - ++registry.calls; - call_depth = 1; } void sd_dart_exit_call_end(void) { - if (call_depth == 0 || --call_depth > 0) { + if (call_depth == 0) { + return; + } + if (call_depth > 1) { + --call_depth; return; } Registry& registry = state(); std::unique_lock lock(registry.mutex); - --registry.calls; - registry.idle.notify_all(); - if (admit(registry, lock)) { - registry.last_call_end = std::chrono::steady_clock::now(); - } + leave_call(registry, lock, kPlainCall); } -void sd_dart_exit_set_wait_ms(int32_t wait_ms) { +void sd_dart_exit_set_wait_ms(int32_t wait_ms, int32_t work_wait_ms) { Registry& registry = state(); std::lock_guard lock(registry.mutex); - registry.wait_ms = std::max(wait_ms, 0); + registry.wait_ms = std::max(wait_ms, 0); + registry.work_wait_ms = std::max(work_wait_ms, 0); } void sd_dart_exit_teardown(void) { @@ -280,14 +332,14 @@ void sd_dart_exit_teardown(void) { return; } teardown_thread = true; - // With nothing to free and nothing being created, a call in flight is - // no reason to hold up the exit. - if (registry.creating_calls == 0 && registry.objects.empty()) { + // With nothing to free and no object being created or freed, a call + // in flight is no reason to hold up the exit. + if (registry.object_calls == 0 && registry.objects.empty()) { return; } // A call in flight on this thread cannot end while teardown runs. const int32_t own_calls = call_depth > 0 ? 1 : 0; - const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(registry.wait_ms); + const auto began = std::chrono::steady_clock::now(); while (registry.calls != own_calls) { // generate_image() clears a cancellation that was requested before // it started, so the request is repeated until the calls end. @@ -296,7 +348,12 @@ void sd_dart_exit_teardown(void) { entry.cancel_fn(object); } } - const auto now = std::chrono::steady_clock::now(); + // A load cannot be cancelled, and a generation only between two + // sampling steps, so those get the longer bound for as long as + // one is in flight. + const bool working = registry.work_calls > work_depth; + const auto deadline = began + std::chrono::milliseconds(working ? registry.work_wait_ms : registry.wait_ms); + const auto now = std::chrono::steady_clock::now(); if (now >= deadline) { return; } @@ -318,7 +375,7 @@ void sd_dart_exit_teardown(void) { } sd_ctx_t* sd_dart_new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { - CreatingCall call; + Call call(kObjectCall | kWorkCall); sd_ctx_t* context = new_sd_ctx(sd_ctx_params); insert(context, free_context, cancel_context, SD_DART_EXIT_STAGE_CONTEXT); return context; @@ -328,7 +385,7 @@ bool sd_dart_generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, int* num_images_out) { - Call call; + Call call(kWorkCall); return generate_image(sd_ctx, sd_img_gen_params, images_out, num_images_out); } diff --git a/src/sd_dart_wrapper.h b/src/sd_dart_wrapper.h index bc55914..2edc59a 100644 --- a/src/sd_dart_wrapper.h +++ b/src/sd_dart_wrapper.h @@ -80,13 +80,29 @@ SD_API size_t sd_dart_progress_read(uint64_t after, // Teardown asks every tracked context to cancel its generation, waits a // bounded time for the calls in flight and then frees the tracked objects in // sd_dart_exit_stage order, latest tracked first within a stage. If a call is -// still in flight when the wait ends, it frees nothing. With nothing tracked -// and no sd_dart_new_sd_ctx() in flight, it does not wait. +// still in flight when the wait ends, it frees nothing. With no call in +// flight it does not wait. Neither does it when nothing is tracked and no +// call in flight is creating or freeing a tracked object: sd_dart_new_sd_ctx() +// and sd_dart_exit_free() are waited for whatever the registry holds. // // A call in flight is the time a thread spends inside sd_dart_new_sd_ctx(), // sd_dart_generate_image() or sd_dart_exit_free(), or between // sd_dart_exit_call_begin() and sd_dart_exit_call_end(). // +// The wait ends when the calls do, and it has two bounds: +// - 15 s while a load or a generation is in flight, that is +// sd_dart_new_sd_ctx() or sd_dart_generate_image(). stable-diffusion.cpp +// reads a cancellation only between the phases of a generation: before a +// sampling step and before the decode of each image. A load, the text +// encoder, a sampling step and the VAE decode of an image each run to their +// end, and the decode alone takes seconds: with SD-Turbo on an M4 Max, +// 4.1 s for 768 x 768 pixels and 7.4 s for 1024 x 1024. So a process that +// quits during a large generation can take that much longer to exit. +// - 2 s otherwise: for sd_dart_exit_free() and for calls marked with +// sd_dart_exit_call_begin(). +// A phase that outlasts its bound leaves everything allocated, and the +// process exits as it would have without this registry. +// // Once teardown has begun, the objects a thread holds may be freed as soon as // it has no call in flight. From then on, the end of a thread's outermost call // in flight never returns to its caller. Neither does a function documented @@ -152,9 +168,11 @@ SD_API int32_t sd_dart_exit_tracked_count(void); SD_API void sd_dart_exit_call_begin(void); SD_API void sd_dart_exit_call_end(void); -// Sets how long teardown waits for calls in flight. Negative values are -// treated as zero. The default is 2000 ms. -SD_API void sd_dart_exit_set_wait_ms(int32_t wait_ms); +// Sets how long teardown waits for calls in flight: `work_wait_ms` while a +// load or a generation is among them, `wait_ms` otherwise. Negative values are +// treated as zero. The defaults are 2000 and 15000. A host that would rather +// abort in ggml-metal than exit late passes one value for both. +SD_API void sd_dart_exit_set_wait_ms(int32_t wait_ms, int32_t work_wait_ms); // Runs exit teardown now; later runs do nothing. Afterwards tracked objects // are unusable and other threads that reach the functions above stay blocked, @@ -166,13 +184,14 @@ SD_API void sd_dart_exit_teardown(void); // new_sd_ctx() that tracks the context in the CONTEXT stage before it // returns. Free the context with sd_dart_exit_free(). A load cannot be -// cancelled: teardown waits for it like for any call in flight. Blocks after +// cancelled: teardown waits for it, with the longer bound. Blocks after // teardown. SD_API sd_ctx_t* sd_dart_new_sd_ctx(const sd_ctx_params_t* sd_ctx_params); // generate_image() as a call in flight. Takes and returns what generate_image() // does. Teardown cancels it with SD_CANCEL_ALL, which stable-diffusion.cpp -// honors between sampling steps. Blocks after teardown. +// honors before a sampling step and before the decode of each image, and +// waits for it with the longer bound. Blocks after teardown. SD_API bool sd_dart_generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, diff --git a/tests/native/exit_teardown_test.cpp b/tests/native/exit_teardown_test.cpp index 25fb872..060f16b 100644 --- a/tests/native/exit_teardown_test.cpp +++ b/tests/native/exit_teardown_test.cpp @@ -104,6 +104,10 @@ std::atomic run_generation{nullptr}; std::atomic hold_free{nullptr}; std::atomic contexts_freed{0}; std::atomic last_freed{nullptr}; +// Upstream calls that began after teardown had returned. +std::atomic teardown_returned{false}; +std::atomic loads_begun{0}; +std::atomic calls_after_teardown{0}; // The arguments the stand-ins were last called with. std::atomic last_load_params{nullptr}; std::atomic last_request{nullptr}; @@ -259,7 +263,7 @@ int test_repeat() { int test_in_flight() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); static std::atomic in_call{false}; static std::atomic returned{false}; @@ -293,7 +297,7 @@ int test_in_flight() { // nested calls and to use the registry until that call ends. int test_nested() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); start_heartbeat(); static char inner[] = "inner"; @@ -355,11 +359,11 @@ int test_settle() { return 0; } -// A generation that never ends: teardown gives up and frees nothing, also not -// the objects no call is using. +// A generation that never ends: teardown gives up after the bound for loads +// and generations and frees nothing, also not the objects no call is using. int test_timeout() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(1000); + sd_dart_exit_set_wait_ms(200, 1000); sd_ctx_params_t params{}; static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); static sd_ctx_t* idle = sd_dart_new_sd_ctx(¶ms); @@ -399,13 +403,104 @@ int test_timeout() { return 0; } +// A call that is neither a load nor a generation gets the shorter bound. +int test_timeout_plain() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(1000, 30000); + static std::atomic in_call{false}; + std::thread([] { + sd_dart_exit_call_begin(); + in_call.store(true); + sleep_ms(600000); + }).detach(); + while (!in_call.load()) { + sleep_ms(1); + } + + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + const int64_t waited = elapsed_ms(started); + CHECK(waited >= 900); + CHECK(waited < 10000); + CHECK(recorded().empty()); + return 0; +} + +// A generation that cannot be cancelled and a load, each outlasting the +// shorter bound: teardown waits for them with the longer one and then frees +// everything. +int test_long_work() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(300, 30000); + sd_ctx_params_t params{}; + static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + static std::atomic working{0}; + run_generation.store([](sd_ctx_t*) { + working.fetch_add(1); + sleep_ms(1200); + record("generation-finished"); + return true; + }); + std::thread([] { + sd_img_gen_params_t request{}; + sd_dart_generate_image(context, &request, nullptr, nullptr); + }).detach(); + while (working.load() < 1) { + sleep_ms(1); + } + hold_load.store([] { + working.fetch_add(1); + sleep_ms(1500); + record("load-finished"); + return true; + }); + std::thread([] { + sd_ctx_params_t late_params{}; + sd_dart_new_sd_ctx(&late_params); + }).detach(); + while (working.load() < 2) { + sleep_ms(1); + } + + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + const int64_t waited = elapsed_ms(started); + CHECK(waited >= 1000); + CHECK(waited < 20000); + std::vector expected = {"generation-finished", "load-finished"}; + expected.insert(expected.end(), kStageOrder.begin(), kStageOrder.end()); + CHECK(recorded() == expected); + CHECK(contexts_freed.load() == 2); + CHECK(sd_dart_exit_tracked_count() == 0); + return 0; +} + +// With contexts tracked and nothing in flight, teardown frees them without +// waiting, however long the bounds are. +int test_idle_no_wait() { + sd_dart_exit_set_wait_ms(30000, 60000); + sd_ctx_params_t params{}; + sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + CHECK(sd_dart_new_sd_ctx(¶ms) != nullptr); + sd_img_gen_params_t request{}; + CHECK(sd_dart_generate_image(context, &request, nullptr, nullptr)); + // Past the time a thread gets to finish what follows its last call. + sleep_ms(400); + + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + CHECK(elapsed_ms(started) < 100); + CHECK(contexts_freed.load() == 2); + return 0; +} + // A call that never ends does not hold up an exit that has nothing to free: // what a caller that tracks nothing gets. int test_idle_wait() { std::thread([] { sd_dart_exit_call_begin(); }).join(); sd_ctx_params_t params{}; sd_ctx_t* untracked = new_sd_ctx(¶ms); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); const auto started = std::chrono::steady_clock::now(); sd_dart_exit_teardown(); CHECK(elapsed_ms(started) < 5000); @@ -418,7 +513,7 @@ int test_idle_wait() { // exits the process, does not wait for the call it is in. int test_own_call() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); sd_dart_exit_call_begin(); const auto started = std::chrono::steady_clock::now(); sd_dart_exit_teardown(); @@ -428,6 +523,37 @@ int test_own_call() { return 0; } +// A generation that exits the process from one of its own callbacks is not a +// reason for the longer bound: teardown cannot wait for the call it is in. +int test_own_work_call() { + track_out_of_stage_order(); + sd_dart_exit_set_wait_ms(300, 30000); + sd_ctx_params_t params{}; + sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + static std::atomic in_call{false}; + std::thread([] { + sd_dart_exit_call_begin(); + in_call.store(true); + sleep_ms(600000); + }).detach(); + while (!in_call.load()) { + sleep_ms(1); + } + static std::atomic waited{-1}; + run_generation.store([](sd_ctx_t*) { + const auto started = std::chrono::steady_clock::now(); + sd_dart_exit_teardown(); + waited.store(elapsed_ms(started)); + return true; + }); + sd_img_gen_params_t request{}; + CHECK(sd_dart_generate_image(context, &request, nullptr, nullptr)); + CHECK(waited.load() >= 250); + CHECK(waited.load() < 10000); + CHECK(recorded().empty()); + return 0; +} + int test_blocked() { char* object = name("object"); CHECK(sd_dart_exit_track(object, free_named, SD_DART_EXIT_STAGE_CONTEXT)); @@ -628,7 +754,7 @@ std::atomic free_started{false}; // else meanwhile, and the freeing thread never returns. int test_free_in_flight() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); start_heartbeat(); sd_ctx_params_t params{}; static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); @@ -659,13 +785,71 @@ int test_free_in_flight() { return 0; } +// The context that is being freed is the only tracked object, so nothing is +// tracked when teardown begins. Teardown still waits for the free: it needs +// the statics that are destroyed once teardown returns. +int test_last_free_in_flight() { + sd_dart_exit_set_wait_ms(30000, 30000); + start_heartbeat(); + sd_ctx_params_t params{}; + static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); + hold_free.store([] { + free_started.store(true); + CHECK(wait_for_teardown()); + sleep_ms(50); + record("free-finished"); + }); + static std::atomic returned{false}; + std::thread([] { + sd_dart_exit_free(context); + returned.store(true); + }).detach(); + while (!free_started.load()) { + sleep_ms(1); + } + CHECK(sd_dart_exit_tracked_count() == 0); + + teardown_requested.store(true); + sd_dart_exit_teardown(); + CHECK(recorded() == std::vector({"free-finished"})); + CHECK(contexts_freed.load() == 1); + sleep_ms(200); + CHECK(!returned.load()); + return 0; +} + +// A thread loads and frees in a loop while teardown runs at a moment of its +// own. Whichever instruction the thread is at, no load and no free begins +// once teardown has returned, and nothing stays tracked. One run tries one +// moment, so this scenario is run many times. +int test_load_race() { + std::thread([] { + sd_ctx_params_t params{}; + for (;;) { + sd_dart_exit_free(sd_dart_new_sd_ctx(¶ms)); + } + }).detach(); + while (loads_begun.load() < 1000) { + std::this_thread::yield(); + } + const auto seed = std::chrono::steady_clock::now().time_since_epoch().count(); + for (volatile int i = 0, spins = static_cast(seed % 20000); i < spins; i = i + 1) { + } + sd_dart_exit_teardown(); + teardown_returned.store(true); + sleep_ms(30); + CHECK(calls_after_teardown.load() == 0); + CHECK(sd_dart_exit_tracked_count() == 0); + return 0; +} + std::atomic load_started{false}; // A load in flight when teardown begins, with nothing tracked yet: teardown // waits for it, the context is tracked although teardown has begun, and it is // freed. The loading thread never returns. int test_load_in_flight() { - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); start_heartbeat(); hold_load.store([] { load_started.store(true); @@ -698,7 +882,7 @@ int test_load_in_flight() { // cancellation the way generate_image() clears one that precedes its start. int test_generate_in_flight() { track_out_of_stage_order(); - sd_dart_exit_set_wait_ms(30000); + sd_dart_exit_set_wait_ms(30000, 30000); start_heartbeat(); sd_ctx_params_t params{}; static sd_ctx_t* context = sd_dart_new_sd_ctx(¶ms); @@ -747,6 +931,8 @@ int test_generate_in_flight() { } // namespace sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { + loads_begun.fetch_add(1); + calls_after_teardown.fetch_add(teardown_returned.load() ? 1 : 0); last_load_params.store(sd_ctx_params); if (bool (*hold)() = hold_load.load()) { hold(); @@ -756,6 +942,7 @@ sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { void free_sd_ctx(sd_ctx_t* sd_ctx) { CHECK(!sd_ctx->generating.load()); + calls_after_teardown.fetch_add(teardown_returned.load() ? 1 : 0); if (void (*hold)() = hold_free.load()) { hold(); } @@ -801,14 +988,20 @@ int main(int argc, char** argv) { {"nested", test_nested}, {"settle", test_settle}, {"timeout", test_timeout}, + {"timeout-plain", test_timeout_plain}, + {"long-work", test_long_work}, + {"idle-no-wait", test_idle_no_wait}, {"idle-wait", test_idle_wait}, {"own-call", test_own_call}, + {"own-work-call", test_own_work_call}, {"blocked", test_blocked}, {"exit", test_exit}, {"late-static", test_late_static}, {"context", test_context}, {"context-stage", test_context_stage}, {"free-in-flight", test_free_in_flight}, + {"last-free-in-flight", test_last_free_in_flight}, + {"load-race", test_load_race}, {"cancel", test_cancel}, {"load-in-flight", test_load_in_flight}, {"generate-in-flight", test_generate_in_flight}, diff --git a/tests/test_exit_teardown.py b/tests/test_exit_teardown.py index 4163124..7d8a7a1 100644 --- a/tests/test_exit_teardown.py +++ b/tests/test_exit_teardown.py @@ -5,6 +5,7 @@ Each scenario runs in its own process, as teardown runs once per process. """ +import concurrent.futures import os import shutil import subprocess @@ -24,6 +25,11 @@ ] COMPILER = os.environ.get("CXX") or shutil.which("c++") SANITIZERS = ("address", "thread") +# Scenarios that try one moment per run, and how often the unsanitized build +# runs them, four at a time. With the two halves of a creating call's +# bookkeeping in separate critical sections, 2 to 5 runs in 100 failed. A +# sanitized build runs them once; it is slower and looks for something else. +REPEATED = {"load-race": 400} SANITIZER_ENV = { # The scenarios leave contexts and blocked threads behind on purpose. "ASAN_OPTIONS": "detect_leaks=0", @@ -59,12 +65,23 @@ def build_and_run(self, *flags: str) -> None: scenarios = subprocess.run([str(binary), "list"], check=True, capture_output=True, text=True).stdout.split() self.assertIn("generate-in-flight", scenarios) + + def run(scenario: str) -> subprocess.CompletedProcess: + return subprocess.run( + [str(binary), scenario], capture_output=True, text=True, timeout=600, + env={**os.environ, **SANITIZER_ENV}) + for scenario in scenarios: with self.subTest(scenario=scenario): - result = subprocess.run( - [str(binary), scenario], capture_output=True, text=True, timeout=600, - env={**os.environ, **SANITIZER_ENV}) + result = run(scenario) self.assertEqual(0, result.returncode, result.stdout + result.stderr) + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + for scenario, runs in ({} if flags else REPEATED).items(): + failures = [result.stdout + result.stderr + for result in pool.map(run, [scenario] * runs) + if result.returncode != 0] + with self.subTest(scenario=scenario, runs=runs): + self.assertEqual([], failures[:3], f"{len(failures)} of {runs} runs failed") def supports(self, flag: str) -> bool: with tempfile.TemporaryDirectory() as directory: From 605f0f4f8cfb8ed3a42f421fe80704af194511b4 Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 6 Oct 2026 14:40:37 -0400 Subject: [PATCH 5/5] Test a quit during the last free, and require Metal on request free-quit exits while the only tracked context is being freed, on the built runtime: with the early return of 2de1e1a it dies of SIGBUS or SIGSEGV on Metal, also without residency sets and without a generation, which is what the hosted jobs run. The Dart VM harness gets the same race as exit-in-free, takes a real model and its image size, and its probe no longer assumes a TAESD decoder in the model file. The fallback for a default device that cannot generate was only visible as a note. SD_REQUIRE_METAL_GENERATION=1 makes it a failure, for runs on a Mac. validate_artifacts.py says what its __cxa_atexit check does not show. Refs #10 --- tests/dart/exit_teardown.dart | 48 ++++++++++++++++----- tests/native/exit_teardown_dart_probe.cpp | 29 +++++++++---- tests/native/exit_teardown_runtime_test.cpp | 34 ++++++++++++++- tests/test_exit_teardown_runtime.py | 45 ++++++++++++------- tools/validate_artifacts.py | 4 +- 5 files changed, 123 insertions(+), 37 deletions(-) diff --git a/tests/dart/exit_teardown.dart b/tests/dart/exit_teardown.dart index 6f7bdfb..6502e89 100644 --- a/tests/dart/exit_teardown.dart +++ b/tests/dart/exit_teardown.dart @@ -1,7 +1,7 @@ // Exit teardown under a real Dart VM. // // Usage: dart exit_teardown.dart -// |default +// |default // // Each scenario leaves the process with a context that no Dart code frees, // the way a Flutter quit or hot restart does: @@ -15,11 +15,15 @@ // exit-in-run main calls C exit() while a worker isolate generates, // as a Flutter quit does. dart:io's exit() would not do: // it leaves without running static destructors +// exit-in-free main tells a worker isolate to free its context, the +// only one, and calls C exit() the given number of +// microseconds later: a quit that races a dispose // -// With api `tracked` the harness calls sd_dart_new_sd_ctx and -// sd_dart_generate_image, and the process must exit with code 0. With api -// `raw` it calls new_sd_ctx and generate_image: the control, which aborts in -// ggml-metal where Metal residency sets are live. +// With api `tracked` the harness calls sd_dart_new_sd_ctx, +// sd_dart_generate_image and sd_dart_exit_free, and the process must exit +// with code 0. With api `raw` it calls new_sd_ctx, generate_image and +// free_sd_ctx: the control, which aborts in ggml-metal where Metal residency +// sets are live. import 'dart:async'; import 'dart:ffi'; import 'dart:io'; @@ -42,6 +46,7 @@ typedef GenerateDart = bool Function( final class Runtime { final Pointer Function(Pointer) load; + final void Function(Pointer) free; final GenerateDart generate; final Pointer contextParams; final Pointer generationParams; @@ -50,6 +55,7 @@ final class Runtime { Runtime._( this.load, + this.free, this.generate, this.contextParams, this.generationParams, @@ -58,7 +64,7 @@ final class Runtime { ); factory Runtime(List arguments) { - final [runtimePath, probePath, model, _, api, backend] = arguments; + final [runtimePath, probePath, model, _, api, backend, size, _] = arguments; final tracked = api == 'tracked'; final runtime = DynamicLibrary.open(runtimePath); final probe = DynamicLibrary.open(probePath); @@ -81,6 +87,10 @@ final class Runtime { runtime.lookupFunction Function(Pointer)>( tracked ? 'sd_dart_new_sd_ctx' : 'new_sd_ctx', ), + runtime.lookupFunction< + Void Function(Pointer), + void Function(Pointer) + >(tracked ? 'sd_dart_exit_free' : 'free_sd_ctx'), runtime.lookupFunction( tracked ? 'sd_dart_generate_image' : 'generate_image', ), @@ -89,9 +99,9 @@ final class Runtime { native(backend), ), probe.lookupFunction< - Pointer Function(Int), - Pointer Function(int) - >('probe_generation_params')(40), + Pointer Function(Int, Int), + Pointer Function(int, int) + >('probe_generation_params')(40, int.parse(size)), malloc(sizeOf>()).cast(), malloc(sizeOf()).cast(), ); @@ -122,6 +132,11 @@ void worker((SendPort, List) arguments) { ReceivePort(); return; } + if (scenario == 'exit-in-free') { + final commands = ReceivePort()..listen((_) => runtime.free(context)); + events.send(commands.sendPort); + return; + } for (;;) { events.send('generating'); runtime.generateOn(context); @@ -142,6 +157,18 @@ Future main(List arguments) async { events.sendPort, arguments, ), onExit: exited.sendPort); + final cExit = DynamicLibrary.process() + .lookupFunction('exit'); + if (scenario == 'exit-in-free') { + final commands = + await events.firstWhere((event) => event is SendPort) as SendPort; + // Past the time teardown gives a thread after its last call in flight. + await Future.delayed(const Duration(milliseconds: 400)); + final clock = Stopwatch()..start(); + commands.send(null); + while (clock.elapsedMicroseconds < int.parse(arguments[7])) {} + cExit(0); + } final waitFor = switch (scenario) { 'leaked' => 'loaded', 'killed-in-load' => 'loading', @@ -150,8 +177,7 @@ Future main(List arguments) async { await events.firstWhere((event) => event == waitFor); if (scenario == 'exit-in-run') { await Future.delayed(const Duration(milliseconds: 20)); - DynamicLibrary.process() - .lookupFunction('exit')(0); + cExit(0); } isolate.kill(priority: Isolate.immediate); await exited.first; diff --git a/tests/native/exit_teardown_dart_probe.cpp b/tests/native/exit_teardown_dart_probe.cpp index 20fdda4..99347b2 100644 --- a/tests/native/exit_teardown_dart_probe.cpp +++ b/tests/native/exit_teardown_dart_probe.cpp @@ -4,8 +4,10 @@ #include "sd_dart_wrapper.h" +#include #include #include +#include #define PROBE_API extern "C" __attribute__((visibility("default"))) @@ -13,6 +15,19 @@ namespace { void quiet_log(enum sd_log_level_t, const char*, void*) {} +// Whether the model holds its TAESD decoder itself, as the one +// exit_teardown_runtime_test writes does. +bool holds_taesd(const char* model) { + std::string header(1 << 16, '\0'); + FILE* file = std::fopen(model, "rb"); + if (file == nullptr) { + return false; + } + header.resize(std::fread(&header[0], 1, header.size(), file)); + std::fclose(file); + return header.find("\"tae.decoder.") != std::string::npos; +} + } // namespace // Keeps the runtime off stdout and stderr without a Dart callback, and gives @@ -20,29 +35,27 @@ void quiet_log(enum sd_log_level_t, const char*, void*) {} PROBE_API void probe_quiet(void) { sd_set_log_callback(quiet_log, nullptr); sd_dart_progress_enable(); - sd_dart_exit_set_wait_ms(120000); + sd_dart_exit_set_wait_ms(120000, 120000); } -// Parameters for a model that holds its TAESD decoder itself, as the one -// exit_teardown_runtime_test writes does. `default` lets the runtime pick its -// device. Never freed. +// `default` lets the runtime pick its device. Never freed. PROBE_API sd_ctx_params_t* probe_context_params(const char* model, const char* backend) { auto* params = static_cast(malloc(sizeof(sd_ctx_params_t))); sd_ctx_params_init(params); params->model_path = model; - params->taesd_path = model; + params->taesd_path = holds_taesd(model) ? model : nullptr; params->backend = std::strcmp(backend, "default") == 0 ? nullptr : backend; params->eager_load = true; return params; } -PROBE_API sd_img_gen_params_t* probe_generation_params(int steps) { +PROBE_API sd_img_gen_params_t* probe_generation_params(int steps, int size) { auto* params = static_cast(malloc(sizeof(sd_img_gen_params_t))); sd_img_gen_params_init(params); params->prompt = "a lighthouse"; params->negative_prompt = ""; - params->width = 64; - params->height = 64; + params->width = size; + params->height = size; params->seed = 7; params->sample_params.sample_steps = steps; params->sample_params.guidance.txt_cfg = 1.0f; diff --git a/tests/native/exit_teardown_runtime_test.cpp b/tests/native/exit_teardown_runtime_test.cpp index eb1e666..19997a4 100644 --- a/tests/native/exit_teardown_runtime_test.cpp +++ b/tests/native/exit_teardown_runtime_test.cpp @@ -496,6 +496,35 @@ int test_load_wait(const char* model, const char* backend) { return 0; } +std::atomic freeing{false}; + +// The process exits while the only tracked context is being freed. Nothing is +// tracked then, and teardown still waits for the free: the free needs the +// statics that are destroyed once teardown is over. +int test_free_quit(const char* model, const char* backend) { + CHECK(atexit([] { + if (!is_freed(witness.engine.load())) { + std::fprintf(stderr, "exit did not wait for the context that was being freed\n"); + std::_Exit(1); + } + }) == 0); + static sd_ctx_t* context = load(model, backend); + CHECK(load_only || generate(context, 1) == 1); + witness.engine.store(engine_of(context)); + std::thread([] { + freeing.store(true); + sd_dart_exit_free(context); + // Reached when the free was over before the exit began. + for (;;) { + sleep_ms(1000); + } + }).detach(); + // Without a sleep: the exit has to begin while the free is running. + while (!freeing.load()) { + } + return 0; +} + // A context on another device, loaded after the first one was tracked, makes // the runtime create statics that did not exist when teardown was registered. // They are destroyed after teardown all the same. @@ -520,7 +549,7 @@ int main(int argc, char** argv) { sd_set_progress_callback(collect_progress, nullptr); // The scenarios are about what teardown does once the calls in flight // have ended. A busy machine must not turn that into the timeout. - sd_dart_exit_set_wait_ms(120000); + sd_dart_exit_set_wait_ms(120000, 120000); if (argc == 3 && scenario == "make-model") { return make_model(argv[2]); } @@ -547,6 +576,9 @@ int main(int argc, char** argv) { if (scenario == "generate-wait") { return test_generate_wait(model, backend); } + if (scenario == "free-quit") { + return test_free_quit(model, backend); + } if (scenario == "load-wait") { return test_load_wait(model, backend); } diff --git a/tests/test_exit_teardown_runtime.py b/tests/test_exit_teardown_runtime.py index 8a27646..1c6af2f 100644 --- a/tests/test_exit_teardown_runtime.py +++ b/tests/test_exit_teardown_runtime.py @@ -11,6 +11,9 @@ SD_REQUIRE_RUNTIME=1 fail instead of skipping SD_REQUIRE_DART=1 fail when no Dart SDK is on PATH; the Dart VM harness skips without one + SD_REQUIRE_METAL_GENERATION=1 fail when the default device cannot run a + generation, instead of only loading + models on it; set it on a Mac SD_EXIT_TEARDOWN_SANITIZER=address build the runtime with AddressSanitizer under build/ and test that one instead @@ -44,7 +47,10 @@ SOURCE = REPO_ROOT / "tests" / "native" / "exit_teardown_runtime_test.cpp" DART_PROBE = REPO_ROOT / "tests" / "native" / "exit_teardown_dart_probe.cpp" DART_HARNESS = REPO_ROOT / "tests" / "dart" / "exit_teardown.dart" -DART_SCENARIOS = ("idle", "leaked", "killed-in-load", "killed-in-run", "exit-in-run") +DART_SCENARIOS = ("idle", "leaked", "killed-in-load", "killed-in-run", "exit-in-run", + "exit-in-free") +# How long after telling a worker to free its context `exit-in-free` exits. +DART_FREE_DELAYS_US = (0, 2000, 6000) COMPILER = os.environ.get("CXX") or shutil.which("c++") DART = shutil.which("dart") TWO_LIBRARIES = REPO_ROOT / "tests" / "native" / "exit_teardown_two_libraries_test.cpp" @@ -54,7 +60,7 @@ RUNS = int(os.environ.get("SD_EXIT_TEARDOWN_RUNS", "1")) REAL_MODEL = os.environ.get("SD_EXIT_TEARDOWN_MODEL", "") REAL_MODEL_SIZE = os.environ.get("SD_EXIT_TEARDOWN_MODEL_SIZE", "256") -SCENARIOS = ("idle", "dispose", "cancel", "generate-wait", "load-wait", "late-load") +SCENARIOS = ("idle", "dispose", "free-quit", "cancel", "generate-wait", "load-wait", "late-load") # What the others do on a device that cannot compute is load, and free. GENERATING_SCENARIOS = ("cancel", "generate-wait") # `default` lets the runtime pick its device, which is Metal where there is one. @@ -156,6 +162,8 @@ def setUpClass(cls) -> None: env={**os.environ, **SCENARIO_ENV}).returncode != 0} if "cpu" in cls.load_only: raise AssertionError("the runtime cannot generate an image on the CPU") + if cls.load_only and os.environ.get("SD_REQUIRE_METAL_GENERATION"): + raise AssertionError("the default device cannot run a generation on this machine") if cls.load_only: message = ("the default device cannot run a generation on this machine: " "there the scenarios only load a model, and those that need a " @@ -224,24 +232,29 @@ def test_dart_vm_exits_cleanly_with_contexts_no_dart_code_frees(self) -> None: self.compile(DART_PROBE, probe, "-shared", "-fPIC") backend = "cpu" if self.load_only else "default" + models = [(str(self.model), "64")] + ([(REAL_MODEL, REAL_MODEL_SIZE)] if REAL_MODEL else []) - def run(scenario: str, api: str) -> subprocess.CompletedProcess: + def run(scenario: str, api: str, model: str, size: str, + attempt: int = 0) -> subprocess.CompletedProcess: + delay = DART_FREE_DELAYS_US[attempt % len(DART_FREE_DELAYS_US)] return subprocess.run( - [DART, str(DART_HARNESS), str(self.library), str(probe), str(self.model), - scenario, api, backend], + [DART, str(DART_HARNESS), str(self.library), str(probe), model, scenario, api, + backend, size, str(delay)], capture_output=True, text=True, timeout=600, errors="replace") - for scenario in DART_SCENARIOS: - for attempt in range(RUNS): - with self.subTest(scenario=scenario, run=attempt): - result = run(scenario, "tracked") - self.assertEqual(0, result.returncode, result.stdout + result.stderr) - # The control: with upstream's functions the same exits abort where - # Metal residency sets are live, and nowhere else. - aborted = [scenario for scenario in DART_SCENARIOS - if METAL_ABORT in run(scenario, "raw").stderr] - print(f"note: with new_sd_ctx and generate_image on {backend}, {len(aborted)} of " - f"{len(DART_SCENARIOS)} Dart scenarios aborted in ggml-metal", flush=True) + for model, size in models: + for scenario in DART_SCENARIOS: + for attempt in range(RUNS): + with self.subTest(model=Path(model).name, scenario=scenario, run=attempt): + result = run(scenario, "tracked", model, size, attempt) + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + # The control: with upstream's functions the same exits abort where + # Metal residency sets are live, and nowhere else. + aborted = [scenario for scenario in DART_SCENARIOS + if METAL_ABORT in run(scenario, "raw", model, size).stderr] + print(f"note: with upstream's functions on {backend}, {len(aborted)} of " + f"{len(DART_SCENARIOS)} Dart scenarios aborted in ggml-metal " + f"({Path(model).name})", flush=True) @unittest.skipUnless(LLAMADART, "set SD_EXIT_TEARDOWN_LLAMADART to a libllamadart.dylib") def test_shares_a_process_with_libllamadart(self) -> None: diff --git a/tools/validate_artifacts.py b/tools/validate_artifacts.py index 47f0b12..303c9f7 100644 --- a/tools/validate_artifacts.py +++ b/tools/validate_artifacts.py @@ -33,7 +33,9 @@ # `src/sd_dart_exit.cpp` defines `__cxa_atexit`, so that every static destructor # of an Apple library runs exit teardown first. A library that imports the # system function registers its destructors past that definition, and exit -# teardown then runs after the statics it needs are gone. +# teardown then runs after the statics it needs are gone. This only shows that +# the definition is there. That it wraps each destructor is what +# `tests/test_exit_teardown.py` and `tests/test_exit_teardown_runtime.py` show. APPLE_STATIC_DESTRUCTOR_IMPORT = "___cxa_atexit" ELF_DEPENDENCIES = { "android": {"libc.so", "libm.so", "libdl.so", "liblog.so"},