1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Stream from "effect/Stream";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { GcpEnvironment } from "../Environment.ts";
10import { createInternalLabels } from "../Labels.ts";
11import type { Providers } from "../Providers.ts";
12import { listLocations } from "./names.ts";
13import { hasStudyOwnership } from "./Study.ts";
14import {
15 lastSegment,
16 locationOf,
17 locationParent,
18 parentOf,
19} from "./ownership.ts";
20
21export type TrialParameter = {
22 /** Parameter id defined on the parent Study. */
23 parameterId: string;
24 /** Assigned value (`number` for numeric parameters, `string` for categorical). */
25 value?: unknown;
26};
27
28export type StudiesTrialProps = {
29 /**
30 * Parent Study resource name
31 * (`projects/{project}/locations/{location}/studies/{study}`).
32 * Immutable — changing it replaces the Trial.
33 */
34 parent: string;
35 /**
36 * User-provided parameter assignments. Immutable after create.
37 */
38 parameters?: TrialParameter[];
39 /**
40 * Client identifier recorded on the Trial. Alchemy stores the logical
41 * id here when omitted so `list` / nuke can find user-provided trials.
42 */
43 clientId?: string;
44};
45
46export type StudiesTrial = Resource<
47 "GCP.AIPlatform.StudiesTrial",
48 StudiesTrialProps,
49 {
50 /** Full resource name. */
51 name: string;
52 /** Trial id assigned by Vizier. */
53 trialId: string;
54 /** Parent Study resource name. */
55 parent: string;
56 /** Region id. */
57 location: string;
58 /** Project id. */
59 project: string;
60 /** Client identifier. */
61 clientId: string | undefined;
62 /** Parameter assignments. */
63 parameters: TrialParameter[];
64 /** Server-reported state. */
65 state: string | undefined;
66 /** Final measurement, if the trial completed. */
67 finalMeasurement: aiplatform.GoogleCloudAiplatformV1Measurement | undefined;
68 /** RFC3339 start timestamp. */
69 startTime: string | undefined;
70 /** RFC3339 end timestamp. */
71 endTime: string | undefined;
72 },
73 never,
74 Providers
75>;
76
77/**
78 * A user-provided Vertex AI Vizier Trial attached to a Study.
79 *
80 * Trials have no labels and no update RPC. Alchemy records the logical
81 * id in `clientId` so `list` can find them. Parent and parameters are
82 * immutable.
83 *
84 * ### Creating a Trial
85 * **Example:** User-provided trial
86 * ```typescript
87 * const trial = yield* GCP.AIPlatform.StudiesTrial("Seed", {
88 * parent: study.name,
89 * parameters: [{ parameterId: "learning_rate", value: 0.01 }],
90 * });
91 * ```
92 *
93 * @resource
94 * @category AIPlatform
95 */
96export const StudiesTrial = Resource<StudiesTrial>(
97 "GCP.AIPlatform.StudiesTrial",
98);
99
100export class StudiesTrialNotResolved extends Data.TaggedError(
101 "GCP.AIPlatform.StudiesTrialNotResolved",
102)<{
103 name: string;
104}> {}
105
106const toAttrs = (
107 trial: aiplatform.GoogleCloudAiplatformV1Trial,
108 project: string,
109) => {
110 const name = trial.name ?? "";
111 return {
112 name,
113 trialId: trial.id ?? lastSegment(name),
114 parent: parentOf(name),
115 location: locationOf(name),
116 project,
117 clientId: trial.clientId,
118 parameters: (trial.parameters ?? []).map((parameter) => ({
119 parameterId: parameter.parameterId ?? "",
120 value: parameter.value,
121 })),
122 state: trial.state,
123 finalMeasurement: trial.finalMeasurement,
124 startTime: trial.startTime,
125 endTime: trial.endTime,
126 };
127};
128
129const getByName = (name: string) =>
130 name.length === 0
131 ? Effect.succeed(undefined)
132 : aiplatform
133 .getProjectsLocationsStudiesTrials({ name })
134 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
135
136const listStudies = (parent: string) =>
137 aiplatform.listProjectsLocationsStudies.pages({ parent, pageSize: 100 }).pipe(
138 Stream.flatMap((page) => Stream.fromIterable(page.studies ?? [])),
139 Stream.filter((study) => hasStudyOwnership(study.displayName)),
140 Stream.map((study) => study.name ?? ""),
141 Stream.filter((name) => name.length > 0),
142 Stream.runCollect,
143 Effect.map((chunk) => Array.from(chunk)),
144 Effect.catchTag("NotFound", () => Effect.succeed([] as string[])),
145 );
146
147const listAtParent = (parent: string, project: string, clientId?: string) =>
148 aiplatform.listProjectsLocationsStudiesTrials
149 .pages({ parent, pageSize: 100 })
150 .pipe(
151 Stream.flatMap((page) => Stream.fromIterable(page.trials ?? [])),
152 Stream.filter((trial) =>
153 clientId !== undefined ? trial.clientId === clientId : true,
154 ),
155 Stream.map((trial) => toAttrs(trial, project)),
156 Stream.runCollect,
157 Effect.map((chunk) => Array.from(chunk)),
158 Effect.catchTag("NotFound", () => Effect.succeed([])),
159 );
160
161const findByClientId = (parent: string, clientId: string, project: string) =>
162 listAtParent(parent, project, clientId).pipe(
163 Effect.map((trials) => trials[0]),
164 );
165
166export const StudiesTrialProvider = () =>
167 Provider.succeed(StudiesTrial, {
168 stables: ["name", "trialId", "parent", "location", "project", "startTime"],
169
170 diff: Effect.fn(function* ({ news, olds, output }) {
171 if (!isResolved(news)) return undefined;
172 const previousParent = olds?.parent ?? output?.parent;
173 if (previousParent !== undefined && news.parent !== previousParent) {
174 return { action: "replace" as const, deleteFirst: false };
175 }
176 return undefined;
177 }),
178
179 read: Effect.fn(function* ({ id, olds, output }) {
180 const env = yield* GcpEnvironment.current;
181 const existing = yield* getByName(output?.name ?? "");
182 if (existing !== undefined) {
183 return toAttrs(existing, env.project);
184 }
185 if (olds?.parent === undefined) return undefined;
186 const ownership = yield* createInternalLabels(id);
187 const clientId = olds.clientId ?? `alchemy-${ownership["alchemy-id"]}`;
188 const match = yield* findByClientId(olds.parent, clientId, env.project);
189 if (match === undefined) return undefined;
190 const fetched = yield* getByName(match.name);
191 if (fetched === undefined) return undefined;
192 const attrs = toAttrs(fetched, env.project);
193 return fetched.clientId === clientId ? attrs : Unowned(attrs);
194 }),
195
196 list: () =>
197 Effect.gen(function* () {
198 const env = yield* GcpEnvironment.current;
199 const studies = (yield* Effect.forEach(
200 listLocations(env.region),
201 (location) => listStudies(locationParent(env.project, location)),
202 )).flat();
203 const pages = yield* Effect.forEach(
204 studies,
205 (parent) => listAtParent(parent, env.project),
206 { concurrency: 4 },
207 );
208 return pages.flat();
209 }),
210
211 reconcile: Effect.fn(function* ({ id, news, output }) {
212 const env = yield* GcpEnvironment.current;
213 const ownership = yield* createInternalLabels(id);
214 const clientId = news.clientId ?? `alchemy-${ownership["alchemy-id"]}`;
215
216 let current = yield* getByName(output?.name ?? "");
217 if (current === undefined) {
218 const match = yield* findByClientId(news.parent, clientId, env.project);
219 if (match !== undefined) {
220 current = yield* getByName(match.name);
221 }
222 }
223
224 if (current === undefined) {
225 const created = yield* aiplatform
226 .createProjectsLocationsStudiesTrials({
227 parent: news.parent,
228 body: {
229 clientId,
230 // Vizier rejects user trials without a status ("Can not add
231 // trial with status: UNKNOWN").
232 state: "ACTIVE",
233 parameters: news.parameters?.map((parameter) => ({
234 parameterId: parameter.parameterId,
235 value: parameter.value,
236 })),
237 },
238 })
239 .pipe(
240 Effect.catchTag("Conflict", () =>
241 findByClientId(news.parent, clientId, env.project).pipe(
242 Effect.flatMap((match) =>
243 match ? getByName(match.name) : Effect.succeed(undefined),
244 ),
245 ),
246 ),
247 );
248 current = created ?? undefined;
249 }
250
251 if (current === undefined) {
252 return yield* new StudiesTrialNotResolved({
253 name: output?.name ?? `${news.parent}/trials/-`,
254 });
255 }
256
257 return toAttrs(current, env.project);
258 }),
259
260 delete: Effect.fn(function* ({ output }) {
261 yield* aiplatform
262 .deleteProjectsLocationsStudiesTrials({ name: output.name })
263 .pipe(Effect.catchTag("NotFound", () => Effect.void));
264 }),
265 });