stemflow.utils.execution_monitor

Install stemflow[monitoring] for psutil-backed process-memory measurements. The usual entry point is model.monitor_execution(...); ExecutionMonitor can also wrap a standalone workflow.

with model.monitor_execution(
    sample_interval=1.0,
    log_interval=30.0,
    output_path="data/execution_monitor.json",
    detailed_memory=False,
    max_samples=0,
) as monitor:
    model.fit(X_train, y_train, temporal_window_prequery=True)
    with monitor.phase("held_out_prediction"):
        probability = model.predict_proba(X_test, n_jobs=2)

report = model.execution_monitor_

Set log_interval=None to disable live logging. A new context replaces the previous model report. Save the model after leaving the context to retain the finalized report; the live thread and locks are not serialized with it. JSON output replaces one report file rather than accumulating snapshots. max_samples=0 retains summaries only; a positive value bounds the recent raw samples retained.

Report interpretation

Field	Meaning
wall_seconds	Elapsed time for the monitored context.
stages	Inclusive stage timings and sampled stage peaks; nested durations are not additive.
peak_tree_rss_bytes	Maximum sampled simultaneous sum of parent and descendant RSS on this node.
peak_process_count	Maximum sampled process count in the monitored tree.
sampling_error, output_error	Measurement/reporting failures, reported independently of model errors.
latest_sample, recent_samples	Latest reading and optionally bounded raw history.

RSS can count shared pages in multiple processes. With detailed_memory=True, supported USS, Linux PSS and macOS physical-footprint readings remain separate fields. Inaccessible/disappearing workers and incomplete coverage are explicit. Sampling can miss short peaks. These values are not GPU memory, remote-node measurements, Slurm/cgroup accounting, an exact allocator breakdown, or a memory limit.

See the executed monitoring example and measurement guide.

Monitor fitting, prediction and custom stages in the calling process.

Memory includes descendants on this node, including native estimator memory. RSS counts shared pages repeatedly. Optional USS, Linux PSS and macOS physical footprint are recorded separately when supported and accessible. These are not interchangeable or equal to a scheduler's memory charge. This is a sampler, not a memory limiter, allocation profiler or remote/GPU monitor.

Parameters:

• sample_interval (float, default: 1.0 ) –

  Seconds between background memory samples (default 1).

• log_interval (float | None, default: 30.0 ) –

  Seconds between flushed stdout updates; None disables logs. Stage completions and workflow start/end are also logged when enabled.

• output_path (str | Path | None, default: None ) –

  Optional compact JSON report. Updated alongside live logs and on exit, replacing one file rather than accumulating snapshots.

• detailed_memory (bool, default: False ) –

  Collect USS/PSS and physical footprint when available. Default False keeps cheap RSS sampling; detailed reads can be slower.

• max_samples (int, default: 0 ) –

  Retain at most this many recent samples; default 0 stores summaries only. Stage records and peaks survive trace truncation.

A context is single-use. Use stages from its calling thread. AdaSTEM models retain a plain report in execution_monitor_, updated after each outer model call and on context exit. The live sampler is never stored on the model.

Source code in stemflow/utils/execution_monitor.py

