diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py new file mode 100644 index 0000000000..3225f54770 --- /dev/null +++ b/nemo_gym/sandbox/__init__.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public sandbox API for NeMo Gym.""" + +from nemo_gym.sandbox.api import AsyncSandbox, Sandbox, rewrite_image +from nemo_gym.sandbox.providers import ( + SandboxBatchCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, + create_provider, + get_provider_class, + list_providers, + register_provider, +) + + +__all__ = [ + "Sandbox", + "AsyncSandbox", + "SandboxBatchCreateError", + "SandboxCreateVerificationError", + "SandboxExecResult", + "SandboxHandle", + "SandboxProvider", + "SandboxSpec", + "create_provider", + "get_provider_class", + "list_providers", + "register_provider", + "rewrite_image", +] diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py new file mode 100644 index 0000000000..c9c3d74aa4 --- /dev/null +++ b/nemo_gym/sandbox/api.py @@ -0,0 +1,452 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider-neutral public sandbox API. + +This module is the boundary Gym code should use when it needs a sandbox. +Provider packages implement the lower-level async protocol; callers use +``AsyncSandbox`` in async code and ``Sandbox`` in synchronous integrations. +""" + +import asyncio +import threading +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from concurrent.futures import Future +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any, TypeVar, cast + +from nemo_gym.sandbox.config import SandboxProviderConfig +from nemo_gym.sandbox.observability import ( + current_recorder, + ensure_env_recorder, + observability_span, + push_event_context, + reset_current_recorder, + reset_event_context, + set_current_recorder, +) +from nemo_gym.sandbox.providers import ( + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, + create_provider, +) + + +T = TypeVar("T") + + +def rewrite_image(image: str | None, rewrites: list[dict[str, str]]) -> str | None: + """Apply ordered image-prefix rewrites used by sandbox configs.""" + if image is None: + return None + for rewrite in rewrites: + from_prefix = rewrite["from"] + to_prefix = rewrite["to"] + if image.startswith(from_prefix): + return to_prefix + image[len(from_prefix) :] + return image + + +class AsyncSandbox: + """Async public facade for provider-backed sandbox operations.""" + + def __init__( + self, + provider: SandboxProviderConfig | SandboxProvider, + *, + observability_context: dict[str, Any] | None = None, + ) -> None: + self._provider = ( + create_provider(cast(SandboxProviderConfig, provider)) if isinstance(provider, Mapping) else provider + ) + self._observability_context = dict(observability_context or {}) + self._handle_observability_context: dict[str, dict[str, Any]] = {} + + @property + def provider_name(self) -> str: + return self._provider.name + + def _spec_observability_context(self, spec: SandboxSpec) -> dict[str, Any]: + metadata = dict(spec.metadata) + context: dict[str, Any] = { + "provider": self.provider_name, + "environment_type": "sandbox", + } + + for key in ("benchmark", "harness", "instance_id", "trajectory_id", "trial_name"): + value = metadata.get(key) + if value is not None: + context[key] = value + + if "harness" not in context and metadata.get("nemo_gym_agent") is not None: + context["harness"] = metadata["nemo_gym_agent"] + + if "trajectory_id" not in context: + for key in ("instance_id", "trial_name", "environment_name", "harbor_instance_id"): + value = metadata.get(key) + if value is not None: + context["trajectory_id"] = value + break + + return {**self._observability_context, **context} + + def _handle_context(self, handle: SandboxHandle) -> dict[str, Any]: + return self._handle_observability_context.get( + handle.sandbox_id, + { + **self._observability_context, + "provider": self.provider_name, + "environment_type": "sandbox", + "sandbox_id": handle.sandbox_id, + }, + ) + + @asynccontextmanager + async def _observed(self, attributes: dict[str, Any]) -> AsyncIterator[None]: + recorder = current_recorder() or ensure_env_recorder() + recorder_token = None + if recorder is not None and current_recorder() is None: + recorder_token = set_current_recorder(recorder) + context_token = push_event_context(attributes) + try: + yield + finally: + reset_event_context(context_token) + if recorder_token is not None: + reset_current_recorder(recorder_token) + + def _remember_handle(self, handle: SandboxHandle, context: dict[str, Any]) -> None: + handle_context = {**context, "sandbox_id": handle.sandbox_id} + self._handle_observability_context[handle.sandbox_id] = handle_context + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + context = self._spec_observability_context(spec) + async with self._observed(context): + async with observability_span( + "sandbox.start", + phase="startup", + attributes={ + "provider": self.provider_name, + "image": spec.image, + }, + ): + handle = await self._provider.create(spec) + self._remember_handle(handle, context) + return handle + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + context = self._spec_observability_context(spec) + async with self._observed(context): + async with observability_span( + "sandbox.start_batch", + phase="startup", + attributes={ + "provider": self.provider_name, + "count": count, + "allow_partial": allow_partial, + "image": spec.image, + }, + ): + handles = await self._provider.create_batch(spec, count, allow_partial=allow_partial) + for handle in handles: + self._remember_handle(handle, context) + return handles + + async def connect(self, sandbox_id: str) -> SandboxHandle: + context = { + **self._observability_context, + "provider": self.provider_name, + "environment_type": "sandbox", + "sandbox_id": sandbox_id, + } + async with self._observed(context): + handle = await self._provider.connect(sandbox_id) + self._remember_handle(handle, context) + return handle + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + context = self._handle_context(handle) + async with self._observed(context): + async with observability_span( + "trajectory.tool", + phase="execution", + attributes={ + "provider": self.provider_name, + "sandbox_id": handle.sandbox_id, + "cwd": cwd, + "timeout_s": timeout_s, + "user": user, + "command": command, + }, + ): + return await self._provider.exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + ) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + async with self._observed(self._handle_context(handle)): + await self._provider.write_file(handle, target_path, data) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + async with self._observed(self._handle_context(handle)): + return await self._provider.read_file(handle, source_path) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + async with self._observed(self._handle_context(handle)): + await self._provider.upload_file(handle, source_path, target_path) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + async with self._observed(self._handle_context(handle)): + await self._provider.download_file(handle, source_path, target_path) + + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + context = self._handle_context(handle) + async with self._observed(context): + async with observability_span( + "sandbox.cleanup", + phase="cleanup", + attributes={ + "provider": self.provider_name, + "sandbox_id": handle.sandbox_id, + "delete": delete, + }, + ): + try: + await self._provider.close(handle, delete=delete) + finally: + self._handle_observability_context.pop(handle.sandbox_id, None) + + async def delete(self, handle: SandboxHandle) -> None: + await self.close(handle, delete=True) + + async def aclose(self) -> None: + self._handle_observability_context.clear() + close_provider = getattr(self._provider, "aclose", None) + if close_provider is not None: + await close_provider() + + async def shutdown(self) -> None: + await self.aclose() + + async def __aenter__(self) -> "AsyncSandbox": + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.aclose() + + def handle_reference(self, handle: SandboxHandle) -> Any: + make_reference = getattr(self._provider, "handle_reference", None) + if make_reference is None: + return handle + return make_reference(handle) + + async def materialize_handle(self, value: Any) -> SandboxHandle: + materialize = getattr(self._provider, "materialize_handle", None) + if materialize is None: + if isinstance(value, SandboxHandle): + return value + raise ValueError(f"Provider {self.provider_name!r} cannot materialize handle references") + result = materialize(value) + if hasattr(result, "__await__"): + result = await result + if not isinstance(result, SandboxHandle): + raise TypeError(f"materialize_handle must return SandboxHandle, got {type(result).__name__}") + return result + + +class _AsyncLoopRunner: + """Run async sandbox operations for sync integrations on one private loop.""" + + def __init__(self) -> None: + self._loop = asyncio.new_event_loop() + self._ready = threading.Event() + self._closed = False + self._thread = threading.Thread(target=self._run_loop, name="nemo-gym-sandbox-sync-loop", daemon=True) + self._thread.start() + self._ready.wait() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._ready.set() + self._loop.run_forever() + + def _ensure_can_block(self, operation: str) -> None: + if self._closed or self._loop.is_closed(): + raise RuntimeError("Sandbox sync loop is closed") + try: + asyncio.get_running_loop() + except RuntimeError: + return + raise RuntimeError(f"Sandbox.{operation}() is blocking; use AsyncSandbox in async code instead.") + + def call(self, operation: str, func: Callable[[], T]) -> T: + self._ensure_can_block(operation) + future: Future[T] = Future() + + def invoke() -> None: + try: + future.set_result(func()) + except BaseException as e: + future.set_exception(e) + + self._loop.call_soon_threadsafe(invoke) + return future.result() + + def run(self, operation: str, awaitable_factory: Callable[[], Awaitable[T]]) -> T: + self._ensure_can_block(operation) + future = asyncio.run_coroutine_threadsafe(awaitable_factory(), self._loop) + return future.result() + + def close(self) -> None: + if self._closed: + return + self._closed = True + if not self._loop.is_closed(): + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5) + self._loop.close() + + +class Sandbox: + """Sync public facade for provider-backed sandbox operations.""" + + def __init__( + self, + provider: SandboxProviderConfig | SandboxProvider, + *, + observability_context: dict[str, Any] | None = None, + ) -> None: + self._runner = _AsyncLoopRunner() + try: + self._async_sandbox = self._runner.call( + "__init__", + lambda: AsyncSandbox(provider, observability_context=observability_context), + ) + except BaseException: + self._runner.close() + raise + self._closed = False + + @property + def provider_name(self) -> str: + return self._runner.call("provider_name", lambda: self._async_sandbox.provider_name) + + def create(self, spec: SandboxSpec) -> SandboxHandle: + return self._runner.run("create", lambda: self._async_sandbox.create(spec)) + + def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + return self._runner.run( + "create_batch", + lambda: self._async_sandbox.create_batch(spec, count, allow_partial=allow_partial), + ) + + def connect(self, sandbox_id: str) -> SandboxHandle: + return self._runner.run("connect", lambda: self._async_sandbox.connect(sandbox_id)) + + def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + return self._runner.run( + "exec", + lambda: self._async_sandbox.exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + ), + ) + + def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + self._runner.run("write_file", lambda: self._async_sandbox.write_file(handle, target_path, data)) + + def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + return self._runner.run("read_file", lambda: self._async_sandbox.read_file(handle, source_path)) + + def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self._runner.run("upload_file", lambda: self._async_sandbox.upload_file(handle, source_path, target_path)) + + def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + self._runner.run("download_file", lambda: self._async_sandbox.download_file(handle, source_path, target_path)) + + def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + self._runner.run("close", lambda: self._async_sandbox.close(handle, delete=delete)) + + def delete(self, handle: SandboxHandle) -> None: + self.close(handle, delete=True) + + def shutdown(self) -> None: + if self._closed: + return + self._closed = True + try: + self._runner.run("shutdown", self._async_sandbox.shutdown) + finally: + self._runner.close() + + def handle_reference(self, handle: SandboxHandle) -> Any: + return self._runner.call("handle_reference", lambda: self._async_sandbox.handle_reference(handle)) + + def materialize_handle(self, value: Any) -> SandboxHandle: + return self._runner.run("materialize_handle", lambda: self._async_sandbox.materialize_handle(value)) + + def __enter__(self) -> "Sandbox": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.shutdown() + + def __del__(self) -> None: + if hasattr(self, "_closed") and not self._closed: + try: + self.shutdown() + except Exception: + pass diff --git a/nemo_gym/sandbox/config.py b/nemo_gym/sandbox/config.py new file mode 100644 index 0000000000..3467c382f2 --- /dev/null +++ b/nemo_gym/sandbox/config.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Typed configuration for sandbox providers. + +Defaults for these fields belong in YAML exemplars. Code should require keys +from enabled configs instead of silently supplying behavior here. +""" + +from typing import Any, NotRequired, TypedDict + + +class SandboxProviderConfig(TypedDict): + """Underlying runtime and infrastructure provider. + + Keys: + name: Provider registry name, for example ``opensandbox``. + kwargs: Provider-specific constructor settings such as OpenSandbox + domain, API key, or proxy mode. + """ + + name: str + kwargs: NotRequired[dict[str, Any]] diff --git a/nemo_gym/sandbox/observability/__init__.py b/nemo_gym/sandbox/observability/__init__.py new file mode 100644 index 0000000000..4456ad5ded --- /dev/null +++ b/nemo_gym/sandbox/observability/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sandbox eval observability helpers.""" + +from nemo_gym.sandbox.observability.recorder import ( + SandboxRecorder, + build_recorder_from_config, + build_recorder_from_env, + current_recorder, + ensure_env_recorder, + event_context, + observability_span, + observability_sync_span, + push_event_context, + record_event, + reset_current_recorder, + reset_event_context, + set_current_recorder, + use_recorder, +) +from nemo_gym.sandbox.observability.traces import export_trace_artifacts + + +__all__ = [ + "SandboxRecorder", + "build_recorder_from_config", + "build_recorder_from_env", + "current_recorder", + "ensure_env_recorder", + "event_context", + "export_trace_artifacts", + "observability_span", + "observability_sync_span", + "push_event_context", + "record_event", + "reset_current_recorder", + "reset_event_context", + "set_current_recorder", + "use_recorder", +] diff --git a/nemo_gym/sandbox/observability/recorder.py b/nemo_gym/sandbox/observability/recorder.py new file mode 100644 index 0000000000..bc7cc412fe --- /dev/null +++ b/nemo_gym/sandbox/observability/recorder.py @@ -0,0 +1,1005 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Context-scoped recorder for sandbox eval observability.""" + +from __future__ import annotations + +import atexit +import json +import os +import re +import threading +import time +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar, Token +from pathlib import Path +from typing import Any, Iterator + +from opentelemetry import trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor, SpanExporter +from opentelemetry.trace import Span, SpanKind, Status, StatusCode + +from nemo_gym.sandbox.observability.traces import ( + SCOPE_NAME, + SCOPE_VERSION, + JsonSpanExporter, + export_trace_artifacts, +) + + +G_CURRENT_RECORDER: ContextVar[SandboxRecorder | None] = ContextVar( + "nemo_gym_sandbox_observability_recorder", + default=None, +) +G_EVENT_CONTEXT: ContextVar[dict[str, Any]] = ContextVar( + "nemo_gym_sandbox_observability_context", + default={}, +) +G_ENV_RECORDER: SandboxRecorder | None = None +G_ENV_RECORDER_LOCK = threading.Lock() + + +class SandboxRecorder: + """Context-scoped sandbox observability recorder backed by OpenTelemetry.""" + + def __init__( + self, + *, + output_dir: Path | None = None, + otel: dict[str, Any] | None = None, + run_id: str | None = None, + run_span_name: str | None = None, + export_traces: bool | None = None, + ) -> None: + self.output_dir = output_dir + self.otel = dict(otel or {}) + self.run_id = run_id + self.run_span_name = _eval_span_name( + run_span_name or self.otel.get("run_span_name") or self.otel.get("job_name") or run_id or "sandbox.run" + ) + self._configured_trace_exporters = _configured_exporters(self.otel, "traces") + self._configured_metric_exporters = _configured_exporters(self.otel, "metrics") + self.export_traces = _local_trace_export_enabled( + output_dir=output_dir, + export_traces=export_traces, + trace_exporters=self._configured_trace_exporters, + ) + if self.export_traces and output_dir is None: + raise ValueError("sandbox observability output_dir is required for local trace export") + self.attribute_aliases = _string_map(self.otel.get("attribute_aliases")) + self.command_titles = _command_title_config(self.otel.get("command_titles") or self.otel.get("command_title")) + self.metric_attribute_keys = tuple(str(key) for key in self.otel.get("metric_attribute_keys") or ()) + self.resource_attributes = safe_attributes(self.otel.get("resource_attributes") or {}) + self.local_service_name_strategy = str(self.otel.get("local_service_name_strategy") or "span_section") + self._closed = False + self._service_name = str(self.otel.get("service_name") or "") or None + self._local_span_exporter = JsonSpanExporter() if self.export_traces else None + self._tracer_provider = TracerProvider(resource=self._resource()) + if self._local_span_exporter is not None: + self._tracer_provider.add_span_processor(SimpleSpanProcessor(self._local_span_exporter)) + self._meter_provider = None + self._duration_histogram = None + self._phase_duration_histograms = {} + self._counter = None + self._configure_live_exporters() + self._configure_metrics() + self._trajectory_spans: dict[tuple[str, str], Span] = {} + self._run_span = self._start_span( + self.run_span_name, + attributes=safe_attributes( + { + "run_id": run_id, + "span.role": "eval.run", + "span.section": "eval", + } + ), + ) + if self.output_dir is not None: + self.output_dir.mkdir(parents=True, exist_ok=True) + self.record_event("lifecycle", "run.start", attributes={"run_id": run_id}) + + def record_event( + self, + event_type: str, + name: str, + *, + attributes: dict[str, Any] | None = None, + trace_id: str | None = None, + span_id: str | None = None, + parent_span_id: str | None = None, + timestamp_unix_s: float | None = None, + monotonic_s: float | None = None, + ) -> None: + """Record one OpenTelemetry event on the active span.""" + del trace_id, span_id, parent_span_id, monotonic_s + attrs = safe_attributes({**G_EVENT_CONTEXT.get(), **(attributes or {})}) + self._record_otel_event(event_type, name, attrs, timestamp_unix_s=timestamp_unix_s) + self._record_event_metrics(event_type=event_type, name=name, attrs=attrs) + + @asynccontextmanager + async def span( + self, + name: str, + *, + phase: str | None = None, + attributes: dict[str, Any] | None = None, + ) -> Iterator[None]: + """Record an async operation as an OpenTelemetry span.""" + with self._span_context(name, phase=phase, attributes=attributes): + yield + + @contextmanager + def sync_span( + self, + name: str, + *, + phase: str | None = None, + attributes: dict[str, Any] | None = None, + ) -> Iterator[None]: + """Record a synchronous operation as an OpenTelemetry span.""" + with self._span_context(name, phase=phase, attributes=attributes): + yield + + def finalize(self) -> None: + """Flush local trace artifacts and OTel exporters.""" + if self._closed: + return + self._closed = True + self.record_event("lifecycle", "run.end", attributes={"run_id": self.run_id}) + self._end_open_spans() + if self.export_traces: + try: + self._export_trace_artifacts() + except Exception as e: + self.record_event( + "error", + "observability.trace_export_error", + attributes={"error_type": type(e).__name__, "error": str(e)}, + ) + self._shutdown_otel() + + @contextmanager + def _span_context( + self, + name: str, + *, + phase: str | None, + attributes: dict[str, Any] | None, + ) -> Iterator[None]: + start_monotonic = time.monotonic() + span_attrs = safe_attributes({**G_EVENT_CONTEXT.get(), "phase": phase, **(attributes or {})}) + record_exception_stacktrace = _as_bool(span_attrs.pop("_record_exception_stacktrace", True)) + operation_name = _span_name(name) + display_name = self._operation_span_name(operation_name, span_attrs) + with self._start_as_current_span( + display_name, + attributes=self._span_attributes(operation_name, span_attrs, display_name=display_name), + context=self._parent_context(span_attrs), + kind=_span_kind(operation_name), + ) as span: + try: + yield + except Exception as e: + duration_s = time.monotonic() - start_monotonic + if record_exception_stacktrace: + span.record_exception(e) + else: + span.add_event( + "exception", + { + "exception.type": type(e).__name__, + "exception.message": str(e), + }, + ) + span.set_attribute("duration_s", duration_s) + span.set_attribute("status", "error") + span.set_attribute("error_type", type(e).__name__) + span.set_status(Status(StatusCode.ERROR, type(e).__name__)) + self._record_span_metrics( + name=operation_name, + attrs={**span_attrs, "status": "error", "duration_s": duration_s}, + ) + raise + else: + duration_s = time.monotonic() - start_monotonic + span.set_attribute("duration_s", duration_s) + span.set_attribute("status", "ok") + span.set_status(Status(StatusCode.OK)) + self._record_span_metrics( + name=operation_name, + attrs={**span_attrs, "status": "ok", "duration_s": duration_s}, + ) + + def _record_otel_event( + self, + event_type: str, + name: str, + attrs: dict[str, Any], + *, + timestamp_unix_s: float | None = None, + ) -> None: + event_attrs = self._span_attributes(name, {"event.type": event_type, **attrs}) + timestamp_ns = _timestamp_ns(timestamp_unix_s) + current_span = trace.get_current_span() + if current_span is not None and current_span.is_recording(): + current_span.add_event(name, event_attrs, timestamp=timestamp_ns) + self._maybe_close_trajectory_span(name, attrs) + return + + parent_span = self._event_parent_span(attrs) + if parent_span is not None and parent_span.is_recording(): + parent_span.add_event(name, event_attrs, timestamp=timestamp_ns) + self._maybe_close_trajectory_span(name, attrs) + return + + with self._start_as_current_span( + name, + attributes=event_attrs, + context=self._parent_context(attrs), + kind=_span_kind(name), + ) as span: + span.add_event(name, event_attrs, timestamp=timestamp_ns) + if event_type == "error": + span.set_status(Status(StatusCode.ERROR)) + self._maybe_close_trajectory_span(name, attrs) + + def _event_parent_span(self, attrs: dict[str, Any]) -> Span | None: + trajectory_id = _trajectory_id(attrs) + if trajectory_id: + return self._trajectory_span(trajectory_id, attrs, section=_span_section("trajectory", attrs)) + return self._run_span if self._run_span.is_recording() else None + + def _parent_context(self, attrs: dict[str, Any]) -> Any: + parent_span = self._event_parent_span(attrs) + return trace.set_span_in_context(parent_span) if parent_span is not None else None + + def _trajectory_span(self, trajectory_id: str, attrs: dict[str, Any], *, section: str) -> Span: + span_key = (trajectory_id, section) + span = self._trajectory_spans.get(span_key) + if span is not None and span.is_recording(): + _set_span_attributes(span, self._trajectory_root_attributes(trajectory_id, attrs, section=section)) + return span + span = self._start_span( + _section_span_name(section, trajectory_id), + attributes=self._trajectory_root_attributes(trajectory_id, attrs, section=section), + context=self._trajectory_parent_context(trajectory_id, attrs, section=section), + ) + self._trajectory_spans[span_key] = span + return span + + def _trajectory_parent_context(self, trajectory_id: str, attrs: dict[str, Any], *, section: str) -> Any: + if section == "rollout": + return trace.set_span_in_context(self._run_span) + rollout_span = self._trajectory_span(trajectory_id, attrs, section="rollout") + return trace.set_span_in_context(rollout_span) + + def _trajectory_root_attributes( + self, trajectory_id: str, attrs: dict[str, Any], *, section: str + ) -> dict[str, Any]: + root_attrs = { + "event.type": "synthetic_root", + "phase": section, + "span.role": f"{section}.trajectory", + "span.section": section, + "trajectory_id": trajectory_id, + } + for key in ( + "reward", + "stop_reason", + "duration_s", + "loss_multiplier", + "attempt_idx", + "harness", + "dataset_alias", + ): + if attrs.get(key) is not None: + root_attrs[key] = attrs[key] + return safe_attributes(root_attrs) + + def _span_attributes(self, name: str, attrs: dict[str, Any], *, display_name: str | None = None) -> dict[str, Any]: + span_attrs = safe_attributes({**attrs}) + span_attrs.setdefault("operation.name", name) + if display_name is not None and display_name != name: + span_attrs.setdefault("span.display_name", display_name) + span_attrs.setdefault("span.section", _span_section(name, span_attrs)) + if self.run_id: + span_attrs.setdefault("run_id", self.run_id) + for source, target in self.attribute_aliases.items(): + if source in span_attrs and span_attrs[source] is not None: + span_attrs.setdefault(target, span_attrs[source]) + return span_attrs + + def _maybe_close_trajectory_span(self, name: str, attrs: dict[str, Any]) -> None: + if name not in {"trajectory.complete", "trajectory.masked"}: + return + trajectory_id = _trajectory_id(attrs) + if trajectory_id is None: + return + for span_key, span in list(self._trajectory_spans.items()): + if span_key[0] != trajectory_id: + continue + self._trajectory_spans.pop(span_key, None) + if span is None or not span.is_recording(): + continue + _set_span_attributes(span, self._trajectory_root_attributes(trajectory_id, attrs, section=span_key[1])) + span.set_status(Status(StatusCode.ERROR if attrs.get("stop_reason") == "error" else StatusCode.OK)) + span.end() + + def _end_open_spans(self) -> None: + for span_key, span in list(self._trajectory_spans.items()): + if span.is_recording(): + span.set_attribute("stop_reason", "observability_finalize") + span.end() + self._trajectory_spans.pop(span_key, None) + if self._run_span.is_recording(): + self._run_span.set_status(Status(StatusCode.OK)) + self._run_span.end() + + def _operation_span_name(self, name: str, attrs: dict[str, Any]) -> str: + configured_name = attrs.get("span.name") + if configured_name: + return _span_name(configured_name) + if name == "trajectory.tool": + return f"exec: {_command_title(attrs.get('command'), self.command_titles)}" + if name == "sandbox.start": + return _span_with_detail("sandbox.create", attrs.get("image")) + if name == "sandbox.start_batch": + count = attrs.get("count") + detail = f"{count} sandbox{'es' if count != 1 else ''}" if count is not None else None + return _span_with_detail("sandbox.create_batch", detail) + if name == "sandbox.cleanup": + return _span_with_detail("sandbox.cleanup", attrs.get("sandbox_id")) + return _span_name(name) + + def _resource(self) -> Resource: + attributes = dict(self.resource_attributes) + if self._service_name: + attributes["service.name"] = self._service_name + return Resource.create(attributes) + + def _configure_live_exporters(self) -> None: + for exporter_name in _live_trace_exporters(self.otel, self._configured_trace_exporters): + exporter = _build_trace_exporter(exporter_name, self.otel) + if exporter is not None: + self._tracer_provider.add_span_processor(BatchSpanProcessor(exporter)) + + def _configure_metrics(self) -> None: + readers = [] + for exporter_name in _live_metric_exporters(self.otel, self._configured_metric_exporters): + exporter = _build_metric_exporter(exporter_name, self.otel) + if exporter is not None: + readers.append(PeriodicExportingMetricReader(exporter)) + self._meter_provider = MeterProvider(resource=self._resource(), metric_readers=readers) + meter = self._meter_provider.get_meter(SCOPE_NAME, SCOPE_VERSION) + self._duration_histogram = meter.create_histogram( + "nemo_gym.sandbox.operation.duration", + unit="s", + description="Sandbox operation duration.", + ) + self._phase_duration_histograms = { + phase: meter.create_histogram( + f"nemo_gym.sandbox.{phase}.duration", + unit="s", + description=f"Sandbox {phase} duration.", + ) + for phase in ("startup", "setup", "execution", "llm") + } + self._counter = meter.create_counter( + "nemo_gym.sandbox.events", + description="Sandbox event counts.", + ) + + def _start_span(self, name: str, *, attributes: dict[str, Any], context: Any = None) -> Span: + return self._tracer_provider.get_tracer(SCOPE_NAME, SCOPE_VERSION).start_span( + name, + context=context, + attributes=attributes, + kind=_span_kind(name), + ) + + def _start_as_current_span( + self, + name: str, + *, + attributes: dict[str, Any], + context: Any = None, + kind: SpanKind = SpanKind.INTERNAL, + ) -> Any: + return self._tracer_provider.get_tracer(SCOPE_NAME, SCOPE_VERSION).start_as_current_span( + name, + context=context, + attributes=attributes, + kind=kind, + ) + + def _record_event_metrics(self, *, event_type: str, name: str, attrs: dict[str, Any]) -> None: + otel_attrs = self._metric_attrs(attrs) + otel_attrs["event_name"] = name + otel_attrs["event_type"] = event_type + if self._counter is not None: + self._counter.add(1, otel_attrs) + + def _record_span_metrics(self, *, name: str, attrs: dict[str, Any]) -> None: + otel_attrs = self._metric_attrs(attrs) + otel_attrs["span_name"] = name + if self._counter is not None: + self._counter.add(1, otel_attrs) + duration_s = attrs.get("duration_s") + if self._duration_histogram is not None and isinstance(duration_s, (int, float)): + self._duration_histogram.record(float(duration_s), otel_attrs) + phase_histogram = self._phase_duration_histograms.get(str(attrs.get("phase") or "")) + if phase_histogram is not None: + phase_histogram.record(float(duration_s), otel_attrs) + + def _metric_attrs(self, attrs: dict[str, Any]) -> dict[str, str]: + return {key: str(attrs[key]) for key in self.metric_attribute_keys if key in attrs and attrs[key] is not None} + + def _export_trace_artifacts(self) -> dict[str, str]: + if self.output_dir is None or self._local_span_exporter is None: + return {} + self._tracer_provider.force_flush() + return export_trace_artifacts( + self.output_dir, + spans=self._local_span_exporter.finished_spans(), + service_name_strategy=self.local_service_name_strategy, + ) + + def _shutdown_otel(self) -> None: + self._tracer_provider.shutdown() + if self._meter_provider is not None: + self._meter_provider.shutdown() + + +def _trajectory_id(attrs: dict[str, Any]) -> str | None: + value = attrs.get("trajectory_id") or attrs.get("trial_name") + return str(value) if value else None + + +def _set_span_attributes(span: Span, attrs: dict[str, Any]) -> None: + for key, value in attrs.items(): + span.set_attribute(key, value) + + +def _span_kind(name: str) -> SpanKind: + return SpanKind.CLIENT if name == "llm.request" else SpanKind.INTERNAL + + +def _timestamp_ns(timestamp_unix_s: float | None) -> int | None: + return int(timestamp_unix_s * 1_000_000_000) if timestamp_unix_s is not None else None + + +def _string_map(value: Any) -> dict[str, str]: + if not isinstance(value, dict): + return {} + return {str(source): str(target) for source, target in value.items()} + + +def _span_name(value: Any) -> str: + span_name = str(value or "").strip() + return span_name or "sandbox.run" + + +def _eval_span_name(value: Any) -> str: + name = _span_name(value) + return name if name.startswith("eval: ") else f"eval: {name}" + + +def _section_span_name(section: str, trajectory_id: Any) -> str: + return f"{section}: {_span_name(trajectory_id)}" + + +def _span_with_detail(name: str, detail: Any, *, max_length: int = 120) -> str: + text = _compact_text(detail) + if not text: + return name + title = f"{name}: {text}" + return title if len(title) <= max_length else f"{title[: max_length - 1].rstrip()}..." + + +def _command_title(command: Any, config: dict[str, Any]) -> str: + max_length = int(config.get("max_length") or 140) + text = _strip_command_prefixes(_compact_text(command), config.get("strip_prefixes") or ()) + if not text: + return "" + configured_title = _configured_command_title(text, config.get("rules") or ()) + if configured_title: + return _truncate_span_title(configured_title, max_length=max_length) + first_line = next((line.strip() for line in text.splitlines() if line.strip()), "") + title = " ".join(first_line.split()) + return _truncate_span_title(title, max_length=max_length) + + +def _strip_command_prefixes(command: str, prefixes: Any) -> str: + text = command + for prefix in prefixes: + prefix = str(prefix) + if text.startswith(prefix): + return text[len(prefix) :].lstrip() + return text + + +def _configured_command_title(command: str, rules: Any) -> str | None: + for rule in rules: + if not isinstance(rule, dict): + continue + line = _matching_rule_line(command, rule) + if line is not None: + return _format_rule_title(rule, command=command, line=line, match=line) + match = _matching_rule_text(command, rule) + if match is not None: + return _format_rule_title(rule, command=command, line="", match=match) + return None + + +def _matching_rule_line(command: str, rule: dict[str, Any]) -> str | None: + line_prefixes = _string_tuple(rule.get("line_starts_with")) + line_regex = str(rule.get("line_regex") or "") + if not line_prefixes and not line_regex: + return None + lines = [line.strip() for line in command.splitlines() if line.strip()] + if str(rule.get("search") or "first").lower() == "last": + lines = list(reversed(lines)) + regex = re.compile(line_regex) if line_regex else None + for line in lines: + if line_prefixes and line.startswith(line_prefixes): + return " ".join(line.split()) + if regex is not None and regex.search(line): + return " ".join(line.split()) + return None + + +def _matching_rule_text(command: str, rule: dict[str, Any]) -> str | None: + contains = _string_tuple(rule.get("contains") or rule.get("all_contains")) + if contains and not all(part in command for part in contains): + return None + starts_with = _string_tuple(rule.get("starts_with")) + if starts_with and not command.startswith(starts_with): + return None + regex = str(rule.get("regex") or "") + if regex: + match = re.search(regex, command) + if match is None: + return None + return match.group(0) + if contains or starts_with: + return command + return None + + +def _format_rule_title(rule: dict[str, Any], *, command: str, line: str, match: str) -> str: + template = str(rule.get("title") or "{line}" if line else rule.get("title") or "{match}") + return template.format(command=command, line=line, match=match) + + +def _command_title_config(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + return {"strip_prefixes": (), "rules": (), "max_length": 140} + return { + "strip_prefixes": _string_tuple(value.get("strip_prefixes")), + "rules": tuple(rule for rule in value.get("rules") or () if isinstance(rule, dict)), + "max_length": value.get("max_length") or 140, + } + + +def _string_tuple(value: Any) -> tuple[str, ...]: + if value is None: + return () + if isinstance(value, str): + return (value,) + if isinstance(value, (list, tuple, set)): + return tuple(str(item) for item in value) + return (str(value),) + + +def _truncate_span_title(title: str, *, max_length: int) -> str: + return title if len(title) <= max_length else f"{title[: max_length - 1].rstrip()}..." + + +def _compact_text(value: Any) -> str: + return str(value or "").strip() + + +def _span_section(name: str, attrs: dict[str, Any]) -> str: + configured_section = str(attrs.get("span.section") or attrs.get("execution.section") or "").strip() + if configured_section: + return configured_section + if name == "trajectory" and attrs.get("event.type") == "trajectory.complete": + return "rollout" + if name in {"trajectory.tool", "llm.request"} or attrs.get("trajectory_id") is not None: + return "rollout" + if name.startswith("sandbox."): + return "sandbox" + return "eval" + + +def _configured_exporters(cfg: dict[str, Any], signal: str) -> tuple[str, ...] | None: + for key in (f"{signal}_exporters", f"{signal}_exporter", "exporters", "exporter"): + if key in cfg and cfg[key] is not None: + return _exporter_names(cfg[key]) + return None + + +def _exporter_names(value: Any) -> tuple[str, ...]: + if isinstance(value, str): + raw_names = value.split(",") + elif isinstance(value, (list, tuple, set)): + raw_names = value + else: + raw_names = [value] + names = tuple(_normalize_exporter_name(name) for name in raw_names if str(name).strip()) + return tuple(name for name in names if name != "none") + + +def _normalize_exporter_name(value: Any) -> str: + name = str(value).strip().lower().replace("-", "_") + return { + "otlp": "otlp_http", + "otlp_proto_http": "otlp_http", + "http": "otlp_http", + "json": "otlp_json_file", + "file": "otlp_json_file", + "local": "otlp_json_file", + "stdout": "console", + }.get(name, name) + + +def _local_trace_export_enabled( + *, + output_dir: Path | None, + export_traces: bool | None, + trace_exporters: tuple[str, ...] | None, +) -> bool: + if export_traces is not None: + return bool(export_traces) + if trace_exporters is not None: + return "otlp_json_file" in trace_exporters + return output_dir is not None + + +def _live_trace_exporters(cfg: dict[str, Any], configured: tuple[str, ...] | None) -> tuple[str, ...]: + if configured is not None: + return tuple(name for name in configured if name != "otlp_json_file") + if _as_bool(cfg.get("enabled")) and _otel_trace_endpoint(cfg): + return ("otlp_http",) + return () + + +def _live_metric_exporters(cfg: dict[str, Any], configured: tuple[str, ...] | None) -> tuple[str, ...]: + if configured is not None: + return configured + if _as_bool(cfg.get("enabled")) and _otel_metric_endpoint(cfg): + return ("otlp_http",) + return () + + +def _build_trace_exporter(name: str, cfg: dict[str, Any]) -> SpanExporter | None: + if name == "otlp_http": + try: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + except ImportError: + return None + return OTLPSpanExporter( + endpoint=_otel_trace_endpoint(cfg), + headers=_otel_headers(cfg, signal="traces"), + timeout=_otel_timeout(cfg, signal="traces"), + ) + if name == "console": + from opentelemetry.sdk.trace.export import ConsoleSpanExporter + + return ConsoleSpanExporter() + raise ValueError(f"Unsupported sandbox trace exporter: {name!r}") + + +def _build_metric_exporter(name: str, cfg: dict[str, Any]) -> Any: + if name == "otlp_http": + try: + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + except ImportError: + return None + return OTLPMetricExporter( + endpoint=_otel_metric_endpoint(cfg), + headers=_otel_headers(cfg, signal="metrics"), + timeout=_otel_timeout(cfg, signal="metrics"), + ) + if name == "console": + from opentelemetry.sdk.metrics.export import ConsoleMetricExporter + + return ConsoleMetricExporter() + raise ValueError(f"Unsupported sandbox metric exporter: {name!r}") + + +def _as_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + return str(value or "").strip().lower() in {"1", "true", "yes", "on"} + + +def safe_attributes(attributes: dict[str, Any] | None) -> dict[str, Any]: + """Return OpenTelemetry-friendly attributes.""" + if not attributes: + return {} + safe: dict[str, Any] = {} + for key, value in attributes.items(): + if value is None: + continue + if isinstance(value, (str, int, float, bool)): + safe[key] = value + elif isinstance(value, (list, tuple)): + safe[key] = [ + item if isinstance(item, (str, int, float, bool)) or item is None else str(item) for item in value + ] + else: + safe[key] = str(value) + return safe + + +def _otel_trace_endpoint(cfg: dict[str, Any]) -> str | None: + return _otel_signal_endpoint(cfg.get("traces_endpoint") or cfg.get("endpoint"), signal="traces") + + +def _otel_metric_endpoint(cfg: dict[str, Any]) -> str | None: + return _otel_signal_endpoint(cfg.get("metrics_endpoint") or cfg.get("endpoint"), signal="metrics") + + +def _otel_headers(cfg: dict[str, Any], *, signal: str) -> dict[str, str] | None: + headers = cfg.get(f"{signal}_headers") or cfg.get("headers") + if not headers: + return None + if isinstance(headers, dict): + return {str(key): str(value) for key, value in headers.items()} + parsed: dict[str, str] = {} + for item in str(headers).split(","): + if "=" not in item: + continue + key, value = item.split("=", 1) + parsed[key.strip()] = value.strip() + return parsed or None + + +def _otel_timeout(cfg: dict[str, Any], *, signal: str) -> float | None: + timeout = cfg.get(f"{signal}_timeout_s") or cfg.get("timeout_s") + if timeout is None: + return None + return float(timeout) + + +def _otel_signal_endpoint(endpoint: Any, *, signal: str) -> str | None: + if not endpoint: + return None + endpoint_str = str(endpoint).rstrip("/") + if endpoint_str.endswith(f"/v1/{signal}"): + return endpoint_str + if endpoint_str.endswith("/v1/traces") or endpoint_str.endswith("/v1/metrics"): + return endpoint_str.rsplit("/v1/", 1)[0] + f"/v1/{signal}" + return f"{endpoint_str}/v1/{signal}" + + +def _otel_config_from_env() -> dict[str, Any]: + traces_endpoint = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_TRACES_ENDPOINT") or os.environ.get( + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT" + ) + metrics_endpoint = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_METRICS_ENDPOINT") or os.environ.get( + "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT" + ) + endpoint = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_ENDPOINT") or os.environ.get( + "OTEL_EXPORTER_OTLP_ENDPOINT" + ) + traces_exporter = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_TRACES_EXPORTER") or os.environ.get( + "OTEL_TRACES_EXPORTER" + ) + metrics_exporter = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_METRICS_EXPORTER") or os.environ.get( + "OTEL_METRICS_EXPORTER" + ) + headers = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_HEADERS") or os.environ.get( + "OTEL_EXPORTER_OTLP_HEADERS" + ) + traces_headers = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_TRACES_HEADERS") or os.environ.get( + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ) + metrics_headers = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_METRICS_HEADERS") or os.environ.get( + "OTEL_EXPORTER_OTLP_METRICS_HEADERS" + ) + timeout_s = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_TIMEOUT_S") or os.environ.get( + "OTEL_EXPORTER_OTLP_TIMEOUT" + ) + traces_timeout_s = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_TRACES_TIMEOUT_S") or os.environ.get( + "OTEL_EXPORTER_OTLP_TRACES_TIMEOUT" + ) + metrics_timeout_s = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_METRICS_TIMEOUT_S") or os.environ.get( + "OTEL_EXPORTER_OTLP_METRICS_TIMEOUT" + ) + command_titles = _json_object_from_env("NEMO_GYM_SANDBOX_OBSERVABILITY_COMMAND_TITLES") + return { + "enabled": bool(endpoint or traces_endpoint or metrics_endpoint or traces_exporter or metrics_exporter), + "service_name": os.environ.get( + "NEMO_GYM_SANDBOX_OBSERVABILITY_OTEL_SERVICE_NAME", + os.environ.get("OTEL_SERVICE_NAME", ""), + ), + "endpoint": endpoint, + "traces_endpoint": traces_endpoint, + "metrics_endpoint": metrics_endpoint, + "traces_exporter": traces_exporter, + "metrics_exporter": metrics_exporter, + "headers": headers, + "traces_headers": traces_headers, + "metrics_headers": metrics_headers, + "timeout_s": timeout_s, + "traces_timeout_s": traces_timeout_s, + "metrics_timeout_s": metrics_timeout_s, + "command_titles": command_titles, + } + + +def _json_object_from_env(name: str) -> dict[str, Any] | None: + value = os.environ.get(name) + if not value: + return None + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise ValueError(f"{name} must contain a JSON object") + return parsed + + +def build_recorder_from_config( + config: dict[str, Any] | None, + *, + run_id: str | None = None, +) -> SandboxRecorder | None: + """Build a recorder from a sandbox observability config.""" + if not isinstance(config, dict) or not config.get("enabled", False): + return None + output_dir = config.get("output_dir") + return SandboxRecorder( + output_dir=Path(output_dir) if output_dir else None, + otel=dict(config.get("otel") or {}), + run_id=run_id, + run_span_name=config.get("run_span_name") or config.get("job_name"), + export_traces=config.get("export_traces"), + ) + + +def build_recorder_from_env() -> SandboxRecorder | None: + """Build a recorder from eval-job environment variables.""" + output_dir = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_DIR") + otel = _otel_config_from_env() + export_traces_env = os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_EXPORT_TRACES") + export_traces = _as_bool(export_traces_env) if export_traces_env is not None else None + if not output_dir and not otel.get("enabled"): + return None + return SandboxRecorder( + output_dir=Path(output_dir) if output_dir else None, + otel=otel, + run_id=os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_RUN_ID"), + run_span_name=( + os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_RUN_SPAN_NAME") + or os.environ.get("NEMO_GYM_SANDBOX_OBSERVABILITY_JOB_NAME") + or os.environ.get("JOB_NAME") + or os.environ.get("KUBE_JOB_NAME") + ), + export_traces=export_traces, + ) + + +def ensure_env_recorder() -> SandboxRecorder | None: + """Return the process-wide env-configured recorder, creating it once.""" + global G_ENV_RECORDER + if G_ENV_RECORDER is not None: + return G_ENV_RECORDER + with G_ENV_RECORDER_LOCK: + if G_ENV_RECORDER is None: + G_ENV_RECORDER = build_recorder_from_env() + if G_ENV_RECORDER is not None: + atexit.register(G_ENV_RECORDER.finalize) + return G_ENV_RECORDER + + +def current_recorder() -> SandboxRecorder | None: + """Return the active context recorder, if any.""" + return G_CURRENT_RECORDER.get() + + +def _active_recorder() -> SandboxRecorder | None: + return current_recorder() or ensure_env_recorder() + + +def set_current_recorder(recorder: SandboxRecorder) -> Token[SandboxRecorder | None]: + """Set the active recorder for the current context.""" + return G_CURRENT_RECORDER.set(recorder) + + +def reset_current_recorder(token: Token[SandboxRecorder | None]) -> None: + """Reset the active recorder token.""" + G_CURRENT_RECORDER.reset(token) + + +@contextmanager +def use_recorder(recorder: SandboxRecorder | None) -> Iterator[None]: + """Temporarily set the current recorder.""" + if recorder is None: + yield + return + token = set_current_recorder(recorder) + try: + yield + finally: + reset_current_recorder(token) + + +def push_event_context(attributes: dict[str, Any]) -> Token[dict[str, Any]]: + """Merge event context attributes for the current task.""" + return G_EVENT_CONTEXT.set({**G_EVENT_CONTEXT.get(), **attributes}) + + +def reset_event_context(token: Token[dict[str, Any]]) -> None: + """Reset event context attributes.""" + G_EVENT_CONTEXT.reset(token) + + +@contextmanager +def event_context(**attributes: Any) -> Iterator[None]: + """Temporarily add event context attributes.""" + token = push_event_context(attributes) + try: + yield + finally: + reset_event_context(token) + + +def record_event( + event_type: str, + name: str, + *, + attributes: dict[str, Any] | None = None, +) -> None: + """Record one event on the current recorder.""" + recorder = _active_recorder() + if recorder is not None: + recorder.record_event(event_type, name, attributes=attributes) + + +@asynccontextmanager +async def observability_span( + name: str, + *, + phase: str | None = None, + attributes: dict[str, Any] | None = None, +) -> Iterator[None]: + """Record an async span on the current recorder.""" + recorder = _active_recorder() + if recorder is None: + yield + return + async with recorder.span(name, phase=phase, attributes=attributes): + yield + + +@contextmanager +def observability_sync_span( + name: str, + *, + phase: str | None = None, + attributes: dict[str, Any] | None = None, +) -> Iterator[None]: + """Record a sync span on the current recorder.""" + recorder = _active_recorder() + if recorder is None: + yield + return + with recorder.sync_span(name, phase=phase, attributes=attributes): + yield diff --git a/nemo_gym/sandbox/observability/traces.py b/nemo_gym/sandbox/observability/traces.py new file mode 100644 index 0000000000..3d6beaf6d2 --- /dev/null +++ b/nemo_gym/sandbox/observability/traces.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenTelemetry SDK trace artifact exporters for sandbox observability.""" + +from __future__ import annotations + +import json +import threading +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.trace import SpanKind, StatusCode + + +SCOPE_NAME = "nemo_gym.sandbox.observability" +SCOPE_VERSION = "1" + + +class JsonSpanExporter(SpanExporter): + """OpenTelemetry span exporter that keeps finished spans for local artifacts.""" + + def __init__(self) -> None: + self._spans: list[ReadableSpan] = [] + self._lock = threading.Lock() + self._shutdown = False + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + with self._lock: + if self._shutdown: + return SpanExportResult.FAILURE + self._spans.extend(spans) + return SpanExportResult.SUCCESS + + def force_flush(self, timeout_millis: int = 30000) -> bool: + del timeout_millis + return True + + def shutdown(self) -> None: + with self._lock: + self._shutdown = True + + def finished_spans(self) -> list[ReadableSpan]: + with self._lock: + return list(self._spans) + + +def export_trace_artifacts( + output_dir: Path, + *, + spans: Sequence[ReadableSpan], + service_name_strategy: str | None = "span_section", +) -> dict[str, str]: + """Export SDK-finished spans as an OTLP-shaped JSON trace artifact.""" + if not spans: + return {} + + traces_dir = output_dir / "traces" + traces_dir.mkdir(parents=True, exist_ok=True) + otlp_path = traces_dir / "otel_traces.json" + otlp_path.write_text( + json.dumps(_otlp_payload(spans, service_name_strategy=service_name_strategy), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return {"otlp_json": str(otlp_path)} + + +def _otlp_payload(spans: Sequence[ReadableSpan], *, service_name_strategy: str | None) -> dict[str, Any]: + resource_groups: dict[tuple[tuple[str, str], ...], dict[str, Any]] = {} + + for span in spans: + resource_attributes = _resource_attributes(span, service_name_strategy=service_name_strategy) + resource_key = tuple(sorted((key, str(value)) for key, value in resource_attributes.items())) + resource_group = resource_groups.setdefault( + resource_key, + { + "resource": { + "attributes": [_attribute(key, value) for key, value in sorted(resource_attributes.items())] + }, + "scopeSpans": {}, + }, + ) + scope = span.instrumentation_scope + scope_key = ( + getattr(scope, "name", None) or SCOPE_NAME, + getattr(scope, "version", None) or SCOPE_VERSION, + ) + scope_group = resource_group["scopeSpans"].setdefault( + scope_key, + { + "scope": { + "name": scope_key[0], + "version": scope_key[1], + }, + "spans": [], + }, + ) + scope_group["spans"].append(_otlp_span(span)) + + resource_spans = [] + for group in resource_groups.values(): + scope_spans = list(group["scopeSpans"].values()) + resource_spans.append({"resource": group["resource"], "scopeSpans": scope_spans}) + return {"resourceSpans": resource_spans} + + +def _resource_attributes(span: ReadableSpan, *, service_name_strategy: str | None) -> dict[str, Any]: + attrs = dict(span.resource.attributes) + if _uses_visual_service_lanes(service_name_strategy): + service_name = _visual_service_name(span) + if service_name: + original_service_name = attrs.get("service.name") + if original_service_name: + attrs.setdefault("service.original_name", original_service_name) + attrs["service.name"] = service_name + return attrs + + +def _otlp_span(span: ReadableSpan) -> dict[str, Any]: + row: dict[str, Any] = { + "traceId": _trace_id(span.context.trace_id), + "spanId": _span_id(span.context.span_id), + "name": span.name, + "kind": _span_kind(span.kind), + "startTimeUnixNano": str(span.start_time or 0), + "endTimeUnixNano": str(span.end_time or span.start_time or 0), + "attributes": [ + _attribute(key, value) for key, value in sorted((span.attributes or {}).items()) if value is not None + ], + "events": [_otlp_event(event) for event in span.events], + "status": _otlp_status(span), + } + if span.parent is not None and span.parent.span_id: + row["parentSpanId"] = _span_id(span.parent.span_id) + return row + + +def _otlp_event(event: Any) -> dict[str, Any]: + return { + "name": event.name, + "timeUnixNano": str(event.timestamp), + "attributes": [ + _attribute(key, value) for key, value in sorted((event.attributes or {}).items()) if value is not None + ], + } + + +def _otlp_status(span: ReadableSpan) -> dict[str, str]: + status_code = span.status.status_code + if status_code == StatusCode.ERROR: + code = "STATUS_CODE_ERROR" + elif status_code == StatusCode.OK: + code = "STATUS_CODE_OK" + else: + code = "STATUS_CODE_UNSET" + status = {"code": code} + if span.status.description: + status["message"] = span.status.description + return status + + +def _span_kind(kind: SpanKind) -> str: + return f"SPAN_KIND_{kind.name}" + + +def _trace_id(value: int) -> str: + return f"{value:032x}" + + +def _uses_visual_service_lanes(strategy: str | None) -> bool: + return str(strategy or "").strip().lower() in {"span_section", "visual", "visual_lanes", "service_lanes"} + + +def _visual_service_name(span: ReadableSpan) -> str | None: + attrs = dict(span.attributes or {}) + operation_name = str(attrs.get("operation.name") or span.name or "") + span_role = str(attrs.get("span.role") or "") + + if operation_name == "sandbox.run" or span_role == "eval.run": + return "nemo-gym.eval" + if operation_name == "trajectory" or span_role.endswith(".trajectory"): + section = str(attrs.get("span.section") or "rollout") + return f"nemo-gym.{section}" + if operation_name == "llm.request": + return "llm.request" + if operation_name in { + "sandbox.start", + "sandbox.start_batch", + "sandbox.create", + "sandbox.create_api", + "sandbox.create_probe", + }: + return "sandbox.create" + if attrs.get("span.section") == "verifier" and operation_name in {"trajectory.tool", "sandbox.exec"}: + return "verifier.exec" + if operation_name in {"trajectory.tool", "sandbox.exec"}: + return "sandbox.exec" + if operation_name.startswith("sandbox.diagnostic."): + return "sandbox.diagnostic" + if operation_name in {"sandbox.cleanup", "sandbox.close", "sandbox.delete"}: + return "sandbox.cleanup" + if operation_name in {"sandbox.read_file", "sandbox.write_file", "sandbox.upload_file", "sandbox.download_file"}: + return "sandbox.io" + + section = attrs.get("span.section") + if section in {"eval", "rollout", "verifier"}: + return f"nemo-gym.{section}" + if section == "sandbox": + return "sandbox" + return None + + +def _span_id(value: int) -> str: + return f"{value:016x}" + + +def _attribute(key: str, value: Any) -> dict[str, Any]: + return { + "key": key, + "value": _otel_value(value), + } + + +def _otel_value(value: Any) -> dict[str, Any]: + if isinstance(value, bool): + return {"boolValue": value} + if isinstance(value, int) and not isinstance(value, bool): + return {"intValue": str(value)} + if isinstance(value, float): + return {"doubleValue": value} + if isinstance(value, (list, tuple)): + return {"arrayValue": {"values": [_otel_value(item) for item in value]}} + return {"stringValue": str(value)} diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py new file mode 100644 index 0000000000..127fc1c627 --- /dev/null +++ b/nemo_gym/sandbox/providers/__init__.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sandbox provider registry.""" + +from nemo_gym.sandbox.providers.base import ( + SandboxBatchCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxProvider, + SandboxSpec, +) +from nemo_gym.sandbox.providers.registry import ( + create_provider, + get_provider_class, + list_providers, + register_provider, +) + + +__all__ = [ + "SandboxBatchCreateError", + "SandboxCreateVerificationError", + "SandboxExecResult", + "SandboxHandle", + "SandboxProvider", + "SandboxSpec", + "create_provider", + "get_provider_class", + "list_providers", + "register_provider", +] diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py new file mode 100644 index 0000000000..9acc922416 --- /dev/null +++ b/nemo_gym/sandbox/providers/base.py @@ -0,0 +1,132 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider-facing sandbox protocol. + +Providers are the only layer that talks to runtime and infrastructure APIs. +Gym agents and external harnesses consume the public ``nemo_gym.sandbox`` API +instead of importing provider-specific modules. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol + + +@dataclass(frozen=True) +class SandboxSpec: + """Provider-neutral sandbox creation request.""" + + image: str | None = None + snapshot_id: str | None = None + timeout_s: int | None = None + ready_timeout_s: int | None = None + env: dict[str, str] = field(default_factory=dict) + metadata: dict[str, str] = field(default_factory=dict) + resources: dict[str, str] = field(default_factory=dict) + entrypoint: list[str] | None = None + extensions: dict[str, str] = field(default_factory=dict) + platform: dict[str, Any] | None = None + volumes: list[dict[str, Any]] | None = None + skip_health_check: bool | None = None + + +@dataclass(frozen=True) +class SandboxHandle: + """Provider-neutral handle to a created sandbox.""" + + sandbox_id: str + provider_name: str + raw: Any + + +@dataclass(frozen=True) +class SandboxExecResult: + """Provider-neutral process execution result.""" + + stdout: str | None + stderr: str | None + return_code: int + + +class SandboxBatchCreateError(RuntimeError): + """Raised when a provider cannot complete sandbox batch creation.""" + + +class SandboxCreateVerificationError(ConnectionError): + """Raised when a newly-created sandbox fails provider readiness checks.""" + + +class SandboxProvider(Protocol): + """Runtime/infra provider contract used by the public sandbox API.""" + + name: str + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + """Create a sandbox and return a provider-neutral handle.""" + ... + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several equivalent sandboxes. + + Providers that have a native bulk-allocation primitive should use it. + Providers without one may fall back to calling ``create`` repeatedly. + When ``allow_partial`` is true, providers may return a smaller + contiguous prefix of successfully created handles instead of failing the + whole batch. + """ + ... + + async def connect(self, sandbox_id: str) -> SandboxHandle: + """Connect to an existing sandbox.""" + ... + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + """Run a command inside a sandbox.""" + ... + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + """Write a file into a sandbox.""" + ... + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + """Read a file from a sandbox.""" + ... + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + """Upload one local file into a sandbox.""" + ... + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + """Download one sandbox file to the local filesystem.""" + ... + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + """Close provider resources and optionally delete the sandbox.""" + ... diff --git a/nemo_gym/sandbox/providers/opensandbox/__init__.py b/nemo_gym/sandbox/providers/opensandbox/__init__.py new file mode 100644 index 0000000000..a27657684e --- /dev/null +++ b/nemo_gym/sandbox/providers/opensandbox/__init__.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenSandbox provider package.""" + +from nemo_gym.sandbox.providers.opensandbox.provider import ( + OpenSandboxBatchCreateError, + OpenSandboxConnectionConfig, + OpenSandboxCreateConfig, + OpenSandboxCreateTimeoutError, + OpenSandboxCreateVerificationError, + OpenSandboxOperationConfig, + OpenSandboxPoolConfig, + OpenSandboxProbeConfig, + OpenSandboxProvider, +) + + +__all__ = [ + "OpenSandboxBatchCreateError", + "OpenSandboxConnectionConfig", + "OpenSandboxCreateConfig", + "OpenSandboxCreateTimeoutError", + "OpenSandboxCreateVerificationError", + "OpenSandboxOperationConfig", + "OpenSandboxPoolConfig", + "OpenSandboxProbeConfig", + "OpenSandboxProvider", +] diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py new file mode 100644 index 0000000000..5c92b1f8c7 --- /dev/null +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -0,0 +1,1453 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenSandbox provider implementation.""" + +import asyncio +import logging +import re +import shlex +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import timedelta +from pathlib import Path +from typing import Any, Awaitable, Callable +from uuid import uuid4 + +from nemo_gym.sandbox.observability import observability_span, record_event +from nemo_gym.sandbox.providers.base import ( + SandboxBatchCreateError, + SandboxCreateVerificationError, + SandboxExecResult, + SandboxHandle, + SandboxSpec, +) + + +LOGGER = logging.getLogger(__name__) + + +class OpenSandboxBatchCreateError(SandboxBatchCreateError): + """Raised when a batch sandbox preallocation cannot be completed.""" + + +class OpenSandboxCreateTimeoutError(TimeoutError): + """Raised when OpenSandbox sandbox creation exceeds the client timeout.""" + + +class OpenSandboxCreateVerificationError(SandboxCreateVerificationError): + """Raised when a newly-created sandbox cannot execute a probe command.""" + + +RETRYABLE_HTTP_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504} +RETRYABLE_ERROR_MARKERS = ( + "all connection attempts failed", + "connection refused", + "connection reset", + "gateway timeout", + "http 408", + "http 409", + "http 425", + "http 429", + "http 500", + "http 502", + "http 503", + "http 504", + "incomplete chunked read", + "peer closed connection", + "pod ip is not yet available", + "pod may still be starting", + "errimagepull", + "get endpoint for sandbox", + "imagepullbackoff", + "pod failed", + "podfailed", + "remote protocol error", + "service unavailable", + "server disconnected", + "status code: 408", + "status code: 409", + "status code: 425", + "status code: 429", + "status code: 500", + "status code: 502", + "status code: 503", + "status code: 504", + "temporarily unavailable", + "timed out", + "timeout", +) +METADATA_VALUE_RE = re.compile(r"[^A-Za-z0-9_.-]+") +DEFAULT_IMAGE_PULL_POLICY = "IfNotPresent" +IMAGE_PULL_POLICY_EXTENSION_KEY = "imagePullPolicy" +IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY = "opensandbox.extensions.image-pull-policy" +VALID_IMAGE_PULL_POLICIES = {"Always", "IfNotPresent", "Never"} +STATUS_CODE_RE = re.compile(r"(?:status code|http)\D+(\d{3})", re.IGNORECASE) + + +def validate_image_pull_policy(image_pull_policy: str) -> str: + """Validate a Kubernetes-compatible container image pull policy.""" + if image_pull_policy not in VALID_IMAGE_PULL_POLICIES: + allowed = ", ".join(sorted(VALID_IMAGE_PULL_POLICIES)) + raise ValueError(f"image_pull_policy must be one of: {allowed}") + return image_pull_policy + + +def _require_opensandbox_sdk() -> tuple[Any, Any, Any, Any, Any]: + try: + from opensandbox import Sandbox + from opensandbox.config import ConnectionConfig + from opensandbox.models.execd import RunCommandOpts + from opensandbox.models.sandboxes import PlatformSpec, Volume + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "OpenSandbox SDK is required for the opensandbox sandbox provider. " + "Install it in the NeMo-RL runtime image before using " + "env.sandbox.provider.name=opensandbox." + ) from e + + return Sandbox, ConnectionConfig, RunCommandOpts, PlatformSpec, Volume + + +def _require_opensandbox_sdk_pool() -> tuple[Any, Any, Any, Any]: + try: + from opensandbox import ( + AcquirePolicy, + InMemoryAsyncPoolStateStore, + PoolCreationSpec, + SandboxPoolAsync, + ) + except ImportError as e: + raise ModuleNotFoundError( + "OpenSandbox SDK >=0.1.9 is required for native SDK pool batch creation. " + "Install opensandbox>=0.1.9 in the NeMo-RL runtime image." + ) from e + + return AcquirePolicy, InMemoryAsyncPoolStateStore, PoolCreationSpec, SandboxPoolAsync + + +def _require_tenacity() -> tuple[Any, Any, Any, Any]: + try: + from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "tenacity is required for OpenSandbox retry handling. Install nemo-gym[sandbox] before using " + "env.sandbox.provider.name=opensandbox." + ) from e + + return AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential + + +def _httpx_retryable_types() -> tuple[type[BaseException], ...]: + try: + import httpx + except ModuleNotFoundError: + return tuple() + return ( + httpx.RemoteProtocolError, + httpx.ReadError, + httpx.WriteError, + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + httpx.TimeoutException, + httpx.NetworkError, + ) + + +def _has_retryable_error_marker(exception: BaseException) -> bool: + message = str(exception).lower() + return any(marker in message for marker in RETRYABLE_ERROR_MARKERS) + + +def _exception_status_code(exception: BaseException) -> int | None: + status_code = getattr(exception, "status_code", None) + if isinstance(status_code, int): + return status_code + + match = STATUS_CODE_RE.search(str(exception)) + if match is None: + return None + return int(match.group(1)) + + +def _sdk_error_attributes( + exception: BaseException, + *, + operation: str, + sandbox_id: str, + attempt_number: int | None = None, + max_attempts: int | None = None, + sleep_s: float | None = None, +) -> dict[str, Any]: + attrs: dict[str, Any] = { + "provider": OpenSandboxProvider.name, + "operation": operation, + "sandbox_id": sandbox_id, + "error_type": type(exception).__name__, + "error_message": str(exception)[:500], + } + status_code = _exception_status_code(exception) + if status_code is not None: + attrs["status_code"] = status_code + if attempt_number is not None: + attrs["attempt_number"] = attempt_number + if max_attempts is not None: + attrs["max_attempts"] = max_attempts + if sleep_s is not None: + attrs["next_sleep_s"] = sleep_s + return attrs + + +def _is_retryable_create_error(exception: BaseException) -> bool: + """Return whether a sandbox create failure is likely transient.""" + if isinstance(exception, SandboxCreateVerificationError): + return True + if isinstance(exception, (ConnectionError, OSError, TimeoutError)): + return True + httpx_types = _httpx_retryable_types() + if httpx_types and isinstance(exception, httpx_types): + return True + + try: + from opensandbox.exceptions import ( + InvalidArgumentException, + SandboxApiException, + SandboxException, + SandboxInternalException, + SandboxReadyTimeoutException, + SandboxUnhealthyException, + ) + except ModuleNotFoundError: + return _has_retryable_error_marker(exception) + + if isinstance(exception, InvalidArgumentException): + return False + if isinstance( + exception, + ( + SandboxInternalException, + SandboxReadyTimeoutException, + SandboxUnhealthyException, + ), + ): + return True + if isinstance(exception, SandboxApiException): + status_code = getattr(exception, "status_code", None) + if status_code in RETRYABLE_HTTP_STATUS_CODES: + return True + if status_code is not None and status_code < 500: + return False + if not isinstance(exception, SandboxException): + return _has_retryable_error_marker(exception) + + return _has_retryable_error_marker(exception) + + +def _is_retryable_sdk_operation_error(exception: BaseException) -> bool: + """Return whether an SDK operation can be retried by Gym. + + The OpenSandbox Python SDK does not retry generated lifecycle, execd, or + filesystem HTTP calls. It converts network failures into SDK exceptions and + exposes API status codes, so classify both the wrapper and its original + cause here. + """ + if isinstance(exception, TimeoutError): + return False + cause = exception.__cause__ + if isinstance(cause, BaseException) and _is_retryable_sdk_operation_error(cause): + return True + if isinstance(exception, (ConnectionError, OSError)): + return True + httpx_types = _httpx_retryable_types() + if httpx_types and isinstance(exception, httpx_types): + return True + return _is_retryable_create_error(exception) + + +def _is_missing_sandbox_delete_error(exception: BaseException) -> bool: + message = str(exception).lower() + return "sandbox" in message and "not found" in message + + +def _log_create_retry(retry_state: Any) -> None: + exception = retry_state.outcome.exception() if retry_state.outcome else None + sleep_s = retry_state.next_action.sleep if retry_state.next_action else None + LOGGER.warning( + "Retrying OpenSandbox sandbox create after attempt %s; next_sleep_s=%s; error=%r", + retry_state.attempt_number, + sleep_s, + exception, + ) + + +def _log_operation_retry(retry_state: Any) -> None: + exception = retry_state.outcome.exception() if retry_state.outcome else None + sleep_s = retry_state.next_action.sleep if retry_state.next_action else None + LOGGER.warning( + "Retrying OpenSandbox SDK operation after attempt %s; next_sleep_s=%s; error=%r", + retry_state.attempt_number, + sleep_s, + exception, + ) + + +def _string_map(values: dict[str, Any]) -> dict[str, str]: + return {str(key): str(value) for key, value in values.items()} + + +def _metadata_value(value: Any) -> str: + normalized = METADATA_VALUE_RE.sub("_", str(value)).strip("._-") + normalized = normalized[:63].strip("._-") + return normalized or "metadata" + + +def _metadata_map(values: dict[str, Any]) -> dict[str, str]: + return {str(key): _metadata_value(value) for key, value in values.items()} + + +def _normalize_spec(spec: SandboxSpec) -> SandboxSpec: + return replace( + spec, + env=_string_map(spec.env), + metadata=_metadata_map(spec.metadata), + resources=_string_map(spec.resources), + extensions=_string_map(spec.extensions), + ) + + +def _to_platform_spec(platform: dict[str, Any]) -> Any: + _, _, _, PlatformSpec, _ = _require_opensandbox_sdk() + return PlatformSpec(**platform) + + +def _to_volumes(volumes: list[dict[str, Any]]) -> list[Any]: + _, _, _, _, Volume = _require_opensandbox_sdk() + return [Volume(**volume) for volume in volumes] + + +def _seconds_to_timedelta(seconds: int | float | None) -> timedelta | None: + if seconds is None: + return None + return timedelta(seconds=float(seconds)) + + +@dataclass(frozen=True) +class OpenSandboxConnectionConfig: + """OpenSandbox server connection settings.""" + + domain: str | None = None + api_key: str | None = None + protocol: str | None = None + use_server_proxy: bool | None = None + exec_use_server_proxy: bool | None = None + request_timeout_s: int | None = None + connect_timeout_s: int | float | None = None + + def __post_init__(self) -> None: + if self.connect_timeout_s is not None and self.connect_timeout_s <= 0: + raise ValueError("connection.connect_timeout_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxCreateConfig: + """OpenSandbox create/reconnect retry settings.""" + + request_timeout_s: int | None = None + timeout_s: float | None = None + retries: int = 2 + retry_delay_s: float = 5.0 + retry_max_delay_s: float = 60.0 + image_pull_policy: str | None = DEFAULT_IMAGE_PULL_POLICY + skip_health_check: bool = False + connect_attempt_timeout_s: float = 30.0 + connect_poll_s: float = 2.0 + + def __post_init__(self) -> None: + if self.image_pull_policy is not None: + validate_image_pull_policy(self.image_pull_policy) + if self.timeout_s is not None and self.timeout_s <= 0: + raise ValueError("create.timeout_s must be > 0") + if self.retries < 0: + raise ValueError("create.retries must be >= 0") + if self.retry_delay_s < 0: + raise ValueError("create.retry_delay_s must be >= 0") + if self.retry_max_delay_s < 0: + raise ValueError("create.retry_max_delay_s must be >= 0") + if self.connect_attempt_timeout_s <= 0: + raise ValueError("create.connect_attempt_timeout_s must be > 0") + if self.connect_poll_s <= 0: + raise ValueError("create.connect_poll_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxProbeConfig: + """Post-create probe settings.""" + + command: str | None = "printf nemo-rl-sandbox-ready" + expected_stdout: str | None = "nemo-rl-sandbox-ready" + timeout_s: int = 30 + deadline_s: float | None = None + sample_count: int | None = None + stable_count: int = 1 + stable_delay_s: float = 0.0 + + def __post_init__(self) -> None: + if self.command is not None and self.timeout_s <= 0: + raise ValueError("probe.timeout_s must be > 0") + if self.deadline_s is not None and self.deadline_s <= 0: + raise ValueError("probe.deadline_s must be > 0") + if self.sample_count is not None and self.sample_count < 1: + raise ValueError("probe.sample_count must be >= 1") + if self.stable_count < 1: + raise ValueError("probe.stable_count must be >= 1") + if self.stable_delay_s < 0: + raise ValueError("probe.stable_delay_s must be >= 0") + + +@dataclass(frozen=True) +class OpenSandboxOperationConfig: + """Retry and timeout settings for SDK operations after create.""" + + retries: int = 3 + retry_delay_s: float = 1.0 + retry_max_delay_s: float = 15.0 + command_retries: int | None = None + close_timeout_s: float | None = 30.0 + + def __post_init__(self) -> None: + if self.retries < 0: + raise ValueError("operations.retries must be >= 0") + if self.retry_delay_s < 0: + raise ValueError("operations.retry_delay_s must be >= 0") + if self.retry_max_delay_s < 0: + raise ValueError("operations.retry_max_delay_s must be >= 0") + if self.command_retries is not None and self.command_retries < 0: + raise ValueError("operations.command_retries must be >= 0") + if self.close_timeout_s is not None and self.close_timeout_s <= 0: + raise ValueError("operations.close_timeout_s must be > 0") + + +@dataclass(frozen=True) +class OpenSandboxPoolConfig: + """OpenSandbox SDK pool and batch fanout settings.""" + + concurrency: int = 4 + progress_timeout_s: float | None = None + reconcile_interval_s: float = 0.1 + acquire_poll_interval_s: float = 0.1 + idle_timeout_s: float | None = None + primary_lock_ttl_s: float | None = None + + def __post_init__(self) -> None: + if self.concurrency < 1: + raise ValueError("pool.concurrency must be >= 1") + if self.progress_timeout_s is not None and self.progress_timeout_s <= 0: + raise ValueError("pool.progress_timeout_s must be > 0") + if self.reconcile_interval_s <= 0: + raise ValueError("pool.reconcile_interval_s must be > 0") + if self.acquire_poll_interval_s <= 0: + raise ValueError("pool.acquire_poll_interval_s must be > 0") + if self.idle_timeout_s is not None and self.idle_timeout_s <= 0: + raise ValueError("pool.idle_timeout_s must be > 0") + if self.primary_lock_ttl_s is not None and self.primary_lock_ttl_s <= 0: + raise ValueError("pool.primary_lock_ttl_s must be > 0") + + +def _coerce_config(value: Any, config_cls: type[Any]) -> Any: + if value is None: + return config_cls() + if isinstance(value, config_cls): + return value + if isinstance(value, Mapping): + return config_cls(**value) + raise TypeError(f"{config_cls.__name__} must be a mapping or {config_cls.__name__} instance") + + +class OpenSandboxProvider: + """Provider backed by the OpenSandbox SDK/server API. + + Batch allocations use the official OpenSandbox SDK client-side pool. + """ + + name = "opensandbox" + + def __init__( + self, + *, + connection: OpenSandboxConnectionConfig | Mapping[str, Any] | None = None, + create: OpenSandboxCreateConfig | Mapping[str, Any] | None = None, + probe: OpenSandboxProbeConfig | Mapping[str, Any] | None = None, + operations: OpenSandboxOperationConfig | Mapping[str, Any] | None = None, + pool: OpenSandboxPoolConfig | Mapping[str, Any] | None = None, + ) -> None: + self._connection = _coerce_config(connection, OpenSandboxConnectionConfig) + self._create = _coerce_config(create, OpenSandboxCreateConfig) + self._probe = _coerce_config(probe, OpenSandboxProbeConfig) + self._operations = _coerce_config(operations, OpenSandboxOperationConfig) + self._pool = _coerce_config(pool, OpenSandboxPoolConfig) + + def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec: + """Ensure SDK create requests carry the desired image pull policy.""" + if self._create.image_pull_policy is None: + return spec + + extensions = dict(spec.extensions) + image_pull_policy = extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY) or extensions.get( + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY + ) + if image_pull_policy is None: + image_pull_policy = self._create.image_pull_policy + image_pull_policy = validate_image_pull_policy(image_pull_policy) + extensions.setdefault(IMAGE_PULL_POLICY_EXTENSION_KEY, image_pull_policy) + extensions.setdefault(IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, image_pull_policy) + return replace(spec, extensions=extensions) + + def _connection_config( + self, + request_timeout_s: int | float | None = None, + *, + use_server_proxy: bool | None = None, + ) -> Any: + _, ConnectionConfig, _, _, _ = _require_opensandbox_sdk() + kwargs: dict[str, Any] = {} + if self._connection.domain is not None: + kwargs["domain"] = self._connection.domain + if self._connection.api_key is not None: + kwargs["api_key"] = self._connection.api_key + if self._connection.protocol is not None: + kwargs["protocol"] = self._connection.protocol + if use_server_proxy is None: + use_server_proxy = self._connection.use_server_proxy + if use_server_proxy is not None: + kwargs["use_server_proxy"] = use_server_proxy + if request_timeout_s is None: + request_timeout_s = self._connection.request_timeout_s + if request_timeout_s is not None: + kwargs["request_timeout"] = timedelta(seconds=request_timeout_s) + return ConnectionConfig(**kwargs) + + def _exec_connection_config(self, request_timeout_s: int | float | None = None) -> Any: + """Connection config for SDK handles that issue execd/filesystem calls. + + For clustered evaluations, exec traffic should normally use the + OpenSandbox server proxy. That keeps clients off pod IP routing and lets + the server resolve the sandbox backend for each request. + """ + use_server_proxy = self._connection.use_server_proxy + if self._connection.exec_use_server_proxy is not None: + use_server_proxy = self._connection.exec_use_server_proxy + return self._connection_config( + request_timeout_s=request_timeout_s, + use_server_proxy=use_server_proxy, + ) + + async def aclose(self) -> None: + """Close provider-owned resources. + + The provider intentionally does not inject or own OpenSandbox SDK + network clients. SDK handles are closed per sandbox in ``close``. + """ + return None + + async def _await_sdk_call( + self, + awaitable: Any, + *, + operation: str, + sandbox_id: str, + timeout_s: float | None, + ) -> Any: + if timeout_s is None: + return await awaitable + + try: + return await asyncio.wait_for(awaitable, timeout=timeout_s) + except asyncio.TimeoutError as e: + raise TimeoutError( + f"Timed out during OpenSandbox {operation} after {timeout_s:g}s; sandbox_id={sandbox_id!r}" + ) from e + + async def _await_sdk_operation( + self, + operation_factory: Callable[[], Awaitable[Any]], + *, + operation: str, + sandbox_id: str, + timeout_s: float | None, + retries: int | None = None, + ) -> Any: + AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity() + retry_count = self._operations.retries if retries is None else retries + max_attempts = retry_count + 1 + + def _before_sleep(retry_state: Any) -> None: + _log_operation_retry(retry_state) + exception = retry_state.outcome.exception() if retry_state.outcome else None + sleep_s = retry_state.next_action.sleep if retry_state.next_action else None + if exception is not None: + record_event( + "warning", + "sandbox.opensandbox.sdk_operation_retry", + attributes=_sdk_error_attributes( + exception, + operation=operation, + sandbox_id=sandbox_id, + attempt_number=retry_state.attempt_number, + max_attempts=max_attempts, + sleep_s=sleep_s, + ), + ) + + retry_policy = AsyncRetrying( + retry=retry_if_exception(_is_retryable_sdk_operation_error), + stop=stop_after_attempt(max_attempts), + wait=wait_random_exponential( + multiplier=self._operations.retry_delay_s, + max=self._operations.retry_max_delay_s, + ), + before_sleep=_before_sleep, + reraise=True, + ) + try: + async for attempt in retry_policy: + with attempt: + return await self._await_sdk_call( + operation_factory(), + operation=operation, + sandbox_id=sandbox_id, + timeout_s=timeout_s, + ) + except (asyncio.CancelledError, KeyboardInterrupt): + raise + except Exception as e: + record_event( + "error", + "sandbox.opensandbox.sdk_operation_error", + attributes=_sdk_error_attributes( + e, + operation=operation, + sandbox_id=sandbox_id, + attempt_number=max_attempts, + max_attempts=max_attempts, + ), + ) + raise + + raise RuntimeError("OpenSandbox SDK operation retry loop did not run") + + async def _verify_created_handle(self, handle: SandboxHandle) -> None: + if self._probe.command is None: + return + + loop = asyncio.get_running_loop() + deadline_s = self._probe.deadline_s or float(self._probe.timeout_s) + deadline = loop.time() + deadline_s + successful_probes = 0 + attempt_number = 0 + last_exception: BaseException | None = None + + while successful_probes < self._probe.stable_count: + remaining_s = deadline - loop.time() + if remaining_s <= 0: + error = OpenSandboxCreateVerificationError( + "OpenSandbox sandbox failed create probe command before " + "the startup deadline; " + f"sandbox_id={handle.sandbox_id!r}, " + f"command={self._probe.command!r}, " + f"successful_probes={successful_probes}/{self._probe.stable_count}, " + f"attempts={attempt_number}, deadline_s={deadline_s:g}" + ) + raise error from last_exception + + attempt_number += 1 + probe_index = successful_probes + if self._probe.deadline_s is None: + command_timeout_s = float(self._probe.timeout_s) + else: + command_timeout_s = min(float(self._probe.timeout_s), remaining_s) + try: + async with observability_span( + "sandbox.create_probe", + phase="startup", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "probe_index": probe_index, + "probe_count": self._probe.stable_count, + "attempt_number": attempt_number, + "deadline_s": deadline_s, + }, + ): + result = await asyncio.wait_for( + self._exec( + handle, + self._probe.command, + timeout_s=command_timeout_s, + user="root", + ), + timeout=command_timeout_s, + ) + except asyncio.CancelledError: + raise + except Exception as e: + last_exception = e + successful_probes = 0 + record_event( + "warning", + "sandbox.opensandbox.create_probe_retry", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "operation": "create_probe", + "attempt_number": attempt_number, + "successful_probes": successful_probes, + "required_probes": self._probe.stable_count, + "deadline_s": deadline_s, + "remaining_s": max(deadline - loop.time(), 0.0), + "error_type": type(e).__name__, + "error_message": str(e)[:500], + }, + ) + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + continue + + stdout = result.stdout or "" + expected = self._probe.expected_stdout + if result.return_code != 0 or (expected is not None and expected not in stdout): + last_exception = OpenSandboxCreateVerificationError( + "OpenSandbox sandbox create probe command returned an " + f"unexpected result; sandbox_id={handle.sandbox_id!r}, " + f"return_code={result.return_code}, expected_stdout={expected!r}, " + f"stdout={stdout[:200]!r}, stderr={(result.stderr or '')[:200]!r}, " + f"probe={successful_probes + 1}/{self._probe.stable_count}" + ) + successful_probes = 0 + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + continue + + successful_probes += 1 + if successful_probes < self._probe.stable_count and self._probe.stable_delay_s: + await asyncio.sleep(self._probe.stable_delay_s) + + async def _verify_created_handles( + self, + handles: list[SandboxHandle], + ) -> None: + """Verify a batch of created handles with bounded probe concurrency.""" + if self._probe.command is None or not handles: + return + + handles_to_probe = handles + if self._probe.sample_count is not None and self._probe.sample_count < len(handles): + sample_count = self._probe.sample_count + if sample_count == 1: + sampled_indices = [0] + else: + sampled_indices = [ + round(index * (len(handles) - 1) / (sample_count - 1)) for index in range(sample_count) + ] + handles_to_probe = [handles[index] for index in sampled_indices] + + semaphore = asyncio.Semaphore(self._pool.concurrency) + + async def _verify_one(handle: SandboxHandle) -> None: + async with semaphore: + await self._verify_created_handle(handle) + + results = await asyncio.gather( + *(_verify_one(handle) for handle in handles_to_probe), + return_exceptions=True, + ) + errors = [result for result in results if isinstance(result, Exception)] + if errors: + raise OpenSandboxCreateVerificationError( + "One or more OpenSandbox sandboxes failed create probe " + f"verification; failed={len(errors)}, total={len(handles)}" + ) from errors[0] + + async def _cleanup_failed_create_handle(self, handle: SandboxHandle) -> None: + try: + await self.close(handle, delete=True) + except Exception as e: + LOGGER.warning( + "Failed to clean up OpenSandbox sandbox after create probe failure; sandbox_id=%s; error=%r", + handle.sandbox_id, + e, + ) + + async def _connect_after_create(self, handle: SandboxHandle, spec: SandboxSpec) -> SandboxHandle: + """Reconnect after SDK create so follow-up calls use a fresh SDK handle.""" + timeout_s = spec.ready_timeout_s + if timeout_s is None: + timeout_s = self._create.timeout_s + if timeout_s is None: + timeout_s = self._create.connect_attempt_timeout_s + + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + loop = asyncio.get_running_loop() + deadline = loop.time() + float(timeout_s) + last_exception: BaseException | None = None + + while True: + remaining_s = deadline - loop.time() + if remaining_s <= 0: + error = OpenSandboxCreateTimeoutError( + "Timed out connecting to OpenSandbox sandbox after SDK create; " + f"sandbox_id={handle.sandbox_id!r}, timeout_s={timeout_s:g}" + ) + raise error from last_exception + + attempt_timeout_s = min(self._create.connect_attempt_timeout_s, remaining_s) + try: + sandbox = await asyncio.wait_for( + Sandbox.connect( + handle.sandbox_id, + connection_config=self._exec_connection_config(request_timeout_s=attempt_timeout_s), + connect_timeout=timedelta(seconds=attempt_timeout_s), + skip_health_check=True, + ), + timeout=attempt_timeout_s, + ) + return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + except asyncio.CancelledError: + raise + except BaseException as e: + last_exception = e + if not _is_retryable_create_error(e): + raise + sleep_s = min(self._create.connect_poll_s, max(deadline - loop.time(), 0.0)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + + async def _create_once(self, spec: SandboxSpec) -> SandboxHandle: + """Create a sandbox through ``opensandbox.Sandbox.create``.""" + if spec.extensions.get("poolRef") and self._connection.use_server_proxy is False: + raise ValueError( + "OpenSandbox pooled creation requires " + "use_server_proxy=True so SDK calls are routed through the " + "server proxy and do not rely on stale cached pod endpoints." + ) + + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + + kwargs: dict[str, Any] = { + "env": spec.env, + "metadata": spec.metadata, + "resource": spec.resources, + "extensions": spec.extensions, + "connection_config": self._exec_connection_config(request_timeout_s=self._create.request_timeout_s), + } + if spec.image is not None: + kwargs["image"] = spec.image + if spec.snapshot_id is not None: + kwargs["snapshot_id"] = spec.snapshot_id + if spec.timeout_s is not None: + kwargs["timeout"] = timedelta(seconds=spec.timeout_s) + if spec.ready_timeout_s is not None: + kwargs["ready_timeout"] = timedelta(seconds=spec.ready_timeout_s) + if spec.entrypoint is not None: + kwargs["entrypoint"] = spec.entrypoint + if spec.platform is not None: + kwargs["platform"] = _to_platform_spec(spec.platform) + if spec.volumes is not None: + kwargs["volumes"] = _to_volumes(spec.volumes) + if self._create.skip_health_check: + kwargs["skip_health_check"] = True + elif spec.skip_health_check is not None: + kwargs["skip_health_check"] = spec.skip_health_check + + timeout_s = self._create.timeout_s + if timeout_s is None and self._connection.request_timeout_s is not None: + timeout_s = float(self._connection.request_timeout_s) + + sandbox_id: str | None = None + sandbox: Any | None = None + try: + async with observability_span( + "sandbox.create_api", + phase="startup", + attributes={ + "provider": self.name, + "image": spec.image, + "pool_ref": spec.extensions.get("poolRef"), + "skip_health_check": self._create.skip_health_check, + "exec_use_server_proxy": self._connection.exec_use_server_proxy, + }, + ): + if timeout_s is None: + sandbox = await Sandbox.create(**kwargs) + else: + sandbox = await asyncio.wait_for( + Sandbox.create(**kwargs), + timeout=timeout_s, + ) + sandbox_id = str(sandbox.id) + except TimeoutError as e: + error = OpenSandboxCreateTimeoutError( + "Timed out creating OpenSandbox sandbox after " + f"{timeout_s:g}s; image={spec.image!r}, " + f"poolRef={spec.extensions.get('poolRef')!r}, " + f"ready_timeout_s={spec.ready_timeout_s!r}" + ) + raise error from e + if sandbox is None or sandbox_id is None: + raise RuntimeError("OpenSandbox SDK create returned no sandbox handle") + created_handle = SandboxHandle( + sandbox_id=sandbox_id, + provider_name=self.name, + raw=sandbox, + ) + handle = created_handle + try: + if self._create.skip_health_check: + handle = await self._connect_after_create(created_handle, spec) + await self._verify_created_handle(handle) + except Exception: + await self._cleanup_failed_create_handle(created_handle) + raise + return handle + + async def _create_with_retries( + self, + spec: SandboxSpec, + *, + semaphore: asyncio.Semaphore | None = None, + ) -> SandboxHandle: + AsyncRetrying, retry_if_exception, stop_after_attempt, wait_random_exponential = _require_tenacity() + retry_policy = AsyncRetrying( + retry=retry_if_exception(_is_retryable_create_error), + stop=stop_after_attempt(self._create.retries + 1), + wait=wait_random_exponential( + multiplier=self._create.retry_delay_s, + max=self._create.retry_max_delay_s, + ), + before_sleep=_log_create_retry, + reraise=True, + ) + async for attempt in retry_policy: + with attempt: + if semaphore is None: + return await self._create_once(spec) + async with semaphore: + return await self._create_once(spec) + + raise OpenSandboxBatchCreateError("OpenSandbox create retry loop did not run") + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + """Create one sandbox through the configured OpenSandbox path.""" + spec = self._with_default_image_pull_policy(_normalize_spec(spec)) + async with observability_span( + "sandbox.create", + phase="startup", + attributes={ + "provider": self.name, + "image": spec.image, + "image_pull_policy": spec.extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY), + "pool_ref": spec.extensions.get("poolRef"), + }, + ): + return await self._create_with_retries(spec) + + async def _close_many( + self, + handles: list[SandboxHandle], + *, + delete: bool, + ) -> list[Any]: + semaphore = asyncio.Semaphore(self._pool.concurrency) + + async def _close_one(handle: SandboxHandle) -> Any: + async with semaphore: + return await self.close(handle, delete=delete) + + return list( + await asyncio.gather( + *(_close_one(handle) for handle in handles), + return_exceptions=True, + ) + ) + + def _validate_sdk_pool_spec(self, spec: SandboxSpec) -> None: + if spec.image is None: + raise ValueError("OpenSandbox SDK pool requires SandboxSpec.image") + if spec.snapshot_id is not None: + raise ValueError("OpenSandbox SDK pool does not support snapshot_id") + + def _to_pool_creation_spec(self, spec: SandboxSpec) -> Any: + self._validate_sdk_pool_spec(spec) + _, _, PoolCreationSpec, _ = _require_opensandbox_sdk_pool() + return PoolCreationSpec( + image=spec.image, + entrypoint=spec.entrypoint, + resource=spec.resources or None, + env=spec.env or None, + metadata=spec.metadata or None, + extensions=spec.extensions or None, + platform=_to_platform_spec(spec.platform) if spec.platform is not None else None, + volumes=_to_volumes(spec.volumes) if spec.volumes is not None else None, + ) + + async def _wait_sdk_pool_idle( + self, + pool: Any, + *, + spec: SandboxSpec, + requested: int, + timeout_s: float, + allow_partial: bool, + ) -> int: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + last_progress_at = loop.time() + last_idle = 0 + last_snapshot: Any = None + + while True: + last_snapshot = await pool.snapshot() + idle_count = int(getattr(last_snapshot, "idle_count", 0) or 0) + if idle_count >= requested: + return requested + if idle_count > last_idle: + last_idle = idle_count + last_progress_at = loop.time() + pool_config = getattr(pool, "_config", None) + record_event( + "sample", + "sandbox.opensandbox.sdk_pool.readiness", + attributes={ + "provider": self.name, + "pool_name": getattr(pool_config, "pool_name", None), + "requested": requested, + "idle_count": idle_count, + "state": getattr(getattr(last_snapshot, "state", None), "value", None), + }, + ) + + now = loop.time() + progress_timeout_s = self._pool.progress_timeout_s + if progress_timeout_s is not None and now - last_progress_at >= progress_timeout_s: + if allow_partial and idle_count > 0: + return idle_count + error = OpenSandboxCreateTimeoutError( + "Timed out waiting for OpenSandbox SDK pool warmup progress " + f"after {progress_timeout_s:g}s; requested={requested}, " + f"idle={idle_count}, snapshot={last_snapshot!r}" + ) + raise error + if now >= deadline: + if allow_partial and idle_count > 0: + return idle_count + error = OpenSandboxCreateTimeoutError( + "Timed out waiting for OpenSandbox SDK pool warmup after " + f"{timeout_s:g}s; requested={requested}, idle={idle_count}, " + f"snapshot={last_snapshot!r}" + ) + raise error + await asyncio.sleep(self._pool.acquire_poll_interval_s) + + async def _direct_exec_handle_for_acquired_sandbox(self, sandbox: Any, spec: SandboxSpec) -> SandboxHandle: + handle = SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + if self._connection.exec_use_server_proxy is None and not self._create.skip_health_check: + return handle + return await self._connect_after_create(handle, spec) + + async def _create_batch_sdk_pool( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool, + ) -> list[SandboxHandle]: + AcquirePolicy, InMemoryAsyncPoolStateStore, _, SandboxPoolAsync = _require_opensandbox_sdk_pool() + ready_timeout_s = float( + spec.ready_timeout_s or self._create.timeout_s or self._connection.request_timeout_s or 300.0 + ) + idle_timeout_s = float(self._pool.idle_timeout_s or spec.timeout_s or max(ready_timeout_s * 2.0, 3600.0)) + primary_lock_ttl_s = float(self._pool.primary_lock_ttl_s or max(ready_timeout_s + 60.0, 60.0)) + pool_name = f"nemo-gym-{uuid4().hex[:12]}" + + async def _warmup_preparer(sandbox: Any) -> None: + if self._probe.command is None: + return + handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) + try: + await self._verify_created_handle(handle) + finally: + if handle.raw is not sandbox: + try: + await self._await_sdk_call( + handle.raw.close(), + operation="close warmup direct handle", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + LOGGER.warning( + "Failed to close temporary OpenSandbox direct exec handle for sandbox %r: %r", + handle.sandbox_id, + e, + ) + + pool = SandboxPoolAsync( + pool_name=pool_name, + max_idle=count, + warmup_concurrency=self._pool.concurrency, + state_store=InMemoryAsyncPoolStateStore(), + connection_config=self._connection_config(request_timeout_s=self._create.request_timeout_s), + creation_spec=self._to_pool_creation_spec(spec), + reconcile_interval=timedelta(seconds=self._pool.reconcile_interval_s), + primary_lock_ttl=timedelta(seconds=primary_lock_ttl_s), + acquire_ready_timeout=timedelta(seconds=ready_timeout_s), + warmup_ready_timeout=timedelta(seconds=ready_timeout_s), + warmup_sandbox_preparer=_warmup_preparer, + acquire_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), + warmup_skip_health_check=bool(self._create.skip_health_check or spec.skip_health_check), + idle_timeout=timedelta(seconds=idle_timeout_s), + ) + handles: list[SandboxHandle] = [] + async with observability_span( + "sandbox.sdk_pool.create_batch", + phase="startup", + attributes={ + "provider": self.name, + "count": count, + "pool_name": pool_name, + "pool_ref": spec.extensions.get("poolRef"), + "exec_use_server_proxy": self._connection.exec_use_server_proxy, + }, + ): + try: + await pool.start() + ready_count = await self._wait_sdk_pool_idle( + pool, + spec=spec, + requested=count, + timeout_s=ready_timeout_s, + allow_partial=allow_partial, + ) + await pool.resize(0) + sandbox_timeout = _seconds_to_timedelta(spec.timeout_s) + for index in range(ready_count): + sandbox = await pool.acquire( + sandbox_timeout=sandbox_timeout, + policy=AcquirePolicy.FAIL_FAST, + ) + handle = await self._direct_exec_handle_for_acquired_sandbox(sandbox, spec) + handles.append(handle) + LOGGER.info( + "Acquired OpenSandbox SDK pool sandbox %s/%s: %s", + index + 1, + ready_count, + sandbox.id, + ) + return handles + except Exception: + await self._close_many(handles, delete=True) + raise + finally: + try: + await pool.shutdown(graceful=False) + finally: + await pool.release_all_idle() + + async def _create_batch_sdk( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several sandboxes through the OpenSandbox SDK pool.""" + if count < 1: + raise ValueError("count must be >= 1") + return await self._create_batch_sdk_pool( + spec, + count, + allow_partial=allow_partial, + ) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + """Create several equivalent OpenSandbox sandboxes.""" + if count < 1: + raise ValueError("count must be >= 1") + spec = self._with_default_image_pull_policy(_normalize_spec(spec)) + async with observability_span( + "sandbox.create_batch", + phase="startup", + attributes={ + "provider": self.name, + "count": count, + "allow_partial": allow_partial, + "image": spec.image, + "image_pull_policy": spec.extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY), + "pool_ref": spec.extensions.get("poolRef"), + }, + ): + return await self._create_batch_sdk( + spec, + count, + allow_partial=allow_partial, + ) + + def handle_reference(self, handle: SandboxHandle) -> dict[str, Any]: + """Build a loop-neutral reference for a sandbox handle. + + OpenSandbox SDK handles are bound to the event loop where they were + created. Prewarmed handles may cross from a FastAPI prewarm request into + a thread-pool runner, so only pass a serializable reference across that + boundary and re-materialize SDK adapters in the consuming event loop. + """ + return { + "kind": "sandbox_id", + "provider": self.name, + "sandbox_id": handle.sandbox_id, + } + + async def materialize_handle(self, reference: dict[str, Any]) -> SandboxHandle: + """Create a loop-local handle from ``handle_reference`` output.""" + kind = reference.get("kind") + if kind == "sandbox_id": + return await self.connect(str(reference["sandbox_id"])) + raise ValueError(f"Unsupported OpenSandbox handle reference kind: {kind!r}") + + async def connect(self, sandbox_id: str) -> SandboxHandle: + """Connect to an existing OpenSandbox sandbox.""" + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + kwargs: dict[str, Any] = { + "connection_config": self._exec_connection_config(), + } + if self._connection.connect_timeout_s is not None: + kwargs["connect_timeout"] = timedelta(seconds=self._connection.connect_timeout_s) + sandbox = await Sandbox.connect(sandbox_id, **kwargs) + return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + + def _command_retry_count(self) -> int: + return ( + self._operations.retries if self._operations.command_retries is None else self._operations.command_retries + ) + + async def _exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + retries: int | None = None, + ) -> SandboxExecResult: + """Run a command inside an OpenSandbox sandbox.""" + _, _, RunCommandOpts, _, _ = _require_opensandbox_sdk() + + opts_kwargs: dict[str, Any] = {} + if cwd is not None: + opts_kwargs["working_directory"] = cwd + if env is not None: + opts_kwargs["envs"] = env + if timeout_s is not None: + opts_kwargs["timeout"] = timedelta(seconds=timeout_s) + + effective_command = command + if isinstance(user, int): + opts_kwargs["uid"] = user + elif isinstance(user, str) and user != "root": + effective_command = f"su -s /bin/sh -c {shlex.quote(command)} {shlex.quote(user)}" + + sdk_timeout_s = ( + float(timeout_s) + 60.0 + if timeout_s is not None + else ( + float(self._connection.request_timeout_s) if self._connection.request_timeout_s is not None else None + ) + ) + effective_retries = self._command_retry_count() if retries is None else retries + async with observability_span( + "sandbox.exec", + phase="execution", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "sdk_timeout_s": sdk_timeout_s, + "retries": effective_retries, + "command": command, + }, + ): + execution = await self._await_sdk_operation( + lambda: handle.raw.commands.run(effective_command, opts=RunCommandOpts(**opts_kwargs)), + operation="command run", + sandbox_id=handle.sandbox_id, + timeout_s=sdk_timeout_s, + retries=effective_retries, + ) + stdout = "\n".join(msg.text for msg in execution.logs.stdout) or None + stderr_parts = [msg.text for msg in execution.logs.stderr] + if execution.error is not None: + stderr_parts.append(f"{execution.error.name}: {execution.error.value}") + stderr = "\n".join(stderr_parts) or None + if execution.exit_code is not None: + return_code = execution.exit_code + elif execution.error is not None: + return_code = 1 + else: + return_code = 0 + + return SandboxExecResult(stdout=stdout, stderr=stderr, return_code=return_code) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + """Run a command inside an OpenSandbox sandbox.""" + return await self._exec( + handle, + command, + cwd=cwd, + env=env, + timeout_s=timeout_s, + user=user, + retries=self._command_retry_count(), + ) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + """Write one file into an OpenSandbox sandbox.""" + async with observability_span( + "sandbox.write_file", + phase="setup", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "target_path": target_path, + "bytes": len(data), + }, + ): + await self._await_sdk_operation( + lambda: handle.raw.files.write_file(target_path, data), + operation=f"write_file({target_path})", + sandbox_id=handle.sandbox_id, + timeout_s=float(self._connection.request_timeout_s) + if self._connection.request_timeout_s is not None + else None, + ) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + """Read one file from an OpenSandbox sandbox.""" + async with observability_span( + "sandbox.read_file", + phase="execution", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "source_path": source_path, + }, + ): + return await self._await_sdk_operation( + lambda: handle.raw.files.read_bytes(source_path), + operation=f"read_file({source_path})", + sandbox_id=handle.sandbox_id, + timeout_s=float(self._connection.request_timeout_s) + if self._connection.request_timeout_s is not None + else None, + ) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + """Upload one local file into an OpenSandbox sandbox.""" + await self.write_file(handle, target_path, source_path.read_bytes()) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + """Download one file from an OpenSandbox sandbox.""" + target_path.parent.mkdir(parents=True, exist_ok=True) + target_path.write_bytes(await self.read_file(handle, source_path)) + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + """Close local SDK resources and optionally terminate the sandbox.""" + async with observability_span( + "sandbox.close", + phase="cleanup", + attributes={ + "provider": self.name, + "sandbox_id": handle.sandbox_id, + "delete": delete, + }, + ): + kill_error: Exception | None = None + if delete: + try: + await self._await_sdk_operation( + lambda: handle.raw.kill(), + operation="kill", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + if not _is_missing_sandbox_delete_error(e): + kill_error = e + else: + LOGGER.info( + "OpenSandbox sandbox %r was already deleted during close", + handle.sandbox_id, + ) + + close_error: Exception | None = None + try: + await self._await_sdk_call( + handle.raw.close(), + operation="close", + sandbox_id=handle.sandbox_id, + timeout_s=self._operations.close_timeout_s, + ) + except Exception as e: + close_error = e + LOGGER.warning( + "Timed out or failed while closing local OpenSandbox SDK handle for sandbox %r: %r", + handle.sandbox_id, + e, + ) + record_event( + "warning", + "sandbox.opensandbox.sdk_handle_close_error", + attributes=_sdk_error_attributes(e, operation="close", sandbox_id=handle.sandbox_id), + ) + + if kill_error is not None: + if close_error is not None: + raise RuntimeError( + "Failed to delete and close OpenSandbox sandbox " + f"{handle.sandbox_id!r}: delete_error={kill_error!r}, " + f"close_error={close_error!r}" + ) from kill_error + raise kill_error + if close_error is not None: + if delete: + return + raise close_error diff --git a/nemo_gym/sandbox/providers/registry.py b/nemo_gym/sandbox/providers/registry.py new file mode 100644 index 0000000000..43d6aebfce --- /dev/null +++ b/nemo_gym/sandbox/providers/registry.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider registration utilities.""" + +from typing import TypeAlias + +from nemo_gym.sandbox.config import SandboxProviderConfig +from nemo_gym.sandbox.providers.base import SandboxProvider + + +ProviderClass: TypeAlias = type[SandboxProvider] + +G_PROVIDER_REGISTRY: dict[str, ProviderClass] = {} + + +def register_provider(name: str, provider_class: ProviderClass) -> None: + """Register a sandbox provider class.""" + if not name: + raise ValueError("Provider name must be non-empty") + if name in G_PROVIDER_REGISTRY: + raise ValueError(f"Sandbox provider {name!r} is already registered") + G_PROVIDER_REGISTRY[name] = provider_class + + +def get_provider_class(name: str) -> ProviderClass: + """Return a registered provider class.""" + try: + return G_PROVIDER_REGISTRY[name] + except KeyError as e: + available = ", ".join(sorted(G_PROVIDER_REGISTRY)) or "" + raise ValueError(f"Unknown sandbox provider {name!r}. Available providers: {available}") from e + + +def create_provider(config: SandboxProviderConfig) -> SandboxProvider: + """Instantiate a provider from ``env.sandbox.provider`` config.""" + provider_class = get_provider_class(config["name"]) + if "kwargs" in config: + return provider_class(**config["kwargs"]) + return provider_class() + + +def list_providers() -> list[str]: + """List registered provider names.""" + return sorted(G_PROVIDER_REGISTRY) + + +def _register_builtins() -> None: + from nemo_gym.sandbox.providers.opensandbox import OpenSandboxProvider + + if "opensandbox" not in G_PROVIDER_REGISTRY: + register_provider("opensandbox", OpenSandboxProvider) + + +_register_builtins() diff --git a/pyproject.toml b/pyproject.toml index 3230b17533..914093fa31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -215,6 +215,25 @@ docs = [ ] [project.optional-dependencies] +sandbox = [ + # Tenacity: Retry helpers used by sandbox providers. + # Updated: Sat May 09, 2026 with tenacity==9.1.4 + # License: Apache 2.0 https://github.com/jd/tenacity/blob/master/LICENSE + "tenacity>=9.1.4", + + # OpenSandbox SDK: used by the OpenSandbox sandbox provider for create/exec/delete and SDK pool creation. + # Updated: Sat May 16, 2026 with opensandbox>=0.1.9 + # License: Apache 2.0 + "opensandbox>=0.1.9", + + # OpenTelemetry Python: SDK-backed sandbox traces, metrics, and OTLP HTTP export. + # Updated: Mon May 18, 2026 with opentelemetry-sdk>=1.36.0 and OTLP HTTP exporter. + # License: Apache 2.0 https://github.com/open-telemetry/opentelemetry-python/blob/main/LICENSE + "opentelemetry-api>=1.36.0", + "opentelemetry-sdk>=1.36.0", + "opentelemetry-exporter-otlp-proto-http>=1.36.0", +] + # We include dev dependencies as an extra since technically each server module is a consumer (which means we cannot use dependency groups, which are intended to be within a project). dev = [ # Pre-commit: Used for pre-commit hooks. @@ -388,7 +407,7 @@ ng_reinstall = "nemo_gym.cli:reinstall" [tool.setuptools.packages.find] where = ["."] -include = ["benchmarks", "resources_servers", "responses_api_agents", "responses_api_models", "nemo_gym"] +include = ["benchmarks", "resources_servers", "responses_api_agents", "responses_api_models", "nemo_gym", "nemo_gym.*"] ################################################ # Testing @@ -441,6 +460,8 @@ omit = [ "results/*", "/tmp/*", "benchmarks/*", + "nemo_gym/sandbox/observability/recorder.py", + "nemo_gym/sandbox/providers/opensandbox/provider.py", ] data_file = "results/.coverage" concurrency = ["thread", "multiprocessing"] diff --git a/responses_api_agents/mini_swe_agent_2/.gitignore b/responses_api_agents/mini_swe_agent_2/.gitignore new file mode 100644 index 0000000000..68bcbc9609 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/.gitignore @@ -0,0 +1 @@ +results/ \ No newline at end of file diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md new file mode 100644 index 0000000000..0058e909aa --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -0,0 +1,428 @@ +# Mini-SWE-Agent 2 Sandbox Agent + +A NeMo Gym Responses API agent that integrates +[mini-swe-agent](https://github.com/SWE-agent/mini-swe-agent) v2 for evaluating +language models on SWE-bench style software engineering tasks through the public +`nemo_gym.sandbox` API. + +This agent intentionally keeps only the sandbox-backed path. It does not carry +over the older Docker/Singularity mini-SWE integration. + +## Contents + +- [Mini-SWE-Agent 2 Sandbox Agent](#mini-swe-agent-2-sandbox-agent) + - [Contents](#contents) + - [Overview](#overview) + - [Dataset Information](#dataset-information) + - [Configuration](#configuration) + - [Agent Configuration](#agent-configuration) + - [Model Parameters](#model-parameters) + - [Usage](#usage) + - [Server](#server) + - [Collect Rollouts](#collect-rollouts) + - [Observability](#observability) + - [Sandbox Environment Adapter](#sandbox-environment-adapter) + - [Environment Lifecycle](#environment-lifecycle) + - [Contributing](#contributing) + - [Licensing Information](#licensing-information) + - [Dependencies](#dependencies) + +## Overview + +`mini_swe_agent_2` runs mini-swe-agent's synchronous SWE-bench harness while +creating and executing each task environment through Gym's provider-neutral +sandbox facade. The validated path in this directory is: + +- mini-swe-agent `2.1.0` +- SWE-bench task rows, including SWE-bench Verified +- `env: sandbox` +- `responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment` +- OpenSandbox through `nemo_gym.sandbox.providers.opensandbox` +- OpenTelemetry sandbox observability through `nemo_gym.sandbox.observability` + +For each `/run` request, `MiniSWEAgent.run()` loads mini-swe-agent's built-in +`swebench.yaml`, injects sandbox settings, runs mini-swe-agent in a Ray remote +task, evaluates the generated patch with the SWE-bench harness, and returns a +Gym verify response with reward `1.0` only when the instance is resolved and the +evaluation report includes test status. + +`MiniSWEAgent.setup_webserver()` also registers `/v1/responses`, but +`MiniSWEAgent.responses()` is intentionally not implemented in this agent. The +supported eval path is `/run`, typically via `ng_collect_rollouts`. + +## Dataset Information + +- Eval data - [princeton-nlp/SWE-bench_Verified](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Verified) + is the primary validation target. It contains 500 human-validated SWE-bench + test instances. +- The rollout input JSONL should preserve the SWE-bench instance fields needed + by `swegym` and `swebench`, such as `instance_id`, `repo`, `base_commit`, + `problem_statement`, `patch`, `test_patch`, `FAIL_TO_PASS`, `PASS_TO_PASS`, + and related version fields. +- Each row must also include `responses_create_params`. Extra top-level + SWE-bench fields are accepted by the agent request model and passed into + mini-swe-agent as the instance dictionary. + +Example row shape: + +```json +{ + "instance_id": "django__django-13410", + "repo": "django/django", + "base_commit": "...", + "problem_statement": "...", + "patch": "...", + "test_patch": "...", + "FAIL_TO_PASS": ["..."], + "PASS_TO_PASS": ["..."], + "responses_create_params": { + "input": [], + "temperature": 0.6, + "top_p": 1.0, + "max_output_tokens": 16384 + } +} +``` + +When `image_name` is present on a row, the agent uses it directly. Otherwise it +derives the SWE-bench image from `instance_id` and `subset`: + +- `subset: verified` uses `swebench/sweb.eval.x86_64.:latest` with `__` + replaced by `_1776_`. +- Other subsets use `xingyaoww/sweb.eval.x86_64.:latest` with `__` replaced + by `_s_`. + +Configured `sandbox_spec.image_rewrites` are applied inside +`MiniSWESandboxEnvironment`; the default OpenSandbox config rewrites +`swebench/` to `mirror.gcr.io/swebench/`. + +## Configuration + +### Agent Configuration + +Path - `responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml` + +```yaml +mini_swe_agent_2: + responses_api_agents: + mini_swe_agent_2: + entrypoint: app.py + domain: coding + description: Software engineering tasks driven by mini-swe-agent harness on OpenSandbox. + value: Improve agentic software engineering capabilities. + model_server: + type: responses_api_models + name: policy_model + concurrency: 64 + env: sandbox + sandbox_provider: + name: opensandbox + kwargs: + connection: + domain: opensandbox-server.opensandbox-system.svc.cluster.local + api_key: ${oc.env:OPENSANDBOX_API_KEY} + protocol: http + use_server_proxy: true + exec_use_server_proxy: true + request_timeout_s: 300 + create: + request_timeout_s: 1200 + timeout_s: 1200 + skip_health_check: true + retries: 10 + retry_delay_s: 5.0 + retry_max_delay_s: 90.0 + probe: + timeout_s: 60 + deadline_s: 180 + stable_count: 2 + stable_delay_s: 1.0 + operations: + retries: 5 + retry_delay_s: 1.0 + retry_max_delay_s: 45.0 + command_retries: 3 + close_timeout_s: 30 + sandbox_spec: + timeout_s: 18000 + ready_timeout_s: 1200 + resources: + cpu: "1" + memory: 8Gi + ephemeral-storage: 20Gi + platform: + os: linux + arch: amd64 + image_rewrites: + - from: swebench/ + to: mirror.gcr.io/swebench/ + metadata: + benchmark: swebench-verified + harness: mini-swe-agent + sandbox-api: opensandbox-sdk + sandbox_environment_kwargs: + cwd: /testbed + conda_env: testbed + activate_conda: true + user: root + delete: true + run_golden: false + step_timeout: 600 + eval_timeout: 1800 + skip_if_exists: false + step_limit: 250 + observability: + enabled: false + output_dir: results/mini_swe_agent_2_observability/{trajectory_id} + export_traces: true + run_id: mini-swe-agent-2 + job_name: mini-swe-agent-2 + otel: + service_name: mini-swe-agent-2 + resource_attributes: + benchmark: swebench-verified + harness: mini_swe_agent_2 + command_titles: + strip_prefixes: + - "cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed &&" + rules: + - line_starts_with: + - "pytest " + - "python -m pytest " + - "./tests/runtests.py " + search: last + title: "run verifier: {line}" +``` + +Optional `sandbox_resource_profiles` can be configured as a list of resource +maps. When present, the agent hashes `instance_id` and deterministically merges +one profile into `sandbox_spec.resources`. This is useful for spreading +SWE-bench tasks across a small set of resource sizes without changing the input +data. + +### Model Parameters + +`MiniSWEAgent.run()` maps supported Responses API fields into mini-swe-agent +chat-completions kwargs: + +- `temperature`, `top_p`, `top_logprobs`, and `parallel_tool_calls` pass through. +- `max_output_tokens` becomes `max_tokens`. +- `responses_create_params.metadata.extra_body` must be a JSON object and is + passed as `extra_body`. +- `responses_create_params.metadata.chat_template_kwargs` must be a JSON object + and is nested under `extra_body.chat_template_kwargs`. +- `tool_choice` comes from the agent config when set, otherwise from the request. + The special value `bash` expands to the OpenAI function choice for the `bash` + tool. + +Keep the requested generation budget compatible with the live vLLM deployment. +For example, a deployment served with `--max-model-len 32768` will reject +`max_output_tokens=49152`. In earlier smoke testing, that upstream vLLM rejection +surfaced in mini-swe-agent as repeated: + +```text +No tool calls found in the response. Every response MUST include at least one tool call. +``` + +That symptom was not a sandbox failure and was not a reason to force the `bash` +tool. The successful smoke kept `tool_choice=auto` and lowered +`max_output_tokens` to `16384`. + +## Usage + +### Server + +Set the policy model endpoint in `env.yaml` or with equivalent Hydra overrides: + +```yaml +policy_base_url: http://terryk-dgd-pd-1p9d-frontend.default.svc.cluster.local:8000/v1 +policy_api_key: dummy-key +policy_model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +``` + +Start the mini-swe-agent 2 server with the OpenSandbox provider and a policy +model server. The values below mirror the current +`hemild-mini-swe2-nemo-4-pd1p9d-p8-r16` eval job: + +```bash +CONFIG_PATHS="responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml" + +ng_run "+config_paths=[$CONFIG_PATHS]" \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.concurrency=64 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.step_timeout=600 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.eval_timeout=1800 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.step_limit=50 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=false \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.enabled=true \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.output_dir=results/mini_swe_agent_2_observability/{trajectory_id}' \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.run_id=mini-swe2-nemotron-4-pd-1p9d-pass8-r16 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.job_name=mini-swe2-nemotron-4-pd-1p9d-pass8-r16 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.otel.service_name=mini-swe2-nemotron-4-pd-1p9d \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.otel.resource_attributes.benchmark=swebench-verified \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.otel.resource_attributes.harness=mini_swe_agent_2 \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.otel.resource_attributes.endpoint_label=4-dgd-pd-1p9d \ + +mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.observability.otel.resource_attributes.run_family=mini-swe2-nemotron-4-pd-1p9d-pass8-r16 \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.resources={cpu: 500m, memory: 4Gi, ephemeral-storage: 8Gi}' \ + '+mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.sandbox_spec.metadata={benchmark: swebench-verified, harness: mini_swe_agent_2, endpoint_label: 4-dgd-pd-1p9d, run_family: mini-swe2-nemotron-4-pd-1p9d-pass8-r16}' +``` + +Use a model server config that matches the policy endpoint you are serving. The +example above uses `vllm_model`, which is the common path for hosted vLLM +`/v1/chat/completions` endpoints. + +### Collect Rollouts + +Collect eval rollouts from a SWE-bench-style JSONL file: + +```bash +ng_collect_rollouts \ + +agent_name=mini_swe_agent_2 \ + +input_jsonl_fpath=/mnt/rl-workspace/hemild/gym_eval/refactor/inputs/mini_swe_verified_smoke8.jsonl \ + +output_jsonl_fpath=results/mini_swe_agent_2_nemotron_4_pd_1p9d_pass8.jsonl \ + +limit=8 \ + +num_repeats=8 \ + +num_samples_in_parallel=64 \ + '+responses_create_params={max_output_tokens: 32768, temperature: 0.6, top_p: 0.95, metadata: {chat_template_kwargs: "{\"enable_thinking\": true}"}}' +``` + +After collecting repeated rollouts, run `ng_reward_profile` on the collected +output: + +```bash +ng_reward_profile \ + +input_jsonl_fpath=/mnt/rl-workspace/hemild/gym_eval/refactor/inputs/mini_swe_verified_smoke8.jsonl \ + +rollouts_jsonl_fpath=results/mini_swe_agent_2_nemotron_4_pd_1p9d_pass8.jsonl \ + +output_jsonl_fpath=results/mini_swe_agent_2_nemotron_4_pd_1p9d_pass8_profiled.jsonl \ + +pass_threshold=1.0 +``` + +The agent writes per-instance mini-swe-agent configs and result artifacts under +`results///`. + +These settings are adapted from the Kubernetes job family +`hemild-mini-swe2-nemo-4*`. The direct Kubernetes orchestrator also used +`SAMPLE_TIMEOUT_S=2400` as an outer per-sample guard; `ng_collect_rollouts` +does not have an exact equivalent, so use the agent's `step_timeout` and +`eval_timeout` overrides above to bound tool and verifier execution. + +## Observability + +Mini SWE Agent 2 uses the shared sandbox observability recorder for sandbox +lifecycle spans, command execution spans, verifier spans, and model request +spans. `app.py` wraps mini-swe-agent model calls in an `llm.request` span, and +`sandbox_environment.py` marks SWE-bench evaluation commands as verifier work. + +Configure observability in the agent config. Use `{trajectory_id}`, +`{instance_id}`, `{task_index}`, and `{rollout_index}` placeholders in +`output_dir`, `run_id`, `job_name`, and `otel.service_name` when you want +per-rollout trace artifacts: + +```yaml +observability: + enabled: true + output_dir: results/mini_swe_agent_2_observability/{trajectory_id} + export_traces: true + run_id: mini-swe2-nemotron-4-pd-1p9d-pass8-r16 + job_name: mini-swe2-nemotron-4-pd-1p9d-pass8-r16 + otel: + service_name: mini-swe2-nemotron-4-pd-1p9d + resource_attributes: + benchmark: swebench-verified + harness: mini_swe_agent_2 + endpoint_label: 4-dgd-pd-1p9d + run_family: mini-swe2-nemotron-4-pd-1p9d-pass8-r16 + command_titles: + strip_prefixes: + - "cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed &&" + rules: + - line_starts_with: + - "pytest " + - "python -m pytest " + - "./tests/runtests.py " + search: last + title: "run verifier: {line}" +``` + +For each trajectory, the recorder writes OpenTelemetry JSON under: + +```text +/traces/otel_traces.json +``` + +The default config leaves observability disabled. Turn it on with a Hydra +override, as shown in the server launch example, or edit the YAML directly. + +## Sandbox Environment Adapter + +`MiniSWESandboxEnvironment` adapts mini-swe-agent's synchronous environment +contract to `nemo_gym.sandbox.Sandbox`. + +When `env` is `sandbox`, Gym injects this environment config before calling +mini-swe-agent: + +```yaml +environment: + environment_class: responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment + image: + provider: + name: opensandbox + kwargs: ... + spec: + resources: ... + platform: ... + metadata: ... +``` + +### Environment Lifecycle + +`MiniSWESandboxEnvironment.__init__()`: + +- Validates that a sandbox provider was configured. +- Builds a `SandboxSpec` from the task image, environment variables, metadata, + resources, platform, volumes, provider-specific extensions, and health-check + settings. +- Applies Gym image rewrites before creating the sandbox. +- Adds standard metadata such as `nemo_gym_agent=mini_swe_agent_2` and + `instance_id`. +- Creates a `Sandbox` facade and calls `Sandbox.create(...)`. + +`execute()`: + +- Receives mini-swe-agent's command action. +- Applies the configured working directory and timeout. +- Optionally wraps the command in `conda activate ` for SWE-bench images + that expect a prebuilt conda environment. +- Calls `Sandbox.exec(...)` as the configured user, root by default. +- Returns mini-swe-agent's expected sync response shape: + +```python +{ + "output": "...", + "returncode": 0, + "exception_info": "", +} +``` + +`_check_finished()` preserves mini-swe-agent's submit sentinel behavior. If the +command output begins with `COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` and the +command succeeded, it raises `minisweagent.exceptions.Submitted` with the final +submission payload. + +`cleanup()` calls `Sandbox.close(..., delete=config.delete)` and then +`Sandbox.shutdown()` to release provider-owned async resources and stop the sync +facade's private loop. + +## Contributing + +Please refer to the main NeMo Gym documentation for contributing guidelines. + +## Licensing Information + +- **Code**: Apache 2.0 +- **SWE-bench Verified**: MIT + +### Dependencies + +- **nemo_gym**: Apache 2.0 +- **mini-swe-agent**: MIT +- **SWE-Bench-Package / swegym**: MIT diff --git a/responses_api_agents/mini_swe_agent_2/__init__.py b/responses_api_agents/mini_swe_agent_2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_agents/mini_swe_agent_2/app.py b/responses_api_agents/mini_swe_agent_2/app.py new file mode 100644 index 0000000000..419a5e89de --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/app.py @@ -0,0 +1,767 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import hashlib +import json +import sys +import time +import traceback +from asyncio import Semaphore +from pathlib import Path +from typing import Any, Callable, Literal, Optional, cast +from uuid import uuid4 + +import ray +import yaml +from fastapi import Body, FastAPI +from minisweagent.config import builtin_config_dir, get_config_path +from pydantic import ConfigDict + +from nemo_gym.base_resources_server import ( + BaseRunRequest, + BaseVerifyRequest, + BaseVerifyResponse, +) +from nemo_gym.base_responses_api_agent import ( + BaseResponsesAPIAgentConfig, + SimpleResponsesAPIAgent, +) +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import ( + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.sandbox.observability import ( + build_recorder_from_config, + event_context, + observability_sync_span, + use_recorder, +) +from nemo_gym.server_utils import ( + ServerClient, + get_first_server_config_dict, +) + + +class MiniSWEAgentConfig(BaseResponsesAPIAgentConfig): + model_server: ModelServerRef + env: Literal["sandbox"] + concurrency: int + sandbox_provider: Optional[dict[str, Any]] = None + sandbox_spec: Optional[dict[str, Any]] = None + sandbox_environment_kwargs: Optional[dict[str, Any]] = None + run_golden: bool = False + step_timeout: int = 600 + eval_timeout: int = 1800 + skip_if_exists: bool = False + step_limit: int = 250 + tool_choice: Optional[str | dict[str, Any]] = None + sandbox_resource_profiles: Optional[list[dict[str, str]]] = None + observability: Optional[dict[str, Any]] = None + + +class MiniSWEAgentRunRequest(BaseRunRequest): + model_config = ConfigDict(extra="allow") + + +class MiniSWEAgentVerifyRequest(BaseVerifyRequest): + model_config = ConfigDict(extra="allow") + + +class MiniSWEAgentVerifyResponse(BaseVerifyResponse): + model_config = ConfigDict(extra="allow") + + +@ray.remote( + scheduling_strategy="SPREAD", + runtime_env={ + "py_executable": sys.executable, + }, +) +def runner_ray_remote(runner: Callable, params: dict[str, Any]) -> Any: + return runner(**params) + + +def _json_dict_from_metadata(value: Any, *, field_name: str) -> dict[str, Any]: + if value is None: + return {} + if isinstance(value, dict): + return value + if isinstance(value, str): + parsed = json.loads(value) + if isinstance(parsed, dict): + return parsed + raise ValueError(f"responses_create_params.metadata.{field_name} must be a JSON object") + + +def _responses_create_params_to_model_kwargs( + params: dict[str, Any], + *, + default_tool_choice: Any = None, +) -> dict[str, Any]: + """Convert Gym Responses API rollout params into mini-swe-agent chat-completions kwargs.""" + model_kwargs: dict[str, Any] = {} + for key in ("temperature", "top_p", "top_logprobs", "parallel_tool_calls"): + value = params.get(key) + if value is not None: + model_kwargs[key] = value + + max_output_tokens = params.get("max_output_tokens") + if max_output_tokens is not None: + model_kwargs["max_tokens"] = max_output_tokens + + metadata = params.get("metadata") or {} + extra_body = _json_dict_from_metadata(metadata.get("extra_body"), field_name="extra_body") + chat_template_kwargs = _json_dict_from_metadata( + metadata.get("chat_template_kwargs"), + field_name="chat_template_kwargs", + ) + if chat_template_kwargs: + extra_body["chat_template_kwargs"] = chat_template_kwargs + if extra_body: + model_kwargs["extra_body"] = extra_body + + tool_choice = default_tool_choice if default_tool_choice is not None else params.get("tool_choice") + if tool_choice == "bash": + model_kwargs["tool_choice"] = _bash_tool_choice() + elif tool_choice is not None: + model_kwargs["tool_choice"] = tool_choice + + return model_kwargs + + +def _bash_tool_choice() -> dict[str, Any]: + return {"type": "function", "function": {"name": "bash"}} + + +class _ObservedModel: + """Add an OTel span around each mini-SWE model query.""" + + def __init__(self, model: Any, *, model_name: str) -> None: + self._model = model + self._model_name = model_name + + def __getattr__(self, name: str) -> Any: + return getattr(self._model, name) + + def query(self, messages: list[dict[str, Any]], **kwargs: Any) -> dict[str, Any]: + model_kwargs = getattr(getattr(self._model, "config", None), "model_kwargs", None) + model_kwargs = model_kwargs if isinstance(model_kwargs, dict) else {} + attributes = { + "model": self._model_name, + "message_count": len(messages), + "temperature": kwargs.get("temperature", model_kwargs.get("temperature")), + "top_p": kwargs.get("top_p", model_kwargs.get("top_p")), + "max_tokens": kwargs.get( + "max_output_tokens", + kwargs.get("max_tokens", model_kwargs.get("max_output_tokens", model_kwargs.get("max_tokens"))), + ), + "tool_choice": kwargs.get("tool_choice", model_kwargs.get("tool_choice")), + "_record_exception_stacktrace": False, + } + with observability_sync_span("llm.request", phase="llm", attributes=attributes): + return self._model.query(messages, **kwargs) + + +def _sandbox_spec_for_instance( + spec: dict[str, Any] | None, + *, + resource_profiles: list[dict[str, str]] | None, + instance_id: str, +) -> dict[str, Any]: + instance_spec = dict(spec or {}) + if not resource_profiles: + return instance_spec + + resources = dict(instance_spec.get("resources") or {}) + digest = hashlib.sha256(instance_id.encode("utf-8")).digest() + profile = resource_profiles[int.from_bytes(digest[:4], "big") % len(resource_profiles)] + resources.update(profile) + instance_spec["resources"] = resources + return instance_spec + + +def _format_template(value: Any, context: dict[str, Any]) -> Any: + if not isinstance(value, str): + return value + try: + return value.format(**context) + except (KeyError, ValueError, IndexError): + return value + + +def _observability_config_for_instance( + config: dict[str, Any] | None, + *, + instance_id: str, + task_index: Any = None, + rollout_index: Any = None, +) -> dict[str, Any] | None: + if not isinstance(config, dict): + return None + + trajectory_id = str(config.get("trajectory_id") or instance_id) + if rollout_index is not None: + try: + rollout_suffix = f"{int(rollout_index) + 1:02d}" + except (TypeError, ValueError): + rollout_suffix = str(rollout_index) + trajectory_id = f"{trajectory_id}__rollout{rollout_suffix}" + + context = { + "instance_id": instance_id, + "task_index": "" if task_index is None else task_index, + "rollout_index": "" if rollout_index is None else rollout_index, + "trajectory_id": trajectory_id, + } + formatted = dict(config) + formatted["trajectory_id"] = trajectory_id + for key in ("output_dir", "run_id", "run_span_name", "job_name"): + if key in formatted: + formatted[key] = _format_template(formatted[key], context) + + otel = dict(formatted.get("otel") or {}) + for key in ("service_name", "run_span_name", "job_name"): + if key in otel: + otel[key] = _format_template(otel[key], context) + if "resource_attributes" in otel and isinstance(otel["resource_attributes"], dict): + otel["resource_attributes"] = { + resource_key: _format_template(resource_value, context) + for resource_key, resource_value in otel["resource_attributes"].items() + } + if otel: + formatted["otel"] = otel + return formatted + + +def _swebench_config_path() -> Path: + for candidate in ( + builtin_config_dir / "extra" / "swebench.yaml", + builtin_config_dir / "benchmarks" / "swebench.yaml", + ): + if candidate.exists(): + return candidate + return builtin_config_dir / "extra" / "swebench.yaml" + + +def _swebench_image_name(instance: dict[str, Any], subset: str) -> str: + image_name = instance.get("image_name") + if image_name: + return str(image_name) + + instance_id = instance["instance_id"] + if subset == "verified": + docker_compatible_id = instance_id.replace("__", "_1776_") + return f"swebench/sweb.eval.x86_64.{docker_compatible_id}:latest".lower() + + docker_compatible_id = instance_id.replace("__", "_s_") + return f"xingyaoww/sweb.eval.x86_64.{docker_compatible_id}:latest".lower() + + +def _message_content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for item in content: + if isinstance(item, dict): + parts.append(str(item.get("text") or item.get("content") or "")) + else: + parts.append(str(item)) + return "\n".join(part for part in parts if part) + return "" if content is None else str(content) + + +def _strip_extra(item: Any) -> dict[str, Any]: + if hasattr(item, "model_dump"): + item = item.model_dump() + if not isinstance(item, dict): + return {"type": "message", "role": "user", "content": str(item)} + return {key: value for key, value in item.items() if key != "extra"} + + +def _split_trajectory_for_responses( + messages: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]: + input_messages: list[dict[str, Any]] = [] + output_items: list[dict[str, Any]] = [] + raw_responses: list[dict[str, Any]] = [] + in_initial_prompt = True + + for message in messages: + role = message.get("role") + if in_initial_prompt and role in {"system", "user"}: + input_messages.append( + {"type": "message", "role": role, "content": _message_content_to_text(message.get("content"))} + ) + continue + + in_initial_prompt = False + if message.get("object") == "response": + response = _strip_extra(message) + raw_responses.append(response) + output_items.extend(_strip_extra(item) for item in response.get("output", [])) + elif role == "assistant": + content = _message_content_to_text(message.get("content")) + if content: + output_items.append( + { + "id": message.get("id") or f"msg_{uuid4()}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content, "annotations": []}], + } + ) + for tool_call in message.get("tool_calls") or []: + function = tool_call.get("function") or {} + output_items.append( + { + "id": tool_call.get("id") or f"fc_{uuid4()}", + "type": "function_call", + "name": function.get("name") or tool_call.get("name") or "", + "call_id": tool_call.get("id") or tool_call.get("call_id") or "", + "arguments": function.get("arguments") or tool_call.get("arguments") or "{}", + } + ) + elif role == "tool": + output_items.append( + { + "type": "function_call_output", + "call_id": message.get("tool_call_id") or message.get("call_id") or "", + "output": _message_content_to_text(message.get("content")), + } + ) + elif message.get("type") == "function_call_output": + output_items.append(_strip_extra(message)) + + return input_messages, output_items, raw_responses + + +def _default_response_object() -> dict[str, Any]: + return { + "id": f"resp_{str(uuid4())}", + "created_at": int(time.time()), + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "object": "response", + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "background": False, + "max_output_tokens": None, + "max_tool_calls": None, + "previous_response_id": None, + "prompt": None, + "reasoning": { + "effort": None, + "generate_summary": None, + "summary": None, + }, + "service_tier": "default", + "status": "completed", + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "top_logprobs": 0, + "truncation": "disabled", + "usage": { + "input_tokens": 0, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 0, + }, + "user": None, + "prompt_cache_key": None, + "safety_identifier": None, + "store": True, + } + + +def _is_resolved(instance_id: str, eval_report: dict[str, Any]) -> bool: + try: + if not eval_report: + return False + report = eval_report["eval_report"][instance_id] + resolved = bool(report["resolved"]) + if not report.get("tests_status"): + return False + + tests_status = report["tests_status"] + f2f = tests_status.get("FAIL_TO_PASS", {}) + p2p = tests_status.get("PASS_TO_PASS", {}) + total_reported = ( + len(f2f.get("success", [])) + + len(f2f.get("failure", [])) + + len(p2p.get("success", [])) + + len(p2p.get("failure", [])) + ) + return resolved and total_reported > 0 + except Exception as exc: + print(f"Error in _is_resolved: {exc}", flush=True) + return False + + +def _run_eval_v2( + *, + instance: dict[str, Any], + env: Any, + model_patch: str, + instance_dir: Path, + run_id: str, + is_golden: bool, +) -> dict[str, Any]: + from swegym.harness.constants import SWEbenchInstance + from swegym.harness.docker_build import setup_logger + from swegym.harness.grading import get_eval_report + from swegym.harness.test_spec import make_test_spec + + swebench_instance = cast(SWEbenchInstance, instance) + test_spec = make_test_spec(swebench_instance) + pred = {"instance_id": test_spec.instance_id, "model_patch": model_patch} + + instance_dir.mkdir(parents=True, exist_ok=True) + log_file = instance_dir / f"run_instance_{run_id}.log" + report_path = instance_dir / f"report_{run_id}.json" + patch_file = instance_dir / f"patch_{run_id}.diff" + patch_file.write_text(model_patch) + + logger = setup_logger(test_spec.instance_id, log_file) + logger.info(f"DEBUG test_spec {test_spec}") + logger.info(f"DEBUG eval_script {test_spec.eval_script}") + + if is_golden: + env.execute(f"cat > patch.diff <<'EOF'\n{model_patch}\n\nEOF") + env.execute("git status --porcelain") + env.execute("git apply --check patch.diff") + env.execute("git apply patch.diff") + + eval_script = test_spec.eval_script.replace("#!/bin/bash", "") + result = env.execute(eval_script, is_eval=True) + test_output = result["output"] + returncode = result["returncode"] + print(f"[EVAL]{test_spec.instance_id} returncode: {returncode}", flush=True) + + test_output_path = instance_dir / f"test_output_{run_id}.txt" + test_output_path.write_text(test_output) + print(f"[EVAL]{test_spec.instance_id} Test output written to {test_output_path}", flush=True) + + report = get_eval_report( + test_spec=test_spec, + prediction=pred, + log_path=test_output_path, + include_tests_status=True, + ) + print(f"[EVAL]{test_spec.instance_id} Result: resolved: {report[test_spec.instance_id]['resolved']}", flush=True) + + report_path.write_text(json.dumps(report, indent=4)) + return { + "instance_id": test_spec.instance_id, + "model_patch": model_patch, + "eval_report": report, + } + + +def _run_swegym_v2(**params: Any) -> dict[str, Any]: + from minisweagent.agents.default import DefaultAgent + from minisweagent.environments import get_environment + from minisweagent.models import get_model + + instance = params.get("instance_dict") + if isinstance(instance, str): + instance = json.loads(instance) + if not isinstance(instance, dict): + raise ValueError("mini-swe-agent v2 path requires instance_dict") + + instance = dict(instance) + instance_id = str(params.get("instance_id") or instance["instance_id"]).lower() + instance["instance_id"] = instance_id + + output_dir = Path(params["output"]) + instance_dir = output_dir / instance_id + output_dir.mkdir(parents=True, exist_ok=True) + instance_dir.mkdir(parents=True, exist_ok=True) + + config = yaml.safe_load(get_config_path(params["config"]).read_text()) + model_config = config.setdefault("model", {}) + model_config["model_class"] = "litellm" + model_config["model_name"] = params["model"] + model_config.setdefault("cost_tracking", "ignore_errors") + model_kwargs = model_config.setdefault("model_kwargs", {}) + model_kwargs["api_key"] = params["api_key"] + model_kwargs["base_url"] = params["base_url"] + model_kwargs.pop("api_base", None) + max_output_tokens = model_kwargs.pop("max_output_tokens", None) + if max_output_tokens is not None and "max_tokens" not in model_kwargs: + model_kwargs["max_tokens"] = max_output_tokens + + environment_config = config.setdefault("environment", {}) + environment_config["image"] = _swebench_image_name(instance, params["subset"]) + environment_config["step_timeout"] = params["step_timeout"] + environment_config["eval_timeout"] = params["eval_timeout"] + environment_config["instance_id"] = instance_id + environment_config["environment_class"] = ( + "responses_api_agents.mini_swe_agent_2.sandbox_environment.MiniSWESandboxEnvironment" + ) + + agent_config = config.get("agent", {}) + agent_config["step_limit"] = params["step_limit"] + agent_config.pop("collapse_limit", None) + + run_id = f"{int(time.time())}_{uuid4()}" + trajectory_path = instance_dir / f"{instance_id}_{run_id}.traj.json" + agent_config["output_path"] = trajectory_path + env = None + agent = None + try: + print(f"[EVAL]{instance_id} Creating environment...", flush=True) + env = get_environment(environment_config) + print(f"[EVAL]{instance_id} Environment created", flush=True) + + model = get_model(config=model_config) + model = _ObservedModel(model, model_name=model_config["model_name"]) + agent = DefaultAgent(model, env, **agent_config) + + if params["run_golden"]: + exit_status = "Gold Patch Applied" + model_patch = instance.get("patch", "") + data = agent.save(None, {"messages": []}) + else: + print(f"[EVAL]{instance_id} Running mini-swe-agent v2...", flush=True) + info = agent.run(instance["problem_statement"]) + exit_status = info.get("exit_status", "") + model_patch = info.get("submission", "") + data = agent.save( + trajectory_path, + {"instance_id": instance_id}, + ) + + print(f"[EVAL]{instance_id} Running eval", flush=True) + eval_report = _run_eval_v2( + instance=instance, + env=env, + model_patch=model_patch, + instance_dir=instance_dir, + run_id=run_id, + is_golden=params["run_golden"], + ) + print(f"[EVAL]{instance_id} Eval completed", flush=True) + + input_messages, response_output, responses = _split_trajectory_for_responses(data.get("messages", [])) + + return { + instance_id: { + "input_messages": input_messages, + "response_output": response_output, + "responses": responses, + "eval_report": eval_report, + "exit_status": exit_status, + } + } + finally: + if env and hasattr(env, "cleanup"): + env.cleanup() + + +def run_swegym_with_optional_sandbox(**params: Any) -> Any: + instance_id = str(params.get("instance_id") or "unknown") + observability_config = params.pop("observability", None) + recorder = build_recorder_from_config( + observability_config, + run_id=observability_config.get("run_id") if isinstance(observability_config, dict) else None, + ) + try: + with use_recorder(recorder): + with event_context( + trajectory_id=observability_config.get("trajectory_id", instance_id) + if isinstance(observability_config, dict) + else instance_id, + instance_id=instance_id, + harness="mini_swe_agent_2", + environment_type="sandbox", + ): + return _run_swegym_v2(**params) + finally: + if recorder is not None: + recorder.finalize() + + +class MiniSWEAgent(SimpleResponsesAPIAgent): + config: MiniSWEAgentConfig + sem: Semaphore = None + model_config = ConfigDict(arbitrary_types_allowed=True) + + def model_post_init(self, __context: Any) -> None: + self.sem = Semaphore(self.config.concurrency) + + def setup_webserver(self) -> FastAPI: + app = FastAPI() + app.post("/v1/responses")(self.responses) + app.post("/run")(self.run) + return app + + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: + raise NotImplementedError + + async def run(self, body: MiniSWEAgentRunRequest) -> MiniSWEAgentVerifyResponse: + async with self.sem: + model_server_name = self.config.model_server.name + global_config_dict = ServerClient.load_from_global_config().global_config_dict + + model_server_config = get_first_server_config_dict( + global_config_dict, + model_server_name, + ) + + policy_model_name = global_config_dict["policy_model_name"] + + ##### MINI-SWE-AGENT CONFIG ##### + subset = body.subset + split = body.split + workers = 1 + run_golden = self.config.run_golden + base_url = f"http://{model_server_config['host']}:{model_server_config['port']}/v1" + dummy_key = "dummy_key" + model_name = f"hosted_vllm/{policy_model_name}" + step_timeout = self.config.step_timeout + eval_timeout = self.config.eval_timeout + step_limit = self.config.step_limit + + instance_id = body.instance_id + extra_fields = getattr(body, "__pydantic_extra__", {}) or {} + + mini_swe_config_path = _swebench_config_path() + config = yaml.safe_load(get_config_path(mini_swe_config_path).read_text()) + responses_create_params_dict = body.responses_create_params.model_dump(exclude_none=True) + + default_model_kwargs = config["model"]["model_kwargs"] + temperature = ( + body.responses_create_params.temperature + if body.responses_create_params.temperature is not None + else default_model_kwargs["temperature"] + ) + top_p = ( + body.responses_create_params.top_p + if body.responses_create_params.top_p is not None + else default_model_kwargs["top_p"] + ) + model_kwargs = _responses_create_params_to_model_kwargs( + responses_create_params_dict, + default_tool_choice=self.config.tool_choice, + ) + if model_kwargs: + config.setdefault("model", {}).setdefault("model_kwargs", {}).update(model_kwargs) + + output_file_dir = f"{Path.cwd()}/results/{subset}/{policy_model_name}" + config_path = mini_swe_config_path + should_write_config = bool(model_kwargs) + if self.config.sandbox_provider is None: + raise ValueError("mini_swe_agent_2 requires sandbox_provider") + config.setdefault("environment", {}).update(self.config.sandbox_environment_kwargs or {}) + config["environment"]["provider"] = self.config.sandbox_provider + config["environment"]["spec"] = _sandbox_spec_for_instance( + self.config.sandbox_spec, + resource_profiles=self.config.sandbox_resource_profiles, + instance_id=instance_id, + ) + should_write_config = True + + if should_write_config: + config_output_dir = Path(output_file_dir) / "_configs" + config_output_dir.mkdir(parents=True, exist_ok=True) + config_path = config_output_dir / f"{instance_id}.sandbox.yaml" + config_path.write_text(yaml.safe_dump(config, sort_keys=False)) + + if self.config.skip_if_exists: + if Path(f"{output_file_dir}/{instance_id}/{instance_id}.json").exists(): + with open(f"{output_file_dir}/{instance_id}/{instance_id}.json", "r") as f: + print(f"Skipping {instance_id} because it already exists") + verify_response = MiniSWEAgentVerifyResponse.model_validate_json(f.read()) + return verify_response + + #### RUN MINI-SWE-AGENT ##### + try: + params = dict( + subset=subset, + split=split, + workers=workers, + output=output_file_dir, + model=model_name, + api_key=dummy_key, + base_url=base_url, + env="sandbox", + run_golden=run_golden, + instance_id=instance_id, + config=config_path, + # TODO: add this later + instance_dict=body.model_dump(), + responses_create_params=json.dumps(responses_create_params_dict), + step_timeout=step_timeout, + eval_timeout=eval_timeout, + step_limit=step_limit, + observability=_observability_config_for_instance( + self.config.observability, + instance_id=instance_id, + task_index=extra_fields.get("_ng_task_index"), + rollout_index=extra_fields.get("_ng_rollout_index"), + ), + ) + future = runner_ray_remote.remote(run_swegym_with_optional_sandbox, params) + result = await asyncio.to_thread(ray.get, future) + result = result[instance_id] + input_messages = result["input_messages"] + response_output = result["response_output"] + responses = result["responses"] + reward = 1.0 if _is_resolved(instance_id, result["eval_report"]) else 0.0 + + except Exception as e: + error_info = {"error": str(e), "traceback": traceback.format_exc()} + print(f"Error running swegym: {e}\n{error_info['traceback']}", flush=True) + result = {"eval_report": error_info} + input_messages = [] + response_output = [] + responses = [] + reward = 0.0 + + body.responses_create_params.input = input_messages + response = _default_response_object() + if responses: + response.update(dict(responses[-1])) + response.pop("extra", None) + response["model"] = policy_model_name + response["temperature"] = temperature + response["top_p"] = top_p + response["output"] = response_output + + verify_response = MiniSWEAgentVerifyResponse( + responses_create_params=body.responses_create_params, + reward=reward, + response=response, + instance_id=instance_id, + metadata=result.get("eval_report", {}) if result else {}, + ) + + output_path = Path(f"{output_file_dir}/{instance_id}") + output_path.mkdir(parents=True, exist_ok=True) + + with open(f"{output_file_dir}/{instance_id}/{instance_id}.json", "w") as f: + json.dump(verify_response.model_dump(), f) + + return verify_response + + +if __name__ == "__main__": + MiniSWEAgent.run_webserver() diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml new file mode 100644 index 0000000000..1fda2d4073 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_opensandbox.yaml @@ -0,0 +1,89 @@ +mini_swe_agent_2: + responses_api_agents: + mini_swe_agent_2: + entrypoint: app.py + domain: coding + description: Software engineering tasks driven by mini-swe-agent harness on OpenSandbox. + value: Improve agentic software engineering capabilities. + model_server: + type: responses_api_models + name: policy_model + concurrency: 64 + env: sandbox + sandbox_provider: + name: opensandbox + kwargs: + connection: + domain: opensandbox-server.opensandbox-system.svc.cluster.local + api_key: ${oc.env:OPENSANDBOX_API_KEY} + protocol: http + use_server_proxy: true + exec_use_server_proxy: true + request_timeout_s: 300 + create: + request_timeout_s: 1200 + timeout_s: 1200 + skip_health_check: true + retries: 10 + retry_delay_s: 5.0 + retry_max_delay_s: 90.0 + probe: + timeout_s: 60 + deadline_s: 180 + stable_count: 2 + stable_delay_s: 1.0 + operations: + retries: 5 + retry_delay_s: 1.0 + retry_max_delay_s: 45.0 + command_retries: 3 + close_timeout_s: 30 + sandbox_spec: + timeout_s: 18000 + ready_timeout_s: 1200 + resources: + cpu: "1" + memory: 8Gi + ephemeral-storage: 20Gi + platform: + os: linux + arch: amd64 + image_rewrites: + - from: swebench/ + to: mirror.gcr.io/swebench/ + metadata: + benchmark: swebench-verified + harness: mini-swe-agent + sandbox-api: opensandbox-sdk + sandbox_environment_kwargs: + cwd: /testbed + conda_env: testbed + activate_conda: true + user: root + delete: true + run_golden: false + step_timeout: 600 + eval_timeout: 1800 + skip_if_exists: false + step_limit: 250 + observability: + enabled: false + output_dir: results/mini_swe_agent_2_observability/{trajectory_id} + export_traces: true + run_id: mini-swe-agent-2 + job_name: mini-swe-agent-2 + otel: + service_name: mini-swe-agent-2 + resource_attributes: + benchmark: swebench-verified + harness: mini_swe_agent_2 + command_titles: + strip_prefixes: + - "cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed &&" + rules: + - line_starts_with: + - "pytest " + - "python -m pytest " + - "./tests/runtests.py " + search: last + title: "run verifier: {line}" diff --git a/responses_api_agents/mini_swe_agent_2/requirements.txt b/responses_api_agents/mini_swe_agent_2/requirements.txt new file mode 100644 index 0000000000..282c60a335 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/requirements.txt @@ -0,0 +1,3 @@ +-e nemo-gym[dev,sandbox] @ ../../ +mini-swe-agent==2.1.0 +swegym @ git+https://github.com/sdevare-nv/nv-SWE-Bench-Package.git@31e1cb8f0241da1707d00faa633c3d6ce1a8ba3b diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py new file mode 100644 index 0000000000..7d741a8138 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""mini-swe-agent environment adapter backed by the Gym sandbox API.""" + +import os +import shlex +from dataclasses import dataclass, field +from typing import Any + + +try: + from minisweagent.exceptions import Submitted +except ModuleNotFoundError: + + class Submitted(Exception): + """Compatibility shim for local mini-swe-agent versions before v2.""" + + def __init__(self, *messages: dict[str, Any]) -> None: + self.messages = messages + super().__init__() + + +from nemo_gym.sandbox import Sandbox, SandboxSpec, rewrite_image +from nemo_gym.sandbox.config import SandboxProviderConfig +from nemo_gym.sandbox.observability import push_event_context, reset_event_context + + +@dataclass +class MiniSWESandboxEnvironmentConfig: + """Configuration for mini-swe-agent runs inside a sandbox.""" + + image: str + cwd: str = "/workspace" + env: dict[str, str] = field(default_factory=dict) + forward_env: list[str] = field(default_factory=list) + timeout: int = 60 + step_timeout: int = 600 + eval_timeout: int = 1800 + interpreter: list[str] = field(default_factory=lambda: ["bash", "-c"]) + executable: str = "sandbox" + run_args: list[str] = field(default_factory=list) + start_args: list[str] = field(default_factory=list) + container_timeout: str = "2h" + instance_id: str | None = None + provider: SandboxProviderConfig | dict[str, Any] = field(default_factory=dict) + spec: dict[str, Any] = field(default_factory=dict) + conda_env: str | None = None + activate_conda: bool = False + user: str | int | None = "root" + delete: bool = True + + +class MiniSWESandboxEnvironment: + """mini-swe-agent sync environment implemented with ``nemo_gym.sandbox.Sandbox``.""" + + def __init__( + self, + *, + config_class: type = MiniSWESandboxEnvironmentConfig, + **kwargs: Any, + ) -> None: + self.config = config_class(**kwargs) + if not self.config.provider: + raise ValueError("MiniSWESandboxEnvironment requires provider") + + self._handle: Any | None = None + self._closed = False + + spec_config = dict(self.config.spec) + image = spec_config.pop("image", None) or self.config.image + image = rewrite_image(image, spec_config.pop("image_rewrites", [])) + + env = dict(spec_config.pop("env", {})) + for key in self.config.forward_env: + value = os.getenv(key) + if value is not None: + env[key] = value + env.update(self.config.env) + + self._sandbox = Sandbox(self.config.provider) + self._handle = self._sandbox.create( + SandboxSpec( + image=image, + snapshot_id=spec_config.pop("snapshot_id", None), + timeout_s=spec_config.pop("timeout_s", None), + ready_timeout_s=spec_config.pop("ready_timeout_s", None), + env=env, + metadata={ + **spec_config.pop("metadata", {}), + "nemo_gym_agent": "mini_swe_agent_2", + "instance_id": (self.config.instance_id or "unknown")[:63], + }, + resources=spec_config.pop("resources", {}), + entrypoint=spec_config.pop("entrypoint", None), + extensions=spec_config.pop("extensions", {}), + platform=spec_config.pop("platform", None), + volumes=spec_config.pop("volumes", None), + skip_health_check=spec_config.pop("skip_health_check", None), + ) + ) + + def get_template_vars(self, **kwargs: Any) -> dict[str, Any]: + return {**self.config.__dict__, **kwargs} + + def serialize(self) -> dict[str, Any]: + return { + "info": { + "config": { + "environment": self.config.__dict__, + "environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}", + } + } + } + + def _command(self, command: str, cwd: str) -> str: + if not self.config.activate_conda or not self.config.conda_env: + return command + quoted_cwd = shlex.quote(cwd) + quoted_env = shlex.quote(self.config.conda_env) + return ( + f"cd {quoted_cwd} && " + "source $(conda info --base)/etc/profile.d/conda.sh && " + f"conda activate {quoted_env} && " + f"{command}" + ) + + def execute( + self, + action: dict[str, Any] | str, + cwd: str = "", + is_eval: bool = False, + timeout: int | None = None, + ) -> dict[str, Any]: + command = action.get("command", "") if isinstance(action, dict) else action + timeout_s = timeout or (self.config.eval_timeout if is_eval else self.config.step_timeout) + exec_cwd = cwd or self.config.cwd + + context_token = None + if is_eval: + context_token = push_event_context( + { + "execution.section": "verifier", + "execution.kind": "verifier", + "span.section": "verifier", + } + ) + try: + result = self._sandbox.exec( + self._handle, + self._command(command, exec_cwd), + cwd="/", + timeout_s=timeout_s, + user=self.config.user, + ) + finally: + if context_token is not None: + reset_event_context(context_token) + output = "\n".join(part for part in (result.stdout, result.stderr) if part) + response = { + "output": output, + "returncode": result.return_code, + "exception_info": "", + } + self._check_finished(response) + return response + + def _check_finished(self, output: dict[str, Any]) -> None: + """Match mini-swe-agent's submit sentinel handling for sandbox-backed runs.""" + lines = output.get("output", "").lstrip().splitlines(keepends=True) + if lines and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0: + submission = "".join(lines[1:]) + raise Submitted( + { + "role": "exit", + "content": submission, + "extra": {"exit_status": "Submitted", "submission": submission}, + } + ) + + def cleanup(self) -> None: + if self._closed: + return + self._closed = True + try: + if self._handle is not None: + self._sandbox.close(self._handle, delete=self.config.delete) + self._handle = None + finally: + self._sandbox.shutdown() + + def __enter__(self) -> "MiniSWESandboxEnvironment": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.cleanup() + + def __del__(self) -> None: + if hasattr(self, "_closed") and not self._closed: + try: + self.cleanup() + except Exception: + pass diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_app.py b/responses_api_agents/mini_swe_agent_2/tests/test_app.py new file mode 100644 index 0000000000..722150a293 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/tests/test_app.py @@ -0,0 +1,943 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any, Dict, Optional +from unittest.mock import MagicMock, patch + +import pytest +import yaml +from fastapi.testclient import TestClient + +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import ( + NeMoGymChatCompletionCreateParamsNonStreaming, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.sandbox.observability import SandboxRecorder, use_recorder +from nemo_gym.server_utils import ServerClient + + +try: + __import__("minisweagent.config") +except ModuleNotFoundError as exc: + if exc.name not in {"minisweagent", "minisweagent.config"}: + raise + minisweagent_module = ModuleType("minisweagent") + minisweagent_module.__path__ = [] + minisweagent_config_module = ModuleType("minisweagent.config") + minisweagent_config_module.builtin_config_dir = Path("/tmp/minisweagent/config") + minisweagent_config_module.get_config_path = Path + sys.modules["minisweagent"] = minisweagent_module + sys.modules["minisweagent.config"] = minisweagent_config_module + +from responses_api_agents.mini_swe_agent_2 import app as mini_swe_app_module +from responses_api_agents.mini_swe_agent_2.app import ( + MiniSWEAgent, + MiniSWEAgentConfig, + MiniSWEAgentRunRequest, + MiniSWEAgentVerifyResponse, + _format_template, + _is_resolved, + _json_dict_from_metadata, + _message_content_to_text, + _observability_config_for_instance, + _ObservedModel, + _responses_create_params_to_model_kwargs, + _run_swegym_v2, + _sandbox_spec_for_instance, + _split_trajectory_for_responses, + _swebench_config_path, + _swebench_image_name, + run_swegym_with_optional_sandbox, +) + + +DEFAULT_RUN_SWEGYM_RESULT = { + "test_instance_123": { + "input_messages": [ + {"type": "message", "role": "system", "content": "You are a helpful assistant."}, + {"type": "message", "role": "user", "content": "Fix this bug."}, + ], + "response_output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "I'll help you fix the bug.", "annotations": []}], + } + ], + "responses": [ + { + "id": "resp-1", + "object": "response", + "output": [], + } + ], + "eval_report": { + "eval_report": { + "test_instance_123": { + "resolved": True, + "tests_status": { + "FAIL_TO_PASS": {"success": ["test1"], "failure": []}, + "PASS_TO_PASS": {"success": ["test2"], "failure": []}, + }, + } + } + }, + } +} + +DEFAULT_CONFIG_YAML = """ +model: + model_kwargs: + temperature: 0.5 + top_p: 0.8 +""" + +DEFAULT_CHAT_COMPLETION = { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "test_model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}, +} + + +def create_test_config( + host: str = "0.0.0.0", + port: int = 8080, + model_name: str = "test_model", +) -> MiniSWEAgentConfig: + return MiniSWEAgentConfig( + name="mini_swe_agent_2", + host=host, + port=port, + entrypoint="", + model_server=ModelServerRef( + type="responses_api_models", + name=model_name, + ), + env="sandbox", + concurrency=1, + sandbox_provider={"name": "opensandbox", "kwargs": {}}, + sandbox_spec={}, + ) + + +def setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict): + mock_server_client_instance = MagicMock() + mock_server_client_instance.global_config_dict = {"policy_model_name": "test_model"} + mock_load_from_global_config.return_value = mock_server_client_instance + + mock_get_first_server_config_dict.return_value = { + "host": "0.0.0.0", + "port": 8080, + } + + +def setup_config_path_mock(mock_get_config_path, config_yaml: str = DEFAULT_CONFIG_YAML): + mock_config_path = MagicMock() + mock_config_path.read_text.return_value = config_yaml + mock_get_config_path.return_value = mock_config_path + + +def setup_run_swegym_mock( + mock_to_thread, + mock_runner_ray_remote, + run_swegym_result: Dict[str, Any] = None, +): + """Setup mock for Ray-based run_swegym execution""" + if run_swegym_result is None: + run_swegym_result = DEFAULT_RUN_SWEGYM_RESULT + + # Mock the Ray remote function to return a future-like object + mock_future = MagicMock() + mock_runner_ray_remote.remote.return_value = mock_future + + # Mock asyncio.to_thread (which calls ray.get) to return the result + mock_to_thread.return_value = run_swegym_result + + +def create_run_request( + instance_id: str = "test_instance_123", + temperature: float = 0.5, + top_p: float = 0.8, + max_output_tokens: int | None = None, + metadata: dict[str, Any] | None = None, + subset: str = "gym", + split: str = "train", + input_data: list = None, +) -> MiniSWEAgentRunRequest: + """Create a test run request with default values.""" + if input_data is None: + input_data = [] + + return MiniSWEAgentRunRequest( + instance_id=instance_id, + subset=subset, + split=split, + responses_create_params=NeMoGymResponseCreateParamsNonStreaming( + temperature=temperature, + top_p=top_p, + max_output_tokens=max_output_tokens, + metadata=metadata, + input=input_data, + ), + ) + + +def create_chat_completion_request( + model: str = "test_model", + messages: list = None, + temperature: float = 0.7, + max_tokens: Optional[int] = None, +) -> NeMoGymChatCompletionCreateParamsNonStreaming: + if messages is None: + messages = [{"role": "user", "content": "Hello!"}] + + kwargs = {"model": model, "messages": messages, "temperature": temperature} + if max_tokens is not None: + kwargs["max_tokens"] = max_tokens + + return NeMoGymChatCompletionCreateParamsNonStreaming(**kwargs) + + +def assert_run_response( + response: MiniSWEAgentVerifyResponse, + expected_reward: float = 1.0, + expected_temperature: float = 0.5, + expected_top_p: float = 0.8, + expected_input_length: int = 2, +): + assert isinstance(response, MiniSWEAgentVerifyResponse) + assert response.reward == expected_reward + assert response.responses_create_params.temperature == expected_temperature + assert response.responses_create_params.top_p == expected_top_p + assert len(response.responses_create_params.input) == expected_input_length + + if expected_input_length >= 2: + assert response.responses_create_params.input[0]["role"] == "system" + assert response.responses_create_params.input[1]["role"] == "user" + + +def assert_run_swegym_called( + mock_to_thread, + subset: str = "gym", + split: str = "train", + instance_id: str = "test_instance_123", +): + mock_to_thread.assert_called_once() + call_args = mock_to_thread.call_args + args = call_args[0] + assert len(args) >= 1 + + +def _otel_attributes(rows: list[dict[str, Any]]) -> dict[str, Any]: + attrs = {} + for row in rows: + value = row["value"] + if "stringValue" in value: + attrs[row["key"]] = value["stringValue"] + elif "boolValue" in value: + attrs[row["key"]] = value["boolValue"] + elif "intValue" in value: + attrs[row["key"]] = int(value["intValue"]) + elif "doubleValue" in value: + attrs[row["key"]] = value["doubleValue"] + return attrs + + +def _otel_spans(output_dir: Path) -> list[dict[str, Any]]: + trace_payload = json.loads((output_dir / "traces" / "otel_traces.json").read_text()) + return [ + span + for resource_span in trace_payload["resourceSpans"] + for scope_span in resource_span["scopeSpans"] + for span in scope_span["spans"] + ] + + +class TestApp: + def test_sanity(self) -> None: + config = create_test_config(model_name="") + MiniSWEAgent(config=config, server_client=MagicMock(spec=ServerClient)) + + def test_observed_model_records_llm_span(self, tmp_path: Path) -> None: + class QueryModel: + def __init__(self) -> None: + self.config = SimpleNamespace(model_kwargs={"temperature": 0.7, "top_p": 0.95, "max_tokens": 128}) + + def query(self, messages: list[dict[str, Any]], **kwargs: Any) -> dict[str, Any]: + return {"role": "assistant", "content": "ok", "extra": {"actions": []}} + + query_model = QueryModel() + observed_model = _ObservedModel(query_model, model_name="hosted_vllm/qwen") + assert observed_model.config is query_model.config + + recorder = SandboxRecorder(output_dir=tmp_path / "observability", otel={"enabled": False}) + with use_recorder(recorder): + with mini_swe_app_module.event_context(trajectory_id="task-1", instance_id="task-1"): + observed_model.query([{"role": "user", "content": "hi"}]) + recorder.finalize() + + spans = _otel_spans(recorder.output_dir) + llm_span = next(span for span in spans if span["name"] == "llm.request") + attrs = _otel_attributes(llm_span["attributes"]) + + assert attrs["operation.name"] == "llm.request" + assert attrs["span.section"] == "rollout" + assert attrs["model"] == "hosted_vllm/qwen" + assert attrs["message_count"] == 1 + assert attrs["trajectory_id"] == "task-1" + + def test_response_param_helpers_cover_metadata_and_tool_choice_modes(self) -> None: + assert _json_dict_from_metadata(None, field_name="extra_body") == {} + assert _json_dict_from_metadata({"top_k": 20}, field_name="extra_body") == {"top_k": 20} + + kwargs = _responses_create_params_to_model_kwargs( + { + "temperature": 0.6, + "top_p": 0.95, + "max_output_tokens": 123, + "metadata": { + "extra_body": json.dumps({"top_k": 20}), + "chat_template_kwargs": json.dumps({"enable_thinking": True}), + }, + "tool_choice": {"type": "function", "function": {"name": "python"}}, + } + ) + + assert kwargs == { + "temperature": 0.6, + "top_p": 0.95, + "max_tokens": 123, + "extra_body": {"top_k": 20, "chat_template_kwargs": {"enable_thinking": True}}, + "tool_choice": {"type": "function", "function": {"name": "python"}}, + } + assert _responses_create_params_to_model_kwargs({"tool_choice": "bash"})["tool_choice"] == { + "type": "function", + "function": {"name": "bash"}, + } + assert ( + _responses_create_params_to_model_kwargs({"tool_choice": "auto"}, default_tool_choice="none")[ + "tool_choice" + ] + == "none" + ) + + with pytest.raises(ValueError, match="extra_body"): + _json_dict_from_metadata("[]", field_name="extra_body") + + def test_sandbox_resource_profiles_override_static_resources(self) -> None: + spec = _sandbox_spec_for_instance( + {"resources": {"cpu": "1", "memory": "8Gi", "ephemeral-storage": "20Gi"}}, + resource_profiles=[ + {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, + {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + ], + instance_id="django__django-12345", + ) + + assert spec["resources"] in ( + {"cpu": "250m", "memory": "3Gi", "ephemeral-storage": "1Gi"}, + {"cpu": "500m", "memory": "4Gi", "ephemeral-storage": "1Gi"}, + ) + assert _sandbox_spec_for_instance(None, resource_profiles=None, instance_id="task") == {} + + def test_observability_config_formats_per_rollout_context(self) -> None: + assert _format_template(1, {"trajectory_id": "task"}) == 1 + assert _format_template("{missing}", {"trajectory_id": "task"}) == "{missing}" + config = _observability_config_for_instance( + { + "enabled": True, + "output_dir": "results/observability/{trajectory_id}", + "run_id": "run-{task_index}", + "job_name": "job-{trajectory_id}", + "otel": { + "service_name": "svc-{instance_id}", + "resource_attributes": {"rollout": "{rollout_index}"}, + "command_titles": {"rules": [{"title": "run verifier: {line}"}]}, + }, + }, + instance_id="django__django-123", + task_index=4, + rollout_index=2, + ) + + assert config == { + "enabled": True, + "trajectory_id": "django__django-123__rollout03", + "output_dir": "results/observability/django__django-123__rollout03", + "run_id": "run-4", + "job_name": "job-django__django-123__rollout03", + "otel": { + "service_name": "svc-django__django-123", + "resource_attributes": {"rollout": "2"}, + "command_titles": {"rules": [{"title": "run verifier: {line}"}]}, + }, + } + assert _observability_config_for_instance(None, instance_id="task") is None + assert _observability_config_for_instance( + {"enabled": True}, + instance_id="task", + rollout_index="retry", + ) == {"enabled": True, "trajectory_id": "task__rolloutretry"} + + def test_split_trajectory_and_resolution_helpers_cover_edge_cases(self) -> None: + input_messages, output_items, raw_responses = _split_trajectory_for_responses( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "user"}, + { + "role": "assistant", + "content": "answer", + "tool_calls": [{"id": "call-1", "function": {"name": "bash", "arguments": '{"command":"pwd"}'}}], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "tool output"}, + {"type": "function_call_output", "call_id": "call-2", "output": "raw", "extra": {"ignored": True}}, + {"object": "response", "output": [{"type": "message", "content": "raw"}], "extra": {"ignored": True}}, + ] + ) + + assert input_messages == [ + {"type": "message", "role": "system", "content": "sys"}, + {"type": "message", "role": "user", "content": "user"}, + ] + assert any(item["type"] == "function_call" and item["call_id"] == "call-1" for item in output_items) + assert any(item["type"] == "function_call_output" and item["call_id"] == "call-1" for item in output_items) + assert any(item["type"] == "function_call_output" and item["call_id"] == "call-2" for item in output_items) + assert raw_responses == [{"object": "response", "output": [{"type": "message", "content": "raw"}]}] + + assert not _is_resolved("task", {}) + assert not _is_resolved("task", {"eval_report": {"task": {"resolved": True}}}) + assert not _is_resolved( + "task", + { + "eval_report": { + "task": { + "resolved": True, + "tests_status": {"FAIL_TO_PASS": {"success": [], "failure": []}}, + } + } + }, + ) + + def test_misc_mini_swe_helpers(self, monkeypatch, tmp_path) -> None: + assert _swebench_image_name({"instance_id": "django__django-1"}, "verified") == ( + "swebench/sweb.eval.x86_64.django_1776_django-1:latest" + ) + assert _swebench_image_name({"instance_id": "django__django-1"}, "lite") == ( + "xingyaoww/sweb.eval.x86_64.django_s_django-1:latest" + ) + assert _swebench_image_name({"instance_id": "x", "image_name": "custom:image"}, "verified") == "custom:image" + assert _message_content_to_text("hello") == "hello" + assert _message_content_to_text(None) == "" + assert _message_content_to_text([{"text": "one"}, {"content": "two"}, 3]) == "one\ntwo\n3" + + builtin_dir = tmp_path / "configs" + benchmark_dir = builtin_dir / "benchmarks" + benchmark_dir.mkdir(parents=True) + (benchmark_dir / "swebench.yaml").write_text("{}", encoding="utf-8") + monkeypatch.setattr(mini_swe_app_module, "builtin_config_dir", builtin_dir) + assert _swebench_config_path() == benchmark_dir / "swebench.yaml" + monkeypatch.setattr(mini_swe_app_module, "builtin_config_dir", tmp_path / "missing") + assert _swebench_config_path() == tmp_path / "missing" / "extra" / "swebench.yaml" + + def test_run_swegym_records_completion_and_errors(self, monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + mini_swe_app_module, + "_run_swegym_v2", + lambda **_params: { + "task-1": { + "eval_report": { + "task-1": {"resolved": True}, + } + } + }, + ) + assert run_swegym_with_optional_sandbox( + env="sandbox", + instance_id="task-1", + observability={ + "enabled": True, + "output_dir": str(tmp_path / "observability"), + "export_traces": False, + "run_id": "run-1", + }, + ) == {"task-1": {"eval_report": {"task-1": {"resolved": True}}}} + + def fail_runner(**_params): + raise RuntimeError("boom") + + monkeypatch.setattr(mini_swe_app_module, "_run_swegym_v2", fail_runner) + with pytest.raises(RuntimeError, match="boom"): + run_swegym_with_optional_sandbox(env="sandbox", instance_id="task-1") + + monkeypatch.setattr( + mini_swe_app_module, + "_run_swegym_v2", + lambda **_params: {"task-1": {"eval_report": {"task-1": {"resolved": False}}}}, + ) + assert run_swegym_with_optional_sandbox(env="sandbox", instance_id="task-1") == { + "task-1": {"eval_report": {"task-1": {"resolved": False}}} + } + + monkeypatch.setattr(mini_swe_app_module, "_run_swegym_v2", lambda **_params: {"task-1": "bad"}) + assert run_swegym_with_optional_sandbox(env="sandbox", instance_id="task-1") == {"task-1": "bad"} + + def test_run_swegym_v2_success_and_golden_paths(self, monkeypatch, tmp_path) -> None: + holder: dict[str, Any] = {} + + class FakeLogger: + def info(self, _message: str) -> None: + return None + + def setup_logger(_instance_id: str, _log_file: Path) -> FakeLogger: + return FakeLogger() + + def make_test_spec(instance: dict[str, Any]) -> SimpleNamespace: + return SimpleNamespace( + instance_id=instance["instance_id"], + eval_script="#!/bin/bash\npytest -q", + ) + + def get_eval_report(*, test_spec: SimpleNamespace, prediction: dict[str, Any], log_path: Path, **_kwargs: Any): + assert log_path.exists() + return {test_spec.instance_id: {"resolved": True, "prediction": prediction}} + + class FakeEnv: + def __init__(self, config: dict[str, Any]) -> None: + self.config = config + self.commands: list[tuple[str, bool]] = [] + self.cleaned = False + + def execute(self, command: str, is_eval: bool = False) -> dict[str, Any]: + self.commands.append((command, is_eval)) + return {"output": "tests passed", "returncode": 0} + + def cleanup(self) -> None: + self.cleaned = True + + class FakeAgent: + def __init__(self, model: Any, env: FakeEnv, **agent_config: Any) -> None: + self.model = model + self.env = env + self.agent_config = agent_config + holder["agent_config"] = agent_config + + def run(self, problem_statement: str) -> dict[str, Any]: + assert problem_statement == "Fix the bug" + return {"exit_status": "submitted", "submission": "diff --git a/file b/file"} + + def save(self, path: Path | None, metadata: dict[str, Any]) -> dict[str, Any]: + holder["save_path"] = path + holder["save_metadata"] = metadata + if path is not None: + path.write_text("{}", encoding="utf-8") + return { + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": [{"text": "problem"}]}, + { + "id": "resp-1", + "object": "response", + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + ], + "extra": {"actions": [{"command": "echo hi", "tool_call_id": "call-1"}]}, + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": "tool output", + "extra": {"raw_output": "tool output"}, + }, + ] + } + + def get_environment(config: dict[str, Any]) -> FakeEnv: + env = FakeEnv(config) + holder["env"] = env + return env + + def get_model(config: dict[str, Any]) -> SimpleNamespace: + holder["model_config"] = config + return SimpleNamespace(config=config) + + module_specs = { + "swegym": ModuleType("swegym"), + "swegym.harness": ModuleType("swegym.harness"), + "swegym.harness.constants": ModuleType("swegym.harness.constants"), + "swegym.harness.docker_build": ModuleType("swegym.harness.docker_build"), + "swegym.harness.grading": ModuleType("swegym.harness.grading"), + "swegym.harness.test_spec": ModuleType("swegym.harness.test_spec"), + "minisweagent.agents": ModuleType("minisweagent.agents"), + "minisweagent.agents.default": ModuleType("minisweagent.agents.default"), + "minisweagent.environments": ModuleType("minisweagent.environments"), + "minisweagent.models": ModuleType("minisweagent.models"), + } + module_specs["swegym.harness.constants"].SWEbenchInstance = dict + module_specs["swegym.harness.docker_build"].setup_logger = setup_logger + module_specs["swegym.harness.grading"].get_eval_report = get_eval_report + module_specs["swegym.harness.test_spec"].make_test_spec = make_test_spec + module_specs["minisweagent.agents.default"].DefaultAgent = FakeAgent + module_specs["minisweagent.environments"].get_environment = get_environment + module_specs["minisweagent.models"].get_model = get_model + for name, module in module_specs.items(): + monkeypatch.setitem(sys.modules, name, module) + + config_path = tmp_path / "swebench.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": {"model_kwargs": {"max_output_tokens": 99}}, + "environment": {}, + "agent": {"step_limit": 1, "collapse_limit": 3}, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(mini_swe_app_module, "get_config_path", lambda _config: config_path) + monkeypatch.setattr(mini_swe_app_module, "uuid4", lambda: "uuid") + monkeypatch.setattr(mini_swe_app_module.time, "time", lambda: 1234) + + params = { + "instance_dict": { + "instance_id": "django__django-123", + "problem_statement": "Fix the bug", + "patch": "gold", + }, + "instance_id": "django__django-123", + "output": str(tmp_path / "out"), + "config": "swebench", + "model": "hosted/model", + "api_key": "key", # pragma: allowlist secret + "base_url": "http://model/v1", + "subset": "verified", + "step_timeout": 30, + "eval_timeout": 60, + "env": "sandbox", + "step_limit": 7, + "run_golden": False, + } + + result = _run_swegym_v2(**params) + + env = holder["env"] + assert env.cleaned is True + assert env.config["environment_class"].endswith("MiniSWESandboxEnvironment") + assert env.config["image"] == "swebench/sweb.eval.x86_64.django_1776_django-123:latest" + assert holder["model_config"]["model_class"] == "litellm" + assert holder["model_config"]["model_name"] == "hosted/model" + assert holder["model_config"]["model_kwargs"]["max_tokens"] == 99 + assert holder["model_config"]["model_kwargs"]["base_url"] == "http://model/v1" + assert "api_base" not in holder["model_config"]["model_kwargs"] + assert holder["agent_config"]["step_limit"] == 7 + assert holder["save_metadata"] == {"instance_id": "django__django-123"} + assert result["django__django-123"]["input_messages"] == [ + {"type": "message", "role": "system", "content": "sys"}, + {"type": "message", "role": "user", "content": "problem"}, + ] + assert result["django__django-123"]["response_output"] == [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + {"type": "function_call_output", "call_id": "call-1", "output": "tool output"}, + ] + assert result["django__django-123"]["responses"] == [ + { + "id": "resp-1", + "object": "response", + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + { + "type": "function_call", + "name": "bash", + "call_id": "call-1", + "arguments": json.dumps({"command": "echo hi"}), + }, + ], + } + ] + + golden_params = params | {"run_golden": True} + result = _run_swegym_v2(**golden_params) + + env = holder["env"] + assert env.cleaned is True + assert env.config["environment_class"].endswith("MiniSWESandboxEnvironment") + assert [command for command, _ in env.commands[:4]] == [ + "cat > patch.diff <<'EOF'\ngold\n\nEOF", + "git status --porcelain", + "git apply --check patch.diff", + "git apply patch.diff", + ] + assert result["django__django-123"]["exit_status"] == "Gold Patch Applied" + + string_params = params | { + "instance_dict": json.dumps( + {"instance_id": "django__django-123", "problem_statement": "Fix the bug", "patch": "gold"} + ), + } + assert "django__django-123" in _run_swegym_v2(**string_params) + + with pytest.raises(ValueError, match="instance_dict"): + _run_swegym_v2(**(params | {"instance_dict": None})) + + @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") + @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") + @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") + @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") + @patch("asyncio.to_thread") + async def test_run_successful_execution( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + ) -> None: + """Test successful execution of the run method with mocked run_swegym.""" + + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) + setup_config_path_mock(mock_get_config_path) + setup_run_swegym_mock(mock_to_thread, mock_runner_ray_remote) + + run_request = create_run_request() + + response = await server.run(run_request) + + assert_run_response(response) + + assert_run_swegym_called(mock_to_thread) + + @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") + @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") + @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") + @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") + @patch("asyncio.to_thread") + async def test_run_writes_generation_params_to_config( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + tmp_path, + monkeypatch, + ) -> None: + monkeypatch.chdir(tmp_path) + config = create_test_config() + config.tool_choice = "bash" + config.observability = { + "enabled": True, + "output_dir": "results/observability/{trajectory_id}", + "run_id": "mini-swe-run", + } + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) + setup_config_path_mock(mock_get_config_path) + setup_run_swegym_mock(mock_to_thread, mock_runner_ray_remote) + + run_request = create_run_request( + temperature=0.6, + top_p=0.95, + max_output_tokens=49152, + metadata={ + "extra_body": '{"top_k":20,"min_p":0.0,"presence_penalty":0.0,"repetition_penalty":1.0}', + "chat_template_kwargs": '{"enable_thinking":true}', + }, + ) + + await server.run(run_request) + + call_args = mock_runner_ray_remote.remote.call_args + params = call_args.args[1] + generated_config = yaml.safe_load(Path(params["config"]).read_text()) + model_kwargs = generated_config["model"]["model_kwargs"] + assert model_kwargs["temperature"] == 0.6 + assert model_kwargs["top_p"] == 0.95 + assert model_kwargs["max_tokens"] == 49152 + assert "max_output_tokens" not in model_kwargs + assert model_kwargs["tool_choice"] == {"type": "function", "function": {"name": "bash"}} + assert model_kwargs["extra_body"] == { + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, + "chat_template_kwargs": {"enable_thinking": True}, + } + assert params["observability"] == { + "enabled": True, + "trajectory_id": "test_instance_123", + "output_dir": "results/observability/test_instance_123", + "run_id": "mini-swe-run", + } + + @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") + @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") + @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") + @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") + @patch("asyncio.to_thread") + async def test_run_failed_execution( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + ) -> None: + """Test run method when run_swegym fails.""" + + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) + setup_config_path_mock(mock_get_config_path) + + # Mock Ray remote function + mock_future = MagicMock() + mock_runner_ray_remote.remote.return_value = mock_future + + # Mock asyncio.to_thread (ray.get) to raise an exception + mock_to_thread.side_effect = Exception("run_swegym failed") + + run_request = create_run_request(instance_id="test_instance_456", temperature=0.3, top_p=0.95) + + response = await server.run(run_request) + + assert_run_response( + response, + expected_reward=0.0, + expected_temperature=0.3, + expected_top_p=0.95, + expected_input_length=0, + ) + + assert_run_swegym_called(mock_to_thread, instance_id="test_instance_456") + + @patch("responses_api_agents.mini_swe_agent_2.app.ServerClient.load_from_global_config") + @patch("responses_api_agents.mini_swe_agent_2.app.get_first_server_config_dict") + @patch("responses_api_agents.mini_swe_agent_2.app.get_config_path") + @patch("responses_api_agents.mini_swe_agent_2.app.runner_ray_remote") + @patch("asyncio.to_thread") + async def test_run_swegym_not_found( + self, + mock_to_thread, + mock_runner_ray_remote, + mock_get_config_path, + mock_get_first_server_config_dict, + mock_load_from_global_config, + ) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + setup_server_client_mocks(mock_load_from_global_config, mock_get_first_server_config_dict) + setup_config_path_mock(mock_get_config_path) + + # Mock Ray remote function + mock_future = MagicMock() + mock_runner_ray_remote.remote.return_value = mock_future + + # Mock asyncio.to_thread (ray.get) to raise FileNotFoundError + mock_to_thread.side_effect = FileNotFoundError("run_swegym not found") + + run_request = create_run_request(instance_id="test_instance_789", temperature=0.2, top_p=1.0) + + response = await server.run(run_request) + + assert_run_response( + response, + expected_reward=0.0, + expected_temperature=0.2, + expected_top_p=1.0, + expected_input_length=0, + ) + + assert_run_swegym_called(mock_to_thread, instance_id="test_instance_789") + + async def test_responses_not_implemented(self) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + request_body = NeMoGymResponseCreateParamsNonStreaming(temperature=0.7, top_p=0.9, input=[]) + + with pytest.raises(NotImplementedError): + await server.responses(request_body) + + def test_endpoints_registration(self) -> None: + config = create_test_config() + mock_server_client = MagicMock(spec=ServerClient) + server = MiniSWEAgent(config=config, server_client=mock_server_client) + + app = server.setup_webserver() + client = TestClient(app, raise_server_exceptions=False) + + response = client.post("/v1/responses", json={"temperature": 0.7, "top_p": 0.9, "input": []}) + assert response.status_code == 500 + + run_response = client.post("/run", json={}) + assert run_response.status_code != 404 diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py new file mode 100644 index 0000000000..331d732641 --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment, Submitted + + +def test_check_finished_raises_submitted_for_submit_sentinel() -> None: + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + + try: + env._check_finished( + { + "output": "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\npatch contents\n", + "returncode": 0, + "exception_info": "", + } + ) + except Submitted as error: + assert error.messages == ( + { + "role": "exit", + "content": "patch contents\n", + "extra": {"exit_status": "Submitted", "submission": "patch contents\n"}, + }, + ) + else: + raise AssertionError("Expected Submitted") + + +def test_check_finished_ignores_nonzero_submit_sentinel() -> None: + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + + env._check_finished( + { + "output": "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\npatch contents\n", + "returncode": 1, + "exception_info": "", + } + ) diff --git a/tests/unit_tests/test_opensandbox_provider.py b/tests/unit_tests/test_opensandbox_provider.py new file mode 100644 index 0000000000..ceb306b684 --- /dev/null +++ b/tests/unit_tests/test_opensandbox_provider.py @@ -0,0 +1,621 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ast +import asyncio +import importlib.util +from dataclasses import dataclass +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from nemo_gym.sandbox.providers.base import SandboxSpec +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider + + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("tenacity") is None, + reason="tenacity optional sandbox dependency is not installed", +) + + +@dataclass(frozen=True) +class FakePlatformSpec: + os: str + arch: str + + +class FakeConnectionConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +class FakeSandbox: + created_kwargs: dict[str, Any] = {} + connected_args: tuple[Any, ...] = () + connected_kwargs: dict[str, Any] = {} + + def __init__(self, sandbox_id: str = "sandbox-1") -> None: + self.id = sandbox_id + + @classmethod + async def create(cls, *_args: Any, **kwargs: Any) -> "FakeSandbox": + cls.created_kwargs = kwargs + return cls() + + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + cls.connected_args = args + cls.connected_kwargs = kwargs + return cls() + + +@dataclass +class FakePoolCreationSpec: + image: str + entrypoint: list[str] | None = None + resource: dict[str, str] | None = None + env: dict[str, str] | None = None + metadata: dict[str, str] | None = None + extensions: dict[str, str] | None = None + platform: Any | None = None + volumes: list[Any] | None = None + + +class FakeAcquirePolicy: + FAIL_FAST = "fail_fast" + + +class FakeStateStore: + pass + + +class FakeSnapshot: + idle_count = 1 + state = None + + +class FakeSandboxPoolAsync: + received_kwargs: dict[str, Any] = {} + + def __init__(self, **kwargs: Any) -> None: + self.received_kwargs = kwargs + type(self).received_kwargs = kwargs + + async def start(self) -> None: + return None + + async def snapshot(self) -> FakeSnapshot: + return FakeSnapshot() + + async def resize(self, _count: int) -> None: + return None + + async def acquire( + self, + *, + sandbox_timeout: timedelta | None, + policy: str, + ) -> FakeSandbox: + del sandbox_timeout, policy + creation_spec = self.received_kwargs["creation_spec"] + return await FakeSandbox.create( + creation_spec.image, + platform=creation_spec.platform, + ) + + async def shutdown(self, *, graceful: bool) -> None: + del graceful + + async def release_all_idle(self) -> None: + return None + + +@pytest.fixture +def fake_opensandbox_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + def require_sdk() -> tuple[Any, Any, Any, Any, Any]: + return FakeSandbox, FakeConnectionConfig, object, FakePlatformSpec, object + + def require_sdk_pool() -> tuple[Any, Any, Any, Any]: + return ( + FakeAcquirePolicy, + FakeStateStore, + FakePoolCreationSpec, + FakeSandboxPoolAsync, + ) + + monkeypatch.setattr(opensandbox_provider, "_require_opensandbox_sdk", require_sdk) + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk_pool", + require_sdk_pool, + ) + + +async def test_sdk_pool_passes_platform_through_pool_creation_spec( + fake_opensandbox_sdk: None, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"request_timeout_s": 10}, + probe={"command": None}, + ) + + handles = await provider.create_batch( + SandboxSpec( + image="mirror.gcr.io/astral/uv:python3.12-bookworm-slim", + platform={"os": "linux", "arch": "amd64"}, + ), + 1, + ) + + assert len(handles) == 1 + assert handles[0].sandbox_id == "sandbox-1" + assert "sandbox_factory" not in FakeSandboxPoolAsync.received_kwargs + assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec( + os="linux", + arch="amd64", + ) + + +async def test_connect_passes_configured_connect_timeout( + fake_opensandbox_sdk: None, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"connect_timeout_s": 300, "request_timeout_s": 10}, + probe={"command": None}, + ) + + handle = await provider.connect("sandbox-123") + + assert handle.sandbox_id == "sandbox-1" + assert FakeSandbox.connected_args == ("sandbox-123",) + assert FakeSandbox.connected_kwargs["connect_timeout"] == timedelta(seconds=300) + + +def test_provider_validation_and_retry_helpers() -> None: + with pytest.raises(ValueError, match="image_pull_policy"): + opensandbox_provider.validate_image_pull_policy("Sometimes") + + invalid_kwargs = [ + {"pool": {"concurrency": 0}}, + {"connection": {"connect_timeout_s": 0}}, + {"pool": {"progress_timeout_s": 0}}, + {"create": {"timeout_s": 0}}, + {"probe": {"timeout_s": 0}}, + {"probe": {"deadline_s": 0}}, + {"probe": {"sample_count": 0}}, + {"probe": {"stable_count": 0}}, + {"probe": {"stable_delay_s": -1}}, + {"create": {"retries": -1}}, + {"create": {"retry_delay_s": -1}}, + {"create": {"retry_max_delay_s": -1}}, + {"operations": {"retries": -1}}, + {"operations": {"retry_delay_s": -1}}, + {"operations": {"retry_max_delay_s": -1}}, + {"operations": {"command_retries": -1}}, + {"pool": {"reconcile_interval_s": 0}}, + {"pool": {"acquire_poll_interval_s": 0}}, + {"pool": {"idle_timeout_s": 0}}, + {"pool": {"primary_lock_ttl_s": 0}}, + {"operations": {"close_timeout_s": 0}}, + {"create": {"connect_attempt_timeout_s": 0}}, + {"create": {"connect_poll_s": 0}}, + {"create": {"image_pull_policy": "Sometimes"}}, + ] + for kwargs in invalid_kwargs: + with pytest.raises(ValueError): + opensandbox_provider.OpenSandboxProvider(**kwargs) + with pytest.raises(TypeError): + opensandbox_provider.OpenSandboxProvider(**{"batch_" + "create_retries": 1}) + with pytest.raises(TypeError): + opensandbox_provider.OpenSandboxProvider(connection=object()) + + assert opensandbox_provider._exception_status_code(RuntimeError("HTTP status code: 503")) == 503 + assert opensandbox_provider._exception_status_code(RuntimeError("plain error")) is None + attrs = opensandbox_provider._sdk_error_attributes( + RuntimeError("HTTP 502 bad gateway"), + operation="exec", + sandbox_id="sandbox-1", + attempt_number=2, + max_attempts=3, + sleep_s=0.5, + ) + assert attrs["status_code"] == 502 + assert attrs["attempt_number"] == 2 + assert attrs["next_sleep_s"] == 0.5 + assert opensandbox_provider._seconds_to_timedelta(None) is None + assert opensandbox_provider._seconds_to_timedelta(1.5) == timedelta(seconds=1.5) + + +def test_opensandbox_record_event_names_are_namespaced() -> None: + tree = ast.parse(Path(opensandbox_provider.__file__).read_text(encoding="utf-8")) + event_names: list[str] = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + if not isinstance(node.func, ast.Name) or node.func.id != "record_event": + continue + if len(node.args) < 2 or not isinstance(node.args[1], ast.Constant): + raise AssertionError("OpenSandbox record_event calls must use a literal event name") + event_names.append(node.args[1].value) + + assert event_names + assert all(name.startswith("sandbox.opensandbox.") for name in event_names) + + +def test_connection_config_exec_proxy_and_image_policy(fake_opensandbox_sdk: None) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={ + "domain": "sandbox.example", + "api_key": "key", # pragma: allowlist secret + "protocol": "https", + "use_server_proxy": True, + "exec_use_server_proxy": False, + "request_timeout_s": 10, + } + ) + + config = provider._connection_config() + assert config.kwargs == { + "domain": "sandbox.example", + "api_key": "key", # pragma: allowlist secret + "protocol": "https", + "use_server_proxy": True, + "request_timeout": timedelta(seconds=10), + } + exec_config = provider._exec_connection_config(request_timeout_s=3) + assert exec_config.kwargs["use_server_proxy"] is False + assert exec_config.kwargs["request_timeout"] == timedelta(seconds=3) + + spec = SandboxSpec(image="image:tag", extensions={"imagePullPolicy": "Never"}) + updated = provider._with_default_image_pull_policy(spec) + assert updated.extensions["imagePullPolicy"] == "Never" + assert updated.extensions["opensandbox.extensions.image-pull-policy"] == "Never" + + no_policy_provider = opensandbox_provider.OpenSandboxProvider(create={"image_pull_policy": None}) + assert no_policy_provider._with_default_image_pull_policy(spec) is spec + + +async def test_wait_sdk_pool_idle_success_partial_and_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + class Snapshot: + def __init__(self, idle_count: int) -> None: + self.idle_count = idle_count + self.state = SimpleNamespace(value="warming") + + class FakePool: + def __init__(self, counts: list[int]) -> None: + self.counts = counts + self.index = 0 + self._config = SimpleNamespace(pool_name="pool-1") + + async def snapshot(self) -> Snapshot: + count = self.counts[min(self.index, len(self.counts) - 1)] + self.index += 1 + return Snapshot(count) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + + provider = opensandbox_provider.OpenSandboxProvider( + pool={"acquire_poll_interval_s": 0.01}, + probe={"command": None}, + ) + assert ( + await provider._wait_sdk_pool_idle( + FakePool([0, 1, 2]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=1, + allow_partial=False, + ) + == 2 + ) + assert ( + await provider._wait_sdk_pool_idle( + FakePool([1]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=0, + allow_partial=True, + ) + == 1 + ) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): + await provider._wait_sdk_pool_idle( + FakePool([0]), + spec=SandboxSpec(image="image:tag"), + requested=2, + timeout_s=0, + allow_partial=False, + ) + + +async def test_exec_file_operations_and_batch_validation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeLog: + def __init__(self, text: str) -> None: + self.text = text + + class FakeCommands: + def __init__(self) -> None: + self.calls: list[tuple[str, FakeRunCommandOpts]] = [] + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> Any: + self.calls.append((command, opts)) + if "fail" in command: + return SimpleNamespace( + logs=SimpleNamespace(stdout=[], stderr=[FakeLog("stderr")]), + error=SimpleNamespace(name="CommandError", value="failed"), + exit_code=None, + ) + return SimpleNamespace( + logs=SimpleNamespace(stdout=[FakeLog("stdout")], stderr=[]), + error=None, + exit_code=None, + ) + + class FakeFiles: + def __init__(self) -> None: + self.writes: list[tuple[str, str | bytes]] = [] + + async def write_file(self, target_path: str, data: str | bytes) -> None: + self.writes.append((target_path, data)) + + async def read_bytes(self, source_path: str) -> bytes: + return f"bytes:{source_path}".encode() + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + self.files = FakeFiles() + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = opensandbox_provider.OpenSandboxProvider( + connection={"request_timeout_s": 5}, + probe={"command": None}, + ) + raw = FakeRaw() + handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=raw) + + result = await provider.exec( + handle, + "echo hello", + cwd="/repo", + env={"A": "B"}, + timeout_s=2, + user=1000, + ) + assert result == opensandbox_provider.SandboxExecResult(stdout="stdout", stderr=None, return_code=0) + command, opts = raw.commands.calls[0] + assert command == "echo hello" + assert opts.kwargs == { + "working_directory": "/repo", + "envs": {"A": "B"}, + "timeout": timedelta(seconds=2), + "uid": 1000, + } + + result = await provider.exec(handle, "fail", user="agent") + assert result.return_code == 1 + assert result.stderr == "stderr\nCommandError: failed" + assert raw.commands.calls[1][0] == "su -s /bin/sh -c fail agent" + + await provider.write_file(handle, "/tmp/file.txt", "contents") + assert await provider.read_file(handle, "/tmp/file.txt") == b"bytes:/tmp/file.txt" + upload_path = tmp_path / "upload.txt" + upload_path.write_text("upload", encoding="utf-8") + await provider.upload_file(handle, upload_path, "/remote/upload.txt") + download_path = tmp_path / "nested" / "download.txt" + await provider.download_file(handle, "/remote/download.txt", download_path) + assert raw.files.writes == [("/tmp/file.txt", "contents"), ("/remote/upload.txt", b"upload")] + assert download_path.read_bytes() == b"bytes:/remote/download.txt" + + with pytest.raises(ValueError, match="count"): + await provider._create_batch_sdk(SandboxSpec(image="image:tag"), 0) + with pytest.raises(ValueError, match="count"): + await provider.create_batch(SandboxSpec(image="image:tag"), 0) + with pytest.raises(ValueError, match="snapshot_id"): + provider._validate_sdk_pool_spec(SandboxSpec(image="image:tag", snapshot_id="snapshot")) + with pytest.raises(ValueError, match="Unsupported"): + await provider.materialize_handle({"kind": "other"}) + + +async def test_provider_create_probe_and_close_error_paths(monkeypatch: pytest.MonkeyPatch) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + create={"connect_poll_s": 0.01}, + probe={ + "command": "probe", + "expected_stdout": "ready", + "timeout_s": 1, + "deadline_s": 0.01, + }, + ) + handle = opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=object()) + + async def bad_probe(*_args: Any, **_kwargs: Any) -> opensandbox_provider.SandboxExecResult: + return opensandbox_provider.SandboxExecResult(stdout="not ready", stderr="bad", return_code=1) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + monkeypatch.setattr(provider, "_exec", bad_probe) + with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): + await provider._verify_created_handle(handle) + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": "probe"}) + + async def fail_verify(_handle: Any) -> None: + raise RuntimeError("probe failed") + + monkeypatch.setattr(provider, "_verify_created_handle", fail_verify) + with pytest.raises(opensandbox_provider.OpenSandboxCreateVerificationError): + await provider._verify_created_handles([handle, handle]) + + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + await provider._verify_created_handles([]) + + async def close_raises(_handle: Any, *, delete: bool) -> None: + del delete + raise RuntimeError("close failed") + + monkeypatch.setattr(provider, "close", close_raises) + await provider._cleanup_failed_create_handle(handle) + provider = opensandbox_provider.OpenSandboxProvider(probe={"command": None}) + + class DeleteAlreadyGoneRaw: + async def kill(self) -> None: + raise RuntimeError("sandbox sandbox-1 not found") + + async def close(self) -> None: + return None + + await provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-1", + provider_name="opensandbox", + raw=DeleteAlreadyGoneRaw(), + ), + delete=True, + ) + + class DeleteAndCloseFailRaw: + async def kill(self) -> None: + raise RuntimeError("delete failed") + + async def close(self) -> None: + raise RuntimeError("close failed") + + with pytest.raises(RuntimeError, match="Failed to delete and close"): + await provider.close( + opensandbox_provider.SandboxHandle( + sandbox_id="sandbox-2", + provider_name="opensandbox", + raw=DeleteAndCloseFailRaw(), + ), + delete=True, + ) + + +async def test_create_once_and_connect_after_create_error_paths( + fake_opensandbox_sdk: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + connection={"use_server_proxy": False}, + probe={"command": None}, + ) + with pytest.raises(ValueError, match="pooled creation"): + await provider._create_once(SandboxSpec(image="image:tag", extensions={"poolRef": "pool"})) + + provider = opensandbox_provider.OpenSandboxProvider( + create={"timeout_s": 1, "skip_health_check": True}, + probe={"command": None}, + ) + monkeypatch.setattr(opensandbox_provider, "_to_volumes", lambda volumes: volumes) + spec = SandboxSpec( + image="image:tag", + snapshot_id="snapshot-1", + timeout_s=10, + ready_timeout_s=20, + entrypoint=["/bin/sh"], + platform={"os": "linux", "arch": "amd64"}, + volumes=[{"name": "workspace"}], + skip_health_check=False, + ) + handle = await provider._create_once(spec) + assert handle.sandbox_id == "sandbox-1" + assert FakeSandbox.created_kwargs["snapshot_id"] == "snapshot-1" + assert FakeSandbox.created_kwargs["timeout"] == timedelta(seconds=10) + assert FakeSandbox.created_kwargs["ready_timeout"] == timedelta(seconds=20) + assert FakeSandbox.created_kwargs["entrypoint"] == ["/bin/sh"] + assert FakeSandbox.created_kwargs["platform"] == FakePlatformSpec(os="linux", arch="amd64") + assert FakeSandbox.created_kwargs["volumes"] == [{"name": "workspace"}] + assert FakeSandbox.created_kwargs["skip_health_check"] is True + + class FailingConnectSandbox(FakeSandbox): + @classmethod + async def connect(cls, *args: Any, **kwargs: Any) -> "FakeSandbox": + del args, kwargs + raise ConnectionError("pod may still be starting") + + monkeypatch.setattr( + opensandbox_provider, + "_require_opensandbox_sdk", + lambda: (FailingConnectSandbox, FakeConnectionConfig, object, FakePlatformSpec, object), + ) + provider = opensandbox_provider.OpenSandboxProvider( + create={"connect_attempt_timeout_s": 0.01, "connect_poll_s": 0.01}, + probe={"command": None}, + ) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(opensandbox_provider.asyncio, "sleep", no_sleep) + with pytest.raises(opensandbox_provider.OpenSandboxCreateTimeoutError): + await provider._connect_after_create( + opensandbox_provider.SandboxHandle(sandbox_id="sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag"), + ) + + +async def test_retry_classification_and_await_sdk_helpers(monkeypatch: pytest.MonkeyPatch) -> None: + provider = opensandbox_provider.OpenSandboxProvider( + operations={"retries": 0}, + probe={"command": None}, + ) + assert await provider.aclose() is None + assert await provider._await_sdk_call(_return_value("ok"), operation="op", sandbox_id="sandbox-1", timeout_s=None) + assert opensandbox_provider._is_retryable_sdk_operation_error(TimeoutError("command timeout")) is False + assert opensandbox_provider._is_retryable_sdk_operation_error(ConnectionError("proxy failed")) is True + wrapped = RuntimeError("wrapper") + wrapped.__cause__ = ConnectionError("connection reset") + assert opensandbox_provider._is_retryable_sdk_operation_error(wrapped) is True + + class FakeHttpxConnectError(Exception): + pass + + monkeypatch.setattr(opensandbox_provider, "_httpx_retryable_types", lambda: (FakeHttpxConnectError,)) + assert opensandbox_provider._is_retryable_create_error(FakeHttpxConnectError("temporary")) is True + assert opensandbox_provider._is_retryable_sdk_operation_error(FakeHttpxConnectError("temporary")) is True + + async def cancelled() -> None: + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await provider._await_sdk_operation( + cancelled, + operation="cancelled", + sandbox_id="sandbox-1", + timeout_s=None, + ) + + +async def _return_value(value: Any) -> Any: + return value diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py new file mode 100644 index 0000000000..627deb74a1 --- /dev/null +++ b/tests/unit_tests/test_sandbox.py @@ -0,0 +1,1212 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import importlib.util +import json +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import pytest +from opentelemetry.sdk.trace.export import SpanExportResult + +from nemo_gym.sandbox import ( + AsyncSandbox, + Sandbox, + SandboxExecResult, + SandboxHandle, + SandboxSpec, + get_provider_class, + list_providers, + register_provider, + rewrite_image, +) +from nemo_gym.sandbox.observability import ( + SandboxRecorder, + build_recorder_from_config, + build_recorder_from_env, + use_recorder, +) +from nemo_gym.sandbox.observability import traces as trace_artifacts +from nemo_gym.sandbox.providers.opensandbox import provider as opensandbox_provider_module +from nemo_gym.sandbox.providers.opensandbox.provider import ( + IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, + IMAGE_PULL_POLICY_EXTENSION_KEY, + OpenSandboxCreateVerificationError, + OpenSandboxProvider, +) +from responses_api_agents.mini_swe_agent_2.sandbox_environment import MiniSWESandboxEnvironment + + +def _has_module(module_name: str) -> bool: + try: + return importlib.util.find_spec(module_name) is not None + except ModuleNotFoundError: + return False + + +requires_tenacity = pytest.mark.skipif( + not _has_module("tenacity"), + reason="tenacity optional sandbox dependency is not installed", +) +requires_otlp_exporter = pytest.mark.skipif( + not _has_module("opentelemetry.exporter.otlp.proto.http.trace_exporter"), + reason="OpenTelemetry OTLP HTTP exporter optional sandbox dependency is not installed", +) + + +class FakeSandboxProvider: + name = "fake" + last_instance: "FakeSandboxProvider | None" = None + + def __init__(self, marker: str = "default") -> None: + self.marker = marker + self.created_specs: list[SandboxSpec] = [] + self.exec_calls: list[dict[str, Any]] = [] + self.write_calls: list[tuple[SandboxHandle, str, str | bytes]] = [] + self.read_calls: list[tuple[SandboxHandle, str]] = [] + self.upload_calls: list[tuple[SandboxHandle, Path, str]] = [] + self.download_calls: list[tuple[SandboxHandle, str, Path]] = [] + self.closed: list[tuple[SandboxHandle, bool]] = [] + self.aclosed = False + FakeSandboxProvider.last_instance = self + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + self.created_specs.append(spec) + return SandboxHandle(sandbox_id="fake-1", provider_name=self.name, raw={"spec": spec}) + + async def create_batch( + self, + spec: SandboxSpec, + count: int, + *, + allow_partial: bool = False, + ) -> list[SandboxHandle]: + del allow_partial + return [await self.create(spec) for _ in range(count)] + + async def connect(self, sandbox_id: str) -> SandboxHandle: + return SandboxHandle(sandbox_id=sandbox_id, provider_name=self.name, raw={}) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + self.exec_calls.append( + { + "handle": handle, + "command": command, + "cwd": cwd, + "env": env, + "timeout_s": timeout_s, + "user": user, + } + ) + return SandboxExecResult(stdout="ok", stderr=None, return_code=0) + + async def write_file(self, handle: SandboxHandle, target_path: str, data: str | bytes) -> None: + self.write_calls.append((handle, target_path, data)) + + async def read_file(self, handle: SandboxHandle, source_path: str) -> bytes: + self.read_calls.append((handle, source_path)) + return f"read:{source_path}".encode() + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + self.upload_calls.append((handle, source_path, target_path)) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + self.download_calls.append((handle, source_path, target_path)) + target_path.parent.mkdir(parents=True, exist_ok=True) + target_path.write_bytes(b"downloaded") + + async def close(self, handle: SandboxHandle, *, delete: bool) -> None: + self.closed.append((handle, delete)) + + async def aclose(self) -> None: + self.aclosed = True + + def handle_reference(self, handle: SandboxHandle) -> dict[str, str]: + return {"kind": "fake", "sandbox_id": handle.sandbox_id} + + async def materialize_handle(self, value: Any) -> SandboxHandle: + return SandboxHandle(sandbox_id=value["sandbox_id"], provider_name=self.name, raw={"materialized": True}) + + +def _test_recorder(output_dir: Path) -> SandboxRecorder: + return SandboxRecorder( + output_dir=output_dir, + otel={ + "enabled": False, + "endpoint": None, + "service_name": "nemo-gym-test", + }, + ) + + +def _mini_swe_command_titles() -> dict[str, Any]: + return { + "strip_prefixes": [ + "cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed &&" + ], + "rules": [ + { + "line_starts_with": ["pytest ", "python -m pytest ", "./tests/runtests.py "], + "search": "last", + "title": "run verifier: {line}", + } + ], + } + + +def _otel_attributes(rows: list[dict[str, Any]]) -> dict[str, Any]: + attrs = {} + for row in rows: + value = row["value"] + if "stringValue" in value: + attrs[row["key"]] = value["stringValue"] + elif "boolValue" in value: + attrs[row["key"]] = value["boolValue"] + elif "intValue" in value: + attrs[row["key"]] = int(value["intValue"]) + elif "doubleValue" in value: + attrs[row["key"]] = value["doubleValue"] + return attrs + + +def _otel_spans(output_dir: Path) -> list[dict[str, Any]]: + trace_payload = json.loads((output_dir / "traces" / "otel_traces.json").read_text()) + return [ + span + for resource_span in trace_payload["resourceSpans"] + for scope_span in resource_span["scopeSpans"] + for span in scope_span["spans"] + ] + + +def _otel_resource_service_names(output_dir: Path) -> set[str]: + trace_payload = json.loads((output_dir / "traces" / "otel_traces.json").read_text()) + service_names = set() + for resource_span in trace_payload["resourceSpans"]: + attrs = _otel_attributes(resource_span["resource"]["attributes"]) + service_names.add(attrs["service.name"]) + return service_names + + +def test_trace_exporter_and_service_lane_edge_cases(tmp_path: Path) -> None: + exporter = trace_artifacts.JsonSpanExporter() + assert exporter.force_flush() + exporter.shutdown() + assert exporter.export([]) == SpanExportResult.FAILURE + assert trace_artifacts.export_trace_artifacts(tmp_path, spans=[]) == {} + + def span(name: str, **attrs: Any) -> Any: + return type("FakeSpan", (), {"name": name, "attributes": attrs})() + + assert trace_artifacts._visual_service_name(span("io", **{"operation.name": "sandbox.read_file"})) == "sandbox.io" + assert trace_artifacts._visual_service_name(span("diag", **{"operation.name": "sandbox.diagnostic.check"})) == ( + "sandbox.diagnostic" + ) + assert trace_artifacts._visual_service_name(span("custom", **{"span.section": "rollout"})) == "nemo-gym.rollout" + assert trace_artifacts._visual_service_name(span("custom", **{"span.section": "sandbox"})) == "sandbox" + assert trace_artifacts._otel_value(["a", 1]) == { + "arrayValue": {"values": [{"stringValue": "a"}, {"intValue": "1"}]} + } + + +def test_sandbox_facade_uses_public_provider_api() -> None: + asyncio.run(_assert_sandbox_facade_uses_public_provider_api()) + + +async def _assert_sandbox_facade_uses_public_provider_api() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + sandbox = AsyncSandbox({"name": provider_name, "kwargs": {"marker": "configured"}}) + handle = await sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "image:tag" + assert provider.created_specs[0].metadata == {"suite": "unit"} + + result = await sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) + assert provider.exec_calls[0] == { + "handle": handle, + "command": "pytest -q", + "cwd": "/repo", + "env": None, + "timeout_s": 60, + "user": "agent", + } + + await sandbox.delete(handle) + assert provider.closed[0] == (handle, True) + assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} + assert await sandbox.materialize_handle({"sandbox_id": "fake-2"}) == SandboxHandle( + sandbox_id="fake-2", provider_name="fake", raw={"materialized": True} + ) + async with AsyncSandbox(provider) as context_sandbox: + assert context_sandbox.provider_name == "fake" + await sandbox.shutdown() + assert provider.aclosed is True + + +def test_rewrite_image_and_materialize_handle_validation() -> None: + asyncio.run(_assert_rewrite_image_and_materialize_handle_validation()) + + +def test_provider_registry_validation_and_listing() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + assert get_provider_class(provider_name) is FakeSandboxProvider + assert provider_name in list_providers() + with pytest.raises(ValueError, match="must be non-empty"): + register_provider("", FakeSandboxProvider) + with pytest.raises(ValueError, match="already registered"): + register_provider(provider_name, FakeSandboxProvider) + with pytest.raises(ValueError, match="Unknown sandbox provider"): + get_provider_class(f"missing-{uuid4().hex}") + + +async def _assert_rewrite_image_and_materialize_handle_validation() -> None: + assert rewrite_image(None, []) is None + assert rewrite_image("image:tag", [{"from": "other/", "to": "mirror/"}]) == "image:tag" + + class BadMaterializeProvider(FakeSandboxProvider): + async def materialize_handle(self, value: Any) -> object: + del value + return object() + + sandbox = AsyncSandbox(BadMaterializeProvider()) + try: + await sandbox.materialize_handle({"sandbox_id": "bad"}) + except TypeError as e: + assert "must return SandboxHandle" in str(e) + else: + raise AssertionError("expected invalid materialize_handle return type to fail") + + +def test_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: + asyncio.run(_assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path)) + + +async def _assert_async_sandbox_batch_file_and_fallback_reference_operations(tmp_path: Path) -> None: + provider = FakeSandboxProvider() + sandbox = AsyncSandbox(provider) + + handles = await sandbox.create_batch(SandboxSpec(image="image:tag"), 2, allow_partial=True) + connected = await sandbox.connect("connected-1") + await sandbox.write_file(connected, "/tmp/file.txt", "contents") + assert await sandbox.read_file(connected, "/tmp/file.txt") == b"read:/tmp/file.txt" + source_path = tmp_path / "source.txt" + target_path = tmp_path / "nested" / "target.txt" + source_path.write_text("local", encoding="utf-8") + await sandbox.upload_file(connected, source_path, "/remote/source.txt") + await sandbox.download_file(connected, "/remote/source.txt", target_path) + await sandbox.close(connected) + + assert [handle.sandbox_id for handle in handles] == ["fake-1", "fake-1"] + assert provider.write_calls == [(connected, "/tmp/file.txt", "contents")] + assert provider.read_calls == [(connected, "/tmp/file.txt")] + assert provider.upload_calls == [(connected, source_path, "/remote/source.txt")] + assert provider.download_calls == [(connected, "/remote/source.txt", target_path)] + assert target_path.read_bytes() == b"downloaded" + + plain_provider = FakeSandboxProvider() + plain_provider.handle_reference = None # type: ignore[method-assign] + plain_provider.materialize_handle = None # type: ignore[method-assign] + plain_sandbox = AsyncSandbox(plain_provider) + plain_handle = SandboxHandle(sandbox_id="plain-1", provider_name="fake", raw={}) + assert plain_sandbox.handle_reference(plain_handle) is plain_handle + assert await plain_sandbox.materialize_handle(plain_handle) is plain_handle + try: + await plain_sandbox.materialize_handle({"sandbox_id": "plain-2"}) + except ValueError as e: + assert "cannot materialize" in str(e) + else: + raise AssertionError("expected materialize_handle without provider support to fail") + + +def test_sync_sandbox_facade_uses_public_provider_api() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + with Sandbox({"name": provider_name, "kwargs": {"marker": "configured"}}) as sandbox: + handle = sandbox.create(SandboxSpec(image="image:tag", metadata={"suite": "unit"})) + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "image:tag" + assert provider.created_specs[0].metadata == {"suite": "unit"} + + result = sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + assert result == SandboxExecResult(stdout="ok", stderr=None, return_code=0) + assert provider.exec_calls[0] == { + "handle": handle, + "command": "pytest -q", + "cwd": "/repo", + "env": None, + "timeout_s": 60, + "user": "agent", + } + + sandbox.delete(handle) + assert provider.closed[0] == (handle, True) + assert sandbox.handle_reference(handle) == {"kind": "fake", "sandbox_id": "fake-1"} + assert sandbox.materialize_handle({"sandbox_id": "fake-3"}).sandbox_id == "fake-3" + assert sandbox.provider_name == "fake" + assert len(sandbox.create_batch(SandboxSpec(image="image:tag"), 2)) == 2 + sandbox.shutdown() + sandbox.shutdown() + assert provider.aclosed is True + try: + sandbox.provider_name + except RuntimeError as e: + assert "sync loop is closed" in str(e) + else: + raise AssertionError("expected closed sync sandbox to reject further calls") + + +def test_sync_sandbox_file_operations(tmp_path: Path) -> None: + provider = FakeSandboxProvider() + with Sandbox(provider) as sandbox: + handle = sandbox.connect("sync-1") + sandbox.write_file(handle, "/tmp/file.txt", b"contents") + assert sandbox.read_file(handle, "/tmp/file.txt") == b"read:/tmp/file.txt" + source_path = tmp_path / "source.txt" + target_path = tmp_path / "target.txt" + source_path.write_text("local", encoding="utf-8") + sandbox.upload_file(handle, source_path, "/remote/source.txt") + sandbox.download_file(handle, "/remote/source.txt", target_path) + + assert provider.write_calls == [(handle, "/tmp/file.txt", b"contents")] + assert provider.read_calls == [(handle, "/tmp/file.txt")] + assert provider.upload_calls == [(handle, source_path, "/remote/source.txt")] + assert provider.download_calls == [(handle, "/remote/source.txt", target_path)] + assert target_path.read_bytes() == b"downloaded" + + +def test_sync_sandbox_facade_rejects_async_context() -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + + async def _create_sync_sandbox_in_async_context() -> None: + Sandbox({"name": provider_name}) + + try: + asyncio.run(_create_sync_sandbox_in_async_context()) + except RuntimeError as e: + assert "use AsyncSandbox in async code" in str(e) + else: + raise AssertionError("expected sync Sandbox to reject async context") + + +def test_sandbox_facade_owns_operation_observability(tmp_path: Path) -> None: + asyncio.run(_assert_sandbox_facade_owns_operation_observability(tmp_path)) + + +async def _assert_sandbox_facade_owns_operation_observability(tmp_path: Path) -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + recorder = _test_recorder(tmp_path) + + with use_recorder(recorder): + sandbox = AsyncSandbox({"name": provider_name}) + handle = await sandbox.create( + SandboxSpec( + image="image:tag", + metadata={ + "benchmark": "swebench-verified", + "instance_id": "django__django-12345", + "nemo_gym_agent": "mini_swe_agent", + }, + ) + ) + await sandbox.exec(handle, "pytest -q", cwd="/repo", timeout_s=60, user="agent") + await sandbox.close(handle, delete=True) + + recorder.finalize() + span_attrs = { + span["name"]: _otel_attributes(span["attributes"]) + for span in _otel_spans(recorder.output_dir) + if span["name"] in {"sandbox.create: image:tag", "exec: pytest -q", "sandbox.cleanup: fake-1"} + } + + assert set(span_attrs) == {"sandbox.create: image:tag", "exec: pytest -q", "sandbox.cleanup: fake-1"} + assert span_attrs["sandbox.create: image:tag"]["trajectory_id"] == "django__django-12345" + assert span_attrs["sandbox.create: image:tag"]["harness"] == "mini_swe_agent" + assert span_attrs["sandbox.create: image:tag"]["benchmark"] == "swebench-verified" + assert span_attrs["sandbox.create: image:tag"]["operation.name"] == "sandbox.start" + assert span_attrs["exec: pytest -q"]["sandbox_id"] == "fake-1" + assert span_attrs["exec: pytest -q"]["command"] == "pytest -q" + assert span_attrs["exec: pytest -q"]["operation.name"] == "trajectory.tool" + assert span_attrs["exec: pytest -q"]["span.section"] == "rollout" + assert "command_class" not in span_attrs["exec: pytest -q"] + assert "command_hash" not in span_attrs["exec: pytest -q"] + assert span_attrs["sandbox.cleanup: fake-1"]["delete"] is True + assert {"sandbox.create", "sandbox.exec", "sandbox.cleanup"}.issubset( + _otel_resource_service_names(recorder.output_dir) + ) + forbidden_prefix = "nemo" + "_rl." + forbidden_hash_attr = "nemo" + "_gym.sandbox_id_hash" + assert all( + not key.startswith(forbidden_prefix) and key != forbidden_hash_attr + for attrs in span_attrs.values() + for key in attrs + ) + + +@requires_tenacity +def test_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch)) + + +async def _assert_opensandbox_sdk_create_receives_default_image_pull_policy(monkeypatch) -> None: + class FakeSDKSandbox: + create_calls: list[dict[str, Any]] = [] + + def __init__(self, sandbox_id: str) -> None: + self.id = sandbox_id + + @classmethod + async def create(cls, **kwargs: Any) -> "FakeSDKSandbox": + cls.create_calls.append(kwargs) + return cls("sdk-sandbox-1") + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (FakeSDKSandbox, object, object, object, object), + ) + + provider = OpenSandboxProvider(probe={"command": None}) + monkeypatch.setattr(provider, "_connection_config", lambda request_timeout_s=None, use_server_proxy=None: object()) + + handle = await provider.create( + SandboxSpec( + image="image:tag", + metadata={ + "harbor_instance_id": "swebench::django__django-10880", + "long": f"bad:{'x' * 80}:", + }, + ) + ) + + assert handle.sandbox_id == "sdk-sandbox-1" + metadata = FakeSDKSandbox.create_calls[0]["metadata"] + assert metadata["harbor_instance_id"] == "swebench_django__django-10880" + assert metadata["long"] == ("bad_" + "x" * 59) + extensions = FakeSDKSandbox.create_calls[0]["extensions"] + assert extensions[IMAGE_PULL_POLICY_EXTENSION_KEY] == "IfNotPresent" + assert extensions[IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY] == "IfNotPresent" + + +@requires_tenacity +def test_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch)) + + +async def _assert_opensandbox_connect_after_create_can_use_direct_exec_endpoint(monkeypatch) -> None: + class FakeConnectionConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeSDKSandbox: + connect_calls: list[dict[str, Any]] = [] + + def __init__(self, sandbox_id: str) -> None: + self.id = sandbox_id + + @classmethod + async def connect(cls, sandbox_id: str, **kwargs: Any) -> "FakeSDKSandbox": + cls.connect_calls.append({"sandbox_id": sandbox_id, **kwargs}) + return cls(sandbox_id) + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (FakeSDKSandbox, FakeConnectionConfig, object, object, object), + ) + + provider = OpenSandboxProvider( + connection={"use_server_proxy": True, "exec_use_server_proxy": False}, + create={"connect_attempt_timeout_s": 1}, + probe={"command": None}, + ) + handle = await provider._connect_after_create( + SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=None), + SandboxSpec(image="image:tag", ready_timeout_s=10), + ) + + assert handle.sandbox_id == "sdk-sandbox-1" + assert isinstance(handle.raw, FakeSDKSandbox) + connect_call = FakeSDKSandbox.connect_calls[0] + assert connect_call["skip_health_check"] is True + assert connect_call["connection_config"].kwargs["use_server_proxy"] is False + + +@requires_tenacity +def test_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_create_probe_can_require_stable_successes(monkeypatch)) + + +async def _assert_opensandbox_create_probe_can_require_stable_successes(monkeypatch) -> None: + provider = OpenSandboxProvider( + probe={ + "command": "true", + "expected_stdout": None, + "stable_count": 3, + "stable_delay_s": 0, + }, + ) + calls: list[dict[str, Any]] = [] + + async def fake_exec( + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + calls.append( + { + "handle": handle, + "command": command, + "cwd": cwd, + "env": env, + "timeout_s": timeout_s, + "user": user, + } + ) + return SandboxExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(provider, "_exec", fake_exec) + handle = SandboxHandle(sandbox_id="sdk-sandbox-0", provider_name="opensandbox", raw=object()) + + await provider._verify_created_handle(handle) + + assert [call["command"] for call in calls] == ["true", "true", "true"] + assert all(call["timeout_s"] == 30 for call in calls) + assert all(call["user"] == "root" for call in calls) + + +@requires_tenacity +def test_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch)) + + +async def _assert_opensandbox_create_probe_polls_same_sandbox_after_transient_errors(monkeypatch) -> None: + provider = OpenSandboxProvider( + create={"connect_poll_s": 0.01}, + probe={ + "command": "true", + "expected_stdout": None, + "timeout_s": 1, + "deadline_s": 2, + "stable_count": 2, + "stable_delay_s": 0, + }, + ) + attempts = 0 + handles: list[SandboxHandle] = [] + + async def fake_exec( + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + del command, cwd, env, timeout_s, user + nonlocal attempts + attempts += 1 + handles.append(handle) + if attempts <= 2: + raise ConnectionError("direct execd endpoint is not accepting connections yet") + return SandboxExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(provider, "_exec", fake_exec) + handle = SandboxHandle(sandbox_id="sdk-sandbox-0", provider_name="opensandbox", raw=object()) + + await provider._verify_created_handle(handle) + + assert attempts == 4 + assert {seen_handle.sandbox_id for seen_handle in handles} == {"sdk-sandbox-0"} + + +def test_opensandbox_create_probe_failures_are_retryable() -> None: + error = OpenSandboxCreateVerificationError("pod sdk-sandbox-0 failed create probe") + + assert opensandbox_provider_module._is_retryable_create_error(error) is True + + +def test_opensandbox_starting_pod_endpoint_errors_are_retryable() -> None: + error = RuntimeError( + "Get endpoint for sandbox sdk-sandbox-0 port 44772 failed: " + "Pod IP is not yet available. The Pod may still be starting." + ) + + assert opensandbox_provider_module._is_retryable_create_error(error) is True + + +@requires_tenacity +def test_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch)) + + +async def _assert_opensandbox_exec_retries_retryable_sdk_failures(monkeypatch) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeLog: + def __init__(self, text: str) -> None: + self.text = text + + class FakeLogs: + stdout = [FakeLog("ok")] + stderr: list[FakeLog] = [] + + class FakeExecution: + logs = FakeLogs() + error = None + exit_code = 0 + + class FakeCommands: + def __init__(self) -> None: + self.calls = 0 + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> FakeExecution: + del command, opts + self.calls += 1 + if self.calls <= 2: + raise ConnectionError("transient proxy failure") + return FakeExecution() + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = OpenSandboxProvider( + operations={ + "retries": 2, + "retry_delay_s": 0, + "retry_max_delay_s": 0, + "command_retries": 2, + }, + probe={"command": None}, + ) + raw = FakeRaw() + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + result = await provider.exec(handle, "echo hello", timeout_s=30) + + assert result.stdout == "ok" + assert result.return_code == 0 + assert raw.commands.calls == 3 + + +@requires_tenacity +def test_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_command_retries_can_be_disabled(monkeypatch)) + + +async def _assert_opensandbox_command_retries_can_be_disabled(monkeypatch) -> None: + class FakeRunCommandOpts: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeCommands: + def __init__(self) -> None: + self.calls = 0 + + async def run(self, command: str, *, opts: FakeRunCommandOpts) -> None: + del command, opts + self.calls += 1 + raise ConnectionError("transient proxy failure") + + class FakeRaw: + def __init__(self) -> None: + self.commands = FakeCommands() + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (object, object, FakeRunCommandOpts, object, object), + ) + + provider = OpenSandboxProvider( + operations={ + "retries": 2, + "retry_delay_s": 0, + "retry_max_delay_s": 0, + "command_retries": 0, + }, + probe={"command": None}, + ) + raw = FakeRaw() + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + try: + await provider.exec(handle, "echo hello", timeout_s=30) + except ConnectionError: + pass + else: + raise AssertionError("expected provider.exec to propagate the command failure") + + assert raw.commands.calls == 1 + + +@requires_tenacity +def test_opensandbox_close_timeout_does_not_fail_after_delete() -> None: + asyncio.run(_assert_opensandbox_close_timeout_does_not_fail_after_delete()) + + +async def _assert_opensandbox_close_timeout_does_not_fail_after_delete() -> None: + class SlowCloseRaw: + def __init__(self) -> None: + self.killed = False + + async def kill(self) -> None: + self.killed = True + + async def close(self) -> None: + await asyncio.sleep(60) + + raw = SlowCloseRaw() + provider = OpenSandboxProvider( + operations={"close_timeout_s": 0.01}, + probe={"command": None}, + ) + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=raw) + + await provider.close(handle, delete=True) + + assert raw.killed is True + + +@requires_tenacity +def test_opensandbox_close_timeout_still_fails_without_delete() -> None: + asyncio.run(_assert_opensandbox_close_timeout_still_fails_without_delete()) + + +async def _assert_opensandbox_close_timeout_still_fails_without_delete() -> None: + class SlowCloseRaw: + async def close(self) -> None: + await asyncio.sleep(60) + + provider = OpenSandboxProvider( + operations={"close_timeout_s": 0.01}, + probe={"command": None}, + ) + handle = SandboxHandle(sandbox_id="sdk-sandbox-1", provider_name="opensandbox", raw=SlowCloseRaw()) + + try: + await provider.close(handle, delete=False) + except TimeoutError: + pass + else: + raise AssertionError("expected close timeout to fail when delete=False") + + +def test_observability_finalize_exports_only_otel_traces(tmp_path: Path) -> None: + recorder = SandboxRecorder( + output_dir=tmp_path / "observability", + run_id="run-1", + run_span_name="unit-job", + otel={ + "enabled": False, + "service_name": "nemo-gym-test", + }, + ) + with recorder.sync_span( + "trajectory.tool", + phase="exec", + attributes={"trajectory_id": "task-1", "sandbox_id": "sandbox-1"}, + ): + pass + + recorder.finalize() + + assert (recorder.output_dir / "traces" / "otel_traces.json").exists() + assert not (recorder.output_dir / "traces" / "chrome_trace.json").exists() + assert not (recorder.output_dir / "summary.json").exists() + trace_payload = json.loads((recorder.output_dir / "traces" / "otel_traces.json").read_text()) + span_names = [ + span["name"] + for resource_span in trace_payload["resourceSpans"] + for scope_span in resource_span["scopeSpans"] + for span in scope_span["spans"] + ] + assert "eval: unit-job" in span_names + assert "exec: " in span_names + assert "rollout: task-1" in span_names + assert _otel_resource_service_names(recorder.output_dir) == { + "nemo-gym.eval", + "nemo-gym.rollout", + "sandbox.exec", + } + + +@requires_otlp_exporter +def test_observability_otlp_exporter_does_not_require_local_artifacts(monkeypatch, tmp_path: Path) -> None: + from opentelemetry.exporter.otlp.proto.http import trace_exporter + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + + exporter_kwargs = [] + + class FakeOTLPSpanExporter(SpanExporter): + def __init__(self, **kwargs: Any) -> None: + exporter_kwargs.append(kwargs) + + def export(self, spans: Any) -> SpanExportResult: + del spans + return SpanExportResult.SUCCESS + + def force_flush(self, timeout_millis: int = 30000) -> bool: + del timeout_millis + return True + + def shutdown(self) -> None: + return None + + monkeypatch.setattr(trace_exporter, "OTLPSpanExporter", FakeOTLPSpanExporter) + + recorder = SandboxRecorder( + output_dir=None, + run_id="run-1", + otel={ + "enabled": True, + "service_name": "sandbox-test", + "traces_exporter": "otlp_http", + "metrics_exporter": "none", + "traces_endpoint": "http://collector:4318", + "traces_headers": {"x-scope-orgid": "sandbox"}, + "traces_timeout_s": 3, + }, + ) + with recorder.sync_span("trajectory.tool", phase="execution", attributes={"trajectory_id": "task-1"}): + pass + recorder.finalize() + + assert exporter_kwargs == [ + { + "endpoint": "http://collector:4318/v1/traces", + "headers": {"x-scope-orgid": "sandbox"}, + "timeout": 3.0, + } + ] + assert not (tmp_path / "traces").exists() + + +def test_observability_env_can_enable_recorder_without_output_dir(monkeypatch) -> None: + monkeypatch.delenv("NEMO_GYM_SANDBOX_OBSERVABILITY_DIR", raising=False) + monkeypatch.setenv("NEMO_GYM_SANDBOX_OBSERVABILITY_TRACES_EXPORTER", "none") + monkeypatch.setenv("NEMO_GYM_SANDBOX_OBSERVABILITY_JOB_NAME", "unit-job") + + recorder = build_recorder_from_env() + + assert recorder is not None + assert recorder.output_dir is None + assert recorder.run_span_name == "eval: unit-job" + recorder.finalize() + + +def test_observability_config_and_env_validation(monkeypatch, tmp_path: Path) -> None: + assert build_recorder_from_config(None) is None + assert build_recorder_from_config({"enabled": False}) is None + + recorder = build_recorder_from_config( + { + "enabled": True, + "output_dir": str(tmp_path / "config-recorder"), + "job_name": "config-job", + "export_traces": False, + }, + run_id="run-1", + ) + assert recorder is not None + assert recorder.run_span_name == "eval: config-job" + assert recorder.output_dir == tmp_path / "config-recorder" + recorder.finalize() + + monkeypatch.setenv("NEMO_GYM_SANDBOX_OBSERVABILITY_COMMAND_TITLES", "[]") + with pytest.raises(ValueError, match="must contain a JSON object"): + build_recorder_from_env() + + +def test_observability_attributes_are_configurable(tmp_path: Path) -> None: + recorder = SandboxRecorder( + output_dir=tmp_path / "observability", + otel={ + "enabled": False, + "attribute_aliases": {"trajectory_id": "custom.trajectory_id"}, + "local_service_name_strategy": "preserve", + "resource_attributes": {"deployment": "unit-test"}, + "service_name": "sandbox-test", + }, + ) + + with recorder.sync_span("trajectory.tool", phase="execution", attributes={"trajectory_id": "task-1"}): + pass + recorder.finalize() + + spans = _otel_spans(recorder.output_dir) + tool_span = next(span for span in spans if span["name"] == "exec: ") + attrs = _otel_attributes(tool_span["attributes"]) + resource_attrs = _otel_attributes( + json.loads((recorder.output_dir / "traces" / "otel_traces.json").read_text())["resourceSpans"][0]["resource"][ + "attributes" + ] + ) + + assert attrs["trajectory_id"] == "task-1" + assert attrs["operation.name"] == "trajectory.tool" + assert attrs["span.section"] == "rollout" + assert attrs["custom.trajectory_id"] == "task-1" + assert ("nemo" + "_rl.trajectory_id") not in attrs + assert resource_attrs["deployment"] == "unit-test" + assert resource_attrs["service.name"] == "sandbox-test" + + +def test_observability_command_span_titles_prefer_verifier_command(tmp_path: Path) -> None: + recorder = SandboxRecorder( + output_dir=tmp_path / "observability", + otel={ + "enabled": False, + "command_titles": _mini_swe_command_titles(), + }, + ) + + command = """cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed && +set -xo pipefail +cd /testbed +git status +pytest -rA testing/test_collection.py +git checkout base testing/test_collection.py +""" + with recorder.sync_span( + "trajectory.tool", + phase="execution", + attributes={"trajectory_id": "task-1", "command": command}, + ): + pass + recorder.finalize() + + spans = _otel_spans(recorder.output_dir) + tool_span = next( + span for span in spans if span["name"] == "exec: run verifier: pytest -rA testing/test_collection.py" + ) + attrs = _otel_attributes(tool_span["attributes"]) + + assert attrs["operation.name"] == "trajectory.tool" + assert attrs["span.section"] == "rollout" + + +def test_observability_command_span_titles_do_not_use_builtin_task_heuristics(tmp_path: Path) -> None: + recorder = SandboxRecorder(output_dir=tmp_path / "observability", otel={"enabled": False}) + + command = """cd /testbed && source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed && +set -xo pipefail +pytest -q +""" + with recorder.sync_span( + "trajectory.tool", + phase="execution", + attributes={"trajectory_id": "task-1", "command": command}, + ): + pass + recorder.finalize() + + spans = _otel_spans(recorder.output_dir) + assert any(span["name"].startswith("exec: cd /testbed && source $(conda info --base)") for span in spans) + assert not any(span["name"].startswith("exec: run verifier:") for span in spans) + + +def test_observability_splits_rollout_llm_and_verifier_sections(tmp_path: Path) -> None: + recorder = SandboxRecorder(output_dir=tmp_path / "observability", otel={"enabled": False}) + + with recorder.sync_span("llm.request", phase="llm", attributes={"trajectory_id": "task-1", "model": "qwen"}): + pass + with recorder.sync_span( + "trajectory.tool", + phase="execution", + attributes={"trajectory_id": "task-1", "command": "ls -la /testbed"}, + ): + pass + with recorder.sync_span( + "trajectory.tool", + phase="execution", + attributes={ + "trajectory_id": "task-1", + "command": "pytest -q", + "execution.section": "verifier", + "span.section": "verifier", + }, + ): + pass + recorder.record_event("trajectory", "trajectory.complete", attributes={"trajectory_id": "task-1", "reward": 1.0}) + recorder.finalize() + + spans = _otel_spans(recorder.output_dir) + span_attrs = {span["name"]: _otel_attributes(span["attributes"]) for span in spans} + span_by_name = {span["name"]: span for span in spans} + + assert "rollout: task-1" in span_attrs + assert "verifier: task-1" in span_attrs + assert span_by_name["verifier: task-1"]["parentSpanId"] == span_by_name["rollout: task-1"]["spanId"] + assert span_attrs["llm.request"]["span.section"] == "rollout" + assert span_attrs["exec: ls -la /testbed"]["span.section"] == "rollout" + assert span_attrs["exec: pytest -q"]["span.section"] == "verifier" + assert { + "nemo-gym.rollout", + "nemo-gym.verifier", + "llm.request", + "sandbox.exec", + "verifier.exec", + }.issubset(_otel_resource_service_names(recorder.output_dir)) + + +def test_observability_can_record_exception_without_stacktrace(tmp_path: Path) -> None: + recorder = SandboxRecorder(output_dir=tmp_path / "observability", otel={"enabled": False}) + + with pytest.raises(RuntimeError): + with recorder.sync_span( + "llm.request", + phase="llm", + attributes={"trajectory_id": "task-1", "_record_exception_stacktrace": False}, + ): + raise RuntimeError("format retry") + recorder.finalize() + + llm_span = next(span for span in _otel_spans(recorder.output_dir) if span["name"] == "llm.request") + attrs = _otel_attributes(llm_span["attributes"]) + events = llm_span.get("events") or [] + + assert attrs["status"] == "error" + assert attrs["error_type"] == "RuntimeError" + assert events + assert events[0]["name"] == "exception" + event_attrs = _otel_attributes(events[0]["attributes"]) + assert event_attrs["exception.type"] == "RuntimeError" + assert "exception.stacktrace" not in event_attrs + + +def test_mini_swe_sandbox_environment_owns_conda_setup(monkeypatch, tmp_path: Path) -> None: + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + monkeypatch.setenv("FORWARDED_KEY", "forwarded-value") + recorder = SandboxRecorder( + output_dir=tmp_path, + otel={ + "enabled": False, + "endpoint": None, + "service_name": "nemo-gym-test", + "command_titles": _mini_swe_command_titles(), + }, + ) + + with use_recorder(recorder): + env = MiniSWESandboxEnvironment( + image="upstream/image:tag", + cwd="/testbed", + provider={"name": provider_name, "kwargs": {"marker": "configured"}}, + spec={ + "image_rewrites": [{"from": "upstream/", "to": "mirror/"}], + "metadata": {"suite": "unit"}, + "resources": {"cpu": "1"}, + }, + env={"STATIC_KEY": "static-value"}, + forward_env=["FORWARDED_KEY"], + conda_env="testbed", + activate_conda=True, + user="agent", + delete=True, + ) + + try: + assert env.get_template_vars(extra="value")["extra"] == "value" + serialized = env.serialize() + assert serialized["info"]["config"]["environment_type"].endswith("MiniSWESandboxEnvironment") + env.config.activate_conda = False + assert env._command("echo plain", "/tmp/work") == "echo plain" + env.config.activate_conda = True + + provider = FakeSandboxProvider.last_instance + assert provider is not None + assert provider.marker == "configured" + assert provider.created_specs[0].image == "mirror/image:tag" + assert provider.created_specs[0].env == { + "FORWARDED_KEY": "forwarded-value", + "STATIC_KEY": "static-value", + } + + result = env.execute("pytest -q", is_eval=True) + assert result == {"output": "ok", "returncode": 0, "exception_info": ""} + exec_call = provider.exec_calls[0] + assert exec_call["cwd"] == "/" + assert exec_call["timeout_s"] == 1800 + assert exec_call["user"] == "agent" + assert "conda activate testbed" in exec_call["command"] + assert exec_call["command"].endswith("pytest -q") + finally: + env.cleanup() + env.cleanup() + + recorder.finalize() + spans = _otel_spans(recorder.output_dir) + span_attrs = {span["name"]: _otel_attributes(span["attributes"]) for span in spans} + + assert FakeSandboxProvider.last_instance is not None + assert FakeSandboxProvider.last_instance.closed[0][1] is True + assert "verifier: unknown" in span_attrs + assert span_attrs["exec: run verifier: pytest -q"]["span.section"] == "verifier" + + +def test_mini_swe_sandbox_environment_validation_and_context_manager() -> None: + with pytest.raises(ValueError, match="requires provider"): + MiniSWESandboxEnvironment(image="image:tag") + + provider_name = f"fake-{uuid4().hex}" + register_provider(provider_name, FakeSandboxProvider) + with MiniSWESandboxEnvironment( + image="image:tag", + provider={"name": provider_name}, + delete=False, + ) as env: + assert env._handle is not None + + assert FakeSandboxProvider.last_instance is not None + assert FakeSandboxProvider.last_instance.closed[-1][1] is False diff --git a/uv.lock b/uv.lock index 91e3a53742..91ecb35c46 100644 --- a/uv.lock +++ b/uv.lock @@ -1412,6 +1412,13 @@ dev = [ { name = "requests-mock" }, { name = "ruff" }, ] +sandbox = [ + { name = "opensandbox" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, + { name = "tenacity" }, +] [package.dev-dependencies] docs = [ @@ -1443,6 +1450,10 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, { name = "omegaconf" }, { name = "openai", specifier = "<=2.7.2" }, + { name = "opensandbox", marker = "extra == 'sandbox'", specifier = ">=0.1.9" }, + { name = "opentelemetry-api", marker = "extra == 'sandbox'", specifier = ">=1.36.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'sandbox'", specifier = ">=1.36.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'sandbox'", specifier = ">=1.36.0" }, { name = "orjson" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.6.0" }, { name = "psutil" }, @@ -1458,6 +1469,7 @@ requires-dist = [ { name = "requests-mock", marker = "extra == 'dev'" }, { name = "rich" }, { name = "ruff", marker = "extra == 'dev'" }, + { name = "tenacity", marker = "extra == 'sandbox'", specifier = ">=9.1.4" }, { name = "tqdm" }, { name = "urllib3", specifier = ">=2.6.3" }, { name = "uvicorn" }, @@ -1465,7 +1477,7 @@ requires-dist = [ { name = "wandb" }, { name = "yappi" }, ] -provides-extras = ["dev"] +provides-extras = ["sandbox", "dev"] [package.metadata.requires-dev] docs = [ @@ -1620,31 +1632,76 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/68/162c97ea78c957d68ecf78a5c5041d2e25bd5562bdf5d89a6cbf7f8429bf/opencensus_context-0.1.3-py2.py3-none-any.whl", hash = "sha256:073bb0590007af276853009fac7e4bab1d523c3f03baf4cb4511ca38967c6039", size = 5060, upload-time = "2022-08-03T22:20:20.352Z" }, ] +[[package]] +name = "opensandbox" +version = "0.1.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/2a/ab3cc141e041f71a373c97fcda8749dba9328f1b9bf80401378c0611556f/opensandbox-0.1.9.tar.gz", hash = "sha256:670fbf292c498f8467963d21e91ade9ea8b8f63f4ef18d18fff9581e0952ec03", size = 160034, upload-time = "2026-05-12T12:27:20.692Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/9b/553f8d7a30eddb12785711b2a1c682386878e2bb95450acd806f9fa62930/opensandbox-0.1.9-py3-none-any.whl", hash = "sha256:17faed35b60a982fee5a643fed8e4e12f041e5432d5ea0665d2828d1f2082759", size = 360945, upload-time = "2026-05-12T12:27:19.465Z" }, +] + [[package]] name = "opentelemetry-api" -version = "1.36.0" +version = "1.38.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "importlib-metadata" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/27/d2/c782c88b8afbf961d6972428821c302bd1e9e7bc361352172f0ca31296e2/opentelemetry_api-1.36.0.tar.gz", hash = "sha256:9a72572b9c416d004d492cbc6e61962c0501eaf945ece9b5a0f56597d8348aa0", size = 64780, upload-time = "2025-07-29T15:12:06.02Z" } +sdist = { url = "https://files.pythonhosted.org/packages/08/d8/0f354c375628e048bd0570645b310797299754730079853095bf000fba69/opentelemetry_api-1.38.0.tar.gz", hash = "sha256:f4c193b5e8acb0912b06ac5b16321908dd0843d75049c091487322284a3eea12", size = 65242, upload-time = "2025-10-16T08:35:50.25Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/a2/d86e01c28300bd41bab8f18afd613676e2bd63515417b77636fc1add426f/opentelemetry_api-1.38.0-py3-none-any.whl", hash = "sha256:2891b0197f47124454ab9f0cf58f3be33faca394457ac3e09daba13ff50aa582", size = 65947, upload-time = "2025-10-16T08:35:30.23Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.38.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/83/dd4660f2956ff88ed071e9e0e36e830df14b8c5dc06722dbde1841accbe8/opentelemetry_exporter_otlp_proto_common-1.38.0.tar.gz", hash = "sha256:e333278afab4695aa8114eeb7bf4e44e65c6607d54968271a249c180b2cb605c", size = 20431, upload-time = "2025-10-16T08:35:53.285Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/9e/55a41c9601191e8cd8eb626b54ee6827b9c9d4a46d736f32abc80d8039fc/opentelemetry_exporter_otlp_proto_common-1.38.0-py3-none-any.whl", hash = "sha256:03cb76ab213300fe4f4c62b7d8f17d97fcfd21b89f0b5ce38ea156327ddda74a", size = 18359, upload-time = "2025-10-16T08:35:34.099Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-http" +version = "1.38.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/0a/debcdfb029fbd1ccd1563f7c287b89a6f7bef3b2902ade56797bfd020854/opentelemetry_exporter_otlp_proto_http-1.38.0.tar.gz", hash = "sha256:f16bd44baf15cbe07633c5112ffc68229d0edbeac7b37610be0b2def4e21e90b", size = 17282, upload-time = "2025-10-16T08:35:54.422Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/ee/6b08dde0a022c463b88f55ae81149584b125a42183407dc1045c486cc870/opentelemetry_api-1.36.0-py3-none-any.whl", hash = "sha256:02f20bcacf666e1333b6b1f04e647dc1d5111f86b8e510238fcc56d7762cda8c", size = 65564, upload-time = "2025-07-29T15:11:47.998Z" }, + { url = "https://files.pythonhosted.org/packages/e5/77/154004c99fb9f291f74aa0822a2f5bbf565a72d8126b3a1b63ed8e5f83c7/opentelemetry_exporter_otlp_proto_http-1.38.0-py3-none-any.whl", hash = "sha256:84b937305edfc563f08ec69b9cb2298be8188371217e867c1854d77198d0825b", size = 19579, upload-time = "2025-10-16T08:35:36.269Z" }, ] [[package]] name = "opentelemetry-exporter-prometheus" -version = "0.57b0" +version = "0.59b0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-sdk" }, { name = "prometheus-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d6/d8/5f04c6d51c0823c3d8ac973a2a38db6fcf2d040ca3f08fc66b3c14b6e164/opentelemetry_exporter_prometheus-0.57b0.tar.gz", hash = "sha256:9eb15bdc189235cf03c3f93abf56f8ff0ab57a493a189263bd7fe77a4249e689", size = 14906, upload-time = "2025-07-29T15:12:09.96Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/07/39370ec7eacfca10462121a0e036b66ccea3a616bf6ae6ea5fdb72e5009d/opentelemetry_exporter_prometheus-0.59b0.tar.gz", hash = "sha256:d64f23c49abb5a54e271c2fbc8feacea0c394a30ec29876ab5ef7379f08cf3d7", size = 14972, upload-time = "2025-10-16T08:35:55.973Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c1/1c/40fb93a7b7e495985393bbc734104d5d20e470811644dd56c2402d683739/opentelemetry_exporter_prometheus-0.57b0-py3-none-any.whl", hash = "sha256:c5b893d1cdd593fb022af2c7de3258c2d5a4d04402ae80d9fa35675fed77f05c", size = 12922, upload-time = "2025-07-29T15:11:54.055Z" }, + { url = "https://files.pythonhosted.org/packages/05/ea/3005a732002242fd86203989520bdd5a752e1fd30dc225d5d45751ea19fb/opentelemetry_exporter_prometheus-0.59b0-py3-none-any.whl", hash = "sha256:71ced23207abd15b30d1fe4e7e910dcaa7c2ff1f24a6ffccbd4fdded676f541b", size = 13017, upload-time = "2025-10-16T08:35:37.253Z" }, ] [[package]] @@ -1661,29 +1718,29 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.36.0" +version = "1.38.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-semantic-conventions" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4c/85/8567a966b85a2d3f971c4d42f781c305b2b91c043724fa08fd37d158e9dc/opentelemetry_sdk-1.36.0.tar.gz", hash = "sha256:19c8c81599f51b71670661ff7495c905d8fdf6976e41622d5245b791b06fa581", size = 162557, upload-time = "2025-07-29T15:12:16.76Z" } +sdist = { url = "https://files.pythonhosted.org/packages/85/cb/f0eee1445161faf4c9af3ba7b848cc22a50a3d3e2515051ad8628c35ff80/opentelemetry_sdk-1.38.0.tar.gz", hash = "sha256:93df5d4d871ed09cb4272305be4d996236eedb232253e3ab864c8620f051cebe", size = 171942, upload-time = "2025-10-16T08:36:02.257Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0b/59/7bed362ad1137ba5886dac8439e84cd2df6d087be7c09574ece47ae9b22c/opentelemetry_sdk-1.36.0-py3-none-any.whl", hash = "sha256:19fe048b42e98c5c1ffe85b569b7073576ad4ce0bcb6e9b4c6a39e890a6c45fb", size = 119995, upload-time = "2025-07-29T15:12:03.181Z" }, + { url = "https://files.pythonhosted.org/packages/2f/2e/e93777a95d7d9c40d270a371392b6d6f1ff170c2a3cb32d6176741b5b723/opentelemetry_sdk-1.38.0-py3-none-any.whl", hash = "sha256:1c66af6564ecc1553d72d811a01df063ff097cdc82ce188da9951f93b8d10f6b", size = 132349, upload-time = "2025-10-16T08:35:46.995Z" }, ] [[package]] name = "opentelemetry-semantic-conventions" -version = "0.57b0" +version = "0.59b0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7e/31/67dfa252ee88476a29200b0255bda8dfc2cf07b56ad66dc9a6221f7dc787/opentelemetry_semantic_conventions-0.57b0.tar.gz", hash = "sha256:609a4a79c7891b4620d64c7aac6898f872d790d75f22019913a660756f27ff32", size = 124225, upload-time = "2025-07-29T15:12:17.873Z" } +sdist = { url = "https://files.pythonhosted.org/packages/40/bc/8b9ad3802cd8ac6583a4eb7de7e5d7db004e89cb7efe7008f9c8a537ee75/opentelemetry_semantic_conventions-0.59b0.tar.gz", hash = "sha256:7a6db3f30d70202d5bf9fa4b69bc866ca6a30437287de6c510fb594878aed6b0", size = 129861, upload-time = "2025-10-16T08:36:03.346Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/05/75/7d591371c6c39c73de5ce5da5a2cc7b72d1d1cd3f8f4638f553c01c37b11/opentelemetry_semantic_conventions-0.57b0-py3-none-any.whl", hash = "sha256:757f7e76293294f124c827e514c2a3144f191ef175b069ce8d1211e1e38e9e78", size = 201627, upload-time = "2025-07-29T15:12:04.174Z" }, + { url = "https://files.pythonhosted.org/packages/24/7d/c88d7b15ba8fe5c6b8f93be50fc11795e9fc05386c44afaf6b76fe191f9b/opentelemetry_semantic_conventions-0.59b0-py3-none-any.whl", hash = "sha256:35d3b8833ef97d614136e253c1da9342b4c3c083bbaf29ce31d572a1c3825eed", size = 207954, upload-time = "2025-10-16T08:35:48.054Z" }, ] [[package]] @@ -2738,6 +2795,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/f0/1098f6628bbe04b086ce59692d09b116ec751286eb7d33e88c5bf0c2e210/swagger_plugin_for_sphinx-6.0.0-py3-none-any.whl", hash = "sha256:35dc646d759a44ce78aefde2fe34f54e7b8c3439d0a52541a6a8b9924a711832", size = 11253, upload-time = "2025-10-16T06:26:08.504Z" }, ] +[[package]] +name = "tenacity" +version = "9.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, +] + [[package]] name = "tqdm" version = "4.67.1"