feat(api):expose TRACT model behind a REST API (tract api)#32
Conversation
rocklambros
left a comment
rocklambros
left a comment
There was a problem hiding this comment.
Hey Parth! 👋 First off — really nice PR. The factory pattern with create_app, the clean separation between routes / schemas / settings / lifecycle, the dependency-injection-based testing with StubPredictor, and the explicit --reload vs --workers mutex are all solid choices. You also wrote a great PR description with the mermaid diagram and the file-by-file breakdown — that made reviewing this much easier. Thank you.
I went deep on this one — multi-perspective review across security, MLOps, data-science, and governance lenses — and I left a bunch of inline comments. Don't be intimidated by the count. Most of them are educational and the actual fixes are small. The goal isn't "pass review", it's that you leave this PR knowing more about FastAPI security, dependency hygiene, and ML-API patterns than you walked in with.
How to read the comments
I tagged each one with a tier:
- 🚫 Tier 0 (blocker, R1–R7) — please address before merge. These are the things that, if shipped as-is, would create real failure modes (silent DoS knobs, CORS data exfil, supply-chain drift, etc.).
- 🟡 Tier 1 (should land close to merge, R8–R16) — not strictly blocking, but high-value and small enough that batching them with this PR is the right call.
Each comment follows the same shape: what I see → what to do → why it matters → how to verify. The "why" is the most important part — feel free to push back on any of it. If my reasoning doesn't hold for your use case, say so and we'll talk it through.
One process note that doesn't have a code anchor (Tier 1, R9)
The repo doesn't have a .github/CODEOWNERS file. For a project that ends up in compliance / audit decision chains, having no enforced-review policy is a SOC 2 CC8.1 change-management gap. That's a maintainer task (not yours to merge yourself), but flagging it here so it lands alongside this PR. Suggested:
/tract/api/ @rocklambros
/tract/inference.py @rocklambros
/tract/calibration/ @rocklambros
/pyproject.toml @rocklambros
*.md @rocklambros
Then enable "Require review from Code Owners" in branch protection. This isn't bureaucracy — it's the documented evidence that a maintainer eyeballed the security-relevant paths. Without it, contributor PRs are merged on trust; with it, they're merged on documented review.
On the bigger picture
The deepest thread across all of these comments is this: the PR sits in a real tension between research dev tool (where the API is a convenience for batch researchers) and production decision-support service (where the API is called by GRC platforms making audit decisions). Several findings dissolve if you commit to one framing. Picking "dev tool" simplifies auth, audit log, and model-card binding. Picking "decision support" requires all of those. The README currently does both, and that's where the regulatory-exposure findings come from. Worth thinking about which one you (and the maintainer) really want this to be.
I'd love to chat through anything that's unclear. Leave a @rocklambros pls clarify on any thread and we'll respond. Don't merge anything you don't understand — asking is way better than shipping.
Excited to see the next iteration!
| cors_origins: list[str] = Field(default_factory=lambda: ["*"]) | ||
| max_batch_size: int = 256 | ||
| max_text_length: int = 8192 | ||
| max_top_k: int = 50 |
There was a problem hiding this comment.
🚫 Tier 0 (R1) — these max_* fields are silently inert. This is the kind of bug that ships looking fine and bites somebody hard six months later.
What I see: these three fields (max_batch_size, max_text_length, max_top_k) look like they let an operator tune the limits via env vars — TRACT_API_MAX_BATCH_SIZE=8 and so on. But check the corresponding constants in schemas.py:7-9 and how they're used at lines 27, 33, 40, 44-45:
# schemas.py
MAX_TEXT_LENGTH = 8192
MAX_BATCH_SIZE = 256
MAX_TOP_K = 50
class AssignRequest(_ApiModel):
text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
top_k: int = Field(default=DEFAULT_TOP_K, ge=1, le=MAX_TOP_K)The MAX_* are module-level constants. Pydantic's Field(..., max_length=MAX_TEXT_LENGTH) resolves the value when the class is defined — at import time, not at request time. The settings fields with the same names are read by nothing. An operator sets TRACT_API_MAX_BATCH_SIZE=8, restarts the API, sends a 9-item batch — and the request goes through. The "knob" did nothing.
What to do (I'd pick option A):
Option A — delete the dead fields. They're misleading. The static module constants in schemas.py are the real caps; document them as static in the README.
# settings.py — delete these three lines
# max_batch_size: int = 256
# max_text_length: int = 8192
# max_top_k: int = 50Option B — make them dynamic. Wire them via a pydantic @model_validator(mode="after") on each request class that reads from get_settings() at validation time. More flexible, more code, more failure modes (e.g., what happens if lru_cache returns stale settings between requests?).
Why this matters — the deep version:
In security work, a knob that pretends to work is worse than no knob at all. The threat model: an operator who tightens max_batch_size to defend against DoS believes they're protected. They don't put a real defense (reverse-proxy rate limit, body-size cap, auth) in front. First abusive traffic shows the limit was never enforced. Now you have an outage AND a violated trust in the API contract.
The general principle: no fake interfaces. If a knob is exposed, it must work. If it can't be made to work cheaply, delete it. This applies to env vars, CLI flags, config fields, optional features — anything an operator can configure.
How to verify: add a test in tests/api/test_routes_assign.py:
def test_max_batch_size_env_override(monkeypatch):
monkeypatch.setenv("TRACT_API_MAX_BATCH_SIZE", "8")
from tract.api.settings import get_settings
get_settings.cache_clear()
# submit a 9-item batch, assert 422If you pick Option A, this test should NOT exist — instead, document the static cap in the README so operators know the env var was removed.
There was a problem hiding this comment.
Yes sir , You are absolutely right ! , I missed it , Let's pick option A :
Changes needed -->
1- in settings.py
# settings.py — delete these three lines
# max_batch_size: int = 256
# max_text_length: int = 8192
# max_top_k: int = 50
2- In README --> Removing the Caps , or make them static
Removing the below 2 lines
| `TRACT_API_MAX_BATCH_SIZE` | `256` | Hard cap on `/v1/assign/batch` payload |
| `TRACT_API_MAX_TEXT_LENGTH` | `8192` | Hard cap per control text |
Adding these -->
# Static API limits. Update the README if these values change.
MAX_BATCH_SIZE = 256
MAX_TEXT_LENGTH = 8192
MAX_TOP_K = 50
| version="1.0", | ||
| lifespan=build_lifespan(settings), | ||
| ) | ||
| app.add_middleware( |
There was a problem hiding this comment.
🚫 Tier 0 (R2) — CORS ["*"] default plus no TrustedHostMiddleware is a real data-exfil vector, even for the "localhost-only" deployment pattern the README endorses.
This one's worth slowing down for — CORS is one of those topics that everyone half-knows, and the wrong defaults are everywhere (including in many FastAPI tutorials).
What I see: in settings.py:24 the default is cors_origins: list[str] = Field(default_factory=lambda: ["*"]), and here in app.py:23-28 you wire CORSMiddleware(allow_origins=settings.cors_origins, allow_methods=["GET","POST"], allow_headers=["*"]). No TrustedHostMiddleware anywhere.
The attack chain you might not have considered — DNS rebinding:
- Victim runs
tract apilocally, bound to127.0.0.1:8000(the safe-looking default). - Victim visits
evil.examplein their browser while running TRACT. - The
evil.exampleserver is configured to respond to DNS lookups with a normal IP for the first lookup, then127.0.0.1for subsequent lookups (TTL=0). - JS on
evil.exampledoesfetch("http://evil.example:8000/v1/assign", {...}). Browser resolves the second time →127.0.0.1. Browser sends the request withHost: evil.example. - Without
TrustedHostMiddleware, FastAPI happily serves the request — it doesn't know the Host header is wrong. - With
cors_origins=["*"], the response includesAccess-Control-Allow-Origin: *, so JS onevil.examplecan read the response. - The victim's internal control text just got exfiltrated to
evil.examplefrom a "localhost-only" deployment.
This isn't theoretical. Real DNS rebinding attacks have shipped against Ethereum wallets, Plex, Slack desktop, Transmission — it's a regular finding against any local HTTP service that isn't Host-header-strict.
What to do:
# settings.py
cors_origins: list[str] = Field(
default_factory=lambda: ["http://localhost", "http://127.0.0.1"],
description="Explicit origins only. Set TRACT_API_CORS_ORIGINS to a JSON list to widen.",
)
trusted_hosts: list[str] = Field(
default_factory=lambda: ["localhost", "127.0.0.1"],
description="HTTP Host header allowlist. Add your public hostname when fronted by a proxy.",
)# app.py
from starlette.middleware.trustedhost import TrustedHostMiddleware
app.add_middleware(TrustedHostMiddleware, allowed_hosts=settings.trusted_hosts)
app.add_middleware(
CORSMiddleware,
allow_origins=settings.cors_origins,
allow_credentials=False, # explicit — never echo Authorization cross-origin
allow_methods=["GET", "POST"],
allow_headers=["content-type", "x-tract-auth-token"], # tighten from "*"
)In the README, replace any "CORS to *" guidance with: "For cross-origin clients (e.g. a web tool at https://app.example.com), set TRACT_API_CORS_ORIGINS='[\"https://app.example.com\"]'. Never set to * in production. When fronted by a reverse proxy with a public hostname, also set TRACT_API_TRUSTED_HOSTS='[\"api.example.com\"]'."
Why this matters — the deep version:
Secure-by-default means the dangerous setting requires an explicit operator choice, not the safe one. The right cognitive frame is: "what if someone deploys this without reading the README?" The default has to be safe under that assumption.
CORSMiddleware(allow_origins=["*"]) is a very common FastAPI tutorial default. Tutorials optimize for "make it work in the browser console"; that's not the same as "make it safe to deploy." When you read a tutorial that does this, mentally translate "this is fine for dev" into "this is a footgun in prod" — and never copy the line verbatim.
How to verify:
# DNS rebinding (simulated): Host header that doesn't match the allowlist
curl -H "Host: evil.example" http://127.0.0.1:8000/v1/health
# → 400 Bad Request (from TrustedHostMiddleware)
# CORS: cross-origin preflight from a non-allowed origin
curl -H "Origin: https://evil.example" -H "Access-Control-Request-Method: POST" \
-X OPTIONS http://127.0.0.1:8000/v1/assign -v
# → no Access-Control-Allow-Origin header echoed back| "fastapi>=0.115", | ||
| "uvicorn[standard]>=0.32", | ||
| "pydantic-settings>=2.6", | ||
| "httpx>=0.27", |
There was a problem hiding this comment.
🚫 Tier 0 (R3) — [api] deps are unpinned (violates project QC.1), and httpx is in the wrong group.
Supply-chain hygiene time. This is one of the things that quietly distinguishes "I shipped a feature" from "I shipped a feature that won't compromise my employer next year."
What I see — two problems:
api = [
"fastapi>=0.115",
"uvicorn[standard]>=0.32",
"pydantic-settings>=2.6",
"httpx>=0.27",
]Problem 1 — pure >= with no upper bound. Compare to the phase0 group right above (lines 30-40), which uses torch>=2.2,<3 and similar bounded ranges. The project's CLAUDE.md and the Quality Contract QC.1 (referencing NIST SP 800-218) explicitly mandate pinned dependencies. With >= alone, pip install -e ".[api]" six months from now silently picks up whatever's latest, including breaking changes or compromised releases.
The threat isn't theoretical: Sonatype's State of the Software Supply Chain reports flag 100+ malicious PyPI releases per year since 2022. Both pydantic-settings and uvicorn have transitive trees (httptools, watchfiles, websockets, anyio, sniffio, h11) that compound the attack surface. Just one of those getting compromised → arbitrary code execution on every TRACT API deployment that did a fresh install during the compromise window.
Problem 2 — httpx doesn't belong here. Quick check: grep -rn "import httpx" tract/ returns zero hits. httpx is only used by fastapi.testclient.TestClient in your test suite. Shipping it as a runtime dependency means every production deploy carries httpx + httpcore + h11 + anyio + sniffio + certifi for no runtime benefit — just extra CVE surface in the deployment.
What to do:
api = [
"fastapi>=0.115,<0.120",
"uvicorn[standard]>=0.32,<0.40",
"pydantic-settings>=2.6,<3.0",
]
test-api = [
"httpx>=0.27,<0.30",
]Then update the README install instruction: pip install -e ".[api]" for production, pip install -e ".[api,test-api]" for development (and update CI to use the dev form).
Why this matters — the deep version:
Two ideas worth internalizing:
-
Least privilege applied to dependencies. Production runtime should carry only what production needs. Every dep is a future CVE you're committing to monitoring and patching. Test-only stuff lives in test-only groups.
-
>=is for application code (requirements.txt+ lockfile);>=X,<Yis for libraries (pyproject.toml). The reasoning: an application gets pinned all the way down to a lockfile that pip-installs reproducibly. A library declares compatible-with ranges so downstream apps can resolve their own pins. TRACT publishes to HF (and effectively acts as a library when used aspip install -e .), so it's library-shaped — bounded ranges, not pure floors. The upper bound is your assertion of "I have tested with versions up to here; beyond that, you're on your own."
How to verify:
pip install -e ".[api]" --dry-run --report -
# Look at the resolved versions. fastapi should be < 0.120, uvicorn < 0.40, etc.Bonus — a CI gate so the rule becomes enforcement, not advice:
python -c "
import tomllib
with open('pyproject.toml','rb') as f: data = tomllib.load(f)
api = data['project']['optional-dependencies']['api']
unbounded = [d for d in api if '<' not in d and not d.startswith('tract')]
if unbounded:
raise SystemExit(f'Unpinned [api] deps: {unbounded}')
"There was a problem hiding this comment.
@rocklambros Sir please clarify , in ci.yml file the jobs are using requirements.txt file as a source to download the dependencies , and in the above comment you asked me pip install -e ".[api,test-api]" for development (and update CI to use the dev form). , which means (please confirm) , that I need to remove requirements.txt file from it and make use of pyproject.toml file .
Means , I need to change pip install -r requirements.txt from lint , test and audit job and replace it with ->
pip install -e ".[dev,phase0,llm]" for lint
pip install -e ".[dev,phase0,api,test-api,llm]" for test and audit
Actually while working on ci.yml file , I found real gaps , which are out of scope of this PR but are worth to attend , I have mentioned them in the issue 55 , please have a look.
Another thing in which I want your attention , similar issue of unbounded range of packages in seen in llm group inside pyproject.toml file -
llm = [
"anthropic>=0.20.0",
"pdfplumber>=0.10.0",
]
An unbounded floor means a future malicious/compromised release on either package gets pulled automatically.
Also , anthropic pkg of different versions in being used in 2 groups phase0 and llm -->
llm = ["anthropic>=0.20.0", ...] # no ceiling
phase0 = ["anthropic>=0.25,<1", ...] # bounded
Tell me which to keep and which to remove , so a bare tract[llm] install and a tract[llm,phase0] install can land on materially different anthropic versions for the exact same feature (tract prepare's LLM extraction) -> one constrained, one not. That directly violates CLAUDE.md's own "Deterministic reproducibility" rule: the same code path (llm_extractor.py) behaves differently depending on which unrelated extra happens to also be installed alongside it.
| router = APIRouter(prefix="/v1") | ||
|
|
||
|
|
||
| def get_predictor(request: Request) -> TRACTPredictor: |
There was a problem hiding this comment.
🚫 Tier 0 (R4) — unauthenticated GPU amplifier. The PR is honest that auth is out of scope for v1, but the README's "safe deployment patterns" include HF Spaces, which is a public deployment. Those two postures contradict each other. One of them has to give.
Stay with me — this one's meaty and worth understanding deeply.
What I see: get_predictor here returns request.app.state.predictor with no caller context. ANY HTTP request can drive the model. Combined with /v1/assign/batch accepting 256 controls × 8192 chars each (schemas.py:43-45), this is what security folks call an amplifier:
- Attacker cost per request: ~2 KB of HTTP body (one TCP connection)
- Defender cost per request: 256 BGE-large GPU forward passes + a 256 × 522 matmul + 256 conformal-set computations
- Asymmetry: roughly 1 : 100,000
A single residential-connection laptop can sustain dozens of these per second. Within hours, you burn through cloud GPU quotas. HuggingFace Spaces has shut down Spaces for runaway cost — there's no platform protection against this, the README is mistaken about that.
There's also no request-body size cap, no concurrency limit, no global timeout. Pydantic validates after the body is buffered — so a 2 GB body of invalid JSON costs 2 GB of RAM before pydantic returns 422.
What to do — defense in depth, pick at least two:
Layer 1 (highly recommended): refuse to start in dangerous configurations.
# tract/cli.py near _cmd_api
LOOPBACK_HOSTS = {"127.0.0.1", "localhost", "::1"}
def _cmd_api(args: argparse.Namespace) -> None:
...
settings = get_settings()
if settings.host not in LOOPBACK_HOSTS and not settings.auth_token:
raise SystemExit(
"tract api: refusing to bind to a non-loopback host without auth.\n"
"Pick one:\n"
" - Local-only dev: TRACT_API_HOST=127.0.0.1 (this is the safe default)\n"
" - Shared-secret auth: TRACT_API_AUTH_TOKEN=<random-secret-32-bytes>\n"
" - Behind reverse proxy: bind to 127.0.0.1, let the proxy do auth + TLS termination"
)Layer 2: shared-secret auth dependency.
# routes.py
import hmac
from fastapi import Depends, HTTPException, Request
def require_auth(request: Request, settings: ApiSettings = Depends(get_settings)) -> None:
expected = settings.auth_token
if expected is None:
return # no auth configured — dev mode; safe only on loopback (Layer 1 enforces)
provided = request.headers.get("X-Tract-Auth-Token", "")
if not hmac.compare_digest(expected, provided):
raise HTTPException(status_code=401, detail="missing or invalid auth token")Then Depends(require_auth) on each route.
The hmac.compare_digest is intentional and important. expected == provided is timing-attack vulnerable because Python's string equality short-circuits on the first mismatching character. compare_digest runs in constant time regardless of where the mismatch is. The attack: a remote adversary measures response time to recover the token one byte at a time. Slow but feasible across the open internet.
Layer 3: body-size cap before buffering.
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
class MaxBodySizeMiddleware(BaseHTTPMiddleware):
def __init__(self, app, max_body_bytes: int = 4 * 1024 * 1024): # 4 MiB
super().__init__(app)
self.max_body_bytes = max_body_bytes
async def dispatch(self, request, call_next):
cl = request.headers.get("content-length")
if cl and int(cl) > self.max_body_bytes:
return JSONResponse(
status_code=413,
content={"error": "request body too large", "max_bytes": self.max_body_bytes},
)
return await call_next(request)Pair with uvicorn.run(..., limit_max_requests=10000) (recycles workers periodically — defense against slow leaks).
Why this matters — the deep version:
Two principles:
-
Enforcement should live at the cheapest layer. Body-size check on
Content-Lengthcosts ~10 nanoseconds; pydantic validation after buffering costs N megabytes of RAM. Auth check before route dispatch costs one hash compare; running a full BGE inference and then deciding "should I have run that?" costs hundreds of ms of GPU time. Push the gates outward. -
Reframe "out of scope" as "deferred to operator." Saying "no auth in v1" is fine if the API hard-fails when an operator deploys without auth. Saying "no auth in v1 but here's how to deploy publicly" without that hard-fail is the documented failure mode. The README's "safe deployment patterns" includes HF Spaces — HF Spaces does not provide per-IP rate limiting. So at least one of the three documented patterns is unsafe, and the README is the only thing telling the operator otherwise.
How to verify:
# Layer 1
TRACT_API_HOST=0.0.0.0 tract api
# → SystemExit with clear message
TRACT_API_HOST=0.0.0.0 TRACT_API_AUTH_TOKEN=$(openssl rand -hex 32) tract api &
curl -X POST http://localhost:8000/v1/assign -H 'content-type: application/json' -d '{"text":"hi"}'
# → 401
curl -X POST http://localhost:8000/v1/assign \
-H 'content-type: application/json' \
-H "X-Tract-Auth-Token: $TRACT_API_AUTH_TOKEN" \
-d '{"text":"hi"}'
# → 200
# Layer 3
curl -X POST http://localhost:8000/v1/assign -H "content-length: 99999999" -d 'x' --max-time 5
# → 413 (without ever buffering the body)| rank: int = Field(..., description="1-indexed rank within the top-k results.") | ||
|
|
||
|
|
||
| class AssignResponse(_ApiModel): |
There was a problem hiding this comment.
🚫 Tier 0 (R5) — the deepest finding in the review, and it's not technical: the API doesn't disclose what its predictions are for.
This one's about contract design, not code. Worth thinking carefully about.
What I see: AssignResponse returns assignments with calibrated_confidence in [0, 1]. The PRD reports the model at hit@1 = 0.531 — roughly 47% of top-1 predictions are wrong on the held-out LOFO eval. The schema doesn't surface this anywhere. The README's REST API section doesn't either. A GRC tool consuming /v1/assign has no schema-level signal that the top result is a research-grade suggestion, not a compliance decision of record.
Why this matters more than it might seem:
Six months from now, a customer wires /v1/assign into their SOC 2 / ISO 27001 audit-package generation. Their auditor traces a control mapping: "TRACT API said control_X maps to hub_Y." No record of human review. The auditor's finding lands on the customer. The customer's first call is to the project. The project's defense is "we never said it was authoritative." But the API response says calibrated_confidence: 0.83, and the README markets the API for "GRC platforms" — the absence of a "this is advisory" stamp is the gap the auditor will point at.
This is contract law for ML systems. If the schema says "answer", consumers treat it as an answer. If the schema says "advisory suggestion", consumers know they need to layer HITL on top.
What to do — two small changes:
1. Add a decision_class field to every prediction envelope:
from typing import Literal
class AssignResponse(_ApiModel):
decision_class: Literal["advisory_non_authoritative"] = "advisory_non_authoritative"
model_card_url: str = Field(
default="https://huggingface.co/rockCO78/tract-cre-assignment",
description="Link to the documented intended use, accuracy bounds, and known limitations.",
)
assignments: list[HubAssignment]
ood_flag: bool
ood_score: float
model_version: str
# …existing fields…It's a single literal field. Costs nothing at runtime. Changes the semantic posture of the contract from "here is THE answer" to "here is a suggestion, please apply human judgment." Apply the same to AssignBatchResponse and DuplicateResponse.
2. Add a "Suitable use" section to the README's REST API section:
### Suitable use
The TRACT API returns hub assignments that are **advisory, not authoritative**.
The current model achieves hit@1 = 0.531 on the LOFO evaluation
(see the [model card](https://huggingface.co/rockCO78/tract-cre-assignment)).
- ✅ Appropriate: surfacing candidate hubs to a human reviewer,
accelerating manual control mapping, populating a draft crosswalk
for expert review, exploratory analysis of control families.
- ❌ Inappropriate: directly populating audit packages, compliance
attestations, regulatory disclosures, or any decision that is not
subsequently reviewed by a qualified person.
For any use in a regulated compliance workflow (SOC 2, ISO 27001,
FedRAMP, EU AI Act), the human reviewer remains the decision-maker
of record. The TRACT prediction is an input to that decision, not
the decision itself.Why this matters — the deep version:
Three things to internalize:
-
The schema is the legal posture. A schema field
decision_class: "advisory_non_authoritative"is the project's documented evidence that this was always positioned as a suggestion. Without it, "we said it in a README" is a weaker defense than "the API contract itself said so." -
Calibrated confidence is not the same as accuracy. A 0.83
calibrated_confidencedoes NOT mean 83% chance of being correct. It means "after temperature scaling, the model assigned 83% of probability mass to this hub given the candidate set." Empirical accuracy is what the LOFO eval measures, and it depends on the framework, the control's vocabulary, OOD status, and many other things. Surface BOTH — let consumers pick the right threshold for their use case. -
Don't let your API be more confident than your model is. The PRD says hit@1 = 0.531 honestly. The model card says so honestly. The API doesn't — it just returns a probability-shaped number with no "by the way, ~half of top-1 predictions are wrong" disclaimer. Bring the API's posture in line with the model's actual capability.
| """Factory: returns a lifespan async-cm bound to these settings.""" | ||
|
|
||
| @asynccontextmanager | ||
| async def lifespan(app: FastAPI) -> AsyncIterator[None]: |
There was a problem hiding this comment.
🟡 Tier 1 (R12) — lifespan has no try/except. On predictor load failure, uvicorn enters a silent respawn loop while the CLI parent process happily exits 0.
What I see: here, TRACTPredictor(settings.model_dir) runs with no exception handling. The predictor's constructor raises FileNotFoundError, ValueError (hash mismatch in inference.py:124), or RuntimeError (health check fail in inference.py:139) for various failure modes. If any fire, the lifespan throws, uvicorn kills the worker, then respawns it immediately. New worker hits the same exception. Respawn. Loop.
From the outside, tract api looks running. From the inside, every worker is dying. The CLI parent returns 0 because uvicorn returned cleanly to it.
Failure mode: operator points TRACT_API_MODEL_DIR at the wrong path during a deploy. CI runs tract api --check and gets exit code 0. The "smoke test" passes. The k8s pod healthcheck eventually catches up but on-call has lost 30 minutes trying to figure out why the process is up but not serving.
What to do:
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
settings = get_settings()
try:
logger.info("Loading TRACTPredictor") # don't log the full model_dir path — see note below
app.state.predictor = TRACTPredictor(settings.model_dir)
logger.info(
"Predictor loaded — adapter_hash=%s, t_deploy=%.4f, ood_threshold=%.4f",
app.state.predictor.model_adapter_hash[:12],
app.state.predictor.t_deploy,
app.state.predictor.ood_threshold,
)
except Exception as e:
logger.error("Predictor load failed: %s", e, exc_info=True)
# Fast-fail the worker. Under uvicorn, SystemExit propagates to the master,
# which will exit non-zero rather than tight-looping (uvicorn >= 0.23).
raise SystemExit(1) from e
yield
# cleanup on shutdown — important if you ever add GPU resources
logger.info("Shutting down predictor")Why this matters — the deep version:
Two principles:
-
Fail loud and fast. The opposite of what's happening now (silent respawn loop) is what you want. Loud means: stderr message, structured log line with exception details. Fast means: master process exits non-zero quickly so CI / k8s / systemd can react.
-
Exit codes are an API. A CLI's exit code is part of its contract with whoever runs it (CI, systemd, k8s, supervisord).
tract apiexiting 0 on a config error is a lie. Treat exit codes as carefully as you'd treat HTTP status codes.
Tangent (one-line freebie while you're in this file): the logger.info("Loading TRACTPredictor from %s", settings.model_dir) line at lifecycle.py:24 also leaks the absolute filesystem path to stderr, which then ends up in log aggregators. Strip to basename or just don't log it. Low-impact, but free to fix here.
How to verify:
TRACT_API_MODEL_DIR=/nonexistent tract api
echo $?
# → 1, with a clear error message above|
|
||
|
|
||
| class ControlInput(_ApiModel): | ||
| id: str = Field(..., description="Caller-supplied identifier echoed back in the response.") |
There was a problem hiding this comment.
🟡 Tier 1 (R13) — ControlInput.id has no length cap and no character constraint. The 500 detail string echoes it verbatim.
What I see:
# schemas.py:39
class ControlInput(_ApiModel):
id: str = Field(..., description="Caller-supplied identifier echoed back in the response.")
text: str = ...No max_length, no pattern. Now check routes.py:98-101:
if not preds:
raise HTTPException(status_code=500, detail=f"no predictions for control {control_id}")The attacker-supplied control_id gets reflected into the error message. Same pattern at routes.py:164-167 for the hub-not-found path.
Attack scenarios:
-
Log injection:
control_idcontains\n2026-06-09 23:59:59 [CRITICAL] DATABASE BREACHED —. The next line in the log file looks like a real alert. Mucks up incident-response timelines. -
ANSI escape mischief:
control_idcontains\x1b[31m...\x1b[0mor terminal cursor manipulation. Anyone tailing the log in a real terminal sees corrupted output. -
Response inflation:
control_idis 10 MB. Pydantic accepts it (no length cap). The 500 response now echoes 10 MB. Cheap amplification.
What to do:
class ControlInput(_ApiModel):
id: str = Field(
...,
min_length=1,
max_length=256,
pattern=r"^[A-Za-z0-9._:\-/]+$", # whitelist — adjust to match your real ID formats
description="Caller-supplied identifier. Alphanumeric + ._:-/ only, max 256 chars.",
)
text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)Same treatment for the hub_id path parameter in routes.py:
@router.get("/hubs/{hub_id}", response_model=HubDetail)
def get_hub(
hub_id: str = Path(..., max_length=128, pattern=r"^[A-Za-z0-9._\-/]+$"),
predictor: TRACTPredictor = Depends(get_predictor),
) -> HubDetail:
...Why this matters — the deep version:
The general rule: never trust input that ends up in any output channel — logs, error messages, response bodies, response headers, anything. This is "context-aware output encoding." For each output channel, you either escape the input for that channel (HTML escaping for HTML output, JSON escaping for JSON, etc.) OR restrict the input upfront so escaping isn't needed.
For identifiers, restriction is almost always better than escaping. An ID that's [A-Za-z0-9._:\-/]+ survives any logging context — no characters need escaping, no maximum length surprises.
For free-text fields (text here), restriction isn't possible, so escaping or hashing is the right move. That's why R7 (audit log) logs sha256(text) rather than text itself — same principle.
How to verify:
curl -X POST http://localhost:8000/v1/assign/batch -H 'content-type: application/json' -d '{
"controls": [{"id": "../../etc/passwd\n[FAKE CRITICAL]", "text": "test"}],
"top_k": 5
}'
# → 422 with a field-level validation error, not a reflective 500| settings: Optional API settings. If not provided, default settings will be used. | ||
| """ | ||
| settings = get_settings() if settings is None else settings | ||
| app = FastAPI( |
There was a problem hiding this comment.
🟡 Tier 1 (R14) — no global exception handler. CUDA OOM or any uncaught error from the predictor leaks default FastAPI 500s (potentially with traceback), and CUDA OOM specifically leaves the GPU context fragmented and unrecoverable.
What I see: create_app here adds only CORSMiddleware. No app.add_exception_handler(Exception, ...). If predictor.predict raises torch.cuda.OutOfMemoryError (which happens — adversarial inputs that maximize tokenizer output can OOM a small GPU), FastAPI's default handler returns a generic 500. Under --reload or debug log level, the response can include the traceback.
Worse: CUDA OOM leaves the GPU's memory allocator in a fragmented state. Subsequent requests that try to allocate any tensor will also OOM, until the worker process restarts. With workers=1 (the GPU-recommended setting), one crafted batch → permanent outage until manual restart.
What to do:
import logging
from fastapi import Request
from fastapi.responses import JSONResponse
logger = logging.getLogger(__name__)
def _general_exception_handler(debug: bool):
async def handler(request: Request, exc: Exception):
request_id = getattr(request.state, "request_id", "unknown")
logger.exception("Unhandled exception (request_id=%s)", request_id)
payload = {"error": "internal", "request_id": request_id}
if debug:
payload["exception_type"] = type(exc).__name__
return JSONResponse(status_code=500, content=payload)
return handler
def create_app() -> FastAPI:
settings = get_settings()
app = FastAPI(...)
app.add_exception_handler(Exception, _general_exception_handler(debug=settings.debug))
# CUDA OOM is special — fragmented GPU memory means the worker can't recover
try:
import torch
if torch.cuda.is_available():
import os
async def cuda_oom_handler(request: Request, exc: torch.cuda.OutOfMemoryError):
logger.error("CUDA OOM — exiting worker for clean restart", exc_info=True)
os._exit(1) # bypass cleanup, let uvicorn respawn fresh
app.add_exception_handler(torch.cuda.OutOfMemoryError, cuda_oom_handler)
except ImportError:
pass # CPU-only deploy
# …existing middleware…
return appThe os._exit(1) (not sys.exit) is intentional. sys.exit raises SystemExit which goes through finally blocks, which on a fragmented CUDA context can themselves hang. os._exit is the immediate-terminate option. The worker dies hard, uvicorn respawns it with a fresh CUDA context.
Why this matters — the deep version:
Three principles:
-
Don't leak diagnostic info to attackers. Default exception responses can include stack traces with file paths, package versions, sometimes config values. The handler above scrubs all of that in non-debug mode.
-
Some failures are recoverable, some aren't. General Python exceptions: recoverable, return 500. CUDA OOM: not recoverable in-process, terminate the worker. The right response depends on the failure class. "Catch everything and return 500" is wrong for the second class.
-
A fragmented runtime is worse than a dead process. A dead worker gets respawned cleanly in seconds. A fragmented worker keeps returning 500s for minutes until on-call gets paged. Always prefer fast-fail over zombie-state.
How to verify:
# add a test endpoint temporarily
@router.post("/_debug/boom")
def boom():
raise RuntimeError("test")curl -X POST http://localhost:8000/_debug/boom
# → {"error": "internal", "request_id": "..."} (no traceback in body)| assignments: list[HubAssignment] | ||
| ood_flag: bool = Field( | ||
| ..., | ||
| description="True if the input is out-of-distribution. Assignments below are unreliable when this is True.", |
There was a problem hiding this comment.
🟡 Tier 1 (R15) — bind the calibration triple to every assignment response, so a future investigator can reproduce a decision from one JSON file alone.
What I see: AssignResponse carries ood_flag, ood_score, and model_version. It does NOT carry model_adapter_hash, t_deploy, or ood_threshold. Those live only on /v1/health and /v1/version. To reproduce a six-month-old decision, an investigator needs:
- The original response JSON
- A
/v1/versionsnapshot from the same instant - Trust that the two are consistent (no
model_dirswap between them)
That's three things; we can make it one.
What to do:
class AssignResponse(_ApiModel):
assignments: list[HubAssignment]
ood_flag: bool
ood_score: float
# NEW — the calibration triple, copied from predictor at response time
model_adapter_hash: str # full SHA-256, not short prefix — this is the audit-grade identity
t_deploy: float
ood_threshold: float
request_id: str
# plus R5 additions
decision_class: Literal["advisory_non_authoritative"] = "advisory_non_authoritative"And in the route:
return AssignResponse(
assignments=[_to_hub_assignment(p) for p in preds],
ood_flag=preds[0].is_ood,
ood_score=preds[0].raw_similarity,
model_adapter_hash=predictor.model_adapter_hash,
t_deploy=predictor.t_deploy,
ood_threshold=predictor.ood_threshold,
request_id=request.state.request_id,
)Why this matters — the deep version:
The principle: every audit-grade output should be self-describing. Given the response JSON alone, an investigator should be able to answer:
- What model decided this? (
adapter_hash) - With what calibration parameters? (
t_deploy,ood_threshold) - Was the decision OOD-flagged? (
ood_flag,ood_score) - Could it be reproduced? (
input_shafrom R7's audit log +adapter_hashfrom the response)
Same idea as "include the schema version in the data" or "include the git SHA in every CI artifact." The artifact carries its own provenance — it doesn't depend on a separate side-channel that might be missing or stale.
This pairs with R6 (model_version derived from artifact) and R7 (audit log). Together: the response is the audit record.
| return HealthResponse( | ||
| status="ok", | ||
| model_adapter_hash=predictor.model_adapter_hash, | ||
| t_deploy=predictor.t_deploy, |
There was a problem hiding this comment.
🟡 Tier 1 (R16) — once R15 lands (calibration on the per-response envelope), /v1/health doesn't need to disclose calibration params to every unauthenticated caller. Tighten it.
What I see: /v1/health returns t_deploy, ood_threshold, and model_adapter_hash to anyone. Once R15 is in place, those values are on every assignment response anyway — so there's no operational reason to expose them on /v1/health.
Why disclosing them publicly is mildly bad: an attacker who knows t_deploy=0.074 and ood_threshold=0.568 can locally compute softmax(sims / 0.074) and craft inputs that:
- Maximize confidence in an attacker-chosen wrong hub (poisoning downstream automation)
- Just-clear the OOD threshold (bypassing the API's own "this is OOD, don't trust me" signal)
- Fingerprint the exact adapter version once any CVE-class issue exists for a specific TRACT release
What to do: trim /v1/health to the minimum needed to prove readiness:
@router.get("/health", response_model=HealthResponse)
def health(predictor: TRACTPredictor = Depends(get_predictor)) -> HealthResponse:
"""Readiness probe."""
return HealthResponse(
status="ok",
adapter_hash_short=predictor.model_adapter_hash[:12], # enough to detect a swap, not enough to extract
# REMOVED: t_deploy, ood_threshold
)Move calibration disclosure to /v1/version instead, and gate /v1/version either behind auth (if you implemented R4) or behind an opt-in env var:
# settings.py
disclose_calibration_on_version: bool = Field(
default=False,
description="If True, /v1/version returns t_deploy and ood_threshold. "
"Off by default to reduce adversarial-input construction surface.",
)Why this matters — the deep version:
The general principle is minimum-necessary-disclosure for unauthenticated surfaces. Health endpoints exist to prove "the process is up and serving"; everything beyond that is bonus data that benefits attackers more than legitimate consumers (legitimate consumers can read the per-response envelope instead).
A useful frame: ask "what would change for a legitimate consumer if I removed this field from /v1/health?" If the answer is "nothing, they read it from the per-response envelope," delete it from health. If the answer is "they couldn't function," figure out the per-response path first, then delete it.
How to verify:
curl http://localhost:8000/v1/health | jq
# → {"status":"ok","adapter_hash_short":"abc123def456"} (no calibration)
curl -X POST http://localhost:8000/v1/assign -H 'content-type: application/json' -d '{"text":"test"}' | jq
# → calibration triple is on the response envelope
2 participants
FIXES - issue (#29)
Summary -
tract/api/package andtract apiCLI subcommand that exposes the existing fine-tuned BGE assignment model as a FastAPI REST service./v1/covering single + batch assignment, duplicate detection, hub taxonomy reads, health, and version metadata. The API is read-only and does not write tocrosswalk.db.Motivation
The CLI is sufficient for batch researchers and reproducible pipelines. It does not serve programmatic integration — OpenCRE tooling, GRC platforms, or any third-party application that needs to call TRACT over the wire cannot do so without spawning a Python process per request.
This PR adds the HTTP surface the issue asked for, without changing the assignment logic or any existing CLI command. The API is an optional install (
pip install -e ".[api]") so the base CLI footprint is unchanged.How a single request flows through the stack
flowchart TD Client["HTTP client<br/>(curl, Slack bot, GRC tool)"] Uvicorn["uvicorn worker<br/>(network entry)"] Pydantic["AssignRequest<br/>(tract/api/schemas.py)"] Reject422{"Valid JSON<br/>+ fields?"} Handler["assign() handler<br/>(tract/api/routes.py)"] Dep["get_predictor()<br/>FastAPI dependency"] Predictor["TRACTPredictor.predict()<br/>(tract/inference.py — unchanged)"] Embed["BGE bi-encoder<br/>encodes text → 1024-d vector"] Score["Cosine vs 522 hub<br/>embeddings"] Cal["Temperature scaling<br/>T=0.074"] OOD["OOD check<br/>(max sim < 0.568?)"] Wrap["_to_hub_assignment()<br/>HubPrediction → HubAssignment<br/>(+ rank, drop is_ood)"] Envelope["AssignResponse envelope<br/>(lift is_ood + ood_score<br/>+ model_version)"] Serialize["FastAPI serializes<br/>to JSON"] Client -->|"POST /v1/assign<br/>{text, top_k}"| Uvicorn Uvicorn --> Pydantic Pydantic --> Reject422 Reject422 -->|"No"| Return422["HTTP 422<br/>(handler never runs)"] Reject422 -->|"Yes"| Handler Handler --> Dep Dep -->|"reads app.state.predictor<br/>(loaded once by lifespan)"| Predictor Predictor --> Embed Embed --> Score Score --> Cal Cal --> OOD OOD --> Wrap Wrap --> Envelope Envelope --> Serialize Serialize -->|"200 OK<br/>{assignments, ood_flag, model_version}"| Client Return422 --> Client style Predictor fill:#e1f5e1 style Pydantic fill:#fff4e1 style Envelope fill:#fff4e1 style Return422 fill:#ffe1e1Two things to notice. First, the green box (
TRACTPredictor.predict()) is unchanged existing code — the new layer is everything around it. Second, pydantic validation runs before the handler executes; a malformed payload never touches the model. That's how every "bad input" test in the suite gets a free 422 without us writing input-validation code.Auto-generated OpenAPI documentation
FastAPI serves a Swagger UI at
/docsand a ReDoc rendering at/redocfor free. Every endpoint, every request schema, every response shape, every validation rule is auto-documented from the pydantic models.The raw OpenAPI spec is at
/openapi.jsonfor client-codegen tools.What changed, by directory
New package —
tract/api/__init__.pycreate_appfor uvicorn'sfactory=Truemodeapp.pylifecycle.pyTRACTPredictoronce per worker on startup, lets failures kill the worker (no degraded mode)routes.pyget_predictordependency + two adapter functions (_to_hub_assignment,_to_dup_out)schemas.py_ApiModelbase (disables protected namespaces) and seventeen v1 contract models. Includes a cross-field validator onDuplicateRequestsettings.pyApiSettingspydantic-settings class, env-var-driven withTRACT_API_*prefix;get_settings()factory withlru_cacheModified —
tract/cli.py— addstract apisubparser (line 64),_cmd_apihandler (line 470), and one entry in the dispatch dict (line 1660). CLI flags--host/--port/--workers/--reloadtranslate to env vars so uvicorn workers inherit them. Mutual-exclusion check between--reloadand--workers > 1.inference.py— adds five read-only@propertyaccessors onTRACTPredictor:hierarchy,model_adapter_hash,t_deploy,ood_threshold,deployment_timestamp. Zero behavior change — these expose existing private state so the API layer doesn't reach into underscores. No existing test or CLI command relies on these names being private.Modified — top-level
pyproject.toml— new optional dependency group[api]containingfastapi,uvicorn[standard],pydantic-settings,httpx.README.md— adds a new## REST APIsection after CLI Overview with endpoint table, configuration table, security guidance, and when-to-use comparison. Bumps subcommand count 19 → 20.New —
tests/api/A self-contained subdirectory with its own
conftest.py(so api fixtures don't leak into the global test suite). 33 unit tests use aStubPredictorinjected viaapp.dependency_overrides. 3 integration tests are marked@pytest.mark.integrationand skipped unless real deployment artifacts exist locally.conftest.pyStubPredictor, app builder, TestClienttest_app_factory.pytest_routes_assign.pytest_routes_hubs.pytest_routes_duplicates.pytest_routes_health_version.pytest_ood.pytest_integration.pyEndpoints
All under
/v1/. OpenAPI docs auto-served at/docs.POST/v1/assignTRACTPredictor.predictPOST/v1/assign/batchTRACTPredictor.predict_batchPOST/v1/duplicatesTRACTPredictor.find_duplicatesGET/v1/hubspredictor.hierarchy.hubs.values()(paginated)GET/v1/hubs/{hub_id}predictor.hierarchy.hubs.get(...)GET/v1/healthGET/v1/versionDesign decisions
tract api --host XsetsTRACT_API_HOST=Xbefore invokinguvicorn.run(...)so spawned workers inherit the value. Required becausefactory=Trueworkers can't share Python state with the parent.create_app. The lifespan would otherwise try to load the real model. We overrideget_predictorandget_settingsvia FastAPI's dependency-override mechanism, which doesn't require touching the lifespan.@propertyaccessors onTRACTPredictorare the only changes to existing code. Net diff intract/inference.pyis +20 lines, all read-only properties exposing existing state.Test plan
pytest tests/api/ -q -m "not integration"— 33 unit tests pass in under 1 second.pytest tests/api/ -m integration— 3 integration tests pass against the real deployment model (load + assign + OOD + health).pytest tests/ -q --ignore=tests/api/test_integration.py— confirm no regression in the existing 866-test baseline (only changedtract/inference.pyto add public properties, no behavior change).mypy tract/api/ --strict— all type checks pass.ruff check tract/api/— clean.tract api:Access-Control-Allow-Origin: *and the expected method/header allowlist.DuplicateRequestrejects inverted thresholds with the messageduplicate_threshold (0.5) must be >= similar_threshold (0.9).tract --helplists the newapisubcommand.tract api --reload --workers 4exits withSystemExit: --reload and --workers > 1 are mutually exclusive.tract api --host 127.0.0.1 --port 8765actually binds to the requested port (verified with curl).Files changed
Added
tract/api/__init__.pytract/api/app.pytract/api/lifecycle.pytract/api/routes.pytract/api/schemas.pytract/api/settings.pytests/api/__init__.pytests/api/conftest.pytests/api/test_app_factory.pytests/api/test_routes_assign.pytests/api/test_routes_duplicates.pytests/api/test_routes_hubs.pytests/api/test_routes_health_version.pytests/api/test_ood.pytests/api/test_integration.pyModified
tract/cli.py— new subparser,_cmd_apihandler, dispatch entrytract/inference.py— five read-only@propertyaccessors onTRACTPredictorpyproject.toml— new[api]optional dependency groupREADME.md— REST API section, CLI overview update, navigation table entryOut of scope
@rocklambros Sir , please review it and tell me if any improvements are needed ....
I also added detailed comments in each file to make the code easier to understand.