diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index eb54b5148..9e22233dc 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -155,10 +155,17 @@ jobs: if: ${{ matrix.python-minor-version == '10' }} run: | python -m pip install 'torch>=2.0.0' --index-url https://download.pytorch.org/whl/cpu + - name: Upload Linux wheel for diagnostics + if: ${{ matrix.python-minor-version == '10' }} + uses: actions/upload-artifact@v7 + with: + name: python-linux-wheel + path: target/wheels/lancedb-*.whl + retention-days: 1 + if-no-files-found: error - uses: ./.github/workflows/run_tests env: PYTHONFAULTHANDLER: "1" - PYTHONMALLOC: debug with: integration: true - name: Test without pylance or pandas diff --git a/python/python/tests/test_torch.py b/python/python/tests/test_torch.py index db8982984..09a3c8de9 100644 --- a/python/python/tests/test_torch.py +++ b/python/python/tests/test_torch.py @@ -9,8 +9,11 @@ import json import multiprocessing as mp import os import pickle +import signal import sys import threading +import time +from pathlib import Path import lancedb import pyarrow as pa @@ -317,11 +320,54 @@ def test_permutation_with_builder_is_picklable(tmp_db): def _fork_dataloader_diagnostics(worker_id=None): - faulthandler.enable(all_threads=True) - faulthandler.dump_traceback_later(20) + if worker_id is None: + os.setsid() print( f"DataLoader pid={os.getpid()} worker={worker_id}", file=sys.stderr, flush=True ) + faulthandler.enable(all_threads=True) + faulthandler.register(signal.SIGUSR1, all_threads=True, chain=False) + + +def _join_fork_dataloader(proc): + deadline = time.monotonic() + 30 + proc.join(timeout=28) + if proc.is_alive(): + children = Path(f"/proc/{proc.pid}/task/{proc.pid}/children") + try: + worker_pids = [int(pid) for pid in children.read_text().split()] + except FileNotFoundError: + worker_pids = [] + for pid in [*worker_pids, proc.pid]: + print( + f"Requesting DataLoader traceback pid={pid}", + file=sys.stderr, + flush=True, + ) + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGUSR1) + time.sleep(0.1) + proc.join(timeout=max(0, deadline - time.monotonic())) + + +def _stop_fork_dataloader(proc): + try: + owns_group = os.getpgid(proc.pid) == proc.pid + except ProcessLookupError: + proc.join() + return + if owns_group: + with contextlib.suppress(ProcessLookupError): + os.killpg(proc.pid, signal.SIGTERM) + else: + proc.terminate() + proc.join(timeout=5) + if owns_group: + with contextlib.suppress(ProcessLookupError): + os.killpg(proc.pid, signal.SIGKILL) + if proc.is_alive(): + proc.kill() + proc.join() def _multiworker_dataloader_target(db_uri: str, result_queue): @@ -441,14 +487,10 @@ def test_permutation_dataloader_fork_workers(tmp_path): queue = ctx.Queue() proc = ctx.Process(target=_multiworker_dataloader_target, args=(db_uri, queue)) proc.start() - proc.join(timeout=30) + _join_fork_dataloader(proc) if proc.is_alive(): - proc.terminate() - proc.join(timeout=5) - if proc.is_alive(): - proc.kill() - proc.join() + _stop_fork_dataloader(proc) pytest.fail("Permutation hung when iterated in a fork-based DataLoader worker") assert proc.exitcode == 0, f"child exited with code {proc.exitcode}" @@ -481,14 +523,10 @@ def test_lazy_permutation_reopens_inherited_table_in_fork_worker(tmp_path): args=(db_uri, queue), ) proc.start() - proc.join(timeout=30) + _join_fork_dataloader(proc) if proc.is_alive(): - proc.terminate() - proc.join(timeout=5) - if proc.is_alive(): - proc.kill() - proc.join() + _stop_fork_dataloader(proc) pytest.fail("Lazy Permutation hung in a fork-based DataLoader worker") assert proc.exitcode == 0, f"child exited with code {proc.exitcode}" @@ -518,14 +556,10 @@ def test_remote_permutation_dataloader_fork_workers(): args=(port, queue), ) proc.start() - proc.join(timeout=30) + _join_fork_dataloader(proc) if proc.is_alive(): - proc.terminate() - proc.join(timeout=5) - if proc.is_alive(): - proc.kill() - proc.join() + _stop_fork_dataloader(proc) pytest.fail( "Remote permutation hung when iterated in a fork-based " "DataLoader worker"