Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
194 changes: 184 additions & 10 deletions pkg/cmd/gpucreate/gpucreate.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"encoding/json"
"fmt"
"io"
"net/url"
"os"
"strconv"
"strings"
Expand Down Expand Up @@ -99,6 +100,7 @@
CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error)
DeleteWorkspace(workspaceID string) (*entity.Workspace, error)
GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error)
GetLaunchable(launchableID string) (*store.LaunchableResponse, error)
}

// Default filter values for automatic GPU selection
Expand Down Expand Up @@ -140,7 +142,7 @@
}

// NewCmdGPUCreate creates the gpu-create command
func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra.Command {

Check failure on line 145 in pkg/cmd/gpucreate/gpucreate.go

View workflow job for this annotation

GitHub Actions / ci (ubuntu-22.04)

cognitive complexity 68 of func `NewCmdGPUCreate` is high (> 30) (gocognit)

Check failure on line 145 in pkg/cmd/gpucreate/gpucreate.go

View workflow job for this annotation

GitHub Actions / ci (ubuntu-22.04)

cognitive complexity 68 of func `NewCmdGPUCreate` is high (> 30) (gocognit)
var name string
var instanceTypes string
var count int
Expand All @@ -153,6 +155,7 @@
var jupyter bool
var containerImage string
var composeFile string
var launchable string
var filters searchFilterFlags

cmd := &cobra.Command{
Expand All @@ -169,8 +172,21 @@
name = args[0]
}

if err := validateBuildMode(mode, containerImage, composeFile); err != nil {
return err
launchableID := parseLaunchableID(launchable)

if launchableID != "" {
// Warn about build mode flags that are ignored with launchable
buildFlagsSet := cmd.Flags().Changed("mode") || cmd.Flags().Changed("container-image") ||
cmd.Flags().Changed("compose-file") || cmd.Flags().Changed("startup-script") ||
cmd.Flags().Changed("jupyter")
if buildFlagsSet {
t.Vprintf("Warning: Build config flags (--mode, --container-image, --compose-file, --startup-script) are ignored when deploying a launchable.\n")
t.Vprintf("The launchable defines its own build configuration.\n\n")
}
} else {
if err := validateBuildMode(mode, containerImage, composeFile); err != nil {
return err
}
}

// Parse instance types from flag or stdin
Expand All @@ -179,12 +195,54 @@
return breverrors.WrapAndTrace(err)
}

// Fetch and display launchable info before dry-run or creation
var launchableInfo *store.LaunchableResponse
if launchableID != "" {
var launchErr error
launchableInfo, launchErr = gpuCreateStore.GetLaunchable(launchableID)
if launchErr != nil {
return fmt.Errorf("failed to fetch launchable %q: %w", launchableID, launchErr)
}
t.Vprintf("Deploying launchable: %q\n", launchableInfo.Name)
if launchableInfo.Description != "" {
t.Vprintf("Description: %s\n", launchableInfo.Description)
}
if launchableInfo.CreateWorkspaceRequest.InstanceType != "" {
t.Vprintf("Instance type: %s\n", launchableInfo.CreateWorkspaceRequest.InstanceType)
}
if launchableInfo.CreateWorkspaceRequest.Storage != "" {
t.Vprintf("Storage: %s\n", launchableInfo.CreateWorkspaceRequest.Storage)
}
buildMode := "VM"
if launchableInfo.BuildRequest.CustomContainer != nil {
buildMode = "Container"
} else if launchableInfo.BuildRequest.DockerCompose != nil {
buildMode = "Docker Compose"
}
t.Vprintf("Build mode: %s\n\n", buildMode)
}

if dryRun {
if launchableID != "" {
return nil // launchable info already displayed above
}
return runDryRun(t, gpuCreateStore, types, &filters)
}

// If no types provided, use search filters (or defaults) to find suitable GPUs
if len(types) == 0 {
// If deploying a launchable and no types provided, use the launchable's
// instance type so we can resolve the correct workspace group
if launchableID != "" && len(types) == 0 && !cmd.Flags().Changed("type") {
launchableInstanceType := ""
if launchableInfo != nil {
launchableInstanceType = launchableInfo.CreateWorkspaceRequest.InstanceType
}
if launchableInstanceType != "" {
types = []InstanceSpec{{Type: launchableInstanceType}}
} else {
return breverrors.NewValidationError("launchable has no instance type configured and no --type was specified")
}
} else if len(types) == 0 {
// If no types provided, use search filters (or defaults) to find suitable GPUs
types, err = getFilteredInstanceTypes(gpuCreateStore, &filters)
if err != nil {
return breverrors.WrapAndTrace(err)
Expand All @@ -195,6 +253,12 @@
}
}

// Warn if overriding launchable instance config
if launchableID != "" && (cmd.Flags().Changed("type") || cmd.Flags().Changed("gpu-name") ||
cmd.Flags().Changed("provider") || cmd.Flags().Changed("min-vram")) {
t.Vprintf("Warning: Overriding the launchable's recommended instance configuration. This is not the recommended path and may cause issues.\n\n")
}

if err := names.ValidateNodeName(name); err != nil {
return breverrors.WrapAndTrace(err)
}
Expand Down Expand Up @@ -228,6 +292,8 @@
JupyterSet: jupyterSet,
ContainerImage: containerImage,
ComposeFile: composeFile,
LaunchableID: launchableID,
LaunchableInfo: launchableInfo,
}

err = RunGPUCreate(t, gpuCreateStore, opts)
Expand All @@ -238,13 +304,13 @@
},
}

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

