Skip to content

Commit 497d8b7

Browse files
DwyaneShiDwyaneShi
authored
fix(rm): fix provision store apis (vllm-project#2195)
Signed-off-by: Haiyang Shi <haiyang.shi@bytedance.com>
1 parent 3c36553 commit 497d8b7

4 files changed

Lines changed: 103 additions & 43 deletions

File tree

‎apps/console/api/resource_manager/provisioner/kubernetes.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,8 @@ func (p *K8sProvisioner) Provision(ctx context.Context, req *types.ResourceProvi
104104
Kubernetes: &types.KubernetesProvisionDetail{},
105105
}
106106

107-
if err := p.store.InsertProvision(ctx, req.IdempotencyKey, result); err != nil {
108-
return nil, fmt.Errorf("insert provision: %w", err)
107+
if err := p.store.UpsertProvision(ctx, result); err != nil {
108+
return nil, fmt.Errorf("upsert provision: %w", err)
109109
}
110110

111111
return result, nil

‎apps/console/api/store/gorm.go‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -678,13 +678,19 @@ func (s *GORMStore) GetProvisionByIdempotencyKey(ctx context.Context, idempotenc
678678
return res, nil
679679
}
680680

681-
func (s *GORMStore) InsertProvision(ctx context.Context, idempotencyKey string, result *types.ProvisionResult) error {
681+
func (s *GORMStore) UpsertProvision(ctx context.Context, result *types.ProvisionResult) error {
682+
if result == nil {
683+
return fmt.Errorf("provision result is required")
684+
}
685+
if result.IdempotencyKey == "" {
686+
return fmt.Errorf("idempotency key is required")
687+
}
682688
record, err := result.ToProvisionRecord()
683689
if err != nil {
684690
return fmt.Errorf("failed to convert provision result to record: %w", err)
685691
}
686692
rec := models.ProvisionResult{
687-
IdempotencyKey: idempotencyKey,
693+
IdempotencyKey: result.IdempotencyKey,
688694
ProvisionID: record.ProvisionID,
689695
Region: record.Region,
690696
Status: record.Status,
@@ -693,8 +699,11 @@ func (s *GORMStore) InsertProvision(ctx context.Context, idempotencyKey string,
693699
UpdatedAt: record.UpdatedAt,
694700
Deleted: false,
695701
}
696-
if err := s.db.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "idempotency_key"}}, DoUpdates: clause.AssignmentColumns([]string{"provision_id", "region", "status", "payload", "created_at", "updated_at", "deleted"})}).Create(&rec).Error; err != nil {
697-
return fmt.Errorf("failed to set provision result: %w", err)
702+
if err := s.db.WithContext(ctx).Clauses(clause.OnConflict{
703+
Columns: []clause.Column{{Name: "idempotency_key"}},
704+
DoUpdates: clause.AssignmentColumns([]string{"provision_id", "region", "status", "payload", "updated_at", "deleted"}),
705+
}).Create(&rec).Error; err != nil {
706+
return fmt.Errorf("failed to upsert provision result: %w", err)
698707
}
699708
return nil
700709
}

‎apps/console/api/store/store.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,8 @@ type Store interface {
8383
// Returns nil and a NotFound error if the key doesn't exist.
8484
GetProvisionByIdempotencyKey(ctx context.Context, idempotencyKey string) (*types.ProvisionResult, error)
8585

86-
// InsertProvision stores a provision result with the given idempotency key.
87-
InsertProvision(ctx context.Context, idempotencyKey string, result *types.ProvisionResult) error
86+
// UpsertProvision inserts/updates a provision result with the given idempotency key.
87+
UpsertProvision(ctx context.Context, result *types.ProvisionResult) error
8888

8989
// UpdateProvisionStatus updates the status of a provision result.
9090
UpdateProvisionStatus(ctx context.Context, provisionId string, status types.ProvisionStatus) error

‎apps/console/api/store/store_test.go‎

Lines changed: 86 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -863,25 +863,70 @@ func TestMemoryStore(t *testing.T) {
863863
})
864864

865865
t.Run("Provision", func(t *testing.T) {
866-
t.Run("InsertProvision", func(t *testing.T) {
866+
t.Run("UpsertProvision_Insert", func(t *testing.T) {
867867
result := &types.ProvisionResult{
868-
ProvisionID: "prov-123",
869-
Status: types.ProvisionStatusPending,
868+
IdempotencyKey: "idem-key-1",
869+
ProvisionID: "prov-123",
870+
Status: types.ProvisionStatusPending,
870871
}
871-
err := s.InsertProvision(ctx, "idem-key-1", result)
872+
err := s.UpsertProvision(ctx, result)
872873
if err != nil {
873-
t.Fatalf("InsertProvision failed: %v", err)
874+
t.Fatalf("UpsertProvision failed: %v", err)
875+
}
876+
})
877+
878+
t.Run("UpsertProvision_Update", func(t *testing.T) {
879+
// First insert
880+
result := &types.ProvisionResult{
881+
IdempotencyKey: "idem-key-upsert",
882+
ProvisionID: "prov-upsert-1",
883+
Status: types.ProvisionStatusPending,
884+
}
885+
err := s.UpsertProvision(ctx, result)
886+
if err != nil {
887+
t.Fatalf("UpsertProvision (insert) failed: %v", err)
888+
}
889+
890+
// Retrieve to get the original CreatedAt
891+
original, err := s.GetProvision(ctx, "prov-upsert-1")
892+
if err != nil {
893+
t.Fatalf("GetProvision failed: %v", err)
894+
}
895+
896+
// Update with same idempotency key but different provision_id and status
897+
updated := &types.ProvisionResult{
898+
IdempotencyKey: "idem-key-upsert",
899+
ProvisionID: "prov-upsert-2",
900+
Status: types.ProvisionStatusRunning,
901+
}
902+
err = s.UpsertProvision(ctx, updated)
903+
if err != nil {
904+
t.Fatalf("UpsertProvision (update) failed: %v", err)
905+
}
906+
907+
// Retrieve the updated record
908+
retrieved, err := s.GetProvision(ctx, "prov-upsert-2")
909+
if err != nil {
910+
t.Fatalf("GetProvision after update failed: %v", err)
911+
}
912+
if retrieved.Status != types.ProvisionStatusRunning {
913+
t.Errorf("expected status 'running', got %q", retrieved.Status)
914+
}
915+
// Verify CreatedAt is preserved
916+
if !retrieved.CreatedAt.Equal(original.CreatedAt) {
917+
t.Errorf("CreatedAt should be preserved: original=%v, retrieved=%v", original.CreatedAt, retrieved.CreatedAt)
874918
}
875919
})
876920

877921
t.Run("GetProvision", func(t *testing.T) {
878922
result := &types.ProvisionResult{
879-
ProvisionID: "prov-456",
880-
Status: types.ProvisionStatusPending,
923+
IdempotencyKey: "idem-key-2",
924+
ProvisionID: "prov-456",
925+
Status: types.ProvisionStatusPending,
881926
}
882-
err := s.InsertProvision(ctx, "idem-key-2", result)
927+
err := s.UpsertProvision(ctx, result)
883928
if err != nil {
884-
t.Fatalf("InsertProvision for get test failed: %v", err)
929+
t.Fatalf("UpsertProvision for get test failed: %v", err)
885930
}
886931

887932
retrieved, err := s.GetProvision(ctx, "prov-456")
@@ -909,11 +954,12 @@ func TestMemoryStore(t *testing.T) {
909954

910955
t.Run("ExistsProvision", func(t *testing.T) {
911956
result := &types.ProvisionResult{
912-
ProvisionID: "prov-789",
913-
Status: types.ProvisionStatusPending,
957+
IdempotencyKey: "idem-key-3",
958+
ProvisionID: "prov-789",
959+
Status: types.ProvisionStatusPending,
914960
}
915-
if err := s.InsertProvision(ctx, "idem-key-3", result); err != nil {
916-
t.Fatalf("InsertProvision failed: %v", err)
961+
if err := s.UpsertProvision(ctx, result); err != nil {
962+
t.Fatalf("UpsertProvision failed: %v", err)
917963
}
918964

919965
exists, err := s.ExistsProvision(ctx, "prov-789")
@@ -935,11 +981,12 @@ func TestMemoryStore(t *testing.T) {
935981

936982
t.Run("UpdateProvisionStatus", func(t *testing.T) {
937983
result := &types.ProvisionResult{
938-
ProvisionID: "prov-update",
939-
Status: types.ProvisionStatusPending,
984+
IdempotencyKey: "idem-key-update",
985+
ProvisionID: "prov-update",
986+
Status: types.ProvisionStatusPending,
940987
}
941-
if err := s.InsertProvision(ctx, "idem-key-update", result); err != nil {
942-
t.Fatalf("InsertProvision failed: %v", err)
988+
if err := s.UpsertProvision(ctx, result); err != nil {
989+
t.Fatalf("UpsertProvision failed: %v", err)
943990
}
944991

945992
err := s.UpdateProvisionStatus(ctx, "prov-update", types.ProvisionStatusRunning)
@@ -958,11 +1005,12 @@ func TestMemoryStore(t *testing.T) {
9581005

9591006
t.Run("DeleteProvision", func(t *testing.T) {
9601007
result := &types.ProvisionResult{
961-
ProvisionID: "prov-delete",
962-
Status: types.ProvisionStatusPending,
1008+
IdempotencyKey: "idem-key-delete",
1009+
ProvisionID: "prov-delete",
1010+
Status: types.ProvisionStatusPending,
9631011
}
964-
if err := s.InsertProvision(ctx, "idem-key-delete", result); err != nil {
965-
t.Fatalf("InsertProvision failed: %v", err)
1012+
if err := s.UpsertProvision(ctx, result); err != nil {
1013+
t.Fatalf("UpsertProvision failed: %v", err)
9661014
}
9671015

9681016
err := s.DeleteProvision(ctx, "prov-delete")
@@ -981,11 +1029,12 @@ func TestMemoryStore(t *testing.T) {
9811029
// Insert multiple provisions
9821030
for i := 1; i <= 3; i++ {
9831031
result := &types.ProvisionResult{
984-
ProvisionID: fmt.Sprintf("prov-list-%d", i),
985-
Status: types.ProvisionStatusPending,
1032+
IdempotencyKey: fmt.Sprintf("idem-key-list-%d", i),
1033+
ProvisionID: fmt.Sprintf("prov-list-%d", i),
1034+
Status: types.ProvisionStatusPending,
9861035
}
987-
if err := s.InsertProvision(ctx, fmt.Sprintf("idem-key-list-%d", i), result); err != nil {
988-
t.Fatalf("InsertProvision failed: %v", err)
1036+
if err := s.UpsertProvision(ctx, result); err != nil {
1037+
t.Fatalf("UpsertProvision failed: %v", err)
9891038
}
9901039
}
9911040

@@ -1024,21 +1073,23 @@ func TestMemoryStore(t *testing.T) {
10241073
}
10251074

10261075
resultA := &types.ProvisionResult{
1027-
ProvisionID: "prov-region-a",
1028-
Status: types.ProvisionStatusPending,
1029-
Region: string(regionABytes),
1076+
IdempotencyKey: "idem-key-region-a",
1077+
ProvisionID: "prov-region-a",
1078+
Status: types.ProvisionStatusPending,
1079+
Region: string(regionABytes),
10301080
}
1031-
if err := s.InsertProvision(ctx, "idem-key-region-a", resultA); err != nil {
1032-
t.Fatalf("InsertProvision regionA failed: %v", err)
1081+
if err := s.UpsertProvision(ctx, resultA); err != nil {
1082+
t.Fatalf("UpsertProvision regionA failed: %v", err)
10331083
}
10341084

10351085
resultB := &types.ProvisionResult{
1036-
ProvisionID: "prov-region-b",
1037-
Status: types.ProvisionStatusPending,
1038-
Region: string(regionBBytes),
1086+
IdempotencyKey: "idem-key-region-b",
1087+
ProvisionID: "prov-region-b",
1088+
Status: types.ProvisionStatusPending,
1089+
Region: string(regionBBytes),
10391090
}
1040-
if err := s.InsertProvision(ctx, "idem-key-region-b", resultB); err != nil {
1041-
t.Fatalf("InsertProvision regionB failed: %v", err)
1091+
if err := s.UpsertProvision(ctx, resultB); err != nil {
1092+
t.Fatalf("UpsertProvision regionB failed: %v", err)
10421093
}
10431094

10441095
regions := []types.RegionSpec{regionA}

0 commit comments

Comments
 (0)