Skip to content

Commit

Permalink
Add job queue status to describe API (#2464)
Browse files Browse the repository at this point in the history
* Add support for job queue metrics in describe API

* add test

* update documentation for describe API

* Retain describe API metrics implementation

---------

Co-authored-by: Naman Nandan <namannan@amazon.com>
Co-authored-by: Ankith Gunapal <agunapal@ischool.Berkeley.edu>
  • Loading branch information
3 people committed Aug 2, 2023
1 parent fdcc223 commit 53ff6a5
Show file tree
Hide file tree
Showing 8 changed files with 183 additions and 9 deletions.
28 changes: 24 additions & 4 deletions docs/management_api.md
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,11 @@ curl http://localhost:8081/models/noop
"gpu": false,
"memoryUsage": 89247744
}
]
],
"jobQueueStatus": {
"remainingCapacity": 100,
"pendingRequests": 0
}
}
]
```
Expand Down Expand Up @@ -234,7 +238,11 @@ curl http://localhost:8081/models/noop/2.0
"gpu": false,
"memoryUsage": 89247744
}
]
],
"jobQueueStatus": {
"remainingCapacity": 100,
"pendingRequests": 0
}
}
]
```
Expand Down Expand Up @@ -264,7 +272,11 @@ curl http://localhost:8081/models/noop/all
"gpu": false,
"memoryUsage": 89247744
}
]
],
"jobQueueStatus": {
"remainingCapacity": 100,
"pendingRequests": 0
}
},
{
"modelName": "noop",
Expand All @@ -284,7 +296,11 @@ curl http://localhost:8081/models/noop/all
"gpu": false,
"memoryUsage": 89247744
}
]
],
"jobQueueStatus": {
"remainingCapacity": 100,
"pendingRequests": 0
}
}
]
```
Expand Down Expand Up @@ -401,6 +417,10 @@ curl http://localhost:8081/models/noop-customized/1.0?customized=true
"gpuUsage": "N/A"
}
],
"jobQueueStatus": {
"remainingCapacity": 100,
"pendingRequests": 0
},
"customizedMetadata": "{\n \"data1\": \"1\",\n \"data2\": \"2\"\n}"
}
]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ public class DescribeModelResponse {

private List<Worker> workers;
private Metrics metrics;
private JobQueueStatus jobQueueStatus;
private String customizedMetadata;

public DescribeModelResponse() {
Expand Down Expand Up @@ -150,6 +151,14 @@ public void setMetrics(Metrics metrics) {
this.metrics = metrics;
}

public JobQueueStatus getJobQueueStatus() {
return jobQueueStatus;
}

public void setJobQueueStatus(JobQueueStatus jobQueueStatus) {
this.jobQueueStatus = jobQueueStatus;
}

public void setCustomizedMetadata(byte[] customizedMetadata) {
this.customizedMetadata = new String(customizedMetadata, Charset.forName("UTF-8"));
}
Expand Down Expand Up @@ -257,4 +266,26 @@ public void setRequests(int requests) {
this.requests = requests;
}
}

public static final class JobQueueStatus {

private int remainingCapacity;
private int pendingRequests;

public int getRemainingCapacity() {
return remainingCapacity;
}

public void setRemainingCapacity(int remainingCapacity) {
this.remainingCapacity = remainingCapacity;
}

public int getPendingRequests() {
return pendingRequests;
}

public void setPendingRequests(int pendingRequests) {
this.pendingRequests = pendingRequests;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -433,7 +433,6 @@ private static Operation getDescribeModelOperation(boolean version) {
workerStatus.setEnumeration(status);
worker.addProperty("status", workerStatus, true);
workers.setItems(worker);

schema.addProperty("workers", workers, true);
Schema metrics = new Schema("object");
metrics.addProperty(
Expand All @@ -449,6 +448,16 @@ private static Operation getDescribeModelOperation(boolean version) {
new Schema("integer", "Number requests processed in last 10 minutes."),
true);
schema.addProperty("metrics", metrics, true);
Schema jobQueueStatus = new Schema("object");
jobQueueStatus.addProperty(
"remainingCapacity",
new Schema("integer", "Number of new requests that can be queued."),
true);
jobQueueStatus.addProperty(
"pendingRequests",
new Schema("integer", "Number of requests waiting in the queue."),
true);
schema.addProperty("jobQueueStatus", jobQueueStatus, true);

MediaType mediaType = new MediaType(HttpHeaderValues.APPLICATION_JSON.toString(), schema);
MediaType error = getErrorResponse();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -416,6 +416,12 @@ private static DescribeModelResponse createModelResponse(
resp.addWorker(workerId, startTime, isRunning, gpuId, memory, pid, gpuUsage);
}

DescribeModelResponse.JobQueueStatus jobQueueStatus =
new DescribeModelResponse.JobQueueStatus();
jobQueueStatus.setRemainingCapacity(model.getJobQueueRemainingCapacity());
jobQueueStatus.setPendingRequests(model.getPendingRequestsInJobQueue());
resp.setJobQueueStatus(jobQueueStatus);

return resp;
}

Expand Down
18 changes: 18 additions & 0 deletions frontend/server/src/main/java/org/pytorch/serve/wlm/Model.java
Original file line number Diff line number Diff line change
Expand Up @@ -402,4 +402,22 @@ public synchronized boolean getJobTickets() {
this.numJobTickets.decrementAndGet();
return true;
}

public int getJobQueueRemainingCapacity() {
LinkedBlockingDeque<Job> jobsQueue = jobsDb.get(DEFAULT_DATA_QUEUE);
if (jobsQueue != null) {
return jobsQueue.remainingCapacity();
}

return 0;
}

public int getPendingRequestsInJobQueue() {
LinkedBlockingDeque<Job> jobsQueue = jobsDb.get(DEFAULT_DATA_QUEUE);
if (jobsQueue != null) {
return jobsQueue.size();
}

return 0;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,24 @@ public void testDescribeSpecificModelVersion() throws InterruptedException {
@Test(
alwaysRun = true,
dependsOnMethods = {"testDescribeSpecificModelVersion"})
public void testDescribeModelJobQueueStatus() throws InterruptedException {
testLoadModelWithInitialWorkers("noop.mar", "noop_describe", "1.11");

Channel channel = TestUtils.getManagementChannel(configManager);
TestUtils.setResult(null);
TestUtils.setLatch(new CountDownLatch(1));
TestUtils.describeModel(channel, "noop_describe", "1.11", false);
TestUtils.getLatch().await();

DescribeModelResponse[] resp =
JsonUtils.GSON.fromJson(TestUtils.getResult(), DescribeModelResponse[].class);
Assert.assertEquals(resp[0].getJobQueueStatus().getRemainingCapacity(), 100);
Assert.assertEquals(resp[0].getJobQueueStatus().getPendingRequests(), 0);
}

@Test(
alwaysRun = true,
dependsOnMethods = {"testDescribeModelJobQueueStatus"})
public void testNoopVersionedPrediction() throws InterruptedException {
testPredictions("noopversioned", "OK", "1.11");
}
Expand Down
40 changes: 38 additions & 2 deletions frontend/server/src/test/resources/management_open_api.json
Original file line number Diff line number Diff line change
Expand Up @@ -479,7 +479,8 @@
"maxWorkers",
"status",
"workers",
"metrics"
"metrics",
"jobQueueStatus"
],
"properties": {
"modelName": {
Expand Down Expand Up @@ -570,6 +571,23 @@
"description": "Number requests processed in last 10 minutes."
}
}
},
"jobQueueStatus": {
"type": "object",
"required": [
"remainingCapacity",
"pendingRequests"
],
"properties": {
"remainingCapacity": {
"type": "integer",
"description": "Number of new requests that can be queued."
},
"pendingRequests": {
"type": "integer",
"description": "Number of requests waiting in the queue."
}
}
}
}
}
Expand Down Expand Up @@ -1049,7 +1067,8 @@
"maxWorkers",
"status",
"workers",
"metrics"
"metrics",
"jobQueueStatus"
],
"properties": {
"modelName": {
Expand Down Expand Up @@ -1140,6 +1159,23 @@
"description": "Number requests processed in last 10 minutes."
}
}
},
"jobQueueStatus": {
"type": "object",
"required": [
"remainingCapacity",
"pendingRequests"
],
"properties": {
"remainingCapacity": {
"type": "integer",
"description": "Number of new requests that can be queued."
},
"pendingRequests": {
"type": "integer",
"description": "Number of requests waiting in the queue."
}
}
}
}
}
Expand Down
40 changes: 38 additions & 2 deletions frontend/server/src/test/resources/model_management_api.json
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,8 @@
"maxWorkers",
"status",
"workers",
"metrics"
"metrics",
"jobQueueStatus"
],
"properties": {
"modelName": {
Expand Down Expand Up @@ -126,6 +127,23 @@
"description": "Number requests processed in last 10 minutes."
}
}
},
"jobQueueStatus": {
"type": "object",
"required": [
"remainingCapacity",
"pendingRequests"
],
"properties": {
"remainingCapacity": {
"type": "integer",
"description": "Number of new requests that can be queued."
},
"pendingRequests": {
"type": "integer",
"description": "Number of requests waiting in the queue."
}
}
}
}
}
Expand Down Expand Up @@ -605,7 +623,8 @@
"maxWorkers",
"status",
"workers",
"metrics"
"metrics",
"jobQueueStatus"
],
"properties": {
"modelName": {
Expand Down Expand Up @@ -696,6 +715,23 @@
"description": "Number requests processed in last 10 minutes."
}
}
},
"jobQueueStatus": {
"type": "object",
"required": [
"remainingCapacity",
"pendingRequests"
],
"properties": {
"remainingCapacity": {
"type": "integer",
"description": "Number of new requests that can be queued."
},
"pendingRequests": {
"type": "integer",
"description": "Number of requests waiting in the queue."
}
}
}
}
}
Expand Down

0 comments on commit 53ff6a5

Please sign in to comment.