Make all debug prints conditional on get_debug() - suppress debug output when --debug not specified

vidai/worker_training.py

@@ -27,11 +27,12 @@ import shutil
 import json
 import time
 from .comm import SocketCommunicator, Message
-from .config import get_comm_type, get_backend_worker_port
+from .config import get_comm_type, get_backend_worker_port, get_debug
 
 def train_model(train_path, output_model, description):
     """Perform training."""
-    print(f"DEBUG: Starting training with videotrain for output_model {output_model}")
+    if get_debug():
+        print(f"DEBUG: Starting training with videotrain for output_model {output_model}")
     desc_file = os.path.join(train_path, "description.txt")
     with open(desc_file, "w") as f:
         f.write(description)
@@ -39,7 +40,8 @@ def train_model(train_path, output_model, description):
     # Assume videotrain is available
     cmd = ["python", "videotrain", train_path, "--output_dir", output_model]
     result = subprocess.run(cmd, capture_output=True, text=True)
-    print(f"DEBUG: Training subprocess completed with returncode {result.returncode}")
+    if get_debug():
+        print(f"DEBUG: Training subprocess completed with returncode {result.returncode}")
     if result.returncode == 0:
         return "Training completed!"
     else:
@@ -47,7 +49,8 @@ def train_model(train_path, output_model, description):
 
 def worker_process(backend_type: str):
     """Main worker process."""
-    print(f"Starting Training Worker for {backend_type}...")
+    if get_debug():
+        print(f"Starting Training Worker for {backend_type}...")
 
     # Workers use TCP for interprocess communication
     comm = SocketCommunicator(host='127.0.0.1', port=get_backend_worker_port(), comm_type='tcp')
@@ -60,10 +63,11 @@ def worker_process(backend_type: str):
     while True:
         try:
             message = comm.receive_message()
-            if message:
+            if message and get_debug():
                 print(f"DEBUG: Worker {os.getpid()} received message: {message}")
             if message and message.msg_type == 'train_request':
-                print(f"DEBUG: Worker received train_request: {message.msg_id}")
+                if get_debug():
+                    print(f"DEBUG: Worker received train_request: {message.msg_id}")
                 data = message.data
                 output_model = data.get('output_model', './VideoModel')
                 description = data.get('description', '')
@@ -71,16 +75,20 @@ def worker_process(backend_type: str):
 
                 if train_dir and os.path.isdir(train_dir):
                     print(f"PROGRESS: Job {message.msg_id} accepted - Starting training")
-                    print(f"DEBUG: Starting training for job {message.msg_id}")
+                    if get_debug():
+                        print(f"DEBUG: Starting training for job {message.msg_id}")
                     result = train_model(train_dir, output_model, description)
                     print(f"PROGRESS: Job {message.msg_id} - 100% - Training completed")
-                    print(f"DEBUG: Training completed for job {message.msg_id}")
+                    if get_debug():
+                        print(f"DEBUG: Training completed for job {message.msg_id}")
                 else:
                     result = "No valid training directory provided"
-                    print(f"DEBUG: No valid training directory for job {message.msg_id}")
+                    if get_debug():
+                        print(f"DEBUG: No valid training directory for job {message.msg_id}")
 
                 response = Message('train_response', message.msg_id, {'message': result})
-                print(f"DEBUG: Sending train_response for job {message.msg_id}")
+                if get_debug():
+                    print(f"DEBUG: Sending train_response for job {message.msg_id}")
                 comm.send_message(response)
             time.sleep(0.1)
         except Exception as e: