mirror of
https://github.com/XuehaiPan/nvitop.git
synced 2026-05-15 14:15:55 -06:00
111 lines
3.1 KiB
Python
Executable file
111 lines
3.1 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
# To install `nvitop` with specific version of `nvidia-ml-py`, use:
|
|
#
|
|
# pip install nvidia-ml-py==xx.yyy.zz nvitop
|
|
#
|
|
# or
|
|
#
|
|
# pip install 'nvitop[cudaXX]'
|
|
#
|
|
|
|
"""Setup script for ``nvitop``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import re
|
|
import sys
|
|
from importlib.util import module_from_spec, spec_from_file_location
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
from setuptools import setup
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Generator
|
|
from types import ModuleType
|
|
|
|
|
|
HERE = Path(__file__).absolute().parent
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def vcs_version(name: str, path: Path | str) -> Generator[ModuleType]:
|
|
"""Context manager to update version string in a version module."""
|
|
path = Path(path).absolute()
|
|
assert path.is_file()
|
|
module_spec = spec_from_file_location(name=name, location=path)
|
|
assert module_spec is not None
|
|
assert module_spec.loader is not None
|
|
module = sys.modules.get(name)
|
|
if module is None:
|
|
module = module_from_spec(module_spec)
|
|
sys.modules[name] = module
|
|
module_spec.loader.exec_module(module)
|
|
|
|
if module.__release__:
|
|
yield module
|
|
return
|
|
|
|
content = None
|
|
try:
|
|
try:
|
|
content = path.read_text(encoding='utf-8')
|
|
path.write_text(
|
|
data=re.sub(
|
|
r"""__version__\s*=\s*('[^']+'|"[^"]+")""",
|
|
f'__version__ = {module.__version__!r}',
|
|
string=content,
|
|
),
|
|
encoding='utf-8',
|
|
)
|
|
except OSError:
|
|
content = None
|
|
|
|
yield module
|
|
finally:
|
|
if content is not None:
|
|
with path.open(mode='wt', encoding='utf-8', newline='') as file:
|
|
file.write(content)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
extra_requirements = {
|
|
'exporter': ['nvitop-exporter'],
|
|
'prometheus': ['nvitop-exporter'],
|
|
'lint': [
|
|
'ruff',
|
|
'pylint[spelling]',
|
|
'xdoctest',
|
|
'mypy',
|
|
'typing-extensions',
|
|
'pre-commit',
|
|
],
|
|
'cuda10': ['nvidia-ml-py == 11.450.51'],
|
|
}
|
|
|
|
with vcs_version(name='nvitop.version', path=HERE / 'nvitop' / 'version.py') as version:
|
|
for pynvml_major in sorted(
|
|
{int(pynvml.partition('.')[0]) for pynvml in version.PYNVML_VERSION_CANDIDATES},
|
|
):
|
|
pynvml_range = [
|
|
pynvml
|
|
for pynvml in version.PYNVML_VERSION_CANDIDATES
|
|
if pynvml.startswith(f'{pynvml_major}.')
|
|
]
|
|
if len(pynvml_range) == 1:
|
|
extra_requirements[f'cuda{pynvml_major}'] = [
|
|
f'nvidia-ml-py == {pynvml_range[0]}',
|
|
]
|
|
elif len(pynvml_range) >= 2:
|
|
extra_requirements[f'cuda{pynvml_major}'] = [
|
|
f'nvidia-ml-py >= {pynvml_range[0]}, <= {pynvml_range[-1]}',
|
|
]
|
|
|
|
setup(
|
|
name='nvitop',
|
|
version=version.__version__,
|
|
extras_require=extra_requirements,
|
|
)
|