Skip to content

Commit a8d61ba

Browse files
patelspratikpatelspratik
committed
initial uptake
1 parent 5e6c288 commit a8d61ba

6 files changed

Lines changed: 245 additions & 30 deletions

File tree

‎pkg/cmd/gpucreate/gpucreate.go‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ type GPUCreateStore interface {
105105
GetWorkspace(workspaceID string) (*entity.Workspace, error)
106106
CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error)
107107
DeleteWorkspace(workspaceID string) (*entity.Workspace, error)
108-
GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error)
108+
GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error)
109109
GetLaunchable(launchableID string) (*store.LaunchableResponse, error)
110110
GetLaunchableLifeCycleScript(launchableID, scriptID string) (*store.LifeCycleScriptResponse, error)
111111
RedeemCouponCode(organizationID string, code string) (*store.RedeemCouponCodeResponse, error)
@@ -768,11 +768,11 @@ func newCreateContext(t *terminal.Terminal, store GPUCreateStore, opts GPUCreate
768768
}
769769
ctx.org = org
770770

771-
// Fetch instance types with workspace groups
772-
allInstanceTypes, err := store.GetAllInstanceTypesWithWorkspaceGroups(org.ID)
771+
// Fetch instance types with cloud credentials.
772+
allInstanceTypes, err := store.GetAllInstanceTypesWithCloudCreds(org.ID)
773773
if err != nil {
774-
ctx.logf("Warning: could not fetch instance types with workspace groups: %s\n", err.Error())
775-
ctx.logf("Falling back to default workspace group\n")
774+
ctx.logf("Warning: could not fetch instance types with cloud credentials: %s\n", err.Error())
775+
ctx.logf("Falling back to default cloud credential\n")
776776
}
777777
ctx.allInstanceTypes = allInstanceTypes
778778

