From e0ed0b11021b7537528237945b05998216dc97ea Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 11 Jan 2026 05:39:00 +0000 Subject: [PATCH 01/13] Add gpu-search command to search and filter GPU instance types Introduces a new `brev gpu-search` command (also aliased as `brev gpu`, `brev gpus`, and `brev gpu-list`) that allows users to search and filter GPU instance types from the Brev API. Features include: - Filter by GPU name (case-insensitive partial match) - Filter by minimum VRAM per GPU (in GB) - Filter by minimum total VRAM (GPU count * VRAM) - Filter by minimum GPU compute capability (e.g., 8.0 for Ampere) - Sort by price, gpu-count, vram, total-vram, vcpu, type, or capability - Support for ascending/descending sort order The command displays results in a formatted table showing instance type, GPU name, count, VRAM per GPU, total VRAM, compute capability, vCPUs, and hourly price. Includes comprehensive unit tests for filtering, sorting, and data processing functionality. --- pkg/cmd/cmd.go | 2 + pkg/cmd/gpusearch/gpusearch.go | 429 ++++++++++++++++++++++++++++ pkg/cmd/gpusearch/gpusearch_test.go | 388 +++++++++++++++++++++++++ pkg/store/instancetypes.go | 48 ++++ 4 files changed, 867 insertions(+) create mode 100644 pkg/cmd/gpusearch/gpusearch.go create mode 100644 pkg/cmd/gpusearch/gpusearch_test.go create mode 100644 pkg/store/instancetypes.go diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 2980b0ce..04564f9c 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -14,6 +14,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/delete" "github.com/brevdev/brev-cli/pkg/cmd/envvars" "github.com/brevdev/brev-cli/pkg/cmd/fu" + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/healthcheck" "github.com/brevdev/brev-cli/pkg/cmd/hello" "github.com/brevdev/brev-cli/pkg/cmd/importideconfig" @@ -270,6 +271,7 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor } cmd.AddCommand(workspacegroups.NewCmdWorkspaceGroups(t, loginCmdStore)) cmd.AddCommand(scale.NewCmdScale(t, noLoginCmdStore)) + cmd.AddCommand(gpusearch.NewCmdGPUSearch(t, noLoginCmdStore)) cmd.AddCommand(configureenvvars.NewCmdConfigureEnvVars(t, loginCmdStore)) cmd.AddCommand(importideconfig.NewCmdImportIDEConfig(t, noLoginCmdStore)) cmd.AddCommand(shell.NewCmdShell(t, loginCmdStore, noLoginCmdStore)) diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go new file mode 100644 index 00000000..d66a5545 --- /dev/null +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -0,0 +1,429 @@ +// Package gpusearch provides a command to search and filter GPU instance types +package gpusearch + +import ( + "fmt" + "os" + "regexp" + "sort" + "strconv" + "strings" + + breverrors "github.com/brevdev/brev-cli/pkg/errors" + "github.com/brevdev/brev-cli/pkg/terminal" + "github.com/jedib0t/go-pretty/v6/table" + "github.com/spf13/cobra" +) + +// MemoryBytes represents the memory size with value and unit +type MemoryBytes struct { + Value int64 `json:"value"` + Unit string `json:"unit"` +} + +// GPU represents a GPU configuration within an instance type +type GPU struct { + Count int `json:"count"` + Name string `json:"name"` + Manufacturer string `json:"manufacturer"` + Memory string `json:"memory"` + MemoryBytes MemoryBytes `json:"memory_bytes"` +} + +// BasePrice represents the pricing information +type BasePrice struct { + Currency string `json:"currency"` + Amount string `json:"amount"` +} + +// InstanceType represents an instance type from the API +type InstanceType struct { + Type string `json:"type"` + SupportedGPUs []GPU `json:"supported_gpus"` + SupportedStorage []interface{} `json:"supported_storage"` // Complex objects, not used in filtering + Memory string `json:"memory"` + VCPU int `json:"vcpu"` + BasePrice BasePrice `json:"base_price"` + Location string `json:"location"` + SubLocation string `json:"sub_location"` + AvailableLocations []string `json:"available_locations"` +} + +// InstanceTypesResponse represents the API response +type InstanceTypesResponse struct { + Items []InstanceType `json:"items"` +} + +// GPUSearchStore defines the interface for fetching instance types +type GPUSearchStore interface { + GetInstanceTypes() (*InstanceTypesResponse, error) +} + +var ( + long = `Search and filter GPU instance types available on Brev. + +Filter instances by GPU name, VRAM, total VRAM, and GPU compute capability. +Sort results by various columns to find the best instance for your needs.` + + example = ` + # List all GPU instances + brev gpu-search + + # Filter by GPU name (case-insensitive, partial match) + brev gpu-search --gpu-name A100 + brev gpu-search --gpu-name "L40S" + + # Filter by minimum VRAM per GPU (in GB) + brev gpu-search --min-vram 24 + + # Filter by minimum total VRAM (in GB) + brev gpu-search --min-total-vram 80 + + # Filter by minimum GPU compute capability + brev gpu-search --min-capability 8.0 + + # Sort by different columns (price, gpu-count, vram, total-vram, vcpu) + brev gpu-search --sort price + brev gpu-search --sort total-vram --desc + + # Combine filters + brev gpu-search --gpu-name A100 --min-vram 40 --sort price +` +) + +// NewCmdGPUSearch creates the gpu-search command +func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command { + var gpuName string + var minVRAM float64 + var minTotalVRAM float64 + var minCapability float64 + var sortBy string + var descending bool + + cmd := &cobra.Command{ + Annotations: map[string]string{"workspace": ""}, + Use: "gpu-search", + Aliases: []string{"gpu", "gpus", "gpu-list"}, + DisableFlagsInUseLine: true, + Short: "Search and filter GPU instance types", + Long: long, + Example: example, + RunE: func(cmd *cobra.Command, args []string) error { + err := RunGPUSearch(t, store, gpuName, minVRAM, minTotalVRAM, minCapability, sortBy, descending) + if err != nil { + return breverrors.WrapAndTrace(err) + } + return nil + }, + } + + cmd.Flags().StringVarP(&gpuName, "gpu-name", "g", "", "Filter by GPU name (case-insensitive, partial match)") + cmd.Flags().Float64VarP(&minVRAM, "min-vram", "v", 0, "Minimum VRAM per GPU in GB") + cmd.Flags().Float64VarP(&minTotalVRAM, "min-total-vram", "t", 0, "Minimum total VRAM (GPU count * VRAM) in GB") + cmd.Flags().Float64VarP(&minCapability, "min-capability", "c", 0, "Minimum GPU compute capability (e.g., 8.0 for Ampere)") + cmd.Flags().StringVarP(&sortBy, "sort", "s", "price", "Sort by: price, gpu-count, vram, total-vram, vcpu, type") + cmd.Flags().BoolVarP(&descending, "desc", "d", false, "Sort in descending order") + + return cmd +} + +// GPUInstanceInfo holds processed GPU instance information for display +type GPUInstanceInfo struct { + Type string + GPUName string + GPUCount int + VRAMPerGPU float64 // in GB + TotalVRAM float64 // in GB + Capability float64 + VCPUs int + Memory string + PricePerHour float64 + Manufacturer string +} + +// RunGPUSearch executes the GPU search with filters and sorting +func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName string, minVRAM, minTotalVRAM, minCapability float64, sortBy string, descending bool) error { + response, err := store.GetInstanceTypes() + if err != nil { + return breverrors.WrapAndTrace(err) + } + + if response == nil || len(response.Items) == 0 { + t.Vprint(t.Yellow("No instance types found")) + return nil + } + + // Process and filter instances + instances := processInstances(response.Items) + + // Apply filters + filtered := filterInstances(instances, gpuName, minVRAM, minTotalVRAM, minCapability) + + if len(filtered) == 0 { + t.Vprint(t.Yellow("No GPU instances match the specified filters")) + return nil + } + + // Sort instances + sortInstances(filtered, sortBy, descending) + + // Display results + displayGPUTable(t, filtered) + + t.Vprintf("\n%s\n", t.Green(fmt.Sprintf("Found %d GPU instance types", len(filtered)))) + + return nil +} + +// parseMemoryToGB converts memory string like "22GiB360MiB" or "40GiB" to GB +func parseMemoryToGB(memory string) float64 { + // Handle memory_bytes if provided (in MiB) + // Otherwise parse the string format + + var totalGB float64 + + // Match GiB values + gibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GiB`) + gibMatches := gibRe.FindAllStringSubmatch(memory, -1) + for _, match := range gibMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val + } + } + + // Match MiB values and convert to GB + mibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*MiB`) + mibMatches := mibRe.FindAllStringSubmatch(memory, -1) + for _, match := range mibMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val / 1024 + } + } + + // Match GB values (in case API uses GB instead of GiB) + gbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GB`) + gbMatches := gbRe.FindAllStringSubmatch(memory, -1) + for _, match := range gbMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val + } + } + + return totalGB +} + +// gpuCapabilityEntry represents a GPU pattern and its compute capability +type gpuCapabilityEntry struct { + pattern string + capability float64 +} + +// getGPUCapability returns the compute capability for known GPU types +func getGPUCapability(gpuName string) float64 { + gpuName = strings.ToUpper(gpuName) + + // Order matters: more specific patterns must come before less specific ones + // (e.g., "A100" before "A10", "L40S" before "L40") + capabilities := []gpuCapabilityEntry{ + // NVIDIA Hopper + {"H100", 9.0}, + {"H200", 9.0}, + + // NVIDIA Ada Lovelace (L40S before L40, L4) + {"L40S", 8.9}, + {"L40", 8.9}, + {"L4", 8.9}, + {"RTX4090", 8.9}, + {"RTX4080", 8.9}, + + // NVIDIA Ampere (A100 before A10G, A10) + {"A100", 8.0}, + {"A10G", 8.6}, + {"A10", 8.6}, + {"A40", 8.6}, + {"A30", 8.0}, + {"A16", 8.6}, + {"RTX3090", 8.6}, + {"RTX3080", 8.6}, + + // NVIDIA Turing + {"T4", 7.5}, + {"RTX2080", 7.5}, + + // NVIDIA Volta + {"V100", 7.0}, + + // NVIDIA Pascal (P100 before P40, P4) + {"P100", 6.0}, + {"P40", 6.1}, + {"P4", 6.1}, + + // NVIDIA Kepler + {"K80", 3.7}, + + // Gaudi (Habana) - not CUDA compatible + {"HL-205", 0}, + {"GAUDI3", 0}, + {"GAUDI2", 0}, + {"GAUDI", 0}, + } + + for _, entry := range capabilities { + if strings.Contains(gpuName, entry.pattern) { + return entry.capability + } + } + return 0 +} + +// processInstances converts raw instance types to GPUInstanceInfo +func processInstances(items []InstanceType) []GPUInstanceInfo { + var instances []GPUInstanceInfo + + for _, item := range items { + if len(item.SupportedGPUs) == 0 { + continue // Skip non-GPU instances + } + + for _, gpu := range item.SupportedGPUs { + vramPerGPU := parseMemoryToGB(gpu.Memory) + // Also check memory_bytes as fallback + if vramPerGPU == 0 && gpu.MemoryBytes.Value > 0 { + // Convert based on unit + if gpu.MemoryBytes.Unit == "MiB" { + vramPerGPU = float64(gpu.MemoryBytes.Value) / 1024 // MiB to GiB + } else if gpu.MemoryBytes.Unit == "GiB" { + vramPerGPU = float64(gpu.MemoryBytes.Value) + } + } + + totalVRAM := vramPerGPU * float64(gpu.Count) + capability := getGPUCapability(gpu.Name) + + price := 0.0 + if item.BasePrice.Amount != "" { + price, _ = strconv.ParseFloat(item.BasePrice.Amount, 64) + } + + instances = append(instances, GPUInstanceInfo{ + Type: item.Type, + GPUName: gpu.Name, + GPUCount: gpu.Count, + VRAMPerGPU: vramPerGPU, + TotalVRAM: totalVRAM, + Capability: capability, + VCPUs: item.VCPU, + Memory: item.Memory, + PricePerHour: price, + Manufacturer: gpu.Manufacturer, + }) + } + } + + return instances +} + +// filterInstances applies all filters to the instance list +func filterInstances(instances []GPUInstanceInfo, gpuName string, minVRAM, minTotalVRAM, minCapability float64) []GPUInstanceInfo { + var filtered []GPUInstanceInfo + + for _, inst := range instances { + // Filter by GPU name (case-insensitive partial match) + if gpuName != "" && !strings.Contains(strings.ToLower(inst.GPUName), strings.ToLower(gpuName)) { + continue + } + + // Filter by minimum VRAM per GPU + if minVRAM > 0 && inst.VRAMPerGPU < minVRAM { + continue + } + + // Filter by minimum total VRAM + if minTotalVRAM > 0 && inst.TotalVRAM < minTotalVRAM { + continue + } + + // Filter by minimum GPU capability + if minCapability > 0 && inst.Capability < minCapability { + continue + } + + filtered = append(filtered, inst) + } + + return filtered +} + +// sortInstances sorts the instance list by the specified column +func sortInstances(instances []GPUInstanceInfo, sortBy string, descending bool) { + sort.Slice(instances, func(i, j int) bool { + var less bool + switch strings.ToLower(sortBy) { + case "price": + less = instances[i].PricePerHour < instances[j].PricePerHour + case "gpu-count": + less = instances[i].GPUCount < instances[j].GPUCount + case "vram": + less = instances[i].VRAMPerGPU < instances[j].VRAMPerGPU + case "total-vram": + less = instances[i].TotalVRAM < instances[j].TotalVRAM + case "vcpu": + less = instances[i].VCPUs < instances[j].VCPUs + case "type": + less = instances[i].Type < instances[j].Type + case "capability": + less = instances[i].Capability < instances[j].Capability + default: + less = instances[i].PricePerHour < instances[j].PricePerHour + } + + if descending { + return !less + } + return less + }) +} + +// getBrevTableOptions returns table styling options +func getBrevTableOptions() table.Options { + options := table.OptionsDefault + options.DrawBorder = false + options.SeparateColumns = false + options.SeparateRows = false + options.SeparateHeader = false + return options +} + +// 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", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "VCPUs", "$/HR"} + ta.AppendHeader(header) + + for _, inst := range instances { + vramStr := fmt.Sprintf("%.0f GB", inst.VRAMPerGPU) + totalVramStr := fmt.Sprintf("%.0f GB", inst.TotalVRAM) + capStr := "-" + if inst.Capability > 0 { + capStr = fmt.Sprintf("%.1f", inst.Capability) + } + priceStr := fmt.Sprintf("$%.2f", inst.PricePerHour) + + row := table.Row{ + inst.Type, + t.Green(inst.GPUName), + inst.GPUCount, + vramStr, + totalVramStr, + capStr, + inst.VCPUs, + priceStr, + } + ta.AppendRow(row) + } + + ta.Render() +} diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go new file mode 100644 index 00000000..0714874f --- /dev/null +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -0,0 +1,388 @@ +package gpusearch + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// MockGPUSearchStore is a mock implementation of GPUSearchStore for testing +type MockGPUSearchStore struct { + Response *InstanceTypesResponse + Err error +} + +func (m *MockGPUSearchStore) GetInstanceTypes() (*InstanceTypesResponse, error) { + if m.Err != nil { + return nil, m.Err + } + return m.Response, nil +} + +func createTestInstanceTypes() *InstanceTypesResponse { + return &InstanceTypesResponse{ + Items: []InstanceType{ + { + Type: "g5.xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "1.006"}, + }, + { + Type: "g5.2xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, + }, + Memory: "32GiB", + VCPU: 8, + BasePrice: BasePrice{Currency: "USD", Amount: "1.212"}, + }, + { + Type: "p3.2xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "V100", Manufacturer: "NVIDIA", Memory: "16GiB"}, + }, + Memory: "61GiB", + VCPU: 8, + BasePrice: BasePrice{Currency: "USD", Amount: "3.06"}, + }, + { + Type: "p3.8xlarge", + SupportedGPUs: []GPU{ + {Count: 4, Name: "V100", Manufacturer: "NVIDIA", Memory: "16GiB"}, + }, + Memory: "244GiB", + VCPU: 32, + BasePrice: BasePrice{Currency: "USD", Amount: "12.24"}, + }, + { + Type: "p4d.24xlarge", + SupportedGPUs: []GPU{ + {Count: 8, Name: "A100", Manufacturer: "NVIDIA", Memory: "40GiB"}, + }, + Memory: "1152GiB", + VCPU: 96, + BasePrice: BasePrice{Currency: "USD", Amount: "32.77"}, + }, + { + Type: "g4dn.xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "T4", Manufacturer: "NVIDIA", Memory: "16GiB"}, + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "0.526"}, + }, + { + Type: "g6.xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "L4", Manufacturer: "NVIDIA", Memory: "24GiB"}, + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "0.805"}, + }, + }, + } +} + +func TestParseMemoryToGB(t *testing.T) { + tests := []struct { + name string + input string + expected float64 + }{ + {"Simple GiB", "24GiB", 24}, + {"GiB with MiB", "22GiB360MiB", 22.3515625}, + {"Simple GB", "16GB", 16}, + {"Large GiB", "1152GiB", 1152}, + {"Empty string", "", 0}, + {"MiB only", "512MiB", 0.5}, + {"With spaces", "24 GiB", 24}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := parseMemoryToGB(tt.input) + assert.InDelta(t, tt.expected, result, 0.01, "Memory parsing failed for %s", tt.input) + }) + } +} + +func TestGetGPUCapability(t *testing.T) { + tests := []struct { + name string + gpuName string + expected float64 + }{ + {"A100", "A100", 8.0}, + {"A10G", "A10G", 8.6}, + {"V100", "V100", 7.0}, + {"T4", "T4", 7.5}, + {"L4", "L4", 8.9}, + {"L40S", "L40S", 8.9}, + {"H100", "H100", 9.0}, + {"Unknown GPU", "Unknown", 0}, + {"Case insensitive", "a100", 8.0}, + {"Gaudi", "HL-205", 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := getGPUCapability(tt.gpuName) + assert.Equal(t, tt.expected, result, "GPU capability mismatch for %s", tt.gpuName) + }) + } +} + +func TestProcessInstances(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + assert.Len(t, instances, 7, "Expected 7 GPU instances") + + // Check specific instance + var a10gInstance *GPUInstanceInfo + for i := range instances { + if instances[i].Type == "g5.xlarge" { + a10gInstance = &instances[i] + break + } + } + + assert.NotNil(t, a10gInstance, "g5.xlarge instance should exist") + assert.Equal(t, "A10G", a10gInstance.GPUName) + assert.Equal(t, 1, a10gInstance.GPUCount) + assert.Equal(t, 24.0, a10gInstance.VRAMPerGPU) + assert.Equal(t, 24.0, a10gInstance.TotalVRAM) + assert.Equal(t, 8.6, a10gInstance.Capability) + assert.Equal(t, 4, a10gInstance.VCPUs) + assert.InDelta(t, 1.006, a10gInstance.PricePerHour, 0.001) +} + +func TestFilterInstancesByGPUName(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Filter by A10G + filtered := filterInstances(instances, "A10G", 0, 0, 0) + assert.Len(t, filtered, 2, "Should have 2 A10G instances") + + // Filter by V100 + filtered = filterInstances(instances, "V100", 0, 0, 0) + assert.Len(t, filtered, 2, "Should have 2 V100 instances") + + // Filter by lowercase (case-insensitive) + filtered = filterInstances(instances, "v100", 0, 0, 0) + assert.Len(t, filtered, 2, "Should have 2 V100 instances (case-insensitive)") + + // Filter by partial match + filtered = filterInstances(instances, "A1", 0, 0, 0) + assert.Len(t, filtered, 3, "Should have 3 instances matching 'A1' (A10G and A100)") +} + +func TestFilterInstancesByMinVRAM(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Filter by min VRAM 24GB + filtered := filterInstances(instances, "", 24, 0, 0) + assert.Len(t, filtered, 4, "Should have 4 instances with >= 24GB VRAM") + + // Filter by min VRAM 40GB + filtered = filterInstances(instances, "", 40, 0, 0) + assert.Len(t, filtered, 1, "Should have 1 instance with >= 40GB VRAM") + assert.Equal(t, "A100", filtered[0].GPUName) +} + +func TestFilterInstancesByMinTotalVRAM(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Filter by min total VRAM 60GB + filtered := filterInstances(instances, "", 0, 60, 0) + assert.Len(t, filtered, 2, "Should have 2 instances with >= 60GB total VRAM") + + // Filter by min total VRAM 300GB + filtered = filterInstances(instances, "", 0, 300, 0) + assert.Len(t, filtered, 1, "Should have 1 instance with >= 300GB total VRAM") + assert.Equal(t, "p4d.24xlarge", filtered[0].Type) +} + +func TestFilterInstancesByMinCapability(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Filter by capability >= 8.0 + filtered := filterInstances(instances, "", 0, 0, 8.0) + assert.Len(t, filtered, 4, "Should have 4 instances with capability >= 8.0") + + // Filter by capability >= 8.5 + filtered = filterInstances(instances, "", 0, 0, 8.5) + assert.Len(t, filtered, 3, "Should have 3 instances with capability >= 8.5") +} + +func TestFilterInstancesCombined(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Filter by GPU name and min VRAM + filtered := filterInstances(instances, "A10G", 24, 0, 0) + assert.Len(t, filtered, 2, "Should have 2 A10G instances with >= 24GB VRAM") + + // Filter by GPU name, min VRAM, and capability + filtered = filterInstances(instances, "", 24, 0, 8.5) + assert.Len(t, filtered, 3, "Should have 3 instances with >= 24GB VRAM and capability >= 8.5") +} + +func TestSortInstancesByPrice(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by price ascending + sortInstances(instances, "price", false) + assert.Equal(t, "g4dn.xlarge", instances[0].Type, "Cheapest should be g4dn.xlarge") + assert.Equal(t, "p4d.24xlarge", instances[len(instances)-1].Type, "Most expensive should be p4d.24xlarge") + + // Sort by price descending + sortInstances(instances, "price", true) + assert.Equal(t, "p4d.24xlarge", instances[0].Type, "Most expensive should be first when descending") +} + +func TestSortInstancesByGPUCount(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by GPU count ascending + sortInstances(instances, "gpu-count", false) + assert.Equal(t, 1, instances[0].GPUCount, "Instances with 1 GPU should be first") + + // Sort by GPU count descending + sortInstances(instances, "gpu-count", true) + assert.Equal(t, 8, instances[0].GPUCount, "Instance with 8 GPUs should be first when descending") +} + +func TestSortInstancesByVRAM(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by VRAM ascending + sortInstances(instances, "vram", false) + assert.Equal(t, 16.0, instances[0].VRAMPerGPU, "Instances with 16GB VRAM should be first") + + // Sort by VRAM descending + sortInstances(instances, "vram", true) + assert.Equal(t, 40.0, instances[0].VRAMPerGPU, "Instance with 40GB VRAM should be first when descending") +} + +func TestSortInstancesByTotalVRAM(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by total VRAM ascending + sortInstances(instances, "total-vram", false) + assert.Equal(t, 16.0, instances[0].TotalVRAM, "Instances with 16GB total VRAM should be first") + + // Sort by total VRAM descending + sortInstances(instances, "total-vram", true) + assert.Equal(t, 320.0, instances[0].TotalVRAM, "Instance with 320GB total VRAM should be first when descending") +} + +func TestSortInstancesByVCPU(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by vCPU ascending + sortInstances(instances, "vcpu", false) + assert.Equal(t, 4, instances[0].VCPUs, "Instances with 4 vCPUs should be first") + + // Sort by vCPU descending + sortInstances(instances, "vcpu", true) + assert.Equal(t, 96, instances[0].VCPUs, "Instance with 96 vCPUs should be first when descending") +} + +func TestSortInstancesByCapability(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by capability ascending + sortInstances(instances, "capability", false) + assert.Equal(t, 7.0, instances[0].Capability, "Instances with capability 7.0 should be first") + + // Sort by capability descending + sortInstances(instances, "capability", true) + assert.Equal(t, 8.9, instances[0].Capability, "Instance with capability 8.9 should be first when descending") +} + +func TestSortInstancesByType(t *testing.T) { + response := createTestInstanceTypes() + instances := processInstances(response.Items) + + // Sort by type ascending + sortInstances(instances, "type", false) + assert.Equal(t, "g4dn.xlarge", instances[0].Type, "g4dn.xlarge should be first alphabetically") + + // Sort by type descending + sortInstances(instances, "type", true) + assert.Equal(t, "p4d.24xlarge", instances[0].Type, "p4d.24xlarge should be first when descending") +} + +func TestEmptyInstanceTypes(t *testing.T) { + response := &InstanceTypesResponse{Items: []InstanceType{}} + instances := processInstances(response.Items) + + assert.Len(t, instances, 0, "Should have 0 instances") + + filtered := filterInstances(instances, "A100", 0, 0, 0) + assert.Len(t, filtered, 0, "Filtered should also be empty") +} + +func TestNonGPUInstancesAreFiltered(t *testing.T) { + response := &InstanceTypesResponse{ + Items: []InstanceType{ + { + Type: "m5.xlarge", + SupportedGPUs: []GPU{}, // No GPUs + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "0.192"}, + }, + { + Type: "g5.xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A10G", Manufacturer: "NVIDIA", Memory: "24GiB"}, + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "1.006"}, + }, + }, + } + + instances := processInstances(response.Items) + assert.Len(t, instances, 1, "Should only have 1 GPU instance, non-GPU instances should be filtered") + assert.Equal(t, "g5.xlarge", instances[0].Type) +} + +func TestMemoryBytesAsFallback(t *testing.T) { + response := &InstanceTypesResponse{ + Items: []InstanceType{ + { + Type: "test.xlarge", + SupportedGPUs: []GPU{ + {Count: 1, Name: "TestGPU", Manufacturer: "NVIDIA", Memory: "", MemoryBytes: MemoryBytes{Value: 24576, Unit: "MiB"}}, // 24GB in MiB + }, + Memory: "16GiB", + VCPU: 4, + BasePrice: BasePrice{Currency: "USD", Amount: "1.00"}, + }, + }, + } + + instances := processInstances(response.Items) + assert.Len(t, instances, 1) + assert.Equal(t, 24.0, instances[0].VRAMPerGPU, "Should fall back to MemoryBytes when Memory string is empty") +} diff --git a/pkg/store/instancetypes.go b/pkg/store/instancetypes.go new file mode 100644 index 00000000..4f12710a --- /dev/null +++ b/pkg/store/instancetypes.go @@ -0,0 +1,48 @@ +package store + +import ( + "encoding/json" + + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" + breverrors "github.com/brevdev/brev-cli/pkg/errors" + resty "github.com/go-resty/resty/v2" +) + +const ( + instanceTypesAPIURL = "https://api.brev.dev" + instanceTypesAPIPath = "v1/instance/types" +) + +// GetInstanceTypes fetches all available instance types from the public API +func (s NoAuthHTTPStore) GetInstanceTypes() (*gpusearch.InstanceTypesResponse, error) { + return fetchInstanceTypes() +} + +// GetInstanceTypes fetches all available instance types from the public API +func (s AuthHTTPStore) GetInstanceTypes() (*gpusearch.InstanceTypesResponse, error) { + return fetchInstanceTypes() +} + +// fetchInstanceTypes fetches instance types from the public Brev API +func fetchInstanceTypes() (*gpusearch.InstanceTypesResponse, error) { + client := resty.New() + client.SetBaseURL(instanceTypesAPIURL) + + res, err := client.R(). + SetHeader("Accept", "application/json"). + Get(instanceTypesAPIPath) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + if res.IsError() { + return nil, NewHTTPResponseError(res) + } + + var result gpusearch.InstanceTypesResponse + err = json.Unmarshal(res.Body(), &result) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + + return &result, nil +} From 78d51e3c82291fc4ecad440755ac8a23caa61ae0 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Sat, 10 Jan 2026 23:14:43 -0800 Subject: [PATCH 02/13] Add compute capabilities for additional NVIDIA GPUs Add capability mappings for: - RTXPro6000 (12.0), B200 and RTX5090 (10.0 Blackwell) - RTX6000Ada, RTX4000Ada (8.9 Ada Lovelace) - A6000, A5000, A4000 (8.6 Ampere) - RTX6000 (7.5 Turing) - M60 (5.2 Maxwell) Co-Authored-By: Claude Opus 4.5 --- pkg/cmd/gpusearch/gpusearch.go | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index d66a5545..c9a3e88f 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -225,14 +225,23 @@ func getGPUCapability(gpuName string) float64 { // Order matters: more specific patterns must come before less specific ones // (e.g., "A100" before "A10", "L40S" before "L40") capabilities := []gpuCapabilityEntry{ + // NVIDIA Professional (before other RTX patterns) + {"RTXPRO6000", 12.0}, + + // NVIDIA Blackwell + {"B200", 10.0}, + {"RTX5090", 10.0}, + // NVIDIA Hopper {"H100", 9.0}, {"H200", 9.0}, - // NVIDIA Ada Lovelace (L40S before L40, L4) + // NVIDIA Ada Lovelace (L40S before L40, L4; RTX*Ada before RTX*) {"L40S", 8.9}, {"L40", 8.9}, {"L4", 8.9}, + {"RTX6000ADA", 8.9}, + {"RTX4000ADA", 8.9}, {"RTX4090", 8.9}, {"RTX4080", 8.9}, @@ -241,6 +250,9 @@ func getGPUCapability(gpuName string) float64 { {"A10G", 8.6}, {"A10", 8.6}, {"A40", 8.6}, + {"A6000", 8.6}, + {"A5000", 8.6}, + {"A4000", 8.6}, {"A30", 8.0}, {"A16", 8.6}, {"RTX3090", 8.6}, @@ -248,6 +260,7 @@ func getGPUCapability(gpuName string) float64 { // NVIDIA Turing {"T4", 7.5}, + {"RTX6000", 7.5}, {"RTX2080", 7.5}, // NVIDIA Volta @@ -258,6 +271,9 @@ func getGPUCapability(gpuName string) float64 { {"P40", 6.1}, {"P4", 6.1}, + // NVIDIA Maxwell + {"M60", 5.2}, + // NVIDIA Kepler {"K80", 3.7}, From 212d51bb7c58378c6fcc19176ca2dee10684f4e1 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Sat, 10 Jan 2026 23:17:34 -0800 Subject: [PATCH 03/13] Filter out non-NVIDIA GPUs from gpu-search results Only show NVIDIA GPUs (exclude AMD Radeon, Intel Gaudi, etc.) since compute capability is NVIDIA-specific. Co-Authored-By: Claude Opus 4.5 --- pkg/cmd/gpusearch/gpusearch.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index c9a3e88f..54d1bd37 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -344,6 +344,11 @@ func filterInstances(instances []GPUInstanceInfo, gpuName string, minVRAM, minTo var filtered []GPUInstanceInfo for _, inst := range instances { + // Filter out non-NVIDIA GPUs (AMD, Intel/Habana, etc.) + if !strings.Contains(strings.ToUpper(inst.Manufacturer), "NVIDIA") { + continue + } + // Filter by GPU name (case-insensitive partial match) if gpuName != "" && !strings.Contains(strings.ToLower(inst.GPUName), strings.ToLower(gpuName)) { continue From 00744b206e38b2c736c0e73f2cf2b0f422305d34 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 11 Jan 2026 05:56:43 +0000 Subject: [PATCH 04/13] Add gpu-create command for retry-enabled GPU instance creation This command allows creating GPU instances with automatic retry across multiple instance types. Key features: - Accept instance types from --type flag or piped from 'brev gpus' - Create multiple instances with --count flag - Parallel creation attempts with --parallel flag - Automatic cleanup of extra instances beyond requested count - Detached mode to not wait for instances to be ready Usage examples: brev gpu-create --name my-instance --type g5.xlarge brev gpus --min-vram 24 | brev gpu-create --name my-instance brev gpus --gpu-name A100 | brev gpu-create --name cluster --count 3 --parallel 5 --- pkg/cmd/cmd.go | 2 + pkg/cmd/gpucreate/gpucreate.go | 471 ++++++++++++++++++++++++++++ pkg/cmd/gpucreate/gpucreate_test.go | 346 ++++++++++++++++++++ 3 files changed, 819 insertions(+) create mode 100644 pkg/cmd/gpucreate/gpucreate.go create mode 100644 pkg/cmd/gpucreate/gpucreate_test.go diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 04564f9c..c7253f26 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -14,6 +14,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/delete" "github.com/brevdev/brev-cli/pkg/cmd/envvars" "github.com/brevdev/brev-cli/pkg/cmd/fu" + "github.com/brevdev/brev-cli/pkg/cmd/gpucreate" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/healthcheck" "github.com/brevdev/brev-cli/pkg/cmd/hello" @@ -272,6 +273,7 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(workspacegroups.NewCmdWorkspaceGroups(t, loginCmdStore)) cmd.AddCommand(scale.NewCmdScale(t, noLoginCmdStore)) cmd.AddCommand(gpusearch.NewCmdGPUSearch(t, noLoginCmdStore)) + cmd.AddCommand(gpucreate.NewCmdGPUCreate(t, loginCmdStore)) cmd.AddCommand(configureenvvars.NewCmdConfigureEnvVars(t, loginCmdStore)) cmd.AddCommand(importideconfig.NewCmdImportIDEConfig(t, noLoginCmdStore)) cmd.AddCommand(shell.NewCmdShell(t, loginCmdStore, noLoginCmdStore)) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go new file mode 100644 index 00000000..ccd11d23 --- /dev/null +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -0,0 +1,471 @@ +// Package gpucreate provides a command to create GPU instances with retry logic +package gpucreate + +import ( + "bufio" + "context" + "fmt" + "os" + "strings" + "sync" + "time" + + "github.com/brevdev/brev-cli/pkg/cmd/util" + "github.com/brevdev/brev-cli/pkg/config" + "github.com/brevdev/brev-cli/pkg/entity" + breverrors "github.com/brevdev/brev-cli/pkg/errors" + "github.com/brevdev/brev-cli/pkg/featureflag" + "github.com/brevdev/brev-cli/pkg/store" + "github.com/brevdev/brev-cli/pkg/terminal" + "github.com/spf13/cobra" +) + +var ( + long = `Create GPU instances with automatic retry across multiple instance types. + +This command attempts to create GPU instances, trying different instance types +until the desired number of instances are successfully created. Instance types +can be specified directly or piped from 'brev gpus'. + +The command will: +1. Try to create instances using the provided instance types (in order) +2. Continue until the desired count is reached +3. Optionally try multiple instance types in parallel +4. Clean up any extra instances that were created beyond the requested count` + + example = ` + # Create a single instance with a specific GPU type + brev gpu-create --name my-instance --type g5.xlarge + + # Pipe instance types from brev gpus (tries each type until one succeeds) + brev gpus --min-vram 24 | brev gpu-create --name my-instance + + # Create 3 instances, trying types in parallel + brev gpus --gpu-name A100 | brev gpu-create --name my-cluster --count 3 --parallel 5 + + # Try multiple specific types in order + brev gpu-create --name my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge +` +) + +// GPUCreateStore defines the interface for GPU create operations +type GPUCreateStore interface { + util.GetWorkspaceByNameOrIDErrStore + GetActiveOrganizationOrDefault() (*entity.Organization, error) + GetCurrentUser() (*entity.User, error) + GetWorkspace(workspaceID string) (*entity.Workspace, error) + CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error) + DeleteWorkspace(workspaceID string) (*entity.Workspace, error) +} + +// CreateResult holds the result of a workspace creation attempt +type CreateResult struct { + Workspace *entity.Workspace + InstanceType string + Error error +} + +// NewCmdGPUCreate creates the gpu-create command +func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra.Command { + var name string + var instanceTypes string + var count int + var parallel int + var detached bool + var timeout int + + cmd := &cobra.Command{ + Annotations: map[string]string{"workspace": ""}, + Use: "gpu-create", + Aliases: []string{"gpu-retry", "gcreate"}, + DisableFlagsInUseLine: true, + Short: "Create GPU instances with automatic retry", + Long: long, + Example: example, + RunE: func(cmd *cobra.Command, args []string) error { + // Parse instance types from flag or stdin + types, err := parseInstanceTypes(instanceTypes) + if err != nil { + return breverrors.WrapAndTrace(err) + } + + if len(types) == 0 { + return breverrors.NewValidationError("no instance types provided. Use --type flag or pipe from 'brev gpus'") + } + + if name == "" { + return breverrors.NewValidationError("--name flag is required") + } + + if count < 1 { + return breverrors.NewValidationError("--count must be at least 1") + } + + if parallel < 1 { + parallel = 1 + } + + opts := GPUCreateOptions{ + Name: name, + InstanceTypes: types, + Count: count, + Parallel: parallel, + Detached: detached, + Timeout: time.Duration(timeout) * time.Second, + } + + err = RunGPUCreate(t, gpuCreateStore, opts) + if err != nil { + return breverrors.WrapAndTrace(err) + } + return nil + }, + } + + cmd.Flags().StringVarP(&name, "name", "n", "", "Base name for the instances (required)") + cmd.Flags().StringVarP(&instanceTypes, "type", "t", "", "Comma-separated list of instance types to try") + cmd.Flags().IntVarP(&count, "count", "c", 1, "Number of instances to create") + cmd.Flags().IntVarP(¶llel, "parallel", "p", 1, "Number of parallel creation attempts") + cmd.Flags().BoolVarP(&detached, "detached", "d", false, "Don't wait for instances to be ready") + cmd.Flags().IntVar(&timeout, "timeout", 300, "Timeout in seconds for each instance to become ready") + + return cmd +} + +// GPUCreateOptions holds the options for GPU instance creation +type GPUCreateOptions struct { + Name string + InstanceTypes []string + Count int + Parallel int + Detached bool + Timeout time.Duration +} + +// parseInstanceTypes parses instance types from flag value or stdin +func parseInstanceTypes(flagValue string) ([]string, error) { + var types []string + + // First check if there's a flag value + if flagValue != "" { + parts := strings.Split(flagValue, ",") + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + types = append(types, p) + } + } + } + + // Check if there's piped input from stdin + stat, _ := os.Stdin.Stat() + if (stat.Mode() & os.ModeCharDevice) == 0 { + // Data is being piped to stdin + scanner := bufio.NewScanner(os.Stdin) + lineNum := 0 + for scanner.Scan() { + line := scanner.Text() + lineNum++ + + // Skip header line (first line typically contains column names) + if lineNum == 1 && (strings.Contains(line, "TYPE") || strings.Contains(line, "GPU")) { + continue + } + + // Skip empty lines + line = strings.TrimSpace(line) + if line == "" { + continue + } + + // Skip summary lines (e.g., "Found X GPU instance types") + if strings.HasPrefix(line, "Found ") { + continue + } + + // Extract the first column (TYPE) from the table output + // The format is: TYPE GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR + fields := strings.Fields(line) + if len(fields) > 0 { + instanceType := fields[0] + // Validate it looks like an instance type (contains letters and possibly numbers/dots) + if isValidInstanceType(instanceType) { + types = append(types, instanceType) + } + } + } + + if err := scanner.Err(); err != nil { + return nil, breverrors.WrapAndTrace(err) + } + } + + return types, nil +} + +// isValidInstanceType checks if a string looks like a valid instance type +func isValidInstanceType(s string) bool { + // Instance types typically have formats like: + // g5.xlarge, p4d.24xlarge, n1-highmem-4:nvidia-tesla-t4:1 + if len(s) < 2 { + return false + } + + // Should contain alphanumeric characters + hasLetter := false + hasNumber := false + for _, c := range s { + if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') { + hasLetter = true + } + if c >= '0' && c <= '9' { + hasNumber = true + } + } + + return hasLetter && hasNumber +} + +// RunGPUCreate executes the GPU create with retry logic +func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions) error { + user, err := gpuCreateStore.GetCurrentUser() + if err != nil { + return breverrors.WrapAndTrace(err) + } + + org, err := gpuCreateStore.GetActiveOrganizationOrDefault() + if err != nil { + return breverrors.WrapAndTrace(err) + } + if org == nil { + return breverrors.NewValidationError("no organization found") + } + + t.Vprintf("Attempting to create %d instance(s) with %d parallel attempts\n", opts.Count, opts.Parallel) + t.Vprintf("Instance types to try: %s\n\n", strings.Join(opts.InstanceTypes, ", ")) + + // Track successful creations + var successfulWorkspaces []*entity.Workspace + var mu sync.Mutex + var wg sync.WaitGroup + + // Create a context for cancellation + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Channel to coordinate attempts + typesChan := make(chan string, len(opts.InstanceTypes)) + for _, it := range opts.InstanceTypes { + typesChan <- it + } + close(typesChan) + + // Results channel + resultsChan := make(chan CreateResult, len(opts.InstanceTypes)) + + // Track instance index for naming + instanceIndex := 0 + var indexMu sync.Mutex + + // Start parallel workers + workerCount := opts.Parallel + if workerCount > len(opts.InstanceTypes) { + workerCount = len(opts.InstanceTypes) + } + + for i := 0; i < workerCount; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + + for instanceType := range typesChan { + // Check if we've already created enough + mu.Lock() + if len(successfulWorkspaces) >= opts.Count { + mu.Unlock() + return + } + mu.Unlock() + + // Check context + select { + case <-ctx.Done(): + return + default: + } + + // Get unique instance name + indexMu.Lock() + currentIndex := instanceIndex + instanceIndex++ + indexMu.Unlock() + + instanceName := opts.Name + if opts.Count > 1 || currentIndex > 0 { + instanceName = fmt.Sprintf("%s-%d", opts.Name, currentIndex+1) + } + + t.Vprintf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, instanceType, instanceName) + + // Attempt to create the workspace + workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user) + + result := CreateResult{ + Workspace: workspace, + InstanceType: instanceType, + Error: err, + } + + if err != nil { + t.Vprintf("[Worker %d] %s Failed: %s\n", workerID+1, t.Yellow(instanceType), err.Error()) + } else { + t.Vprintf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, t.Green(instanceType), instanceName) + mu.Lock() + successfulWorkspaces = append(successfulWorkspaces, workspace) + if len(successfulWorkspaces) >= opts.Count { + cancel() // Signal other workers to stop + } + mu.Unlock() + } + + resultsChan <- result + } + }(i) + } + + // Wait for all workers to finish + go func() { + wg.Wait() + close(resultsChan) + }() + + // Collect results + for range resultsChan { + // Just drain the channel + } + + // Check if we created enough instances + if len(successfulWorkspaces) < opts.Count { + t.Vprintf("\n%s Only created %d/%d instances\n", t.Yellow("Warning:"), len(successfulWorkspaces), opts.Count) + + if len(successfulWorkspaces) > 0 { + t.Vprintf("Successfully created instances:\n") + for _, ws := range successfulWorkspaces { + t.Vprintf(" - %s (ID: %s)\n", ws.Name, ws.ID) + } + } + + return breverrors.NewValidationError(fmt.Sprintf("could only create %d/%d instances", len(successfulWorkspaces), opts.Count)) + } + + // If we created more than needed, clean up extras + if len(successfulWorkspaces) > opts.Count { + extras := successfulWorkspaces[opts.Count:] + t.Vprintf("\nCleaning up %d extra instance(s)...\n", len(extras)) + + for _, ws := range extras { + t.Vprintf(" Deleting %s...", ws.Name) + _, err := gpuCreateStore.DeleteWorkspace(ws.ID) + if err != nil { + t.Vprintf(" %s\n", t.Red("Failed")) + } else { + t.Vprintf(" %s\n", t.Green("Done")) + } + } + + successfulWorkspaces = successfulWorkspaces[:opts.Count] + } + + // Wait for instances to be ready (unless detached) + if !opts.Detached { + t.Vprintf("\nWaiting for instance(s) to be ready...\n") + t.Vprintf("You can safely ctrl+c to exit\n") + + for _, ws := range successfulWorkspaces { + err := pollUntilReady(t, ws.ID, gpuCreateStore, opts.Timeout) + if err != nil { + t.Vprintf(" %s: %s\n", ws.Name, t.Yellow("Timeout waiting for ready state")) + } + } + } + + // Print summary + fmt.Print("\n") + t.Vprint(t.Green(fmt.Sprintf("Successfully created %d instance(s)!\n\n", len(successfulWorkspaces)))) + + for _, ws := range successfulWorkspaces { + t.Vprintf("Instance: %s\n", t.Green(ws.Name)) + t.Vprintf(" ID: %s\n", ws.ID) + t.Vprintf(" Type: %s\n", ws.InstanceType) + displayConnectBreadCrumb(t, ws) + fmt.Print("\n") + } + + return nil +} + +// createWorkspaceWithType creates a workspace with the specified instance type +func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User) (*entity.Workspace, error) { + clusterID := config.GlobalConfig.GetDefaultClusterID() + cwOptions := store.NewCreateWorkspacesOptions(clusterID, name) + cwOptions.WithInstanceType(instanceType) + cwOptions = resolveWorkspaceUserOptions(cwOptions, user) + + workspace, err := gpuCreateStore.CreateWorkspace(orgID, cwOptions) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + + return workspace, nil +} + +// resolveWorkspaceUserOptions sets workspace template and class based on user type +func resolveWorkspaceUserOptions(options *store.CreateWorkspacesOptions, user *entity.User) *store.CreateWorkspacesOptions { + if options.WorkspaceTemplateID == "" { + if featureflag.IsAdmin(user.GlobalUserType) { + options.WorkspaceTemplateID = store.DevWorkspaceTemplateID + } else { + options.WorkspaceTemplateID = store.UserWorkspaceTemplateID + } + } + if options.WorkspaceClassID == "" { + if featureflag.IsAdmin(user.GlobalUserType) { + options.WorkspaceClassID = store.DevWorkspaceClassID + } else { + options.WorkspaceClassID = store.UserWorkspaceClassID + } + } + return options +} + +// pollUntilReady waits for a workspace to reach the running state +func pollUntilReady(t *terminal.Terminal, wsID string, gpuCreateStore GPUCreateStore, timeout time.Duration) error { + deadline := time.Now().Add(timeout) + + for time.Now().Before(deadline) { + ws, err := gpuCreateStore.GetWorkspace(wsID) + if err != nil { + return breverrors.WrapAndTrace(err) + } + + if ws.Status == entity.Running { + t.Vprintf(" %s: %s\n", ws.Name, t.Green("Ready")) + return nil + } + + if ws.Status == entity.Failure { + return breverrors.NewValidationError(fmt.Sprintf("instance %s failed", ws.Name)) + } + + time.Sleep(5 * time.Second) + } + + return breverrors.NewValidationError("timeout waiting for instance to be ready") +} + +// displayConnectBreadCrumb shows connection instructions +func displayConnectBreadCrumb(t *terminal.Terminal, workspace *entity.Workspace) { + t.Vprintf(" Connect:\n") + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("brev open %s", workspace.Name))) + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("brev shell %s", workspace.Name))) +} diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go new file mode 100644 index 00000000..8c9b935d --- /dev/null +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -0,0 +1,346 @@ +package gpucreate + +import ( + "testing" + + "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/store" + "github.com/stretchr/testify/assert" +) + +// MockGPUCreateStore is a mock implementation of GPUCreateStore for testing +type MockGPUCreateStore struct { + User *entity.User + Org *entity.Organization + Workspaces map[string]*entity.Workspace + CreateError error + CreateErrorTypes map[string]error // Errors for specific instance types + DeleteError error + CreatedWorkspaces []*entity.Workspace + DeletedWorkspaceIDs []string +} + +func NewMockGPUCreateStore() *MockGPUCreateStore { + return &MockGPUCreateStore{ + User: &entity.User{ + ID: "user-123", + GlobalUserType: "Standard", + }, + Org: &entity.Organization{ + ID: "org-123", + Name: "test-org", + }, + Workspaces: make(map[string]*entity.Workspace), + CreateErrorTypes: make(map[string]error), + CreatedWorkspaces: []*entity.Workspace{}, + DeletedWorkspaceIDs: []string{}, + } +} + +func (m *MockGPUCreateStore) GetCurrentUser() (*entity.User, error) { + return m.User, nil +} + +func (m *MockGPUCreateStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) { + return m.Org, nil +} + +func (m *MockGPUCreateStore) GetWorkspace(workspaceID string) (*entity.Workspace, error) { + if ws, ok := m.Workspaces[workspaceID]; ok { + return ws, nil + } + return &entity.Workspace{ + ID: workspaceID, + Status: entity.Running, + }, nil +} + +func (m *MockGPUCreateStore) CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error) { + // Check for type-specific errors first + if err, ok := m.CreateErrorTypes[options.InstanceType]; ok { + return nil, err + } + + if m.CreateError != nil { + return nil, m.CreateError + } + + ws := &entity.Workspace{ + ID: "ws-" + options.Name, + Name: options.Name, + InstanceType: options.InstanceType, + Status: entity.Running, + } + m.Workspaces[ws.ID] = ws + m.CreatedWorkspaces = append(m.CreatedWorkspaces, ws) + return ws, nil +} + +func (m *MockGPUCreateStore) DeleteWorkspace(workspaceID string) (*entity.Workspace, error) { + if m.DeleteError != nil { + return nil, m.DeleteError + } + + m.DeletedWorkspaceIDs = append(m.DeletedWorkspaceIDs, workspaceID) + ws := m.Workspaces[workspaceID] + delete(m.Workspaces, workspaceID) + return ws, nil +} + +func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID string) ([]entity.Workspace, error) { + return []entity.Workspace{}, nil +} + +func TestIsValidInstanceType(t *testing.T) { + tests := []struct { + name string + input string + expected bool + }{ + {"Valid AWS instance type", "g5.xlarge", true}, + {"Valid AWS large instance", "p4d.24xlarge", true}, + {"Valid GCP instance type", "n1-highmem-4:nvidia-tesla-t4:1", true}, + {"Single letter", "a", false}, + {"No numbers", "xlarge", false}, + {"No letters", "12345", false}, + {"Empty string", "", false}, + {"Single character", "1", false}, + {"Valid with colon", "g5:xlarge", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := isValidInstanceType(tt.input) + assert.Equal(t, tt.expected, result, "Validation failed for %s", tt.input) + }) + } +} + +func TestParseInstanceTypesFromFlag(t *testing.T) { + tests := []struct { + name string + input string + expected []string + }{ + {"Single type", "g5.xlarge", []string{"g5.xlarge"}}, + {"Multiple types comma separated", "g5.xlarge,g5.2xlarge,p3.2xlarge", []string{"g5.xlarge", "g5.2xlarge", "p3.2xlarge"}}, + {"With spaces", "g5.xlarge, g5.2xlarge, p3.2xlarge", []string{"g5.xlarge", "g5.2xlarge", "p3.2xlarge"}}, + {"Empty string", "", []string{}}, + {"Only spaces", " ", []string{}}, + {"Trailing comma", "g5.xlarge,", []string{"g5.xlarge"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := parseInstanceTypes(tt.input) + assert.NoError(t, err) + + // Handle nil vs empty slice + if len(tt.expected) == 0 { + assert.Empty(t, result) + } else { + assert.Equal(t, tt.expected, result) + } + }) + } +} + +func TestGPUCreateOptions(t *testing.T) { + opts := GPUCreateOptions{ + Name: "my-instance", + InstanceTypes: []string{"g5.xlarge", "g5.2xlarge"}, + Count: 2, + Parallel: 3, + Detached: true, + } + + assert.Equal(t, "my-instance", opts.Name) + assert.Len(t, opts.InstanceTypes, 2) + assert.Equal(t, 2, opts.Count) + assert.Equal(t, 3, opts.Parallel) + assert.True(t, opts.Detached) +} + +func TestResolveWorkspaceUserOptionsStandard(t *testing.T) { + user := &entity.User{ + ID: "user-123", + GlobalUserType: "Standard", + } + + options := &store.CreateWorkspacesOptions{} + result := resolveWorkspaceUserOptions(options, user) + + assert.Equal(t, store.UserWorkspaceTemplateID, result.WorkspaceTemplateID) + assert.Equal(t, store.UserWorkspaceClassID, result.WorkspaceClassID) +} + +func TestResolveWorkspaceUserOptionsAdmin(t *testing.T) { + user := &entity.User{ + ID: "user-123", + GlobalUserType: "Admin", + } + + options := &store.CreateWorkspacesOptions{} + result := resolveWorkspaceUserOptions(options, user) + + assert.Equal(t, store.DevWorkspaceTemplateID, result.WorkspaceTemplateID) + assert.Equal(t, store.DevWorkspaceClassID, result.WorkspaceClassID) +} + +func TestResolveWorkspaceUserOptionsPreserveExisting(t *testing.T) { + user := &entity.User{ + ID: "user-123", + GlobalUserType: "Standard", + } + + options := &store.CreateWorkspacesOptions{ + WorkspaceTemplateID: "custom-template", + WorkspaceClassID: "custom-class", + } + result := resolveWorkspaceUserOptions(options, user) + + // Should preserve existing values + assert.Equal(t, "custom-template", result.WorkspaceTemplateID) + assert.Equal(t, "custom-class", result.WorkspaceClassID) +} + +func TestMockGPUCreateStoreBasics(t *testing.T) { + mock := NewMockGPUCreateStore() + + user, err := mock.GetCurrentUser() + assert.NoError(t, err) + assert.Equal(t, "user-123", user.ID) + + org, err := mock.GetActiveOrganizationOrDefault() + assert.NoError(t, err) + assert.Equal(t, "org-123", org.ID) +} + +func TestMockGPUCreateStoreCreateWorkspace(t *testing.T) { + mock := NewMockGPUCreateStore() + + options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") + options.WithInstanceType("g5.xlarge") + + ws, err := mock.CreateWorkspace("org-123", options) + assert.NoError(t, err) + assert.Equal(t, "test-instance", ws.Name) + assert.Equal(t, "g5.xlarge", ws.InstanceType) + assert.Len(t, mock.CreatedWorkspaces, 1) +} + +func TestMockGPUCreateStoreDeleteWorkspace(t *testing.T) { + mock := NewMockGPUCreateStore() + + // First create a workspace + options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") + ws, _ := mock.CreateWorkspace("org-123", options) + + // Then delete it + _, err := mock.DeleteWorkspace(ws.ID) + assert.NoError(t, err) + assert.Contains(t, mock.DeletedWorkspaceIDs, ws.ID) +} + +func TestMockGPUCreateStoreTypeSpecificError(t *testing.T) { + mock := NewMockGPUCreateStore() + mock.CreateErrorTypes["g5.xlarge"] = assert.AnError + + options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") + options.WithInstanceType("g5.xlarge") + + _, err := mock.CreateWorkspace("org-123", options) + assert.Error(t, err) + + // Different type should work + options2 := store.NewCreateWorkspacesOptions("cluster-1", "test-instance-2") + options2.WithInstanceType("g5.2xlarge") + + ws, err := mock.CreateWorkspace("org-123", options2) + assert.NoError(t, err) + assert.NotNil(t, ws) +} + +func TestParseInstanceTypesFromTableOutput(t *testing.T) { + // Simulated table output from brev gpus command + // Note: This tests the parsing logic, not actual stdin reading + tableLines := []string{ + "TYPE GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR", + "g5.xlarge A10G 1 24 GB 24 GB 8.6 4 $1.01", + "g5.2xlarge A10G 1 24 GB 24 GB 8.6 8 $1.21", + "p4d.24xlarge A100 8 40 GB 320 GB 8.0 96 $32.77", + "", + "Found 3 GPU instance types", + } + + // Test parsing each line (simulating the scanner behavior) + var types []string + lineNum := 0 + for _, line := range tableLines { + lineNum++ + + // Skip header line + if lineNum == 1 && (contains(line, "TYPE") || contains(line, "GPU")) { + continue + } + + // Skip empty lines and summary + if line == "" || startsWith(line, "Found ") { + continue + } + + // Extract first column + fields := splitFields(line) + if len(fields) > 0 && isValidInstanceType(fields[0]) { + types = append(types, fields[0]) + } + } + + assert.Len(t, types, 3) + assert.Contains(t, types, "g5.xlarge") + assert.Contains(t, types, "g5.2xlarge") + assert.Contains(t, types, "p4d.24xlarge") +} + +// Helper functions for testing +func contains(s, substr string) bool { + return len(s) >= len(substr) && findSubstring(s, substr) >= 0 +} + +func findSubstring(s, substr string) int { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return i + } + } + return -1 +} + +func startsWith(s, prefix string) bool { + return len(s) >= len(prefix) && s[:len(prefix)] == prefix +} + +func splitFields(s string) []string { + var fields []string + var current string + inField := false + + for _, c := range s { + if c == ' ' || c == '\t' { + if inField { + fields = append(fields, current) + current = "" + inField = false + } + } else { + current += string(c) + inField = true + } + } + + if inField { + fields = append(fields, current) + } + + return fields +} From 1b3641333b3d70f541c188b36065d98649388820 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 11 Jan 2026 06:08:12 +0000 Subject: [PATCH 05/13] Add 'provision' as alias for gpu-create command --- pkg/cmd/gpucreate/gpucreate.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index ccd11d23..f6c5e14d 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -77,7 +77,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra cmd := &cobra.Command{ Annotations: map[string]string{"workspace": ""}, Use: "gpu-create", - Aliases: []string{"gpu-retry", "gcreate"}, + Aliases: []string{"gpu-retry", "gcreate", "provision"}, DisableFlagsInUseLine: true, Short: "Create GPU instances with automatic retry", Long: long, From 2f19c56543a595ca947a80b0949d7d1f2bff2eb9 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Tue, 27 Jan 2026 20:55:59 -0800 Subject: [PATCH 06/13] Fix gpu-create to dynamically lookup workspaceGroupId for all providers Previously gpu-create hardcoded workspaceGroupId to "GCP", causing instance creation to fail for non-GCP providers (shadeform, nebius, crusoe, lambda, etc.) with "instance type not found" errors. This change: - Adds GetAllInstanceTypesWithWorkspaceGroups API call to fetch instance types with their associated workspace groups from the authenticated endpoint (/api/instances/alltypesavailable/{orgId}) - Updates gpu-create to lookup the correct workspaceGroupId for each instance type before creating the workspace - Adds WorkspaceGroup type and AllInstanceTypesResponse for the new API - Adds provider filtering, disk size filtering, and enhanced display columns to gpu-search command Tested with: GCP, shadeform (massedcompute, hyperstack, vultr, scaleway), nebius, crusoe, lambda, and devplane providers. --- pkg/cmd/gpucreate/gpucreate.go | 22 ++- pkg/cmd/gpusearch/gpusearch.go | 213 +++++++++++++++++++++++++--- pkg/cmd/gpusearch/gpusearch_test.go | 26 ++-- pkg/store/instancetypes.go | 25 ++++ 4 files changed, 253 insertions(+), 33 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index f6c5e14d..a42b333a 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -10,6 +10,7 @@ import ( "sync" "time" + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/util" "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" @@ -56,6 +57,7 @@ type GPUCreateStore interface { GetWorkspace(workspaceID string) (*entity.Workspace, error) CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error) DeleteWorkspace(workspaceID string) (*entity.Workspace, error) + GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) } // CreateResult holds the result of a workspace creation attempt @@ -241,6 +243,14 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC return breverrors.NewValidationError("no organization found") } + // Fetch instance types with workspace groups to determine correct workspace group ID + allInstanceTypes, err := gpuCreateStore.GetAllInstanceTypesWithWorkspaceGroups(org.ID) + if err != nil { + t.Vprintf("Warning: could not fetch instance types with workspace groups: %s\n", err.Error()) + t.Vprintf("Falling back to default workspace group\n") + allInstanceTypes = nil + } + t.Vprintf("Attempting to create %d instance(s) with %d parallel attempts\n", opts.Count, opts.Parallel) t.Vprintf("Instance types to try: %s\n\n", strings.Join(opts.InstanceTypes, ", ")) @@ -308,7 +318,7 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC t.Vprintf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, instanceType, instanceName) // Attempt to create the workspace - workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user) + workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user, allInstanceTypes) result := CreateResult{ Workspace: workspace, @@ -405,12 +415,20 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC } // createWorkspaceWithType creates a workspace with the specified instance type -func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User) (*entity.Workspace, error) { +func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User, allInstanceTypes *gpusearch.AllInstanceTypesResponse) (*entity.Workspace, error) { clusterID := config.GlobalConfig.GetDefaultClusterID() cwOptions := store.NewCreateWorkspacesOptions(clusterID, name) cwOptions.WithInstanceType(instanceType) cwOptions = resolveWorkspaceUserOptions(cwOptions, user) + // Look up the workspace group ID for this instance type + if allInstanceTypes != nil { + workspaceGroupID := allInstanceTypes.GetWorkspaceGroupID(instanceType) + if workspaceGroupID != "" { + cwOptions.WorkspaceGroupID = workspaceGroupID + } + } + workspace, err := gpuCreateStore.CreateWorkspace(orgID, cwOptions) if err != nil { return nil, breverrors.WrapAndTrace(err) diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index 54d1bd37..e8fa5a3a 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -36,17 +36,36 @@ type BasePrice struct { Amount string `json:"amount"` } +// Storage represents a storage configuration within an instance type +type Storage struct { + Count int `json:"count"` + Size string `json:"size"` + Type string `json:"type"` + MinSize string `json:"min_size"` + MaxSize string `json:"max_size"` + SizeBytes MemoryBytes `json:"size_bytes"` +} + +// WorkspaceGroup represents a workspace group that can run an instance type +type WorkspaceGroup struct { + ID string `json:"id"` + Name string `json:"name"` + PlatformType string `json:"platformType"` +} + // InstanceType represents an instance type from the API type InstanceType struct { - Type string `json:"type"` - SupportedGPUs []GPU `json:"supported_gpus"` - SupportedStorage []interface{} `json:"supported_storage"` // Complex objects, not used in filtering - Memory string `json:"memory"` - VCPU int `json:"vcpu"` - BasePrice BasePrice `json:"base_price"` - Location string `json:"location"` - SubLocation string `json:"sub_location"` - AvailableLocations []string `json:"available_locations"` + Type string `json:"type"` + SupportedGPUs []GPU `json:"supported_gpus"` + SupportedStorage []Storage `json:"supported_storage"` + Memory string `json:"memory"` + VCPU int `json:"vcpu"` + BasePrice BasePrice `json:"base_price"` + Location string `json:"location"` + SubLocation string `json:"sub_location"` + AvailableLocations []string `json:"available_locations"` + Provider string `json:"provider"` + WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"` } // InstanceTypesResponse represents the API response @@ -54,6 +73,23 @@ type InstanceTypesResponse struct { Items []InstanceType `json:"items"` } +// AllInstanceTypesResponse represents the authenticated API response with workspace groups +type AllInstanceTypesResponse struct { + AllInstanceTypes []InstanceType `json:"allInstanceTypes"` +} + +// GetWorkspaceGroupID returns the workspace group ID for an instance type, or empty string if not found +func (r *AllInstanceTypesResponse) GetWorkspaceGroupID(instanceType string) string { + for _, it := range r.AllInstanceTypes { + if it.Type == instanceType { + if len(it.WorkspaceGroups) > 0 { + return it.WorkspaceGroups[0].ID + } + } + } + return "" +} + // GPUSearchStore defines the interface for fetching instance types type GPUSearchStore interface { GetInstanceTypes() (*InstanceTypesResponse, error) @@ -62,7 +98,7 @@ type GPUSearchStore interface { var ( long = `Search and filter GPU instance types available on Brev. -Filter instances by GPU name, VRAM, total VRAM, and GPU compute capability. +Filter instances by GPU name, provider, VRAM, total VRAM, GPU compute capability, and disk size. Sort results by various columns to find the best instance for your needs.` example = ` @@ -73,6 +109,10 @@ Sort results by various columns to find the best instance for your needs.` brev gpu-search --gpu-name A100 brev gpu-search --gpu-name "L40S" + # Filter by provider/cloud (case-insensitive, partial match) + brev gpu-search --provider aws + brev gpu-search --provider gcp + # Filter by minimum VRAM per GPU (in GB) brev gpu-search --min-vram 24 @@ -82,21 +122,28 @@ Sort results by various columns to find the best instance for your needs.` # Filter by minimum GPU compute capability brev gpu-search --min-capability 8.0 - # Sort by different columns (price, gpu-count, vram, total-vram, vcpu) + # Filter by minimum disk size (in GB) + brev gpu-search --min-disk 500 + + # Sort by different columns (price, gpu-count, vram, total-vram, vcpu, provider, disk) brev gpu-search --sort price - brev gpu-search --sort total-vram --desc + brev gpu-search --sort provider + brev gpu-search --sort disk --desc # Combine filters brev gpu-search --gpu-name A100 --min-vram 40 --sort price + brev gpu-search --provider aws --gpu-name A100 --sort price ` ) // NewCmdGPUSearch creates the gpu-search command func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command { var gpuName string + var provider string var minVRAM float64 var minTotalVRAM float64 var minCapability float64 + var minDisk float64 var sortBy string var descending bool @@ -109,7 +156,7 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command Long: long, Example: example, RunE: func(cmd *cobra.Command, args []string) error { - err := RunGPUSearch(t, store, gpuName, minVRAM, minTotalVRAM, minCapability, sortBy, descending) + err := RunGPUSearch(t, store, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, sortBy, descending) if err != nil { return breverrors.WrapAndTrace(err) } @@ -118,10 +165,12 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command } cmd.Flags().StringVarP(&gpuName, "gpu-name", "g", "", "Filter by GPU name (case-insensitive, partial match)") + cmd.Flags().StringVarP(&provider, "provider", "p", "", "Filter by provider/cloud (case-insensitive, partial match)") cmd.Flags().Float64VarP(&minVRAM, "min-vram", "v", 0, "Minimum VRAM per GPU in GB") cmd.Flags().Float64VarP(&minTotalVRAM, "min-total-vram", "t", 0, "Minimum total VRAM (GPU count * VRAM) in GB") cmd.Flags().Float64VarP(&minCapability, "min-capability", "c", 0, "Minimum GPU compute capability (e.g., 8.0 for Ampere)") - cmd.Flags().StringVarP(&sortBy, "sort", "s", "price", "Sort by: price, gpu-count, vram, total-vram, vcpu, type") + cmd.Flags().Float64Var(&minDisk, "min-disk", 0, "Minimum disk size in GB") + cmd.Flags().StringVarP(&sortBy, "sort", "s", "price", "Sort by: price, gpu-count, vram, total-vram, vcpu, type, provider, disk") cmd.Flags().BoolVarP(&descending, "desc", "d", false, "Sort in descending order") return cmd @@ -130,6 +179,7 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command // GPUInstanceInfo holds processed GPU instance information for display type GPUInstanceInfo struct { Type string + Provider string GPUName string GPUCount int VRAMPerGPU float64 // in GB @@ -137,12 +187,14 @@ type GPUInstanceInfo struct { Capability float64 VCPUs int Memory string + DiskMin float64 // in GB (min disk size, same as DiskMax if fixed) + DiskMax float64 // in GB (max disk size, same as DiskMin if fixed) PricePerHour float64 Manufacturer string } // RunGPUSearch executes the GPU search with filters and sorting -func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName string, minVRAM, minTotalVRAM, minCapability float64, sortBy string, descending bool) error { +func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, sortBy string, descending bool) error { response, err := store.GetInstanceTypes() if err != nil { return breverrors.WrapAndTrace(err) @@ -157,7 +209,7 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName string, mi instances := processInstances(response.Items) // Apply filters - filtered := filterInstances(instances, gpuName, minVRAM, minTotalVRAM, minCapability) + filtered := filterInstances(instances, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk) if len(filtered) == 0 { t.Vprint(t.Yellow("No GPU instances match the specified filters")) @@ -212,6 +264,86 @@ func parseMemoryToGB(memory string) float64 { return totalGB } +// parseSizeToGB parses size strings like "16TiB", "10GiB", "2TiB768GiB" to GB +func parseSizeToGB(size string) float64 { + var totalGB float64 + + // Match TiB values and convert to GB (1 TiB = 1024 GiB) + tibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*TiB`) + tibMatches := tibRe.FindAllStringSubmatch(size, -1) + for _, match := range tibMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val * 1024 + } + } + + // Match GiB values + gibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GiB`) + gibMatches := gibRe.FindAllStringSubmatch(size, -1) + for _, match := range gibMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val + } + } + + // Match TB values (1 TB = 1000 GB) + tbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*TB`) + tbMatches := tbRe.FindAllStringSubmatch(size, -1) + for _, match := range tbMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val * 1000 + } + } + + // Match GB values + gbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GB`) + gbMatches := gbRe.FindAllStringSubmatch(size, -1) + for _, match := range gbMatches { + if val, err := strconv.ParseFloat(match[1], 64); err == nil { + totalGB += val + } + } + + return totalGB +} + +// extractDiskSize extracts min and max disk size from storage configuration +// Returns (minGB, maxGB). For fixed-size storage, both values are the same. +func extractDiskSize(storage []Storage) (float64, float64) { + if len(storage) == 0 { + return 0, 0 + } + + // Use the first storage entry + s := storage[0] + + // Check if it's flexible storage (has min_size and max_size) + if s.MinSize != "" && s.MaxSize != "" { + minGB := parseSizeToGB(s.MinSize) + maxGB := parseSizeToGB(s.MaxSize) + return minGB, maxGB + } + + // Fixed storage - use size or size_bytes + var sizeGB float64 + if s.Size != "" && s.Size != "0B" { + sizeGB = parseSizeToGB(s.Size) + } + + // Fallback to size_bytes + if sizeGB == 0 && s.SizeBytes.Value > 0 { + if s.SizeBytes.Unit == "GiB" || s.SizeBytes.Unit == "GB" { + sizeGB = float64(s.SizeBytes.Value) + } else if s.SizeBytes.Unit == "MiB" || s.SizeBytes.Unit == "MB" { + sizeGB = float64(s.SizeBytes.Value) / 1024 + } else if s.SizeBytes.Unit == "TiB" || s.SizeBytes.Unit == "TB" { + sizeGB = float64(s.SizeBytes.Value) * 1024 + } + } + + return sizeGB, sizeGB +} + // gpuCapabilityEntry represents a GPU pattern and its compute capability type gpuCapabilityEntry struct { pattern string @@ -229,6 +361,7 @@ func getGPUCapability(gpuName string) float64 { {"RTXPRO6000", 12.0}, // NVIDIA Blackwell + {"B300", 10.3}, {"B200", 10.0}, {"RTX5090", 10.0}, @@ -301,6 +434,9 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { continue // Skip non-GPU instances } + // Extract disk size info from first storage entry + diskMin, diskMax := extractDiskSize(item.SupportedStorage) + for _, gpu := range item.SupportedGPUs { vramPerGPU := parseMemoryToGB(gpu.Memory) // Also check memory_bytes as fallback @@ -323,6 +459,7 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { instances = append(instances, GPUInstanceInfo{ Type: item.Type, + Provider: item.Provider, GPUName: gpu.Name, GPUCount: gpu.Count, VRAMPerGPU: vramPerGPU, @@ -330,6 +467,8 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { Capability: capability, VCPUs: item.VCPU, Memory: item.Memory, + DiskMin: diskMin, + DiskMax: diskMax, PricePerHour: price, Manufacturer: gpu.Manufacturer, }) @@ -340,7 +479,7 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { } // filterInstances applies all filters to the instance list -func filterInstances(instances []GPUInstanceInfo, gpuName string, minVRAM, minTotalVRAM, minCapability float64) []GPUInstanceInfo { +func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64) []GPUInstanceInfo { var filtered []GPUInstanceInfo for _, inst := range instances { @@ -354,6 +493,11 @@ func filterInstances(instances []GPUInstanceInfo, gpuName string, minVRAM, minTo continue } + // Filter by provider (case-insensitive partial match) + if provider != "" && !strings.Contains(strings.ToLower(inst.Provider), strings.ToLower(provider)) { + continue + } + // Filter by minimum VRAM per GPU if minVRAM > 0 && inst.VRAMPerGPU < minVRAM { continue @@ -369,6 +513,11 @@ func filterInstances(instances []GPUInstanceInfo, gpuName string, minVRAM, minTo continue } + // Filter by minimum disk size (use max available size for comparison) + if minDisk > 0 && inst.DiskMax < minDisk { + continue + } + filtered = append(filtered, inst) } @@ -394,6 +543,10 @@ func sortInstances(instances []GPUInstanceInfo, sortBy string, descending bool) less = instances[i].Type < instances[j].Type case "capability": less = instances[i].Capability < instances[j].Capability + case "provider": + less = instances[i].Provider < instances[j].Provider + case "disk": + less = instances[i].DiskMax < instances[j].DiskMax default: less = instances[i].PricePerHour < instances[j].PricePerHour } @@ -415,13 +568,34 @@ func getBrevTableOptions() table.Options { return options } +// formatDiskSize formats the disk size for display +func formatDiskSize(minGB, maxGB float64) string { + if minGB == 0 && maxGB == 0 { + return "-" + } + + formatSize := func(gb float64) string { + if gb >= 1000 { + return fmt.Sprintf("%.0fTB", gb/1000) + } + return fmt.Sprintf("%.0fGB", gb) + } + + if minGB == maxGB { + // Fixed size + return formatSize(minGB) + } + // Range + return fmt.Sprintf("%s-%s", formatSize(minGB), formatSize(maxGB)) +} + // 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", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -431,15 +605,18 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { if inst.Capability > 0 { capStr = fmt.Sprintf("%.1f", inst.Capability) } + diskStr := formatDiskSize(inst.DiskMin, inst.DiskMax) priceStr := fmt.Sprintf("$%.2f", inst.PricePerHour) row := table.Row{ inst.Type, + inst.Provider, t.Green(inst.GPUName), inst.GPUCount, vramStr, totalVramStr, capStr, + diskStr, inst.VCPUs, priceStr, } diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index 0714874f..15456a79 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -168,19 +168,19 @@ func TestFilterInstancesByGPUName(t *testing.T) { instances := processInstances(response.Items) // Filter by A10G - filtered := filterInstances(instances, "A10G", 0, 0, 0) + filtered := filterInstances(instances, "A10G", "", 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances") // Filter by V100 - filtered = filterInstances(instances, "V100", 0, 0, 0) + filtered = filterInstances(instances, "V100", "", 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances") // Filter by lowercase (case-insensitive) - filtered = filterInstances(instances, "v100", 0, 0, 0) + filtered = filterInstances(instances, "v100", "", 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances (case-insensitive)") // Filter by partial match - filtered = filterInstances(instances, "A1", 0, 0, 0) + filtered = filterInstances(instances, "A1", "", 0, 0, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances matching 'A1' (A10G and A100)") } @@ -189,11 +189,11 @@ func TestFilterInstancesByMinVRAM(t *testing.T) { instances := processInstances(response.Items) // Filter by min VRAM 24GB - filtered := filterInstances(instances, "", 24, 0, 0) + filtered := filterInstances(instances, "", "", 24, 0, 0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with >= 24GB VRAM") // Filter by min VRAM 40GB - filtered = filterInstances(instances, "", 40, 0, 0) + filtered = filterInstances(instances, "", "", 40, 0, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 40GB VRAM") assert.Equal(t, "A100", filtered[0].GPUName) } @@ -203,11 +203,11 @@ func TestFilterInstancesByMinTotalVRAM(t *testing.T) { instances := processInstances(response.Items) // Filter by min total VRAM 60GB - filtered := filterInstances(instances, "", 0, 60, 0) + filtered := filterInstances(instances, "", "", 0, 60, 0, 0) assert.Len(t, filtered, 2, "Should have 2 instances with >= 60GB total VRAM") // Filter by min total VRAM 300GB - filtered = filterInstances(instances, "", 0, 300, 0) + filtered = filterInstances(instances, "", "", 0, 300, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 300GB total VRAM") assert.Equal(t, "p4d.24xlarge", filtered[0].Type) } @@ -217,11 +217,11 @@ func TestFilterInstancesByMinCapability(t *testing.T) { instances := processInstances(response.Items) // Filter by capability >= 8.0 - filtered := filterInstances(instances, "", 0, 0, 8.0) + filtered := filterInstances(instances, "", "", 0, 0, 8.0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with capability >= 8.0") // Filter by capability >= 8.5 - filtered = filterInstances(instances, "", 0, 0, 8.5) + filtered = filterInstances(instances, "", "", 0, 0, 8.5, 0) assert.Len(t, filtered, 3, "Should have 3 instances with capability >= 8.5") } @@ -230,11 +230,11 @@ func TestFilterInstancesCombined(t *testing.T) { instances := processInstances(response.Items) // Filter by GPU name and min VRAM - filtered := filterInstances(instances, "A10G", 24, 0, 0) + filtered := filterInstances(instances, "A10G", "", 24, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances with >= 24GB VRAM") // Filter by GPU name, min VRAM, and capability - filtered = filterInstances(instances, "", 24, 0, 8.5) + filtered = filterInstances(instances, "", "", 24, 0, 8.5, 0) assert.Len(t, filtered, 3, "Should have 3 instances with >= 24GB VRAM and capability >= 8.5") } @@ -336,7 +336,7 @@ func TestEmptyInstanceTypes(t *testing.T) { assert.Len(t, instances, 0, "Should have 0 instances") - filtered := filterInstances(instances, "A100", 0, 0, 0) + filtered := filterInstances(instances, "A100", "", 0, 0, 0, 0) assert.Len(t, filtered, 0, "Filtered should also be empty") } diff --git a/pkg/store/instancetypes.go b/pkg/store/instancetypes.go index 4f12710a..e2daa047 100644 --- a/pkg/store/instancetypes.go +++ b/pkg/store/instancetypes.go @@ -2,6 +2,8 @@ package store import ( "encoding/json" + "fmt" + "runtime" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" breverrors "github.com/brevdev/brev-cli/pkg/errors" @@ -11,6 +13,8 @@ import ( const ( instanceTypesAPIURL = "https://api.brev.dev" instanceTypesAPIPath = "v1/instance/types" + // Authenticated API for instance types with workspace groups + allInstanceTypesPathPattern = "api/instances/alltypesavailable/%s" ) // GetInstanceTypes fetches all available instance types from the public API @@ -46,3 +50,24 @@ func fetchInstanceTypes() (*gpusearch.InstanceTypesResponse, error) { return &result, nil } + +// GetAllInstanceTypesWithWorkspaceGroups fetches instance types with workspace groups from the authenticated API +func (s AuthHTTPStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { + path := fmt.Sprintf(allInstanceTypesPathPattern, orgID) + + var result gpusearch.AllInstanceTypesResponse + res, err := s.authHTTPClient.restyClient.R(). + SetHeader("Content-Type", "application/json"). + SetQueryParam("utm_source", "cli"). + SetQueryParam("os", runtime.GOOS). + SetResult(&result). + Get(path) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + if res.IsError() { + return nil, NewHTTPResponseError(res) + } + + return &result, nil +} From bbc250df252cf2050f5059ffaf1df58a9db45ec5 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Tue, 27 Jan 2026 22:28:30 -0800 Subject: [PATCH 07/13] Add startup script support to gpu-create command Adds --startup-script (-s) flag to attach a startup script that runs when the instance boots. 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 Also improves CLI documentation with examples for startup script usage. --- pkg/cmd/gpucreate/gpucreate.go | 80 ++++++++-- pkg/cmd/gpucreate/gpucreate_test.go | 54 ++----- pkg/cmd/gpusearch/gpusearch.go | 230 +++++++++++++++++----------- pkg/cmd/gpusearch/gpusearch_test.go | 26 ++-- 4 files changed, 226 insertions(+), 164 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index a42b333a..033ad72b 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -9,6 +9,7 @@ import ( "strings" "sync" "time" + "unicode" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/util" @@ -32,7 +33,14 @@ The command will: 1. Try to create instances using the provided instance types (in order) 2. Continue until the desired count is reached 3. Optionally try multiple instance types in parallel -4. Clean up any extra instances that were created beyond the requested count` +4. Clean up any extra instances that were created beyond the requested count + +Startup Scripts: +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` example = ` # Create a single instance with a specific GPU type @@ -46,6 +54,15 @@ The command will: # Try multiple specific types in order brev gpu-create --name my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge + + # Attach a startup script from a file + brev gpu-create --name my-instance --type g5.xlarge --startup-script @setup.sh + + # Attach an inline startup script + brev gpu-create --name my-instance --type g5.xlarge --startup-script 'pip install torch transformers' + + # Combine: find cheapest A100, attach setup script + brev gpus --gpu-name A100 --sort price | brev gpu-create --name ml-box --startup-script @ml-setup.sh ` ) @@ -75,6 +92,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra var parallel int var detached bool var timeout int + var startupScript string cmd := &cobra.Command{ Annotations: map[string]string{"workspace": ""}, @@ -107,6 +125,12 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra parallel = 1 } + // Parse startup script (can be a string or @filepath) + scriptContent, err := parseStartupScript(startupScript) + if err != nil { + return breverrors.WrapAndTrace(err) + } + opts := GPUCreateOptions{ Name: name, InstanceTypes: types, @@ -114,6 +138,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra Parallel: parallel, Detached: detached, Timeout: time.Duration(timeout) * time.Second, + StartupScript: scriptContent, } err = RunGPUCreate(t, gpuCreateStore, opts) @@ -130,6 +155,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra cmd.Flags().IntVarP(¶llel, "parallel", "p", 1, "Number of parallel creation attempts") cmd.Flags().BoolVarP(&detached, "detached", "d", false, "Don't wait for instances to be ready") cmd.Flags().IntVar(&timeout, "timeout", 300, "Timeout in seconds for each instance to become ready") + cmd.Flags().StringVarP(&startupScript, "startup-script", "s", "", "Startup script to run on instance (string or @filepath)") return cmd } @@ -142,6 +168,28 @@ type GPUCreateOptions struct { Parallel int Detached bool Timeout time.Duration + StartupScript string +} + +// parseStartupScript parses the startup script from a string or file path +// If the value starts with @, it's treated as a file path +func parseStartupScript(value string) (string, error) { + if value == "" { + return "", nil + } + + // Check if it's a file path (prefixed with @) + if strings.HasPrefix(value, "@") { + filePath := strings.TrimPrefix(value, "@") + content, err := os.ReadFile(filePath) + if err != nil { + return "", breverrors.WrapAndTrace(err) + } + return string(content), nil + } + + // Otherwise, treat it as the script content directly + return value, nil } // parseInstanceTypes parses instance types from flag value or stdin @@ -205,27 +253,24 @@ func parseInstanceTypes(flagValue string) ([]string, error) { return types, nil } -// isValidInstanceType checks if a string looks like a valid instance type +// isValidInstanceType checks if a string looks like a valid instance type. +// Instance types typically have formats like: g5.xlarge, p4d.24xlarge, n1-highmem-4:nvidia-tesla-t4:1 func isValidInstanceType(s string) bool { - // Instance types typically have formats like: - // g5.xlarge, p4d.24xlarge, n1-highmem-4:nvidia-tesla-t4:1 if len(s) < 2 { return false } - - // Should contain alphanumeric characters - hasLetter := false - hasNumber := false + var hasLetter, hasDigit bool for _, c := range s { - if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') { + if unicode.IsLetter(c) { hasLetter = true + } else if unicode.IsDigit(c) { + hasDigit = true } - if c >= '0' && c <= '9' { - hasNumber = true + if hasLetter && hasDigit { + return true } } - - return hasLetter && hasNumber + return hasLetter && hasDigit } // RunGPUCreate executes the GPU create with retry logic @@ -318,7 +363,7 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC t.Vprintf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, instanceType, instanceName) // Attempt to create the workspace - workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user, allInstanceTypes) + workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user, allInstanceTypes, opts.StartupScript) result := CreateResult{ Workspace: workspace, @@ -415,7 +460,7 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC } // createWorkspaceWithType creates a workspace with the specified instance type -func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User, allInstanceTypes *gpusearch.AllInstanceTypesResponse) (*entity.Workspace, error) { +func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User, allInstanceTypes *gpusearch.AllInstanceTypesResponse, startupScript string) (*entity.Workspace, error) { clusterID := config.GlobalConfig.GetDefaultClusterID() cwOptions := store.NewCreateWorkspacesOptions(clusterID, name) cwOptions.WithInstanceType(instanceType) @@ -429,6 +474,11 @@ func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanc } } + // Set startup script if provided + if startupScript != "" { + cwOptions.StartupScript = startupScript + } + workspace, err := gpuCreateStore.CreateWorkspace(orgID, cwOptions) if err != nil { return nil, breverrors.WrapAndTrace(err) diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index 8c9b935d..059eff2a 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -1,8 +1,10 @@ package gpucreate import ( + "strings" "testing" + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/store" "github.com/stretchr/testify/assert" @@ -91,6 +93,10 @@ func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID strin return []entity.Workspace{}, nil } +func (m *MockGPUCreateStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { + return nil, nil +} + func TestIsValidInstanceType(t *testing.T) { tests := []struct { name string @@ -280,17 +286,17 @@ func TestParseInstanceTypesFromTableOutput(t *testing.T) { lineNum++ // Skip header line - if lineNum == 1 && (contains(line, "TYPE") || contains(line, "GPU")) { + if lineNum == 1 && (strings.Contains(line, "TYPE") || strings.Contains(line, "GPU")) { continue } // Skip empty lines and summary - if line == "" || startsWith(line, "Found ") { + if line == "" || strings.HasPrefix(line, "Found ") { continue } // Extract first column - fields := splitFields(line) + fields := strings.Fields(line) if len(fields) > 0 && isValidInstanceType(fields[0]) { types = append(types, fields[0]) } @@ -302,45 +308,3 @@ func TestParseInstanceTypesFromTableOutput(t *testing.T) { assert.Contains(t, types, "p4d.24xlarge") } -// Helper functions for testing -func contains(s, substr string) bool { - return len(s) >= len(substr) && findSubstring(s, substr) >= 0 -} - -func findSubstring(s, substr string) int { - for i := 0; i <= len(s)-len(substr); i++ { - if s[i:i+len(substr)] == substr { - return i - } - } - return -1 -} - -func startsWith(s, prefix string) bool { - return len(s) >= len(prefix) && s[:len(prefix)] == prefix -} - -func splitFields(s string) []string { - var fields []string - var current string - inField := false - - for _, c := range s { - if c == ' ' || c == '\t' { - if inField { - fields = append(fields, current) - current = "" - inField = false - } - } else { - current += string(c) - inField = true - } - } - - if inField { - fields = append(fields, current) - } - - return fields -} diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index e8fa5a3a..d435bd43 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -2,6 +2,7 @@ package gpusearch import ( + "encoding/json" "fmt" "os" "regexp" @@ -55,17 +56,18 @@ type WorkspaceGroup struct { // InstanceType represents an instance type from the API type InstanceType struct { - Type string `json:"type"` - SupportedGPUs []GPU `json:"supported_gpus"` - SupportedStorage []Storage `json:"supported_storage"` - Memory string `json:"memory"` - VCPU int `json:"vcpu"` - BasePrice BasePrice `json:"base_price"` - Location string `json:"location"` - SubLocation string `json:"sub_location"` - AvailableLocations []string `json:"available_locations"` - Provider string `json:"provider"` - WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"` + Type string `json:"type"` + SupportedGPUs []GPU `json:"supported_gpus"` + SupportedStorage []Storage `json:"supported_storage"` + Memory string `json:"memory"` + VCPU int `json:"vcpu"` + BasePrice BasePrice `json:"base_price"` + Location string `json:"location"` + SubLocation string `json:"sub_location"` + AvailableLocations []string `json:"available_locations"` + Provider string `json:"provider"` + WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"` + EstimatedDeployTime string `json:"estimated_deploy_time"` } // InstanceTypesResponse represents the API response @@ -98,7 +100,7 @@ type GPUSearchStore interface { var ( long = `Search and filter GPU instance types available on Brev. -Filter instances by GPU name, provider, VRAM, total VRAM, GPU compute capability, and disk size. +Filter instances by GPU name, provider, VRAM, total VRAM, GPU compute capability, disk size, and boot time. Sort results by various columns to find the best instance for your needs.` example = ` @@ -125,14 +127,17 @@ Sort results by various columns to find the best instance for your needs.` # Filter by minimum disk size (in GB) brev gpu-search --min-disk 500 - # Sort by different columns (price, gpu-count, vram, total-vram, vcpu, provider, disk) + # Filter by maximum boot time (in minutes) + brev gpu-search --max-boot-time 5 + + # Sort by different columns (price, gpu-count, vram, total-vram, vcpu, provider, disk, boot-time) brev gpu-search --sort price - brev gpu-search --sort provider + brev gpu-search --sort boot-time brev gpu-search --sort disk --desc # Combine filters brev gpu-search --gpu-name A100 --min-vram 40 --sort price - brev gpu-search --provider aws --gpu-name A100 --sort price + brev gpu-search --gpu-name H100 --max-boot-time 3 --sort price ` ) @@ -144,8 +149,10 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command var minTotalVRAM float64 var minCapability float64 var minDisk float64 + var maxBootTime int var sortBy string var descending bool + var jsonOutput bool cmd := &cobra.Command{ Annotations: map[string]string{"workspace": ""}, @@ -156,7 +163,7 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command Long: long, Example: example, RunE: func(cmd *cobra.Command, args []string) error { - err := RunGPUSearch(t, store, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, sortBy, descending) + err := RunGPUSearch(t, store, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, maxBootTime, sortBy, descending, jsonOutput) if err != nil { return breverrors.WrapAndTrace(err) } @@ -170,37 +177,44 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command cmd.Flags().Float64VarP(&minTotalVRAM, "min-total-vram", "t", 0, "Minimum total VRAM (GPU count * VRAM) in GB") cmd.Flags().Float64VarP(&minCapability, "min-capability", "c", 0, "Minimum GPU compute capability (e.g., 8.0 for Ampere)") cmd.Flags().Float64Var(&minDisk, "min-disk", 0, "Minimum disk size in GB") - cmd.Flags().StringVarP(&sortBy, "sort", "s", "price", "Sort by: price, gpu-count, vram, total-vram, vcpu, type, provider, disk") + cmd.Flags().IntVar(&maxBootTime, "max-boot-time", 0, "Maximum boot time in minutes") + cmd.Flags().StringVarP(&sortBy, "sort", "s", "price", "Sort by: price, gpu-count, vram, total-vram, vcpu, type, provider, disk, boot-time") cmd.Flags().BoolVarP(&descending, "desc", "d", false, "Sort in descending order") + cmd.Flags().BoolVar(&jsonOutput, "json", false, "Output results as JSON") return cmd } // GPUInstanceInfo holds processed GPU instance information for display type GPUInstanceInfo struct { - Type string - Provider string - GPUName string - GPUCount int - VRAMPerGPU float64 // in GB - TotalVRAM float64 // in GB - Capability float64 - VCPUs int - Memory string - DiskMin float64 // in GB (min disk size, same as DiskMax if fixed) - DiskMax float64 // in GB (max disk size, same as DiskMin if fixed) - PricePerHour float64 - Manufacturer string + Type string `json:"type"` + Provider string `json:"provider"` + 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"` + DiskMin float64 `json:"disk_min_gb"` + DiskMax float64 `json:"disk_max_gb"` + BootTime int `json:"boot_time_seconds"` + PricePerHour float64 `json:"price_per_hour"` + Manufacturer string `json:"-"` // exclude from JSON output } // RunGPUSearch executes the GPU search with filters and sorting -func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, sortBy string, descending bool) error { +func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, maxBootTime int, sortBy string, descending, jsonOutput bool) error { response, err := store.GetInstanceTypes() if err != nil { return breverrors.WrapAndTrace(err) } if response == nil || len(response.Items) == 0 { + if jsonOutput { + fmt.Println("[]") + return nil + } t.Vprint(t.Yellow("No instance types found")) return nil } @@ -209,9 +223,13 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider instances := processInstances(response.Items) // Apply filters - filtered := filterInstances(instances, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk) + filtered := filterInstances(instances, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, maxBootTime) if len(filtered) == 0 { + if jsonOutput { + fmt.Println("[]") + return nil + } t.Vprint(t.Yellow("No GPU instances match the specified filters")) return nil } @@ -220,91 +238,87 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider sortInstances(filtered, sortBy, descending) // Display results - displayGPUTable(t, filtered) + if jsonOutput { + return displayGPUJSON(filtered) + } + displayGPUTable(t, filtered) t.Vprintf("\n%s\n", t.Green(fmt.Sprintf("Found %d GPU instance types", len(filtered)))) return nil } -// parseMemoryToGB converts memory string like "22GiB360MiB" or "40GiB" to GB -func parseMemoryToGB(memory string) float64 { - // Handle memory_bytes if provided (in MiB) - // Otherwise parse the string format - - var totalGB float64 - - // Match GiB values - gibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GiB`) - gibMatches := gibRe.FindAllStringSubmatch(memory, -1) - for _, match := range gibMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val - } +// displayGPUJSON outputs the GPU instances as JSON +func displayGPUJSON(instances []GPUInstanceInfo) error { + output, err := json.MarshalIndent(instances, "", " ") + if err != nil { + return breverrors.WrapAndTrace(err) } + fmt.Println(string(output)) + return nil +} - // Match MiB values and convert to GB - mibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*MiB`) - mibMatches := mibRe.FindAllStringSubmatch(memory, -1) - for _, match := range mibMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val / 1024 - } - } +// unitMultipliers maps size units to their GB equivalent +var unitMultipliers = map[string]float64{ + "TiB": 1024, + "TB": 1000, + "GiB": 1, + "GB": 1, + "MiB": 1.0 / 1024, + "MB": 1.0 / 1000, +} - // Match GB values (in case API uses GB instead of GiB) - gbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GB`) - gbMatches := gbRe.FindAllStringSubmatch(memory, -1) - for _, match := range gbMatches { +// parseToGB converts size/memory strings like "22GiB360MiB", "16TiB", "2TiB768GiB" to GB +func parseToGB(s string) float64 { + var totalGB float64 + re := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*(TiB|TB|GiB|GB|MiB|MB)`) + for _, match := range re.FindAllStringSubmatch(s, -1) { if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val + totalGB += val * unitMultipliers[match[2]] } } - return totalGB } +// parseMemoryToGB converts memory string like "22GiB360MiB" or "40GiB" to GB +func parseMemoryToGB(memory string) float64 { + return parseToGB(memory) +} + // parseSizeToGB parses size strings like "16TiB", "10GiB", "2TiB768GiB" to GB func parseSizeToGB(size string) float64 { - var totalGB float64 + return parseToGB(size) +} - // Match TiB values and convert to GB (1 TiB = 1024 GiB) - tibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*TiB`) - tibMatches := tibRe.FindAllStringSubmatch(size, -1) - for _, match := range tibMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val * 1024 - } - } +// parseDurationToSeconds parses Go duration strings like "7m0s", "1m30s" to seconds +func parseDurationToSeconds(duration string) int { + var totalSeconds int - // Match GiB values - gibRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GiB`) - gibMatches := gibRe.FindAllStringSubmatch(size, -1) - for _, match := range gibMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val + // Match hours + hRe := regexp.MustCompile(`(\d+)h`) + if match := hRe.FindStringSubmatch(duration); len(match) > 1 { + if val, err := strconv.Atoi(match[1]); err == nil { + totalSeconds += val * 3600 } } - // Match TB values (1 TB = 1000 GB) - tbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*TB`) - tbMatches := tbRe.FindAllStringSubmatch(size, -1) - for _, match := range tbMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val * 1000 + // Match minutes + mRe := regexp.MustCompile(`(\d+)m`) + if match := mRe.FindStringSubmatch(duration); len(match) > 1 { + if val, err := strconv.Atoi(match[1]); err == nil { + totalSeconds += val * 60 } } - // Match GB values - gbRe := regexp.MustCompile(`(\d+(?:\.\d+)?)\s*GB`) - gbMatches := gbRe.FindAllStringSubmatch(size, -1) - for _, match := range gbMatches { - if val, err := strconv.ParseFloat(match[1], 64); err == nil { - totalGB += val + // Match seconds + sRe := regexp.MustCompile(`(\d+)s`) + if match := sRe.FindStringSubmatch(duration); len(match) > 1 { + if val, err := strconv.Atoi(match[1]); err == nil { + totalSeconds += val } } - return totalGB + return totalSeconds } // extractDiskSize extracts min and max disk size from storage configuration @@ -437,6 +451,9 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { // Extract disk size info from first storage entry diskMin, diskMax := extractDiskSize(item.SupportedStorage) + // Extract boot time + bootTime := parseDurationToSeconds(item.EstimatedDeployTime) + for _, gpu := range item.SupportedGPUs { vramPerGPU := parseMemoryToGB(gpu.Memory) // Also check memory_bytes as fallback @@ -469,6 +486,7 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { Memory: item.Memory, DiskMin: diskMin, DiskMax: diskMax, + BootTime: bootTime, PricePerHour: price, Manufacturer: gpu.Manufacturer, }) @@ -479,7 +497,7 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { } // filterInstances applies all filters to the instance list -func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64) []GPUInstanceInfo { +func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, maxBootTime int) []GPUInstanceInfo { var filtered []GPUInstanceInfo for _, inst := range instances { @@ -518,6 +536,11 @@ func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minV continue } + // Filter by maximum boot time (convert minutes to seconds for comparison) + if maxBootTime > 0 && inst.BootTime > maxBootTime*60 { + continue + } + filtered = append(filtered, inst) } @@ -547,6 +570,16 @@ func sortInstances(instances []GPUInstanceInfo, sortBy string, descending bool) less = instances[i].Provider < instances[j].Provider case "disk": less = instances[i].DiskMax < instances[j].DiskMax + case "boot-time": + // Instances with no boot time (0) should always appear last + if instances[i].BootTime == 0 && instances[j].BootTime == 0 { + return false // both unknown, equal + } else if instances[i].BootTime == 0 { + return false // i unknown goes after j + } else if instances[j].BootTime == 0 { + return true // j unknown goes after i + } + less = instances[i].BootTime < instances[j].BootTime default: less = instances[i].PricePerHour < instances[j].PricePerHour } @@ -589,13 +622,26 @@ func formatDiskSize(minGB, maxGB float64) string { return fmt.Sprintf("%s-%s", formatSize(minGB), formatSize(maxGB)) } +// formatBootTime formats boot time in seconds to a human-readable string +func formatBootTime(seconds int) string { + if seconds == 0 { + return "-" + } + minutes := seconds / 60 + secs := seconds % 60 + if secs == 0 { + return fmt.Sprintf("%dm", minutes) + } + return fmt.Sprintf("%dm%ds", minutes, secs) +} + // 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", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "BOOT", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -606,6 +652,7 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { capStr = fmt.Sprintf("%.1f", inst.Capability) } diskStr := formatDiskSize(inst.DiskMin, inst.DiskMax) + bootStr := formatBootTime(inst.BootTime) priceStr := fmt.Sprintf("$%.2f", inst.PricePerHour) row := table.Row{ @@ -617,6 +664,7 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { totalVramStr, capStr, diskStr, + bootStr, inst.VCPUs, priceStr, } diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index 15456a79..1e56637e 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -168,19 +168,19 @@ func TestFilterInstancesByGPUName(t *testing.T) { instances := processInstances(response.Items) // Filter by A10G - filtered := filterInstances(instances, "A10G", "", 0, 0, 0, 0) + filtered := filterInstances(instances, "A10G", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances") // Filter by V100 - filtered = filterInstances(instances, "V100", "", 0, 0, 0, 0) + filtered = filterInstances(instances, "V100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances") // Filter by lowercase (case-insensitive) - filtered = filterInstances(instances, "v100", "", 0, 0, 0, 0) + filtered = filterInstances(instances, "v100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances (case-insensitive)") // Filter by partial match - filtered = filterInstances(instances, "A1", "", 0, 0, 0, 0) + filtered = filterInstances(instances, "A1", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances matching 'A1' (A10G and A100)") } @@ -189,11 +189,11 @@ func TestFilterInstancesByMinVRAM(t *testing.T) { instances := processInstances(response.Items) // Filter by min VRAM 24GB - filtered := filterInstances(instances, "", "", 24, 0, 0, 0) + filtered := filterInstances(instances, "", "", 24, 0, 0, 0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with >= 24GB VRAM") // Filter by min VRAM 40GB - filtered = filterInstances(instances, "", "", 40, 0, 0, 0) + filtered = filterInstances(instances, "", "", 40, 0, 0, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 40GB VRAM") assert.Equal(t, "A100", filtered[0].GPUName) } @@ -203,11 +203,11 @@ func TestFilterInstancesByMinTotalVRAM(t *testing.T) { instances := processInstances(response.Items) // Filter by min total VRAM 60GB - filtered := filterInstances(instances, "", "", 0, 60, 0, 0) + filtered := filterInstances(instances, "", "", 0, 60, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 instances with >= 60GB total VRAM") // Filter by min total VRAM 300GB - filtered = filterInstances(instances, "", "", 0, 300, 0, 0) + filtered = filterInstances(instances, "", "", 0, 300, 0, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 300GB total VRAM") assert.Equal(t, "p4d.24xlarge", filtered[0].Type) } @@ -217,11 +217,11 @@ func TestFilterInstancesByMinCapability(t *testing.T) { instances := processInstances(response.Items) // Filter by capability >= 8.0 - filtered := filterInstances(instances, "", "", 0, 0, 8.0, 0) + filtered := filterInstances(instances, "", "", 0, 0, 8.0, 0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with capability >= 8.0") // Filter by capability >= 8.5 - filtered = filterInstances(instances, "", "", 0, 0, 8.5, 0) + filtered = filterInstances(instances, "", "", 0, 0, 8.5, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances with capability >= 8.5") } @@ -230,11 +230,11 @@ func TestFilterInstancesCombined(t *testing.T) { instances := processInstances(response.Items) // Filter by GPU name and min VRAM - filtered := filterInstances(instances, "A10G", "", 24, 0, 0, 0) + filtered := filterInstances(instances, "A10G", "", 24, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances with >= 24GB VRAM") // Filter by GPU name, min VRAM, and capability - filtered = filterInstances(instances, "", "", 24, 0, 8.5, 0) + filtered = filterInstances(instances, "", "", 24, 0, 8.5, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances with >= 24GB VRAM and capability >= 8.5") } @@ -336,7 +336,7 @@ func TestEmptyInstanceTypes(t *testing.T) { assert.Len(t, instances, 0, "Should have 0 instances") - filtered := filterInstances(instances, "A100", "", 0, 0, 0, 0) + filtered := filterInstances(instances, "A100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 0, "Filtered should also be empty") } From 0f45d145376c7cc4c5ee95b8ccff13ae958748e8 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 00:50:55 -0800 Subject: [PATCH 08/13] Add piping support and -c flag to brev shell for scriptable workflows Enable seamless command chaining between provision, shell, and open commands: - Add stdin support to `brev shell` and `brev open` to read instance names from piped input, enabling `brev provision | brev shell -c "cmd"` - Add `-c` flag to `brev shell` for running non-interactive commands with stdout/stderr piped back (supports inline commands and @filepath syntax) - Fix error output to go to stderr instead of stdout for proper piping - Add default GPU selection for `brev provision` when no type specified (min 20GB VRAM, 500GB disk, capability 8.0+, boot time <7min) - Skip version check output for gpu-create/provision commands when piped - Simplify gpusearch code by consolidating parseMemoryToGB/parseSizeToGB - Export ProcessInstances, FilterInstances, SortInstances from gpusearch - Fix duplicate workspace name error to fail immediately instead of retry - Remove dead code: unused -r/--remote and -d/--dir flags from brev shell - Remove debug print statement from workspace creation Example workflows now supported: brev provision --name my-instance | brev shell -c "nvidia-smi" brev provision --name my-instance | brev open brev shell $(brev provision --name my-instance) # interactive --- pkg/cmd/cmd.go | 9 +- pkg/cmd/cmderrors/cmderrors.go | 6 +- pkg/cmd/gpucreate/gpucreate.go | 353 ++++++++++++++++++++++------ pkg/cmd/gpucreate/gpucreate_test.go | 121 +++++++++- pkg/cmd/gpusearch/gpusearch.go | 120 +++++++--- pkg/cmd/gpusearch/gpusearch_test.go | 141 +++++++---- pkg/cmd/open/open.go | 65 ++++- pkg/cmd/shell/shell.go | 112 ++++++++- pkg/store/workspace.go | 19 +- 9 files changed, 764 insertions(+), 182 deletions(-) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index c7253f26..9529daee 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -159,8 +159,13 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin PersistentPreRunE: func(cmd *cobra.Command, args []string) error { breverrors.GetDefaultErrorReporter().AddTag("command", cmd.Name()) // version info gets in the way of the output for - // configure-env-vars, since shells are going to eval it - if featureflag.ShowVersionOnRun() && !printVersion && cmd.Name() != "configure-env-vars" { + // configure-env-vars (shells eval it) and gpu-create/provision (piped to other commands) + skipVersionCommands := map[string]bool{ + "configure-env-vars": true, + "gpu-create": true, + "provision": true, + } + if featureflag.ShowVersionOnRun() && !printVersion && !skipVersionCommands[cmd.Name()] { v, err := remoteversion.BuildCheckLatestVersionString(t, noLoginCmdStore) // todo this should not be fatal when it errors if err != nil { diff --git a/pkg/cmd/cmderrors/cmderrors.go b/pkg/cmd/cmderrors/cmderrors.go index 6c290e59..4b07989a 100644 --- a/pkg/cmd/cmderrors/cmderrors.go +++ b/pkg/cmd/cmderrors/cmderrors.go @@ -39,7 +39,7 @@ func DisplayAndHandleError(err error) { case *breverrors.NvidiaMigrationError: // Handle nvidia migration error if nvErr, ok := errors.Cause(err).(*breverrors.NvidiaMigrationError); ok { - fmt.Println("\n This account has been migrated to NVIDIA Auth. Attempting to log in with NVIDIA account...") + fmt.Fprintln(os.Stderr, "\n This account has been migrated to NVIDIA Auth. Attempting to log in with NVIDIA account...") brevBin, err1 := os.Executable() if err1 == nil { cmd := exec.Command(brevBin, "login", "--auth", "nvidia") // #nosec G204 @@ -68,9 +68,9 @@ func DisplayAndHandleError(err error) { } } if featureflag.Debug() || featureflag.IsDev() { - fmt.Println(err) + fmt.Fprintln(os.Stderr, err) } else { - fmt.Println(prettyErr) + fmt.Fprintln(os.Stderr, prettyErr) } } } diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index 033ad72b..429477e1 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -2,9 +2,10 @@ package gpucreate import ( - "bufio" "context" + "encoding/json" "fmt" + "io" "os" "strings" "sync" @@ -27,7 +28,16 @@ var ( This command attempts to create GPU instances, trying different instance types until the desired number of instances are successfully created. Instance types -can be specified directly or piped from 'brev gpus'. +can be specified directly, piped from 'brev gpus', or auto-selected using defaults. + +Default Behavior: +If no instance types are specified (no --type flag and no piped input), the command +automatically searches for GPUs matching these criteria: + - Minimum 20GB total VRAM + - Minimum 500GB disk + - Compute capability 8.0+ (Ampere or newer) + - Boot time under 7 minutes +Results are sorted by price (cheapest first). The command will: 1. Try to create instances using the provided instance types (in order) @@ -43,6 +53,18 @@ You can attach a startup script that runs when the instance boots using the - An absolute file path: --startup-script @/path/to/setup.sh` example = ` + # Quick start: create an instance using smart defaults (sorted by price) + brev provision --name my-instance + + # Create and immediately open in VS Code (reads instance name from stdin) + brev provision --name my-instance | brev open + + # Create and SSH into the instance (use command substitution for interactive shell) + brev shell $(brev provision --name my-instance) + + # Create and run a command (non-interactive, reads instance name from stdin) + brev provision --name my-instance | brev shell -c "nvidia-smi" + # Create a single instance with a specific GPU type brev gpu-create --name my-instance --type g5.xlarge @@ -69,6 +91,7 @@ You can attach a startup script that runs when the instance boots using the // GPUCreateStore defines the interface for GPU create operations type GPUCreateStore interface { util.GetWorkspaceByNameOrIDErrStore + gpusearch.GPUSearchStore GetActiveOrganizationOrDefault() (*entity.Organization, error) GetCurrentUser() (*entity.User, error) GetWorkspace(workspaceID string) (*entity.Workspace, error) @@ -77,6 +100,14 @@ type GPUCreateStore interface { GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) } +// Default filter values for automatic GPU selection +const ( + defaultMinTotalVRAM = 20.0 // GB + defaultMinDisk = 500.0 // GB + defaultMinCapability = 8.0 + defaultMaxBootTime = 7 // minutes +) + // CreateResult holds the result of a workspace creation attempt type CreateResult struct { Workspace *entity.Workspace @@ -103,14 +134,33 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra Long: long, Example: example, RunE: func(cmd *cobra.Command, args []string) error { + // Check if output is being piped (for chaining with brev shell/open) + piped := isStdoutPiped() + // Parse instance types from flag or stdin types, err := parseInstanceTypes(instanceTypes) if err != nil { return breverrors.WrapAndTrace(err) } + // If no types provided, use default filters to find suitable GPUs if len(types) == 0 { - return breverrors.NewValidationError("no instance types provided. Use --type flag or pipe from 'brev gpus'") + msg := fmt.Sprintf("No instance types specified, using defaults: min-total-vram=%.0fGB, min-disk=%.0fGB, min-capability=%.1f, max-boot-time=%dm\n\n", + defaultMinTotalVRAM, defaultMinDisk, defaultMinCapability, defaultMaxBootTime) + if piped { + fmt.Fprint(os.Stderr, msg) + } else { + t.Vprint(msg) + } + + types, err = getDefaultInstanceTypes(gpuCreateStore) + if err != nil { + return breverrors.WrapAndTrace(err) + } + + if len(types) == 0 { + return breverrors.NewValidationError("no GPU instances match the default filters. Try 'brev gpus' to see available options") + } } if name == "" { @@ -160,10 +210,16 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra return cmd } +// 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 +} + // GPUCreateOptions holds the options for GPU instance creation type GPUCreateOptions struct { Name string - InstanceTypes []string + InstanceTypes []InstanceSpec Count int Parallel int Detached bool @@ -192,9 +248,40 @@ func parseStartupScript(value string) (string, error) { return value, nil } +// getDefaultInstanceTypes fetches GPU instance types using default filters +func getDefaultInstanceTypes(store GPUCreateStore) ([]InstanceSpec, error) { + response, err := store.GetInstanceTypes() + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + + if response == nil || len(response.Items) == 0 { + return nil, nil + } + + // Use gpusearch package to process, filter, and sort instances + instances := gpusearch.ProcessInstances(response.Items) + filtered := gpusearch.FilterInstances(instances, "", "", 0, defaultMinTotalVRAM, defaultMinCapability, defaultMinDisk, defaultMaxBootTime) + gpusearch.SortInstances(filtered, "price", false) + + // Convert to InstanceSpec with disk info + var specs []InstanceSpec + for _, inst := range filtered { + // For defaults, use the minimum disk size that meets the filter + diskGB := inst.DiskMin + if inst.DiskMin != inst.DiskMax && defaultMinDisk > inst.DiskMin && defaultMinDisk <= inst.DiskMax { + diskGB = defaultMinDisk + } + specs = append(specs, InstanceSpec{Type: inst.Type, DiskGB: diskGB}) + } + + return specs, nil +} + // parseInstanceTypes parses instance types from flag value or stdin -func parseInstanceTypes(flagValue string) ([]string, error) { - var types []string +// Returns InstanceSpec with type and optional disk size (from JSON input) +func parseInstanceTypes(flagValue string) ([]InstanceSpec, error) { + var specs []InstanceSpec // First check if there's a flag value if flagValue != "" { @@ -202,7 +289,7 @@ func parseInstanceTypes(flagValue string) ([]string, error) { for _, p := range parts { p = strings.TrimSpace(p) if p != "" { - types = append(types, p) + specs = append(specs, InstanceSpec{Type: p}) } } } @@ -210,47 +297,87 @@ func parseInstanceTypes(flagValue string) ([]string, error) { // Check if there's piped input from stdin stat, _ := os.Stdin.Stat() if (stat.Mode() & os.ModeCharDevice) == 0 { - // Data is being piped to stdin - scanner := bufio.NewScanner(os.Stdin) - lineNum := 0 - for scanner.Scan() { - line := scanner.Text() - lineNum++ - - // Skip header line (first line typically contains column names) - if lineNum == 1 && (strings.Contains(line, "TYPE") || strings.Contains(line, "GPU")) { - continue - } + // Data is being piped to stdin - read all input first + input, err := io.ReadAll(os.Stdin) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } - // Skip empty lines - line = strings.TrimSpace(line) - if line == "" { - continue - } + inputStr := strings.TrimSpace(string(input)) + if inputStr == "" { + return specs, nil + } - // Skip summary lines (e.g., "Found X GPU instance types") - if strings.HasPrefix(line, "Found ") { - continue + // Check if input is JSON (starts with '[') + if strings.HasPrefix(inputStr, "[") { + jsonSpecs, err := parseJSONInput(inputStr) + if err != nil { + return nil, breverrors.WrapAndTrace(err) } + specs = append(specs, jsonSpecs...) + } else { + // Parse as table format + tableSpecs := parseTableInput(inputStr) + specs = append(specs, tableSpecs...) + } + } - // Extract the first column (TYPE) from the table output - // The format is: TYPE GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR - fields := strings.Fields(line) - if len(fields) > 0 { - instanceType := fields[0] - // Validate it looks like an instance type (contains letters and possibly numbers/dots) - if isValidInstanceType(instanceType) { - types = append(types, instanceType) - } - } + return specs, nil +} + +// parseJSONInput parses JSON array input from gpu-search --json +func parseJSONInput(input string) ([]InstanceSpec, error) { + var instances []gpusearch.GPUInstanceInfo + if err := json.Unmarshal([]byte(input), &instances); err != nil { + return nil, breverrors.WrapAndTrace(err) + } + + var specs []InstanceSpec + for _, inst := range instances { + spec := InstanceSpec{ + Type: inst.Type, + DiskGB: inst.TargetDisk, } + specs = append(specs, spec) + } + return specs, nil +} - if err := scanner.Err(); err != nil { - return nil, breverrors.WrapAndTrace(err) +// parseTableInput parses table format input from gpu-search +func parseTableInput(input string) []InstanceSpec { + var specs []InstanceSpec + lines := strings.Split(input, "\n") + + for i, line := range lines { + // Skip header line (first line typically contains column names) + if i == 0 && (strings.Contains(line, "TYPE") || strings.Contains(line, "GPU")) { + continue + } + + // Skip empty lines + line = strings.TrimSpace(line) + if line == "" { + continue + } + + // Skip summary lines (e.g., "Found X GPU instance types") + if strings.HasPrefix(line, "Found ") { + continue + } + + // Extract the first column (TYPE) from the table output + // The format is: TYPE GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR + fields := strings.Fields(line) + if len(fields) > 0 { + instanceType := fields[0] + // Validate it looks like an instance type (contains letters and possibly numbers/dots) + if isValidInstanceType(instanceType) { + specs = append(specs, InstanceSpec{Type: instanceType}) + } } } - return types, nil + return specs } // isValidInstanceType checks if a string looks like a valid instance type. @@ -273,8 +400,44 @@ func isValidInstanceType(s string) bool { return hasLetter && hasDigit } +// isStdoutPiped returns true if stdout is being piped (not a terminal) +func isStdoutPiped() bool { + stat, _ := os.Stdout.Stat() + return (stat.Mode() & os.ModeCharDevice) == 0 +} + +// stderrPrintf prints to stderr (used when stdout is piped) +func stderrPrintf(format string, a ...interface{}) { + fmt.Fprintf(os.Stderr, format, a...) +} + +// formatInstanceSpecs formats a slice of InstanceSpec for display +func formatInstanceSpecs(specs []InstanceSpec) string { + var parts []string + for _, spec := range specs { + if spec.DiskGB > 0 { + parts = append(parts, fmt.Sprintf("%s (%.0fGB disk)", spec.Type, spec.DiskGB)) + } else { + parts = append(parts, spec.Type) + } + } + return strings.Join(parts, ", ") +} + // RunGPUCreate executes the GPU create with retry logic func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions) error { + // Check if output is being piped (for chaining with brev shell/open) + piped := isStdoutPiped() + + // Helper to print progress - uses stderr when piped so only instance name goes to stdout + logf := func(format string, a ...interface{}) { + if piped { + fmt.Fprintf(os.Stderr, format, a...) + } else { + t.Vprintf(format, a...) + } + } + user, err := gpuCreateStore.GetCurrentUser() if err != nil { return breverrors.WrapAndTrace(err) @@ -291,13 +454,13 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC // Fetch instance types with workspace groups to determine correct workspace group ID allInstanceTypes, err := gpuCreateStore.GetAllInstanceTypesWithWorkspaceGroups(org.ID) if err != nil { - t.Vprintf("Warning: could not fetch instance types with workspace groups: %s\n", err.Error()) - t.Vprintf("Falling back to default workspace group\n") + logf("Warning: could not fetch instance types with workspace groups: %s\n", err.Error()) + logf("Falling back to default workspace group\n") allInstanceTypes = nil } - t.Vprintf("Attempting to create %d instance(s) with %d parallel attempts\n", opts.Count, opts.Parallel) - t.Vprintf("Instance types to try: %s\n\n", strings.Join(opts.InstanceTypes, ", ")) + logf("Attempting to create %d instance(s) with %d parallel attempts\n", opts.Count, opts.Parallel) + logf("Instance types to try: %s\n\n", formatInstanceSpecs(opts.InstanceTypes)) // Track successful creations var successfulWorkspaces []*entity.Workspace @@ -309,18 +472,17 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC defer cancel() // Channel to coordinate attempts - typesChan := make(chan string, len(opts.InstanceTypes)) - for _, it := range opts.InstanceTypes { - typesChan <- it + specsChan := make(chan InstanceSpec, len(opts.InstanceTypes)) + for _, spec := range opts.InstanceTypes { + specsChan <- spec } - close(typesChan) + close(specsChan) // Results channel resultsChan := make(chan CreateResult, len(opts.InstanceTypes)) - // Track instance index for naming + // Track instance index for naming (incremented only on successful creation) instanceIndex := 0 - var indexMu sync.Mutex // Start parallel workers workerCount := opts.Parallel @@ -333,13 +495,15 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC go func(workerID int) { defer wg.Done() - for instanceType := range typesChan { + for spec := range specsChan { // Check if we've already created enough mu.Lock() if len(successfulWorkspaces) >= opts.Count { mu.Unlock() return } + // Get current index for naming (only increment on success) + currentIndex := instanceIndex mu.Unlock() // Check context @@ -349,34 +513,47 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC default: } - // Get unique instance name - indexMu.Lock() - currentIndex := instanceIndex - instanceIndex++ - indexMu.Unlock() - + // Determine instance name based on current successful count instanceName := opts.Name if opts.Count > 1 || currentIndex > 0 { instanceName = fmt.Sprintf("%s-%d", opts.Name, currentIndex+1) } - t.Vprintf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, instanceType, instanceName) + logf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, spec.Type, instanceName) // Attempt to create the workspace - workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, instanceType, user, allInstanceTypes, opts.StartupScript) + workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, spec.Type, spec.DiskGB, user, allInstanceTypes, opts.StartupScript) result := CreateResult{ Workspace: workspace, - InstanceType: instanceType, + InstanceType: spec.Type, Error: err, } if err != nil { - t.Vprintf("[Worker %d] %s Failed: %s\n", workerID+1, t.Yellow(instanceType), err.Error()) + errStr := err.Error() + if piped { + logf("[Worker %d] %s Failed: %s\n", workerID+1, spec.Type, errStr) + } else { + logf("[Worker %d] %s Failed: %s\n", workerID+1, t.Yellow(spec.Type), errStr) + } + + // Check for fatal errors that should stop all workers + if strings.Contains(errStr, "duplicate workspace") { + logf("\nError: Workspace '%s' already exists. Use a different name or delete the existing workspace.\n", instanceName) + cancel() // Stop all workers + resultsChan <- result + return + } } else { - t.Vprintf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, t.Green(instanceType), instanceName) + if piped { + logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, spec.Type, instanceName) + } else { + logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, t.Green(spec.Type), instanceName) + } mu.Lock() successfulWorkspaces = append(successfulWorkspaces, workspace) + instanceIndex++ // Only increment on success if len(successfulWorkspaces) >= opts.Count { cancel() // Signal other workers to stop } @@ -401,12 +578,12 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC // Check if we created enough instances if len(successfulWorkspaces) < opts.Count { - t.Vprintf("\n%s Only created %d/%d instances\n", t.Yellow("Warning:"), len(successfulWorkspaces), opts.Count) + logf("\nWarning: Only created %d/%d instances\n", len(successfulWorkspaces), opts.Count) if len(successfulWorkspaces) > 0 { - t.Vprintf("Successfully created instances:\n") + logf("Successfully created instances:\n") for _, ws := range successfulWorkspaces { - t.Vprintf(" - %s (ID: %s)\n", ws.Name, ws.ID) + logf(" - %s (ID: %s)\n", ws.Name, ws.ID) } } @@ -416,15 +593,15 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC // If we created more than needed, clean up extras if len(successfulWorkspaces) > opts.Count { extras := successfulWorkspaces[opts.Count:] - t.Vprintf("\nCleaning up %d extra instance(s)...\n", len(extras)) + logf("\nCleaning up %d extra instance(s)...\n", len(extras)) for _, ws := range extras { - t.Vprintf(" Deleting %s...", ws.Name) + logf(" Deleting %s...", ws.Name) _, err := gpuCreateStore.DeleteWorkspace(ws.ID) if err != nil { - t.Vprintf(" %s\n", t.Red("Failed")) + logf(" Failed\n") } else { - t.Vprintf(" %s\n", t.Green("Done")) + logf(" Done\n") } } @@ -433,17 +610,25 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC // Wait for instances to be ready (unless detached) if !opts.Detached { - t.Vprintf("\nWaiting for instance(s) to be ready...\n") - t.Vprintf("You can safely ctrl+c to exit\n") + logf("\nWaiting for instance(s) to be ready...\n") + logf("You can safely ctrl+c to exit\n") for _, ws := range successfulWorkspaces { - err := pollUntilReady(t, ws.ID, gpuCreateStore, opts.Timeout) + err := pollUntilReady(t, ws.ID, gpuCreateStore, opts.Timeout, piped, logf) if err != nil { - t.Vprintf(" %s: %s\n", ws.Name, t.Yellow("Timeout waiting for ready state")) + logf(" %s: Timeout waiting for ready state\n", ws.Name) } } } + // If output is piped, just print instance name(s) for chaining with brev shell/open + if piped { + for _, ws := range successfulWorkspaces { + fmt.Println(ws.Name) + } + return nil + } + // Print summary fmt.Print("\n") t.Vprint(t.Green(fmt.Sprintf("Successfully created %d instance(s)!\n\n", len(successfulWorkspaces)))) @@ -460,12 +645,17 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC } // createWorkspaceWithType creates a workspace with the specified instance type -func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, user *entity.User, allInstanceTypes *gpusearch.AllInstanceTypesResponse, startupScript string) (*entity.Workspace, error) { +func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanceType string, diskGB float64, user *entity.User, allInstanceTypes *gpusearch.AllInstanceTypesResponse, startupScript string) (*entity.Workspace, error) { clusterID := config.GlobalConfig.GetDefaultClusterID() cwOptions := store.NewCreateWorkspacesOptions(clusterID, name) cwOptions.WithInstanceType(instanceType) cwOptions = resolveWorkspaceUserOptions(cwOptions, user) + // Set disk size if specified (convert GB to Gi format) + if diskGB > 0 { + cwOptions.DiskStorage = fmt.Sprintf("%.0fGi", diskGB) + } + // Look up the workspace group ID for this instance type if allInstanceTypes != nil { workspaceGroupID := allInstanceTypes.GetWorkspaceGroupID(instanceType) @@ -474,9 +664,14 @@ func createWorkspaceWithType(gpuCreateStore GPUCreateStore, orgID, name, instanc } } - // Set startup script if provided + // Set startup script if provided using VMBuild lifecycle script if startupScript != "" { - cwOptions.StartupScript = startupScript + cwOptions.VMBuild = &store.VMBuild{ + ForceJupyterInstall: true, + LifeCycleScriptAttr: &store.LifeCycleScriptAttr{ + Script: startupScript, + }, + } } workspace, err := gpuCreateStore.CreateWorkspace(orgID, cwOptions) @@ -507,7 +702,7 @@ func resolveWorkspaceUserOptions(options *store.CreateWorkspacesOptions, user *e } // pollUntilReady waits for a workspace to reach the running state -func pollUntilReady(t *terminal.Terminal, wsID string, gpuCreateStore GPUCreateStore, timeout time.Duration) error { +func pollUntilReady(t *terminal.Terminal, wsID string, gpuCreateStore GPUCreateStore, timeout time.Duration, piped bool, logf func(string, ...interface{})) error { deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { @@ -517,7 +712,11 @@ func pollUntilReady(t *terminal.Terminal, wsID string, gpuCreateStore GPUCreateS } if ws.Status == entity.Running { - t.Vprintf(" %s: %s\n", ws.Name, t.Green("Ready")) + if piped { + logf(" %s: Ready\n", ws.Name) + } else { + logf(" %s: %s\n", ws.Name, t.Green("Ready")) + } return nil } diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index 059eff2a..2ad98680 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -97,6 +97,27 @@ func (m *MockGPUCreateStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string return nil, nil } +func (m *MockGPUCreateStore) GetInstanceTypes() (*gpusearch.InstanceTypesResponse, error) { + // Return a default set of instance types for testing + return &gpusearch.InstanceTypesResponse{ + Items: []gpusearch.InstanceType{ + { + Type: "g5.xlarge", + 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", + }, + }, + }, nil +} + func TestIsValidInstanceType(t *testing.T) { tests := []struct { name string @@ -145,7 +166,12 @@ func TestParseInstanceTypesFromFlag(t *testing.T) { if len(tt.expected) == 0 { assert.Empty(t, result) } else { - assert.Equal(t, tt.expected, result) + // Compare just the Type field of each InstanceSpec + var resultTypes []string + for _, spec := range result { + resultTypes = append(resultTypes, spec.Type) + } + assert.Equal(t, tt.expected, resultTypes) } }) } @@ -153,15 +179,22 @@ func TestParseInstanceTypesFromFlag(t *testing.T) { func TestGPUCreateOptions(t *testing.T) { opts := GPUCreateOptions{ - Name: "my-instance", - InstanceTypes: []string{"g5.xlarge", "g5.2xlarge"}, - Count: 2, - Parallel: 3, - Detached: true, + Name: "my-instance", + InstanceTypes: []InstanceSpec{ + {Type: "g5.xlarge", DiskGB: 500}, + {Type: "g5.2xlarge"}, + }, + Count: 2, + Parallel: 3, + Detached: true, } assert.Equal(t, "my-instance", opts.Name) assert.Len(t, opts.InstanceTypes, 2) + assert.Equal(t, "g5.xlarge", opts.InstanceTypes[0].Type) + assert.Equal(t, 500.0, opts.InstanceTypes[0].DiskGB) + assert.Equal(t, "g5.2xlarge", opts.InstanceTypes[1].Type) + assert.Equal(t, 0.0, opts.InstanceTypes[1].DiskGB) assert.Equal(t, 2, opts.Count) assert.Equal(t, 3, opts.Parallel) assert.True(t, opts.Detached) @@ -267,6 +300,32 @@ func TestMockGPUCreateStoreTypeSpecificError(t *testing.T) { assert.NotNil(t, ws) } +func TestGetDefaultInstanceTypes(t *testing.T) { + mock := NewMockGPUCreateStore() + + // Get default instance types - the mock returns a g5.xlarge which has: + // - 24GB VRAM (>= 20GB total VRAM requirement) + // - 500GB disk (>= 500GB requirement) + // - A10G GPU = 8.6 capability (>= 8.0 requirement) + // - 5m boot time (< 7m requirement) + specs, err := getDefaultInstanceTypes(mock) + assert.NoError(t, err) + assert.Len(t, specs, 1) + assert.Equal(t, "g5.xlarge", specs[0].Type) + assert.Equal(t, 500.0, specs[0].DiskGB) // Should use the instance's disk size +} + +func TestGetDefaultInstanceTypesFiltersOut(t *testing.T) { + // The mock returns a g5.xlarge which meets all requirements + mock := NewMockGPUCreateStore() + + specs, err := getDefaultInstanceTypes(mock) + assert.NoError(t, err) + // Should return the A10G instance which meets all requirements + assert.Len(t, specs, 1) + assert.Equal(t, "g5.xlarge", specs[0].Type) +} + func TestParseInstanceTypesFromTableOutput(t *testing.T) { // Simulated table output from brev gpus command // Note: This tests the parsing logic, not actual stdin reading @@ -308,3 +367,53 @@ func TestParseInstanceTypesFromTableOutput(t *testing.T) { assert.Contains(t, types, "p4d.24xlarge") } +func TestParseJSONInput(t *testing.T) { + // Simulated JSON output from gpu-search --json + jsonInput := `[ + { + "type": "g5.xlarge", + "provider": "aws", + "gpu_name": "A10G", + "target_disk_gb": 1000 + }, + { + "type": "p4d.24xlarge", + "provider": "aws", + "gpu_name": "A100", + "target_disk_gb": 500 + }, + { + "type": "g6.xlarge", + "provider": "aws", + "gpu_name": "L4" + } + ]` + + specs, err := parseJSONInput(jsonInput) + assert.NoError(t, err) + assert.Len(t, specs, 3) + + // Check first instance with disk + assert.Equal(t, "g5.xlarge", specs[0].Type) + assert.Equal(t, 1000.0, specs[0].DiskGB) + + // Check second instance with different disk + assert.Equal(t, "p4d.24xlarge", specs[1].Type) + assert.Equal(t, 500.0, specs[1].DiskGB) + + // Check third instance without disk (should be 0) + assert.Equal(t, "g6.xlarge", specs[2].Type) + assert.Equal(t, 0.0, specs[2].DiskGB) +} + +func TestFormatInstanceSpecs(t *testing.T) { + specs := []InstanceSpec{ + {Type: "g5.xlarge", DiskGB: 1000}, + {Type: "p4d.24xlarge", DiskGB: 0}, + {Type: "g6.xlarge", DiskGB: 500}, + } + + result := formatInstanceSpecs(specs) + assert.Equal(t, "g5.xlarge (1000GB disk), p4d.24xlarge, g6.xlarge (500GB disk)", result) +} + diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index d435bd43..b2d22471 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -39,12 +39,13 @@ type BasePrice struct { // Storage represents a storage configuration within an instance type type Storage struct { - Count int `json:"count"` - Size string `json:"size"` - Type string `json:"type"` - MinSize string `json:"min_size"` - MaxSize string `json:"max_size"` - SizeBytes MemoryBytes `json:"size_bytes"` + Count int `json:"count"` + Size string `json:"size"` + Type string `json:"type"` + MinSize string `json:"min_size"` + MaxSize string `json:"max_size"` + SizeBytes MemoryBytes `json:"size_bytes"` + PricePerGBHr float64 `json:"price_per_gb_hr"` } // WorkspaceGroup represents a workspace group that can run an instance type @@ -56,18 +57,21 @@ type WorkspaceGroup struct { // InstanceType represents an instance type from the API type InstanceType struct { - Type string `json:"type"` - SupportedGPUs []GPU `json:"supported_gpus"` - SupportedStorage []Storage `json:"supported_storage"` - Memory string `json:"memory"` - VCPU int `json:"vcpu"` - BasePrice BasePrice `json:"base_price"` - Location string `json:"location"` - SubLocation string `json:"sub_location"` - AvailableLocations []string `json:"available_locations"` - Provider string `json:"provider"` - WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"` - EstimatedDeployTime string `json:"estimated_deploy_time"` + Type string `json:"type"` + SupportedGPUs []GPU `json:"supported_gpus"` + SupportedStorage []Storage `json:"supported_storage"` + Memory string `json:"memory"` + VCPU int `json:"vcpu"` + BasePrice BasePrice `json:"base_price"` + Location string `json:"location"` + SubLocation string `json:"sub_location"` + AvailableLocations []string `json:"available_locations"` + Provider string `json:"provider"` + WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"` + EstimatedDeployTime string `json:"estimated_deploy_time"` + Stoppable bool `json:"stoppable"` + Rebootable bool `json:"rebootable"` + CanModifyFirewallRules bool `json:"can_modify_firewall_rules"` } // InstanceTypesResponse represents the API response @@ -101,7 +105,12 @@ var ( long = `Search and filter GPU instance types available on Brev. Filter instances by GPU name, provider, VRAM, total VRAM, GPU compute capability, disk size, and boot time. -Sort results by various columns to find the best instance for your needs.` +Sort results by various columns to find the best instance for your needs. + +Features column shows instance capabilities: + S = Stoppable (can stop and restart without losing data) + R = Rebootable (can reboot the instance) + P = Flex Ports (can modify firewall/port rules)` example = ` # List all GPU instances @@ -199,12 +208,27 @@ type GPUInstanceInfo struct { DiskMin float64 `json:"disk_min_gb"` DiskMax float64 `json:"disk_max_gb"` 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 } +// isStdoutPiped returns true if stdout is being piped (not a terminal) +func isStdoutPiped() bool { + stat, _ := os.Stdout.Stat() + return (stat.Mode() & os.ModeCharDevice) == 0 +} + // RunGPUSearch executes the GPU search with filters and sorting func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, maxBootTime int, sortBy string, descending, jsonOutput bool) error { + // Auto-switch to JSON when stdout is piped (for chaining with provision) + if !jsonOutput && isStdoutPiped() { + jsonOutput = true + } + response, err := store.GetInstanceTypes() if err != nil { return breverrors.WrapAndTrace(err) @@ -220,10 +244,10 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider } // Process and filter instances - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Apply filters - filtered := filterInstances(instances, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, maxBootTime) + filtered := FilterInstances(instances, gpuName, provider, minVRAM, minTotalVRAM, minCapability, minDisk, maxBootTime) if len(filtered) == 0 { if jsonOutput { @@ -234,8 +258,19 @@ func RunGPUSearch(t *terminal.Terminal, store GPUSearchStore, gpuName, provider return nil } + // Set target disk for each instance + // For flexible storage, use minDisk if specified and within range + for i := range filtered { + inst := &filtered[i] + if inst.DiskMin != inst.DiskMax && minDisk > 0 && minDisk >= inst.DiskMin && minDisk <= inst.DiskMax { + inst.TargetDisk = minDisk + } else { + inst.TargetDisk = inst.DiskMin + } + } + // Sort instances - sortInstances(filtered, sortBy, descending) + SortInstances(filtered, sortBy, descending) // Display results if jsonOutput { @@ -439,8 +474,8 @@ func getGPUCapability(gpuName string) float64 { return 0 } -// processInstances converts raw instance types to GPUInstanceInfo -func processInstances(items []InstanceType) []GPUInstanceInfo { +// ProcessInstances converts raw instance types to GPUInstanceInfo +func ProcessInstances(items []InstanceType) []GPUInstanceInfo { var instances []GPUInstanceInfo for _, item := range items { @@ -487,6 +522,9 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { DiskMin: diskMin, DiskMax: diskMax, BootTime: bootTime, + Stoppable: item.Stoppable, + Rebootable: item.Rebootable, + FlexPorts: item.CanModifyFirewallRules, PricePerHour: price, Manufacturer: gpu.Manufacturer, }) @@ -496,8 +534,8 @@ func processInstances(items []InstanceType) []GPUInstanceInfo { return instances } -// filterInstances applies all filters to the instance list -func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, maxBootTime int) []GPUInstanceInfo { +// FilterInstances applies all filters to the instance list +func FilterInstances(instances []GPUInstanceInfo, gpuName, provider string, minVRAM, minTotalVRAM, minCapability, minDisk float64, maxBootTime int) []GPUInstanceInfo { var filtered []GPUInstanceInfo for _, inst := range instances { @@ -537,7 +575,8 @@ func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minV } // Filter by maximum boot time (convert minutes to seconds for comparison) - if maxBootTime > 0 && inst.BootTime > maxBootTime*60 { + // Exclude instances with unknown boot time (0) when filter is specified + if maxBootTime > 0 && (inst.BootTime == 0 || inst.BootTime > maxBootTime*60) { continue } @@ -547,8 +586,8 @@ func filterInstances(instances []GPUInstanceInfo, gpuName, provider string, minV return filtered } -// sortInstances sorts the instance list by the specified column -func sortInstances(instances []GPUInstanceInfo, sortBy string, descending bool) { +// SortInstances sorts the instance list by the specified column +func SortInstances(instances []GPUInstanceInfo, sortBy string, descending bool) { sort.Slice(instances, func(i, j int) bool { var less bool switch strings.ToLower(sortBy) { @@ -635,13 +674,32 @@ func formatBootTime(seconds int) string { return fmt.Sprintf("%dm%ds", minutes, secs) } +// formatFeatures formats feature flags as abbreviated string +// S=stoppable, R=rebootable, P=flex ports (can modify firewall) +func formatFeatures(stoppable, rebootable, flexPorts bool) string { + var features []string + if stoppable { + features = append(features, "S") + } + if rebootable { + features = append(features, "R") + } + if flexPorts { + features = append(features, "P") + } + if len(features) == 0 { + return "-" + } + return strings.Join(features, "") +} + // 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", "BOOT", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -653,6 +711,7 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { } diskStr := formatDiskSize(inst.DiskMin, inst.DiskMax) bootStr := formatBootTime(inst.BootTime) + featuresStr := formatFeatures(inst.Stoppable, inst.Rebootable, inst.FlexPorts) priceStr := fmt.Sprintf("$%.2f", inst.PricePerHour) row := table.Row{ @@ -665,6 +724,7 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { capStr, diskStr, bootStr, + featuresStr, inst.VCPUs, priceStr, } diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index 1e56637e..76b7088d 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -140,7 +140,7 @@ func TestGetGPUCapability(t *testing.T) { func TestProcessInstances(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) assert.Len(t, instances, 7, "Expected 7 GPU instances") @@ -165,178 +165,178 @@ func TestProcessInstances(t *testing.T) { func TestFilterInstancesByGPUName(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Filter by A10G - filtered := filterInstances(instances, "A10G", "", 0, 0, 0, 0, 0) + filtered := FilterInstances(instances, "A10G", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances") // Filter by V100 - filtered = filterInstances(instances, "V100", "", 0, 0, 0, 0, 0) + filtered = FilterInstances(instances, "V100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances") // Filter by lowercase (case-insensitive) - filtered = filterInstances(instances, "v100", "", 0, 0, 0, 0, 0) + filtered = FilterInstances(instances, "v100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 V100 instances (case-insensitive)") // Filter by partial match - filtered = filterInstances(instances, "A1", "", 0, 0, 0, 0, 0) + filtered = FilterInstances(instances, "A1", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances matching 'A1' (A10G and A100)") } func TestFilterInstancesByMinVRAM(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Filter by min VRAM 24GB - filtered := filterInstances(instances, "", "", 24, 0, 0, 0, 0) + filtered := FilterInstances(instances, "", "", 24, 0, 0, 0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with >= 24GB VRAM") // Filter by min VRAM 40GB - filtered = filterInstances(instances, "", "", 40, 0, 0, 0, 0) + filtered = FilterInstances(instances, "", "", 40, 0, 0, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 40GB VRAM") assert.Equal(t, "A100", filtered[0].GPUName) } func TestFilterInstancesByMinTotalVRAM(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Filter by min total VRAM 60GB - filtered := filterInstances(instances, "", "", 0, 60, 0, 0, 0) + filtered := FilterInstances(instances, "", "", 0, 60, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 instances with >= 60GB total VRAM") // Filter by min total VRAM 300GB - filtered = filterInstances(instances, "", "", 0, 300, 0, 0, 0) + filtered = FilterInstances(instances, "", "", 0, 300, 0, 0, 0) assert.Len(t, filtered, 1, "Should have 1 instance with >= 300GB total VRAM") assert.Equal(t, "p4d.24xlarge", filtered[0].Type) } func TestFilterInstancesByMinCapability(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Filter by capability >= 8.0 - filtered := filterInstances(instances, "", "", 0, 0, 8.0, 0, 0) + filtered := FilterInstances(instances, "", "", 0, 0, 8.0, 0, 0) assert.Len(t, filtered, 4, "Should have 4 instances with capability >= 8.0") // Filter by capability >= 8.5 - filtered = filterInstances(instances, "", "", 0, 0, 8.5, 0, 0) + filtered = FilterInstances(instances, "", "", 0, 0, 8.5, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances with capability >= 8.5") } func TestFilterInstancesCombined(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Filter by GPU name and min VRAM - filtered := filterInstances(instances, "A10G", "", 24, 0, 0, 0, 0) + filtered := FilterInstances(instances, "A10G", "", 24, 0, 0, 0, 0) assert.Len(t, filtered, 2, "Should have 2 A10G instances with >= 24GB VRAM") // Filter by GPU name, min VRAM, and capability - filtered = filterInstances(instances, "", "", 24, 0, 8.5, 0, 0) + filtered = FilterInstances(instances, "", "", 24, 0, 8.5, 0, 0) assert.Len(t, filtered, 3, "Should have 3 instances with >= 24GB VRAM and capability >= 8.5") } func TestSortInstancesByPrice(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by price ascending - sortInstances(instances, "price", false) + SortInstances(instances, "price", false) assert.Equal(t, "g4dn.xlarge", instances[0].Type, "Cheapest should be g4dn.xlarge") assert.Equal(t, "p4d.24xlarge", instances[len(instances)-1].Type, "Most expensive should be p4d.24xlarge") // Sort by price descending - sortInstances(instances, "price", true) + SortInstances(instances, "price", true) assert.Equal(t, "p4d.24xlarge", instances[0].Type, "Most expensive should be first when descending") } func TestSortInstancesByGPUCount(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by GPU count ascending - sortInstances(instances, "gpu-count", false) + SortInstances(instances, "gpu-count", false) assert.Equal(t, 1, instances[0].GPUCount, "Instances with 1 GPU should be first") // Sort by GPU count descending - sortInstances(instances, "gpu-count", true) + SortInstances(instances, "gpu-count", true) assert.Equal(t, 8, instances[0].GPUCount, "Instance with 8 GPUs should be first when descending") } func TestSortInstancesByVRAM(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by VRAM ascending - sortInstances(instances, "vram", false) + SortInstances(instances, "vram", false) assert.Equal(t, 16.0, instances[0].VRAMPerGPU, "Instances with 16GB VRAM should be first") // Sort by VRAM descending - sortInstances(instances, "vram", true) + SortInstances(instances, "vram", true) assert.Equal(t, 40.0, instances[0].VRAMPerGPU, "Instance with 40GB VRAM should be first when descending") } func TestSortInstancesByTotalVRAM(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by total VRAM ascending - sortInstances(instances, "total-vram", false) + SortInstances(instances, "total-vram", false) assert.Equal(t, 16.0, instances[0].TotalVRAM, "Instances with 16GB total VRAM should be first") // Sort by total VRAM descending - sortInstances(instances, "total-vram", true) + SortInstances(instances, "total-vram", true) assert.Equal(t, 320.0, instances[0].TotalVRAM, "Instance with 320GB total VRAM should be first when descending") } func TestSortInstancesByVCPU(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by vCPU ascending - sortInstances(instances, "vcpu", false) + SortInstances(instances, "vcpu", false) assert.Equal(t, 4, instances[0].VCPUs, "Instances with 4 vCPUs should be first") // Sort by vCPU descending - sortInstances(instances, "vcpu", true) + SortInstances(instances, "vcpu", true) assert.Equal(t, 96, instances[0].VCPUs, "Instance with 96 vCPUs should be first when descending") } func TestSortInstancesByCapability(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by capability ascending - sortInstances(instances, "capability", false) + SortInstances(instances, "capability", false) assert.Equal(t, 7.0, instances[0].Capability, "Instances with capability 7.0 should be first") // Sort by capability descending - sortInstances(instances, "capability", true) + SortInstances(instances, "capability", true) assert.Equal(t, 8.9, instances[0].Capability, "Instance with capability 8.9 should be first when descending") } func TestSortInstancesByType(t *testing.T) { response := createTestInstanceTypes() - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) // Sort by type ascending - sortInstances(instances, "type", false) + SortInstances(instances, "type", false) assert.Equal(t, "g4dn.xlarge", instances[0].Type, "g4dn.xlarge should be first alphabetically") // Sort by type descending - sortInstances(instances, "type", true) + SortInstances(instances, "type", true) assert.Equal(t, "p4d.24xlarge", instances[0].Type, "p4d.24xlarge should be first when descending") } func TestEmptyInstanceTypes(t *testing.T) { response := &InstanceTypesResponse{Items: []InstanceType{}} - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) assert.Len(t, instances, 0, "Should have 0 instances") - filtered := filterInstances(instances, "A100", "", 0, 0, 0, 0, 0) + filtered := FilterInstances(instances, "A100", "", 0, 0, 0, 0, 0) assert.Len(t, filtered, 0, "Filtered should also be empty") } @@ -362,7 +362,7 @@ func TestNonGPUInstancesAreFiltered(t *testing.T) { }, } - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) assert.Len(t, instances, 1, "Should only have 1 GPU instance, non-GPU instances should be filtered") assert.Equal(t, "g5.xlarge", instances[0].Type) } @@ -382,7 +382,62 @@ func TestMemoryBytesAsFallback(t *testing.T) { }, } - instances := processInstances(response.Items) + instances := ProcessInstances(response.Items) assert.Len(t, instances, 1) assert.Equal(t, 24.0, instances[0].VRAMPerGPU, "Should fall back to MemoryBytes when Memory string is empty") } + +func TestFilterByMaxBootTimeExcludesUnknown(t *testing.T) { + response := &InstanceTypesResponse{ + Items: []InstanceType{ + { + Type: "fast-boot", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A100", Manufacturer: "NVIDIA", Memory: "40GiB"}, + }, + Memory: "64GiB", + VCPU: 8, + BasePrice: BasePrice{Currency: "USD", Amount: "3.00"}, + EstimatedDeployTime: "5m0s", + }, + { + Type: "slow-boot", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A100", Manufacturer: "NVIDIA", Memory: "40GiB"}, + }, + Memory: "64GiB", + VCPU: 8, + BasePrice: BasePrice{Currency: "USD", Amount: "2.50"}, + EstimatedDeployTime: "15m0s", + }, + { + Type: "unknown-boot", + SupportedGPUs: []GPU{ + {Count: 1, Name: "A100", Manufacturer: "NVIDIA", Memory: "40GiB"}, + }, + Memory: "64GiB", + VCPU: 8, + BasePrice: BasePrice{Currency: "USD", Amount: "2.00"}, + EstimatedDeployTime: "", // Unknown boot time + }, + }, + } + + instances := ProcessInstances(response.Items) + assert.Len(t, instances, 3, "Should have 3 instances before filtering") + + // Filter by max boot time of 10 minutes - should exclude unknown and slow-boot + filtered := FilterInstances(instances, "", "", 0, 0, 0, 0, 10) + assert.Len(t, filtered, 1, "Should have 1 instance with boot time <= 10 minutes") + assert.Equal(t, "fast-boot", filtered[0].Type, "Only fast-boot should match") + + // Verify unknown boot time instance is excluded + for _, inst := range filtered { + assert.NotEqual(t, "unknown-boot", inst.Type, "Unknown boot time should be excluded") + assert.NotEqual(t, 0, inst.BootTime, "Instances with 0 boot time should be excluded") + } + + // Without filter, all instances should be included + noFilter := FilterInstances(instances, "", "", 0, 0, 0, 0, 0) + assert.Len(t, noFilter, 3, "Without filter, all 3 instances should be included") +} diff --git a/pkg/cmd/open/open.go b/pkg/cmd/open/open.go index f0691ece..85c0f7d7 100644 --- a/pkg/cmd/open/open.go +++ b/pkg/cmd/open/open.go @@ -1,6 +1,7 @@ package open import ( + "bufio" "errors" "fmt" "os" @@ -38,7 +39,25 @@ const ( var ( openLong = "[command in beta] This will open VS Code, Cursor, Windsurf, or tmux SSH-ed in to your instance. You must have the editor installed in your path." - openExample = "brev open instance_id_or_name\nbrev open instance\nbrev open instance code\nbrev open instance cursor\nbrev open instance windsurf\nbrev open instance tmux\nbrev open --set-default cursor\nbrev open --set-default windsurf\nbrev open --set-default tmux" + openExample = ` # Open an instance by name or ID + brev open instance_id_or_name + brev open my-instance + + # Open with a specific editor + brev open my-instance code + brev open my-instance cursor + brev open my-instance windsurf + brev open my-instance tmux + + # Set a default editor + brev open --set-default cursor + brev open --set-default windsurf + + # Create a GPU instance and open it immediately (reads instance name from stdin) + brev provision --name my-instance | brev open + + # Create with specific GPU and open in Cursor + brev gpus --gpu-name A100 | brev gpu-create --name ml-box | brev open cursor` ) type OpenStore interface { @@ -72,7 +91,8 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto if setDefaultFlag != "" { return cobra.NoArgs(cmd, args) } - return cobra.RangeArgs(1, 2)(cmd, args) + // Allow 0-2 args: instance name can come from stdin + return cobra.RangeArgs(0, 2)(cmd, args) }), ValidArgsFunction: completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t), RunE: func(cmd *cobra.Command, args []string) error { @@ -80,17 +100,23 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto return handleSetDefault(t, setDefault) } + // Get instance name from args or stdin + instanceName, remainingArgs, err := getInstanceNameOpen(args) + if err != nil { + return breverrors.WrapAndTrace(err) + } + setupDoneString := "------ Git repo cloned ------" if waitForSetupToFinish { setupDoneString = "------ Done running execs ------" } - editorType, err := determineEditorType(args) + editorType, err := determineEditorType(remainingArgs) if err != nil { return breverrors.WrapAndTrace(err) } - err = runOpenCommand(t, store, args[0], setupDoneString, directory, host, editorType) + err = runOpenCommand(t, store, instanceName, setupDoneString, directory, host, editorType) if err != nil { return breverrors.WrapAndTrace(err) } @@ -105,6 +131,32 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto return cmd } +// getInstanceNameOpen gets the instance name from args or stdin, returning remaining args for editor type +func getInstanceNameOpen(args []string) (string, []string, error) { + // Check if stdin is piped + stat, _ := os.Stdin.Stat() + stdinPiped := (stat.Mode() & os.ModeCharDevice) == 0 + + if stdinPiped { + // Read instance name from stdin + scanner := bufio.NewScanner(os.Stdin) + if scanner.Scan() { + name := strings.TrimSpace(scanner.Text()) + if name != "" { + // All args are for editor type + return name, args, nil + } + } + } + + // Instance name from args + if len(args) > 0 { + return args[0], args[1:], nil + } + + return "", nil, breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") +} + func handleSetDefault(t *terminal.Terminal, editorType string) error { if editorType != EditorVSCode && editorType != EditorCursor && editorType != EditorWindsurf && editorType != EditorTmux { return fmt.Errorf("invalid editor type: %s. Must be 'code', 'cursor', 'windsurf', or 'tmux'", editorType) @@ -129,8 +181,9 @@ func handleSetDefault(t *terminal.Terminal, editorType string) error { } func determineEditorType(args []string) (string, error) { - if len(args) == 2 { - editorType := args[1] + // args now contains only the editor type (if provided), not the instance name + if len(args) >= 1 { + editorType := args[0] if editorType != EditorVSCode && editorType != EditorCursor && editorType != EditorWindsurf && editorType != EditorTmux { return "", fmt.Errorf("invalid editor type: %s. Must be 'code', 'cursor', 'windsurf', or 'tmux'", editorType) } diff --git a/pkg/cmd/shell/shell.go b/pkg/cmd/shell/shell.go index 49474ad3..ce1fef60 100644 --- a/pkg/cmd/shell/shell.go +++ b/pkg/cmd/shell/shell.go @@ -1,6 +1,7 @@ package shell import ( + "bufio" "errors" "fmt" "os" @@ -26,7 +27,28 @@ import ( var ( openLong = "[command in beta] This will shell in to your instance" - openExample = "brev shell instance_id_or_name\nbrev shell instance\nbrev open h9fp5vxwe" + openExample = ` # SSH into an instance by name or ID + brev shell instance_id_or_name + brev shell my-instance + + # Run a command on the instance (non-interactive, pipes stdout/stderr) + brev shell my-instance -c "nvidia-smi" + brev shell my-instance -c "python train.py" + + # Run a script file on the instance + brev shell my-instance -c @setup.sh + + # Chain: provision and run a command (reads instance name from stdin) + brev provision --name my-instance | brev shell -c "nvidia-smi" + + # Create a GPU instance and SSH into it immediately (use command substitution for interactive shell) + brev shell $(brev provision --name my-instance) + + # Create with specific GPU and connect + brev shell $(brev gpus --gpu-name A100 | brev gpu-create --name ml-box) + + # SSH into the host machine instead of the container + brev shell my-instance --host` ) type ShellStore interface { @@ -40,9 +62,8 @@ type ShellStore interface { } func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore ShellStore) *cobra.Command { - var runRemoteCMD bool - var directory string var host bool + var command string cmd := &cobra.Command{ Annotations: map[string]string{"access": ""}, Use: "shell", @@ -51,10 +72,22 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell Short: "[beta] Open a shell in your instance", Long: openLong, Example: openExample, - Args: cmderrors.TransformToValidationError(cmderrors.TransformToValidationError(cobra.ExactArgs(1))), + Args: cmderrors.TransformToValidationError(cobra.RangeArgs(0, 1)), ValidArgsFunction: completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t), RunE: func(cmd *cobra.Command, args []string) error { - err := runShellCommand(t, store, args[0], directory, host) + // Get instance name from args or stdin + instanceName, err := getInstanceName(args) + if err != nil { + return breverrors.WrapAndTrace(err) + } + + // Parse command (can be inline or @filepath) + cmdToRun, err := parseCommand(command) + if err != nil { + return breverrors.WrapAndTrace(err) + } + + err = runShellCommand(t, store, instanceName, host, cmdToRun) if err != nil { return breverrors.WrapAndTrace(err) } @@ -62,13 +95,53 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell }, } cmd.Flags().BoolVarP(&host, "host", "", false, "ssh into the host machine instead of the container") - cmd.Flags().BoolVarP(&runRemoteCMD, "remote", "r", true, "run remote commands") - cmd.Flags().StringVarP(&directory, "dir", "d", "", "override directory to launch shell") + cmd.Flags().StringVarP(&command, "command", "c", "", "command to run on the instance (use @filename to run a script file)") return cmd } -func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID, directory string, host bool) error { +// getInstanceName gets the instance name from args or stdin +func getInstanceName(args []string) (string, error) { + if len(args) > 0 { + return args[0], nil + } + + // Check if stdin is piped + stat, _ := os.Stdin.Stat() + if (stat.Mode() & os.ModeCharDevice) == 0 { + // Stdin is piped, read instance name + scanner := bufio.NewScanner(os.Stdin) + if scanner.Scan() { + name := strings.TrimSpace(scanner.Text()) + if name != "" { + return name, nil + } + } + } + + return "", breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") +} + +// parseCommand parses the command string, loading from file if prefixed with @ +func parseCommand(command string) (string, error) { + if command == "" { + return "", nil + } + + // If prefixed with @, read from file + if strings.HasPrefix(command, "@") { + filePath := strings.TrimPrefix(command, "@") + content, err := os.ReadFile(filePath) + if err != nil { + return "", breverrors.WrapAndTrace(err) + } + return string(content), nil + } + + return command, nil +} + +func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID string, host bool, command string) error { s := t.NewSpinner() workspace, err := util.GetUserWorkspaceByNameOrIDErr(sstore, workspaceNameOrID) if err != nil { @@ -114,7 +187,7 @@ func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID, // legacy environments wont support this and cause errrors, // but we don't want to block the user from using the shell _ = writeconnectionevent.WriteWCEOnEnv(sstore, workspace.DNS) - err = runSSH(workspace, sshName, directory) + err = runSSH(workspace, sshName, command) if err != nil { return breverrors.WrapAndTrace(err) } @@ -162,14 +235,29 @@ func waitForSSHToBeAvailable(sshAlias string, s *spinner.Spinner) error { } } -func runSSH(_ *entity.Workspace, sshAlias, _ string) error { +func runSSH(_ *entity.Workspace, sshAlias string, command string) error { sshAgentEval := "eval $(ssh-agent -s)" - cmd := fmt.Sprintf("ssh %s", sshAlias) + + var cmd string + if command != "" { + // Non-interactive: run command and pipe stdout/stderr + // Escape the command for passing to SSH + escapedCmd := strings.ReplaceAll(command, "'", "'\\''") + cmd = fmt.Sprintf("ssh %s '%s'", sshAlias, escapedCmd) + } else { + // Interactive shell + cmd = fmt.Sprintf("ssh %s", sshAlias) + } + cmd = fmt.Sprintf("%s && %s", sshAgentEval, cmd) sshCmd := exec.Command("bash", "-c", cmd) //nolint:gosec //cmd is user input sshCmd.Stderr = os.Stderr sshCmd.Stdout = os.Stdout - sshCmd.Stdin = os.Stdin + + // Only attach stdin for interactive sessions + if command == "" { + sshCmd.Stdin = os.Stdin + } err := hello.SetHasRunShell(true) if err != nil { diff --git a/pkg/store/workspace.go b/pkg/store/workspace.go index 5190d313..5d119f7b 100644 --- a/pkg/store/workspace.go +++ b/pkg/store/workspace.go @@ -34,6 +34,17 @@ type ModifyWorkspaceRequest struct { InstanceType string `json:"instanceType,omitempty"` } +// LifeCycleScriptAttr holds the lifecycle script configuration +type LifeCycleScriptAttr struct { + Script string `json:"script,omitempty"` +} + +// VMBuild holds VM-specific build configuration +type VMBuild struct { + ForceJupyterInstall bool `json:"forceJupyterInstall,omitempty"` + LifeCycleScriptAttr *LifeCycleScriptAttr `json:"lifeCycleScriptAttr,omitempty"` +} + type CreateWorkspacesOptions struct { Name string `json:"name"` WorkspaceGroupID string `json:"workspaceGroupId"` @@ -57,6 +68,7 @@ type CreateWorkspacesOptions struct { DiskStorage string `json:"diskStorage"` BaseImage string `json:"baseImage"` VMOnlyMode bool `json:"vmOnlyMode"` + VMBuild *VMBuild `json:"vmBuild,omitempty"` PortMappings map[string]string `json:"portMappings"` Files interface{} `json:"files"` Labels interface{} `json:"labels"` @@ -88,6 +100,7 @@ var ( var DefaultApplicationList = []entity.Application{DefaultApplication} func NewCreateWorkspacesOptions(clusterID, name string) *CreateWorkspacesOptions { + isStoppable := false return &CreateWorkspacesOptions{ BaseImage: "", Description: "", @@ -95,12 +108,12 @@ func NewCreateWorkspacesOptions(clusterID, name string) *CreateWorkspacesOptions ExecsV1: &entity.ExecsV1{}, Files: nil, InstanceType: "", - IsStoppable: nil, + IsStoppable: &isStoppable, Labels: nil, LaunchJupyterOnStart: false, Name: name, - PortMappings: nil, - ReposV1: nil, + PortMappings: map[string]string{}, + ReposV1: &entity.ReposV1{}, VMOnlyMode: true, WorkspaceGroupID: "GCP", WorkspaceTemplateID: DefaultWorkspaceTemplateID, From 2877337ddea2c604a9f9a74003474c433e4f5721 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 01:18:37 -0800 Subject: [PATCH 09/13] Rename commands and add disk pricing to search Command renaming: - Rename 'gpu-search' to 'search' with 'gpu-search' as alias - Rename 'gpu-create' to 'create' with 'gpu-create' as alias - Support positional argument for instance name: brev create my-instance - Update all examples to use new command names - Keep backwards compatibility with old command names as aliases Search enhancements: - Add $/GB/MO column showing monthly disk storage pricing - Extract disk price from API's price_per_gb_hr field (converted to monthly) - Add cloud field to distinguish underlying cloud from provider/aggregator - Display cloud:provider format (e.g., hyperstack:shadeform) when different - Auto-switch to JSON output when stdout is piped (for provision chaining) - Add target_disk_gb to JSON for passing --min-disk through pipeline - Add tests for extractCloud, extractDiskInfo, and cloud extraction --- pkg/cmd/gpucreate/gpucreate.go | 52 ++++++---- pkg/cmd/gpusearch/gpusearch.go | 149 ++++++++++++++++++---------- pkg/cmd/gpusearch/gpusearch_test.go | 107 ++++++++++++++++++++ pkg/cmd/open/open.go | 7 +- pkg/cmd/shell/shell.go | 10 +- 5 files changed, 246 insertions(+), 79 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index 429477e1..f10f2218 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -28,7 +28,7 @@ var ( This command attempts to create GPU instances, trying different instance types until the desired number of instances are successfully created. Instance types -can be specified directly, piped from 'brev gpus', or auto-selected using defaults. +can be specified directly, piped from 'brev search', or auto-selected using defaults. Default Behavior: If no instance types are specified (no --type flag and no piped input), the command @@ -54,37 +54,40 @@ You can attach a startup script that runs when the instance boots using the example = ` # Quick start: create an instance using smart defaults (sorted by price) - brev provision --name my-instance + brev create my-instance - # Create and immediately open in VS Code (reads instance name from stdin) - brev provision --name my-instance | brev open + # Create with explicit --name flag + brev create --name my-instance - # Create and SSH into the instance (use command substitution for interactive shell) - brev shell $(brev provision --name my-instance) + # Create and immediately open in VS Code + brev create my-instance | brev open - # Create and run a command (non-interactive, reads instance name from stdin) - brev provision --name my-instance | brev shell -c "nvidia-smi" + # Create and SSH into the instance + brev shell $(brev create my-instance) - # Create a single instance with a specific GPU type - brev gpu-create --name my-instance --type g5.xlarge + # Create and run a command + brev create my-instance | brev shell -c "nvidia-smi" - # Pipe instance types from brev gpus (tries each type until one succeeds) - brev gpus --min-vram 24 | brev gpu-create --name my-instance + # Create with a specific GPU type + brev create my-instance --type g5.xlarge + + # Pipe instance types from brev search (tries each type until one succeeds) + brev search --min-vram 24 | brev create my-instance # Create 3 instances, trying types in parallel - brev gpus --gpu-name A100 | brev gpu-create --name my-cluster --count 3 --parallel 5 + brev search --gpu-name A100 | brev create my-cluster --count 3 --parallel 5 # Try multiple specific types in order - brev gpu-create --name my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge + brev create my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge # Attach a startup script from a file - brev gpu-create --name my-instance --type g5.xlarge --startup-script @setup.sh + brev create my-instance --type g5.xlarge --startup-script @setup.sh # Attach an inline startup script - brev gpu-create --name my-instance --type g5.xlarge --startup-script 'pip install torch transformers' + brev create my-instance --startup-script 'pip install torch' # Combine: find cheapest A100, attach setup script - brev gpus --gpu-name A100 --sort price | brev gpu-create --name ml-box --startup-script @ml-setup.sh + brev search --gpu-name A100 --sort price | brev create ml-box -s @ml-setup.sh ` ) @@ -127,13 +130,18 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra cmd := &cobra.Command{ Annotations: map[string]string{"workspace": ""}, - Use: "gpu-create", - Aliases: []string{"gpu-retry", "gcreate", "provision"}, + Use: "create [name]", + Aliases: []string{"provision", "gpu-create", "gpu-retry", "gcreate"}, DisableFlagsInUseLine: true, Short: "Create GPU instances with automatic retry", Long: long, Example: example, RunE: func(cmd *cobra.Command, args []string) error { + // Accept name as positional arg or --name flag + if len(args) > 0 && name == "" { + name = args[0] + } + // Check if output is being piped (for chaining with brev shell/open) piped := isStdoutPiped() @@ -159,12 +167,12 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra } if len(types) == 0 { - return breverrors.NewValidationError("no GPU instances match the default filters. Try 'brev gpus' to see available options") + return breverrors.NewValidationError("no GPU instances match the default filters. Try 'brev search' to see available options") } } if name == "" { - return breverrors.NewValidationError("--name flag is required") + return breverrors.NewValidationError("name is required (as argument or --name flag)") } if count < 1 { @@ -199,7 +207,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra }, } - cmd.Flags().StringVarP(&name, "name", "n", "", "Base name for the instances (required)") + cmd.Flags().StringVarP(&name, "name", "n", "", "Base name for the instances (or pass as first argument)") cmd.Flags().StringVarP(&instanceTypes, "type", "t", "", "Comma-separated list of instance types to try") cmd.Flags().IntVarP(&count, "count", "c", 1, "Number of instances to create") cmd.Flags().IntVarP(¶llel, "parallel", "p", 1, "Number of parallel creation attempts") diff --git a/pkg/cmd/gpusearch/gpusearch.go b/pkg/cmd/gpusearch/gpusearch.go index b2d22471..59cfa3b6 100644 --- a/pkg/cmd/gpusearch/gpusearch.go +++ b/pkg/cmd/gpusearch/gpusearch.go @@ -45,7 +45,7 @@ type Storage struct { MinSize string `json:"min_size"` MaxSize string `json:"max_size"` SizeBytes MemoryBytes `json:"size_bytes"` - PricePerGBHr float64 `json:"price_per_gb_hr"` + PricePerGBHr BasePrice `json:"price_per_gb_hr"` // Uses BasePrice since API returns {currency, amount} } // WorkspaceGroup represents a workspace group that can run an instance type @@ -114,43 +114,43 @@ Features column shows instance capabilities: example = ` # List all GPU instances - brev gpu-search + brev search # Filter by GPU name (case-insensitive, partial match) - brev gpu-search --gpu-name A100 - brev gpu-search --gpu-name "L40S" + brev search --gpu-name A100 + brev search --gpu-name "L40S" # Filter by provider/cloud (case-insensitive, partial match) - brev gpu-search --provider aws - brev gpu-search --provider gcp + brev search --provider aws + brev search --provider gcp # Filter by minimum VRAM per GPU (in GB) - brev gpu-search --min-vram 24 + brev search --min-vram 24 # Filter by minimum total VRAM (in GB) - brev gpu-search --min-total-vram 80 + brev search --min-total-vram 80 # Filter by minimum GPU compute capability - brev gpu-search --min-capability 8.0 + brev search --min-capability 8.0 # Filter by minimum disk size (in GB) - brev gpu-search --min-disk 500 + brev search --min-disk 500 # Filter by maximum boot time (in minutes) - brev gpu-search --max-boot-time 5 + brev search --max-boot-time 5 # Sort by different columns (price, gpu-count, vram, total-vram, vcpu, provider, disk, boot-time) - brev gpu-search --sort price - brev gpu-search --sort boot-time - brev gpu-search --sort disk --desc + brev search --sort price + brev search --sort boot-time + brev search --sort disk --desc # Combine filters - brev gpu-search --gpu-name A100 --min-vram 40 --sort price - brev gpu-search --gpu-name H100 --max-boot-time 3 --sort price + brev search --gpu-name A100 --min-vram 40 --sort price + brev search --gpu-name H100 --max-boot-time 3 --sort price ` ) -// NewCmdGPUSearch creates the gpu-search command +// NewCmdGPUSearch creates the search command func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command { var gpuName string var provider string @@ -165,8 +165,8 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command cmd := &cobra.Command{ Annotations: map[string]string{"workspace": ""}, - Use: "gpu-search", - Aliases: []string{"gpu", "gpus", "gpu-list"}, + Use: "search", + Aliases: []string{"gpu-search", "gpu", "gpus", "gpu-list"}, DisableFlagsInUseLine: true, Short: "Search and filter GPU instance types", Long: long, @@ -197,7 +197,8 @@ func NewCmdGPUSearch(t *terminal.Terminal, store GPUSearchStore) *cobra.Command // GPUInstanceInfo holds processed GPU instance information for display type GPUInstanceInfo struct { Type string `json:"type"` - Provider string `json:"provider"` + 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"` @@ -205,9 +206,10 @@ type GPUInstanceInfo struct { Capability float64 `json:"capability"` VCPUs int `json:"vcpus"` Memory string `json:"memory"` - DiskMin float64 `json:"disk_min_gb"` - DiskMax float64 `json:"disk_max_gb"` - BootTime int `json:"boot_time_seconds"` + 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"` @@ -356,21 +358,30 @@ func parseDurationToSeconds(duration string) int { return totalSeconds } -// extractDiskSize extracts min and max disk size from storage configuration -// Returns (minGB, maxGB). For fixed-size storage, both values are the same. -func extractDiskSize(storage []Storage) (float64, float64) { +// extractDiskInfo extracts min/max disk size and price from storage configuration +// Returns (minGB, maxGB, pricePerGBMonth). For fixed-size storage, min==max and price is 0. +func extractDiskInfo(storage []Storage) (float64, float64, float64) { if len(storage) == 0 { - return 0, 0 + return 0, 0, 0 } // Use the first storage entry s := storage[0] + // Convert price per GB per hour to per month (730 hours average) + var pricePerGBMonth float64 + if s.PricePerGBHr.Amount != "" { + pricePerHr, _ := strconv.ParseFloat(s.PricePerGBHr.Amount, 64) + if pricePerHr > 0 { + pricePerGBMonth = pricePerHr * 730 + } + } + // Check if it's flexible storage (has min_size and max_size) if s.MinSize != "" && s.MaxSize != "" { minGB := parseSizeToGB(s.MinSize) maxGB := parseSizeToGB(s.MaxSize) - return minGB, maxGB + return minGB, maxGB, pricePerGBMonth } // Fixed storage - use size or size_bytes @@ -390,7 +401,32 @@ func extractDiskSize(storage []Storage) (float64, float64) { } } - return sizeGB, sizeGB + // Fixed storage doesn't have separate pricing - it's included in base price + return sizeGB, sizeGB, 0 +} + +// extractCloud extracts the underlying cloud from the instance type and provider +// For aggregators like shadeform, the cloud is in the type name prefix (e.g., "hyperstack_H100" -> "hyperstack") +// For direct providers like aws/gcp, the provider IS the cloud +func extractCloud(instanceType, provider string) string { + // Direct cloud providers - provider is the cloud + directProviders := map[string]bool{ + "aws": true, "gcp": true, "azure": true, "oci": true, + "nebius": true, "crusoe": true, "lambda-labs": true, "launchpad": true, + } + + if directProviders[strings.ToLower(provider)] { + return provider + } + + // For aggregators, try to extract cloud from type name prefix + // Pattern: cloudname_GPUtype (e.g., "hyperstack_H100", "cudo_A40", "latitude_H100") + if idx := strings.Index(instanceType, "_"); idx > 0 { + return instanceType[:idx] + } + + // Fallback to provider + return provider } // gpuCapabilityEntry represents a GPU pattern and its compute capability @@ -483,8 +519,8 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { continue // Skip non-GPU instances } - // Extract disk size info from first storage entry - diskMin, diskMax := extractDiskSize(item.SupportedStorage) + // Extract disk size and price info from first storage entry + diskMin, diskMax, diskPricePerMo := extractDiskInfo(item.SupportedStorage) // Extract boot time bootTime := parseDurationToSeconds(item.EstimatedDeployTime) @@ -510,23 +546,25 @@ func ProcessInstances(items []InstanceType) []GPUInstanceInfo { } instances = append(instances, GPUInstanceInfo{ - Type: item.Type, - Provider: item.Provider, - GPUName: gpu.Name, - GPUCount: gpu.Count, - VRAMPerGPU: vramPerGPU, - TotalVRAM: totalVRAM, - Capability: capability, - VCPUs: item.VCPU, - Memory: item.Memory, - DiskMin: diskMin, - DiskMax: diskMax, - 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, + DiskMin: diskMin, + DiskMax: diskMax, + DiskPricePerMo: diskPricePerMo, + BootTime: bootTime, + Stoppable: item.Stoppable, + Rebootable: item.Rebootable, + FlexPorts: item.CanModifyFirewallRules, + PricePerHour: price, + Manufacturer: gpu.Manufacturer, }) } } @@ -699,7 +737,7 @@ func displayGPUTable(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", "DISK", "BOOT", "FEATURES", "VCPUs", "$/HR"} + header := table.Row{"TYPE", "PROVIDER", "GPU", "COUNT", "VRAM/GPU", "TOTAL VRAM", "CAPABILITY", "DISK", "$/GB/MO", "BOOT", "FEATURES", "VCPUs", "$/HR"} ta.AppendHeader(header) for _, inst := range instances { @@ -710,19 +748,30 @@ func displayGPUTable(t *terminal.Terminal, instances []GPUInstanceInfo) { capStr = fmt.Sprintf("%.1f", inst.Capability) } diskStr := formatDiskSize(inst.DiskMin, inst.DiskMax) + diskPriceStr := "-" + if inst.DiskPricePerMo > 0 { + diskPriceStr = fmt.Sprintf("$%.2f", inst.DiskPricePerMo) + } bootStr := formatBootTime(inst.BootTime) featuresStr := formatFeatures(inst.Stoppable, inst.Rebootable, inst.FlexPorts) priceStr := fmt.Sprintf("$%.2f", inst.PricePerHour) + // Format cloud:provider - only show both if different + providerStr := inst.Provider + if inst.Cloud != "" && inst.Cloud != inst.Provider { + providerStr = fmt.Sprintf("%s:%s", inst.Cloud, inst.Provider) + } + row := table.Row{ inst.Type, - inst.Provider, + providerStr, t.Green(inst.GPUName), inst.GPUCount, vramStr, totalVramStr, capStr, diskStr, + diskPriceStr, bootStr, featuresStr, inst.VCPUs, diff --git a/pkg/cmd/gpusearch/gpusearch_test.go b/pkg/cmd/gpusearch/gpusearch_test.go index 76b7088d..93110c9e 100644 --- a/pkg/cmd/gpusearch/gpusearch_test.go +++ b/pkg/cmd/gpusearch/gpusearch_test.go @@ -441,3 +441,110 @@ func TestFilterByMaxBootTimeExcludesUnknown(t *testing.T) { noFilter := FilterInstances(instances, "", "", 0, 0, 0, 0, 0) assert.Len(t, noFilter, 3, "Without filter, all 3 instances should be included") } + +func TestExtractCloud(t *testing.T) { + tests := []struct { + name string + instanceType string + provider string + expectedCloud string + }{ + // Direct providers - cloud equals provider + {"AWS direct", "g5.xlarge", "aws", "aws"}, + {"GCP direct", "n1-highmem-8:nvidia-tesla-v100:8", "gcp", "gcp"}, + {"Nebius direct", "gpu-h100-sxm.1gpu-16vcpu-200gb", "nebius", "nebius"}, + {"OCI direct", "oci.h100x8.sxm", "oci", "oci"}, + {"Lambda Labs direct", "gpu_1x_h100_sxm5", "lambda-labs", "lambda-labs"}, + {"Crusoe direct", "l40s-48gb.1x", "crusoe", "crusoe"}, + {"Launchpad direct", "dmz.h100x2.pcie", "launchpad", "launchpad"}, + + // Aggregators - extract cloud from type name prefix + {"Shadeform hyperstack", "hyperstack_H100", "shadeform", "hyperstack"}, + {"Shadeform latitude", "latitude_H100x4", "shadeform", "latitude"}, + {"Shadeform cudo", "cudo_A40", "shadeform", "cudo"}, + {"Shadeform horizon", "horizon_H100x8", "shadeform", "horizon"}, + {"Shadeform paperspace", "paperspace_H100", "shadeform", "paperspace"}, + + // Edge cases + {"Unknown aggregator no underscore", "someinstance", "unknown-agg", "unknown-agg"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := extractCloud(tt.instanceType, tt.provider) + assert.Equal(t, tt.expectedCloud, result) + }) + } +} + +func TestExtractDiskInfoWithPricing(t *testing.T) { + // Test flexible storage with pricing + storageWithPrice := []Storage{ + { + MinSize: "50GiB", + MaxSize: "2560GiB", + PricePerGBHr: BasePrice{Currency: "USD", Amount: "0.00014"}, + }, + } + + minGB, maxGB, pricePerMo := extractDiskInfo(storageWithPrice) + assert.Equal(t, 50.0, minGB) + assert.Equal(t, 2560.0, maxGB) + assert.InDelta(t, 0.1022, pricePerMo, 0.001, "Price should be ~$0.10/GB/mo (0.00014 * 730)") + + // Test fixed storage (no pricing) + fixedStorage := []Storage{ + { + Size: "500GiB", + }, + } + + minGB, maxGB, pricePerMo = extractDiskInfo(fixedStorage) + assert.Equal(t, 500.0, minGB) + assert.Equal(t, 500.0, maxGB) + assert.Equal(t, 0.0, pricePerMo, "Fixed storage should have no separate price") + + // Test empty storage + minGB, maxGB, pricePerMo = extractDiskInfo([]Storage{}) + assert.Equal(t, 0.0, minGB) + assert.Equal(t, 0.0, maxGB) + assert.Equal(t, 0.0, pricePerMo) +} + +func TestProcessInstancesCloudExtraction(t *testing.T) { + response := &InstanceTypesResponse{ + Items: []InstanceType{ + { + Type: "hyperstack_H100", + Provider: "shadeform", + SupportedGPUs: []GPU{ + {Count: 1, Name: "H100", Manufacturer: "NVIDIA", Memory: "80GiB"}, + }, + Memory: "180GiB", + VCPU: 28, + BasePrice: BasePrice{Currency: "USD", Amount: "2.28"}, + }, + { + Type: "gpu-h100-sxm.1gpu-16vcpu-200gb", + Provider: "nebius", + SupportedGPUs: []GPU{ + {Count: 1, Name: "H100", Manufacturer: "NVIDIA", Memory: "80GiB"}, + }, + Memory: "200GiB", + VCPU: 16, + BasePrice: BasePrice{Currency: "USD", Amount: "3.54"}, + }, + }, + } + + instances := ProcessInstances(response.Items) + assert.Len(t, instances, 2) + + // Shadeform instance should have cloud extracted from type name + assert.Equal(t, "hyperstack", instances[0].Cloud) + assert.Equal(t, "shadeform", instances[0].Provider) + + // Nebius instance should have cloud = provider + assert.Equal(t, "nebius", instances[1].Cloud) + assert.Equal(t, "nebius", instances[1].Provider) +} diff --git a/pkg/cmd/open/open.go b/pkg/cmd/open/open.go index 85c0f7d7..d05026c7 100644 --- a/pkg/cmd/open/open.go +++ b/pkg/cmd/open/open.go @@ -54,10 +54,13 @@ var ( brev open --set-default windsurf # Create a GPU instance and open it immediately (reads instance name from stdin) - brev provision --name my-instance | brev open + brev create my-instance | brev open # Create with specific GPU and open in Cursor - brev gpus --gpu-name A100 | brev gpu-create --name ml-box | brev open cursor` + brev search --gpu-name A100 | brev create ml-box | brev open cursor + + # Note: tmux requires command substitution (not pipes) since it needs interactive stdin + brev open $(brev create my-instance) tmux` ) type OpenStore interface { diff --git a/pkg/cmd/shell/shell.go b/pkg/cmd/shell/shell.go index ce1fef60..7d172217 100644 --- a/pkg/cmd/shell/shell.go +++ b/pkg/cmd/shell/shell.go @@ -38,14 +38,14 @@ var ( # Run a script file on the instance brev shell my-instance -c @setup.sh - # Chain: provision and run a command (reads instance name from stdin) - brev provision --name my-instance | brev shell -c "nvidia-smi" + # Chain: create and run a command (reads instance name from stdin) + brev create my-instance | brev shell -c "nvidia-smi" - # Create a GPU instance and SSH into it immediately (use command substitution for interactive shell) - brev shell $(brev provision --name my-instance) + # Create a GPU instance and SSH into it (use command substitution for interactive shell) + brev shell $(brev create my-instance) # Create with specific GPU and connect - brev shell $(brev gpus --gpu-name A100 | brev gpu-create --name ml-box) + brev shell $(brev search --gpu-name A100 | brev create ml-box) # SSH into the host machine instead of the container brev shell my-instance --host` From 68cb68be4cc073d52213905cb8a6b378146ab7aa Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 01:31:40 -0800 Subject: [PATCH 10/13] Add multi-instance support to brev shell and brev open Both commands now support multiple instance names: brev shell: - Run commands on multiple instances: brev shell i1 i2 i3 -c "nvidia-smi" - Read multiple instances from stdin (one per line) - Interactive shell still only supports one instance brev open: - Open multiple instances in separate editor windows: brev open i1 i2 i3 - Read multiple instances from stdin (one per line) - tmux still only supports one instance (requires interactive stdin) Enables workflows like: brev create my-cluster --count 3 | brev shell -c "nvidia-smi" brev create my-cluster --count 3 | brev open --- pkg/cmd/open/open.go | 137 ++++++++++++++++++++++++++--------------- pkg/cmd/shell/shell.go | 66 ++++++++++++++------ 2 files changed, 137 insertions(+), 66 deletions(-) diff --git a/pkg/cmd/open/open.go b/pkg/cmd/open/open.go index d05026c7..7a2edf16 100644 --- a/pkg/cmd/open/open.go +++ b/pkg/cmd/open/open.go @@ -43,12 +43,22 @@ var ( brev open instance_id_or_name brev open my-instance + # Open multiple instances (each in separate editor window) + brev open instance1 instance2 instance3 + # Open with a specific editor brev open my-instance code brev open my-instance cursor brev open my-instance windsurf brev open my-instance tmux + # Open multiple instances with specific editor (flag is explicit) + brev open instance1 instance2 --editor cursor + brev open instance1 instance2 -e cursor + + # Or use positional arg (last arg is editor if it matches code/cursor/windsurf/tmux) + brev open instance1 instance2 cursor + # Set a default editor brev open --set-default cursor brev open --set-default windsurf @@ -56,6 +66,9 @@ var ( # Create a GPU instance and open it immediately (reads instance name from stdin) brev create my-instance | brev open + # Open a cluster (multiple instances from stdin) + brev create my-cluster --count 3 | brev open + # Create with specific GPU and open in Cursor brev search --gpu-name A100 | brev create ml-box | brev open cursor @@ -81,10 +94,11 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto var directory string var host bool var setDefault string + var editor string cmd := &cobra.Command{ Annotations: map[string]string{"access": ""}, - Use: "open", + Use: "open [instance...] [editor]", DisableFlagsInUseLine: true, Short: "[beta] Open VSCode, Cursor, Windsurf, or tmux to your instance", Long: openLong, @@ -94,8 +108,8 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto if setDefaultFlag != "" { return cobra.NoArgs(cmd, args) } - // Allow 0-2 args: instance name can come from stdin - return cobra.RangeArgs(0, 2)(cmd, args) + // Allow arbitrary args: instance names can come from stdin, last arg might be editor + return nil }), ValidArgsFunction: completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t), RunE: func(cmd *cobra.Command, args []string) error { @@ -103,25 +117,45 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto return handleSetDefault(t, setDefault) } - // Get instance name from args or stdin - instanceName, remainingArgs, err := getInstanceNameOpen(args) + // Validate editor flag if provided + if editor != "" && !isEditorType(editor) { + return breverrors.NewValidationError(fmt.Sprintf("invalid editor: %s. Must be 'code', 'cursor', 'windsurf', or 'tmux'", editor)) + } + + // Get instance names and editor type from args or stdin + instanceNames, editorType, err := getInstanceNamesAndEditor(args, editor) if err != nil { return breverrors.WrapAndTrace(err) } + // tmux only supports one instance (interactive) + if editorType == EditorTmux && len(instanceNames) > 1 { + return breverrors.NewValidationError("tmux only supports one instance since it requires interactive stdin") + } + setupDoneString := "------ Git repo cloned ------" if waitForSetupToFinish { setupDoneString = "------ Done running execs ------" } - editorType, err := determineEditorType(remainingArgs) - if err != nil { - return breverrors.WrapAndTrace(err) + // Open each instance + var lastErr error + for _, instanceName := range instanceNames { + if len(instanceNames) > 1 { + fmt.Fprintf(os.Stderr, "Opening %s...\n", instanceName) + } + err = runOpenCommand(t, store, instanceName, setupDoneString, directory, host, editorType) + if err != nil { + if len(instanceNames) > 1 { + fmt.Fprintf(os.Stderr, "Error opening %s: %v\n", instanceName, err) + lastErr = err + continue + } + return breverrors.WrapAndTrace(err) + } } - - err = runOpenCommand(t, store, instanceName, setupDoneString, directory, host, editorType) - if err != nil { - return breverrors.WrapAndTrace(err) + if lastErr != nil { + return breverrors.NewValidationError("one or more instances failed to open") } return nil }, @@ -130,34 +164,64 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto cmd.Flags().BoolVarP(&waitForSetupToFinish, "wait", "w", false, "wait for setup to finish") cmd.Flags().StringVarP(&directory, "dir", "d", "", "directory to open") cmd.Flags().StringVar(&setDefault, "set-default", "", "set default editor (code, cursor, windsurf, or tmux)") + cmd.Flags().StringVarP(&editor, "editor", "e", "", "editor to use (code, cursor, windsurf, or tmux)") return cmd } -// getInstanceNameOpen gets the instance name from args or stdin, returning remaining args for editor type -func getInstanceNameOpen(args []string) (string, []string, error) { +// isEditorType checks if a string is a valid editor type +func isEditorType(s string) bool { + return s == EditorVSCode || s == EditorCursor || s == EditorWindsurf || s == EditorTmux +} + +// getInstanceNamesAndEditor gets instance names from args/stdin and determines editor type +// editorFlag takes precedence, otherwise last arg may be an editor type (code, cursor, windsurf, tmux) +func getInstanceNamesAndEditor(args []string, editorFlag string) ([]string, string, error) { + var names []string + editorType := editorFlag + + // If no editor flag, check if last arg is an editor type + if editorType == "" && len(args) > 0 && isEditorType(args[len(args)-1]) { + editorType = args[len(args)-1] + args = args[:len(args)-1] + } + + // Add names from remaining args + names = append(names, args...) + // Check if stdin is piped stat, _ := os.Stdin.Stat() - stdinPiped := (stat.Mode() & os.ModeCharDevice) == 0 - - if stdinPiped { - // Read instance name from stdin + if (stat.Mode() & os.ModeCharDevice) == 0 { + // Stdin is piped, read instance names (one per line) scanner := bufio.NewScanner(os.Stdin) - if scanner.Scan() { + for scanner.Scan() { name := strings.TrimSpace(scanner.Text()) if name != "" { - // All args are for editor type - return name, args, nil + names = append(names, name) } } } - // Instance name from args - if len(args) > 0 { - return args[0], args[1:], nil + if len(names) == 0 { + return nil, "", breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") + } + + // If no editor specified, get default + if editorType == "" { + homeDir, err := os.UserHomeDir() + if err != nil { + editorType = EditorVSCode + } else { + settings, err := files.ReadPersonalSettings(files.AppFs, homeDir) + if err != nil { + editorType = EditorVSCode + } else { + editorType = settings.DefaultEditor + } + } } - return "", nil, breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") + return names, editorType, nil } func handleSetDefault(t *terminal.Terminal, editorType string) error { @@ -183,29 +247,6 @@ func handleSetDefault(t *terminal.Terminal, editorType string) error { return nil } -func determineEditorType(args []string) (string, error) { - // args now contains only the editor type (if provided), not the instance name - if len(args) >= 1 { - editorType := args[0] - if editorType != EditorVSCode && editorType != EditorCursor && editorType != EditorWindsurf && editorType != EditorTmux { - return "", fmt.Errorf("invalid editor type: %s. Must be 'code', 'cursor', 'windsurf', or 'tmux'", editorType) - } - return editorType, nil - } - - homeDir, err := os.UserHomeDir() - if err != nil { - return EditorVSCode, nil - } - - settings, err := files.ReadPersonalSettings(files.AppFs, homeDir) - if err != nil { - return EditorVSCode, nil - } - - return settings.DefaultEditor, nil -} - // Fetch workspace info, then open code editor func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, setupDoneString string, directory string, host bool, editorType string) error { //nolint:funlen,gocyclo // define brev command // todo check if workspace is stopped and start if it if it is stopped diff --git a/pkg/cmd/shell/shell.go b/pkg/cmd/shell/shell.go index 7d172217..c24837f3 100644 --- a/pkg/cmd/shell/shell.go +++ b/pkg/cmd/shell/shell.go @@ -10,7 +10,6 @@ import ( "time" "github.com/brevdev/brev-cli/pkg/analytics" - "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" "github.com/brevdev/brev-cli/pkg/cmd/completions" "github.com/brevdev/brev-cli/pkg/cmd/hello" "github.com/brevdev/brev-cli/pkg/cmd/refresh" @@ -35,12 +34,18 @@ var ( brev shell my-instance -c "nvidia-smi" brev shell my-instance -c "python train.py" + # Run a command on multiple instances + brev shell instance1 instance2 instance3 -c "nvidia-smi" + # Run a script file on the instance brev shell my-instance -c @setup.sh - # Chain: create and run a command (reads instance name from stdin) + # Chain: create and run a command (reads instance names from stdin) brev create my-instance | brev shell -c "nvidia-smi" + # Run command on a cluster (multiple instances from stdin) + brev create my-cluster --count 3 | brev shell -c "nvidia-smi" + # Create a GPU instance and SSH into it (use command substitution for interactive shell) brev shell $(brev create my-instance) @@ -66,17 +71,17 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell var command string cmd := &cobra.Command{ Annotations: map[string]string{"access": ""}, - Use: "shell", + Use: "shell [instance...]", Aliases: []string{"ssh"}, DisableFlagsInUseLine: true, Short: "[beta] Open a shell in your instance", Long: openLong, Example: openExample, - Args: cmderrors.TransformToValidationError(cobra.RangeArgs(0, 1)), + Args: cobra.ArbitraryArgs, ValidArgsFunction: completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t), RunE: func(cmd *cobra.Command, args []string) error { - // Get instance name from args or stdin - instanceName, err := getInstanceName(args) + // Get instance names from args or stdin + instanceNames, err := getInstanceNames(args) if err != nil { return breverrors.WrapAndTrace(err) } @@ -87,9 +92,29 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell return breverrors.WrapAndTrace(err) } - err = runShellCommand(t, store, instanceName, host, cmdToRun) - if err != nil { - return breverrors.WrapAndTrace(err) + // Interactive shell only supports one instance + if cmdToRun == "" && len(instanceNames) > 1 { + return breverrors.NewValidationError("interactive shell only supports one instance; use -c to run a command on multiple instances") + } + + // Run on each instance + var lastErr error + for _, instanceName := range instanceNames { + if len(instanceNames) > 1 { + fmt.Fprintf(os.Stderr, "\n=== %s ===\n", instanceName) + } + err = runShellCommand(t, store, instanceName, host, cmdToRun) + if err != nil { + if len(instanceNames) > 1 { + fmt.Fprintf(os.Stderr, "Error on %s: %v\n", instanceName, err) + lastErr = err + continue + } + return breverrors.WrapAndTrace(err) + } + } + if lastErr != nil { + return breverrors.NewValidationError("one or more instances failed") } return nil }, @@ -100,26 +125,31 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell return cmd } -// getInstanceName gets the instance name from args or stdin -func getInstanceName(args []string) (string, error) { - if len(args) > 0 { - return args[0], nil - } +// getInstanceNames gets instance names from args or stdin (supports multiple) +func getInstanceNames(args []string) ([]string, error) { + var names []string + + // Add names from args + names = append(names, args...) // Check if stdin is piped stat, _ := os.Stdin.Stat() if (stat.Mode() & os.ModeCharDevice) == 0 { - // Stdin is piped, read instance name + // Stdin is piped, read instance names (one per line) scanner := bufio.NewScanner(os.Stdin) - if scanner.Scan() { + for scanner.Scan() { name := strings.TrimSpace(scanner.Text()) if name != "" { - return name, nil + names = append(names, name) } } } - return "", breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") + if len(names) == 0 { + return nil, breverrors.NewValidationError("instance name required: provide as argument or pipe from another command") + } + + return names, nil } // parseCommand parses the command string, loading from file if prefixed with @ From eece297dc3fbfd8c8bc88195b67ea4d5869b9a50 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 01:40:17 -0800 Subject: [PATCH 11/13] Add Terminal.app support to brev open (cross-platform) - Add 'terminal' editor type: opens new terminal window with SSH - Update 'tmux' to open in new terminal window with tmux session - Both now support multiple instances (each opens new window) - Add --editor / -e flag for explicit editor selection Cross-platform terminal support: macOS: Terminal.app (via osascript) Linux: gnome-terminal, konsole, or xterm (tries in order) WSL: Windows Terminal (wt.exe) Windows: Windows Terminal or cmd Editor options: code - VS Code cursor - Cursor windsurf - Windsurf terminal - New terminal window with SSH tmux - New terminal window with SSH + tmux --- pkg/cmd/open/open.go | 138 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 112 insertions(+), 26 deletions(-) diff --git a/pkg/cmd/open/open.go b/pkg/cmd/open/open.go index 7a2edf16..f5a5fa3f 100644 --- a/pkg/cmd/open/open.go +++ b/pkg/cmd/open/open.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "os/exec" + "runtime" "strings" "time" @@ -31,14 +32,30 @@ import ( ) const ( - EditorVSCode = "code" - EditorCursor = "cursor" - EditorWindsurf = "windsurf" - EditorTmux = "tmux" + EditorVSCode = "code" + EditorCursor = "cursor" + EditorWindsurf = "windsurf" + EditorTerminal = "terminal" + EditorTmux = "tmux" ) var ( - openLong = "[command in beta] This will open VS Code, Cursor, Windsurf, or tmux SSH-ed in to your instance. You must have the editor installed in your path." + openLong = `[command in beta] This will open an editor SSH-ed in to your instance. + +Supported editors: + code - VS Code + cursor - Cursor + windsurf - Windsurf + terminal - Opens a new terminal window with SSH + tmux - Opens a new terminal window with SSH + tmux session + +Terminal support by platform: + macOS: Terminal.app + Linux: gnome-terminal, konsole, or xterm + WSL: Windows Terminal (wt.exe) + Windows: Windows Terminal or cmd + +You must have the editor installed in your path.` openExample = ` # Open an instance by name or ID brev open instance_id_or_name brev open my-instance @@ -50,6 +67,7 @@ var ( brev open my-instance code brev open my-instance cursor brev open my-instance windsurf + brev open my-instance terminal brev open my-instance tmux # Open multiple instances with specific editor (flag is explicit) @@ -72,8 +90,11 @@ var ( # Create with specific GPU and open in Cursor brev search --gpu-name A100 | brev create ml-box | brev open cursor - # Note: tmux requires command substitution (not pipes) since it needs interactive stdin - brev open $(brev create my-instance) tmux` + # Open in a new terminal window with SSH + brev create my-instance | brev open terminal + + # Open in a new terminal window with tmux (supports multiple instances) + brev create my-cluster --count 3 | brev open tmux` ) type OpenStore interface { @@ -128,10 +149,6 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto return breverrors.WrapAndTrace(err) } - // tmux only supports one instance (interactive) - if editorType == EditorTmux && len(instanceNames) > 1 { - return breverrors.NewValidationError("tmux only supports one instance since it requires interactive stdin") - } setupDoneString := "------ Git repo cloned ------" if waitForSetupToFinish { @@ -163,15 +180,15 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto cmd.Flags().BoolVarP(&host, "host", "", false, "ssh into the host machine instead of the container") cmd.Flags().BoolVarP(&waitForSetupToFinish, "wait", "w", false, "wait for setup to finish") cmd.Flags().StringVarP(&directory, "dir", "d", "", "directory to open") - cmd.Flags().StringVar(&setDefault, "set-default", "", "set default editor (code, cursor, windsurf, or tmux)") - cmd.Flags().StringVarP(&editor, "editor", "e", "", "editor to use (code, cursor, windsurf, or tmux)") + cmd.Flags().StringVar(&setDefault, "set-default", "", "set default editor (code, cursor, windsurf, terminal, or tmux)") + cmd.Flags().StringVarP(&editor, "editor", "e", "", "editor to use (code, cursor, windsurf, terminal, or tmux)") return cmd } // isEditorType checks if a string is a valid editor type func isEditorType(s string) bool { - return s == EditorVSCode || s == EditorCursor || s == EditorWindsurf || s == EditorTmux + return s == EditorVSCode || s == EditorCursor || s == EditorWindsurf || s == EditorTerminal || s == EditorTmux } // getInstanceNamesAndEditor gets instance names from args/stdin and determines editor type @@ -225,8 +242,8 @@ func getInstanceNamesAndEditor(args []string, editorFlag string) ([]string, stri } func handleSetDefault(t *terminal.Terminal, editorType string) error { - if editorType != EditorVSCode && editorType != EditorCursor && editorType != EditorWindsurf && editorType != EditorTmux { - return fmt.Errorf("invalid editor type: %s. Must be 'code', 'cursor', 'windsurf', or 'tmux'", editorType) + if !isEditorType(editorType) { + return fmt.Errorf("invalid editor type: %s. Must be 'code', 'cursor', 'windsurf', 'terminal', or 'tmux'", editorType) } homeDir, err := os.UserHomeDir() @@ -457,6 +474,8 @@ func getEditorName(editorType string) string { return "Cursor" case EditorWindsurf: return "Windsurf" + case EditorTerminal: + return "Terminal" case EditorTmux: return "tmux" default: @@ -487,8 +506,10 @@ func openEditorByType(t *terminal.Terminal, editorType string, sshAlias string, case EditorWindsurf: tryToInstallWindsurfExtensions(t, extensions) return openWindsurf(sshAlias, path, tstore) + case EditorTerminal: + return openTerminal(sshAlias, path, tstore) case EditorTmux: - return openTmux(sshAlias, path, tstore) + return openTerminalWithTmux(sshAlias, path, tstore) default: tryToInstallExtensions(t, extensions) return openVsCode(sshAlias, path, tstore) @@ -636,8 +657,75 @@ func getWindowsWindsurfPaths(store vscodePathStore) []string { return paths } -func openTmux(sshAlias string, path string, store OpenStore) error { +// openInNewTerminalWindow opens a command in a new terminal window based on the platform +// macOS: Terminal.app via osascript +// Linux: gnome-terminal, konsole, or xterm (tries in order) +// Windows/WSL: Windows Terminal (wt.exe) +func openInNewTerminalWindow(command string) error { + switch runtime.GOOS { + case "darwin": + // macOS: use osascript to open Terminal.app + script := fmt.Sprintf(`tell application "Terminal" + activate + do script "%s" +end tell`, command) + cmd := exec.Command("osascript", "-e", script) // #nosec G204 + return cmd.Run() + + case "linux": + // Check if we're in WSL by looking for wt.exe + if _, err := exec.LookPath("wt.exe"); err == nil { + // WSL: use Windows Terminal + cmd := exec.Command("wt.exe", "new-tab", "bash", "-c", command) // #nosec G204 + return cmd.Run() + } + // Try gnome-terminal first (Ubuntu/GNOME) + if _, err := exec.LookPath("gnome-terminal"); err == nil { + cmd := exec.Command("gnome-terminal", "--", "bash", "-c", command+"; exec bash") // #nosec G204 + return cmd.Run() + } + // Try konsole (KDE) + if _, err := exec.LookPath("konsole"); err == nil { + cmd := exec.Command("konsole", "-e", "bash", "-c", command+"; exec bash") // #nosec G204 + return cmd.Run() + } + // Try xterm as fallback + if _, err := exec.LookPath("xterm"); err == nil { + cmd := exec.Command("xterm", "-e", "bash", "-c", command+"; exec bash") // #nosec G204 + return cmd.Run() + } + return breverrors.NewValidationError("no supported terminal emulator found. Install gnome-terminal, konsole, or xterm") + + case "windows": + // Windows: use Windows Terminal + if _, err := exec.LookPath("wt.exe"); err == nil { + cmd := exec.Command("wt.exe", "new-tab", "cmd", "/c", command) // #nosec G204 + return cmd.Run() + } + // Fallback to start cmd + cmd := exec.Command("cmd", "/c", "start", "cmd", "/k", command) // #nosec G204 + return cmd.Run() + + default: + return breverrors.NewValidationError(fmt.Sprintf("'terminal' editor is not supported on %s", runtime.GOOS)) + } +} + +func openTerminal(sshAlias string, path string, store OpenStore) error { + _ = store // unused parameter required by interface + _ = path // unused, just opens SSH + + sshCmd := fmt.Sprintf("ssh %s", sshAlias) + err := openInNewTerminalWindow(sshCmd) + if err != nil { + return breverrors.WrapAndTrace(err) + } + return nil +} + +func openTerminalWithTmux(sshAlias string, path string, store OpenStore) error { _ = store // unused parameter required by interface + err := ensureTmuxInstalled(sshAlias) if err != nil { return breverrors.WrapAndTrace(err) @@ -645,23 +733,21 @@ func openTmux(sshAlias string, path string, store OpenStore) error { sessionName := "brev" + // Check if tmux session exists checkCmd := fmt.Sprintf("ssh %s 'tmux has-session -t %s 2>/dev/null'", sshAlias, sessionName) checkExec := exec.Command("bash", "-c", checkCmd) // #nosec G204 - err = checkExec.Run() + checkErr := checkExec.Run() var tmuxCmd string - if err == nil { + if checkErr == nil { + // Session exists, attach to it tmuxCmd = fmt.Sprintf("ssh -t %s 'tmux attach-session -t %s'", sshAlias, sessionName) } else { + // Create new session tmuxCmd = fmt.Sprintf("ssh -t %s 'cd %s && tmux new-session -s %s'", sshAlias, path, sessionName) } - sshCmd := exec.Command("bash", "-c", tmuxCmd) // #nosec G204 - sshCmd.Stderr = os.Stderr - sshCmd.Stdout = os.Stdout - sshCmd.Stdin = os.Stdin - - err = sshCmd.Run() + err = openInNewTerminalWindow(tmuxCmd) if err != nil { return breverrors.WrapAndTrace(err) } From f0b37bed4967866c9e08bff63d97bc1933b89138 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 01:50:57 -0800 Subject: [PATCH 12/13] Fix WSL exec format error when opening editors In WSL, Windows .exe files cannot be executed directly - they need to go through cmd.exe. This fixes the "Exec format error" when running `brev open cursor` or similar commands in WSL. Changes: - Add isWSL() helper to detect Windows Subsystem for Linux - Add wslPathToWindows() to convert /mnt/c/... paths to C:\... - Add runWindowsExeInWSL() to run Windows executables via cmd.exe - Update runVsCodeCommand, runCursorCommand, runWindsurfCommand to detect WSL and use cmd.exe for Windows executables This allows users in WSL to seamlessly open VS Code, Cursor, and Windsurf installed on the Windows side. --- pkg/util/util.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/pkg/util/util.go b/pkg/util/util.go index d9004e1d..299ba5b1 100644 --- a/pkg/util/util.go +++ b/pkg/util/util.go @@ -6,6 +6,7 @@ import ( "os" "os/exec" "path/filepath" + "runtime" "strings" breverrors "github.com/brevdev/brev-cli/pkg/errors" @@ -13,6 +14,53 @@ import ( "github.com/hashicorp/go-multierror" ) +// isWSL returns true if running in Windows Subsystem for Linux +func isWSL() bool { + if runtime.GOOS != "linux" { + return false + } + // Check for WSL-specific indicators + if _, err := os.Stat("/proc/sys/fs/binfmt_misc/WSLInterop"); err == nil { + return true + } + // Also check /proc/version for "microsoft" or "WSL" + if data, err := os.ReadFile("/proc/version"); err == nil { + lower := strings.ToLower(string(data)) + if strings.Contains(lower, "microsoft") || strings.Contains(lower, "wsl") { + return true + } + } + return false +} + +// wslPathToWindows converts a WSL path like /mnt/c/Users/... to C:\Users\... +func wslPathToWindows(wslPath string) string { + if strings.HasPrefix(wslPath, "/mnt/") && len(wslPath) > 6 { + // Extract drive letter: /mnt/c/... -> c + drive := strings.ToUpper(string(wslPath[5])) + // Get rest of path: /mnt/c/Users/... -> /Users/... + rest := wslPath[6:] + // Convert to Windows path: C:\Users\... + windowsPath := drive + ":" + strings.ReplaceAll(rest, "/", "\\") + return windowsPath + } + return wslPath +} + +// runWindowsExeInWSL runs a Windows executable from WSL using cmd.exe +func runWindowsExeInWSL(exePath string, args []string) ([]byte, error) { + // Convert WSL path to Windows path + windowsPath := wslPathToWindows(exePath) + + // Build the command string for cmd.exe + // We need to quote the path and args properly for Windows + cmdArgs := []string{"/c", windowsPath} + cmdArgs = append(cmdArgs, args...) + + cmd := exec.Command("cmd.exe", cmdArgs...) // #nosec G204 + return cmd.CombinedOutput() +} + // This package should only be used as a holding pattern to be later moved into more specific packages func MapAppend(m map[string]interface{}, n ...map[string]interface{}) map[string]interface{} { @@ -205,6 +253,15 @@ func runManyCursorCommand(cursorpaths []string, args []string) ([]byte, error) { } func runVsCodeCommand(vscodepath string, args []string) ([]byte, error) { + // In WSL, Windows .exe files need to be run through cmd.exe + if isWSL() && (strings.HasSuffix(vscodepath, ".exe") || strings.HasPrefix(vscodepath, "/mnt/")) { + res, err := runWindowsExeInWSL(vscodepath, args) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + return res, nil + } + cmd := exec.Command(vscodepath, args...) // #nosec G204 res, err := cmd.CombinedOutput() if err != nil { @@ -214,6 +271,15 @@ func runVsCodeCommand(vscodepath string, args []string) ([]byte, error) { } func runCursorCommand(cursorpath string, args []string) ([]byte, error) { + // In WSL, Windows .exe files need to be run through cmd.exe + if isWSL() && (strings.HasSuffix(cursorpath, ".exe") || strings.HasPrefix(cursorpath, "/mnt/")) { + res, err := runWindowsExeInWSL(cursorpath, args) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + return res, nil + } + cmd := exec.Command(cursorpath, args...) // #nosec G204 res, err := cmd.CombinedOutput() if err != nil { @@ -236,6 +302,15 @@ func runManyWindsurfCommand(windsurfpaths []string, args []string) ([]byte, erro } func runWindsurfCommand(windsurfpath string, args []string) ([]byte, error) { + // In WSL, Windows .exe files need to be run through cmd.exe + if isWSL() && (strings.HasSuffix(windsurfpath, ".exe") || strings.HasPrefix(windsurfpath, "/mnt/")) { + res, err := runWindowsExeInWSL(windsurfpath, args) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + return res, nil + } + cmd := exec.Command(windsurfpath, args...) // #nosec G204 res, err := cmd.CombinedOutput() if err != nil { From 23e28b21b33e4b1b9b66e5e6562fc7421334cba1 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Wed, 28 Jan 2026 01:53:02 -0800 Subject: [PATCH 13/13] Improve gpu-create retry logic to exhaust each type before fallback Previously, workers would each grab different instance types in parallel, leading to mixed type usage even when capacity was available. Now the command tries to create ALL instances with the first type before moving to the next, resulting in more consistent instance configurations. Changes: - Refactor parallel workers to focus on one type at a time - Remove unused context import - Update documentation to explain the new retry behavior with examples --- pkg/cmd/gpucreate/gpucreate.go | 224 ++++++++++++++++++--------------- 1 file changed, 122 insertions(+), 102 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index f10f2218..2af8f96f 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -2,7 +2,6 @@ package gpucreate import ( - "context" "encoding/json" "fmt" "io" @@ -39,11 +38,25 @@ automatically searches for GPUs matching these criteria: - Boot time under 7 minutes Results are sorted by price (cheapest first). -The command will: -1. Try to create instances using the provided instance types (in order) -2. Continue until the desired count is reached -3. Optionally try multiple instance types in parallel -4. Clean up any extra instances that were created beyond the requested count +Retry and Fallback Logic: +When multiple instance types are provided (via --type or piped input), the command +tries to create ALL instances using the first type before falling back to the next: + + 1. Try first type for all instances (using --parallel workers if specified) + 2. If first type succeeds for all instances, done + 3. If first type fails for some instances, try second type for remaining instances + 4. Continue until all instances are created or all types are exhausted + +Example with --count 2 and types [A, B]: + - Try A for instance-1 → success + - Try A for instance-2 → success + - Done! (both instances use type A) + +If type A fails for instance-2: + - Try A for instance-1 → success + - Try A for instance-2 → fail + - Try B for instance-2 → success + - Done! (instance-1 uses A, instance-2 uses B) Startup Scripts: You can attach a startup script that runs when the instance boots using the @@ -71,13 +84,21 @@ 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 - # Pipe instance types from brev search (tries each type until one succeeds) + # Pipe instance types from brev search (tries first type, falls back if needed) brev search --min-vram 24 | brev create my-instance - # Create 3 instances, trying types in parallel - brev search --gpu-name A100 | brev create my-cluster --count 3 --parallel 5 + # Create multiple instances (all use same type, with fallback) + brev create my-cluster --count 3 --type g5.xlarge + # Creates: my-cluster-1, my-cluster-2, my-cluster-3 (all g5.xlarge) + + # Create multiple instances with fallback types + brev search --gpu-name A100 | brev create my-cluster --count 2 + # Tries first A100 type for both instances, falls back to next type if needed + + # Create instances in parallel (faster, but may use more types on partial failures) + brev search --gpu-name A100 | brev create my-cluster --count 3 --parallel 3 - # Try multiple specific types in order + # Try multiple specific types in order (fallback chain) brev create my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge # Attach a startup script from a file @@ -472,116 +493,115 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC // Track successful creations var successfulWorkspaces []*entity.Workspace - var mu sync.Mutex - var wg sync.WaitGroup + var fatalError error - // Create a context for cancellation - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - // Channel to coordinate attempts - specsChan := make(chan InstanceSpec, len(opts.InstanceTypes)) + // Try each instance type in order, attempting to create ALL instances with that type + // before falling back to the next type for _, spec := range opts.InstanceTypes { - specsChan <- spec - } - close(specsChan) - - // Results channel - resultsChan := make(chan CreateResult, len(opts.InstanceTypes)) - - // Track instance index for naming (incremented only on successful creation) - instanceIndex := 0 - - // Start parallel workers - workerCount := opts.Parallel - if workerCount > len(opts.InstanceTypes) { - workerCount = len(opts.InstanceTypes) - } + // Check if we've created enough instances + if len(successfulWorkspaces) >= opts.Count { + break + } - for i := 0; i < workerCount; i++ { - wg.Add(1) - go func(workerID int) { - defer wg.Done() + remaining := opts.Count - len(successfulWorkspaces) + logf("Trying %s for %d instance(s)...\n", spec.Type, remaining) - for spec := range specsChan { - // Check if we've already created enough - mu.Lock() - if len(successfulWorkspaces) >= opts.Count { - mu.Unlock() - return - } - // Get current index for naming (only increment on success) - currentIndex := instanceIndex - mu.Unlock() - - // Check context - select { - case <-ctx.Done(): - return - default: - } + // Create instances with this type (in parallel if requested) + var mu sync.Mutex + var wg sync.WaitGroup - // Determine instance name based on current successful count - instanceName := opts.Name - if opts.Count > 1 || currentIndex > 0 { - instanceName = fmt.Sprintf("%s-%d", opts.Name, currentIndex+1) - } + // Determine how many parallel workers to use + workerCount := opts.Parallel + if workerCount > remaining { + workerCount = remaining + } - logf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, spec.Type, instanceName) + // Track which instance indices need to be created + indicesToCreate := make(chan int, remaining) + for i := len(successfulWorkspaces); i < opts.Count; i++ { + indicesToCreate <- i + } + close(indicesToCreate) - // Attempt to create the workspace - workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, spec.Type, spec.DiskGB, user, allInstanceTypes, opts.StartupScript) + // Track results for this type + var typeSuccesses []*entity.Workspace + var typeHadFailure bool - result := CreateResult{ - Workspace: workspace, - InstanceType: spec.Type, - Error: err, - } + for i := 0; i < workerCount; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() - if err != nil { - errStr := err.Error() - if piped { - logf("[Worker %d] %s Failed: %s\n", workerID+1, spec.Type, errStr) - } else { - logf("[Worker %d] %s Failed: %s\n", workerID+1, t.Yellow(spec.Type), errStr) + for idx := range indicesToCreate { + // Check if we've already created enough + mu.Lock() + currentSuccessCount := len(successfulWorkspaces) + len(typeSuccesses) + if currentSuccessCount >= opts.Count { + mu.Unlock() + return } + mu.Unlock() - // Check for fatal errors that should stop all workers - if strings.Contains(errStr, "duplicate workspace") { - logf("\nError: Workspace '%s' already exists. Use a different name or delete the existing workspace.\n", instanceName) - cancel() // Stop all workers - resultsChan <- result - return + // Determine instance name + instanceName := opts.Name + if opts.Count > 1 { + instanceName = fmt.Sprintf("%s-%d", opts.Name, idx+1) } - } else { - if piped { - logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, spec.Type, instanceName) + + logf("[Worker %d] Trying %s for instance '%s'...\n", workerID+1, spec.Type, instanceName) + + // Attempt to create the workspace + workspace, err := createWorkspaceWithType(gpuCreateStore, org.ID, instanceName, spec.Type, spec.DiskGB, user, allInstanceTypes, opts.StartupScript) + + if err != nil { + errStr := err.Error() + if piped { + logf("[Worker %d] %s Failed: %s\n", workerID+1, spec.Type, errStr) + } else { + logf("[Worker %d] %s Failed: %s\n", workerID+1, t.Yellow(spec.Type), errStr) + } + + mu.Lock() + typeHadFailure = true + // Check for fatal errors + if strings.Contains(errStr, "duplicate workspace") { + fatalError = fmt.Errorf("workspace '%s' already exists. Use a different name or delete the existing workspace", instanceName) + } + mu.Unlock() } else { - logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, t.Green(spec.Type), instanceName) + if piped { + logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, spec.Type, instanceName) + } else { + logf("[Worker %d] %s Success! Created instance '%s'\n", workerID+1, t.Green(spec.Type), instanceName) + } + mu.Lock() + typeSuccesses = append(typeSuccesses, workspace) + mu.Unlock() } - mu.Lock() - successfulWorkspaces = append(successfulWorkspaces, workspace) - instanceIndex++ // Only increment on success - if len(successfulWorkspaces) >= opts.Count { - cancel() // Signal other workers to stop - } - mu.Unlock() } + }(i) + } - resultsChan <- result - } - }(i) - } - - // Wait for all workers to finish - go func() { wg.Wait() - close(resultsChan) - }() - // Collect results - for range resultsChan { - // Just drain the channel + // Add successful creations from this type + successfulWorkspaces = append(successfulWorkspaces, typeSuccesses...) + + // Check for fatal error + if fatalError != nil { + logf("\nError: %s\n", fatalError.Error()) + break + } + + // If this type worked for all remaining instances, we're done + if !typeHadFailure && len(successfulWorkspaces) >= opts.Count { + break + } + + // If we still need more instances and this type had failures, try the next type + if len(successfulWorkspaces) < opts.Count && typeHadFailure { + logf("\nType %s had failures, trying next type...\n\n", spec.Type) + } } // Check if we created enough instances