diff options
| -rw-r--r-- | scripts/run_sweep.py | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/scripts/run_sweep.py b/scripts/run_sweep.py index 34cccf7..7b75921 100644 --- a/scripts/run_sweep.py +++ b/scripts/run_sweep.py @@ -60,7 +60,10 @@ def run_one(exp: str, job: dict, post_only: bool) -> dict: if not post_only: cmd = [PY, "-m", "src.train", job["model"], job["seed"], "--out_dir", f"runs/{exp}/{job['job']}"] + job["flags"] - env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1") + # CPU-resident jobs get 2 threads (sequential tied cell benefits); GPU jobs + # stay at 1 so the feeder process doesn't contend with concurrent CPU jobs. + n_threads = "2" if "--device cpu" in job["flags"] else "1" + env = dict(os.environ, OMP_NUM_THREADS=n_threads, MKL_NUM_THREADS=n_threads) with open(logpath, "w") as lf: p = subprocess.run(cmd, cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT) if p.returncode != 0: |