Skip to content

Commit 9b3a566

Browse files
DavidoniumDavidonium
authored
feat: add support for gpu and multi-az (#13)
## Add multi-AZ and GPU support to cluster package commands Surfaces two fields from the booking API that were previously not exposed in the CLI. **`cluster package list`** - New GPU column showing the GPU resource value (e.g. `1000m`), or `n/a` if the package has no GPU - New MULTI-AZ column showing `yes`/`no` **`cluster create`** - New `--gpu` flag to select a package by GPU resource (same pattern as `--cpu`/`--ram`/`--disk`) - New `--multi-az` flag to filter packages to multi-AZ only; acts as a narrowing filter alongside at least one resource flag - Shell completion for `--gpu`, aware of other resource flags already set **`cluster describe`** - GPU resource allocation (base/reserved/available) is shown in the Resources section when present
1 parent 48c6e9f commit 9b3a566

9 files changed

Lines changed: 635 additions & 23 deletions

File tree

‎internal/cmd/cluster/completion.go‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,182 @@ func versionCompletion(s *state.State) func(*cobra.Command, []string, string) ([
128128
}
129129
}
130130

131+
// filteredPackages fetches active packages filtered by non-empty cpu/ram/disk/gpu values and the multiAz flag.
132+
// Returns nil (no completions) if --cloud-provider is not set.
133+
func filteredPackages(cmd *cobra.Command, s *state.State, cpu, ram, disk, gpu string, multiAz bool) ([]*bookingv1.Package, error) {
134+
provider, _ := cmd.Flags().GetString("cloud-provider")
135+
if provider == "" {
136+
return nil, nil
137+
}
138+
139+
ctx := cmd.Context()
140+
client, err := s.Client(ctx)
141+
if err != nil {
142+
return nil, err
143+
}
144+
145+
accountID, err := s.AccountID()
146+
if err != nil {
147+
return nil, err
148+
}
149+
150+
region, _ := cmd.Flags().GetString("cloud-region")
151+
req := &bookingv1.ListPackagesRequest{
152+
AccountId: accountID,
153+
CloudProviderId: provider,
154+
Statuses: []bookingv1.PackageStatus{bookingv1.PackageStatus_PACKAGE_STATUS_ACTIVE},
155+
}
156+
if region != "" {
157+
req.CloudProviderRegionId = &region
158+
}
159+
if multiAz {
160+
req.MultiAz = new(true)
161+
}
162+
163+
resp, err := client.Booking().ListPackages(ctx, req)
164+
if err != nil {
165+
return nil, err
166+
}
167+
168+
var result []*bookingv1.Package
169+
for _, p := range resp.GetItems() {
170+
rc := p.GetResourceConfiguration()
171+
if cpu != "" && rc.GetCpu() != cpu {
172+
continue
173+
}
174+
if ram != "" && rc.GetRam() != ram {
175+
continue
176+
}
177+
if disk != "" && rc.GetDisk() != disk {
178+
continue
179+
}
180+
if gpu != "" && rc.GetGpu() != gpu {
181+
continue
182+
}
183+
result = append(result, p)
184+
}
185+
return result, nil
186+
}
187+
188+
// cpuCompletion returns a completion function for the --cpu flag.
189+
func cpuCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
190+
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
191+
provider, _ := cmd.Flags().GetString("cloud-provider")
192+
if provider == "" {
193+
return nil, cobra.ShellCompDirectiveNoFileComp
194+
}
195+
ram, _ := cmd.Flags().GetString("ram")
196+
disk, _ := cmd.Flags().GetString("disk")
197+
gpu, _ := cmd.Flags().GetString("gpu")
198+
multiAz, _ := cmd.Flags().GetBool("multi-az")
199+
pkgs, err := filteredPackages(cmd, s, "", ram, disk, gpu, multiAz)
200+
if err != nil {
201+
return nil, cobra.ShellCompDirectiveError
202+
}
203+
204+
seen := make(map[string]struct{})
205+
var completions []string
206+
for _, p := range pkgs {
207+
v := p.GetResourceConfiguration().GetCpu()
208+
if _, ok := seen[v]; !ok {
209+
seen[v] = struct{}{}
210+
completions = append(completions, v)
211+
}
212+
}
213+
return completions, cobra.ShellCompDirectiveNoFileComp
214+
}
215+
}
216+
217+
// ramCompletion returns a completion function for the --ram flag.
218+
func ramCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
219+
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
220+
provider, _ := cmd.Flags().GetString("cloud-provider")
221+
if provider == "" {
222+
return nil, cobra.ShellCompDirectiveNoFileComp
223+
}
224+
cpu, _ := cmd.Flags().GetString("cpu")
225+
disk, _ := cmd.Flags().GetString("disk")
226+
gpu, _ := cmd.Flags().GetString("gpu")
227+
multiAz, _ := cmd.Flags().GetBool("multi-az")
228+
pkgs, err := filteredPackages(cmd, s, cpu, "", disk, gpu, multiAz)
229+
if err != nil {
230+
return nil, cobra.ShellCompDirectiveError
231+
}
232+
233+
seen := make(map[string]struct{})
234+
var completions []string
235+
for _, p := range pkgs {
236+
v := p.GetResourceConfiguration().GetRam()
237+
if _, ok := seen[v]; !ok {
238+
seen[v] = struct{}{}
239+
completions = append(completions, v)
240+
}
241+
}
242+
return completions, cobra.ShellCompDirectiveNoFileComp
243+
}
244+
}
245+
246+
// diskCompletion returns a completion function for the --disk flag.
247+
func diskCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
248+
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
249+
provider, _ := cmd.Flags().GetString("cloud-provider")
250+
if provider == "" {
251+
return nil, cobra.ShellCompDirectiveNoFileComp
252+
}
253+
cpu, _ := cmd.Flags().GetString("cpu")
254+
ram, _ := cmd.Flags().GetString("ram")
255+
gpu, _ := cmd.Flags().GetString("gpu")
256+
multiAz, _ := cmd.Flags().GetBool("multi-az")
257+
pkgs, err := filteredPackages(cmd, s, cpu, ram, "", gpu, multiAz)
258+
if err != nil {
259+
return nil, cobra.ShellCompDirectiveError
260+
}
261+
262+
seen := make(map[string]struct{})
263+
var completions []string
264+
for _, p := range pkgs {
265+
v := p.GetResourceConfiguration().GetDisk()
266+
if _, ok := seen[v]; !ok {
267+
seen[v] = struct{}{}
268+
completions = append(completions, v)
269+
}
270+
}
271+
return completions, cobra.ShellCompDirectiveNoFileComp
272+
}
273+
}
274+
275+
// gpuCompletion returns a completion function for the --gpu flag.
276+
func gpuCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
277+
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
278+
provider, _ := cmd.Flags().GetString("cloud-provider")
279+
if provider == "" {
280+
return nil, cobra.ShellCompDirectiveNoFileComp
281+
}
282+
cpu, _ := cmd.Flags().GetString("cpu")
283+
ram, _ := cmd.Flags().GetString("ram")
284+
disk, _ := cmd.Flags().GetString("disk")
285+
multiAz, _ := cmd.Flags().GetBool("multi-az")
286+
pkgs, err := filteredPackages(cmd, s, cpu, ram, disk, "", multiAz)
287+
if err != nil {
288+
return nil, cobra.ShellCompDirectiveError
289+
}
290+
291+
seen := make(map[string]struct{})
292+
var completions []string
293+
for _, p := range pkgs {
294+
v := p.GetResourceConfiguration().GetGpu()
295+
if v == "" {
296+
continue
297+
}
298+
if _, ok := seen[v]; !ok {
299+
seen[v] = struct{}{}
300+
completions = append(completions, v)
301+
}
302+
}
303+
return completions, cobra.ShellCompDirectiveNoFileComp
304+
}
305+
}
306+
131307
// packageCompletion returns a completion function for the --package flag.
132308
func packageCompletion(s *state.State) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
133309
return func(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {

‎internal/cmd/cluster/completion_test.go‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cluster_test
22

33
import (
4+
"strings"
45
"testing"
56

67
"github.com/stretchr/testify/assert"
@@ -124,6 +125,99 @@ func TestVersionCompletion(t *testing.T) {
124125
assert.NotContains(t, stdout, "1.12.0")
125126
}
126127

128+
func TestCPUCompletion(t *testing.T) {
129+
env := testutil.NewTestEnv(t)
130+
131+
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
132+
Items: []*bookingv1.Package{
133+
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
134+
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
135+
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "2GiB", Disk: "200GiB"}},
136+
},
137+
}, nil)
138+
139+
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "")
140+
require.NoError(t, err)
141+
assert.Contains(t, stdout, "500m")
142+
assert.Contains(t, stdout, "1000m")
143+
// Deduplication: 1000m appears twice in packages but once in completions.
144+
count := strings.Count(stdout, "1000m")
145+
assert.Equal(t, 1, count, "1000m should appear only once")
146+
}
147+
148+
func TestCPUCompletion_FilteredByRAM(t *testing.T) {
149+
env := testutil.NewTestEnv(t)
150+
151+
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
152+
Items: []*bookingv1.Package{
153+
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
154+
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
155+
},
156+
}, nil)
157+
158+
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--ram", "1GiB", "--cpu", "")
159+
require.NoError(t, err)
160+
assert.Contains(t, stdout, "1000m")
161+
assert.NotContains(t, stdout, "500m")
162+
}
163+
164+
func TestRAMCompletion_FilteredByCPU(t *testing.T) {
165+
env := testutil.NewTestEnv(t)
166+
167+
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
168+
Items: []*bookingv1.Package{
169+
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
170+
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
171+
},
172+
}, nil)
173+
174+
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "500m", "--ram", "")
175+
require.NoError(t, err)
176+
assert.Contains(t, stdout, "512MiB")
177+
assert.NotContains(t, stdout, "1GiB")
178+
}
179+
180+
func TestDiskCompletion_FilteredByCPUAndRAM(t *testing.T) {
181+
env := testutil.NewTestEnv(t)
182+
183+
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
184+
Items: []*bookingv1.Package{
185+
{Id: "pkg-1", Name: "starter", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m", Ram: "512MiB", Disk: "50GiB"}},
186+
{Id: "pkg-2", Name: "business", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "100GiB"}},
187+
{Id: "pkg-3", Name: "enterprise", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Ram: "1GiB", Disk: "200GiB"}},
188+
},
189+
}, nil)
190+
191+
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--cpu", "1000m", "--ram", "1GiB", "--disk", "")
192+
require.NoError(t, err)
193+
assert.Contains(t, stdout, "100GiB")
194+
assert.Contains(t, stdout, "200GiB")
195+
assert.NotContains(t, stdout, "50GiB")
196+
}
197+
198+
func TestGPUCompletion(t *testing.T) {
199+
env := testutil.NewTestEnv(t)
200+
201+
env.BookingServer.ListPackagesCalls.Returns(&bookingv1.ListPackagesResponse{
202+
Items: []*bookingv1.Package{
203+
{Id: "pkg-1", Name: "gpu-small", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "1000m", Gpu: new("1000m")}},
204+
{Id: "pkg-2", Name: "gpu-large", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "2000m", Gpu: new("2000m")}},
205+
{Id: "pkg-3", Name: "cpu-only", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "500m"}},
206+
{Id: "pkg-4", Name: "gpu-dup", ResourceConfiguration: &bookingv1.ResourceConfiguration{Cpu: "4000m", Gpu: new("1000m")}},
207+
},
208+
}, nil)
209+
210+
stdout, _, err := testutil.Exec(t, env, "__complete", "cluster", "create", "--cloud-provider", "aws", "--gpu", "")
211+
require.NoError(t, err)
212+
assert.Contains(t, stdout, "1000m")
213+
assert.Contains(t, stdout, "2000m")
214+
// Packages without GPU should not appear.
215+
assert.NotContains(t, stdout, "500m")
216+
// Deduplication: 1000m appears in two packages but once in completions.
217+
count := strings.Count(stdout, "1000m")
218+
assert.Equal(t, 1, count, "1000m should appear only once")
219+
}
220+
127221
func TestVersionCompletion_UnavailableExcluded(t *testing.T) {
128222
env := testutil.NewTestEnv(t)
129223

‎internal/cmd/cluster/create.go‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,19 @@ func newCreateCommand(s *state.State) *cobra.Command {
2727
cmd.Flags().String("cloud-region", "", "Cloud provider region ID (required, see 'cluster cloud-region list --cloud-provider <provider_id>)")
2828
cmd.Flags().String("version", "", "Qdrant version (default latest)")
2929
cmd.Flags().Uint32("nodes", 1, "Number of nodes (default 1)")
30-
cmd.Flags().String("package", "", "Booking package name or ID (required, see 'cluster package list')")
30+
cmd.Flags().String("package", "", "Booking package name or ID (see 'cluster package list')")
31+
cmd.Flags().String("cpu", "", "CPU to select a package (e.g. \"1000m\"); must match a value from 'cluster package list'")
32+
cmd.Flags().String("ram", "", "RAM to select a package (e.g. \"1GiB\"); must match a value from 'cluster package list'")
33+
cmd.Flags().String("disk", "", "Disk size to select a package (e.g. \"100GiB\"); must match a value from 'cluster package list'")
34+
cmd.Flags().String("gpu", "", "GPU resource to select a package (e.g. \"1000m\"); must match a value from 'cluster package list'")
35+
cmd.Flags().Bool("multi-az", false, "Require a multi-AZ package")
3136
cmd.Flags().StringToString("label", nil, "Label to apply to the cluster ('key=value'), can be specified multiple times")
3237
cmd.Flags().Bool("wait", false, "Wait for the cluster to become healthy")
3338
cmd.Flags().Duration("wait-timeout", 10*time.Minute, "Maximum time to wait for cluster health")
3439
cmd.Flags().Duration("wait-poll-interval", 5*time.Second, "How often to poll for cluster health")
3540
_ = cmd.Flags().MarkHidden("wait-poll-interval")
3641
_ = cmd.MarkFlagRequired("cloud-provider")
3742
_ = cmd.MarkFlagRequired("cloud-region")
38-
_ = cmd.MarkFlagRequired("package")
3943
return cmd
4044
},
4145
Run: func(s *state.State, cmd *cobra.Command, args []string) (*clusterv1.Cluster, error) {
@@ -65,8 +69,17 @@ func newCreateCommand(s *state.State) *cobra.Command {
6569
version, _ := cmd.Flags().GetString("version")
6670
nodes, _ := cmd.Flags().GetUint32("nodes")
6771
packageValue, _ := cmd.Flags().GetString("package")
72+
cpu, _ := cmd.Flags().GetString("cpu")
73+
ram, _ := cmd.Flags().GetString("ram")
74+
disk, _ := cmd.Flags().GetString("disk")
75+
gpu, _ := cmd.Flags().GetString("gpu")
76+
multiAz, _ := cmd.Flags().GetBool("multi-az")
6877
labelMap, _ := cmd.Flags().GetStringToString("label")
6978

79+
if packageValue == "" && cpu == "" && ram == "" && disk == "" && gpu == "" {
80+
return nil, fmt.Errorf("either --package or at least one of --cpu, --ram, --disk, --gpu is required")
81+
}
82+
7083
var packageID string
7184
if packageValue != "" {
7285
if isUUID(packageValue) {
@@ -78,6 +91,12 @@ func newCreateCommand(s *state.State) *cobra.Command {
7891
}
7992
packageID = pkg.GetId()
8093
}
94+
} else {
95+
pkg, err := resolvePackageByResources(ctx, client.Booking(), accountID, cloudProvider, cloudRegion, cpu, ram, disk, gpu, multiAz)
96+
if err != nil {
97+
return nil, err
98+
}
99+
packageID = pkg.GetId()
81100
}
82101

83102
cluster := &clusterv1.Cluster{
@@ -130,5 +149,9 @@ func newCreateCommand(s *state.State) *cobra.Command {
130149
_ = cmd.RegisterFlagCompletionFunc("cloud-region", cloudRegionCompletion(s))
131150
_ = cmd.RegisterFlagCompletionFunc("package", packageCompletion(s))
132151
_ = cmd.RegisterFlagCompletionFunc("version", versionCompletion(s))
152+
_ = cmd.RegisterFlagCompletionFunc("cpu", cpuCompletion(s))
153+
_ = cmd.RegisterFlagCompletionFunc("ram", ramCompletion(s))
154+
_ = cmd.RegisterFlagCompletionFunc("disk", diskCompletion(s))
155+
_ = cmd.RegisterFlagCompletionFunc("gpu", gpuCompletion(s))
133156
return cmd
134157
}

0 commit comments

Comments
 (0)