From 58f815f5d3e9ecdb47a2093f567a687529fd9f0f Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Sat, 22 Aug 2026 14:41:28 -0700 Subject: [PATCH 1/2] feat: add region-aware instance search and creation --- pkg/cmd/gpucreate/gpucreate.go | 135 +++++++++++++++++-- pkg/cmd/gpucreate/gpucreate_test.go | 152 ++++++++++++++++++++- pkg/cmd/gpusearch/gpusearch.go | 197 +++++++++++++++++----------- pkg/cmd/gpusearch/gpusearch_test.go | 40 +++++- 4 files changed, 431 insertions(+), 93 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index c4179824..a0adbc66 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -83,6 +83,9 @@ You can attach a startup script that runs when the instance boots using the # Create with a specific GPU type brev create my-instance --type g5.xlarge + # Create in a specific region + brev create my-instance --type g5.xlarge --region us-east-1 + # Try multiple types in order (fallback chain) brev create my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge @@ -134,6 +137,7 @@ type CreateResult struct { // searchFilterFlags holds the search filter flag values for create type searchFilterFlags struct { gpuName string + region string provider string minVRAM float64 minTotalVRAM float64 @@ -145,11 +149,13 @@ type searchFilterFlags struct { flexPorts bool sortBy string descending bool + catalogItems []gpusearch.InstanceType + catalogLoaded bool } // hasUserFilters returns true if the user specified any search filter flags func (f *searchFilterFlags) hasUserFilters() bool { - return f.gpuName != "" || f.provider != "" || f.minVRAM > 0 || f.minTotalVRAM > 0 || + return f.gpuName != "" || f.region != "" || f.provider != "" || f.minVRAM > 0 || f.minTotalVRAM > 0 || f.minCapability > 0 || f.minDisk > 0 || f.maxBootTime > 0 || f.stoppable || f.rebootable || f.flexPorts } @@ -231,6 +237,21 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra return breverrors.WrapAndTrace(err) } + if strings.TrimSpace(filters.region) != "" { + response, err := gpuCreateStore.GetInstanceTypes(false) + if err != nil { + return breverrors.WrapAndTrace(err) + } + filters.catalogLoaded = true + if response != nil { + filters.catalogItems = response.Items + } + filters.region, err = canonicalRegion(filters.region, filters.catalogItems) + if err != nil { + return err + } + } + scriptContent, err := parseStartupScript(startupScript) if err != nil { return breverrors.WrapAndTrace(err) @@ -252,12 +273,16 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra LaunchableID: launchableID, LaunchableInfo: launchableInfo, ParameterBindings: parameterBindings, + Region: filters.region, } opts.InstanceTypes, err = resolveInstanceTypes(cmd, gpuCreateStore, opts, types, &filters) if err != nil { return err } + if err := validateTypesSupportRegion(opts.InstanceTypes, filters.region, filters.catalogItems); err != nil { + return err + } if dryRun { return runDryRun(t, gpuCreateStore, opts.InstanceTypes, &filters) @@ -302,6 +327,7 @@ func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, cmd.Flags().StringVarP(launchable, "launchable", "l", "", "Launchable ID or URL to deploy (e.g., env-XXX or console URL)") cmd.Flags().StringArrayVar(launchableParams, "param", nil, "Launchable setup value NAME=VALUE (repeatable)") + cmd.Flags().StringVarP(&filters.region, "region", "r", "", "Region/location to deploy the instance (e.g., us-east-1, us-central1)") cmd.Flags().StringVarP(&filters.gpuName, "gpu-name", "g", "", "Filter by GPU name (e.g., A100, H100)") cmd.Flags().StringVar(&filters.provider, "provider", "", "Filter by provider/cloud (e.g., aws, gcp)") cmd.Flags().Float64VarP(&filters.minVRAM, "min-vram", "v", 0, "Minimum VRAM per GPU in GB") @@ -339,6 +365,7 @@ type GPUCreateOptions struct { LaunchableID string LaunchableInfo *store.LaunchableResponse // populated when LaunchableID is set ParameterBindings []store.ParameterBinding + Region string } // parseLaunchableID extracts a launchable ID from either a raw ID (env-XXX) or @@ -481,7 +508,8 @@ func warnLaunchableFlagConflicts(cmd *cobra.Command, t *terminal.Terminal, launc } instanceFlagsSet := cmd.Flags().Changed("type") || cmd.Flags().Changed("gpu-name") || - cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram") + cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram") || + cmd.Flags().Changed("region") if instanceFlagsSet { t.Vprintf("Warning: Overriding the launchable's recommended instance configuration. This is not the recommended path and may cause issues.\n\n") } @@ -674,12 +702,18 @@ func parseStartupScript(value string) (string, error) { // searchInstances fetches and filters GPU instances using user-provided filters merged with defaults func searchInstances(s GPUCreateStore, filters *searchFilterFlags) ([]gpusearch.GPUInstanceInfo, float64, error) { - response, err := s.GetInstanceTypes(false) - if err != nil { - return nil, 0, breverrors.WrapAndTrace(err) + items := filters.catalogItems + if !filters.catalogLoaded { + response, err := s.GetInstanceTypes(false) + if err != nil { + return nil, 0, breverrors.WrapAndTrace(err) + } + if response != nil { + items = response.Items + } } - if response == nil || len(response.Items) == 0 { + if len(items) == 0 { return nil, 0, nil } @@ -695,9 +729,10 @@ func searchInstances(s GPUCreateStore, filters *searchFilterFlags) ([]gpusearch. sortBy = "price" } - instances := gpusearch.ProcessInstances(response.Items) + instances := gpusearch.ProcessInstances(items) filtered := gpusearch.FilterInstances(instances, filters.gpuName, filters.provider, "", filters.minVRAM, minTotalVRAM, minCapability, 0, minDisk, 0, maxBootTime, filters.stoppable, filters.rebootable, filters.flexPorts, true) + filtered = gpusearch.FilterInstancesByRegion(filtered, filters.region) gpusearch.SortInstances(filtered, sortBy, filters.descending) return filtered, minDisk, nil @@ -742,6 +777,70 @@ func runDryRun(t *terminal.Terminal, s GPUCreateStore, specs []InstanceSpec, fil return nil } +// canonicalRegion returns the catalog spelling for an exact, case-insensitive +// region match. Create must not forward a partial region name to the provider. +func canonicalRegion(region string, items []gpusearch.InstanceType) (string, error) { + region = strings.TrimSpace(region) + for _, item := range items { + for _, availableRegion := range item.AvailableLocations { + if strings.EqualFold(availableRegion, region) { + return availableRegion, nil + } + } + } + + return "", breverrors.NewValidationError(fmt.Sprintf( + "region %q is not available; run 'brev search --json' to list valid regions", region, + )) +} + +func validateTypesSupportRegion(types []InstanceSpec, region string, items []gpusearch.InstanceType) error { + if region == "" || len(types) == 0 || len(items) == 0 { + return nil + } + + var unsupported []string + for _, spec := range types { + foundType := false + supported := false + for _, item := range items { + if item.Type != spec.Type { + continue + } + foundType = true + if instanceTypeSupportsRegion(item, region) { + supported = true + break + } + } + // A fallback chain remains usable when at least one type supports the + // region. Types missing from the public catalog are left for the + // authenticated create path to validate. + if !foundType || supported { + return nil + } + unsupported = append(unsupported, spec.Type) + } + if len(unsupported) == 0 { + return nil + } + + sort.Strings(unsupported) + return breverrors.NewValidationError(fmt.Sprintf( + "region %q is not available for instance type(s) %s; run 'brev search --region %s' to find compatible types", + region, strings.Join(unsupported, ", "), region, + )) +} + +func instanceTypeSupportsRegion(item gpusearch.InstanceType, region string) bool { + for _, availableRegion := range item.AvailableLocations { + if strings.EqualFold(availableRegion, region) { + return true + } + } + return false +} + // orDefault returns val if it's non-zero, otherwise returns def func orDefault(val, def float64) float64 { if val > 0 { @@ -968,7 +1067,7 @@ func (c *createContext) validateInstanceTypeAvailability(instanceType string) er if c.allInstanceTypes == nil { return nil } - if c.allInstanceTypes.GetCloudCredID(instanceType) != "" { + if c.allInstanceTypes.GetCloudCredIDForRegion(instanceType, c.opts.Region) != "" { return nil } if !c.allInstanceTypes.HasInstanceType(instanceType) { @@ -977,6 +1076,12 @@ func (c *createContext) validateInstanceTypeAvailability(instanceType string) er instanceType, )) } + if c.opts.Region != "" { + return breverrors.NewValidationError(fmt.Sprintf( + "instance type %q is unavailable in region %q; run 'brev search --region %s' to find a compatible type", + instanceType, c.opts.Region, c.opts.Region, + )) + } return breverrors.NewValidationError(fmt.Sprintf( "instance type %q is currently unavailable (no capacity); try again later or run 'brev search' to find another type", instanceType, @@ -1219,8 +1324,13 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity } if c.allInstanceTypes != nil { - if cloudCredID := c.allInstanceTypes.GetCloudCredID(spec.Type); cloudCredID != "" { + if cloudCredID := c.allInstanceTypes.GetCloudCredIDForRegion(spec.Type, c.opts.Region); cloudCredID != "" { cwOptions.WithCloudCredID(cloudCredID) + } else if c.opts.Region != "" && c.allInstanceTypes.HasInstanceType(spec.Type) { + return nil, breverrors.NewValidationError(fmt.Sprintf( + "instance type %q has no available cloud credential in region %q; run 'brev search --region %s' to find a compatible type", + spec.Type, c.opts.Region, c.opts.Region, + )) } } @@ -1234,6 +1344,13 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity } } + if c.opts.Region != "" { + cwOptions.Location = c.opts.Region + // A launchable sub-location belongs to its original location. Let the + // provider choose a compatible zone when the CLI overrides the region. + cwOptions.SubLocation = "" + } + if cwOptions.CloudCredID == "" { if c.allInstanceTypes == nil { return nil, breverrors.NewValidationError(fmt.Sprintf( diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index a812a8ec..3b9badf5 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -29,6 +29,7 @@ type MockGPUCreateStore struct { CreatedWorkspaces []*entity.Workspace DeletedWorkspaceIDs []string FetchedLifeCycleScriptIDs []string + AllInstanceTypes *gpusearch.AllInstanceTypesResponse } func NewMockGPUCreateStore() *MockGPUCreateStore { @@ -108,7 +109,7 @@ func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID strin } func (m *MockGPUCreateStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { - return nil, nil + return m.AllInstanceTypes, nil } func (m *MockGPUCreateStore) GetLaunchable(launchableID string) (*store.LaunchableResponse, error) { @@ -134,7 +135,8 @@ func (m *MockGPUCreateStore) GetInstanceTypes(_ bool) (*gpusearch.InstanceTypesR return &gpusearch.InstanceTypesResponse{ Items: []gpusearch.InstanceType{ { - Type: "g5.xlarge", + Type: "g5.xlarge", + AvailableLocations: []string{"us-east-1", "us-west-2"}, SupportedGPUs: []gpusearch.GPU{ {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, }, @@ -893,6 +895,72 @@ func TestCreateDryRunWithExplicitTypesDoesNotProvision(t *testing.T) { assert.Empty(t, mock.CreatedWorkspaces) } +func TestCreateDryRunAcceptsCanonicalRegion(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"dry-run-region", "--type", "g5.xlarge", "--region", "US-EAST-1", "--dry-run"}) + + err := cmd.Execute() + + assert.NoError(t, err) + assert.Empty(t, mock.CreatedWorkspaces) +} + +func TestCreateAutoSelectsTypeAndCredentialForRegion(t *testing.T) { + mock := NewMockGPUCreateStore() + mock.AllInstanceTypes = &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + { + Type: "g5.xlarge", + CloudCredID: "cc-west", + AvailableLocations: []string{"us-west-2"}, + }, + { + Type: "g5.xlarge", + CloudCredID: "cc-east", + AvailableLocations: []string{"us-east-1"}, + }, + }, + } + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"auto-region", "--region", "US-EAST-1", "--detached"}) + + err := cmd.Execute() + + require.NoError(t, err) + require.Len(t, mock.CreatedOptions, 1) + assert.Equal(t, "g5.xlarge", mock.CreatedOptions[0].InstanceType) + assert.Equal(t, "us-east-1", mock.CreatedOptions[0].Location) + assert.Equal(t, "cc-east", mock.CreatedOptions[0].CloudCredID) +} + +func TestCreateRejectsUnknownRegion(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"bad-region", "--type", "g5.xlarge", "--region", "moon-1", "--dry-run"}) + + err := cmd.Execute() + + assert.ErrorContains(t, err, `region "moon-1" is not available`) + assert.Empty(t, mock.CreatedWorkspaces) +} + +func TestRegionValidation(t *testing.T) { + items := []gpusearch.InstanceType{ + {Type: "g5.xlarge", AvailableLocations: []string{"us-west-2"}}, + {Type: "g5.xlarge", AvailableLocations: []string{"us-east-1"}}, + {Type: "g6.xlarge", AvailableLocations: []string{"eu-west-1"}}, + } + + region, err := canonicalRegion(" US-EAST-1 ", items) + require.NoError(t, err) + assert.Equal(t, "us-east-1", region) + assert.NoError(t, validateTypesSupportRegion([]InstanceSpec{{Type: "g5.xlarge"}}, region, items)) + assert.ErrorContains(t, validateTypesSupportRegion([]InstanceSpec{{Type: "g6.xlarge"}}, region, items), "g6.xlarge") + assert.NoError(t, validateTypesSupportRegion([]InstanceSpec{{Type: "g6.xlarge"}, {Type: "g5.xlarge"}}, region, items), "a compatible fallback keeps the chain usable") + assert.NoError(t, validateTypesSupportRegion([]InstanceSpec{{Type: "private-type"}}, region, items), "types absent from the public catalog are validated by the authenticated create path") +} + func TestGetFilteredInstanceTypesDefaults(t *testing.T) { mock := NewMockGPUCreateStore() @@ -1163,6 +1231,24 @@ func TestValidateInstanceTypeAvailability(t *testing.T) { assert.Contains(t, err.Error(), "brev search") }) + t.Run("returns a region-specific error when credentials exist only elsewhere", func(t *testing.T) { + ctx := &createContext{ + opts: GPUCreateOptions{Region: "us-east-1"}, + allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + { + Type: "hyperstack_H100_sxm5x8", + CloudCredID: "cc-west", + AvailableLocations: []string{"us-west-2"}, + }, + }, + }, + } + err := ctx.validateInstanceTypeAvailability("hyperstack_H100_sxm5x8") + assert.ErrorContains(t, err, `unavailable in region "us-east-1"`) + assert.Contains(t, err.Error(), "brev search --region us-east-1") + }) + t.Run("error type is ValidationError so no stack trace is appended", func(t *testing.T) { ctx := &createContext{ allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ @@ -1234,15 +1320,21 @@ func TestCreateInstancesWithTypeSetsCloudCredIDFromCatalog(t *testing.T) { ctx := &createContext{ t: terminal.New(), store: mock, - opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4"}, + opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4", Region: "us-east-1"}, org: mock.Org, user: mock.User, piped: true, allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ AllInstanceTypes: []gpusearch.InstanceType{ { - Type: "hyperstack_H100_sxm5x8", - CloudCredID: "cc-shadeform", + Type: "hyperstack_H100_sxm5x8", + CloudCredID: "cc-shadeform-west", + AvailableLocations: []string{"us-west-2"}, + }, + { + Type: "hyperstack_H100_sxm5x8", + CloudCredID: "cc-shadeform-east", + AvailableLocations: []string{"us-east-1"}, }, }, }, @@ -1253,7 +1345,8 @@ func TestCreateInstancesWithTypeSetsCloudCredIDFromCatalog(t *testing.T) { assert.False(t, result.hadFailure) require.Len(t, mock.CreatedOptions, 1) - assert.Equal(t, "cc-shadeform", mock.CreatedOptions[0].CloudCredID) + assert.Equal(t, "cc-shadeform-east", mock.CreatedOptions[0].CloudCredID) + assert.Equal(t, "us-east-1", mock.CreatedOptions[0].Location) } func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { @@ -1266,12 +1359,15 @@ func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { Parallel: 1, Name: "jt-4", LaunchableID: "env-abc", + Region: "us-east-1", LaunchableInfo: &store.LaunchableResponse{ ID: "env-abc", Name: "test-launchable", CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{ CloudCredID: "cc-from-launchable", InstanceType: "n2-standard-4", + Location: "eu-west-1", + SubLocation: "eu-west-1a", }, }, }, @@ -1289,4 +1385,48 @@ func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { assert.False(t, result.hadFailure, "launchable should not be blocked by pre-flight validation") assert.Len(t, result.successes, 1, "expected the launchable instance to be created") assert.Len(t, mock.CreatedWorkspaces, 1) + assert.Equal(t, "cc-from-launchable", mock.CreatedOptions[0].CloudCredID, "private launchable types absent from the catalog keep their configured credential") + assert.Equal(t, "us-east-1", mock.CreatedOptions[0].Location, "explicit region should override the launchable default") + assert.Empty(t, mock.CreatedOptions[0].SubLocation, "a sub-location from the launchable's original region must not leak into the override") +} + +func TestCreateLaunchableRejectsCredentialFromDifferentRegion(t *testing.T) { + mock := NewMockGPUCreateStore() + ctx := &createContext{ + t: terminal.New(), + store: mock, + opts: GPUCreateOptions{ + Count: 1, + Parallel: 1, + Name: "jt-4", + LaunchableID: "env-abc", + Region: "us-east-1", + LaunchableInfo: &store.LaunchableResponse{ + ID: "env-abc", + CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{ + CloudCredID: "cc-west", + InstanceType: "g5.xlarge", + Location: "us-west-2", + }, + }, + }, + org: mock.Org, + user: mock.User, + piped: true, + allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + { + Type: "g5.xlarge", + CloudCredID: "cc-west", + AvailableLocations: []string{"us-west-2"}, + }, + }, + }, + } + ctx.logf = func(_ string, _ ...interface{}) {} + + result := ctx.createInstancesWithType(InstanceSpec{Type: "g5.xlarge"}, 0, 1) + + assert.True(t, result.hadFailure) + assert.Empty(t, mock.CreatedOptions) } diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index a00aa3ce..0d7676e5 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -90,19 +90,37 @@ type AllInstanceTypesResponse struct { // GetCloudCredID returns the cloud credential ID for an instance type, or empty string if not found. func (r *AllInstanceTypesResponse) GetCloudCredID(instanceType string) string { + return r.GetCloudCredIDForRegion(instanceType, "") +} + +// GetCloudCredIDForRegion returns the cloud credential ID for an instance type +// available in region. Region matching is exact and case-insensitive. An empty +// region preserves the existing first-match behavior. +func (r *AllInstanceTypesResponse) GetCloudCredIDForRegion(instanceType, region string) string { + region = strings.TrimSpace(region) for _, it := range r.AllInstanceTypes { - if it.Type == instanceType { - if it.CloudCredID != "" { - return it.CloudCredID - } - if len(it.CloudCreds) > 0 { - return it.CloudCreds[0].ID - } + if it.Type != instanceType || (region != "" && !instanceTypeAvailableInRegion(it, region)) { + continue + } + if it.CloudCredID != "" { + return it.CloudCredID + } + if len(it.CloudCreds) > 0 { + return it.CloudCreds[0].ID } } return "" } +func instanceTypeAvailableInRegion(instanceType InstanceType, region string) bool { + for _, availableRegion := range instanceType.AvailableLocations { + if strings.EqualFold(availableRegion, region) { + return true + } + } + return false +} + // HasInstanceType reports whether the type exists in the API listing, independent of capacity. func (r *AllInstanceTypesResponse) HasInstanceType(instanceType string) bool { for _, it := range r.AllInstanceTypes { @@ -137,6 +155,9 @@ Features column shows instance capabilities: # Filter by GPU name (case-insensitive, partial match) brev search gpu --gpu-name A100 + # Filter by region/location + brev search gpu --region us-east-1 + # Filter by minimum VRAM per GPU (in GB) brev search gpu --min-vram 24 @@ -155,6 +176,9 @@ Features column shows instance capabilities: # Filter by provider brev search cpu --provider aws + # Filter by region/location + brev search cpu --region us-east-1 + # Filter by minimum RAM brev search cpu --min-ram 64 @@ -168,6 +192,7 @@ Features column shows instance capabilities: // sharedFlags holds flags shared between gpu and cpu subcommands type sharedFlags struct { + region string provider string arch string minVCPU int @@ -184,6 +209,7 @@ type sharedFlags struct { // addSharedFlags adds common flags to a command func addSharedFlags(cmd *cobra.Command, f *sharedFlags) { + cmd.Flags().StringVarP(&f.region, "region", "r", "", "Filter by region/location (case-insensitive, partial match)") cmd.Flags().StringVarP(&f.provider, "provider", "p", "", "Filter by provider/cloud (case-insensitive, partial match)") cmd.Flags().StringVar(&f.arch, "arch", "", "Filter by architecture (e.g., x86_64, arm64)") cmd.Flags().IntVar(&f.minVCPU, "min-vcpu", 0, "Minimum number of vCPUs") @@ -218,7 +244,7 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command Example: gpuExample, RunE: func(cmd *cobra.Command, args []string) error { // Default behavior: GPU search - return RunGPUSearch(t, store, gpuName, shared.provider, shared.arch, minVRAM, minTotalVRAM, minCapability, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput, wide) + return RunGPUSearch(t, store, gpuName, shared.region, shared.provider, shared.arch, minVRAM, minTotalVRAM, minCapability, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput, wide) }, } @@ -252,7 +278,7 @@ func newCmdGPUSubcommand(t *terminal.Terminal, store GPUSearchStore) *cobra.Comm Short: "Search GPU instance types", Example: gpuExample, RunE: func(cmd *cobra.Command, args []string) error { - return RunGPUSearch(t, store, gpuName, shared.provider, shared.arch, minVRAM, minTotalVRAM, minCapability, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput, wide) + return RunGPUSearch(t, store, gpuName, shared.region, shared.provider, shared.arch, minVRAM, minTotalVRAM, minCapability, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput, wide) }, } @@ -276,7 +302,7 @@ func newCmdCPUSubcommand(t *terminal.Terminal, store GPUSearchStore) *cobra.Comm Short: "Search CPU-only instance types", Example: cpuExample, RunE: func(cmd *cobra.Command, args []string) error { - return RunCPUSearch(t, store, shared.provider, shared.arch, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput) + return RunCPUSearch(t, store, shared.region, shared.provider, shared.arch, shared.minRAM, shared.minDisk, shared.minVCPU, shared.maxBootTime, shared.stoppable, shared.rebootable, shared.flexPorts, shared.sortBy, shared.descending, shared.jsonOutput) }, } @@ -287,28 +313,29 @@ func newCmdCPUSubcommand(t *terminal.Terminal, store GPUSearchStore) *cobra.Comm // GPUInstanceInfo holds processed GPU instance information for display type GPUInstanceInfo struct { - Type string `json:"type"` - Cloud string `json:"cloud"` // Underlying cloud (e.g., hyperstack, aws, gcp) - Provider string `json:"provider"` // Provider/aggregator (e.g., shadeform, aws, gcp) - GPUName string `json:"gpu_name"` - GPUCount int `json:"gpu_count"` - VRAMPerGPU float64 `json:"vram_per_gpu_gb"` - TotalVRAM float64 `json:"total_vram_gb"` - Capability float64 `json:"capability"` - VCPUs int `json:"vcpus"` - Memory string `json:"memory"` - RAMInGB float64 `json:"ram_gb"` - Arch string `json:"arch"` - DiskMin float64 `json:"disk_min_gb"` - DiskMax float64 `json:"disk_max_gb"` - DiskPricePerMo float64 `json:"disk_price_per_gb_mo,omitempty"` // $/GB/month for flexible storage - BootTime int `json:"boot_time_seconds"` - Stoppable bool `json:"stoppable"` - Rebootable bool `json:"rebootable"` - FlexPorts bool `json:"flex_ports"` - TargetDisk float64 `json:"target_disk_gb,omitempty"` - PricePerHour float64 `json:"price_per_hour"` - Manufacturer string `json:"-"` // exclude from JSON output + Type string `json:"type"` + Cloud string `json:"cloud"` // Underlying cloud (e.g., hyperstack, aws, gcp) + Provider string `json:"provider"` // Provider/aggregator (e.g., shadeform, aws, gcp) + GPUName string `json:"gpu_name"` + GPUCount int `json:"gpu_count"` + VRAMPerGPU float64 `json:"vram_per_gpu_gb"` + TotalVRAM float64 `json:"total_vram_gb"` + Capability float64 `json:"capability"` + VCPUs int `json:"vcpus"` + Memory string `json:"memory"` + RAMInGB float64 `json:"ram_gb"` + Arch string `json:"arch"` + DiskMin float64 `json:"disk_min_gb"` + DiskMax float64 `json:"disk_max_gb"` + DiskPricePerMo float64 `json:"disk_price_per_gb_mo,omitempty"` // $/GB/month for flexible storage + BootTime int `json:"boot_time_seconds"` + Stoppable bool `json:"stoppable"` + Rebootable bool `json:"rebootable"` + FlexPorts bool `json:"flex_ports"` + AvailableRegions []string `json:"available_regions,omitempty"` + TargetDisk float64 `json:"target_disk_gb,omitempty"` + PricePerHour float64 `json:"price_per_hour"` + Manufacturer string `json:"-"` // exclude from JSON output } // IsStdoutPiped returns true if stdout is being piped (not a terminal) @@ -318,7 +345,7 @@ func IsStdoutPiped() bool { } // RunGPUSearch executes the GPU search with filters and sorting -func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider, arch string, minVRAM, minTotalVRAM, minCapability, minRAM, minDisk float64, minVCPU, maxBootTime int, stoppable, rebootable, flexPorts bool, sortBy string, descending, jsonOutput, wide bool) error { +func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, region, provider, arch string, minVRAM, minTotalVRAM, minCapability, minRAM, minDisk float64, minVCPU, maxBootTime int, stoppable, rebootable, flexPorts bool, sortBy string, descending, jsonOutput, wide bool) error { if err := validateSortOption(sortBy); err != nil { return err } @@ -338,6 +365,7 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider, // Filter to GPU-only instances filtered := FilterInstances(instances, gpuName, provider, arch, minVRAM, minTotalVRAM, minCapability, minRAM, minDisk, minVCPU, maxBootTime, stoppable, rebootable, flexPorts, false) + filtered = FilterInstancesByRegion(filtered, region) if len(filtered) == 0 { return displayEmptyResults(t, "No GPU instances match the specified filters", jsonOutput, piped) @@ -349,7 +377,7 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider, } // RunCPUSearch executes the CPU search with filters and sorting -func RunCPUSearch(t *terminal.Terminal, store GPUSearchStore, provider, arch string, minRAM, minDisk float64, minVCPU, maxBootTime int, stoppable, rebootable, flexPorts bool, sortBy string, descending, jsonOutput bool) error { +func RunCPUSearch(t *terminal.Terminal, store GPUSearchStore, region, provider, arch string, minRAM, minDisk float64, minVCPU, maxBootTime int, stoppable, rebootable, flexPorts bool, sortBy string, descending, jsonOutput bool) error { if err := validateSortOption(sortBy); err != nil { return err } @@ -369,6 +397,7 @@ func RunCPUSearch(t *terminal.Terminal, store GPUSearchStore, provider, arch str // Filter to CPU-only instances filtered := FilterCPUInstances(instances, provider, arch, minRAM, minDisk, minVCPU, maxBootTime, stoppable, rebootable, flexPorts) + filtered = FilterInstancesByRegion(filtered, region) if len(filtered) == 0 { return displayEmptyResults(t, "No CPU instances match the specified filters", jsonOutput, piped) @@ -746,24 +775,25 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { if len(item.SupportedGPUs) == 0 { // CPU-only instance instances = append(instances, GPUInstanceInfo{ - Type: item.Type, - Cloud: extractCloud(item.Type, item.Provider), - Provider: item.Provider, - GPUName: "-", - GPUCount: 0, - VCPUs: item.VCPU, - Memory: item.Memory, - RAMInGB: ramInGB, - Arch: arch, - DiskMin: diskMin, - DiskMax: diskMax, - DiskPricePerMo: diskPricePerMo, - BootTime: bootTime, - Stoppable: item.Stoppable, - Rebootable: item.Rebootable, - FlexPorts: item.CanModifyFirewallRules, - PricePerHour: price, - Manufacturer: "cpu", + Type: item.Type, + Cloud: extractCloud(item.Type, item.Provider), + Provider: item.Provider, + GPUName: "-", + GPUCount: 0, + VCPUs: item.VCPU, + Memory: item.Memory, + RAMInGB: ramInGB, + Arch: arch, + DiskMin: diskMin, + DiskMax: diskMax, + DiskPricePerMo: diskPricePerMo, + BootTime: bootTime, + Stoppable: item.Stoppable, + Rebootable: item.Rebootable, + FlexPorts: item.CanModifyFirewallRules, + AvailableRegions: item.AvailableLocations, + PricePerHour: price, + Manufacturer: "cpu", }) continue } @@ -779,27 +809,28 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { capability := getGPUCapability(gpu.Name) instances = append(instances, GPUInstanceInfo{ - Type: item.Type, - Cloud: extractCloud(item.Type, item.Provider), - Provider: item.Provider, - GPUName: gpu.Name, - GPUCount: gpu.Count, - VRAMPerGPU: vramPerGPU, - TotalVRAM: totalVRAM, - Capability: capability, - VCPUs: item.VCPU, - Memory: item.Memory, - RAMInGB: ramInGB, - Arch: arch, - DiskMin: diskMin, - DiskMax: diskMax, - DiskPricePerMo: diskPricePerMo, - BootTime: bootTime, - Stoppable: item.Stoppable, - Rebootable: item.Rebootable, - FlexPorts: item.CanModifyFirewallRules, - PricePerHour: price, - Manufacturer: gpu.Manufacturer, + Type: item.Type, + Cloud: extractCloud(item.Type, item.Provider), + Provider: item.Provider, + GPUName: gpu.Name, + GPUCount: gpu.Count, + VRAMPerGPU: vramPerGPU, + TotalVRAM: totalVRAM, + Capability: capability, + VCPUs: item.VCPU, + Memory: item.Memory, + RAMInGB: ramInGB, + Arch: arch, + DiskMin: diskMin, + DiskMax: diskMax, + DiskPricePerMo: diskPricePerMo, + BootTime: bootTime, + Stoppable: item.Stoppable, + Rebootable: item.Rebootable, + FlexPorts: item.CanModifyFirewallRules, + AvailableRegions: item.AvailableLocations, + PricePerHour: price, + Manufacturer: gpu.Manufacturer, }) } } @@ -923,6 +954,26 @@ func FilterInstances(instances []GPUInstanceInfo, gpuName, provider, arch string return filtered } +// FilterInstancesByRegion returns instances available in a region. Search uses +// case-insensitive partial matching so users can discover provider region names. +func FilterInstancesByRegion(instances []GPUInstanceInfo, region string) []GPUInstanceInfo { + region = strings.TrimSpace(region) + if region == "" { + return instances + } + + var filtered []GPUInstanceInfo + for _, inst := range instances { + for _, availableRegion := range inst.AvailableRegions { + if strings.Contains(strings.ToLower(availableRegion), strings.ToLower(region)) { + filtered = append(filtered, inst) + break + } + } + } + return filtered +} + // FilterCPUInstances filters to CPU-only instances using shared filter logic func FilterCPUInstances(instances []GPUInstanceInfo, provider, arch string, minRAM, minDisk float64, minVCPU, maxBootTime int, stoppable, rebootable, flexPorts bool) []GPUInstanceInfo { // Filter out GPU instances first, then apply shared filters diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index cbc40f1d..aa4db977 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // MockGPUSearchStore is a mock implementation of GPUSearchStore for testing @@ -23,7 +24,8 @@ func createTestInstanceTypes() *InstanceTypesResponse { return &InstanceTypesResponse{ Items: []InstanceType{ { - Type: "g5.xlarge", + Type: "g5.xlarge", + AvailableLocations: []string{"us-east-1", "us-west-2"}, SupportedGPUs: []GPU{ {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, }, @@ -32,7 +34,8 @@ func createTestInstanceTypes() *InstanceTypesResponse { BasePrice: BasePrice{Currency: "USD", Amount: "1.006"}, }, { - Type: "g5.2xlarge", + Type: "g5.2xlarge", + AvailableLocations: []string{"eu-west-1"}, SupportedGPUs: []GPU{ {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, }, @@ -160,9 +163,25 @@ func TestProcessInstances(t *testing.T) { assert.Equal(t, 24.0, a10gInstance.TotalVRAM) assert.Equal(t, 8.6, a10gInstance.Capability) assert.Equal(t, 4, a10gInstance.VCPUs) + assert.Equal(t, []string{"us-east-1", "us-west-2"}, a10gInstance.AvailableRegions) assert.InDelta(t, 1.006, a10gInstance.PricePerHour, 0.001) } +func TestFilterInstancesByRegion(t *testing.T) { + instances := ProcessInstances(createTestInstanceTypes().Items) + + filtered := FilterInstancesByRegion(instances, "US-EAST") + require.Len(t, filtered, 1) + assert.Equal(t, "g5.xlarge", filtered[0].Type) + + filtered = FilterInstancesByRegion(instances, "west") + require.Len(t, filtered, 2) + assert.ElementsMatch(t, []string{"g5.xlarge", "g5.2xlarge"}, []string{filtered[0].Type, filtered[1].Type}) + + assert.Len(t, FilterInstancesByRegion(instances, ""), len(instances)) + assert.Empty(t, FilterInstancesByRegion(instances, "ap-south")) +} + func TestFilterInstancesByGPUName(t *testing.T) { response := createTestInstanceTypes() instances := ProcessInstances(response.Items) @@ -590,8 +609,14 @@ func TestAllInstanceTypesResponseLookup(t *testing.T) { resp := &AllInstanceTypesResponse{ AllInstanceTypes: []InstanceType{ { - Type: "hyperstack_H100_sxm5x8", - CloudCredID: "cc-shadeform", + Type: "hyperstack_H100_sxm5x8", + CloudCredID: "cc-shadeform-east", + AvailableLocations: []string{"us-east-1"}, + }, + { + Type: "hyperstack_H100_sxm5x8", + CloudCredID: "cc-shadeform-west", + AvailableLocations: []string{"us-west-2"}, }, { Type: "hyperstack_H100x8_NVLINK", @@ -603,7 +628,12 @@ func TestAllInstanceTypesResponseLookup(t *testing.T) { } t.Run("GetCloudCredID returns the cloud credential instead of the workspace group", func(t *testing.T) { - assert.Equal(t, "cc-shadeform", resp.GetCloudCredID("hyperstack_H100_sxm5x8")) + assert.Equal(t, "cc-shadeform-east", resp.GetCloudCredID("hyperstack_H100_sxm5x8")) + }) + + t.Run("GetCloudCredIDForRegion selects the credential offering the requested region", func(t *testing.T) { + assert.Equal(t, "cc-shadeform-west", resp.GetCloudCredIDForRegion("hyperstack_H100_sxm5x8", "US-WEST-2")) + assert.Empty(t, resp.GetCloudCredIDForRegion("hyperstack_H100_sxm5x8", "eu-west-1")) }) t.Run("GetCloudCredID returns empty when no cloud credential is available", func(t *testing.T) { From 81c94d8ed8439b3955ac62462b68a47cf68d390d Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Sat, 22 Aug 2026 15:52:37 -0700 Subject: [PATCH 2/2] feat: complete location-aware instance creation --- pkg/cmd/gpucreate/gpucreate.go | 212 +++++++++++++++++++---- pkg/cmd/gpucreate/gpucreate_test.go | 249 +++++++++++++++++++++++----- pkg/cmd/gpusearch/gpusearch.go | 39 ++++- pkg/cmd/gpusearch/gpusearch_test.go | 112 +++++++++++++ 4 files changed, 532 insertions(+), 80 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index a0adbc66..a289a642 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -74,7 +74,12 @@ You can attach a startup script that runs when the instance boots using the --startup-script flag. The script can be provided as: - An inline string: --startup-script 'pip install torch' - A file path (prefix with @): --startup-script @setup.sh - - An absolute file path: --startup-script @/path/to/setup.sh` + - An absolute file path: --startup-script @/path/to/setup.sh + +Placement: +Use --location (or --region) to choose a catalog location and --sub-location +to choose a provider zone within it. The location is validated before create, +and dry runs show either the selected placement or the catalog default.` example = ` # Create an instance using smart defaults (sorted by price) @@ -83,8 +88,14 @@ You can attach a startup script that runs when the instance boots using the # Create with a specific GPU type brev create my-instance --type g5.xlarge - # Create in a specific region - brev create my-instance --type g5.xlarge --region us-east-1 + # Create in a specific location + brev create my-instance --type g5.xlarge --location us-east-1 + + # Create a CPU instance in a specific location and zone + brev create my-cpu --type n2d-standard-2 --location us-west2 --sub-location us-west2-a + + # Preview the type, disk, and selected/default placement without provisioning + brev create my-instance --type g5.xlarge --dry-run # Try multiple types in order (fallback chain) brev create my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge @@ -138,6 +149,8 @@ type CreateResult struct { type searchFilterFlags struct { gpuName string region string + location string + subLocation string provider string minVRAM float64 minTotalVRAM float64 @@ -155,7 +168,7 @@ type searchFilterFlags struct { // hasUserFilters returns true if the user specified any search filter flags func (f *searchFilterFlags) hasUserFilters() bool { - return f.gpuName != "" || f.region != "" || f.provider != "" || f.minVRAM > 0 || f.minTotalVRAM > 0 || + return f.gpuName != "" || f.region != "" || f.location != "" || f.subLocation != "" || f.provider != "" || f.minVRAM > 0 || f.minTotalVRAM > 0 || f.minCapability > 0 || f.minDisk > 0 || f.maxBootTime > 0 || f.stoppable || f.rebootable || f.flexPorts } @@ -237,8 +250,17 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra return breverrors.WrapAndTrace(err) } - if strings.TrimSpace(filters.region) != "" { - response, err := gpuCreateStore.GetInstanceTypes(false) + filters.region, err = resolveCreateLocation(filters.region, filters.location) + if err != nil { + return err + } + filters.subLocation = strings.TrimSpace(filters.subLocation) + if filters.subLocation != "" && filters.region == "" { + return breverrors.NewValidationError("--sub-location requires --location or --region") + } + + if filters.region != "" || dryRun { + response, err := gpuCreateStore.GetInstanceTypes(true) if err != nil { return breverrors.WrapAndTrace(err) } @@ -246,9 +268,11 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra if response != nil { filters.catalogItems = response.Items } - filters.region, err = canonicalRegion(filters.region, filters.catalogItems) - if err != nil { - return err + if filters.region != "" { + filters.region, err = canonicalRegion(filters.region, types, filters.catalogItems) + if err != nil { + return err + } } } @@ -274,6 +298,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra LaunchableInfo: launchableInfo, ParameterBindings: parameterBindings, Region: filters.region, + SubLocation: filters.subLocation, } opts.InstanceTypes, err = resolveInstanceTypes(cmd, gpuCreateStore, opts, types, &filters) @@ -285,7 +310,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra } if dryRun { - return runDryRun(t, gpuCreateStore, opts.InstanceTypes, &filters) + return runDryRun(t, gpuCreateStore, opts, &filters) } return RunGPUCreate(t, gpuCreateStore, opts) @@ -308,6 +333,21 @@ func validateArgs(name string, count int) error { return nil } +func resolveCreateLocation(region, location string) (string, error) { + region = strings.TrimSpace(region) + location = strings.TrimSpace(location) + if region != "" && location != "" && !strings.EqualFold(region, location) { + return "", breverrors.NewValidationError(fmt.Sprintf( + "--region %q conflicts with --location %q; specify only one location", + region, location, + )) + } + if location != "" { + return location, nil + } + return region, nil +} + // registerCreateFlags registers all flags for the create command func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, parallel *int, detached *bool, timeout *int, startupScript *string, dryRun *bool, mode *string, jupyter *bool, containerImage, composeFile, launchable *string, launchableParams *[]string, filters *searchFilterFlags) { cmd.Flags().StringVarP(name, "name", "n", "", "Base name for the instances (or pass as first argument)") @@ -328,6 +368,8 @@ func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, cmd.Flags().StringArrayVar(launchableParams, "param", nil, "Launchable setup value NAME=VALUE (repeatable)") cmd.Flags().StringVarP(&filters.region, "region", "r", "", "Region/location to deploy the instance (e.g., us-east-1, us-central1)") + cmd.Flags().StringVar(&filters.location, "location", "", "Location/region to deploy the instance (alias of --region)") + cmd.Flags().StringVar(&filters.subLocation, "sub-location", "", "Provider sub-location or zone (requires --location or --region)") cmd.Flags().StringVarP(&filters.gpuName, "gpu-name", "g", "", "Filter by GPU name (e.g., A100, H100)") cmd.Flags().StringVar(&filters.provider, "provider", "", "Filter by provider/cloud (e.g., aws, gcp)") cmd.Flags().Float64VarP(&filters.minVRAM, "min-vram", "v", 0, "Minimum VRAM per GPU in GB") @@ -344,8 +386,10 @@ func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, // InstanceSpec holds an instance type and its target disk size type InstanceSpec struct { - Type string - DiskGB float64 // Target disk size in GB, 0 means use default + Type string + DiskGB float64 // Target disk size in GB, 0 means use default + Location string + SubLocation string } // GPUCreateOptions holds the options for GPU instance creation @@ -366,6 +410,7 @@ type GPUCreateOptions struct { LaunchableInfo *store.LaunchableResponse // populated when LaunchableID is set ParameterBindings []store.ParameterBinding Region string + SubLocation string } // parseLaunchableID extracts a launchable ID from either a raw ID (env-XXX) or @@ -509,7 +554,8 @@ func warnLaunchableFlagConflicts(cmd *cobra.Command, t *terminal.Terminal, launc instanceFlagsSet := cmd.Flags().Changed("type") || cmd.Flags().Changed("gpu-name") || cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram") || - cmd.Flags().Changed("region") + cmd.Flags().Changed("region") || cmd.Flags().Changed("location") || + cmd.Flags().Changed("sub-location") if instanceFlagsSet { t.Vprintf("Warning: Overriding the launchable's recommended instance configuration. This is not the recommended path and may cause issues.\n\n") } @@ -752,15 +798,32 @@ func getFilteredInstanceTypes(s GPUCreateStore, filters *searchFilterFlags) ([]I if inst.DiskMin != inst.DiskMax && minDisk > inst.DiskMin && minDisk <= inst.DiskMax { diskGB = minDisk } - specs = append(specs, InstanceSpec{Type: inst.Type, DiskGB: diskGB}) + specs = append(specs, InstanceSpec{ + Type: inst.Type, + DiskGB: diskGB, + Location: inst.Location, + SubLocation: inst.SubLocation, + }) } return specs, nil } // runDryRun shows the instance types that would be used without creating anything -func runDryRun(t *terminal.Terminal, s GPUCreateStore, specs []InstanceSpec, filters *searchFilterFlags) error { - if len(specs) > 0 { +func runDryRun(t *terminal.Terminal, s GPUCreateStore, opts GPUCreateOptions, filters *searchFilterFlags) error { + items := filters.catalogItems + if !filters.catalogLoaded { + response, err := s.GetInstanceTypes(true) + if err != nil { + return breverrors.WrapAndTrace(err) + } + if response != nil { + items = response.Items + } + } + + if len(opts.InstanceTypes) > 0 { + specs := annotateInstanceSpecs(opts.InstanceTypes, opts.Region, opts.SubLocation, items) t.Print(formatInstanceSpecs(specs)) return nil } @@ -779,7 +842,7 @@ func runDryRun(t *terminal.Terminal, s GPUCreateStore, specs []InstanceSpec, fil // canonicalRegion returns the catalog spelling for an exact, case-insensitive // region match. Create must not forward a partial region name to the provider. -func canonicalRegion(region string, items []gpusearch.InstanceType) (string, error) { +func canonicalRegion(region string, types []InstanceSpec, items []gpusearch.InstanceType) (string, error) { region = strings.TrimSpace(region) for _, item := range items { for _, availableRegion := range item.AvailableLocations { @@ -789,8 +852,13 @@ func canonicalRegion(region string, items []gpusearch.InstanceType) (string, err } } + typeNames := make([]string, 0, len(types)) + for _, spec := range types { + typeNames = append(typeNames, spec.Type) + } return "", breverrors.NewValidationError(fmt.Sprintf( - "region %q is not available; run 'brev search --json' to list valid regions", region, + "location %q is not available; available locations: %s", + region, formatAvailableLocations(availableLocationsForTypes(typeNames, items)), )) } @@ -827,11 +895,51 @@ func validateTypesSupportRegion(types []InstanceSpec, region string, items []gpu sort.Strings(unsupported) return breverrors.NewValidationError(fmt.Sprintf( - "region %q is not available for instance type(s) %s; run 'brev search --region %s' to find compatible types", - region, strings.Join(unsupported, ", "), region, + "location %q is not available for instance type(s) %s; available locations: %s", + region, strings.Join(unsupported, ", "), formatAvailableLocations(availableLocationsForTypes(unsupported, items)), )) } +func availableLocationsForTypes(types []string, items []gpusearch.InstanceType) []string { + typeFilter := make(map[string]struct{}, len(types)) + for _, instanceType := range types { + typeFilter[instanceType] = struct{}{} + } + + locations := make(map[string]string) + for _, item := range items { + if len(typeFilter) > 0 { + if _, ok := typeFilter[item.Type]; !ok { + continue + } + } + for _, location := range item.AvailableLocations { + location = strings.TrimSpace(location) + if location == "" { + continue + } + key := strings.ToLower(location) + if _, exists := locations[key]; !exists { + locations[key] = location + } + } + } + + result := make([]string, 0, len(locations)) + for _, location := range locations { + result = append(result, location) + } + sort.Strings(result) + return result +} + +func formatAvailableLocations(locations []string) string { + if len(locations) == 0 { + return "none reported by the catalog" + } + return strings.Join(locations, ", ") +} + func instanceTypeSupportsRegion(item gpusearch.InstanceType, region string) bool { for _, availableRegion := range item.AvailableLocations { if strings.EqualFold(availableRegion, region) { @@ -915,7 +1023,7 @@ func parseJSONInput(input string) ([]InstanceSpec, error) { } // parseTableInput parses table format input from gpu-search -// Table format: TYPE TARGET_DISK PROVIDER GPU COUNT ... +// Table format: TYPE TARGET_DISK PROVIDER DEFAULT_LOCATION GPU/CPU fields ... func parseTableInput(input string) []InstanceSpec { var specs []InstanceSpec lines := strings.Split(input, "\n") @@ -938,7 +1046,7 @@ func parseTableInput(input string) []InstanceSpec { } // Extract TYPE (column 0) and TARGET_DISK (column 1) from the table output - // The format is: TYPE TARGET_DISK PROVIDER GPU COUNT ... + // The format starts with TYPE and TARGET_DISK; later discovery columns are ignored. fields := strings.Fields(line) if len(fields) > 0 { instanceType := fields[0] @@ -983,15 +1091,46 @@ func isValidInstanceType(s string) bool { func formatInstanceSpecs(specs []InstanceSpec) string { var parts []string for _, spec := range specs { + var details []string if spec.DiskGB > 0 { - parts = append(parts, fmt.Sprintf("%s (%.0fGB disk)", spec.Type, spec.DiskGB)) - } else { + details = append(details, fmt.Sprintf("%.0fGB disk", spec.DiskGB)) + } + if spec.Location != "" { + details = append(details, "location: "+spec.Location) + } + if spec.SubLocation != "" { + details = append(details, "sub-location: "+spec.SubLocation) + } + if len(details) == 0 { parts = append(parts, spec.Type) + continue } + parts = append(parts, fmt.Sprintf("%s (%s)", spec.Type, strings.Join(details, ", "))) } return strings.Join(parts, ", ") } +func annotateInstanceSpecs(specs []InstanceSpec, location, subLocation string, items []gpusearch.InstanceType) []InstanceSpec { + annotated := append([]InstanceSpec(nil), specs...) + for i := range annotated { + if location != "" { + annotated[i].Location = location + annotated[i].SubLocation = subLocation + continue + } + + for _, item := range items { + if item.Type != annotated[i].Type { + continue + } + annotated[i].Location = item.Location + annotated[i].SubLocation = item.SubLocation + break + } + } + return annotated +} + // createContext holds shared state for instance creation type createContext struct { t *terminal.Terminal @@ -1077,9 +1216,10 @@ func (c *createContext) validateInstanceTypeAvailability(instanceType string) er )) } if c.opts.Region != "" { + locations := availableLocationsForTypes([]string{instanceType}, c.allInstanceTypes.AllInstanceTypes) return breverrors.NewValidationError(fmt.Sprintf( - "instance type %q is unavailable in region %q; run 'brev search --region %s' to find a compatible type", - instanceType, c.opts.Region, c.opts.Region, + "instance type %q is unavailable in location %q; available locations: %s", + instanceType, c.opts.Region, formatAvailableLocations(locations), )) } return breverrors.NewValidationError(fmt.Sprintf( @@ -1254,6 +1394,15 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC if err != nil { return err } + if ctx.allInstanceTypes != nil { + opts.InstanceTypes = annotateInstanceSpecs( + opts.InstanceTypes, + opts.Region, + opts.SubLocation, + ctx.allInstanceTypes.AllInstanceTypes, + ) + ctx.opts = opts + } // Auto-redeem coupon code if the launchable has one attached. // This is silent — the UI doesn't surface coupon redemption to the user either. @@ -1327,9 +1476,10 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity if cloudCredID := c.allInstanceTypes.GetCloudCredIDForRegion(spec.Type, c.opts.Region); cloudCredID != "" { cwOptions.WithCloudCredID(cloudCredID) } else if c.opts.Region != "" && c.allInstanceTypes.HasInstanceType(spec.Type) { + locations := availableLocationsForTypes([]string{spec.Type}, c.allInstanceTypes.AllInstanceTypes) return nil, breverrors.NewValidationError(fmt.Sprintf( - "instance type %q has no available cloud credential in region %q; run 'brev search --region %s' to find a compatible type", - spec.Type, c.opts.Region, c.opts.Region, + "instance type %q has no available cloud credential in location %q; available locations: %s", + spec.Type, c.opts.Region, formatAvailableLocations(locations), )) } } @@ -1346,9 +1496,9 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity if c.opts.Region != "" { cwOptions.Location = c.opts.Region - // A launchable sub-location belongs to its original location. Let the - // provider choose a compatible zone when the CLI overrides the region. - cwOptions.SubLocation = "" + // A launchable sub-location belongs to its original location. Replace it + // with the explicit CLI value, or clear it so the provider can choose. + cwOptions.SubLocation = c.opts.SubLocation } if cwOptions.CloudCredID == "" { diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index 3b9badf5..1da18271 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -2,8 +2,10 @@ package gpucreate import ( "encoding/json" + "io" "net/http" "net/http/httptest" + "os" "strings" "testing" "time" @@ -30,6 +32,7 @@ type MockGPUCreateStore struct { DeletedWorkspaceIDs []string FetchedLifeCycleScriptIDs []string AllInstanceTypes *gpusearch.AllInstanceTypesResponse + InstanceTypesIncludeCPU []bool } func NewMockGPUCreateStore() *MockGPUCreateStore { @@ -130,28 +133,84 @@ func (m *MockGPUCreateStore) RedeemCouponCode(organizationID string, code string return &store.RedeemCouponCodeResponse{}, nil } -func (m *MockGPUCreateStore) GetInstanceTypes(_ bool) (*gpusearch.InstanceTypesResponse, error) { - // Return a default set of instance types for testing - return &gpusearch.InstanceTypesResponse{ - Items: []gpusearch.InstanceType{ - { - Type: "g5.xlarge", - AvailableLocations: []string{"us-east-1", "us-west-2"}, - SupportedGPUs: []gpusearch.GPU{ - {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, - }, - SupportedStorage: []gpusearch.Storage{ - {Size: "500GiB"}, - }, - Memory: "16GiB", - VCPU: 4, - BasePrice: gpusearch.BasePrice{Currency: "USD", Amount: "1.006"}, - EstimatedDeployTime: "5m0s", +func (m *MockGPUCreateStore) GetInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, error) { + m.InstanceTypesIncludeCPU = append(m.InstanceTypesIncludeCPU, includeCPU) + items := []gpusearch.InstanceType{ + { + Type: "g5.xlarge", + Location: "us-east-1", + SubLocation: "us-east-1a", + AvailableLocations: []string{"us-east-1", "us-west-2"}, + SupportedGPUs: []gpusearch.GPU{ + {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, }, + SupportedStorage: []gpusearch.Storage{ + {Size: "500GiB"}, + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: gpusearch.BasePrice{Currency: "USD", Amount: "1.006"}, + EstimatedDeployTime: "5m0s", }, + } + if includeCPU { + items = append(items, gpusearch.InstanceType{ + Type: "n2d-standard-2", + Location: "asia-south1", + SubLocation: "asia-south1-a", + AvailableLocations: []string{"asia-south1", "us-west2"}, + SupportedStorage: []gpusearch.Storage{ + {MinSize: "10GiB", MaxSize: "65536GiB"}, + }, + Memory: "8GiB", + VCPU: 2, + BasePrice: gpusearch.BasePrice{Currency: "USD", Amount: "0.097"}, + EstimatedDeployTime: "5m0s", + }) + } + + // Return a default set of instance types for testing + return &gpusearch.InstanceTypesResponse{ + Items: items, }, nil } +func executeCreateAndCaptureOutput(t *testing.T, mock *MockGPUCreateStore, args ...string) string { + t.Helper() + + oldStdout := os.Stdout + oldStderr := os.Stderr + stdoutReader, stdoutWriter, err := os.Pipe() + require.NoError(t, err) + stderrReader, stderrWriter, err := os.Pipe() + require.NoError(t, err) + os.Stdout = stdoutWriter + os.Stderr = stderrWriter + defer func() { + os.Stdout = oldStdout + os.Stderr = oldStderr + _ = stdoutReader.Close() + _ = stdoutWriter.Close() + _ = stderrReader.Close() + _ = stderrWriter.Close() + }() + + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs(args) + executeErr := cmd.Execute() + require.NoError(t, executeErr) + require.NoError(t, stdoutWriter.Close()) + require.NoError(t, stderrWriter.Close()) + os.Stdout = oldStdout + os.Stderr = oldStderr + + stdout, readErr := io.ReadAll(stdoutReader) + require.NoError(t, readErr) + stderr, readErr := io.ReadAll(stderrReader) + require.NoError(t, readErr) + return string(stdout) + string(stderr) +} + func TestIsValidInstanceType(t *testing.T) { tests := []struct { name string @@ -885,25 +944,109 @@ func TestMockGPUCreateStoreTypeSpecificError(t *testing.T) { func TestCreateDryRunWithExplicitTypesDoesNotProvision(t *testing.T) { mock := NewMockGPUCreateStore() - term := terminal.New() - - cmd := NewCmdGPUCreate(term, mock) - cmd.SetArgs([]string{"dry-run-test", "--type", "g5.xlarge", "--dry-run"}) + output := executeCreateAndCaptureOutput(t, mock, "dry-run-test", "--type", "g5.xlarge", "--dry-run") - err := cmd.Execute() - assert.NoError(t, err) assert.Empty(t, mock.CreatedWorkspaces) + assert.Contains(t, output, "g5.xlarge") + assert.Contains(t, output, "location: us-east-1") + assert.Contains(t, output, "sub-location: us-east-1a") + assert.Equal(t, []bool{true}, mock.InstanceTypesIncludeCPU, "dry-run catalog must include CPU types") } func TestCreateDryRunAcceptsCanonicalRegion(t *testing.T) { mock := NewMockGPUCreateStore() - cmd := NewCmdGPUCreate(terminal.New(), mock) - cmd.SetArgs([]string{"dry-run-region", "--type", "g5.xlarge", "--region", "US-EAST-1", "--dry-run"}) + output := executeCreateAndCaptureOutput(t, mock, "dry-run-region", "--type", "g5.xlarge", "--region", "US-EAST-1", "--dry-run") - err := cmd.Execute() + assert.Empty(t, mock.CreatedWorkspaces) + assert.Contains(t, output, "location: us-east-1") + assert.NotContains(t, output, "sub-location:", "an explicit location without a zone must clear the catalog default zone") +} + +func TestCreateCPUDryRunShowsDefaultLocation(t *testing.T) { + mock := NewMockGPUCreateStore() + output := executeCreateAndCaptureOutput(t, mock, "cpu-dry-run", "--type", "n2d-standard-2", "--dry-run") - assert.NoError(t, err) assert.Empty(t, mock.CreatedWorkspaces) + assert.Contains(t, output, "n2d-standard-2") + assert.Contains(t, output, "location: asia-south1") + assert.Contains(t, output, "sub-location: asia-south1-a") +} + +func TestCreateCPUWithLocationAndSubLocation(t *testing.T) { + mock := NewMockGPUCreateStore() + mock.AllInstanceTypes = &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + { + Type: "n2d-standard-2", + CloudCredID: "cc-gcp-west", + Location: "asia-south1", + SubLocation: "asia-south1-a", + AvailableLocations: []string{"asia-south1", "us-west2"}, + }, + }, + } + + output := executeCreateAndCaptureOutput( + t, + mock, + "cpu-west", + "--type", "n2d-standard-2", + "--location", "us-west2", + "--sub-location", "us-west2-a", + "--detached", + ) + + require.Len(t, mock.CreatedOptions, 1) + assert.Equal(t, "n2d-standard-2", mock.CreatedOptions[0].InstanceType) + assert.Equal(t, "us-west2", mock.CreatedOptions[0].Location) + assert.Equal(t, "us-west2-a", mock.CreatedOptions[0].SubLocation) + assert.Equal(t, "cc-gcp-west", mock.CreatedOptions[0].CloudCredID) + assert.Contains(t, output, "location: us-west2") + assert.Contains(t, output, "sub-location: us-west2-a") + payload, marshalErr := json.Marshal(mock.CreatedOptions[0]) + require.NoError(t, marshalErr) + assert.Contains(t, string(payload), `"location":"us-west2"`) + assert.Contains(t, string(payload), `"subLocation":"us-west2-a"`) +} + +func TestCreateLocationFlagValidation(t *testing.T) { + t.Run("conflicting aliases are rejected", func(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"conflict", "--type", "g5.xlarge", "--region", "us-east-1", "--location", "us-west-2", "--dry-run"}) + + err := cmd.Execute() + + assert.ErrorContains(t, err, `--region "us-east-1" conflicts with --location "us-west-2"`) + assert.Empty(t, mock.CreatedWorkspaces) + }) + + t.Run("sub-location requires a location", func(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"orphan-zone", "--type", "n2d-standard-2", "--sub-location", "us-west2-a", "--dry-run"}) + + err := cmd.Execute() + + assert.ErrorContains(t, err, "--sub-location requires --location or --region") + assert.Empty(t, mock.CreatedWorkspaces) + }) + + t.Run("equivalent aliases are accepted", func(t *testing.T) { + mock := NewMockGPUCreateStore() + output := executeCreateAndCaptureOutput( + t, + mock, + "same-location", + "--type", "g5.xlarge", + "--region", "US-EAST-1", + "--location", "us-east-1", + "--dry-run", + ) + + assert.Contains(t, output, "location: us-east-1") + assert.Empty(t, mock.CreatedWorkspaces) + }) } func TestCreateAutoSelectsTypeAndCredentialForRegion(t *testing.T) { @@ -941,7 +1084,23 @@ func TestCreateRejectsUnknownRegion(t *testing.T) { err := cmd.Execute() - assert.ErrorContains(t, err, `region "moon-1" is not available`) + assert.ErrorContains(t, err, `location "moon-1" is not available`) + assert.Contains(t, err.Error(), "us-east-1") + assert.Contains(t, err.Error(), "us-west-2") + assert.Empty(t, mock.CreatedWorkspaces) +} + +func TestCreateRejectsUnsupportedCPULocationWithSupportedList(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"cpu-wrong-location", "--type", "n2d-standard-2", "--location", "us-east-1", "--dry-run"}) + + err := cmd.Execute() + + assert.ErrorContains(t, err, `location "us-east-1" is not available for instance type(s) n2d-standard-2`) + assert.Contains(t, err.Error(), "asia-south1") + assert.Contains(t, err.Error(), "us-west2") + assert.NotContains(t, err.Error(), "us-west-2", "the error should list locations for the requested CPU type, not the entire catalog") assert.Empty(t, mock.CreatedWorkspaces) } @@ -952,7 +1111,7 @@ func TestRegionValidation(t *testing.T) { {Type: "g6.xlarge", AvailableLocations: []string{"eu-west-1"}}, } - region, err := canonicalRegion(" US-EAST-1 ", items) + region, err := canonicalRegion(" US-EAST-1 ", []InstanceSpec{{Type: "g5.xlarge"}}, items) require.NoError(t, err) assert.Equal(t, "us-east-1", region) assert.NoError(t, validateTypesSupportRegion([]InstanceSpec{{Type: "g5.xlarge"}}, region, items)) @@ -997,10 +1156,10 @@ func TestGetFilteredInstanceTypesNoMatch(t *testing.T) { func TestParseTableInput(t *testing.T) { tableInput := strings.Join([]string{ - "TYPE TARGET_DISK GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR", - "g5.xlarge 500 A10G 1 24 GB 24 GB 8.6 4 $1.01", - "g5.2xlarge 500 A10G 1 24 GB 24 GB 8.6 8 $1.21", - "p4d.24xlarge 1000 A100 8 40 GB 320 GB 8.0 96 $32.77", + "TYPE TARGET_DISK PROVIDER DEFAULT_LOCATION GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR", + "g5.xlarge 500 aws us-east-1/us-east-1a A10G 1 24 GB 24 GB 8.6 4 $1.01", + "g5.2xlarge 500 aws us-west-2 A10G 1 24 GB 24 GB 8.6 8 $1.21", + "p4d.24xlarge 1000 aws us-east-1 A100 8 40 GB 320 GB 8.0 96 $32.77", "", "Found 3 GPU instance types", }, "\n") @@ -1019,10 +1178,10 @@ func TestParseTableInput(t *testing.T) { func TestParseTableInputCPU(t *testing.T) { // Simulated plain table output from `brev search cpu` tableInput := strings.Join([]string{ - " TYPE TARGET_DISK PROVIDER VCPUS RAM ARCH DISK $/GB/MO BOOT FEATURES $/HR", - " n2d-highcpu-2 10 gcp 2 2 x86_64 10GB-16TB $0.13 7m SP $0.05", - " n1-standard-1 10 gcp 1 4 x86_64 10GB-16TB $0.14 7m SP $0.06", - " m8i-flex.8xlarge 500 aws 32 128 x86_64 10GB-16TB $0.10 7m SRP $1.93", + " TYPE TARGET_DISK PROVIDER DEFAULT_LOCATION VCPUS RAM ARCH DISK $/GB/MO BOOT FEATURES $/HR", + " n2d-standard-2 10 gcp asia-south1/asia-south1-a 2 8 x86_64 10GB-16TB $0.13 7m SP $0.05", + " n1-standard-1 10 gcp us-west2 1 4 x86_64 10GB-16TB $0.14 7m SP $0.06", + " m8i-flex.8xlarge 500 aws us-west-2 32 128 x86_64 10GB-16TB $0.10 7m SRP $1.93", "", "Found 3 CPU instance types", }, "\n") @@ -1030,7 +1189,7 @@ func TestParseTableInputCPU(t *testing.T) { specs := parseTableInput(tableInput) assert.Len(t, specs, 3) - assert.Equal(t, "n2d-highcpu-2", specs[0].Type) + assert.Equal(t, "n2d-standard-2", specs[0].Type) assert.Equal(t, 10.0, specs[0].DiskGB) assert.Equal(t, "n1-standard-1", specs[1].Type) assert.Equal(t, 10.0, specs[1].DiskGB) @@ -1045,6 +1204,8 @@ func TestParseJSONInput(t *testing.T) { "type": "g5.xlarge", "provider": "aws", "gpu_name": "A10G", + "location": "us-east-1", + "sub_location": "us-east-1a", "target_disk_gb": 1000 }, { @@ -1067,6 +1228,8 @@ func TestParseJSONInput(t *testing.T) { // Check first instance with disk assert.Equal(t, "g5.xlarge", specs[0].Type) assert.Equal(t, 1000.0, specs[0].DiskGB) + assert.Empty(t, specs[0].Location, "search placement metadata must not silently become an explicit create override") + assert.Empty(t, specs[0].SubLocation) // Check second instance with different disk assert.Equal(t, "p4d.24xlarge", specs[1].Type) @@ -1080,12 +1243,12 @@ func TestParseJSONInput(t *testing.T) { func TestFormatInstanceSpecs(t *testing.T) { specs := []InstanceSpec{ {Type: "g5.xlarge", DiskGB: 1000}, - {Type: "p4d.24xlarge", DiskGB: 0}, - {Type: "g6.xlarge", DiskGB: 500}, + {Type: "p4d.24xlarge", DiskGB: 0, Location: "us-west-2"}, + {Type: "g6.xlarge", DiskGB: 500, Location: "us-west2", SubLocation: "us-west2-a"}, } result := formatInstanceSpecs(specs) - assert.Equal(t, "g5.xlarge (1000GB disk), p4d.24xlarge, g6.xlarge (500GB disk)", result) + assert.Equal(t, "g5.xlarge (1000GB disk), p4d.24xlarge (location: us-west-2), g6.xlarge (500GB disk, location: us-west2, sub-location: us-west2-a)", result) } func TestPollUntilReadyReportsWorkspaceFailureMessage(t *testing.T) { @@ -1245,8 +1408,8 @@ func TestValidateInstanceTypeAvailability(t *testing.T) { }, } err := ctx.validateInstanceTypeAvailability("hyperstack_H100_sxm5x8") - assert.ErrorContains(t, err, `unavailable in region "us-east-1"`) - assert.Contains(t, err.Error(), "brev search --region us-east-1") + assert.ErrorContains(t, err, `unavailable in location "us-east-1"`) + assert.Contains(t, err.Error(), "available locations: us-west-2") }) t.Run("error type is ValidationError so no stack trace is appended", func(t *testing.T) { diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index 0d7676e5..fa511b7c 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -141,6 +141,8 @@ var ( Use 'brev search gpu' (default) to find GPU instances. Use 'brev search cpu' to find CPU-only instances. +Search output includes each type's default location, and JSON output also +includes its default location, sub-location, and all available regions. Features column shows instance capabilities: S = Stoppable (can stop and restart without losing data) @@ -332,6 +334,8 @@ type GPUInstanceInfo struct { Stoppable bool `json:"stoppable"` Rebootable bool `json:"rebootable"` FlexPorts bool `json:"flex_ports"` + Location string `json:"location,omitempty"` + SubLocation string `json:"sub_location,omitempty"` AvailableRegions []string `json:"available_regions,omitempty"` TargetDisk float64 `json:"target_disk_gb,omitempty"` PricePerHour float64 `json:"price_per_hour"` @@ -791,6 +795,8 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { Stoppable: item.Stoppable, Rebootable: item.Rebootable, FlexPorts: item.CanModifyFirewallRules, + Location: item.Location, + SubLocation: item.SubLocation, AvailableRegions: item.AvailableLocations, PricePerHour: price, Manufacturer: "cpu", @@ -828,6 +834,8 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { Stoppable: item.Stoppable, Rebootable: item.Rebootable, FlexPorts: item.CanModifyFirewallRules, + Location: item.Location, + SubLocation: item.SubLocation, AvailableRegions: item.AvailableLocations, PricePerHour: price, Manufacturer: gpu.Manufacturer, @@ -1163,13 +1171,26 @@ func formatInstanceFields(inst GPUInstanceInfo, includeUnits bool) formattedInst } } +func formatDefaultLocation(location, subLocation string) string { + switch { + case location != "" && subLocation != "": + return location + "/" + subLocation + case location != "": + return location + case subLocation != "": + return subLocation + default: + return "-" + } +} + // displayGPUTable renders the GPU instances as a table func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { ta := table.NewWriter() ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "DEFAULT LOCATION", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1177,6 +1198,7 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { row := table.Row{ inst.Type, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), t.Green(inst.GPUName), inst.GPUCount, f.VRAM, @@ -1203,7 +1225,7 @@ func displayGPUTablePlain(instances []GPUInstanceInfo) { ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL_VRAM", "CAPABILITY", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "DEFAULT_LOCATION", "GPU", "COUNT", "VRAM/GPU", "TOTAL_VRAM", "CAPABILITY", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1212,6 +1234,7 @@ func displayGPUTablePlain(instances []GPUInstanceInfo) { inst.Type, f.TargetDisk, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), inst.GPUName, inst.GPUCount, f.VRAM, @@ -1236,7 +1259,7 @@ func displayGPUTableWide(t *terminal.Terminal, instances []GPUInstanceInfo) { ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "DEFAULT LOCATION", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1244,6 +1267,7 @@ func displayGPUTableWide(t *terminal.Terminal, instances []GPUInstanceInfo) { row := table.Row{ inst.Type, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), t.Green(inst.GPUName), inst.GPUCount, f.VRAM, @@ -1270,7 +1294,7 @@ func displayGPUTablePlainWide(instances []GPUInstanceInfo) { ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL_VRAM", "CAPABILITY", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "DEFAULT_LOCATION", "GPU", "COUNT", "VRAM/GPU", "TOTAL_VRAM", "CAPABILITY", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1279,6 +1303,7 @@ func displayGPUTablePlainWide(instances []GPUInstanceInfo) { inst.Type, f.TargetDisk, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), inst.GPUName, inst.GPUCount, f.VRAM, @@ -1305,7 +1330,7 @@ func displayCPUTable(instances []GPUInstanceInfo) { ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "PROVIDER", "VCPUs", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "DEFAULT LOCATION", "VCPUs", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1313,6 +1338,7 @@ func displayCPUTable(instances []GPUInstanceInfo) { row := table.Row{ inst.Type, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), inst.VCPUs, f.RAM, inst.Arch, @@ -1334,7 +1360,7 @@ func displayCPUTablePlain(instances []GPUInstanceInfo) { ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() - header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "VCPUs", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "$/HR"} + header := table.Row{"TYPE", "TARGET_DISK", "PROVIDER", "DEFAULT_LOCATION", "VCPUs", "RAM", "ARCH", "DISK", "$/GB/MO", "BOOT", "FEATURES", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -1343,6 +1369,7 @@ func displayCPUTablePlain(instances []GPUInstanceInfo) { inst.Type, f.TargetDisk, f.Provider, + formatDefaultLocation(inst.Location, inst.SubLocation), inst.VCPUs, f.RAM, inst.Arch, diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index aa4db977..2248f2a6 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -1,6 +1,10 @@ package gpusearch import ( + "encoding/json" + "io" + "os" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -20,11 +24,46 @@ func (m *MockGPUSearchStore) GetInstanceTypes(_ bool) (*InstanceTypesResponse, e return m.Response, nil } +func captureSearchStdout(t *testing.T, fn func()) string { + t.Helper() + + oldStdout := os.Stdout + readPipe, writePipe, err := os.Pipe() + require.NoError(t, err) + os.Stdout = writePipe + defer func() { + os.Stdout = oldStdout + _ = readPipe.Close() + _ = writePipe.Close() + }() + + fn() + require.NoError(t, writePipe.Close()) + os.Stdout = oldStdout + output, err := io.ReadAll(readPipe) + require.NoError(t, err) + return string(output) +} + +func tableHeaderFields(t *testing.T, output string) []string { + t.Helper() + for _, line := range strings.Split(output, "\n") { + fields := strings.Fields(line) + if len(fields) > 0 && fields[0] == "TYPE" { + return fields + } + } + t.Fatalf("table header not found in output:\n%s", output) + return nil +} + func createTestInstanceTypes() *InstanceTypesResponse { return &InstanceTypesResponse{ Items: []InstanceType{ { Type: "g5.xlarge", + Location: "us-east-1", + SubLocation: "us-east-1a", AvailableLocations: []string{"us-east-1", "us-west-2"}, SupportedGPUs: []GPU{ {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, @@ -35,6 +74,7 @@ func createTestInstanceTypes() *InstanceTypesResponse { }, { Type: "g5.2xlarge", + Location: "eu-west-1", AvailableLocations: []string{"eu-west-1"}, SupportedGPUs: []GPU{ {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, @@ -163,10 +203,82 @@ func TestProcessInstances(t *testing.T) { assert.Equal(t, 24.0, a10gInstance.TotalVRAM) assert.Equal(t, 8.6, a10gInstance.Capability) assert.Equal(t, 4, a10gInstance.VCPUs) + assert.Equal(t, "us-east-1", a10gInstance.Location) + assert.Equal(t, "us-east-1a", a10gInstance.SubLocation) assert.Equal(t, []string{"us-east-1", "us-west-2"}, a10gInstance.AvailableRegions) assert.InDelta(t, 1.006, a10gInstance.PricePerHour, 0.001) } +func TestFormatDefaultLocation(t *testing.T) { + assert.Equal(t, "us-west2/us-west2-a", formatDefaultLocation("us-west2", "us-west2-a")) + assert.Equal(t, "us-west2", formatDefaultLocation("us-west2", "")) + assert.Equal(t, "us-west2-a", formatDefaultLocation("", "us-west2-a")) + assert.Equal(t, "-", formatDefaultLocation("", "")) +} + +func TestProcessCPUIncludesDefaultLocationInJSON(t *testing.T) { + instances := ProcessInstances([]InstanceType{ + { + Type: "n2d-standard-2", + Location: "asia-south1", + SubLocation: "asia-south1-a", + AvailableLocations: []string{"asia-south1", "us-west2"}, + Memory: "8GiB", + VCPU: 2, + }, + }) + + require.Len(t, instances, 1) + assert.Equal(t, "n2d-standard-2", instances[0].Type) + assert.Equal(t, "asia-south1", instances[0].Location) + assert.Equal(t, "asia-south1-a", instances[0].SubLocation) + assert.Equal(t, []string{"asia-south1", "us-west2"}, instances[0].AvailableRegions) + + payload, err := json.Marshal(instances) + require.NoError(t, err) + assert.Contains(t, string(payload), `"location":"asia-south1"`) + assert.Contains(t, string(payload), `"sub_location":"asia-south1-a"`) + assert.Contains(t, string(payload), `"available_regions":["asia-south1","us-west2"]`) +} + +func TestPlainSearchTablesPreservePipeColumnsAndShowDefaultLocation(t *testing.T) { + gpu := GPUInstanceInfo{ + Type: "g5.xlarge", + Provider: "aws", + GPUName: "A10G", + GPUCount: 1, + Location: "us-east-1", + SubLocation: "us-east-1a", + } + cpu := GPUInstanceInfo{ + Type: "n2d-standard-2", + Provider: "gcp", + Location: "asia-south1", + SubLocation: "asia-south1-a", + Manufacturer: "cpu", + } + + tests := []struct { + name string + location string + render func() + }{ + {name: "gpu", location: "us-east-1/us-east-1a", render: func() { displayGPUTablePlain([]GPUInstanceInfo{gpu}) }}, + {name: "gpu wide", location: "us-east-1/us-east-1a", render: func() { displayGPUTablePlainWide([]GPUInstanceInfo{gpu}) }}, + {name: "cpu", location: "asia-south1/asia-south1-a", render: func() { displayCPUTablePlain([]GPUInstanceInfo{cpu}) }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + output := captureSearchStdout(t, tt.render) + header := tableHeaderFields(t, output) + require.GreaterOrEqual(t, len(header), 4) + assert.Equal(t, []string{"TYPE", "TARGET_DISK", "PROVIDER", "DEFAULT_LOCATION"}, header[:4]) + assert.Contains(t, output, tt.location, "default location and sub-location should render as one pipe-safe field") + }) + } +} + func TestFilterInstancesByRegion(t *testing.T) { instances := ProcessInstances(createTestInstanceTypes().Items)