Skip to content

Commit adaa842

Browse files
davidmontoyagodavidmontoyago
committed
ensure endpoint lookup with full name
1 parent d8df543 commit adaa842

2 files changed

Lines changed: 213 additions & 3 deletions

File tree

‎pkg/resources/provider_test.go‎

Lines changed: 206 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1110,7 +1110,9 @@ func TestVertexModelDeploymentUpdate_ModelOnly(t *testing.T) {
11101110
func TestVertexModelDeploymentRead_ModelWithEndpoint(t *testing.T) {
11111111
ctx := context.Background()
11121112

1113-
// Test state for model with endpoint deployment read
1113+
// Test state for model with endpoint deployment read using SHORT endpoint name
1114+
// This test verifies that when the state contains a short endpoint name,
1115+
// the VertexEndpointModelGetter properly converts it to a fully qualified name for the API call
11141116
projectID := testProjectID
11151117
region := testRegion
11161118
endpointID := testEndpointID
@@ -1145,7 +1147,7 @@ func TestVertexModelDeploymentRead_ModelWithEndpoint(t *testing.T) {
11451147
},
11461148
ModelName: modelName,
11471149
DeployedModelID: deployedModelID,
1148-
EndpointName: testEndpointPath, // Use full path format
1150+
EndpointName: endpointID, // Use SHORT endpoint name (just the ID) to test name conversion
11491151
CreateTime: createTime,
11501152
}
11511153

@@ -1305,3 +1307,205 @@ func TestVertexModelDeploymentRead_ModelWithEndpoint(t *testing.T) {
13051307
t.Errorf("Expected Inputs.Region %s, got %s", region, result.Inputs.Region)
13061308
}
13071309
}
1310+
1311+
//nolint:paralleltest,tparallel // Cannot run in parallel due to shared testFactoryRegistry
1312+
func TestVertexModelDeploymentRead_ModelWithEndpointNameFullyQualified(t *testing.T) {
1313+
ctx := context.Background()
1314+
1315+
// Test state for model with endpoint deployment read using FULLY QUALIFIED endpoint name
1316+
// This test verifies that when the state already contains a fully qualified endpoint name,
1317+
// it is passed through correctly to the endpoint client
1318+
projectID := testProjectID
1319+
region := testRegion
1320+
endpointID := testEndpointID
1321+
fullEndpointName := testEndpointPath
1322+
modelName := testModelName
1323+
deployedModelID := "deployed-model-id-456"
1324+
modelImageURL := testModelImageURL
1325+
modelArtifactsBucketURI := testModelArtifactsBucketURI
1326+
modelPredictionInputSchemaURI := testModelPredictionInputSchemaURI
1327+
modelPredictionOutputSchemaURI := testModelPredictionOutputSchemaURI
1328+
machineType := "n1-standard-8"
1329+
minReplicas := int32(3)
1330+
maxReplicas := int32(10)
1331+
trafficPercent := int32(50)
1332+
createTime := testCreateTime
1333+
1334+
// Create initial state with endpoint deployment using FULLY QUALIFIED endpoint name
1335+
state := VertexModelDeploymentState{
1336+
VertexModelDeploymentArgs: VertexModelDeploymentArgs{
1337+
ProjectID: projectID,
1338+
Region: region,
1339+
ModelImageURL: "old-image-url", // Different from mock to test update
1340+
ModelArtifactsBucketURI: "old-bucket-uri",
1341+
ModelPredictionInputSchemaURI: "old-input-schema",
1342+
EndpointModelDeployment: &EndpointModelDeploymentArgs{
1343+
EndpointID: endpointID,
1344+
MachineType: "old-machine-type", // Different from mock to test update
1345+
MinReplicas: 1, // Different from mock
1346+
MaxReplicas: 3, // Different from mock
1347+
TrafficPercent: 100, // Different from mock
1348+
},
1349+
},
1350+
ModelName: modelName,
1351+
DeployedModelID: deployedModelID,
1352+
EndpointName: testEndpointPath, // Use FULLY QUALIFIED endpoint name - this is the key difference
1353+
CreateTime: createTime,
1354+
}
1355+
1356+
// Mock model response
1357+
mockModel := &aiplatformpb.Model{
1358+
Name: modelName,
1359+
ArtifactUri: modelArtifactsBucketURI,
1360+
ContainerSpec: &aiplatformpb.ModelContainerSpec{
1361+
ImageUri: modelImageURL,
1362+
},
1363+
PredictSchemata: &aiplatformpb.PredictSchemata{
1364+
InstanceSchemaUri: modelPredictionInputSchemaURI,
1365+
PredictionSchemaUri: modelPredictionOutputSchemaURI,
1366+
},
1367+
Labels: map[string]string{"env": "staging", "team": "ai"},
1368+
}
1369+
1370+
// Mock endpoint response with deployed model
1371+
mockEndpoint := &aiplatformpb.Endpoint{
1372+
Name: fullEndpointName,
1373+
DeployedModels: []*aiplatformpb.DeployedModel{
1374+
{
1375+
Id: deployedModelID,
1376+
PredictionResources: &aiplatformpb.DeployedModel_DedicatedResources{
1377+
DedicatedResources: &aiplatformpb.DedicatedResources{
1378+
MachineSpec: &aiplatformpb.MachineSpec{
1379+
MachineType: machineType,
1380+
},
1381+
MinReplicaCount: minReplicas,
1382+
MaxReplicaCount: maxReplicas,
1383+
},
1384+
},
1385+
},
1386+
},
1387+
TrafficSplit: map[string]int32{
1388+
deployedModelID: trafficPercent,
1389+
},
1390+
}
1391+
1392+
req := infer.ReadRequest[VertexModelDeploymentArgs, VertexModelDeploymentState]{
1393+
ID: "test-model-with-endpoint",
1394+
Inputs: state.VertexModelDeploymentArgs,
1395+
State: state,
1396+
}
1397+
1398+
// Variables to capture request parameters
1399+
var capturedGetModelRequest *aiplatformpb.GetModelRequest
1400+
var capturedGetEndpointRequest *aiplatformpb.GetEndpointRequest
1401+
1402+
modelClientFactory := MockModelClientFactory(&MockModelClient{
1403+
GetModelFunc: func(_ context.Context, req *aiplatformpb.GetModelRequest, _ ...gax.CallOption) (*aiplatformpb.Model, error) {
1404+
capturedGetModelRequest = req
1405+
1406+
return mockModel, nil
1407+
},
1408+
})
1409+
1410+
endpointClientFactory := MockEndpointClientFactory(&MockEndpointClient{
1411+
GetEndpointFunc: func(_ context.Context, req *aiplatformpb.GetEndpointRequest, _ ...gax.CallOption) (*aiplatformpb.Endpoint, error) {
1412+
capturedGetEndpointRequest = req
1413+
1414+
return mockEndpoint, nil
1415+
},
1416+
})
1417+
1418+
provider := &VertexModelDeployment{}
1419+
testFactoryRegistry.modelClientFactory = modelClientFactory
1420+
testFactoryRegistry.endpointClientFactory = endpointClientFactory
1421+
1422+
// Execute Read
1423+
result, err := provider.Read(ctx, req)
1424+
if err != nil {
1425+
t.Fatalf("Expected nil error from Read, but got %v", err)
1426+
}
1427+
1428+
// Assert that the resource ID is preserved
1429+
if result.ID != req.ID {
1430+
t.Errorf("Expected resource ID %s, got %s", req.ID, result.ID)
1431+
}
1432+
1433+
// Validate GetModelRequest was captured and has correct parameters
1434+
if capturedGetModelRequest == nil {
1435+
t.Fatal("GetModelRequest was not captured")
1436+
}
1437+
if capturedGetModelRequest.Name != modelName {
1438+
t.Errorf("Expected model name %s, got %s", modelName, capturedGetModelRequest.Name)
1439+
}
1440+
1441+
// Validate GetEndpointRequest was captured and has correct parameters
1442+
if capturedGetEndpointRequest == nil {
1443+
t.Fatal("GetEndpointRequest was not captured")
1444+
}
1445+
expectedEndpointPath := testEndpointPath
1446+
if capturedGetEndpointRequest.Name != expectedEndpointPath {
1447+
t.Errorf("Expected endpoint path %s, got %s", expectedEndpointPath, capturedGetEndpointRequest.Name)
1448+
}
1449+
1450+
// Validate the returned state reflects the mock model data
1451+
resultState := result.State
1452+
if resultState.ModelName != modelName {
1453+
t.Errorf("Expected ModelName %s, got %s", modelName, resultState.ModelName)
1454+
}
1455+
if resultState.ModelImageURL != modelImageURL {
1456+
t.Errorf("Expected ModelImageURL %s, got %s", modelImageURL, resultState.ModelImageURL)
1457+
}
1458+
if resultState.ModelArtifactsBucketURI != modelArtifactsBucketURI {
1459+
t.Errorf("Expected ModelArtifactsBucketURI %s, got %s", modelArtifactsBucketURI, resultState.ModelArtifactsBucketURI)
1460+
}
1461+
if resultState.ModelPredictionInputSchemaURI != modelPredictionInputSchemaURI {
1462+
t.Errorf("Expected ModelPredictionInputSchemaURI %s, got %s", modelPredictionInputSchemaURI, resultState.ModelPredictionInputSchemaURI)
1463+
}
1464+
if resultState.ModelPredictionOutputSchemaURI != modelPredictionOutputSchemaURI {
1465+
t.Errorf("Expected ModelPredictionOutputSchemaURI %s, got %s", modelPredictionOutputSchemaURI, resultState.ModelPredictionOutputSchemaURI)
1466+
}
1467+
1468+
// Validate labels are updated from model
1469+
if len(resultState.Labels) != 2 {
1470+
t.Errorf("Expected 2 labels, got %d", len(resultState.Labels))
1471+
}
1472+
if resultState.Labels["env"] != "staging" {
1473+
t.Errorf("Expected label env=staging, got %s", resultState.Labels["env"])
1474+
}
1475+
if resultState.Labels["team"] != "ai" {
1476+
t.Errorf("Expected label team=ai, got %s", resultState.Labels["team"])
1477+
}
1478+
1479+
// Validate endpoint-related fields are updated from endpoint response
1480+
if resultState.EndpointName != fullEndpointName {
1481+
t.Errorf("Expected EndpointName %s, got %s", fullEndpointName, resultState.EndpointName)
1482+
}
1483+
if resultState.DeployedModelID != deployedModelID {
1484+
t.Errorf("Expected DeployedModelID %s, got %s", deployedModelID, resultState.DeployedModelID)
1485+
}
1486+
1487+
// Validate endpoint deployment configuration is updated from live endpoint
1488+
if resultState.EndpointModelDeployment == nil {
1489+
t.Fatal("EndpointModelDeployment should not be nil")
1490+
}
1491+
if resultState.EndpointModelDeployment.MachineType != machineType {
1492+
t.Errorf("Expected MachineType %s, got %s", machineType, resultState.EndpointModelDeployment.MachineType)
1493+
}
1494+
if resultState.EndpointModelDeployment.MinReplicas != int(minReplicas) {
1495+
t.Errorf("Expected MinReplicas %d, got %d", minReplicas, resultState.EndpointModelDeployment.MinReplicas)
1496+
}
1497+
if resultState.EndpointModelDeployment.MaxReplicas != int(maxReplicas) {
1498+
t.Errorf("Expected MaxReplicas %d, got %d", maxReplicas, resultState.EndpointModelDeployment.MaxReplicas)
1499+
}
1500+
if resultState.EndpointModelDeployment.TrafficPercent != int(trafficPercent) {
1501+
t.Errorf("Expected TrafficPercent %d, got %d", trafficPercent, resultState.EndpointModelDeployment.TrafficPercent)
1502+
}
1503+
1504+
// Validate original inputs are preserved
1505+
if result.Inputs.ProjectID != projectID {
1506+
t.Errorf("Expected Inputs.ProjectID %s, got %s", projectID, result.Inputs.ProjectID)
1507+
}
1508+
if result.Inputs.Region != region {
1509+
t.Errorf("Expected Inputs.Region %s, got %s", region, result.Inputs.Region)
1510+
}
1511+
}

‎pkg/services/modelendpointget.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package services
33
import (
44
"context"
55
"fmt"
6+
"strings"
67

78
"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
89
)
@@ -32,8 +33,13 @@ func NewVertexEndpointModelGetter(endpointClient VertexEndpointClient, projectID
3233
// Get retrieves an endpoint and finds the specified deployed model within it.
3334
// Returns the endpoint, the deployed model (if found), and any error.
3435
func (g *VertexEndpointModelGetter) Get(ctx context.Context, endpointName, deployedModelID string) (*aiplatformpb.Endpoint, *aiplatformpb.DeployedModel, error) {
36+
endpointFullName := endpointName
37+
if !strings.HasPrefix(endpointName, "projects/") {
38+
endpointFullName = fmt.Sprintf("projects/%s/locations/%s/endpoints/%s",
39+
g.projectID, g.region, endpointName)
40+
}
3541
getReq := &aiplatformpb.GetEndpointRequest{
36-
Name: endpointName,
42+
Name: endpointFullName,
3743
}
3844

3945
endpoint, err := g.endpointClient.GetEndpoint(ctx, getReq)

0 commit comments

Comments
 (0)