Add latest changes from gitlab-org/gitlab@15-5-stable-eev15.5.0-rc42

1 files changed, 137 insertions, 54 deletions

diff --git a/lib/api/ml/mlflow.rb b/lib/api/ml/mlflow.rb
index 4f5bd42f8f9..2ffb04ebcbd 100644
--- a/lib/api/ml/mlflow.rb
+++ b/lib/api/ml/mlflow.rb
@@ -9,20 +9,28 @@ module API
       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/mflow/api/2.0/mlflow/'
+      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
 
-      feature_category :mlops
-
-      content_type :json, 'application/json'
-      default_format :json
+      rescue_from ActiveRecord::ActiveRecordError do |e|
+        invalid_parameter!(e.message)
+      end
 
       helpers do
         def resource_not_found!
@@ -32,6 +40,34 @@ module API
         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
       end
 
       params do
@@ -44,33 +80,35 @@ module API
         namespace MLFLOW_API_PREFIX do
           resource :experiments do
             desc 'Fetch experiment by experiment_id' do
-              success Entities::Ml::Mlflow::Experiment
+              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
-              experiment = ::Ml::Experiment.by_project_id_and_iid(user_project.id, params[:experiment_id])
-
-              resource_not_found! unless experiment
-
-              present experiment, with: Entities::Ml::Mlflow::Experiment
+              present experiment, with: Entities::Ml::Mlflow::GetExperiment
             end
 
             desc 'Fetch experiment by experiment_name' do
-              success Entities::Ml::Mlflow::Experiment
+              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
-              experiment = ::Ml::Experiment.by_project_id_and_name(user_project, params[:experiment_name])
+              present experiment, with: Entities::Ml::Mlflow::GetExperiment
+            end
 
-              resource_not_found! unless experiment
+            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 experiment, with: Entities::Ml::Mlflow::Experiment
+              present response, with: Entities::Ml::Mlflow::ListExperiment
             end
 
             desc 'Create experiment' do
@@ -83,33 +121,13 @@ module API
               optional :tags, type: Array, desc: 'This will be ignored'
             end
             post 'create', urgency: :low do
-              resource_already_exists! if ::Ml::Experiment.has_record?(user_project.id, params[:name])
-
-              experiment = ::Ml::Experiment.create!(name: params[:name],
-                                                    user: current_user,
-                                                    project: user_project)
-
-              present experiment, with: Entities::Ml::Mlflow::NewExperiment
+              present experiment_repository.create!(params[:name]), with: Entities::Ml::Mlflow::NewExperiment
+            rescue ActiveRecord::RecordInvalid
+              resource_already_exists!
             end
           end
 
           resource :runs do
-            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
-              optional :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
-              candidate = ::Ml::Candidate.with_project_id_and_iid(user_project.id, params[:run_id])
-
-              resource_not_found! unless candidate
-
-              present candidate, with: Entities::Ml::Mlflow::Run
-            end
-
             desc 'Creates a Run.' do
               success Entities::Ml::Mlflow::Run
               detail  ['https://www.mlflow.org/docs/1.28.0/rest-api.html#create-run',
@@ -125,16 +143,18 @@ module API
               optional :tags, type: Array, desc: 'This will be ignored'
             end
             post 'create', urgency: :low do
-              experiment = ::Ml::Experiment.by_project_id_and_iid(user_project.id, params[:experiment_id].to_i)
-
-              resource_not_found! unless experiment
-
-              candidate = ::Ml::Candidate.create!(
-                experiment: experiment,
-                user: current_user,
-                start_time: params[:start_time] || 0
-              )
+              present candidate_repository.create!(experiment, params[:start_time]), with: Entities::Ml::Mlflow::Run
+            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
             end
 
@@ -144,7 +164,7 @@ module API
                        'MLFlow Runs map to GitLab Candidates']
             end
             params do
-              optional :run_id, type: String, desc: 'UUID of the candidate.'
+              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: " \
@@ -152,16 +172,79 @@ module API
               optional :end_time, type: Integer, desc: 'Ending time of the run'
             end
             post 'update', urgency: :low do
-              candidate = ::Ml::Candidate.with_project_id_and_iid(user_project.id, params[:run_id])
+              candidate_repository.update(candidate, params[:status], params[:end_time])
 
-              resource_not_found! unless candidate
+              present candidate, with: Entities::Ml::Mlflow::UpdateRun
+            end
 
-              candidate.status = params[:status].downcase if params[:status]
-              candidate.end_time = params[:end_time] if params[:end_time]
+            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
 
-              candidate.save if candidate.valid?
+            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'
 
-              present candidate, with: Entities::Ml::Mlflow::UpdateRun
+              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 '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])
+
+              {}
             end
           end
         end