Skip to content

Commit 05aa321

Browse files
jbern0rdjbern0rd
authored
feat: retrieve and cache TEE sessions metadata during task preflight checks (#664)
1 parent cbde5d2 commit 05aa321

21 files changed

+481
-469
lines changed

‎src/main/java/com/iexec/worker/compute/ComputeManagerService.java‎

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,8 @@
1919
import com.iexec.common.replicate.ReplicateStatusCause;
2020
import com.iexec.common.result.ComputedFile;
2121
import com.iexec.common.utils.FileHelper;
22-
import com.iexec.commons.poco.chain.WorkerpoolAuthorization;
2322
import com.iexec.commons.poco.dapp.DappType;
2423
import com.iexec.commons.poco.task.TaskDescription;
25-
import com.iexec.sms.api.TeeSessionGenerationResponse;
2624
import com.iexec.worker.compute.app.AppComputeResponse;
2725
import com.iexec.worker.compute.app.AppComputeService;
2826
import com.iexec.worker.compute.post.PostComputeResponse;
@@ -91,8 +89,7 @@ public boolean downloadApp(TaskDescription taskDescription) {
9189
if (taskDescription == null || taskDescription.getAppType() == null) {
9290
return false;
9391
}
94-
boolean isDockerType =
95-
taskDescription.getAppType().equals(DappType.DOCKER);
92+
final boolean isDockerType = taskDescription.getAppType() == DappType.DOCKER;
9693
if (!isDockerType || taskDescription.getAppUri() == null) {
9794
return false;
9895
}
@@ -153,23 +150,21 @@ public boolean isAppDownloaded(String imageUri) {
153150
* Execute pre-compute stage for standard and TEE tasks.
154151
* <ul>
155152
* <li>Standard tasks: Nothing is executed, an empty result is returned
156-
* <li>TEE tasks: Call {@link PreComputeService#runTeePreCompute(TaskDescription, WorkerpoolAuthorization)}
153+
* <li>TEE tasks: Call {@link PreComputeService#runTeePreCompute(TaskDescription)}
157154
* </ul>
158155
* TEE tasks: download pre-compute and post-compute images,
159156
* create SCONE secure session, and run pre-compute container.
160157
*
161158
* @param taskDescription Description of the task
162-
* @param workerpoolAuth Authorization to contribute delivered by the scheduler for the given task
163159
* @return {@code PreComputeResponse} instance
164-
* @see PreComputeService#runTeePreCompute(TaskDescription, WorkerpoolAuthorization)
160+
* @see PreComputeService#runTeePreCompute(TaskDescription)
165161
*/
166-
public PreComputeResponse runPreCompute(final TaskDescription taskDescription,
167-
final WorkerpoolAuthorization workerpoolAuth) {
162+
public PreComputeResponse runPreCompute(final TaskDescription taskDescription) {
168163
log.info("Running pre-compute [chainTaskId:{}, isTee:{}]",
169164
taskDescription.getChainTaskId(), taskDescription.isTeeTask());
170165

171166
if (taskDescription.isTeeTask()) {
172-
return preComputeService.runTeePreCompute(taskDescription, workerpoolAuth);
167+
return preComputeService.runTeePreCompute(taskDescription);
173168
}
174169
return PreComputeResponse.builder().build();
175170
}
@@ -178,18 +173,15 @@ public PreComputeResponse runPreCompute(final TaskDescription taskDescription,
178173
* Execute application stage for standard and TEE tasks.
179174
*
180175
* @param taskDescription Description of the task
181-
* @param secureSession Session ID and session storage URL for TEE tasks
182176
* @return {@code AppComputeResponse} instance
183-
* @see AppComputeService#runCompute(TaskDescription, TeeSessionGenerationResponse)
177+
* @see AppComputeService#runCompute(TaskDescription)
184178
*/
185-
public AppComputeResponse runCompute(final TaskDescription taskDescription,
186-
final TeeSessionGenerationResponse secureSession) {
179+
public AppComputeResponse runCompute(final TaskDescription taskDescription) {
187180
final String chainTaskId = taskDescription.getChainTaskId();
188181
log.info("Running compute [chainTaskId:{}, isTee:{}]",
189182
chainTaskId, taskDescription.isTeeTask());
190183

191-
final AppComputeResponse appComputeResponse =
192-
appComputeService.runCompute(taskDescription, secureSession);
184+
final AppComputeResponse appComputeResponse = appComputeService.runCompute(taskDescription);
193185

194186
if (appComputeResponse.isSuccessful()) {
195187
writeLogs(chainTaskId, STDOUT_FILENAME, appComputeResponse.getStdout());
@@ -213,26 +205,20 @@ private void writeLogs(String chainTaskId, String filename, String logs) {
213205
* This method calls methods from {@code PostComputeService} depending on the Task type.
214206
*
215207
* @param taskDescription Description of the task
216-
* @param secureSession Session ID and session storage URL for TEE tasks
217208
* @return {@code PostComputeResponse} instance
218209
* @see PostComputeService#runStandardPostCompute(TaskDescription)
219-
* @see PostComputeService#runTeePostCompute(TaskDescription, TeeSessionGenerationResponse)
210+
* @see PostComputeService#runTeePostCompute(TaskDescription)
220211
*/
221-
public PostComputeResponse runPostCompute(final TaskDescription taskDescription,
222-
final TeeSessionGenerationResponse secureSession) {
212+
public PostComputeResponse runPostCompute(final TaskDescription taskDescription) {
223213
final String chainTaskId = taskDescription.getChainTaskId();
224214
log.info("Running post-compute [chainTaskId:{}, isTee:{}]",
225215
chainTaskId, taskDescription.isTeeTask());
226216

227217
final PostComputeResponse postComputeResponse;
228218
if (!taskDescription.isTeeTask()) {
229219
postComputeResponse = postComputeService.runStandardPostCompute(taskDescription);
230-
} else if (secureSession != null) {
231-
postComputeResponse = postComputeService.runTeePostCompute(taskDescription, secureSession);
232220
} else {
233-
postComputeResponse = PostComputeResponse.builder()
234-
.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.POST_COMPUTE_FAILED_UNKNOWN_ISSUE)))
235-
.build();
221+
postComputeResponse = postComputeService.runTeePostCompute(taskDescription);
236222
}
237223
if (!postComputeResponse.isSuccessful()) {
238224
return postComputeResponse;

‎src/main/java/com/iexec/worker/compute/app/AppComputeResponse.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
@Value
2727
@Builder
2828
public class AppComputeResponse implements ComputeResponse {
29-
3029
@Builder.Default
3130
List<WorkflowError> exitCauses = List.of();
3231
String stdout;

‎src/main/java/com/iexec/worker/compute/app/AppComputeService.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.iexec.commons.containers.DockerRunResponse;
2626
import com.iexec.commons.containers.SgxDriverMode;
2727
import com.iexec.commons.poco.task.TaskDescription;
28-
import com.iexec.sms.api.TeeSessionGenerationResponse;
2928
import com.iexec.worker.config.WorkerConfigurationService;
3029
import com.iexec.worker.docker.DockerService;
3130
import com.iexec.worker.metric.ComputeDurationsService;
@@ -61,8 +60,7 @@ public AppComputeService(
6160
this.appComputeDurationsService = appComputeDurationsService;
6261
}
6362

64-
public AppComputeResponse runCompute(final TaskDescription taskDescription,
65-
final TeeSessionGenerationResponse secureSession) {
63+
public AppComputeResponse runCompute(final TaskDescription taskDescription) {
6664
final String chainTaskId = taskDescription.getChainTaskId();
6765

6866
final List<Bind> binds = new ArrayList<>();
@@ -73,7 +71,7 @@ public AppComputeResponse runCompute(final TaskDescription taskDescription,
7371
final List<String> env;
7472
if (taskDescription.isTeeTask()) {
7573
final TeeService teeService = teeServicesManager.getTeeService(taskDescription.getTeeFramework());
76-
env = teeService.buildComputeDockerEnv(taskDescription, secureSession);
74+
env = teeService.buildComputeDockerEnv(taskDescription);
7775
binds.addAll(teeService.getAdditionalBindings().stream().map(Bind::parse).toList());
7876
sgxDriverMode = sgxService.getSgxDriverMode();
7977
} else {

‎src/main/java/com/iexec/worker/compute/post/PostComputeResponse.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
@Value
2727
@Builder
2828
public class PostComputeResponse implements ComputeResponse {
29-
3029
@Builder.Default
3130
List<WorkflowError> exitCauses = List.of();
3231
String stdout;

‎src/main/java/com/iexec/worker/compute/post/PostComputeService.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -161,8 +161,7 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
161161
return Optional.empty();
162162
}
163163

164-
public PostComputeResponse runTeePostCompute(TaskDescription taskDescription,
165-
TeeSessionGenerationResponse secureSession) {
164+
public PostComputeResponse runTeePostCompute(final TaskDescription taskDescription) {
166165
String chainTaskId = taskDescription.getChainTaskId();
167166

168167
TeeServicesProperties properties =
@@ -179,7 +178,7 @@ public PostComputeResponse runTeePostCompute(TaskDescription taskDescription,
179178
}
180179
TeeService teeService = teeServicesManager.getTeeService(taskDescription.getTeeFramework());
181180
List<String> env = teeService
182-
.buildPostComputeDockerEnv(taskDescription, secureSession);
181+
.buildPostComputeDockerEnv(taskDescription);
183182
List<Bind> binds = Stream.of(
184183
Collections.singletonList(dockerService.getIexecOutBind(chainTaskId)),
185184
teeService.getAdditionalBindings())

‎src/main/java/com/iexec/worker/compute/pre/PreComputeResponse.java‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package com.iexec.worker.compute.pre;
1818

19-
import com.iexec.sms.api.TeeSessionGenerationResponse;
2019
import com.iexec.worker.compute.ComputeResponse;
2120
import com.iexec.worker.workflow.WorkflowError;
2221
import lombok.Builder;
@@ -27,19 +26,8 @@
2726
@Value
2827
@Builder
2928
public class PreComputeResponse implements ComputeResponse {
30-
3129
@Builder.Default
3230
List<WorkflowError> exitCauses = List.of();
33-
boolean isTeeTask;
34-
TeeSessionGenerationResponse secureSession;
3531
String stdout;
3632
String stderr;
37-
38-
@Override
39-
public boolean isSuccessful() {
40-
if (isTeeTask) {
41-
return ComputeResponse.super.isSuccessful() && secureSession != null;
42-
}
43-
return ComputeResponse.super.isSuccessful();
44-
}
4533
}

‎src/main/java/com/iexec/worker/compute/pre/PreComputeService.java‎

Lines changed: 10 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,8 @@
2222
import com.iexec.commons.containers.DockerRunFinalStatus;
2323
import com.iexec.commons.containers.DockerRunRequest;
2424
import com.iexec.commons.containers.DockerRunResponse;
25-
import com.iexec.commons.poco.chain.WorkerpoolAuthorization;
2625
import com.iexec.commons.poco.task.TaskDescription;
2726
import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
28-
import com.iexec.sms.api.TeeSessionGenerationError;
29-
import com.iexec.sms.api.TeeSessionGenerationResponse;
3027
import com.iexec.sms.api.config.TeeAppProperties;
3128
import com.iexec.sms.api.config.TeeServicesProperties;
3229
import com.iexec.worker.compute.ComputeExitCauseService;
@@ -35,8 +32,6 @@
3532
import com.iexec.worker.docker.DockerService;
3633
import com.iexec.worker.metric.ComputeDurationsService;
3734
import com.iexec.worker.sgx.SgxService;
38-
import com.iexec.worker.sms.SmsService;
39-
import com.iexec.worker.sms.TeeSessionGenerationException;
4035
import com.iexec.worker.tee.TeeServicesManager;
4136
import com.iexec.worker.tee.TeeServicesPropertiesService;
4237
import com.iexec.worker.workflow.WorkflowError;
@@ -49,13 +44,10 @@
4944
import java.util.List;
5045
import java.util.concurrent.TimeoutException;
5146

52-
import static com.iexec.sms.api.TeeSessionGenerationError.UNKNOWN_ISSUE;
53-
5447
@Slf4j
5548
@Service
5649
public class PreComputeService {
5750

58-
private final SmsService smsService;
5951
private final DockerService dockerService;
6052
private final TeeServicesManager teeServicesManager;
6153
private final WorkerConfigurationService workerConfigService;
@@ -65,15 +57,13 @@ public class PreComputeService {
6557
private final ComputeDurationsService preComputeDurationsService;
6658

6759
public PreComputeService(
68-
SmsService smsService,
6960
DockerService dockerService,
7061
TeeServicesManager teeServicesManager,
7162
WorkerConfigurationService workerConfigService,
7263
SgxService sgxService,
7364
ComputeExitCauseService computeExitCauseService,
7465
TeeServicesPropertiesService teeServicesPropertiesService,
7566
ComputeDurationsService preComputeDurationsService) {
76-
this.smsService = smsService;
7767
this.dockerService = dockerService;
7868
this.teeServicesManager = teeServicesManager;
7969
this.workerConfigService = workerConfigService;
@@ -90,13 +80,11 @@ public PreComputeService(
9080
* is started to handle them.
9181
*
9282
* @param taskDescription Task description read on-chain
93-
* @param workerpoolAuth Workerpool authorization provided by scheduler
9483
* @return PreComputeResponse
9584
*/
96-
public PreComputeResponse runTeePreCompute(TaskDescription taskDescription, WorkerpoolAuthorization workerpoolAuth) {
85+
public PreComputeResponse runTeePreCompute(final TaskDescription taskDescription) {
9786
final String chainTaskId = taskDescription.getChainTaskId();
98-
final PreComputeResponse.PreComputeResponseBuilder preComputeResponseBuilder = PreComputeResponse.builder()
99-
.isTeeTask(taskDescription.isTeeTask());
87+
final PreComputeResponse.PreComputeResponseBuilder preComputeResponseBuilder = PreComputeResponse.builder();
10088

10189
// verify enclave configuration for compute stage
10290
final TeeEnclaveConfiguration enclaveConfig = taskDescription.getAppEnclaveConfiguration();
@@ -122,37 +110,21 @@ public PreComputeResponse runTeePreCompute(TaskDescription taskDescription, Work
122110
preComputeResponseBuilder.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_HEAP_CONFIGURATION)));
123111
return preComputeResponseBuilder.build();
124112
}
125-
// create secure session
126-
final TeeSessionGenerationResponse secureSession;
127-
try {
128-
secureSession = smsService.createTeeSession(workerpoolAuth);
129-
if (secureSession == null) {
130-
throw new TeeSessionGenerationException(UNKNOWN_ISSUE);
131-
}
132-
preComputeResponseBuilder.secureSession(secureSession);
133-
} catch (TeeSessionGenerationException e) {
134-
log.error("Failed to create TEE secure session [chainTaskId:{}]", chainTaskId, e);
135-
return preComputeResponseBuilder
136-
.exitCauses(List.of(new WorkflowError(teeSessionGenerationErrorToReplicateStatusCause(e.getTeeSessionGenerationError()))))
137-
.build();
138-
}
139113

140114
// run TEE pre-compute container if needed
141115
if (taskDescription.requiresPreCompute()) {
142116
log.info("Task contains TEE input data [chainTaskId:{}, containsDataset:{}, containsInputFiles:{}, isBulkRequest:{}]",
143117
chainTaskId, taskDescription.containsDataset(), taskDescription.containsInputFiles(), taskDescription.isBulkRequest());
144-
final List<WorkflowError> exitCauses = downloadDatasetAndFiles(taskDescription, secureSession);
118+
final List<WorkflowError> exitCauses = downloadDatasetAndFiles(taskDescription);
145119
preComputeResponseBuilder.exitCauses(exitCauses);
146120
}
147121

148122
return preComputeResponseBuilder.build();
149123
}
150124

151-
private List<WorkflowError> downloadDatasetAndFiles(
152-
final TaskDescription taskDescription,
153-
final TeeSessionGenerationResponse secureSession) {
125+
private List<WorkflowError> downloadDatasetAndFiles(final TaskDescription taskDescription) {
154126
try {
155-
final Integer exitCode = prepareTeeInputData(taskDescription, secureSession);
127+
final Integer exitCode = prepareTeeInputData(taskDescription);
156128
if (exitCode == null || exitCode != 0) {
157129
final String chainTaskId = taskDescription.getChainTaskId();
158130
final List<WorkflowError> exitCauses = getExitCauses(chainTaskId, exitCode);
@@ -179,33 +151,14 @@ private List<WorkflowError> getExitCauses(final String chainTaskId, final Intege
179151
};
180152
}
181153

182-
183-
/**
184-
* {@link TeeSessionGenerationError} and {@link ReplicateStatusCause} are dynamically bound
185-
* such as {@code TeeSessionGenerationError.MEMBER_X == ReplicateStatusCause.TEE_SESSION_GENERATION_MEMBER_X}.
186-
*
187-
* @return {@literal null} if no member of {@link ReplicateStatusCause} matches,
188-
* the matching member otherwise.
189-
*/
190-
ReplicateStatusCause teeSessionGenerationErrorToReplicateStatusCause(TeeSessionGenerationError error) {
191-
try {
192-
return ReplicateStatusCause.valueOf("TEE_SESSION_GENERATION_" + error.name());
193-
} catch (IllegalArgumentException e) {
194-
return null;
195-
}
196-
}
197-
198154
/**
199155
* Run tee-worker-pre-compute docker image. The pre-compute enclave downloads
200156
* the dataset and decrypts it for the compute stage. It also downloads input
201157
* files if requested.
202158
*
203159
* @return pre-compute exit code
204160
*/
205-
private Integer prepareTeeInputData(
206-
TaskDescription taskDescription,
207-
TeeSessionGenerationResponse secureSession)
208-
throws TimeoutException {
161+
private Integer prepareTeeInputData(final TaskDescription taskDescription) throws TimeoutException {
209162
String chainTaskId = taskDescription.getChainTaskId();
210163
log.info("Preparing tee input data [chainTaskId:{}]", chainTaskId);
211164

@@ -221,7 +174,7 @@ private Integer prepareTeeInputData(
221174
}
222175
// run container
223176
List<String> env = teeServicesManager.getTeeService(taskDescription.getTeeFramework())
224-
.buildPreComputeDockerEnv(taskDescription, secureSession);
177+
.buildPreComputeDockerEnv(taskDescription);
225178
List<Bind> binds = Collections.singletonList(Bind.parse(dockerService.getInputBind(chainTaskId)));
226179
HostConfig hostConfig = HostConfig.newHostConfig()
227180
.withBinds(binds)
@@ -244,15 +197,14 @@ private Integer prepareTeeInputData(
244197
}
245198
final DockerRunFinalStatus finalStatus = dockerResponse.getFinalStatus();
246199
if (finalStatus == DockerRunFinalStatus.TIMEOUT) {
247-
log.error("Tee pre-compute container timed out" +
248-
" [chainTaskId:{}, maxExecutionTime:{}]",
200+
log.error("Tee pre-compute container timed out [chainTaskId:{}, maxExecutionTime:{}]",
249201
chainTaskId, taskDescription.getMaxExecutionTime());
250202
throw new TimeoutException("Tee pre-compute container timed out");
251203
}
252204
if (finalStatus == DockerRunFinalStatus.FAILED) {
253205
int exitCode = dockerResponse.getContainerExitCode();
254-
log.error("Tee pre-compute container failed [chainTaskId:{}, " +
255-
"exitCode:{}]", chainTaskId, exitCode);
206+
log.error("Tee pre-compute container failed [chainTaskId:{}, exitCode:{}]",
207+
chainTaskId, exitCode);
256208
return dockerResponse.getContainerExitCode();
257209
}
258210
log.info("Prepared tee input data successfully [chainTaskId:{}]", chainTaskId);

0 commit comments

Comments
 (0)