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 FeatureOnlineStoreBigtableAutoScaling = {
33 /** Minimum node count (>= 1). */
34 minNodeCount?: number;
35 /** Maximum node count (<= 10 * minNodeCount). */
36 maxNodeCount?: number;
37 /** Target CPU percent (10–80). Defaults to 50. */
38 cpuUtilizationTarget?: number;
39};
40
41export type FeatureOnlineStoreBigtable = {
42 /** Autoscaling applied to the backing Bigtable instance. */
43 autoScaling?: FeatureOnlineStoreBigtableAutoScaling;
44 /** Allow direct Bigtable access. */
45 enableDirectBigtableAccess?: boolean;
46 /** Zone for the primary Bigtable cluster (e.g. `us-central1-a`). Immutable. */
47 zone?: string;
48};
49
50export type FeatureOnlineStoreDedicatedServingEndpoint = {
51 /** Private Service Connect config (Optimized storage only). */
52 privateServiceConnectConfig?: {
53 enablePrivateServiceConnect?: boolean;
54 projectAllowlist?: string[];
55 };
56};
57
58export type FeatureOnlineStoreProps = {
59 /**
60 * Feature Online Store id. Valid characters `[a-z0-9_]`, must start
61 * with a letter, max 60. Immutable — changing it replaces the store.
62 */
63 featureOnlineStoreId?: string;
64 /**
65 * Region. Immutable — changing it replaces the store.
66 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
67 */
68 location?: string;
69 /**
70 * User labels. Alchemy ownership labels are merged in automatically.
71 */
72 labels?: Record<string, string>;
73 /**
74 * Bigtable storage. Mutually exclusive with `optimized`. Immutable
75 * storage type — switching to Optimized replaces the store.
76 * @default { autoScaling: { minNodeCount: 1, maxNodeCount: 1 } } when `optimized` is unset
77 */
78 bigtable?: FeatureOnlineStoreBigtable;
79 /**
80 * Optimized storage. When true or `{}`, Vertex provisions the
81 * Optimized serving backend. Immutable storage type.
82 */
83 optimized?: boolean | Record<string, never>;
84 /**
85 * Dedicated serving endpoint (Optimized storage).
86 */
87 dedicatedServingEndpoint?: FeatureOnlineStoreDedicatedServingEndpoint;
88 /**
89 * Customer-managed encryption. Immutable.
90 */
91 encryptionSpec?: { kmsKeyName?: string };
92};
93
94export type FeatureOnlineStore = Resource<
95 "GCP.AIPlatform.FeatureOnlineStore",
96 FeatureOnlineStoreProps,
97 {
98 /** Full resource name. */
99 name: string;
100 /** Store id. */
101 featureOnlineStoreId: string;
102 /** Project id. */
103 project: string;
104 /** Location id. */
105 location: string;
106 /** User labels (Alchemy ownership labels stripped). */
107 labels: Record<string, string>;
108 /** Bigtable config, if used. */
109 bigtable: FeatureOnlineStoreBigtable | undefined;
110 /** Whether Optimized storage is enabled. */
111 optimized: boolean;
112 /** Store state (`STABLE`, `UPDATING`). */
113 state: string | undefined;
114 /** Dedicated serving endpoint. */
115 dedicatedServingEndpoint:
116 | {
117 publicEndpointDomainName: string | undefined;
118 serviceAttachment: string | undefined;
119 privateServiceConnectConfig:
120 | {
121 enablePrivateServiceConnect?: boolean;
122 projectAllowlist?: string[];
123 }
124 | undefined;
125 }
126 | undefined;
127 /** Customer-managed encryption key. */
128 encryptionSpec: { kmsKeyName?: string } | undefined;
129 /** RFC3339 creation timestamp. */
130 createTime: string | undefined;
131 /** RFC3339 last-update timestamp. */
132 updateTime: string | undefined;
133 /** Server etag. */
134 etag: string | undefined;
135 },
136 never,
137 Providers
138>;
139
140/**
141 * A Vertex AI Feature Online Store for low-latency feature and embedding
142 * serving.
143 *
144 * Id, location, storage type (Bigtable vs Optimized), Bigtable zone, and
145 * encryption are identity. Labels and Bigtable autoscaling update in
146 * place. Defaults to a one-node Bigtable backend when neither backend is
147 * set (Vertex AI no longer creates new Optimized stores for most projects).
148 *
149 * ### Creating a Feature Online Store
150 * **Example:** Default one-node Bigtable store
151 * ```typescript
152 * const store = yield* GCP.AIPlatform.FeatureOnlineStore("Serving", {
153 * labels: { env: "prod" },
154 * });
155 * ```
156 *
157 * **Example:** Bigtable-backed store
158 * ```typescript
159 * const store = yield* GCP.AIPlatform.FeatureOnlineStore("Serving", {
160 * bigtable: {
161 * autoScaling: { minNodeCount: 1, maxNodeCount: 2 },
162 * },
163 * });
164 * ```
165 *
166 * @resource
167 * @category AIPlatform
168 */
169export const FeatureOnlineStore = Resource<FeatureOnlineStore>(
170 "GCP.AIPlatform.FeatureOnlineStore",
171);
172
173export class FeatureOnlineStoreNotResolved extends Data.TaggedError(
174 "GCP.AIPlatform.FeatureOnlineStoreNotResolved",
175)<{
176 name: string;
177}> {}
178
179export class FeatureOnlineStoreStillExists extends Data.TaggedError(
180 "GCP.AIPlatform.FeatureOnlineStoreStillExists",
181)<{
182 name: string;
183}> {}
184
185const resourceName = (project: string, location: string, storeId: string) =>
186 `projects/${project}/locations/${location}/featureOnlineStores/${storeId}`;
187
188// Vertex no longer creates Optimized stores unless asked explicitly, so
189// the default backend is the smallest Bigtable one.
190const wantsOptimized = (news: FeatureOnlineStoreProps) =>
191 news.bigtable === undefined &&
192 news.optimized !== undefined &&
193 news.optimized !== false;
194
195const DEFAULT_BIGTABLE: FeatureOnlineStoreBigtable = {
196 autoScaling: { minNodeCount: 1, maxNodeCount: 1 },
197};
198
199const bigtableOf = (news: FeatureOnlineStoreProps) =>
200 news.bigtable ?? DEFAULT_BIGTABLE;
201
202const toBigtable = (
203 config:
204 | aiplatform.GoogleCloudAiplatformV1FeatureOnlineStoreBigtable
205 | FeatureOnlineStoreBigtable
206 | undefined,
207): FeatureOnlineStoreBigtable | undefined => {
208 if (config === undefined) return undefined;
209 return {
210 autoScaling: config.autoScaling
211 ? {
212 minNodeCount: config.autoScaling.minNodeCount,
213 maxNodeCount: config.autoScaling.maxNodeCount,
214 cpuUtilizationTarget: config.autoScaling.cpuUtilizationTarget,
215 }
216 : undefined,
217 enableDirectBigtableAccess: config.enableDirectBigtableAccess,
218 zone: config.zone,
219 };
220};
221
222const toAttrs = (
223 store: aiplatform.GoogleCloudAiplatformV1FeatureOnlineStore,
224 project: string,
225) => {
226 const name = store.name ?? "";
227 const parsed = parseResourceName(name, "featureOnlineStores");
228 const endpoint = store.dedicatedServingEndpoint;
229 return {
230 name,
231 featureOnlineStoreId: parsed.id,
232 project: parsed.project || project,
233 location: parsed.location,
234 labels: userLabels(store.labels),
235 bigtable: toBigtable(store.bigtable),
236 optimized: store.optimized !== undefined,
237 state: store.state,
238 dedicatedServingEndpoint:
239 endpoint === undefined
240 ? undefined
241 : {
242 publicEndpointDomainName: endpoint.publicEndpointDomainName,
243 serviceAttachment: endpoint.serviceAttachment,
244 privateServiceConnectConfig: endpoint.privateServiceConnectConfig
245 ? {
246 enablePrivateServiceConnect:
247 endpoint.privateServiceConnectConfig
248 .enablePrivateServiceConnect === true,
249 projectAllowlist:
250 endpoint.privateServiceConnectConfig.projectAllowlist,
251 }
252 : undefined,
253 },
254 encryptionSpec: store.encryptionSpec?.kmsKeyName
255 ? { kmsKeyName: store.encryptionSpec.kmsKeyName }
256 : undefined,
257 createTime: store.createTime,
258 updateTime: store.updateTime,
259 etag: store.etag,
260 };
261};
262
263const getByName = (name: string) =>
264 aiplatform
265 .getProjectsLocationsFeatureOnlineStores({ name })
266 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
267
268const waitUntilExists = (name: string) =>
269 getByName(name).pipe(
270 Effect.flatMap((store) =>
271 store
272 ? Effect.succeed(store)
273 : Effect.fail(new FeatureOnlineStoreNotResolved({ name })),
274 ),
275 Effect.retry({
276 while: (error) =>
277 error._tag === "GCP.AIPlatform.FeatureOnlineStoreNotResolved",
278 times: 8,
279 schedule: Schedule.spaced("1 second"),
280 }),
281 );
282
283const waitUntilGone = (name: string) =>
284 getByName(name).pipe(
285 Effect.flatMap((store) =>
286 store === undefined
287 ? Effect.void
288 : Effect.fail(new FeatureOnlineStoreStillExists({ name })),
289 ),
290 Effect.retry({
291 while: (error) =>
292 error._tag === "GCP.AIPlatform.FeatureOnlineStoreStillExists",
293 times: 10,
294 schedule: Schedule.spaced("2 seconds"),
295 }),
296 );
297
298const isReady = (state: string | undefined) => {
299 const value = (state ?? "").toUpperCase();
300 return value === "STABLE" || value === "";
301};
302
303const waitUntilReady = (name: string) =>
304 getByName(name).pipe(
305 Effect.filterOrFail(
306 (store): store is aiplatform.GoogleCloudAiplatformV1FeatureOnlineStore =>
307 store !== undefined,
308 () => new FeatureOnlineStoreNotResolved({ name }),
309 ),
310 Effect.filterOrFail(
311 (store) => isReady(store.state),
312 () => new FeatureOnlineStoreNotResolved({ name }),
313 ),
314 Effect.retry({
315 while: (error) =>
316 error._tag === "GCP.AIPlatform.FeatureOnlineStoreNotResolved",
317 times: 10,
318 schedule: Schedule.spaced("8 seconds"),
319 }),
320 );
321
322export const FeatureOnlineStoreProvider = () =>
323 Provider.succeed(FeatureOnlineStore, {
324 stables: [
325 "name",
326 "featureOnlineStoreId",
327 "project",
328 "location",
329 "createTime",
330 ],
331
332 diff: Effect.fn(function* ({ news, olds, output }) {
333 const env = yield* GcpEnvironment.current;
334 if (!isResolved(news)) return undefined;
335 const previousId =
336 olds?.featureOnlineStoreId ?? output?.featureOnlineStoreId;
337 const nextId = news.featureOnlineStoreId ?? previousId;
338 const previousLocation = normalizeLocation(
339 olds?.location ?? output?.location,
340 env.region,
341 );
342 const nextLocation = normalizeLocation(
343 news.location ?? output?.location,
344 env.region,
345 );
346 const previousOptimized =
347 olds?.optimized !== undefined
348 ? olds.optimized !== false
349 : (output?.optimized ?? false);
350 const nextOptimized =
351 news.bigtable !== undefined
352 ? false
353 : news.optimized !== undefined
354 ? news.optimized !== false
355 : previousOptimized;
356 const previousZone = olds?.bigtable?.zone ?? output?.bigtable?.zone ?? "";
357 const nextZone = news.bigtable?.zone ?? previousZone;
358 const previousKey =
359 olds?.encryptionSpec?.kmsKeyName ??
360 output?.encryptionSpec?.kmsKeyName ??
361 "";
362 const nextKey = news.encryptionSpec?.kmsKeyName ?? previousKey;
363 const replace =
364 (previousId !== undefined &&
365 nextId !== undefined &&
366 nextId !== previousId) ||
367 previousLocation !== nextLocation ||
368 previousOptimized !== nextOptimized ||
369 previousZone !== nextZone ||
370 previousKey !== nextKey;
371 if (!replace) return undefined;
372 return {
373 action: "replace" as const,
374 deleteFirst:
375 previousLocation === nextLocation &&
376 previousId !== undefined &&
377 nextId === previousId,
378 };
379 }),
380
381 read: Effect.fn(function* ({ id, olds, output }) {
382 const env = yield* GcpEnvironment.current;
383 const storeId = yield* toPhysicalSnake(
384 id,
385 olds?.featureOnlineStoreId,
386 output?.featureOnlineStoreId,
387 MAX_NAME_LENGTH,
388 );
389 const location = normalizeLocation(
390 olds?.location ?? output?.location,
391 env.region,
392 );
393 const name = output?.name ?? resourceName(env.project, location, storeId);
394 const existing = yield* getByName(name);
395 if (existing === undefined) return undefined;
396 const attrs = toAttrs(existing, env.project);
397 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
398 ? attrs
399 : Unowned(attrs);
400 }),
401
402 list: () =>
403 Effect.gen(function* () {
404 const env = yield* GcpEnvironment.current;
405 return yield* Stream.fromIterable(listLocations(env.region))
406 .pipe(
407 Stream.flatMap((location) =>
408 aiplatform.listProjectsLocationsFeatureOnlineStores.pages({
409 parent: `projects/${env.project}/locations/${location}`,
410 pageSize: 100,
411 }),
412 ),
413 )
414 .pipe(
415 Stream.flatMap((page) =>
416 Stream.fromIterable(page.featureOnlineStores ?? []),
417 ),
418 Stream.filter((store) => hasAlchemyLabelMap(store.labels)),
419 Stream.map((store) => toAttrs(store, env.project)),
420 Stream.runCollect,
421 Effect.map((chunk) => Array.from(chunk)),
422 Effect.catchTag("NotFound", () => Effect.succeed([])),
423 );
424 }),
425
426 reconcile: Effect.fn(function* ({ id, news, output }) {
427 const env = yield* GcpEnvironment.current;
428 const storeId = yield* toPhysicalSnake(
429 id,
430 news.featureOnlineStoreId,
431 output?.featureOnlineStoreId,
432 MAX_NAME_LENGTH,
433 );
434 const location = normalizeLocation(
435 news.location ?? output?.location,
436 env.region,
437 );
438 const name = resourceName(env.project, location, storeId);
439 const desiredLabels = {
440 ...toLabels(news.labels),
441 ...(yield* createInternalLabels(id)),
442 };
443 const optimized = wantsOptimized(news);
444
445 let current = yield* getByName(output?.name ?? name);
446
447 if (current === undefined) {
448 const created = yield* aiplatform
449 .createProjectsLocationsFeatureOnlineStores({
450 parent: `projects/${env.project}/locations/${location}`,
451 featureOnlineStoreId: storeId,
452 body: {
453 labels: desiredLabels,
454 bigtable: optimized ? undefined : bigtableOf(news),
455 optimized: optimized ? {} : undefined,
456 dedicatedServingEndpoint: news.dedicatedServingEndpoint,
457 encryptionSpec: news.encryptionSpec,
458 },
459 })
460 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
461 if (created !== undefined) {
462 yield* waitForOperation(created);
463 }
464 current = yield* waitUntilExists(name);
465 }
466
467 if (current === undefined) {
468 return yield* new FeatureOnlineStoreNotResolved({ name });
469 }
470
471 if (!isReady(current.state)) {
472 current = yield* waitUntilReady(name);
473 }
474
475 const observedLabels = tagRecord(current.labels);
476 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
477 const labelsChanged = upsert.length > 0 || removed.length > 0;
478 const scalingChanged =
479 news.bigtable?.autoScaling !== undefined &&
480 !specifiedEquals(
481 news.bigtable.autoScaling,
482 current.bigtable?.autoScaling,
483 );
484 const directChanged =
485 news.bigtable?.enableDirectBigtableAccess !== undefined &&
486 (current.bigtable?.enableDirectBigtableAccess === true) !==
487 (news.bigtable.enableDirectBigtableAccess === true);
488
489 if (labelsChanged || scalingChanged || directChanged) {
490 const updateMask = [
491 labelsChanged ? "labels" : undefined,
492 scalingChanged ? "bigtable.auto_scaling" : undefined,
493 directChanged ? "bigtable" : undefined,
494 ].filter((field): field is string => field !== undefined);
495
496 const patched = yield* aiplatform
497 .patchProjectsLocationsFeatureOnlineStores({
498 name,
499 updateMask: updateMask.join(","),
500 body: {
501 name,
502 labels: desiredLabels,
503 bigtable: news.bigtable,
504 },
505 })
506 .pipe(
507 Effect.retry({
508 while: (error) => error._tag === "Conflict",
509 times: 8,
510 schedule: Schedule.spaced("5 seconds"),
511 }),
512 );
513 yield* waitForOperation(patched);
514 current = yield* waitUntilReady(name);
515 }
516
517 return toAttrs(current, env.project);
518 }),
519
520 delete: Effect.fn(function* ({ output }) {
521 const operation = yield* aiplatform
522 .deleteProjectsLocationsFeatureOnlineStores({
523 name: output.name,
524 force: true,
525 })
526 .pipe(
527 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
528 Effect.retry({
529 while: (error) => error._tag === "Conflict",
530 times: 8,
531 schedule: Schedule.spaced("5 seconds"),
532 }),
533 );
534 if (operation !== undefined) {
535 yield* waitForOperation(operation, { notFoundOk: true });
536 }
537 yield* waitUntilGone(output.name);
538 }),
539 });