Skip to content

Commit 44914fc

Browse files
alberto-mirandaalberto-miranda
authored
feat: Add support for instance templates in GCP CEs (#571)
This commit adds support for two new google-batch flags: - `--head-job-template`: Allows specifying which instance template to use for the head job - `--compute-job-template`: Allows specifying which instance template to use for compute tasks. Signed-off-by: Alberto Miranda <alberto.miranda@seqera.io>
1 parent 334e36c commit 44914fc

2 files changed

Lines changed: 80 additions & 1 deletion

File tree

‎src/main/java/io/seqera/tower/cli/commands/computeenvs/platforms/GoogleBatchPlatform.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,9 @@ public GoogleBatchConfig computeConfig() throws ApiException, IOException {
6868
.bootDiskSizeGb(adv.bootDiskSizeGb)
6969
.headJobCpus(adv.headJobCpus)
7070
.headJobMemoryMb(adv.headJobMemoryMb)
71-
.serviceAccount(adv.serviceAccountEmail);
71+
.serviceAccount(adv.serviceAccountEmail)
72+
.headJobInstanceTemplate(adv.headJobInstanceTemplate)
73+
.computeJobsInstanceTemplate(adv.computeJobInstanceTemplate);
7274
}
7375

7476
// Common
@@ -97,5 +99,10 @@ public static class AdvancedOptions {
9799
@Option(names = {"--service-account-email"}, description = "The service account email address used when deploying pipeline executions with this compute environment.")
98100
public String serviceAccountEmail;
99101

102+
@Option(names = {"--head-job-template"}, description = "The name or fully qualified reference of the instance template to use for the head job.")
103+
public String headJobInstanceTemplate;
104+
105+
@Option(names = {"--compute-job-template"}, description = "The name or fully qualified reference of the instance template to use for the compute jobs.")
106+
public String computeJobInstanceTemplate;
100107
}
101108
}

‎src/test/java/io/seqera/tower/cli/computeenvs/platforms/GoogleBatchPlatformTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,76 @@ void testAddWithAdvancedOptions(MockServerClient mock) throws IOException {
8585
assertEquals(0, out.exitCode);
8686
}
8787

