1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import { createHash } from "node:crypto";
3import * as Data from "effect/Data";
4import * as Effect from "effect/Effect";
5import * as Option from "effect/Option";
6import * as Stream from "effect/Stream";
7import { Unowned } from "../../AdoptPolicy.ts";
8import { isResolved } from "../../Diff.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { createInternalLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import { listLocations } from "./names.ts";
15import { lastSegment, locationOf, locationParent } from "./ownership.ts";
16
17export type StudyMetricSpec = {
18 /** Metric id. Must be unique among metrics. */
19 metricId: string;
20 /** Optimization goal. */
21 goal: "MAXIMIZE" | "MINIMIZE" | (string & {});
22};
23
24export type StudyParameterSpec = {
25 /** Parameter id. Must be unique among parameters. */
26 parameterId: string;
27 /** Inclusive range for a `DOUBLE` parameter. */
28 doubleValueSpec?: {
29 minValue?: number;
30 maxValue?: number;
31 defaultValue?: number;
32 };
33 /** Inclusive range for an `INTEGER` parameter. */
34 integerValueSpec?: {
35 minValue?: string;
36 maxValue?: string;
37 defaultValue?: string;
38 };
39 /** Discrete feasible values. */
40 discreteValueSpec?: { values?: number[]; defaultValue?: number };
41 /** Categorical feasible values. */
42 categoricalValueSpec?: { values?: string[]; defaultValue?: string };
43 /** Scaling. Leave unset for categorical parameters. */
44 scaleType?:
45 | "UNIT_LINEAR_SCALE"
46 | "UNIT_LOG_SCALE"
47 | "UNIT_REVERSE_LOG_SCALE"
48 | (string & {});
49};
50
51export type StudySpec = {
52 /** Metrics to optimize. */
53 metrics: StudyMetricSpec[];
54 /** Parameters to tune. */
55 parameters: StudyParameterSpec[];
56 /** Search algorithm. */
57 algorithm?:
58 | "ALGORITHM_UNSPECIFIED"
59 | "GRID_SEARCH"
60 | "RANDOM_SEARCH"
61 | (string & {});
62 /** Which measurement to use when a trial reports several. */
63 measurementSelectionType?:
64 | "LAST_MEASUREMENT"
65 | "BEST_MEASUREMENT"
66 | (string & {});
67};
68
69export type StudyProps = {
70 /**
71 * Region (`us-central1`, …). Immutable — changing it replaces the Study.
72 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
73 */
74 location?: string;
75 /**
76 * User-facing display name: letters, digits, and underscores, starting
77 * with a letter. Vertex AI Studies have no labels field, so Alchemy
78 * appends an `_alchemy_{hash}` ownership suffix (stripped from
79 * attributes). Display name is the lookup key.
80 * @default "study"
81 */
82 displayName?: string;
83 /**
84 * Study configuration. Immutable — changing it replaces the Study.
85 */
86 studySpec: StudySpec;
87};
88
89export type Study = Resource<
90 "GCP.AIPlatform.Study",
91 StudyProps,
92 {
93 /** Full resource name `projects/{project}/locations/{location}/studies/{study}`. */
94 name: string;
95 /** Study id (last path segment). */
96 studyId: string;
97 /** Region id. */
98 location: string;
99 /** Project id. */
100 project: string;
101 /** User display name with the Alchemy ownership prefix stripped. */
102 displayName: string | undefined;
103 /** Study configuration. */
104 studySpec: aiplatform.GoogleCloudAiplatformV1StudySpec | undefined;
105 /** Server-reported state. */
106 state: string | undefined;
107 /** Why the study is inactive, if set. */
108 inactiveReason: string | undefined;
109 /** RFC3339 creation timestamp. */
110 createTime: string | undefined;
111 },
112 never,
113 Providers
114>;
115
116/**
117 * A Vertex AI Vizier Study for hyperparameter search.
118 *
119 * Studies have no labels field and no update RPC — Alchemy stamps an
120 * ownership hash onto the display name. Location and study spec are immutable.
121 *
122 * ### Creating a Study
123 * **Example:** Maximize accuracy over a double parameter
124 * ```typescript
125 * const study = yield* GCP.AIPlatform.Study("Tune", {
126 * studySpec: {
127 * metrics: [{ metricId: "accuracy", goal: "MAXIMIZE" }],
128 * parameters: [
129 * {
130 * parameterId: "learning_rate",
131 * doubleValueSpec: { minValue: 0.001, maxValue: 0.1 },
132 * },
133 * ],
134 * algorithm: "RANDOM_SEARCH",
135 * },
136 * });
137 * ```
138 *
139 * @resource
140 * @category AIPlatform
141 */
142export const Study = Resource<Study>("GCP.AIPlatform.Study");
143
144export class StudyNotResolved extends Data.TaggedError(
145 "GCP.AIPlatform.StudyNotResolved",
146)<{
147 name: string;
148}> {}
149
150const toStudySpec = (
151 spec: StudySpec,
152): aiplatform.GoogleCloudAiplatformV1StudySpec => ({
153 metrics: spec.metrics.map((metric) => ({
154 metricId: metric.metricId,
155 goal: metric.goal,
156 })),
157 parameters: spec.parameters.map((parameter) => ({
158 parameterId: parameter.parameterId,
159 doubleValueSpec: parameter.doubleValueSpec,
160 integerValueSpec: parameter.integerValueSpec,
161 discreteValueSpec: parameter.discreteValueSpec,
162 categoricalValueSpec: parameter.categoricalValueSpec,
163 scaleType: parameter.scaleType,
164 })),
165 algorithm: spec.algorithm,
166 measurementSelectionType: spec.measurementSelectionType,
167});
168
169const specKey = (spec: StudySpec) =>
170 JSON.stringify({
171 metrics: spec.metrics,
172 parameters: spec.parameters,
173 algorithm: spec.algorithm ?? "",
174 measurementSelectionType: spec.measurementSelectionType ?? "",
175 });
176
177/**
178 * Study display names allow only `[A-Za-z0-9_]`, so ownership is a
179 * `_alchemy_{hash}` suffix over the stack, stage, and logical id rather
180 * than the bracketed marker other Vertex AI resources use.
181 */
182const OWNER_SEPARATOR = "_alchemy_";
183const DEFAULT_DISPLAY_NAME = "study";
184
185const ownerSuffix = (id: string) =>
186 createInternalLabels(id).pipe(
187 Effect.flatMap((labels) =>
188 Effect.sync(
189 () =>
190 OWNER_SEPARATOR +
191 createHash("sha256")
192 .update(JSON.stringify(labels))
193 .digest("hex")
194 .slice(0, 20),
195 ),
196 ),
197 );
198
199const encodeDisplayName = (suffix: string, displayName: string | undefined) =>
200 `${displayName ?? DEFAULT_DISPLAY_NAME}${suffix}`;
201
202const parseDisplayName = (displayName: string | undefined) => {
203 const at = displayName?.lastIndexOf(OWNER_SEPARATOR) ?? -1;
204 return displayName !== undefined && at > 0
205 ? displayName.slice(0, at)
206 : displayName;
207};
208
209/** Whether a study's display name carries an Alchemy ownership suffix. */
210export const hasStudyOwnership = (displayName: string | undefined) =>
211 (displayName?.lastIndexOf(OWNER_SEPARATOR) ?? -1) > 0;
212
213const ownedByAlchemy = (id: string, displayName: string | undefined) =>
214 ownerSuffix(id).pipe(
215 Effect.map((suffix) => displayName?.endsWith(suffix) === true),
216 );
217
218const toAttrs = (
219 study: aiplatform.GoogleCloudAiplatformV1Study,
220 project: string,
221) => {
222 const name = study.name ?? "";
223 const parsed = { text: parseDisplayName(study.displayName) };
224 return {
225 name,
226 studyId: lastSegment(name),
227 location: locationOf(name),
228 project,
229 displayName: parsed.text,
230 studySpec: study.studySpec,
231 state: study.state,
232 inactiveReason: study.inactiveReason,
233 createTime: study.createTime,
234 };
235};
236
237const getByName = (name: string) =>
238 name.length === 0
239 ? Effect.succeed(undefined)
240 : aiplatform
241 .getProjectsLocationsStudies({ name })
242 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
243
244// `lookupProjectsLocationsStudies` answers a missing display name with a
245// 500 ("There was an issue when sending 'LookupStudyRequest'"), so scan
246// the listing instead.
247const lookupByDisplayName = (parent: string, displayName: string) =>
248 aiplatform.listProjectsLocationsStudies.pages({ parent, pageSize: 100 }).pipe(
249 Stream.flatMap((page) => Stream.fromIterable(page.studies ?? [])),
250 Stream.filter((study) => study.displayName === displayName),
251 Stream.runHead,
252 Effect.map(Option.getOrUndefined),
253 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
254 );
255
256const listAt = (parent: string, project: string) =>
257 aiplatform.listProjectsLocationsStudies.pages({ parent, pageSize: 100 }).pipe(
258 Stream.flatMap((page) => Stream.fromIterable(page.studies ?? [])),
259 Stream.filter((study) => hasStudyOwnership(study.displayName)),
260 Stream.map((study) => toAttrs(study, project)),
261 Stream.runCollect,
262 Effect.map((chunk) => Array.from(chunk)),
263 Effect.catchTag("NotFound", () => Effect.succeed([])),
264 );
265
266export const StudyProvider = () =>
267 Provider.succeed(Study, {
268 stables: ["name", "studyId", "location", "project", "createTime"],
269
270 diff: Effect.fn(function* ({ news, olds, output }) {
271 if (!isResolved(news)) return undefined;
272 const previousLocation = olds?.location ?? output?.location;
273 const nextLocation = news.location ?? previousLocation;
274 if (previousLocation !== undefined && previousLocation !== nextLocation) {
275 return { action: "replace" as const, deleteFirst: false };
276 }
277 if (olds?.studySpec !== undefined) {
278 if (specKey(olds.studySpec) !== specKey(news.studySpec)) {
279 return { action: "replace" as const, deleteFirst: false };
280 }
281 }
282 return undefined;
283 }),
284
285 read: Effect.fn(function* ({ id, olds, output }) {
286 const env = yield* GcpEnvironment.current;
287 const existing = yield* getByName(output?.name ?? "");
288 if (existing !== undefined) {
289 const attrs = toAttrs(existing, env.project);
290 return (yield* ownedByAlchemy(id, existing.displayName))
291 ? attrs
292 : Unowned(attrs);
293 }
294 const location = olds?.location ?? output?.location ?? env.region;
295 const parent = locationParent(env.project, location);
296 const displayName = encodeDisplayName(
297 yield* ownerSuffix(id),
298 olds?.displayName,
299 );
300 const lookedUp = yield* lookupByDisplayName(parent, displayName);
301 if (lookedUp === undefined) return undefined;
302 const attrs = toAttrs(lookedUp, env.project);
303 return (yield* ownedByAlchemy(id, lookedUp.displayName))
304 ? attrs
305 : Unowned(attrs);
306 }),
307
308 list: () =>
309 Effect.gen(function* () {
310 const env = yield* GcpEnvironment.current;
311 return (yield* Effect.forEach(listLocations(env.region), (location) =>
312 listAt(locationParent(env.project, location), env.project),
313 )).flat();
314 }),
315
316 reconcile: Effect.fn(function* ({ id, news, output }) {
317 const env = yield* GcpEnvironment.current;
318 const location = news.location ?? output?.location ?? env.region;
319 const parent = locationParent(env.project, location);
320 const displayName = encodeDisplayName(
321 yield* ownerSuffix(id),
322 news.displayName,
323 );
324 const studySpec = toStudySpec(news.studySpec);
325
326 let current = yield* getByName(output?.name ?? "");
327 if (current === undefined) {
328 current = yield* lookupByDisplayName(parent, displayName);
329 }
330
331 if (current === undefined) {
332 const created = yield* aiplatform
333 .createProjectsLocationsStudies({
334 parent,
335 body: {
336 displayName,
337 studySpec,
338 },
339 })
340 .pipe(
341 Effect.catchTag("Conflict", () =>
342 lookupByDisplayName(parent, displayName),
343 ),
344 );
345 current = created ?? undefined;
346 }
347
348 if (current === undefined) {
349 return yield* new StudyNotResolved({
350 name: output?.name ?? `${parent}/studies/-`,
351 });
352 }
353
354 return toAttrs(current, env.project);
355 }),
356
357 delete: Effect.fn(function* ({ output }) {
358 yield* aiplatform
359 .deleteProjectsLocationsStudies({ name: output.name })
360 .pipe(Effect.catchTag("NotFound", () => Effect.void));
361 }),
362 });