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 SchedulerResourceStatus = sagemaker.SchedulerResourceStatus;
20
21export interface ClusterSchedulerConfigProps {
22 /**
23 * Name of the cluster policy. Maximum 63 characters, alphanumeric and
24 * hyphens.
25 * @default ${app}-${stage}-${id}
26 */
27 name?: string;
28 /**
29 * ARN of the EKS-orchestrated HyperPod cluster the policy applies to.
30 * Changing the cluster replaces the policy. AWS allows ONE cluster
31 * policy per cluster — creating a second fails with the typed
32 * `ClusterSchedulerConfigAlreadyExists` error.
33 */
34 clusterArn: string;
35 /**
36 * The policy itself: task priority classes, fair-share allocation, and
37 * idle-resource sharing.
38 */
39 schedulerConfig?: sagemaker.SchedulerConfig;
40 /**
41 * A description of the cluster policy.
42 */
43 description?: string;
44 /**
45 * Tags to associate with the cluster policy. Merged with internal Alchemy
46 * tags.
47 */
48 tags?: Record<string, string>;
49}
50
51export interface ClusterSchedulerConfig extends Resource<
52 "AWS.SageMaker.ClusterSchedulerConfig",
53 ClusterSchedulerConfigProps,
54 {
55 /**
56 * The cluster policy's ID.
57 */
58 clusterSchedulerConfigId: string;
59 /**
60 * ARN of the cluster policy.
61 */
62 clusterSchedulerConfigArn: string;
63 /**
64 * The cluster policy's name.
65 */
66 name: string;
67 /**
68 * ARN of the HyperPod cluster the policy applies to.
69 */
70 clusterArn: string;
71 /**
72 * The policy's current version (incremented on every update).
73 */
74 clusterSchedulerConfigVersion: number;
75 },
76 never,
77 Providers
78> {}
79
80/**
81 * A SageMaker HyperPod cluster policy (task governance) — configures how an
82 * EKS-orchestrated HyperPod cluster prioritizes tasks and allocates idle
83 * compute across teams via priority classes and fair-share weights.
84 * ### Creating Cluster Policies
85 * **Example:** Priority Classes with Fair-Share
86 * ```typescript
87 * import * as AWS from "alchemy/AWS";
88 *
89 * const policy = yield* AWS.SageMaker.ClusterSchedulerConfig("Scheduler", {
90 * clusterArn: hyperpod.clusterArn,
91 * schedulerConfig: {
92 * PriorityClasses: [
93 * { Name: "inference", Weight: 100 },
94 * { Name: "training", Weight: 75 },
95 * ],
96 * FairShare: "Enabled",
97 * },
98 * description: "Prioritize inference over training",
99 * });
100 * ```
101 *
102 * @resource
103 */
104export const ClusterSchedulerConfig = Resource<ClusterSchedulerConfig>(
105 "AWS.SageMaker.ClusterSchedulerConfig",
106);
107
108const createConfigName = (id: string, props: { name?: string | undefined }) =>
109 props.name
110 ? Effect.succeed(props.name)
111 : createPhysicalName({ id, maxLength: 63 });
112
113const describeConfigOrUndefined = (configId: string) =>
114 sagemaker
115 .describeClusterSchedulerConfig({ ClusterSchedulerConfigId: configId })
116 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.succeed(undefined)));
117
118/**
119 * Look a cluster policy up by exact name — used when state was lost (read
120 * without output) or a create raced.
121 */
122const findConfigByName = Effect.fn(function* (name: string) {
123 const summaries = yield* sagemaker.listClusterSchedulerConfigs
124 .pages({ NameContains: name })
125 .pipe(
126 EffectStream.runCollect,
127 Effect.map((chunk) =>
128 Array.from(chunk).flatMap(
129 (page) => page.ClusterSchedulerConfigSummaries ?? [],
130 ),
131 ),
132 );
133 return summaries.find((s) => s.Name === name && s.Status !== "Deleted");
134});
135
136const fetchConfigTags = Effect.fn(function* (arn: string) {
137 const response = yield* sagemaker
138 .listTags({ ResourceArn: arn })
139 .pipe(
140 Effect.catchTag("AccessDeniedException", () => Effect.succeed(undefined)),
141 );
142 return Object.fromEntries(
143 (response?.Tags ?? []).flatMap((tag) =>
144 tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [],
145 ),
146 );
147});
148
149const toAttrs = (
150 described: sagemaker.DescribeClusterSchedulerConfigResponse,
151): ClusterSchedulerConfig["Attributes"] => ({
152 clusterSchedulerConfigId: described.ClusterSchedulerConfigId,
153 clusterSchedulerConfigArn: described.ClusterSchedulerConfigArn,
154 name: described.Name,
155 clusterArn: described.ClusterArn ?? "",
156 clusterSchedulerConfigVersion: described.ClusterSchedulerConfigVersion,
157});
158
159/**
160 * The cluster policy is still transitioning toward the awaited state —
161 * retried by the bounded wait schedule.
162 */
163class SchedulerConfigNotReady extends Data.TaggedError(
164 "SchedulerConfigNotReady",
165)<{
166 readonly configId: string;
167 readonly status: string | undefined;
168}> {}
169
170/**
171 * The cluster policy converged to a terminal failed status.
172 */
173export class SchedulerConfigFailed extends Data.TaggedError(
174 "SchedulerConfigFailed",
175)<{
176 readonly configId: string;
177 readonly status: string | undefined;
178 readonly message: string | undefined;
179}> {}
180
181const FAILED_STATUSES: string[] = [
182 "CreateFailed",
183 "CreateRollbackFailed",
184 "UpdateFailed",
185 "UpdateRollbackFailed",
186 "DeleteFailed",
187 "DeleteRollbackFailed",
188];
189
190// Explicitly-typed retry wrapper — an inline `Effect.retry` in provider
191// lifecycle code leaks `Retry.Return`'s conditional type into declaration
192// emit and widens the provider layer to `unknown` for every consumer of
193// `AWS.providers()`.
194const retryWhileNotReady = <A, E extends { readonly _tag: string }, R>(
195 self: Effect.Effect<A, E, R>,
196): Effect.Effect<A, E, R> =>
197 Effect.retry(self, {
198 while: (e) => e._tag === "SchedulerConfigNotReady",
199 schedule: Schedule.max([Schedule.spaced("5 seconds"), Schedule.recurs(60)]),
200 });
201
202const waitForConfig = (configId: string, target: "Ready" | "Gone") =>
203 retryWhileNotReady(
204 Effect.gen(function* () {
205 const described = yield* describeConfigOrUndefined(configId);
206 if (target === "Gone") {
207 if (described === undefined || described.Status === "Deleted") return;
208 if (FAILED_STATUSES.includes(described.Status)) {
209 return yield* Effect.fail(
210 new SchedulerConfigFailed({
211 configId,
212 status: described.Status,
213 message: described.FailureReason,
214 }),
215 );
216 }
217 return yield* Effect.fail(
218 new SchedulerConfigNotReady({ configId, status: described.Status }),
219 );
220 }
221 if (described?.Status === "Created" || described?.Status === "Updated") {
222 return;
223 }
224 if (
225 described !== undefined &&
226 FAILED_STATUSES.includes(described.Status)
227 ) {
228 return yield* Effect.fail(
229 new SchedulerConfigFailed({
230 configId,
231 status: described.Status,
232 message: described.FailureReason,
233 }),
234 );
235 }
236 return yield* Effect.fail(
237 new SchedulerConfigNotReady({ configId, status: described?.Status }),
238 );
239 }),
240 );
241
242export const ClusterSchedulerConfigProvider = () =>
243 Provider.effect(
244 ClusterSchedulerConfig,
245 Effect.gen(function* () {
246 return {
247 stables: [
248 "clusterSchedulerConfigId",
249 "clusterSchedulerConfigArn",
250 "name",
251 "clusterArn",
252 ],
253 list: () =>
254 Effect.gen(function* () {
255 const summaries = yield* sagemaker.listClusterSchedulerConfigs
256 .pages({})
257 .pipe(
258 EffectStream.runCollect,
259 Effect.map((chunk) =>
260 Array.from(chunk).flatMap(
261 (page) => page.ClusterSchedulerConfigSummaries ?? [],
262 ),
263 ),
264 );
265 return summaries.flatMap((s) =>
266 s.ClusterSchedulerConfigId !== undefined && s.Status !== "Deleted"
267 ? [
268 {
269 clusterSchedulerConfigId: s.ClusterSchedulerConfigId,
270 clusterSchedulerConfigArn:
271 s.ClusterSchedulerConfigArn ?? "",
272 name: s.Name ?? "",
273 clusterArn: s.ClusterArn ?? "",
274 clusterSchedulerConfigVersion:
275 s.ClusterSchedulerConfigVersion ?? 1,
276 },
277 ]
278 : [],
279 );
280 }),
281 read: Effect.fn(function* ({ id, olds, output }) {
282 const configId =
283 output?.clusterSchedulerConfigId ??
284 (yield* findConfigByName(yield* createConfigName(id, olds ?? {})))
285 ?.ClusterSchedulerConfigId;
286 if (configId === undefined) return undefined;
287 const described = yield* describeConfigOrUndefined(configId);
288 if (
289 !described ||
290 described.Status === "Deleting" ||
291 described.Status === "Deleted"
292 ) {
293 return undefined;
294 }
295 const attrs = toAttrs(described);
296 const tags = yield* fetchConfigTags(attrs.clusterSchedulerConfigArn);
297 return (yield* hasAlchemyTags(id, tags as Tags))
298 ? attrs
299 : Unowned(attrs);
300 }),
301 diff: Effect.fn(function* ({ id, news, olds }) {
302 if (!isResolved(news)) return;
303 if (olds === undefined) return;
304 const oldName = yield* createConfigName(id, olds);
305 const newName = yield* createConfigName(id, news);
306 // The name and target cluster are fixed at creation.
307 if (oldName !== newName || olds.clusterArn !== news.clusterArn) {
308 return { action: "replace" } as const;
309 }
310 }),
311 reconcile: Effect.fn(function* ({ id, news, output, session }) {
312 if (!news) {
313 return yield* Effect.fail(
314 new Error("SageMaker ClusterSchedulerConfig requires props"),
315 );
316 }
317 const name = yield* createConfigName(id, news);
318 const internalTags = yield* createInternalTags(id);
319 const desiredTags = { ...internalTags, ...news.tags };
320
321 // Observe — by cached id first, then by name (lost state or race).
322 let configId = output?.clusterSchedulerConfigId;
323 let described =
324 configId !== undefined
325 ? yield* describeConfigOrUndefined(configId)
326 : undefined;
327 if (described === undefined) {
328 const found = yield* findConfigByName(name);
329 described =
330 found?.ClusterSchedulerConfigId !== undefined
331 ? yield* describeConfigOrUndefined(
332 found.ClusterSchedulerConfigId,
333 )
334 : undefined;
335 }
336
337 // Ensure — create if missing; a Conflict means a concurrent
338 // create won the race, so re-observe by name.
339 if (described === undefined) {
340 const created = yield* sagemaker
341 .createClusterSchedulerConfig({
342 Name: name,
343 ClusterArn: news.clusterArn,
344 SchedulerConfig: news.schedulerConfig,
345 Description: news.description,
346 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
347 Key,
348 Value,
349 })),
350 })
351 .pipe(
352 Effect.catchTag("ConflictException", () =>
353 Effect.succeed(undefined),
354 ),
355 );
356 configId =
357 created?.ClusterSchedulerConfigId ??
358 (yield* findConfigByName(name))?.ClusterSchedulerConfigId;
359 if (configId === undefined) {
360 return yield* Effect.fail(
361 new Error(`failed to create cluster scheduler config ${name}`),
362 );
363 }
364 yield* session.note(`Creating cluster policy ${name}...`);
365 yield* waitForConfig(configId, "Ready");
366 } else {
367 configId = described.ClusterSchedulerConfigId;
368 // Wait out any in-flight transition before diffing.
369 yield* waitForConfig(configId, "Ready");
370 described = yield* describeConfigOrUndefined(configId);
371 // Sync — diff observed policy against desired.
372 if (
373 described !== undefined &&
374 (JSON.stringify(described.SchedulerConfig) !==
375 JSON.stringify(news.schedulerConfig) ||
376 (described.Description ?? undefined) !==
377 (news.description ?? undefined))
378 ) {
379 yield* sagemaker.updateClusterSchedulerConfig({
380 ClusterSchedulerConfigId: configId,
381 TargetVersion: described.ClusterSchedulerConfigVersion,
382 SchedulerConfig: news.schedulerConfig,
383 Description: news.description,
384 });
385 yield* session.note(`Updating cluster policy ${name}...`);
386 yield* waitForConfig(configId, "Ready");
387 }
388 }
389
390 described = yield* describeConfigOrUndefined(configId);
391 if (described === undefined) {
392 return yield* Effect.fail(
393 new Error(
394 `failed to read reconciled cluster scheduler config ${name}`,
395 ),
396 );
397 }
398 const attrs = toAttrs(described);
399
400 // Sync tags — diff against OBSERVED cloud tags.
401 const currentTags = yield* fetchConfigTags(
402 attrs.clusterSchedulerConfigArn,
403 );
404 const { removed, upsert } = diffTags(currentTags, desiredTags);
405 if (removed.length > 0) {
406 yield* sagemaker.deleteTags({
407 ResourceArn: attrs.clusterSchedulerConfigArn,
408 TagKeys: removed,
409 });
410 }
411 if (upsert.length > 0) {
412 yield* sagemaker.addTags({
413 ResourceArn: attrs.clusterSchedulerConfigArn,
414 Tags: upsert.map(({ Key, Value }) => ({ Key, Value })),
415 });
416 }
417
418 yield* session.note(attrs.clusterSchedulerConfigArn);
419 return attrs;
420 }),
421 delete: Effect.fn(function* ({ output }) {
422 yield* sagemaker
423 .deleteClusterSchedulerConfig({
424 ClusterSchedulerConfigId: output.clusterSchedulerConfigId,
425 })
426 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.void));
427 yield* waitForConfig(output.clusterSchedulerConfigId, "Gone");
428 }),
429 };
430 }),
431 );