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 HyperparameterTuningJobProps = {
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 * Desired total number of Trials.
48 */
49 maxTrialCount: number;
50 /**
51 * Desired number of Trials to run in parallel.
52 */
53 parallelTrialCount: number;
54 /**
55 * Number of failed Trials before the job fails. `0` lets Vertex decide.
56 */
57 maxFailedTrialCount?: number;
58 /**
59 * Study configuration (metrics and parameters).
60 */
61 studySpec: aiplatform.GoogleCloudAiplatformV1StudySpec;
62 /**
63 * Spec of the CustomJob backing each Trial.
64 */
65 trialJobSpec: aiplatform.GoogleCloudAiplatformV1CustomJobSpec;
66 /**
67 * Customer-managed encryption key. Immutable.
68 */
69 encryptionSpec?: aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
70};
71
72export type HyperparameterTuningJob = Resource<
73 "GCP.AIPlatform.HyperparameterTuningJob",
74 HyperparameterTuningJobProps,
75 {
76 /** Full resource name. */
77 name: string;
78 /** Job id (last path segment). */
79 hyperparameterTuningJobId: string;
80 /** Project id. */
81 project: string;
82 /** Location id. */
83 location: string;
84 /** Display name. */
85 displayName: string | undefined;
86 /** User labels (Alchemy ownership labels stripped). */
87 labels: Record<string, string>;
88 /** Job state. */
89 state: string | undefined;
90 /** Desired total trial count. */
91 maxTrialCount: number | undefined;
92 /** Parallel trial count. */
93 parallelTrialCount: number | undefined;
94 /** RFC3339 creation timestamp. */
95 createTime: string | undefined;
96 /** RFC3339 start timestamp. */
97 startTime: string | undefined;
98 /** RFC3339 end timestamp. */
99 endTime: string | undefined;
100 },
101 never,
102 Providers
103>;
104
105/**
106 * A Vertex AI HyperparameterTuningJob that runs CustomJobs for each trial.
107 *
108 * The API assigns the job id. There is no update method — changing display
109 * name, labels, trial counts, or specs replaces the job. Delete cancels a
110 * running job first.
111 *
112 * ### Creating a HyperparameterTuningJob
113 * **Example:** Maximize accuracy
114 * ```typescript
115 * const job = yield* GCP.AIPlatform.HyperparameterTuningJob("Tune", {
116 * maxTrialCount: 2,
117 * parallelTrialCount: 1,
118 * studySpec: {
119 * metrics: [{ metricId: "accuracy", goal: "MAXIMIZE" }],
120 * parameters: [{
121 * parameterId: "lr",
122 * doubleValueSpec: { minValue: 0.001, maxValue: 0.1 },
123 * }],
124 * },
125 * trialJobSpec: {
126 * workerPoolSpecs: [{
127 * machineSpec: { machineType: "n1-standard-4" },
128 * replicaCount: "1",
129 * containerSpec: { imageUri: "gcr.io/cloud-aiplatform/training/tf-cpu.2-8:latest" },
130 * }],
131 * },
132 * });
133 * ```
134 *
135 * @resource
136 * @category AIPlatform
137 */
138export const HyperparameterTuningJob = Resource<HyperparameterTuningJob>(
139 "GCP.AIPlatform.HyperparameterTuningJob",
140);
141
142export class HyperparameterTuningJobNotResolved extends Data.TaggedError(
143 "GCP.AIPlatform.HyperparameterTuningJobNotResolved",
144)<{
145 name: string;
146}> {}
147
148export class HyperparameterTuningJobStillExists extends Data.TaggedError(
149 "GCP.AIPlatform.HyperparameterTuningJobStillExists",
150)<{
151 name: string;
152}> {}
153
154const toAttrs = (
155 job: aiplatform.GoogleCloudAiplatformV1HyperparameterTuningJob,
156 project: string,
157) => {
158 const name = job.name ?? "";
159 return {
160 name,
161 hyperparameterTuningJobId: lastSegment(name),
162 project: projectOf(name, project),
163 location: locationOf(name),
164 displayName: job.displayName,
165 labels: userLabels(job.labels),
166 state: job.state,
167 maxTrialCount: job.maxTrialCount,
168 parallelTrialCount: job.parallelTrialCount,
169 createTime: job.createTime,
170 startTime: job.startTime,
171 endTime: job.endTime,
172 };
173};
174
175const getByName = (name: string) =>
176 aiplatform
177 .getProjectsLocationsHyperparameterTuningJobs({ name })
178 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
179
180const listPage = (parent: string, filter?: string) =>
181 aiplatform.listProjectsLocationsHyperparameterTuningJobs
182 .pages({ parent, pageSize: 100, filter })
183 .pipe(
184 Stream.runCollect,
185 Effect.map((pages) =>
186 Array.from(pages).flatMap(
187 (page) => page.hyperparameterTuningJobs ?? [],
188 ),
189 ),
190 Effect.catchTag("NotFound", () =>
191 Effect.succeed(
192 [] as aiplatform.GoogleCloudAiplatformV1HyperparameterTuningJob[],
193 ),
194 ),
195 );
196
197const findOwned = (
198 project: string,
199 location: string,
200 labels: Record<string, string>,
201) =>
202 listPage(
203 `projects/${project}/locations/${location}`,
204 alchemyIdFilter(labels),
205 ).pipe(
206 Effect.map(
207 (items) =>
208 items.find((item) => hasAlchemyPrefix(item.labels)) ?? undefined,
209 ),
210 );
211
212const waitUntilGone = (name: string) =>
213 getByName(name).pipe(
214 Effect.filterOrFail(
215 (job) => job === undefined,
216 () => new HyperparameterTuningJobStillExists({ name }),
217 ),
218 Effect.asVoid,
219 Effect.retry({
220 while: (error) =>
221 error._tag === "GCP.AIPlatform.HyperparameterTuningJobStillExists",
222 times: 10,
223 schedule: Schedule.spaced("2 seconds"),
224 }),
225 );
226
227export const HyperparameterTuningJobProvider = () =>
228 Provider.succeed(HyperparameterTuningJob, {
229 stables: [
230 "name",
231 "hyperparameterTuningJobId",
232 "project",
233 "location",
234 "createTime",
235 ],
236
237 diff: Effect.fn(function* ({ news, olds, output }) {
238 const env = yield* GcpEnvironment.current;
239 if (!isResolved(news)) return undefined;
240 const previousLocation = normalizeLocation(
241 olds?.location ?? output?.location,
242 env.region,
243 );
244 const nextLocation = normalizeLocation(
245 news.location ?? olds?.location ?? output?.location,
246 env.region,
247 );
248 const previousDisplay = olds?.displayName ?? output?.displayName;
249 const displayChanged =
250 news.displayName !== undefined &&
251 previousDisplay !== undefined &&
252 news.displayName !== previousDisplay;
253 const previousMax = olds?.maxTrialCount ?? output?.maxTrialCount;
254 const previousParallel =
255 olds?.parallelTrialCount ?? output?.parallelTrialCount;
256 const specsChanged =
257 (olds !== undefined &&
258 stableJson(olds.studySpec) !== stableJson(news.studySpec)) ||
259 (olds !== undefined &&
260 stableJson(olds.trialJobSpec) !== stableJson(news.trialJobSpec));
261 if (
262 previousLocation !== nextLocation ||
263 displayChanged ||
264 (previousMax !== undefined && news.maxTrialCount !== previousMax) ||
265 (previousParallel !== undefined &&
266 news.parallelTrialCount !== previousParallel) ||
267 specsChanged
268 ) {
269 return { action: "replace" as const, deleteFirst: false };
270 }
271 return undefined;
272 }),
273
274 read: Effect.fn(function* ({ id, olds, output }) {
275 const env = yield* GcpEnvironment.current;
276 const location = normalizeLocation(
277 olds?.location ?? output?.location,
278 env.region,
279 );
280 const ownership = yield* createInternalLabels(id);
281 const existing =
282 (output?.name !== undefined
283 ? yield* getByName(output.name)
284 : undefined) ?? (yield* findOwned(env.project, location, ownership));
285 if (existing === undefined) return undefined;
286 const attrs = toAttrs(existing, env.project);
287 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
288 ? attrs
289 : Unowned(attrs);
290 }),
291
292 list: () =>
293 Effect.gen(function* () {
294 const env = yield* GcpEnvironment.current;
295 const pages = yield* Effect.forEach(
296 listLocations(env.region),
297 (location) =>
298 listPage(`projects/${env.project}/locations/${location}`),
299 { concurrency: 4 },
300 );
301 return pages
302 .flat()
303 .filter((job) => hasAlchemyPrefix(job.labels))
304 .map((job) => toAttrs(job, env.project));
305 }),
306
307 reconcile: Effect.fn(function* ({ id, news, output }) {
308 const env = yield* GcpEnvironment.current;
309 const location = normalizeLocation(
310 news.location ?? output?.location,
311 env.region,
312 );
313 const displayName = yield* toDisplayName(
314 id,
315 news.displayName,
316 output?.displayName,
317 );
318 const desiredLabels = {
319 ...toLabels(news.labels),
320 ...(yield* createInternalLabels(id)),
321 };
322
323 let current =
324 (output?.name !== undefined
325 ? yield* getByName(output.name)
326 : undefined) ??
327 (yield* findOwned(env.project, location, desiredLabels));
328
329 if (current === undefined) {
330 const created = yield* aiplatform
331 .createProjectsLocationsHyperparameterTuningJobs({
332 parent: `projects/${env.project}/locations/${location}`,
333 body: {
334 displayName,
335 labels: desiredLabels,
336 maxTrialCount: news.maxTrialCount,
337 parallelTrialCount: news.parallelTrialCount,
338 maxFailedTrialCount: news.maxFailedTrialCount,
339 studySpec: news.studySpec,
340 trialJobSpec: news.trialJobSpec,
341 encryptionSpec: news.encryptionSpec,
342 },
343 })
344 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
345 current =
346 created ?? (yield* findOwned(env.project, location, desiredLabels));
347 }
348
349 if (current === undefined || current.name === undefined) {
350 return yield* new HyperparameterTuningJobNotResolved({
351 name: output?.name ?? `${location}/hyperparameterTuningJobs`,
352 });
353 }
354
355 return toAttrs(current, env.project);
356 }),
357
358 delete: Effect.fn(function* ({ output }) {
359 yield* aiplatform
360 .cancelProjectsLocationsHyperparameterTuningJobs({
361 name: output.name,
362 body: {},
363 })
364 .pipe(
365 Effect.catchTag(
366 ["NotFound", "BadRequest", "Conflict"],
367 () => Effect.void,
368 ),
369 );
370 // Running jobs reject deletes; wait for the cancel to land.
371 yield* getByName(output.name).pipe(
372 Effect.repeat({
373 schedule: Schedule.spaced("4 seconds"),
374 until: (job) => job === undefined || isJobTerminal(job.state),
375 times: 45,
376 }),
377 );
378 const operation = yield* aiplatform
379 .deleteProjectsLocationsHyperparameterTuningJobs({
380 name: output.name,
381 })
382 .pipe(
383 Effect.retry({
384 while: (error) => error._tag === "Conflict",
385 times: 8,
386 schedule: Schedule.spaced("2 seconds"),
387 }),
388 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
389 );
390 if (operation !== undefined) {
391 yield* waitForOperation(operation, { notFoundOk: true });
392 }
393 yield* waitUntilGone(output.name);
394 }),
395 });