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

chore(internal): minor test changes #1874

Merged
merged 1 commit into from
Nov 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ dev-dependencies = [
"dirty-equals>=0.6.0",
"importlib-metadata>=6.7.0",
"rich>=13.7.1",
"nest_asyncio==1.6.0"
]

[tool.rye.scripts]
Expand Down
1 change: 1 addition & 0 deletions requirements-dev.lock
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ mdurl==0.1.2
mypy==1.13.0
mypy-extensions==1.0.0
# via mypy
nest-asyncio==1.6.0
nodeenv==1.8.0
# via pyright
nox==2023.4.22
Expand Down
90 changes: 40 additions & 50 deletions src/openai/_utils/_sync.py
Original file line number Diff line number Diff line change
@@ -1,56 +1,62 @@
from __future__ import annotations

import sys
import asyncio
import functools
from typing import TypeVar, Callable, Awaitable
import contextvars
from typing import Any, TypeVar, Callable, Awaitable
from typing_extensions import ParamSpec

import anyio
import anyio.to_thread

from ._reflection import function_has_argument

T_Retval = TypeVar("T_Retval")
T_ParamSpec = ParamSpec("T_ParamSpec")


# copied from `asyncer`, https://github.com/tiangolo/asyncer
def asyncify(
function: Callable[T_ParamSpec, T_Retval],
*,
cancellable: bool = False,
limiter: anyio.CapacityLimiter | None = None,
) -> Callable[T_ParamSpec, Awaitable[T_Retval]]:
if sys.version_info >= (3, 9):
to_thread = asyncio.to_thread
else:
# backport of https://docs.python.org/3/library/asyncio-task.html#asyncio.to_thread
# for Python 3.8 support
async def to_thread(
func: Callable[T_ParamSpec, T_Retval], /, *args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs
) -> Any:
"""Asynchronously run function *func* in a separate thread.

Any *args and **kwargs supplied for this function are directly passed
to *func*. Also, the current :class:`contextvars.Context` is propagated,
allowing context variables from the main thread to be accessed in the
separate thread.

Returns a coroutine that can be awaited to get the eventual result of *func*.
"""
loop = asyncio.events.get_running_loop()
ctx = contextvars.copy_context()
func_call = functools.partial(ctx.run, func, *args, **kwargs)
return await loop.run_in_executor(None, func_call)


# inspired by `asyncer`, https://github.com/tiangolo/asyncer
def asyncify(function: Callable[T_ParamSpec, T_Retval]) -> Callable[T_ParamSpec, Awaitable[T_Retval]]:
"""
Take a blocking function and create an async one that receives the same
positional and keyword arguments, and that when called, calls the original function
in a worker thread using `anyio.to_thread.run_sync()`. Internally,
`asyncer.asyncify()` uses the same `anyio.to_thread.run_sync()`, but it supports
keyword arguments additional to positional arguments and it adds better support for
autocompletion and inline errors for the arguments of the function called and the
return value.

If the `cancellable` option is enabled and the task waiting for its completion is
cancelled, the thread will still run its course but its return value (or any raised
exception) will be ignored.
positional and keyword arguments. For python version 3.9 and above, it uses
asyncio.to_thread to run the function in a separate thread. For python version
3.8, it uses locally defined copy of the asyncio.to_thread function which was
introduced in python 3.9.

Use it like this:
Usage:

```Python
def do_work(arg1, arg2, kwarg1="", kwarg2="") -> str:
# Do work
return "Some result"
```python
def blocking_func(arg1, arg2, kwarg1=None):
# blocking code
return result


result = await to_thread.asyncify(do_work)("spam", "ham", kwarg1="a", kwarg2="b")
print(result)
result = asyncify(blocking_function)(arg1, arg2, kwarg1=value1)
```

## Arguments

`function`: a blocking regular callable (e.g. a function)
`cancellable`: `True` to allow cancellation of the operation
`limiter`: capacity limiter to use to limit the total amount of threads running
(if omitted, the default limiter is used)

## Return

Expand All @@ -60,22 +66,6 @@ def do_work(arg1, arg2, kwarg1="", kwarg2="") -> str:
"""

async def wrapper(*args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs) -> T_Retval:
partial_f = functools.partial(function, *args, **kwargs)

# In `v4.1.0` anyio added the `abandon_on_cancel` argument and deprecated the old
# `cancellable` argument, so we need to use the new `abandon_on_cancel` to avoid
# surfacing deprecation warnings.
if function_has_argument(anyio.to_thread.run_sync, "abandon_on_cancel"):
return await anyio.to_thread.run_sync(
partial_f,
abandon_on_cancel=cancellable,
limiter=limiter,
)

return await anyio.to_thread.run_sync(
partial_f,
cancellable=cancellable,
limiter=limiter,
)
return await to_thread(function, *args, **kwargs)

return wrapper
38 changes: 38 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,14 @@

import gc
import os
import sys
import json
import asyncio
import inspect
import subprocess
import tracemalloc
from typing import Any, Union, cast
from textwrap import dedent
from unittest import mock
from typing_extensions import Literal

Expand Down Expand Up @@ -1766,3 +1769,38 @@ def retry_handler(_request: httpx.Request) -> httpx.Response:
) as response:
assert response.retries_taken == failures_before_success
assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success

def test_get_platform(self) -> None:
# A previous implementation of asyncify could leave threads unterminated when
# used with nest_asyncio.
#
# Since nest_asyncio.apply() is global and cannot be un-applied, this
# test is run in a separate process to avoid affecting other tests.
test_code = dedent("""
import asyncio
import nest_asyncio
import threading

from openai._utils import asyncify
from openai._base_client import get_platform

async def test_main() -> None:
result = await asyncify(get_platform)()
print(result)
for thread in threading.enumerate():
print(thread.name)

nest_asyncio.apply()
asyncio.run(test_main())
""")
with subprocess.Popen(
[sys.executable, "-c", test_code],
text=True,
) as process:
try:
process.wait(2)
if process.returncode:
raise AssertionError("calling get_platform using asyncify resulted in a non-zero exit code")
except subprocess.TimeoutExpired as e:
process.kill()
raise AssertionError("calling get_platform using asyncify resulted in a hung process") from e