Reload defunct runners

In case a runner becomes defunct, e.g. as a result of a backend crash it
would be neat to be able to reload it. So, if the loader finds runner,
have it check if the runner is still alive, and create a new one if the
runner is defunct.

Signed-off-by: Piotr Stankiewicz <piotr.stankiewicz@docker.com>

Author: Piotr Stankiewicz

pkg/inference/scheduling/loader.go

@@ -154,15 +154,22 @@ func (l *loader) broadcast() {
 }
 
 // evict evicts all unused runners from the loader. If idleOnly is true, then
-// only those unused runners which are considered "idle" (based on usage
-// timestamp) are evicted. The caller must hold the loader lock. It returns the
-// number of remaining runners.
+// only those unused, but functioning, runners which are considered "idle" (based
+// on usage timestamp) are evicted. Defunct (e.g. crashed) runners will be evicted
+// regardless of whether they are considered "idle". The caller must hold the loader
+// lock. It returns the number of remaining runners.
 func (l *loader) evict(idleOnly bool) int {
 	now := time.Now()
 	for r, slot := range l.runners {
 		unused := l.references[slot] == 0
 		idle := unused && now.Sub(l.timestamps[slot]) > runnerIdleTimeout
-		if unused && (!idleOnly || idle) {
+		defunct := false
+		select {
+		case <-l.slots[slot].done:
+			defunct = true
+		default:
+		}
+		if unused && (!idleOnly || idle || defunct) {
 			l.log.Infof("Evicting %s backend runner with model %s in %s mode",
 				r.backend, r.model, r.mode,
 			)
@@ -364,6 +371,8 @@ func (l *loader) load(ctx context.Context, backendName, model string, mode infer
 
 	// Loop until we can satisfy the request or an error occurs.
 	for {
+		slot := -1
+
 		// If loads are disabled, then there's nothing we can do.
 		if !l.loadsEnabled {
 			return nil, errLoadsDisabled
@@ -372,9 +381,15 @@ func (l *loader) load(ctx context.Context, backendName, model string, mode infer
 		// See if we can satisfy the request with an existing runner.
 		existing, ok := l.runners[runnerKey{backendName, model, mode}]
 		if ok {
-			l.references[existing] += 1
-			l.timestamps[existing] = time.Time{}
-			return l.slots[existing], nil
+			select {
+			case <-l.slots[existing].done:
+				l.log.Warnf("%s runner for %s is defunct. Waiting for it to be evicted.", backendName, model)
+				goto WaitForChange
+			default:
+				l.references[existing] += 1
+				l.timestamps[existing] = time.Time{}
+				return l.slots[existing], nil
+			}
 		}
 
 		// If there's not sufficient memory or all slots are full, then try
@@ -384,7 +399,6 @@ func (l *loader) load(ctx context.Context, backendName, model string, mode infer
 		}
 
 		// If there's sufficient memory and a free slot, then find the slot.
-		slot := -1
 		if memory <= l.availableMemory && len(l.runners) < len(l.slots) {
 			for s, runner := range l.slots {
 				if runner == nil {
@@ -432,6 +446,7 @@ func (l *loader) load(ctx context.Context, backendName, model string, mode infer
 		// Wait for something to change. Note that we always re-lock with
 		// context.Background() because we need to ensure we hold the lock by
 		// the time we return.
+	WaitForChange:
 		l.unlock()
 		select {
 		case <-ctx.Done():
@@ -455,13 +470,19 @@ func (l *loader) release(runner *runner) {
 	// Decrement the runner's reference count.
 	l.references[slot] -= 1
 
-	// If the runner's reference count is now zero, then record now as its idle
-	// start time and signal the idle checker.
+	// If the runner's reference count is now zero, then check if it is still
+	// active, and record now as its idle start time and signal the idle
+	// checker.
 	if l.references[slot] == 0 {
-		l.timestamps[slot] = time.Now()
 		select {
-		case l.idleCheck <- struct{}{}:
+		case <-runner.done:
+			l.evictRunner(runner.backend.Name(), runner.model)
 		default:
+			l.timestamps[slot] = time.Now()
+			select {
+			case l.idleCheck <- struct{}{}:
+			default:
+			}
 		}
 	}