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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
| import logging
import paramiko
import subprocess
from typing import Optional, List, Tuple, Dict, Any
from src.logging_config import SHARED_LOGGER_NAME
from src.connect import (
_execute_remote_command,
_connect_to_ssh_server,
)
logger = logging.getLogger(SHARED_LOGGER_NAME)
def _parse_overall_gpu_stats_output(output_str: str) -> List[Tuple[int, int, int]]:
"""
Parses the output string from nvidia-smi command for overall GPU stats
into a list of GPU stats tuples (utilization, used_memory, total_memory).
"""
results = []
if not output_str:
logger.warning(
"No output received from nvidia-smi command to parse for overall stats."
)
return results
for line in output_str.split("\n"):
if not line:
continue
try:
parts = line.split(",")
if len(parts) == 3:
util = int(parts[0].strip())
mem_used = int(parts[1].strip())
mem_total = int(parts[2].strip())
results.append((util, mem_used, mem_total))
else:
logger.warning(
f"Skipping malformed nvidia-smi line: '{line}'. Expected 3 comma-separated values for overall stats."
)
except (ValueError, IndexError) as e:
logger.warning(
f"Could not parse nvidia-smi line '{line}' for overall stats: {e}"
)
return results
def _get_command_output(
command: str, ssh_client: Optional[paramiko.SSHClient] = None
) -> Optional[str]:
"""
Executes a shell command, either locally or remotely via SSH, and returns its stdout.
Returns None if the command fails or no output is received.
"""
if ssh_client:
output_str = _execute_remote_command(ssh_client, command)
if output_str is None:
logger.error(f"Failed to retrieve output for remote command: {command}")
return output_str
else:
try:
result = subprocess.run(
command, shell=True, capture_output=True, text=True, check=True
)
return result.stdout.strip()
except subprocess.CalledProcessError as e:
logger.error(f"Local command failed: {command}\nStderr: {e.stderr}")
return None
except FileNotFoundError:
logger.error(f"Local command not found: {command}")
return None
def _get_gpu_info(ssh_client: Optional[paramiko.SSHClient] = None) -> Dict[str, str]:
"""Get the mapping of GPU UUID to name, either locally or remotely."""
command = "nvidia-smi --query-gpu=uuid,name --format=csv,noheader"
try:
output = _get_command_output(command, ssh_client)
if output is None:
return {}
gpu_map = {}
for line in output.split("\n"):
if line:
uuid, name = line.split(",", 1)
gpu_map[uuid.strip()] = name.strip()
return gpu_map
except Exception as e:
logger.error(f"Failed to get GPU info: {e}")
return {}
def _get_pid_to_username_map(
ssh_client: Optional[paramiko.SSHClient] = None,
) -> Dict[int, str]:
"""Get a mapping from all system process PIDs to usernames, either locally or remotely."""
command = "ps -eo pid,user --no-headers"
try:
output = _get_command_output(command, ssh_client)
if output is None:
return {}
pid_username_map = {}
for line in output.split("\n"):
if line:
parts = line.strip().split()
if len(parts) == 2:
try:
pid = int(parts[0])
username = parts[1]
pid_username_map[pid] = username
except ValueError:
pass
return pid_username_map
except Exception as e:
logger.error(f"Failed to get process usernames: {e}")
return {}
def get_gpu_data(
gpu_server_ip: Optional[str] = None,
gpu_server_port: Optional[int] = None,
gpu_server_user: Optional[str] = None,
gpu_server_password: Optional[str] = None,
target_username: Optional[str] = None,
) -> Any:
"""
Retrieves either overall GPU statistics or specific user's GPU process information.
Can operate locally or on a specified remote server.
Args:
gpu_server_ip (str, optional): The IP address or hostname of the GPU server.
If None, attempts to get local GPU data.
gpu_server_port (int, optional): The SSH port of the GPU server. Required if gpu_server_ip is provided.
gpu_server_user (str, optional): The username for SSH connection. Required if gpu_server_ip is provided.
gpu_server_password (str, optional): The password for SSH connection. Required if gpu_server_ip is provided.
target_username (str, optional): The username whose GPU processes are to be retrieved.
If None, returns overall GPU stats.
If provided, returns process details for that user.
Returns:
If target_username is None:
A list of tuples, where each tuple contains (compute_utilization_percent, used_memory_MiB, total_memory_MiB).
Returns an empty list on failure.
If target_username is provided:
A dictionary where keys are GPU names (e.g., "NVIDIA GeForce RTX 3090") and values
are lists of dictionaries, each representing a process:
`{"pid": int, "process_name": str, "username": str, "used_memory": str}`.
Returns an empty dictionary on failure.
"""
ssh_client = None
try:
if gpu_server_ip:
if not all([gpu_server_port, gpu_server_user, gpu_server_password]):
logger.error(
"SSH connection details (port, user, password) are required for remote access."
)
return [] if target_username is None else {}
ssh_client = _connect_to_ssh_server(
gpu_server_ip, gpu_server_port, gpu_server_user, gpu_server_password
)
if not ssh_client:
logger.error(f"Failed to establish SSH connection to {gpu_server_ip}.")
return [] if target_username is None else {}
logger.info(f"Connected to {gpu_server_ip} for GPU data.")
else:
logger.info("Getting GPU data from local machine.")
if target_username is None:
command = "nvidia-smi --query-gpu=utilization.gpu,memory.used,memory.total --format=csv,noheader,nounits"
output_str = _get_command_output(command, ssh_client)
if output_str is None:
return []
return _parse_overall_gpu_stats_output(output_str)
else:
nvidia_smi_command = "nvidia-smi --query-compute-apps=gpu_uuid,pid,process_name,used_memory --format=csv,noheader"
gpu_process_lines_str = _get_command_output(nvidia_smi_command, ssh_client)
if gpu_process_lines_str is None:
return {}
gpu_process_lines = gpu_process_lines_str.split("\n")
pid_username_map = _get_pid_to_username_map(ssh_client)
gpu_name_map = _get_gpu_info(ssh_client)
user_gpu_data = {}
for line in gpu_process_lines:
if not line:
continue
parts = line.split(",")
if len(parts) == 4:
try:
gpu_uuid = parts[0].strip()
pid = int(parts[1].strip())
process_name = parts[2].strip()
used_memory = parts[3].strip()
username = pid_username_map.get(pid, "unknown_user")
if username == target_username:
if gpu_uuid not in user_gpu_data:
user_gpu_data[gpu_uuid] = []
user_gpu_data[gpu_uuid].append(
{
"id": gpu_server_ip,
"pid": pid,
"process_name": process_name,
"used_memory": used_memory,
}
)
except ValueError as e:
logger.warning(f"Error parsing NVIDIA-SMI line '{line}': {e}")
else:
logger.warning(
f"Skipping malformed nvidia-smi process line: '{line}'. Expected 4 comma-separated values."
)
final_display_data = {}
for uuid, processes in user_gpu_data.items():
display_name = gpu_name_map.get(uuid, f"Unknown GPU ({uuid})")
final_display_data[display_name] = processes
return final_display_data
except Exception as e:
logger.error(f"An error occurred while getting GPU data: {e}", exc_info=True)
return [] if target_username is None else {}
finally:
if ssh_client:
ssh_client.close()
logger.info("SSH connection closed for GPU data retrieval.")
|