1import { Role } from "@/AWS/IAM/Role.ts";
2import * as Lambda from "@/AWS/Lambda";
3import * as SageMaker from "@/AWS/SageMaker";
4import * as Duration from "effect/Duration";
5import * as Effect from "effect/Effect";
6import * as Layer from "effect/Layer";
7import { HttpServerRequest } from "effect/http/HttpServerRequest";
8import * as HttpServerResponse from "effect/http/HttpServerResponse";
9import path from "pathe";
10
11const main = path.resolve(import.meta.dirname, "endpoint-handler.ts");
12
13export class SageMakerEndpointTestFunction extends Lambda.Function<Lambda.Function>()(
14 "SageMakerEndpointTestFunction",
15) {}
16
17// Gated fixture — deploying a live endpoint takes minutes and requires a
18// working serving container (AWS_TEST_SAGEMAKER_IMAGE). Only deployed by
19// EndpointBindings.test.ts when AWS_TEST_SAGEMAKER_ENDPOINT=1.
20export default SageMakerEndpointTestFunction.make(
21 {
22 main,
23 functionUrl: true,
24 timeout: Duration.seconds(30),
25 },
26 Effect.gen(function* () {
27 const image = process.env.AWS_TEST_SAGEMAKER_IMAGE!;
28 const modelData = process.env.AWS_TEST_SAGEMAKER_MODEL_DATA;
29
30 const role = yield* Role("SageMakerBindingsEndpointRole", {
31 assumeRolePolicyDocument: {
32 Version: "2012-10-17",
33 Statement: [
34 {
35 Effect: "Allow",
36 Principal: { Service: "sagemaker.amazonaws.com" },
37 Action: ["sts:AssumeRole"],
38 },
39 ],
40 },
41 managedPolicyArns: ["arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"],
42 });
43 const model = yield* SageMaker.Model("BindingsEndpointModel", {
44 executionRoleArn: role.roleArn,
45 primaryContainer: {
46 Image: image,
47 ...(modelData ? { ModelDataUrl: modelData } : {}),
48 },
49 });
50 const config = yield* SageMaker.EndpointConfig("BindingsEndpointConfig", {
51 productionVariants: [
52 {
53 VariantName: "AllTraffic",
54 ModelName: model.modelName,
55 ServerlessConfig: { MemorySizeInMB: 2048, MaxConcurrency: 1 },
56 },
57 ],
58 });
59 const endpoint = yield* SageMaker.Endpoint("BindingsEndpoint", {
60 endpointConfigName: config.endpointConfigName,
61 });
62
63 const describeEndpoint = yield* SageMaker.DescribeEndpoint(endpoint);
64 const updateWeights =
65 yield* SageMaker.UpdateEndpointWeightsAndCapacities(endpoint);
66
67 const bound = { describeEndpoint, updateWeights };
68
69 return {
70 fetch: Effect.gen(function* () {
71 const request = yield* HttpServerRequest;
72 const url = new URL(request.originalUrl);
73 const pathname = url.pathname;
74
75 if (request.method === "GET" && pathname === "/bindings") {
76 return yield* HttpServerResponse.json({ bound: Object.keys(bound) });
77 }
78
79 if (request.method === "GET" && pathname === "/describe-endpoint") {
80 const described = yield* describeEndpoint();
81 return yield* HttpServerResponse.json({
82 endpointName: described.EndpointName,
83 status: described.EndpointStatus,
84 variants: (described.ProductionVariants ?? []).map(
85 (v) => v.VariantName,
86 ),
87 });
88 }
89
90 if (request.method === "GET" && pathname === "/health") {
91 return yield* HttpServerResponse.json({ ok: true });
92 }
93
94 return yield* HttpServerResponse.json(
95 { error: "Not found", pathname },
96 { status: 404 },
97 );
98 }).pipe(Effect.orDie),
99 };
100 }).pipe(
101 Effect.provide(
102 Layer.mergeAll(
103 SageMaker.DescribeEndpointHttp,
104 SageMaker.UpdateEndpointWeightsAndCapacitiesHttp,
105 ),
106 ),
107 ),
108);