Add missing tags to tests config (#312)

* Add missing tags to tests config

* Add revision

* Add regression files

* update test_memory_tracking

---------

Co-authored-by: Pierre Delaunay <[email protected]>

benchmarks/llava/main.py

@@ -62,8 +62,12 @@ def main():
         "llava-hf/llava-1.5-7b-hf",
         torch_dtype=torch.bfloat16,
         device_map=compat.device_type,
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
+    )
+    processor = AutoProcessor.from_pretrained(
+        "llava-hf/llava-1.5-7b-hf",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
     )
-    processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
 
     # Load dataset and create DataLoader
     dataset = load_dataset("HuggingFaceM4/the_cauldron", "aokvqa")["train"]

benchmarks/llava/prepare.py

@@ -11,8 +11,12 @@ def main():
         "llava-hf/llava-1.5-7b-hf",
         torch_dtype=torch.float32,  # Change to float32
         device_map="auto",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
+    )
+    _ = AutoProcessor.from_pretrained(
+        "llava-hf/llava-1.5-7b-hf",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
     )
-    _ = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
 
     # Load dataset and create DataLoader
     _ = load_dataset("HuggingFaceM4/the_cauldron", "aokvqa")["train"]

milabench/sizer.py

@@ -53,9 +53,10 @@ def to_octet(value: str) -> float:
 class Sizer:
     """Automatically scale the batch size to match GPU spec"""
 
-    def __init__(self, scaling_config=None):
+    def __init__(self, sizer=None, scaling_config=option("sizer.config", etype=str)):
         self.path = scaling_config
-
+        self.sizer_override = sizer
+
         if scaling_config is None:
             scaling_config = default_scaling_config
 
@@ -64,6 +65,8 @@ def __init__(self, scaling_config=None):
 
     @property
     def options(self):
+        if self.sizer_override:
+            return self.sizer_override
         return SizerOptions()
 
     def benchscaling(self, benchmark):
@@ -252,7 +255,7 @@ class MemoryUsageExtractor(ValidationLayer):
     def __init__(self):
 
         self.filepath = option("sizer.save", str, None)
-        sizer = batch_sizer()
+        sizer = Sizer()
         self.memory = deepcopy(sizer.scaling_config)
         self.scaling = None
         self.benchname = None

milabench/system.py

@@ -121,6 +121,11 @@ def apply_system(config: dict):
     system = system_global.get()
     old = deepcopy(system)
 
+    if system is None:
+        system = dict()
+        system_global.set(system)
+        system = system_global.get()
+
     for k, v in config.items():
         frags = k.split(".")
 

tests/config/argerror.yaml

@@ -9,3 +9,5 @@ benchio:
     n: 1
   argv:
     --start: 0
+  tags:
+    - monogpu

tests/config/benchio.yaml

@@ -4,4 +4,6 @@ benchio:
   weight: 2
   plan:
     method: njobs
-    n: 2
+    n: 2
+  tags:
+    - monogpu

tests/config/benchio_bad.yaml

@@ -8,4 +8,5 @@ benchio:
 
   argv:
     --bad: true
-
+  tags:
+    - monogpu

tests/config/scaling.yaml

@@ -5,3 +5,5 @@ benchio:
     64: 12Go
     128: 24Go
     256: 48Go
+  tags:
+    - monogpu

tests/test_scaler.py

@@ -76,28 +76,27 @@ def fakeexec(pack):
 
 def test_scaler_enabled(multipack, config):
     from milabench.system import system_global
-    import contextvars
-
-    ctx = contextvars.copy_context()
-
-    def update_ctx():
-        sizer = Sizer(
-            SizerOptions(
-                size=None,
-                autoscale=True,
-                multiple=8,
-            ),
-            config("scaling"),
-        )
-        sizer_global.set(sizer)
-        system = system_global.get()
-        gpu = system.setdefault("gpu", dict())
-        gpu["capacity"] = "41920 MiB"
-
-    ctx.run(update_ctx)
+    from milabench.system import apply_system
+
+    conf = {
+        "gpu": {
+            "capacity": "41920 MiB"
+        },
+        "options": {
+            "sizer": {
+                "multiple": 8
+            }
+        }
+    }
 
     for k, pack in multipack.packs.items():
