fix: coalesce concurrent per-AZ placement group creation in mpicohort

CHANGELOG.md

@@ -39,6 +39,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   returns `(cost float64, measured bool)`; the caller only logs a rate when
   `measured` is true, and logs an explicit "unknown (lookup failed or denied)"
   line otherwise.
+- **`pkg/mpicohort`'s `Actuator.ensurePlacementGroup` had a check-then-act
+  race**: the mutex was released between checking whether an AZ's placement
+  group was already created and calling `CreatePlacementGroup`, so concurrent
+  cohort members in the same newly-visited AZ (the normal case — a round
+  launches N members at once) could each observe "not created yet" and each
+  call `CreatePlacementGroup`, defeating the once-per-AZ design that exists to
+  avoid the ~30s availability poll for every member of a round (#514,
+  surfaced by a CI flake in `TestActuator_PerAZPlacementGroup`). Concurrent
+  callers for the same AZ now coalesce onto a single in-flight
+  `CreatePlacementGroup` call instead of each issuing their own; callers for
+  different AZs are not serialized against each other.
 
 ## [0.100.2] - 2026-08-18
 

pkg/mpicohort/adapter.go

@@ -85,8 +85,21 @@
 	// on demand. Empty → no PG (or the config already carries a fixed one).
 	PlacementGroupPrefix string
 
