Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Disable registration of new workflows through workflow management api #3388

Merged
merged 2 commits into from
Feb 13, 2025
Merged

Disable registration of new workflows through workflow management api #3388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8c5cdf6
Disable registration of new workflows through workflow management api
mreso Feb 13, 2025
bc71e09
Disable management tests
mreso Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 2 additions & 111 deletions frontend/server/src/main/java/org/pytorch/serve/workflow/WorkflowManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,117 +122,8 @@ public StatusResponse registerWorkflow(
boolean s3SseKms)
throws WorkflowException {

if (url == null) {
throw new BadRequestException("Parameter url is required.");
}

StatusResponse status = new StatusResponse();

ExecutorService executorService = Executors.newFixedThreadPool(4);
CompletionService<ModelRegistrationResult> executorCompletionService =
new ExecutorCompletionService<>(executorService);
boolean failed = false;
ArrayList<String> failedMessages = new ArrayList<>();
ArrayList<String> successNodes = new ArrayList<>();
try {
WorkflowArchive archive = createWorkflowArchive(workflowName, url);
WorkFlow workflow = createWorkflow(archive);

if (workflowMap.get(workflow.getWorkflowArchive().getWorkflowName()) != null) {
throw new ConflictStatusException(
"Workflow "
+ workflow.getWorkflowArchive().getWorkflowName()
+ " is already registered.");
}

Map<String, Node> nodes = workflow.getDag().getNodes();

List<Future<ModelRegistrationResult>> futures = new ArrayList<>();

for (Map.Entry<String, Node> entry : nodes.entrySet()) {
Node node = entry.getValue();
WorkflowModel wfm = node.getWorkflowModel();

futures.add(
executorCompletionService.submit(
() ->
registerModelWrapper(
wfm,
responseTimeout,
startupTimeout,
synchronous)));
}

int i = 0;
while (i < futures.size()) {
i++;
Future<ModelRegistrationResult> future = executorCompletionService.take();
ModelRegistrationResult result = future.get();
if (result.getResponse().getHttpResponseCode() != HttpURLConnection.HTTP_OK) {
failed = true;
String msg;
if (result.getResponse().getStatus() == null) {
msg =
"Failed to register the model "
+ result.getModelName()
+ ". Check error logs.";
} else {
msg = result.getResponse().getStatus();
}
failedMessages.add(msg);
} else {
successNodes.add(result.getModelName());
}
}

if (failed) {
String rollbackFailure = null;
try {
removeArtifacts(workflowName, workflow, successNodes);
} catch (Exception e) {
rollbackFailure =
"Error while doing rollback of failed workflow. Details"
+ e.getMessage();
}

if (rollbackFailure != null) {
failedMessages.add(rollbackFailure);
}
status.setHttpResponseCode(HttpURLConnection.HTTP_INTERNAL_ERROR);
String message =
String.format(
"Workflow %s has failed to register. Failures: %s",
workflow.getWorkflowArchive().getWorkflowName(),
failedMessages.toString());
status.setStatus(message);
status.setE(new WorkflowException(message));

} else {
status.setHttpResponseCode(HttpURLConnection.HTTP_OK);
status.setStatus(
String.format(
"Workflow %s has been registered and scaled successfully.",
workflow.getWorkflowArchive().getWorkflowName()));

workflowMap.putIfAbsent(workflow.getWorkflowArchive().getWorkflowName(), workflow);
}

} catch (DownloadArchiveException e) {
status.setHttpResponseCode(HttpURLConnection.HTTP_BAD_REQUEST);
status.setStatus("Failed to download workflow archive file");
status.setE(e);
} catch (InvalidDAGException e) {
status.setHttpResponseCode(HttpURLConnection.HTTP_BAD_REQUEST);
status.setStatus("Invalid workflow specification");
status.setE(e);
} catch (InterruptedException | ExecutionException | IOException e) {
status.setHttpResponseCode(HttpURLConnection.HTTP_INTERNAL_ERROR);
status.setStatus("Failed to register workflow.");
status.setE(e);
} finally {
executorService.shutdown();
}
return status;
// Permanently disable registeration of new workflows
throw new BadRequestException("Registering new workflows is disabled.");
}

public ModelRegistrationResult registerModelWrapper(
Expand Down
5 changes: 2 additions & 3 deletions ts_scripts/api_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,8 +396,8 @@ def trigger_all():
exit_code8 = trigger_inference_tests_kfv2()
exit_code9 = trigger_https_tests_kfv2()
exit_code10 = trigger_explanation_tests()
exit_code11 = trigger_workflow_tests()
# Skipping as this test is flaky
# Skipping as these tests are flaky
# exit_code11 = trigger_workflow_tests()
# exit_code12 = trigger_workflow_inference_tests()
return (
1
Expand All @@ -414,7 +414,6 @@ def trigger_all():
exit_code8,
exit_code9,
exit_code10,
exit_code11,
]
)
else 0
Expand Down
Loading