88+
@Test
89+
void testAddWithInstanceTemplates(MockServerClient mock) throws IOException {
90+
91+
mock.reset();
92+
93+
mock.when(
94+
request().withMethod("GET").withPath("/credentials").withQueryStringParameter("platformId", "google-batch"), exactly(1)
95+
).respond(
96+
response().withStatusCode(200).withBody("{\"credentials\":[{\"id\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"description\":null,\"discriminator\":\"google\",\"baseUrl\":null,\"category\":null,\"deleted\":null,\"lastUsed\":\"2021-09-08T18:20:46Z\",\"dateCreated\":\"2021-09-08T12:57:04Z\",\"lastUpdated\":\"2021-09-08T12:57:04Z\"}]}").withContentType(MediaType.APPLICATION_JSON)
97+
);
98+
99+
mock.when(
100+
request().withMethod("POST").withPath("/compute-envs")
101+
.withBody(JsonBody.json("{\"computeEnv\":{\"credentialsId\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"platform\":\"google-batch\",\"config\":{\"location\":\"europe\",\"workDir\":\"gs://workdir\",\"fusion2Enabled\":false,\"waveEnabled\":false,\"headJobInstanceTemplate\":\"projects/my-project/global/instanceTemplates/head-template\",\"computeJobsInstanceTemplate\":\"projects/my-project/global/instanceTemplates/compute-template\"}}}")), exactly(1)
102+
).respond(
103+
response().withStatusCode(200).withBody("{\"computeEnvId\":\"isnEDBLvHDAIteOEF44ow\"}").withContentType(MediaType.APPLICATION_JSON)
104+
);
105+
106+
ExecOut out = exec(mock, "compute-envs", "add", "google-batch", "-n", "google", "--work-dir", "gs://workdir", "-l", "europe", "--head-job-template", "projects/my-project/global/instanceTemplates/head-template", "--compute-job-template", "projects/my-project/global/instanceTemplates/compute-template");
107+
assertEquals("", out.stdErr);
108+
assertEquals(new ComputeEnvAdded("google-batch", "isnEDBLvHDAIteOEF44ow", "google", null, USER_WORKSPACE_NAME).toString(), out.stdOut);
109+
assertEquals(0, out.exitCode);
110+
}
111+
112+
@Test
113+
void testAddWithOnlyHeadJobTemplate(MockServerClient mock) throws IOException {
114+
115+
mock.reset();
116+
117+
mock.when(
118+
request().withMethod("GET").withPath("/credentials").withQueryStringParameter("platformId", "google-batch"), exactly(1)
119+
).respond(
120+
response().withStatusCode(200).withBody("{\"credentials\":[{\"id\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"description\":null,\"discriminator\":\"google\",\"baseUrl\":null,\"category\":null,\"deleted\":null,\"lastUsed\":\"2021-09-08T18:20:46Z\",\"dateCreated\":\"2021-09-08T12:57:04Z\",\"lastUpdated\":\"2021-09-08T12:57:04Z\"}]}").withContentType(MediaType.APPLICATION_JSON)
121+
);
122+
123+
mock.when(
124+
request().withMethod("POST").withPath("/compute-envs")
125+
.withBody(JsonBody.json("{\"computeEnv\":{\"credentialsId\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"platform\":\"google-batch\",\"config\":{\"location\":\"europe\",\"workDir\":\"gs://workdir\",\"fusion2Enabled\":false,\"waveEnabled\":false,\"headJobInstanceTemplate\":\"projects/my-project/global/instanceTemplates/head-template\"}}}")), exactly(1)
126+
).respond(
127+
response().withStatusCode(200).withBody("{\"computeEnvId\":\"isnEDBLvHDAIteOEF44ow\"}").withContentType(MediaType.APPLICATION_JSON)
128+
);
129+
130+
ExecOut out = exec(mock, "compute-envs", "add", "google-batch", "-n", "google", "--work-dir", "gs://workdir", "-l", "europe", "--head-job-template", "projects/my-project/global/instanceTemplates/head-template");
131+
assertEquals("", out.stdErr);
132+
assertEquals(new ComputeEnvAdded("google-batch", "isnEDBLvHDAIteOEF44ow", "google", null, USER_WORKSPACE_NAME).toString(), out.stdOut);
133+
assertEquals(0, out.exitCode);
134+
}
135+
136+
@Test
137+
void testAddWithOnlyComputeJobTemplate(MockServerClient mock) throws IOException {
138+
139+
mock.reset();
140+
141+
mock.when(
142+
request().withMethod("GET").withPath("/credentials").withQueryStringParameter("platformId", "google-batch"), exactly(1)
143+
).respond(
144+
response().withStatusCode(200).withBody("{\"credentials\":[{\"id\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"description\":null,\"discriminator\":\"google\",\"baseUrl\":null,\"category\":null,\"deleted\":null,\"lastUsed\":\"2021-09-08T18:20:46Z\",\"dateCreated\":\"2021-09-08T12:57:04Z\",\"lastUpdated\":\"2021-09-08T12:57:04Z\"}]}").withContentType(MediaType.APPLICATION_JSON)
145+
);
146+
147+
mock.when(
148+
request().withMethod("POST").withPath("/compute-envs")
149+
.withBody(JsonBody.json("{\"computeEnv\":{\"credentialsId\":\"6XfOhoztUq6de3Dw3X9LSb\",\"name\":\"google\",\"platform\":\"google-batch\",\"config\":{\"location\":\"europe\",\"workDir\":\"gs://workdir\",\"fusion2Enabled\":false,\"waveEnabled\":false,\"computeJobsInstanceTemplate\":\"projects/my-project/global/instanceTemplates/compute-template\"}}}")), exactly(1)
150+
).respond(
151+
response().withStatusCode(200).withBody("{\"computeEnvId\":\"isnEDBLvHDAIteOEF44ow\"}").withContentType(MediaType.APPLICATION_JSON)
152+
);
153+
154+
ExecOut out = exec(mock, "compute-envs", "add", "google-batch", "-n", "google", "--work-dir", "gs://workdir", "-l", "europe", "--compute-job-template", "projects/my-project/global/instanceTemplates/compute-template");
155+
assertEquals("", out.stdErr);
156+
assertEquals(new ComputeEnvAdded("google-batch", "isnEDBLvHDAIteOEF44ow", "google", null, USER_WORKSPACE_NAME).toString(), out.stdOut);
157+
assertEquals(0, out.exitCode);
158+
}
159+
88160
}

0 commit comments

Comments
 (0)