Add async support

[mattt]

This PR adds support for async operations to the Replicate client. Namespace operations like predictions.list and models.create now have async variants with the async_ prefix (predictions.async_list and models.async_create).

Here's an example of what that looks like in practice:

import replicate

model = await replicate.models.async_get("stability-ai/sdxl")

input = {
  "prompt": "A chariot pulled by a team of rainbow unicorns, driven by an astronaut, dramatic lighting",
}

output = await replicate.async_run(f"stability-ai/sdxl:{model.latest_version.id}", input)

Output

One of the most common questions I hear is how to run a bunch of predictions in parallel. The async functionality provided by this PR makes it really straightforward:

import asyncio
import replicate

# https://replicate.com/stability-ai/sdxl
model_version = "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b"
prompts = [
    f"A chariot pulled by a team of {count} rainbow unicorns"
    for count in ["two", "four", "six", "eight"]
]

async with asyncio.TaskGroup() as tg:
    tasks = [
        tg.create_task(replicate.async_run(model_version, input={"prompt": prompt}))
        for prompt in prompts
    ]

results = await asyncio.gather(*tasks)
print(results)

Under the hood, Client manages an httpx.Client and an httpx.AsyncClient, which handle calls to _request and _async_request, respectively. Both are created lazily, so API consumers using only sync or only async functionality won't be affected by functionality they aren't using.

Implementation-wise, sync and async variants have separate code paths. This creates nontrivial amounts of duplication, but its benefits to clarity and performance justify those costs. For instance, it'd have been nice if the sync variant were implemented as a blocking call to the async variant, but that would require starting an event loop, which has additional overhead and causes problems if done within an existing event loop.

Alternative to #76
Resolves #145
Resolves #107
Resolves #74

[zeke]

What does good-names mean?

[mattt]

They're good names, Zeke.

[mattt]

(But seriously, this is to stop pylint from complaining about variables named id)

[zeke]

Wow this looks like it was a lot of work.

async variants with the async_ prefix

Is this a cow-path convention in Python land for async variants of synchronous methods?

What are your thoughts on documenting these methods? Do something now or kick the can down the road until we have legit API reference docs for this client?

[mattt]

@zeke Thanks for the review!

Wow this looks like it was a lot of work.

It sure was 🫠

async variants with the async_ prefix

Is this a cow-path convention in Python land for async variants of synchronous methods?

As I understand it, yes. Another convention I've seen is prepending just a (so aget instead of async_get), but I think you'll agree with my choice here.

The other (anti) pattern I've seen is to make it a separate package / module. So long as sync and async can coexist in replicate without causing problems, then I'm keen to keep them together.

What are your thoughts on documenting these methods? Do something now or kick the can down the road until we have legit API reference docs for this client?

I added a callout towards the top of the README about methods having async variants. As for API reference documentation, the docstrings in the code itself are in good shape. Some Python projects I like use mkdocstrings to generate those automatically. So that's something to consider ¹.

Footnotes

1. I'm a fan of publishing docs where people expect to see them. For our Swift library, I recently set up docs hosted on Swift Package Index. For Go, I recently linked to pkg.go.dev. ↩

[antont]

The other (anti) pattern I've seen is to make it a separate package / module. So long as sync and async can coexist in replicate without causing problems, then I'm keen to keep them together.

+1

Some libs use an unasync tool to only maintain the async code, and generate the sync code from it using a script. For example https://github.com/ioxiocom/firedantic/blob/main/README.md#about-sync-and-async-versions-of-library documents their use which is implemented in https://github.com/ioxiocom/firedantic/blob/main/unasync.py

[mattt]

@antont Interesting! Thank you for sharing that. The code base is at a manageable size now so that doing things manually isn't onerous, but I'll definitely want to set up automation like this once things get out of hand.