1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Stream from "effect/Stream";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { GcpEnvironment } from "../Environment.ts";
10import { createInternalLabels } from "../Labels.ts";
11import type { Providers } from "../Providers.ts";
12import { listLocations } from "./names.ts";
13import { waitForOperation } from "./operations.ts";
14import {
15 encodeOwnershipLine,
16 hasOwnershipMarker,
17 lastSegment,
18 locationOf,
19 locationParent,
20 ownedByAlchemy,
21 parseOwnership,
22} from "./ownership.ts";
23
24export type PipelineJobRequest = {
25 /**
26 * Location parent for created pipeline jobs
27 * (`projects/{project}/locations/{location}`). Defaults to the
28 * Schedule's location.
29 */
30 parent?: string;
31 /** Optional pipeline job id. */
32 pipelineJobId?: string;
33 /** Pipeline job to create on each run. */
34 pipelineJob?: {
35 displayName?: string;
36 templateUri?: string;
37 pipelineSpec?: Record<string, unknown>;
38 serviceAccount?: string;
39 network?: string;
40 labels?: Record<string, string>;
41 runtimeConfig?: {
42 gcsOutputDirectory?: string;
43 parameterValues?: Record<string, unknown>;
44 };
45 };
46};
47
48export type NotebookExecutionJobRequest = {
49 parent?: string;
50 notebookExecutionJobId?: string;
51 notebookExecutionJob?: Record<string, unknown>;
52};
53
54export type ScheduleProps = {
55 /**
56 * Region (`us-central1`, …). Immutable — changing it replaces the
57 * Schedule.
58 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
59 */
60 location?: string;
61 /**
62 * User-facing display name. Vertex AI Schedules have no labels field,
63 * so Alchemy ownership is stored in a `[alchemy …]` prefix.
64 */
65 displayName?: string;
66 /**
67 * Cron schedule. Prefix with `CRON_TZ=${IANA_TIME_ZONE}` to set a
68 * timezone (for example `CRON_TZ=UTC 0 0 * * *`).
69 */
70 cron: string;
71 /**
72 * Maximum number of runs that can be started concurrently.
73 * @default "1"
74 */
75 maxConcurrentRunCount?: string;
76 /**
77 * Maximum number of active (non-terminal) pipeline runs. Only applies
78 * to `createPipelineJobRequest`.
79 */
80 maxConcurrentActiveRunCount?: string;
81 /**
82 * Maximum number of runs before the schedule completes.
83 */
84 maxRunCount?: string;
85 /**
86 * RFC3339 timestamp after which the first run may be scheduled.
87 */
88 startTime?: string;
89 /**
90 * RFC3339 timestamp after which no new runs are scheduled.
91 */
92 endTime?: string;
93 /**
94 * Queue new runs when `maxConcurrentRunCount` is reached instead of
95 * skipping them.
96 * @default false
97 */
98 allowQueueing?: boolean;
99 /**
100 * Pipeline job created on each tick. Exactly one of
101 * `createPipelineJobRequest` or `createNotebookExecutionJobRequest`
102 * must be set.
103 */
104 createPipelineJobRequest?: PipelineJobRequest;
105 /**
106 * Notebook execution job created on each tick.
107 */
108 createNotebookExecutionJobRequest?: NotebookExecutionJobRequest;
109 /**
110 * When true, the schedule exists but does not start new runs (`PAUSED`).
111 * Synced via `schedules.pause` / `schedules.resume`.
112 * @default false
113 */
114 paused?: boolean;
115};
116
117export type Schedule = Resource<
118 "GCP.AIPlatform.Schedule",
119 ScheduleProps,
120 {
121 /** Full resource name `projects/{project}/locations/{location}/schedules/{schedule}`. */
122 name: string;
123 /** Schedule id (last path segment). */
124 scheduleId: string;
125 /** Region id. */
126 location: string;
127 /** Project id. */
128 project: string;
129 /** User display name with the Alchemy ownership prefix stripped. */
130 displayName: string | undefined;
131 /** Cron expression. */
132 cron: string | undefined;
133 /** Server-reported state. */
134 state: string | undefined;
135 /** Whether the schedule is paused. */
136 paused: boolean;
137 /** Maximum concurrent starts. */
138 maxConcurrentRunCount: string | undefined;
139 /** Maximum concurrent active pipeline runs. */
140 maxConcurrentActiveRunCount: string | undefined;
141 /** Maximum run count. */
142 maxRunCount: string | undefined;
143 /** RFC3339 start bound. */
144 startTime: string | undefined;
145 /** RFC3339 end bound. */
146 endTime: string | undefined;
147 /** Whether new runs may queue. */
148 allowQueueing: boolean;
149 /** Number of runs started so far. */
150 startedRunCount: string | undefined;
151 /** Next scheduled run (RFC3339). */
152 nextRunTime: string | undefined;
153 /** RFC3339 creation timestamp. */
154 createTime: string | undefined;
155 /** RFC3339 last-update timestamp. */
156 updateTime: string | undefined;
157 },
158 never,
159 Providers
160>;
161
162/**
163 * A Vertex AI Schedule that periodically starts pipeline or notebook jobs.
164 *
165 * Schedules have no labels field — Alchemy stamps ownership into the
166 * display name. Location is immutable. Cron, run limits, and pause state
167 * update in place.
168 *
169 * ### Creating a Schedule
170 * **Example:** Daily paused pipeline schedule
171 * ```typescript
172 * const schedule = yield* GCP.AIPlatform.Schedule("Nightly", {
173 * cron: "CRON_TZ=UTC 0 8 * * *",
174 * paused: true,
175 * maxRunCount: "1",
176 * createPipelineJobRequest: {
177 * pipelineJob: {
178 * displayName: "nightly",
179 * templateUri: "https://us-kfp.pkg.dev/ml-pipeline/google-cloud-registry/hello-world/latest",
180 * },
181 * },
182 * });
183 * ```
184 *
185 * @resource
186 * @category AIPlatform
187 */
188export const Schedule = Resource<Schedule>("GCP.AIPlatform.Schedule");
189
190export class ScheduleNotResolved extends Data.TaggedError(
191 "GCP.AIPlatform.ScheduleNotResolved",
192)<{
193 name: string;
194}> {}
195
196const toAttrs = (
197 schedule: aiplatform.GoogleCloudAiplatformV1Schedule,
198 project: string,
199) => {
200 const name = schedule.name ?? "";
201 const parsed = parseOwnership(schedule.displayName);
202 return {
203 name,
204 scheduleId: lastSegment(name),
205 location: locationOf(name),
206 project,
207 displayName: parsed.text,
208 cron: schedule.cron,
209 state: schedule.state,
210 paused: schedule.state === "PAUSED",
211 maxConcurrentRunCount: schedule.maxConcurrentRunCount,
212 maxConcurrentActiveRunCount: schedule.maxConcurrentActiveRunCount,
213 maxRunCount: schedule.maxRunCount,
214 startTime: schedule.startTime,
215 endTime: schedule.endTime,
216 allowQueueing: schedule.allowQueueing === true,
217 startedRunCount: schedule.startedRunCount,
218 nextRunTime: schedule.nextRunTime,
219 createTime: schedule.createTime,
220 updateTime: schedule.updateTime,
221 };
222};
223
224const getByName = (name: string) =>
225 name.length === 0
226 ? Effect.succeed(undefined)
227 : aiplatform
228 .getProjectsLocationsSchedules({ name })
229 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
230
231const listAt = (parent: string, project: string) =>
232 aiplatform.listProjectsLocationsSchedules
233 .pages({ parent, pageSize: 100 })
234 .pipe(
235 Stream.flatMap((page) => Stream.fromIterable(page.schedules ?? [])),
236 Stream.filter((schedule) => hasOwnershipMarker(schedule.displayName)),
237 Stream.map((schedule) => toAttrs(schedule, project)),
238 Stream.runCollect,
239 Effect.map((chunk) => Array.from(chunk)),
240 Effect.catchTag("NotFound", () => Effect.succeed([])),
241 );
242
243const toPipelineRequest = (
244 request: PipelineJobRequest | undefined,
245 parent: string,
246): aiplatform.GoogleCloudAiplatformV1CreatePipelineJobRequest | undefined => {
247 if (request === undefined) return undefined;
248 const job = request.pipelineJob;
249 return {
250 parent: request.parent ?? parent,
251 pipelineJobId: request.pipelineJobId,
252 pipelineJob: job
253 ? {
254 displayName: job.displayName,
255 templateUri: job.templateUri,
256 pipelineSpec: job.pipelineSpec,
257 serviceAccount: job.serviceAccount,
258 network: job.network,
259 labels: job.labels,
260 runtimeConfig: job.runtimeConfig,
261 }
262 : undefined,
263 };
264};
265
266const toNotebookRequest = (
267 request: NotebookExecutionJobRequest | undefined,
268):
269 | aiplatform.GoogleCloudAiplatformV1CreateNotebookExecutionJobRequest
270 | undefined => {
271 if (request === undefined) return undefined;
272 return {
273 parent: request.parent,
274 notebookExecutionJobId: request.notebookExecutionJobId,
275 notebookExecutionJob:
276 request.notebookExecutionJob as aiplatform.GoogleCloudAiplatformV1NotebookExecutionJob,
277 };
278};
279
280export const ScheduleProvider = () =>
281 Provider.succeed(Schedule, {
282 stables: ["name", "scheduleId", "location", "project", "createTime"],
283
284 diff: Effect.fn(function* ({ news, olds, output }) {
285 if (!isResolved(news)) return undefined;
286 const previousLocation = olds?.location ?? output?.location;
287 const nextLocation = news.location ?? previousLocation;
288 if (previousLocation !== undefined && previousLocation !== nextLocation) {
289 return { action: "replace" as const, deleteFirst: false };
290 }
291 return undefined;
292 }),
293
294 read: Effect.fn(function* ({ id, olds, output }) {
295 const env = yield* GcpEnvironment.current;
296 const existing = yield* getByName(output?.name ?? "");
297 if (existing !== undefined) {
298 const attrs = toAttrs(existing, env.project);
299 return (yield* ownedByAlchemy(id, existing.displayName))
300 ? attrs
301 : Unowned(attrs);
302 }
303 const location = olds?.location ?? output?.location ?? env.region;
304 const parent = locationParent(env.project, location);
305 const ownership = yield* createInternalLabels(id);
306 const displayName = encodeOwnershipLine(ownership, olds?.displayName);
307 const found = yield* aiplatform.listProjectsLocationsSchedules
308 .pages({ parent, pageSize: 100 })
309 .pipe(
310 Stream.flatMap((page) => Stream.fromIterable(page.schedules ?? [])),
311 Stream.filter((schedule) => schedule.displayName === displayName),
312 Stream.runHead,
313 Effect.map((option) =>
314 option._tag === "Some" ? option.value : undefined,
315 ),
316 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
317 );
318 if (found === undefined) return undefined;
319 const attrs = toAttrs(found, env.project);
320 return (yield* ownedByAlchemy(id, found.displayName))
321 ? attrs
322 : Unowned(attrs);
323 }),
324
325 list: () =>
326 Effect.gen(function* () {
327 const env = yield* GcpEnvironment.current;
328 return (yield* Effect.forEach(listLocations(env.region), (location) =>
329 listAt(locationParent(env.project, location), env.project),
330 )).flat();
331 }),
332
333 reconcile: Effect.fn(function* ({ id, news, output }) {
334 const env = yield* GcpEnvironment.current;
335 const location = news.location ?? output?.location ?? env.region;
336 const parent = locationParent(env.project, location);
337 const ownership = yield* createInternalLabels(id);
338 const displayName = encodeOwnershipLine(ownership, news.displayName);
339 const maxConcurrentRunCount = news.maxConcurrentRunCount ?? "1";
340 const desiredPaused = news.paused === true;
341 const pipelineRequest = toPipelineRequest(
342 news.createPipelineJobRequest,
343 parent,
344 );
345 const notebookRequest = toNotebookRequest(
346 news.createNotebookExecutionJobRequest,
347 );
348
349 let current = yield* getByName(output?.name ?? "");
350 if (current === undefined) {
351 current = yield* aiplatform.listProjectsLocationsSchedules
352 .pages({ parent, pageSize: 100 })
353 .pipe(
354 Stream.flatMap((page) => Stream.fromIterable(page.schedules ?? [])),
355 Stream.filter((schedule) => schedule.displayName === displayName),
356 Stream.runHead,
357 Effect.map((option) =>
358 option._tag === "Some" ? option.value : undefined,
359 ),
360 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
361 );
362 }
363
364 if (current === undefined) {
365 const created = yield* aiplatform
366 .createProjectsLocationsSchedules({
367 parent,
368 body: {
369 displayName,
370 cron: news.cron,
371 maxConcurrentRunCount,
372 maxConcurrentActiveRunCount: news.maxConcurrentActiveRunCount,
373 maxRunCount: news.maxRunCount,
374 startTime: news.startTime,
375 endTime: news.endTime,
376 allowQueueing: news.allowQueueing === true ? true : undefined,
377 createPipelineJobRequest: pipelineRequest,
378 createNotebookExecutionJobRequest: notebookRequest,
379 },
380 })
381 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
382 current = created ?? undefined;
383 if (current === undefined) {
384 current = yield* aiplatform.listProjectsLocationsSchedules
385 .pages({ parent, pageSize: 100 })
386 .pipe(
387 Stream.flatMap((page) =>
388 Stream.fromIterable(page.schedules ?? []),
389 ),
390 Stream.filter((schedule) => schedule.displayName === displayName),
391 Stream.runHead,
392 Effect.map((option) =>
393 option._tag === "Some" ? option.value : undefined,
394 ),
395 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
396 );
397 }
398 }
399
400 if (current === undefined) {
401 return yield* new ScheduleNotResolved({
402 name: output?.name ?? `${parent}/schedules/-`,
403 });
404 }
405
406 const name = current.name ?? "";
407 const cronChanged = (current.cron ?? "") !== news.cron;
408 const displayChanged = (current.displayName ?? "") !== displayName;
409 const concurrentChanged =
410 (current.maxConcurrentRunCount ?? "") !== maxConcurrentRunCount;
411 const activeChanged =
412 (current.maxConcurrentActiveRunCount ?? "") !==
413 (news.maxConcurrentActiveRunCount ?? "");
414 const maxRunChanged =
415 (current.maxRunCount ?? "") !== (news.maxRunCount ?? "");
416 const startChanged = (current.startTime ?? "") !== (news.startTime ?? "");
417 const endChanged = (current.endTime ?? "") !== (news.endTime ?? "");
418 const queueChanged =
419 (current.allowQueueing === true) !== (news.allowQueueing === true);
420
421 if (
422 cronChanged ||
423 displayChanged ||
424 concurrentChanged ||
425 activeChanged ||
426 maxRunChanged ||
427 startChanged ||
428 endChanged ||
429 queueChanged
430 ) {
431 current = yield* aiplatform.patchProjectsLocationsSchedules({
432 name,
433 updateMask: [
434 cronChanged ? "cron" : undefined,
435 displayChanged ? "display_name" : undefined,
436 concurrentChanged ? "max_concurrent_run_count" : undefined,
437 activeChanged ? "max_concurrent_active_run_count" : undefined,
438 maxRunChanged ? "max_run_count" : undefined,
439 startChanged ? "start_time" : undefined,
440 endChanged ? "end_time" : undefined,
441 queueChanged ? "allow_queueing" : undefined,
442 ]
443 .filter((field): field is string => field !== undefined)
444 .join(","),
445 body: {
446 name,
447 displayName,
448 cron: news.cron,
449 maxConcurrentRunCount,
450 maxConcurrentActiveRunCount: news.maxConcurrentActiveRunCount,
451 maxRunCount: news.maxRunCount,
452 startTime: news.startTime,
453 endTime: news.endTime,
454 allowQueueing: news.allowQueueing === true,
455 },
456 });
457 }
458
459 const isPaused = current.state === "PAUSED";
460 if (desiredPaused && !isPaused) {
461 yield* aiplatform.pauseProjectsLocationsSchedules({ name });
462 current = (yield* getByName(name)) ?? current;
463 } else if (!desiredPaused && isPaused) {
464 yield* aiplatform.resumeProjectsLocationsSchedules({
465 name,
466 body: { catchUp: false },
467 });
468 current = (yield* getByName(name)) ?? current;
469 }
470
471 return toAttrs(current, env.project);
472 }),
473
474 delete: Effect.fn(function* ({ output }) {
475 const existing = yield* getByName(output.name);
476 if (existing === undefined) return;
477 const operation = yield* aiplatform
478 .deleteProjectsLocationsSchedules({ name: output.name })
479 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
480 if (operation !== undefined) {
481 yield* waitForOperation(operation, { notFoundOk: true });
482 }
483 }),
484 });