diff --git a/args.py b/args.py index 4d83145..6dd283b 100644 --- a/args.py +++ b/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 // directories are removed." + ), + ) parser.add_argument( "--confirm-cloud-delete-job-id", default="", diff --git a/core/data_manager/cloud_delete_guard.py b/core/data_manager/cloud_delete_guard.py index 8a915ef..a08ff12 100644 --- a/core/data_manager/cloud_delete_guard.py +++ b/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) diff --git a/core/data_manager/manager.py b/core/data_manager/manager.py index c1a90f8..534bb82 100644 --- a/core/data_manager/manager.py +++ b/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, diff --git a/core/data_manager/strategy/base_strategy.py b/core/data_manager/strategy/base_strategy.py index 7682226..c34c2ed 100644 --- a/core/data_manager/strategy/base_strategy.py +++ b/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.""" diff --git a/core/data_manager/strategy/cloud_strategy_impl.py b/core/data_manager/strategy/cloud_strategy_impl.py index 119a4aa..16a62b7 100644 --- a/core/data_manager/strategy/cloud_strategy_impl.py +++ b/core/data_manager/strategy/cloud_strategy_impl.py @@ -1,5 +1,6 @@ import asyncio from importlib import import_module +import inspect import json import logging import os @@ -211,21 +212,6 @@ def _meta_json_object(meta_json: Any) -> Dict[str, Any]: return meta -def _truthy_bool(value: Any) -> bool: - if isinstance(value, bool): - return value - if value is None: - return False - if isinstance(value, str): - return value.strip().lower() in {"1", "true", "yes", "y", "on"} - try: - if value != value: - return False - except Exception: - pass - return bool(value) - - def _response_text(value: Any) -> str: """Extract training text without discarding non-text model output.""" if value is None: @@ -524,19 +510,18 @@ async def get_environment_by_env_id(self, env_id: str) -> Optional[Dict[str, Any self._env_configs[str(config["env_id"])] = config return dict(config) + async def clear_environment_cache(self, env_ids: List[str]) -> int: + """Clear cached environment configs without changing persisted state.""" + removed = 0 + for env_id in env_ids: + if self._env_configs.pop(env_id, None) is not None: + removed += 1 + return removed + async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, Any]]: """Read environment rows from the authoritative config store.""" await self.init() - clauses = [] - if query.job_id: - clauses.append(f"job_id = '{_escape_sql_literal(query.job_id)}'") - if query.env_id: - clauses.append(f"env_id = '{_escape_sql_literal(query.env_id)}'") - if query.after_id: - clauses.append(f"id > {int(query.after_id)}") - if query.finished is not None: - clauses.append(f"finished = {str(query.finished).lower()}") - filter_query = " AND ".join(clauses) + filter_query = self._environment_filter_query(query) page_size = max(100, query.limit or 1000) effective_offset = max(0, query.offset) normalized: List[Dict[str, Any]] = [] @@ -558,8 +543,6 @@ async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, "cloud environment pagination requires EnvConfigManager " "to return the physical id column" ) - if query.is_deleted is not None and _truthy_bool(row.get("is_deleted")) != query.is_deleted: - continue env_id = str(row.get("env_id") or "") if env_id: self._env_configs[env_id] = row @@ -571,6 +554,59 @@ async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, break return normalized + async def list_environment_refs(self, query: EnvironmentQuery) -> List[Dict[str, Any]]: + """Read resume identity fields without populating the full config cache.""" + await self.init() + # wt_sdk v0.6.2 get_env_configs materializes every column before it + # paginates. Use its projection-capable query boundary for resume. + filter_table = getattr(self.env_manager, "_filter_table", None) + if not callable(filter_table): + raise RuntimeError( + "Cloud resume requires the projection-capable EnvConfigManager " + "provided by wt-data-platform-sdk>=0.6.2" + ) + kwargs: Dict[str, Any] = { + "query": self._environment_filter_query(query), + "limit": query.limit if not query.after_id and not query.offset else None, + "columns": ["id", "env_id", "env_name"], + } + if "checkout_latest" in inspect.signature(filter_table).parameters: + kwargs["checkout_latest"] = True + result = await asyncio.to_thread(filter_table, **kwargs) + rows = result.to_dict(orient="records") if hasattr(result, "to_dict") else result + refs: List[Dict[str, Any]] = [] + for config in rows or []: + row = dict(config) + if row.get("id") is None: + raise RuntimeError( + "cloud environment pagination requires EnvConfigManager " + "to return the physical id column" + ) + refs.append({ + "id": row["id"], + "env_id": row.get("env_id"), + "env_name": row.get("env_name"), + }) + refs.sort(key=lambda row: int(row["id"])) + if query.offset: + refs = refs[query.offset:] + if query.limit is not None: + refs = refs[:query.limit] + return refs + + @staticmethod + def _environment_filter_query(query: EnvironmentQuery) -> str: + clauses = [] + if query.job_id: + clauses.append(f"job_id = '{_escape_sql_literal(query.job_id)}'") + if query.env_id: + clauses.append(f"env_id = '{_escape_sql_literal(query.env_id)}'") + if query.after_id: + clauses.append(f"id > {int(query.after_id)}") + if query.finished is not None: + clauses.append(f"finished = {str(query.finished).lower()}") + return " AND ".join(clauses) + async def insert_environment_rows(self, rows: List[Dict[str, Any]]) -> List[str]: await self.init() if not rows: @@ -588,7 +624,6 @@ async def insert_environment_rows(self, rows: List[Dict[str, Any]]) -> List[str] "image": str(row.get("image") or ""), "group_id": str(row.get("group_id") or ""), "finished": bool(row.get("finished", False)), - "is_deleted": bool(row.get("is_deleted", False)), "created_at": int(row.get("created_at") or time.time()), } configs.append(config) @@ -602,7 +637,7 @@ async def update_environment_rows( updates: Dict[str, Any], ) -> int: await self.init() - allowed = {"env_name", "env_params", "image", "group_id", "finished", "is_deleted"} + allowed = {"env_name", "env_params", "image", "group_id", "finished"} unknown = set(updates) - allowed if unknown: raise ValueError(f"Unknown environment update fields: {sorted(unknown)}") @@ -611,8 +646,11 @@ async def update_environment_rows( for row in rows: env_id = str(row.get("env_id") or "") if env_id and await asyncio.to_thread(self.env_manager.update_config, env_id, updates): - cached = self._env_configs.setdefault(env_id, dict(row)) - cached.update(updates) + if updates.get("finished") is True: + self._env_configs.pop(env_id, None) + else: + cached = self._env_configs.setdefault(env_id, dict(row)) + cached.update(updates) updated += 1 return updated @@ -622,7 +660,7 @@ async def _preflight_destructive_delete( operation: str, job_id: str, landing_filter: str, - ) -> List[Dict[str, Any]]: + ) -> int: guard = CloudDeleteGuard( client=self.client, db_uri=self.db_url, @@ -650,13 +688,12 @@ async def delete_session_step_rows(self, query: SessionStepQuery) -> int: if job_id: clauses.insert(0, f"job_id = '{_escape_sql_literal(job_id)}'") landing_filter = " AND ".join(clauses) - rows = await self._preflight_destructive_delete( + selected = await self._preflight_destructive_delete( operation="delete_session_step_rows", job_id=job_id, landing_filter=landing_filter, ) - await asyncio.to_thread(self.client.delete_landing, landing_filter) - return len(rows) + return await self._delete_landing_rows(landing_filter, selected) session_ids = list(query.session_ids) if query.session_id: session_ids.append(query.session_id) @@ -672,13 +709,29 @@ async def delete_session_step_rows(self, query: SessionStepQuery) -> int: if not clauses: raise ValueError("job_id or session_ids is required for cloud deletion") landing_filter = " AND ".join(clauses) - rows = await self._preflight_destructive_delete( + selected = await self._preflight_destructive_delete( operation="delete_session_step_rows", job_id=job_id, landing_filter=landing_filter, ) - await asyncio.to_thread(self.client.delete_landing, landing_filter) - return len(rows) + return await self._delete_landing_rows(landing_filter, selected) + + async def _delete_landing_rows(self, landing_filter: str, selected: int) -> int: + if not selected: + return 0 + result = await asyncio.to_thread(self.client.delete_landing, landing_filter) + deleted = ( + int(result) + if isinstance(result, int) and not isinstance(result, bool) + else selected + ) + log.info( + "Cloud landing delete completed: selected_rows=%d deleted_rows=%d filter=%s", + selected, + deleted, + landing_filter, + ) + return deleted async def delete_job_rows(self, job_id: str) -> None: await self.init() diff --git a/core/data_manager/strategy/sqlite_strategy_impl.py b/core/data_manager/strategy/sqlite_strategy_impl.py index 1d9836f..106afbd 100644 --- a/core/data_manager/strategy/sqlite_strategy_impl.py +++ b/core/data_manager/strategy/sqlite_strategy_impl.py @@ -251,6 +251,21 @@ async def get_environment_by_env_id(self, env_id: str) -> Optional[Dict[str, Any async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, Any]]: await self.init() + rows = self._environment_query(query) + rows = rows.order_by("id").offset(query.offset) + if query.limit is not None: + rows = rows.limit(query.limit) + return [self._environment_to_dict(env) for env in await rows] + + async def list_environment_refs(self, query: EnvironmentQuery) -> List[Dict[str, Any]]: + await self.init() + rows = self._environment_query(query).order_by("id").offset(query.offset) + if query.limit is not None: + rows = rows.limit(query.limit) + return list(await rows.values("id", "env_id", "env_name")) + + @staticmethod + def _environment_query(query: EnvironmentQuery): rows = JobEnvironment.all() if query.job_id: rows = rows.filter(job_id=query.job_id) @@ -262,10 +277,7 @@ async def list_environment_rows(self, query: EnvironmentQuery) -> List[Dict[str, rows = rows.filter(finished=query.finished) if query.is_deleted is not None: rows = rows.filter(is_deleted=query.is_deleted) - rows = rows.order_by("id").offset(query.offset) - if query.limit is not None: - rows = rows.limit(query.limit) - return [self._environment_to_dict(env) for env in await rows] + return rows async def insert_environment_rows(self, rows: List[Dict[str, Any]]) -> List[str]: await self.init() @@ -326,7 +338,14 @@ async def delete_session_step_rows(self, query: SessionStepQuery) -> int: rows = rows.filter(record_id=query.record_id) if query.record_ids: rows = rows.filter(record_id__in=query.record_ids) - return await rows.delete() + selected = await rows.count() + log.info( + "Session-step delete preflight: job_id=%s sessions=%d rows=%d", + query.job_id or self.job_id, + len(query.session_ids) + int(bool(query.session_id)), + selected, + ) + return await rows.delete() if selected else 0 async def delete_job_rows(self, job_id: str) -> None: await self.init() diff --git a/core/data_manager/upsert_batcher.py b/core/data_manager/upsert_batcher.py new file mode 100644 index 0000000..502449a --- /dev/null +++ b/core/data_manager/upsert_batcher.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Dict, List + + +log = logging.getLogger("core.data_manager.upsert_batcher") + +UpsertWriter = Callable[[List[Dict[str, Any]]], Awaitable[List[str]]] +_STOP = object() + + +@dataclass(slots=True) +class _UpsertRequest: + rows: List[Dict[str, Any]] + future: asyncio.Future[List[str]] + + +class SessionStepUpsertBatcher: + """Serialize complete-row upserts and acknowledge them after persistence.""" + + def __init__( + self, + writer: UpsertWriter, + *, + flush_interval: float = 10.0, + batch_size: int = 100, + queue_size: int = 1000, + ) -> None: + self._writer = writer + self._flush_interval = max(0.0, float(flush_interval)) + self._batch_size = max(1, int(batch_size)) + self._queue: asyncio.Queue[_UpsertRequest | object] = asyncio.Queue( + maxsize=max(1, int(queue_size)), + ) + self._state_lock = asyncio.Lock() + self._task: asyncio.Task[None] | None = None + self._closing = False + + async def submit(self, rows: List[Dict[str, Any]]) -> List[str]: + if not rows: + return [] + + future: asyncio.Future[List[str]] = asyncio.get_running_loop().create_future() + request = _UpsertRequest(rows=rows, future=future) + async with self._state_lock: + if self._closing: + raise RuntimeError("session-step upsert batcher is closed") + if self._task is None: + self._task = asyncio.create_task( + self._run(), + name="session-step-upsert-batcher", + ) + await self._queue.put(request) + return await future + + async def close(self) -> None: + async with self._state_lock: + if not self._closing: + self._closing = True + if self._task is not None: + await self._queue.put(_STOP) + task = self._task + if task is not None: + await task + + async def _run(self) -> None: + while True: + item = await self._queue.get() + if item is _STOP: + self._queue.task_done() + return + + requests = [item] + row_count = len(item.rows) + stop_after_flush = False + deadline = asyncio.get_running_loop().time() + self._flush_interval + + while row_count < self._batch_size: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + break + try: + item = await asyncio.wait_for(self._queue.get(), timeout=remaining) + except asyncio.TimeoutError: + break + if item is _STOP: + self._queue.task_done() + stop_after_flush = True + break + requests.append(item) + row_count += len(item.rows) + + try: + await self._flush(requests) + finally: + for _ in requests: + self._queue.task_done() + + if stop_after_flush: + return + + async def _flush(self, requests: List[_UpsertRequest]) -> None: + # Complete rows use last-enqueued-wins semantics for duplicate keys. + rows_by_key: Dict[tuple[str, str], Dict[str, Any]] = {} + for request in requests: + for row in request.rows: + key = (str(row["job_id"]), str(row["record_id"])) + rows_by_key[key] = row + + try: + await self._writer(list(rows_by_key.values())) + except Exception as exc: + log.exception( + "Session-step batch upsert failed: requests=%d rows=%d", + len(requests), + len(rows_by_key), + ) + for request in requests: + if not request.future.done(): + request.future.set_exception(exc) + return + + for request in requests: + if not request.future.done(): + request.future.set_result([ + str(row["record_id"]) + for row in request.rows + ]) diff --git a/core/data_manager/yaml_aggregator.py b/core/data_manager/yaml_aggregator.py index 724646f..b68a910 100644 --- a/core/data_manager/yaml_aggregator.py +++ b/core/data_manager/yaml_aggregator.py @@ -15,6 +15,7 @@ _insert_tasks: Set[asyncio.Task] = set() _job_db_processing_done: Dict[str, bool] = {} +_RESUME_CLEANUP_BATCH_SIZE = 100 def set_job_db_processing_done(job_id: str, done: bool) -> None: @@ -129,36 +130,40 @@ async def sync_configs_to_db( *, rebuild_table: bool = False, resume: bool = False, -) -> None: +) -> List[Dict[str, Any]]: """Synchronize configs without leaking a connection or backend client.""" if storage_type not in {"sqlite", "cloud"}: raise ValueError(f"Unknown storage type: {storage_type}") if rebuild_table and resume: raise ValueError("--rebuild-table and --resume cannot be used together") + job_id = str(data_manager.job_id or "").strip() + if resume and not job_id: + raise ValueError("resume requires an explicit job_id") await data_manager.init() - job_id = data_manager.job_id set_job_db_processing_done(job_id, False) try: - existing = await data_manager.list_environment_rows(job_id=job_id, limit=1) - if existing and resume: - unfinished = await data_manager.list_environment_rows( + existing = await data_manager.list_environment_refs(job_id=job_id, limit=1) + if resume: + if not existing: + raise RuntimeError( + f"cannot resume job_id={job_id!r}: no existing environment rows found" + ) + + unfinished = await data_manager.list_environment_refs( job_id=job_id, finished=False, + is_deleted=False, ) - unfinished_ids = [ - str(row.get("env_id") or "") - for row in unfinished - if row.get("env_id") - ] - if unfinished_ids: - await data_manager.delete_session_step_rows( - job_id=job_id, - session_ids=unfinished_ids, - ) + set_job_db_processing_done(job_id, True) - log.info("Resuming existing job_id=%s; finished environments will be skipped", job_id) - return + log.info( + "Resume environment preflight: job_id=%s unfinished=%d; " + "finished environments will be skipped", + job_id, + len(unfinished), + ) + return unfinished if existing and not rebuild_table: raise RuntimeError( f"job_id={job_id!r} already exists; use --resume to continue it " @@ -188,11 +193,52 @@ async def sync_configs_to_db( len(remaining), job_id, ) + return [] except Exception: set_job_db_processing_done(job_id, True) raise +async def delete_resume_session_steps( + data_manager: Any, + *, + job_id: str, + environment_rows: List[Dict[str, Any]], +) -> int: + """Count and delete unfinished-session trajectories in bounded batches.""" + deleted_steps = 0 + batch_number = 0 + for index in range(0, len(environment_rows), _RESUME_CLEANUP_BATCH_SIZE): + batch = environment_rows[index:index + _RESUME_CLEANUP_BATCH_SIZE] + session_ids = list(dict.fromkeys( + str(row.get("env_id") or "").strip() + for row in batch + if str(row.get("env_id") or "").strip() + )) + if not session_ids: + continue + batch_number += 1 + deleted = await data_manager.delete_session_step_rows( + job_id=job_id, + session_ids=session_ids, + ) + deleted_steps += deleted + log.info( + "Resume trajectory cleanup batch: job_id=%s batch=%d sessions=%d rows=%d", + job_id, + batch_number, + len(session_ids), + deleted, + ) + log.info( + "Resume trajectory cleanup completed: job_id=%s sessions=%d rows=%d", + job_id, + len(environment_rows), + deleted_steps, + ) + return deleted_steps + + def _expand_environment_rows(job_id: str, yaml_configs: List[Dict]) -> List[Dict[str, Any]]: rows: List[Dict[str, Any]] = [] for config in yaml_configs: diff --git a/gateway/storage.py b/gateway/storage.py index 0f1777a..de442b9 100644 --- a/gateway/storage.py +++ b/gateway/storage.py @@ -194,36 +194,17 @@ async def _query_session_environment(self, session_id: str) -> _SessionEnvironme return resolved async def _load_environment_config(self, env_id: str) -> dict[str, Any] | None: - lookup = getattr(self.data_manager, "get_environment_by_env_id", None) - if callable(lookup): - maybe_environment = lookup(env_id) - environment = await maybe_environment if inspect.isawaitable(maybe_environment) else maybe_environment - return environment if isinstance(environment, dict) else None - - get_all = getattr(self.data_manager, "get_all_environments", None) - if not callable(get_all): - return None - - maybe_environments = get_all() - environments = await maybe_environments if inspect.isawaitable(maybe_environments) else maybe_environments - if not isinstance(environments, list): - return None - - for environment in environments: - if not isinstance(environment, dict): - continue - if str(environment.get("env_id") or "") == env_id: - return environment - return None + environment = await self.data_manager.get_environment_by_env_id(env_id) + return environment if isinstance(environment, dict) else None async def count_environment_rows(self, job_id: str) -> int: """Return the number of unfinished, active environment rows for a job.""" - rows = await self.data_manager.list_environment_rows( + refs = await self.data_manager.list_environment_refs( job_id=job_id, finished=False, is_deleted=False, ) - return len(rows) + return len(refs) @staticmethod def _environment_from_mapping(environment: dict[str, Any] | None) -> _SessionEnvironment | None: @@ -405,14 +386,6 @@ async def record_inference_steps_batch( trace.emit_summary(status="failed", error_type=type(exc).__name__, error=str(exc)) raise - async def flush_session(self, binding: GatewaySessionBinding) -> None: - """Flush any DAO buffer and force a latest-snapshot read for one session.""" - await self.data_manager.list_session_steps( - binding.session_id, - job_id=binding.job_id, - checkout_latest=True, - ) - async def clear_session_cache(self, session_ids: list[str]) -> int: targets = set(session_ids) async with self._lock: @@ -429,7 +402,8 @@ async def clear_session_cache(self, session_ids: list[str]) -> int: self._environments.pop(session_id, None) for session_id in patched: self._patched_environment_sessions.discard(session_id) - return len(session_keys) + len(environments) + len(patched) + removed = len(session_keys) + len(environments) + len(patched) + return removed + await self.data_manager.clear_environment_cache(list(targets)) async def close(self) -> None: log.info("Gateway storage close begin") diff --git a/gateway/telemetry.py b/gateway/telemetry.py index 5268118..ee83100 100644 --- a/gateway/telemetry.py +++ b/gateway/telemetry.py @@ -200,9 +200,6 @@ async def wait_for_session_flush(self, binding: GatewaySessionBinding) -> None: _SessionFlushBarrier(binding.session_id, future) ) await future - flush_session = getattr(self.storage, "flush_session", None) - if callable(flush_session): - await flush_session(binding) async def latest_success_step_id(self, session_id: str, model: str) -> int | None: async with self._lock: diff --git a/manager/episode_common.py b/manager/episode_common.py index 7ab99fd..113b359 100644 --- a/manager/episode_common.py +++ b/manager/episode_common.py @@ -101,30 +101,13 @@ def containerize_local_gateway_url(url: str) -> str: def result_artifact_path(request: SimulationStartRequest) -> str: - env_params = request.env_params if isinstance(request.env_params, dict) else {} - return _result_artifact_path(request.job_id, request.session_id, env_params) - + return _result_artifact_path(request.job_id, request.session_id) -def _result_artifact_path(job_id: str, session_id: str, env_params: Dict[str, Any]) -> str: - dataset = env_params.get("dataset") if isinstance(env_params.get("dataset"), dict) else {} - explicit = first_text( - dataset.get("safactory_result_path"), - env_params.get("safactory_result_path"), - ) - if explicit: - return explicit - - root = first_text( - dataset.get("safactory_results_root"), - env_params.get("safactory_results_root"), - dataset.get("results_root"), - env_params.get("results_root"), - DEFAULT_RESULT_ROOT, - ).rstrip("/") +def _result_artifact_path(job_id: str, session_id: str) -> str: return "/".join( [ - root or DEFAULT_RESULT_ROOT, + DEFAULT_RESULT_ROOT, safe_path_part(job_id), safe_path_part(session_id), RESULT_FILENAME, @@ -144,13 +127,8 @@ def result_session_dir_candidates( env_params: Dict[str, Any] | None = None, ) -> list[Path]: """Return launcher-visible candidates for results//.""" - params = env_params if isinstance(env_params, dict) else {} - dataset = params.get("dataset") if isinstance(params.get("dataset"), dict) else {} - explicit = first_text( - dataset.get("safactory_result_path"), - params.get("safactory_result_path"), - ) - artifact = _result_artifact_path(job_id, session_id, params) + del env_params + artifact = _result_artifact_path(job_id, session_id) candidates: list[Path] = [] def add(path: Path) -> None: @@ -158,7 +136,7 @@ def add(path: Path) -> None: candidates.append(path) for path in _result_path_candidates(artifact): - add(path if explicit else path.parent) + add(path.parent) add(Path.cwd() / "results" / safe_path_part(job_id) / safe_path_part(session_id)) return candidates diff --git a/manager/resume_cleanup.py b/manager/resume_cleanup.py index 0f207e8..7de8cdd 100644 --- a/manager/resume_cleanup.py +++ b/manager/resume_cleanup.py @@ -4,29 +4,77 @@ import logging import shutil from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict, List, Sequence from clusters.rjob_cluster import RJobClusterBackend -from .episode_common import result_session_dir_candidates +from .episode_common import safe_path_part log = logging.getLogger("manager.resume_cleanup") +def select_generated_resume_environments( + *, + job_id: str, + results_root: Path, + environment_rows: Sequence[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Return unique unfinished environments that have generated result paths.""" + root = _resolve_results_root(results_root) + targets: Dict[Path, Dict[str, Any]] = {} + for row in environment_rows: + session_id = str(row.get("env_id") or "").strip() + if not session_id: + continue + result_path = root / safe_path_part(job_id) / safe_path_part(session_id) + if not (result_path.exists() or result_path.is_symlink()): + continue + existing = targets.get(result_path) + if existing is not None: + existing_session_id = str(existing.get("env_id") or "").strip() + if existing_session_id != session_id: + raise RuntimeError( + "resume result path collision: " + f"{existing_session_id!r} and {session_id!r} map to {result_path}" + ) + continue + target = dict(row) + target["job_id"] = job_id + target["env_id"] = session_id + targets[result_path] = target + return list(targets.values()) + + async def cleanup_resume_artifacts( *, job_id: str, model: str, data_manager: Any, manager_cfg: Dict[str, Any], + results_root: Path, + environment_rows: Sequence[Dict[str, Any]] | None = None, rjob_backend: RJobClusterBackend | None = None, ) -> List[Path]: """Remove stale RJobs and result paths for unfinished resume sessions.""" - rows = await data_manager.list_environment_rows( + rows = ( + list(environment_rows) + if environment_rows is not None + else await _load_environment_refs(data_manager, job_id=job_id) + ) + root = _resolve_results_root(results_root) + targets = select_generated_resume_environments( job_id=job_id, - finished=False, - is_deleted=False, + results_root=root, + environment_rows=rows, ) + if not targets: + log.info( + "resume cleanup skipped: job_id=%s unfinished=%d generated=0", + job_id, + len(rows), + ) + return [] + owned_backend = rjob_backend is None backend = rjob_backend or RJobClusterBackend( cluster_cfg=dict(manager_cfg.get("cluster") or {}) @@ -34,23 +82,8 @@ async def cleanup_resume_artifacts( removed: List[Path] = [] try: - for row in rows: + for row in targets: session_id = str(row.get("env_id") or "").strip() - if not session_id: - continue - env_params = row.get("env_params") if isinstance(row.get("env_params"), dict) else {} - result_paths = [ - path - for path in result_session_dir_candidates( - job_id=job_id, - session_id=session_id, - env_params=env_params, - ) - if path.exists() or path.is_symlink() - ] - if not result_paths: - continue - agent_name = str(row.get("env_name") or "").strip() if not agent_name: raise RuntimeError( @@ -63,29 +96,53 @@ async def cleanup_resume_artifacts( job_id=job_id, session_id=session_id, ) - for path in result_paths: - await asyncio.to_thread(_remove_result_path, path) - removed.append(path) + result_path = ( + root + / safe_path_part(job_id) + / safe_path_part(session_id) + ) + if result_path.exists() or result_path.is_symlink(): + await asyncio.to_thread(_remove_result_path, result_path) + removed.append(result_path) log.info( - "resume cleanup completed: job_id=%s session_id=%s rjobs=%s result_paths=%s", + "resume cleanup completed: job_id=%s session_id=%s rjobs=%s result_path=%s", job_id, session_id, cleaned_jobs, - [str(path) for path in result_paths], + str(result_path), ) finally: if owned_backend: await backend.close() log.info( - "resume result preflight completed: job_id=%s unfinished=%d removed_paths=%d", + "resume artifact cleanup completed: job_id=%s root=%s unfinished=%d " + "generated=%d removed_paths=%d", job_id, + root, len(rows), + len(targets), len(removed), ) return removed +async def _load_environment_refs(data_manager: Any, *, job_id: str) -> List[Dict[str, Any]]: + """Lightweight path for direct callers; the launcher reuses rows it already read.""" + return await data_manager.list_environment_refs( + job_id=job_id, + finished=False, + is_deleted=False, + ) + + +def _resolve_results_root(results_root: Path) -> Path: + root = Path(results_root).expanduser().resolve(strict=False) + if root.parent == root or not root.is_dir(): + raise ValueError(f"invalid resume results root: {root}") + return root + + def _remove_result_path(path: Path) -> None: if path.is_symlink() or path.is_file(): path.unlink() diff --git a/manager/simulation_config.py b/manager/simulation_config.py index d9a7ba7..0fccc59 100644 --- a/manager/simulation_config.py +++ b/manager/simulation_config.py @@ -190,6 +190,28 @@ def load_simulation_run_config(args: Any) -> SimulationRunConfig: if mode not in {"docker", "rjob", "sandbox"}: raise ValueError(f"Unsupported simulation mode: {mode!r}") + requested_job_id = str(args.job_id or "").strip() + resume = bool(getattr(args, "resume", False)) + if resume and not requested_job_id: + raise ValueError("--resume requires an explicit --job-id") + resume_clean_files_root = str( + getattr(args, "resume_clean_files_root", "") or "" + ).strip() + if resume and mode == "rjob": + if not resume_clean_files_root: + raise ValueError( + "RJob --resume requires --resume-clean-files-root" + ) + root = Path(resume_clean_files_root).expanduser().resolve(strict=False) + if root.parent == root: + raise ValueError("--resume-clean-files-root must not be the filesystem root") + if not root.is_dir(): + raise ValueError( + "--resume-clean-files-root must be an existing directory: " + f"{root}" + ) + resume_clean_files_root = str(root) + rjob_section = ( load_rjob_global_config(str(getattr(args, "rjob_config", "") or "")) if mode == "rjob" @@ -207,7 +229,7 @@ def load_simulation_run_config(args: Any) -> SimulationRunConfig: multiplier=float(args.multiplier), ) - job_id = str(args.job_id or "").strip() or uuid.uuid4().hex + job_id = requested_job_id or uuid.uuid4().hex max_workers = int(args.max_workers) if int(args.max_workers or 0) > 0 else None _validate_gateway_route_key(str(args.llm_model), arg_name="--llm-model") @@ -326,7 +348,8 @@ def load_simulation_run_config(args: Any) -> SimulationRunConfig: cleanup_stale_docker_containers=bool(getattr(args, "cleanup_stale_docker_containers", True)), max_workers=max_workers, rebuild_table=bool(args.rebuild_table), - resume=bool(getattr(args, "resume", False)), + resume=resume, + resume_clean_files_root=resume_clean_files_root, confirm_cloud_delete_job_id=str( getattr(args, "confirm_cloud_delete_job_id", "") or "" ).strip(), diff --git a/manager/simulation_flow.py b/manager/simulation_flow.py index 50a6c03..66d05f9 100644 --- a/manager/simulation_flow.py +++ b/manager/simulation_flow.py @@ -5,7 +5,7 @@ import logging import re from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from urllib.parse import quote, urlsplit import requests @@ -13,6 +13,7 @@ from core.data_manager.manager import DataManager from core.data_manager.yaml_aggregator import ( all_env_yaml_load, + delete_resume_session_steps, is_job_db_processing_done, sync_configs_to_db, wait_for_pending_inserts, @@ -24,7 +25,10 @@ from evaluator.service import EvaluationService from .agent_start_client import AgentStartClient from .manager import AgentPoolManager -from .resume_cleanup import cleanup_resume_artifacts +from .resume_cleanup import ( + cleanup_resume_artifacts, + select_generated_resume_environments, +) from .simulation_config import ( build_manager_runtime_config, expand_rl_epoch, @@ -38,6 +42,7 @@ _GATEWAY_ENV_ROW_POLL_INTERVAL_S = 30.0 _GATEWAY_ENV_ROW_MAX_ATTEMPTS = 10 +_UPSERT_FLUSH_INTERVAL_S = 10.0 class SimulationFlow: @@ -52,6 +57,7 @@ def __init__(self, cfg: SimulationRunConfig) -> None: self.gateway_client: Optional[GatewayClient] = None self.evaluation_service: Optional[EvaluationService] = None self.reward_committer: Optional[RewardCommitter] = None + self._resume_session_ids: List[str] = [] self._shutdown_started = False async def run(self) -> SimulationRunSummary: @@ -102,6 +108,10 @@ async def prepare_storage(self) -> None: "enable_buffer": self.cfg.enable_buffer, "buffer_size": self.cfg.buffer_size, "flush_interval": self.cfg.flush_interval, + "enable_upsert_batching": self.cfg.enable_buffer, + "upsert_batch_size": self.cfg.buffer_size, + "upsert_flush_interval": _UPSERT_FLUSH_INTERVAL_S, + "upsert_queue_size": max(100, self.cfg.buffer_size * 2), } if self.cfg.storage_type == "sqlite": storage_config["db_url"] = self.cfg.db_url @@ -117,11 +127,16 @@ async def prepare_storage(self) -> None: **storage_config, ) - yaml_config_list = all_env_yaml_load(env_root=self.cfg.agent_root, env_config=self.cfg.agent_config) - yaml_config_list = expand_rl_group_size(yaml_config_list, self.cfg.rl_group_size) - yaml_config_list = expand_rl_epoch(yaml_config_list, self.cfg.rl_epoch) + yaml_config_list = [] + if not self.cfg.resume: + yaml_config_list = all_env_yaml_load( + env_root=self.cfg.agent_root, + env_config=self.cfg.agent_config, + ) + yaml_config_list = expand_rl_group_size(yaml_config_list, self.cfg.rl_group_size) + yaml_config_list = expand_rl_epoch(yaml_config_list, self.cfg.rl_epoch) - await sync_configs_to_db( + resume_environments = await sync_configs_to_db( self.data_manager, yaml_config_list, self.cfg.storage_type, @@ -131,13 +146,33 @@ async def prepare_storage(self) -> None: resume=self.cfg.resume, ) self.manager_cfg = build_manager_runtime_config(self.cfg) + cleanup_environments = resume_environments if self.cfg.resume and self.cfg.mode == "rjob": + cleanup_environments = select_generated_resume_environments( + job_id=self.cfg.job_id, + results_root=Path(self.cfg.resume_clean_files_root), + environment_rows=resume_environments, + ) await cleanup_resume_artifacts( job_id=self.cfg.job_id, model=self.cfg.llm_model, data_manager=self.data_manager, manager_cfg=self.manager_cfg, + results_root=Path(self.cfg.resume_clean_files_root), + environment_rows=cleanup_environments, + ) + self._resume_session_ids = list(dict.fromkeys( + str(row.get("env_id") or "").strip() + for row in cleanup_environments + if str(row.get("env_id") or "").strip() + )) + if self.cfg.resume: + await delete_resume_session_steps( + self.data_manager, + job_id=self.cfg.job_id, + environment_rows=cleanup_environments, ) + resume_environments.clear() log.info( "storage prepared: job_id=%s base_pool_size=%d warm_pool_size=%d startup_submit_count=%d followup_submit_batch=%d", self.cfg.job_id, @@ -213,18 +248,7 @@ def _fetch_count() -> tuple[int, str]: ) async def clear_resume_gateway_session_cache(self) -> None: - if self.data_manager is None: - raise RuntimeError("data manager is not prepared") - rows = await self.data_manager.list_environment_rows( - job_id=self.cfg.job_id, - finished=False, - is_deleted=False, - ) - session_ids = [ - str(row.get("env_id")) - for row in rows - if row.get("env_id") - ] + session_ids = self._resume_session_ids if not session_ids: return client = GatewayClient(gateway_base_url=self.cfg.gateway_base_url) @@ -232,6 +256,7 @@ async def clear_resume_gateway_session_cache(self) -> None: result = await client.clear_session_cache(session_ids) finally: await client.aclose() + self._resume_session_ids = [] log.info( "gateway resume session cache cleared: job_id=%s sessions=%d removed=%s", self.cfg.job_id, diff --git a/manager/simulation_worker.py b/manager/simulation_worker.py index 16cd023..82369b3 100644 --- a/manager/simulation_worker.py +++ b/manager/simulation_worker.py @@ -5,10 +5,9 @@ import logging import time from collections import deque +from dataclasses import dataclass from typing import Any, Dict, Optional -import httpx - from core.data_manager.load_yaml import materialize_dataset_env_params from core.data_manager.contracts import SessionContext from core.data_manager.manager import DataManager @@ -34,6 +33,12 @@ log = logging.getLogger("manager.simulation_worker") +@dataclass(frozen=True, slots=True) +class _StoredResult: + status: str + total_reward: Optional[float] + + class _SimulationCircuitBreaker: _TIMEOUT_MARKERS = ( "timed out", @@ -145,8 +150,7 @@ def __init__( self.evaluation_service = evaluation_service self.reward_committer = reward_committer self.worker_count = self._derive_worker_count() - self._results: Dict[str, SimulationStartResult] = {} - self._results_lock = asyncio.Lock() + self._results: Dict[str, _StoredResult] = {} self._circuit_breaker = _SimulationCircuitBreaker(cfg) async def run_all(self) -> SimulationRunSummary: @@ -174,8 +178,7 @@ async def run_all(self) -> SimulationRunSummary: await asyncio.gather(*tasks, return_exceptions=True) raise - async with self._results_lock: - results = dict(self._results) + results = self._results if not results: return SimulationRunSummary( @@ -290,7 +293,8 @@ async def _worker_loop(self, worker_id: int) -> None: if self.gateway_client is not None: with trace.span("gateway_finalize"): gateway_finalized = await self._finalize_gateway_session( - result, + result.session_id, + rollout_status=result.status, completion_mode=completion_mode, worker_id=worker_id, agent_key=agent_key, @@ -319,10 +323,7 @@ async def _worker_loop(self, worker_id: int) -> None: ) result.total_reward = 0.0 with trace.span("mark_environment_finished"): - await self._mark_environment_finished_and_clean_gateway( - lease.agent_id, - result.session_id, - ) + await self.data_manager.mark_environment_finished(lease.agent_id) release_reusable = False elif self.evaluation_service is not None and self.reward_committer is not None: with trace.span("eval_discover_rule"): @@ -365,10 +366,7 @@ async def _worker_loop(self, worker_id: int) -> None: ) result.total_reward = eval_result.normalized_score_10 with trace.span("mark_environment_finished"): - await self._mark_environment_finished_and_clean_gateway( - lease.agent_id, - result.session_id, - ) + await self.data_manager.mark_environment_finished(lease.agent_id) else: result.status = "failed" result.error_text = eval_result.error_text or eval_result.reason @@ -386,15 +384,11 @@ async def _worker_loop(self, worker_id: int) -> None: job_id=self.cfg.job_id, llm_model=self.cfg.llm_model, ) - await self._mark_environment_finished_and_clean_gateway( - lease.agent_id, - result.session_id, - ) + await self.data_manager.mark_environment_finished(lease.agent_id) release_reusable = None with trace.span("store_result"): - async with self._results_lock: - self._results[agent_key] = result + self._results[agent_key] = _StoredResult(result.status, result.total_reward) except asyncio.CancelledError: cancelled = True release_reusable = False @@ -419,10 +413,20 @@ async def _worker_loop(self, worker_id: int) -> None: result.error_text = str(exc) release_reusable = False with trace.span("store_failed_result"): - async with self._results_lock: - self._results[agent_key] = result + self._results[agent_key] = _StoredResult(result.status, result.total_reward) finally: try: + if self.gateway_client is not None and session is not None: + if not gateway_finalized: + gateway_finalized = await self._finalize_gateway_session( + session.session_id, + rollout_status=result.status if result is not None else "cancelled", + completion_mode="abort", + worker_id=worker_id, + agent_key=agent_key, + trace=trace, + ) + await self._clean_gateway_session(session.session_id) if result is not None: with trace.span("record_circuit_result"): await self._record_circuit_result(result, worker_id=worker_id, agent_key=agent_key) @@ -521,8 +525,9 @@ async def _run_one_episode( async def _finalize_gateway_session( self, - result: SimulationStartResult, + session_id: str, *, + rollout_status: str, completion_mode: str, worker_id: int, agent_key: str, @@ -538,52 +543,43 @@ async def _finalize_gateway_session( try: if trace is None: await self.gateway_client.close_session( - result.session_id, + session_id, reason=reason, completion_mode=completion_mode, ) else: with trace.span("gateway_close_session"): await self.gateway_client.close_session( - result.session_id, + session_id, reason=reason, completion_mode=completion_mode, ) return True - except httpx.HTTPError as exc: + except Exception as exc: log.warning( "worker=%d agent=%s gateway session finalization failed; preserving rollout status=%s " "session_id=%s error=%s", worker_id, agent_key, - result.status, - result.session_id, + rollout_status, + session_id, exc, ) return False - async def _mark_environment_finished_and_clean_gateway( - self, - env_id: str, - session_id: str, - ) -> None: - await self.data_manager.mark_environment_finished(env_id) + async def _clean_gateway_session(self, session_id: str) -> None: if self.gateway_client is None: return try: cleaned = await self.gateway_client.clean_session(session_id) log.info( - "gateway session cleaned after environment completion: " - "env_id=%s session_id=%s status=%s", - env_id, + "gateway session cleaned after worker completion: session_id=%s status=%s", session_id, cleaned.get("status") if isinstance(cleaned, dict) else "cleaned", ) except Exception as exc: log.warning( - "gateway session clean failed after environment completion: " - "env_id=%s session_id=%s error=%s", - env_id, + "gateway session clean failed after worker completion: session_id=%s error=%s", session_id, exc, ) diff --git a/manager/types.py b/manager/types.py index 10c5f2b..31d9586 100644 --- a/manager/types.py +++ b/manager/types.py @@ -81,6 +81,7 @@ class SimulationRunConfig: max_workers: Optional[int] = None rebuild_table: bool = False resume: bool = False + resume_clean_files_root: str = "" confirm_cloud_delete_job_id: str = "" confirm_production: bool = False enable_buffer: bool = True