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 { listLocations } from "./names.ts";
7import { Unowned } from "../../AdoptPolicy.ts";
8import { isResolved } from "../../Diff.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 { waitForOperation } from "./operations.ts";
21import {
22 hasAlchemyLabelMap,
23 normalizeLocation,
24 parseResourceName,
25 specifiedEquals,
26 toPhysicalSnake,
27 userLabels,
28} from "./helpers.ts";
29
30const MAX_NAME_LENGTH = 60;
31
32export type FeaturestoreOnlineServingScaling = {
33 minNodeCount?: number;
34 maxNodeCount?: number;
35 cpuUtilizationTarget?: number;
36};
37
38export type FeaturestoreOnlineServingConfig = {
39 /** Fixed node count. `0` disables online serving. */
40 fixedNodeCount?: number;
41 /** Autoscaling. Mutually exclusive with `fixedNodeCount`. */
42 scaling?: FeaturestoreOnlineServingScaling;
43};
44
45export type FeaturestoreProps = {
46 /**
47 * Featurestore id. Valid characters `[a-z0-9_]`, must start with a
48 * letter, max 60. Immutable — changing it replaces the store.
49 */
50 featurestoreId?: string;
51 /**
52 * Region. Immutable — changing it replaces the store.
53 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
54 */
55 location?: string;
56 /**
57 * User labels. Alchemy ownership labels are merged in automatically.
58 */
59 labels?: Record<string, string>;
60 /**
61 * Online serving resources. Omit for an offline-only store.
62 */
63 onlineServingConfig?: FeaturestoreOnlineServingConfig;
64 /**
65 * TTL in days for online storage. Defaults to 4000.
66 */
67 onlineStorageTtlDays?: number;
68 /**
69 * Customer-managed encryption. Immutable.
70 */
71 encryptionSpec?: { kmsKeyName?: string };
72};
73
74export type Featurestore = Resource<
75 "GCP.AIPlatform.Featurestore",
76 FeaturestoreProps,
77 {
78 /** Full resource name. */
79 name: string;
80 /** Featurestore id. */
81 featurestoreId: string;
82 /** Project id. */
83 project: string;
84 /** Location id. */
85 location: string;
86 /** User labels (Alchemy ownership labels stripped). */
87 labels: Record<string, string>;
88 /** Online serving config. */
89 onlineServingConfig: FeaturestoreOnlineServingConfig | undefined;
90 /** Online storage TTL in days. */
91 onlineStorageTtlDays: number | undefined;
92 /** Store state (`STABLE`, `UPDATING`). */
93 state: string | undefined;
94 /** Customer-managed encryption key. */
95 encryptionSpec: { kmsKeyName?: string } | undefined;
96 /** RFC3339 creation timestamp. */
97 createTime: string | undefined;
98 /** RFC3339 last-update timestamp. */
99 updateTime: string | undefined;
100 /** Server etag. */
101 etag: string | undefined;
102 },
103 never,
104 Providers
105>;
106
107/**
108 * A Vertex AI Feature Store (legacy) — a container for Entity Types and
109 * Features.
110 *
111 * Id, location, and encryption are identity. Labels, online serving, and
112 * online TTL update in place. Omit `onlineServingConfig` for an
113 * offline-only store.
114 *
115 * ### Creating a Featurestore
116 * **Example:** Offline-only store
117 * ```typescript
118 * const store = yield* GCP.AIPlatform.Featurestore("Features", {
119 * labels: { env: "prod" },
120 * });
121 * ```
122 *
123 * **Example:** Online serving with a fixed node
124 * ```typescript
125 * const store = yield* GCP.AIPlatform.Featurestore("Features", {
126 * onlineServingConfig: { fixedNodeCount: 1 },
127 * });
128 * ```
129 *
130 * @resource
131 * @category AIPlatform
132 */
133export const Featurestore = Resource<Featurestore>(
134 "GCP.AIPlatform.Featurestore",
135);
136
137export class FeaturestoreNotResolved extends Data.TaggedError(
138 "GCP.AIPlatform.FeaturestoreNotResolved",
139)<{
140 name: string;
141}> {}
142
143export class FeaturestoreStillExists extends Data.TaggedError(
144 "GCP.AIPlatform.FeaturestoreStillExists",
145)<{
146 name: string;
147}> {}
148
149const resourceName = (project: string, location: string, storeId: string) =>
150 `projects/${project}/locations/${location}/featurestores/${storeId}`;
151
152const toServing = (
153 config:
154 | aiplatform.GoogleCloudAiplatformV1FeaturestoreOnlineServingConfig
155 | FeaturestoreOnlineServingConfig
156 | undefined,
157): FeaturestoreOnlineServingConfig | undefined => {
158 if (config === undefined) return undefined;
159 return {
160 fixedNodeCount: config.fixedNodeCount,
161 scaling: config.scaling
162 ? {
163 minNodeCount: config.scaling.minNodeCount,
164 maxNodeCount: config.scaling.maxNodeCount,
165 cpuUtilizationTarget: config.scaling.cpuUtilizationTarget,
166 }
167 : undefined,
168 };
169};
170
171const toAttrs = (
172 store: aiplatform.GoogleCloudAiplatformV1Featurestore,
173 project: string,
174) => {
175 const name = store.name ?? "";
176 const parsed = parseResourceName(name, "featurestores");
177 return {
178 name,
179 featurestoreId: parsed.id,
180 project: parsed.project || project,
181 location: parsed.location,
182 labels: userLabels(store.labels),
183 onlineServingConfig: toServing(store.onlineServingConfig),
184 onlineStorageTtlDays: store.onlineStorageTtlDays,
185 state: store.state,
186 encryptionSpec: store.encryptionSpec?.kmsKeyName
187 ? { kmsKeyName: store.encryptionSpec.kmsKeyName }
188 : undefined,
189 createTime: store.createTime,
190 updateTime: store.updateTime,
191 etag: store.etag,
192 };
193};
194
195const getByName = (name: string) =>
196 aiplatform
197 .getProjectsLocationsFeaturestores({ name })
198 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
199
200const waitUntilExists = (name: string) =>
201 getByName(name).pipe(
202 Effect.flatMap((store) =>
203 store
204 ? Effect.succeed(store)
205 : Effect.fail(new FeaturestoreNotResolved({ name })),
206 ),
207 Effect.retry({
208 while: (error) => error._tag === "GCP.AIPlatform.FeaturestoreNotResolved",
209 times: 8,
210 schedule: Schedule.spaced("1 second"),
211 }),
212 );
213
214const waitUntilGone = (name: string) =>
215 getByName(name).pipe(
216 Effect.flatMap((store) =>
217 store === undefined
218 ? Effect.void
219 : Effect.fail(new FeaturestoreStillExists({ name })),
220 ),
221 Effect.retry({
222 while: (error) => error._tag === "GCP.AIPlatform.FeaturestoreStillExists",
223 times: 10,
224 schedule: Schedule.spaced("2 seconds"),
225 }),
226 );
227
228const isReady = (state: string | undefined) => {
229 const value = (state ?? "").toUpperCase();
230 return value === "STABLE" || value === "";
231};
232
233const waitUntilReady = (name: string) =>
234 getByName(name).pipe(
235 Effect.filterOrFail(
236 (store): store is aiplatform.GoogleCloudAiplatformV1Featurestore =>
237 store !== undefined,
238 () => new FeaturestoreNotResolved({ name }),
239 ),
240 Effect.filterOrFail(
241 (store) => isReady(store.state),
242 () => new FeaturestoreNotResolved({ name }),
243 ),
244 Effect.retry({
245 while: (error) => error._tag === "GCP.AIPlatform.FeaturestoreNotResolved",
246 times: 10,
247 schedule: Schedule.spaced("8 seconds"),
248 }),
249 );
250
251export const FeaturestoreProvider = () =>
252 Provider.succeed(Featurestore, {
253 stables: ["name", "featurestoreId", "project", "location", "createTime"],
254
255 diff: Effect.fn(function* ({ news, olds, output }) {
256 const env = yield* GcpEnvironment.current;
257 if (!isResolved(news)) return undefined;
258 const previousId = olds?.featurestoreId ?? output?.featurestoreId;
259 const nextId = news.featurestoreId ?? previousId;
260 const previousLocation = normalizeLocation(
261 olds?.location ?? output?.location,
262 env.region,
263 );
264 const nextLocation = normalizeLocation(
265 news.location ?? output?.location,
266 env.region,
267 );
268 const previousKey =
269 olds?.encryptionSpec?.kmsKeyName ??
270 output?.encryptionSpec?.kmsKeyName ??
271 "";
272 const nextKey = news.encryptionSpec?.kmsKeyName ?? previousKey;
273 const replace =
274 (previousId !== undefined &&
275 nextId !== undefined &&
276 nextId !== previousId) ||
277 previousLocation !== nextLocation ||
278 previousKey !== nextKey;
279 if (!replace) return undefined;
280 return {
281 action: "replace" as const,
282 deleteFirst:
283 previousLocation === nextLocation &&
284 previousId !== undefined &&
285 nextId === previousId,
286 };
287 }),
288
289 read: Effect.fn(function* ({ id, olds, output }) {
290 const env = yield* GcpEnvironment.current;
291 const storeId = yield* toPhysicalSnake(
292 id,
293 olds?.featurestoreId,
294 output?.featurestoreId,
295 MAX_NAME_LENGTH,
296 );
297 const location = normalizeLocation(
298 olds?.location ?? output?.location,
299 env.region,
300 );
301 const name = output?.name ?? resourceName(env.project, location, storeId);
302 const existing = yield* getByName(name);
303 if (existing === undefined) return undefined;
304 const attrs = toAttrs(existing, env.project);
305 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
306 ? attrs
307 : Unowned(attrs);
308 }),
309
310 list: () =>
311 Effect.gen(function* () {
312 const env = yield* GcpEnvironment.current;
313 return yield* Stream.fromIterable(listLocations(env.region))
314 .pipe(
315 Stream.flatMap((location) =>
316 aiplatform.listProjectsLocationsFeaturestores.pages({
317 parent: `projects/${env.project}/locations/${location}`,
318 pageSize: 100,
319 }),
320 ),
321 )
322 .pipe(
323 Stream.flatMap((page) =>
324 Stream.fromIterable(page.featurestores ?? []),
325 ),
326 Stream.filter((store) => hasAlchemyLabelMap(store.labels)),
327 Stream.map((store) => toAttrs(store, env.project)),
328 Stream.runCollect,
329 Effect.map((chunk) => Array.from(chunk)),
330 Effect.catchTag("NotFound", () => Effect.succeed([])),
331 );
332 }),
333
334 reconcile: Effect.fn(function* ({ id, news, output }) {
335 const env = yield* GcpEnvironment.current;
336 const storeId = yield* toPhysicalSnake(
337 id,
338 news.featurestoreId,
339 output?.featurestoreId,
340 MAX_NAME_LENGTH,
341 );
342 const location = normalizeLocation(
343 news.location ?? output?.location,
344 env.region,
345 );
346 const name = resourceName(env.project, location, storeId);
347 const desiredLabels = {
348 ...toLabels(news.labels),
349 ...(yield* createInternalLabels(id)),
350 };
351
352 let current = yield* getByName(output?.name ?? name);
353
354 if (current === undefined) {
355 const created = yield* aiplatform
356 .createProjectsLocationsFeaturestores({
357 parent: `projects/${env.project}/locations/${location}`,
358 featurestoreId: storeId,
359 body: {
360 labels: desiredLabels,
361 onlineServingConfig: news.onlineServingConfig,
362 onlineStorageTtlDays: news.onlineStorageTtlDays,
363 encryptionSpec: news.encryptionSpec,
364 },
365 })
366 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
367 if (created !== undefined) {
368 yield* waitForOperation(created);
369 }
370 current = yield* waitUntilExists(name);
371 }
372
373 if (current === undefined) {
374 return yield* new FeaturestoreNotResolved({ name });
375 }
376
377 if (!isReady(current.state)) {
378 current = yield* waitUntilReady(name);
379 }
380
381 const observedLabels = tagRecord(current.labels);
382 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
383 const labelsChanged = upsert.length > 0 || removed.length > 0;
384 const servingChanged =
385 news.onlineServingConfig !== undefined &&
386 !specifiedEquals(
387 news.onlineServingConfig,
388 toServing(current.onlineServingConfig),
389 );
390 const ttlChanged =
391 news.onlineStorageTtlDays !== undefined &&
392 (current.onlineStorageTtlDays ?? 0) !== news.onlineStorageTtlDays;
393
394 if (labelsChanged || servingChanged || ttlChanged) {
395 const updateMask = [
396 labelsChanged ? "labels" : undefined,
397 servingChanged && news.onlineServingConfig?.scaling
398 ? "online_serving_config.scaling"
399 : servingChanged
400 ? "online_serving_config.fixed_node_count"
401 : undefined,
402 ttlChanged ? "online_storage_ttl_days" : undefined,
403 ].filter((field): field is string => field !== undefined);
404
405 const patched = yield* aiplatform
406 .patchProjectsLocationsFeaturestores({
407 name,
408 updateMask: updateMask.join(","),
409 body: {
410 name,
411 labels: desiredLabels,
412 onlineServingConfig: news.onlineServingConfig,
413 onlineStorageTtlDays: news.onlineStorageTtlDays,
414 },
415 })
416 .pipe(
417 Effect.retry({
418 while: (error) => error._tag === "Conflict",
419 times: 8,
420 schedule: Schedule.spaced("5 seconds"),
421 }),
422 );
423 yield* waitForOperation(patched);
424 current = yield* waitUntilReady(name);
425 }
426
427 return toAttrs(current, env.project);
428 }),
429
430 delete: Effect.fn(function* ({ output }) {
431 const operation = yield* aiplatform
432 .deleteProjectsLocationsFeaturestores({
433 name: output.name,
434 force: true,
435 })
436 .pipe(
437 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
438 Effect.retry({
439 while: (error) => error._tag === "Conflict",
440 times: 8,
441 schedule: Schedule.spaced("5 seconds"),
442 }),
443 );
444 if (operation !== undefined) {
445 yield* waitForOperation(operation, { notFoundOk: true });
446 }
447 yield* waitUntilGone(output.name);
448 }),
449 });