Revert bad commit

benchmarks/diffusion/main.py

@@ -57,11 +57,6 @@ def models(accelerator, args: Arguments):
     unet = UNet2DConditionModel.from_pretrained(
         args.model, subfolder="unet", revision=args.revision, variant=args.variant
     )
-
-    from benchmate.models import model_size
-    print(model_size(unet))
-    print(model_size(encoder))
-    print(model_size(vae))
 
     vae.requires_grad_(False)
     encoder.requires_grad_(False)

benchmarks/flops/benchfile.py

@@ -8,15 +8,11 @@ class FlopsBenchmarch(Package):
 
     def build_run_plan(self) -> "Command":
         import milabench.commands as cmd
-        main = self.dirs.code / self.main_script
+
         pack = cmd.PackCommand(self, *self.argv, lazy=True)
-
-        use_stdout = True
-
-        if use_stdout:
-            return pack.use_stdout()
-        else:
-            pack = cmd.VoirCommand(pack, cwd=main.parent)
-            return pack
+        # pack = cmd.VoirCommand(pack, cwd=main.parent)
+        pack = cmd.ActivatorCommand(pack)
+        return pack.use_stdout()
+
 
 __pack__ = FlopsBenchmarch

benchmarks/flops/main.py

@@ -109,12 +109,11 @@ def main():
 
     log, monitor = setupvoir()
 
-    # FIXME
-    #with monitor:
     f(args.number, args.repeat, args.m, args.n, TERA, dtypes[args.dtype], log)
 
     monitor.stop()
 
+
 if __name__ == "__main__":
     main()
     print("done")

benchmarks/purejaxrl/dqn.py

@@ -98,11 +98,6 @@ def train(rng):
         init_x = jnp.zeros(env.observation_space(env_params).shape)
         network_params = network.init(_rng, init_x)
 
-
-        param_count = sum(x.size for x in jax.tree.leaves(network_params))
-        print("PARAM COUNT", param_count)
-
-
         def linear_schedule(count):
             frac = 1.0 - (count / config["NUM_UPDATES"])
             return config["LR"] * frac

benchmarks/purejaxrl/ppo.py

@@ -107,15 +107,9 @@ def train(rng):
         network = ActorCritic(
             env.action_space(env_params).shape[0], activation=config["ACTIVATION"]
         )
-
-
         rng, _rng = jax.random.split(rng)
         init_x = jnp.zeros(env.observation_space(env_params).shape)
         network_params = network.init(_rng, init_x)
-
-        param_count = sum(x.size for x in jax.tree.leaves(network_params))
-        print("PARAM COUNT", param_count)
-
         if config["ANNEAL_LR"]:
             tx = optax.chain(
                 optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),

benchmarks/recursiongfn/main.py

@@ -92,13 +92,11 @@ def __init__(
         self.num_cond_dim = self.temperature_conditional.encoding_size()
 
     def _load_task_models(self):
-        xdg_cache = os.environ.get("XDG_CACHE_HOME")
+        xdg_cache = os.environ["XDG_CACHE_HOME"]
         model = bengio2021flow.load_original_model(
             cache=True,
             location=Path(os.path.join(xdg_cache, "bengio2021flow_proxy.pkl.gz")),
         )
-        from benchmate.models import model_size
-        print(model_size(model))
         model.to(get_worker_device())
         model = self._wrap_model(model)
         return {"seh": model}

benchmarks/torchatari/main.py

@@ -201,10 +201,7 @@ def main():
     envs = RecordEpisodeStatistics(envs)
     assert isinstance(envs.action_space, gym.spaces.Discrete), "only discrete action space is supported"
 
-
-    from benchmate.models import model_size
     agent = Agent(envs).to(device)
-    print(model_size(agent))
     optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)
 
     # ALGO Logic: Storage setup

benchmate/benchmate/monitor.py

@@ -13,30 +13,7 @@
 from voir.instruments.io import io_monitor
 from voir.instruments.network import network_monitor
 from voir.instruments.monitor import monitor
-from voir.helpers import current_overseer
 
-from .metrics import sumggle_push, give_push, file_push
-
-
-def auto_push():
-    # use_stdout = int(os.getenv("MILABENCH_USE_STDOUT", 0))
-    mb_managed = int(os.getenv("MILABENCH_MANAGED", 0))
-
-    # Milabench managed: we need to push metrics to it
-    if mb_managed == 1:
-        # Using voir, DATA_FD is defined as well
-        ov = current_overseer.get()
-        if ov is not None:
-            return ov.give
-
-        # Not using Voir, using structured stdout
-        if int(os.getenv("MILABENCH_USE_STDOUT", 0)) == 1:
-            return sumggle_push()
-
-        raise RuntimeError("Could not find something to push to")
-
-    # Not using milabench; using stdout
-    return file_push()
 
 
 @instrument_definition
@@ -64,10 +41,16 @@ def monitor_node(ov, poll_interval=1, arch=None):
 
 
 def _smuggle_monitor(poll_interval=10, worker_init=None, **monitors):
-    log = auto_push()
-
+    data_file = SmuggleWriter(sys.stdout)
     def mblog(data):
-        log(**data)
+        nonlocal data_file
+
+        if data_file is not None:
+            try:
+                print(json.dumps(data), file=data_file)
+            except ValueError:
+                pass
+                # print("Is bench ending?, ignoring ValueError")
 
     def get():
         t = time.time()

milabench/_version.py

@@ -1,5 +1,5 @@
 """This file is generated, do not modify"""
 
-__tag__ = "v0.1.0-129-ga60a3aa"
-__commit__ = "a60a3aae21e87e46bcce403620a3f56c12878554"
-__date__ = "2024-11-06 22:52:12 -0500"
+__tag__ = "v1.0.0_RC1-9-g6d1e1140"
+__commit__ = "6d1e114000cc4200ea307330032234db6696e40d"
+__date__ = "2024-09-30 14:39:43 -0400"

milabench/commands/executors.py

@@ -32,9 +32,6 @@ async def execute(pack, *args, cwd=None, env={}, external=False, use_stdout=Fals
     sized_args = scale_argv(pack, args)
     final_args = resolve_argv(pack, sized_args)
 
-    if use_stdout:
-        exec_env["MILABENCH_USE_STDOUT"] = "1"
-
     return await run(
         final_args,
         **kwargs,

milabench/pack.py

@@ -335,8 +335,6 @@ def make_env(self):
             f"MILABENCH_DIR_{name.upper()}": path
             for name, path in self.config["dirs"].items()
         }
-
-        env["MILABENCH_MANAGED"] = "1"
 
         env["OMP_NUM_THREADS"] = resolve_placeholder(self, "{cpu_per_gpu}")
 

milabench/report.py

@@ -525,12 +525,12 @@ def pandas_to_string(df, formatters=_formatters):
     # Compute column size
     col_size = defaultdict(int)
     for index, row in df.iterrows():
-        col_size["bench"] = max(col_size["bench"], len(index), len("bench"))
+        col_size["bench"] = max(col_size["bench"], len(index))
         for col, val in zip(columns, row):
             fmt = formatters.get(col)
             if fmt is not None:
                 val = fmt(val)
-                col_size[col] = max(col_size[col], len(val), len(col))
+                col_size[col] = max(col_size[col], len(val))
 
     # Generate report
     sep = " | "