Reach past Modal's 150s ceiling with background jobs

connect/index.html

@@ -347,6 +347,15 @@ <h2>Connect in three steps.</h2>
           <li class="prompt-card"><span class="prompt-tool">frbus_shock</span><code>What happens to US GDP and unemployment if the Fed raises rates by 100 basis points?</code><button data-copy>copy</button></li>
           <li class="prompt-card"><span class="prompt-tool">hank_shock</span><code>How does a 25 basis point US rate cut hit households across the wealth distribution?</code><button data-copy>copy</button></li>
         </ul>
+        <p class="opt-note">
+          <strong>Slow calls run as jobs.</strong> The hosted server sits behind
+          a transport that abandons any request after 150 seconds, and
+          <code>score_reform</code> over its default five-year window takes
+          longer. Ask for it normally &mdash; the model should reach for
+          <code>start_job</code> and then <code>get_job_result</code>, which
+          hands the work to a 30-minute worker and polls it. Running locally
+          there is no such limit and no job step.
+        </p>
       </div>
       <details class="more compact-more">
         <summary>Terminal setup</summary>

integration/README.md

@@ -212,8 +212,12 @@ benefits the MCP server).
 ## MCP server
 
 Runs over stdio via `python -m policyengine_macro.mcp_server`, exposing
-26 tools:
+28 tools:
 `score_reform` (a PolicyEngine reform through a chosen macro model),
