Skip to content

feat: add --gpu flag on cluster create that filters packages properly #22

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

Merged
Davidonium merged 1 commit into from
Mar 18, 2026
Merged

feat: add --gpu flag on cluster create that filters packages properly #22

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
2 changes: 1 addition & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ require (
github.com/creativeprojects/go-selfupdate v1.5.2
github.com/dustin/go-humanize v1.0.1
github.com/jedib0t/go-pretty/v6 v6.7.8
github.com/qdrant/qdrant-cloud-public-api v0.110.0
github.com/qdrant/qdrant-cloud-public-api v0.111.0
github.com/robfig/cron/v3 v3.0.1
github.com/spf13/cobra v1.10.2
github.com/spf13/pflag v1.0.10
Expand Down
4 changes: 2 additions & 2 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0
github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/qdrant/qdrant-cloud-public-api v0.110.0 h1:WFA9PB1UIRfciW5mdshwmKIo46vAuRI1vhBznmfOJp4=
github.com/qdrant/qdrant-cloud-public-api v0.110.0/go.mod h1:XqgxNTOYlCx1d6UKbPjttSyORF3WCvZefOpYTlBcJ8g=
github.com/qdrant/qdrant-cloud-public-api v0.111.0 h1:9sB4+2i9NVNihr8DtbZf1Pw1fcd6FRqj7JutLR/dfL8=
github.com/qdrant/qdrant-cloud-public-api v0.111.0/go.mod h1:XqgxNTOYlCx1d6UKbPjttSyORF3WCvZefOpYTlBcJ8g=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
Expand Down
142 changes: 130 additions & 12 deletions internal/cmd/cluster/completion.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,9 +128,18 @@ func versionCompletion(s *state.State) func(*cobra.Command, []string, string) ([
}
}

// filteredPackages fetches active packages filtered by non-empty cpu/ram values and the multiAz flag.
// packageFilter holds the parameters for filtering packages.
type packageFilter struct {
CPU string
RAM string
GPU string
IncludeGPU bool
MultiAz bool
}

// filteredPackages fetches active packages matching the given filter.
// Returns nil (no completions) if --cloud-provider is not set.
func filteredPackages(cmd *cobra.Command, s *state.State, cpu, ram string, multiAz bool) ([]*bookingv1.Package, error) {
func filteredPackages(cmd *cobra.Command, s *state.State, f packageFilter) ([]*bookingv1.Package, error) {
provider, _ := cmd.Flags().GetString("cloud-provider")
if provider == "" {
return nil, nil
Expand All @@ -152,11 +161,12 @@ func filteredPackages(cmd *cobra.Command, s *state.State, cpu, ram string, multi
AccountId: accountID,
CloudProviderId: provider,
Statuses: []bookingv1.PackageStatus{bookingv1.PackageStatus_PACKAGE_STATUS_ACTIVE},
Gpu: &f.IncludeGPU,
}
if region != "" {
req.CloudProviderRegionId = &region
}
if multiAz {
if f.MultiAz {
req.MultiAz = new(true)
}

Expand All @@ -165,14 +175,37 @@ func filteredPackages(cmd *cobra.Command, s *state.State, cpu, ram string, multi
return nil, err
}

var wantCPU, wantRAM, wantGPU int64
if f.CPU != "" {
wantCPU, _ = parseCPUMillicores(f.CPU)
}
if f.RAM != "" {
wantRAM, _ = parseRAMGiB(f.RAM)
}
if f.GPU != "" {
wantGPU, _ = parseGPUMillicores(f.GPU)
}

var result []*bookingv1.Package
for _, p := range resp.GetItems() {
rc := p.GetResourceConfiguration()
if cpu != "" && rc.GetCpu() != cpu {
continue
if f.CPU != "" {
pkgCPU, _ := parseCPUMillicores(rc.GetCpu())
if pkgCPU != wantCPU {
continue
}
}
if ram != "" && rc.GetRam() != ram {
continue
if f.RAM != "" {
pkgRAM, _ := parseRAMGiB(rc.GetRam())
if pkgRAM != wantRAM {
continue
}
}
if f.GPU != "" {
pkgGPU, _ := parseGPUMillicores(rc.GetGpu())
if pkgGPU != wantGPU {
continue
}
}
result = append(result, p)
}
Expand All @@ -187,16 +220,31 @@ func cpuCompletion(s *state.State) func(*cobra.Command, []string, string) ([]str
return nil, cobra.ShellCompDirectiveNoFileComp
}
ram, _ := cmd.Flags().GetString("ram")
gpu, err := gpuFlagToMillicores(cmd)
if err != nil {
cobra.CompErrorln(err.Error())
return nil, cobra.ShellCompDirectiveError
}
multiAz, _ := cmd.Flags().GetBool("multi-az")
pkgs, err := filteredPackages(cmd, s, "", ram, multiAz)
filter := packageFilter{
RAM: ram,
GPU: gpu,
IncludeGPU: gpu != "",
MultiAz: multiAz,
}
pkgs, err := filteredPackages(cmd, s, filter)
if err != nil {
return nil, cobra.ShellCompDirectiveError
}

seen := make(map[string]struct{})
var completions []string
for _, p := range pkgs {
v := p.GetResourceConfiguration().GetCpu()
v, err := normalizeMillicores(p.GetResourceConfiguration().GetCpu())
if err != nil {
cobra.CompErrorln(fmt.Sprintf("package %s: %v", p.GetName(), err))
continue
}
if _, ok := seen[v]; !ok {
seen[v] = struct{}{}
completions = append(completions, v)
Expand All @@ -214,16 +262,31 @@ func ramCompletion(s *state.State) func(*cobra.Command, []string, string) ([]str
return nil, cobra.ShellCompDirectiveNoFileComp
}
cpu, _ := cmd.Flags().GetString("cpu")
gpu, err := gpuFlagToMillicores(cmd)
if err != nil {
cobra.CompErrorln(err.Error())
return nil, cobra.ShellCompDirectiveError
}
multiAz, _ := cmd.Flags().GetBool("multi-az")
pkgs, err := filteredPackages(cmd, s, cpu, "", multiAz)
filter := packageFilter{
CPU: cpu,
GPU: gpu,
IncludeGPU: gpu != "",
MultiAz: multiAz,
}
pkgs, err := filteredPackages(cmd, s, filter)
if err != nil {
return nil, cobra.ShellCompDirectiveError
}

seen := make(map[string]struct{})
var completions []string
for _, p := range pkgs {
v := p.GetResourceConfiguration().GetRam()
v, err := normalizeRAM(p.GetResourceConfiguration().GetRam())
if err != nil {
cobra.CompErrorln(fmt.Sprintf("package %s: %v", p.GetName(), err))
continue
}
if _, ok := seen[v]; !ok {
seen[v] = struct{}{}
completions = append(completions, v)
Expand All @@ -242,8 +305,20 @@ func diskCompletion(s *state.State) func(*cobra.Command, []string, string) ([]st
}
cpu, _ := cmd.Flags().GetString("cpu")
ram, _ := cmd.Flags().GetString("ram")
gpu, err := gpuFlagToMillicores(cmd)
if err != nil {
cobra.CompErrorln(err.Error())
return nil, cobra.ShellCompDirectiveError
}
multiAz, _ := cmd.Flags().GetBool("multi-az")
pkgs, err := filteredPackages(cmd, s, cpu, ram, multiAz)
filter := packageFilter{
CPU: cpu,
RAM: ram,
GPU: gpu,
IncludeGPU: gpu != "",
MultiAz: multiAz,
}
pkgs, err := filteredPackages(cmd, s, filter)
if err != nil {
return nil, cobra.ShellCompDirectiveError
}
Expand All @@ -261,6 +336,49 @@ func diskCompletion(s *state.State) func(*cobra.Command, []string, string) ([]st
}
}

// gpuCompletion returns a completion function for the --gpu flag.
func gpuCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
provider, _ := cmd.Flags().GetString("cloud-provider")
if provider == "" {
return nil, cobra.ShellCompDirectiveNoFileComp
}
cpu, _ := cmd.Flags().GetString("cpu")
ram, _ := cmd.Flags().GetString("ram")
multiAz, _ := cmd.Flags().GetBool("multi-az")
filter := packageFilter{
CPU: cpu,
RAM: ram,
IncludeGPU: true,
MultiAz: multiAz,
}
pkgs, err := filteredPackages(cmd, s, filter)
if err != nil {
return nil, cobra.ShellCompDirectiveError
}

seen := make(map[string]struct{})
var completions []string
for _, p := range pkgs {
v := p.GetResourceConfiguration().GetGpu()
if v == "" {
continue
}
// Convert millicores (e.g. "1000m") to integer GPUs (e.g. "1")
// to match the --gpu flag's Int type.
intVal := gpuMillicoresToCount(v)
if intVal == "" {
continue
}
if _, ok := seen[intVal]; !ok {
seen[intVal] = struct{}{}
completions = append(completions, intVal)
}
}
return completions, cobra.ShellCompDirectiveNoFileComp
}
}

// packageCompletion returns a completion function for the --package flag.
func packageCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
Expand Down
116 changes: 102 additions & 14 deletions internal/cmd/cluster/completion_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,9 +130,9 @@ func TestCPUCompletion(t *testing.T) {

env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "2GiB", Disk: "200GiB"}},
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "1Gi", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "2Gi", Disk: "100GiB"}},
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "200GiB"}},
},
}, nil)

