diff --git a/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java b/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java index 6a7499959d..03532d0244 100644 --- a/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java +++ b/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java @@ -68,6 +68,7 @@ public final class ConfigManager { private static final String TS_NUMBER_OF_GPU = "number_of_gpu"; private static final String TS_MIN_FREE_GPU_MEMORY = "min_free_gpu_memory"; private static final String TS_MAX_SHARE_GPU_FAILURES = "max_share_gpu_failures"; + private static final String TS_OVERRIDE_GPU_ID = "override_gpu_id"; private static final String TS_METRICS_CONFIG = "metrics_config"; private static final String TS_METRICS_MODE = "metrics_mode"; private static final String TS_DISABLE_SYSTEM_METRICS = "disable_system_metrics"; @@ -380,8 +381,8 @@ public int getMinFreeGpuMemory() { return getIntProperty(TS_MIN_FREE_GPU_MEMORY, 4096); } - public float getMaxShareGpuFailures() { - return getFloatProperty(TS_MAX_SHARE_GPU_FAILURES, 0.90f); + public int getOverrideGpuId() { + return getIntProperty(TS_OVERRIDE_GPU_ID, -1); } public String getMetricsConfigPath() { @@ -685,7 +686,9 @@ public String dumpConfigurations() { + "\nWorkflow Store: " + (getWorkflowStore() == null ? "N/A" : getWorkflowStore()) + "\nModel config: " - + prop.getProperty(MODEL_CONFIG, "N/A"); + + prop.getProperty(MODEL_CONFIG, "N/A") + + "\nOverride GPU ID: " + + getOverrideGpuId(); } public boolean useNativeIo() { diff --git a/frontend/server/src/main/java/org/pytorch/serve/util/GPUManager.java b/frontend/server/src/main/java/org/pytorch/serve/util/GPUManager.java index 410a34bdcc..482d05240c 100644 --- a/frontend/server/src/main/java/org/pytorch/serve/util/GPUManager.java +++ b/frontend/server/src/main/java/org/pytorch/serve/util/GPUManager.java @@ -18,24 +18,21 @@ public final class GPUManager { private static final Logger logger = LoggerFactory.getLogger(GPUManager.class); - private static final int nFailureHistory = 100; private static GPUManager instance; private final int nGPUs; private final int minFreeMemory; - private final float maxShareFailures; + private final int overrideGpuId; private AtomicInteger[] freeMemory; private HashMap workerIds; - private ArrayDeque gpuFailureHistory; - private GPUManager(int nGPUs, int minFreeMemory, float maxShareFailures) { + private GPUManager(int nGPUs, int minFreeMemory, int overrideGpuId) { this.nGPUs = nGPUs; this.minFreeMemory = minFreeMemory; - this.maxShareFailures = maxShareFailures; + this.overrideGpuId = overrideGpuId; - this.gpuFailureHistory = new ArrayDeque<> (); this.workerIds = new HashMap<> (); if (nGPUs > 0) { @@ -84,8 +81,8 @@ private int queryNvidiaSmiFreeMemory(int gpuId) { public static synchronized void init(ConfigManager configManager) { int nGPUs = configManager.getNumberOfGpu(); int minFreeMemory = configManager.getMinFreeGpuMemory(); - float maxShareFailures = configManager.getMaxShareGpuFailures(); - instance = new GPUManager(nGPUs, minFreeMemory, maxShareFailures); + int override_gpu_id = configManager.getOverrideGpuId(); + instance = new GPUManager(nGPUs, minFreeMemory, override_gpu_id); } public static synchronized GPUManager getInstance() { @@ -95,44 +92,24 @@ public static synchronized GPUManager getInstance() { public synchronized int getGPU(String workerId) { // return -1 if there are no gpus if (this.nGPUs == 0) { + logger.error("No eligible GPUs available, falling back to CPU"); return -1; } - int failedGpuId; - // if the worker was previously assigned to a GPU and now requests a new one, it has likely failed - // add failed gpu id to failure history, removing old entries to make space if necessary - if (this.workerIds.containsKey(workerId)) { - failedGpuId = this.workerIds.get(workerId); - while (this.gpuFailureHistory.size() > nFailureHistory - 1) { - this.gpuFailureHistory.removeFirst(); - } - this.gpuFailureHistory.addLast(failedGpuId); + // return override if given + if (this.overrideGpuId > -1) { + return this.overrideGpuId; } // get free memory per GPU for (int i = 0; i < this.nGPUs; i++) { this.freeMemory[i].set(queryNvidiaSmiFreeMemory(i)); } - // get failures for share calculation - int[] nFailures = new int[this.nGPUs]; - for (Iterator iter = this.gpuFailureHistory.iterator(); iter.hasNext();) { - failedGpuId = iter.next(); - nFailures[failedGpuId]++; - } + // get free memory for all eligible GPUs HashMap eligibleIdFreeMems = new HashMap (); for (int i = 0; i < this.nGPUs; i++) { // check that free memory is available and exceeds minimum if (this.freeMemory[i].intValue() > this.minFreeMemory) { - if (this.gpuFailureHistory.size() > 1) { - // check that share of failures is smaller than maximum - float shareFailures = (float) nFailures[i] / (float) this.gpuFailureHistory.size(); - if (shareFailures < this.maxShareFailures) { - eligibleIdFreeMems.put(i, this.freeMemory[i].intValue()); - } else { - logger.warn("GPU ID {} deemed ineligible since it accounts for {} out of failures {}", i, nFailures[i], this.maxShareFailures); - } - } else { - eligibleIdFreeMems.put(i, this.freeMemory[i].intValue()); - } + eligibleIdFreeMems.put(i, this.freeMemory[i].intValue()); logger.info("eligibleIdFreeMems[{}] {}", i, this.freeMemory[i].intValue()); } @@ -161,8 +138,6 @@ public synchronized int getGPU(String workerId) { int freeMem = entry.getValue(); cumProb += (float) freeMem / (float) eligibleIdFreeMemSum; cumProbIds.put(cumProb, i); - // TODO Simon: This log should maybe have been removed during rebase - logger.info("cumProbIds[{}] {} because of freeMem {}", cumProb, i, freeMem); } // make random selection float randFloat = ThreadLocalRandom.current().nextFloat(); @@ -172,7 +147,6 @@ public synchronized int getGPU(String workerId) { } logger.info("Assigning gpuId " + gpuId + " with free memory " + eligibleIdFreeMems.get(gpuId) + - " with number of failures " + nFailures[gpuId] + " to workerId " + workerId); return gpuId; }