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 { resourceNameFromOperation, waitForOperation } from "./operations.ts";
15import {
16 AiPlatformNotResolved,
17 AiPlatformStillExists,
18 collectPages,
19 locationParent,
20 normalizeLocation,
21 parseResourceName,
22 toPhysicalId,
23 userLabels,
24 type EncryptionSpec,
25 type MachineSpec,
26 type NetworkSpec,
27 type PersistentDiskSpec,
28} from "./shared.ts";
29
30const COLLECTION = "notebookExecutionJobs";
31
32export type DirectNotebookSource = {
33 /** Base64-encoded contents of the input notebook. */
34 content?: string;
35};
36
37export type GcsNotebookSource = {
38 /** Cloud Storage URI of the ipynb (`gs://bucket/notebook.ipynb`). */
39 uri?: string;
40 /** Object generation to read. Unset reads the live object. */
41 generation?: string;
42};
43
44export type DataformRepositorySource = {
45 /** Dataform repository resource name. */
46 dataformRepositoryResourceName?: string;
47 /** Commit SHA. Unset reads HEAD. */
48 commitSha?: string;
49};
50
51export type CustomEnvironmentSpec = {
52 /** Machine spec for the execution job. */
53 machineSpec?: MachineSpec;
54 /** Network spec. */
55 networkSpec?: NetworkSpec;
56 /** Persistent disk attached to the job. */
57 persistentDiskSpec?: PersistentDiskSpec;
58};
59
60export type NotebookExecutionJobProps = {
61 /**
62 * Job id. If omitted, a unique RFC1035 name is generated. Immutable —
63 * changing it replaces the job (a new execution is started).
64 */
65 notebookExecutionJobId?: string;
66 /**
67 * Vertex AI location. Immutable — changing it replaces the job.
68 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
69 */
70 location?: string;
71 /**
72 * Display name (max 128 UTF-8 characters). Defaults to the job id.
73 */
74 displayName?: string;
75 /**
76 * User labels. Alchemy ownership labels are merged in automatically.
77 */
78 labels?: Record<string, string>;
79 /**
80 * NotebookRuntimeTemplate to source compute from.
81 */
82 notebookRuntimeTemplateResourceName?: string;
83 /**
84 * Custom compute configuration. Mutually exclusive with a template.
85 */
86 customEnvironmentSpec?: CustomEnvironmentSpec;
87 /**
88 * Inline notebook contents. Mutually exclusive with GCS / Dataform.
89 */
90 directNotebookSource?: DirectNotebookSource;
91 /**
92 * Cloud Storage notebook URI.
93 */
94 gcsNotebookSource?: GcsNotebookSource;
95 /**
96 * Dataform repository notebook source.
97 */
98 dataformRepositorySource?: DataformRepositorySource;
99 /**
100 * Cloud Storage location for results (`gs://bucket-name`).
101 */
102 gcsOutputUri?: string;
103 /**
104 * Service account the execution runs as.
105 */
106 serviceAccount?: string;
107 /**
108 * User email to run as (Colab runtimes only).
109 */
110 executionUser?: string;
111 /**
112 * Kernel name. Unset uses the default kernel.
113 */
114 kernelName?: string;
115 /**
116 * Max running time (e.g. `"86400s"`).
117 */
118 executionTimeout?: string;
119 /**
120 * Customer-managed encryption key.
121 */
122 encryptionSpec?: EncryptionSpec;
123};
124
125export type NotebookExecutionJob = Resource<
126 "GCP.AIPlatform.NotebookExecutionJob",
127 NotebookExecutionJobProps,
128 {
129 /** Full resource name. */
130 name: string;
131 /** Job id (last path segment). */
132 notebookExecutionJobId: string;
133 /** Project id. */
134 project: string;
135 /** Location id. */
136 location: string;
137 /** Display name. */
138 displayName: string | undefined;
139 /** User labels (Alchemy ownership labels stripped). */
140 labels: Record<string, string>;
141 /** Job state (`JOB_STATE_RUNNING`, `JOB_STATE_SUCCEEDED`, …). */
142 jobState: string | undefined;
143 /** Cloud Storage output URI. */
144 gcsOutputUri: string | undefined;
145 /** NotebookRuntimeTemplate resource name, if used. */
146 notebookRuntimeTemplateResourceName: string | undefined;
147 /** Service account. */
148 serviceAccount: string | undefined;
149 /** RFC3339 creation timestamp. */
150 createTime: string | undefined;
151 /** RFC3339 last-update timestamp. */
152 updateTime: string | undefined;
153 },
154 never,
155 Providers
156>;
157
158/**
159 * A Vertex AI notebook execution job — runs an ipynb once against a
160 * runtime template or custom environment.
161 *
162 * There is no update API. Changing identity (`notebookExecutionJobId`,
163 * `location`) or the notebook source replaces the job. Other fields are
164 * create-time only.
165 *
166 * ### Creating a Job
167 * **Example:** Execute a GCS notebook
168 * ```typescript
169 * const job = yield* GCP.AIPlatform.NotebookExecutionJob("Nightly", {
170 * notebookRuntimeTemplateResourceName: template.name,
171 * gcsNotebookSource: { uri: "gs://bucket/notebook.ipynb" },
172 * gcsOutputUri: "gs://bucket/output",
173 * });
174 * ```
175 *
176 * @resource
177 * @category AIPlatform
178 */
179export const NotebookExecutionJob = Resource<NotebookExecutionJob>(
180 "GCP.AIPlatform.NotebookExecutionJob",
181);
182
183export class NotebookExecutionJobNotResolved extends Data.TaggedError(
184 "GCP.AIPlatform.NotebookExecutionJobNotResolved",
185)<{
186 name: string;
187}> {}
188
189const resourceName = (project: string, location: string, id: string) =>
190 `${locationParent(project, location)}/${COLLECTION}/${id}`;
191
192const toAttrs = (
193 job: aiplatform.GoogleCloudAiplatformV1NotebookExecutionJob,
194 project: string,
195) => {
196 const name = job.name ?? "";
197 const parsed = parseResourceName(name, COLLECTION);
198 return {
199 name,
200 notebookExecutionJobId: parsed.id,
201 project: parsed.project || project,
202 location: parsed.location,
203 displayName: job.displayName,
204 labels: userLabels(job.labels),
205 jobState: job.jobState,
206 gcsOutputUri: job.gcsOutputUri,
207 notebookRuntimeTemplateResourceName:
208 job.notebookRuntimeTemplateResourceName,
209 serviceAccount: job.serviceAccount,
210 createTime: job.createTime,
211 updateTime: job.updateTime,
212 };
213};
214
215const getByName = (name: string) =>
216 aiplatform
217 .getProjectsLocationsNotebookExecutionJobs({ name })
218 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
219
220const waitUntilExists = (name: string) =>
221 getByName(name).pipe(
222 Effect.filterOrFail(
223 (job): job is aiplatform.GoogleCloudAiplatformV1NotebookExecutionJob =>
224 job !== undefined,
225 () => new AiPlatformNotResolved({ name }),
226 ),
227 Effect.retry({
228 while: (error) => error._tag === "GCP.AIPlatform.NotResolved",
229 times: 8,
230 schedule: Schedule.spaced("2 seconds"),
231 }),
232 );
233
234const waitUntilGone = (name: string) =>
235 getByName(name).pipe(
236 Effect.filterOrFail(
237 (job) => job === undefined,
238 () => new AiPlatformStillExists({ name }),
239 ),
240 Effect.asVoid,
241 Effect.retry({
242 while: (error) => error._tag === "GCP.AIPlatform.StillExists",
243 times: 10,
244 schedule: Schedule.spaced("2 seconds"),
245 }),
246 );
247
248export const NotebookExecutionJobProvider = () =>
249 Provider.succeed(NotebookExecutionJob, {
250 stables: [
251 "name",
252 "notebookExecutionJobId",
253 "project",
254 "location",
255 "createTime",
256 ],
257
258 diff: Effect.fn(function* ({ news, olds, output }) {
259 const env = yield* GcpEnvironment.current;
260 if (!isResolved(news)) return undefined;
261 const previousId =
262 olds?.notebookExecutionJobId ?? output?.notebookExecutionJobId;
263 const nextId = news.notebookExecutionJobId ?? previousId;
264 const previousLocation = normalizeLocation(
265 olds?.location ?? output?.location,
266 env.region,
267 );
268 const nextLocation = normalizeLocation(
269 news.location ?? output?.location,
270 env.region,
271 );
272 const sourceChanged =
273 (news.gcsNotebookSource?.uri ?? "") !==
274 (olds?.gcsNotebookSource?.uri ?? "") ||
275 (news.directNotebookSource?.content ?? "") !==
276 (olds?.directNotebookSource?.content ?? "") ||
277 (news.dataformRepositorySource?.dataformRepositoryResourceName ??
278 "") !==
279 (olds?.dataformRepositorySource?.dataformRepositoryResourceName ??
280 "");
281 const replace =
282 (previousId !== undefined &&
283 nextId !== undefined &&
284 nextId !== previousId) ||
285 previousLocation !== nextLocation ||
286 (olds !== undefined && sourceChanged);
287 if (!replace) return undefined;
288 return {
289 action: "replace" as const,
290 deleteFirst:
291 previousLocation === nextLocation &&
292 previousId !== undefined &&
293 nextId === previousId,
294 };
295 }),
296
297 read: Effect.fn(function* ({ id, olds, output }) {
298 const env = yield* GcpEnvironment.current;
299 const jobId = yield* toPhysicalId(
300 id,
301 olds?.notebookExecutionJobId,
302 output?.notebookExecutionJobId,
303 );
304 const location = normalizeLocation(
305 olds?.location ?? output?.location,
306 env.region,
307 );
308 const name = output?.name ?? resourceName(env.project, location, jobId);
309 const existing = yield* getByName(name);
310 if (existing === undefined) return undefined;
311 const attrs = toAttrs(existing, env.project);
312 // Vertex AI replaces a job's labels with its runtime template's, so a
313 // job at the generated (instance-unique) id is ours even when the
314 // ownership labels name the template.
315 const generated = olds?.notebookExecutionJobId === undefined;
316 return generated ||
317 (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
318 ? attrs
319 : Unowned(attrs);
320 }),
321
322 list: () =>
323 Effect.gen(function* () {
324 const env = yield* GcpEnvironment.current;
325 const pages = (yield* Effect.forEach(
326 listLocations(env.region),
327 (location) =>
328 collectPages(
329 aiplatform.listProjectsLocationsNotebookExecutionJobs.pages({
330 parent: locationParent(env.project, location),
331 pageSize: 100,
332 }),
333 ).pipe(Effect.catchTag("NotFound", () => Effect.succeed([]))),
334 )).flat();
335 return pages.flatMap((page) =>
336 (page.notebookExecutionJobs ?? [])
337 .filter((job) =>
338 Object.keys(job.labels ?? {}).some((key) =>
339 key.startsWith("alchemy-"),
340 ),
341 )
342 .map((job) => toAttrs(job, env.project)),
343 );
344 }),
345
346 reconcile: Effect.fn(function* ({ id, news, output }) {
347 const env = yield* GcpEnvironment.current;
348 const jobId = yield* toPhysicalId(
349 id,
350 news.notebookExecutionJobId,
351 output?.notebookExecutionJobId,
352 );
353 const location = normalizeLocation(
354 news.location ?? output?.location,
355 env.region,
356 );
357 const name = resourceName(env.project, location, jobId);
358 const desiredLabels = {
359 ...toLabels(news.labels),
360 ...(yield* createInternalLabels(id)),
361 };
362
363 let current = yield* getByName(output?.name ?? name);
364
365 if (current === undefined) {
366 const created = yield* aiplatform
367 .createProjectsLocationsNotebookExecutionJobs({
368 parent: locationParent(env.project, location),
369 notebookExecutionJobId: jobId,
370 body: {
371 displayName: news.displayName ?? jobId,
372 labels: desiredLabels,
373 notebookRuntimeTemplateResourceName:
374 news.notebookRuntimeTemplateResourceName,
375 customEnvironmentSpec: news.customEnvironmentSpec,
376 directNotebookSource: news.directNotebookSource,
377 gcsNotebookSource: news.gcsNotebookSource,
378 dataformRepositorySource: news.dataformRepositorySource,
379 gcsOutputUri: news.gcsOutputUri,
380 serviceAccount: news.serviceAccount,
381 executionUser: news.executionUser,
382 kernelName: news.kernelName,
383 executionTimeout: news.executionTimeout,
384 encryptionSpec: news.encryptionSpec,
385 },
386 })
387 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
388 if (created !== undefined) {
389 yield* waitForOperation(created, { alreadyExistsOk: true });
390 }
391 const createdName =
392 resourceNameFromOperation(created ?? {}) ?? output?.name ?? name;
393 current = yield* waitUntilExists(createdName);
394 }
395
396 if (current === undefined) {
397 return yield* new NotebookExecutionJobNotResolved({ name });
398 }
399
400 return toAttrs(current, env.project);
401 }),
402
403 delete: Effect.fn(function* ({ output }) {
404 const operation = yield* aiplatform
405 .deleteProjectsLocationsNotebookExecutionJobs({ name: output.name })
406 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
407 if (operation !== undefined) {
408 yield* waitForOperation(operation, { notFoundOk: true });
409 }
410 yield* waitUntilGone(output.name);
411 }),
412 });