ate-env-api: wait for a suspended environment to wake instead of failing its first call

[Vicente Ferrara (vicentefb)]

Why

A suspended environment is resumed by the first request that reaches the router, and the router holds that request only for its own parking budget (a few seconds). A wake that pulls the environment's image onto the node or reads its snapshot across a region takes longer: 18 s for a cold image pull, 7 s for a 1 GiB snapshot across regions (measured on GKE this week). The caller got Unavailable while the wake succeeded behind it. A guest not yet listening when its first request arrived produced the same error from the worker. Every client made one attempt, so one such window failed a training step and, in one run, crashed the trainer and leaked its fleet.

What

Before forwarding any guest call, ate-env-api sends one trivial request to the guest (GetProcess with a probe id) and repeats it with backoff (200 ms → 2 s) while the router answers Unavailable, bounded by --wake-budget (default 2m). Between attempts it checks the actor's state in the control plane and fails at once with FailedPrecondition if it crashed. Anything other than Unavailable means the guest answered; the real call then proceeds on the same connection. On a running environment the probe is one extra round trip.

Server-side on purpose: the Python client, the CLI and the MCP handler all get it without changes, and the server already owns the per-actor router connection.

Tests

TestGuestCallWaitsForWake: environment held SUSPENDED, marked RUNNING after 600 ms by a stand-in for the resume; StartProcess succeeds and returns after the wake rather than failing with the router's 503. TestGuestCallFailsFastOnCrash: environment marked CRASHED; the call fails with FailedPrecondition in milliseconds. The fake control plane gains SetState. Existing proxy tests pass with the probe in the path.

Verified on GKE: an actor whose wake takes 18 s (cold task image on the node) now answers its first command after the wake instead of failing at the 10 s parking budget.

[Tomer Glottmann (tomergee)]

Review at ee8eabe. I read the diff, ran the two new tests three times with -race, ran the package suite, and trial-merged the branch against #82, #83 and #84 (clean). The design is sound and the server is the right place for it. One behaviour change to bound before merge, one cost to note.

Blocking: none.

Before merge

1. internal/apiservice/wakeretry.go:59-83. The loop retries on each Unavailable for the full budget, also when the control plane reports the actor RUNNING. Today a router outage, or a worker the router cannot reach, fails a call at once; with this PR each call hangs for two minutes and holds a goroutine and a connection while it waits. RUNNING plus Unavailable has two causes: a guest that is not listening yet, which resolves in seconds, and a path that is down, which does not. Use the state you already read in actorCrashed: for RUNNING cap the wait at a short budget (10 to 30 s, a second flag or a constant), and keep the two minutes for SUSPENDED, RESUMING and the other transitional states. Put the state in the final error message, thus an operator can tell the two cases apart. The same loop calls client.Get on each attempt, about 60 control-plane reads per waiting call over the two minutes; the shorter cap for RUNNING bounds that too.

Nits

• internal/apiservice/wakeretry_test.go. The budget-expiry path has no test: Unavailable for the whole budget must end in Unavailable with the "did not wake within" message and the attempt count. WakeBudget is a public field, thus the test needs only the *apiservice.Server on testEnv (it is not kept today) and a budget of a few hundred milliseconds.
• internal/apiservice/server.go:497. The probe is one extra round trip on each of the three guest paths, on a running environment too. Env.shell() is two guest calls, thus two probes; a rollout loop pays it on each exec, and the nemo-gym readiness poll pays it on each attempt. Not a reason to hold the PR. A follow-up can remember "awake" per (atespace, id) for a few seconds after a probe or a successful call, so a loop probes once.
• internal/apiservice/wakeretry.go:93-96. A deleted environment comes back from actorCrashed as NotFound through toGRPCError, which is the right answer; a comment there saves the next reader the trace.
• The budget is reachable only by the --wake-budget flag; ate-env manifest workerpool does not expose it. Fine for now, because the default is the intended value. One line in docs/release.md or the README's deploy section when flags get documented.

Verified

• TestGuestCallWaitsForWake and TestGuestCallFailsFastOnCrash pass three times with -race; the full internal/apiservice suite passes.
• The fake router returns 503 from the Running hook, thus the wait test exercises the real path: 503 → Unavailable → retry → state flips → probe answers NotFound from the real guest → call proceeds.
• The probe id reaches the guest as GetProcess, which answers NotFound without a log line (guest/process/service.go:229), thus no noise per call.
• Merges clean with #82, #83, #84 on main at 490130e.

[Vicente Ferrara (vicentefb)]

