separate model from endpoint lookup

Author: davidmontoyago

pkg/resources/provider.go

@@ -7,7 +7,6 @@ import (
 	"log"
 	"time"
 
-	"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
 	"github.com/pulumi/pulumi-go-provider/infer"
 
 	"github.com/davidmontoyago/pulumi-gcp-vertex-model-deployment/pkg/services"
@@ -218,110 +217,125 @@ func (v VertexModelDeployment) Read(
 	state := req.State
 
 	if req.State.ModelName != "" {
-		// Lookup the model
+		// Read the model from the registry
+
 		modelClientFactory := v.getModelClientFactory()
 		modelClient, err := modelClientFactory(ctx, req.State.Region)
 		if err != nil {
-			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, err
+			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, fmt.Errorf("failed to create model client: %w", err)
 		}
 		defer func() {
 			if closeErr := modelClient.Close(); closeErr != nil {
 				log.Printf("failed to close model client: %v", closeErr)
 			}
 		}()
 
-		modelGetter := services.NewVertexModelGet(ctx, modelClient, req.State.ModelName)
-		model, err := modelGetter.Get(ctx, req.State.ModelName)
+		err = readRegistryModel(ctx, modelClient, req, &state)
 		if err != nil {
-			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, err
-		}
-
-		// Update state with current model values
-		state.ModelName = model.Name
-		state.ModelArtifactsBucketURI = model.ArtifactUri
-		state.Labels = model.Labels
-
-		// Safely access ContainerSpec fields
-		if model.ContainerSpec != nil {
-			state.ModelImageURL = model.ContainerSpec.ImageUri
-			state.PredictRoute = model.ContainerSpec.PredictRoute
-			state.HealthRoute = model.ContainerSpec.HealthRoute
-		}
-
-		// Safely access PredictSchemata fields
-		if model.PredictSchemata != nil {
-			state.ModelPredictionInputSchemaURI = model.PredictSchemata.InstanceSchemaUri
-			state.ModelPredictionOutputSchemaURI = model.PredictSchemata.PredictionSchemaUri
+			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, fmt.Errorf("failed to read model from registry: %w", err)
 		}
 	}
 
 	if req.State.DeployedModelID != "" && req.State.EndpointName != "" {
-		// Lookup the endpoint if model is deployed to an endpoint
+		// Read the endpoint if model is deployed to an endpoint
 
 		endpointClientFactory := v.getEndpointClientFactory()
 		endpointClient, err := endpointClientFactory(ctx, req.State.Region)
 		if err != nil {
-			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, err
+			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, fmt.Errorf("failed to create endpoint client: %w", err)
 		}
 		defer func() {
 			if closeErr := endpointClient.Close(); closeErr != nil {
 				log.Printf("failed to close endpoint client: %v", closeErr)
 			}
 		}()
 
-		getReq := &aiplatformpb.GetEndpointRequest{
-			Name: fmt.Sprintf("projects/%s/locations/%s/endpoints/%s",
-				req.State.ProjectID, req.State.Region, req.State.EndpointName),
-		}
-
-		endpoint, err := endpointClient.GetEndpoint(ctx, getReq)
+		err = readEndpointModel(ctx, endpointClient, req, &state)
 		if err != nil {
-			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, err
+			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, fmt.Errorf("failed to read model endpoint: %w", err)
 		}
+	}
 
-		// Verify the deployed model still exists and update its properties
-		var foundDeployedModel *aiplatformpb.DeployedModel
-		for _, deployedModel := range endpoint.DeployedModels {
-			if deployedModel.Id == req.State.DeployedModelID {
-				foundDeployedModel = deployedModel
+	return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{
+		Inputs: req.Inputs,
+		State:  state,
+	}, nil
+}
 
-				break
-			}
+func readEndpointModel(ctx context.Context,
+	endpointClient services.VertexEndpointClient,
+	req infer.ReadRequest[VertexModelDeploymentArgs, VertexModelDeploymentState],
+	state *VertexModelDeploymentState) error {
+
+	endpointGetter := services.NewVertexEndpointModelGetter(endpointClient, req.State.ProjectID, req.State.Region)
+	endpoint, foundDeployedModel, err := endpointGetter.Get(ctx, req.State.EndpointName, req.State.DeployedModelID)
+	if err != nil {
+		return err
+	}
+
+	if foundDeployedModel == nil {
+		// Model is no longer deployed - return empty response to indicate resource doesn't exist
+		return nil
+	}
+
+	// Update state with current endpoint and deployed model information
+	state.EndpointName = endpoint.Name
+	state.DeployedModelID = foundDeployedModel.Id
+
+	// Update endpoint deployment configuration with current values if available
+	if state.EndpointModelDeployment == nil {
+		return nil
+	}
+
+	// Extract current deployment configuration from the deployed model
+	if dedicatedResources := foundDeployedModel.GetDedicatedResources(); dedicatedResources != nil {
+		if machineSpec := dedicatedResources.MachineSpec; machineSpec != nil {
+			state.EndpointModelDeployment.MachineType = machineSpec.MachineType
 		}
+		state.EndpointModelDeployment.MinReplicas = int(dedicatedResources.MinReplicaCount)
+		state.EndpointModelDeployment.MaxReplicas = int(dedicatedResources.MaxReplicaCount)
+	}
 
-		if foundDeployedModel == nil {
-			// Model is no longer deployed - return empty response to indicate resource doesn't exist
-			return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{}, nil
+	// Update traffic percentage from endpoint's traffic split if available
+	if endpoint.TrafficSplit != nil {
+		if trafficPercent, exists := endpoint.TrafficSplit[foundDeployedModel.Id]; exists {
+			state.EndpointModelDeployment.TrafficPercent = int(trafficPercent)
 		}
+	}
 
-		// Update state with current endpoint and deployed model information
-		state.EndpointName = endpoint.Name
-		state.DeployedModelID = foundDeployedModel.Id
-
-		// Update endpoint deployment configuration with current values if available
-		if state.EndpointModelDeployment != nil && foundDeployedModel.PredictionResources != nil {
-			// Extract current deployment configuration from the deployed model
-			if dedicatedResources := foundDeployedModel.GetDedicatedResources(); dedicatedResources != nil {
-				if machineSpec := dedicatedResources.MachineSpec; machineSpec != nil {
-					state.EndpointModelDeployment.MachineType = machineSpec.MachineType
-				}
-				state.EndpointModelDeployment.MinReplicas = int(dedicatedResources.MinReplicaCount)
-				state.EndpointModelDeployment.MaxReplicas = int(dedicatedResources.MaxReplicaCount)
-			}
+	return nil
+}
 
-			// Update traffic percentage from endpoint's traffic split if available
-			if endpoint.TrafficSplit != nil {
-				if trafficPercent, exists := endpoint.TrafficSplit[foundDeployedModel.Id]; exists {
-					state.EndpointModelDeployment.TrafficPercent = int(trafficPercent)
-				}
-			}
-		}
+func readRegistryModel(ctx context.Context,
+	modelClient services.VertexModelClient,
+	req infer.ReadRequest[VertexModelDeploymentArgs,
+		VertexModelDeploymentState], state *VertexModelDeploymentState) error {
+
+	modelGetter := services.NewVertexModelGet(ctx, modelClient, req.State.ModelName)
+	model, err := modelGetter.Get(ctx, req.State.ModelName)
+	if err != nil {
+		return fmt.Errorf("failed to get model: %w", err)
 	}
 
-	return infer.ReadResponse[VertexModelDeploymentArgs, VertexModelDeploymentState]{
-		Inputs: req.Inputs,
-		State:  state,
-	}, nil
+	// Update state with current model values
+	state.ModelName = model.Name
+	state.ModelArtifactsBucketURI = model.ArtifactUri
+	state.Labels = model.Labels
+
+	// Safely access ContainerSpec fields
+	if model.ContainerSpec != nil {
+		state.ModelImageURL = model.ContainerSpec.ImageUri
+		state.PredictRoute = model.ContainerSpec.PredictRoute
+		state.HealthRoute = model.ContainerSpec.HealthRoute
+	}
+
+	// Safely access PredictSchemata fields
+	if model.PredictSchemata != nil {
+		state.ModelPredictionInputSchemaURI = model.PredictSchemata.InstanceSchemaUri
+		state.ModelPredictionOutputSchemaURI = model.PredictSchemata.PredictionSchemaUri
+	}
+
+	return nil
 }
 
 // testFactoryRegistry holds test factories for dependency injection during testing

pkg/resources/provider_test.go

@@ -14,16 +14,16 @@ import (
 )
 
 const (
-	testProjectID                     = "test-project"
-	testRegion                        = "us-central1"
-	testEndpointID                    = "test-endpoint"
-	testModelImageURL                 = "gcr.io/test-project/custom-model:latest"
-	testModelArtifactsBucketURI       = "gs://test-bucket/model-artifacts/"
-	testModelPredictionInputSchemaURI = "gs://test-bucket/schemas/input_schema.json"
+	testProjectID                      = "test-project"
+	testRegion                         = "us-central1"
+	testEndpointID                     = "test-endpoint"
+	testModelImageURL                  = "gcr.io/test-project/custom-model:latest"
+	testModelArtifactsBucketURI        = "gs://test-bucket/model-artifacts/"
+	testModelPredictionInputSchemaURI  = "gs://test-bucket/schemas/input_schema.json"
 	testModelPredictionOutputSchemaURI = "gs://test-bucket/schemas/output_schema.json"
-	testEndpointPath                  = "projects/test-project/locations/us-central1/endpoints/test-endpoint"
-	testModelName                     = "projects/test-project/locations/us-central1/models/1234567890"
-	testCreateTime                    = "2023-10-15T10:30:00Z"
+	testEndpointPath                   = "projects/test-project/locations/us-central1/endpoints/test-endpoint"
+	testModelName                      = "projects/test-project/locations/us-central1/models/1234567890"
+	testCreateTime                     = "2023-10-15T10:30:00Z"
 )
 
 //nolint:paralleltest,tparallel // Cannot run in parallel due to shared testFactoryRegistry

pkg/services/modeldeploy.go renamed to pkg/services/modelendpointdeploy.go

@@ -0,0 +1,61 @@
+package services
+
+import (
+	"context"
+	"fmt"
+
+	"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
+)
+
+// EndpointModelGetter allows getting endpoints and their deployed models from the registry.
+type EndpointModelGetter interface {
+	Get(ctx context.Context, endpointName, deployedModelID string) (*aiplatformpb.Endpoint, *aiplatformpb.DeployedModel, error)
+	Close() error
+}
+
+// VertexEndpointModelGetter implements the EndpointModelGetter interface for Vertex AI.
+type VertexEndpointModelGetter struct {
+	endpointClient VertexEndpointClient
+	projectID      string
+	region         string
+}
+
+// NewVertexEndpointModelGetter creates a new VertexEndpointModelGetter with the provided endpoint client.
+func NewVertexEndpointModelGetter(endpointClient VertexEndpointClient, projectID, region string) *VertexEndpointModelGetter {
+	return &VertexEndpointModelGetter{
+		endpointClient: endpointClient,
+		projectID:      projectID,
+		region:         region,
+	}
+}
+
+// Get retrieves an endpoint and finds the specified deployed model within it.
+// Returns the endpoint, the deployed model (if found), and any error.
+func (g *VertexEndpointModelGetter) Get(ctx context.Context, endpointName, deployedModelID string) (*aiplatformpb.Endpoint, *aiplatformpb.DeployedModel, error) {
+	getReq := &aiplatformpb.GetEndpointRequest{
+		Name: fmt.Sprintf("projects/%s/locations/%s/endpoints/%s",
+			g.projectID, g.region, endpointName),
+	}
+
+	endpoint, err := g.endpointClient.GetEndpoint(ctx, getReq)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to get endpoint: %w", err)
+	}
+
+	// Verify the deployed model still exists and update its properties
+	var foundDeployedModel *aiplatformpb.DeployedModel
+	for _, deployedModel := range endpoint.DeployedModels {
+		if deployedModel.Id == deployedModelID {
+			foundDeployedModel = deployedModel
+
+			break
+		}
+	}
+
+	return endpoint, foundDeployedModel, nil
+}
+
+// Close closes the endpoint client.
+func (g *VertexEndpointModelGetter) Close() error {
+	return g.endpointClient.Close()
+}

pkg/services/modelendpointget.go

@@ -2,6 +2,7 @@ package services
 
 import (
 	"context"
+	"fmt"
 
 	"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
 )
@@ -34,7 +35,7 @@ func (g *VertexModelGet) Get(ctx context.Context, modelName string) (*aiplatform
 
 	model, err := g.modelClient.GetModel(ctx, getReq)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to get model: %w", err)
 	}
 
 	return model, nil

pkg/services/modelundeploy.go renamed to pkg/services/modelendpointundeploy.go

@@ -4,7 +4,7 @@ go 1.24.5
 
 require (
 	github.com/blang/semver v3.5.1+incompatible
-	github.com/pulumi/pulumi/sdk/v3 v3.196.0
+	github.com/pulumi/pulumi/sdk/v3 v3.197.0
 )
 
 require (

pkg/services/modelget.go

@@ -155,8 +155,8 @@ github.com/pulumi/appdash v0.0.0-20231130102222-75f619a67231 h1:vkHw5I/plNdTr435
 github.com/pulumi/appdash v0.0.0-20231130102222-75f619a67231/go.mod h1:murToZ2N9hNJzewjHBgfFdXhZKjY3z5cYC1VXk+lbFE=
 github.com/pulumi/esc v0.17.0 h1:oaVOIyFTENlYDuqc3pW75lQT9jb2cd6ie/4/Twxn66w=
 github.com/pulumi/esc v0.17.0/go.mod h1:XnSxlt5NkmuAj304l/gK4pRErFbtqq6XpfX1tYT9Jbc=
-github.com/pulumi/pulumi/sdk/v3 v3.196.0 h1:OwD+S4udFwxrdfw9n4dHv6gToF+SQNtggQJIfacSBYQ=
-github.com/pulumi/pulumi/sdk/v3 v3.196.0/go.mod h1:aV0+c5xpSYccWKmOjTZS9liYCqh7+peu3cQgSXu7CJw=
+github.com/pulumi/pulumi/sdk/v3 v3.197.0 h1:ZNKda7CQpfVbRS2r/7U5F+s4iejfL9HK39bXl5CCTpY=
+github.com/pulumi/pulumi/sdk/v3 v3.197.0/go.mod h1:aV0+c5xpSYccWKmOjTZS9liYCqh7+peu3cQgSXu7CJw=
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=