Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
128 changes: 128 additions & 0 deletions builder/deploy/deployer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1250,3 +1250,131 @@ func TestDeployer_Wakeup(t *testing.T) {
time.Sleep(100 * time.Millisecond)
})
}

// TestDeployer_serverlessDeploy_PD verifies that PD config is correctly
// propagated from DeployRequest to the database Deploy record in
// serverlessDeploy. Before the fix, field shadowing caused dr.PD to always
// be nil even when DeployExtend.PD was set.
func TestDeployer_serverlessDeploy_PD(t *testing.T) {
t.Run("deploy model with PD", func(t *testing.T) {
var oldDeploy database.Deploy
oldDeploy.ID = 1

pdConfig := &types.PDConfig{
Enabled: true,
PrefillReplicas: 2,
DecodeReplicas: 2,
Prefill: &types.PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
Decode: &types.PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
}

dr := types.DeployRequest{
RepoID: 1,
Type: types.InferenceType,
UserUUID: "1",
SKU: "1",
DeployExtend: types.DeployExtend{
PD: pdConfig,
},
}

newDeploy := oldDeploy
newDeploy.UserUUID = dr.UserUUID
newDeploy.SKU = dr.SKU
newDeploy.PD = dr.PD

mockTaskStore := mockdb.NewMockDeployTaskStore(t)
mockTaskStore.EXPECT().GetServerlessDeployByRepID(mock.Anything, dr.RepoID).Return(&oldDeploy, nil)
mockTaskStore.EXPECT().UpdateDeploy(mock.Anything, &newDeploy).Return(nil)

d := &deployer{
deployTaskStore: mockTaskStore,
}
dbdeploy, err := d.serverlessDeploy(context.TODO(), dr)
require.Nil(t, err)
require.NotNil(t, dbdeploy.PD)
require.True(t, dbdeploy.PD.Enabled)
require.Equal(t, 2, dbdeploy.PD.PrefillReplicas)
require.Equal(t, 2, dbdeploy.PD.DecodeReplicas)
require.Same(t, pdConfig, dbdeploy.PD)
})

t.Run("deploy model without PD (nil)", func(t *testing.T) {
var oldDeploy database.Deploy
oldDeploy.ID = 1

dr := types.DeployRequest{
RepoID: 1,
Type: types.InferenceType,
UserUUID: "1",
SKU: "1",
}

newDeploy := oldDeploy
newDeploy.UserUUID = dr.UserUUID
newDeploy.SKU = dr.SKU
newDeploy.PD = nil

mockTaskStore := mockdb.NewMockDeployTaskStore(t)
mockTaskStore.EXPECT().GetServerlessDeployByRepID(mock.Anything, dr.RepoID).Return(&oldDeploy, nil)
mockTaskStore.EXPECT().UpdateDeploy(mock.Anything, &newDeploy).Return(nil)

d := &deployer{
deployTaskStore: mockTaskStore,
}
dbdeploy, err := d.serverlessDeploy(context.TODO(), dr)
require.Nil(t, err)
require.Nil(t, dbdeploy.PD)
})
}

// TestDeployer_dedicatedDeploy_PD verifies that PD config is correctly
// propagated from DeployRequest to the database Deploy record in
// dedicatedDeploy.
func TestDeployer_dedicatedDeploy_PD(t *testing.T) {
pdConfig := &types.PDConfig{
Enabled: true,
PrefillReplicas: 1,
DecodeReplicas: 1,
Prefill: &types.PDRoleRuntimeConfig{
TP: 2, EP: 2, DP: 1, TotalGPUs: 2,
},
Decode: &types.PDRoleRuntimeConfig{
TP: 2, EP: 2, DP: 1, TotalGPUs: 2,
},
}

dr := types.DeployRequest{
Path: "namespace/name",
Type: types.InferenceType,
DeployExtend: types.DeployExtend{
PD: pdConfig,
},
}

var capturedDeploy *database.Deploy
mockTaskStore := mockdb.NewMockDeployTaskStore(t)
mockTaskStore.EXPECT().CreateDeploy(mock.Anything, mock.MatchedBy(func(d *database.Deploy) bool {
capturedDeploy = d
return true
})).Return(nil)

node, _ := snowflake.NewNode(1)
d := &deployer{
snowflakeNode: node,
deployTaskStore: mockTaskStore,
}

_, err := d.dedicatedDeploy(context.TODO(), dr)
require.Nil(t, err)
require.NotNil(t, capturedDeploy)
require.NotNil(t, capturedDeploy.PD)
require.True(t, capturedDeploy.PD.Enabled)
require.Equal(t, 1, capturedDeploy.PD.PrefillReplicas)
require.Equal(t, 1, capturedDeploy.PD.DecodeReplicas)
require.Same(t, pdConfig, capturedDeploy.PD)
}
73 changes: 73 additions & 0 deletions common/types/deploy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -241,3 +241,76 @@ func TestSVCRequest_PDThroughDeployExtend(t *testing.T) {
require.Equal(t, 2, req.PD.PrefillReplicas)
require.Equal(t, 2, req.PD.DecodeReplicas)
}

