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 GcsSource = {
29 /** GCS URIs of input files. Wildcards allowed. */
30 uris?: string[];
31};
32
33export type GcsDestination = {
34 /** GCS URI prefix of the output directory. */
35 outputUriPrefix?: string;
36};
37
38export type BigQuerySource = {
39 /** BigQuery table URI, for example `bq://project.dataset.table`. */
40 inputUri?: string;
41};
42
43export type BigQueryDestination = {
44 /** BigQuery project, dataset, or table URI (`bq://...`). */
45 outputUri?: string;
46};
47
48export type BatchPredictionInputConfig = {
49 /** Instance format; must be one of the Model's supported input formats. */
50 instancesFormat?: string;
51 /** GCS input. */
52 gcsSource?: GcsSource;
53 /** BigQuery input. */
54 bigquerySource?: BigQuerySource;
55 /** Vertex multimodal dataset input. */
56 vertexMultimodalDatasetSource?: { datasetName?: string };
57};
58
59export type BatchPredictionOutputConfig = {
60 /** Predictions format; must be one of the Model's supported output formats. */
61 predictionsFormat?: string;
62 /** GCS output directory. */
63 gcsDestination?: GcsDestination;
64 /** BigQuery output. */
65 bigqueryDestination?: BigQueryDestination;
66};
67
68export type BatchDedicatedResources = {
69 /** Max machine replicas. Defaults to 10. */
70 maxReplicaCount?: number;
71 /** Starting replica count. */
72 startingReplicaCount?: number;
73 /** Single-machine spec. */
74 machineSpec?: {
75 machineType?: string;
76 acceleratorType?: string;
77 acceleratorCount?: number;
78 };
79};
80
81export type BatchPredictionJobProps = {
82 /**
83 * Region. Immutable — changing it replaces the job.
84 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
85 */
86 location?: string;
87 /**
88 * User-facing display name. Defaults to a generated id. Immutable after
89 * create — BatchPredictionJob has no update API.
90 */
91 displayName?: string;
92 /**
93 * Model resource that produces predictions
94 * (`projects/{project}/locations/{location}/models/{model}` or a
95 * publisher model). Exactly one of `model` or `endpoint` should be set.
96 */
97 model?: string;
98 /**
99 * Input instances. Required.
100 */
101 inputConfig: BatchPredictionInputConfig;
102 /**
103 * Output predictions location. Required.
104 */
105 outputConfig: BatchPredictionOutputConfig;
106 /**
107 * Dedicated machine resources for the prediction workers.
108 */
109 dedicatedResources?: BatchDedicatedResources;
110 /**
111 * Model parameters forwarded to the predictor.
112 */
113 modelParameters?: unknown;
114 /**
115 * Generate explanations with the predictions.
116 */
117 generateExplanation?: boolean;
118 /**
119 * Disable container logging.
120 */
121 disableContainerLogging?: boolean;
122 /**
123 * Runtime service account email.
124 */
125 serviceAccount?: string;
126 /**
127 * User labels. Alchemy ownership labels are merged in automatically.
128 */
129 labels?: Record<string, string>;
130 /**
131 * Customer-managed encryption.
132 */
133 encryptionSpec?: EncryptionSpec;
134};
135
136export type BatchPredictionJob = Resource<
137 "GCP.AIPlatform.BatchPredictionJob",
138 BatchPredictionJobProps,
139 {
140 /** Full resource name `.../batchPredictionJobs/{job}`. */
141 name: string;
142 /** Job id (last path segment). */
143 batchPredictionJobId: string;
144 /** Project id. */
145 project: string;
146 /** Location id. */
147 location: string;
148 /** Display name. */
149 displayName: string | undefined;
150 /** Model resource name. */
151 model: string | undefined;
152 /** Model version id. */
153 modelVersionId: string | undefined;
154 /** Detailed job state. */
155 state: string | undefined;
156 /** User labels (Alchemy ownership labels stripped). */
157 labels: Record<string, string>;
158 /** GCS output directory, if any. */
159 gcsOutputDirectory: string | undefined;
160 /** BigQuery output dataset, if any. */
161 bigqueryOutputDataset: string | undefined;
162 /** Customer-managed KMS key, if any. */
163 kmsKeyName: string | undefined;
164 /** RFC3339 creation timestamp. */
165 createTime: string | undefined;
166 /** RFC3339 start timestamp. */
167 startTime: string | undefined;
168 /** RFC3339 end timestamp. */
169 endTime: string | undefined;
170 /** Failure/cancel status, if any. */
171 error: { code?: number; message?: string } | undefined;
172 },
173 never,
174 Providers
175>;
176
177/**
178 * A Vertex AI BatchPredictionJob — offline predictions over many
179 * instances.
180 *
181 * Creating a job starts it immediately. There is no update API, so
182 * reconcile is observe-ensure. Delete can only run after the job
183 * finishes; Alchemy cancels first. Alchemy ownership labels are merged
184 * into `labels` so `list` / nuke can find the job.
185 *
186 * ### Creating a Batch Prediction Job
187 * **Example:** JSONL in, JSONL out
188 * ```typescript
189 * const job = yield* GCP.AIPlatform.BatchPredictionJob("Nightly", {
190 * displayName: "nightly-scores",
191 * model: "projects/my-project/locations/us-central1/models/my-model",
192 * inputConfig: {
193 * instancesFormat: "jsonl",
194 * gcsSource: { uris: ["gs://bucket/input.jsonl"] },
195 * },
196 * outputConfig: {
197 * predictionsFormat: "jsonl",
198 * gcsDestination: { outputUriPrefix: "gs://bucket/out/" },
199 * },
200 * });
201 * ```
202 *
203 * @resource
204 * @category AIPlatform
205 */
206export const BatchPredictionJob = Resource<BatchPredictionJob>(
207 "GCP.AIPlatform.BatchPredictionJob",
208);
209
210export class BatchPredictionJobNotResolved extends Data.TaggedError(
211 "GCP.AIPlatform.BatchPredictionJobNotResolved",
212)<{
213 name: string;
214}> {}
215
216const toId = (id: string, existing?: string) =>
217 Effect.gen(function* () {
218 if (existing !== undefined) return existing;
219 return rfc1035(
220 yield* createPhysicalName({
221 id,
222 maxLength: 63,
223 lowercase: true,
224 }),
225 );
226 });
227
228const toAttrs = (
229 job: aiplatform.GoogleCloudAiplatformV1BatchPredictionJob,
230 project: string,
231) => {
232 const name = job.name ?? "";
233 const parsed = parseResourceName(name, "batchPredictionJobs");
234 return {
235 name,
236 batchPredictionJobId: parsed.id,
237 project: parsed.project || project,
238 location: parsed.location,
239 displayName: job.displayName,
240 model: job.model,
241 modelVersionId: job.modelVersionId,
242 state: job.state,
243 labels: userLabels(job.labels),
244 gcsOutputDirectory: job.outputInfo?.gcsOutputDirectory,
245 bigqueryOutputDataset: job.outputInfo?.bigqueryOutputDataset,
246 kmsKeyName: job.encryptionSpec?.kmsKeyName,
247 createTime: job.createTime,
248 startTime: job.startTime,
249 endTime: job.endTime,
250 error: job.error
251 ? { code: job.error.code, message: job.error.message }
252 : undefined,
253 };
254};
255
256const getByName = (name: string) =>
257 aiplatform
258 .getProjectsLocationsBatchPredictionJobs({ name })
259 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
260
261const listJobs = (project: string, region: string) => {
262 const collect = (parent: string) =>
263 aiplatform.listProjectsLocationsBatchPredictionJobs
264 .pages({ parent, pageSize: 100 })
265 .pipe(
266 Stream.flatMap((page) =>
267 Stream.fromIterable(page.batchPredictionJobs ?? []),
268 ),
269 Stream.filter((job) => hasAlchemyLabelKeys(job.labels)),
270 Stream.runCollect,
271 Effect.map((chunk) => Array.from(chunk)),
272 );
273 const fallback = Effect.forEach(listLocations(region), (location) =>
274 collect(`projects/${project}/locations/${location}`),
275 ).pipe(Effect.map((pages) => pages.flat()));
276 // Vertex AI has no `locations/-` wildcard; scan known locations.
277 return fallback.pipe(Effect.catchTag("NotFound", () => Effect.succeed([])));
278};
279
280const findOwned = (
281 id: string,
282 project: string,
283 region: string,
284 hinted?: string,
285) =>
286 Effect.gen(function* () {
287 if (hinted !== undefined && hinted.length > 0) {
288 const existing = yield* getByName(hinted);
289 if (existing !== undefined) return existing;
290 }
291 const jobs = yield* listJobs(project, region);
292 for (const job of jobs) {
293 if (yield* hasAlchemyLabels(id, tagRecord(job.labels))) return job;
294 }
295 return undefined as
296 | aiplatform.GoogleCloudAiplatformV1BatchPredictionJob
297 | undefined;
298 });
299
300const cancelAndDelete = (name: string) =>
301 Effect.gen(function* () {
302 const existing = yield* getByName(name);
303 if (existing === undefined) return;
304 if (!isJobTerminal(existing.state)) {
305 yield* aiplatform
306 .cancelProjectsLocationsBatchPredictionJobs({ name })
307 .pipe(Effect.catchTag("NotFound", () => Effect.void));
308 yield* getByName(name).pipe(
309 Effect.filterOrFail(
310 (job) => job === undefined || isJobTerminal(job.state),
311 () => new BatchPredictionJobNotResolved({ name }),
312 ),
313 Effect.retry({
314 while: (error) =>
315 error._tag === "GCP.AIPlatform.BatchPredictionJobNotResolved",
316 times: 8,
317 schedule: Schedule.spaced("4 seconds"),
318 }),
319 Effect.catchTag(
320 "GCP.AIPlatform.BatchPredictionJobNotResolved",
321 () => Effect.void,
322 ),
323 );
324 }
325 const operation = yield* aiplatform
326 .deleteProjectsLocationsBatchPredictionJobs({ name })
327 .pipe(
328 Effect.retry({
329 while: (error) => error._tag === "Conflict",
330 times: 8,
331 schedule: Schedule.spaced("3 seconds"),
332 }),
333 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
334 Effect.catchTag("Conflict", () => Effect.succeed(undefined)),
335 );
336 if (operation !== undefined) {
337 yield* waitForOperation(operation, { notFoundOk: true });
338 }
339 });
340
341export const BatchPredictionJobProvider = () =>
342 Provider.succeed(BatchPredictionJob, {
343 stables: [
344 "name",
345 "batchPredictionJobId",
346 "project",
347 "location",
348 "createTime",
349 ],
350
351 diff: Effect.fn(function* ({ news, olds, output }) {
352 const env = yield* GcpEnvironment.current;
353 if (!isResolved(news)) return undefined;
354 const previousLocation = normalizeLocation(
355 olds?.location ?? output?.location,
356 env.region,
357 );
358 const nextLocation = normalizeLocation(
359 news.location ?? olds?.location ?? output?.location,
360 env.region,
361 );
362 if (previousLocation !== nextLocation) {
363 return { action: "replace" as const, deleteFirst: false };
364 }
365 return undefined;
366 }),
367
368 read: Effect.fn(function* ({ id, olds, output }) {
369 const env = yield* GcpEnvironment.current;
370 const existing = yield* findOwned(
371 id,
372 env.project,
373 env.region,
374 output?.name,
375 );
376 if (existing === undefined) return undefined;
377 const attrs = toAttrs(existing, env.project);
378 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
379 ? attrs
380 : Unowned(attrs);
381 }),
382
383 list: () =>
384 Effect.gen(function* () {
385 const env = yield* GcpEnvironment.current;
386 const jobs = yield* listJobs(env.project, env.region);
387 return jobs.map((job) => toAttrs(job, env.project));
388 }),
389
390 reconcile: Effect.fn(function* ({ id, news, output }) {
391 const env = yield* GcpEnvironment.current;
392 const location = normalizeLocation(
393 news.location ?? output?.location,
394 env.region,
395 );
396 const displayName =
397 news.displayName ?? (yield* toId(id, output?.batchPredictionJobId));
398 const desiredLabels = {
399 ...toLabels(news.labels),
400 ...(yield* createInternalLabels(id)),
401 };
402
403 let current = yield* findOwned(id, env.project, env.region, output?.name);
404
405 if (current === undefined) {
406 const created = yield* aiplatform
407 .createProjectsLocationsBatchPredictionJobs({
408 parent: parentOf(env.project, location),
409 body: {
410 displayName,
411 model: news.model,
412 inputConfig: news.inputConfig,
413 outputConfig: news.outputConfig,
414 dedicatedResources: news.dedicatedResources,
415 modelParameters: news.modelParameters,
416 generateExplanation: news.generateExplanation,
417 disableContainerLogging: news.disableContainerLogging,
418 serviceAccount: news.serviceAccount,
419 labels: desiredLabels,
420 encryptionSpec: news.encryptionSpec,
421 },
422 })
423 .pipe(
424 Effect.catchTag("Conflict", () =>
425 findOwned(id, env.project, env.region),
426 ),
427 );
428 current = created ?? undefined;
429 }
430
431 if (current === undefined) {
432 return yield* new BatchPredictionJobNotResolved({
433 name: output?.name ?? parentOf(env.project, location),
434 });
435 }
436
437 return toAttrs(current, env.project);
438 }),
439
440 delete: Effect.fn(function* ({ output }) {
441 yield* cancelAndDelete(output.name);
442 }),
443 });