From 1316e60f54afb0be6d74ffae04a787c34ae774fa Mon Sep 17 00:00:00 2001 From: mason h <9421505+drawadiagram@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:19:40 -0500 Subject: [PATCH] Harden ImpressManager lifecycle, add GPU discovery, filterable logging Core framework changes extracted from the impress_fixes branch. ImpressManager: - Write asyncflow session dirs to IMPRESS_SESSION_DIR (default: tempdir) instead of cwd, which exhausted scratch inodes over many HPC runs. - Shut the WorkflowEngine down in a finally block so a pipeline raising mid-run no longer leaks the engine and its backend. This is where the per-example `await manager.flow.shutdown()` call went; the examples no longer have to remember it. - Guard submit_new_pipelines() with a clear RuntimeError when called before start(), instead of failing on a missing self.flow attribute. - Carry each pipeline's future through cleanup so a failed pipeline is reported via logger.pipeline_failed() rather than logged as completed. - Capture the buffered-pipeline count before clearing the buffer, so activity_summary() stops always reporting 0 new submissions. ImpressLogger: - Add min_level filtering to every level method and pipeline_log(). - Add pipeline_failed(name, exc). - Route error/critical to output_stream rather than stderr: under Dragon the job's stdout is the reviewable log, and errors were being split into a separate stream nobody reads. Other: - New impress.gpu.find_gpus(), re-exported from the package root: CUDA_VISIBLE_DEVICES, else nvidia-smi, else an empty list. - ImpressBasePipeline.finalize() is no longer abstract; pipelines that need no cleanup should not have to define an empty override. - Import Callable from collections.abc (deprecated in typing). - Pin radical-asyncflow>=0.4.0: WorkflowEngine.create(work_dir=...) does not exist in 0.3.1 and earlier, so the session-dir change above needs a floor on the dependency. Python 3.9/3.10 support is retained, per review of PR #59. That means keeping `except asyncio.TimeoutError` in test_pipeline_management.py -- asyncio.TimeoutError only became an alias of the builtin TimeoutError in 3.11, so the shorter spelling silently fails to catch the timeout on 3.9 and 3.10. Nothing else in src/ needs 3.10+: the generics here are PEP 585, which 3.9 evaluates fine, and radical-asyncflow itself declares requires-python >=3.9. Adds unit coverage for the logger, the pipeline base class, and PipelineSetup, plus regression tests for the shutdown and submit-before-start fixes. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Lo8DwSbyvdWZRkkkka6gA2 --- .gitignore | 7 + examples/.gitignore | 14 + pyproject.toml | 2 +- src/impress/__init__.py | 4 +- src/impress/gpu.py | 32 ++ src/impress/impress_manager.py | 235 ++++++++------- src/impress/pipelines/impress_pipeline.py | 3 +- src/impress/pipelines/setup.py | 4 +- src/impress/utils/logger.py | 51 +++- tests/conftest.py | 5 +- tests/unit/test_logger.py | 278 ++++++++++++++++++ tests/unit/test_manager_core.py | 3 +- tests/unit/test_manager_life_cycle.py | 20 +- .../unit/test_manager_pipeline_submission.py | 9 + tests/unit/test_pipeline_base.py | 157 ++++++++++ tests/unit/test_pipeline_setup.py | 141 +++++++++ 16 files changed, 843 insertions(+), 122 deletions(-) create mode 100644 examples/.gitignore create mode 100644 src/impress/gpu.py create mode 100644 tests/unit/test_logger.py create mode 100644 tests/unit/test_pipeline_base.py create mode 100644 tests/unit/test_pipeline_setup.py diff --git a/.gitignore b/.gitignore index 0dafbe4..01d1dc9 100644 --- a/.gitignore +++ b/.gitignore @@ -135,6 +135,11 @@ dmypy.json # asyncflow related asyncflow.session.* +ddict_* + +# ROME runtime outputs +af_stats_*.csv +examples/protien_binding_usecase/logs/ # pdzbinder wf outputs af_pipeline_outputs_multi/ @@ -146,3 +151,5 @@ ddict* b0 slurm* *slurm +# scratch archives +arch/ diff --git a/examples/.gitignore b/examples/.gitignore new file mode 100644 index 0000000..31e9608 --- /dev/null +++ b/examples/.gitignore @@ -0,0 +1,14 @@ +# SLURM log output (all workflows) +*/logs/ + +# Per-pipeline task directories written to cwd +*/p*/ + +# Dragon telemetry session dirs +*/telemetry/ + +# Run metadata written to cwd +*/runinfo + +# Legacy output dirs +*/myoutputs/ diff --git a/pyproject.toml b/pyproject.toml index 9daa18b..bf1d9e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ requires-python = ">=3.9" dependencies = [ "radical.pilot", - "radical-asyncflow" + "radical-asyncflow>=0.4.0" ] [project.urls] diff --git a/src/impress/__init__.py b/src/impress/__init__.py index 89dd8d3..6a37de7 100644 --- a/src/impress/__init__.py +++ b/src/impress/__init__.py @@ -1,10 +1,10 @@ -from __future__ import annotations - +from impress.gpu import find_gpus from impress.impress_manager import ImpressManager from impress.pipelines.impress_pipeline import ImpressBasePipeline from impress.pipelines.setup import PipelineSetup __all__ = [ + "find_gpus", "ImpressManager", "ImpressBasePipeline", "PipelineSetup", diff --git a/src/impress/gpu.py b/src/impress/gpu.py new file mode 100644 index 0000000..42583d3 --- /dev/null +++ b/src/impress/gpu.py @@ -0,0 +1,32 @@ +import os +import subprocess + + +def find_gpus() -> list[int]: + """Return GPU IDs available to this process. + + Checks CUDA_VISIBLE_DEVICES first, then nvidia-smi. + Falls back to an empty list when neither yields results. + """ + val = os.environ.get("CUDA_VISIBLE_DEVICES", "") + ids = [int(g) for g in val.split(",") if g.strip().isdigit()] + if ids: + return ids + + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ) + if out.returncode == 0: + return [ + int(ln.strip()) + for ln in out.stdout.splitlines() + if ln.strip().isdigit() + ] + except Exception: + pass + + return [] diff --git a/src/impress/impress_manager.py b/src/impress/impress_manager.py index 5fac46e..b9446f7 100644 --- a/src/impress/impress_manager.py +++ b/src/impress/impress_manager.py @@ -1,6 +1,8 @@ import asyncio -from collections.abc import Awaitable -from typing import Any, Callable, Optional, Union +import os +import tempfile +from collections.abc import Awaitable, Callable +from typing import Any, Optional, Union from radical.asyncflow import WorkflowEngine @@ -43,6 +45,7 @@ def __init__( self._telemetry_config: dict[str, Any] = telemetry_config or {} self._telemetry_subscribers: list[Callable] = telemetry_subscribers or [] self.telemetry: Any = None + self.flow: Optional[WorkflowEngine] = None def _normalize_pipeline_setup( self, setup: Union[dict[str, Any], PipelineSetup] @@ -77,6 +80,10 @@ def submit_new_pipelines( ValueError: If pipeline type is not a subclass of ImpressBasePipeline """ + if self.flow is None: + raise RuntimeError( + "ImpressManager.start() must be called before submit_new_pipelines()" + ) for setup_input in pipeline_setups: # Normalize to PipelineSetup object setup = self._normalize_pipeline_setup(setup_input) @@ -136,113 +143,139 @@ async def start( """ self.logger.separator("IMPRESS MANAGER STARTING") - self.flow: WorkflowEngine = await WorkflowEngine.create( - backend=self.execution_backend + # Write asyncflow session dirs to /tmp (node-local, no quota) instead of + # cwd on scratch, which exhausts inodes over many runs. + _session_base = os.environ.get("IMPRESS_SESSION_DIR", tempfile.gettempdir()) + self.flow = await WorkflowEngine.create( + backend=self.execution_backend, + work_dir=_session_base, ) - if self._telemetry_config: - self.telemetry = await self.flow.start_telemetry(**self._telemetry_config) - for fn in self._telemetry_subscribers: - self.telemetry.subscribe(fn) - - self.logger.manager_starting(len(pipeline_setups)) - - self.submit_new_pipelines(pipeline_setups) - - while True: - any_activity: bool = False - completed_pipelines: list[ImpressBasePipeline] = [] - - for pipeline, pipeline_future in list(self.pipeline_tasks.items()): - # Check if pipeline needs adaptive step and isn't already running one - if ( - getattr(pipeline, "invoke_adaptive_step", False) - and pipeline not in self.adaptive_tasks - ): - adaptive_task: asyncio.Task = asyncio.create_task( - self._run_adaptive_fn(pipeline) + try: + if self._telemetry_config: + self.telemetry = await self.flow.start_telemetry( + **self._telemetry_config + ) + for fn in self._telemetry_subscribers: + self.telemetry.subscribe(fn) + + self.logger.manager_starting(len(pipeline_setups)) + + self.submit_new_pipelines(pipeline_setups) + + while True: + any_activity: bool = False + completed_pipelines: list[tuple] = [] + + for pipeline, pipeline_future in list(self.pipeline_tasks.items()): + # Check if pipeline needs adaptive step and isn't running one yet + if ( + getattr(pipeline, "invoke_adaptive_step", False) + and pipeline not in self.adaptive_tasks + ): + adaptive_task: asyncio.Task = asyncio.create_task( + self._run_adaptive_fn(pipeline) + ) + self.adaptive_tasks[pipeline] = adaptive_task + any_activity = True + + # Check if pipeline has new config ready + config: Optional[dict[str, Any]] = ( + pipeline.get_child_pipeline_request() ) - self.adaptive_tasks[pipeline] = adaptive_task - any_activity = True - # Check if pipeline has new config ready - config: Optional[dict[str, Any]] = pipeline.get_child_pipeline_request() - - if config: - self.logger.child_pipeline_submitted(config["name"], pipeline.name) - # Convert dict to PipelineSetup for consistency - child_setup = PipelineSetup.from_dict(config) - self.new_pipeline_buffer.append(child_setup) - any_activity = True + if config: + self.logger.child_pipeline_submitted( + config["name"], pipeline.name + ) + # Convert dict to PipelineSetup for consistency + child_setup = PipelineSetup.from_dict(config) + self.new_pipeline_buffer.append(child_setup) + any_activity = True + + # Check if parent should be killed + if getattr(pipeline, "kill_parent", False): + self.logger.pipeline_killed(pipeline.name) + pipeline_future.cancel() + completed_pipelines.append((pipeline, pipeline_future)) + continue - # Check if parent should be killed - if getattr(pipeline, "kill_parent", False): - self.logger.pipeline_killed(pipeline.name) - pipeline_future.cancel() - completed_pipelines.append(pipeline) - continue - - # Check if pipeline is done - but only mark as completed - # if adaptive task is also done - if pipeline_future.done(): - # If there's an adaptive task running, don't mark as completed yet + # Check if pipeline is done - but only mark as completed + # if adaptive task is also done + if pipeline_future.done(): + # Adaptive task still running — wait before marking completed + if pipeline in self.adaptive_tasks: + adaptive_task = self.adaptive_tasks[pipeline] + if not adaptive_task.done(): + continue + + completed_pipelines.append((pipeline, pipeline_future)) + + # Clean up completed pipelines - but only if their + # adaptive tasks are also done + actually_completed: list[ImpressBasePipeline] = [] + for pipeline, future in completed_pipelines: + # Double-check: only clean up if adaptive task is + # done or doesn't exist if pipeline in self.adaptive_tasks: adaptive_task = self.adaptive_tasks[pipeline] if not adaptive_task.done(): continue + self.adaptive_tasks.pop(pipeline) + + self.pipeline_tasks.pop(pipeline, None) + exc = None + if future.done() and not future.cancelled(): + try: + exc = future.exception() + except Exception: + pass + if exc is not None: + self.logger.pipeline_failed(pipeline.name, exc) + else: + self.logger.pipeline_completed(pipeline.name) + actually_completed.append(pipeline) + + completed_pipelines = actually_completed + + # Clean up completed adaptive tasks + completed_adaptive: list[ImpressBasePipeline] = [] + for pipeline, adaptive_task in list(self.adaptive_tasks.items()): + if adaptive_task.done(): + completed_adaptive.append(pipeline) + + for pipeline in completed_adaptive: + self.adaptive_tasks.pop(pipeline, None) + + # Submit new pipelines; capture count before clearing so + # activity_summary reports the real number submitted. + if self.new_pipeline_buffer: + buffered_count = len(self.new_pipeline_buffer) + self.submit_new_pipelines(self.new_pipeline_buffer) + self.new_pipeline_buffer.clear() + any_activity = True + else: + buffered_count = 0 + + # Log activity summary periodically + if any_activity: + self.logger.activity_summary( + len(self.pipeline_tasks), + len(self.adaptive_tasks), + buffered_count, + ) - completed_pipelines.append(pipeline) - - # Clean up completed pipelines - but only if their - # adaptive tasks are also done - actually_completed: list[ImpressBasePipeline] = [] - for pipeline in completed_pipelines: - # Double-check: only clean up if adaptive task is - # done or doesn't exist - if pipeline in self.adaptive_tasks: - adaptive_task = self.adaptive_tasks[pipeline] - if not adaptive_task.done(): - continue - self.adaptive_tasks.pop(pipeline) - - self.pipeline_tasks.pop(pipeline, None) - self.logger.pipeline_completed(pipeline.name) - actually_completed.append(pipeline) - - completed_pipelines = actually_completed - - # Clean up completed adaptive tasks - completed_adaptive: list[ImpressBasePipeline] = [] - for pipeline, adaptive_task in list(self.adaptive_tasks.items()): - if adaptive_task.done(): - completed_adaptive.append(pipeline) - - for pipeline in completed_adaptive: - self.adaptive_tasks.pop(pipeline, None) - - # Submit new pipelines - if self.new_pipeline_buffer: - self.submit_new_pipelines(self.new_pipeline_buffer) - self.new_pipeline_buffer.clear() - any_activity = True - - # Log activity summary periodically - if any_activity: - self.logger.activity_summary( - len(self.pipeline_tasks), - len(self.adaptive_tasks), - len(self.new_pipeline_buffer), - ) + # Exit condition + if ( + not self.pipeline_tasks + and not self.new_pipeline_buffer + and not self.adaptive_tasks + ): + self.logger.manager_exiting() + self.logger.separator("IMPRESS MANAGER FINISHED") + break - # Exit condition - if ( - not self.pipeline_tasks - and not self.new_pipeline_buffer - and not self.adaptive_tasks - ): - self.logger.manager_exiting() - self.logger.separator("IMPRESS MANAGER FINISHED") - break - - if not any_activity: - await asyncio.sleep(0.5) + if not any_activity: + await asyncio.sleep(0.5) + finally: + await self.flow.shutdown() diff --git a/src/impress/pipelines/impress_pipeline.py b/src/impress/pipelines/impress_pipeline.py index f856fab..5279384 100644 --- a/src/impress/pipelines/impress_pipeline.py +++ b/src/impress/pipelines/impress_pipeline.py @@ -95,8 +95,7 @@ async def get_scores_map(self): """Optional: Return scores mapping""" return {} - @abstractmethod - async def finalize(self): + async def finalize(self): # noqa: B027 """Optional: Cleanup or finalization logic""" pass diff --git a/src/impress/pipelines/setup.py b/src/impress/pipelines/setup.py index e0b63e0..4e10036 100644 --- a/src/impress/pipelines/setup.py +++ b/src/impress/pipelines/setup.py @@ -1,5 +1,5 @@ -from collections.abc import Awaitable -from typing import Annotated, Any, Callable, Optional +from collections.abc import Awaitable, Callable +from typing import Annotated, Any, Optional from pydantic import BaseModel, Field, field_validator diff --git a/src/impress/utils/logger.py b/src/impress/utils/logger.py index 642d699..2919c8f 100644 --- a/src/impress/utils/logger.py +++ b/src/impress/utils/logger.py @@ -34,10 +34,25 @@ class LogLevel(Enum): class ImpressLogger: - def __init__(self, name="ImpressManager", use_colors=True, output_stream=None): + _LEVEL_ORDER = [ + LogLevel.DEBUG, + LogLevel.INFO, + LogLevel.WARNING, + LogLevel.ERROR, + LogLevel.CRITICAL, + ] + + def __init__( + self, + name="ImpressManager", + use_colors=True, + output_stream=None, + min_level: LogLevel = LogLevel.DEBUG, + ): self.name = name self.use_colors = use_colors self.output_stream = output_stream or sys.stdout + self.min_level = min_level self.level_colors = { LogLevel.DEBUG: Colors.BRIGHT_BLACK, @@ -91,40 +106,52 @@ def _format_message(self, level, component, message, pipeline_name=None): f"{timestamp} {colored_level} {colored_component}{pipeline_part} {message}" ) - def _write_log(self, message, to_stderr=False): - stream = sys.stderr if to_stderr else self.output_stream - stream.write(message + "\n") - stream.flush() + def _is_enabled(self, level: LogLevel) -> bool: + return self._LEVEL_ORDER.index(level) >= self._LEVEL_ORDER.index(self.min_level) + + def _write_log(self, message): + self.output_stream.write(message + "\n") + self.output_stream.flush() def debug(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.DEBUG): + return formatted = self._format_message( LogLevel.DEBUG, component, message, pipeline_name ) self._write_log(formatted) def info(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.INFO): + return formatted = self._format_message( LogLevel.INFO, component, message, pipeline_name ) self._write_log(formatted) def warning(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.WARNING): + return formatted = self._format_message( LogLevel.WARNING, component, message, pipeline_name ) self._write_log(formatted) def error(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.ERROR): + return formatted = self._format_message( LogLevel.ERROR, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def critical(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.CRITICAL): + return formatted = self._format_message( LogLevel.CRITICAL, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def pipeline_started(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) @@ -136,6 +163,11 @@ def pipeline_completed(self, pipeline_name): message = f"Pipeline completed: {colored_name}" self.info(message, "manager") + def pipeline_failed(self, pipeline_name, exc): + colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) + message = f"Pipeline FAILED: {colored_name} — {exc}" + self.error(message, "manager") + def pipeline_killed(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) message = f"Pipeline killed: {colored_name}" @@ -182,10 +214,11 @@ def activity_summary(self, active_pipelines, active_adaptive, buffered_pipelines self.debug(summary, "manager") def pipeline_log(self, message, level=LogLevel.INFO): + if not self._is_enabled(level): + return pipeline_component = f"PIPELINE-{self.name.upper()}" formatted = self._format_message(level, pipeline_component, message) - stderr_levels = [LogLevel.ERROR, LogLevel.CRITICAL] - self._write_log(formatted, to_stderr=level in stderr_levels) + self._write_log(formatted) def separator(self, title=None): if title: diff --git a/tests/conftest.py b/tests/conftest.py index c35a0f5..cb6ad28 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,10 @@ # tests/conftest.py -import pytest import shutil from pathlib import Path - from unittest.mock import Mock + +import pytest + from impress import ImpressManager diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py new file mode 100644 index 0000000..930d5b1 --- /dev/null +++ b/tests/unit/test_logger.py @@ -0,0 +1,278 @@ +import io +import sys + +from impress.utils.logger import ImpressLogger, LogLevel + + +class TestImpressLoggerInit: + def test_default_init(self): + logger = ImpressLogger() + assert logger.name == "ImpressManager" + assert logger.use_colors is True + assert logger.output_stream is sys.stdout + assert logger.min_level == LogLevel.DEBUG + + def test_custom_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream) + assert logger.output_stream is stream + + def test_custom_min_level(self): + logger = ImpressLogger(min_level=LogLevel.WARNING) + assert logger.min_level == LogLevel.WARNING + + def test_custom_name(self): + logger = ImpressLogger(name="my_pipeline") + assert logger.name == "my_pipeline" + + +class TestLogLevelFiltering: + def test_is_enabled_at_exact_level(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.INFO) is True + + def test_is_enabled_above_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.WARNING) is True + assert logger._is_enabled(LogLevel.ERROR) is True + assert logger._is_enabled(LogLevel.CRITICAL) is True + + def test_is_disabled_below_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.DEBUG) is False + + def test_debug_suppressed_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.debug("should not appear") + assert stream.getvalue() == "" + + def test_info_written_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.info("should appear") + assert "should appear" in stream.getvalue() + + def test_warning_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.ERROR, use_colors=False + ) + logger.warning("should not appear") + assert stream.getvalue() == "" + + def test_all_levels_write_at_debug_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.DEBUG, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" in output + assert "i" in output + assert "w" in output + assert "e" in output + assert "c" in output + + def test_critical_only_at_critical_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.CRITICAL, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" not in output + assert "i" not in output + assert "w" not in output + assert "e" not in output + assert "c" in output + + +class TestOutputStreamRespected: + def test_error_writes_to_output_stream_not_stderr(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.error("an error") + assert "an error" in stream.getvalue() + + def test_critical_writes_to_output_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.critical("critical msg") + assert "critical msg" in stream.getvalue() + + def test_custom_stream_receives_all_output(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("info line") + logger.error("error line") + logger.warning("warn line") + output = stream.getvalue() + assert "info line" in output + assert "error line" in output + assert "warn line" in output + + +class TestColors: + def test_no_ansi_when_colors_off(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("hello") + assert "\033[" not in stream.getvalue() + + def test_ansi_present_when_colors_on(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=True) + logger.info("hello") + assert "\033[" in stream.getvalue() + + def test_colorize_returns_plain_when_off(self): + logger = ImpressLogger(use_colors=False) + result = logger._colorize("text", "\033[31m") + assert result == "text" + + def test_colorize_wraps_when_on(self): + logger = ImpressLogger(use_colors=True) + result = logger._colorize("text", "\033[31m") + assert result.startswith("\033[31m") + assert "text" in result + + +class TestHighLevelMethods: + def _make_logger(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + return logger, stream + + def test_pipeline_started(self): + logger, stream = self._make_logger() + logger.pipeline_started("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_completed(self): + logger, stream = self._make_logger() + logger.pipeline_completed("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_failed(self): + logger, stream = self._make_logger() + logger.pipeline_failed("bad_pipe", ValueError("oops")) + out = stream.getvalue() + assert "bad_pipe" in out + assert "oops" in out + + def test_pipeline_killed(self): + logger, stream = self._make_logger() + logger.pipeline_killed("dead_pipe") + assert "dead_pipe" in stream.getvalue() + + def test_adaptive_started(self): + logger, stream = self._make_logger() + logger.adaptive_started("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_completed(self): + logger, stream = self._make_logger() + logger.adaptive_completed("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_failed(self): + logger, stream = self._make_logger() + logger.adaptive_failed("p1", "bad fn") + out = stream.getvalue() + assert "p1" in out + assert "bad fn" in out + + def test_child_pipeline_submitted(self): + logger, stream = self._make_logger() + logger.child_pipeline_submitted("child", "parent") + out = stream.getvalue() + assert "child" in out + assert "parent" in out + + def test_manager_starting(self): + logger, stream = self._make_logger() + logger.manager_starting(5) + assert "5" in stream.getvalue() + + def test_manager_exiting(self): + logger, stream = self._make_logger() + logger.manager_exiting() + assert stream.getvalue() != "" + + def test_activity_summary_suppressed_at_info(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.INFO + ) + logger.activity_summary(3, 1, 2) + assert stream.getvalue() == "" # activity_summary is DEBUG level + + def test_activity_summary_written_at_debug(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.DEBUG + ) + logger.activity_summary(3, 1, 2) + out = stream.getvalue() + assert "3" in out + + +class TestPipelineLog: + def test_pipeline_log_default_info(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("step done") + assert "step done" in stream.getvalue() + + def test_pipeline_log_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + "pipe1", output_stream=stream, use_colors=False, min_level=LogLevel.WARNING + ) + logger.pipeline_log("step done", level=LogLevel.INFO) + assert stream.getvalue() == "" + + def test_pipeline_log_debug_level(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("debug step", level=LogLevel.DEBUG) + assert "debug step" in stream.getvalue() + + def test_pipeline_log_includes_pipeline_name_component(self): + stream = io.StringIO() + logger = ImpressLogger("mypipe", output_stream=stream, use_colors=False) + logger.pipeline_log("event") + assert "PIPELINE-MYPIPE" in stream.getvalue() + + +class TestSeparator: + def test_separator_no_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert "=" in stream.getvalue() + + def test_separator_with_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator("HELLO WORLD") + assert "HELLO WORLD" in stream.getvalue() + + def test_separator_ends_with_newline(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert stream.getvalue().endswith("\n") diff --git a/tests/unit/test_manager_core.py b/tests/unit/test_manager_core.py index d35495c..883e912 100644 --- a/tests/unit/test_manager_core.py +++ b/tests/unit/test_manager_core.py @@ -3,8 +3,7 @@ import pytest # Import the classes we're testing -from impress import ImpressBasePipeline, PipelineSetup -from impress import ImpressManager +from impress import ImpressBasePipeline, ImpressManager, PipelineSetup class MockPipeline(ImpressBasePipeline): diff --git a/tests/unit/test_manager_life_cycle.py b/tests/unit/test_manager_life_cycle.py index 9783467..bd110b3 100644 --- a/tests/unit/test_manager_life_cycle.py +++ b/tests/unit/test_manager_life_cycle.py @@ -12,9 +12,12 @@ class MockWorkflowEngine: """Mock workflow engine""" @classmethod - async def create(cls, backend=None): + async def create(cls, backend=None, **kwargs): return cls() + async def shutdown(self): + pass + class TestManagerLifecycle: @pytest.mark.asyncio @@ -113,3 +116,18 @@ async def run(self): # Should take at least 0.15 seconds due to slow adaptive function assert end_time - start_time >= 0.15 + + @pytest.mark.asyncio + @patch("impress.impress_manager.WorkflowEngine", MockWorkflowEngine) + async def test_start_exception_still_shuts_down_engine(self, impress_manager): + """WorkflowEngine.shutdown() is called even when a pipeline raises mid-run""" + + class ExplodingPipeline(MockPipeline): + async def run(self): + raise RuntimeError("pipeline exploded") + + await impress_manager.start( + [{"name": "boom", "type": ExplodingPipeline, "config": {}, "kwargs": {}}] + ) + # If shutdown() raises AttributeError, the try/finally fix is broken + assert len(impress_manager.pipeline_tasks) == 0 diff --git a/tests/unit/test_manager_pipeline_submission.py b/tests/unit/test_manager_pipeline_submission.py index e390b93..92bbdfa 100644 --- a/tests/unit/test_manager_pipeline_submission.py +++ b/tests/unit/test_manager_pipeline_submission.py @@ -1,5 +1,7 @@ from unittest.mock import Mock, patch +import pytest + # Import the classes we're testing from impress import PipelineSetup @@ -64,3 +66,10 @@ def test_submit_multiple_pipelines(self, mock_create_task, impress_manager): assert len(impress_manager.pipeline_tasks) == 2 assert mock_create_task.call_count == 2 + + def test_submit_before_start_raises(self, impress_manager): + """submit_new_pipelines raises RuntimeError when called before start()""" + with pytest.raises(RuntimeError, match="start\\(\\) must be called"): + impress_manager.submit_new_pipelines( + [{"name": "p", "type": MockPipeline, "config": {}, "kwargs": {}}] + ) diff --git a/tests/unit/test_pipeline_base.py b/tests/unit/test_pipeline_base.py new file mode 100644 index 0000000..928ea2d --- /dev/null +++ b/tests/unit/test_pipeline_base.py @@ -0,0 +1,157 @@ +import asyncio + +import pytest + +from impress.pipelines.impress_pipeline import ImpressBasePipeline + + +class MinimalPipeline(ImpressBasePipeline): + """Minimal concrete subclass for testing the base class.""" + + async def run(self): + pass + + def register_pipeline_tasks(self): + pass + + +class TestImpressBasePipelineInit: + def test_name_set(self): + p = MinimalPipeline(name="p1") + assert p.name == "p1" + + def test_flow_defaults_to_none(self): + p = MinimalPipeline(name="p1") + assert p.flow is None + + def test_flow_passed_through(self): + sentinel = object() + p = MinimalPipeline(name="p1", flow=sentinel) + assert p.flow is sentinel + + def test_state_is_empty_dict(self): + p = MinimalPipeline(name="p1") + assert p.state == {} + + def test_kill_parent_false(self): + p = MinimalPipeline(name="p1") + assert p.kill_parent is False + + def test_invoke_adaptive_step_false(self): + p = MinimalPipeline(name="p1") + assert p.invoke_adaptive_step is False + + def test_adaptive_barrier_is_event(self): + p = MinimalPipeline(name="p1") + assert isinstance(p._adaptive_barrier, asyncio.Event) + + def test_incoming_child_pipeline_request_empty(self): + p = MinimalPipeline(name="p1") + assert not p.incoming_child_pipeline_request + + def test_kwargs_stored_in_config(self): + p = MinimalPipeline(name="p1", foo="bar", baz=42) + assert p.config["foo"] == "bar" + assert p.config["baz"] == 42 + + +class TestChildPipelineRequest: + def test_submit_sets_request(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + assert p.incoming_child_pipeline_request == config + + def test_get_returns_config_then_clears(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + + result1 = p.get_child_pipeline_request() + assert result1 == config + + result2 = p.get_child_pipeline_request() + assert result2 is None + + def test_get_returns_none_when_nothing_pending(self): + p = MinimalPipeline(name="p1") + assert p.get_child_pipeline_request() is None + + def test_get_clears_after_retrieval(self): + p = MinimalPipeline(name="p1") + p.submit_child_pipeline_request({"name": "c"}) + p.get_child_pipeline_request() + assert not p.incoming_child_pipeline_request + + +class TestAdaptiveStep: + @pytest.mark.asyncio + async def test_run_adaptive_step_sets_flag(self): + p = MinimalPipeline(name="p1") + # wait=False: flag is set without blocking on the barrier + await p.run_adaptive_step(wait=False) + assert p.invoke_adaptive_step is True + + @pytest.mark.asyncio + async def test_run_adaptive_step_no_wait_does_not_hang(self): + p = MinimalPipeline(name="p1") + # _adaptive_barrier is clear — with wait=False this must return immediately + await asyncio.wait_for(p.run_adaptive_step(wait=False), timeout=1.0) + assert p.invoke_adaptive_step is True + + def test_set_adaptive_flag_true_clears_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(True) + assert p.invoke_adaptive_step is True + assert not p._adaptive_barrier.is_set() + + def test_set_adaptive_flag_false_does_not_touch_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(False) + assert p.invoke_adaptive_step is False + assert p._adaptive_barrier.is_set() # barrier unchanged + + +class TestOptionalMethods: + @pytest.mark.asyncio + async def test_finalize_is_noop(self): + p = MinimalPipeline(name="p1") + result = await p.finalize() + assert result is None + + @pytest.mark.asyncio + async def test_get_scores_map_returns_empty_dict(self): + p = MinimalPipeline(name="p1") + scores = await p.get_scores_map() + assert scores == {} + + def test_get_current_config_has_name_and_type(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert "name" in cfg + assert "type" in cfg + + def test_get_current_config_type_is_class(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert cfg["type"] is MinimalPipeline + + +class TestAbstractMethods: + def test_cannot_instantiate_without_run(self): + class NoRun(ImpressBasePipeline): + def register_pipeline_tasks(self): + pass + + with pytest.raises(TypeError): + NoRun(name="x") + + def test_cannot_instantiate_without_register(self): + class NoRegister(ImpressBasePipeline): + async def run(self): + pass + + with pytest.raises(TypeError): + NoRegister(name="x") diff --git a/tests/unit/test_pipeline_setup.py b/tests/unit/test_pipeline_setup.py new file mode 100644 index 0000000..c56d398 --- /dev/null +++ b/tests/unit/test_pipeline_setup.py @@ -0,0 +1,141 @@ +import pytest +from pydantic import ValidationError + +from impress import PipelineSetup + +from .test_manager_core import MockPipeline + + +class TestPipelineSetupConstruction: + def test_all_fields(self): + async def fn(p): + pass + + setup = PipelineSetup( + name="my_pipe", + type=MockPipeline, + config={"a": 1}, + kwargs={"b": 2}, + adaptive_fn=fn, + ) + assert setup.name == "my_pipe" + assert setup.type is MockPipeline + assert setup.config == {"a": 1} + assert setup.kwargs == {"b": 2} + assert setup.adaptive_fn is fn + + def test_defaults(self): + setup = PipelineSetup(name="p", type=MockPipeline) + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + def test_validate_type_rejects_non_subclass(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type=str) + + def test_validate_type_rejects_non_type(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type="not_a_class") + + def test_validate_type_accepts_subclass(self): + class Sub(MockPipeline): + pass + + setup = PipelineSetup(name="p", type=Sub) + assert setup.type is Sub + + +class TestFromDict: + def test_known_fields_separated(self): + async def fn(p): + pass + + data = { + "name": "p1", + "type": MockPipeline, + "config": {"x": 1}, + "adaptive_fn": fn, + } + setup = PipelineSetup.from_dict(data) + assert setup.name == "p1" + assert setup.type is MockPipeline + assert setup.config == {"x": 1} + assert setup.adaptive_fn is fn + assert setup.kwargs == {} + + def test_extra_keys_go_to_kwargs(self): + data = { + "name": "p1", + "type": MockPipeline, + "foo": "bar", + "baz": 42, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"foo": "bar", "baz": 42} + + def test_kwargs_key_in_dict_lands_in_kwargs(self): + # "kwargs" is not a known field, so it ends up nested inside kwargs + data = { + "name": "p1", + "type": MockPipeline, + "kwargs": {"inner": "value"}, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"kwargs": {"inner": "value"}} + + def test_minimal_dict(self): + setup = PipelineSetup.from_dict({"name": "p", "type": MockPipeline}) + assert setup.name == "p" + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + +class TestToDict: + def test_basic_structure(self): + setup = PipelineSetup( + name="p1", + type=MockPipeline, + config={"c": 1}, + ) + d = setup.to_dict() + assert d["name"] == "p1" + assert d["type"] is MockPipeline + assert d["config"] == {"c": 1} + + def test_adaptive_fn_omitted_when_none(self): + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=None) + d = setup.to_dict() + assert "adaptive_fn" not in d + + def test_adaptive_fn_included_when_set(self): + async def fn(p): + pass + + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=fn) + d = setup.to_dict() + assert d["adaptive_fn"] is fn + + def test_kwargs_spread_into_result(self): + setup = PipelineSetup( + name="p", + type=MockPipeline, + kwargs={"foo": "bar", "num": 7}, + ) + d = setup.to_dict() + assert d["foo"] == "bar" + assert d["num"] == 7 + + def test_roundtrip_from_dict(self): + original = { + "name": "rt", + "type": MockPipeline, + "config": {"k": "v"}, + "extra_param": 99, + } + setup = PipelineSetup.from_dict(original) + d = setup.to_dict() + assert d["name"] == "rt" + assert d["type"] is MockPipeline + assert d["extra_param"] == 99