fix: fix regression caused by join on trials view

Author: salonig23

master/internal/api_experiment_intg_test.go

@@ -516,14 +516,25 @@ func TestPutExperimentRetainLogs(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, res)
 
-	var logRetentionDays []int
+	var taskLogRetentionDays []int
 	err = db.Bun().NewSelect().Table("tasks").
 		Column("log_retention_days").
 		Where("task_id IN (?)", bun.In(taskIDs)).
-		Scan(ctx, &logRetentionDays)
+		Scan(ctx, &taskLogRetentionDays)
 	require.NoError(t, err)
 
-	for _, v := range logRetentionDays {
+	for _, v := range taskLogRetentionDays {
+		require.Equal(t, v, numDays)
+	}
+
+	var trialLogRetentionDays []int
+	err = db.Bun().NewSelect().Table("runs").
+		Column("log_retention_days").
+		Where("id IN (?)", bun.In(trialIDs)).
+		Scan(ctx, &trialLogRetentionDays)
+	require.NoError(t, err)
+
+	for _, v := range trialLogRetentionDays {
 		require.Equal(t, v, numDays)
 	}
 }
@@ -561,14 +572,25 @@ func TestPutExperimentsRetainLogs(t *testing.T) {
 	_, taskIDs, err := db.ExperimentsTrialAndTaskIDs(ctx, db.Bun(), intExpIDS)
 	require.NoError(t, err)
 
-	var logRetentionDays []int
+	var taskLogRetentionDays []int
 	err = db.Bun().NewSelect().Table("tasks").
 		Column("log_retention_days").
 		Where("task_id IN (?)", bun.In(taskIDs)).
-		Scan(ctx, &logRetentionDays)
+		Scan(ctx, &taskLogRetentionDays)
+	require.NoError(t, err)
+
+	for _, v := range taskLogRetentionDays {
+		require.Equal(t, v, numDays)
+	}
+
+	var trialLogRetentionDays []int
+	err = db.Bun().NewSelect().Table("runs").
+		Column("log_retention_days").
+		Where("id IN (?)", bun.In(trialIDs)).
+		Scan(ctx, &trialLogRetentionDays)
 	require.NoError(t, err)
 
-	for _, v := range logRetentionDays {
+	for _, v := range trialLogRetentionDays {
 		require.Equal(t, v, numDays)
 	}
 }

master/internal/api_trials.go

@@ -717,6 +717,12 @@ func (a *apiServer) PutTrialRetainLogs(
 								Exec(ctx); err != nil {
 			return fmt.Errorf("updating log retention days for tasks: %w", err)
 		}
+		if _, err := tx.NewUpdate().Table("runs").
+			Set("log_retention_days = ?", req.NumDays).
+			Where("id = ?", req.TrialId).
+			Exec(ctx); err != nil {
+			return fmt.Errorf("updating log retention days for trial: %w", err)
+		}
 		return nil
 	})
 	if err != nil {

master/internal/experiment/bulk_action.go

@@ -807,7 +807,7 @@ func BulkUpdateLogRentention(ctx context.Context, database db.DB,
 		}
 	}
 
-	_, taskIDs, err := db.ExperimentsTrialAndTaskIDs(ctx, db.Bun(), intExpIDs)
+	trialIDs, taskIDs, err := db.ExperimentsTrialAndTaskIDs(ctx, db.Bun(), intExpIDs)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to gather trial IDs for experiments")
 	}
@@ -823,8 +823,15 @@ func BulkUpdateLogRentention(ctx context.Context, database db.DB,
 			Exec(ctx); err != nil {
 			return fmt.Errorf("updating log retention days for tasks: %w", err)
 		}
+		if _, err := tx.NewUpdate().Table("runs").
+			Set("log_retention_days = ?", numDays).
+			Where("id IN (?)", bun.In(trialIDs)).
+			Exec(ctx); err != nil {
+			return fmt.Errorf("updating log retention days for tasks: %w", err)
+		}
 		return nil
 	})
+
 	if err != nil {
 		return nil, err
 	}

master/pkg/model/experiment.go