class ExecutionMonitor:
    """Monitor fitting, prediction and custom stages in the calling process.

    Memory includes descendants on this node, including native estimator memory.
    RSS counts shared pages repeatedly. Optional USS, Linux PSS and macOS
    physical footprint are recorded separately when supported and accessible.
    These are not interchangeable or equal to a scheduler's memory charge. This is a
    sampler, not a memory limiter, allocation profiler or remote/GPU monitor.

    Args:
        sample_interval (float): Seconds between background memory samples (default 1).
        log_interval (float | None): Seconds between flushed stdout updates; None disables logs.
            Stage completions and workflow start/end are also logged when enabled.
        output_path (str | pathlib.Path | None): Optional compact JSON report. Updated alongside live logs
            and on exit, replacing one file rather than accumulating snapshots.
        detailed_memory (bool): Collect USS/PSS and physical footprint when available.
            Default False keeps cheap RSS sampling; detailed reads can be slower.
        max_samples (int): Retain at most this many recent samples; default 0 stores
            summaries only. Stage records and peaks survive trace truncation.

    A context is single-use. Use stages from its calling thread. AdaSTEM models
    retain a plain report in execution_monitor_, updated after each outer model
    call and on context exit. The live sampler is never stored on the model.
    """

    def __init__(self, sample_interval=1.0, log_interval=30.0,
                 output_path=None, max_samples=0, detailed_memory=False):
        for name, value in [("sample_interval", sample_interval), ("log_interval", log_interval)]:
            if value is None and name == "log_interval":
                continue
            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0:
                raise ValueError(f"{name} must be a positive finite number" + (" or None." if name == "log_interval" else "."))
        if isinstance(max_samples, bool) or not isinstance(max_samples, int) or max_samples < 0:
            raise ValueError("max_samples must be a nonnegative integer.")
        if not isinstance(detailed_memory, bool):
            raise ValueError("detailed_memory must be a bool.")
        self.detailed_memory = detailed_memory
        self.sample_interval = float(sample_interval)
        self.log_interval = log_interval
        self.output_path = Path(output_path) if output_path is not None else None
        self.max_samples = max_samples
        self._lock = threading.RLock()
        self._stop = threading.Event()
        self._trace = deque(maxlen=max_samples)
        self._stages, self._stack = [], []
        self._started = None
        self._ended = None
        self._error = None
        self._sampling_error = None
        self._output_error = None
        self._thread = None
        self._models = []
        self._model_call_depth = 0

    def _track_model(self, model):
        if not any(reference() is model for reference in self._models):
            self._models.append(weakref.ref(model))

    def _publish_reports(self):
        report = self.summary()
        for reference in self._models:
            model = reference()
            if model is not None:
                model.execution_monitor_ = copy.deepcopy(report)

    @staticmethod
    def _statistics():
        return dict(samples=0, partial_rss_samples=0, peak_process_count=0,
                    **{f"peak_{name}": None for name in _memory_fields})

    def __enter__(self):
        if self._started is not None:
            raise RuntimeError("An ExecutionMonitor context can only be entered once.")
        if current_monitor() is not None:
            raise RuntimeError("Use monitor.phase(...) for nested stages, rather than nested monitors.")
        self.pid, self._owner_thread = os.getpid(), threading.get_ident()
        self._reader = _ProcessMemoryReader(self.detailed_memory)
        self._sampling_seconds = 0.0
        self._sampling_attempts = 0
        self._max_sampling_seconds = 0.0
        self._max_sample_gap_seconds = 0.0
        self._statistics_total = self._statistics()
        self._started = time.perf_counter()
        self._token = _active_monitor.set(self)
        with self._lock:
            self._sample()
        self._thread = threading.Thread(target=self._run, name="stemflow-memory-monitor", daemon=True)
        try:
            self._thread.start()
        except BaseException:
            _active_monitor.reset(self._token)
            raise
        self._log("started")
        return self

    def __exit__(self, error_type, error, traceback):
        self._stop.set()
        self._thread.join()
        with self._lock:
            self._sample()
            self._ended = time.perf_counter()
            self._error = error_type.__name__ if error_type is not None else None
        _active_monitor.reset(self._token)
        self._log("finished" if error_type is None else f"failed ({self._error})")
        if self.log_interval is None:
            self._save_report(self.summary())
        self._publish_reports()
        return False

    def _save_report(self, report):
        if self.output_path is None or self._output_error is not None:
            return
        with self._lock:
            temporary = self.output_path.with_name(self.output_path.name + ".partial")
            try:
                self.output_path.parent.mkdir(parents=True, exist_ok=True)
                temporary.write_text(json.dumps(report, indent=2), encoding="utf-8")
                temporary.replace(self.output_path)
            except OSError as caught:
                self._output_error = f"{type(caught).__name__}: {caught}"
                self._warn(f"Could not save execution monitor report: {caught}")

    @staticmethod
    def _warn(message):
        try:
            warnings.warn(message, RuntimeWarning, stacklevel=2)
        except Warning:
            pass  # A monitoring warning must not replace the model's exception.

    def _sample(self):
        started = time.perf_counter()
        self._sampling_attempts += 1
        try:
            sample = self._reader.snapshot(self.pid)
        except Exception as caught:
            self._sampling_error = f"{type(caught).__name__}: {caught}"
            return
        finally:
            duration = time.perf_counter() - started
            self._sampling_seconds += duration
            self._max_sampling_seconds = max(self._max_sampling_seconds, duration)
        sample["elapsed_seconds"] = time.perf_counter() - self._started
        if hasattr(self, "_latest"):
            self._max_sample_gap_seconds = max(self._max_sample_gap_seconds,
                sample["elapsed_seconds"] - self._latest["elapsed_seconds"])
        sample["stage"] = "/".join(stage["name"] for stage in self._stack) or "workflow"
        self._latest = sample
        self._trace.append(sample)
        for record in [self._statistics_total, *self._stack]:
            record["samples"] += 1
            record["partial_rss_samples"] += int(not sample["rss_complete"])
            record["peak_process_count"] = max(record["peak_process_count"], sample["process_count"])
            for name in _memory_fields:
                value, key = sample[name], f"peak_{name}"
                if value is not None and (record[key] is None or value > record[key]):
                    record[key] = value
                    if record is self._statistics_total and name == "tree_rss_bytes":
                        record["peak_rss_stage"] = sample["stage"]
                        record["peak_rss_elapsed_seconds"] = sample["elapsed_seconds"]

    def _run(self):
        next_log = time.perf_counter() + (self.log_interval or float("inf"))
        while not self._stop.wait(self.sample_interval):
            with self._lock:
                self._sample()
            now = time.perf_counter()
            if now >= next_log:
                self._log("running")
                next_log = now + self.log_interval

    def _log(self, message):
        if self.log_interval is None:
            return
        report = self.summary()
        self._save_report(report)
        sample = report["latest_sample"] or {}
        def mib(value):
            return "unavailable" if value is None else f"{value / 2**20:.1f} MiB"
        text = (f"[stemflow] {message}; elapsed={report['wall_seconds']:.2f}s; "
                f"stage={sample.get('stage', 'workflow')}; "
                f"RSS={mib(sample.get('tree_rss_bytes'))}; "
                f"RSS coverage={sample.get('rss_processes', 0)}/{sample.get('process_count', '?')}"
                f"{' (partial tree)' if not sample.get('rss_complete') else ''}; "
                f"peak RSS={mib(report['peak_tree_rss_bytes'])}; "
                f"parent={mib(sample.get('parent_rss_bytes'))}; "
                f"children={mib(sample.get('children_rss_bytes'))}")
        try:
            print(text, flush=True)
        except (OSError, ValueError):
            pass

    @contextmanager
    def phase(self, name):
        """Record an inclusive stage, including nested calls and worker memory."""
        if current_monitor() is not self or self._ended is not None:
            raise RuntimeError("Stages must run inside this monitor in its calling thread/process.")
        if not isinstance(name, str) or not name:
            raise ValueError("A stage name must be a nonempty string.")
        with self._lock:
            stage = dict(id=len(self._stages), name=name,
                         parent_id=self._stack[-1]["id"] if self._stack else None,
                         start_seconds=time.perf_counter() - self._started,
                         wall_seconds=None, error=None, **self._statistics())
            self._stages.append(stage)
            self._stack.append(stage)
            self._sample()
        try:
            yield self
        except BaseException as caught:
            stage["error"] = type(caught).__name__
            raise
        finally:
            with self._lock:
                self._sample()
                stage["wall_seconds"] = time.perf_counter() - self._started - stage["start_seconds"]
                self._stack.pop()
            self._log(f"completed {name} ({stage['wall_seconds']:.2f}s)" if stage["error"] is None
                      else f"failed {name} ({stage['error']})")

    def summary(self):
        """Return a JSON-compatible snapshot, also available during execution."""
        if self._started is None:
            raise RuntimeError("Enter the monitor context before requesting its summary.")
        with self._lock:
            now = self._ended or time.perf_counter()
            stages = [dict(stage) for stage in self._stages]
            for stage in stages:
                if stage["wall_seconds"] is None:
                    stage["wall_seconds"] = now - self._started - stage["start_seconds"]
            return dict(pid=self.pid, running=self._ended is None, error=self._error,
                        wall_seconds=now - self._started, sample_interval=self.sample_interval,
                        platform=platform.platform(), psutil_version=self._reader.psutil.__version__,
                        detailed_memory=self.detailed_memory,
                        sampling_attempts=self._sampling_attempts,
                        sampling_seconds=self._sampling_seconds,
                        max_sampling_seconds=self._max_sampling_seconds,
                        max_sample_gap_seconds=self._max_sample_gap_seconds,
                        sampling_error=self._sampling_error, output_error=self._output_error,
                        **self._statistics_total,
                        stages=stages, latest_sample=dict(self._latest) if hasattr(self, "_latest") else None,
                        recent_samples=[dict(sample) for sample in self._trace])

