Specify Device and Environment (#4)

README.md

@@ -12,12 +12,35 @@ $ which construe
 $ construe --help
 ```
 
+There are several top-level configurations that you can specify either as an environment variable or a command line option before the command. The environment variables are as follows:
+
+- `$CONSTRUE_ENV` or `$ENV`: specify the name of the experimental environment for comparison purposes.
+- `$CONSTRUE_DEVICE` or `$TORCH_DEVICE`: specify the name of the default device to use with PyTorch e.g. cpu, mps, or cuda.
+
+The command line utility help is as follows:
+
+```
+Usage: construe [OPTIONS] COMMAND [ARGS]...
+
+Options:
+  --version          Show the version and exit.
+  -d, --device TEXT  specify the pytorch device to run on e.g. cpu, mps or
+                     cuda
+  -e, --env TEXT     name of the experimental environment for comparison
+                     (default is hostname)
+  -h, --help         Show this message and exit.
+
+Commands:
+  basic
+  moondream
+```
+
 ## Basic Benchmarks
 
 The basic benchmarks implement dot product benchmarks from the [PyTorch documentation](https://pytorch.org/tutorials/recipes/recipes/benchmark.html). These benchmarks can be run using `construe basic`; for example by running:
 
 ```
-$ construe basic -e "MacBook Pro 2022 M1" -o results-macbook.pickle
+$ construe -e "MacBook Pro 2022 M1" basic -o results-macbook.pickle
 ```
 
 The `-e` flag specifies the environment for comparison purposes and the `-o` flag saves the measurements out to disk as a Pickle file that can be loaded for comparison to other environments later.

construe/__main__.py

@@ -3,7 +3,10 @@
 """
 
 import click
+import torch
+import platform
 
+from click import ClickException
 from .version import get_version
 from .basic import BasicBenchmark
 from .moondream import MoonDreamBenchmark
@@ -16,17 +19,40 @@
 
 @click.group(context_settings=CONTEXT_SETTINGS)
 @click.version_option(get_version(), message="%(prog)s v%(version)s")
-def main():
-    pass
-
-
-@main.command()
+@click.option(
+    "-d",
+    "--device",
+    default=None,
+    type=str,
+    envvar=["CONSTRUE_DEVICE", "TORCH_DEVICE"],
+    help="specify the pytorch device to run on e.g. cpu, mps or cuda",
+)
 @click.option(
     "-e",
     "--env",
     default=None,
+    envvar=["CONSTRUE_ENV", "ENV"],
     help="name of the experimental environment for comparison (default is hostname)",
 )
+@click.pass_context
+def main(ctx, env=None, device=None):
+    if device is not None:
+        try:
+            torch.set_default_device(device)
+        except RuntimeError as e:
+            raise ClickException(str(e))
+
+        click.echo(f"using torch.device(\"{device}\")")
+
+    if env is None:
+        env = platform.node()
+
+    ctx.ensure_object(dict)
+    ctx.obj["device"] = device
+    ctx.obj["env"] = env
+
+
+@main.command()
 @click.option(
     "-o",
     "--saveto",
@@ -53,18 +79,23 @@ def main():
     type=int,
     help="set the random seed for random generation",
 )
-def basic(**kwargs):
+@click.pass_context
+def basic(ctx, **kwargs):
+    kwargs["env"] = ctx.obj["env"]
     benchmark = BasicBenchmark(**kwargs)
     benchmark.run()
 
 
 @main.command()
-def moondream(**kwargs):
+@click.pass_context
+def moondream(ctx, **kwargs):
+    kwargs["env"] = ctx["env"]
     benchmark = MoonDreamBenchmark(**kwargs)
     benchmark.run()
 
 
 if __name__ == "__main__":
     main(
+        obj={},
         prog_name="construe",
     )

construe/basic.py

@@ -4,9 +4,9 @@
 See: https://pytorch.org/tutorials/recipes/recipes/benchmark.html
 """
 
+import tqdm
 import torch
 import pickle
-import platform
 import torch.utils.benchmark as benchmark
 
 from itertools import product
@@ -35,9 +35,6 @@ def __init__(self, env=None, saveto=None, num_threads=None, fuzz=False, seed=Non
         if num_threads is None:
             num_threads = torch.get_num_threads()
 
-        if env is None:
-            env = platform.node()
-
         self.env = env
         self.saveto = saveto
         self.num_threads = num_threads
@@ -46,15 +43,15 @@ def __init__(self, env=None, saveto=None, num_threads=None, fuzz=False, seed=Non
 
     def run(self):
         results = []
-        dataset = self.fuzzer().take(10) if self.fuzz else self.static()
+        dataset = self.fuzzer().take(10) if self.fuzz else list(self.static())
 
         kwargs = {
             "label": "Batched Dot",
             "num_threads": self.num_threads,
             "env": self.env,
         }
 
-        for tensors, tensor_params, params in dataset:
+        for tensors, tensor_params, params in tqdm.tqdm(dataset, leave=False):
             sub_label = f"{params['k0']:<6} x {params['k1']:<4} {'' if tensor_params['x']['is_contiguous'] else '(discontiguous)'}"  # noqa
             results.append(
                 benchmark.Timer(

construe/moondream.py

@@ -27,7 +27,7 @@ def __init__(self):
     @profile
     def run(self):
         results = []
-        for path in tqdm.tqdm(self.dataset):
+        for path in tqdm.tqdm(self.dataset, leave=False):
             encoded, encode_time = self.moondream.encode_image(path)
             inference_time = self.moondream.inference(encoded)
             results.append((encode_time, inference_time))