-	pgMu      sync.Mutex          // guards pgCreated
-	pgCreated map[string]struct{} // AZ → placement group ensured (create-once per AZ)
+	pgMu      sync.Mutex               // guards pgCreated and pgOnce
+	pgCreated map[string]struct{}      // AZ → placement group ensured (create-once per AZ)
+	pgOnce    map[string]*pgCreateOnce // AZ → in-flight/completed CreatePlacementGroup call
+}
+
+// pgCreateOnce coalesces concurrent ensurePlacementGroup calls for the same AZ
+// into a single CreatePlacementGroup call. A round launches N cohort members
+// concurrently (see TestActuator_PerAZPlacementGroup), so multiple goroutines
+// commonly race into ensurePlacementGroup for the same AZ at once; without
+// this, more than one could observe pgCreated[az] unset before any of them
+// finished creating it, each paying the ~30s availability poll the once-per-AZ
+// design exists to avoid (spawn#514).
+type pgCreateOnce struct {
+	done chan struct{}
+	err  error
 }
 
 func (a *Actuator) Launch(ctx context.Context, intent cohort.EntityIntent) (cohort.Observation, error) {
@@ -136,26 +149,67 @@
 // ensurePlacementGroup creates the per-AZ cluster placement group at most once
 // per AZ (create is idempotent, but this avoids the redundant 30s availability
 // poll for every member of a round in the same AZ) and returns its name.
+//
+// Concurrent callers for the SAME AZ coalesce onto one CreatePlacementGroup
+// call via pgOnce: the first caller in does the create and the rest wait on
+// its result, rather than each independently observing pgCreated[az] unset
+// and calling CreatePlacementGroup themselves (spawn#514 — the lock used to be
+// released between the check and the create, so a round's concurrent members
+// could each pay the 30s poll for a newly-visited AZ). Callers for DIFFERENT
+// AZs are not serialized against each other.
 func (a *Actuator) ensurePlacementGroup(ctx context.Context, az string) (string, error) {
 	name := PlacementGroupName(a.PlacementGroupPrefix, az)
 
 	a.pgMu.Lock()
 	if a.pgCreated == nil {
 		a.pgCreated = make(map[string]struct{})
 	}
-	_, done := a.pgCreated[az]
-	a.pgMu.Unlock()
-	if done {
+	if _, done := a.pgCreated[az]; done {
+		a.pgMu.Unlock()
 		return name, nil
 	}
+	if a.pgOnce == nil {
+		a.pgOnce = make(map[string]*pgCreateOnce)
+	}
+	once, inFlight := a.pgOnce[az]
+	if !inFlight {
+		once = &pgCreateOnce{done: make(chan struct{})}
+		a.pgOnce[az] = once
+	}
+	a.pgMu.Unlock()
 
-	if err := a.Client.CreatePlacementGroup(ctx, name, a.Region); err != nil {
-		return "", fmt.Errorf("create placement group %q: %w", name, err)
+	if inFlight {
+		// Another goroutine is already creating this AZ's placement group;
+		// wait for it rather than issuing a redundant CreatePlacementGroup.
+		select {
+		case <-once.done:
+			if once.err != nil {
+				return "", fmt.Errorf("create placement group %q: %w", name, once.err)
+			}
+			return name, nil
+		case <-ctx.Done():
+			return "", ctx.Err()
+		}
 	}
 
+	// We won the race to create this AZ's group.
+	err := a.Client.CreatePlacementGroup(ctx, name, a.Region)
+	once.err = err
+	close(once.done)
+
 	a.pgMu.Lock()
-	a.pgCreated[az] = struct{}{}
+	// Drop the once-entry regardless of outcome: on success pgCreated below is
+	// the source of truth; on failure this lets a future call retry cleanly
+	// instead of replaying the same error forever.
+	delete(a.pgOnce, az)
+	if err == nil {
+		a.pgCreated[az] = struct{}{}
+	}
 	a.pgMu.Unlock()
+
+	if err != nil {
+		return "", fmt.Errorf("create placement group %q: %w", name, err)
+	}
 	return name, nil
 }
 

pkg/mpicohort/adapter_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
@@ -27,8 +28,9 @@ type fakeLauncher struct {
 	failAZs   map[string]bool   // if set, any launch into one of these AZs ICEs (multi-AZ exhaustion)
 	launchLog []launchRec
 
-	pgCreated map[string]int // placement group name → CreatePlacementGroup call count
-	pgDeleted []string       // placement group names passed to DeletePlacementGroup
+	pgCreated     map[string]int // placement group name → CreatePlacementGroup call count
+	pgDeleted     []string       // placement group names passed to DeletePlacementGroup
+	pgCreateDelay time.Duration  // artificial delay inside CreatePlacementGroup, widens race windows in tests
 
 	ssmCmds       map[string]string // instanceID → last RunShellScript command
 	ssmFailIDs    map[string]bool   // instanceIDs whose RunShellScript returns Failed
@@ -97,6 +99,14 @@ func (f *fakeLauncher) StopInstance(_ context.Context, _, _ string, _ bool) erro
 func (f *fakeLauncher) StartInstance(_ context.Context, _, _ string) error        { return nil }
 
 func (f *fakeLauncher) CreatePlacementGroup(_ context.Context, name, _ string) error {
+	// Model the real ~30s availability poll's latency by yielding before
+	// recording the call — without this, a race between the check and the
+	// create (spawn#514) is very unlikely to actually interleave in a fast
+	// in-memory fake, even though it's exactly what production hits under
+	// load. Sleeping briefly (rather than 30s) widens the window so
+	// concurrent ensurePlacementGroup callers for the same AZ reliably race
+	// into this method if the caller doesn't coalesce them.
+	f.createPGDelay()
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	if f.pgCreated == nil {
@@ -106,6 +116,20 @@ func (f *fakeLauncher) CreatePlacementGroup(_ context.Context, name, _ string) e
 	return nil
 }
 
+// createPGDelay yields to the scheduler (and optionally sleeps, if
+// pgCreateDelay is set) before CreatePlacementGroup records its call, to
+// widen the check-then-act race window for TestEnsurePlacementGroup_Concurrent.
+func (f *fakeLauncher) createPGDelay() {
+	f.mu.Lock()
+	d := f.pgCreateDelay
+	f.mu.Unlock()
+	if d > 0 {
+		time.Sleep(d)
+	} else {
+		runtime.Gosched()
+	}
+}
+
 func (f *fakeLauncher) DeletePlacementGroup(_ context.Context, name string) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
@@ -561,3 +585,97 @@ func TestActuator_ConfigsLookup(t *testing.T) {
 		t.Errorf("node-9 user-data = %q, want BASE-UD (BaseConfig fallback)", got["node-9"])
 	}
 }
+
+// TestEnsurePlacementGroup_Concurrent is the direct regression test for
+// spawn#514: ensurePlacementGroup had a check-then-act race — the lock was
+// released between checking pgCreated[az] and calling CreatePlacementGroup —
+// so concurrent callers for the same AZ (the normal case: a round launches N
+// cohort members at once) could each observe "not created yet" and each call
+// CreatePlacementGroup, defeating the once-per-AZ design that exists
+// specifically to avoid the ~30s availability poll for every member of a
+// round in a newly-visited AZ.
+//
+// This drives many goroutines at once for a SINGLE az, with an artificial
+// delay inside CreatePlacementGroup to widen the race window (a fast in-memory
+// fake without it may not interleave even with a real race present — this is
+// how the CI flake surfaced but a quiet local run didn't, per the issue).
+func TestEnsurePlacementGroup_Concurrent(t *testing.T) {
+	f := newFakeLauncher()
+	f.pgCreateDelay = 5 * time.Millisecond
+
+	act := &Actuator{Client: f, Region: "us-east-1", PlacementGroupPrefix: "spawn-mpi-train"}
+
+	const n = 50
+	var wg sync.WaitGroup
+	errs := make([]error, n)
+	names := make([]string, n)
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			names[i], errs[i] = act.ensurePlacementGroup(context.Background(), "us-east-1a")
+		}(i)
+	}
+	wg.Wait()
+
+	for i, err := range errs {
+		if err != nil {
+			t.Fatalf("goroutine %d: ensurePlacementGroup error: %v", i, err)
+		}
+		if names[i] != "spawn-mpi-train-us-east-1a" {
+			t.Errorf("goroutine %d: name = %q, want spawn-mpi-train-us-east-1a", i, names[i])
+		}
+	}
+
+	f.mu.Lock()
+	got := f.pgCreated["spawn-mpi-train-us-east-1a"]
+	f.mu.Unlock()
+	if got != 1 {
+		t.Errorf("CreatePlacementGroup called %d times for one AZ under %d concurrent callers, want exactly 1", got, n)
+	}
+}
+
+// TestEnsurePlacementGroup_ConcurrentDifferentAZsNotSerialized confirms the
+// fix doesn't over-serialize: concurrent callers for DIFFERENT AZs must not
+// block on each other (only same-AZ callers coalesce).
+func TestEnsurePlacementGroup_ConcurrentDifferentAZsNotSerialized(t *testing.T) {
+	f := newFakeLauncher()
+	f.pgCreateDelay = 5 * time.Millisecond
+
+	act := &Actuator{Client: f, Region: "us-east-1", PlacementGroupPrefix: "spawn-mpi-train"}
+
+	azs := []string{"us-east-1a", "us-east-1b", "us-east-1c"}
+	const perAZ = 10
+	var wg sync.WaitGroup
+	start := time.Now()
+	for _, az := range azs {
+		for i := 0; i < perAZ; i++ {
+			wg.Add(1)
+			go func(az string) {
+				defer wg.Done()
+				if _, err := act.ensurePlacementGroup(context.Background(), az); err != nil {
+					t.Errorf("ensurePlacementGroup(%s): %v", az, err)
+				}
+			}(az)
+		}
+	}
+	wg.Wait()
+	elapsed := time.Since(start)
+
+	// If AZs were serialized against each other, this would take roughly
+	// len(azs) * pgCreateDelay; since they're independent, it should complete
+	// in roughly one delay's worth of time. Generous bound to avoid flaking
+	// on a loaded CI runner.
+	if elapsed > 500*time.Millisecond {
+		t.Errorf("ensurePlacementGroup for different AZs took %v, suggests they are serialized against each other", elapsed)
+	}
+
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	for _, az := range azs {
+		name := "spawn-mpi-train-" + az
+		if got := f.pgCreated[name]; got != 1 {
+			t.Errorf("CreatePlacementGroup(%s) called %d times, want 1", name, got)
+		}
+	}
+}