1import * as sagemaker from "@distilled.cloud/aws/sagemaker";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as EffectStream 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 {
12 createInternalTags,
13 diffTags,
14 hasAlchemyTags,
15 type Tags,
16} from "../../Tags.ts";
17import type { Providers } from "../Providers.ts";
18
19export type FeatureGroupStatus = sagemaker.FeatureGroupStatus;
20
21export interface FeatureGroupProps {
22 /**
23 * Name of the feature group. Maximum 64 characters, alphanumeric and
24 * hyphens.
25 * @default ${app}-${stage}-${id}
26 */
27 featureGroupName?: string;
28 /**
29 * The name of the feature whose value uniquely identifies a record. Must
30 * be one of the `featureDefinitions`.
31 */
32 recordIdentifierFeatureName: string;
33 /**
34 * The name of the feature that stores each record's event time (an ISO
35 * 8601 string or unix seconds). Must be one of the `featureDefinitions`.
36 */
37 eventTimeFeatureName: string;
38 /**
39 * The schema: every feature's name and type (`String`, `Integral`,
40 * `Fractional`).
41 */
42 featureDefinitions: sagemaker.FeatureDefinition[];
43 /**
44 * Enable the low-latency online store (required for
45 * `GetRecord`/`PutRecord` runtime bindings).
46 */
47 onlineStoreConfig?: sagemaker.OnlineStoreConfig;
48 /**
49 * Replicate writes to an S3-backed offline store (for training). Requires
50 * `roleArn`.
51 */
52 offlineStoreConfig?: sagemaker.OfflineStoreConfig;
53 /**
54 * Provisioned or on-demand read/write throughput for the online store.
55 */
56 throughputConfig?: sagemaker.ThroughputConfig;
57 /**
58 * ARN of the IAM role SageMaker assumes to write the offline store to S3.
59 * Only required when `offlineStoreConfig` is set.
60 */
61 roleArn?: string;
62 /**
63 * A description of the feature group.
64 */
65 description?: string;
66 /**
67 * Tags to associate with the feature group. Merged with internal Alchemy
68 * tags.
69 */
70 tags?: Record<string, string>;
71}
72
73export interface FeatureGroup extends Resource<
74 "AWS.SageMaker.FeatureGroup",
75 FeatureGroupProps,
76 {
77 /**
78 * The feature group's name.
79 */
80 featureGroupName: string;
81 /**
82 * ARN of the feature group.
83 */
84 featureGroupArn: string;
85 /**
86 * The record-identifier feature name.
87 */
88 recordIdentifierFeatureName: string;
89 /**
90 * The event-time feature name.
91 */
92 eventTimeFeatureName: string;
93 },
94 never,
95 Providers
96> {}
97
98/**
99 * An Amazon SageMaker Feature Store FeatureGroup — a typed, versioned table
100 * of ML features with an optional low-latency online store (for inference
101 * lookups) and an S3-backed offline store (for training).
102 *
103 * With the online store enabled, functions read and write records at runtime
104 * via the `AWS.SageMaker.GetRecord` / `AWS.SageMaker.PutRecord` bindings.
105 * ### Creating Feature Groups
106 * **Example:** Online-Store Feature Group
107 * ```typescript
108 * import * as AWS from "alchemy/AWS";
109 *
110 * const features = yield* AWS.SageMaker.FeatureGroup("UserFeatures", {
111 * recordIdentifierFeatureName: "user_id",
112 * eventTimeFeatureName: "event_time",
113 * featureDefinitions: [
114 * { FeatureName: "user_id", FeatureType: "String" },
115 * { FeatureName: "event_time", FeatureType: "String" },
116 * { FeatureName: "clicks", FeatureType: "Integral" },
117 * ],
118 * onlineStoreConfig: { EnableOnlineStore: true },
119 * });
120 * ```
121 *
122 * ### Runtime Access
123 * **Example:** Read and write records from a Lambda function
124 * ```typescript
125 * // init
126 * const putRecord = yield* AWS.SageMaker.PutRecord(features);
127 * const getRecord = yield* AWS.SageMaker.GetRecord(features);
128 *
129 * // runtime
130 * yield* putRecord({
131 * Record: [
132 * { FeatureName: "user_id", ValueAsString: "user-123" },
133 * { FeatureName: "event_time", ValueAsString: new Date().toISOString() },
134 * { FeatureName: "clicks", ValueAsString: "42" },
135 * ],
136 * });
137 * const { Record } = yield* getRecord({
138 * RecordIdentifierValueAsString: "user-123",
139 * });
140 * ```
141 *
142 * @resource
143 */
144export const FeatureGroup = Resource<FeatureGroup>(
145 "AWS.SageMaker.FeatureGroup",
146);
147
148const createFeatureGroupName = (
149 id: string,
150 props: { featureGroupName?: string | undefined },
151) =>
152 props.featureGroupName
153 ? Effect.succeed(props.featureGroupName)
154 : createPhysicalName({ id, maxLength: 64 });
155
156const fetchFeatureGroupTags = Effect.fn(function* (arn: string) {
157 const response = yield* sagemaker
158 .listTags({ ResourceArn: arn })
159 .pipe(
160 Effect.catchTag("AccessDeniedException", () => Effect.succeed(undefined)),
161 );
162 return Object.fromEntries(
163 (response?.Tags ?? []).flatMap((tag) =>
164 tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [],
165 ),
166 );
167});
168
169const describeFeatureGroupOrUndefined = (name: string) =>
170 sagemaker
171 .describeFeatureGroup({ FeatureGroupName: name })
172 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.succeed(undefined)));
173
174const toAttrs = (
175 described: sagemaker.DescribeFeatureGroupResponse,
176): FeatureGroup["Attributes"] => ({
177 featureGroupName: described.FeatureGroupName,
178 featureGroupArn: described.FeatureGroupArn,
179 recordIdentifierFeatureName: described.RecordIdentifierFeatureName,
180 eventTimeFeatureName: described.EventTimeFeatureName,
181});
182
183/**
184 * The feature group is still transitioning toward the awaited state —
185 * retried by the bounded wait schedule.
186 */
187class FeatureGroupNotReady extends Data.TaggedError("FeatureGroupNotReady")<{
188 readonly featureGroupName: string;
189 readonly status: string | undefined;
190}> {}
191
192/**
193 * The feature group's asynchronous creation converged to the terminal
194 * `CreateFailed` status.
195 */
196export class FeatureGroupCreateFailed extends Data.TaggedError(
197 "FeatureGroupCreateFailed",
198)<{
199 readonly featureGroupName: string;
200 readonly message: string | undefined;
201}> {}
202
203// Explicitly-typed retry wrapper — an inline `Effect.retry` in provider
204// lifecycle code leaks `Retry.Return`'s conditional type into declaration
205// emit and widens the provider layer to `unknown` for every consumer of
206// `AWS.providers()`.
207const retryWhileNotReady = <A, E extends { readonly _tag: string }, R>(
208 self: Effect.Effect<A, E, R>,
209): Effect.Effect<A, E, R> =>
210 Effect.retry(self, {
211 while: (e) => e._tag === "FeatureGroupNotReady",
212 // Online-store creation typically converges in well under 2 minutes;
213 // offline-store (Glue) creation can take a few. Poll 5s up to ~5 min.
214 schedule: Schedule.max([Schedule.spaced("5 seconds"), Schedule.recurs(60)]),
215 });
216
217const waitForFeatureGroup = (name: string, target: "Created" | "Gone") =>
218 retryWhileNotReady(
219 Effect.gen(function* () {
220 const described = yield* describeFeatureGroupOrUndefined(name);
221 if (target === "Gone") {
222 if (described === undefined) return;
223 return yield* Effect.fail(
224 new FeatureGroupNotReady({
225 featureGroupName: name,
226 status: described.FeatureGroupStatus,
227 }),
228 );
229 }
230 if (described?.FeatureGroupStatus === "Created") return;
231 if (
232 described?.FeatureGroupStatus === "CreateFailed" ||
233 described?.FeatureGroupStatus === "DeleteFailed"
234 ) {
235 return yield* Effect.fail(
236 new FeatureGroupCreateFailed({
237 featureGroupName: name,
238 message: described.FailureReason,
239 }),
240 );
241 }
242 return yield* Effect.fail(
243 new FeatureGroupNotReady({
244 featureGroupName: name,
245 status: described?.FeatureGroupStatus,
246 }),
247 );
248 }),
249 );
250
251export const FeatureGroupProvider = () =>
252 Provider.effect(
253 FeatureGroup,
254 Effect.gen(function* () {
255 return {
256 stables: [
257 "featureGroupName",
258 "featureGroupArn",
259 "recordIdentifierFeatureName",
260 "eventTimeFeatureName",
261 ],
262 list: () =>
263 Effect.gen(function* () {
264 const summaries = yield* sagemaker.listFeatureGroups.pages({}).pipe(
265 EffectStream.runCollect,
266 Effect.map((chunk) =>
267 Array.from(chunk).flatMap(
268 (page) => page.FeatureGroupSummaries ?? [],
269 ),
270 ),
271 );
272 const hydrated = yield* Effect.forEach(
273 summaries.flatMap((s) =>
274 s.FeatureGroupName !== undefined ? [s.FeatureGroupName] : [],
275 ),
276 (name) =>
277 describeFeatureGroupOrUndefined(name).pipe(
278 Effect.map((d) => (d === undefined ? undefined : toAttrs(d))),
279 ),
280 { concurrency: 5 },
281 );
282 return hydrated.filter(
283 (attrs): attrs is FeatureGroup["Attributes"] =>
284 attrs !== undefined,
285 );
286 }),
287 read: Effect.fn(function* ({ id, olds, output }) {
288 const name =
289 output?.featureGroupName ??
290 (yield* createFeatureGroupName(id, olds ?? {}));
291 const described = yield* describeFeatureGroupOrUndefined(name);
292 if (!described || described.FeatureGroupStatus === "Deleting") {
293 return undefined;
294 }
295 const attrs = toAttrs(described);
296 const tags = yield* fetchFeatureGroupTags(attrs.featureGroupArn);
297 return (yield* hasAlchemyTags(id, tags as Tags))
298 ? attrs
299 : Unowned(attrs);
300 }),
301 diff: Effect.fn(function* ({ id, news, olds }) {
302 if (!isResolved(news)) return;
303 if (olds === undefined) return;
304 // The feature group schema and stores are fixed at creation —
305 // everything except tags replaces.
306 const identity = (props: FeatureGroupProps) => [
307 props.recordIdentifierFeatureName,
308 props.eventTimeFeatureName,
309 props.featureDefinitions,
310 props.onlineStoreConfig,
311 props.offlineStoreConfig,
312 props.throughputConfig,
313 props.roleArn,
314 props.description,
315 ];
316 const oldName = yield* createFeatureGroupName(id, olds);
317 const newName = yield* createFeatureGroupName(id, news);
318 if (
319 oldName !== newName ||
320 JSON.stringify(identity(olds)) !== JSON.stringify(identity(news))
321 ) {
322 return { action: "replace" } as const;
323 }
324 }),
325 reconcile: Effect.fn(function* ({ id, news, output, session }) {
326 if (!news) {
327 return yield* Effect.fail(
328 new Error("SageMaker FeatureGroup requires props"),
329 );
330 }
331 const name =
332 output?.featureGroupName ??
333 (yield* createFeatureGroupName(id, news));
334 const internalTags = yield* createInternalTags(id);
335 const desiredTags = { ...internalTags, ...news.tags };
336
337 // Observe.
338 let described = yield* describeFeatureGroupOrUndefined(name);
339
340 // Ensure — create if missing; tolerate the already-exists race.
341 if (described === undefined) {
342 yield* sagemaker
343 .createFeatureGroup({
344 FeatureGroupName: name,
345 RecordIdentifierFeatureName: news.recordIdentifierFeatureName,
346 EventTimeFeatureName: news.eventTimeFeatureName,
347 FeatureDefinitions: news.featureDefinitions,
348 OnlineStoreConfig: news.onlineStoreConfig,
349 OfflineStoreConfig: news.offlineStoreConfig,
350 ThroughputConfig: news.throughputConfig,
351 RoleArn: news.roleArn,
352 Description: news.description,
353 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
354 Key,
355 Value,
356 })),
357 })
358 .pipe(Effect.catchTag("ResourceInUse", () => Effect.void));
359 yield* session.note(`Creating feature group ${name}...`);
360 }
361
362 // Converge to Created (creation is asynchronous).
363 yield* waitForFeatureGroup(name, "Created");
364 described = yield* describeFeatureGroupOrUndefined(name);
365 if (described === undefined) {
366 return yield* Effect.fail(
367 new Error(`failed to read reconciled feature group ${name}`),
368 );
369 }
370 const attrs = toAttrs(described);
371
372 // Sync tags — diff against OBSERVED cloud tags.
373 const currentTags = yield* fetchFeatureGroupTags(
374 attrs.featureGroupArn,
375 );
376 const { removed, upsert } = diffTags(currentTags, desiredTags);
377 if (removed.length > 0) {
378 yield* sagemaker.deleteTags({
379 ResourceArn: attrs.featureGroupArn,
380 TagKeys: removed,
381 });
382 }
383 if (upsert.length > 0) {
384 yield* sagemaker.addTags({
385 ResourceArn: attrs.featureGroupArn,
386 Tags: upsert.map(({ Key, Value }) => ({ Key, Value })),
387 });
388 }
389
390 yield* session.note(attrs.featureGroupArn);
391 return attrs;
392 }),
393 delete: Effect.fn(function* ({ output }) {
394 yield* sagemaker
395 .deleteFeatureGroup({
396 FeatureGroupName: output.featureGroupName,
397 })
398 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.void));
399 yield* waitForFeatureGroup(output.featureGroupName, "Gone");
400 }),
401 };
402 }),
403 );