Skip to content

feat(coglet): restore Scope.context for per-prediction context#2853

Merged
markphelps merged 4 commits intomainfrom
mphelps/restore-scope-context
Mar 20, 2026
Merged

feat(coglet): restore Scope.context for per-prediction context#2853
markphelps merged 4 commits intomainfrom
mphelps/restore-scope-context

Conversation

@markphelps
Copy link
Copy Markdown
Contributor

@markphelps markphelps commented Mar 20, 2026

Summary

  • Re-implements current_scope().context on coglet's Rust Scope that was dropped during the coglet rewrite (feature: Coglet (Rust) HTTP Server Only #2714)
  • Threads context: HashMap<String, String> from the HTTP request body through the full stack to the Python-visible Scope.context getter
  • Fixes hello-context in replicate/cog-examples which relies on this API

Background

The original Scope.context was introduced in #2330 (v0.14.11) to allow per-prediction configuration to be passed alongside user input:

POST /predictions
{
  "input": { ... },
  "context": {
    "procedure_source_url": "file:///foo/bar/baz",
    "replicate_api_token": "r8_foop"
  }
}
def predict():
    context = cog.current_scope().context
    token = context["replicate_api_token"]

When the Python server was replaced by coglet (Rust) in #2714, python/cog/server/scope.py was deleted and the new Rust Scope only had metrics and record_metric() — the context field was not carried over.

Changes

File Change
metric_scope.rs Add context: Py<PyDict> to Scope, expose as #[getter], accept in Scope::new() and ScopeGuard::enter()
worker_bridge.rs Pass context from predict() to ScopeGuard::enter()
protocol.rs Add context: HashMap<String, String> to SlotRequest::Predict, update rehydrate_input()
routes.rs Add context to PredictionRequest, thread through create_prediction_with_id()
service.rs Thread context through predict() and build_slot_request()
worker.rs Add context to PredictHandler::predict trait and run_prediction()
__init__.pyi Add context property to Scope type stub

Testing

  • All 199 existing Rust tests pass
  • Clippy clean, rustfmt clean
  • Snapshot tests updated for new context: {} field in serialized SlotRequest

Closes #2852

@markphelps
feat(coglet): restore Scope.context for per-prediction context
a6d8bd4 
Re-implement the context property on coglet's Scope that was lost when
the Python server was replaced by the Rust coglet in #2714. The original
feature (#2330) allowed predictors to access per-prediction context via
current_scope().context.

Thread context (HashMap<String, String>) from the HTTP request body
through the full stack: PredictionRequest -> service.predict ->
SlotRequest::Predict -> worker -> ScopeGuard -> Scope.context getter.

Closes #2852
@markphelps markphelps requested a review from a team as a code owner March 20, 2026 16:43
@markphelps markphelps mentioned this pull request Mar 20, 2026
Merged
@markphelps markphelps requested a review from michaeldwan March 20, 2026 17:42
michaeldwan
michaeldwan approved these changes Mar 20, 2026
View reviewed changes
Copy link
Copy Markdown
Member

@michaeldwan michaeldwan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice find!

@markphelps markphelps merged commit ec70f54 into main Mar 20, 2026
34 checks passed
@markphelps markphelps deleted the mphelps/restore-scope-context branch March 20, 2026 18:09
michaeldwan added a commit that referenced this pull request Apr 7, 2026
@michaeldwan
fix: restore dict[str, str] return type for Scope.context stub
9978a6f 
Regression from #2853: the context property returns Py<PyDict> which
pyo3-stub-gen infers as bare `dict`. Add an override_return_type
annotation so the generated stub preserves the intended dict[str, str]
contract.
@michaeldwan michaeldwan mentioned this pull request Apr 7, 2026
Merged
github-merge-queue bot pushed a commit that referenced this pull request Apr 8, 2026
@michaeldwan
ci: enforce stub freshness in CI, fix existing stub drift (#2912)
93c5569 
* ci: add stub freshness check to CI

The generate:stubs mise task and stub:check already exist but weren't
enforced in CI, so .pyi stubs could drift when Rust source changed.

Adds a check-stubs job that runs `mise run --force stub:check` on Rust
changes, using --force to bypass sources/outputs mtime freshness (which
would skip regeneration in CI's fresh checkout).

* chore: regenerate coglet Python stubs

Stubs drifted after the BuildInfo.dirty field was added in #2829 and
docstring updates in recent PRs. Regenerated with `mise run generate:stubs`.

* fix(ci): set LD_LIBRARY_PATH for stub_gen in check-stubs

stub_gen is a PyO3 binary that dynamically links libpython. setup-uv
installs Python to a non-standard path, so the shared library isn't on
LD_LIBRARY_PATH by default. Query sysconfig for the correct LIBDIR.

* fix: restore dict[str, str] return type for Scope.context stub

Regression from #2853: the context property returns Py<PyDict> which
pyo3-stub-gen infers as bare `dict`. Add an override_return_type
annotation so the generated stub preserves the intended dict[str, str]
contract.

* fix(ci): use setup-python for shared libpython in check-stubs

stub_gen needs libpython3.x.so at runtime (PyO3 auto-initialize).
setup-uv's python-build-standalone is statically linked and doesn't
ship the .so, causing 'cannot open shared object file' at runtime.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@michaeldwan michaeldwan michaeldwan approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Re-implement Scope.context in coglet (Rust)

2 participants

@markphelps @michaeldwan