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 { createPhysicalName } from "../../PhysicalName.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { tagRecord } from "../../Tags.ts";
12import { GcpEnvironment } from "../Environment.ts";
13import { createInternalLabels, hasAlchemyLabels, toLabels } from "../Labels.ts";
14import type { Providers } from "../Providers.ts";
15import { listLocations } from "./names.ts";
16import {
17 hasAlchemyLabelKeys,
18 isJobTerminal,
19 normalizeLocation,
20 parentOf,
21 parseResourceName,
22 rfc1035,
23 userLabels,
24} from "./internal.ts";
25
26export type TuningJobProps = {
27 /**
28 * Vertex AI location. Immutable — changing it replaces the job.
29 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
30 */
31 location?: string;
32 /**
33 * Description of the tuning job.
34 */
35 description?: string;
36 /**
37 * Display name of the tuned model (max 128 Unicode characters).
38 */
39 tunedModelDisplayName?: string;
40 /**
41 * Base model being tuned (for example `gemini-2.0-flash-001`).
42 */
43 baseModel?: string;
44 /**
45 * Supervised fine-tuning spec. Required unless `preferenceOptimizationSpec`
46 * or `preTunedModel` is set.
47 */
48 supervisedTuningSpec?: aiplatform.GoogleCloudAiplatformV1SupervisedTuningSpec;
49 /**
50 * Preference-optimization spec.
51 */
52 preferenceOptimizationSpec?: aiplatform.GoogleCloudAiplatformV1PreferenceOptimizationSpec;
53 /**
54 * Pre-tuned model for continuous tuning.
55 */
56 preTunedModel?: aiplatform.GoogleCloudAiplatformV1PreTunedModel;
57 /**
58 * Runtime service account email.
59 */
60 serviceAccount?: string;
61 /**
62 * User labels. Alchemy ownership labels are merged in automatically.
63 */
64 labels?: Record<string, string>;
65 /**
66 * Customer-managed encryption key. Immutable.
67 */
68 encryptionSpec?: aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
69};
70
71export type TuningJob = Resource<
72 "GCP.AIPlatform.TuningJob",
73 TuningJobProps,
74 {
75 /** Full resource name. */
76 name: string;
77 /** Job id (last path segment). */
78 tuningJobId: string;
79 /** Project id. */
80 project: string;
81 /** Location. */
82 location: string;
83 /** Description. */
84 description: string | undefined;
85 /** Tuned-model display name. */
86 tunedModelDisplayName: string | undefined;
87 /** Base model. */
88 baseModel: string | undefined;
89 /** User labels (Alchemy ownership labels stripped). */
90 labels: Record<string, string>;
91 /** Job state. */
92 state: string | undefined;
93 /** RFC3339 creation timestamp. */
94 createTime: string | undefined;
95 /** RFC3339 start timestamp. */
96 startTime: string | undefined;
97 /** RFC3339 end timestamp. */
98 endTime: string | undefined;
99 /** RFC3339 last-update timestamp. */
100 updateTime: string | undefined;
101 /** Failure/cancel status, if any. */
102 error: { code?: number; message?: string } | undefined;
103 },
104 never,
105 Providers
106>;
107
108/**
109 * A Vertex AI TuningJob (supervised fine-tune / preference optimization).
110 *
111 * Creating a TuningJob starts it immediately. There is no update API, so
112 * reconcile is observe-ensure (create if missing). There is no delete RPC —
113 * destroy cancels the job. Alchemy ownership labels are merged into `labels`
114 * so `list` / nuke can find it.
115 *
116 * ### Creating a Tuning Job
117 * **Example:** Supervised fine-tune
118 * ```typescript
119 * const job = yield* GCP.AIPlatform.TuningJob("Tune", {
120 * baseModel: "gemini-2.0-flash-001",
121 * supervisedTuningSpec: {
122 * trainingDatasetUri: "gs://bucket/train.jsonl",
123 * },
124 * });
125 * ```
126 *
127 * @resource
128 * @category AIPlatform
129 */
130export const TuningJob = Resource<TuningJob>("GCP.AIPlatform.TuningJob");
131
132export class TuningJobNotResolved extends Data.TaggedError(
133 "GCP.AIPlatform.TuningJobNotResolved",
134)<{
135 name: string;
136}> {}
137
138const toDisplayName = (id: string, existing?: string) =>
139 Effect.gen(function* () {
140 if (existing !== undefined) return existing;
141 return rfc1035(
142 yield* createPhysicalName({
143 id,
144 maxLength: 63,
145 lowercase: true,
146 }),
147 );
148 });
149
150const toAttrs = (
151 job: aiplatform.GoogleCloudAiplatformV1TuningJob,
152 project: string,
153) => {
154 const name = job.name ?? "";
155 const parsed = parseResourceName(name, "tuningJobs");
156 return {
157 name,
158 tuningJobId: parsed.id,
159 project: parsed.project || project,
160 location: parsed.location,
161 description: job.description,
162 tunedModelDisplayName: job.tunedModelDisplayName,
163 baseModel: job.baseModel,
164 labels: userLabels(job.labels),
165 state: typeof job.state === "string" ? job.state : undefined,
166 createTime: job.createTime,
167 startTime: job.startTime,
168 endTime: job.endTime,
169 updateTime: job.updateTime,
170 error: job.error
171 ? { code: job.error.code, message: job.error.message }
172 : undefined,
173 };
174};
175
176const getByName = (name: string) =>
177 aiplatform
178 .getProjectsLocationsTuningJobs({ name })
179 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
180
181const listJobs = (project: string, region: string) => {
182 const collect = (parent: string) =>
183 aiplatform.listProjectsLocationsTuningJobs
184 .pages({ parent, pageSize: 100 })
185 .pipe(
186 Stream.flatMap((page) => Stream.fromIterable(page.tuningJobs ?? [])),
187 Stream.filter((job) => hasAlchemyLabelKeys(job.labels)),
188 Stream.runCollect,
189 Effect.map((chunk) => Array.from(chunk)),
190 );
191 const fallback = Effect.forEach(listLocations(region), (location) =>
192 collect(`projects/${project}/locations/${location}`),
193 ).pipe(Effect.map((pages) => pages.flat()));
194 // Vertex AI has no `locations/-` wildcard; scan known locations.
195 return fallback.pipe(Effect.catchTag("NotFound", () => Effect.succeed([])));
196};
197
198const findOwned = (
199 id: string,
200 project: string,
201 region: string,
202 hinted?: string,
203) =>
204 Effect.gen(function* () {
205 if (hinted !== undefined && hinted.length > 0) {
206 const existing = yield* getByName(hinted);
207 if (existing !== undefined) return existing;
208 }
209 const jobs = yield* listJobs(project, region);
210 for (const job of jobs) {
211 if (yield* hasAlchemyLabels(id, tagRecord(job.labels))) return job;
212 }
213 return undefined as aiplatform.GoogleCloudAiplatformV1TuningJob | undefined;
214 });
215
216const cancelJob = (name: string) =>
217 Effect.gen(function* () {
218 const existing = yield* getByName(name);
219 if (existing === undefined) return;
220 if (
221 !isJobTerminal(
222 typeof existing.state === "string" ? existing.state : undefined,
223 )
224 ) {
225 yield* aiplatform
226 .cancelProjectsLocationsTuningJobs({ name })
227 .pipe(Effect.catchTag("NotFound", () => Effect.void));
228 yield* getByName(name).pipe(
229 Effect.filterOrFail(
230 (job) =>
231 job === undefined ||
232 isJobTerminal(
233 typeof job.state === "string" ? job.state : undefined,
234 ),
235 () => new TuningJobNotResolved({ name }),
236 ),
237 Effect.retry({
238 while: (error) =>
239 error._tag === "GCP.AIPlatform.TuningJobNotResolved",
240 times: 8,
241 schedule: Schedule.spaced("4 seconds"),
242 }),
243 Effect.catchTag(
244 "GCP.AIPlatform.TuningJobNotResolved",
245 () => Effect.void,
246 ),
247 );
248 }
249 });
250
251export const TuningJobProvider = () =>
252 Provider.succeed(TuningJob, {
253 stables: ["name", "tuningJobId", "project", "location", "createTime"],
254
255 diff: Effect.fn(function* ({ news, olds, output }) {
256 const env = yield* GcpEnvironment.current;
257 if (!isResolved(news)) return undefined;
258 const previousLocation = normalizeLocation(
259 olds?.location ?? output?.location,
260 env.region,
261 );
262 const nextLocation = normalizeLocation(
263 news.location ?? olds?.location ?? output?.location,
264 env.region,
265 );
266 if (previousLocation !== nextLocation) {
267 return { action: "replace" as const, deleteFirst: false };
268 }
269 return undefined;
270 }),
271
272 read: Effect.fn(function* ({ id, olds, output }) {
273 const env = yield* GcpEnvironment.current;
274 const existing = yield* findOwned(
275 id,
276 env.project,
277 env.region,
278 output?.name,
279 );
280 if (existing === undefined) return undefined;
281 const attrs = toAttrs(existing, env.project);
282 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
283 ? attrs
284 : Unowned(attrs);
285 }),
286
287 list: () =>
288 Effect.gen(function* () {
289 const env = yield* GcpEnvironment.current;
290 const jobs = yield* listJobs(env.project, env.region);
291 return jobs.map((job) => toAttrs(job, env.project));
292 }),
293
294 reconcile: Effect.fn(function* ({ id, news, output }) {
295 const env = yield* GcpEnvironment.current;
296 const location = normalizeLocation(
297 news.location ?? output?.location,
298 env.region,
299 );
300 const tunedModelDisplayName =
301 news.tunedModelDisplayName ??
302 (yield* toDisplayName(id, output?.tuningJobId));
303 const desiredLabels = {
304 ...toLabels(news.labels),
305 ...(yield* createInternalLabels(id)),
306 };
307
308 let current = yield* findOwned(id, env.project, env.region, output?.name);
309
310 if (current === undefined) {
311 const created = yield* aiplatform
312 .createProjectsLocationsTuningJobs({
313 parent: parentOf(env.project, location),
314 body: {
315 description: news.description,
316 tunedModelDisplayName,
317 baseModel: news.baseModel,
318 supervisedTuningSpec: news.supervisedTuningSpec,
319 preferenceOptimizationSpec: news.preferenceOptimizationSpec,
320 preTunedModel: news.preTunedModel,
321 serviceAccount: news.serviceAccount,
322 labels: desiredLabels,
323 encryptionSpec: news.encryptionSpec,
324 },
325 })
326 .pipe(
327 Effect.catchTag("Conflict", () =>
328 findOwned(id, env.project, env.region),
329 ),
330 );
331 current = created ?? undefined;
332 }
333
334 if (current === undefined) {
335 return yield* new TuningJobNotResolved({
336 name: output?.name ?? parentOf(env.project, location),
337 });
338 }
339
340 return toAttrs(current, env.project);
341 }),
342
343 delete: Effect.fn(function* ({ output }) {
344 yield* cancelJob(output.name);
345 }),
346 });