@@ -791,7 +791,7 @@ func (c *createContext) validateInstanceTypeAvailability(instanceType string) er
791791
if c.allInstanceTypes == nil {
792792
return nil
793793
}
794-
if c.allInstanceTypes.GetWorkspaceGroupID(instanceType) != "" {
794+
if c.allInstanceTypes.GetCloudCredID(instanceType) != "" {
795795
return nil
796796
}
797797
if !c.allInstanceTypes.HasInstanceType(instanceType) {
@@ -1042,8 +1042,8 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity
10421042
}
10431043

10441044
if c.allInstanceTypes != nil {
1045-
if wgID := c.allInstanceTypes.GetWorkspaceGroupID(spec.Type); wgID != "" {
1046-
cwOptions.WorkspaceGroupID = wgID
1045+
if cloudCredID := c.allInstanceTypes.GetCloudCredID(spec.Type); cloudCredID != "" {
1046+
cwOptions.WithCloudCredID(cloudCredID)
10471047
}
10481048
}
10491049

@@ -1060,7 +1060,7 @@ func (c *createContext) createWorkspace(name string, spec InstanceSpec) (*entity
10601060
if cwOptions.WorkspaceGroupID == "" {
10611061
if c.allInstanceTypes == nil {
10621062
return nil, breverrors.NewValidationError(fmt.Sprintf(
1063-
"could not resolve workspace group for %q (instance-type listing was unavailable); please retry",
1063+
"could not resolve cloud credential for %q (instance-type listing was unavailable); please retry",
10641064
spec.Type,
10651065
))
10661066
}
@@ -1178,9 +1178,15 @@ func applyLaunchableConfig(cwOptions *store.CreateWorkspacesOptions, launchableI
11781178

11791179
wsReq := info.CreateWorkspaceRequest
11801180

1181-
// Use launchable's workspace group if not already resolved from instance types
1181+
// Use launchable's cloud credential if not already resolved from instance types.
1182+
if cwOptions.CloudCredID == "" && wsReq.CloudCredID != "" {
1183+
cwOptions.WithCloudCredID(wsReq.CloudCredID)
1184+
}
11821185
if cwOptions.WorkspaceGroupID == "" && wsReq.WorkspaceGroupID != "" {
11831186
cwOptions.WorkspaceGroupID = wsReq.WorkspaceGroupID
1187+
if cwOptions.CloudCredID == "" {
1188+
cwOptions.CloudCredID = wsReq.WorkspaceGroupID
1189+
}
11841190
}
11851191

11861192
// Location / sub-location
@@ -1240,6 +1246,7 @@ func applyLaunchableConfig(cwOptions *store.CreateWorkspacesOptions, launchableI
12401246
}
12411247
labels["launchableId"] = launchableID
12421248
labels["launchableInstanceType"] = wsReq.InstanceType
1249+
labels["cloudCredId"] = cwOptions.CloudCredID
12431250
labels["workspaceGroupId"] = cwOptions.WorkspaceGroupID
12441251
labels["launchableCreatedByUserId"] = info.CreatedByUserID
12451252
labels["launchableCreatedByOrgId"] = info.CreatedByOrgID

‎pkg/cmd/gpucreate/gpucreate_test.go‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/brevdev/brev-cli/pkg/store"
1515
"github.com/brevdev/brev-cli/pkg/terminal"
1616
"github.com/stretchr/testify/assert"
17+
"github.com/stretchr/testify/require"
1718
)
1819

1920
// MockGPUCreateStore is a mock implementation of GPUCreateStore for testing
@@ -24,6 +25,7 @@ type MockGPUCreateStore struct {
2425
CreateError error
2526
CreateErrorTypes map[string]error // Errors for specific instance types
2627
DeleteError error
28+
CreatedOptions []*store.CreateWorkspacesOptions
2729
CreatedWorkspaces []*entity.Workspace
2830
DeletedWorkspaceIDs []string
2931
FetchedLifeCycleScriptIDs []string
@@ -85,6 +87,7 @@ func (m *MockGPUCreateStore) CreateWorkspace(organizationID string, options *sto
8587
Status: entity.Running,
8688
}
8789
m.Workspaces[ws.ID] = ws
90+
m.CreatedOptions = append(m.CreatedOptions, options)
8891
m.CreatedWorkspaces = append(m.CreatedWorkspaces, ws)
8992
return ws, nil
9093
}
@@ -104,7 +107,7 @@ func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID strin
104107
return []entity.Workspace{}, nil
105108
}
106109

107-
func (m *MockGPUCreateStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) {
110+
func (m *MockGPUCreateStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) {
108111
return nil, nil
109112
}
110113

@@ -247,7 +250,8 @@ func TestApplyLaunchableConfig(t *testing.T) { //nolint:funlen // test
247250

248251
applyLaunchableConfig(cwOptions, "env-abc123", info)
249252

250-
// Workspace group from launchable
253+
// Cloud credential from launchable compatibility input.
254+
assert.Equal(t, "GCP", cwOptions.CloudCredID)
251255
assert.Equal(t, "GCP", cwOptions.WorkspaceGroupID)
252256
// Location / sub-location
253257
assert.Equal(t, "us-west1", cwOptions.Location)
@@ -274,13 +278,15 @@ func TestApplyLaunchableConfig(t *testing.T) { //nolint:funlen // test
274278
assert.True(t, ok)
275279
assert.Equal(t, "env-abc123", labels["launchableId"])
276280
assert.Equal(t, "n2-standard-4", labels["launchableInstanceType"])
281+
assert.Equal(t, "GCP", labels["cloudCredId"])
277282
assert.Equal(t, "GCP", labels["workspaceGroupId"])
278283
assert.Equal(t, "user-1", labels["launchableCreatedByUserId"])
279284
assert.Equal(t, "org-1", labels["launchableCreatedByOrgId"])
280285
})
281286

282287
t.Run("preserves existing workspace group", func(t *testing.T) {
283288
cwOptions := &store.CreateWorkspacesOptions{
289+
CloudCredID: "existing-wg",
284290
WorkspaceGroupID: "existing-wg",
285291
}
286292
info := &store.LaunchableResponse{
@@ -293,6 +299,7 @@ func TestApplyLaunchableConfig(t *testing.T) { //nolint:funlen // test
293299
applyLaunchableConfig(cwOptions, "env-abc", info)
294300

295301
assert.Equal(t, "existing-wg", cwOptions.WorkspaceGroupID)
302+
assert.Equal(t, "existing-wg", cwOptions.CloudCredID)
296303
})
297304

298305
t.Run("storage already has Gi suffix", func(t *testing.T) {
@@ -1116,6 +1123,34 @@ func TestCreateInstancesWithTypeSkipsUnavailableType(t *testing.T) {
11161123
assert.Empty(t, mock.CreatedWorkspaces, "CreateWorkspace must not be called when no workspace group is available")
11171124
}
11181125

1126+
func TestCreateInstancesWithTypeSetsCloudCredIDFromCatalog(t *testing.T) {
1127+
mock := NewMockGPUCreateStore()
1128+
ctx := &createContext{
1129+
t: terminal.New(),
1130+
store: mock,
1131+
opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4"},
1132+
org: mock.Org,
1133+
user: mock.User,
1134+
piped: true,
1135+
allInstanceTypes: &gpusearch.AllInstanceTypesResponse{
1136+
AllInstanceTypes: []gpusearch.InstanceType{
1137+
{
1138+
Type: "hyperstack_H100_sxm5x8",
1139+
CloudCredID: "cc-shadeform",
1140+
},
1141+
},
1142+
},
1143+
}
1144+
ctx.logf = func(_ string, _ ...interface{}) {}
1145+
1146+
result := ctx.createInstancesWithType(InstanceSpec{Type: "hyperstack_H100_sxm5x8"}, 0, 1)
1147+
1148+
assert.False(t, result.hadFailure)
1149+
require.Len(t, mock.CreatedOptions, 1)
1150+
assert.Equal(t, "cc-shadeform", mock.CreatedOptions[0].CloudCredID)
1151+
assert.Equal(t, "cc-shadeform", mock.CreatedOptions[0].WorkspaceGroupID)
1152+
}
1153+
11191154
func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) {
11201155
mock := NewMockGPUCreateStore()
11211156
ctx := &createContext{

‎pkg/cmd/gpusearch/gpusearch.go‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,14 @@ type WorkspaceGroup struct {
5555
PlatformType string `json:"platformType"`
5656
}
5757

58+
// CloudCred represents a cloud credential that can run an instance type.
59+
type CloudCred struct {
60+
ID string `json:"id"`
61+
Name string `json:"name"`
62+
PlatformType string `json:"platformType"`
63+
TenantType string `json:"tenantType"`
64+
}
65+
5866
// InstanceType represents an instance type from the API
5967
type InstanceType struct {
6068
Type string `json:"type"`
@@ -69,6 +77,8 @@ type InstanceType struct {
6977
SubLocation string `json:"sub_location"`
7078
AvailableLocations []string `json:"available_locations"`
7179
Provider string `json:"provider"`
80+
CloudCredID string `json:"cloud_cred_id"`
81+
CloudCreds []CloudCred `json:"cloud_creds"`
7282
WorkspaceGroups []WorkspaceGroup `json:"workspace_groups"`
7383
EstimatedDeployTime string `json:"estimated_deploy_time"`
7484
Stoppable bool `json:"stoppable"`
@@ -81,15 +91,21 @@ type InstanceTypesResponse struct {
8191
Items []InstanceType `json:"items"`
8292
}
8393

84-
// AllInstanceTypesResponse represents the authenticated API response with workspace groups
94+
// AllInstanceTypesResponse represents the authenticated API response with cloud credentials.
8595
type AllInstanceTypesResponse struct {
8696
AllInstanceTypes []InstanceType `json:"allInstanceTypes"`
8797
}
8898

89-
// GetWorkspaceGroupID returns the workspace group ID for an instance type, or empty string if not found
90-
func (r *AllInstanceTypesResponse) GetWorkspaceGroupID(instanceType string) string {
99+
// GetCloudCredID returns the cloud credential ID for an instance type, or empty string if not found.
100+
func (r *AllInstanceTypesResponse) GetCloudCredID(instanceType string) string {
91101
for _, it := range r.AllInstanceTypes {
92102
if it.Type == instanceType {
103+
if it.CloudCredID != "" {
104+
return it.CloudCredID
105+
}
106+
if len(it.CloudCreds) > 0 {
107+
return it.CloudCreds[0].ID
108+
}
93109
if len(it.WorkspaceGroups) > 0 {
94110
return it.WorkspaceGroups[0].ID
95111
}
@@ -98,6 +114,11 @@ func (r *AllInstanceTypesResponse) GetWorkspaceGroupID(instanceType string) stri
98114
return ""
99115
}
100116

117+
// GetWorkspaceGroupID is a compatibility alias while create still accepts workspaceGroupId.
118+
func (r *AllInstanceTypesResponse) GetWorkspaceGroupID(instanceType string) string {
119+
return r.GetCloudCredID(instanceType)
120+
}
121+
101122
// HasInstanceType reports whether the type exists in the API listing, independent of capacity.
102123
func (r *AllInstanceTypesResponse) HasInstanceType(instanceType string) bool {
103124
for _, it := range r.AllInstanceTypes {

‎pkg/store/instancetypes.go‎

Lines changed: 95 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,19 @@ package store
22

33
import (
44
"encoding/json"
5-
"fmt"
6-
"runtime"
75

6+
devplaneapiv1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
87
"github.com/brevdev/brev-cli/pkg/cmd/gpusearch"
98
"github.com/brevdev/brev-cli/pkg/config"
109
breverrors "github.com/brevdev/brev-cli/pkg/errors"
10+
"google.golang.org/protobuf/encoding/protojson"
1111
)
1212

1313
const (
14-
instanceTypesAPIPath = "v1/instance/types"
15-
// Authenticated API for instance types with workspace groups
16-
allInstanceTypesPathPattern = "api/instances/alltypesavailable/%s"
14+
instanceTypesAPIPath = "v1/instance/types"
15+
organizationAvailableTypesRPCPath = "devplaneapi.v1.InstanceService/ListOrganizationAvailableInstanceTypes"
16+
tenantTypeShared = "shared"
17+
tenantTypeIsolated = "isolated"
1718
)
1819

1920
// GetInstanceTypes fetches all available instance types from the public API
@@ -53,23 +54,102 @@ func fetchInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, erro
5354
return &result, nil
5455
}
5556

56-
// GetAllInstanceTypesWithWorkspaceGroups fetches instance types with workspace groups from the authenticated API
57-
func (s AuthHTTPStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) {
58-
path := fmt.Sprintf(allInstanceTypesPathPattern, orgID)
57+
// GetAllInstanceTypesWithCloudCreds fetches org-scoped instance types from dev-plane's public Connect API.
58+
func (s AuthHTTPStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) {
59+
publicClient := NewAuthHTTPClient(s.authHTTPClient.auth, config.NewConstants().GetBrevPublicAPIURL()).restyClient
5960

60-
var result gpusearch.AllInstanceTypesResponse
61-
res, err := s.authHTTPClient.restyClient.R().
61+
res, err := publicClient.R().
6262
SetHeader("Content-Type", "application/json").
63-
SetQueryParam("utm_source", "cli").
64-
SetQueryParam("os", runtime.GOOS).
65-
SetResult(&result).
66-
Get(path)
63+
SetBody(map[string]interface{}{
64+
"organizationId": orgID,
65+
"options": map[string]interface{}{
66+
"includeUnavailable": false,
67+
"includePreemptible": false,
68+
"includeCpu": true,
69+
"uniqueInstanceType": true,
70+
"skipAccessFilter": false,
71+
},
72+
}).
73+
Post(organizationAvailableTypesRPCPath)
6774
if err != nil {
6875
return nil, breverrors.WrapAndTrace(err)
6976
}
7077
if res.IsError() {
7178
return nil, NewHTTPResponseError(res)
7279
}
7380

74-
return &result, nil
81+
var protoResp devplaneapiv1.ListInstanceTypeResponse
82+
if err := protojson.Unmarshal(res.Body(), &protoResp); err != nil {
83+
return nil, breverrors.WrapAndTrace(err)
84+
}
85+
86+
return mapProtoInstanceTypesToGPUSearchResponse(protoResp.Items), nil
87+
}
88+
89+
// GetAllInstanceTypesWithWorkspaceGroups is a compatibility alias while create still accepts workspaceGroupId.
90+
func (s AuthHTTPStore) GetAllInstanceTypesWithWorkspaceGroups(orgID string) (*gpusearch.AllInstanceTypesResponse, error) {
91+
return s.GetAllInstanceTypesWithCloudCreds(orgID)
92+
}
93+
94+
func mapProtoInstanceTypesToGPUSearchResponse(instanceTypes []*devplaneapiv1.InstanceType) *gpusearch.AllInstanceTypesResponse {
95+
items := make([]gpusearch.InstanceType, 0, len(instanceTypes))
96+
for _, instanceType := range instanceTypes {
97+
if instanceType == nil {
98+
continue
99+
}
100+
item := gpusearch.InstanceType{
101+
Type: instanceType.GetType(),
102+
VCPU: int(instanceType.GetVcpu()),
103+
Location: instanceType.GetLocation(),
104+
SubLocation: instanceType.GetSubLocation(),
105+
AvailableLocations: instanceType.GetAvailableLocations(),
106+
Provider: instanceType.GetProvider(),
107+
CloudCredID: instanceType.GetCloudCredId(),
108+
EstimatedDeployTime: instanceType.GetEstimatedDeployTime(),
109+
Stoppable: instanceType.GetStoppable(),
110+
Rebootable: instanceType.GetRebootable(),
111+
CanModifyFirewallRules: instanceType.GetCanModifyFirewallRules(),
112+
}
113+
if basePrice := instanceType.GetBasePrice(); basePrice != nil {
114+
item.BasePrice = gpusearch.BasePrice{
115+
Currency: basePrice.GetCurrency(),
116+
Amount: basePrice.GetAmount(),
117+
}
118+
}
119+
for _, gpu := range instanceType.GetSupportedGpus() {
120+
item.SupportedGPUs = append(item.SupportedGPUs, gpusearch.GPU{
121+
Count: int(gpu.GetCount()),
122+
Name: gpu.GetName(),
123+
Manufacturer: gpu.GetManufacturer(),
124+
Memory: gpu.GetMemory(),
125+
})
126+
}
127+
if cloudCred := instanceType.GetCloudCred(); cloudCred != nil {
128+
mappedCloudCred := gpusearch.CloudCred{
129+
ID: cloudCred.GetCloudCredId(),
130+
Name: cloudCred.GetName(),
131+
PlatformType: cloudCred.GetProviderId(),
132+
TenantType: mapTenantType(cloudCred.GetTenantType()),
133+
}
134+
item.CloudCreds = []gpusearch.CloudCred{mappedCloudCred}
135+
item.WorkspaceGroups = []gpusearch.WorkspaceGroup{{
136+
ID: mappedCloudCred.ID,
137+
Name: mappedCloudCred.Name,
138+
PlatformType: mappedCloudCred.PlatformType,
139+
}}
140+
}
141+
items = append(items, item)
142+
}
143+
return &gpusearch.AllInstanceTypesResponse{AllInstanceTypes: items}
144+
}
145+
146+
func mapTenantType(tenantType devplaneapiv1.TenantType) string {
147+
switch tenantType {
148+
case devplaneapiv1.TenantType_TENANT_TYPE_ISOLATED:
149+
return tenantTypeIsolated
150+
case devplaneapiv1.TenantType_TENANT_TYPE_SHARED:
151+
return tenantTypeShared
152+
default:
153+
return ""
154+
}
75155
}

0 commit comments

Comments
 (0)