feat(worker): add HasCapability method (#73)

This commit adds a HasCapability method to the worker that can be used to
check if the worker has enough capacity to take a request.

worker/docker.go

@@ -130,6 +130,24 @@ func (m *DockerManager) Return(rc *RunnerContainer) {
 	m.containers[dockerContainerName(rc.Pipeline, rc.ModelID)] = rc
 }
 
+// HasCapacity checks if an unused managed container exists or if a GPU is available for a new container.
+func (m *DockerManager) HasCapacity(ctx context.Context, pipeline, modelID string) bool {
+	containerName := dockerContainerName(pipeline, modelID)
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// Check if unused managed container exists for the requested model.
+	_, ok := m.containers[containerName]
+	if ok {
+		return true
+	}
+
+	// Check for available GPU to allocate for a new container for the requested model.
+	_, err := m.allocGPU(ctx)
+	return err == nil
+}
+
 func (m *DockerManager) createContainer(ctx context.Context, pipeline string, modelID string, keepWarm bool, optimizationFlags OptimizationFlags) (*RunnerContainer, error) {
 	containerName := dockerContainerName(pipeline, modelID)
 

worker/worker.go

@@ -239,6 +239,22 @@ func (w *Worker) Stop(ctx context.Context) error {
 	return nil
 }
 
+// HasCapacity returns true if the worker has capacity for the given pipeline and model ID.
+func (w *Worker) HasCapacity(pipeline, modelID string) bool {
+	managedCapacity := w.manager.HasCapacity(context.Background(), pipeline, modelID)
+	if managedCapacity {
+		return true
+	}
+
+	// Check if we have capacity for external containers.
+	name := dockerContainerName(pipeline, modelID)
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	_, ok := w.externalContainers[name]
+
+	return ok
+}
+
 func (w *Worker) borrowContainer(ctx context.Context, pipeline, modelID string) (*RunnerContainer, error) {
 	w.mu.Lock()