diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index fbc46ac0ee3..75e2b00f538 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -1303,8 +1303,7 @@ class ProxyExtrasDBManager: def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool: """Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the migrations target. Idempotent and never raises; False when an index is still - missing or invalid, so the migration job reports it and gets rerun instead of - leaving the table unindexed until the next deploy.""" + missing or invalid, which the builder logs for the next index build to pick up.""" database_url: Final = os.environ.get("DATABASE_URL") if not database_url: return True @@ -1322,9 +1321,13 @@ class ProxyExtrasDBManager: build: Callable[[], bool] = build_request_log_indexes, ) -> bool: """The migration job's whole run: `setup_database`, then the request-log indexes, - built synchronously so the job exits only once they are in place. False when the - migrations failed or an index could not be built, so the Job is rerun.""" - return setup(use_migrate, use_v2_resolver) and build() + built synchronously so the job exits only once the build has finished. The result is + the migrations' alone, so an index that could not be built never fails an upgrade + whose migrations passed.""" + if not setup(use_migrate, use_v2_resolver): + return False + build() + return True @staticmethod def start_request_log_index_build(build: Callable[[], bool] = build_request_log_indexes) -> threading.Thread: diff --git a/migrations/run.py b/migrations/run.py index 94e7cc7c0f0..3620f73e3e2 100644 --- a/migrations/run.py +++ b/migrations/run.py @@ -5,9 +5,10 @@ recovery logic in `litellm_proxy_extras.ProxyExtrasDBManager.setup_database` (P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.), then builds the request-log indexes the migrations leave out (`litellm_proxy_extras.request_log_indexes`), waiting for them. The job exits -non-zero when an index could not be built so that it is rerun. A serving proxy -that runs the migrations itself builds the same indexes in the background once -it serves. +by the migrations' result, so an index that could not be built is logged and +left to the next run instead of failing an upgrade whose migrations passed. A +serving proxy that runs the migrations itself builds the same indexes in the +background once it serves. Env vars: DATABASE_URL required unless it can be assembled at diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index fe1c934e680..8907face226 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1587,9 +1587,7 @@ class TestBuildRequestLogIndexes: assert builds.get_nowait() == ("postgresql://u:p@primary:5432/db", "public") assert builds.empty() - def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_a_build_that_leaves_indexes_missing_is_reported(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db") assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False @@ -1628,8 +1626,8 @@ class TestStartRequestLogIndexBuild: class TestRunMigrationJob: - """`run_migration_job` is `setup_database` followed by the index build, each step's - result deciding whether the job reports success.""" + """`run_migration_job` is `setup_database` followed by the index build, the migrations' + result alone deciding whether the job reports success.""" @pytest.fixture def calls(self) -> SimpleQueue[tuple[str, bool, bool] | tuple[str]]: @@ -1681,13 +1679,13 @@ class TestRunMigrationJob: assert calls.get_nowait() == ("setup", True, True) assert calls.empty() - def test_the_job_fails_when_an_index_could_not_be_built( + def test_the_job_succeeds_by_the_migrations_when_an_index_could_not_be_built( self, calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]], setup: Callable[[bool], Callable[[bool, bool], bool]], build: Callable[[bool], Callable[[], bool]], ) -> None: - assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is False + assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is True assert (calls.get_nowait(), calls.get_nowait()) == (("setup", True, True), ("build",)) assert calls.empty() diff --git a/tests/proxy_migration_tests/test_request_log_indexes.py b/tests/proxy_migration_tests/test_request_log_indexes.py index 015492c7040..f8f34a39768 100644 --- a/tests/proxy_migration_tests/test_request_log_indexes.py +++ b/tests/proxy_migration_tests/test_request_log_indexes.py @@ -581,15 +581,14 @@ def test_a_process_holding_the_migration_lock_makes_the_build_wait_for_the_next_ @requires_db @RESOLVERS -def test_a_migration_job_that_could_not_build_the_indexes_reports_failure_and_succeeds_when_rerun( +def test_a_migration_job_that_could_not_build_the_indexes_succeeds_and_a_rerun_builds_them( scratch_database: str, use_v2_resolver: bool ) -> None: - """The migration job waits for the build and exits by run_migration_job's result; a job - that exits 0 with the indexes missing would leave the table unindexed until the next - deploy or until a serving proxy's background build gets to them.""" + """The migration job exits by the migrations' result, so an index another replica's lock kept + it from building never fails an upgrade, and the next run builds it.""" with psycopg.connect(scratch_database, autocommit=True) as other_replica: other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,)) - assert _migration_job(use_v2_resolver) is False + assert _migration_job(use_v2_resolver) is True _assert_every_ledger_row_is_finished(scratch_database) assert _index_validity(scratch_database, "litellm_call_id_idx") == {}