path: root/lib/api/ml/mlflow.rb

# frozen_string_literal: true

require 'mime/types'

module API
  # MLFlow integration API, replicating the Rest API https://www.mlflow.org/docs/latest/rest-api.html#rest-api
  module Ml
    class Mlflow < ::API::Base
      include APIGuard

      # The first part of the url is the namespace, the second part of the URL is what the MLFlow client calls
      MLFLOW_API_PREFIX = ':id/ml/mlflow/api/2.0/mlflow/'

      allow_access_with_scope :api
      allow_access_with_scope :read_api, if: -> (request) { request.get? || request.head? }

      feature_category :mlops

      content_type :json, 'application/json'
      default_format :json

      before do
        # MLFlow Client considers any status code different than 200 an error, even 201
        status 200

        authenticate!

        not_found! unless Feature.enabled?(:ml_experiment_tracking, user_project)
      end

      rescue_from ActiveRecord::ActiveRecordError do |e|
        invalid_parameter!(e.message)
      end

      helpers do
        def resource_not_found!
          render_structured_api_error!({ error_code: 'RESOURCE_DOES_NOT_EXIST' }, 404)
        end

        def resource_already_exists!
          render_structured_api_error!({ error_code: 'RESOURCE_ALREADY_EXISTS' }, 400)
        end

        def invalid_parameter!(message = nil)
          render_structured_api_error!({ error_code: 'INVALID_PARAMETER_VALUE', message: message }, 400)
        end

        def experiment_repository
          ::Ml::ExperimentTracking::ExperimentRepository.new(user_project, current_user)
        end

        def candidate_repository
          ::Ml::ExperimentTracking::CandidateRepository.new(user_project, current_user)
        end

        def experiment
          @experiment ||= find_experiment!(params[:experiment_id], params[:experiment_name])
        end

        def candidate
          @candidate ||= find_candidate!(params[:run_id])
        end

        def find_experiment!(iid, name)
          experiment_repository.by_iid_or_name(iid: iid, name: name) || resource_not_found!
        end

        def find_candidate!(iid)
          candidate_repository.by_iid(iid) || resource_not_found!
        end

        def packages_url
          path = api_v4_projects_packages_generic_package_version_path(
            id: user_project.id, package_name: '', file_name: ''
          )
          path = path.delete_suffix('/package_version')

          "#{request.base_url}#{path}"
        end
      end

      params do
        requires :id, types: [String, Integer], desc: 'The ID or URL-encoded path of the project'
      end
      resource :projects, requirements: API::NAMESPACE_OR_PROJECT_REQUIREMENTS do
        desc 'API to interface with MLFlow Client, REST API version 1.28.0' do
          detail 'This feature is gated by :ml_experiment_tracking.'
        end
        namespace MLFLOW_API_PREFIX do
          resource :experiments do
            desc 'Fetch experiment by experiment_id' do
              success Entities::Ml::Mlflow::GetExperiment
              detail 'https://www.mlflow.org/docs/1.28.0/rest-api.html#get-experiment'
            end
            params do
              optional :experiment_id, type: String, default: '', desc: 'Experiment ID, in reference to the project'
            end
            get 'get', urgency: :low do
              present experiment, with: Entities::Ml::Mlflow::GetExperiment
            end

            desc 'Fetch experiment by experiment_name' do
              success Entities::Ml::Mlflow::GetExperiment
              detail 'https://www.mlflow.org/docs/1.28.0/rest-api.html#get-experiment-by-name'
            end
            params do
              optional :experiment_name, type: String, default: '', desc: 'Experiment name'
            end
            get 'get-by-name', urgency: :low do
              present experiment, with: Entities::Ml::Mlflow::GetExperiment
            end

            desc 'List experiments' do
              success Entities::Ml::Mlflow::ListExperiment
              detail 'https://www.mlflow.org/docs/latest/rest-api.html#list-experiments'
            end
            get 'list', urgency: :low do
              response = { experiments: experiment_repository.all }

              present response, with: Entities::Ml::Mlflow::ListExperiment
            end

            desc 'Create experiment' do
              success Entities::Ml::Mlflow::NewExperiment
              detail 'https://www.mlflow.org/docs/1.28.0/rest-api.html#create-experiment'
            end
            params do
              requires :name, type: String, desc: 'Experiment name'
              optional :tags, type: Array, desc: 'Tags with information about the experiment'
              optional :artifact_location, type: String, desc: 'This will be ignored'
            end
            post 'create', urgency: :low do
              present experiment_repository.create!(params[:name], params[:tags]),
                      with: Entities::Ml::Mlflow::NewExperiment
            rescue ActiveRecord::RecordInvalid
              resource_already_exists!
            end

            desc 'Sets a tag for an experiment.' do
              summary 'Sets a tag for an experiment. '

              detail  'https://www.mlflow.org/docs/1.28.0/rest-api.html#set-experiment-tag'
            end
            params do
              requires :experiment_id, type: String, desc: 'ID of the experiment.'
              requires :key, type: String, desc: 'Name for the tag.'
              requires :value, type: String, desc: 'Value for the tag.'
            end
            post 'set-experiment-tag', urgency: :low do
              bad_request! unless experiment_repository.add_tag!(experiment, params[:key], params[:value])

              {}
            end
          end

          resource :runs do
            desc 'Creates a Run.' do
              success Entities::Ml::Mlflow::Run
              detail 'MLFlow Runs map to GitLab Candidates. https://www.mlflow.org/docs/1.28.0/rest-api.html#create-run'
            end
            params do
              requires :experiment_id, type: Integer,
                                       desc: 'Id for the experiment, relative to the project'
              optional :start_time, type: Integer,
                                    desc: 'Unix timestamp in milliseconds of when the run started.',
                                    default: 0
              optional :user_id, type: String, desc: 'This will be ignored'
              optional :tags, type: Array, desc: 'Tags are stored, but not displayed'
            end
            post 'create', urgency: :low do
              present candidate_repository.create!(experiment, params[:start_time], params[:tags]),
                      with: Entities::Ml::Mlflow::Run, packages_url: packages_url
            end

            desc 'Gets an MLFlow Run, which maps to GitLab Candidates' do
              success Entities::Ml::Mlflow::Run
              detail 'https://www.mlflow.org/docs/1.28.0/rest-api.html#get-run'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the candidate.'
              optional :run_uuid, type: String, desc: 'This parameter is ignored'
            end
            get 'get', urgency: :low do
              present candidate, with: Entities::Ml::Mlflow::Run, packages_url: packages_url
            end

            desc 'Updates a Run.' do
              success Entities::Ml::Mlflow::UpdateRun
              detail 'MLFlow Runs map to GitLab Candidates. https://www.mlflow.org/docs/1.28.0/rest-api.html#update-run'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the candidate.'
              optional :status, type: String,
                                values: ::Ml::Candidate.statuses.keys.map(&:upcase),
                                desc: "Status of the run. Accepts: " \
                                      "#{::Ml::Candidate.statuses.keys.map(&:upcase)}."
              optional :end_time, type: Integer, desc: 'Ending time of the run'
            end
            post 'update', urgency: :low do
              candidate_repository.update(candidate, params[:status], params[:end_time])

              present candidate, with: Entities::Ml::Mlflow::UpdateRun, packages_url: packages_url
            end

            desc 'Logs a metric to a run.' do
              summary 'Log a metric for a run. A metric is a key-value pair (string key, float value) with an '\
                      'associated timestamp. Examples include the various metrics that represent ML model accuracy. '\
                      'A metric can be logged multiple times.'
              detail  'https://www.mlflow.org/docs/1.28.0/rest-api.html#log-metric'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the run.'
              requires :key, type: String, desc: 'Name for the metric.'
              requires :value, type: Float, desc: 'Value of the metric.'
              requires :timestamp, type: Integer, desc: 'Unix timestamp in milliseconds when metric was recorded'
              optional :step, type: Integer, desc: 'Step at which the metric was recorded'
            end
            post 'log-metric', urgency: :low do
              candidate_repository.add_metric!(
                candidate,
                params[:key],
                params[:value],
                params[:timestamp],
                params[:step]
              )

              {}
            end

            desc 'Logs a parameter to a run.' do
              summary 'Log a param used for a run. A param is a key-value pair (string key, string value). '\
                      'Examples include hyperparameters used for ML model training and constant dates and values '\
                      'used in an ETL pipeline. A param can be logged only once for a run, duplicate will be .'\
                      'ignored'

              detail  'https://www.mlflow.org/docs/1.28.0/rest-api.html#log-param'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the run.'
              requires :key, type: String, desc: 'Name for the parameter.'
              requires :value, type: String, desc: 'Value for the parameter.'
            end
            post 'log-parameter', urgency: :low do
              bad_request! unless candidate_repository.add_param!(candidate, params[:key], params[:value])

              {}
            end

            desc 'Sets a tag for a run.' do
              summary 'Sets a tag for a run. '

              detail  'https://www.mlflow.org/docs/1.28.0/rest-api.html#set-tag'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the run.'
              requires :key, type: String, desc: 'Name for the tag.'
              requires :value, type: String, desc: 'Value for the tag.'
            end
            post 'set-tag', urgency: :low do
              bad_request! unless candidate_repository.add_tag!(candidate, params[:key], params[:value])

              {}
            end

            desc 'Logs multiple parameters and metrics.' do
              summary 'Log a batch of metrics and params for a run. Validation errors will block the entire batch, '\
                      'duplicate errors will be ignored.'

              detail  'https://www.mlflow.org/docs/1.28.0/rest-api.html#log-param'
            end
            params do
              requires :run_id, type: String, desc: 'UUID of the run.'
              optional :metrics, type: Array, default: [] do
                requires :key, type: String, desc: 'Name for the metric.'
                requires :value, type: Float, desc: 'Value of the metric.'
                requires :timestamp, type: Integer, desc: 'Unix timestamp in milliseconds when metric was recorded'
                optional :step, type: Integer, desc: 'Step at which the metric was recorded'
              end
              optional :params, type: Array, default: [] do
                requires :key, type: String, desc: 'Name for the metric.'
                requires :value, type: String, desc: 'Value of the metric.'
              end
            end
            post 'log-batch', urgency: :low do
              candidate_repository.add_metrics(candidate, params[:metrics])
              candidate_repository.add_params(candidate, params[:params])
              candidate_repository.add_tags(candidate, params[:tags])

              {}
            end
          end
        end
      end
    end
  end
end