Skip to content

Commit 7c21d75

Browse files
james-zhou-inspire11james-zhou-inspire11
committed
feat(aiPlatformNotebook) add AI Platform Notebook
1 parent 402ddf7 commit 7c21d75

16 files changed

Lines changed: 393 additions & 0 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
"@google-cloud/kms": "^2.10.0",
4646
"@google-cloud/logging": "^9.6.4",
4747
"@google-cloud/monitoring": "^2.3.5",
48+
"@google-cloud/notebooks": "^1.3.1",
4849
"@google-cloud/resource-manager": "^3.0.0",
4950
"@google-cloud/secret-manager": "^3.10.1",
5051
"@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,5 +47,6 @@ export default {
4747
[services.dataprocAutoscalingPolicy]: 'gcpDataprocAutoscalingPolicy',
4848
[services.dataprocJob]: 'gcpDataprocJob',
4949
[services.dataprocWorkflowTemplate]: 'gcpDataprocWorkflowTemplate',
50+
[services.aiPlatformNotebook]: 'gcpAiPlatformNotebook',
5051
tag: 'gcpTag',
5152
}

‎src/enums/serviceMap.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ import GcpDataprocCluster from '../services/dataprocCluster'
4343
import GcpDataprocAutoscalingPolicy from '../services/dataprocAutoscalingPolicy'
4444
import GcpDataprocJob from '../services/dataprocJob'
4545
import GcpDataprocWorkflowTemplate from '../services/dataprocWorkflowTemplate'
46+
import GcpAiPlatformNotebook from '../services/aiPlatformNotebook'
47+
4648
/**
4749
* serviceMap is an object that contains all currently supported services
4850
* serviceMap is used by the serviceFactory to produce instances of service classes
@@ -91,5 +93,6 @@ export default {
9193
[services.dataprocAutoscalingPolicy]: GcpDataprocAutoscalingPolicy,
9294
[services.dataprocJob]: GcpDataprocJob,
9395
[services.dataprocWorkflowTemplate]: GcpDataprocWorkflowTemplate,
96+
[services.aiPlatformNotebook]: GcpAiPlatformNotebook,
9497
tag: GcpTag,
9598
}

‎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+
}`;

‎src/services/aiPlatformNotebook/schema.graphql‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
type gcpAiPlatformNotebookDisk
2+
@generate(
3+
query: { get: true, query: true, aggregate: true }
4+
mutation: { add: true, delete: false }
5+
)
6+
@key(fields: "id") {
7+
id: String! @id
8+
autoDelete: Boolean @search
9+
boot: Boolean @search
10+
deviceName: String @search(by: [hash, regexp])
11+
diskSizeGb: String @search(by: [hash, regexp])
12+
guestOsFeaturesTypes: [String] @search(by: [hash])
13+
index: String @search(by: [hash, regexp])
14+
kind: String @search(by: [hash, regexp])
15+
licenses: [String] @search(by: [hash])
16+
mode: String @search(by: [hash, regexp])
17+
source: String @search(by: [hash, regexp])
18+
type: String @search(by: [hash, regexp])
19+
}
20+
21+
type gcpAiPlatformNotebook implements gcpBaseResource
22+
@generate(
23+
query: { get: true, query: true, aggregate: true }
24+
mutation: { add: true, delete: false }
25+
)
26+
@key(fields: "id") {
27+
vmImageProject: String @search(by: [hash, regexp])
28+
vmImageImageName: String @search(by: [hash, regexp])
29+
vmImageImageFamily: String @search(by: [hash, regexp])
30+
containerImageRepository: String @search(by: [hash, regexp])
31+
containerImageTag: String @search(by: [hash, regexp])
32+
postStartupScript: String @search(by: [hash, regexp])
33+
proxyUri: String @search(by: [hash, regexp])
34+
instanceOwners: [String] @search(by: [hash])
35+
serviceAccount: String @search(by: [hash, regexp])
36+
machineType: String @search(by: [hash, regexp])
37+
acceleratorConfigType: String @search(by: [hash, regexp])
38+
acceleratorConfigCoreCount: String @search(by: [hash, regexp])
39+
state: String @search(by: [hash, regexp])
40+
installGpuDriver: Boolean @search
41+
customGpuDriverPath: String @search(by: [hash, regexp])
42+
bootDiskType: String @search(by: [hash, regexp])
43+
bootDiskSizeGb: String @search(by: [hash, regexp])
44+
dataDiskType: String @search(by: [hash, regexp])
45+
dataDiskSizeGb: String @search(by: [hash, regexp])
46+
noRemoveDataDisk: Boolean @search
47+
diskEncryption: String @search(by: [hash, regexp])
48+
noPublicIp: Boolean @search
49+
noProxyAccess: Boolean @search
50+
metadata: [gcpKeyValue]
51+
createTime: String @search(by: [hash, regexp])
52+
updateTime: String @search(by: [hash, regexp])
53+
network: [gcpNetwork] @hasInverse(field: aiPlatformNotebooks)
54+
subnet: [gcpSubnet] @hasInverse(field: aiPlatformNotebooks)
55+
kmsCryptoKeys: [gcpKmsCryptoKey] @hasInverse(field: aiPlatformNotebooks)
56+
project: [gcpProject] @hasInverse(field: aiPlatformNotebooks)
57+
}

0 commit comments

Comments
 (0)