Automatically remove dangling tasks from shim (#3036)

- Terminate and remove dangling tasks from shim
  when processing the instance. Tasks can become
  dangling because shim was unavailable when the
  job was in the `terminating` status.
- Adjust the shim API to return task statuses to
  avoid redundant termination. Dangling task
  termination and removal works with the old API
  too, but may result in redundant termination
  requests and warnings if the task is already
  terminated, but not removed. This is
  particularly noticeable if
  `DSTACK_SERVER_KEEP_SHIM_TASKS` is set.

Author: jvstme

runner/docs/shim.openapi.yaml

@@ -396,13 +396,21 @@ components:
       title: shim.api.TaskListResponse
       type: object
       properties:
-        ids:
+        tasks:
           type: array
           items:
-            $ref: "#/components/schemas/TaskID"
-          description: A list of all task IDs tracked by shim
+            type: object
+            properties:
+              id:
+                $ref: "#/components/schemas/TaskID"
+              status:
+                $ref: "#/components/schemas/TaskStatus"
+            required:
+              - id
+              - status
+          description: A list of all tasks tracked by shim, each with its ID and status
       required:
-        - ids
+        - tasks
       additionalProperties: false
 
     TaskInfoResponse:

runner/internal/shim/api/api_test.go

@@ -34,8 +34,8 @@ func (ds *DummyRunner) Remove(context.Context, string) error {
 	return nil
 }
 
-func (ds *DummyRunner) TaskIDs() []string {
-	return []string{}
+func (ds *DummyRunner) TaskList() []*shim.TaskListItem {
+	return []*shim.TaskListItem{}
 }
 
 func (ds *DummyRunner) TaskInfo(taskID string) shim.TaskInfo {

runner/internal/shim/api/handlers.go

@@ -36,7 +36,8 @@ func (s *ShimServer) InstanceHealthHandler(w http.ResponseWriter, r *http.Reques
 }
 
 func (s *ShimServer) TaskListHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) {
-	return &TaskListResponse{IDs: s.runner.TaskIDs()}, nil
+	tasks := s.runner.TaskList()
+	return &TaskListResponse{tasks}, nil
 }
 
 func (s *ShimServer) TaskInfoHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) {

runner/internal/shim/api/schemas.go

@@ -15,7 +15,7 @@ type InstanceHealthResponse struct {
 }
 
 type TaskListResponse struct {
-	IDs []string `json:"ids"`
+	Tasks []*shim.TaskListItem `json:"tasks"`
 }
 
 type TaskInfoResponse struct {

runner/internal/shim/api/server.go

@@ -19,7 +19,7 @@ type TaskRunner interface {
 	Remove(ctx context.Context, taskID string) error
 
 	Resources(context.Context) shim.Resources
-	TaskIDs() []string
+	TaskList() []*shim.TaskListItem
 	TaskInfo(taskID string) shim.TaskInfo
 }
 

runner/internal/shim/docker.go

@@ -216,8 +216,13 @@ func (d *DockerRunner) Resources(ctx context.Context) Resources {
 	}
 }
 
-func (d *DockerRunner) TaskIDs() []string {
-	return d.tasks.IDs()
+func (d *DockerRunner) TaskList() []*TaskListItem {
+	tasks := d.tasks.List()
+	result := make([]*TaskListItem, 0, len(tasks))
+	for _, task := range tasks {
+		result = append(result, &TaskListItem{ID: task.ID, Status: task.Status})
+	}
+	return result
 }
 
 func (d *DockerRunner) TaskInfo(taskID string) TaskInfo {

runner/internal/shim/models.go

@@ -104,6 +104,11 @@ type TaskConfig struct {
 	ContainerSshKeys []string `json:"container_ssh_keys"`
 }
 
+type TaskListItem struct {
+	ID     string     `json:"id"`
+	Status TaskStatus `json:"status"`
+}
+
 type TaskInfo struct {
 	ID                 string
 	Status             TaskStatus

runner/internal/shim/task.go

@@ -148,14 +148,15 @@ type TaskStorage struct {
 	mu    sync.RWMutex
 }
 
-func (ts *TaskStorage) IDs() []string {
+// Get a _copy_ of all tasks. To "commit" changes, use Update()
+func (ts *TaskStorage) List() []Task {
 	ts.mu.RLock()
 	defer ts.mu.RUnlock()
-	ids := make([]string, 0, len(ts.tasks))
-	for id := range ts.tasks {
-		ids = append(ids, id)
+	tasks := make([]Task, 0, len(ts.tasks))
+	for _, task := range ts.tasks {
+		tasks = append(tasks, task)
 	}
-	return ids
+	return tasks
 }
 
 // Get a _copy_ of the task. To "commit" changes, use Update()

src/dstack/_internal/server/background/tasks/process_instances.py

@@ -85,8 +85,10 @@
     get_instance_provisioning_data,
     get_instance_requirements,
     get_instance_ssh_private_keys,
+    remove_dangling_tasks_from_instance,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import is_divisible_into_blocks
 from dstack._internal.server.services.placement import (
     get_fleet_placement_group_models,
@@ -789,6 +791,7 @@ async def _check_instance(session: AsyncSession, instance: InstanceModel) -> Non
         ssh_private_keys,
         job_provisioning_data,
         None,
+        instance=instance,
         check_instance_health=check_instance_health,
     )
     if instance_check is False:
@@ -935,7 +938,7 @@ async def _wait_for_instance_provisioning_data(
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT], retries=1)
 def _check_instance_inner(
-    ports: Dict[int, int], *, check_instance_health: bool = False
+    ports: Dict[int, int], *, instance: InstanceModel, check_instance_health: bool = False
 ) -> InstanceCheck:
     instance_health_response: Optional[InstanceHealthResponse] = None
     shim_client = runner_client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
@@ -955,6 +958,10 @@ def _check_instance_inner(
         args = (method.__func__.__name__, e.__class__.__name__, e)
         logger.exception(template, *args)
         return InstanceCheck(reachable=False, message=template % args)
+    try:
+        remove_dangling_tasks_from_instance(shim_client, instance)
+    except Exception as e:
+        logger.exception("%s: error removing dangling tasks: %s", fmt(instance), e)
     return runner_client.healthcheck_response_to_instance_check(
         healthcheck_response, instance_health_response
     )

src/dstack/_internal/server/schemas/runner.py

@@ -159,6 +159,16 @@ class GPUDevice(CoreModel):
     path_in_container: str
 
 
+class TaskListItem(CoreModel):
+    id: str
+    status: TaskStatus
+
+
+class TaskListResponse(CoreModel):
+    ids: Optional[list[str]] = None  # returned by pre-0.19.26 shim
+    tasks: Optional[list[TaskListItem]] = None  # returned by 0.19.26+ shim
+
+
 class TaskInfoResponse(CoreModel):
     id: str
     status: TaskStatus