feat(sandbox): expire failed Ray submissions immediately

Author: zhangjaycee

rock/actions/sandbox/sandbox_info.py

@@ -29,8 +29,8 @@ class SandboxInfo(TypedDict, total=False):
     archive_time: str
     auto_transition_state: State
     auto_transition_time: str
-    auto_archive_seconds: int
-    auto_delete_seconds: int
+    auto_archive_seconds: int | None
+    auto_delete_seconds: int | None
     archive_prefix: str
     registry_namespace: str
     extended_params: dict[str, str]

rock/sandbox/sandbox_manager.py

@@ -6,7 +6,7 @@
 
 from fastapi import UploadFile
 
-from rock import env_vars
+from rock import InternalServerRockError, env_vars
 from rock.actions import (
     BashObservation,
     CloseBashSessionResponse,
@@ -182,8 +182,21 @@ async def start_async(
         with StageTimer("startup_timing", f"[{sandbox_id}] Meta store policy update", logger):
             await self._meta_store.update(sandbox_id, sandbox_info)
 
-        with StageTimer("startup_timing", f"[{sandbox_id}] Operator submit", logger):
-            submitted_info = await self._operator.submit(docker_deployment_config, user_info)
+        try:
+            with StageTimer("startup_timing", f"[{sandbox_id}] Operator submit", logger):
+                submitted_info = await self._operator.submit(docker_deployment_config, user_info)
+        except Exception as e:
+            immediate_timeout_info = SandboxTimeoutHelper.make_timeout_info(0)
+            sandbox_info["auto_archive_seconds"] = None
+            sandbox_info["auto_delete_seconds"] = 0
+            try:
+                await self._meta_store.update(sandbox_id, sandbox_info)
+                await self._meta_store.update_timeout(sandbox_id, immediate_timeout_info)
+            except Exception:
+                logger.exception("[%s] failed to update immediate cleanup metadata", sandbox_id)
+            raise InternalServerRockError(
+                f"Sandbox {sandbox_id} submission failed and has been marked for immediate deletion: {e}"
+            ) from e
         create_time = sandbox_info.get("create_time")
         sandbox_info.update(submitted_info)
         await self._build_sandbox_info_metadata(sandbox_info, user_info, cluster_info)

rock/sandbox/sandbox_statemachine.py

@@ -67,8 +67,9 @@ def apply_effective_auto_transition_policy(info: dict[str, Any], config: DockerD
 
     @staticmethod
     def _resolve_user_seconds(info: dict[str, Any], field: str) -> int | None:
-        raw = info.get(field)
-        if raw is None:
+        if field in info:
+            raw = info[field]
+        else:
             spec = info.get("spec") or {}
             raw = spec.get(field)
         return SandboxLifecycleHelper._coerce_seconds(raw)

tests/unit/sandbox/test_sandbox_statemachine.py

@@ -18,6 +18,25 @@
 from rock.config import AutoTransitionConfig
 from rock.sandbox.sandbox_statemachine import SandboxLifecycleHelper, SandboxStateMachine
 
+
+class TestLifecyclePolicyResolution:
+    def test_missing_metadata_field_falls_back_to_spec(self):
+        info = {"spec": {"auto_archive_seconds": 600, "auto_delete_seconds": 1200}}
+
+        assert SandboxLifecycleHelper.resolve_auto_archive_seconds(info) == 600
+        assert SandboxLifecycleHelper.resolve_auto_delete_seconds(info) == 1200
+
+    def test_explicit_none_disables_spec_policy(self):
+        info = {
+            "auto_archive_seconds": None,
+            "auto_delete_seconds": None,
+            "spec": {"auto_archive_seconds": 600, "auto_delete_seconds": 1200},
+        }
+
+        assert SandboxLifecycleHelper.resolve_auto_archive_seconds(info) is None
+        assert SandboxLifecycleHelper.resolve_auto_delete_seconds(info) is None
+
+
 # ---------------------------------------------------------------------------
 # Transitions
 # ---------------------------------------------------------------------------

tests/unit/sandbox/test_sandbox_transitions.py

@@ -8,11 +8,13 @@
 
 import pytest
 
+from rock import InternalServerRockError
 from rock.actions.sandbox.response import State
 from rock.admin.proto.response import SandboxStartResponse
 from rock.common.constants import StopReason
 from rock.config import AutoTransitionConfig, SandboxLifecycleConfig
 from rock.sandbox.sandbox_manager import SandboxManager
+from rock.sandbox.sandbox_statemachine import SandboxLifecycleHelper
 from rock.sdk.common.exceptions import BadRequestRockError
 
 
@@ -428,22 +430,31 @@ async def assert_pending_exists(*args, **kwargs):
         assert mock_meta_store.update.await_count == 2
 
     @pytest.mark.asyncio
-    async def test_submit_failure_keeps_pending_record_with_spec_timeout_and_effective_policy(
+    async def test_submit_failure_marks_pending_record_for_immediate_deletion(
         self, mgr_start, mock_meta_store, mock_operator, mock_docker_config
     ):
-        mock_docker_config.auto_delete_seconds = 7200
+        mock_docker_config.auto_archive_seconds = 600
         mock_operator.submit.side_effect = RuntimeError("ray submit failed")
 
-        with pytest.raises(RuntimeError, match="ray submit failed"):
+        with pytest.raises(
+            InternalServerRockError,
+            match="Sandbox sb-1 submission failed and has been marked for immediate deletion: ray submit failed",
+        ):
             await mgr_start.start_async(MagicMock(image="python:3.11"))
 
         mock_meta_store.create.assert_awaited_once()
         assert mock_meta_store.create.await_args.kwargs["timeout_info"]
-        assert mgr_start.created_spec["auto_delete_seconds"] == 7200
-        mock_meta_store.update.assert_awaited_once()
-        pending_info = mock_meta_store.update.await_args.args[1]
+        assert mgr_start.created_spec["auto_archive_seconds"] == 600
+        assert mock_meta_store.update.await_count == 2
+        pending_info = mock_meta_store.update.await_args_list[-1].args[1]
         assert pending_info["state"] == State.PENDING
-        assert pending_info["auto_delete_seconds"] == 3600
+        assert pending_info["auto_delete_seconds"] == 0
+        assert pending_info["auto_archive_seconds"] is None
+        assert SandboxLifecycleHelper.resolve_auto_archive_seconds(pending_info) is None
+        assert SandboxLifecycleHelper.resolve_auto_delete_seconds(pending_info) == 0
+        mock_meta_store.update_timeout.assert_awaited_once()
+        timeout_info = mock_meta_store.update_timeout.await_args.args[1]
+        assert timeout_info["auto_clear_time"] == "0"
 
     @pytest.mark.asyncio
     async def test_uses_cluster_auto_delete_when_user_policy_is_unspecified(