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 interface ComputeQuotaProps {
20 /**
21 * Name of the compute allocation. Maximum 63 characters, alphanumeric and
22 * hyphens.
23 * @default ${app}-${stage}-${id}
24 */
25 name?: string;
26 /**
27 * ARN of the EKS-orchestrated HyperPod cluster the allocation applies to.
28 * Changing the cluster replaces the allocation.
29 */
30 clusterArn: string;
31 /**
32 * The team the quota is allocated to (its name maps to a Kubernetes
33 * namespace `hyperpod-ns-<team-name>`) and its fair-share weight.
34 */
35 computeQuotaTarget: sagemaker.ComputeQuotaTarget;
36 /**
37 * The allocation itself: instance-type quotas, the resource-sharing
38 * strategy for idle compute, and whether the team's own tasks can preempt
39 * each other.
40 */
41 computeQuotaConfig?: sagemaker.ComputeQuotaConfig;
42 /**
43 * Whether the quota is enforced.
44 * @default "Enabled"
45 */
46 activationState?: sagemaker.ActivationState;
47 /**
48 * A description of the compute allocation.
49 */
50 description?: string;
51 /**
52 * Tags to associate with the compute allocation. Merged with internal
53 * Alchemy tags.
54 */
55 tags?: Record<string, string>;
56}
57
58/**
59 * Kueue label selecting the task-governance queue. Set it on a
60 * `Kubernetes.Job` / `Kubernetes.Deployment`'s `labels` with the quota's
61 * `queueName` attribute to submit the workload through HyperPod task
62 * governance.
63 */
64export const KUEUE_QUEUE_NAME_LABEL = "kueue.x-k8s.io/queue-name";
65
66/**
67 * Kueue label selecting the task-governance priority class. Its value is
68 * `<PriorityClass name>-priority` for a class declared on the cluster's
69 * `AWS.SageMaker.ClusterSchedulerConfig`.
70 */
71export const KUEUE_PRIORITY_CLASS_LABEL = "kueue.x-k8s.io/priority-class";
72
73export interface ComputeQuota extends Resource<
74 "AWS.SageMaker.ComputeQuota",
75 ComputeQuotaProps,
76 {
77 /**
78 * The compute allocation's ID.
79 */
80 computeQuotaId: string;
81 /**
82 * ARN of the compute allocation.
83 */
84 computeQuotaArn: string;
85 /**
86 * The compute allocation's name.
87 */
88 name: string;
89 /**
90 * ARN of the HyperPod cluster the allocation applies to.
91 */
92 clusterArn: string;
93 /**
94 * The allocation's current version (incremented on every update).
95 */
96 computeQuotaVersion: number;
97 /**
98 * The team the quota is allocated to. Task governance materializes the
99 * `hyperpod-ns-<teamName>` namespace and its Kueue LocalQueue from it.
100 */
101 teamName: string;
102 /**
103 * The Kubernetes namespace task governance materializes for the team
104 * (`hyperpod-ns-<teamName>`). Pass it as a governed
105 * `Kubernetes.Job` / `Kubernetes.Deployment`'s `namespace`.
106 */
107 namespace: string;
108 /**
109 * The team's Kueue LocalQueue name
110 * (`hyperpod-ns-<teamName>-localqueue`). Set it on a governed
111 * workload's `labels` under {@link KUEUE_QUEUE_NAME_LABEL}.
112 */
113 queueName: string;
114 },
115 never,
116 Providers
117> {}
118
119/**
120 * A SageMaker HyperPod compute allocation (task governance) — reserves
121 * instance capacity on an EKS-orchestrated HyperPod cluster for a team,
122 * with fair-share weights and borrow/lend rules for idle compute.
123 * ### Creating Compute Allocations
124 * **Example:** Team Quota with Borrowing
125 * ```typescript
126 * import * as AWS from "alchemy/AWS";
127 *
128 * const quota = yield* AWS.SageMaker.ComputeQuota("ResearchQuota", {
129 * clusterArn: hyperpod.clusterArn,
130 * computeQuotaTarget: { TeamName: "research", FairShareWeight: 10 },
131 * computeQuotaConfig: {
132 * ComputeQuotaResources: [{ InstanceType: "ml.g5.xlarge", Count: 2 }],
133 * ResourceSharingConfig: { Strategy: "Lend", BorrowLimit: 50 },
134 * },
135 * });
136 * ```
137 *
138 * @resource
139 */
140export const ComputeQuota = Resource<ComputeQuota>(
141 "AWS.SageMaker.ComputeQuota",
142);
143
144const createQuotaName = (id: string, props: { name?: string | undefined }) =>
145 props.name
146 ? Effect.succeed(props.name)
147 : createPhysicalName({ id, maxLength: 63 });
148
149const describeQuotaOrUndefined = (quotaId: string) =>
150 sagemaker
151 .describeComputeQuota({ ComputeQuotaId: quotaId })
152 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.succeed(undefined)));
153
154/**
155 * Look a compute allocation up by exact name — used when state was lost
156 * (read without output) or a create raced.
157 */
158const findQuotaByName = Effect.fn(function* (name: string) {
159 const summaries = yield* sagemaker.listComputeQuotas
160 .pages({ NameContains: name })
161 .pipe(
162 EffectStream.runCollect,
163 Effect.map((chunk) =>
164 Array.from(chunk).flatMap((page) => page.ComputeQuotaSummaries ?? []),
165 ),
166 );
167 return summaries.find((s) => s.Name === name && s.Status !== "Deleted");
168});
169
170const fetchQuotaTags = Effect.fn(function* (arn: string) {
171 const response = yield* sagemaker
172 .listTags({ ResourceArn: arn })
173 .pipe(
174 Effect.catchTag("AccessDeniedException", () => Effect.succeed(undefined)),
175 );
176 return Object.fromEntries(
177 (response?.Tags ?? []).flatMap((tag) =>
178 tag.Key !== undefined ? [[tag.Key, tag.Value ?? ""]] : [],
179 ),
180 );
181});
182
183const toAttrs = (
184 described: sagemaker.DescribeComputeQuotaResponse,
185): ComputeQuota["Attributes"] => {
186 const teamName = described.ComputeQuotaTarget?.TeamName ?? "";
187 return {
188 computeQuotaId: described.ComputeQuotaId,
189 computeQuotaArn: described.ComputeQuotaArn,
190 name: described.Name,
191 clusterArn: described.ClusterArn ?? "",
192 computeQuotaVersion: described.ComputeQuotaVersion,
193 teamName,
194 // Task-governance conventions: the quota materializes the team's
195 // namespace and Kueue LocalQueue under these derived names.
196 namespace: `hyperpod-ns-${teamName}`,
197 queueName: `hyperpod-ns-${teamName}-localqueue`,
198 };
199};
200
201/**
202 * The compute allocation is still transitioning toward the awaited state —
203 * retried by the bounded wait schedule.
204 */
205class ComputeQuotaNotReady extends Data.TaggedError("ComputeQuotaNotReady")<{
206 readonly quotaId: string;
207 readonly status: string | undefined;
208}> {}
209
210/**
211 * The compute allocation converged to a terminal failed status.
212 */
213export class ComputeQuotaFailed extends Data.TaggedError("ComputeQuotaFailed")<{
214 readonly quotaId: string;
215 readonly status: string | undefined;
216 readonly message: string | undefined;
217}> {}
218
219const FAILED_STATUSES: string[] = [
220 "CreateFailed",
221 "CreateRollbackFailed",
222 "UpdateFailed",
223 "UpdateRollbackFailed",
224 "DeleteFailed",
225 "DeleteRollbackFailed",
226];
227
228// Explicitly-typed retry wrapper — an inline `Effect.retry` in provider
229// lifecycle code leaks `Retry.Return`'s conditional type into declaration
230// emit and widens the provider layer to `unknown` for every consumer of
231// `AWS.providers()`.
232const retryWhileNotReady = <A, E extends { readonly _tag: string }, R>(
233 self: Effect.Effect<A, E, R>,
234): Effect.Effect<A, E, R> =>
235 Effect.retry(self, {
236 while: (e) => e._tag === "ComputeQuotaNotReady",
237 schedule: Schedule.max([Schedule.spaced("5 seconds"), Schedule.recurs(60)]),
238 });
239
240const waitForQuota = (quotaId: string, target: "Ready" | "Gone") =>
241 retryWhileNotReady(
242 Effect.gen(function* () {
243 const described = yield* describeQuotaOrUndefined(quotaId);
244 if (target === "Gone") {
245 if (described === undefined || described.Status === "Deleted") return;
246 if (FAILED_STATUSES.includes(described.Status)) {
247 return yield* Effect.fail(
248 new ComputeQuotaFailed({
249 quotaId,
250 status: described.Status,
251 message: described.FailureReason,
252 }),
253 );
254 }
255 return yield* Effect.fail(
256 new ComputeQuotaNotReady({ quotaId, status: described.Status }),
257 );
258 }
259 if (described?.Status === "Created" || described?.Status === "Updated") {
260 return;
261 }
262 if (
263 described !== undefined &&
264 FAILED_STATUSES.includes(described.Status)
265 ) {
266 return yield* Effect.fail(
267 new ComputeQuotaFailed({
268 quotaId,
269 status: described.Status,
270 message: described.FailureReason,
271 }),
272 );
273 }
274 return yield* Effect.fail(
275 new ComputeQuotaNotReady({ quotaId, status: described?.Status }),
276 );
277 }),
278 );
279
280export const ComputeQuotaProvider = () =>
281 Provider.effect(
282 ComputeQuota,
283 Effect.gen(function* () {
284 return {
285 stables: ["computeQuotaId", "computeQuotaArn", "name", "clusterArn"],
286 list: () =>
287 Effect.gen(function* () {
288 const summaries = yield* sagemaker.listComputeQuotas.pages({}).pipe(
289 EffectStream.runCollect,
290 Effect.map((chunk) =>
291 Array.from(chunk).flatMap(
292 (page) => page.ComputeQuotaSummaries ?? [],
293 ),
294 ),
295 );
296 return summaries.flatMap((s) => {
297 if (s.ComputeQuotaId === undefined || s.Status === "Deleted") {
298 return [];
299 }
300 const teamName = s.ComputeQuotaTarget?.TeamName ?? "";
301 return [
302 {
303 computeQuotaId: s.ComputeQuotaId,
304 computeQuotaArn: s.ComputeQuotaArn ?? "",
305 name: s.Name ?? "",
306 clusterArn: s.ClusterArn ?? "",
307 computeQuotaVersion: s.ComputeQuotaVersion ?? 1,
308 teamName,
309 namespace: `hyperpod-ns-${teamName}`,
310 queueName: `hyperpod-ns-${teamName}-localqueue`,
311 },
312 ];
313 });
314 }),
315 read: Effect.fn(function* ({ id, olds, output }) {
316 const quotaId =
317 output?.computeQuotaId ??
318 (yield* findQuotaByName(yield* createQuotaName(id, olds ?? {})))
319 ?.ComputeQuotaId;
320 if (quotaId === undefined) return undefined;
321 const described = yield* describeQuotaOrUndefined(quotaId);
322 if (
323 !described ||
324 described.Status === "Deleting" ||
325 described.Status === "Deleted"
326 ) {
327 return undefined;
328 }
329 const attrs = toAttrs(described);
330 const tags = yield* fetchQuotaTags(attrs.computeQuotaArn);
331 return (yield* hasAlchemyTags(id, tags as Tags))
332 ? attrs
333 : Unowned(attrs);
334 }),
335 diff: Effect.fn(function* ({ id, news, olds }) {
336 if (!isResolved(news)) return;
337 if (olds === undefined) return;
338 const oldName = yield* createQuotaName(id, olds);
339 const newName = yield* createQuotaName(id, news);
340 // The name and target cluster are fixed at creation.
341 if (oldName !== newName || olds.clusterArn !== news.clusterArn) {
342 return { action: "replace" } as const;
343 }
344 }),
345 reconcile: Effect.fn(function* ({ id, news, output, session }) {
346 if (!news) {
347 return yield* Effect.fail(
348 new Error("SageMaker ComputeQuota requires props"),
349 );
350 }
351 const name = yield* createQuotaName(id, news);
352 const internalTags = yield* createInternalTags(id);
353 const desiredTags = { ...internalTags, ...news.tags };
354
355 // Observe — by cached id first, then by name (lost state or race).
356 let quotaId = output?.computeQuotaId;
357 let described =
358 quotaId !== undefined
359 ? yield* describeQuotaOrUndefined(quotaId)
360 : undefined;
361 if (described === undefined) {
362 const found = yield* findQuotaByName(name);
363 described =
364 found?.ComputeQuotaId !== undefined
365 ? yield* describeQuotaOrUndefined(found.ComputeQuotaId)
366 : undefined;
367 }
368
369 // Ensure — create if missing; a Conflict means a concurrent
370 // create won the race, so re-observe by name.
371 if (described === undefined) {
372 const created = yield* sagemaker
373 .createComputeQuota({
374 Name: name,
375 ClusterArn: news.clusterArn,
376 ComputeQuotaConfig: news.computeQuotaConfig,
377 ComputeQuotaTarget: news.computeQuotaTarget,
378 ActivationState: news.activationState,
379 Description: news.description,
380 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
381 Key,
382 Value,
383 })),
384 })
385 .pipe(
386 Effect.catchTag("ConflictException", () =>
387 Effect.succeed(undefined),
388 ),
389 );
390 quotaId =
391 created?.ComputeQuotaId ??
392 (yield* findQuotaByName(name))?.ComputeQuotaId;
393 if (quotaId === undefined) {
394 return yield* Effect.fail(
395 new Error(`failed to create compute quota ${name}`),
396 );
397 }
398 yield* session.note(`Creating compute quota ${name}...`);
399 yield* waitForQuota(quotaId, "Ready");
400 } else {
401 quotaId = described.ComputeQuotaId;
402 // Wait out any in-flight transition before diffing.
403 yield* waitForQuota(quotaId, "Ready");
404 described = yield* describeQuotaOrUndefined(quotaId);
405 // Sync — diff observed allocation against desired.
406 if (
407 described !== undefined &&
408 (JSON.stringify(described.ComputeQuotaConfig) !==
409 JSON.stringify(news.computeQuotaConfig) ||
410 JSON.stringify(described.ComputeQuotaTarget) !==
411 JSON.stringify(news.computeQuotaTarget) ||
412 (news.activationState !== undefined &&
413 described.ActivationState !== news.activationState) ||
414 (described.Description ?? undefined) !==
415 (news.description ?? undefined))
416 ) {
417 yield* sagemaker.updateComputeQuota({
418 ComputeQuotaId: quotaId,
419 TargetVersion: described.ComputeQuotaVersion,
420 ComputeQuotaConfig: news.computeQuotaConfig,
421 ComputeQuotaTarget: news.computeQuotaTarget,
422 ActivationState: news.activationState,
423 Description: news.description,
424 });
425 yield* session.note(`Updating compute quota ${name}...`);
426 yield* waitForQuota(quotaId, "Ready");
427 }
428 }
429
430 described = yield* describeQuotaOrUndefined(quotaId);
431 if (described === undefined) {
432 return yield* Effect.fail(
433 new Error(`failed to read reconciled compute quota ${name}`),
434 );
435 }
436 const attrs = toAttrs(described);
437
438 // Sync tags — diff against OBSERVED cloud tags.
439 const currentTags = yield* fetchQuotaTags(attrs.computeQuotaArn);
440 const { removed, upsert } = diffTags(currentTags, desiredTags);
441 if (removed.length > 0) {
442 yield* sagemaker.deleteTags({
443 ResourceArn: attrs.computeQuotaArn,
444 TagKeys: removed,
445 });
446 }
447 if (upsert.length > 0) {
448 yield* sagemaker.addTags({
449 ResourceArn: attrs.computeQuotaArn,
450 Tags: upsert.map(({ Key, Value }) => ({ Key, Value })),
451 });
452 }
453
454 yield* session.note(attrs.computeQuotaArn);
455 return attrs;
456 }),
457 delete: Effect.fn(function* ({ output }) {
458 yield* sagemaker
459 .deleteComputeQuota({ ComputeQuotaId: output.computeQuotaId })
460 .pipe(Effect.catchTag("ResourceNotFound", () => Effect.void));
461 yield* waitForQuota(output.computeQuotaId, "Gone");
462 }),
463 };
464 }),
465 );