Add idle model unload for mediapipe LLM graphs - #4486
Open
atobiszei wants to merge 25 commits into
Open
Conversation
…vable_loading # Conflicts: # src/mediapipe_internal/mediapipefactory.cpp
Cherry-picked from idle-models branch (65f2bab). Adds UNLOADED state machine, per-graph idle unload, ActiveInferenceGuard, config/schema/metrics additions, and related tests.
Cherry-picked from idle-models branch (ceb9f4d). Adds ModelGroupManager, group_name config field, CLI idle_unload_timeout_seconds, status codes, and model_group_manager tests.
atobiszei
commented
Aug 28, 2026
atobiszei
commented
Aug 28, 2026
atobiszei
commented
Aug 28, 2026
atobiszei
commented
Aug 28, 2026
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Aug 31, 2026
Comment on lines
+183
to
+188
| // Idle unload: timestamp (nanoseconds from steady_clock epoch) of the last | ||
| // inference activity. Updated in create() on every inference acquisition and | ||
| // when an in-flight inference finishes (via ActiveInferenceGuard destructor). | ||
| // Held as shared_ptr so executors can safely write to it even after the | ||
| // definition is retired/destroyed — the atomic outlives the definition. | ||
| std::shared_ptr<std::atomic<int64_t>> lastActivityTimeNs; |
Collaborator
Author
There was a problem hiding this comment.
This possibly could be moved to servable definition?
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Aug 31, 2026
atobiszei
commented
Sep 1, 2026
| // Expect old stream to continue responding until closure | ||
| // Expect new stream to be rejected | ||
| TEST_F(MediapipeStreamFlowAddTest, InferOnUnloadedGraph) { | ||
| TEST_F(MediapipeStreamFlowAddTest, InferOnSleepingGraph) { |
atobiszei
commented
Sep 1, 2026
atobiszei
force-pushed
the
atobisze_idle_servable_loading
branch
from
September 1, 2026 16:16
9a971bb to
069d8c1
Compare
atobiszei
commented
Sep 1, 2026
atobiszei
force-pushed
the
atobisze_idle_servable_loading
branch
from
September 2, 2026 09:39
b07f19d to
7885431
Compare
atobiszei
force-pushed
the
atobisze_idle_servable_loading
branch
from
September 3, 2026 10:06
04f5c3a to
bae1e86
Compare
atobiszei
force-pushed
the
atobisze_idle_servable_loading
branch
from
September 3, 2026 13:37
ce0456c to
426797f
Compare
atobiszei
changed the base branch from
atobisze_refactor_servable_loading
to
main
September 3, 2026 14:21
atobiszei
marked this pull request as ready for review
September 3, 2026 14:21
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces substantial new lifecycle and concurrency behavior in core serving components and currently contains several concrete correctness/robustness issues that should be addressed before approval.
Pull request overview
This PR introduces an idle-unload (“sleep”) lifecycle for MediaPipe LLM graphs (and related servables) so heavy resources can be freed after inactivity and reloaded on-demand, building on the concurrent-safe servable loading queue work from PR #4476. It also adds servable grouping for coordinated load/unload, new SLEEPING states, supporting config/schema/CLI wiring, metrics, and extensive unit/integration tests.
Changes:
- Add a serialized
ServableLoadingQueueand use it to route model/mediapipe load, wake, sleep, and retire operations. - Implement “SLEEPING” lifecycle for models and MediaPipe graphs (idle unload + lazy wake-up), plus group-based idle management (
ServableGroupManager). - Extend schema/CLI/docs/metrics and add new tests validating lifecycle, concurrency, and config behavior.
File summaries
| File | Description |
|---|---|
| src/test/test_utils.hpp | Update includes; add test seams for idle-unload lifecycle. |
| src/test/test_utils.cpp | Fix include path for servable manager module. |
| src/test/stress_test_utils.hpp | Fix include path for servable manager module. |
| src/test/streaming_test.cpp | Fix include path for servable manager module. |
| src/test/server_test.cpp | Fix include paths for model/servable manager modules. |
| src/test/servable_loading_queue_test.cpp | New unit tests for loading queue behavior. |
| src/test/servable_group_manager_test.cpp | New tests for group swap/load/unload semantics. |
| src/test/schema_test.cpp | Add schema tests for mediapipe idle unload timeout + group name. |
| src/test/pythonnode_test.cpp | Fix include path for servable manager module. |
| src/test/pipelinedefinitionstatus_test.cpp | Add tests for new SLEEPING pipeline state transitions. |
| src/test/multipart_calculator_test.cpp | Fix include path for servable manager module. |
| src/test/modelmanager_test.cpp | Update mocks for lazy-load signature; add wake-up/retire correctness test. |
| src/test/model_test.cpp | Fix include path for model manager. |
| src/test/model_cache_test.cpp | Fix include path for model manager. |
| src/test/mockmodelinstancechangingstates.hpp | Support lazy-load path in test mock instance. |
| src/test/metrics_flow_test.cpp | Fix include path; ensure module shutdown ordering in test server. |
| src/test/mediapipeflow_test.cpp | Add idle-unload/wake/retire tests; add guards-related tests. |
| src/test/mediapipe_framework_test.cpp | Fix include path for servable manager module. |
| src/test/llm/llmnode_test.cpp | Add LLM graph idle-unload lifecycle & concurrency tests. |
| src/test/kfs_rest_test.cpp | Fix include path for servable manager module. |
| src/test/kfs_metadata_test.cpp | Update overridden loadModel signature for lazy-load compatibility. |
| src/test/idle_model_test.cpp | New tests for model sleeping/wake/retire concurrency. |
| src/test/idle_mediapipe_test.cpp | New tests for mediapipe sleeping/wake/retire behaviors. |
| src/test/http_rest_api_handler_test.cpp | Fix include paths for model/servable manager modules. |
| src/test/http_openai_handler_test.cpp | Fix include path for servable manager module. |
| src/test/get_mediapipe_graph_metadata_response_test.cpp | Fix include path for model manager. |
| src/test/environment.hpp | Add env-gated switch for running all idle tests. |
| src/test/environment.cpp | Implement RUN_ALL_IDLE env handling. |
| src/test/ensemble_config_change_stress.cpp | Fix include path for servable manager module. |
| src/test/embeddingsnode_test.cpp | Fix include path for servable manager module. |
| src/test/constructor_enabled_model_manager.hpp | Add ctor for idle group manager; expose loading queue for tests. |
| src/test/constructor_enabled_model_manager.cpp | Initialize ServableGroupManager in new ctor. |
| src/test/c_api_tests.cpp | Fix include path; update overridden loadModel signature. |
| src/test/c_api_stress_tests.cpp | Fix include path for servable manager module. |
| src/status.hpp | Add new status codes for sleep/unload and group management. |
| src/status.cpp | Add status messages for new status codes. |
| src/single_version_servable_definition.hpp | Rename request counter to pendingCreateExecutorCount. |
| src/server.cpp | Switch includes to new servable_management paths. |
| src/servable_management/servablemanagermodule.hpp | Fix include path for module base header. |
| src/servable_management/servablemanagermodule.cpp | Fix include paths to src/...; align with new build layout. |
| src/servable_management/servable_loading_task.hpp | New task type struct for loading queue. |
| src/servable_management/servable_loading_queue.hpp | New serialized task queue interface. |
| src/servable_management/servable_loading_queue.cpp | New queue implementation with worker thread. |
| src/servable_management/servable_group_manager.hpp | New group manager interface for idle group swaps/unloads. |
| src/servable_management/servable_group_manager.cpp | New group manager implementation, load/swap/unload logic. |
| src/servable_management/modelmanager.hpp | Add loading queue + group manager hooks; new wake/sleep APIs. |
| src/servable_management/modelmanager.cpp | Route config loads through queue; add idle unload sweeps and wake/sleep scheduling. |
| src/servable_management/BUILD | New Bazel targets for servable management components. |
| src/schema.cpp | Add idle_unload_timeout_seconds and group_name to schema. |
| src/modelversionstatus.hpp | Add ModelVersionState::SLEEPING and helper methods. |
| src/modelversionstatus.cpp | Implement SLEEPING status behaviors and string conversion. |
| src/modelinstance.hpp | Extend loadModel signature; add wakeUpIfSleeping/putToSleep API. |
| src/modelinstance.cpp | Implement lazy load (sleep), wake-up, and sleep transitions. |
| src/modelconfig.hpp | Add model groupName config field + accessors. |
| src/modelconfig.cpp | Parse group_name; include it in reload-required checks. |
| src/model.hpp | Add lazyLoad parameters and wake/sleep model helpers. |
| src/model.cpp | Implement lazy-load addVersion/addVersions and model wake/sleep methods. |
| src/model_metric_reporter.hpp | Add ovms_graph_loaded gauge for mediapipe graphs. |
| src/model_metric_reporter.cpp | Register and populate new ovms_graph_loaded metric. |
| src/metrics/metric_config.hpp | Add metric family name for graph loaded gauge. |
| src/metrics/metric_config.cpp | Define ovms_graph_loaded metric constant. |
| src/mediapipe_internal/mediapipegraphexecutor.hpp | Add ActiveInferenceGuard and pass counters/timestamps to executor. |
| src/mediapipe_internal/mediapipegraphexecutor.cpp | Wire ActiveInferenceGuard into executor construction. |
| src/mediapipe_internal/mediapipegraphdefinition.hpp | Add idle-unload APIs, counters, timestamp tracking, lifecycle mutex. |
| src/mediapipe_internal/mediapipegraphdefinition.cpp | Implement putToSleep/wakeUpIfSleeping/idle detection; integrate metric + serialization. |
| src/mediapipe_internal/mediapipegraphconfig.hpp | Add idleUnloadTimeoutSeconds and groupName fields. |
| src/mediapipe_internal/mediapipegraphconfig.cpp | Parse new mediapipe config fields with validation. |
| src/mediapipe_internal/mediapipefactory.hpp | Add wake/sleep/retire APIs; support lazyLoad createDefinition. |
| src/mediapipe_internal/mediapipefactory.cpp | Implement wake/sleep/retire; handle lazyLoad and alias registration. |
| src/kfs_frontend/kfs_grpc_inference_service.cpp | Treat SLEEPING as “ready” via appearsAvailable; fix includes. |
| src/http_server.cpp | Map SERVER_SHUTTING_DOWN and new mediapipe sleep errors to HTTP. |
| src/http_rest_api_handler.cpp | Fix includes; include group manager header. |
| src/grpcservermodule.cpp | Fix includes to new servable_management paths. |
| src/grpc_utils.cpp | Map SERVER_SHUTTING_DOWN to UNAVAILABLE. |
| src/dags/pipelinedefinitionstatus.hpp | Add PipelineDefinitionStateCode::SLEEPING + SleepEvent and helpers. |
| src/dags/pipelinedefinitionstatus.cpp | Implement SLEEPING state transitions and appearsAvailable behavior. |
| src/dags/pipelinedefinition.cpp | Update pending counter name used during reload/retire. |
| src/config.hpp | Add server setting getter for idle unload timeout seconds. |
| src/config.cpp | Allow group_name in config export add/remove model settings; expose idle timeout. |
| src/config_export_module/config_export.cpp | Export optional group_name to config JSON. |
| src/cli_parser.cpp | Add CLI flags for idle unload timeout and group_name. |
| src/capi_frontend/server_settings.hpp | Add idle unload timeout + groupName to server/models settings. |
| src/capi_frontend/capi.cpp | Fix includes to new servable_management paths. |
| src/capi_frontend/capi_dag_utils.cpp | Map SLEEPING pipeline state to AVAILABLE for C API. |
| src/BUILD | Move model manager + servable module targets; add new tests to build. |
| spelling-whitelist.txt | Add whitelist entries for new idle-unload-related identifiers/files. |
| docs/llm/reference.md | Document idle unload option and ovms_graph_loaded metric. |
Review details
Suppressed comments (1)
src/servable_management/modelmanager.cpp:1800
- createPipeline() only triggers wake-up when the ServableGroupManager feature is enabled, but Mediapipe graphs can enter SLEEPING due to per-graph idle_unload_timeout_seconds even when group management is disabled. In that case, MediapipeGraphDefinition::create() will just wait until timeout because nothing schedules a wake-up. Wake SLEEPING graphs on-demand here before calling mediapipeFactory->create().
- Files reviewed: 87/87 changed files
- Comments generated: 11
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+67
to
+82
| std::future<Status> ServableLoadingQueue::scheduleTask(ServableLoadingTask task) { | ||
| auto future = task.completion.get_future(); | ||
| { | ||
| std::lock_guard<std::mutex> lock(this->mutex); | ||
| if (this->taskObserver) { | ||
| this->taskObserver(TaskEvent::Scheduled, task); | ||
| } | ||
| if (task.urgent) { | ||
| this->queue.push_front(std::move(task)); | ||
| } else { | ||
| this->queue.push_back(std::move(task)); | ||
| } | ||
| } | ||
| this->cv.notify_one(); | ||
| return future; | ||
| } |
Comment on lines
+546
to
+549
| int64_t lastActivity = lastActivityTimeNs->load(std::memory_order_relaxed); | ||
| int64_t nowNs = std::chrono::steady_clock::now().time_since_epoch().count(); | ||
| int64_t timeoutNs = timeoutSeconds * 1'000'000'000LL; | ||
| return (nowNs - lastActivity) >= timeoutNs; |
Comment on lines
+92
to
+94
| void recordActivity(int64_t timestampNs = std::chrono::steady_clock::now().time_since_epoch().count()) { | ||
| lastActivityTimeNs->store(timestampNs, std::memory_order_relaxed); | ||
| } |
Comment on lines
+80
to
+83
| lastActivityTimeNs->store( | ||
| std::chrono::steady_clock::now().time_since_epoch().count(), | ||
| std::memory_order_relaxed); | ||
| } |
Comment on lines
+1357
to
+1359
| SPDLOG_INFO("Waking up model: {}, version: {} ...", getName(), getVersion()); | ||
| auto loadStatus = loadModelImpl(this->config); | ||
| if (!loadStatus.ok()) { |
Comment on lines
+284
to
+298
| for (auto& [name, future] : futures) { | ||
| auto status = future.get(); | ||
| if (!status.ok()) { | ||
| SPDLOG_LOGGER_WARN(modelmanager_logger, "Failed to unload '{}' in group '{}': {}", name, groupName, status.string()); | ||
| } else { | ||
| SPDLOG_LOGGER_INFO(modelmanager_logger, "Unloaded '{}' in group '{}'", name, groupName); | ||
| } | ||
| } | ||
|
|
||
| if (isActiveGroup(groupName)) { | ||
| setActiveGroup(""); | ||
| } | ||
| SPDLOG_LOGGER_INFO(modelmanager_logger, "Model group '{}' unloaded successfully", groupName); | ||
| return StatusCode::OK; | ||
| } |
Comment on lines
+381
to
+384
| int64_t lastActivity = lastActivityTimeNs->load(std::memory_order_relaxed); | ||
| int64_t nowNs = std::chrono::steady_clock::now().time_since_epoch().count(); | ||
| int64_t timeoutNs = static_cast<int64_t>(idleTimeoutMicroseconds) * 1'000LL; | ||
| if ((nowNs - lastActivity) < timeoutNs) { |
Comment on lines
+16
to
+19
| #include <condition_variable> | ||
| #include <mutex> | ||
| #include <string> | ||
| #include <vector> |
Comment on lines
+45
to
+49
| src/mediapipe_internal/mediapipegraphdefinition.cpp | ||
| src/mediapipe_internal/mediapipegraphdefinition.hpp | ||
| src/model_group_manager.cpp | ||
| src/test/llm/llmnode_test.cpp | ||
| nowNs ==> knowns, nouns |
Comment on lines
+782
to
+783
| // Identity of the sidePacketMaps shared_ptr, so a test can detect whether it | ||
| // was reset/swapped (unload uses clear(), not reset(), so the pointer must be stable). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-picked from idle-models branch (65f2bab).
Based on top of #4476