Add cancellation support for generation jobs

api.py

@@ -1,6 +1,7 @@
 import os
 import json
 import time
+import threading
 import uuid
 import shutil
 import subprocess
@@ -23,11 +24,29 @@ BUNDLE_DIR = APP_ROOT / "bundle"
 OUTPUTS_DIR = APP_ROOT / "outputs"
 JOBS_DIR = OUTPUTS_DIR / "jobs"
 
-PUBLIC_BASE_URL = os.getenv("PUBLIC_BASE_URL", "http://localhost:8000").rstrip("/")
+PUBLIC_BASE_URL = os.getenv("PUBLIC_BASE_URL", "").strip().rstrip("/")
 DEEPPURPOSE_URL = os.getenv("DEEPPURPOSE_URL", "http://127.0.0.1:7860/reinvent_predict")
 REINVENT_DEVICE = os.getenv("REINVENT_DEVICE", "cpu")
 ADGPU_BIN = os.getenv("ADGPU_BIN", "")
 
+def get_public_base_url(request: Request | None = None) -> str:
+    configured = PUBLIC_BASE_URL.strip().rstrip("/")
+
+    if configured and "localhost" not in configured and "127.0.0.1" not in configured:
+        return configured
+
+    if request is not None:
+        forwarded_proto = request.headers.get("x-forwarded-proto")
+        forwarded_host = request.headers.get("x-forwarded-host") or request.headers.get("host")
+
+        if forwarded_host:
+            proto = forwarded_proto or request.url.scheme or "https"
+            return f"{proto}://{forwarded_host}".rstrip("/")
+
+        return str(request.base_url).rstrip("/")
+
+    return configured or "http://localhost:8000"
+
 JOBS_DIR.mkdir(parents=True, exist_ok=True)
 
 # Attempt to load DeepPurpose affinity service (optional)
@@ -80,16 +99,22 @@ def new_job_id(prefix: str) -> str:
     return f"{prefix}_{stamp}_{short}"
 
 
-def make_file_meta(job_id: str, file_path: Path) -> Dict[str, str]:
+def make_file_meta(
+    job_id: str,
+    file_path: Path,
+    request: Request | None = None,
+    base_url: str | None = None,
+) -> Dict[str, str]:
     rel_path = file_path.relative_to(JOBS_DIR / job_id).as_posix()
     relative_url = f"/files/jobs/{job_id}/{rel_path}"
+    resolved_base_url = (base_url or get_public_base_url(request)).rstrip("/")
+
     return {
         "filename": file_path.name,
         "relative_url": relative_url,
-        "download_url": f"{PUBLIC_BASE_URL}{relative_url}",
+        "download_url": f"{resolved_base_url}{relative_url}",
     }
 
