edits

e2e_tests/tests/cluster/test_master_restart.py

@@ -90,8 +90,6 @@ def test_master_agent_restart_reattach_recover_experiment(
             None,
         )
 
-        # exp.wait_for_experiment_state(sess, exp_id, bindings.experimentv1State.RUNNING)
-
         # Kill the agent & master
         restartable_managed_cluster.kill_agent()
         restartable_managed_cluster.kill_master()

master/internal/trial.go

@@ -747,17 +747,16 @@ func (t *trial) transition(s model.StateWithReason) error {
 				InformationalReason: s.InformationalReason,
 			})
 		default:
-			if action, ok := map[model.State]task.AllocationSignal{
-				model.StoppingCompletedState: task.TerminateAllocation,
-				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]: