crfm-helm 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/benchmark/slurm_runner.py

@@ -10,6 +10,10 @@ import sys
 
 from helm.common.codec import from_json, to_json
 from helm.common.general import write
+from helm.benchmark.config_registry import (
+    register_configs_from_directory,
+    register_builtin_configs_from_helm_package,
+)
 from helm.benchmark.executor import ExecutionSpec
 from helm.benchmark.runner import Runner, RunSpec, RunnerError
 from helm.benchmark.slurm_jobs import (
@@ -24,10 +28,11 @@ from helm.benchmark.slurm_jobs import (
 from helm.common.general import ensure_directory_exists
 from helm.common.hierarchical_logger import hlog, htrack_block
 
+from helm.benchmark.runner_config_registry import RUNNER_CONFIG
 
-_DEFAULT_MAX_CONCURRENT_WORKER_SLURM_JOBS = 8
 _MAX_CONCURRENT_WORKER_SLURM_JOBS_ENV_NAME = "HELM_MAX_CONCURRENT_WORKER_SLURM_JOBS"
 _SLURM_NODE_NAMES_ENV_NAME = "HELM_SLURM_NODE_NAMES"
+_DEFAULT_MAX_CONCURRENT_WORKER_SLURM = 8
 
 
 @dataclass
@@ -89,12 +94,15 @@ class SlurmRunner(Runner):
         self.slurm_runner_spec_path = os.path.join(self.slurm_base_dir, "slurm_runner_spec.json")
 
         # Configure max concurrent worker Slurm jobs from the environment variable.
-        # TODO: Read from a configuration file instead
         env_max_concurrent_worker_slurm_jobs = os.getenv(_MAX_CONCURRENT_WORKER_SLURM_JOBS_ENV_NAME)
         self.max_concurrent_worker_slurm_jobs = (
             int(env_max_concurrent_worker_slurm_jobs)
             if env_max_concurrent_worker_slurm_jobs
-            else _DEFAULT_MAX_CONCURRENT_WORKER_SLURM_JOBS
+            else (
+                RUNNER_CONFIG.helm_max_concurrent_workers
+                if RUNNER_CONFIG.helm_max_concurrent_workers > 0
+                else _DEFAULT_MAX_CONCURRENT_WORKER_SLURM
+            )
         )
 
     def run_all(self, run_specs: List[RunSpec]):
@@ -222,8 +230,7 @@ class SlurmRunner(Runner):
                         break
 
                     # Refresh every minute
-                    # TODO: Make this period configurable
-                    time.sleep(60)
+                    time.sleep(RUNNER_CONFIG.slurm_monitor_interval)
         finally:
             # Cleanup by cancelling all jobs during program termination or if an exception is raised.
             cancel_all_jobs()
@@ -261,34 +268,48 @@ class SlurmRunner(Runner):
                 run_spec_path,
             ]
         )
-        # TODO: Make default Slurm arguments configurable.
-        raw_slurm_args: Dict[str, str] = {
-            "account": "nlp",
-            "cpus_per_task": "4",
-            "mem": "32G",
-            "gres": "gpu:0",
-            "open_mode": "append",
-            "partition": "john",
-            "time": "14-0",  # Deadline of 14 days
-            "mail_type": "FAIL",
-            "job_name": run_name,
-            "output": log_path,
-            "chdir": os.getcwd(),
-        }
-        # TODO: Move resource requirements into RunSpec.
-        slurm_node_names = os.getenv(_SLURM_NODE_NAMES_ENV_NAME)
-        if run_spec.name.startswith("msmarco:"):
-            raw_slurm_args["mem"] = "64G"
-        if "device=cuda" in run_spec.name:
-            raw_slurm_args["gres"] = "gpu:1"
-            raw_slurm_args["partition"] = "jag-hi"
-        if "model=huggingface" in run_spec.name:
-            raw_slurm_args["gres"] = "gpu:1"
-            raw_slurm_args["partition"] = "sphinx"
-            if not slurm_node_names or "sphinx" not in slurm_node_names:
-                raise Exception(f"Environment variable {_SLURM_NODE_NAMES_ENV_NAME} must be set to sphinx node names")
-        if slurm_node_names:
-            raw_slurm_args["nodelist"] = slurm_node_names
+        if RUNNER_CONFIG.slurm_args is None:
+            raw_slurm_args: Dict[str, str] = {
+                "account": "nlp",
+                "cpus_per_task": "4",
+                "mem": "32G",
+                "gres": "gpu:0",
+                "open_mode": "append",
+                "partition": "john",
+                "time": "14-0",  # Deadline of 14 days
+                "mail_type": "FAIL",
+                "job_name": run_name,
+                "output": log_path,
+                "chdir": os.getcwd(),
+            }
+            # TODO: Move resource requirements into RunSpec.
+            slurm_node_names = os.getenv(_SLURM_NODE_NAMES_ENV_NAME)
+            if run_spec.name.startswith("msmarco:"):
+                raw_slurm_args["mem"] = "64G"
+            if "device=cuda" in run_spec.name:
+                raw_slurm_args["gres"] = "gpu:1"
+                raw_slurm_args["partition"] = "jag-hi"
+            if "model=huggingface" in run_spec.name:
+                raw_slurm_args["gres"] = "gpu:1"
+                raw_slurm_args["partition"] = "sphinx"
+                if not slurm_node_names or "sphinx" not in slurm_node_names:
+                    raise Exception(
+                        f"Environment variable {_SLURM_NODE_NAMES_ENV_NAME} must be set to sphinx node names"
+                    )
+            if slurm_node_names:
+                raw_slurm_args["nodelist"] = slurm_node_names
+
+        else:
+            raw_slurm_args = RUNNER_CONFIG.slurm_args
+
+            dynamic_slurm_args = {
+                "job_name": run_name,
+                "output": log_path,
+                "chdir": os.getcwd(),
+            }
+
+            # User should not set these manually, overwrite them if necessary
+            raw_slurm_args.update(dynamic_slurm_args)
 
         slurm_args: Dict[str, str] = {key: shlex.quote(value) for key, value in raw_slurm_args.items()}
         # Uncomment this to get notification emails from Slurm for Slurm worker jobs.