-        assert ctx.run(lambda: fakeexec(pack)) == ["--batch_size", "232"]
+        # Sizer is only enabled when config is applied
+        assert fakeexec(pack) == []
+
+    with apply_system(conf):
+        for k, pack in multipack.packs.items():
+            fakeexec(pack) == ["--batch_size", "232"]
 
-        # Sizer is only enabled inside the context
+    for k, pack in multipack.packs.items():
+        # Sizer is only enabled when config is applied
         assert fakeexec(pack) == []

tests/test_summary/test_compare.txt

@@ -1,5 +1,5 @@
                                        |   rijubigo |   sedumoje
-                                       | 2023-03-24 | 2023-03-24
-bench                |          metric |   13:45:27 |   13:57:35
+                                       | 2024-08-23 | 2024-08-23
+bench                |          metric |   09:22:03 |   09:22:03
 ----------------------------------------------------------------
 benchio              |      train_rate |    8780.41 |    8286.03

tests/test_summary/test_report.txt

@@ -11,4 +11,4 @@ benchio |    0 |   4 |    0 |    7979.82 |   2.9% |  17.2% |         nan |    79
 Scores
 ------
 Failure rate:       0.00% (PASS)
-Score:            7979.82
+Score:            7980.82

tests/test_summary/test_report_folder_does_average.txt

@@ -11,4 +11,4 @@ benchio |    0 |   6 |    0 |    7878.45 |   2.5% |  18.0% |       24456 |    78
 Scores
 ------
 Failure rate:       0.00% (PASS)
-Score:            7878.45
+Score:            7879.45

tests/test_validation.py

@@ -76,39 +76,46 @@ def test_planning_layer_per_gpu_bad(replayfolder, monkeypatch):
 
 def test_memory_tracking(replayfolder, config, tmp_path):
     import contextvars
-
-    from milabench.sizer import (
-        MemoryUsageExtractor,
-        Sizer,
-        SizerOptions,
-        sizer_global,
-        system_global,
-    )
-
-    ctx = contextvars.copy_context()
-
-    def update_ctx():
-        sizer = Sizer(
-            SizerOptions(
-                size=None,
-                autoscale=True,
-                multiple=8,
-            ),
-            config("scaling"),
+    import yaml
+    from milabench.system import apply_system, option
+
+    conf = {
+        "gpu": {
+            "capacity": "41920 MiB"
+        },
+        "options": {
+            "sizer": {
+                "multiple": 8,
+                "autoscale": 1
+            }
+        }
+    }
+
+    with apply_system(conf):
+        from milabench.sizer import (
+            MemoryUsageExtractor,
+            Sizer,
+            SizerOptions,
+            sizer_global,
+            system_global,
         )
-        sizer_global.set(sizer)
-        system_global.set({"gpu": {"capacity": "41920 MiB"}})
-
-    ctx.run(update_ctx)
-    layer = ctx.run(lambda: MemoryUsageExtractor())
-
-    layer.filepath = f"{tmp_path}/dummy"
-
-    assert 123 not in layer.memory["benchio"]["model"]
-
-    ctx.run(lambda: replay_validation_scenario(replayfolder, layer, filename="usage"))
-
-    assert 123 in layer.memory["benchio"]["model"]
+
+        layer = MemoryUsageExtractor()
+        with open(config("scaling"), "r") as sconf:
+            layer.memory = yaml.safe_load(sconf)
+
+        layer.filepath = f"{tmp_path}/dummy"
+
+        print(system_global.get())
+        # print(option("sizer.multiple", etype=int))
+        # print(option("sizer.config", etype=str))
+        # print(Sizer().scaling_config)
+        assert 123 not in layer.memory["benchio"]["model"]
+
+        replay_validation_scenario(replayfolder, layer, filename="usage")
+
+        # print(layer.memory)
+        assert 123 in layer.memory["benchio"]["model"]
 
 
 def test_exception_tracking(replayfolder, file_regression, capsys):