From 9a902c067fc8205ee1ad86343cb0eead3344eea7 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 21 Mar 2025 11:45:36 -0600 Subject: [PATCH] Show the console log on the job overview in the ui. --- extensions_built_in/sd_trainer/UITrainer.py | 2 +- toolkit/timer.py | 15 +++-- ui/src/app/api/jobs/[jobID]/log/route.ts | 35 +++++++++++ ui/src/app/api/jobs/[jobID]/start/route.ts | 1 + ui/src/components/JobOverview.tsx | 64 ++++++++++++++++++++- ui/src/hooks/useJobLog.tsx | 60 +++++++++++++++++++ 6 files changed, 169 insertions(+), 8 deletions(-) create mode 100644 ui/src/app/api/jobs/[jobID]/log/route.ts create mode 100644 ui/src/hooks/useJobLog.tsx diff --git a/extensions_built_in/sd_trainer/UITrainer.py b/extensions_built_in/sd_trainer/UITrainer.py index 124f80a8..df940ced 100644 --- a/extensions_built_in/sd_trainer/UITrainer.py +++ b/extensions_built_in/sd_trainer/UITrainer.py @@ -149,7 +149,7 @@ class UITrainer(SDTrainer): try: await asyncio.gather(*self._async_tasks) except Exception as e: - print(f"Error waiting for async operations: {e}") + pass finally: # Clear the task list after completion self._async_tasks.clear() diff --git a/toolkit/timer.py b/toolkit/timer.py index 3592ba5e..e849ba5f 100644 --- a/toolkit/timer.py +++ b/toolkit/timer.py @@ -1,6 +1,10 @@ import time from collections import OrderedDict, deque +import sys +import os +# check if is ui process will have IS_AI_TOOLKIT_UI in env +is_ui = os.environ.get("IS_AI_TOOLKIT_UI", "0") == "1" class Timer: def __init__(self, name='Timer', max_buffer=10): @@ -39,18 +43,21 @@ class Timer: self._after_print_hooks.append(hook) def print(self): - print(f"\nTimer '{self.name}':") + if not is_ui: + print(f"\nTimer '{self.name}':") timing_dict = {} # sort by longest at top for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True): avg_time = sum(timings) / len(timings) - print(f" - {avg_time:.4f}s avg - {timer_name}, num = {len(timings)}") + + if not is_ui: + print(f" - {avg_time:.4f}s avg - {timer_name}, num = {len(timings)}") timing_dict[timer_name] = avg_time for hook in self._after_print_hooks: hook(timing_dict) - - print('') + if not is_ui: + print('') def reset(self): self.timers.clear() diff --git a/ui/src/app/api/jobs/[jobID]/log/route.ts b/ui/src/app/api/jobs/[jobID]/log/route.ts new file mode 100644 index 00000000..10ccbdaa --- /dev/null +++ b/ui/src/app/api/jobs/[jobID]/log/route.ts @@ -0,0 +1,35 @@ +import { NextRequest, NextResponse } from 'next/server'; +import { PrismaClient } from '@prisma/client'; +import path from 'path'; +import fs from 'fs'; +import { getTrainingFolder } from '@/server/settings'; + +const prisma = new PrismaClient(); + +export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) { + const { jobID } = await params; + + const job = await prisma.job.findUnique({ + where: { id: jobID }, + }); + + if (!job) { + return NextResponse.json({ error: 'Job not found' }, { status: 404 }); + } + + const trainingFolder = await getTrainingFolder(); + const jobFolder = path.join(trainingFolder, job.name); + const logPath = path.join(jobFolder, 'log.txt'); + + if (!fs.existsSync(logPath)) { + return NextResponse.json({ log: '' }); + } + let log = ''; + try { + log = fs.readFileSync(logPath, 'utf-8'); + } catch (error) { + console.error('Error reading log file:', error); + log = 'Error reading log file'; + } + return NextResponse.json({ log: log }); +} diff --git a/ui/src/app/api/jobs/[jobID]/start/route.ts b/ui/src/app/api/jobs/[jobID]/start/route.ts index ba1feaf9..e26c1e49 100644 --- a/ui/src/app/api/jobs/[jobID]/start/route.ts +++ b/ui/src/app/api/jobs/[jobID]/start/route.ts @@ -96,6 +96,7 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s const additionalEnv: any = { AITK_JOB_ID: jobID, CUDA_VISIBLE_DEVICES: `${job.gpu_ids}`, + IS_AI_TOOLKIT_UI: '1' }; // HF_TOKEN diff --git a/ui/src/components/JobOverview.tsx b/ui/src/components/JobOverview.tsx index a29d4fb0..9f304ca3 100644 --- a/ui/src/components/JobOverview.tsx +++ b/ui/src/components/JobOverview.tsx @@ -4,7 +4,8 @@ import GPUWidget from '@/components/GPUWidget'; import FilesWidget from '@/components/FilesWidget'; import { getTotalSteps } from '@/utils/jobs'; import { Cpu, HardDrive, Info, Gauge } from 'lucide-react'; -import { useMemo } from 'react'; +import { useEffect, useMemo, useRef, useState } from 'react'; +import useJobLog from '@/hooks/useJobLog'; interface JobOverviewProps { job: Job; @@ -12,12 +13,50 @@ interface JobOverviewProps { export default function JobOverview({ job }: JobOverviewProps) { const gpuIds = useMemo(() => job.gpu_ids.split(',').map(id => parseInt(id)), [job.gpu_ids]); + const { log, setLog, status: statusLog, refresh: refreshLog } = useJobLog(job.id, 2000); + const logRef = useRef(null); + // Track whether we should auto-scroll to bottom + const [isScrolledToBottom, setIsScrolledToBottom] = useState(true); const { gpuList, isGPUInfoLoaded } = useGPUInfo(gpuIds, 5000); const totalSteps = getTotalSteps(job); const progress = (job.step / totalSteps) * 100; const isStopping = job.stop && job.status === 'running'; + const logLines: string[] = useMemo(() => { + // split at line breaks on \n or \r\n but not \r + let splits: string[] = log.split(/\n|\r\n/); + + splits = splits.map(line => { + return line.split(/\r/).pop(); + }) as string[]; + + // only return last 100 lines max + const maxLines = 1000; + if (splits.length > maxLines) { + splits = splits.slice(splits.length - maxLines); + } + + return splits; + }, [log]); + + // Handle scroll events to determine if user has scrolled away from bottom + const handleScroll = () => { + if (logRef.current) { + const { scrollTop, scrollHeight, clientHeight } = logRef.current; + // Consider "at bottom" if within 10 pixels of the bottom + const isAtBottom = scrollHeight - scrollTop - clientHeight < 10; + setIsScrolledToBottom(isAtBottom); + } + }; + + // Auto-scroll to bottom only if we were already at the bottom + useEffect(() => { + if (logRef.current && isScrolledToBottom) { + logRef.current.scrollTop = logRef.current.scrollHeight; + } + }, [log, isScrolledToBottom]); + const getStatusColor = (status: string) => { switch (status.toLowerCase()) { case 'running': @@ -43,7 +82,7 @@ export default function JobOverview({ job }: JobOverviewProps) { return (
{/* Job Information Panel */} -
+

