in ailab/benchmark/db_controller.py [0:0]
def get_payload(req):
action = req.get("action")
identifier = req.get("identifier")
benchmarks = req.get("benchmarks")
ids_str = req.get("ids")
id = req.get("id")
claimer = req.get("claimer")
result = req.get("result")
devices_str = req.get("devices")
user = req.get("user")
job_queue = req.get("job_queue")
assert action is not None, "action not provided"
devices = []
if devices_str is not None:
devices = devices_str.split(",")
if len(devices) == 1:
if devices[0] == "":
devices = []
ids = []
if ids_str is None and (action == "run" or action == "release"):
raise AssertionError("ids must be specified for {}".format(action))
elif ids_str is not None:
ids = ids_str.split(",")
if action == "add":
for device in devices:
r = BenchmarkInfo(
status="QUEUE",
identifier=identifier,
benchmarks=benchmarks,
device=device,
user=user,
job_queue=job_queue,
)
r.save()
return {"status": "success"}
elif action == "claim":
for device in devices:
queryset = BenchmarkInfo.objects.filter(
status="CLAIMED", device=device, claimer=claimer, job_queue=job_queue
)
if not queryset:
queryset = BenchmarkInfo.objects.filter(
status="QUEUE", device=device, job_queue=job_queue
)
if queryset:
r = queryset[0]
r.status = "CLAIMED"
r.claimed_time = timezone.now()
r.claimer = claimer
r.save()
Device.objects.filter(device__in=devices).update(heartbeat_time=timezone.now())
queryset = BenchmarkInfo.objects.filter(status="CLAIMED", claimer=claimer)
return {"status": "success", "values": list(queryset.values())}
elif action == "run":
BenchmarkInfo.objects.filter(
id__in=ids, claimer=claimer, job_queue=job_queue
).update(status="RUNNING", start_time=timezone.now())
return {"status": "success"}
elif action == "release":
BenchmarkInfo.objects.filter(id__in=ids, job_queue=job_queue).update(
status="QUEUE", claimed_time=None, start_time=None, claimer=None
)
return {"status": "success"}
elif action == "done":
assert id is not None, "id must be specified for " + action
log = req.get("log")
assert log is not None, "Log field is empty"
status = req.get("status")
assert status == "DONE" or status == "FAILED" or status == "USER_ERROR", (
"Unknown status " + status
)
BenchmarkInfo.objects.filter(id=id, job_queue=job_queue).update(
status=status, done_time=timezone.now(), result=result, log=log
)
return {"status": "success"}
elif action == "status":
queryset = BenchmarkInfo.objects.filter(
identifier=identifier, job_queue=job_queue
)
return {
"status": "success",
"values": list(queryset.values("status", "id", "device")),
}
elif action == "get":
queryset = BenchmarkInfo.objects.filter(id__in=ids, job_queue=job_queue)
return {"status": "success", "values": list(queryset.values())}
elif action == "list_devices":
# Two options: (1) If job_queue=*, it will query all non DISABLED devices;
# (2) Otherwise, it will query non DISABLED devices from the specified
# job_queue.
queryset = Device.objects.exclude(status="DISABLED")
if job_queue != "*":
queryset = Device.objects.filter(job_queue=job_queue)
return {"status": "success", "values": list(queryset.values())}
elif action == "update_devices":
reset = req.get("reset")
if reset == "true":
Device.objects.filter(claimer=claimer).update(
status="DISABLED", update_time=timezone.now()
)
for device in devices:
d = device.split("|")
assert len(d) == 3, "Must have three elements in the input"
data = {
"device": d[0],
"hash": d[1],
"status": "OCCUPIED"
if d[2] == "0"
else "AVAILABLE"
if d[2] == "1"
else "OFFLINE",
"claimer": claimer,
"job_queue": job_queue,
}
defaults = {
"device": data["device"],
"status": data["status"],
"claimer": data["claimer"],
"job_queue": data["job_queue"],
"update_time": timezone.now(),
"heartbeat_time": timezone.now(),
}
obj, created = Device.objects.update_or_create(
hash=data["hash"],
defaults=defaults,
)
obj.save()
return {"status": "success"}
else:
raise AssertionError("action {} not recognized".format(action))