1import * as GCP from "@/GCP";
2import { GcpEnvironment } from "@/GCP/Environment";
3import * as Output from "@/Output";
4import * as Test from "@/Test/Alchemy";
5import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
6import { expect } from "alchemy-test";
7import * as Effect from "effect/Effect";
8import { MinimumLogLevel } from "effect/References";
9import * as ScheduleLib from "effect/Schedule";
10
11const { test } = Test.make({ providers: GCP.providers() });
12
13const logLevel = Effect.provideService(
14 MinimumLogLevel,
15 process.env.DEBUG ? "Debug" : "Info",
16);
17
18const waitUntilGone = (name: string) =>
19 aiplatform.getProjectsLocationsSchedules({ name }).pipe(
20 Effect.as("found" as const),
21 Effect.catchTag("NotFound", () => Effect.succeed("gone" as const)),
22 Effect.repeat({
23 schedule: ScheduleLib.spaced("2 seconds"),
24 until: (status) => status === "gone",
25 times: 10,
26 }),
27 );
28
29test.provider(
30 "getProjectsLocationsSchedules on a missing schedule fails with a typed tag",
31 (stack) =>
32 Effect.gen(function* () {
33 const { project } = yield* GcpEnvironment.current;
34 yield* stack.destroy();
35
36 const error = yield* Effect.flip(
37 aiplatform.getProjectsLocationsSchedules({
38 name: `projects/${project}/locations/us-central1/schedules/1234567890123456789`,
39 }),
40 );
41 expect(error._tag).toEqual("NotFound");
42
43 yield* stack.destroy();
44 }).pipe(logLevel),
45 {
46 tags: ["provider:gcp", "provider:gcp:aiplatform", "live"],
47 timeout: 90_000,
48 },
49);
50
51test.provider(
52 "create, pause, and delete a vertex schedule",
53 (stack) =>
54 Effect.gen(function* () {
55 yield* stack.destroy();
56
57 const created = yield* stack.deploy(
58 Effect.gen(function* () {
59 const bucket = yield* GCP.Storage.Bucket("PipelineRoot", {
60 location: "US-CENTRAL1",
61 forceDestroy: true,
62 });
63 return yield* GCP.AIPlatform.Schedule("Nightly", {
64 location: "us-central1",
65 displayName: "nightly",
66 cron: "0 8 * * *",
67 paused: true,
68 maxRunCount: "1",
69 createPipelineJobRequest: {
70 pipelineJob: {
71 displayName: "nightly-hello",
72 templateUri:
73 "https://us-kfp.pkg.dev/ml-pipeline/google-cloud-registry/hello-world/latest",
74 runtimeConfig: {
75 gcsOutputDirectory: Output.interpolate`gs://${bucket.bucketName}/runs`,
76 },
77 },
78 },
79 });
80 }),
81 );
82
83 expect(created.name).toContain("/schedules/");
84 expect(created.paused).toEqual(true);
85 expect(created.cron).toContain("0 8 * * *");
86
87 const fetched = yield* aiplatform.getProjectsLocationsSchedules({
88 name: created.name,
89 });
90 expect(fetched.name).toEqual(created.name);
91 expect(fetched.displayName).toContain("alchemy-id=");
92
93 const updated = yield* stack.deploy(
94 Effect.gen(function* () {
95 const bucket = yield* GCP.Storage.Bucket("PipelineRoot", {
96 location: "US-CENTRAL1",
97 forceDestroy: true,
98 });
99 return yield* GCP.AIPlatform.Schedule("Nightly", {
100 location: "us-central1",
101 displayName: "nightly-v2",
102 cron: "0 9 * * *",
103 paused: true,
104 maxRunCount: "1",
105 createPipelineJobRequest: {
106 pipelineJob: {
107 displayName: "nightly-hello",
108 templateUri:
109 "https://us-kfp.pkg.dev/ml-pipeline/google-cloud-registry/hello-world/latest",
110 runtimeConfig: {
111 gcsOutputDirectory: Output.interpolate`gs://${bucket.bucketName}/runs`,
112 },
113 },
114 },
115 });
116 }),
117 );
118 expect(updated.name).toEqual(created.name);
119 expect(updated.cron).toContain("0 9 * * *");
120
121 yield* stack.destroy();
122 const gone = yield* waitUntilGone(created.name);
123 expect(gone).toEqual("gone");
124 }).pipe(logLevel),
125 {
126 tags: ["provider:gcp", "provider:gcp:aiplatform", "live"],
127 timeout: 120_000,
128 },
129);