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 { hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import {
15 createInternalLabels,
16 hasAlchemyPrefix,
17 labelsDiffer,
18 listLocations,
19 locationOf,
20 lastSegment,
21 projectOf,
22 stableJson,
23 toDisplayName,
24 toLabels,
25 toResourceId,
26 userLabels,
27} from "./names.ts";
28import { waitForOperation } from "./operations.ts";
29
30const DEFAULT_SCHEMA_TITLE = "system.Run";
31
32export type MetadataStoresExecutionProps = {
33 /**
34 * Parent MetadataStore resource name
35 * `projects/{project}/locations/{location}/metadataStores/{metadatastore}`.
36 * Immutable — changing it replaces the execution.
37 */
38 metadataStore: string;
39 /**
40 * Execution id. If omitted, a unique id is generated. Must be 4-128
41 * characters of `a-z` and `-`. Immutable.
42 */
43 executionId?: string;
44 /**
45 * Display name (max 128 Unicode characters).
46 */
47 displayName?: string;
48 /**
49 * Human-readable description.
50 */
51 description?: string;
52 /**
53 * User labels. Alchemy ownership labels are merged in automatically.
54 */
55 labels?: Record<string, string>;
56 /**
57 * Execution lifecycle state.
58 */
59 state?: aiplatform.GoogleCloudAiplatformV1ExecutionStateEnum | (string & {});
60 /**
61 * Schema title registered in the MetadataStore. Vertex AI rejects
62 * entries without one.
63 * @default "system.Run"
64 */
65 schemaTitle?: string;
66 /**
67 * Schema version registered in the MetadataStore.
68 */
69 schemaVersion?: string;
70 /**
71 * Arbitrary metadata properties.
72 */
73 metadata?: aiplatform.DocumentMap;
74};
75
76export type MetadataStoresExecution = Resource<
77 "GCP.AIPlatform.MetadataStoresExecution",
78 MetadataStoresExecutionProps,
79 {
80 /** Full resource name. */
81 name: string;
82 /** Execution id (last path segment). */
83 executionId: string;
84 /** Parent MetadataStore resource name. */
85 metadataStore: string;
86 /** Project id. */
87 project: string;
88 /** Location id. */
89 location: string;
90 /** Display name. */
91 displayName: string | undefined;
92 /** User-provided description. */
93 description: string | undefined;
94 /** User labels (Alchemy ownership labels stripped). */
95 labels: Record<string, string>;
96 /** Execution state. */
97 state: string | undefined;
98 /** Schema title. */
99 schemaTitle: string | undefined;
100 /** Schema version. */
101 schemaVersion: string | undefined;
102 /** RFC3339 creation timestamp. */
103 createTime: string | undefined;
104 /** RFC3339 last-update timestamp. */
105 updateTime: string | undefined;
106 },
107 never,
108 Providers
109>;
110
111/**
112 * A Vertex ML Metadata Execution associated with a MetadataStore.
113 *
114 * Changing `metadataStore` or `executionId` replaces the execution. Display
115 * name, description, labels, state, schema, and metadata update in place.
116 *
117 * ### Creating an Execution
118 * **Example:** Training execution
119 * ```typescript
120 * const execution = yield* GCP.AIPlatform.MetadataStoresExecution("Train", {
121 * metadataStore: store.name,
122 * displayName: "train-step",
123 * state: "RUNNING",
124 * labels: { env: "prod" },
125 * });
126 * ```
127 *
128 * @resource
129 * @category AIPlatform
130 */
131export const MetadataStoresExecution = Resource<MetadataStoresExecution>(
132 "GCP.AIPlatform.MetadataStoresExecution",
133);
134
135export class MetadataStoresExecutionNotResolved extends Data.TaggedError(
136 "GCP.AIPlatform.MetadataStoresExecutionNotResolved",
137)<{
138 name: string;
139}> {}
140
141export class MetadataStoresExecutionStillExists extends Data.TaggedError(
142 "GCP.AIPlatform.MetadataStoresExecutionStillExists",
143)<{
144 name: string;
145}> {}
146
147const parentOf = (name: string) => {
148 const at = name.lastIndexOf("/executions/");
149 return at >= 0 ? name.slice(0, at) : name;
150};
151
152const resourceName = (metadataStore: string, executionId: string) =>
153 `${metadataStore}/executions/${executionId}`;
154
155const toAttrs = (
156 execution: aiplatform.GoogleCloudAiplatformV1Execution,
157 project: string,
158) => {
159 const name = execution.name ?? "";
160 return {
161 name,
162 executionId: lastSegment(name),
163 metadataStore: parentOf(name),
164 project: projectOf(name, project),
165 location: locationOf(name),
166 displayName: execution.displayName,
167 description: execution.description,
168 labels: userLabels(execution.labels),
169 state: execution.state,
170 schemaTitle: execution.schemaTitle,
171 schemaVersion: execution.schemaVersion,
172 createTime: execution.createTime,
173 updateTime: execution.updateTime,
174 };
175};
176
177const getByName = (name: string) =>
178 aiplatform
179 .getProjectsLocationsMetadataStoresExecutions({ name })
180 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
181
182const listExecutions = (parent: string) =>
183 aiplatform.listProjectsLocationsMetadataStoresExecutions
184 .pages({ parent, pageSize: 100 })
185 .pipe(
186 Stream.runCollect,
187 Effect.map((pages) =>
188 Array.from(pages).flatMap((page) => page.executions ?? []),
189 ),
190 Effect.catchTag("NotFound", () =>
191 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1Execution[]),
192 ),
193 );
194
195const listStores = (parent: string) =>
196 aiplatform.listProjectsLocationsMetadataStores
197 .pages({ parent, pageSize: 100 })
198 .pipe(
199 Stream.runCollect,
200 Effect.map((pages) =>
201 Array.from(pages).flatMap((page) => page.metadataStores ?? []),
202 ),
203 Effect.catchTag("NotFound", () =>
204 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1MetadataStore[]),
205 ),
206 );
207
208const waitUntilGone = (name: string) =>
209 getByName(name).pipe(
210 Effect.filterOrFail(
211 (execution) => execution === undefined,
212 () => new MetadataStoresExecutionStillExists({ name }),
213 ),
214 Effect.asVoid,
215 Effect.retry({
216 while: (error) =>
217 error._tag === "GCP.AIPlatform.MetadataStoresExecutionStillExists",
218 times: 10,
219 schedule: Schedule.spaced("1 second"),
220 }),
221 );
222
223export const MetadataStoresExecutionProvider = () =>
224 Provider.succeed(MetadataStoresExecution, {
225 stables: [
226 "name",
227 "executionId",
228 "metadataStore",
229 "project",
230 "location",
231 "createTime",
232 ],
233
234 diff: Effect.fn(function* ({ news, olds, output }) {
235 if (!isResolved(news)) return undefined;
236 const previousParent = olds?.metadataStore ?? output?.metadataStore;
237 const nextParent = news.metadataStore ?? previousParent;
238 const previousId = olds?.executionId ?? output?.executionId;
239 const nextId = news.executionId ?? previousId;
240 if (
241 (previousParent !== undefined &&
242 nextParent !== undefined &&
243 nextParent !== previousParent) ||
244 (previousId !== undefined &&
245 nextId !== undefined &&
246 nextId !== previousId)
247 ) {
248 return {
249 action: "replace" as const,
250 deleteFirst:
251 previousParent === nextParent &&
252 previousId !== undefined &&
253 nextId === previousId,
254 };
255 }
256 return undefined;
257 }),
258
259 read: Effect.fn(function* ({ id, olds, output }) {
260 const env = yield* GcpEnvironment.current;
261 const executionId = yield* toResourceId(
262 id,
263 olds?.executionId,
264 output?.executionId,
265 );
266 const metadataStore = olds?.metadataStore ?? output?.metadataStore;
267 const name =
268 output?.name ??
269 (metadataStore !== undefined
270 ? resourceName(metadataStore, executionId)
271 : "");
272 if (name.length === 0) return undefined;
273 const existing = yield* getByName(name);
274 if (existing === undefined) return undefined;
275 const attrs = toAttrs(existing, env.project);
276 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
277 ? attrs
278 : Unowned(attrs);
279 }),
280
281 list: () =>
282 Effect.gen(function* () {
283 const env = yield* GcpEnvironment.current;
284 const stores = yield* Effect.forEach(
285 listLocations(env.region),
286 (location) =>
287 listStores(`projects/${env.project}/locations/${location}`),
288 { concurrency: 4 },
289 );
290 const executions = yield* Effect.forEach(
291 stores.flat().filter((store) => store.name !== undefined),
292 (store) => listExecutions(store.name!),
293 { concurrency: 4 },
294 );
295 return executions
296 .flat()
297 .filter((execution) => hasAlchemyPrefix(execution.labels))
298 .map((execution) => toAttrs(execution, env.project));
299 }),
300
301 reconcile: Effect.fn(function* ({ id, news, output }) {
302 const env = yield* GcpEnvironment.current;
303 const executionId = yield* toResourceId(
304 id,
305 news.executionId,
306 output?.executionId,
307 );
308 const metadataStore = news.metadataStore;
309 const name = resourceName(metadataStore, executionId);
310 const displayName = yield* toDisplayName(
311 id,
312 news.displayName,
313 output?.displayName,
314 );
315 const desiredLabels = {
316 ...toLabels(news.labels),
317 ...(yield* createInternalLabels(id)),
318 };
319
320 let current = yield* getByName(output?.name ?? name);
321
322 if (current === undefined) {
323 const created = yield* aiplatform
324 .createProjectsLocationsMetadataStoresExecutions({
325 parent: metadataStore,
326 executionId,
327 body: {
328 displayName,
329 description: news.description,
330 labels: desiredLabels,
331 state: news.state,
332 schemaTitle: news.schemaTitle ?? DEFAULT_SCHEMA_TITLE,
333 schemaVersion: news.schemaVersion,
334 metadata: news.metadata,
335 },
336 })
337 .pipe(Effect.catchTag("Conflict", () => getByName(name)));
338 current = created ?? undefined;
339 }
340
341 if (current === undefined) {
342 return yield* new MetadataStoresExecutionNotResolved({ name });
343 }
344
345 const descriptionChanged =
346 (current.description ?? "") !== (news.description ?? "");
347 const displayChanged = (current.displayName ?? "") !== displayName;
348 const labelsChanged = labelsDiffer(current.labels, desiredLabels);
349 const stateChanged = (current.state ?? "") !== (news.state ?? "");
350 const schemaTitleChanged =
351 (current.schemaTitle ?? "") !==
352 (news.schemaTitle ?? DEFAULT_SCHEMA_TITLE);
353 const schemaVersionChanged =
354 (current.schemaVersion ?? "") !== (news.schemaVersion ?? "");
355 const metadataChanged =
356 news.metadata !== undefined &&
357 stableJson(current.metadata) !== stableJson(news.metadata);
358
359 if (
360 descriptionChanged ||
361 displayChanged ||
362 labelsChanged ||
363 stateChanged ||
364 schemaTitleChanged ||
365 schemaVersionChanged ||
366 metadataChanged
367 ) {
368 current =
369 yield* aiplatform.patchProjectsLocationsMetadataStoresExecutions({
370 name,
371 updateMask: [
372 descriptionChanged ? "description" : undefined,
373 displayChanged ? "display_name" : undefined,
374 labelsChanged ? "labels" : undefined,
375 stateChanged ? "state" : undefined,
376 schemaTitleChanged ? "schema_title" : undefined,
377 schemaVersionChanged ? "schema_version" : undefined,
378 metadataChanged ? "metadata" : undefined,
379 ]
380 .filter((field): field is string => field !== undefined)
381 .join(","),
382 body: {
383 name,
384 displayName,
385 description: news.description,
386 labels: desiredLabels,
387 state: news.state,
388 schemaTitle: news.schemaTitle ?? DEFAULT_SCHEMA_TITLE,
389 schemaVersion: news.schemaVersion,
390 metadata: news.metadata,
391 etag: current.etag,
392 },
393 });
394 }
395
396 return toAttrs(current, env.project);
397 }),
398
399 delete: Effect.fn(function* ({ output }) {
400 const operation = yield* aiplatform
401 .deleteProjectsLocationsMetadataStoresExecutions({
402 name: output.name,
403 })
404 .pipe(
405 Effect.retry({
406 while: (error) => error._tag === "Conflict",
407 times: 8,
408 schedule: Schedule.spaced("2 seconds"),
409 }),
410 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
411 );
412 if (operation !== undefined) {
413 yield* waitForOperation(operation, { notFoundOk: true });
414 }
415 yield* waitUntilGone(output.name);
416 }),
417 });