Remove per-record queries from MLflow list endpoints

Eduardo Bonetrequested to merge
635042-preload-mlflow-list-associations into master

What does this MR do and why?

The MLflow list endpoints now run a fixed number of queries per request. Before, experiments/list, experiments/search, runs/search and registered-models/get and search loaded metadata and candidates one experiment, run, model or model version at a time, so one request could run thousands of queries. Behind the mlflow_bounded_list_responses flag (default off), latest_versions holds only the latest model version, as MLflow defines it, and experiments/list returns pages of at most 1,000 experiments, so the response size is bounded too.

Design notes, MLflow client compatibility, and what stays out

Preloads (no flag, rollback is a revert). They change only how many queries run, not the response.

  • Ml::Experiment.including_metadata and Ml::Candidate.including_metadata (new scopes, preload(:metadata)), used by experiments/list, experiments/search and runs/search.
  • I did not add :metadata to Ml::Candidate.including_relationships. Its other caller, Projects::Ml::ExperimentsController#show (HTML and CSV), does not render candidate metadata, and the CSV export is unpaginated, so it would load every metadata row of the experiment.
  • preload_registered_models in lib/api/ml/mlflow/api_helpers.rb preloads model metadata, and metadata and candidate of the versions the entity renders (latest_version with the flag, versions without). get, update and search use it.

Flag mlflow_bounded_list_responses (gitlab_com_derisk, actor: project). Rollout issue: https://gitlab.com/gitlab-org/gitlab/-/issues/635108. Rollback: disable the flag.

  • latest_versions: the MLflow REST reference defines it as "Collection of latest model versions for each stage" (RegisteredModel), and the MLflow SQL store keeps one version per stage. GitLab reports every version with current_stage: development, so the MLflow answer is one version. The entity uses the existing latest_version association (highest semantic version), the same version get-latest-versions returns. Without the flag it keeps returning every version.
  • MLflow Python client (checked in 2.19.0): get_registered_model and search_registered_models only return the entity. The client reads latest_versions only in docstring examples that print it. models:/<name>/latest resolves through get_latest_versions (POST registered-models/get-latest-versions), which this MR does not change.
  • Visible change with the flag: a script that reads latest_versions to list every version sees only the latest one. On master latest_versions[0] is the first version created; with the flag it is the latest, which is what the docs example in doc/user/project/ml/experiment_tracking/mlflow_client.md ("Loading a model version") expects.
  • experiments/list: MLflow 2.0 removed list_experiments; 2.x clients call experiments/search, which is already paginated. The MLflow 1.x ListExperiments request has max_results and page_token, and the response has next_page_token (List Experiments). The 1.x REST store docstring says the server picks the page size. mlflow.list_experiments() always sends max_results=1000 and follows next_page_token (checked in 1.20 to 1.28), so it still gets every experiment. The cap matches experiments/search and MLflow's own SEARCH_MAX_RESULTS_DEFAULT. With the flag the response adds next_page_token (documented on ListExperiment and in the OpenAPI document).
    • Deviation from the REST reference: it says an unspecified max_results returns every experiment. With the flag, GitLab returns 1,000 and a token, so a raw REST script or a MlflowClient().list_experiments() call that ignores the token sees the first 1,000.
    • Validation with the flag: a page_token that this endpoint did not issue (not a forward cursor with one id in the bigint range) and a max_results below 1 return MLflow's INVALID_PARAMETER_VALUE 400, as MLflow's own server does. A max_results above 1,000 is capped at 1,000, as the REST reference describes.
    • max_results is now declared as an integer on list, so even with the flag off ?max_results=abc returns 400 instead of 200 (the value was ignored before).
    • Follow-up, not changed here: experiments/search, runs/search and registered-models/search still return 500 for an undecodable page_token and a database error for a negative max_results.
    • Follow-up, not changed here: an index on ml_experiments (project_id, id) would cut the flag-on experiments/list first page from about 18,700 to 516 buffers on projects with older experiments (see query 1). Tracked in https://gitlab.com/gitlab-org/gitlab/-/work_items/635210.
  • update is flagged too because it renders the same entity, so latest_versions means the same thing on get, update and search.
  • Docs: history entries and flag notes in doc/user/project/ml/experiment_tracking/mlflow_client.md, in "Fetching a model" (latest_versions) and in a new "Listing experiments" topic (experiments/list pages).

