feat: 适配前端PD 分离部署逻辑

CementZhang · Rader · commit a0ccabbfead3 · 2026-07-01T16:12:28.000+08:00
diff --git a/builder/deploy/deployer_test.go b/builder/deploy/deployer_test.go
@@ -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)
+}
diff --git a/common/types/deploy_test.go b/common/types/deploy_test.go
@@ -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)
+}
diff --git a/common/types/model.go b/common/types/model.go
@@ -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"`
 }
diff --git a/common/types/repo.go b/common/types/repo.go
@@ -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 {
diff --git a/component/model.go b/component/model.go
@@ -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{
diff --git a/component/repo_deploy.go b/component/repo_deploy.go
@@ -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
 }
diff --git a/component/repo_test.go b/component/repo_test.go
@@ -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)

Original file line number	Diff line number	Diff line change
`@@ -575,6 +575,7 @@ func (c *repoComponentImpl) DeployDetail(ctx context.Context, detailReq types.De`
`575`	`575`	`UserUUID: deploy.UserUUID,`
`576`	`576`	`OwnerNamespace: deploy.OwnerNamespace,`
`577`	`577`	`}`
	`578`	`+ resDeploy.PD = deploy.PD`
`578`	`579`
`579`	`580`	`return &resDeploy, nil`
`580`	`581`	`}`