diff options
Diffstat (limited to 'experiments/transformer_crossover_t2.py')
| -rw-r--r-- | experiments/transformer_crossover_t2.py | 62 |
1 files changed, 24 insertions, 38 deletions
diff --git a/experiments/transformer_crossover_t2.py b/experiments/transformer_crossover_t2.py index 17f61a2..4f12e56 100644 --- a/experiments/transformer_crossover_t2.py +++ b/experiments/transformer_crossover_t2.py @@ -10,6 +10,15 @@ import time ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) +from experiments.crossover_hardware import ( + DEFAULT_PROFILE, + physical_gpu_report, + policy_for_report, + profile_choices, +) + + PROTOCOL = os.path.join(ROOT, "TRANSFORMER_CROSSOVER.md") RESULT_ROOT = os.path.join(ROOT, "results", "transformer_crossover") DEFAULT_SELECTOR = os.path.join(RESULT_ROOT, "p1_selector.json") @@ -123,6 +132,7 @@ def source_report(selector_path): os.path.join( ROOT, "experiments", "transformer_feedback_smoke.py" ), + os.path.join(ROOT, "experiments", "crossover_hardware.py"), os.path.join(ROOT, "sdil", "transformer.py"), os.path.abspath(selector_path), ] @@ -235,7 +245,7 @@ def registry_sha256(jobs): return hashlib.sha256(encoded).hexdigest() -def ensure_launch(source, selector, jobs): +def ensure_launch(source, selector, jobs, hardware_policy): path = os.path.join(RESULT_ROOT, "t2_launch.json") expected = { "stage": "t2", @@ -243,7 +253,9 @@ def ensure_launch(source, selector, jobs): "selector": selector, "registry_sha256": registry_sha256(jobs), "num_jobs": len(jobs), - "allowed_physical_gpus": [5, 7], + "hardware_policy": hardware_policy, + "allowed_physical_gpus": + hardware_policy["allowed_physical_gpu_indices"], } if os.path.isfile(path): if read_json(path) != expected: @@ -274,40 +286,6 @@ def assert_source_unchanged(source): raise RuntimeError("Transformer T2 environment drift") -def physical_gpu_report(dry_run=False): - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if dry_run: - return { - "cuda_visible_devices": visible, - "physical_gpu_index": None, - "physical_gpu_uuid": None, - } - if visible not in {"5", "7"}: - raise RuntimeError("Transformer T2 requires physical GPU 5 or 7") - query = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=index,uuid,name", - "--format=csv,noheader,nounits", - ], - check=True, - capture_output=True, - text=True, - ).stdout.splitlines() - rows = {} - for line in query: - index, uuid, name = [value.strip() for value in line.split(",", 2)] - rows[index] = {"uuid": uuid, "name": name} - if visible not in rows: - raise RuntimeError(f"physical GPU {visible} not found") - return { - "cuda_visible_devices": visible, - "physical_gpu_index": int(visible), - "physical_gpu_uuid": rows[visible]["uuid"], - "physical_gpu_name": rows[visible]["name"], - } - - def run_job(job, source, selector, gpu, dry_run): if not dry_run: assert_source_unchanged(source) @@ -361,16 +339,24 @@ def main(): parser.add_argument("--num-shards", type=int, default=1) parser.add_argument("--method", choices=METHODS) parser.add_argument("--depth", type=int, choices=DEPTHS) + parser.add_argument( + "--hardware-profile", + choices=profile_choices(), + default=os.environ.get("SDIL_HARDWARE_PROFILE", DEFAULT_PROFILE), + ) parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() if not 0 <= args.shard_index < args.num_shards: raise ValueError("invalid Transformer T2 shard") selector = selector_report(args.selector) source = source_report(args.selector) - gpu = physical_gpu_report(args.dry_run) + gpu = physical_gpu_report(args.hardware_profile, args.dry_run) complete_jobs = t2_jobs(selector["selected_rates"]) if not args.dry_run: - launch = ensure_launch(source, selector, complete_jobs) + hardware_policy = policy_for_report(args.hardware_profile, gpu) + launch = ensure_launch( + source, selector, complete_jobs, hardware_policy + ) print(f"T2 launch lock: {launch}", flush=True) jobs = complete_jobs if args.method: |