Merge branch 'model_registry/update_endpoints' into 'master'

Improves ml_model package endpoints

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/138787



Merged-by: John Mason <9717668-johnmason@users.noreply.gitlab.com>
Approved-by: Corinna Gogolok <cgogolok@gitlab.com>
Approved-by: John Mason <9717668-johnmason@users.noreply.gitlab.com>
Reviewed-by: Corinna Gogolok <cgogolok@gitlab.com>
Reviewed-by: Eduardo Bonet <ebonet@gitlab.com>
Co-authored-by: Eduardo Bonet <ebonet@gitlab.com>
Co-authored-by: Corinna Gogolok <cgogolok@gitlab.com>

app/services/ml/find_or_create_model_version_service.rb

@@ -8,21 +8,20 @@ def initialize(project, params = {})
       @version = params[:version]
       @package = params[:package]
       @description = params[:description]
+      @user = params[:user]
+      @params = params
     end
 
     def execute
-      model = Ml::FindOrCreateModelService.new(@project, @name).execute
+      model_version = Ml::ModelVersion.by_project_id_name_and_version(@project.id, @name, @version)
 
-      model_version = Ml::ModelVersion.find_or_create!(model, @version, @package, @description)
+      return model_version if model_version
 
-      unless model_version.candidate
-        model_version.candidate = ::Ml::CreateCandidateService.new(
-          model.default_experiment,
-          { model_version: model_version }
-        ).execute
-      end
+      model = Ml::Model.by_project_id_and_name(@project.id, @name)
 
-      model_version
+      return unless model
+
+      Ml::CreateModelVersionService.new(model, @params).execute
     end
   end
 end

app/services/packages/ml_model/create_package_file_service.rb

@@ -4,47 +4,27 @@ module Packages
   module MlModel
     class CreatePackageFileService < BaseService
       def execute
-        ::Packages::Package.transaction do
-          package = find_or_create_package
-          find_or_create_model_version(package)
+        @package = params[:model_version]&.package
+
+        return unless @package
 
-          create_package_file(package)
+        ::Packages::Package.transaction do
+          update_package
+          create_package_file
         end
       end
 
       private
 
-      def find_or_create_package
-        package_params = {
-          name: params[:package_name],
-          version: params[:package_version],
-          build: params[:build],
-          status: params[:status]
-        }
-
-        package = ::Packages::MlModel::FindOrCreatePackageService
-                    .new(project, current_user, package_params)
-                    .execute
+      attr_reader :package
 
+      def update_package
         package.update_column(:status, params[:status]) if params[:status] && params[:status] != package.status
 
         package.create_build_infos!(params[:build])
-
-        package
-      end
-
-      def find_or_create_model_version(package)
-        model_version_params = {
-          model_name: package.name,
-          version: package.version,
-          package: package,
-          user: current_user
-        }
-
-        Ml::FindOrCreateModelVersionService.new(project, model_version_params).execute
       end
 