Expand All @@ -150,12 +150,12 @@ func TestCPUCompletion_FilteredByRAM(t *testing.T) {

env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "2Gi", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "100GiB"}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--ram", "1GiB", "--cpu", "")
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--ram", "4GiB", "--cpu", "")
require.NoError(t, err)
assert.Contains(t, stdout, "1000m")
assert.NotContains(t, stdout, "500m")
Expand All @@ -166,35 +166,123 @@ func TestRAMCompletion_FilteredByCPU(t *testing.T) {

env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "2Gi", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "100GiB"}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "500m", "--ram", "")
require.NoError(t, err)
assert.Contains(t, stdout, "512MiB")
assert.NotContains(t, stdout, "1GiB")
assert.Contains(t, stdout, "2GiB")
assert.NotContains(t, stdout, "4GiB")
}

func TestDiskCompletion_FilteredByCPUAndRAM(t *testing.T) {
env := testutil.NewTestEnv(t)

env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "200GiB"}},
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "2Gi", Disk: "50GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "100GiB"}},
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "200GiB"}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "1000m", "--ram", "1GiB", "--disk", "")
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "1000m", "--ram", "4GiB", "--disk", "")
require.NoError(t, err)
assert.Contains(t, stdout, "100GiB")
assert.Contains(t, stdout, "200GiB")
assert.NotContains(t, stdout, "50GiB")
}