{job.info} @@ -51,7 +90,7 @@ export default function JobOverview({ job }: JobOverviewProps) { {job.status}

-
+
{/* Progress Bar */}
@@ -91,6 +130,25 @@ export default function JobOverview({ job }: JobOverviewProps) {
+ + {/* Log - Now using flex-grow to fill remaining space */} +
+
+ {statusLog === 'loading' && 'Loading log...'} + {statusLog === 'error' && 'Error loading log'} + {['success', 'refreshing'].includes(statusLog) && ( +
+ {logLines.map((line, index) => { + return 
{line}
; + })} +
+ )} +
+
diff --git a/ui/src/hooks/useJobLog.tsx b/ui/src/hooks/useJobLog.tsx new file mode 100644 index 00000000..e7770418 --- /dev/null +++ b/ui/src/hooks/useJobLog.tsx @@ -0,0 +1,60 @@ +'use client'; + +import { useEffect, useState, useRef } from 'react'; +import { apiClient } from '@/utils/api'; + +interface FileObject { + path: string; + size: number; +} + +const clean = (text: string): string => { + // remove \x1B[A\x1B[A + text = text.replace(/\x1B\[A/g, ''); + return text; +}; + +export default function useJobLog(jobID: string, reloadInterval: null | number = null) { + const [log, setLog] = useState(''); + const didInitialLoadRef = useRef(false); + const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error' | 'refreshing'>('idle'); + + const refresh = () => { + let loadStatus: 'loading' | 'refreshing' = 'loading'; + if (didInitialLoadRef.current) { + loadStatus = 'refreshing'; + } + setStatus(loadStatus); + apiClient + .get(`/api/jobs/${jobID}/log`) + .then(res => res.data) + .then(data => { + if (data.log) { + let cleanLog = clean(data.log); + setLog(cleanLog); + } + setStatus('success'); + didInitialLoadRef.current = true; + }) + .catch(error => { + console.error('Error fetching log:', error); + setStatus('error'); + }); + }; + + useEffect(() => { + refresh(); + + if (reloadInterval) { + const interval = setInterval(() => { + refresh(); + }, reloadInterval); + + return () => { + clearInterval(interval); + }; + } + }, [jobID]); + + return { log, setLog, status, refresh }; +}