Refactor (#179)

* Refactor

* Ruff

* Add missing updates

---------

Co-authored-by: pierre.delaunay <[email protected]>

benchmarks/accelerate_opt/benchfile.py

@@ -1,9 +1,9 @@
-from milabench.executors import (
-    AccelerateLaunchExecutor,
-    CmdExecutor,
-    DockerRunExecutor,
-    ListExecutor,
-    SSHExecutor,
+from milabench.commands import (
+    AccelerateLaunchCommand,
+    CmdCommand,
+    DockerRunCommand,
+    ListCommand,
+    SSHCommand,
 )
 from milabench.pack import Package
 from milabench.utils import select_nodes
@@ -18,7 +18,7 @@ def make_env(self):
         return env
 
     def build_prepare_plan(self):
-        return CmdExecutor(
+        return CmdCommand(
             self,
             "accelerate",
             "launch",
@@ -59,19 +59,19 @@ def build_run_plan(self):
                 tags.append("nolog")
 
             pack = self.copy({"tag": tags})
-            worker = SSHExecutor(
+            worker = SSHCommand(
                 host=host,
                 user=user,
                 key=key,
-                executor=DockerRunExecutor(
-                    AccelerateLaunchExecutor(pack, rank=rank),
+                executor=DockerRunCommand(
+                    AccelerateLaunchCommand(pack, rank=rank),
                     self.config["system"].get("docker_image"),
                 ),
                 **options
             )
             plans.append(worker)
 
-        return ListExecutor(*plans)
+        return ListCommand(*plans)
 
 
 __pack__ = AccelerateBenchmark

benchmarks/accelerate_opt/main.py

@@ -66,20 +66,18 @@ def arguments():
 from dataclasses import dataclass
 from datetime import timedelta
 from itertools import chain
-from pathlib import Path
 from typing import Optional
 
 import datasets
 import rich.logging
 import torch
 import transformers
-from accelerate import Accelerator, DistributedType
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DummyOptim, DummyScheduler
 from accelerate.utils.dataclasses import InitProcessGroupKwargs
 from datasets import load_dataset
 from torch.utils.data import DataLoader
-from tqdm.auto import tqdm
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -88,7 +86,6 @@ def arguments():
     default_data_collator,
     get_scheduler,
 )
-from transformers.utils import get_full_repo_name
 from voir.smuggle import SmuggleWriter
 from voir.instruments.gpu import get_gpu_info
 from voir.instruments.utils import Monitor
@@ -374,7 +371,6 @@ def group_texts(examples):
 
     completed_steps = 0
     starting_epoch = 0
-    best_metric = None
     last_log_time = time.time()
 
     for epoch in range(starting_epoch, num_train_epochs):

benchmarks/flops/benchfile.py

@@ -6,13 +6,12 @@ class FlopsBenchmarch(Package):
     prepare_script = "prepare.py"
     main_script = "main.py"
 
-    def build_run_plan(self) -> "execs.Executor":
-        import milabench.executors as execs
+    def build_run_plan(self) -> "Command":
+        import milabench.commands as cmd
 
-        main = self.dirs.code / self.main_script
-        pack = execs.PackExecutor(self, *self.argv, lazy=True)
-        # pack = execs.VoirExecutor(pack, cwd=main.parent)
-        pack = execs.ActivatorExecutor(pack, use_stdout=True)
+        pack = cmd.PackCommand(self, *self.argv, lazy=True)
+        # pack = cmd.VoirCommand(pack, cwd=main.parent)
+        pack = cmd.ActivatorCommand(pack, use_stdout=True)
         return pack
 
 

benchmarks/llama/benchfile.py

@@ -1,6 +1,5 @@
-import uuid
 
-from milabench.executors import CmdExecutor
+from milabench.commands import CmdCommand
 from milabench.pack import Package
 
 
@@ -18,7 +17,7 @@ async def install(self):
         await super().install()
 
     def build_prepare_plan(self):
-        return CmdExecutor(
+        return CmdCommand(
             self,
             "python",
             str(self.dirs.code / "main.py"),
@@ -29,7 +28,7 @@ def build_prepare_plan(self):
         )
 
     def build_run_plan(self):
-        return CmdExecutor(
+        return CmdCommand(
             self,
             "python",
             str(self.dirs.code / "main.py"),

benchmarks/rwkv/rwkv-v4neo/chat.py

@@ -188,7 +188,7 @@ def load_all_stat(srv, name):
 ########################################################################################################
 
 # Run inference
-print(f"\nRun prompt...")
+print("\nRun prompt...")
 
 out = run_rnn(tokenizer.tokenizer.encode(init_prompt))
 gc.collect()

benchmarks/rwkv/rwkv-v4neo/img_demoAE.py

@@ -2,7 +2,7 @@
 # The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
 ########################################################################################################
 
-import torch, types, os
+import torch, types
 import numpy as np
 from PIL import Image
 import torch.nn as nn
@@ -11,7 +11,7 @@
 import torchvision.transforms as transforms
 
 np.set_printoptions(precision=4, suppress=True, linewidth=200)
-print(f"loading...")
+print("loading...")
 
 ########################################################################################################
 
@@ -138,7 +138,7 @@ def forward(self, code):
 
 ########################################################################################################
 
-print(f"building model...")
+print("building model...")
 args = types.SimpleNamespace()
 args.my_img_bit = 13
 encoder = R_ENCODER(args).eval().cuda()
@@ -151,7 +151,7 @@ def forward(self, code):
 
 ########################################################################################################
 
-print(f"test image...")
+print("test image...")
 img_transform = transforms.Compose(
     [
         transforms.PILToTensor(),

benchmarks/rwkv/rwkv-v4neo/run.py

@@ -3,7 +3,7 @@
 ########################################################################################################
 
 import numpy as np
-import math, os, sys, types, time, gc
+import os, sys, types, time, gc
 import torch
 from src.utils import TOKENIZER
 
@@ -137,7 +137,7 @@
 
 model = RWKV_RNN(args)
 
-print(f"\nOptimizing speed...")
+print("\nOptimizing speed...")
 out, _ = model.forward([187], None)
 # print(out)
 gc.collect()

benchmarks/rwkv/rwkv-v4neo/src/binidx.py

@@ -1,8 +1,6 @@
-from lib2to3.pgen2 import token
 import os
 import torch
 import numpy as np
-import shutil
 import struct
 from functools import lru_cache
 from itertools import accumulate

benchmarks/rwkv/rwkv-v4neo/src/dataset.py

@@ -2,7 +2,7 @@
 # The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
 ########################################################################################################
 
-import json, math, random, os, sys
+import json, math, random
 import numpy as np
 import torch
 from torch.utils.data import Dataset

benchmarks/rwkv/rwkv-v4neo/src/model.py

@@ -10,7 +10,6 @@
 import torch.nn as nn
 from torch.nn import functional as F
 import pytorch_lightning as pl
-from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
 from pytorch_lightning.strategies import DeepSpeedStrategy
 
 if importlib.util.find_spec("deepspeed"):
@@ -713,7 +712,7 @@ def training_step_end(self, batch_parts):
 
     def generate_init_weight(self):
         print(
-            f"""
+            """
 ############################################################################
 #
 # Init model weight (slow for large models)...

benchmarks/rwkv/rwkv-v4neo/src/model_img.py

@@ -3,15 +3,13 @@
 ########################################################################################################
 
 import numpy as np
-import os, math, gc
+import os, gc
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision as vision
 import pytorch_lightning as pl
-from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
 from pytorch_lightning.strategies import DeepSpeedStrategy
-import deepspeed
 from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
 
 # from pytorch_msssim import MS_SSIM
@@ -425,7 +423,7 @@ def training_step_end(self, batch_parts):
 
     def generate_init_weight(self):
         print(
-            f"""
+            """
 ############################################################################
 #
 # Init model weight (slow for large models)...

benchmarks/rwkv/rwkv-v4neo/src/model_run.py

@@ -4,10 +4,9 @@
 
 import types
 import torch
-import math, os, gc
+import os, gc
 from torch.nn import functional as F
 import torch.nn as nn
-from typing import List, Dict
 
 MyModule = nn.Module
 

benchmarks/rwkv/rwkv-v4neo/src/trainer.py

@@ -1,8 +1,7 @@
-import os, math, time, datetime, subprocess
+import math, time, datetime
 import torch
-from torch.utils.data import DataLoader
 import pytorch_lightning as pl
-from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
+from pytorch_lightning.utilities import rank_zero_only
 from giving import give
 
 

benchmarks/rwkv/rwkv-v4neo/train.py

@@ -7,7 +7,7 @@
     import os
     from argparse import ArgumentParser
     from pytorch_lightning import Trainer
-    from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
+    from pytorch_lightning.utilities import rank_zero_info
 
     rank_zero_info("########## work in progress ##########")
 
@@ -152,7 +152,7 @@
 
     ########################################################################################################
 
-    import os, warnings, math, datetime, sys, time, importlib
+    import os, warnings, datetime, importlib
     import numpy as np
     import torch
     from torch.utils.data import DataLoader

benchmarks/stargan/stargan/model.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import numpy as np
 
 

benchmarks/stargan/stargan/solver.py

@@ -1,6 +1,5 @@
 from model import Generator
 from model import Discriminator
-from torch.autograd import Variable
 from torchvision.utils import save_image
 import torch
 import torch.nn.functional as F

benchmarks/stargan/stargan/synth.py

@@ -1,4 +1,3 @@
-import torch
 
 
 class SyntheticData:

benchmarks/super-slomo/slomo/model.py

@@ -1,7 +1,4 @@
 import torch
-import torchvision
-import torchvision.transforms as transforms
-import torch.optim as optim
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np

benchmarks/super-slomo/slomo/train.py

@@ -4,12 +4,10 @@
 import argparse
 import torch
 import torchvision
-import torchvision.transforms as transforms
 import torch.optim as optim
 import torch.nn as nn
 import torch.nn.functional as F
 import model
-import dataloader
 from giving import give
 import voir
 from synth import SyntheticData

benchmarks/timm/benchfile.py

@@ -1,6 +1,5 @@
-import uuid
 
-from milabench.executors import TorchRunExecutor
+from milabench.commands import TorchRunCommand
 from milabench.pack import Package
 
 
@@ -43,7 +42,7 @@ async def install(self):
     def build_run_plan(self):
         # self.config is not the right config for this
         plan = super().build_run_plan()
-        return TorchRunExecutor(plan, use_stdout=True)
+        return TorchRunCommand(plan, use_stdout=True)
 
 
 __pack__ = TimmBenchmarkPack

milabench/cli.py

@@ -3,10 +3,8 @@
 import re
 import io
 import runpy
-import shutil
 import subprocess
 import sys
-import tempfile
 import traceback
 from datetime import datetime
 import getpass