1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
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 {
12 createInternalLabels,
13 diffLabels,
14 hasAlchemyLabels,
15 stripInternalLabels,
16 toLabels,
17} from "../Labels.ts";
18import type { Providers } from "../Providers.ts";
19import { listLocations } from "./names.ts";
20import { waitForOperation } from "./operations.ts";
21import {
22 lastSegment,
23 locationOf,
24 locationParent,
25 parentOf,
26 toResourceId,
27} from "./ownership.ts";
28
29export type TensorboardsExperimentProps = {
30 /**
31 * Parent Tensorboard resource name
32 * (`projects/{project}/locations/{location}/tensorboards/{tensorboard}`).
33 * Immutable — changing it replaces the experiment.
34 */
35 parent: string;
36 /**
37 * Experiment id (the `{experiment}` segment). If omitted, a unique id is
38 * generated from the stack, stage, and logical id. Immutable.
39 */
40 experimentId?: string;
41 /**
42 * User-facing display name.
43 */
44 displayName?: string;
45 /**
46 * Human-readable description.
47 */
48 description?: string;
49 /**
50 * User labels. Alchemy ownership labels are merged in automatically.
51 */
52 labels?: Record<string, string>;
53 /**
54 * Immutable source of the experiment (for example a custom training job).
55 */
56 source?: string;
57};
58
59export type TensorboardsExperiment = Resource<
60 "GCP.AIPlatform.TensorboardsExperiment",
61 TensorboardsExperimentProps,
62 {
63 /** Full resource name. */
64 name: string;
65 /** Experiment id (last path segment). */
66 experimentId: string;
67 /** Parent Tensorboard resource name. */
68 parent: string;
69 /** Region id. */
70 location: string;
71 /** Project id. */
72 project: string;
73 /** User-facing display name. */
74 displayName: string | undefined;
75 /** Description. */
76 description: string | undefined;
77 /** User labels (Alchemy ownership labels stripped). */
78 labels: Record<string, string>;
79 /** Immutable source, if set. */
80 source: string | undefined;
81 /** RFC3339 creation timestamp. */
82 createTime: string | undefined;
83 /** RFC3339 last-update timestamp. */
84 updateTime: string | undefined;
85 },
86 never,
87 Providers
88>;
89
90/**
91 * A Vertex AI TensorboardExperiment grouping runs from a training job.
92 *
93 * Parent Tensorboard, experiment id, and source are immutable. Display
94 * name, description, and labels update in place.
95 *
96 * ### Creating an Experiment
97 * **Example:** Experiment under a Tensorboard
98 * ```typescript
99 * const experiment = yield* GCP.AIPlatform.TensorboardsExperiment("RunGroup", {
100 * parent: board.name,
101 * displayName: "baseline",
102 * });
103 * ```
104 *
105 * @resource
106 * @category AIPlatform
107 */
108export const TensorboardsExperiment = Resource<TensorboardsExperiment>(
109 "GCP.AIPlatform.TensorboardsExperiment",
110);
111
112export class TensorboardsExperimentNotResolved extends Data.TaggedError(
113 "GCP.AIPlatform.TensorboardsExperimentNotResolved",
114)<{
115 name: string;
116}> {}
117
118const userLabels = (
119 labels: Record<string, string | undefined> | null | undefined,
120): Record<string, string> => stripInternalLabels(tagRecord(labels));
121
122const resourceName = (parent: string, experimentId: string) =>
123 `${parent}/experiments/${experimentId}`;
124
125const toAttrs = (
126 experiment: aiplatform.GoogleCloudAiplatformV1TensorboardExperiment,
127 project: string,
128) => {
129 const name = experiment.name ?? "";
130 return {
131 name,
132 experimentId: lastSegment(name),
133 parent: parentOf(name),
134 location: locationOf(name),
135 project,
136 displayName: experiment.displayName,
137 description: experiment.description,
138 labels: userLabels(experiment.labels),
139 source: experiment.source,
140 createTime: experiment.createTime,
141 updateTime: experiment.updateTime,
142 };
143};
144
145const getByName = (name: string) =>
146 name.length === 0
147 ? Effect.succeed(undefined)
148 : aiplatform
149 .getProjectsLocationsTensorboardsExperiments({ name })
150 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
151
152const listParents = (project: string, location: string) =>
153 aiplatform.listProjectsLocationsTensorboards
154 .pages({
155 parent: locationParent(project, location),
156 pageSize: 100,
157 })
158 .pipe(
159 Stream.flatMap((page) => Stream.fromIterable(page.tensorboards ?? [])),
160 Stream.map((board) => board.name ?? ""),
161 Stream.filter((name) => name.length > 0),
162 Stream.runCollect,
163 Effect.map((chunk) => Array.from(chunk)),
164 Effect.catchTag("NotFound", () => Effect.succeed([] as string[])),
165 );
166
167const listAtParent = (parent: string, project: string) =>
168 aiplatform.listProjectsLocationsTensorboardsExperiments
169 .pages({ parent, pageSize: 100 })
170 .pipe(
171 Stream.flatMap((page) =>
172 Stream.fromIterable(page.tensorboardExperiments ?? []),
173 ),
174 Stream.filter((experiment) =>
175 Object.keys(experiment.labels ?? {}).some((key) =>
176 key.startsWith("alchemy-"),
177 ),
178 ),
179 Stream.map((experiment) => toAttrs(experiment, project)),
180 Stream.runCollect,
181 Effect.map((chunk) => Array.from(chunk)),
182 Effect.catchTag("NotFound", () => Effect.succeed([])),
183 );
184
185export const TensorboardsExperimentProvider = () =>
186 Provider.succeed(TensorboardsExperiment, {
187 stables: [
188 "name",
189 "experimentId",
190 "parent",
191 "location",
192 "project",
193 "createTime",
194 ],
195
196 diff: Effect.fn(function* ({ news, olds, output }) {
197 if (!isResolved(news)) return undefined;
198 const previousParent = olds?.parent ?? output?.parent;
199 if (previousParent !== undefined && news.parent !== previousParent) {
200 return { action: "replace" as const, deleteFirst: false };
201 }
202 const previousId = olds?.experimentId ?? output?.experimentId;
203 if (
204 previousId !== undefined &&
205 news.experimentId !== undefined &&
206 news.experimentId !== previousId
207 ) {
208 return { action: "replace" as const, deleteFirst: true };
209 }
210 return undefined;
211 }),
212
213 read: Effect.fn(function* ({ id, olds, output }) {
214 const env = yield* GcpEnvironment.current;
215 const experimentId = yield* toResourceId(
216 id,
217 olds?.experimentId,
218 output?.experimentId,
219 );
220 const name =
221 output?.name ??
222 (olds?.parent !== undefined
223 ? resourceName(olds.parent, experimentId)
224 : "");
225 const existing = yield* getByName(name);
226 if (existing === undefined) return undefined;
227 const attrs = toAttrs(existing, env.project);
228 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
229 ? attrs
230 : Unowned(attrs);
231 }),
232
233 list: () =>
234 Effect.gen(function* () {
235 const env = yield* GcpEnvironment.current;
236 const parents = (yield* Effect.forEach(
237 listLocations(env.region),
238 (location) => listParents(env.project, location),
239 )).flat();
240 const pages = yield* Effect.forEach(
241 parents,
242 (parent) => listAtParent(parent, env.project),
243 { concurrency: 4 },
244 );
245 return pages.flat();
246 }),
247
248 reconcile: Effect.fn(function* ({ id, news, output }) {
249 const env = yield* GcpEnvironment.current;
250 const experimentId = yield* toResourceId(
251 id,
252 news.experimentId,
253 output?.experimentId,
254 );
255 const name = resourceName(news.parent, experimentId);
256 const desiredLabels = {
257 ...toLabels(news.labels),
258 ...(yield* createInternalLabels(id)),
259 };
260 const displayName = news.displayName ?? experimentId;
261
262 let current = yield* getByName(output?.name ?? name);
263
264 if (current === undefined) {
265 const created = yield* aiplatform
266 .createProjectsLocationsTensorboardsExperiments({
267 parent: news.parent,
268 tensorboardExperimentId: experimentId,
269 body: {
270 displayName,
271 description: news.description,
272 labels: desiredLabels,
273 source: news.source,
274 },
275 })
276 .pipe(Effect.catchTag("Conflict", () => getByName(name)));
277 current = created ?? undefined;
278 }
279
280 if (current === undefined) {
281 return yield* new TensorboardsExperimentNotResolved({ name });
282 }
283
284 const observedLabels = tagRecord(current.labels);
285 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
286 const labelsChanged = upsert.length > 0 || removed.length > 0;
287 const displayChanged = (current.displayName ?? "") !== displayName;
288 const descriptionChanged =
289 (current.description ?? "") !== (news.description ?? "");
290
291 if (labelsChanged || displayChanged || descriptionChanged) {
292 current =
293 yield* aiplatform.patchProjectsLocationsTensorboardsExperiments({
294 name,
295 updateMask: [
296 labelsChanged ? "labels" : undefined,
297 displayChanged ? "displayName" : undefined,
298 descriptionChanged ? "description" : undefined,
299 ]
300 .filter((field): field is string => field !== undefined)
301 .join(","),
302 body: {
303 name,
304 displayName,
305 description: news.description,
306 labels: desiredLabels,
307 },
308 });
309 }
310
311 return toAttrs(current, env.project);
312 }),
313
314 delete: Effect.fn(function* ({ output }) {
315 const existing = yield* getByName(output.name);
316 if (existing === undefined) return;
317 const operation = yield* aiplatform
318 .deleteProjectsLocationsTensorboardsExperiments({ name: output.name })
319 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
320 if (operation !== undefined) {
321 yield* waitForOperation(operation, { notFoundOk: true });
322 }
323 }),
324 });