return cmd
}

// 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 *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, 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 @@ -259,6 +325,7 @@
cmd.Flags().BoolVar(jupyter, "jupyter", true, "Install Jupyter (default true for vm/k8s modes)")
cmd.Flags().StringVar(containerImage, "container-image", "", "Container image URL (required for container mode)")
cmd.Flags().StringVar(composeFile, "compose-file", "", "Docker compose file path or URL (required for compose mode)")
cmd.Flags().StringVarP(launchable, "launchable", "l", "", "Launchable ID or URL to deploy (e.g., env-XXX or console URL)")

cmd.Flags().StringVarP(&filters.gpuName, "gpu-name", "g", "", "Filter by GPU name (e.g., A100, H100)")
cmd.Flags().StringVar(&filters.provider, "provider", "", "Filter by provider/cloud (e.g., aws, gcp)")
Expand Down Expand Up @@ -294,6 +361,36 @@
JupyterSet bool // whether --jupyter was explicitly set
ContainerImage string
ComposeFile string
LaunchableID string
LaunchableInfo *store.LaunchableResponse // populated when LaunchableID is set
}

// parseLaunchableID extracts a launchable ID from either a raw ID (env-XXX) or
// a console URL (https://console.brev.dev/launchable/deploy?launchableID=env-XXX)
func parseLaunchableID(input string) string {
if input == "" {
return ""
}
// Check if it looks like a URL
if strings.HasPrefix(input, "http://") || strings.HasPrefix(input, "https://") {
u, err := url.Parse(input)
if err != nil {
return input
}
if id := u.Query().Get("launchableID"); id != "" {
return id
}
// Check path for launchable ID (e.g., /launchables/env-XXX)
parts := strings.Split(strings.TrimRight(u.Path, "/"), "/")
if len(parts) > 0 {
last := parts[len(parts)-1]
if strings.HasPrefix(last, "env-") {
return last
}
}
return input
}
return input
}

// parseStartupScript parses the startup script from a string or file path
Expand Down Expand Up @@ -827,12 +924,19 @@
}
}

// Apply build mode
err := applyBuildMode(cwOptions, c.opts)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
// Apply launchable config or build mode
if c.opts.LaunchableID != "" {
applyLaunchableConfig(cwOptions, c.opts.LaunchableID, c.opts.LaunchableInfo)
} else {
err := applyBuildMode(cwOptions, c.opts)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}
}

c.logf(" Creating workspace: instanceType=%s workspaceGroupID=%s launchable=%v\n",
cwOptions.InstanceType, cwOptions.WorkspaceGroupID, cwOptions.LaunchableConfig != nil)

workspace, err := c.store.CreateWorkspace(c.org.ID, cwOptions)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
Expand Down Expand Up @@ -930,6 +1034,76 @@
return nil
}

