1import * as sagemaker from "@distilled.cloud/aws/sagemaker";
2import * as Effect from "effect/Effect";
3import * as Schedule from "effect/Schedule";
4import * as EffectStream from "effect/Stream";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import { createPhysicalName } from "../../PhysicalName.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import {
11 createInternalTags,
12 diffTags,
13 hasAlchemyTags,
14 type Tags,
15} from "../../Tags.ts";
16import type { Providers } from "../Providers.ts";
17
18export interface ModelProps {
19 /**
20 * Name of the model. Must be unique within the account/region.
21 * Maximum 63 characters, alphanumeric and hyphens.
22 * @default ${app}-${stage}-${id}
23 */
24 modelName?: string;
25 /**
26 * The primary inference container: the ECR image (and optionally model
27 * artifacts in S3) SageMaker runs when the model is deployed to an
28 * endpoint. Exactly one of `primaryContainer` or `containers` is required.
29 */
30 primaryContainer?: sagemaker.ContainerDefinition;
31 /**
32 * Containers for an inference pipeline (executed as a sequence) or a
33 * multi-container model. Mutually exclusive with `primaryContainer`.
34 */
35 containers?: sagemaker.ContainerDefinition[];
36 /**
37 * How containers in a multi-container model are run (`Serial` pipeline or
38 * `Direct` invocation of a specific container).
39 */
40 inferenceExecutionConfig?: sagemaker.InferenceExecutionConfig;
41 /**
42 * ARN of the IAM role SageMaker assumes to pull the container image and
43 * model artifacts. The role must trust `sagemaker.amazonaws.com`.
44 */
45 executionRoleArn: string;
46 /**
47 * VPC configuration for the model's containers.
48 */
49 vpcConfig?: sagemaker.VpcConfig;
50 /**
51 * Isolate the model containers from the network (no outbound calls).
52 * @default false
53 */
54 enableNetworkIsolation?: boolean;
55 /**
56 * Tags to associate with the model. Merged with internal Alchemy tags.
57 */
58 tags?: Record<string, string>;
59}
60
61export interface Model extends Resource<
62 "AWS.SageMaker.Model",
63 ModelProps,
64 {
65 /**
66 * The model's name.
67 */
68 modelName: string;
69 /**
70 * ARN of the model.
71 */
72 modelArn: string;
73 },
74 never,
75 Providers
76> {}
77
78/**
79 * An Amazon SageMaker Model — the immutable pairing of an inference
80 * container image (and optional S3 model artifacts) with an execution role.
81 * A model is pure configuration: it costs nothing until it is deployed to an
82 * endpoint via an `EndpointConfig` + `Endpoint`.
83 *
84 * SageMaker models are immutable — any change other than tags replaces the
85 * model.
86 * ### Creating Models
87 * **Example:** Model from an ECR image
88 * ```typescript
89 * import * as AWS from "alchemy/AWS";
90 *
91 * const role = yield* AWS.IAM.Role("SageMakerRole", {
92 * assumeRolePolicyDocument: {
93 * Version: "2012-10-17",
94 * Statement: [{
95 * Effect: "Allow",
96 * Principal: { Service: "sagemaker.amazonaws.com" },
97 * Action: ["sts:AssumeRole"],
98 * }],
99 * },
100 * managedPolicyArns: ["arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"],
101 * });
102 *
103 * const model = yield* AWS.SageMaker.Model("MyModel", {
104 * executionRoleArn: role.roleArn,
105 * primaryContainer: {
106 * Image: "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-inference:latest",
107 * ModelDataUrl: "s3://my-bucket/model.tar.gz",
108 * },
109 * });
110 * ```
111 *
112 * **Example:** Serverless deployment (Model → EndpointConfig → Endpoint)
113 * ```typescript
114 * const config = yield* AWS.SageMaker.EndpointConfig("MyConfig", {
115 * productionVariants: [{
116 * VariantName: "AllTraffic",
117 * ModelName: model.modelName,
118 * ServerlessConfig: { MemorySizeInMB: 2048, MaxConcurrency: 5 },
119 * }],
120 * });
121 * const endpoint = yield* AWS.SageMaker.Endpoint("MyEndpoint", {
122 * endpointConfigName: config.endpointConfigName,
123 * });
124 * ```
125 *
126 * @resource
127 */
128export const Model = Resource<Model>("AWS.SageMaker.Model");
129
130const createModelName = (
131 id: string,
132 props: { modelName?: string | undefined },
133) =>
134 props.modelName
135 ? Effect.succeed(props.modelName)
136 : createPhysicalName({ id, maxLength: 63 });
137
138const fetchModelTags = Effect.fn(function* (arn: string) {
139 const tags = yield* sagemaker.listTags.items({ ResourceArn: arn }).pipe(
140 EffectStream.runCollect,
141 Effect.map((chunk) => Array.from(chunk)),
142 // A just-deleted (or foreign) resource ARN surfaces as AccessDenied.
143 Effect.catchTag("AccessDeniedException", () =>
144 Effect.succeed<sagemaker.Tag[]>([]),
145 ),
146 );
147 return Object.fromEntries(
148 tags.flatMap((tag) =>
149 tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [],
150 ),
151 );
152});
153
154const readModel = Effect.fn(function* (modelName: string) {
155 const described = yield* sagemaker
156 .describeModel({ ModelName: modelName })
157 .pipe(Effect.catchTag("ModelNotFound", () => Effect.succeed(undefined)));
158 if (!described) return undefined;
159 return {
160 modelName: described.ModelName,
161 modelArn: described.ModelArn,
162 };
163});
164
165// CreateModel validates the execution role (existence + sagemaker trust); a
166// freshly-created IAM role takes a few seconds to propagate and surfaces as
167// ValidationException. Bounded retry through the propagation window.
168// Explicitly typed so Retry.Return's conditional type never reaches
169// declaration emit (it would widen AWS.providers() to unknown downstream).
170const retryThroughIamPropagation = <A, E extends { readonly _tag: string }, R>(
171 self: Effect.Effect<A, E, R>,
172): Effect.Effect<A, E, R> =>
173 Effect.retry(self, {
174 while: (e) => e._tag === "ValidationException",
175 schedule: Schedule.max([Schedule.spaced("5 seconds"), Schedule.recurs(10)]),
176 });
177
178export const ModelProvider = () =>
179 Provider.effect(
180 Model,
181 Effect.gen(function* () {
182 return {
183 stables: ["modelName", "modelArn"],
184 list: () =>
185 Effect.gen(function* () {
186 const summaries = yield* sagemaker.listModels.pages({}).pipe(
187 EffectStream.runCollect,
188 Effect.map((chunk) =>
189 Array.from(chunk).flatMap((page) => page.Models ?? []),
190 ),
191 );
192 return summaries.flatMap((s) =>
193 s.ModelName !== undefined && s.ModelArn !== undefined
194 ? [{ modelName: s.ModelName, modelArn: s.ModelArn }]
195 : [],
196 );
197 }),
198 read: Effect.fn(function* ({ id, olds, output }) {
199 const modelName =
200 output?.modelName ?? (yield* createModelName(id, olds ?? {}));
201 const attrs = yield* readModel(modelName);
202 if (!attrs) return undefined;
203 const tags = yield* fetchModelTags(attrs.modelArn);
204 return (yield* hasAlchemyTags(id, tags as Tags))
205 ? attrs
206 : Unowned(attrs);
207 }),
208 diff: Effect.fn(function* ({ id, news, olds }) {
209 if (!isResolved(news)) return;
210 if (olds === undefined) return;
211 // SageMaker models are immutable — everything except tags replaces.
212 const identity = (props: ModelProps) => [
213 props.primaryContainer,
214 props.containers,
215 props.inferenceExecutionConfig,
216 props.executionRoleArn,
217 props.vpcConfig,
218 props.enableNetworkIsolation ?? false,
219 ];
220 const oldName = yield* createModelName(id, olds);
221 const newName = yield* createModelName(id, news);
222 if (
223 oldName !== newName ||
224 JSON.stringify(identity(olds)) !== JSON.stringify(identity(news))
225 ) {
226 return { action: "replace" } as const;
227 }
228 }),
229 reconcile: Effect.fn(function* ({ id, news, output, session }) {
230 if (!news) {
231 return yield* Effect.fail(
232 new Error("SageMaker Model requires props"),
233 );
234 }
235 const modelName =
236 output?.modelName ?? (yield* createModelName(id, news));
237 const internalTags = yield* createInternalTags(id);
238 const desiredTags = { ...internalTags, ...news.tags };
239
240 // Observe — cloud state is authoritative.
241 let attrs = yield* readModel(modelName);
242
243 // Ensure — create if missing; tolerate the already-exists race.
244 if (attrs === undefined) {
245 yield* retryThroughIamPropagation(
246 sagemaker.createModel({
247 ModelName: modelName,
248 PrimaryContainer: news.primaryContainer,
249 Containers: news.containers,
250 InferenceExecutionConfig: news.inferenceExecutionConfig,
251 ExecutionRoleArn: news.executionRoleArn,
252 VpcConfig: news.vpcConfig,
253 EnableNetworkIsolation: news.enableNetworkIsolation,
254 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
255 Key,
256 Value,
257 })),
258 }),
259 ).pipe(Effect.catchTag("ModelAlreadyExists", () => Effect.void));
260 attrs = yield* readModel(modelName);
261 if (attrs === undefined) {
262 return yield* Effect.fail(
263 new Error(`failed to read created model ${modelName}`),
264 );
265 }
266 }
267
268 // Sync tags — diff against OBSERVED cloud tags.
269 const currentTags = yield* fetchModelTags(attrs.modelArn);
270 const { removed, upsert } = diffTags(currentTags, desiredTags);
271 if (removed.length > 0) {
272 yield* sagemaker.deleteTags({
273 ResourceArn: attrs.modelArn,
274 TagKeys: removed,
275 });
276 }
277 if (upsert.length > 0) {
278 yield* sagemaker.addTags({
279 ResourceArn: attrs.modelArn,
280 Tags: upsert.map(({ Key, Value }) => ({ Key, Value })),
281 });
282 }
283
284 yield* session.note(attrs.modelArn);
285 return attrs;
286 }),
287 delete: Effect.fn(function* ({ output }) {
288 yield* sagemaker
289 .deleteModel({ ModelName: output.modelName })
290 .pipe(Effect.catchTag("ModelNotFound", () => Effect.void));
291 }),
292 };
293 }),
294 );