mlflow-artifacts/artifacts list stays unbounded. The artifacts proxy ListArtifacts contract (mlflow_artifacts.proto) has only path in the request and files in the response, with no page token. HttpArtifactRepository.list_artifacts reads every entry, and download_artifacts downloads what the list returns, so a truncated list would silently drop files from downloads. The endpoint already runs a constant number of queries (one for the files); its cost grows with files per version, so the bound belongs on the upload side. I noted this as a follow-up on the task.

Sibling MR. The metric-order fix for https://gitlab.com/gitlab-org/gitlab/-/work_items/635038 also edits app/models/ml/candidate.rb (order_by_metric). This MR adds one scope line next to including_relationships. Whoever lands second rebases.

References

Screenshots or screen recordings

No UI change.

How to set up and validate locally

Use a project you own on GDK (<id> below) and a personal access token with the api scope (export TOKEN=...). Set export API="http://gdk.test:3000/api/v4/projects/<id>/ml/mlflow/api/2.0/mlflow".

  1. Create a model with three versions:

    curl -s -H "Authorization: Bearer $TOKEN" -X POST "$API/registered-models/create" -d name=my-model
for v in 1.0.0 3.0.0 2.0.0; do curl -s -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/json' -X POST "$API/model-versions/create" -d "{\"name\":\"my-model\",\"tags\":[{\"key\":\"gitlab.version\",\"value\":\"$v\"}]}"; done

  2. Get the model: curl -s -H "Authorization: Bearer $TOKEN" "$API/registered-models/get?name=my-model" | jq '[.registered_model.latest_versions[].aliases[0]]'. Expect: ["1.0.0", "3.0.0", "2.0.0"] (flag off, every version).

  3. Create three experiments: for n in a b c; do curl -s -H "Authorization: Bearer $TOKEN" -X POST "$API/experiments/create" -d name=exp-$n; done.

  4. List them: curl -s -H "Authorization: Bearer $TOKEN" "$API/experiments/list?max_results=2" | jq '{count: (.experiments | length), next_page_token}'. Expect: every experiment of the project and no next_page_token (flag off ignores max_results).

  5. Enable the flag in the rails console: Feature.enable(:mlflow_bounded_list_responses).

  6. Repeat step 2. Expect: ["3.0.0"].

  7. Repeat step 4. Expect: count is 2 and next_page_token is set. Pass it as &page_token=<token>. Expect: the remaining experiments and next_page_token: null.

Query counts, specs, and query plans for database review

Specs. A new shared example, MLflow|a request with constant queries, takes the control on the second request (the first one runs one-off queries that hide an N+1) and then adds records with the same associations. It covers experiments/list (flag on and off), experiments/search, runs/search, registered-models/get and search (flag on and off), and registered-models/update (flag off). Each example checks that the response renders more records after the additions. Every example except registered-models/get with the flag on fails without the preload (checked by removing it). With the flag on, get renders one version, so its example adds versions and tags and checks that the query count stays the same; it passes with or without the preload. update with the flag on renders one version for the same reason and has no query-count example. Run them with bin/rspec spec/requests/api/ml/mlflow/ -e 'does not run queries per record'.

Queries. Each preload replaces one WHERE <fk> = $1 query per record with one IN (...) query per page. All of them use an existing index on the foreign key. Plans are from Database Lab (gitlab-production-main), cold clone first, then the warm execution time. Ids are replaced by placeholders.

  1. experiments/list page (new for this endpoint, flag on). The project with the most experiments (about 13,000), first page and a later page:

    SELECT "ml_experiments".* FROM "ml_experiments" WHERE "ml_experiments"."project_id" = $1 ORDER BY "ml_experiments"."id" DESC LIMIT 1001
    Limit  (cost=0.29..172.39 rows=1001 width=97) (actual time=0.856..17.307 rows=1001 loops=1)
  Buffers: shared hit=977 read=66
  ->  Index Scan Backward using ml_experiments_pkey on ml_experiments  (cost=0.29..2200.43 rows=12797 width=97) (actual time=0.855..17.232 rows=1001 loops=1)
        Filter: (project_id = $1)
        Rows Removed by Filter: 69
