Integrate context filtering logic into the Aim GetRunMetrics and SearchAlignedMetrics endpoints.

pkg/api/aim/runs.go

@@ -160,17 +160,18 @@ func GetRunMetrics(c *fiber.Ctx) error {
 		return fiber.NewError(fiber.StatusUnprocessableEntity, err.Error())
 	}
 
-	b := []struct {
-		Name    string    `json:"name"`
-		Context fiber.Map `json:"context"`
-	}{}
+	var b []struct {
+		Name    string            `json:"name"`
+		Context map[string]string `json:"context"`
+	}
 
 	if err := c.BodyParser(&b); err != nil {
 		return fiber.NewError(fiber.StatusUnprocessableEntity, err.Error())
 	}
 
-	metricKeysMap := make(fiber.Map, len(b))
+	metricKeysMap, contexts := make(fiber.Map, len(b)), make([]map[string]string, 0, len(b))
 	for _, m := range b {
+		contexts = append(contexts, m.Context)
 		metricKeysMap[m.Name] = nil
 	}
 	metricKeys := make([]string, len(metricKeysMap))
@@ -181,39 +182,57 @@ func GetRunMetrics(c *fiber.Ctx) error {
 		i++
 	}
 
-	r := database.Run{
+	// check that requested run actually exists.
+	if err := database.DB.Select(
+		"ID",
+	).InnerJoins(
+		"Experiment",
+		database.DB.Select(
+			"ID",
+		).Where(
+			&models.Experiment{NamespaceID: ns.ID},
+		),
+	).First(&database.Run{
 		ID: p.ID,
-	}
-	if err := database.DB.
-		Select("ID").
-		InnerJoins(
-			"Experiment",
-			database.DB.Select(
-				"ID",
-			).Where(
-				&models.Experiment{NamespaceID: ns.ID},
-			),
-		).
-		Preload("Metrics", func(db *gorm.DB) *gorm.DB {
-			return db.
-				Where("key IN ?", metricKeys).
-				Order("iter")
-		}).
-		Preload("Metrics.Context").
-		First(&r).Error; err != nil {
+	}).Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return fiber.ErrNotFound
 		}
 		return fmt.Errorf("unable to find run %q: %w", p.ID, err)
 	}
 
+	// fetch run metrics based on provided criteria.
+	var data []database.Metric
+	if err := database.DB.Where(
+		"run_uuid = ?", p.ID,
+	).InnerJoins(
+		"Context",
+		func() *gorm.DB {
+			subQuery := database.DB
+			for _, context := range contexts {
+				sql, args := repositories.BuildJsonCondition(database.DB.Dialector.Name(), "json", context)
+				subQuery = subQuery.Or(sql, args...)
+			}
+			return subQuery
+		}(),
+	).Where(
+		"key IN ?", metricKeys,
+	).Order(
+		"iter",
+	).Find(
+		&data,
+	).Error; err != nil {
+		return fmt.Errorf("unable to find run metrics: %w", err)
+	}
+
 	metrics := make(map[string]struct {
 		name    string
 		iters   []int
 		values  []*float64
 		context datatypes.JSON
 	}, len(metricKeys))
