Merge branch 'stacc' of https://github.com/mila-iqia/milabench into stacc

Author: Delaunay

.github/workflows/docker.yml

@@ -69,11 +69,9 @@ jobs:
         run: |
           REGEX="(.*)v(.*)\.(.*)\.(.*)"
           IMAGE_TAG="nightly"
-          if [[ "${GITHUB_REF_NAME_ENV}" =~ $REGEX ]]; then
-              IMAGE_TAG="${GITHUB_REF_NAME##*/}"
-          fi
+          IMAGE_TAG="${GITHUB_REF_NAME##*/}"
           echo "IMAGE_TAG=$IMAGE_TAG" >> $GITHUB_ENV
-    
+
       - name: Log in to the registry
         uses: docker/login-action@v2
         with:

docker/Dockerfile-cuda

@@ -41,7 +41,7 @@ COPY . /milabench/milabench/
 
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update -y &&\
-    apt-get install -y --no-install-recommends ssh rsync git build-essential curl python3 python3-virtualenv python-is-python3 python3-pip libgl1 libglib2.0-0 gcc g++ python3-dev &&\
+    apt-get install -y --no-install-recommends zip ssh rsync git build-essential curl python3 python3-virtualenv python-is-python3 python3-pip libgl1 libglib2.0-0 gcc g++ python3-dev &&\
     curl -o /etc/apt/trusted.gpg.d/mellanox.asc https://content.mellanox.com/ofed/RPM-GPG-KEY-Mellanox &&\
     curl -o /etc/apt/sources.list.d/mellanox.list https://linux.mellanox.com/public/repo/mlnx_ofed/${MOFED_VERSION}/ubuntu22.04/mellanox_mlnx_ofed.list &&\
     apt-get update -y &&\

milabench/cli/run.py

@@ -75,7 +75,7 @@ def _fetch_arch(mp):
         return None
 
 
-def run(mp, args, name):
+def run(mp, args, name, common_args):
     layers = validation_names(args.validations)
 
     dash_class = {
@@ -116,9 +116,10 @@ def run(mp, args, name):
             summary = make_summary(reports)
             assert len(summary) != 0, "No summaries"
 
-            margs = multipack_args()
-            margs.config = args.config
-            config = _get_multipack(margs, return_config=True)
+            # This makes the config smaller but we want the full config for the weights
+            common_args.select = ""
+            common_args.exclude = ""
+            config = _get_multipack(args=common_args, return_config=True)
 
             make_report(
                 summary,
@@ -136,28 +137,32 @@ def run(mp, args, name):
 
 
 @tooled
-def cli_run(args=None):
+def cli_run(run_args=None):
     """Run the benchmarks."""
-    if args is None:
-        args = arguments()
+
+    if run_args is None:
+        run_args = arguments()
+
+    common_args = multipack_args()
 
     # Load the configuration and system
-    mp = get_multipack(run_name=args.run_name)
+    mp = get_multipack(args=common_args, run_name=run_args.run_name)
+
     arch = _fetch_arch(mp)
 
     # Initialize the backend here so we can retrieve GPU stats
     init_arch(arch)
 
     success = 0
     for name, conf in multirun():
-        run_name = name or args.run_name
+        run_name = name or run_args.run_name
 
         # Note that this function overrides the system config
-        mp = get_multipack(run_name=run_name)
+        mp = get_multipack(args=common_args, run_name=run_name)
 
         with apply_system(conf):
             try:
-                success += run(mp, args, run_name)
+                success += run(mp, run_args, run_name, common_args)
             except AssertionError as err:
                 print(err)
 

milabench/common.py

@@ -210,7 +210,7 @@ def _get_multipack(
     if args.config is None:
         sys.exit("Error: CONFIG argument not provided and no $MILABENCH_CONFIG")
 
-    if args.select:
+    if args.select and isinstance(args.select, str):
         args.select = set(args.select.split(","))
 
     if args.exclude: