Move GET /mlflow/metrics/get-history endpoint. (#85)

Signed-off-by: Software Developer <7852635+dsuhinin@users.noreply.github.com>
mlflow · Nov 12, 2024 · 5733c95 · 5733c95
1 parent 6070234
commit 5733c95
Show file tree

Hide file tree

Showing 10 changed files with 174 additions and 12 deletions.
diff --git a/magefiles/generate/endpoints.go b/magefiles/generate/endpoints.go
@@ -32,7 +32,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			"deleteTag",
 			"searchRuns",
 			// "listArtifacts",
-			// "getMetricHistory",
+			"getMetricHistory",
 			// "getMetricHistoryBulkInterval",
 			"logBatch",
 			// "logModel",

diff --git a/mlflow_go/store/tracking.py b/mlflow_go/store/tracking.py
@@ -4,6 +4,7 @@
 
 from mlflow.entities import (
     Experiment,
+    Metric,
     Run,
     RunInfo,
     TraceInfo,
@@ -24,6 +25,7 @@
     EndTrace,
     GetExperiment,
     GetExperimentByName,
+    GetMetricHistory,
     GetRun,
     GetTraceInfo,
     LogBatch,
@@ -311,6 +313,13 @@ def delete_traces(
         response = self.service.call_endpoint(get_lib().TrackingServiceDeleteTraces, request)
         return response.traces_deleted
 
+    def get_metric_history(self, run_id, metric_key, max_results=None, page_token=None):
+        request = GetMetricHistory(
+            run_id=run_id, metric_key=metric_key, max_results=max_results, page_token=page_token
+        )
+        response = self.service.call_endpoint(get_lib().TrackingServiceGetMetricHistory, request)
+        return PagedList([Metric.from_proto(metric) for metric in response.metrics], None)
+
 
 def TrackingStore(cls):
     return type(cls.__name__, (_TrackingStore, cls), {})

diff --git a/pkg/contract/service/tracking.g.go b/pkg/contract/service/tracking.g.go
diff --git a/pkg/lib/tracking.g.go b/pkg/lib/tracking.g.go
diff --git a/pkg/server/routes/tracking.g.go b/pkg/server/routes/tracking.g.go
diff --git a/pkg/tracking/service/metrics.go b/pkg/tracking/service/metrics.go
@@ -2,6 +2,7 @@ package service
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
 	"github.com/mlflow/mlflow-go/pkg/entities"
@@ -28,3 +29,38 @@ func (ts TrackingService) LogParam(
 
 	return &protos.LogParam_Response{}, nil
 }
+
+func (ts TrackingService) GetMetricHistory(
+	ctx context.Context, input *protos.GetMetricHistory,
+) (*protos.GetMetricHistory_Response, *contract.Error) {
+	if input.PageToken != nil {
+		//nolint:lll
+		return nil, contract.NewError(
+			protos.ErrorCode_INVALID_PARAMETER_VALUE,
+			fmt.Sprintf(
+				"The SQLAlchemyStore backend does not support pagination for the `get_metric_history` API. Supplied argument `page_token` '%s' must be `None`.",
+				*input.PageToken,
+			),
+		)
+	}
+
+	runID := input.GetRunId()
+	if input.RunUuid != nil {
+		runID = input.GetRunUuid()
+	}
+
+	metrics, err := ts.Store.GetMetricHistory(ctx, runID, input.GetMetricKey())
+	if err != nil {
+		return nil, err
+	}
+
+	response := protos.GetMetricHistory_Response{
+		Metrics: make([]*protos.Metric, len(metrics)),
+	}
+
+	for i, metric := range metrics {
+		response.Metrics[i] = metric.ToProto()
+	}
+
+	return &response, nil
+}
diff --git a/pkg/tracking/store/mock_tracking_store.go b/pkg/tracking/store/mock_tracking_store.go
diff --git a/pkg/tracking/store/sql/metrics.go b/pkg/tracking/store/sql/metrics.go
@@ -191,3 +191,27 @@ func (s TrackingSQLStore) LogMetric(ctx context.Context, runID string, metric *e
 
 	return nil
 }
+
+func (s TrackingSQLStore) GetMetricHistory(
+	ctx context.Context, runID, metricKey string,
+) ([]*entities.Metric, *contract.Error) {
+	var metrics []*models.Metric
+	if err := s.db.WithContext(
+		ctx,
+	).Where(
+		"run_uuid = ?", runID,
+	).Where(
+		"key = ?", metricKey,
+	).Find(&metrics).Error; err != nil {
+		return nil, contract.NewError(
+			protos.ErrorCode_INTERNAL_ERROR, fmt.Sprintf("error getting metric history: %v", err),
+		)
+	}
+
+	entityMetrics := make([]*entities.Metric, len(metrics))
+	for i, metric := range metrics {
+		entityMetrics[i] = metric.ToEntity()
+	}
+
+	return entityMetrics, nil
+}
diff --git a/pkg/tracking/store/sql/models/metrics.go b/pkg/tracking/store/sql/models/metrics.go
@@ -45,6 +45,16 @@ func NewMetricFromEntity(runID string, metric *entities.Metric) *Metric {
 	return &model
 }
 
+func (m Metric) ToEntity() *entities.Metric {
+	return &entities.Metric{
+		Key:       m.Key,
+		Value:     m.Value,
+		Timestamp: m.Timestamp,
+		Step:      m.Step,
+		IsNaN:     m.IsNaN,
+	}
+}
+
 func (m Metric) NewLatestMetricFromProto() LatestMetric {
 	return LatestMetric{
 		RunID:     m.RunID,

diff --git a/pkg/tracking/store/store.go b/pkg/tracking/store/store.go
@@ -80,6 +80,7 @@ type (
 
 		LogMetric(ctx context.Context, runID string, metric *entities.Metric) *contract.Error
 		LogParam(ctx context.Context, runID string, metric *entities.Param) *contract.Error
+		GetMetricHistory(ctx context.Context, runID, metricKey string) ([]*entities.Metric, *contract.Error)
 	}
 
 	ExperimentTrackingStore interface {