| 1 | import * as sagemaker from "@distilled.cloud/aws/sagemaker"; |
| 2 | import * as Data from "effect/Data"; |
| 3 | import * as Effect from "effect/Effect"; |
| 4 | import * as Schedule from "effect/Schedule"; |
| 5 | import * as EffectStream from "effect/Stream"; |
| 6 | import { Unowned } from "../../AdoptPolicy.ts"; |
| 7 | import { isResolved } from "../../Diff.ts"; |
| 8 | import type { Input } from "../../Input.ts"; |
| 9 | import * as Output from "../../Output.ts"; |
| 10 | import { createPhysicalName } from "../../PhysicalName.ts"; |
| 11 | import * as Provider from "../../Provider.ts"; |
| 12 | import { Resource } from "../../Resource.ts"; |
| 13 | import { |
| 14 | createInternalTags, |
| 15 | diffTags, |
| 16 | hasAlchemyTags, |
| 17 | type Tags, |
| 18 | } from "../../Tags.ts"; |
| 19 | import type { Providers } from "../Providers.ts"; |
| 20 | |
| 21 | export type ClusterStatus = sagemaker.ClusterStatus; |
| 22 | |
| 23 | /** |
| 24 | * The well-known node label carrying the HyperPod instance-group name. |
| 25 | * HyperPod nodes join the orchestrating EKS cluster as ordinary Kubernetes |
| 26 | * nodes carrying this label. |
| 27 | */ |
| 28 | export const HYPERPOD_INSTANCE_GROUP_LABEL = |
| 29 | "sagemaker.amazonaws.com/instance-group-name"; |
| 30 | |
| 31 | /** The well-known node label carrying HyperPod's node health verdict. */ |
| 32 | export const HYPERPOD_NODE_HEALTH_LABEL = |
| 33 | "sagemaker.amazonaws.com/node-health-status"; |
| 34 | |
| 35 | /** |
| 36 | * An instance group surfaced on the cluster's attributes — Kubernetes |
| 37 | * workloads pin themselves to the group by referencing its |
| 38 | * {@link ClusterInstanceGroupRef.nodeSelector}. |
| 39 | */ |
| 40 | export interface ClusterInstanceGroupRef { |
| 41 | /** The group's name (the `sagemaker.amazonaws.com/instance-group-name` node label). */ |
| 42 | InstanceGroupName: string; |
| 43 | /** The group's instance type (e.g. `ml.g5.xlarge`). */ |
| 44 | InstanceType: string | undefined; |
| 45 | /** Instances currently in service. */ |
| 46 | CurrentCount: number | undefined; |
| 47 | /** Instances the group is converging toward. */ |
| 48 | TargetCount: number | undefined; |
| 49 | /** |
| 50 | * A Kubernetes node selector pinning pods onto this group's |
| 51 | * health-checked nodes — pass it to a `Kubernetes.Job` / |
| 52 | * `Kubernetes.Deployment` through the `podTemplate` escape hatch: |
| 53 | * |
| 54 | * ```ts |
| 55 | * podTemplate: { |
| 56 | * spec: { nodeSelector: hyperpod.instanceGroups.workers.nodeSelector }, |
| 57 | * } |
| 58 | * ``` |
| 59 | */ |
| 60 | nodeSelector: Record<string, string>; |
| 61 | } |
| 62 | |
| 63 | /** |
| 64 | * One instance group's specification, sans name — the group's name is the |
| 65 | * key in {@link ClusterProps.instanceGroups}. |
| 66 | */ |
| 67 | export type ClusterInstanceGroup = Omit< |
| 68 | sagemaker.ClusterInstanceGroupSpecification, |
| 69 | "InstanceGroupName" |
| 70 | >; |
| 71 | |
| 72 | /** |
| 73 | * One restricted instance group's specification, sans name — the group's |
| 74 | * name is the key in {@link ClusterProps.restrictedInstanceGroups}. |
| 75 | */ |
| 76 | export type ClusterRestrictedInstanceGroup = Omit< |
| 77 | sagemaker.ClusterRestrictedInstanceGroupSpecification, |
| 78 | "InstanceGroupName" |
| 79 | >; |
| 80 | |
| 81 | export interface ClusterProps { |
| 82 | /** |
| 83 | * Name of the HyperPod cluster. Maximum 63 characters, alphanumeric and |
| 84 | * hyphens. |
| 85 | * @default ${app}-${stage}-${id} |
| 86 | */ |
| 87 | clusterName?: string; |
| 88 | /** |
| 89 | * The instance groups of the cluster, keyed by group name. Each group |
| 90 | * specifies its instance type, instance count, lifecycle config (S3 URI |
| 91 | * + `on_create` script) and execution role. |
| 92 | * |
| 93 | * Groups are updated in place; removing a key deletes the group from the |
| 94 | * cluster. The keys carry through to the cluster's `instanceGroups` |
| 95 | * attribute, so `hyperpod.instanceGroups.workers` is typed per key. |
| 96 | */ |
| 97 | instanceGroups?: Record<string, ClusterInstanceGroup>; |
| 98 | /** |
| 99 | * Restricted instance groups for HyperPod clusters running Amazon-managed |
| 100 | * workloads (e.g. Nova model customization), keyed by group name. |
| 101 | */ |
| 102 | restrictedInstanceGroups?: Record<string, ClusterRestrictedInstanceGroup>; |
| 103 | /** |
| 104 | * Shared environment configuration (e.g. FSx for Lustre) for restricted |
| 105 | * instance groups. |
| 106 | */ |
| 107 | restrictedInstanceGroupsConfig?: sagemaker.ClusterRestrictedInstanceGroupsConfig; |
| 108 | /** |
| 109 | * VPC to launch the cluster's nodes into. Changing the VPC replaces the |
| 110 | * cluster. |
| 111 | */ |
| 112 | vpcConfig?: sagemaker.VpcConfig; |
| 113 | /** |
| 114 | * The cluster's orchestrator. Set `{ Eks: { ClusterArn } }` to attach the |
| 115 | * HyperPod cluster to an EKS control plane; defaults to Slurm when |
| 116 | * omitted. |
| 117 | */ |
| 118 | orchestrator?: sagemaker.ClusterOrchestrator; |
| 119 | /** |
| 120 | * Whether HyperPod automatically replaces faulty nodes. |
| 121 | * @default "Automatic" |
| 122 | */ |
| 123 | nodeRecovery?: sagemaker.ClusterNodeRecovery; |
| 124 | /** |
| 125 | * Managed tiered storage (memory + local NVMe) configuration for |
| 126 | * checkpointing. |
| 127 | */ |
| 128 | tieredStorageConfig?: sagemaker.ClusterTieredStorageConfig; |
| 129 | /** |
| 130 | * Set to `"Continuous"` to enable continuous node provisioning, where |
| 131 | * HyperPod provisions capacity as it becomes available instead of |
| 132 | * all-or-nothing. |
| 133 | */ |
| 134 | nodeProvisioningMode?: sagemaker.ClusterNodeProvisioningMode; |
| 135 | /** |
| 136 | * ARN of the IAM role HyperPod assumes for cluster-level operations |
| 137 | * (required for continuous provisioning and autoscaling). |
| 138 | */ |
| 139 | clusterRole?: string; |
| 140 | /** |
| 141 | * Cluster autoscaling configuration (e.g. Karpenter) for EKS-orchestrated |
| 142 | * clusters. |
| 143 | */ |
| 144 | autoScaling?: sagemaker.ClusterAutoScalingConfig; |
| 145 | /** |
| 146 | * Tags to associate with the cluster. Merged with internal Alchemy tags. |
| 147 | */ |
| 148 | tags?: Record<string, string>; |
| 149 | } |
| 150 | |
| 151 | export interface Cluster extends Resource< |
| 152 | "AWS.SageMaker.Cluster", |
| 153 | ClusterProps, |
| 154 | { |
| 155 | /** |
| 156 | * The cluster's name. |
| 157 | */ |
| 158 | clusterName: string; |
| 159 | /** |
| 160 | * ARN of the cluster. |
| 161 | */ |
| 162 | clusterArn: string; |
| 163 | /** |
| 164 | * The cluster's status (`InService` once reconciled). |
| 165 | */ |
| 166 | clusterStatus: ClusterStatus; |
| 167 | /** |
| 168 | * ARN of the orchestrating EKS cluster, when EKS-orchestrated. |
| 169 | */ |
| 170 | orchestratorEksClusterArn: string | undefined; |
| 171 | /** |
| 172 | * The cluster's instance groups, keyed by group name. Pin a |
| 173 | * Kubernetes workload onto a group through the resource graph with |
| 174 | * its node selector: |
| 175 | * `hyperpod.instanceGroups.workers.nodeSelector`. |
| 176 | */ |
| 177 | instanceGroups: Record<string, ClusterInstanceGroupRef>; |
| 178 | }, |
| 179 | never, |
| 180 | Providers |
| 181 | > {} |
| 182 | |
| 183 | const ClusterResource = Resource<Cluster>("AWS.SageMaker.Cluster"); |
| 184 | |
| 185 | /** |
| 186 | * A `Cluster` narrowed to the instance-group keys declared in its props — |
| 187 | * `hyperpod.instanceGroups.workers` is typed per key, and a typo'd key is |
| 188 | * a compile error. |
| 189 | */ |
| 190 | export type ClusterOf<Groups> = Omit<Cluster, "instanceGroups"> & { |
| 191 | instanceGroups: Output.ObjectExpr< |
| 192 | { [K in keyof Groups]: ClusterInstanceGroupRef }, |
| 193 | never |
| 194 | >; |
| 195 | }; |
| 196 | |
| 197 | /** |
| 198 | * An Amazon SageMaker HyperPod cluster — a resilient, persistent cluster of |
| 199 | * ML compute for distributed training and inference, orchestrated by Slurm |
| 200 | * or EKS, with automatic faulty-node recovery and deep health checks. |
| 201 | * |
| 202 | * Provisioning a HyperPod cluster takes 10–25 minutes; instance groups are |
| 203 | * updated in place and removing a group from `instanceGroups` deletes it |
| 204 | * from the cluster. |
| 205 | * ### Creating Clusters |
| 206 | * **Example:** Slurm-Orchestrated Cluster |
| 207 | * ```typescript |
| 208 | * import * as AWS from "alchemy/AWS"; |
| 209 | * |
| 210 | * const cluster = yield* AWS.SageMaker.Cluster("TrainingCluster", { |
| 211 | * instanceGroups: { |
| 212 | * controller: { |
| 213 | * InstanceType: "ml.t3.medium", |
| 214 | * InstanceCount: 1, |
| 215 | * ExecutionRole: role.roleArn, |
| 216 | * LifeCycleConfig: { |
| 217 | * SourceS3Uri: `s3://${bucket.bucketName}/lifecycle`, |
| 218 | * OnCreate: "on_create.sh", |
| 219 | * }, |
| 220 | * }, |
| 221 | * }, |
| 222 | * }); |
| 223 | * ``` |
| 224 | * |
| 225 | * **Example:** EKS-Orchestrated Cluster |
| 226 | * ```typescript |
| 227 | * // The EKS cluster must use the `API` (or `API_AND_CONFIG_MAP`) |
| 228 | * // authentication mode — pass `accessConfig` explicitly, EKS's own |
| 229 | * // CONFIG_MAP default is rejected. LifeCycleConfig is required for |
| 230 | * // EKS-orchestrated instance groups too. |
| 231 | * const hyperpod = yield* AWS.SageMaker.Cluster("EksHyperPod", { |
| 232 | * orchestrator: { Eks: { ClusterArn: eksCluster.clusterArn } }, |
| 233 | * vpcConfig: { |
| 234 | * SecurityGroupIds: [securityGroupId], |
| 235 | * Subnets: network.privateSubnetIds, |
| 236 | * }, |
| 237 | * instanceGroups: { |
| 238 | * workers: { |
| 239 | * InstanceType: "ml.g5.xlarge", |
| 240 | * InstanceCount: 2, |
| 241 | * ExecutionRole: role.roleArn, |
| 242 | * LifeCycleConfig: { |
| 243 | * SourceS3Uri: `s3://${bucket.bucketName}/lifecycle`, |
| 244 | * OnCreate: "on_create.sh", |
| 245 | * }, |
| 246 | * }, |
| 247 | * }, |
| 248 | * nodeRecovery: "Automatic", |
| 249 | * }); |
| 250 | * |
| 251 | * // The keys carry through to the attributes — typed per key: |
| 252 | * const workers = hyperpod.instanceGroups.workers; |
| 253 | * ``` |
| 254 | * |
| 255 | * ### Running Workloads (Slurm) |
| 256 | * **Example:** Submit jobs from the login node over SSM |
| 257 | * ```sh |
| 258 | * # Slurm jobs are submitted on the cluster itself. Each node is an SSM |
| 259 | * # target named sagemaker-cluster:<cluster-id>_<instance-group>-<instance-id> |
| 260 | * # (list nodes with `aws sagemaker list-cluster-nodes`). |
| 261 | * aws ssm start-session \ |
| 262 | * --target sagemaker-cluster:6wl4at0i68c6_controller-i-0123456789abcdef0 |
| 263 | * # then, on the node: |
| 264 | * sbatch --nodes=4 train.sbatch |
| 265 | * ``` |
| 266 | * |
| 267 | * ### Running Workloads (EKS) |
| 268 | * **Example:** Low level: apply any Kubernetes manifest to the orchestrator |
| 269 | * ```typescript |
| 270 | * // HyperPod nodes are ordinary EKS nodes — target them from a raw |
| 271 | * // manifest (a PyTorchJob CRD, a batch/v1 Job, ...) with the well-known |
| 272 | * // node labels. |
| 273 | * const job = yield* AWS.EKS.Manifest("RawTrainJob", { |
| 274 | * cluster: eksCluster, |
| 275 | * manifest: { |
| 276 | * apiVersion: "batch/v1", |
| 277 | * kind: "Job", |
| 278 | * metadata: { name: "raw-train", namespace: "default" }, |
| 279 | * spec: { |
| 280 | * template: { |
| 281 | * spec: { |
| 282 | * nodeSelector: { |
| 283 | * "sagemaker.amazonaws.com/node-health-status": "Schedulable", |
| 284 | * "sagemaker.amazonaws.com/instance-group-name": "workers", |
| 285 | * }, |
| 286 | * containers: [{ name: "train", image: "ghcr.io/acme/train:v3" }], |
| 287 | * restartPolicy: "Never", |
| 288 | * }, |
| 289 | * }, |
| 290 | * }, |
| 291 | * }, |
| 292 | * }); |
| 293 | * ``` |
| 294 | * |
| 295 | * **Example:** High level: an effectful Job pinned to HyperPod nodes |
| 296 | * ```typescript |
| 297 | * // Kubernetes.Job / Kubernetes.Deployment run on HyperPod via the |
| 298 | * // orchestrating EKS cluster in plain Kubernetes vocabulary — the |
| 299 | * // HyperPod resources expose the derived values as attributes: the |
| 300 | * // group's `nodeSelector`, the quota's governed `namespace` and Kueue |
| 301 | * // `queueName`. |
| 302 | * const evaluate = yield* Kubernetes.Job( |
| 303 | * "Evaluate", |
| 304 | * { |
| 305 | * cluster: eksCluster, |
| 306 | * main: import.meta.url, |
| 307 | * namespace: quota.namespace, |
| 308 | * labels: { |
| 309 | * [AWS.SageMaker.KUEUE_QUEUE_NAME_LABEL]: quota.queueName, |
| 310 | * [AWS.SageMaker.KUEUE_PRIORITY_CLASS_LABEL]: "training-priority", |
| 311 | * }, |
| 312 | * podTemplate: { |
| 313 | * spec: { |
| 314 | * nodeSelector: hyperpod.instanceGroups.workers.nodeSelector, |
| 315 | * }, |
| 316 | * }, |
| 317 | * }, |
| 318 | * Effect.gen(function* () { |
| 319 | * const putItem = yield* AWS.DynamoDB.PutItem(resultsTable); |
| 320 | * return { |
| 321 | * run: Effect.gen(function* () { |
| 322 | * // evaluation logic; bindings land IAM on the pod-identity role |
| 323 | * }), |
| 324 | * }; |
| 325 | * }).pipe(Effect.provide(AWS.DynamoDB.PutItemHttp)), |
| 326 | * ); |
| 327 | * ``` |
| 328 | * |
| 329 | * ### Task Governance |
| 330 | * **Example:** Prioritize workloads with a scheduler policy and team quotas |
| 331 | * ```typescript |
| 332 | * // Requires the amazon-sagemaker-hyperpod-taskgovernance EKS add-on. |
| 333 | * const policy = yield* AWS.SageMaker.ClusterSchedulerConfig("Scheduler", { |
| 334 | * clusterArn: hyperpod.clusterArn, |
| 335 | * schedulerConfig: { |
| 336 | * PriorityClasses: [{ Name: "training", Weight: 90 }], |
| 337 | * FairShare: "Enabled", |
| 338 | * }, |
| 339 | * }); |
| 340 | * |
| 341 | * // Creates the hyperpod-ns-research namespace + Kueue LocalQueue — |
| 342 | * // exposed as `quota.namespace` / `quota.queueName` for governed |
| 343 | * // Kubernetes workloads to reference. |
| 344 | * const quota = yield* AWS.SageMaker.ComputeQuota("ResearchQuota", { |
| 345 | * clusterArn: hyperpod.clusterArn, |
| 346 | * computeQuotaTarget: { TeamName: "research", FairShareWeight: 10 }, |
| 347 | * computeQuotaConfig: { |
| 348 | * ComputeQuotaResources: [ |
| 349 | * { InstanceType: "ml.g5.xlarge", Count: 1 }, |
| 350 | * ], |
| 351 | * }, |
| 352 | * }); |
| 353 | * ``` |
| 354 | * |
| 355 | * @resource |
| 356 | */ |
| 357 | export const Cluster: { |
| 358 | < |
| 359 | const Props extends { |
| 360 | [prop in keyof ClusterProps]: Input<ClusterProps[prop]>; |
| 361 | }, |
| 362 | >( |
| 363 | id: string, |
| 364 | props: Props | Effect.Effect<Props>, |
| 365 | ): Effect.Effect< |
| 366 | ClusterOf<NonNullable<Props["instanceGroups"]>>, |
| 367 | never, |
| 368 | Providers |
| 369 | >; |
| 370 | } & typeof ClusterResource = ClusterResource as never; |
| 371 | |
| 372 | const createClusterName = ( |
| 373 | id: string, |
| 374 | props: { clusterName?: string | undefined }, |
| 375 | ) => |
| 376 | props.clusterName |
| 377 | ? Effect.succeed(props.clusterName) |
| 378 | : createPhysicalName({ id, maxLength: 63 }); |
| 379 | |
| 380 | const describeClusterOrUndefined = (name: string) => |
| 381 | sagemaker |
| 382 | .describeCluster({ ClusterName: name }) |
| 383 | .pipe(Effect.catchTag("ResourceNotFound", () => Effect.succeed(undefined))); |
| 384 | |
| 385 | const fetchClusterTags = Effect.fn(function* (arn: string) { |
| 386 | const response = yield* sagemaker |
| 387 | .listTags({ ResourceArn: arn }) |
| 388 | .pipe( |
| 389 | Effect.catchTag("AccessDeniedException", () => Effect.succeed(undefined)), |
| 390 | ); |
| 391 | return Object.fromEntries( |
| 392 | (response?.Tags ?? []).flatMap((tag) => |
| 393 | tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [], |
| 394 | ), |
| 395 | ); |
| 396 | }); |
| 397 | |
| 398 | const toAttrs = ( |
| 399 | described: sagemaker.DescribeClusterResponse, |
| 400 | ): Cluster["Attributes"] => ({ |
| 401 | clusterName: described.ClusterName ?? "", |
| 402 | clusterArn: described.ClusterArn, |
| 403 | clusterStatus: described.ClusterStatus, |
| 404 | orchestratorEksClusterArn: described.Orchestrator?.Eks?.ClusterArn, |
| 405 | instanceGroups: Object.fromEntries( |
| 406 | (described.InstanceGroups ?? []).flatMap((group) => |
| 407 | group.InstanceGroupName !== undefined |
| 408 | ? [ |
| 409 | [ |
| 410 | group.InstanceGroupName, |
| 411 | { |
| 412 | InstanceGroupName: group.InstanceGroupName, |
| 413 | InstanceType: group.InstanceType, |
| 414 | CurrentCount: group.CurrentCount, |
| 415 | TargetCount: group.TargetCount, |
| 416 | nodeSelector: { |
| 417 | [HYPERPOD_NODE_HEALTH_LABEL]: "Schedulable", |
| 418 | [HYPERPOD_INSTANCE_GROUP_LABEL]: group.InstanceGroupName, |
| 419 | }, |
| 420 | }, |
| 421 | ], |
| 422 | ] |
| 423 | : [], |
| 424 | ), |
| 425 | ), |
| 426 | }); |
| 427 | |
| 428 | /** |
| 429 | * The cluster is still transitioning toward the awaited state — retried by |
| 430 | * the bounded wait schedule. |
| 431 | */ |
| 432 | class ClusterNotReady extends Data.TaggedError("ClusterNotReady")<{ |
| 433 | readonly clusterName: string; |
| 434 | readonly status: string | undefined; |
| 435 | }> {} |
| 436 | |
| 437 | /** |
| 438 | * The cluster converged to the terminal `Failed` status. |
| 439 | */ |
| 440 | export class ClusterFailed extends Data.TaggedError("ClusterFailed")<{ |
| 441 | readonly clusterName: string; |
| 442 | readonly message: string | undefined; |
| 443 | }> {} |
| 444 | |
| 445 | /** |
| 446 | * The cluster is stuck in `Deleting` because SageMaker's internal teardown |
| 447 | * cannot proceed — e.g. the instance group's execution role (which HyperPod |
| 448 | * assumes to delete node ENIs) was deleted or lost its |
| 449 | * `ec2:DeleteNetworkInterface` / `ec2:DeleteNetworkInterfacePermission` |
| 450 | * permissions mid-teardown. Retrying `deleteCluster` can never fix this; |
| 451 | * the role and its permissions must be restored out-of-band, after which |
| 452 | * SageMaker's own retry completes the deletion within a minute. |
| 453 | */ |
| 454 | export class ClusterTeardownBlocked extends Data.TaggedError( |
| 455 | "ClusterTeardownBlocked", |
| 456 | )<{ |
| 457 | readonly clusterName: string; |
| 458 | readonly message: string; |
| 459 | }> {} |
| 460 | |
| 461 | // Explicitly-typed retry wrapper — an inline `Effect.retry` in provider |
| 462 | // lifecycle code leaks `Retry.Return`'s conditional type into declaration |
| 463 | // emit and widens the provider layer to `unknown` for every consumer of |
| 464 | // `AWS.providers()`. |
| 465 | const retryWhileNotReady = <A, E extends { readonly _tag: string }, R>( |
| 466 | self: Effect.Effect<A, E, R>, |
| 467 | ): Effect.Effect<A, E, R> => |
| 468 | Effect.retry(self, { |
| 469 | while: (e) => e._tag === "ClusterNotReady", |
| 470 | // HyperPod cluster provisioning routinely takes 10–25 minutes. |
| 471 | // Poll 15s up to ~35 min. |
| 472 | schedule: Schedule.max([ |
| 473 | Schedule.spaced("15 seconds"), |
| 474 | Schedule.recurs(140), |
| 475 | ]), |
| 476 | }); |
| 477 | |
| 478 | // A freshly created execution role isn't assumable by SageMaker for a few |
| 479 | // seconds (IAM propagation) — createCluster rejects it with a |
| 480 | // ValidationException until it is. |
| 481 | const retryWhileRoleUnassumable = < |
| 482 | A, |
| 483 | E extends { readonly _tag: string; readonly message?: string }, |
| 484 | R, |
| 485 | >( |
| 486 | self: Effect.Effect<A, E, R>, |
| 487 | ): Effect.Effect<A, E, R> => |
| 488 | Effect.retry(self, { |
| 489 | while: (e): boolean => |
| 490 | e._tag === "ValidationException" && |
| 491 | (e.message?.includes("cannot assume the execution role") ?? false), |
| 492 | schedule: Schedule.max([Schedule.spaced("5 seconds"), Schedule.recurs(12)]), |
| 493 | }); |
| 494 | |
| 495 | const waitForCluster = (name: string, target: "InService" | "Gone") => |
| 496 | retryWhileNotReady( |
| 497 | Effect.gen(function* () { |
| 498 | const described = yield* describeClusterOrUndefined(name); |
| 499 | if (target === "Gone") { |
| 500 | if (described === undefined) return; |
| 501 | // SageMaker's internal teardown deletes node ENIs by assuming the |
| 502 | // instance group's execution role. If that role (or its EC2 |
| 503 | // permissions) was deleted mid-teardown, every internal retry |
| 504 | // fails and the cluster stays `Deleting` forever — surface the |
| 505 | // node's permission failure immediately instead of polling the |
| 506 | // full wait budget on a delete that can never finish. |
| 507 | const nodes = yield* sagemaker |
| 508 | .listClusterNodes({ ClusterName: name }) |
| 509 | .pipe( |
| 510 | Effect.catchTag("ResourceNotFound", () => |
| 511 | Effect.succeed(undefined), |
| 512 | ), |
| 513 | ); |
| 514 | const blocked = (nodes?.ClusterNodeSummaries ?? []).find((node) => |
| 515 | node.InstanceStatus?.Message?.includes( |
| 516 | "does not have permission to perform", |
| 517 | ), |
| 518 | ); |
| 519 | if (blocked !== undefined) { |
| 520 | return yield* Effect.fail( |
| 521 | new ClusterTeardownBlocked({ |
| 522 | clusterName: name, |
| 523 | message: blocked.InstanceStatus?.Message ?? "", |
| 524 | }), |
| 525 | ); |
| 526 | } |
| 527 | return yield* Effect.fail( |
| 528 | new ClusterNotReady({ |
| 529 | clusterName: name, |
| 530 | status: described.ClusterStatus, |
| 531 | }), |
| 532 | ); |
| 533 | } |
| 534 | if (described?.ClusterStatus === "InService") return; |
| 535 | if (described?.ClusterStatus === "Failed") { |
| 536 | return yield* Effect.fail( |
| 537 | new ClusterFailed({ |
| 538 | clusterName: name, |
| 539 | message: described.FailureMessage, |
| 540 | }), |
| 541 | ); |
| 542 | } |
| 543 | return yield* Effect.fail( |
| 544 | new ClusterNotReady({ |
| 545 | clusterName: name, |
| 546 | status: described?.ClusterStatus, |
| 547 | }), |
| 548 | ); |
| 549 | }), |
| 550 | ); |
| 551 | |
| 552 | /** |
| 553 | * The subset of an instance group's specification that `updateCluster` |
| 554 | * can change and `describeCluster` reports back — used to diff desired |
| 555 | * groups against observed groups. |
| 556 | */ |
| 557 | const observedGroupIdentity = ( |
| 558 | group: sagemaker.DescribeClusterResponse["InstanceGroups"][number], |
| 559 | ) => ({ |
| 560 | count: group.TargetCount ?? group.CurrentCount, |
| 561 | type: group.InstanceType, |
| 562 | }); |
| 563 | |
| 564 | const desiredGroupIdentity = (group: ClusterInstanceGroup) => ({ |
| 565 | count: group.InstanceCount, |
| 566 | type: group.InstanceType, |
| 567 | }); |
| 568 | |
| 569 | /** Fold the keyed props form into the API's named-array form. */ |
| 570 | const toGroupSpecs = <G extends { InstanceGroupName?: string }>( |
| 571 | groups: Record<string, Omit<G, "InstanceGroupName">> | undefined, |
| 572 | ): G[] | undefined => |
| 573 | groups === undefined |
| 574 | ? undefined |
| 575 | : Object.entries(groups).map( |
| 576 | ([name, group]) => ({ InstanceGroupName: name, ...group }) as G, |
| 577 | ); |
| 578 | |
| 579 | export const ClusterProvider = () => |
| 580 | Provider.effect( |
| 581 | Cluster, |
| 582 | Effect.gen(function* () { |
| 583 | return { |
| 584 | stables: ["clusterName", "clusterArn"], |
| 585 | // HyperPod's internal teardown deletes node ENIs by assuming the |
| 586 | // instance group's execution role inside the cluster's VPC — |
| 587 | // deleting the role or network mid-teardown wedges the cluster in |
| 588 | // `Deleting` permanently, so nuke must fully delete clusters |
| 589 | // before touching these types. |
| 590 | nuke: { |
| 591 | dependsOn: [ |
| 592 | "AWS.IAM.Role", |
| 593 | "AWS.IAM.Policy", |
| 594 | "AWS.IAM.InstanceProfile", |
| 595 | "AWS.EC2.Vpc", |
| 596 | "AWS.EC2.Subnet", |
| 597 | "AWS.EC2.SecurityGroup", |
| 598 | ], |
| 599 | }, |
| 600 | list: () => |
| 601 | Effect.gen(function* () { |
| 602 | const summaries = yield* sagemaker.listClusters.pages({}).pipe( |
| 603 | EffectStream.runCollect, |
| 604 | Effect.map((chunk) => |
| 605 | Array.from(chunk).flatMap( |
| 606 | (page) => page.ClusterSummaries ?? [], |
| 607 | ), |
| 608 | ), |
| 609 | ); |
| 610 | return summaries.flatMap((s) => |
| 611 | s.ClusterName !== undefined && s.ClusterArn !== undefined |
| 612 | ? [ |
| 613 | { |
| 614 | clusterName: s.ClusterName, |
| 615 | clusterArn: s.ClusterArn, |
| 616 | clusterStatus: s.ClusterStatus ?? "InService", |
| 617 | orchestratorEksClusterArn: undefined, |
| 618 | instanceGroups: {}, |
| 619 | }, |
| 620 | ] |
| 621 | : [], |
| 622 | ); |
| 623 | }), |
| 624 | read: Effect.fn(function* ({ id, olds, output }) { |
| 625 | const name = |
| 626 | output?.clusterName ?? (yield* createClusterName(id, olds ?? {})); |
| 627 | const described = yield* describeClusterOrUndefined(name); |
| 628 | if (!described || described.ClusterStatus === "Deleting") { |
| 629 | return undefined; |
| 630 | } |
| 631 | const attrs = toAttrs(described); |
| 632 | const tags = yield* fetchClusterTags(attrs.clusterArn); |
| 633 | return (yield* hasAlchemyTags(id, tags as Tags)) |
| 634 | ? attrs |
| 635 | : Unowned(attrs); |
| 636 | }), |
| 637 | diff: Effect.fn(function* ({ id, news, olds }) { |
| 638 | if (!isResolved(news)) return; |
| 639 | if (olds === undefined) return; |
| 640 | const oldName = yield* createClusterName(id, olds); |
| 641 | const newName = yield* createClusterName(id, news); |
| 642 | // The VPC is fixed at creation — changing it (or the name) |
| 643 | // replaces the cluster. Everything else updates in place. |
| 644 | if ( |
| 645 | oldName !== newName || |
| 646 | JSON.stringify(olds.vpcConfig) !== JSON.stringify(news.vpcConfig) |
| 647 | ) { |
| 648 | return { action: "replace" } as const; |
| 649 | } |
| 650 | }), |
| 651 | reconcile: Effect.fn(function* ({ id, news, output, session }) { |
| 652 | if (!news) { |
| 653 | return yield* Effect.fail( |
| 654 | new Error("SageMaker HyperPod Cluster requires props"), |
| 655 | ); |
| 656 | } |
| 657 | const name = |
| 658 | output?.clusterName ?? (yield* createClusterName(id, news)); |
| 659 | const internalTags = yield* createInternalTags(id); |
| 660 | const desiredTags = { ...internalTags, ...news.tags }; |
| 661 | |
| 662 | // Observe. |
| 663 | let described = yield* describeClusterOrUndefined(name); |
| 664 | |
| 665 | // Ensure — create if missing; tolerate the already-exists race. |
| 666 | if (described === undefined) { |
| 667 | yield* sagemaker |
| 668 | .createCluster({ |
| 669 | ClusterName: name, |
| 670 | InstanceGroups: toGroupSpecs(news.instanceGroups), |
| 671 | RestrictedInstanceGroups: toGroupSpecs( |
| 672 | news.restrictedInstanceGroups, |
| 673 | ), |
| 674 | RestrictedInstanceGroupsConfig: |
| 675 | news.restrictedInstanceGroupsConfig, |
| 676 | VpcConfig: news.vpcConfig, |
| 677 | Orchestrator: news.orchestrator, |
| 678 | NodeRecovery: news.nodeRecovery, |
| 679 | TieredStorageConfig: news.tieredStorageConfig, |
| 680 | NodeProvisioningMode: news.nodeProvisioningMode, |
| 681 | ClusterRole: news.clusterRole, |
| 682 | AutoScaling: news.autoScaling, |
| 683 | Tags: Object.entries(desiredTags).map(([Key, Value]) => ({ |
| 684 | Key, |
| 685 | Value, |
| 686 | })), |
| 687 | }) |
| 688 | .pipe( |
| 689 | retryWhileRoleUnassumable, |
| 690 | Effect.catchTag("ResourceInUse", () => Effect.void), |
| 691 | ); |
| 692 | yield* session.note( |
| 693 | `Creating HyperPod cluster ${name} (typically 10-25 minutes)...`, |
| 694 | ); |
| 695 | } else { |
| 696 | // Sync — diff observed instance groups + mutable settings |
| 697 | // against desired and apply one update call for the delta. |
| 698 | const observedGroups = new Map( |
| 699 | (described.InstanceGroups ?? []).flatMap((g) => |
| 700 | g.InstanceGroupName !== undefined |
| 701 | ? [[g.InstanceGroupName, observedGroupIdentity(g)]] |
| 702 | : [], |
| 703 | ), |
| 704 | ); |
| 705 | const desiredGroups = new Map( |
| 706 | Object.entries(news.instanceGroups ?? {}).map( |
| 707 | ([groupName, g]) => [groupName, desiredGroupIdentity(g)], |
| 708 | ), |
| 709 | ); |
| 710 | const groupsToDelete = [...observedGroups.keys()].filter( |
| 711 | (groupName) => !desiredGroups.has(groupName), |
| 712 | ); |
| 713 | const groupsChanged = |
| 714 | groupsToDelete.length > 0 || |
| 715 | [...desiredGroups].some( |
| 716 | ([groupName, identity]) => |
| 717 | JSON.stringify(observedGroups.get(groupName)) !== |
| 718 | JSON.stringify(identity), |
| 719 | ); |
| 720 | const settingsChanged = |
| 721 | (news.nodeRecovery !== undefined && |
| 722 | news.nodeRecovery !== described.NodeRecovery) || |
| 723 | (news.nodeProvisioningMode !== undefined && |
| 724 | news.nodeProvisioningMode !== described.NodeProvisioningMode) || |
| 725 | (news.tieredStorageConfig !== undefined && |
| 726 | JSON.stringify(news.tieredStorageConfig) !== |
| 727 | JSON.stringify(described.TieredStorageConfig)) || |
| 728 | (news.autoScaling !== undefined && |
| 729 | news.autoScaling.Mode !== described.AutoScaling?.Mode); |
| 730 | |
| 731 | if (groupsChanged || settingsChanged) { |
| 732 | yield* sagemaker |
| 733 | .updateCluster({ |
| 734 | ClusterName: name, |
| 735 | InstanceGroups: toGroupSpecs(news.instanceGroups), |
| 736 | RestrictedInstanceGroups: toGroupSpecs( |
| 737 | news.restrictedInstanceGroups, |
| 738 | ), |
| 739 | RestrictedInstanceGroupsConfig: |
| 740 | news.restrictedInstanceGroupsConfig, |
| 741 | TieredStorageConfig: news.tieredStorageConfig, |
| 742 | NodeRecovery: news.nodeRecovery, |
| 743 | InstanceGroupsToDelete: |
| 744 | groupsToDelete.length > 0 ? groupsToDelete : undefined, |
| 745 | NodeProvisioningMode: news.nodeProvisioningMode, |
| 746 | ClusterRole: news.clusterRole, |
| 747 | AutoScaling: news.autoScaling, |
| 748 | }) |
| 749 | .pipe(Effect.catchTag("ResourceNotFound", () => Effect.void)); |
| 750 | yield* session.note(`Updating HyperPod cluster ${name}...`); |
| 751 | } |
| 752 | } |
| 753 | |
| 754 | // Converge to InService (provisioning is asynchronous). |
| 755 | yield* waitForCluster(name, "InService"); |
| 756 | described = yield* describeClusterOrUndefined(name); |
| 757 | if (described === undefined) { |
| 758 | return yield* Effect.fail( |
| 759 | new Error(`failed to read reconciled HyperPod cluster ${name}`), |
| 760 | ); |
| 761 | } |
| 762 | const attrs = toAttrs(described); |
| 763 | |
| 764 | // Sync tags — diff against OBSERVED cloud tags. |
| 765 | const currentTags = yield* fetchClusterTags(attrs.clusterArn); |
| 766 | const { removed, upsert } = diffTags(currentTags, desiredTags); |
| 767 | if (removed.length > 0) { |
| 768 | yield* sagemaker.deleteTags({ |
| 769 | ResourceArn: attrs.clusterArn, |
| 770 | TagKeys: removed, |
| 771 | }); |
| 772 | } |
| 773 | if (upsert.length > 0) { |
| 774 | yield* sagemaker.addTags({ |
| 775 | ResourceArn: attrs.clusterArn, |
| 776 | Tags: upsert.map(({ Key, Value }) => ({ Key, Value })), |
| 777 | }); |
| 778 | } |
| 779 | |
| 780 | yield* session.note(attrs.clusterArn); |
| 781 | return attrs; |
| 782 | }), |
| 783 | delete: Effect.fn(function* ({ output }) { |
| 784 | const described = yield* describeClusterOrUndefined( |
| 785 | output.clusterName, |
| 786 | ); |
| 787 | if (described === undefined) return; |
| 788 | // A cluster mid-create/update rejects deletion with a Conflict — |
| 789 | // wait for it to settle first, then delete. A cluster stuck in |
| 790 | // `Failed` (or one that never converges) is still deletable. One |
| 791 | // already `Deleting` (a previous attempt) skips straight to the |
| 792 | // Gone wait instead of burning the full InService wait budget. |
| 793 | if (described.ClusterStatus !== "Deleting") { |
| 794 | yield* waitForCluster(output.clusterName, "InService").pipe( |
| 795 | Effect.catchTag( |
| 796 | ["ClusterFailed", "ClusterNotReady"], |
| 797 | () => Effect.void, |
| 798 | ), |
| 799 | ); |
| 800 | yield* sagemaker |
| 801 | .deleteCluster({ ClusterName: output.clusterName }) |
| 802 | .pipe( |
| 803 | Effect.catchTag("ResourceNotFound", () => Effect.void), |
| 804 | Effect.catchTag("ConflictException", () => Effect.void), |
| 805 | ); |
| 806 | } |
| 807 | yield* waitForCluster(output.clusterName, "Gone"); |
| 808 | }), |
| 809 | }; |
| 810 | }), |
| 811 | ); |