chore: bunify db/postgres_tasks.go

master/internal/api_checkpoint_intg_test.go

@@ -42,7 +42,7 @@ func createVersionTwoCheckpoint(
 		ResourcePool: "somethingelse",
 		StartTime:    ptrs.Ptr(time.Now().UTC().Truncate(time.Millisecond)),
 	}
-	require.NoError(t, api.m.db.AddAllocation(aIn))
+	require.NoError(t, db.AddAllocation(ctx, aIn))
 
 	checkpoint := &model.CheckpointV2{
 		ID:           0,

master/internal/api_experiment.go

@@ -2737,7 +2737,7 @@ func (a *apiServer) createTrialTx(
 		nil,
 		0)
 
-	if err := a.m.db.AddTask(&model.Task{
+	if err := db.AddTask(ctx, &model.Task{
 		TaskID:     taskID,
 		TaskType:   model.TaskTypeTrial,
 		StartTime:  time.Now(),

master/internal/api_experiment_intg_test.go

@@ -209,7 +209,7 @@ func TestGetTaskContextDirectoryExperiment(t *testing.T) {
 func TestGetTaskContextDirectoryTask(t *testing.T) {
 	api, _, ctx := setupAPITest(t, nil)
 	task := &model.Task{TaskType: model.TaskTypeNotebook, TaskID: model.NewTaskID()}
-	require.NoError(t, api.m.db.AddTask(task))
+	require.NoError(t, db.AddTask(ctx, task))
 
 	expectedContextDirectory := []byte("expectedContextDirectory")
 	_, err := db.Bun().NewInsert().Model(&model.TaskContextDirectory{
@@ -539,7 +539,7 @@ func TestGetExperiments(t *testing.T) {
 	require.NoError(t, api.m.db.AddExperiment(exp0, activeConfig0))
 	for i := 0; i < 3; i++ {
 		task := &model.Task{TaskType: model.TaskTypeTrial, TaskID: model.NewTaskID()}
-		require.NoError(t, api.m.db.AddTask(task))
+		require.NoError(t, db.AddTask(ctx, task))
 		require.NoError(t, db.AddTrial(ctx, &model.Trial{
 			State:        model.PausedState,
 			ExperimentID: exp0.ID,
@@ -791,7 +791,7 @@ func TestSearchExperiments(t *testing.T) {
 	// Trial without validations doesn't cause issues.
 	noValidationsExp := createTestExpWithProjectID(t, api, curUser, projectIDInt)
 	task := &model.Task{TaskType: model.TaskTypeTrial, TaskID: model.NewTaskID()}
-	require.NoError(t, api.m.db.AddTask(task))
+	require.NoError(t, db.AddTask(ctx, task))
 	require.NoError(t, db.AddTrial(ctx, &model.Trial{
 		State:        model.PausedState,
 		ExperimentID: noValidationsExp.ID,

master/internal/api_tasks_intg_test.go

@@ -113,10 +113,10 @@ func mockNotebookWithWorkspaceID(
 		TaskID:   model.NewTaskID(),
 		TaskType: model.TaskTypeNotebook,
 	}
-	require.NoError(t, api.m.db.AddTask(nb))
+	require.NoError(t, db.AddTask(ctx, nb))
 
 	allocationID := model.AllocationID(string(nb.TaskID) + ".1")
-	require.NoError(t, api.m.db.AddAllocation(&model.Allocation{
+	require.NoError(t, db.AddAllocation(ctx, &model.Allocation{
 		TaskID:       nb.TaskID,
 		AllocationID: allocationID,
 	}))
@@ -324,7 +324,7 @@ func TestAddAllocationAcceleratorData(t *testing.T) {
 		TaskType:  model.TaskTypeTrial,
 		StartTime: time.Now().UTC().Truncate(time.Millisecond),
 	}
-	require.NoError(t, api.m.db.AddTask(task), "failed to add task")
+	require.NoError(t, db.AddTask(ctx, task), "failed to add task")
 
 	aID := tID + "-1"
 	a := &model.Allocation{
@@ -333,7 +333,7 @@ func TestAddAllocationAcceleratorData(t *testing.T) {
 		Slots:        1,
 		ResourcePool: "default",
 	}
-	require.NoError(t, api.m.db.AddAllocation(a), "failed to add allocation")
+	require.NoError(t, db.AddAllocation(ctx, a), "failed to add allocation")
 	accData := &model.AcceleratorData{
 		ContainerID:      uuid.NewString(),
 		AllocationID:     model.AllocationID(aID),
@@ -362,7 +362,7 @@ func TestGetAllocationAcceleratorDataWithNoData(t *testing.T) {
 		TaskType:  model.TaskTypeTrial,
 		StartTime: time.Now().UTC().Truncate(time.Millisecond),
 	}
-	require.NoError(t, api.m.db.AddTask(task), "failed to add task")
+	require.NoError(t, db.AddTask(ctx, task), "failed to add task")
 
 	aID := tID + "-1"
 	a := &model.Allocation{
@@ -371,7 +371,7 @@ func TestGetAllocationAcceleratorDataWithNoData(t *testing.T) {
 		Slots:        1,
 		ResourcePool: "default",
 	}
-	require.NoError(t, api.m.db.AddAllocation(a), "failed to add allocation")
+	require.NoError(t, db.AddAllocation(ctx, a), "failed to add allocation")
 
 	resp, err := api.GetTaskAcceleratorData(ctx,
 		&apiv1.GetTaskAcceleratorDataRequest{TaskId: tID.String()})
@@ -390,7 +390,7 @@ func TestGetAllocationAcceleratorData(t *testing.T) {
 		TaskType:  model.TaskTypeTrial,
 		StartTime: time.Now().UTC().Truncate(time.Millisecond),
 	}
-	require.NoError(t, api.m.db.AddTask(task), "failed to add task")
+	require.NoError(t, db.AddTask(ctx, task), "failed to add task")
 
 	aID1 := tID + "-1"
 	a1 := &model.Allocation{
@@ -399,7 +399,7 @@ func TestGetAllocationAcceleratorData(t *testing.T) {
 		Slots:        1,
 		ResourcePool: "default",
 	}
-	require.NoError(t, api.m.db.AddAllocation(a1), "failed to add allocation")
+	require.NoError(t, db.AddAllocation(ctx, a1), "failed to add allocation")
 	accData := &model.AcceleratorData{
 		ContainerID:      uuid.NewString(),
 		AllocationID:     model.AllocationID(aID1),
@@ -418,7 +418,7 @@ func TestGetAllocationAcceleratorData(t *testing.T) {
 		Slots:        1,
 		ResourcePool: "default",
 	}
-	require.NoError(t, api.m.db.AddAllocation(a2), "failed to add allocation")
+	require.NoError(t, db.AddAllocation(ctx, a2), "failed to add allocation")
 
 	resp, err := api.GetTaskAcceleratorData(ctx,
 		&apiv1.GetTaskAcceleratorDataRequest{TaskId: tID.String()})

master/internal/api_trials_intg_test.go

@@ -38,6 +38,7 @@ var inferenceMetricGroup = "inference"
 func createTestTrial(
 	t *testing.T, api *apiServer, curUser model.User,
 ) (*model.Trial, *model.Task) {
+	ctx := context.Background()
 	exp := createTestExpWithProjectID(t, api, curUser, 1)
 
 	task := &model.Task{
@@ -46,17 +47,17 @@ func createTestTrial(
 		StartTime:  time.Now(),
 		TaskID:     trialTaskID(exp.ID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task))
+	require.NoError(t, db.AddTask(ctx, task))
 
 	trial := &model.Trial{
 		StartTime:    time.Now(),
 		State:        model.PausedState,
 		ExperimentID: exp.ID,
 	}
-	require.NoError(t, db.AddTrial(context.TODO(), trial, task.TaskID))
+	require.NoError(t, db.AddTrial(ctx, trial, task.TaskID))
 
 	// Return trial exactly the way the API will generally get it.
-	outTrial, err := db.TrialByID(context.TODO(), trial.ID)
+	outTrial, err := db.TrialByID(ctx, trial.ID)
 	require.NoError(t, err)
 	return outTrial, task
 }
@@ -751,15 +752,15 @@ func TestTrialProtoTaskIDs(t *testing.T) {
 		StartTime:  task0.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task1))
+	require.NoError(t, db.AddTask(ctx, task1))
 
 	task2 := &model.Task{
 		TaskType:   model.TaskTypeTrial,
 		LogVersion: model.TaskLogVersion1,
 		StartTime:  task1.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task2))
+	require.NoError(t, db.AddTask(ctx, task2))
 
 	_, err = db.Bun().NewInsert().Model(&[]model.TrialTaskID{
 		{TrialID: trial.ID, TaskID: task1.TaskID},
@@ -834,7 +835,7 @@ func TestExperimentIDFromTrialTaskID(t *testing.T) {
 		StartTime:  time.Now(),
 		TaskID:     model.TaskID(uuid.New().String()),
 	}
-	require.NoError(t, api.m.db.AddTask(task))
+	require.NoError(t, db.AddTask(context.Background(), task))
 	_, err = experimentIDFromTrialTaskID(notTrialTask.TaskID)
 	require.ErrorIs(t, err, errIsNotTrialTaskID)
 
@@ -852,7 +853,7 @@ func TestTrialLogsBackported(t *testing.T) {
 		StartTime:  time.Now(),
 		TaskID:     model.TaskID(fmt.Sprintf("backported.%d", exp.ID)),
 	}
-	require.NoError(t, api.m.db.AddTask(task))
+	require.NoError(t, db.AddTask(ctx, task))
 
 	trial := &model.Trial{
 		StartTime:    time.Now(),
@@ -889,15 +890,15 @@ func TestTrialLogs(t *testing.T) {
 		StartTime:  task0.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task1))
+	require.NoError(t, db.AddTask(ctx, task1))
 
 	task2 := &model.Task{
 		TaskType:   model.TaskTypeTrial,
 		LogVersion: model.TaskLogVersion1,
 		StartTime:  task1.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task2))
+	require.NoError(t, db.AddTask(ctx, task2))
 
 	_, err := db.Bun().NewInsert().Model(&[]model.TrialTaskID{
 		{TrialID: trial.ID, TaskID: task1.TaskID},
@@ -985,15 +986,15 @@ func TestTrialLogFields(t *testing.T) {
 		StartTime:  task0.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task1))
+	require.NoError(t, db.AddTask(ctx, task1))
 
 	task2 := &model.Task{
 		TaskType:   model.TaskTypeTrial,
 		LogVersion: model.TaskLogVersion1,
 		StartTime:  task1.StartTime.Add(time.Second),
 		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
 	}
-	require.NoError(t, api.m.db.AddTask(task2))
+	require.NoError(t, db.AddTask(ctx, task2))
 
 	_, err := db.Bun().NewInsert().Model(&[]model.TrialTaskID{
 		{TrialID: trial.ID, TaskID: task1.TaskID},

master/internal/checkpoint_gc.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -26,7 +27,7 @@ const fullDeleteGlob = "**/*"
 
 func runCheckpointGCTask(
 	rm rm.ResourceManager,
-	db *db.PgDB,
+	pgDB *db.PgDB,
 	taskID model.TaskID,
 	jobID model.JobID,
 	jobSubmissionTime time.Time,
@@ -82,7 +83,8 @@ func runCheckpointGCTask(
 	})
 	syslog := logrus.WithField("component", "checkpointgc").WithFields(logCtx.Fields())
 
-	if err := db.AddTask(&model.Task{
+	ctx := context.Background()
+	if err := db.AddTask(ctx, &model.Task{
 		TaskID:     taskID,
 		TaskType:   model.TaskTypeCheckpointGC,
 		StartTime:  time.Now().UTC(),
@@ -97,7 +99,7 @@ func runCheckpointGCTask(
 
 	resultChan := make(chan error, 1)
 	onExit := func(ae *task.AllocationExited) {
-		if err := db.CompleteTask(taskID, time.Now().UTC()); err != nil {
+		if err := db.CompleteTask(ctx, taskID, time.Now().UTC()); err != nil {
 			syslog.WithError(err).Error("marking GC task complete")
 		}
 		if err := tasklist.GroupPriorityChangeRegistry.Delete(gcJobID); err != nil {
@@ -119,7 +121,7 @@ func runCheckpointGCTask(
 			SingleAgent: true,
 		},
 		ResourcePool: rp,
-	}, db, rm, gcSpec, onExit)
+	}, pgDB, rm, gcSpec, onExit)
 	if err != nil {
 		return err
 	}

master/internal/command/command.go

@@ -230,10 +230,11 @@ func (c *Command) OnExit(ae *task.AllocationExited) {
 
 	c.exitStatus = ae
 
-	if err := c.db.CompleteTask(c.taskID, time.Now().UTC()); err != nil {
+	ctx := context.Background()
+	if err := db.CompleteTask(ctx, c.taskID, time.Now().UTC()); err != nil {
 		c.syslog.WithError(err).Error("marking task complete")
 	}
-	if err := user.DeleteSessionByToken(context.TODO(), c.GenericCommandSpec.Base.UserSessionToken); err != nil {
+	if err := user.DeleteSessionByToken(ctx, c.GenericCommandSpec.Base.UserSessionToken); err != nil {
 		c.syslog.WithError(err).Errorf(
 			"failure to delete user session for task: %v", c.taskID)
 	}
@@ -251,7 +252,7 @@ func (c *Command) garbageCollect() {
 	}
 
 	if c.exitStatus == nil {
-		if err := c.db.CompleteTask(c.taskID, time.Now().UTC()); err != nil {
+		if err := db.CompleteTask(context.Background(), c.taskID, time.Now().UTC()); err != nil {
 			c.syslog.WithError(err).Error("marking task complete")
 		}
 	}

master/internal/core.go

@@ -791,9 +791,9 @@ func (m *Master) restoreNonTerminalExperiments() error {
 	return nil
 }
 
-func (m *Master) closeOpenAllocations() error {
+func (m *Master) closeOpenAllocations(ctx context.Context) error {
 	allocationIds := task.DefaultService.GetAllAllocationIDs()
-	if err := m.db.CloseOpenAllocations(allocationIds); err != nil {
+	if err := db.CloseOpenAllocations(ctx, allocationIds); err != nil {
 		return err
 	}
 	return nil
@@ -1081,11 +1081,11 @@ func (m *Master) Run(ctx context.Context, gRPCLogInitDone chan struct{}) error {
 		return err
 	}
 
-	if err = m.closeOpenAllocations(); err != nil {
+	if err = m.closeOpenAllocations(ctx); err != nil {
 		return err
 	}
 
-	if err = m.db.EndAllTaskStats(); err != nil {
+	if err = db.EndAllTaskStats(ctx); err != nil {
 		return err
 	}
 

master/internal/db/database.go

@@ -43,13 +43,8 @@ type DB interface {
 		id int,
 		experimentBest, trialBest, trialLatest int,
 	) ([]uuid.UUID, error)
-	AddTask(t *model.Task) error
-	UpdateTrial(id int, newState model.State) error
 	UpdateTrialRunnerState(id int, state string) error
 	UpdateTrialRunnerMetadata(id int, md *trialv1.TrialRunnerMetadata) error
-	AddAllocation(a *model.Allocation) error
-	CompleteAllocation(a *model.Allocation) error
-	CompleteAllocationTelemetry(aID model.AllocationID) ([]byte, error)
 	TrialRunIDAndRestarts(trialID int) (int, int, error)
 	UpdateTrialRunID(id, runID int) error
 	UpdateTrialRestarts(id, restarts int) error
@@ -89,11 +84,6 @@ type DB interface {
 		trials []*apiv1.TrialsSnapshotResponse_Trial, endTime time.Time, err error)
 	TopTrialsByTrainingLength(experimentID int, maxTrials int, metric string,
 		smallerIsBetter bool) (trials []int32, err error)
-	StartAllocationSession(allocationID model.AllocationID, owner *model.User) (string, error)
-	DeleteAllocationSession(allocationID model.AllocationID) error
-	UpdateAllocationState(allocation model.Allocation) error
-	UpdateAllocationStartTime(allocation model.Allocation) error
-	UpdateAllocationProxyAddress(allocation model.Allocation) error
 	ExperimentSnapshot(experimentID int) ([]byte, int, error)
 	SaveSnapshot(
 		experimentID int, version int, experimentSnapshot []byte,
@@ -114,9 +104,6 @@ type DB interface {
 	RecordInstanceStats(a *model.InstanceStats) error
 	EndInstanceStats(a *model.InstanceStats) error
 	EndAllInstanceStats() error
-	EndAllTaskStats() error
-	RecordTaskEndStats(stats *model.TaskStats) error
-	RecordTaskStats(stats *model.TaskStats) error
 	UpdateJobPosition(jobID model.JobID, position decimal.Decimal) error
 }