1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { tagRecord } from "../../Tags.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import {
15 alchemyIdFilter,
16 createInternalLabels,
17 hasAlchemyPrefix,
18 listLocations,
19 locationOf,
20 lastSegment,
21 normalizeLocation,
22 projectOf,
23 stableJson,
24 toDisplayName,
25 toLabels,
26 userLabels,
27} from "./names.ts";
28import { isJobTerminal } from "./internal.ts";
29import { waitForOperation } from "./operations.ts";
30
31export type NasJobProps = {
32 /**
33 * Vertex AI location. Immutable — changing it replaces the job.
34 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
35 */
36 location?: string;
37 /**
38 * Display name (max 128 Unicode characters). Generated from the stack,
39 * stage, and logical id when omitted.
40 */
41 displayName?: string;
42 /**
43 * User labels. Alchemy ownership labels are merged in automatically.
44 */
45 labels?: Record<string, string>;
46 /**
47 * Neural Architecture Search specification.
48 */
49 nasJobSpec: aiplatform.GoogleCloudAiplatformV1NasJobSpec;
50 /**
51 * Enable restricted-image training for the tenant project.
52 * @default false
53 */
54 enableRestrictedImageTraining?: boolean;
55 /**
56 * Customer-managed encryption key. Immutable.
57 */
58 encryptionSpec?: aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
59};
60
61export type NasJob = Resource<
62 "GCP.AIPlatform.NasJob",
63 NasJobProps,
64 {
65 /** Full resource name. */
66 name: string;
67 /** Job id (last path segment). */
68 nasJobId: string;
69 /** Project id. */
70 project: string;
71 /** Location id. */
72 location: string;
73 /** Display name. */
74 displayName: string | undefined;
75 /** User labels (Alchemy ownership labels stripped). */
76 labels: Record<string, string>;
77 /** Job state. */
78 state: string | undefined;
79 /** Whether restricted-image training is enabled. */
80 enableRestrictedImageTraining: boolean;
81 /** RFC3339 creation timestamp. */
82 createTime: string | undefined;
83 /** RFC3339 start timestamp. */
84 startTime: string | undefined;
85 /** RFC3339 end timestamp. */
86 endTime: string | undefined;
87 },
88 never,
89 Providers
90>;
91
92/**
93 * A Vertex AI Neural Architecture Search job.
94 *
95 * The API assigns the job id. There is no update method — changing display
96 * name, labels, or `nasJobSpec` replaces the job. Delete cancels a running
97 * job first.
98 *
99 * ### Creating a NasJob
100 * **Example:** Multi-trial search
101 * ```typescript
102 * const job = yield* GCP.AIPlatform.NasJob("Search", {
103 * nasJobSpec: {
104 * searchSpaceSpec: "{}",
105 * multiTrialAlgorithmSpec: {
106 * metric: { metricId: "accuracy", goal: "MAXIMIZE" },
107 * searchTrialSpec: {
108 * maxTrialCount: 1,
109 * maxParallelTrialCount: 1,
110 * searchTrialJobSpec: {
111 * workerPoolSpecs: [{
112 * machineSpec: { machineType: "n1-standard-4" },
113 * replicaCount: "1",
114 * containerSpec: { imageUri: "gcr.io/cloud-aiplatform/training/tf-cpu.2-8:latest" },
115 * }],
116 * },
117 * },
118 * },
119 * },
120 * });
121 * ```
122 *
123 * @resource
124 * @category AIPlatform
125 */
126export const NasJob = Resource<NasJob>("GCP.AIPlatform.NasJob");
127
128export class NasJobNotResolved extends Data.TaggedError(
129 "GCP.AIPlatform.NasJobNotResolved",
130)<{
131 name: string;
132}> {}
133
134export class NasJobStillExists extends Data.TaggedError(
135 "GCP.AIPlatform.NasJobStillExists",
136)<{
137 name: string;
138}> {}
139
140const toAttrs = (
141 job: aiplatform.GoogleCloudAiplatformV1NasJob,
142 project: string,
143) => {
144 const name = job.name ?? "";
145 return {
146 name,
147 nasJobId: lastSegment(name),
148 project: projectOf(name, project),
149 location: locationOf(name),
150 displayName: job.displayName,
151 labels: userLabels(job.labels),
152 state: job.state,
153 enableRestrictedImageTraining: job.enableRestrictedImageTraining === true,
154 createTime: job.createTime,
155 startTime: job.startTime,
156 endTime: job.endTime,
157 };
158};
159
160const getByName = (name: string) =>
161 aiplatform
162 .getProjectsLocationsNasJobs({ name })
163 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
164
165const listPage = (parent: string, filter?: string) =>
166 aiplatform.listProjectsLocationsNasJobs
167 .pages({ parent, pageSize: 100, filter })
168 .pipe(
169 Stream.runCollect,
170 Effect.map((pages) =>
171 Array.from(pages).flatMap((page) => page.nasJobs ?? []),
172 ),
173 Effect.catchTag("NotFound", () =>
174 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1NasJob[]),
175 ),
176 );
177
178const findOwned = (
179 project: string,
180 location: string,
181 labels: Record<string, string>,
182) =>
183 listPage(
184 `projects/${project}/locations/${location}`,
185 alchemyIdFilter(labels),
186 ).pipe(
187 Effect.map(
188 (items) =>
189 items.find((item) => hasAlchemyPrefix(item.labels)) ?? undefined,
190 ),
191 );
192
193const waitUntilGone = (name: string) =>
194 getByName(name).pipe(
195 Effect.filterOrFail(
196 (job) => job === undefined,
197 () => new NasJobStillExists({ name }),
198 ),
199 Effect.asVoid,
200 Effect.retry({
201 while: (error) => error._tag === "GCP.AIPlatform.NasJobStillExists",
202 times: 10,
203 schedule: Schedule.spaced("2 seconds"),
204 }),
205 );
206
207export const NasJobProvider = () =>
208 Provider.succeed(NasJob, {
209 stables: ["name", "nasJobId", "project", "location", "createTime"],
210
211 diff: Effect.fn(function* ({ news, olds, output }) {
212 const env = yield* GcpEnvironment.current;
213 if (!isResolved(news)) return undefined;
214 const previousLocation = normalizeLocation(
215 olds?.location ?? output?.location,
216 env.region,
217 );
218 const nextLocation = normalizeLocation(
219 news.location ?? olds?.location ?? output?.location,
220 env.region,
221 );
222 const previousDisplay = olds?.displayName ?? output?.displayName;
223 const displayChanged =
224 news.displayName !== undefined &&
225 previousDisplay !== undefined &&
226 news.displayName !== previousDisplay;
227 const specsChanged =
228 olds !== undefined &&
229 stableJson(olds.nasJobSpec) !== stableJson(news.nasJobSpec);
230 if (previousLocation !== nextLocation || displayChanged || specsChanged) {
231 return { action: "replace" as const, deleteFirst: false };
232 }
233 return undefined;
234 }),
235
236 read: Effect.fn(function* ({ id, olds, output }) {
237 const env = yield* GcpEnvironment.current;
238 const location = normalizeLocation(
239 olds?.location ?? output?.location,
240 env.region,
241 );
242 const ownership = yield* createInternalLabels(id);
243 const existing =
244 (output?.name !== undefined
245 ? yield* getByName(output.name)
246 : undefined) ?? (yield* findOwned(env.project, location, ownership));
247 if (existing === undefined) return undefined;
248 const attrs = toAttrs(existing, env.project);
249 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
250 ? attrs
251 : Unowned(attrs);
252 }),
253
254 list: () =>
255 Effect.gen(function* () {
256 const env = yield* GcpEnvironment.current;
257 const pages = yield* Effect.forEach(
258 listLocations(env.region),
259 (location) =>
260 listPage(`projects/${env.project}/locations/${location}`),
261 { concurrency: 4 },
262 );
263 return pages
264 .flat()
265 .filter((job) => hasAlchemyPrefix(job.labels))
266 .map((job) => toAttrs(job, env.project));
267 }),
268
269 reconcile: Effect.fn(function* ({ id, news, output }) {
270 const env = yield* GcpEnvironment.current;
271 const location = normalizeLocation(
272 news.location ?? output?.location,
273 env.region,
274 );
275 const displayName = yield* toDisplayName(
276 id,
277 news.displayName,
278 output?.displayName,
279 );
280 const desiredLabels = {
281 ...toLabels(news.labels),
282 ...(yield* createInternalLabels(id)),
283 };
284
285 let current =
286 (output?.name !== undefined
287 ? yield* getByName(output.name)
288 : undefined) ??
289 (yield* findOwned(env.project, location, desiredLabels));
290
291 if (current === undefined) {
292 const created = yield* aiplatform
293 .createProjectsLocationsNasJobs({
294 parent: `projects/${env.project}/locations/${location}`,
295 body: {
296 displayName,
297 labels: desiredLabels,
298 nasJobSpec: news.nasJobSpec,
299 enableRestrictedImageTraining:
300 news.enableRestrictedImageTraining === true,
301 encryptionSpec: news.encryptionSpec,
302 },
303 })
304 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
305 current =
306 created ?? (yield* findOwned(env.project, location, desiredLabels));
307 }
308
309 if (current === undefined || current.name === undefined) {
310 return yield* new NasJobNotResolved({
311 name: output?.name ?? `${location}/nasJobs`,
312 });
313 }
314
315 return toAttrs(current, env.project);
316 }),
317
318 delete: Effect.fn(function* ({ output }) {
319 yield* aiplatform
320 .cancelProjectsLocationsNasJobs({
321 name: output.name,
322 body: {},
323 })
324 .pipe(
325 Effect.catchTag(
326 ["NotFound", "BadRequest", "Conflict"],
327 () => Effect.void,
328 ),
329 );
330 // Running jobs reject deletes; wait for the cancel to land.
331 yield* getByName(output.name).pipe(
332 Effect.repeat({
333 schedule: Schedule.spaced("4 seconds"),
334 until: (job) => job === undefined || isJobTerminal(job.state),
335 times: 45,
336 }),
337 );
338 const operation = yield* aiplatform
339 .deleteProjectsLocationsNasJobs({ name: output.name })
340 .pipe(
341 Effect.retry({
342 while: (error) => error._tag === "Conflict",
343 times: 8,
344 schedule: Schedule.spaced("2 seconds"),
345 }),
346 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
347 );
348 if (operation !== undefined) {
349 yield* waitForOperation(operation, { notFoundOk: true });
350 }
351 yield* waitUntilGone(output.name);
352 }),
353 });