refactor color adjustment to use batch processing for improved performance

Author: binary-husky

ajet/context_tracker/multiagent_tracking.py

@@ -17,7 +17,7 @@
 )
 from ajet.schema.extended_msg import INVALID_LOG_PROB_VALUE
 from ajet.schema.trajectory import Reward
-from ajet.utils.color_hsl import adjust_color_hsl
+from ajet.utils.color_hsl import adjust_color_hsl_batch
 from ajet.utils.compute_madness import compute_string_madness
 from ajet.utils.tokenizer import ajet_apply_chat_template
 
@@ -503,15 +503,9 @@ def generate_log(self, task_id=None, global_step="NA"):
             logprobs = [INVALID_LOG_PROB_VALUE] * len(
                 tracker_tokenized["prompt_ids"]
             ) + tracker_tokenized["response_logprobs"]
-            # Create adjusted color array
-            loss_mask_color_abl_arr = [
-                (
-                    adjust_color_hsl("#09ABCF", logprob)
-                    if mask == 1
-                    else adjust_color_hsl("#D98510", logprob)
-                )
-                for mask, logprob in zip(tracker_tokenized["loss_mask"], logprobs)
-            ]
+            # Create adjusted color array using batch processing for better performance
+            base_colors = ["#09ABCF" if mask == 1 else "#D98510" for mask in tracker_tokenized["loss_mask"]]
+            loss_mask_color_abl_arr = adjust_color_hsl_batch(base_colors, logprobs)
             logprob_text_arr = [
                 (f"{logprob:.4f}" if logprob != INVALID_LOG_PROB_VALUE else "N/A")
                 for logprob in logprobs

ajet/copilot/train-complex-blackbox/SKILL.md

@@ -155,6 +155,7 @@ Finally, you can start training.
 
 Run `ajet-swarm start` to start training server (if the user has already installed agentjet swarm environment),
 if the user has docker environment, you can also refer to `docs/en/ajet-swarm-docker.md` to start a AgentSwarm docker container.
+If the user can provider the ssh connection to the GPU server / cluster, you can send the `ajet-swarm start` command to the remote server via ssh to start the swarm server, the port forward `10086` port (default agentjet swarm port) to user local machine.
 
 Create a duplication of `agent_roll.py` named `agent_roll_one_episode_debug.py`, and modify it to only run one episode, this can help you debug whether the episode runner and reward function work as expected.
 

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -437,7 +437,7 @@ def start_engine(self):
         self._wait_until_status_change_to(desired_status="ENGINE.ROLLING")
         logger.success("Training engine is now ROLLING and ready.")
 
-    def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=True, timeout=1800):
+    def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=True, timeout=3600):
         """
         Poll engine status until it reaches desired_status.
         Reports status every 5 seconds while waiting.

ajet/utils/color_hsl.py

@@ -1,6 +1,9 @@
 import colorsys
+import numpy as np
+from functools import lru_cache
 
 
+@lru_cache(maxsize=2048)
 def adjust_color_hsl(base_color, logprob):
     """
     Adjust color saturation using the HSL color space based on log probability.
@@ -39,3 +42,58 @@ def adjust_color_hsl(base_color, logprob):
 
     # Convert back to hexadecimal
     return f"#{int(r_adjusted*255):02x}{int(g_adjusted*255):02x}{int(b_adjusted*255):02x}"
+
+
+def adjust_color_hsl_batch(base_colors, logprobs):
+    """
+    Vectorized version of adjust_color_hsl for batch processing.
+    Args:
+        base_colors (list[str]): List of hexadecimal color strings.
+        logprobs (list[float]): List of log probability values.
+    Returns:
+        list[str]: List of adjusted hexadecimal color strings.
+    """
+    if not base_colors or not logprobs:
+        return []
+
+    # Constants
+    sat_min = 0.333
+    sat_max = 1.0
+    lp_min = -7
+    lp_max = 0
+
+    # Convert to numpy arrays for vectorized operations
+    logprobs_arr = np.array(logprobs, dtype=np.float32)
+
+    # Vectorized saturation factor calculation
+    saturation_factors = np.where(
+        logprobs_arr <= lp_min,
+        sat_min,
+        np.where(
+            logprobs_arr >= 0,
+            sat_max,
+            sat_min + (logprobs_arr - lp_min) / (lp_max - lp_min) * (sat_max - sat_min)
+        )
+    )
+
+    # Pre-convert unique base colors to RGB and HSL
+    unique_colors = list(set(base_colors))
+    color_to_hls = {}
+
+    for color in unique_colors:
+        r = int(color[1:3], 16) / 255.0
+        g = int(color[3:5], 16) / 255.0
+        b = int(color[5:7], 16) / 255.0
+        h, l, s = colorsys.rgb_to_hls(r, g, b)
+        color_to_hls[color] = (h, l, s)
+
+    # Process each color
+    result = []
+    for base_color, sat_factor in zip(base_colors, saturation_factors):
+        h, l, s = color_to_hls[base_color]
+        s_adjusted = s * sat_factor
+        r_adj, g_adj, b_adj = colorsys.hls_to_rgb(h, l, s_adjusted)
+        hex_color = f"#{int(r_adj*255):02x}{int(g_adj*255):02x}{int(b_adj*255):02x}"
+        result.append(hex_color)
+
+    return result