Merge pull request #1459 from transformerlab/fix/improve-live-status-remote-trap

Set Job to RUNNING when remote trap starts and introduce a lab_init live status

Author: deep1401

api/pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
   "soundfile==0.13.1",
   "tensorboardX==2.6.2.2",
   "timm==1.0.15",
-  "transformerlab==0.0.90",
+  "transformerlab==0.0.91",
   "transformerlab-inference==0.2.52",
   "transformers==4.57.1",
   "wandb==0.23.1",

api/transformerlab/routers/compute_provider.py

@@ -1862,8 +1862,8 @@ async def check_provider_job_status(
             "launch_progress": launch_progress,
         }
 
-    # Only check provider status for jobs in LAUNCHING state
-    if job_status != "LAUNCHING":
+    # Only check provider status for jobs that are still launching or running
+    if job_status not in ("LAUNCHING", "RUNNING"):
         return {
             "status": "success",
             "job_id": job_id,

lab-sdk/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab"
-version = "0.0.90"
+version = "0.0.91"
 description = "Python SDK for Transformer Lab"
 readme = "README.md"
 requires-python = ">=3.10"

lab-sdk/src/lab/lab_facade.py

@@ -117,6 +117,17 @@ def init(self, experiment_id: str | None = None, config: Optional[Dict[str, Any]
         # Update status to RUNNING for both cases
         _run_async(self._job.update_status("RUNNING"))
 
+        # Best-effort marker so UIs can distinguish jobs where lab has been
+        # explicitly initialized. This reuses the existing live_status field
+        # that remote_trap also writes to (started/finished/crashed). Depending
+        # on ordering, live_status may briefly be "lab_init" before or after
+        # "started" is set by tfl-remote-trap.
+        try:
+            _run_async(self._job.update_job_data_field("live_status", "lab_init"))
+        except Exception:
+            # Never let status marker failures break user code.
+            logger.debug("Failed to set live_status=lab_init on job", exc_info=True)
+
         # Check for wandb integration and capture URL if available
         self._detect_and_capture_wandb_url()
 

lab-sdk/src/lab/remote_trap.py

@@ -25,6 +25,18 @@ async def _set_live_status_async(job_id: str, status: str) -> None:
         return
 
 
+async def _set_status_async(job_id: str, status: str) -> None:
+    """Async helper to set the high-level job status."""
+    try:
+        job = await Job.get(job_id)
+        if job is None:
+            return
+        await job.update_status(status)
+    except Exception:
+        # This helper should never cause the wrapped command to fail.
+        return
+
+
 def _set_live_status(status: str) -> None:
     """Set live_status on the current remote job, if _TFL_JOB_ID is available."""
     job_id = os.environ.get("_TFL_JOB_ID")
@@ -47,6 +59,26 @@ def _set_live_status(status: str) -> None:
             return
 
 
+def _set_status(status: str) -> None:
+    """Set high-level job status for the current remote job, if _TFL_JOB_ID is available."""
+    job_id = os.environ.get("_TFL_JOB_ID")
+    if not job_id:
+        return
+
+    try:
+        asyncio.run(_set_status_async(job_id, status))
+    except RuntimeError:
+        # Fallback in case an event loop already exists.
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop.create_task(_set_status_async(job_id, status))
+            else:
+                loop.run_until_complete(_set_status_async(job_id, status))
+        except Exception:
+            return
+
+
 async def _write_provider_logs_async(job_id: str, logs_text: str) -> None:
     """
     Best-effort helper to write combined stdout/stderr logs to the job directory.
@@ -124,6 +156,7 @@ def main(argv: List[str] | None = None) -> int:
 
     # Mark job as started.
     _set_live_status("started")
+    _set_status("RUNNING")
 
     # Run the original command in the shell so it behaves exactly as submitted.
     # Capture stdout/stderr so we can save a copy to provider_logs.txt while still

src/renderer/components/Experiment/Tasks/JobProgress.tsx

@@ -122,6 +122,14 @@ export default function JobProgress({
     const liveStatus = job?.job_data?.live_status;
     if (!liveStatus) return null;
 
+    if (liveStatus === 'lab_init') {
+      return (
+        <Typography level="body-xs" color="neutral">
+          Lab instance initialized inside remote job.
+        </Typography>
+      );
+    }
+
     if (liveStatus === 'started') {
       return (
         <Typography level="body-xs" color="neutral">

src/renderer/components/Experiment/Tasks/Tasks.tsx

@@ -308,8 +308,12 @@ export default function Tasks({ subtype }: { subtype?: string }) {
         return false;
       }
 
-      // Always check LAUNCHING and WAITING jobs (for launch progress)
-      if (job.status === 'LAUNCHING' || job.status === 'WAITING') {
+      // Always check LAUNCHING, RUNNING, and WAITING jobs (for launch progress and live status)
+      if (
+        job.status === 'LAUNCHING' ||
+        job.status === 'RUNNING' ||
+        job.status === 'WAITING'
+      ) {
         return true;
       }
 
@@ -356,11 +360,15 @@ export default function Tasks({ subtype }: { subtype?: string }) {
       }
     };
 
-    // Check immediately and then every 2s when there are LAUNCHING/WAITING jobs (for progress), else 10s
-    const hasLaunching = jobsToCheck.some(
-      (j: any) => j.status === 'LAUNCHING' || j.status === 'WAITING',
+    // Check immediately and then every 2s when there are active jobs (LAUNCHING/RUNNING/WAITING),
+    // else every 10s (mainly for recent COMPLETE jobs to ensure quota is recorded).
+    const hasActiveRemoteJobs = jobsToCheck.some(
+      (j: any) =>
+        j.status === 'LAUNCHING' ||
+        j.status === 'RUNNING' ||
+        j.status === 'WAITING',
     );
-    const intervalMs = hasLaunching ? 2000 : 10000;
+    const intervalMs = hasActiveRemoteJobs ? 2000 : 10000;
     checkJobs();
     const interval = setInterval(checkJobs, intervalMs);