oom related bug fix

args.py

@@ -56,6 +56,14 @@ def parse_simulation_args(argv: Sequence[str] | None = None) -> argparse.Namespa
             "requires the same destructive-operation safeguards as rebuild."
         ),
     )
+    parser.add_argument(
+        "--resume-clean-files-root",
+        default="",
+        help=(
+            "Launcher-visible results root required by RJob resume cleanup. "
+            "Only <root>/<job_id>/<session_id> directories are removed."
+        ),
+    )
     parser.add_argument(
         "--confirm-cloud-delete-job-id",
         default="",

core/data_manager/cloud_delete_guard.py

@@ -5,7 +5,7 @@
 import asyncio
 import logging
 import os
-from typing import Any, Dict, List
+from typing import Any
 
 
 log = logging.getLogger("core.data_manager.cloud_delete_guard")
@@ -39,8 +39,8 @@ async def preflight(
         operation: str,
         job_id: str,
         landing_filter: str,
-    ) -> List[Dict[str, Any]]:
-        """Verify the exact landing target before deletion."""
+    ) -> int:
+        """Verify the exact target and count matching rows using only their ids."""
         normalized_job_id = str(job_id or "").strip()
         if not normalized_job_id:
             raise CloudDestructiveOperationError(
@@ -52,20 +52,21 @@ async def preflight(
                 "could not both be resolved explicitly"
             )
 
-        landing_rows = _rows_as_dicts(await asyncio.to_thread(
+        profile = str(os.environ.get("WT_SDK_PROFILE") or "").strip().lower()
+        production = (
+            profile in {"prod", "production"}
+            or self.landing_table == "wind_tunnel_landing"
+        )
+        landing_row_count = _row_count(await asyncio.to_thread(
             self.client.query_data,
             filter_query=landing_filter,
             limit=None,
+            columns=["id"],
             partition=normalized_job_id,
             checkout_latest=True,
             deserialize_json=False,
             table=self.landing_table,
         ))
-        profile = str(os.environ.get("WT_SDK_PROFILE") or "").strip().lower()
-        production = (
-            profile in {"prod", "production"}
-            or self.landing_table == "wind_tunnel_landing"
-        )
         log.warning(
             "Cloud delete preflight: operation=%s profile=%s db_uri=%s "
             "landing_table=%s job_id=%s landing_rows=%d filter=%s",
@@ -74,10 +75,9 @@ async def preflight(
             self.db_uri,
             self.landing_table,
             normalized_job_id,
-            len(landing_rows),
+            landing_row_count,
             landing_filter,
         )
-
         if self.confirmed_job_id != normalized_job_id:
             raise CloudDestructiveOperationError(
                 f"{operation} refused for cloud job_id={normalized_job_id!r}; pass "
@@ -90,17 +90,15 @@ async def preflight(
                 f"{operation} refused for production target "
                 f"{self.landing_table!r}; --confirm-production is required"
             )
-        return landing_rows
+        return landing_row_count
 
 
-def _rows_as_dicts(value: Any) -> List[Dict[str, Any]]:
+def _row_count(value: Any) -> int:
     if value is None:
-        return []
-    if hasattr(value, "to_dict"):
-        try:
-            value = value.to_dict(orient="records")
-        except TypeError:
-            value = value.to_dict()
+        return 0
     if isinstance(value, dict):
-        return [dict(value)]
-    return [dict(row) for row in value]
+        return 1
+    try:
+        return len(value)
+    except TypeError:
+        return sum(1 for _ in value)

core/data_manager/manager.py

@@ -6,6 +6,7 @@
 from core.data_manager.contracts import EnvironmentQuery, SessionContext, SessionStepQuery
 from core.data_manager.strategy.base_strategy import StorageStrategy
 from core.data_manager.strategy_factory import StorageFactory
+from core.data_manager.upsert_batcher import SessionStepUpsertBatcher
 
 log = logging.getLogger("core.data_manager.manager")
 
@@ -19,6 +20,11 @@ def __init__(
         self,
         job_id: str,
         storage_type: str = "sqlite",
+        *,
+        enable_upsert_batching: bool = False,
+        upsert_batch_size: int = 100,
+        upsert_flush_interval: float = 10.0,
+        upsert_queue_size: int = 1000,
         **storage_config
     ):
         self.job_id = job_id
@@ -45,6 +51,17 @@ def __init__(
             log.error("%s Original Error: %s", error_msg, e)
             raise RuntimeError(error_msg) from e
 
+        self._upsert_batcher = (
+            SessionStepUpsertBatcher(
+                self._strategy.upsert_session_step_rows,
+                batch_size=upsert_batch_size,
+                flush_interval=upsert_flush_interval,
+                queue_size=upsert_queue_size,
+            )
+            if enable_upsert_batching
+            else None
+        )
+
     async def init(self) -> None:
         """Initialize the storage strategy"""
         await self._strategy.init()
@@ -93,6 +110,10 @@ async def get_environment_by_env_id(self, env_id: str) -> Optional[Dict]:
         """Retrieve one environment config by env_id."""
         return await self._strategy.get_environment_by_env_id(env_id)
 
+    async def clear_environment_cache(self, env_ids: List[str]) -> int:
+        """Clear backend-local environment cache entries."""
+        return await self._strategy.clear_environment_cache(env_ids)
+
     async def mark_environment_finished(self, env_id: str) -> int:
         """Mark one environment completed for this job."""
         updated = await self._strategy.update_environment_rows(
@@ -128,6 +149,24 @@ async def list_environment_rows(
             is_deleted=is_deleted,
         ))
 
+    async def list_environment_refs(
+        self,
+        *,
+        job_id: Optional[str] = None,
+        after_id: int = 0,
+        limit: Optional[int] = None,
+        finished: Optional[bool] = None,
+        is_deleted: Optional[bool] = None,
+    ) -> List[Dict[str, Any]]:
+        """Query only the environment identity fields needed by resume cleanup."""
+        return await self._strategy.list_environment_refs(EnvironmentQuery(
+            job_id=job_id or self.job_id,
+            after_id=max(0, int(after_id)),
+            limit=None if limit is None else max(0, int(limit)),
+            finished=finished,
+            is_deleted=is_deleted,
+        ))
+
     async def insert_environment_rows(self, rows: List[Dict[str, Any]]) -> List[str]:
         """Insert environment rows through the configured DAO."""
         normalized = []
@@ -312,6 +351,8 @@ async def upsert_session_step_rows(
             seen_keys.add(key)
             item["meta_json"] = _metadata_object(item.get("meta_json"))
             normalized.append(item)
+        if self._upsert_batcher is not None:
+            return await self._upsert_batcher.submit(normalized)
         return await self._strategy.upsert_session_step_rows(normalized)
 
     async def mark_records_completed(self, record_ids: List[str]) -> int:
@@ -442,7 +483,11 @@ async def mark_latest_session_completed(
 
     async def close(self) -> None:
         """Close the storage strategy"""
-        await self._strategy.close()
+        try:
+            if self._upsert_batcher is not None:
+                await self._upsert_batcher.close()
+        finally:
+            await self._strategy.close()
 
     async def fetch_done_steps_with_context(
         self,

core/data_manager/strategy/base_strategy.py

@@ -63,11 +63,20 @@ async def get_environment_by_env_id(self, env_id: str) -> Optional[Dict[str, Any
         """Retrieve one active environment by env_id."""
         pass
 
+    async def clear_environment_cache(self, env_ids: List[str]) -> int:
+        """Clear backend-local environment cache entries, if any."""
+        return 0
+
     @abstractmethod
     async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, Any]]:
         """List environment rows using backend-neutral filters."""
         pass
 
+    @abstractmethod
+    async def list_environment_refs(self, query: EnvironmentQuery) -> List[Dict[str, Any]]:
+        """List only id/env_id/env_name for lightweight workflow coordination."""
+        pass
+
     @abstractmethod
     async def insert_environment_rows(self, rows: List[Dict[str, Any]]) -> List[str]:
         """Insert environment rows and return their env_ids."""