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 waitForOperation,
25} from "./internal.ts";
26import type { EncryptionSpec } from "./shared.ts";
27
28export type CustomJobSpec = aiplatform.GoogleCloudAiplatformV1CustomJobSpec;
29
30export type CustomJobProps = {
31 /**
32 * Region. Immutable — changing it replaces the job.
33 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
34 */
35 location?: string;
36 /**
37 * User-facing display name (max 128 characters). Defaults to a generated
38 * id. Immutable after create — CustomJob has no update API.
39 */
40 displayName?: string;
41 /**
42 * Worker pools, container/python spec, networking, and scheduling.
43 * Required. Immutable after create.
44 */
45 jobSpec: CustomJobSpec;
46 /**
47 * User labels. Alchemy ownership labels are merged in automatically.
48 * Set at create; CustomJob cannot patch labels.
49 */
50 labels?: Record<string, string>;
51 /**
52 * Customer-managed encryption. Immutable after create.
53 */
54 encryptionSpec?: EncryptionSpec;
55};
56
57export type CustomJob = Resource<
58 "GCP.AIPlatform.CustomJob",
59 CustomJobProps,
60 {
61 /** Full resource name `.../customJobs/{custom_job}`. */
62 name: string;
63 /** Custom job id (last path segment). */
64 customJobId: string;
65 /** Project id. */
66 project: string;
67 /** Location id. */
68 location: string;
69 /** Display name. */
70 displayName: string | undefined;
71 /** Job spec as stored by Vertex. */
72 jobSpec: CustomJobSpec | undefined;
73 /** User labels (Alchemy ownership labels stripped). */
74 labels: Record<string, string>;
75 /** Detailed job state. */
76 state: string | undefined;
77 /** Customer-managed KMS key, if any. */
78 kmsKeyName: string | undefined;
79 /** RFC3339 creation timestamp. */
80 createTime: string | undefined;
81 /** RFC3339 start timestamp. */
82 startTime: string | undefined;
83 /** RFC3339 end timestamp. */
84 endTime: string | undefined;
85 /** RFC3339 last-update timestamp. */
86 updateTime: string | undefined;
87 /** Failure/cancel status, if any. */
88 error: { code?: number; message?: string } | undefined;
89 },
90 never,
91 Providers
92>;
93
94/**
95 * A Vertex AI CustomJob — a container or Python training workload.
96 *
97 * Creating a CustomJob starts it immediately. There is no update API, so
98 * reconcile is observe-ensure (create if missing). Delete cancels a
99 * running job, then deletes it. Alchemy ownership labels are merged into
100 * `labels` so `list` / nuke can find the job.
101 *
102 * ### Creating a Custom Job
103 * **Example:** Single-worker container job
104 * ```typescript
105 * const job = yield* GCP.AIPlatform.CustomJob("Train", {
106 * displayName: "echo-train",
107 * jobSpec: {
108 * workerPoolSpecs: [
109 * {
110 * machineSpec: { machineType: "n1-standard-4" },
111 * replicaCount: "1",
112 * containerSpec: {
113 * imageUri: "gcr.io/google-samples/hello-app:1.0",
114 * command: ["echo"],
115 * args: ["ok"],
116 * },
117 * },
118 * ],
119 * scheduling: { timeout: "600s" },
120 * },
121 * labels: { env: "test" },
122 * });
123 * ```
124 *
125 * @resource
126 * @category AIPlatform
127 */
128export const CustomJob = Resource<CustomJob>("GCP.AIPlatform.CustomJob");
129
130export class CustomJobNotResolved extends Data.TaggedError(
131 "GCP.AIPlatform.CustomJobNotResolved",
132)<{
133 name: string;
134}> {}
135
136const toId = (id: string, existing?: string) =>
137 Effect.gen(function* () {
138 if (existing !== undefined) return existing;
139 return rfc1035(
140 yield* createPhysicalName({
141 id,
142 maxLength: 63,
143 lowercase: true,
144 }),
145 );
146 });
147
148const toAttrs = (
149 job: aiplatform.GoogleCloudAiplatformV1CustomJob,
150 project: string,
151) => {
152 const name = job.name ?? "";
153 const parsed = parseResourceName(name, "customJobs");
154 return {
155 name,
156 customJobId: parsed.id,
157 project: parsed.project || project,
158 location: parsed.location,
159 displayName: job.displayName,
160 jobSpec: job.jobSpec,
161 labels: userLabels(job.labels),
162 state: job.state,
163 kmsKeyName: job.encryptionSpec?.kmsKeyName,
164 createTime: job.createTime,
165 startTime: job.startTime,
166 endTime: job.endTime,
167 updateTime: job.updateTime,
168 error: job.error
169 ? { code: job.error.code, message: job.error.message }
170 : undefined,
171 };
172};
173
174const getByName = (name: string) =>
175 aiplatform
176 .getProjectsLocationsCustomJobs({ name })
177 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
178
179const listJobs = (project: string, region: string) => {
180 const collect = (parent: string) =>
181 aiplatform.listProjectsLocationsCustomJobs
182 .pages({ parent, pageSize: 100 })
183 .pipe(
184 Stream.flatMap((page) => Stream.fromIterable(page.customJobs ?? [])),
185 Stream.filter((job) => hasAlchemyLabelKeys(job.labels)),
186 Stream.runCollect,
187 Effect.map((chunk) => Array.from(chunk)),
188 );
189 const fallback = Effect.forEach(listLocations(region), (location) =>
190 collect(`projects/${project}/locations/${location}`),
191 ).pipe(Effect.map((pages) => pages.flat()));
192 // Vertex AI has no `locations/-` wildcard; scan known locations.
193 return fallback.pipe(Effect.catchTag("NotFound", () => Effect.succeed([])));
194};
195
196const findOwned = (
197 id: string,
198 project: string,
199 region: string,
200 hinted?: string,
201) =>
202 Effect.gen(function* () {
203 if (hinted !== undefined && hinted.length > 0) {
204 const existing = yield* getByName(hinted);
205 if (existing !== undefined) return existing;
206 }
207 const jobs = yield* listJobs(project, region);
208 for (const job of jobs) {
209 if (yield* hasAlchemyLabels(id, tagRecord(job.labels))) return job;
210 }
211 return undefined as aiplatform.GoogleCloudAiplatformV1CustomJob | undefined;
212 });
213
214const cancelAndDelete = (name: string) =>
215 Effect.gen(function* () {
216 const existing = yield* getByName(name);
217 if (existing === undefined) return;
218 if (!isJobTerminal(existing.state)) {
219 yield* aiplatform
220 .cancelProjectsLocationsCustomJobs({ name })
221 .pipe(Effect.catchTag("NotFound", () => Effect.void));
222 yield* getByName(name).pipe(
223 Effect.filterOrFail(
224 (job) => job === undefined || isJobTerminal(job.state),
225 () => new CustomJobNotResolved({ name }),
226 ),
227 Effect.retry({
228 while: (error) =>
229 error._tag === "GCP.AIPlatform.CustomJobNotResolved",
230 times: 8,
231 schedule: Schedule.spaced("4 seconds"),
232 }),
233 Effect.catchTag(
234 "GCP.AIPlatform.CustomJobNotResolved",
235 () => Effect.void,
236 ),
237 );
238 }
239 const operation = yield* aiplatform
240 .deleteProjectsLocationsCustomJobs({ name })
241 .pipe(
242 Effect.retry({
243 while: (error) => error._tag === "Conflict",
244 times: 8,
245 schedule: Schedule.spaced("3 seconds"),
246 }),
247 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
248 Effect.catchTag("Conflict", () => Effect.succeed(undefined)),
249 );
250 if (operation !== undefined) {
251 yield* waitForOperation(operation, { notFoundOk: true });
252 }
253 });
254
255export const CustomJobProvider = () =>
256 Provider.succeed(CustomJob, {
257 stables: ["name", "customJobId", "project", "location", "createTime"],
258
259 diff: Effect.fn(function* ({ news, olds, output }) {
260 const env = yield* GcpEnvironment.current;
261 if (!isResolved(news)) return undefined;
262 const previousLocation = normalizeLocation(
263 olds?.location ?? output?.location,
264 env.region,
265 );
266 const nextLocation = normalizeLocation(
267 news.location ?? olds?.location ?? output?.location,
268 env.region,
269 );
270 if (previousLocation !== nextLocation) {
271 return { action: "replace" as const, deleteFirst: false };
272 }
273 return undefined;
274 }),
275
276 read: Effect.fn(function* ({ id, olds, output }) {
277 const env = yield* GcpEnvironment.current;
278 const existing = yield* findOwned(
279 id,
280 env.project,
281 env.region,
282 output?.name,
283 );
284 if (existing === undefined) return undefined;
285 const attrs = toAttrs(existing, env.project);
286 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
287 ? attrs
288 : Unowned(attrs);
289 }),
290
291 list: () =>
292 Effect.gen(function* () {
293 const env = yield* GcpEnvironment.current;
294 const jobs = yield* listJobs(env.project, env.region);
295 return jobs.map((job) => toAttrs(job, env.project));
296 }),
297
298 reconcile: Effect.fn(function* ({ id, news, output }) {
299 const env = yield* GcpEnvironment.current;
300 const location = normalizeLocation(
301 news.location ?? output?.location,
302 env.region,
303 );
304 const displayName =
305 news.displayName ?? (yield* toId(id, output?.customJobId));
306 const desiredLabels = {
307 ...toLabels(news.labels),
308 ...(yield* createInternalLabels(id)),
309 };
310
311 let current = yield* findOwned(id, env.project, env.region, output?.name);
312
313 if (current === undefined) {
314 const created = yield* aiplatform
315 .createProjectsLocationsCustomJobs({
316 parent: parentOf(env.project, location),
317 body: {
318 displayName,
319 jobSpec: news.jobSpec,
320 labels: desiredLabels,
321 encryptionSpec: news.encryptionSpec,
322 },
323 })
324 .pipe(
325 Effect.catchTag("Conflict", () =>
326 findOwned(id, env.project, env.region),
327 ),
328 );
329 current = created ?? undefined;
330 }
331
332 if (current === undefined) {
333 return yield* new CustomJobNotResolved({
334 name: output?.name ?? parentOf(env.project, location),
335 });
336 }
337
338 return toAttrs(current, env.project);
339 }),
340
341 delete: Effect.fn(function* ({ output }) {
342 yield* cancelAndDelete(output.name);
343 }),
344 });