1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
|
#!/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[pynvml-xx.yyy.zz]'
#
"""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 = {
'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]}',
]
extra_requirements.update(
{
# The identifier could not start with numbers, add a prefix `pynvml-`
f'pynvml-{pynvml}': [f'nvidia-ml-py == {pynvml}']
for pynvml in version.PYNVML_VERSION_CANDIDATES
},
)
setup(
name='nvitop',
version=version.__version__,
extras_require=extra_requirements,
)
|
