feat(router): add deterministic lifecycle (close()/ aclose() + context managers) #14551

grahama1970 · 2025-09-14T14:11:47Z

Title: Router Lifecycle: Deterministic close()/aclose() + Context Managers (Tests Only)

Summary

Adds a minimal, deterministic lifecycle API to litellm.Router:
- Router.close() and Router.aclose() — idempotent teardown that unhooks Router-managed callbacks.
- Sync and async context managers (__enter__/__exit__, __aenter__/__aexit__) that call close/aclose.
Scope is intentionally tight: implementation + targeted tests only. No repo-wide Docker/CI or docs changes in this PR.

Motivation

Short‑lived scripts and tests can leave Router-managed globals/callbacks attached, leading to lingering state across runs and flakiness. A deterministic, idempotent lifecycle:
- Ensures clean teardown without requiring process exit.
- Prevents double-close errors by design.
- Improves test isolation and makes resource cleanup explicit.

Behavior & API

Router.close():
- No-op if already closed (_closed flag).
- Calls existing discard() to unhook callbacks from global managers.
- Swallows exceptions from discard() to keep teardown resilient.
Router.aclose():
- Async variant delegating to close() for symmetry.
Context managers:
- with Router(...) as r: calls close() on exit.
- async with Router(...) as r: calls aclose() on exit.

Backwards Compatibility

No breaking changes to Router construction or routing semantics.
New methods are additive. Existing codepaths are untouched unless a caller opts into the new API.

Performance & Threading

Negligible overhead: a _closed boolean flag and a call into discard() during teardown.
No persistent threads or background loops introduced. This change reduces the chance of lingering state between tests.

Security

Lifecycle only. No changes to auth, network surfaces, or request parsing.

Files Touched (and why)

litellm/router.py — Adds _closed flag, close()/aclose(), and sync/async context manager methods. Uses existing discard() for unhooking.
tests/test_litellm/test_router.py — Adds tests for:
- Idempotent close()
- Sync context manager closes on exit
- Async context manager closes on exit

Developer Feedback Addressed (from earlier review)

“Why are Docker files modified?” — Reverted. This PR no longer changes repo Dockerfiles or compose. All local/dev Docker remains under local/ in the fork and is excluded from the PR.
Out-of-scope changes (e.g., http_parsing_utils.py, openai_endpoint_utils.py, organization_endpoints.py, files endpoints guards) — Removed from this PR. If needed, we’ll submit targeted follow-ups with focused tests.

Testing

Validated locally by running only tests/test_litellm/test_router.py in a clean Python container.
Result: 27 passed, 0 failed (Python 3.12).
Exact command is included in a PR comment for reproducibility.

Reviewer Notes

The PR is intentionally small and reversible: a few lines in Router plus targeted tests. No docs included to keep scope tight; happy to add a brief “Lifecycle” section in Router docs if desired.
Fork CI may show red checks unrelated to these changes (matrix jobs needing org secrets). The tests‑only run above validates the changed areas; maintainers can “Approve and run” to execute the full matrix with org secrets.

Examples

import litellm

