From 6adfbb463b34f3d3ef463ffd21c32df689c06908 Mon Sep 17 00:00:00 2001 From: "Roger Liu (Content Tech)" Date: Sun, 14 Jun 2026 23:31:21 -0700 Subject: [PATCH] Add logging filter to exclude per-task subprocess output from sflow.log - Introduced a new logging filter, `_DropTaskStreamFilter`, to prevent per-task subprocess output from being written to sflow.log, ensuring that only orchestration lines and command/status hints are logged. - Updated logging configuration to apply this filter to file handlers. - Enhanced the `SubprocessLauncher` to conditionally echo per-task output to the root logger based on whether the output is being streamed to an interactive console. - Added unit tests to verify the correct routing of per-task output and the functionality of the logging filter. --- src/sflow/cli/batch.py | 26 ++++- src/sflow/core/launcher.py | 64 +++++++++-- src/sflow/logging.py | 23 ++++ tests/unit/test_cli_batch.py | 47 ++++++++ tests/unit/test_launcher_logging_routing.py | 118 ++++++++++++++++++++ tests/unit/test_logging_sflow_log_filter.py | 47 ++++++++ 6 files changed, 317 insertions(+), 8 deletions(-) create mode 100644 tests/unit/test_launcher_logging_routing.py create mode 100644 tests/unit/test_logging_sflow_log_filter.py diff --git a/src/sflow/cli/batch.py b/src/sflow/cli/batch.py index be87ad1..f65fca4 100644 --- a/src/sflow/cli/batch.py +++ b/src/sflow/cli/batch.py @@ -361,6 +361,16 @@ def _generate_sbatch_script( "", "set -x", "", + "# sbatch defaults to --export=ALL, so an activated virtualenv from the submitting", + "# shell leaks into the job. Drop it from PATH and the environment so the venv is", + "# bootstrapped with a real system interpreter for the compute node's architecture", + "# (a wrong-arch caller venv otherwise fails with 'Exec format error').", + 'if [ -n "${VIRTUAL_ENV:-}" ]; then', + " PATH=$(printf '%s' \"$PATH\" | tr ':' '\\n' | grep -vxF \"$VIRTUAL_ENV/bin\" | paste -sd ':' -)", + " export PATH", + "fi", + "unset VIRTUAL_ENV VIRTUAL_ENV_PROMPT PYTHONHOME", + "", ] if sflow_venv_path: @@ -400,7 +410,21 @@ def _generate_sbatch_script( "else", " # Venv not found; create from scratch and install sflow", f" cd {venv_parent}", - " python3 -m venv .sflow_venv", + " # Resolve a real system python3 for venv creation. Prefer well-known", + " # absolute locations, then fall back to PATH (already cleaned of the", + " # caller venv above) so nodes that install python outside /usr/bin work.", + ' SFLOW_BOOTSTRAP_PYTHON=""', + ' for _candidate in /usr/bin/python3 /usr/local/bin/python3 "$(command -v python3 || true)"; do', + ' if [ -n "$_candidate" ] && [ -x "$_candidate" ]; then', + ' SFLOW_BOOTSTRAP_PYTHON="$_candidate"', + " break", + " fi", + " done", + ' if [ -z "$SFLOW_BOOTSTRAP_PYTHON" ]; then', + ' echo "ERROR: could not locate a system python3 to bootstrap the sflow venv"', + " exit 1", + " fi", + ' "$SFLOW_BOOTSTRAP_PYTHON" -m venv .sflow_venv', " source .sflow_venv/bin/activate", ' "$VIRTUAL_ENV/bin/pip" install uv', f' "$VIRTUAL_ENV/bin/uv" pip install {sflow_install_cmd}', diff --git a/src/sflow/core/launcher.py b/src/sflow/core/launcher.py index 9ff3ad2..a3d33cb 100644 --- a/src/sflow/core/launcher.py +++ b/src/sflow/core/launcher.py @@ -7,15 +7,50 @@ import pty import shlex import subprocess +import sys from typing import Mapping, Optional -from sflow.logging import get_logger +from sflow.logging import SFLOW_TASK_STREAM_ATTR, get_logger from .command import Command, format_command _logger = get_logger(__name__) +def _console_streams_task_output() -> bool: + """Whether per-task subprocess output should be echoed to the root logger. + + True only for interactive TTY sessions. In batch mode (Slurm redirects + stdout/stderr to files) or when output is piped to a file, this returns + False so task content never floods sflow.log or the Slurm stdout/err files. + Per-task content is always written to the per-task log regardless. + """ + try: + return bool(sys.stdout.isatty()) + except Exception: + return False + + +def _emit_task_output_line( + line_str: str, + *, + pfx: str, + output_logger: Optional[logging.Logger], + stream_console: bool, +) -> None: + """Route a single line of task (subprocess) output. + + Always writes to the per-task logger (``output_logger`` -> ``.log``). + Only additionally echoes to the root ``sflow`` logger (console / Slurm + stdout) when ``stream_console`` is True, tagging the record with + ``SFLOW_TASK_STREAM_ATTR`` so sflow.log's file handler drops it. + """ + if output_logger: + output_logger.info(line_str) + if stream_console: + _logger.info(f"{pfx}{line_str}", extra={SFLOW_TASK_STREAM_ATTR: True}) + + def _strip_ansi(text: str) -> str: """Strip ANSI escape sequences from text.""" import re @@ -76,9 +111,18 @@ async def run_async( task_name: str | None = None, ) -> int: pfx = self._console_prefix(task_name) + # Orchestration hint: goes to sflow.log + console (and Slurm stdout/err). + # Intentionally NOT written to the per-task log: LogWatchProbe scans the + # per-task log file and would false-match patterns inside the command text. _logger.info(f"{pfx}========== Command ==========") _logger.info(f"{pfx}{format_command(command)}") _logger.info(f"{pfx}=============================") + + # Decide once whether to echo per-task output to the root logger/console. + # Per-task content always goes to the per-task log (output_logger); it is + # only additionally streamed to the console for interactive TTY sessions. + stream_console = _console_streams_task_output() + if isinstance(command, Command): command = command.as_list() @@ -160,9 +204,12 @@ async def run_async( # Strip ANSI escape sequences for cleaner logs line_str = _strip_ansi(line_str).rstrip() if line_str: - _logger.info(f"{pfx}{line_str}") - if output_logger: - output_logger.info(line_str) + _emit_task_output_line( + line_str, + pfx=pfx, + output_logger=output_logger, + stream_console=stream_console, + ) if ret is not None and not chunk: # Process exited and no more data @@ -171,9 +218,12 @@ async def run_async( buffer.decode("utf-8", errors="replace") ).rstrip() if line_str: - _logger.info(f"{pfx}{line_str}") - if output_logger: - output_logger.info(line_str) + _emit_task_output_line( + line_str, + pfx=pfx, + output_logger=output_logger, + stream_console=stream_console, + ) break if not chunk: diff --git a/src/sflow/logging.py b/src/sflow/logging.py index 6ed096a..b8f5386 100644 --- a/src/sflow/logging.py +++ b/src/sflow/logging.py @@ -11,6 +11,25 @@ # Default width for non-interactive terminals (piping, file output, etc.) _DEFAULT_NON_TTY_WIDTH = 200 +# Log records carrying this attribute (set to True) are per-task subprocess +# output. They may be streamed to an interactive console, but must never be +# written to sflow.log (or the Slurm stdout/err files). Those sinks are reserved +# for orchestration lines and command/status hints; full per-task content lives +# in /.log. +SFLOW_TASK_STREAM_ATTR = "sflow_task_stream" + + +class _DropTaskStreamFilter(logging.Filter): + """Drop per-task subprocess output records so they never reach a file handler. + + Records produced by streaming a task's stdout/stderr are tagged with + ``SFLOW_TASK_STREAM_ATTR``. This filter keeps them out of sflow.log while + still allowing them through to the interactive console handler. + """ + + def filter(self, record: logging.LogRecord) -> bool: + return not getattr(record, SFLOW_TASK_STREAM_ATTR, False) + def configure_logging( level: str = "INFO", log_file: Optional[str] = None, *, console: bool = True @@ -46,6 +65,8 @@ def configure_logging( file_handler.setFormatter( logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") ) + # Per-task subprocess output must never land in the log file. + file_handler.addFilter(_DropTaskStreamFilter()) handlers.append(file_handler) # Configure the sflow logger @@ -82,6 +103,8 @@ def add_log_file(log_file: str) -> None: fh.setFormatter( logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") ) + # Keep per-task subprocess output out of sflow.log (orchestration only). + fh.addFilter(_DropTaskStreamFilter()) logger.addHandler(fh) # Ensure the logger itself accepts INFO messages even if the console diff --git a/tests/unit/test_cli_batch.py b/tests/unit/test_cli_batch.py index b9ad628..60d1bd3 100644 --- a/tests/unit/test_cli_batch.py +++ b/tests/unit/test_cli_batch.py @@ -3250,3 +3250,50 @@ def test_batch_defaults_sflow_version_from_execution_env( "git+https://github.com/NVIDIA/nv-sflow.git@feature/infmax_v3" in script_content ) + + +def test_batch_script_bootstraps_venv_with_resolved_system_python( + mock_sflow_app, temp_workflow_file, tmp_path +): + """The generated sbatch venv bootstrap must not use a bare ``python3``. + + A bare ``python3`` resolves through PATH, so if the submitting shell has a + virtualenv activated, sbatch's default ``--export=ALL`` leaks that venv into + the job and ``python3 -m venv`` runs the caller's interpreter -- which fails + with "Exec format error" when the login node and compute node differ in + architecture (e.g. x86 login vs aarch64 Grace compute). + + The script must instead clear the inherited virtualenv and resolve a real + system python3: prefer well-known absolute locations, then fall back to a + PATH-resolved interpreter so nodes that install python outside /usr/bin work. + """ + sbatch_path = tmp_path / "test.sh" + result = runner.invoke( + app, + [ + "batch", + "--file", + str(temp_workflow_file), + "--partition", + "batch", + "--account", + "testaccount", + "--nodes", + "1", + "--sbatch-path", + str(sbatch_path), + ], + ) + + assert result.exit_code == 0, f"CLI failed: {result.output}" + script_content = sbatch_path.read_text() + # Venv is created via a resolved interpreter, never a bare python3 from PATH. + assert '"$SFLOW_BOOTSTRAP_PYTHON" -m venv .sflow_venv' in script_content + assert " python3 -m venv .sflow_venv" not in script_content + # Well-known absolute location is tried first, with a PATH fallback for nodes + # that install python elsewhere. + assert "/usr/bin/python3" in script_content + assert "command -v python3" in script_content + # An inherited (possibly wrong-arch) virtualenv must be neutralized. + assert "unset VIRTUAL_ENV" in script_content + assert 'grep -vxF "$VIRTUAL_ENV/bin"' in script_content diff --git a/tests/unit/test_launcher_logging_routing.py b/tests/unit/test_launcher_logging_routing.py new file mode 100644 index 0000000..1cfa383 --- /dev/null +++ b/tests/unit/test_launcher_logging_routing.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for how SubprocessLauncher routes per-task output to loggers. + +Per-task subprocess output must always be written to the per-task logger +(``output_logger``), but only additionally echoed to the root ``sflow`` logger +(console / slurm stdout) for interactive TTY sessions, and only when tagged with +``SFLOW_TASK_STREAM_ATTR`` so sflow.log's file handler can drop it. +""" + +import logging +from contextlib import contextmanager + +import sflow.core.launcher as launcher_mod +from sflow.core.launcher import ( + _console_streams_task_output, + _emit_task_output_line, +) +from sflow.logging import SFLOW_TASK_STREAM_ATTR + + +class _CaptureHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.NOTSET) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +@contextmanager +def _capture_root_logger(): + """Temporarily replace the ``sflow.core.launcher`` logger handlers with a + capture handler to observe exactly what reaches the root (console/sflow.log). + """ + logger = launcher_mod._logger + saved = (logger.handlers, logger.level, logger.propagate) + cap = _CaptureHandler() + logger.handlers = [cap] + logger.setLevel(logging.INFO) + logger.propagate = False + try: + yield cap + finally: + logger.handlers, logger.level, logger.propagate = saved + + +def _make_output_logger() -> tuple[logging.Logger, _CaptureHandler]: + out = logging.getLogger("sflow.task._test_routing") + out.handlers = [] + cap = _CaptureHandler() + out.addHandler(cap) + out.setLevel(logging.INFO) + out.propagate = False + return out, cap + + +def test_task_line_not_echoed_to_root_when_not_streaming(): + out_logger, task_cap = _make_output_logger() + + with _capture_root_logger() as root_cap: + _emit_task_output_line( + "TASK_OUTPUT_LINE", + pfx="[t] ", + output_logger=out_logger, + stream_console=False, + ) + + # Per-task content is always written to the per-task logger. + assert [r.getMessage() for r in task_cap.records] == ["TASK_OUTPUT_LINE"] + # Non-TTY (batch): nothing echoed to the root logger. + assert root_cap.records == [] + + +def test_task_line_echoed_to_root_with_marker_when_streaming(): + out_logger, task_cap = _make_output_logger() + + with _capture_root_logger() as root_cap: + _emit_task_output_line( + "TASK_OUTPUT_LINE", + pfx="[t] ", + output_logger=out_logger, + stream_console=True, + ) + + # Per-task content still written to the per-task logger. + assert any("TASK_OUTPUT_LINE" in r.getMessage() for r in task_cap.records) + # TTY: echoed to the root logger and tagged as task-stream (so file handlers + # drop it), with the console task prefix applied. + assert len(root_cap.records) == 1 + rec = root_cap.records[0] + assert getattr(rec, SFLOW_TASK_STREAM_ATTR, False) is True + assert rec.getMessage() == "[t] TASK_OUTPUT_LINE" + + +def test_emit_without_output_logger_is_safe(): + # Should not raise when there is no per-task logger. + with _capture_root_logger() as root_cap: + _emit_task_output_line( + "X", pfx="", output_logger=None, stream_console=True + ) + assert len(root_cap.records) == 1 + + +def test_console_streams_task_output_follows_tty(monkeypatch): + class _FakeStdout: + def __init__(self, tty: bool) -> None: + self._tty = tty + + def isatty(self) -> bool: + return self._tty + + monkeypatch.setattr(launcher_mod.sys, "stdout", _FakeStdout(True)) + assert _console_streams_task_output() is True + + monkeypatch.setattr(launcher_mod.sys, "stdout", _FakeStdout(False)) + assert _console_streams_task_output() is False diff --git a/tests/unit/test_logging_sflow_log_filter.py b/tests/unit/test_logging_sflow_log_filter.py new file mode 100644 index 0000000..50fdd52 --- /dev/null +++ b/tests/unit/test_logging_sflow_log_filter.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests that the sflow.log file handler drops per-task subprocess output. + +``add_log_file`` attaches a filter so that records tagged with +``SFLOW_TASK_STREAM_ATTR`` (per-task subprocess output) never land in sflow.log, +which is reserved for orchestration lines and command/status hints. +""" + +import logging + +from sflow.logging import SFLOW_TASK_STREAM_ATTR, add_log_file + + +def test_add_log_file_drops_task_stream_records(tmp_path): + log_file = tmp_path / "sflow.log" + logger = logging.getLogger("sflow") + saved = (list(logger.handlers), logger.level, logger.propagate) + + try: + add_log_file(str(log_file)) + logger.setLevel(logging.INFO) + logger.propagate = False + + logger.info("ORCHESTRATION_MARKER from orchestrator") + logger.info( + "TASKSTREAM_MARKER from a noisy task", + extra={SFLOW_TASK_STREAM_ATTR: True}, + ) + + for h in logger.handlers: + h.flush() + content = log_file.read_text() + finally: + for h in list(logger.handlers): + if ( + isinstance(h, logging.FileHandler) + and getattr(h, "baseFilename", None) == str(log_file) + ): + logger.removeHandler(h) + h.close() + logger.handlers, logger.level, logger.propagate = saved + + # Orchestration lines are kept; per-task stream lines are dropped. + assert "ORCHESTRATION_MARKER" in content + assert "TASKSTREAM_MARKER" not in content