[FEATURE] uv provider

[burtenshaw]

This PR adds a uv provider to open env core which can be used to run envs on spaces without docker. The main intention here is to make a pure python option for notebook environments like colab.

This is a working example based on a specific env, which is the result of #160 .

!pip install git+https://huggingface.co/spaces/burtenshaw/echo-cli

from envs.echo_env.client import EchoEnv, EchoAction

# Pulls `openenv/echo_env` from the Hub and runs it via `uv run`.
client = EchoEnv.from_hub(
    "burtenshaw/echo-cli",
    use_docker=False,
)

# use it
reset_result = client.reset()
step_result = client.step(EchoAction(message="Hello from the Hub!"))
print("Echoed:", step_result.observation.echoed_message)
print("Reward:", step_result.reward)

# close it
client.close()  # stops the uv subprocess cleanly

We can interact directly with the uv provided env like so:

from src.core.containers.runtime.uv_provider import UVProvider

# Pulls `openenv/echo_env` from the Hub and runs it via `uv run`.
client = UVProvider("burtenshaw/echo-cli")

# use it
base_url = client.start()
client.wait_for_ready(base_url)

print(base_url)
print(client.base_url)

# close it
client.stop()  # stops the uv subprocess cleanly

[burtenshaw]

@Wauplin is the maintain of hf and he agreed to take a look at this PR.

[Wauplin]

Hey @burtenshaw , thanks for the ping! Happy to review a PR :) I've left quite a lot of comments on how I would implement things if starting the project right now. Feel free to ignore when not applicable or if I did not take into account other relevant parts of the library. In general my main comment is that I always tend to reduce the number of available parameters ("if no-one asked for it, let's not implement it") and isolate logic ("let's not have Hub-logic or docker-logic in uv provider").

Given that it's the second provider implemented, I think it's good to take some time to think about which API makes the most sense for the future of the library.

Let me know if you have any questions / comments 🤗

[Wauplin]

Since uv is not per-se a container, I would rename the abstract class RuntimeProvider (or simply Runtime?) and the abstract methods start, stop, and wait_for_ready (instead of start_container, etc.). I feel that semantically it would be more accurate.

(just my 2c, feel free to ignore 🤗 )

[burtenshaw]

I agree with you @Wauplin , we should clarify naming here. However, I'm reluctant to let this PR swell whilst we're trying to move fast.

For now, I've added a new ABC named RuntimeProvider which has no container specific methods and is used by the UVProvider.

[Wauplin]

Understandable! Works like this as well :)

[Wauplin]

Suggested change

	class UVProvider(ContainerProvider):
	class UVRuntime(RuntimeProvider):

(naming suggestion to be aligned with the "runtime/" folder)

[Wauplin]

I find this signature very hard to understand without context since it's mixing kwargs for docker and for uv. Also I don't think one needs use_docker, provider, and runner.

[Wauplin]

My suggestion would be to remove provider and runner. Also remove project_url and connect_host. Have a single env (instead of env for docker and extra_env for uv). And use some typed dict + typing.overload so that IDEs can have correct autocompletion. Here is a simplified example:

from typing import Any, Dict, NotRequired, TypedDict, Unpack, overload


class DockerKwargs(TypedDict, total=False):
    tag: NotRequired[str]
    env: NotRequired[Dict[str, str]]


class UVProvider(TypedDict):
    host: NotRequired[str]
    port: NotRequired[int]
    reload: NotRequired[bool]
    timeout_s: NotRequired[float]
    env: NotRequired[Dict[str, str]]


@overload
def from_hub(repo_id: str, *, use_docker: bool = True, **kwargs: Unpack[DockerKwargs]) -> str: ...


@overload
def from_hub(repo_id: str, *, use_docker: bool = False, **kwargs: Unpack[UVProvider]) -> str: ...


def from_hub(repo_id: str, *, use_docker: bool = False, **kwargs: Any) -> str:
    raise NotImplementedError()

[Wauplin]

Note that I'm not a fan of overloads and typed dict but it's the only solution I see to correctly document the signature while keeping a single method.

Another solution is to have from_hub_docker and from_hub_uv (more explicit but less elegant)

[burtenshaw]

Thanks for this. I've simplified the signatures right down, but I haven't added type overloading in this PR.

[Wauplin]

Once the runner is defined and instantiated, I feel that the start and wait calls should be made in the __init__ instead. This way you get the same behavior both in from_hub and from_docker. With current implementation, one stops the runner if failing to start, the other don't (which is not consistent).

[Wauplin]

Suggested change

	uv_runner = runner or UVProvider(
	repo_id=repo_id,
	host=host,
	port=port,
	reload=reload,
	project_url=project_url,
	connect_host=connect_host,
	extra_env=extra_env,
	)
	runner = UVProvider(
	project=f"git+https://huggingface.co/spaces/{repo_id}",
	host=host,
	port=port,
	reload=reload,
	env=env,
	)

With all the suggestions above, that would be the runner call (i.e. provide a project url, remove connect_host, remove project_url, rename extra_env to env, and remove runner). I don't think final user loose much in this simplification.

[burtenshaw]

This PR now exposes uvicorn's workers and contributes to concurrency as mentioned in #194