refactor container and fix suspend time (#38)

tylerharter · web-flow · commit d3dc54ea1ee8 · 2025-12-03T22:04:56.000-06:00
diff --git a/eudoxia/executor/container.py b/eudoxia/executor/container.py
@@ -1,16 +1,16 @@
 import logging
-from typing import List
+from typing import List, Generator, Optional
 from eudoxia.utils import DISK_SCAN_GB_SEC, Priority
 from eudoxia.workload import Operator
 
 logger = logging.getLogger(__name__)
 
 
-class Container: 
+class Container:
     """
-    An encapsulation of CPU, RAM, and a list of operators. A container is
-    created and then calculates how many ticks it will need to run with
-    resources provided.
+    An encapsulation of CPU, RAM, and a list of operators. A container executes
+    operators tick-by-tick, tracking memory usage and allowing suspension only
+    between operator boundaries.
     """
     next_container_num = 1
 
@@ -21,18 +21,22 @@ def __init__(self, ram, cpu, ops, prty: Priority, pool_id: int, ticks_per_second
         self.ram = ram
         self.cpu = cpu
         self.operators: List[Operator] = ops
-        self.segment_tick_boundaries = []
         self.suspend_ticks = None
         self._suspend_ticks_left = None
         self.priority = prty
-        self.error: str = None
+        self.error: Optional[str] = None
         self.ticks_per_second = ticks_per_second
         self.tick_length_secs = 1.0 / ticks_per_second
 
-        self.num_ticks = self._compute_ticks() 
-        self._num_ticks_left = self.num_ticks
-        self.num_secs = self.num_ticks * self.tick_length_secs
-        
+        # Tick state (updated by generator)
+        self._current_memory: float = 0.0
+        self._can_suspend: bool = False
+        self._completed: bool = False
+        self._ticks_elapsed: int = 0
+
+        # Start the generator
+        self._tick_iter = self._tick_generator()
+
     def get_pipeline_id(self):
         """Get the pipeline ID from operators, handling mixed pipelines"""
         pipeline_ids = set()
@@ -49,62 +53,98 @@ def get_pipeline_id(self):
 
     def __repr__(self):
         num_ops = len(self.operators)
-        return f"container={self.container_id} pipeline={self.get_pipeline_id()} ops={num_ops} cpus={self.cpu} ram_gb={self.ram} runtime_secs={self.num_secs:.1f}"
+        return f"container={self.container_id} pipeline={self.get_pipeline_id()} ops={num_ops} cpus={self.cpu} ram_gb={self.ram}"
 
-    def _compute_ticks(self) -> int:
+    def _tick_generator(self) -> Generator[None, None, None]:
         """
-        This function utilizes functions provided by Segment to calculate the
-        amount of CPU and RAM ticks that are needed across all segments.  It also
-        computes the boundaries between segments.  The reason is that we can only
-        suspend containers between segments (not in the middle of while a segment
-        is running).
-        Returns:
-            int: number of ticks this container will need to run for
+        Generator that drives tick-by-tick execution, updating container state directly.
+
+        Updates self._current_memory and self._can_suspend on each yield.
+
+        Memory usage is determined by the current segment:
+        - If segment.memory_gb is set: fixed memory usage
+        - If segment.memory_gb is None: memory grows linearly with I/O progress
+
+        Suspension is only allowed between operators (not between segments).
         """
-        total_ticks = 0
-        for op in self.operators:
-            for seg in op.get_segments():
-                # will it OOM, and if so, when?
-                oom_seconds = seg.get_seconds_until_oom(self.ram)
-
-                if oom_seconds is not None:
-                    # compute how long it it will be until the OOM occurs
-                    self.error = "OOM"
-                    seg_ticks_before_OOM = int(oom_seconds / self.tick_length_secs)
-                    total_ticks += seg_ticks_before_OOM
-                    return total_ticks # after OOM, we cannot run other segments
-                else:
-                    # there is no OOM.  We will spend all the time
-                    # expected on I/O (first), then CPU (second)
-                    io_time_secs = seg.get_io_seconds()
-                    cpu_time_secs = seg.get_cpu_time(self.cpu)
-                    total_ticks += int((io_time_secs + cpu_time_secs) / self.tick_length_secs)
-                    self.segment_tick_boundaries.append(total_ticks)
-        return total_ticks
+
+        # loop over every tick of every segment of every op
+        #
+        # at each iteration, determine memory usage, whether there is
+        # an OOM, and whether suspension is possible
+        for op_idx, op in enumerate(self.operators):
+            segments = op.get_segments()
+            for seg_idx, seg in enumerate(segments):
+                # Calculate ticks for I/O phase and CPU phase
+                io_secs = seg.get_io_seconds()
+                cpu_secs = seg.get_cpu_time(self.cpu)
+                io_ticks = int(io_secs / self.tick_length_secs)
+                cpu_ticks = int(cpu_secs / self.tick_length_secs)
+                total_seg_ticks = io_ticks+cpu_ticks
+
+                for i in range(total_seg_ticks):
+                    # determine current memory consumption
+                    if i < io_ticks:
+                        if seg.memory_gb is not None:
+                            self._current_memory = seg.memory_gb
+                        else:
+                            # Memory grows linearly with I/O progress
+                            io_progress_secs = (i + 1) * self.tick_length_secs
+                            self._current_memory = io_progress_secs * DISK_SCAN_GB_SEC
+                    else:
+                        self._current_memory = seg.get_peak_memory_gb()
+
+                    # have we OOM'd?
+                    if self._current_memory > self.ram:
+                        self.error = "OOM"
+                        self._completed = True
+                        yield
+                        return
+
+                    # are we at the end of the op (last tick of last
+                    # seg)?  if so, we're either completed, or we can
+                    # suspend, depending on whether this is the last
+                    # op.
+                    self._can_suspend = False
+                    if seg_idx == len(segments)-1 and i == total_seg_ticks - 1:
+                        if op_idx == len(self.operators) - 1:
+                            self._current_memory = 0.0
+                            self._completed = True
+                        else:
+                            self._can_suspend = True
+                    yield
 
     def tick(self):
-        self._num_ticks_left -= 1
+        """
+        Execute one tick. Advances to next state (OOM checked in _advance_tick).
+        """
+        if self._completed:
+            return
+        next(self._tick_iter)
+        self._ticks_elapsed += 1
 
     def is_completed(self):
-        # Use <= to handle edge case: immediate OOM results in 0 ticks,
-        # but tick() may still be called, causing _num_ticks_left to go negative
-        return (self._num_ticks_left <= 0)
+        return self._completed
+
+    def ticks_elapsed(self) -> int:
+        """Returns the number of ticks that have been executed."""
+        return self._ticks_elapsed
 
     def can_suspend_container(self) -> bool:
-        """Can only suspend a container if between execution of
-        segments. Must wait for Segment to complete before the
-        container is pausable"""
-        elapsed = self.num_ticks - self._num_ticks_left
-        if elapsed in self.segment_tick_boundaries:
-            return True
-        return False
+        """Can only suspend a container between operators."""
+        return self._can_suspend
+
+    def get_current_memory_usage(self) -> float:
+        """Returns current memory usage in GB."""
+        return self._current_memory
 
     def suspend_container(self):
         """
         Suspend container execution, free CPUs and RAM. Requires writing
-        current data to disk.  
+        current data to disk.
         """
-        write_to_disk_ticks = self.ram / DISK_SCAN_GB_SEC
+        write_to_disk_secs = self.ram / DISK_SCAN_GB_SEC
+        write_to_disk_ticks = int(write_to_disk_secs / self.tick_length_secs)
         self.suspend_ticks = write_to_disk_ticks
         self._suspend_ticks_left = write_to_disk_ticks
 
diff --git a/eudoxia/executor/resource_pool.py b/eudoxia/executor/resource_pool.py
@@ -138,7 +138,6 @@ def run_one_tick(self, suspensions: List[Suspend],
                 self.avail_cpu_pool -= a.cpu
                 self.avail_ram_pool -= a.ram
                 self.active_containers.append(container)
-                self.container_tick_times.append(container.num_ticks)
                 logger.info(f"start container {container}")
 
         to_remove = []
@@ -168,6 +167,9 @@ def run_one_tick(self, suspensions: List[Suspend],
                 logger.info(result)
                 results.append(result)
 
+                # Record actual ticks executed for stats
+                self.container_tick_times.append(c.ticks_elapsed())
+
                 # Track successful completions
                 if c.error is None:
                     self.num_completed += 1
diff --git a/tests/test_container.py b/tests/test_container.py
@@ -35,13 +35,14 @@ def test_container_oom():
 
         assert ticks_executed <= 1000, "Container should complete within 1000 ticks"
 
-    # the first 3 containers should run (30s/30 ticks), then we should
-    assert ticks_executed == 30, \
-        f"Container should stop at OOM at 30 ticks, but ran for {ticks_executed} ticks"
+    # the first 3 operators should run (30 ticks), then the 4th operator (40GB)
+    # triggers OOM on tick 31 when we try to execute it
+    assert ticks_executed == 31, \
+        f"Container should stop at OOM at 31 ticks, but ran for {ticks_executed} ticks"
     assert container.error == "OOM", "Container should have OOM error after completion"
 
 
-def test_container_tick0_oom():
+def test_container_immediate_oom():
     """Test that immediate OOM (when segment needs more memory than container has) completes properly"""
 
     # Create a segment that needs 200GB of memory
@@ -63,13 +64,146 @@ def test_container_tick0_oom():
         ticks_per_second=10
     )
 
-    # Container should have 0 ticks since it OOMs immediately
-    assert container.num_ticks == 0, f"Expected 0 ticks for immediate OOM, got {container.num_ticks}"
-    assert container.error == "OOM", "Container should have OOM error"
+    # Container is not yet completed - OOM detected on first tick
+    assert not container.is_completed(), "Container should not be completed before first tick"
 
-    # Container should already be completed (0 ticks to run)
-    assert container.is_completed(), "Container with 0 ticks should be immediately completed"
+    # First tick should detect OOM
+    container.tick()
+    assert container.ticks_elapsed() == 1, f"Expected 1 tick before OOM, got {container.ticks_elapsed()}"
+    assert container.error == "OOM", "Container should have OOM error"
+    assert container.is_completed(), "Container should be completed after OOM"
 
-    # Calling tick() shouldn't break the is_completed() check
+    # Calling tick() again shouldn't break anything
     container.tick()
-    assert container.is_completed(), "Container should still be completed after tick()"
+    assert container.is_completed(), "Container should still be completed after extra tick()"
+
+
+def test_container_suspension_ticks():
+    """Test that suspension takes the correct number of ticks to write memory to disk"""
+
+    op = Operator()
+    op.add_segment(Segment(
+        baseline_cpu_seconds=1.0,
+        cpu_scaling="const",
+        memory_gb=10,
+    ))
+
+    # 100 ticks per second, so tick_length = 0.01 seconds
+    container = Container(
+        ram=100,  # 100GB RAM
+        cpu=10,
+        ops=[op],
+        prty=Priority.BATCH_PIPELINE,
+        pool_id=0,
+        ticks_per_second=100
+    )
+
+    # Trigger suspension
+    container.suspend_container()
+
+    # Suspension should take: ram / DISK_SCAN_GB_SEC seconds = 100 / 20 = 5 seconds
+    # At 100 ticks/second, that's 500 ticks
+    expected_suspend_ticks = int((100 / DISK_SCAN_GB_SEC) / container.tick_length_secs)
+    assert container.suspend_ticks == expected_suspend_ticks, \
+        f"Expected {expected_suspend_ticks} suspend ticks, got {container.suspend_ticks}"
+    assert container.suspend_ticks == 500, \
+        f"Expected 500 suspend ticks, got {container.suspend_ticks}"
+
+    # Tick through suspension
+    suspend_ticks_counted = 0
+    while not container.is_suspended():
+        container.suspend_container_tick()
+        suspend_ticks_counted += 1
+
+    assert suspend_ticks_counted == 500, \
+        f"Expected 500 ticks to suspend, counted {suspend_ticks_counted}"
+
+
+def test_container_memory_over_time():
+    """Test that memory usage is tracked correctly over time.
+
+    Op1: single segment with fixed memory (10GB), 1 second CPU
+    Op2: two segments with I/O-based memory growth
+        - seg1: 2 seconds I/O (memory grows 0->40GB), 1 second CPU (stays at 40GB)
+        - seg2: 1 second I/O (memory grows 0->20GB), 0.5 second CPU (stays at 20GB)
+
+    At 10 ticks/second:
+    - Op1 seg1: 10 CPU ticks at 10GB fixed
+    - Op2 seg1: 20 I/O ticks (growing 2->40GB), then 10 CPU ticks at 40GB
+    - Op2 seg2: 10 I/O ticks (growing 2->20GB), then 5 CPU ticks at 20GB
+    """
+    # Op1: fixed memory
+    op1 = Operator()
+    op1.add_segment(Segment(
+        baseline_cpu_seconds=1.0,
+        cpu_scaling="const",
+        memory_gb=10,
+        storage_read_gb=0,
+    ))
+
+    # Op2: I/O-based memory (memory_gb=None means memory grows with I/O)
+    op2 = Operator()
+    op2.add_segment(Segment(
+        baseline_cpu_seconds=1.0,
+        cpu_scaling="const",
+        memory_gb=None,
+        storage_read_gb=40,  # 40GB read at 20GB/sec = 2 sec I/O, peak memory = 40GB
+    ))
+    op2.add_segment(Segment(
+        baseline_cpu_seconds=0.5,
+        cpu_scaling="const",
+        memory_gb=None,
+        storage_read_gb=20,  # 20GB read at 20GB/sec = 1 sec I/O, peak memory = 20GB
+    ))
+
+    container = Container(
+        ram=100,
+        cpu=1,  # 1 CPU so baseline_cpu_seconds = actual seconds
+        ops=[op1, op2],
+        prty=Priority.BATCH_PIPELINE,
+        pool_id=0,
+        ticks_per_second=10,
+    )
+
+    tick_length = 0.1  # 10 ticks/sec
+    expected_memories = []
+
+    # Op1 seg1: 10 CPU ticks at 10GB (no I/O since storage_read_gb=0)
+    for _ in range(10):
+        expected_memories.append(10.0)
+
+    # Op2 seg1: 20 I/O ticks with memory growing linearly
+    # Memory grows at DISK_SCAN_GB_SEC (20 GB/sec)
+    # After tick i, memory = (i+1) * tick_length * 20 = (i+1) * 2
+    for i in range(20):
+        expected_memories.append((i + 1) * tick_length * DISK_SCAN_GB_SEC)
+
+    # Op2 seg1: 10 CPU ticks at peak memory (40GB)
+    for _ in range(10):
+        expected_memories.append(40.0)
+
+    # Op2 seg2: 10 I/O ticks with memory growing linearly
+    for i in range(10):
+        expected_memories.append((i + 1) * tick_length * DISK_SCAN_GB_SEC)
+
+    # Op2 seg2: 5 CPU ticks at peak memory (20GB)
+    # Last tick will show 0.0 since container completes
+    for i in range(5):
+        if i < 4:
+            expected_memories.append(20.0)
+        else:
+            expected_memories.append(0.0)  # Container completes, memory cleared
+
+    # Run through all ticks and verify memory
+    tick = 0
+    while not container.is_completed():
+        container.tick()
+        actual_memory = container.get_current_memory_usage()
+        expected_memory = expected_memories[tick]
+        assert actual_memory == expected_memory, \
+            f"Tick {tick}: expected memory {expected_memory}GB, got {actual_memory}GB"
+        tick += 1
+
+    assert tick == len(expected_memories), \
+        f"Expected {len(expected_memories)} ticks, got {tick}"
+    assert tick == 55, f"Expected 55 total ticks, got {tick}"
