Report MLflow API database errors and stop returning PostgreSQL details

Eduardo Bonetrequested to merge
635049-mlflow-database-error-handling into master

What does this MR do and why?

The MLflow API now reports unexpected database errors to error tracking and no longer returns PostgreSQL error text. Today every ActiveRecord error becomes 400 INVALID_PARAMETER_VALUE with the raw database message, so statement timeouts never reach error tracking and constraint errors return the submitted row. Statement timeouts stay 400 so the MLflow client does not retry an expensive query, and transient errors such as a lost connection now return 503 so the client retries the write.

The change is behind the mlflow_database_error_handling feature flag (gitlab_com_derisk, default off). Rollback: disable the flag.

Design notes: error mapping and MLflow client retries

Responses for every endpoint under lib/api/ml/mlflow/:

Error Flag off (today) Flag on Error tracking
QueryCanceled (statement timeout) 400 INVALID_PARAMETER_VALUE, PostgreSQL message 400 BAD_REQUEST, Request timed out yes
Transient: ConnectionNotEstablished (and its subclasses such as ConnectionTimeoutError), ConnectionFailed, Deadlocked, SerializationFailure, LockWaitTimeout same as above 503 TEMPORARILY_UNAVAILABLE, Temporary database error yes
Column constraints: NotNullViolation, ValueTooLong, RangeError, StatementInvalid caused by PG::CheckViolation same as above 400 INVALID_PARAMETER_VALUE, A parameter value is missing, too long, or out of range no
RecordNotUnique same as above 400 RESOURCE_ALREADY_EXISTS no
Other ActiveRecordError same as above 400 BAD_REQUEST, Database error yes
RecordInvalid, RecordNotFound 400 INVALID_PARAMETER_VALUE, model message unchanged no

model-versions/create and the mlModelVersionCreate GraphQL mutation: when two requests create the same next version at the same time, Ml::CreateModelVersionService returned the RecordNotUnique text. With the flag on it returns Model version already exists. The endpoint keeps 400 INVALID_PARAMETER_VALUE, the same response as its validation path for a duplicate version (Version has already been taken).

Why these statuses:

  • The MLflow Python client retries 408, 429, 500, 502, 503 and 504 (_TRANSIENT_FAILURE_RESPONSE_CODES in mlflow/utils/request_utils.py), up to 7 times with backoff factor 2 (MLFLOW_HTTP_REQUEST_MAX_RETRIES, MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR).
  • Transient errors return 503 so that a training job's log_metric call retries during a database failover, instead of failing at once. These requests are cheap writes.
  • Statement timeouts and other database errors return 400, which the client does not retry, so the client does not repeat a query that already ran to the timeout. lib/api/merge_requests.rb and lib/api/search.rb answer QueryCanceled with 408, which MLflow retries, so that status is not used here. MLflow's own code for this case, INTERNAL_ERROR, maps to 500 (ERROR_CODE_TO_HTTP_STATUS in mlflow/exceptions.py), which the client also retries.
  • MLflow maps a plain 400 to BAD_REQUEST, its generic code. Server-side database errors use it instead of INVALID_PARAMETER_VALUE, which pointed at a parameter that was not the cause.
  • Column constraint errors come from the request. runs/log-batch uses insert_all, which skips model validations, so a param value over 250 characters, a tag value over 5,000 characters or a missing value reaches the database constraint. The same input on log-parameter fails the model validation. These errors are not tracked: they are client input, and the exception's DETAIL holds the submitted values.
  • RESOURCE_ALREADY_EXISTS for RecordNotUnique: mlflow.set_experiment catches that code when two processes create the same experiment (mlflow/tracking/fluent.py). experiments/create already maps the uniqueness validation to it, but a race past the validation hits the unique index on ml_experiments (project_id, name) and returned INVALID_PARAMETER_VALUE.
  • RecordNotFound keeps its current response. A change to 404 would be a status change, and no MLflow code path raises RecordNotFound today.

Implementation notes:

  • lib/api/ml/mlflow/entrypoint.rb keeps one rescue_from ActiveRecord::ActiveRecordError and delegates to render_active_record_error! in lib/api/ml/mlflow/api_helpers.rb. Grape uses the first matching handler of the mounted API before lib/api/api.rb, which is why rescue_from :all in api.rb never ran for these endpoints.
  • The flag check uses @project, not user_project. When the error came from the token or project lookup, user_project would run that lookup again inside the handler, and a second error would escape the handler as a 500. When @project is nil, the global flag state applies.
  • Tracked errors set env[API::Helpers::API_EXCEPTION_ENV], as handle_api_exception does, so api_json.log records exception.class and the normalized SQL. The user and remote IP of the error tracking event come from the application context that the before block in lib/api/api.rb pushes. The handler does not push them again, because current_user would repeat the token lookup.
  • lib/api/ml/mlflow_artifacts/entrypoint.rb has no such handler, and its errors already reach lib/api/api.rb. Not changed.

References

Screenshots or screen recordings

No UI change.

How to set up and validate locally

  1. Create a personal access token with the api scope and pick a project you own. Set TOKEN, PROJECT_ID and BASE=http://gdk.test:3000/api/v4/projects/$PROJECT_ID/ml/mlflow/api/2.0/mlflow.
  2. Create an experiment and a run: 
    curl -s -X POST -H "PRIVATE-TOKEN: $TOKEN" -H "Content-Type: application/json" -d '{"name":"db-errors"}' "$BASE/experiments/create"
curl -s -X POST -H "PRIVATE-TOKEN: $TOKEN" -H "Content-Type: application/json" -d '{"experiment_id":"<experiment_id>"}' "$BASE/runs/create"
  3. Log a param value of 251 characters, one more than the column allows: 
    VALUE=$(printf 'a%.0s' {1..251})
curl -s -X POST -H "PRIVATE-TOKEN: $TOKEN" -H "Content-Type: application/json" -d "{\"run_id\":\"<run_id>\",\"params\":[{\"key\":\"p\",\"value\":\"$VALUE\"}]}" "$BASE/runs/log-batch"
    Expect (flag off): 400 with INVALID_PARAMETER_VALUE and a PG::CheckViolation message that contains DETAIL: Failing row contains (...).
  4. In the rails console: Feature.enable(:mlflow_database_error_handling).
  5. Repeat step 3. Expect: 400 with {"error_code":"INVALID_PARAMETER_VALUE","message":"A parameter value is missing, too long, or out of range"}, and no new entry in log/exceptions_json.log.
Statement timeouts, transient errors and the other error classes

These are hard to cause on GDK. spec/requests/api/ml/mlflow/entrypoint_spec.rb stubs the lookup on runs/get (experiment tracking) and registered-models/get (model registry) to raise each error class in the table, and checks the status, the body, and whether Gitlab::ErrorTracking.track_exception is called, with the flag on and off. It also covers:

  • a statement timeout in the token lookup: the handler does not repeat the lookup and returns 400,
  • exception.class in the API log line,
  • the column constraint path through runs/log-batch without stubs, in spec/requests/api/ml/mlflow/entrypoint_column_constraint_spec.rb. That spec uses table deletion because the database error aborts the transaction that wraps a normal spec.

spec/services/ml/create_model_version_service_spec.rb covers the RecordNotUnique message with the flag on and off.

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