@@ -479,6 +479,7 @@ func (t *Trial) ToRunAndTrialV2(experimentsProjectID int) (*Run, *TrialV2) {
 		Restarts:              t.Restarts,
 		RunnerState:           t.RunnerState,
 		LastActivity:          t.LastActivity,
+		LogRetentionDays:      t.LogRetentionDays,
 	}
 	v2 := &TrialV2{
 		RunID:     t.ID,
@@ -516,6 +517,7 @@ type Run struct {
 	Restarts              int            `db:"restarts"`
 	RunnerState           string         `db:"runner_state"`
 	LastActivity          *time.Time     `db:"last_activity"`
+	LogRetentionDays      *int16         `db:"log_retention_days"`
 }
 
 // RunTaskID represents a row from the `run_id_task_id` table.

master/static/migrations/20240403070639_fix-trials-view.tx.down.sql

@@ -0,0 +1,58 @@
+ALTER TABLE RUNS DROP COLUMN log_retention_days;
+
+DROP VIEW trials;
+CREATE VIEW trials AS
+WITH task_log_retention AS (
+  SELECT  
+    r.run_id, 
+    -- This is written with the assumption that every task related 
+    -- to a trial will have the same log_retention_days. MIN() is 
+    -- just used to aggregate the number of days for each trial.
+    MIN(t.log_retention_days) as log_retention_days
+  FROM tasks t
+  JOIN run_id_task_id as r ON t.task_id = r.task_id
+  GROUP BY r.run_id
+)
+SELECT
+  t.run_id AS id,
+
+  -- metrics
+  r.summary_metrics AS summary_metrics,
+  r.summary_metrics_timestamp AS summary_metrics_timestamp,
+  r.latest_validation_id AS latest_validation_id,
+  r.total_batches AS total_batches,
+
+  -- metadata fields
+  r.state AS state,
+  r.tags AS tags,
+  r.external_run_id AS external_trial_id,
+  r.restart_id AS run_id,
+  r.last_activity AS last_activity,
+  r.start_time AS start_time,
+  r.end_time AS end_time,
+  r.restarts AS restarts,
+  l.log_retention_days AS log_retention_days,
+
+  -- run_hp_search_stuff
+  r.hparams AS hparams,
+  r.searcher_metric_value AS searcher_metric_value,
+  r.searcher_metric_value_signed AS searcher_metric_value_signed,
+  r.best_validation_id AS best_validation_id,
+
+  -- run_checkpoint_stats
+  r.checkpoint_size AS checkpoint_size,
+  r.checkpoint_count AS checkpoint_count,
+
+  -- trial_v2 table.
+  t.request_id AS request_id,
+  t.seed AS seed,
+
+  r.experiment_id AS experiment_id,
+  r.warm_start_checkpoint_id AS warm_start_checkpoint_id,
+
+  -- eventually delete runner state.
+  r.runner_state AS runner_state
+FROM trials_v2 t
+JOIN runs r ON t.run_id = r.id
+JOIN task_log_retention l ON l.run_id = t.run_id;
+

master/static/migrations/20240403070639_fix-trials-view.tx.up.sql

@@ -0,0 +1,46 @@
+ALTER TABLE runs ADD COLUMN log_retention_days SMALLINT;
+
+DROP VIEW trials;
+
+CREATE VIEW trials AS
+SELECT
+  t.run_id AS id,
+
+  -- metrics
+  r.summary_metrics AS summary_metrics,
+  r.summary_metrics_timestamp AS summary_metrics_timestamp,
+  r.latest_validation_id AS latest_validation_id,
+  r.total_batches AS total_batches,
+
+  -- metadata fields
+  r.state AS state,
+  r.tags AS tags,
+  r.external_run_id AS external_trial_id,
+  r.restart_id AS run_id,
+  r.last_activity AS last_activity,
+  r.start_time AS start_time,
+  r.end_time AS end_time,
+  r.restarts AS restarts,
+  r.log_retention_days AS log_retention_days,
+  
+  -- run_hp_search_stuff
+  r.hparams AS hparams,
+  r.searcher_metric_value AS searcher_metric_value,
+  r.searcher_metric_value_signed AS searcher_metric_value_signed,
+  r.best_validation_id AS best_validation_id,
+
+  -- run_checkpoint_stats
+  r.checkpoint_size AS checkpoint_size,
+  r.checkpoint_count AS checkpoint_count,
+
+  -- trial_v2 table.
+  t.request_id AS request_id,
+  t.seed AS seed,
+
+  r.experiment_id AS experiment_id,
+  r.warm_start_checkpoint_id AS warm_start_checkpoint_id,
+
+  -- eventually delete runner state.
+  r.runner_state AS runner_state
+FROM trials_v2 t
+JOIN runs r ON t.run_id = r.id;

proto/pkg/runv1/run.pb.go

@@ -111,4 +111,6 @@ message FlatRun {
   bool parent_archived = 18;
   // Data related the the experiment associated with this run.
   optional FlatRunExperiment experiment = 19;
+  // Number of days to retain logs for.
+  optional int32 log_retention_days = 20;
 }