feature: trainer logic

adds a exception type that triggers a node restart

learning_loop_node/trainer/exceptions.py

@@ -1,2 +1,12 @@
 class CriticalError(Exception):
-    pass
+    '''
+    CriticalError is raised when the training cannot be continued.
+    In this case the trainer jumps to the TrainerState.ReadyForCleanup and tries to upload the latest model.
+    '''
+
+
+class NodeNeedsRestartError(Exception):
+    '''
+    NodeNeedsRestartError is raised when the node needs to be restarted.
+    This is e.g. the case when the GPU is not available anymore.
+    '''

learning_loop_node/trainer/trainer_logic_generic.py

@@ -14,7 +14,7 @@
                             TrainingOut, TrainingStateData)
 from ..helpers.misc import create_project_folder, delete_all_training_folders, generate_training, is_valid_uuid4
 from .downloader import TrainingsDownloader
-from .exceptions import CriticalError
+from .exceptions import CriticalError, NodeNeedsRestartError
 from .io_helpers import ActiveTrainingIO, EnvironmentVars, LastTrainingIO
 
 if TYPE_CHECKING:
@@ -294,6 +294,9 @@ async def _perform_state(self, error_key: str, state_during: TrainerState, state
             logger.error('CriticalError in %s - Exception: %s', state_during, e)
             self.errors.set(error_key, str(e))
             self.training.training_state = TrainerState.ReadyForCleanup
+        except NodeNeedsRestartError:
+            logger.error('Node Restart Requested')
+            sys.exit(0)
         except Exception as e:
             self.errors.set(error_key, str(e))
             logger.exception('Error in %s - Exception: %s', state_during, e)