Enforce MLflow request limits on batch logging and tags

Eduardo Bonetrequested to merge
635039-limit-mlflow-batch-sizes into master

What does this MR do and why?

The MLflow-compatible API now rejects runs/log-batch requests over MLflow's own LogBatch limits (1,000 metrics, 100 params, 100 tags, 1,000 items in total), runs/create and experiments/create requests with more than 1,000 tags, and registered-models/create and model-versions/create requests with more than 100 tags. Today these endpoints accept any number of items, and each item is a row written in the request, so large requests hold a Puma worker and a database transaction for a long time. The limits are behind the mlflow_enforce_request_limits feature flag (gitlab_com_derisk, default off, project actor); rollback is to disable the flag.

Design notes
  • The checks live in lib/api/ml/mlflow/api_helpers.rb (check_log_batch_limits!, check_tags_limit!) and run first in each endpoint body, after authentication and before any record lookup or write.
  • Rejected requests get 400 with error_code: INVALID_PARAMETER_VALUE, the same shape as the other MLflow validation errors in this API (invalid_parameter!). Batch limits use the MLflow tracking server's message (A batch logging request can contain at most 1000 metrics. Got 1001 metrics. Please split up metrics across multiple requests and try again.); the create-request cap uses A request can contain at most 100 tags. Got 101 tags.
  • The batch limits are MLflow's (MAX_*_PER_BATCH). The create-request caps are separate GitLab constants (MAX_TAGS_PER_BULK_CREATE_REQUEST, MAX_TAGS_PER_CREATE_REQUEST), because MLflow does not limit tags on create requests.
  • log-batch read tags raw from params[:tags]; it was not declared. It is now declared as an array of key/value strings, like metrics and params, so Grape validates each tag. This part is not behind the flag (Grape validation cannot be gated per project). Behavior change outside the flag: a tag without key returned 400 INVALID_PARAMETER_VALUE from the database NOT NULL error, and now returns Grape's 400 {"error": "tags[0][key] is missing"}; a non-array tags returned 500, and now returns 400. The OpenAPI document gains the tags parameter.
  • Create-request caps: MLflow documents no count limit for CreateRun, CreateExperiment, CreateRegisteredModel or CreateModelVersion tags. runs/create and experiments/create insert all tags with one insert_all, so they allow 1,000 tags, MLflow's LogBatch total. registered-models/create and model-versions/create insert one row per tag, so they keep 100, MLflow's LogBatch tag limit. The MLflow client adds only a few system tags to create_run (mlflow.user, mlflow.source.*, mlflow.runName).
  • runs/create writes to ml_candidate_metadata, the table that the log-batch tag limit covers, and experiments/create writes to ml_experiment_metadata; both now have the cap.
  • The per-tag create! loop in Ml::CreateModelService#add_metadata and Ml::ModelVersion#add_metadata is unchanged; the cap alone bounds the transaction. Because each tag costs its own queries, the registered model and model version specs stub MAX_TAGS_PER_CREATE_REQUEST to 2 to stay under the spec query limiter; the runs/create and experiments/create specs use the real 1,000-tag cap.
  • A spec pins the flag actor: enabling the flag for another project does not apply the limits to this project.
  • No new per-user request-rate limit in this MR.
  • Follow-ups, not in this MR: bulk-insert the tags in the two add_metadata loops; Grape validation errors for malformed metrics, params and tags items return {"error": ...} without an MLflow error_code (already true for metrics and params).
  • Docs: the rollout issue tracks documenting the limits on the MLflow client compatibility page before global rollout.

MLflow client compatibility

The MLflow Python client already splits log_batch calls into requests within these limits, and the MLflow tracking server enforces the same limits, so MLflow clients are not affected. A production data check found no request over these create caps; the details are in the confidential rollout issue.

MLflow source references (v3.17.2)

References

Screenshots or screen recordings

No UI change.

How to set up and validate locally

  1. In a Rails console, enable the flag: Feature.enable(:mlflow_enforce_request_limits).

  2. Create a personal access token with the api scope for a user with the Developer role on any project, then set TOKEN=<your_access_token> and MLFLOW=http://gdk.test:3000/api/v4/projects/<project_id>/ml/mlflow/api/2.0/mlflow.

  3. Create an experiment and a run:

    curl --request POST --header "Authorization: Bearer $TOKEN" --header "Content-Type: application/json" \
  --url "$MLFLOW/experiments/create" --data '{"name": "limits"}'
curl --request POST --header "Authorization: Bearer $TOKEN" --header "Content-Type: application/json" \
  --url "$MLFLOW/runs/create" --data '{"experiment_id": "<experiment_id>"}'

  4. Build a batch with 1,001 metrics and send it:

    ruby -rjson -e 'puts({ run_id: ARGV[0], metrics: Array.new(1001) { |i| { key: "loss", value: 0.5, timestamp: 0, step: i } } }.to_json)' <run_id> > batch.json
curl --request POST --header "Authorization: Bearer $TOKEN" --header "Content-Type: application/json" \
  --url "$MLFLOW/runs/log-batch" --data @batch.json

    Expect: status 400 with {"error_code":"INVALID_PARAMETER_VALUE","message":"A batch logging request can contain at most 1000 metrics. Got 1001 metrics. Please split up metrics across multiple requests and try again."}.

  5. Change 1001 to 1000 and send again. Expect: {} with status 200.

  6. Run Feature.disable(:mlflow_enforce_request_limits) and send the 1,001-metric batch again. Expect: {} with status 200.

More scenarios
  • Tags on model creation, flag on: curl ... --url "$MLFLOW/registered-models/create" --data "$(ruby -rjson -e 'puts({ name: "m1", tags: Array.new(101) { |i| { key: "t#{i}", value: "v" } } }.to_json)')". Expect 400 with A request can contain at most 100 tags. Got 101 tags. With 100 tags, expect 200. model-versions/create (with name: "m1") behaves the same. runs/create and experiments/create behave the same at 1,000 and 1,001 tags.
  • MLflow Python client, flag on: MlflowClient().log_batch(run_id, metrics=[Metric("loss", 0.5, 0, i) for i in range(2500)]) succeeds, because the client sends three requests of at most 1,000 metrics.
  • Malformed tags in log-batch: "tags": [{"value": "v"}] returns 400 tags[0][key] is missing, with or without the flag.

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.

Database review

Not needed: the change validates request input before any query runs. It adds and changes no query, index or migration.

Edited by Eduardo Bonet

Merge request reports