24 lines
736 B
Plaintext
24 lines
736 B
Plaintext
# -*- coding:utf-8 -*-
|
|
import json, subprocess
|
|
|
|
result = {
|
|
'service': 'clip-embedding',
|
|
'model': 'laion/CLIP-ViT-H-14-laion2B-s32B-b79K',
|
|
'projection_dim': 1024,
|
|
'gpu_id': int(__import__('os').environ.get('CLIP_GPU_ID', '2')),
|
|
'gpus': []
|
|
}
|
|
try:
|
|
out = subprocess.check_output(
|
|
['nvidia-smi', '--query-gpu=index,utilization.gpu,memory.used,memory.total',
|
|
'--format=csv,noheader,nounits'], timeout=5
|
|
).decode().strip()
|
|
for line in out.split(chr(10)):
|
|
p = [x.strip() for x in line.split(',')]
|
|
result['gpus'].append({'id': int(p[0]), 'util': int(p[1]), 'mem_used': int(p[2]), 'mem_total': int(p[3])})
|
|
except Exception:
|
|
pass
|
|
|
|
return json.dumps(result, ensure_ascii=False)
|
|
|