Skip to content

Commit 415a3c3

Browse files
omerdemirokactions-user
omerdemirok
authored and
actions-user
committed
Refactor gcp adapter files (#2815)
Refactor GCP dynamic adapters and their associated tests into individual files to improve maintainability. --- Linear Issue: [ENG-1578](https://linear.app/overmind/issue/ENG-1578/refactor-adapter-meta-and-blast-propagations-into-individual-files) <a href="https://cursor.com/background-agent?bcId=bc-e411289d-9650-4181-89ac-5b3996f6a0b7"><picture><source media="(prefers-color-scheme: dark)" srcset="https://cursor.com/open-in-cursor-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://cursor.com/open-in-cursor-light.svg"><img alt="Open in Cursor" src="https://cursor.com/open-in-cursor.svg"></picture></a>&nbsp;<a href="https://cursor.com/agents?id=bc-e411289d-9650-4181-89ac-5b3996f6a0b7"><picture><source media="(prefers-color-scheme: dark)" srcset="https://cursor.com/open-in-web-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://cursor.com/open-in-web-light.svg"><img alt="Open in Web" src="https://cursor.com/open-in-web.svg"></picture></a> GitOrigin-RevId: 264916f5a3b26a7f86af5297773ce3b5824c80e7
1 parent 7c5a742 commit 415a3c3

File tree

84 files changed

+7298
-2495
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

84 files changed

+7298
-2495
lines changed

‎sources/gcp/dynamic/adapter_test.go‎

Lines changed: 5 additions & 869 deletions
Large diffs are not rendered by default.

‎sources/gcp/dynamic/adapters/ai-platform-custom-job.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package adapters
2+
3+
import (
4+
"github.com/overmindtech/cli/sdp-go"
5+
gcpshared "github.com/overmindtech/cli/sources/gcp/shared"
6+
)
7+
8+
// AI Platform Custom Job adapter for Vertex AI custom training jobs
9+
// There are multiple service endpoints: https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints
10+
// We stick to the default one for now: https://aiplatform.googleapis.com
11+
// Other endpoints are in the form of https://{region}-aiplatform.googleapis.com
12+
// If we use the default endpoint the location must be set to `global`.
13+
// So, for simplicity, we can get custom jobs by their name globally, list globally,
14+
// otherwise we have to check the validity of the location and use the regional endpoint.
15+
var _ = registerableAdapter{
16+
sdpType: gcpshared.AIPlatformCustomJob,
17+
meta: gcpshared.AdapterMeta{
18+
SDPAdapterCategory: sdp.AdapterCategory_ADAPTER_CATEGORY_AI,
19+
Scope: gcpshared.ScopeProject,
20+
// Vertex AI API must be enabled for the project!
21+
// Reference: https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.customJobs/get
22+
// https://aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/customJobs/{customJob}
23+
GetEndpointBaseURLFunc: gcpshared.ProjectLevelEndpointFuncWithSingleQuery("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/customJobs/%s"),
24+
// Reference: https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.customJobs/list
25+
// https://aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/customJobs
26+
// Expected location is `global` for the default endpoint.
27+
ListEndpointFunc: gcpshared.ProjectLevelListFunc("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/customJobs"),
28+
UniqueAttributeKeys: []string{"customJobs"},
29+
IAMPermissions: []string{"aiplatform.customJobs.get", "aiplatform.customJobs.list"},
30+
PredefinedRole: "roles/aiplatform.viewer",
31+
},
32+
blastPropagation: map[string]*gcpshared.Impact{
33+
// The Cloud KMS key that will be used to encrypt the output artifacts.
34+
"encryptionSpec.kmsKeyName": {
35+
Description: "If the Cloud KMS CryptoKey is updated: The CustomJob may not be able to access encrypted output artifacts. If the CustomJob is updated: The CryptoKey remains unaffected.",
36+
ToSDPItemType: gcpshared.CloudKMSCryptoKey,
37+
BlastPropagation: gcpshared.ImpactInOnly,
38+
},
39+
// The full name of the network to which the job should be peered.
40+
"jobSpec.network": {
41+
Description: "If the Compute Network is deleted or updated: The CustomJob may lose connectivity or fail to run as expected. If the CustomJob is updated: The network remains unaffected.",
42+
ToSDPItemType: gcpshared.ComputeNetwork,
43+
BlastPropagation: gcpshared.ImpactInOnly,
44+
},
45+
// The service account that the job runs as.
46+
"jobSpec.serviceAccount": {
47+
Description: "If the IAM Service Account is deleted or updated: The CustomJob may fail to run or lose permissions. If the CustomJob is updated: The service account remains unaffected.",
48+
ToSDPItemType: gcpshared.IAMServiceAccount,
49+
BlastPropagation: gcpshared.ImpactInOnly,
50+
},
51+
// The Cloud Storage location to store the output of this CustomJob.
52+
"jobSpec.baseOutputDirectory.gcsOutputDirectory": {
53+
Description: "If the Storage Bucket is deleted or updated: The CustomJob may fail to write outputs. If the CustomJob is updated: The bucket remains unaffected.",
54+
ToSDPItemType: gcpshared.StorageBucket,
55+
BlastPropagation: gcpshared.ImpactInOnly,
56+
},
57+
// Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs.
58+
"jobSpec.tensorboard": {
59+
Description: "If the Vertex AI Tensorboard is deleted or updated: The CustomJob may fail to upload logs or lose access to previous logs. If the CustomJob is updated: The tensorboard remains unaffected.",
60+
ToSDPItemType: gcpshared.AIPlatformTensorBoard,
61+
BlastPropagation: gcpshared.ImpactInOnly,
62+
},
63+
// Optional. The name of an experiment to associate with the CustomJob.
64+
"jobSpec.experiment": {
65+
Description: "If the Vertex AI Experiment is deleted or updated: The CustomJob may lose experiment tracking or association. If the CustomJob is updated: The experiment remains unaffected.",
66+
ToSDPItemType: gcpshared.AIPlatformExperiment,
67+
BlastPropagation: gcpshared.ImpactInOnly,
68+
},
69+
// Optional. The name of an experiment run to associate with the CustomJob.
70+
"jobSpec.experimentRun": {
71+
Description: "If the Vertex AI ExperimentRun is deleted or updated: The CustomJob may lose run tracking or association. If the CustomJob is updated: The experiment run remains unaffected.",
72+
ToSDPItemType: gcpshared.AIPlatformExperimentRun,
73+
BlastPropagation: gcpshared.ImpactInOnly,
74+
},
75+
// Optional. The name of a model to upload the trained Model to upon job completion.
76+
"jobSpec.models": {
77+
Description: "If the Vertex AI Model is deleted or updated: The CustomJob may fail to upload or associate the trained model. If the CustomJob is updated: The model remains unaffected.",
78+
ToSDPItemType: gcpshared.AIPlatformModel,
79+
BlastPropagation: gcpshared.ImpactInOnly,
80+
},
81+
},
82+
terraformMapping: gcpshared.TerraformMapping{
83+
Description: "There is no terraform resource for this type.",
84+
},
85+
}.Register()

‎sources/gcp/dynamic/adapters/ai-platform-custom-job_test.go‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
package adapters_test
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"net/http"
7+
"testing"
8+
9+
"google.golang.org/api/aiplatform/v1"
10+
11+
"github.com/overmindtech/cli/discovery"
12+
"github.com/overmindtech/cli/sdp-go"
13+
"github.com/overmindtech/cli/sources/gcp/dynamic"
14+
gcpshared "github.com/overmindtech/cli/sources/gcp/shared"
15+
"github.com/overmindtech/cli/sources/shared"
16+
)
17+
18+
func TestAIPlatformCustomJob(t *testing.T) {
19+
ctx := context.Background()
20+
projectID := "test-project"
21+
linker := gcpshared.NewLinker()
22+
jobID := "test-job"
23+
24+
customJob := &aiplatform.GoogleCloudAiplatformV1CustomJob{
25+
Name: fmt.Sprintf("projects/%s/locations/global/customJobs/%s", projectID, jobID),
26+
JobSpec: &aiplatform.GoogleCloudAiplatformV1CustomJobSpec{
27+
ServiceAccount: "[email protected]",
28+
Network: fmt.Sprintf("projects/%s/global/networks/default", projectID),
29+
},
30+
EncryptionSpec: &aiplatform.GoogleCloudAiplatformV1EncryptionSpec{
31+
KmsKeyName: "projects/test-project/locations/global/keyRings/my-keyring/cryptoKeys/my-key",
32+
},
33+
}
34+
35+
jobList := &aiplatform.GoogleCloudAiplatformV1ListCustomJobsResponse{
36+
CustomJobs: []*aiplatform.GoogleCloudAiplatformV1CustomJob{customJob},
37+
}
38+
39+
sdpItemType := gcpshared.AIPlatformCustomJob
40+
41+
expectedCallAndResponses := map[string]shared.MockResponse{
42+
fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/customJobs/%s", projectID, jobID): {
43+
StatusCode: http.StatusOK,
44+
Body: customJob,
45+
},
46+
fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/customJobs", projectID): {
47+
StatusCode: http.StatusOK,
48+
Body: jobList,
49+
},
50+
}
51+
52+
t.Run("Get", func(t *testing.T) {
53+
httpCli := shared.NewMockHTTPClientProvider(expectedCallAndResponses)
54+
adapter, err := dynamic.MakeAdapter(sdpItemType, linker, httpCli, projectID)
55+
if err != nil {
56+
t.Fatalf("Failed to create adapter for %s: %v", sdpItemType, err)
57+
}
58+
59+
sdpItem, err := adapter.Get(ctx, projectID, jobID, true)
60+
if err != nil {
61+
t.Fatalf("Failed to get custom job: %v", err)
62+
}
63+
64+
if sdpItem.GetType() != sdpItemType.String() {
65+
t.Errorf("Expected type %s, got %s", sdpItemType.String(), sdpItem.GetType())
66+
}
67+
68+
t.Run("StaticTests", func(t *testing.T) {
69+
queryTests := shared.QueryTests{
70+
{
71+
// encryptionSpec.kmsKeyName
72+
ExpectedType: gcpshared.CloudKMSCryptoKey.String(),
73+
ExpectedMethod: sdp.QueryMethod_GET,
74+
ExpectedQuery: shared.CompositeLookupKey("global", "my-keyring", "my-key"),
75+
ExpectedScope: projectID,
76+
ExpectedBlastPropagation: &sdp.BlastPropagation{
77+
In: true,
78+
Out: false,
79+
},
80+
},
81+
{
82+
// jobSpec.network
83+
ExpectedType: gcpshared.ComputeNetwork.String(),
84+
ExpectedMethod: sdp.QueryMethod_GET,
85+
ExpectedQuery: "default",
86+
ExpectedScope: projectID,
87+
ExpectedBlastPropagation: &sdp.BlastPropagation{
88+
In: true,
89+
Out: false,
90+
},
91+
},
92+
{
93+
// jobSpec.serviceAccount
94+
ExpectedType: gcpshared.IAMServiceAccount.String(),
95+
ExpectedMethod: sdp.QueryMethod_GET,
96+
ExpectedQuery: "[email protected]",
97+
ExpectedScope: projectID,
98+
ExpectedBlastPropagation: &sdp.BlastPropagation{
99+
In: true,
100+
Out: false,
101+
},
102+
},
103+
}
104+
105+
shared.RunStaticTests(t, adapter, sdpItem, queryTests)
106+
})
107+
})
108+
109+
t.Run("List", func(t *testing.T) {
110+
httpCli := shared.NewMockHTTPClientProvider(expectedCallAndResponses)
111+
adapter, err := dynamic.MakeAdapter(sdpItemType, linker, httpCli, projectID)
112+
if err != nil {
113+
t.Fatalf("Failed to create adapter for %s: %v", sdpItemType, err)
114+
}
115+
116+
listable, ok := adapter.(discovery.ListableAdapter)
117+
if !ok {
118+
t.Fatalf("Adapter for %s does not implement ListableAdapter", sdpItemType)
119+
}
120+
121+
sdpItems, err := listable.List(ctx, projectID, true)
122+
if err != nil {
123+
t.Fatalf("Failed to list custom jobs: %v", err)
124+
}
125+
126+
if len(sdpItems) != 1 {
127+
t.Errorf("Expected 1 custom job, got %d", len(sdpItems))
128+
}
129+
})
130+
131+
t.Run("ErrorHandling", func(t *testing.T) {
132+
errorResponses := map[string]shared.MockResponse{
133+
fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/customJobs/%s", projectID, jobID): {
134+
StatusCode: http.StatusNotFound,
135+
Body: map[string]interface{}{"error": "Custom job not found"},
136+
},
137+
}
138+
139+
httpCli := shared.NewMockHTTPClientProvider(errorResponses)
140+
adapter, err := dynamic.MakeAdapter(sdpItemType, linker, httpCli, projectID)
141+
if err != nil {
142+
t.Fatalf("Failed to create adapter for %s: %v", sdpItemType, err)
143+
}
144+
145+
_, err = adapter.Get(ctx, projectID, jobID, true)
146+
if err == nil {
147+
t.Error("Expected error when getting non-existent custom job, but got nil")
148+
}
149+
})
150+
}

