From 4d14a93b88fb92097b3c19550c795f8eaa1eb0b1 Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Mon, 13 May 2019 16:31:26 +0800
Subject: [PATCH] add GPU share support

---
 src/AllocatorFIFO.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/AllocatorFIFO.go b/src/AllocatorFIFO.go
index b382d73..abdf1b8 100644
--- a/src/AllocatorFIFO.go
+++ b/src/AllocatorFIFO.go
@@ -3,7 +3,7 @@ package main
 import (
 	"sync"
 	"time"
-	)
+)
 
 type AllocatorFIFO struct {
 	history    []*Job
@@ -85,7 +85,7 @@ func (allocator *AllocatorFIFO) requestResource(task Task) NodeStatus {
 	for id, node := range pool.nodes {
 		var available []GPUStatus
 		for _, status := range node.Status {
-			if status.MemoryAllocated == 0 {
+			if status.MemoryTotal-status.MemoryAllocated >= task.MemoryGPU {
 				available = append(available, status)
 			}
 		}
@@ -97,7 +97,8 @@ func (allocator *AllocatorFIFO) requestResource(task Task) NodeStatus {
 			for i := range res.Status {
 				for j := range node.Status {
 					if res.Status[i].UUID == node.Status[j].UUID {
-						node.Status[j].MemoryAllocated = task.MemoryGPU
+						node.Status[j].MemoryAllocated += task.MemoryGPU
+						res.Status[i].MemoryTotal = task.MemoryGPU
 					}
 				}
 			}
@@ -114,7 +115,7 @@ func (allocator *AllocatorFIFO) returnResource(agent NodeStatus) {
 	for _, gpu := range agent.Status {
 		for j := range nodes.Status {
 			if gpu.UUID == nodes.Status[j].UUID {
-				nodes.Status[j].MemoryAllocated = 0
+				nodes.Status[j].MemoryAllocated -= gpu.MemoryTotal
 			}
 		}
 	}