Skip to content

Commit 524612f

Browse files
committed
Disable registration of new workflows through workflow management api
1 parent 1a3b18b commit 524612f

File tree

1 file changed

+2
-111
lines changed

1 file changed

+2
-111
lines changed

frontend/server/src/main/java/org/pytorch/serve/workflow/WorkflowManager.java

+2-111
Original file line numberDiff line numberDiff line change
@@ -122,117 +122,8 @@ public StatusResponse registerWorkflow(
122122
boolean s3SseKms)
123123
throws WorkflowException {
124124

125-
if (url == null) {
126-
throw new BadRequestException("Parameter url is required.");
127-
}
128-
129-
StatusResponse status = new StatusResponse();
130-
131-
ExecutorService executorService = Executors.newFixedThreadPool(4);
132-
CompletionService<ModelRegistrationResult> executorCompletionService =
133-
new ExecutorCompletionService<>(executorService);
134-
boolean failed = false;
135-
ArrayList<String> failedMessages = new ArrayList<>();
136-
ArrayList<String> successNodes = new ArrayList<>();
137-
try {
138-
WorkflowArchive archive = createWorkflowArchive(workflowName, url);
139-
WorkFlow workflow = createWorkflow(archive);
140-
141-
if (workflowMap.get(workflow.getWorkflowArchive().getWorkflowName()) != null) {
142-
throw new ConflictStatusException(
143-
"Workflow "
144-
+ workflow.getWorkflowArchive().getWorkflowName()
145-
+ " is already registered.");
146-
}
147-
148-
Map<String, Node> nodes = workflow.getDag().getNodes();
149-
150-
List<Future<ModelRegistrationResult>> futures = new ArrayList<>();
151-
152-
for (Map.Entry<String, Node> entry : nodes.entrySet()) {
153-
Node node = entry.getValue();
154-
WorkflowModel wfm = node.getWorkflowModel();
155-
156-
futures.add(
157-
executorCompletionService.submit(
158-
() ->
159-
registerModelWrapper(
160-
wfm,
161-
responseTimeout,
162-
startupTimeout,
163-
synchronous)));
164-
}
165-
166-
int i = 0;
167-
while (i < futures.size()) {
168-
i++;
169-
Future<ModelRegistrationResult> future = executorCompletionService.take();
170-
ModelRegistrationResult result = future.get();
171-
if (result.getResponse().getHttpResponseCode() != HttpURLConnection.HTTP_OK) {
172-
failed = true;
173-
String msg;
174-
if (result.getResponse().getStatus() == null) {
175-
msg =
176-
"Failed to register the model "
177-
+ result.getModelName()
178-
+ ". Check error logs.";
179-
} else {
180-
msg = result.getResponse().getStatus();
181-
}
182-
failedMessages.add(msg);
183-
} else {
184-
successNodes.add(result.getModelName());
185-
}
186-
}
187-
188-
if (failed) {
189-
String rollbackFailure = null;
190-
try {
191-
removeArtifacts(workflowName, workflow, successNodes);
192-
} catch (Exception e) {
193-
rollbackFailure =
194-
"Error while doing rollback of failed workflow. Details"
195-
+ e.getMessage();
196-
}
197-
198-
if (rollbackFailure != null) {
199-
failedMessages.add(rollbackFailure);
200-
}
201-
status.setHttpResponseCode(HttpURLConnection.HTTP_INTERNAL_ERROR);
202-
String message =
203-
String.format(
204-
"Workflow %s has failed to register. Failures: %s",
205-
workflow.getWorkflowArchive().getWorkflowName(),
206-
failedMessages.toString());
207-
status.setStatus(message);
208-
status.setE(new WorkflowException(message));
209-
210-
} else {
211-
status.setHttpResponseCode(HttpURLConnection.HTTP_OK);
212-
status.setStatus(
213-
String.format(
214-
"Workflow %s has been registered and scaled successfully.",
215-
workflow.getWorkflowArchive().getWorkflowName()));
216-
217-
workflowMap.putIfAbsent(workflow.getWorkflowArchive().getWorkflowName(), workflow);
218-
}
219-
220-
} catch (DownloadArchiveException e) {
221-
status.setHttpResponseCode(HttpURLConnection.HTTP_BAD_REQUEST);
222-
status.setStatus("Failed to download workflow archive file");
223-
status.setE(e);
224-
} catch (InvalidDAGException e) {
225-
status.setHttpResponseCode(HttpURLConnection.HTTP_BAD_REQUEST);
226-
status.setStatus("Invalid workflow specification");
227-
status.setE(e);
228-
} catch (InterruptedException | ExecutionException | IOException e) {
229-
status.setHttpResponseCode(HttpURLConnection.HTTP_INTERNAL_ERROR);
230-
status.setStatus("Failed to register workflow.");
231-
status.setE(e);
232-
} finally {
233-
executorService.shutdown();
234-
}
235-
return status;
125+
// Permanently disable registeration of new workflows
126+
throw new BadRequestException("Registering new workflows is disabled.");
236127
}
237128

238129
public ModelRegistrationResult registerModelWrapper(

0 commit comments

Comments
 (0)