1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { tagRecord } from "../../Tags.ts";
10import { GcpEnvironment } from "../Environment.ts";
11import { createInternalLabels, hasAlchemyLabels, toLabels } from "../Labels.ts";
12import type { Providers } from "../Providers.ts";
13import { listLocations } from "./names.ts";
14import { resourceNameFromOperation, waitForOperation } from "./operations.ts";
15import {
16 AiPlatformNotResolved,
17 AiPlatformStillExists,
18 collectPages,
19 jsonEqual,
20 locationParent,
21 normalizeLocation,
22 parseResourceName,
23 toPhysicalId,
24 userLabels,
25 type EncryptionSpec,
26 type MachineSpec,
27 type NetworkSpec,
28 type PersistentDiskSpec,
29} from "./shared.ts";
30
31const COLLECTION = "notebookRuntimeTemplates";
32const DEFAULT_MACHINE_TYPE = "e2-standard-4";
33
34export type NotebookEnvVar = {
35 /** Environment variable name. */
36 name?: string;
37 /** Environment variable value. */
38 value?: string;
39};
40
41export type ColabImage = {
42 /** Colab Enterprise release name (e.g. `colab-enterprise-release`). */
43 releaseName?: string;
44};
45
46export type PostStartupScriptConfig = {
47 /** Inline post-startup script. */
48 postStartupScript?: string;
49 /** Cloud Storage URI of a post-startup script. */
50 postStartupScriptUrl?: string;
51 /** When the script runs (`RUN_ONCE`, `RUN_EVERY_START`). */
52 postStartupScriptBehavior?: string;
53};
54
55export type NotebookSoftwareConfig = {
56 /** Google-managed Colab image. */
57 colabImage?: ColabImage;
58 /** Environment variables passed to the runtime. */
59 env?: NotebookEnvVar[];
60 /** Post-startup script configuration. */
61 postStartupScriptConfig?: PostStartupScriptConfig;
62};
63
64export type NotebookReservationAffinity = {
65 /** Affinity type (`NO_RESERVATION`, `ANY_RESERVATION`, `SPECIFIC_RESERVATION`). */
66 reservationAffinityType?: string;
67 /** Label key of the reservation. */
68 key?: string;
69 /** Reservation resource names. */
70 values?: string[];
71};
72
73export type NotebookRuntimeTemplateProps = {
74 /**
75 * Template id (the `{notebook_runtime_template}` segment). If omitted,
76 * a unique RFC1035 name is generated. Immutable — changing it replaces
77 * the template.
78 */
79 notebookRuntimeTemplateId?: string;
80 /**
81 * Vertex AI location. Immutable — changing it replaces the template.
82 * `US-CENTRAL1` is accepted and normalized to `us-central1`.
83 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
84 */
85 location?: string;
86 /**
87 * Display name (max 128 UTF-8 characters). Defaults to the template id.
88 */
89 displayName?: string;
90 /**
91 * Human-readable description. Vertex AI cannot update it, so changing
92 * it replaces the template.
93 */
94 description?: string;
95 /**
96 * User labels. Alchemy ownership labels are merged in automatically.
97 * Vertex AI cannot update them, so changing them replaces the template.
98 */
99 labels?: Record<string, string>;
100 /**
101 * Runtime type. Immutable.
102 * `USER_DEFINED` or `ONE_CLICK`.
103 */
104 notebookRuntimeType?: string;
105 /**
106 * Machine spec. Immutable — changing it replaces the template.
107 * @default { machineType: "e2-standard-4" }
108 */
109 machineSpec?: MachineSpec;
110 /**
111 * Persistent data disk attached to the runtime.
112 */
113 dataPersistentDiskSpec?: PersistentDiskSpec;
114 /**
115 * Network spec.
116 */
117 networkSpec?: NetworkSpec;
118 /**
119 * Notebook software configuration.
120 */
121 softwareConfig?: NotebookSoftwareConfig;
122 /**
123 * Customer-managed encryption key.
124 */
125 encryptionSpec?: EncryptionSpec;
126 /**
127 * Compute Engine network tags.
128 */
129 networkTags?: string[];
130 /**
131 * Idle shutdown configuration.
132 */
133 idleShutdownConfig?: {
134 idleTimeout?: string;
135 idleShutdownDisabled?: boolean;
136 };
137 /**
138 * Shielded VM config. Immutable.
139 */
140 shieldedVmConfig?: { enableSecureBoot?: boolean };
141 /**
142 * Reservation affinity.
143 */
144 reservationAffinity?: NotebookReservationAffinity;
145};
146
147export type NotebookRuntimeTemplate = Resource<
148 "GCP.AIPlatform.NotebookRuntimeTemplate",
149 NotebookRuntimeTemplateProps,
150 {
151 /** Full resource name. */
152 name: string;
153 /** Template id (last path segment). */
154 notebookRuntimeTemplateId: string;
155 /** Project id. */
156 project: string;
157 /** Location id. */
158 location: string;
159 /** Display name. */
160 displayName: string | undefined;
161 /** Description. */
162 description: string | undefined;
163 /** User labels (Alchemy ownership labels stripped). */
164 labels: Record<string, string>;
165 /** Runtime type. */
166 notebookRuntimeType: string | undefined;
167 /** Machine type. */
168 machineType: string | undefined;
169 /** RFC3339 creation timestamp. */
170 createTime: string | undefined;
171 /** RFC3339 last-update timestamp. */
172 updateTime: string | undefined;
173 /** Etag for read-modify-write. */
174 etag: string | undefined;
175 },
176 never,
177 Providers
178>;
179
180/**
181 * A Vertex AI notebook runtime template — machine, disk, network, and
182 * software defaults used to create Colab Enterprise runtimes.
183 *
184 * Only `displayName` and `softwareConfig` update in place; changing
185 * `notebookRuntimeTemplateId`, `location`, `machineSpec`,
186 * `notebookRuntimeType`, `description`, or `labels` replaces the template.
187 *
188 * ### Creating a Template
189 * **Example:** Generated name
190 * ```typescript
191 * const template = yield* GCP.AIPlatform.NotebookRuntimeTemplate("Runtime", {
192 * machineSpec: { machineType: "e2-standard-4" },
193 * });
194 * ```
195 *
196 * **Example:** Named template with labels
197 * ```typescript
198 * const template = yield* GCP.AIPlatform.NotebookRuntimeTemplate("Runtime", {
199 * notebookRuntimeTemplateId: "colab-default",
200 * displayName: "colab default",
201 * labels: { env: "prod" },
202 * machineSpec: { machineType: "e2-standard-4" },
203 * networkSpec: { enableInternetAccess: true },
204 * });
205 * ```
206 *
207 * @resource
208 * @category AIPlatform
209 */
210export const NotebookRuntimeTemplate = Resource<NotebookRuntimeTemplate>(
211 "GCP.AIPlatform.NotebookRuntimeTemplate",
212);
213
214export class NotebookRuntimeTemplateNotResolved extends Data.TaggedError(
215 "GCP.AIPlatform.NotebookRuntimeTemplateNotResolved",
216)<{
217 name: string;
218}> {}
219
220const resourceName = (project: string, location: string, id: string) =>
221 `${locationParent(project, location)}/${COLLECTION}/${id}`;
222
223const toAttrs = (
224 template: aiplatform.GoogleCloudAiplatformV1NotebookRuntimeTemplate,
225 project: string,
226) => {
227 const name = template.name ?? "";
228 const parsed = parseResourceName(name, COLLECTION);
229 return {
230 name,
231 notebookRuntimeTemplateId: parsed.id,
232 project: parsed.project || project,
233 location: parsed.location,
234 displayName: template.displayName,
235 description: template.description,
236 labels: userLabels(template.labels),
237 notebookRuntimeType: template.notebookRuntimeType,
238 machineType: template.machineSpec?.machineType,
239 createTime: template.createTime,
240 updateTime: template.updateTime,
241 etag: template.etag,
242 };
243};
244
245const getByName = (name: string) =>
246 aiplatform
247 .getProjectsLocationsNotebookRuntimeTemplates({ name })
248 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
249
250const waitUntilExists = (name: string) =>
251 getByName(name).pipe(
252 Effect.filterOrFail(
253 (
254 template,
255 ): template is aiplatform.GoogleCloudAiplatformV1NotebookRuntimeTemplate =>
256 template !== undefined,
257 () => new AiPlatformNotResolved({ name }),
258 ),
259 Effect.retry({
260 while: (error) => error._tag === "GCP.AIPlatform.NotResolved",
261 times: 8,
262 schedule: Schedule.spaced("2 seconds"),
263 }),
264 );
265
266const waitUntilGone = (name: string) =>
267 getByName(name).pipe(
268 Effect.filterOrFail(
269 (template) => template === undefined,
270 () => new AiPlatformStillExists({ name }),
271 ),
272 Effect.asVoid,
273 Effect.retry({
274 while: (error) => error._tag === "GCP.AIPlatform.StillExists",
275 times: 10,
276 schedule: Schedule.spaced("2 seconds"),
277 }),
278 );
279
280const desiredMachine = (spec: MachineSpec | undefined): MachineSpec => ({
281 machineType: spec?.machineType ?? DEFAULT_MACHINE_TYPE,
282 acceleratorType: spec?.acceleratorType,
283 acceleratorCount: spec?.acceleratorCount,
284 gpuPartitionSize: spec?.gpuPartitionSize,
285 tpuTopology: spec?.tpuTopology,
286});
287
288export const NotebookRuntimeTemplateProvider = () =>
289 Provider.succeed(NotebookRuntimeTemplate, {
290 stables: [
291 "name",
292 "notebookRuntimeTemplateId",
293 "project",
294 "location",
295 "createTime",
296 ],
297
298 diff: Effect.fn(function* ({ news, olds, output }) {
299 const env = yield* GcpEnvironment.current;
300 if (!isResolved(news)) return undefined;
301 const previousId =
302 olds?.notebookRuntimeTemplateId ?? output?.notebookRuntimeTemplateId;
303 const nextId = news.notebookRuntimeTemplateId ?? previousId;
304 const previousLocation = normalizeLocation(
305 olds?.location ?? output?.location,
306 env.region,
307 );
308 const nextLocation = normalizeLocation(
309 news.location ?? output?.location,
310 env.region,
311 );
312 const previousMachine =
313 olds?.machineSpec?.machineType ?? output?.machineType;
314 const nextMachine = desiredMachine(news.machineSpec).machineType;
315 const typeChanged =
316 (news.notebookRuntimeType ?? olds?.notebookRuntimeType) !==
317 (olds?.notebookRuntimeType ?? output?.notebookRuntimeType);
318 const replace =
319 (previousId !== undefined &&
320 nextId !== undefined &&
321 nextId !== previousId) ||
322 previousLocation !== nextLocation ||
323 (previousMachine !== undefined && previousMachine !== nextMachine) ||
324 (olds?.notebookRuntimeType !== undefined && typeChanged) ||
325 (olds !== undefined &&
326 ((news.description ?? "") !== (olds.description ?? "") ||
327 JSON.stringify(news.labels ?? {}) !==
328 JSON.stringify(olds.labels ?? {})));
329 if (!replace) return undefined;
330 return {
331 action: "replace" as const,
332 deleteFirst:
333 previousLocation === nextLocation &&
334 previousId !== undefined &&
335 nextId === previousId,
336 };
337 }),
338
339 read: Effect.fn(function* ({ id, olds, output }) {
340 const env = yield* GcpEnvironment.current;
341 const templateId = yield* toPhysicalId(
342 id,
343 olds?.notebookRuntimeTemplateId,
344 output?.notebookRuntimeTemplateId,
345 );
346 const location = normalizeLocation(
347 olds?.location ?? output?.location,
348 env.region,
349 );
350 const name =
351 output?.name ?? resourceName(env.project, location, templateId);
352 const existing = yield* getByName(name);
353 if (existing === undefined) return undefined;
354 const attrs = toAttrs(existing, env.project);
355 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
356 ? attrs
357 : Unowned(attrs);
358 }),
359
360 list: () =>
361 Effect.gen(function* () {
362 const env = yield* GcpEnvironment.current;
363 const pages = (yield* Effect.forEach(
364 listLocations(env.region),
365 (location) =>
366 collectPages(
367 aiplatform.listProjectsLocationsNotebookRuntimeTemplates.pages({
368 parent: locationParent(env.project, location),
369 pageSize: 100,
370 }),
371 ).pipe(Effect.catchTag("NotFound", () => Effect.succeed([]))),
372 )).flat();
373 return pages.flatMap((page) =>
374 (page.notebookRuntimeTemplates ?? [])
375 .filter((template) =>
376 Object.keys(template.labels ?? {}).some((key) =>
377 key.startsWith("alchemy-"),
378 ),
379 )
380 .map((template) => toAttrs(template, env.project)),
381 );
382 }),
383
384 reconcile: Effect.fn(function* ({ id, news, output }) {
385 const env = yield* GcpEnvironment.current;
386 const templateId = yield* toPhysicalId(
387 id,
388 news.notebookRuntimeTemplateId,
389 output?.notebookRuntimeTemplateId,
390 );
391 const location = normalizeLocation(
392 news.location ?? output?.location,
393 env.region,
394 );
395 const name = resourceName(env.project, location, templateId);
396 const desiredLabels = {
397 ...toLabels(news.labels),
398 ...(yield* createInternalLabels(id)),
399 };
400 const displayName = news.displayName ?? templateId;
401 const machineSpec = desiredMachine(news.machineSpec);
402
403 let current = yield* getByName(output?.name ?? name);
404
405 if (current === undefined) {
406 const created = yield* aiplatform
407 .createProjectsLocationsNotebookRuntimeTemplates({
408 parent: locationParent(env.project, location),
409 notebookRuntimeTemplateId: templateId,
410 body: {
411 displayName,
412 description: news.description,
413 labels: desiredLabels,
414 notebookRuntimeType: news.notebookRuntimeType,
415 machineSpec,
416 dataPersistentDiskSpec: news.dataPersistentDiskSpec,
417 networkSpec: news.networkSpec,
418 softwareConfig: news.softwareConfig,
419 encryptionSpec: news.encryptionSpec,
420 networkTags: news.networkTags,
421 idleShutdownConfig: news.idleShutdownConfig,
422 shieldedVmConfig: news.shieldedVmConfig,
423 reservationAffinity: news.reservationAffinity,
424 },
425 })
426 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
427 if (created !== undefined) {
428 yield* waitForOperation(created, { alreadyExistsOk: true });
429 }
430 const createdName =
431 resourceNameFromOperation(created ?? {}) ?? output?.name ?? name;
432 current = yield* waitUntilExists(createdName);
433 }
434
435 if (current === undefined) {
436 return yield* new NotebookRuntimeTemplateNotResolved({ name });
437 }
438
439 const observedName = current.name ?? name;
440 // Only display name and software config are updatable; label and
441 // description changes replace the template (see `diff`).
442 const displayChanged = (current.displayName ?? "") !== displayName;
443 const softwareChanged =
444 news.softwareConfig !== undefined &&
445 !jsonEqual(current.softwareConfig, news.softwareConfig);
446
447 if (displayChanged || softwareChanged) {
448 const updateMask = [
449 displayChanged ? "display_name" : undefined,
450 softwareChanged ? "software_config" : undefined,
451 ].filter((field): field is string => field !== undefined);
452 const patched =
453 yield* aiplatform.patchProjectsLocationsNotebookRuntimeTemplates({
454 name: observedName,
455 updateMask: updateMask.join(","),
456 body: {
457 name: observedName,
458 displayName,
459 softwareConfig: news.softwareConfig,
460 etag: current.etag,
461 },
462 });
463 current = patched ?? current;
464 }
465
466 return toAttrs(current, env.project);
467 }),
468
469 delete: Effect.fn(function* ({ output }) {
470 const operation = yield* aiplatform
471 .deleteProjectsLocationsNotebookRuntimeTemplates({
472 name: output.name,
473 })
474 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
475 if (operation !== undefined) {
476 yield* waitForOperation(operation, { notFoundOk: true });
477 }
478 yield* waitUntilGone(output.name);
479 }),
480 });