@@ -300,27 +321,15 @@ class SlurmRunner(Runner):
         return slurm_job_id
 
 
-def run_as_worker(slurm_runner_spec_path: str, run_spec_path: str):
-    """Deserialize SlurmRunner and RunSpec from the given files, then run the RunSpec with the SlurmRunner.
-
-    Used by the worker Slurm jobs only."""
-    with open(slurm_runner_spec_path, "r") as f:
-        slurm_runner_spec = from_json(f.read(), SlurmRunnerSpec)
-    with open(run_spec_path, "r") as f:
-        run_spec = from_json(f.read(), RunSpec)
-    slurm_runner = SlurmRunner(**slurm_runner_spec.to_kwargs())
-    slurm_runner.run_one(run_spec)
-
-
 def main():
     """Entry point for the SlurmRunner's worker Slurm jobs that run a single RunSpec.
 
     This entry point should only be used by SlurmRunner. Users should use `helm-run` instead.
     SlurmRunner has to use this entry point instead of helm-run because there is no way to
     specify the worker Slurm job parameters through `helm-run`. In particular, there is no way
-    to run a specific `RunSpec` using the `--run-specs` parameter of `helm-run`, because the
-    `run-specs` argument is a `RunSpec` description (not a `RunSpec`), and there is no way to
-    convert a `RunSpec` into a `RunSpec` description."""
+    to run a specific `RunSpec` using the `--run-entries` parameter of `helm-run`, because the
+    `run-entries` argument contains `RunEntry` description (not `RunSpec`s), and there is no way to
+    convert a `RunSpec` into a `RunEntry` description."""
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--slurm-runner-spec-path",
@@ -335,7 +344,19 @@ def main():
         required=True,
     )
     args = parser.parse_args()
-    run_as_worker(slurm_runner_spec_path=args.slurm_runner_spec_path, run_spec_path=args.run_spec_path)
+
+    # Deserialize SlurmRunner and RunSpec from the given files, then run the RunSpec with the SlurmRunner.
+    with open(args.slurm_runner_spec_path, "r") as f:
+        slurm_runner_spec = from_json(f.read(), SlurmRunnerSpec)
+    with open(args.run_spec_path, "r") as f:
+        run_spec = from_json(f.read(), RunSpec)
+
+    register_builtin_configs_from_helm_package()
+    if slurm_runner_spec.execution_spec.local_path is not None:
+        register_configs_from_directory(slurm_runner_spec.execution_spec.local_path)
+
+    slurm_runner = SlurmRunner(**slurm_runner_spec.to_kwargs())
+    slurm_runner.run_one(run_spec)
 
 
 if __name__ == "__main__":

helm/benchmark/static/benchmarking.js

@@ -492,7 +492,7 @@ $(function () {
               {{~#if perturbation~}}
                 {{highlightNewWords input.text ../unperturbedInstance.input.text}}
               {{~else~}}
-                {{input.text}}
+                {{{input.text}}}
               {{~/if~}}
             </div>
           {{/if}}