diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 6a7fc88a7ea..275c25c0069 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -42,6 +42,7 @@ create_provider, ) from nemo_gym.telemetry._fallbacks import is_span_group_enabled, managed_span, safe_set_span_attributes +from nemo_gym.telemetry.metrics import record_sandbox_active from nemo_gym.telemetry.span_groups import GymSpanGroup @@ -393,6 +394,9 @@ def __init__( self._handle: SandboxHandle | None = None self._stopped = True self._closed = False + # Whether this instance added itself to `gym.sandbox.active`; a `connect()`ed sandbox + # was counted by the process that started it and must not be subtracted here. + self._counted_active = False self.pty = SandboxPty(self) def _telemetry_provider_name(self) -> str: @@ -423,6 +427,8 @@ async def start( **{"nemo.gym.sandbox.provider": self._telemetry_provider_name()}, ): handle = await self._provider.create(requested_spec) + record_sandbox_active(1, provider=self._telemetry_provider_name()) + self._counted_active = True else: handle = await self._provider.create(requested_spec) self._handle = handle @@ -573,8 +579,13 @@ async def stop(self) -> None: return try: if self._handle is not None and not self._stopped: - await self._provider.close(self._handle) - self._stopped = True + try: + await self._provider.close(self._handle) + finally: + self._stopped = True + if self._counted_active: + self._counted_active = False + record_sandbox_active(-1, provider=self._telemetry_provider_name()) finally: if self._owns_provider: await self._provider.aclose() diff --git a/nemo_gym/telemetry/metrics.py b/nemo_gym/telemetry/metrics.py index 298cf7b0d0e..1d4314abc49 100644 --- a/nemo_gym/telemetry/metrics.py +++ b/nemo_gym/telemetry/metrics.py @@ -47,6 +47,11 @@ ``http.request.method`` and ``http.response.status_code``, which is strictly better. Use that instead; see ``fern/versions/latest/pages/observability/metrics.mdx``. +``gym.sandbox.active`` (up-down counter, Gym-owned) + Used. Created on the lens meter by :func:`record_sandbox_active`: ``+1`` when a provider + hands back a handle, ``-1`` when Gym releases it, per ``nemo.gym.sandbox.provider``. A + counter rather than a gauge so the backend sums it across the processes holding sandboxes. + Every function here is a no-op unless telemetry is initialised *and* exporting, so call sites do not need their own guards for correctness — though they should still sit under a span-group gate to stay free when disabled. @@ -65,6 +70,12 @@ _VERIFY_TOTAL = 0 _VERIFY_SUCCEEDED = 0 +SANDBOX_ACTIVE_INSTRUMENT = "gym.sandbox.active" +SANDBOX_PROVIDER_ATTRIBUTE = "nemo.gym.sandbox.provider" +#: (meter, counter) so a re-initialised telemetry handle gets a fresh instrument. +_SANDBOX_ACTIVE_LOCK = threading.Lock() +_SANDBOX_ACTIVE: tuple[object, object] | None = None + def _record(**kwargs) -> None: """Forward to ``record_gym_metrics`` when a meter is available; never raise.""" @@ -116,6 +127,33 @@ def record_active_servers(count: int) -> None: _record(active_servers=count) +def _sandbox_active_counter(meter): + """The process's ``gym.sandbox.active`` instrument on ``meter``, created once per meter.""" + global _SANDBOX_ACTIVE + with _SANDBOX_ACTIVE_LOCK: + if _SANDBOX_ACTIVE is None or _SANDBOX_ACTIVE[0] is not meter: + counter = meter.create_up_down_counter( + SANDBOX_ACTIVE_INSTRUMENT, + unit="{sandbox}", + description="Sandboxes this process currently holds, from provider create to release.", + ) + _SANDBOX_ACTIVE = (meter, counter) + return _SANDBOX_ACTIVE[1] + + +def record_sandbox_active(delta: int, *, provider: str) -> None: + """Add ``delta`` (``+1`` on start, ``-1`` on stop) to ``gym.sandbox.active`` for ``provider``.""" + from nemo_gym.telemetry.setup import get_telemetry + + telemetry = get_telemetry() + if telemetry is None or not telemetry.is_exporting: + return + try: + _sandbox_active_counter(telemetry.meter).add(delta, {SANDBOX_PROVIDER_ATTRIBUTE: provider}) + except Exception: + logger.debug("nemo-lens: failed to record %s", SANDBOX_ACTIVE_INSTRUMENT, exc_info=True) + + def _reset_verify_tally_for_testing() -> None: """Reset the cumulative verify tally. Test-only.""" global _VERIFY_TOTAL, _VERIFY_SUCCEEDED diff --git a/tests/unit_tests/telemetry/test_sandbox_active.py b/tests/unit_tests/telemetry/test_sandbox_active.py new file mode 100644 index 00000000000..a618ebd65ac --- /dev/null +++ b/tests/unit_tests/telemetry/test_sandbox_active.py @@ -0,0 +1,189 @@ +# 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. + +"""``gym.sandbox.active``: the Gym-owned up-down counter behind "sandboxes running now". + +Asserted against a real in-memory OTel metric reader. The counter is created on the lens +meter directly (it is not one of lens's ``gym.*`` instruments), so these run wherever the +OTel SDK is importable, with or without nemo-lens. +""" + +import pytest + +from nemo_gym.sandbox import api as sandbox_api +from nemo_gym.sandbox.api import AsyncSandbox, SandboxSpec +from nemo_gym.telemetry import metrics as telemetry_metrics +from nemo_gym.telemetry import setup as telemetry_setup +from tests.unit_tests.test_sandbox import FakeSandboxProvider + + +pytest.importorskip("opentelemetry.sdk.metrics") + +PROVIDER_ATTR = telemetry_metrics.SANDBOX_PROVIDER_ATTRIBUTE +ACTIVE = telemetry_metrics.SANDBOX_ACTIVE_INSTRUMENT + + +@pytest.fixture +def collected_metrics(monkeypatch): + """A live meter on an in-memory reader, installed as the process handle.""" + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + + class _Handle: + is_exporting = True + meter = provider.get_meter("test") + + monkeypatch.setattr(telemetry_setup, "_TELEMETRY_HANDLE", _Handle()) + monkeypatch.setattr(telemetry_metrics, "_SANDBOX_ACTIVE", None) + + def collect(): + data = reader.get_metrics_data() + out = {} + # The reader hands back None when no instrument has recorded yet. + for resource_metric in data.resource_metrics if data is not None else (): + for scope_metric in resource_metric.scope_metrics: + for metric in scope_metric.metrics: + out[metric.name] = list(metric.data.data_points) + return out + + return collect + + +def _active_by_provider(points) -> dict[str, int]: + return {p.attributes[PROVIDER_ATTR]: p.value for p in points} + + +@pytest.fixture +def sandbox_group_on(monkeypatch): + monkeypatch.setattr(sandbox_api, "is_span_group_enabled", lambda group: True) + + +def test_recording_without_telemetry_is_a_no_op(monkeypatch): + monkeypatch.setattr(telemetry_setup, "_TELEMETRY_HANDLE", None) + telemetry_metrics.record_sandbox_active(1, provider="fake") + + +def test_recording_errors_never_reach_the_caller(monkeypatch): + class _Broken: + is_exporting = True + + @property + def meter(self): + raise RuntimeError("meter is gone") + + monkeypatch.setattr(telemetry_setup, "_TELEMETRY_HANDLE", _Broken()) + telemetry_metrics.record_sandbox_active(1, provider="fake") + + +def test_counter_is_non_monotonic_and_dimensioned_by_provider(collected_metrics): + telemetry_metrics.record_sandbox_active(1, provider="opensandbox") + telemetry_metrics.record_sandbox_active(1, provider="opensandbox") + telemetry_metrics.record_sandbox_active(1, provider="local") + telemetry_metrics.record_sandbox_active(-1, provider="opensandbox") + + points = collected_metrics()[ACTIVE] + assert _active_by_provider(points) == {"opensandbox": 1, "local": 1} + # A second reader collection must still see the same instrument, not a duplicate one. + telemetry_metrics.record_sandbox_active(-1, provider="local") + assert _active_by_provider(collected_metrics()[ACTIVE])["local"] == 0 + + +def test_instrument_is_created_once_per_meter(collected_metrics): + handle = telemetry_setup.get_telemetry() + first = telemetry_metrics._sandbox_active_counter(handle.meter) + assert telemetry_metrics._sandbox_active_counter(handle.meter) is first + + +@pytest.mark.asyncio +async def test_sandbox_start_and_stop_move_the_counter(collected_metrics, sandbox_group_on): + provider = FakeSandboxProvider() + sandbox = AsyncSandbox(provider, SandboxSpec(image="img")) + + await sandbox.start() + assert _active_by_provider(collected_metrics()[ACTIVE]) == {"fake": 1} + + await sandbox.stop() + assert _active_by_provider(collected_metrics()[ACTIVE]) == {"fake": 0} + + # A second stop (idempotent path) must not subtract again. + await sandbox.stop() + assert _active_by_provider(collected_metrics()[ACTIVE]) == {"fake": 0} + + +@pytest.mark.asyncio +async def test_failed_start_does_not_count(collected_metrics, sandbox_group_on): + class _Refusing(FakeSandboxProvider): + async def create(self, spec): + raise RuntimeError("no capacity") + + sandbox = AsyncSandbox(_Refusing(), SandboxSpec(image="img")) + with pytest.raises(RuntimeError): + await sandbox.start() + assert ACTIVE not in collected_metrics() + + +@pytest.mark.asyncio +async def test_failed_initial_upload_releases_the_count(collected_metrics, sandbox_group_on): + class _FailingUpload(FakeSandboxProvider): + async def upload_file(self, handle, source_path, target_path): + raise RuntimeError("upload failed") + + sandbox = AsyncSandbox(_FailingUpload(), SandboxSpec(image="img", files={"/tmp/a": "x"})) + with pytest.raises(RuntimeError): + await sandbox.start() + assert _active_by_provider(collected_metrics()[ACTIVE]) == {"fake": 0} + + +@pytest.mark.asyncio +async def test_counter_is_released_even_when_provider_close_fails(collected_metrics, sandbox_group_on): + class _StickyClose(FakeSandboxProvider): + async def close(self, handle): + raise RuntimeError("close failed") + + sandbox = AsyncSandbox(_StickyClose(), SandboxSpec(image="img")) + await sandbox.start() + with pytest.raises(RuntimeError): + await sandbox.stop() + assert _active_by_provider(collected_metrics()[ACTIVE]) == {"fake": 0} + + +@pytest.mark.asyncio +async def test_sandbox_group_off_records_nothing(collected_metrics, monkeypatch): + monkeypatch.setattr(sandbox_api, "is_span_group_enabled", lambda group: False) + sandbox = AsyncSandbox(FakeSandboxProvider(), SandboxSpec(image="img")) + await sandbox.start() + await sandbox.stop() + assert ACTIVE not in collected_metrics() + + +@pytest.mark.asyncio +async def test_connected_sandbox_is_not_subtracted_on_stop(collected_metrics, sandbox_group_on): + """A sandbox rebuilt via connect() was counted by the process that started it.""" + from nemo_gym.sandbox.providers.base import ConnectableProvider, SandboxHandle + + class _Connectable(FakeSandboxProvider): + async def connect(self, descriptor): + return SandboxHandle(sandbox_id=descriptor["sandbox_id"], provider_name=self.name, raw={}) + + async def serialize_handle(self, handle, *, scope=None): + return {"sandbox_id": handle.sandbox_id} + + assert isinstance(_Connectable(), ConnectableProvider) + sandbox = await AsyncSandbox.connect({"sandbox_id": "elsewhere"}, provider=_Connectable()) + await sandbox.stop() + assert ACTIVE not in collected_metrics()