‎sources/gcp/dynamic/adapters/ai-platform-pipeline-job.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package adapters
2+
3+
import (
4+
"github.com/overmindtech/cli/sdp-go"
5+
gcpshared "github.com/overmindtech/cli/sources/gcp/shared"
6+
)
7+
8+
// AI Platform Pipeline Job adapter for Vertex AI pipeline jobs
9+
var _ = registerableAdapter{
10+
sdpType: gcpshared.AIPlatformPipelineJob,
11+
meta: gcpshared.AdapterMeta{
12+
SDPAdapterCategory: sdp.AdapterCategory_ADAPTER_CATEGORY_AI,
13+
Scope: gcpshared.ScopeProject,
14+
// When using the default endpoint, the location must be set to `global`.
15+
// Format: projects/{project}/locations/{location}/pipelineJobs/{pipelineJob}
16+
GetEndpointBaseURLFunc: gcpshared.ProjectLevelEndpointFuncWithSingleQuery("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/pipelineJobs/%s"),
17+
// Reference: https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs/list
18+
ListEndpointFunc: gcpshared.ProjectLevelListFunc("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/pipelineJobs"),
19+
UniqueAttributeKeys: []string{"pipelineJobs"},
20+
IAMPermissions: []string{"aiplatform.pipelineJobs.get", "aiplatform.pipelineJobs.list"},
21+
PredefinedRole: "roles/aiplatform.viewer",
22+
},
23+
blastPropagation: map[string]*gcpshared.Impact{
24+
// The service account that the pipeline workload runs as (root-level).
25+
"serviceAccount": gcpshared.IAMServiceAccountImpactInOnly,
26+
// The full name of the network to which the job should be peered (root-level).
27+
"network": gcpshared.ComputeNetworkImpactInOnly,
28+
// The Cloud KMS key used to encrypt PipelineJob outputs.
29+
"encryptionSpec.kmsKeyName": gcpshared.CryptoKeyImpactInOnly,
30+
// The Cloud Storage location to store the output of this PipelineJob.
31+
"runtimeConfig.gcsOutputDirectory": {
32+
Description: "If the Storage Bucket is deleted or updated: The PipelineJob may fail to write outputs. If the PipelineJob is updated: The bucket remains unaffected.",
33+
ToSDPItemType: gcpshared.StorageBucket,
34+
BlastPropagation: gcpshared.ImpactInOnly,
35+
},
36+
},
37+
terraformMapping: gcpshared.TerraformMapping{
38+
Description: "There is no terraform resource for this type.",
39+
},
40+
}.Register()

0 commit comments

Comments
 (0)