Wait for a running task instead of a running job

Author: jfmlima

jupyterhub_nomad_spawner/nomad/nomad_model.py

@@ -1571,6 +1571,20 @@ class AllocatedResources(BaseModel):
     Shared: Optional[AllocatedSharedResources] = None
     Tasks: Optional[Dict[str, AllocatedTaskResources]] = None
 
+class TaskEvent(BaseModel):
+    Type: str
+    Time: int
+    DisplayMessage: str
+    Details: Dict[str, Any]
+    FailsTask: bool
+    DriverMessage: str
+
+class TaskState(BaseModel):
+    State: str
+    Failed: bool
+    StartedAt: Optional[str]
+    FinishedAt: Optional[str]
+    Events: List[TaskEvent]
 
 class AllocationListStub(BaseModel):
     AllocatedResources: Optional[AllocatedResources] = None

jupyterhub_nomad_spawner/nomad/nomad_service.py

@@ -11,6 +11,7 @@
     CSIVolumeCapability,
     CSIVolumeCreateRequest,
     JobsParseRequest,
+    TaskState,
 )
 
 
@@ -125,6 +126,43 @@ async def job_status(self, job_id) -> str:
 
         job_detail = response.json()
         return job_detail.get("Status", "")
+    
+    async def task_status(self, job_name: str) -> str:
+        """Get detailed task status from most recent allocation"""
+        allocs = await self.client.get(f"/v1/job/{job_name}/allocations")
+        if not allocs:
+            return "pending"
+
+        allocs = allocs.json()
+        latest_alloc = max(allocs, key=lambda x: x["CreateTime"])
+        if not latest_alloc:
+            return "pending"
+        
+        task_states = latest_alloc.get("TaskStates", {}) or {}
+        task_states = {name: TaskState(**state) for name, state in task_states.items()}
+        
+        if not task_states:
+            return "pending"
+
+        for task in task_states.values():
+            if task.State == "dead" and task.Failed:
+                return "dead"
+            if task.State != "running":
+                return self._get_task_state_from_event(task)
+
+        return "running"
+    
+    def _get_task_state_from_event(self, task: TaskState) -> str:
+        """Determine task state from latest event"""
+        events = task.Events
+        if not events:
+            return "pending"
+        
+        latest_event = events[-1]
+        if latest_event.Type in ["Driver", "Task Setup"]:
+            return "starting"
+        return "pending"
+
 
     async def job_allocations(self, job_id) -> list[dict[str, Any]]:
         response = await self.client.get(f"/v1/job/{job_id}/allocations")

jupyterhub_nomad_spawner/spawner.py

@@ -544,16 +544,24 @@ def _get_csi_extra_parameters(self) -> Optional[Dict]:
     async def _ensure_running(self, nomad_service: NomadService):
         while True:
             try:
-                status = await nomad_service.job_status(self.job_name)
-            except Exception:
-                self.log.exception("Failed to get job status")
-            if status == "running":
-                break
-            elif status == "dead":
-                raise Exception(f"Job (name={self.job_name}) is dead already")
-            else:
-                self.log.info("Waiting for %s...", self.job_name)
-                await asyncio.sleep(5)
+                job_status = await nomad_service.job_status(self.job_name)
+                if job_status == "dead":
+                    raise Exception(f"Job (name={self.job_name}) is dead")
+                elif job_status == "running":
+                    task_status = await nomad_service.task_status(self.job_name)
+                    if task_status == "running":
+                        break
+                    elif task_status == "dead":
+                        raise Exception(f"Task for job (name={self.job_name}) is dead")
+                    else:
+                        self.log.info("Task for %s is %s, waiting...", self.job_name, task_status)
+                else:
+                    self.log.info("Job %s is %s, waiting...", self.job_name, job_status)
+            except Exception as e:
+                self.log.exception("Failed to get job/task status")
+                raise e
+            
+            await asyncio.sleep(5)
 
     @retry(wait=wait_fixed(3), stop=stop_after_attempt(5))
     async def address_and_port_from_consul(