feat(sandbox): validate OpenSandbox provider_options via a frozen dataclass

nemo_gym/sandbox/providers/opensandbox/provider.py

@@ -19,7 +19,7 @@
 import re
 import shlex
 from collections.abc import Mapping
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from datetime import timedelta
 from pathlib import Path
 from typing import Any, Awaitable, Callable
@@ -92,10 +92,6 @@ class OpenSandboxCreateVerificationError(SandboxCreateVerificationError):
 DEFAULT_IMAGE_PULL_POLICY = "IfNotPresent"
 IMAGE_PULL_POLICY_EXTENSION_KEY = "imagePullPolicy"
 IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY = "opensandbox.extensions.image-pull-policy"
-PROVIDER_OPTION_PLATFORM = "platform"
-PROVIDER_OPTION_SKIP_HEALTH_CHECK = "skip_health_check"
-PROVIDER_OPTION_SNAPSHOT_ID = "snapshot_id"
-PROVIDER_OPTION_VOLUMES = "volumes"
 VALID_IMAGE_PULL_POLICIES = {"Always", "IfNotPresent", "Never"}
 STATUS_CODE_RE = re.compile(r"(?:status code|http)\D+(\d{3})", re.IGNORECASE)
 
@@ -323,26 +319,6 @@ def _to_volumes(volumes: list[Mapping[str, Any]]) -> list[Any]:
     return [Volume(**dict(volume)) for volume in volumes]
 
 
-def _spec_volumes(spec: SandboxSpec) -> list[Mapping[str, Any]] | None:
-    return spec.provider_options.get(PROVIDER_OPTION_VOLUMES)
-
-
-def _spec_extensions(spec: SandboxSpec) -> dict[str, str]:
-    value = spec.provider_options.get("extensions", {})
-    if not isinstance(value, Mapping):
-        raise TypeError("OpenSandbox provider option 'extensions' must be a mapping")
-    return _string_map(dict(value))
-
-
-def _provider_option_bool(provider_options: dict[str, Any], key: str) -> bool | None:
-    value = provider_options.get(key)
-    if value is None:
-        return None
-    if not isinstance(value, bool):
-        raise TypeError(f"OpenSandbox provider option {key!r} must be a bool")
-    return value
-
-
 def _to_sandbox_status(state: Any) -> SandboxStatus:
     normalized = str(state or "").lower()
     if normalized in {"active", "ready", "running"}:
@@ -453,6 +429,60 @@ def _coerce_config(value: Any, config_cls: type[Any]) -> Any:
     raise TypeError(f"{config_cls.__name__} must be a mapping or {config_cls.__name__} instance")
 
 
