(
model: List[str],
task: List[str],
batch_size: int,
limit: Optional[int],
parallelize: bool,
gpus: str,
context_length: int,
cutting_context: bool,
answer_length: int,
output_dir: str,
)
| 70 | @click.option("--cutting_context", is_flag=True) |
| 71 | @click.option("--output_dir", default="output", type=str) |
| 72 | def main( |
| 73 | model: List[str], |
| 74 | task: List[str], |
| 75 | batch_size: int, |
| 76 | limit: Optional[int], |
| 77 | parallelize: bool, |
| 78 | gpus: str, |
| 79 | context_length: int, |
| 80 | cutting_context: bool, |
| 81 | answer_length: int, |
| 82 | output_dir: str, |
| 83 | ): |
| 84 | if limit < 0: limit = None |
| 85 | |
| 86 | if gpus is not None: |
| 87 | os.environ["CUDA_VISIBLE_DEVICES"] = gpus |
| 88 | |
| 89 | # Load the given Python file as a module |
| 90 | configs = [ |
| 91 | {"model": m, "task": t} for m in model for t in task |
| 92 | ] |
| 93 | |
| 94 | use_ray = parallelize and len(configs) > 0 |
| 95 | if use_ray: |
| 96 | import ray |
| 97 | # ray was killing workers due to OOM, but it didn't seem to be necessary |
| 98 | os.environ["RAY_memory_monitor_refresh_ms"] = "0" |
| 99 | ray.init(ignore_reinit_error=True, log_to_driver=True) |
| 100 | |
| 101 | print(f"Running sweep with {len(configs)} configs") |
| 102 | |
| 103 | output_dir = f"{output_dir}/{datetime.now().strftime('%y-%m-%d_%H-%M')}" |
| 104 | |
| 105 | # Run each script in parallel using Ray |
| 106 | if not use_ray: |
| 107 | for config in configs: |
| 108 | execute_config( |
| 109 | **config, |
| 110 | batch_size=batch_size, |
| 111 | limit=limit, |
| 112 | output_dir=output_dir, |
| 113 | context_length=context_length, |
| 114 | answer_length=answer_length, |
| 115 | cutting_context=cutting_context |
| 116 | ) |
| 117 | else: |
| 118 | completed = 0 |
| 119 | total = len(configs) |
| 120 | print(f"Completed: {completed} ({completed / total:0.1%}) | Total: {total}") |
| 121 | |
| 122 | remote = ray.remote(num_gpus=(1 // MAX_WORKERS_PER_GPU))(execute_config) |
| 123 | futures = [remote.remote( |
| 124 | **config, batch_size=batch_size, limit=limit, output_dir=output_dir, |
| 125 | answer_length=answer_length, cutting_context=cutting_context, context_length=context_length, |
| 126 | ) for config in configs] |
| 127 | |
| 128 | while futures: |
| 129 | complete, futures = ray.wait(futures) |
no test coverage detected