Repository navigation
LCORE-4610: e2e tests for quota management - #3135
Open
max-svistunov wants to merge 5 commits into
Open
max-svistunov wants to merge 5 commits into
max-svistunov wants to merge 5 commits into
Conversation
The quota scheduler and the token usage history had no direct tests: nothing failed if a limiter type mapped to the wrong subject, if the scheduler started without storage, or if the PostgreSQL statements took different parameters from the SQLite ones. Cover the subject mapping, the choice of statements, connection checks, the startup conditions and reconnection attempts of the scheduler, and the parameters of the PostgreSQL statements, which no test can run against a real server here. For the token usage history cover the running totals per user, provider and model, reconnecting after a closed connection, running without storage, and a write failing loudly when its table is gone.
LCORE-4606 showed that the renewal statements were never executed by any test: on SQLite a quota was renewed at every scheduler run. These tests run the real statements and scheduler functions on a real database. They pin that a quota is renewed only once its period has elapsed and once per period, for both the reset and the increase; that the time a limiter stamps on a new row is read correctly by SQLite; that a zero value skips its statement; that a renewal touches only the rows of its own subject; and that the scheduler loop creates the table, survives a failing limiter and reconnects after its connection died. With initial_quota and quota_increase both set, the increase runs first and stamps revoked_at, so the reset of the same run no longer matches and the balance is not what the quota_limiter module documents. That is pinned as a strict xfail rather than fixed here.
Add integration tests that send requests through the real endpoints with real limiters on SQLite, so the client's view and the database are compared. They cover /v1/query, /v1/streaming_query, /v1/responses (blocking and streamed) and /v1/infer: a turn is charged and the balance reported, an exhausted user is rejected with the 429 body, a shield-blocked streamed turn costs nothing, and /v1/infer is not enforced without a quota subject. They also cover the shared cluster balance, which limiter is named when both are exhausted, the token history, a limiter reconnecting after its connection died, and an unreachable database giving a 500 instead of a served turn, also when the error comes after the model answered.
The quota e2e steps need to read a streamed response and parse the end event of /v1/streaming_query, which the token counter steps and the LLM response steps each did with their own copy. Move the stream reader, the end event parser and DEFAULT_LLM_TIMEOUT to tests/e2e/utils/utils.py and use them from those step modules. Behaviour does not change.
Quota management had no e2e coverage. Add tests/e2e/features/quota.feature with its steps and configurations for library and server mode. The scenarios check that /v1/query, /v1/streaming_query and /v1/responses (blocking and streamed) report the balances and then reject the request with the quota error, that /v1/infer does the same when the quota subject is the user id, that the cluster balance is paid by every user and a rejected request costs nothing, and that the scheduler renews an exhausted balance while the service runs. The feature has the new @cfg_quota tag, added to the "other" shard of the e2e workflow and to the Makefile's tag expression, and to test_list.txt. An integration test loads the e2e configurations and pins the values the feature relies on, and renewal from a configuration file. The e2e scenario list and testing docs are updated.
Contributor
|
Warning Review limit reachedYou've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Next included review available in 31 minutes. View limit details
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. CodeRabbit Verify
|
4 of 20 tasks
This branch has not been deployed
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.
Description
Part 2 of 2 for LCORE-4610: the quota e2e feature. It is stacked on #3134, so this PR also shows that PR's three commits until it merges; please review only the last two commits here:
share the e2e stream and usage helpersandadd the quota e2e feature and its configurations. Tests, their configuration and CI only; nothing undersrc/changes.E2E (
tests/e2e/features/quota.feature,@cfg_quota, steps insteps/quota.py): withlightspeed-stack-quota.yaml(user quota of one token, cluster quota,/v1/inferenforced per user id) each ofquery,streaming_query,responses, streamedresponsesandinferis served once and then rejected with the exact quota error body; balances are checked against the response's own token counts; the cluster balance is shared across users and a rejected request costs nothing. Withlightspeed-stack-quota-renewal.yaml(empty quotas renewed every two seconds) a fresh user is rejected and then served, observed by bounded polling. Every scenario uses a fresh user id, so the persisted quota database never matters. The feature is added to theotherGitHub shard (both places in the workflow), the Makefile tag expression andtests/e2e/test_list.txt.The quota database of the e2e fixtures lives at
/tmp/quota.db, not under/tmp/data: nothing on the quota path creates a parent directory (sqlite3.connectdoes not), the image does not ship/tmp/data, and in theothershard it exists only becauseconversation-compaction.featurewrote transcripts there first. With the file directly under/tmpthe feature does not depend on which feature ran before it.The stream-reading and
streaming_queryend-event parsing the new steps need were private copies insteps/llm_query_response.pyandsteps/token_counters.py; they now live intests/e2e/utils/utils.py(read_streamed_response,parse_streaming_query_end_event, and theDEFAULT_LLM_TIMEOUTthata2a.pyalso copied), imported by all of them.tests/integration/test_quota_configuration.pygoes with the feature: quota handlers loaded from a YAML file (tests/configuration/lightspeed-stack-quota-sqlite.yaml, db path from the environment) build the limiters and the token history and renew that database; a period SQLite cannot use is rejected at load; both e2e quota fixtures load in both deployment modes with the values the feature relies on, so fixture drift fails here and not in the e2e shard.Docs:
quota.featureand the@cfg_quotatag are added todocs/testing/e2e_scenarios.mdanddocs/testing/e2e_testing.md, the new fixtures totests/e2e/configuration/README.md, and the new test modules to the README of each test directory.Type of change
pyproject.toml+uv.lock]requirements.*.txtfor Konflux]Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Run the configuration test:
Expected: 7 tests pass.
Run the e2e feature alone, in either mode:
What was run here: the configuration test passes; the full
tests/unitandtests/integrationsuites (withoutcontainer_lifecycle) pass; ruff, black, pylint and pyright report nothing new on the touched files;behave --dry-runbinds all 112 steps of the feature with no undefined step;make schemaleavesdocs/devel_doc/openapi.jsonunchanged.What was not run here: the e2e feature against a running stack (no Docker on the machine). Its first real runs are the
othershard of this PR's e2e workflow, in server and library mode, and the Konflux pipeline. Please also run the feature once on its own with the command above: the shard always has another feature running before it, and the standalone run is the one that proves the fixture carries no hidden precondition. The renewal scenario depends on the scheduler thread running in the deployed service (src/lightspeed_stack.pystarts it before uvicorn).The nightly providers workflow runs the whole test list, so it will run this feature against azure, google-vertex, watsonx and aws-bedrock. The streamed scenarios require the provider to report a positive token usage for a streamed turn (the existing token-counter steps accept zero there). A red
quota.featurein that job means the provider reports no usage on streams, which is quota not being enforced for that provider: a bug to file, not an assertion to relax.