Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support run tokens #2105

Draft
wants to merge 13 commits into
base: main
Choose a base branch
from
Draft

Support run tokens #2105

wants to merge 13 commits into from

Conversation

andreasjansson
Copy link
Member

"Run tokens" are used to run downstream Replicate models from within a model.

They are not yet exposed in an API, but can be found inside a Cog predictor with some ugly frame inspection:

def _find_api_token() -> str:
    frame = inspect.currentframe()
    while frame:
        if "self" in frame.f_locals:
            self = frame.f_locals["self"]
            if hasattr(self, "_current_run_token"):
                token = self._current_run_token
                return token
        frame = frame.f_back
    raise ValueError("No run token found in call stack")

Copy link
Contributor

@philandstuff philandstuff left a comment

Choose a reason for hiding this comment

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

I've pushed a commit to expose run_token via current_scope, which we added as a way to expose per-prediction functionality like this.


def include(model_path: str) -> Callable[..., Any]:
def run(**inputs: dict[str, Any]) -> Any:
client = replicate.Client(api_token=_find_api_token())
Copy link
Contributor

Choose a reason for hiding this comment

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

This can use current_scope().run_token, which is the affordance that exposes the run token to the model. In future, if the prototype is successful, we should make replicate-python automatically detect this and use it if no token was explicitly provided.

Further, I think the whole include function here probably belongs on replicate-python, not in cog?

andreasjansson and others added 8 commits February 7, 2025 15:00

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning
"Run tokens" are used to run downstream Replicate models from within a model.

They can be retrieved using scope:

```python
from cog import current_scope

token = current_scope()._run_token
```

Co-authored-by: Nick Stenning <[email protected]>

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning
@nickstenning nickstenning force-pushed the run-tokens branch 2 times, most recently from d8717f1 to a36c27d Compare February 7, 2025 14:35

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning
Use `Optional[T]` rather than `T | None` and use types from `typing`
rather than subscripting native types.
nickstenning and others added 3 commits February 7, 2025 16:29

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning
On replicate.com, a model with return type `Path` can return a URL, and it is handled as though the model returned a local file path.

Locally, `cog predict` will fail with

```
ⅹ Failed to write output: Failed to decode dataurl: missing data prefix
```

This change makes `cog predict` download and save the output URL.

The `useExtensionIfUnknownContentType` is mostly needed because .webp is incorrectly returned by replicate.delivery as application/octet-stream.

Verified

This commit was signed with the committer’s verified signature.
nickstenning Nick Stenning
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants