Skip to content

feat(create): add --disk-size flag #439

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
theFong wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat(create): add --disk-size flag #439

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 27 additions & 8 deletions pkg/cmd/gpucreate/gpucreate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"encoding/json"
"fmt"
"io"
"math"
"math/rand/v2"
"net"
"net/http"
Expand Down Expand Up @@ -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)
Expand All @@ -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

Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand All @@ -248,7 +265,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
},
}

registerCreateFlags(cmd, &name, &instanceTypes, &count, &parallel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &launchable, &filters)
registerCreateFlags(cmd, &name, &instanceTypes, &count, &parallel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &launchable, &diskSize, &filters)

return cmd
}
Expand All @@ -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")
Expand All @@ -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")
Expand Down Expand Up @@ -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")
}
Expand Down Expand Up @@ -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)
Expand All @@ -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(
Expand Down
36 changes: 34 additions & 2 deletions pkg/cmd/gpucreate/gpucreate_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ type MockGPUCreateStore struct {
CreatedWorkspaces []*entity.Workspace
DeletedWorkspaceIDs []string
FetchedLifeCycleScriptIDs []string
AllInstanceTypes *gpusearch.AllInstanceTypesResponse
}

func NewMockGPUCreateStore() *MockGPUCreateStore {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -1180,6 +1210,7 @@ func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) {
CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{
CloudCredID: "cc-from-launchable",
InstanceType: "n2-standard-4",
Storage: "256",
},
},
},
Expand All @@ -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")
}