Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions docs/user-guide/extending.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,15 @@ leftmost class in the inheritance list wins.

### Making a mixin CLI-compatible

A mixin with `cli_compatible = True` appears in `--show-mixins` and can be
selected with `--mixin`. That attribute alone is enough for mixins that need
no configuration:
Document CLI compatibility in the docstring, and add it to cli/registry.py

```python
class MBMachineType:
"""Capture the machine architecture (e.g. x86_64, arm64)."""
"""Capture the machine architecture (e.g. x86_64, arm64).

cli_compatible = True
Note:
CLI compatible.
"""

def capture_machine_type(self, bm_data):
import platform
Expand Down
2 changes: 1 addition & 1 deletion microbench/cli/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from microbench.mixins.vcs import MBFileHash, MBGitInfo

# Maps kebab-case CLI name → mixin class.
# Only mixins with ``cli_compatible = True`` should be listed here.
# Only CLI-compatible mixins (noted in their docstring) should be listed here.
MIXIN_REGISTRY: dict = {
'python-info': MBPythonInfo,
'host-info': MBHostInfo,
Expand Down
4 changes: 3 additions & 1 deletion microbench/mixins/gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,11 @@ class MBNvidiaSmi:
nvidia_gpus (tuple): GPU IDs to poll — zero-based indexes, UUIDs,
or PCI bus IDs. GPU UUIDs are recommended (indexes can change
after a reboot). Omit to poll all installed GPUs.

Note:
CLI compatible.
"""

cli_compatible = True
_nvidia_default_attributes = ('gpu_name', 'memory.total')
_nvidia_gpu_regex = _NVIDIA_GPU_REGEX
cli_args = [
Expand Down
4 changes: 2 additions & 2 deletions microbench/mixins/profiling.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@ class MBPeakMemory:
which covers Python objects and most C-extension allocations. Memory
allocated directly via ``malloc`` in C extensions (e.g. some large
NumPy arrays) is not tracked.
"""

cli_compatible = True
CLI compatible.
"""

def capture_peak_memory(self, bm_data):
import tracemalloc
Expand Down
13 changes: 9 additions & 4 deletions microbench/mixins/python.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,10 @@ class MBPythonInfo:

This mixin is included in :class:`MicroBench` by default (Python API)
and in the CLI default mixin set. It supersedes the former ``MBPythonVersion``.
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_python_info(self, bm_data):
python = bm_data.setdefault('python', {})
Expand Down Expand Up @@ -97,9 +98,11 @@ class MBInstalledPackages:
capture_paths (bool): Also record the installation path of each
package under ``python.installed_package_paths``. Defaults to
``False``.

Note:
CLI compatible.
"""

cli_compatible = True
capture_paths = False

def capture_packages(self, bm_data):
Expand Down Expand Up @@ -140,9 +143,11 @@ class MBCondaPackages:
Defaults to ``True``.
include_channels (bool): Include the channel name in the version.
Defaults to ``False``.

Note:
CLI compatible.
"""

cli_compatible = True
include_builds = True
include_channels = False

Expand Down
25 changes: 15 additions & 10 deletions microbench/mixins/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,10 @@ class MBHostInfo:

This mixin supersedes the former ``MBHostCpuCores`` and ``MBHostRamTotal``
mixins, which have been removed.
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_hostname(self, bm_data):
bm_data.setdefault('host', {})['hostname'] = socket.gethostname()
Expand All @@ -52,9 +53,10 @@ class MBWorkingDir:

Records the current working directory as ``call.working_dir``. This is
per-call data since the working directory can change between calls.
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_working_dir(self, bm_data):
bm_data.setdefault('call', {})['working_dir'] = os.getcwd()
Expand All @@ -77,9 +79,10 @@ class MBSlurmInfo:
"cpus_per_task": "4"
}
}
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_slurm(self, bm_data):
bm_data['slurm'] = {
Expand Down Expand Up @@ -108,9 +111,10 @@ class MBLoadedModules:

Module entries without a version (e.g. ``null``) are stored with an
empty string as the version.
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_loaded_modules(self, bm_data):
loaded = os.environ.get('LOADEDMODULES', '')
Expand Down Expand Up @@ -236,9 +240,10 @@ class MBCgroupLimits:
"version": 2
}
}
"""

cli_compatible = True
Note:
CLI compatible.
"""

def capture_cgroup_limits(self, bm_data):
if sys.platform != 'linux':
Expand Down
8 changes: 6 additions & 2 deletions microbench/mixins/vcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,11 @@ class MBGitInfo:
"dirty": false
}
}

Note:
CLI compatible.
"""

cli_compatible = True
cli_args = [
CLIArg(
flags=['--git-repo'],
Expand Down Expand Up @@ -168,9 +170,11 @@ class MBFileHash:
"run_experiment.py": "e3b0c44298fc1c14..."
}
}

Note:
CLI compatible.
"""

cli_compatible = True
cli_args = [
CLIArg(
flags=['--hash-file'],
Expand Down
2 changes: 1 addition & 1 deletion tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ def bad_capture(self, bm_data):


def test_cli_all_flag_includes_all_mixins():
"""--all includes every cli_compatible mixin."""
"""--all includes every mixin in the CLI registry."""
from microbench.__main__ import _get_mixin_map

all_names = set(_get_mixin_map())
Expand Down
Loading