Add cancellable job execution
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
Lihatoo
parent
11dc85d7ad
commit
048d4b0c07
2 changed files with 292 additions and 59 deletions
|
|
@ -1,6 +1,8 @@
|
|||
import json
|
||||
import mimetypes
|
||||
import multiprocessing
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -80,6 +82,9 @@ IUPAC_CODES = "ACGTUWSMKRYBDHVN"
|
|||
IUPAC_CONSTRAINT_TOKEN = re.compile(rf"[{IUPAC_CODES}](?:\d+)?")
|
||||
|
||||
VALID_COMPUTE = {"pfunc", "pairs", "mfe", "sample", "subopt", "ensemble_size"}
|
||||
CANCEL_REQUESTED_STATUS = "cancel_requested"
|
||||
CANCELED_STATUS = "canceled"
|
||||
TERMINAL_JOB_STATUSES = {"success", "error", CANCELED_STATUS}
|
||||
JOB_STORE = {}
|
||||
JOB_LOCK = threading.Lock()
|
||||
JOB_TTL_SECONDS = int(os.environ.get("NP_JOB_TTL_SECONDS", "3600"))
|
||||
|
|
@ -1321,12 +1326,14 @@ def run_job_payload(payload):
|
|||
|
||||
if workflow == "design":
|
||||
design_options = parse_design_options(payload)
|
||||
design_job_options = DesignOptions(
|
||||
f_stop=design_options["stop_condition"],
|
||||
seed=design_options["seed"],
|
||||
wobble_mutations=design_options["wobble_mutations"],
|
||||
max_time=design_options["max_time_seconds"],
|
||||
)
|
||||
design_job_option_kwargs = {
|
||||
"f_stop": design_options["stop_condition"],
|
||||
"seed": design_options["seed"],
|
||||
"wobble_mutations": design_options["wobble_mutations"],
|
||||
}
|
||||
if design_options["max_time_seconds"] > 0:
|
||||
design_job_option_kwargs["max_time"] = design_options["max_time_seconds"]
|
||||
design_job_options = DesignOptions(**design_job_option_kwargs)
|
||||
design_domain_map, design_domains = build_design_domains(payload.get("design_domains") or [])
|
||||
target_strand_map, design_strands = build_design_strands(
|
||||
payload.get("strands") or [],
|
||||
|
|
@ -1562,21 +1569,64 @@ def set_job_data(job):
|
|||
JOB_STORE[job["job_id"]] = dict(job)
|
||||
|
||||
|
||||
def get_job_data(job_id, include_payload=False):
|
||||
if redis_enabled():
|
||||
raw = redis_client().get(job_key(job_id))
|
||||
if raw is None:
|
||||
return None
|
||||
job = json.loads(raw)
|
||||
if not include_payload:
|
||||
job.pop("payload", None)
|
||||
return job
|
||||
with JOB_LOCK:
|
||||
prune_jobs()
|
||||
job = JOB_STORE.get(job_id)
|
||||
if job is None:
|
||||
return None
|
||||
output = dict(job)
|
||||
if not include_payload:
|
||||
output.pop("payload", None)
|
||||
return output
|
||||
|
||||
|
||||
def should_keep_cancel_state(current, updates):
|
||||
current_status = current.get("status") if current else None
|
||||
next_status = updates.get("status")
|
||||
if current_status in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}:
|
||||
return next_status not in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}
|
||||
return False
|
||||
|
||||
|
||||
def update_job_data(job_id, **updates):
|
||||
if redis_enabled():
|
||||
client = redis_client()
|
||||
current = get_job(job_id)
|
||||
if current is None:
|
||||
current = {"job_id": job_id, "created_at": time.time()}
|
||||
current.update(updates)
|
||||
current["updated_at"] = time.time()
|
||||
client.setex(job_key(job_id), JOB_TTL_SECONDS, json.dumps(current, ensure_ascii=False))
|
||||
return current
|
||||
key = job_key(job_id)
|
||||
while True:
|
||||
pipe = client.pipeline()
|
||||
try:
|
||||
pipe.watch(key)
|
||||
raw = pipe.get(key)
|
||||
current = json.loads(raw) if raw is not None else {"job_id": job_id, "created_at": time.time()}
|
||||
if should_keep_cancel_state(current, updates):
|
||||
pipe.unwatch()
|
||||
return current
|
||||
current.update(updates)
|
||||
current["updated_at"] = time.time()
|
||||
pipe.multi()
|
||||
pipe.setex(key, JOB_TTL_SECONDS, json.dumps(current, ensure_ascii=False))
|
||||
pipe.execute()
|
||||
return current
|
||||
except redis.WatchError:
|
||||
continue
|
||||
finally:
|
||||
pipe.reset()
|
||||
|
||||
with JOB_LOCK:
|
||||
current = JOB_STORE.get(job_id)
|
||||
if current is None:
|
||||
current = {"job_id": job_id, "created_at": time.time()}
|
||||
if should_keep_cancel_state(current, updates):
|
||||
return dict(current)
|
||||
current.update(updates)
|
||||
current["updated_at"] = time.time()
|
||||
prune_jobs(current["updated_at"])
|
||||
|
|
@ -1607,52 +1657,161 @@ def create_job(payload):
|
|||
return job_id
|
||||
|
||||
|
||||
def _job_process_entry(payload, result_queue):
|
||||
apply_thread_limits()
|
||||
try:
|
||||
result_queue.put({"status": "success", "result": run_job_payload(payload)})
|
||||
except BaseException as exc:
|
||||
result_queue.put(
|
||||
{
|
||||
"status": "error",
|
||||
"error": {
|
||||
"message": str(exc),
|
||||
"traceback": traceback.format_exc(),
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def terminate_process(process):
|
||||
if not process.is_alive():
|
||||
return
|
||||
if hasattr(process, "kill"):
|
||||
process.kill()
|
||||
else:
|
||||
process.terminate()
|
||||
process.join(timeout=2)
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(timeout=2)
|
||||
|
||||
|
||||
def _run_job(job_id, payload):
|
||||
started_at = time.time()
|
||||
current = get_job_data(job_id)
|
||||
if current and current.get("status") in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}:
|
||||
elapsed = round(time.time() - started_at, 3)
|
||||
update_job_data(job_id, status=CANCELED_STATUS, payload=None, elapsed_seconds=elapsed)
|
||||
return
|
||||
|
||||
log_event(f"running job_id={job_id}")
|
||||
update_job_data(job_id, status="running")
|
||||
if redis_enabled():
|
||||
redis_client().sadd(JOB_RUNNING_KEY, job_id)
|
||||
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
result_queue = ctx.Queue(maxsize=1)
|
||||
process = ctx.Process(target=_job_process_entry, args=(payload, result_queue), daemon=True)
|
||||
message = None
|
||||
|
||||
try:
|
||||
result = run_job_payload(payload)
|
||||
process.start()
|
||||
while process.is_alive():
|
||||
current = get_job_data(job_id)
|
||||
if current and current.get("status") in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}:
|
||||
terminate_process(process)
|
||||
elapsed = round(time.time() - started_at, 3)
|
||||
update_job_data(
|
||||
job_id,
|
||||
status=CANCELED_STATUS,
|
||||
error={"message": "Job canceled by user."},
|
||||
result=None,
|
||||
payload=None,
|
||||
elapsed_seconds=elapsed,
|
||||
)
|
||||
log_event(f"canceled job_id={job_id} elapsed={elapsed}s")
|
||||
return
|
||||
try:
|
||||
message = result_queue.get_nowait()
|
||||
break
|
||||
except queue.Empty:
|
||||
time.sleep(0.25)
|
||||
|
||||
process.join(timeout=2)
|
||||
if message is None:
|
||||
try:
|
||||
message = result_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
message = None
|
||||
|
||||
current = get_job_data(job_id)
|
||||
if current and current.get("status") in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}:
|
||||
elapsed = round(time.time() - started_at, 3)
|
||||
update_job_data(
|
||||
job_id,
|
||||
status=CANCELED_STATUS,
|
||||
error={"message": "Job canceled by user."},
|
||||
result=None,
|
||||
payload=None,
|
||||
elapsed_seconds=elapsed,
|
||||
)
|
||||
log_event(f"canceled job_id={job_id} elapsed={elapsed}s")
|
||||
return
|
||||
|
||||
elapsed = round(time.time() - started_at, 3)
|
||||
update_job_data(job_id, status="success", result=result, payload=None, elapsed_seconds=elapsed)
|
||||
log_event(f"success job_id={job_id} elapsed={elapsed}s")
|
||||
except Exception as exc:
|
||||
elapsed = round(time.time() - started_at, 3)
|
||||
update_job_data(
|
||||
job_id,
|
||||
status="error",
|
||||
error={
|
||||
"message": str(exc),
|
||||
"traceback": traceback.format_exc(),
|
||||
},
|
||||
payload=None,
|
||||
elapsed_seconds=elapsed,
|
||||
)
|
||||
log_event(f"error job_id={job_id} elapsed={elapsed}s message={exc}")
|
||||
if message and message.get("status") == "success":
|
||||
update_job_data(job_id, status="success", result=message["result"], payload=None, elapsed_seconds=elapsed)
|
||||
log_event(f"success job_id={job_id} elapsed={elapsed}s")
|
||||
elif message and message.get("status") == "error":
|
||||
update_job_data(
|
||||
job_id,
|
||||
status="error",
|
||||
error=message.get("error") or {"message": "Job failed."},
|
||||
payload=None,
|
||||
elapsed_seconds=elapsed,
|
||||
)
|
||||
log_event(f"error job_id={job_id} elapsed={elapsed}s message={message.get('error', {}).get('message')}")
|
||||
else:
|
||||
update_job_data(
|
||||
job_id,
|
||||
status="error",
|
||||
error={"message": f"Job process exited with code {process.exitcode}."},
|
||||
payload=None,
|
||||
elapsed_seconds=elapsed,
|
||||
)
|
||||
log_event(f"error job_id={job_id} elapsed={elapsed}s exitcode={process.exitcode}")
|
||||
finally:
|
||||
terminate_process(process)
|
||||
result_queue.close()
|
||||
result_queue.join_thread()
|
||||
if redis_enabled():
|
||||
redis_client().srem(JOB_RUNNING_KEY, job_id)
|
||||
|
||||
|
||||
def get_job(job_id):
|
||||
if redis_enabled():
|
||||
raw = redis_client().get(job_key(job_id))
|
||||
if raw is None:
|
||||
return None
|
||||
job = json.loads(raw)
|
||||
job.pop("payload", None)
|
||||
return job
|
||||
with JOB_LOCK:
|
||||
prune_jobs()
|
||||
job = JOB_STORE.get(job_id)
|
||||
if job is None:
|
||||
return None
|
||||
output = dict(job)
|
||||
def cancel_job(job_id):
|
||||
current = get_job_data(job_id, include_payload=True)
|
||||
if current is None:
|
||||
return None
|
||||
status = current.get("status")
|
||||
if status in TERMINAL_JOB_STATUSES:
|
||||
output = dict(current)
|
||||
output.pop("payload", None)
|
||||
return output
|
||||
|
||||
if redis_enabled() and status == "queued":
|
||||
redis_client().lrem(JOB_QUEUE_KEY, 0, job_id)
|
||||
|
||||
if status == "queued":
|
||||
updated = update_job_data(
|
||||
job_id,
|
||||
status=CANCELED_STATUS,
|
||||
error={"message": "Job canceled by user."},
|
||||
result=None,
|
||||
payload=None,
|
||||
)
|
||||
else:
|
||||
updated = update_job_data(
|
||||
job_id,
|
||||
status=CANCEL_REQUESTED_STATUS,
|
||||
error={"message": "Cancellation requested."},
|
||||
)
|
||||
updated.pop("payload", None)
|
||||
return updated
|
||||
|
||||
|
||||
def get_job(job_id):
|
||||
return get_job_data(job_id, include_payload=False)
|
||||
|
||||
|
||||
def prune_shares(now=None):
|
||||
if SHARE_MAX_COUNT < 1:
|
||||
|
|
@ -1759,6 +1918,8 @@ def run_worker_loop():
|
|||
if raw is None:
|
||||
continue
|
||||
job = json.loads(raw)
|
||||
if job.get("status") in {CANCEL_REQUESTED_STATUS, CANCELED_STATUS}:
|
||||
continue
|
||||
payload = job.get("payload")
|
||||
if payload is None:
|
||||
continue
|
||||
|
|
@ -1975,6 +2136,14 @@ class AppHandler(BaseHTTPRequestHandler):
|
|||
parsed = urlparse(self.path)
|
||||
|
||||
if parsed.path != "/api/analyze":
|
||||
if parsed.path.startswith("/api/jobs/") and parsed.path.endswith("/cancel"):
|
||||
job_id = parsed.path.split("/")[-2]
|
||||
job = cancel_job(job_id)
|
||||
if job is None:
|
||||
self._respond(*json_bytes({"error": "Job not found"}, status=HTTPStatus.NOT_FOUND))
|
||||
return
|
||||
self._respond(*json_bytes({"status": "success", "job": job}))
|
||||
return
|
||||
if parsed.path == "/api/jobs":
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue