Staging (#367)

* Update sharedsetup to handle tar archive
* Update docker container generation to handle UV correctly
* Tweak dashboard utilities

Author: Delaunay

benchmarks/dinov2/prepare.py

@@ -37,7 +37,8 @@ def loop_on(iterable:list):
         device=None,
         image_size=[3, 384, 384],
         val=0.1,
-        test=0.1
+        test=0.1,
+        output=os.getenv("MILABENCH_DIR_DATA", None),
     )
     # 
     generate_fakeimagenet(args)

benchmate/benchmate/datagen.py

@@ -6,6 +6,8 @@
 import os
 from collections import defaultdict
 from pathlib import Path
+import time
+import random
 
 import torchcompat.core as acc
 import torch
@@ -17,11 +19,18 @@ def write(args):
 
     offset, outdir, prefix, size = args
 
-    img = torch.randn(*size)
-    target = offset % 1000  # torch.randint(0, 1000, size=(1,), dtype=torch.long)[0]
+    seed = int(time.time() + offset)
+    
+    torch.manual_seed(seed)
+    random.seed(seed)
+    
+    img = torch.randint(0, 256, size, dtype=torch.uint8)
+    # img = torch.randn(*size)
+    
     img = transforms.ToPILImage()(img)
+    target = offset % 1000  # torch.randint(0, 1000, size=(1,), dtype=torch.long)[0]
     class_val = int(target)
-
+    
     # Some benches need filenames to match those of imagenet:
     # https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/data/datasets/image_net.py#L40-L43
     if not prefix:  # train
@@ -110,6 +119,7 @@ def fakeimagenet_args():
     parser.add_argument("--image-size", default=[3, 384, 384], type=int, nargs="+")
     parser.add_argument("--val", default=0.1, type=float, nargs="+")
     parser.add_argument("--test", default=0.1, type=float, nargs="+")
+    parser.add_argument("--output", default=os.getenv("MILABENCH_DIR_DATA", None), type=str)
     args, _ = parser.parse_known_args()
     return args
 
@@ -124,7 +134,8 @@ def generate_fakeimagenet(args=None):
         bs, bc = overrides.split(",")
         args.batch_size, args.batch_count = int(bs), int(bc)
 
-    data_directory = os.environ["MILABENCH_DIR_DATA"]
+    assert args.output is not None, "Output directory is required"
+    data_directory = args.output
 
     dest = os.path.join(data_directory, f"FakeImageNet")
     print(f"Generating fake data into {dest}...")

config/slurm.yaml

@@ -50,16 +50,18 @@ single-node-full-A100l:
 - --export=ALL
 - -w cn-d[003-004]
 single-node-full-L40S:
-- -w cn-l[001-091]
-- --ntasks=1
-- --gpus-per-task=4
-- --exclusive
+- --job-name=run_l40s
+- --partition=staff-idt
 - --nodes=1
+- --ntasks=1
 - --cpus-per-task=48
-- --time=1:30:00
-- --ntasks-per-node=1
 - --mem=0
+- --time=4:30:00
+- --gpus-per-task=4
+- --ntasks-per-node=1
+- --exclusive
 - --export=ALL,MILABENCH_SIZER_AUTO=0
+- -w cn-l[001-092]
 single-node-small:
 - --partition=staff-idt
 - --ntasks=1
@@ -172,3 +174,14 @@ run_shared_install_prepare:
 - --gpus-per-task=1
 - --ntasks-per-node=1
 - --export=ALL
+tar_prepare:
+- --job-name=tar
+- --partition=staff-idt
+- --nodes=1
+- --ntasks=1
+- --cpus-per-task=2
+- --mem=16G
+- --time=06:00:00
+- --gpus-per-task=0
+- --ntasks-per-node=1
+- --export=ALL

docker/Dockerfile-cuda

@@ -1,7 +1,7 @@
 # FROM ubuntu:22.04
 
-# For cuda-gdb
-FROM nvidia/cuda:12.9.1-devel-ubuntu24.04
+FROM nvidia/cuda:12.9.1-runtime-ubuntu24.04
+
 
 # Arguments
 # ---------
@@ -59,11 +59,12 @@ ENV CUDA_HOME=/usr/local/cuda-12/
 
 ENV MILAENCH_ENV=/milabench/.env
 
-RUN virtualenv $MILAENCH_ENV                                            &&\
+RUN RUN curl -LsSf https://astral.sh/uv/install.sh | sh                 &&\
+    virtualenv $MILAENCH_ENV                                            &&\
     $MILAENCH_ENV/bin/python -m pip install -U pip                      &&\
     $MILAENCH_ENV/bin/python -m pip install -U setuptools               &&\
     $MILAENCH_ENV/bin/python -m pip install -U poetry                   &&\
-    $MILAENCH_ENV/bin/python -m pip install -e /milabench/milabench/    &&\
+    $MILAENCH_ENV/bin/python -m pip install -e /milabench/milabench[cuda]    &&\
     $MILAENCH_ENV/bin/python -m pip cache purge                       
 
 # Prepare bench