Feature/probinso/fast completion detection

.gitignore

@@ -1,3 +1,5 @@
 __pycache__/
 *.log*
 *out
+*.swp
+*.swo

ats/atsMachines/fluxScheduled.py

@@ -175,6 +175,15 @@ def get_physical_node(self, rel_index):
             raise IndexError(f"Relative index {rel_index} out of range (0-{len(nodes)-1})")
         return nodes[rel_index]
 
+    def getNumberOfProcessors(self):
+        """Return total schedulable processor slots in the Flux allocation.
+
+        Returns:
+            int: Total processor capacity ATS may schedule inside the current
+            Flux allocation.
+        """
+        return self.maxCores
+
 
     def kill(self, test):
         """

ats/completion_detector.py

@@ -0,0 +1,143 @@
+"""Completion-detection strategy base class and factory helpers."""
+
+from abc import ABC, abstractmethod
+
+from ats.atsut import AtsError
+
+
+def normalize_completion_detection_mode(mode):
+    """Normalize a completion-detection mode string.
+
+    Args:
+        mode (str|None): Requested mode name.
+
+    Returns:
+        str: Normalized lowercase mode name, defaulting to
+        ``"per_test_watcher"``.
+    """
+    return str(mode or "per_test_watcher").strip().lower() or "per_test_watcher"
+
+
+def _validate_completion_detection_mode_for_machine(machine, normalized_mode):
+    """Reject detector modes unsupported by one machine implementation.
+
+    Args:
+        machine: Machine instance that would own the detector.
+        normalized_mode (str): Normalized detector mode name.
+
+    Returns:
+        None: Validation succeeds without modifying ``machine``.
+
+    Raises:
+        AtsError: If ``normalized_mode`` is unsupported for ``machine``.
+    """
+    machine_class = machine.__class__.__name__
+    machine_module = machine.__class__.__module__
+    if (
+        normalized_mode in ("waitpid_reaper", "per_test_watcher")
+        and machine_class == "FluxDirect"
+        and machine_module.endswith("flux_direct")
+    ):
+        raise AtsError(
+            "threaded completion detection is unsupported for FluxDirect. "
+            "Use 'poll' for this experimental machine."
+        )
+
+
+class CompletionDetector(ABC):
+    """Abstract policy object with shared ATS completion-detection helpers."""
+
+    mode_name = ""
+
+    def __init__(self, machine):
+        """Bind one completion detector to one ATS machine.
+
+        Args:
+            machine: Machine instance that owns completion helpers and running
+                test state.
+
+        Returns:
+            None: The detector stores a reference to ``machine``.
+        """
+        self.machine = machine
+
+    def register_launched_test(self, test):
+        """Prepare one launched test for detector-specific completion work.
+
+        Args:
+            test: ATS test object whose child process has just been launched.
+
+        Returns:
+            None: The default detector implementation needs no launch-time
+            setup.
+        """
+
+    def unregister_finished_test(self, test):
+        """Release detector-owned state associated with one finished test.
+
+        Args:
+            test: ATS test object that may own detector-specific wait state.
+
+        Returns:
+            None: The default detector implementation needs no cleanup.
+        """
+
+    def owns_child_reaping(self):
+        """Return whether this detector is responsible for child reaping.
+
+        Returns:
+            bool: ``False`` for the default detector behavior.
+        """
+        return False
+
+    @abstractmethod
+    def logCompletionWarnings(self, logger):
+        """Print any detector-specific end-of-run completion warnings.
+
+        Args:
+            logger (callable): Logging function compatible with ``ats.log.log``.
+
+        Returns:
+            None: Implementations may emit zero or more warnings.
+        """
+
+    @abstractmethod
+    def check_running(self):
+        """Update machine running state according to one detector strategy.
+
+        Returns:
+            None: Implementations may finish tests and update ``machine.running``.
+        """
+
+
+def create_completion_detector(machine, mode):
+    """Create one completion detector instance for the requested mode.
+
+    Args:
+        machine: Machine instance that will own the detector.
+        mode (str|None): Requested detector mode.
+
+    Returns:
+        CompletionDetector: Strategy instance for the requested mode.
+
+    Raises:
+        AtsError: If ``mode`` is not one of the supported detector modes.
+    """
+    normalized_mode = normalize_completion_detection_mode(mode)
+    _validate_completion_detection_mode_for_machine(machine, normalized_mode)
+    if normalized_mode == "waitpid_reaper":
+        from ats.completion_waitpid_reaper import WaitpidReaperCompletionDetector
+
+        return WaitpidReaperCompletionDetector(machine)
+    if normalized_mode == "per_test_watcher":
+        from ats.completion_per_test_watcher import PerTestWatcherCompletionDetector
+
+        return PerTestWatcherCompletionDetector(machine)
+    if normalized_mode == "poll":
+        from ats.completion_poll import PollingCompletionDetector
+
+        return PollingCompletionDetector(machine)
+    raise AtsError(
+        "Unknown completion detection mode %r. Expected one of: "
+        "'waitpid_reaper', 'per_test_watcher', 'poll'." % mode
+    )

ats/completion_per_test_watcher.py

@@ -0,0 +1,57 @@
+"""Per-test watcher completion detector for ATS machines."""
+
+import threading
+
+from ats.completion_waitpid_reaper import WaitpidReaperCompletionDetector
+
+
+class PerTestWatcherCompletionDetector(WaitpidReaperCompletionDetector):
+    """Use one ``Popen.wait`` watcher thread per launched test.
+
+    Each watcher
+    waits on the exact child process owned by its test, then enqueues that
+    test for the shared fast-path drain.
+    This detector is thread safe, but may not scale beyond a few nodes of
+    concurrently running tests.
+    """
+
+    mode_name = "per_test_watcher"
+
+    def owns_child_reaping(self):
+        """Simple queue mode leaves child reaping to ``subprocess.Popen``."""
+        return False
+
+    def register_launched_test(self, test):
+        """Start one watcher thread that waits for this child to exit."""
+        child = getattr(test, "child", None)
+        if child is None:
+            return
+        watcher = getattr(test, "_completionWatcher", None)
+        if watcher is not None:
+            self.machine._incrementCompletionStat("per_test_watcher_already_registered")
+            return
+
+        def watch_for_completion():
+            """Wait for one child to exit and then enqueue its completion."""
+            try:
+                child.wait()
+            except Exception:
+                self.machine._incrementCompletionStat("per_test_watcher_wait_error")
+                return
+            self.machine._incrementCompletionStat("per_test_watcher_wait_completed")
+            self.record_completion_signal(test)
+
+        watcher = threading.Thread(
+            target=watch_for_completion,
+            name=f"ats-completion-{getattr(child, 'pid', 'unknown')}",
+            daemon=True,
+        )
+        test._completionWatcher = watcher
+        watcher.start()
+        self.machine._incrementCompletionStat("per_test_watcher_registered")
+
+    def unregister_finished_test(self, test):
+        """Clear watcher bookkeeping for one finished test."""
+        if hasattr(test, "_completionWatcher"):
+            test._completionWatcher = None
+            self.machine._incrementCompletionStat("per_test_watcher_cleared")

ats/completion_poll.py

@@ -0,0 +1,126 @@
+"""Polling completion detector for ATS machines."""
+
+import time
+
+from ats.completion_detector import CompletionDetector
+from ats.atsut import AtsError, PASSED
+
+
+class PollingCompletionDetector(CompletionDetector):
+    """Sleep-then-poll completion behavior."""
+
+    mode_name = "poll"
+
+    def preserve_new_running_tests(self, remaining, seen_ids):
+        """Keep tests appended to ``machine.running`` during completion callbacks.
+
+        Args:
+            remaining (list): Running tests that should remain after the current
+                polling pass.
+            seen_ids (set): Object ids already considered in the polling pass.
+
+        Returns:
+            None: ``remaining`` is updated in place.
+        """
+        remaining_ids = {id(test) for test in remaining}
+        for test in self.machine.running:
+            test_id = id(test)
+            if test_id in seen_ids or test_id in remaining_ids:
+                continue
+            remaining.append(test)
+            remaining_ids.add(test_id)
+
+    def poll_running_tests(
+        self,
+        allow_running_checks,
+        completion_limit=None,
+    ):
+        """Poll running tests in scheduler order.
+
+        Args:
+            allow_running_checks (bool): When ``False``, skip timeout and
+                runtime error checks for children that have not yet exited.
+            completion_limit (int|None): Maximum number of completions to
+                process before returning control to the scheduler.
+
+        Returns:
+            int: Number of completed tests processed in this polling pass.
+        """
+        from ats import configuration
+
+        machine = self.machine
+        start_us = time.time_ns() // 1000
+        machine._incrementCompletionStat("_pollRunningTests_called")
+        if allow_running_checks:
+            machine._incrementCompletionStat("_pollRunningTests_allow_running_checks_true")
+        else:
+            machine._incrementCompletionStat("_pollRunningTests_allow_running_checks_false")
+
+        ordered_count = 0
+        completed = 0
+        result_kind = "completed_none"
+        try:
+            ordered = list(machine.running)
+            seen_ids = {id(test) for test in ordered}
+            ordered_count = len(ordered)
+            machine._incrementCompletionStat("_pollRunningTests_total_ordered", ordered_count)
+
+            remaining = []
+            for index, test in enumerate(ordered):
+                done = machine.getStatus(test, allow_running_checks=allow_running_checks)
+                if not done:
+                    remaining.append(test)
+                    continue
+                completed += 1
+                if test.status is not PASSED and configuration.options.oneFailure:
+                    raise AtsError("Test failed in oneFailure mode.")
+                if completion_limit is not None and completed >= completion_limit:
+                    remaining.extend(ordered[index + 1:])
+                    self.preserve_new_running_tests(remaining, seen_ids)
+                    machine.running = remaining
+                    result_kind = "stopped_after_completion_limit"
+                    machine._incrementCompletionStat("_pollRunningTests_stopped_after_completion")
+                    machine._incrementCompletionStat("_pollRunningTests_total_completed", completed)
+                    return completed
+
+            self.preserve_new_running_tests(remaining, seen_ids)
+            machine.running = remaining
+            machine._incrementCompletionStat("_pollRunningTests_total_completed", completed)
+            if completed:
+                result_kind = "completed"
+                machine._incrementCompletionStat("_pollRunningTests_completed")
+            else:
+                machine._incrementCompletionStat("_pollRunningTests_completed_none")
+            return completed
+        finally:
+            machine._recordCompletionInternalSpan(
+                "_pollRunningTests",
+                start_us,
+                time.time_ns() // 1000,
+                metadata={
+                    "mode": getattr(machine, "completion_detection_mode", ""),
+                    "allow_running_checks": bool(allow_running_checks),
+                    "ordered_count": ordered_count,
+                    "completion_limit": completion_limit,
+                    "completed_count": completed,
+                    "result": result_kind,
+                },
+            )
+
+
+    def check_running(self):
+        """Advance machine state using plain sleep-then-poll behavior.
+
+        Returns:
+            None: Completed tests may be finalized and removed from
+            ``machine.running``.
+        """
+        machine = self.machine
+        time.sleep(machine.naptime)
+        self.poll_running_tests(
+            allow_running_checks=True,
+        )
+
+    def logCompletionWarnings(self, logger):
+        """Polling mode has no detector-specific completion warnings."""
+        return None