feat(infrastructure): lazy backend proxy — start backends on first request, stop after idle TTL - #474
Conversation
… lifecycle Implements Phase 1.5 'Lazy lifecycle wrappers' from the framework integration bridge design spec. LazyBackendProxy is a transparent TCP proxy that: - Listens on a user-facing port without starting the backend - On first inbound connection, runs start_cmd to launch the real backend (sd-server, llama-server, whisper.cpp server, etc.) - Bidirectionally forwards all TCP traffic to the backend - Stops the subprocess after idle_timeout_seconds with no connections - Handles concurrent requests during cold start via a lock - Returns HTTP 503 when the backend is unavailable ~180 LOC — well under the ~50 LOC per-backend estimate since the proxy is fully generic (host:port + start_cmd is all that differs per backend). 9 tests covering: lifecycle start/stop, bidirectional forwarding, 503 on backend failure, real subprocess cold start, subprocess exit handling, idle timeout, and keepalive via active requests. Refs jaylfc#63
📝 WalkthroughWalkthroughAdds LazyBackendProxy: an asyncio TCP proxy that listens without starting the backend, launches a configured subprocess on first connection with health polling, forwards bytes bidirectionally, and stops the backend after an idle timeout. Includes tests for lifecycle, forwarding, cold-start/failure, and idle eviction. ChangesLazy Backend Proxy
Sequence DiagramsequenceDiagram
participant Client
participant Proxy
participant HealthEndpoint
participant Backend
Client->>Proxy: First connection
Proxy->>Proxy: _ensure_backend (locked)
Proxy->>Proxy: subprocess.Popen(start_cmd)
Proxy->>HealthEndpoint: GET health_url (polling)
HealthEndpoint-->>Proxy: 2xx OK
Client->>Proxy: Send request bytes
Proxy->>Backend: Connect and forward bytes
Backend-->>Proxy: Response bytes
Proxy->>Client: Forward response
Proxy->>Proxy: _restart_idle_timer
Note over Proxy: Idle timeout expires
Proxy->>Backend: terminate() / kill()
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tinyagentos/lazy_backend_proxy.py`:
- Around line 107-155: The idle-timer must be canceled while any connection is
active: in _handle_connection, cancel the idle eviction timer when a new
connection is accepted (before await self._ensure_backend()) and increment an
active-connection counter (e.g., self._active_connections); when the piping
finishes (in the finally block) decrement that counter and only call
_restart_idle_timer when the counter drops to zero. Add a helper (or reuse
existing) methods to cancel the timer (e.g., _cancel_idle_timer) and to
atomically manage self._active_connections so long-lived streams (SSE) prevent
the backend from being stopped mid-stream; ensure all early-return/error paths
also decrement the counter and avoid leaving the timer canceled permanently.
- Around line 168-196: The health-check loop doesn't log the cold-start
duration; record start = time.monotonic() just before invoking
subprocess.Popen() and compute elapsed = time.monotonic() - start where you
currently log "starting" and "healthy" and where you raise errors. Update the
logger.info call that notes starting to include elapsed (even zero), include
elapsed in the success log inside the httpx check (the "healthy" message), and
include elapsed in the RuntimeError and TimeoutError messages (referencing
self._start_cmd, self._proc.returncode, self._health_url, _COLD_START_TIMEOUT
and _HEALTH_POLL_INTERVAL to locate the code). Ensure elapsed is formatted
succinctly (e.g., seconds with one or two decimals).
- Around line 198-232: _stop_subprocess is async but uses blocking
subprocess.wait and also drops self._proc early when _stop_cmd is used; change
it to use asyncio-friendly waiting (e.g., await asyncio.create_subprocess_shell
for _stop_cmd or await asyncio.to_thread(self._stop_proc.wait) wrapped in
asyncio.wait_for) and after running the helper ensure you also terminate/await
the actual backend process (self._proc.terminate(), await its exit with
asyncio.wait_for(awaitable to wait on self._proc.wait via to_thread), then kill
if timeout) before setting self._proc = None; likewise make _kill_subprocess
async (or replace its blocking wait with await
asyncio.to_thread(self._proc.wait) wrapped in asyncio.wait_for) so calls from
async _ensure_backend() do not block the event loop and only clear self._proc
after the subprocess has actually exited; use the symbols _stop_subprocess,
_kill_subprocess, self._proc, _stop_cmd, and _STOP_GRACE_PERIOD in your changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ba9abd56-8f67-41a6-981d-1e2a8bb7f3d2
📒 Files selected for processing (2)
tests/test_lazy_backend_proxy.pytinyagentos/lazy_backend_proxy.py
Code Review SummaryStatus: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
Other Observations (not in diff)Issues found in unchanged code that cannot receive inline comments:
Files Reviewed (2 files)
Reviewed by nemotron-3-super-120b-a12b-20230311:free · 459,788 tokens |
- Track active connections count; cancel idle timer on accept, only restart when connections hit zero (CodeRabbit critical, kilo line 129) - Remove idle restart on failed backend connections (kilo line 129) - Log cold-start duration in health check (CodeRabbit major) - Use asyncio.to_thread for subprocess waits — no event-loop blocking in _stop_subprocess / _kill_subprocess (CodeRabbit major) - Fix _stop_subprocess to always terminate/wait backend after stop_cmd helper, not just the helper process (CodeRabbit major) - Health check: require 200-299, not <500 (kilo line 185)
hognek
left a comment
There was a problem hiding this comment.
All review feedback addressed in 688802f:
-
Active connection tracking — added
_active_connectionscounter. Idle timer canceled on accept, only restarted when counter hits zero. SSE/long-lived streams won't get killed mid-stream. Removed idle restart on failed backend connections. -
Cold-start logging — captures
time.monotonic()beforePopen(), logs elapsed in healthy/timeout messages (e.g. "healthy in 3.21s", "timed out after 120.00s"). -
Async shutdown —
_stop_subprocessand_kill_subprocessnow useasyncio.to_thread()for all subprocess waits._stop_subprocessalways terminates/awaits the actual backend process after runningstop_cmdhelper — no more orphaned backends. -
Health check — now requires
200 <= status < 300instead of< 500.
Tests: 9/9 pass.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tinyagentos/lazy_backend_proxy.py (1)
118-136:⚠️ Potential issue | 🔴 Critical | ⚡ Quick win
_active_connectionsleaks on early-return error paths.When
_ensure_backend()(line 127) orasyncio.open_connection()(line 136) fails, the handler returns early, bypassing thefinallyblock that decrements_active_connections. After any such failure, the counter stays inflated forever, preventing idle timeout from ever triggering.🐛 Proposed fix: wrap the entire handler body in a single try/finally
async def _handle_connection( self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter, ) -> None: if not self._running: client_writer.close() return self._active_connections += 1 self._cancel_idle() backend_writer: asyncio.StreamWriter | None = None try: - await self._ensure_backend() - except Exception: - _write_503(client_writer) - return - - try: - backend_reader, backend_writer = await asyncio.wait_for( - asyncio.open_connection(self._backend_host, self._backend_port), - timeout=5.0, - ) - except Exception: - _write_503(client_writer) - return - - # Bidirectional copy. - async def _pipe(src: asyncio.StreamReader, dst: asyncio.StreamWriter): - try: - while True: - data = await src.read(65536) - if not data: - break - dst.write(data) - await dst.drain() - except (ConnectionResetError, BrokenPipeError, OSError): - pass - - try: - await asyncio.gather( - _pipe(client_reader, backend_writer), - _pipe(backend_reader, client_writer), - ) - except Exception: - pass + try: + await self._ensure_backend() + except Exception: + _write_503(client_writer) + return + + try: + backend_reader, backend_writer = await asyncio.wait_for( + asyncio.open_connection(self._backend_host, self._backend_port), + timeout=5.0, + ) + except Exception: + _write_503(client_writer) + return + + # Bidirectional copy. + async def _pipe(src: asyncio.StreamReader, dst: asyncio.StreamWriter): + try: + while True: + data = await src.read(65536) + if not data: + break + dst.write(data) + await dst.drain() + except (ConnectionResetError, BrokenPipeError, OSError): + pass + + try: + await asyncio.gather( + _pipe(client_reader, backend_writer), + _pipe(backend_reader, client_writer), + ) + except Exception: + pass finally: if backend_writer is not None: backend_writer.close() client_writer.close() self._active_connections -= 1 if self._running and self._active_connections == 0: await self._restart_idle_timer()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/lazy_backend_proxy.py` around lines 118 - 136, The handler increments self._active_connections then returns early on errors in _ensure_backend() or asyncio.open_connection(), leaking the counter; wrap the entire handler body that runs after incrementing (including calls to _ensure_backend(), asyncio.open_connection(), the main proxy loop, and any writes like _write_503(client_writer)) in a try/finally so that the finally always decrements self._active_connections and calls any cleanup (e.g., close backend_writer if set, cancel tasks) and still calls _cancel_idle() as appropriate; locate the logic around _active_connections, _cancel_idle, _ensure_backend, asyncio.open_connection, backend_writer and _write_503 to implement the try/finally cleanup.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tinyagentos/lazy_backend_proxy.py`:
- Around line 118-136: The handler increments self._active_connections then
returns early on errors in _ensure_backend() or asyncio.open_connection(),
leaking the counter; wrap the entire handler body that runs after incrementing
(including calls to _ensure_backend(), asyncio.open_connection(), the main proxy
loop, and any writes like _write_503(client_writer)) in a try/finally so that
the finally always decrements self._active_connections and calls any cleanup
(e.g., close backend_writer if set, cancel tasks) and still calls _cancel_idle()
as appropriate; locate the logic around _active_connections, _cancel_idle,
_ensure_backend, asyncio.open_connection, backend_writer and _write_503 to
implement the try/finally cleanup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 89421f3e-0370-4c03-8dfa-5a63900e2be6
📒 Files selected for processing (1)
tinyagentos/lazy_backend_proxy.py
… leak _active_connections was leaked on early returns from _ensure_backend() or asyncio.open_connection() errors — the counter never decremented, preventing idle timeout from ever firing.
hognek
left a comment
There was a problem hiding this comment.
Fixed in c7a0a9f — wrapped the entire handler body (after counter increment) in a single outer try/finally so the counter is always decremented, even on early returns from _ensure_backend() or open_connection() errors.
|
Really nice feature 🙌 Two things worth tightening before merge (flagged by the review and I agree): the idle timer resets on failed backend connections — it should only reset on a successful/active request, or an unreachable backend keeps itself alive; and the health check treats 4xx as healthy — only 2xx/3xx should pass (a 404/401 from the backend isn't "up"). Want to push a fix, or shall I? |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tinyagentos/lazy_backend_proxy.py (1)
177-178:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't log the raw
start_cmd.Line 178 writes the full shell command to application logs. If the configured command includes tokens, credentials, or other sensitive args, they get exposed on every cold start. Prefer a generic message or a redacted executable name instead.
🔒 Minimal change
- logger.info("lazy-proxy :%d → starting: %r", self._proxy_port, self._start_cmd) + logger.info("lazy-proxy :%d → starting backend", self._proxy_port)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/lazy_backend_proxy.py` around lines 177 - 178, The log currently prints the full sensitive shell command via self._start_cmd in the LazyBackendProxy cold-start path; change the logger.info call to avoid exposing raw args by logging a generic message or a redacted executable name instead (e.g., derive and log only the executable basename from self._start_cmd or replace args with "<REDACTED>"), updating the logger line that references self._start_cmd so it no longer emits full command tokens.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tinyagentos/lazy_backend_proxy.py`:
- Around line 177-178: The log currently prints the full sensitive shell command
via self._start_cmd in the LazyBackendProxy cold-start path; change the
logger.info call to avoid exposing raw args by logging a generic message or a
redacted executable name instead (e.g., derive and log only the executable
basename from self._start_cmd or replace args with "<REDACTED>"), updating the
logger line that references self._start_cmd so it no longer emits full command
tokens.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a736a5c9-7399-41d9-a29e-6486be4ccac5
📒 Files selected for processing (1)
tinyagentos/lazy_backend_proxy.py
Problem
Backend servers that load models at process start (sd-server, llama-server,
whisper.cpp) have no built-in idle eviction, wasting GPU VRAM when idle.
Issue #63.
Solution
LazyBackendProxy— an asyncio TCP proxy that starts the backend subprocesson first inbound connection and stops it after all connections have been
idle for a configurable timeout.
Key behaviours:
are done (no mid-stream kills for SSE)
asyncio.to_thread(no event-loop blocking)Files
tinyagentos/lazy_backend_proxy.py(~290 LOC)tests/test_lazy_backend_proxy.py(9 tests)Verification
cold-start, cold-start failure, idle timeout, sustained-request timer reset
asyncio.to_threadfor all subprocess waits_stop_subprocessalways terminates backend process, not just stop helperFuture work (out of scope for this PR)
sd_proxy,llama_proxy,whisper_proxyRefs #63
Summary by CodeRabbit
New Features
Tests