phase(name)

Record an inclusive stage, including nested calls and worker memory.

Source code in stemflow/utils/execution_monitor.py

@contextmanager
def phase(self, name):
    """Record an inclusive stage, including nested calls and worker memory."""
    if current_monitor() is not self or self._ended is not None:
        raise RuntimeError("Stages must run inside this monitor in its calling thread/process.")
    if not isinstance(name, str) or not name:
        raise ValueError("A stage name must be a nonempty string.")
    with self._lock:
        stage = dict(id=len(self._stages), name=name,
                     parent_id=self._stack[-1]["id"] if self._stack else None,
                     start_seconds=time.perf_counter() - self._started,
                     wall_seconds=None, error=None, **self._statistics())
        self._stages.append(stage)
        self._stack.append(stage)
        self._sample()
    try:
        yield self
    except BaseException as caught:
        stage["error"] = type(caught).__name__
        raise
    finally:
        with self._lock:
            self._sample()
            stage["wall_seconds"] = time.perf_counter() - self._started - stage["start_seconds"]
            self._stack.pop()
        self._log(f"completed {name} ({stage['wall_seconds']:.2f}s)" if stage["error"] is None
                  else f"failed {name} ({stage['error']})")

summary()

Return a JSON-compatible snapshot, also available during execution.

Source code in stemflow/utils/execution_monitor.py

def summary(self):
    """Return a JSON-compatible snapshot, also available during execution."""
    if self._started is None:
        raise RuntimeError("Enter the monitor context before requesting its summary.")
    with self._lock:
        now = self._ended or time.perf_counter()
        stages = [dict(stage) for stage in self._stages]
        for stage in stages:
            if stage["wall_seconds"] is None:
                stage["wall_seconds"] = now - self._started - stage["start_seconds"]
        return dict(pid=self.pid, running=self._ended is None, error=self._error,
                    wall_seconds=now - self._started, sample_interval=self.sample_interval,
                    platform=platform.platform(), psutil_version=self._reader.psutil.__version__,
                    detailed_memory=self.detailed_memory,
                    sampling_attempts=self._sampling_attempts,
                    sampling_seconds=self._sampling_seconds,
                    max_sampling_seconds=self._max_sampling_seconds,
                    max_sample_gap_seconds=self._max_sample_gap_seconds,
                    sampling_error=self._sampling_error, output_error=self._output_error,
                    **self._statistics_total,
                    stages=stages, latest_sample=dict(self._latest) if hasattr(self, "_latest") else None,
                    recent_samples=[dict(sample) for sample in self._trace])