def f(*args, **kwargs):
task = args[0]
task_id = task.request.id
start = time.time()
traceback_str = None
res = {}
try:
outputs = func(*args, **kwargs)
if task.name.endswith("sim"):
version = outputs.pop("version", OUTPUTS_VERSION)
if version == "v0":
res["model_version"] = "NA"
res.update(dict(outputs, **{"version": version}))
else:
outputs = cs_storage.write(task_id, outputs)
res.update({
"model_version": functions.get_version(),
"outputs": outputs,
"version": version,
})
else:
res.update(outputs)
except Exception:
traceback_str = traceback.format_exc()
finish = time.time()
if "meta" not in res:
res["meta"] = {}
res["meta"]["task_times"] = [finish - start]
if traceback_str is None:
res["status"] = "SUCCESS"
else:
res["status"] = "FAIL"
res["traceback"] = traceback_str
return res
def done_callback(future, job_id, comp_url, comp_api_token, start_time):
"""
This should be called like:
callback = functools.partial(
done_callback,
job_id=job_id,
comp_url=os.environ.get("COMP_URL"),
comp_api_token=os.environ.get("comp_api_token"),
start_time=time.time()
)
This is because this function needs the job id, comp url,
api token, and start time arguments, but dask only passes the
future object.
"""
finish = time.time()
print(f"job_id: {job_id}")
print(f"from dask")
print(f"state: {future.status}")
res = {}
traceback_str = None
try:
outputs = future.result()
outputs = cs_storage.write(job_id, outputs)
res.update({
"model_version": functions.get_version(),
"outputs": outputs,
"version": "v1",
})
except Exception:
traceback_str = traceback.format_exc()
print(f"exception in callback with job_id: {job_id}")
print(traceback_str)
if "meta" not in res:
res["meta"] = {}
res["meta"]["task_times"] = [finish - start_time]
if traceback_str is None:
res["status"] = "SUCCESS"
else:
res["status"] = "FAIL"
res["traceback"] = traceback_str
res["job_id"] = job_id
print("got result", res)
print(f"posting data to {comp_url}/outputs/api/")
resp = requests.put(
f"{comp_url}/outputs/api/",
json=res,
headers={"Authorization": f"Token {comp_api_token}"},
)
print("resp", resp.status_code)
if resp.status_code == 400:
print("errors", resp.json())
async def task_wrapper(callback_url, task_name, func, task_kwargs=None):
print("async task", callback_url, func, task_kwargs)
start = time.time()
traceback_str = None
res = {
"task_name": task_name,
}
try:
if task_kwargs is None:
print("getting task_kwargs")
resp = await get_task_kwargs(callback_url)
task_kwargs = resp.json()["inputs"]
print("got task_kwargs", task_kwargs)
outputs = func(**(task_kwargs or {}))
res.update(
{
"model_version": functions.get_version(),
"outputs": outputs,
"version": "v1",
}
)
except Exception:
traceback_str = traceback.format_exc()
finish = time.time()
if "meta" not in res:
res["meta"] = {}
res["meta"]["task_times"] = [finish - start]
if traceback_str is None:
res["status"] = "SUCCESS"
else:
res["status"] = "FAIL"
res["traceback"] = traceback_str
print("saving results...")
async with httpx.AsyncClient() as client:
resp = await client.post(callback_url, json=res, timeout=120)
print("resp", resp.status_code, resp.url)
assert resp.status_code in (200, 201), f"Got code: {resp.status_code} ({resp.text})"
return res
def inputs_version(self):
return {"version": functions.get_version()}
def version(**task_kwargs):
return {"version": functions.get_version()}