1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import { createHash } from "node:crypto";
3import * as Data from "effect/Data";
4import * as Effect from "effect/Effect";
5import * as Schedule from "effect/Schedule";
6import * as Stream from "effect/Stream";
7import { Unowned } from "../../AdoptPolicy.ts";
8import { isResolved } from "../../Diff.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { createInternalLabels, hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import { listLocations } from "./names.ts";
15import {
16 hasAlchemyLabelKeys,
17 parseOwnership,
18 parseResourceName,
19 resourceNameFromOperation,
20 waitForOperation,
21} from "./internal.ts";
22
23export type DatasetsDatasetVersionProps = {
24 /**
25 * Parent Dataset resource name
26 * `projects/{project}/locations/{location}/datasets/{dataset}`.
27 * Immutable — changing it replaces the version.
28 */
29 dataset: string;
30 /**
31 * User-facing display name (max 128 characters). Alchemy stamps
32 * ownership into this field (versions have no labels) so `list` /
33 * nuke can find the version.
34 */
35 displayName?: string;
36};
37
38export type DatasetsDatasetVersion = Resource<
39 "GCP.AIPlatform.DatasetsDatasetVersion",
40 DatasetsDatasetVersionProps,
41 {
42 /** Full resource name `.../datasets/{dataset}/datasetVersions/{version}`. */
43 name: string;
44 /** Dataset version id (last path segment). */
45 datasetVersionId: string;
46 /** Parent dataset resource name. */
47 dataset: string;
48 /** Project id. */
49 project: string;
50 /** Location id. */
51 location: string;
52 /** User display name with the Alchemy ownership prefix stripped. */
53 displayName: string | undefined;
54 /** Schema-specific metadata. */
55 metadata: unknown;
56 /** Associated BigQuery dataset name, if any. */
57 bigQueryDatasetName: string | undefined;
58 /** Public base model reference, if any. */
59 modelReference: string | undefined;
60 /** RFC3339 creation timestamp. */
61 createTime: string | undefined;
62 /** RFC3339 last-update timestamp. */
63 updateTime: string | undefined;
64 },
65 never,
66 Providers
67>;
68
69/**
70 * A Vertex AI DatasetVersion — a snapshot of a Dataset.
71 *
72 * Versions have no labels field, so Alchemy stamps ownership into
73 * `displayName` for `list` / nuke. The parent dataset is immutable;
74 * display name updates in place.
75 *
76 * ### Creating a Dataset Version
77 * **Example:** Snapshot a dataset
78 * ```typescript
79 * const dataset = yield* GCP.AIPlatform.Dataset("Samples", {
80 * metadataSchemaUri:
81 * "gs://google-cloud-aiplatform/schema/dataset/metadata/image_1.0.0.yaml",
82 * });
83 * const version = yield* GCP.AIPlatform.DatasetsDatasetVersion("V1", {
84 * dataset: dataset.name,
85 * displayName: "v1",
86 * });
87 * ```
88 *
89 * ### Updating a Dataset Version
90 * **Example:** Rename the snapshot
91 * ```typescript
92 * // Same logical id, changed props: the engine updates it in place.
93 * const version = yield* GCP.AIPlatform.DatasetsDatasetVersion("V1", {
94 * dataset: dataset.name,
95 * displayName: "v1-final",
96 * });
97 * ```
98 *
99 * @resource
100 * @category AIPlatform
101 */
102export const DatasetsDatasetVersion = Resource<DatasetsDatasetVersion>(
103 "GCP.AIPlatform.DatasetsDatasetVersion",
104);
105
106export class DatasetsDatasetVersionNotResolved extends Data.TaggedError(
107 "GCP.AIPlatform.DatasetsDatasetVersionNotResolved",
108)<{
109 name: string;
110}> {}
111
112export class DatasetsDatasetVersionStillExists extends Data.TaggedError(
113 "GCP.AIPlatform.DatasetsDatasetVersionStillExists",
114)<{
115 name: string;
116}> {}
117
118const parentDataset = (name: string) => {
119 const marker = "/datasetVersions/";
120 const index = name.indexOf(marker);
121 return index >= 0 ? name.slice(0, index) : name;
122};
123
124const parseVersionName = (name: string) => {
125 const parsed = parseResourceName(name, "datasetVersions");
126 const datasets = parseResourceName(name, "datasets");
127 return {
128 ...parsed,
129 datasetId: datasets.id,
130 dataset: parentDataset(name),
131 };
132};
133
134const getByName = (name: string) =>
135 name.length === 0
136 ? Effect.succeed(undefined)
137 : aiplatform
138 .getDatasetsDatasetVersions({ name })
139 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
140
141/**
142 * Dataset versions carry ownership only in their display name, which
143 * Vertex AI caps at 128 characters — too short for the full
144 * `[alchemy stack=… stage=… id=…]` marker of long stack names. Stamp a
145 * fixed-width hash of the ownership labels instead.
146 */
147const ownerTag = (id: string) =>
148 createInternalLabels(id).pipe(
149 Effect.flatMap((labels) =>
150 Effect.sync(
151 () =>
152 `[alchemy #${createHash("sha256")
153 .update(JSON.stringify(labels))
154 .digest("hex")
155 .slice(0, 16)}]`,
156 ),
157 ),
158 );
159
160const encodeDisplayName = (tag: string, displayName: string | undefined) =>
161 displayName ? `${tag} ${displayName}` : tag;
162
163const ownedByAlchemy = (id: string, displayName: string | undefined) =>
164 Effect.gen(function* () {
165 if (displayName?.startsWith(yield* ownerTag(id))) return true;
166 return yield* hasAlchemyLabels(id, parseOwnership(displayName).labels);
167 });
168
169const toAttrs = (
170 version: aiplatform.GoogleCloudAiplatformV1DatasetVersion,
171 project: string,
172) => {
173 const name = version.name ?? "";
174 const parsed = parseVersionName(name);
175 const ownership = parseOwnership(version.displayName);
176 return {
177 name,
178 datasetVersionId: parsed.id,
179 dataset: parsed.dataset,
180 project: parsed.project || project,
181 location: parsed.location,
182 displayName: ownership.text,
183 metadata: version.metadata,
184 bigQueryDatasetName: version.bigQueryDatasetName,
185 modelReference: version.modelReference,
186 createTime: version.createTime,
187 updateTime: version.updateTime,
188 };
189};
190
191const listVersions = (dataset: string) =>
192 aiplatform.listProjectsLocationsDatasetsDatasetVersions
193 .pages({ parent: dataset, pageSize: 100 })
194 .pipe(
195 Stream.flatMap((page) => Stream.fromIterable(page.datasetVersions ?? [])),
196 Stream.runCollect,
197 Effect.map((chunk) => Array.from(chunk)),
198 Effect.catchTag("NotFound", () => Effect.succeed([])),
199 );
200
201const listDatasets = (project: string, region: string) =>
202 Effect.forEach(listLocations(region), (location) =>
203 listDatasetsAt(`projects/${project}/locations/${location}`),
204 ).pipe(Effect.map((pages) => pages.flat()));
205
206const listDatasetsAt = (parent: string) =>
207 aiplatform.listProjectsLocationsDatasets
208 .pages({
209 parent,
210 pageSize: 100,
211 })
212 .pipe(
213 Stream.flatMap((page) => Stream.fromIterable(page.datasets ?? [])),
214 Stream.filter((dataset) => hasAlchemyLabelKeys(dataset.labels)),
215 Stream.runCollect,
216 Effect.map((chunk) => Array.from(chunk)),
217 Effect.catchTag("NotFound", () => Effect.succeed([])),
218 );
219
220const findOwned = (id: string, dataset: string, hinted?: string) =>
221 Effect.gen(function* () {
222 if (hinted !== undefined && hinted.length > 0) {
223 const existing = yield* getByName(hinted);
224 if (existing !== undefined) return existing;
225 }
226 const versions = yield* listVersions(dataset);
227 for (const version of versions) {
228 if (yield* ownedByAlchemy(id, version.displayName)) return version;
229 }
230 return undefined as
231 | aiplatform.GoogleCloudAiplatformV1DatasetVersion
232 | undefined;
233 });
234
235const waitUntilExists = (name: string) =>
236 getByName(name).pipe(
237 Effect.flatMap((version) =>
238 version
239 ? Effect.succeed(version)
240 : Effect.fail(new DatasetsDatasetVersionNotResolved({ name })),
241 ),
242 Effect.retry({
243 while: (error) =>
244 error._tag === "GCP.AIPlatform.DatasetsDatasetVersionNotResolved",
245 times: 8,
246 schedule: Schedule.spaced("1 second"),
247 }),
248 );
249
250const waitUntilGone = (name: string) =>
251 getByName(name).pipe(
252 Effect.flatMap((version) =>
253 version === undefined
254 ? Effect.void
255 : Effect.fail(new DatasetsDatasetVersionStillExists({ name })),
256 ),
257 Effect.retry({
258 while: (error) =>
259 error._tag === "GCP.AIPlatform.DatasetsDatasetVersionStillExists",
260 times: 10,
261 schedule: Schedule.spaced("2 seconds"),
262 }),
263 );
264
265export const DatasetsDatasetVersionProvider = () =>
266 Provider.succeed(DatasetsDatasetVersion, {
267 stables: [
268 "name",
269 "datasetVersionId",
270 "dataset",
271 "project",
272 "location",
273 "createTime",
274 ],
275
276 diff: Effect.fn(function* ({ news, olds, output }) {
277 if (!isResolved(news)) return undefined;
278 const previous = olds?.dataset ?? output?.dataset;
279 if (previous !== undefined && news.dataset !== previous) {
280 return { action: "replace" as const, deleteFirst: false };
281 }
282 return undefined;
283 }),
284
285 read: Effect.fn(function* ({ id, olds, output }) {
286 const env = yield* GcpEnvironment.current;
287 const dataset = olds?.dataset ?? output?.dataset;
288 const existing =
289 output?.name !== undefined
290 ? yield* getByName(output.name)
291 : dataset !== undefined
292 ? yield* findOwned(id, dataset)
293 : undefined;
294 if (existing === undefined) return undefined;
295 const attrs = toAttrs(existing, env.project);
296 return (yield* ownedByAlchemy(id, existing.displayName))
297 ? attrs
298 : Unowned(attrs);
299 }),
300
301 list: () =>
302 Effect.gen(function* () {
303 const env = yield* GcpEnvironment.current;
304 const datasets = yield* listDatasets(env.project, env.region);
305 const versions: ReturnType<typeof toAttrs>[] = [];
306 for (const dataset of datasets) {
307 if (dataset.name === undefined) continue;
308 const page = yield* listVersions(dataset.name);
309 for (const version of page) {
310 if (version.displayName?.startsWith("[alchemy ")) {
311 versions.push(toAttrs(version, env.project));
312 }
313 }
314 }
315 return versions;
316 }),
317
318 reconcile: Effect.fn(function* ({ id, news, output }) {
319 const env = yield* GcpEnvironment.current;
320 const desiredDisplayName = encodeDisplayName(
321 yield* ownerTag(id),
322 news.displayName,
323 );
324
325 let current = yield* findOwned(id, news.dataset, output?.name);
326
327 if (current === undefined) {
328 const created = yield* aiplatform
329 .createDatasetsDatasetVersions({
330 parent: news.dataset,
331 body: { displayName: desiredDisplayName },
332 })
333 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
334 if (created !== undefined) {
335 const done = yield* waitForOperation(created);
336 const createdName =
337 resourceNameFromOperation(done) ??
338 (yield* findOwned(id, news.dataset))?.name;
339 if (createdName !== undefined && createdName.length > 0) {
340 current = yield* waitUntilExists(createdName);
341 }
342 }
343 if (current === undefined) {
344 current = yield* findOwned(id, news.dataset);
345 }
346 }
347
348 if (current === undefined) {
349 return yield* new DatasetsDatasetVersionNotResolved({
350 name: output?.name ?? `${news.dataset}/datasetVersions/-`,
351 });
352 }
353
354 const name = current.name ?? "";
355 if ((current.displayName ?? "") !== desiredDisplayName) {
356 current =
357 yield* aiplatform.patchProjectsLocationsDatasetsDatasetVersions({
358 name,
359 updateMask: "display_name",
360 body: { name, displayName: desiredDisplayName },
361 });
362 }
363
364 return toAttrs(current, env.project);
365 }),
366
367 delete: Effect.fn(function* ({ output }) {
368 const existing = yield* getByName(output.name);
369 if (existing === undefined) return;
370 const operation = yield* aiplatform
371 .deleteDatasetsDatasetVersions({ name: output.name })
372 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
373 if (operation !== undefined) {
374 yield* waitForOperation(operation, { notFoundOk: true });
375 }
376 yield* waitUntilGone(output.name);
377 }),
378 });