diff --git a/doc/user/project/ml/experiment_tracking/mlflow_client.md b/doc/user/project/ml/experiment_tracking/mlflow_client.md
index 477a492f09d35b20d2f8c8feb811d9638f3f5043..c46a084d0833885fbf9e4ee17e6c03f7a3c6604f 100644
--- a/doc/user/project/ml/experiment_tracking/mlflow_client.md
+++ b/doc/user/project/ml/experiment_tracking/mlflow_client.md
@@ -130,7 +130,26 @@ with mlflow.start_run():
model.fit(X_train, y_train)
# Log the model using MLflow sklearn mode flavour
- mlflow.sklearn.log_model(model, artifact_path="model")
+ mlflow.sklearn.log_model(model, artifact_path="")
+```
+
+### Loading a run
+
+> - [Introduced](https://gitlab.com/gitlab-org/gitlab/-/issues/509595) in GitLab 17.9.
+
+You can load a run from the GitLab model registry to, for example, make predictions.
+
+```python
+import mlflow
+import mlflow.pyfunc
+
+run_id = "<your_run_id>"
+download_path = "models" # Local folder to download to
+
+mlflow.pyfunc.load_model(f"runs:/{run_id}/", dst_path=download_path)
+
+sample_input = [[1,0,3,4],[2,0,1,2]]
+model.predict(data=sample_input)
```
### Associating a run to a CI/CD job
diff --git a/lib/api/entities/ml/mlflow/run_info.rb b/lib/api/entities/ml/mlflow/run_info.rb
index ee694c16fc9038ca6ac1156234a4652e1aaccb7e..aa6fad4d1509bf1a12f82685916505c225eadf7c 100644
--- a/lib/api/entities/ml/mlflow/run_info.rb
+++ b/lib/api/entities/ml/mlflow/run_info.rb
@@ -21,21 +21,22 @@ class RunInfo < Grape::Entity
private
CANDIDATE_PREFIX = 'candidate:'
+ MLFLOW_ARTIFACTS_PREFIX = 'mlflow-artifacts'
def run_id
object.eid.to_s
end
def artifact_uri
- uri = if object.package&.generic?
- generic_package_uri
- elsif object.model_version_id
- model_version_uri
- else
- ml_model_candidate_uri
- end
-
- expose_url(uri)
+ if object.package&.generic?
+ expose_url(generic_package_uri)
+ elsif object.model_version_id
+ expose_url(model_version_uri)
+ elsif object.package&.version&.start_with?('candidate_')
+ "#{MLFLOW_ARTIFACTS_PREFIX}:/#{CANDIDATE_PREFIX}#{object.iid}"
+ else
+ expose_url(ml_model_candidate_uri)
+ end
end
# Example: http://127.0.0.1:3000/api/v4/projects/20/packages/ml_models/1/files/
diff --git a/lib/api/ml/mlflow/api_helpers.rb b/lib/api/ml/mlflow/api_helpers.rb
index 255388cac80b036ddb6849857165644ef5daa479..ad50bf99c9e0788c7beecd2af4a13b5216636f81 100644
--- a/lib/api/ml/mlflow/api_helpers.rb
+++ b/lib/api/ml/mlflow/api_helpers.rb
@@ -156,6 +156,11 @@ def find_model_artifact(project, version, file_path)
::Packages::PackageFileFinder.new(package, file_path).execute || resource_not_found!
end
+ def find_run_artifact(project, version, file_path)
+ package = ::Ml::Candidate.with_project_id_and_iid(project, version).package
+ ::Packages::PackageFileFinder.new(package, file_path).execute || resource_not_found!
+ end
+
def list_model_artifacts(project, version)
model_version = ::Ml::ModelVersion.by_project_id_and_id(project, version)
resource_not_found! unless model_version && model_version.package
@@ -163,9 +168,23 @@ def list_model_artifacts(project, version)
model_version.package.installable_package_files
end
+ def list_run_artifacts(project, version)
+ run = ::Ml::Candidate.with_project_id_and_iid(project, version)
+
+ resource_not_found! unless run&.package
+
+ run.package.installable_package_files
+ end
+
def model
@model ||= find_model(user_project, params[:name])
end
+
+ def candidate_version?(model_version)
+ return false unless model_version
+
+ model_version&.start_with?('candidate:')
+ end
end
end
end
diff --git a/lib/api/ml/mlflow_artifacts/artifacts.rb b/lib/api/ml/mlflow_artifacts/artifacts.rb
index ff16e39fc2c6650f9af22fc25f1ea0404f314759..f7184bb4a6128a3b97e94ced35090b930ecb52e1 100644
--- a/lib/api/ml/mlflow_artifacts/artifacts.rb
+++ b/lib/api/ml/mlflow_artifacts/artifacts.rb
@@ -6,6 +6,8 @@ module API
# MLFlow integration API, replicating the Rest API https://www.mlflow.org/docs/latest/rest-api.html#rest-api
module Ml
module MlflowArtifacts
+ CANDIDATE_PREFIX = 'candidate:'
+
class Artifacts < ::API::Base
feature_category :mlops
helpers ::API::Helpers::PackagesHelpers
@@ -30,15 +32,25 @@ class Artifacts < ::API::Base
# MLflow handles directories differently than GitLab does so when MLflow checks if a path is a directory
# we return an empty array as 404s would cause issues for MLflow
- files = path.present? ? [] : list_model_artifacts(user_project, model_version).all
+ files = if candidate_version?(model_version)
+ run_version = model_version.delete_prefix(CANDIDATE_PREFIX)
+ path.present? ? [] : list_run_artifacts(user_project, run_version).all
+ else
+ path.present? ? [] : list_model_artifacts(user_project, model_version).all
+ end
package_files = { files: files }
present package_files, with: Entities::Ml::MlflowArtifacts::ArtifactsList
end
get 'artifacts/:model_version/*file_path', format: false, urgency: :low do
- present_package_file!(find_model_artifact(user_project, params[:model_version],
- CGI.escape(params[:file_path])))
+ if candidate_version?(params[:model_version])
+ version = params[:model_version].delete_prefix(CANDIDATE_PREFIX)
+ present_package_file!(find_run_artifact(user_project, version, CGI.escape(params[:file_path])))
+ else
+ present_package_file!(find_model_artifact(user_project, params[:model_version],
+ CGI.escape(params[:file_path])))
+ end
end
end
end
diff --git a/spec/lib/api/entities/ml/mlflow/run_info_spec.rb b/spec/lib/api/entities/ml/mlflow/run_info_spec.rb
index 19011b19eaf9fdd117dac70eac842de520746ba9..49a235c25d566bf2d78b61fd8fc75b215a484173 100644
--- a/spec/lib/api/entities/ml/mlflow/run_info_spec.rb
+++ b/spec/lib/api/entities/ml/mlflow/run_info_spec.rb
@@ -88,6 +88,14 @@
expect(subject[:artifact_uri]).to eq("http://localhost/api/v4/projects/#{candidate.project_id}/packages/generic#{candidate.artifact_root}")
end
end
+
+ context 'when candidate has no file or generic package' do
+ let!(:candidate) { create(:ml_candidates, :with_ml_model, name: 'candidate_1') }
+
+ it 'returns a string with no package' do
+ expect(subject[:artifact_uri]).to eq("mlflow-artifacts:/candidate:#{candidate.iid}")
+ end
+ end
end
describe 'lifecycle_stage' do
diff --git a/spec/lib/api/ml/mlflow/api_helpers_spec.rb b/spec/lib/api/ml/mlflow/api_helpers_spec.rb
index f45fccfba4cf2bf0008b7118b3e6b3994b793e5c..faafc6f2dd08a0771226f0cc2df1203eefbabc25 100644
--- a/spec/lib/api/ml/mlflow/api_helpers_spec.rb
+++ b/spec/lib/api/ml/mlflow/api_helpers_spec.rb
@@ -123,4 +123,52 @@
end
end
end
+
+ describe '#icandidate_version?' do
+ describe 'when version is nil' do
+ let(:version) { nil }
+
+ it 'returns false' do
+ expect(candidate_version?(version)).to be false
+ end
+ end
+
+ describe 'when version has candidate prefix' do
+ let(:version) { 'candidate:1' }
+
+ it 'returns true' do
+ expect(candidate_version?(version)).to be true
+ end
+ end
+
+ describe 'when version does not have candidate prefix' do
+ let(:version) { '1' }
+
+ it 'returns false' do
+ expect(candidate_version?(version)).to be false
+ end
+ end
+ end
+
+ describe '#find_run_artifact' do
+ let_it_be(:project) { create(:project) }
+ let_it_be(:candidate) { create(:ml_candidates, :with_ml_model, project: project) }
+ let_it_be(:candidate_package_file) { create(:package_file, :ml_model, package: candidate.package) }
+
+ it 'returns list of files' do
+ expect(find_run_artifact(project, candidate.iid, candidate_package_file.file_name)).to eq candidate_package_file
+ end
+ end
+
+ describe '#list_run_artifacts' do
+ let_it_be(:project) { create(:project) }
+ let_it_be(:candidate) { create(:ml_candidates, :with_ml_model, project: project) }
+ let_it_be(:candidate_package_file) { create(:package_file, :ml_model, package: candidate.package) }
+ let_it_be(:candidate_package_file_2) { create(:package_file, :ml_model, package: candidate.package) }
+
+ it 'returns list of files' do
+ expect(list_run_artifacts(project,
+ candidate.iid)).to match_array [candidate_package_file, candidate_package_file_2]
+ end
+ end
end
diff --git a/spec/requests/api/ml/mlflow_artifacts/artifacts_spec.rb b/spec/requests/api/ml/mlflow_artifacts/artifacts_spec.rb
index 7de6ce2e15df7208e574f48dc95cdc669069ec31..8ac8a3414c510aa41a76592ee370ca49f7984b7e 100644
--- a/spec/requests/api/ml/mlflow_artifacts/artifacts_spec.rb
+++ b/spec/requests/api/ml/mlflow_artifacts/artifacts_spec.rb
@@ -13,6 +13,8 @@
let_it_be(:model_version) { create(:ml_model_versions, :with_package, model: model, version: version) }
let_it_be(:package_file) { create(:package_file, :ml_model, package: model_version.package) }
let_it_be(:model_version_no_package) { create(:ml_model_versions, model: model, version: '0.0.2') }
+ let_it_be(:candidate) { create(:ml_candidates, :with_ml_model, project: project) }
+ let_it_be(:candidate_package_file) { create(:package_file, :ml_model, package: candidate.package) }
let_it_be(:tokens) do
{
@@ -89,6 +91,19 @@
end
end
+ context 'when the model version is a candidate version' do
+ let(:route) do
+ "/projects/#{project_id}/ml/mlflow/api/2.0/mlflow-artifacts/artifacts?path=candidate:#{candidate.iid}/MlModel"
+ end
+
+ it 'returns an empty list of artifacts', :aggregate_failures do
+ is_expected.to have_gitlab_http_status(:ok)
+ expect(json_response).to have_key('files')
+ expect(json_response['files']).to be_an_instance_of(Array)
+ expect(json_response['files']).to be_empty
+ end
+ end
+
it_behaves_like 'MLflow|an authenticated resource'
it_behaves_like 'MLflow|a read-only model registry resource'
end
@@ -109,6 +124,22 @@
end
end
+ context 'when the model version is a candidate' do
+ let_it_be(:file) { candidate_package_file.file_name }
+ let(:route) do
+ "/projects/#{project_id}/ml/mlflow/api/2.0/mlflow-artifacts/artifacts/candidate:#{candidate.iid}/#{file}"
+ end
+
+ it 'returns the artifact file', :aggregate_failures do
+ is_expected.to have_gitlab_http_status(:ok)
+ expect(response.headers['Content-Disposition']).to match(
+ "attachment; filename=\"#{candidate_package_file.file_name}\""
+ )
+ expect(response.body).to eq(candidate_package_file.file.read)
+ expect(response.headers['Content-Length']).to eq(candidate_package_file.size.to_s)
+ end
+ end
+
context 'when the file does not exist' do
let(:file_path) { 'non_existent_file.txt' }