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 { createPhysicalName } from "../../PhysicalName.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { tagRecord } from "../../Tags.ts";
12import { GcpEnvironment } from "../Environment.ts";
13import {
14 createInternalLabels,
15 diffLabels,
16 hasAlchemyLabels,
17 toLabels,
18} from "../Labels.ts";
19import type { Providers } from "../Providers.ts";
20import { listLocations } from "./names.ts";
21import {
22 DEFAULT_TABULAR_METADATA_SCHEMA_URI,
23 hasAlchemyLabelKeys,
24 normalizeLocation,
25 parentOf,
26 parseResourceName,
27 rfc1035,
28 resourceNameFromOperation,
29 userLabels,
30 waitForOperation,
31} from "./internal.ts";
32import type { EncryptionSpec } from "./shared.ts";
33
34export type DatasetProps = {
35 /**
36 * Dataset id (the `{dataset}` segment of
37 * `projects/{project}/locations/{location}/datasets/{dataset}`). If
38 * omitted, Vertex AI assigns an id. Immutable — changing it replaces
39 * the dataset.
40 */
41 datasetId?: string;
42 /**
43 * Region (`us-central1`, `us-east1`, …). Immutable — changing it
44 * replaces the dataset. `US-CENTRAL1` is accepted and normalized to
45 * `us-central1`.
46 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
47 */
48 location?: string;
49 /**
50 * User-facing display name (max 128 characters). Defaults to the
51 * dataset id.
52 */
53 displayName?: string;
54 /**
55 * Human-readable description.
56 */
57 description?: string;
58 /**
59 * User labels. Alchemy ownership labels are merged in automatically.
60 */
61 labels?: Record<string, string>;
62 /**
63 * GCS URI of the dataset metadata schema YAML (OpenAPI 3.0.2 Schema
64 * Object). Immutable — changing it replaces the dataset.
65 * @default "gs://google-cloud-aiplatform/schema/dataset/metadata/tabular_1.0.0.yaml"
66 */
67 metadataSchemaUri?: string;
68 /**
69 * Additional schema-specific metadata. For an empty tabular dataset,
70 * `{}` is sufficient.
71 */
72 metadata?: unknown;
73 /**
74 * Customer-managed encryption. Immutable — changing it replaces the
75 * dataset.
76 */
77 encryptionSpec?: EncryptionSpec;
78 /**
79 * Public base model last used by a prompt dataset.
80 */
81 modelReference?: string;
82 /**
83 * Saved queries created with the dataset (max one on create). Required
84 * for DatasetVersion — Vertex rejects versioning an empty dataset with
85 * no SavedQuery (`Unsupported dataset version creation with empty set
86 * of saved queries`).
87 */
88 savedQueries?: Array<{
89 /** User-facing saved query name. */
90 displayName?: string;
91 /**
92 * Problem type (`IMAGE_CLASSIFICATION_SINGLE_LABEL`,
93 * `TEXT_CLASSIFICATION_SINGLE_LABEL`, …).
94 */
95 problemType?: string;
96 /** Additional saved-query metadata. */
97 metadata?: unknown;
98 }>;
99};
100
101export type Dataset = Resource<
102 "GCP.AIPlatform.Dataset",
103 DatasetProps,
104 {
105 /** Full resource name `projects/{project}/locations/{location}/datasets/{dataset}`. */
106 name: string;
107 /** Dataset id (last path segment). */
108 datasetId: string;
109 /** Project id. */
110 project: string;
111 /** Location id (`us-central1`, …). */
112 location: string;
113 /** User-facing display name. */
114 displayName: string | undefined;
115 /** Description. */
116 description: string | undefined;
117 /** User labels (Alchemy ownership labels stripped). */
118 labels: Record<string, string>;
119 /** Metadata schema URI. */
120 metadataSchemaUri: string | undefined;
121 /** Schema-specific metadata. */
122 metadata: unknown;
123 /** Number of DataItems (non-structured datasets). */
124 dataItemCount: string | undefined;
125 /** MetadataStore artifact created with the dataset. */
126 metadataArtifact: string | undefined;
127 /** Public base model reference, if any. */
128 modelReference: string | undefined;
129 /** Customer-managed KMS key, if any. */
130 kmsKeyName: string | undefined;
131 /** RFC3339 creation timestamp. */
132 createTime: string | undefined;
133 /** RFC3339 last-update timestamp. */
134 updateTime: string | undefined;
135 },
136 never,
137 Providers
138>;
139
140/**
141 * A Vertex AI Dataset — a collection of DataItems and Annotations.
142 *
143 * Location, dataset id, metadata schema, and CMEK are immutable. Display
144 * name, description, and labels update in place. Alchemy ownership labels
145 * are merged into `labels` so `list` / `pnpm nuke:gcp` can find the
146 * dataset.
147 *
148 * ### Creating a Dataset
149 * **Example:** Generated id, empty tabular dataset
150 * ```typescript
151 * const dataset = yield* GCP.AIPlatform.Dataset("Samples", {
152 * displayName: "sample rows",
153 * labels: { env: "test" },
154 * });
155 * ```
156 *
157 * **Example:** Explicit id and image metadata schema
158 * ```typescript
159 * const dataset = yield* GCP.AIPlatform.Dataset("Images", {
160 * datasetId: "app-images",
161 * location: "us-central1",
162 * metadataSchemaUri:
163 * "gs://google-cloud-aiplatform/schema/dataset/metadata/image_1.0.0.yaml",
164 * metadata: {},
165 * });
166 * ```
167 *
168 * ### Updating a Dataset
169 * **Example:** Rename and relabel
170 * ```typescript
171 * // Same logical id, changed props: the engine updates it in place.
172 * const dataset = yield* GCP.AIPlatform.Dataset("Samples", {
173 * displayName: "prod rows",
174 * labels: { env: "prod" },
175 * });
176 * ```
177 *
178 * @resource
179 * @category AIPlatform
180 */
181export const Dataset = Resource<Dataset>("GCP.AIPlatform.Dataset");
182
183export class DatasetNotResolved extends Data.TaggedError(
184 "GCP.AIPlatform.DatasetNotResolved",
185)<{
186 name: string;
187}> {}
188
189export class DatasetStillExists extends Data.TaggedError(
190 "GCP.AIPlatform.DatasetStillExists",
191)<{
192 name: string;
193}> {}
194
195const resourceName = (project: string, location: string, datasetId: string) =>
196 `projects/${project}/locations/${location}/datasets/${datasetId}`;
197
198const toId = (id: string, datasetId: string | undefined, existing?: string) =>
199 Effect.gen(function* () {
200 if (datasetId !== undefined) return rfc1035(datasetId);
201 if (existing !== undefined) return existing;
202 return rfc1035(
203 yield* createPhysicalName({
204 id,
205 maxLength: 63,
206 lowercase: true,
207 }),
208 );
209 });
210
211const schemaUriOf = (uri: string | undefined) =>
212 uri && uri.length > 0 ? uri : DEFAULT_TABULAR_METADATA_SCHEMA_URI;
213
214const toAttrs = (
215 dataset: aiplatform.GoogleCloudAiplatformV1Dataset,
216 project: string,
217) => {
218 const name = dataset.name ?? "";
219 const parsed = parseResourceName(name, "datasets");
220 return {
221 name,
222 datasetId: parsed.id,
223 project: parsed.project || project,
224 location: parsed.location,
225 displayName: dataset.displayName,
226 description: dataset.description,
227 labels: userLabels(dataset.labels),
228 metadataSchemaUri: dataset.metadataSchemaUri,
229 metadata: dataset.metadata,
230 dataItemCount: dataset.dataItemCount,
231 metadataArtifact: dataset.metadataArtifact,
232 modelReference: dataset.modelReference,
233 kmsKeyName: dataset.encryptionSpec?.kmsKeyName,
234 createTime: dataset.createTime,
235 updateTime: dataset.updateTime,
236 };
237};
238
239const getByName = (name: string) =>
240 name.length === 0
241 ? Effect.succeed(undefined)
242 : aiplatform
243 .getDatasets({ name })
244 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
245
246const listDatasets = (project: string, region: string) =>
247 Effect.forEach(listLocations(region), (location) =>
248 listDatasetsAt(parentOf(project, location)),
249 ).pipe(Effect.map((pages) => pages.flat()));
250
251const listDatasetsAt = (parent: string) =>
252 aiplatform.listProjectsLocationsDatasets
253 .pages({
254 parent,
255 pageSize: 100,
256 })
257 .pipe(
258 Stream.flatMap((page) => Stream.fromIterable(page.datasets ?? [])),
259 Stream.filter((dataset) => hasAlchemyLabelKeys(dataset.labels)),
260 Stream.runCollect,
261 Effect.map((chunk) => Array.from(chunk)),
262 Effect.catchTag("NotFound", () => Effect.succeed([])),
263 );
264
265const findOwned = (
266 id: string,
267 project: string,
268 region: string,
269 hinted?: string,
270) =>
271 Effect.gen(function* () {
272 if (hinted !== undefined && hinted.length > 0) {
273 const existing = yield* getByName(hinted);
274 if (existing !== undefined) return existing;
275 }
276 const owned = yield* listDatasets(project, region);
277 for (const dataset of owned) {
278 if (yield* hasAlchemyLabels(id, tagRecord(dataset.labels))) {
279 return dataset;
280 }
281 }
282 return undefined as aiplatform.GoogleCloudAiplatformV1Dataset | undefined;
283 });
284
285const waitUntilExists = (name: string) =>
286 getByName(name).pipe(
287 Effect.flatMap((dataset) =>
288 dataset
289 ? Effect.succeed(dataset)
290 : Effect.fail(new DatasetNotResolved({ name })),
291 ),
292 Effect.retry({
293 while: (error) => error._tag === "GCP.AIPlatform.DatasetNotResolved",
294 times: 8,
295 schedule: Schedule.spaced("1 second"),
296 }),
297 );
298
299const waitUntilGone = (name: string) =>
300 getByName(name).pipe(
301 Effect.flatMap((dataset) =>
302 dataset === undefined
303 ? Effect.void
304 : Effect.fail(new DatasetStillExists({ name })),
305 ),
306 Effect.retry({
307 while: (error) => error._tag === "GCP.AIPlatform.DatasetStillExists",
308 times: 10,
309 schedule: Schedule.spaced("2 seconds"),
310 }),
311 );
312
313export const DatasetProvider = () =>
314 Provider.succeed(Dataset, {
315 stables: [
316 "name",
317 "datasetId",
318 "project",
319 "location",
320 "metadataSchemaUri",
321 "createTime",
322 ],
323
324 diff: Effect.fn(function* ({ news, olds, output }) {
325 const env = yield* GcpEnvironment.current;
326 if (!isResolved(news)) return undefined;
327 const previousId = olds?.datasetId ?? output?.datasetId;
328 const nextId = news.datasetId ?? previousId;
329 const previousLocation = normalizeLocation(
330 olds?.location ?? output?.location,
331 env.region,
332 );
333 const nextLocation = normalizeLocation(
334 news.location ?? olds?.location ?? output?.location,
335 env.region,
336 );
337 const previousSchema = schemaUriOf(
338 olds?.metadataSchemaUri ?? output?.metadataSchemaUri,
339 );
340 const nextSchema = schemaUriOf(
341 news.metadataSchemaUri ??
342 olds?.metadataSchemaUri ??
343 output?.metadataSchemaUri,
344 );
345 const previousKey =
346 olds?.encryptionSpec?.kmsKeyName ?? output?.kmsKeyName ?? "";
347 const nextKey = news.encryptionSpec?.kmsKeyName ?? previousKey;
348 const replace =
349 (previousId !== undefined &&
350 nextId !== undefined &&
351 nextId !== previousId) ||
352 previousLocation !== nextLocation ||
353 previousSchema !== nextSchema ||
354 previousKey !== nextKey;
355 if (!replace) return undefined;
356 return {
357 action: "replace" as const,
358 deleteFirst:
359 previousLocation === nextLocation &&
360 previousId !== undefined &&
361 nextId === previousId,
362 };
363 }),
364
365 read: Effect.fn(function* ({ id, olds, output }) {
366 const env = yield* GcpEnvironment.current;
367 const existing = yield* findOwned(
368 id,
369 env.project,
370 env.region,
371 output?.name,
372 );
373 if (existing === undefined) {
374 if (output?.name) return undefined;
375 // Dataset ids are server-assigned numbers; only a caller-supplied
376 // id can be looked up by name.
377 const datasetId = olds?.datasetId;
378 if (datasetId === undefined) return undefined;
379 const location = normalizeLocation(
380 olds?.location ?? output?.location,
381 env.region,
382 );
383 const named = yield* getByName(
384 resourceName(env.project, location, datasetId),
385 );
386 if (named === undefined) return undefined;
387 const attrs = toAttrs(named, env.project);
388 return (yield* hasAlchemyLabels(id, tagRecord(named.labels)))
389 ? attrs
390 : Unowned(attrs);
391 }
392 const attrs = toAttrs(existing, env.project);
393 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
394 ? attrs
395 : Unowned(attrs);
396 }),
397
398 list: () =>
399 Effect.gen(function* () {
400 const env = yield* GcpEnvironment.current;
401 const datasets = yield* listDatasets(env.project, env.region);
402 return datasets.map((dataset) => toAttrs(dataset, env.project));
403 }),
404
405 reconcile: Effect.fn(function* ({ id, news, output }) {
406 const env = yield* GcpEnvironment.current;
407 const location = normalizeLocation(
408 news.location ?? output?.location,
409 env.region,
410 );
411 const datasetId = yield* toId(id, news.datasetId, output?.datasetId);
412 const desiredLabels = {
413 ...toLabels(news.labels),
414 ...(yield* createInternalLabels(id)),
415 };
416 const displayName = news.displayName ?? datasetId;
417 const metadataSchemaUri = schemaUriOf(news.metadataSchemaUri);
418
419 let current = yield* findOwned(id, env.project, env.region, output?.name);
420 if (current === undefined && news.datasetId !== undefined) {
421 current = yield* getByName(
422 resourceName(env.project, location, news.datasetId),
423 );
424 }
425
426 if (current === undefined) {
427 const created = yield* aiplatform
428 .createProjectsLocationsDatasets({
429 parent: parentOf(env.project, location),
430 body: {
431 displayName,
432 description: news.description,
433 labels: desiredLabels,
434 metadataSchemaUri,
435 metadata: news.metadata ?? {},
436 encryptionSpec: news.encryptionSpec,
437 modelReference: news.modelReference,
438 savedQueries: news.savedQueries,
439 },
440 })
441 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
442 if (created !== undefined) {
443 const done = yield* waitForOperation(created);
444 const createdName =
445 resourceNameFromOperation(done) ??
446 (yield* findOwned(id, env.project, env.region))?.name;
447 if (createdName !== undefined && createdName.length > 0) {
448 current = yield* waitUntilExists(createdName);
449 }
450 }
451 if (current === undefined) {
452 current = yield* findOwned(id, env.project, env.region);
453 }
454 }
455
456 if (current === undefined) {
457 return yield* new DatasetNotResolved({
458 name: output?.name ?? resourceName(env.project, location, datasetId),
459 });
460 }
461
462 const name = current.name ?? "";
463 const observedLabels = tagRecord(current.labels);
464 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
465 const labelsChanged = upsert.length > 0 || removed.length > 0;
466 const displayNameChanged = (current.displayName ?? "") !== displayName;
467 const descriptionChanged =
468 (current.description ?? "") !== (news.description ?? "");
469
470 if (labelsChanged || displayNameChanged || descriptionChanged) {
471 current = yield* aiplatform.patchProjectsLocationsDatasets({
472 name,
473 updateMask: [
474 labelsChanged ? "labels" : undefined,
475 displayNameChanged ? "display_name" : undefined,
476 descriptionChanged ? "description" : undefined,
477 ]
478 .filter((field): field is string => field !== undefined)
479 .join(","),
480 body: {
481 name,
482 labels: desiredLabels,
483 displayName,
484 description: news.description,
485 },
486 });
487 }
488
489 return toAttrs(current, env.project);
490 }),
491
492 delete: Effect.fn(function* ({ output }) {
493 const existing = yield* getByName(output.name);
494 if (existing === undefined) return;
495 const operation = yield* aiplatform
496 .deleteDatasets({ name: output.name })
497 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
498 if (operation !== undefined) {
499 yield* waitForOperation(operation, { notFoundOk: true });
500 }
501 yield* waitUntilGone(output.name);
502 }),
503 });