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 ActiveLearningConfig =
29 aiplatform.GoogleCloudAiplatformV1ActiveLearningConfig;
30
31export type DataLabelingJobProps = {
32 /**
33 * Region. Immutable — changing it replaces the job.
34 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
35 */
36 location?: string;
37 /**
38 * User-facing display name. Defaults to a generated id. Immutable after
39 * create — DataLabelingJob has no update API.
40 */
41 displayName?: string;
42 /**
43 * Dataset resource names to label. Currently a single dataset.
44 * Format `projects/{project}/locations/{location}/datasets/{dataset}`.
45 */
46 datasets: string[];
47 /**
48 * Number of labelers per DataItem.
49 */
50 labelerCount: number;
51 /**
52 * GCS URI of the instruction PDF shared with labelers.
53 */
54 instructionUri: string;
55 /**
56 * GCS URI of the YAML schema describing the job config. Schema files
57 * live under
58 * `https://storage.googleapis.com/google-cloud-aiplatform/schema/datalabelingjob/inputs/`.
59 */
60 inputsSchemaUri: string;
61 /**
62 * Input config parameters for the job (schema-specific).
63 */
64 inputs: unknown;
65 /**
66 * Specialist pool resource names associated with this job.
67 */
68 specialistPools?: string[];
69 /**
70 * Labels applied to annotations created by this job.
71 */
72 annotationLabels?: Record<string, string>;
73 /**
74 * Active learning pipeline configuration.
75 */
76 activeLearningConfig?: ActiveLearningConfig;
77 /**
78 * User labels. Alchemy ownership labels are merged in automatically.
79 */
80 labels?: Record<string, string>;
81 /**
82 * Customer-managed encryption.
83 */
84 encryptionSpec?: EncryptionSpec;
85};
86
87export type DataLabelingJob = Resource<
88 "GCP.AIPlatform.DataLabelingJob",
89 DataLabelingJobProps,
90 {
91 /** Full resource name `.../dataLabelingJobs/{job}`. */
92 name: string;
93 /** Job id (last path segment). */
94 dataLabelingJobId: string;
95 /** Project id. */
96 project: string;
97 /** Location id. */
98 location: string;
99 /** Display name. */
100 displayName: string | undefined;
101 /** Dataset resource names. */
102 datasets: string[];
103 /** Labelers per DataItem. */
104 labelerCount: number | undefined;
105 /** Instruction PDF URI. */
106 instructionUri: string | undefined;
107 /** Inputs schema URI. */
108 inputsSchemaUri: string | undefined;
109 /** Detailed job state. */
110 state: string | undefined;
111 /** Labeling progress in `[0, 100]`. */
112 labelingProgress: number | undefined;
113 /** User labels (Alchemy ownership labels stripped). */
114 labels: Record<string, string>;
115 /** Customer-managed KMS key, if any. */
116 kmsKeyName: string | undefined;
117 /** RFC3339 creation timestamp. */
118 createTime: string | undefined;
119 /** RFC3339 last-update timestamp. */
120 updateTime: string | undefined;
121 /** Failure/cancel status, if any. */
122 error: { code?: number; message?: string } | undefined;
123 },
124 never,
125 Providers
126>;
127
128/**
129 * A Vertex AI DataLabelingJob — human labeling of Dataset DataItems.
130 *
131 * Creating a job starts labeling immediately. There is no update API, so
132 * reconcile is observe-ensure. Delete cancels a running job, then deletes
133 * it. Alchemy ownership labels are merged into `labels` so `list` / nuke
134 * can find the job.
135 *
136 * ### Creating a Data Labeling Job
137 * **Example:** Single-label image classification
138 * ```typescript
139 * const job = yield* GCP.AIPlatform.DataLabelingJob("Label", {
140 * datasets: [dataset.name],
141 * displayName: "label-images",
142 * labelerCount: 1,
143 * instructionUri: "gs://bucket/instructions.pdf",
144 * inputsSchemaUri:
145 * "gs://google-cloud-aiplatform/schema/datalabelingjob/inputs/image_classification_1.0.0.yaml",
146 * inputs: { annotationSpecs: ["cat", "dog"] },
147 * });
148 * ```
149 *
150 * @resource
151 * @category AIPlatform
152 */
153export const DataLabelingJob = Resource<DataLabelingJob>(
154 "GCP.AIPlatform.DataLabelingJob",
155);
156
157export class DataLabelingJobNotResolved extends Data.TaggedError(
158 "GCP.AIPlatform.DataLabelingJobNotResolved",
159)<{
160 name: string;
161}> {}
162
163const toId = (id: string, existing?: string) =>
164 Effect.gen(function* () {
165 if (existing !== undefined) return existing;
166 return rfc1035(
167 yield* createPhysicalName({
168 id,
169 maxLength: 63,
170 lowercase: true,
171 }),
172 );
173 });
174
175const toAttrs = (
176 job: aiplatform.GoogleCloudAiplatformV1DataLabelingJob,
177 project: string,
178) => {
179 const name = job.name ?? "";
180 const parsed = parseResourceName(name, "dataLabelingJobs");
181 return {
182 name,
183 dataLabelingJobId: parsed.id,
184 project: parsed.project || project,
185 location: parsed.location,
186 displayName: job.displayName,
187 datasets: job.datasets ?? [],
188 labelerCount: job.labelerCount,
189 instructionUri: job.instructionUri,
190 inputsSchemaUri: job.inputsSchemaUri,
191 state: job.state,
192 labelingProgress: job.labelingProgress,
193 labels: userLabels(job.labels),
194 kmsKeyName: job.encryptionSpec?.kmsKeyName,
195 createTime: job.createTime,
196 updateTime: job.updateTime,
197 error: job.error
198 ? { code: job.error.code, message: job.error.message }
199 : undefined,
200 };
201};
202
203const getByName = (name: string) =>
204 aiplatform
205 .getProjectsLocationsDataLabelingJobs({ name })
206 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
207
208const listJobs = (project: string, region: string) => {
209 const collect = (parent: string) =>
210 aiplatform.listProjectsLocationsDataLabelingJobs
211 .pages({ parent, pageSize: 100 })
212 .pipe(
213 Stream.flatMap((page) =>
214 Stream.fromIterable(page.dataLabelingJobs ?? []),
215 ),
216 Stream.filter((job) => hasAlchemyLabelKeys(job.labels)),
217 Stream.runCollect,
218 Effect.map((chunk) => Array.from(chunk)),
219 );
220 const fallback = Effect.forEach(listLocations(region), (location) =>
221 collect(`projects/${project}/locations/${location}`),
222 ).pipe(Effect.map((pages) => pages.flat()));
223 // Vertex AI has no `locations/-` wildcard; scan known locations.
224 return fallback.pipe(Effect.catchTag("NotFound", () => Effect.succeed([])));
225};
226
227const findOwned = (
228 id: string,
229 project: string,
230 region: string,
231 hinted?: string,
232) =>
233 Effect.gen(function* () {
234 if (hinted !== undefined && hinted.length > 0) {
235 const existing = yield* getByName(hinted);
236 if (existing !== undefined) return existing;
237 }
238 const jobs = yield* listJobs(project, region);
239 for (const job of jobs) {
240 if (yield* hasAlchemyLabels(id, tagRecord(job.labels))) return job;
241 }
242 return undefined as
243 | aiplatform.GoogleCloudAiplatformV1DataLabelingJob
244 | undefined;
245 });
246
247const cancelAndDelete = (name: string) =>
248 Effect.gen(function* () {
249 const existing = yield* getByName(name);
250 if (existing === undefined) return;
251 if (!isJobTerminal(existing.state)) {
252 yield* aiplatform
253 .cancelProjectsLocationsDataLabelingJobs({ name })
254 .pipe(Effect.catchTag("NotFound", () => Effect.void));
255 yield* getByName(name).pipe(
256 Effect.filterOrFail(
257 (job) => job === undefined || isJobTerminal(job.state),
258 () => new DataLabelingJobNotResolved({ name }),
259 ),
260 Effect.retry({
261 while: (error) =>
262 error._tag === "GCP.AIPlatform.DataLabelingJobNotResolved",
263 times: 8,
264 schedule: Schedule.spaced("4 seconds"),
265 }),
266 Effect.catchTag(
267 "GCP.AIPlatform.DataLabelingJobNotResolved",
268 () => Effect.void,
269 ),
270 );
271 }
272 const operation = yield* aiplatform
273 .deleteProjectsLocationsDataLabelingJobs({ name })
274 .pipe(
275 Effect.retry({
276 while: (error) => error._tag === "Conflict",
277 times: 8,
278 schedule: Schedule.spaced("3 seconds"),
279 }),
280 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
281 Effect.catchTag("Conflict", () => Effect.succeed(undefined)),
282 );
283 if (operation !== undefined) {
284 yield* waitForOperation(operation, { notFoundOk: true });
285 }
286 });
287
288export const DataLabelingJobProvider = () =>
289 Provider.succeed(DataLabelingJob, {
290 stables: ["name", "dataLabelingJobId", "project", "location", "createTime"],
291
292 diff: Effect.fn(function* ({ news, olds, output }) {
293 const env = yield* GcpEnvironment.current;
294 if (!isResolved(news)) return undefined;
295 const previousLocation = normalizeLocation(
296 olds?.location ?? output?.location,
297 env.region,
298 );
299 const nextLocation = normalizeLocation(
300 news.location ?? olds?.location ?? output?.location,
301 env.region,
302 );
303 if (previousLocation !== nextLocation) {
304 return { action: "replace" as const, deleteFirst: false };
305 }
306 return undefined;
307 }),
308
309 read: Effect.fn(function* ({ id, olds, output }) {
310 const env = yield* GcpEnvironment.current;
311 const existing = yield* findOwned(
312 id,
313 env.project,
314 env.region,
315 output?.name,
316 );
317 if (existing === undefined) return undefined;
318 const attrs = toAttrs(existing, env.project);
319 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
320 ? attrs
321 : Unowned(attrs);
322 }),
323
324 list: () =>
325 Effect.gen(function* () {
326 const env = yield* GcpEnvironment.current;
327 const jobs = yield* listJobs(env.project, env.region);
328 return jobs.map((job) => toAttrs(job, env.project));
329 }),
330
331 reconcile: Effect.fn(function* ({ id, news, output }) {
332 const env = yield* GcpEnvironment.current;
333 const location = normalizeLocation(
334 news.location ?? output?.location,
335 env.region,
336 );
337 const displayName =
338 news.displayName ?? (yield* toId(id, output?.dataLabelingJobId));
339 const desiredLabels = {
340 ...toLabels(news.labels),
341 ...(yield* createInternalLabels(id)),
342 };
343
344 let current = yield* findOwned(id, env.project, env.region, output?.name);
345
346 if (current === undefined) {
347 const created = yield* aiplatform
348 .createProjectsLocationsDataLabelingJobs({
349 parent: parentOf(env.project, location),
350 body: {
351 displayName,
352 datasets: news.datasets,
353 labelerCount: news.labelerCount,
354 instructionUri: news.instructionUri,
355 inputsSchemaUri: news.inputsSchemaUri,
356 inputs: news.inputs,
357 specialistPools: news.specialistPools,
358 annotationLabels: news.annotationLabels,
359 activeLearningConfig: news.activeLearningConfig,
360 labels: desiredLabels,
361 encryptionSpec: news.encryptionSpec,
362 },
363 })
364 .pipe(
365 Effect.catchTag("Conflict", () =>
366 findOwned(id, env.project, env.region),
367 ),
368 );
369 current = created ?? undefined;
370 }
371
372 if (current === undefined) {
373 return yield* new DataLabelingJobNotResolved({
374 name: output?.name ?? parentOf(env.project, location),
375 });
376 }
377
378 return toAttrs(current, env.project);
379 }),
380
381 delete: Effect.fn(function* ({ output }) {
382 yield* cancelAndDelete(output.name);
383 }),
384 });