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 DiskSpec,
26 type EncryptionSpec,
27 type MachineSpec,
28} from "./shared.ts";
29
30const COLLECTION = "persistentResources";
31const DEFAULT_MACHINE_TYPE = "n1-standard-4";
32const DEFAULT_REPLICA_COUNT = "1";
33
34export type ResourcePoolAutoscalingSpec = {
35 /** Minimum replicas (must be > 0 for Persistent Resource). */
36 minReplicaCount?: string;
37 /** Maximum replicas. */
38 maxReplicaCount?: string;
39};
40
41export type ResourcePool = {
42 /**
43 * Unique id within the PersistentResource. Generated if omitted.
44 */
45 id?: string;
46 /** Machine spec. Immutable. */
47 machineSpec?: MachineSpec;
48 /** Replica count (string). @default "1" */
49 replicaCount?: string;
50 /** Disk spec. */
51 diskSpec?: DiskSpec;
52 /** Autoscaling spec. */
53 autoscalingSpec?: ResourcePoolAutoscalingSpec;
54};
55
56export type RaySpec = {
57 /** Default Ray image URI. */
58 imageUri?: string;
59 /** Per-pool Ray images keyed by resource pool id. */
60 resourcePoolImages?: Record<string, string>;
61 /** Resource pool that serves as the Ray head node. */
62 headNodeResourcePoolId?: string;
63 /** Disable Ray OSS log export. */
64 rayLogsDisabled?: boolean;
65 /** Disable Ray metrics. */
66 rayMetricsDisabled?: boolean;
67};
68
69export type ServiceAccountSpec = {
70 /** Enforce a custom service account for workloads. */
71 enableCustomServiceAccount?: boolean;
72 /** Service account email. */
73 serviceAccount?: string;
74};
75
76export type PersistentResourceProps = {
77 /**
78 * Persistent resource id. If omitted, a unique RFC1035 name is
79 * generated. Must match `/^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$/`.
80 * Immutable — changing it replaces the resource.
81 */
82 persistentResourceId?: string;
83 /**
84 * Vertex AI location. Immutable — changing it replaces the resource.
85 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
86 */
87 location?: string;
88 /**
89 * Display name (max 128 UTF-8 characters). Defaults to the resource id.
90 * Changing it replaces the resource.
91 */
92 displayName?: string;
93 /**
94 * User labels. Alchemy ownership labels are merged in automatically.
95 * Changing them replaces the resource.
96 */
97 labels?: Record<string, string>;
98 /**
99 * Resource pools. At least one pool is required. Pool machine specs
100 * are immutable — changing them replaces the resource. Replica counts
101 * update in place only for Ray clusters (`raySpec`); GCP rejects updates
102 * to other persistent resources, so a replica change replaces them.
103 */
104 resourcePools?: ResourcePool[];
105 /**
106 * VPC network to peer (`projects/{project}/global/networks/{network}`).
107 * Immutable.
108 */
109 network?: string;
110 /**
111 * Reserved IP ranges under the VPC.
112 */
113 reservedIpRanges?: string[];
114 /**
115 * Customer-managed encryption key. Immutable.
116 */
117 encryptionSpec?: EncryptionSpec;
118 /**
119 * Ray cluster configuration.
120 */
121 raySpec?: RaySpec;
122 /**
123 * Workload identity / custom service account.
124 */
125 serviceAccountSpec?: ServiceAccountSpec;
126};
127
128export type PersistentResource = Resource<
129 "GCP.AIPlatform.PersistentResource",
130 PersistentResourceProps,
131 {
132 /** Full resource name. */
133 name: string;
134 /** Persistent resource id (last path segment). */
135 persistentResourceId: string;
136 /** Project id. */
137 project: string;
138 /** Location id. */
139 location: string;
140 /** Display name. */
141 displayName: string | undefined;
142 /** User labels (Alchemy ownership labels stripped). */
143 labels: Record<string, string>;
144 /** Server-reported state (`PROVISIONING`, `RUNNING`, …). */
145 state: string | undefined;
146 /** VPC network. */
147 network: string | undefined;
148 /** Resource pool ids. */
149 resourcePoolIds: string[];
150 /** RFC3339 creation timestamp. */
151 createTime: string | undefined;
152 /** RFC3339 last-update timestamp. */
153 updateTime: string | undefined;
154 /** Time the resource first entered `RUNNING`. */
155 startTime: string | undefined;
156 },
157 never,
158 Providers
159>;
160
161/**
162 * A Vertex AI Persistent Resource — dedicated node pools for custom
163 * training and Ray-on-Vertex workloads.
164 *
165 * Changing `persistentResourceId`, `location`, `network`,
166 * `encryptionSpec`, or pool machine specs replaces the resource. Replica
167 * counts, labels, and display name update in place.
168 *
169 * Provisioning typically takes several minutes.
170 *
171 * ### Creating a Persistent Resource
172 * **Example:** Single n1-standard-4 pool
173 * ```typescript
174 * const pool = yield* GCP.AIPlatform.PersistentResource("Train", {
175 * resourcePools: [
176 * {
177 * id: "worker",
178 * replicaCount: "1",
179 * machineSpec: { machineType: "n1-standard-4" },
180 * },
181 * ],
182 * });
183 * ```
184 *
185 * @resource
186 * @category AIPlatform
187 */
188export const PersistentResource = Resource<PersistentResource>(
189 "GCP.AIPlatform.PersistentResource",
190);
191
192export class PersistentResourceNotResolved extends Data.TaggedError(
193 "GCP.AIPlatform.PersistentResourceNotResolved",
194)<{
195 name: string;
196}> {}
197
198const resourceName = (project: string, location: string, id: string) =>
199 `${locationParent(project, location)}/${COLLECTION}/${id}`;
200
201const toPools = (
202 pools: ResourcePool[] | undefined,
203): aiplatform.GoogleCloudAiplatformV1ResourcePool[] =>
204 (pools ?? [{}]).map((pool, index) => ({
205 id: pool.id ?? `pool-${index}`,
206 replicaCount: pool.replicaCount ?? DEFAULT_REPLICA_COUNT,
207 machineSpec: {
208 machineType: pool.machineSpec?.machineType ?? DEFAULT_MACHINE_TYPE,
209 acceleratorType: pool.machineSpec?.acceleratorType,
210 acceleratorCount: pool.machineSpec?.acceleratorCount,
211 gpuPartitionSize: pool.machineSpec?.gpuPartitionSize,
212 tpuTopology: pool.machineSpec?.tpuTopology,
213 },
214 diskSpec: pool.diskSpec,
215 autoscalingSpec: pool.autoscalingSpec,
216 }));
217
218const machineKey = (pools: ResourcePool[] | undefined) =>
219 JSON.stringify(
220 toPools(pools).map((pool) => ({
221 id: pool.id,
222 machineSpec: pool.machineSpec,
223 })),
224 );
225
226const toAttrs = (
227 resource: aiplatform.GoogleCloudAiplatformV1PersistentResource,
228 project: string,
229) => {
230 const name = resource.name ?? "";
231 const parsed = parseResourceName(name, COLLECTION);
232 return {
233 name,
234 persistentResourceId: parsed.id,
235 project: parsed.project || project,
236 location: parsed.location,
237 displayName: resource.displayName,
238 labels: userLabels(resource.labels),
239 state: resource.state,
240 network: resource.network,
241 resourcePoolIds: (resource.resourcePools ?? []).map(
242 (pool) => pool.id ?? "",
243 ),
244 createTime: resource.createTime,
245 updateTime: resource.updateTime,
246 startTime: resource.startTime,
247 };
248};
249
250const getByName = (name: string) =>
251 aiplatform
252 .getProjectsLocationsPersistentResources({ name })
253 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
254
255const waitUntilExists = (name: string) =>
256 getByName(name).pipe(
257 Effect.filterOrFail(
258 (
259 resource,
260 ): resource is aiplatform.GoogleCloudAiplatformV1PersistentResource =>
261 resource !== undefined,
262 () => new AiPlatformNotResolved({ name }),
263 ),
264 Effect.retry({
265 while: (error) => error._tag === "GCP.AIPlatform.NotResolved",
266 times: 8,
267 schedule: Schedule.spaced("2 seconds"),
268 }),
269 );
270
271const waitUntilGone = (name: string) =>
272 getByName(name).pipe(
273 Effect.filterOrFail(
274 (resource) => resource === undefined,
275 () => new AiPlatformStillExists({ name }),
276 ),
277 Effect.asVoid,
278 Effect.retry({
279 while: (error) => error._tag === "GCP.AIPlatform.StillExists",
280 // Tearing down the cluster takes several minutes.
281 times: 75,
282 schedule: Schedule.spaced("8 seconds"),
283 }),
284 );
285
286export const PersistentResourceProvider = () =>
287 Provider.succeed(PersistentResource, {
288 stables: [
289 "name",
290 "persistentResourceId",
291 "project",
292 "location",
293 "createTime",
294 ],
295
296 diff: Effect.fn(function* ({ news, olds, output }) {
297 const env = yield* GcpEnvironment.current;
298 if (!isResolved(news)) return undefined;
299 const previousId =
300 olds?.persistentResourceId ?? output?.persistentResourceId;
301 const nextId = news.persistentResourceId ?? previousId;
302 const previousLocation = normalizeLocation(
303 olds?.location ?? output?.location,
304 env.region,
305 );
306 const nextLocation = normalizeLocation(
307 news.location ?? output?.location,
308 env.region,
309 );
310 const networkChanged =
311 (news.network ?? olds?.network ?? "") !== (olds?.network ?? "");
312 const machineChanged =
313 olds !== undefined &&
314 machineKey(news.resourcePools) !== machineKey(olds.resourcePools);
315 // PATCH only accepts `resource_pools.replica_count`; `display_name` and
316 // `labels` are rejected ("Unrecognized path"), so changing them
317 // replaces the resource.
318 const displayNameChanged =
319 olds !== undefined &&
320 (news.displayName ?? "") !== (olds.displayName ?? "");
321 const labelsChanged =
322 olds !== undefined &&
323 JSON.stringify(Object.entries(news.labels ?? {}).sort()) !==
324 JSON.stringify(Object.entries(olds.labels ?? {}).sort());
325 // GCP only updates Ray-cluster persistent resources ("Currently we only
326 // support the update function on Ray cluster"); any other resource is
327 // replaced when its replica counts change.
328 const replicasChanged =
329 olds !== undefined &&
330 news.raySpec === undefined &&
331 JSON.stringify(
332 (news.resourcePools ?? []).map((p) => p.replicaCount),
333 ) !==
334 JSON.stringify((olds.resourcePools ?? []).map((p) => p.replicaCount));
335 const replace =
336 replicasChanged ||
337 (previousId !== undefined &&
338 nextId !== undefined &&
339 nextId !== previousId) ||
340 previousLocation !== nextLocation ||
341 (olds !== undefined && networkChanged) ||
342 machineChanged ||
343 displayNameChanged ||
344 labelsChanged;
345 if (!replace) return undefined;
346 return {
347 action: "replace" as const,
348 deleteFirst:
349 previousLocation === nextLocation &&
350 previousId !== undefined &&
351 nextId === previousId,
352 };
353 }),
354
355 read: Effect.fn(function* ({ id, olds, output }) {
356 const env = yield* GcpEnvironment.current;
357 const resourceId = yield* toPhysicalId(
358 id,
359 olds?.persistentResourceId,
360 output?.persistentResourceId,
361 );
362 const location = normalizeLocation(
363 olds?.location ?? output?.location,
364 env.region,
365 );
366 const name =
367 output?.name ?? resourceName(env.project, location, resourceId);
368 const existing = yield* getByName(name);
369 if (existing === undefined) return undefined;
370 const attrs = toAttrs(existing, env.project);
371 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
372 ? attrs
373 : Unowned(attrs);
374 }),
375
376 list: () =>
377 Effect.gen(function* () {
378 const env = yield* GcpEnvironment.current;
379 const pages = (yield* Effect.forEach(
380 listLocations(env.region),
381 (location) =>
382 collectPages(
383 aiplatform.listProjectsLocationsPersistentResources.pages({
384 parent: locationParent(env.project, location),
385 pageSize: 100,
386 }),
387 ).pipe(Effect.catchTag("NotFound", () => Effect.succeed([]))),
388 )).flat();
389 return pages.flatMap((page) =>
390 (page.persistentResources ?? [])
391 .filter((resource) =>
392 Object.keys(resource.labels ?? {}).some((key) =>
393 key.startsWith("alchemy-"),
394 ),
395 )
396 .map((resource) => toAttrs(resource, env.project)),
397 );
398 }),
399
400 reconcile: Effect.fn(function* ({ id, news, output }) {
401 const env = yield* GcpEnvironment.current;
402 const resourceId = yield* toPhysicalId(
403 id,
404 news.persistentResourceId,
405 output?.persistentResourceId,
406 );
407 const location = normalizeLocation(
408 news.location ?? output?.location,
409 env.region,
410 );
411 const name = resourceName(env.project, location, resourceId);
412 const desiredLabels = {
413 ...toLabels(news.labels),
414 ...(yield* createInternalLabels(id)),
415 };
416 const displayName = news.displayName ?? resourceId;
417 const resourcePools = toPools(news.resourcePools);
418 const resourceRuntimeSpec =
419 news.raySpec || news.serviceAccountSpec
420 ? {
421 raySpec: news.raySpec
422 ? {
423 imageUri: news.raySpec.imageUri,
424 resourcePoolImages: news.raySpec.resourcePoolImages,
425 headNodeResourcePoolId: news.raySpec.headNodeResourcePoolId,
426 rayLogsSpec: news.raySpec.rayLogsDisabled
427 ? { disabled: true }
428 : undefined,
429 rayMetricSpec: news.raySpec.rayMetricsDisabled
430 ? { disabled: true }
431 : undefined,
432 }
433 : undefined,
434 serviceAccountSpec: news.serviceAccountSpec,
435 }
436 : undefined;
437
438 let current = yield* getByName(output?.name ?? name);
439
440 if (current === undefined) {
441 const created = yield* aiplatform
442 .createProjectsLocationsPersistentResources({
443 parent: locationParent(env.project, location),
444 persistentResourceId: resourceId,
445 body: {
446 displayName,
447 labels: desiredLabels,
448 resourcePools,
449 network: news.network,
450 reservedIpRanges: news.reservedIpRanges,
451 encryptionSpec: news.encryptionSpec,
452 resourceRuntimeSpec,
453 },
454 })
455 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
456 if (created !== undefined) {
457 yield* waitForOperation(created, { alreadyExistsOk: true });
458 }
459 const createdName =
460 resourceNameFromOperation(created ?? {}) ?? output?.name ?? name;
461 current = yield* waitUntilExists(createdName);
462 }
463
464 if (current === undefined) {
465 return yield* new PersistentResourceNotResolved({ name });
466 }
467
468 const observedName = current.name ?? name;
469 // Reads fill in server defaults (disk spec, used replicas), and a pool's
470 // machine spec is immutable (a change replaces), so only replica counts
471 // are compared and patched.
472 const replicaCounts = (
473 pools:
474 | ReadonlyArray<{ id?: string; replicaCount?: string }>
475 | undefined,
476 ) =>
477 Object.fromEntries(
478 (pools ?? []).map((pool, index) => [
479 pool.id ?? String(index),
480 String(pool.replicaCount ?? ""),
481 ]),
482 );
483 const poolsChanged = !jsonEqual(
484 replicaCounts(current.resourcePools),
485 replicaCounts(resourcePools),
486 );
487
488 // Display name and labels are create-only; `diff` replaces the resource
489 // when they change.
490 // Only Ray clusters accept updates; `diff` replaces other resources.
491 if (poolsChanged && news.raySpec !== undefined) {
492 const patched =
493 yield* aiplatform.patchProjectsLocationsPersistentResources({
494 name: observedName,
495 updateMask: "resource_pools.replica_count",
496 body: { name: observedName, resourcePools },
497 });
498 yield* waitForOperation(patched);
499 current = yield* getByName(observedName);
500 }
501
502 if (current === undefined) {
503 return yield* new PersistentResourceNotResolved({ name: observedName });
504 }
505 return toAttrs(current, env.project);
506 }),
507
508 delete: Effect.fn(function* ({ output }) {
509 // A PROVISIONING resource rejects deletes ("is being created thus can
510 // not be deleted now"), so wait for provisioning to settle first.
511 yield* getByName(output.name).pipe(
512 Effect.repeat({
513 schedule: Schedule.spaced("10 seconds"),
514 until: (resource) => resource?.state !== "PROVISIONING",
515 times: 60,
516 }),
517 );
518 const operation = yield* aiplatform
519 .deleteProjectsLocationsPersistentResources({ name: output.name })
520 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
521 if (operation !== undefined) {
522 yield* waitForOperation(operation, { notFoundOk: true });
523 }
524 yield* waitUntilGone(output.name);
525 }),
526 });