+@dataclass(frozen=True)
+class OpenSandboxProviderOptions:
+    """Recognized per-sandbox create options read from ``SandboxSpec.provider_options``.
+
+    ``platform`` and ``volumes`` entries are passed through to the OpenSandbox SDK,
+    so their inner fields are validated by the SDK rather than here.
+    """
+
+    platform: Mapping[str, Any] | None = None
+    snapshot_id: str | None = None
+    volumes: tuple[Mapping[str, Any], ...] = ()
+    skip_health_check: bool | None = None
+    extensions: Mapping[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def from_mapping(cls, options: Mapping[str, Any] | None) -> "OpenSandboxProviderOptions":
+        if options is None:
+            return cls()
+        if not isinstance(options, Mapping):
+            raise TypeError("OpenSandbox provider_options must be a mapping")
+
+        allowed = set(cls.__dataclass_fields__)
+        unknown = set(options) - allowed
+        if unknown:
+            raise ValueError(
+                f"Unknown OpenSandbox provider option(s): {', '.join(sorted(unknown))}. "
+                f"Supported: {', '.join(sorted(allowed))}"
+            )
+
+        platform = options.get("platform")
+        if platform is not None and not isinstance(platform, Mapping):
+            raise TypeError("OpenSandbox provider option 'platform' must be a mapping")
+        snapshot_id = options.get("snapshot_id")
+        if snapshot_id is not None and not isinstance(snapshot_id, str):
+            raise TypeError("OpenSandbox provider option 'snapshot_id' must be a string")
+        volumes = options.get("volumes") or ()
+        if not isinstance(volumes, (list, tuple)) or not all(isinstance(volume, Mapping) for volume in volumes):
+            raise TypeError("OpenSandbox provider option 'volumes' must be a list of mappings")
+        skip_health_check = options.get("skip_health_check")
+        if skip_health_check is not None and not isinstance(skip_health_check, bool):
+            raise TypeError("OpenSandbox provider option 'skip_health_check' must be a bool")
+        extensions = options.get("extensions", {})
+        if not isinstance(extensions, Mapping):
+            raise TypeError("OpenSandbox provider option 'extensions' must be a mapping")
+
+        return cls(
+            platform=dict(platform) if platform is not None else None,
+            snapshot_id=snapshot_id,
+            volumes=tuple(dict(volume) for volume in volumes),
+            skip_health_check=skip_health_check,
+            extensions=_string_map(dict(extensions)),
+        )
+
+
 class OpenSandboxProvider:
     """Provider backed by the OpenSandbox SDK/server API."""
 
@@ -471,23 +501,21 @@ def __init__(
         self._probe = _coerce_config(probe, OpenSandboxProbeConfig)
         self._operations = _coerce_config(operations, OpenSandboxOperationConfig)
 
-    def _with_default_image_pull_policy(self, spec: SandboxSpec) -> SandboxSpec:
-        """Ensure SDK create requests carry the desired image pull policy."""
+    def _resolve_extensions(self, extensions: Mapping[str, str]) -> dict[str, str]:
+        """Add the configured default image pull policy to SDK create extensions."""
+        resolved = dict(extensions)
         if self._create.image_pull_policy is None:
-            return spec
+            return resolved
 
-        provider_options = dict(spec.provider_options)
-        extensions = _spec_extensions(spec)
-        image_pull_policy = extensions.get(IMAGE_PULL_POLICY_EXTENSION_KEY) or extensions.get(
-            IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY
+        image_pull_policy = (
+            resolved.get(IMAGE_PULL_POLICY_EXTENSION_KEY)
+            or resolved.get(IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY)
+            or self._create.image_pull_policy
         )
-        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)
-        provider_options["extensions"] = extensions
-        return replace(spec, provider_options=provider_options)
+        resolved.setdefault(IMAGE_PULL_POLICY_EXTENSION_KEY, image_pull_policy)
+        resolved.setdefault(IMAGE_PULL_POLICY_ANNOTATION_EXTENSION_KEY, image_pull_policy)
+        return resolved
 
     def _connection_config(
         self,
@@ -694,37 +722,33 @@ async def _connect_after_create(self, handle: SandboxHandle, spec: SandboxSpec)
     async def _create_once(self, spec: SandboxSpec) -> SandboxHandle:
         """Create a sandbox through ``opensandbox.Sandbox.create``."""
         Sandbox, _, _, _, _ = _require_opensandbox_sdk()
+        options = OpenSandboxProviderOptions.from_mapping(spec.provider_options)
 
         kwargs: dict[str, Any] = {
             "env": spec.env,
             "metadata": spec.metadata,
             "resource": _resource_map(spec.resources),
-            "extensions": _spec_extensions(spec),
+            "extensions": self._resolve_extensions(options.extensions),
             "connection_config": self._connection_config(request_timeout_s=self._create.request_timeout_s),
         }
         if spec.image is not None:
             kwargs["image"] = spec.image
-        snapshot_id = spec.provider_options.get(PROVIDER_OPTION_SNAPSHOT_ID)
-        if snapshot_id is not None:
-            kwargs["snapshot_id"] = snapshot_id
+        if options.snapshot_id is not None:
+            kwargs["snapshot_id"] = options.snapshot_id
         if spec.ttl_s is not None:
             kwargs["timeout"] = timedelta(seconds=spec.ttl_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
-        platform = spec.provider_options.get(PROVIDER_OPTION_PLATFORM)
-        volumes = _spec_volumes(spec)
-        if platform is not None:
-            kwargs["platform"] = _to_platform_spec(platform)
-        if volumes is not None:
-            kwargs["volumes"] = _to_volumes(volumes)
+        if options.platform is not None:
+            kwargs["platform"] = _to_platform_spec(options.platform)
+        if options.volumes:
+            kwargs["volumes"] = _to_volumes(list(options.volumes))
         if self._create.skip_health_check:
             kwargs["skip_health_check"] = True
-        else:
-            skip_health_check = _provider_option_bool(spec.provider_options, PROVIDER_OPTION_SKIP_HEALTH_CHECK)
-            if skip_health_check is not None:
-                kwargs["skip_health_check"] = skip_health_check
+        elif options.skip_health_check is not None:
+            kwargs["skip_health_check"] = options.skip_health_check
 
         timeout_s = self._create.timeout_s
         if timeout_s is None and self._connection.request_timeout_s is not None:
@@ -790,8 +814,7 @@ async def _create_with_retries(
 
     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))
-        return await self._create_with_retries(spec)
+        return await self._create_with_retries(_normalize_spec(spec))
 
     async def status(self, handle: SandboxHandle) -> SandboxStatus:
         """Return the current OpenSandbox lifecycle status."""

tests/unit_tests/test_opensandbox_provider.py

@@ -190,11 +190,9 @@ def test_provider_validation_and_retry_helpers() -> None:
     with pytest.raises(ValueError, match="image_pull_policy"):
         opensandbox_provider.validate_image_pull_policy("Sometimes")
     with pytest.raises(TypeError, match="extensions"):
-        opensandbox_provider._spec_extensions(
-            SandboxSpec(image="image:tag", provider_options={"extensions": ["not", "a", "mapping"]})
-        )
+        opensandbox_provider.OpenSandboxProviderOptions.from_mapping({"extensions": ["not", "a", "mapping"]})
     with pytest.raises(TypeError, match="must be a bool"):
-        opensandbox_provider._provider_option_bool({"skip_health_check": "true"}, "skip_health_check")
+        opensandbox_provider.OpenSandboxProviderOptions.from_mapping({"skip_health_check": "true"})
 
     assert opensandbox_provider._resource_map(SandboxResources(cpu=2.0))["cpu"] == "2"
     assert opensandbox_provider._to_sandbox_status("starting") == SandboxStatus.STARTING
@@ -243,6 +241,38 @@ def test_provider_validation_and_retry_helpers() -> None:
     assert attrs["next_sleep_s"] == 0.5
 
 
+def test_provider_options_from_mapping() -> None:
+    options_cls = opensandbox_provider.OpenSandboxProviderOptions
+
+    assert options_cls.from_mapping(None) == options_cls()
+
+    parsed = options_cls.from_mapping(
+        {
+            "platform": {"os": "linux", "arch": "amd64"},
+            "snapshot_id": "snap-1",
+            "volumes": [{"name": "workspace"}],
+            "skip_health_check": True,
+            "extensions": {"imagePullPolicy": "Never"},
+        }
+    )
+    assert parsed.platform == {"os": "linux", "arch": "amd64"}
+    assert parsed.snapshot_id == "snap-1"
+    assert parsed.volumes == ({"name": "workspace"},)
+    assert parsed.skip_health_check is True
+    assert parsed.extensions == {"imagePullPolicy": "Never"}
+
+    with pytest.raises(ValueError, match="Unknown OpenSandbox provider option"):
+        options_cls.from_mapping({"bogus": 1})
+    with pytest.raises(TypeError, match="provider_options must be a mapping"):
+        options_cls.from_mapping(["not", "a", "mapping"])
+    with pytest.raises(TypeError, match="'platform' must be a mapping"):
+        options_cls.from_mapping({"platform": "linux/amd64"})
+    with pytest.raises(TypeError, match="'snapshot_id' must be a string"):
+        options_cls.from_mapping({"snapshot_id": 123})
+    with pytest.raises(TypeError, match="'volumes' must be a list of mappings"):
+        options_cls.from_mapping({"volumes": ["workspace"]})
+
+
 def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None:
     provider = opensandbox_provider.OpenSandboxProvider(
         connection={
@@ -265,14 +295,12 @@ def test_connection_config_and_image_policy(fake_opensandbox_sdk: None) -> None:
     short_timeout_config = provider._connection_config(request_timeout_s=3)
     assert short_timeout_config.kwargs["request_timeout"] == timedelta(seconds=3)
 
-    spec = SandboxSpec(image="image:tag", provider_options={"extensions": {"imagePullPolicy": "Never"}})
-    updated = provider._with_default_image_pull_policy(spec)
-    extensions = updated.provider_options["extensions"]
+    extensions = provider._resolve_extensions({"imagePullPolicy": "Never"})
     assert extensions["imagePullPolicy"] == "Never"
     assert 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
+    assert no_policy_provider._resolve_extensions({"imagePullPolicy": "Never"}) == {"imagePullPolicy": "Never"}
 
 
 async def test_exec_file_operations_and_reference_validation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: