ai/worker: Watch containers continuously even when not in use (#3416)

* ai/worker: Auto-restart containers (only 3x if not warm)

* ai/worker: Remove manual restart of containers

* ai/worker: Run watchContainer only once per container

We should keep watching containers even when they are returned

* ai/worker: Use context.Background instead of nil

* ai/worker: Fix tests

* ai/worker: Increase tests for new logic

* doc: Add a little documentation on ai-worker container lifefycle

* ai/worker: Remove AutoRemove flag

Cannot have it together with a restart policy

* Revert "ai/worker: Remove manual restart of containers"

This reverts commit b346d37.

* ai/docker: Add a little comment on borrowctx usage

Author: victorges

ai/worker/container.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"log/slog"
+	"sync"
 	"time"
 
 	"github.com/deepmap/oapi-codegen/v2/pkg/securityprovider"
@@ -21,6 +22,9 @@ type RunnerContainer struct {
 	Name     string
 	Client   *ClientWithResponses
 	Hardware *HardwareInformation
+
+	BorrowCtx context.Context
+	sync.RWMutex
 }
 
 type RunnerEndpoint struct {

ai/worker/doc.go

@@ -13,7 +13,7 @@ Listens for inference requests from the Livepeer AI subnet and routes them to th
 
 - Docker Manager (./docker.go):
 
-Manages AI runner containers.
+Manages AI runner containers. For a state diagram showing the lifecycle of a container, see the /doc/worker.md file.
 
 [AI runner containers]: https://github.com/livepeer/ai-runner
 */

ai/worker/docker.go

@@ -153,14 +153,11 @@ func (m *DockerManager) Warm(ctx context.Context, pipeline string, modelID strin
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	rc, err := m.createContainer(ctx, pipeline, modelID, true, optimizationFlags)
+	_, err := m.createContainer(ctx, pipeline, modelID, true, optimizationFlags)
 	if err != nil {
 		return err
 	}
 
-	// Watch with a background context since we're not borrowing the container.
-	go m.watchContainer(rc, context.Background())
-
 	return nil
 }
 
@@ -200,19 +197,25 @@ func (m *DockerManager) Borrow(ctx context.Context, pipeline, modelID string) (*
 		}
 	}
 
-	// Remove container so it is unavailable until Return() is called
+	// Remove container and set the BorrowCtx so it is unavailable until returnContainer() is called by watchContainer()
 	delete(m.containers, rc.Name)
-	// watch container to return when request completed
-	go m.watchContainer(rc, ctx)
+	rc.Lock()
+	rc.BorrowCtx = ctx
+	rc.Unlock()
 
 	return rc, nil
 }
 
 // returnContainer returns a container to the pool so it can be reused. It is called automatically by watchContainer
-// when the context used to borrow the container is done.
+// when the BorrowCtx of the container is done or the container is IDLE.
 func (m *DockerManager) returnContainer(rc *RunnerContainer) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+
+	rc.Lock()
+	rc.BorrowCtx = nil
+	rc.Unlock()
+
 	m.containers[rc.Name] = rc
 }
 
@@ -348,6 +351,14 @@ func (m *DockerManager) createContainer(ctx context.Context, pipeline string, mo
 		gpuOpts.Set("device=" + gpu)
 	}
 
+	restartPolicy := container.RestartPolicy{
+		Name:              "on-failure",
+		MaximumRetryCount: 3,
+	}
+	if keepWarm {
+		restartPolicy = container.RestartPolicy{Name: "always"}
+	}
+
 	hostConfig := &container.HostConfig{
 		Resources: container.Resources{
 			DeviceRequests: gpuOpts.Value(),
@@ -367,7 +378,7 @@ func (m *DockerManager) createContainer(ctx context.Context, pipeline string, mo
 				},
 			},
 		},
-		AutoRemove: true,
+		RestartPolicy: restartPolicy,
 	}
 
 	resp, err := m.dockerClient.ContainerCreate(ctx, containerConfig, hostConfig, nil, nil, containerName)
