fix: resolve indefinitely queued (STOPPING_COMPLETED) trials (#9605)

determined-ai · Jul 18, 2024 · c70dd8c · c70dd8c
1 parent f4ecd91
commit c70dd8c
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 12 deletions.
diff --git a/e2e_tests/tests/cluster/test_master_restart.py b/e2e_tests/tests/cluster/test_master_restart.py
@@ -75,6 +75,40 @@ def test_master_restart_reattach_recover_experiment_k8s(
     _test_master_restart_reattach_recover_experiment(k8s_managed_cluster, downtime)
 
 
+@pytest.mark.managed_devcluster
+def test_master_agent_restart_reattach_recover_experiment(
+    restartable_managed_cluster: managed_cluster.ManagedCluster,
+) -> None:
+    sess = api_utils.user_session()
+
+    try:
+        # Start an experiment
+        exp_id = exp.create_experiment(
+            sess,
+            conf.fixtures_path("no_op/single-medium-train-step.yaml"),
+            conf.fixtures_path("no_op"),
+            None,
+        )
+
+        # Kill the agent & master
+        restartable_managed_cluster.kill_agent()
+        restartable_managed_cluster.kill_master()
+
+        # Restart the agent & master
+        restartable_managed_cluster.restart_master()
+        restartable_managed_cluster.restart_agent(True)
+
+        exp.wait_for_experiment_state(
+            sess, exp_id, bindings.experimentv1State.COMPLETED, max_wait_secs=60
+        )
+        trials = exp.experiment_trials(sess, exp_id)
+        assert (trials[0].trial.state) == bindings.trialv1State.COMPLETED
+    except Exception:
+        restartable_managed_cluster.restart_master()
+        restartable_managed_cluster.restart_agent()
+        raise
+
+
 @pytest.mark.managed_devcluster
 def test_master_restart_generic_task(
     managed_cluster_restarts: managed_cluster.ManagedCluster,

diff --git a/master/internal/trial.go b/master/internal/trial.go
@@ -747,16 +747,16 @@ func (t *trial) transition(s model.StateWithReason) error {
 				InformationalReason: s.InformationalReason,
 			})
 		default:
-			if action, ok := map[model.State]task.AllocationSignal{
-				model.StoppingCanceledState: task.TerminateAllocation,
-				model.StoppingKilledState:   task.KillAllocation,
-				model.StoppingErrorState:    task.KillAllocation,
-			}[t.state]; ok {
-				t.syslog.Infof("decided to %s trial", action)
-				err := task.DefaultService.Signal(*t.allocationID, action, s.InformationalReason)
-				if err != nil {
-					t.syslog.WithError(err).Warnf("could not %s allocation during stop", action)
-				}
+			// For StoppingKilled and StoppingErrorState, default to KillAllocation.
+			action := task.KillAllocation
+
+			if (t.state == model.StoppingCompletedState) || (t.state == model.StoppingCanceledState) {
+				action = task.TerminateAllocation
+			}
+
+			t.syslog.Infof("decided to %s trial", action)
+			if err := task.DefaultService.Signal(*t.allocationID, action, s.InformationalReason); err != nil {
+				t.syslog.WithError(err).Warnf("could not %s allocation during stop", action)
 			}
 		}
 	case model.TerminalStates[t.state]:

diff --git a/master/internal/trial_intg_test.go b/master/internal/trial_intg_test.go
@@ -44,8 +44,6 @@ func TestTrial(t *testing.T) {
 		Complete: false,
 		Closed:   true,
 	}))
-	require.True(t, alloc.AssertExpectations(t))
-	require.NotNil(t, tr.allocationID)
 
 	// Running stage.
 	require.NoError(t, tr.PatchSearcherState(experiment.TrialSearcherState{
@@ -57,6 +55,8 @@ func TestTrial(t *testing.T) {
 		Complete: true,
 		Closed:   true,
 	}))
+	require.True(t, alloc.AssertExpectations(t))
+	require.NotNil(t, tr.allocationID)
 
 	dbTrial, err := internaldb.TrialByID(context.TODO(), tr.id)
 	require.NoError(t, err)
@@ -135,6 +135,7 @@ func setup(t *testing.T) (
 		"StartAllocation", mock.Anything, mock.Anything, mock.Anything,
 		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
 	).Return(nil)
+	as.On("Signal", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 	a, _, _ := setupAPITest(t, nil)
 	j := &model.Job{JobID: model.NewJobID(), JobType: model.JobTypeExperiment}