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) |