@@ -421,6 +432,8 @@ func (m *DockerManager) createContainer(ctx context.Context, pipeline string, mo
 	m.containers[containerName] = rc
 	m.gpuContainers[gpu] = containerName
 
+	go m.watchContainer(rc)
+
 	return rc, nil
 }
 
@@ -479,7 +492,7 @@ func (m *DockerManager) destroyContainer(rc *RunnerContainer, locked bool) error
 // watchContainer monitors a container's running state and automatically cleans
 // up the internal state when the container stops. It will also monitor the
 // borrowCtx to return the container to the pool when it is done.
-func (m *DockerManager) watchContainer(rc *RunnerContainer, borrowCtx context.Context) {
+func (m *DockerManager) watchContainer(rc *RunnerContainer) {
 	defer func() {
 		if r := recover(); r != nil {
 			slog.Error("Panic in container watch routine",
@@ -509,10 +522,19 @@ func (m *DockerManager) watchContainer(rc *RunnerContainer, borrowCtx context.Co
 			return
 		}
 
+		rc.RLock()
+		// The BorrowCtx is set when the container has been borrowed for a request/stream. If it is not set (nil) it means
+		// that it's not currently borrowed, so we don't need to wait for it to be done (hence using the background context).
+		borrowCtx := rc.BorrowCtx
+		if borrowCtx == nil {
+			borrowCtx = context.Background()
+		}
+		rc.RUnlock()
+
 		select {
 		case <-borrowCtx.Done():
 			m.returnContainer(rc)
-			return
+			continue
 		case <-ticker.C:
 			ctx, cancel := context.WithTimeout(context.Background(), containerWatchInterval)
 			health, err := rc.Client.HealthWithResponse(ctx)
@@ -543,7 +565,7 @@ func (m *DockerManager) watchContainer(rc *RunnerContainer, borrowCtx context.Co
 				if time.Since(startTime) > pipelineStartGracePeriod {
 					slog.Info("Container is idle, returning to pool", slog.String("container", rc.Name))
 					m.returnContainer(rc)
-					return
+					continue
 				}
 				fallthrough
 			case OK:

ai/worker/docker_test.go

@@ -741,17 +741,43 @@ func TestDockerManager_watchContainer(t *testing.T) {
 		defer mockServer.AssertExpectations(t)
 		defer mockDockerClient.AssertNotCalled(t, "ContainerRemove", mock.Anything, rc.Name, mock.Anything)
 
+		mockServer.On("ServeHTTP", "GET", "/health", mock.Anything).
+			Return(200, "application/json", `{"status":"OK"}`)
+
 		borrowCtx, cancel := context.WithCancel(context.Background())
+		rc.BorrowCtx = borrowCtx
 
-		go dockerManager.watchContainer(rc, borrowCtx)
+		watchingCtx, stopWatching := context.WithCancel(context.Background())
+		go func() {
+			dockerManager.watchContainer(rc)
+			stopWatching()
+		}()
 		cancel()                          // Cancel the context.
 		time.Sleep(50 * time.Millisecond) // Ensure the ticker triggers.
 
 		// Verify that the container was returned.
 		_, exists := dockerManager.containers[rc.Name]
 		require.True(t, exists)
 
-		mockServer.AssertNotCalled(t, "ServeHTTP", "GET", "/health", mock.Anything)
+		// Watch should keep running
+		require.Nil(t, watchingCtx.Err())
+
+		// Borrow the container again
+		borrowCtx, cancel = context.WithCancel(context.Background())
+		delete(dockerManager.containers, rc.Name)
+		rc.Lock()
+		rc.BorrowCtx = borrowCtx
+		rc.Unlock()
+
+		cancel()                          // Cancel the borrow context
+		time.Sleep(50 * time.Millisecond) // Ensure the ticker triggers.
+
+		// Verify that the container was returned.
+		_, exists = dockerManager.containers[rc.Name]
+		require.True(t, exists)
+
+		// Watch should still keep running
+		require.Nil(t, watchingCtx.Err())
 	})
 
 	notHealthyTestCases := []struct {
@@ -818,8 +844,6 @@ func TestDockerManager_watchContainer(t *testing.T) {
 			defer mockDockerClient.AssertExpectations(t)
 			defer mockServer.AssertExpectations(t)
 
-			borrowCtx := context.Background()
-
 			tt.mockServerSetup(mockServer)
 
 			// Mock destroyContainer to verify it is called.
@@ -829,7 +853,7 @@ func TestDockerManager_watchContainer(t *testing.T) {
 			done := make(chan struct{})
 			go func() {
 				defer close(done)
-				dockerManager.watchContainer(rc, borrowCtx)
+				dockerManager.watchContainer(rc)
 			}()
 			select {
 			case <-done:
@@ -852,14 +876,12 @@ func TestDockerManager_watchContainer(t *testing.T) {
 		defer mockDockerClient.AssertExpectations(t)
 		defer mockServer.AssertExpectations(t)
 
-		borrowCtx := context.Background()
-
 		// Must fail twice before the watch routine gives up on it
 		mockServer.On("ServeHTTP", "GET", "/health", mock.Anything).
 			Return(200, "application/json", `{"status":"ERROR"}`).
 			Times(4) // 4 calls during the grace period (first call only after 10ms)
 
-		go dockerManager.watchContainer(rc, borrowCtx)
+		go dockerManager.watchContainer(rc)
 		time.Sleep(40 * time.Millisecond) // Almost the entire grace period
 
 		// Make sure container wasn't destroyed yet
@@ -885,8 +907,6 @@ func TestDockerManager_watchContainer(t *testing.T) {
 		defer mockServer.AssertExpectations(t)
 		defer mockDockerClient.AssertNotCalled(t, "ContainerRemove", mock.Anything, rc.Name, mock.Anything)
 
-		borrowCtx := context.Background()
-
 		// Schedule:
 		// - Return IDLE for the first 3 times during grace period (should not return)
 		// - Then OK for the next 5 times (stayin' alive)
@@ -908,7 +928,7 @@ func TestDockerManager_watchContainer(t *testing.T) {
 				time.Sleep(dur)
 			}
 		}
-		go dockerManager.watchContainer(rc, borrowCtx)
+		go dockerManager.watchContainer(rc)
 		sleepUntil(30 * time.Millisecond) // Almost the entire grace period
 
 		// Verify that the container was not returned yet.

doc/assets/ai-runner-container-lifecycle.jpg

@@ -0,0 +1,31 @@
+# AI Worker
+
+The AI Worker node manages [`ai-runner`](https://github.com/livepeer/ai-runner) containers running on the host system.
+These containers are started, monitored and stopped dynamically depending on the usage.
+
+This diagram describes the lifecycle of a container:
+
+![ai-runner container lifecycle](./assets/ai-runner-container-lifecycle.jpg)
+
+Source: [Miro Board](https://miro.com/app/board/uXjVIZ0vO4k=/?share_link_id=987855784886)
+
+It can also be described by the following mermaid chart, but the rendered version is more confusing:
+```
+stateDiagram-v2
+    direction TB
+    [*] --> OFFLINE
+    OFFLINE --> IDLE: Warm()->createCont()
+    OFFLINE --> BORROWED: Borrow(ctx)->createCont()
+    state RUNNING {
+        [*] --> IDLE
+        IDLE --> BORROWED: Borrow(ctx)
+        BORROWED --> IDLE: BorrowCtx.Done()
+    }
+    hc: GET /health
+    RUNNING --> hc
+    state healthcheck <<choice>>
+    hc --> healthcheck
+    healthcheck --> OFFLINE: if error x2
+    healthcheck --> RUNNING: if state=OK
+    healthcheck --> IDLE: if state=IDLE
+```