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 * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { tagRecord } from "../../Tags.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import {
13 createInternalLabels,
14 diffLabels,
15 hasAlchemyLabels,
16 stripInternalLabels,
17 toLabels,
18} from "../Labels.ts";
19import type { Providers } from "../Providers.ts";
20import { listLocations } from "./names.ts";
21import { resourceNameFromOperation, waitForOperation } from "./operations.ts";
22import { lastSegment, locationOf, locationParent } from "./ownership.ts";
23import type { EncryptionSpec } from "./shared.ts";
24
25export type TensorboardProps = {
26 /**
27 * Region (`us-central1`, …). Immutable — changing it replaces the
28 * Tensorboard.
29 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
30 */
31 location?: string;
32 /**
33 * User-facing display name. If omitted, a unique name is generated from
34 * the stack, stage, and logical id.
35 */
36 displayName?: string;
37 /**
38 * Human-readable description.
39 */
40 description?: string;
41 /**
42 * User labels. Alchemy ownership labels are merged in automatically.
43 */
44 labels?: Record<string, string>;
45 /**
46 * Customer-managed encryption key. Immutable.
47 */
48 encryptionSpec?: EncryptionSpec;
49 /**
50 * Mark this Tensorboard as the default for the project and region.
51 * @default false
52 */
53 isDefault?: boolean;
54};
55
56export type Tensorboard = Resource<
57 "GCP.AIPlatform.Tensorboard",
58 TensorboardProps,
59 {
60 /** Full resource name `projects/{project}/locations/{location}/tensorboards/{tensorboard}`. */
61 name: string;
62 /** Tensorboard id (last path segment). */
63 tensorboardId: string;
64 /** Region id. */
65 location: string;
66 /** Project id. */
67 project: string;
68 /** User-facing display name. */
69 displayName: string | undefined;
70 /** Description. */
71 description: string | undefined;
72 /** User labels (Alchemy ownership labels stripped). */
73 labels: Record<string, string>;
74 /** CMEK resource name, if any. */
75 encryptionSpec: EncryptionSpec | undefined;
76 /** Whether this is the default Tensorboard in the region. */
77 isDefault: boolean;
78 /** Cloud Storage prefix used for blob data. */
79 blobStoragePathPrefix: string | undefined;
80 /** Number of runs stored in this Tensorboard. */
81 runCount: number | undefined;
82 /** RFC3339 creation timestamp. */
83 createTime: string | undefined;
84 /** RFC3339 last-update timestamp. */
85 updateTime: string | undefined;
86 },
87 never,
88 Providers
89>;
90
91/**
92 * A Vertex AI Tensorboard, the physical store for training metrics.
93 *
94 * Location and encryption key are immutable. Display name, description,
95 * labels, and the default flag update in place.
96 *
97 * ### Creating a Tensorboard
98 * **Example:** Generated display name
99 * ```typescript
100 * const board = yield* GCP.AIPlatform.Tensorboard("Metrics", {});
101 * ```
102 *
103 * **Example:** Named Tensorboard with labels
104 * ```typescript
105 * const board = yield* GCP.AIPlatform.Tensorboard("Metrics", {
106 * displayName: "training-metrics",
107 * description: "experiment store",
108 * labels: { env: "prod" },
109 * });
110 * ```
111 *
112 * @resource
113 * @category AIPlatform
114 */
115export const Tensorboard = Resource<Tensorboard>("GCP.AIPlatform.Tensorboard");
116
117export class TensorboardNotResolved extends Data.TaggedError(
118 "GCP.AIPlatform.TensorboardNotResolved",
119)<{
120 name: string;
121}> {}
122
123const userLabels = (
124 labels: Record<string, string | undefined> | null | undefined,
125): Record<string, string> => stripInternalLabels(tagRecord(labels));
126
127const toAttrs = (
128 board: aiplatform.GoogleCloudAiplatformV1Tensorboard,
129 project: string,
130) => {
131 const name = board.name ?? "";
132 return {
133 name,
134 tensorboardId: lastSegment(name),
135 location: locationOf(name),
136 project,
137 displayName: board.displayName,
138 description: board.description,
139 labels: userLabels(board.labels),
140 encryptionSpec: board.encryptionSpec?.kmsKeyName
141 ? { kmsKeyName: board.encryptionSpec.kmsKeyName }
142 : undefined,
143 isDefault: board.isDefault === true,
144 blobStoragePathPrefix: board.blobStoragePathPrefix,
145 runCount: board.runCount,
146 createTime: board.createTime,
147 updateTime: board.updateTime,
148 };
149};
150
151const getByName = (name: string) =>
152 name.length === 0
153 ? Effect.succeed(undefined)
154 : aiplatform
155 .getProjectsLocationsTensorboards({ name })
156 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
157
158const waitUntilExists = (name: string) =>
159 getByName(name).pipe(
160 Effect.flatMap((board) =>
161 board
162 ? Effect.succeed(board)
163 : Effect.fail(new TensorboardNotResolved({ name })),
164 ),
165 Effect.retry({
166 while: (error) => error._tag === "GCP.AIPlatform.TensorboardNotResolved",
167 times: 8,
168 schedule: Schedule.spaced("2 seconds"),
169 }),
170 );
171
172const listAt = (parent: string, project: string) =>
173 aiplatform.listProjectsLocationsTensorboards
174 .pages({ parent, pageSize: 100 })
175 .pipe(
176 Stream.flatMap((page) => Stream.fromIterable(page.tensorboards ?? [])),
177 Stream.filter((board) =>
178 Object.keys(board.labels ?? {}).some((key) =>
179 key.startsWith("alchemy-"),
180 ),
181 ),
182 Stream.map((board) => toAttrs(board, project)),
183 Stream.runCollect,
184 Effect.map((chunk) => Array.from(chunk)),
185 Effect.catchTag("NotFound", () => Effect.succeed([])),
186 );
187
188const findByDisplayName = (
189 parent: string,
190 displayName: string,
191 project: string,
192) =>
193 aiplatform.listProjectsLocationsTensorboards
194 .pages({ parent, pageSize: 100 })
195 .pipe(
196 Stream.flatMap((page) => Stream.fromIterable(page.tensorboards ?? [])),
197 Stream.filter((board) => board.displayName === displayName),
198 Stream.runHead,
199 Effect.map((option) =>
200 option._tag === "Some" ? toAttrs(option.value, project) : undefined,
201 ),
202 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
203 );
204
205export const TensorboardProvider = () =>
206 Provider.succeed(Tensorboard, {
207 stables: ["name", "tensorboardId", "location", "project", "createTime"],
208
209 diff: Effect.fn(function* ({ news, olds, output }) {
210 if (!isResolved(news)) return undefined;
211 const previousLocation = olds?.location ?? output?.location;
212 const nextLocation = news.location ?? previousLocation;
213 if (previousLocation !== undefined && previousLocation !== nextLocation) {
214 return { action: "replace" as const, deleteFirst: false };
215 }
216 const previousKey =
217 olds?.encryptionSpec?.kmsKeyName ?? output?.encryptionSpec?.kmsKeyName;
218 const nextKey = news.encryptionSpec?.kmsKeyName;
219 if (
220 previousKey !== undefined &&
221 nextKey !== undefined &&
222 previousKey !== nextKey
223 ) {
224 return { action: "replace" as const, deleteFirst: false };
225 }
226 return undefined;
227 }),
228
229 read: Effect.fn(function* ({ id, olds, output }) {
230 const env = yield* GcpEnvironment.current;
231 const existing = yield* getByName(output?.name ?? "");
232 if (existing === undefined) {
233 if (output?.name !== undefined) return undefined;
234 const location = olds?.location ?? output?.location ?? env.region;
235 const parent = locationParent(env.project, location);
236 const match = yield* findByDisplayName(
237 parent,
238 olds?.displayName ?? "",
239 env.project,
240 );
241 if (match === undefined) return undefined;
242 const fetched = yield* getByName(match.name);
243 if (fetched === undefined) return undefined;
244 const attrs = toAttrs(fetched, env.project);
245 return (yield* hasAlchemyLabels(id, tagRecord(fetched.labels)))
246 ? attrs
247 : Unowned(attrs);
248 }
249 const attrs = toAttrs(existing, env.project);
250 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
251 ? attrs
252 : Unowned(attrs);
253 }),
254
255 list: () =>
256 Effect.gen(function* () {
257 const env = yield* GcpEnvironment.current;
258 return (yield* Effect.forEach(listLocations(env.region), (location) =>
259 listAt(locationParent(env.project, location), env.project),
260 )).flat();
261 }),
262
263 reconcile: Effect.fn(function* ({ id, news, output }) {
264 const env = yield* GcpEnvironment.current;
265 const location = news.location ?? output?.location ?? env.region;
266 const parent = locationParent(env.project, location);
267 const desiredLabels = {
268 ...toLabels(news.labels),
269 ...(yield* createInternalLabels(id)),
270 };
271 const displayName =
272 news.displayName ?? output?.displayName ?? lastSegment(id);
273 const desiredDefault = news.isDefault === true;
274 const encryptionSpec = news.encryptionSpec?.kmsKeyName
275 ? { kmsKeyName: news.encryptionSpec.kmsKeyName }
276 : undefined;
277
278 let current = yield* getByName(output?.name ?? "");
279
280 if (current === undefined) {
281 const created = yield* aiplatform
282 .createProjectsLocationsTensorboards({
283 parent,
284 body: {
285 displayName,
286 description: news.description,
287 labels: desiredLabels,
288 encryptionSpec,
289 isDefault: desiredDefault ? true : undefined,
290 },
291 })
292 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
293 if (created !== undefined) {
294 const done = yield* waitForOperation(created, {
295 alreadyExistsOk: true,
296 });
297 const createdName =
298 resourceNameFromOperation(done) ?? output?.name ?? "";
299 if (createdName.length > 0) {
300 current = yield* waitUntilExists(createdName);
301 }
302 }
303 if (current === undefined) {
304 const match = yield* findByDisplayName(
305 parent,
306 displayName,
307 env.project,
308 );
309 if (match !== undefined) {
310 current = yield* getByName(match.name);
311 }
312 }
313 }
314
315 if (current === undefined) {
316 return yield* new TensorboardNotResolved({
317 name: output?.name ?? `${parent}/tensorboards/-`,
318 });
319 }
320
321 const name = current.name ?? "";
322 const observedLabels = tagRecord(current.labels);
323 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
324 const labelsChanged = upsert.length > 0 || removed.length > 0;
325 const displayChanged = (current.displayName ?? "") !== displayName;
326 const descriptionChanged =
327 (current.description ?? "") !== (news.description ?? "");
328 const defaultChanged = (current.isDefault === true) !== desiredDefault;
329
330 if (
331 labelsChanged ||
332 displayChanged ||
333 descriptionChanged ||
334 defaultChanged
335 ) {
336 const updateMask = [
337 labelsChanged ? "labels" : undefined,
338 displayChanged ? "displayName" : undefined,
339 descriptionChanged ? "description" : undefined,
340 defaultChanged ? "isDefault" : undefined,
341 ].filter((field): field is string => field !== undefined);
342 const patched = yield* aiplatform.patchProjectsLocationsTensorboards({
343 name,
344 updateMask: updateMask.join(","),
345 body: {
346 name,
347 displayName,
348 description: news.description,
349 labels: desiredLabels,
350 isDefault: desiredDefault,
351 },
352 });
353 yield* waitForOperation(patched);
354 current = yield* waitUntilExists(name);
355 }
356
357 return toAttrs(current, env.project);
358 }),
359
360 delete: Effect.fn(function* ({ output }) {
361 const existing = yield* getByName(output.name);
362 if (existing === undefined) return;
363 const operation = yield* aiplatform
364 .deleteProjectsLocationsTensorboards({ name: output.name })
365 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
366 if (operation !== undefined) {
367 yield* waitForOperation(operation, { notFoundOk: true });
368 }
369 }),
370 });