Merge pull request #115 from KhrulkovV/perf/benchmark-squeeze

perf: 1.3-1.5x engine speedup via hot-path optimizations

Author: KhrulkovV

gigaevo/database/program_storage.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 from gigaevo.programs.program import Program
+from gigaevo.programs.program_state import ProgramState, validate_transition
 
 
 class PopulationSnapshot:
@@ -151,6 +152,43 @@ async def renew_instance_lock(self) -> bool:
         """
         ...
 
+    async def fast_state_transition(
+        self, program: Program, old_state: str, new_state: str
+    ) -> None:
+        """Fast state transition: 2 RT (INCR + pipeline) instead of ~5 RT.
+
+        Safe only when the caller holds exclusive single-process ownership
+        (e.g., asyncio.Lock in ProgramStateManager). Does NOT provide cross-process
+        safety — assumes each program is processed by exactly one engine instance.
+        Default falls back to atomic_state_transition.
+        """
+        await self.atomic_state_transition(program, old_state, new_state)
+
+    async def batch_transition_state(
+        self,
+        programs: list[Program],
+        old_state: str,
+        new_state: str,
+    ) -> int:
+        """Batch-transition programs between states.
+
+        Default implementation falls back to individual transitions.
+        Subclasses may override with pipelined operations.
+        Returns the number of programs transitioned.
+        """
+        old_enum = ProgramState(old_state)
+        new_enum = ProgramState(new_state)
+        count = 0
+        for prog in programs:
+            validate_transition(old_enum, new_enum)
+            prog.state = new_enum
+            await self.atomic_state_transition(prog, old_state, new_state)
+            count += 1
+        return count
+
+    async def remove_ids_from_status_set(self, status: str, ids: list[str]) -> None:
+        """Remove specific IDs from a status set. No-op by default."""
+
     async def wait_for_activity(self, timeout: float) -> None:
         """
         Block up to `timeout` seconds until storage observes activity (e.g., new

gigaevo/database/redis_program_storage.py

@@ -22,7 +22,7 @@
 )
 from gigaevo.exceptions import StorageError
 from gigaevo.programs.program import Program
-from gigaevo.programs.program_state import ProgramState
+from gigaevo.programs.program_state import ProgramState, validate_transition
 from gigaevo.utils.json import dumps as _dumps
 from gigaevo.utils.json import loads as _loads
 from gigaevo.utils.trackers.base import LogWriter
@@ -447,6 +447,109 @@ async def _atomic(r: aioredis.Redis) -> None:
 
         await self._conn.execute("atomic_state_transition", _atomic)
 
+    async def fast_state_transition(
+        self, program: Program, old_state: str, new_state: str
+    ) -> None:
+        """Fast state transition: 2 RT (INCR + pipeline) instead of ~5 RT.
+
+        Safe only when the caller holds exclusive single-process ownership
+        (e.g., asyncio.Lock in ProgramStateManager). Does NOT provide cross-process
+        safety — assumes each program is processed by exactly one engine instance.
+        Unlike atomic_state_transition, does not WATCH/GET/MERGE.
+        """
+        self._check_write_allowed("fast_state_transition")
+
+        async def _fast(r: aioredis.Redis) -> None:
+            key = self._keys.program(program.id)
+            counter = await r.incr(self._keys.timestamp())
+            data = program.to_dict()
+            data["atomic_counter"] = int(counter)
+
+            pipe = r.pipeline(transaction=False)
+            pipe.set(key, _dumps(data))
+            if old_state != new_state:
+                pipe.srem(self._keys.status_set(old_state), program.id)
+            pipe.sadd(self._keys.status_set(new_state), program.id)
+            pipe.xadd(
+                self._keys.status_stream(),
+                {"id": program.id, "status": new_state, "event": "transition"},
+                maxlen=STREAM_MAX_LEN,
+                approximate=True,
+            )
+            await pipe.execute()
+
+        await self._conn.execute("fast_state_transition", _fast)
+
+    async def batch_transition_state(
+        self,
+        programs: list[Program],
+        old_state: str,
+        new_state: str,
+    ) -> int:
+        """Batch-transition programs between states using pipelined ops.
+
+        Much faster than individual atomic_state_transition calls for large
+        batches (e.g., refresh phase with 5000 programs). Assumes exclusive
+        ownership — no WATCH/MERGE needed.
+
+        Returns the number of programs transitioned.
+        """
+        self._check_write_allowed("batch_transition_state")
+        if not programs:
+            return 0
+
+        old_enum = ProgramState(old_state)
+        new_enum = ProgramState(new_state)
+        for prog in programs:
+            validate_transition(old_enum, new_enum)
+
+        async def _batch(r: aioredis.Redis) -> int:
+            old_set_key = self._keys.status_set(old_state)
+            new_set_key = self._keys.status_set(new_state)
+            stream_key = self._keys.status_stream()
+            ts_key = self._keys.timestamp()
+
+            count = 0
+            for chunk in self._chunks(programs, MGET_CHUNK_SIZE):
+                n = len(chunk)
+                # Reserve N counters in one call
+                end_counter = await r.incrby(ts_key, n)
+                start_counter = end_counter - n + 1
+
+                pipe = r.pipeline(transaction=False)
+                for i, prog in enumerate(chunk):
+                    prog.state = ProgramState(new_state)
+                    data = prog.to_dict()
+                    data["atomic_counter"] = int(start_counter + i)
+
+                    pipe.set(self._keys.program(prog.id), _dumps(data))
+                    pipe.srem(old_set_key, prog.id)
+                    pipe.sadd(new_set_key, prog.id)
+
+                pipe.xadd(
+                    stream_key,
+                    {"id": "batch", "status": new_state, "event": "batch_transition"},
+                    maxlen=STREAM_MAX_LEN,
+                    approximate=True,
+                )
+                await pipe.execute()
+                count += n
+
+            return count
+
+        return await self._conn.execute("batch_transition_state", _batch)
+
+    async def remove_ids_from_status_set(self, status: str, ids: list[str]) -> None:
+        """Remove specific IDs from a status set using SREM."""
+        if not ids:
+            return
+        self._check_write_allowed("remove_ids_from_status_set")
+
+        async def _srem(r: aioredis.Redis) -> None:
+            await r.srem(self._keys.status_set(status), *ids)
+
+        await self._conn.execute("remove_ids_from_status_set", _srem)
+
     # --------------------- Run State (resume support) ---------------------
 
     async def save_run_state(self, field: str, value: int) -> None:

