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 {
13 createInternalLabels,
14 hasAlchemyLabels,
15 stripInternalLabels,
16 toLabels,
17} from "../Labels.ts";
18import type { Providers } from "../Providers.ts";
19import {
20 compact,
21 normalizeLocation,
22 parentOf,
23 parseName,
24 toPhysicalId,
25 listLocations,
26} from "./names.ts";
27import { waitForOperation } from "./operations.ts";
28
29export type TrainingPipelineEncryptionSpec =
30 aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
31export type TrainingPipelineInputDataConfig =
32 aiplatform.GoogleCloudAiplatformV1InputDataConfig;
33export type TrainingPipelineModel = aiplatform.GoogleCloudAiplatformV1Model;
34
35export type TrainingPipelineProps = {
36 /**
37 * Training pipeline id (the `{training_pipeline}` segment of
38 * `projects/{project}/locations/{location}/trainingPipelines/{training_pipeline}`).
39 * Vertex assigns this on create. Pass the value from a previous deploy
40 * to target the same pipeline. Immutable — changing it replaces the
41 * pipeline.
42 */
43 trainingPipelineId?: string;
44 /**
45 * Vertex AI location (`us-central1`, `us-east1`, …). Immutable —
46 * changing it replaces the pipeline. `US-CENTRAL1` is accepted and
47 * normalized to `us-central1`.
48 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
49 */
50 location?: string;
51 /**
52 * User-defined display name. Required by Vertex. If omitted, a unique
53 * name is generated from the stack, stage, and logical id.
54 */
55 displayName?: string;
56 /**
57 * Cloud Storage URI of the training-task YAML
58 * (`gs://google-cloud-aiplatform/schema/trainingjob/definition/…`).
59 * Immutable — changing it replaces the pipeline.
60 */
61 trainingTaskDefinition: string;
62 /**
63 * Training-task inputs as specified by `trainingTaskDefinition`.
64 * Immutable — changing it replaces the pipeline.
65 */
66 trainingTaskInputs?: unknown;
67 /**
68 * Vertex Dataset split and export configuration used by the training
69 * task.
70 */
71 inputDataConfig?: TrainingPipelineInputDataConfig;
72 /**
73 * Model uploaded when the pipeline succeeds.
74 */
75 modelToUpload?: TrainingPipelineModel;
76 /**
77 * Existing model to version instead of uploading a new model.
78 */
79 parentModel?: string;
80 /**
81 * Id of the uploaded model (last path segment). Letters, digits,
82 * underscores, and hyphens; cannot start with a digit or hyphen.
83 */
84 modelId?: string;
85 /**
86 * Customer-managed encryption key for the pipeline (and the uploaded
87 * model when `modelToUpload` is omitted).
88 */
89 encryptionSpec?: TrainingPipelineEncryptionSpec;
90 /**
91 * User labels. Alchemy ownership labels are merged in automatically.
92 */
93 labels?: Record<string, string>;
94};
95
96export type TrainingPipeline = Resource<
97 "GCP.AIPlatform.TrainingPipeline",
98 TrainingPipelineProps,
99 {
100 /** Full resource name `projects/{project}/locations/{location}/trainingPipelines/{training_pipeline}`. */
101 name: string;
102 /** Training pipeline id (last path segment). */
103 trainingPipelineId: string;
104 /** Location id (`us-central1`, …). */
105 location: string;
106 /** Project id. */
107 project: string;
108 /** User-defined display name. */
109 displayName: string | undefined;
110 /** Training-task YAML URI. */
111 trainingTaskDefinition: string | undefined;
112 /** Training-task inputs. */
113 trainingTaskInputs: unknown;
114 /** Training-task metadata populated while the pipeline runs. */
115 trainingTaskMetadata: unknown;
116 /** Dataset split and export configuration. */
117 inputDataConfig: TrainingPipelineInputDataConfig | undefined;
118 /** Model uploaded by the pipeline, if any. */
119 modelToUpload: TrainingPipelineModel | undefined;
120 /** Parent model versioned by this pipeline. */
121 parentModel: string | undefined;
122 /** Uploaded model id. */
123 modelId: string | undefined;
124 /** CMEK spec. */
125 encryptionSpec: TrainingPipelineEncryptionSpec | undefined;
126 /** User labels (Alchemy ownership labels stripped). */
127 labels: Record<string, string>;
128 /** Pipeline state (`PIPELINE_STATE_RUNNING`, …). */
129 state: string | undefined;
130 /** Error when the pipeline failed or was cancelled. */
131 error: aiplatform.GoogleRpcStatus | undefined;
132 /** RFC3339 creation timestamp. */
133 createTime: string | undefined;
134 /** RFC3339 start timestamp. */
135 startTime: string | undefined;
136 /** RFC3339 end timestamp. */
137 endTime: string | undefined;
138 /** RFC3339 last-update timestamp. */
139 updateTime: string | undefined;
140 },
141 never,
142 Providers
143>;
144
145/**
146 * A Vertex AI TrainingPipeline that runs a training task and optionally
147 * uploads a Model.
148 *
149 * Creating a TrainingPipeline starts it immediately. Vertex assigns the
150 * resource id; there is no update RPC — changing `location` or
151 * `trainingTaskDefinition` replaces the pipeline. Labels brand the
152 * resource for `list` / nuke.
153 *
154 * ### Creating a TrainingPipeline
155 * **Example:** Custom training task
156 * ```typescript
157 * const pipeline = yield* GCP.AIPlatform.TrainingPipeline("Train", {
158 * location: "us-central1",
159 * displayName: "custom-train",
160 * trainingTaskDefinition:
161 * "gs://google-cloud-aiplatform/schema/trainingjob/definition/custom_task_1.0.0.yaml",
162 * trainingTaskInputs: {
163 * workerPoolSpecs: [
164 * {
165 * machineSpec: { machineType: "n1-standard-4" },
166 * replicaCount: "1",
167 * containerSpec: { imageUri: "gcr.io/my-project/trainer:latest" },
168 * },
169 * ],
170 * },
171 * labels: { env: "dev" },
172 * });
173 * ```
174 *
175 * @resource
176 * @category AIPlatform
177 */
178export const TrainingPipeline = Resource<TrainingPipeline>(
179 "GCP.AIPlatform.TrainingPipeline",
180);
181
182export class TrainingPipelineNotResolved extends Data.TaggedError(
183 "GCP.AIPlatform.TrainingPipelineNotResolved",
184)<{
185 name: string;
186}> {}
187
188export class TrainingPipelineStillExists extends Data.TaggedError(
189 "GCP.AIPlatform.TrainingPipelineStillExists",
190)<{
191 name: string;
192}> {}
193
194const resourceName = (
195 project: string,
196 location: string,
197 trainingPipelineId: string,
198) =>
199 `projects/${project}/locations/${location}/trainingPipelines/${trainingPipelineId}`;
200
201const userLabels = (
202 labels: Record<string, string | undefined> | null | undefined,
203): Record<string, string> => stripInternalLabels(tagRecord(labels));
204
205const toAttrs = (
206 pipeline: aiplatform.GoogleCloudAiplatformV1TrainingPipeline,
207 project: string,
208) => {
209 const name = pipeline.name ?? "";
210 const parsed = parseName(name, "trainingPipelines");
211 return {
212 name,
213 trainingPipelineId: parsed.resourceId,
214 location: parsed.location,
215 project: parsed.project || project,
216 displayName: pipeline.displayName,
217 trainingTaskDefinition: pipeline.trainingTaskDefinition,
218 trainingTaskInputs: pipeline.trainingTaskInputs,
219 trainingTaskMetadata: pipeline.trainingTaskMetadata,
220 inputDataConfig: pipeline.inputDataConfig,
221 modelToUpload: pipeline.modelToUpload,
222 parentModel: pipeline.parentModel,
223 modelId: pipeline.modelId,
224 encryptionSpec: pipeline.encryptionSpec,
225 labels: userLabels(pipeline.labels),
226 state: pipeline.state,
227 error: pipeline.error,
228 createTime: pipeline.createTime,
229 startTime: pipeline.startTime,
230 endTime: pipeline.endTime,
231 updateTime: pipeline.updateTime,
232 };
233};
234
235const getByName = (name: string) =>
236 aiplatform
237 .getProjectsLocationsTrainingPipelines({ name })
238 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
239
240const listAt = (parent: string) =>
241 aiplatform.listProjectsLocationsTrainingPipelines
242 .pages({ parent, pageSize: 100 })
243 .pipe(
244 Stream.flatMap((page) =>
245 Stream.fromIterable(page.trainingPipelines ?? []),
246 ),
247 Stream.take(500),
248 Stream.runCollect,
249 Effect.map((chunk) => Array.from(chunk)),
250 Effect.catchTag("NotFound", () => Effect.succeed([])),
251 );
252
253const findOwned = (id: string, parent: string) =>
254 Effect.gen(function* () {
255 const expected = yield* createInternalLabels(id);
256 const pipelines = yield* listAt(parent);
257 return pipelines.find((pipeline) =>
258 Object.entries(expected).every(
259 ([key, value]) => (pipeline.labels ?? {})[key] === value,
260 ),
261 );
262 });
263
264const waitUntilGone = (name: string) =>
265 getByName(name).pipe(
266 Effect.flatMap((pipeline) =>
267 pipeline === undefined
268 ? Effect.void
269 : Effect.fail(new TrainingPipelineStillExists({ name })),
270 ),
271 Effect.retry({
272 while: (error) =>
273 error._tag === "GCP.AIPlatform.TrainingPipelineStillExists",
274 times: 30,
275 schedule: Schedule.spaced("5 seconds"),
276 }),
277 );
278
279const toBody = (
280 news: TrainingPipelineProps,
281 displayName: string,
282 labels: Record<string, string>,
283): aiplatform.GoogleCloudAiplatformV1TrainingPipeline =>
284 compact({
285 displayName,
286 trainingTaskDefinition: news.trainingTaskDefinition,
287 trainingTaskInputs: news.trainingTaskInputs,
288 inputDataConfig: news.inputDataConfig,
289 modelToUpload: news.modelToUpload,
290 parentModel: news.parentModel,
291 modelId: news.modelId,
292 encryptionSpec: news.encryptionSpec,
293 labels,
294 });
295
296export const TrainingPipelineProvider = () =>
297 Provider.succeed(TrainingPipeline, {
298 stables: [
299 "name",
300 "trainingPipelineId",
301 "location",
302 "project",
303 "createTime",
304 ],
305
306 diff: Effect.fn(function* ({ news, olds, output }) {
307 const env = yield* GcpEnvironment.current;
308 if (!isResolved(news)) return undefined;
309 const previousId = olds?.trainingPipelineId ?? output?.trainingPipelineId;
310 const nextId = news.trainingPipelineId ?? previousId;
311 const idChanged =
312 previousId !== undefined &&
313 nextId !== undefined &&
314 nextId !== previousId;
315 const previousLocation = normalizeLocation(
316 olds?.location ?? output?.location,
317 env.region,
318 );
319 const nextLocation = normalizeLocation(
320 news.location ?? output?.location,
321 env.region,
322 );
323 const definitionChanged =
324 (olds?.trainingTaskDefinition ?? output?.trainingTaskDefinition) !==
325 undefined &&
326 news.trainingTaskDefinition !==
327 (olds?.trainingTaskDefinition ?? output?.trainingTaskDefinition);
328 if (idChanged || previousLocation !== nextLocation || definitionChanged) {
329 return { action: "replace" as const, deleteFirst: false };
330 }
331 return undefined;
332 }),
333
334 read: Effect.fn(function* ({ id, olds, output }) {
335 const env = yield* GcpEnvironment.current;
336 const location = normalizeLocation(
337 olds?.location ?? output?.location,
338 env.region,
339 );
340 // Pipeline ids are server-assigned numbers, so only a recorded name
341 // or a caller-supplied id can be fetched directly.
342 const trainingPipelineId =
343 olds?.trainingPipelineId ?? output?.trainingPipelineId;
344 const name =
345 output?.name ??
346 (trainingPipelineId !== undefined
347 ? resourceName(env.project, location, trainingPipelineId)
348 : undefined);
349 const existing =
350 (name !== undefined ? yield* getByName(name) : undefined) ??
351 (yield* findOwned(id, parentOf(env.project, location)));
352 if (existing === undefined) return undefined;
353 const attrs = toAttrs(existing, env.project);
354 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
355 ? attrs
356 : Unowned(attrs);
357 }),
358
359 list: () =>
360 Effect.gen(function* () {
361 const env = yield* GcpEnvironment.current;
362 const pipelines = (yield* Effect.forEach(
363 listLocations(env.region),
364 (location) => listAt(parentOf(env.project, location)),
365 )).flat();
366 return pipelines
367 .filter((pipeline) =>
368 Object.keys(pipeline.labels ?? {}).some((key) =>
369 key.startsWith("alchemy-"),
370 ),
371 )
372 .map((pipeline) => toAttrs(pipeline, env.project));
373 }),
374
375 reconcile: Effect.fn(function* ({ id, news, output }) {
376 const env = yield* GcpEnvironment.current;
377 const location = normalizeLocation(
378 news.location ?? output?.location,
379 env.region,
380 );
381 const parent = parentOf(env.project, location);
382 const desiredLabels = {
383 ...toLabels(news.labels),
384 ...(yield* createInternalLabels(id)),
385 };
386 const displayName =
387 news.displayName ??
388 (yield* toPhysicalId(
389 id,
390 news.trainingPipelineId,
391 output?.trainingPipelineId,
392 ));
393
394 let current =
395 (output?.name !== undefined
396 ? yield* getByName(output.name)
397 : undefined) ?? (yield* findOwned(id, parent));
398
399 if (current === undefined) {
400 const created = yield* aiplatform
401 .createProjectsLocationsTrainingPipelines({
402 parent,
403 body: toBody(news, displayName, desiredLabels),
404 })
405 .pipe(Effect.catchTag("Conflict", () => findOwned(id, parent)));
406 current = created ?? undefined;
407 }
408
409 if (current === undefined) {
410 return yield* new TrainingPipelineNotResolved({
411 name: output?.name ?? parent,
412 });
413 }
414
415 return toAttrs(current, env.project);
416 }),
417
418 delete: Effect.fn(function* ({ output }) {
419 const name = output.name;
420 // Custom training runs in a backing CustomJob that outlives the
421 // pipeline; remember it so it can be deleted too.
422 const metadata = (yield* getByName(name))?.trainingTaskMetadata;
423 const backingJob =
424 metadata !== null &&
425 typeof metadata === "object" &&
426 "backingCustomJob" in metadata
427 ? (metadata as { backingCustomJob?: unknown }).backingCustomJob
428 : undefined;
429 yield* aiplatform
430 .cancelProjectsLocationsTrainingPipelines({ name, body: {} })
431 .pipe(
432 Effect.catchTag("NotFound", () => Effect.void),
433 Effect.catchTag("BadRequest", () => Effect.void),
434 Effect.catchTag("Conflict", () => Effect.void),
435 );
436 // Delete is rejected until the cancellation reaches a terminal state,
437 // which can take a couple of minutes.
438 const operation = yield* aiplatform
439 .deleteProjectsLocationsTrainingPipelines({ name })
440 .pipe(
441 Effect.retry({
442 while: (error) =>
443 error._tag === "Conflict" || error._tag === "BadRequest",
444 times: 24,
445 schedule: Schedule.spaced("10 seconds"),
446 }),
447 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
448 );
449 if (operation !== undefined) {
450 yield* waitForOperation(operation, { notFoundOk: true });
451 }
452 yield* waitUntilGone(name);
453 if (typeof backingJob === "string" && backingJob.length > 0) {
454 const jobDelete = yield* aiplatform
455 .deleteProjectsLocationsCustomJobs({ name: backingJob })
456 .pipe(
457 Effect.retry({
458 while: (error) => error._tag === "BadRequest",
459 times: 8,
460 schedule: Schedule.spaced("3 seconds"),
461 }),
462 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
463 );
464 if (jobDelete !== undefined) {
465 yield* waitForOperation(jobDelete, { notFoundOk: true });
466 }
467 }
468 }),
469 });