# Sync context
with litellm.Router(model_list=[{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo"}}]) as r:
    ...  # use r
# here r.close() has been called

# Async context
async def main():
    async with litellm.Router(model_list=[{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo"}}]) as r:
        ...
    # r.aclose() has been called

# Manual close (idempotent)
r = litellm.Router(model_list=[{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo"}}])
r.close()
r.close()  # safe

Changelog

Added: Router lifecycle API (close/aclose + context managers) with tests.
No docs/CI/Docker changes in this PR.

vercel · 2025-09-14T14:11:52Z

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Preview	Comments	Updated (UTC)
litellm	Error			Sep 17, 2025 5:05pm

krrishdholakia · 2025-09-16T04:09:37Z

docker-compose.yml

@grahama1970 modifying our docker-compose seems out of scope for this PR.

Can this be reverted please? @grahama1970

Thanks for flagging. Reverted — no repo Docker/compose changes remain. Scope is limited to router lifecycle + tests.

krrishdholakia

please revert the docker-compose changes

krrishdholakia · 2025-09-16T04:10:32Z

litellm/proxy/common_utils/http_parsing_utils.py


                    try:
-                        parsed_body = json.loads(body_str)
+                        parsed_body = json.loads(body_str) if body_str else {}


why change this? the change seems like it would be harder to catch invalid json

Good catch — reverted. No parsing changes in this PR.

krrishdholakia · 2025-09-16T04:10:56Z

litellm/proxy/common_utils/openai_endpoint_utils.py

is this related to this pr?

Reverted — unrelated to router lifecycle.

krrishdholakia · 2025-09-16T04:11:21Z

litellm/proxy/management_endpoints/organization_endpoints.py

same comment - why does this file exist?

Reverted — unrelated to this PR; will handle separately if needed.

krrishdholakia · 2025-09-16T04:11:38Z

litellm/proxy/openai_files_endpoints/files_endpoints.py

i assume this is to fix the ci / cd?

yes, I'll back to this PR in 48 hours

Reverted — unrelated to this PR; will handle separately if needed.

…t managers)\n\n- Add Router.close()/aclose() with idempotent shutdown flag and managed thread joins\n- Add sync/async context manager support\n- Hook into existing discard() to unhook callbacks\n- Tests for idempotency and context managers

…se/aclose + context managers only

…agers); tests for idempotent close + sync/async context; fix(org): restore transform_organization_update_request alias; chore(make): ci-parity-local lints only changed files and add preflight-local

…o avoid TypeError in files endpoint; chore(makefile): add smoke-proxy-files target and clearer progress output

…purpose to str; prevent MagicMock TypeError in CI

…efile targets inside CI container; add Dockerfile.ci and docker-compose.yml

…docker only

grahama1970 · 2025-09-17T18:12:20Z

Scope narrowed to exactly the router lifecycle + tests. Summary for reviewers:

What changed (now)

Files touched: 2 total
1. litellm/router.py — adds close()/aclose() and sync/async context managers (idempotent)
2. tests/test_litellm/test_router.py — adds lifecycle tests (idempotent close, sync/async context)
Removed Docker bits from this PR: Dockerfile.ci deleted in HEAD (commit b68df6f). No repo docker-compose edits. All local/dev Docker remains only under local/ in the fork and is excluded from this PR.

Validation (tests-only in container as repo user)

Command used (from the repo root on host):
docker compose --env-file local/docker/.env.user -f local/docker/compose.ci.yml run --rm
-v /home/graham/workspace/experiments/litellm_router_wt:/wt
-w /wt
ci-python bash -lc '
set -euo pipefail &&
python -m venv /tmp/poetry_venv &&
/tmp/poetry_venv/bin/pip install --upgrade pip &&
/tmp/poetry_venv/bin/pip install poetry &&
export PATH=/tmp/poetry_venv/bin:$PATH &&
poetry install --with dev,proxy-dev --extras "proxy semantic-router" &&
poetry run pip install "pytest-retry==1.6.3" pytest-xdist pytest-timeout
"google-genai==1.22.0" "google-cloud-aiplatform>=1.38"
"fastapi-offline==1.7.3" &&
(cd enterprise && python -m pip install -e .) &&
poetry run pytest -vv -n 0 tests/test_litellm/test_router.py
--timeout=120 --timeout-method=thread'
Result: 27 passed, 0 failed (2 warnings), Python 3.12 — no root-owned files created (container runs as ${UID}:${GID}).

Notes

If Files changed still shows Dockerfile.ci, that’s likely GitHub cache catching up. Latest HEAD (b68df6f) removes it; only the two files above should be in scope.

Request

Please review just these two files and, if acceptable, “Approve and run” CI (full matrix needs org secrets). Happy to squash after review or follow preferred commit style.

Reviewer comments addressed

"Modifying our docker-compose seems out of scope" — Agreed and reverted. No repo docker-compose or Dockerfile changes remain in this PR. Dockerfile.ci added earlier is removed in b68df6f; net PR scope is router.py + tests only.
http_parsing_utils.py: Concern about making JSON parsing harder to catch invalid JSON — Reverted to upstream behavior; no changes to parsing remain in this PR.
openai_endpoint_utils.py and organization_endpoints.py: Not related to router lifecycle — Reverted/removed from PR scope.
files_endpoints.py: CI-only guard change — Removed from PR; we will raise a separate targeted fix if still needed under CI conditions.

vercel bot deployed to Preview September 14, 2025 14:16 View deployment

vercel bot deployed to Preview September 14, 2025 14:35 View deployment

vercel bot deployed to Preview September 14, 2025 15:02 View deployment

vercel bot deployed to Preview September 14, 2025 15:20 View deployment

krrishdholakia reviewed Sep 16, 2025

View reviewed changes

krrishdholakia requested changes Sep 16, 2025

View reviewed changes

krrishdholakia reviewed Sep 16, 2025

View reviewed changes

grahama1970 added 7 commits September 17, 2025 11:52

refactor(router): simplify lifecycle API — no options, idempotent clo…

a151793

…se/aclose + context managers only

feat(router): add deterministic lifecycle (close/aclose + context man…

39ebe63

…agers); tests for idempotent close + sync/async context; fix(org): restore transform_organization_update_request alias; chore(make): ci-parity-local lints only changed files and add preflight-local

fix(proxy): harden request body parsing for multipart/mocked bodies t…

0b82938

…o avoid TypeError in files endpoint; chore(makefile): add smoke-proxy-files target and clearer progress output

fix(proxy/files): guard multipart/custom provider parsing and coerce …

26575f7

…purpose to str; prevent MagicMock TypeError in CI

chore(makefile/docker): add docker-up/down and docker-exec to run Mak…

85a5d1d

…efile targets inside CI container; add Dockerfile.ci and docker-compose.yml

scope(router): restrict PR to router lifecycle and tests only

8c65aec

grahama1970 force-pushed the feat/router-lifecycle-clean branch from 58ac17c to 8c65aec Compare September 17, 2025 15:54

vercel bot had a problem deploying to Preview September 17, 2025 15:55 Failure

chore(ci): drop Dockerfile.ci from PR; keep local docker under local/…

b68df6f

…docker only

vercel bot had a problem deploying to Preview September 17, 2025 17:05 Failure

grahama1970 changed the title ~~feat(router): deterministic shutdown lifecycle (close/aclose + context manager)~~ feat(router): add deterministic lifecycle (close()/ aclose() + context managers) Sep 17, 2025

grahama1970 requested a review from krrishdholakia September 17, 2025 18:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

feat(router): add deterministic lifecycle (close()/ aclose() + context managers) #14551

feat(router): add deterministic lifecycle (close()/ aclose() + context managers) #14551

grahama1970 commented Sep 14, 2025 •

edited

Loading

Uh oh!

vercel bot commented Sep 14, 2025 •

edited

Loading

Uh oh!

krrishdholakia Sep 16, 2025

Uh oh!

grahama1970 Sep 17, 2025

Uh oh!

krrishdholakia left a comment

Uh oh!

krrishdholakia Sep 16, 2025

Uh oh!

grahama1970 Sep 17, 2025

Uh oh!

krrishdholakia Sep 16, 2025

Uh oh!

grahama1970 Sep 17, 2025

Uh oh!

krrishdholakia Sep 16, 2025

Uh oh!

grahama1970 Sep 17, 2025

Uh oh!

krrishdholakia Sep 16, 2025

Uh oh!

grahama1970 Sep 16, 2025

Uh oh!

grahama1970 Sep 17, 2025

Uh oh!

grahama1970 commented Sep 17, 2025

Uh oh!

Uh oh!

Uh oh!

feat(router): add deterministic lifecycle (close()/ aclose() + context managers) #14551

Are you sure you want to change the base?

feat(router): add deterministic lifecycle (close()/ aclose() + context managers) #14551

Conversation

grahama1970 commented Sep 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vercel bot commented Sep 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

krrishdholakia left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

grahama1970 commented Sep 17, 2025

Uh oh!

Uh oh!

grahama1970 commented Sep 14, 2025 •

edited

Loading

vercel bot commented Sep 14, 2025 •

edited

Loading