From c35bee50827498e2fd1e1b87560d61c46a98e851 Mon Sep 17 00:00:00 2001 From: Alec Fong Date: Fri, 7 Aug 2026 00:42:25 -0700 Subject: [PATCH] feat(create): add --disk-size flag --- pkg/cmd/gpucreate/gpucreate.go | 35 +++++++++++++++++++++------- pkg/cmd/gpucreate/gpucreate_test.go | 36 +++++++++++++++++++++++++++-- 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index 05194a57..da037b18 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "math" "math/rand/v2" "net" "net/http" @@ -72,7 +73,11 @@ You can attach a startup script that runs when the instance boots using the --startup-script flag. The script can be provided as: - An inline string: --startup-script 'pip install torch' - A file path (prefix with @): --startup-script @setup.sh - - An absolute file path: --startup-script @/path/to/setup.sh` + - An absolute file path: --startup-script @/path/to/setup.sh + +Disk Size: +Use --disk-size to request a disk size in GB. The requested size applies to +every instance type in the fallback chain and must be supported by the provider.` example = ` # Create an instance using smart defaults (sorted by price) @@ -81,6 +86,9 @@ You can attach a startup script that runs when the instance boots using the # Create with a specific GPU type brev create my-instance --type g5.xlarge + # Create with a 1TB disk + brev create my-instance --type g5.xlarge --disk-size 1000 + # Try multiple types in order (fallback chain) brev create my-instance --type g5.xlarge,g5.2xlarge,g4dn.xlarge @@ -164,6 +172,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra var containerImage string var composeFile string var launchable string + var diskSize float64 var filters searchFilterFlags cmd := &cobra.Command{ @@ -207,6 +216,9 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra if err != nil { return err } + if cmd.Flags().Changed("disk-size") && (diskSize <= 0 || math.IsNaN(diskSize) || math.IsInf(diskSize, 0)) { + return breverrors.NewValidationError("--disk-size must be greater than 0") + } types, err := parseInstanceTypes(instanceTypes) if err != nil { @@ -239,6 +251,11 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra if err != nil { return err } + if cmd.Flags().Changed("disk-size") { + for i := range opts.InstanceTypes { + opts.InstanceTypes[i].DiskGB = diskSize + } + } if dryRun { return runDryRun(t, gpuCreateStore, opts.InstanceTypes, &filters) @@ -248,7 +265,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra }, } - registerCreateFlags(cmd, &name, &instanceTypes, &count, ¶llel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &launchable, &filters) + registerCreateFlags(cmd, &name, &instanceTypes, &count, ¶llel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &launchable, &diskSize, &filters) return cmd } @@ -265,7 +282,7 @@ func validateArgs(name string, count int) error { } // registerCreateFlags registers all flags for the create command -func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, parallel *int, detached *bool, timeout *int, startupScript *string, dryRun *bool, mode *string, jupyter *bool, containerImage, composeFile, launchable *string, filters *searchFilterFlags) { +func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, parallel *int, detached *bool, timeout *int, startupScript *string, dryRun *bool, mode *string, jupyter *bool, containerImage, composeFile, launchable *string, diskSize *float64, filters *searchFilterFlags) { 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") @@ -274,6 +291,7 @@ func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count, 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)") cmd.Flags().BoolVar(dryRun, "dry-run", false, "Show matching instance types without creating anything") + cmd.Flags().Float64Var(diskSize, "disk-size", 0, "Disk size to provision in GB") // Build mode flags cmd.Flags().StringVarP(mode, "mode", "m", "vm", "Build mode: vm (default), k8s, container, compose") @@ -380,7 +398,8 @@ func warnLaunchableFlagConflicts(cmd *cobra.Command, t *terminal.Terminal, launc } instanceFlagsSet := cmd.Flags().Changed("type") || cmd.Flags().Changed("gpu-name") || - cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram") + cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram") || + cmd.Flags().Changed("disk-size") if instanceFlagsSet { t.Vprintf("Warning: Overriding the launchable's recommended instance configuration. This is not the recommended path and may cause issues.\n\n") } @@ -1037,10 +1056,6 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity cwOptions.WithInstanceType(spec.Type) cwOptions = resolveWorkspaceUserOptions(cwOptions, c.user) - if spec.DiskGB > 0 { - cwOptions.DiskStorage = fmt.Sprintf("%.0fGi", spec.DiskGB) - } - if c.allInstanceTypes != nil { if cloudCredID := c.allInstanceTypes.GetCloudCredID(spec.Type); cloudCredID != "" { cwOptions.WithCloudCredID(cloudCredID) @@ -1057,6 +1072,10 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity } } + if spec.DiskGB > 0 { + cwOptions.DiskStorage = fmt.Sprintf("%sGi", strconv.FormatFloat(spec.DiskGB, 'f', -1, 64)) + } + if cwOptions.CloudCredID == "" { if c.allInstanceTypes == nil { return nil, breverrors.NewValidationError(fmt.Sprintf( diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index 2c076a1d..c944c0a1 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -29,6 +29,7 @@ type MockGPUCreateStore struct { CreatedWorkspaces []*entity.Workspace DeletedWorkspaceIDs []string FetchedLifeCycleScriptIDs []string + AllInstanceTypes *gpusearch.AllInstanceTypesResponse } func NewMockGPUCreateStore() *MockGPUCreateStore { @@ -108,7 +109,7 @@ func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID strin } func (m *MockGPUCreateStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { - return nil, nil + return m.AllInstanceTypes, nil } func (m *MockGPUCreateStore) GetLaunchable(launchableID string) (*store.LaunchableResponse, error) { @@ -801,6 +802,35 @@ func TestCreateDryRunWithExplicitTypesDoesNotProvision(t *testing.T) { assert.Empty(t, mock.CreatedWorkspaces) } +func TestCreateWithExplicitDiskSize(t *testing.T) { + mock := NewMockGPUCreateStore() + mock.AllInstanceTypes = &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + {Type: "g5.xlarge", CloudCredID: "cc-1"}, + }, + } + + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"disk-size-test", "--type", "g5.xlarge", "--disk-size", "750", "--detached"}) + + err := cmd.Execute() + require.NoError(t, err) + require.Len(t, mock.CreatedOptions, 1) + assert.Equal(t, "750Gi", mock.CreatedOptions[0].DiskStorage) +} + +func TestCreateRejectsInvalidDiskSize(t *testing.T) { + mock := NewMockGPUCreateStore() + cmd := NewCmdGPUCreate(terminal.New(), mock) + cmd.SetArgs([]string{"disk-size-test", "--type", "g5.xlarge", "--disk-size=0"}) + + err := cmd.Execute() + + assert.Error(t, err) + assert.Contains(t, err.Error(), "--disk-size must be greater than 0") + assert.Empty(t, mock.CreatedWorkspaces) +} + func TestGetFilteredInstanceTypesDefaults(t *testing.T) { mock := NewMockGPUCreateStore() @@ -1180,6 +1210,7 @@ func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{ CloudCredID: "cc-from-launchable", InstanceType: "n2-standard-4", + Storage: "256", }, }, }, @@ -1192,9 +1223,10 @@ func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { } ctx.logf = func(_ string, _ ...interface{}) {} - result := ctx.createInstancesWithType(InstanceSpec{Type: "n2-standard-4"}, 0, 1) + result := ctx.createInstancesWithType(InstanceSpec{Type: "n2-standard-4", DiskGB: 750}, 0, 1) assert.False(t, result.hadFailure, "launchable should not be blocked by pre-flight validation") assert.Len(t, result.successes, 1, "expected the launchable instance to be created") assert.Len(t, mock.CreatedWorkspaces, 1) + assert.Equal(t, "750Gi", mock.CreatedOptions[0].DiskStorage, "explicit disk size should override launchable storage") }