Make priority values categorical

frontend/server/src/main/java/org/pytorch/serve/grpcimpl/InferenceImpl.java

@@ -103,7 +103,7 @@ public void predictions(
 
         try {
             if (!ModelManager.getInstance().addJob(job)) {
-                int priority = job.getPriority();
+                String priority = job.getPriority().toString();
                 String responseMessage =
                         ApiUtils.getInferenceErrorResponseMessage(modelName, modelVersion, priority);
                 InternalServerException e = new InternalServerException(responseMessage);

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/InferenceRequestHandler.java

@@ -274,7 +274,7 @@ private static RequestInput parseRequest(
 
         CharSequence contentType = HttpUtil.getMimeType(req);
         for (Map.Entry<String, String> entry : req.headers().entries()) {
-            inputData.updateHeaders(entry.getKey(), entry.getValue());
+            inputData.updateHeaders(entry.getKey().toLowerCase(), entry.getValue());
         }
 
         if (HttpPostRequestDecoder.isMultipart(req)

frontend/server/src/main/java/org/pytorch/serve/job/GRPCJob.java

@@ -86,8 +86,8 @@ public void response(
                 "{}",
                 new Metric(
                         "RequestPriority",
-                        String.valueOf(this.getPriority()),
-                        "int",
+                        this.getPriority().toString(),
+                        "category",
                         ConfigManager.getInstance().getHostName(),
                         DIMENSION));
         } else if (this.getCmd() == WorkerCommands.DESCRIBE) {

frontend/server/src/main/java/org/pytorch/serve/job/Job.java

@@ -4,14 +4,15 @@
 import org.pytorch.serve.util.messages.RequestInput;
 import org.pytorch.serve.util.messages.WorkerCommands;
 import org.pytorch.serve.util.Prioritisable;
+import org.pytorch.serve.util.Priority;
 
 public abstract class Job implements Prioritisable {
 
     private String modelName;
     private String modelVersion;
     private WorkerCommands cmd; // Else its data msg or inf requests
     private RequestInput input;
-    private int priority;
+    private Priority priority;
     private long begin;
     private long scheduled;
 
@@ -22,20 +23,14 @@ public Job(String modelName, String version, WorkerCommands cmd, RequestInput in
         this.modelVersion = version;
         begin = System.nanoTime();
         scheduled = begin;
-
-        Map<String, String> headers = input.getHeaders();
-        if (headers.containsKey("X-TS-Priority")) {
-            this.priority = Integer.parseInt(headers.get("X-TS-Priority"));
-        } else {
-            this.priority = 0;
-        }
+        this.priority = Priority.valueOf(input.getHeaders().getOrDefault("x-ts-priority", "MAX").toUpperCase());
     }
 
-    public int getPriority() {
+    public Priority getPriority() {
         return this.priority;
     }
 
-    public void setPriority(int priority) {
+    public void setPriority(Priority priority) {
         this.priority = priority;
     }
 

frontend/server/src/main/java/org/pytorch/serve/job/RestJob.java

@@ -163,8 +163,8 @@ private void responseInference(
             "{}",
             new Metric(
                     "RequestPriority",
-                    String.valueOf(this.getPriority()),
-                    "int",
+                    this.getPriority().toString(),
+                    "category",
                     ConfigManager.getInstance().getHostName(),
                     DIMENSION));
     }

frontend/server/src/main/java/org/pytorch/serve/util/ApiUtils.java

@@ -394,22 +394,22 @@ public static RestJob addRESTInferenceJob(
             throws ModelNotFoundException, ModelVersionNotFoundException {
         RestJob job = new RestJob(ctx, modelName, version, WorkerCommands.PREDICT, input);
         if (!ModelManager.getInstance().addJob(job)) {
-            int priority = job.getPriority();
+            String priority = job.getPriority().toString();
             String responseMessage = getInferenceErrorResponseMessage(modelName, version, priority);
             throw new ServiceUnavailableException(responseMessage);
         }
         return job;
     }
 
     @SuppressWarnings("PMD")
-    public static String getInferenceErrorResponseMessage(String modelName, String modelVersion, int jobPriority) {
+    public static String getInferenceErrorResponseMessage(String modelName, String modelVersion, String jobPriority) {
         String responseMessage = "Model: " + modelName + "\n";
 
         if (modelVersion != null) {
             responseMessage += "Version: " + modelVersion + "\n";
         }
 
-        responseMessage += "Priority: " + String.valueOf(jobPriority) + "\n";
+        responseMessage += "Priority: " + jobPriority + "\n";
 
         responseMessage +=
                 "Reason: queue full";

frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java

@@ -64,7 +64,7 @@ public final class ConfigManager {
     private static final String TS_NUMBER_OF_NETTY_THREADS = "number_of_netty_threads";
     private static final String TS_NETTY_CLIENT_THREADS = "netty_client_threads";
     private static final String TS_JOB_QUEUE_SIZE = "job_queue_size";
-    private static final String TS_NUMBER_OF_PRIORITIES = "n_priorities";
+    private static final String TS_HIGH_PRIORITY_PROBABILITY = "high_prio_prob";
     private static final String TS_NUMBER_OF_GPU = "number_of_gpu";
     private static final String TS_METRICS_CONFIG = "metrics_config";
 
@@ -362,8 +362,13 @@ public int getJobQueueSize() {
         return getIntProperty(TS_JOB_QUEUE_SIZE, 100);
     }
 
-    public int getNumberOfPriorities() {
-        return getIntProperty(TS_NUMBER_OF_PRIORITIES, 1);
+    public float getHighPrioProb() throws IllegalArgumentException {
+        float highPrioProb = getFloatProperty(TS_HIGH_PRIORITY_PROBABILITY, 0.67f);
+        if (highPrioProb < 0.00f || highPrioProb > 1.00f){
+            throw new IllegalArgumentException("highPrioProb " + String.valueOf(highPrioProb) + 
+                " is not a valid probability!");
+        }
+        return highPrioProb;
     }
 
     public int getNumberOfGpu() {
@@ -685,6 +690,14 @@ private int getIntProperty(String key, int def) {
         return Integer.parseInt(value);
     }
 
+    private float getFloatProperty(String key, float def) {
+        String value = prop.getProperty(key);
+        if (value == null) {
+            return def;
+        }
+        return Float.parseFloat(value);
+    }
+
     public int getDefaultResponseTimeout() {
         return Integer.parseInt(prop.getProperty(TS_DEFAULT_RESPONSE_TIMEOUT, "120"));
     }

frontend/server/src/main/java/org/pytorch/serve/util/Prioritisable.java

@@ -2,7 +2,7 @@
 
 public interface Prioritisable {
 
-    public int getPriority();
-    public void setPriority(int priority);
+    public Priority getPriority();
+    public void setPriority(Priority priority);
 
 }

frontend/server/src/main/java/org/pytorch/serve/util/Priority.java

@@ -0,0 +1,5 @@
+package org.pytorch.serve.util;
+
+public enum Priority {
+    LOW, HIGH, MAX
+}

frontend/server/src/main/java/org/pytorch/serve/util/PriorityLinkedBlockingDeque.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.Enumeration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -19,72 +20,49 @@ public class PriorityLinkedBlockingDeque<T extends Prioritisable> {
     final ReentrantLock lock = new ReentrantLock();
     private final Condition notEmpty = lock.newCondition();
 
-    private int nPriorities;
-    private int sumPriorityWeights;
-    private int[] weightedPriorityMap;
-    private ConcurrentHashMap<Integer, LinkedBlockingDeque<T>> priorityDeques;
+    private final float highPrioProb;
+    private final ConcurrentHashMap<Priority, LinkedBlockingDeque<T>> priorityDeques;
 
-    public PriorityLinkedBlockingDeque(int nPriorities, int queueSize) {
+    public PriorityLinkedBlockingDeque(int queueSize, float highPrioProb) {
 
-        this.nPriorities = nPriorities;
-        this.priorityDeques = new ConcurrentHashMap<Integer, LinkedBlockingDeque<T>>();
+        this.highPrioProb = highPrioProb;
+        this.priorityDeques = new ConcurrentHashMap<Priority, LinkedBlockingDeque<T>>();
 
-        // note: the weight calculation logic may be adapted to the user's needs
-        // sum of 0 + 1 + 2 + ... + nPriorities - 1 via triangular sum
-        this.sumPriorityWeights = ((this.nPriorities - 1) * this.nPriorities) / 2;
-        this.weightedPriorityMap = new int[sumPriorityWeights];
-
-        // initialize priority deques and weight map
-        int keyStart = 0;
-        for (int priority = 0; priority < this.nPriorities; priority++) {
+        // initialize priority deques
+        for (Priority priority : Priority.values()) { 
             this.priorityDeques.put(priority, new LinkedBlockingDeque<T>(queueSize));
-            if (priority > 0) {
-                // priority weights are inverse priority values for now
-                int priorityWeight = this.nPriorities - priority;
-                for (int key = keyStart; key < keyStart + priorityWeight; key++) {
-                    this.weightedPriorityMap[key] = priority;
-                }
-                keyStart += priorityWeight;
-            }
         }
     }
 
     private LinkedBlockingDeque<T> getDequeForExtraction() {
 
-        // always select deque 0 first if non-empty
-        if (this.nPriorities == 1 || !this.priorityDeques.get(0).isEmpty()) {
-            return this.priorityDeques.get(0);
+        // always select deque max first if non-empty
+        if (!this.priorityDeques.get(Priority.MAX).isEmpty()) {
+            return this.priorityDeques.get(Priority.MAX);
         }
 
-        // sample according to weight map
-        int randInt = ThreadLocalRandom.current().nextInt(sumPriorityWeights);
-        int randPriority = this.weightedPriorityMap[randInt];
-        LinkedBlockingDeque<T> dequeForExtraction = this.priorityDeques.get(randPriority);
-
-        // if sampled deque is empty, scan deques according to priority
-        if (dequeForExtraction.isEmpty()) {
-            for (int priority = 1; priority < this.nPriorities; priority++) {
-                LinkedBlockingDeque<T> priorityDeque = this.priorityDeques.get(priority);
-                if (!priorityDeque.isEmpty()) {
-                    return priorityDeque;
-                }
+        boolean highNonEmpty = !this.priorityDeques.get(Priority.HIGH).isEmpty();
+
+        // if both high and low are non-empty, make random selection
+        if (highNonEmpty && !this.priorityDeques.get(Priority.LOW).isEmpty()) {
+            if (ThreadLocalRandom.current().nextFloat() < this.highPrioProb) {
+                return this.priorityDeques.get(Priority.HIGH);
+            } else {
+                return this.priorityDeques.get(Priority.LOW);
             }
+        // if only high is non-empty, return high
+        } else if (highNonEmpty) {
+            return this.priorityDeques.get(Priority.HIGH);
         }
-        return dequeForExtraction;
+
+        // if both empty or only low non-empty, return low
+        return this.priorityDeques.get(Priority.LOW);
+
     }
 
     private LinkedBlockingDeque<T> getDequeForInsertion(T p) {
-        int priority = p.getPriority();
+        Priority priority = p.getPriority();
         LinkedBlockingDeque<T> dequeForInsertion = this.priorityDeques.get(priority);
-
-        if (dequeForInsertion == null) {
-            logger.warn("Priority value "  + String.valueOf(priority) + " not valid, setting to highest valid priority value " 
-                + String.valueOf(this.nPriorities - 1) + ".");
-            int newPriority = this.nPriorities - 1;
-            p.setPriority(newPriority);
-            dequeForInsertion = this.priorityDeques.get(newPriority);
-        }
-
         return dequeForInsertion;
     }
 
@@ -98,10 +76,8 @@ private T unlinkFirst() {
     }
 
     public boolean isEmpty() {
-        Function<LinkedBlockingDeque<T>, Boolean> getIsEmpty = (LinkedBlockingDeque<T> deque) -> deque.isEmpty();
-        BiFunction<Boolean, Boolean, Boolean> logicalAnd = (Boolean a, Boolean b) -> a && b;
         // return true iff all deques are empty
-        return this.priorityDeques.reduceValues(Long.MAX_VALUE, getIsEmpty, logicalAnd);
+        return this.priorityDeques.reduceValues(Long.MAX_VALUE, LinkedBlockingDeque::isEmpty, Boolean::logicalAnd);
     }
 
     public boolean offer(T p) {

frontend/server/src/main/java/org/pytorch/serve/wlm/Model.java

@@ -40,7 +40,7 @@ public class Model {
     private ReentrantLock lock;
     private int responseTimeout;
     private int queueSize;
-    private int nPriorities;
+    private float highPrioProb;
     private ModelVersionName modelVersionName;
 
     private boolean isWorkflowModel;
@@ -51,15 +51,15 @@ public class Model {
     // Per worker thread job queue. This separates out the control queue from data queue
     private ConcurrentMap<String, PriorityLinkedBlockingDeque<Job>> jobsDb;
 
-    public Model(ModelArchive modelArchive, int queueSize, int nPriorities) {
+    public Model(ModelArchive modelArchive, int queueSize, float highPrioProb) {
         this.modelArchive = modelArchive;
         this.queueSize = queueSize;
-        this.nPriorities = nPriorities;
+        this.highPrioProb = highPrioProb;
         batchSize = 1;
         maxBatchDelay = 100;
         jobsDb = new ConcurrentHashMap<>();
         // Always have a queue for data
-        jobsDb.putIfAbsent(DEFAULT_DATA_QUEUE, new PriorityLinkedBlockingDeque<>(this.nPriorities, this.queueSize));
+        jobsDb.putIfAbsent(DEFAULT_DATA_QUEUE, new PriorityLinkedBlockingDeque<>(this.queueSize, this.highPrioProb));
         failedInfReqs = new AtomicInteger(0);
         lock = new ReentrantLock();
         modelVersionName =
@@ -156,7 +156,7 @@ public void setWorkflowModel(boolean workflowModel) {
     public void addJob(String threadId, Job job) {
         PriorityLinkedBlockingDeque<Job> blockingDeque = jobsDb.get(threadId);
         if (blockingDeque == null) {
-            blockingDeque = new PriorityLinkedBlockingDeque<>(this.nPriorities, this.queueSize);
+            blockingDeque = new PriorityLinkedBlockingDeque<>(this.queueSize, this.highPrioProb);
             jobsDb.put(threadId, blockingDeque);
         }
         blockingDeque.offer(job);

frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java

@@ -264,7 +264,7 @@ private Model createModel(
             int maxBatchDelay,
             int responseTimeout,
             boolean isWorkflowModel) {
-        Model model = new Model(archive, configManager.getJobQueueSize(), configManager.getNumberOfPriorities());
+        Model model = new Model(archive, configManager.getJobQueueSize(), configManager.getHighPrioProb());
 
         model.setBatchSize(
                 configManager.getJsonIntValue(
@@ -290,7 +290,7 @@ private Model createModel(
     }
 
     private Model createModel(ModelArchive archive, JsonObject modelInfo) {
-        Model model = new Model(archive, configManager.getJobQueueSize(), configManager.getNumberOfPriorities());
+        Model model = new Model(archive, configManager.getJobQueueSize(), configManager.getHighPrioProb());
         model.setModelState(modelInfo);
         model.setWorkflowModel(false);