1import * as tpu from "@distilled.cloud/gcp/tpu_v2";
2import * as Effect from "effect/Effect";
3import * as Schedule from "effect/Schedule";
4import * as Stream from "effect/Stream";
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 {
14 encodeDescription,
15 hasAlchemyLabelMap,
16 hasOwnershipMarker,
17 normalizeLocation,
18 parseDescription,
19 parseName,
20 replaceOnIdentity,
21 ResourceNotResolved,
22 toPhysicalId,
23 waitForOperation,
24 waitUntilExists,
25 waitUntilGone,
26} from "./internal.ts";
27import { toNodeBody, type NodeProps } from "./Node.ts";
28
29export type MultisliceParams = {
30 /**
31 * Prefix for generated node ids (`np` → `np-0`, `np-1`, …). Defaults to
32 * the queued resource id.
33 */
34 nodeIdPrefix?: string;
35 /**
36 * Number of nodes to provision. Must be greater than 1.
37 */
38 nodeCount?: number;
39};
40
41export type NodeSpec = {
42 /**
43 * Parent location
44 * (`projects/{project}/locations/{location}`). Defaults to the queued
45 * resource parent.
46 */
47 parent?: string;
48 /**
49 * Unqualified node id when requesting a single node. Mutually exclusive
50 * with `multisliceParams`.
51 */
52 nodeId?: string;
53 /**
54 * Multislice request. Mutually exclusive with `nodeId`.
55 */
56 multisliceParams?: MultisliceParams;
57 /**
58 * Node configuration. Alchemy ownership labels are merged into
59 * `node.labels` (QueuedResource has no labels field of its own).
60 */
61 node?: NodeProps;
62};
63
64export type QueueingPolicy = {
65 /** Absolute time after which the request fails if still unfulfilled. */
66 validUntilTime?: string;
67 /** Relative duration after which resources may be created. */
68 validAfterDuration?: string;
69 /** Absolute interval during which resources may be created. */
70 validInterval?: {
71 /** Inclusive start (RFC3339). */
72 startTime?: string;
73 /** Exclusive end (RFC3339). */
74 endTime?: string;
75 };
76 /** Absolute time after which resources may be created. */
77 validAfterTime?: string;
78 /** Relative duration after which the request fails if still unfulfilled. */
79 validUntilDuration?: string;
80};
81
82export type Guaranteed = {
83 /**
84 * Minimum duration the resources must be allocated (`3600s`, …).
85 */
86 minDuration?: string;
87};
88
89export type QueuedResourceProps = {
90 /**
91 * Queued resource id (the `{queuedResource}` segment of
92 * `projects/{project}/locations/{location}/queuedResources/{queuedResource}`).
93 * If omitted, a unique RFC1035 name is generated from the stack, stage,
94 * and logical id. Immutable — changing it replaces the resource.
95 */
96 queuedResourceId?: string;
97 /**
98 * Zone (`us-central1-c`, …). Immutable — changing it replaces the
99 * resource.
100 * @default "us-central1-c"
101 */
102 location?: string;
103 /**
104 * Reservation to provision in
105 * (`projects/{project}/locations/{zone}/reservations/{reservation}`).
106 * Immutable.
107 */
108 reservationName?: string;
109 /**
110 * Request the Spot tier. Immutable.
111 * @default false
112 */
113 spot?: boolean;
114 /**
115 * Guaranteed tier. Immutable.
116 */
117 guaranteed?: Guaranteed;
118 /**
119 * When the request may be fulfilled. Immutable.
120 */
121 queueingPolicy?: QueueingPolicy;
122 /**
123 * TPU node specification. Omit to request a default v2-8 node with the
124 * standard runtime. Alchemy ownership is stamped on each nested node's
125 * labels and description (QueuedResource has no labels API).
126 */
127 nodeSpec?: NodeSpec[];
128 /**
129 * Delete running nodes belonging to this request before deleting the
130 * queued resource.
131 * @default true
132 */
133 forceDestroy?: boolean;
134};
135
136export type QueuedResource = Resource<
137 "GCP.TPU.QueuedResource",
138 QueuedResourceProps,
139 {
140 /** Full resource name `projects/{project}/locations/{location}/queuedResources/{queuedResource}`. */
141 name: string;
142 /** Queued resource id (last path segment). */
143 queuedResourceId: string;
144 /** Project id. */
145 project: string;
146 /** Zone id. */
147 location: string;
148 /** Reservation name, if any. */
149 reservationName: string | undefined;
150 /** Whether the Spot tier is set. */
151 spot: boolean;
152 /** Guaranteed tier configuration. */
153 guaranteed: Guaranteed | undefined;
154 /** Queueing policy. */
155 queueingPolicy: QueueingPolicy | undefined;
156 /** Node specification currently stored on the request. */
157 nodeSpec: NodeSpec[];
158 /** Server-reported state (`ACCEPTED`, `WAITING_FOR_RESOURCES`, …). */
159 state: string | undefined;
160 /** Who initiated the current state (`USER`, `SERVICE`). */
161 stateInitiator: string | undefined;
162 /** Error that caused `FAILED`, if any. */
163 failedError: string | undefined;
164 /** RFC3339 creation timestamp. */
165 createTime: string | undefined;
166 },
167 never,
168 Providers
169>;
170
171/**
172 * A Cloud TPU queued resource request. The request sits in a queue until
173 * capacity is available, then provisions the requested TPU node(s).
174 *
175 * Queued resources have no labels field — Alchemy stamps ownership onto
176 * each nested `nodeSpec.node` labels and description so `list` / nuke can
177 * find them. The resource is create-only: changing identity or the node
178 * spec replaces it.
179 *
180 * ### Creating a Queued Resource
181 * **Example:** Default v2-8 request
182 * ```typescript
183 * const request = yield* GCP.TPU.QueuedResource("Trainer", {});
184 * ```
185 *
186 * **Example:** Spot request with an explicit node spec
187 * ```typescript
188 * const request = yield* GCP.TPU.QueuedResource("Trainer", {
189 * location: "us-central1-c",
190 * spot: true,
191 * nodeSpec: [{
192 * nodeId: "app-tpu",
193 * node: {
194 * acceleratorType: "v2-8",
195 * runtimeVersion: "tpu-ubuntu2204-base",
196 * labels: { env: "prod" },
197 * },
198 * }],
199 * });
200 * ```
201 *
202 * @resource
203 * @category TPU
204 */
205export const QueuedResource = Resource<QueuedResource>(
206 "GCP.TPU.QueuedResource",
207);
208
209const resourceName = (
210 project: string,
211 location: string,
212 queuedResourceId: string,
213) =>
214 `projects/${project}/locations/${location}/queuedResources/${queuedResourceId}`;
215
216const parentOf = (project: string, location: string) =>
217 `projects/${project}/locations/${location}`;
218
219const intervalOf = (
220 interval: tpu.Interval | QueueingPolicy["validInterval"] | undefined,
221) => {
222 if (interval === undefined) return undefined;
223 return { startTime: interval.startTime, endTime: interval.endTime };
224};
225
226const queueingOf = (
227 policy: tpu.QueueingPolicy | QueueingPolicy | undefined,
228): QueueingPolicy | undefined => {
229 if (policy === undefined) return undefined;
230 return {
231 validUntilTime: policy.validUntilTime,
232 validAfterDuration: policy.validAfterDuration,
233 validInterval: intervalOf(policy.validInterval),
234 validAfterTime: policy.validAfterTime,
235 validUntilDuration: policy.validUntilDuration,
236 };
237};
238
239const queueingKey = (policy: QueueingPolicy | undefined) =>
240 JSON.stringify({
241 validUntilTime: policy?.validUntilTime ?? "",
242 validAfterDuration: policy?.validAfterDuration ?? "",
243 validIntervalStart: policy?.validInterval?.startTime ?? "",
244 validIntervalEnd: policy?.validInterval?.endTime ?? "",
245 validAfterTime: policy?.validAfterTime ?? "",
246 validUntilDuration: policy?.validUntilDuration ?? "",
247 });
248
249const guaranteedOf = (
250 guaranteed: tpu.Guaranteed | Guaranteed | undefined,
251): Guaranteed | undefined => {
252 if (guaranteed === undefined) return undefined;
253 return { minDuration: guaranteed.minDuration };
254};
255
256const networkTagsOf = (node: tpu.Node | NodeProps) => {
257 if ("networkTags" in node && node.networkTags) {
258 return [...node.networkTags];
259 }
260 if ("tags" in node && node.tags) {
261 return [...node.tags];
262 }
263 return undefined;
264};
265
266const stringRecordOf = (map: Record<string, string | undefined> | undefined) =>
267 map
268 ? Object.fromEntries(
269 Object.entries(map).filter(
270 (entry): entry is [string, string] => entry[1] !== undefined,
271 ),
272 )
273 : undefined;
274
275const nodeSpecOf = (
276 spec: tpu.NodeSpec | NodeSpec,
277 parent: string,
278): NodeSpec => ({
279 parent: spec.parent ?? parent,
280 nodeId: spec.nodeId,
281 multisliceParams: spec.multisliceParams
282 ? {
283 nodeIdPrefix: spec.multisliceParams.nodeIdPrefix,
284 nodeCount: spec.multisliceParams.nodeCount,
285 }
286 : undefined,
287 node: spec.node
288 ? {
289 description: spec.node.description,
290 labels: stringRecordOf(spec.node.labels),
291 metadata: stringRecordOf(spec.node.metadata),
292 networkTags: networkTagsOf(spec.node),
293 runtimeVersion: spec.node.runtimeVersion,
294 acceleratorType: spec.node.acceleratorType,
295 acceleratorConfig: spec.node.acceleratorConfig,
296 cidrBlock: spec.node.cidrBlock,
297 serviceAccount: spec.node.serviceAccount,
298 networkConfig: spec.node.networkConfig,
299 networkConfigs: spec.node.networkConfigs,
300 bootDiskConfig: spec.node.bootDiskConfig,
301 dataDisks: spec.node.dataDisks,
302 schedulingConfig: spec.node.schedulingConfig,
303 shieldedInstanceConfig: spec.node.shieldedInstanceConfig,
304 }
305 : undefined,
306});
307
308const nodeSpecKey = (spec: NodeSpec) =>
309 JSON.stringify({
310 parent: spec.parent ?? "",
311 nodeId: spec.nodeId ?? "",
312 prefix: spec.multisliceParams?.nodeIdPrefix ?? "",
313 count: spec.multisliceParams?.nodeCount ?? "",
314 runtime: spec.node?.runtimeVersion ?? "",
315 accelerator: spec.node?.acceleratorType ?? "",
316 topology: spec.node?.acceleratorConfig?.topology ?? "",
317 cidr: spec.node?.cidrBlock ?? "",
318 });
319
320const specsKey = (specs: readonly NodeSpec[]) =>
321 JSON.stringify(specs.map(nodeSpecKey));
322
323const nestedNodes = (resource: tpu.QueuedResource) =>
324 resource.tpu?.nodeSpec ?? [];
325
326const nestedOwnershipLabels = (resource: tpu.QueuedResource) => {
327 const labels: Record<string, string> = {};
328 for (const spec of nestedNodes(resource)) {
329 Object.assign(labels, tagRecord(spec.node?.labels));
330 Object.assign(labels, parseDescription(spec.node?.description).labels);
331 }
332 return labels;
333};
334
335const hasNestedOwnership = (resource: tpu.QueuedResource) => {
336 for (const spec of nestedNodes(resource)) {
337 if (hasAlchemyLabelMap(spec.node?.labels)) return true;
338 if (hasOwnershipMarker(spec.node?.description)) return true;
339 }
340 return false;
341};
342
343const desiredSpecs = (
344 news: QueuedResourceProps,
345 parent: string,
346 desiredLabels: Record<string, string>,
347): tpu.NodeSpec[] => {
348 const specs =
349 news.nodeSpec !== undefined && news.nodeSpec.length > 0
350 ? news.nodeSpec
351 : [{ node: {} as NodeProps }];
352 return specs.map((spec) => {
353 const node = spec.node ?? {};
354 const description = encodeDescription(desiredLabels, node.description);
355 const labels = { ...toLabels(node.labels), ...desiredLabels };
356 return {
357 parent: spec.parent ?? parent,
358 nodeId: spec.nodeId,
359 multisliceParams: spec.multisliceParams,
360 node: toNodeBody({ ...node, labels }, labels, { description }),
361 };
362 });
363};
364
365const toAttrs = (resource: tpu.QueuedResource, project: string) => {
366 const name = resource.name ?? "";
367 const parsed = parseName(name, "queuedResources");
368 const parent = parentOf(parsed.project || project, parsed.location);
369 return {
370 name,
371 queuedResourceId: parsed.id,
372 project: parsed.project || project,
373 location: parsed.location,
374 reservationName: resource.reservationName,
375 spot: resource.spot !== undefined,
376 guaranteed: guaranteedOf(resource.guaranteed),
377 queueingPolicy: queueingOf(resource.queueingPolicy),
378 nodeSpec: (resource.tpu?.nodeSpec ?? []).map((spec) => {
379 const mapped = nodeSpecOf(spec, parent);
380 if (mapped.node?.description !== undefined) {
381 mapped.node = {
382 ...mapped.node,
383 description: parseDescription(mapped.node.description).description,
384 labels: mapped.node.labels
385 ? Object.fromEntries(
386 Object.entries(mapped.node.labels).filter(
387 ([key]) => !key.startsWith("alchemy-"),
388 ),
389 )
390 : undefined,
391 };
392 }
393 return mapped;
394 }),
395 state: resource.state?.state,
396 stateInitiator: resource.state?.stateInitiator,
397 failedError: resource.state?.failedData?.error?.message,
398 createTime: resource.createTime,
399 };
400};
401
402const isPlaceholder = (resource: tpu.QueuedResource) => {
403 const name = resource.name ?? "";
404 return (
405 name.length === 0 ||
406 name.endsWith("/queuedResources/-") ||
407 name.endsWith("/queuedResources/")
408 );
409};
410
411const getByName = (name: string) =>
412 tpu
413 .getProjectsLocationsQueuedResources({ name })
414 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
415
416const specIdentityChanged = (
417 previous: QueuedResourceProps | { nodeSpec?: NodeSpec[] } | undefined,
418 next: QueuedResourceProps,
419 parent: string,
420) => {
421 const previousSpecs = (previous?.nodeSpec ?? []).map((spec) =>
422 nodeSpecOf(spec, parent),
423 );
424 const nextSpecs =
425 next.nodeSpec !== undefined && next.nodeSpec.length > 0
426 ? next.nodeSpec.map((spec) => nodeSpecOf(spec, parent))
427 : previousSpecs;
428 if (previousSpecs.length === 0) return false;
429 return specsKey(previousSpecs) !== specsKey(nextSpecs);
430};
431
432export const QueuedResourceProvider = () =>
433 Provider.succeed(QueuedResource, {
434 stables: ["name", "queuedResourceId", "project", "location", "createTime"],
435
436 diff: Effect.fn(function* ({ news, olds, output }) {
437 if (!isResolved(news)) return undefined;
438
439 const previousId = olds?.queuedResourceId ?? output?.queuedResourceId;
440 const nextId = news.queuedResourceId ?? previousId;
441 const previousLocation = normalizeLocation(
442 olds?.location ?? output?.location,
443 );
444 const nextLocation = normalizeLocation(news.location ?? output?.location);
445 const parent = `projects/_/locations/${nextLocation}`;
446 const previousReservation =
447 olds?.reservationName ?? output?.reservationName ?? "";
448 const nextReservation = news.reservationName ?? previousReservation;
449 const previousSpot = olds?.spot ?? output?.spot ?? false;
450 const nextSpot = news.spot ?? previousSpot;
451 const previousGuaranteed =
452 guaranteedOf(olds?.guaranteed ?? output?.guaranteed)?.minDuration ?? "";
453 const nextGuaranteed =
454 guaranteedOf(news.guaranteed ?? olds?.guaranteed ?? output?.guaranteed)
455 ?.minDuration ?? previousGuaranteed;
456 const previousQueueing = queueingKey(
457 queueingOf(olds?.queueingPolicy ?? output?.queueingPolicy),
458 );
459 const nextQueueing = queueingKey(
460 queueingOf(
461 news.queueingPolicy ?? olds?.queueingPolicy ?? output?.queueingPolicy,
462 ),
463 );
464
465 return replaceOnIdentity({
466 previousId,
467 nextId,
468 previousLocation,
469 nextLocation,
470 extra:
471 previousReservation !== nextReservation ||
472 previousSpot !== nextSpot ||
473 previousGuaranteed !== nextGuaranteed ||
474 previousQueueing !== nextQueueing ||
475 specIdentityChanged(olds ?? output, news, parent),
476 });
477 }),
478
479 read: Effect.fn(function* ({ id, olds, output }) {
480 const env = yield* GcpEnvironment.current;
481 const queuedResourceId = yield* toPhysicalId(
482 id,
483 olds?.queuedResourceId,
484 output?.queuedResourceId,
485 );
486 const location = normalizeLocation(olds?.location ?? output?.location);
487 const name =
488 output?.name ?? resourceName(env.project, location, queuedResourceId);
489 const existing = yield* getByName(name);
490 if (existing === undefined) return undefined;
491 const attrs = toAttrs(existing, env.project);
492 const nested = nestedOwnershipLabels(existing);
493 return (yield* hasAlchemyLabels(id, nested)) ? attrs : Unowned(attrs);
494 }),
495
496 list: () =>
497 Effect.gen(function* () {
498 const env = yield* GcpEnvironment.current;
499 return yield* tpu.listProjectsLocationsQueuedResources
500 .pages({
501 parent: `projects/${env.project}/locations/-`,
502 pageSize: 1000,
503 })
504 .pipe(
505 Stream.flatMap((page) =>
506 Stream.fromIterable(page.queuedResources ?? []),
507 ),
508 Stream.filter(
509 (resource) =>
510 !isPlaceholder(resource) && hasNestedOwnership(resource),
511 ),
512 Stream.map((resource) => toAttrs(resource, env.project)),
513 Stream.runCollect,
514 Effect.map((chunk) => Array.from(chunk)),
515 Effect.catchTag("NotFound", () => Effect.succeed([])),
516 );
517 }),
518
519 reconcile: Effect.fn(function* ({ id, news, output }) {
520 const env = yield* GcpEnvironment.current;
521 const queuedResourceId = yield* toPhysicalId(
522 id,
523 news.queuedResourceId,
524 output?.queuedResourceId,
525 );
526 const location = normalizeLocation(news.location ?? output?.location);
527 const parent = parentOf(env.project, location);
528 const name = resourceName(env.project, location, queuedResourceId);
529 const desiredLabels = yield* createInternalLabels(id);
530
531 let current = yield* getByName(output?.name ?? name);
532
533 if (current === undefined) {
534 const created = yield* tpu
535 .createProjectsLocationsQueuedResources({
536 parent,
537 queuedResourceId,
538 body: {
539 reservationName: news.reservationName,
540 spot: news.spot === true ? {} : undefined,
541 guaranteed: news.guaranteed,
542 queueingPolicy: news.queueingPolicy,
543 tpu: { nodeSpec: desiredSpecs(news, parent, desiredLabels) },
544 },
545 })
546 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
547 if (created !== undefined) {
548 yield* waitForOperation(created);
549 }
550 current = yield* waitUntilExists(getByName(name), name);
551 }
552
553 if (current === undefined) {
554 return yield* new ResourceNotResolved({ name });
555 }
556
557 return toAttrs(current, env.project);
558 }),
559
560 delete: Effect.fn(function* ({ output, olds }) {
561 const force = olds?.forceDestroy !== false;
562 const operation = yield* tpu
563 .deleteProjectsLocationsQueuedResources({
564 name: output.name,
565 force,
566 })
567 .pipe(
568 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
569 Effect.retry({
570 while: (error) => error._tag === "Conflict",
571 times: 8,
572 schedule: Schedule.spaced("5 seconds"),
573 }),
574 );
575 if (operation !== undefined) {
576 yield* waitForOperation(operation, { notFoundOk: true });
577 }
578 yield* waitUntilGone(getByName(output.name), output.name);
579 }),
580 });