// TestDeployRequest_PDNoShadow verifies that DeployRequest.PD (from embedded
// DeployExtend) is not shadowed by an explicit PD field. Before the fix,
// DeployRequest had both an explicit PD field and an embedded DeployExtend.PD,
// causing field shadowing: setting DeployExtend.PD was invisible when reading
// DeployRequest.PD. This test ensures they are the same field.
func TestDeployRequest_PDNoShadow(t *testing.T) {
pdConfig := &PDConfig{
Enabled: true,
PrefillReplicas: 2,
DecodeReplicas: 3,
Prefill: &PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
Decode: &PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
}

dr := DeployRequest{}
// Set PD via DeployExtend (the embedded field)
dr.DeployExtend.PD = pdConfig

// Reading dr.PD should return the same pointer (no shadowing)
require.NotNil(t, dr.PD)
require.Same(t, pdConfig, dr.PD)
require.True(t, dr.PD.Enabled)
require.Equal(t, 2, dr.PD.PrefillReplicas)
require.Equal(t, 3, dr.PD.DecodeReplicas)

// Setting dr.PD should also set DeployExtend.PD (same field)
dr.PD.Enabled = false
require.False(t, dr.DeployExtend.PD.Enabled)
}

// TestDeployRequest_PDJSONBinding verifies that JSON unmarshaling correctly
// populates DeployRequest.PD through the embedded DeployExtend field.
func TestDeployRequest_PDJSONBinding(t *testing.T) {
jsonStr := `{
"deploy_name": "test-deploy",
"pd": {
"enabled": true,
"prefill_replicas": 2,
"decode_replicas": 2,
"prefill": {"tp": 2, "ep": 1, "dp": 1, "total_gpus": 2},
"decode": {"tp": 2, "ep": 1, "dp": 1, "total_gpus": 2}
}
}`

var dr DeployRequest
err := json.Unmarshal([]byte(jsonStr), &dr)
require.NoError(t, err)
require.Equal(t, "test-deploy", dr.DeployName)

// PD should be populated via embedded DeployExtend
require.NotNil(t, dr.PD)
require.True(t, dr.PD.Enabled)
require.Equal(t, 2, dr.PD.PrefillReplicas)
require.Equal(t, 2, dr.PD.DecodeReplicas)
require.NotNil(t, dr.PD.Prefill)
require.Equal(t, 2, dr.PD.Prefill.TP)
require.NotNil(t, dr.PD.Decode)
require.Equal(t, 2, dr.PD.Decode.TP)

// Verify JSON marshaling round-trips correctly
out, err := json.Marshal(dr)
require.NoError(t, err)
var dr2 DeployRequest
err = json.Unmarshal(out, &dr2)
require.NoError(t, err)
require.NotNil(t, dr2.PD)
require.True(t, dr2.PD.Enabled)
}
4 changes: 4 additions & 0 deletions common/types/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,10 @@ type InstanceRunReq struct {
Revision string `json:"revision"`
OrderDetailID int64 `json:"order_detail_id"`
EngineArgs string `json:"engine_args"`
// EnablePD enables PD (Prefill-Decode) disaggregation inference architecture.
// When true, the system checks the model metadata for PD recommendation,
// validates hardware resources, and splits resources between prefill and decode.
EnablePD bool `json:"enable_pd"`
// OwnerNamespace is optional. If set, the finetune is created under this namespace (user or org); path {namespace} remains the model's owner.
OwnerNamespace string `json:"owner_namespace,omitempty"`
}
Expand Down
17 changes: 9 additions & 8 deletions common/types/repo.go
Original file line number Diff line number Diff line change
Expand Up @@ -168,14 +168,15 @@ type Repository struct {
}

