gitlab-ce/lib/api/ml_model_packages.rb

# frozen_string_literal: true

module API
  class MlModelPackages < ::API::Base
    include APIGuard
    include ::API::Helpers::Authentication

    ML_MODEL_PACKAGES_REQUIREMENTS = {
      model_name: API::NO_SLASH_URL_PART_REGEX,
      file_name: API::NO_SLASH_URL_PART_REGEX
    }.freeze

    FAILURES = [
      { code: 401, message: 'Unauthorized' },
      { code: 403, message: 'Forbidden' },
      { code: 404, message: 'Not Found' }
    ].freeze

    ALLOWED_STATUSES = %w[default hidden].freeze

    CANDIDATE_PREFIX = 'candidate:'

    feature_category :mlops
    urgency :low

    after_validation do
      require_packages_enabled!
      authenticate_non_get!

      not_found! unless can?(current_user, :read_model_registry, user_project)
    end

    authenticate_with do |accept|
      accept.token_types(:personal_access_token, :job_token)
            .sent_through(:http_bearer_token)
    end

    helpers do
      include ::API::Helpers::PackagesHelpers
      include ::API::Helpers::Packages::BasicAuthHelpers

      def project
        authorized_user_project
      end

      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_and_id(user_project.id, params[:model_version_id]) || not_found!
      end

      def find_candidate!
        candidate_iid = params[:model_version_id].delete_prefix(CANDIDATE_PREFIX)
        candidate = ::Ml::Candidate.with_project_id_and_iid(user_project.id, candidate_iid)

        candidate || not_found!
      end

      def model_version
        @model_version ||= find_model_version!
      end

      def candidate
        @candidate ||= find_candidate!
      end

      def candidate_package
        ::Packages::MlModel::PackageForCandidateService
           .new(candidate.project, current_user, { candidate: candidate })
           .execute
      end

      def package
        return candidate_package if params[:model_version_id].starts_with?(CANDIDATE_PREFIX)

        model_version.package
      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
      params do
        requires :file_name, type: String, desc: 'File name', file_path: true,
          regexp: Gitlab::Regex.ml_model_file_name_regex
        optional :path, type: String, desc: 'File directory path'
        optional :status, type: String, values: ALLOWED_STATUSES, desc: 'Package status'
        requires :model_version_id, type: String, desc: 'Model version id'
      end
      namespace ':id/packages/ml_models/:model_version_id/files/(*path/):file_name',
        requirements: ML_MODEL_PACKAGES_REQUIREMENTS do
        desc 'Workhorse authorize model package file' do
          detail 'Introduced in GitLab 16.8'
          success code: 200
          failure FAILURES
          tags %w[ml_model_registry]
        end
        put 'authorize' do
          authorize_workhorse!(subject: project, maximum_size: project.actual_limits.ml_model_max_file_size)
        end

        desc 'Workhorse upload model package file' do
          detail 'Introduced in GitLab 16.8'
          success code: 201
          failure FAILURES
          tags %w[ml_model_registry]
        end
        params do
          requires :file,
            type: ::API::Validations::Types::WorkhorseFile,
            desc: 'The package file to be published (generated by Multipart middleware)',
            documentation: { type: 'file' }
        end
        put do
          authorize_upload!(project)
          not_found! unless can?(current_user, :write_model_registry, project)

          bad_request!(s_('MlModelRegistry|Artifact file is too large')) if max_file_size_exceeded?

          bad_request!(s_('MlModelRegistry|Package creation failed')) unless package

          create_package_file_params = declared(params).merge(
            package: package,
            build: current_authenticated_job,
            package_name: package.name,
            package_version: package.version,
            file_name: [params[:path], params[:file_name]].compact.join('/')
          )

          package_file = ::Packages::MlModel::CreatePackageFileService
                           .new(project, current_user, create_package_file_params)
                           .execute

          bad_request!(s_('MlModelRegistry|Artifact file creation failed')) unless package_file

          track_package_event('push_package', :ml_model, project: project, namespace: project.namespace)

          created!
        rescue ObjectStorage::RemoteStoreError => e
          Gitlab::ErrorTracking.track_exception(e, extra: { file_name: params[:file_name], project_id: project.id })

          forbidden!
        end

        desc 'Download an ml_model package file' do
          detail 'This feature was introduced in GitLab 16.8'
          success code: 200
          failure FAILURES
          tags %w[ml_model_registry]
        end

        get do
          authorize_read_package!(project)

          file_name = URI.encode_uri_component([params[:path], params[:file_name]].compact.join('/'))

          package_file = ::Packages::PackageFileFinder.new(package, file_name).execute!

          track_package_event('pull_package', :ml_model, project: project, namespace: project.namespace)

          present_package_file!(package_file)
        end
      end
    end
  end
end