+`start_job` and `get_job_result` (run a slow tool in the background and
+collect it — the hosted transport abandons any request after 150s, and
+`score_reform` over its default five-year window needs longer; there is
+no such limit locally, and no job backend, so call tools directly here),
 `frbus_shock_incidence`, `hank_shock_incidence` and
 `svar_inflation_incidence` (macro-to-microsimulation incidence overlays),
 `dynamic_reform_impact` (the OG-UK overlay dynamic score; local-only —

integration/modal_app.py

@@ -2,8 +2,8 @@
 
     modal deploy integration/modal_app.py
 
-Serves the FastMCP instance from `policyengine_macro.mcp_server` (26 tools:
-list_model_capabilities, get_model_status, recommend_model,
+Serves the FastMCP instance from `policyengine_macro.mcp_server` (28 tools:
+start_job, get_job_result, list_model_capabilities, get_model_status, recommend_model,
 format_score_report, score_reform, obr_shock, list_reform_variables, frbus_shock,
 frbus_list_variables, frbus_summary, frbus_shock_incidence, hank_shock,
 hank_summary, hank_shock_incidence, forecast_uk,
@@ -298,6 +298,37 @@
 pe_data_volume = modal.Volume.from_name("policyengine-macro-pe-data", create_if_missing=True)
 
 
+@app.function(
+    image=image.env({
+        "HF_HOME": f"{CACHE_DIR}/huggingface",
+        "POLICYENGINE_MACRO_PE_DATA_DIR": f"{CACHE_DIR}/policyengine-data",
+    }),
+    cpu=4,
+    memory=8192,
+    # 30 minutes. This function is NOT behind the HTTP proxy, so the 150s
+    # web-endpoint ceiling does not apply to it -- that is the entire point.
+    # The ceiling is why it exists: score_reform over its default five-year
+    # window is two full 372-equation solves plus one PolicyEngine static
+    # costing per year, which does not fit in 150s on this hardware and
+    # cannot be made to.
+    timeout=1800,
+    min_containers=0,
+    scaledown_window=300,
+    max_containers=3,
+    secrets=[modal.Secret.from_name("macromod-hf")],
+    volumes={CACHE_DIR: pe_data_volume},
+)
+def run_tool_job(tool: str, arguments: dict) -> dict:
+    """Worker for start_job: run one allow-listed adapter call to completion."""
+    import os
+
+    if "HUGGING_FACE_TOKEN" not in os.environ and os.environ.get("HF_TOKEN"):
+        os.environ["HUGGING_FACE_TOKEN"] = os.environ["HF_TOKEN"]
+    from policyengine_macro import jobs
+
+    return jobs.run(tool, arguments)
+
+
 @app.function(
     image=image.env({
         "HF_HOME": f"{CACHE_DIR}/huggingface",
@@ -320,9 +351,24 @@ def serve():
     if "HUGGING_FACE_TOKEN" not in os.environ and os.environ.get("HF_TOKEN"):
         os.environ["HUGGING_FACE_TOKEN"] = os.environ["HF_TOKEN"]
 
-    from policyengine_macro import core
+    from policyengine_macro import core, jobs
     from policyengine_macro.mcp_server import mcp
 
+    # Wire the job tools to Modal. Done here rather than at import time so the
+    # local stdio server and the CLI keep no Modal dependency and report the
+    # absence plainly instead of queueing work nothing will run.
+    def _spawn(tool: str, arguments: dict) -> str:
+        return run_tool_job.spawn(tool, arguments).object_id
+
+    def _poll(job_id: str, wait_seconds: int):
+        call = modal.FunctionCall.from_id(job_id)
+        try:
+            return True, call.get(timeout=wait_seconds)
+        except TimeoutError:
+            return False, None
+
+    jobs.set_backend(_spawn, _poll)
+
     # Warm the cheap in-process cache (parses committed results/*.md only —
     # NOT a model estimation, which would make cold starts take minutes).
     try:

integration/src/policyengine_macro/jobs.py

@@ -0,0 +1,138 @@
+"""Background jobs for adapter calls that outlive an HTTP request.
+
+Modal enforces a hard **150-second** ceiling on any HTTP request to a web
+endpoint. Past that the proxy abandons the request and returns a 303 pointing
+at a polling URL. That escape hatch does not work here: a 303 tells a
+standards-compliant client to re-issue the request as a GET, and Modal's
+polling URL rejects GET with ``400 modal-http: bad redirect method``. So every
+HTTP client that follows redirects correctly -- curl -L, fetch, httpx with
+follow_redirects -- gets a 400, and every client that does not gets a bare 303.
+There is no client-side fix.
+
+Measured against the deployed server:
+
+    obr_shock    ~103s warm  -> succeeds; cold -> exceeds 150s and fails
+    score_reform ~111s local, more on Modal -> exceeds 150s, always fails
+
+score_reform cannot be squeezed under the ceiling: the work is two full
+372-equation solves plus one PolicyEngine static costing per year of the
+window. The ceiling is a property of the transport, so the fix is to stop
+doing the work inside the request.
+
+``start_job`` hands the call to a worker with a 30-minute budget and returns
+immediately; ``get_job_result`` polls it, blocking for a bounded interval that
+is itself comfortably inside the 150s ceiling.
+
+This module is deliberately backend-agnostic. The hosted server installs a
+Modal backend at startup (see ``modal_app.serve``); anywhere else -- the local
+stdio server, the CLI, tests -- no backend is installed and the job tools
+report that plainly instead of pretending to queue work that nothing will run.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from policyengine_macro import core
+
+# Adapter calls allowed to run as jobs. An allow-list, not getattr on a
+# caller-supplied name: `tool` arrives from an MCP client, and resolving it
+# straight onto a module would expose every callable in `core`.
+# Names are the ADAPTER function in `core`, which is not always the MCP tool
+# name (the tool `population_reform_impact` is `core.pe_population_impact`,
+# and `dynamic_reform_impact` is `core.dynamic_population_reform_impact`).
+# test_every_allow_listed_tool_resolves_to_a_real_adapter pins that these
+# resolve -- it caught two names that did not.
+JOB_TOOLS: tuple[str, ...] = (
+    "score_reform",
+    "obr_shock",
+    "dynamic_population_reform_impact",
+    "frbus_shock",
+    "hank_shock",
+    "pe_population_impact",
+)
+
+# The longest a get_job_result call may block. The transport ceiling is 150s;
+# this leaves room for request overhead so the poll itself never becomes the
+# thing that times out.
+MAX_WAIT_SECONDS = 120
+
+_spawn: Callable[[str, dict], str] | None = None
+_poll: Callable[[str, int], Any] | None = None
+
+
+def set_backend(spawn, poll) -> None:
+    """Install the job backend. Called once, by the hosted server at startup."""
+    global _spawn, _poll
+    _spawn, _poll = spawn, poll
+
+
+def backend_available() -> bool:
+    return _spawn is not None and _poll is not None
+
+
+class NoBackend(RuntimeError):
+    """Raised where jobs are not available -- i.e. everywhere but the hosted server."""
+
+
+_NO_BACKEND_MESSAGE = (
+    "Background jobs run only on the hosted PolicyEngine Macro server. This "
+    "process has no job backend, so there is nothing to queue the work onto. "
+    "Call the tool directly instead ({tools}) -- running locally there is no "
+    "150-second HTTP ceiling to work around."
+)
+
+
+def resolve(tool: str) -> Callable[..., Any]:
+    """Look up an allow-listed adapter callable by tool name."""
+    if tool not in JOB_TOOLS:
+        raise ValueError(
+            f"{tool!r} cannot be run as a job. Allowed: {', '.join(JOB_TOOLS)}. "
+            "Fast tools should be called directly -- they return well inside "
+            "the transport's 150-second ceiling."
+        )
+    return getattr(core, tool)
+
+
+def run(tool: str, arguments: dict | None) -> dict:
+    """Execute an allow-listed adapter call. This is what the worker runs."""
+    return resolve(tool)(**(arguments or {}))
+
+
+def start(tool: str, arguments: dict | None = None) -> dict:
+    """Queue an adapter call and return a handle without waiting for it."""
+    resolve(tool)  # validate before spawning, so a typo fails in milliseconds
+    if not backend_available():
+        raise NoBackend(_NO_BACKEND_MESSAGE.format(tools=", ".join(JOB_TOOLS)))
+    job_id = _spawn(tool, arguments or {})
+    return {
+        "job_id": job_id,
+        "tool": tool,
+        "status": "running",
+        "next_step": (
+            f"Call get_job_result(job_id={job_id!r}). It blocks until the job "
+            f"finishes or up to wait_seconds (max {MAX_WAIT_SECONDS}); if it "
+            "returns status 'running', call it again with the same job_id."
+        ),
+    }
+
+
+def result(job_id: str, wait_seconds: int = 60) -> dict:
+    """Poll a job. Blocks up to wait_seconds, then reports back either way."""
+    if not backend_available():
+        raise NoBackend(_NO_BACKEND_MESSAGE.format(tools=", ".join(JOB_TOOLS)))
+    wait = max(0, min(int(wait_seconds), MAX_WAIT_SECONDS))
+    done, payload = _poll(job_id, wait)
+    if done:
+        return {"job_id": job_id, "status": "done", "result": payload}
+    return {
+        "job_id": job_id,
+        "status": "running",
+        "waited_seconds": wait,
+        "next_step": (
+            "Not finished yet. Call get_job_result again with the same "
+            "job_id; a score_reform over the default five-year window "
+            "typically needs two or three polls."
+        ),
+    }

integration/src/policyengine_macro/mcp_server.py

@@ -14,11 +14,53 @@
 
 from policyengine_macro import core
 from policyengine_macro import capabilities
+from policyengine_macro import jobs
 from policyengine_macro import reporting
 
 mcp = FastMCP("policyengine-macro")
 
 
+@mcp.tool()
+def start_job(
+    tool: Annotated[str, Field(description="Adapter tool to run as a job")],
+    arguments: Annotated[
+        dict | None, Field(description="Arguments for that tool")
+    ] = None,
+) -> dict:
+    """Run a slow tool in the background and get a job id back immediately.
+
+    USE THIS FOR score_reform. The hosted server sits behind a transport that
+    abandons any HTTP request after 150 seconds, and score_reform over its
+    default five-year window takes longer than that, so calling it directly
+    on the hosted server fails every time. obr_shock, frbus_shock, hank_shock,
+    dynamic_reform_impact and population_reform_impact are also close enough
+    to the ceiling to fail on a cold server.
+
+    Returns a job id straight away. Poll it with get_job_result. The job
+    itself has a 30-minute budget, so the 150-second ceiling stops applying.
+
+    Running locally (stdio server or the pe-macro CLI) there is no such
+    ceiling and no job backend: call the tool directly instead.
+    """
+    return jobs.start(tool, arguments)
+
+
+@mcp.tool()
+def get_job_result(
+    job_id: Annotated[str, Field(description="Job id from start_job")],
+    wait_seconds: Annotated[
+        int, Field(description="Seconds to block waiting, max 120")
+    ] = 60,
+) -> dict:
+    """Fetch a started job's result, waiting up to wait_seconds for it.
+
+    Returns status 'done' with the result, or status 'running' -- in which
+    case call again with the same job_id. A score_reform over the default
+    five-year window typically needs two or three polls.
+    """
+    return jobs.result(job_id, wait_seconds)
+
+
 @mcp.tool()
 def list_model_capabilities() -> list[dict]:
     """List supported questions, outputs, access, runtime, evidence status,