Improve logging readability

Author: chinganc

opto/features/async_search/async_priority_search.py

@@ -199,7 +199,7 @@ async def _initialize_search_parameters(self, *,
         """
         # Validate and adjust num_candidates
         if num_candidates < len(self._optimizers):
-            print(f"Warning: num_candidates {num_candidates} is less than the number of optimizers {len(self._optimizers)}. Setting num_candidates to {len(self._optimizers)}.")
+            print(f"[AsyncPrioritySearch] Warning: num_candidates {num_candidates} is less than the number of optimizers {len(self._optimizers)}. Setting num_candidates to {len(self._optimizers)}.")
             num_candidates = len(self._optimizers)
 
         # Set core parameters
@@ -230,15 +230,15 @@ async def _initialize_search_parameters(self, *,
 
         # Initialize memory structures
         if memory_update_frequency is None:
-            print("AsyncPrioritySearch initialized with only short-term memory.")
+            print("[AsyncPrioritySearch] AsyncPrioritySearch initialized with only short-term memory.")
             assert short_term_memory_size is None or short_term_memory_size > 0, \
                 "short_term_memory_size must be None or greater than 0 when memory_update_frequency is None."
         elif memory_update_frequency == 0:
-            print("AsyncPrioritySearch initialized with only long-term memory.")
+            print("[AsyncPrioritySearch] AsyncPrioritySearch initialized with only long-term memory.")
             assert long_term_memory_size is None or long_term_memory_size > 0, \
                 "long_term_memory_size must be None or greater than 0 when memory_update_frequency is 0."
         else:
-            print(f"AsyncPrioritySearch initialized with both short-term and long-term memory. Candidates will be merged into long-term memory every {memory_update_frequency} iterations.")
+            print(f"[AsyncPrioritySearch] AsyncPrioritySearch initialized with both short-term and long-term memory. Candidates will be merged into long-term memory every {memory_update_frequency} iterations.")
 
         self.long_term_memory = HeapMemory(size=long_term_memory_size, processing_fun=self.compress_candidate_memory)
         self.short_term_memory = HeapMemory(size=short_term_memory_size)
@@ -270,7 +270,7 @@ def memory(self):
         if self.n_iters % self.memory_update_frequency == 0:
             # Merge short-term memory into long-term memory
             if len(self.short_term_memory) > 0:
-                print('Merging short-term memory into long-term memory.')
+                print('[AsyncPrioritySearch] Merging short-term memory into long-term memory.')
                 self.long_term_memory.append(self.short_term_memory)
                 self.short_term_memory.reset()
             return self.long_term_memory
@@ -371,7 +371,7 @@ async def propose(self, samples: Samples, exploration_candidates: List[ModuleCan
         Returns:
             List of proposed ModuleCandidate objects
         """
-        print("--- Proposing new parameters...") if verbose else None
+        print("[AsyncPrioritySearch] --- Proposing new parameters...") if verbose else None
         assert isinstance(samples, Samples), "samples must be an instance of Samples."
         samples_list = samples.samples  # list of BatchRollout objects
         n_proposals = self.num_proposals
@@ -405,7 +405,7 @@ async def _backward(n):
             return optimizer
 
         # Run backward passes concurrently
-        print(f"Running backward on {n_batches} batches...") if verbose else None
+        print(f"[AsyncPrioritySearch] Running backward on {n_batches} batches...") if verbose else None
         optimizers = await asyncio.gather(*[_backward(n) for n in range(n_batches)])
         assert len(optimizers) == n_batches, "Number of optimizers must match number of batch rollouts."
 
@@ -432,7 +432,7 @@ async def _step(n):
             return update_dict
 
         # Run step operations concurrently
-        print(f"Generating {n_proposals} proposals for each of {n_batches} batches...") if verbose else None
+        print(f"[AsyncPrioritySearch] Generating {n_proposals} proposals for each of {n_batches} batches...") if verbose else None
         update_dicts = await asyncio.gather(*[_step(n) for n in range(n_batches * n_proposals)])
 
         # Clear optimizer state TODO is this for preventing memory leak?
@@ -467,7 +467,7 @@ async def validate(self,
         Returns:
             Dictionary mapping ModuleCandidate to list of rollouts
         """
-        print("--- Validating candidates...") if verbose else None
+        print("[AsyncPrioritySearch] --- Validating candidates...") if verbose else None
         assert isinstance(samples, Samples), "samples must be an instance of Samples."
         assert exploration_candidates is not None, "exploration_candidate must be set."
 
@@ -537,7 +537,7 @@ async def update_memory(self, validate_results: Dict[ModuleCandidate, List[Dict[
             verbose: Whether to print verbose output
             **kwargs: Additional keyword arguments
         """
-        print("--- Updating memory with validation results...") if verbose else None
+        print("[AsyncPrioritySearch] --- Updating memory with validation results...") if verbose else None
         for candidate, rollouts in validate_results.items():
             candidate.add_rollouts(rollouts)
             priority = self.compute_exploration_priority(candidate)
@@ -560,7 +560,7 @@ async def explore(self,
             priorities: List of priorities for candidates
             info_dict: Dictionary with logging information
         """
-        print(f"--- Generating {min(len(self.memory), num_candidates)} exploration candidates...") if verbose else None
+        print(f"[AsyncPrioritySearch] --- Generating {min(len(self.memory), num_candidates)} exploration candidates...") if verbose else None
 
         top_candidates = [self._best_candidate] if use_best_candidate_to_explore else []
         priorities = [self._best_candidate_priority] if use_best_candidate_to_explore else []
@@ -599,7 +599,7 @@ async def exploit(self, verbose: bool = False, **kwargs) -> Tuple[ModuleCandidat
             priority: Priority of best candidate
             info_dict: Dictionary with logging information
         """
-        print("--- Exploiting the best candidate...") if verbose else None
+        print("[AsyncPrioritySearch] --- Exploiting the best candidate...") if verbose else None
         if not self.memory:
             raise ValueError("The priority queue is empty. Cannot exploit.")
 

opto/features/async_search/async_search.py

@@ -245,7 +245,7 @@ async def create_task(self) -> Optional[Coroutine[Any, Any, Any]]:
                     return self.eval_task(task_id=task_id)
 
         # Otherwise, run worker task
-        print(f"Epoch: {self.n_epochs}. Iteration: {self.n_iters}")
+        print(f"[AsyncSearch] Epoch: {self.n_epochs}. Iteration: {self.n_iters}")
         task_id = WORKER_TASK + ':' + str(self.n_iters) + ':' + uuid.uuid4().hex
         proposal, task_state = await self.get_task_state()
         self._running_worker_tasks[task_id] = (proposal, task_state)
@@ -283,9 +283,11 @@ async def process_result(self, result: Tuple[str, Any]) -> None:
 
         task_id, result = result
 
+        print(f"[AsyncSearch] Processing result for task: {task_id}")
+
         if task_id.startswith(EVAL_TASK):
             info_test = result
-            self.log(info_test, prefix="Test/")
+            self.log(info_test, prefix="Test/", color='green')
             return
 
         # Else, worker task is done
@@ -356,7 +358,7 @@ async def sample(self, agents, verbose=False, **kwargs):
         }
         # check if the scores are within the score range
         if hasattr(self, '_score_range') and not (self.min_score <= log_info['mean_score'] <= self.max_score):
-            print(f"Warning: Mean score {log_info['mean_score']} is out of the range {self._score_range}.")
+            print(f"[AsyncSearch] Warning: Mean score {log_info['mean_score']} is out of the range {self._score_range}.")
         return samples, log_info
 
 
@@ -370,21 +372,21 @@ async def evaluate_agent(self, test_dataset, guide):
         test_score = safe_mean(test_scores)
         # check if the test_score is within the score range
         if hasattr(self, '_score_range') and not (self.min_score <= test_score <= self.max_score):
-            print(f"Warning: Test score {test_score} is out of the range {self._score_range}.")
+            print(f"[AsyncSearch] Warning: Test score {test_score} is out of the range {self._score_range}.")
         return {'test_score': test_score}
 
 
-    def log(self, info_log: Dict[str, Any], prefix=""):
+    def log(self, info_log: Dict[str, Any], prefix="", color=None) -> None:
         """Log information from the algorithm."""
         for key, value in info_log.items():
             if value is not None and self.logger:
                 try:
-                    self.logger.log(f"{prefix}{key}", value, self.n_iters)
+                    self.logger.log(f"{prefix}{key}", value, self.n_iters, color=color)
                 except Exception as e:
-                    print(f"Logging failed for key {key}: {e}")
+                    print(f"[AsyncSearch] Logging failed for key {key}: {e}")
 
     def save(self, save_path):
-        print(f"Saving algorithm state to {save_path} at iteration {self.n_iters}.")
+        print(f"[AsyncSearch] Saving algorithm state to {save_path} at iteration {self.n_iters}.")
         if not os.path.exists(save_path):
             os.makedirs(save_path)
         obj = copy.deepcopy(self)  # to detach nodes from the computation graph
@@ -430,7 +432,7 @@ def resume(self, *,
         last_train_kwargs['train_dataset'] = train_dataset
         last_train_kwargs['validate_dataset'] = validate_dataset
         last_train_kwargs['test_dataset'] = test_dataset
-        print(f"Resuming training with parameters: {last_train_kwargs}")
+        print(f"[AsyncSearch] Resuming training with parameters: {last_train_kwargs}")
         return self.train(*last_train_args, **last_train_kwargs)
 
 

opto/features/async_search/controller.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from typing import Any, Coroutine, Optional
 
 
@@ -30,7 +31,8 @@ async def run(self, *args: Any, **kwargs: Any) -> None:
             done, pending = await asyncio.wait(worker_tasks, return_when=asyncio.FIRST_COMPLETED)
             for task in done:
                 worker_tasks.remove(task)
-                result = task.result()
+                result, elapsed_time = task.result()
+                print(f"[Controller] Task completed in {elapsed_time:.4f} seconds")
                 await self.process_result(result)
                 if self.should_stop():
                     self.post_process()
@@ -47,7 +49,15 @@ async def _create_new_task(self) -> Optional[asyncio.Task[Any]]:
         new_task = await self.create_task()
         if new_task is None:
             return None
-        return asyncio.create_task(new_task)
+
+        # Wrap the task with timing
+        async def timed_task() -> tuple[Any, float]:
+            start_time = time.time()
+            result = await new_task
+            elapsed_time = time.time() - start_time
+            return result, elapsed_time
+
+        return asyncio.create_task(timed_task())
 
     #### Methods to be implemented by subclasses ####