-
 def clean_record(record: Dict[str, Any]) -> Dict[str, Any]:
     """
     Public response cleaner.
@@ -128,19 +153,112 @@ def dataframe_to_public_records(df: pd.DataFrame) -> List[Dict[str, Any]]:
         records.append(item)
     return records
 
+class JobCancelled(Exception):
+    pass
+
+
+ACTIVE_PROCESSES: Dict[str, subprocess.Popen] = {}
+ACTIVE_PROCESSES_LOCK = threading.Lock()
+
+
+def cancel_flag_path(job_id: str) -> Path:
+    return JOBS_DIR / job_id / "cancel.requested"
 
-def run_cmd(cmd: List[str], cwd: Optional[Path] = None, timeout: Optional[int] = None):
-    result = subprocess.run(
+
+def is_cancel_requested(job_id: str) -> bool:
+    return cancel_flag_path(job_id).exists()
+
+
+def request_job_cancel(job_id: str) -> None:
+    job_dir = JOBS_DIR / job_id
+    job_dir.mkdir(parents=True, exist_ok=True)
+    cancel_flag_path(job_id).touch(exist_ok=True)
+
+
+def raise_if_cancelled(job_id: str) -> None:
+    if is_cancel_requested(job_id):
+        raise JobCancelled(f"Job {job_id} was cancelled.")
+
+
+def terminate_active_process(job_id: str) -> bool:
+    with ACTIVE_PROCESSES_LOCK:
+        proc = ACTIVE_PROCESSES.get(job_id)
+
+    if proc is None or proc.poll() is not None:
+        return False
+
+    proc.terminate()
+    return True
+
+def run_cmd(
+    cmd: List[str],
+    cwd: Optional[Path] = None,
+    timeout: Optional[int] = None,
+    job_id: Optional[str] = None,
+):
+    if job_id:
+        raise_if_cancelled(job_id)
+
+    if not job_id:
+        return subprocess.run(
+            cmd,
+            cwd=str(cwd) if cwd else None,
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            timeout=timeout,
+            check=False,
+        )
+
+    started_at = time.time()
+
+    proc = subprocess.Popen(
         cmd,
         cwd=str(cwd) if cwd else None,
         text=True,
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
-        timeout=timeout,
-        check=False,
     )
-    return result
 
+    with ACTIVE_PROCESSES_LOCK:
+        ACTIVE_PROCESSES[job_id] = proc
+
+    stdout = ""
+
+    try:
+        while True:
+            try:
+                stdout, _ = proc.communicate(timeout=1)
+                break
+            except subprocess.TimeoutExpired:
+                if is_cancel_requested(job_id):
+                    proc.terminate()
+                    try:
+                        stdout, _ = proc.communicate(timeout=10)
+                    except subprocess.TimeoutExpired:
+                        proc.kill()
+                        stdout, _ = proc.communicate()
+
+                    raise JobCancelled(f"Job {job_id} was cancelled while command was running.")
+
+                if timeout is not None and (time.time() - started_at) > timeout:
+                    proc.kill()
+                    stdout, _ = proc.communicate()
+                    raise subprocess.TimeoutExpired(cmd, timeout, output=stdout)
+
+        if is_cancel_requested(job_id):
+            raise JobCancelled(f"Job {job_id} was cancelled.")
+
+        return subprocess.CompletedProcess(
+            args=cmd,
+            returncode=proc.returncode,
+            stdout=stdout,
+            stderr=None,
+        )
+
+    finally:
+        with ACTIVE_PROCESSES_LOCK:
+            ACTIVE_PROCESSES.pop(job_id, None)
 
 def canonicalize_smiles(smiles: str):
     mol = Chem.MolFromSmiles(smiles)
@@ -231,20 +349,26 @@ def job_status_path(job_id: str) -> Path:
     return JOBS_DIR / job_id / "job_status.json"
 
 
-def make_api_url(path: str) -> str:
+def make_api_url(
+    path: str,
+    request: Request | None = None,
+    base_url: str | None = None,
+) -> str:
     if not path.startswith("/"):
         path = "/" + path
-    return f"{PUBLIC_BASE_URL}{path}"
+
+    resolved_base_url = (base_url or get_public_base_url(request)).rstrip("/")
+    return f"{resolved_base_url}{path}"
 
 
-def write_job_status(job_id: str, payload: Dict[str, Any]):
+def write_job_status(job_id: str, payload: Dict[str, Any], base_url: str | None = None):
     job_dir = JOBS_DIR / job_id
     job_dir.mkdir(parents=True, exist_ok=True)
 
     base = {
         "job_id": job_id,
-        "status_url": make_api_url(f"/jobs/{job_id}"),
-        "result_url": make_api_url(f"/jobs/{job_id}/result"),
+        "status_url": make_api_url(f"/jobs/{job_id}", base_url=base_url),
+        "result_url": make_api_url(f"/jobs/{job_id}/result", base_url=base_url),
     }
     base.update(payload)
 
@@ -259,7 +383,7 @@ def write_job_status(job_id: str, payload: Dict[str, Any]):
     tmp_path.replace(path)
 
 
-def read_job_status(job_id: str) -> Dict[str, Any]:
+def read_job_status(job_id: str, base_url: str | None = None) -> Dict[str, Any]:
     path = job_status_path(job_id)
     if not path.exists():
         raise HTTPException(status_code=404, detail="Job not found.")
@@ -279,8 +403,8 @@ def read_job_status(job_id: str) -> Dict[str, Any]:
     if last_error:
         return {
             "job_id": job_id,
-            "status_url": make_api_url(f"/jobs/{job_id}"),
-            "result_url": make_api_url(f"/jobs/{job_id}/result"),
+            "status_url": make_api_url(f"/jobs/{job_id}", base_url=base_url),
+            "result_url": make_api_url(f"/jobs/{job_id}/result", base_url=base_url),
             "status": "running",
             "stage": "status_update",
             "message": "Job status is being updated. Retry shortly."
@@ -288,47 +412,56 @@ def read_job_status(job_id: str) -> Dict[str, Any]:
 
     return {
         "job_id": job_id,
-        "status_url": make_api_url(f"/jobs/{job_id}"),
-        "result_url": make_api_url(f"/jobs/{job_id}/result"),
+        "status_url": make_api_url(f"/jobs/{job_id}", base_url=base_url),
+        "result_url": make_api_url(f"/jobs/{job_id}/result", base_url=base_url),
         "status": "running",
         "stage": "status_update",
         "message": "Job status is not ready yet."
     }
 
-def run_generate_job(job_id: str, req_data: Dict[str, Any]):
+def run_generate_job(job_id: str, req_data: Dict[str, Any], base_url: str):
     job_dir = JOBS_DIR / job_id
 
     try:
         req = GenerateRequest(**req_data)
 
+        raise_if_cancelled(job_id)
+
         write_job_status(job_id, {
             "status": "running",
             "stage": "sampling",
             "message": "Running REINVENT molecule generation",
             "request": req_data,
-        })
+        }, base_url=base_url)
 
         runtime_config = write_runtime_sampling_config(job_dir, req.num_molecules)
 
-        run_reinvent_sampling(runtime_config, job_dir)
+        raise_if_cancelled(job_id)
+        run_reinvent_sampling(runtime_config, job_dir, job_id)
+
+        raise_if_cancelled(job_id)
 
         write_job_status(job_id, {
             "status": "running",
             "stage": "enrichment",
             "message": "Running RDKit descriptors and DeepPurpose predictions",
             "request": req_data,
-        })
+        }, base_url=base_url)
+
+        enriched_csv = run_enrichment(job_dir, req.return_top_k, job_id)
 
-        enriched_csv = run_enrichment(job_dir, req.return_top_k)
+        raise_if_cancelled(job_id)
 
         write_job_status(job_id, {
             "status": "running",
             "stage": "docking",
             "message": f"Running docking mode: {req.docking_mode}",
             "request": req_data,
-        })
+        }, base_url=base_url)
+
+        final_csv = run_optional_docking(job_dir, enriched_csv, req, job_id)
 
-        final_csv = run_optional_docking(job_dir, enriched_csv, req)
+        raise_if_cancelled(job_id)
 
         df = pd.read_csv(final_csv)
 
@@ -355,8 +488,8 @@ def run_generate_job(job_id: str, req_data: Dict[str, Any]):
                 "num_docked": count_docked(df),
             },
             "files": {
-                "csv": make_file_meta(job_id, clean_csv),
-                "json": make_file_meta(job_id, clean_json),
+                "csv": make_file_meta(job_id, clean_csv, base_url=base_url),
+                "json": make_file_meta(job_id, clean_json, base_url=base_url),
             },
             "results": results,
             "warnings": [
@@ -376,7 +509,15 @@ def run_generate_job(job_id: str, req_data: Dict[str, Any]):
             "request": req_data,
             "summary": result_payload["summary"],
             "files": result_payload["files"],
-        })
+        }, base_url=base_url)
+
+    except JobCancelled:
+        write_job_status(job_id, {
+            "status": "cancelled",
+            "stage": "cancelled",
+            "message": "Generation job cancelled by user.",
+            "request": req_data,
+        }, base_url=base_url)
 
     except Exception as e:
         if isinstance(e, HTTPException):
@@ -390,21 +531,46 @@ def run_generate_job(job_id: str, req_data: Dict[str, Any]):
             "message": "Generation job failed",
             "request": req_data,
             "error": detail,
-        })
+        }, base_url=base_url)
+
+def rewrite_public_urls(payload: Dict[str, Any], base_url: str) -> Dict[str, Any]:
+    base_url = base_url.rstrip("/")
+
+    for key in ("status_url", "result_url"):
+        value = payload.get(key)
+        if isinstance(value, str) and value.startswith(("http://localhost", "http://127.0.0.1")):
+            path = "/" + value.split("/", 3)[3] if "/" in value[8:] else ""
+            payload[key] = f"{base_url}{path}"
+
+    files = payload.get("files")
+    if isinstance(files, dict):
+        for meta in files.values():
+            if isinstance(meta, dict):
+                relative_url = meta.get("relative_url")
+                if relative_url:
+                    meta["download_url"] = f"{base_url}{relative_url}"
+
+    file_meta = payload.get("file")
+    if isinstance(file_meta, dict):
+        relative_url = file_meta.get("relative_url")
+        if relative_url:
+            file_meta["download_url"] = f"{base_url}{relative_url}"
+
+    return payload
 
 # -----------------------------
 # Endpoints
 # -----------------------------
 
 @app.get("/health")
-def health():
+def health(request: Request):
     return {
         "status": "ok",
         "bundle_exists": BUNDLE_DIR.exists(),
         "models_generator_exists": (BUNDLE_DIR / "models" / "generator" / "egfr_generator.chkpt").exists(),
         "models_affinity_exists": (BUNDLE_DIR / "models" / "affinity" / "model.pt").exists(),
         "docking_grid_exists": (BUNDLE_DIR / "docking" / "maps_current" / "4WKQ_receptor_v5_SBr.maps.fld").exists(),
-        "public_base_url": PUBLIC_BASE_URL,
+        "public_base_url": get_public_base_url(request),
         "reinvent_device": REINVENT_DEVICE,
         "adgpu_bin": ADGPU_BIN or None,
     }
@@ -442,7 +608,7 @@ def get_job_file(job_id: str, path: str):
 
 
 @app.post("/ligands/export")
-def export_ligand(req: LigandExportRequest):
+def export_ligand(req: LigandExportRequest, request: Request):
     job_id = new_job_id("lig")
     job_dir = JOBS_DIR / job_id
     job_dir.mkdir(parents=True, exist_ok=True)
@@ -467,7 +633,7 @@ def export_ligand(req: LigandExportRequest):
     else:
         raise HTTPException(status_code=400, detail="Unsupported format.")
 
-    meta = make_file_meta(job_id, out_path)
+    meta = make_file_meta(job_id, out_path, request)
     meta["format"] = req.format
 
     return {
@@ -526,10 +692,10 @@ def find_reinvent_command() -> List[str]:
     return [os.getenv("PYTHON", "python"), "-m", "reinvent.Reinvent"]
 
 
-def run_reinvent_sampling(runtime_config: Path, job_dir: Path):
+def run_reinvent_sampling(runtime_config: Path, job_dir: Path, job_id: str):
     cmd = find_reinvent_command() + ["-l", str(job_dir / "reinvent_sampling.log"), str(runtime_config)]
 
-    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=1800)
+    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=1800, job_id=job_id)
 
     log_path = job_dir / "reinvent_command_output.log"
     log_path.write_text(result.stdout or "", encoding="utf-8", errors="replace")
@@ -545,7 +711,7 @@ def run_reinvent_sampling(runtime_config: Path, job_dir: Path):
         )
 
 
-def run_enrichment(job_dir: Path, return_top_k: int) -> Path:
+def run_enrichment(job_dir: Path, return_top_k: int, job_id: str) -> Path:
     input_csv = job_dir / "generated_smiles.csv"
     enriched_csv = job_dir / "generated_smiles_enriched.csv"
 
@@ -565,7 +731,7 @@ def run_enrichment(job_dir: Path, return_top_k: int) -> Path:
         "--top-k", str(return_top_k),
     ]
 
-    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=1800)
+    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=1800, job_id=job_id)
 
     (job_dir / "enrichment.log").write_text(result.stdout or "", encoding="utf-8", errors="replace")
 
@@ -582,7 +748,7 @@ def run_enrichment(job_dir: Path, return_top_k: int) -> Path:
     return enriched_csv
 
 
-def run_optional_docking(job_dir: Path, enriched_csv: Path, req: GenerateRequest) -> Path:
+def run_optional_docking(job_dir: Path, enriched_csv: Path, req: GenerateRequest, job_id: str) -> Path:
     final_csv = job_dir / "generated_results.csv"
 
     script = BUNDLE_DIR / "tools" / "dock_enriched.py"
@@ -610,7 +776,7 @@ def run_optional_docking(job_dir: Path, enriched_csv: Path, req: GenerateRequest
             )
         cmd += ["--adgpu-bin", ADGPU_BIN]
 
-    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=7200)
+    result = run_cmd(cmd, cwd=BUNDLE_DIR, timeout=7200, job_id=job_id)
 
     (job_dir / "docking.log").write_text(result.stdout or "", encoding="utf-8", errors="replace")
 
@@ -628,48 +794,79 @@ def run_optional_docking(job_dir: Path, enriched_csv: Path, req: GenerateRequest
 
 
 @app.post("/generate", status_code=202)
-def submit_generate(req: GenerateRequest, background_tasks: BackgroundTasks):
+def submit_generate(req: GenerateRequest, background_tasks: BackgroundTasks, request: Request):
     job_id = new_job_id("gen")
     job_dir = JOBS_DIR / job_id
     job_dir.mkdir(parents=True, exist_ok=True)
 
     req_data = req.model_dump()
-
+    base_url = get_public_base_url(request)
     write_job_status(job_id, {
         "status": "queued",
         "stage": "queued",
         "message": "Generation job accepted and queued",
         "request": req_data,
-    })
+    }, base_url=base_url)
 
-    background_tasks.add_task(run_generate_job, job_id, req_data)
+    background_tasks.add_task(run_generate_job, job_id, req_data, base_url)
 
     return {
         "job_id": job_id,
         "status": "queued",
         "message": "Generation job accepted. Poll status_url until completed.",
-        "status_url": make_api_url(f"/jobs/{job_id}"),
-        "result_url": make_api_url(f"/jobs/{job_id}/result"),
+        "status_url": make_api_url(f"/jobs/{job_id}", base_url=base_url),
+        "result_url": make_api_url(f"/jobs/{job_id}/result", base_url=base_url),
     }
 
 
 @app.get("/jobs/{job_id}")
-def get_job_status(job_id: str):
-    return read_job_status(job_id)
+def get_job_status(job_id: str, request: Request):
+    base_url = get_public_base_url(request)
+    status = read_job_status(job_id, base_url=base_url)
+    return rewrite_public_urls(status, base_url)
+
+@app.post("/jobs/{job_id}/cancel", status_code=202)
+def cancel_job(job_id: str, request: Request):
+    base_url = get_public_base_url(request)
+    job_dir = JOBS_DIR / job_id
+
+    if not job_dir.exists():
+        raise HTTPException(status_code=404, detail="Job not found.")
+
+    status = read_job_status(job_id, base_url=base_url)
+    current_status = status.get("status")
+
+    if current_status in {"completed", "failed", "cancelled"}:
+        status["message"] = f"Job already {current_status}."
+        return rewrite_public_urls(status, base_url)
+
+    request_job_cancel(job_id)
+    process_terminated = terminate_active_process(job_id)
+
+    write_job_status(job_id, {
+        "status": "cancel_requested",
+        "stage": status.get("stage", "cancel_requested"),
+        "message": "Cancellation requested. The job will stop at the next safe checkpoint.",
+        "process_terminated": process_terminated,
+        "request": status.get("request"),
+    }, base_url=base_url)
 
+    return rewrite_public_urls(read_job_status(job_id, base_url=base_url), base_url)
 
 @app.get("/jobs/{job_id}/result")
-def get_job_result(job_id: str):
-    status = read_job_status(job_id)
+def get_job_result(job_id: str, request: Request):
+    base_url = get_public_base_url(request)
+    status = read_job_status(job_id, base_url=base_url)
 
     if status.get("status") == "failed":
-        return JSONResponse(status_code=500, content=status)
+        return JSONResponse(status_code=500, content=rewrite_public_urls(status, base_url))
 
     if status.get("status") != "completed":
-        return JSONResponse(status_code=202, content=status)
+        return JSONResponse(status_code=202, content=rewrite_public_urls(status, base_url))
 
     result_path = JOBS_DIR / job_id / "generated_results.json"
     if not result_path.exists():
         raise HTTPException(status_code=404, detail="Result file not found.")
 
-    return json.loads(result_path.read_text(encoding="utf-8"))
+    payload = json.loads(result_path.read_text(encoding="utf-8"))
+    return rewrite_public_urls(payload, base_url)