Delete ML experiments and models in batches in the background

Eduardo Bonetrequested to merge
635041-delete-ml-experiments-in-batches into 635041-ml-experiment-deletion-worker

What does this MR do and why?

Deleting an ML experiment or model now hides the experiment at once and leaves its runs, with their metrics, parameters and metadata, to the background worker from !264622, which deletes them in statements of bounded size. Today the request runs one DELETE and PostgreSQL cascades through every run of the experiment in that statement, so the work grows with the experiment and a large one cannot finish within the statement timeout. The new path is behind the ml_delete_experiments_in_batches flag (default off).

Depends on !264622.

Design notes

Hiding without a schema change. ml_experiments.deleted_on already exists and the MLflow entity already maps it to lifecycle_stage: deleted, but no code wrote it. Ml::Experiment#mark_for_deletion sets it and renames the experiment to deletion_scheduled-<id>-<16 random hex characters>, so a new experiment or model can take the name right away and no user-chosen name can collide with the placeholder. Ml::Experiment.by_project_id_and_iid, .by_project_id_and_name, Projects::Ml::ExperimentFinder and the GraphQL mlExperiment(id:) resolver skip experiments with deleted_on set. With the flag on, the run lookups (Ml::Candidate.with_project_id_and_eid, _iid, _id) also return nothing for runs of a marked experiment, so MLflow run reads and writes, run artifacts, the run page and model version promotion answer 404. That costs one primary key read on ml_experiments when a run is found; callers usually load the experiment anyway.

No job per deletion. Marking enqueues nothing. The cron and LimitedCapacity::Worker from !264622 pick marked experiments up within 5 minutes and delete them in bounded slices: at most 10,000 child rows or 100 runs per statement, and 100,000 rows or 2 minutes per job. Those workers do not check the flag: with the flag off nothing is marked, and after a rollback the experiments already marked still get deleted.

Experiment deletion (Ml::DestroyExperimentService, used by MLflow experiments/delete and the experiment page): with the flag on, it marks the experiment instead of destroying it. An experiment that belongs to a model is still refused with the same message.

Model deletion (Ml::DestroyModelService, used by MLflow registered-models/delete and GraphQL mlModelDelete and mlModelDestroy): with the flag on, it validates the experiment's new state first, then marks the packages for destruction (unchanged), then in one transaction unlinks the versions' runs (model_version_id = NULL, in batches of 1,000 plucked version ids), marks the default experiment and deletes the model. The cascade then deletes only the versions and their metadata, and the worker deletes the runs. Errors from the experiment are returned in the response. The version ids are plucked (at most 1,000 per batch) because the IN (subquery) form plans a merge join over index_ml_candidates_on_model_version_id.

Known limit: the model request still does work proportional to the number of versions, which has no cap: about 1 s cold and 12 MB of WAL per 1,000 versions, about 3 s cold for the model with the most versions on GitLab.com (7.8 s today). Moving the version unlinking to the worker needs a column on ml_models; see the follow-ups.

Not changed: model version deletion (Ml::DestroyModelVersionService, mlModelVersionDelete) and MLflow runs/delete still cascade one run in the request.

