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