fix(trainer): change trainer function

api/internal_services/background_worker.py

@@ -4,6 +4,7 @@ from threading import Thread
 import grpc
 
 from api.internal_services.logger import logger
+from api.internal_services.trainer import start_training
 from api.models.Job import JobType
 from api.protos.trainer import trainer_pb2_grpc, trainer_pb2
 
@@ -45,19 +46,7 @@ def process_trainer():
         if job is None:
             break
 
-        with grpc.insecure_channel(job.job_data['server_url']) as channel:
-            stub = trainer_pb2_grpc.TrainerStub(channel)
-            request = trainer_pb2.TrainingInput(
-                training_data=[],
-                fondation_model_id=job.job_data['fondation_model_id'],
-                finetuned_repo_name=job.job_data['finetuned_repo_name'],
-                huggingface_token=job.job_data['huggingface_token'],
-            )
-            logger.debug(request)
-            responses = stub.StartTraining(request)
-            for response in responses:
-                logger.debug(f"gRPC message : {response.status}")
-        logger.debug(f"fin de la connexion gRPC")
+        start_training(job)
 
         logger.info(f"Ending of the job {job.job_id}")
         trainer_queue.task_done()

api/internal_services/trainer.py

+import grpc
+from api.internal_services.logger import logger
+from api.protos.trainer import trainer_pb2_grpc, trainer_pb2
+
+
+def start_training(job):
+    with grpc.insecure_channel(job.job_data['server_url']) as channel:
+        stub = trainer_pb2_grpc.TrainerStub(channel)
+        request = trainer_pb2.TrainingInput(
+            training_data=[],
+            fondation_model_id=job.job_data['fondation_model_id'],
+            finetuned_repo_name=job.job_data['finetuned_repo_name'],
+            huggingface_token=job.job_data['huggingface_token'],
+        )
+        logger.debug(request)
+        responses = stub.StartTraining(request)
+        for response in responses:
+            logger.debug(f"gRPC message : {response.status}")
+    logger.debug(f"fin de la connexion gRPC")
\ No newline at end of file

microservices/trainer/trainer.py

@@ -30,12 +30,15 @@ class TrainerServicer(trainer_pb2_grpc.TrainerServicer):
             if not is_busy:
                 is_busy = True
                 print(f"incoming request : {request}")
-                training_process(
-                    training_data=request.training_data,
-                    fondation_model_id=request.fondation_model_id,
-                    finetuned_repo_name=request.finetuned_repo_name,
-                    huggingface_token=request.huggingface_token,
-                )
+                try:
+                    training_process(
+                        training_data=request.training_data,
+                        fondation_model_id=request.fondation_model_id,
+                        finetuned_repo_name=request.finetuned_repo_name,
+                        huggingface_token=request.huggingface_token,
+                    )
+                except Exception as e:
+                    print(f"Error : {e}")
                 is_busy = False
             else:
                 print(f"gRPC server is already busy")