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 { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { tagRecord } from "../../Tags.ts";
10import { GcpEnvironment } from "../Environment.ts";
11import { createInternalLabels, hasAlchemyLabels, toLabels } from "../Labels.ts";
12import type { Providers } from "../Providers.ts";
13import { listLocations } from "./names.ts";
14import { waitForOperation } from "./operations.ts";
15import {
16 AiPlatformStillExists,
17 MAX_PIPELINE_JOB_ID_LENGTH,
18 collectPages,
19 jsonEqual,
20 locationParent,
21 normalizeLocation,
22 parseResourceName,
23 toPhysicalId,
24 userLabels,
25 type EncryptionSpec,
26} from "./shared.ts";
27
28const COLLECTION = "pipelineJobs";
29
30export type PipelineRuntimeConfig = {
31 /** Cloud Storage root for pipeline artifacts. */
32 gcsOutputDirectory?: string;
33 /** Runtime parameters (schema 2.1.0+). */
34 parameterValues?: Record<string, unknown>;
35 /** Failure policy (`PIPELINE_FAILURE_POLICY_FAIL_SLOW` / `_FAIL_FAST`). */
36 failurePolicy?: string;
37 /** Input artifacts keyed by name. */
38 inputArtifacts?: Record<string, { artifactId?: string }>;
39};
40
41export type PipelineJobProps = {
42 /**
43 * Job id. If omitted, a unique RFC1035 name is generated. Less than
44 * 128 characters, matching `/[a-z0-9-]+/`. Immutable — changing it
45 * replaces the job (a new run is started).
46 */
47 pipelineJobId?: string;
48 /**
49 * Vertex AI location. Immutable — changing it replaces the job.
50 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
51 */
52 location?: string;
53 /**
54 * Display name (max 128 UTF-8 characters). Defaults to the job id.
55 */
56 displayName?: string;
57 /**
58 * User labels. Alchemy ownership labels are merged in automatically.
59 */
60 labels?: Record<string, string>;
61 /**
62 * Compiled Kubeflow pipeline spec as a JSON object.
63 */
64 pipelineSpec?: Record<string, unknown>;
65 /**
66 * Template URI (Vertex Template Registry or Gallery). Used when
67 * `pipelineSpec` is omitted.
68 */
69 templateUri?: string;
70 /**
71 * Runtime config (output directory, parameters, failure policy).
72 */
73 runtimeConfig?: PipelineRuntimeConfig;
74 /**
75 * Service account the pipeline runs as.
76 */
77 serviceAccount?: string;
78 /**
79 * VPC network to peer.
80 */
81 network?: string;
82 /**
83 * Reserved IP ranges under the VPC.
84 */
85 reservedIpRanges?: string[];
86 /**
87 * Customer-managed encryption key.
88 */
89 encryptionSpec?: EncryptionSpec;
90 /**
91 * Run component-level validations before create.
92 */
93 preflightValidations?: boolean;
94};
95
96export type PipelineJob = Resource<
97 "GCP.AIPlatform.PipelineJob",
98 PipelineJobProps,
99 {
100 /** Full resource name. */
101 name: string;
102 /** Job id (last path segment). */
103 pipelineJobId: string;
104 /** Project id. */
105 project: string;
106 /** Location id. */
107 location: string;
108 /** Display name. */
109 displayName: string | undefined;
110 /** User labels (Alchemy ownership labels stripped). */
111 labels: Record<string, string>;
112 /** Pipeline state (`PIPELINE_STATE_RUNNING`, …). */
113 state: string | undefined;
114 /** Template URI. */
115 templateUri: string | undefined;
116 /** Service account. */
117 serviceAccount: string | undefined;
118 /** RFC3339 creation timestamp. */
119 createTime: string | undefined;
120 /** RFC3339 start timestamp. */
121 startTime: string | undefined;
122 /** RFC3339 end timestamp. */
123 endTime: string | undefined;
124 /** RFC3339 last-update timestamp. */
125 updateTime: string | undefined;
126 },
127 never,
128 Providers
129>;
130
131/**
132 * A Vertex AI PipelineJob — a compiled Kubeflow pipeline that runs
133 * immediately on create.
134 *
135 * There is no update API. Changing identity (`pipelineJobId`, `location`)
136 * or the pipeline spec / template URI replaces the job.
137 *
138 * ### Creating a Pipeline Job
139 * **Example:** Run a compiled spec
140 * ```typescript
141 * const job = yield* GCP.AIPlatform.PipelineJob("Train", {
142 * pipelineSpec: compiled,
143 * runtimeConfig: { gcsOutputDirectory: "gs://bucket/pipeline-out" },
144 * labels: { env: "prod" },
145 * });
146 * ```
147 *
148 * @resource
149 * @category AIPlatform
150 */
151export const PipelineJob = Resource<PipelineJob>("GCP.AIPlatform.PipelineJob");
152
153export class PipelineJobNotResolved extends Data.TaggedError(
154 "GCP.AIPlatform.PipelineJobNotResolved",
155)<{
156 name: string;
157}> {}
158
159const resourceName = (project: string, location: string, id: string) =>
160 `${locationParent(project, location)}/${COLLECTION}/${id}`;
161
162const toAttrs = (
163 job: aiplatform.GoogleCloudAiplatformV1PipelineJob,
164 project: string,
165) => {
166 const name = job.name ?? "";
167 const parsed = parseResourceName(name, COLLECTION);
168 return {
169 name,
170 pipelineJobId: parsed.id,
171 project: parsed.project || project,
172 location: parsed.location,
173 displayName: job.displayName,
174 labels: userLabels(job.labels),
175 state: job.state,
176 templateUri: job.templateUri,
177 serviceAccount: job.serviceAccount,
178 createTime: job.createTime,
179 startTime: job.startTime,
180 endTime: job.endTime,
181 updateTime: job.updateTime,
182 };
183};
184
185const getByName = (name: string) =>
186 aiplatform
187 .getProjectsLocationsPipelineJobs({ name })
188 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
189
190const TERMINAL_STATES = new Set([
191 "PIPELINE_STATE_SUCCEEDED",
192 "PIPELINE_STATE_FAILED",
193 "PIPELINE_STATE_CANCELLED",
194]);
195
196class PipelineJobNotTerminal extends Data.TaggedError(
197 "GCP.AIPlatform.PipelineJobNotTerminal",
198)<{ name: string; state: string | undefined }> {}
199
200/**
201 * A pending or running job cannot be deleted: cancel it and wait until it
202 * settles (cancellation takes up to a few minutes).
203 */
204const cancelUntilTerminal = (name: string) =>
205 Effect.gen(function* () {
206 const job = yield* getByName(name);
207 if (job === undefined || TERMINAL_STATES.has(job.state ?? "")) return;
208 yield* aiplatform
209 .cancelProjectsLocationsPipelineJobs({ name, body: {} })
210 .pipe(
211 Effect.catchTag("NotFound", () => Effect.void),
212 // The job can reach a terminal state between the read above and the
213 // cancel ("… is in state JOB_STATE_FAILED and cannot be canceled");
214 // the wait below re-reads the state and still fails if it is not
215 // terminal.
216 Effect.catchTag("BadRequest", () => Effect.void),
217 );
218 yield* getByName(name).pipe(
219 Effect.flatMap((current) =>
220 current === undefined || TERMINAL_STATES.has(current.state ?? "")
221 ? Effect.void
222 : Effect.fail(
223 new PipelineJobNotTerminal({ name, state: current.state }),
224 ),
225 ),
226 Effect.retry({
227 while: (error) =>
228 error._tag === "GCP.AIPlatform.PipelineJobNotTerminal",
229 times: 36,
230 schedule: Schedule.spaced("5 seconds"),
231 }),
232 );
233 });
234
235const waitUntilGone = (name: string) =>
236 getByName(name).pipe(
237 Effect.filterOrFail(
238 (job) => job === undefined,
239 () => new AiPlatformStillExists({ name }),
240 ),
241 Effect.asVoid,
242 Effect.retry({
243 while: (error) => error._tag === "GCP.AIPlatform.StillExists",
244 times: 10,
245 schedule: Schedule.spaced("2 seconds"),
246 }),
247 );
248
249export const PipelineJobProvider = () =>
250 Provider.succeed(PipelineJob, {
251 stables: ["name", "pipelineJobId", "project", "location", "createTime"],
252
253 diff: Effect.fn(function* ({ news, olds, output }) {
254 const env = yield* GcpEnvironment.current;
255 if (!isResolved(news)) return undefined;
256 const previousId = olds?.pipelineJobId ?? output?.pipelineJobId;
257 const nextId = news.pipelineJobId ?? previousId;
258 const previousLocation = normalizeLocation(
259 olds?.location ?? output?.location,
260 env.region,
261 );
262 const nextLocation = normalizeLocation(
263 news.location ?? output?.location,
264 env.region,
265 );
266 const specChanged =
267 olds !== undefined &&
268 (!jsonEqual(news.pipelineSpec, olds.pipelineSpec) ||
269 (news.templateUri ?? "") !== (olds.templateUri ?? ""));
270 const replace =
271 (previousId !== undefined &&
272 nextId !== undefined &&
273 nextId !== previousId) ||
274 previousLocation !== nextLocation ||
275 specChanged;
276 if (!replace) return undefined;
277 return {
278 action: "replace" as const,
279 deleteFirst:
280 previousLocation === nextLocation &&
281 previousId !== undefined &&
282 nextId === previousId,
283 };
284 }),
285
286 read: Effect.fn(function* ({ id, olds, output }) {
287 const env = yield* GcpEnvironment.current;
288 const jobId = yield* toPhysicalId(
289 id,
290 olds?.pipelineJobId,
291 output?.pipelineJobId,
292 MAX_PIPELINE_JOB_ID_LENGTH,
293 );
294 const location = normalizeLocation(
295 olds?.location ?? output?.location,
296 env.region,
297 );
298 const name = output?.name ?? resourceName(env.project, location, jobId);
299 const existing = yield* getByName(name);
300 if (existing === undefined) return undefined;
301 const attrs = toAttrs(existing, env.project);
302 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
303 ? attrs
304 : Unowned(attrs);
305 }),
306
307 list: () =>
308 Effect.gen(function* () {
309 const env = yield* GcpEnvironment.current;
310 const pages = (yield* Effect.forEach(
311 listLocations(env.region),
312 (location) =>
313 collectPages(
314 aiplatform.listProjectsLocationsPipelineJobs.pages({
315 parent: locationParent(env.project, location),
316 pageSize: 100,
317 }),
318 ).pipe(Effect.catchTag("NotFound", () => Effect.succeed([]))),
319 )).flat();
320 return pages.flatMap((page) =>
321 (page.pipelineJobs ?? [])
322 .filter((job) =>
323 Object.keys(job.labels ?? {}).some((key) =>
324 key.startsWith("alchemy-"),
325 ),
326 )
327 .map((job) => toAttrs(job, env.project)),
328 );
329 }),
330
331 reconcile: Effect.fn(function* ({ id, news, output }) {
332 const env = yield* GcpEnvironment.current;
333 const jobId = yield* toPhysicalId(
334 id,
335 news.pipelineJobId,
336 output?.pipelineJobId,
337 MAX_PIPELINE_JOB_ID_LENGTH,
338 );
339 const location = normalizeLocation(
340 news.location ?? output?.location,
341 env.region,
342 );
343 const name = resourceName(env.project, location, jobId);
344 const desiredLabels = {
345 ...toLabels(news.labels),
346 ...(yield* createInternalLabels(id)),
347 };
348
349 let current = yield* getByName(output?.name ?? name);
350
351 if (current === undefined) {
352 const created = yield* aiplatform
353 .createProjectsLocationsPipelineJobs({
354 parent: locationParent(env.project, location),
355 pipelineJobId: jobId,
356 body: {
357 displayName: news.displayName ?? jobId,
358 labels: desiredLabels,
359 pipelineSpec: news.pipelineSpec,
360 templateUri: news.templateUri,
361 runtimeConfig: news.runtimeConfig,
362 serviceAccount: news.serviceAccount,
363 network: news.network,
364 reservedIpRanges: news.reservedIpRanges,
365 encryptionSpec: news.encryptionSpec,
366 preflightValidations: news.preflightValidations,
367 },
368 })
369 .pipe(Effect.catchTag("Conflict", () => getByName(name)));
370 current = created ?? undefined;
371 }
372
373 if (current === undefined) {
374 return yield* new PipelineJobNotResolved({ name });
375 }
376 return toAttrs(current, env.project);
377 }),
378
379 delete: Effect.fn(function* ({ output }) {
380 yield* cancelUntilTerminal(output.name);
381 const operation = yield* aiplatform
382 .deleteProjectsLocationsPipelineJobs({ name: output.name })
383 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
384 if (operation !== undefined) {
385 yield* waitForOperation(operation, { notFoundOk: true });
386 }
387 yield* waitUntilGone(output.name);
388 }),
389 });