NVIDIA NeMo Gym

import math
import shlex
import uuid
from pathlib import Path

import superserve
from superserve import AsyncSandbox, NetworkConfig, NotFoundError

from nemo_gym.sandbox import (
    SandboxCreateError,
    SandboxExecResult,
    SandboxHandle,
    SandboxSpec,
    SandboxStatus,
)

STATUSES = {
    superserve.SandboxStatus.STARTING: SandboxStatus.STARTING,
    superserve.SandboxStatus.RESUMING: SandboxStatus.STARTING,
    superserve.SandboxStatus.ACTIVE: SandboxStatus.RUNNING,
    superserve.SandboxStatus.MIGRATING: SandboxStatus.RUNNING,
    superserve.SandboxStatus.PAUSING: SandboxStatus.PAUSED,
    superserve.SandboxStatus.PAUSED: SandboxStatus.PAUSED,
    superserve.SandboxStatus.FAILED: SandboxStatus.ERROR,
    superserve.SandboxStatus.DELETED: SandboxStatus.STOPPED,
}


class SuperserveProvider:
    """Runs each NeMo Gym sandbox as a Superserve sandbox."""

    name = "superserve"

    def __init__(
        self, template: str = "superserve/base", deny_out: list[str] | None = None
    ) -> None:
        self._template = template
        self._network = NetworkConfig(deny_out=deny_out) if deny_out else None

    async def create(self, spec: SandboxSpec) -> SandboxHandle:
        try:
            sandbox = await AsyncSandbox.create(
                name=f"nemo-gym-{uuid.uuid4().hex[:8]}",
                # SandboxSpec.image names a Superserve template. The template
                # also sets the sandbox's CPU, memory, and disk.
                from_template=spec.image or self._template,
                # Pause the sandbox after ttl_s, then delete it at once.
                timeout_seconds=int(spec.ttl_s or 3600),
                auto_delete_seconds=0,
                metadata=spec.metadata or None,
                env_vars=spec.env or None,
                network=self._network,
            )
        except superserve.SandboxError as error:
            raise SandboxCreateError(str(error)) from error
        try:
            if spec.workdir:
                mkdir = f"mkdir -p {shlex.quote(spec.workdir)}"
                result = await sandbox.commands.run(mkdir)
                if result.exit_code != 0:
                    raise SandboxCreateError(
                        f"Could not create workdir {spec.workdir!r}: {result.stderr}"
                    )
        except BaseException:
            await sandbox.kill()
            raise
        return SandboxHandle(sandbox.id, self.name, raw=sandbox)

    async def exec(
        self,
        handle: SandboxHandle,
        command: str,
        *,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout_s: int | float | None = None,
        user: str | int | None = None,
    ) -> SandboxExecResult:
        sandbox: AsyncSandbox = handle.raw
        result = await sandbox.commands.run(
            command,
            cwd=cwd,
            env=env,
            user=None if user is None else str(user),
            # A command that runs past the timeout is stopped with exit code 124.
            timeout_seconds=math.ceil(timeout_s) if timeout_s else None,
        )
        return SandboxExecResult(
            stdout=result.stdout, stderr=result.stderr, return_code=result.exit_code
        )

    async def upload_file(
        self, handle: SandboxHandle, source_path: Path, target_path: str
    ) -> None:
        sandbox: AsyncSandbox = handle.raw
        await sandbox.files.write(target_path, source_path.read_bytes())

    async def download_file(
        self, handle: SandboxHandle, source_path: str, target_path: Path
    ) -> None:
        sandbox: AsyncSandbox = handle.raw
        target_path.parent.mkdir(parents=True, exist_ok=True)
        target_path.write_bytes(await sandbox.files.read(source_path))

    async def status(self, handle: SandboxHandle) -> SandboxStatus:
        if handle.raw is None:  # Closed by close().
            return SandboxStatus.STOPPED
        sandbox: AsyncSandbox = handle.raw
        try:
            info = await sandbox.get_info()
        except NotFoundError:
            return SandboxStatus.STOPPED
        return STATUSES.get(info.status, SandboxStatus.UNKNOWN)

    async def close(self, handle: SandboxHandle) -> None:
        sandbox: AsyncSandbox | None = handle.raw
        if sandbox is None:
            return
        try:
            await sandbox.kill()
        except NotFoundError:
            pass  # Already deleted.
        handle.raw = None

    async def aclose(self) -> None:
        return None

import re
from decimal import Decimal

from fastapi import FastAPI
from pydantic import BaseModel

from nemo_gym.base_resources_server import (
    BaseResourcesServerConfig,
    BaseRunRequest,
    BaseVerifyRequest,
    BaseVerifyResponse,
    SimpleResourcesServer,
)
from nemo_gym.global_config import get_global_config_dict
from nemo_gym.sandbox import (
    AsyncSandbox,
    SandboxSpec,
    register_provider,
    resolve_provider_config,
)
from resources_servers.superserve_python.superserve_provider import SuperserveProvider

register_provider("superserve", SuperserveProvider)


class SuperservePythonConfig(BaseResourcesServerConfig):
    sandbox_provider: str  # Name of the sandbox config block, for example "sandbox".


class RunPythonRequest(BaseModel):
    code: str


class RunPythonResponse(BaseModel):
    exit_code: int
    stdout: str
    stderr: str


class SuperservePythonRunRequest(BaseRunRequest):
    expected_answer: str


class SuperservePythonVerifyRequest(SuperservePythonRunRequest, BaseVerifyRequest):
    pass


class SuperservePythonServer(SimpleResourcesServer):
    config: SuperservePythonConfig

    def setup_webserver(self) -> FastAPI:
        app = super().setup_webserver()
        app.post("/run_python")(self.run_python)
        return app

    async def run_python(self, body: RunPythonRequest) -> RunPythonResponse:
        provider = resolve_provider_config(
            self.config.sandbox_provider, get_global_config_dict()
        )
        spec = SandboxSpec(
            workdir="/workspace", ttl_s=600, files={"/workspace/main.py": body.code}
        )
        # Each tool call gets a new sandbox, which is deleted when the block exits.
        async with AsyncSandbox(provider, spec) as sandbox:
            await sandbox.start()
            result = await sandbox.exec("python3 main.py", timeout_s=30)
        return RunPythonResponse(
            exit_code=result.return_code,
            stdout=result.stdout or "",
            stderr=result.stderr or "",
        )

    async def verify(  # type: ignore[override]
        self, body: SuperservePythonVerifyRequest
    ) -> BaseVerifyResponse:
        # Reward 1.0 when the last number in the final answer equals the expected one.
        text = body.response.output_text.replace(",", "")
        # Match whole numbers such as 249, -3, 249.0, .249, and 2.49e2.
        numbers = re.findall(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?", text)
        correct = numbers and Decimal(numbers[-1]) == Decimal(body.expected_answer)
        return BaseVerifyResponse(**body.model_dump(), reward=1.0 if correct else 0.0)


if __name__ == "__main__":
    SuperservePythonServer.run_webserver()