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 5 commits into
base: main
Choose a base branch
from
Draft

Support run tokens #2105

wants to merge 5 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")

andreasjansson and others added 2 commits January 8, 2025 18:29
"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:

```python
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")
```
This makes run_token available via current_scope. Example usage:

```python
from cog import current_scope

class MyPredictor:
    def predict(self):
        token = current_scope().run_token()
        if token:
            # use token
```
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.

Eventually this should just be a value.

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?

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.

2 participants