5 files changed, 109 insertions, 0 deletions
diff --git a/lib/api/entities/ml/mlflow/experiment.rb b/lib/api/entities/ml/mlflow/experiment.rb
new file mode 100644
index 00000000000..cfe366feaab
--- /dev/null
+++ b/lib/api/entities/ml/mlflow/experiment.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module API
+  module Entities
+    module Ml
+      module Mlflow
+        class Experiment < Grape::Entity
+          expose :experiment do
+            expose :experiment_id
+            expose :name
+            expose :lifecycle_stage
+            expose :artifact_location
+          end
+
+          private
+
+          def lifecycle_stage
+            object.deleted_on? ? 'deleted' : 'active'
+          end
+
+          def experiment_id
+            object.iid.to_s
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/api/entities/ml/mlflow/new_experiment.rb b/lib/api/entities/ml/mlflow/new_experiment.rb
new file mode 100644
index 00000000000..09791839850
--- /dev/null
+++ b/lib/api/entities/ml/mlflow/new_experiment.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module API
+  module Entities
+    module Ml
+      module Mlflow
+        class NewExperiment < Grape::Entity
+          expose :experiment_id
+
+          private
+
+          def experiment_id
+            object.iid.to_s
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/api/entities/ml/mlflow/run.rb b/lib/api/entities/ml/mlflow/run.rb
new file mode 100644
index 00000000000..c679330206e
--- /dev/null
+++ b/lib/api/entities/ml/mlflow/run.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module API
+  module Entities
+    module Ml
+      module Mlflow
+        class Run < Grape::Entity
+          expose :run do
+            expose(:info) { |candidate| RunInfo.represent(candidate) }
+            expose(:data) { |candidate| {} }
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/api/entities/ml/mlflow/run_info.rb b/lib/api/entities/ml/mlflow/run_info.rb
new file mode 100644
index 00000000000..096950e349d
--- /dev/null
+++ b/lib/api/entities/ml/mlflow/run_info.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module API
+  module Entities
+    module Ml
+      module Mlflow
+        class RunInfo < Grape::Entity
+          expose :run_id
+          expose :run_id, as: :run_uuid
+          expose(:experiment_id) { |candidate| candidate.experiment.iid.to_s }
+          expose(:start_time) { |candidate| candidate.start_time || 0 }
+          expose :end_time, expose_nil: false
+          expose(:status) { |candidate| candidate.status.to_s.upcase }
+          expose(:artifact_uri) { |candidate| 'not_implemented' }
+          expose(:lifecycle_stage) { |candidate| 'active' }
+          expose(:user_id) { |candidate| candidate.user_id.to_s }
+
+          private
+
+          def run_id
+            object.iid.to_s
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/api/entities/ml/mlflow/update_run.rb b/lib/api/entities/ml/mlflow/update_run.rb
new file mode 100644
index 00000000000..5acdaab0e33
--- /dev/null
+++ b/lib/api/entities/ml/mlflow/update_run.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module API
+  module Entities
+    module Ml
+      module Mlflow
+        class UpdateRun < Grape::Entity
+          expose :run_info
+
+          private
+
+          def run_info
+            ::API::Entities::Ml::Mlflow::RunInfo.represent object
+          end
+        end
+      end
+    end
+  end
+end