From 4e360a5702841b78d64c052511737c064bb382d1 Mon Sep 17 00:00:00 2001 From: Moto Hira Date: Tue, 2 Jun 2026 10:46:55 -0700 Subject: [PATCH] Move per-module tests under each module (#1500) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/facebookresearch/spdl/pull/1500 Reorganize SPDL tests so each module owns its own tests. Tests previously lived under `fbcode/spdl/tests//` and depended on a shared root for fixtures and Buck macros. After this change, each module's tests live in `/tests/`, which makes `buck test fbcode//spdl//...` run every test for that module without enumerating sibling test paths. New layout: - `spdl/io/tests/` (was `spdl/tests/io/`) — also hosts the shared `fixture.py` + `fb/__init__.py` (FFmpeg helpers), since only `io` and `cuda` tests use them and `cuda` tests already depend only on `spdl.io`. - `spdl/cuda/tests/` (was `spdl/tests/cuda/`) — `PACKAGE` moved with it; depends on `//spdl/io/tests:fixture` and imports `from spdl.io.tests.fixture`. - `spdl/pipeline/tests/` (was `spdl/tests/pipeline/`, including `fb/`). - `spdl/dataloader/tests/` (was `spdl/tests/dataloader/`). - `spdl/autoresearch/tests/` (was `spdl/tests/autoresearch/`, including `fb/`). - `spdl/def.bzl` — new file holding the `spdl_tests_with_ffmpeg_variants` and `spdl_tests_with_python_variants` macros that were previously in `spdl/tests/def.bzl`. - `fbcode/spdl/tests/` deleted entirely. Other updates: - `from ..fixture` (relative) is now `from .fixture` for `io` tests and `from spdl.io.tests.fixture` (absolute) for `cuda` tests. - `spdl.io.tests/fb/__init__.py` resolves the FFmpeg resource via `importlib.resources.path("spdl.io.tests", "ffmpeg")`. - Visibility lists in module BUCKs (`spdl/io/...`, `spdl/io/lib/...`, `spdl/io/utils/...`, `spdl/pipeline/...`, `spdl/pipeline/_iter_utils/...`, `spdl/pipeline/fb/...`, `spdl/pipeline/fb/lib/...`, `spdl/autoresearch/...`, `spdl/autoresearch/_app/...`, `spdl/autoresearch/_common/fb/...`, `spdl/autoresearch/core/...`, `spdl/autoresearch/pipeline_optimization/...`) updated to point to the new per-module test paths. - `fbcode/spdl/PACKAGE` comments updated to reference the new `tests/` paths. - `fbcode/spdl/.llms/skills/code-authoring/SKILL.md` updated to describe the new layout convention. - GitHub workflows (`_build_linux.yml`, `_build_linux_cuda.yml`, `_build_macos.yml`, `_build_windows.yml`) updated so each pytest invocation lists the new per-module paths under the OSS `src/spdl/` prefix; `docs.yml` `paths-ignore` updated to `src/spdl/**/tests/**`. - Hydra `_target_` string in `pipeline_hydra_test.py` updated from `spdl.tests.pipeline.fb...` to `spdl.pipeline.tests.fb...`. Differential Revision: D107140706 --- .github/workflows/_build_linux.yml | 2 +- .github/workflows/_build_linux_cuda.yml | 3 +- .github/workflows/_build_macos.yml | 2 +- .github/workflows/_build_windows.yml | 2 +- .github/workflows/docs.yml | 2 + src/spdl/autoresearch/tests/app_test.py | 283 ++ src/spdl/autoresearch/tests/factory_test.py | 205 ++ .../autoresearch/tests/orchestrator_test.py | 128 + .../autoresearch/tests/persistence_test.py | 168 + src/spdl/autoresearch/tests/platform_test.py | 228 ++ src/spdl/autoresearch/tests/prompts_test.py | 129 + src/spdl/autoresearch/tests/state_test.py | 27 + .../autoresearch/tests/supervisor_test.py | 34 + src/spdl/autoresearch/tests/types_test.py | 80 + .../autoresearch/tests/visualization_test.py | 44 + src/spdl/autoresearch/tests/workflow_test.py | 1049 +++++++ .../dataloader/tests/cache_dataloader_test.py | 108 + src/spdl/dataloader/tests/dataloader_test.py | 167 + src/spdl/dataloader/tests/iterator_test.py | 370 +++ src/spdl/dataloader/tests/sampler_test.py | 343 ++ src/spdl/dataloader/tests/source_test.py | 89 + .../dataloader/tests/source_utils_test.py | 92 + .../spdl/io/tests/core}/array_test.py | 0 .../spdl/io/tests/core}/async_test.py | 3 +- .../io/tests/core}/audio_decoding_test.py | 3 +- .../io/tests/core}/audio_encoding_test.py | 3 +- .../core}/buffer_conversion_refcount_test.py | 3 +- .../spdl/io/tests/core}/configs_test.py | 3 +- .../spdl/io/tests/core}/demuxer_test.py | 3 +- .../spdl/io/tests/core}/encoding_test.py | 3 +- .../spdl/io/tests/core}/filter_test.py | 3 +- .../spdl/io/tests/core}/frames_clone_test.py | 3 +- .../spdl/io/tests/core}/frames_test.py | 3 +- .../io/tests/core}/image_decoding_test.py | 9 +- .../spdl/io/tests/core}/memoryview_test.py | 0 .../spdl/io/tests/core}/packets_test.py | 3 +- .../io/tests/core}/reference_frames_test.py | 0 .../spdl/io/tests/core}/serialization_test.py | 3 +- .../io/tests/core}/streaming_decoding_test.py | 3 +- .../core}/subprocess_serialization_test.py | 3 +- .../io => src/spdl/io/tests/core}/tar_test.py | 0 .../spdl/io/tests/core}/transfer_test.py | 0 .../spdl/io/tests/core}/utils_test.py | 0 .../io/tests/core}/video_encoding_test.py | 3 +- .../io/tests/core}/video_frame_slice_test.py | 3 +- .../io => src/spdl/io/tests/core}/wav_test.py | 0 .../core}/zero_copy_bytes_passing_test.py | 3 +- .../io => src/spdl/io/tests/core}/zip_test.py | 0 .../io/tests}/cuda/buffer_transfer_test.py | 3 +- .../spdl/io/tests/cuda/cuda_transfer_test.py | 0 .../tests}/cuda/nvdec_video_decoding_test.py | 3 +- .../spdl/io/tests}/cuda/nvjpeg_decode_test.py | 3 +- .../spdl/io/tests}/cuda/pin_memory_test.py | 3 +- .../cuda/streaming_load_video_nvdec_test.py | 3 +- .../io/tests}/cuda/subprocess_nvdec_test.py | 3 +- {tests => src/spdl/io/tests}/fixture.py | 0 src/spdl/pipeline/tests/aggregate_test.py | 326 ++ .../pipeline/tests/background_task_test.py | 278 ++ .../pipeline/tests/build_pipeline_test.py | 52 + src/spdl/pipeline/tests/compact_log_test.py | 279 ++ src/spdl/pipeline/tests/config_test.py | 210 ++ .../tests/continuous_pipeline_test.py | 623 ++++ src/spdl/pipeline/tests/defs_repr_test.py | 355 +++ src/spdl/pipeline/tests/failure_rate_test.py | 804 +++++ src/spdl/pipeline/tests/merge_config_test.py | 475 +++ src/spdl/pipeline/tests/path_variants_test.py | 704 +++++ .../pipeline/tests/percentile_stats_test.py | 453 +++ src/spdl/pipeline/tests/pgrp_stats_test.py | 496 +++ .../pipeline/tests/pipeline_builder_test.py | 2763 ++++++++++++++++ .../pipeline/tests/pipeline_cleanup_test.py | 203 ++ src/spdl/pipeline/tests/pipeline_def_test.py | 168 + .../tests/pipeline_failure_exceptstar_test.py | 52 + src/spdl/pipeline/tests/pipeline_node_test.py | 299 ++ .../pipeline/tests/pipeline_profiling_test.py | 321 ++ .../pipeline/tests/priority_executor_test.py | 937 ++++++ .../pipeline/tests/source_locator_test.py | 200 ++ .../pipeline/tests/subinterpreter_test.py | 278 ++ .../tests/subprocess_break_reiterate_test.py | 109 + src/spdl/pipeline/tests/subprocess_test.py | 506 +++ tests/__init__.py | 0 tests/autoresearch/app_test.py | 284 +- tests/autoresearch/factory_test.py | 206 +- tests/autoresearch/orchestrator_test.py | 129 +- tests/autoresearch/persistence_test.py | 169 +- tests/autoresearch/platform_test.py | 229 +- tests/autoresearch/prompts_test.py | 130 +- tests/autoresearch/state_test.py | 28 +- tests/autoresearch/supervisor_test.py | 35 +- tests/autoresearch/types_test.py | 81 +- tests/autoresearch/visualization_test.py | 45 +- tests/autoresearch/workflow_test.py | 1050 +------ tests/conftest.py | 46 + tests/cuda/__init__.py | 0 tests/dataloader/cache_dataloader_test.py | 109 +- tests/dataloader/dataloader_test.py | 168 +- tests/dataloader/iterator_test.py | 371 +-- tests/dataloader/sampler_test.py | 344 +- tests/dataloader/source_test.py | 90 +- tests/dataloader/source_utils_test.py | 93 +- tests/io/__init__.py | 0 tests/io/core | 1 + tests/io/cuda | 1 + tests/io/fixture.py | 1 + tests/pipeline/__init__.py | 0 tests/pipeline/aggregate_test.py | 327 +- tests/pipeline/background_task_test.py | 279 +- tests/pipeline/build_pipeline_test.py | 53 +- tests/pipeline/compact_log_test.py | 280 +- tests/pipeline/config_test.py | 211 +- tests/pipeline/continuous_pipeline_test.py | 624 +--- tests/pipeline/defs_repr_test.py | 356 +-- tests/pipeline/failure_rate_test.py | 805 +---- tests/pipeline/merge_config_test.py | 476 +-- tests/pipeline/path_variants_test.py | 705 +---- tests/pipeline/percentile_stats_test.py | 454 +-- tests/pipeline/pgrp_stats_test.py | 497 +-- tests/pipeline/pipeline_builder_test.py | 2764 +---------------- tests/pipeline/pipeline_cleanup_test.py | 204 +- tests/pipeline/pipeline_def_test.py | 169 +- .../pipeline_failure_exceptstar_test.py | 53 +- tests/pipeline/pipeline_node_test.py | 300 +- tests/pipeline/pipeline_profiling_test.py | 322 +- tests/pipeline/priority_executor_test.py | 938 +----- tests/pipeline/source_locator_test.py | 201 +- tests/pipeline/subinterpreter_test.py | 279 +- .../subprocess_break_reiterate_test.py | 110 +- tests/pipeline/subprocess_test.py | 507 +-- 127 files changed, 14560 insertions(+), 14488 deletions(-) create mode 100644 src/spdl/autoresearch/tests/app_test.py create mode 100644 src/spdl/autoresearch/tests/factory_test.py create mode 100644 src/spdl/autoresearch/tests/orchestrator_test.py create mode 100644 src/spdl/autoresearch/tests/persistence_test.py create mode 100644 src/spdl/autoresearch/tests/platform_test.py create mode 100644 src/spdl/autoresearch/tests/prompts_test.py create mode 100644 src/spdl/autoresearch/tests/state_test.py create mode 100644 src/spdl/autoresearch/tests/supervisor_test.py create mode 100644 src/spdl/autoresearch/tests/types_test.py create mode 100644 src/spdl/autoresearch/tests/visualization_test.py create mode 100644 src/spdl/autoresearch/tests/workflow_test.py create mode 100644 src/spdl/dataloader/tests/cache_dataloader_test.py create mode 100644 src/spdl/dataloader/tests/dataloader_test.py create mode 100644 src/spdl/dataloader/tests/iterator_test.py create mode 100644 src/spdl/dataloader/tests/sampler_test.py create mode 100644 src/spdl/dataloader/tests/source_test.py create mode 100644 src/spdl/dataloader/tests/source_utils_test.py rename {tests/io => src/spdl/io/tests/core}/array_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/async_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/audio_decoding_test.py (97%) rename {tests/io => src/spdl/io/tests/core}/audio_encoding_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/buffer_conversion_refcount_test.py (96%) rename {tests/io => src/spdl/io/tests/core}/configs_test.py (97%) rename {tests/io => src/spdl/io/tests/core}/demuxer_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/encoding_test.py (98%) rename {tests/io => src/spdl/io/tests/core}/filter_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/frames_clone_test.py (98%) rename {tests/io => src/spdl/io/tests/core}/frames_test.py (96%) rename {tests/io => src/spdl/io/tests/core}/image_decoding_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/memoryview_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/packets_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/reference_frames_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/serialization_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/streaming_decoding_test.py (99%) rename {tests/io => src/spdl/io/tests/core}/subprocess_serialization_test.py (98%) rename {tests/io => src/spdl/io/tests/core}/tar_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/transfer_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/utils_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/video_encoding_test.py (98%) rename {tests/io => src/spdl/io/tests/core}/video_frame_slice_test.py (98%) rename {tests/io => src/spdl/io/tests/core}/wav_test.py (100%) rename {tests/io => src/spdl/io/tests/core}/zero_copy_bytes_passing_test.py (97%) rename {tests/io => src/spdl/io/tests/core}/zip_test.py (100%) rename {tests => src/spdl/io/tests}/cuda/buffer_transfer_test.py (99%) rename tests/cuda/transfer_test.py => src/spdl/io/tests/cuda/cuda_transfer_test.py (100%) rename {tests => src/spdl/io/tests}/cuda/nvdec_video_decoding_test.py (99%) rename {tests => src/spdl/io/tests}/cuda/nvjpeg_decode_test.py (98%) rename {tests => src/spdl/io/tests}/cuda/pin_memory_test.py (98%) rename {tests => src/spdl/io/tests}/cuda/streaming_load_video_nvdec_test.py (99%) rename {tests => src/spdl/io/tests}/cuda/subprocess_nvdec_test.py (98%) rename {tests => src/spdl/io/tests}/fixture.py (100%) create mode 100644 src/spdl/pipeline/tests/aggregate_test.py create mode 100644 src/spdl/pipeline/tests/background_task_test.py create mode 100644 src/spdl/pipeline/tests/build_pipeline_test.py create mode 100644 src/spdl/pipeline/tests/compact_log_test.py create mode 100644 src/spdl/pipeline/tests/config_test.py create mode 100644 src/spdl/pipeline/tests/continuous_pipeline_test.py create mode 100644 src/spdl/pipeline/tests/defs_repr_test.py create mode 100644 src/spdl/pipeline/tests/failure_rate_test.py create mode 100644 src/spdl/pipeline/tests/merge_config_test.py create mode 100644 src/spdl/pipeline/tests/path_variants_test.py create mode 100644 src/spdl/pipeline/tests/percentile_stats_test.py create mode 100644 src/spdl/pipeline/tests/pgrp_stats_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_builder_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_cleanup_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_def_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_failure_exceptstar_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_node_test.py create mode 100644 src/spdl/pipeline/tests/pipeline_profiling_test.py create mode 100644 src/spdl/pipeline/tests/priority_executor_test.py create mode 100644 src/spdl/pipeline/tests/source_locator_test.py create mode 100644 src/spdl/pipeline/tests/subinterpreter_test.py create mode 100644 src/spdl/pipeline/tests/subprocess_break_reiterate_test.py create mode 100644 src/spdl/pipeline/tests/subprocess_test.py delete mode 100644 tests/__init__.py mode change 100644 => 120000 tests/autoresearch/app_test.py mode change 100644 => 120000 tests/autoresearch/factory_test.py mode change 100644 => 120000 tests/autoresearch/orchestrator_test.py mode change 100644 => 120000 tests/autoresearch/persistence_test.py mode change 100644 => 120000 tests/autoresearch/platform_test.py mode change 100644 => 120000 tests/autoresearch/prompts_test.py mode change 100644 => 120000 tests/autoresearch/state_test.py mode change 100644 => 120000 tests/autoresearch/supervisor_test.py mode change 100644 => 120000 tests/autoresearch/types_test.py mode change 100644 => 120000 tests/autoresearch/visualization_test.py mode change 100644 => 120000 tests/autoresearch/workflow_test.py create mode 100644 tests/conftest.py delete mode 100644 tests/cuda/__init__.py mode change 100644 => 120000 tests/dataloader/cache_dataloader_test.py mode change 100644 => 120000 tests/dataloader/dataloader_test.py mode change 100644 => 120000 tests/dataloader/iterator_test.py mode change 100644 => 120000 tests/dataloader/sampler_test.py mode change 100644 => 120000 tests/dataloader/source_test.py mode change 100644 => 120000 tests/dataloader/source_utils_test.py delete mode 100644 tests/io/__init__.py create mode 120000 tests/io/core create mode 120000 tests/io/cuda create mode 120000 tests/io/fixture.py delete mode 100644 tests/pipeline/__init__.py mode change 100644 => 120000 tests/pipeline/aggregate_test.py mode change 100644 => 120000 tests/pipeline/background_task_test.py mode change 100644 => 120000 tests/pipeline/build_pipeline_test.py mode change 100644 => 120000 tests/pipeline/compact_log_test.py mode change 100644 => 120000 tests/pipeline/config_test.py mode change 100644 => 120000 tests/pipeline/continuous_pipeline_test.py mode change 100644 => 120000 tests/pipeline/defs_repr_test.py mode change 100644 => 120000 tests/pipeline/failure_rate_test.py mode change 100644 => 120000 tests/pipeline/merge_config_test.py mode change 100644 => 120000 tests/pipeline/path_variants_test.py mode change 100644 => 120000 tests/pipeline/percentile_stats_test.py mode change 100644 => 120000 tests/pipeline/pgrp_stats_test.py mode change 100644 => 120000 tests/pipeline/pipeline_builder_test.py mode change 100644 => 120000 tests/pipeline/pipeline_cleanup_test.py mode change 100644 => 120000 tests/pipeline/pipeline_def_test.py mode change 100644 => 120000 tests/pipeline/pipeline_failure_exceptstar_test.py mode change 100644 => 120000 tests/pipeline/pipeline_node_test.py mode change 100644 => 120000 tests/pipeline/pipeline_profiling_test.py mode change 100644 => 120000 tests/pipeline/priority_executor_test.py mode change 100644 => 120000 tests/pipeline/source_locator_test.py mode change 100644 => 120000 tests/pipeline/subinterpreter_test.py mode change 100644 => 120000 tests/pipeline/subprocess_break_reiterate_test.py mode change 100644 => 120000 tests/pipeline/subprocess_test.py diff --git a/.github/workflows/_build_linux.yml b/.github/workflows/_build_linux.yml index ec9df8633..6dd8c2705 100644 --- a/.github/workflows/_build_linux.yml +++ b/.github/workflows/_build_linux.yml @@ -177,7 +177,7 @@ jobs: python -c 'import spdl.io.utils;assert not spdl.io.utils.built_with_cuda()' pytest -v -n ${{ inputs.test-concurrency }} \ - tests/io/ \ + tests/io/core/ \ tests/dataloader/ \ tests/pipeline/ \ tests/autoresearch/ diff --git a/.github/workflows/_build_linux_cuda.yml b/.github/workflows/_build_linux_cuda.yml index 7e94f65e1..f23faa4ec 100644 --- a/.github/workflows/_build_linux_cuda.yml +++ b/.github/workflows/_build_linux_cuda.yml @@ -221,7 +221,6 @@ jobs: python -c 'import spdl.io.utils;assert spdl.io.utils.built_with_nvcodec()' fi pytest -v -n ${{ inputs.test-concurrency }} \ - tests/cuda/ \ tests/io/ test-cpu: @@ -284,7 +283,7 @@ jobs: set -ex pytest -v -n ${{ inputs.test-concurrency }} \ - tests/io \ + tests/io/core \ tests/dataloader \ tests/pipeline \ tests/autoresearch diff --git a/.github/workflows/_build_macos.yml b/.github/workflows/_build_macos.yml index d8705b365..67297b01b 100644 --- a/.github/workflows/_build_macos.yml +++ b/.github/workflows/_build_macos.yml @@ -128,7 +128,7 @@ jobs: run: | set -ex pytest -v -n ${{ inputs.test-concurrency }} \ - tests/io \ + tests/io/core \ tests/dataloader \ tests/pipeline \ tests/autoresearch diff --git a/.github/workflows/_build_windows.yml b/.github/workflows/_build_windows.yml index 0dae88b68..0030320de 100644 --- a/.github/workflows/_build_windows.yml +++ b/.github/workflows/_build_windows.yml @@ -134,7 +134,7 @@ jobs: run: | set -ex pytest -v -n ${{ inputs.test-concurrency }} \ - tests/io \ + tests/io/core \ tests/dataloader \ tests/pipeline \ tests/autoresearch diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 6ca411d3a..40540ced9 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -5,6 +5,7 @@ on: pull_request: paths-ignore: - "tests/**" + - "src/spdl/**/tests/**" - "third_party/**" - "*.md" branches: @@ -12,6 +13,7 @@ on: push: paths-ignore: - "tests/**" + - "src/spdl/**/tests/**" - "third_party/**" - "*.md" branches: diff --git a/src/spdl/autoresearch/tests/app_test.py b/src/spdl/autoresearch/tests/app_test.py new file mode 100644 index 000000000..08b8c2c37 --- /dev/null +++ b/src/spdl/autoresearch/tests/app_test.py @@ -0,0 +1,283 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import importlib +import sys +import tempfile +import unittest +from collections.abc import Callable +from pathlib import Path +from typing import get_origin + +from spdl.autoresearch._app._engine import _parse_engine_args +from spdl.autoresearch._app._spec import ( + _read_workflow_factory, + _record_workflow_factory, + _resolve_workflow, +) +from spdl.autoresearch._app._supervisor import ( + _build_engine_command, + _parse_supervisor_args, +) + +__all__: list[str] = [] + + +def _identity_factory(argv: list[str], workdir: Path | None) -> object: + """A trivial factory used as a resolution target by the tests below.""" + return (argv, workdir) + + +class _ResolveWorkflowTest(unittest.TestCase): + def test_module_factory_form(self) -> None: + """_resolve_workflow imports module.path:factory_name and returns the callable.""" + factory = _resolve_workflow(f"{__name__}:_identity_factory") + self.assertIs(factory, _identity_factory) + + def test_empty_specifier_raises(self) -> None: + """An empty string is rejected with ValueError, not silently importing.""" + with self.assertRaises(ValueError): + _resolve_workflow("") + + def test_malformed_specifier_raises(self) -> None: + """A specifier with a colon but missing one half is rejected.""" + for bad in (":factory", "module.path:", ":"): + with self.subTest(bad=bad): + with self.assertRaises(ValueError): + _resolve_workflow(bad) + + def test_unknown_module_raises_import_error(self) -> None: + """A non-existent module surfaces ModuleNotFoundError to the caller.""" + with self.assertRaises(ModuleNotFoundError): + _resolve_workflow("definitely.not.a.real.module:create") + + def test_missing_attribute_raises(self) -> None: + """An existing module with a missing attribute raises AttributeError.""" + with self.assertRaises(AttributeError): + _resolve_workflow(f"{__name__}:does_not_exist") + + def test_short_name_lookup_misses_cleanly(self) -> None: + """Short-name lookup raises LookupError when no entry point matches.""" + with self.assertRaises(LookupError): + _resolve_workflow("not_registered_workflow_xyz") + + +class _WorkflowFactoryRecordTest(unittest.TestCase): + def test_round_trip(self) -> None: + """_record_workflow_factory followed by _read_workflow_factory returns the spec.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + _record_workflow_factory(workdir, "pkg.mod:factory") + self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") + + def test_read_returns_none_when_missing(self) -> None: + """_read_workflow_factory on a fresh workdir returns None instead of raising.""" + with tempfile.TemporaryDirectory() as tmp: + self.assertIsNone(_read_workflow_factory(Path(tmp))) + + def test_read_rejects_malformed_record(self) -> None: + """Reading a malformed record file raises ValueError.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + (workdir / "workflow_factory.json").write_text("[]\n") + with self.assertRaises(ValueError): + _read_workflow_factory(workdir) + + def test_record_creates_workdir(self) -> None: + """_record_workflow_factory creates the workdir if it does not yet exist.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) / "nested" + _record_workflow_factory(workdir, "pkg.mod:factory") + self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") + + +class _ArgvSplitTest(unittest.TestCase): + def test_supervisor_splits_at_double_dash(self) -> None: + """Tokens after '--' are forwarded as the workflow tail, not parsed by the framework.""" + ns, tail = _parse_supervisor_args( + [ + "/tmp/workdir", + "--workflow", + "pkg.mod:factory", + "--max-concurrency", + "5", + "--", + "--pipeline-script", + "x.py", + ] + ) + self.assertEqual(ns.workdir, "/tmp/workdir") + self.assertEqual(ns.workflow, "pkg.mod:factory") + self.assertEqual(ns.max_concurrency, 5) + self.assertEqual(tail, ["--pipeline-script", "x.py"]) + + def test_supervisor_workdir_optional(self) -> None: + """The supervisor accepts no workdir during initial config gathering.""" + ns, tail = _parse_supervisor_args(["--workflow", "pkg.mod:factory"]) + self.assertIsNone(ns.workdir) + self.assertEqual(tail, []) + + def test_engine_requires_workflow_and_workdir(self) -> None: + """The engine refuses to start without --workflow and --workdir.""" + with self.assertRaises(SystemExit): + _parse_engine_args(["--workdir", "/tmp/wd"]) + with self.assertRaises(SystemExit): + _parse_engine_args(["--workflow", "pkg.mod:factory"]) + + def test_engine_passes_tail_through(self) -> None: + """The engine surfaces the workflow tail unchanged to the caller.""" + ns, tail = _parse_engine_args( + [ + "--workflow", + "pkg.mod:factory", + "--workdir", + "/tmp/wd", + "--", + "--build-command", + "make", + ] + ) + self.assertEqual(ns.workflow, "pkg.mod:factory") + self.assertEqual(ns.workdir, "/tmp/wd") + self.assertEqual(tail, ["--build-command", "make"]) + + def test_engine_max_concurrency_defaults_to_none(self) -> None: + """Omitting --max-concurrency leaves ns.max_concurrency as None. + + The engine treats this as "use the workflow-supplied default" so + the WorkflowSpec.max_concurrency value is honored when the user + does not override it on the CLI. + """ + ns, _ = _parse_engine_args( + ["--workflow", "pkg.mod:factory", "--workdir", "/tmp/wd"] + ) + self.assertIsNone(ns.max_concurrency) + + def test_engine_max_concurrency_accepts_explicit_value(self) -> None: + """An explicit --max-concurrency value is preserved as an int.""" + ns, _ = _parse_engine_args( + [ + "--workflow", + "pkg.mod:factory", + "--workdir", + "/tmp/wd", + "--max-concurrency", + "7", + ] + ) + self.assertEqual(ns.max_concurrency, 7) + + +class _EngineCommandTest(unittest.TestCase): + def test_default_uses_spdl_autoresearch_engine(self) -> None: + """Without an override, the engine prefix is 'spdl autoresearch engine'.""" + cmd = _build_engine_command( + engine_command_override=None, + workflow_spec="pkg.mod:factory", + workdir=Path("/tmp/wd"), + framework_flags=["--max-concurrency", "3"], + workflow_argv_tail=["--build-command", "make"], + ) + self.assertEqual( + cmd, + [ + "spdl", + "autoresearch", + "engine", + "--workflow", + "pkg.mod:factory", + "--workdir", + "/tmp/wd", + "--max-concurrency", + "3", + "--", + "--build-command", + "make", + ], + ) + + def test_override_replaces_prefix(self) -> None: + """An --engine-command override replaces the default argv[0] prefix.""" + cmd = _build_engine_command( + engine_command_override="buck run //x:engine --", + workflow_spec="pkg.mod:factory", + workdir=Path("/tmp/wd"), + framework_flags=[], + workflow_argv_tail=[], + ) + self.assertEqual( + cmd, + [ + "buck", + "run", + "//x:engine", + "--", + "--workflow", + "pkg.mod:factory", + "--workdir", + "/tmp/wd", + ], + ) + + def test_no_tail_omits_double_dash(self) -> None: + """An empty workflow tail does not append a stray '--'.""" + cmd = _build_engine_command( + engine_command_override=None, + workflow_spec="pkg.mod:factory", + workdir=Path("/tmp/wd"), + framework_flags=[], + workflow_argv_tail=[], + ) + self.assertNotIn("--", cmd[2:]) + + +class _CoreWorkflowExportTest(unittest.TestCase): + def test_workflow_spec_is_protocol(self) -> None: + """``WorkflowSpec`` re-exported from core is a ``Protocol`` subclass. + + ``Protocol`` subclasses are flagged with ``_is_protocol = True`` by + the typing machinery; this guards against accidentally weakening + ``WorkflowSpec`` to a regular class (which would silently change + the runtime semantics for workflow authors). + """ + from spdl.autoresearch.core import WorkflowSpec + + self.assertTrue(getattr(WorkflowSpec, "_is_protocol", False)) + + def test_workflow_factory_is_callable_alias(self) -> None: + """``WorkflowFactory`` re-exported from core is a ``Callable`` alias.""" + from spdl.autoresearch.core import WorkflowFactory + + self.assertIs(get_origin(WorkflowFactory), Callable) + + +class _MainImportTest(unittest.TestCase): + def test_main_import_does_not_load_app(self) -> None: + """Importing spdl.autoresearch.__main__ as a module is a no-op. + + The framework dispatcher (under spdl.autoresearch._app) must + NOT be transitively loaded by ``import + spdl.autoresearch.__main__``. _app is reachable only when + __main__.py runs as a script (i.e. via ``python -m + spdl.autoresearch``), at which point ``__name__ == + "__main__"`` and the lazy import inside the guard fires. + """ + removed = {} + for mod_name in [ + name + for name in list(sys.modules) + if name == "spdl.autoresearch.__main__" + or name.startswith("spdl.autoresearch._app") + ]: + removed[mod_name] = sys.modules.pop(mod_name) + self.addCleanup(sys.modules.update, removed) + + importlib.import_module("spdl.autoresearch.__main__") + + self.assertNotIn("spdl.autoresearch._app", sys.modules) + self.assertNotIn("spdl.autoresearch._app._main", sys.modules) diff --git a/src/spdl/autoresearch/tests/factory_test.py b/src/spdl/autoresearch/tests/factory_test.py new file mode 100644 index 000000000..95d37a511 --- /dev/null +++ b/src/spdl/autoresearch/tests/factory_test.py @@ -0,0 +1,205 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +from spdl.autoresearch.pipeline_optimization import create_workflow +from spdl.autoresearch.pipeline_optimization._ops._analysis_ops import ( + MASTER_TABLE_HEADERS, +) +from spdl.autoresearch.pipeline_optimization._ops._policy import write_state + +__all__: list[str] = [] + + +def _full_argv() -> list[str]: + return [ + "--pipeline-script", + "/tmp/pipeline.py", + "--source-dir", + "/tmp/src", + "--build-command", + "make image", + "--base-launch-command", + "torchx run --image $IMAGE", + "--notes", + "smoke", + "--max-iterations", + "5", + "--patience", + "2", + "--job-timeout", + "300", + ] + + +class _CreateWorkflowTest(unittest.TestCase): + def test_returns_workflow_spec(self) -> None: + """create_workflow returns an object exposing the WorkflowSpec surface.""" + spec = create_workflow(_full_argv(), None) + for attr in ( + "engine_argv_tail", + "description", + "supervisor_known_config", + "supervisor_missing_config", + "setup", + "build_workflow", + ): + self.assertTrue(callable(getattr(spec, attr)), attr) + self.assertEqual(spec.max_concurrency, 3) + + def test_engine_argv_tail_round_trips_supplied_flags(self) -> None: + """Every value passed in survives in engine_argv_tail in flag/value order.""" + spec = create_workflow(_full_argv(), None) + tail = spec.engine_argv_tail() + self.assertEqual(tail[tail.index("--pipeline-script") + 1], "/tmp/pipeline.py") + self.assertEqual(tail[tail.index("--source-dir") + 1], "/tmp/src") + self.assertEqual(tail[tail.index("--build-command") + 1], "make image") + self.assertEqual( + tail[tail.index("--base-launch-command") + 1], + "torchx run --image $IMAGE", + ) + self.assertEqual(tail[tail.index("--max-iterations") + 1], "5") + self.assertEqual(tail[tail.index("--patience") + 1], "2") + self.assertEqual(tail[tail.index("--max-concurrency") + 1], "3") + self.assertEqual(tail[tail.index("--job-timeout") + 1], "300") + self.assertEqual(tail[tail.index("--platform") + 1], "auto") + + def test_engine_argv_tail_omits_unset_options(self) -> None: + """Unset optional flags are not emitted at all.""" + spec = create_workflow([], None) + tail = spec.engine_argv_tail() + self.assertNotIn("--pipeline-script", tail) + self.assertNotIn("--build-command", tail) + self.assertNotIn("--base-launch-command", tail) + self.assertNotIn("--source-dir", tail) + self.assertNotIn("--notes", tail) + + def test_engine_argv_tail_emits_boolean_flags(self) -> None: + """Boolean flags appear by themselves with no value when set.""" + spec = create_workflow( + [ + "--skip-instrument", + "--dangerously-skip-permissions", + ], + None, + ) + tail = spec.engine_argv_tail() + self.assertIn("--skip-instrument", tail) + self.assertIn("--dangerously-skip-permissions", tail) + + def test_supervisor_missing_config_lists_required_fields(self) -> None: + """A bare invocation reports all four required fields as missing.""" + spec = create_workflow([], None) + missing = spec.supervisor_missing_config() + self.assertIn("pipeline script", missing) + self.assertIn("source directory", missing) + self.assertIn("build command", missing) + self.assertIn("launch command template", missing) + + def test_supervisor_missing_config_empty_when_all_supplied(self) -> None: + """A fully-configured invocation reports no missing fields.""" + spec = create_workflow(_full_argv(), None) + self.assertEqual(spec.supervisor_missing_config(), []) + + def test_supervisor_known_config_reflects_argv(self) -> None: + """supervisor_known_config exposes the parsed values for the supervisor prompt.""" + spec = create_workflow(_full_argv(), None) + known = spec.supervisor_known_config() + self.assertEqual(known["pipeline_script"], "/tmp/pipeline.py") + self.assertEqual(known["build_command"], "make image") + self.assertEqual(known["local_execution_mode"], "full") + + def test_max_concurrency_reflects_argv(self) -> None: + """A non-default --max-concurrency is visible on spec.max_concurrency.""" + spec = create_workflow([*_full_argv(), "--max-concurrency", "7"], None) + self.assertEqual(spec.max_concurrency, 7) + + def test_description_contains_supervisor_and_platform_content(self) -> None: + """description() joins the supervisor and platform prompt directories.""" + spec = create_workflow(_full_argv(), None) + description = spec.description() + self.assertIsNotNone(description) + assert description is not None # for type checker + self.assertIn("---", description) + self.assertIn("Automated SPDL Pipeline Optimization", description) + + +def _write_minimal_workdir(workdir: Path) -> None: + """Create the minimum files PipelineOptimizationWorkflow.summarize reads.""" + workdir.mkdir(parents=True, exist_ok=True) + (workdir / "config.json").write_text( + json.dumps( + { + "schema_version": 1, + "pipeline_script": "", + "source_dir": "", + "scm": "", + "build_command": "", + "base_launch_command": "", + "stopping_criteria": {"max_iterations": 1, "patience": 1}, + "max_concurrency": 1, + "job_timeout_s": 60, + "poll_interval": 0, + "platform": "auto", + "agent": "claude", + "local_execution_mode": "full", + } + ) + ) + write_state( + workdir, + { + "iteration": 0, + "status": "looping", + "baseline_job": None, + "current_best": None, + "best_metric": None, + "plateau_count": 0, + "best_practices_tried": [], + "history": [], + }, + ) + (workdir / "master_table.tsv").write_text("\t".join(MASTER_TABLE_HEADERS) + "\n") + + +class _SummarizeTest(unittest.TestCase): + def test_summarize_returns_markdown_for_empty_workdir(self) -> None: + """summarize handles a freshly initialised workdir without raising.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + _write_minimal_workdir(workdir) + spec = create_workflow([], workdir) + workflow = spec.build_workflow(workdir) + + output = workflow.summarize(workdir) + + self.assertIn("# Autoresearch summary", output) + self.assertIn(str(workdir), output) + self.assertIn("## Failures", output) + + def test_summarize_includes_master_table_and_live_summary(self) -> None: + """summarize renders master-table rows and summary.md content.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + _write_minimal_workdir(workdir) + with open(workdir / "master_table.tsv", "a") as f: + f.write("run_001\tbaseline\t0.5\n") + (workdir / "summary.md").write_text("Best SM util improved to 85%") + spec = create_workflow([], workdir) + workflow = spec.build_workflow(workdir) + + output = workflow.summarize(workdir) + + self.assertIn("## Master table", output) + self.assertIn("run_001", output) + self.assertIn("## Live summary", output) + self.assertIn("Best SM util improved to 85%", output) diff --git a/src/spdl/autoresearch/tests/orchestrator_test.py b/src/spdl/autoresearch/tests/orchestrator_test.py new file mode 100644 index 000000000..b72e6ac5f --- /dev/null +++ b/src/spdl/autoresearch/tests/orchestrator_test.py @@ -0,0 +1,128 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import asyncio +import unittest +from pathlib import Path + +from spdl.autoresearch.core import Orchestrator, TaskResult, TaskSpec + +__all__: list[str] = [] + + +class _FakeAdapter: + def __init__(self, specs: list[TaskSpec]) -> None: + self.specs = specs + self.started: list[str] = [] + self.checkpoints: list[tuple[list[str], list[str], str]] = [] + self.children: dict[str, list[TaskSpec]] = {} + self.block = False + + def load(self) -> list[TaskSpec]: + return self.specs + + def checkpoint( + self, + queued: list[TaskSpec], + running: list[TaskSpec], + status: str, + ) -> None: + self.checkpoints.append( + ([spec.id for spec in queued], [spec.id for spec in running], status) + ) + + async def make_coro(self, spec: TaskSpec) -> TaskResult: + self.started.append(spec.id) + if self.block: + await asyncio.sleep(60) + return TaskResult(children=self.children.get(spec.id, [])) + + async def on_result(self, spec: TaskSpec, result: TaskResult) -> list[TaskSpec]: + return result.children + + def summarize(self, workdir: Path) -> str: + return f"_FakeAdapter summary at {workdir}" + + +class OrchestratorTest(unittest.IsolatedAsyncioTestCase): + async def test_priority_order_lowest_first(self) -> None: + """Specs are executed in ascending priority order (lowest value first).""" + adapter = _FakeAdapter( + [ + TaskSpec(id="slow", priority=10), + TaskSpec(id="first", priority=-1), + TaskSpec(id="middle", priority=5), + ] + ) + + await Orchestrator(workflow=adapter, max_concurrency=1).run() + + self.assertEqual(["first", "middle", "slow"], adapter.started) + self.assertEqual(([], [], "stopped"), adapter.checkpoints[-1]) + + async def test_completion_enqueues_children(self) -> None: + """Child specs returned by a completed item are enqueued and executed.""" + adapter = _FakeAdapter([TaskSpec(id="root", priority=0)]) + adapter.children["root"] = [ + TaskSpec(id="child_a", priority=1), + TaskSpec(id="child_b", priority=2), + ] + + await Orchestrator(workflow=adapter, max_concurrency=1).run() + + self.assertEqual(["root", "child_a", "child_b"], adapter.started) + self.assertEqual(([], [], "stopped"), adapter.checkpoints[-1]) + + async def test_cancelled_error_persists_interrupted_state(self) -> None: + """Cancellation checkpoints running specs with 'interrupted' status.""" + adapter = _FakeAdapter([TaskSpec(id="running", priority=0)]) + adapter.block = True + task = asyncio.create_task( + Orchestrator(workflow=adapter, max_concurrency=1).run() + ) + + while not adapter.checkpoints: + await asyncio.sleep(0) + task.cancel() + await task + + self.assertEqual(([], ["running"], "interrupted"), adapter.checkpoints[-1]) + + async def test_checkpoint_resume_golden_lifecycle(self) -> None: + """An interrupted engine can resume from checkpointed state and complete.""" + first = _FakeAdapter( + [ + TaskSpec(id="running", priority=0), + TaskSpec(id="queued", priority=1), + ] + ) + first.block = True + task = asyncio.create_task( + Orchestrator(workflow=first, max_concurrency=1).run() + ) + + while not first.checkpoints: + await asyncio.sleep(0) + task.cancel() + await task + + queued_ids, running_ids, status = first.checkpoints[-1] + self.assertEqual( + (["queued"], ["running"], "interrupted"), first.checkpoints[-1] + ) + + resumed_specs = [ + TaskSpec(id=spec_id, priority=0 if spec_id == "running" else 1) + for spec_id in running_ids + queued_ids + ] + resumed = _FakeAdapter(resumed_specs) + await Orchestrator(workflow=resumed, max_concurrency=1).run() + + self.assertEqual("interrupted", status) + self.assertEqual(["running", "queued"], resumed.started) + self.assertEqual(([], [], "stopped"), resumed.checkpoints[-1]) diff --git a/src/spdl/autoresearch/tests/persistence_test.py b/src/spdl/autoresearch/tests/persistence_test.py new file mode 100644 index 000000000..57a0b72d1 --- /dev/null +++ b/src/spdl/autoresearch/tests/persistence_test.py @@ -0,0 +1,168 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import json +import tempfile +import unittest +import unittest.mock +from pathlib import Path + +from spdl.autoresearch.core import ( + load_or_init, + read_engine_state, + TaskSpec, + write_engine_state, +) + +__all__: list[str] = [] + + +class _PersistenceTest(unittest.TestCase): + def test_round_trip_preserves_spec_fields(self) -> None: + """write_engine_state followed by read_engine_state preserves all TaskSpec fields.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + queued = [ + TaskSpec( + id="exp_001", + priority=-1.5, + kind="experiment", + payload={"node": {"node_id": "exp_001"}, "extra": [1, 2, 3]}, + ), + TaskSpec(id="exp_002", priority=0.0, kind="default", payload={}), + ] + running = [TaskSpec(id="exp_003", priority=2.0, kind="experiment")] + + write_engine_state( + workdir, queued=queued, running=running, status="running" + ) + result = read_engine_state(workdir) + + self.assertIsNotNone(result) + assert result is not None # for type checker + got_queued, got_running, status = result + self.assertEqual(status, "running") + self.assertEqual([spec.id for spec in got_queued], ["exp_001", "exp_002"]) + self.assertEqual(got_queued[0].priority, -1.5) + self.assertEqual(got_queued[0].kind, "experiment") + self.assertEqual( + got_queued[0].payload, + {"node": {"node_id": "exp_001"}, "extra": [1, 2, 3]}, + ) + self.assertEqual([spec.id for spec in got_running], ["exp_003"]) + + def test_read_returns_none_when_missing(self) -> None: + """read_engine_state on a fresh workdir returns None instead of raising.""" + with tempfile.TemporaryDirectory() as tmp: + self.assertIsNone(read_engine_state(Path(tmp))) + + def test_status_round_trips(self) -> None: + """All three orchestrator status values survive persistence.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + for status in ("running", "stopped", "interrupted"): + write_engine_state(workdir, queued=[], running=[], status=status) + result = read_engine_state(workdir) + assert result is not None + self.assertEqual(result[2], status) + + def test_write_creates_workdir(self) -> None: + """write_engine_state creates the workdir if it does not yet exist.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) / "nested" / "fresh" + self.assertFalse(workdir.exists()) + write_engine_state(workdir, queued=[], running=[], status="running") + self.assertTrue((workdir / "engine_state.json").exists()) + + def test_load_or_init_uses_factory_on_fresh_run(self) -> None: + """load_or_init calls the factory exactly once when no checkpoint exists.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + calls: list[int] = [] + + def factory() -> list[TaskSpec]: + calls.append(1) + return [TaskSpec(id="seed", priority=0.0)] + + specs = load_or_init(workdir, factory) + + self.assertEqual(len(calls), 1) + self.assertEqual([spec.id for spec in specs], ["seed"]) + + def test_load_or_init_resumes_from_checkpoint(self) -> None: + """load_or_init returns queued+running and skips the factory on resume.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + queued = [TaskSpec(id="q1", priority=-1.0)] + running = [TaskSpec(id="r1", priority=0.0)] + write_engine_state( + workdir, queued=queued, running=running, status="interrupted" + ) + + def factory() -> list[TaskSpec]: + self.fail("factory must not be invoked when checkpoint exists") + + specs = load_or_init(workdir, factory) + + self.assertEqual([spec.id for spec in specs], ["q1", "r1"]) + + def test_read_rejects_malformed_json_object(self) -> None: + """A non-object JSON file raises ValueError instead of silently misparsing.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + (workdir / "engine_state.json").write_text("[]\n") + with self.assertRaises(ValueError): + read_engine_state(workdir) + + def test_read_rejects_non_list_field(self) -> None: + """A non-list queued/running field raises ValueError.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + (workdir / "engine_state.json").write_text( + json.dumps({"status": "running", "queued": "oops", "running": []}) + ) + with self.assertRaises(ValueError): + read_engine_state(workdir) + + def test_write_does_not_truncate_existing_checkpoint_on_failure(self) -> None: + """A failing write leaves the previous engine_state.json intact. + + Simulates a mid-write interruption by patching the temp file's + ``write_text`` to raise after the previous checkpoint has been + written successfully. The reader must still see the previous + valid checkpoint, never a truncated or partial file. + """ + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + write_engine_state( + workdir, + queued=[TaskSpec(id="prev", priority=0.0)], + running=[], + status="running", + ) + + original_write_text = Path.write_text + + def _fail_on_tmp(self: Path, *args: object, **kwargs: object) -> int: + if ".tmp." in self.name: + raise OSError("simulated mid-write failure") + return original_write_text(self, *args, **kwargs) # type: ignore[arg-type] + + with unittest.mock.patch.object(Path, "write_text", _fail_on_tmp): + with self.assertRaises(OSError): + write_engine_state( + workdir, + queued=[TaskSpec(id="new", priority=0.0)], + running=[], + status="running", + ) + + result = read_engine_state(workdir) + assert result is not None + queued, _, _ = result + self.assertEqual([spec.id for spec in queued], ["prev"]) diff --git a/src/spdl/autoresearch/tests/platform_test.py b/src/spdl/autoresearch/tests/platform_test.py new file mode 100644 index 000000000..235feffcd --- /dev/null +++ b/src/spdl/autoresearch/tests/platform_test.py @@ -0,0 +1,228 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import os +import shlex +import sys +import tempfile +import time +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + +from spdl.autoresearch.pipeline_optimization._ops._store import _set_queued_priority +from spdl.autoresearch.pipeline_optimization._platform import ( + _MetricsEvidence, + AutoresearchPlatform, + create_platform, +) +from spdl.autoresearch.pipeline_optimization._platform._agents import ( + _MockAgent, + _parse_agent_result, +) + +__all__: list[str] = [] + + +# The GitHub Actions Windows runner ships a conda Python whose PATH/DLL setup +# breaks nested ``cmd.exe`` invocation: ``subprocess.Popen(shell=True, ...)`` +# returns NT status ``0xC0000142`` (``STATUS_DLL_INIT_FAILED``) before the +# child shell can run anything, regardless of what command we give it. This is +# an environment bug in the runner, not in the production code under test, so +# we skip the two subprocess-launching tests there. Other CI matrices (Linux, +# macOS, internal Windows) still exercise this code path. +_SKIP_WINDOWS_GHA: bool = ( + sys.platform == "win32" and os.environ.get("GITHUB_ACTIONS") == "true" +) +_SKIP_REASON: str = ( + "GitHub Actions Windows runner cannot spawn nested cmd.exe (STATUS_DLL_INIT_FAILED)" +) + + +def _echo_marker_command(workdir: Path, line: str) -> str: + """Build a cross-platform shell command that prints ``line`` to stdout. + + We write a tiny shell script on disk and return its path as the command, + so the launched process is just the shell executing one builtin (``echo``) + — no external interpreter is loaded. + + Why not invoke a second ``python.exe``? On the GitHub Actions + Windows-miniconda runner, spawning a nested ``python.exe`` from + ``subprocess.Popen(shell=True, ...)`` returns NT status + ``0xC0000142`` (``STATUS_DLL_INIT_FAILED``) before the script runs. + cmd's ``echo`` is an internal command, so no DLL initialization happens + in the child process. + """ + workdir.mkdir(parents=True, exist_ok=True) + if sys.platform == "win32": + script = workdir / f"_marker_{uuid.uuid4().hex}.cmd" + script.write_text(f"@echo off\r\necho {line}\r\n", encoding="ascii") + return f'"{script}"' + script = workdir / f"_marker_{uuid.uuid4().hex}.sh" + script.write_text(f"#!/bin/sh\necho {shlex.quote(line)}\n", encoding="ascii") + script.chmod(0o755) + return shlex.quote(str(script)) + + +class _PlatformTest(unittest.TestCase): + def test_default_platform_has_capability_parts(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + with patch.dict( + "os.environ", + {"SPDL_AUTORESEARCH_PLATFORM_PROVIDERS": ""}, + ): + platform = create_platform("auto", Path(tmp)) + + self.assertIsInstance(platform, AutoresearchPlatform) + self.assertTrue(hasattr(platform, "workspace")) + self.assertTrue(hasattr(platform, "artifacts")) + self.assertTrue(hasattr(platform, "execution")) + self.assertTrue(hasattr(platform, "evidence")) + self.assertTrue(hasattr(platform, "agent")) + + @unittest.skipIf(_SKIP_WINDOWS_GHA, _SKIP_REASON) + def test_local_platform_runs_subprocess_and_collects_log_evidence(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + platform = create_platform("local", workdir) + + job_id = platform.execution.launch( + _echo_marker_command(workdir, "[autoresearch] step=1"), + workdir, + ) + self.assertIsNotNone(job_id) + assert job_id is not None + + for _ in range(50): + status = platform.execution.status(job_id) + if status == "SUCCEEDED": + break + time.sleep(0.05) + self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) + self.assertEqual( + "[autoresearch] step=1", platform.execution.progress(job_id) + ) + + metrics_dir = workdir / "runs" / "000_baseline" / "metrics" + evidence = platform.evidence.collect(job_id, metrics_dir) + + self.assertIsInstance(evidence, _MetricsEvidence) + self.assertIn("system metrics unavailable", evidence.system_metrics) + self.assertIn("[autoresearch] step=1", evidence.pipeline_stats_log) + + def test_local_dry_run_completes_without_launching_subprocess(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + platform = create_platform( + {"platform": "local", "local_execution_mode": "dry_run"}, + workdir, + ) + + job_id = platform.execution.launch( + _echo_marker_command(workdir, "not executed"), workdir + ) + self.assertIsNotNone(job_id) + assert job_id is not None + + self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) + self.assertIn("dry_run", platform.execution.progress(job_id) or "") + + @unittest.skipIf(_SKIP_WINDOWS_GHA, _SKIP_REASON) + def test_local_dataloader_only_uses_dataloader_command(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + platform = create_platform( + { + "platform": "local", + "local_execution_mode": "dataloader_only", + "local_dataloader_command": _echo_marker_command( + workdir, "[autoresearch] dataloader" + ), + }, + workdir, + ) + + job_id = platform.execution.launch( + _echo_marker_command(workdir, "training"), workdir + ) + self.assertIsNotNone(job_id) + assert job_id is not None + + for _ in range(50): + status = platform.execution.status(job_id) + if status == "SUCCEEDED": + break + time.sleep(0.05) + + self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) + self.assertEqual( + "[autoresearch] dataloader", platform.execution.progress(job_id) + ) + + def test_mock_agent_is_selected_independently_of_platform(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + platform = create_platform( + {"platform": "local", "agent": "mock"}, + Path(tmp), + ) + + self.assertEqual("", platform.agent.run("prompt", Path(tmp), "phase")) + + def test_platform_config_validation_rejects_bad_local_mode(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaisesRegex(ValueError, "local execution mode"): + create_platform( + {"platform": "local", "local_execution_mode": "unknown"}, + Path(tmp), + ) + + def test_unknown_remote_provider_is_explicit(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + with patch.dict( + "os.environ", + {"SPDL_AUTORESEARCH_PLATFORM_PROVIDERS": ""}, + ): + with self.assertRaisesRegex( + ValueError, "Unknown autoresearch platform" + ): + create_platform( + {"platform": "unknown_remote", "agent": "mock"}, + Path(tmp), + ) + + def test_agent_result_reports_parse_errors_without_llm(self) -> None: + agent = _MockAgent() + + parsed = _parse_agent_result(agent, '```json\n{"action": "stop"}\n```') + failed = _parse_agent_result(agent, "not json") + + self.assertEqual({"action": "stop"}, parsed.json) + self.assertIsNone(parsed.parse_error) + self.assertIsNone(failed.json) + self.assertEqual("No JSON object found", failed.parse_error) + + def test_queue_command_updates_checkpoint_priority(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + engine = workdir / "engine" + engine.mkdir(parents=True) + (engine / "checkpoint.json").write_text( + '{"status": "interrupted", "queued": [' + '{"id": "001_a", "priority": 10, "kind": "experiment", ' + '"payload": {"node": {"node_id": "001_a"}}}], "running": []}\n' + ) + + _set_queued_priority(workdir, "001_a", -5) + + text = (engine / "checkpoint.json").read_text() + self.assertIn('"priority": -5.0', text) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/spdl/autoresearch/tests/prompts_test.py b/src/spdl/autoresearch/tests/prompts_test.py new file mode 100644 index 000000000..e75177044 --- /dev/null +++ b/src/spdl/autoresearch/tests/prompts_test.py @@ -0,0 +1,129 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest + +from spdl.autoresearch.pipeline_optimization._prompts import ( + load_knowledge, + load_prompt, + load_prompt_directory, +) + + +class LoadPromptTest(unittest.TestCase): + def test_loads_existing_prompt(self) -> None: + """A valid prompt name returns non-empty template content.""" + result = load_prompt("analyze", KNOWLEDGE="test-knowledge") + + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + def test_substitutes_placeholders(self) -> None: + """All __KEY__ placeholders are replaced with the supplied values.""" + result = load_prompt( + "headspace", + KNOWLEDGE="INJECTED_KNOWLEDGE", + PIPELINE_SCRIPT="/tmp/test.py", + PIPELINE_CODE="def main(): pass", + ) + + self.assertIn("INJECTED_KNOWLEDGE", result) + self.assertIn("/tmp/test.py", result) + self.assertIn("def main(): pass", result) + self.assertNotIn("__KNOWLEDGE__", result) + self.assertNotIn("__PIPELINE_SCRIPT__", result) + self.assertNotIn("__PIPELINE_CODE__", result) + + def test_missing_prompt_exits(self) -> None: + """Requesting a nonexistent prompt template triggers SystemExit.""" + with self.assertRaises(SystemExit): + load_prompt("nonexistent_prompt_that_does_not_exist") + + def test_headspace_prompt_requires_stop_after(self) -> None: + """The headspace prompt instructs the agent to include stop_after=500.""" + prompt = load_prompt( + "headspace", + KNOWLEDGE="", + PIPELINE_SCRIPT="/tmp/pipeline.py", + PIPELINE_CODE="def main():\n pass\n", + ) + + self.assertIn("stop_after=500", prompt) + self.assertIn("must include `stop_after=500`", prompt) + + def test_all_phase_prompts_loadable(self) -> None: + """Every phase prompt shipped with the package loads without error.""" + phase_prompts = [ + "analyze", + "apply_changes", + "apply_startup_repair", + "assess", + "headspace", + "instrument", + "plan_next", + ] + for name in phase_prompts: + with self.subTest(prompt=name): + result = load_prompt(name, KNOWLEDGE="k") + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + +class LoadPromptDirectoryTest(unittest.TestCase): + def test_loads_knowledge_directory(self) -> None: + """The knowledge directory contains at least one .md file.""" + result = load_prompt_directory("knowledge") + + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + def test_loads_supervisor_directory(self) -> None: + """The supervisor directory contains at least one .md file.""" + result = load_prompt_directory("supervisor") + + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + def test_loads_platform_directory(self) -> None: + """The platform directory contains at least one .md file.""" + result = load_prompt_directory("platform") + + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + def test_nonexistent_directory_returns_empty(self) -> None: + """A missing directory returns an empty string instead of raising.""" + result = load_prompt_directory("nonexistent_dir") + + self.assertEqual(result, "") + + def test_deterministic_order(self) -> None: + """Repeated loads produce identical output (sorted path order).""" + first = load_prompt_directory("knowledge") + second = load_prompt_directory("knowledge") + + self.assertEqual(first, second) + + +class LoadKnowledgeTest(unittest.TestCase): + def test_returns_nonempty_string(self) -> None: + """The combined knowledge + platform content is non-empty.""" + result = load_knowledge() + + self.assertIsInstance(result, str) + self.assertGreater(len(result), 0) + + def test_includes_knowledge_and_platform_content(self) -> None: + """The result contains the full text of both knowledge and platform directories.""" + result = load_knowledge() + knowledge_only = load_prompt_directory("knowledge") + platform_only = load_prompt_directory("platform") + + for section in (knowledge_only, platform_only): + if section: + self.assertIn(section, result) diff --git a/src/spdl/autoresearch/tests/state_test.py b/src/spdl/autoresearch/tests/state_test.py new file mode 100644 index 000000000..83f186da4 --- /dev/null +++ b/src/spdl/autoresearch/tests/state_test.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest + +from spdl.autoresearch._common._state import SCHEMA_VERSION +from spdl.autoresearch.pipeline_optimization._ops._policy import ( + _normalize_config, + _normalize_state, +) + + +class StateTest(unittest.TestCase): + def test_schema_normalizers_add_versions_and_defaults(self) -> None: + """Empty dicts are normalized with schema version and default values.""" + config = _normalize_config({}) + state = _normalize_state({}) + + self.assertEqual(SCHEMA_VERSION, config["schema_version"]) + self.assertEqual(SCHEMA_VERSION, state["schema_version"]) + self.assertEqual("auto", config["platform"]) + self.assertEqual([], state["history"]) diff --git a/src/spdl/autoresearch/tests/supervisor_test.py b/src/spdl/autoresearch/tests/supervisor_test.py new file mode 100644 index 000000000..b89bb3cca --- /dev/null +++ b/src/spdl/autoresearch/tests/supervisor_test.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest + +from spdl.autoresearch._common._supervisor import ( + _ClaudeSupervisor, + _CodexSupervisor, +) + + +class SupervisorTest(unittest.TestCase): + def test_claude_supervisor_builds_system_prompt_command(self) -> None: + """Claude supervisor passes system prompt and initial request as separate args.""" + command = _ClaudeSupervisor().command("SYSTEM", "INITIAL") + + self.assertEqual(["claude", "--system-prompt", "SYSTEM", "INITIAL"], command) + + def test_codex_supervisor_merges_system_and_request_into_single_prompt( + self, + ) -> None: + """Codex supervisor combines system prompt and request into one argument.""" + command = _CodexSupervisor().command("SYSTEM", "INITIAL") + + self.assertEqual("codex", command[0]) + self.assertEqual(2, len(command)) + self.assertIn("SYSTEM", command[1]) + self.assertIn("## User Request", command[1]) + self.assertIn("INITIAL", command[1]) diff --git a/src/spdl/autoresearch/tests/types_test.py b/src/spdl/autoresearch/tests/types_test.py new file mode 100644 index 000000000..6eb48db61 --- /dev/null +++ b/src/spdl/autoresearch/tests/types_test.py @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest + +from spdl.autoresearch.core import ( + FailureKind, + FailurePhase, + FailureRecord, + HypothesisNode, +) + + +class FailureRecordTest(unittest.TestCase): + def test_round_trips_through_dict(self) -> None: + """FailureRecord survives to_dict/from_dict serialization.""" + record = FailureRecord( + kind=FailureKind.JOB_STARTUP_FAILED, + phase=FailurePhase.JOB, + message="MTP failed during initialization", + details={"component": "mtp"}, + job_id="job123", + created_at="2026-01-01T00:00:00", + ) + + loaded = FailureRecord.from_dict(record.to_dict()) + + self.assertEqual(FailureKind.JOB_STARTUP_FAILED, loaded.kind) + self.assertEqual(FailurePhase.JOB, loaded.phase) + self.assertEqual("MTP failed during initialization", loaded.message) + self.assertEqual({"component": "mtp"}, loaded.details) + self.assertEqual("job123", loaded.job_id) + + +class HypothesisNodeTest(unittest.TestCase): + def test_round_trips_through_dict(self) -> None: + """HypothesisNode with a failure survives to_dict/from_dict.""" + node = HypothesisNode( + node_id="001_bad_mtp", + name="bad_mtp", + status="failed", + failure=FailureRecord( + kind=FailureKind.JOB_STARTUP_FAILED, + phase=FailurePhase.JOB, + message="MTP failed during initialization", + details={"component": "mtp"}, + job_id="job123", + created_at="2026-01-01T00:00:00", + ), + ) + + loaded = HypothesisNode.from_dict(node.to_dict()) + + self.assertEqual("001_bad_mtp", loaded.node_id) + self.assertEqual("failed", loaded.status) + failure = loaded.failure + assert failure is not None + self.assertEqual(FailureKind.JOB_STARTUP_FAILED, failure.kind) + self.assertEqual({"component": "mtp"}, failure.details) + + def test_round_trips_without_failure(self) -> None: + """HypothesisNode without a failure round-trips cleanly.""" + node = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="completed", + priority=-1000, + ) + + loaded = HypothesisNode.from_dict(node.to_dict()) + + self.assertEqual("000_baseline", loaded.node_id) + self.assertEqual("completed", loaded.status) + self.assertIsNone(loaded.failure) + self.assertEqual(-1000, loaded.priority) diff --git a/src/spdl/autoresearch/tests/visualization_test.py b/src/spdl/autoresearch/tests/visualization_test.py new file mode 100644 index 000000000..b7db34ef3 --- /dev/null +++ b/src/spdl/autoresearch/tests/visualization_test.py @@ -0,0 +1,44 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import unittest + +from spdl.autoresearch._common._visualization import ( + _edge_label, + _tree_font_sizes, +) + + +class VisualizationTest(unittest.TestCase): + def test_tree_edge_label_describes_experiment_evolution(self) -> None: + """Edge labels show change summary or startup repair attempt number.""" + parent = {"node_id": "001_parent", "name": "parent", "spec": {}} + child = { + "node_id": "002_child", + "name": "child", + "spec": { + "change_summary": "raise decode threads", + "description": "increase decode thread count", + }, + } + retry = { + "node_id": "003_retry", + "name": "retry", + "spec": {"_startup_retry_attempt": 2, "description": "repair"}, + } + + self.assertEqual("raise decode threads", _edge_label(parent, child)) + self.assertEqual("startup repair #2", _edge_label(parent, retry)) + + def test_tree_font_sizes_grow_with_tree_size(self) -> None: + """Larger trees get proportionally larger font sizes.""" + small = _tree_font_sizes(4, 1) + large = _tree_font_sizes(120, 10) + + self.assertGreater(large["title"], small["title"]) + self.assertGreater(large["legend"], small["legend"]) diff --git a/src/spdl/autoresearch/tests/workflow_test.py b/src/spdl/autoresearch/tests/workflow_test.py new file mode 100644 index 000000000..0eb66e8a6 --- /dev/null +++ b/src/spdl/autoresearch/tests/workflow_test.py @@ -0,0 +1,1049 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import asyncio +import json +import tempfile +import unittest +from pathlib import Path + +from spdl.autoresearch._common._state import _append_master_row +from spdl.autoresearch._common._visualization import _load_tsv +from spdl.autoresearch.core import ( + AnalysisResult, + FailureKind, + FailurePhase, + HypothesisNode, + TaskSpec, +) +from spdl.autoresearch.pipeline_optimization._ops import ( + _WorkflowStateStore, + PipelineOptimizationWorkflow, +) +from spdl.autoresearch.pipeline_optimization._ops._analysis_ops import ( + _update_on_complete, + MASTER_TABLE_HEADERS, +) +from spdl.autoresearch.pipeline_optimization._ops._failures import ( + _classify_terminal_job_failure, + _FAILURE_POLICIES, + _failure_summary, + _make_failure, + _read_failures, +) +from spdl.autoresearch.pipeline_optimization._ops._policy import ( + _build_change_set, + _change_summary_for_spec, + _compare_metric_value, + _extract_default_executor_concurrency, + _extract_param_changes, + _extract_total_threads, + _is_duplicate_spec, + _node_from_spec, + _retry_policy_for_failure, + _select_planning_node, + _spec_from_node, + _startup_retry_spec, + _validate_thread_budget, + write_state, +) +from spdl.autoresearch.pipeline_optimization._ops._source_ops import _build_apply_prompt +from spdl.autoresearch.pipeline_optimization._ops._store import _write_text_atomic +from spdl.autoresearch.pipeline_optimization._platform import ( + _MetricsEvidence, + create_platform, +) +from spdl.autoresearch.pipeline_optimization._platform._agents import _MockAgent +from spdl.autoresearch.pipeline_optimization._platform._local import _summarize_error + +__all__: list[str] = [] + + +def _config() -> dict: + return { + "pipeline_script": "", + "source_dir": "", + "scm": "", + "build_command": "", + "base_launch_command": "torchx run example --num-fetch-threads 8", + "stopping_criteria": { + "max_iterations": 20, + "patience": 5, + }, + "max_concurrency": 4, + "job_timeout_s": 600, + "poll_interval": 0, + } + + +def _state() -> dict: + return { + "iteration": 0, + "status": "looping", + "baseline_job": None, + "current_best": None, + "best_metric": None, + "plateau_count": 0, + "best_practices_tried": [], + "anchor_commit": "", + "history": [], + } + + +class _AutoresearchWorkflowTest(unittest.TestCase): + def test_fresh_load_creates_initial_must_run_specs(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + adapter = self._adapter(Path(tmp)) + + specs = adapter.load() + + self.assertEqual( + ["000_baseline", "000_headspace", "001_mtp"], + [spec.id for spec in specs], + ) + self.assertEqual([-1000, -999, -998], [spec.priority for spec in specs]) + + def test_checkpoint_writes_compatibility_files(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + specs = adapter.load() + + adapter.checkpoint(queued=specs[1:], running=specs[:1], status="running") + + engine_state = json.loads( + (workdir / "engine" / "engine_state.json").read_text() + ) + queue = json.loads((workdir / "engine" / "queue.json").read_text()) + active = json.loads((workdir / "engine" / "active.json").read_text()) + baseline_status = ( + workdir / "engine" / "nodes" / "000_baseline" / "status.txt" + ).read_text() + + self.assertEqual("running", engine_state["status"]) + self.assertEqual(2, engine_state["queued"]) + self.assertEqual(1, engine_state["running"]) + self.assertEqual( + ["000_headspace", "001_mtp"], [q["node_id"] for q in queue] + ) + self.assertEqual([], active) + self.assertEqual("queued\n", baseline_status) + + def test_checkpoint_round_trips_specs(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + specs = adapter.load() + adapter.checkpoint(queued=specs[1:], running=specs[:1], status="running") + + loaded = self._adapter(workdir).load() + + self.assertEqual( + ["000_headspace", "001_mtp", "000_baseline"], + [spec.id for spec in loaded], + ) + + def test_planning_is_blocked_until_must_run_experiments_finish(self) -> None: + baseline = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="completed", + ) + headspace = HypothesisNode( + node_id="000_headspace", + name="headspace_cache", + status="queued", + ) + mtp = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="completed", + ) + + selected = _select_planning_node( + baseline, + { + baseline.node_id: baseline, + headspace.node_id: headspace, + mtp.node_id: mtp, + }, + ) + + self.assertIsNone(selected) + + def _load_node(self, spec: TaskSpec) -> HypothesisNode: + """Extract node from a TaskSpec with a safe dict cast for Pyre.""" + node_data = spec.payload["node"] + assert isinstance(node_data, dict) + return HypothesisNode.from_dict(node_data) + + def test_child_spec_parents_under_baseline_when_goto_is_null(self) -> None: + """Experiments that start from anchor (goto=null) should be parented + under the baseline node, not whichever node triggered planning.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + specs = adapter.load() + # Simulate baseline completed with a commit. + baseline_node = self._load_node(specs[0]) + baseline_node.status = "completed" + baseline_node.commit = "baseline_commit_abc" + adapter._store.upsert_node(baseline_node) + + # Simulate MTP completed with a different commit. + mtp_node = self._load_node(specs[2]) + mtp_node.status = "completed" + mtp_node.commit = "mtp_commit_xyz" + adapter._store.upsert_node(mtp_node) + + # Create a child with goto=None (should parent under baseline). + child_spec = adapter._create_child_spec( + mtp_node, + {"name": "nvdec_decode", "goto": None}, + ) + child_node = self._load_node(child_spec) + + self.assertEqual("000_baseline", child_node.parent_id) + self.assertEqual("baseline_commit_abc", child_node.commit) + + def test_child_spec_parents_under_goto_commit_owner(self) -> None: + """Experiments with a goto commit should be parented under the node + that produced that commit.""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + specs = adapter.load() + + baseline_node = self._load_node(specs[0]) + baseline_node.status = "completed" + baseline_node.commit = "baseline_commit_abc" + adapter._store.upsert_node(baseline_node) + + mtp_node = self._load_node(specs[2]) + mtp_node.status = "completed" + mtp_node.commit = "mtp_commit_xyz" + adapter._store.upsert_node(mtp_node) + + # Create a child with goto pointing to MTP's commit. + child_spec = adapter._create_child_spec( + baseline_node, # default parent is baseline + {"name": "batch_on_mtp", "goto": "mtp_commit_xyz"}, + ) + child_node = self._load_node(child_spec) + + self.assertEqual(mtp_node.node_id, child_node.parent_id) + self.assertEqual("mtp_commit_xyz", child_node.commit) + + def test_child_spec_parents_under_baseline_without_goto(self) -> None: + """When goto is absent, the spec defaults to anchor (baseline).""" + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + specs = adapter.load() + + mtp_node = self._load_node(specs[2]) + mtp_node.status = "completed" + mtp_node.commit = "mtp_commit_xyz" + adapter._store.upsert_node(mtp_node) + + # Ensure baseline exists so _resolve_parent can find it. + baseline_node = self._load_node(specs[0]) + baseline_node.status = "completed" + adapter._store.upsert_node(baseline_node) + + # No goto key at all — .get("goto") returns None, which means + # "start from anchor". Parent should be baseline. + child_spec = adapter._create_child_spec( + mtp_node, + {"name": "some_exp"}, + ) + child_node = self._load_node(child_spec) + self.assertEqual("000_baseline", child_node.parent_id) + + def test_headspace_completion_selects_mtp_after_must_run_finish(self) -> None: + baseline = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="completed", + ) + headspace = HypothesisNode( + node_id="000_headspace", + name="headspace_cache", + status="completed", + spec={"_is_headspace": True}, + ) + mtp = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="completed", + ) + + selected = _select_planning_node( + headspace, + { + baseline.node_id: baseline, + headspace.node_id: headspace, + mtp.node_id: mtp, + }, + ) + + self.assertEqual(mtp, selected) + + def test_policy_helpers_cover_metric_and_thread_decisions(self) -> None: + self.assertEqual( + ("step_ms", -12.5), + _compare_metric_value({"steady_step_time_ms": 12.5, "duration_s": 100}), + ) + # Only --num-threads → use it directly. + self.assertEqual( + 12, + _extract_total_threads("--num-threads 12"), + ) + # Stage concurrency flags are not additive thread budgets. + self.assertIsNone( + _extract_total_threads("--num-fetch-threads 8 --num-decode-threads 16"), + ) + self.assertEqual( + 16, + _extract_default_executor_concurrency( + "--num-fetch-threads 8 --num-decode-threads 16" + ), + ) + self.assertEqual( + 8, + _extract_default_executor_concurrency("--num-fetch-threads 8"), + ) + self.assertEqual( + 16, + _extract_default_executor_concurrency("--num-decode-threads 16"), + ) + # No num_threads flag at all → None (unknown budget). + self.assertIsNone( + _extract_total_threads("--num_workers 8 --num_epochs 3"), + ) + self.assertIsNone( + _extract_total_threads(""), + ) + self.assertEqual( + ["ok", "fits_concurrency"], + [ + spec["name"] + for spec in _validate_thread_budget( + [ + {"name": "ok", "launch_command": "--num-threads 8"}, + {"name": "too_many", "launch_command": "--num-threads 64"}, + { + "name": "too_small", + "launch_command": ( + "--num-threads 8 --num-decode-threads 16" + ), + }, + { + "name": "fits_concurrency", + "launch_command": ( + "--num-threads 16 --num-decode-threads 16" + ), + }, + ], + 16, + ) + ], + ) + # Commands with no thread flags pass validation (unknown budget). + self.assertEqual( + ["pass_through"], + [ + spec["name"] + for spec in _validate_thread_budget( + [ + {"name": "pass_through", "launch_command": "--num_workers 8"}, + ], + 16, + ) + ], + ) + + # -- _extract_param_changes / _build_change_set / _is_duplicate_spec ------ + + def test_extract_param_changes_detects_flag_diffs(self) -> None: + base = "torchx run app --image $IMAGE --num_workers 8 --num_epochs 3" + + # New flag added. + self.assertEqual( + ["batch_size=48"], + _extract_param_changes(base + " --batch_size 48", base), + ) + + # Flag value changed. + self.assertEqual( + ["num_workers=16"], + _extract_param_changes( + "torchx run app --image $IMAGE --num_workers 16 --num_epochs 3", + base, + ), + ) + + # No diff when identical. + self.assertEqual([], _extract_param_changes(base, base)) + + # Empty commands return nothing. + self.assertEqual([], _extract_param_changes("", base)) + self.assertEqual([], _extract_param_changes(base, "")) + + def test_extract_param_changes_normalizes_dashes(self) -> None: + base = "torchx run app --num-fetch-threads 8" + exp = "torchx run app --num-fetch-threads 16" + changes = _extract_param_changes(exp, base) + self.assertEqual(["num_fetch_threads=16"], changes) + + def test_extract_param_changes_handles_negative_values(self) -> None: + base = "torchx run app --max_steps -1" + exp = "torchx run app --max_steps -2" + changes = _extract_param_changes(exp, base) + self.assertEqual(["max_steps=-2"], changes) + + def test_build_change_set_merges_explicit_and_param_changes(self) -> None: + base = "torchx run --image $IMAGE --num_workers 8" + spec = { + "changes": ["torch_compile"], + "launch_command": base + " --batch_size 48", + } + result = _build_change_set(spec, base) + self.assertEqual(frozenset({"torch_compile", "batch_size=48"}), result) + + def test_build_change_set_empty_for_baseline(self) -> None: + base = "torchx run --image $IMAGE --num_workers 8" + spec = {"changes": [], "launch_command": base} + self.assertEqual(frozenset(), _build_change_set(spec, base)) + + def test_build_change_set_normalizes_case(self) -> None: + spec = {"changes": ["Torch_Compile", " FUSED_ADAMW "]} + result = _build_change_set(spec, "") + self.assertEqual(frozenset({"torch_compile", "fused_adamw"}), result) + + def test_duplicate_requires_matching_change_sets(self) -> None: + base = "torchx run --image $IMAGE --num_workers 8" + baseline = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="completed", + spec={"changes": [], "launch_command": base}, + ) + mtp = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="completed", + spec={"changes": ["mtp"], "launch_command": base}, + ) + nodes = [baseline, mtp] + + # torch_compile: different code changes, same launch → NOT duplicate. + self.assertFalse( + _is_duplicate_spec( + {"changes": ["torch_compile"], "launch_command": base}, + nodes, + base, + ) + ) + + # fused_adamw: different code changes, same launch → NOT duplicate. + self.assertFalse( + _is_duplicate_spec( + {"changes": ["fused_adamw"], "launch_command": base}, + nodes, + base, + ) + ) + + # Exact same change set as baseline → IS duplicate. + self.assertTrue( + _is_duplicate_spec( + {"changes": [], "launch_command": base}, + nodes, + base, + ) + ) + + # Exact same change set as MTP → IS duplicate. + self.assertTrue( + _is_duplicate_spec( + {"changes": ["mtp"], "launch_command": base}, + nodes, + base, + ) + ) + + def test_duplicate_distinguishes_param_only_experiments(self) -> None: + base = "torchx run --image $IMAGE --num_workers 8" + batch_48 = HypothesisNode( + node_id="002_batch48", + name="batch_size_48", + status="completed", + spec={ + "changes": [], + "launch_command": base + " --batch_size 48", + }, + ) + nodes = [batch_48] + + # batch_size=64 has different param → NOT duplicate. + self.assertFalse( + _is_duplicate_spec( + {"changes": [], "launch_command": base + " --batch_size 64"}, + nodes, + base, + ) + ) + + # batch_size=48 again → IS duplicate. + self.assertTrue( + _is_duplicate_spec( + {"changes": [], "launch_command": base + " --batch_size 48"}, + nodes, + base, + ) + ) + + def test_duplicate_skips_failed_nodes(self) -> None: + base = "torchx run --image $IMAGE" + failed = HypothesisNode( + node_id="003_oom", + name="batch_64", + status="failed", + spec={ + "changes": [], + "launch_command": base + " --batch_size 64", + }, + ) + # Exact match of a failed node → NOT duplicate (allow retry). + self.assertFalse( + _is_duplicate_spec( + {"changes": [], "launch_command": base + " --batch_size 64"}, + [failed], + base, + ) + ) + + def test_duplicate_combination_vs_individual(self) -> None: + base = "torchx run --image $IMAGE" + mtp_only = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="completed", + spec={"changes": ["mtp"], "launch_command": base}, + ) + # Combination is not a dup of individual. + self.assertFalse( + _is_duplicate_spec( + { + "changes": ["mtp", "torch_compile"], + "launch_command": base, + }, + [mtp_only], + base, + ) + ) + + def test_duplicate_backward_compat_no_changes_field(self) -> None: + base = "torchx run --image $IMAGE --num_workers 8" + # Old spec without changes field — change set is derived purely from + # launch command diffs. + old_node = HypothesisNode( + node_id="002_batch48", + name="batch_size_48", + status="completed", + spec={"launch_command": base + " --batch_size 48"}, + ) + + # Same param diff → duplicate. + self.assertTrue( + _is_duplicate_spec( + {"launch_command": base + " --batch_size 48"}, + [old_node], + base, + ) + ) + + # Different param → not duplicate. + self.assertFalse( + _is_duplicate_spec( + {"launch_command": base + " --batch_size 64"}, + [old_node], + base, + ) + ) + + def test_startup_retry_inherits_changes(self) -> None: + node = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="failed", + spec={ + "name": "mtp", + "changes": ["mtp"], + "description": "try MTP", + "best_practices_tags": ["mtp"], + }, + failure=_make_failure( + FailureKind.JOB_STARTUP_FAILED, + FailurePhase.JOB, + "Tokenizer cannot pickle", + ), + ) + retry = _startup_retry_spec(node, _config()) + self.assertIsNotNone(retry) + assert retry is not None + self.assertEqual(["mtp"], retry["changes"]) + + def test_initial_nodes_have_changes(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + adapter = self._adapter(Path(tmp)) + specs = adapter.load() + nodes = [] + for spec in specs: + node_data = spec.payload["node"] + assert isinstance(node_data, dict) + nodes.append(HypothesisNode.from_dict(node_data)) + + by_name = {node.name: node for node in nodes} + self.assertEqual([], by_name["baseline"].spec["changes"]) + self.assertEqual( + ["cache_dataloader"], by_name["headspace_cache"].spec["changes"] + ) + self.assertEqual(["mtp"], by_name["mtp"].spec["changes"]) + + def test_store_update_spec_refreshes_checkpoint_and_active_view(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + self._write_base_files(workdir) + store = _WorkflowStateStore(workdir, _state()) + node = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="queued", + ) + spec = _spec_from_node(node) + store.save_scheduler_state(queued=[], running=[spec], status="running") + + node.status = "running" + node.job_id = "remote_job" + store.update_spec(spec, node) + + checkpoint = json.loads( + (workdir / "engine" / "checkpoint.json").read_text() + ) + active = json.loads((workdir / "engine" / "active.json").read_text()) + running_node = _node_from_spec(TaskSpec.from_dict(checkpoint["running"][0])) + + self.assertEqual("remote_job", running_node.job_id) + self.assertEqual("000_baseline", active[0]["node_id"]) + self.assertEqual("remote_job", active[0]["job_id"]) + self.assertIn("launched_at_iso", active[0]) + + def test_queue_view_includes_retry_lineage(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + self._write_base_files(workdir) + store = _WorkflowStateStore(workdir, _state()) + node = HypothesisNode( + node_id="002_retry", + name="retry", + status="queued", + spec={ + "_startup_retry_of": "001_mtp", + "_startup_retry_attempt": 1, + }, + ) + + store.save_scheduler_state( + queued=[_spec_from_node(node)], + running=[], + status="running", + ) + + queue = json.loads((workdir / "engine" / "queue.json").read_text()) + self.assertEqual("001_mtp", queue[0]["retry_of"]) + self.assertEqual(1, queue[0]["retry_attempt"]) + + def test_store_rejects_malformed_checkpoint(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + self._write_base_files(workdir) + engine_dir = workdir / "engine" + engine_dir.mkdir() + (engine_dir / "checkpoint.json").write_text( + json.dumps({"queued": [{"id": "bad", "payload": {}}]}) + "\n" + ) + store = _WorkflowStateStore(workdir, _state()) + + with self.assertRaisesRegex(ValueError, "payload.node"): + store.load_checkpoint() + + def test_store_persists_failure_view(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + self._write_base_files(workdir) + store = _WorkflowStateStore(workdir, _state()) + node = HypothesisNode( + node_id="001_failed", + name="failed", + status="failed", + failure=_make_failure( + FailureKind.BUILD_FAILED, + FailurePhase.BUILD, + "Build failed", + ), + ) + + store.upsert_node(node) + store.write_all() + + failure = json.loads( + ( + workdir / "engine" / "nodes" / "001_failed" / "failure.json" + ).read_text() + ) + engine_state = json.loads( + (workdir / "engine" / "engine_state.json").read_text() + ) + self.assertEqual("build_failed", failure["kind"]) + self.assertEqual({"build_failed": 1}, engine_state["failed_by_kind"]) + self.assertIn("build_failed", _failure_summary(workdir)) + self.assertIn("build_failed", _read_failures(workdir)) + + def test_adapter_records_structured_failure(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + adapter = self._adapter(workdir) + node = HypothesisNode(node_id="001_failed", name="failed") + spec = _spec_from_node(node) + failure = _make_failure( + FailureKind.LAUNCH_FAILED, + FailurePhase.LAUNCH, + "No launch command configured", + ) + + asyncio.run(adapter._record_failure(spec, node, failure)) + + stored = json.loads( + ( + workdir / "engine" / "nodes" / "001_failed" / "failure.json" + ).read_text() + ) + master_table = (workdir / "master_table.tsv").read_text() + self.assertEqual("launch_failed", stored["kind"]) + self.assertIn("launch_failed: No launch command configured", master_table) + + def test_completed_failure_history_uses_structured_failure(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + self._write_base_files(workdir) + state = _state() + node = HypothesisNode( + node_id="001_failed", + name="failed", + status="failed", + spec={"name": "failed"}, + ) + result = AnalysisResult( + structured={"metrics": {}, "findings": []}, + failure=_make_failure( + FailureKind.JOB_FAILED, + FailurePhase.JOB, + "Job failed", + ), + ) + + _update_on_complete(workdir, _config(), state, node, result) + + entry = state["history"][0] + self.assertNotIn("failure", entry) + self.assertEqual("job_failed", entry["structured"]["failure"]["kind"]) + + def test_job_failure_classifier_splits_startup_runtime_and_unknown(self) -> None: + startup = _classify_terminal_job_failure( + _MetricsEvidence( + system_metrics="", + pipeline_stats_log="TypeError: cannot pickle local function for MTP", + metrics_summary="", + ), + job_id="job_startup", + progress_seen=False, + ) + runtime = _classify_terminal_job_failure( + _MetricsEvidence( + system_metrics="", + pipeline_stats_log="[autoresearch] step=12\nCUDA out of memory", + metrics_summary="steady_step_time_ms: 12", + ), + job_id="job_runtime", + progress_seen=True, + ) + unknown = _classify_terminal_job_failure( + _MetricsEvidence( + system_metrics="", pipeline_stats_log="", metrics_summary="" + ), + job_id="job_unknown", + progress_seen=False, + ) + + self.assertEqual(FailureKind.JOB_STARTUP_FAILED, startup.kind) + self.assertEqual(FailureKind.JOB_RUNTIME_FAILED, runtime.kind) + self.assertEqual(FailureKind.JOB_FAILED, unknown.kind) + + def test_structured_metrics_evidence_guides_failure_classification(self) -> None: + startup = _classify_terminal_job_failure( + _MetricsEvidence( + system_metrics="", + pipeline_stats_log="", + metrics_summary="", + error_summary="TypeError: can't pickle tokenizer", + log_paths=["stderr.log"], + ), + job_id="job_startup", + progress_seen=False, + ) + runtime = _classify_terminal_job_failure( + _MetricsEvidence( + system_metrics="", + pipeline_stats_log="", + metrics_summary="", + progress_seen=True, + exit_code=1, + ), + job_id="job_runtime", + progress_seen=False, + ) + + self.assertEqual(FailureKind.JOB_STARTUP_FAILED, startup.kind) + self.assertEqual(["stderr.log"], startup.details["log_paths"]) + self.assertEqual(FailureKind.JOB_RUNTIME_FAILED, runtime.kind) + self.assertEqual(1, runtime.details["exit_code"]) + + def test_startup_failure_retry_is_bounded_for_mtp(self) -> None: + node = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="failed", + spec={ + "name": "mtp", + "description": "try MTP", + "best_practices_tags": ["mtp"], + }, + failure=_make_failure( + FailureKind.JOB_STARTUP_FAILED, + FailurePhase.JOB, + "Tokenizer cannot pickle", + ), + ) + + retry = _startup_retry_spec(node, _config()) + self.assertIsNotNone(retry) + assert retry is not None + self.assertEqual("mtp_startup_retry_1", retry["name"]) + self.assertEqual(1, retry["_startup_retry_attempt"]) + self.assertIn("pickling", retry["hypothesis"]) + + node.spec["_startup_retry_attempt"] = 2 + self.assertIsNone(_startup_retry_spec(node, _config())) + + def test_retry_policy_is_kind_specific(self) -> None: + node = HypothesisNode( + node_id="001_mtp", + name="mtp", + spec={"best_practices_tags": ["mtp"]}, + failure=_make_failure( + FailureKind.JOB_STARTUP_FAILED, + FailurePhase.JOB, + "startup", + ), + ) + policy = _retry_policy_for_failure(node, _config()) + self.assertIsNotNone(policy) + assert policy is not None + self.assertEqual(2, policy["max_attempts"]) + + node.failure = _make_failure( + FailureKind.BUILD_FAILED, + FailurePhase.BUILD, + "build", + ) + self.assertIsNone(_retry_policy_for_failure(node, _config())) + + def test_planning_prefers_non_startup_failed_retry(self) -> None: + baseline = HypothesisNode( + node_id="000_baseline", + name="baseline", + status="completed", + ) + headspace = HypothesisNode( + node_id="000_headspace", + name="headspace_cache", + status="completed", + spec={"_is_headspace": True}, + ) + mtp = HypothesisNode( + node_id="001_mtp", + name="mtp", + status="failed", + failure=_make_failure( + FailureKind.JOB_STARTUP_FAILED, + FailurePhase.JOB, + "startup", + ), + ) + retry = HypothesisNode( + node_id="002_mtp_startup_retry_1", + name="mtp_startup_retry_1", + status="failed", + spec={"_startup_retry_of": "001_mtp"}, + failure=_make_failure( + FailureKind.JOB_STARTUP_FAILED, + FailurePhase.JOB, + "startup", + ), + ) + better_candidate = HypothesisNode( + node_id="003_threads", + name="threads", + status="completed", + ) + + selected = _select_planning_node( + retry, + { + node.node_id: node + for node in [baseline, headspace, mtp, retry, better_candidate] + }, + ) + + self.assertEqual(better_candidate, selected) + + def test_startup_retry_uses_repair_prompt(self) -> None: + platform = create_platform({"platform": "local", "agent": "mock"}) + assert isinstance(platform.agent, _MockAgent) + platform.agent.responses["prompt:apply_startup_repair"] = ( + "failed during job startup __STARTUP_FAILURE_JSON__" + ) + prompt = _build_apply_prompt( + platform, + { + "name": "mtp_startup_retry_1", + "description": "repair", + "hypothesis": "fix startup", + "_startup_retry_attempt": 1, + "_startup_failure": {"kind": "job_startup_failed"}, + }, + "002_mtp_startup_retry_1", + "knowledge", + "/tmp/pipeline.py", + "def main():\n pass\n", + ) + + self.assertIn("failed during job startup", prompt) + self.assertIn("job_startup_failed", prompt) + + def test_headspace_node_uses_dedicated_prompt_with_knowledge(self) -> None: + platform = create_platform({"platform": "local", "agent": "mock"}) + assert isinstance(platform.agent, _MockAgent) + platform.agent.responses["prompt:headspace"] = ( + "headspace prompt __KNOWLEDGE__ __PIPELINE_CODE__" + ) + + prompt = _build_apply_prompt( + platform, + { + "name": "headspace_cache", + "description": "Wrap with CacheDataLoader for headspace analysis", + "_is_headspace": True, + }, + "000_headspace", + "knowledge", + "/tmp/pipeline.py", + "def main():\n pass\n", + ) + + self.assertIn("headspace prompt", prompt) + self.assertIn("knowledge", prompt) + self.assertIn("def main", prompt) + + def test_change_summary_is_concise_and_persisted_in_master_table(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + (workdir / "master_table.tsv").write_text( + "\t".join(MASTER_TABLE_HEADERS) + "\n" + ) + + summary = _change_summary_for_spec( + { + "description": ( + "Increase decode thread count while preserving the rest " + "of the pipeline" + ) + } + ) + _append_master_row( + workdir, + { + "run_id": "001_threads", + "name": "threads", + "status": "completed", + "change_summary": summary, + "sm_util_pct": "50", + }, + MASTER_TABLE_HEADERS, + ) + + rows = _load_tsv(workdir / "master_table.tsv") + + self.assertEqual("Increase decode thread count while", summary) + self.assertEqual(summary, rows[0]["change_summary"]) + + def test_every_failure_kind_has_policy(self) -> None: + self.assertEqual(set(FailureKind.__members__.values()), set(_FAILURE_POLICIES)) + self.assertTrue(_FAILURE_POLICIES[FailureKind.JOB_STARTUP_FAILED].retryable) + + def test_error_summary_prefers_traceback_block(self) -> None: + summary = _summarize_error( + "before\n" + "Traceback (most recent call last):\n" + ' File "x.py", line 1, in \n' + "TypeError: cannot pickle tokenizer\n" + "after\n" + ) + + self.assertIsNotNone(summary) + assert summary is not None + self.assertIn("Traceback", summary) + self.assertIn("cannot pickle", summary) + + def test_atomic_write_replaces_json(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "data.json" + + _write_text_atomic(path, '{"a": 1}\n') + _write_text_atomic(path, '{"a": 2}\n') + + self.assertEqual({"a": 2}, json.loads(path.read_text())) + + def _adapter(self, workdir: Path) -> PipelineOptimizationWorkflow: + self._write_base_files(workdir) + return PipelineOptimizationWorkflow( + workdir=workdir, + config=_config(), + state=_state(), + platform=create_platform({"platform": "auto", "agent": "mock"}, workdir), + ) + + def _write_base_files(self, workdir: Path) -> None: + workdir.mkdir(parents=True, exist_ok=True) + (workdir / "config.json").write_text(json.dumps(_config()) + "\n") + write_state(workdir, _state()) + (workdir / "master_table.tsv").write_text( + "\t".join(MASTER_TABLE_HEADERS) + "\n" + ) diff --git a/src/spdl/dataloader/tests/cache_dataloader_test.py b/src/spdl/dataloader/tests/cache_dataloader_test.py new file mode 100644 index 000000000..6893b33e8 --- /dev/null +++ b/src/spdl/dataloader/tests/cache_dataloader_test.py @@ -0,0 +1,108 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from collections.abc import Iterator + +from spdl.dataloader import CacheDataLoader +from spdl.pipeline import cache_iterator + + +class TestCacheIterator(unittest.TestCase): + def test_cache_iterator(self) -> None: + """cache_iterator returns the cached values""" + + ite = iter(cache_iterator(range(5), 3)) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + def test_cache_iterator_cache_return_after(self) -> None: + """cache_iterator returns the cached values""" + + ite = iter(cache_iterator(range(7), 3, return_caches_after=5)) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + self.assertEqual(next(ite), 3) + self.assertEqual(next(ite), 4) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + def test_cache_iterator_cache_return_after_len(self) -> None: + """cache_iterator returns the cached values""" + + ite = iter(cache_iterator(range(7), 3, return_caches_after=5, stop_after=10)) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + self.assertEqual(next(ite), 3) + self.assertEqual(next(ite), 4) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + + with self.assertRaises(StopIteration): + next(ite) + + +class TestCacheDataLoader(unittest.TestCase): + def test_CacheDataLoader(self) -> None: + """Smoke test""" + + class DL: + def __init__(self, n: int) -> None: + self.n = n + + def __iter__(self) -> Iterator[int]: + yield from range(self.n) + + def __len__(self) -> int: + return self.n + + N = 8 + dl = CacheDataLoader(DL(N), num_caches=2, return_caches_after=3, stop_after=N) + + self.assertEqual(dl.n, N) + self.assertEqual(len(dl), N) + + ite = iter(dl) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + + self.assertEqual(next(ite), 0) + + with self.assertRaises(StopIteration): + next(ite) diff --git a/src/spdl/dataloader/tests/dataloader_test.py b/src/spdl/dataloader/tests/dataloader_test.py new file mode 100644 index 000000000..afd5c807b --- /dev/null +++ b/src/spdl/dataloader/tests/dataloader_test.py @@ -0,0 +1,167 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import os +import platform +import time +import unittest + +from spdl.dataloader import DataLoader + + +def get_dl(*args, timeout=3, num_threads=2, **kwargs): + # on default values + # timeout -> so that test would fail rather stack + # num_threads -> keep it minimum but have more than 1 + return DataLoader(*args, **kwargs, num_threads=num_threads, timeout=timeout) + + +class TestDataLoader(unittest.TestCase): + def test_dataloader_iterable(self) -> None: + src = list(range(10)) + + dl = get_dl(src) + + self.assertEqual(sorted(dl), src) + + def test_dataloader_stateful_iterable(self) -> None: + class src: + def __init__(self, num_items: int = 10): + self.num_iter = 0 + self.num_items = num_items + + def __iter__(self): + for i in range(self.num_items): + yield (self.num_iter, i) + self.num_iter += 1 + + dl = get_dl(src()) + + self.assertEqual(sorted(dl), [(0, i) for i in range(10)]) + self.assertEqual(sorted(dl), [(1, i) for i in range(10)]) + self.assertEqual(sorted(dl), [(2, i) for i in range(10)]) + + def test_dataloader_preprocess(self) -> None: + """preprocessor process the value of the source""" + src = list(range(10)) + + def double(x): + time.sleep(0.05 * x) # to reduce flakiness from multi-threading + return 2 * x + + dl = get_dl(src, preprocessor=double) + + self.assertEqual(sorted(dl), [i * 2 for i in range(10)]) + + def test_dataloader_preprocess_in_order(self) -> None: + """When output_order='input', the order must be preserved.""" + src = list(range(10, -1, -1)) + + def delay(x): + time.sleep(0.1 * x) + return x + + dl = get_dl(src, preprocessor=delay, output_order="input") + + self.assertEqual(list(dl), src) + + dl = get_dl(src, preprocessor=delay, output_order="completion") + + self.assertNotEqual(list(dl), src) + + @unittest.skipIf( + platform.system() == "Darwin" and "CI" in os.environ, + "GitHub macOS CI is not timely enough.", + ) + def test_dataloader_buffer_size(self) -> None: + """Bigger buffer_size allows the BG to proceed while FG is not fetching the data""" + src = list(range(12)) + + def delay(x): + time.sleep(0.05) + return x + + def test(dl): + # Kick off the background thread + dli = iter(dl) + self.assertEqual(next(dli), 0) + + # Wait: (simulate foreground load) + time.sleep(1) + + # Iterate the rest + t0 = time.monotonic() + result = list(dli) + elapsed = time.monotonic() - t0 + print(elapsed) + self.assertEqual(result, src[1:]) + return elapsed + + # With buffer_size == 1, then the background thread cannot proceed + # while foreground thread does not fetch any. + dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=1) + elapsed = test(dl) + self.assertGreater(elapsed, 0.3) + + # With bigger buffer_size, the background thread proceed + # while foreground thread does not fetch any. + dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=len(src)) + elapsed = test(dl) + self.assertLess(elapsed, 0.15) + + def test_dataloader_num_threads(self) -> None: + """Increasing the num_threads reduces the overall time.""" + src = list(range(10)) + + def delay(x): + time.sleep(0.1) + return x + + def test(dl): + t0 = time.monotonic() + result = list(dl) + elapsed = time.monotonic() - t0 + print(elapsed) + self.assertEqual(sorted(result), src) + return elapsed + + dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=1) + self.assertGreater(test(dl), 0.8) + + dl = get_dl(src, preprocessor=delay, num_threads=len(src), buffer_size=1) + self.assertLess(test(dl), 0.6) + + def test_dataloader_batch(self) -> None: + """batching works with or without dropping""" + src = list(range(10)) + + dl = get_dl(src, batch_size=3, drop_last=False) + + self.assertEqual(list(dl), [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) + + dl = get_dl(src, batch_size=3, drop_last=True) + + self.assertEqual(list(dl), [[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + + def test_dataloader_aggregate(self) -> None: + """Aggregator processes the batched input""" + src = list(range(10)) + + def agg(vals: list[int]) -> tuple[int, int, int, int]: + return len(vals), min(vals), max(vals), sum(vals) + + dl = get_dl(src, batch_size=3, drop_last=False, aggregator=agg) + + expected = [ + (3, 0, 2, 3), + (3, 3, 5, 12), + (3, 6, 8, 21), + (1, 9, 9, 9), + ] + + self.assertEqual(list(dl), expected) diff --git a/src/spdl/dataloader/tests/iterator_test.py b/src/spdl/dataloader/tests/iterator_test.py new file mode 100644 index 000000000..2c9f6d09e --- /dev/null +++ b/src/spdl/dataloader/tests/iterator_test.py @@ -0,0 +1,370 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import functools +import pickle +import random +import unittest +import warnings +from collections.abc import Iterator +from functools import partial +from unittest.mock import patch + +from spdl.pipeline import iterate_in_subprocess as _iterate_in_subprocess +from spdl.source.utils import ( + embed_shuffle, + IterableWithShuffle, + MergeIterator, + repeat_source, +) + + +def _ignore_fork_warning(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + return wrapper + + +def iterate_in_subprocess(fn, *, timeout=10, **kwargs): + return _iterate_in_subprocess(fn, timeout=timeout, **kwargs) + + +class TestMergeIterator(unittest.TestCase): + def test_mergeiterator_ordered(self) -> None: + """MergeIterator iterates multiple iterators""" + + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables)) + self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2, 12, 22]) + + def test_mergeiterator_ordered_stop_after_first_exhaustion(self) -> None: + """MergeIterator stops after the first exhaustion""" + + iterables = [ + [0], + [10, 11, 12], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables, stop_after=-1)) + self.assertEqual(result, [0, 10, 20]) + + iterables = [ + [0, 1, 2], + [10], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables, stop_after=-1)) + self.assertEqual(result, [0, 10, 20, 1]) + + iterables = [ + [0, 1, 2], + [10, 11], + [20], + ] + + result = list(MergeIterator(iterables, stop_after=-1)) + self.assertEqual(result, [0, 10, 20, 1, 11]) + + def test_mergeiterator_ordered_stop_after_N(self) -> None: + """MergeIterator stops after N items are yielded""" + + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables, stop_after=1)) + self.assertEqual(result, [0]) + + result = list(MergeIterator(iterables, stop_after=5)) + self.assertEqual(result, [0, 10, 20, 1, 11]) + + result = list(MergeIterator(iterables, stop_after=7)) + self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2]) + + def test_mergeiterator_ordered_stop_after_minus1(self) -> None: + """MergeIterator stops after all the iterables are exhausted""" + + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables)) + self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2, 12, 22]) + + iterables = [ + [0, 1, 2], + [10], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables)) + self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) + + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20], + ] + + result = list(MergeIterator(iterables)) + self.assertEqual(result, [0, 10, 20, 1, 11, 2, 12]) + + def test_mergeiterator_ordered_n(self) -> None: + """with stop_after=N, MergeIterator continues iterating after exhaustion.""" + iterables = [ + [0, 1, 2], + [10], + [20, 21, 22], + ] + + result = list(MergeIterator(iterables, stop_after=5)) + self.assertEqual(result, [0, 10, 20, 1, 21]) + + result = list(MergeIterator(iterables, stop_after=7)) + self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) + + result = list(MergeIterator(iterables, stop_after=8)) + self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) + + def test_mergeiterator_stochastic_smoke_test(self) -> None: + """MergeIterator with probabilitiies do not get stuck.""" + + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20, 21, 22], + ] + + weights = [1, 1, 1] + + result = list(MergeIterator(iterables, weights=weights)) + self.assertEqual(set(result), {0, 1, 2, 10, 11, 12, 20, 21, 22}) + + def test_mergeiterator_stochastic_rejects_zero(self) -> None: + """weight=0 is rejected.""" + weights = [1, 0] + + with self.assertRaises(ValueError): + MergeIterator([[1]], weights=weights) + + weights = [1, 0.0] + + with self.assertRaises(ValueError): + MergeIterator([[1]], weights=weights) + + def test_mergeiterator_skip_zero_weight(self) -> None: + """Iterables with zero weight are skipped.""" + iterables = [ + [0, 1, 2], + [10, 11, 12], + [20, 21, 22], + [30, 31, 32], + ] + + weights = [1, 0, 2, 0] + + merge_iter = MergeIterator(iterables, weights=weights) + + self.assertEqual(len(merge_iter.iterables), 2) + self.assertEqual(merge_iter.iterables[0], [0, 1, 2]) + self.assertEqual(merge_iter.iterables[1], [20, 21, 22]) + + self.assertIsNotNone(merge_iter.weights) + # pyre-ignore[16]: weights is not None after assertion + self.assertEqual(len(merge_iter.weights), 2) + # pyre-ignore[16]: weights is not None after assertion + self.assertEqual(merge_iter.weights[0], 1) + # pyre-ignore[16]: weights is not None after assertion + self.assertEqual(merge_iter.weights[1], 2) + + result = list(merge_iter) + self.assertEqual(set(result), {0, 1, 2, 20, 21, 22}) + + def test_mergeiterator_stochastic_stop_after_N(self) -> None: + """Values are taken from iterables with higher weights""" + weights = [1000000, 1] + + iterables = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], + ] + + result = list(MergeIterator(iterables, weights=weights, stop_after=3)) + self.assertEqual(result, [0, 1, 2]) + + def test_mergeiterator_stochastic_stop_after_first_exhaustion(self) -> None: + """Values are taken from iterables with higher weights""" + weights = [1000000, 1] + + iterables = [ + [0, 1, 2, 3], + [10, 11, 12, 13], + ] + + result = list(MergeIterator(iterables, weights=weights, stop_after=-1)) + self.assertEqual(result, [0, 1, 2, 3]) + + +class TestRepeatSource(unittest.TestCase): + def test_repeat_source_iterable_with_shuffle(self) -> None: + """repeat_source repeats source while calling shuffle""" + + class _IteWithShuffle: + def __init__(self) -> None: + self.vals = list(range(3)) + + def shuffle(self, seed: int) -> None: + assert isinstance(seed, int) + self.vals = self.vals[1:] + self.vals[:1] + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + src = _IteWithShuffle() + gen = iter(repeat_source(src, epoch=2)) + + with patch.object(src, "shuffle", side_effect=src.shuffle) as mock_method: + self.assertEqual(next(gen), 1) + mock_method.assert_called_with(seed=0) + self.assertEqual(next(gen), 2) + self.assertEqual(next(gen), 0) + + self.assertEqual(next(gen), 2) + mock_method.assert_called_with(seed=1) + self.assertEqual(next(gen), 0) + self.assertEqual(next(gen), 1) + + self.assertEqual(next(gen), 0) + mock_method.assert_called_with(seed=2) + self.assertEqual(next(gen), 1) + self.assertEqual(next(gen), 2) + + self.assertEqual(next(gen), 1) + mock_method.assert_called_with(seed=3) + self.assertEqual(next(gen), 2) + self.assertEqual(next(gen), 0) + + self.assertEqual(next(gen), 2) + mock_method.assert_called_with(seed=4) + self.assertEqual(next(gen), 0) + self.assertEqual(next(gen), 1) + + def test_repeat_source_iterable(self) -> None: + """repeat_source works Iterable without shuffle method""" + + class _IteWithoutShuffle: + def __init__(self) -> None: + self.vals = list(range(3)) + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + src = _IteWithoutShuffle() + gen = iter(repeat_source(src, epoch=2)) + + for _ in range(100): + self.assertEqual(next(gen), 0) + self.assertEqual(next(gen), 1) + self.assertEqual(next(gen), 2) + + def test_repeat_source_picklable(self) -> None: + """repeat_source is picklable.""" + + src = list(range(10)) + src = repeat_source(src) + + serialized = pickle.dumps(src) + src2 = pickle.loads(serialized) + + for _ in range(3): + for i in range(10): + self.assertEqual(next(src), i) + self.assertEqual(next(src2), i) + + +class IterableWithShuffleSource: + def __init__(self, n: int) -> None: + self.vals = list(range(n)) + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + def shuffle(self, seed: int) -> None: + random.seed(seed) + random.shuffle(self.vals) + + +class SourceIterableWithShuffle(IterableWithShuffle[int]): + def __init__(self, n: int) -> None: + self.i = 0 + self.vals = list(range(n)) + + def shuffle(self, seed: int) -> None: + assert isinstance(seed, int) + self.vals = self.vals[1:] + self.vals[:1] + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + +class TestShuffleAndIterate(unittest.TestCase): + def test_shuffle_and_iterate_picklable(self) -> None: + """The result of embed_shuffle must be pickable (for multiprocessing)""" + + src = embed_shuffle(IterableWithShuffleSource(10)) + state = pickle.dumps(src) + src2 = pickle.loads(state) + + # pyre-ignore[16]: embed_shuffle returns an object with src attribute + self.assertEqual(src.src.vals, src2.src.vals) + + def test_shuffle_and_iterate(self) -> None: + N = 10 + + src = embed_shuffle(IterableWithShuffleSource(N)) + + ref = list(range(N)) + for i in range(3): + random.seed(i) + random.shuffle(ref) + + hyp = list(src) + self.assertEqual(hyp, ref) + + @_ignore_fork_warning + def test_move_iterable_to_subprocess_success_iterable_with_shuffle(self) -> None: + """IterableWithShuffle can be executed in the subprocess.""" + iterator = iterate_in_subprocess( + partial(embed_shuffle, SourceIterableWithShuffle(3)) + ) + + self.assertEqual(list(iterator), [1, 2, 0]) + self.assertEqual(list(iterator), [2, 0, 1]) + self.assertEqual(list(iterator), [0, 1, 2]) diff --git a/src/spdl/dataloader/tests/sampler_test.py b/src/spdl/dataloader/tests/sampler_test.py new file mode 100644 index 000000000..a59b98921 --- /dev/null +++ b/src/spdl/dataloader/tests/sampler_test.py @@ -0,0 +1,343 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import functools +import unittest +import warnings +from collections import Counter +from collections.abc import Callable +from functools import partial +from typing import TypeVar + +import numpy as np +from parameterized import parameterized +from spdl.pipeline import iterate_in_subprocess +from spdl.source import ( + DistributedDeterministicSampler, + DistributedRandomSampler, + SizedIterable, + SizedIterableWithShuffle, +) +from spdl.source.utils import embed_shuffle + +_F = TypeVar("_F", bound=Callable[..., object]) + + +def _ignore_fork_warning(fn: _F) -> _F: + @functools.wraps(fn) + def wrapper(*args: object, **kwargs: object) -> object: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + # pyre-ignore[7] + return wrapper + + +class TestDistributedSamplerInterface(unittest.TestCase): + def test_distributed_sampler_interface(self) -> None: + """samplers conform to Iterable/IterableWithShuffle protocol""" + self.assertIsInstance( + DistributedRandomSampler(9, rank=0, world_size=1), SizedIterableWithShuffle + ) + self.assertIsInstance( + DistributedDeterministicSampler(9, rank=0, world_size=1), SizedIterable + ) + + +class TestDistributedSamplerDeterministic(unittest.TestCase): + def test_deterministic_iter(self) -> None: + """without distributed, deterministic iteration behaves same as `range(N)`""" + N = 30 + sampler = DistributedDeterministicSampler(N, rank=0, world_size=1) + self.assertEqual(len(sampler), N) + self.assertEqual(list(sampler), list(range(N))) + + def test_deterministic_iter_distributed(self) -> None: + """deterministic iteration behaves same as `range(rank, M, world_size)`""" + N = 26 + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"The size of dataset \(\d+\) is not divisible by the " + r"world size \(\d+\)\. Some samples are never visited\." + ), + category=UserWarning, + ) + for world_size in range(1, N + 1): + len_ = N // world_size + max_ = len_ * world_size + c = Counter() + for rank in range(world_size): + print(f"{N=}, {world_size=}, {rank=}, {len_=}, {max_=}") + sampler = DistributedDeterministicSampler( + N, rank=rank, world_size=world_size + ) + self.assertEqual(len(sampler), len_) + + indices = list(sampler) + self.assertEqual(indices, list(range(rank, max_, world_size))) + c.update(indices) + + # Check that together, the samplers covered the whole dataset + num_iters = N // world_size * world_size + self.assertEqual(c.total(), num_iters) + self.assertEqual(len(c.keys()), num_iters) + self.assertEqual(set(c.keys()), set(range(num_iters))) + self.assertTrue(all(v == 1 for v in c.values())) + + def test_deterministic_iter_stable_across_epochs(self) -> None: + """Deterministic sampler produces the same sequence on every iteration.""" + N = 30 + world_size = 4 + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"The size of dataset \(\d+\) is not divisible by the " + r"world size \(\d+\)\. Some samples are never visited\." + ), + category=UserWarning, + ) + for rank in range(world_size): + sampler = DistributedDeterministicSampler( + N, rank=rank, world_size=world_size + ) + first_epoch = list(sampler) + for _ in range(5): + self.assertEqual(list(sampler), first_epoch) + + +class TestDistributedSamplerRandom(unittest.TestCase): + def test_shuffle(self) -> None: + """shuffling makes sampler generates different indices.""" + N = 640 + rank = 3 + world_size = 8 + + previous: list[int] = [] + for epoch in range(100): + sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) + sampler.shuffle(seed=epoch) + + indices = list(sampler) + print(f"{indices=}") + self.assertNotEqual(indices, previous) + previous = indices + + def test_shuffle_epoch_loop(self) -> None: + """Calling shuffle(seed=epoch) on a single sampler produces different sequences each epoch.""" + N = 640 + world_size = 8 + + for rank in range(world_size): + sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) + previous: list[int] = [] + for epoch in range(10): + sampler.shuffle(seed=epoch) + indices = list(sampler) + self.assertEqual(len(indices), N // world_size) + self.assertNotEqual(indices, previous) + previous = indices + + def test_shuffle_is_stateless(self) -> None: + """shuffle(seed) output depends only on the seed, not on prior iteration history.""" + N = 640 + world_size = 8 + + for rank in range(world_size): + sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) + for i in range(10): + sampler.shuffle(seed=i) + list(sampler) + sampler.shuffle(seed=5) + after_many = list(sampler) + + fresh = DistributedRandomSampler(N, rank=rank, world_size=world_size) + fresh.shuffle(seed=5) + from_fresh = list(fresh) + + self.assertEqual(after_many, from_fresh) + + def test_shuffle_epoch_loop_mutual_exclusive(self) -> None: + """All ranks together cover the full dataset at each epoch when using shuffle(seed=epoch).""" + N = 640 + world_size = 8 + + samplers = [ + DistributedRandomSampler(N, rank=rank, world_size=world_size) + for rank in range(world_size) + ] + + for epoch in range(10): + c = Counter() + for sampler in samplers: + sampler.shuffle(seed=epoch) + c.update(sampler) + + self.assertEqual(c.total(), N) + self.assertEqual(len(c.keys()), N) + self.assertEqual(set(c.keys()), set(range(N))) + self.assertTrue(all(v == 1 for v in c.values())) + + @parameterized.expand( + [ + (None,), + (1,), + ] + ) + def test_repeat(self, w: int | None) -> None: + """Without calling shuffle, sampler generates the same sequence.""" + N = 40 + world_size = 8 + + weights = None if w is None else [1.0] * N + for rank in range(world_size): + previous = [] + for i in range(100): + sampler = DistributedRandomSampler( + N, rank=rank, world_size=world_size, weights=weights + ) + + indices = list(sampler) + print(f"{indices=}") + if i > 0: + self.assertEqual(indices, previous) + previous = indices + + @parameterized.expand( + [ + (True,), + (False,), + ] + ) + def test_mutual_exclusive(self, shuffle: bool) -> None: + """Without weights, samplers generate mutually exclusive sets""" + N = 640 + world_size = 8 + + for epoch in range(100): + c = Counter() + for rank in range(world_size): + sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) + if shuffle: + sampler.shuffle(seed=epoch) + c.update(sampler) + + self.assertEqual(c.total(), N) + self.assertEqual(len(c.keys()), N) + self.assertEqual(set(c.keys()), set(range(N))) + self.assertTrue(all(v == 1 for v in c.values())) + + @parameterized.expand( + [ + (True,), + (False,), + ] + ) + def test_mutual_exclusive_num_draws(self, shuffle: bool) -> None: + """Without weights, samplers generate mutually exclusive sets""" + N = 640 + num_draws = 321 + world_size = 8 + + for epoch in range(100): + c = Counter() + for rank in range(world_size): + sampler = DistributedRandomSampler( + N, rank=rank, world_size=world_size, num_draws=num_draws + ) + if shuffle: + sampler.shuffle(seed=epoch) + c.update(sampler) + + m = num_draws // world_size * world_size + self.assertEqual(c.total(), m) + self.assertEqual(len(c.keys()), m) + self.assertTrue(all(v == 1 for v in c.values())) + + +class TestDistributedSamplerWeighted(unittest.TestCase): + def test_weighted_sampling(self) -> None: + """Indices are drawn according to the weights""" + weights = [0.0, 1.0, 3.0, 5.0, 10.0] + N = len(weights) + + sampler = DistributedRandomSampler( + N, rank=0, world_size=1, num_draws=1_000_000, weights=weights + ) + + c = Counter(sampler) + distribution = [c[i] for i in range(N)] + + print(f"{weights=}") + print(f"{distribution=}") + + ref = np.asarray(weights) / np.sum(weights) + hyp = np.asarray(distribution) / np.sum(distribution) + + print(f"{ref=}") + print(f"{hyp=}") + + self.assertTrue(np.allclose(hyp, ref, atol=1e-3)) + + +class TestDistributedSamplerEmbedShuffle(unittest.TestCase): + def test_embed_shuffle(self) -> None: + """DistributedSampler is compatibile with embed_shuffle""" + N = 10 + weights = [1.0 for _ in range(N)] + + s0 = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) + s1 = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) + + s1 = embed_shuffle(s1) + + previous = [] + for i in range(100): + hyp = list(s1) + print(f"{hyp=}") + + s0.shuffle(i) + ref = list(s0) + print(f"{ref=}") + + self.assertEqual(hyp, ref) + self.assertNotEqual(hyp, previous) + previous = hyp + + +class TestDistributedSamplerIterateInSubprocess(unittest.TestCase): + @_ignore_fork_warning + def test_iterate_in_subprocess(self) -> None: + """Iterating in a subprocess generates identical result""" + N = 10 + weights = [1.0 for _ in range(N)] + + sampler = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) + sampler_sub = iterate_in_subprocess(partial(embed_shuffle, sampler)) + sampler = embed_shuffle(sampler) + + previous = [] + for _ in range(100): + hyp = list(sampler_sub) + print(f"{hyp=}") + ref = list(sampler) + print(f"{ref=}") + + self.assertEqual(hyp, ref) + self.assertNotEqual(hyp, previous) + previous = hyp diff --git a/src/spdl/dataloader/tests/source_test.py b/src/spdl/dataloader/tests/source_test.py new file mode 100644 index 000000000..61f31e8a5 --- /dev/null +++ b/src/spdl/dataloader/tests/source_test.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import unittest +from collections.abc import Iterable +from pathlib import Path +from tempfile import TemporaryDirectory + +from spdl.source.imagenet import ImageNet +from spdl.source.local_directory import LocalDirectory + + +def _make_files(paths: Iterable[Path]) -> None: + for path in paths: + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + + +class SourceTest(unittest.TestCase): + def test_LocalDirectory(self) -> None: + """LocalDirectory can traverse specified files""" + with TemporaryDirectory() as root_dir: + root_dir = Path(root_dir) + + targets = { + root_dir / "foo.txt", + root_dir / "dir" / "bar.txt", + root_dir / "dir" / "dir" / "bazz.txt", + } + others = { + root_dir / "foo.dat", + root_dir / "dir" / "bar.dat", + root_dir / "dir" / "dir" / "bazz.dat", + } + _make_files(targets | others) + + src = LocalDirectory(root=root_dir, pattern="**/*.txt") + + vals1 = list(src) + vals2 = list(src) + src.shuffle(seed=0) + vals3 = list(src) + + self.assertEqual(vals1, vals2) + self.assertNotEqual(vals1, vals3) + self.assertEqual(set(vals1), targets) + self.assertEqual(set(vals2), targets) + self.assertEqual(set(vals3), targets) + + def test_ImageNet(self) -> None: + """ImageNet returns image path and class ID""" + with TemporaryDirectory() as root_dir: + root_dir = Path(root_dir) + + vals = { + (root_dir / "val" / "n02110958" / "FOO.JPEG", 254), + (root_dir / "val" / "n02027492" / "FOO.JPEG", 140), + (root_dir / "val" / "n02071294" / "FOO.JPEG", 148), + (root_dir / "val" / "n02088632" / "FOO.JPEG", 164), + } + trains = { + (root_dir / "train" / "n02066245" / "FOO.JPEG", 147), + (root_dir / "train" / "n02277742" / "FOO.JPEG", 322), + (root_dir / "train" / "n02965783" / "FOO.JPEG", 475), + (root_dir / "train" / "n03240683" / "FOO.JPEG", 540), + } + _make_files([v for v, _ in vals]) + _make_files([v for v, _ in trains]) + + src = ImageNet(root=root_dir, split="val") + v1 = list(src) + src.shuffle(0) + v2 = list(src) + self.assertNotEqual(v1, v2) + self.assertEqual(set(v1), vals) + self.assertEqual(set(v2), vals) + + src = ImageNet(root=root_dir, split="train") + v1 = list(src) + src.shuffle(0) + v2 = list(src) + self.assertNotEqual(v1, v2) + self.assertEqual(set(v1), trains) + self.assertEqual(set(v2), trains) diff --git a/src/spdl/dataloader/tests/source_utils_test.py b/src/spdl/dataloader/tests/source_utils_test.py new file mode 100644 index 000000000..c62d51438 --- /dev/null +++ b/src/spdl/dataloader/tests/source_utils_test.py @@ -0,0 +1,92 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import unittest +from collections.abc import Iterator + +from spdl.source.utils import embed_shuffle + + +class IterableWithShuffle_: + def __init__(self, n: int) -> None: + self.vals: list[int] = list(range(n)) + self._seed: int | None = None + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + def shuffle(self, seed: int) -> None: + # rotate + self._seed = seed + self.vals = self.vals[1:] + self.vals[:1] + + +class SourceUtilsTest(unittest.TestCase): + def test_embed_shuffle(self) -> None: + """Iterable created by embed_shuffle calls shuffle automatically""" + + foo = IterableWithShuffle_(3) + self.assertIsNone(foo._seed) + iterable = embed_shuffle(foo) + self.assertEqual(list(iterable), [1, 2, 0]) + self.assertEqual(foo._seed, 0) + self.assertEqual(list(iterable), [2, 0, 1]) + self.assertEqual(foo._seed, 1) + self.assertEqual(list(iterable), [0, 1, 2]) + self.assertEqual(foo._seed, 2) + + def test_embed_shuffle_halt(self) -> None: + """The value is shuffled with different seed even after an iteration is halted.""" + + foo = IterableWithShuffle_(5) + iterable = embed_shuffle(foo) + + iterator = iter(iterable) + self.assertIsNone(foo._seed) + self.assertEqual(next(iterator), 1) + self.assertEqual(foo._seed, 0) + self.assertEqual(next(iterator), 2) + del iterator + + iterator = iter(iterable) + self.assertEqual(next(iterator), 2) + self.assertEqual(foo._seed, 1) + self.assertEqual(next(iterator), 3) + del iterator + + def test_embed_shuffle_shuffle_after(self) -> None: + """Iterable created by embed_shuffle calls shuffle automatically after iteration""" + + foo = IterableWithShuffle_(3) + iterable = embed_shuffle(foo, shuffle_last=True) + self.assertIsNone(foo._seed) + self.assertEqual(list(iterable), [0, 1, 2]) + self.assertEqual(foo._seed, 0) + self.assertEqual(list(iterable), [1, 2, 0]) + self.assertEqual(foo._seed, 1) + self.assertEqual(list(iterable), [2, 0, 1]) + self.assertEqual(foo._seed, 2) + + def test_embed_shuffle_shuffle_after_halt(self) -> None: + """The value is shuffled with different seed even after an iteration is halted.""" + + foo = IterableWithShuffle_(5) + iterable = embed_shuffle(foo, shuffle_last=True) + + iterator = iter(iterable) + self.assertEqual(next(iterator), 0) + self.assertEqual(next(iterator), 1) + self.assertIsNone(foo._seed) + del iterator + self.assertEqual(foo._seed, 0) + + iterator = iter(iterable) + self.assertEqual(next(iterator), 1) + self.assertEqual(next(iterator), 2) + del iterator + self.assertEqual(foo._seed, 1) diff --git a/tests/io/array_test.py b/src/spdl/io/tests/core/array_test.py similarity index 100% rename from tests/io/array_test.py rename to src/spdl/io/tests/core/array_test.py diff --git a/tests/io/async_test.py b/src/spdl/io/tests/core/async_test.py similarity index 99% rename from tests/io/async_test.py rename to src/spdl/io/tests/core/async_test.py index 37afbd99c..2315e677c 100644 --- a/tests/io/async_test.py +++ b/src/spdl/io/tests/core/async_test.py @@ -13,8 +13,7 @@ import spdl.io import spdl.io.utils from spdl.io import get_audio_filter_desc, get_video_filter_desc - -from ..fixture import FFMPEG_CLI, get_sample, get_samples +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, get_samples class TestDemuxer(unittest.TestCase): diff --git a/tests/io/audio_decoding_test.py b/src/spdl/io/tests/core/audio_decoding_test.py similarity index 97% rename from tests/io/audio_decoding_test.py rename to src/spdl/io/tests/core/audio_decoding_test.py index 53f7a9a65..9c3161f92 100644 --- a/tests/io/audio_decoding_test.py +++ b/src/spdl/io/tests/core/audio_decoding_test.py @@ -14,8 +14,7 @@ import spdl.io.utils from parameterized import parameterized from spdl.io import get_audio_filter_desc, get_filter_desc - -from ..fixture import FFMPEG_CLI, get_sample, load_ref_audio +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, load_ref_audio def _get_format(sample_fmt: str): diff --git a/tests/io/audio_encoding_test.py b/src/spdl/io/tests/core/audio_encoding_test.py similarity index 99% rename from tests/io/audio_encoding_test.py rename to src/spdl/io/tests/core/audio_encoding_test.py index f70973594..1e2bcbabe 100644 --- a/tests/io/audio_encoding_test.py +++ b/src/spdl/io/tests/core/audio_encoding_test.py @@ -14,8 +14,7 @@ import numpy as np import spdl.io from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample, load_ref_audio +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, load_ref_audio sample_fmt2dtype = { "s16": np.int16, diff --git a/tests/io/buffer_conversion_refcount_test.py b/src/spdl/io/tests/core/buffer_conversion_refcount_test.py similarity index 96% rename from tests/io/buffer_conversion_refcount_test.py rename to src/spdl/io/tests/core/buffer_conversion_refcount_test.py index 03e39ebf5..8e08d24f0 100644 --- a/tests/io/buffer_conversion_refcount_test.py +++ b/src/spdl/io/tests/core/buffer_conversion_refcount_test.py @@ -14,8 +14,7 @@ import spdl.io import spdl.io.utils from spdl.io import get_video_filter_desc - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestBufferConversionRefcount(unittest.TestCase): diff --git a/tests/io/configs_test.py b/src/spdl/io/tests/core/configs_test.py similarity index 97% rename from tests/io/configs_test.py rename to src/spdl/io/tests/core/configs_test.py index 8f79dfaf8..e04669b2b 100644 --- a/tests/io/configs_test.py +++ b/src/spdl/io/tests/core/configs_test.py @@ -9,8 +9,7 @@ import unittest import spdl.io - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestDemuxConfig(unittest.TestCase): diff --git a/tests/io/demuxer_test.py b/src/spdl/io/tests/core/demuxer_test.py similarity index 99% rename from tests/io/demuxer_test.py rename to src/spdl/io/tests/core/demuxer_test.py index 2712b1618..4c89e55bb 100644 --- a/tests/io/demuxer_test.py +++ b/src/spdl/io/tests/core/demuxer_test.py @@ -11,8 +11,7 @@ import spdl.io from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestDemuxer(unittest.TestCase): diff --git a/tests/io/encoding_test.py b/src/spdl/io/tests/core/encoding_test.py similarity index 98% rename from tests/io/encoding_test.py rename to src/spdl/io/tests/core/encoding_test.py index 81d156fb5..bf4149765 100644 --- a/tests/io/encoding_test.py +++ b/src/spdl/io/tests/core/encoding_test.py @@ -15,8 +15,7 @@ import spdl.io import torch from parameterized import parameterized - -from ..fixture import load_ref_image +from spdl.io.tests.fixture import load_ref_image class TestEncodeImageParity(unittest.TestCase): diff --git a/tests/io/filter_test.py b/src/spdl/io/tests/core/filter_test.py similarity index 99% rename from tests/io/filter_test.py rename to src/spdl/io/tests/core/filter_test.py index 1a27f0e68..997bc1271 100644 --- a/tests/io/filter_test.py +++ b/src/spdl/io/tests/core/filter_test.py @@ -12,8 +12,7 @@ import numpy as np import spdl.io from spdl.io import get_abuffer_desc, get_buffer_desc - -from ..fixture import FFMPEG_CLI, get_sample, load_ref_audio, load_ref_video +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, load_ref_audio, load_ref_video class FilterTest(unittest.TestCase): diff --git a/tests/io/frames_clone_test.py b/src/spdl/io/tests/core/frames_clone_test.py similarity index 98% rename from tests/io/frames_clone_test.py rename to src/spdl/io/tests/core/frames_clone_test.py index 337f4d3b9..5343174f3 100644 --- a/tests/io/frames_clone_test.py +++ b/src/spdl/io/tests/core/frames_clone_test.py @@ -14,8 +14,7 @@ from numpy.typing import NDArray from parameterized import parameterized from spdl.io import AudioFrames, ImageFrames, VideoFrames - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample CMDS = { "audio": f"{FFMPEG_CLI} -hide_banner -y -f lavfi -i sine=frequency=1000:sample_rate=48000:duration=3 -c:a pcm_s16le sample.wav", diff --git a/tests/io/frames_test.py b/src/spdl/io/tests/core/frames_test.py similarity index 96% rename from tests/io/frames_test.py rename to src/spdl/io/tests/core/frames_test.py index 117f3a8b9..faf369bfb 100644 --- a/tests/io/frames_test.py +++ b/src/spdl/io/tests/core/frames_test.py @@ -9,8 +9,7 @@ import unittest import spdl.io - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestFrames(unittest.TestCase): diff --git a/tests/io/image_decoding_test.py b/src/spdl/io/tests/core/image_decoding_test.py similarity index 99% rename from tests/io/image_decoding_test.py rename to src/spdl/io/tests/core/image_decoding_test.py index dd8bf4165..d6b92805e 100644 --- a/tests/io/image_decoding_test.py +++ b/src/spdl/io/tests/core/image_decoding_test.py @@ -12,10 +12,15 @@ import spdl.io import spdl.io.utils from spdl.io import get_video_filter_desc +from spdl.io.tests.fixture import ( + FFMPEG_CLI, + get_sample, + get_samples, + load_ref_data, + load_ref_image, +) from spdl.io.utils import get_ffmpeg_versions -from ..fixture import FFMPEG_CLI, get_sample, get_samples, load_ref_data, load_ref_image - def _load_image(src, filter_desc="format=pix_fmts=rgb24"): return spdl.io.to_numpy(spdl.io.load_image(src, filter_desc=filter_desc)) diff --git a/tests/io/memoryview_test.py b/src/spdl/io/tests/core/memoryview_test.py similarity index 100% rename from tests/io/memoryview_test.py rename to src/spdl/io/tests/core/memoryview_test.py diff --git a/tests/io/packets_test.py b/src/spdl/io/tests/core/packets_test.py similarity index 99% rename from tests/io/packets_test.py rename to src/spdl/io/tests/core/packets_test.py index bd21cae51..0bd8b6692 100644 --- a/tests/io/packets_test.py +++ b/src/spdl/io/tests/core/packets_test.py @@ -14,8 +14,7 @@ import numpy as np import spdl.io from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample CMDS = { "audio": f"{FFMPEG_CLI} -hide_banner -y -f lavfi -i sine=frequency=1000:sample_rate=48000:duration=3 -c:a pcm_s16le sample.wav", diff --git a/tests/io/reference_frames_test.py b/src/spdl/io/tests/core/reference_frames_test.py similarity index 100% rename from tests/io/reference_frames_test.py rename to src/spdl/io/tests/core/reference_frames_test.py diff --git a/tests/io/serialization_test.py b/src/spdl/io/tests/core/serialization_test.py similarity index 99% rename from tests/io/serialization_test.py rename to src/spdl/io/tests/core/serialization_test.py index c4069aefe..22f9a1e06 100644 --- a/tests/io/serialization_test.py +++ b/src/spdl/io/tests/core/serialization_test.py @@ -11,8 +11,7 @@ import numpy as np import spdl.io - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestAudioPacketsSerialization(unittest.TestCase): diff --git a/tests/io/streaming_decoding_test.py b/src/spdl/io/tests/core/streaming_decoding_test.py similarity index 99% rename from tests/io/streaming_decoding_test.py rename to src/spdl/io/tests/core/streaming_decoding_test.py index 59d5d7d3a..ce296b3f6 100644 --- a/tests/io/streaming_decoding_test.py +++ b/src/spdl/io/tests/core/streaming_decoding_test.py @@ -14,8 +14,7 @@ import torch from parameterized import parameterized from spdl.io import VideoPackets - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestDemuxer(unittest.TestCase): diff --git a/tests/io/subprocess_serialization_test.py b/src/spdl/io/tests/core/subprocess_serialization_test.py similarity index 98% rename from tests/io/subprocess_serialization_test.py rename to src/spdl/io/tests/core/subprocess_serialization_test.py index d30eed66c..85f8d489f 100644 --- a/tests/io/subprocess_serialization_test.py +++ b/src/spdl/io/tests/core/subprocess_serialization_test.py @@ -11,8 +11,7 @@ import numpy as np import spdl.io - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample def _demux_audio_worker( diff --git a/tests/io/tar_test.py b/src/spdl/io/tests/core/tar_test.py similarity index 100% rename from tests/io/tar_test.py rename to src/spdl/io/tests/core/tar_test.py diff --git a/tests/io/transfer_test.py b/src/spdl/io/tests/core/transfer_test.py similarity index 100% rename from tests/io/transfer_test.py rename to src/spdl/io/tests/core/transfer_test.py diff --git a/tests/io/utils_test.py b/src/spdl/io/tests/core/utils_test.py similarity index 100% rename from tests/io/utils_test.py rename to src/spdl/io/tests/core/utils_test.py diff --git a/tests/io/video_encoding_test.py b/src/spdl/io/tests/core/video_encoding_test.py similarity index 98% rename from tests/io/video_encoding_test.py rename to src/spdl/io/tests/core/video_encoding_test.py index f92848cec..3081bc9ff 100644 --- a/tests/io/video_encoding_test.py +++ b/src/spdl/io/tests/core/video_encoding_test.py @@ -12,8 +12,7 @@ import numpy as np import spdl.io from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample, load_ref_video +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, load_ref_video class TestEncodeVideoMultiColor(unittest.TestCase): diff --git a/tests/io/video_frame_slice_test.py b/src/spdl/io/tests/core/video_frame_slice_test.py similarity index 98% rename from tests/io/video_frame_slice_test.py rename to src/spdl/io/tests/core/video_frame_slice_test.py index 0409f6823..1e330b809 100644 --- a/tests/io/video_frame_slice_test.py +++ b/src/spdl/io/tests/core/video_frame_slice_test.py @@ -12,8 +12,7 @@ import spdl.io import spdl.io.utils from spdl.io import get_video_filter_desc - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample def _to_numpy(frames): diff --git a/tests/io/wav_test.py b/src/spdl/io/tests/core/wav_test.py similarity index 100% rename from tests/io/wav_test.py rename to src/spdl/io/tests/core/wav_test.py diff --git a/tests/io/zero_copy_bytes_passing_test.py b/src/spdl/io/tests/core/zero_copy_bytes_passing_test.py similarity index 97% rename from tests/io/zero_copy_bytes_passing_test.py rename to src/spdl/io/tests/core/zero_copy_bytes_passing_test.py index d90464c6e..457fc279f 100644 --- a/tests/io/zero_copy_bytes_passing_test.py +++ b/src/spdl/io/tests/core/zero_copy_bytes_passing_test.py @@ -12,8 +12,7 @@ import numpy as np import spdl.io from numpy.typing import NDArray - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample def _decode(media_type: str, src: str | bytes) -> NDArray[Any]: diff --git a/tests/io/zip_test.py b/src/spdl/io/tests/core/zip_test.py similarity index 100% rename from tests/io/zip_test.py rename to src/spdl/io/tests/core/zip_test.py diff --git a/tests/cuda/buffer_transfer_test.py b/src/spdl/io/tests/cuda/buffer_transfer_test.py similarity index 99% rename from tests/cuda/buffer_transfer_test.py rename to src/spdl/io/tests/cuda/buffer_transfer_test.py index d010779b0..b551bfe98 100644 --- a/tests/cuda/buffer_transfer_test.py +++ b/src/spdl/io/tests/cuda/buffer_transfer_test.py @@ -12,8 +12,7 @@ import spdl.io.utils import torch from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample if not spdl.io.utils.built_with_cuda(): raise unittest.SkipTest("SPDL is not compiled with CUDA support") diff --git a/tests/cuda/transfer_test.py b/src/spdl/io/tests/cuda/cuda_transfer_test.py similarity index 100% rename from tests/cuda/transfer_test.py rename to src/spdl/io/tests/cuda/cuda_transfer_test.py diff --git a/tests/cuda/nvdec_video_decoding_test.py b/src/spdl/io/tests/cuda/nvdec_video_decoding_test.py similarity index 99% rename from tests/cuda/nvdec_video_decoding_test.py rename to src/spdl/io/tests/cuda/nvdec_video_decoding_test.py index 629667cca..dd961bda2 100644 --- a/tests/cuda/nvdec_video_decoding_test.py +++ b/src/spdl/io/tests/cuda/nvdec_video_decoding_test.py @@ -12,8 +12,7 @@ import spdl.io import spdl.io.utils import torch - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample if not spdl.io.utils.built_with_nvcodec(): raise unittest.SkipTest( # pyre-ignore: [29] diff --git a/tests/cuda/nvjpeg_decode_test.py b/src/spdl/io/tests/cuda/nvjpeg_decode_test.py similarity index 98% rename from tests/cuda/nvjpeg_decode_test.py rename to src/spdl/io/tests/cuda/nvjpeg_decode_test.py index b2e3e9730..819e9802d 100644 --- a/tests/cuda/nvjpeg_decode_test.py +++ b/src/spdl/io/tests/cuda/nvjpeg_decode_test.py @@ -10,8 +10,7 @@ import spdl.io import spdl.io.utils import torch - -from ..fixture import FFMPEG_CLI, get_sample, get_samples +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, get_samples DEFAULT_CUDA = 0 diff --git a/tests/cuda/pin_memory_test.py b/src/spdl/io/tests/cuda/pin_memory_test.py similarity index 98% rename from tests/cuda/pin_memory_test.py rename to src/spdl/io/tests/cuda/pin_memory_test.py index add25b065..40aa45c54 100644 --- a/tests/cuda/pin_memory_test.py +++ b/src/spdl/io/tests/cuda/pin_memory_test.py @@ -12,8 +12,7 @@ import spdl.io import torch from parameterized import parameterized - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample class TestPinMemory(unittest.TestCase): diff --git a/tests/cuda/streaming_load_video_nvdec_test.py b/src/spdl/io/tests/cuda/streaming_load_video_nvdec_test.py similarity index 99% rename from tests/cuda/streaming_load_video_nvdec_test.py rename to src/spdl/io/tests/cuda/streaming_load_video_nvdec_test.py index 9f8a0f12c..12a7c6e1d 100644 --- a/tests/cuda/streaming_load_video_nvdec_test.py +++ b/src/spdl/io/tests/cuda/streaming_load_video_nvdec_test.py @@ -11,8 +11,7 @@ import spdl.io import spdl.io.utils import torch - -from ..fixture import FFMPEG_CLI, get_sample, SrcInfo +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample, SrcInfo if not spdl.io.utils.built_with_nvcodec(): raise unittest.SkipTest("SPDL is not compiled with NVCODEC support") diff --git a/tests/cuda/subprocess_nvdec_test.py b/src/spdl/io/tests/cuda/subprocess_nvdec_test.py similarity index 98% rename from tests/cuda/subprocess_nvdec_test.py rename to src/spdl/io/tests/cuda/subprocess_nvdec_test.py index 5a83a95bb..b908c435e 100644 --- a/tests/cuda/subprocess_nvdec_test.py +++ b/src/spdl/io/tests/cuda/subprocess_nvdec_test.py @@ -11,8 +11,7 @@ import spdl.io import spdl.io.utils import torch - -from ..fixture import FFMPEG_CLI, get_sample +from spdl.io.tests.fixture import FFMPEG_CLI, get_sample if not spdl.io.utils.built_with_nvcodec(): raise unittest.SkipTest( # pyre-ignore: [29] diff --git a/tests/fixture.py b/src/spdl/io/tests/fixture.py similarity index 100% rename from tests/fixture.py rename to src/spdl/io/tests/fixture.py diff --git a/src/spdl/pipeline/tests/aggregate_test.py b/src/spdl/pipeline/tests/aggregate_test.py new file mode 100644 index 000000000..614c81498 --- /dev/null +++ b/src/spdl/pipeline/tests/aggregate_test.py @@ -0,0 +1,326 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import unittest +from collections.abc import Iterable +from typing import Any + +from spdl.pipeline import AsyncQueue, PipelineBuilder +from spdl.pipeline._components._aggregate import _aggregate +from spdl.pipeline._components._common import _EOF, StageInfo +from spdl.pipeline._components._pipe import _get_fail_counter +from spdl.pipeline.defs import Aggregator + +_TEST_INFO = StageInfo(pipeline_id=0, stage_id="0", stage_name="test") + + +def _put_aqueue(queue: AsyncQueue, vals: Iterable[object], *, eof: bool) -> None: + for val in vals: + queue.put_nowait(val) + if eof: + queue.put_nowait(_EOF) + + +def _flush_aqueue(queue: AsyncQueue) -> list[object]: + ret = [] + while not queue.empty(): + ret.append(queue.get_nowait()) + return ret + + +class _TrackingAggregator(Aggregator): + """Aggregator that collects items into batches of size N and records + accumulate call order for verifying drain behavior.""" + + def __init__(self, batch_size: int) -> None: + self.batch_size = batch_size + self.buffer: list[Any] = [] + self.accumulate_log: list[Any] = [] + + def accumulate(self, item: Any) -> list[Any] | None: + self.accumulate_log.append(item) + self.buffer.append(item) + if len(self.buffer) >= self.batch_size: + result = self.buffer + self.buffer = [] + return result + return None + + def flush(self) -> list[Any] | None: + if self.buffer: + result = self.buffer + self.buffer = [] + return result + return None + + +class AggregatePipeBulkDrainTest(unittest.IsolatedAsyncioTestCase): + async def test_stops_draining_on_emit(self) -> None: + """After the aggregator emits, bulk draining stops and remaining + items stay in the input queue for the next drain cycle.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + agg = _TrackingAggregator(batch_size=3) + # Pre-fill 7 items + EOF. The aggregator emits after every 3 items. + # With stop-on-emit, each drain cycle processes exactly 3 items + # (emit batch), then stops. The remaining items stay in the input + # queue for the next cycle. The last item is flushed at EOF. + _put_aqueue(input_queue, list(range(7)), eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=True, + ) + + results = _flush_aqueue(output_queue) + # Two full batches of 3 + flush of remainder [6] + EOF + self.assertEqual(results, [[0, 1, 2], [3, 4, 5], [6], _EOF]) + # All 7 items were accumulated + self.assertEqual(agg.accumulate_log, [0, 1, 2, 3, 4, 5, 6]) + + async def test_drains_without_blocking_when_no_emit(self) -> None: + """When the aggregator doesn't emit, items are drained from the + queue without blocking (via get_nowait).""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + # batch_size=10 means 5 items won't trigger an emit + agg = _TrackingAggregator(batch_size=10) + + # Pre-fill 5 items + EOF. None will trigger an emit during + # accumulate, but flush will emit the remaining buffer. + _put_aqueue(input_queue, list(range(5)), eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=True, + ) + + results = _flush_aqueue(output_queue) + # flush emits the remaining 5 items + self.assertEqual(results, [[0, 1, 2, 3, 4], _EOF]) + self.assertEqual(agg.accumulate_log, [0, 1, 2, 3, 4]) + + async def test_drop_last(self) -> None: + """When op_requires_eof=False, EOF stops processing and flush + is not called, dropping incomplete batches.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + agg = _TrackingAggregator(batch_size=3) + + # 5 items: batch at 3, then 2 remaining are dropped + _put_aqueue(input_queue, list(range(5)), eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=False, + ) + + results = _flush_aqueue(output_queue) + # Only the complete batch + EOF from queue_stage_hook + self.assertEqual(results, [[0, 1, 2], _EOF]) + + async def test_exception_propagates(self) -> None: + """Exceptions from the aggregator propagate instead of being + silently swallowed.""" + + class FailingAggregator(Aggregator): + def accumulate(self, item: Any) -> None: + raise ValueError("aggregation failed") + + def flush(self) -> None: + return None + + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + _put_aqueue(input_queue, [1], eof=True) + + with self.assertRaises(ValueError, msg="aggregation failed"): + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + FailingAggregator(), + _get_fail_counter()(), + [], + op_requires_eof=False, + ) + + async def test_single_item_no_emit(self) -> None: + """A single item that doesn't trigger emit is flushed at EOF.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + agg = _TrackingAggregator(batch_size=5) + + _put_aqueue(input_queue, [42], eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=True, + ) + + results = _flush_aqueue(output_queue) + self.assertEqual(results, [[42], _EOF]) + + async def test_eof_with_empty_flush(self) -> None: + """When items evenly divide batch_size, flush() returns None at EOF. + The function must still return instead of blocking.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + agg = _TrackingAggregator(batch_size=3) + + # 6 items evenly divides batch_size=3, so flush() returns None + _put_aqueue(input_queue, list(range(6)), eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=True, + ) + + results = _flush_aqueue(output_queue) + self.assertEqual(results, [[0, 1, 2], [3, 4, 5], _EOF]) + + async def test_emit_on_every_item(self) -> None: + """When batch_size=1, every item triggers an emit and each + drain cycle processes exactly one item.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + agg = _TrackingAggregator(batch_size=1) + + _put_aqueue(input_queue, [10, 20, 30], eof=True) + + await _aggregate( + _TEST_INFO, + input_queue, + output_queue, + agg, + _get_fail_counter()(), + [], + op_requires_eof=False, + ) + + results = _flush_aqueue(output_queue) + self.assertEqual(results, [[10], [20], [30], _EOF]) + + +class AggregatePipeEndToEndTest(unittest.TestCase): + def test_aggregate_bulk_drain_correctness(self) -> None: + """End-to-end pipeline test: aggregate produces correct batches.""" + src = list(range(10)) + + pipeline = ( + PipelineBuilder() + .add_source(src) + .aggregate(3) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(results, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) + + def test_aggregate_custom_op_bulk_drain(self) -> None: + """End-to-end: custom aggregator with bulk drain produces correct output.""" + + class SizeAggregator(Aggregator): + def __init__(self, threshold: int) -> None: + self.threshold = threshold + self.buffer: list[str] = [] + self.total: int = 0 + + def accumulate(self, item: str) -> str | None: + self.buffer.append(item) + self.total += len(item) + if self.total >= self.threshold: + result = "".join(self.buffer) + self.buffer = [] + self.total = 0 + return result + return None + + def flush(self) -> str | None: + if self.buffer: + result = "".join(self.buffer) + self.buffer = [] + self.total = 0 + return result + return None + + src = ["a", "bb", "ccc", "dddd", "e", "ff", "ggg", "h"] + + pipeline = ( + PipelineBuilder() + .add_source(src) + .aggregate(SizeAggregator(threshold=10)) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(results, ["abbcccdddd", "effgggh"]) diff --git a/src/spdl/pipeline/tests/background_task_test.py b/src/spdl/pipeline/tests/background_task_test.py new file mode 100644 index 000000000..da6582813 --- /dev/null +++ b/src/spdl/pipeline/tests/background_task_test.py @@ -0,0 +1,278 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import time +import unittest + +from spdl.pipeline import BackgroundTask, build_pipeline +from spdl.pipeline.config import ( + get_default_background_tasks, + set_default_background_tasks, +) +from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig + + +def _simple_cfg() -> PipelineConfig[int]: + return PipelineConfig( + src=SourceConfig(range(5)), + pipes=[Pipe(lambda x: x)], + sink=SinkConfig(3), + ) + + +def _slow_cfg(delay: float = 0.5) -> PipelineConfig[int]: + """Pipeline that takes a while to complete, giving background tasks time to run.""" + + def slow_op(x: int) -> int: + time.sleep(delay) + return x + + return PipelineConfig( + src=SourceConfig(range(3)), + pipes=[Pipe(slow_op)], + sink=SinkConfig(3), + ) + + +class _TrackingTask(BackgroundTask): + """Background task that tracks whether it started and was cancelled.""" + + def __init__(self) -> None: + self.started = False + self.cancelled = False + + async def run(self) -> None: + self.started = True + try: + while True: + await asyncio.sleep(0.01) + except asyncio.CancelledError: + self.cancelled = True + raise + + +class _CountingTask(BackgroundTask): + """Background task that counts iterations.""" + + def __init__(self) -> None: + self.count = 0 + + async def run(self) -> None: + try: + while True: + self.count += 1 + await asyncio.sleep(0.01) + except asyncio.CancelledError: + pass + + +class _FailingTask(BackgroundTask): + async def run(self) -> None: + raise RuntimeError("bg task error") + + +class _ShortTask(BackgroundTask): + """Background task that completes quickly.""" + + def __init__(self) -> None: + self.completed = False + + async def run(self) -> None: + await asyncio.sleep(0.01) + self.completed = True + + +class BackgroundTaskTest(unittest.TestCase): + def setUp(self) -> None: + self._saved = get_default_background_tasks() + set_default_background_tasks(None) + + def tearDown(self) -> None: + set_default_background_tasks(self._saved) + + def test_background_task_runs_and_is_cancelled(self) -> None: + """Background tasks run alongside pipeline and get cancelled on completion.""" + task_instance: _TrackingTask = _TrackingTask() + + def factory() -> BackgroundTask: + return task_instance + + pipeline = build_pipeline( + _simple_cfg(), num_threads=1, background_tasks=[factory] + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertTrue(task_instance.started, "Background task should have started") + self.assertTrue( + task_instance.cancelled, "Background task should have been cancelled" + ) + + def test_background_task_error_does_not_crash_pipeline(self) -> None: + """Background task errors are logged but don't fail the pipeline.""" + pipeline = build_pipeline( + _simple_cfg(), num_threads=1, background_tasks=[_FailingTask] + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + + def test_multiple_background_tasks(self) -> None: + """Multiple background tasks can run concurrently.""" + task_0 = _CountingTask() + task_1 = _CountingTask() + + pipeline = build_pipeline( + _simple_cfg(), + num_threads=1, + background_tasks=[lambda: task_0, lambda: task_1], + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertGreater(task_0.count, 0, "First background task should have run") + self.assertGreater(task_1.count, 0, "Second background task should have run") + + def test_no_background_tasks(self) -> None: + """Pipeline works normally when no background tasks are provided.""" + pipeline = build_pipeline(_simple_cfg(), num_threads=1) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + + def test_empty_background_tasks_list(self) -> None: + """Pipeline works normally with an empty background tasks list.""" + pipeline = build_pipeline(_simple_cfg(), num_threads=1, background_tasks=[]) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + + def test_background_task_completes_before_pipeline(self) -> None: + """A background task that finishes early doesn't affect the pipeline.""" + task_instance = _ShortTask() + + pipeline = build_pipeline( + _slow_cfg(0.1), + num_threads=1, + background_tasks=[lambda: task_instance], + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(sorted(items), [0, 1, 2]) + self.assertTrue( + task_instance.completed, "Background task should have completed" + ) + + def test_background_task_mixed_success_and_failure(self) -> None: + """One failing and one succeeding background task — pipeline still works.""" + good_task = _TrackingTask() + + pipeline = build_pipeline( + _simple_cfg(), + num_threads=1, + background_tasks=[lambda: good_task, _FailingTask], + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertTrue(good_task.started, "Good background task should have run") + + def test_class_as_factory(self) -> None: + """A BackgroundTask class itself can be used as a factory.""" + + class MyTask(BackgroundTask): + ran = False + + async def run(self) -> None: + MyTask.ran = True + try: + while True: + await asyncio.sleep(0.01) + except asyncio.CancelledError: + pass + + MyTask.ran = False + pipeline = build_pipeline( + _simple_cfg(), num_threads=1, background_tasks=[MyTask] + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertTrue(MyTask.ran, "Task class used as factory should have run") + + +class DefaultBackgroundTasksTest(unittest.TestCase): + def setUp(self) -> None: + self._saved = get_default_background_tasks() + set_default_background_tasks(None) + + def tearDown(self) -> None: + set_default_background_tasks(self._saved) + + def test_get_set_default_background_tasks(self) -> None: + """set/get_default_background_tasks round-trips correctly.""" + self.assertIsNone(get_default_background_tasks()) + + set_default_background_tasks([_TrackingTask]) + result = get_default_background_tasks() + self.assertIsNotNone(result) + self.assertEqual(len(result), 1) # pyre-ignore[6] + self.assertIs(result[0], _TrackingTask) # pyre-ignore[16] + + set_default_background_tasks(None) + self.assertIsNone(get_default_background_tasks()) + + def test_default_background_tasks_run_automatically(self) -> None: + """Default background tasks are started without explicit parameter.""" + task_instance = _TrackingTask() + set_default_background_tasks([lambda: task_instance]) + + pipeline = build_pipeline(_simple_cfg(), num_threads=1) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertTrue( + task_instance.started, "Default background task should have run" + ) + + def test_default_and_per_pipeline_tasks_merged(self) -> None: + """Both default and per-pipeline background tasks run.""" + default_task = _TrackingTask() + custom_task = _TrackingTask() + + set_default_background_tasks([lambda: default_task]) + + pipeline = build_pipeline( + _simple_cfg(), num_threads=1, background_tasks=[lambda: custom_task] + ) + + with pipeline.auto_stop(): + items = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) + self.assertTrue(default_task.started, "Default background task should have run") + self.assertTrue(custom_task.started, "Custom background task should have run") diff --git a/src/spdl/pipeline/tests/build_pipeline_test.py b/src/spdl/pipeline/tests/build_pipeline_test.py new file mode 100644 index 000000000..7dd52da5e --- /dev/null +++ b/src/spdl/pipeline/tests/build_pipeline_test.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import os +import unittest +import warnings +from unittest.mock import patch + +from spdl.pipeline import build_pipeline +from spdl.pipeline._profile import _ProfilePipeline +from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig + +# pyre-strict + + +class TestBuildPipeline(unittest.TestCase): + """Test class for build_pipeline functionality.""" + + def test_build_pipeline_diagnostic_mode(self) -> None: + """Test that when SPDL_PIPELINE_DIAGNOSTIC_MODE=1, build_pipeline + calls _build_pipeline_diagnostic_mode and returns _ProfilePipeline. + """ + + def simple_op(i: int) -> int: + return i * 2 + + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="coroutine .* was never awaited", + category=RuntimeWarning, + ) + with patch.dict("os.environ", {"SPDL_PIPELINE_DIAGNOSTIC_MODE": "1"}): + pipeline = build_pipeline(cfg, num_threads=2) + self.assertIsInstance(pipeline, _ProfilePipeline) + + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("SPDL_PIPELINE_DIAGNOSTIC_MODE", None) + + pipeline = build_pipeline(cfg, num_threads=2) + self.assertNotIsInstance(pipeline, _ProfilePipeline) diff --git a/src/spdl/pipeline/tests/compact_log_test.py b/src/spdl/pipeline/tests/compact_log_test.py new file mode 100644 index 000000000..031a9405d --- /dev/null +++ b/src/spdl/pipeline/tests/compact_log_test.py @@ -0,0 +1,279 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import os +import unittest + +from spdl.pipeline._common._misc import _get_compact_log, _set_compact_log, create_task +from spdl.pipeline.config import set_compact_log + + +class DummyException(Exception): + """Test exception for simulating task failures.""" + + pass + + +class CompactLogTest(unittest.TestCase): + """Tests for compact logging mode functionality.""" + + def setUp(self) -> None: + """Reset the global compact log setting before each test.""" + # Reset to None to ensure clean state + _set_compact_log(None) + # Clear any environment variable that might be set + if "SPDL_PIPELINE_COMPACT_LOG" in os.environ: + del os.environ["SPDL_PIPELINE_COMPACT_LOG"] + + def tearDown(self) -> None: + """Clean up after each test.""" + # Reset to None + _set_compact_log(None) + # Clear environment variable + if "SPDL_PIPELINE_COMPACT_LOG" in os.environ: + del os.environ["SPDL_PIPELINE_COMPACT_LOG"] + + def test_get_compact_log_defaults_to_false_when_env_not_set(self) -> None: + """Test that _get_compact_log returns False when environment variable is not set.""" + # Setup: Environment variable is not set (cleared in setUp) + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should default to False + self.assertFalse(result) + + def test_get_compact_log_returns_true_when_env_is_1(self) -> None: + """Test that _get_compact_log returns True when environment variable is '1'.""" + # Setup: Set environment variable to '1' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "1" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_get_compact_log_returns_true_when_env_is_true(self) -> None: + """Test that _get_compact_log returns True when environment variable is 'true'.""" + # Setup: Set environment variable to 'true' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_get_compact_log_returns_true_when_env_is_TRUE(self) -> None: + """Test that _get_compact_log returns True when environment variable is 'TRUE'.""" + # Setup: Set environment variable to 'TRUE' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "TRUE" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_get_compact_log_returns_true_when_env_is_yes(self) -> None: + """Test that _get_compact_log returns True when environment variable is 'yes'.""" + # Setup: Set environment variable to 'yes' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "yes" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_get_compact_log_returns_false_when_env_is_0(self) -> None: + """Test that _get_compact_log returns False when environment variable is '0'.""" + # Setup: Set environment variable to '0' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "0" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return False + self.assertFalse(result) + + def test_get_compact_log_returns_false_when_env_is_false(self) -> None: + """Test that _get_compact_log returns False when environment variable is 'false'.""" + # Setup: Set environment variable to 'false' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return False + self.assertFalse(result) + + def test_get_compact_log_caches_result_after_first_call(self) -> None: + """Test that _get_compact_log caches the result and doesn't re-check env var.""" + # Setup: Set environment variable to 'true' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" + + # Execute: Get compact log setting twice + result1 = _get_compact_log() + # Change environment variable + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" + result2 = _get_compact_log() + + # Assert: Both should return True because value is cached + self.assertTrue(result1) + self.assertTrue(result2) + + def test_set_compact_log_to_true(self) -> None: + """Test that _set_compact_log can set the value to True.""" + # Setup: Set to True + _set_compact_log(True) + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_set_compact_log_to_false(self) -> None: + """Test that _set_compact_log can set the value to False.""" + # Setup: Set to False + _set_compact_log(False) + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return False + self.assertFalse(result) + + def test_set_compact_log_to_none_resets_to_env_check(self) -> None: + """Test that setting to None causes re-check of environment variable.""" + # Setup: Set to True first + _set_compact_log(True) + self.assertTrue(_get_compact_log()) + + # Reset to None + _set_compact_log(None) + # Set environment variable + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" + + # Execute: Get compact log setting after reset + result = _get_compact_log() + + # Assert: Should return False from environment variable + self.assertFalse(result) + + def test_set_compact_log_overrides_env_var(self) -> None: + """Test that programmatically setting the value overrides environment variable.""" + # Setup: Set environment variable to 'true' + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" + # Override with False + _set_compact_log(False) + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return False (overridden value, not env var) + self.assertFalse(result) + + def test_config_module_exposes_set_compact_log(self) -> None: + """Test that set_compact_log is exposed in the config module.""" + # Setup: Set through config module + set_compact_log(True) + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result) + + def test_create_task_properly_calls_compact_log_setting(self) -> None: + """Test that create_task respects the compact log setting.""" + + async def failing_coro() -> None: + raise DummyException("test error") + + async def run() -> None: + # Test with compact=False (default) + _set_compact_log(False) + task1 = create_task(failing_coro(), name="task1") + await asyncio.sleep(0) + try: + await task1 + except DummyException: + pass + # Task completed, no assertion needed - just verify no exceptions + + # Test with compact=True + _set_compact_log(True) + task2 = create_task(failing_coro(), name="task2") + await asyncio.sleep(0) + try: + await task2 + except DummyException: + pass + # Task completed, no assertion needed - just verify no exceptions + + asyncio.run(run()) + + def test_create_task_uses_get_compact_log(self) -> None: + """Test that create_task gets the compact setting from _get_compact_log.""" + + async def failing_coro() -> None: + raise DummyException("test error") + + async def run() -> None: + # Setup: Set compact log via environment variable + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "1" + _set_compact_log(None) # Reset to force env var check + + # Execute: Create task - this should use compact mode from env var + task = create_task(failing_coro(), name="test_task") + await asyncio.sleep(0) + try: + await task + except DummyException: + pass + + # Assert: Verify the getter returns True (from env var) + self.assertTrue(_get_compact_log()) + + asyncio.run(run()) + + def test_get_compact_log_with_various_truthy_env_values(self) -> None: + """Test that _get_compact_log handles various truthy environment values.""" + truthy_values = ["1", "true", "TRUE", "on", "ON", "yes", "YES"] + + for value in truthy_values: + with self.subTest(value=value): + # Setup: Reset and set env var + _set_compact_log(None) + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = value + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return True + self.assertTrue(result, f"Expected True for value '{value}'") + + def test_get_compact_log_with_various_falsy_env_values(self) -> None: + """Test that _get_compact_log handles various falsy environment values.""" + falsy_values = ["0", "false", "FALSE", "off", "OFF", "no", "NO"] + + for value in falsy_values: + with self.subTest(value=value): + # Setup: Reset and set env var + _set_compact_log(None) + os.environ["SPDL_PIPELINE_COMPACT_LOG"] = value + + # Execute: Get compact log setting + result = _get_compact_log() + + # Assert: Should return False + self.assertFalse(result, f"Expected False for value '{value}'") diff --git a/src/spdl/pipeline/tests/config_test.py b/src/spdl/pipeline/tests/config_test.py new file mode 100644 index 000000000..b64274b8e --- /dev/null +++ b/src/spdl/pipeline/tests/config_test.py @@ -0,0 +1,210 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import unittest +from collections.abc import Iterator +from contextlib import contextmanager + +from spdl.pipeline import ( + AsyncQueue, + ProfileHook, + ProfileResult, + StatsQueue, + TaskHook, + TaskStatsHook, +) +from spdl.pipeline.config import ( + get_default_hook_class, + get_default_profile_callback, + get_default_profile_hook, + get_default_queue_class, + set_default_hook_class, + set_default_profile_callback, + set_default_profile_hook, + set_default_queue_class, +) + + +def reset() -> None: + set_default_hook_class() + set_default_queue_class() + set_default_profile_hook() + set_default_profile_callback() + + +class ConfigTest(unittest.TestCase): + """Test the configuration setter/getter functions.""" + + def setUp(self) -> None: + """Reset all configuration state before each test.""" + super().setUp() + reset() + + def tearDown(self) -> None: + """Reset all configuration state after each test.""" + reset() + super().tearDown() + + def test_hook_class_default_is_none(self) -> None: + """Test that default hook class is None when not configured.""" + result = get_default_hook_class() + self.assertIs(result, TaskStatsHook) + + def test_hook_class_can_be_set_and_retrieved(self) -> None: + """Test that hook class can be set and retrieved correctly.""" + + class CustomHook(TaskHook): + pass + + set_default_hook_class(CustomHook) + result = get_default_hook_class() + + self.assertIs(result, CustomHook) + + def test_hook_class_via_config_module(self) -> None: + """Test that hook class can be accessed via _config module.""" + + class CustomHook(TaskHook): + pass + + set_default_hook_class(CustomHook) + result = get_default_hook_class() + + self.assertIs(result, CustomHook) + + def test_queue_class_default_is_none(self) -> None: + """Test that default queue class is None when not configured.""" + result = get_default_queue_class() + self.assertIs(result, StatsQueue) + + def test_queue_class_can_be_set_and_retrieved(self) -> None: + """Test that queue class can be set and retrieved correctly.""" + + class CustomQueue(StatsQueue): + pass + + set_default_queue_class(CustomQueue) + result = get_default_queue_class() + self.assertIs(result, CustomQueue) + + def test_queue_class_via_config_module(self) -> None: + """Test that queue class can be accessed via _config module.""" + + class CustomQueue(StatsQueue): + pass + + set_default_queue_class(CustomQueue) + result = get_default_queue_class() + self.assertIs(result, CustomQueue) + + def test_profile_hook_default_is_none(self) -> None: + """Test that default profile hook is None when not configured.""" + result = get_default_profile_hook() + self.assertIsNone(result) + + def test_profile_hook_can_be_set_and_retrieved(self) -> None: + """Test that profile hook can be set and retrieved correctly.""" + + class MockProfileHook(ProfileHook): + @contextmanager + def stage_profile_hook( + self, + stage: str, # noqa: ARG002 + concurrency: int, # noqa: ARG002 + ) -> Iterator[None]: + yield + + @contextmanager + def pipeline_profile_hook(self) -> Iterator[None]: + yield + + hook_instance = MockProfileHook() + set_default_profile_hook(hook_instance) + result = get_default_profile_hook() + self.assertIs(result, hook_instance) + + def test_profile_hook_via_config_module(self) -> None: + """Test that profile hook can be accessed via _config module.""" + + class MockProfileHook(ProfileHook): + @contextmanager + def stage_profile_hook( + self, + stage: str, # noqa: ARG002 + concurrency: int, # noqa: ARG002 + ) -> Iterator[None]: + yield + + @contextmanager + def pipeline_profile_hook(self) -> Iterator[None]: + yield + + hook_instance = MockProfileHook() + set_default_profile_hook(hook_instance) + result = get_default_profile_hook() + self.assertIs(result, hook_instance) + + def test_profile_callback_default_is_none(self) -> None: + """Test that default profile callback is None when not configured.""" + result = get_default_profile_callback() + self.assertIsNone(result) + + def test_profile_callback_can_be_set_and_retrieved(self) -> None: + """Test that profile callback can be set and retrieved correctly.""" + + def mock_callback(_: ProfileResult) -> None: + pass + + set_default_profile_callback(mock_callback) + result = get_default_profile_callback() + self.assertIs(result, mock_callback) + + def test_profile_callback_via_config_module(self) -> None: + """Test that profile callback can be accessed via _config module.""" + + def mock_callback(_: object) -> None: + pass + + set_default_profile_callback(mock_callback) + result = get_default_profile_callback() + self.assertIs(result, mock_callback) + + def test_multiple_configurations_independent(self) -> None: + """Test that different configuration settings are independent.""" + + class CustomHook(TaskHook): + pass + + class CustomQueue(AsyncQueue): + pass + + def custom_callback(_: object) -> None: + pass + + set_default_hook_class(CustomHook) + set_default_queue_class(CustomQueue) + set_default_profile_callback(custom_callback) + + self.assertIs(get_default_hook_class(), CustomHook) + self.assertIs(get_default_queue_class(), CustomQueue) + self.assertIs(get_default_profile_callback(), custom_callback) + + def test_configuration_can_be_updated(self) -> None: + """Test that configuration can be updated to new values.""" + + class FirstHook(TaskHook): + pass + + class SecondHook(TaskHook): + pass + + # Set first value, then update to second value + set_default_hook_class(FirstHook) + set_default_hook_class(SecondHook) + result = get_default_hook_class() + self.assertIs(result, SecondHook) diff --git a/src/spdl/pipeline/tests/continuous_pipeline_test.py b/src/spdl/pipeline/tests/continuous_pipeline_test.py new file mode 100644 index 000000000..2affb6466 --- /dev/null +++ b/src/spdl/pipeline/tests/continuous_pipeline_test.py @@ -0,0 +1,623 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import functools +import os +import sys +import threading +import time +import unittest +import warnings +import weakref +from collections.abc import Iterator + +from spdl.pipeline import ( + build_pipeline, + PipelineBuilder, + PipelineFailure, + run_pipeline_in_subinterpreter, + run_pipeline_in_subprocess, +) +from spdl.pipeline.defs import Merge, PipelineConfig, SinkConfig + + +def _ignore_fork_warning(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + return wrapper + + +class SourceIterable: + """Reusable iterable that yields range(n) on each iteration.""" + + def __init__(self, n: int) -> None: + self.n = n + + def __iter__(self) -> Iterator[int]: + yield from range(self.n) + + +class TestContinuousPipelineBasic(unittest.TestCase): + def test_continuous_multi_epoch(self) -> None: + """Pipeline with continuous=True can be iterated multiple times.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") + + def test_continuous_single_epoch(self) -> None: + """Continuous pipeline works for a single epoch.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(3), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 1, 2]) + + def test_continuous_empty_epoch(self) -> None: + """Continuous pipeline handles empty iterations.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(0), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [], f"epoch {epoch}") + + def test_continuous_single_item_epoch(self) -> None: + """Continuous pipeline works with single-item epochs.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(1), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0], f"epoch {epoch}") + + +class TestContinuousPipelinePipe(unittest.TestCase): + def test_continuous_pipe_concurrent(self) -> None: + """Pipe with concurrency > 1 handles epoch boundaries correctly.""" + + def double(x: int) -> int: + return x * 2 + + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .pipe(double, concurrency=4) + .add_sink(buffer_size=3) + .build(num_threads=4) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = sorted(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 2, 4, 6, 8], f"epoch {epoch}") + + def test_continuous_pipe_chain(self) -> None: + """EPOCH_END propagates through multiple pipe stages.""" + + def add_one(x: int) -> int: + return x + 1 + + def double(x: int) -> int: + return x * 2 + + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(3), continuous=True) + .pipe(add_one, concurrency=1) + .pipe(double, concurrency=1) + .add_sink(buffer_size=3) + .build(num_threads=2) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [2, 4, 6], f"epoch {epoch}") + + +class TestContinuousPipelineAggregate(unittest.TestCase): + def test_continuous_aggregate_exact_batch(self) -> None: + """Aggregate with exact batch size across epochs.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(6), continuous=True) + .aggregate(3) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [[0, 1, 2], [3, 4, 5]], f"epoch {epoch}") + + def test_continuous_aggregate_partial_batch_flushed(self) -> None: + """Partial batch at epoch end is flushed by the default aggregator.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .aggregate(3) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + # 5 items, batch_size=3: full batch [0,1,2] + partial batch [3,4] + self.assertEqual(result, [[0, 1, 2], [3, 4]], f"epoch {epoch}") + + def test_continuous_aggregate_drop_last(self) -> None: + """With drop_last=True, partial batch at epoch end is discarded.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .aggregate(3, drop_last=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + # 5 items, batch_size=3, drop_last: only [0,1,2], partial [3,4] dropped + self.assertEqual(result, [[0, 1, 2]], f"epoch {epoch}") + + +class TestContinuousPipelineShutdown(unittest.TestCase): + def test_continuous_auto_stop_mid_epoch(self) -> None: + """auto_stop() exits cleanly even if mid-epoch.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(100), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + it = pipeline.get_iterator(timeout=5) + # Consume only a few items + self.assertEqual(next(it), 0) + self.assertEqual(next(it), 1) + # auto_stop exits here — must not hang + + def test_continuous_stop_between_epochs(self) -> None: + """stop() between epochs works cleanly.""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(3), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 1, 2]) + # auto_stop exits here after one epoch — must not hang + + def test_continuous_get_iterator_reuse(self) -> None: + """get_iterator() can be called multiple times within auto_stop().""" + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(3), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + r1 = list(pipeline.get_iterator(timeout=5)) + r2 = list(pipeline.get_iterator(timeout=5)) + r3 = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(r1, [0, 1, 2]) + self.assertEqual(r2, [0, 1, 2]) + self.assertEqual(r3, [0, 1, 2]) + + def test_continuous_stop_with_pipe_stage(self) -> None: + """Continuous pipeline with pipe stage can be stopped after epochs.""" + + def double(x: int) -> int: + return x * 2 + + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .pipe(double, concurrency=2) + .add_sink(buffer_size=3) + .build(num_threads=2) + ) + + with pipeline.auto_stop(): + r1 = sorted(pipeline.get_iterator(timeout=5)) + self.assertEqual(r1, [0, 2, 4, 6, 8]) + # auto_stop exits — pipeline has items buffered for next epoch + # stop() must drain and shut down cleanly + + @_ignore_fork_warning + def test_continuous_stop_with_subprocess_source(self) -> None: + """Continuous pipeline reading from subprocess can be stopped.""" + backend = ( + PipelineBuilder().add_source(SourceIterable(5)).add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + pipeline = ( + PipelineBuilder() + .add_source(source, continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + r1 = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(r1, [0, 1, 2, 3, 4]) + # auto_stop exits — must not hang on subprocess IPC + + @_ignore_fork_warning + def test_continuous_pipeline_in_subprocess_multi_epoch(self) -> None: + """A continuous pipeline running inside a subprocess supports + multi-epoch iteration without recreating the subprocess.""" + backend = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + # Iterate 3 epochs from the parent — subprocess is reused + for epoch in range(3): + result = list(source) + self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") + + @_ignore_fork_warning + def test_continuous_pipeline_in_subprocess_stop_mid_epoch(self) -> None: + """Subprocess with continuous pipeline can be abandoned mid-epoch.""" + backend = ( + PipelineBuilder() + .add_source(SourceIterable(100), continuous=True) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + it = iter(source) + self.assertEqual(next(it), 0) + self.assertEqual(next(it), 1) + # Abandon — must not hang + del it + del source + + @_ignore_fork_warning + def test_continuous_frontend_backend_multi_epoch(self) -> None: + """Frontend continuous pipeline on top of subprocess backend + supports multi-epoch iteration.""" + backend = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + pipeline = ( + PipelineBuilder() + .add_source(source, continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for epoch in range(3): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") + + @_ignore_fork_warning + def test_continuous_frontend_backend_stop_mid_epoch(self) -> None: + """Frontend continuous pipeline on top of subprocess backend + can be stopped mid-epoch without hanging.""" + backend = ( + PipelineBuilder() + .add_source(SourceIterable(100), continuous=True) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + pipeline = ( + PipelineBuilder() + .add_source(source, continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + it = pipeline.get_iterator(timeout=5) + self.assertEqual(next(it), 0) + self.assertEqual(next(it), 1) + # auto_stop exits mid-epoch — must not hang + + @_ignore_fork_warning + def test_continuous_frontend_backend_finalizer_shutdown(self) -> None: + """Frontend+backend pipeline cleaned up via weakref.finalize. + + Simulates the _SPDLDataLoader pattern: pipeline is started, iterated, + then the wrapper goes out of scope. The finalizer calls stop(timeout=10). + Must not hang. + """ + + backend = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + pipeline = ( + PipelineBuilder() + .add_source(source, continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + pipeline.start() + finalizer = weakref.finalize(pipeline, lambda p: p.stop(timeout=10), pipeline) + + # Iterate 2 epochs + for epoch in range(2): + result = list(pipeline.get_iterator(timeout=5)) + self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") + + t0 = time.monotonic() + # Trigger finalizer (simulates going out of scope) + finalizer() + elapsed = time.monotonic() - t0 + self.assertLess(elapsed, 15, f"finalizer took {elapsed:.1f}s — likely hung") + + +class CustomError(ValueError): + pass + + +class TestContinuousPipelineErrors(unittest.TestCase): + def test_continuous_source_failure(self) -> None: + """Source raising mid-epoch propagates error.""" + + class FailingSource: + def __iter__(self) -> Iterator[int]: + yield 0 + raise CustomError("source failed") + + pipeline = ( + PipelineBuilder() + .add_source(FailingSource(), continuous=True) + .add_sink(buffer_size=3) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=5)) + + def test_continuous_pipe_failure(self) -> None: + """Pipe function raising propagates error.""" + + def failing_fn(x: int) -> int: + if x == 2: + raise CustomError("pipe failed") + return x + + pipeline = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .pipe(failing_fn, concurrency=1) + .add_sink(buffer_size=3) + .build(num_threads=1, max_failures=0) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=5)) + + +class TestContinuousPipelineValidation(unittest.TestCase): + def test_mixed_continuous_mode_rejected(self) -> None: + """Mixing continuous and non-continuous sources raises ValueError.""" + plc_continuous = ( + PipelineBuilder() + .add_source(SourceIterable(3), continuous=True) + .add_sink() + .get_config() + ) + plc_normal = ( + PipelineBuilder().add_source(SourceIterable(3)).add_sink().get_config() + ) + + merged_config = PipelineConfig( + src=Merge([plc_continuous, plc_normal]), + pipes=[], + sink=SinkConfig(buffer_size=3), + ) + + with self.assertRaisesRegex(ValueError, "Mixed continuous mode"): + build_pipeline(merged_config, num_threads=1) + + +class _RecordIDs: + """Pipe function that records the thread ID and process ID.""" + + def __init__(self) -> None: + self.thread_ids: list[int] = [] + self.process_ids: list[int] = [] + + def __call__(self, x: int) -> int: + self.thread_ids.append(threading.get_ident()) + self.process_ids.append(os.getpid()) + return x + + +class TestContinuousSubprocessPipelineReuse(unittest.TestCase): + @_ignore_fork_warning + def test_subprocess_pipeline_reused_across_epochs(self) -> None: + """Thread and process IDs stay the same across epochs, proving + the pipeline is reused rather than rebuilt.""" + recorder = _RecordIDs() + + backend = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .pipe(recorder, concurrency=1) + .add_sink(buffer_size=3) + ) + source = run_pipeline_in_subprocess( + backend.get_config(), + num_threads=1, + timeout=10, + ) + + all_thread_ids: list[set[int]] = [] + all_process_ids: list[set[int]] = [] + + for epoch in range(3): + recorder.thread_ids.clear() + recorder.process_ids.clear() + result = list(source) + self.assertEqual(sorted(result), [0, 1, 2, 3, 4], f"epoch {epoch}") + all_thread_ids.append(set(recorder.thread_ids)) + all_process_ids.append(set(recorder.process_ids)) + + # All epochs should use the same thread(s) — pipeline was reused + self.assertEqual( + all_thread_ids[0], + all_thread_ids[1], + "Thread IDs changed between epoch 0 and 1 — pipeline was rebuilt", + ) + self.assertEqual( + all_thread_ids[1], + all_thread_ids[2], + "Thread IDs changed between epoch 1 and 2 — pipeline was rebuilt", + ) + + # All epochs should run in the same subprocess + self.assertEqual( + all_process_ids[0], + all_process_ids[1], + "Process IDs changed between epoch 0 and 1", + ) + self.assertEqual( + all_process_ids[1], + all_process_ids[2], + "Process IDs changed between epoch 1 and 2", + ) + + +@unittest.skipIf( + sys.version_info < (3, 14), + "Subinterpreters require Python 3.14+", +) +class TestContinuousSubinterpreterPipelineReuse(unittest.TestCase): + @_ignore_fork_warning + def test_subinterpreter_pipeline_reused_across_epochs(self) -> None: + """Thread and process IDs stay the same across epochs in subinterpreter.""" + recorder = _RecordIDs() + + config = ( + PipelineBuilder() + .add_source(SourceIterable(5), continuous=True) + .pipe(recorder, concurrency=1) + .add_sink(buffer_size=3) + .get_config() + ) + source = run_pipeline_in_subinterpreter( + config, + num_threads=1, + timeout=10, + ) + + all_thread_ids: list[set[int]] = [] + all_process_ids: list[set[int]] = [] + + for epoch in range(3): + recorder.thread_ids.clear() + recorder.process_ids.clear() + result = list(source) + self.assertEqual(sorted(result), [0, 1, 2, 3, 4], f"epoch {epoch}") + all_thread_ids.append(set(recorder.thread_ids)) + all_process_ids.append(set(recorder.process_ids)) + + # All epochs should use the same thread(s) — pipeline was reused + self.assertEqual( + all_thread_ids[0], + all_thread_ids[1], + "Thread IDs changed between epoch 0 and 1 — pipeline was rebuilt", + ) + self.assertEqual( + all_thread_ids[1], + all_thread_ids[2], + "Thread IDs changed between epoch 1 and 2 — pipeline was rebuilt", + ) + + # All epochs should run in the same process (subinterpreter shares process) + self.assertEqual( + all_process_ids[0], + all_process_ids[1], + "Process IDs changed between epoch 0 and 1", + ) diff --git a/src/spdl/pipeline/tests/defs_repr_test.py b/src/spdl/pipeline/tests/defs_repr_test.py new file mode 100644 index 000000000..20d08a258 --- /dev/null +++ b/src/spdl/pipeline/tests/defs_repr_test.py @@ -0,0 +1,355 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import inspect +import unittest +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence + +from spdl.pipeline import StageInfo +from spdl.pipeline.defs import ( + Aggregate, + Disaggregate, + Merge, + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, +) + + +# Test helper functions and classes +def example_sync_function(x: int) -> int: + """Example synchronous function for testing.""" + return x * 2 + + +async def example_async_function(x: int) -> int: + """Example async function for testing.""" + await asyncio.sleep(0) + return x * 2 + + +def _ln(target: object) -> int: + """Helper to get line number from inspect.getsourcelines.""" + return inspect.getsourcelines(target)[1] # pyre-ignore[6] + + +class ExampleIterable(Iterable[int]): + """Example iterable class for testing.""" + + def __iter__(self) -> Iterator[int]: + return iter([1, 2, 3]) + + +class ExampleAsyncIterable(AsyncIterable[int]): + """Example async iterable class for testing.""" + + async def __aiter__(self) -> AsyncIterator[int]: + for i in [1, 2, 3]: + yield i + + +async def custom_merge_op( + info: StageInfo, + input_queues: Sequence[asyncio.Queue], + output_queue: asyncio.Queue, +) -> None: + """Example custom merge operation for testing.""" + pass + + +class TestSourceConfigRepr(unittest.TestCase): + """Test SourceConfig.__repr__ with source location.""" + + def test_source_config_repr_with_iterable_class(self) -> None: + """Test __repr__ shows source location for iterable class.""" + # Setup: create source config with custom iterable + source = ExampleIterable() + config = SourceConfig(source=source) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains class name + # Note: for class instances, source location may not always be available + self.assertIn("ExampleIterable", result) + self.assertIn("SourceConfig", result) + + def test_source_config_repr_with_generator(self) -> None: + """Test __repr__ shows source location for generator function.""" + # Setup: create source config with generator + source = (x for x in range(10)) + config = SourceConfig(source=source) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains generator class name + self.assertIn("generator", result) + + def test_source_config_repr_with_async_iterable(self) -> None: + """Test __repr__ shows source location for async iterable.""" + # Setup: create source config with async iterable + source = ExampleAsyncIterable() + config = SourceConfig(source=source) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains class name + # Note: for class instances, source location may not always be available + self.assertIn("ExampleAsyncIterable", result) + self.assertIn("SourceConfig", result) + + +class TestPipeConfigRepr(unittest.TestCase): + """Test PipeConfig.__repr__ with source location.""" + + def test_pipe_config_repr_with_sync_function(self) -> None: + """Test __repr__ shows source location for sync function.""" + # Setup: create pipe config with sync function + config = Pipe(example_sync_function, concurrency=4) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains function name, concurrency, and source location + self.assertIn("concurrency=4", result) + self.assertIn("example_sync_function", result) + self.assertIn(__file__, result) + self.assertIn(f":{_ln(example_sync_function)}", result) + + def test_pipe_config_repr_with_async_function(self) -> None: + """Test __repr__ shows source location for async function.""" + # Setup: create pipe config with async function + config = Pipe(example_async_function, concurrency=2) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains function name, concurrency, and source location + self.assertIn("concurrency=2", result) + self.assertIn("example_async_function", result) + self.assertIn(__file__, result) + self.assertIn(f":{_ln(example_async_function)}", result) + + def test_pipe_config_repr_with_lambda(self) -> None: + """Test __repr__ handles lambda functions gracefully.""" + # Setup: create pipe config with lambda + config = Pipe(lambda x: x * 2, concurrency=1) + + # Execute: get repr + result = repr(config) + + # Assert: repr contains lambda and concurrency + self.assertIn("concurrency=1", result) + self.assertIn("lambda", result) + + def test_pipe_config_repr_without_source_location(self) -> None: + """Test __repr__ handles cases where source location cannot be determined.""" + # Setup: create pipe config with built-in function + config = Pipe(len, concurrency=1) + + # Execute: get repr + result = repr(config) + + # Assert: repr still works and contains concurrency + self.assertIn("concurrency=1", result) + self.assertIn("len", result) + + +class TestMergeConfigRepr(unittest.TestCase): + """Test MergeConfig.__repr__ with nested pipelines.""" + + def test_merge_config_repr_with_two_pipelines(self) -> None: + """Test __repr__ shows nested pipeline configs with proper indentation.""" + # Setup: create two pipeline configs and merge them + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=20), + ) + merge_config = Merge([plc1, plc2]) + + # Execute: get repr + result = repr(merge_config) + + # Assert: repr contains merge structure with both pipelines + self.assertIn("MergeConfig(", result) + self.assertIn("Pipeline 1:", result) + self.assertIn("Pipeline 2:", result) + self.assertIn("PipelineConfig", result) + + def test_merge_config_repr_with_custom_op(self) -> None: + """Test __repr__ shows custom merge operation with source location.""" + # Setup: create merge config with custom op + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + merge_config = Merge([plc1, plc2], op=custom_merge_op) + + # Execute: get repr + result = repr(merge_config) + + # Assert: repr contains op info with source location + self.assertIn("op=", result) + self.assertIn("custom_merge_op", result) + self.assertIn(__file__, result) + self.assertIn(f":{_ln(custom_merge_op)}", result) + + def test_merge_config_repr_multiline_structure(self) -> None: + """Test __repr__ creates multi-line output with proper indentation.""" + # Setup: create merge config with pipes + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[Pipe(example_sync_function, concurrency=2)], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + merge_config = Merge([plc1, plc2]) + + # Execute: get repr + result = repr(merge_config) + + # Assert: result is multi-line and properly indented + lines = result.split("\n") + self.assertGreater(len(lines), 5) # Multi-line output + # Check some lines have proper indentation + pipeline_lines = [line for line in lines if "Pipeline" in line] + self.assertGreater(len(pipeline_lines), 0) + + +class TestPipelineConfigRepr(unittest.TestCase): + """Test PipelineConfig.__repr__ with MergeConfig source.""" + + def test_pipeline_config_repr_with_source_config(self) -> None: + """Test __repr__ with SourceConfig shows inline representation.""" + # Setup: create pipeline config with simple source + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[Pipe(example_sync_function, concurrency=4)], + sink=SinkConfig(buffer_size=10), + ) + + # Execute: get repr + result = repr(config) + + # Assert: repr shows source inline + self.assertIn("PipelineConfig", result) + self.assertIn("Source:", result) + self.assertIn("Pipes:", result) + self.assertIn("Sink:", result) + self.assertIn("example_sync_function", result) + + def test_pipeline_config_repr_with_merge_config(self) -> None: + """Test __repr__ with MergeConfig shows proper indentation.""" + # Setup: create pipeline config with merge source + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + merge_config = Merge([plc1, plc2]) + final_config = PipelineConfig( + src=merge_config, + pipes=[Pipe(example_sync_function, concurrency=2)], + sink=SinkConfig(buffer_size=100), + ) + + # Execute: get repr + result = repr(final_config) + + # Assert: repr shows nested structure with proper indentation + self.assertIn("PipelineConfig", result) + self.assertIn("Source:", result) + self.assertIn("MergeConfig(", result) + self.assertIn("Pipeline 1:", result) + self.assertIn("Pipeline 2:", result) + # Check final pipe appears after merge + self.assertIn("example_sync_function", result) + + def test_pipeline_config_repr_indentation_hierarchy(self) -> None: + """Test __repr__ maintains correct indentation hierarchy.""" + # Setup: create nested pipeline config + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[Pipe(example_sync_function, concurrency=1)], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + merge_config = Merge([plc1, plc2]) + final_config = PipelineConfig( + src=merge_config, + pipes=[ + Pipe(example_async_function, concurrency=4), + Aggregate(5), + Disaggregate(), + ], + sink=SinkConfig(buffer_size=100), + ) + + # Execute: get repr + result = repr(final_config) + + # Assert: verify indentation levels exist + lines = result.split("\n") + # Should have various indentation levels + has_no_indent = any(line and not line[0].isspace() for line in lines) + has_some_indent = any(line.startswith(" ") for line in lines) + has_more_indent = any(line.startswith(" ") for line in lines) + + self.assertTrue(has_no_indent, "Should have lines with no indentation") + self.assertTrue(has_some_indent, "Should have lines with 2-space indentation") + self.assertTrue(has_more_indent, "Should have lines with 4+ space indentation") + + def test_pipeline_config_repr_with_aggregate_disaggregate(self) -> None: + """Test __repr__ shows aggregate and disaggregate pipes correctly.""" + # Setup: create pipeline config with various pipe types + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + Pipe(example_sync_function, concurrency=2), + Aggregate(10, drop_last=True), + Disaggregate(), + ], + sink=SinkConfig(buffer_size=10), + ) + + # Execute: get repr + result = repr(config) + + # Assert: repr shows all pipe types + self.assertIn("example_sync_function", result) + self.assertIn("aggregate", result) + self.assertIn("disaggregate", result) diff --git a/src/spdl/pipeline/tests/failure_rate_test.py b/src/spdl/pipeline/tests/failure_rate_test.py new file mode 100644 index 000000000..1e6ecb12f --- /dev/null +++ b/src/spdl/pipeline/tests/failure_rate_test.py @@ -0,0 +1,804 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import functools +import unittest +import warnings +from collections.abc import Callable +from fractions import Fraction +from typing import Type, TypeVar + +from parameterized import parameterized +from spdl.pipeline import PipelineBuilder, PipelineFailure + +_F = TypeVar("_F", bound=Callable[..., object]) +_C = TypeVar("_C", bound=Type[object]) + + +def _ignore_intentional_warnings(fn: _F) -> _F: + """Suppress warnings emitted intentionally by failure-rate tests: + + - the fork() multi-threaded DeprecationWarning from the subprocess + pipeline machinery, and + - "coroutine ... was never awaited" RuntimeWarnings, which surface when + the pipeline is forced to fail mid-iteration and the source coroutine + is dropped before being fully consumed. + """ + + @functools.wraps(fn) + def wrapper(*args: object, **kwargs: object) -> object: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + warnings.filterwarnings( + "ignore", + message="coroutine .* was never awaited", + category=RuntimeWarning, + ) + return fn(*args, **kwargs) + + # pyre-ignore[7] + return wrapper + + +def _ignore_intentional_warnings_in_class(cls: _C) -> _C: + for name, member in list(vars(cls).items()): + if name.startswith("test_") and callable(member): + setattr(cls, name, _ignore_intentional_warnings(member)) + return cls + + +@_ignore_intentional_warnings_in_class +class PipelineFailureRateTest(unittest.TestCase): + """Tests for Fraction-based failure rate thresholds in SPDL pipeline. + + Key design: A fixed probation period of 100 invocations is used before + rate-based checking kicks in. This prevents early false positives when + sample size is too small to be statistically meaningful. + + The pipeline stops when failure rate strictly exceeds the threshold (>). + """ + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_basic(self, output_order: str) -> None: + """Pipeline fails when failure rate exceeds Fraction threshold. + + Uses 1000 items. Fails on multiples of 100 (10 out of 1000 = 1%). + With 0.1% threshold and fixed probation of 100, should fail. + After probation: rate = 1% > 0.1% threshold -> fails. + """ + + def fail_on_hundred(x: int) -> int: + if x % 100 == 0: # Fails on 0, 100, 200, ..., 900 (10 out of 1000) + raise ValueError(f"Multiple of 100: {x}") + return x + + # 0.1% threshold - should fail because actual rate is 1% + pipeline = ( + PipelineBuilder() + .add_source(range(1000)) + .pipe(fail_on_hundred, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 1000)) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=30)) + + all_expected = {x for x in range(1000) if x % 100 != 0} + self.assertTrue(len(vals) > 0) + self.assertTrue(set(vals).issubset(all_expected)) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_passes(self, output_order: str) -> None: + """Pipeline succeeds when failure rate stays below threshold. + + Uses 100 items. Fails on multiples of 10 (10 failures = 10%). + With 15% threshold (Fraction(3, 20)) and fixed probation of 100. + After probation: rate = 10% < 15% -> succeeds. + """ + + def fail_on_ten(x: int) -> int: + if x % 10 == 0: # Fails on 0, 10, 20, ..., 90 (10 out of 100 = 10%) + raise ValueError(f"Multiple of 10: {x}") + return x + + # Allow 15% failure rate - should succeed because actual rate is 10% + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_on_ten, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(3, 20)) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 10 + expected = [x for x in range(100) if x % 10 != 0] + self.assertEqual(expected, vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_just_below_threshold( + self, output_order: str + ) -> None: + """Pipeline succeeds when failure rate is just below threshold. + + Uses 100 items, fails on items >= 90 (10 failures = 10%). + With 11% threshold and fixed probation of 100. + After probation: rate = 10% < 11% -> succeeds. + """ + + def fail_late(x: int) -> int: + if x >= 90: # Fails on 90-99 (10 out of 100 = 10%) + raise ValueError(f"Item >= 90: {x}") + return x + + # Allow 11% failure rate - should succeed because actual rate is 10% + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_late, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(11, 100)) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get values 0-89 (90 successful items) + self.assertEqual(list(range(90)), vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_above_threshold(self, output_order: str) -> None: + """Pipeline fails when failure rate exceeds threshold. + + Uses 100 items, fails on items >= 90 (10 failures = 10%). + With 9% threshold and fixed probation of 100. + After probation: rate = 10% > 9% -> fails. + """ + + def fail_late(x: int) -> int: + if x >= 90: # Fails on 90-99 (10 out of 100 = 10%) + raise ValueError(f"Item >= 90: {x}") + return x + + # 9% threshold - should fail because actual rate is 10% + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_late, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(9, 100)) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get values 0-89 (90 successful items) + self.assertEqual(list(range(90)), vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_probation_period(self, output_order: str) -> None: + """Early failures don't trigger threshold due to fixed probation period (100). + + Uses 20 items. Fails on multiples of 5 (4 out of 20 = 20%). + With 1% threshold but fixed probation of 100, only 20 items processed. + Since probation never completes, pipeline succeeds despite 20% > 1%. + """ + + def fail_on_five(x: int) -> int: + if x % 5 == 0: # Fails on 0, 5, 10, 15 (4 out of 20 = 20%) + raise ValueError(f"Multiple of 5: {x}") + return x + + # 1% threshold with fixed probation=100 + # Since we only have 20 items, probation never completes + pipeline = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(fail_on_five, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 100)) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 5: 1,2,3,4,6,7,8,9,11,12,13,14,16,17,18,19 + expected = [x for x in range(20) if x % 5 != 0] + self.assertEqual(expected, vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_pipe_override(self, output_order: str) -> None: + """Per-pipe Fraction override works correctly. + + Uses 100 items. Fails on multiples of 7 (15 out of 100 = 15%). + Pipe-level: 11% threshold - should fail (15% > 11%). + Global: 19% threshold - would pass. + Since pipe-level is stricter, pipeline should fail. + """ + + def fail_on_seven(x: int) -> int: + if x % 7 == 0: # Fails on 0, 7, 14, ..., 98 (15 out of 100 = 15%) + raise ValueError(f"Multiple of 7: {x}") + return x + + # Global allows 19% but pipe-level restricts to 11% - should fail + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe( + fail_on_seven, + output_order=output_order, + max_failures=Fraction(11, 100), + ) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(19, 100)) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 7 + expected = [x for x in range(100) if x % 7 != 0] + self.assertEqual(expected, vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_multiple_stages(self, output_order: str) -> None: + """Multiple stages with different Fraction thresholds. + + First stage: fails on odd numbers (50%) but allowed unlimited. + Second stage: receives 50 even numbers, fails on divisible by 12 (9/50 = 18%). + With 22% threshold (Fraction(11, 50)), 18% < 22% -> should pass. + Note: probation is fixed at 100, but only 50 items reach second stage. + """ + + def fail_odd(x: int) -> int: + if x % 2: + raise ValueError(f"Odd number: {x}") + return x + + def fail_twelve(x: int) -> int: + if (x % 12) == 0: + raise ValueError(f"Divisible by 12: {x}") + return x + + # First stage fails 50% (odd numbers) but allowed unlimited failures + # Second stage: 18% failure rate with 22% threshold -> should succeed + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_odd, output_order=output_order, max_failures=-1) + .pipe(fail_twelve, output_order=output_order, max_failures=Fraction(11, 50)) + .add_sink(1) + .build(num_threads=1, max_failures=-1) + ) + + # Second stage receives 50 even numbers (0,2,4,...98) + # Fails on 0,12,24,36,48,60,72,84,96 = 9 failures out of 50 = 18% + # With 22% threshold, should succeed + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Even numbers not divisible by 12: 2,4,6,8,10,14,16,... + expected = [x for x in range(100) if x % 2 == 0 and x % 12 != 0] + self.assertEqual(expected, vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_vs_count(self, output_order: str) -> None: + """Verify int count-based behavior unchanged. + + Fails on multiples of 10 (10 failures). + Count-based: Allow 15 failures - should succeed. + Count-based: Allow 5 failures - should fail. + """ + + def fail_on_ten(x: int) -> int: + if x % 10 == 0: # Fails on 0, 10, 20, ..., 90 (10 failures) + raise ValueError(f"Multiple of 10: {x}") + return x + + # Count-based: Allow 15 failures - should succeed + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_on_ten, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=15) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 10 + expected = [x for x in range(100) if x % 10 != 0] + self.assertEqual(expected, vals) + + # Count-based: Allow 5 failures - should fail after 5 failures + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_on_ten, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=5) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Pipeline stops early after 5 failures; vals is a subset of expected + all_expected = {x for x in range(100) if x % 10 != 0} + self.assertTrue(len(vals) > 0) + self.assertTrue(set(vals).issubset(all_expected)) + + def test_pipeline_failure_rate_ordered_pipe(self) -> None: + """Test with output_order='input'. + + Uses 100 items. Fails on multiples of 7 (15 out of 100 = 15%). + With 11% threshold and fixed probation of 100. + After probation: rate = 15% > 11% -> should fail. + """ + + def fail_on_seven(x: int) -> int: + if x % 7 == 0: # Fails on 0, 7, 14, ..., 98 (15 out of 100 = 15%) + raise ValueError(f"Multiple of 7: {x}") + return x + + # 11% threshold - should fail because actual rate is ~15% + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_on_seven, output_order="input", max_failures=Fraction(11, 100)) + .add_sink(1) + .build(num_threads=1) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 7 + expected = [x for x in range(100) if x % 7 != 0] + self.assertEqual(expected, vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_high_concurrency(self, output_order: str) -> None: + """Test failure rate with high concurrency. + + Uses 100 items. Fails on multiples of 5 (20 out of 100 = 20%). + With 15% threshold (Fraction(3, 20)) and fixed probation of 100. + After probation: rate = 20% > 15% -> should fail. + With high concurrency, exact processing order is nondeterministic. + """ + + def fail_on_five(x: int) -> int: + if x % 5 == 0: # Fails on 0, 5, 10, ..., 95 (20 out of 100 = 20%) + raise ValueError(f"Multiple of 5: {x}") + return x + + # 15% threshold - should fail because actual rate is 20% + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe( + fail_on_five, + output_order=output_order, + concurrency=5, + max_failures=Fraction(3, 20), + ) + .add_sink(1) + .build(num_threads=5, max_failures=Fraction(1, 2)) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_zero_failures(self, output_order: str) -> None: + """Test with zero failures. + + 0% failure rate with 10% threshold - should succeed. + """ + + def no_fail(x: int) -> int: + return x * 2 + + # 0% failure rate with 10% threshold - should succeed + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(no_fail, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 10)) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all values doubled + self.assertEqual([x * 2 for x in range(100)], vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_failure_rate_under_probation_always_succeeds( + self, output_order: str + ) -> None: + """With fixed probation of 100, pipelines with < 100 items always succeed. + + Uses 10 items with 40% failure rate (multiples of 3 fail). + With 30% threshold but fixed probation of 100. + Since only 10 items processed, probation not reached -> succeeds. + """ + + def fail_on_three(x: int) -> int: + if x % 3 == 0: # Fails on 0, 3, 6, 9 (4 out of 10 = 40%) + raise ValueError(f"Multiple of 3: {x}") + return x + + # 30% threshold, but only 10 items (under probation) + # Should succeed because probation (100) not reached + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_on_three, output_order=output_order) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(3, 10)) + ) + + # Should succeed despite 40% > 30% because probation not complete + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 3: 1, 2, 4, 5, 7, 8 + expected = [x for x in range(10) if x % 3 != 0] + self.assertEqual(expected, vals) + + def test_pipeline_failure_rate_invalid_fraction_zero(self) -> None: + """Building pipeline with Fraction <= 0 raises ValueError.""" + + def noop(x: int) -> int: + return x + + # Zero Fraction should raise ValueError + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( + num_threads=1, max_failures=Fraction(0, 100) + ) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + def test_pipeline_failure_rate_invalid_fraction_negative(self) -> None: + """Building pipeline with negative Fraction raises ValueError.""" + + def noop(x: int) -> int: + return x + + # Negative Fraction should raise ValueError + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( + num_threads=1, max_failures=Fraction(-1, 10) + ) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + def test_pipeline_failure_rate_invalid_fraction_greater_than_one(self) -> None: + """Building pipeline with Fraction > 1 raises ValueError.""" + + def noop(x: int) -> int: + return x + + # Fraction > 1 (e.g., 150%) should raise ValueError + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( + num_threads=1, max_failures=Fraction(15, 10) + ) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + def test_pipeline_failure_rate_invalid_pipe_fraction_zero(self) -> None: + """Building pipeline with zero Fraction at pipe level raises ValueError.""" + + def noop(x: int) -> int: + return x + + # Zero Fraction at pipe level should raise ValueError + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe( + noop, max_failures=Fraction(0, 100) + ).add_sink(1).build(num_threads=1) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + def test_pipeline_failure_rate_invalid_pipe_fraction_greater_than_one(self) -> None: + """Building pipeline with Fraction > 1 at pipe level raises ValueError.""" + + def noop(x: int) -> int: + return x + + # Fraction > 1 at pipe level should raise ValueError + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe( + noop, max_failures=Fraction(200, 100) + ).add_sink(1).build(num_threads=1) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + def test_pipeline_failure_rate_valid_fraction_one(self) -> None: + """Fraction(1, 1) = 100% is valid (allows all failures). + + With > comparison, rate can never exceed 100%, so pipeline always succeeds. + """ + + def always_fail(x: int) -> int: + raise ValueError(f"Always fail: {x}") + + # 100% failure rate threshold - should never fail + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(always_fail) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 1)) + ) + + # Should complete without PipelineFailure (100% failures allowed) + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # No items should pass through since all fail + self.assertEqual([], vals) + + def test_pipeline_failure_rate_valid_small_fraction(self) -> None: + """Very small Fraction like Fraction(1, 1000) is valid.""" + + def fail_on_hundred(x: int) -> int: + if x % 100 == 0: # Fails on 0, 100, 200, ..., 900 (10 out of 1000 = 1%) + raise ValueError(f"Multiple of 100: {x}") + return x + + # 0.1% threshold (Fraction(1, 1000)) - should fail because actual rate is 1% + pipeline = ( + PipelineBuilder() + .add_source(range(1000)) + .pipe(fail_on_hundred) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 1000)) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=30)) + + all_expected = {x for x in range(1000) if x % 100 != 0} + self.assertTrue(len(vals) > 0) + self.assertTrue(set(vals).issubset(all_expected)) + + @parameterized.expand( + [ + (Fraction(0, 1), "zero numerator"), + (Fraction(0, 100), "zero with large denominator"), + (Fraction(-1, 10), "negative numerator"), + (Fraction(1, -10), "negative denominator"), + (Fraction(-1, -10), "double negative (positive > 0 but > 1)"), + ] + ) + def test_pipeline_failure_rate_invalid_fractions_parameterized( + self, fraction: Fraction, description: str + ) -> None: + """Parameterized test for various invalid Fraction values.""" + + def noop(x: int) -> int: + return x + + # Note: Fraction(-1, -10) normalizes to Fraction(1, 10) which is valid + # But Fraction(0, x) and negative fractions should fail + if fraction <= 0 or fraction > 1: + with self.assertRaises(ValueError): + PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( + num_threads=1, max_failures=fraction + ) + else: + # This should not raise - it's a valid fraction + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(noop) + .add_sink(1) + .build(num_threads=1, max_failures=fraction) + ) + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(list(range(10)), vals) + + @parameterized.expand( + [ + (Fraction(11, 10), "110%"), + (Fraction(2, 1), "200%"), + (Fraction(150, 100), "150%"), + (Fraction(101, 100), "101%"), + ] + ) + def test_pipeline_failure_rate_invalid_fractions_greater_than_one_parameterized( + self, fraction: Fraction, description: str + ) -> None: + """Parameterized test for Fraction values > 1 (greater than 100%).""" + + def noop(x: int) -> int: + return x + + with self.assertRaises(ValueError) as ctx: + PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( + num_threads=1, max_failures=fraction + ) + + self.assertIn("must be in range (0, 1]", str(ctx.exception)) + + @parameterized.expand( + [ + (Fraction(1, 100), "1%"), + (Fraction(1, 10), "10%"), + (Fraction(1, 2), "50%"), + (Fraction(99, 100), "99%"), + (Fraction(1, 1), "100%"), + ] + ) + def test_pipeline_failure_rate_valid_fractions_parameterized( + self, fraction: Fraction, description: str + ) -> None: + """Parameterized test for valid Fraction values in range (0, 1].""" + + def noop(x: int) -> int: + return x + + # Should not raise - these are all valid fractions + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(noop) + .add_sink(1) + .build(num_threads=1, max_failures=fraction) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # All values should pass through + self.assertEqual(list(range(100)), vals) + + def test_pipeline_failure_rate_probation_prevents_early_trigger(self) -> None: + """Probation period (fixed at 100) prevents triggering even when rate is high. + + Uses 10 items, fails on multiples of 5 (2/10 = 20%). + With 1% threshold but fixed probation of 100, only 10 items processed. + Since probation not reached, pipeline succeeds despite 20% > 1%. + """ + + def fail_on_five(x: int) -> int: + if x % 5 == 0: # Fails on 0, 5 (2 out of 10 = 20%) + raise ValueError(f"Multiple of 5: {x}") + return x + + # 1% threshold with fixed probation=100. Only 10 items -> no check runs. + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_on_five) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 100)) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 5: 1, 2, 3, 4, 6, 7, 8, 9 + expected = [x for x in range(10) if x % 5 != 0] + self.assertEqual(expected, vals) + + def test_pipeline_failure_rate_probation_triggers_after_warmup(self) -> None: + """Once probation period (100) is met, threshold check runs. + + Uses 100 items, fails on multiples of 5 (20/100 = 20%). + With 10% threshold and fixed probation of 100. + After 100 invocations: rate = 20% > 10% -> fails. + """ + + def fail_on_five(x: int) -> int: + if x % 5 == 0: # Fails on 0, 5, ..., 95 (20 out of 100 = 20%) + raise ValueError(f"Multiple of 5: {x}") + return x + + # 10% threshold with fixed probation=100. All 100 items are processed. + pipeline = ( + PipelineBuilder() + .add_source(range(100)) + .pipe(fail_on_five) + .add_sink(1) + .build(num_threads=1, max_failures=Fraction(1, 10)) + ) + + vals = [] + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + # Should get all non-multiples of 5 + expected = [x for x in range(100) if x % 5 != 0] + self.assertEqual(expected, vals) diff --git a/src/spdl/pipeline/tests/merge_config_test.py b/src/spdl/pipeline/tests/merge_config_test.py new file mode 100644 index 000000000..22917001e --- /dev/null +++ b/src/spdl/pipeline/tests/merge_config_test.py @@ -0,0 +1,475 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for MergeConfig class.""" + +# pyre-strict + +import asyncio +import unittest +from collections.abc import Sequence + +from spdl.pipeline import build_pipeline, create_task, is_eof, StageInfo +from spdl.pipeline.defs import ( + Merge, + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, +) + + +class MergeConfigTest(unittest.TestCase): + """Test MergeConfig functionality.""" + + def test_merge_config_with_two_simple_pipelines(self) -> None: + """Test MergeConfig merges outputs from two simple pipelines.""" + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 6) + self.assertCountEqual(results, [1, 2, 3, 4, 5, 6]) + + def test_merge_config_with_processed_pipelines(self) -> None: + """Test MergeConfig merges outputs from pipelines with processing.""" + double_pipe = Pipe(lambda x: x * 2) + add_ten_pipe = Pipe(lambda x: x + 10) + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[double_pipe], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[add_ten_pipe], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 6) + # Pipeline 1: [1, 2, 3] -> [2, 4, 6] (doubled) + # Pipeline 2: [4, 5, 6] -> [14, 15, 16] (added 10) + self.assertCountEqual(results, [2, 4, 6, 14, 15, 16]) + + def test_merge_config_with_multiple_pipelines(self) -> None: + """Test MergeConfig can merge outputs from three pipelines.""" + plc1 = PipelineConfig( + src=SourceConfig([1, 2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([10, 20]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc3 = PipelineConfig( + src=SourceConfig([100, 200]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2, plc3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=3) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 6) + self.assertCountEqual(results, [1, 2, 10, 20, 100, 200]) + + def test_merge_config_with_post_processing(self) -> None: + """Test MergeConfig output can be further processed in main pipeline.""" + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + multiply_by_5_pipe = Pipe(lambda x: x * 5) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[multiply_by_5_pipe], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 6) + self.assertCountEqual(results, [5, 10, 15, 20, 25, 30]) + + def test_merge_config_with_async_processing(self) -> None: + """Test MergeConfig works with async processing functions.""" + + async def async_double(x: int) -> int: + await asyncio.sleep(0.01) # Small delay to simulate async work + return x * 2 + + async_pipe = Pipe(async_double) + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[async_pipe], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[async_pipe], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=5)) + + self.assertEqual(len(results), 6) + self.assertCountEqual(results, [2, 4, 6, 8, 10, 12]) + + def test_merge_config_with_different_data_types(self) -> None: + """Test MergeConfig works with different data types.""" + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig(["a", "b", "c"]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 6) + self.assertCountEqual(results, [1, 2, 3, "a", "b", "c"]) + + def test_merge_config_with_empty_pipeline(self) -> None: + """Test MergeConfig handles pipeline with empty source.""" + plc1 = PipelineConfig( + src=SourceConfig([]), # Empty source + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 3) + self.assertCountEqual(results, [1, 2, 3]) + + def test_merge_config_validation_empty_list(self) -> None: + """Test MergeConfig validation fails with empty pipeline list.""" + with self.assertRaises(ValueError) as cm: + Merge([]) + + self.assertIn("at least one upstream pipeline", str(cm.exception)) + + def test_merge_config_with_aggregation(self) -> None: + """Test MergeConfig works with aggregation operations.""" + from spdl.pipeline.defs import Aggregate, Disaggregate + + aggregate_pipe = Aggregate(2) + disaggregate_pipe = Disaggregate() + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3, 4]), + pipes=[aggregate_pipe, disaggregate_pipe], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([10, 20, 30, 40]), + pipes=[aggregate_pipe, disaggregate_pipe], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 8) + self.assertCountEqual(results, [1, 2, 3, 4, 10, 20, 30, 40]) + + def test_merge_config_with_concurrency(self) -> None: + """Test MergeConfig works with concurrent processing.""" + concurrent_pipe = Pipe(lambda x: x + 100, concurrency=3) + + plc1 = PipelineConfig( + src=SourceConfig(list(range(10))), + pipes=[concurrent_pipe], + sink=SinkConfig(buffer_size=20), + ) + + plc2 = PipelineConfig( + src=SourceConfig(list(range(50, 60))), + pipes=[concurrent_pipe], + sink=SinkConfig(buffer_size=20), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[], + sink=SinkConfig(buffer_size=50), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=4) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 20) + expected = list(range(100, 110)) + list(range(150, 160)) + self.assertCountEqual(results, expected) + + def test_merge_config_with_different_pipe_counts_and_post_processing(self) -> None: + """Test MergeConfig merges pipelines with different numbers of pipes and applies post-processing.""" + + # Pipeline 1: single pipe (multiply by 2) + multiply_pipe = Pipe(lambda x: x * 2) + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[multiply_pipe], + sink=SinkConfig(buffer_size=10), + ) + + # Pipeline 2: three pipes (add 10, multiply by 3, subtract 5) + add_ten_pipe = Pipe(lambda x: x + 10) + multiply_by_three_pipe = Pipe(lambda x: x * 3) + subtract_five_pipe = Pipe(lambda x: x - 5) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5]), + pipes=[add_ten_pipe, multiply_by_three_pipe, subtract_five_pipe], + sink=SinkConfig(buffer_size=10), + ) + + # Add post-processing after merge: add 100 to all merged results + post_process_pipe = Pipe(lambda x: x + 100) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[post_process_pipe], + sink=SinkConfig(buffer_size=20), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(len(results), 5) + + # [1, 2, 3] --(multiply by 2)--> [2, 4, 6] --(add 100)--> [102, 104, 106] + pipeline1_expected = [102, 104, 106] + + # [4, 5] --(add 10)--> [14, 15] + # --(multiply by 3)--> [42, 45] + # --(subtract 5)--> [37, 40] + # --(add 100)--> [137, 140] + pipeline2_expected = [137, 140] + + expected_all = pipeline1_expected + pipeline2_expected + self.assertCountEqual(results, expected_all) + + def test_merge_config_with_custom_merge_op(self) -> None: + """Test MergeConfig accepts and uses custom merge operation.""" + + # Track which pipelines contributed items (for verification) + collected_items: list[str] = [] + + async def custom_merge_op( + info: StageInfo, + input_queues: Sequence[asyncio.Queue[object]], + output_queue: asyncio.Queue[object], + ) -> None: + """Custom merge that adds a prefix to each item based on its source pipeline.""" + + async def process_queue( + queue_idx: int, in_q: asyncio.Queue[object] + ) -> None: + while True: + item = await in_q.get() + if is_eof(item): + return + # Add prefix based on source pipeline + prefixed_item = f"p{queue_idx}_{item}" + collected_items.append(prefixed_item) + await output_queue.put(prefixed_item) + + tasks = [ + create_task(process_queue(i, in_q), name=f"{info}:{i}") + for i, in_q in enumerate(input_queues) + ] + await asyncio.wait(tasks) + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2], op=custom_merge_op), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + # Verify we got all items with prefixes + self.assertEqual(len(results), 6) + # Items from pipeline 0 should have p0_ prefix + self.assertIn("p0_1", results) + self.assertIn("p0_2", results) + self.assertIn("p0_3", results) + # Items from pipeline 1 should have p1_ prefix + self.assertIn("p1_4", results) + self.assertIn("p1_5", results) + self.assertIn("p1_6", results) + + def test_merge_config_with_custom_merge_op_early_exit(self) -> None: + """Test MergeConfig accepts and uses custom merge operation.""" + + async def custom_merge_op( + _: StageInfo, + input_queues: Sequence[asyncio.Queue[object]], + output_queue: asyncio.Queue[object], + ) -> None: + """Custom merge that exists when one sub-pipeline completes""" + while True: + for in_q in input_queues: + item = await in_q.get() + if is_eof(item): + return + + await output_queue.put(item) + + plc1 = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + plc2 = PipelineConfig( + src=SourceConfig([4, 5, 6, 7, 8, 9]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + main_pipeline_config = PipelineConfig( + src=Merge([plc1, plc2], op=custom_merge_op), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + + pipeline = build_pipeline(main_pipeline_config, num_threads=2) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=3)) + + self.assertEqual(results, [1, 4, 2, 5, 3, 6]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/spdl/pipeline/tests/path_variants_test.py b/src/spdl/pipeline/tests/path_variants_test.py new file mode 100644 index 000000000..c6c387d93 --- /dev/null +++ b/src/spdl/pipeline/tests/path_variants_test.py @@ -0,0 +1,704 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for PathVariants feature.""" + +# pyre-unsafe + +import asyncio +import unittest + +from spdl.pipeline import build_pipeline +from spdl.pipeline._components._node import PipelineFailure +from spdl.pipeline.defs import ( + Aggregate, + Disaggregate, + Merge, + PathVariants, + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, +) + + +def _run_pipeline(config, num_threads=2, timeout=5): + """Helper to build, run, and collect results from a pipeline.""" + pipeline = build_pipeline(config, num_threads=num_threads) + with pipeline.auto_stop(): + return list(pipeline.get_iterator(timeout=timeout)) + + +async def _slow_source(items, delay=0.1): + """Yield items with a delay between them. + + Each item is fully processed before the next one is dispatched, + making cross-path ordering deterministic (matching input order). + """ + for item in items: + yield item + await asyncio.sleep(delay) + + +class PathVariantsBasicTest(unittest.TestCase): + """Basic functional tests for PathVariants.""" + + def test_basic_routing(self) -> None: + """Even items to path 0 (double), odd items to path 1 (add 100).""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(6))), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 2)], # path 0: evens + [Pipe(lambda x: x + 100)], # path 1: odds + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # Slow source ensures each item is processed before the next arrives, + # so output order matches input order: + # 0→path0→0, 1→path1→101, 2→path0→4, 3→path1→103, 4→path0→8, 5→path1→105 + self.assertEqual(results, [0, 101, 4, 103, 8, 105]) + + def test_async_router(self) -> None: + """Async router function works correctly.""" + + async def async_router(x: int) -> int: + return x % 2 + + config = PipelineConfig( + src=SourceConfig(_slow_source(range(6))), + pipes=[ + PathVariants( + router=async_router, + paths=[ + [Pipe(lambda x: x * 2)], + [Pipe(lambda x: x + 100)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, [0, 101, 4, 103, 8, 105]) + + def test_async_callable_router(self) -> None: + """Class with async __call__ works as router.""" + + class AsyncRouter: + async def __call__(self, x: int) -> int: + return x % 2 + + config = PipelineConfig( + src=SourceConfig(_slow_source(range(6))), + pipes=[ + PathVariants( + router=AsyncRouter(), + paths=[ + [Pipe(lambda x: x * 2)], + [Pipe(lambda x: x + 100)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, [0, 101, 4, 103, 8, 105]) + + def test_all_to_one_path(self) -> None: + """Router always returns 0 — all items go to path 0.""" + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + PathVariants( + router=lambda x: 0, + paths=[ + [Pipe(lambda x: x * 10)], + [Pipe(lambda x: x + 1000)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, [10, 20, 30]) + + def test_multiple_paths_different_processing(self) -> None: + """3 paths with different transforms.""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(9))), + pipes=[ + PathVariants( + router=lambda x: x % 3, + paths=[ + [Pipe(lambda x: x * 2)], # path 0: ×2 + [Pipe(lambda x: x + 10)], # path 1: +10 + [Pipe(lambda x: -x)], # path 2: negate + ], + ), + ], + sink=SinkConfig(buffer_size=20), + ) + results = _run_pipeline(config) + # Slow source ensures interleaved input order: + # 0→path0→0, 1→path1→11, 2→path2→-2, + # 3→path0→6, 4→path1→14, 5→path2→-5, + # 6→path0→12, 7→path1→17, 8→path2→-8 + self.assertEqual(results, [0, 11, -2, 6, 14, -5, 12, 17, -8]) + + def test_identity_path_passthrough(self) -> None: + """A path with an identity pipe passes items through unchanged.""" + config = PipelineConfig( + src=SourceConfig(_slow_source([1, 2, 3, 4])), + pipes=[ + PathVariants( + router=lambda x: 0 if x <= 2 else 1, + paths=[ + [Pipe(lambda x: x * 100)], # path 0: transform + [Pipe(lambda x: x)], # path 1: passthrough + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # 1→path0→100, 2→path0→200, 3→path1→3, 4→path1→4 + self.assertEqual(results, [100, 200, 3, 4]) + + def test_paths_with_different_stage_counts(self) -> None: + """Paths with different numbers of pipe stages.""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(6))), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 10)], # path 0: 1 stage + [ + Pipe(lambda x: (x, x + 100)), + Pipe(lambda t: t[1]), + ], # path 1: 2 stages + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # Slow source ensures input-order interleaving even with different + # stage counts: 0→0, 1→(1,101)→101, 2→20, 3→(3,103)→103, 4→40, 5→105 + self.assertEqual(results, [0, 101, 20, 103, 40, 105]) + + def test_path_with_aggregate(self) -> None: + """Path containing Aggregate inside.""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(6))), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Aggregate(3)], # path 0: batch evens by 3 + [Pipe(lambda x: x + 100)], # path 1: add 100 to odds + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # Aggregate waits for 3 evens (0,2,4). With slow source, odds pass + # through immediately while aggregate buffers: + # 0→agg, 1→101, 2→agg, 3→103, 4→agg→[0,2,4], 5→105 + self.assertEqual(results, [101, 103, [0, 2, 4], 105]) + + def test_path_with_aggregate_and_disaggregate(self) -> None: + """Path containing Aggregate + Disaggregate inside.""" + config = PipelineConfig( + src=SourceConfig(range(6)), + pipes=[ + PathVariants( + router=lambda x: 0, + paths=[ + [Aggregate(2), Disaggregate()], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, [0, 1, 2, 3, 4, 5]) + + def test_nested_path_variants(self) -> None: + """PathVariants inside a path of another PathVariants.""" + inner_variants = PathVariants( + router=lambda x: 0 if x < 50 else 1, + paths=[ + [Pipe(lambda x: x + 1000)], # inner path 0 + [Pipe(lambda x: x + 2000)], # inner path 1 + ], + ) + config = PipelineConfig( + src=SourceConfig(_slow_source([1, 2, 51, 52])), + pipes=[ + PathVariants( + router=lambda x: 0, # all to path 0 + paths=[ + [inner_variants], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # 1→inner path0→1001, 2→inner path0→1002, + # 51→inner path1→2051, 52→inner path1→2052 + self.assertEqual(results, [1001, 1002, 2051, 2052]) + + def test_immediately_nested_path_variants(self) -> None: + """PathVariants as the first and only config in each path of an outer + PathVariants — the inner router reads directly from the outer router's + per-path queue.""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(12))), + pipes=[ + PathVariants( + router=lambda x: x % 2, # outer: evens vs odds + paths=[ + # path 0 (evens): immediately nest another PathVariants + [ + PathVariants( + router=lambda x: 0 if x < 6 else 1, + paths=[ + [Pipe(lambda x: x * 10)], # small evens + [Pipe(lambda x: x * 100)], # large evens + ], + ), + ], + # path 1 (odds): immediately nest another PathVariants + [ + PathVariants( + router=lambda x: 0 if x < 6 else 1, + paths=[ + [Pipe(lambda x: -x)], # small odds + [Pipe(lambda x: -(x * 10))], # large odds + ], + ), + ], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # Items interleaved across outer paths, then inner paths: + # 0→even→small→0, 1→odd→small→-1, 2→even→small→20, 3→odd→small→-3, + # 4→even→small→40, 5→odd→small→-5, 6→even→large→600, 7→odd→large→-70, + # 8→even→large→800, 9→odd→large→-90, 10→even→large→1000, 11→odd→large→-110 + self.assertEqual( + results, + [0, -1, 20, -3, 40, -5, 600, -70, 800, -90, 1000, -110], + ) + + def test_path_variants_before_and_after_pipes(self) -> None: + """Pipes before and after PathVariants in the main pipeline.""" + config = PipelineConfig( + src=SourceConfig(_slow_source([1, 2, 3, 4])), + pipes=[ + Pipe(lambda x: x * 10), # pre-processing + PathVariants( + router=lambda x: 0 if x < 25 else 1, + paths=[ + [Pipe(lambda x: x + 1)], + [Pipe(lambda x: x + 2)], + ], + ), + Pipe(lambda x: x * -1), # post-processing + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # pre: [10,20,30,40] + # path0 (<25): 10->11, 20->21 + # path1 (>=25): 30->32, 40->42 + # post: [-11,-21,-32,-42] + self.assertEqual(results, [-11, -21, -32, -42]) + + def test_path_variants_with_merge_source(self) -> None: + """Merge as source, then PathVariants in pipes.""" + plc1 = PipelineConfig( + src=SourceConfig([1, 2]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + plc2 = PipelineConfig( + src=SourceConfig([3, 4]), + pipes=[], + sink=SinkConfig(buffer_size=10), + ) + config = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 100)], + [Pipe(lambda x: x * -1)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # evens: 2,4 -> 200,400 + # odds: 1,3 -> -1,-3 + self.assertCountEqual(results, [200, 400, -1, -3]) + + def test_path_variants_with_async_pipe(self) -> None: + """Async pipe inside a path.""" + + async def async_double(x: int) -> int: + await asyncio.sleep(0.01) + return x * 2 + + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + PathVariants( + router=lambda x: 0, + paths=[ + [Pipe(async_double)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, [2, 4, 6]) + + def test_path_variants_with_concurrent_pipe(self) -> None: + """Pipe with concurrency > 1 inside a path.""" + config = PipelineConfig( + src=SourceConfig(_slow_source(range(10))), + pipes=[ + PathVariants( + router=lambda x: 0, + paths=[ + [Pipe(lambda x: x + 1, concurrency=3)], + ], + ), + ], + sink=SinkConfig(buffer_size=20), + ) + results = _run_pipeline(config) + self.assertEqual(results, list(range(1, 11))) + + +class PathVariantsValidationTest(unittest.TestCase): + """Validation tests for PathVariants config.""" + + def test_validation_empty_paths(self) -> None: + """PathVariants with no paths raises ValueError.""" + with self.assertRaises(ValueError): + PathVariants(router=lambda x: 0, paths=[]) + + def test_validation_empty_path(self) -> None: + """PathVariants with an empty path raises ValueError.""" + with self.assertRaises(ValueError): + PathVariants( + router=lambda x: 0, + paths=[[Pipe(lambda x: x)], []], + ) + + def test_validation_non_callable_router(self) -> None: + """Non-callable router raises ValueError.""" + with self.assertRaises(ValueError): + # pyre-ignore[6]: Intentionally passing non-callable to test validation. + PathVariants(router=42, paths=[[Pipe(lambda x: x)]]) + + def test_validation_source_in_path(self) -> None: + """SourceConfig in a path raises ValueError at construction time.""" + with self.assertRaises(ValueError): + PathVariants( + router=lambda x: 0, + # pyre-ignore[6]: Intentionally passing SourceConfig. + paths=[[SourceConfig([3, 4])]], + ) + + def test_validation_sink_in_path(self) -> None: + """SinkConfig in a path raises ValueError at construction time.""" + with self.assertRaises(ValueError): + PathVariants( + router=lambda x: 0, + # pyre-ignore[6]: Intentionally passing SinkConfig. + paths=[[SinkConfig(buffer_size=10)]], + ) + + def test_validation_source_in_second_path(self) -> None: + """SourceConfig in a later path position raises ValueError.""" + with self.assertRaises(ValueError): + PathVariants( + router=lambda x: 0, + # pyre-ignore[6]: Intentionally passing SourceConfig. + paths=[ + [Pipe(lambda x: x)], + [Pipe(lambda x: x), SourceConfig([1, 2])], + ], + ) + + def test_validation_sink_in_middle_of_path(self) -> None: + """SinkConfig in the middle of a path raises ValueError.""" + with self.assertRaises(ValueError): + PathVariants( + router=lambda x: 0, + # pyre-ignore[6]: Intentionally passing SinkConfig. + paths=[ + [Pipe(lambda x: x), SinkConfig(buffer_size=10), Pipe(lambda x: x)], + ], + ) + + +class PathVariantsErrorHandlingTest(unittest.TestCase): + """Error handling and edge case tests for PathVariants.""" + + def test_router_returns_negative_index(self) -> None: + """Router returns -1 — pipeline fails with PipelineFailure.""" + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + PathVariants( + router=lambda x: -1, + paths=[ + [Pipe(lambda x: x)], + [Pipe(lambda x: x)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + def test_router_returns_index_equal_to_num_paths(self) -> None: + """Router returns N (== len(paths)) — pipeline fails.""" + config = PipelineConfig( + src=SourceConfig([1]), + pipes=[ + PathVariants( + router=lambda x: 2, # only 2 paths, index 2 is out of range + paths=[ + [Pipe(lambda x: x)], + [Pipe(lambda x: x)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + def test_router_returns_index_greater_than_num_paths(self) -> None: + """Router returns N+5 — pipeline fails.""" + config = PipelineConfig( + src=SourceConfig([1]), + pipes=[ + PathVariants( + router=lambda x: 7, + paths=[ + [Pipe(lambda x: x)], + [Pipe(lambda x: x)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + def test_one_path_fails_other_continues(self) -> None: + """One path's pipe raises; the other path's items still processed. + + With default max_failures=-1, individual task failures are tolerated. + The pipeline completes successfully and the non-failing path's results + are collected. The failing path's items are dropped. + """ + + def fail_on_odd(x): + if x % 2 == 1: + raise ValueError(f"odd item {x}") + return x * 10 + + config = PipelineConfig( + src=SourceConfig(range(6)), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 10)], # path 0: evens succeed + [Pipe(fail_on_odd)], # path 1: odds fail (dropped) + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + # Path 0 items succeed: 0,2,4 -> 0,20,40 + # Path 1 items fail and are dropped + self.assertEqual(results, [0, 20, 40]) + + def test_one_path_fails_with_max_failures(self) -> None: + """One path's pipe raises with max_failures=0; pipeline fails.""" + + def fail_always(x): + raise ValueError("fail") + + config = PipelineConfig( + src=SourceConfig(range(4)), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 10)], + [Pipe(fail_always, max_failures=0)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + def test_all_paths_fail(self) -> None: + """All paths fail with max_failures=0 — pipeline raises PipelineFailure.""" + + def always_fail(x): + raise ValueError("boom") + + config = PipelineConfig( + src=SourceConfig([1, 2]), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(always_fail, max_failures=0)], + [Pipe(always_fail, max_failures=0)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + def test_empty_source_with_path_variants(self) -> None: + """Empty source — clean shutdown with no results.""" + config = PipelineConfig( + src=SourceConfig([]), + pipes=[ + PathVariants( + router=lambda x: 0, + paths=[ + [Pipe(lambda x: x * 2)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + results = _run_pipeline(config) + self.assertEqual(results, []) + + def test_path_failure_cancels_router(self) -> None: + """When a path fails with max_failures=0, the router is cancelled and + the pipeline shuts down cleanly without hanging.""" + + def fail_immediately(x): + raise ValueError("boom") + + config = PipelineConfig( + # Use enough items so the router is still active when path fails. + src=SourceConfig(range(100)), + pipes=[ + PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x, max_failures=0)], # path 0: succeeds + [Pipe(fail_immediately, max_failures=0)], # path 1: fails + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + # The pipeline must raise PipelineFailure (not hang). + # If the router is not cancelled on path failure, this would deadlock + # because the router keeps trying to push items to the dead path's queue. + with self.assertRaises(PipelineFailure): + _run_pipeline(config, timeout=5) + + def test_path_failure_cancels_router_all_to_failing_path(self) -> None: + """All items routed to the failing path — router cancelled, no hang.""" + + def fail_immediately(x): + raise ValueError("boom") + + config = PipelineConfig( + src=SourceConfig(range(100)), + pipes=[ + PathVariants( + router=lambda x: 1, # all items to failing path + paths=[ + [Pipe(lambda x: x)], + [Pipe(fail_immediately, max_failures=0)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config, timeout=5) + + def test_router_raises_exception(self) -> None: + """Router function itself raises — pipeline fails cleanly.""" + + def bad_router(x): + raise RuntimeError("router error") + + config = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + PathVariants( + router=bad_router, + paths=[ + [Pipe(lambda x: x)], + ], + ), + ], + sink=SinkConfig(buffer_size=10), + ) + with self.assertRaises(PipelineFailure): + _run_pipeline(config) + + +class PathVariantsReprTest(unittest.TestCase): + """Repr tests for PathVariants.""" + + def test_repr(self) -> None: + """Verify repr is readable and includes path info.""" + cfg = PathVariants( + router=lambda x: x % 2, + paths=[ + [Pipe(lambda x: x * 2, name="double")], + [Pipe(lambda x: x + 1, name="add_one")], + ], + ) + r = repr(cfg) + self.assertIn("PathVariants", r) + self.assertIn("path0", r) + self.assertIn("path1", r) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/spdl/pipeline/tests/percentile_stats_test.py b/src/spdl/pipeline/tests/percentile_stats_test.py new file mode 100644 index 000000000..83758584f --- /dev/null +++ b/src/spdl/pipeline/tests/percentile_stats_test.py @@ -0,0 +1,453 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import unittest +from collections.abc import Callable + +from spdl.pipeline._components._common import _P2Percentile, _StatsCounter, StageInfo +from spdl.pipeline._components._hook import TaskPerfStats, TaskStatsHook +from spdl.pipeline._components._queue import QueuePerfStats, StatsQueue + + +class P2PercentileTest(unittest.TestCase): + def test_empty(self) -> None: + """A fresh _P2Percentile with no observations reports 0.0.""" + p = _P2Percentile(90) + self.assertEqual(p.value, 0.0) + + def test_single_element(self) -> None: + """With only one observation, the percentile value equals that observation.""" + p = _P2Percentile(90) + p.update(5.0) + self.assertEqual(p.value, 5.0) + + def test_two_elements(self) -> None: + """With two observations, p90 returns the larger value.""" + p = _P2Percentile(90) + p.update(1.0) + p.update(2.0) + self.assertEqual(p.value, 2.0) + + def test_four_elements(self) -> None: + """With fewer than 5 observations, the fallback sorted-lookup is used. + + Checks that p50 of [1, 2, 3, 4] returns the median (3.0). + """ + p = _P2Percentile(50) + for v in [1.0, 2.0, 3.0, 4.0]: + p.update(v) + self.assertEqual(p.value, 3.0) + + def test_five_elements_p90(self) -> None: + """With exactly 5 observations the P² algorithm initializes. + + Checks that the p90 estimate falls within a reasonable range. + """ + p = _P2Percentile(90) + for v in [1.0, 2.0, 3.0, 4.0, 5.0]: + p.update(v) + self.assertGreaterEqual(p.value, 2.0) + self.assertLessEqual(p.value, 5.0) + + def test_hundred_elements_p90(self) -> None: + """P² p90 estimate over 100 sequential values is close to the true p90 (90.0).""" + p = _P2Percentile(90) + for i in range(100): + p.update(float(i)) + self.assertAlmostEqual(p.value, 90.0, delta=3.0) + + def test_hundred_elements_p99(self) -> None: + """P² p99 estimate over 100 sequential values is close to the true p99 (99.0).""" + p = _P2Percentile(99) + for i in range(100): + p.update(float(i)) + self.assertAlmostEqual(p.value, 99.0, delta=3.0) + + def test_hundred_elements_p50(self) -> None: + """P² p50 (median) estimate over 100 sequential values is close to 50.0.""" + p = _P2Percentile(50) + for i in range(100): + p.update(float(i)) + self.assertAlmostEqual(p.value, 50.0, delta=3.0) + + def test_thousand_elements_accuracy(self) -> None: + """With 1000 observations, p90 and p99 estimates stay within tight bounds. + + Checks p90 ≈ 900 (±20) and p99 ≈ 990 (±20). + """ + p90 = _P2Percentile(90) + p99 = _P2Percentile(99) + for i in range(1000): + v = float(i) + p90.update(v) + p99.update(v) + self.assertAlmostEqual(p90.value, 900.0, delta=20.0) + self.assertAlmostEqual(p99.value, 990.0, delta=20.0) + + def test_reset(self) -> None: + """After reset(), the estimator returns to its initial state (value 0.0).""" + p = _P2Percentile(90) + for i in range(20): + p.update(float(i)) + self.assertGreater(p.value, 0.0) + p.reset() + self.assertEqual(p.value, 0.0) + + def test_reset_and_reuse(self) -> None: + """After reset(), feeding new data produces estimates based only on the new data. + + First feeds [0..99], resets, then feeds [100..199] and checks p50 ≈ 150. + """ + p = _P2Percentile(50) + for i in range(100): + p.update(float(i)) + p.reset() + for i in range(100, 200): + p.update(float(i)) + self.assertAlmostEqual(p.value, 150.0, delta=5.0) + + def test_unsorted_input(self) -> None: + """P² produces accurate estimates regardless of input order. + + Feeds 10 values in shuffled order and checks p90 ≈ 8.0 (±2). + """ + p = _P2Percentile(90) + values = [9.0, 1.0, 5.0, 3.0, 7.0, 2.0, 8.0, 4.0, 6.0, 0.0] + for v in values: + p.update(v) + self.assertAlmostEqual(p.value, 8.0, delta=2.0) + + def test_constant_values(self) -> None: + """When all observations are identical, the percentile equals that constant.""" + p = _P2Percentile(90) + for _ in range(20): + p.update(42.0) + self.assertAlmostEqual(p.value, 42.0, delta=0.01) + + +class StatsCounterPercentileTest(unittest.TestCase): + def test_initial_state(self) -> None: + """A fresh _StatsCounter has zero items, zero average, and zero percentiles.""" + counter = _StatsCounter() + self.assertEqual(counter.p90_time, 0.0) + self.assertEqual(counter.p99_time, 0.0) + self.assertEqual(counter.num_items, 0) + self.assertEqual(counter.ave_time, 0.0) + + def test_update_tracks_percentiles(self) -> None: + """After 100 updates, the counter's p90 and p99 properties reflect accurate estimates.""" + counter = _StatsCounter() + for i in range(100): + counter.update(float(i)) + self.assertEqual(counter.num_items, 100) + self.assertAlmostEqual(counter.p90_time, 90.0, delta=3.0) + self.assertAlmostEqual(counter.p99_time, 99.0, delta=3.0) + + def test_count_context_manager(self) -> None: + """The count() context manager records one item with non-negative percentiles.""" + counter = _StatsCounter() + with counter.count(): + pass + self.assertEqual(counter.num_items, 1) + self.assertGreaterEqual(counter.p90_time, 0.0) + self.assertGreaterEqual(counter.p99_time, 0.0) + + def test_consume_lap_percentiles(self) -> None: + """consume_lap_percentiles() returns current lap p90/p99, then resets lap trackers. + + After consuming, a second call returns (0.0, 0.0). + """ + counter = _StatsCounter() + for i in range(100): + counter.update(float(i)) + + p90, p99 = counter.consume_lap_percentiles() + self.assertAlmostEqual(p90, 90.0, delta=3.0) + self.assertAlmostEqual(p99, 99.0, delta=3.0) + + p90_after, p99_after = counter.consume_lap_percentiles() + self.assertEqual(p90_after, 0.0) + self.assertEqual(p99_after, 0.0) + + def test_consume_lap_does_not_affect_overall(self) -> None: + """Consuming lap percentiles does not change the overall (lifetime) p90 value.""" + counter = _StatsCounter() + for i in range(100): + counter.update(float(i)) + + overall_p90_before = counter.p90_time + counter.consume_lap_percentiles() + self.assertEqual(counter.p90_time, overall_p90_before) + + +class TaskPerfStatsTest(unittest.TestCase): + def test_fields_present(self) -> None: + """TaskPerfStats dataclass stores all fields including p90_time and p99_time.""" + stats = TaskPerfStats( + num_tasks=10, + num_failures=1, + ave_time=0.5, + p90_time=0.8, + p99_time=1.2, + ) + self.assertEqual(stats.num_tasks, 10) + self.assertEqual(stats.num_failures, 1) + self.assertEqual(stats.ave_time, 0.5) + self.assertEqual(stats.p90_time, 0.8) + self.assertEqual(stats.p99_time, 1.2) + + +class QueuePerfStatsTest(unittest.TestCase): + def test_fields_present(self) -> None: + """QueuePerfStats dataclass stores p90/p99 fields for both put and get operations.""" + stats = QueuePerfStats( + elapsed=60.0, + num_items=100, + ave_put_time=0.01, + ave_get_time=0.02, + p90_put_time=0.015, + p99_put_time=0.025, + p90_get_time=0.03, + p99_get_time=0.04, + occupancy_rate=0.75, + ) + self.assertEqual(stats.p90_put_time, 0.015) + self.assertEqual(stats.p99_put_time, 0.025) + self.assertEqual(stats.p90_get_time, 0.03) + self.assertEqual(stats.p99_get_time, 0.04) + + def test_qps(self) -> None: + """The qps property computes num_items / elapsed correctly.""" + stats = QueuePerfStats( + elapsed=10.0, + num_items=100, + ave_put_time=0.0, + ave_get_time=0.0, + p90_put_time=0.0, + p99_put_time=0.0, + p90_get_time=0.0, + p99_get_time=0.0, + occupancy_rate=0.0, + ) + self.assertAlmostEqual(stats.qps, 10.0) + + def test_qps_zero_elapsed(self) -> None: + """When elapsed is zero, qps returns 0 to avoid division by zero.""" + stats = QueuePerfStats( + elapsed=0.0, + num_items=100, + ave_put_time=0.0, + ave_get_time=0.0, + p90_put_time=0.0, + p99_put_time=0.0, + p90_get_time=0.0, + p99_get_time=0.0, + occupancy_rate=0.0, + ) + self.assertEqual(stats.qps, 0) + + +class TaskStatsHookPercentileTest(unittest.IsolatedAsyncioTestCase): + async def test_task_hook_records_percentiles(self) -> None: + """Successful tasks are tracked by the P² percentile estimators. + + Checks that after 10 successful tasks, num_tasks/num_success are correct + and both p90/p99 trackers have non-negative values. + """ + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 + ) + for _ in range(10): + async with hook.task_hook(): # pyre-ignore[16] + pass + self.assertEqual(hook.num_tasks, 10) + self.assertEqual(hook.num_success, 10) + self.assertGreaterEqual(hook._p90.value, 0.0) + self.assertGreaterEqual(hook._p99.value, 0.0) + + async def test_failed_task_not_tracked_in_percentiles(self) -> None: + """Failed tasks increment num_tasks but are excluded from percentile tracking. + + Runs 3 tasks (2 succeed, 1 fails). Checks that the P² tracker count is 2. + """ + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 + ) + + async with hook.task_hook(): # pyre-ignore[16] + pass + + try: + async with hook.task_hook(): # pyre-ignore[16] + raise ValueError("fail") + except ValueError: + pass + + async with hook.task_hook(): # pyre-ignore[16] + pass + + self.assertEqual(hook.num_tasks, 3) + self.assertEqual(hook.num_success, 2) + self.assertEqual(hook._p90._count, 2) + + async def test_stage_hook_produces_stats_with_percentiles(self) -> None: + """When stage_hook exits, _log_stats is called with a TaskPerfStats that + includes non-negative p90_time and p99_time values. + """ + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 + ) + logged_stats: list[TaskPerfStats] = [] + original_log: Callable[[TaskPerfStats], None] = hook._log_stats + + def capture_stats(stats: TaskPerfStats) -> None: + logged_stats.append(stats) + original_log(stats) + + hook._log_stats = capture_stats # pyre-ignore[8] + + async with hook.stage_hook(): # pyre-ignore[16] + for _ in range(5): + async with hook.task_hook(): # pyre-ignore[16] + pass + + self.assertEqual(len(logged_stats), 1) + stats = logged_stats[0] + self.assertEqual(stats.num_tasks, 5) + self.assertEqual(stats.num_failures, 0) + self.assertGreater(stats.ave_time, 0.0) + self.assertGreaterEqual(stats.p90_time, 0.0) + self.assertGreaterEqual(stats.p99_time, 0.0) + + async def test_lap_stats_with_percentiles(self) -> None: + """Lap stats report percentiles only for the current interval, then reset. + + First lap covers 10 tasks; second lap covers 5 new tasks. Each lap's + p90/p99 should be positive and independent of the other. + """ + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 + ) + + for _ in range(10): + async with hook.task_hook(): # pyre-ignore[16] + pass + + lap1 = hook._get_lap_stats() + self.assertEqual(lap1.num_tasks, 10) + self.assertGreater(lap1.p90_time, 0.0) + self.assertGreater(lap1.p99_time, 0.0) + + for _ in range(5): + async with hook.task_hook(): # pyre-ignore[16] + pass + + lap2 = hook._get_lap_stats() + self.assertEqual(lap2.num_tasks, 5) + self.assertGreater(lap2.p90_time, 0.0) + self.assertGreater(lap2.p99_time, 0.0) + + async def test_empty_lap_stats(self) -> None: + """When no tasks have run, lap stats report zero for all percentile fields.""" + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 + ) + lap = hook._get_lap_stats() + self.assertEqual(lap.num_tasks, 0) + self.assertEqual(lap.p90_time, 0.0) + self.assertEqual(lap.p99_time, 0.0) + + +class StatsQueuePercentileTest(unittest.IsolatedAsyncioTestCase): + async def test_put_get_records_percentiles(self) -> None: + """After put/get operations, the queue's internal counters have + non-negative p90 percentile values for both put and get. + """ + queue = StatsQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), + buffer_size=10, + interval=-1, + ) + async with queue.stage_hook(): # pyre-ignore[16] + for i in range(5): + await queue.put(i) + for _ in range(5): + await queue.get() + + self.assertGreaterEqual(queue._putc.p90_time, 0.0) + self.assertGreaterEqual(queue._getc.p90_time, 0.0) + + async def test_stage_hook_produces_stats_with_percentiles(self) -> None: + """When stage_hook exits, _log_stats receives a QueuePerfStats with + non-negative p90/p99 values for both put and get operations. + """ + queue = StatsQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), + buffer_size=10, + interval=-1, + ) + logged_stats: list[QueuePerfStats] = [] + original_log: Callable[[QueuePerfStats], None] = queue._log_stats + + def capture_stats(stats: QueuePerfStats) -> None: + logged_stats.append(stats) + original_log(stats) + + queue._log_stats = capture_stats # pyre-ignore[8] + + async with queue.stage_hook(): # pyre-ignore[16] + for i in range(5): + await queue.put(i) + for _ in range(5): + await queue.get() + + self.assertEqual(len(logged_stats), 1) + stats = logged_stats[0] + self.assertEqual(stats.num_items, 5) + self.assertGreaterEqual(stats.p90_put_time, 0.0) + self.assertGreaterEqual(stats.p99_put_time, 0.0) + self.assertGreaterEqual(stats.p90_get_time, 0.0) + self.assertGreaterEqual(stats.p99_get_time, 0.0) + + async def test_lap_stats_with_percentiles(self) -> None: + """Lap stats for the queue report per-interval p90/p99 for put and get, + resetting between laps. + + First lap covers 8 items; second lap covers 3 new items. Each lap's + percentile fields should be non-negative and independent. + """ + queue = StatsQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), + buffer_size=10, + interval=-1, + ) + queue._lap_t0 = asyncio.get_running_loop().time() + queue._empty_t0 = queue._lap_t0 + + for i in range(8): + await queue.put(i) + for _ in range(8): + await queue.get() + + lap1 = queue._get_lap_stats() + self.assertEqual(lap1.num_items, 8) + self.assertGreaterEqual(lap1.p90_put_time, 0.0) + self.assertGreaterEqual(lap1.p90_get_time, 0.0) + self.assertGreaterEqual(lap1.p99_put_time, 0.0) + self.assertGreaterEqual(lap1.p99_get_time, 0.0) + + for i in range(3): + await queue.put(i) + for _ in range(3): + await queue.get() + + lap2 = queue._get_lap_stats() + self.assertEqual(lap2.num_items, 3) + self.assertGreaterEqual(lap2.p90_put_time, 0.0) + self.assertGreaterEqual(lap2.p90_get_time, 0.0) diff --git a/src/spdl/pipeline/tests/pgrp_stats_test.py b/src/spdl/pipeline/tests/pgrp_stats_test.py new file mode 100644 index 000000000..d7ede1829 --- /dev/null +++ b/src/spdl/pipeline/tests/pgrp_stats_test.py @@ -0,0 +1,496 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import multiprocessing +import os +import sys +import tempfile +import unittest +from collections.abc import Awaitable, Callable +from unittest.mock import AsyncMock, MagicMock, patch + +from spdl.pipeline._bg_task import BackgroundTask +from spdl.pipeline._pgrp_stats import ( + _collect_pgrp_stats, + _parse_proc_io, + _parse_proc_stat, + _parse_smaps_rollup, + _pgrp_monitor_subprocess, + _read_file, + _read_network_bytes, + _read_pgrp_stats, + _warned, + ProcessGroupResourceUsage, + ProcessGroupStatsMonitor, +) + +_MODULE = "spdl.pipeline._pgrp_stats" + + +@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") +class LiveProcMonitorTest(unittest.TestCase): + """Integration test that reads real /proc data without mocking.""" + + def test_read_pgrp_stats_returns_valid_data(self) -> None: + """_read_pgrp_stats should find at least this process.""" + _warned.discard("proc_stat") + _warned.discard("proc_io") + _warned.discard("smaps_rollup") + + result = _read_pgrp_stats() + self.assertGreaterEqual(result.num_procs, 1) + self.assertGreaterEqual(result.cpu_usec, 0) + self.assertGreaterEqual(result.rss_bytes, 0) + self.assertGreaterEqual(result.disk_read_bytes, 0) + self.assertGreaterEqual(result.disk_write_bytes, 0) + # smaps_rollup should be available on modern Linux + if result.pss_bytes is not None: + self.assertGreater(result.pss_bytes, 0) + self.assertIsNotNone(result.private_bytes) + self.assertGreater(result.private_bytes, 0) + + def test_collect_pgrp_stats_returns_complete_snapshot(self) -> None: + """_collect_pgrp_stats should return a fully populated snapshot.""" + result, cpu_usec, time_usec, net_rx, net_tx = _collect_pgrp_stats() + self.assertIsInstance(result, ProcessGroupResourceUsage) + self.assertEqual(result.pid, os.getpid()) + self.assertEqual(result.pgid, os.getpgrp()) + + # First call: cpu_percent and net deltas should be None (no previous value). + self.assertIsNone(result.cpu_percent) + self.assertIsNotNone(cpu_usec) + self.assertIsNotNone(result.rss_bytes) + self.assertIsNotNone(result.num_procs) + self.assertIsNone(result.net_rx_bytes) + self.assertIsNone(result.net_tx_bytes) + self.assertIsNotNone(net_rx) + self.assertIsNotNone(net_tx) + + # Second call with prev values: cpu_percent and net deltas should be set. + result2, _, _, _, _ = _collect_pgrp_stats(cpu_usec, time_usec, net_rx, net_tx) + cpu_pct = result2.cpu_percent + assert cpu_pct is not None + self.assertGreaterEqual(cpu_pct, 0.0) + rx = result2.net_rx_bytes + assert rx is not None + self.assertGreaterEqual(rx, 0) + tx = result2.net_tx_bytes + assert tx is not None + self.assertGreaterEqual(tx, 0) + + # Sanity: at least one process (this one) should be counted. + assert result.num_procs is not None + self.assertGreaterEqual(result.num_procs, 1) + # RSS must be positive for a running process. + assert result.rss_bytes is not None + self.assertGreater(result.rss_bytes, 0) + + +class ReadFileTest(unittest.TestCase): + def test_read_existing_file(self) -> None: + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("hello\n") + f.flush() + result = _read_file(f.name) + self.assertEqual(result, "hello") + os.unlink(f.name) + + def test_read_nonexistent_file(self) -> None: + result = _read_file("/nonexistent/path/file.txt") + self.assertIsNone(result) + + +class ReadNetworkTest(unittest.TestCase): + @patch(f"{_MODULE}._read_file") + def test_network_bytes(self, mock_read: MagicMock) -> None: + mock_read.return_value = ( + "Inter-| Receive | Transmit\n" # noqa: B950 + " face |bytes packets errs drop fifo frame compressed multicast|bytes packets errs drop fifo colls carrier compressed\n" # noqa: B950 + " lo: 1000 10 0 0 0 0 0 0 2000 20 0 0 0 0 0 0\n" # noqa: B950 + " eth0: 5000 50 0 0 0 0 0 0 3000 30 0 0 0 0 0 0\n" # noqa: B950 + " eth1: 7000 70 0 0 0 0 0 0 4000 40 0 0 0 0 0 0" # noqa: B950 + ) + result = _read_network_bytes() + # lo is excluded; eth0 + eth1 + self.assertEqual(result.rx_bytes, 12000) + self.assertEqual(result.tx_bytes, 7000) + + @patch(f"{_MODULE}._read_file") + def test_network_file_missing(self, mock_read: MagicMock) -> None: + mock_read.return_value = None + result = _read_network_bytes() + self.assertEqual(result.rx_bytes, 0) + self.assertEqual(result.tx_bytes, 0) + + @patch(f"{_MODULE}._read_file") + def test_network_malformed_line_raises(self, mock_read: MagicMock) -> None: + mock_read.return_value = ( + "Inter-| Receive\n" + " face |bytes\n" + " eth0: bad 0 0 0 0 0 0 0 also_bad 0 0 0 0 0 0 0\n" + ) + with self.assertRaises(RuntimeError, msg="Failed to parse /proc/net/dev"): + _read_network_bytes() + + +class ParseProcStatTest(unittest.TestCase): + def test_normal_comm(self) -> None: + content = ( + "12345 (python3) S 100 200 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" + ) + stat = _parse_proc_stat(content) + self.assertEqual(stat.pgrp, 200) + self.assertEqual(stat.utime, 500) + self.assertEqual(stat.stime, 300) + self.assertEqual(stat.rss, 4096) + + def test_comm_with_spaces_and_parens(self) -> None: + content = "12345 (my (weird) app) S 100 200 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" # noqa: B950 + stat = _parse_proc_stat(content) + self.assertEqual(stat.pgrp, 200) + + def test_malformed_no_parens_raises(self) -> None: + with self.assertRaises(RuntimeError, msg="missing closing paren"): + _parse_proc_stat("no parens here") + + def test_too_few_fields_raises(self) -> None: + with self.assertRaises(RuntimeError, msg="expected >=22 fields"): + _parse_proc_stat("12345 (python3) S 100 200") + + def test_non_numeric_field_raises(self) -> None: + content = ( + "12345 (python3) S 100 abc 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" # noqa: B950 + ) + with self.assertRaises(RuntimeError, msg="Failed to parse"): + _parse_proc_stat(content) + + +class ParseProcIoTest(unittest.TestCase): + def test_normal_io(self) -> None: + content = ( + "rchar: 123456\n" + "wchar: 654321\n" + "syscr: 100\n" + "syscw: 200\n" + "read_bytes: 4096\n" + "write_bytes: 8192\n" + "cancelled_write_bytes: 0" + ) + io = _parse_proc_io(content) + self.assertEqual(io.read_bytes, 4096) + self.assertEqual(io.write_bytes, 8192) + + def test_empty_content(self) -> None: + io = _parse_proc_io("") + self.assertEqual(io.read_bytes, 0) + self.assertEqual(io.write_bytes, 0) + + def test_malformed_value_raises(self) -> None: + content = "read_bytes: not_a_number\n" + with self.assertRaises(RuntimeError, msg="Failed to parse /proc/[pid]/io"): + _parse_proc_io(content) + + +class ParseSmapsRollupTest(unittest.TestCase): + def test_normal_smaps_rollup(self) -> None: + content = ( + "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\n" + "Rss: 123456 kB\n" + "Pss: 98765 kB\n" + "Pss_Dirty: 40000 kB\n" + "Private_Clean: 50000 kB\n" + "Private_Dirty: 30000 kB\n" + "Shared_Clean: 20000 kB\n" + "Shared_Dirty: 10000 kB\n" + ) + result = _parse_smaps_rollup(content) + self.assertEqual(result.pss, 98765 * 1024) + self.assertEqual(result.private_clean, 50000 * 1024) + self.assertEqual(result.private_dirty, 30000 * 1024) + + def test_empty_content(self) -> None: + result = _parse_smaps_rollup("") + self.assertEqual(result.pss, 0) + self.assertEqual(result.private_clean, 0) + self.assertEqual(result.private_dirty, 0) + + def test_malformed_value_raises(self) -> None: + content = "Pss: not_a_number kB\n" + with self.assertRaises(RuntimeError, msg="Failed to parse"): + _parse_smaps_rollup(content) + + +@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") +class ReadPgrpStatsTest(unittest.TestCase): + def setUp(self) -> None: + _warned.discard("proc_stat") + _warned.discard("proc_io") + + @patch(f"{_MODULE}.os.scandir") + @patch(f"{_MODULE}._read_file") + @patch(f"{_MODULE}.os.getpgrp") + def test_sums_processes_in_same_pgrp( + self, + mock_getpgrp: MagicMock, + mock_read: MagicMock, + mock_scandir: MagicMock, + ) -> None: + mock_getpgrp.return_value = 1000 + + # Two processes in pgrp 1000, one in pgrp 9999 + entries = [] + for name in ["101", "102", "103", "not_a_pid"]: + entry = MagicMock() + entry.name = name + entry.is_dir.return_value = True + entries.append(entry) + mock_scandir.return_value = entries + + def read_side_effect(path: str) -> str | None: + if path == "/proc/101/stat": + return "101 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 100 50 0 0 20 0 1 0 0 0 2000" # noqa: B950 + if path == "/proc/101/smaps_rollup": + return "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\nRss: 8000 kB\nPss: 6000 kB\nPrivate_Clean: 3000 kB\nPrivate_Dirty: 2000 kB\n" # noqa: B950 + if path == "/proc/101/io": + return "rchar: 1000\nwchar: 2000\nsyscr: 10\nsyscw: 20\nread_bytes: 4096\nwrite_bytes: 8192\ncancelled_write_bytes: 0" # noqa: B950 + if path == "/proc/102/stat": + return "102 (worker) S 1 1000 1000 0 -1 0 0 0 0 0 200 75 0 0 20 0 1 0 0 0 3000" # noqa: B950 + if path == "/proc/102/smaps_rollup": + return "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\nRss: 12000 kB\nPss: 9000 kB\nPrivate_Clean: 5000 kB\nPrivate_Dirty: 3000 kB\n" # noqa: B950 + if path == "/proc/102/io": + return "rchar: 3000\nwchar: 4000\nsyscr: 30\nsyscw: 40\nread_bytes: 1024\nwrite_bytes: 2048\ncancelled_write_bytes: 0" # noqa: B950 + if path == "/proc/103/stat": + # Different pgrp + return "103 (other) S 1 9999 9999 0 -1 0 0 0 0 0 999 999 0 0 20 0 1 0 0 0 9999" # noqa: B950 + return None + + mock_read.side_effect = read_side_effect + + result = _read_pgrp_stats() + + # utime: 100+200=300, stime: 50+75=125, total_ticks=425 + from spdl.pipeline._pgrp_stats import _get_sc_clk_tck + + expected_cpu_usec = 425 * 1_000_000 // _get_sc_clk_tck() + self.assertEqual(result.cpu_usec, expected_cpu_usec) + + # rss: 2000+3000=5000 pages + from spdl.pipeline._pgrp_stats import _get_page_size + + expected_rss = 5000 * _get_page_size() + self.assertEqual(result.rss_bytes, expected_rss) + + # pss: 6000+9000=15000 kB + self.assertEqual(result.pss_bytes, 15000 * 1024) + + # private: (3000+2000)+(5000+3000)=13000 kB + self.assertEqual(result.private_bytes, 13000 * 1024) + + # disk IO: 4096+1024=5120 read, 8192+2048=10240 write + self.assertEqual(result.disk_read_bytes, 5120) + self.assertEqual(result.disk_write_bytes, 10240) + + self.assertEqual(result.num_procs, 2) + + @patch(f"{_MODULE}.os.scandir") + @patch(f"{_MODULE}.os.getpgrp") + def test_scandir_failure_raises( + self, + mock_getpgrp: MagicMock, + mock_scandir: MagicMock, + ) -> None: + mock_getpgrp.return_value = 1000 + mock_scandir.side_effect = OSError("permission denied") + + with self.assertRaises(RuntimeError, msg="Failed to scan /proc"): + _read_pgrp_stats() + + @patch(f"{_MODULE}.os.scandir") + @patch(f"{_MODULE}._read_file") + @patch(f"{_MODULE}.os.getpgrp") + def test_missing_io_and_smaps_file( + self, + mock_getpgrp: MagicMock, + mock_read: MagicMock, + mock_scandir: MagicMock, + ) -> None: + """Disk IO is 0 and PSS/private are None when files are unreadable.""" + mock_getpgrp.return_value = 1000 + + entry = MagicMock() + entry.name = "101" + mock_scandir.return_value = [entry] + + def read_side_effect(path: str) -> str | None: + if path == "/proc/101/stat": + return "101 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 100 50 0 0 20 0 1 0 0 0 2000" # noqa: B950 + # /proc/101/io and /proc/101/smaps_rollup return None + return None + + mock_read.side_effect = read_side_effect + + result = _read_pgrp_stats() + self.assertEqual(result.disk_read_bytes, 0) + self.assertEqual(result.disk_write_bytes, 0) + self.assertIsNone(result.pss_bytes) + self.assertIsNone(result.private_bytes) + self.assertEqual(result.num_procs, 1) + + @patch(f"{_MODULE}.os.scandir") + @patch(f"{_MODULE}._read_file") + @patch(f"{_MODULE}.os.getpgrp") + def test_malformed_stat_skips_with_warning( + self, + mock_getpgrp: MagicMock, + mock_read: MagicMock, + mock_scandir: MagicMock, + ) -> None: + """A process with malformed stat is skipped and warns once.""" + mock_getpgrp.return_value = 1000 + + entries = [] + for name in ["101", "102"]: + entry = MagicMock() + entry.name = name + entries.append(entry) + mock_scandir.return_value = entries + + def read_side_effect(path: str) -> str | None: + if path == "/proc/101/stat": + return "malformed content" # no parens + if path == "/proc/102/stat": + return "102 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 200 75 0 0 20 0 1 0 0 0 3000" # noqa: B950 + return None + + mock_read.side_effect = read_side_effect + + with self.assertLogs(_MODULE, level="WARNING") as cm: + result = _read_pgrp_stats() + + # Only pid 102 was counted + self.assertEqual(result.num_procs, 1) + self.assertTrue(any("missing closing paren" in m for m in cm.output)) + + +class ProcessGroupStatsMonitorClassTest(unittest.TestCase): + def test_is_background_task(self) -> None: + monitor = ProcessGroupStatsMonitor(callback=AsyncMock()) + self.assertIsInstance(monitor, BackgroundTask) + + +@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") +class ProcessGroupStatsMonitorSubprocessTest(unittest.TestCase): + def test_monitor_spawns_and_cancels_subprocess(self) -> None: + """Verify the monitor spawns a subprocess and terminates it on cancel.""" + mock_proc = MagicMock(spec=multiprocessing.Process) + mock_proc.pid = 12345 + mock_proc.is_alive.side_effect = [True, True, False] + + mock_ctx: MagicMock = MagicMock() + mock_ctx.Process.return_value = mock_proc + + async def run_monitor() -> None: + monitor = ProcessGroupStatsMonitor( + callback=AsyncMock(), mp_context=mock_ctx + ) + task = asyncio.create_task(monitor.run()) + await asyncio.sleep(0.01) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run_monitor()) + + mock_proc.start.assert_called_once() + mock_proc.terminate.assert_called_once() + mock_proc.join.assert_called() + + def test_monitor_warns_on_unexpected_exit(self) -> None: + """Verify warning is logged when the subprocess exits unexpectedly.""" + mock_proc = MagicMock(spec=multiprocessing.Process) + mock_proc.pid = 12345 + mock_proc.exitcode = 1 + mock_proc.is_alive.return_value = False + + mock_ctx: MagicMock = MagicMock() + mock_ctx.Process.return_value = mock_proc + + async def run_monitor() -> None: + monitor = ProcessGroupStatsMonitor( + callback=AsyncMock(), mp_context=mock_ctx + ) + await monitor.run() + + with self.assertLogs(_MODULE, level="WARNING") as cm: + asyncio.run(run_monitor()) + + exit_warnings = [m for m in cm.output if "exited unexpectedly" in m] + self.assertEqual(len(exit_warnings), 1) + self.assertIn("exit code 1", exit_warnings[0]) + + +@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") +class PgrpMonitorSubprocessFunctionTest(unittest.TestCase): + @patch(f"{_MODULE}._collect_pgrp_stats") + def test_subprocess_function_collects_and_calls_callback( + self, + mock_collect: MagicMock, + ) -> None: + """Test the subprocess entry point invokes the callback.""" + usage = ProcessGroupResourceUsage( + pid=os.getpid(), + pgid=os.getpgrp(), + cpu_percent=50.0, + rss_bytes=1048576, + pss_bytes=800000, + private_bytes=600000, + disk_read_bytes=4096, + disk_write_bytes=8192, + num_procs=3, + net_rx_bytes=100, + net_tx_bytes=200, + ) + mock_collect.return_value = (usage, 500000, 1000000, 100, 200) + + mock_callback = AsyncMock() + + call_count = 0 + original_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep + + async def counting_sleep(delay: float) -> None: + nonlocal call_count + call_count += 1 + if call_count >= 2: + raise KeyboardInterrupt + await original_sleep(0) + + with ( + patch("asyncio.sleep", counting_sleep), + patch(f"{_MODULE}.signal.signal"), + ): + try: + _pgrp_monitor_subprocess(0.01, mock_callback) + except KeyboardInterrupt: + pass + + mock_collect.assert_called() + mock_callback.assert_called() + received = mock_callback.call_args[0][0] + self.assertIsInstance(received, ProcessGroupResourceUsage) + self.assertEqual(received.cpu_percent, 50.0) + self.assertEqual(received.rss_bytes, 1048576) + self.assertEqual(received.pss_bytes, 800000) + self.assertEqual(received.private_bytes, 600000) + self.assertEqual(received.disk_read_bytes, 4096) + self.assertEqual(received.disk_write_bytes, 8192) + self.assertEqual(received.num_procs, 3) + self.assertEqual(received.net_rx_bytes, 100) + self.assertEqual(received.net_tx_bytes, 200) diff --git a/src/spdl/pipeline/tests/pipeline_builder_test.py b/src/spdl/pipeline/tests/pipeline_builder_test.py new file mode 100644 index 000000000..d3f92bebd --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_builder_test.py @@ -0,0 +1,2763 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import asyncio +import functools +import os +import platform +import random +import re +import sys +import threading +import time +import unittest +import warnings +from collections.abc import Iterator +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from contextlib import asynccontextmanager +from functools import partial +from multiprocessing import Process +from typing import TypeVar + +from parameterized import parameterized +from spdl.pipeline import ( + AsyncQueue, + PipelineBuilder, + PipelineFailure, + run_pipeline_in_subprocess, + TaskHook, + TaskStatsHook, +) +from spdl.pipeline._components import _get_global_id, _set_global_id +from spdl.pipeline._components._common import _EOF, StageInfo +from spdl.pipeline._components._hook import _periodic_dispatch +from spdl.pipeline._components._pipe import ( + _FailCounter, + _get_fail_counter, + _pipe, + _PipeArgs, +) +from spdl.pipeline._components._sink import _sink +from spdl.pipeline._components._source import _source +from spdl.pipeline.defs import Aggregator +from spdl.source.utils import embed_shuffle + +T = TypeVar("T") + + +def _ignore_warnings(*filters): + """Decorator that wraps a test in `warnings.catch_warnings()` and applies + the given filters. Each ``filter`` is a dict of kwargs forwarded to + ``warnings.filterwarnings``. + """ + + def decorator(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + for f in filters: + warnings.filterwarnings("ignore", **f) + return fn(*args, **kwargs) + + return wrapper + + return decorator + + +_FORK_WARNING = { + "message": ( + r"This process \(pid=\d+\) is multi-threaded, use of fork\(\) " + r"may lead to deadlocks in the child" + ), + "category": DeprecationWarning, +} + +_RUN_PIPELINE_DEPRECATION = { + "message": ( + r"Passing a `PipelineBuilder` object directly to " + r"`run_pipeline_in_subprocess` is now deprecated\..*" + ), + "category": UserWarning, +} + +_UNAWAITED_COROUTINE = { + "message": "coroutine .* was never awaited", + "category": RuntimeWarning, +} + + +def _SI(name: str) -> StageInfo: + """Shorthand for creating a test StageInfo.""" + return StageInfo(pipeline_id=0, stage_id="0", stage_name=name) + + +def _put_aqueue(queue, vals, *, eof): + for val in vals: + queue.put_nowait(val) + if eof: + queue.put_nowait(_EOF) + + +def _flush_aqueue(queue): + ret = [] + while not queue.empty(): + ret.append(queue.get_nowait()) + return ret + + +async def no_op(val): + return val + + +################################################################################ +# _source +################################################################################ + + +class TestSource(unittest.TestCase): + def test_async_enqueue_empty(self) -> None: + """_async_enqueue can handle empty iterator""" + queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 + ) + coro = _source([], queue) + asyncio.run(coro) + self.assertEqual(_flush_aqueue(queue), [_EOF]) + + def test_async_enqueue_simple(self) -> None: + """_async_enqueue should put the values in the queue.""" + src = list(range(6)) + queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 + ) + coro = _source(src, queue) + asyncio.run(coro) + vals = _flush_aqueue(queue) + self.assertEqual(vals, [*src, _EOF]) + + def test_async_enqueue_iterator_failure(self) -> None: + """When `iterator` fails, the exception is propagated.""" + + def src(): + yield from range(10) + raise RuntimeError("Failing the iterator.") + + coro = _source( + src(), + AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 + ), + ) + + with self.assertRaises(RuntimeError): + asyncio.run(coro) # Not raising + + def test_async_enqueue_cancel(self) -> None: + """_async_enqueue is cancellable.""" + + async def _test(): + queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=1 + ) + + src = list(range(3)) + + coro = _source(src, queue) + task = asyncio.create_task(coro) + + await asyncio.sleep(0.1) + + task.cancel() + + with self.assertRaises(asyncio.CancelledError): + await task + + asyncio.run(_test()) + + +################################################################################ +# _sink +################################################################################ + + +class TestSink(unittest.TestCase): + @parameterized.expand( + [ + (False,), + (True,), + ] + ) + def test_async_sink_simple(self, empty: bool) -> None: + """_sink pass the contents from input_queue to output_queue""" + input_queue: AsyncQueue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue: AsyncQueue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + data = [] if empty else list(range(3)) + _put_aqueue(input_queue, data, eof=True) + + coro = _sink(input_queue, output_queue) + + asyncio.run(coro) + results = _flush_aqueue(output_queue) + + self.assertEqual(results, data) + + def test_async_sink_cancel(self) -> None: + """_async_sink is cancellable.""" + + async def _test(): + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input") + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output") + ) + + coro = _sink(input_queue, output_queue) + task = asyncio.create_task(coro) + + await asyncio.sleep(0.1) + + task.cancel() + + with self.assertRaises(asyncio.CancelledError): + await task + + asyncio.run(_test()) + + +################################################################################ +# _pipe +################################################################################ + + +async def adouble(val: int): + return 2 * val + + +async def aplus1(val: int): + return val + 1 + + +async def passthrough(val): + print("passthrough:", val) + return val + + +class TestPipe(unittest.IsolatedAsyncioTestCase): + def test_async_pipe(self) -> None: + """_pipe processes the data in input queue and pass it to output queue.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + async def test(): + ref = list(range(6)) + _put_aqueue(input_queue, ref, eof=True) + + await _pipe( + _SI("adouble"), + input_queue, + output_queue, + _PipeArgs(op=adouble), + _FailCounter(), + [], + False, + ) + + result = _flush_aqueue(output_queue) + + self.assertEqual(result, [v * 2 for v in ref] + [_EOF]) + + asyncio.run(test()) + + def test_async_pipe_skip(self) -> None: + """_pipe skips the result if it's None.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + async def skip_even(v): + if v % 2: + return v + + async def test(): + _put_aqueue(input_queue, range(10), eof=True) + + await _pipe( + _SI("skip_even"), + input_queue, + output_queue, + _PipeArgs(op=skip_even), + _FailCounter(), + [], + False, + ) + + result = _flush_aqueue(output_queue) + + self.assertEqual(result, [*list(range(1, 10, 2)), _EOF]) + + asyncio.run(test()) + + def test_async_pipe_wrong_task_signature(self) -> None: + """_pipe fails immediately if user provided incompatible iterator/afunc.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 + ) + + async def _2args(val: int, _): + return val + + async def test(): + ref = list(range(6)) + _put_aqueue(input_queue, ref, eof=False) + + with self.assertRaises(TypeError): + await _pipe( + _SI("_2args"), + input_queue, + output_queue, + _PipeArgs(op=_2args, concurrency=3), + _FailCounter(), + [], + False, + ) + + remaining = _flush_aqueue(input_queue) + self.assertEqual(remaining, ref[1:]) + + result = _flush_aqueue(output_queue) + self.assertEqual(result, [_EOF]) + + asyncio.run(test()) + + @parameterized.expand( + [ + (False,), + (True,), + ] + ) + def test_async_pipe_cancel(self, full: bool) -> None: + """_pipe is cancellable.""" + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=1 + ) + + _put_aqueue(input_queue, list(range(3)), eof=False) + + if full: + output_queue.put_nowait(None) + + cancelled = False + + async def astuck(i): + try: + await asyncio.sleep(10) + return i + except asyncio.CancelledError: + nonlocal cancelled + cancelled = True + raise + + async def test(): + coro = _pipe( + _SI("astuck"), + input_queue, + output_queue, + _PipeArgs(op=astuck), + _FailCounter(), + [], + False, + ) + task = asyncio.create_task(coro) + + await asyncio.sleep(0.5) + + task.cancel() + + with self.assertRaises(asyncio.CancelledError): + await task + + self.assertFalse(cancelled) + asyncio.run(test()) + self.assertTrue(cancelled) + + def test_async_pipe_concurrency(self) -> None: + """Changing concurrency changes the number of items fetched and processed.""" + + async def delay(val): + await asyncio.sleep(0.5) + return val + + async def test(concurrency): + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), + buffer_size=0, + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), + buffer_size=0, + ) + + ref = [1, 2, 3, 4] + _put_aqueue(input_queue, ref, eof=False) + + coro = _pipe( + _SI("delay"), + input_queue, + output_queue, + _PipeArgs( + op=delay, + concurrency=concurrency, + ), + _FailCounter(), + [], + False, + ) + + task = asyncio.create_task(coro) + await asyncio.sleep(0.8) + task.cancel() + + return _flush_aqueue(input_queue), _flush_aqueue(output_queue) + + # With concurrency==1, there should be + # 1 in output_queue, 2 is in flight, 3 and 4 remain in input_queue + remain, output = asyncio.run(test(1)) + self.assertEqual(remain, [3, 4]) + self.assertEqual(output, [1]) + + # With concurrency==4, there should be + # 1, 2, 3 and 4 in output_queue. + remain, output = asyncio.run(test(4)) + self.assertEqual(remain, []) + self.assertEqual(set(output), {1, 2, 3, 4}) + + def test_async_pipe_concurrency_throughput(self) -> None: + """increasing concurrency improves the throughput.""" + + async def delay(val): + await asyncio.sleep(0.5) + return val + + async def test(concurrency): + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), + buffer_size=0, + ) + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), + buffer_size=0, + ) + + ref = [4, 5, 6, 7, _EOF] + _put_aqueue(input_queue, ref, eof=False) + + t0 = time.monotonic() + await _pipe( + _SI("delay"), + input_queue, + output_queue, + _PipeArgs( + op=delay, + concurrency=concurrency, + ), + _FailCounter(), + [], + False, + ) + elapsed = time.monotonic() - t0 + + result = _flush_aqueue(output_queue) + + self.assertEqual(set(result), set(ref)) + self.assertEqual(result[-1], ref[-1]) + self.assertEqual(result[-1], _EOF) + + return elapsed + + elapsed1 = asyncio.run(test(1)) + elapsed4 = asyncio.run(test(4)) + + self.assertGreater(elapsed1, 1.8) + self.assertLess(elapsed4, 1) + + +################################################################################ +# Pipeline +################################################################################ + + +class TestPipeline(unittest.TestCase): + def test_pipeline_stage_hook_wrong_def1(self) -> None: + """Pipeline fails if stage_hook is not properly overrode.""" + + class _hook(TaskHook): + # missing asynccontextmanager + async def stage_hook(self): + yield + + @asynccontextmanager + async def task_hook(self, input_item=None): + yield + + with self.assertRaises(ValueError): + ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink() + # pyre-ignore + .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) + ) + + def test_pipeline_stage_hook_wrong_def2(self) -> None: + """Pipeline fails if task_hook is not properly overrode.""" + + class _hook(TaskHook): + # missing asynccontextmanager and async keyword + def stage_hook(self): + yield + + @asynccontextmanager + async def task_hook(self, input_item=None): + yield + + with self.assertRaises(ValueError): + ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink() + # pyre-ignore + .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) + ) + + +class CountHook(TaskHook): + def __init__(self): + self._enter_task_called = 0 + self._enter_stage_called = 0 + self._exit_task_called = 0 + self._exit_stage_called = 0 + + @asynccontextmanager + async def stage_hook(self): + self._enter_stage_called += 1 + yield + self._exit_stage_called += 1 + + @asynccontextmanager + async def task_hook(self, input_item=None): + self._enter_task_called += 1 + try: + yield + finally: + self._exit_task_called += 1 + + +class TestPipelineHook(unittest.TestCase): + @parameterized.expand( + [ + (False,), + (True,), + ] + ) + def test_pipeline_hook_drop_last(self, drop_last: bool) -> None: + """Hook is executed properly""" + + h1, h2, h3 = CountHook(), CountHook(), CountHook() + + def hook_factory(name) -> list[TaskHook]: + sname = str(name) + if "adouble" in sname: + return [h1] + if "aggregate" in sname: + return [h2] + if "_fail" in sname: + return [h3] + raise RuntimeError(f"Unexpected name: {sname}") + + async def _fail(_): + raise RuntimeError("Failing") + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(adouble) + .aggregate(5, drop_last=drop_last) + .pipe(_fail) + .add_sink(1000) + .build(num_threads=1, task_hook_factory=hook_factory) + ) + + with pipeline.auto_stop(): + self.assertEqual([], list(pipeline.get_iterator(timeout=10))) + + self.assertEqual(h1._enter_stage_called, 1) + self.assertEqual(h1._exit_stage_called, 1) + self.assertEqual(h1._enter_task_called, 10) + self.assertEqual(h1._exit_task_called, 10) + + self.assertEqual(h2._enter_stage_called, 1) + self.assertEqual(h2._exit_stage_called, 1) + # When drop_last=False, EOF is passed to the aggregation operator (11 calls: 10 items + 1 EOF) + # When drop_last=True, EOF is NOT passed to the aggregation operator (10 calls: 10 items only) + expected_h2_calls = 10 if drop_last else 11 + self.assertEqual(h2._enter_task_called, expected_h2_calls) + self.assertEqual(h2._exit_task_called, expected_h2_calls) + + # Even when the stage task fails, + # the exit_stage and exit_task are still called. + self.assertEqual(h3._enter_stage_called, 1) + self.assertEqual(h3._exit_stage_called, 1) + self.assertEqual(h3._enter_task_called, 2) + self.assertEqual(h3._exit_task_called, 2) + + def test_pipeline_hook_multiple(self) -> None: + """Multiple hooks are executed properly""" + + class _hook(TaskHook): + def __init__(self): + self._enter_task_called = 0 + self._enter_stage_called = 0 + self._exit_task_called = 0 + self._exit_stage_called = 0 + + @asynccontextmanager + async def stage_hook(self): + self._enter_stage_called += 1 + yield + self._exit_stage_called += 1 + + @asynccontextmanager + async def task_hook(self, input_item=None): + self._enter_task_called += 1 + try: + yield + finally: + self._exit_task_called += 1 + + hooks = [_hook(), _hook(), _hook()] + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: hooks) + ) + + with pipeline.auto_stop(): + self.assertEqual(list(range(10)), list(pipeline.get_iterator(timeout=10))) + + for h in hooks: + self.assertEqual(h._enter_stage_called, 1) + self.assertEqual(h._exit_stage_called, 1) + self.assertEqual(h._enter_task_called, 10) + self.assertEqual(h._exit_task_called, 10) + + @_ignore_warnings({"category": RuntimeWarning}) + @_ignore_warnings(_UNAWAITED_COROUTINE) + def test_pipeline_hook_failure_enter_stage(self) -> None: + """If enter_stage fails, the pipeline is aborted.""" + + class _enter_stage_fail(TaskHook): + @asynccontextmanager + async def stage_hook(self): + raise RuntimeError("failing") + + @asynccontextmanager + async def task_hook(self, input_item=None): + yield + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_enter_stage_fail()]) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(vals, []) + + @_ignore_warnings({"category": RuntimeWarning}) + @_ignore_warnings(_UNAWAITED_COROUTINE) + def test_pipeline_hook_failure_exit_stage(self) -> None: + """If exit_stage fails, the error is propagated to the front end.""" + + class _exit_stage_fail(TaskHook): + @asynccontextmanager + async def stage_hook(self): + yield + raise RuntimeError("failing") + + @asynccontextmanager + async def task_hook(self, input_item=None): + yield + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_exit_stage_fail()]) + ) + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(vals, list(range(10))) + + @_ignore_warnings({"category": RuntimeWarning}) + def test_pipeline_hook_failure_enter_task(self) -> None: + """If enter_task fails, the pipeline does not fail.""" + + class _hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + raise RuntimeError("failing enter_task") + + @asynccontextmanager + async def stage_hook(self, *_): + yield + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) + ) + + with pipeline.auto_stop(): + self.assertEqual([], list(pipeline.get_iterator(timeout=10))) + + @_ignore_warnings({"category": RuntimeWarning}) + def test_pipeline_hook_failure_exit_task(self) -> None: + """If exit_task fails, the pipeline does not fail. + + IMPORTANT: The result is dropped. + """ + + class _exit_stage_fail(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + yield + raise RuntimeError("failing exit_task") + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_exit_stage_fail()]) + ) + + with pipeline.auto_stop(): + self.assertEqual(list(pipeline.get_iterator(timeout=10)), []) + + def test_pipeline_hook_exit_task_capture_error(self) -> None: + """If task fails exit_task captures the error.""" + + exc_info = None + + class _capture(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + try: + yield + except Exception as e: + nonlocal exc_info + exc_info = e + + err = RuntimeError("failing") + + async def _fail(_): + raise err + + pipeline = ( + PipelineBuilder() + .add_source([None]) + .pipe(_fail) + .add_sink(100) + .build( + num_threads=1, + # pyre-ignore[6] + task_hook_factory=lambda _: [_capture()], + ) + ) + + with pipeline.auto_stop(): + self.assertEqual(list(pipeline.get_iterator(timeout=10)), []) + + self.assertTrue(exc_info is err) + + def test_pipeline_hook_receives_input_item(self) -> None: + """task_hook receives the input_item being processed.""" + + received_items = [] + + class _item_capture_hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + received_items.append(input_item) + yield + + pipeline = ( + PipelineBuilder() + .add_source(range(5)) + .pipe(passthrough) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_item_capture_hook()]) + ) + + with pipeline.auto_stop(): + output = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(output, list(range(5))) + self.assertEqual(received_items, list(range(5))) + + def test_pipeline_hook_receives_input_item_on_failure(self) -> None: + """task_hook receives input_item even when the task fails.""" + + captured_items_on_failure = [] + + class _failure_capture_hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + try: + yield + except StopAsyncIteration: + raise + except Exception: + captured_items_on_failure.append(input_item) + raise + + def fail_on_even(x: int) -> int: + if x % 2 == 0: + raise RuntimeError(f"fail on {x}") + return x + + pipeline = ( + PipelineBuilder() + .add_source(range(6)) + .pipe(fail_on_even) + .add_sink(1000) + .build( + num_threads=1, + task_hook_factory=lambda _: [_failure_capture_hook()], # pyre-ignore[6] + ) + ) + + with pipeline.auto_stop(): + output = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(output, [1, 3, 5]) + self.assertEqual(captured_items_on_failure, [0, 2, 4]) + + def test_ordered_pipe_hook_receives_input_item(self) -> None: + """task_hook receives input_item in ordered pipe.""" + + received_items = [] + + class _item_capture_hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + received_items.append(input_item) + yield + + pipeline = ( + PipelineBuilder() + .add_source(range(5)) + .pipe(passthrough, output_order="input", concurrency=4) + .add_sink(1000) + # pyre-ignore[6] + .build(num_threads=4, task_hook_factory=lambda _: [_item_capture_hook()]) + ) + + with pipeline.auto_stop(): + output = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(output, list(range(5))) + self.assertEqual(received_items, list(range(5))) + + def test_ordered_pipe_hook_receives_input_item_on_failure(self) -> None: + """task_hook receives input_item on failure in ordered pipe.""" + + captured_items_on_failure = [] + + class _failure_capture_hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + try: + yield + except StopAsyncIteration: + raise + except Exception: + captured_items_on_failure.append(input_item) + raise + + def fail_on_even(x: int) -> int: + if x % 2 == 0: + raise RuntimeError(f"fail on {x}") + return x + + pipeline = ( + PipelineBuilder() + .add_source(range(6)) + .pipe(fail_on_even, output_order="input", concurrency=4) + .add_sink(1000) + .build( + num_threads=4, + task_hook_factory=lambda _: [_failure_capture_hook()], # pyre-ignore[6] + ) + ) + + with pipeline.auto_stop(): + output = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(output, [1, 3, 5]) + self.assertEqual(sorted(captured_items_on_failure), [0, 2, 4]) + + +################################################################################ +# TaskStatsHook +################################################################################ + + +class TestTaskStatsHook(unittest.TestCase): + def test_task_stats(self) -> None: + """TaskStatsHook logs the interval of each task.""" + + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), 1 + ) + + async def _test(): + async with hook.stage_hook(): + for _ in range(3): + async with hook.task_hook(): + await asyncio.sleep(0.5) + + self.assertEqual(hook.num_tasks, 3) + self.assertEqual(hook.num_success, 3) + self.assertGreater(hook.ave_time, 0.3) + self.assertLess(hook.ave_time, 0.7) + + for _ in range(2): + with self.assertRaises(RuntimeError): + async with hook.task_hook(): + await asyncio.sleep(1.0) + raise RuntimeError("failing") + + self.assertEqual(hook.num_tasks, 5) + self.assertEqual(hook.num_success, 3) + self.assertGreater(hook.ave_time, 0.45) + self.assertLess(hook.ave_time, 0.9) + + asyncio.run(_test()) + + +class TestPeriodicDispatch(unittest.TestCase): + def test_periodic_dispatch_smoke_test(self) -> None: + """_periodic_dispatch runs functions with the given interval.""" + + calls = [] + + async def afun(): + print("afun: ", time.time()) + calls.append(time.monotonic()) + + async def _test(): + done = asyncio.Event() + task = asyncio.create_task(_periodic_dispatch(afun, done, 1)) + + await asyncio.sleep(3.2) + + done.set() + await task + + print("start: ", time.time()) + asyncio.run(_test()) + + self.assertEqual(len(calls), 3) + self.assertGreater(calls[1] - calls[0], 0.9) + self.assertLess(calls[1] - calls[0], 1.1) + self.assertGreater(calls[2] - calls[1], 0.9) + self.assertLess(calls[2] - calls[1], 1.1) + + def test_task_stats_log_interval_stats(self) -> None: + """Smoke test for _log_interval_stats.""" + + hook = TaskStatsHook( + StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), 1 + ) + asyncio.run(hook._log_interval_stats()) + + +################################################################################ +# __str__ +################################################################################ + + +class TestPipelineStr(unittest.TestCase): + def test_pipeline_str_smoke(self) -> None: + async def passthrough(i): + return i + + builder = PipelineBuilder() + + print(builder) + + builder = builder.add_source(range(10)) + + print(builder) + + builder = builder.pipe(passthrough) + + print(builder) + + builder = builder.aggregate(1) + + print(builder) + + builder = builder.pipe(passthrough, output_order="input") + + print(builder) + + builder = builder.aggregate(1) + + print(builder) + + builder = builder.add_sink(100) + + print(builder) + + +################################################################################ +# AsyncPipeline - resume +################################################################################ + + +class TestPipelineResume(unittest.TestCase): + def test_pipeline_reiterate(self) -> None: + """Pipeline can be iterated multiple times as long as it's not stopped""" + + pipeline = ( + PipelineBuilder().add_source(range(20)).add_sink(1000).build(num_threads=1) + ) + + with pipeline.auto_stop(): + for i in range(5): + for j, val in enumerate(pipeline.get_iterator(timeout=10)): + self.assertEqual(val, (i * 4) + j) + + # Now it's empty + with self.assertRaises(StopIteration): + next(pipeline.get_iterator(timeout=10)) + + def test_pipeline_resume(self) -> None: + """AsyncPipeline can execute the source partially, then resumed""" + + # Note + # If we pass `range(10)` directly, new iterator is created at every run. + src = iter(range(10)) + + pipeline = PipelineBuilder().add_source(src).add_sink(1000).build(num_threads=1) + + with pipeline.auto_stop(): + iterator = pipeline.get_iterator(timeout=10) + self.assertEqual([0, 1], [next(iterator) for _ in range(2)]) + + iterator = pipeline.get_iterator(timeout=10) + self.assertEqual([2, 3, 4], [next(iterator) for _ in range(3)]) + + iterator = pipeline.get_iterator(timeout=10) + self.assertEqual([5, 6, 7, 8, 9], [next(iterator) for _ in range(5)]) + + with self.assertRaises(StopIteration): + next(iterator) + + def test_pipeline_infinite_loop(self) -> None: + """AsyncPipeline can execute infinite iterable""" + + def src(i=-1): + while True: + yield (i := i + 1) + + pipeline = ( + PipelineBuilder().add_source(src()).add_sink(1000).build(num_threads=1) + ) + + with pipeline.auto_stop(): + i = 0 + for _ in range(10): + num_items = random.randint(1, 128) + for j, item in enumerate(pipeline.get_iterator(timeout=10)): + self.assertEqual(item, i) + i += 1 + + if num_items == j: + break + + +################################################################################ +# AsyncPipeline - order +################################################################################ + + +class TestPipelineOrder(unittest.TestCase): + def test_pipeline_order_complete(self) -> None: + """The output is in the order of completion.""" + + async def _sleep(i): + await asyncio.sleep(i / 10) + return i + + src = list(reversed(range(10))) + pipeline = ( + PipelineBuilder() + .add_source(src) + .pipe(_sleep, concurrency=10, output_order="completion") + .add_sink(100) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + self.assertEqual(list(pipeline.get_iterator(timeout=10)), list(range(10))) + + def test_pipeline_order_input(self) -> None: + """The output is in the order of the input.""" + + async def _sleep(i): + print(f"Sleeping: {i}") + await asyncio.sleep(i / 10) + print(f"Returning: {i}") + return i + + src = list(reversed(range(10))) + pipeline = ( + PipelineBuilder() + .add_source(src) + .pipe(_sleep, concurrency=10, output_order="input") + .add_sink(100) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + self.assertEqual(src, list(pipeline.get_iterator(timeout=10))) + + def test_pipeline_order_input_sync_func(self) -> None: + """The output is in the order of the input.""" + + def _sleep(i): + print(f"Sleeping: {i}") + time.sleep(i / 10) + print(f"Returning: {i}") + return i + + src = list(reversed(range(10))) + pipeline = ( + PipelineBuilder() + .add_source(src) + .pipe(_sleep, concurrency=10, output_order="input") + .add_sink(100) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + self.assertEqual(list(pipeline.get_iterator(timeout=10)), src) + + @_ignore_warnings({"category": RuntimeWarning}) + def test_pipeline_order_input_filter_none(self) -> None: + """Ordered pipe filters out None values returned by the pipe operation.""" + + pipeline = ( + PipelineBuilder() + .add_source(list(range(10))) + .pipe(lambda x: None if x % 2 == 0 else x, output_order="input") + .add_sink(2) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(result, [1, 3, 5, 7, 9]) + + def test_pipeline_order_input_filter_none_async(self) -> None: + """Ordered pipe filters out None values with async function.""" + + async def filter_even(x): + await asyncio.sleep(0.01) + return None if x % 2 == 0 else x + + pipeline = ( + PipelineBuilder() + .add_source(list(range(10))) + .pipe(filter_even, output_order="input", concurrency=3) + .add_sink(2) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(result, [1, 3, 5, 7, 9]) + + def test_pipeline_order_input_filter_none_with_concurrency(self) -> None: + """Ordered pipe filters out None values with high concurrency.""" + + def slow_filter(x): + time.sleep(0.05) + return None if x % 3 == 0 else x + + pipeline = ( + PipelineBuilder() + .add_source(list(range(15))) + .pipe(slow_filter, output_order="input", concurrency=5) + .add_sink(10) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=5)) + # Filters out 0, 3, 6, 9, 12 + self.assertEqual(result, [1, 2, 4, 5, 7, 8, 10, 11, 13, 14]) + + def test_pipeline_order_input_all_none(self) -> None: + """Ordered pipe handles case where all values are None.""" + + pipeline = ( + PipelineBuilder() + .add_source(list(range(5))) + .pipe(lambda _: None, output_order="input") + .add_sink(2) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(result, []) + + def test_pipeline_order_input_mixed_none_and_values(self) -> None: + """Ordered pipe correctly handles mixed None and values in specific pattern.""" + + def pattern_filter(x): + if x < 2: + return None + if x < 5: + return x + if x < 7: + return None + return x + + pipeline = ( + PipelineBuilder() + .add_source(list(range(10))) + .pipe(pattern_filter, output_order="input") + .add_sink(5) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + result = list(pipeline.get_iterator(timeout=10)) + # Returns 2, 3, 4 (x < 5), and 7, 8, 9 (x >= 7) + self.assertEqual(result, [2, 3, 4, 7, 8, 9]) + + +################################################################################ +# AsyncPipeline2 +################################################################################ + + +class TestPipelineNoop(unittest.TestCase): + def test_pipeline_noop(self) -> None: + """AsyncPipeline2 functions without pipe.""" + + apl = PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) + + with apl.auto_stop(): + for i in range(10): + print("fetching", i) + self.assertEqual(i, apl.get_item(timeout=1)) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelinePassthrough(unittest.TestCase): + def test_pipeline_passthrough(self) -> None: + """AsyncPipeline2 can passdown items operation.""" + + apl = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1) + .build(num_threads=1) + ) + + with apl.auto_stop(): + for i in range(10): + print("fetching", i) + self.assertEqual(i, apl.get_item(timeout=1)) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelineSkip(unittest.TestCase): + def test_pipeline_skip(self) -> None: + """AsyncPipeline2 does not output None items.""" + + src = list(range(10)) + + async def odd(i): + if i % 2: + return i + + pipeline = ( + PipelineBuilder() + .add_source(src) + .pipe(odd) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for i in range(5): + self.assertEqual(i * 2 + 1, pipeline.get_item(timeout=10)) + + with self.assertRaises(EOFError): + pipeline.get_item(timeout=10) + + +class TestPipelineLambda(unittest.TestCase): + def test_pipeline_lambda(self) -> None: + """AsyncPipeline2 pipe supports lambda items operation.""" + + apl = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x) + .add_sink(1) + .build(num_threads=1) + ) + + with apl.auto_stop(): + for i in range(10): + print("fetching", i) + self.assertEqual(i, apl.get_item(timeout=1)) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelineSimple(unittest.TestCase): + def test_pipeline_simple(self) -> None: + """AsyncPipeline2 can perform simple operation.""" + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(adouble) + .pipe(aplus1) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + for i, item in enumerate(pipeline.get_iterator(timeout=10)): + self.assertEqual(item, i * 2 + 1) + + +class TestPipelineAggregate(unittest.TestCase): + def test_pipeline_aggregate(self) -> None: + """AsyncPipeline aggregates the input""" + + src = list(range(13)) + + pipeline = ( + PipelineBuilder() + .add_source(src) + .aggregate(4) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual( + results, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12]] + ) + + def test_pipeline_aggregate_drop_last(self) -> None: + """AsyncPipeline aggregates the input and drop the last""" + + src = list(range(13)) + + pipeline = ( + PipelineBuilder() + .add_source(src) + .aggregate(4, drop_last=True) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(results, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]) + + @parameterized.expand([(False,), (True,)]) + def test_pipeline_aggregate_custom_op(self, drop_last: bool) -> None: + """AsyncPipeline aggregates with custom operation that concatenates when threshold is exceeded""" + + # Custom aggregation: concatenate strings when total size exceeds threshold + class CustomAggregator(Aggregator): + def __init__(self, size_threshold: int = 10): + self.size_threshold = size_threshold + self.buffer: list[str] = [] + self.total_size = 0 + + def _flush(self) -> str: + result = "".join(self.buffer) + self.buffer = [] + self.total_size = 0 + return result + + def flush(self) -> str | None: + # Emit remaining buffer when EOF is reached + if self.buffer: + return self._flush() + return None # _SKIP + + def accumulate(self, item: str) -> str | None: + self.buffer.append(item) + self.total_size += len(item) + + if self.total_size >= self.size_threshold: + return self._flush() + return None # _SKIP + + src = ["a", "bb", "ccc", "dddd", "e", "ff", "ggg", "h"] + + pipeline = ( + PipelineBuilder() + .add_source(src) + .aggregate(CustomAggregator(size_threshold=10), drop_last=drop_last) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + # "a", "bb", "ccc", "dddd" = 10 chars -> first result + if drop_last: + # When drop_last=True, flush is not called, + # so the remaining buffer ["e", "ff", "ggg", "h"] is dropped + self.assertEqual(results, ["abbcccdddd"]) + else: + # When drop_last=False, flush is called, + # so remaining buffer is emitted: "e", "ff", "ggg", "h" + self.assertEqual(results, ["abbcccdddd", "effgggh"]) + + +class TestPipelineDisaggregate(unittest.TestCase): + def test_pipeline_disaggregate(self) -> None: + """AsyncPipeline disaggregates the input""" + + src = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12]] + + pipeline = ( + PipelineBuilder() + .add_source(src) + .disaggregate() + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(results, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]) + + +class TestPipelineSource(unittest.TestCase): + def test_pipeline_source_failure(self) -> None: + """AsyncPipeline continues when source fails. + + note: the front end will propagate the error at the end of the `stop`. + before that, the pipeline should continue functioning. + """ + + def failing_range(i): + yield from range(i) + raise ValueError("Iterator failed") + + pipeline = ( + PipelineBuilder() + .add_source(failing_range(10)) + .pipe(adouble) + .pipe(aplus1) + .add_sink(1000) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, [1 + 2 * i for i in range(10)]) + + +class TestPipelineType(unittest.TestCase): + def test_pipeline_type_error(self) -> None: + """AsyncPipeline immediately fails if pipe function has wrong signature""" + + async def wrong_sig(i, _): + return i + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(wrong_sig) + .add_sink(1000) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(vals, []) + + +class TestPipelineTask(unittest.TestCase): + def test_pipeline_task_failure(self) -> None: + """AsyncPipeline is robust against task-level failure.""" + + async def areject_m3(i): + if i % 3 == 0: + raise ValueError(f"Multiple of 3 is prohibited: {i}") + return i + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(areject_m3) + .pipe(adouble) + .pipe(aplus1) + .add_sink(1000) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=10)) + self.assertEqual(results, [1 + 2 * i for i in range(10) if i % 3]) + + +class TestPipelineCancel(unittest.TestCase): + def test_pipeline_cancel_empty(self) -> None: + """AsyncPipeline2 can be cancelled while it's blocked on the pipeline.""" + + apl = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1) + .build(num_threads=1) + ) + # The nuffer and coroutinues will be blocked holding items as follows: + # + # [src] --|queue(1)|--> [passthrough] --|queue(1)|--> [sink] -->|queue(1)| + # i+5 i+4 i+3 i+2 i+1 i + # + + with apl.auto_stop(): + for i in range(5): + print("fetching", i) + self.assertEqual(i, apl.get_item(timeout=1)) + + # Ensure that buffers are filled and the pipeline is blocked. + time.sleep(0.1) + # At this point, the output queue holds 5. + + # Only the "5" is retrievable. + self.assertEqual(5, apl.get_item(timeout=1)) + + # The background thread is stopped, so no more data is coming. + for _ in range(3): + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelineFail(unittest.TestCase): + def test_pipeline_fail_middle(self) -> None: + """When a stage in the middle fails, downstream stages are not failing.""" + + async def fail(i, _): + return i + + class PassthroughWithCache: + def __init__(self): + self.cache = [] + + async def __call__(self, i): + self.cache.append(i) + return i + + pwc = PassthroughWithCache() + + apl = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .pipe(fail) + .pipe(pwc) + .add_sink(1) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with apl.auto_stop(): + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + self.assertEqual(pwc.cache, []) + + # The background thread is stopped, and the output queue is empty. + for _ in range(3): + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelineEof(unittest.TestCase): + def test_pipeline_eof_stop(self) -> None: + """APL2 can be closed after reaching EOF.""" + apl = ( + PipelineBuilder() + .add_source(range(2)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + with apl.auto_stop(): + for i in range(2): + print("fetching", i) + self.assertEqual(i, apl.get_item(timeout=1)) + + with self.assertRaises(EOFError): + apl.get_item(timeout=1) + + +class TestPipelineIterator(unittest.TestCase): + def test_pipeline_iterator(self) -> None: + """Can iterate the pipeline.""" + + apl = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(passthrough) + .add_sink(1) + .build(num_threads=1) + ) + + with apl.auto_stop(): + for i, item in enumerate(apl.get_iterator(timeout=1)): + print(i, item) + self.assertEqual(i, item) + + +class TestPipelineIter(unittest.TestCase): + def test_pipeline_iter_and_next(self) -> None: + """Pipeline `iter` and `next` fulfill the following contracts + + 1. `iter(pipeline)` creates a new iterator. + 2. `next(iterator)` gives the next item in the pipeline. + 3. one ca call `next` multiple times on an iterator. + 4. If iterator is exhausted, it raises `StopIteration`. + 5. Iterator object can be discarded without an side-effect by itself. + (Other side-effects might be happening but it's independent from iterator) + 6. Multiple instances of iterators can be created by + repeatedly calling `iter(pipeline)`. + 7. We do not define the beheviors for multiple iterators exist at the same + time. We only consider the case where one iterator is created and + discarded, then another is created. + 8. A new iterator should return items that are sequel to items generated by + the previous iterator. + """ + + apl = PipelineBuilder().add_source(range(12)).add_sink(1).build(num_threads=1) + + with apl.auto_stop(): + iterator = iter(apl) + self.assertEqual(next(iterator), 0) + self.assertEqual(next(iterator), 1) + self.assertEqual(next(iterator), 2) + + iterator = iter(apl) + self.assertEqual(next(iterator), 3) + self.assertEqual(next(iterator), 4) + self.assertEqual(next(iterator), 5) + + iterator = iter(apl) + self.assertEqual(next(iterator), 6) + self.assertEqual(next(iterator), 7) + self.assertEqual(next(iterator), 8) + + iterator = iter(apl) + self.assertEqual(next(iterator), 9) + self.assertEqual(next(iterator), 10) + self.assertEqual(next(iterator), 11) + + iterator = iter(apl) + with self.assertRaises(StopIteration): + next(iterator) + + +class TestPipelineStuck(unittest.TestCase): + def test_pipeline_stuck(self) -> None: + """`get_item` waits for slow pipeline.""" + + async def delay(i): + print(f"Sleeping: {i}") + await asyncio.sleep(0.5) + print(f"Sleeping: {i} - done") + return i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(delay) + .add_sink(1) + .build(num_threads=1) + ) + + with apl.auto_stop(): + for i, item in enumerate(apl.get_iterator(timeout=10)): + print(i, item) + self.assertEqual(i, item) + + +class TestPipelinePipe(unittest.TestCase): + def test_pipeline_pipe_agen(self) -> None: + """pipe works with async generator function""" + + async def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(expected, output) + + def test_pipeline_pipe_sync_gen(self) -> None: + """pipe works with sync generator function""" + + def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(expected, output) + + +class TestCallableGenerator(unittest.TestCase): + def test_callable_generator(self) -> None: + """pipe works with sync callable class returning generator""" + + class DupIncrement: + def __init__(self) -> None: + pass + + def __call__(self, v: int) -> Iterator[int]: + for i in range(3): + yield v + i + + dup_increment = DupIncrement() + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(expected, output) + + def test_pipeline_pipe_agen_max_failures(self) -> None: + """pipe works with async generator function and max_failure""" + + async def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1, max_failures=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(output, expected) + + def test_pipeline_pipe_gen(self) -> None: + """pipe works with sync generator function""" + + def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(output, expected) + + def test_pipeline_pipe_gen_max_failures(self) -> None: + """pipe works with sync generator function and max_failure""" + + def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1, max_failures=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(output, expected) + + @unittest.skipIf( + platform.system() == "Darwin" and "CI" in os.environ, + reason="GitHub macOS CI is not timely enough.", + ) + def test_pipeline_pipe_gen_incremental(self) -> None: + """pipe returns output of generator function immediately if not in ProcessPoolExecutor""" + + # We introduce delay in each iteration, so that, if the pipeline is returning + # the yielded value immediately, the output will be obtained quickly. + # If the pipeline is not returning the yielded value immediately, the output + # won't be available until the iteration ends, and by that time + # the foreground pipeline should timeout. + def dup_increment(v): + for i in range(3): + time.sleep(0.1) + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + .build(num_threads=1) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=0.2)) + self.assertEqual(output, expected) + + def test_pipeline_pipe_agen_wrong_hook(self) -> None: + """pipe works with async generator function, even when hook abosrb the StopAsyncIteration""" + + class _Hook(TaskHook): + @asynccontextmanager + async def task_hook(self, input_item=None): + try: + yield + except StopAsyncIteration: + pass + + async def dup_increment(v): + for i in range(3): + yield v + i + + apl = ( + PipelineBuilder() + .add_source(range(3)) + .pipe(dup_increment) + .add_sink(1) + # pyre-ignore[6] + .build(num_threads=1, task_hook_factory=lambda _: [_Hook()]) + ) + + expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(output, expected) + + def test_pipeline_source_agen(self) -> None: + """source works with async generator function""" + + async def source(): + for i in range(3): + yield i + + apl = PipelineBuilder().add_source(source()).add_sink(1).build(num_threads=1) + + expected = [0, 1, 2] + with apl.auto_stop(): + output = list(apl.get_iterator(timeout=10)) + self.assertEqual(output, expected) + + +class TestPipelineStart(unittest.TestCase): + def test_pipeline_start_multiple_times(self) -> None: + """`Pipeline.start` cannot be called multiple times.""" + + pipeline = ( + PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) + ) + + with pipeline.auto_stop(): + with self.assertRaises(RuntimeError): + pipeline.start() + + +class TestPipelineStop(unittest.TestCase): + def test_pipeline_stop_multiple_times(self) -> None: + """`Pipeline.stop` can be called multiple times.""" + + pipeline = ( + PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) + ) + + pipeline.stop() + pipeline.stop() + pipeline.stop() + + with pipeline.auto_stop(): + pipeline.stop() + pipeline.stop() + pipeline.stop() + + pipeline.stop() + pipeline.stop() + pipeline.stop() + + +def _run_pipeline_without_closing(): + pipeline = PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) + pipeline.start() + + +def get_pid(_): + import os + + time.sleep(0.5) + pid = os.getpid() + print(f"{pid=}") + return pid + + +def _range(item): + print(item) + for i in range(item): + print(f"yielding {item} - {i}") + yield i + + +class TestPipelineNo(unittest.TestCase): + def test_pipeline_no_close(self) -> None: + """Python interpreter can terminate even when Pipeline is not explicitly closed.""" + + p = Process(target=_run_pipeline_without_closing) + p.start() + p.join(timeout=10) + + if p.exitcode is None: + p.kill() + raise RuntimeError("Process did not self-terminate.") + + +class TestPipelineCustom(unittest.TestCase): + def test_pipeline_custom_pipe_executor(self) -> None: + """`pipe` accepts custom ThreadPoolExecutor. + + The primal goal of custom executor is to make it easy to use + thread local storages. + + So in this test, we initialize a custom executor with some thread + local storages, and then we access it without any check (hasattr) + in pipe function. + """ + num_threads = 10 + sleep = 0.5 + + ref = set(range(num_threads)) + + ref_copy = ref.copy() + thread_local_storage = threading.local() + + def init_storage(): + print("Initializing thread:", threading.get_ident()) + thread_local_storage.value = ref_copy.pop() + + executor = ThreadPoolExecutor( + max_workers=num_threads, + initializer=init_storage, + ) + + def op(i: int) -> int: + # sleep to block this thread, so that + # we use all the threads in the pool + time.sleep(sleep) + print(i, thread_local_storage.value) + return thread_local_storage.value + + pipeline = ( + PipelineBuilder() + .add_source(range(num_threads)) + .pipe(op, executor=executor, concurrency=num_threads) + .add_sink(1) + .build(num_threads=1) + ) + + t0 = time.monotonic() + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + elapsed = time.monotonic() - t0 + + self.assertEqual(0, len(ref_copy)) + self.assertEqual(ref, set(vals)) + self.assertLess(elapsed, sleep * 2) + + def test_pipeline_custom_pipe_executor_process(self) -> None: + """`pipe` accepts custom ProcessPoolExecutor.""" + num_processes = 5 + + executor = ProcessPoolExecutor(max_workers=num_processes) + + pipeline = ( + PipelineBuilder() + .add_source(range(num_processes)) + .pipe(get_pid, executor=executor, concurrency=num_processes) + .add_sink(1) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(num_processes, len(set(vals))) + + def test_pipeline_custom_pipe_executor_process_generator(self) -> None: + """`pipe` accepts custom ProcessPoolExecutor and generator function.""" + executor = ProcessPoolExecutor(max_workers=1) + + pipeline = ( + PipelineBuilder() + .add_source(range(4)) + .pipe(_range, executor=executor, concurrency=1) + .add_sink(1) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual([0, 0, 1, 0, 1, 2], vals) + + def test_pipeline_custom_pipe_executor_async(self) -> None: + """pipe rejects custom executor if op is async""" + + async def op(i: int) -> int: + return i + + with self.assertRaises(ValueError): + PipelineBuilder().add_source(range(10)).pipe( + op, executor=ThreadPoolExecutor() + ).add_sink(1).build(num_threads=1) + + def test_pipeline_pipe_list(self) -> None: + """pipe supports list as op.""" + + op = [i + 1 for i in range(10)] + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + # pyre-ignore[6] + .pipe(op) + .add_sink(1) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual(op, vals) + + def test_pipeline_pipe_tuple(self) -> None: + """pipe supports list as op.""" + + op = tuple(i + 1 for i in range(10)) + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + # pyre-ignore[6] + .pipe(op) + .add_sink(1) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual([i + 1 for i in range(10)], vals) + + def test_pipeline_pipe_dict(self) -> None: + """pipe supports dict as op.""" + + op = {i: i + 1 for i in range(10)} + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + # pyre-ignore[6] + .pipe(op) + .add_sink(1) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual([i + 1 for i in range(10)], vals) + + +class _PicklableSource: + def __init__(self, n: int) -> None: + self.n = n + + def __iter__(self) -> Iterator[int]: + yield from range(self.n) + + +class _ValidatePipelineId: + def __init__(self, val: int) -> None: + self.val = val + + def __iter__(self) -> Iterator[int]: + if (v := _get_global_id()) != self.val: + raise AssertionError(f"_node._PIPELINE_ID={v} != {self.val=}") + yield 0 + + +def plusN(x: int, N: int) -> int: + return x + N + + +def hook_factory(_: StageInfo) -> list[TaskHook]: + return [CountHook()] + + +class TestPipelinebuilderPicklable(unittest.TestCase): + @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) + def test_pipelinebuilder_picklable(self) -> None: + """PipelineBuilder can be passed to subprocess (==picklable)""" + + builder = ( + PipelineBuilder() + .add_source(_PicklableSource(10)) + .pipe( + adouble, + concurrency=3, + ) + .pipe( + aplus1, + concurrency=3, + ) + .pipe(partial(plusN, N=3)) + .pipe(passthrough) + .aggregate(3) + .disaggregate() + .add_sink(10) + ) + + results = list( + run_pipeline_in_subprocess( + # pyre-ignore[6] + builder, + num_threads=5, + buffer_size=-1, + task_hook_factory=hook_factory, + ) + ) + + def _ref(x: int) -> int: + return 2 * x + 1 + 3 + + self.assertEqual([_ref(i) for i in range(10)], sorted(results)) + + +class TestFailureCounter(unittest.TestCase): + def test_failure_counter_global_countes(self) -> None: + """_get_fail_counter creates _FailCounter subclass with different class valiable""" + + FC1 = _get_fail_counter() + FC2 = _get_fail_counter() + + fc1_1 = FC1(-1, -1) + fc1_2 = FC1(-1, -1) + + fc2_1 = FC2(-1, -1) + fc2_2 = FC2(-1, -1) + + self.assertEqual(0, _FailCounter._num_global_failures) + self.assertEqual(0, FC1._num_global_failures) + self.assertEqual(0, FC2._num_global_failures) + self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) + self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) + self.assertEqual(0, fc1_1._num_global_failures) + self.assertEqual(0, fc1_2._num_global_failures) + self.assertEqual(0, fc1_1._num_stage_failures) + self.assertEqual(0, fc1_2._num_stage_failures) + self.assertEqual(0, fc2_1._num_global_failures) + self.assertEqual(0, fc2_2._num_global_failures) + self.assertEqual(0, fc2_1._num_stage_failures) + self.assertEqual(0, fc2_2._num_stage_failures) + + fc1_1.__class__._num_global_failures += 1 + fc1_1._num_stage_failures += 1 + + self.assertEqual(0, _FailCounter._num_global_failures) + self.assertEqual(1, FC1._num_global_failures) + self.assertEqual(0, FC2._num_global_failures) + self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) + self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) + self.assertEqual(1, fc1_1._num_global_failures) + self.assertEqual(1, fc1_2._num_global_failures) + self.assertEqual(1, fc1_1._num_stage_failures) + self.assertEqual(0, fc1_2._num_stage_failures) + self.assertEqual(0, fc2_1._num_global_failures) + self.assertEqual(0, fc2_2._num_global_failures) + self.assertEqual(0, fc2_1._num_stage_failures) + self.assertEqual(0, fc2_2._num_stage_failures) + + fc1_1.__class__._num_global_failures += 1 + fc1_1._num_stage_failures += 1 + + self.assertEqual(0, _FailCounter._num_global_failures) + self.assertEqual(2, FC1._num_global_failures) + self.assertEqual(0, FC2._num_global_failures) + self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) + self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) + self.assertEqual(2, fc1_1._num_global_failures) + self.assertEqual(2, fc1_2._num_global_failures) + self.assertEqual(2, fc1_1._num_stage_failures) + self.assertEqual(0, fc1_2._num_stage_failures) + self.assertEqual(0, fc2_1._num_global_failures) + self.assertEqual(0, fc2_2._num_global_failures) + self.assertEqual(0, fc2_1._num_stage_failures) + self.assertEqual(0, fc2_2._num_stage_failures) + + fc1_2.__class__._num_global_failures += 1 + fc1_2._num_stage_failures += 1 + + self.assertEqual(0, _FailCounter._num_global_failures) + self.assertEqual(3, FC1._num_global_failures) + self.assertEqual(0, FC2._num_global_failures) + self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) + self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) + self.assertEqual(3, fc1_1._num_global_failures) + self.assertEqual(3, fc1_2._num_global_failures) + self.assertEqual(2, fc1_1._num_stage_failures) + self.assertEqual(1, fc1_2._num_stage_failures) + self.assertEqual(0, fc2_1._num_global_failures) + self.assertEqual(0, fc2_2._num_global_failures) + self.assertEqual(0, fc2_1._num_stage_failures) + self.assertEqual(0, fc2_2._num_stage_failures) + + fc2_1.__class__._num_global_failures += 1 + fc2_1._num_stage_failures += 1 + + self.assertEqual(0, _FailCounter._num_global_failures) + self.assertEqual(3, FC1._num_global_failures) + self.assertEqual(1, FC2._num_global_failures) + self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) + self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) + self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) + self.assertEqual(3, fc1_1._num_global_failures) + self.assertEqual(3, fc1_2._num_global_failures) + self.assertEqual(2, fc1_1._num_stage_failures) + self.assertEqual(1, fc1_2._num_stage_failures) + self.assertEqual(1, fc2_1._num_global_failures) + self.assertEqual(1, fc2_2._num_global_failures) + self.assertEqual(1, fc2_1._num_stage_failures) + self.assertEqual(0, fc2_2._num_stage_failures) + + +class TestPipelineMax(unittest.TestCase): + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_max_failures(self, output_order: str) -> None: + """max_failures stop the pipeline.""" + + def fail_odd(x): + if x % 2: + raise ValueError(f"Only evan numbers are allowed. {x}") + return x + + builder = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_odd, output_order=output_order) + .add_sink(1) + ) + + pipeline = builder.build(num_threads=1) + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual([0, 2, 4, 6, 8], vals) + + pipeline = builder.build(num_threads=1, max_failures=3) + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + + self.assertEqual([0, 2, 4, 6], vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_max_failures_multiple_pipeline(self, output_order: str) -> None: + """When using multiple pipelines with different error caps, they work + + Note: FailCounter uses class method to combines the errors + from the different stages. We create a different child class + for each pipeline construction so that class counter are separate for + each pipeline object. This test ensures that. + """ + + def fail_odd(x): + if x % 2: + raise ValueError(f"Only evan numbers are allowed. {x}") + return x + + src = range(10) + + builder = ( + PipelineBuilder() + .add_source(src) + .pipe(fail_odd, output_order=output_order) + .add_sink(1) + ) + + pipeline1 = builder.build(num_threads=1, max_failures=2) + pipeline2 = builder.build(num_threads=1, max_failures=3) + + with self.assertRaises(PipelineFailure): + with pipeline2.auto_stop(): + vals = list(pipeline2.get_iterator(timeout=10)) + + self.assertEqual([0, 2, 4, 6], vals) + + with self.assertRaises(PipelineFailure): + with pipeline1.auto_stop(): + vals = list(pipeline1.get_iterator(timeout=10)) + + self.assertEqual([0, 2, 4], vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_max_failures_pipe_override_strict( + self, output_order: str + ) -> None: + """max_failures at pipe overrides the global threshold.""" + + def fail_odd(x): + if x % 2: + raise ValueError(f"Only evan numbers are allowed. {x}") + return x + + builder = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_odd, output_order=output_order, max_failures=2) + .add_sink(1) + ) + + pipeline = builder.build(num_threads=1, max_failures=-1) + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual([0, 2, 4], vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_max_failures_pipe_override_loose(self, output_order: str) -> None: + """max_failures at pipe overrides the global threshold.""" + + def fail_odd(x): + if x % 2: + raise ValueError(f"Only evan numbers are allowed. {x}") + return x + + builder = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_odd, output_order=output_order, concurrency=3, max_failures=-1) + .add_sink(1) + ) + + pipeline = builder.build(num_threads=1, max_failures=2) + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual([0, 2, 4, 6, 8], vals) + + @parameterized.expand( + [ + ("completion",), + ("input",), + ] + ) + def test_pipeline_max_failures_pipe_override_multiple( + self, output_order: str + ) -> None: + """max_failures at pipe overrides the global threshold.""" + + # Remove odd values + def fail_odd(x): + if x % 2: + raise ValueError(f"Only evan numbers are allowed. {x}") + return x + + # Remove multiplier of 6s + def fail_six(x): + if (x % 6) == 0: + raise ValueError(f"Values divisible by 6 are not allowed. {x}") + return x + + builder = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(fail_odd, output_order=output_order, max_failures=-1) + .pipe(fail_six, output_order=output_order, max_failures=3) + .add_sink(1) + ) + + # fail_odd fails more often, but it is allowed to fail any number of times. + # fail_six fails less often, but at the fourth failure (18), + # it should shutdown the pipeline. + + pipeline = builder.build(num_threads=1, max_failures=2) + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + vals = list(pipeline.get_iterator(timeout=10)) + self.assertEqual([2, 4, 8, 10, 14, 16], vals) + + +class TestPipelinePropagate(unittest.TestCase): + def test_pipeline_propagate_source_failure(self) -> None: + """When source itrator fails, the exception is propagated to the front end""" + + def failure_source(): + raise RuntimeError("Foo") + yield None + + pipeline = ( + PipelineBuilder() + .add_source(failure_source()) + .add_sink() + .build(num_threads=1) + ) + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + pass + + +class TestIterableWithShuffle: + __test__ = False + + def __init__(self, n: int) -> None: + self.vals = list(range(n)) + self.seed = 0 + + def shuffle(self, *, seed: int) -> None: + self.vals = self.vals[1:] + self.vals[:1] + self.seed = seed + + def __iter__(self) -> Iterator[int]: + yield from self.vals + + +class TestRunPipeline(unittest.TestCase): + @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) + def test_run_pipeline_in_subprocess_state(self) -> None: + """The status of the source is maintained and propagated properly in subprocess""" + n = 5 + + # pyre-ignore[6] + src = embed_shuffle(TestIterableWithShuffle(n)) + builder = PipelineBuilder().add_source(src).add_sink() + # pyre-ignore[6] + iterable = run_pipeline_in_subprocess(builder, num_threads=1) + + self.assertEqual([1, 2, 3, 4, 0], list(iterable)) + self.assertEqual([2, 3, 4, 0, 1], list(iterable)) + self.assertEqual([3, 4, 0, 1, 2], list(iterable)) + + # since the src is copied to the subprocess iterating it yields the original state + + self.assertEqual([1, 2, 3, 4, 0], list(src)) + # pyre-ignore[16] + self.assertEqual(0, src.src.seed) + self.assertEqual([2, 3, 4, 0, 1], list(src)) + self.assertEqual(1, src.src.seed) + self.assertEqual([3, 4, 0, 1, 2], list(src)) + self.assertEqual(2, src.src.seed) + + @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) + def test_run_pipeline_in_subprocess_pipeline_id(self) -> None: + """The pipeline construdted in a subprocess inherits the global ID from the main process""" + + # Set to a number that's not zero and something unlikely to happen during the testing + _set_global_id(123456) + ref = _get_global_id() + 1 + + builder = PipelineBuilder().add_source(_ValidatePipelineId(ref)).add_sink() + + # pyre-ignore[6] + iterable = run_pipeline_in_subprocess(builder, num_threads=1) + + for _ in iterable: + pass + + +class TestOverrideStage(unittest.TestCase): + @_ignore_warnings(_UNAWAITED_COROUTINE) + def test_override_stage_id(self) -> None: + """Providing `stage_id` overrides the index of stages.""" + ref = 12345 + + class CheckNameQueue(AsyncQueue): + index = ref + + def __init__(self, name, *, buffer_size: int = 1) -> None: + print(name) + id = re.match(r"\d+:(\d+):.*", str(name)).group(1) + assert id == str(self.index) + CheckNameQueue.index += 1 + super().__init__(name, buffer_size=buffer_size) + + ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x) + .pipe(lambda x: x) + .pipe(lambda x: x) + .add_sink() + .build(num_threads=1, queue_class=CheckNameQueue, stage_id=ref) + ) + + +class TestPipelineFailureStructure(unittest.TestCase): + def _build_failing_pipeline(self): + def failing_range(n): + yield from range(n) + raise ValueError("Iterator failed") + + return ( + PipelineBuilder() + .add_source(failing_range(3)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + + def test_pipeline_failure_is_exception_group(self) -> None: + pipeline = self._build_failing_pipeline() + + with self.assertRaises(PipelineFailure) as ctx: + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pf = ctx.exception + if sys.version_info >= (3, 11): + self.assertIsInstance(pf, ExceptionGroup) + else: + self.assertIsInstance(pf, RuntimeError) + self.assertGreaterEqual(len(pf.exceptions), 1) + exception_types = {type(e) for e in pf.exceptions} + self.assertTrue(exception_types & {ValueError}) + + def test_pipeline_failure_individual_exceptions(self) -> None: + def failing_range(n): + yield from range(n) + raise TypeError("source failed") + + pipeline = ( + PipelineBuilder() + .add_source(failing_range(3)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure) as ctx: + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pf = ctx.exception + if sys.version_info >= (3, 11): + self.assertIsInstance(pf, ExceptionGroup) + else: + self.assertIsInstance(pf, RuntimeError) + self.assertGreaterEqual(len(pf.exceptions), 1) + self.assertTrue( + any(isinstance(e, TypeError) for e in pf.exceptions), + ) + + @unittest.skipIf(sys.version_info < (3, 11), "ExceptionGroup requires Python 3.11+") + def test_pipeline_failure_subgroup(self) -> None: + def failing_range(n): + yield from range(n) + raise ValueError("Iterator failed") + + pipeline = ( + PipelineBuilder() + .add_source(failing_range(3)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure) as ctx: + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pf = ctx.exception + sub = pf.subgroup(ValueError) + self.assertIsNotNone(sub) + self.assertIsInstance(sub, PipelineFailure) + self.assertTrue(all(isinstance(e, ValueError) for e in sub.exceptions)) + + @unittest.skipIf(sys.version_info < (3, 11), "ExceptionGroup requires Python 3.11+") + def test_pipeline_failure_notes_contain_stage_name(self) -> None: + def failing_range(n): + yield from range(n) + raise ValueError("Iterator failed") + + pipeline = ( + PipelineBuilder() + .add_source(failing_range(3)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure) as ctx: + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pf = ctx.exception + for exc in pf.exceptions: + notes = getattr(exc, "__notes__", []) + self.assertTrue( + any(note.startswith("Pipeline stage:") for note in notes), + ) diff --git a/src/spdl/pipeline/tests/pipeline_cleanup_test.py b/src/spdl/pipeline/tests/pipeline_cleanup_test.py new file mode 100644 index 000000000..92ff1adef --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_cleanup_test.py @@ -0,0 +1,203 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import gc +import unittest +import warnings + +from spdl.pipeline import PipelineBuilder + + +class TestPipelineCleanup(unittest.TestCase): + """Test class for Pipeline cleanup functionality.""" + + def test_cleanup_called_on_garbage_collection(self) -> None: + """Test that the pipeline background thread is automatically stopped + on garbage collection without warnings.""" + + # Setup: Create a pipeline and start it + pipeline = ( + PipelineBuilder().add_source(range(100)).add_sink(1000).build(num_threads=1) + ) + + pipeline.start(timeout=3) + + # Verify the pipeline is running + self.assertTrue(pipeline._impl._event_loop.is_started()) + + # Keep a reference to the impl to verify it stopped + impl = pipeline._impl + + # Execute: Delete the pipeline reference without calling stop() + # The facade's finalizer should cleanly stop the pipeline + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + # Delete the pipeline and force garbage collection + del pipeline + gc.collect() + + # Assert: No warning should be issued — the facade handles cleanup + cleanup_warnings = [ + warning + for warning in w + if "Pipeline is running in the background" in str(warning.message) + ] + self.assertEqual(len(cleanup_warnings), 0) + + # Verify the background thread actually stopped + self.assertTrue(impl._event_loop.is_task_completed()) + + def test_cleanup_not_called_when_explicitly_stopped(self) -> None: + """Test that _cleanup_pipeline is not called when Pipeline is explicitly stopped.""" + + # Setup: Create a pipeline and start it + pipeline = ( + PipelineBuilder().add_source(range(100)).add_sink(1000).build(num_threads=1) + ) + + pipeline.start(timeout=3) + + # Execute: Explicitly stop the pipeline + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + pipeline.stop(timeout=3) + + # Delete the pipeline reference and force garbage collection + del pipeline + gc.collect() + + # Assert: Verify that no warning was issued + # Since we explicitly stopped the pipeline, the finalizer should be detached + # and no cleanup warning should be issued + cleanup_warnings = [ + warning + for warning in w + if "Pipeline is running in the background" in str(warning.message) + ] + self.assertEqual(len(cleanup_warnings), 0) + + def test_cleanup_with_auto_stop_context_manager(self) -> None: + """Test that cleanup is not called when using auto_stop context manager.""" + + # Setup: Create a pipeline + pipeline = ( + PipelineBuilder().add_source(range(10)).add_sink(1000).build(num_threads=1) + ) + + # Execute: Use the pipeline with auto_stop context manager + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + with pipeline.auto_stop(timeout=3): + # Get some items from the pipeline + iterator = pipeline.get_iterator(timeout=3) + for _ in range(5): + next(iterator) + + # Delete the pipeline reference and force garbage collection + del pipeline + gc.collect() + + # Assert: Verify that no cleanup warning was issued + # The context manager should have stopped the pipeline properly + cleanup_warnings = [ + warning + for warning in w + if "Pipeline is running in the background" in str(warning.message) + ] + self.assertEqual(len(cleanup_warnings), 0) + + def test_auto_start_and_cleanup_without_explicit_start_stop(self) -> None: + """Test that iterating a pipeline without calling start/stop works: + the background thread is started automatically on first iteration, + and the finalizer cleans it up on garbage collection.""" + + pipeline = ( + PipelineBuilder().add_source(range(10)).add_sink(1000).build(num_threads=1) + ) + + # Pipeline should not be started yet + self.assertFalse(pipeline._impl._event_loop.is_started()) + + # Iterate without explicit start — auto-start should kick in + items = [] + for item in pipeline: + items.append(item) + + # Verify auto-start happened + self.assertTrue(pipeline._impl._event_loop.is_started()) + self.assertEqual(sorted(items), list(range(10))) + + # Keep a reference to the impl to verify cleanup + impl = pipeline._impl + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + del pipeline + gc.collect() + + cleanup_warnings = [ + warning + for warning in w + if "Pipeline is running in the background" in str(warning.message) + ] + self.assertEqual(len(cleanup_warnings), 0) + + # Verify the impl was stopped (stop was called by finalizer) + self.assertTrue(impl._event_loop.is_task_completed()) + + def test_auto_start_and_cleanup_continuous_source(self) -> None: + """Test auto start/stop with a continuous source iterated multiple times. + + With continuous=True the source re-iterates, injecting epoch boundary + sentinels. Each ``for ... in pipeline`` consumes one epoch. The pipeline + should auto-start on the first epoch and remain running across epochs, + then clean up on garbage collection.""" + + pipeline = ( + PipelineBuilder() + .add_source(range(5), continuous=True) + .add_sink(1000) + .build(num_threads=1) + ) + + self.assertFalse(pipeline._impl._event_loop.is_started()) + + num_epochs = 3 + all_epoch_items = [] + for _ in range(num_epochs): + epoch_items = [] + for item in pipeline: + epoch_items.append(item) + all_epoch_items.append(sorted(epoch_items)) + + # Verify auto-start happened and each epoch produced the same items + self.assertTrue(pipeline._impl._event_loop.is_started()) + for epoch_items in all_epoch_items: + self.assertEqual(epoch_items, list(range(5))) + + # Verify cleanup on garbage collection + impl = pipeline._impl + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + del pipeline + gc.collect() + + cleanup_warnings = [ + warning + for warning in w + if "Pipeline is running in the background" in str(warning.message) + ] + self.assertEqual(len(cleanup_warnings), 0) + + self.assertTrue(impl._event_loop.is_task_completed()) diff --git a/src/spdl/pipeline/tests/pipeline_def_test.py b/src/spdl/pipeline/tests/pipeline_def_test.py new file mode 100644 index 000000000..4dc6026a6 --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_def_test.py @@ -0,0 +1,168 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from typing import TypeVar + +from spdl.pipeline import build_pipeline +from spdl.pipeline.defs import ( + Aggregate, + Disaggregate, + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, +) + +T = TypeVar("T") +U = TypeVar("U") + + +# pyre-strict + + +class PipelineDefTest(unittest.TestCase): + def test_source_repr(self) -> None: + """`repr` of SourceConfig should not generate a huge string.""" + + src = SourceConfig(list(range(10000))) + + self.assertGreater(len(repr(src.source)), 10000) + self.assertLess(len(repr(src)), 10000) + + def test_pipe_args_repr(self) -> None: + """`repr` of Pipe should not generate a huge string.""" + + lst = list(range(10000)) + self.assertGreater(len(repr(lst)), 10000) + pipe = Pipe(lst) + self.assertLess(len(repr(pipe._args.op)), 10000) + self.assertLess(len(repr(pipe)), 10000) + + dct = {i: i for i in range(10000)} + self.assertGreater(len(repr(dct)), 10000) + pipe = Pipe(dct) + self.assertLess(len(repr(pipe._args.op)), 10000) + self.assertLess(len(repr(pipe)), 10000) + + def _test_build_pipeline(self, cfg: PipelineConfig[T], expected: list[T]) -> None: + print(cfg) + pipeline = build_pipeline(cfg, num_threads=1) + + with pipeline.auto_stop(): + ite = pipeline.get_iterator(timeout=3) + self.assertEqual(list(ite), expected) + + def test_build_pipeline_simple(self) -> None: + """PipelineConfig and build_pipeline works without pipes.""" + src = range(10) + cfg = PipelineConfig( + src=SourceConfig(src), + pipes=[], + sink=SinkConfig(3), + ) + + self._test_build_pipeline(cfg, list(src)) + + def test_build_pipeline_aggregate(self) -> None: + """Aggregate works""" + cfg = PipelineConfig( + src=SourceConfig(range(8)), + pipes=[ + Aggregate(3, drop_last=False), + ], + sink=SinkConfig(3), + ) + + expected = [[0, 1, 2], [3, 4, 5], [6, 7]] + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_aggregate_drop_last(self) -> None: + """Aggregate works""" + cfg = PipelineConfig( + src=SourceConfig(range(8)), + pipes=[ + Aggregate(3, drop_last=True), + ], + sink=SinkConfig(3), + ) + + expected = [[0, 1, 2], [3, 4, 5]] + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_disaggregate(self) -> None: + """Disaggregate works""" + cfg = PipelineConfig( + src=SourceConfig([[0, 1, 2, 3]]), + pipes=[ + Disaggregate(), + ], + sink=SinkConfig(3), + ) + + expected = [0, 1, 2, 3] + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_pipe_identity(self) -> None: + """Pipe works with identity""" + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(lambda x: x), + ], + sink=SinkConfig(3), + ) + + expected = list(range(5)) + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_pipe_double(self) -> None: + """Pipe works with simple lambda""" + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(lambda x: 2 * x), + ], + sink=SinkConfig(3), + ) + + expected = [2 * i for i in range(5)] + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_pipe_sum(self) -> None: + """Pipe works with aggregated data""" + cfg = PipelineConfig( + src=SourceConfig(range(8)), + pipes=[ + Aggregate(3), + Pipe(sum), + ], + sink=SinkConfig(3), + ) + + expected = [3, 12, 13] + self._test_build_pipeline(cfg, expected) + + def test_build_pipeline_pipe_list(self) -> None: + """Pipe works with list""" + mapping = [i * i for i in range(8)] + cfg = PipelineConfig( + src=SourceConfig(range(8)), + pipes=[Pipe(mapping)], + sink=SinkConfig(3), + ) + self._test_build_pipeline(cfg, mapping) + + def test_build_pipeline_pipe_map(self) -> None: + """Pipe works with map (dict)""" + mapping = {i: i * i for i in range(8)} + cfg = PipelineConfig( + src=SourceConfig(range(8)), + pipes=[Pipe(mapping)], + sink=SinkConfig(3), + ) + + self._test_build_pipeline(cfg, list(mapping.values())) diff --git a/src/spdl/pipeline/tests/pipeline_failure_exceptstar_test.py b/src/spdl/pipeline/tests/pipeline_failure_exceptstar_test.py new file mode 100644 index 000000000..a46baf8f3 --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_failure_exceptstar_test.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Tests for PipelineFailure using ``except*`` syntax (Python 3.11+ only). + +This file is separated from pipeline_builder_test.py because ``except*`` is a +syntactic construct that causes SyntaxError on Python < 3.11 at import time, +which would prevent the entire module from loading. +""" + +import sys +import unittest +from collections.abc import Iterator + +if sys.version_info < (3, 11): + raise unittest.SkipTest("except* syntax requires Python 3.11+") + +from spdl.pipeline import PipelineBuilder + + +def passthrough(x: int) -> int: + return x + + +class TestPipelineFailureExceptStar(unittest.TestCase): + def test_pipeline_failure_except_star(self) -> None: + def failing_range(n: int) -> Iterator[int]: + yield from range(n) + raise ValueError("Iterator failed") + + pipeline = ( + PipelineBuilder() + .add_source(failing_range(3)) + .pipe(passthrough) + .add_sink(1000) + .build(num_threads=1) + ) + + caught = [] + try: + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + except* ValueError as eg: + caught.extend(eg.exceptions) + + self.assertGreaterEqual(len(caught), 1) + self.assertIsInstance(caught[0], ValueError) diff --git a/src/spdl/pipeline/tests/pipeline_node_test.py b/src/spdl/pipeline/tests/pipeline_node_test.py new file mode 100644 index 000000000..eab1bb24b --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_node_test.py @@ -0,0 +1,299 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import unittest + +from spdl.pipeline._components import AsyncQueue +from spdl.pipeline._components._common import StageInfo +from spdl.pipeline._components._node import ( + _cancel_orphaned, + _cancel_recursive, + _FanInNode, + _gather_error, + _Node, + _PathVariantsMergeConfig, + _SourceNode, + _start_tasks, +) +from spdl.pipeline.defs import SinkConfig, SourceConfig + + +class DummyException(Exception): + pass + + +_TTestNode = _SourceNode | _Node | _FanInNode + + +def _node( + name: str, + deps: list[_TTestNode], + exc: Exception | None = None, +) -> _TTestNode: + async def coro() -> None: + if exc: + raise exc + else: + await asyncio.sleep(10) + + info = StageInfo(pipeline_id=0, stage_id="0", stage_name=name) + n: _TTestNode + if not deps: + n = _SourceNode( + info, + SourceConfig(source=[]), + output_queue=AsyncQueue(info), + ) + elif len(deps) > 1: + n = _FanInNode( + info, + _PathVariantsMergeConfig(), + deps, + input_queues=[ + AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name=f"{name}_in_{i}") + ) + for i in range(len(deps)) + ], + output_queue=AsyncQueue(info), + ) + else: + n = _Node( + info, + SinkConfig(buffer_size=1), + deps, + input_queue=AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name=f"{name}_in") + ), + output_queue=AsyncQueue(info), + ) + n._coro = coro() + return n + + +class PipelineNodeTest(unittest.TestCase): + def test_node_chain_start_and_cancel(self) -> None: + # A -> B -> C + + async def run() -> None: + a = _node("A", []) + b = _node("B", [a]) + c = _node("C", [b]) + tasks = _start_tasks(c) + self.assertTrue(all(isinstance(t, asyncio.Task) for t in tasks)) + self.assertEqual(len(tasks), 3) + + _cancel_recursive(c) + await asyncio.wait(tasks) + self.assertTrue(all(t.cancelled() for t in tasks)) + + asyncio.run(run()) + + def test_node_y_shape_upstream(self) -> None: + # A1 B1 + # | | + # A2 B2 + # \ / + # C1 + + async def run() -> None: + a1 = _node("A1", []) + a2 = _node("A2", [a1]) + b1 = _node("B1", []) + b2 = _node("B2", [b1]) + c1 = _node("C1", [a2, b2]) + tasks = _start_tasks(c1) + self.assertEqual(len(tasks), 5) + _cancel_recursive(c1) + await asyncio.wait(tasks) + self.assertTrue(all(t.cancelled() for t in tasks)) + + asyncio.run(run()) + + def test_cancel_error_upstreams_and_gather_error(self) -> None: + # A1 B1 + # | | + # A2 B2 (raises) + # \ / + # C1 + + async def run() -> None: + a1 = _node("A1", []) + a2 = _node("A2", [a1]) + b1 = _node("B1", []) + b2 = _node("B2", [b1], exc=DummyException("fail B2")) + c1 = _node("C1", [a2, b2]) + tasks = _start_tasks(c1) + await asyncio.sleep(0) + + # Let B2 fail + await asyncio.wait([b2.task]) + + _cancel_orphaned(c1) + await asyncio.wait(tasks) + + # Only B1 should be cancelled, since B2 errored + self.assertFalse(a1.task.cancelled()) + self.assertFalse(a2.task.cancelled()) + self.assertTrue(b1.task.cancelled()) + self.assertFalse(b2.task.cancelled()) + self.assertFalse(c1.task.cancelled()) + + errs = _gather_error(c1) + self.assertEqual(len(errs), 1) + name, err = errs[0] + self.assertEqual(name, "0:0:B2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail B2") + + asyncio.run(run()) + + def test_cancel_error_upstreams_and_gather_error_multiple(self) -> None: + # A1 B1 + # | | + # (raises) A2 B2 (raises) + # \ / + # C1 + async def run() -> None: + a1 = _node("A1", []) + a2 = _node("A2", [a1], exc=DummyException("fail A2")) + b1 = _node("B1", []) + b2 = _node("B2", [b1], exc=DummyException("fail B2")) + c1 = _node("C1", [a2, b2]) + + tasks = _start_tasks(c1) + await asyncio.sleep(0) + + # Let A2, B2 fail + await asyncio.wait([a2.task, b2.task]) + + _cancel_orphaned(c1) + await asyncio.wait(tasks) + + self.assertTrue(a1.task.cancelled()) + self.assertFalse(a2.task.cancelled()) + self.assertTrue(b1.task.cancelled()) + self.assertFalse(b2.task.cancelled()) + self.assertFalse(c1.task.cancelled()) + + errs = _gather_error(c1) + self.assertEqual(len(errs), 2) + name, err = errs[0] + self.assertEqual(name, "0:0:A2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail A2") + name, err = errs[1] + self.assertEqual(name, "0:0:B2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail B2") + + asyncio.run(run()) + + def test_cancel_error_upstreams_and_gather_error_complex(self) -> None: + # B1 C1 + # | | + # (raises) B2 C2 (raises) + # \ / + # A1 D1 E1 + # | | | + # (raises) A2 D2 E2 + # \ | / + # F1 + + async def run() -> None: + a1 = _node("A1", []) + a2 = _node("A2", [a1], exc=DummyException("fail A2")) + b1 = _node("B1", []) + b2 = _node("B2", [b1], exc=DummyException("fail B2")) + c1 = _node("C1", []) + c2 = _node("C2", [c1], exc=DummyException("fail C2")) + d1 = _node("D1", [b2, c2]) + d2 = _node("D2", [d1]) + e1 = _node("E1", []) + e2 = _node("E1", [e1]) + f1 = _node("F1", [a2, d2, e2]) + + tasks = _start_tasks(f1) + await asyncio.sleep(0) + + # Let A2, B2, C2 fail + await asyncio.wait([a2.task, b2.task, c2.task]) + + _cancel_orphaned(f1) + + # Let the cancellations propagate + await asyncio.sleep(0) + + self.assertTrue(a1.task.cancelled()) + self.assertFalse(a2.task.cancelled()) + self.assertTrue(b1.task.cancelled()) + self.assertFalse(b2.task.cancelled()) + self.assertTrue(c1.task.cancelled()) + self.assertFalse(c2.task.cancelled()) + self.assertFalse(d1.task.cancelled()) + self.assertFalse(d2.task.cancelled()) + self.assertFalse(e1.task.cancelled()) + self.assertFalse(e2.task.cancelled()) + self.assertFalse(f1.task.cancelled()) + + await asyncio.wait(tasks) + + errs = _gather_error(f1) + self.assertEqual(len(errs), 3) + name, err = errs[0] + self.assertEqual(name, "0:0:A2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail A2") + name, err = errs[1] + self.assertEqual(name, "0:0:B2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail B2") + name, err = errs[2] + self.assertEqual(name, "0:0:C2") + self.assertIsInstance(err, DummyException) + self.assertEqual(err.args[0], "fail C2") + + asyncio.run(run()) + + def test_gather_error_with_cancelled(self) -> None: + # A1 B1 + # | | + # A2 B2 + # \ / + # C1 (raises) + + async def run() -> None: + a1 = _node("A1", []) + a2 = _node("A2", [a1]) + b1 = _node("B1", []) + b2 = _node("B2", [b1]) + c1 = _node("C1", [a2, b2], exc=DummyException("fail C")) + + tasks = _start_tasks(c1) + await asyncio.sleep(0) + + # Let C fail + await asyncio.wait([c1.task]) + + _cancel_orphaned(c1) + await asyncio.wait(tasks) + + self.assertTrue(a1.task.cancelled()) + self.assertTrue(a2.task.cancelled()) + self.assertTrue(b1.task.cancelled()) + self.assertTrue(b2.task.cancelled()) + self.assertFalse(c1.task.cancelled()) + + errs = _gather_error(c1) + self.assertEqual(len(errs), 1) + name, err = errs[0] + self.assertEqual(name, "0:0:C1") + self.assertIsInstance(err, DummyException) + + asyncio.run(run()) diff --git a/src/spdl/pipeline/tests/pipeline_profiling_test.py b/src/spdl/pipeline/tests/pipeline_profiling_test.py new file mode 100644 index 000000000..7e77c2abf --- /dev/null +++ b/src/spdl/pipeline/tests/pipeline_profiling_test.py @@ -0,0 +1,321 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from collections.abc import AsyncIterator, Iterator +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +from spdl.pipeline import ( + config, + profile_pipeline, + ProfileHook, + ProfileResult, +) +from spdl.pipeline._profile import ( + _build_pipeline_config, + _fetch_inputs, +) +from spdl.pipeline.defs import ( + Aggregate, + Disaggregate, + Merge, + Pipe, + PipeConfig, + PipelineConfig, + SinkConfig, + SourceConfig, +) + + +class FetchInputsTest(unittest.TestCase): + """Test _fetch_inputs functionality.""" + + def test_fetch_inputs(self): + """_fetch_inputs collects input items""" + src = SourceConfig(range(10)) + + inputs = _fetch_inputs(src, num_items=3) + self.assertEqual(inputs, list(range(3))) + + def test_fetch_inputs_async(self): + """_fetch_inputs collects input items""" + + async def arange(n: int) -> AsyncIterator[int]: + for i in range(n): + yield i + + src = SourceConfig(arange(10)) + + inputs = _fetch_inputs(src, num_items=3) + self.assertEqual(inputs, list(range(3))) + + +class ProfilePipelineTest(unittest.TestCase): + """Test profile_pipeline functionality.""" + + def setUp(self) -> None: + """Reset all configuration state before each test.""" + config.set_default_profile_hook() + config.set_default_profile_callback() + + def test_profile_pipeline(self): + def foo(i: int) -> int: + return 2 * i + + def bar(items: list[int]) -> list[int]: + return [sum(items)] + + def bazz(i: int) -> int: + return i * i + + N, m = 25, 3 + + plc = PipelineConfig( + src=SourceConfig(range(N)), + pipes=[ + Pipe(foo), + Aggregate(m), + Pipe(bar), + Disaggregate(), + Pipe(bazz), + ], + sink=SinkConfig(3), + ) + + class Intercept_: + def __init__(self) -> None: + self.i = 0 + + def __call__(self, inputs, pipe, concurrency): + num_inputs = N if self.i < 2 else (N + m - 1) // m + self.assertEqual(len(inputs), num_inputs) + self.assertEqual(pipe, plc.pipes[self.i]) + ret = _build_pipeline_config(inputs, pipe, concurrency) + self.assertEqual(len(ret.pipes), 1) + if isinstance(pipe, PipeConfig): + self.assertIs(ret.pipes[0]._args.op, plc.pipes[self.i]._args.op) + self.i += 1 + return ret + + mock = Intercept_() + mock.assertEqual = self.assertEqual + mock.assertIs = self.assertIs + with patch("spdl.pipeline._profile._build_pipeline_config", mock): + profile_pipeline(plc) + + self.assertEqual(mock.i, 5) + + def test_profile_pipeline_callback(self): + """Test that profile_pipeline calls the callback for each pipe stage.""" + + def simple_op(i: int) -> int: + return i + 1 + + cfg = PipelineConfig( + src=SourceConfig(range(10)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + callback_mock = MagicMock() + results = profile_pipeline(cfg, num_inputs=5, callback=callback_mock) + + callback_mock.assert_called_once() + called_args = callback_mock.call_args[0] + self.assertEqual(len(called_args), 1) + called_result = called_args[0] + + self.assertIsInstance(called_result, ProfileResult) + self.assertEqual(called_result.name, "simple_op") + self.assertGreater(len(called_result.stats), 0) + + self.assertEqual(len(results), 1) + self.assertEqual(results[0].name, called_result.name) + self.assertEqual(len(results[0].stats), len(called_result.stats)) + + def test_profile_pipeline_no_callback(self): + """Test that profile_pipeline works correctly when no callback is provided.""" + + def simple_op(i: int) -> int: + return i * 2 + + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + results = profile_pipeline(cfg, num_inputs=3, callback=None) + + self.assertEqual(len(results), 1) + self.assertEqual(results[0].name, "simple_op") + self.assertGreater(len(results[0].stats), 0) + + +class ProfileHookTest(unittest.TestCase): + """Test class for ProfileHook functionality.""" + + def setUp(self) -> None: + """Reset all configuration state before each test.""" + config.set_default_profile_hook() + config.set_default_profile_callback() + + def test_profile_pipeline_custom_hook_methods_called(self): + """Test that custom ProfileHook's stage_profile_hook and + pipeline_profile_hook methods are called. + """ + + def simple_op(i: int) -> int: + return i + 10 + + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + stage_hook_mock = MagicMock() + pipeline_hook_mock = MagicMock() + + class MockProfileHook(ProfileHook): + @contextmanager + def stage_profile_hook( + self, _stage: str, _concurrency: int + ) -> Iterator[None]: + stage_hook_mock() + try: + yield None + finally: + stage_hook_mock() + + @contextmanager + def pipeline_profile_hook(self) -> Iterator[None]: + pipeline_hook_mock() + try: + yield + finally: + pipeline_hook_mock() + + custom_hook = MockProfileHook() + + results = profile_pipeline(cfg, num_inputs=3, hook=custom_hook) + + self.assertGreater(len(results), 0) + self.assertEqual(pipeline_hook_mock.call_count, 2) + self.assertEqual(stage_hook_mock.call_count, 10) + + def test_profile_pipeline_skips_when_local_rank_not_zero(self): + """Test that profiling is skipped if LOCAL_RANK is not '0'.""" + + def simple_op(i: int) -> int: + return i * 3 + + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + with patch("spdl.pipeline._profile._get_local_rank", return_value=1): + results = profile_pipeline(cfg, num_inputs=5) + + self.assertEqual(results, []) + + def test_profile_pipeline_runs_when_local_rank_zero(self): + """Test that profiling runs normally when LOCAL_RANK is '0'.""" + + def simple_op(i: int) -> int: + return i * 2 + + cfg = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(simple_op), + ], + sink=SinkConfig(1), + ) + + with patch("spdl.pipeline._profile._get_local_rank", return_value=0): + results = profile_pipeline(cfg, num_inputs=3) + + self.assertGreater(len(results), 0) + self.assertEqual(results[0].name, "simple_op") + self.assertGreater(len(results[0].stats), 0) + + +class MergeConfigTest(unittest.TestCase): + """Test class for profile_pipeline with Merge configurations.""" + + def setUp(self) -> None: + """Reset all configuration state before each test.""" + config.set_default_profile_hook() + config.set_default_profile_callback() + + def test_profile_pipeline_with_merge_config_and_post_merge_stages(self): + """Test that profile_pipeline profiles all stages including + those in Merge and post-merge stages. + """ + + def double(i: int) -> int: + return i * 2 + + def triple(i: int) -> int: + return i * 3 + + def add_ten(i: int) -> int: + return i + 10 + + def square(i: int) -> int: + return i * i + + plc1 = PipelineConfig( + src=SourceConfig(range(5)), + pipes=[ + Pipe(double, name="double"), + ], + sink=SinkConfig(1), + ) + + plc2 = PipelineConfig( + src=SourceConfig(range(10, 15)), + pipes=[ + Pipe(triple, name="triple"), + ], + sink=SinkConfig(1), + ) + + main_cfg = PipelineConfig( + src=Merge([plc1, plc2]), + pipes=[ + Pipe(add_ten, name="add_ten"), + Pipe(square, name="square"), + ], + sink=SinkConfig(1), + ) + results = profile_pipeline(main_cfg, num_inputs=3) + + self.assertEqual(len(results), 4) + + self.assertEqual(results[0].name, "double") + self.assertEqual(results[1].name, "triple") + self.assertEqual(results[2].name, "add_ten") + self.assertEqual(results[3].name, "square") + + for result in results: + self.assertGreater(len(result.stats), 0) + for stat in result.stats: + self.assertTrue(hasattr(stat, "concurrency")) + self.assertTrue(hasattr(stat, "qps")) + self.assertTrue(hasattr(stat, "occupancy_rate")) diff --git a/src/spdl/pipeline/tests/priority_executor_test.py b/src/spdl/pipeline/tests/priority_executor_test.py new file mode 100644 index 000000000..462568d0a --- /dev/null +++ b/src/spdl/pipeline/tests/priority_executor_test.py @@ -0,0 +1,937 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import functools +import gc +import itertools +import pickle +import sys +import threading +import time +import unittest +import warnings +import weakref +from collections.abc import Callable +from concurrent.futures import Executor, Future, ThreadPoolExecutor +from queue import Empty +from typing import Type, TypeVar + +from spdl.pipeline import ( + PipelineBuilder, + PipelineFailure, + PriorityExecutorEntrypoint, + PriorityProcessPoolExecutor, + PriorityThreadPoolExecutor, +) +from spdl.pipeline._priority_executor import _OWNER_REGISTRY, _PriorityQueueAdapter + +_F = TypeVar("_F", bound=Callable[..., object]) +_C = TypeVar("_C", bound=Type[object]) + + +def _ignore_fork_warning(fn: _F) -> _F: + @functools.wraps(fn) + def wrapper(*args: object, **kwargs: object) -> object: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + # pyre-ignore[7] + return wrapper + + +def _ignore_fork_warning_in_class(cls: _C) -> _C: + for name, member in list(vars(cls).items()): + if name.startswith("test_") and callable(member): + setattr(cls, name, _ignore_fork_warning(member)) + return cls + + +def _raise_value_error() -> None: + raise ValueError("process boom") + + +# ─── _PriorityQueueAdapter unit tests ─── + + +class TestPriorityQueueAdapter(unittest.TestCase): + def test_priority_ordering(self) -> None: + q = _PriorityQueueAdapter() + q.put("low") + q.put("high") + q.put("mid") + # All inserted without priority context → same default → FIFO + self.assertEqual(q.get(), "low") + self.assertEqual(q.get(), "high") + self.assertEqual(q.get(), "mid") + + def test_priority_ordering_with_explicit_priority(self) -> None: + q = _PriorityQueueAdapter() + + q.put("low", priority=(10, 0)) + q.put("high", priority=(0, 0)) + q.put("mid", priority=(5, 0)) + + self.assertEqual(q.get(), "high") + self.assertEqual(q.get(), "mid") + self.assertEqual(q.get(), "low") + + def test_fifo_within_same_priority(self) -> None: + q = _PriorityQueueAdapter() + for i in range(5): + q.put(f"item_{i}", priority=(1, i)) + + for i in range(5): + self.assertEqual(q.get(), f"item_{i}") + + def test_sentinel_none_processed_first(self) -> None: + q = _PriorityQueueAdapter() + q.put("work", priority=(0, 0)) + q.put(None) # shutdown sentinel, no priority + + self.assertIsNone(q.get()) + self.assertEqual(q.get(), "work") + + def test_get_nowait_empty_raises(self) -> None: + q = _PriorityQueueAdapter() + with self.assertRaises(Empty): + q.get_nowait() + + def test_empty_and_qsize(self) -> None: + q = _PriorityQueueAdapter() + self.assertTrue(q.empty()) + self.assertEqual(q.qsize(), 0) + q.put("x") + self.assertFalse(q.empty()) + self.assertEqual(q.qsize(), 1) + + +# ─── PriorityExecutorEntrypoint unit tests ─── + + +class TestPriorityExecutorEntrypoint(unittest.TestCase): + def test_submit_returns_future(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + fut = stage.submit(lambda: 42) + self.assertIsInstance(fut, Future) + self.assertEqual(fut.result(timeout=5), 42) + executor.shutdown() + + def test_is_executor_compatible(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + self.assertIsInstance(stage, Executor) + executor.shutdown() + + def test_shutdown_is_noop(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + stage.shutdown() # should not affect the pool + fut = stage.submit(lambda: 1) + self.assertEqual(fut.result(timeout=5), 1) + executor.shutdown() + + +# ─── PriorityThreadPoolExecutor ordering tests ─── + + +class TestPriorityThreadPoolExecutorOrdering(unittest.TestCase): + def test_downstream_stage_runs_first(self) -> None: + """With 1 worker, tasks are executed in priority order.""" + barrier = threading.Barrier(2) + results: list[str] = [] + + executor = PriorityThreadPoolExecutor(max_workers=1) + upstream = executor.get_executor(priority=0) + downstream = executor.get_executor(priority=2) + + # Block the single worker so we can enqueue both tasks + executor.get_executor().submit(lambda: barrier.wait(timeout=5)) + + # Enqueue upstream first, then downstream + upstream.submit(lambda: results.append("upstream")) + downstream.submit(lambda: results.append("downstream")) + + # Release the worker + barrier.wait(timeout=5) + executor.shutdown(wait=True) + + self.assertEqual(results, ["downstream", "upstream"]) + + def test_fifo_within_stage(self) -> None: + barrier = threading.Barrier(2) + results: list[int] = [] + + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + + executor.get_executor().submit(lambda: barrier.wait(timeout=5)) + + for i in range(5): + stage.submit(lambda i=i: results.append(i)) + + barrier.wait(timeout=5) + executor.shutdown(wait=True) + + self.assertEqual(results, [0, 1, 2, 3, 4]) + + def test_multiple_stages_interleaved(self) -> None: + """3 stages, items interleaved — should sort by priority then FIFO.""" + barrier = threading.Barrier(2) + results: list[tuple[int, int]] = [] + + executor = PriorityThreadPoolExecutor(max_workers=1) + stages = [executor.get_executor(priority=p) for p in [0, 1, 2]] + + executor.get_executor().submit(lambda: barrier.wait(timeout=5)) + + # Submit: stage0, stage1, stage2, stage0, stage1, stage2 + for round_idx in range(2): + for idx, stage in enumerate(stages): + stage.submit(lambda i=idx, ri=round_idx: results.append((i, ri))) + + barrier.wait(timeout=5) + executor.shutdown(wait=True) + + # priority 2 (highest) first, then priority 1, then priority 0 + # Within same priority, FIFO by round + self.assertEqual( + results, + [(2, 0), (2, 1), (1, 0), (1, 1), (0, 0), (0, 1)], + ) + + def test_basic_execution(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=2) + stage = executor.get_executor() + futs = [stage.submit(lambda x=x: x * 2, x) for x in range(10)] + results = {f.result(timeout=5) for f in futs} + self.assertEqual(results, {x * 2 for x in range(10)}) + executor.shutdown() + + def test_exception_propagation(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + + def fail() -> None: + raise ValueError("boom") + + fut = stage.submit(fail) + with self.assertRaises(ValueError): + fut.result(timeout=5) + executor.shutdown() + + def test_submit_after_shutdown_raises(self) -> None: + executor = PriorityThreadPoolExecutor(max_workers=1) + stage = executor.get_executor() + executor.shutdown() + with self.assertRaises(RuntimeError): + executor._submit_with_priority((0, 0), lambda: None, (), {}) + + def test_multiple_workers_all_complete(self) -> None: + """With multiple workers, all tasks must complete.""" + executor = PriorityThreadPoolExecutor(max_workers=4) + stages = [executor.get_executor() for _ in range(3)] + + counter: itertools.count[int] = itertools.count() + results: list[int] = [] + lock: threading.Lock = threading.Lock() + + def work() -> None: + val = next(counter) + with lock: + results.append(val) + + futs = [] + for stage in stages: + for _ in range(10): + futs.append(stage.submit(work)) + + for f in futs: + f.result(timeout=10) + + executor.shutdown() + self.assertEqual(len(results), 30) + + +# ─── PriorityProcessPoolExecutor tests ─── + + +@_ignore_fork_warning_in_class +class TestPriorityProcessPoolExecutor(unittest.TestCase): + def test_basic_execution(self) -> None: + executor = PriorityProcessPoolExecutor(max_workers=2) + stage = executor.get_executor() + futs = [stage.submit(pow, 2, x) for x in range(10)] + results = {f.result(timeout=10) for f in futs} + self.assertEqual(results, {2**x for x in range(10)}) + executor.shutdown() + + def test_exception_propagation(self) -> None: + executor = PriorityProcessPoolExecutor(max_workers=1) + stage = executor.get_executor() + + fut = stage.submit(_raise_value_error) + with self.assertRaises(ValueError): + fut.result(timeout=10) + executor.shutdown() + + +# ─── Pipeline integration: correctness ─── + + +class TestPriorityExecutorPipelineCorrectness(unittest.TestCase): + def test_two_stage_pipeline(self) -> None: + """All items flow through correctly with a shared priority executor.""" + pool = PriorityThreadPoolExecutor(max_workers=4) + s1 = pool.get_executor() + s2 = pool.get_executor() + + pipeline = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(lambda x: x * 2, executor=s1, concurrency=4) + .pipe(lambda x: x + 1, executor=s2, concurrency=4) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) + pool.shutdown() + + def test_three_stage_pipeline(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=4) + s1 = pool.get_executor() + s2 = pool.get_executor() + s3 = pool.get_executor() + + pipeline = ( + PipelineBuilder() + .add_source(range(15)) + .pipe(lambda x: x + 1, executor=s1, concurrency=2) + .pipe(lambda x: x * 3, executor=s2, concurrency=2) + .pipe(lambda x: x - 1, executor=s3, concurrency=2) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted((x + 1) * 3 - 1 for x in range(15))) + pool.shutdown() + + def test_mixed_priority_and_default_executor(self) -> None: + """Some stages use priority executor, others use default.""" + pool = PriorityThreadPoolExecutor(max_workers=2) + s1 = pool.get_executor() + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x * 2, executor=s1, concurrency=2) + .pipe(lambda x: x + 1, concurrency=2) + .add_sink(3) + .build(num_threads=2) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted(x * 2 + 1 for x in range(10))) + pool.shutdown() + + def test_exception_propagates_through_pipeline(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=2) + s1 = pool.get_executor() + s2 = pool.get_executor() + + def fail_on_five(x: int) -> int: + if x == 5: + raise ValueError("boom on 5") + return x + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_on_five, executor=s1, concurrency=1, max_failures=0) + .pipe(lambda x: x, executor=s2, concurrency=1) + .add_sink(3) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pool.shutdown() + + +# ─── Pipeline integration: priority ordering ─── + + +class TestPriorityExecutorPipelineOrdering(unittest.TestCase): + def test_downstream_not_starved(self) -> None: + """With 1 shared worker, downstream should interleave with upstream, + not wait until all upstream completes.""" + execution_log: list[tuple[str, int]] = [] + lock: threading.Lock = threading.Lock() + + pool = PriorityThreadPoolExecutor(max_workers=1) + upstream_exec = pool.get_executor() + downstream_exec = pool.get_executor() + + def upstream_op(item: int) -> int: + with lock: + execution_log.append(("up", item)) + # Sleep so the event loop has time to submit the downstream task + # before the worker picks the next item. + time.sleep(0.03) + return item + + def downstream_op(item: int) -> int: + with lock: + execution_log.append(("down", item)) + return item + + pipeline = ( + PipelineBuilder() + .add_source(range(8)) + .pipe(upstream_op, executor=upstream_exec, concurrency=1) + .pipe(downstream_op, executor=downstream_exec, concurrency=1) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=30)) + + self.assertEqual(len(results), 8) + pool.shutdown() + + # With priority: up/down alternate → max consecutive upstream ≤ 2. + # Without priority (FIFO): upstream dominates → all upstream first. + max_consec_up = 0 + consec = 0 + for stage, _ in execution_log: + if stage == "up": + consec += 1 + max_consec_up = max(max_consec_up, consec) + else: + consec = 0 + + self.assertLessEqual( + max_consec_up, + 2, + f"Upstream ran {max_consec_up} consecutive times — " + f"downstream was starved. Full log: {execution_log}", + ) + + def test_three_stage_priority_order(self) -> None: + """With 3 stages sharing 1 worker, the most downstream stage + with pending work should run first.""" + execution_log: list[tuple[str, int]] = [] + lock: threading.Lock = threading.Lock() + + pool = PriorityThreadPoolExecutor(max_workers=1) + s1_exec = pool.get_executor() + s2_exec = pool.get_executor() + s3_exec = pool.get_executor() + + def make_op(name: str): # pyre-ignore[3] + def op(item: int) -> int: + with lock: + execution_log.append((name, item)) + time.sleep(0.03) + return item + + return op + + pipeline = ( + PipelineBuilder() + .add_source(range(6)) + .pipe(make_op("s1"), executor=s1_exec, concurrency=1) + .pipe(make_op("s2"), executor=s2_exec, concurrency=1) + .pipe(make_op("s3"), executor=s3_exec, concurrency=1) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = list(pipeline.get_iterator(timeout=30)) + + self.assertEqual(len(results), 6) + pool.shutdown() + + # After downstream stages have pending work, upstream should not + # run consecutively. + stages = [stage for stage, _ in execution_log] + for i in range(len(stages) - 1): + if stages[i] == "s1" and stages[i + 1] == "s1": + prior = stages[:i] + self.assertNotIn( + "s2", + prior, + f"Consecutive s1 at positions {i},{i + 1} after s2 " + f"already had work. Log: {execution_log}", + ) + + def test_priority_vs_fifo_comparison(self) -> None: + """Priority executor interleaves downstream at least as well as FIFO.""" + + def run_pipeline( + upstream_exec: Executor, + downstream_exec: Executor, + num_threads: int, + ) -> list[str]: + log: list[str] = [] + lock: threading.Lock = threading.Lock() + + def up_op(item: int) -> int: + with lock: + log.append("up") + time.sleep(0.03) + return item + + def down_op(item: int) -> int: + with lock: + log.append("down") + return item + + pipeline = ( + PipelineBuilder() + .add_source(range(8)) + .pipe(up_op, executor=upstream_exec, concurrency=1) + .pipe(down_op, executor=downstream_exec, concurrency=1) + .add_sink(3) + .build(num_threads=num_threads) + ) + + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=30)) + return log + + # Priority executor + pool = PriorityThreadPoolExecutor(max_workers=1) + priority_log = run_pipeline( + pool.get_executor(), + pool.get_executor(), + num_threads=1, + ) + pool.shutdown() + + # Plain FIFO executor + plain = ThreadPoolExecutor(max_workers=1) + fifo_log = run_pipeline(plain, plain, num_threads=1) + plain.shutdown() + + # With priority, downstream entries should appear at least as early. + halfway = len(priority_log) // 2 + priority_down_first_half = priority_log[:halfway].count("down") + fifo_down_first_half = fifo_log[:halfway].count("down") + + self.assertGreaterEqual( + priority_down_first_half, + fifo_down_first_half, + f"Priority executor should interleave downstream tasks at least " + f"as well as FIFO.\n" + f"Priority log: {priority_log}\n" + f"FIFO log: {fifo_log}", + ) + + +# ─── Drop-in compatibility ─── + + +class TestDropInCompatibility(unittest.TestCase): + def test_stage_executor_with_as_completed(self) -> None: + from concurrent.futures import as_completed + + pool = PriorityThreadPoolExecutor(max_workers=2) + stage = pool.get_executor() + futs = [stage.submit(pow, 2, i) for i in range(5)] + results = set() + for f in as_completed(futs, timeout=5): + results.add(f.result()) + self.assertEqual(results, {1, 2, 4, 8, 16}) + pool.shutdown() + + def test_context_manager(self) -> None: + with PriorityThreadPoolExecutor(max_workers=2) as executor: + stage = executor.get_executor() + self.assertEqual(stage.submit(lambda: 99).result(timeout=5), 99) + + +# ─── Mixed priority + regular ThreadPoolExecutor pipelines ─── + + +class TestMixedExecutorPipeline(unittest.TestCase): + def test_priority_upstream_regular_downstream(self) -> None: + """Priority executor on upstream stages, regular ThreadPoolExecutor + on downstream stage.""" + pool = PriorityThreadPoolExecutor(max_workers=2) + regular = ThreadPoolExecutor(max_workers=2) + + pipeline = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(lambda x: x * 2, executor=pool.get_executor(), concurrency=2) + .pipe(lambda x: x + 1, executor=regular, concurrency=2) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) + pool.shutdown() + regular.shutdown() + + def test_regular_upstream_priority_downstream(self) -> None: + """Regular ThreadPoolExecutor on upstream, priority executor on + downstream.""" + pool = PriorityThreadPoolExecutor(max_workers=2) + regular = ThreadPoolExecutor(max_workers=2) + + pipeline = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(lambda x: x + 10, executor=regular, concurrency=2) + .pipe(lambda x: x * 3, executor=pool.get_executor(), concurrency=2) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted((x + 10) * 3 for x in range(20))) + pool.shutdown() + regular.shutdown() + + def test_three_stage_alternating_executors(self) -> None: + """Three stages: priority, regular, priority — verifying they + compose correctly.""" + pool = PriorityThreadPoolExecutor(max_workers=3) + regular = ThreadPoolExecutor(max_workers=2) + + pipeline = ( + PipelineBuilder() + .add_source(range(15)) + .pipe(lambda x: x + 1, executor=pool.get_executor(), concurrency=2) + .pipe(lambda x: x * 2, executor=regular, concurrency=2) + .pipe(lambda x: x - 1, executor=pool.get_executor(), concurrency=2) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted((x + 1) * 2 - 1 for x in range(15))) + pool.shutdown() + regular.shutdown() + + def test_two_priority_pools_and_regular(self) -> None: + """Two independent priority pools plus a regular pool, all in one + pipeline.""" + pool_a = PriorityThreadPoolExecutor(max_workers=2) + pool_b = PriorityThreadPoolExecutor(max_workers=2) + regular = ThreadPoolExecutor(max_workers=2) + + pipeline = ( + PipelineBuilder() + .add_source(range(12)) + .pipe(lambda x: x + 1, executor=pool_a.get_executor(), concurrency=2) + .pipe(lambda x: x * 2, executor=regular, concurrency=2) + .pipe(lambda x: x - 1, executor=pool_b.get_executor(), concurrency=2) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted((x + 1) * 2 - 1 for x in range(12))) + pool_a.shutdown() + pool_b.shutdown() + regular.shutdown() + + def test_exception_in_regular_stage(self) -> None: + """Exception in the regular executor stage propagates correctly.""" + pool = PriorityThreadPoolExecutor(max_workers=2) + regular = ThreadPoolExecutor(max_workers=2) + + def fail_on_five(x: int) -> int: + if x == 5: + raise ValueError("mixed boom") + return x + + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x, executor=pool.get_executor(), concurrency=2) + .pipe(fail_on_five, executor=regular, concurrency=1, max_failures=0) + .add_sink(3) + .build(num_threads=1) + ) + + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=10)) + + pool.shutdown() + regular.shutdown() + + +# ─── Pickle support ─── + + +class TestPriorityExecutorPickle(unittest.TestCase): + def test_thread_executor_round_trip(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=4) + data = pickle.dumps(pool) + pool.shutdown() + + restored = pickle.loads(data) + stage = restored.get_executor() + fut = stage.submit(pow, 2, 10) + self.assertEqual(fut.result(timeout=5), 1024) + restored.shutdown() + + def test_entrypoint_round_trip(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=2) + ep = pool.get_executor(priority=5) + data = pickle.dumps(ep) + pool.shutdown() + + restored_ep = pickle.loads(data) + fut = restored_ep.submit(pow, 2, 3) + self.assertEqual(fut.result(timeout=5), 8) + restored_ep._owner.shutdown() + + def test_multiple_entrypoints_share_master(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=2) + ep1 = pool.get_executor(priority=1) + ep2 = pool.get_executor(priority=2) + + data1 = pickle.dumps(ep1) + data2 = pickle.dumps(ep2) + + # Remove original pool from registry to simulate new process + pool_id = pool._id + pool.shutdown() + _OWNER_REGISTRY.pop(pool_id, None) + + restored1 = pickle.loads(data1) + restored2 = pickle.loads(data2) + + self.assertIs(restored1._owner, restored2._owner) + + fut1 = restored1.submit(pow, 2, 3) + fut2 = restored2.submit(pow, 3, 2) + self.assertEqual(fut1.result(timeout=5), 8) + self.assertEqual(fut2.result(timeout=5), 9) + restored1._owner.shutdown() + + @_ignore_fork_warning + def test_process_executor_round_trip(self) -> None: + pool = PriorityProcessPoolExecutor(max_workers=2) + data = pickle.dumps(pool) + pool.shutdown() + + restored = pickle.loads(data) + stage = restored.get_executor() + fut = stage.submit(pow, 2, 10) + self.assertEqual(fut.result(timeout=10), 1024) + restored.shutdown() + + def test_entrypoint_is_executor_subclass_after_unpickle(self) -> None: + pool = PriorityThreadPoolExecutor(max_workers=1) + ep = pool.get_executor() + data = pickle.dumps(ep) + pool.shutdown() + + restored = pickle.loads(data) + self.assertIsInstance(restored, Executor) + self.assertIsInstance(restored, PriorityExecutorEntrypoint) + restored._owner.shutdown() + + def test_entrypoint_priority_preserved(self) -> None: + """Priority ordering is preserved across pickle round-trip.""" + pool = PriorityThreadPoolExecutor(max_workers=1) + upstream = pool.get_executor(priority=0) + downstream = pool.get_executor(priority=2) + + data_up = pickle.dumps(upstream) + data_down = pickle.dumps(downstream) + pool.shutdown() + _OWNER_REGISTRY.pop(pool._id, None) + + restored_up = pickle.loads(data_down) + restored_down = pickle.loads(data_up) + + # Submit and wait for results + fut_up = restored_up.submit(pow, 2, 3) + fut_down = restored_down.submit(pow, 3, 2) + self.assertEqual(fut_up.result(timeout=5), 8) + self.assertEqual(fut_down.result(timeout=5), 9) + + # Verify both share the same owner + self.assertIs(restored_up._owner, restored_down._owner) + # Verify priority values were preserved + self.assertEqual(restored_up._priority, -2) + self.assertEqual(restored_down._priority, 0) + restored_up._owner.shutdown() + + +# ─── Garbage collection ─── + + +class TestPriorityExecutorGarbageCollection(unittest.TestCase): + def test_owner_gc_after_all_references_dropped(self) -> None: + """Owner is garbage-collected once all entrypoints and user refs are gone.""" + pool = PriorityThreadPoolExecutor(max_workers=1) + owner_id = pool._id + weak = weakref.ref(pool) + stage = pool.get_executor() + + self.assertIn(owner_id, _OWNER_REGISTRY) + self.assertIsNotNone(weak()) + + # Drop the user reference — entrypoint still holds a strong ref + del pool + gc.collect() + self.assertIsNotNone(weak()) + + # Drop the entrypoint — last strong ref gone + del stage + gc.collect() + self.assertIsNone(weak()) + self.assertNotIn(owner_id, _OWNER_REGISTRY) + + def test_owner_gc_multiple_entrypoints(self) -> None: + """Owner survives until ALL entrypoints are dropped.""" + pool = PriorityThreadPoolExecutor(max_workers=1) + weak = weakref.ref(pool) + s1 = pool.get_executor() + s2 = pool.get_executor() + del pool + gc.collect() + + self.assertIsNotNone(weak()) + del s1 + gc.collect() + self.assertIsNotNone(weak()) + del s2 + gc.collect() + self.assertIsNone(weak()) + + def test_owner_gc_after_shutdown(self) -> None: + """Shutdown + dropping all refs allows GC.""" + pool = PriorityThreadPoolExecutor(max_workers=1) + weak = weakref.ref(pool) + stage = pool.get_executor() + fut = stage.submit(lambda: 42) + self.assertEqual(fut.result(timeout=5), 42) + + pool.shutdown() + del pool + gc.collect() + # Entrypoint still holds a strong ref + self.assertIsNotNone(weak()) + + del stage + gc.collect() + self.assertIsNone(weak()) + + +# ─── PriorityInterpreterPoolExecutor tests (Python 3.14+ only) ─── + +_has_interpreter_pool: bool = sys.version_info >= (3, 14) + + +@unittest.skipUnless( + _has_interpreter_pool, "InterpreterPoolExecutor requires Python 3.14+" +) +class TestPriorityInterpreterPoolExecutor(unittest.TestCase): + def test_basic_execution(self) -> None: + from spdl.pipeline import PriorityInterpreterPoolExecutor + + executor = PriorityInterpreterPoolExecutor(max_workers=2) + stage = executor.get_executor() + futs = [stage.submit(pow, 2, x) for x in range(10)] + results = {f.result(timeout=10) for f in futs} + self.assertEqual(results, {2**x for x in range(10)}) + executor.shutdown() + + def test_exception_propagation(self) -> None: + from spdl.pipeline import PriorityInterpreterPoolExecutor + + executor = PriorityInterpreterPoolExecutor(max_workers=1) + stage = executor.get_executor() + + def fail() -> None: + raise ValueError("interpreter boom") + + fut = stage.submit(fail) + with self.assertRaises(ValueError): + fut.result(timeout=10) + executor.shutdown() + + def test_submit_after_shutdown_raises(self) -> None: + from spdl.pipeline import PriorityInterpreterPoolExecutor + + executor = PriorityInterpreterPoolExecutor(max_workers=1) + stage = executor.get_executor() + executor.shutdown() + with self.assertRaises(RuntimeError): + executor._submit_with_priority( # pyre-ignore[16] + (0, 0), lambda: None, (), {} + ) + + def test_is_executor_compatible(self) -> None: + from spdl.pipeline import PriorityInterpreterPoolExecutor + + executor = PriorityInterpreterPoolExecutor(max_workers=1) + stage = executor.get_executor() + self.assertIsInstance(stage, Executor) + executor.shutdown() + + def test_pipeline_two_stage(self) -> None: + from spdl.pipeline import PriorityInterpreterPoolExecutor + + pool = PriorityInterpreterPoolExecutor(max_workers=4) + s1 = pool.get_executor() + s2 = pool.get_executor() + + pipeline = ( + PipelineBuilder() + .add_source(range(20)) + .pipe(lambda x: x * 2, executor=s1, concurrency=4) + .pipe(lambda x: x + 1, executor=s2, concurrency=4) + .add_sink(3) + .build(num_threads=1) + ) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator(timeout=10)) + + self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) + pool.shutdown() diff --git a/src/spdl/pipeline/tests/source_locator_test.py b/src/spdl/pipeline/tests/source_locator_test.py new file mode 100644 index 000000000..fb4e1ba2f --- /dev/null +++ b/src/spdl/pipeline/tests/source_locator_test.py @@ -0,0 +1,200 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +""" +Unit tests for the source_locator module. +""" + +import asyncio +import functools +import inspect +import unittest +from typing import Any + +from spdl.pipeline._common._source_locator import locate_source + + +def regular_function(x: int, y: int) -> int: + """A regular function for testing.""" + return x + y + + +def generator_function(n: int) -> Any: + """A generator function for testing.""" + for i in range(n): + yield i + + +async def async_function(x: int) -> int: + """An async function for testing.""" + await asyncio.sleep(0) + return x * 2 + + +async def async_generator_function(n: int) -> Any: + """An async generator function for testing.""" + for i in range(n): + await asyncio.sleep(0) + yield i + + +class SimpleCallable: + """A simple callable class for testing.""" + + def __call__(self, x: int) -> int: + return x * 3 + + +def _ln(target: object) -> int: + return inspect.getsourcelines(target)[1] # pyre-ignore[6] + + +class TestSourceLocator(unittest.TestCase): + """Test cases for the locate_source function.""" + + def test_regular_function(self) -> None: + """Test locating source for a regular function.""" + loc = locate_source(regular_function) + + self.assertEqual(loc.name, f"{__name__}.regular_function") + self.assertEqual(loc.file_path, __file__) + self.assertEqual(loc.line_number, _ln(regular_function)) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_generator_function(self) -> None: + """Test locating source for a generator function.""" + loc = locate_source(generator_function) + + self.assertEqual(loc.name, f"{__name__}.generator_function") + self.assertEqual(loc.file_path, __file__) + self.assertEqual(loc.line_number, _ln(generator_function)) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_async_function(self) -> None: + """Test locating source for an async function.""" + loc = locate_source(async_function) + + self.assertEqual(loc.name, f"{__name__}.async_function") + self.assertEqual(loc.file_path, __file__) + self.assertEqual(loc.line_number, _ln(async_function)) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_async_generator_function(self) -> None: + """Test locating source for an async generator function.""" + loc = locate_source(async_generator_function) + + self.assertEqual(loc.name, f"{__name__}.async_generator_function") + self.assertEqual(loc.file_path, __file__) + self.assertEqual(loc.line_number, _ln(async_generator_function)) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_callable_class_object(self) -> None: + """Test locating source for a callable class object.""" + obj = SimpleCallable() + loc = locate_source(obj) + + self.assertEqual(loc.name, f"{__name__}.SimpleCallable") + self.assertEqual(loc.file_path, __file__) + self.assertEqual(loc.line_number, _ln(SimpleCallable)) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_builtin_function(self) -> None: + """Test locating source for a built-in function.""" + loc = locate_source(len) + + self.assertEqual(loc.name, "builtins.len") + self.assertIsNone(loc.file_path) + self.assertIsNone(loc.line_number) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {}) + + def test_partial_with_positional_args(self) -> None: + """ + Test locating source for a function wrapped with functools.partial + (positional args). + """ + partial_func = functools.partial(regular_function, 5) + loc = locate_source(partial_func) + + self.assertEqual(loc.name, f"{__name__}.regular_function") + self.assertEqual(loc.file_path, __file__) + self.assertIsNotNone(loc.line_number) + self.assertGreater(loc.line_number, 0) + self.assertEqual(loc.partial_args, (5,)) + self.assertEqual(loc.partial_kwargs, {}) + + def test_partial_with_keyword_args(self) -> None: + """ + Test locating source for a function wrapped with functools.partial + (keyword args). + """ + partial_func = functools.partial(regular_function, y=10) + loc = locate_source(partial_func) + + self.assertEqual(loc.name, f"{__name__}.regular_function") + self.assertEqual(loc.file_path, __file__) + self.assertIsNotNone(loc.line_number) + self.assertGreater(loc.line_number, 0) + self.assertEqual(loc.partial_args, ()) + self.assertEqual(loc.partial_kwargs, {"y": 10}) + + def test_callable_object_wrapped_with_partial(self) -> None: + """ + Test locating source for a callable object wrapped with + functools.partial. + """ + obj = SimpleCallable() + partial_obj = functools.partial(obj, 7) + loc = locate_source(partial_obj) + + self.assertEqual(loc.name, f"{__name__}.SimpleCallable") + self.assertEqual(loc.file_path, __file__) + self.assertIsNotNone(loc.line_number) + self.assertGreater(loc.line_number, 0) + self.assertEqual(loc.partial_args, (7,)) + self.assertEqual(loc.partial_kwargs, {}) + + def test_nested_partial(self) -> None: + """Test locating source for nested functools.partial wrapping.""" + partial_func1 = functools.partial(regular_function, 3) + partial_func2 = functools.partial(partial_func1, y=8) + loc = locate_source(partial_func2) + + self.assertEqual(loc.name, f"{__name__}.regular_function") + self.assertEqual(loc.file_path, __file__) + self.assertIsNotNone(loc.line_number) + self.assertGreater(loc.line_number, 0) + self.assertEqual(loc.partial_args, (3,)) + self.assertEqual(loc.partial_kwargs, {"y": 8}) + + def test_nested_partial_positional_args_order(self) -> None: + """Test that nested partial positional args are in correct order.""" + # partial(partial(f, 1), 2) should produce args (1, 2) + partial_func1 = functools.partial(regular_function, 1) + partial_func2 = functools.partial(partial_func1, 2) + loc = locate_source(partial_func2) + + self.assertEqual(loc.partial_args, (1, 2)) + # Verify the actual behavior matches + self.assertEqual(partial_func2(), regular_function(1, 2)) + + def test_nested_partial_keyword_override(self) -> None: + """Test that outer partial keywords override inner ones.""" + # partial(partial(f, x=1), x=2) should use x=2 + partial_func1 = functools.partial(regular_function, x=1, y=5) + partial_func2 = functools.partial(partial_func1, x=2) + loc = locate_source(partial_func2) + + self.assertEqual(loc.partial_kwargs, {"x": 2, "y": 5}) + # Verify the actual behavior matches + self.assertEqual(partial_func2(), regular_function(x=2, y=5)) diff --git a/src/spdl/pipeline/tests/subinterpreter_test.py b/src/spdl/pipeline/tests/subinterpreter_test.py new file mode 100644 index 000000000..1f07a93a5 --- /dev/null +++ b/src/spdl/pipeline/tests/subinterpreter_test.py @@ -0,0 +1,278 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import sys +import unittest +from collections.abc import Callable, Iterable, Iterator, Sequence +from typing import Generic, TypeVar + +from parameterized import parameterized +from spdl.pipeline import PipelineBuilder, run_pipeline_in_subinterpreter +from spdl.pipeline._components import _get_global_id, _set_global_id +from spdl.pipeline._iter_utils._subinterpreter import ( + iterate_in_subinterpreter as _iterate_in_subinterpreter, +) + +T = TypeVar("T") + + +def iterate_in_subinterpreter( + fn: Callable[[], Iterable[T]], + *, + buffer_size: int = 3, + initializer: Callable[[], None] | Sequence[Callable[[], None]] | None = None, + timeout: float = 5, +) -> Iterable[T]: + """Set timeout for unittest""" + return _iterate_in_subinterpreter( + fn, buffer_size=buffer_size, initializer=initializer, timeout=timeout + ) + + +class _Wrap(Generic[T]): + """Helper class to wrap an iterable as a callable. + + This class wraps an iterable object and makes it callable, which is useful + for testing iterate_in_subinterpreter. It can optionally execute a pre-flight + function before returning the iterable, allowing for assertions to verify + that initializers have run correctly in the subinterpreter. + + Args: + obj: The iterable object to wrap. + pre: Optional callable to execute before returning the iterable. + Typically used for assertions in tests. + """ + + def __init__(self, obj: Iterable[T], pre: Callable[[], None] | None = None) -> None: + self.obj = obj + self.pre = pre + + def __call__(self) -> Iterable[T]: + if self.pre is not None: + self.pre() + return self.obj + + +_FLAGS: list[int] = [] + + +def _init_flag0() -> None: + _FLAGS.append(0) + + +def _init_flag1() -> None: + _FLAGS.append(1) + + +def _check_flag0and1() -> None: + ref = [0, 1] + assert _FLAGS == ref, f"{_FLAGS=} != {ref=}" + + +if sys.version_info >= (3, 14): + + class TestIterateInSubinterpreter(unittest.TestCase): + """Test cases for iterate_in_subinterpreter function.""" + + @parameterized.expand( + [ + ("basic_iteration", list(range(5))), + ("string_iteration", ["hello", "world", "test"]), + ("empty_iterator", []), + ], + ) + def test_iteration(self, name: str, ref: list[object]) -> None: # noqa: ARG002 + """Test iteration with various input types.""" + iterable = iterate_in_subinterpreter(_Wrap(ref)) + result = list(iterable) + self.assertEqual(result, ref) + result2 = list(iterable) + self.assertEqual(result2, ref) + + def test_buffer_size(self) -> None: + """Test with custom buffer size.""" + ref = list(range(10)) + result = list(iterate_in_subinterpreter(_Wrap(ref), buffer_size=5)) + self.assertEqual(result, ref) + + def test_with_initializers(self) -> None: + """Test with multiple initializer functions.""" + ref = list(range(10)) + result = list( + iterate_in_subinterpreter( + _Wrap(ref, pre=_check_flag0and1), + initializer=[_init_flag0, _init_flag1], + timeout=5.0, + ) + ) + self.assertEqual(result, ref) + # The flag should not be set in the main interpreter + self.assertEqual(_FLAGS, []) + + def test_partial_iteration(self) -> None: + """Test partial iteration by breaking early.""" + iterable = iterate_in_subinterpreter(_Wrap(range(10))) + result = [] + for i, item in enumerate(iterable): + result.append(item) + if i >= 4: + break + self.assertEqual(result, [0, 1, 2, 3, 4]) + + +# Module-level functions and classes (required for pickling/subinterpreter compatibility) +def _double(x: int) -> int: + """Helper function to double a value.""" + return x * 2 + + +def _only_even(x: int) -> int | None: + """Helper function to filter only even numbers.""" + return x if x % 2 == 0 else None + + +class _StatefulSource: + """Stateful source that tracks iteration calls.""" + + def __init__(self, n: int) -> None: + self.n = n + self.calls = 0 + + def __iter__(self) -> Iterator[int]: + start = self.calls * self.n + self.calls += 1 + yield from range(start, start + self.n) + + +class _validate_pipeline_id: + """Helper class to validate that the pipeline ID is as expected.""" + + def __init__(self, val: int) -> None: + self.val = val + + def __iter__(self) -> Iterator[int]: + if (v := _get_global_id()) != self.val: + raise AssertionError(f"_node._PIPELINE_ID={v} != {self.val=}") + yield 0 + + +if sys.version_info >= (3, 14): + + class TestRunPipelineInSubinterpreter(unittest.TestCase): + """Test cases for run_pipeline_in_subinterpreter function.""" + + def test_basic_pipeline(self) -> None: + """Test basic pipeline execution in subinterpreter.""" + condig = PipelineBuilder().add_source(range(5)).add_sink().get_config() + iterable = run_pipeline_in_subinterpreter(condig, num_threads=1, timeout=5) + result = list(iterable) + self.assertEqual(result, [0, 1, 2, 3, 4]) + + def test_pipeline_with_pipe(self) -> None: + """Test pipeline with pipe operation in subinterpreter.""" + config = ( + PipelineBuilder() + .add_source(range(5)) + .pipe(_double) + .add_sink() + .get_config() + ) + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + result = list(iterable) + self.assertEqual(result, [0, 2, 4, 6, 8]) + + def test_pipeline_multiple_iterations(self) -> None: + """Test that the pipeline can be iterated multiple times.""" + config = ( + PipelineBuilder() + .add_source(_StatefulSource(3)) + .add_sink(buffer_size=10) + .get_config() + ) + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + + # First iteration + result1 = list(iterable) + self.assertEqual(result1, [0, 1, 2]) + + # Second iteration should start from where we left off + result2 = list(iterable) + self.assertEqual(result2, [3, 4, 5]) + + def test_pipeline_with_aggregate(self) -> None: + """Test pipeline with aggregation in subinterpreter.""" + config = ( + PipelineBuilder() + .add_source(range(10)) + .aggregate(3) + .add_sink(buffer_size=10) + .get_config() + ) + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + result = list(iterable) + self.assertEqual(result, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) + + def test_pipeline_empty_source(self) -> None: + """Test pipeline with empty source in subinterpreter.""" + config = PipelineBuilder().add_source([]).add_sink().get_config() + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + result = list(iterable) + self.assertEqual(result, []) + + def test_pipeline_with_filter(self) -> None: + """Test pipeline with filter operation (returning None to skip items).""" + config = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(_only_even) + .add_sink() + .get_config() + ) + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + result = list(iterable) + self.assertEqual(result, [0, 2, 4, 6, 8]) + + def test_pipeline_with_buffer_size(self) -> None: + """Test run_pipeline_in_subinterpreter with custom buffer_size.""" + config = PipelineBuilder().add_source(range(10)).add_sink().get_config() + iterable = run_pipeline_in_subinterpreter( + config, num_threads=1, buffer_size=5, timeout=5 + ) + result = list(iterable) + self.assertEqual(result, list(range(10))) + + def test_run_pipeline_in_subinterpreter_pipeline_id(self) -> None: + """Test pipeline inherits global ID in subinterpreter.""" + # Set to a number that's not zero and something unlikely to + # happen during testing + _set_global_id(123456) + ref = _get_global_id() + 1 + + config = ( + PipelineBuilder() + .add_source(_validate_pipeline_id(ref)) + .add_sink() + .get_config() + ) + + iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) + + for _ in iterable: + pass + +else: + + class TestRunPipelineInSubinterpreter(unittest.TestCase): + """Placeholder tests for Python < 3.14.""" + + def test_requires_python_3_14(self) -> None: + """Test that run_pipeline_in_subinterpreter requires Python 3.14+.""" + config = PipelineBuilder().add_source([1, 2, 3]).add_sink().get_config() + with self.assertRaises(RuntimeError) as cm: + run_pipeline_in_subinterpreter(config, num_threads=1) + self.assertIn("Python 3.14", str(cm.exception)) diff --git a/src/spdl/pipeline/tests/subprocess_break_reiterate_test.py b/src/spdl/pipeline/tests/subprocess_break_reiterate_test.py new file mode 100644 index 000000000..6bf18e1c1 --- /dev/null +++ b/src/spdl/pipeline/tests/subprocess_break_reiterate_test.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +"""Regression test for D101554675: breaking out of a subprocess iterable +must not kill the worker, so subsequent iterations still work.""" + +import functools +import unittest +import warnings +from collections.abc import Iterator +from functools import partial + +from spdl.pipeline import iterate_in_subprocess + + +def _ignore_fork_warning(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + return wrapper + + +def _ignore_fork_warning_in_class(cls): + for name, member in list(vars(cls).items()): + if name.startswith("test_") and callable(member): + setattr(cls, name, _ignore_fork_warning(member)) + return cls + + +class SourceIterable: + def __init__(self, n: int) -> None: + self.n = n + + def __iter__(self) -> Iterator[int]: + yield from range(self.n) + + +@_ignore_fork_warning_in_class +class TestSubprocessBreakAndReiterate(unittest.TestCase): + def test_break_then_reiterate(self) -> None: + """Breaking out of a subprocess iterable must not prevent re-iteration. + + This is a regression test for the BaseException widening in D101554675. + When a consumer `break`s out of `for ... in iterable`, Python sends + GeneratorExit into the generator. Prior to D101554675, only + (Exception, KeyboardInterrupt) triggered _shutdown(). After D101554675, + BaseException (which includes GeneratorExit) triggers _shutdown(), + making subsequent iter() calls raise RuntimeError. + """ + src = iterate_in_subprocess(partial(SourceIterable, 10), timeout=10) + + # First iteration: consume only 3 items, then break + count = 0 + for _item in src: + count += 1 + if count >= 3: + break + + # Second iteration: must succeed (worker should still be alive) + result = list(src) + self.assertEqual(result, list(range(10))) + + def test_break_then_reiterate_multiple_times(self) -> None: + """Multiple break-then-reiterate cycles must all succeed.""" + src = iterate_in_subprocess(partial(SourceIterable, 5), timeout=10) + + for cycle in range(3): + # Break after 2 items + count = 0 + for _item in src: + count += 1 + if count >= 2: + break + + # Full iteration must still work + result = list(src) + self.assertEqual(result, list(range(5)), f"cycle {cycle}") + + def test_partial_iteration_via_zip(self) -> None: + """Partial iteration via zip() (which breaks implicitly) must not kill worker.""" + src = iterate_in_subprocess(partial(SourceIterable, 100), timeout=10) + + # zip stops when the shorter iterable is exhausted, causing + # GeneratorExit on the longer one + partial_result = list(zip(range(3), src)) + self.assertEqual(partial_result, [(0, 0), (1, 1), (2, 2)]) + + # Subsequent full iteration must work + result = list(src) + self.assertEqual(result, list(range(100))) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/spdl/pipeline/tests/subprocess_test.py b/src/spdl/pipeline/tests/subprocess_test.py new file mode 100644 index 000000000..36149bb59 --- /dev/null +++ b/src/spdl/pipeline/tests/subprocess_test.py @@ -0,0 +1,506 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +import functools +import multiprocessing as mp +import os.path +import random +import tempfile +import threading +import time +import unittest +import warnings +from collections.abc import Iterable, Iterator +from functools import partial + +from spdl.pipeline import iterate_in_subprocess as _iterate_in_subprocess +from spdl.pipeline._iter_utils._common import _Cmd, _execute_iterable, _Status + + +def _ignore_fork_warning(fn): + """Suppress the multi-threaded fork() DeprecationWarning emitted by + multiprocessing.popen_fork when starting subprocesses while pipeline + worker threads are alive. The warning is intentional in these tests. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=( + r"This process \(pid=\d+\) is multi-threaded, use of " + r"fork\(\) may lead to deadlocks in the child" + ), + category=DeprecationWarning, + ) + return fn(*args, **kwargs) + + return wrapper + + +def _ignore_fork_warning_in_class(cls): + """Apply ``_ignore_fork_warning`` to every ``test_*`` method on a class.""" + for name, member in list(vars(cls).items()): + if name.startswith("test_") and callable(member): + setattr(cls, name, _ignore_fork_warning(member)) + return cls + + +def iterate_in_subprocess(fn, *, timeout=10, **kwargs): + return _iterate_in_subprocess(fn, timeout=timeout, **kwargs) + + +def iter_range(n: int) -> Iterable[int]: + yield from range(n) + + +def initializer(path: str, val: str) -> None: + with open(path, "w") as f: + f.write(val) + + +@_ignore_fork_warning_in_class +class TestIterateInSubprocess(unittest.TestCase): + def test_iterate_in_subprocess(self) -> None: + """iterate_in_subprocess iterates""" + N = 10 + + src = iterate_in_subprocess(fn=partial(iter_range, n=N)) + self.assertEqual(list(src), list(range(N))) + + def test_iterate_in_subprocess_initializer(self) -> None: + """iterate_in_subprocess initializer is called before iteration starts""" + + N = 10 + val = str(random.random()) + with tempfile.TemporaryDirectory() as dir: + path = os.path.join(dir, "foo.txt") + + self.assertFalse(os.path.exists(path)) + src = iterate_in_subprocess( + fn=partial(iter_range, n=N), + initializer=partial(initializer, path=path, val=val), + buffer_size=1, + ) + self.assertTrue(os.path.exists(path)) + + ite = iter(src) + self.assertEqual(next(ite), 0) + + with open(path, "r") as f: + self.assertEqual(f.read(), val) + + for i in range(1, N): + self.assertEqual(next(ite), i) + + with self.assertRaises(StopIteration): + next(ite) + + def test_iterate_in_subprocess_multiple_initializer(self) -> None: + """iterate_in_subprocess accepts multiple iterators""" + N = 10 + val1 = str(random.random()) + val2 = str(random.random()) + with tempfile.TemporaryDirectory() as dir: + path1 = os.path.join(dir, "foo.txt") + path2 = os.path.join(dir, "bar.txt") + + self.assertFalse(os.path.exists(path1)) + self.assertFalse(os.path.exists(path2)) + src = iterate_in_subprocess( + fn=partial(iter_range, n=N), + initializer=[ + partial(initializer, path=path1, val=val1), + partial(initializer, path=path2, val=val2), + ], + buffer_size=1, + ) + self.assertTrue(os.path.exists(path1)) + self.assertTrue(os.path.exists(path2)) + + ite = iter(src) + self.assertEqual(next(ite), 0) + + with open(path1, "r") as f: + self.assertEqual(f.read(), val1) + + with open(path2, "r") as f: + self.assertEqual(f.read(), val2) + + for i in range(1, N): + self.assertEqual(next(ite), i) + + with self.assertRaises(StopIteration): + next(ite) + + +def iter_range_and_store_with_sync(n: int, sync_queue: mp.Queue) -> Iterable[int]: + """Generator that synchronizes with main process via queue.""" + yield 0 + for i in range(n): + yield i + # Signal main process that we've yielded this value + sync_queue.put(i) + + +@_ignore_fork_warning_in_class +class TestIterateInSubprocessBufferSize(unittest.TestCase): + def test_iterate_in_subprocess_buffer_size_1(self) -> None: + """buffer_size=1 makes iterate_in_subprocess works sort-of interactively""" + + N = 10 + + # Use queue for synchronization between processes + sync_queue = mp.Queue() + + src = iterate_in_subprocess( + fn=partial(iter_range_and_store_with_sync, n=N, sync_queue=sync_queue), + daemon=True, + buffer_size=1, + ) + ite = iter(src) + self.assertEqual(next(ite), 0) + + for i in range(N): + # Wait for subprocess to signal it has yielded value i + # Use timeout to avoid hanging if subprocess fails + subprocess_value = sync_queue.get(timeout=5) + self.assertEqual( + subprocess_value, i, f"Expected {i}, got {subprocess_value}" + ) + + # With buffer_size=1, the queue should be empty after fetching + self.assertTrue( + sync_queue.empty(), f"Queue should be empty after fetching item {i}" + ) + + # Now fetch the value from the iterator + self.assertEqual(next(ite), i) + + with self.assertRaises(StopIteration): + next(ite) + + def test_iterate_in_subprocess_buffer_size_64(self) -> None: + """big buffer_size makes iterate_in_subprocess processes data in one go""" + + N = 10 + + # Use queue for synchronization between processes + sync_queue = mp.Queue() + + src = iterate_in_subprocess( + fn=partial(iter_range_and_store_with_sync, n=N, sync_queue=sync_queue), + daemon=True, + buffer_size=64, + ) + ite = iter(src) + self.assertEqual(next(ite), 0) + + # With buffer_size=64, subprocess should process all data without waiting + # Wait for subprocess to signal all values have been processed + for expected_i in range(N): + # Use timeout to avoid hanging + subprocess_value = sync_queue.get(timeout=5) + self.assertEqual( + subprocess_value, + expected_i, + f"Expected {expected_i}, got {subprocess_value}", + ) + + # Now all data should be available in the buffer, fetch them + for i in range(N): + self.assertEqual(next(ite), i) + + with self.assertRaises(StopIteration): + next(ite) + + +class SourceIterable: + def __init__(self, n: int) -> None: + self.n = n + + def __iter__(self) -> Iterator[int]: + yield from range(self.n) + + +def noop() -> None: + pass + + +class TestExecuteIterable(unittest.TestCase): + def test_execute_iterable_initializer_failure(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + def src_fn() -> Iterable[int]: + return SourceIterable(10) + + def fail() -> None: + raise ValueError("Failed!") + + _execute_iterable(msg_queue, data_queue, src_fn, [fail]) + + self.assertTrue(msg_queue.empty()) + + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.INITIALIZATION_FAILED) + self.assertIn("Failed!", result.message) + self.assertTrue(data_queue.empty()) + + def test_execute_iterable_iterator_initialize_failure(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + def src_fn() -> Iterator[int]: + raise ValueError("Failed!") + return SourceIterable(10) + + _execute_iterable(msg_queue, data_queue, src_fn, [noop]) + + self.assertTrue(msg_queue.empty()) + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.INITIALIZATION_FAILED) + self.assertIn("Failed!", result.message) + self.assertTrue(data_queue.empty()) + + def test_execute_iterable_quite_immediately(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + msg_queue.put(_Cmd.ABORT) + time.sleep(1) + + def src_fn() -> Iterable[int]: + return SourceIterable(10) + + _execute_iterable(msg_queue, data_queue, src_fn, [noop]) + time.sleep(1) + + self.assertTrue(msg_queue.empty()) + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) + self.assertTrue(data_queue.empty()) + + def test_execute_iterable_generator_fail(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + class SourceIterableFails(SourceIterable): + def __iter__(self) -> Iterator[int]: + raise ValueError("Failed!") + yield from range(self.n) + + def src_fn() -> Iterable[int]: + return SourceIterableFails(10) + + msg_queue.put(_Cmd.START_ITERATION) + _execute_iterable(msg_queue, data_queue, src_fn, [noop]) + + self.assertTrue(msg_queue.empty()) + + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.ITERATION_STARTED) + + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.ITERATOR_FAILED) + self.assertIn("Failed!", result.message) + self.assertTrue(data_queue.empty()) + + def test_execute_iterable_generator_fail_after_n(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + class SourceIterableFails(SourceIterable): + def __iter__(self) -> Iterator[int]: + for v in range(self.n): + yield v + if v == 2: + raise ValueError("Failed!") + + def src_fn() -> Iterable[int]: + return SourceIterableFails(10) + + msg_queue.put(_Cmd.START_ITERATION) + _execute_iterable(msg_queue, data_queue, src_fn, [noop]) + + self.assertTrue(msg_queue.empty()) + + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.ITERATION_STARTED) + for i in range(3): + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.ITERATOR_SUCCESS) + self.assertEqual(result.message, i) + + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.ITERATOR_FAILED) + self.assertIn("Failed!", result.message) + self.assertTrue(data_queue.empty()) + + def test_execute_iterator_generator_success(self) -> None: + msg_queue, data_queue = mp.Queue(), mp.Queue() + + def src_fn() -> Iterable[int]: + return SourceIterable(3) + + msg_queue.put(_Cmd.START_ITERATION) + + # Add abort with delay, so that _execute_iterable can exit after + # the iteration + def done(): + time.sleep(3) + msg_queue.put(_Cmd.ABORT) + + t = threading.Thread(target=done) + t.start() + _execute_iterable(msg_queue, data_queue, src_fn, [noop]) + t.join() + + self.assertTrue(msg_queue.empty()) + + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) + ack = data_queue.get(timeout=1) + self.assertEqual(ack.status, _Status.ITERATION_STARTED) + for i in range(3): + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.ITERATOR_SUCCESS) + self.assertEqual(result.message, i) + + result = data_queue.get(timeout=1) + self.assertEqual(result.status, _Status.ITERATION_FINISHED) + + +def _src1() -> Iterable[int]: + return SourceIterable(10) + + +def _init1() -> None: + raise ValueError("Failed!") + + +def _src2() -> Iterator[int]: + if True: + raise ValueError("Failed!") + return SourceIterable(10) + + +def _src3() -> Iterable[int]: + class SourceIterableFails(SourceIterable): + def __iter__(self) -> Iterator[int]: + raise ValueError("Failed!") + yield from range(self.n) + + return SourceIterableFails(10) + + +def _src4() -> Iterable[int]: + class SourceIterableFails(SourceIterable): + def __iter__(self) -> Iterator[int]: + for v in range(self.n): + yield v + if v == 2: + raise ValueError("Failed!") + + return SourceIterableFails(10) + + +def _src5(N) -> Iterable[int]: + return SourceIterable(N) + + +class SleepSourceIterable(SourceIterable): + def __iter__(self): + time.sleep(10) + yield 0 + + +def _src6() -> Iterable[int]: + return SleepSourceIterable(3) + + +def _fail_initializer(): + raise RuntimeError("Failed!") + + +_VERY_BAD_REFERENCE = None + + +@_ignore_fork_warning_in_class +class TestIterateInSubprocessFailures(unittest.TestCase): + def test_iterate_in_subprocess_initializer_failure(self) -> None: + with self.assertRaisesRegex(RuntimeError, r"Initializer failed"): + iterate_in_subprocess(_src1, buffer_size=1, timeout=3, initializer=_init1) + + def test_iterate_in_subprocess_iterator_initialize_failure(self) -> None: + with self.assertRaisesRegex(RuntimeError, r"Failed to create the iterable"): + iterate_in_subprocess(_src2, buffer_size=1, timeout=3) + + def test_iterate_in_subprocess_generator_fail(self) -> None: + ite = iter(iterate_in_subprocess(_src3, buffer_size=1, timeout=3)) + + with self.assertRaisesRegex(RuntimeError, r"Failed to fetch the next item"): + next(ite) + + def test_iterate_in_subprocess_fail_after_n(self) -> None: + ite = iter(iterate_in_subprocess(_src4, buffer_size=1, timeout=3)) + self.assertEqual(next(ite), 0) + self.assertEqual(next(ite), 1) + self.assertEqual(next(ite), 2) + + with self.assertRaisesRegex(RuntimeError, r"Failed to fetch the next item"): + next(ite) + + def test_iterate_in_subprocess_success(self) -> None: + N = 3 + + hyp = list(iterate_in_subprocess(partial(_src5, N), buffer_size=-1, timeout=3)) + self.assertEqual(hyp, list(range(N))) + + def test_iterate_in_subprocess_timeout(self) -> None: + iterable = iterate_in_subprocess(_src6, buffer_size=-1, timeout=3) + iterator = iter(iterable) + with self.assertRaisesRegex( + RuntimeError, r"The worker subprocess did not produce any data for" + ): + next(iterator) + + def test_iterate_in_subprocess_initializer_fail(self) -> None: + """The initialization failure is propagated to the main process""" + + with self.assertRaisesRegex(RuntimeError, r"Initializer failed"): + iterate_in_subprocess(SourceIterable, initializer=_fail_initializer) + + def test_iterate_in_subprocess_iterable_creation_fail(self) -> None: + """The initialization failure is propagated to the main process""" + + with self.assertRaisesRegex(RuntimeError, r"Failed to create the iterable"): + iterate_in_subprocess(SourceIterable) + + def test_iterate_in_subprocess_success_simple_iterable(self) -> None: + iterator = iterate_in_subprocess(partial(SourceIterable, 3)) + + self.assertEqual(list(iterator), [0, 1, 2]) + self.assertEqual(list(iterator), [0, 1, 2]) + self.assertEqual(list(iterator), [0, 1, 2]) + + def test_iterate_in_subprocess_fail_not_stuck(self) -> None: + """An exception does not make Python stack. + + If a (non-daemon) subprocess is launched without a context manager + that ensures its clean exit, raising an exception while the reference + to the process object is held causes the Python interpreter to get + stuck at the exit. + + To avoid this, we register atexit function, which push the ABORT + command to the command queue, which will be received by the subprocess + if the subprocess is not shut down. This test ensures that behavior. + """ + + global _VERY_BAD_REFERENCE + _VERY_BAD_REFERENCE = iterate_in_subprocess(partial(SourceIterable, 3)) diff --git a/tests/__init__.py b/tests/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/autoresearch/app_test.py b/tests/autoresearch/app_test.py deleted file mode 100644 index 08b8c2c37..000000000 --- a/tests/autoresearch/app_test.py +++ /dev/null @@ -1,283 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import importlib -import sys -import tempfile -import unittest -from collections.abc import Callable -from pathlib import Path -from typing import get_origin - -from spdl.autoresearch._app._engine import _parse_engine_args -from spdl.autoresearch._app._spec import ( - _read_workflow_factory, - _record_workflow_factory, - _resolve_workflow, -) -from spdl.autoresearch._app._supervisor import ( - _build_engine_command, - _parse_supervisor_args, -) - -__all__: list[str] = [] - - -def _identity_factory(argv: list[str], workdir: Path | None) -> object: - """A trivial factory used as a resolution target by the tests below.""" - return (argv, workdir) - - -class _ResolveWorkflowTest(unittest.TestCase): - def test_module_factory_form(self) -> None: - """_resolve_workflow imports module.path:factory_name and returns the callable.""" - factory = _resolve_workflow(f"{__name__}:_identity_factory") - self.assertIs(factory, _identity_factory) - - def test_empty_specifier_raises(self) -> None: - """An empty string is rejected with ValueError, not silently importing.""" - with self.assertRaises(ValueError): - _resolve_workflow("") - - def test_malformed_specifier_raises(self) -> None: - """A specifier with a colon but missing one half is rejected.""" - for bad in (":factory", "module.path:", ":"): - with self.subTest(bad=bad): - with self.assertRaises(ValueError): - _resolve_workflow(bad) - - def test_unknown_module_raises_import_error(self) -> None: - """A non-existent module surfaces ModuleNotFoundError to the caller.""" - with self.assertRaises(ModuleNotFoundError): - _resolve_workflow("definitely.not.a.real.module:create") - - def test_missing_attribute_raises(self) -> None: - """An existing module with a missing attribute raises AttributeError.""" - with self.assertRaises(AttributeError): - _resolve_workflow(f"{__name__}:does_not_exist") - - def test_short_name_lookup_misses_cleanly(self) -> None: - """Short-name lookup raises LookupError when no entry point matches.""" - with self.assertRaises(LookupError): - _resolve_workflow("not_registered_workflow_xyz") - - -class _WorkflowFactoryRecordTest(unittest.TestCase): - def test_round_trip(self) -> None: - """_record_workflow_factory followed by _read_workflow_factory returns the spec.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - _record_workflow_factory(workdir, "pkg.mod:factory") - self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") - - def test_read_returns_none_when_missing(self) -> None: - """_read_workflow_factory on a fresh workdir returns None instead of raising.""" - with tempfile.TemporaryDirectory() as tmp: - self.assertIsNone(_read_workflow_factory(Path(tmp))) - - def test_read_rejects_malformed_record(self) -> None: - """Reading a malformed record file raises ValueError.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - (workdir / "workflow_factory.json").write_text("[]\n") - with self.assertRaises(ValueError): - _read_workflow_factory(workdir) - - def test_record_creates_workdir(self) -> None: - """_record_workflow_factory creates the workdir if it does not yet exist.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) / "nested" - _record_workflow_factory(workdir, "pkg.mod:factory") - self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") - - -class _ArgvSplitTest(unittest.TestCase): - def test_supervisor_splits_at_double_dash(self) -> None: - """Tokens after '--' are forwarded as the workflow tail, not parsed by the framework.""" - ns, tail = _parse_supervisor_args( - [ - "/tmp/workdir", - "--workflow", - "pkg.mod:factory", - "--max-concurrency", - "5", - "--", - "--pipeline-script", - "x.py", - ] - ) - self.assertEqual(ns.workdir, "/tmp/workdir") - self.assertEqual(ns.workflow, "pkg.mod:factory") - self.assertEqual(ns.max_concurrency, 5) - self.assertEqual(tail, ["--pipeline-script", "x.py"]) - - def test_supervisor_workdir_optional(self) -> None: - """The supervisor accepts no workdir during initial config gathering.""" - ns, tail = _parse_supervisor_args(["--workflow", "pkg.mod:factory"]) - self.assertIsNone(ns.workdir) - self.assertEqual(tail, []) - - def test_engine_requires_workflow_and_workdir(self) -> None: - """The engine refuses to start without --workflow and --workdir.""" - with self.assertRaises(SystemExit): - _parse_engine_args(["--workdir", "/tmp/wd"]) - with self.assertRaises(SystemExit): - _parse_engine_args(["--workflow", "pkg.mod:factory"]) - - def test_engine_passes_tail_through(self) -> None: - """The engine surfaces the workflow tail unchanged to the caller.""" - ns, tail = _parse_engine_args( - [ - "--workflow", - "pkg.mod:factory", - "--workdir", - "/tmp/wd", - "--", - "--build-command", - "make", - ] - ) - self.assertEqual(ns.workflow, "pkg.mod:factory") - self.assertEqual(ns.workdir, "/tmp/wd") - self.assertEqual(tail, ["--build-command", "make"]) - - def test_engine_max_concurrency_defaults_to_none(self) -> None: - """Omitting --max-concurrency leaves ns.max_concurrency as None. - - The engine treats this as "use the workflow-supplied default" so - the WorkflowSpec.max_concurrency value is honored when the user - does not override it on the CLI. - """ - ns, _ = _parse_engine_args( - ["--workflow", "pkg.mod:factory", "--workdir", "/tmp/wd"] - ) - self.assertIsNone(ns.max_concurrency) - - def test_engine_max_concurrency_accepts_explicit_value(self) -> None: - """An explicit --max-concurrency value is preserved as an int.""" - ns, _ = _parse_engine_args( - [ - "--workflow", - "pkg.mod:factory", - "--workdir", - "/tmp/wd", - "--max-concurrency", - "7", - ] - ) - self.assertEqual(ns.max_concurrency, 7) - - -class _EngineCommandTest(unittest.TestCase): - def test_default_uses_spdl_autoresearch_engine(self) -> None: - """Without an override, the engine prefix is 'spdl autoresearch engine'.""" - cmd = _build_engine_command( - engine_command_override=None, - workflow_spec="pkg.mod:factory", - workdir=Path("/tmp/wd"), - framework_flags=["--max-concurrency", "3"], - workflow_argv_tail=["--build-command", "make"], - ) - self.assertEqual( - cmd, - [ - "spdl", - "autoresearch", - "engine", - "--workflow", - "pkg.mod:factory", - "--workdir", - "/tmp/wd", - "--max-concurrency", - "3", - "--", - "--build-command", - "make", - ], - ) - - def test_override_replaces_prefix(self) -> None: - """An --engine-command override replaces the default argv[0] prefix.""" - cmd = _build_engine_command( - engine_command_override="buck run //x:engine --", - workflow_spec="pkg.mod:factory", - workdir=Path("/tmp/wd"), - framework_flags=[], - workflow_argv_tail=[], - ) - self.assertEqual( - cmd, - [ - "buck", - "run", - "//x:engine", - "--", - "--workflow", - "pkg.mod:factory", - "--workdir", - "/tmp/wd", - ], - ) - - def test_no_tail_omits_double_dash(self) -> None: - """An empty workflow tail does not append a stray '--'.""" - cmd = _build_engine_command( - engine_command_override=None, - workflow_spec="pkg.mod:factory", - workdir=Path("/tmp/wd"), - framework_flags=[], - workflow_argv_tail=[], - ) - self.assertNotIn("--", cmd[2:]) - - -class _CoreWorkflowExportTest(unittest.TestCase): - def test_workflow_spec_is_protocol(self) -> None: - """``WorkflowSpec`` re-exported from core is a ``Protocol`` subclass. - - ``Protocol`` subclasses are flagged with ``_is_protocol = True`` by - the typing machinery; this guards against accidentally weakening - ``WorkflowSpec`` to a regular class (which would silently change - the runtime semantics for workflow authors). - """ - from spdl.autoresearch.core import WorkflowSpec - - self.assertTrue(getattr(WorkflowSpec, "_is_protocol", False)) - - def test_workflow_factory_is_callable_alias(self) -> None: - """``WorkflowFactory`` re-exported from core is a ``Callable`` alias.""" - from spdl.autoresearch.core import WorkflowFactory - - self.assertIs(get_origin(WorkflowFactory), Callable) - - -class _MainImportTest(unittest.TestCase): - def test_main_import_does_not_load_app(self) -> None: - """Importing spdl.autoresearch.__main__ as a module is a no-op. - - The framework dispatcher (under spdl.autoresearch._app) must - NOT be transitively loaded by ``import - spdl.autoresearch.__main__``. _app is reachable only when - __main__.py runs as a script (i.e. via ``python -m - spdl.autoresearch``), at which point ``__name__ == - "__main__"`` and the lazy import inside the guard fires. - """ - removed = {} - for mod_name in [ - name - for name in list(sys.modules) - if name == "spdl.autoresearch.__main__" - or name.startswith("spdl.autoresearch._app") - ]: - removed[mod_name] = sys.modules.pop(mod_name) - self.addCleanup(sys.modules.update, removed) - - importlib.import_module("spdl.autoresearch.__main__") - - self.assertNotIn("spdl.autoresearch._app", sys.modules) - self.assertNotIn("spdl.autoresearch._app._main", sys.modules) diff --git a/tests/autoresearch/app_test.py b/tests/autoresearch/app_test.py new file mode 120000 index 000000000..6f7d2f100 --- /dev/null +++ b/tests/autoresearch/app_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/app_test.py \ No newline at end of file diff --git a/tests/autoresearch/factory_test.py b/tests/autoresearch/factory_test.py deleted file mode 100644 index 95d37a511..000000000 --- a/tests/autoresearch/factory_test.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import json -import tempfile -import unittest -from pathlib import Path - -from spdl.autoresearch.pipeline_optimization import create_workflow -from spdl.autoresearch.pipeline_optimization._ops._analysis_ops import ( - MASTER_TABLE_HEADERS, -) -from spdl.autoresearch.pipeline_optimization._ops._policy import write_state - -__all__: list[str] = [] - - -def _full_argv() -> list[str]: - return [ - "--pipeline-script", - "/tmp/pipeline.py", - "--source-dir", - "/tmp/src", - "--build-command", - "make image", - "--base-launch-command", - "torchx run --image $IMAGE", - "--notes", - "smoke", - "--max-iterations", - "5", - "--patience", - "2", - "--job-timeout", - "300", - ] - - -class _CreateWorkflowTest(unittest.TestCase): - def test_returns_workflow_spec(self) -> None: - """create_workflow returns an object exposing the WorkflowSpec surface.""" - spec = create_workflow(_full_argv(), None) - for attr in ( - "engine_argv_tail", - "description", - "supervisor_known_config", - "supervisor_missing_config", - "setup", - "build_workflow", - ): - self.assertTrue(callable(getattr(spec, attr)), attr) - self.assertEqual(spec.max_concurrency, 3) - - def test_engine_argv_tail_round_trips_supplied_flags(self) -> None: - """Every value passed in survives in engine_argv_tail in flag/value order.""" - spec = create_workflow(_full_argv(), None) - tail = spec.engine_argv_tail() - self.assertEqual(tail[tail.index("--pipeline-script") + 1], "/tmp/pipeline.py") - self.assertEqual(tail[tail.index("--source-dir") + 1], "/tmp/src") - self.assertEqual(tail[tail.index("--build-command") + 1], "make image") - self.assertEqual( - tail[tail.index("--base-launch-command") + 1], - "torchx run --image $IMAGE", - ) - self.assertEqual(tail[tail.index("--max-iterations") + 1], "5") - self.assertEqual(tail[tail.index("--patience") + 1], "2") - self.assertEqual(tail[tail.index("--max-concurrency") + 1], "3") - self.assertEqual(tail[tail.index("--job-timeout") + 1], "300") - self.assertEqual(tail[tail.index("--platform") + 1], "auto") - - def test_engine_argv_tail_omits_unset_options(self) -> None: - """Unset optional flags are not emitted at all.""" - spec = create_workflow([], None) - tail = spec.engine_argv_tail() - self.assertNotIn("--pipeline-script", tail) - self.assertNotIn("--build-command", tail) - self.assertNotIn("--base-launch-command", tail) - self.assertNotIn("--source-dir", tail) - self.assertNotIn("--notes", tail) - - def test_engine_argv_tail_emits_boolean_flags(self) -> None: - """Boolean flags appear by themselves with no value when set.""" - spec = create_workflow( - [ - "--skip-instrument", - "--dangerously-skip-permissions", - ], - None, - ) - tail = spec.engine_argv_tail() - self.assertIn("--skip-instrument", tail) - self.assertIn("--dangerously-skip-permissions", tail) - - def test_supervisor_missing_config_lists_required_fields(self) -> None: - """A bare invocation reports all four required fields as missing.""" - spec = create_workflow([], None) - missing = spec.supervisor_missing_config() - self.assertIn("pipeline script", missing) - self.assertIn("source directory", missing) - self.assertIn("build command", missing) - self.assertIn("launch command template", missing) - - def test_supervisor_missing_config_empty_when_all_supplied(self) -> None: - """A fully-configured invocation reports no missing fields.""" - spec = create_workflow(_full_argv(), None) - self.assertEqual(spec.supervisor_missing_config(), []) - - def test_supervisor_known_config_reflects_argv(self) -> None: - """supervisor_known_config exposes the parsed values for the supervisor prompt.""" - spec = create_workflow(_full_argv(), None) - known = spec.supervisor_known_config() - self.assertEqual(known["pipeline_script"], "/tmp/pipeline.py") - self.assertEqual(known["build_command"], "make image") - self.assertEqual(known["local_execution_mode"], "full") - - def test_max_concurrency_reflects_argv(self) -> None: - """A non-default --max-concurrency is visible on spec.max_concurrency.""" - spec = create_workflow([*_full_argv(), "--max-concurrency", "7"], None) - self.assertEqual(spec.max_concurrency, 7) - - def test_description_contains_supervisor_and_platform_content(self) -> None: - """description() joins the supervisor and platform prompt directories.""" - spec = create_workflow(_full_argv(), None) - description = spec.description() - self.assertIsNotNone(description) - assert description is not None # for type checker - self.assertIn("---", description) - self.assertIn("Automated SPDL Pipeline Optimization", description) - - -def _write_minimal_workdir(workdir: Path) -> None: - """Create the minimum files PipelineOptimizationWorkflow.summarize reads.""" - workdir.mkdir(parents=True, exist_ok=True) - (workdir / "config.json").write_text( - json.dumps( - { - "schema_version": 1, - "pipeline_script": "", - "source_dir": "", - "scm": "", - "build_command": "", - "base_launch_command": "", - "stopping_criteria": {"max_iterations": 1, "patience": 1}, - "max_concurrency": 1, - "job_timeout_s": 60, - "poll_interval": 0, - "platform": "auto", - "agent": "claude", - "local_execution_mode": "full", - } - ) - ) - write_state( - workdir, - { - "iteration": 0, - "status": "looping", - "baseline_job": None, - "current_best": None, - "best_metric": None, - "plateau_count": 0, - "best_practices_tried": [], - "history": [], - }, - ) - (workdir / "master_table.tsv").write_text("\t".join(MASTER_TABLE_HEADERS) + "\n") - - -class _SummarizeTest(unittest.TestCase): - def test_summarize_returns_markdown_for_empty_workdir(self) -> None: - """summarize handles a freshly initialised workdir without raising.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - _write_minimal_workdir(workdir) - spec = create_workflow([], workdir) - workflow = spec.build_workflow(workdir) - - output = workflow.summarize(workdir) - - self.assertIn("# Autoresearch summary", output) - self.assertIn(str(workdir), output) - self.assertIn("## Failures", output) - - def test_summarize_includes_master_table_and_live_summary(self) -> None: - """summarize renders master-table rows and summary.md content.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - _write_minimal_workdir(workdir) - with open(workdir / "master_table.tsv", "a") as f: - f.write("run_001\tbaseline\t0.5\n") - (workdir / "summary.md").write_text("Best SM util improved to 85%") - spec = create_workflow([], workdir) - workflow = spec.build_workflow(workdir) - - output = workflow.summarize(workdir) - - self.assertIn("## Master table", output) - self.assertIn("run_001", output) - self.assertIn("## Live summary", output) - self.assertIn("Best SM util improved to 85%", output) diff --git a/tests/autoresearch/factory_test.py b/tests/autoresearch/factory_test.py new file mode 120000 index 000000000..b40dba8c9 --- /dev/null +++ b/tests/autoresearch/factory_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/factory_test.py \ No newline at end of file diff --git a/tests/autoresearch/orchestrator_test.py b/tests/autoresearch/orchestrator_test.py deleted file mode 100644 index b72e6ac5f..000000000 --- a/tests/autoresearch/orchestrator_test.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import asyncio -import unittest -from pathlib import Path - -from spdl.autoresearch.core import Orchestrator, TaskResult, TaskSpec - -__all__: list[str] = [] - - -class _FakeAdapter: - def __init__(self, specs: list[TaskSpec]) -> None: - self.specs = specs - self.started: list[str] = [] - self.checkpoints: list[tuple[list[str], list[str], str]] = [] - self.children: dict[str, list[TaskSpec]] = {} - self.block = False - - def load(self) -> list[TaskSpec]: - return self.specs - - def checkpoint( - self, - queued: list[TaskSpec], - running: list[TaskSpec], - status: str, - ) -> None: - self.checkpoints.append( - ([spec.id for spec in queued], [spec.id for spec in running], status) - ) - - async def make_coro(self, spec: TaskSpec) -> TaskResult: - self.started.append(spec.id) - if self.block: - await asyncio.sleep(60) - return TaskResult(children=self.children.get(spec.id, [])) - - async def on_result(self, spec: TaskSpec, result: TaskResult) -> list[TaskSpec]: - return result.children - - def summarize(self, workdir: Path) -> str: - return f"_FakeAdapter summary at {workdir}" - - -class OrchestratorTest(unittest.IsolatedAsyncioTestCase): - async def test_priority_order_lowest_first(self) -> None: - """Specs are executed in ascending priority order (lowest value first).""" - adapter = _FakeAdapter( - [ - TaskSpec(id="slow", priority=10), - TaskSpec(id="first", priority=-1), - TaskSpec(id="middle", priority=5), - ] - ) - - await Orchestrator(workflow=adapter, max_concurrency=1).run() - - self.assertEqual(["first", "middle", "slow"], adapter.started) - self.assertEqual(([], [], "stopped"), adapter.checkpoints[-1]) - - async def test_completion_enqueues_children(self) -> None: - """Child specs returned by a completed item are enqueued and executed.""" - adapter = _FakeAdapter([TaskSpec(id="root", priority=0)]) - adapter.children["root"] = [ - TaskSpec(id="child_a", priority=1), - TaskSpec(id="child_b", priority=2), - ] - - await Orchestrator(workflow=adapter, max_concurrency=1).run() - - self.assertEqual(["root", "child_a", "child_b"], adapter.started) - self.assertEqual(([], [], "stopped"), adapter.checkpoints[-1]) - - async def test_cancelled_error_persists_interrupted_state(self) -> None: - """Cancellation checkpoints running specs with 'interrupted' status.""" - adapter = _FakeAdapter([TaskSpec(id="running", priority=0)]) - adapter.block = True - task = asyncio.create_task( - Orchestrator(workflow=adapter, max_concurrency=1).run() - ) - - while not adapter.checkpoints: - await asyncio.sleep(0) - task.cancel() - await task - - self.assertEqual(([], ["running"], "interrupted"), adapter.checkpoints[-1]) - - async def test_checkpoint_resume_golden_lifecycle(self) -> None: - """An interrupted engine can resume from checkpointed state and complete.""" - first = _FakeAdapter( - [ - TaskSpec(id="running", priority=0), - TaskSpec(id="queued", priority=1), - ] - ) - first.block = True - task = asyncio.create_task( - Orchestrator(workflow=first, max_concurrency=1).run() - ) - - while not first.checkpoints: - await asyncio.sleep(0) - task.cancel() - await task - - queued_ids, running_ids, status = first.checkpoints[-1] - self.assertEqual( - (["queued"], ["running"], "interrupted"), first.checkpoints[-1] - ) - - resumed_specs = [ - TaskSpec(id=spec_id, priority=0 if spec_id == "running" else 1) - for spec_id in running_ids + queued_ids - ] - resumed = _FakeAdapter(resumed_specs) - await Orchestrator(workflow=resumed, max_concurrency=1).run() - - self.assertEqual("interrupted", status) - self.assertEqual(["running", "queued"], resumed.started) - self.assertEqual(([], [], "stopped"), resumed.checkpoints[-1]) diff --git a/tests/autoresearch/orchestrator_test.py b/tests/autoresearch/orchestrator_test.py new file mode 120000 index 000000000..4c5f72ce6 --- /dev/null +++ b/tests/autoresearch/orchestrator_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/orchestrator_test.py \ No newline at end of file diff --git a/tests/autoresearch/persistence_test.py b/tests/autoresearch/persistence_test.py deleted file mode 100644 index 57a0b72d1..000000000 --- a/tests/autoresearch/persistence_test.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import json -import tempfile -import unittest -import unittest.mock -from pathlib import Path - -from spdl.autoresearch.core import ( - load_or_init, - read_engine_state, - TaskSpec, - write_engine_state, -) - -__all__: list[str] = [] - - -class _PersistenceTest(unittest.TestCase): - def test_round_trip_preserves_spec_fields(self) -> None: - """write_engine_state followed by read_engine_state preserves all TaskSpec fields.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - queued = [ - TaskSpec( - id="exp_001", - priority=-1.5, - kind="experiment", - payload={"node": {"node_id": "exp_001"}, "extra": [1, 2, 3]}, - ), - TaskSpec(id="exp_002", priority=0.0, kind="default", payload={}), - ] - running = [TaskSpec(id="exp_003", priority=2.0, kind="experiment")] - - write_engine_state( - workdir, queued=queued, running=running, status="running" - ) - result = read_engine_state(workdir) - - self.assertIsNotNone(result) - assert result is not None # for type checker - got_queued, got_running, status = result - self.assertEqual(status, "running") - self.assertEqual([spec.id for spec in got_queued], ["exp_001", "exp_002"]) - self.assertEqual(got_queued[0].priority, -1.5) - self.assertEqual(got_queued[0].kind, "experiment") - self.assertEqual( - got_queued[0].payload, - {"node": {"node_id": "exp_001"}, "extra": [1, 2, 3]}, - ) - self.assertEqual([spec.id for spec in got_running], ["exp_003"]) - - def test_read_returns_none_when_missing(self) -> None: - """read_engine_state on a fresh workdir returns None instead of raising.""" - with tempfile.TemporaryDirectory() as tmp: - self.assertIsNone(read_engine_state(Path(tmp))) - - def test_status_round_trips(self) -> None: - """All three orchestrator status values survive persistence.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - for status in ("running", "stopped", "interrupted"): - write_engine_state(workdir, queued=[], running=[], status=status) - result = read_engine_state(workdir) - assert result is not None - self.assertEqual(result[2], status) - - def test_write_creates_workdir(self) -> None: - """write_engine_state creates the workdir if it does not yet exist.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) / "nested" / "fresh" - self.assertFalse(workdir.exists()) - write_engine_state(workdir, queued=[], running=[], status="running") - self.assertTrue((workdir / "engine_state.json").exists()) - - def test_load_or_init_uses_factory_on_fresh_run(self) -> None: - """load_or_init calls the factory exactly once when no checkpoint exists.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - calls: list[int] = [] - - def factory() -> list[TaskSpec]: - calls.append(1) - return [TaskSpec(id="seed", priority=0.0)] - - specs = load_or_init(workdir, factory) - - self.assertEqual(len(calls), 1) - self.assertEqual([spec.id for spec in specs], ["seed"]) - - def test_load_or_init_resumes_from_checkpoint(self) -> None: - """load_or_init returns queued+running and skips the factory on resume.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - queued = [TaskSpec(id="q1", priority=-1.0)] - running = [TaskSpec(id="r1", priority=0.0)] - write_engine_state( - workdir, queued=queued, running=running, status="interrupted" - ) - - def factory() -> list[TaskSpec]: - self.fail("factory must not be invoked when checkpoint exists") - - specs = load_or_init(workdir, factory) - - self.assertEqual([spec.id for spec in specs], ["q1", "r1"]) - - def test_read_rejects_malformed_json_object(self) -> None: - """A non-object JSON file raises ValueError instead of silently misparsing.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - (workdir / "engine_state.json").write_text("[]\n") - with self.assertRaises(ValueError): - read_engine_state(workdir) - - def test_read_rejects_non_list_field(self) -> None: - """A non-list queued/running field raises ValueError.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - (workdir / "engine_state.json").write_text( - json.dumps({"status": "running", "queued": "oops", "running": []}) - ) - with self.assertRaises(ValueError): - read_engine_state(workdir) - - def test_write_does_not_truncate_existing_checkpoint_on_failure(self) -> None: - """A failing write leaves the previous engine_state.json intact. - - Simulates a mid-write interruption by patching the temp file's - ``write_text`` to raise after the previous checkpoint has been - written successfully. The reader must still see the previous - valid checkpoint, never a truncated or partial file. - """ - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - write_engine_state( - workdir, - queued=[TaskSpec(id="prev", priority=0.0)], - running=[], - status="running", - ) - - original_write_text = Path.write_text - - def _fail_on_tmp(self: Path, *args: object, **kwargs: object) -> int: - if ".tmp." in self.name: - raise OSError("simulated mid-write failure") - return original_write_text(self, *args, **kwargs) # type: ignore[arg-type] - - with unittest.mock.patch.object(Path, "write_text", _fail_on_tmp): - with self.assertRaises(OSError): - write_engine_state( - workdir, - queued=[TaskSpec(id="new", priority=0.0)], - running=[], - status="running", - ) - - result = read_engine_state(workdir) - assert result is not None - queued, _, _ = result - self.assertEqual([spec.id for spec in queued], ["prev"]) diff --git a/tests/autoresearch/persistence_test.py b/tests/autoresearch/persistence_test.py new file mode 120000 index 000000000..e5f8e1877 --- /dev/null +++ b/tests/autoresearch/persistence_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/persistence_test.py \ No newline at end of file diff --git a/tests/autoresearch/platform_test.py b/tests/autoresearch/platform_test.py deleted file mode 100644 index 235feffcd..000000000 --- a/tests/autoresearch/platform_test.py +++ /dev/null @@ -1,228 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import os -import shlex -import sys -import tempfile -import time -import unittest -import uuid -from pathlib import Path -from unittest.mock import patch - -from spdl.autoresearch.pipeline_optimization._ops._store import _set_queued_priority -from spdl.autoresearch.pipeline_optimization._platform import ( - _MetricsEvidence, - AutoresearchPlatform, - create_platform, -) -from spdl.autoresearch.pipeline_optimization._platform._agents import ( - _MockAgent, - _parse_agent_result, -) - -__all__: list[str] = [] - - -# The GitHub Actions Windows runner ships a conda Python whose PATH/DLL setup -# breaks nested ``cmd.exe`` invocation: ``subprocess.Popen(shell=True, ...)`` -# returns NT status ``0xC0000142`` (``STATUS_DLL_INIT_FAILED``) before the -# child shell can run anything, regardless of what command we give it. This is -# an environment bug in the runner, not in the production code under test, so -# we skip the two subprocess-launching tests there. Other CI matrices (Linux, -# macOS, internal Windows) still exercise this code path. -_SKIP_WINDOWS_GHA: bool = ( - sys.platform == "win32" and os.environ.get("GITHUB_ACTIONS") == "true" -) -_SKIP_REASON: str = ( - "GitHub Actions Windows runner cannot spawn nested cmd.exe (STATUS_DLL_INIT_FAILED)" -) - - -def _echo_marker_command(workdir: Path, line: str) -> str: - """Build a cross-platform shell command that prints ``line`` to stdout. - - We write a tiny shell script on disk and return its path as the command, - so the launched process is just the shell executing one builtin (``echo``) - — no external interpreter is loaded. - - Why not invoke a second ``python.exe``? On the GitHub Actions - Windows-miniconda runner, spawning a nested ``python.exe`` from - ``subprocess.Popen(shell=True, ...)`` returns NT status - ``0xC0000142`` (``STATUS_DLL_INIT_FAILED``) before the script runs. - cmd's ``echo`` is an internal command, so no DLL initialization happens - in the child process. - """ - workdir.mkdir(parents=True, exist_ok=True) - if sys.platform == "win32": - script = workdir / f"_marker_{uuid.uuid4().hex}.cmd" - script.write_text(f"@echo off\r\necho {line}\r\n", encoding="ascii") - return f'"{script}"' - script = workdir / f"_marker_{uuid.uuid4().hex}.sh" - script.write_text(f"#!/bin/sh\necho {shlex.quote(line)}\n", encoding="ascii") - script.chmod(0o755) - return shlex.quote(str(script)) - - -class _PlatformTest(unittest.TestCase): - def test_default_platform_has_capability_parts(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - with patch.dict( - "os.environ", - {"SPDL_AUTORESEARCH_PLATFORM_PROVIDERS": ""}, - ): - platform = create_platform("auto", Path(tmp)) - - self.assertIsInstance(platform, AutoresearchPlatform) - self.assertTrue(hasattr(platform, "workspace")) - self.assertTrue(hasattr(platform, "artifacts")) - self.assertTrue(hasattr(platform, "execution")) - self.assertTrue(hasattr(platform, "evidence")) - self.assertTrue(hasattr(platform, "agent")) - - @unittest.skipIf(_SKIP_WINDOWS_GHA, _SKIP_REASON) - def test_local_platform_runs_subprocess_and_collects_log_evidence(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - platform = create_platform("local", workdir) - - job_id = platform.execution.launch( - _echo_marker_command(workdir, "[autoresearch] step=1"), - workdir, - ) - self.assertIsNotNone(job_id) - assert job_id is not None - - for _ in range(50): - status = platform.execution.status(job_id) - if status == "SUCCEEDED": - break - time.sleep(0.05) - self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) - self.assertEqual( - "[autoresearch] step=1", platform.execution.progress(job_id) - ) - - metrics_dir = workdir / "runs" / "000_baseline" / "metrics" - evidence = platform.evidence.collect(job_id, metrics_dir) - - self.assertIsInstance(evidence, _MetricsEvidence) - self.assertIn("system metrics unavailable", evidence.system_metrics) - self.assertIn("[autoresearch] step=1", evidence.pipeline_stats_log) - - def test_local_dry_run_completes_without_launching_subprocess(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - platform = create_platform( - {"platform": "local", "local_execution_mode": "dry_run"}, - workdir, - ) - - job_id = platform.execution.launch( - _echo_marker_command(workdir, "not executed"), workdir - ) - self.assertIsNotNone(job_id) - assert job_id is not None - - self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) - self.assertIn("dry_run", platform.execution.progress(job_id) or "") - - @unittest.skipIf(_SKIP_WINDOWS_GHA, _SKIP_REASON) - def test_local_dataloader_only_uses_dataloader_command(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - platform = create_platform( - { - "platform": "local", - "local_execution_mode": "dataloader_only", - "local_dataloader_command": _echo_marker_command( - workdir, "[autoresearch] dataloader" - ), - }, - workdir, - ) - - job_id = platform.execution.launch( - _echo_marker_command(workdir, "training"), workdir - ) - self.assertIsNotNone(job_id) - assert job_id is not None - - for _ in range(50): - status = platform.execution.status(job_id) - if status == "SUCCEEDED": - break - time.sleep(0.05) - - self.assertEqual("SUCCEEDED", platform.execution.status(job_id)) - self.assertEqual( - "[autoresearch] dataloader", platform.execution.progress(job_id) - ) - - def test_mock_agent_is_selected_independently_of_platform(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - platform = create_platform( - {"platform": "local", "agent": "mock"}, - Path(tmp), - ) - - self.assertEqual("", platform.agent.run("prompt", Path(tmp), "phase")) - - def test_platform_config_validation_rejects_bad_local_mode(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - with self.assertRaisesRegex(ValueError, "local execution mode"): - create_platform( - {"platform": "local", "local_execution_mode": "unknown"}, - Path(tmp), - ) - - def test_unknown_remote_provider_is_explicit(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - with patch.dict( - "os.environ", - {"SPDL_AUTORESEARCH_PLATFORM_PROVIDERS": ""}, - ): - with self.assertRaisesRegex( - ValueError, "Unknown autoresearch platform" - ): - create_platform( - {"platform": "unknown_remote", "agent": "mock"}, - Path(tmp), - ) - - def test_agent_result_reports_parse_errors_without_llm(self) -> None: - agent = _MockAgent() - - parsed = _parse_agent_result(agent, '```json\n{"action": "stop"}\n```') - failed = _parse_agent_result(agent, "not json") - - self.assertEqual({"action": "stop"}, parsed.json) - self.assertIsNone(parsed.parse_error) - self.assertIsNone(failed.json) - self.assertEqual("No JSON object found", failed.parse_error) - - def test_queue_command_updates_checkpoint_priority(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - engine = workdir / "engine" - engine.mkdir(parents=True) - (engine / "checkpoint.json").write_text( - '{"status": "interrupted", "queued": [' - '{"id": "001_a", "priority": 10, "kind": "experiment", ' - '"payload": {"node": {"node_id": "001_a"}}}], "running": []}\n' - ) - - _set_queued_priority(workdir, "001_a", -5) - - text = (engine / "checkpoint.json").read_text() - self.assertIn('"priority": -5.0', text) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/autoresearch/platform_test.py b/tests/autoresearch/platform_test.py new file mode 120000 index 000000000..4c0b4cf7c --- /dev/null +++ b/tests/autoresearch/platform_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/platform_test.py \ No newline at end of file diff --git a/tests/autoresearch/prompts_test.py b/tests/autoresearch/prompts_test.py deleted file mode 100644 index e75177044..000000000 --- a/tests/autoresearch/prompts_test.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import unittest - -from spdl.autoresearch.pipeline_optimization._prompts import ( - load_knowledge, - load_prompt, - load_prompt_directory, -) - - -class LoadPromptTest(unittest.TestCase): - def test_loads_existing_prompt(self) -> None: - """A valid prompt name returns non-empty template content.""" - result = load_prompt("analyze", KNOWLEDGE="test-knowledge") - - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - def test_substitutes_placeholders(self) -> None: - """All __KEY__ placeholders are replaced with the supplied values.""" - result = load_prompt( - "headspace", - KNOWLEDGE="INJECTED_KNOWLEDGE", - PIPELINE_SCRIPT="/tmp/test.py", - PIPELINE_CODE="def main(): pass", - ) - - self.assertIn("INJECTED_KNOWLEDGE", result) - self.assertIn("/tmp/test.py", result) - self.assertIn("def main(): pass", result) - self.assertNotIn("__KNOWLEDGE__", result) - self.assertNotIn("__PIPELINE_SCRIPT__", result) - self.assertNotIn("__PIPELINE_CODE__", result) - - def test_missing_prompt_exits(self) -> None: - """Requesting a nonexistent prompt template triggers SystemExit.""" - with self.assertRaises(SystemExit): - load_prompt("nonexistent_prompt_that_does_not_exist") - - def test_headspace_prompt_requires_stop_after(self) -> None: - """The headspace prompt instructs the agent to include stop_after=500.""" - prompt = load_prompt( - "headspace", - KNOWLEDGE="", - PIPELINE_SCRIPT="/tmp/pipeline.py", - PIPELINE_CODE="def main():\n pass\n", - ) - - self.assertIn("stop_after=500", prompt) - self.assertIn("must include `stop_after=500`", prompt) - - def test_all_phase_prompts_loadable(self) -> None: - """Every phase prompt shipped with the package loads without error.""" - phase_prompts = [ - "analyze", - "apply_changes", - "apply_startup_repair", - "assess", - "headspace", - "instrument", - "plan_next", - ] - for name in phase_prompts: - with self.subTest(prompt=name): - result = load_prompt(name, KNOWLEDGE="k") - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - -class LoadPromptDirectoryTest(unittest.TestCase): - def test_loads_knowledge_directory(self) -> None: - """The knowledge directory contains at least one .md file.""" - result = load_prompt_directory("knowledge") - - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - def test_loads_supervisor_directory(self) -> None: - """The supervisor directory contains at least one .md file.""" - result = load_prompt_directory("supervisor") - - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - def test_loads_platform_directory(self) -> None: - """The platform directory contains at least one .md file.""" - result = load_prompt_directory("platform") - - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - def test_nonexistent_directory_returns_empty(self) -> None: - """A missing directory returns an empty string instead of raising.""" - result = load_prompt_directory("nonexistent_dir") - - self.assertEqual(result, "") - - def test_deterministic_order(self) -> None: - """Repeated loads produce identical output (sorted path order).""" - first = load_prompt_directory("knowledge") - second = load_prompt_directory("knowledge") - - self.assertEqual(first, second) - - -class LoadKnowledgeTest(unittest.TestCase): - def test_returns_nonempty_string(self) -> None: - """The combined knowledge + platform content is non-empty.""" - result = load_knowledge() - - self.assertIsInstance(result, str) - self.assertGreater(len(result), 0) - - def test_includes_knowledge_and_platform_content(self) -> None: - """The result contains the full text of both knowledge and platform directories.""" - result = load_knowledge() - knowledge_only = load_prompt_directory("knowledge") - platform_only = load_prompt_directory("platform") - - for section in (knowledge_only, platform_only): - if section: - self.assertIn(section, result) diff --git a/tests/autoresearch/prompts_test.py b/tests/autoresearch/prompts_test.py new file mode 120000 index 000000000..e8bdeffd2 --- /dev/null +++ b/tests/autoresearch/prompts_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/prompts_test.py \ No newline at end of file diff --git a/tests/autoresearch/state_test.py b/tests/autoresearch/state_test.py deleted file mode 100644 index 83f186da4..000000000 --- a/tests/autoresearch/state_test.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import unittest - -from spdl.autoresearch._common._state import SCHEMA_VERSION -from spdl.autoresearch.pipeline_optimization._ops._policy import ( - _normalize_config, - _normalize_state, -) - - -class StateTest(unittest.TestCase): - def test_schema_normalizers_add_versions_and_defaults(self) -> None: - """Empty dicts are normalized with schema version and default values.""" - config = _normalize_config({}) - state = _normalize_state({}) - - self.assertEqual(SCHEMA_VERSION, config["schema_version"]) - self.assertEqual(SCHEMA_VERSION, state["schema_version"]) - self.assertEqual("auto", config["platform"]) - self.assertEqual([], state["history"]) diff --git a/tests/autoresearch/state_test.py b/tests/autoresearch/state_test.py new file mode 120000 index 000000000..c1ed09991 --- /dev/null +++ b/tests/autoresearch/state_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/state_test.py \ No newline at end of file diff --git a/tests/autoresearch/supervisor_test.py b/tests/autoresearch/supervisor_test.py deleted file mode 100644 index b89bb3cca..000000000 --- a/tests/autoresearch/supervisor_test.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import unittest - -from spdl.autoresearch._common._supervisor import ( - _ClaudeSupervisor, - _CodexSupervisor, -) - - -class SupervisorTest(unittest.TestCase): - def test_claude_supervisor_builds_system_prompt_command(self) -> None: - """Claude supervisor passes system prompt and initial request as separate args.""" - command = _ClaudeSupervisor().command("SYSTEM", "INITIAL") - - self.assertEqual(["claude", "--system-prompt", "SYSTEM", "INITIAL"], command) - - def test_codex_supervisor_merges_system_and_request_into_single_prompt( - self, - ) -> None: - """Codex supervisor combines system prompt and request into one argument.""" - command = _CodexSupervisor().command("SYSTEM", "INITIAL") - - self.assertEqual("codex", command[0]) - self.assertEqual(2, len(command)) - self.assertIn("SYSTEM", command[1]) - self.assertIn("## User Request", command[1]) - self.assertIn("INITIAL", command[1]) diff --git a/tests/autoresearch/supervisor_test.py b/tests/autoresearch/supervisor_test.py new file mode 120000 index 000000000..8a3179b1f --- /dev/null +++ b/tests/autoresearch/supervisor_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/supervisor_test.py \ No newline at end of file diff --git a/tests/autoresearch/types_test.py b/tests/autoresearch/types_test.py deleted file mode 100644 index 6eb48db61..000000000 --- a/tests/autoresearch/types_test.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import unittest - -from spdl.autoresearch.core import ( - FailureKind, - FailurePhase, - FailureRecord, - HypothesisNode, -) - - -class FailureRecordTest(unittest.TestCase): - def test_round_trips_through_dict(self) -> None: - """FailureRecord survives to_dict/from_dict serialization.""" - record = FailureRecord( - kind=FailureKind.JOB_STARTUP_FAILED, - phase=FailurePhase.JOB, - message="MTP failed during initialization", - details={"component": "mtp"}, - job_id="job123", - created_at="2026-01-01T00:00:00", - ) - - loaded = FailureRecord.from_dict(record.to_dict()) - - self.assertEqual(FailureKind.JOB_STARTUP_FAILED, loaded.kind) - self.assertEqual(FailurePhase.JOB, loaded.phase) - self.assertEqual("MTP failed during initialization", loaded.message) - self.assertEqual({"component": "mtp"}, loaded.details) - self.assertEqual("job123", loaded.job_id) - - -class HypothesisNodeTest(unittest.TestCase): - def test_round_trips_through_dict(self) -> None: - """HypothesisNode with a failure survives to_dict/from_dict.""" - node = HypothesisNode( - node_id="001_bad_mtp", - name="bad_mtp", - status="failed", - failure=FailureRecord( - kind=FailureKind.JOB_STARTUP_FAILED, - phase=FailurePhase.JOB, - message="MTP failed during initialization", - details={"component": "mtp"}, - job_id="job123", - created_at="2026-01-01T00:00:00", - ), - ) - - loaded = HypothesisNode.from_dict(node.to_dict()) - - self.assertEqual("001_bad_mtp", loaded.node_id) - self.assertEqual("failed", loaded.status) - failure = loaded.failure - assert failure is not None - self.assertEqual(FailureKind.JOB_STARTUP_FAILED, failure.kind) - self.assertEqual({"component": "mtp"}, failure.details) - - def test_round_trips_without_failure(self) -> None: - """HypothesisNode without a failure round-trips cleanly.""" - node = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="completed", - priority=-1000, - ) - - loaded = HypothesisNode.from_dict(node.to_dict()) - - self.assertEqual("000_baseline", loaded.node_id) - self.assertEqual("completed", loaded.status) - self.assertIsNone(loaded.failure) - self.assertEqual(-1000, loaded.priority) diff --git a/tests/autoresearch/types_test.py b/tests/autoresearch/types_test.py new file mode 120000 index 000000000..af6fba6ef --- /dev/null +++ b/tests/autoresearch/types_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/types_test.py \ No newline at end of file diff --git a/tests/autoresearch/visualization_test.py b/tests/autoresearch/visualization_test.py deleted file mode 100644 index b7db34ef3..000000000 --- a/tests/autoresearch/visualization_test.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import unittest - -from spdl.autoresearch._common._visualization import ( - _edge_label, - _tree_font_sizes, -) - - -class VisualizationTest(unittest.TestCase): - def test_tree_edge_label_describes_experiment_evolution(self) -> None: - """Edge labels show change summary or startup repair attempt number.""" - parent = {"node_id": "001_parent", "name": "parent", "spec": {}} - child = { - "node_id": "002_child", - "name": "child", - "spec": { - "change_summary": "raise decode threads", - "description": "increase decode thread count", - }, - } - retry = { - "node_id": "003_retry", - "name": "retry", - "spec": {"_startup_retry_attempt": 2, "description": "repair"}, - } - - self.assertEqual("raise decode threads", _edge_label(parent, child)) - self.assertEqual("startup repair #2", _edge_label(parent, retry)) - - def test_tree_font_sizes_grow_with_tree_size(self) -> None: - """Larger trees get proportionally larger font sizes.""" - small = _tree_font_sizes(4, 1) - large = _tree_font_sizes(120, 10) - - self.assertGreater(large["title"], small["title"]) - self.assertGreater(large["legend"], small["legend"]) diff --git a/tests/autoresearch/visualization_test.py b/tests/autoresearch/visualization_test.py new file mode 120000 index 000000000..772ee6161 --- /dev/null +++ b/tests/autoresearch/visualization_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/visualization_test.py \ No newline at end of file diff --git a/tests/autoresearch/workflow_test.py b/tests/autoresearch/workflow_test.py deleted file mode 100644 index 0eb66e8a6..000000000 --- a/tests/autoresearch/workflow_test.py +++ /dev/null @@ -1,1049 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import asyncio -import json -import tempfile -import unittest -from pathlib import Path - -from spdl.autoresearch._common._state import _append_master_row -from spdl.autoresearch._common._visualization import _load_tsv -from spdl.autoresearch.core import ( - AnalysisResult, - FailureKind, - FailurePhase, - HypothesisNode, - TaskSpec, -) -from spdl.autoresearch.pipeline_optimization._ops import ( - _WorkflowStateStore, - PipelineOptimizationWorkflow, -) -from spdl.autoresearch.pipeline_optimization._ops._analysis_ops import ( - _update_on_complete, - MASTER_TABLE_HEADERS, -) -from spdl.autoresearch.pipeline_optimization._ops._failures import ( - _classify_terminal_job_failure, - _FAILURE_POLICIES, - _failure_summary, - _make_failure, - _read_failures, -) -from spdl.autoresearch.pipeline_optimization._ops._policy import ( - _build_change_set, - _change_summary_for_spec, - _compare_metric_value, - _extract_default_executor_concurrency, - _extract_param_changes, - _extract_total_threads, - _is_duplicate_spec, - _node_from_spec, - _retry_policy_for_failure, - _select_planning_node, - _spec_from_node, - _startup_retry_spec, - _validate_thread_budget, - write_state, -) -from spdl.autoresearch.pipeline_optimization._ops._source_ops import _build_apply_prompt -from spdl.autoresearch.pipeline_optimization._ops._store import _write_text_atomic -from spdl.autoresearch.pipeline_optimization._platform import ( - _MetricsEvidence, - create_platform, -) -from spdl.autoresearch.pipeline_optimization._platform._agents import _MockAgent -from spdl.autoresearch.pipeline_optimization._platform._local import _summarize_error - -__all__: list[str] = [] - - -def _config() -> dict: - return { - "pipeline_script": "", - "source_dir": "", - "scm": "", - "build_command": "", - "base_launch_command": "torchx run example --num-fetch-threads 8", - "stopping_criteria": { - "max_iterations": 20, - "patience": 5, - }, - "max_concurrency": 4, - "job_timeout_s": 600, - "poll_interval": 0, - } - - -def _state() -> dict: - return { - "iteration": 0, - "status": "looping", - "baseline_job": None, - "current_best": None, - "best_metric": None, - "plateau_count": 0, - "best_practices_tried": [], - "anchor_commit": "", - "history": [], - } - - -class _AutoresearchWorkflowTest(unittest.TestCase): - def test_fresh_load_creates_initial_must_run_specs(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - adapter = self._adapter(Path(tmp)) - - specs = adapter.load() - - self.assertEqual( - ["000_baseline", "000_headspace", "001_mtp"], - [spec.id for spec in specs], - ) - self.assertEqual([-1000, -999, -998], [spec.priority for spec in specs]) - - def test_checkpoint_writes_compatibility_files(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - specs = adapter.load() - - adapter.checkpoint(queued=specs[1:], running=specs[:1], status="running") - - engine_state = json.loads( - (workdir / "engine" / "engine_state.json").read_text() - ) - queue = json.loads((workdir / "engine" / "queue.json").read_text()) - active = json.loads((workdir / "engine" / "active.json").read_text()) - baseline_status = ( - workdir / "engine" / "nodes" / "000_baseline" / "status.txt" - ).read_text() - - self.assertEqual("running", engine_state["status"]) - self.assertEqual(2, engine_state["queued"]) - self.assertEqual(1, engine_state["running"]) - self.assertEqual( - ["000_headspace", "001_mtp"], [q["node_id"] for q in queue] - ) - self.assertEqual([], active) - self.assertEqual("queued\n", baseline_status) - - def test_checkpoint_round_trips_specs(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - specs = adapter.load() - adapter.checkpoint(queued=specs[1:], running=specs[:1], status="running") - - loaded = self._adapter(workdir).load() - - self.assertEqual( - ["000_headspace", "001_mtp", "000_baseline"], - [spec.id for spec in loaded], - ) - - def test_planning_is_blocked_until_must_run_experiments_finish(self) -> None: - baseline = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="completed", - ) - headspace = HypothesisNode( - node_id="000_headspace", - name="headspace_cache", - status="queued", - ) - mtp = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="completed", - ) - - selected = _select_planning_node( - baseline, - { - baseline.node_id: baseline, - headspace.node_id: headspace, - mtp.node_id: mtp, - }, - ) - - self.assertIsNone(selected) - - def _load_node(self, spec: TaskSpec) -> HypothesisNode: - """Extract node from a TaskSpec with a safe dict cast for Pyre.""" - node_data = spec.payload["node"] - assert isinstance(node_data, dict) - return HypothesisNode.from_dict(node_data) - - def test_child_spec_parents_under_baseline_when_goto_is_null(self) -> None: - """Experiments that start from anchor (goto=null) should be parented - under the baseline node, not whichever node triggered planning.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - specs = adapter.load() - # Simulate baseline completed with a commit. - baseline_node = self._load_node(specs[0]) - baseline_node.status = "completed" - baseline_node.commit = "baseline_commit_abc" - adapter._store.upsert_node(baseline_node) - - # Simulate MTP completed with a different commit. - mtp_node = self._load_node(specs[2]) - mtp_node.status = "completed" - mtp_node.commit = "mtp_commit_xyz" - adapter._store.upsert_node(mtp_node) - - # Create a child with goto=None (should parent under baseline). - child_spec = adapter._create_child_spec( - mtp_node, - {"name": "nvdec_decode", "goto": None}, - ) - child_node = self._load_node(child_spec) - - self.assertEqual("000_baseline", child_node.parent_id) - self.assertEqual("baseline_commit_abc", child_node.commit) - - def test_child_spec_parents_under_goto_commit_owner(self) -> None: - """Experiments with a goto commit should be parented under the node - that produced that commit.""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - specs = adapter.load() - - baseline_node = self._load_node(specs[0]) - baseline_node.status = "completed" - baseline_node.commit = "baseline_commit_abc" - adapter._store.upsert_node(baseline_node) - - mtp_node = self._load_node(specs[2]) - mtp_node.status = "completed" - mtp_node.commit = "mtp_commit_xyz" - adapter._store.upsert_node(mtp_node) - - # Create a child with goto pointing to MTP's commit. - child_spec = adapter._create_child_spec( - baseline_node, # default parent is baseline - {"name": "batch_on_mtp", "goto": "mtp_commit_xyz"}, - ) - child_node = self._load_node(child_spec) - - self.assertEqual(mtp_node.node_id, child_node.parent_id) - self.assertEqual("mtp_commit_xyz", child_node.commit) - - def test_child_spec_parents_under_baseline_without_goto(self) -> None: - """When goto is absent, the spec defaults to anchor (baseline).""" - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - specs = adapter.load() - - mtp_node = self._load_node(specs[2]) - mtp_node.status = "completed" - mtp_node.commit = "mtp_commit_xyz" - adapter._store.upsert_node(mtp_node) - - # Ensure baseline exists so _resolve_parent can find it. - baseline_node = self._load_node(specs[0]) - baseline_node.status = "completed" - adapter._store.upsert_node(baseline_node) - - # No goto key at all — .get("goto") returns None, which means - # "start from anchor". Parent should be baseline. - child_spec = adapter._create_child_spec( - mtp_node, - {"name": "some_exp"}, - ) - child_node = self._load_node(child_spec) - self.assertEqual("000_baseline", child_node.parent_id) - - def test_headspace_completion_selects_mtp_after_must_run_finish(self) -> None: - baseline = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="completed", - ) - headspace = HypothesisNode( - node_id="000_headspace", - name="headspace_cache", - status="completed", - spec={"_is_headspace": True}, - ) - mtp = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="completed", - ) - - selected = _select_planning_node( - headspace, - { - baseline.node_id: baseline, - headspace.node_id: headspace, - mtp.node_id: mtp, - }, - ) - - self.assertEqual(mtp, selected) - - def test_policy_helpers_cover_metric_and_thread_decisions(self) -> None: - self.assertEqual( - ("step_ms", -12.5), - _compare_metric_value({"steady_step_time_ms": 12.5, "duration_s": 100}), - ) - # Only --num-threads → use it directly. - self.assertEqual( - 12, - _extract_total_threads("--num-threads 12"), - ) - # Stage concurrency flags are not additive thread budgets. - self.assertIsNone( - _extract_total_threads("--num-fetch-threads 8 --num-decode-threads 16"), - ) - self.assertEqual( - 16, - _extract_default_executor_concurrency( - "--num-fetch-threads 8 --num-decode-threads 16" - ), - ) - self.assertEqual( - 8, - _extract_default_executor_concurrency("--num-fetch-threads 8"), - ) - self.assertEqual( - 16, - _extract_default_executor_concurrency("--num-decode-threads 16"), - ) - # No num_threads flag at all → None (unknown budget). - self.assertIsNone( - _extract_total_threads("--num_workers 8 --num_epochs 3"), - ) - self.assertIsNone( - _extract_total_threads(""), - ) - self.assertEqual( - ["ok", "fits_concurrency"], - [ - spec["name"] - for spec in _validate_thread_budget( - [ - {"name": "ok", "launch_command": "--num-threads 8"}, - {"name": "too_many", "launch_command": "--num-threads 64"}, - { - "name": "too_small", - "launch_command": ( - "--num-threads 8 --num-decode-threads 16" - ), - }, - { - "name": "fits_concurrency", - "launch_command": ( - "--num-threads 16 --num-decode-threads 16" - ), - }, - ], - 16, - ) - ], - ) - # Commands with no thread flags pass validation (unknown budget). - self.assertEqual( - ["pass_through"], - [ - spec["name"] - for spec in _validate_thread_budget( - [ - {"name": "pass_through", "launch_command": "--num_workers 8"}, - ], - 16, - ) - ], - ) - - # -- _extract_param_changes / _build_change_set / _is_duplicate_spec ------ - - def test_extract_param_changes_detects_flag_diffs(self) -> None: - base = "torchx run app --image $IMAGE --num_workers 8 --num_epochs 3" - - # New flag added. - self.assertEqual( - ["batch_size=48"], - _extract_param_changes(base + " --batch_size 48", base), - ) - - # Flag value changed. - self.assertEqual( - ["num_workers=16"], - _extract_param_changes( - "torchx run app --image $IMAGE --num_workers 16 --num_epochs 3", - base, - ), - ) - - # No diff when identical. - self.assertEqual([], _extract_param_changes(base, base)) - - # Empty commands return nothing. - self.assertEqual([], _extract_param_changes("", base)) - self.assertEqual([], _extract_param_changes(base, "")) - - def test_extract_param_changes_normalizes_dashes(self) -> None: - base = "torchx run app --num-fetch-threads 8" - exp = "torchx run app --num-fetch-threads 16" - changes = _extract_param_changes(exp, base) - self.assertEqual(["num_fetch_threads=16"], changes) - - def test_extract_param_changes_handles_negative_values(self) -> None: - base = "torchx run app --max_steps -1" - exp = "torchx run app --max_steps -2" - changes = _extract_param_changes(exp, base) - self.assertEqual(["max_steps=-2"], changes) - - def test_build_change_set_merges_explicit_and_param_changes(self) -> None: - base = "torchx run --image $IMAGE --num_workers 8" - spec = { - "changes": ["torch_compile"], - "launch_command": base + " --batch_size 48", - } - result = _build_change_set(spec, base) - self.assertEqual(frozenset({"torch_compile", "batch_size=48"}), result) - - def test_build_change_set_empty_for_baseline(self) -> None: - base = "torchx run --image $IMAGE --num_workers 8" - spec = {"changes": [], "launch_command": base} - self.assertEqual(frozenset(), _build_change_set(spec, base)) - - def test_build_change_set_normalizes_case(self) -> None: - spec = {"changes": ["Torch_Compile", " FUSED_ADAMW "]} - result = _build_change_set(spec, "") - self.assertEqual(frozenset({"torch_compile", "fused_adamw"}), result) - - def test_duplicate_requires_matching_change_sets(self) -> None: - base = "torchx run --image $IMAGE --num_workers 8" - baseline = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="completed", - spec={"changes": [], "launch_command": base}, - ) - mtp = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="completed", - spec={"changes": ["mtp"], "launch_command": base}, - ) - nodes = [baseline, mtp] - - # torch_compile: different code changes, same launch → NOT duplicate. - self.assertFalse( - _is_duplicate_spec( - {"changes": ["torch_compile"], "launch_command": base}, - nodes, - base, - ) - ) - - # fused_adamw: different code changes, same launch → NOT duplicate. - self.assertFalse( - _is_duplicate_spec( - {"changes": ["fused_adamw"], "launch_command": base}, - nodes, - base, - ) - ) - - # Exact same change set as baseline → IS duplicate. - self.assertTrue( - _is_duplicate_spec( - {"changes": [], "launch_command": base}, - nodes, - base, - ) - ) - - # Exact same change set as MTP → IS duplicate. - self.assertTrue( - _is_duplicate_spec( - {"changes": ["mtp"], "launch_command": base}, - nodes, - base, - ) - ) - - def test_duplicate_distinguishes_param_only_experiments(self) -> None: - base = "torchx run --image $IMAGE --num_workers 8" - batch_48 = HypothesisNode( - node_id="002_batch48", - name="batch_size_48", - status="completed", - spec={ - "changes": [], - "launch_command": base + " --batch_size 48", - }, - ) - nodes = [batch_48] - - # batch_size=64 has different param → NOT duplicate. - self.assertFalse( - _is_duplicate_spec( - {"changes": [], "launch_command": base + " --batch_size 64"}, - nodes, - base, - ) - ) - - # batch_size=48 again → IS duplicate. - self.assertTrue( - _is_duplicate_spec( - {"changes": [], "launch_command": base + " --batch_size 48"}, - nodes, - base, - ) - ) - - def test_duplicate_skips_failed_nodes(self) -> None: - base = "torchx run --image $IMAGE" - failed = HypothesisNode( - node_id="003_oom", - name="batch_64", - status="failed", - spec={ - "changes": [], - "launch_command": base + " --batch_size 64", - }, - ) - # Exact match of a failed node → NOT duplicate (allow retry). - self.assertFalse( - _is_duplicate_spec( - {"changes": [], "launch_command": base + " --batch_size 64"}, - [failed], - base, - ) - ) - - def test_duplicate_combination_vs_individual(self) -> None: - base = "torchx run --image $IMAGE" - mtp_only = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="completed", - spec={"changes": ["mtp"], "launch_command": base}, - ) - # Combination is not a dup of individual. - self.assertFalse( - _is_duplicate_spec( - { - "changes": ["mtp", "torch_compile"], - "launch_command": base, - }, - [mtp_only], - base, - ) - ) - - def test_duplicate_backward_compat_no_changes_field(self) -> None: - base = "torchx run --image $IMAGE --num_workers 8" - # Old spec without changes field — change set is derived purely from - # launch command diffs. - old_node = HypothesisNode( - node_id="002_batch48", - name="batch_size_48", - status="completed", - spec={"launch_command": base + " --batch_size 48"}, - ) - - # Same param diff → duplicate. - self.assertTrue( - _is_duplicate_spec( - {"launch_command": base + " --batch_size 48"}, - [old_node], - base, - ) - ) - - # Different param → not duplicate. - self.assertFalse( - _is_duplicate_spec( - {"launch_command": base + " --batch_size 64"}, - [old_node], - base, - ) - ) - - def test_startup_retry_inherits_changes(self) -> None: - node = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="failed", - spec={ - "name": "mtp", - "changes": ["mtp"], - "description": "try MTP", - "best_practices_tags": ["mtp"], - }, - failure=_make_failure( - FailureKind.JOB_STARTUP_FAILED, - FailurePhase.JOB, - "Tokenizer cannot pickle", - ), - ) - retry = _startup_retry_spec(node, _config()) - self.assertIsNotNone(retry) - assert retry is not None - self.assertEqual(["mtp"], retry["changes"]) - - def test_initial_nodes_have_changes(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - adapter = self._adapter(Path(tmp)) - specs = adapter.load() - nodes = [] - for spec in specs: - node_data = spec.payload["node"] - assert isinstance(node_data, dict) - nodes.append(HypothesisNode.from_dict(node_data)) - - by_name = {node.name: node for node in nodes} - self.assertEqual([], by_name["baseline"].spec["changes"]) - self.assertEqual( - ["cache_dataloader"], by_name["headspace_cache"].spec["changes"] - ) - self.assertEqual(["mtp"], by_name["mtp"].spec["changes"]) - - def test_store_update_spec_refreshes_checkpoint_and_active_view(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - self._write_base_files(workdir) - store = _WorkflowStateStore(workdir, _state()) - node = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="queued", - ) - spec = _spec_from_node(node) - store.save_scheduler_state(queued=[], running=[spec], status="running") - - node.status = "running" - node.job_id = "remote_job" - store.update_spec(spec, node) - - checkpoint = json.loads( - (workdir / "engine" / "checkpoint.json").read_text() - ) - active = json.loads((workdir / "engine" / "active.json").read_text()) - running_node = _node_from_spec(TaskSpec.from_dict(checkpoint["running"][0])) - - self.assertEqual("remote_job", running_node.job_id) - self.assertEqual("000_baseline", active[0]["node_id"]) - self.assertEqual("remote_job", active[0]["job_id"]) - self.assertIn("launched_at_iso", active[0]) - - def test_queue_view_includes_retry_lineage(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - self._write_base_files(workdir) - store = _WorkflowStateStore(workdir, _state()) - node = HypothesisNode( - node_id="002_retry", - name="retry", - status="queued", - spec={ - "_startup_retry_of": "001_mtp", - "_startup_retry_attempt": 1, - }, - ) - - store.save_scheduler_state( - queued=[_spec_from_node(node)], - running=[], - status="running", - ) - - queue = json.loads((workdir / "engine" / "queue.json").read_text()) - self.assertEqual("001_mtp", queue[0]["retry_of"]) - self.assertEqual(1, queue[0]["retry_attempt"]) - - def test_store_rejects_malformed_checkpoint(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - self._write_base_files(workdir) - engine_dir = workdir / "engine" - engine_dir.mkdir() - (engine_dir / "checkpoint.json").write_text( - json.dumps({"queued": [{"id": "bad", "payload": {}}]}) + "\n" - ) - store = _WorkflowStateStore(workdir, _state()) - - with self.assertRaisesRegex(ValueError, "payload.node"): - store.load_checkpoint() - - def test_store_persists_failure_view(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - self._write_base_files(workdir) - store = _WorkflowStateStore(workdir, _state()) - node = HypothesisNode( - node_id="001_failed", - name="failed", - status="failed", - failure=_make_failure( - FailureKind.BUILD_FAILED, - FailurePhase.BUILD, - "Build failed", - ), - ) - - store.upsert_node(node) - store.write_all() - - failure = json.loads( - ( - workdir / "engine" / "nodes" / "001_failed" / "failure.json" - ).read_text() - ) - engine_state = json.loads( - (workdir / "engine" / "engine_state.json").read_text() - ) - self.assertEqual("build_failed", failure["kind"]) - self.assertEqual({"build_failed": 1}, engine_state["failed_by_kind"]) - self.assertIn("build_failed", _failure_summary(workdir)) - self.assertIn("build_failed", _read_failures(workdir)) - - def test_adapter_records_structured_failure(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - adapter = self._adapter(workdir) - node = HypothesisNode(node_id="001_failed", name="failed") - spec = _spec_from_node(node) - failure = _make_failure( - FailureKind.LAUNCH_FAILED, - FailurePhase.LAUNCH, - "No launch command configured", - ) - - asyncio.run(adapter._record_failure(spec, node, failure)) - - stored = json.loads( - ( - workdir / "engine" / "nodes" / "001_failed" / "failure.json" - ).read_text() - ) - master_table = (workdir / "master_table.tsv").read_text() - self.assertEqual("launch_failed", stored["kind"]) - self.assertIn("launch_failed: No launch command configured", master_table) - - def test_completed_failure_history_uses_structured_failure(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - self._write_base_files(workdir) - state = _state() - node = HypothesisNode( - node_id="001_failed", - name="failed", - status="failed", - spec={"name": "failed"}, - ) - result = AnalysisResult( - structured={"metrics": {}, "findings": []}, - failure=_make_failure( - FailureKind.JOB_FAILED, - FailurePhase.JOB, - "Job failed", - ), - ) - - _update_on_complete(workdir, _config(), state, node, result) - - entry = state["history"][0] - self.assertNotIn("failure", entry) - self.assertEqual("job_failed", entry["structured"]["failure"]["kind"]) - - def test_job_failure_classifier_splits_startup_runtime_and_unknown(self) -> None: - startup = _classify_terminal_job_failure( - _MetricsEvidence( - system_metrics="", - pipeline_stats_log="TypeError: cannot pickle local function for MTP", - metrics_summary="", - ), - job_id="job_startup", - progress_seen=False, - ) - runtime = _classify_terminal_job_failure( - _MetricsEvidence( - system_metrics="", - pipeline_stats_log="[autoresearch] step=12\nCUDA out of memory", - metrics_summary="steady_step_time_ms: 12", - ), - job_id="job_runtime", - progress_seen=True, - ) - unknown = _classify_terminal_job_failure( - _MetricsEvidence( - system_metrics="", pipeline_stats_log="", metrics_summary="" - ), - job_id="job_unknown", - progress_seen=False, - ) - - self.assertEqual(FailureKind.JOB_STARTUP_FAILED, startup.kind) - self.assertEqual(FailureKind.JOB_RUNTIME_FAILED, runtime.kind) - self.assertEqual(FailureKind.JOB_FAILED, unknown.kind) - - def test_structured_metrics_evidence_guides_failure_classification(self) -> None: - startup = _classify_terminal_job_failure( - _MetricsEvidence( - system_metrics="", - pipeline_stats_log="", - metrics_summary="", - error_summary="TypeError: can't pickle tokenizer", - log_paths=["stderr.log"], - ), - job_id="job_startup", - progress_seen=False, - ) - runtime = _classify_terminal_job_failure( - _MetricsEvidence( - system_metrics="", - pipeline_stats_log="", - metrics_summary="", - progress_seen=True, - exit_code=1, - ), - job_id="job_runtime", - progress_seen=False, - ) - - self.assertEqual(FailureKind.JOB_STARTUP_FAILED, startup.kind) - self.assertEqual(["stderr.log"], startup.details["log_paths"]) - self.assertEqual(FailureKind.JOB_RUNTIME_FAILED, runtime.kind) - self.assertEqual(1, runtime.details["exit_code"]) - - def test_startup_failure_retry_is_bounded_for_mtp(self) -> None: - node = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="failed", - spec={ - "name": "mtp", - "description": "try MTP", - "best_practices_tags": ["mtp"], - }, - failure=_make_failure( - FailureKind.JOB_STARTUP_FAILED, - FailurePhase.JOB, - "Tokenizer cannot pickle", - ), - ) - - retry = _startup_retry_spec(node, _config()) - self.assertIsNotNone(retry) - assert retry is not None - self.assertEqual("mtp_startup_retry_1", retry["name"]) - self.assertEqual(1, retry["_startup_retry_attempt"]) - self.assertIn("pickling", retry["hypothesis"]) - - node.spec["_startup_retry_attempt"] = 2 - self.assertIsNone(_startup_retry_spec(node, _config())) - - def test_retry_policy_is_kind_specific(self) -> None: - node = HypothesisNode( - node_id="001_mtp", - name="mtp", - spec={"best_practices_tags": ["mtp"]}, - failure=_make_failure( - FailureKind.JOB_STARTUP_FAILED, - FailurePhase.JOB, - "startup", - ), - ) - policy = _retry_policy_for_failure(node, _config()) - self.assertIsNotNone(policy) - assert policy is not None - self.assertEqual(2, policy["max_attempts"]) - - node.failure = _make_failure( - FailureKind.BUILD_FAILED, - FailurePhase.BUILD, - "build", - ) - self.assertIsNone(_retry_policy_for_failure(node, _config())) - - def test_planning_prefers_non_startup_failed_retry(self) -> None: - baseline = HypothesisNode( - node_id="000_baseline", - name="baseline", - status="completed", - ) - headspace = HypothesisNode( - node_id="000_headspace", - name="headspace_cache", - status="completed", - spec={"_is_headspace": True}, - ) - mtp = HypothesisNode( - node_id="001_mtp", - name="mtp", - status="failed", - failure=_make_failure( - FailureKind.JOB_STARTUP_FAILED, - FailurePhase.JOB, - "startup", - ), - ) - retry = HypothesisNode( - node_id="002_mtp_startup_retry_1", - name="mtp_startup_retry_1", - status="failed", - spec={"_startup_retry_of": "001_mtp"}, - failure=_make_failure( - FailureKind.JOB_STARTUP_FAILED, - FailurePhase.JOB, - "startup", - ), - ) - better_candidate = HypothesisNode( - node_id="003_threads", - name="threads", - status="completed", - ) - - selected = _select_planning_node( - retry, - { - node.node_id: node - for node in [baseline, headspace, mtp, retry, better_candidate] - }, - ) - - self.assertEqual(better_candidate, selected) - - def test_startup_retry_uses_repair_prompt(self) -> None: - platform = create_platform({"platform": "local", "agent": "mock"}) - assert isinstance(platform.agent, _MockAgent) - platform.agent.responses["prompt:apply_startup_repair"] = ( - "failed during job startup __STARTUP_FAILURE_JSON__" - ) - prompt = _build_apply_prompt( - platform, - { - "name": "mtp_startup_retry_1", - "description": "repair", - "hypothesis": "fix startup", - "_startup_retry_attempt": 1, - "_startup_failure": {"kind": "job_startup_failed"}, - }, - "002_mtp_startup_retry_1", - "knowledge", - "/tmp/pipeline.py", - "def main():\n pass\n", - ) - - self.assertIn("failed during job startup", prompt) - self.assertIn("job_startup_failed", prompt) - - def test_headspace_node_uses_dedicated_prompt_with_knowledge(self) -> None: - platform = create_platform({"platform": "local", "agent": "mock"}) - assert isinstance(platform.agent, _MockAgent) - platform.agent.responses["prompt:headspace"] = ( - "headspace prompt __KNOWLEDGE__ __PIPELINE_CODE__" - ) - - prompt = _build_apply_prompt( - platform, - { - "name": "headspace_cache", - "description": "Wrap with CacheDataLoader for headspace analysis", - "_is_headspace": True, - }, - "000_headspace", - "knowledge", - "/tmp/pipeline.py", - "def main():\n pass\n", - ) - - self.assertIn("headspace prompt", prompt) - self.assertIn("knowledge", prompt) - self.assertIn("def main", prompt) - - def test_change_summary_is_concise_and_persisted_in_master_table(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - workdir = Path(tmp) - (workdir / "master_table.tsv").write_text( - "\t".join(MASTER_TABLE_HEADERS) + "\n" - ) - - summary = _change_summary_for_spec( - { - "description": ( - "Increase decode thread count while preserving the rest " - "of the pipeline" - ) - } - ) - _append_master_row( - workdir, - { - "run_id": "001_threads", - "name": "threads", - "status": "completed", - "change_summary": summary, - "sm_util_pct": "50", - }, - MASTER_TABLE_HEADERS, - ) - - rows = _load_tsv(workdir / "master_table.tsv") - - self.assertEqual("Increase decode thread count while", summary) - self.assertEqual(summary, rows[0]["change_summary"]) - - def test_every_failure_kind_has_policy(self) -> None: - self.assertEqual(set(FailureKind.__members__.values()), set(_FAILURE_POLICIES)) - self.assertTrue(_FAILURE_POLICIES[FailureKind.JOB_STARTUP_FAILED].retryable) - - def test_error_summary_prefers_traceback_block(self) -> None: - summary = _summarize_error( - "before\n" - "Traceback (most recent call last):\n" - ' File "x.py", line 1, in \n' - "TypeError: cannot pickle tokenizer\n" - "after\n" - ) - - self.assertIsNotNone(summary) - assert summary is not None - self.assertIn("Traceback", summary) - self.assertIn("cannot pickle", summary) - - def test_atomic_write_replaces_json(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - path = Path(tmp) / "data.json" - - _write_text_atomic(path, '{"a": 1}\n') - _write_text_atomic(path, '{"a": 2}\n') - - self.assertEqual({"a": 2}, json.loads(path.read_text())) - - def _adapter(self, workdir: Path) -> PipelineOptimizationWorkflow: - self._write_base_files(workdir) - return PipelineOptimizationWorkflow( - workdir=workdir, - config=_config(), - state=_state(), - platform=create_platform({"platform": "auto", "agent": "mock"}, workdir), - ) - - def _write_base_files(self, workdir: Path) -> None: - workdir.mkdir(parents=True, exist_ok=True) - (workdir / "config.json").write_text(json.dumps(_config()) + "\n") - write_state(workdir, _state()) - (workdir / "master_table.tsv").write_text( - "\t".join(MASTER_TABLE_HEADERS) + "\n" - ) diff --git a/tests/autoresearch/workflow_test.py b/tests/autoresearch/workflow_test.py new file mode 120000 index 000000000..2a6f67900 --- /dev/null +++ b/tests/autoresearch/workflow_test.py @@ -0,0 +1 @@ +../../../src/spdl/autoresearch/tests/workflow_test.py \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..d6800e544 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,46 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# OSS-side test bootstrap. +# +# 1. Make ``spdl.io.tests`` importable by registering a synthetic namespace +# package whose ``__path__`` points at the symlinked ``tests/io/`` shim. +# Tests do ``from spdl.io.tests.fixture import ...``; the wheel ships +# ``spdl`` and ``spdl.io`` but excludes the ``tests`` subpackage, so we +# splice it in at pytest startup. +# +# 2. Some pipeline tests pickle local helper classes (e.g. ``_Wrap``) and +# unpickle them inside a fresh subinterpreter. A new subinterpreter is +# initialized from scratch and only sees ``PYTHONPATH`` — pytest's runtime +# ``sys.path`` additions are not inherited. Prepend ``tests/pipeline/`` to +# ``PYTHONPATH`` so the subinterpreter can resolve the originating test +# module when unpickling. +# +# This file is unused in fbcode (Buck wires up ``spdl.io.tests`` directly via +# the ``//spdl/io/tests:fixture`` library). + +import importlib +import os +import sys +import types + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_TESTS_IO = os.path.join(_HERE, "io") +_TESTS_PIPELINE = os.path.join(_HERE, "pipeline") + +# Splice ``spdl.io.tests`` to point at our symlink shim. ``spdl.io`` itself is +# loaded from the installed wheel; we attach a synthetic ``tests`` submodule. +_spdl_io = importlib.import_module("spdl.io") +_tests_pkg = types.ModuleType("spdl.io.tests") +_tests_pkg.__path__ = [_TESTS_IO] +sys.modules["spdl.io.tests"] = _tests_pkg +_spdl_io.tests = _tests_pkg # type: ignore[attr-defined] + +# Subinterpreter helper: ensure pickled test classes resolve in fresh interps. +_existing = os.environ.get("PYTHONPATH", "") +os.environ["PYTHONPATH"] = ( + os.pathsep.join([_TESTS_PIPELINE, _existing]) if _existing else _TESTS_PIPELINE +) diff --git a/tests/cuda/__init__.py b/tests/cuda/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/dataloader/cache_dataloader_test.py b/tests/dataloader/cache_dataloader_test.py deleted file mode 100644 index 6893b33e8..000000000 --- a/tests/dataloader/cache_dataloader_test.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import unittest -from collections.abc import Iterator - -from spdl.dataloader import CacheDataLoader -from spdl.pipeline import cache_iterator - - -class TestCacheIterator(unittest.TestCase): - def test_cache_iterator(self) -> None: - """cache_iterator returns the cached values""" - - ite = iter(cache_iterator(range(5), 3)) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - def test_cache_iterator_cache_return_after(self) -> None: - """cache_iterator returns the cached values""" - - ite = iter(cache_iterator(range(7), 3, return_caches_after=5)) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - self.assertEqual(next(ite), 3) - self.assertEqual(next(ite), 4) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - def test_cache_iterator_cache_return_after_len(self) -> None: - """cache_iterator returns the cached values""" - - ite = iter(cache_iterator(range(7), 3, return_caches_after=5, stop_after=10)) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - self.assertEqual(next(ite), 3) - self.assertEqual(next(ite), 4) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - - with self.assertRaises(StopIteration): - next(ite) - - -class TestCacheDataLoader(unittest.TestCase): - def test_CacheDataLoader(self) -> None: - """Smoke test""" - - class DL: - def __init__(self, n: int) -> None: - self.n = n - - def __iter__(self) -> Iterator[int]: - yield from range(self.n) - - def __len__(self) -> int: - return self.n - - N = 8 - dl = CacheDataLoader(DL(N), num_caches=2, return_caches_after=3, stop_after=N) - - self.assertEqual(dl.n, N) - self.assertEqual(len(dl), N) - - ite = iter(dl) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - - self.assertEqual(next(ite), 0) - - with self.assertRaises(StopIteration): - next(ite) diff --git a/tests/dataloader/cache_dataloader_test.py b/tests/dataloader/cache_dataloader_test.py new file mode 120000 index 000000000..29e236ed7 --- /dev/null +++ b/tests/dataloader/cache_dataloader_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/cache_dataloader_test.py \ No newline at end of file diff --git a/tests/dataloader/dataloader_test.py b/tests/dataloader/dataloader_test.py deleted file mode 100644 index afd5c807b..000000000 --- a/tests/dataloader/dataloader_test.py +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -import os -import platform -import time -import unittest - -from spdl.dataloader import DataLoader - - -def get_dl(*args, timeout=3, num_threads=2, **kwargs): - # on default values - # timeout -> so that test would fail rather stack - # num_threads -> keep it minimum but have more than 1 - return DataLoader(*args, **kwargs, num_threads=num_threads, timeout=timeout) - - -class TestDataLoader(unittest.TestCase): - def test_dataloader_iterable(self) -> None: - src = list(range(10)) - - dl = get_dl(src) - - self.assertEqual(sorted(dl), src) - - def test_dataloader_stateful_iterable(self) -> None: - class src: - def __init__(self, num_items: int = 10): - self.num_iter = 0 - self.num_items = num_items - - def __iter__(self): - for i in range(self.num_items): - yield (self.num_iter, i) - self.num_iter += 1 - - dl = get_dl(src()) - - self.assertEqual(sorted(dl), [(0, i) for i in range(10)]) - self.assertEqual(sorted(dl), [(1, i) for i in range(10)]) - self.assertEqual(sorted(dl), [(2, i) for i in range(10)]) - - def test_dataloader_preprocess(self) -> None: - """preprocessor process the value of the source""" - src = list(range(10)) - - def double(x): - time.sleep(0.05 * x) # to reduce flakiness from multi-threading - return 2 * x - - dl = get_dl(src, preprocessor=double) - - self.assertEqual(sorted(dl), [i * 2 for i in range(10)]) - - def test_dataloader_preprocess_in_order(self) -> None: - """When output_order='input', the order must be preserved.""" - src = list(range(10, -1, -1)) - - def delay(x): - time.sleep(0.1 * x) - return x - - dl = get_dl(src, preprocessor=delay, output_order="input") - - self.assertEqual(list(dl), src) - - dl = get_dl(src, preprocessor=delay, output_order="completion") - - self.assertNotEqual(list(dl), src) - - @unittest.skipIf( - platform.system() == "Darwin" and "CI" in os.environ, - "GitHub macOS CI is not timely enough.", - ) - def test_dataloader_buffer_size(self) -> None: - """Bigger buffer_size allows the BG to proceed while FG is not fetching the data""" - src = list(range(12)) - - def delay(x): - time.sleep(0.05) - return x - - def test(dl): - # Kick off the background thread - dli = iter(dl) - self.assertEqual(next(dli), 0) - - # Wait: (simulate foreground load) - time.sleep(1) - - # Iterate the rest - t0 = time.monotonic() - result = list(dli) - elapsed = time.monotonic() - t0 - print(elapsed) - self.assertEqual(result, src[1:]) - return elapsed - - # With buffer_size == 1, then the background thread cannot proceed - # while foreground thread does not fetch any. - dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=1) - elapsed = test(dl) - self.assertGreater(elapsed, 0.3) - - # With bigger buffer_size, the background thread proceed - # while foreground thread does not fetch any. - dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=len(src)) - elapsed = test(dl) - self.assertLess(elapsed, 0.15) - - def test_dataloader_num_threads(self) -> None: - """Increasing the num_threads reduces the overall time.""" - src = list(range(10)) - - def delay(x): - time.sleep(0.1) - return x - - def test(dl): - t0 = time.monotonic() - result = list(dl) - elapsed = time.monotonic() - t0 - print(elapsed) - self.assertEqual(sorted(result), src) - return elapsed - - dl = get_dl(src, preprocessor=delay, num_threads=1, buffer_size=1) - self.assertGreater(test(dl), 0.8) - - dl = get_dl(src, preprocessor=delay, num_threads=len(src), buffer_size=1) - self.assertLess(test(dl), 0.6) - - def test_dataloader_batch(self) -> None: - """batching works with or without dropping""" - src = list(range(10)) - - dl = get_dl(src, batch_size=3, drop_last=False) - - self.assertEqual(list(dl), [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) - - dl = get_dl(src, batch_size=3, drop_last=True) - - self.assertEqual(list(dl), [[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - - def test_dataloader_aggregate(self) -> None: - """Aggregator processes the batched input""" - src = list(range(10)) - - def agg(vals: list[int]) -> tuple[int, int, int, int]: - return len(vals), min(vals), max(vals), sum(vals) - - dl = get_dl(src, batch_size=3, drop_last=False, aggregator=agg) - - expected = [ - (3, 0, 2, 3), - (3, 3, 5, 12), - (3, 6, 8, 21), - (1, 9, 9, 9), - ] - - self.assertEqual(list(dl), expected) diff --git a/tests/dataloader/dataloader_test.py b/tests/dataloader/dataloader_test.py new file mode 120000 index 000000000..7ee304e7e --- /dev/null +++ b/tests/dataloader/dataloader_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/dataloader_test.py \ No newline at end of file diff --git a/tests/dataloader/iterator_test.py b/tests/dataloader/iterator_test.py deleted file mode 100644 index 2c9f6d09e..000000000 --- a/tests/dataloader/iterator_test.py +++ /dev/null @@ -1,370 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -import functools -import pickle -import random -import unittest -import warnings -from collections.abc import Iterator -from functools import partial -from unittest.mock import patch - -from spdl.pipeline import iterate_in_subprocess as _iterate_in_subprocess -from spdl.source.utils import ( - embed_shuffle, - IterableWithShuffle, - MergeIterator, - repeat_source, -) - - -def _ignore_fork_warning(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - return wrapper - - -def iterate_in_subprocess(fn, *, timeout=10, **kwargs): - return _iterate_in_subprocess(fn, timeout=timeout, **kwargs) - - -class TestMergeIterator(unittest.TestCase): - def test_mergeiterator_ordered(self) -> None: - """MergeIterator iterates multiple iterators""" - - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables)) - self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2, 12, 22]) - - def test_mergeiterator_ordered_stop_after_first_exhaustion(self) -> None: - """MergeIterator stops after the first exhaustion""" - - iterables = [ - [0], - [10, 11, 12], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables, stop_after=-1)) - self.assertEqual(result, [0, 10, 20]) - - iterables = [ - [0, 1, 2], - [10], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables, stop_after=-1)) - self.assertEqual(result, [0, 10, 20, 1]) - - iterables = [ - [0, 1, 2], - [10, 11], - [20], - ] - - result = list(MergeIterator(iterables, stop_after=-1)) - self.assertEqual(result, [0, 10, 20, 1, 11]) - - def test_mergeiterator_ordered_stop_after_N(self) -> None: - """MergeIterator stops after N items are yielded""" - - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables, stop_after=1)) - self.assertEqual(result, [0]) - - result = list(MergeIterator(iterables, stop_after=5)) - self.assertEqual(result, [0, 10, 20, 1, 11]) - - result = list(MergeIterator(iterables, stop_after=7)) - self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2]) - - def test_mergeiterator_ordered_stop_after_minus1(self) -> None: - """MergeIterator stops after all the iterables are exhausted""" - - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables)) - self.assertEqual(result, [0, 10, 20, 1, 11, 21, 2, 12, 22]) - - iterables = [ - [0, 1, 2], - [10], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables)) - self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) - - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20], - ] - - result = list(MergeIterator(iterables)) - self.assertEqual(result, [0, 10, 20, 1, 11, 2, 12]) - - def test_mergeiterator_ordered_n(self) -> None: - """with stop_after=N, MergeIterator continues iterating after exhaustion.""" - iterables = [ - [0, 1, 2], - [10], - [20, 21, 22], - ] - - result = list(MergeIterator(iterables, stop_after=5)) - self.assertEqual(result, [0, 10, 20, 1, 21]) - - result = list(MergeIterator(iterables, stop_after=7)) - self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) - - result = list(MergeIterator(iterables, stop_after=8)) - self.assertEqual(result, [0, 10, 20, 1, 21, 2, 22]) - - def test_mergeiterator_stochastic_smoke_test(self) -> None: - """MergeIterator with probabilitiies do not get stuck.""" - - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20, 21, 22], - ] - - weights = [1, 1, 1] - - result = list(MergeIterator(iterables, weights=weights)) - self.assertEqual(set(result), {0, 1, 2, 10, 11, 12, 20, 21, 22}) - - def test_mergeiterator_stochastic_rejects_zero(self) -> None: - """weight=0 is rejected.""" - weights = [1, 0] - - with self.assertRaises(ValueError): - MergeIterator([[1]], weights=weights) - - weights = [1, 0.0] - - with self.assertRaises(ValueError): - MergeIterator([[1]], weights=weights) - - def test_mergeiterator_skip_zero_weight(self) -> None: - """Iterables with zero weight are skipped.""" - iterables = [ - [0, 1, 2], - [10, 11, 12], - [20, 21, 22], - [30, 31, 32], - ] - - weights = [1, 0, 2, 0] - - merge_iter = MergeIterator(iterables, weights=weights) - - self.assertEqual(len(merge_iter.iterables), 2) - self.assertEqual(merge_iter.iterables[0], [0, 1, 2]) - self.assertEqual(merge_iter.iterables[1], [20, 21, 22]) - - self.assertIsNotNone(merge_iter.weights) - # pyre-ignore[16]: weights is not None after assertion - self.assertEqual(len(merge_iter.weights), 2) - # pyre-ignore[16]: weights is not None after assertion - self.assertEqual(merge_iter.weights[0], 1) - # pyre-ignore[16]: weights is not None after assertion - self.assertEqual(merge_iter.weights[1], 2) - - result = list(merge_iter) - self.assertEqual(set(result), {0, 1, 2, 20, 21, 22}) - - def test_mergeiterator_stochastic_stop_after_N(self) -> None: - """Values are taken from iterables with higher weights""" - weights = [1000000, 1] - - iterables = [ - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], - ] - - result = list(MergeIterator(iterables, weights=weights, stop_after=3)) - self.assertEqual(result, [0, 1, 2]) - - def test_mergeiterator_stochastic_stop_after_first_exhaustion(self) -> None: - """Values are taken from iterables with higher weights""" - weights = [1000000, 1] - - iterables = [ - [0, 1, 2, 3], - [10, 11, 12, 13], - ] - - result = list(MergeIterator(iterables, weights=weights, stop_after=-1)) - self.assertEqual(result, [0, 1, 2, 3]) - - -class TestRepeatSource(unittest.TestCase): - def test_repeat_source_iterable_with_shuffle(self) -> None: - """repeat_source repeats source while calling shuffle""" - - class _IteWithShuffle: - def __init__(self) -> None: - self.vals = list(range(3)) - - def shuffle(self, seed: int) -> None: - assert isinstance(seed, int) - self.vals = self.vals[1:] + self.vals[:1] - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - src = _IteWithShuffle() - gen = iter(repeat_source(src, epoch=2)) - - with patch.object(src, "shuffle", side_effect=src.shuffle) as mock_method: - self.assertEqual(next(gen), 1) - mock_method.assert_called_with(seed=0) - self.assertEqual(next(gen), 2) - self.assertEqual(next(gen), 0) - - self.assertEqual(next(gen), 2) - mock_method.assert_called_with(seed=1) - self.assertEqual(next(gen), 0) - self.assertEqual(next(gen), 1) - - self.assertEqual(next(gen), 0) - mock_method.assert_called_with(seed=2) - self.assertEqual(next(gen), 1) - self.assertEqual(next(gen), 2) - - self.assertEqual(next(gen), 1) - mock_method.assert_called_with(seed=3) - self.assertEqual(next(gen), 2) - self.assertEqual(next(gen), 0) - - self.assertEqual(next(gen), 2) - mock_method.assert_called_with(seed=4) - self.assertEqual(next(gen), 0) - self.assertEqual(next(gen), 1) - - def test_repeat_source_iterable(self) -> None: - """repeat_source works Iterable without shuffle method""" - - class _IteWithoutShuffle: - def __init__(self) -> None: - self.vals = list(range(3)) - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - src = _IteWithoutShuffle() - gen = iter(repeat_source(src, epoch=2)) - - for _ in range(100): - self.assertEqual(next(gen), 0) - self.assertEqual(next(gen), 1) - self.assertEqual(next(gen), 2) - - def test_repeat_source_picklable(self) -> None: - """repeat_source is picklable.""" - - src = list(range(10)) - src = repeat_source(src) - - serialized = pickle.dumps(src) - src2 = pickle.loads(serialized) - - for _ in range(3): - for i in range(10): - self.assertEqual(next(src), i) - self.assertEqual(next(src2), i) - - -class IterableWithShuffleSource: - def __init__(self, n: int) -> None: - self.vals = list(range(n)) - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - def shuffle(self, seed: int) -> None: - random.seed(seed) - random.shuffle(self.vals) - - -class SourceIterableWithShuffle(IterableWithShuffle[int]): - def __init__(self, n: int) -> None: - self.i = 0 - self.vals = list(range(n)) - - def shuffle(self, seed: int) -> None: - assert isinstance(seed, int) - self.vals = self.vals[1:] + self.vals[:1] - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - -class TestShuffleAndIterate(unittest.TestCase): - def test_shuffle_and_iterate_picklable(self) -> None: - """The result of embed_shuffle must be pickable (for multiprocessing)""" - - src = embed_shuffle(IterableWithShuffleSource(10)) - state = pickle.dumps(src) - src2 = pickle.loads(state) - - # pyre-ignore[16]: embed_shuffle returns an object with src attribute - self.assertEqual(src.src.vals, src2.src.vals) - - def test_shuffle_and_iterate(self) -> None: - N = 10 - - src = embed_shuffle(IterableWithShuffleSource(N)) - - ref = list(range(N)) - for i in range(3): - random.seed(i) - random.shuffle(ref) - - hyp = list(src) - self.assertEqual(hyp, ref) - - @_ignore_fork_warning - def test_move_iterable_to_subprocess_success_iterable_with_shuffle(self) -> None: - """IterableWithShuffle can be executed in the subprocess.""" - iterator = iterate_in_subprocess( - partial(embed_shuffle, SourceIterableWithShuffle(3)) - ) - - self.assertEqual(list(iterator), [1, 2, 0]) - self.assertEqual(list(iterator), [2, 0, 1]) - self.assertEqual(list(iterator), [0, 1, 2]) diff --git a/tests/dataloader/iterator_test.py b/tests/dataloader/iterator_test.py new file mode 120000 index 000000000..346daba10 --- /dev/null +++ b/tests/dataloader/iterator_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/iterator_test.py \ No newline at end of file diff --git a/tests/dataloader/sampler_test.py b/tests/dataloader/sampler_test.py deleted file mode 100644 index a59b98921..000000000 --- a/tests/dataloader/sampler_test.py +++ /dev/null @@ -1,343 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import functools -import unittest -import warnings -from collections import Counter -from collections.abc import Callable -from functools import partial -from typing import TypeVar - -import numpy as np -from parameterized import parameterized -from spdl.pipeline import iterate_in_subprocess -from spdl.source import ( - DistributedDeterministicSampler, - DistributedRandomSampler, - SizedIterable, - SizedIterableWithShuffle, -) -from spdl.source.utils import embed_shuffle - -_F = TypeVar("_F", bound=Callable[..., object]) - - -def _ignore_fork_warning(fn: _F) -> _F: - @functools.wraps(fn) - def wrapper(*args: object, **kwargs: object) -> object: - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - # pyre-ignore[7] - return wrapper - - -class TestDistributedSamplerInterface(unittest.TestCase): - def test_distributed_sampler_interface(self) -> None: - """samplers conform to Iterable/IterableWithShuffle protocol""" - self.assertIsInstance( - DistributedRandomSampler(9, rank=0, world_size=1), SizedIterableWithShuffle - ) - self.assertIsInstance( - DistributedDeterministicSampler(9, rank=0, world_size=1), SizedIterable - ) - - -class TestDistributedSamplerDeterministic(unittest.TestCase): - def test_deterministic_iter(self) -> None: - """without distributed, deterministic iteration behaves same as `range(N)`""" - N = 30 - sampler = DistributedDeterministicSampler(N, rank=0, world_size=1) - self.assertEqual(len(sampler), N) - self.assertEqual(list(sampler), list(range(N))) - - def test_deterministic_iter_distributed(self) -> None: - """deterministic iteration behaves same as `range(rank, M, world_size)`""" - N = 26 - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"The size of dataset \(\d+\) is not divisible by the " - r"world size \(\d+\)\. Some samples are never visited\." - ), - category=UserWarning, - ) - for world_size in range(1, N + 1): - len_ = N // world_size - max_ = len_ * world_size - c = Counter() - for rank in range(world_size): - print(f"{N=}, {world_size=}, {rank=}, {len_=}, {max_=}") - sampler = DistributedDeterministicSampler( - N, rank=rank, world_size=world_size - ) - self.assertEqual(len(sampler), len_) - - indices = list(sampler) - self.assertEqual(indices, list(range(rank, max_, world_size))) - c.update(indices) - - # Check that together, the samplers covered the whole dataset - num_iters = N // world_size * world_size - self.assertEqual(c.total(), num_iters) - self.assertEqual(len(c.keys()), num_iters) - self.assertEqual(set(c.keys()), set(range(num_iters))) - self.assertTrue(all(v == 1 for v in c.values())) - - def test_deterministic_iter_stable_across_epochs(self) -> None: - """Deterministic sampler produces the same sequence on every iteration.""" - N = 30 - world_size = 4 - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"The size of dataset \(\d+\) is not divisible by the " - r"world size \(\d+\)\. Some samples are never visited\." - ), - category=UserWarning, - ) - for rank in range(world_size): - sampler = DistributedDeterministicSampler( - N, rank=rank, world_size=world_size - ) - first_epoch = list(sampler) - for _ in range(5): - self.assertEqual(list(sampler), first_epoch) - - -class TestDistributedSamplerRandom(unittest.TestCase): - def test_shuffle(self) -> None: - """shuffling makes sampler generates different indices.""" - N = 640 - rank = 3 - world_size = 8 - - previous: list[int] = [] - for epoch in range(100): - sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) - sampler.shuffle(seed=epoch) - - indices = list(sampler) - print(f"{indices=}") - self.assertNotEqual(indices, previous) - previous = indices - - def test_shuffle_epoch_loop(self) -> None: - """Calling shuffle(seed=epoch) on a single sampler produces different sequences each epoch.""" - N = 640 - world_size = 8 - - for rank in range(world_size): - sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) - previous: list[int] = [] - for epoch in range(10): - sampler.shuffle(seed=epoch) - indices = list(sampler) - self.assertEqual(len(indices), N // world_size) - self.assertNotEqual(indices, previous) - previous = indices - - def test_shuffle_is_stateless(self) -> None: - """shuffle(seed) output depends only on the seed, not on prior iteration history.""" - N = 640 - world_size = 8 - - for rank in range(world_size): - sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) - for i in range(10): - sampler.shuffle(seed=i) - list(sampler) - sampler.shuffle(seed=5) - after_many = list(sampler) - - fresh = DistributedRandomSampler(N, rank=rank, world_size=world_size) - fresh.shuffle(seed=5) - from_fresh = list(fresh) - - self.assertEqual(after_many, from_fresh) - - def test_shuffle_epoch_loop_mutual_exclusive(self) -> None: - """All ranks together cover the full dataset at each epoch when using shuffle(seed=epoch).""" - N = 640 - world_size = 8 - - samplers = [ - DistributedRandomSampler(N, rank=rank, world_size=world_size) - for rank in range(world_size) - ] - - for epoch in range(10): - c = Counter() - for sampler in samplers: - sampler.shuffle(seed=epoch) - c.update(sampler) - - self.assertEqual(c.total(), N) - self.assertEqual(len(c.keys()), N) - self.assertEqual(set(c.keys()), set(range(N))) - self.assertTrue(all(v == 1 for v in c.values())) - - @parameterized.expand( - [ - (None,), - (1,), - ] - ) - def test_repeat(self, w: int | None) -> None: - """Without calling shuffle, sampler generates the same sequence.""" - N = 40 - world_size = 8 - - weights = None if w is None else [1.0] * N - for rank in range(world_size): - previous = [] - for i in range(100): - sampler = DistributedRandomSampler( - N, rank=rank, world_size=world_size, weights=weights - ) - - indices = list(sampler) - print(f"{indices=}") - if i > 0: - self.assertEqual(indices, previous) - previous = indices - - @parameterized.expand( - [ - (True,), - (False,), - ] - ) - def test_mutual_exclusive(self, shuffle: bool) -> None: - """Without weights, samplers generate mutually exclusive sets""" - N = 640 - world_size = 8 - - for epoch in range(100): - c = Counter() - for rank in range(world_size): - sampler = DistributedRandomSampler(N, rank=rank, world_size=world_size) - if shuffle: - sampler.shuffle(seed=epoch) - c.update(sampler) - - self.assertEqual(c.total(), N) - self.assertEqual(len(c.keys()), N) - self.assertEqual(set(c.keys()), set(range(N))) - self.assertTrue(all(v == 1 for v in c.values())) - - @parameterized.expand( - [ - (True,), - (False,), - ] - ) - def test_mutual_exclusive_num_draws(self, shuffle: bool) -> None: - """Without weights, samplers generate mutually exclusive sets""" - N = 640 - num_draws = 321 - world_size = 8 - - for epoch in range(100): - c = Counter() - for rank in range(world_size): - sampler = DistributedRandomSampler( - N, rank=rank, world_size=world_size, num_draws=num_draws - ) - if shuffle: - sampler.shuffle(seed=epoch) - c.update(sampler) - - m = num_draws // world_size * world_size - self.assertEqual(c.total(), m) - self.assertEqual(len(c.keys()), m) - self.assertTrue(all(v == 1 for v in c.values())) - - -class TestDistributedSamplerWeighted(unittest.TestCase): - def test_weighted_sampling(self) -> None: - """Indices are drawn according to the weights""" - weights = [0.0, 1.0, 3.0, 5.0, 10.0] - N = len(weights) - - sampler = DistributedRandomSampler( - N, rank=0, world_size=1, num_draws=1_000_000, weights=weights - ) - - c = Counter(sampler) - distribution = [c[i] for i in range(N)] - - print(f"{weights=}") - print(f"{distribution=}") - - ref = np.asarray(weights) / np.sum(weights) - hyp = np.asarray(distribution) / np.sum(distribution) - - print(f"{ref=}") - print(f"{hyp=}") - - self.assertTrue(np.allclose(hyp, ref, atol=1e-3)) - - -class TestDistributedSamplerEmbedShuffle(unittest.TestCase): - def test_embed_shuffle(self) -> None: - """DistributedSampler is compatibile with embed_shuffle""" - N = 10 - weights = [1.0 for _ in range(N)] - - s0 = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) - s1 = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) - - s1 = embed_shuffle(s1) - - previous = [] - for i in range(100): - hyp = list(s1) - print(f"{hyp=}") - - s0.shuffle(i) - ref = list(s0) - print(f"{ref=}") - - self.assertEqual(hyp, ref) - self.assertNotEqual(hyp, previous) - previous = hyp - - -class TestDistributedSamplerIterateInSubprocess(unittest.TestCase): - @_ignore_fork_warning - def test_iterate_in_subprocess(self) -> None: - """Iterating in a subprocess generates identical result""" - N = 10 - weights = [1.0 for _ in range(N)] - - sampler = DistributedRandomSampler(N, rank=0, world_size=1, weights=weights) - sampler_sub = iterate_in_subprocess(partial(embed_shuffle, sampler)) - sampler = embed_shuffle(sampler) - - previous = [] - for _ in range(100): - hyp = list(sampler_sub) - print(f"{hyp=}") - ref = list(sampler) - print(f"{ref=}") - - self.assertEqual(hyp, ref) - self.assertNotEqual(hyp, previous) - previous = hyp diff --git a/tests/dataloader/sampler_test.py b/tests/dataloader/sampler_test.py new file mode 120000 index 000000000..224fd6612 --- /dev/null +++ b/tests/dataloader/sampler_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/sampler_test.py \ No newline at end of file diff --git a/tests/dataloader/source_test.py b/tests/dataloader/source_test.py deleted file mode 100644 index 61f31e8a5..000000000 --- a/tests/dataloader/source_test.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import unittest -from collections.abc import Iterable -from pathlib import Path -from tempfile import TemporaryDirectory - -from spdl.source.imagenet import ImageNet -from spdl.source.local_directory import LocalDirectory - - -def _make_files(paths: Iterable[Path]) -> None: - for path in paths: - path.parent.mkdir(parents=True, exist_ok=True) - path.touch() - - -class SourceTest(unittest.TestCase): - def test_LocalDirectory(self) -> None: - """LocalDirectory can traverse specified files""" - with TemporaryDirectory() as root_dir: - root_dir = Path(root_dir) - - targets = { - root_dir / "foo.txt", - root_dir / "dir" / "bar.txt", - root_dir / "dir" / "dir" / "bazz.txt", - } - others = { - root_dir / "foo.dat", - root_dir / "dir" / "bar.dat", - root_dir / "dir" / "dir" / "bazz.dat", - } - _make_files(targets | others) - - src = LocalDirectory(root=root_dir, pattern="**/*.txt") - - vals1 = list(src) - vals2 = list(src) - src.shuffle(seed=0) - vals3 = list(src) - - self.assertEqual(vals1, vals2) - self.assertNotEqual(vals1, vals3) - self.assertEqual(set(vals1), targets) - self.assertEqual(set(vals2), targets) - self.assertEqual(set(vals3), targets) - - def test_ImageNet(self) -> None: - """ImageNet returns image path and class ID""" - with TemporaryDirectory() as root_dir: - root_dir = Path(root_dir) - - vals = { - (root_dir / "val" / "n02110958" / "FOO.JPEG", 254), - (root_dir / "val" / "n02027492" / "FOO.JPEG", 140), - (root_dir / "val" / "n02071294" / "FOO.JPEG", 148), - (root_dir / "val" / "n02088632" / "FOO.JPEG", 164), - } - trains = { - (root_dir / "train" / "n02066245" / "FOO.JPEG", 147), - (root_dir / "train" / "n02277742" / "FOO.JPEG", 322), - (root_dir / "train" / "n02965783" / "FOO.JPEG", 475), - (root_dir / "train" / "n03240683" / "FOO.JPEG", 540), - } - _make_files([v for v, _ in vals]) - _make_files([v for v, _ in trains]) - - src = ImageNet(root=root_dir, split="val") - v1 = list(src) - src.shuffle(0) - v2 = list(src) - self.assertNotEqual(v1, v2) - self.assertEqual(set(v1), vals) - self.assertEqual(set(v2), vals) - - src = ImageNet(root=root_dir, split="train") - v1 = list(src) - src.shuffle(0) - v2 = list(src) - self.assertNotEqual(v1, v2) - self.assertEqual(set(v1), trains) - self.assertEqual(set(v2), trains) diff --git a/tests/dataloader/source_test.py b/tests/dataloader/source_test.py new file mode 120000 index 000000000..3c371d7ee --- /dev/null +++ b/tests/dataloader/source_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/source_test.py \ No newline at end of file diff --git a/tests/dataloader/source_utils_test.py b/tests/dataloader/source_utils_test.py deleted file mode 100644 index c62d51438..000000000 --- a/tests/dataloader/source_utils_test.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import unittest -from collections.abc import Iterator - -from spdl.source.utils import embed_shuffle - - -class IterableWithShuffle_: - def __init__(self, n: int) -> None: - self.vals: list[int] = list(range(n)) - self._seed: int | None = None - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - def shuffle(self, seed: int) -> None: - # rotate - self._seed = seed - self.vals = self.vals[1:] + self.vals[:1] - - -class SourceUtilsTest(unittest.TestCase): - def test_embed_shuffle(self) -> None: - """Iterable created by embed_shuffle calls shuffle automatically""" - - foo = IterableWithShuffle_(3) - self.assertIsNone(foo._seed) - iterable = embed_shuffle(foo) - self.assertEqual(list(iterable), [1, 2, 0]) - self.assertEqual(foo._seed, 0) - self.assertEqual(list(iterable), [2, 0, 1]) - self.assertEqual(foo._seed, 1) - self.assertEqual(list(iterable), [0, 1, 2]) - self.assertEqual(foo._seed, 2) - - def test_embed_shuffle_halt(self) -> None: - """The value is shuffled with different seed even after an iteration is halted.""" - - foo = IterableWithShuffle_(5) - iterable = embed_shuffle(foo) - - iterator = iter(iterable) - self.assertIsNone(foo._seed) - self.assertEqual(next(iterator), 1) - self.assertEqual(foo._seed, 0) - self.assertEqual(next(iterator), 2) - del iterator - - iterator = iter(iterable) - self.assertEqual(next(iterator), 2) - self.assertEqual(foo._seed, 1) - self.assertEqual(next(iterator), 3) - del iterator - - def test_embed_shuffle_shuffle_after(self) -> None: - """Iterable created by embed_shuffle calls shuffle automatically after iteration""" - - foo = IterableWithShuffle_(3) - iterable = embed_shuffle(foo, shuffle_last=True) - self.assertIsNone(foo._seed) - self.assertEqual(list(iterable), [0, 1, 2]) - self.assertEqual(foo._seed, 0) - self.assertEqual(list(iterable), [1, 2, 0]) - self.assertEqual(foo._seed, 1) - self.assertEqual(list(iterable), [2, 0, 1]) - self.assertEqual(foo._seed, 2) - - def test_embed_shuffle_shuffle_after_halt(self) -> None: - """The value is shuffled with different seed even after an iteration is halted.""" - - foo = IterableWithShuffle_(5) - iterable = embed_shuffle(foo, shuffle_last=True) - - iterator = iter(iterable) - self.assertEqual(next(iterator), 0) - self.assertEqual(next(iterator), 1) - self.assertIsNone(foo._seed) - del iterator - self.assertEqual(foo._seed, 0) - - iterator = iter(iterable) - self.assertEqual(next(iterator), 1) - self.assertEqual(next(iterator), 2) - del iterator - self.assertEqual(foo._seed, 1) diff --git a/tests/dataloader/source_utils_test.py b/tests/dataloader/source_utils_test.py new file mode 120000 index 000000000..98541729d --- /dev/null +++ b/tests/dataloader/source_utils_test.py @@ -0,0 +1 @@ +../../../src/spdl/dataloader/tests/source_utils_test.py \ No newline at end of file diff --git a/tests/io/__init__.py b/tests/io/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/io/core b/tests/io/core new file mode 120000 index 000000000..c7eb44eb0 --- /dev/null +++ b/tests/io/core @@ -0,0 +1 @@ +../../src/spdl/io/tests/core \ No newline at end of file diff --git a/tests/io/cuda b/tests/io/cuda new file mode 120000 index 000000000..8dfa45fc5 --- /dev/null +++ b/tests/io/cuda @@ -0,0 +1 @@ +../../src/spdl/io/tests/cuda \ No newline at end of file diff --git a/tests/io/fixture.py b/tests/io/fixture.py new file mode 120000 index 000000000..1de6b1fb3 --- /dev/null +++ b/tests/io/fixture.py @@ -0,0 +1 @@ +../../src/spdl/io/tests/fixture.py \ No newline at end of file diff --git a/tests/pipeline/__init__.py b/tests/pipeline/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/pipeline/aggregate_test.py b/tests/pipeline/aggregate_test.py deleted file mode 100644 index 614c81498..000000000 --- a/tests/pipeline/aggregate_test.py +++ /dev/null @@ -1,326 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import unittest -from collections.abc import Iterable -from typing import Any - -from spdl.pipeline import AsyncQueue, PipelineBuilder -from spdl.pipeline._components._aggregate import _aggregate -from spdl.pipeline._components._common import _EOF, StageInfo -from spdl.pipeline._components._pipe import _get_fail_counter -from spdl.pipeline.defs import Aggregator - -_TEST_INFO = StageInfo(pipeline_id=0, stage_id="0", stage_name="test") - - -def _put_aqueue(queue: AsyncQueue, vals: Iterable[object], *, eof: bool) -> None: - for val in vals: - queue.put_nowait(val) - if eof: - queue.put_nowait(_EOF) - - -def _flush_aqueue(queue: AsyncQueue) -> list[object]: - ret = [] - while not queue.empty(): - ret.append(queue.get_nowait()) - return ret - - -class _TrackingAggregator(Aggregator): - """Aggregator that collects items into batches of size N and records - accumulate call order for verifying drain behavior.""" - - def __init__(self, batch_size: int) -> None: - self.batch_size = batch_size - self.buffer: list[Any] = [] - self.accumulate_log: list[Any] = [] - - def accumulate(self, item: Any) -> list[Any] | None: - self.accumulate_log.append(item) - self.buffer.append(item) - if len(self.buffer) >= self.batch_size: - result = self.buffer - self.buffer = [] - return result - return None - - def flush(self) -> list[Any] | None: - if self.buffer: - result = self.buffer - self.buffer = [] - return result - return None - - -class AggregatePipeBulkDrainTest(unittest.IsolatedAsyncioTestCase): - async def test_stops_draining_on_emit(self) -> None: - """After the aggregator emits, bulk draining stops and remaining - items stay in the input queue for the next drain cycle.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - agg = _TrackingAggregator(batch_size=3) - # Pre-fill 7 items + EOF. The aggregator emits after every 3 items. - # With stop-on-emit, each drain cycle processes exactly 3 items - # (emit batch), then stops. The remaining items stay in the input - # queue for the next cycle. The last item is flushed at EOF. - _put_aqueue(input_queue, list(range(7)), eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=True, - ) - - results = _flush_aqueue(output_queue) - # Two full batches of 3 + flush of remainder [6] + EOF - self.assertEqual(results, [[0, 1, 2], [3, 4, 5], [6], _EOF]) - # All 7 items were accumulated - self.assertEqual(agg.accumulate_log, [0, 1, 2, 3, 4, 5, 6]) - - async def test_drains_without_blocking_when_no_emit(self) -> None: - """When the aggregator doesn't emit, items are drained from the - queue without blocking (via get_nowait).""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - # batch_size=10 means 5 items won't trigger an emit - agg = _TrackingAggregator(batch_size=10) - - # Pre-fill 5 items + EOF. None will trigger an emit during - # accumulate, but flush will emit the remaining buffer. - _put_aqueue(input_queue, list(range(5)), eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=True, - ) - - results = _flush_aqueue(output_queue) - # flush emits the remaining 5 items - self.assertEqual(results, [[0, 1, 2, 3, 4], _EOF]) - self.assertEqual(agg.accumulate_log, [0, 1, 2, 3, 4]) - - async def test_drop_last(self) -> None: - """When op_requires_eof=False, EOF stops processing and flush - is not called, dropping incomplete batches.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - agg = _TrackingAggregator(batch_size=3) - - # 5 items: batch at 3, then 2 remaining are dropped - _put_aqueue(input_queue, list(range(5)), eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=False, - ) - - results = _flush_aqueue(output_queue) - # Only the complete batch + EOF from queue_stage_hook - self.assertEqual(results, [[0, 1, 2], _EOF]) - - async def test_exception_propagates(self) -> None: - """Exceptions from the aggregator propagate instead of being - silently swallowed.""" - - class FailingAggregator(Aggregator): - def accumulate(self, item: Any) -> None: - raise ValueError("aggregation failed") - - def flush(self) -> None: - return None - - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - _put_aqueue(input_queue, [1], eof=True) - - with self.assertRaises(ValueError, msg="aggregation failed"): - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - FailingAggregator(), - _get_fail_counter()(), - [], - op_requires_eof=False, - ) - - async def test_single_item_no_emit(self) -> None: - """A single item that doesn't trigger emit is flushed at EOF.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - agg = _TrackingAggregator(batch_size=5) - - _put_aqueue(input_queue, [42], eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=True, - ) - - results = _flush_aqueue(output_queue) - self.assertEqual(results, [[42], _EOF]) - - async def test_eof_with_empty_flush(self) -> None: - """When items evenly divide batch_size, flush() returns None at EOF. - The function must still return instead of blocking.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - agg = _TrackingAggregator(batch_size=3) - - # 6 items evenly divides batch_size=3, so flush() returns None - _put_aqueue(input_queue, list(range(6)), eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=True, - ) - - results = _flush_aqueue(output_queue) - self.assertEqual(results, [[0, 1, 2], [3, 4, 5], _EOF]) - - async def test_emit_on_every_item(self) -> None: - """When batch_size=1, every item triggers an emit and each - drain cycle processes exactly one item.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - agg = _TrackingAggregator(batch_size=1) - - _put_aqueue(input_queue, [10, 20, 30], eof=True) - - await _aggregate( - _TEST_INFO, - input_queue, - output_queue, - agg, - _get_fail_counter()(), - [], - op_requires_eof=False, - ) - - results = _flush_aqueue(output_queue) - self.assertEqual(results, [[10], [20], [30], _EOF]) - - -class AggregatePipeEndToEndTest(unittest.TestCase): - def test_aggregate_bulk_drain_correctness(self) -> None: - """End-to-end pipeline test: aggregate produces correct batches.""" - src = list(range(10)) - - pipeline = ( - PipelineBuilder() - .add_source(src) - .aggregate(3) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(results, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) - - def test_aggregate_custom_op_bulk_drain(self) -> None: - """End-to-end: custom aggregator with bulk drain produces correct output.""" - - class SizeAggregator(Aggregator): - def __init__(self, threshold: int) -> None: - self.threshold = threshold - self.buffer: list[str] = [] - self.total: int = 0 - - def accumulate(self, item: str) -> str | None: - self.buffer.append(item) - self.total += len(item) - if self.total >= self.threshold: - result = "".join(self.buffer) - self.buffer = [] - self.total = 0 - return result - return None - - def flush(self) -> str | None: - if self.buffer: - result = "".join(self.buffer) - self.buffer = [] - self.total = 0 - return result - return None - - src = ["a", "bb", "ccc", "dddd", "e", "ff", "ggg", "h"] - - pipeline = ( - PipelineBuilder() - .add_source(src) - .aggregate(SizeAggregator(threshold=10)) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(results, ["abbcccdddd", "effgggh"]) diff --git a/tests/pipeline/aggregate_test.py b/tests/pipeline/aggregate_test.py new file mode 120000 index 000000000..6bd0fb488 --- /dev/null +++ b/tests/pipeline/aggregate_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/aggregate_test.py \ No newline at end of file diff --git a/tests/pipeline/background_task_test.py b/tests/pipeline/background_task_test.py deleted file mode 100644 index da6582813..000000000 --- a/tests/pipeline/background_task_test.py +++ /dev/null @@ -1,278 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import time -import unittest - -from spdl.pipeline import BackgroundTask, build_pipeline -from spdl.pipeline.config import ( - get_default_background_tasks, - set_default_background_tasks, -) -from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig - - -def _simple_cfg() -> PipelineConfig[int]: - return PipelineConfig( - src=SourceConfig(range(5)), - pipes=[Pipe(lambda x: x)], - sink=SinkConfig(3), - ) - - -def _slow_cfg(delay: float = 0.5) -> PipelineConfig[int]: - """Pipeline that takes a while to complete, giving background tasks time to run.""" - - def slow_op(x: int) -> int: - time.sleep(delay) - return x - - return PipelineConfig( - src=SourceConfig(range(3)), - pipes=[Pipe(slow_op)], - sink=SinkConfig(3), - ) - - -class _TrackingTask(BackgroundTask): - """Background task that tracks whether it started and was cancelled.""" - - def __init__(self) -> None: - self.started = False - self.cancelled = False - - async def run(self) -> None: - self.started = True - try: - while True: - await asyncio.sleep(0.01) - except asyncio.CancelledError: - self.cancelled = True - raise - - -class _CountingTask(BackgroundTask): - """Background task that counts iterations.""" - - def __init__(self) -> None: - self.count = 0 - - async def run(self) -> None: - try: - while True: - self.count += 1 - await asyncio.sleep(0.01) - except asyncio.CancelledError: - pass - - -class _FailingTask(BackgroundTask): - async def run(self) -> None: - raise RuntimeError("bg task error") - - -class _ShortTask(BackgroundTask): - """Background task that completes quickly.""" - - def __init__(self) -> None: - self.completed = False - - async def run(self) -> None: - await asyncio.sleep(0.01) - self.completed = True - - -class BackgroundTaskTest(unittest.TestCase): - def setUp(self) -> None: - self._saved = get_default_background_tasks() - set_default_background_tasks(None) - - def tearDown(self) -> None: - set_default_background_tasks(self._saved) - - def test_background_task_runs_and_is_cancelled(self) -> None: - """Background tasks run alongside pipeline and get cancelled on completion.""" - task_instance: _TrackingTask = _TrackingTask() - - def factory() -> BackgroundTask: - return task_instance - - pipeline = build_pipeline( - _simple_cfg(), num_threads=1, background_tasks=[factory] - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertTrue(task_instance.started, "Background task should have started") - self.assertTrue( - task_instance.cancelled, "Background task should have been cancelled" - ) - - def test_background_task_error_does_not_crash_pipeline(self) -> None: - """Background task errors are logged but don't fail the pipeline.""" - pipeline = build_pipeline( - _simple_cfg(), num_threads=1, background_tasks=[_FailingTask] - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - - def test_multiple_background_tasks(self) -> None: - """Multiple background tasks can run concurrently.""" - task_0 = _CountingTask() - task_1 = _CountingTask() - - pipeline = build_pipeline( - _simple_cfg(), - num_threads=1, - background_tasks=[lambda: task_0, lambda: task_1], - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertGreater(task_0.count, 0, "First background task should have run") - self.assertGreater(task_1.count, 0, "Second background task should have run") - - def test_no_background_tasks(self) -> None: - """Pipeline works normally when no background tasks are provided.""" - pipeline = build_pipeline(_simple_cfg(), num_threads=1) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - - def test_empty_background_tasks_list(self) -> None: - """Pipeline works normally with an empty background tasks list.""" - pipeline = build_pipeline(_simple_cfg(), num_threads=1, background_tasks=[]) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - - def test_background_task_completes_before_pipeline(self) -> None: - """A background task that finishes early doesn't affect the pipeline.""" - task_instance = _ShortTask() - - pipeline = build_pipeline( - _slow_cfg(0.1), - num_threads=1, - background_tasks=[lambda: task_instance], - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(sorted(items), [0, 1, 2]) - self.assertTrue( - task_instance.completed, "Background task should have completed" - ) - - def test_background_task_mixed_success_and_failure(self) -> None: - """One failing and one succeeding background task — pipeline still works.""" - good_task = _TrackingTask() - - pipeline = build_pipeline( - _simple_cfg(), - num_threads=1, - background_tasks=[lambda: good_task, _FailingTask], - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertTrue(good_task.started, "Good background task should have run") - - def test_class_as_factory(self) -> None: - """A BackgroundTask class itself can be used as a factory.""" - - class MyTask(BackgroundTask): - ran = False - - async def run(self) -> None: - MyTask.ran = True - try: - while True: - await asyncio.sleep(0.01) - except asyncio.CancelledError: - pass - - MyTask.ran = False - pipeline = build_pipeline( - _simple_cfg(), num_threads=1, background_tasks=[MyTask] - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertTrue(MyTask.ran, "Task class used as factory should have run") - - -class DefaultBackgroundTasksTest(unittest.TestCase): - def setUp(self) -> None: - self._saved = get_default_background_tasks() - set_default_background_tasks(None) - - def tearDown(self) -> None: - set_default_background_tasks(self._saved) - - def test_get_set_default_background_tasks(self) -> None: - """set/get_default_background_tasks round-trips correctly.""" - self.assertIsNone(get_default_background_tasks()) - - set_default_background_tasks([_TrackingTask]) - result = get_default_background_tasks() - self.assertIsNotNone(result) - self.assertEqual(len(result), 1) # pyre-ignore[6] - self.assertIs(result[0], _TrackingTask) # pyre-ignore[16] - - set_default_background_tasks(None) - self.assertIsNone(get_default_background_tasks()) - - def test_default_background_tasks_run_automatically(self) -> None: - """Default background tasks are started without explicit parameter.""" - task_instance = _TrackingTask() - set_default_background_tasks([lambda: task_instance]) - - pipeline = build_pipeline(_simple_cfg(), num_threads=1) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertTrue( - task_instance.started, "Default background task should have run" - ) - - def test_default_and_per_pipeline_tasks_merged(self) -> None: - """Both default and per-pipeline background tasks run.""" - default_task = _TrackingTask() - custom_task = _TrackingTask() - - set_default_background_tasks([lambda: default_task]) - - pipeline = build_pipeline( - _simple_cfg(), num_threads=1, background_tasks=[lambda: custom_task] - ) - - with pipeline.auto_stop(): - items = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(sorted(items), [0, 1, 2, 3, 4]) - self.assertTrue(default_task.started, "Default background task should have run") - self.assertTrue(custom_task.started, "Custom background task should have run") diff --git a/tests/pipeline/background_task_test.py b/tests/pipeline/background_task_test.py new file mode 120000 index 000000000..7f75d7db3 --- /dev/null +++ b/tests/pipeline/background_task_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/background_task_test.py \ No newline at end of file diff --git a/tests/pipeline/build_pipeline_test.py b/tests/pipeline/build_pipeline_test.py deleted file mode 100644 index 7dd52da5e..000000000 --- a/tests/pipeline/build_pipeline_test.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import os -import unittest -import warnings -from unittest.mock import patch - -from spdl.pipeline import build_pipeline -from spdl.pipeline._profile import _ProfilePipeline -from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig - -# pyre-strict - - -class TestBuildPipeline(unittest.TestCase): - """Test class for build_pipeline functionality.""" - - def test_build_pipeline_diagnostic_mode(self) -> None: - """Test that when SPDL_PIPELINE_DIAGNOSTIC_MODE=1, build_pipeline - calls _build_pipeline_diagnostic_mode and returns _ProfilePipeline. - """ - - def simple_op(i: int) -> int: - return i * 2 - - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message="coroutine .* was never awaited", - category=RuntimeWarning, - ) - with patch.dict("os.environ", {"SPDL_PIPELINE_DIAGNOSTIC_MODE": "1"}): - pipeline = build_pipeline(cfg, num_threads=2) - self.assertIsInstance(pipeline, _ProfilePipeline) - - with patch.dict("os.environ", {}, clear=False): - os.environ.pop("SPDL_PIPELINE_DIAGNOSTIC_MODE", None) - - pipeline = build_pipeline(cfg, num_threads=2) - self.assertNotIsInstance(pipeline, _ProfilePipeline) diff --git a/tests/pipeline/build_pipeline_test.py b/tests/pipeline/build_pipeline_test.py new file mode 120000 index 000000000..15ee0df12 --- /dev/null +++ b/tests/pipeline/build_pipeline_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/build_pipeline_test.py \ No newline at end of file diff --git a/tests/pipeline/compact_log_test.py b/tests/pipeline/compact_log_test.py deleted file mode 100644 index 031a9405d..000000000 --- a/tests/pipeline/compact_log_test.py +++ /dev/null @@ -1,279 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import os -import unittest - -from spdl.pipeline._common._misc import _get_compact_log, _set_compact_log, create_task -from spdl.pipeline.config import set_compact_log - - -class DummyException(Exception): - """Test exception for simulating task failures.""" - - pass - - -class CompactLogTest(unittest.TestCase): - """Tests for compact logging mode functionality.""" - - def setUp(self) -> None: - """Reset the global compact log setting before each test.""" - # Reset to None to ensure clean state - _set_compact_log(None) - # Clear any environment variable that might be set - if "SPDL_PIPELINE_COMPACT_LOG" in os.environ: - del os.environ["SPDL_PIPELINE_COMPACT_LOG"] - - def tearDown(self) -> None: - """Clean up after each test.""" - # Reset to None - _set_compact_log(None) - # Clear environment variable - if "SPDL_PIPELINE_COMPACT_LOG" in os.environ: - del os.environ["SPDL_PIPELINE_COMPACT_LOG"] - - def test_get_compact_log_defaults_to_false_when_env_not_set(self) -> None: - """Test that _get_compact_log returns False when environment variable is not set.""" - # Setup: Environment variable is not set (cleared in setUp) - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should default to False - self.assertFalse(result) - - def test_get_compact_log_returns_true_when_env_is_1(self) -> None: - """Test that _get_compact_log returns True when environment variable is '1'.""" - # Setup: Set environment variable to '1' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "1" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_get_compact_log_returns_true_when_env_is_true(self) -> None: - """Test that _get_compact_log returns True when environment variable is 'true'.""" - # Setup: Set environment variable to 'true' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_get_compact_log_returns_true_when_env_is_TRUE(self) -> None: - """Test that _get_compact_log returns True when environment variable is 'TRUE'.""" - # Setup: Set environment variable to 'TRUE' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "TRUE" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_get_compact_log_returns_true_when_env_is_yes(self) -> None: - """Test that _get_compact_log returns True when environment variable is 'yes'.""" - # Setup: Set environment variable to 'yes' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "yes" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_get_compact_log_returns_false_when_env_is_0(self) -> None: - """Test that _get_compact_log returns False when environment variable is '0'.""" - # Setup: Set environment variable to '0' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "0" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return False - self.assertFalse(result) - - def test_get_compact_log_returns_false_when_env_is_false(self) -> None: - """Test that _get_compact_log returns False when environment variable is 'false'.""" - # Setup: Set environment variable to 'false' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return False - self.assertFalse(result) - - def test_get_compact_log_caches_result_after_first_call(self) -> None: - """Test that _get_compact_log caches the result and doesn't re-check env var.""" - # Setup: Set environment variable to 'true' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" - - # Execute: Get compact log setting twice - result1 = _get_compact_log() - # Change environment variable - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" - result2 = _get_compact_log() - - # Assert: Both should return True because value is cached - self.assertTrue(result1) - self.assertTrue(result2) - - def test_set_compact_log_to_true(self) -> None: - """Test that _set_compact_log can set the value to True.""" - # Setup: Set to True - _set_compact_log(True) - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_set_compact_log_to_false(self) -> None: - """Test that _set_compact_log can set the value to False.""" - # Setup: Set to False - _set_compact_log(False) - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return False - self.assertFalse(result) - - def test_set_compact_log_to_none_resets_to_env_check(self) -> None: - """Test that setting to None causes re-check of environment variable.""" - # Setup: Set to True first - _set_compact_log(True) - self.assertTrue(_get_compact_log()) - - # Reset to None - _set_compact_log(None) - # Set environment variable - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "false" - - # Execute: Get compact log setting after reset - result = _get_compact_log() - - # Assert: Should return False from environment variable - self.assertFalse(result) - - def test_set_compact_log_overrides_env_var(self) -> None: - """Test that programmatically setting the value overrides environment variable.""" - # Setup: Set environment variable to 'true' - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "true" - # Override with False - _set_compact_log(False) - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return False (overridden value, not env var) - self.assertFalse(result) - - def test_config_module_exposes_set_compact_log(self) -> None: - """Test that set_compact_log is exposed in the config module.""" - # Setup: Set through config module - set_compact_log(True) - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result) - - def test_create_task_properly_calls_compact_log_setting(self) -> None: - """Test that create_task respects the compact log setting.""" - - async def failing_coro() -> None: - raise DummyException("test error") - - async def run() -> None: - # Test with compact=False (default) - _set_compact_log(False) - task1 = create_task(failing_coro(), name="task1") - await asyncio.sleep(0) - try: - await task1 - except DummyException: - pass - # Task completed, no assertion needed - just verify no exceptions - - # Test with compact=True - _set_compact_log(True) - task2 = create_task(failing_coro(), name="task2") - await asyncio.sleep(0) - try: - await task2 - except DummyException: - pass - # Task completed, no assertion needed - just verify no exceptions - - asyncio.run(run()) - - def test_create_task_uses_get_compact_log(self) -> None: - """Test that create_task gets the compact setting from _get_compact_log.""" - - async def failing_coro() -> None: - raise DummyException("test error") - - async def run() -> None: - # Setup: Set compact log via environment variable - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = "1" - _set_compact_log(None) # Reset to force env var check - - # Execute: Create task - this should use compact mode from env var - task = create_task(failing_coro(), name="test_task") - await asyncio.sleep(0) - try: - await task - except DummyException: - pass - - # Assert: Verify the getter returns True (from env var) - self.assertTrue(_get_compact_log()) - - asyncio.run(run()) - - def test_get_compact_log_with_various_truthy_env_values(self) -> None: - """Test that _get_compact_log handles various truthy environment values.""" - truthy_values = ["1", "true", "TRUE", "on", "ON", "yes", "YES"] - - for value in truthy_values: - with self.subTest(value=value): - # Setup: Reset and set env var - _set_compact_log(None) - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = value - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return True - self.assertTrue(result, f"Expected True for value '{value}'") - - def test_get_compact_log_with_various_falsy_env_values(self) -> None: - """Test that _get_compact_log handles various falsy environment values.""" - falsy_values = ["0", "false", "FALSE", "off", "OFF", "no", "NO"] - - for value in falsy_values: - with self.subTest(value=value): - # Setup: Reset and set env var - _set_compact_log(None) - os.environ["SPDL_PIPELINE_COMPACT_LOG"] = value - - # Execute: Get compact log setting - result = _get_compact_log() - - # Assert: Should return False - self.assertFalse(result, f"Expected False for value '{value}'") diff --git a/tests/pipeline/compact_log_test.py b/tests/pipeline/compact_log_test.py new file mode 120000 index 000000000..9962540f8 --- /dev/null +++ b/tests/pipeline/compact_log_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/compact_log_test.py \ No newline at end of file diff --git a/tests/pipeline/config_test.py b/tests/pipeline/config_test.py deleted file mode 100644 index b64274b8e..000000000 --- a/tests/pipeline/config_test.py +++ /dev/null @@ -1,210 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import unittest -from collections.abc import Iterator -from contextlib import contextmanager - -from spdl.pipeline import ( - AsyncQueue, - ProfileHook, - ProfileResult, - StatsQueue, - TaskHook, - TaskStatsHook, -) -from spdl.pipeline.config import ( - get_default_hook_class, - get_default_profile_callback, - get_default_profile_hook, - get_default_queue_class, - set_default_hook_class, - set_default_profile_callback, - set_default_profile_hook, - set_default_queue_class, -) - - -def reset() -> None: - set_default_hook_class() - set_default_queue_class() - set_default_profile_hook() - set_default_profile_callback() - - -class ConfigTest(unittest.TestCase): - """Test the configuration setter/getter functions.""" - - def setUp(self) -> None: - """Reset all configuration state before each test.""" - super().setUp() - reset() - - def tearDown(self) -> None: - """Reset all configuration state after each test.""" - reset() - super().tearDown() - - def test_hook_class_default_is_none(self) -> None: - """Test that default hook class is None when not configured.""" - result = get_default_hook_class() - self.assertIs(result, TaskStatsHook) - - def test_hook_class_can_be_set_and_retrieved(self) -> None: - """Test that hook class can be set and retrieved correctly.""" - - class CustomHook(TaskHook): - pass - - set_default_hook_class(CustomHook) - result = get_default_hook_class() - - self.assertIs(result, CustomHook) - - def test_hook_class_via_config_module(self) -> None: - """Test that hook class can be accessed via _config module.""" - - class CustomHook(TaskHook): - pass - - set_default_hook_class(CustomHook) - result = get_default_hook_class() - - self.assertIs(result, CustomHook) - - def test_queue_class_default_is_none(self) -> None: - """Test that default queue class is None when not configured.""" - result = get_default_queue_class() - self.assertIs(result, StatsQueue) - - def test_queue_class_can_be_set_and_retrieved(self) -> None: - """Test that queue class can be set and retrieved correctly.""" - - class CustomQueue(StatsQueue): - pass - - set_default_queue_class(CustomQueue) - result = get_default_queue_class() - self.assertIs(result, CustomQueue) - - def test_queue_class_via_config_module(self) -> None: - """Test that queue class can be accessed via _config module.""" - - class CustomQueue(StatsQueue): - pass - - set_default_queue_class(CustomQueue) - result = get_default_queue_class() - self.assertIs(result, CustomQueue) - - def test_profile_hook_default_is_none(self) -> None: - """Test that default profile hook is None when not configured.""" - result = get_default_profile_hook() - self.assertIsNone(result) - - def test_profile_hook_can_be_set_and_retrieved(self) -> None: - """Test that profile hook can be set and retrieved correctly.""" - - class MockProfileHook(ProfileHook): - @contextmanager - def stage_profile_hook( - self, - stage: str, # noqa: ARG002 - concurrency: int, # noqa: ARG002 - ) -> Iterator[None]: - yield - - @contextmanager - def pipeline_profile_hook(self) -> Iterator[None]: - yield - - hook_instance = MockProfileHook() - set_default_profile_hook(hook_instance) - result = get_default_profile_hook() - self.assertIs(result, hook_instance) - - def test_profile_hook_via_config_module(self) -> None: - """Test that profile hook can be accessed via _config module.""" - - class MockProfileHook(ProfileHook): - @contextmanager - def stage_profile_hook( - self, - stage: str, # noqa: ARG002 - concurrency: int, # noqa: ARG002 - ) -> Iterator[None]: - yield - - @contextmanager - def pipeline_profile_hook(self) -> Iterator[None]: - yield - - hook_instance = MockProfileHook() - set_default_profile_hook(hook_instance) - result = get_default_profile_hook() - self.assertIs(result, hook_instance) - - def test_profile_callback_default_is_none(self) -> None: - """Test that default profile callback is None when not configured.""" - result = get_default_profile_callback() - self.assertIsNone(result) - - def test_profile_callback_can_be_set_and_retrieved(self) -> None: - """Test that profile callback can be set and retrieved correctly.""" - - def mock_callback(_: ProfileResult) -> None: - pass - - set_default_profile_callback(mock_callback) - result = get_default_profile_callback() - self.assertIs(result, mock_callback) - - def test_profile_callback_via_config_module(self) -> None: - """Test that profile callback can be accessed via _config module.""" - - def mock_callback(_: object) -> None: - pass - - set_default_profile_callback(mock_callback) - result = get_default_profile_callback() - self.assertIs(result, mock_callback) - - def test_multiple_configurations_independent(self) -> None: - """Test that different configuration settings are independent.""" - - class CustomHook(TaskHook): - pass - - class CustomQueue(AsyncQueue): - pass - - def custom_callback(_: object) -> None: - pass - - set_default_hook_class(CustomHook) - set_default_queue_class(CustomQueue) - set_default_profile_callback(custom_callback) - - self.assertIs(get_default_hook_class(), CustomHook) - self.assertIs(get_default_queue_class(), CustomQueue) - self.assertIs(get_default_profile_callback(), custom_callback) - - def test_configuration_can_be_updated(self) -> None: - """Test that configuration can be updated to new values.""" - - class FirstHook(TaskHook): - pass - - class SecondHook(TaskHook): - pass - - # Set first value, then update to second value - set_default_hook_class(FirstHook) - set_default_hook_class(SecondHook) - result = get_default_hook_class() - self.assertIs(result, SecondHook) diff --git a/tests/pipeline/config_test.py b/tests/pipeline/config_test.py new file mode 120000 index 000000000..4e90094d2 --- /dev/null +++ b/tests/pipeline/config_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/config_test.py \ No newline at end of file diff --git a/tests/pipeline/continuous_pipeline_test.py b/tests/pipeline/continuous_pipeline_test.py deleted file mode 100644 index 2affb6466..000000000 --- a/tests/pipeline/continuous_pipeline_test.py +++ /dev/null @@ -1,623 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -import functools -import os -import sys -import threading -import time -import unittest -import warnings -import weakref -from collections.abc import Iterator - -from spdl.pipeline import ( - build_pipeline, - PipelineBuilder, - PipelineFailure, - run_pipeline_in_subinterpreter, - run_pipeline_in_subprocess, -) -from spdl.pipeline.defs import Merge, PipelineConfig, SinkConfig - - -def _ignore_fork_warning(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - return wrapper - - -class SourceIterable: - """Reusable iterable that yields range(n) on each iteration.""" - - def __init__(self, n: int) -> None: - self.n = n - - def __iter__(self) -> Iterator[int]: - yield from range(self.n) - - -class TestContinuousPipelineBasic(unittest.TestCase): - def test_continuous_multi_epoch(self) -> None: - """Pipeline with continuous=True can be iterated multiple times.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") - - def test_continuous_single_epoch(self) -> None: - """Continuous pipeline works for a single epoch.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(3), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 1, 2]) - - def test_continuous_empty_epoch(self) -> None: - """Continuous pipeline handles empty iterations.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(0), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [], f"epoch {epoch}") - - def test_continuous_single_item_epoch(self) -> None: - """Continuous pipeline works with single-item epochs.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(1), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0], f"epoch {epoch}") - - -class TestContinuousPipelinePipe(unittest.TestCase): - def test_continuous_pipe_concurrent(self) -> None: - """Pipe with concurrency > 1 handles epoch boundaries correctly.""" - - def double(x: int) -> int: - return x * 2 - - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .pipe(double, concurrency=4) - .add_sink(buffer_size=3) - .build(num_threads=4) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = sorted(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 2, 4, 6, 8], f"epoch {epoch}") - - def test_continuous_pipe_chain(self) -> None: - """EPOCH_END propagates through multiple pipe stages.""" - - def add_one(x: int) -> int: - return x + 1 - - def double(x: int) -> int: - return x * 2 - - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(3), continuous=True) - .pipe(add_one, concurrency=1) - .pipe(double, concurrency=1) - .add_sink(buffer_size=3) - .build(num_threads=2) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [2, 4, 6], f"epoch {epoch}") - - -class TestContinuousPipelineAggregate(unittest.TestCase): - def test_continuous_aggregate_exact_batch(self) -> None: - """Aggregate with exact batch size across epochs.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(6), continuous=True) - .aggregate(3) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [[0, 1, 2], [3, 4, 5]], f"epoch {epoch}") - - def test_continuous_aggregate_partial_batch_flushed(self) -> None: - """Partial batch at epoch end is flushed by the default aggregator.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .aggregate(3) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - # 5 items, batch_size=3: full batch [0,1,2] + partial batch [3,4] - self.assertEqual(result, [[0, 1, 2], [3, 4]], f"epoch {epoch}") - - def test_continuous_aggregate_drop_last(self) -> None: - """With drop_last=True, partial batch at epoch end is discarded.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .aggregate(3, drop_last=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - # 5 items, batch_size=3, drop_last: only [0,1,2], partial [3,4] dropped - self.assertEqual(result, [[0, 1, 2]], f"epoch {epoch}") - - -class TestContinuousPipelineShutdown(unittest.TestCase): - def test_continuous_auto_stop_mid_epoch(self) -> None: - """auto_stop() exits cleanly even if mid-epoch.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(100), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - it = pipeline.get_iterator(timeout=5) - # Consume only a few items - self.assertEqual(next(it), 0) - self.assertEqual(next(it), 1) - # auto_stop exits here — must not hang - - def test_continuous_stop_between_epochs(self) -> None: - """stop() between epochs works cleanly.""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(3), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 1, 2]) - # auto_stop exits here after one epoch — must not hang - - def test_continuous_get_iterator_reuse(self) -> None: - """get_iterator() can be called multiple times within auto_stop().""" - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(3), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - r1 = list(pipeline.get_iterator(timeout=5)) - r2 = list(pipeline.get_iterator(timeout=5)) - r3 = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(r1, [0, 1, 2]) - self.assertEqual(r2, [0, 1, 2]) - self.assertEqual(r3, [0, 1, 2]) - - def test_continuous_stop_with_pipe_stage(self) -> None: - """Continuous pipeline with pipe stage can be stopped after epochs.""" - - def double(x: int) -> int: - return x * 2 - - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .pipe(double, concurrency=2) - .add_sink(buffer_size=3) - .build(num_threads=2) - ) - - with pipeline.auto_stop(): - r1 = sorted(pipeline.get_iterator(timeout=5)) - self.assertEqual(r1, [0, 2, 4, 6, 8]) - # auto_stop exits — pipeline has items buffered for next epoch - # stop() must drain and shut down cleanly - - @_ignore_fork_warning - def test_continuous_stop_with_subprocess_source(self) -> None: - """Continuous pipeline reading from subprocess can be stopped.""" - backend = ( - PipelineBuilder().add_source(SourceIterable(5)).add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - pipeline = ( - PipelineBuilder() - .add_source(source, continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - r1 = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(r1, [0, 1, 2, 3, 4]) - # auto_stop exits — must not hang on subprocess IPC - - @_ignore_fork_warning - def test_continuous_pipeline_in_subprocess_multi_epoch(self) -> None: - """A continuous pipeline running inside a subprocess supports - multi-epoch iteration without recreating the subprocess.""" - backend = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - # Iterate 3 epochs from the parent — subprocess is reused - for epoch in range(3): - result = list(source) - self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") - - @_ignore_fork_warning - def test_continuous_pipeline_in_subprocess_stop_mid_epoch(self) -> None: - """Subprocess with continuous pipeline can be abandoned mid-epoch.""" - backend = ( - PipelineBuilder() - .add_source(SourceIterable(100), continuous=True) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - it = iter(source) - self.assertEqual(next(it), 0) - self.assertEqual(next(it), 1) - # Abandon — must not hang - del it - del source - - @_ignore_fork_warning - def test_continuous_frontend_backend_multi_epoch(self) -> None: - """Frontend continuous pipeline on top of subprocess backend - supports multi-epoch iteration.""" - backend = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - pipeline = ( - PipelineBuilder() - .add_source(source, continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for epoch in range(3): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") - - @_ignore_fork_warning - def test_continuous_frontend_backend_stop_mid_epoch(self) -> None: - """Frontend continuous pipeline on top of subprocess backend - can be stopped mid-epoch without hanging.""" - backend = ( - PipelineBuilder() - .add_source(SourceIterable(100), continuous=True) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - pipeline = ( - PipelineBuilder() - .add_source(source, continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - it = pipeline.get_iterator(timeout=5) - self.assertEqual(next(it), 0) - self.assertEqual(next(it), 1) - # auto_stop exits mid-epoch — must not hang - - @_ignore_fork_warning - def test_continuous_frontend_backend_finalizer_shutdown(self) -> None: - """Frontend+backend pipeline cleaned up via weakref.finalize. - - Simulates the _SPDLDataLoader pattern: pipeline is started, iterated, - then the wrapper goes out of scope. The finalizer calls stop(timeout=10). - Must not hang. - """ - - backend = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - pipeline = ( - PipelineBuilder() - .add_source(source, continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - pipeline.start() - finalizer = weakref.finalize(pipeline, lambda p: p.stop(timeout=10), pipeline) - - # Iterate 2 epochs - for epoch in range(2): - result = list(pipeline.get_iterator(timeout=5)) - self.assertEqual(result, [0, 1, 2, 3, 4], f"epoch {epoch}") - - t0 = time.monotonic() - # Trigger finalizer (simulates going out of scope) - finalizer() - elapsed = time.monotonic() - t0 - self.assertLess(elapsed, 15, f"finalizer took {elapsed:.1f}s — likely hung") - - -class CustomError(ValueError): - pass - - -class TestContinuousPipelineErrors(unittest.TestCase): - def test_continuous_source_failure(self) -> None: - """Source raising mid-epoch propagates error.""" - - class FailingSource: - def __iter__(self) -> Iterator[int]: - yield 0 - raise CustomError("source failed") - - pipeline = ( - PipelineBuilder() - .add_source(FailingSource(), continuous=True) - .add_sink(buffer_size=3) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=5)) - - def test_continuous_pipe_failure(self) -> None: - """Pipe function raising propagates error.""" - - def failing_fn(x: int) -> int: - if x == 2: - raise CustomError("pipe failed") - return x - - pipeline = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .pipe(failing_fn, concurrency=1) - .add_sink(buffer_size=3) - .build(num_threads=1, max_failures=0) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=5)) - - -class TestContinuousPipelineValidation(unittest.TestCase): - def test_mixed_continuous_mode_rejected(self) -> None: - """Mixing continuous and non-continuous sources raises ValueError.""" - plc_continuous = ( - PipelineBuilder() - .add_source(SourceIterable(3), continuous=True) - .add_sink() - .get_config() - ) - plc_normal = ( - PipelineBuilder().add_source(SourceIterable(3)).add_sink().get_config() - ) - - merged_config = PipelineConfig( - src=Merge([plc_continuous, plc_normal]), - pipes=[], - sink=SinkConfig(buffer_size=3), - ) - - with self.assertRaisesRegex(ValueError, "Mixed continuous mode"): - build_pipeline(merged_config, num_threads=1) - - -class _RecordIDs: - """Pipe function that records the thread ID and process ID.""" - - def __init__(self) -> None: - self.thread_ids: list[int] = [] - self.process_ids: list[int] = [] - - def __call__(self, x: int) -> int: - self.thread_ids.append(threading.get_ident()) - self.process_ids.append(os.getpid()) - return x - - -class TestContinuousSubprocessPipelineReuse(unittest.TestCase): - @_ignore_fork_warning - def test_subprocess_pipeline_reused_across_epochs(self) -> None: - """Thread and process IDs stay the same across epochs, proving - the pipeline is reused rather than rebuilt.""" - recorder = _RecordIDs() - - backend = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .pipe(recorder, concurrency=1) - .add_sink(buffer_size=3) - ) - source = run_pipeline_in_subprocess( - backend.get_config(), - num_threads=1, - timeout=10, - ) - - all_thread_ids: list[set[int]] = [] - all_process_ids: list[set[int]] = [] - - for epoch in range(3): - recorder.thread_ids.clear() - recorder.process_ids.clear() - result = list(source) - self.assertEqual(sorted(result), [0, 1, 2, 3, 4], f"epoch {epoch}") - all_thread_ids.append(set(recorder.thread_ids)) - all_process_ids.append(set(recorder.process_ids)) - - # All epochs should use the same thread(s) — pipeline was reused - self.assertEqual( - all_thread_ids[0], - all_thread_ids[1], - "Thread IDs changed between epoch 0 and 1 — pipeline was rebuilt", - ) - self.assertEqual( - all_thread_ids[1], - all_thread_ids[2], - "Thread IDs changed between epoch 1 and 2 — pipeline was rebuilt", - ) - - # All epochs should run in the same subprocess - self.assertEqual( - all_process_ids[0], - all_process_ids[1], - "Process IDs changed between epoch 0 and 1", - ) - self.assertEqual( - all_process_ids[1], - all_process_ids[2], - "Process IDs changed between epoch 1 and 2", - ) - - -@unittest.skipIf( - sys.version_info < (3, 14), - "Subinterpreters require Python 3.14+", -) -class TestContinuousSubinterpreterPipelineReuse(unittest.TestCase): - @_ignore_fork_warning - def test_subinterpreter_pipeline_reused_across_epochs(self) -> None: - """Thread and process IDs stay the same across epochs in subinterpreter.""" - recorder = _RecordIDs() - - config = ( - PipelineBuilder() - .add_source(SourceIterable(5), continuous=True) - .pipe(recorder, concurrency=1) - .add_sink(buffer_size=3) - .get_config() - ) - source = run_pipeline_in_subinterpreter( - config, - num_threads=1, - timeout=10, - ) - - all_thread_ids: list[set[int]] = [] - all_process_ids: list[set[int]] = [] - - for epoch in range(3): - recorder.thread_ids.clear() - recorder.process_ids.clear() - result = list(source) - self.assertEqual(sorted(result), [0, 1, 2, 3, 4], f"epoch {epoch}") - all_thread_ids.append(set(recorder.thread_ids)) - all_process_ids.append(set(recorder.process_ids)) - - # All epochs should use the same thread(s) — pipeline was reused - self.assertEqual( - all_thread_ids[0], - all_thread_ids[1], - "Thread IDs changed between epoch 0 and 1 — pipeline was rebuilt", - ) - self.assertEqual( - all_thread_ids[1], - all_thread_ids[2], - "Thread IDs changed between epoch 1 and 2 — pipeline was rebuilt", - ) - - # All epochs should run in the same process (subinterpreter shares process) - self.assertEqual( - all_process_ids[0], - all_process_ids[1], - "Process IDs changed between epoch 0 and 1", - ) diff --git a/tests/pipeline/continuous_pipeline_test.py b/tests/pipeline/continuous_pipeline_test.py new file mode 120000 index 000000000..7536396cf --- /dev/null +++ b/tests/pipeline/continuous_pipeline_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/continuous_pipeline_test.py \ No newline at end of file diff --git a/tests/pipeline/defs_repr_test.py b/tests/pipeline/defs_repr_test.py deleted file mode 100644 index 20d08a258..000000000 --- a/tests/pipeline/defs_repr_test.py +++ /dev/null @@ -1,355 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import inspect -import unittest -from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence - -from spdl.pipeline import StageInfo -from spdl.pipeline.defs import ( - Aggregate, - Disaggregate, - Merge, - Pipe, - PipelineConfig, - SinkConfig, - SourceConfig, -) - - -# Test helper functions and classes -def example_sync_function(x: int) -> int: - """Example synchronous function for testing.""" - return x * 2 - - -async def example_async_function(x: int) -> int: - """Example async function for testing.""" - await asyncio.sleep(0) - return x * 2 - - -def _ln(target: object) -> int: - """Helper to get line number from inspect.getsourcelines.""" - return inspect.getsourcelines(target)[1] # pyre-ignore[6] - - -class ExampleIterable(Iterable[int]): - """Example iterable class for testing.""" - - def __iter__(self) -> Iterator[int]: - return iter([1, 2, 3]) - - -class ExampleAsyncIterable(AsyncIterable[int]): - """Example async iterable class for testing.""" - - async def __aiter__(self) -> AsyncIterator[int]: - for i in [1, 2, 3]: - yield i - - -async def custom_merge_op( - info: StageInfo, - input_queues: Sequence[asyncio.Queue], - output_queue: asyncio.Queue, -) -> None: - """Example custom merge operation for testing.""" - pass - - -class TestSourceConfigRepr(unittest.TestCase): - """Test SourceConfig.__repr__ with source location.""" - - def test_source_config_repr_with_iterable_class(self) -> None: - """Test __repr__ shows source location for iterable class.""" - # Setup: create source config with custom iterable - source = ExampleIterable() - config = SourceConfig(source=source) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains class name - # Note: for class instances, source location may not always be available - self.assertIn("ExampleIterable", result) - self.assertIn("SourceConfig", result) - - def test_source_config_repr_with_generator(self) -> None: - """Test __repr__ shows source location for generator function.""" - # Setup: create source config with generator - source = (x for x in range(10)) - config = SourceConfig(source=source) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains generator class name - self.assertIn("generator", result) - - def test_source_config_repr_with_async_iterable(self) -> None: - """Test __repr__ shows source location for async iterable.""" - # Setup: create source config with async iterable - source = ExampleAsyncIterable() - config = SourceConfig(source=source) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains class name - # Note: for class instances, source location may not always be available - self.assertIn("ExampleAsyncIterable", result) - self.assertIn("SourceConfig", result) - - -class TestPipeConfigRepr(unittest.TestCase): - """Test PipeConfig.__repr__ with source location.""" - - def test_pipe_config_repr_with_sync_function(self) -> None: - """Test __repr__ shows source location for sync function.""" - # Setup: create pipe config with sync function - config = Pipe(example_sync_function, concurrency=4) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains function name, concurrency, and source location - self.assertIn("concurrency=4", result) - self.assertIn("example_sync_function", result) - self.assertIn(__file__, result) - self.assertIn(f":{_ln(example_sync_function)}", result) - - def test_pipe_config_repr_with_async_function(self) -> None: - """Test __repr__ shows source location for async function.""" - # Setup: create pipe config with async function - config = Pipe(example_async_function, concurrency=2) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains function name, concurrency, and source location - self.assertIn("concurrency=2", result) - self.assertIn("example_async_function", result) - self.assertIn(__file__, result) - self.assertIn(f":{_ln(example_async_function)}", result) - - def test_pipe_config_repr_with_lambda(self) -> None: - """Test __repr__ handles lambda functions gracefully.""" - # Setup: create pipe config with lambda - config = Pipe(lambda x: x * 2, concurrency=1) - - # Execute: get repr - result = repr(config) - - # Assert: repr contains lambda and concurrency - self.assertIn("concurrency=1", result) - self.assertIn("lambda", result) - - def test_pipe_config_repr_without_source_location(self) -> None: - """Test __repr__ handles cases where source location cannot be determined.""" - # Setup: create pipe config with built-in function - config = Pipe(len, concurrency=1) - - # Execute: get repr - result = repr(config) - - # Assert: repr still works and contains concurrency - self.assertIn("concurrency=1", result) - self.assertIn("len", result) - - -class TestMergeConfigRepr(unittest.TestCase): - """Test MergeConfig.__repr__ with nested pipelines.""" - - def test_merge_config_repr_with_two_pipelines(self) -> None: - """Test __repr__ shows nested pipeline configs with proper indentation.""" - # Setup: create two pipeline configs and merge them - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=20), - ) - merge_config = Merge([plc1, plc2]) - - # Execute: get repr - result = repr(merge_config) - - # Assert: repr contains merge structure with both pipelines - self.assertIn("MergeConfig(", result) - self.assertIn("Pipeline 1:", result) - self.assertIn("Pipeline 2:", result) - self.assertIn("PipelineConfig", result) - - def test_merge_config_repr_with_custom_op(self) -> None: - """Test __repr__ shows custom merge operation with source location.""" - # Setup: create merge config with custom op - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - merge_config = Merge([plc1, plc2], op=custom_merge_op) - - # Execute: get repr - result = repr(merge_config) - - # Assert: repr contains op info with source location - self.assertIn("op=", result) - self.assertIn("custom_merge_op", result) - self.assertIn(__file__, result) - self.assertIn(f":{_ln(custom_merge_op)}", result) - - def test_merge_config_repr_multiline_structure(self) -> None: - """Test __repr__ creates multi-line output with proper indentation.""" - # Setup: create merge config with pipes - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[Pipe(example_sync_function, concurrency=2)], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - merge_config = Merge([plc1, plc2]) - - # Execute: get repr - result = repr(merge_config) - - # Assert: result is multi-line and properly indented - lines = result.split("\n") - self.assertGreater(len(lines), 5) # Multi-line output - # Check some lines have proper indentation - pipeline_lines = [line for line in lines if "Pipeline" in line] - self.assertGreater(len(pipeline_lines), 0) - - -class TestPipelineConfigRepr(unittest.TestCase): - """Test PipelineConfig.__repr__ with MergeConfig source.""" - - def test_pipeline_config_repr_with_source_config(self) -> None: - """Test __repr__ with SourceConfig shows inline representation.""" - # Setup: create pipeline config with simple source - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[Pipe(example_sync_function, concurrency=4)], - sink=SinkConfig(buffer_size=10), - ) - - # Execute: get repr - result = repr(config) - - # Assert: repr shows source inline - self.assertIn("PipelineConfig", result) - self.assertIn("Source:", result) - self.assertIn("Pipes:", result) - self.assertIn("Sink:", result) - self.assertIn("example_sync_function", result) - - def test_pipeline_config_repr_with_merge_config(self) -> None: - """Test __repr__ with MergeConfig shows proper indentation.""" - # Setup: create pipeline config with merge source - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - merge_config = Merge([plc1, plc2]) - final_config = PipelineConfig( - src=merge_config, - pipes=[Pipe(example_sync_function, concurrency=2)], - sink=SinkConfig(buffer_size=100), - ) - - # Execute: get repr - result = repr(final_config) - - # Assert: repr shows nested structure with proper indentation - self.assertIn("PipelineConfig", result) - self.assertIn("Source:", result) - self.assertIn("MergeConfig(", result) - self.assertIn("Pipeline 1:", result) - self.assertIn("Pipeline 2:", result) - # Check final pipe appears after merge - self.assertIn("example_sync_function", result) - - def test_pipeline_config_repr_indentation_hierarchy(self) -> None: - """Test __repr__ maintains correct indentation hierarchy.""" - # Setup: create nested pipeline config - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[Pipe(example_sync_function, concurrency=1)], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - merge_config = Merge([plc1, plc2]) - final_config = PipelineConfig( - src=merge_config, - pipes=[ - Pipe(example_async_function, concurrency=4), - Aggregate(5), - Disaggregate(), - ], - sink=SinkConfig(buffer_size=100), - ) - - # Execute: get repr - result = repr(final_config) - - # Assert: verify indentation levels exist - lines = result.split("\n") - # Should have various indentation levels - has_no_indent = any(line and not line[0].isspace() for line in lines) - has_some_indent = any(line.startswith(" ") for line in lines) - has_more_indent = any(line.startswith(" ") for line in lines) - - self.assertTrue(has_no_indent, "Should have lines with no indentation") - self.assertTrue(has_some_indent, "Should have lines with 2-space indentation") - self.assertTrue(has_more_indent, "Should have lines with 4+ space indentation") - - def test_pipeline_config_repr_with_aggregate_disaggregate(self) -> None: - """Test __repr__ shows aggregate and disaggregate pipes correctly.""" - # Setup: create pipeline config with various pipe types - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[ - Pipe(example_sync_function, concurrency=2), - Aggregate(10, drop_last=True), - Disaggregate(), - ], - sink=SinkConfig(buffer_size=10), - ) - - # Execute: get repr - result = repr(config) - - # Assert: repr shows all pipe types - self.assertIn("example_sync_function", result) - self.assertIn("aggregate", result) - self.assertIn("disaggregate", result) diff --git a/tests/pipeline/defs_repr_test.py b/tests/pipeline/defs_repr_test.py new file mode 120000 index 000000000..1648db541 --- /dev/null +++ b/tests/pipeline/defs_repr_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/defs_repr_test.py \ No newline at end of file diff --git a/tests/pipeline/failure_rate_test.py b/tests/pipeline/failure_rate_test.py deleted file mode 100644 index 1e6ecb12f..000000000 --- a/tests/pipeline/failure_rate_test.py +++ /dev/null @@ -1,804 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import functools -import unittest -import warnings -from collections.abc import Callable -from fractions import Fraction -from typing import Type, TypeVar - -from parameterized import parameterized -from spdl.pipeline import PipelineBuilder, PipelineFailure - -_F = TypeVar("_F", bound=Callable[..., object]) -_C = TypeVar("_C", bound=Type[object]) - - -def _ignore_intentional_warnings(fn: _F) -> _F: - """Suppress warnings emitted intentionally by failure-rate tests: - - - the fork() multi-threaded DeprecationWarning from the subprocess - pipeline machinery, and - - "coroutine ... was never awaited" RuntimeWarnings, which surface when - the pipeline is forced to fail mid-iteration and the source coroutine - is dropped before being fully consumed. - """ - - @functools.wraps(fn) - def wrapper(*args: object, **kwargs: object) -> object: - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - warnings.filterwarnings( - "ignore", - message="coroutine .* was never awaited", - category=RuntimeWarning, - ) - return fn(*args, **kwargs) - - # pyre-ignore[7] - return wrapper - - -def _ignore_intentional_warnings_in_class(cls: _C) -> _C: - for name, member in list(vars(cls).items()): - if name.startswith("test_") and callable(member): - setattr(cls, name, _ignore_intentional_warnings(member)) - return cls - - -@_ignore_intentional_warnings_in_class -class PipelineFailureRateTest(unittest.TestCase): - """Tests for Fraction-based failure rate thresholds in SPDL pipeline. - - Key design: A fixed probation period of 100 invocations is used before - rate-based checking kicks in. This prevents early false positives when - sample size is too small to be statistically meaningful. - - The pipeline stops when failure rate strictly exceeds the threshold (>). - """ - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_basic(self, output_order: str) -> None: - """Pipeline fails when failure rate exceeds Fraction threshold. - - Uses 1000 items. Fails on multiples of 100 (10 out of 1000 = 1%). - With 0.1% threshold and fixed probation of 100, should fail. - After probation: rate = 1% > 0.1% threshold -> fails. - """ - - def fail_on_hundred(x: int) -> int: - if x % 100 == 0: # Fails on 0, 100, 200, ..., 900 (10 out of 1000) - raise ValueError(f"Multiple of 100: {x}") - return x - - # 0.1% threshold - should fail because actual rate is 1% - pipeline = ( - PipelineBuilder() - .add_source(range(1000)) - .pipe(fail_on_hundred, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 1000)) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=30)) - - all_expected = {x for x in range(1000) if x % 100 != 0} - self.assertTrue(len(vals) > 0) - self.assertTrue(set(vals).issubset(all_expected)) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_passes(self, output_order: str) -> None: - """Pipeline succeeds when failure rate stays below threshold. - - Uses 100 items. Fails on multiples of 10 (10 failures = 10%). - With 15% threshold (Fraction(3, 20)) and fixed probation of 100. - After probation: rate = 10% < 15% -> succeeds. - """ - - def fail_on_ten(x: int) -> int: - if x % 10 == 0: # Fails on 0, 10, 20, ..., 90 (10 out of 100 = 10%) - raise ValueError(f"Multiple of 10: {x}") - return x - - # Allow 15% failure rate - should succeed because actual rate is 10% - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_on_ten, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(3, 20)) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 10 - expected = [x for x in range(100) if x % 10 != 0] - self.assertEqual(expected, vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_just_below_threshold( - self, output_order: str - ) -> None: - """Pipeline succeeds when failure rate is just below threshold. - - Uses 100 items, fails on items >= 90 (10 failures = 10%). - With 11% threshold and fixed probation of 100. - After probation: rate = 10% < 11% -> succeeds. - """ - - def fail_late(x: int) -> int: - if x >= 90: # Fails on 90-99 (10 out of 100 = 10%) - raise ValueError(f"Item >= 90: {x}") - return x - - # Allow 11% failure rate - should succeed because actual rate is 10% - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_late, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(11, 100)) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get values 0-89 (90 successful items) - self.assertEqual(list(range(90)), vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_above_threshold(self, output_order: str) -> None: - """Pipeline fails when failure rate exceeds threshold. - - Uses 100 items, fails on items >= 90 (10 failures = 10%). - With 9% threshold and fixed probation of 100. - After probation: rate = 10% > 9% -> fails. - """ - - def fail_late(x: int) -> int: - if x >= 90: # Fails on 90-99 (10 out of 100 = 10%) - raise ValueError(f"Item >= 90: {x}") - return x - - # 9% threshold - should fail because actual rate is 10% - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_late, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(9, 100)) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get values 0-89 (90 successful items) - self.assertEqual(list(range(90)), vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_probation_period(self, output_order: str) -> None: - """Early failures don't trigger threshold due to fixed probation period (100). - - Uses 20 items. Fails on multiples of 5 (4 out of 20 = 20%). - With 1% threshold but fixed probation of 100, only 20 items processed. - Since probation never completes, pipeline succeeds despite 20% > 1%. - """ - - def fail_on_five(x: int) -> int: - if x % 5 == 0: # Fails on 0, 5, 10, 15 (4 out of 20 = 20%) - raise ValueError(f"Multiple of 5: {x}") - return x - - # 1% threshold with fixed probation=100 - # Since we only have 20 items, probation never completes - pipeline = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(fail_on_five, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 100)) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 5: 1,2,3,4,6,7,8,9,11,12,13,14,16,17,18,19 - expected = [x for x in range(20) if x % 5 != 0] - self.assertEqual(expected, vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_pipe_override(self, output_order: str) -> None: - """Per-pipe Fraction override works correctly. - - Uses 100 items. Fails on multiples of 7 (15 out of 100 = 15%). - Pipe-level: 11% threshold - should fail (15% > 11%). - Global: 19% threshold - would pass. - Since pipe-level is stricter, pipeline should fail. - """ - - def fail_on_seven(x: int) -> int: - if x % 7 == 0: # Fails on 0, 7, 14, ..., 98 (15 out of 100 = 15%) - raise ValueError(f"Multiple of 7: {x}") - return x - - # Global allows 19% but pipe-level restricts to 11% - should fail - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe( - fail_on_seven, - output_order=output_order, - max_failures=Fraction(11, 100), - ) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(19, 100)) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 7 - expected = [x for x in range(100) if x % 7 != 0] - self.assertEqual(expected, vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_multiple_stages(self, output_order: str) -> None: - """Multiple stages with different Fraction thresholds. - - First stage: fails on odd numbers (50%) but allowed unlimited. - Second stage: receives 50 even numbers, fails on divisible by 12 (9/50 = 18%). - With 22% threshold (Fraction(11, 50)), 18% < 22% -> should pass. - Note: probation is fixed at 100, but only 50 items reach second stage. - """ - - def fail_odd(x: int) -> int: - if x % 2: - raise ValueError(f"Odd number: {x}") - return x - - def fail_twelve(x: int) -> int: - if (x % 12) == 0: - raise ValueError(f"Divisible by 12: {x}") - return x - - # First stage fails 50% (odd numbers) but allowed unlimited failures - # Second stage: 18% failure rate with 22% threshold -> should succeed - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_odd, output_order=output_order, max_failures=-1) - .pipe(fail_twelve, output_order=output_order, max_failures=Fraction(11, 50)) - .add_sink(1) - .build(num_threads=1, max_failures=-1) - ) - - # Second stage receives 50 even numbers (0,2,4,...98) - # Fails on 0,12,24,36,48,60,72,84,96 = 9 failures out of 50 = 18% - # With 22% threshold, should succeed - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Even numbers not divisible by 12: 2,4,6,8,10,14,16,... - expected = [x for x in range(100) if x % 2 == 0 and x % 12 != 0] - self.assertEqual(expected, vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_vs_count(self, output_order: str) -> None: - """Verify int count-based behavior unchanged. - - Fails on multiples of 10 (10 failures). - Count-based: Allow 15 failures - should succeed. - Count-based: Allow 5 failures - should fail. - """ - - def fail_on_ten(x: int) -> int: - if x % 10 == 0: # Fails on 0, 10, 20, ..., 90 (10 failures) - raise ValueError(f"Multiple of 10: {x}") - return x - - # Count-based: Allow 15 failures - should succeed - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_on_ten, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=15) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 10 - expected = [x for x in range(100) if x % 10 != 0] - self.assertEqual(expected, vals) - - # Count-based: Allow 5 failures - should fail after 5 failures - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_on_ten, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=5) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Pipeline stops early after 5 failures; vals is a subset of expected - all_expected = {x for x in range(100) if x % 10 != 0} - self.assertTrue(len(vals) > 0) - self.assertTrue(set(vals).issubset(all_expected)) - - def test_pipeline_failure_rate_ordered_pipe(self) -> None: - """Test with output_order='input'. - - Uses 100 items. Fails on multiples of 7 (15 out of 100 = 15%). - With 11% threshold and fixed probation of 100. - After probation: rate = 15% > 11% -> should fail. - """ - - def fail_on_seven(x: int) -> int: - if x % 7 == 0: # Fails on 0, 7, 14, ..., 98 (15 out of 100 = 15%) - raise ValueError(f"Multiple of 7: {x}") - return x - - # 11% threshold - should fail because actual rate is ~15% - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_on_seven, output_order="input", max_failures=Fraction(11, 100)) - .add_sink(1) - .build(num_threads=1) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 7 - expected = [x for x in range(100) if x % 7 != 0] - self.assertEqual(expected, vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_high_concurrency(self, output_order: str) -> None: - """Test failure rate with high concurrency. - - Uses 100 items. Fails on multiples of 5 (20 out of 100 = 20%). - With 15% threshold (Fraction(3, 20)) and fixed probation of 100. - After probation: rate = 20% > 15% -> should fail. - With high concurrency, exact processing order is nondeterministic. - """ - - def fail_on_five(x: int) -> int: - if x % 5 == 0: # Fails on 0, 5, 10, ..., 95 (20 out of 100 = 20%) - raise ValueError(f"Multiple of 5: {x}") - return x - - # 15% threshold - should fail because actual rate is 20% - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe( - fail_on_five, - output_order=output_order, - concurrency=5, - max_failures=Fraction(3, 20), - ) - .add_sink(1) - .build(num_threads=5, max_failures=Fraction(1, 2)) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_zero_failures(self, output_order: str) -> None: - """Test with zero failures. - - 0% failure rate with 10% threshold - should succeed. - """ - - def no_fail(x: int) -> int: - return x * 2 - - # 0% failure rate with 10% threshold - should succeed - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(no_fail, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 10)) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all values doubled - self.assertEqual([x * 2 for x in range(100)], vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_failure_rate_under_probation_always_succeeds( - self, output_order: str - ) -> None: - """With fixed probation of 100, pipelines with < 100 items always succeed. - - Uses 10 items with 40% failure rate (multiples of 3 fail). - With 30% threshold but fixed probation of 100. - Since only 10 items processed, probation not reached -> succeeds. - """ - - def fail_on_three(x: int) -> int: - if x % 3 == 0: # Fails on 0, 3, 6, 9 (4 out of 10 = 40%) - raise ValueError(f"Multiple of 3: {x}") - return x - - # 30% threshold, but only 10 items (under probation) - # Should succeed because probation (100) not reached - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_on_three, output_order=output_order) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(3, 10)) - ) - - # Should succeed despite 40% > 30% because probation not complete - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 3: 1, 2, 4, 5, 7, 8 - expected = [x for x in range(10) if x % 3 != 0] - self.assertEqual(expected, vals) - - def test_pipeline_failure_rate_invalid_fraction_zero(self) -> None: - """Building pipeline with Fraction <= 0 raises ValueError.""" - - def noop(x: int) -> int: - return x - - # Zero Fraction should raise ValueError - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( - num_threads=1, max_failures=Fraction(0, 100) - ) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - def test_pipeline_failure_rate_invalid_fraction_negative(self) -> None: - """Building pipeline with negative Fraction raises ValueError.""" - - def noop(x: int) -> int: - return x - - # Negative Fraction should raise ValueError - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( - num_threads=1, max_failures=Fraction(-1, 10) - ) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - def test_pipeline_failure_rate_invalid_fraction_greater_than_one(self) -> None: - """Building pipeline with Fraction > 1 raises ValueError.""" - - def noop(x: int) -> int: - return x - - # Fraction > 1 (e.g., 150%) should raise ValueError - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( - num_threads=1, max_failures=Fraction(15, 10) - ) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - def test_pipeline_failure_rate_invalid_pipe_fraction_zero(self) -> None: - """Building pipeline with zero Fraction at pipe level raises ValueError.""" - - def noop(x: int) -> int: - return x - - # Zero Fraction at pipe level should raise ValueError - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe( - noop, max_failures=Fraction(0, 100) - ).add_sink(1).build(num_threads=1) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - def test_pipeline_failure_rate_invalid_pipe_fraction_greater_than_one(self) -> None: - """Building pipeline with Fraction > 1 at pipe level raises ValueError.""" - - def noop(x: int) -> int: - return x - - # Fraction > 1 at pipe level should raise ValueError - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe( - noop, max_failures=Fraction(200, 100) - ).add_sink(1).build(num_threads=1) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - def test_pipeline_failure_rate_valid_fraction_one(self) -> None: - """Fraction(1, 1) = 100% is valid (allows all failures). - - With > comparison, rate can never exceed 100%, so pipeline always succeeds. - """ - - def always_fail(x: int) -> int: - raise ValueError(f"Always fail: {x}") - - # 100% failure rate threshold - should never fail - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(always_fail) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 1)) - ) - - # Should complete without PipelineFailure (100% failures allowed) - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # No items should pass through since all fail - self.assertEqual([], vals) - - def test_pipeline_failure_rate_valid_small_fraction(self) -> None: - """Very small Fraction like Fraction(1, 1000) is valid.""" - - def fail_on_hundred(x: int) -> int: - if x % 100 == 0: # Fails on 0, 100, 200, ..., 900 (10 out of 1000 = 1%) - raise ValueError(f"Multiple of 100: {x}") - return x - - # 0.1% threshold (Fraction(1, 1000)) - should fail because actual rate is 1% - pipeline = ( - PipelineBuilder() - .add_source(range(1000)) - .pipe(fail_on_hundred) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 1000)) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=30)) - - all_expected = {x for x in range(1000) if x % 100 != 0} - self.assertTrue(len(vals) > 0) - self.assertTrue(set(vals).issubset(all_expected)) - - @parameterized.expand( - [ - (Fraction(0, 1), "zero numerator"), - (Fraction(0, 100), "zero with large denominator"), - (Fraction(-1, 10), "negative numerator"), - (Fraction(1, -10), "negative denominator"), - (Fraction(-1, -10), "double negative (positive > 0 but > 1)"), - ] - ) - def test_pipeline_failure_rate_invalid_fractions_parameterized( - self, fraction: Fraction, description: str - ) -> None: - """Parameterized test for various invalid Fraction values.""" - - def noop(x: int) -> int: - return x - - # Note: Fraction(-1, -10) normalizes to Fraction(1, 10) which is valid - # But Fraction(0, x) and negative fractions should fail - if fraction <= 0 or fraction > 1: - with self.assertRaises(ValueError): - PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( - num_threads=1, max_failures=fraction - ) - else: - # This should not raise - it's a valid fraction - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(noop) - .add_sink(1) - .build(num_threads=1, max_failures=fraction) - ) - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(list(range(10)), vals) - - @parameterized.expand( - [ - (Fraction(11, 10), "110%"), - (Fraction(2, 1), "200%"), - (Fraction(150, 100), "150%"), - (Fraction(101, 100), "101%"), - ] - ) - def test_pipeline_failure_rate_invalid_fractions_greater_than_one_parameterized( - self, fraction: Fraction, description: str - ) -> None: - """Parameterized test for Fraction values > 1 (greater than 100%).""" - - def noop(x: int) -> int: - return x - - with self.assertRaises(ValueError) as ctx: - PipelineBuilder().add_source(range(10)).pipe(noop).add_sink(1).build( - num_threads=1, max_failures=fraction - ) - - self.assertIn("must be in range (0, 1]", str(ctx.exception)) - - @parameterized.expand( - [ - (Fraction(1, 100), "1%"), - (Fraction(1, 10), "10%"), - (Fraction(1, 2), "50%"), - (Fraction(99, 100), "99%"), - (Fraction(1, 1), "100%"), - ] - ) - def test_pipeline_failure_rate_valid_fractions_parameterized( - self, fraction: Fraction, description: str - ) -> None: - """Parameterized test for valid Fraction values in range (0, 1].""" - - def noop(x: int) -> int: - return x - - # Should not raise - these are all valid fractions - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(noop) - .add_sink(1) - .build(num_threads=1, max_failures=fraction) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # All values should pass through - self.assertEqual(list(range(100)), vals) - - def test_pipeline_failure_rate_probation_prevents_early_trigger(self) -> None: - """Probation period (fixed at 100) prevents triggering even when rate is high. - - Uses 10 items, fails on multiples of 5 (2/10 = 20%). - With 1% threshold but fixed probation of 100, only 10 items processed. - Since probation not reached, pipeline succeeds despite 20% > 1%. - """ - - def fail_on_five(x: int) -> int: - if x % 5 == 0: # Fails on 0, 5 (2 out of 10 = 20%) - raise ValueError(f"Multiple of 5: {x}") - return x - - # 1% threshold with fixed probation=100. Only 10 items -> no check runs. - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_on_five) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 100)) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 5: 1, 2, 3, 4, 6, 7, 8, 9 - expected = [x for x in range(10) if x % 5 != 0] - self.assertEqual(expected, vals) - - def test_pipeline_failure_rate_probation_triggers_after_warmup(self) -> None: - """Once probation period (100) is met, threshold check runs. - - Uses 100 items, fails on multiples of 5 (20/100 = 20%). - With 10% threshold and fixed probation of 100. - After 100 invocations: rate = 20% > 10% -> fails. - """ - - def fail_on_five(x: int) -> int: - if x % 5 == 0: # Fails on 0, 5, ..., 95 (20 out of 100 = 20%) - raise ValueError(f"Multiple of 5: {x}") - return x - - # 10% threshold with fixed probation=100. All 100 items are processed. - pipeline = ( - PipelineBuilder() - .add_source(range(100)) - .pipe(fail_on_five) - .add_sink(1) - .build(num_threads=1, max_failures=Fraction(1, 10)) - ) - - vals = [] - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - # Should get all non-multiples of 5 - expected = [x for x in range(100) if x % 5 != 0] - self.assertEqual(expected, vals) diff --git a/tests/pipeline/failure_rate_test.py b/tests/pipeline/failure_rate_test.py new file mode 120000 index 000000000..dc8662dfc --- /dev/null +++ b/tests/pipeline/failure_rate_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/failure_rate_test.py \ No newline at end of file diff --git a/tests/pipeline/merge_config_test.py b/tests/pipeline/merge_config_test.py deleted file mode 100644 index 22917001e..000000000 --- a/tests/pipeline/merge_config_test.py +++ /dev/null @@ -1,475 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Tests for MergeConfig class.""" - -# pyre-strict - -import asyncio -import unittest -from collections.abc import Sequence - -from spdl.pipeline import build_pipeline, create_task, is_eof, StageInfo -from spdl.pipeline.defs import ( - Merge, - Pipe, - PipelineConfig, - SinkConfig, - SourceConfig, -) - - -class MergeConfigTest(unittest.TestCase): - """Test MergeConfig functionality.""" - - def test_merge_config_with_two_simple_pipelines(self) -> None: - """Test MergeConfig merges outputs from two simple pipelines.""" - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 6) - self.assertCountEqual(results, [1, 2, 3, 4, 5, 6]) - - def test_merge_config_with_processed_pipelines(self) -> None: - """Test MergeConfig merges outputs from pipelines with processing.""" - double_pipe = Pipe(lambda x: x * 2) - add_ten_pipe = Pipe(lambda x: x + 10) - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[double_pipe], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[add_ten_pipe], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 6) - # Pipeline 1: [1, 2, 3] -> [2, 4, 6] (doubled) - # Pipeline 2: [4, 5, 6] -> [14, 15, 16] (added 10) - self.assertCountEqual(results, [2, 4, 6, 14, 15, 16]) - - def test_merge_config_with_multiple_pipelines(self) -> None: - """Test MergeConfig can merge outputs from three pipelines.""" - plc1 = PipelineConfig( - src=SourceConfig([1, 2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([10, 20]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc3 = PipelineConfig( - src=SourceConfig([100, 200]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2, plc3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=3) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 6) - self.assertCountEqual(results, [1, 2, 10, 20, 100, 200]) - - def test_merge_config_with_post_processing(self) -> None: - """Test MergeConfig output can be further processed in main pipeline.""" - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - multiply_by_5_pipe = Pipe(lambda x: x * 5) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[multiply_by_5_pipe], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 6) - self.assertCountEqual(results, [5, 10, 15, 20, 25, 30]) - - def test_merge_config_with_async_processing(self) -> None: - """Test MergeConfig works with async processing functions.""" - - async def async_double(x: int) -> int: - await asyncio.sleep(0.01) # Small delay to simulate async work - return x * 2 - - async_pipe = Pipe(async_double) - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[async_pipe], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[async_pipe], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=5)) - - self.assertEqual(len(results), 6) - self.assertCountEqual(results, [2, 4, 6, 8, 10, 12]) - - def test_merge_config_with_different_data_types(self) -> None: - """Test MergeConfig works with different data types.""" - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig(["a", "b", "c"]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 6) - self.assertCountEqual(results, [1, 2, 3, "a", "b", "c"]) - - def test_merge_config_with_empty_pipeline(self) -> None: - """Test MergeConfig handles pipeline with empty source.""" - plc1 = PipelineConfig( - src=SourceConfig([]), # Empty source - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 3) - self.assertCountEqual(results, [1, 2, 3]) - - def test_merge_config_validation_empty_list(self) -> None: - """Test MergeConfig validation fails with empty pipeline list.""" - with self.assertRaises(ValueError) as cm: - Merge([]) - - self.assertIn("at least one upstream pipeline", str(cm.exception)) - - def test_merge_config_with_aggregation(self) -> None: - """Test MergeConfig works with aggregation operations.""" - from spdl.pipeline.defs import Aggregate, Disaggregate - - aggregate_pipe = Aggregate(2) - disaggregate_pipe = Disaggregate() - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3, 4]), - pipes=[aggregate_pipe, disaggregate_pipe], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([10, 20, 30, 40]), - pipes=[aggregate_pipe, disaggregate_pipe], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 8) - self.assertCountEqual(results, [1, 2, 3, 4, 10, 20, 30, 40]) - - def test_merge_config_with_concurrency(self) -> None: - """Test MergeConfig works with concurrent processing.""" - concurrent_pipe = Pipe(lambda x: x + 100, concurrency=3) - - plc1 = PipelineConfig( - src=SourceConfig(list(range(10))), - pipes=[concurrent_pipe], - sink=SinkConfig(buffer_size=20), - ) - - plc2 = PipelineConfig( - src=SourceConfig(list(range(50, 60))), - pipes=[concurrent_pipe], - sink=SinkConfig(buffer_size=20), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[], - sink=SinkConfig(buffer_size=50), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=4) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 20) - expected = list(range(100, 110)) + list(range(150, 160)) - self.assertCountEqual(results, expected) - - def test_merge_config_with_different_pipe_counts_and_post_processing(self) -> None: - """Test MergeConfig merges pipelines with different numbers of pipes and applies post-processing.""" - - # Pipeline 1: single pipe (multiply by 2) - multiply_pipe = Pipe(lambda x: x * 2) - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[multiply_pipe], - sink=SinkConfig(buffer_size=10), - ) - - # Pipeline 2: three pipes (add 10, multiply by 3, subtract 5) - add_ten_pipe = Pipe(lambda x: x + 10) - multiply_by_three_pipe = Pipe(lambda x: x * 3) - subtract_five_pipe = Pipe(lambda x: x - 5) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5]), - pipes=[add_ten_pipe, multiply_by_three_pipe, subtract_five_pipe], - sink=SinkConfig(buffer_size=10), - ) - - # Add post-processing after merge: add 100 to all merged results - post_process_pipe = Pipe(lambda x: x + 100) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[post_process_pipe], - sink=SinkConfig(buffer_size=20), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(len(results), 5) - - # [1, 2, 3] --(multiply by 2)--> [2, 4, 6] --(add 100)--> [102, 104, 106] - pipeline1_expected = [102, 104, 106] - - # [4, 5] --(add 10)--> [14, 15] - # --(multiply by 3)--> [42, 45] - # --(subtract 5)--> [37, 40] - # --(add 100)--> [137, 140] - pipeline2_expected = [137, 140] - - expected_all = pipeline1_expected + pipeline2_expected - self.assertCountEqual(results, expected_all) - - def test_merge_config_with_custom_merge_op(self) -> None: - """Test MergeConfig accepts and uses custom merge operation.""" - - # Track which pipelines contributed items (for verification) - collected_items: list[str] = [] - - async def custom_merge_op( - info: StageInfo, - input_queues: Sequence[asyncio.Queue[object]], - output_queue: asyncio.Queue[object], - ) -> None: - """Custom merge that adds a prefix to each item based on its source pipeline.""" - - async def process_queue( - queue_idx: int, in_q: asyncio.Queue[object] - ) -> None: - while True: - item = await in_q.get() - if is_eof(item): - return - # Add prefix based on source pipeline - prefixed_item = f"p{queue_idx}_{item}" - collected_items.append(prefixed_item) - await output_queue.put(prefixed_item) - - tasks = [ - create_task(process_queue(i, in_q), name=f"{info}:{i}") - for i, in_q in enumerate(input_queues) - ] - await asyncio.wait(tasks) - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2], op=custom_merge_op), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - # Verify we got all items with prefixes - self.assertEqual(len(results), 6) - # Items from pipeline 0 should have p0_ prefix - self.assertIn("p0_1", results) - self.assertIn("p0_2", results) - self.assertIn("p0_3", results) - # Items from pipeline 1 should have p1_ prefix - self.assertIn("p1_4", results) - self.assertIn("p1_5", results) - self.assertIn("p1_6", results) - - def test_merge_config_with_custom_merge_op_early_exit(self) -> None: - """Test MergeConfig accepts and uses custom merge operation.""" - - async def custom_merge_op( - _: StageInfo, - input_queues: Sequence[asyncio.Queue[object]], - output_queue: asyncio.Queue[object], - ) -> None: - """Custom merge that exists when one sub-pipeline completes""" - while True: - for in_q in input_queues: - item = await in_q.get() - if is_eof(item): - return - - await output_queue.put(item) - - plc1 = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - plc2 = PipelineConfig( - src=SourceConfig([4, 5, 6, 7, 8, 9]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - main_pipeline_config = PipelineConfig( - src=Merge([plc1, plc2], op=custom_merge_op), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - - pipeline = build_pipeline(main_pipeline_config, num_threads=2) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=3)) - - self.assertEqual(results, [1, 4, 2, 5, 3, 6]) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/pipeline/merge_config_test.py b/tests/pipeline/merge_config_test.py new file mode 120000 index 000000000..98a45acc2 --- /dev/null +++ b/tests/pipeline/merge_config_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/merge_config_test.py \ No newline at end of file diff --git a/tests/pipeline/path_variants_test.py b/tests/pipeline/path_variants_test.py deleted file mode 100644 index c6c387d93..000000000 --- a/tests/pipeline/path_variants_test.py +++ /dev/null @@ -1,704 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Tests for PathVariants feature.""" - -# pyre-unsafe - -import asyncio -import unittest - -from spdl.pipeline import build_pipeline -from spdl.pipeline._components._node import PipelineFailure -from spdl.pipeline.defs import ( - Aggregate, - Disaggregate, - Merge, - PathVariants, - Pipe, - PipelineConfig, - SinkConfig, - SourceConfig, -) - - -def _run_pipeline(config, num_threads=2, timeout=5): - """Helper to build, run, and collect results from a pipeline.""" - pipeline = build_pipeline(config, num_threads=num_threads) - with pipeline.auto_stop(): - return list(pipeline.get_iterator(timeout=timeout)) - - -async def _slow_source(items, delay=0.1): - """Yield items with a delay between them. - - Each item is fully processed before the next one is dispatched, - making cross-path ordering deterministic (matching input order). - """ - for item in items: - yield item - await asyncio.sleep(delay) - - -class PathVariantsBasicTest(unittest.TestCase): - """Basic functional tests for PathVariants.""" - - def test_basic_routing(self) -> None: - """Even items to path 0 (double), odd items to path 1 (add 100).""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(6))), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 2)], # path 0: evens - [Pipe(lambda x: x + 100)], # path 1: odds - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # Slow source ensures each item is processed before the next arrives, - # so output order matches input order: - # 0→path0→0, 1→path1→101, 2→path0→4, 3→path1→103, 4→path0→8, 5→path1→105 - self.assertEqual(results, [0, 101, 4, 103, 8, 105]) - - def test_async_router(self) -> None: - """Async router function works correctly.""" - - async def async_router(x: int) -> int: - return x % 2 - - config = PipelineConfig( - src=SourceConfig(_slow_source(range(6))), - pipes=[ - PathVariants( - router=async_router, - paths=[ - [Pipe(lambda x: x * 2)], - [Pipe(lambda x: x + 100)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, [0, 101, 4, 103, 8, 105]) - - def test_async_callable_router(self) -> None: - """Class with async __call__ works as router.""" - - class AsyncRouter: - async def __call__(self, x: int) -> int: - return x % 2 - - config = PipelineConfig( - src=SourceConfig(_slow_source(range(6))), - pipes=[ - PathVariants( - router=AsyncRouter(), - paths=[ - [Pipe(lambda x: x * 2)], - [Pipe(lambda x: x + 100)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, [0, 101, 4, 103, 8, 105]) - - def test_all_to_one_path(self) -> None: - """Router always returns 0 — all items go to path 0.""" - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[ - PathVariants( - router=lambda x: 0, - paths=[ - [Pipe(lambda x: x * 10)], - [Pipe(lambda x: x + 1000)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, [10, 20, 30]) - - def test_multiple_paths_different_processing(self) -> None: - """3 paths with different transforms.""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(9))), - pipes=[ - PathVariants( - router=lambda x: x % 3, - paths=[ - [Pipe(lambda x: x * 2)], # path 0: ×2 - [Pipe(lambda x: x + 10)], # path 1: +10 - [Pipe(lambda x: -x)], # path 2: negate - ], - ), - ], - sink=SinkConfig(buffer_size=20), - ) - results = _run_pipeline(config) - # Slow source ensures interleaved input order: - # 0→path0→0, 1→path1→11, 2→path2→-2, - # 3→path0→6, 4→path1→14, 5→path2→-5, - # 6→path0→12, 7→path1→17, 8→path2→-8 - self.assertEqual(results, [0, 11, -2, 6, 14, -5, 12, 17, -8]) - - def test_identity_path_passthrough(self) -> None: - """A path with an identity pipe passes items through unchanged.""" - config = PipelineConfig( - src=SourceConfig(_slow_source([1, 2, 3, 4])), - pipes=[ - PathVariants( - router=lambda x: 0 if x <= 2 else 1, - paths=[ - [Pipe(lambda x: x * 100)], # path 0: transform - [Pipe(lambda x: x)], # path 1: passthrough - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # 1→path0→100, 2→path0→200, 3→path1→3, 4→path1→4 - self.assertEqual(results, [100, 200, 3, 4]) - - def test_paths_with_different_stage_counts(self) -> None: - """Paths with different numbers of pipe stages.""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(6))), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 10)], # path 0: 1 stage - [ - Pipe(lambda x: (x, x + 100)), - Pipe(lambda t: t[1]), - ], # path 1: 2 stages - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # Slow source ensures input-order interleaving even with different - # stage counts: 0→0, 1→(1,101)→101, 2→20, 3→(3,103)→103, 4→40, 5→105 - self.assertEqual(results, [0, 101, 20, 103, 40, 105]) - - def test_path_with_aggregate(self) -> None: - """Path containing Aggregate inside.""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(6))), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Aggregate(3)], # path 0: batch evens by 3 - [Pipe(lambda x: x + 100)], # path 1: add 100 to odds - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # Aggregate waits for 3 evens (0,2,4). With slow source, odds pass - # through immediately while aggregate buffers: - # 0→agg, 1→101, 2→agg, 3→103, 4→agg→[0,2,4], 5→105 - self.assertEqual(results, [101, 103, [0, 2, 4], 105]) - - def test_path_with_aggregate_and_disaggregate(self) -> None: - """Path containing Aggregate + Disaggregate inside.""" - config = PipelineConfig( - src=SourceConfig(range(6)), - pipes=[ - PathVariants( - router=lambda x: 0, - paths=[ - [Aggregate(2), Disaggregate()], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, [0, 1, 2, 3, 4, 5]) - - def test_nested_path_variants(self) -> None: - """PathVariants inside a path of another PathVariants.""" - inner_variants = PathVariants( - router=lambda x: 0 if x < 50 else 1, - paths=[ - [Pipe(lambda x: x + 1000)], # inner path 0 - [Pipe(lambda x: x + 2000)], # inner path 1 - ], - ) - config = PipelineConfig( - src=SourceConfig(_slow_source([1, 2, 51, 52])), - pipes=[ - PathVariants( - router=lambda x: 0, # all to path 0 - paths=[ - [inner_variants], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # 1→inner path0→1001, 2→inner path0→1002, - # 51→inner path1→2051, 52→inner path1→2052 - self.assertEqual(results, [1001, 1002, 2051, 2052]) - - def test_immediately_nested_path_variants(self) -> None: - """PathVariants as the first and only config in each path of an outer - PathVariants — the inner router reads directly from the outer router's - per-path queue.""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(12))), - pipes=[ - PathVariants( - router=lambda x: x % 2, # outer: evens vs odds - paths=[ - # path 0 (evens): immediately nest another PathVariants - [ - PathVariants( - router=lambda x: 0 if x < 6 else 1, - paths=[ - [Pipe(lambda x: x * 10)], # small evens - [Pipe(lambda x: x * 100)], # large evens - ], - ), - ], - # path 1 (odds): immediately nest another PathVariants - [ - PathVariants( - router=lambda x: 0 if x < 6 else 1, - paths=[ - [Pipe(lambda x: -x)], # small odds - [Pipe(lambda x: -(x * 10))], # large odds - ], - ), - ], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # Items interleaved across outer paths, then inner paths: - # 0→even→small→0, 1→odd→small→-1, 2→even→small→20, 3→odd→small→-3, - # 4→even→small→40, 5→odd→small→-5, 6→even→large→600, 7→odd→large→-70, - # 8→even→large→800, 9→odd→large→-90, 10→even→large→1000, 11→odd→large→-110 - self.assertEqual( - results, - [0, -1, 20, -3, 40, -5, 600, -70, 800, -90, 1000, -110], - ) - - def test_path_variants_before_and_after_pipes(self) -> None: - """Pipes before and after PathVariants in the main pipeline.""" - config = PipelineConfig( - src=SourceConfig(_slow_source([1, 2, 3, 4])), - pipes=[ - Pipe(lambda x: x * 10), # pre-processing - PathVariants( - router=lambda x: 0 if x < 25 else 1, - paths=[ - [Pipe(lambda x: x + 1)], - [Pipe(lambda x: x + 2)], - ], - ), - Pipe(lambda x: x * -1), # post-processing - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # pre: [10,20,30,40] - # path0 (<25): 10->11, 20->21 - # path1 (>=25): 30->32, 40->42 - # post: [-11,-21,-32,-42] - self.assertEqual(results, [-11, -21, -32, -42]) - - def test_path_variants_with_merge_source(self) -> None: - """Merge as source, then PathVariants in pipes.""" - plc1 = PipelineConfig( - src=SourceConfig([1, 2]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - plc2 = PipelineConfig( - src=SourceConfig([3, 4]), - pipes=[], - sink=SinkConfig(buffer_size=10), - ) - config = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 100)], - [Pipe(lambda x: x * -1)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # evens: 2,4 -> 200,400 - # odds: 1,3 -> -1,-3 - self.assertCountEqual(results, [200, 400, -1, -3]) - - def test_path_variants_with_async_pipe(self) -> None: - """Async pipe inside a path.""" - - async def async_double(x: int) -> int: - await asyncio.sleep(0.01) - return x * 2 - - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[ - PathVariants( - router=lambda x: 0, - paths=[ - [Pipe(async_double)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, [2, 4, 6]) - - def test_path_variants_with_concurrent_pipe(self) -> None: - """Pipe with concurrency > 1 inside a path.""" - config = PipelineConfig( - src=SourceConfig(_slow_source(range(10))), - pipes=[ - PathVariants( - router=lambda x: 0, - paths=[ - [Pipe(lambda x: x + 1, concurrency=3)], - ], - ), - ], - sink=SinkConfig(buffer_size=20), - ) - results = _run_pipeline(config) - self.assertEqual(results, list(range(1, 11))) - - -class PathVariantsValidationTest(unittest.TestCase): - """Validation tests for PathVariants config.""" - - def test_validation_empty_paths(self) -> None: - """PathVariants with no paths raises ValueError.""" - with self.assertRaises(ValueError): - PathVariants(router=lambda x: 0, paths=[]) - - def test_validation_empty_path(self) -> None: - """PathVariants with an empty path raises ValueError.""" - with self.assertRaises(ValueError): - PathVariants( - router=lambda x: 0, - paths=[[Pipe(lambda x: x)], []], - ) - - def test_validation_non_callable_router(self) -> None: - """Non-callable router raises ValueError.""" - with self.assertRaises(ValueError): - # pyre-ignore[6]: Intentionally passing non-callable to test validation. - PathVariants(router=42, paths=[[Pipe(lambda x: x)]]) - - def test_validation_source_in_path(self) -> None: - """SourceConfig in a path raises ValueError at construction time.""" - with self.assertRaises(ValueError): - PathVariants( - router=lambda x: 0, - # pyre-ignore[6]: Intentionally passing SourceConfig. - paths=[[SourceConfig([3, 4])]], - ) - - def test_validation_sink_in_path(self) -> None: - """SinkConfig in a path raises ValueError at construction time.""" - with self.assertRaises(ValueError): - PathVariants( - router=lambda x: 0, - # pyre-ignore[6]: Intentionally passing SinkConfig. - paths=[[SinkConfig(buffer_size=10)]], - ) - - def test_validation_source_in_second_path(self) -> None: - """SourceConfig in a later path position raises ValueError.""" - with self.assertRaises(ValueError): - PathVariants( - router=lambda x: 0, - # pyre-ignore[6]: Intentionally passing SourceConfig. - paths=[ - [Pipe(lambda x: x)], - [Pipe(lambda x: x), SourceConfig([1, 2])], - ], - ) - - def test_validation_sink_in_middle_of_path(self) -> None: - """SinkConfig in the middle of a path raises ValueError.""" - with self.assertRaises(ValueError): - PathVariants( - router=lambda x: 0, - # pyre-ignore[6]: Intentionally passing SinkConfig. - paths=[ - [Pipe(lambda x: x), SinkConfig(buffer_size=10), Pipe(lambda x: x)], - ], - ) - - -class PathVariantsErrorHandlingTest(unittest.TestCase): - """Error handling and edge case tests for PathVariants.""" - - def test_router_returns_negative_index(self) -> None: - """Router returns -1 — pipeline fails with PipelineFailure.""" - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[ - PathVariants( - router=lambda x: -1, - paths=[ - [Pipe(lambda x: x)], - [Pipe(lambda x: x)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - def test_router_returns_index_equal_to_num_paths(self) -> None: - """Router returns N (== len(paths)) — pipeline fails.""" - config = PipelineConfig( - src=SourceConfig([1]), - pipes=[ - PathVariants( - router=lambda x: 2, # only 2 paths, index 2 is out of range - paths=[ - [Pipe(lambda x: x)], - [Pipe(lambda x: x)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - def test_router_returns_index_greater_than_num_paths(self) -> None: - """Router returns N+5 — pipeline fails.""" - config = PipelineConfig( - src=SourceConfig([1]), - pipes=[ - PathVariants( - router=lambda x: 7, - paths=[ - [Pipe(lambda x: x)], - [Pipe(lambda x: x)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - def test_one_path_fails_other_continues(self) -> None: - """One path's pipe raises; the other path's items still processed. - - With default max_failures=-1, individual task failures are tolerated. - The pipeline completes successfully and the non-failing path's results - are collected. The failing path's items are dropped. - """ - - def fail_on_odd(x): - if x % 2 == 1: - raise ValueError(f"odd item {x}") - return x * 10 - - config = PipelineConfig( - src=SourceConfig(range(6)), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 10)], # path 0: evens succeed - [Pipe(fail_on_odd)], # path 1: odds fail (dropped) - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - # Path 0 items succeed: 0,2,4 -> 0,20,40 - # Path 1 items fail and are dropped - self.assertEqual(results, [0, 20, 40]) - - def test_one_path_fails_with_max_failures(self) -> None: - """One path's pipe raises with max_failures=0; pipeline fails.""" - - def fail_always(x): - raise ValueError("fail") - - config = PipelineConfig( - src=SourceConfig(range(4)), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 10)], - [Pipe(fail_always, max_failures=0)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - def test_all_paths_fail(self) -> None: - """All paths fail with max_failures=0 — pipeline raises PipelineFailure.""" - - def always_fail(x): - raise ValueError("boom") - - config = PipelineConfig( - src=SourceConfig([1, 2]), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(always_fail, max_failures=0)], - [Pipe(always_fail, max_failures=0)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - def test_empty_source_with_path_variants(self) -> None: - """Empty source — clean shutdown with no results.""" - config = PipelineConfig( - src=SourceConfig([]), - pipes=[ - PathVariants( - router=lambda x: 0, - paths=[ - [Pipe(lambda x: x * 2)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - results = _run_pipeline(config) - self.assertEqual(results, []) - - def test_path_failure_cancels_router(self) -> None: - """When a path fails with max_failures=0, the router is cancelled and - the pipeline shuts down cleanly without hanging.""" - - def fail_immediately(x): - raise ValueError("boom") - - config = PipelineConfig( - # Use enough items so the router is still active when path fails. - src=SourceConfig(range(100)), - pipes=[ - PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x, max_failures=0)], # path 0: succeeds - [Pipe(fail_immediately, max_failures=0)], # path 1: fails - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - # The pipeline must raise PipelineFailure (not hang). - # If the router is not cancelled on path failure, this would deadlock - # because the router keeps trying to push items to the dead path's queue. - with self.assertRaises(PipelineFailure): - _run_pipeline(config, timeout=5) - - def test_path_failure_cancels_router_all_to_failing_path(self) -> None: - """All items routed to the failing path — router cancelled, no hang.""" - - def fail_immediately(x): - raise ValueError("boom") - - config = PipelineConfig( - src=SourceConfig(range(100)), - pipes=[ - PathVariants( - router=lambda x: 1, # all items to failing path - paths=[ - [Pipe(lambda x: x)], - [Pipe(fail_immediately, max_failures=0)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config, timeout=5) - - def test_router_raises_exception(self) -> None: - """Router function itself raises — pipeline fails cleanly.""" - - def bad_router(x): - raise RuntimeError("router error") - - config = PipelineConfig( - src=SourceConfig([1, 2, 3]), - pipes=[ - PathVariants( - router=bad_router, - paths=[ - [Pipe(lambda x: x)], - ], - ), - ], - sink=SinkConfig(buffer_size=10), - ) - with self.assertRaises(PipelineFailure): - _run_pipeline(config) - - -class PathVariantsReprTest(unittest.TestCase): - """Repr tests for PathVariants.""" - - def test_repr(self) -> None: - """Verify repr is readable and includes path info.""" - cfg = PathVariants( - router=lambda x: x % 2, - paths=[ - [Pipe(lambda x: x * 2, name="double")], - [Pipe(lambda x: x + 1, name="add_one")], - ], - ) - r = repr(cfg) - self.assertIn("PathVariants", r) - self.assertIn("path0", r) - self.assertIn("path1", r) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/pipeline/path_variants_test.py b/tests/pipeline/path_variants_test.py new file mode 120000 index 000000000..3662710e9 --- /dev/null +++ b/tests/pipeline/path_variants_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/path_variants_test.py \ No newline at end of file diff --git a/tests/pipeline/percentile_stats_test.py b/tests/pipeline/percentile_stats_test.py deleted file mode 100644 index 83758584f..000000000 --- a/tests/pipeline/percentile_stats_test.py +++ /dev/null @@ -1,453 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import unittest -from collections.abc import Callable - -from spdl.pipeline._components._common import _P2Percentile, _StatsCounter, StageInfo -from spdl.pipeline._components._hook import TaskPerfStats, TaskStatsHook -from spdl.pipeline._components._queue import QueuePerfStats, StatsQueue - - -class P2PercentileTest(unittest.TestCase): - def test_empty(self) -> None: - """A fresh _P2Percentile with no observations reports 0.0.""" - p = _P2Percentile(90) - self.assertEqual(p.value, 0.0) - - def test_single_element(self) -> None: - """With only one observation, the percentile value equals that observation.""" - p = _P2Percentile(90) - p.update(5.0) - self.assertEqual(p.value, 5.0) - - def test_two_elements(self) -> None: - """With two observations, p90 returns the larger value.""" - p = _P2Percentile(90) - p.update(1.0) - p.update(2.0) - self.assertEqual(p.value, 2.0) - - def test_four_elements(self) -> None: - """With fewer than 5 observations, the fallback sorted-lookup is used. - - Checks that p50 of [1, 2, 3, 4] returns the median (3.0). - """ - p = _P2Percentile(50) - for v in [1.0, 2.0, 3.0, 4.0]: - p.update(v) - self.assertEqual(p.value, 3.0) - - def test_five_elements_p90(self) -> None: - """With exactly 5 observations the P² algorithm initializes. - - Checks that the p90 estimate falls within a reasonable range. - """ - p = _P2Percentile(90) - for v in [1.0, 2.0, 3.0, 4.0, 5.0]: - p.update(v) - self.assertGreaterEqual(p.value, 2.0) - self.assertLessEqual(p.value, 5.0) - - def test_hundred_elements_p90(self) -> None: - """P² p90 estimate over 100 sequential values is close to the true p90 (90.0).""" - p = _P2Percentile(90) - for i in range(100): - p.update(float(i)) - self.assertAlmostEqual(p.value, 90.0, delta=3.0) - - def test_hundred_elements_p99(self) -> None: - """P² p99 estimate over 100 sequential values is close to the true p99 (99.0).""" - p = _P2Percentile(99) - for i in range(100): - p.update(float(i)) - self.assertAlmostEqual(p.value, 99.0, delta=3.0) - - def test_hundred_elements_p50(self) -> None: - """P² p50 (median) estimate over 100 sequential values is close to 50.0.""" - p = _P2Percentile(50) - for i in range(100): - p.update(float(i)) - self.assertAlmostEqual(p.value, 50.0, delta=3.0) - - def test_thousand_elements_accuracy(self) -> None: - """With 1000 observations, p90 and p99 estimates stay within tight bounds. - - Checks p90 ≈ 900 (±20) and p99 ≈ 990 (±20). - """ - p90 = _P2Percentile(90) - p99 = _P2Percentile(99) - for i in range(1000): - v = float(i) - p90.update(v) - p99.update(v) - self.assertAlmostEqual(p90.value, 900.0, delta=20.0) - self.assertAlmostEqual(p99.value, 990.0, delta=20.0) - - def test_reset(self) -> None: - """After reset(), the estimator returns to its initial state (value 0.0).""" - p = _P2Percentile(90) - for i in range(20): - p.update(float(i)) - self.assertGreater(p.value, 0.0) - p.reset() - self.assertEqual(p.value, 0.0) - - def test_reset_and_reuse(self) -> None: - """After reset(), feeding new data produces estimates based only on the new data. - - First feeds [0..99], resets, then feeds [100..199] and checks p50 ≈ 150. - """ - p = _P2Percentile(50) - for i in range(100): - p.update(float(i)) - p.reset() - for i in range(100, 200): - p.update(float(i)) - self.assertAlmostEqual(p.value, 150.0, delta=5.0) - - def test_unsorted_input(self) -> None: - """P² produces accurate estimates regardless of input order. - - Feeds 10 values in shuffled order and checks p90 ≈ 8.0 (±2). - """ - p = _P2Percentile(90) - values = [9.0, 1.0, 5.0, 3.0, 7.0, 2.0, 8.0, 4.0, 6.0, 0.0] - for v in values: - p.update(v) - self.assertAlmostEqual(p.value, 8.0, delta=2.0) - - def test_constant_values(self) -> None: - """When all observations are identical, the percentile equals that constant.""" - p = _P2Percentile(90) - for _ in range(20): - p.update(42.0) - self.assertAlmostEqual(p.value, 42.0, delta=0.01) - - -class StatsCounterPercentileTest(unittest.TestCase): - def test_initial_state(self) -> None: - """A fresh _StatsCounter has zero items, zero average, and zero percentiles.""" - counter = _StatsCounter() - self.assertEqual(counter.p90_time, 0.0) - self.assertEqual(counter.p99_time, 0.0) - self.assertEqual(counter.num_items, 0) - self.assertEqual(counter.ave_time, 0.0) - - def test_update_tracks_percentiles(self) -> None: - """After 100 updates, the counter's p90 and p99 properties reflect accurate estimates.""" - counter = _StatsCounter() - for i in range(100): - counter.update(float(i)) - self.assertEqual(counter.num_items, 100) - self.assertAlmostEqual(counter.p90_time, 90.0, delta=3.0) - self.assertAlmostEqual(counter.p99_time, 99.0, delta=3.0) - - def test_count_context_manager(self) -> None: - """The count() context manager records one item with non-negative percentiles.""" - counter = _StatsCounter() - with counter.count(): - pass - self.assertEqual(counter.num_items, 1) - self.assertGreaterEqual(counter.p90_time, 0.0) - self.assertGreaterEqual(counter.p99_time, 0.0) - - def test_consume_lap_percentiles(self) -> None: - """consume_lap_percentiles() returns current lap p90/p99, then resets lap trackers. - - After consuming, a second call returns (0.0, 0.0). - """ - counter = _StatsCounter() - for i in range(100): - counter.update(float(i)) - - p90, p99 = counter.consume_lap_percentiles() - self.assertAlmostEqual(p90, 90.0, delta=3.0) - self.assertAlmostEqual(p99, 99.0, delta=3.0) - - p90_after, p99_after = counter.consume_lap_percentiles() - self.assertEqual(p90_after, 0.0) - self.assertEqual(p99_after, 0.0) - - def test_consume_lap_does_not_affect_overall(self) -> None: - """Consuming lap percentiles does not change the overall (lifetime) p90 value.""" - counter = _StatsCounter() - for i in range(100): - counter.update(float(i)) - - overall_p90_before = counter.p90_time - counter.consume_lap_percentiles() - self.assertEqual(counter.p90_time, overall_p90_before) - - -class TaskPerfStatsTest(unittest.TestCase): - def test_fields_present(self) -> None: - """TaskPerfStats dataclass stores all fields including p90_time and p99_time.""" - stats = TaskPerfStats( - num_tasks=10, - num_failures=1, - ave_time=0.5, - p90_time=0.8, - p99_time=1.2, - ) - self.assertEqual(stats.num_tasks, 10) - self.assertEqual(stats.num_failures, 1) - self.assertEqual(stats.ave_time, 0.5) - self.assertEqual(stats.p90_time, 0.8) - self.assertEqual(stats.p99_time, 1.2) - - -class QueuePerfStatsTest(unittest.TestCase): - def test_fields_present(self) -> None: - """QueuePerfStats dataclass stores p90/p99 fields for both put and get operations.""" - stats = QueuePerfStats( - elapsed=60.0, - num_items=100, - ave_put_time=0.01, - ave_get_time=0.02, - p90_put_time=0.015, - p99_put_time=0.025, - p90_get_time=0.03, - p99_get_time=0.04, - occupancy_rate=0.75, - ) - self.assertEqual(stats.p90_put_time, 0.015) - self.assertEqual(stats.p99_put_time, 0.025) - self.assertEqual(stats.p90_get_time, 0.03) - self.assertEqual(stats.p99_get_time, 0.04) - - def test_qps(self) -> None: - """The qps property computes num_items / elapsed correctly.""" - stats = QueuePerfStats( - elapsed=10.0, - num_items=100, - ave_put_time=0.0, - ave_get_time=0.0, - p90_put_time=0.0, - p99_put_time=0.0, - p90_get_time=0.0, - p99_get_time=0.0, - occupancy_rate=0.0, - ) - self.assertAlmostEqual(stats.qps, 10.0) - - def test_qps_zero_elapsed(self) -> None: - """When elapsed is zero, qps returns 0 to avoid division by zero.""" - stats = QueuePerfStats( - elapsed=0.0, - num_items=100, - ave_put_time=0.0, - ave_get_time=0.0, - p90_put_time=0.0, - p99_put_time=0.0, - p90_get_time=0.0, - p99_get_time=0.0, - occupancy_rate=0.0, - ) - self.assertEqual(stats.qps, 0) - - -class TaskStatsHookPercentileTest(unittest.IsolatedAsyncioTestCase): - async def test_task_hook_records_percentiles(self) -> None: - """Successful tasks are tracked by the P² percentile estimators. - - Checks that after 10 successful tasks, num_tasks/num_success are correct - and both p90/p99 trackers have non-negative values. - """ - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 - ) - for _ in range(10): - async with hook.task_hook(): # pyre-ignore[16] - pass - self.assertEqual(hook.num_tasks, 10) - self.assertEqual(hook.num_success, 10) - self.assertGreaterEqual(hook._p90.value, 0.0) - self.assertGreaterEqual(hook._p99.value, 0.0) - - async def test_failed_task_not_tracked_in_percentiles(self) -> None: - """Failed tasks increment num_tasks but are excluded from percentile tracking. - - Runs 3 tasks (2 succeed, 1 fails). Checks that the P² tracker count is 2. - """ - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 - ) - - async with hook.task_hook(): # pyre-ignore[16] - pass - - try: - async with hook.task_hook(): # pyre-ignore[16] - raise ValueError("fail") - except ValueError: - pass - - async with hook.task_hook(): # pyre-ignore[16] - pass - - self.assertEqual(hook.num_tasks, 3) - self.assertEqual(hook.num_success, 2) - self.assertEqual(hook._p90._count, 2) - - async def test_stage_hook_produces_stats_with_percentiles(self) -> None: - """When stage_hook exits, _log_stats is called with a TaskPerfStats that - includes non-negative p90_time and p99_time values. - """ - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 - ) - logged_stats: list[TaskPerfStats] = [] - original_log: Callable[[TaskPerfStats], None] = hook._log_stats - - def capture_stats(stats: TaskPerfStats) -> None: - logged_stats.append(stats) - original_log(stats) - - hook._log_stats = capture_stats # pyre-ignore[8] - - async with hook.stage_hook(): # pyre-ignore[16] - for _ in range(5): - async with hook.task_hook(): # pyre-ignore[16] - pass - - self.assertEqual(len(logged_stats), 1) - stats = logged_stats[0] - self.assertEqual(stats.num_tasks, 5) - self.assertEqual(stats.num_failures, 0) - self.assertGreater(stats.ave_time, 0.0) - self.assertGreaterEqual(stats.p90_time, 0.0) - self.assertGreaterEqual(stats.p99_time, 0.0) - - async def test_lap_stats_with_percentiles(self) -> None: - """Lap stats report percentiles only for the current interval, then reset. - - First lap covers 10 tasks; second lap covers 5 new tasks. Each lap's - p90/p99 should be positive and independent of the other. - """ - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 - ) - - for _ in range(10): - async with hook.task_hook(): # pyre-ignore[16] - pass - - lap1 = hook._get_lap_stats() - self.assertEqual(lap1.num_tasks, 10) - self.assertGreater(lap1.p90_time, 0.0) - self.assertGreater(lap1.p99_time, 0.0) - - for _ in range(5): - async with hook.task_hook(): # pyre-ignore[16] - pass - - lap2 = hook._get_lap_stats() - self.assertEqual(lap2.num_tasks, 5) - self.assertGreater(lap2.p90_time, 0.0) - self.assertGreater(lap2.p99_time, 0.0) - - async def test_empty_lap_stats(self) -> None: - """When no tasks have run, lap stats report zero for all percentile fields.""" - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), interval=-1 - ) - lap = hook._get_lap_stats() - self.assertEqual(lap.num_tasks, 0) - self.assertEqual(lap.p90_time, 0.0) - self.assertEqual(lap.p99_time, 0.0) - - -class StatsQueuePercentileTest(unittest.IsolatedAsyncioTestCase): - async def test_put_get_records_percentiles(self) -> None: - """After put/get operations, the queue's internal counters have - non-negative p90 percentile values for both put and get. - """ - queue = StatsQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), - buffer_size=10, - interval=-1, - ) - async with queue.stage_hook(): # pyre-ignore[16] - for i in range(5): - await queue.put(i) - for _ in range(5): - await queue.get() - - self.assertGreaterEqual(queue._putc.p90_time, 0.0) - self.assertGreaterEqual(queue._getc.p90_time, 0.0) - - async def test_stage_hook_produces_stats_with_percentiles(self) -> None: - """When stage_hook exits, _log_stats receives a QueuePerfStats with - non-negative p90/p99 values for both put and get operations. - """ - queue = StatsQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), - buffer_size=10, - interval=-1, - ) - logged_stats: list[QueuePerfStats] = [] - original_log: Callable[[QueuePerfStats], None] = queue._log_stats - - def capture_stats(stats: QueuePerfStats) -> None: - logged_stats.append(stats) - original_log(stats) - - queue._log_stats = capture_stats # pyre-ignore[8] - - async with queue.stage_hook(): # pyre-ignore[16] - for i in range(5): - await queue.put(i) - for _ in range(5): - await queue.get() - - self.assertEqual(len(logged_stats), 1) - stats = logged_stats[0] - self.assertEqual(stats.num_items, 5) - self.assertGreaterEqual(stats.p90_put_time, 0.0) - self.assertGreaterEqual(stats.p99_put_time, 0.0) - self.assertGreaterEqual(stats.p90_get_time, 0.0) - self.assertGreaterEqual(stats.p99_get_time, 0.0) - - async def test_lap_stats_with_percentiles(self) -> None: - """Lap stats for the queue report per-interval p90/p99 for put and get, - resetting between laps. - - First lap covers 8 items; second lap covers 3 new items. Each lap's - percentile fields should be non-negative and independent. - """ - queue = StatsQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="test"), - buffer_size=10, - interval=-1, - ) - queue._lap_t0 = asyncio.get_running_loop().time() - queue._empty_t0 = queue._lap_t0 - - for i in range(8): - await queue.put(i) - for _ in range(8): - await queue.get() - - lap1 = queue._get_lap_stats() - self.assertEqual(lap1.num_items, 8) - self.assertGreaterEqual(lap1.p90_put_time, 0.0) - self.assertGreaterEqual(lap1.p90_get_time, 0.0) - self.assertGreaterEqual(lap1.p99_put_time, 0.0) - self.assertGreaterEqual(lap1.p99_get_time, 0.0) - - for i in range(3): - await queue.put(i) - for _ in range(3): - await queue.get() - - lap2 = queue._get_lap_stats() - self.assertEqual(lap2.num_items, 3) - self.assertGreaterEqual(lap2.p90_put_time, 0.0) - self.assertGreaterEqual(lap2.p90_get_time, 0.0) diff --git a/tests/pipeline/percentile_stats_test.py b/tests/pipeline/percentile_stats_test.py new file mode 120000 index 000000000..23320ce25 --- /dev/null +++ b/tests/pipeline/percentile_stats_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/percentile_stats_test.py \ No newline at end of file diff --git a/tests/pipeline/pgrp_stats_test.py b/tests/pipeline/pgrp_stats_test.py deleted file mode 100644 index d7ede1829..000000000 --- a/tests/pipeline/pgrp_stats_test.py +++ /dev/null @@ -1,496 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import multiprocessing -import os -import sys -import tempfile -import unittest -from collections.abc import Awaitable, Callable -from unittest.mock import AsyncMock, MagicMock, patch - -from spdl.pipeline._bg_task import BackgroundTask -from spdl.pipeline._pgrp_stats import ( - _collect_pgrp_stats, - _parse_proc_io, - _parse_proc_stat, - _parse_smaps_rollup, - _pgrp_monitor_subprocess, - _read_file, - _read_network_bytes, - _read_pgrp_stats, - _warned, - ProcessGroupResourceUsage, - ProcessGroupStatsMonitor, -) - -_MODULE = "spdl.pipeline._pgrp_stats" - - -@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") -class LiveProcMonitorTest(unittest.TestCase): - """Integration test that reads real /proc data without mocking.""" - - def test_read_pgrp_stats_returns_valid_data(self) -> None: - """_read_pgrp_stats should find at least this process.""" - _warned.discard("proc_stat") - _warned.discard("proc_io") - _warned.discard("smaps_rollup") - - result = _read_pgrp_stats() - self.assertGreaterEqual(result.num_procs, 1) - self.assertGreaterEqual(result.cpu_usec, 0) - self.assertGreaterEqual(result.rss_bytes, 0) - self.assertGreaterEqual(result.disk_read_bytes, 0) - self.assertGreaterEqual(result.disk_write_bytes, 0) - # smaps_rollup should be available on modern Linux - if result.pss_bytes is not None: - self.assertGreater(result.pss_bytes, 0) - self.assertIsNotNone(result.private_bytes) - self.assertGreater(result.private_bytes, 0) - - def test_collect_pgrp_stats_returns_complete_snapshot(self) -> None: - """_collect_pgrp_stats should return a fully populated snapshot.""" - result, cpu_usec, time_usec, net_rx, net_tx = _collect_pgrp_stats() - self.assertIsInstance(result, ProcessGroupResourceUsage) - self.assertEqual(result.pid, os.getpid()) - self.assertEqual(result.pgid, os.getpgrp()) - - # First call: cpu_percent and net deltas should be None (no previous value). - self.assertIsNone(result.cpu_percent) - self.assertIsNotNone(cpu_usec) - self.assertIsNotNone(result.rss_bytes) - self.assertIsNotNone(result.num_procs) - self.assertIsNone(result.net_rx_bytes) - self.assertIsNone(result.net_tx_bytes) - self.assertIsNotNone(net_rx) - self.assertIsNotNone(net_tx) - - # Second call with prev values: cpu_percent and net deltas should be set. - result2, _, _, _, _ = _collect_pgrp_stats(cpu_usec, time_usec, net_rx, net_tx) - cpu_pct = result2.cpu_percent - assert cpu_pct is not None - self.assertGreaterEqual(cpu_pct, 0.0) - rx = result2.net_rx_bytes - assert rx is not None - self.assertGreaterEqual(rx, 0) - tx = result2.net_tx_bytes - assert tx is not None - self.assertGreaterEqual(tx, 0) - - # Sanity: at least one process (this one) should be counted. - assert result.num_procs is not None - self.assertGreaterEqual(result.num_procs, 1) - # RSS must be positive for a running process. - assert result.rss_bytes is not None - self.assertGreater(result.rss_bytes, 0) - - -class ReadFileTest(unittest.TestCase): - def test_read_existing_file(self) -> None: - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - f.write("hello\n") - f.flush() - result = _read_file(f.name) - self.assertEqual(result, "hello") - os.unlink(f.name) - - def test_read_nonexistent_file(self) -> None: - result = _read_file("/nonexistent/path/file.txt") - self.assertIsNone(result) - - -class ReadNetworkTest(unittest.TestCase): - @patch(f"{_MODULE}._read_file") - def test_network_bytes(self, mock_read: MagicMock) -> None: - mock_read.return_value = ( - "Inter-| Receive | Transmit\n" # noqa: B950 - " face |bytes packets errs drop fifo frame compressed multicast|bytes packets errs drop fifo colls carrier compressed\n" # noqa: B950 - " lo: 1000 10 0 0 0 0 0 0 2000 20 0 0 0 0 0 0\n" # noqa: B950 - " eth0: 5000 50 0 0 0 0 0 0 3000 30 0 0 0 0 0 0\n" # noqa: B950 - " eth1: 7000 70 0 0 0 0 0 0 4000 40 0 0 0 0 0 0" # noqa: B950 - ) - result = _read_network_bytes() - # lo is excluded; eth0 + eth1 - self.assertEqual(result.rx_bytes, 12000) - self.assertEqual(result.tx_bytes, 7000) - - @patch(f"{_MODULE}._read_file") - def test_network_file_missing(self, mock_read: MagicMock) -> None: - mock_read.return_value = None - result = _read_network_bytes() - self.assertEqual(result.rx_bytes, 0) - self.assertEqual(result.tx_bytes, 0) - - @patch(f"{_MODULE}._read_file") - def test_network_malformed_line_raises(self, mock_read: MagicMock) -> None: - mock_read.return_value = ( - "Inter-| Receive\n" - " face |bytes\n" - " eth0: bad 0 0 0 0 0 0 0 also_bad 0 0 0 0 0 0 0\n" - ) - with self.assertRaises(RuntimeError, msg="Failed to parse /proc/net/dev"): - _read_network_bytes() - - -class ParseProcStatTest(unittest.TestCase): - def test_normal_comm(self) -> None: - content = ( - "12345 (python3) S 100 200 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" - ) - stat = _parse_proc_stat(content) - self.assertEqual(stat.pgrp, 200) - self.assertEqual(stat.utime, 500) - self.assertEqual(stat.stime, 300) - self.assertEqual(stat.rss, 4096) - - def test_comm_with_spaces_and_parens(self) -> None: - content = "12345 (my (weird) app) S 100 200 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" # noqa: B950 - stat = _parse_proc_stat(content) - self.assertEqual(stat.pgrp, 200) - - def test_malformed_no_parens_raises(self) -> None: - with self.assertRaises(RuntimeError, msg="missing closing paren"): - _parse_proc_stat("no parens here") - - def test_too_few_fields_raises(self) -> None: - with self.assertRaises(RuntimeError, msg="expected >=22 fields"): - _parse_proc_stat("12345 (python3) S 100 200") - - def test_non_numeric_field_raises(self) -> None: - content = ( - "12345 (python3) S 100 abc 200 0 -1 0 0 0 0 0 500 300 0 0 20 0 1 0 0 0 4096" # noqa: B950 - ) - with self.assertRaises(RuntimeError, msg="Failed to parse"): - _parse_proc_stat(content) - - -class ParseProcIoTest(unittest.TestCase): - def test_normal_io(self) -> None: - content = ( - "rchar: 123456\n" - "wchar: 654321\n" - "syscr: 100\n" - "syscw: 200\n" - "read_bytes: 4096\n" - "write_bytes: 8192\n" - "cancelled_write_bytes: 0" - ) - io = _parse_proc_io(content) - self.assertEqual(io.read_bytes, 4096) - self.assertEqual(io.write_bytes, 8192) - - def test_empty_content(self) -> None: - io = _parse_proc_io("") - self.assertEqual(io.read_bytes, 0) - self.assertEqual(io.write_bytes, 0) - - def test_malformed_value_raises(self) -> None: - content = "read_bytes: not_a_number\n" - with self.assertRaises(RuntimeError, msg="Failed to parse /proc/[pid]/io"): - _parse_proc_io(content) - - -class ParseSmapsRollupTest(unittest.TestCase): - def test_normal_smaps_rollup(self) -> None: - content = ( - "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\n" - "Rss: 123456 kB\n" - "Pss: 98765 kB\n" - "Pss_Dirty: 40000 kB\n" - "Private_Clean: 50000 kB\n" - "Private_Dirty: 30000 kB\n" - "Shared_Clean: 20000 kB\n" - "Shared_Dirty: 10000 kB\n" - ) - result = _parse_smaps_rollup(content) - self.assertEqual(result.pss, 98765 * 1024) - self.assertEqual(result.private_clean, 50000 * 1024) - self.assertEqual(result.private_dirty, 30000 * 1024) - - def test_empty_content(self) -> None: - result = _parse_smaps_rollup("") - self.assertEqual(result.pss, 0) - self.assertEqual(result.private_clean, 0) - self.assertEqual(result.private_dirty, 0) - - def test_malformed_value_raises(self) -> None: - content = "Pss: not_a_number kB\n" - with self.assertRaises(RuntimeError, msg="Failed to parse"): - _parse_smaps_rollup(content) - - -@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") -class ReadPgrpStatsTest(unittest.TestCase): - def setUp(self) -> None: - _warned.discard("proc_stat") - _warned.discard("proc_io") - - @patch(f"{_MODULE}.os.scandir") - @patch(f"{_MODULE}._read_file") - @patch(f"{_MODULE}.os.getpgrp") - def test_sums_processes_in_same_pgrp( - self, - mock_getpgrp: MagicMock, - mock_read: MagicMock, - mock_scandir: MagicMock, - ) -> None: - mock_getpgrp.return_value = 1000 - - # Two processes in pgrp 1000, one in pgrp 9999 - entries = [] - for name in ["101", "102", "103", "not_a_pid"]: - entry = MagicMock() - entry.name = name - entry.is_dir.return_value = True - entries.append(entry) - mock_scandir.return_value = entries - - def read_side_effect(path: str) -> str | None: - if path == "/proc/101/stat": - return "101 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 100 50 0 0 20 0 1 0 0 0 2000" # noqa: B950 - if path == "/proc/101/smaps_rollup": - return "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\nRss: 8000 kB\nPss: 6000 kB\nPrivate_Clean: 3000 kB\nPrivate_Dirty: 2000 kB\n" # noqa: B950 - if path == "/proc/101/io": - return "rchar: 1000\nwchar: 2000\nsyscr: 10\nsyscw: 20\nread_bytes: 4096\nwrite_bytes: 8192\ncancelled_write_bytes: 0" # noqa: B950 - if path == "/proc/102/stat": - return "102 (worker) S 1 1000 1000 0 -1 0 0 0 0 0 200 75 0 0 20 0 1 0 0 0 3000" # noqa: B950 - if path == "/proc/102/smaps_rollup": - return "00400000-ffffffff ---p 00000000 00:00 0 [rollup]\nRss: 12000 kB\nPss: 9000 kB\nPrivate_Clean: 5000 kB\nPrivate_Dirty: 3000 kB\n" # noqa: B950 - if path == "/proc/102/io": - return "rchar: 3000\nwchar: 4000\nsyscr: 30\nsyscw: 40\nread_bytes: 1024\nwrite_bytes: 2048\ncancelled_write_bytes: 0" # noqa: B950 - if path == "/proc/103/stat": - # Different pgrp - return "103 (other) S 1 9999 9999 0 -1 0 0 0 0 0 999 999 0 0 20 0 1 0 0 0 9999" # noqa: B950 - return None - - mock_read.side_effect = read_side_effect - - result = _read_pgrp_stats() - - # utime: 100+200=300, stime: 50+75=125, total_ticks=425 - from spdl.pipeline._pgrp_stats import _get_sc_clk_tck - - expected_cpu_usec = 425 * 1_000_000 // _get_sc_clk_tck() - self.assertEqual(result.cpu_usec, expected_cpu_usec) - - # rss: 2000+3000=5000 pages - from spdl.pipeline._pgrp_stats import _get_page_size - - expected_rss = 5000 * _get_page_size() - self.assertEqual(result.rss_bytes, expected_rss) - - # pss: 6000+9000=15000 kB - self.assertEqual(result.pss_bytes, 15000 * 1024) - - # private: (3000+2000)+(5000+3000)=13000 kB - self.assertEqual(result.private_bytes, 13000 * 1024) - - # disk IO: 4096+1024=5120 read, 8192+2048=10240 write - self.assertEqual(result.disk_read_bytes, 5120) - self.assertEqual(result.disk_write_bytes, 10240) - - self.assertEqual(result.num_procs, 2) - - @patch(f"{_MODULE}.os.scandir") - @patch(f"{_MODULE}.os.getpgrp") - def test_scandir_failure_raises( - self, - mock_getpgrp: MagicMock, - mock_scandir: MagicMock, - ) -> None: - mock_getpgrp.return_value = 1000 - mock_scandir.side_effect = OSError("permission denied") - - with self.assertRaises(RuntimeError, msg="Failed to scan /proc"): - _read_pgrp_stats() - - @patch(f"{_MODULE}.os.scandir") - @patch(f"{_MODULE}._read_file") - @patch(f"{_MODULE}.os.getpgrp") - def test_missing_io_and_smaps_file( - self, - mock_getpgrp: MagicMock, - mock_read: MagicMock, - mock_scandir: MagicMock, - ) -> None: - """Disk IO is 0 and PSS/private are None when files are unreadable.""" - mock_getpgrp.return_value = 1000 - - entry = MagicMock() - entry.name = "101" - mock_scandir.return_value = [entry] - - def read_side_effect(path: str) -> str | None: - if path == "/proc/101/stat": - return "101 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 100 50 0 0 20 0 1 0 0 0 2000" # noqa: B950 - # /proc/101/io and /proc/101/smaps_rollup return None - return None - - mock_read.side_effect = read_side_effect - - result = _read_pgrp_stats() - self.assertEqual(result.disk_read_bytes, 0) - self.assertEqual(result.disk_write_bytes, 0) - self.assertIsNone(result.pss_bytes) - self.assertIsNone(result.private_bytes) - self.assertEqual(result.num_procs, 1) - - @patch(f"{_MODULE}.os.scandir") - @patch(f"{_MODULE}._read_file") - @patch(f"{_MODULE}.os.getpgrp") - def test_malformed_stat_skips_with_warning( - self, - mock_getpgrp: MagicMock, - mock_read: MagicMock, - mock_scandir: MagicMock, - ) -> None: - """A process with malformed stat is skipped and warns once.""" - mock_getpgrp.return_value = 1000 - - entries = [] - for name in ["101", "102"]: - entry = MagicMock() - entry.name = name - entries.append(entry) - mock_scandir.return_value = entries - - def read_side_effect(path: str) -> str | None: - if path == "/proc/101/stat": - return "malformed content" # no parens - if path == "/proc/102/stat": - return "102 (python3) S 1 1000 1000 0 -1 0 0 0 0 0 200 75 0 0 20 0 1 0 0 0 3000" # noqa: B950 - return None - - mock_read.side_effect = read_side_effect - - with self.assertLogs(_MODULE, level="WARNING") as cm: - result = _read_pgrp_stats() - - # Only pid 102 was counted - self.assertEqual(result.num_procs, 1) - self.assertTrue(any("missing closing paren" in m for m in cm.output)) - - -class ProcessGroupStatsMonitorClassTest(unittest.TestCase): - def test_is_background_task(self) -> None: - monitor = ProcessGroupStatsMonitor(callback=AsyncMock()) - self.assertIsInstance(monitor, BackgroundTask) - - -@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") -class ProcessGroupStatsMonitorSubprocessTest(unittest.TestCase): - def test_monitor_spawns_and_cancels_subprocess(self) -> None: - """Verify the monitor spawns a subprocess and terminates it on cancel.""" - mock_proc = MagicMock(spec=multiprocessing.Process) - mock_proc.pid = 12345 - mock_proc.is_alive.side_effect = [True, True, False] - - mock_ctx: MagicMock = MagicMock() - mock_ctx.Process.return_value = mock_proc - - async def run_monitor() -> None: - monitor = ProcessGroupStatsMonitor( - callback=AsyncMock(), mp_context=mock_ctx - ) - task = asyncio.create_task(monitor.run()) - await asyncio.sleep(0.01) - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - - asyncio.run(run_monitor()) - - mock_proc.start.assert_called_once() - mock_proc.terminate.assert_called_once() - mock_proc.join.assert_called() - - def test_monitor_warns_on_unexpected_exit(self) -> None: - """Verify warning is logged when the subprocess exits unexpectedly.""" - mock_proc = MagicMock(spec=multiprocessing.Process) - mock_proc.pid = 12345 - mock_proc.exitcode = 1 - mock_proc.is_alive.return_value = False - - mock_ctx: MagicMock = MagicMock() - mock_ctx.Process.return_value = mock_proc - - async def run_monitor() -> None: - monitor = ProcessGroupStatsMonitor( - callback=AsyncMock(), mp_context=mock_ctx - ) - await monitor.run() - - with self.assertLogs(_MODULE, level="WARNING") as cm: - asyncio.run(run_monitor()) - - exit_warnings = [m for m in cm.output if "exited unexpectedly" in m] - self.assertEqual(len(exit_warnings), 1) - self.assertIn("exit code 1", exit_warnings[0]) - - -@unittest.skipUnless(sys.platform == "linux", "Requires Linux /proc filesystem") -class PgrpMonitorSubprocessFunctionTest(unittest.TestCase): - @patch(f"{_MODULE}._collect_pgrp_stats") - def test_subprocess_function_collects_and_calls_callback( - self, - mock_collect: MagicMock, - ) -> None: - """Test the subprocess entry point invokes the callback.""" - usage = ProcessGroupResourceUsage( - pid=os.getpid(), - pgid=os.getpgrp(), - cpu_percent=50.0, - rss_bytes=1048576, - pss_bytes=800000, - private_bytes=600000, - disk_read_bytes=4096, - disk_write_bytes=8192, - num_procs=3, - net_rx_bytes=100, - net_tx_bytes=200, - ) - mock_collect.return_value = (usage, 500000, 1000000, 100, 200) - - mock_callback = AsyncMock() - - call_count = 0 - original_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep - - async def counting_sleep(delay: float) -> None: - nonlocal call_count - call_count += 1 - if call_count >= 2: - raise KeyboardInterrupt - await original_sleep(0) - - with ( - patch("asyncio.sleep", counting_sleep), - patch(f"{_MODULE}.signal.signal"), - ): - try: - _pgrp_monitor_subprocess(0.01, mock_callback) - except KeyboardInterrupt: - pass - - mock_collect.assert_called() - mock_callback.assert_called() - received = mock_callback.call_args[0][0] - self.assertIsInstance(received, ProcessGroupResourceUsage) - self.assertEqual(received.cpu_percent, 50.0) - self.assertEqual(received.rss_bytes, 1048576) - self.assertEqual(received.pss_bytes, 800000) - self.assertEqual(received.private_bytes, 600000) - self.assertEqual(received.disk_read_bytes, 4096) - self.assertEqual(received.disk_write_bytes, 8192) - self.assertEqual(received.num_procs, 3) - self.assertEqual(received.net_rx_bytes, 100) - self.assertEqual(received.net_tx_bytes, 200) diff --git a/tests/pipeline/pgrp_stats_test.py b/tests/pipeline/pgrp_stats_test.py new file mode 120000 index 000000000..af668b8f8 --- /dev/null +++ b/tests/pipeline/pgrp_stats_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pgrp_stats_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_builder_test.py b/tests/pipeline/pipeline_builder_test.py deleted file mode 100644 index d3f92bebd..000000000 --- a/tests/pipeline/pipeline_builder_test.py +++ /dev/null @@ -1,2763 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -import asyncio -import functools -import os -import platform -import random -import re -import sys -import threading -import time -import unittest -import warnings -from collections.abc import Iterator -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor -from contextlib import asynccontextmanager -from functools import partial -from multiprocessing import Process -from typing import TypeVar - -from parameterized import parameterized -from spdl.pipeline import ( - AsyncQueue, - PipelineBuilder, - PipelineFailure, - run_pipeline_in_subprocess, - TaskHook, - TaskStatsHook, -) -from spdl.pipeline._components import _get_global_id, _set_global_id -from spdl.pipeline._components._common import _EOF, StageInfo -from spdl.pipeline._components._hook import _periodic_dispatch -from spdl.pipeline._components._pipe import ( - _FailCounter, - _get_fail_counter, - _pipe, - _PipeArgs, -) -from spdl.pipeline._components._sink import _sink -from spdl.pipeline._components._source import _source -from spdl.pipeline.defs import Aggregator -from spdl.source.utils import embed_shuffle - -T = TypeVar("T") - - -def _ignore_warnings(*filters): - """Decorator that wraps a test in `warnings.catch_warnings()` and applies - the given filters. Each ``filter`` is a dict of kwargs forwarded to - ``warnings.filterwarnings``. - """ - - def decorator(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - with warnings.catch_warnings(): - for f in filters: - warnings.filterwarnings("ignore", **f) - return fn(*args, **kwargs) - - return wrapper - - return decorator - - -_FORK_WARNING = { - "message": ( - r"This process \(pid=\d+\) is multi-threaded, use of fork\(\) " - r"may lead to deadlocks in the child" - ), - "category": DeprecationWarning, -} - -_RUN_PIPELINE_DEPRECATION = { - "message": ( - r"Passing a `PipelineBuilder` object directly to " - r"`run_pipeline_in_subprocess` is now deprecated\..*" - ), - "category": UserWarning, -} - -_UNAWAITED_COROUTINE = { - "message": "coroutine .* was never awaited", - "category": RuntimeWarning, -} - - -def _SI(name: str) -> StageInfo: - """Shorthand for creating a test StageInfo.""" - return StageInfo(pipeline_id=0, stage_id="0", stage_name=name) - - -def _put_aqueue(queue, vals, *, eof): - for val in vals: - queue.put_nowait(val) - if eof: - queue.put_nowait(_EOF) - - -def _flush_aqueue(queue): - ret = [] - while not queue.empty(): - ret.append(queue.get_nowait()) - return ret - - -async def no_op(val): - return val - - -################################################################################ -# _source -################################################################################ - - -class TestSource(unittest.TestCase): - def test_async_enqueue_empty(self) -> None: - """_async_enqueue can handle empty iterator""" - queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 - ) - coro = _source([], queue) - asyncio.run(coro) - self.assertEqual(_flush_aqueue(queue), [_EOF]) - - def test_async_enqueue_simple(self) -> None: - """_async_enqueue should put the values in the queue.""" - src = list(range(6)) - queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 - ) - coro = _source(src, queue) - asyncio.run(coro) - vals = _flush_aqueue(queue) - self.assertEqual(vals, [*src, _EOF]) - - def test_async_enqueue_iterator_failure(self) -> None: - """When `iterator` fails, the exception is propagated.""" - - def src(): - yield from range(10) - raise RuntimeError("Failing the iterator.") - - coro = _source( - src(), - AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=0 - ), - ) - - with self.assertRaises(RuntimeError): - asyncio.run(coro) # Not raising - - def test_async_enqueue_cancel(self) -> None: - """_async_enqueue is cancellable.""" - - async def _test(): - queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), buffer_size=1 - ) - - src = list(range(3)) - - coro = _source(src, queue) - task = asyncio.create_task(coro) - - await asyncio.sleep(0.1) - - task.cancel() - - with self.assertRaises(asyncio.CancelledError): - await task - - asyncio.run(_test()) - - -################################################################################ -# _sink -################################################################################ - - -class TestSink(unittest.TestCase): - @parameterized.expand( - [ - (False,), - (True,), - ] - ) - def test_async_sink_simple(self, empty: bool) -> None: - """_sink pass the contents from input_queue to output_queue""" - input_queue: AsyncQueue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue: AsyncQueue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - data = [] if empty else list(range(3)) - _put_aqueue(input_queue, data, eof=True) - - coro = _sink(input_queue, output_queue) - - asyncio.run(coro) - results = _flush_aqueue(output_queue) - - self.assertEqual(results, data) - - def test_async_sink_cancel(self) -> None: - """_async_sink is cancellable.""" - - async def _test(): - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input") - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output") - ) - - coro = _sink(input_queue, output_queue) - task = asyncio.create_task(coro) - - await asyncio.sleep(0.1) - - task.cancel() - - with self.assertRaises(asyncio.CancelledError): - await task - - asyncio.run(_test()) - - -################################################################################ -# _pipe -################################################################################ - - -async def adouble(val: int): - return 2 * val - - -async def aplus1(val: int): - return val + 1 - - -async def passthrough(val): - print("passthrough:", val) - return val - - -class TestPipe(unittest.IsolatedAsyncioTestCase): - def test_async_pipe(self) -> None: - """_pipe processes the data in input queue and pass it to output queue.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - async def test(): - ref = list(range(6)) - _put_aqueue(input_queue, ref, eof=True) - - await _pipe( - _SI("adouble"), - input_queue, - output_queue, - _PipeArgs(op=adouble), - _FailCounter(), - [], - False, - ) - - result = _flush_aqueue(output_queue) - - self.assertEqual(result, [v * 2 for v in ref] + [_EOF]) - - asyncio.run(test()) - - def test_async_pipe_skip(self) -> None: - """_pipe skips the result if it's None.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - async def skip_even(v): - if v % 2: - return v - - async def test(): - _put_aqueue(input_queue, range(10), eof=True) - - await _pipe( - _SI("skip_even"), - input_queue, - output_queue, - _PipeArgs(op=skip_even), - _FailCounter(), - [], - False, - ) - - result = _flush_aqueue(output_queue) - - self.assertEqual(result, [*list(range(1, 10, 2)), _EOF]) - - asyncio.run(test()) - - def test_async_pipe_wrong_task_signature(self) -> None: - """_pipe fails immediately if user provided incompatible iterator/afunc.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=0 - ) - - async def _2args(val: int, _): - return val - - async def test(): - ref = list(range(6)) - _put_aqueue(input_queue, ref, eof=False) - - with self.assertRaises(TypeError): - await _pipe( - _SI("_2args"), - input_queue, - output_queue, - _PipeArgs(op=_2args, concurrency=3), - _FailCounter(), - [], - False, - ) - - remaining = _flush_aqueue(input_queue) - self.assertEqual(remaining, ref[1:]) - - result = _flush_aqueue(output_queue) - self.assertEqual(result, [_EOF]) - - asyncio.run(test()) - - @parameterized.expand( - [ - (False,), - (True,), - ] - ) - def test_async_pipe_cancel(self, full: bool) -> None: - """_pipe is cancellable.""" - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), buffer_size=0 - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), buffer_size=1 - ) - - _put_aqueue(input_queue, list(range(3)), eof=False) - - if full: - output_queue.put_nowait(None) - - cancelled = False - - async def astuck(i): - try: - await asyncio.sleep(10) - return i - except asyncio.CancelledError: - nonlocal cancelled - cancelled = True - raise - - async def test(): - coro = _pipe( - _SI("astuck"), - input_queue, - output_queue, - _PipeArgs(op=astuck), - _FailCounter(), - [], - False, - ) - task = asyncio.create_task(coro) - - await asyncio.sleep(0.5) - - task.cancel() - - with self.assertRaises(asyncio.CancelledError): - await task - - self.assertFalse(cancelled) - asyncio.run(test()) - self.assertTrue(cancelled) - - def test_async_pipe_concurrency(self) -> None: - """Changing concurrency changes the number of items fetched and processed.""" - - async def delay(val): - await asyncio.sleep(0.5) - return val - - async def test(concurrency): - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), - buffer_size=0, - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), - buffer_size=0, - ) - - ref = [1, 2, 3, 4] - _put_aqueue(input_queue, ref, eof=False) - - coro = _pipe( - _SI("delay"), - input_queue, - output_queue, - _PipeArgs( - op=delay, - concurrency=concurrency, - ), - _FailCounter(), - [], - False, - ) - - task = asyncio.create_task(coro) - await asyncio.sleep(0.8) - task.cancel() - - return _flush_aqueue(input_queue), _flush_aqueue(output_queue) - - # With concurrency==1, there should be - # 1 in output_queue, 2 is in flight, 3 and 4 remain in input_queue - remain, output = asyncio.run(test(1)) - self.assertEqual(remain, [3, 4]) - self.assertEqual(output, [1]) - - # With concurrency==4, there should be - # 1, 2, 3 and 4 in output_queue. - remain, output = asyncio.run(test(4)) - self.assertEqual(remain, []) - self.assertEqual(set(output), {1, 2, 3, 4}) - - def test_async_pipe_concurrency_throughput(self) -> None: - """increasing concurrency improves the throughput.""" - - async def delay(val): - await asyncio.sleep(0.5) - return val - - async def test(concurrency): - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input"), - buffer_size=0, - ) - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output"), - buffer_size=0, - ) - - ref = [4, 5, 6, 7, _EOF] - _put_aqueue(input_queue, ref, eof=False) - - t0 = time.monotonic() - await _pipe( - _SI("delay"), - input_queue, - output_queue, - _PipeArgs( - op=delay, - concurrency=concurrency, - ), - _FailCounter(), - [], - False, - ) - elapsed = time.monotonic() - t0 - - result = _flush_aqueue(output_queue) - - self.assertEqual(set(result), set(ref)) - self.assertEqual(result[-1], ref[-1]) - self.assertEqual(result[-1], _EOF) - - return elapsed - - elapsed1 = asyncio.run(test(1)) - elapsed4 = asyncio.run(test(4)) - - self.assertGreater(elapsed1, 1.8) - self.assertLess(elapsed4, 1) - - -################################################################################ -# Pipeline -################################################################################ - - -class TestPipeline(unittest.TestCase): - def test_pipeline_stage_hook_wrong_def1(self) -> None: - """Pipeline fails if stage_hook is not properly overrode.""" - - class _hook(TaskHook): - # missing asynccontextmanager - async def stage_hook(self): - yield - - @asynccontextmanager - async def task_hook(self, input_item=None): - yield - - with self.assertRaises(ValueError): - ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink() - # pyre-ignore - .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) - ) - - def test_pipeline_stage_hook_wrong_def2(self) -> None: - """Pipeline fails if task_hook is not properly overrode.""" - - class _hook(TaskHook): - # missing asynccontextmanager and async keyword - def stage_hook(self): - yield - - @asynccontextmanager - async def task_hook(self, input_item=None): - yield - - with self.assertRaises(ValueError): - ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink() - # pyre-ignore - .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) - ) - - -class CountHook(TaskHook): - def __init__(self): - self._enter_task_called = 0 - self._enter_stage_called = 0 - self._exit_task_called = 0 - self._exit_stage_called = 0 - - @asynccontextmanager - async def stage_hook(self): - self._enter_stage_called += 1 - yield - self._exit_stage_called += 1 - - @asynccontextmanager - async def task_hook(self, input_item=None): - self._enter_task_called += 1 - try: - yield - finally: - self._exit_task_called += 1 - - -class TestPipelineHook(unittest.TestCase): - @parameterized.expand( - [ - (False,), - (True,), - ] - ) - def test_pipeline_hook_drop_last(self, drop_last: bool) -> None: - """Hook is executed properly""" - - h1, h2, h3 = CountHook(), CountHook(), CountHook() - - def hook_factory(name) -> list[TaskHook]: - sname = str(name) - if "adouble" in sname: - return [h1] - if "aggregate" in sname: - return [h2] - if "_fail" in sname: - return [h3] - raise RuntimeError(f"Unexpected name: {sname}") - - async def _fail(_): - raise RuntimeError("Failing") - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(adouble) - .aggregate(5, drop_last=drop_last) - .pipe(_fail) - .add_sink(1000) - .build(num_threads=1, task_hook_factory=hook_factory) - ) - - with pipeline.auto_stop(): - self.assertEqual([], list(pipeline.get_iterator(timeout=10))) - - self.assertEqual(h1._enter_stage_called, 1) - self.assertEqual(h1._exit_stage_called, 1) - self.assertEqual(h1._enter_task_called, 10) - self.assertEqual(h1._exit_task_called, 10) - - self.assertEqual(h2._enter_stage_called, 1) - self.assertEqual(h2._exit_stage_called, 1) - # When drop_last=False, EOF is passed to the aggregation operator (11 calls: 10 items + 1 EOF) - # When drop_last=True, EOF is NOT passed to the aggregation operator (10 calls: 10 items only) - expected_h2_calls = 10 if drop_last else 11 - self.assertEqual(h2._enter_task_called, expected_h2_calls) - self.assertEqual(h2._exit_task_called, expected_h2_calls) - - # Even when the stage task fails, - # the exit_stage and exit_task are still called. - self.assertEqual(h3._enter_stage_called, 1) - self.assertEqual(h3._exit_stage_called, 1) - self.assertEqual(h3._enter_task_called, 2) - self.assertEqual(h3._exit_task_called, 2) - - def test_pipeline_hook_multiple(self) -> None: - """Multiple hooks are executed properly""" - - class _hook(TaskHook): - def __init__(self): - self._enter_task_called = 0 - self._enter_stage_called = 0 - self._exit_task_called = 0 - self._exit_stage_called = 0 - - @asynccontextmanager - async def stage_hook(self): - self._enter_stage_called += 1 - yield - self._exit_stage_called += 1 - - @asynccontextmanager - async def task_hook(self, input_item=None): - self._enter_task_called += 1 - try: - yield - finally: - self._exit_task_called += 1 - - hooks = [_hook(), _hook(), _hook()] - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: hooks) - ) - - with pipeline.auto_stop(): - self.assertEqual(list(range(10)), list(pipeline.get_iterator(timeout=10))) - - for h in hooks: - self.assertEqual(h._enter_stage_called, 1) - self.assertEqual(h._exit_stage_called, 1) - self.assertEqual(h._enter_task_called, 10) - self.assertEqual(h._exit_task_called, 10) - - @_ignore_warnings({"category": RuntimeWarning}) - @_ignore_warnings(_UNAWAITED_COROUTINE) - def test_pipeline_hook_failure_enter_stage(self) -> None: - """If enter_stage fails, the pipeline is aborted.""" - - class _enter_stage_fail(TaskHook): - @asynccontextmanager - async def stage_hook(self): - raise RuntimeError("failing") - - @asynccontextmanager - async def task_hook(self, input_item=None): - yield - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_enter_stage_fail()]) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(vals, []) - - @_ignore_warnings({"category": RuntimeWarning}) - @_ignore_warnings(_UNAWAITED_COROUTINE) - def test_pipeline_hook_failure_exit_stage(self) -> None: - """If exit_stage fails, the error is propagated to the front end.""" - - class _exit_stage_fail(TaskHook): - @asynccontextmanager - async def stage_hook(self): - yield - raise RuntimeError("failing") - - @asynccontextmanager - async def task_hook(self, input_item=None): - yield - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_exit_stage_fail()]) - ) - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(vals, list(range(10))) - - @_ignore_warnings({"category": RuntimeWarning}) - def test_pipeline_hook_failure_enter_task(self) -> None: - """If enter_task fails, the pipeline does not fail.""" - - class _hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - raise RuntimeError("failing enter_task") - - @asynccontextmanager - async def stage_hook(self, *_): - yield - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_hook()]) - ) - - with pipeline.auto_stop(): - self.assertEqual([], list(pipeline.get_iterator(timeout=10))) - - @_ignore_warnings({"category": RuntimeWarning}) - def test_pipeline_hook_failure_exit_task(self) -> None: - """If exit_task fails, the pipeline does not fail. - - IMPORTANT: The result is dropped. - """ - - class _exit_stage_fail(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - yield - raise RuntimeError("failing exit_task") - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_exit_stage_fail()]) - ) - - with pipeline.auto_stop(): - self.assertEqual(list(pipeline.get_iterator(timeout=10)), []) - - def test_pipeline_hook_exit_task_capture_error(self) -> None: - """If task fails exit_task captures the error.""" - - exc_info = None - - class _capture(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - try: - yield - except Exception as e: - nonlocal exc_info - exc_info = e - - err = RuntimeError("failing") - - async def _fail(_): - raise err - - pipeline = ( - PipelineBuilder() - .add_source([None]) - .pipe(_fail) - .add_sink(100) - .build( - num_threads=1, - # pyre-ignore[6] - task_hook_factory=lambda _: [_capture()], - ) - ) - - with pipeline.auto_stop(): - self.assertEqual(list(pipeline.get_iterator(timeout=10)), []) - - self.assertTrue(exc_info is err) - - def test_pipeline_hook_receives_input_item(self) -> None: - """task_hook receives the input_item being processed.""" - - received_items = [] - - class _item_capture_hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - received_items.append(input_item) - yield - - pipeline = ( - PipelineBuilder() - .add_source(range(5)) - .pipe(passthrough) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_item_capture_hook()]) - ) - - with pipeline.auto_stop(): - output = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(output, list(range(5))) - self.assertEqual(received_items, list(range(5))) - - def test_pipeline_hook_receives_input_item_on_failure(self) -> None: - """task_hook receives input_item even when the task fails.""" - - captured_items_on_failure = [] - - class _failure_capture_hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - try: - yield - except StopAsyncIteration: - raise - except Exception: - captured_items_on_failure.append(input_item) - raise - - def fail_on_even(x: int) -> int: - if x % 2 == 0: - raise RuntimeError(f"fail on {x}") - return x - - pipeline = ( - PipelineBuilder() - .add_source(range(6)) - .pipe(fail_on_even) - .add_sink(1000) - .build( - num_threads=1, - task_hook_factory=lambda _: [_failure_capture_hook()], # pyre-ignore[6] - ) - ) - - with pipeline.auto_stop(): - output = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(output, [1, 3, 5]) - self.assertEqual(captured_items_on_failure, [0, 2, 4]) - - def test_ordered_pipe_hook_receives_input_item(self) -> None: - """task_hook receives input_item in ordered pipe.""" - - received_items = [] - - class _item_capture_hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - received_items.append(input_item) - yield - - pipeline = ( - PipelineBuilder() - .add_source(range(5)) - .pipe(passthrough, output_order="input", concurrency=4) - .add_sink(1000) - # pyre-ignore[6] - .build(num_threads=4, task_hook_factory=lambda _: [_item_capture_hook()]) - ) - - with pipeline.auto_stop(): - output = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(output, list(range(5))) - self.assertEqual(received_items, list(range(5))) - - def test_ordered_pipe_hook_receives_input_item_on_failure(self) -> None: - """task_hook receives input_item on failure in ordered pipe.""" - - captured_items_on_failure = [] - - class _failure_capture_hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - try: - yield - except StopAsyncIteration: - raise - except Exception: - captured_items_on_failure.append(input_item) - raise - - def fail_on_even(x: int) -> int: - if x % 2 == 0: - raise RuntimeError(f"fail on {x}") - return x - - pipeline = ( - PipelineBuilder() - .add_source(range(6)) - .pipe(fail_on_even, output_order="input", concurrency=4) - .add_sink(1000) - .build( - num_threads=4, - task_hook_factory=lambda _: [_failure_capture_hook()], # pyre-ignore[6] - ) - ) - - with pipeline.auto_stop(): - output = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(output, [1, 3, 5]) - self.assertEqual(sorted(captured_items_on_failure), [0, 2, 4]) - - -################################################################################ -# TaskStatsHook -################################################################################ - - -class TestTaskStatsHook(unittest.TestCase): - def test_task_stats(self) -> None: - """TaskStatsHook logs the interval of each task.""" - - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), 1 - ) - - async def _test(): - async with hook.stage_hook(): - for _ in range(3): - async with hook.task_hook(): - await asyncio.sleep(0.5) - - self.assertEqual(hook.num_tasks, 3) - self.assertEqual(hook.num_success, 3) - self.assertGreater(hook.ave_time, 0.3) - self.assertLess(hook.ave_time, 0.7) - - for _ in range(2): - with self.assertRaises(RuntimeError): - async with hook.task_hook(): - await asyncio.sleep(1.0) - raise RuntimeError("failing") - - self.assertEqual(hook.num_tasks, 5) - self.assertEqual(hook.num_success, 3) - self.assertGreater(hook.ave_time, 0.45) - self.assertLess(hook.ave_time, 0.9) - - asyncio.run(_test()) - - -class TestPeriodicDispatch(unittest.TestCase): - def test_periodic_dispatch_smoke_test(self) -> None: - """_periodic_dispatch runs functions with the given interval.""" - - calls = [] - - async def afun(): - print("afun: ", time.time()) - calls.append(time.monotonic()) - - async def _test(): - done = asyncio.Event() - task = asyncio.create_task(_periodic_dispatch(afun, done, 1)) - - await asyncio.sleep(3.2) - - done.set() - await task - - print("start: ", time.time()) - asyncio.run(_test()) - - self.assertEqual(len(calls), 3) - self.assertGreater(calls[1] - calls[0], 0.9) - self.assertLess(calls[1] - calls[0], 1.1) - self.assertGreater(calls[2] - calls[1], 0.9) - self.assertLess(calls[2] - calls[1], 1.1) - - def test_task_stats_log_interval_stats(self) -> None: - """Smoke test for _log_interval_stats.""" - - hook = TaskStatsHook( - StageInfo(pipeline_id=0, stage_id="0", stage_name="foo"), 1 - ) - asyncio.run(hook._log_interval_stats()) - - -################################################################################ -# __str__ -################################################################################ - - -class TestPipelineStr(unittest.TestCase): - def test_pipeline_str_smoke(self) -> None: - async def passthrough(i): - return i - - builder = PipelineBuilder() - - print(builder) - - builder = builder.add_source(range(10)) - - print(builder) - - builder = builder.pipe(passthrough) - - print(builder) - - builder = builder.aggregate(1) - - print(builder) - - builder = builder.pipe(passthrough, output_order="input") - - print(builder) - - builder = builder.aggregate(1) - - print(builder) - - builder = builder.add_sink(100) - - print(builder) - - -################################################################################ -# AsyncPipeline - resume -################################################################################ - - -class TestPipelineResume(unittest.TestCase): - def test_pipeline_reiterate(self) -> None: - """Pipeline can be iterated multiple times as long as it's not stopped""" - - pipeline = ( - PipelineBuilder().add_source(range(20)).add_sink(1000).build(num_threads=1) - ) - - with pipeline.auto_stop(): - for i in range(5): - for j, val in enumerate(pipeline.get_iterator(timeout=10)): - self.assertEqual(val, (i * 4) + j) - - # Now it's empty - with self.assertRaises(StopIteration): - next(pipeline.get_iterator(timeout=10)) - - def test_pipeline_resume(self) -> None: - """AsyncPipeline can execute the source partially, then resumed""" - - # Note - # If we pass `range(10)` directly, new iterator is created at every run. - src = iter(range(10)) - - pipeline = PipelineBuilder().add_source(src).add_sink(1000).build(num_threads=1) - - with pipeline.auto_stop(): - iterator = pipeline.get_iterator(timeout=10) - self.assertEqual([0, 1], [next(iterator) for _ in range(2)]) - - iterator = pipeline.get_iterator(timeout=10) - self.assertEqual([2, 3, 4], [next(iterator) for _ in range(3)]) - - iterator = pipeline.get_iterator(timeout=10) - self.assertEqual([5, 6, 7, 8, 9], [next(iterator) for _ in range(5)]) - - with self.assertRaises(StopIteration): - next(iterator) - - def test_pipeline_infinite_loop(self) -> None: - """AsyncPipeline can execute infinite iterable""" - - def src(i=-1): - while True: - yield (i := i + 1) - - pipeline = ( - PipelineBuilder().add_source(src()).add_sink(1000).build(num_threads=1) - ) - - with pipeline.auto_stop(): - i = 0 - for _ in range(10): - num_items = random.randint(1, 128) - for j, item in enumerate(pipeline.get_iterator(timeout=10)): - self.assertEqual(item, i) - i += 1 - - if num_items == j: - break - - -################################################################################ -# AsyncPipeline - order -################################################################################ - - -class TestPipelineOrder(unittest.TestCase): - def test_pipeline_order_complete(self) -> None: - """The output is in the order of completion.""" - - async def _sleep(i): - await asyncio.sleep(i / 10) - return i - - src = list(reversed(range(10))) - pipeline = ( - PipelineBuilder() - .add_source(src) - .pipe(_sleep, concurrency=10, output_order="completion") - .add_sink(100) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - self.assertEqual(list(pipeline.get_iterator(timeout=10)), list(range(10))) - - def test_pipeline_order_input(self) -> None: - """The output is in the order of the input.""" - - async def _sleep(i): - print(f"Sleeping: {i}") - await asyncio.sleep(i / 10) - print(f"Returning: {i}") - return i - - src = list(reversed(range(10))) - pipeline = ( - PipelineBuilder() - .add_source(src) - .pipe(_sleep, concurrency=10, output_order="input") - .add_sink(100) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - self.assertEqual(src, list(pipeline.get_iterator(timeout=10))) - - def test_pipeline_order_input_sync_func(self) -> None: - """The output is in the order of the input.""" - - def _sleep(i): - print(f"Sleeping: {i}") - time.sleep(i / 10) - print(f"Returning: {i}") - return i - - src = list(reversed(range(10))) - pipeline = ( - PipelineBuilder() - .add_source(src) - .pipe(_sleep, concurrency=10, output_order="input") - .add_sink(100) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - self.assertEqual(list(pipeline.get_iterator(timeout=10)), src) - - @_ignore_warnings({"category": RuntimeWarning}) - def test_pipeline_order_input_filter_none(self) -> None: - """Ordered pipe filters out None values returned by the pipe operation.""" - - pipeline = ( - PipelineBuilder() - .add_source(list(range(10))) - .pipe(lambda x: None if x % 2 == 0 else x, output_order="input") - .add_sink(2) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(result, [1, 3, 5, 7, 9]) - - def test_pipeline_order_input_filter_none_async(self) -> None: - """Ordered pipe filters out None values with async function.""" - - async def filter_even(x): - await asyncio.sleep(0.01) - return None if x % 2 == 0 else x - - pipeline = ( - PipelineBuilder() - .add_source(list(range(10))) - .pipe(filter_even, output_order="input", concurrency=3) - .add_sink(2) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(result, [1, 3, 5, 7, 9]) - - def test_pipeline_order_input_filter_none_with_concurrency(self) -> None: - """Ordered pipe filters out None values with high concurrency.""" - - def slow_filter(x): - time.sleep(0.05) - return None if x % 3 == 0 else x - - pipeline = ( - PipelineBuilder() - .add_source(list(range(15))) - .pipe(slow_filter, output_order="input", concurrency=5) - .add_sink(10) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=5)) - # Filters out 0, 3, 6, 9, 12 - self.assertEqual(result, [1, 2, 4, 5, 7, 8, 10, 11, 13, 14]) - - def test_pipeline_order_input_all_none(self) -> None: - """Ordered pipe handles case where all values are None.""" - - pipeline = ( - PipelineBuilder() - .add_source(list(range(5))) - .pipe(lambda _: None, output_order="input") - .add_sink(2) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(result, []) - - def test_pipeline_order_input_mixed_none_and_values(self) -> None: - """Ordered pipe correctly handles mixed None and values in specific pattern.""" - - def pattern_filter(x): - if x < 2: - return None - if x < 5: - return x - if x < 7: - return None - return x - - pipeline = ( - PipelineBuilder() - .add_source(list(range(10))) - .pipe(pattern_filter, output_order="input") - .add_sink(5) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - result = list(pipeline.get_iterator(timeout=10)) - # Returns 2, 3, 4 (x < 5), and 7, 8, 9 (x >= 7) - self.assertEqual(result, [2, 3, 4, 7, 8, 9]) - - -################################################################################ -# AsyncPipeline2 -################################################################################ - - -class TestPipelineNoop(unittest.TestCase): - def test_pipeline_noop(self) -> None: - """AsyncPipeline2 functions without pipe.""" - - apl = PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) - - with apl.auto_stop(): - for i in range(10): - print("fetching", i) - self.assertEqual(i, apl.get_item(timeout=1)) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelinePassthrough(unittest.TestCase): - def test_pipeline_passthrough(self) -> None: - """AsyncPipeline2 can passdown items operation.""" - - apl = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1) - .build(num_threads=1) - ) - - with apl.auto_stop(): - for i in range(10): - print("fetching", i) - self.assertEqual(i, apl.get_item(timeout=1)) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelineSkip(unittest.TestCase): - def test_pipeline_skip(self) -> None: - """AsyncPipeline2 does not output None items.""" - - src = list(range(10)) - - async def odd(i): - if i % 2: - return i - - pipeline = ( - PipelineBuilder() - .add_source(src) - .pipe(odd) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for i in range(5): - self.assertEqual(i * 2 + 1, pipeline.get_item(timeout=10)) - - with self.assertRaises(EOFError): - pipeline.get_item(timeout=10) - - -class TestPipelineLambda(unittest.TestCase): - def test_pipeline_lambda(self) -> None: - """AsyncPipeline2 pipe supports lambda items operation.""" - - apl = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(lambda x: x) - .add_sink(1) - .build(num_threads=1) - ) - - with apl.auto_stop(): - for i in range(10): - print("fetching", i) - self.assertEqual(i, apl.get_item(timeout=1)) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelineSimple(unittest.TestCase): - def test_pipeline_simple(self) -> None: - """AsyncPipeline2 can perform simple operation.""" - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(adouble) - .pipe(aplus1) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - for i, item in enumerate(pipeline.get_iterator(timeout=10)): - self.assertEqual(item, i * 2 + 1) - - -class TestPipelineAggregate(unittest.TestCase): - def test_pipeline_aggregate(self) -> None: - """AsyncPipeline aggregates the input""" - - src = list(range(13)) - - pipeline = ( - PipelineBuilder() - .add_source(src) - .aggregate(4) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual( - results, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12]] - ) - - def test_pipeline_aggregate_drop_last(self) -> None: - """AsyncPipeline aggregates the input and drop the last""" - - src = list(range(13)) - - pipeline = ( - PipelineBuilder() - .add_source(src) - .aggregate(4, drop_last=True) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(results, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]) - - @parameterized.expand([(False,), (True,)]) - def test_pipeline_aggregate_custom_op(self, drop_last: bool) -> None: - """AsyncPipeline aggregates with custom operation that concatenates when threshold is exceeded""" - - # Custom aggregation: concatenate strings when total size exceeds threshold - class CustomAggregator(Aggregator): - def __init__(self, size_threshold: int = 10): - self.size_threshold = size_threshold - self.buffer: list[str] = [] - self.total_size = 0 - - def _flush(self) -> str: - result = "".join(self.buffer) - self.buffer = [] - self.total_size = 0 - return result - - def flush(self) -> str | None: - # Emit remaining buffer when EOF is reached - if self.buffer: - return self._flush() - return None # _SKIP - - def accumulate(self, item: str) -> str | None: - self.buffer.append(item) - self.total_size += len(item) - - if self.total_size >= self.size_threshold: - return self._flush() - return None # _SKIP - - src = ["a", "bb", "ccc", "dddd", "e", "ff", "ggg", "h"] - - pipeline = ( - PipelineBuilder() - .add_source(src) - .aggregate(CustomAggregator(size_threshold=10), drop_last=drop_last) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - # "a", "bb", "ccc", "dddd" = 10 chars -> first result - if drop_last: - # When drop_last=True, flush is not called, - # so the remaining buffer ["e", "ff", "ggg", "h"] is dropped - self.assertEqual(results, ["abbcccdddd"]) - else: - # When drop_last=False, flush is called, - # so remaining buffer is emitted: "e", "ff", "ggg", "h" - self.assertEqual(results, ["abbcccdddd", "effgggh"]) - - -class TestPipelineDisaggregate(unittest.TestCase): - def test_pipeline_disaggregate(self) -> None: - """AsyncPipeline disaggregates the input""" - - src = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12]] - - pipeline = ( - PipelineBuilder() - .add_source(src) - .disaggregate() - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(results, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]) - - -class TestPipelineSource(unittest.TestCase): - def test_pipeline_source_failure(self) -> None: - """AsyncPipeline continues when source fails. - - note: the front end will propagate the error at the end of the `stop`. - before that, the pipeline should continue functioning. - """ - - def failing_range(i): - yield from range(i) - raise ValueError("Iterator failed") - - pipeline = ( - PipelineBuilder() - .add_source(failing_range(10)) - .pipe(adouble) - .pipe(aplus1) - .add_sink(1000) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, [1 + 2 * i for i in range(10)]) - - -class TestPipelineType(unittest.TestCase): - def test_pipeline_type_error(self) -> None: - """AsyncPipeline immediately fails if pipe function has wrong signature""" - - async def wrong_sig(i, _): - return i - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(wrong_sig) - .add_sink(1000) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(vals, []) - - -class TestPipelineTask(unittest.TestCase): - def test_pipeline_task_failure(self) -> None: - """AsyncPipeline is robust against task-level failure.""" - - async def areject_m3(i): - if i % 3 == 0: - raise ValueError(f"Multiple of 3 is prohibited: {i}") - return i - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(areject_m3) - .pipe(adouble) - .pipe(aplus1) - .add_sink(1000) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=10)) - self.assertEqual(results, [1 + 2 * i for i in range(10) if i % 3]) - - -class TestPipelineCancel(unittest.TestCase): - def test_pipeline_cancel_empty(self) -> None: - """AsyncPipeline2 can be cancelled while it's blocked on the pipeline.""" - - apl = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1) - .build(num_threads=1) - ) - # The nuffer and coroutinues will be blocked holding items as follows: - # - # [src] --|queue(1)|--> [passthrough] --|queue(1)|--> [sink] -->|queue(1)| - # i+5 i+4 i+3 i+2 i+1 i - # - - with apl.auto_stop(): - for i in range(5): - print("fetching", i) - self.assertEqual(i, apl.get_item(timeout=1)) - - # Ensure that buffers are filled and the pipeline is blocked. - time.sleep(0.1) - # At this point, the output queue holds 5. - - # Only the "5" is retrievable. - self.assertEqual(5, apl.get_item(timeout=1)) - - # The background thread is stopped, so no more data is coming. - for _ in range(3): - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelineFail(unittest.TestCase): - def test_pipeline_fail_middle(self) -> None: - """When a stage in the middle fails, downstream stages are not failing.""" - - async def fail(i, _): - return i - - class PassthroughWithCache: - def __init__(self): - self.cache = [] - - async def __call__(self, i): - self.cache.append(i) - return i - - pwc = PassthroughWithCache() - - apl = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .pipe(fail) - .pipe(pwc) - .add_sink(1) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with apl.auto_stop(): - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - self.assertEqual(pwc.cache, []) - - # The background thread is stopped, and the output queue is empty. - for _ in range(3): - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelineEof(unittest.TestCase): - def test_pipeline_eof_stop(self) -> None: - """APL2 can be closed after reaching EOF.""" - apl = ( - PipelineBuilder() - .add_source(range(2)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - with apl.auto_stop(): - for i in range(2): - print("fetching", i) - self.assertEqual(i, apl.get_item(timeout=1)) - - with self.assertRaises(EOFError): - apl.get_item(timeout=1) - - -class TestPipelineIterator(unittest.TestCase): - def test_pipeline_iterator(self) -> None: - """Can iterate the pipeline.""" - - apl = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(passthrough) - .add_sink(1) - .build(num_threads=1) - ) - - with apl.auto_stop(): - for i, item in enumerate(apl.get_iterator(timeout=1)): - print(i, item) - self.assertEqual(i, item) - - -class TestPipelineIter(unittest.TestCase): - def test_pipeline_iter_and_next(self) -> None: - """Pipeline `iter` and `next` fulfill the following contracts - - 1. `iter(pipeline)` creates a new iterator. - 2. `next(iterator)` gives the next item in the pipeline. - 3. one ca call `next` multiple times on an iterator. - 4. If iterator is exhausted, it raises `StopIteration`. - 5. Iterator object can be discarded without an side-effect by itself. - (Other side-effects might be happening but it's independent from iterator) - 6. Multiple instances of iterators can be created by - repeatedly calling `iter(pipeline)`. - 7. We do not define the beheviors for multiple iterators exist at the same - time. We only consider the case where one iterator is created and - discarded, then another is created. - 8. A new iterator should return items that are sequel to items generated by - the previous iterator. - """ - - apl = PipelineBuilder().add_source(range(12)).add_sink(1).build(num_threads=1) - - with apl.auto_stop(): - iterator = iter(apl) - self.assertEqual(next(iterator), 0) - self.assertEqual(next(iterator), 1) - self.assertEqual(next(iterator), 2) - - iterator = iter(apl) - self.assertEqual(next(iterator), 3) - self.assertEqual(next(iterator), 4) - self.assertEqual(next(iterator), 5) - - iterator = iter(apl) - self.assertEqual(next(iterator), 6) - self.assertEqual(next(iterator), 7) - self.assertEqual(next(iterator), 8) - - iterator = iter(apl) - self.assertEqual(next(iterator), 9) - self.assertEqual(next(iterator), 10) - self.assertEqual(next(iterator), 11) - - iterator = iter(apl) - with self.assertRaises(StopIteration): - next(iterator) - - -class TestPipelineStuck(unittest.TestCase): - def test_pipeline_stuck(self) -> None: - """`get_item` waits for slow pipeline.""" - - async def delay(i): - print(f"Sleeping: {i}") - await asyncio.sleep(0.5) - print(f"Sleeping: {i} - done") - return i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(delay) - .add_sink(1) - .build(num_threads=1) - ) - - with apl.auto_stop(): - for i, item in enumerate(apl.get_iterator(timeout=10)): - print(i, item) - self.assertEqual(i, item) - - -class TestPipelinePipe(unittest.TestCase): - def test_pipeline_pipe_agen(self) -> None: - """pipe works with async generator function""" - - async def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(expected, output) - - def test_pipeline_pipe_sync_gen(self) -> None: - """pipe works with sync generator function""" - - def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(expected, output) - - -class TestCallableGenerator(unittest.TestCase): - def test_callable_generator(self) -> None: - """pipe works with sync callable class returning generator""" - - class DupIncrement: - def __init__(self) -> None: - pass - - def __call__(self, v: int) -> Iterator[int]: - for i in range(3): - yield v + i - - dup_increment = DupIncrement() - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(expected, output) - - def test_pipeline_pipe_agen_max_failures(self) -> None: - """pipe works with async generator function and max_failure""" - - async def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1, max_failures=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(output, expected) - - def test_pipeline_pipe_gen(self) -> None: - """pipe works with sync generator function""" - - def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(output, expected) - - def test_pipeline_pipe_gen_max_failures(self) -> None: - """pipe works with sync generator function and max_failure""" - - def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1, max_failures=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(output, expected) - - @unittest.skipIf( - platform.system() == "Darwin" and "CI" in os.environ, - reason="GitHub macOS CI is not timely enough.", - ) - def test_pipeline_pipe_gen_incremental(self) -> None: - """pipe returns output of generator function immediately if not in ProcessPoolExecutor""" - - # We introduce delay in each iteration, so that, if the pipeline is returning - # the yielded value immediately, the output will be obtained quickly. - # If the pipeline is not returning the yielded value immediately, the output - # won't be available until the iteration ends, and by that time - # the foreground pipeline should timeout. - def dup_increment(v): - for i in range(3): - time.sleep(0.1) - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - .build(num_threads=1) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=0.2)) - self.assertEqual(output, expected) - - def test_pipeline_pipe_agen_wrong_hook(self) -> None: - """pipe works with async generator function, even when hook abosrb the StopAsyncIteration""" - - class _Hook(TaskHook): - @asynccontextmanager - async def task_hook(self, input_item=None): - try: - yield - except StopAsyncIteration: - pass - - async def dup_increment(v): - for i in range(3): - yield v + i - - apl = ( - PipelineBuilder() - .add_source(range(3)) - .pipe(dup_increment) - .add_sink(1) - # pyre-ignore[6] - .build(num_threads=1, task_hook_factory=lambda _: [_Hook()]) - ) - - expected = [0, 1, 2, 1, 2, 3, 2, 3, 4] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(output, expected) - - def test_pipeline_source_agen(self) -> None: - """source works with async generator function""" - - async def source(): - for i in range(3): - yield i - - apl = PipelineBuilder().add_source(source()).add_sink(1).build(num_threads=1) - - expected = [0, 1, 2] - with apl.auto_stop(): - output = list(apl.get_iterator(timeout=10)) - self.assertEqual(output, expected) - - -class TestPipelineStart(unittest.TestCase): - def test_pipeline_start_multiple_times(self) -> None: - """`Pipeline.start` cannot be called multiple times.""" - - pipeline = ( - PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) - ) - - with pipeline.auto_stop(): - with self.assertRaises(RuntimeError): - pipeline.start() - - -class TestPipelineStop(unittest.TestCase): - def test_pipeline_stop_multiple_times(self) -> None: - """`Pipeline.stop` can be called multiple times.""" - - pipeline = ( - PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) - ) - - pipeline.stop() - pipeline.stop() - pipeline.stop() - - with pipeline.auto_stop(): - pipeline.stop() - pipeline.stop() - pipeline.stop() - - pipeline.stop() - pipeline.stop() - pipeline.stop() - - -def _run_pipeline_without_closing(): - pipeline = PipelineBuilder().add_source(range(10)).add_sink(1).build(num_threads=1) - pipeline.start() - - -def get_pid(_): - import os - - time.sleep(0.5) - pid = os.getpid() - print(f"{pid=}") - return pid - - -def _range(item): - print(item) - for i in range(item): - print(f"yielding {item} - {i}") - yield i - - -class TestPipelineNo(unittest.TestCase): - def test_pipeline_no_close(self) -> None: - """Python interpreter can terminate even when Pipeline is not explicitly closed.""" - - p = Process(target=_run_pipeline_without_closing) - p.start() - p.join(timeout=10) - - if p.exitcode is None: - p.kill() - raise RuntimeError("Process did not self-terminate.") - - -class TestPipelineCustom(unittest.TestCase): - def test_pipeline_custom_pipe_executor(self) -> None: - """`pipe` accepts custom ThreadPoolExecutor. - - The primal goal of custom executor is to make it easy to use - thread local storages. - - So in this test, we initialize a custom executor with some thread - local storages, and then we access it without any check (hasattr) - in pipe function. - """ - num_threads = 10 - sleep = 0.5 - - ref = set(range(num_threads)) - - ref_copy = ref.copy() - thread_local_storage = threading.local() - - def init_storage(): - print("Initializing thread:", threading.get_ident()) - thread_local_storage.value = ref_copy.pop() - - executor = ThreadPoolExecutor( - max_workers=num_threads, - initializer=init_storage, - ) - - def op(i: int) -> int: - # sleep to block this thread, so that - # we use all the threads in the pool - time.sleep(sleep) - print(i, thread_local_storage.value) - return thread_local_storage.value - - pipeline = ( - PipelineBuilder() - .add_source(range(num_threads)) - .pipe(op, executor=executor, concurrency=num_threads) - .add_sink(1) - .build(num_threads=1) - ) - - t0 = time.monotonic() - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - elapsed = time.monotonic() - t0 - - self.assertEqual(0, len(ref_copy)) - self.assertEqual(ref, set(vals)) - self.assertLess(elapsed, sleep * 2) - - def test_pipeline_custom_pipe_executor_process(self) -> None: - """`pipe` accepts custom ProcessPoolExecutor.""" - num_processes = 5 - - executor = ProcessPoolExecutor(max_workers=num_processes) - - pipeline = ( - PipelineBuilder() - .add_source(range(num_processes)) - .pipe(get_pid, executor=executor, concurrency=num_processes) - .add_sink(1) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(num_processes, len(set(vals))) - - def test_pipeline_custom_pipe_executor_process_generator(self) -> None: - """`pipe` accepts custom ProcessPoolExecutor and generator function.""" - executor = ProcessPoolExecutor(max_workers=1) - - pipeline = ( - PipelineBuilder() - .add_source(range(4)) - .pipe(_range, executor=executor, concurrency=1) - .add_sink(1) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual([0, 0, 1, 0, 1, 2], vals) - - def test_pipeline_custom_pipe_executor_async(self) -> None: - """pipe rejects custom executor if op is async""" - - async def op(i: int) -> int: - return i - - with self.assertRaises(ValueError): - PipelineBuilder().add_source(range(10)).pipe( - op, executor=ThreadPoolExecutor() - ).add_sink(1).build(num_threads=1) - - def test_pipeline_pipe_list(self) -> None: - """pipe supports list as op.""" - - op = [i + 1 for i in range(10)] - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - # pyre-ignore[6] - .pipe(op) - .add_sink(1) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual(op, vals) - - def test_pipeline_pipe_tuple(self) -> None: - """pipe supports list as op.""" - - op = tuple(i + 1 for i in range(10)) - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - # pyre-ignore[6] - .pipe(op) - .add_sink(1) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual([i + 1 for i in range(10)], vals) - - def test_pipeline_pipe_dict(self) -> None: - """pipe supports dict as op.""" - - op = {i: i + 1 for i in range(10)} - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - # pyre-ignore[6] - .pipe(op) - .add_sink(1) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual([i + 1 for i in range(10)], vals) - - -class _PicklableSource: - def __init__(self, n: int) -> None: - self.n = n - - def __iter__(self) -> Iterator[int]: - yield from range(self.n) - - -class _ValidatePipelineId: - def __init__(self, val: int) -> None: - self.val = val - - def __iter__(self) -> Iterator[int]: - if (v := _get_global_id()) != self.val: - raise AssertionError(f"_node._PIPELINE_ID={v} != {self.val=}") - yield 0 - - -def plusN(x: int, N: int) -> int: - return x + N - - -def hook_factory(_: StageInfo) -> list[TaskHook]: - return [CountHook()] - - -class TestPipelinebuilderPicklable(unittest.TestCase): - @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) - def test_pipelinebuilder_picklable(self) -> None: - """PipelineBuilder can be passed to subprocess (==picklable)""" - - builder = ( - PipelineBuilder() - .add_source(_PicklableSource(10)) - .pipe( - adouble, - concurrency=3, - ) - .pipe( - aplus1, - concurrency=3, - ) - .pipe(partial(plusN, N=3)) - .pipe(passthrough) - .aggregate(3) - .disaggregate() - .add_sink(10) - ) - - results = list( - run_pipeline_in_subprocess( - # pyre-ignore[6] - builder, - num_threads=5, - buffer_size=-1, - task_hook_factory=hook_factory, - ) - ) - - def _ref(x: int) -> int: - return 2 * x + 1 + 3 - - self.assertEqual([_ref(i) for i in range(10)], sorted(results)) - - -class TestFailureCounter(unittest.TestCase): - def test_failure_counter_global_countes(self) -> None: - """_get_fail_counter creates _FailCounter subclass with different class valiable""" - - FC1 = _get_fail_counter() - FC2 = _get_fail_counter() - - fc1_1 = FC1(-1, -1) - fc1_2 = FC1(-1, -1) - - fc2_1 = FC2(-1, -1) - fc2_2 = FC2(-1, -1) - - self.assertEqual(0, _FailCounter._num_global_failures) - self.assertEqual(0, FC1._num_global_failures) - self.assertEqual(0, FC2._num_global_failures) - self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) - self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) - self.assertEqual(0, fc1_1._num_global_failures) - self.assertEqual(0, fc1_2._num_global_failures) - self.assertEqual(0, fc1_1._num_stage_failures) - self.assertEqual(0, fc1_2._num_stage_failures) - self.assertEqual(0, fc2_1._num_global_failures) - self.assertEqual(0, fc2_2._num_global_failures) - self.assertEqual(0, fc2_1._num_stage_failures) - self.assertEqual(0, fc2_2._num_stage_failures) - - fc1_1.__class__._num_global_failures += 1 - fc1_1._num_stage_failures += 1 - - self.assertEqual(0, _FailCounter._num_global_failures) - self.assertEqual(1, FC1._num_global_failures) - self.assertEqual(0, FC2._num_global_failures) - self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) - self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) - self.assertEqual(1, fc1_1._num_global_failures) - self.assertEqual(1, fc1_2._num_global_failures) - self.assertEqual(1, fc1_1._num_stage_failures) - self.assertEqual(0, fc1_2._num_stage_failures) - self.assertEqual(0, fc2_1._num_global_failures) - self.assertEqual(0, fc2_2._num_global_failures) - self.assertEqual(0, fc2_1._num_stage_failures) - self.assertEqual(0, fc2_2._num_stage_failures) - - fc1_1.__class__._num_global_failures += 1 - fc1_1._num_stage_failures += 1 - - self.assertEqual(0, _FailCounter._num_global_failures) - self.assertEqual(2, FC1._num_global_failures) - self.assertEqual(0, FC2._num_global_failures) - self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) - self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) - self.assertEqual(2, fc1_1._num_global_failures) - self.assertEqual(2, fc1_2._num_global_failures) - self.assertEqual(2, fc1_1._num_stage_failures) - self.assertEqual(0, fc1_2._num_stage_failures) - self.assertEqual(0, fc2_1._num_global_failures) - self.assertEqual(0, fc2_2._num_global_failures) - self.assertEqual(0, fc2_1._num_stage_failures) - self.assertEqual(0, fc2_2._num_stage_failures) - - fc1_2.__class__._num_global_failures += 1 - fc1_2._num_stage_failures += 1 - - self.assertEqual(0, _FailCounter._num_global_failures) - self.assertEqual(3, FC1._num_global_failures) - self.assertEqual(0, FC2._num_global_failures) - self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) - self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) - self.assertEqual(3, fc1_1._num_global_failures) - self.assertEqual(3, fc1_2._num_global_failures) - self.assertEqual(2, fc1_1._num_stage_failures) - self.assertEqual(1, fc1_2._num_stage_failures) - self.assertEqual(0, fc2_1._num_global_failures) - self.assertEqual(0, fc2_2._num_global_failures) - self.assertEqual(0, fc2_1._num_stage_failures) - self.assertEqual(0, fc2_2._num_stage_failures) - - fc2_1.__class__._num_global_failures += 1 - fc2_1._num_stage_failures += 1 - - self.assertEqual(0, _FailCounter._num_global_failures) - self.assertEqual(3, FC1._num_global_failures) - self.assertEqual(1, FC2._num_global_failures) - self.assertTrue(fc1_1._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc1_2._num_global_failures is FC1._num_global_failures) - self.assertTrue(fc2_1._num_global_failures is FC2._num_global_failures) - self.assertTrue(fc2_2._num_global_failures is FC2._num_global_failures) - self.assertEqual(3, fc1_1._num_global_failures) - self.assertEqual(3, fc1_2._num_global_failures) - self.assertEqual(2, fc1_1._num_stage_failures) - self.assertEqual(1, fc1_2._num_stage_failures) - self.assertEqual(1, fc2_1._num_global_failures) - self.assertEqual(1, fc2_2._num_global_failures) - self.assertEqual(1, fc2_1._num_stage_failures) - self.assertEqual(0, fc2_2._num_stage_failures) - - -class TestPipelineMax(unittest.TestCase): - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_max_failures(self, output_order: str) -> None: - """max_failures stop the pipeline.""" - - def fail_odd(x): - if x % 2: - raise ValueError(f"Only evan numbers are allowed. {x}") - return x - - builder = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_odd, output_order=output_order) - .add_sink(1) - ) - - pipeline = builder.build(num_threads=1) - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual([0, 2, 4, 6, 8], vals) - - pipeline = builder.build(num_threads=1, max_failures=3) - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - - self.assertEqual([0, 2, 4, 6], vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_max_failures_multiple_pipeline(self, output_order: str) -> None: - """When using multiple pipelines with different error caps, they work - - Note: FailCounter uses class method to combines the errors - from the different stages. We create a different child class - for each pipeline construction so that class counter are separate for - each pipeline object. This test ensures that. - """ - - def fail_odd(x): - if x % 2: - raise ValueError(f"Only evan numbers are allowed. {x}") - return x - - src = range(10) - - builder = ( - PipelineBuilder() - .add_source(src) - .pipe(fail_odd, output_order=output_order) - .add_sink(1) - ) - - pipeline1 = builder.build(num_threads=1, max_failures=2) - pipeline2 = builder.build(num_threads=1, max_failures=3) - - with self.assertRaises(PipelineFailure): - with pipeline2.auto_stop(): - vals = list(pipeline2.get_iterator(timeout=10)) - - self.assertEqual([0, 2, 4, 6], vals) - - with self.assertRaises(PipelineFailure): - with pipeline1.auto_stop(): - vals = list(pipeline1.get_iterator(timeout=10)) - - self.assertEqual([0, 2, 4], vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_max_failures_pipe_override_strict( - self, output_order: str - ) -> None: - """max_failures at pipe overrides the global threshold.""" - - def fail_odd(x): - if x % 2: - raise ValueError(f"Only evan numbers are allowed. {x}") - return x - - builder = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_odd, output_order=output_order, max_failures=2) - .add_sink(1) - ) - - pipeline = builder.build(num_threads=1, max_failures=-1) - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual([0, 2, 4], vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_max_failures_pipe_override_loose(self, output_order: str) -> None: - """max_failures at pipe overrides the global threshold.""" - - def fail_odd(x): - if x % 2: - raise ValueError(f"Only evan numbers are allowed. {x}") - return x - - builder = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_odd, output_order=output_order, concurrency=3, max_failures=-1) - .add_sink(1) - ) - - pipeline = builder.build(num_threads=1, max_failures=2) - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual([0, 2, 4, 6, 8], vals) - - @parameterized.expand( - [ - ("completion",), - ("input",), - ] - ) - def test_pipeline_max_failures_pipe_override_multiple( - self, output_order: str - ) -> None: - """max_failures at pipe overrides the global threshold.""" - - # Remove odd values - def fail_odd(x): - if x % 2: - raise ValueError(f"Only evan numbers are allowed. {x}") - return x - - # Remove multiplier of 6s - def fail_six(x): - if (x % 6) == 0: - raise ValueError(f"Values divisible by 6 are not allowed. {x}") - return x - - builder = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(fail_odd, output_order=output_order, max_failures=-1) - .pipe(fail_six, output_order=output_order, max_failures=3) - .add_sink(1) - ) - - # fail_odd fails more often, but it is allowed to fail any number of times. - # fail_six fails less often, but at the fourth failure (18), - # it should shutdown the pipeline. - - pipeline = builder.build(num_threads=1, max_failures=2) - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - vals = list(pipeline.get_iterator(timeout=10)) - self.assertEqual([2, 4, 8, 10, 14, 16], vals) - - -class TestPipelinePropagate(unittest.TestCase): - def test_pipeline_propagate_source_failure(self) -> None: - """When source itrator fails, the exception is propagated to the front end""" - - def failure_source(): - raise RuntimeError("Foo") - yield None - - pipeline = ( - PipelineBuilder() - .add_source(failure_source()) - .add_sink() - .build(num_threads=1) - ) - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - pass - - -class TestIterableWithShuffle: - __test__ = False - - def __init__(self, n: int) -> None: - self.vals = list(range(n)) - self.seed = 0 - - def shuffle(self, *, seed: int) -> None: - self.vals = self.vals[1:] + self.vals[:1] - self.seed = seed - - def __iter__(self) -> Iterator[int]: - yield from self.vals - - -class TestRunPipeline(unittest.TestCase): - @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) - def test_run_pipeline_in_subprocess_state(self) -> None: - """The status of the source is maintained and propagated properly in subprocess""" - n = 5 - - # pyre-ignore[6] - src = embed_shuffle(TestIterableWithShuffle(n)) - builder = PipelineBuilder().add_source(src).add_sink() - # pyre-ignore[6] - iterable = run_pipeline_in_subprocess(builder, num_threads=1) - - self.assertEqual([1, 2, 3, 4, 0], list(iterable)) - self.assertEqual([2, 3, 4, 0, 1], list(iterable)) - self.assertEqual([3, 4, 0, 1, 2], list(iterable)) - - # since the src is copied to the subprocess iterating it yields the original state - - self.assertEqual([1, 2, 3, 4, 0], list(src)) - # pyre-ignore[16] - self.assertEqual(0, src.src.seed) - self.assertEqual([2, 3, 4, 0, 1], list(src)) - self.assertEqual(1, src.src.seed) - self.assertEqual([3, 4, 0, 1, 2], list(src)) - self.assertEqual(2, src.src.seed) - - @_ignore_warnings(_RUN_PIPELINE_DEPRECATION, _FORK_WARNING, _UNAWAITED_COROUTINE) - def test_run_pipeline_in_subprocess_pipeline_id(self) -> None: - """The pipeline construdted in a subprocess inherits the global ID from the main process""" - - # Set to a number that's not zero and something unlikely to happen during the testing - _set_global_id(123456) - ref = _get_global_id() + 1 - - builder = PipelineBuilder().add_source(_ValidatePipelineId(ref)).add_sink() - - # pyre-ignore[6] - iterable = run_pipeline_in_subprocess(builder, num_threads=1) - - for _ in iterable: - pass - - -class TestOverrideStage(unittest.TestCase): - @_ignore_warnings(_UNAWAITED_COROUTINE) - def test_override_stage_id(self) -> None: - """Providing `stage_id` overrides the index of stages.""" - ref = 12345 - - class CheckNameQueue(AsyncQueue): - index = ref - - def __init__(self, name, *, buffer_size: int = 1) -> None: - print(name) - id = re.match(r"\d+:(\d+):.*", str(name)).group(1) - assert id == str(self.index) - CheckNameQueue.index += 1 - super().__init__(name, buffer_size=buffer_size) - - ( - PipelineBuilder() - .add_source(range(10)) - .pipe(lambda x: x) - .pipe(lambda x: x) - .pipe(lambda x: x) - .add_sink() - .build(num_threads=1, queue_class=CheckNameQueue, stage_id=ref) - ) - - -class TestPipelineFailureStructure(unittest.TestCase): - def _build_failing_pipeline(self): - def failing_range(n): - yield from range(n) - raise ValueError("Iterator failed") - - return ( - PipelineBuilder() - .add_source(failing_range(3)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - - def test_pipeline_failure_is_exception_group(self) -> None: - pipeline = self._build_failing_pipeline() - - with self.assertRaises(PipelineFailure) as ctx: - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pf = ctx.exception - if sys.version_info >= (3, 11): - self.assertIsInstance(pf, ExceptionGroup) - else: - self.assertIsInstance(pf, RuntimeError) - self.assertGreaterEqual(len(pf.exceptions), 1) - exception_types = {type(e) for e in pf.exceptions} - self.assertTrue(exception_types & {ValueError}) - - def test_pipeline_failure_individual_exceptions(self) -> None: - def failing_range(n): - yield from range(n) - raise TypeError("source failed") - - pipeline = ( - PipelineBuilder() - .add_source(failing_range(3)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure) as ctx: - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pf = ctx.exception - if sys.version_info >= (3, 11): - self.assertIsInstance(pf, ExceptionGroup) - else: - self.assertIsInstance(pf, RuntimeError) - self.assertGreaterEqual(len(pf.exceptions), 1) - self.assertTrue( - any(isinstance(e, TypeError) for e in pf.exceptions), - ) - - @unittest.skipIf(sys.version_info < (3, 11), "ExceptionGroup requires Python 3.11+") - def test_pipeline_failure_subgroup(self) -> None: - def failing_range(n): - yield from range(n) - raise ValueError("Iterator failed") - - pipeline = ( - PipelineBuilder() - .add_source(failing_range(3)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure) as ctx: - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pf = ctx.exception - sub = pf.subgroup(ValueError) - self.assertIsNotNone(sub) - self.assertIsInstance(sub, PipelineFailure) - self.assertTrue(all(isinstance(e, ValueError) for e in sub.exceptions)) - - @unittest.skipIf(sys.version_info < (3, 11), "ExceptionGroup requires Python 3.11+") - def test_pipeline_failure_notes_contain_stage_name(self) -> None: - def failing_range(n): - yield from range(n) - raise ValueError("Iterator failed") - - pipeline = ( - PipelineBuilder() - .add_source(failing_range(3)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure) as ctx: - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pf = ctx.exception - for exc in pf.exceptions: - notes = getattr(exc, "__notes__", []) - self.assertTrue( - any(note.startswith("Pipeline stage:") for note in notes), - ) diff --git a/tests/pipeline/pipeline_builder_test.py b/tests/pipeline/pipeline_builder_test.py new file mode 120000 index 000000000..9b0dd1d82 --- /dev/null +++ b/tests/pipeline/pipeline_builder_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_builder_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_cleanup_test.py b/tests/pipeline/pipeline_cleanup_test.py deleted file mode 100644 index 92ff1adef..000000000 --- a/tests/pipeline/pipeline_cleanup_test.py +++ /dev/null @@ -1,203 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import gc -import unittest -import warnings - -from spdl.pipeline import PipelineBuilder - - -class TestPipelineCleanup(unittest.TestCase): - """Test class for Pipeline cleanup functionality.""" - - def test_cleanup_called_on_garbage_collection(self) -> None: - """Test that the pipeline background thread is automatically stopped - on garbage collection without warnings.""" - - # Setup: Create a pipeline and start it - pipeline = ( - PipelineBuilder().add_source(range(100)).add_sink(1000).build(num_threads=1) - ) - - pipeline.start(timeout=3) - - # Verify the pipeline is running - self.assertTrue(pipeline._impl._event_loop.is_started()) - - # Keep a reference to the impl to verify it stopped - impl = pipeline._impl - - # Execute: Delete the pipeline reference without calling stop() - # The facade's finalizer should cleanly stop the pipeline - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - # Delete the pipeline and force garbage collection - del pipeline - gc.collect() - - # Assert: No warning should be issued — the facade handles cleanup - cleanup_warnings = [ - warning - for warning in w - if "Pipeline is running in the background" in str(warning.message) - ] - self.assertEqual(len(cleanup_warnings), 0) - - # Verify the background thread actually stopped - self.assertTrue(impl._event_loop.is_task_completed()) - - def test_cleanup_not_called_when_explicitly_stopped(self) -> None: - """Test that _cleanup_pipeline is not called when Pipeline is explicitly stopped.""" - - # Setup: Create a pipeline and start it - pipeline = ( - PipelineBuilder().add_source(range(100)).add_sink(1000).build(num_threads=1) - ) - - pipeline.start(timeout=3) - - # Execute: Explicitly stop the pipeline - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - pipeline.stop(timeout=3) - - # Delete the pipeline reference and force garbage collection - del pipeline - gc.collect() - - # Assert: Verify that no warning was issued - # Since we explicitly stopped the pipeline, the finalizer should be detached - # and no cleanup warning should be issued - cleanup_warnings = [ - warning - for warning in w - if "Pipeline is running in the background" in str(warning.message) - ] - self.assertEqual(len(cleanup_warnings), 0) - - def test_cleanup_with_auto_stop_context_manager(self) -> None: - """Test that cleanup is not called when using auto_stop context manager.""" - - # Setup: Create a pipeline - pipeline = ( - PipelineBuilder().add_source(range(10)).add_sink(1000).build(num_threads=1) - ) - - # Execute: Use the pipeline with auto_stop context manager - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - with pipeline.auto_stop(timeout=3): - # Get some items from the pipeline - iterator = pipeline.get_iterator(timeout=3) - for _ in range(5): - next(iterator) - - # Delete the pipeline reference and force garbage collection - del pipeline - gc.collect() - - # Assert: Verify that no cleanup warning was issued - # The context manager should have stopped the pipeline properly - cleanup_warnings = [ - warning - for warning in w - if "Pipeline is running in the background" in str(warning.message) - ] - self.assertEqual(len(cleanup_warnings), 0) - - def test_auto_start_and_cleanup_without_explicit_start_stop(self) -> None: - """Test that iterating a pipeline without calling start/stop works: - the background thread is started automatically on first iteration, - and the finalizer cleans it up on garbage collection.""" - - pipeline = ( - PipelineBuilder().add_source(range(10)).add_sink(1000).build(num_threads=1) - ) - - # Pipeline should not be started yet - self.assertFalse(pipeline._impl._event_loop.is_started()) - - # Iterate without explicit start — auto-start should kick in - items = [] - for item in pipeline: - items.append(item) - - # Verify auto-start happened - self.assertTrue(pipeline._impl._event_loop.is_started()) - self.assertEqual(sorted(items), list(range(10))) - - # Keep a reference to the impl to verify cleanup - impl = pipeline._impl - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - del pipeline - gc.collect() - - cleanup_warnings = [ - warning - for warning in w - if "Pipeline is running in the background" in str(warning.message) - ] - self.assertEqual(len(cleanup_warnings), 0) - - # Verify the impl was stopped (stop was called by finalizer) - self.assertTrue(impl._event_loop.is_task_completed()) - - def test_auto_start_and_cleanup_continuous_source(self) -> None: - """Test auto start/stop with a continuous source iterated multiple times. - - With continuous=True the source re-iterates, injecting epoch boundary - sentinels. Each ``for ... in pipeline`` consumes one epoch. The pipeline - should auto-start on the first epoch and remain running across epochs, - then clean up on garbage collection.""" - - pipeline = ( - PipelineBuilder() - .add_source(range(5), continuous=True) - .add_sink(1000) - .build(num_threads=1) - ) - - self.assertFalse(pipeline._impl._event_loop.is_started()) - - num_epochs = 3 - all_epoch_items = [] - for _ in range(num_epochs): - epoch_items = [] - for item in pipeline: - epoch_items.append(item) - all_epoch_items.append(sorted(epoch_items)) - - # Verify auto-start happened and each epoch produced the same items - self.assertTrue(pipeline._impl._event_loop.is_started()) - for epoch_items in all_epoch_items: - self.assertEqual(epoch_items, list(range(5))) - - # Verify cleanup on garbage collection - impl = pipeline._impl - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - - del pipeline - gc.collect() - - cleanup_warnings = [ - warning - for warning in w - if "Pipeline is running in the background" in str(warning.message) - ] - self.assertEqual(len(cleanup_warnings), 0) - - self.assertTrue(impl._event_loop.is_task_completed()) diff --git a/tests/pipeline/pipeline_cleanup_test.py b/tests/pipeline/pipeline_cleanup_test.py new file mode 120000 index 000000000..7e3b02f08 --- /dev/null +++ b/tests/pipeline/pipeline_cleanup_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_cleanup_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_def_test.py b/tests/pipeline/pipeline_def_test.py deleted file mode 100644 index 4dc6026a6..000000000 --- a/tests/pipeline/pipeline_def_test.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import unittest -from typing import TypeVar - -from spdl.pipeline import build_pipeline -from spdl.pipeline.defs import ( - Aggregate, - Disaggregate, - Pipe, - PipelineConfig, - SinkConfig, - SourceConfig, -) - -T = TypeVar("T") -U = TypeVar("U") - - -# pyre-strict - - -class PipelineDefTest(unittest.TestCase): - def test_source_repr(self) -> None: - """`repr` of SourceConfig should not generate a huge string.""" - - src = SourceConfig(list(range(10000))) - - self.assertGreater(len(repr(src.source)), 10000) - self.assertLess(len(repr(src)), 10000) - - def test_pipe_args_repr(self) -> None: - """`repr` of Pipe should not generate a huge string.""" - - lst = list(range(10000)) - self.assertGreater(len(repr(lst)), 10000) - pipe = Pipe(lst) - self.assertLess(len(repr(pipe._args.op)), 10000) - self.assertLess(len(repr(pipe)), 10000) - - dct = {i: i for i in range(10000)} - self.assertGreater(len(repr(dct)), 10000) - pipe = Pipe(dct) - self.assertLess(len(repr(pipe._args.op)), 10000) - self.assertLess(len(repr(pipe)), 10000) - - def _test_build_pipeline(self, cfg: PipelineConfig[T], expected: list[T]) -> None: - print(cfg) - pipeline = build_pipeline(cfg, num_threads=1) - - with pipeline.auto_stop(): - ite = pipeline.get_iterator(timeout=3) - self.assertEqual(list(ite), expected) - - def test_build_pipeline_simple(self) -> None: - """PipelineConfig and build_pipeline works without pipes.""" - src = range(10) - cfg = PipelineConfig( - src=SourceConfig(src), - pipes=[], - sink=SinkConfig(3), - ) - - self._test_build_pipeline(cfg, list(src)) - - def test_build_pipeline_aggregate(self) -> None: - """Aggregate works""" - cfg = PipelineConfig( - src=SourceConfig(range(8)), - pipes=[ - Aggregate(3, drop_last=False), - ], - sink=SinkConfig(3), - ) - - expected = [[0, 1, 2], [3, 4, 5], [6, 7]] - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_aggregate_drop_last(self) -> None: - """Aggregate works""" - cfg = PipelineConfig( - src=SourceConfig(range(8)), - pipes=[ - Aggregate(3, drop_last=True), - ], - sink=SinkConfig(3), - ) - - expected = [[0, 1, 2], [3, 4, 5]] - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_disaggregate(self) -> None: - """Disaggregate works""" - cfg = PipelineConfig( - src=SourceConfig([[0, 1, 2, 3]]), - pipes=[ - Disaggregate(), - ], - sink=SinkConfig(3), - ) - - expected = [0, 1, 2, 3] - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_pipe_identity(self) -> None: - """Pipe works with identity""" - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(lambda x: x), - ], - sink=SinkConfig(3), - ) - - expected = list(range(5)) - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_pipe_double(self) -> None: - """Pipe works with simple lambda""" - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(lambda x: 2 * x), - ], - sink=SinkConfig(3), - ) - - expected = [2 * i for i in range(5)] - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_pipe_sum(self) -> None: - """Pipe works with aggregated data""" - cfg = PipelineConfig( - src=SourceConfig(range(8)), - pipes=[ - Aggregate(3), - Pipe(sum), - ], - sink=SinkConfig(3), - ) - - expected = [3, 12, 13] - self._test_build_pipeline(cfg, expected) - - def test_build_pipeline_pipe_list(self) -> None: - """Pipe works with list""" - mapping = [i * i for i in range(8)] - cfg = PipelineConfig( - src=SourceConfig(range(8)), - pipes=[Pipe(mapping)], - sink=SinkConfig(3), - ) - self._test_build_pipeline(cfg, mapping) - - def test_build_pipeline_pipe_map(self) -> None: - """Pipe works with map (dict)""" - mapping = {i: i * i for i in range(8)} - cfg = PipelineConfig( - src=SourceConfig(range(8)), - pipes=[Pipe(mapping)], - sink=SinkConfig(3), - ) - - self._test_build_pipeline(cfg, list(mapping.values())) diff --git a/tests/pipeline/pipeline_def_test.py b/tests/pipeline/pipeline_def_test.py new file mode 120000 index 000000000..9413d3663 --- /dev/null +++ b/tests/pipeline/pipeline_def_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_def_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_failure_exceptstar_test.py b/tests/pipeline/pipeline_failure_exceptstar_test.py deleted file mode 100644 index a46baf8f3..000000000 --- a/tests/pipeline/pipeline_failure_exceptstar_test.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -"""Tests for PipelineFailure using ``except*`` syntax (Python 3.11+ only). - -This file is separated from pipeline_builder_test.py because ``except*`` is a -syntactic construct that causes SyntaxError on Python < 3.11 at import time, -which would prevent the entire module from loading. -""" - -import sys -import unittest -from collections.abc import Iterator - -if sys.version_info < (3, 11): - raise unittest.SkipTest("except* syntax requires Python 3.11+") - -from spdl.pipeline import PipelineBuilder - - -def passthrough(x: int) -> int: - return x - - -class TestPipelineFailureExceptStar(unittest.TestCase): - def test_pipeline_failure_except_star(self) -> None: - def failing_range(n: int) -> Iterator[int]: - yield from range(n) - raise ValueError("Iterator failed") - - pipeline = ( - PipelineBuilder() - .add_source(failing_range(3)) - .pipe(passthrough) - .add_sink(1000) - .build(num_threads=1) - ) - - caught = [] - try: - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - except* ValueError as eg: - caught.extend(eg.exceptions) - - self.assertGreaterEqual(len(caught), 1) - self.assertIsInstance(caught[0], ValueError) diff --git a/tests/pipeline/pipeline_failure_exceptstar_test.py b/tests/pipeline/pipeline_failure_exceptstar_test.py new file mode 120000 index 000000000..1cc58bfa5 --- /dev/null +++ b/tests/pipeline/pipeline_failure_exceptstar_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_failure_exceptstar_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_node_test.py b/tests/pipeline/pipeline_node_test.py deleted file mode 100644 index eab1bb24b..000000000 --- a/tests/pipeline/pipeline_node_test.py +++ /dev/null @@ -1,299 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import unittest - -from spdl.pipeline._components import AsyncQueue -from spdl.pipeline._components._common import StageInfo -from spdl.pipeline._components._node import ( - _cancel_orphaned, - _cancel_recursive, - _FanInNode, - _gather_error, - _Node, - _PathVariantsMergeConfig, - _SourceNode, - _start_tasks, -) -from spdl.pipeline.defs import SinkConfig, SourceConfig - - -class DummyException(Exception): - pass - - -_TTestNode = _SourceNode | _Node | _FanInNode - - -def _node( - name: str, - deps: list[_TTestNode], - exc: Exception | None = None, -) -> _TTestNode: - async def coro() -> None: - if exc: - raise exc - else: - await asyncio.sleep(10) - - info = StageInfo(pipeline_id=0, stage_id="0", stage_name=name) - n: _TTestNode - if not deps: - n = _SourceNode( - info, - SourceConfig(source=[]), - output_queue=AsyncQueue(info), - ) - elif len(deps) > 1: - n = _FanInNode( - info, - _PathVariantsMergeConfig(), - deps, - input_queues=[ - AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name=f"{name}_in_{i}") - ) - for i in range(len(deps)) - ], - output_queue=AsyncQueue(info), - ) - else: - n = _Node( - info, - SinkConfig(buffer_size=1), - deps, - input_queue=AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name=f"{name}_in") - ), - output_queue=AsyncQueue(info), - ) - n._coro = coro() - return n - - -class PipelineNodeTest(unittest.TestCase): - def test_node_chain_start_and_cancel(self) -> None: - # A -> B -> C - - async def run() -> None: - a = _node("A", []) - b = _node("B", [a]) - c = _node("C", [b]) - tasks = _start_tasks(c) - self.assertTrue(all(isinstance(t, asyncio.Task) for t in tasks)) - self.assertEqual(len(tasks), 3) - - _cancel_recursive(c) - await asyncio.wait(tasks) - self.assertTrue(all(t.cancelled() for t in tasks)) - - asyncio.run(run()) - - def test_node_y_shape_upstream(self) -> None: - # A1 B1 - # | | - # A2 B2 - # \ / - # C1 - - async def run() -> None: - a1 = _node("A1", []) - a2 = _node("A2", [a1]) - b1 = _node("B1", []) - b2 = _node("B2", [b1]) - c1 = _node("C1", [a2, b2]) - tasks = _start_tasks(c1) - self.assertEqual(len(tasks), 5) - _cancel_recursive(c1) - await asyncio.wait(tasks) - self.assertTrue(all(t.cancelled() for t in tasks)) - - asyncio.run(run()) - - def test_cancel_error_upstreams_and_gather_error(self) -> None: - # A1 B1 - # | | - # A2 B2 (raises) - # \ / - # C1 - - async def run() -> None: - a1 = _node("A1", []) - a2 = _node("A2", [a1]) - b1 = _node("B1", []) - b2 = _node("B2", [b1], exc=DummyException("fail B2")) - c1 = _node("C1", [a2, b2]) - tasks = _start_tasks(c1) - await asyncio.sleep(0) - - # Let B2 fail - await asyncio.wait([b2.task]) - - _cancel_orphaned(c1) - await asyncio.wait(tasks) - - # Only B1 should be cancelled, since B2 errored - self.assertFalse(a1.task.cancelled()) - self.assertFalse(a2.task.cancelled()) - self.assertTrue(b1.task.cancelled()) - self.assertFalse(b2.task.cancelled()) - self.assertFalse(c1.task.cancelled()) - - errs = _gather_error(c1) - self.assertEqual(len(errs), 1) - name, err = errs[0] - self.assertEqual(name, "0:0:B2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail B2") - - asyncio.run(run()) - - def test_cancel_error_upstreams_and_gather_error_multiple(self) -> None: - # A1 B1 - # | | - # (raises) A2 B2 (raises) - # \ / - # C1 - async def run() -> None: - a1 = _node("A1", []) - a2 = _node("A2", [a1], exc=DummyException("fail A2")) - b1 = _node("B1", []) - b2 = _node("B2", [b1], exc=DummyException("fail B2")) - c1 = _node("C1", [a2, b2]) - - tasks = _start_tasks(c1) - await asyncio.sleep(0) - - # Let A2, B2 fail - await asyncio.wait([a2.task, b2.task]) - - _cancel_orphaned(c1) - await asyncio.wait(tasks) - - self.assertTrue(a1.task.cancelled()) - self.assertFalse(a2.task.cancelled()) - self.assertTrue(b1.task.cancelled()) - self.assertFalse(b2.task.cancelled()) - self.assertFalse(c1.task.cancelled()) - - errs = _gather_error(c1) - self.assertEqual(len(errs), 2) - name, err = errs[0] - self.assertEqual(name, "0:0:A2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail A2") - name, err = errs[1] - self.assertEqual(name, "0:0:B2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail B2") - - asyncio.run(run()) - - def test_cancel_error_upstreams_and_gather_error_complex(self) -> None: - # B1 C1 - # | | - # (raises) B2 C2 (raises) - # \ / - # A1 D1 E1 - # | | | - # (raises) A2 D2 E2 - # \ | / - # F1 - - async def run() -> None: - a1 = _node("A1", []) - a2 = _node("A2", [a1], exc=DummyException("fail A2")) - b1 = _node("B1", []) - b2 = _node("B2", [b1], exc=DummyException("fail B2")) - c1 = _node("C1", []) - c2 = _node("C2", [c1], exc=DummyException("fail C2")) - d1 = _node("D1", [b2, c2]) - d2 = _node("D2", [d1]) - e1 = _node("E1", []) - e2 = _node("E1", [e1]) - f1 = _node("F1", [a2, d2, e2]) - - tasks = _start_tasks(f1) - await asyncio.sleep(0) - - # Let A2, B2, C2 fail - await asyncio.wait([a2.task, b2.task, c2.task]) - - _cancel_orphaned(f1) - - # Let the cancellations propagate - await asyncio.sleep(0) - - self.assertTrue(a1.task.cancelled()) - self.assertFalse(a2.task.cancelled()) - self.assertTrue(b1.task.cancelled()) - self.assertFalse(b2.task.cancelled()) - self.assertTrue(c1.task.cancelled()) - self.assertFalse(c2.task.cancelled()) - self.assertFalse(d1.task.cancelled()) - self.assertFalse(d2.task.cancelled()) - self.assertFalse(e1.task.cancelled()) - self.assertFalse(e2.task.cancelled()) - self.assertFalse(f1.task.cancelled()) - - await asyncio.wait(tasks) - - errs = _gather_error(f1) - self.assertEqual(len(errs), 3) - name, err = errs[0] - self.assertEqual(name, "0:0:A2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail A2") - name, err = errs[1] - self.assertEqual(name, "0:0:B2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail B2") - name, err = errs[2] - self.assertEqual(name, "0:0:C2") - self.assertIsInstance(err, DummyException) - self.assertEqual(err.args[0], "fail C2") - - asyncio.run(run()) - - def test_gather_error_with_cancelled(self) -> None: - # A1 B1 - # | | - # A2 B2 - # \ / - # C1 (raises) - - async def run() -> None: - a1 = _node("A1", []) - a2 = _node("A2", [a1]) - b1 = _node("B1", []) - b2 = _node("B2", [b1]) - c1 = _node("C1", [a2, b2], exc=DummyException("fail C")) - - tasks = _start_tasks(c1) - await asyncio.sleep(0) - - # Let C fail - await asyncio.wait([c1.task]) - - _cancel_orphaned(c1) - await asyncio.wait(tasks) - - self.assertTrue(a1.task.cancelled()) - self.assertTrue(a2.task.cancelled()) - self.assertTrue(b1.task.cancelled()) - self.assertTrue(b2.task.cancelled()) - self.assertFalse(c1.task.cancelled()) - - errs = _gather_error(c1) - self.assertEqual(len(errs), 1) - name, err = errs[0] - self.assertEqual(name, "0:0:C1") - self.assertIsInstance(err, DummyException) - - asyncio.run(run()) diff --git a/tests/pipeline/pipeline_node_test.py b/tests/pipeline/pipeline_node_test.py new file mode 120000 index 000000000..d0f5157a9 --- /dev/null +++ b/tests/pipeline/pipeline_node_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_node_test.py \ No newline at end of file diff --git a/tests/pipeline/pipeline_profiling_test.py b/tests/pipeline/pipeline_profiling_test.py deleted file mode 100644 index 7e77c2abf..000000000 --- a/tests/pipeline/pipeline_profiling_test.py +++ /dev/null @@ -1,321 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import unittest -from collections.abc import AsyncIterator, Iterator -from contextlib import contextmanager -from unittest.mock import MagicMock, patch - -from spdl.pipeline import ( - config, - profile_pipeline, - ProfileHook, - ProfileResult, -) -from spdl.pipeline._profile import ( - _build_pipeline_config, - _fetch_inputs, -) -from spdl.pipeline.defs import ( - Aggregate, - Disaggregate, - Merge, - Pipe, - PipeConfig, - PipelineConfig, - SinkConfig, - SourceConfig, -) - - -class FetchInputsTest(unittest.TestCase): - """Test _fetch_inputs functionality.""" - - def test_fetch_inputs(self): - """_fetch_inputs collects input items""" - src = SourceConfig(range(10)) - - inputs = _fetch_inputs(src, num_items=3) - self.assertEqual(inputs, list(range(3))) - - def test_fetch_inputs_async(self): - """_fetch_inputs collects input items""" - - async def arange(n: int) -> AsyncIterator[int]: - for i in range(n): - yield i - - src = SourceConfig(arange(10)) - - inputs = _fetch_inputs(src, num_items=3) - self.assertEqual(inputs, list(range(3))) - - -class ProfilePipelineTest(unittest.TestCase): - """Test profile_pipeline functionality.""" - - def setUp(self) -> None: - """Reset all configuration state before each test.""" - config.set_default_profile_hook() - config.set_default_profile_callback() - - def test_profile_pipeline(self): - def foo(i: int) -> int: - return 2 * i - - def bar(items: list[int]) -> list[int]: - return [sum(items)] - - def bazz(i: int) -> int: - return i * i - - N, m = 25, 3 - - plc = PipelineConfig( - src=SourceConfig(range(N)), - pipes=[ - Pipe(foo), - Aggregate(m), - Pipe(bar), - Disaggregate(), - Pipe(bazz), - ], - sink=SinkConfig(3), - ) - - class Intercept_: - def __init__(self) -> None: - self.i = 0 - - def __call__(self, inputs, pipe, concurrency): - num_inputs = N if self.i < 2 else (N + m - 1) // m - self.assertEqual(len(inputs), num_inputs) - self.assertEqual(pipe, plc.pipes[self.i]) - ret = _build_pipeline_config(inputs, pipe, concurrency) - self.assertEqual(len(ret.pipes), 1) - if isinstance(pipe, PipeConfig): - self.assertIs(ret.pipes[0]._args.op, plc.pipes[self.i]._args.op) - self.i += 1 - return ret - - mock = Intercept_() - mock.assertEqual = self.assertEqual - mock.assertIs = self.assertIs - with patch("spdl.pipeline._profile._build_pipeline_config", mock): - profile_pipeline(plc) - - self.assertEqual(mock.i, 5) - - def test_profile_pipeline_callback(self): - """Test that profile_pipeline calls the callback for each pipe stage.""" - - def simple_op(i: int) -> int: - return i + 1 - - cfg = PipelineConfig( - src=SourceConfig(range(10)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - callback_mock = MagicMock() - results = profile_pipeline(cfg, num_inputs=5, callback=callback_mock) - - callback_mock.assert_called_once() - called_args = callback_mock.call_args[0] - self.assertEqual(len(called_args), 1) - called_result = called_args[0] - - self.assertIsInstance(called_result, ProfileResult) - self.assertEqual(called_result.name, "simple_op") - self.assertGreater(len(called_result.stats), 0) - - self.assertEqual(len(results), 1) - self.assertEqual(results[0].name, called_result.name) - self.assertEqual(len(results[0].stats), len(called_result.stats)) - - def test_profile_pipeline_no_callback(self): - """Test that profile_pipeline works correctly when no callback is provided.""" - - def simple_op(i: int) -> int: - return i * 2 - - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - results = profile_pipeline(cfg, num_inputs=3, callback=None) - - self.assertEqual(len(results), 1) - self.assertEqual(results[0].name, "simple_op") - self.assertGreater(len(results[0].stats), 0) - - -class ProfileHookTest(unittest.TestCase): - """Test class for ProfileHook functionality.""" - - def setUp(self) -> None: - """Reset all configuration state before each test.""" - config.set_default_profile_hook() - config.set_default_profile_callback() - - def test_profile_pipeline_custom_hook_methods_called(self): - """Test that custom ProfileHook's stage_profile_hook and - pipeline_profile_hook methods are called. - """ - - def simple_op(i: int) -> int: - return i + 10 - - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - stage_hook_mock = MagicMock() - pipeline_hook_mock = MagicMock() - - class MockProfileHook(ProfileHook): - @contextmanager - def stage_profile_hook( - self, _stage: str, _concurrency: int - ) -> Iterator[None]: - stage_hook_mock() - try: - yield None - finally: - stage_hook_mock() - - @contextmanager - def pipeline_profile_hook(self) -> Iterator[None]: - pipeline_hook_mock() - try: - yield - finally: - pipeline_hook_mock() - - custom_hook = MockProfileHook() - - results = profile_pipeline(cfg, num_inputs=3, hook=custom_hook) - - self.assertGreater(len(results), 0) - self.assertEqual(pipeline_hook_mock.call_count, 2) - self.assertEqual(stage_hook_mock.call_count, 10) - - def test_profile_pipeline_skips_when_local_rank_not_zero(self): - """Test that profiling is skipped if LOCAL_RANK is not '0'.""" - - def simple_op(i: int) -> int: - return i * 3 - - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - with patch("spdl.pipeline._profile._get_local_rank", return_value=1): - results = profile_pipeline(cfg, num_inputs=5) - - self.assertEqual(results, []) - - def test_profile_pipeline_runs_when_local_rank_zero(self): - """Test that profiling runs normally when LOCAL_RANK is '0'.""" - - def simple_op(i: int) -> int: - return i * 2 - - cfg = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(simple_op), - ], - sink=SinkConfig(1), - ) - - with patch("spdl.pipeline._profile._get_local_rank", return_value=0): - results = profile_pipeline(cfg, num_inputs=3) - - self.assertGreater(len(results), 0) - self.assertEqual(results[0].name, "simple_op") - self.assertGreater(len(results[0].stats), 0) - - -class MergeConfigTest(unittest.TestCase): - """Test class for profile_pipeline with Merge configurations.""" - - def setUp(self) -> None: - """Reset all configuration state before each test.""" - config.set_default_profile_hook() - config.set_default_profile_callback() - - def test_profile_pipeline_with_merge_config_and_post_merge_stages(self): - """Test that profile_pipeline profiles all stages including - those in Merge and post-merge stages. - """ - - def double(i: int) -> int: - return i * 2 - - def triple(i: int) -> int: - return i * 3 - - def add_ten(i: int) -> int: - return i + 10 - - def square(i: int) -> int: - return i * i - - plc1 = PipelineConfig( - src=SourceConfig(range(5)), - pipes=[ - Pipe(double, name="double"), - ], - sink=SinkConfig(1), - ) - - plc2 = PipelineConfig( - src=SourceConfig(range(10, 15)), - pipes=[ - Pipe(triple, name="triple"), - ], - sink=SinkConfig(1), - ) - - main_cfg = PipelineConfig( - src=Merge([plc1, plc2]), - pipes=[ - Pipe(add_ten, name="add_ten"), - Pipe(square, name="square"), - ], - sink=SinkConfig(1), - ) - results = profile_pipeline(main_cfg, num_inputs=3) - - self.assertEqual(len(results), 4) - - self.assertEqual(results[0].name, "double") - self.assertEqual(results[1].name, "triple") - self.assertEqual(results[2].name, "add_ten") - self.assertEqual(results[3].name, "square") - - for result in results: - self.assertGreater(len(result.stats), 0) - for stat in result.stats: - self.assertTrue(hasattr(stat, "concurrency")) - self.assertTrue(hasattr(stat, "qps")) - self.assertTrue(hasattr(stat, "occupancy_rate")) diff --git a/tests/pipeline/pipeline_profiling_test.py b/tests/pipeline/pipeline_profiling_test.py new file mode 120000 index 000000000..419cbd3bf --- /dev/null +++ b/tests/pipeline/pipeline_profiling_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/pipeline_profiling_test.py \ No newline at end of file diff --git a/tests/pipeline/priority_executor_test.py b/tests/pipeline/priority_executor_test.py deleted file mode 100644 index 462568d0a..000000000 --- a/tests/pipeline/priority_executor_test.py +++ /dev/null @@ -1,937 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import functools -import gc -import itertools -import pickle -import sys -import threading -import time -import unittest -import warnings -import weakref -from collections.abc import Callable -from concurrent.futures import Executor, Future, ThreadPoolExecutor -from queue import Empty -from typing import Type, TypeVar - -from spdl.pipeline import ( - PipelineBuilder, - PipelineFailure, - PriorityExecutorEntrypoint, - PriorityProcessPoolExecutor, - PriorityThreadPoolExecutor, -) -from spdl.pipeline._priority_executor import _OWNER_REGISTRY, _PriorityQueueAdapter - -_F = TypeVar("_F", bound=Callable[..., object]) -_C = TypeVar("_C", bound=Type[object]) - - -def _ignore_fork_warning(fn: _F) -> _F: - @functools.wraps(fn) - def wrapper(*args: object, **kwargs: object) -> object: - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - # pyre-ignore[7] - return wrapper - - -def _ignore_fork_warning_in_class(cls: _C) -> _C: - for name, member in list(vars(cls).items()): - if name.startswith("test_") and callable(member): - setattr(cls, name, _ignore_fork_warning(member)) - return cls - - -def _raise_value_error() -> None: - raise ValueError("process boom") - - -# ─── _PriorityQueueAdapter unit tests ─── - - -class TestPriorityQueueAdapter(unittest.TestCase): - def test_priority_ordering(self) -> None: - q = _PriorityQueueAdapter() - q.put("low") - q.put("high") - q.put("mid") - # All inserted without priority context → same default → FIFO - self.assertEqual(q.get(), "low") - self.assertEqual(q.get(), "high") - self.assertEqual(q.get(), "mid") - - def test_priority_ordering_with_explicit_priority(self) -> None: - q = _PriorityQueueAdapter() - - q.put("low", priority=(10, 0)) - q.put("high", priority=(0, 0)) - q.put("mid", priority=(5, 0)) - - self.assertEqual(q.get(), "high") - self.assertEqual(q.get(), "mid") - self.assertEqual(q.get(), "low") - - def test_fifo_within_same_priority(self) -> None: - q = _PriorityQueueAdapter() - for i in range(5): - q.put(f"item_{i}", priority=(1, i)) - - for i in range(5): - self.assertEqual(q.get(), f"item_{i}") - - def test_sentinel_none_processed_first(self) -> None: - q = _PriorityQueueAdapter() - q.put("work", priority=(0, 0)) - q.put(None) # shutdown sentinel, no priority - - self.assertIsNone(q.get()) - self.assertEqual(q.get(), "work") - - def test_get_nowait_empty_raises(self) -> None: - q = _PriorityQueueAdapter() - with self.assertRaises(Empty): - q.get_nowait() - - def test_empty_and_qsize(self) -> None: - q = _PriorityQueueAdapter() - self.assertTrue(q.empty()) - self.assertEqual(q.qsize(), 0) - q.put("x") - self.assertFalse(q.empty()) - self.assertEqual(q.qsize(), 1) - - -# ─── PriorityExecutorEntrypoint unit tests ─── - - -class TestPriorityExecutorEntrypoint(unittest.TestCase): - def test_submit_returns_future(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - fut = stage.submit(lambda: 42) - self.assertIsInstance(fut, Future) - self.assertEqual(fut.result(timeout=5), 42) - executor.shutdown() - - def test_is_executor_compatible(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - self.assertIsInstance(stage, Executor) - executor.shutdown() - - def test_shutdown_is_noop(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - stage.shutdown() # should not affect the pool - fut = stage.submit(lambda: 1) - self.assertEqual(fut.result(timeout=5), 1) - executor.shutdown() - - -# ─── PriorityThreadPoolExecutor ordering tests ─── - - -class TestPriorityThreadPoolExecutorOrdering(unittest.TestCase): - def test_downstream_stage_runs_first(self) -> None: - """With 1 worker, tasks are executed in priority order.""" - barrier = threading.Barrier(2) - results: list[str] = [] - - executor = PriorityThreadPoolExecutor(max_workers=1) - upstream = executor.get_executor(priority=0) - downstream = executor.get_executor(priority=2) - - # Block the single worker so we can enqueue both tasks - executor.get_executor().submit(lambda: barrier.wait(timeout=5)) - - # Enqueue upstream first, then downstream - upstream.submit(lambda: results.append("upstream")) - downstream.submit(lambda: results.append("downstream")) - - # Release the worker - barrier.wait(timeout=5) - executor.shutdown(wait=True) - - self.assertEqual(results, ["downstream", "upstream"]) - - def test_fifo_within_stage(self) -> None: - barrier = threading.Barrier(2) - results: list[int] = [] - - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - - executor.get_executor().submit(lambda: barrier.wait(timeout=5)) - - for i in range(5): - stage.submit(lambda i=i: results.append(i)) - - barrier.wait(timeout=5) - executor.shutdown(wait=True) - - self.assertEqual(results, [0, 1, 2, 3, 4]) - - def test_multiple_stages_interleaved(self) -> None: - """3 stages, items interleaved — should sort by priority then FIFO.""" - barrier = threading.Barrier(2) - results: list[tuple[int, int]] = [] - - executor = PriorityThreadPoolExecutor(max_workers=1) - stages = [executor.get_executor(priority=p) for p in [0, 1, 2]] - - executor.get_executor().submit(lambda: barrier.wait(timeout=5)) - - # Submit: stage0, stage1, stage2, stage0, stage1, stage2 - for round_idx in range(2): - for idx, stage in enumerate(stages): - stage.submit(lambda i=idx, ri=round_idx: results.append((i, ri))) - - barrier.wait(timeout=5) - executor.shutdown(wait=True) - - # priority 2 (highest) first, then priority 1, then priority 0 - # Within same priority, FIFO by round - self.assertEqual( - results, - [(2, 0), (2, 1), (1, 0), (1, 1), (0, 0), (0, 1)], - ) - - def test_basic_execution(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=2) - stage = executor.get_executor() - futs = [stage.submit(lambda x=x: x * 2, x) for x in range(10)] - results = {f.result(timeout=5) for f in futs} - self.assertEqual(results, {x * 2 for x in range(10)}) - executor.shutdown() - - def test_exception_propagation(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - - def fail() -> None: - raise ValueError("boom") - - fut = stage.submit(fail) - with self.assertRaises(ValueError): - fut.result(timeout=5) - executor.shutdown() - - def test_submit_after_shutdown_raises(self) -> None: - executor = PriorityThreadPoolExecutor(max_workers=1) - stage = executor.get_executor() - executor.shutdown() - with self.assertRaises(RuntimeError): - executor._submit_with_priority((0, 0), lambda: None, (), {}) - - def test_multiple_workers_all_complete(self) -> None: - """With multiple workers, all tasks must complete.""" - executor = PriorityThreadPoolExecutor(max_workers=4) - stages = [executor.get_executor() for _ in range(3)] - - counter: itertools.count[int] = itertools.count() - results: list[int] = [] - lock: threading.Lock = threading.Lock() - - def work() -> None: - val = next(counter) - with lock: - results.append(val) - - futs = [] - for stage in stages: - for _ in range(10): - futs.append(stage.submit(work)) - - for f in futs: - f.result(timeout=10) - - executor.shutdown() - self.assertEqual(len(results), 30) - - -# ─── PriorityProcessPoolExecutor tests ─── - - -@_ignore_fork_warning_in_class -class TestPriorityProcessPoolExecutor(unittest.TestCase): - def test_basic_execution(self) -> None: - executor = PriorityProcessPoolExecutor(max_workers=2) - stage = executor.get_executor() - futs = [stage.submit(pow, 2, x) for x in range(10)] - results = {f.result(timeout=10) for f in futs} - self.assertEqual(results, {2**x for x in range(10)}) - executor.shutdown() - - def test_exception_propagation(self) -> None: - executor = PriorityProcessPoolExecutor(max_workers=1) - stage = executor.get_executor() - - fut = stage.submit(_raise_value_error) - with self.assertRaises(ValueError): - fut.result(timeout=10) - executor.shutdown() - - -# ─── Pipeline integration: correctness ─── - - -class TestPriorityExecutorPipelineCorrectness(unittest.TestCase): - def test_two_stage_pipeline(self) -> None: - """All items flow through correctly with a shared priority executor.""" - pool = PriorityThreadPoolExecutor(max_workers=4) - s1 = pool.get_executor() - s2 = pool.get_executor() - - pipeline = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(lambda x: x * 2, executor=s1, concurrency=4) - .pipe(lambda x: x + 1, executor=s2, concurrency=4) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) - pool.shutdown() - - def test_three_stage_pipeline(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=4) - s1 = pool.get_executor() - s2 = pool.get_executor() - s3 = pool.get_executor() - - pipeline = ( - PipelineBuilder() - .add_source(range(15)) - .pipe(lambda x: x + 1, executor=s1, concurrency=2) - .pipe(lambda x: x * 3, executor=s2, concurrency=2) - .pipe(lambda x: x - 1, executor=s3, concurrency=2) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted((x + 1) * 3 - 1 for x in range(15))) - pool.shutdown() - - def test_mixed_priority_and_default_executor(self) -> None: - """Some stages use priority executor, others use default.""" - pool = PriorityThreadPoolExecutor(max_workers=2) - s1 = pool.get_executor() - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(lambda x: x * 2, executor=s1, concurrency=2) - .pipe(lambda x: x + 1, concurrency=2) - .add_sink(3) - .build(num_threads=2) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted(x * 2 + 1 for x in range(10))) - pool.shutdown() - - def test_exception_propagates_through_pipeline(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=2) - s1 = pool.get_executor() - s2 = pool.get_executor() - - def fail_on_five(x: int) -> int: - if x == 5: - raise ValueError("boom on 5") - return x - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(fail_on_five, executor=s1, concurrency=1, max_failures=0) - .pipe(lambda x: x, executor=s2, concurrency=1) - .add_sink(3) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pool.shutdown() - - -# ─── Pipeline integration: priority ordering ─── - - -class TestPriorityExecutorPipelineOrdering(unittest.TestCase): - def test_downstream_not_starved(self) -> None: - """With 1 shared worker, downstream should interleave with upstream, - not wait until all upstream completes.""" - execution_log: list[tuple[str, int]] = [] - lock: threading.Lock = threading.Lock() - - pool = PriorityThreadPoolExecutor(max_workers=1) - upstream_exec = pool.get_executor() - downstream_exec = pool.get_executor() - - def upstream_op(item: int) -> int: - with lock: - execution_log.append(("up", item)) - # Sleep so the event loop has time to submit the downstream task - # before the worker picks the next item. - time.sleep(0.03) - return item - - def downstream_op(item: int) -> int: - with lock: - execution_log.append(("down", item)) - return item - - pipeline = ( - PipelineBuilder() - .add_source(range(8)) - .pipe(upstream_op, executor=upstream_exec, concurrency=1) - .pipe(downstream_op, executor=downstream_exec, concurrency=1) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=30)) - - self.assertEqual(len(results), 8) - pool.shutdown() - - # With priority: up/down alternate → max consecutive upstream ≤ 2. - # Without priority (FIFO): upstream dominates → all upstream first. - max_consec_up = 0 - consec = 0 - for stage, _ in execution_log: - if stage == "up": - consec += 1 - max_consec_up = max(max_consec_up, consec) - else: - consec = 0 - - self.assertLessEqual( - max_consec_up, - 2, - f"Upstream ran {max_consec_up} consecutive times — " - f"downstream was starved. Full log: {execution_log}", - ) - - def test_three_stage_priority_order(self) -> None: - """With 3 stages sharing 1 worker, the most downstream stage - with pending work should run first.""" - execution_log: list[tuple[str, int]] = [] - lock: threading.Lock = threading.Lock() - - pool = PriorityThreadPoolExecutor(max_workers=1) - s1_exec = pool.get_executor() - s2_exec = pool.get_executor() - s3_exec = pool.get_executor() - - def make_op(name: str): # pyre-ignore[3] - def op(item: int) -> int: - with lock: - execution_log.append((name, item)) - time.sleep(0.03) - return item - - return op - - pipeline = ( - PipelineBuilder() - .add_source(range(6)) - .pipe(make_op("s1"), executor=s1_exec, concurrency=1) - .pipe(make_op("s2"), executor=s2_exec, concurrency=1) - .pipe(make_op("s3"), executor=s3_exec, concurrency=1) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = list(pipeline.get_iterator(timeout=30)) - - self.assertEqual(len(results), 6) - pool.shutdown() - - # After downstream stages have pending work, upstream should not - # run consecutively. - stages = [stage for stage, _ in execution_log] - for i in range(len(stages) - 1): - if stages[i] == "s1" and stages[i + 1] == "s1": - prior = stages[:i] - self.assertNotIn( - "s2", - prior, - f"Consecutive s1 at positions {i},{i + 1} after s2 " - f"already had work. Log: {execution_log}", - ) - - def test_priority_vs_fifo_comparison(self) -> None: - """Priority executor interleaves downstream at least as well as FIFO.""" - - def run_pipeline( - upstream_exec: Executor, - downstream_exec: Executor, - num_threads: int, - ) -> list[str]: - log: list[str] = [] - lock: threading.Lock = threading.Lock() - - def up_op(item: int) -> int: - with lock: - log.append("up") - time.sleep(0.03) - return item - - def down_op(item: int) -> int: - with lock: - log.append("down") - return item - - pipeline = ( - PipelineBuilder() - .add_source(range(8)) - .pipe(up_op, executor=upstream_exec, concurrency=1) - .pipe(down_op, executor=downstream_exec, concurrency=1) - .add_sink(3) - .build(num_threads=num_threads) - ) - - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=30)) - return log - - # Priority executor - pool = PriorityThreadPoolExecutor(max_workers=1) - priority_log = run_pipeline( - pool.get_executor(), - pool.get_executor(), - num_threads=1, - ) - pool.shutdown() - - # Plain FIFO executor - plain = ThreadPoolExecutor(max_workers=1) - fifo_log = run_pipeline(plain, plain, num_threads=1) - plain.shutdown() - - # With priority, downstream entries should appear at least as early. - halfway = len(priority_log) // 2 - priority_down_first_half = priority_log[:halfway].count("down") - fifo_down_first_half = fifo_log[:halfway].count("down") - - self.assertGreaterEqual( - priority_down_first_half, - fifo_down_first_half, - f"Priority executor should interleave downstream tasks at least " - f"as well as FIFO.\n" - f"Priority log: {priority_log}\n" - f"FIFO log: {fifo_log}", - ) - - -# ─── Drop-in compatibility ─── - - -class TestDropInCompatibility(unittest.TestCase): - def test_stage_executor_with_as_completed(self) -> None: - from concurrent.futures import as_completed - - pool = PriorityThreadPoolExecutor(max_workers=2) - stage = pool.get_executor() - futs = [stage.submit(pow, 2, i) for i in range(5)] - results = set() - for f in as_completed(futs, timeout=5): - results.add(f.result()) - self.assertEqual(results, {1, 2, 4, 8, 16}) - pool.shutdown() - - def test_context_manager(self) -> None: - with PriorityThreadPoolExecutor(max_workers=2) as executor: - stage = executor.get_executor() - self.assertEqual(stage.submit(lambda: 99).result(timeout=5), 99) - - -# ─── Mixed priority + regular ThreadPoolExecutor pipelines ─── - - -class TestMixedExecutorPipeline(unittest.TestCase): - def test_priority_upstream_regular_downstream(self) -> None: - """Priority executor on upstream stages, regular ThreadPoolExecutor - on downstream stage.""" - pool = PriorityThreadPoolExecutor(max_workers=2) - regular = ThreadPoolExecutor(max_workers=2) - - pipeline = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(lambda x: x * 2, executor=pool.get_executor(), concurrency=2) - .pipe(lambda x: x + 1, executor=regular, concurrency=2) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) - pool.shutdown() - regular.shutdown() - - def test_regular_upstream_priority_downstream(self) -> None: - """Regular ThreadPoolExecutor on upstream, priority executor on - downstream.""" - pool = PriorityThreadPoolExecutor(max_workers=2) - regular = ThreadPoolExecutor(max_workers=2) - - pipeline = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(lambda x: x + 10, executor=regular, concurrency=2) - .pipe(lambda x: x * 3, executor=pool.get_executor(), concurrency=2) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted((x + 10) * 3 for x in range(20))) - pool.shutdown() - regular.shutdown() - - def test_three_stage_alternating_executors(self) -> None: - """Three stages: priority, regular, priority — verifying they - compose correctly.""" - pool = PriorityThreadPoolExecutor(max_workers=3) - regular = ThreadPoolExecutor(max_workers=2) - - pipeline = ( - PipelineBuilder() - .add_source(range(15)) - .pipe(lambda x: x + 1, executor=pool.get_executor(), concurrency=2) - .pipe(lambda x: x * 2, executor=regular, concurrency=2) - .pipe(lambda x: x - 1, executor=pool.get_executor(), concurrency=2) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted((x + 1) * 2 - 1 for x in range(15))) - pool.shutdown() - regular.shutdown() - - def test_two_priority_pools_and_regular(self) -> None: - """Two independent priority pools plus a regular pool, all in one - pipeline.""" - pool_a = PriorityThreadPoolExecutor(max_workers=2) - pool_b = PriorityThreadPoolExecutor(max_workers=2) - regular = ThreadPoolExecutor(max_workers=2) - - pipeline = ( - PipelineBuilder() - .add_source(range(12)) - .pipe(lambda x: x + 1, executor=pool_a.get_executor(), concurrency=2) - .pipe(lambda x: x * 2, executor=regular, concurrency=2) - .pipe(lambda x: x - 1, executor=pool_b.get_executor(), concurrency=2) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted((x + 1) * 2 - 1 for x in range(12))) - pool_a.shutdown() - pool_b.shutdown() - regular.shutdown() - - def test_exception_in_regular_stage(self) -> None: - """Exception in the regular executor stage propagates correctly.""" - pool = PriorityThreadPoolExecutor(max_workers=2) - regular = ThreadPoolExecutor(max_workers=2) - - def fail_on_five(x: int) -> int: - if x == 5: - raise ValueError("mixed boom") - return x - - pipeline = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(lambda x: x, executor=pool.get_executor(), concurrency=2) - .pipe(fail_on_five, executor=regular, concurrency=1, max_failures=0) - .add_sink(3) - .build(num_threads=1) - ) - - with self.assertRaises(PipelineFailure): - with pipeline.auto_stop(): - list(pipeline.get_iterator(timeout=10)) - - pool.shutdown() - regular.shutdown() - - -# ─── Pickle support ─── - - -class TestPriorityExecutorPickle(unittest.TestCase): - def test_thread_executor_round_trip(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=4) - data = pickle.dumps(pool) - pool.shutdown() - - restored = pickle.loads(data) - stage = restored.get_executor() - fut = stage.submit(pow, 2, 10) - self.assertEqual(fut.result(timeout=5), 1024) - restored.shutdown() - - def test_entrypoint_round_trip(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=2) - ep = pool.get_executor(priority=5) - data = pickle.dumps(ep) - pool.shutdown() - - restored_ep = pickle.loads(data) - fut = restored_ep.submit(pow, 2, 3) - self.assertEqual(fut.result(timeout=5), 8) - restored_ep._owner.shutdown() - - def test_multiple_entrypoints_share_master(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=2) - ep1 = pool.get_executor(priority=1) - ep2 = pool.get_executor(priority=2) - - data1 = pickle.dumps(ep1) - data2 = pickle.dumps(ep2) - - # Remove original pool from registry to simulate new process - pool_id = pool._id - pool.shutdown() - _OWNER_REGISTRY.pop(pool_id, None) - - restored1 = pickle.loads(data1) - restored2 = pickle.loads(data2) - - self.assertIs(restored1._owner, restored2._owner) - - fut1 = restored1.submit(pow, 2, 3) - fut2 = restored2.submit(pow, 3, 2) - self.assertEqual(fut1.result(timeout=5), 8) - self.assertEqual(fut2.result(timeout=5), 9) - restored1._owner.shutdown() - - @_ignore_fork_warning - def test_process_executor_round_trip(self) -> None: - pool = PriorityProcessPoolExecutor(max_workers=2) - data = pickle.dumps(pool) - pool.shutdown() - - restored = pickle.loads(data) - stage = restored.get_executor() - fut = stage.submit(pow, 2, 10) - self.assertEqual(fut.result(timeout=10), 1024) - restored.shutdown() - - def test_entrypoint_is_executor_subclass_after_unpickle(self) -> None: - pool = PriorityThreadPoolExecutor(max_workers=1) - ep = pool.get_executor() - data = pickle.dumps(ep) - pool.shutdown() - - restored = pickle.loads(data) - self.assertIsInstance(restored, Executor) - self.assertIsInstance(restored, PriorityExecutorEntrypoint) - restored._owner.shutdown() - - def test_entrypoint_priority_preserved(self) -> None: - """Priority ordering is preserved across pickle round-trip.""" - pool = PriorityThreadPoolExecutor(max_workers=1) - upstream = pool.get_executor(priority=0) - downstream = pool.get_executor(priority=2) - - data_up = pickle.dumps(upstream) - data_down = pickle.dumps(downstream) - pool.shutdown() - _OWNER_REGISTRY.pop(pool._id, None) - - restored_up = pickle.loads(data_down) - restored_down = pickle.loads(data_up) - - # Submit and wait for results - fut_up = restored_up.submit(pow, 2, 3) - fut_down = restored_down.submit(pow, 3, 2) - self.assertEqual(fut_up.result(timeout=5), 8) - self.assertEqual(fut_down.result(timeout=5), 9) - - # Verify both share the same owner - self.assertIs(restored_up._owner, restored_down._owner) - # Verify priority values were preserved - self.assertEqual(restored_up._priority, -2) - self.assertEqual(restored_down._priority, 0) - restored_up._owner.shutdown() - - -# ─── Garbage collection ─── - - -class TestPriorityExecutorGarbageCollection(unittest.TestCase): - def test_owner_gc_after_all_references_dropped(self) -> None: - """Owner is garbage-collected once all entrypoints and user refs are gone.""" - pool = PriorityThreadPoolExecutor(max_workers=1) - owner_id = pool._id - weak = weakref.ref(pool) - stage = pool.get_executor() - - self.assertIn(owner_id, _OWNER_REGISTRY) - self.assertIsNotNone(weak()) - - # Drop the user reference — entrypoint still holds a strong ref - del pool - gc.collect() - self.assertIsNotNone(weak()) - - # Drop the entrypoint — last strong ref gone - del stage - gc.collect() - self.assertIsNone(weak()) - self.assertNotIn(owner_id, _OWNER_REGISTRY) - - def test_owner_gc_multiple_entrypoints(self) -> None: - """Owner survives until ALL entrypoints are dropped.""" - pool = PriorityThreadPoolExecutor(max_workers=1) - weak = weakref.ref(pool) - s1 = pool.get_executor() - s2 = pool.get_executor() - del pool - gc.collect() - - self.assertIsNotNone(weak()) - del s1 - gc.collect() - self.assertIsNotNone(weak()) - del s2 - gc.collect() - self.assertIsNone(weak()) - - def test_owner_gc_after_shutdown(self) -> None: - """Shutdown + dropping all refs allows GC.""" - pool = PriorityThreadPoolExecutor(max_workers=1) - weak = weakref.ref(pool) - stage = pool.get_executor() - fut = stage.submit(lambda: 42) - self.assertEqual(fut.result(timeout=5), 42) - - pool.shutdown() - del pool - gc.collect() - # Entrypoint still holds a strong ref - self.assertIsNotNone(weak()) - - del stage - gc.collect() - self.assertIsNone(weak()) - - -# ─── PriorityInterpreterPoolExecutor tests (Python 3.14+ only) ─── - -_has_interpreter_pool: bool = sys.version_info >= (3, 14) - - -@unittest.skipUnless( - _has_interpreter_pool, "InterpreterPoolExecutor requires Python 3.14+" -) -class TestPriorityInterpreterPoolExecutor(unittest.TestCase): - def test_basic_execution(self) -> None: - from spdl.pipeline import PriorityInterpreterPoolExecutor - - executor = PriorityInterpreterPoolExecutor(max_workers=2) - stage = executor.get_executor() - futs = [stage.submit(pow, 2, x) for x in range(10)] - results = {f.result(timeout=10) for f in futs} - self.assertEqual(results, {2**x for x in range(10)}) - executor.shutdown() - - def test_exception_propagation(self) -> None: - from spdl.pipeline import PriorityInterpreterPoolExecutor - - executor = PriorityInterpreterPoolExecutor(max_workers=1) - stage = executor.get_executor() - - def fail() -> None: - raise ValueError("interpreter boom") - - fut = stage.submit(fail) - with self.assertRaises(ValueError): - fut.result(timeout=10) - executor.shutdown() - - def test_submit_after_shutdown_raises(self) -> None: - from spdl.pipeline import PriorityInterpreterPoolExecutor - - executor = PriorityInterpreterPoolExecutor(max_workers=1) - stage = executor.get_executor() - executor.shutdown() - with self.assertRaises(RuntimeError): - executor._submit_with_priority( # pyre-ignore[16] - (0, 0), lambda: None, (), {} - ) - - def test_is_executor_compatible(self) -> None: - from spdl.pipeline import PriorityInterpreterPoolExecutor - - executor = PriorityInterpreterPoolExecutor(max_workers=1) - stage = executor.get_executor() - self.assertIsInstance(stage, Executor) - executor.shutdown() - - def test_pipeline_two_stage(self) -> None: - from spdl.pipeline import PriorityInterpreterPoolExecutor - - pool = PriorityInterpreterPoolExecutor(max_workers=4) - s1 = pool.get_executor() - s2 = pool.get_executor() - - pipeline = ( - PipelineBuilder() - .add_source(range(20)) - .pipe(lambda x: x * 2, executor=s1, concurrency=4) - .pipe(lambda x: x + 1, executor=s2, concurrency=4) - .add_sink(3) - .build(num_threads=1) - ) - - with pipeline.auto_stop(): - results = sorted(pipeline.get_iterator(timeout=10)) - - self.assertEqual(results, sorted(x * 2 + 1 for x in range(20))) - pool.shutdown() diff --git a/tests/pipeline/priority_executor_test.py b/tests/pipeline/priority_executor_test.py new file mode 120000 index 000000000..c287fb222 --- /dev/null +++ b/tests/pipeline/priority_executor_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/priority_executor_test.py \ No newline at end of file diff --git a/tests/pipeline/source_locator_test.py b/tests/pipeline/source_locator_test.py deleted file mode 100644 index fb4e1ba2f..000000000 --- a/tests/pipeline/source_locator_test.py +++ /dev/null @@ -1,200 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -""" -Unit tests for the source_locator module. -""" - -import asyncio -import functools -import inspect -import unittest -from typing import Any - -from spdl.pipeline._common._source_locator import locate_source - - -def regular_function(x: int, y: int) -> int: - """A regular function for testing.""" - return x + y - - -def generator_function(n: int) -> Any: - """A generator function for testing.""" - for i in range(n): - yield i - - -async def async_function(x: int) -> int: - """An async function for testing.""" - await asyncio.sleep(0) - return x * 2 - - -async def async_generator_function(n: int) -> Any: - """An async generator function for testing.""" - for i in range(n): - await asyncio.sleep(0) - yield i - - -class SimpleCallable: - """A simple callable class for testing.""" - - def __call__(self, x: int) -> int: - return x * 3 - - -def _ln(target: object) -> int: - return inspect.getsourcelines(target)[1] # pyre-ignore[6] - - -class TestSourceLocator(unittest.TestCase): - """Test cases for the locate_source function.""" - - def test_regular_function(self) -> None: - """Test locating source for a regular function.""" - loc = locate_source(regular_function) - - self.assertEqual(loc.name, f"{__name__}.regular_function") - self.assertEqual(loc.file_path, __file__) - self.assertEqual(loc.line_number, _ln(regular_function)) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_generator_function(self) -> None: - """Test locating source for a generator function.""" - loc = locate_source(generator_function) - - self.assertEqual(loc.name, f"{__name__}.generator_function") - self.assertEqual(loc.file_path, __file__) - self.assertEqual(loc.line_number, _ln(generator_function)) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_async_function(self) -> None: - """Test locating source for an async function.""" - loc = locate_source(async_function) - - self.assertEqual(loc.name, f"{__name__}.async_function") - self.assertEqual(loc.file_path, __file__) - self.assertEqual(loc.line_number, _ln(async_function)) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_async_generator_function(self) -> None: - """Test locating source for an async generator function.""" - loc = locate_source(async_generator_function) - - self.assertEqual(loc.name, f"{__name__}.async_generator_function") - self.assertEqual(loc.file_path, __file__) - self.assertEqual(loc.line_number, _ln(async_generator_function)) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_callable_class_object(self) -> None: - """Test locating source for a callable class object.""" - obj = SimpleCallable() - loc = locate_source(obj) - - self.assertEqual(loc.name, f"{__name__}.SimpleCallable") - self.assertEqual(loc.file_path, __file__) - self.assertEqual(loc.line_number, _ln(SimpleCallable)) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_builtin_function(self) -> None: - """Test locating source for a built-in function.""" - loc = locate_source(len) - - self.assertEqual(loc.name, "builtins.len") - self.assertIsNone(loc.file_path) - self.assertIsNone(loc.line_number) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {}) - - def test_partial_with_positional_args(self) -> None: - """ - Test locating source for a function wrapped with functools.partial - (positional args). - """ - partial_func = functools.partial(regular_function, 5) - loc = locate_source(partial_func) - - self.assertEqual(loc.name, f"{__name__}.regular_function") - self.assertEqual(loc.file_path, __file__) - self.assertIsNotNone(loc.line_number) - self.assertGreater(loc.line_number, 0) - self.assertEqual(loc.partial_args, (5,)) - self.assertEqual(loc.partial_kwargs, {}) - - def test_partial_with_keyword_args(self) -> None: - """ - Test locating source for a function wrapped with functools.partial - (keyword args). - """ - partial_func = functools.partial(regular_function, y=10) - loc = locate_source(partial_func) - - self.assertEqual(loc.name, f"{__name__}.regular_function") - self.assertEqual(loc.file_path, __file__) - self.assertIsNotNone(loc.line_number) - self.assertGreater(loc.line_number, 0) - self.assertEqual(loc.partial_args, ()) - self.assertEqual(loc.partial_kwargs, {"y": 10}) - - def test_callable_object_wrapped_with_partial(self) -> None: - """ - Test locating source for a callable object wrapped with - functools.partial. - """ - obj = SimpleCallable() - partial_obj = functools.partial(obj, 7) - loc = locate_source(partial_obj) - - self.assertEqual(loc.name, f"{__name__}.SimpleCallable") - self.assertEqual(loc.file_path, __file__) - self.assertIsNotNone(loc.line_number) - self.assertGreater(loc.line_number, 0) - self.assertEqual(loc.partial_args, (7,)) - self.assertEqual(loc.partial_kwargs, {}) - - def test_nested_partial(self) -> None: - """Test locating source for nested functools.partial wrapping.""" - partial_func1 = functools.partial(regular_function, 3) - partial_func2 = functools.partial(partial_func1, y=8) - loc = locate_source(partial_func2) - - self.assertEqual(loc.name, f"{__name__}.regular_function") - self.assertEqual(loc.file_path, __file__) - self.assertIsNotNone(loc.line_number) - self.assertGreater(loc.line_number, 0) - self.assertEqual(loc.partial_args, (3,)) - self.assertEqual(loc.partial_kwargs, {"y": 8}) - - def test_nested_partial_positional_args_order(self) -> None: - """Test that nested partial positional args are in correct order.""" - # partial(partial(f, 1), 2) should produce args (1, 2) - partial_func1 = functools.partial(regular_function, 1) - partial_func2 = functools.partial(partial_func1, 2) - loc = locate_source(partial_func2) - - self.assertEqual(loc.partial_args, (1, 2)) - # Verify the actual behavior matches - self.assertEqual(partial_func2(), regular_function(1, 2)) - - def test_nested_partial_keyword_override(self) -> None: - """Test that outer partial keywords override inner ones.""" - # partial(partial(f, x=1), x=2) should use x=2 - partial_func1 = functools.partial(regular_function, x=1, y=5) - partial_func2 = functools.partial(partial_func1, x=2) - loc = locate_source(partial_func2) - - self.assertEqual(loc.partial_kwargs, {"x": 2, "y": 5}) - # Verify the actual behavior matches - self.assertEqual(partial_func2(), regular_function(x=2, y=5)) diff --git a/tests/pipeline/source_locator_test.py b/tests/pipeline/source_locator_test.py new file mode 120000 index 000000000..975aecf0e --- /dev/null +++ b/tests/pipeline/source_locator_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/source_locator_test.py \ No newline at end of file diff --git a/tests/pipeline/subinterpreter_test.py b/tests/pipeline/subinterpreter_test.py deleted file mode 100644 index 1f07a93a5..000000000 --- a/tests/pipeline/subinterpreter_test.py +++ /dev/null @@ -1,278 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import sys -import unittest -from collections.abc import Callable, Iterable, Iterator, Sequence -from typing import Generic, TypeVar - -from parameterized import parameterized -from spdl.pipeline import PipelineBuilder, run_pipeline_in_subinterpreter -from spdl.pipeline._components import _get_global_id, _set_global_id -from spdl.pipeline._iter_utils._subinterpreter import ( - iterate_in_subinterpreter as _iterate_in_subinterpreter, -) - -T = TypeVar("T") - - -def iterate_in_subinterpreter( - fn: Callable[[], Iterable[T]], - *, - buffer_size: int = 3, - initializer: Callable[[], None] | Sequence[Callable[[], None]] | None = None, - timeout: float = 5, -) -> Iterable[T]: - """Set timeout for unittest""" - return _iterate_in_subinterpreter( - fn, buffer_size=buffer_size, initializer=initializer, timeout=timeout - ) - - -class _Wrap(Generic[T]): - """Helper class to wrap an iterable as a callable. - - This class wraps an iterable object and makes it callable, which is useful - for testing iterate_in_subinterpreter. It can optionally execute a pre-flight - function before returning the iterable, allowing for assertions to verify - that initializers have run correctly in the subinterpreter. - - Args: - obj: The iterable object to wrap. - pre: Optional callable to execute before returning the iterable. - Typically used for assertions in tests. - """ - - def __init__(self, obj: Iterable[T], pre: Callable[[], None] | None = None) -> None: - self.obj = obj - self.pre = pre - - def __call__(self) -> Iterable[T]: - if self.pre is not None: - self.pre() - return self.obj - - -_FLAGS: list[int] = [] - - -def _init_flag0() -> None: - _FLAGS.append(0) - - -def _init_flag1() -> None: - _FLAGS.append(1) - - -def _check_flag0and1() -> None: - ref = [0, 1] - assert _FLAGS == ref, f"{_FLAGS=} != {ref=}" - - -if sys.version_info >= (3, 14): - - class TestIterateInSubinterpreter(unittest.TestCase): - """Test cases for iterate_in_subinterpreter function.""" - - @parameterized.expand( - [ - ("basic_iteration", list(range(5))), - ("string_iteration", ["hello", "world", "test"]), - ("empty_iterator", []), - ], - ) - def test_iteration(self, name: str, ref: list[object]) -> None: # noqa: ARG002 - """Test iteration with various input types.""" - iterable = iterate_in_subinterpreter(_Wrap(ref)) - result = list(iterable) - self.assertEqual(result, ref) - result2 = list(iterable) - self.assertEqual(result2, ref) - - def test_buffer_size(self) -> None: - """Test with custom buffer size.""" - ref = list(range(10)) - result = list(iterate_in_subinterpreter(_Wrap(ref), buffer_size=5)) - self.assertEqual(result, ref) - - def test_with_initializers(self) -> None: - """Test with multiple initializer functions.""" - ref = list(range(10)) - result = list( - iterate_in_subinterpreter( - _Wrap(ref, pre=_check_flag0and1), - initializer=[_init_flag0, _init_flag1], - timeout=5.0, - ) - ) - self.assertEqual(result, ref) - # The flag should not be set in the main interpreter - self.assertEqual(_FLAGS, []) - - def test_partial_iteration(self) -> None: - """Test partial iteration by breaking early.""" - iterable = iterate_in_subinterpreter(_Wrap(range(10))) - result = [] - for i, item in enumerate(iterable): - result.append(item) - if i >= 4: - break - self.assertEqual(result, [0, 1, 2, 3, 4]) - - -# Module-level functions and classes (required for pickling/subinterpreter compatibility) -def _double(x: int) -> int: - """Helper function to double a value.""" - return x * 2 - - -def _only_even(x: int) -> int | None: - """Helper function to filter only even numbers.""" - return x if x % 2 == 0 else None - - -class _StatefulSource: - """Stateful source that tracks iteration calls.""" - - def __init__(self, n: int) -> None: - self.n = n - self.calls = 0 - - def __iter__(self) -> Iterator[int]: - start = self.calls * self.n - self.calls += 1 - yield from range(start, start + self.n) - - -class _validate_pipeline_id: - """Helper class to validate that the pipeline ID is as expected.""" - - def __init__(self, val: int) -> None: - self.val = val - - def __iter__(self) -> Iterator[int]: - if (v := _get_global_id()) != self.val: - raise AssertionError(f"_node._PIPELINE_ID={v} != {self.val=}") - yield 0 - - -if sys.version_info >= (3, 14): - - class TestRunPipelineInSubinterpreter(unittest.TestCase): - """Test cases for run_pipeline_in_subinterpreter function.""" - - def test_basic_pipeline(self) -> None: - """Test basic pipeline execution in subinterpreter.""" - condig = PipelineBuilder().add_source(range(5)).add_sink().get_config() - iterable = run_pipeline_in_subinterpreter(condig, num_threads=1, timeout=5) - result = list(iterable) - self.assertEqual(result, [0, 1, 2, 3, 4]) - - def test_pipeline_with_pipe(self) -> None: - """Test pipeline with pipe operation in subinterpreter.""" - config = ( - PipelineBuilder() - .add_source(range(5)) - .pipe(_double) - .add_sink() - .get_config() - ) - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - result = list(iterable) - self.assertEqual(result, [0, 2, 4, 6, 8]) - - def test_pipeline_multiple_iterations(self) -> None: - """Test that the pipeline can be iterated multiple times.""" - config = ( - PipelineBuilder() - .add_source(_StatefulSource(3)) - .add_sink(buffer_size=10) - .get_config() - ) - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - - # First iteration - result1 = list(iterable) - self.assertEqual(result1, [0, 1, 2]) - - # Second iteration should start from where we left off - result2 = list(iterable) - self.assertEqual(result2, [3, 4, 5]) - - def test_pipeline_with_aggregate(self) -> None: - """Test pipeline with aggregation in subinterpreter.""" - config = ( - PipelineBuilder() - .add_source(range(10)) - .aggregate(3) - .add_sink(buffer_size=10) - .get_config() - ) - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - result = list(iterable) - self.assertEqual(result, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) - - def test_pipeline_empty_source(self) -> None: - """Test pipeline with empty source in subinterpreter.""" - config = PipelineBuilder().add_source([]).add_sink().get_config() - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - result = list(iterable) - self.assertEqual(result, []) - - def test_pipeline_with_filter(self) -> None: - """Test pipeline with filter operation (returning None to skip items).""" - config = ( - PipelineBuilder() - .add_source(range(10)) - .pipe(_only_even) - .add_sink() - .get_config() - ) - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - result = list(iterable) - self.assertEqual(result, [0, 2, 4, 6, 8]) - - def test_pipeline_with_buffer_size(self) -> None: - """Test run_pipeline_in_subinterpreter with custom buffer_size.""" - config = PipelineBuilder().add_source(range(10)).add_sink().get_config() - iterable = run_pipeline_in_subinterpreter( - config, num_threads=1, buffer_size=5, timeout=5 - ) - result = list(iterable) - self.assertEqual(result, list(range(10))) - - def test_run_pipeline_in_subinterpreter_pipeline_id(self) -> None: - """Test pipeline inherits global ID in subinterpreter.""" - # Set to a number that's not zero and something unlikely to - # happen during testing - _set_global_id(123456) - ref = _get_global_id() + 1 - - config = ( - PipelineBuilder() - .add_source(_validate_pipeline_id(ref)) - .add_sink() - .get_config() - ) - - iterable = run_pipeline_in_subinterpreter(config, num_threads=1, timeout=5) - - for _ in iterable: - pass - -else: - - class TestRunPipelineInSubinterpreter(unittest.TestCase): - """Placeholder tests for Python < 3.14.""" - - def test_requires_python_3_14(self) -> None: - """Test that run_pipeline_in_subinterpreter requires Python 3.14+.""" - config = PipelineBuilder().add_source([1, 2, 3]).add_sink().get_config() - with self.assertRaises(RuntimeError) as cm: - run_pipeline_in_subinterpreter(config, num_threads=1) - self.assertIn("Python 3.14", str(cm.exception)) diff --git a/tests/pipeline/subinterpreter_test.py b/tests/pipeline/subinterpreter_test.py new file mode 120000 index 000000000..2dfba40c7 --- /dev/null +++ b/tests/pipeline/subinterpreter_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/subinterpreter_test.py \ No newline at end of file diff --git a/tests/pipeline/subprocess_break_reiterate_test.py b/tests/pipeline/subprocess_break_reiterate_test.py deleted file mode 100644 index 6bf18e1c1..000000000 --- a/tests/pipeline/subprocess_break_reiterate_test.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -"""Regression test for D101554675: breaking out of a subprocess iterable -must not kill the worker, so subsequent iterations still work.""" - -import functools -import unittest -import warnings -from collections.abc import Iterator -from functools import partial - -from spdl.pipeline import iterate_in_subprocess - - -def _ignore_fork_warning(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - return wrapper - - -def _ignore_fork_warning_in_class(cls): - for name, member in list(vars(cls).items()): - if name.startswith("test_") and callable(member): - setattr(cls, name, _ignore_fork_warning(member)) - return cls - - -class SourceIterable: - def __init__(self, n: int) -> None: - self.n = n - - def __iter__(self) -> Iterator[int]: - yield from range(self.n) - - -@_ignore_fork_warning_in_class -class TestSubprocessBreakAndReiterate(unittest.TestCase): - def test_break_then_reiterate(self) -> None: - """Breaking out of a subprocess iterable must not prevent re-iteration. - - This is a regression test for the BaseException widening in D101554675. - When a consumer `break`s out of `for ... in iterable`, Python sends - GeneratorExit into the generator. Prior to D101554675, only - (Exception, KeyboardInterrupt) triggered _shutdown(). After D101554675, - BaseException (which includes GeneratorExit) triggers _shutdown(), - making subsequent iter() calls raise RuntimeError. - """ - src = iterate_in_subprocess(partial(SourceIterable, 10), timeout=10) - - # First iteration: consume only 3 items, then break - count = 0 - for _item in src: - count += 1 - if count >= 3: - break - - # Second iteration: must succeed (worker should still be alive) - result = list(src) - self.assertEqual(result, list(range(10))) - - def test_break_then_reiterate_multiple_times(self) -> None: - """Multiple break-then-reiterate cycles must all succeed.""" - src = iterate_in_subprocess(partial(SourceIterable, 5), timeout=10) - - for cycle in range(3): - # Break after 2 items - count = 0 - for _item in src: - count += 1 - if count >= 2: - break - - # Full iteration must still work - result = list(src) - self.assertEqual(result, list(range(5)), f"cycle {cycle}") - - def test_partial_iteration_via_zip(self) -> None: - """Partial iteration via zip() (which breaks implicitly) must not kill worker.""" - src = iterate_in_subprocess(partial(SourceIterable, 100), timeout=10) - - # zip stops when the shorter iterable is exhausted, causing - # GeneratorExit on the longer one - partial_result = list(zip(range(3), src)) - self.assertEqual(partial_result, [(0, 0), (1, 1), (2, 2)]) - - # Subsequent full iteration must work - result = list(src) - self.assertEqual(result, list(range(100))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/pipeline/subprocess_break_reiterate_test.py b/tests/pipeline/subprocess_break_reiterate_test.py new file mode 120000 index 000000000..1c9e5ed93 --- /dev/null +++ b/tests/pipeline/subprocess_break_reiterate_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/subprocess_break_reiterate_test.py \ No newline at end of file diff --git a/tests/pipeline/subprocess_test.py b/tests/pipeline/subprocess_test.py deleted file mode 100644 index 36149bb59..000000000 --- a/tests/pipeline/subprocess_test.py +++ /dev/null @@ -1,506 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -import functools -import multiprocessing as mp -import os.path -import random -import tempfile -import threading -import time -import unittest -import warnings -from collections.abc import Iterable, Iterator -from functools import partial - -from spdl.pipeline import iterate_in_subprocess as _iterate_in_subprocess -from spdl.pipeline._iter_utils._common import _Cmd, _execute_iterable, _Status - - -def _ignore_fork_warning(fn): - """Suppress the multi-threaded fork() DeprecationWarning emitted by - multiprocessing.popen_fork when starting subprocesses while pipeline - worker threads are alive. The warning is intentional in these tests. - """ - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - r"This process \(pid=\d+\) is multi-threaded, use of " - r"fork\(\) may lead to deadlocks in the child" - ), - category=DeprecationWarning, - ) - return fn(*args, **kwargs) - - return wrapper - - -def _ignore_fork_warning_in_class(cls): - """Apply ``_ignore_fork_warning`` to every ``test_*`` method on a class.""" - for name, member in list(vars(cls).items()): - if name.startswith("test_") and callable(member): - setattr(cls, name, _ignore_fork_warning(member)) - return cls - - -def iterate_in_subprocess(fn, *, timeout=10, **kwargs): - return _iterate_in_subprocess(fn, timeout=timeout, **kwargs) - - -def iter_range(n: int) -> Iterable[int]: - yield from range(n) - - -def initializer(path: str, val: str) -> None: - with open(path, "w") as f: - f.write(val) - - -@_ignore_fork_warning_in_class -class TestIterateInSubprocess(unittest.TestCase): - def test_iterate_in_subprocess(self) -> None: - """iterate_in_subprocess iterates""" - N = 10 - - src = iterate_in_subprocess(fn=partial(iter_range, n=N)) - self.assertEqual(list(src), list(range(N))) - - def test_iterate_in_subprocess_initializer(self) -> None: - """iterate_in_subprocess initializer is called before iteration starts""" - - N = 10 - val = str(random.random()) - with tempfile.TemporaryDirectory() as dir: - path = os.path.join(dir, "foo.txt") - - self.assertFalse(os.path.exists(path)) - src = iterate_in_subprocess( - fn=partial(iter_range, n=N), - initializer=partial(initializer, path=path, val=val), - buffer_size=1, - ) - self.assertTrue(os.path.exists(path)) - - ite = iter(src) - self.assertEqual(next(ite), 0) - - with open(path, "r") as f: - self.assertEqual(f.read(), val) - - for i in range(1, N): - self.assertEqual(next(ite), i) - - with self.assertRaises(StopIteration): - next(ite) - - def test_iterate_in_subprocess_multiple_initializer(self) -> None: - """iterate_in_subprocess accepts multiple iterators""" - N = 10 - val1 = str(random.random()) - val2 = str(random.random()) - with tempfile.TemporaryDirectory() as dir: - path1 = os.path.join(dir, "foo.txt") - path2 = os.path.join(dir, "bar.txt") - - self.assertFalse(os.path.exists(path1)) - self.assertFalse(os.path.exists(path2)) - src = iterate_in_subprocess( - fn=partial(iter_range, n=N), - initializer=[ - partial(initializer, path=path1, val=val1), - partial(initializer, path=path2, val=val2), - ], - buffer_size=1, - ) - self.assertTrue(os.path.exists(path1)) - self.assertTrue(os.path.exists(path2)) - - ite = iter(src) - self.assertEqual(next(ite), 0) - - with open(path1, "r") as f: - self.assertEqual(f.read(), val1) - - with open(path2, "r") as f: - self.assertEqual(f.read(), val2) - - for i in range(1, N): - self.assertEqual(next(ite), i) - - with self.assertRaises(StopIteration): - next(ite) - - -def iter_range_and_store_with_sync(n: int, sync_queue: mp.Queue) -> Iterable[int]: - """Generator that synchronizes with main process via queue.""" - yield 0 - for i in range(n): - yield i - # Signal main process that we've yielded this value - sync_queue.put(i) - - -@_ignore_fork_warning_in_class -class TestIterateInSubprocessBufferSize(unittest.TestCase): - def test_iterate_in_subprocess_buffer_size_1(self) -> None: - """buffer_size=1 makes iterate_in_subprocess works sort-of interactively""" - - N = 10 - - # Use queue for synchronization between processes - sync_queue = mp.Queue() - - src = iterate_in_subprocess( - fn=partial(iter_range_and_store_with_sync, n=N, sync_queue=sync_queue), - daemon=True, - buffer_size=1, - ) - ite = iter(src) - self.assertEqual(next(ite), 0) - - for i in range(N): - # Wait for subprocess to signal it has yielded value i - # Use timeout to avoid hanging if subprocess fails - subprocess_value = sync_queue.get(timeout=5) - self.assertEqual( - subprocess_value, i, f"Expected {i}, got {subprocess_value}" - ) - - # With buffer_size=1, the queue should be empty after fetching - self.assertTrue( - sync_queue.empty(), f"Queue should be empty after fetching item {i}" - ) - - # Now fetch the value from the iterator - self.assertEqual(next(ite), i) - - with self.assertRaises(StopIteration): - next(ite) - - def test_iterate_in_subprocess_buffer_size_64(self) -> None: - """big buffer_size makes iterate_in_subprocess processes data in one go""" - - N = 10 - - # Use queue for synchronization between processes - sync_queue = mp.Queue() - - src = iterate_in_subprocess( - fn=partial(iter_range_and_store_with_sync, n=N, sync_queue=sync_queue), - daemon=True, - buffer_size=64, - ) - ite = iter(src) - self.assertEqual(next(ite), 0) - - # With buffer_size=64, subprocess should process all data without waiting - # Wait for subprocess to signal all values have been processed - for expected_i in range(N): - # Use timeout to avoid hanging - subprocess_value = sync_queue.get(timeout=5) - self.assertEqual( - subprocess_value, - expected_i, - f"Expected {expected_i}, got {subprocess_value}", - ) - - # Now all data should be available in the buffer, fetch them - for i in range(N): - self.assertEqual(next(ite), i) - - with self.assertRaises(StopIteration): - next(ite) - - -class SourceIterable: - def __init__(self, n: int) -> None: - self.n = n - - def __iter__(self) -> Iterator[int]: - yield from range(self.n) - - -def noop() -> None: - pass - - -class TestExecuteIterable(unittest.TestCase): - def test_execute_iterable_initializer_failure(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - def src_fn() -> Iterable[int]: - return SourceIterable(10) - - def fail() -> None: - raise ValueError("Failed!") - - _execute_iterable(msg_queue, data_queue, src_fn, [fail]) - - self.assertTrue(msg_queue.empty()) - - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.INITIALIZATION_FAILED) - self.assertIn("Failed!", result.message) - self.assertTrue(data_queue.empty()) - - def test_execute_iterable_iterator_initialize_failure(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - def src_fn() -> Iterator[int]: - raise ValueError("Failed!") - return SourceIterable(10) - - _execute_iterable(msg_queue, data_queue, src_fn, [noop]) - - self.assertTrue(msg_queue.empty()) - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.INITIALIZATION_FAILED) - self.assertIn("Failed!", result.message) - self.assertTrue(data_queue.empty()) - - def test_execute_iterable_quite_immediately(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - msg_queue.put(_Cmd.ABORT) - time.sleep(1) - - def src_fn() -> Iterable[int]: - return SourceIterable(10) - - _execute_iterable(msg_queue, data_queue, src_fn, [noop]) - time.sleep(1) - - self.assertTrue(msg_queue.empty()) - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) - self.assertTrue(data_queue.empty()) - - def test_execute_iterable_generator_fail(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - class SourceIterableFails(SourceIterable): - def __iter__(self) -> Iterator[int]: - raise ValueError("Failed!") - yield from range(self.n) - - def src_fn() -> Iterable[int]: - return SourceIterableFails(10) - - msg_queue.put(_Cmd.START_ITERATION) - _execute_iterable(msg_queue, data_queue, src_fn, [noop]) - - self.assertTrue(msg_queue.empty()) - - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.ITERATION_STARTED) - - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.ITERATOR_FAILED) - self.assertIn("Failed!", result.message) - self.assertTrue(data_queue.empty()) - - def test_execute_iterable_generator_fail_after_n(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - class SourceIterableFails(SourceIterable): - def __iter__(self) -> Iterator[int]: - for v in range(self.n): - yield v - if v == 2: - raise ValueError("Failed!") - - def src_fn() -> Iterable[int]: - return SourceIterableFails(10) - - msg_queue.put(_Cmd.START_ITERATION) - _execute_iterable(msg_queue, data_queue, src_fn, [noop]) - - self.assertTrue(msg_queue.empty()) - - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.ITERATION_STARTED) - for i in range(3): - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.ITERATOR_SUCCESS) - self.assertEqual(result.message, i) - - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.ITERATOR_FAILED) - self.assertIn("Failed!", result.message) - self.assertTrue(data_queue.empty()) - - def test_execute_iterator_generator_success(self) -> None: - msg_queue, data_queue = mp.Queue(), mp.Queue() - - def src_fn() -> Iterable[int]: - return SourceIterable(3) - - msg_queue.put(_Cmd.START_ITERATION) - - # Add abort with delay, so that _execute_iterable can exit after - # the iteration - def done(): - time.sleep(3) - msg_queue.put(_Cmd.ABORT) - - t = threading.Thread(target=done) - t.start() - _execute_iterable(msg_queue, data_queue, src_fn, [noop]) - t.join() - - self.assertTrue(msg_queue.empty()) - - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.INITIALIZATION_SUCCEEDED) - ack = data_queue.get(timeout=1) - self.assertEqual(ack.status, _Status.ITERATION_STARTED) - for i in range(3): - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.ITERATOR_SUCCESS) - self.assertEqual(result.message, i) - - result = data_queue.get(timeout=1) - self.assertEqual(result.status, _Status.ITERATION_FINISHED) - - -def _src1() -> Iterable[int]: - return SourceIterable(10) - - -def _init1() -> None: - raise ValueError("Failed!") - - -def _src2() -> Iterator[int]: - if True: - raise ValueError("Failed!") - return SourceIterable(10) - - -def _src3() -> Iterable[int]: - class SourceIterableFails(SourceIterable): - def __iter__(self) -> Iterator[int]: - raise ValueError("Failed!") - yield from range(self.n) - - return SourceIterableFails(10) - - -def _src4() -> Iterable[int]: - class SourceIterableFails(SourceIterable): - def __iter__(self) -> Iterator[int]: - for v in range(self.n): - yield v - if v == 2: - raise ValueError("Failed!") - - return SourceIterableFails(10) - - -def _src5(N) -> Iterable[int]: - return SourceIterable(N) - - -class SleepSourceIterable(SourceIterable): - def __iter__(self): - time.sleep(10) - yield 0 - - -def _src6() -> Iterable[int]: - return SleepSourceIterable(3) - - -def _fail_initializer(): - raise RuntimeError("Failed!") - - -_VERY_BAD_REFERENCE = None - - -@_ignore_fork_warning_in_class -class TestIterateInSubprocessFailures(unittest.TestCase): - def test_iterate_in_subprocess_initializer_failure(self) -> None: - with self.assertRaisesRegex(RuntimeError, r"Initializer failed"): - iterate_in_subprocess(_src1, buffer_size=1, timeout=3, initializer=_init1) - - def test_iterate_in_subprocess_iterator_initialize_failure(self) -> None: - with self.assertRaisesRegex(RuntimeError, r"Failed to create the iterable"): - iterate_in_subprocess(_src2, buffer_size=1, timeout=3) - - def test_iterate_in_subprocess_generator_fail(self) -> None: - ite = iter(iterate_in_subprocess(_src3, buffer_size=1, timeout=3)) - - with self.assertRaisesRegex(RuntimeError, r"Failed to fetch the next item"): - next(ite) - - def test_iterate_in_subprocess_fail_after_n(self) -> None: - ite = iter(iterate_in_subprocess(_src4, buffer_size=1, timeout=3)) - self.assertEqual(next(ite), 0) - self.assertEqual(next(ite), 1) - self.assertEqual(next(ite), 2) - - with self.assertRaisesRegex(RuntimeError, r"Failed to fetch the next item"): - next(ite) - - def test_iterate_in_subprocess_success(self) -> None: - N = 3 - - hyp = list(iterate_in_subprocess(partial(_src5, N), buffer_size=-1, timeout=3)) - self.assertEqual(hyp, list(range(N))) - - def test_iterate_in_subprocess_timeout(self) -> None: - iterable = iterate_in_subprocess(_src6, buffer_size=-1, timeout=3) - iterator = iter(iterable) - with self.assertRaisesRegex( - RuntimeError, r"The worker subprocess did not produce any data for" - ): - next(iterator) - - def test_iterate_in_subprocess_initializer_fail(self) -> None: - """The initialization failure is propagated to the main process""" - - with self.assertRaisesRegex(RuntimeError, r"Initializer failed"): - iterate_in_subprocess(SourceIterable, initializer=_fail_initializer) - - def test_iterate_in_subprocess_iterable_creation_fail(self) -> None: - """The initialization failure is propagated to the main process""" - - with self.assertRaisesRegex(RuntimeError, r"Failed to create the iterable"): - iterate_in_subprocess(SourceIterable) - - def test_iterate_in_subprocess_success_simple_iterable(self) -> None: - iterator = iterate_in_subprocess(partial(SourceIterable, 3)) - - self.assertEqual(list(iterator), [0, 1, 2]) - self.assertEqual(list(iterator), [0, 1, 2]) - self.assertEqual(list(iterator), [0, 1, 2]) - - def test_iterate_in_subprocess_fail_not_stuck(self) -> None: - """An exception does not make Python stack. - - If a (non-daemon) subprocess is launched without a context manager - that ensures its clean exit, raising an exception while the reference - to the process object is held causes the Python interpreter to get - stuck at the exit. - - To avoid this, we register atexit function, which push the ABORT - command to the command queue, which will be received by the subprocess - if the subprocess is not shut down. This test ensures that behavior. - """ - - global _VERY_BAD_REFERENCE - _VERY_BAD_REFERENCE = iterate_in_subprocess(partial(SourceIterable, 3)) diff --git a/tests/pipeline/subprocess_test.py b/tests/pipeline/subprocess_test.py new file mode 120000 index 000000000..c774f2ec2 --- /dev/null +++ b/tests/pipeline/subprocess_test.py @@ -0,0 +1 @@ +../../../src/spdl/pipeline/tests/subprocess_test.py \ No newline at end of file