Merge remote-tracking branch 'upstream/main' into rhoai-3.2

Author: m-rafeeq

tests/trainer/resources/mnist.ipynb

@@ -505,9 +505,39 @@
    },
    "outputs": [],
    "source": [
-    "# Wait for the running status, then completion.\n",
-    "client.wait_for_job_status(name=job_name, status={\"Running\"})\n",
-    "client.wait_for_job_status(name=job_name, status={\"Complete\"})"
+    "# Wait for the running status, then wait for completion or failure\n",
+    "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n",
+    "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=900)\n",
+    "\n",
+    "# Get job details and logs\n",
+    "job = client.get_job(name=job_name)\n",
+    "pod_logs = client.get_job_logs(name=job_name, follow=False)\n",
+    "\n",
+    "# Collect all log lines from the generator into a list\n",
+    "logs = list(pod_logs)\n",
+    "log_text = \"\\n\".join(str(line) for line in logs)\n",
+    "\n",
+    "\n",
+    "print(f\"Training job final status: {job.status}\")\n",
+    "\n",
+    "# Check 1: Job status must not be \"Failed\"  \n",
+    "if job.status == \"Failed\":\n",
+    "    print(f\"ERROR: Training job '{job_name}' has Failed status\")\n",
+    "    print(\"Last 30 lines of logs:\")\n",
+    "    for line in logs[-30:]:\n",
+    "        print(line)\n",
+    "    raise RuntimeError(f\"Training job '{job_name}' failed\")\n",
+    "\n",
+    "# Check 2: Look for the training completion message in logs\n",
+    "# This is critical because the training script may catch exceptions and exit 0\n",
+    "if \"Training is finished\" not in log_text:\n",
+    "    print(f\"ERROR: Training completion message not found in logs\")\n",
+    "    print(\"Last 50 lines of logs:\")\n",
+    "    for line in logs[-50:]:\n",
+    "        print(line)\n",
+    "    raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n",
+    "\n",
+    "print(f\"✓ Training job '{job_name}' completed successfully\")"
    ]
   },
   {