Planning Time: 7.422 ms
Execution Time: 17.375 ms (warm 0.386 ms)

    Next page (AND "ml_experiments"."id" < $2): Index Scan Backward using ml_experiments_pkey, Index Cond: (id < $2), 1,006 shared hits, 0.469 ms. A project with three experiments uses index_ml_experiments_on_project_id_and_iid and a sort, 0.038 ms. The table has about 30,000 rows. A project whose experiments are older (about 2,900 experiments) still scans ml_experiments_pkey backward: the first page removes 21,771 rows and reads 18,700 buffers, 86 ms cold, 4.6 ms warm. An index on (project_id, id) brings that page to 516 buffers; it is a follow-up (https://gitlab.com/gitlab-org/gitlab/-/work_items/635210), not part of this MR.

  2. Experiment metadata (experiments/list, experiments/search):

    SELECT "ml_experiment_metadata".* FROM "ml_experiment_metadata" WHERE "ml_experiment_metadata"."experiment_id" IN ($1, ...)
    Index Scan using index_ml_experiment_metadata_on_experiment_id_and_name on ml_experiment_metadata  (cost=0.28..140.37 rows=1418 width=83) (actual time=0.069..0.615 rows=1352 loops=1)
  Index Cond: (experiment_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=150 read=24
Planning Time: 0.235 ms
Execution Time: 0.686 ms (warm 0.276 ms)

    With 1,000 ids (a full page of the largest project): 2 buffers read, 1.6 ms cold, planning 22.8 ms cold. With the flag off, experiments/list sends every id of the project in one list: 13,144 ids plan in 117 ms cold and 93 ms warm, and execute in 10 ms cold.

  3. Candidate metadata (runs/search), the 1,000 newest runs (max_results=1000; the default page is 200) of the experiment with the most metadata per page:

    SELECT "ml_candidate_metadata".* FROM "ml_candidate_metadata" WHERE "ml_candidate_metadata"."candidate_id" IN ($1, ...)
    Index Scan using index_ml_candidate_metadata_on_candidate_id_and_name on ml_candidate_metadata  (cost=0.43..6964.25 rows=7448 width=76) (actual time=2.597..234.456 rows=17648 loops=1)
  Index Cond: (candidate_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=4345 read=1107
  I/O Timings: shared read=227.142
Planning Time: 19.910 ms
Execution Time: 235.276 ms (warm 5.715 ms)

    A second experiment with a similar page (16,464 rows) took 1,354 ms cold (2,365 buffers read) and 6.1 ms warm. No flag covers this preload, so it runs on GitLab.com right after the merge. Master loads the same rows with one query per run (1,000 queries for this page).

  4. Model metadata (registered-models/get, update, search), every model of the project with the most model metadata (8 models, 32 rows; the table has about 100 rows):

    SELECT "ml_model_metadata".* FROM "ml_model_metadata" WHERE "ml_model_metadata"."model_id" IN ($1, ...)
    Index Scan using unique_index_ml_model_metadata_name on ml_model_metadata  (cost=0.14..4.57 rows=35 width=65) (actual time=5.342..6.193 rows=32 loops=1)
  Index Cond: (model_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=3 read=2
  I/O Timings: shared read=6.167
Planning Time: 3.483 ms
Execution Time: 6.213 ms (warm 0.019 ms)

  5. Model versions of a page of models (registered-models/search, flag off), same 1,000 models:

    SELECT "ml_model_versions".* FROM "ml_model_versions" WHERE "ml_model_versions"."model_id" IN ($1, ...)
    Index Scan using index_ml_model_versions_on_created_at_on_model_id on ml_model_versions  (cost=0.29..819.06 rows=6296 width=162) (actual time=8.692..233.415 rows=7048 loops=1)
  Index Cond: (model_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=6717 read=302 dirtied=3
  I/O Timings: shared read=229.369
Planning Time: 36.643 ms
Execution Time: 233.785 ms (warm 2.834 ms)

  6. Version metadata and candidates of the rendered versions. Flag off, search: the 7,048 versions of the page above. Flag off, get: the model with the most versions (3,161). Flag on: one latest version per model (1,000).

    SELECT "ml_model_version_metadata".* FROM "ml_model_version_metadata" WHERE "ml_model_version_metadata"."model_version_id" IN ($1, ...)
SELECT "ml_candidates".* FROM "ml_candidates" WHERE "ml_candidates"."model_version_id" IN ($1, ...)
    Index Scan using unique_index_ml_model_version_metadata_name on ml_model_version_metadata  (cost=0.29..521.29 rows=7104 width=64) (actual time=4.903..29.577 rows=7048 loops=1)
  Index Cond: (model_version_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=4110 read=203
Planning Time: 120.085 ms
Execution Time: 29.920 ms (warm 3.295 ms)

Index Scan using index_ml_candidates_on_model_version_id on ml_candidates  (cost=0.43..4580.91 rows=7048 width=141) (actual time=7.334..1507.065 rows=7048 loops=1)
  Index Cond: (model_version_id = ANY ('{...}'::bigint[]))
  Buffers: shared hit=3946 read=1932 dirtied=8
  I/O Timings: shared read=1493.889
Planning Time: 15.188 ms
Execution Time: 1507.744 ms (warm 5.794 ms)
    • Flag off, search, 7,048 ids: version metadata plans in 118 to 126 ms; candidates take 1.5 s cold (1,932 buffers read) and 5.8 ms warm. An earlier run on a warmer cache took 183 ms.
    • Flag off, get, 3,161 ids: version metadata 0.8 ms (planning 48.6 ms); candidates take 2.3 s cold (2,730 buffers read) and 3.1 ms warm. An earlier run on a warmer cache took 21.6 ms.
    • Flag on, 1,000 latest versions: version metadata 880 shared hits, 0.7 ms, planning 14 to 18 ms; candidates 941 shared hits, 0.6 ms warm.

Flag-off cost of the new preloads. With the flag off (the default after merge), experiments/list, registered-models/get and registered-models/search send one IN list with every id of an unbounded set. On the largest cases above, planning reaches 93 to 126 ms and a cold read reaches 1.5 to 2.3 s for one query. Each query is still cheaper than master, which reads the same rows with one query per record (13,144 queries for that experiments/list, 3,161 candidate queries for that get). The flag bounds both: one version per model and 1,000 experiments per page.

The latest_version query already runs on master for search, through Projects::Ml::ModelFinder, and does not change there. With the flag on, get and update also run it for the one model they render, which they did not do on master. The model with the most versions (3,161):

SELECT DISTINCT ON (model_id) * FROM "ml_model_versions" WHERE "ml_model_versions"."model_id" = $1 ORDER BY "ml_model_versions"."model_id" DESC, "ml_model_versions"."semver_major" DESC, "ml_model_versions"."semver_minor" DESC, "ml_model_versions"."semver_patch" DESC
Limit  (cost=642.35..642.36 rows=1 width=164) (actual time=131.057..131.059 rows=1 loops=1)
  Buffers: shared hit=744 read=205
  I/O Timings: read=129.191
  ->  Sort  (cost=642.35..649.99 rows=3055 width=164) (actual time=131.056..131.057 rows=1 loops=1)
        Sort Key: semver_major DESC, semver_minor DESC, semver_patch DESC
        Sort Method: top-N heapsort  Memory: 25kB
        Buffers: shared hit=744 read=205
        ->  Index Scan using unique_ml_model_versions_on_model_id_and_id on ml_model_versions  (cost=0.29..627.08 rows=3055 width=164) (actual time=2.288..130.326 rows=3161 loops=1)
              Index Cond: (model_id = $1)
              Buffers: shared hit=741 read=205
Planning Time: 0.985 ms
Execution Time: 131.117 ms (warm 1.498 ms)

No index covers the semver_* order, so the query reads every version of the model and keeps the highest one. The flag-off get reads the same version rows, and then the metadata and candidates of all 3,161 versions (2.3 s cold for the candidates, see query 6).

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Eduardo Bonet

Merge request reports