gigaevo/database/state_manager.py

@@ -106,7 +106,14 @@ async def set_program_state(
 
             program.state = new_state
             old = old_state.value if old_state else None
-            await self.storage.atomic_state_transition(program, old, new_state.value)
+            # Use fast path (2 RT) when old state is known — safe because
+            # per-program asyncio.Lock above prevents concurrent writes.
+            if old is not None:
+                await self.storage.fast_state_transition(program, old, new_state.value)
+            else:
+                await self.storage.atomic_state_transition(
+                    program, old, new_state.value
+                )
 
             logger.debug(
                 "[ProgramStateManager] {} {} → {}",

gigaevo/evolution/engine/core.py

@@ -313,14 +313,52 @@ def _format_eta(self) -> str:
     async def _await_idle(self) -> None:
         """Block until there are no programs in QUEUED or RUNNING."""
         t0 = time.monotonic()
-        while await self._has_active_dags():
+        ghost_checked = False
+        while True:
+            has_active = await self._has_active_dags()
+            if not has_active:
+                break
+
             elapsed = time.monotonic() - t0
             if elapsed > 30 and int(elapsed) % 60 < self.config.loop_interval:
                 logger.info(
                     "[EvolutionEngine] gen={} Waiting for idle ({:.0f}s elapsed)",
                     self.metrics.total_generations,
                     elapsed,
                 )
+            # Ghost safety: after 30s, verify counts with full fetch (once)
+            if elapsed > 30 and not ghost_checked:
+                ghost_checked = True
+                real_q = len(
+                    await self.storage.get_all_by_status(ProgramState.QUEUED.value)
+                )
+                real_r = len(
+                    await self.storage.get_all_by_status(ProgramState.RUNNING.value)
+                )
+                if real_q == 0 and real_r == 0:
+                    # Clean up ghost IDs from status sets
+                    queued_ids = await self.storage.get_ids_by_status(
+                        ProgramState.QUEUED.value
+                    )
+                    running_ids = await self.storage.get_ids_by_status(
+                        ProgramState.RUNNING.value
+                    )
+                    if queued_ids:
+                        await self.storage.remove_ids_from_status_set(
+                            ProgramState.QUEUED.value, queued_ids
+                        )
+                    if running_ids:
+                        await self.storage.remove_ids_from_status_set(
+                            ProgramState.RUNNING.value, running_ids
+                        )
+                    ghost_count = len(queued_ids) + len(running_ids)
+                    logger.warning(
+                        "[EvolutionEngine] Ghost IDs detected — SCARD says active "
+                        "but no real programs found. Cleaned {} ghost ID(s) from "
+                        "status sets. Breaking idle wait.",
+                        ghost_count,
+                    )
+                    break
             await asyncio.sleep(self.config.loop_interval)
 
     async def _select_elites_for_mutation(self) -> list[Program]:
@@ -359,18 +397,40 @@ async def _ingest_completed_programs(self) -> None:
         Programs already in the archive stay DONE (they arrived from a refresh DAG).
         New programs are added if accepted, otherwise discarded.
         """
-        completed = await self.storage.get_all_by_status(ProgramState.DONE.value)
-        if not completed:
+        # Fetch only IDs first (SMEMBERS — no deserialization), then filter
+        # out archive programs before doing the expensive mget+deserialize.
+        done_ids = await self.storage.get_ids_by_status(ProgramState.DONE.value)
+        if not done_ids:
             logger.debug(
                 "[EvolutionEngine] gen={} No completed programs to ingest",
                 self.metrics.total_generations,
             )
             return
 
+        archive_program_ids = set(await self.strategy.get_program_ids())
+        new_ids = [pid for pid in done_ids if pid not in archive_program_ids]
+
+        if not new_ids:
+            logger.debug(
+                "[EvolutionEngine] gen={} {} DONE programs all in archive, skipping",
+                self.metrics.total_generations,
+                len(done_ids),
+            )
+            return
+
+        # Only deserialize the new (non-archive) programs
+        completed = await self.storage.mget(new_ids)
+        # Filter to actual DONE state (mget may return stale status)
+        completed = [p for p in completed if p.state == ProgramState.DONE]
+
+        if not completed:
+            return
+
         logger.info(
-            "[EvolutionEngine] gen={} Ingest {} program(s)",
+            "[EvolutionEngine] gen={} Ingest {} program(s) ({} in archive skipped)",
             self.metrics.total_generations,
             len(completed),
+            len(done_ids) - len(new_ids),
         )
         logger.debug(
             "[EvolutionEngine] Program IDs: {}",
@@ -383,14 +443,10 @@ async def _ingest_completed_programs(self) -> None:
         rej_strategy = 0
 
         state_tasks: list[asyncio.Task] = []
-        archive_program_ids = set(await self.strategy.get_program_ids())
 
         for prog in completed:
             try:
-                if prog.id in archive_program_ids:
-                    # already in archive (arrived from a refresh DAG) — stays DONE
-                    continue
-                elif not self.config.program_acceptor.is_accepted(prog):
+                if not self.config.program_acceptor.is_accepted(prog):
                     # rejected by basic checks
                     rej_valid += 1
                     logger.info(
@@ -460,18 +516,25 @@ async def _refresh_archive_programs(self) -> int:
             return 0
 
         programs_to_refresh = await self.storage.mget(program_ids_to_refresh)
+        done_programs = [p for p in programs_to_refresh if p.state == ProgramState.DONE]
 
-        tasks: list[asyncio.Task] = []
-        for program in programs_to_refresh:
-            if program.state == ProgramState.DONE:
-                tasks.append(
-                    asyncio.create_task(self._set_state(program, ProgramState.QUEUED))
-                )
+        if not done_programs:
+            return 0
 
-        if tasks:
-            await asyncio.gather(*tasks, return_exceptions=True)
+        try:
+            count = await self.storage.batch_transition_state(
+                done_programs,
+                ProgramState.DONE.value,
+                ProgramState.QUEUED.value,
+            )
+        except Exception as e:
+            logger.error(
+                "[EvolutionEngine] gen={} batch_transition_state failed: {}",
+                self.metrics.total_generations,
+                e,
+            )
+            return 0
 
-        count = len(tasks)
         if count:
             logger.info(
                 "[EvolutionEngine] gen={} Submitted {} program(s) for refresh",
@@ -484,16 +547,14 @@ async def _refresh_archive_programs(self) -> int:
     async def _has_active_dags(self) -> bool:
         """True if any programs are QUEUED or RUNNING (i.e., engine not idle).
 
-        Uses get_all_by_status (SMEMBERS + MGET) rather than count_by_status
-        (raw SCARD) so that ghost IDs — set members with no backing program
-        data — are filtered out and cannot permanently stall _await_idle().
+        Uses count_by_status (SCARD, O(1)) for the fast path.  Falls back to
+        the expensive get_all_by_status after 30s of continuous waiting to
+        detect ghost IDs that would otherwise stall _await_idle forever.
         """
-        queued_progs, running_progs = await asyncio.gather(
-            self.storage.get_all_by_status(ProgramState.QUEUED.value),
-            self.storage.get_all_by_status(ProgramState.RUNNING.value),
+        queued, running = await asyncio.gather(
+            self.storage.count_by_status(ProgramState.QUEUED.value),
+            self.storage.count_by_status(ProgramState.RUNNING.value),
         )
-        queued = len(queued_progs)
-        running = len(running_progs)
 
         if queued or running:
             current_counts = (queued, running)