limit val_max_num_task

Author: binary-husky

ajet/backbone/trainer_verl.py

@@ -1042,5 +1042,24 @@ def get_val_dataset(self):
             self.config.ajet.task_reader,
         )
         tasks = task_reader.get_validation_tasks()
+
+        # clip validation tasks if val_max_num_task_each_validation is set
+        val_max_num_task = self.config.ajet.trainer_common.val_max_num_task_each_validation
+        if val_max_num_task is not None and len(tasks) > val_max_num_task:
+            original_size = len(tasks)
+            clip_method = self.config.ajet.trainer_common.val_max_num_task_clip_method
+            if clip_method == "fix_seed_random_n":
+                rng = np.random.RandomState(seed=42)
+                indices = rng.choice(len(tasks), val_max_num_task, replace=False)
+                tasks = [tasks[i] for i in sorted(indices)]
+            elif clip_method == "random_n":
+                indices = np.random.choice(len(tasks), val_max_num_task, replace=False)
+                tasks = [tasks[i] for i in sorted(indices)]
+            elif clip_method == "first_n":
+                tasks = tasks[:val_max_num_task]
+            else:
+                raise ValueError(f"Unknown val_max_num_task_clip_method: {clip_method}, expected 'fix_seed_random_n', 'random_n', or 'first_n'")
+            logger.info(f"Clipped validation dataset from {original_size} to {val_max_num_task} tasks using '{clip_method}'")
+
         self.main_val_dataset = tasks
         return self.main_val_dataset

ajet/default_config/ajet_default.yaml

@@ -235,6 +235,8 @@ ajet:
     val_before_train: False
     val_pass_n: 4
     val_only: False
+    val_max_num_task_each_validation: null
+    val_max_num_task_clip_method: fix_seed_random_n  # fix_seed_random_n, random_n, or first_n
     val_print_to_markdown_file_path: null
     train_print_to_markdown_file_path: null
 

ajet/utils/launch_utils.py

@@ -219,6 +219,9 @@ def start_ray_service(args, env, cluster=False):
     # Get the current Python interpreter directory
     python_dir = os.path.dirname(sys.executable)
     ray_path = os.path.join(python_dir, "ray")
+    # if CUDA_VISIBLE_DEVICES is set, remove it from `env`
+    if "CUDA_VISIBLE_DEVICES" in env:
+        del env["CUDA_VISIBLE_DEVICES"]
     if not cluster:
         companion = LaunchCommandWhenAbsent(
             full_argument_list=[f"{ray_path} start --head --block"],
@@ -252,7 +255,7 @@ def start_ray_service(args, env, cluster=False):
                 tag="ray_service_worker",
                 use_pty=True,
             )
-            launch_wait_time = 9999999999
+            launch_wait_time = 9999999999999
             # success_std_string = "Connected to Ray cluster"
             success_std_string = "Just wait here forever"
     companion.launch(