// applyLaunchableConfig populates the workspace create request with all launchable
// configuration, mirroring what the web UI sends when deploying a launchable.
func applyLaunchableConfig(cwOptions *store.CreateWorkspacesOptions, launchableID string, info *store.LaunchableResponse) {
cwOptions.LaunchableConfig = &store.LaunchableConfig{ID: launchableID}

if info == nil {
return
}

wsReq := info.CreateWorkspaceRequest

// Use launchable's workspace group if not already resolved from instance types
if cwOptions.WorkspaceGroupID == "" && wsReq.WorkspaceGroupID != "" {
cwOptions.WorkspaceGroupID = wsReq.WorkspaceGroupID
}

// Location
if wsReq.Location != "" {
cwOptions.Location = wsReq.Location
}

// Disk storage — ensure Gi suffix
if wsReq.Storage != "" {
storage := wsReq.Storage
if !strings.HasSuffix(storage, "Gi") {
storage += "Gi"
}
cwOptions.DiskStorage = storage
}

// Build configuration from launchable
build := info.BuildRequest
if build.VMBuild != nil {

Check failure on line 1069 in pkg/cmd/gpucreate/gpucreate.go

View workflow job for this annotation

GitHub Actions / ci (ubuntu-22.04)

ifElseChain: rewrite if-else to switch statement (gocritic)

Check failure on line 1069 in pkg/cmd/gpucreate/gpucreate.go

View workflow job for this annotation

GitHub Actions / ci (ubuntu-22.04)

ifElseChain: rewrite if-else to switch statement (gocritic)
cwOptions.VMBuild = build.VMBuild
} else if build.CustomContainer != nil {
cwOptions.VMBuild = nil
cwOptions.CustomContainer = build.CustomContainer
} else if build.DockerCompose != nil {
cwOptions.VMBuild = nil
cwOptions.DockerCompose = build.DockerCompose
}

// Port mappings from build request ports
if len(build.Ports) > 0 {
portMappings := make(map[string]string)
for _, p := range build.Ports {
portMappings[p.Name] = p.Port
}
cwOptions.PortMappings = portMappings
}

// Files from launchable
if info.File != nil {
cwOptions.Files = []map[string]string{
{"url": info.File.URL, "path": info.File.Path},
}
}

// Labels for tracking and UI rendering
labels := map[string]string{
"launchableId": launchableID,
"launchableInstanceType": wsReq.InstanceType,
"workspaceGroupId": cwOptions.WorkspaceGroupID,
"launchableCreatedByUserId": info.CreatedByUserID,
"launchableCreatedByOrgId": info.CreatedByOrgID,
"launchableRawURL": "/launchable/deploy/now?launchableID=" + launchableID,
}
cwOptions.Labels = labels
}

// resolveWorkspaceUserOptions sets workspace template and class based on user type
func resolveWorkspaceUserOptions(options *store.CreateWorkspacesOptions, user *entity.User) *store.CreateWorkspacesOptions {
isAdmin := featureflag.IsAdmin(user.GlobalUserType)
Expand Down
30 changes: 30 additions & 0 deletions pkg/cmd/gpucreate/gpucreate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,13 @@ func (m *MockGPUCreateStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string
return nil, nil
}

func (m *MockGPUCreateStore) GetLaunchable(launchableID string) (*store.LaunchableResponse, error) {
return &store.LaunchableResponse{
ID: launchableID,
Name: "test-launchable",
}, nil
}

func (m *MockGPUCreateStore) GetInstanceTypes(_ bool) (*gpusearch.InstanceTypesResponse, error) {
// Return a default set of instance types for testing
return &gpusearch.InstanceTypesResponse{
Expand Down Expand Up @@ -144,6 +151,29 @@ func TestIsValidInstanceType(t *testing.T) {
}
}

func TestParseLaunchableID(t *testing.T) {
tests := []struct {
name string
input string
expected string
}{
{"Empty string", "", ""},
{"Raw ID", "env-2jeVokEK44iJZzleTF8yKjt3hh7", "env-2jeVokEK44iJZzleTF8yKjt3hh7"},
{"Console URL with query param", "https://console.brev.dev/launchable/deploy?launchableID=env-abc123", "env-abc123"},
{"Console URL with extra params", "https://console.brev.dev/launchable/deploy?userID=u1&launchableID=env-abc123&name=test", "env-abc123"},
{"URL with env- in path", "https://console.brev.dev/launchables/env-abc123", "env-abc123"},
{"URL without launchableID param", "https://console.brev.dev/launchable/deploy", "https://console.brev.dev/launchable/deploy"},
{"Non-env ID", "some-other-id", "some-other-id"},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := parseLaunchableID(tt.input)
assert.Equal(t, tt.expected, result)
})
}
}

func TestParseInstanceTypesFromFlag(t *testing.T) {
tests := []struct {
name string
Expand Down
Loading
Loading