1import * as AWS from "alchemy/AWS";
2import * as Kubernetes from "alchemy/Kubernetes";
3import * as Output from "alchemy/Output";
4import * as Effect from "effect/Effect";
5import { FetchHyperPodChart } from "./hyperpod-chart.ts";
6import { UploadLifecycleScript } from "./lifecycle.ts";
7
8/**
9 * Shared infrastructure for the EKS-orchestrated HyperPod stack:
10 * network → EKS control plane → HyperPod cluster → task governance.
11 *
12 * Resources are memoized by logical id, so this Effect can be yielded from
13 * the stack program and from workload modules and always converges on the
14 * same instances.
15 */
16export const HyperPodEksInfra = Effect.gen(function* () {
17 // HyperPod nodes must live in private subnets; NAT gives them a path to
18 // pull container images.
19 const network = yield* AWS.EC2.Network("HyperPodNetwork", {
20 cidrBlock: "10.72.0.0/16",
21 availabilityZones: 2,
22 nat: "single",
23 });
24
25 // The control-plane role for the orchestrating EKS cluster.
26 const eksRole = yield* AWS.IAM.Role("OrchestratorRole", {
27 assumeRolePolicyDocument: {
28 Version: "2012-10-17",
29 Statement: [
30 {
31 Effect: "Allow",
32 Principal: { Service: "eks.amazonaws.com" },
33 Action: ["sts:AssumeRole"],
34 },
35 ],
36 },
37 managedPolicyArns: ["arn:aws:iam::aws:policy/AmazonEKSClusterPolicy"],
38 });
39
40 // The orchestrator: a plain EKS control plane. HyperPod requires the
41 // `API` (or `API_AND_CONFIG_MAP`) authentication mode — EKS's own
42 // CONFIG_MAP default is rejected by CreateCluster. HyperPod provides the
43 // nodes, so there is no compute config here.
44 const eks = yield* AWS.EKS.Cluster("Orchestrator", {
45 roleArn: eksRole.roleArn,
46 // HyperPod trails the newest EKS Kubernetes version — pin one it
47 // supports (1.28-1.35 today) rather than taking the EKS default.
48 version: "1.34",
49 resourcesVpcConfig: { subnetIds: network.privateSubnetIds },
50 accessConfig: {
51 authenticationMode: "API",
52 bootstrapClusterCreatorAdminPermissions: true,
53 },
54 });
55
56 // Lifecycle scripts are REQUIRED for EKS-orchestrated instance groups
57 // (and Slurm continuous provisioning): every node runs its `OnCreate`
58 // script from S3 when it boots.
59 const lifecycleBucket = yield* AWS.S3.Bucket("EksLifecycleScripts", {
60 forceDestroy: true,
61 });
62 const script = yield* UploadLifecycleScript({
63 bucketName: lifecycleBucket.bucketName,
64 });
65
66 // The HyperPod instance role: the AWS-managed cluster-instance policy
67 // plus the EKS-orchestration networking/ECR/pod-identity permissions
68 // from the HyperPod IAM guide.
69 const role = yield* AWS.IAM.Role("HyperPodEksInstanceRole", {
70 assumeRolePolicyDocument: {
71 Version: "2012-10-17",
72 Statement: [
73 {
74 Effect: "Allow",
75 Principal: { Service: "sagemaker.amazonaws.com" },
76 Action: ["sts:AssumeRole"],
77 },
78 ],
79 },
80 managedPolicyArns: [
81 "arn:aws:iam::aws:policy/AmazonSageMakerClusterInstanceRolePolicy",
82 ],
83 inlinePolicies: {
84 "hyperpod-eks-nodes": {
85 Version: "2012-10-17",
86 Statement: [
87 {
88 Effect: "Allow",
89 Action: [
90 "ec2:AssignPrivateIpAddresses",
91 "ec2:AttachNetworkInterface",
92 "ec2:CreateNetworkInterface",
93 "ec2:CreateNetworkInterfacePermission",
94 "ec2:DeleteNetworkInterface",
95 "ec2:DeleteNetworkInterfacePermission",
96 "ec2:DescribeInstances",
97 "ec2:DescribeInstanceTypes",
98 "ec2:DescribeNetworkInterfaces",
99 "ec2:DescribeTags",
100 "ec2:DescribeVpcs",
101 "ec2:DescribeDhcpOptions",
102 "ec2:DescribeSubnets",
103 "ec2:DescribeSecurityGroups",
104 "ec2:DetachNetworkInterface",
105 "ec2:ModifyNetworkInterfaceAttribute",
106 "ec2:UnassignPrivateIpAddresses",
107 "ecr:BatchCheckLayerAvailability",
108 "ecr:BatchGetImage",
109 "ecr:GetAuthorizationToken",
110 "ecr:GetDownloadUrlForLayer",
111 "eks-auth:AssumeRoleForPodIdentity",
112 ],
113 Resource: "*",
114 },
115 {
116 Effect: "Allow",
117 Action: ["ec2:CreateTags"],
118 Resource: ["arn:aws:ec2:*:*:network-interface/*"],
119 },
120 {
121 Effect: "Allow",
122 Action: ["s3:GetObject", "s3:ListBucket"],
123 Resource: [
124 lifecycleBucket.bucketArn,
125 Output.interpolate`${lifecycleBucket.bucketArn}/*`,
126 ],
127 },
128 ],
129 },
130 },
131 });
132
133 // The HyperPod dependencies Helm chart (health-monitoring agent,
134 // training operators, device plugins, RBAC) is MANDATORY — SageMaker
135 // validates it and fails the cluster with "missing one or more required
136 // dependencies" otherwise. Rendered locally and applied to the EKS
137 // cluster before the HyperPod cluster attaches.
138 const chart = yield* FetchHyperPodChart({
139 repo: "https://github.com/aws/sagemaker-hyperpod-cli.git",
140 });
141 const dependencies = yield* Kubernetes.HelmChart("HyperPodDependencies", {
142 cluster: eks,
143 chart: chart.chartPath,
144 releaseName: "dependencies",
145 namespace: "kube-system",
146 values: {
147 // The vendored legacy mpi-operator CRD fails Kubernetes >=1.34
148 // strict server-side-apply validation, and MPI jobs aren't part of
149 // this example.
150 "mpi-operator": { enabled: false },
151 },
152 });
153
154 // The HyperPod cluster, attached to the EKS control plane. SageMaker
155 // creates the EKS access entry for the node role automatically.
156 // `Output.all` holds the attach until the dependencies chart is applied.
157 // Instance-group keys carry through to the cluster's attributes, so
158 // workloads reference `hyperpod.instanceGroups.workers` — typed per key.
159 const hyperpod = yield* AWS.SageMaker.Cluster("HyperPod", {
160 orchestrator: {
161 Eks: {
162 ClusterArn: Output.map(
163 Output.all(eks.clusterArn, dependencies.releaseName),
164 ([clusterArn]) => clusterArn,
165 ),
166 },
167 },
168 vpcConfig: {
169 // The EKS-managed cluster security group already allows node ↔
170 // control-plane and intra-cluster traffic.
171 SecurityGroupIds: [
172 Output.map(
173 eks.resourcesVpcConfig,
174 (vpc) => vpc.clusterSecurityGroupId!,
175 ),
176 ],
177 Subnets: network.privateSubnetIds,
178 },
179 instanceGroups: {
180 workers: {
181 InstanceType: "ml.t3.medium",
182 InstanceCount: 1,
183 ExecutionRole: role.roleArn,
184 LifeCycleConfig: {
185 SourceS3Uri: script.sourceS3Uri,
186 OnCreate: script.onCreate,
187 },
188 },
189 },
190 // Node auto-replacement is not supported for CPU instances.
191 nodeRecovery: "None",
192 tags: { app: "aws-hyperpod-example" },
193 });
194
195 // Task governance (Kueue + the HyperPod scheduler) ships as an EKS
196 // add-on. Deriving the cluster name from the HyperPod cluster's
197 // orchestrator ARN (instead of using eks.clusterName directly) makes the
198 // add-on deploy AFTER the HyperPod nodes join — its controllers need a
199 // schedulable node.
200 const governance = yield* AWS.EKS.Addon("TaskGovernance", {
201 clusterName: Output.map(hyperpod.orchestratorEksClusterArn, (arn) =>
202 arn!.split("/").pop()!,
203 ),
204 addonName: "amazon-sagemaker-hyperpod-taskgovernance",
205 });
206
207 // The cluster policy: priority classes + fair-share arbitration.
208 // `Output.all` makes the policy wait for the governance add-on (which
209 // hosts the controllers that reconcile it) while still reading the
210 // HyperPod cluster's ARN.
211 const scheduler = yield* AWS.SageMaker.ClusterSchedulerConfig("Scheduler", {
212 clusterArn: Output.map(
213 Output.all(hyperpod.clusterArn, governance.addonArn),
214 ([clusterArn]) => clusterArn,
215 ),
216 schedulerConfig: {
217 PriorityClasses: [
218 { Name: "inference", Weight: 100 },
219 { Name: "training", Weight: 75 },
220 ],
221 FairShare: "Enabled",
222 },
223 });
224
225 // The research team's quota. Creating it materializes the
226 // `hyperpod-ns-research` namespace and its Kueue LocalQueue.
227 const researchQuota = yield* AWS.SageMaker.ComputeQuota("ResearchQuota", {
228 clusterArn: scheduler.clusterArn,
229 computeQuotaTarget: { TeamName: "research", FairShareWeight: 10 },
230 computeQuotaConfig: {
231 ComputeQuotaResources: [{ InstanceType: "ml.t3.medium", Count: 1 }],
232 },
233 });
234
235 return {
236 network,
237 eks,
238 role,
239 hyperpod,
240 governance,
241 scheduler,
242 researchQuota,
243 };
244});