1import * as GCP from "@/GCP";
2import * as Test from "@/Test/Alchemy";
3import * as Core from "@/Test/Core";
4import * as resourcemanager from "@distilled.cloud/gcp/cloudresourcemanager_v3";
5import * as tpu from "@distilled.cloud/gcp/tpu_v2";
6import { describe, expect } from "alchemy-test";
7import * as Effect from "effect/Effect";
8import { dockerAvailable, expectProbe } from "../bindingHost.ts";
9import TpuBindingsHost, {
10 QueuedTrainer,
11 Trainer,
12} from "./fixtures/bindings-host.ts";
13
14const testOptions = { providers: GCP.providers() };
15const { test, beforeAll, afterAll } = Test.make(testOptions);
16const sharedStack = Core.scratchStack(testOptions, "TpuBindings");
17
18// Cloud TPU is disabled on the testing project (ServiceDisabled "Cloud TPU API
19// has not been used in project ...") and TPU VMs bill by the hour (dollars
20// per run) and need TPU quota; set GCP_TEST_TPU=1 to opt in.
21const runLifecycle = !!process.env.GCP_TEST_TPU && !process.env.FAST;
22
23let baseUrl: string;
24let hostAccount: string;
25let project: string;
26let nodeName: string;
27let queuedResourceName: string;
28
29/** Project-level roles held by the host's service account. */
30const projectRoles = () =>
31 resourcemanager
32 .getIamPolicyProjects({
33 resource: `projects/${project}`,
34 body: { options: { requestedPolicyVersion: 3 } },
35 })
36 .pipe(
37 Effect.map((policy) =>
38 (policy.bindings ?? [])
39 .filter((binding) =>
40 (binding.members ?? []).includes(`serviceAccount:${hostAccount}`),
41 )
42 .map((binding) => ({
43 role: binding.role,
44 condition: binding.condition?.expression,
45 })),
46 ),
47 );
48
49describe.skipIf(!dockerAvailable || !runLifecycle)(
50 "TPU Bindings",
51 { tags: ["provider:gcp", "provider:gcp:tpu", "provider:gcp:run", "live"] },
52 () => {
53 beforeAll(
54 Effect.gen(function* () {
55 yield* sharedStack.destroy();
56 const out = yield* sharedStack.deploy(
57 Effect.gen(function* () {
58 const host = yield* TpuBindingsHost;
59 const node = yield* Trainer;
60 const queued = yield* QueuedTrainer;
61 return {
62 uri: host.uri,
63 serviceAccount: host.serviceAccount,
64 project: host.project,
65 node: node.name,
66 queued: queued.name,
67 };
68 }),
69 );
70 baseUrl = out.uri!;
71 hostAccount = out.serviceAccount!;
72 project = out.project;
73 nodeName = out.node;
74 queuedResourceName = out.queued;
75 }),
76 { timeout: 1_800_000 },
77 );
78
79 afterAll(sharedStack.destroy(), { timeout: 1_200_000 });
80
81 // Cloud TPU has no resource-level IAM: both bindings grant tpu.viewer
82 // on the project.
83 const expectTpuViewerOnly = Effect.gen(function* () {
84 expect(yield* projectRoles()).toEqual([
85 { role: "roles/tpu.viewer", condition: undefined },
86 ]);
87 });
88
89 describe("GetNode", () => {
90 test.provider(
91 "reads the bound node as the host",
92 (_stack) =>
93 Effect.gen(function* () {
94 const out = yield* expectProbe<{
95 name: string;
96 acceleratorType: string;
97 }>(baseUrl, "getNode");
98 const live = yield* tpu.getProjectsLocationsNodes({
99 name: nodeName,
100 });
101 expect(out.name).toEqual(live.name);
102 expect(out.acceleratorType).toEqual("v2-8");
103 yield* expectTpuViewerOnly;
104 }),
105 {
106 tags: ["provider:gcp", "provider:gcp:tpu", "live"],
107 timeout: 600_000,
108 },
109 );
110 });
111
112 describe("GetQueuedResource", () => {
113 test.provider(
114 "reads the bound queued resource as the host",
115 (_stack) =>
116 Effect.gen(function* () {
117 const out = yield* expectProbe<{ name: string }>(
118 baseUrl,
119 "getQueuedResource",
120 );
121 const live = yield* tpu.getProjectsLocationsQueuedResources({
122 name: queuedResourceName,
123 });
124 expect(out.name).toEqual(live.name);
125 yield* expectTpuViewerOnly;
126 }),
127 {
128 tags: ["provider:gcp", "provider:gcp:tpu", "live"],
129 timeout: 600_000,
130 },
131 );
132 });
133 },
134);