From a9bd4c5bc84f78972f2df4a8455f554ee1b64ce9 Mon Sep 17 00:00:00 2001
From: Eduardo Bonet <ebonet@gitlab.com>
Date: Mon, 17 Jul 2023 12:59:37 +0200
Subject: [PATCH] Creates Ml::ModelVersion on package creation
When a package of type ml_model_package is created through the api,
creates also a companion Ml::ModelVersion
Changelog: added
FeatureFlag: model_registry
---
app/models/ml/model_version.rb | 4 ++--
.../find_or_create_model_version_service.rb | 2 +-
.../ml_model/create_package_file_service.rb | 15 ++++++++++++-
spec/models/ml/model_version_spec.rb | 4 ++--
.../create_package_file_service_spec.rb | 21 +++++++++++--------
5 files changed, 31 insertions(+), 15 deletions(-)
diff --git a/app/models/ml/model_version.rb b/app/models/ml/model_version.rb
index 3450dc64b79f3..934b924e4afe8 100644
--- a/app/models/ml/model_version.rb
+++ b/app/models/ml/model_version.rb
@@ -19,8 +19,8 @@ class ModelVersion < ApplicationRecord
delegate :name, to: :model
class << self
- def find_or_create(model, version, package)
- create_with(package: package).find_or_create_by(project: model.project, model: model, version: version)
+ def find_or_create!(model, version, package)
+ create_with(package: package).find_or_create_by!(project: model.project, model: model, version: version)
end
end
diff --git a/app/services/ml/find_or_create_model_version_service.rb b/app/services/ml/find_or_create_model_version_service.rb
index 558374ad57bd0..1316b2546b9d8 100644
--- a/app/services/ml/find_or_create_model_version_service.rb
+++ b/app/services/ml/find_or_create_model_version_service.rb
@@ -12,7 +12,7 @@ def initialize(project, params = {})
def execute
model = Ml::FindOrCreateModelService.new(project, name).execute
- Ml::ModelVersion.find_or_create(model, version, package)
+ Ml::ModelVersion.find_or_create!(model, version, package)
end
private
diff --git a/app/services/packages/ml_model/create_package_file_service.rb b/app/services/packages/ml_model/create_package_file_service.rb
index 574f70940fcf3..b1e8e8140159d 100644
--- a/app/services/packages/ml_model/create_package_file_service.rb
+++ b/app/services/packages/ml_model/create_package_file_service.rb
@@ -5,7 +5,10 @@ module MlModel
class CreatePackageFileService < BaseService
def execute
::Packages::Package.transaction do
- create_package_file(find_or_create_package)
+ package = find_or_create_package
+ find_or_create_model_version(package)
+
+ create_package_file(package)
end
end
@@ -30,6 +33,16 @@ def find_or_create_package
package
end
+ def find_or_create_model_version(package)
+ model_version_params = {
+ model_name: package.name,
+ version: package.version,
+ package: package
+ }
+
+ Ml::FindOrCreateModelVersionService.new(project, model_version_params).execute
+ end
+
def create_package_file(package)
file_params = {
file: params[:file],
diff --git a/spec/models/ml/model_version_spec.rb b/spec/models/ml/model_version_spec.rb
index 9af9a2a56c6c7..7b8410f99a889 100644
--- a/spec/models/ml/model_version_spec.rb
+++ b/spec/models/ml/model_version_spec.rb
@@ -88,13 +88,13 @@
end
end
- describe '.find_or_create' do
+ describe '#find_or_create!' do
let_it_be(:existing_model_version) { create(:ml_model_versions, model: model, version: 'abc') }
let(:version) { existing_model_version.version }
let(:package) { nil }
- subject(:find_or_create) { described_class.find_or_create(model, version, package) }
+ subject(:find_or_create) { described_class.find_or_create!(model, version, package) }
context 'if model version exists' do
it 'returns the model version', :aggregate_failures do
diff --git a/spec/services/packages/ml_model/create_package_file_service_spec.rb b/spec/services/packages/ml_model/create_package_file_service_spec.rb
index d749aee227a70..32754279e1724 100644
--- a/spec/services/packages/ml_model/create_package_file_service_spec.rb
+++ b/spec/services/packages/ml_model/create_package_file_service_spec.rb
@@ -41,19 +41,22 @@
.to change { project.packages.ml_model.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 = project.packages.ml_model.last
package_file = new_model.package_files.last
+ new_model_version = Ml::ModelVersion.last
- aggregate_failures do
- 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)
- end
+ 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)
end
end
--
GitLab