type Metadata struct {
ModelParams float32 `json:"model_params"`
TensorType string `json:"tensor_type"`
Architecture string `json:"architecture"`
MiniGPUMemoryGB float32 `json:"mini_gpu_memory_gb"`
MiniGPUFinetuneGB float32 `json:"mini_gpu_finetune_gb"`
ModelType string `json:"model_type"`
ClassName string `json:"class_name"`
Quantizations []Quantization `json:"quantizations,omitempty"`
ModelParams float32 `json:"model_params"`
TensorType string `json:"tensor_type"`
Architecture string `json:"architecture"`
MiniGPUMemoryGB float32 `json:"mini_gpu_memory_gb"`
MiniGPUFinetuneGB float32 `json:"mini_gpu_finetune_gb"`
ModelType string `json:"model_type"`
ClassName string `json:"class_name"`
Quantizations []Quantization `json:"quantizations,omitempty"`
PDRecommendation *PDRecommendation `json:"pd_recommendation,omitempty"`
}

type RepoPageOpts struct {
Expand Down
1 change: 1 addition & 0 deletions component/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,7 @@ func (c *modelComponentImpl) Show(ctx context.Context, namespace, name, currentU
ClassName: model.Repository.Metadata.ClassName,
Quantizations: model.Repository.Metadata.Quantizations,
MiniGPUFinetuneGB: model.Repository.Metadata.MiniGPUFinetuneGB,
PDRecommendation: model.Repository.Metadata.PDRecommendation,
},

MultiSource: types.MultiSource{
Expand Down
1 change: 1 addition & 0 deletions component/repo_deploy.go
Original file line number Diff line number Diff line change
Expand Up @@ -575,6 +575,7 @@ func (c *repoComponentImpl) DeployDetail(ctx context.Context, detailReq types.De
UserUUID: deploy.UserUUID,
OwnerNamespace: deploy.OwnerNamespace,
}
resDeploy.PD = deploy.PD

return &resDeploy, nil
}
Expand Down
69 changes: 69 additions & 0 deletions component/repo_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1430,6 +1430,75 @@ func TestRepoComponent_DeployDetail(t *testing.T) {

}

func TestRepoComponent_DeployDetailWithPD(t *testing.T) {
ctx := context.TODO()
repo := initializeTestRepoComponent(ctx, t)
mockUserRepoAdminPermission(ctx, repo.mocks.stores, "user")

pdConfig := &types.PDConfig{
Enabled: true,
PrefillReplicas: 2,
DecodeReplicas: 2,
Prefill: &types.PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
Decode: &types.PDRoleRuntimeConfig{
TP: 2, EP: 1, DP: 1, TotalGPUs: 2,
},
}

repo.mocks.stores.ClusterInfoMock().EXPECT().ByClusterID(ctx, "cluster").Return(database.ClusterInfo{
Zone: "z",
}, nil)
repo.mocks.stores.DeployTaskMock().EXPECT().GetDeployByID(ctx, int64(1)).Return(&database.Deploy{
RepoID: 1,
UserUUID: "uuid",
OrderDetailID: 11,
ClusterID: "cluster",
SvcName: "svc",
Status: deployStatus.Running,
PD: pdConfig,
}, nil)

repo.mocks.deployer.EXPECT().GetReplica(ctx, types.DeployRequest{
Namespace: "ns",
Name: "n",
ClusterID: "cluster",
SvcName: "svc",
}).Return(1, 2, []types.Instance{{Name: "i1"}}, nil)

repo.mocks.deployer.EXPECT().Status(ctx, types.DeployRequest{
DeployID: 0,
SpaceID: 0,
ModelID: 0,
Namespace: "ns",
Name: "n",
SvcName: "svc",
ClusterID: "cluster",
}, false).Return("svc", 23, nil, nil)

dp, err := repo.DeployDetail(ctx, types.DeployActReq{
RepoType: types.ModelRepo,
Namespace: "ns",
Name: "n",
CurrentUser: "user",
DeployID: 1,
DeployType: 2,
InstanceName: "i1",
})
require.Nil(t, err)

// PD config should be populated in the response
require.NotNil(t, dp.PD)
require.True(t, dp.PD.Enabled)
require.Equal(t, 2, dp.PD.PrefillReplicas)
require.Equal(t, 2, dp.PD.DecodeReplicas)
require.NotNil(t, dp.PD.Prefill)
require.Equal(t, 2, dp.PD.Prefill.TP)
require.NotNil(t, dp.PD.Decode)
require.Equal(t, 2, dp.PD.Decode.TP)
}

func TestRepoComponent_DeployInstanceLogs(t *testing.T) {
ctx := context.TODO()
repo := initializeTestRepoComponent(ctx, t)
Expand Down
Loading