-

benchmarks/torchvision/main.py

@@ -165,7 +165,10 @@ def main():
             if data_directory:
                 args.data = os.path.join(data_directory, "FakeImageNet")
 
-    use_cuda = not args.no_cuda and torch.cuda.is_available()
+    if not args.no_cuda:
+        assert torch.cuda.is_available(), "Why is CUDA not available"
+
+    use_cuda = not args.no_cuda
 
     torch.manual_seed(args.seed)
     if use_cuda:

milabench/cli.py

@@ -741,15 +741,21 @@ def publish():
 
     def schedule():
         """Launch a slurm job to run milabench"""
+        # milabench schedule --sync -- --select resnet50
 
         # tail -f on the slurm job
         sync: Option & bool = False
 
         # Print the command and return without running it
         dry: Option & bool = False
 
+        # pip arguments
+        # [remainder]
+        args: Option = []
+
         launch_milabench(
-            None,
-            dry,
-            sync
+            args,
+            sbatch_args=None,
+            dry=dry,
+            sync=sync
         )

milabench/report.py

@@ -49,6 +49,7 @@ def _make_row(summary, compare, weights):
     score = (acc if acc > 0 else row["perf"]) * success_ratio
 
     row["score"] = score
+    print(score)
     row["weight"] = weights.get("weight", summary["weight"])
     # ----
 
@@ -224,10 +225,12 @@ def make_report(
         weights = dict()
 
     df = make_dataframe(summary, compare, weights)
-
+    print(df)
+
     # Reorder columns
     df = df[sorted(df.columns, key=lambda k: columns_order.get(k, 0))]
-
+    print(df)
+
     out = Outputter(stdout=sys.stdout, html=html)
 
     if sources:

milabench/schedule.py

@@ -125,7 +125,7 @@ def arguments(self):
         ]
 
 
-def launch_milabench(sbatch_args=None, dry: bool = False, sync: bool = False):
+def launch_milabench(args, sbatch_args=None, dry: bool = False, sync: bool = False):
     sbatch_script = importlib_resources.files(__name__) / "scripts" / "milabench.bash"
     sbatch_script = str(sbatch_script)
 
@@ -143,7 +143,7 @@ def launch_milabench(sbatch_args=None, dry: bool = False, sync: bool = False):
     script_args.deduce_from_repository()
     script_args = script_args.arguments()
 
-    cmd = sbatch_args + [sbatch_script] + script_args
+    cmd = sbatch_args + [sbatch_script] + script_args + args
 
     if dry:
         print("sbatch " + ' '.join(cmd))

milabench/scripts/milabench.bash

@@ -68,6 +68,7 @@ CONDA_EXEC="$(which conda)"
 CONDA_BASE=$(dirname $CONDA_EXEC)
 source $CONDA_BASE/../etc/profile.d/conda.sh
 
+cd $SLURM_TMPDIR
 #
 #   Create a new environment
 #
@@ -89,12 +90,10 @@ export PYTHONUNBUFFERED=1
 #
 # Fetch the repo
 #
-cd $SLURM_TMPDIR
 git clone --single-branch --depth 1 -b $BRANCH $ORIGIN
 python -m pip install ./milabench
 
 SYSTEM="$SLURM_TMPDIR/system.yaml"
-unset CUDA_VISIBLE_DEVICES
 
 echo ""
 echo "System"