1import * as GCP from "@/GCP";
2import * as Test from "@/Test/Alchemy";
3import * as tpu from "@distilled.cloud/gcp/tpu_v2";
4import { expect } from "alchemy-test";
5import * as Effect from "effect/Effect";
6import { MinimumLogLevel } from "effect/References";
7import * as Schedule from "effect/Schedule";
8import { GcpEnvironment } from "@/GCP/Environment";
9
10const { test } = Test.make({ providers: GCP.providers() });
11
12const logLevel = Effect.provideService(
13 MinimumLogLevel,
14 process.env.DEBUG ? "Debug" : "Info",
15);
16
17// Cloud TPU is disabled on the testing project (ServiceDisabled "Cloud TPU API
18// has not been used in project ...") and TPU VMs bill by the hour (dollars
19// per run), need TPU quota, and take
20// several minutes to provision; set GCP_TEST_TPU=1 to opt in.
21const runLifecycle = !!process.env.GCP_TEST_TPU && !process.env.FAST;
22
23const waitUntilGone = (name: string) =>
24 tpu.getProjectsLocationsNodes({ name }).pipe(
25 Effect.as("found" as const),
26 Effect.catchTag("NotFound", () => Effect.succeed("gone" as const)),
27 Effect.repeat({
28 schedule: Schedule.spaced("2 seconds"),
29 until: (status) => status === "gone",
30 times: 10,
31 }),
32 );
33
34test.provider(
35 "getProjectsLocationsNodes on a missing node fails with a typed tag",
36 (stack) =>
37 Effect.gen(function* () {
38 const { project } = yield* GcpEnvironment.current;
39 yield* stack.destroy();
40
41 const error = yield* Effect.flip(
42 tpu.getProjectsLocationsNodes({
43 name: `projects/${project}/locations/us-central1-c/nodes/alchemy-tpu-missing`,
44 }),
45 );
46 expect(error._tag).toEqual("ServiceDisabled");
47
48 yield* stack.destroy();
49 }).pipe(logLevel),
50 { tags: ["provider:gcp", "provider:gcp:tpu", "live"], timeout: 90_000 },
51);
52
53test.provider.skipIf(!runLifecycle)(
54 "create, update, and delete a TPU node",
55 (stack) =>
56 Effect.gen(function* () {
57 yield* stack.destroy();
58
59 const created = yield* stack.deploy(
60 Effect.gen(function* () {
61 return yield* GCP.TPU.Node("Trainer", {
62 location: "us-central1-c",
63 acceleratorType: "v2-8",
64 runtimeVersion: "tpu-ubuntu2204-base",
65 description: "alchemy-test-tpu",
66 labels: { env: "test" },
67 });
68 }),
69 );
70
71 expect(created.name).toContain("/nodes/");
72 expect(created.nodeId).toEqual(expect.any(String));
73 expect(created.location).toEqual("us-central1-c");
74 expect(created.acceleratorType).toEqual("v2-8");
75 expect(created.description).toEqual("alchemy-test-tpu");
76 expect(created.labels).toMatchObject({ env: "test" });
77
78 const fetched = yield* tpu.getProjectsLocationsNodes({
79 name: created.name,
80 });
81 expect(fetched.name).toEqual(created.name);
82 expect(fetched.acceleratorType).toEqual("v2-8");
83 expect(fetched.description).toEqual("alchemy-test-tpu");
84 expect(fetched.labels?.env).toEqual("test");
85
86 const updated = yield* stack.deploy(
87 Effect.gen(function* () {
88 return yield* GCP.TPU.Node("Trainer", {
89 nodeId: created.nodeId,
90 location: "us-central1-c",
91 acceleratorType: "v2-8",
92 runtimeVersion: "tpu-ubuntu2204-base",
93 description: "alchemy-prod-tpu",
94 labels: { env: "prod", role: "tpu" },
95 });
96 }),
97 );
98
99 expect(updated.name).toEqual(created.name);
100 expect(updated.description).toEqual("alchemy-prod-tpu");
101 expect(updated.labels).toMatchObject({ env: "prod", role: "tpu" });
102
103 const refetched = yield* tpu.getProjectsLocationsNodes({
104 name: created.name,
105 });
106 expect(refetched.description).toEqual("alchemy-prod-tpu");
107 expect(refetched.labels?.env).toEqual("prod");
108 expect(refetched.labels?.role).toEqual("tpu");
109
110 yield* stack.destroy();
111
112 const gone = yield* waitUntilGone(created.name);
113 expect(gone).toEqual("gone");
114 }).pipe(logLevel),
115 { tags: ["provider:gcp", "provider:gcp:tpu", "live"], timeout: 120_000 },
116);