pass jobid in start script, when loading tensorboard

transformerlab · Jan 3, 2024 · df92979 · df92979
1 parent 1fd7ba0
commit df92979
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/src/renderer/components/Experiment/Train/TensorboardModal.tsx b/src/renderer/components/Experiment/Train/TensorboardModal.tsx
@@ -4,13 +4,19 @@ import * as chatAPI from 'renderer/lib/transformerlab-api-sdk';
 
 const fetcher = (url) => fetch(url).then((res) => res.json());
 
-export default function TensorboardModal({ open, setOpen }) {
+export default function TensorboardModal({
+  currentTensorboard,
+  setCurrentTensorboard,
+}) {
   const [iframeReady, setIframeReady] = useState(false);
 
   useEffect(() => {
-    if (open) {
+    if (currentTensorboard !== -1) {
       console.log('starting tensorboard');
-      fetcher(chatAPI.API_URL() + 'train/tensorboard/start').then((res) => {
+      var job_id = currentTensorboard;
+      fetcher(
+        chatAPI.API_URL() + 'train/tensorboard/start?job_id=' + job_id
+      ).then((res) => {
         console.log(res);
       });
 
@@ -22,20 +28,23 @@ export default function TensorboardModal({ open, setOpen }) {
       }, 3000);
     }
 
-    if (!open) {
+    if (currentTensorboard == -1) {
       console.log('stopping tensorboard');
       fetcher(chatAPI.API_URL() + 'train/tensorboard/stop').then((res) => {
         console.log(res);
       });
     }
-  }, [open]);
+  }, [currentTensorboard]);
 
   var currentServerURL = window.TransformerLab.API_URL;
   // If there is a port number, remove it:
   currentServerURL = currentServerURL.replace(/:[0-9]+\/$/, '');
 
   return (
-    <Modal open={open} onClose={() => setOpen(false)}>
+    <Modal
+      open={currentTensorboard !== -1}
+      onClose={() => setCurrentTensorboard(-1)}
+    >
       <ModalDialog
         sx={{
           height: '80vh',

diff --git a/src/renderer/components/Experiment/Train/TrainLoRA.tsx b/src/renderer/components/Experiment/Train/TrainLoRA.tsx
@@ -76,7 +76,8 @@ const fetcher = (url) => fetch(url).then((res) => res.json());
 
 export default function TrainLoRA({ experimentInfo }) {
   const [open, setOpen] = useState(false);
-  const [openTensorboardModal, setOpenTensorboardModal] = useState(false);
+  const [currentTensorboardForModal, setCurrentTensorboardForModal] =
+    useState(-1);
   const [viewOutputFromJob, setViewOutputFromJob] = useState(-1);
 
   const { data, error, isLoading, mutate } = useSWR(
@@ -108,8 +109,8 @@ export default function TrainLoRA({ experimentInfo }) {
         experimentInfo={experimentInfo}
       />
       <TensorboardModal
-        open={openTensorboardModal}
-        setOpen={setOpenTensorboardModal}
+        currentTensorboard={currentTensorboardForModal}
+        setCurrentTensorboard={setCurrentTensorboardForModal}
       />
       <ViewOutputModal
         jobId={viewOutputFromJob}
@@ -284,7 +285,7 @@ export default function TrainLoRA({ experimentInfo }) {
                       <Button
                         size="sm"
                         onClick={() => {
-                          setOpenTensorboardModal(true);
+                          setCurrentTensorboardForModal(job?.id);
                         }}
                         startDecorator={<LineChartIcon />}
                       >