Skip to content

Added MLflow API to search for Registered Models

Darby Frey requested to merge registered-models-search-endpoint into master

What does this MR do and why?

Continuing on the work to build the MLFlow Registered Models API, this MR adds an endpoint to search for registered models.

The Search RegisteredModel API accepts filter, max_results, order_by, and page_token and it returns an array of RegisteredModels or an error message if a validation failure occurs.

Note: This implementation provides a limited set of features compatible with the MLflow API. Future iteration will add more capabilities.

Database Review

This change introduces a new updated_at as a new order_by param to the ModelFinder. Below is a sample query:

SELECT
  ml_models.*,
  count(ml_model_versions.id) as version_count
FROM
  "ml_models"
  LEFT OUTER JOIN "ml_model_versions" ON "ml_model_versions"."model_id" = "ml_models"."id"
WHERE
  "ml_models"."project_id" = 20
GROUP BY
  "ml_models"."id"
ORDER BY
  "ml_models"."updated_at" DESC,
  "ml_models"."id" DESC

https://console.postgres.ai/gitlab/gitlab-production-tunnel-pg12/sessions/24340/commands/77757

How to set up and validate locally

  1. In the Rails console, ensure the feature flag is enabled

    Feature.enable(:ml_experiment_tracking)

  2. In the Rails console, run the following to create multiple registered models for the given project:

    project = Project.last # or find a project locally you want to use
model = Ml::CreateModelService.new(project, 'my-model').execute
model = Ml::CreateModelService.new(project, 'my-model-2').execute
model = Ml::CreateModelService.new(project, 'my-model-3').execute

  3. Use the following cURL command to return an array of Registerd Models from the API

    curl -X "GET" "http://GDKHOST/api/v4/projects/PROJECT_ID/ml/mlflow/api/2.0/mlflow/registered-models/search?filter=name%3Dmy" \
    -H 'Authorization: Bearer PERSONAL_ACCESS_TOKEN' \
    -H 'Content-Type: application/json; charset=utf-8'

    This request should return all three Registered Models.

MR acceptance checklist

This checklist encourages us to confirm any changes have been analyzed to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Darby Frey

Merge request reports