Related MRs and merge order (https://gitlab.com/gitlab-org/gitlab/-/work_items/635041):

  1. !264622 adds the deletion worker and cron. This MR targets its branch.
  2. This MR merges after it. Both MRs squash on merge, so a retarget or a plain rebase onto master conflicts. After that MR merges, replay only this MR's commits onto master (today: git rebase --onto origin/master dc7ed3f1ea22) and force-push.
  3. !264621 (merged) routes the web delete actions through the services. It is independent and can merge at any point; once both are in, the web actions also hide experiments when the flag is on.

Follow-ups, not in this MR:

  • Project deletion still cascades through the ten ml_* tables in one statement: https://gitlab.com/gitlab-org/gitlab/-/work_items/635215.
  • Mark the model and move the version unlinking to the worker (needs a column on ml_models): https://gitlab.com/gitlab-org/gitlab/-/work_items/635216.
  • MLflow runs/delete and model version deletion cascade one run in one statement; a run with very many metrics needs the same child batching.
  • count_ml_experiments_metric, count_projects_with_ml_experiments_metric and the compliance check in ee/lib/compliance_management/compliance_requirements/project_fields.rb still count marked experiments until the worker deletes them. Each is a one-scope change, but each brings another review domain for a short window, so they are left for later.

Sibling MRs that touch the same files: !263896 (merged) edits app/models/ml/experiment.rb, lib/api/ml/mlflow/api_helpers.rb and spec/requests/api/ml/mlflow/experiments_spec.rb, and !263900 (merged) edits app/models/ml/candidate.rb and spec/models/ml/candidate_spec.rb. The changes are on different lines and do not depend on each other; whichever merges second rebases.

References

Screenshots or screen recordings

No visual change.

How to set up and validate locally

  1. In the rails console, enable the flag and create an experiment with runs and metrics:

    Feature.enable(:ml_delete_experiments_in_batches)
project = Project.find_by_full_path('gitlab-org/gitlab-test'); user = User.find_by_username('root')
exp = Ml::Experiment.create!(project: project, user: user, name: 'delete-me')
3.times { |i| Ml::Candidate.create!(experiment: exp, project: project, user: user, name: "run#{i}").metrics.create!(name: 'loss', value: i, step: 0, tracked_at: 0, project: project) }

  2. Delete it through the MLflow API (the experiment page uses the service only after !264621 (merged)): curl -X POST -H "Authorization: Bearer <token>" "http://gdk.test:3000/api/v4/projects/<project_id>/ml/mlflow/api/2.0/mlflow/experiments/delete" -d experiment_id=<exp.iid>.

  3. Expect: {}. delete-me is no longer listed at http://gdk.test:3000/gitlab-org/gitlab-test/-/ml/experiments, and experiments/get for the same iid returns 404. In the console, exp.reload.deleted_on is set and the name starts with deletion_scheduled-.

  4. Run Ml::ScheduleExperimentDeletionCronWorker.new.perform (or wait up to 5 minutes for the cron), then Ml::Experiment.find_by(id: exp.id). Expect: nil, and Ml::Candidate.where(experiment_id: exp.id).count is 0.

  5. Create a model at http://gdk.test:3000/gitlab-org/gitlab-test/-/ml/models, then delete it from its page. Expect: the model is gone at once, and its [model]<name> experiment is removed by the next worker run.

Other paths and flag-off check
  • runs/get for one of the deleted experiment's runs returns 404 until the worker deletes it.
  • Deleting a model's experiment through the API returns 400 Cannot delete an experiment associated to a model, as before.
  • Feature.disable(:ml_delete_experiments_in_batches): the same steps delete the experiment in the request, and nothing is marked.
Database review: queries and plans (Database Lab, gitlab-production-main)

Ids and values are placeholders. Every write ran in a rolled-back transaction. "Cold" means pages read from disk on the clone; a worker never reads the same batch twice, so the batch size is chosen on cold numbers. Table sizes: ml_candidate_metrics 7.5 GB, ml_candidate_metadata 1.2 GB, ml_candidate_params 552 MB, ml_candidates 403 MB, ml_experiments 11 MB. No migration.

Before (current code). The run cascade and the worker statements (pick, claim, batch boundaries, child and run deletes, WAL per batch) are in !264622. DELETE FROM ml_models WHERE id = <model_id> on the model with the most versions (its runs have no metrics), cold:

Trigger for constraint fk_e86e0bfa5a on ml_model_versions: time=1423.563 calls=3161
Trigger for constraint fk_rails_5b67cc9107 on ml_model_versions: time=56.737 calls=3161
Trigger for constraint fk_ml_candidate_metrics_on_candidate_id on ml_candidates: time=1329.795 calls=3161
Trigger for constraint fk_ml_candidate_params_on_candidate_id on ml_candidates: time=374.271 calls=3161
Trigger for constraint fk_rails_5117dddf22 on ml_candidates: time=4406.612 calls=3161
Execution Time: 7772.683 ms

1. Request: mark the experiment

UPDATE "ml_experiments" SET "updated_at" = '<now>', "name" = 'deletion_scheduled-<id>-<hex>', "deleted_on" = '<now>' WHERE "ml_experiments"."id" = <id>
Update on ml_experiments  (cost=0.29..3.31 rows=0 width=0) (actual time=15.526..15.527 rows=0 loops=1)
  Buffers: shared hit=27 read=25 dirtied=8
  ->  Index Scan using ml_experiments_pkey on ml_experiments  (cost=0.29..3.31 rows=1 width=54) (actual time=2.370..2.373 rows=1 loops=1)
        Index Cond: (id = <id>)
Execution Time: 17.474 ms

The model path adds "model_id" = NULL. The uniqueness validation before it:

SELECT 1 AS one FROM "ml_experiments" WHERE "ml_experiments"."name" = 'deletion_scheduled-<id>-<hex>' AND "ml_experiments"."id" != <id> AND "ml_experiments"."project_id" = <project_id> LIMIT 1
Limit  (cost=0.41..3.44 rows=1 width=4) (actual time=2.991..2.992 rows=0 loops=1)
  Buffers: shared hit=3 read=3
  ->  Index Scan using index_ml_experiments_on_project_id_and_name on ml_experiments  (cost=0.41..3.44 rows=1 width=4) (actual time=2.990..2.991 rows=0 loops=1)
        Index Cond: ((project_id = <project_id>) AND (name = 'deletion_scheduled-<id>-<hex>'::text))
        Filter: (id <> <id>)
Execution Time: 3.057 ms

2. Lookups and lists with deleted_on IS NULL (project with the most experiments). pg_stats.null_frac is 1 for the column, so the planner keeps the same plans; the filter removes 0 rows.

SELECT "ml_experiments".* FROM "ml_experiments" WHERE "ml_experiments"."deleted_on" IS NULL AND "ml_experiments"."project_id" = <project_id> AND "ml_experiments"."iid" = <iid> LIMIT 1
Limit  (cost=0.29..3.31 rows=1 width=97) (actual time=0.030..0.031 rows=1 loops=1)
  Buffers: shared hit=6
  ->  Index Scan using index_ml_experiments_on_project_id_and_iid on public.ml_experiments  (cost=0.29..3.31 rows=1 width=97) (actual time=0.029..0.030 rows=1 loops=1)
        Index Cond: ((ml_experiments.project_id = <project_id>) AND (ml_experiments.iid = <iid>))
        Filter: (ml_experiments.deleted_on IS NULL)
Execution Time: 0.061 ms
SELECT "ml_experiments".* FROM "ml_experiments" WHERE "ml_experiments"."deleted_on" IS NULL AND "ml_experiments"."project_id" = <project_id> AND "ml_experiments"."name" = '<name>' LIMIT 1
Limit  (cost=0.41..3.43 rows=1 width=97) (actual time=5.946..5.946 rows=1 loops=1)
  Buffers: shared hit=4 read=3
  ->  Index Scan using index_ml_experiments_on_project_id_and_name on public.ml_experiments  (cost=0.41..3.43 rows=1 width=97) (actual time=5.945..5.945 rows=1 loops=1)
        Index Cond: ((ml_experiments.project_id = <project_id>) AND (ml_experiments.name = '<name>'::text))
        Filter: (ml_experiments.deleted_on IS NULL)
Execution Time: 5.980 ms

Projects::Ml::ExperimentFinder, id DESC (GraphQL, MLflow experiments/search with order_by=id):

SELECT "ml_experiments".* FROM "ml_experiments" WHERE "ml_experiments"."project_id" = <project_id> AND "ml_experiments"."deleted_on" IS NULL ORDER BY "ml_experiments"."id" DESC LIMIT 101
Limit  (cost=0.29..17.85 rows=101 width=97) (actual time=3.080..17.202 rows=101 loops=1)
  Buffers: shared hit=78 read=27
  ->  Index Scan Backward using ml_experiments_pkey on public.ml_experiments  (cost=0.29..2193.08 rows=12613 width=97) (actual time=3.079..17.189 rows=101 loops=1)
        Filter: ((ml_experiments.deleted_on IS NULL) AND (ml_experiments.project_id = <project_id>))
        Rows Removed by Filter: 4
Execution Time: 17.268 ms

created_at DESC (MLflow experiments/search default), warm; the plan is identical without the filter. LOWER(name) has the same shape (index scan on (project_id, iid) and a top-N sort, 451 ms cold, also identical without the filter). Reading all of a project's experiments for these orders predates this MR.

SELECT "ml_experiments".* FROM "ml_experiments" WHERE "ml_experiments"."project_id" = <project_id> AND "ml_experiments"."deleted_on" IS NULL ORDER BY "ml_experiments"."created_at" DESC, "ml_experiments"."id" DESC LIMIT 201
Limit  (cost=1552.60..1553.10 rows=201 width=97) (actual time=6.781..6.804 rows=201 loops=1)
  Buffers: shared hit=11240
  ->  Sort  (cost=1552.60..1584.13 rows=12613 width=97) (actual time=6.780..6.789 rows=201 loops=1)
        Sort Key: ml_experiments.created_at DESC, ml_experiments.id DESC
        Sort Method: top-N heapsort  Memory: 79kB
        ->  Index Scan using index_ml_experiments_on_project_id_and_iid on public.ml_experiments  (cost=0.29..1007.02 rows=12613 width=97) (actual time=0.032..4.394 rows=13144 loops=1)
              Index Cond: (ml_experiments.project_id = <project_id>)
              Filter: (ml_experiments.deleted_on IS NULL)
Execution Time: 6.847 ms

3. Model request (model with the most versions, cold, one rolled-back transaction)

SELECT "ml_model_versions"."id" FROM "ml_model_versions" WHERE "ml_model_versions"."model_id" = <model_id> AND "ml_model_versions"."id" >= <version_id> AND "ml_model_versions"."id" < <next_version_id>
Index Only Scan using unique_ml_model_versions_on_model_id_and_id on ml_model_versions  (cost=0.29..25.13 rows=704 width=8) (actual time=4.429..9.457 rows=1000 loops=1)
  Index Cond: ((model_id = <model_id>) AND (id >= <version_id>) AND (id < <next_version_id>))
  Heap Fetches: 0
  Buffers: shared read=10
Execution Time: 9.518 ms
UPDATE "ml_candidates" SET "model_version_id" = NULL WHERE "ml_candidates"."model_version_id" IN (<1,000 version ids>)
Update on ml_candidates  (cost=0.43..1443.68 rows=0 width=0) (actual time=1198.064..1198.065 rows=0 loops=1)
  Buffers: shared hit=30490 read=1827 dirtied=1736 written=3
  ->  Index Scan using index_ml_candidates_on_model_version_id on ml_candidates  (cost=0.43..1443.68 rows=1000 width=14) (actual time=4.930..701.221 rows=1000 loops=1)
        Index Cond: (model_version_id = ANY ('{<1,000 version ids>}'::bigint[]))
Execution Time: 1199.729 ms

The four unlink batches took 1,101 + 920 + 741 + 51 ms cold; warm, 1,000 rows take 28 ms and 11.7 MB of WAL (the update changes an indexed column, so it is not a HOT update). The IN (subquery) form planned a merge join from the first entry of index_ml_candidates_on_model_version_id (3,875 ms cold for the first batch), hence the plucked ids.

DELETE FROM "ml_models" WHERE "ml_models"."id" = <model_id>
Delete on ml_models  (cost=0.29..3.31 rows=0 width=0) (actual time=3.947..3.948 rows=0 loops=1)
Trigger for constraint fk_4e8b59e7a8 on ml_models: time=129.660 calls=1
Trigger for constraint fk_rails_97194a054e on ml_models: time=1.766 calls=1
Trigger for constraint fk_rails_d907835e01 on ml_models: time=2.497 calls=1
Trigger for constraint fk_e86e0bfa5a on ml_model_versions: time=15.245 calls=3161
Trigger for constraint fk_rails_5b67cc9107 on ml_model_versions: time=46.430 calls=3161
Execution Time: 200.123 ms

Total for this model: about 3 s cold and 37 MB of WAL, against 7.8 s for the current single DELETE; it no longer grows with the runs' metrics.

Feature flag and rollback

ml_delete_experiments_in_batches (gitlab_com_derisk, default off, project actor). Rollback: disable the flag. New deletions go back to the synchronous path, and the cron and worker from !264622 still finish the experiments that are already marked.

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