def test_publish_device_state_without_channel_name(self):
device = self.device.serialize()
publish_device_state(device)
device.update({"gpus": []})
data = {
'type': 'update_info',
'message': json.dumps(device),
}
send_function_mock.assert_called_with(self.device.name, data)
def publish_device_states(self, gpu_data=None):
if gpu_data is None:
# we do not care about the GPU data of our devices
for device in Device.objects.all():
gpus = []
for _ in range(random.randint(0, 2)):
gpu = get_gpu_template()
gpu['uuid'] = uuid.uuid4().hex
gpus.append(gpu)
device_data = {"name": device.name, "gpus": gpus}
publish_device_state(device_data)
else:
for data in gpu_data:
publish_device_state(data)
def connect(self):
self.user = self.scope['user']
self.device_name = self.scope['url_route']['kwargs']['device_name']
device = Device.objects.get(name=self.device_name)
if device.can_be_used_by(self.user):
async_to_sync(self.channel_layer.group_add)(
self.device_name,
self.channel_name,
)
publish_device_state(device, self.channel_name)
self.accept()
else:
self.close()
def update_gpu_info(request):
data = json.loads(request.read().decode("utf-8"))
device_name = data["device_name"]
device = Device.objects.get(
name=device_name) # Device should exist because it's authorized
gpus = []
for gpu_data in data["gpu_data"]:
uuid = gpu_data['uuid']
gpu, created = GPU.objects.get_or_create(uuid=uuid,
model_name=gpu_data['name'],
device=device)
gpu_in_use = True if gpu_data.get("in_use", "na") == "yes" else False
processes = []
if gpu_in_use:
# push processes if this is supported by the GPU
for process in gpu_data.get('processes', []):
processes.append({
"name":
process.get("name", "Unknown"),
"pid":
int(process.get("pid", "0")),
"memory_usage":
process.get("used_memory", "Unknown"),
"username":
process.get("username", "Unknown"),
})
serialized_gpu = gpu.serialize()
gpu = {
"used_memory": gpu_data["memory"]["used"],
"total_memory": gpu_data["memory"]["total"],
"utilization": gpu_data["gpu_util"],
"in_use": gpu_in_use,
"marked_as_failed": False,
"processes": processes,
}
gpu.update(serialized_gpu)
gpus.append(gpu)
device_data = device.serialize()
device_data["gpus"] = gpus
publish_device_state(device_data)
return HttpResponse()