-	for _, m := range r.Metrics {
+
+	for _, m := range data {
 		v := m.Value
 		pv := &v
 		if m.IsNan {
@@ -231,13 +250,12 @@ func GetRunMetrics(c *fiber.Ctx) error {
 
 	resp := make([]fiber.Map, 0, len(metrics))
 	for _, m := range metrics {
-		data := fiber.Map{
+		resp = append(resp, fiber.Map{
 			"name":    m.name,
 			"iters":   m.iters,
 			"values":  m.values,
 			"context": m.context,
-		}
-		resp = append(resp, data)
+		})
 	}
 
 	return c.JSON(resp)
@@ -921,9 +939,9 @@ func SearchAlignedMetrics(c *fiber.Ctx) error {
 		Runs    []struct {
 			ID     string `json:"run_id"`
 			Traces []struct {
-				Name    string    `json:"name"`
-				Slice   [3]int    `json:"slice"`
-				Context fiber.Map `json:"context"`
+				Name    string            `json:"name"`
+				Slice   [3]int            `json:"slice"`
+				Context map[string]string `json:"context"`
 			} `json:"traces"`
 		} `json:"runs"`
 	}{}
@@ -932,15 +950,17 @@ func SearchAlignedMetrics(c *fiber.Ctx) error {
 		return fiber.NewError(fiber.StatusUnprocessableEntity, err.Error())
 	}
 
-	var capacity int
-	var values []any
+	values, capacity, contexts := []any{}, 0, []map[string]string{}
 	for _, r := range b.Runs {
 		for _, t := range r.Traces {
 			l := t.Slice[2]
 			if l > capacity {
 				capacity = l
 			}
 			values = append(values, r.ID, t.Name, float32(l))
+			if len(t.Context) > 0 {
+				contexts = append(contexts, t.Context)
+			}
 		}
 	}
 
@@ -953,9 +973,24 @@ func SearchAlignedMetrics(c *fiber.Ctx) error {
 		}
 	}
 
-	// TODO this should probably be batched
-
 	values = append(values, ns.ID, b.AlignBy)
+
+	// process context information.
+	var contextStmt strings.Builder
+	if len(contexts) > 0 {
+		for i, context := range contexts {
+			sql, args := repositories.BuildJsonCondition(database.DB.Dialector.Name(), "c.json", context)
+			contextStmt.WriteString(sql)
+			if i < len(contexts)-1 {
+				contextStmt.WriteString(" OR ")
+			}
+			values = append(values, args...)
+		}
+	} else {
+		contextStmt.WriteString(" 1=1 ")
+	}
+
+	// TODO this should probably be batched
 	rows, err := database.DB.Raw(
 		fmt.Sprintf("WITH params(run_uuid, key, steps) AS (VALUES %s)", &valuesStmt)+
 			"        SELECT m.run_uuid, rm.key, m.iter, m.value, m.is_nan, c.json AS context_json FROM metrics AS m"+
@@ -970,6 +1005,7 @@ func SearchAlignedMetrics(c *fiber.Ctx) error {
 			"        WHERE m.key = ?"+
 			"          AND m.iter <= rm.max"+
 			"          AND MOD(m.iter + 1 + rm.interval / 2, rm.interval) < 1"+
+			"		   AND ("+contextStmt.String()+")"+
 			"        ORDER BY m.run_uuid, rm.key, m.iter",
 		values...,
 	).Rows()

tests/integration/golang/aim/metric/search_aligned_test.go

@@ -70,6 +70,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run1.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key11": "value11"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric1Run1, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -90,6 +93,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run1.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key12": "value12"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric2Run1, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -110,6 +116,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run1.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key13": "value13"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric3Run1, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -149,6 +158,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run2.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key21": "value21"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric1Run2, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -169,6 +181,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run2.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key22": "value22"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric2Run2, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -189,6 +204,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run2.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key23": "value23"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric3Run2, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -228,6 +246,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run3.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key31": "value31"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric1Run3, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -248,6 +269,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run3.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key32": "value32"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric2Run3, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -268,6 +292,9 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		IsNan:     false,
 		RunID:     run3.ID,
 		Iter:      1,
+		Context: models.Context{
+			Json: []byte(`{"key33": "value33"}`),
+		},
 	})
 	s.Require().Nil(err)
 	metric3Run3, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{
@@ -288,6 +315,93 @@ func (s *SearchAlignedMetricsTestSuite) Test_Ok() {
 		request  *request.GetAlignedMetricRequest
 		response []float64
 	}{
+		{
+			name: "TestSearchAlignedByMetric1AndContext",
+			request: &request.GetAlignedMetricRequest{
+				AlignBy: "TestMetric1",
+				Runs: []request.AlignedMetricRunRequest{
+					{
+						ID: run1.ID,
+						Traces: []request.AlignedMetricTraceRequest{
+							{
+								Name:  "TestMetric1",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key11": "value11",
+								},
+							},
+							{
+								Name:  "TestMetric2",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key11": "value11",
+								},
+							},
+							{
+								Name:  "TestMetric3",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key11": "value11",
+								},
+							},
+						},
+					},
+					{
+						ID: run2.ID,
+						Traces: []request.AlignedMetricTraceRequest{
+							{
+								Name:  "TestMetric1",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key21": "value21",
+								},
+							},
+							{
+								Name:  "TestMetric2",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key21": "value21",
+								},
+							},
+							{
+								Name:  "TestMetric3",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key21": "value21",
+								},
+							},
+						},
+					},
+					{
+						ID: run3.ID,
+						Traces: []request.AlignedMetricTraceRequest{
+							{
+								Name:  "TestMetric1",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key31": "value31",
+								},
+							},
+							{
+								Name:  "TestMetric2",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key31": "value31",
+								},
+							},
+							{
+								Name:  "TestMetric3",
+								Slice: []int{0, 0, 500},
+								Context: map[string]interface{}{
+									"key31": "value31",
+								},
+							},
+						},
+					},
+				},
+			},
+			response: []float64{metric1Run1.Value, metric1Run2.Value, metric1Run3.Value},
+		},
 		{
 			name: "TestSearchAlignedByMetric1",
 			request: &request.GetAlignedMetricRequest{