fix: revert: fix: resolve indefinitely queued (STOPPING_COMPLETED) tr…

…ials (#10207)

e2e_tests/tests/cluster/test_master_restart.py

@@ -116,33 +116,6 @@ 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_ref = noop.create_experiment(sess, [noop.Sleep(10)])
-
-        # 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)
-
-        assert exp_ref.wait(interval=0.01) == client.ExperimentState.COMPLETED
-        trials = exp.experiment_trials(sess, exp_ref.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,

master/internal/trial.go

@@ -766,16 +766,16 @@ func (t *trial) transition(s model.StateWithReason) error {
 				InformationalReason: s.InformationalReason,
 			})
 		default:
-			// 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)
+			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)
+				}
 			}
 		}
 	case model.TerminalStates[t.state]:

master/internal/trial_intg_test.go

@@ -41,15 +41,15 @@ func TestTrial(t *testing.T) {
 		EarlyStoppedBySearcher: false,
 		EarlyExitedByUserCode:  false,
 	}))
+	require.True(t, alloc.AssertExpectations(t))
+	require.NotNil(t, tr.allocationID)
 
 	// Running stage.
 	require.NoError(t, tr.PatchSearcherState(experiment.TrialSearcherState{
 		Create:                 searcher.Create{},
 		EarlyStoppedBySearcher: true,
 		EarlyExitedByUserCode:  false,
 	}))
-	require.True(t, alloc.AssertExpectations(t))
-	require.NotNil(t, tr.allocationID)
 
 	dbTrial, err := internaldb.TrialByID(context.TODO(), tr.id)
 	require.NoError(t, err)
@@ -123,7 +123,6 @@ 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}