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 { GcpEnvironment } from "../Environment.ts";
11import { createInternalLabels } from "../Labels.ts";
12import type { Providers } from "../Providers.ts";
13import { listLocations } from "./names.ts";
14import { resourceNameFromOperation, waitForOperation } from "./operations.ts";
15import {
16 encodeOwnershipLine,
17 hasOwnershipMarker,
18 lastSegment,
19 locationOf,
20 locationParent,
21 ownedByAlchemy,
22 parseOwnership,
23} from "./ownership.ts";
24
25export type SpecialistPoolProps = {
26 /**
27 * Region (`us-central1`, …). Immutable — changing it replaces the pool.
28 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
29 */
30 location?: string;
31 /**
32 * User-facing display name. Must be unique per project. Specialist pools
33 * have no labels field, so Alchemy ownership is stored in a `[alchemy …]`
34 * prefix.
35 */
36 displayName?: string;
37 /**
38 * Email addresses of specialist managers.
39 */
40 specialistManagerEmails?: string[];
41 /**
42 * Email addresses of specialist workers.
43 */
44 specialistWorkerEmails?: string[];
45};
46
47export type SpecialistPool = Resource<
48 "GCP.AIPlatform.SpecialistPool",
49 SpecialistPoolProps,
50 {
51 /** Full resource name `projects/{project}/locations/{location}/specialistPools/{specialist_pool}`. */
52 name: string;
53 /** Pool id (last path segment). */
54 specialistPoolId: string;
55 /** Region id. */
56 location: string;
57 /** Project id. */
58 project: string;
59 /** User display name with the Alchemy ownership prefix stripped. */
60 displayName: string | undefined;
61 /** Manager emails. */
62 specialistManagerEmails: string[];
63 /** Worker emails. */
64 specialistWorkerEmails: string[];
65 /** Number of managers. */
66 specialistManagersCount: number | undefined;
67 /** Pending data labeling jobs. */
68 pendingDataLabelingJobs: string[];
69 },
70 never,
71 Providers
72>;
73
74/**
75 * A Vertex AI SpecialistPool of managers and workers for data labeling.
76 *
77 * Specialist pools have no labels field — Alchemy stamps ownership into
78 * the display name. Location is immutable. Manager and worker emails
79 * update in place.
80 *
81 * ### Creating a Pool
82 * **Example:** Empty pool with a display name
83 * ```typescript
84 * const pool = yield* GCP.AIPlatform.SpecialistPool("Labelers", {
85 * displayName: "labelers",
86 * });
87 * ```
88 *
89 * @resource
90 * @category AIPlatform
91 */
92export const SpecialistPool = Resource<SpecialistPool>(
93 "GCP.AIPlatform.SpecialistPool",
94);
95
96export class SpecialistPoolNotResolved extends Data.TaggedError(
97 "GCP.AIPlatform.SpecialistPoolNotResolved",
98)<{
99 name: string;
100}> {}
101
102const sorted = (values: readonly string[] | undefined) =>
103 [...(values ?? [])].sort((left, right) => left.localeCompare(right));
104
105const sameEmails = (
106 left: readonly string[] | undefined,
107 right: readonly string[] | undefined,
108) => JSON.stringify(sorted(left)) === JSON.stringify(sorted(right));
109
110const toAttrs = (
111 pool: aiplatform.GoogleCloudAiplatformV1SpecialistPool,
112 project: string,
113) => {
114 const name = pool.name ?? "";
115 const parsed = parseOwnership(pool.displayName);
116 return {
117 name,
118 specialistPoolId: lastSegment(name),
119 location: locationOf(name),
120 project,
121 displayName: parsed.text,
122 specialistManagerEmails: pool.specialistManagerEmails ?? [],
123 specialistWorkerEmails: pool.specialistWorkerEmails ?? [],
124 specialistManagersCount: pool.specialistManagersCount,
125 pendingDataLabelingJobs: pool.pendingDataLabelingJobs ?? [],
126 };
127};
128
129const getByName = (name: string) =>
130 name.length === 0
131 ? Effect.succeed(undefined)
132 : aiplatform
133 .getProjectsLocationsSpecialistPools({ name })
134 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
135
136const waitUntilExists = (name: string) =>
137 getByName(name).pipe(
138 Effect.flatMap((pool) =>
139 pool
140 ? Effect.succeed(pool)
141 : Effect.fail(new SpecialistPoolNotResolved({ name })),
142 ),
143 Effect.retry({
144 while: (error) =>
145 error._tag === "GCP.AIPlatform.SpecialistPoolNotResolved",
146 times: 8,
147 schedule: Schedule.spaced("2 seconds"),
148 }),
149 );
150
151const listAt = (parent: string, project: string) =>
152 aiplatform.listProjectsLocationsSpecialistPools
153 .pages({ parent, pageSize: 100 })
154 .pipe(
155 Stream.flatMap((page) => Stream.fromIterable(page.specialistPools ?? [])),
156 Stream.filter((pool) => hasOwnershipMarker(pool.displayName)),
157 Stream.map((pool) => toAttrs(pool, project)),
158 Stream.runCollect,
159 Effect.map((chunk) => Array.from(chunk)),
160 Effect.catchTag("NotFound", () => Effect.succeed([])),
161 );
162
163const findByDisplayName = (parent: string, displayName: string) =>
164 aiplatform.listProjectsLocationsSpecialistPools
165 .pages({ parent, pageSize: 100 })
166 .pipe(
167 Stream.flatMap((page) => Stream.fromIterable(page.specialistPools ?? [])),
168 Stream.filter((pool) => pool.displayName === displayName),
169 Stream.runHead,
170 Effect.map((option) =>
171 option._tag === "Some" ? option.value : undefined,
172 ),
173 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
174 );
175
176export const SpecialistPoolProvider = () =>
177 Provider.succeed(SpecialistPool, {
178 stables: ["name", "specialistPoolId", "location", "project"],
179
180 diff: Effect.fn(function* ({ news, olds, output }) {
181 if (!isResolved(news)) return undefined;
182 const previousLocation = olds?.location ?? output?.location;
183 const nextLocation = news.location ?? previousLocation;
184 if (previousLocation !== undefined && previousLocation !== nextLocation) {
185 return { action: "replace" as const, deleteFirst: false };
186 }
187 return undefined;
188 }),
189
190 read: Effect.fn(function* ({ id, olds, output }) {
191 const env = yield* GcpEnvironment.current;
192 const existing = yield* getByName(output?.name ?? "");
193 if (existing !== undefined) {
194 const attrs = toAttrs(existing, env.project);
195 return (yield* ownedByAlchemy(id, existing.displayName))
196 ? attrs
197 : Unowned(attrs);
198 }
199 const location = olds?.location ?? output?.location ?? env.region;
200 const parent = locationParent(env.project, location);
201 const ownership = yield* createInternalLabels(id);
202 const displayName = encodeOwnershipLine(ownership, olds?.displayName);
203 const found = yield* findByDisplayName(parent, displayName);
204 if (found === undefined) return undefined;
205 const attrs = toAttrs(found, env.project);
206 return (yield* ownedByAlchemy(id, found.displayName))
207 ? attrs
208 : Unowned(attrs);
209 }),
210
211 list: () =>
212 Effect.gen(function* () {
213 const env = yield* GcpEnvironment.current;
214 return (yield* Effect.forEach(listLocations(env.region), (location) =>
215 listAt(locationParent(env.project, location), env.project),
216 )).flat();
217 }),
218
219 reconcile: Effect.fn(function* ({ id, news, output }) {
220 const env = yield* GcpEnvironment.current;
221 const location = news.location ?? output?.location ?? env.region;
222 const parent = locationParent(env.project, location);
223 const ownership = yield* createInternalLabels(id);
224 const displayName = encodeOwnershipLine(ownership, news.displayName);
225 const managers = news.specialistManagerEmails;
226 const workers = news.specialistWorkerEmails;
227
228 let current = yield* getByName(output?.name ?? "");
229 if (current === undefined) {
230 current = yield* findByDisplayName(parent, displayName);
231 }
232
233 if (current === undefined) {
234 const created = yield* aiplatform
235 .createProjectsLocationsSpecialistPools({
236 parent,
237 body: {
238 displayName,
239 specialistManagerEmails: managers,
240 specialistWorkerEmails: workers,
241 },
242 })
243 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
244 if (created !== undefined) {
245 const done = yield* waitForOperation(created, {
246 alreadyExistsOk: true,
247 });
248 const createdName = resourceNameFromOperation(done) ?? "";
249 if (createdName.length > 0) {
250 current = yield* waitUntilExists(createdName);
251 }
252 }
253 if (current === undefined) {
254 current = yield* findByDisplayName(parent, displayName);
255 }
256 }
257
258 if (current === undefined) {
259 return yield* new SpecialistPoolNotResolved({
260 name: output?.name ?? `${parent}/specialistPools/-`,
261 });
262 }
263
264 const name = current.name ?? "";
265 const displayChanged = (current.displayName ?? "") !== displayName;
266 const managersChanged = !sameEmails(
267 current.specialistManagerEmails,
268 managers,
269 );
270 const workersChanged = !sameEmails(
271 current.specialistWorkerEmails,
272 workers,
273 );
274
275 if (displayChanged || managersChanged || workersChanged) {
276 const patched = yield* aiplatform.patchProjectsLocationsSpecialistPools(
277 {
278 name,
279 updateMask: [
280 displayChanged ? "displayName" : undefined,
281 managersChanged ? "specialistManagerEmails" : undefined,
282 workersChanged ? "specialistWorkerEmails" : undefined,
283 ]
284 .filter((field): field is string => field !== undefined)
285 .join(","),
286 body: {
287 name,
288 displayName,
289 specialistManagerEmails: managers,
290 specialistWorkerEmails: workers,
291 },
292 },
293 );
294 yield* waitForOperation(patched);
295 current = yield* waitUntilExists(name);
296 }
297
298 return toAttrs(current, env.project);
299 }),
300
301 delete: Effect.fn(function* ({ output }) {
302 const existing = yield* getByName(output.name);
303 if (existing === undefined) return;
304 const operation = yield* aiplatform
305 .deleteProjectsLocationsSpecialistPools({
306 name: output.name,
307 force: true,
308 })
309 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
310 if (operation !== undefined) {
311 yield* waitForOperation(operation, { notFoundOk: true });
312 }
313 }),
314 });