Update app.py

app.py

@@ -32,8 +32,11 @@ def start_train():
     # Handles CUDA OOM errors.
     os.system(f"export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True")
     os.system("echo 'Okay, trying training.'")
+    status = gr.Textbox(value='TRAINING')
     os.system(f"cd pytorch-image-models; ./train.sh 4 --dataset hfds/datacomp/imagenet-1k-random-0.0-frac-1over4 --log-wandb --wandb-project ImageNetTraining0.0-frac-1over4 --experiment ImageNetTraining0.0-frac-1over4 --model seresnet34 --sched cosine --epochs 150 --warmup-epochs 5 --lr 0.4 --reprob 0.5 --remode pixel --batch-size 256 --amp -j 4")
     os.system("echo 'Done'.")
+    API.add_space_variable['STATUS'] = "IDLE"
+    status.value = "IDLE"
     os.system("ls")
     # Upload output to repository
     os.system("echo 'trying to upload...'")
@@ -44,8 +47,12 @@ def run():
     with gr.Blocks() as app:
         gr.Markdown(f"Randomization: 0.0")
         gr.Markdown(f"Subset: frac-1over4")
-        start = gr.Button("Start")
-        start.click(start_train)
+        space_variables = API.get_space_variables(repo_id=experiment_repo)
+        if 'STATUS' not in space_variables or space_variables['STATUS'].value != "TRAINING":
+            start = gr.Button("Start")
+            start.click(start_train)
+        else:
+            gr.Markdown(f"Status: Training.")
     app.launch(server_name="0.0.0.0", server_port=7860)
     
 if __name__ == '__main__':