Addressed. RUNNING now has its own budget (--wake-running-budget, default 20 s) so a down path fails in seconds while suspended and transitional states keep the 2 m budget; the final error names the state and the attempt count; the state is read once per attempt and the 2 s backoff cap bounds that at ~60 reads per two-minute wait. Added TestGuestCallFailsWhenWakeBudgetExpires and TestGuestCallCapsWaitWhenRunning (the fixture now keeps the server so budgets can be set), a comment on the NotFound path for deleted environments, and the two flags in the README deploy section. The per-call probe cost is a follow-up as you suggested: a short-lived awake cache per (atespace, id).

[Tomer Glottmann (tomergee)]

Review at 3df2928. All five items from the first pass are done: the RUNNING cap with its own flag, the state and attempt count in the error, the budget-expiry test, the NotFound comment, the flags in the README. I read the diff, ran the package suite and the four wake tests three times with -race, ran gofmt, vet and staticcheck (clean), and trial-merged against main at df24cef (clean). One defect in the new accounting blocks.

Blocking

B1. internal/apiservice/wakeretry.go:90-98. The running budget is measured from start, the beginning of the whole wait, not from the attempt that first sees the actor RUNNING. A wake whose suspended phase lasts longer than WakeRunningBudget fails on the first attempt after the state flips, unless the guest already listens at that moment: remaining = runningBudget - time.Since(start) is negative. That is the slow-wake case this PR exists for. 18 s for a cold image pull plus 7 s for a cross-region snapshot read exceeds the 20 s default, and the attempts are 2 s apart.

Proof against the branch (test below): suspended for 2.5 s, RUNNING at 2.5 s, guest listening at 4.5 s, WakeRunningBudget 3 s, WakeBudget 1 min. The call fails at 3.0 s with did not answer within 3s while RUNNING (5 attempts). The actor was RUNNING for 0.5 s of its 3 s budget.

Fix: start the running budget at the first attempt that sees RUNNING, and keep the wake budget as the bound on the total. With this change the five wake tests pass twice with -race.

	start := time.Now()
	var runningSince time.Time // first attempt that saw the actor RUNNING
	...
		allowed, since := budget, start
		if state == ateapipb.ActorState_ACTOR_STATE_RUNNING {
			if runningSince.IsZero() {
				runningSince = time.Now()
			}
			allowed, since = runningBudget, runningSince
		} else {
			runningSince = time.Time{}
		}
		remaining := allowed - time.Since(since)
		if overall := budget - time.Since(start); overall < remaining {
			remaining = overall
		}

Test (fails on the branch, passes with the fix):

// A wake whose suspended phase outlasts the running budget must still get
// the running budget once the actor is RUNNING: the guest listens a moment
// after the state flips.
func TestGuestCallRunningBudgetStartsAtRunning(t *testing.T) {
	te := newFullTestEnv(t)
	envCtx := registerGuest(t, te, "late-running-test")
	te.control.SetState("late-running-test", ateapipb.ActorState_ACTOR_STATE_SUSPENDED)
	te.srv.WakeBudget = time.Minute
	te.srv.WakeRunningBudget = 3 * time.Second

	start := time.Now()
	flip := start.Add(2500 * time.Millisecond)   // the control plane reports RUNNING; the next attempt is at ~3.0 s
	listen := start.Add(4500 * time.Millisecond) // the guest answers 2 s later, inside the 3 s running budget
	te.router.Running = func(string) bool { return time.Now().After(listen) }
	go func() {
		time.Sleep(time.Until(flip))
		te.control.SetState("late-running-test", ateapipb.ActorState_ACTOR_STATE_RUNNING)
	}()

	_, err := te.procClient.StartProcess(envCtx, &ateenvv1alpha.StartProcessRequest{Command: []string{"sh", "-c", "true"}})
	if err != nil {
		t.Fatalf("StartProcess after a 2.5 s suspended phase and a 2 s listen lag: %v", err)
	}
	if elapsed := time.Since(start); elapsed < 4400*time.Millisecond {
		t.Errorf("call returned after %s, before the guest was listening", elapsed)
	}
}

Before merge: none.

Nit

• wakeretry.go:92-93. With the fix, the total bound can trip while the actor is RUNNING, and the message then names the running budget. Name the budget that tripped.

Verified

• go test -race ./internal/apiservice/... passes. The four wake tests pass three times with -race. gofmt, go vet and staticcheck are clean.
• Trial merge with main at df24cef: clean. CI at 3df2928: green.
• TestGuestCallCapsWaitWhenRunning holds the router at 503 with the state RUNNING from the first attempt, thus it does not reach the accounting defect. The test above does.

[Vicente Ferrara (vicentefb)]

Addressed at 4685b2b. The running budget now starts at the first attempt that sees RUNNING and resets if the state leaves RUNNING; the wake budget remains the bound on the total, and the message and the first-attempt log name whichever budget tripped. Added TestGuestCallRunningBudgetStartsAtRunning unchanged; it fails on 3df2928 and passes now, and the five wake tests pass three times with -race along with the package suite.