Skip to content

Commit d70aedc

Browse files
tyler-dunkeltyler-dunkel
committed
Merge branch 'feature/CG-664-ai-platform-notebook' into 'main'
Resolve CG-664 "Feature/ ai platform notebook" Closes CG-664 See merge request auto-cloud/cloudgraph/provider/cloudgraph-provider-gcp!49
2 parents 96a32a6 + b1a69af commit d70aedc

17 files changed

Lines changed: 442 additions & 49 deletions

File tree

‎README.md‎

Lines changed: 50 additions & 49 deletions
Large diffs are not rendered by default.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
"@google-cloud/kms": "^2.10.0",
4949
"@google-cloud/logging": "^9.6.4",
5050
"@google-cloud/monitoring": "^2.3.5",
51+
"@google-cloud/notebooks": "^1.3.1",
5152
"@google-cloud/resource-manager": "^3.0.0",
5253
"@google-cloud/secret-manager": "^3.10.1",
5354
"@google-cloud/storage": "^5.16.1",

‎src/enums/schemasMap.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ export default {
4747
[services.dataprocAutoscalingPolicy]: 'gcpDataprocAutoscalingPolicy',
4848
[services.dataprocJob]: 'gcpDataprocJob',
4949
[services.dataprocWorkflowTemplate]: 'gcpDataprocWorkflowTemplate',
50+
[services.aiPlatformNotebook]: 'gcpAiPlatformNotebook',
5051
[services.apiGatewayGateway]: 'gcpApiGatewayGateway',
5152
[services.apiGatewayApi]: 'gcpApiGatewayApi',
5253
[services.apiGatewayApiConfig]: 'gcpApiGatewayApiConfig',

‎src/enums/serviceMap.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import GcpDataprocCluster from '../services/dataprocCluster'
4242
import GcpDataprocAutoscalingPolicy from '../services/dataprocAutoscalingPolicy'
4343
import GcpDataprocJob from '../services/dataprocJob'
4444
import GcpDataprocWorkflowTemplate from '../services/dataprocWorkflowTemplate'
45+
import GcpAiPlatformNotebook from '../services/aiPlatformNotebook'
4546
import GcpApiGatewayGateway from '../services/apiGatewayGateway'
4647
import GcpApiGatewayApi from '../services/apiGatewayApi'
4748
import GcpApiGatewayApiConfig from '../services/apiGatewayApiConfig'
@@ -97,6 +98,7 @@ export default {
9798
[services.dataprocAutoscalingPolicy]: GcpDataprocAutoscalingPolicy,
9899
[services.dataprocJob]: GcpDataprocJob,
99100
[services.dataprocWorkflowTemplate]: GcpDataprocWorkflowTemplate,
101+
[services.aiPlatformNotebook]: GcpAiPlatformNotebook,
100102
[services.apiGatewayGateway]: GcpApiGatewayGateway,
101103
[services.apiGatewayApi]: GcpApiGatewayApi,
102104
[services.apiGatewayApiConfig]: GcpApiGatewayApiConfig,

‎src/enums/services.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ export default {
44
assets: 'assets',
55
// ai: 'ai',
66
// aiPlatform: 'ai-platform',
7+
aiPlatformNotebook: 'aiPlatformNotebooks',
78
// ml: 'ml',
89
// mlEngine: 'ml-engine',
910
// notebooks: 'notebooks',

‎src/services/aiPlatformNotebook/connections.ts‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import { ServiceConnection } from '@cloudgraph/sdk'
2+
import { RawGcpAiPlatformNotebook } from './data'
3+
import services from '../../enums/services'
4+
import { RawGcpKmsCryptoKey } from '../kmsCryptoKey/data'
5+
import { GLOBAL_REGION } from '../../config/constants'
6+
7+
export default ({
8+
service,
9+
data,
10+
region,
11+
}: {
12+
service: RawGcpAiPlatformNotebook
13+
data: { name: string; data: { [property: string]: any[] } }[]
14+
region: string
15+
}): {
16+
[property: string]: ServiceConnection[]
17+
} => {
18+
const { id, kmsKey } = service
19+
const connections: ServiceConnection[] = []
20+
21+
/**
22+
* Find Kms Crypto Keys
23+
*/
24+
const cryptoKeys: {
25+
name: string
26+
data: { [property: string]: any[] }
27+
} = data.find(({ name }) => name === services.kmsCryptoKeys)
28+
29+
const regions = [region, GLOBAL_REGION]
30+
for (const reg of regions) {
31+
if (cryptoKeys?.data?.[reg]) {
32+
const cryptoKey = cryptoKeys.data[reg].find(
33+
({ id: cryptoKeyId }: RawGcpKmsCryptoKey) => cryptoKeyId === kmsKey
34+
)
35+
36+
if (cryptoKey) {
37+
connections.push({
38+
id: cryptoKey.id,
39+
resourceType: services.kmsCryptoKeys,
40+
relation: 'child',
41+
field: 'kmsCryptoKeys',
42+
})
43+
}
44+
}
45+
}
46+
47+
const result = {
48+
[id]: connections,
49+
}
50+
return result
51+
}

‎src/services/aiPlatformNotebook/data.ts‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
import groupBy from 'lodash/groupBy'
2+
import { NotebookServiceClient } from '@google-cloud/notebooks'
3+
import CloudGraph from '@cloudgraph/sdk'
4+
import { google } from '@google-cloud/notebooks/build/protos/protos'
5+
import gcpLoggerText from '../../properties/logger'
6+
import { GcpServiceInput } from '../../types'
7+
import { generateGcpErrorLog, initTestEndpoint } from '../../utils'
8+
import zones from '../../enums/zones'
9+
10+
const lt = { ...gcpLoggerText }
11+
const { logger } = CloudGraph
12+
const serviceName = 'AI Platform Notebook'
13+
const apiEndpoint = initTestEndpoint(serviceName)
14+
15+
// TODO update SDK to stable version when available
16+
export interface RawGcpAiPlatformNotebook extends
17+
Omit<google.cloud.notebooks.v1beta1.IInstance, 'labels' | 'network' | 'subnet'> {
18+
id: string,
19+
region: string,
20+
projectId: string,
21+
Labels: { [key: string]: string },
22+
network: string[],
23+
subnet: string[],
24+
}
25+
26+
export default async ({
27+
regions,
28+
config,
29+
}: GcpServiceInput): Promise<{
30+
[region: string]: RawGcpAiPlatformNotebook[]
31+
}> => {
32+
const { projectId } = config
33+
34+
const notebookData: RawGcpAiPlatformNotebook[] = []
35+
const notebookClient = new NotebookServiceClient({ ...config, apiEndpoint })
36+
37+
for (const region of regions.split(',')) {
38+
for (const zone of zones.filter(zone => zone.indexOf(region) !== -1)) {
39+
try {
40+
const instanceNnotebookIter = await notebookClient.listInstancesAsync({
41+
parent: `projects/${projectId}/locations/${zone}`,
42+
})
43+
for await (const {name, network, subnet, labels, ...notebook} of instanceNnotebookIter) {
44+
notebookData.push({
45+
id: name,
46+
projectId,
47+
region,
48+
Labels: labels,
49+
network: [network],
50+
subnet: [subnet],
51+
...notebook,
52+
})
53+
}
54+
} catch (error) {
55+
generateGcpErrorLog(serviceName, 'aiPlatformNotebook:listInstancesAsync', error)
56+
}
57+
}
58+
}
59+
60+
logger.debug(lt.foundResources(serviceName, notebookData.length))
61+
62+
return groupBy(notebookData, 'region')
63+
}

‎src/services/aiPlatformNotebook/format.ts‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
import cuid from 'cuid'
2+
import { google } from '@google-cloud/notebooks/build/protos/protos'
3+
import {
4+
GcpAiPlatformNotebook,
5+
GcpAiPlatformNotebookDisk,
6+
} from '../../types/generated'
7+
import { enumKeyToString, formatKeyValueMap, formatLabelsFromMap } from '../../utils/format'
8+
import { toISOString } from '../../utils/dateutils'
9+
import { RawGcpAiPlatformNotebook } from './data'
10+
11+
12+
// TODO disk object is not available in the v1beta1 version.
13+
// This definition is base on unused type of v1 version and SDK document:
14+
// https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.instances
15+
const formatDisk = (disk: google.cloud.notebooks.v1.Instance.IDisk): GcpAiPlatformNotebookDisk => {
16+
const {
17+
autoDelete,
18+
boot,
19+
deviceName,
20+
diskSizeGb,
21+
guestOsFeatures = [],
22+
index,
23+
kind,
24+
licenses = [],
25+
mode,
26+
source,
27+
type,
28+
} = disk
29+
30+
return {
31+
id: cuid(),
32+
autoDelete,
33+
boot,
34+
deviceName,
35+
diskSizeGb: diskSizeGb?.toString() || '0',
36+
guestOsFeaturesTypes: guestOsFeatures?.map(f => f?.type),
37+
index: index?.toString() || '0',
38+
kind,
39+
licenses,
40+
mode,
41+
source,
42+
type,
43+
}
44+
}
45+
46+
// TODO object definition is based on the v1beta1 version
47+
// Couple fields are missing from the document
48+
// Need to update the SDK once a stable version is released
49+
export default ({
50+
service,
51+
account,
52+
region,
53+
}: {
54+
service: RawGcpAiPlatformNotebook
55+
account: string
56+
region: string
57+
}): GcpAiPlatformNotebook => {
58+
const {
59+
id,
60+
projectId,
61+
Labels = {},
62+
name,
63+
vmImage = {},
64+
containerImage = {},
65+
postStartupScript,
66+
proxyUri,
67+
instanceOwners = [],
68+
serviceAccount,
69+
// serviceAccountScopes = [], // v1
70+
machineType,
71+
acceleratorConfig = {},
72+
state,
73+
installGpuDriver,
74+
customGpuDriverPath,
75+
bootDiskType,
76+
bootDiskSizeGb,
77+
dataDiskType,
78+
dataDiskSizeGb,
79+
noRemoveDataDisk,
80+
diskEncryption,
81+
// disks = [], // v1
82+
// shieldedInstanceConfig, // v1
83+
noPublicIp,
84+
noProxyAccess,
85+
// network,
86+
// subnet,
87+
metadata,
88+
// tags = [], // v1
89+
// upgradeHistory, // v1
90+
// nicType, // v1
91+
// reservationAffinity, // v1
92+
createTime,
93+
updateTime,
94+
} = service
95+
96+
return {
97+
id: id || cuid(),
98+
projectId,
99+
region,
100+
labels: formatLabelsFromMap(Labels),
101+
name,
102+
vmImageProject: vmImage?.project || '',
103+
vmImageImageName: vmImage?.imageName || '',
104+
vmImageImageFamily: vmImage?.imageFamily || '',
105+
containerImageRepository: containerImage?.repository || '',
106+
containerImageTag: containerImage?.tag || '',
107+
postStartupScript,
108+
proxyUri,
109+
instanceOwners,
110+
serviceAccount,
111+
machineType,
112+
acceleratorConfigType: enumKeyToString(google.cloud.notebooks.v1beta1.Instance.AcceleratorType, acceleratorConfig?.type),
113+
acceleratorConfigCoreCount: acceleratorConfig?.coreCount?.toString() || '0',
114+
state: enumKeyToString(google.cloud.notebooks.v1.Instance.State, state),
115+
installGpuDriver,
116+
customGpuDriverPath,
117+
bootDiskType: enumKeyToString(google.cloud.notebooks.v1.Instance.DiskType, bootDiskType),
118+
bootDiskSizeGb: bootDiskSizeGb?.toString() || '0',
119+
dataDiskType: enumKeyToString(google.cloud.notebooks.v1.Instance.DiskType, dataDiskType),
120+
dataDiskSizeGb: dataDiskSizeGb?.toString() || '0',
121+
noRemoveDataDisk,
122+
diskEncryption: enumKeyToString(google.cloud.notebooks.v1.Instance.DiskEncryption, diskEncryption),
123+
noPublicIp,
124+
noProxyAccess,
125+
metadata: formatKeyValueMap(metadata),
126+
createTime: toISOString(createTime?.seconds?.toString()),
127+
updateTime: toISOString(updateTime?.seconds?.toString()),
128+
// kmsKey - connection
129+
// network - connection
130+
// subnet - connection
131+
}
132+
}

‎src/services/aiPlatformNotebook/index.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
import { Service } from '@cloudgraph/sdk'
2+
import BaseService from '../base'
3+
import format from './format'
4+
import getData from './data'
5+
import getConnections from './connections'
6+
import mutation from './mutation'
7+
8+
export default class GcpAiPlatformNotebook extends BaseService implements Service {
9+
format = format.bind(this)
10+
11+
getData = getData.bind(this)
12+
13+
getConnections = getConnections.bind(this)
14+
15+
mutation = mutation;
16+
}

‎src/services/aiPlatformNotebook/mutation.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
export default `mutation($input: [AddgcpAiPlatformNotebookInput!]!) {
2+
addgcpAiPlatformNotebook(input: $input, upsert: true) {
3+
numUids
4+
}
5+
}`;

0 commit comments

Comments
 (0)