Merge pull request transformerlab#711 from transformerlab/add/checkpointing-ui

Add/checkpointing UI

Author: aliasaria

src/renderer/components/Experiment/Train/TrainLoRA.tsx

@@ -34,6 +34,7 @@ import {
   StopCircleIcon,
   Trash2Icon,
   UploadIcon,
+  WaypointsIcon,
 } from 'lucide-react';
 
 import dayjs from 'dayjs';
@@ -51,6 +52,7 @@ import TensorboardModal from './TensorboardModal';
 import ViewOutputModal from './ViewOutputModal';
 import ViewEvalImagesModal from './ViewEvalImagesModal';
 import { useExperimentInfo } from 'renderer/lib/ExperimentInfoContext';
+import ViewCheckpointsModal from './ViewCheckpointsModal';
 dayjs.extend(relativeTime);
 var duration = require('dayjs/plugin/duration');
 dayjs.extend(duration);
@@ -106,6 +108,7 @@ export default function TrainLoRA({}) {
   const [viewEvalImagesFromJob, setViewEvalImagesFromJob] = useState(-1);
   const [templateID, setTemplateID] = useState('-1');
   const [currentPlugin, setCurrentPlugin] = useState('');
+  const [viewCheckpointsFromJob, setViewCheckpointsFromJob] = useState(-1);
 
   const { data, error, isLoading, mutate } = useSWR(
     chatAPI.Endpoints.Tasks.ListByTypeInExperiment('TRAIN', experimentInfo?.id),
@@ -190,12 +193,16 @@ export default function TrainLoRA({}) {
         sweeps={viewOutputFromSweepJob}
         setsweepJob={setViewOutputFromSweepJob}
       />
-
       <ViewEvalImagesModal
         open={viewEvalImagesFromJob !== -1}
         onClose={() => setViewEvalImagesFromJob(-1)}
         jobId={viewEvalImagesFromJob}
       />
+      <ViewCheckpointsModal
+        open={viewCheckpointsFromJob !== -1}
+        onClose={() => setViewCheckpointsFromJob(-1)}
+        jobId={viewCheckpointsFromJob}
+      />
       <Sheet
         sx={{
           display: 'flex',
@@ -523,6 +530,18 @@ export default function TrainLoRA({}) {
                               Sweep Output
                             </Button>
                           )}
+                          {job?.job_data?.checkpoints && (
+                            <Button
+                              size="sm"
+                              variant="plain"
+                              onClick={() => {
+                                setViewCheckpointsFromJob(job?.id);
+                              }}
+                              startDecorator={<WaypointsIcon />}
+                            >
+                              Checkpoints
+                            </Button>
+                          )}
                           <IconButton variant="plain">
                             <Trash2Icon
                               onClick={async () => {

src/renderer/components/Experiment/Train/ViewCheckpointsModal.tsx

@@ -0,0 +1,110 @@
+import {
+  Modal,
+  ModalDialog,
+  Typography,
+  ModalClose,
+  Table,
+  Button,
+  Box,
+} from '@mui/joy';
+import { PlayIcon } from 'lucide-react';
+import { useAPI } from 'renderer/lib/transformerlab-api-sdk';
+import { formatBytes } from 'renderer/lib/utils';
+
+export default function ViewCheckpointsModal({ open, onClose, jobId }) {
+  const { data, isLoading: checkpointsLoading } = useAPI(
+    'jobs',
+    ['getCheckpoints'],
+    { jobId },
+  );
+
+  const handleRestartFromCheckpoint = (checkpoint) => {
+    // TODO: Implement restart functionality
+    console.log('Restarting from checkpoint:', checkpoint);
+  };
+
+  let noCheckpoints = false;
+
+  if (!checkpointsLoading && data?.checkpoints?.length === 0) {
+    noCheckpoints = true;
+  }
+
+  return (
+    <Modal open={open} onClose={() => onClose()}>
+      <ModalDialog sx={{ minWidth: '80%' }}>
+        <ModalClose />
+
+        {noCheckpoints ? (
+          <Typography level="body-md" sx={{ textAlign: 'center', py: 4 }}>
+            No checkpoints were saved in this job.
+          </Typography>
+        ) : (
+          <>
+            <Typography level="h4" component="h2">
+              Checkpoints for Job {jobId}
+            </Typography>
+
+            {!checkpointsLoading && data && (
+              <Box sx={{ mb: 2 }}>
+                <Typography level="body-md">
+                  <strong>Model:</strong> {data.model_name}
+                </Typography>
+                <Typography level="body-md">
+                  <strong>Adaptor:</strong> {data.adaptor_name}
+                </Typography>
+              </Box>
+            )}
+
+            {checkpointsLoading ? (
+              <Typography level="body-md">Loading checkpoints...</Typography>
+            ) : (
+              <Box sx={{ maxHeight: 400, overflow: 'auto' }}>
+                <Table>
+                  <thead>
+                    <tr>
+                      <th width="50px">#</th>
+                      <th>Checkpoint</th>
+                      <th>Date</th>
+                      <th width="100px">Size</th>
+                      <th style={{ textAlign: 'right' }}>&nbsp;</th>
+                    </tr>
+                  </thead>
+                  <tbody>
+                    {data?.checkpoints?.map((checkpoint, index) => (
+                      <tr key={index}>
+                        <td>
+                          <Typography level="body-sm">
+                            {data?.checkpoints?.length - index}.
+                          </Typography>
+                        </td>
+                        <td>
+                          <Typography level="title-sm">
+                            {checkpoint.filename}
+                          </Typography>
+                        </td>
+                        <td>{new Date(checkpoint.date).toLocaleString()}</td>
+                        <td>{formatBytes(checkpoint.size)}</td>
+                        <td style={{ textAlign: 'right' }}>
+                          {/* <Button
+                            size="sm"
+                            variant="outlined"
+                            onClick={() =>
+                              handleRestartFromCheckpoint(checkpoint.filename)
+                            }
+                            startDecorator={<PlayIcon />}
+                          >
+                            Restart training from here
+                          </Button> */}
+                        </td>
+                      </tr>
+                    ))}
+                  </tbody>
+                </Table>
+              </Box>
+            )}
+          </>
+        )}
+      </ModalDialog>
+    </Modal>
+  );
+}

src/renderer/lib/api-client/allEndpoints.json

@@ -87,6 +87,10 @@
     "stop": {
       "method": "GET",
       "path": "jobs/{jobId}/stop"
+    },
+    "getCheckpoints": {
+      "method": "GET",
+      "path": "jobs/{jobId}/checkpoints"
     }
   },
   "datasets": {
@@ -234,6 +238,5 @@
       "method": "GET",
       "path": "train/job/{job_id}/sweep_results"
     }
-
   }
 }