From 817c3a9fb5ce96b7a2fd8cb700b3dfc22a5279e9 Mon Sep 17 00:00:00 2001
From: Yondon Fu <yondon.fu@gmail.com>
Date: Thu, 25 Jan 2024 19:38:06 +0000
Subject: [PATCH] worker: Stop() removes all containers

---
 worker/worker.go | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/worker/worker.go b/worker/worker.go
index 4183fe51e..765b80b92 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -185,21 +185,27 @@ func (w *Worker) Warm(ctx context.Context, containerName, modelID string) error
 	return err
 }
 
-func (w *Worker) Stop(ctx context.Context, containerName string) error {
-	c, ok := w.containers[containerName]
-	if !ok {
-		return fmt.Errorf("container %v is not running", containerName)
-	}
+func (w *Worker) Stop(ctx context.Context) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
 
-	// TODO: Handle if container fails to stop or be removed
-	delete(w.containers, containerName)
+	var stopContainerWg sync.WaitGroup
+	for name, rc := range w.containers {
+		stopContainerWg.Add(1)
+		go func(containerID string) {
+			defer stopContainerWg.Done()
+			if err := dockerRemoveContainer(ctx, w.dockerClient, containerID); err != nil {
+				slog.Error("Error removing container", slog.String("name", name), slog.String("id", containerID))
+			}
+		}(rc.ID)
 
-	if err := w.dockerClient.ContainerStop(ctx, c.ID, container.StopOptions{}); err != nil {
-		return err
+		w.gpuLoad[rc.GPU] -= 1
+		delete(w.containers, name)
 	}
 
-	// Is there a reason to not remove the container?
-	return w.dockerClient.ContainerRemove(ctx, c.ID, types.ContainerRemoveOptions{})
+	stopContainerWg.Wait()
+
+	return nil
 }
 
 func (w *Worker) getWarmContainer(ctx context.Context, containerName string, modelID string) (*RunnerContainer, error) {