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 EndpointStatus = sagemaker.EndpointStatus;
20
21export interface EndpointProps {
22 /**
23 * Name of the endpoint. Maximum 63 characters.
24 * @default ${app}-${stage}-${id}
25 */
26 endpointName?: string;
27 /**
28 * Name of the `EndpointConfig` that describes the models and hosting
29 * resources to deploy. Changing it updates the live endpoint in place
30 * (blue/green by default on the SageMaker side).
31 */
32 endpointConfigName: string;
33 /**
34 * Deployment (blue/green or rolling) configuration applied when the
35 * endpoint is updated to a new configuration.
36 */
37 deploymentConfig?: sagemaker.DeploymentConfig;
38 /**
39 * Tags to associate with the endpoint. Merged with internal Alchemy tags.
40 */
41 tags?: Record<string, string>;
42}
43
44export interface Endpoint extends Resource<
45 "AWS.SageMaker.Endpoint",
46 EndpointProps,
47 {
48 /**
49 * The endpoint's name.
50 */
51 endpointName: string;
52 /**
53 * ARN of the endpoint.
54 */
55 endpointArn: string;
56 /**
57 * Lifecycle status of the endpoint after reconcile (`InService`).
58 */
59 endpointStatus: EndpointStatus | undefined;
60 },
61 never,
62 Providers
63> {}
64
65/**
66 * An Amazon SageMaker Endpoint — the live, invocable deployment of an
67 * `EndpointConfig`. Provisioning takes minutes and **bills while the
68 * endpoint exists** (serverless variants bill per request; instance variants
69 * bill per instance-hour). Destroy endpoints promptly.
70 *
71 * Invoke a deployed endpoint from a function with
72 * `AWS.SageMakerRuntime.InvokeEndpoint`.
73 * ### Creating Endpoints
74 * **Example:** Deploy an EndpointConfig
75 * ```typescript
76 * import * as AWS from "alchemy/AWS";
77 *
78 * const endpoint = yield* AWS.SageMaker.Endpoint("MyEndpoint", {
79 * endpointConfigName: config.endpointConfigName,
80 * });
81 * ```
82 *
83 * ### Invoking
84 * **Example:** Invoke from a Lambda function
85 * ```typescript
86 * // init
87 * const invoke = yield* AWS.SageMakerRuntime.InvokeEndpoint(
88 * endpoint.endpointName,
89 * );
90 *
91 * // runtime
92 * const result = yield* invoke({
93 * ContentType: "application/json",
94 * Body: JSON.stringify({ instances: [[1, 2, 3, 4]] }),
95 * });
96 * ```
97 *
98 * @resource
99 */
100export const Endpoint = Resource<Endpoint>("AWS.SageMaker.Endpoint");
101
102const createEndpointName = (
103 id: string,
104 props: { endpointName?: string | undefined },
105) =>
106 props.endpointName
107 ? Effect.succeed(props.endpointName)
108 : createPhysicalName({ id, maxLength: 63 });
109
110const fetchEndpointTags = Effect.fn(function* (arn: string) {
111 const response = yield* sagemaker
112 .listTags({ ResourceArn: arn })
113 .pipe(
114 Effect.catchTag("AccessDeniedException", () => Effect.succeed(undefined)),
115 );
116 return Object.fromEntries(
117 (response?.Tags ?? []).flatMap((tag) =>
118 tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [],
119 ),
120 );
121});
122
123const describeEndpointOrUndefined = (name: string) =>
124 sagemaker
125 .describeEndpoint({ EndpointName: name })
126 .pipe(Effect.catchTag("EndpointNotFound", () => Effect.succeed(undefined)));
127
128/**
129 * The endpoint is still transitioning toward the awaited state — retried by
130 * the bounded wait schedule.
131 */
132class EndpointNotReady extends Data.TaggedError("EndpointNotReady")<{
133 readonly endpointName: string;
134 readonly status: string | undefined;
135}> {}
136
137/**
138 * The endpoint's asynchronous provisioning converged to the terminal
139 * `Failed` status (e.g. the container image could not be pulled or the model
140 * server failed its health checks).
141 */
142export class EndpointProvisioningFailed extends Data.TaggedError(
143 "EndpointProvisioningFailed",
144)<{
145 readonly endpointName: string;
146 readonly message: string | undefined;
147}> {}
148
149// Explicitly-typed retry wrapper — an inline `Effect.retry` in provider
150// lifecycle code leaks `Retry.Return`'s conditional type into declaration
151// emit and widens the provider layer to `unknown` for every consumer of
152// `AWS.providers()`.
153const retryWhileNotReady = <A, E extends { readonly _tag: string }, R>(
154 self: Effect.Effect<A, E, R>,
155): Effect.Effect<A, E, R> =>
156 Effect.retry(self, {
157 while: (e) => e._tag === "EndpointNotReady",
158 // Endpoint provisioning takes ~3-10 min; poll every 15s up to ~20 min.
159 schedule: Schedule.max([
160 Schedule.spaced("15 seconds"),
161 Schedule.recurs(80),
162 ]),
163 });
164
165const waitForEndpoint = (name: string, target: "InService" | "Gone") =>
166 retryWhileNotReady(
167 Effect.gen(function* () {
168 const described = yield* describeEndpointOrUndefined(name);
169 if (target === "Gone") {
170 if (described === undefined) return;
171 return yield* Effect.fail(
172 new EndpointNotReady({
173 endpointName: name,
174 status: described.EndpointStatus,
175 }),
176 );
177 }
178 if (described?.EndpointStatus === "InService") return;
179 if (described?.EndpointStatus === "Failed") {
180 return yield* Effect.fail(
181 new EndpointProvisioningFailed({
182 endpointName: name,
183 message: described.FailureReason,
184 }),
185 );
186 }
187 return yield* Effect.fail(
188 new EndpointNotReady({
189 endpointName: name,
190 status: described?.EndpointStatus,
191 }),
192 );
193 }),
194 );
195
196export const EndpointProvider = () =>
197 Provider.effect(
198 Endpoint,
199 Effect.gen(function* () {
200 return {
201 stables: ["endpointName", "endpointArn"],
202 list: () =>
203 Effect.gen(function* () {
204 const summaries = yield* sagemaker.listEndpoints.pages({}).pipe(
205 EffectStream.runCollect,
206 Effect.map((chunk) =>
207 Array.from(chunk).flatMap((page) => page.Endpoints ?? []),
208 ),
209 );
210 return summaries.flatMap((s) =>
211 s.EndpointName !== undefined && s.EndpointArn !== undefined
212 ? [
213 {
214 endpointName: s.EndpointName,
215 endpointArn: s.EndpointArn,
216 endpointStatus: s.EndpointStatus,
217 },
218 ]
219 : [],
220 );
221 }),
222 read: Effect.fn(function* ({ id, olds, output }) {
223 const name =
224 output?.endpointName ?? (yield* createEndpointName(id, olds ?? {}));
225 const described = yield* describeEndpointOrUndefined(name);
226 if (!described || described.EndpointStatus === "Deleting") {
227 return undefined;
228 }
229 const attrs = {
230 endpointName: described.EndpointName,
231 endpointArn: described.EndpointArn,
232 endpointStatus: described.EndpointStatus,
233 };
234 const tags = yield* fetchEndpointTags(described.EndpointArn);
235 return (yield* hasAlchemyTags(id, tags as Tags))
236 ? attrs
237 : Unowned(attrs);
238 }),
239 diff: Effect.fn(function* ({ id, news, olds }) {
240 if (!isResolved(news)) return;
241 if (olds === undefined) return;
242 const oldName = yield* createEndpointName(id, olds);
243 const newName = yield* createEndpointName(id, news);
244 if (oldName !== newName) {
245 return { action: "replace" } as const;
246 }
247 // endpointConfigName changes update the live endpoint in place.
248 }),
249 reconcile: Effect.fn(function* ({ id, news, output, session }) {
250 if (!news) {
251 return yield* Effect.fail(
252 new Error("SageMaker Endpoint requires props"),
253 );
254 }
255 const name =
256 output?.endpointName ?? (yield* createEndpointName(id, news));
257 const internalTags = yield* createInternalTags(id);
258 const desiredTags = { ...internalTags, ...news.tags };
259
260 // Observe.
261 let described = yield* describeEndpointOrUndefined(name);
262
263 // A terminally-failed endpoint cannot be updated — delete it and
264 // fall through to a fresh create so the reconciler converges.
265 if (described?.EndpointStatus === "Failed") {
266 yield* session.note(
267 `Endpoint ${name} is Failed (${described.FailureReason ?? "unknown"}) — recreating`,
268 );
269 yield* sagemaker
270 .deleteEndpoint({ EndpointName: name })
271 .pipe(Effect.catchTag("EndpointNotFound", () => Effect.void));
272 yield* waitForEndpoint(name, "Gone");
273 described = undefined;
274 }
275
276 // Ensure — create if missing; tolerate the already-exists race.
277 if (described === undefined) {
278 yield* sagemaker
279 .createEndpoint({
280 EndpointName: name,
281 EndpointConfigName: news.endpointConfigName,
282 DeploymentConfig: news.deploymentConfig,
283 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
284 Key,
285 Value,
286 })),
287 })
288 .pipe(
289 Effect.catchTag("EndpointAlreadyExists", () => Effect.void),
290 );
291 yield* session.note(`Creating endpoint ${name}...`);
292 } else if (
293 described.EndpointConfigName !== news.endpointConfigName &&
294 described.EndpointStatus === "InService"
295 ) {
296 // Sync — roll the live endpoint onto the desired configuration.
297 yield* sagemaker.updateEndpoint({
298 EndpointName: name,
299 EndpointConfigName: news.endpointConfigName,
300 DeploymentConfig: news.deploymentConfig,
301 });
302 yield* session.note(
303 `Updating endpoint ${name} to config ${news.endpointConfigName}...`,
304 );
305 }
306
307 // Converge to InService (also settles Creating/Updating states we
308 // observed mid-flight, e.g. after a crashed reconcile).
309 yield* waitForEndpoint(name, "InService");
310 const final = yield* describeEndpointOrUndefined(name);
311 if (final === undefined) {
312 return yield* Effect.fail(
313 new Error(`failed to read reconciled endpoint ${name}`),
314 );
315 }
316
317 // Sync tags — diff against OBSERVED cloud tags.
318 const currentTags = yield* fetchEndpointTags(final.EndpointArn);
319 const { removed, upsert } = diffTags(currentTags, desiredTags);
320 if (removed.length > 0) {
321 yield* sagemaker.deleteTags({
322 ResourceArn: final.EndpointArn,
323 TagKeys: removed,
324 });
325 }
326 if (upsert.length > 0) {
327 yield* sagemaker.addTags({
328 ResourceArn: final.EndpointArn,
329 Tags: upsert.map(({ Key, Value }) => ({ Key, Value })),
330 });
331 }
332
333 yield* session.note(final.EndpointArn);
334 return {
335 endpointName: final.EndpointName,
336 endpointArn: final.EndpointArn,
337 endpointStatus: final.EndpointStatus,
338 };
339 }),
340 delete: Effect.fn(function* ({ output }) {
341 yield* sagemaker
342 .deleteEndpoint({ EndpointName: output.endpointName })
343 .pipe(Effect.catchTag("EndpointNotFound", () => Effect.void));
344 yield* waitForEndpoint(output.endpointName, "Gone");
345 }),
346 };
347 }),
348 );