-      def create_package_file(package)
+      def create_package_file
         file_params = {
           file: params[:file],
           size: params[:file].size,

lib/api/ml_model_packages.rb

@@ -23,8 +23,8 @@ class MlModelPackages < ::API::Base
     end
 
     authenticate_with do |accept|
-      accept.token_types(:personal_access_token, :deploy_token, :job_token)
-            .sent_through(:http_token)
+      accept.token_types(:personal_access_token, :job_token)
+            .sent_through(:http_bearer_token)
     end
 
     helpers do
@@ -38,6 +38,15 @@ def project
       def max_file_size_exceeded?
         project.actual_limits.exceeded?(:ml_model_max_file_size, params[:file].size)
       end
+
+      def find_model_version!
+        ::Ml::ModelVersion.by_project_id_name_and_version(project.id, params[:model_name], params[:model_version]) ||
+          not_found!
+      end
+
+      def model_version
+        @model_version ||= find_model_version!
+      end
     end
 
     params do
@@ -88,10 +97,12 @@ def max_file_size_exceeded?
           end
           put do
             authorize_upload!(project)
+            not_found! unless can?(current_user, :write_model_registry, project)
 
             bad_request!('File is too large') if max_file_size_exceeded?
 
             create_package_file_params = declared(params).merge(
+              model_version: model_version,
               build: current_authenticated_job,
               package_name: params[:model_name],
               package_version: params[:model_version]
@@ -123,9 +134,7 @@ def max_file_size_exceeded?
           get do
             authorize_read_package!(project)
 
-            package = ::Packages::MlModel::PackageFinder.new(project)
-                                                        .execute!(params[:model_name], params[:model_version])
-            package_file = ::Packages::PackageFileFinder.new(package, params[:file_name]).execute!
+            package_file = ::Packages::PackageFileFinder.new(model_version.package, params[:file_name]).execute!
 
             present_package_file!(package_file)
           end

spec/requests/api/ml_model_packages_spec.rb

@@ -16,6 +16,8 @@
   let_it_be(:deploy_token) { create(:deploy_token, read_package_registry: true, write_package_registry: true) }
   let_it_be(:project_deploy_token) { create(:project_deploy_token, deploy_token: deploy_token, project: project) }
   let_it_be(:another_project, reload: true) { create(:project) }
+  let_it_be(:model) { create(:ml_models, user: project.owner, project: project) }
+  let_it_be(:model_version) { create(:ml_model_versions, :with_package, model: model, version: '0.1.0') }
 
   let_it_be(:tokens) do
     {
@@ -70,10 +72,6 @@ def authorize_permissions_table
       :private | :guest      | false | :job_token             | true  | :not_found
       :private | :developer  | false | :job_token             | false | :unauthorized
       :private | :guest      | false | :job_token             | false | :unauthorized
-      :public  | :developer  | true  | :deploy_token          | true  | :success
-      :public  | :developer  | true  | :deploy_token          | false | :unauthorized
-      :private | :developer  | true  | :deploy_token          | true  | :success
-      :private | :developer  | true  | :deploy_token          | false | :unauthorized
     end
 
     # :visibility, :user_role, :member, :token_type, :valid_token, :expected_status
@@ -112,10 +110,6 @@ def download_permissions_tables
       :private | :guest      | false | :job_token             | true  |  :not_found
       :private | :developer  | false | :job_token             | false |  :unauthorized
       :private | :guest      | false | :job_token             | false |  :unauthorized
-      :public  | :developer  | true  | :deploy_token          | true  |  :success
-      :public  | :developer  | true  | :deploy_token          | false |  :unauthorized
-      :private | :developer  | true  | :deploy_token          | true  |  :success
-      :private | :developer  | true  | :deploy_token          | false |  :unauthorized
     end
     # rubocop:enable Metrics/AbcSize
   end
@@ -128,14 +122,15 @@ def download_permissions_tables
     include_context 'ml model authorize permissions table'
 
     let(:token) { tokens[:personal_access_token] }
-    let(:user_headers) { { 'HTTP_AUTHORIZATION' => token } }
+    let(:user_headers) { { 'Authorization' => "Bearer #{token}" } }
     let(:headers) { user_headers.merge(workhorse_headers) }
     let(:request) { authorize_upload_file(headers) }
-    let(:model_name) { 'my_package' }
+    let(:model_name) { model_version.name }
+    let(:version) { model_version.version }
     let(:file_name) { 'myfile.tar.gz' }
 
     subject(:api_response) do
-      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/0.0.1/#{file_name}/authorize"
+      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/#{version}/#{file_name}/authorize"
 
       put api(url), headers: headers
 
@@ -149,7 +144,7 @@ def download_permissions_tables
 
       with_them do
         let(:token) { valid_token ? tokens[token_type] : 'invalid-token123' }
-        let(:user_headers) { user_role == :anonymous ? {} : { 'HTTP_AUTHORIZATION' => token } }
+        let(:user_headers) { user_role == :anonymous ? {} : { 'Authorization' => "Bearer #{token}" } }
 
         before do
           project.update_column(:visibility_level, Gitlab::VisibilityLevel.level_value(visibility.to_s))
@@ -183,15 +178,16 @@ def download_permissions_tables
     let_it_be(:file_name) { 'model.md5' }
 
     let(:token) { tokens[:personal_access_token] }
-    let(:user_headers) { { 'HTTP_AUTHORIZATION' => token } }
+    let(:user_headers) { { 'Authorization' => "Bearer #{token}" } }
     let(:headers) { user_headers.merge(workhorse_headers) }
     let(:params) { { file: temp_file(file_name) } }
     let(:file_key) { :file }
     let(:send_rewritten_field) { true }
-    let(:model_name) { 'my_package' }
+    let(:model_name) { model_version.name }
+    let(:version) { model_version.version }
 
     subject(:api_response) do
-      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/0.0.1/#{file_name}"
+      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/#{version}/#{file_name}"
 
       workhorse_finalize(
         api(url),
@@ -219,7 +215,7 @@ def download_permissions_tables
 
       with_them do
         let(:token) { valid_token ? tokens[token_type] : 'invalid-token123' }
-        let(:user_headers) { user_role == :anonymous ? {} : { 'HTTP_AUTHORIZATION' => token } }
+        let(:user_headers) { user_role == :anonymous ? {} : { 'Authorization' => "Bearer #{token}" } }
 
         before do
           project.update_column(:visibility_level, Gitlab::VisibilityLevel.level_value(visibility.to_s))
@@ -233,25 +229,27 @@ def download_permissions_tables
       end
 
       it_behaves_like 'Endpoint not found if read_model_registry not available'
+      it_behaves_like 'Endpoint not found if write_model_registry not available'
+      it_behaves_like 'Not found when model version does not exist'
     end
   end
 
   describe 'GET /api/v4/projects/:project_id/packages/ml_models/:model_name/:model_version/:file_name' do
     include_context 'ml model authorize permissions table'
 
-    let_it_be(:package) { create(:ml_model_package, project: project, name: 'model', version: '0.0.1') }
+    let_it_be(:package) { model_version.package }
     let_it_be(:package_file) { create(:package_file, :generic, package: package, file_name: 'model.md5') }
 
-    let(:model_name) { package.name }
-    let(:model_version) { package.version }
+    let(:model_name) { model_version.name }
+    let(:version) { model_version.version }
     let(:file_name) { package_file.file_name }
 
     let(:token) { tokens[:personal_access_token] }
-    let(:user_headers) { { 'HTTP_AUTHORIZATION' => token } }
+    let(:user_headers) { { 'Authorization' => "Bearer #{token}" } }
     let(:headers) { user_headers.merge(workhorse_headers) }
 
     subject(:api_response) do
-      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/#{model_version}/#{file_name}"
+      url = "/projects/#{project.id}/packages/ml_models/#{model_name}/#{version}/#{file_name}"
 
       get api(url), headers: headers
 
@@ -265,7 +263,7 @@ def download_permissions_tables
 
       with_them do
         let(:token) { valid_token ? tokens[token_type] : 'invalid-token123' }
-        let(:user_headers) { user_role == :anonymous ? {} : { 'HTTP_AUTHORIZATION' => token } }
+        let(:user_headers) { user_role == :anonymous ? {} : { 'Authorization' => "Bearer #{token}" } }
 
         before do
           project.update_column(:visibility_level, Gitlab::VisibilityLevel.level_value(visibility.to_s))

spec/services/ml/find_or_create_model_version_service_spec.rb

@@ -35,21 +35,29 @@
       end
     end
 
-    context 'when model version does not exist' do
+    context 'when model does not exist' do
       let(:project) { existing_version.project }
       let(:name) { 'a_new_model' }
       let(:version) { '2.0.0' }
+
+      it 'does not create a new model version', :aggregate_failures do
+        expect { model_version }.to change { Ml::ModelVersion.count }.by(0)
+      end
+    end
+
+    context 'when model exists and model version does not' do
+      let(:project) { existing_version.project }
+      let(:name) { existing_version.name }
+      let(:version) { '2.0.0' }
       let(:description) { 'A model version' }
 
       let(:package) { create(:ml_model_package, project: project, name: name, version: version) }
 
       it 'creates a new model version', :aggregate_failures do
-        expect { model_version }.to change { Ml::ModelVersion.count }.by(1).and change { Ml::Candidate.count }.by(1)
+        expect { model_version }.to change { Ml::ModelVersion.count }.by(1)
 
-        expect(model_version.name).to eq(name)
         expect(model_version.version).to eq(version)
-        expect(model_version.package).to eq(package)
-        expect(model_version.candidate.model_version_id).to eq(model_version.id)
+        expect(model_version.model).to eq(existing_version.model)
         expect(model_version.description).to eq(description)
       end
     end

spec/services/packages/ml_model/create_package_file_service_spec.rb

@@ -8,6 +8,8 @@
     let_it_be(:user) { create(:user) }
     let_it_be(:pipeline) { create(:ci_pipeline, user: user, project: project) }
     let_it_be(:file_name) { 'myfile.tar.gz.1' }
+    let_it_be(:model) { create(:ml_models, user: user, project: project) }
+    let_it_be(:model_version) { create(:ml_model_versions, :with_package, model: model, version: '0.1.0') }
 
     let(:build) { instance_double(Ci::Build, pipeline: pipeline) }
 
@@ -26,47 +28,24 @@
       FileUtils.rm_f(temp_file)
     end
 
-    context 'without existing package' do
+    context 'when model version is nil' do
       let(:params) do
         {
-          package_name: 'new_model',
-          package_version: '1.0.0',
+          model_version: nil,
           file: file,
           file_name: file_name
         }
       end
 
-      it 'creates package file', :aggregate_failures do
-        expect { execute_service }
-          .to change { Packages::MlModel::Package.count }.by(1)
-          .and change { Packages::PackageFile.count }.by(1)
-          .and change { Packages::PackageFileBuildInfo.count }.by(0)
-          .and change { Ml::ModelVersion.count }.by(1)
-
-        new_model = Packages::MlModel::Package.last
-        package_file = new_model.package_files.last
-        new_model_version = Ml::ModelVersion.last
-
-        expect(new_model.name).to eq('new_model')
-        expect(new_model.version).to eq('1.0.0')
-        expect(new_model.status).to eq('default')
-        expect(package_file.package).to eq(new_model)
-        expect(package_file.file_name).to eq(file_name)
-        expect(package_file.size).to eq(file.size)
-        expect(package_file.file_sha256).to eq(sha256)
-        expect(new_model_version.name).to eq('new_model')
-        expect(new_model_version.version).to eq('1.0.0')
-        expect(new_model_version.package).to eq(new_model)
+      it 'does not create package file', :aggregate_failures do
+        expect(execute_service).to be(nil)
       end
     end
 
-    context 'with existing package' do
-      let_it_be(:model) { create(:ml_model_package, creator: user, project: project, version: '0.1.0') }
-
+    context 'with existing model version' do
       let(:params) do
         {
-          package_name: model.name,
-          package_version: model.version,
+          model_version: model_version,
           file: file,
           file_name: file_name,
           status: :hidden,
@@ -76,18 +55,16 @@
 
       it 'adds the package file and updates status and ci_build', :aggregate_failures do
         expect { execute_service }
-          .to change { project.packages.ml_model.count }.by(0)
-          .and change { model.package_files.count }.by(1)
+          .to change { model_version.package.package_files.count }.by(1)
           .and change { Packages::PackageFileBuildInfo.count }.by(1)
 
-        model.reload
-
-        package_file = model.package_files.last
+        package = model_version.reload.package
+        package_file = package.package_files.last
 
-        expect(model.build_infos.first.pipeline).to eq(build.pipeline)
-        expect(model.status).to eq('hidden')
+        expect(package.build_infos.first.pipeline).to eq(build.pipeline)
+        expect(package.status).to eq('hidden')
 
-        expect(package_file.package).to eq(model)
+        expect(package_file.package).to eq(package)
         expect(package_file.file_name).to eq(file_name)
         expect(package_file.size).to eq(file.size)
         expect(package_file.file_sha256).to eq(sha256)

spec/support/shared_examples/requests/api/ml_model_packages_shared_examples.rb

@@ -15,11 +15,31 @@
   end
 end
 
-RSpec.shared_examples 'creates model experiments package files' do
+RSpec.shared_examples 'Endpoint not found if write_model_registry not available' do
+  context 'when write_model_registry is disabled for current project' do
+    before do
+      allow(Ability).to receive(:allowed?).and_call_original
+      allow(Ability).to receive(:allowed?)
+                          .with(user, :write_model_registry, project)
+                          .and_return(false)
+    end
+
+    it { is_expected.to have_gitlab_http_status(:not_found) }
+  end
+end
+
+RSpec.shared_examples 'Not found when model version does not exist' do
+  context 'when model version does not exist' do
+    let(:version) { "#{non_existing_record_id}.0.0" }
+
+    it { is_expected.to have_gitlab_http_status(:not_found) }
+  end
+end
+
+RSpec.shared_examples 'creates package files for model versions' do
   it 'creates package files', :aggregate_failures do
     expect { api_response }
-      .to change { project.packages.count }.by(1)
-                                           .and change { Packages::PackageFile.count }.by(1)
+      .to change { Packages::PackageFile.count }.by(1)
     expect(api_response).to have_gitlab_http_status(:created)
 
     package_file = project.packages.last.package_files.reload.last
@@ -59,7 +79,7 @@
 
     context 'with correct params' do
       it_behaves_like 'package workhorse uploads'
-      it_behaves_like 'creates model experiments package files'
+      it_behaves_like 'creates package files for model versions'
       # To be reactivated with https://gitlab.com/gitlab-org/gitlab/-/issues/414270
       # it_behaves_like 'a package tracking event', '::API::MlModelPackages', 'push_package'
     end
@@ -81,7 +101,7 @@
         stub_package_file_object_storage(direct_upload: true)
       end
 
-      it_behaves_like 'creates model experiments package files'
+      it_behaves_like 'creates package files for model versions'
 
       ['123123', '../../123123'].each do |remote_id|
         context "with invalid remote_id: #{remote_id}" do
@@ -102,7 +122,7 @@
         stub_package_file_object_storage(direct_upload: false)
       end
 
-      it_behaves_like 'creates model experiments package files'
+      it_behaves_like 'creates package files for model versions'
     end
   end
 end
@@ -112,13 +132,5 @@
     it { is_expected.to have_gitlab_http_status(:success) }
   end
 
-  context 'when record does not exist' do
-    it 'response is not found' do
-      expect_next_instance_of(::Packages::MlModel::PackageFinder) do |instance|
-        expect(instance).to receive(:execute!).and_raise(ActiveRecord::RecordNotFound)
-      end
-
-      expect(api_response).to have_gitlab_http_status(:not_found)
-    end
-  end
+  it_behaves_like 'Not found when model version does not exist'
 end