Skip to content

feat(create): add build mode flags and fix workspace group resolution #287

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

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

feat(create): add build mode flags and fix workspace group resolution #287

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
3 changes: 3 additions & 0 deletions pkg/cmd/create/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
// Deprecated: This package is superseded by pkg/cmd/gpucreate which is the
// registered "brev create" command. This code is not wired into cmd.go and
// should not be modified. Use gpucreate for all new work.
package create

import (
Expand Down
158 changes: 135 additions & 23 deletions pkg/cmd/gpucreate/gpucreate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,10 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
var timeout int
var startupScript string
var dryRun bool
var mode string
var jupyter bool
var containerImage string
var composeFile string
var filters searchFilterFlags

cmd := &cobra.Command{
Expand All @@ -164,6 +168,10 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
name = args[0]
}

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

// Parse instance types from flag or stdin
types, err := parseInstanceTypes(instanceTypes)
if err != nil {
Expand Down Expand Up @@ -204,14 +212,21 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
return breverrors.WrapAndTrace(err)
}

jupyterSet := cmd.Flags().Changed("jupyter")

opts := GPUCreateOptions{
Name: name,
InstanceTypes: types,
Count: count,
Parallel: parallel,
Detached: detached,
Timeout: time.Duration(timeout) * time.Second,
StartupScript: scriptContent,
Name: name,
InstanceTypes: types,
Count: count,
Parallel: parallel,
Detached: detached,
Timeout: time.Duration(timeout) * time.Second,
StartupScript: scriptContent,
Mode: mode,
Jupyter: jupyter,
JupyterSet: jupyterSet,
ContainerImage: containerImage,
ComposeFile: composeFile,
}

err = RunGPUCreate(t, gpuCreateStore, opts)
Expand All @@ -222,13 +237,13 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
},
}

registerCreateFlags(cmd, &name, &instanceTypes, &count, &parallel, &detached, &timeout, &startupScript, &dryRun, &filters)
registerCreateFlags(cmd, &name, &instanceTypes, &count, &parallel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &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, 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 *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 @@ -238,6 +253,12 @@ func registerCreateFlags(cmd *cobra.Command, name, instanceTypes *string, count,
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")

// Build mode flags
cmd.Flags().StringVarP(mode, "mode", "m", "vm", "Build mode: vm (default), k8s, container, compose")
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(&filters.gpuName, "gpu-name", "g", "", "Filter by GPU name (e.g., A100, H100)")
cmd.Flags().StringVar(&filters.provider, "provider", "", "Filter by provider/cloud (e.g., aws, gcp)")
cmd.Flags().Float64VarP(&filters.minVRAM, "min-vram", "v", 0, "Minimum VRAM per GPU in GB")
Expand All @@ -260,13 +281,18 @@ type InstanceSpec struct {

// GPUCreateOptions holds the options for GPU instance creation
type GPUCreateOptions struct {
Name string
InstanceTypes []InstanceSpec
Count int
Parallel int
Detached bool
Timeout time.Duration
StartupScript string
Name string
InstanceTypes []InstanceSpec
Count int
Parallel int
Detached bool
Timeout time.Duration
StartupScript string
Mode string
Jupyter bool
JupyterSet bool // whether --jupyter was explicitly set
ContainerImage string
ComposeFile string
}

// parseStartupScript parses the startup script from a string or file path
Expand Down Expand Up @@ -795,13 +821,10 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity
}
}

if c.opts.StartupScript != "" {
cwOptions.VMBuild = &store.VMBuild{
ForceJupyterInstall: true,
LifeCycleScriptAttr: &store.LifeCycleScriptAttr{
Script: c.opts.StartupScript,
},
}
// Apply build mode
err := applyBuildMode(cwOptions, c.opts)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}

workspace, err := c.store.CreateWorkspace(c.org.ID, cwOptions)
Expand All @@ -812,6 +835,95 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity
return workspace, nil
}

func validateBuildMode(mode, containerImage, composeFile string) error {
switch mode {
case "vm", "k8s", "container", "compose":
// valid
default:
return breverrors.NewValidationError(fmt.Sprintf("invalid mode %q: must be one of vm, k8s, container, compose", mode))
}
if mode == "container" && containerImage == "" {
return breverrors.NewValidationError("--container-image is required for container mode")
}
if mode == "compose" && composeFile == "" {
return breverrors.NewValidationError("--compose-file is required for compose mode")
}
return nil
}

