From 686f1af182c6766c1ed6a54ee1dc6232bac2b66c Mon Sep 17 00:00:00 2001 From: scttfrdmn <3011922+scttfrdmn@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:06:12 -0700 Subject: [PATCH] feat(resume): --max-concurrent-auto derives concurrency from real quota Extends #492's quota-derived concurrency ceiling to `spawn resume` (#494), which had the same class of gap: its --max-concurrent override was a plain user-typed int, same as launch's was before #492. Refactored the shared quota-derivation core out of resolveAutoMaxConcurrent into deriveMaxConcurrentFromCombos, so a second entry point -- resolveAutoMaxConcurrentFromConfigs -- can reuse it directly from resume's already-built []*aws.LaunchConfig (the pending parameter sets, already defaults-merged and per-entry overridden) instead of re-parsing raw params through a second merge pass. Scope: foreground (local-state) resume only. --detach is explicitly rejected with --max-concurrent-auto: the Lambda-orchestrated path's stored params live in S3 with no download helper yet (only upload exists via UploadParamsToS3), and there's no route today from resume into the orchestrator's stored MaxConcurrent before re-invoking. That gap is real but bigger than this fix -- not silently ignored, just not built yet. Also found and filed (not fixed here, out of scope): resume's region-from-defaults type assertion (`paramFormat.Defaults["region"]. (string)`) panics rather than falling through to its own "us-east-1" default when the params file has no region default and at least one instance was already launched (spawn#495). Fixes #494. --- CHANGELOG.md | 17 +++++- cmd/launch_sweep_quota.go | 71 +++++++++++++++++++++++- cmd/launch_sweep_quota_substrate_test.go | 38 +++++++++++++ cmd/launch_sweep_quota_test.go | 45 +++++++++++++++ cmd/resume.go | 40 ++++++++++++- cmd/resume_test.go | 57 +++++++++++++++++++ docs-gen/resume.md | 1 + 7 files changed, 262 insertions(+), 7 deletions(-) create mode 100644 cmd/resume_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index d25ffac..f678795 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -## [0.99.0] - 2026-08-10 +### Added +- **`spawn resume --max-concurrent-auto`** — the same quota-derived + concurrency ceiling `spawn launch` gained in #492 (v0.99.0), now available + when resuming an interrupted parameter sweep (#494). Re-derives the + ceiling from the account's real AWS quota headroom for the PENDING + parameter sets' instance type(s), instead of the sweep's original ceiling + or a user-typed override — useful when resuming after the account's quota + situation has changed (freed up, or was the reason the sweep stalled in + the first place). Mutually exclusive with `--max-concurrent`. + + Not yet supported with `--detach` (the Lambda-orchestrated resume path): + that path's stored parameters live in S3 with no download helper yet + (only upload exists), and there's currently no route from `resume` to + update the Lambda-orchestrator's stored ceiling before re-invoking. + `--max-concurrent-auto --detach` is rejected with a clear error rather + than silently ignored; tracked as a separate follow-up. ### Added - **`spawn launch --max-concurrent-auto`** derives the parameter-sweep diff --git a/cmd/launch_sweep_quota.go b/cmd/launch_sweep_quota.go index 54ba8ba..fbe2330 100644 --- a/cmd/launch_sweep_quota.go +++ b/cmd/launch_sweep_quota.go @@ -53,14 +53,55 @@ import ( // avoid a second client construction and so tests can inject a // Substrate-backed client. func resolveAutoMaxConcurrent(ctx context.Context, paramFormat *ParamFileFormat, baseConfig *aws.LaunchConfig, region string, awsClient *aws.Client) (int, error) { - if region == "" { - return 0, fmt.Errorf("auto max-concurrent: no region resolved yet — derive the ceiling AFTER region resolution") + combos, err := sweepQuotaCombos(paramFormat, baseConfig) + if err != nil { + return 0, err } + return deriveMaxConcurrentFromCombos(ctx, combos, region, awsClient) +} - combos, err := sweepQuotaCombos(paramFormat, baseConfig) +// resolveAutoMaxConcurrentFromConfigs is the `spawn resume --max-concurrent-auto` +// entry point (#494): unlike resolveAutoMaxConcurrent, the caller already has +// fully-built *aws.LaunchConfig entries (resume reloads the original param file, +// reconciles against live EC2 state, and builds launch configs for exactly the +// PENDING parameter sets before this point) — re-deriving combos from raw +// params via sweepQuotaCombos would apply a second, subtly different merge +// pass over data that's already been merged once. Extracts combos directly +// from the already-resolved configs instead. +func resolveAutoMaxConcurrentFromConfigs(ctx context.Context, launchConfigs []*aws.LaunchConfig, region string, awsClient *aws.Client) (int, error) { + combos, err := combosFromLaunchConfigs(launchConfigs) if err != nil { return 0, err } + return deriveMaxConcurrentFromCombos(ctx, combos, region, awsClient) +} + +// deriveMaxConcurrentFromCombos is the shared core of both +// resolveAutoMaxConcurrent and resolveAutoMaxConcurrentFromConfigs: given the +// distinct (instance type, spot) combinations a sweep will launch and a +// resolved region, query truffle's quota client once for headroom (quota - +// current usage) per family, convert vCPU headroom to an instance count via +// Capabilities.VCPUs (truffle#492 — the real EC2 value, not a guessed +// size-suffix), and return the MINIMUM across every combination — the ceiling +// must respect the tightest-fitting quota, not the loosest, or a +// heterogeneous sweep could still overrun a scarce family while looking fine +// against a roomy one. +// +// A family this can't get a quota/vCPU answer for (missing credentials, +// unsupported family, API error) is skipped with a warning rather than +// aborting the whole sweep — this is a SAFETY DEFAULT, not a hard gate, and a +// caller that explicitly asked for --max-concurrent=auto should still get a +// best-effort number rather than an outright failure when only some rows in a +// mixed sweep are affected. +// +// awsClient is an already-constructed client PINNED TO region; this function +// builds truffle clients from its config rather than constructing its own, +// both to avoid a second client construction and so tests can inject a +// Substrate-backed client. +func deriveMaxConcurrentFromCombos(ctx context.Context, combos []sweepQuotaCombo, region string, awsClient *aws.Client) (int, error) { + if region == "" { + return 0, fmt.Errorf("auto max-concurrent: no region resolved yet — derive the ceiling AFTER region resolution") + } quotaClient := truffleQuotas.NewClientFromConfig(awsClient.Config()) capsClient := truffleaws.NewClientFromConfig(awsClient.Config()) @@ -149,3 +190,27 @@ func sweepQuotaCombos(paramFormat *ParamFileFormat, baseConfig *aws.LaunchConfig } return combos, nil } + +// combosFromLaunchConfigs extracts the distinct (instance type, spot) +// combinations from already-built launch configs (#494's `spawn resume` +// path — the configs are the PENDING parameter sets only, already +// defaults-merged and per-entry-overridden by buildLaunchConfigFromParams, +// so no second merge pass is needed here). +func combosFromLaunchConfigs(launchConfigs []*aws.LaunchConfig) ([]sweepQuotaCombo, error) { + seen := make(map[sweepQuotaCombo]bool) + var combos []sweepQuotaCombo + for i, cfg := range launchConfigs { + if cfg.InstanceType == "" { + return nil, fmt.Errorf("auto max-concurrent: launch config %d has no instance_type", i) + } + c := sweepQuotaCombo{instanceType: cfg.InstanceType, spot: cfg.Spot} + if !seen[c] { + seen[c] = true + combos = append(combos, c) + } + } + if len(combos) == 0 { + return nil, fmt.Errorf("auto max-concurrent: no pending launch configs to derive a ceiling from") + } + return combos, nil +} diff --git a/cmd/launch_sweep_quota_substrate_test.go b/cmd/launch_sweep_quota_substrate_test.go index 0c03a42..fb42f9d 100644 --- a/cmd/launch_sweep_quota_substrate_test.go +++ b/cmd/launch_sweep_quota_substrate_test.go @@ -80,3 +80,41 @@ func TestResolveAutoMaxConcurrent_RequiresResolvedRegion(t *testing.T) { t.Error("want error when region is empty") } } + +// TestResolveAutoMaxConcurrentFromConfigs_DerivesFromRealQuota is the +// spawn#494 (`spawn resume --max-concurrent-auto`) end-to-end guard, mirroring +// TestResolveAutoMaxConcurrent_DerivesFromRealQuota but from already-built +// *aws.LaunchConfig entries (resume's pending configs) rather than raw +// ParamFileFormat params. Same Substrate seeding: 32 vCPU Standard On-Demand +// quota / 4 vCPU per c5.xlarge = 8. +func TestResolveAutoMaxConcurrentFromConfigs_DerivesFromRealQuota(t *testing.T) { + env := testutil.SubstrateServer(t) + ctx := context.Background() + client := aws.NewClientFromConfig(env.AWSConfig) + + launchConfigs := []*aws.LaunchConfig{ + {InstanceType: "c5.xlarge", Region: "us-east-1"}, + {InstanceType: "c5.xlarge", Region: "us-east-1"}, + } + + got, err := resolveAutoMaxConcurrentFromConfigs(ctx, launchConfigs, "us-east-1", client) + if err != nil { + t.Fatalf("resolveAutoMaxConcurrentFromConfigs: %v", err) + } + if got != 8 { + t.Errorf("derived max-concurrent = %d, want 8 (32 vCPU quota / 4 vCPU per c5.xlarge)", got) + } +} + +// TestResolveAutoMaxConcurrentFromConfigs_RequiresResolvedRegion is the +// resume-path counterpart of TestResolveAutoMaxConcurrent_RequiresResolvedRegion. +func TestResolveAutoMaxConcurrentFromConfigs_RequiresResolvedRegion(t *testing.T) { + env := testutil.SubstrateServer(t) + ctx := context.Background() + client := aws.NewClientFromConfig(env.AWSConfig) + + launchConfigs := []*aws.LaunchConfig{{InstanceType: "c5.xlarge"}} + if _, err := resolveAutoMaxConcurrentFromConfigs(ctx, launchConfigs, "", client); err == nil { + t.Error("want error when region is empty") + } +} diff --git a/cmd/launch_sweep_quota_test.go b/cmd/launch_sweep_quota_test.go index 53831f2..7f2bae8 100644 --- a/cmd/launch_sweep_quota_test.go +++ b/cmd/launch_sweep_quota_test.go @@ -76,3 +76,48 @@ func TestSweepQuotaCombos_NoParamsErrors(t *testing.T) { t.Error("want error for an empty param file") } } + +// TestCombosFromLaunchConfigs_Dedupes is the spawn#494 (`spawn resume +// --max-concurrent-auto`) counterpart to TestSweepQuotaCombos_*: extracting +// combos from already-built *aws.LaunchConfig entries (resume's pending +// configs are already fully merged, unlike launch's raw param sets) must +// still dedupe correctly. +func TestCombosFromLaunchConfigs_Dedupes(t *testing.T) { + configs := []*aws.LaunchConfig{ + {InstanceType: "g7e.2xlarge", Spot: true}, + {InstanceType: "g7e.2xlarge", Spot: true}, // duplicate + {InstanceType: "g7e.4xlarge", Spot: true}, + {InstanceType: "c7i.xlarge", Spot: false}, + } + + combos, err := combosFromLaunchConfigs(configs) + if err != nil { + t.Fatalf("combosFromLaunchConfigs: %v", err) + } + want := map[sweepQuotaCombo]bool{ + {instanceType: "g7e.2xlarge", spot: true}: true, + {instanceType: "g7e.4xlarge", spot: true}: true, + {instanceType: "c7i.xlarge", spot: false}: true, + } + if len(combos) != len(want) { + t.Fatalf("got %d combos, want %d: %v", len(combos), len(want), combos) + } + for _, c := range combos { + if !want[c] { + t.Errorf("unexpected combo %+v", c) + } + } +} + +func TestCombosFromLaunchConfigs_MissingInstanceTypeErrors(t *testing.T) { + configs := []*aws.LaunchConfig{{}} + if _, err := combosFromLaunchConfigs(configs); err == nil { + t.Error("want error when a launch config has no instance_type") + } +} + +func TestCombosFromLaunchConfigs_EmptyErrors(t *testing.T) { + if _, err := combosFromLaunchConfigs(nil); err == nil { + t.Error("want error for no launch configs") + } +} diff --git a/cmd/resume.go b/cmd/resume.go index 609833d..a90d912 100644 --- a/cmd/resume.go +++ b/cmd/resume.go @@ -15,9 +15,10 @@ import ( ) var ( - resumeSweepID string - resumeMaxConcurrent int - resumeDetach bool + resumeSweepID string + resumeMaxConcurrent int + resumeMaxConcurrentAuto bool + resumeDetach bool ) var resumeCmd = &cobra.Command{ @@ -46,6 +47,7 @@ func init() { resumeCmd.Flags().StringVar(&resumeSweepID, "sweep-id", "", "Sweep ID to resume (required)") _ = resumeCmd.MarkFlagRequired("sweep-id") resumeCmd.Flags().IntVar(&resumeMaxConcurrent, "max-concurrent", 0, "Override max concurrent instances (0 = use original)") + resumeCmd.Flags().BoolVar(&resumeMaxConcurrentAuto, "max-concurrent-auto", false, "Re-derive max concurrent instances from the account's real AWS quota headroom for the pending parameter sets' instance type(s)/region, instead of the original or a user-supplied number (spawn#494). Not yet supported with --detach.") resumeCmd.Flags().BoolVar(&resumeDetach, "detach", false, "Run sweep orchestration in Lambda") rootCmd.AddCommand(resumeCmd) @@ -54,8 +56,23 @@ func init() { func runResume(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + if resumeMaxConcurrentAuto && resumeMaxConcurrent > 0 { + return fmt.Errorf("--max-concurrent-auto and --max-concurrent are mutually exclusive — pick one") + } + // Check if detached mode requested - handle early since detached sweeps have no local state if resumeDetach { + // --max-concurrent-auto needs the pending parameter sets' instance + // types, which for a detached sweep live in S3 (SweepRecord only + // stores an S3ParamsKey, not the params themselves) — no download + // helper exists yet (only the upload side, UploadParamsToS3), and + // there's currently no override path from resume into the stored + // SweepRecord.MaxConcurrent the Lambda orchestrator reads. spawn#494 + // tracks this as a separate follow-up; reject explicitly for now + // rather than silently ignoring the flag. + if resumeMaxConcurrentAuto { + return fmt.Errorf("--max-concurrent-auto is not yet supported with --detach (spawn#494)") + } return resumeSweepDetached(ctx, resumeSweepID) } @@ -212,6 +229,23 @@ func runResume(cmd *cobra.Command, args []string) error { launchConfigs = append(launchConfigs, &config) } + // --max-concurrent-auto (#494): now that the pending launch configs are + // built (each with a resolved InstanceType/Region/Spot), derive a real + // quota-backed ceiling instead of the original or a user-typed number. + if resumeMaxConcurrentAuto { + regionClient, rerr := aws.NewClientWithRegion(ctx, region) + if rerr != nil { + return fmt.Errorf("--max-concurrent-auto: create AWS client for %s: %w", region, rerr) + } + derived, rerr := resolveAutoMaxConcurrentFromConfigs(ctx, launchConfigs, region, regionClient) + if rerr != nil { + return fmt.Errorf("--max-concurrent-auto: %w", rerr) + } + fmt.Fprintf(os.Stderr, "✓ --max-concurrent-auto: derived %d (account quota headroom in %s)\n", derived, region) + fmt.Fprintf(os.Stderr, "🔧 Overriding max-concurrent: %d -> %d\n\n", state.MaxConcurrent, derived) + maxConcurrent = derived + } + // We need to setup shared resources (AMI, SSH key, IAM role) for the pending configs // Use the first pending config as template if len(launchConfigs) == 0 { diff --git a/cmd/resume_test.go b/cmd/resume_test.go new file mode 100644 index 0000000..4bcccf7 --- /dev/null +++ b/cmd/resume_test.go @@ -0,0 +1,57 @@ +package cmd + +import "testing" + +// resetResumeFlags restores the package-level resume flag vars to their +// zero values after a test mutates them directly (bypassing cobra's flag +// parsing, since these tests only exercise runResume's own validation, not +// the CLI plumbing). +func resetResumeFlags(t *testing.T) { + t.Helper() + orig := struct { + sweepID string + maxConc int + maxConcAut bool + detach bool + }{resumeSweepID, resumeMaxConcurrent, resumeMaxConcurrentAuto, resumeDetach} + t.Cleanup(func() { + resumeSweepID = orig.sweepID + resumeMaxConcurrent = orig.maxConc + resumeMaxConcurrentAuto = orig.maxConcAut + resumeDetach = orig.detach + }) +} + +// TestRunResume_MaxConcurrentAutoAndMaxConcurrentMutuallyExclusive is the +// spawn#494 validation guard: --max-concurrent-auto and --max-concurrent must +// not both be set, matching the pattern launch_sweep.go already has for the +// same two flags. This must be checked BEFORE any file/AWS I/O, which is why +// it's safe to unit-test directly (no Substrate/filesystem needed) — the +// error returns before runResume touches either. +func TestRunResume_MaxConcurrentAutoAndMaxConcurrentMutuallyExclusive(t *testing.T) { + resetResumeFlags(t) + resumeSweepID = "does-not-exist" + resumeMaxConcurrentAuto = true + resumeMaxConcurrent = 5 + + err := runResume(resumeCmd, nil) + if err == nil { + t.Fatal("want error when both --max-concurrent-auto and --max-concurrent are set") + } +} + +// TestRunResume_MaxConcurrentAutoRejectedWithDetach is the spawn#494 guard for +// the not-yet-supported detached path (see the doc comment in resume.go for +// why): --max-concurrent-auto + --detach must fail fast with a clear message +// rather than silently ignoring the flag or reaching into detached-mode I/O. +func TestRunResume_MaxConcurrentAutoRejectedWithDetach(t *testing.T) { + resetResumeFlags(t) + resumeSweepID = "does-not-exist" + resumeMaxConcurrentAuto = true + resumeDetach = true + + err := runResume(resumeCmd, nil) + if err == nil { + t.Fatal("want error when --max-concurrent-auto is combined with --detach") + } +} diff --git a/docs-gen/resume.md b/docs-gen/resume.md index 1c9af73..a6c01f9 100644 --- a/docs-gen/resume.md +++ b/docs-gen/resume.md @@ -27,6 +27,7 @@ spawn resume [flags] | Flag | Short | Type | Default | Description | |------|-------|------|---------|-------------| | `--detach` | | bool | | Run sweep orchestration in Lambda | +| `--max-concurrent-auto` | | bool | | Re-derive max concurrent instances from the account's real AWS quota headroom for the pending parameter sets' instance type(s)/region, instead of the original or a user-supplied number (spawn#494). Not yet supported with --detach. | | `--max-concurrent` | | int | | Override max concurrent instances (0 = use original) | | `--sweep-id` | | string | | Sweep ID to resume (required) |