func TestGPUCompletion(t *testing.T) {
env := testutil.NewTestEnv(t)

gpu1 := "1000m"
gpu2 := "2000m"
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-gpu-1", Name: "gpu-starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "4000m", Ram: "16GiB", Disk: "100GiB", Gpu: &gpu1}},
{Id: "pkg-gpu-2", Name: "gpu-pro", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "8000m", Ram: "32GiB", Disk: "200GiB", Gpu: &gpu2}},
{Id: "pkg-gpu-3", Name: "gpu-starter-dup", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "4000m", Ram: "16GiB", Disk: "200GiB", Gpu: &gpu1}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--gpu", "")
require.NoError(t, err)
// Should return integer counts, not millicores.
assert.Contains(t, stdout, "1")
assert.Contains(t, stdout, "2")
assert.NotContains(t, stdout, "1000m")
assert.NotContains(t, stdout, "2000m")
// Deduplication: "1" appears in two packages but once in completions.
lines := strings.Split(stdout, "\n")
count := 0
for _, l := range lines {
if l == "1" {
count++
}
}
assert.Equal(t, 1, count, "GPU count 1 should appear only once")
}

func TestGPUCompletion_NoProvider(t *testing.T) {
env := testutil.NewTestEnv(t)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--gpu", "")
require.NoError(t, err)
assert.NotContains(t, stdout, "1")
}

func TestRAMCompletion_APIFormatMismatch(t *testing.T) {
env := testutil.NewTestEnv(t)

// API returns "4Gi" but filter uses "4GiB" -- numeric comparison should still match.
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "100GiB"}},
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "2000m", Ram: "8Gi", Disk: "200GiB"}},
},
}, nil)

// Filter by ram=4GiB should match the API's "4Gi" via numeric parsing.
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--ram", "4GiB", "--cpu", "")
require.NoError(t, err)
assert.Contains(t, stdout, "1000m")
assert.NotContains(t, stdout, "2000m")
}

func TestCPUCompletion_NormalizedOutput(t *testing.T) {
env := testutil.NewTestEnv(t)

// API returns "1" for CPU -- completion should normalize to "1000m".
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1", Ram: "4Gi", Disk: "100GiB"}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "")
require.NoError(t, err)
assert.Contains(t, stdout, "1000m")
}

func TestRAMCompletion_NormalizedOutput(t *testing.T) {
env := testutil.NewTestEnv(t)

// API returns "4Gi" -- completion should normalize to "4GiB".
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
Items: []*bookingv1.Package{
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "4Gi", Disk: "100GiB"}},
},
}, nil)

stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--ram", "")
require.NoError(t, err)
assert.Contains(t, stdout, "4GiB")
assert.NotContains(t, stdout, "4Gi\n")
}

func TestVersionCompletion_UnavailableExcluded(t *testing.T) {
env := testutil.NewTestEnv(t)

Expand Down
Loading
Loading