// applyBuildMode configures the workspace options based on the selected build mode.
// The server determines the mode from which build field is present (vmBuild, customContainer,
// dockerCompose) — we do NOT send vmOnlyMode or onContainer, matching the UI behavior.
func applyBuildMode(cwOptions *store.CreateWorkspacesOptions, opts GPUCreateOptions) error {
mode := opts.Mode
if mode == "" {
mode = "vm"
}

switch mode {
case "vm":
jupyter := true
if opts.JupyterSet {
jupyter = opts.Jupyter
}
cwOptions.VMBuild = &store.VMBuild{
ForceJupyterInstall: jupyter,
}
if opts.StartupScript != "" {
cwOptions.VMBuild.LifeCycleScriptAttr = &store.LifeCycleScriptAttr{
Script: opts.StartupScript,
}
}

case "k8s":
jupyter := false // UI defaults to false for k8s
if opts.JupyterSet {
jupyter = opts.Jupyter
}
cwOptions.VMBuild = &store.VMBuild{
ForceJupyterInstall: jupyter,
K8s: &store.K8sConfig{
IsDisabled: false,
IsDashboardEnabled: true,
},
}
if opts.StartupScript != "" {
cwOptions.VMBuild.LifeCycleScriptAttr = &store.LifeCycleScriptAttr{
Script: opts.StartupScript,
}
}

case "container":
cwOptions.VMBuild = nil
cwOptions.CustomContainer = &store.CustomContainer{
ContainerURL: opts.ContainerImage,
}

case "compose":
cwOptions.VMBuild = nil
composeConfig := &store.DockerCompose{}

if strings.HasPrefix(opts.ComposeFile, "http://") || strings.HasPrefix(opts.ComposeFile, "https://") {
composeConfig.FileURL = opts.ComposeFile
} else {
content, err := os.ReadFile(opts.ComposeFile)
if err != nil {
return breverrors.WrapAndTrace(fmt.Errorf("could not read compose file %s: %w", opts.ComposeFile, err))
}
composeConfig.YamlString = string(content)
}

jupyter := false
if opts.JupyterSet {
jupyter = opts.Jupyter
}
composeConfig.JupyterInstall = jupyter
cwOptions.DockerCompose = composeConfig
}

return nil
}

// 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
47 changes: 45 additions & 2 deletions pkg/store/workspace.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,36 @@ type LifeCycleScriptAttr struct {
type VMBuild struct {
ForceJupyterInstall bool `json:"forceJupyterInstall,omitempty"`
LifeCycleScriptAttr *LifeCycleScriptAttr `json:"lifeCycleScriptAttr,omitempty"`
K8s *K8sConfig `json:"k8s,omitempty"`
}

// K8sConfig holds Kubernetes configuration for VM builds
type K8sConfig struct {
IsDisabled bool `json:"isDisabled"`
IsDashboardEnabled bool `json:"isDashboardEnabled"`
}

// CustomContainer holds custom container build configuration
type CustomContainer struct {
ContainerURL string `json:"containerUrl"`
EntryPoint string `json:"entryPoint"`
Registry *Registry `json:"registry,omitempty"`
}

// DockerCompose holds Docker Compose build configuration
type DockerCompose struct {
FileURL string `json:"fileUrl,omitempty"`
YamlString string `json:"yamlString,omitempty"`
JupyterInstall bool `json:"jupyterInstall"`
EnvironmentVariables map[string]string `json:"environmentVariables,omitempty"`
Registries []*Registry `json:"registries,omitempty"`
}

// Registry holds container registry credentials
type Registry struct {
Username string `json:"username"`
Password string `json:"password"`
Url string `json:"url"`
}

type CreateWorkspacesOptions struct {
Expand All @@ -69,6 +99,9 @@ type CreateWorkspacesOptions struct {
BaseImage string `json:"baseImage"`
VMOnlyMode bool `json:"vmOnlyMode"`
VMBuild *VMBuild `json:"vmBuild,omitempty"`
CustomContainer *CustomContainer `json:"customContainer,omitempty"`
DockerCompose *DockerCompose `json:"dockerCompose,omitempty"`
OnContainer bool `json:"onContainer,omitempty"`
PortMappings map[string]string `json:"portMappings"`
Files interface{} `json:"files"`
Labels interface{} `json:"labels"`
Expand Down Expand Up @@ -114,8 +147,8 @@ func NewCreateWorkspacesOptions(clusterID, name string) *CreateWorkspacesOptions
Name: name,
PortMappings: map[string]string{},
ReposV1: &entity.ReposV1{},
VMOnlyMode: true,
WorkspaceGroupID: "GCP",
VMBuild: &VMBuild{ForceJupyterInstall: true},
WorkspaceGroupID: "", // resolved dynamically from instance type
WorkspaceTemplateID: DefaultWorkspaceTemplateID,
WorkspaceVersion: "v1",
}
Expand Down Expand Up @@ -156,6 +189,16 @@ func (c *CreateWorkspacesOptions) WithWorkspaceClassID(workspaceClassID string)
return c
}

func (c *CreateWorkspacesOptions) WithVMBuild(vmBuild *VMBuild) *CreateWorkspacesOptions {
c.VMBuild = vmBuild
return c
}

func (c *CreateWorkspacesOptions) WithWorkspaceGroupID(workspaceGroupID string) *CreateWorkspacesOptions {
c.WorkspaceGroupID = workspaceGroupID
return c
}

func (s AuthHTTPStore) CreateWorkspace(organizationID string, options *CreateWorkspacesOptions) (*entity.Workspace, error) {
if options == nil {
return nil, fmt.Errorf("options can not be nil")
Expand Down
Loading