1import * as tpu from "@distilled.cloud/gcp/tpu_v2";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { waitForOperation as waitForGcpOperation } from "../Operation.ts";
7import { createPhysicalName } from "../../PhysicalName.ts";
8import { tagRecord } from "../../Tags.ts";
9import { alchemyLabelKeys, stripInternalLabels } from "../Labels.ts";
10
11// A zone, not a region: TPUs (and the v2-8 default) exist only in a few zones.
12export const DEFAULT_LOCATION = "us-central1-c";
13export const DEFAULT_ACCELERATOR = "v2-8";
14export const DEFAULT_RUNTIME = "tpu-ubuntu2204-base";
15export const DEFAULT_NETWORK = "default";
16export const MAX_NAME_LENGTH = 63;
17
18export class ResourceNotResolved extends Data.TaggedError(
19 "GCP.TPU.ResourceNotResolved",
20)<{
21 name: string;
22}> {}
23
24export class ResourceStillExists extends Data.TaggedError(
25 "GCP.TPU.ResourceStillExists",
26)<{
27 name: string;
28}> {}
29
30export class ResourceNotReady extends Data.TaggedError(
31 "GCP.TPU.ResourceNotReady",
32)<{
33 name: string;
34 state: string;
35}> {}
36
37export class ResourceFailed extends Data.TaggedError("GCP.TPU.ResourceFailed")<{
38 name: string;
39 state: string;
40 message: string | undefined;
41}> {}
42
43export const lastSegment = (value: string | undefined) => {
44 if (value === undefined || value.length === 0) return "";
45 const trimmed = value.replace(/\/+$/, "");
46 const parts = trimmed.split("/");
47 return parts[parts.length - 1] || trimmed;
48};
49
50export const rfc1035 = (name: string, fallback = "tpu"): string => {
51 let next = name
52 .toLowerCase()
53 .replace(/[^a-z0-9-]/g, "-")
54 .replace(/-+/g, "-")
55 .replace(/^-+|-+$/g, "");
56 if (!/^[a-z]/.test(next)) next = `t${next}`;
57 next = next.slice(0, MAX_NAME_LENGTH).replace(/-+$/g, "");
58 if (next.length === 0) return fallback;
59 if (!/[a-z0-9]$/.test(next)) next = `${next.slice(0, MAX_NAME_LENGTH - 1)}0`;
60 return next.slice(0, MAX_NAME_LENGTH);
61};
62
63export const normalizeLocation = (
64 location: string | undefined,
65 fallback = DEFAULT_LOCATION,
66) => lastSegment(location ?? fallback).toLowerCase();
67
68export const toPhysicalId = (
69 id: string,
70 explicit: string | undefined,
71 existing: string | undefined,
72 fallback = "tpu",
73) =>
74 Effect.gen(function* () {
75 if (explicit !== undefined) return rfc1035(explicit, fallback);
76 if (existing !== undefined) return existing;
77 return rfc1035(
78 yield* createPhysicalName({
79 id,
80 maxLength: MAX_NAME_LENGTH,
81 lowercase: true,
82 }),
83 fallback,
84 );
85 });
86
87export const parseName = (name: string, collection: string) => {
88 const parts = name.split("/").filter((part) => part.length > 0);
89 const collectionAt = parts.lastIndexOf(collection);
90 const locationsAt = parts.lastIndexOf("locations");
91 const projectsAt = parts.lastIndexOf("projects");
92 return {
93 project:
94 projectsAt >= 0 && parts[projectsAt + 1] ? parts[projectsAt + 1]! : "",
95 location:
96 locationsAt >= 0 && parts[locationsAt + 1]
97 ? parts[locationsAt + 1]!
98 : DEFAULT_LOCATION,
99 id:
100 collectionAt >= 0 && parts[collectionAt + 1]
101 ? parts[collectionAt + 1]!
102 : lastSegment(name),
103 };
104};
105
106export const userLabels = (
107 labels: Record<string, string | undefined> | null | undefined,
108): Record<string, string> => stripInternalLabels(tagRecord(labels));
109
110export const hasAlchemyLabelMap = (
111 labels: Record<string, string | undefined> | null | undefined,
112) => Object.keys(labels ?? {}).some((key) => key.startsWith("alchemy-"));
113
114export const stringMapOf = (
115 map: Record<string, string | undefined> | null | undefined,
116): Record<string, string> =>
117 Object.fromEntries(
118 Object.entries(map ?? {}).filter(
119 (entry): entry is [string, string] => entry[1] !== undefined,
120 ),
121 );
122
123export const mapKey = (
124 map: Record<string, string | undefined> | null | undefined,
125) =>
126 JSON.stringify(
127 Object.fromEntries(
128 Object.entries(stringMapOf(map)).sort(([left], [right]) =>
129 left.localeCompare(right),
130 ),
131 ),
132 );
133
134export const stringsOf = (
135 values: ReadonlyArray<string | undefined> | null | undefined,
136): string[] =>
137 (values ?? []).filter((value): value is string => value !== undefined);
138
139export const stringsKey = (
140 values: ReadonlyArray<string | undefined> | null | undefined,
141) => JSON.stringify([...stringsOf(values)].sort());
142
143export const encodeDescription = (
144 labels: Record<string, string>,
145 description: string | undefined,
146): string => {
147 const marker = `[alchemy ${alchemyLabelKeys.stack}=${labels[alchemyLabelKeys.stack]} ${alchemyLabelKeys.stage}=${labels[alchemyLabelKeys.stage]} ${alchemyLabelKeys.id}=${labels[alchemyLabelKeys.id]}]`;
148 return description ? `${marker}\n${description}` : marker;
149};
150
151export const parseDescription = (
152 description: string | undefined,
153): {
154 labels: Record<string, string>;
155 description: string | undefined;
156} => {
157 if (!description?.startsWith("[alchemy ")) {
158 return { labels: {}, description };
159 }
160 const end = description.indexOf("]");
161 if (end < 0) return { labels: {}, description };
162 const labels: Record<string, string> = {};
163 for (const part of description.slice("[alchemy ".length, end).split(/\s+/)) {
164 const eq = part.indexOf("=");
165 if (eq > 0) {
166 labels[part.slice(0, eq)] = part.slice(eq + 1);
167 }
168 }
169 const rest = description.slice(end + 1).replace(/^\n/, "");
170 return { labels, description: rest.length > 0 ? rest : undefined };
171};
172
173export const hasOwnershipMarker = (description: string | undefined) =>
174 Object.keys(parseDescription(description).labels).some((key) =>
175 key.startsWith("alchemy-"),
176 );
177
178export const fieldMask = (fields: Array<string | false | undefined>) =>
179 fields
180 .filter((field): field is string => typeof field === "string")
181 .join(",");
182
183export const replaceOnIdentity = (input: {
184 previousId: string | undefined;
185 nextId: string | undefined;
186 previousLocation: string;
187 nextLocation: string;
188 extra?: boolean;
189}) => {
190 const replace =
191 (input.extra ?? false) ||
192 (input.previousId !== undefined &&
193 input.nextId !== undefined &&
194 input.nextId !== input.previousId) ||
195 input.previousLocation !== input.nextLocation;
196 if (!replace) return undefined;
197 const samePhysical =
198 input.previousLocation === input.nextLocation &&
199 input.previousId !== undefined &&
200 input.nextId === input.previousId;
201 return {
202 action: "replace" as const,
203 deleteFirst: samePhysical,
204 };
205};
206
207/**
208 * Wait for a Cloud TPU operation. TPU VM creates take several minutes
209 * (large slices longer). An operation that finished with `ALREADY_EXISTS`
210 * (6) is a lost create race; `notFoundOk` also accepts `NOT_FOUND` (5).
211 */
212export const waitForOperation = (
213 operation: tpu.Operation,
214 options?: { notFoundOk?: boolean },
215) =>
216 waitForGcpOperation(
217 operation,
218 (name) =>
219 tpu
220 .getProjectsLocationsOperations({ name })
221 .pipe(
222 Effect.catchTag("NotFound", (error) =>
223 options?.notFoundOk === true
224 ? Effect.succeed<tpu.Operation>({ name, done: true })
225 : Effect.fail(error),
226 ),
227 ),
228 { budget: "30 minutes" },
229 ).pipe(
230 Effect.catchIf(
231 (error) =>
232 error._tag === "GCP.OperationFailed" &&
233 (error.code === 6 ||
234 (options?.notFoundOk === true && error.code === 5)),
235 () => Effect.void,
236 ),
237 Effect.asVoid,
238 );
239
240export const waitUntilExists = <A, E extends { readonly _tag: string }, R>(
241 get: Effect.Effect<A | undefined, E, R>,
242 name: string,
243): Effect.Effect<NonNullable<A>, E | ResourceNotResolved, R> =>
244 get.pipe(
245 Effect.filterOrFail(
246 (value): value is NonNullable<A> => value != null,
247 () => new ResourceNotResolved({ name }),
248 ),
249 Effect.retry({
250 while: (error) => error._tag === "GCP.TPU.ResourceNotResolved",
251 times: 8,
252 schedule: Schedule.spaced("1 second"),
253 }),
254 );
255
256export const waitUntilGone = <A, E extends { readonly _tag: string }, R>(
257 get: Effect.Effect<A | undefined, E, R>,
258 name: string,
259): Effect.Effect<void, E | ResourceStillExists, R> =>
260 get.pipe(
261 Effect.filterOrFail(
262 (value) => value === undefined,
263 () => new ResourceStillExists({ name }),
264 ),
265 Effect.retry({
266 while: (error) => error._tag === "GCP.TPU.ResourceStillExists",
267 times: 10,
268 schedule: Schedule.spaced("8 seconds"),
269 }),
270 Effect.asVoid,
271 );
272
273const READY_STATES = new Set(["READY", "ACTIVE", "ACCEPTED"]);
274const FAILED_STATES = new Set([
275 "FAILED",
276 "ERROR",
277 "TERMINATED",
278 "PREEMPTED",
279 "UNKNOWN",
280]);
281
282const isReadyState = (state: string | undefined) =>
283 READY_STATES.has((state ?? "").toUpperCase());
284
285const isFailedState = (state: string | undefined) =>
286 FAILED_STATES.has((state ?? "").toUpperCase());
287
288export const waitUntilReady = <A, E extends { readonly _tag: string }, R>(
289 get: Effect.Effect<A | undefined, E, R>,
290 name: string,
291 stateOf: (value: NonNullable<A>) => string | undefined,
292 messageOf?: (value: NonNullable<A>) => string | undefined,
293): Effect.Effect<
294 NonNullable<A>,
295 E | ResourceNotResolved | ResourceFailed | ResourceNotReady,
296 R
297> =>
298 get.pipe(
299 Effect.filterOrFail(
300 (value): value is NonNullable<A> => value != null,
301 () => new ResourceNotResolved({ name }),
302 ),
303 Effect.filterOrFail(
304 (value) => !isFailedState(stateOf(value)),
305 (value) =>
306 new ResourceFailed({
307 name,
308 state: (stateOf(value) ?? "").toUpperCase(),
309 message: messageOf?.(value),
310 }),
311 ),
312 Effect.filterOrFail(
313 (value) => {
314 const state = stateOf(value) ?? "";
315 return isReadyState(state) || state.length === 0;
316 },
317 (value) =>
318 new ResourceNotReady({
319 name,
320 state: (stateOf(value) ?? "").toUpperCase(),
321 }),
322 ),
323 Effect.retry({
324 while: (error) =>
325 error._tag === "GCP.TPU.ResourceNotReady" ||
326 error._tag === "GCP.TPU.ResourceNotResolved",
327 times: 10,
328 schedule: Schedule.spaced("8 seconds"),
329 }),
330 );
331
332export const listLabeledPages = <Page, A, E, R>(
333 pages: Stream.Stream<Page, E, R>,
334 items: (page: Page) => readonly A[] | undefined,
335 labelsOf: (item: A) => Record<string, string | undefined> | null | undefined,
336) =>
337 pages.pipe(
338 Stream.flatMap((page) => Stream.fromIterable(items(page) ?? [])),
339 Stream.filter((item) => hasAlchemyLabelMap(labelsOf(item))),
340 Stream.runCollect,
341 Effect.map((chunk) => Array.from(chunk)),
342 );