lupine 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- lupine-0.1.0/.gitignore +4 -0
- lupine-0.1.0/PKG-INFO +59 -0
- lupine-0.1.0/README.md +49 -0
- lupine-0.1.0/examples/tensor.py +35 -0
- lupine-0.1.0/lupine/__init__.py +266 -0
- lupine-0.1.0/pyproject.toml +21 -0
- lupine-0.1.0/tests/test_lupine_adapter.py +163 -0
- lupine-0.1.0/uv.lock +963 -0
lupine-0.1.0/.gitignore
ADDED
lupine-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: lupine
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Small PyTorch adapter helpers for LUPINE-backed CUDA devices
|
|
5
|
+
License: MIT
|
|
6
|
+
Requires-Python: >=3.9
|
|
7
|
+
Provides-Extra: torch
|
|
8
|
+
Requires-Dist: torch; extra == 'torch'
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
|
|
11
|
+
# lupine Python adapter
|
|
12
|
+
|
|
13
|
+
This package provides small PyTorch helpers for LUPINE. It intentionally returns
|
|
14
|
+
ordinary `torch.device("cuda:N")` objects so PyTorch continues to use its normal
|
|
15
|
+
CUDA dispatch path while LUPINE handles CUDA driver/NVML calls underneath.
|
|
16
|
+
|
|
17
|
+
Declare all LUPINE hosts before any PyTorch CUDA work:
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
import lupine
|
|
21
|
+
|
|
22
|
+
with lupine.connect(host="<server>:14833") as s:
|
|
23
|
+
device = s.device()
|
|
24
|
+
model = model.to(device)
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
`connect()` loads the LUPINE `libcuda.so.1` from `../build/libcuda.so.1` when
|
|
28
|
+
used from this repository. For an installed package, pass `libcuda=...` or set
|
|
29
|
+
`LUPINE_LIBCUDA` if the library lives somewhere else:
|
|
30
|
+
|
|
31
|
+
```python
|
|
32
|
+
with lupine.connect(host="<server>:14833", libcuda="/opt/lupine/libcuda.so.1") as s:
|
|
33
|
+
device = s.device()
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
For multiple LUPINE servers, pass the full host list in one call. The order
|
|
37
|
+
defines the CUDA ordinals that PyTorch sees:
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
import lupine
|
|
41
|
+
|
|
42
|
+
with lupine.connect(host=["<server-a>:14833", "<server-b>:14833"]) as s:
|
|
43
|
+
gpu0, gpu1 = s.devices()
|
|
44
|
+
model0 = model0.to(gpu0) # cuda:0
|
|
45
|
+
model1 = model1.to(gpu1) # cuda:1
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
Do not add a second host after tensors have already been moved to the first one.
|
|
49
|
+
LUPINE opens connections from `LUPINE_SERVER` when CUDA first initializes, and
|
|
50
|
+
later changes to `LUPINE_SERVER` are not picked up by the current process.
|
|
51
|
+
|
|
52
|
+
Exiting the context restores `LUPINE_SERVER` only if CUDA was not initialized
|
|
53
|
+
inside the block. If CUDA was initialized, the process-global LUPINE connection
|
|
54
|
+
is already active and cannot be disconnected safely.
|
|
55
|
+
|
|
56
|
+
The adapter does not create a new PyTorch backend such as
|
|
57
|
+
`torch.device("lupine")`. A true custom PyTorch device would require registering
|
|
58
|
+
PrivateUse1 kernels and backend support. LUPINE already works best when PyTorch
|
|
59
|
+
sees CUDA tensors and the LUPINE library is selected through the dynamic linker.
|
lupine-0.1.0/README.md
ADDED
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
# lupine Python adapter
|
|
2
|
+
|
|
3
|
+
This package provides small PyTorch helpers for LUPINE. It intentionally returns
|
|
4
|
+
ordinary `torch.device("cuda:N")` objects so PyTorch continues to use its normal
|
|
5
|
+
CUDA dispatch path while LUPINE handles CUDA driver/NVML calls underneath.
|
|
6
|
+
|
|
7
|
+
Declare all LUPINE hosts before any PyTorch CUDA work:
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
import lupine
|
|
11
|
+
|
|
12
|
+
with lupine.connect(host="<server>:14833") as s:
|
|
13
|
+
device = s.device()
|
|
14
|
+
model = model.to(device)
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
`connect()` loads the LUPINE `libcuda.so.1` from `../build/libcuda.so.1` when
|
|
18
|
+
used from this repository. For an installed package, pass `libcuda=...` or set
|
|
19
|
+
`LUPINE_LIBCUDA` if the library lives somewhere else:
|
|
20
|
+
|
|
21
|
+
```python
|
|
22
|
+
with lupine.connect(host="<server>:14833", libcuda="/opt/lupine/libcuda.so.1") as s:
|
|
23
|
+
device = s.device()
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
For multiple LUPINE servers, pass the full host list in one call. The order
|
|
27
|
+
defines the CUDA ordinals that PyTorch sees:
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
import lupine
|
|
31
|
+
|
|
32
|
+
with lupine.connect(host=["<server-a>:14833", "<server-b>:14833"]) as s:
|
|
33
|
+
gpu0, gpu1 = s.devices()
|
|
34
|
+
model0 = model0.to(gpu0) # cuda:0
|
|
35
|
+
model1 = model1.to(gpu1) # cuda:1
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
Do not add a second host after tensors have already been moved to the first one.
|
|
39
|
+
LUPINE opens connections from `LUPINE_SERVER` when CUDA first initializes, and
|
|
40
|
+
later changes to `LUPINE_SERVER` are not picked up by the current process.
|
|
41
|
+
|
|
42
|
+
Exiting the context restores `LUPINE_SERVER` only if CUDA was not initialized
|
|
43
|
+
inside the block. If CUDA was initialized, the process-global LUPINE connection
|
|
44
|
+
is already active and cannot be disconnected safely.
|
|
45
|
+
|
|
46
|
+
The adapter does not create a new PyTorch backend such as
|
|
47
|
+
`torch.device("lupine")`. A true custom PyTorch device would require registering
|
|
48
|
+
PrivateUse1 kernels and backend support. LUPINE already works best when PyTorch
|
|
49
|
+
sees CUDA tensors and the LUPINE library is selected through the dynamic linker.
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
# /// script
|
|
2
|
+
# dependencies = [
|
|
3
|
+
# "numpy",
|
|
4
|
+
# "lupine",
|
|
5
|
+
# "torch",
|
|
6
|
+
# ]
|
|
7
|
+
#
|
|
8
|
+
# [tool.uv.sources]
|
|
9
|
+
# lupine = { path = "..", editable = true }
|
|
10
|
+
# ///
|
|
11
|
+
|
|
12
|
+
import torch
|
|
13
|
+
import lupine
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def prompt_endpoint() -> str:
|
|
17
|
+
host = input("LUPINE server host: ").strip()
|
|
18
|
+
while not host:
|
|
19
|
+
host = input("LUPINE server host: ").strip()
|
|
20
|
+
|
|
21
|
+
port_text = input("LUPINE server port [14833]: ").strip() or "14833"
|
|
22
|
+
port = int(port_text)
|
|
23
|
+
return f"{host}:{port}"
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
with lupine.connect(host=prompt_endpoint()) as session:
|
|
27
|
+
device = session.device()
|
|
28
|
+
props = torch.cuda.get_device_properties(device)
|
|
29
|
+
x = torch.arange(8, device=device, dtype=torch.float32)
|
|
30
|
+
y = (x * 2).cpu()
|
|
31
|
+
print("cuda available:", torch.cuda.is_available())
|
|
32
|
+
print("device:", device)
|
|
33
|
+
print("count:", torch.cuda.device_count())
|
|
34
|
+
print("gpu:", props.name)
|
|
35
|
+
print("result:", y.tolist())
|
|
@@ -0,0 +1,266 @@
|
|
|
1
|
+
"""PyTorch adapter helpers for LUPINE.
|
|
2
|
+
|
|
3
|
+
The adapter returns ordinary ``torch.device("cuda:N")`` objects. PyTorch stays
|
|
4
|
+
on its built-in CUDA dispatch path while LUPINE handles CUDA driver calls below
|
|
5
|
+
it.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import os
|
|
11
|
+
import ctypes
|
|
12
|
+
from collections.abc import Sequence
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
from pathlib import Path
|
|
15
|
+
from typing import Any
|
|
16
|
+
|
|
17
|
+
DEFAULT_PORT = 14833
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class LupineError(RuntimeError):
|
|
21
|
+
"""Raised when the LUPINE adapter cannot select a usable device."""
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _torch() -> Any:
|
|
25
|
+
try:
|
|
26
|
+
import torch
|
|
27
|
+
except ModuleNotFoundError as exc:
|
|
28
|
+
raise LupineError("PyTorch is required to use LUPINE devices.") from exc
|
|
29
|
+
return torch
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _cuda_initialized() -> bool:
|
|
33
|
+
try:
|
|
34
|
+
return bool(_torch().cuda.is_initialized())
|
|
35
|
+
except LupineError:
|
|
36
|
+
return False
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _require_mutable_config() -> None:
|
|
40
|
+
if _cuda_initialized():
|
|
41
|
+
raise LupineError("connect to LUPINE before PyTorch initializes CUDA")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _normalize_server(host: str, port: int | None = None) -> str:
|
|
45
|
+
host = str(host).strip()
|
|
46
|
+
if not host:
|
|
47
|
+
raise LupineError("host must not be empty")
|
|
48
|
+
if port is not None:
|
|
49
|
+
return f"{host}:{int(port)}"
|
|
50
|
+
if host.startswith("[") and "]:" in host:
|
|
51
|
+
return host
|
|
52
|
+
if host.count(":") == 1:
|
|
53
|
+
return host
|
|
54
|
+
return f"{host}:{DEFAULT_PORT}"
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _normalize_hosts(host: str | Sequence[str], port: int | None = None) -> tuple[str, ...]:
|
|
58
|
+
if isinstance(host, str):
|
|
59
|
+
servers = (_normalize_server(host, port),)
|
|
60
|
+
else:
|
|
61
|
+
servers = tuple(_normalize_server(item, port) for item in host)
|
|
62
|
+
if not servers:
|
|
63
|
+
raise LupineError("at least one LUPINE host is required")
|
|
64
|
+
if len(set(servers)) != len(servers):
|
|
65
|
+
raise LupineError("LUPINE hosts must be unique")
|
|
66
|
+
return servers
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _set_server_env(servers: Sequence[str]) -> None:
|
|
70
|
+
os.environ["LUPINE_SERVER"] = ",".join(servers)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _default_libcuda() -> Path | None:
|
|
74
|
+
override = os.environ.get("LUPINE_LIBCUDA")
|
|
75
|
+
if override:
|
|
76
|
+
return Path(override)
|
|
77
|
+
repo_candidate = Path(__file__).resolve().parents[2] / "build" / "libcuda.so.1"
|
|
78
|
+
if repo_candidate.exists():
|
|
79
|
+
return repo_candidate
|
|
80
|
+
return None
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _load_libcuda(path: str | os.PathLike[str] | None) -> None:
|
|
84
|
+
libcuda = Path(path) if path is not None else _default_libcuda()
|
|
85
|
+
if libcuda is None:
|
|
86
|
+
return
|
|
87
|
+
if not libcuda.exists():
|
|
88
|
+
raise LupineError(f"LUPINE libcuda does not exist: {libcuda}")
|
|
89
|
+
ctypes.CDLL(str(libcuda), mode=ctypes.RTLD_GLOBAL)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def _servers_from_env() -> tuple[str, ...]:
|
|
93
|
+
value = os.environ.get("LUPINE_SERVER", "")
|
|
94
|
+
return tuple(server.strip() for server in value.split(",") if server.strip())
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _cuda_device(index: int, *, require_available: bool = False) -> Any:
|
|
98
|
+
torch = _torch()
|
|
99
|
+
index = int(index)
|
|
100
|
+
if require_available:
|
|
101
|
+
count = int(torch.cuda.device_count())
|
|
102
|
+
if count <= 0:
|
|
103
|
+
raise LupineError(
|
|
104
|
+
"PyTorch does not see any CUDA devices. Check that the LUPINE "
|
|
105
|
+
"client library is selected and LUPINE_SERVER is configured."
|
|
106
|
+
)
|
|
107
|
+
if index < 0 or index >= count:
|
|
108
|
+
raise LupineError(f"CUDA device index {index} is out of range for {count} devices")
|
|
109
|
+
return torch.device("cuda", index)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
@dataclass
|
|
113
|
+
class Session:
|
|
114
|
+
"""A process-local LUPINE connection declaration."""
|
|
115
|
+
|
|
116
|
+
servers: tuple[str, ...]
|
|
117
|
+
require_available: bool = False
|
|
118
|
+
libcuda: str | os.PathLike[str] | None = None
|
|
119
|
+
|
|
120
|
+
def __post_init__(self) -> None:
|
|
121
|
+
if not self.servers:
|
|
122
|
+
raise LupineError("at least one LUPINE host is required")
|
|
123
|
+
|
|
124
|
+
def __enter__(self) -> "Session":
|
|
125
|
+
_require_mutable_config()
|
|
126
|
+
self._previous_server = os.environ.get("LUPINE_SERVER")
|
|
127
|
+
configured = _servers_from_env()
|
|
128
|
+
if configured and configured != self.servers:
|
|
129
|
+
raise LupineError(
|
|
130
|
+
"LUPINE_SERVER is already configured differently; start a new "
|
|
131
|
+
"process or pass the same hosts to lupine.connect()."
|
|
132
|
+
)
|
|
133
|
+
if not configured:
|
|
134
|
+
_set_server_env(self.servers)
|
|
135
|
+
_load_libcuda(self.libcuda)
|
|
136
|
+
return self
|
|
137
|
+
|
|
138
|
+
def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
|
|
139
|
+
if not _cuda_initialized():
|
|
140
|
+
self._restore_env()
|
|
141
|
+
return False
|
|
142
|
+
|
|
143
|
+
def _restore_env(self) -> None:
|
|
144
|
+
if getattr(self, "_previous_server", None) is None:
|
|
145
|
+
os.environ.pop("LUPINE_SERVER", None)
|
|
146
|
+
else:
|
|
147
|
+
os.environ["LUPINE_SERVER"] = self._previous_server
|
|
148
|
+
def devices(self, *, require_available: bool | None = None) -> list[Any]:
|
|
149
|
+
"""Return all declared LUPINE GPUs as ``torch.device("cuda:N")``."""
|
|
150
|
+
|
|
151
|
+
check = self.require_available if require_available is None else require_available
|
|
152
|
+
return [
|
|
153
|
+
_cuda_device(index, require_available=check)
|
|
154
|
+
for index in range(len(self.servers))
|
|
155
|
+
]
|
|
156
|
+
|
|
157
|
+
def device(self, index: int = 0, *, require_available: bool | None = None) -> Any:
|
|
158
|
+
"""Return one declared LUPINE GPU as ``torch.device("cuda:N")``."""
|
|
159
|
+
|
|
160
|
+
index = int(index)
|
|
161
|
+
if index < 0 or index >= len(self.servers):
|
|
162
|
+
raise LupineError(
|
|
163
|
+
f"LUPINE device index {index} is out of range for {len(self.servers)} hosts"
|
|
164
|
+
)
|
|
165
|
+
check = self.require_available if require_available is None else require_available
|
|
166
|
+
return _cuda_device(index, require_available=check)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def connect(
|
|
170
|
+
*,
|
|
171
|
+
host: str | Sequence[str],
|
|
172
|
+
port: int | None = None,
|
|
173
|
+
require_available: bool = False,
|
|
174
|
+
libcuda: str | os.PathLike[str] | None = None,
|
|
175
|
+
) -> Session:
|
|
176
|
+
"""Create a LUPINE session for one or more remote GPU hosts.
|
|
177
|
+
|
|
178
|
+
Use the session before any PyTorch CUDA operation:
|
|
179
|
+
|
|
180
|
+
``with lupine.connect(host=["a:14833", "b:14833"]) as s:``
|
|
181
|
+
|
|
182
|
+
``s.devices()`` then returns ``[torch.device("cuda:0"), torch.device("cuda:1")]``.
|
|
183
|
+
"""
|
|
184
|
+
|
|
185
|
+
return Session(
|
|
186
|
+
servers=_normalize_hosts(host, port),
|
|
187
|
+
require_available=require_available,
|
|
188
|
+
libcuda=libcuda,
|
|
189
|
+
)
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def devices(*, require_available: bool = True) -> list[Any]:
|
|
193
|
+
"""Return devices for the current ``LUPINE_SERVER`` environment."""
|
|
194
|
+
|
|
195
|
+
servers = _servers_from_env()
|
|
196
|
+
if not servers:
|
|
197
|
+
raise LupineError("LUPINE_SERVER is not configured")
|
|
198
|
+
return [
|
|
199
|
+
_cuda_device(index, require_available=require_available)
|
|
200
|
+
for index in range(len(servers))
|
|
201
|
+
]
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def device(index: int = 0, *, require_available: bool = True) -> Any:
|
|
205
|
+
"""Return one device for the current ``LUPINE_SERVER`` environment."""
|
|
206
|
+
|
|
207
|
+
return devices(require_available=require_available)[int(index)]
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
def servers() -> tuple[str, ...]:
|
|
211
|
+
"""Return configured LUPINE servers from ``LUPINE_SERVER``."""
|
|
212
|
+
|
|
213
|
+
return _servers_from_env()
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
def is_configured() -> bool:
|
|
217
|
+
"""Return true when ``LUPINE_SERVER`` names at least one server."""
|
|
218
|
+
|
|
219
|
+
return bool(servers())
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
def is_available() -> bool:
|
|
223
|
+
"""Return true when PyTorch sees at least one CUDA device."""
|
|
224
|
+
|
|
225
|
+
try:
|
|
226
|
+
torch = _torch()
|
|
227
|
+
except LupineError:
|
|
228
|
+
return False
|
|
229
|
+
return bool(torch.cuda.is_available())
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
def device_count() -> int:
|
|
233
|
+
"""Return PyTorch's CUDA device count."""
|
|
234
|
+
|
|
235
|
+
torch = _torch()
|
|
236
|
+
return int(torch.cuda.device_count())
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def current_device() -> int:
|
|
240
|
+
"""Return PyTorch's current CUDA device index."""
|
|
241
|
+
|
|
242
|
+
torch = _torch()
|
|
243
|
+
return int(torch.cuda.current_device())
|
|
244
|
+
|
|
245
|
+
|
|
246
|
+
def synchronize(index: int = 0) -> None:
|
|
247
|
+
"""Synchronize a LUPINE-backed CUDA device."""
|
|
248
|
+
|
|
249
|
+
torch = _torch()
|
|
250
|
+
torch.cuda.synchronize(_cuda_device(index, require_available=False))
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
__all__ = [
|
|
254
|
+
"DEFAULT_PORT",
|
|
255
|
+
"LupineError",
|
|
256
|
+
"Session",
|
|
257
|
+
"connect",
|
|
258
|
+
"current_device",
|
|
259
|
+
"device",
|
|
260
|
+
"device_count",
|
|
261
|
+
"devices",
|
|
262
|
+
"is_available",
|
|
263
|
+
"is_configured",
|
|
264
|
+
"servers",
|
|
265
|
+
"synchronize",
|
|
266
|
+
]
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "lupine"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Small PyTorch adapter helpers for LUPINE-backed CUDA devices"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
dependencies = []
|
|
13
|
+
|
|
14
|
+
[project.optional-dependencies]
|
|
15
|
+
torch = ["torch"]
|
|
16
|
+
|
|
17
|
+
[dependency-groups]
|
|
18
|
+
dev = ["pytest>=8.0"]
|
|
19
|
+
|
|
20
|
+
[tool.hatch.build.targets.wheel]
|
|
21
|
+
packages = ["lupine"]
|
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
import importlib
|
|
2
|
+
import os
|
|
3
|
+
import sys
|
|
4
|
+
import types
|
|
5
|
+
|
|
6
|
+
import pytest
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class FakeDevice:
|
|
10
|
+
def __init__(self, kind, index=None):
|
|
11
|
+
self.type = kind
|
|
12
|
+
self.index = index
|
|
13
|
+
|
|
14
|
+
def __eq__(self, other):
|
|
15
|
+
return isinstance(other, FakeDevice) and (
|
|
16
|
+
self.type,
|
|
17
|
+
self.index,
|
|
18
|
+
) == (other.type, other.index)
|
|
19
|
+
|
|
20
|
+
def __repr__(self):
|
|
21
|
+
return f"{self.type}:{self.index}"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class FakeCuda:
|
|
25
|
+
def __init__(self):
|
|
26
|
+
self.initialized = False
|
|
27
|
+
self.count = 2
|
|
28
|
+
self.current = 1
|
|
29
|
+
self.synchronized = None
|
|
30
|
+
|
|
31
|
+
def is_initialized(self):
|
|
32
|
+
return self.initialized
|
|
33
|
+
|
|
34
|
+
def is_available(self):
|
|
35
|
+
return self.count > 0
|
|
36
|
+
|
|
37
|
+
def device_count(self):
|
|
38
|
+
return self.count
|
|
39
|
+
|
|
40
|
+
def current_device(self):
|
|
41
|
+
return self.current
|
|
42
|
+
|
|
43
|
+
def synchronize(self, selected):
|
|
44
|
+
self.synchronized = selected
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class FakeTorch(types.SimpleNamespace):
|
|
48
|
+
def __init__(self):
|
|
49
|
+
super().__init__()
|
|
50
|
+
self.cuda = FakeCuda()
|
|
51
|
+
|
|
52
|
+
def device(self, kind, index=None):
|
|
53
|
+
return FakeDevice(kind, index)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@pytest.fixture
|
|
57
|
+
def lupine_module(monkeypatch):
|
|
58
|
+
fake_torch = FakeTorch()
|
|
59
|
+
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
|
60
|
+
monkeypatch.delenv("LUPINE_SERVER", raising=False)
|
|
61
|
+
monkeypatch.setattr("ctypes.CDLL", lambda *args, **kwargs: None)
|
|
62
|
+
import lupine
|
|
63
|
+
|
|
64
|
+
yield importlib.reload(lupine), fake_torch
|
|
65
|
+
importlib.reload(lupine)
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def test_connect_sets_env_and_returns_devices(lupine_module):
|
|
69
|
+
lupine, _ = lupine_module
|
|
70
|
+
|
|
71
|
+
with lupine.connect(host="host-a") as session:
|
|
72
|
+
assert os.environ["LUPINE_SERVER"] == "host-a:14833"
|
|
73
|
+
assert session.devices() == [FakeDevice("cuda", 0)]
|
|
74
|
+
assert session.device() == FakeDevice("cuda", 0)
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def test_connect_loads_explicit_libcuda(lupine_module, monkeypatch, tmp_path):
|
|
78
|
+
lupine, _ = lupine_module
|
|
79
|
+
loaded = []
|
|
80
|
+
libcuda = tmp_path / "libcuda.so.1"
|
|
81
|
+
libcuda.write_bytes(b"")
|
|
82
|
+
monkeypatch.setattr(lupine.ctypes, "CDLL", lambda *args, **kwargs: loaded.append(args))
|
|
83
|
+
|
|
84
|
+
with lupine.connect(host="host-a", libcuda=libcuda):
|
|
85
|
+
pass
|
|
86
|
+
|
|
87
|
+
assert loaded == [(str(libcuda),)]
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def test_connect_accepts_multiple_hosts_in_order(lupine_module):
|
|
91
|
+
lupine, _ = lupine_module
|
|
92
|
+
|
|
93
|
+
with lupine.connect(host=["host-a:15000", "host-b:16000"]) as session:
|
|
94
|
+
assert session.servers == ("host-a:15000", "host-b:16000")
|
|
95
|
+
assert session.devices() == [FakeDevice("cuda", 0), FakeDevice("cuda", 1)]
|
|
96
|
+
assert session.device(1) == FakeDevice("cuda", 1)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_connect_restores_env_when_cuda_was_not_initialized(lupine_module, monkeypatch):
|
|
100
|
+
lupine, _ = lupine_module
|
|
101
|
+
|
|
102
|
+
with lupine.connect(host="host-a"):
|
|
103
|
+
assert os.environ["LUPINE_SERVER"] == "host-a:14833"
|
|
104
|
+
|
|
105
|
+
assert "LUPINE_SERVER" not in os.environ
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def test_connect_leaves_env_when_cuda_initialized_inside_context(lupine_module):
|
|
109
|
+
lupine, fake_torch = lupine_module
|
|
110
|
+
|
|
111
|
+
with lupine.connect(host="host-a"):
|
|
112
|
+
fake_torch.cuda.initialized = True
|
|
113
|
+
|
|
114
|
+
assert os.environ["LUPINE_SERVER"] == "host-a:14833"
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def test_connect_accepts_matching_preconfigured_env(lupine_module, monkeypatch):
|
|
118
|
+
lupine, _ = lupine_module
|
|
119
|
+
monkeypatch.setenv("LUPINE_SERVER", "host-a:14833")
|
|
120
|
+
|
|
121
|
+
with lupine.connect(host="host-a:14833") as session:
|
|
122
|
+
assert session.devices() == [FakeDevice("cuda", 0)]
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def test_connect_rejects_different_preconfigured_env(lupine_module, monkeypatch):
|
|
126
|
+
lupine, _ = lupine_module
|
|
127
|
+
monkeypatch.setenv("LUPINE_SERVER", "other:14833")
|
|
128
|
+
|
|
129
|
+
with pytest.raises(lupine.LupineError, match="already configured differently"):
|
|
130
|
+
with lupine.connect(host="host-a:14833"):
|
|
131
|
+
pass
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
def test_connect_refuses_after_cuda_init(lupine_module):
|
|
135
|
+
lupine, fake_torch = lupine_module
|
|
136
|
+
fake_torch.cuda.initialized = True
|
|
137
|
+
|
|
138
|
+
with pytest.raises(lupine.LupineError, match="before PyTorch initializes CUDA"):
|
|
139
|
+
with lupine.connect(host="host-a"):
|
|
140
|
+
pass
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def test_devices_use_current_env(lupine_module, monkeypatch):
|
|
144
|
+
lupine, _ = lupine_module
|
|
145
|
+
monkeypatch.setenv("LUPINE_SERVER", "host-a:14833,host-b:14833")
|
|
146
|
+
|
|
147
|
+
assert lupine.devices() == [FakeDevice("cuda", 0), FakeDevice("cuda", 1)]
|
|
148
|
+
assert lupine.device(1) == FakeDevice("cuda", 1)
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def test_device_bounds_check(lupine_module):
|
|
152
|
+
lupine, _ = lupine_module
|
|
153
|
+
|
|
154
|
+
with lupine.connect(host="host-a") as session:
|
|
155
|
+
with pytest.raises(lupine.LupineError, match="out of range"):
|
|
156
|
+
session.device(1)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def test_duplicate_hosts_are_rejected(lupine_module):
|
|
160
|
+
lupine, _ = lupine_module
|
|
161
|
+
|
|
162
|
+
with pytest.raises(lupine.LupineError, match="unique"):
|
|
163
|
+
lupine.connect(host=["host-a:14833", "host-a"])
|