From 31a097a2a96b6343d13c59298e3a2f3a9612e7ac Mon Sep 17 00:00:00 2001
From: ad-astra-video <99882368+ad-astra-video@users.noreply.github.com>
Date: Wed, 4 Dec 2024 07:15:58 -0600
Subject: [PATCH] fix: managed containers return after inference completed
 (#301)

* fix managed container return to start watchContainer monitoring to enabled returning container when inference is completed
---
 worker/docker.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/worker/docker.go b/worker/docker.go
index ae47ed67..9fdbbede 100644
--- a/worker/docker.go
+++ b/worker/docker.go
@@ -170,24 +170,28 @@ func (m *DockerManager) Stop(ctx context.Context) error {
 func (m *DockerManager) Borrow(ctx context.Context, pipeline, modelID string) (*RunnerContainer, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+	var rc *RunnerContainer
+	var err error
 
 	for _, runner := range m.containers {
 		if runner.Pipeline == pipeline && runner.ModelID == modelID {
-			delete(m.containers, runner.Name)
-			return runner, nil
+			rc = runner
+			break
 		}
 	}
 
 	// The container does not exist so try to create it
-	var err error
-	// TODO: Optimization flags for dynamically loaded (borrowed) containers are not currently supported due to startup delays.
-	rc, err := m.createContainer(ctx, pipeline, modelID, false, map[string]EnvValue{})
-	if err != nil {
-		return nil, err
+	if rc == nil {
+		// TODO: Optimization flags for dynamically loaded (borrowed) containers are not currently supported due to startup delays.
+		rc, err = m.createContainer(ctx, pipeline, modelID, false, map[string]EnvValue{})
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	// Remove container so it is unavailable until Return() is called
 	delete(m.containers, rc.Name)
+	// watch container to return when request completed
 	go m.watchContainer(rc, ctx)
 
 	return rc, nil