Pause Space

app_training.py

@@ -105,8 +105,8 @@ def create_training_demo(trainer: Trainer,
                         choices=[_.value for _ in UploadTarget],
                         value=UploadTarget.MODEL_LIBRARY.value)
 
-        remove_gpu_after_training = gr.Checkbox(
-            label='Remove GPU after training',
+        pause_space_after_training = gr.Checkbox(
+            label='Pause this Space after training',
             value=False,
             interactive=bool(os.getenv('SPACE_ID')),
             visible=False)
@@ -143,7 +143,7 @@ def create_training_demo(trainer: Trainer,
                              use_private_repo,
                              delete_existing_repo,
                              upload_to,
-                             remove_gpu_after_training,
+                             pause_space_after_training,
                              hf_token,
                          ])
     return demo

trainer.py

@@ -60,7 +60,7 @@ class Trainer:
         use_private_repo: bool,
         delete_existing_repo: bool,
         upload_to: str,
-        remove_gpu_after_training: bool,
+        pause_space_after_training: bool,
         hf_token: str,
     ) -> None:
         if not torch.cuda.is_available():
@@ -140,9 +140,7 @@ class Trainer:
             with open(self.log_file, 'a') as f:
                 f.write(upload_message)
 
-        if remove_gpu_after_training:
-            space_id = os.getenv('SPACE_ID')
-            if space_id:
+        if pause_space_after_training:
+            if space_id := os.getenv('SPACE_ID'):
                 api = HfApi(token=os.getenv('HF_TOKEN') or hf_token)
-                api.request_space_hardware(repo_id=space_id,
-                                           hardware='cpu-basic')
+                api.pause_space(repo_id=space_id)