From 5a8dc4fd026a7c0fe43900d205fccfb52d69dce6 Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Sun, 26 Jul 2026 13:56:56 -0500 Subject: [PATCH 1/7] ci: don't start the triton server for the job that never uses it test_triton is the only test touching the inference server and it is marked dask_client. The "without dask Client" job selects -m "not dask_client", so the server it starts is never contacted. Pulling the image there costs a multi- gigabyte download on six matrix entries and makes the job fail whenever nvcr.io is unreachable, which then cancels the rest of the matrix. Assisted-by: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b3427537..ea87f42c7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -86,10 +86,6 @@ jobs: if: matrix.os == 'windows-latest' run: | uv pip install '.[dev,caches,dask,dask-awkward]' - - name: Start triton server with example model - if: startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != '3.14t' - run: | - docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models - name: Test with pytest (without dask Client) run: | python -m pytest --cov-report=xml --cov=coffea --ignore=tests/test_taskvine_dask.py --ignore=tests/test_taskvine_virtual.py -m "not dask_client" -n 4 From a08d0c6cadf75918be7453144d15b0066e11d76d Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Sun, 26 Jul 2026 14:06:41 -0500 Subject: [PATCH 2/7] ci: tolerate an unreachable nvcr.io when starting the triton server Pulls of the triton image intermittently fail with a connection timeout to nvcr.io, which failed the whole job before any test ran. Retry the start three times with a widening delay, and on exhaustion record TRITON_UNAVAILABLE so test_triton skips instead. The skip is keyed on that flag rather than on server reachability, so it applies only when the image could not be fetched. A server that starts and then misbehaves still fails the test. Assisted-by: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 14 +++++++++++++- tests/test_ml_tools.py | 6 ++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ea87f42c7..4a08b6166 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -162,7 +162,19 @@ jobs: - name: Start triton server with example model if: startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != '3.14t' run: | - docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models + # Pulls from nvcr.io intermittently time out. Retry, and on exhaustion + # report the server as unavailable so test_triton skips rather than + # failing the job on a registry outage. A server that starts but + # misbehaves still fails the test. + for attempt in 1 2 3; do + if docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models; then + exit 0 + fi + echo "::warning::triton server start failed (attempt ${attempt}/3)" + sleep $((attempt * 15)) + done + echo "::warning::triton server unavailable after 3 attempts; test_triton will skip" + echo "TRITON_UNAVAILABLE=1" >> "$GITHUB_ENV" - name: Test with pytest (with dask Client) run: | python -m pytest --cov-report=xml --cov=coffea --ignore=tests/test_taskvine_dask.py --ignore=tests/test_taskvine_virtual.py -m "dask_client" -n 4 diff --git a/tests/test_ml_tools.py b/tests/test_ml_tools.py index e08abd30d..ca8ff1a08 100644 --- a/tests/test_ml_tools.py +++ b/tests/test_ml_tools.py @@ -1,3 +1,5 @@ +import os + import awkward as ak import numpy as np import pytest @@ -80,6 +82,10 @@ def my_pad(arr): @pytest.mark.dask_client def test_triton(tmp_path, dask_client): _ = pytest.importorskip("tritonclient") + if os.environ.get("TRITON_UNAVAILABLE") == "1": + # Set by CI when the server image could not be pulled; a server that did + # start is still expected to answer, so this does not mask real failures. + pytest.skip("triton inference server image unavailable") from coffea.ml_tools.triton_wrapper import triton_wrapper From c7b06883e95e589e7156ef572edfe9d1a23abce2 Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Mon, 27 Jul 2026 11:08:49 -0500 Subject: [PATCH 3/7] ci: widen the triton start backoff to 15s/1m/5m/20m Three attempts 15s apart cover only a 45-second outage. Five attempts spaced 15s, 1m, 5m and 20m ride out a transient registry problem while bounding a sustained one at roughly 26 minutes of waiting. Assisted-by: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4a08b6166..e3f268632 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -166,14 +166,19 @@ jobs: # report the server as unavailable so test_triton skips rather than # failing the job on a registry outage. A server that starts but # misbehaves still fails the test. - for attempt in 1 2 3; do + # Back off 15s, 1m, 5m, 20m between attempts: long enough to ride out a + # transient registry outage, bounded so a sustained one costs ~26 min. + backoff=(15 60 300 1200) + for attempt in 0 1 2 3 4; do if docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models; then exit 0 fi - echo "::warning::triton server start failed (attempt ${attempt}/3)" - sleep $((attempt * 15)) + if [ "${attempt}" -lt 4 ]; then + echo "::warning::triton server start failed (attempt $((attempt + 1))/5); retrying in ${backoff[$attempt]}s" + sleep "${backoff[$attempt]}" + fi done - echo "::warning::triton server unavailable after 3 attempts; test_triton will skip" + echo "::warning::triton server unavailable after 5 attempts; triton tests will skip" echo "TRITON_UNAVAILABLE=1" >> "$GITHUB_ENV" - name: Test with pytest (with dask Client) run: | From c409feb0eedd2bf2e8d8057735a5e47518cbd45b Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Mon, 27 Jul 2026 11:26:23 -0500 Subject: [PATCH 4/7] Revert "ci: don't start the triton server for the job that never uses it" This reverts commit 5a8dc4fd026a7c0fe43900d205fccfb52d69dce6. --- .github/workflows/ci.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e3f268632..abe574126 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -86,6 +86,10 @@ jobs: if: matrix.os == 'windows-latest' run: | uv pip install '.[dev,caches,dask,dask-awkward]' + - name: Start triton server with example model + if: startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != '3.14t' + run: | + docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models - name: Test with pytest (without dask Client) run: | python -m pytest --cov-report=xml --cov=coffea --ignore=tests/test_taskvine_dask.py --ignore=tests/test_taskvine_virtual.py -m "not dask_client" -n 4 From c76a9442fd3a85a4be714279ea25478b22800d44 Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Mon, 27 Jul 2026 11:30:11 -0500 Subject: [PATCH 5/7] test(ml_tools): cover the eager wrapper path without a dask client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every ml_tools test mixed eager and dask_awkward assertions in one dask_client marked function, so under -m "not dask_client" the eager path of the triton, torch, tensorflow and xgboost wrappers had no coverage at all — the path a user without dask takes, and the default configuration since dask became optional. Split each test in two: an unmarked eager test asserting the wrapper runs and returns the expected lengths, and the dask_client test that compares the two backends and checks necessary_columns. Wrapper construction moves into shared factories so neither copy duplicates it. dask_awkward is now imported per test instead of at module scope. Previously its absence skipped the whole module, including tests that never touch it. Assisted-by: Claude Opus 4.7 (1M context) --- tests/test_ml_tools.py | 236 ++++++++++++++++++++++++++++------------- 1 file changed, 160 insertions(+), 76 deletions(-) diff --git a/tests/test_ml_tools.py b/tests/test_ml_tools.py index ca8ff1a08..88100314c 100644 --- a/tests/test_ml_tools.py +++ b/tests/test_ml_tools.py @@ -4,12 +4,31 @@ import numpy as np import pytest -dak = pytest.importorskip("dask_awkward") +# Columns the jet wrappers below are expected to touch. +EXPECTED_JET_COLUMNS = { + "eta", + "phi", + "pfcands.pt", + "pfcands.phi", + "pfcands.eta", + "pfcands.feat1", + "pfcands.feat2", +} +XGBOOST_FEATURES = [f"feat{i}" for i in range(16)] -def prepare_jets_array(njets, tmp_path): - # Creating jagged Jet-with-constituent array, returning both awkward and lazy - # dask_awkward arrays + +def _dask_awkward(): + """dask_awkward, skipping the calling test when it is not installed. + + Imported per test rather than at module scope so the eager tests run in an + installation without dask_awkward. + """ + return pytest.importorskip("dask_awkward") + + +def prepare_jets_array(njets): + # Creating jagged Jet-with-constituent array NFEAT = 100 jets = ak.zip( { @@ -38,11 +57,14 @@ def prepare_jets_array(njets, tmp_path): pfcands = pfcands[idx < jets.ncands] jets["pfcands"] = pfcands[:] - ak_jets = jets[:] - parquet_path = str(tmp_path / "ml_tools.parquet") - ak.to_parquet(jets, parquet_path) - dak_jets = dak.from_parquet(parquet_path) - return ak_jets, dak_jets + return jets[:] + + +def to_dask(arr, path): + """Round-trip an eager array through parquet into a lazy dask_awkward array.""" + dak = _dask_awkward() + ak.to_parquet(arr, path) + return dak.from_parquet(path) def common_prepare_awkward(jets): @@ -79,9 +101,8 @@ def my_pad(arr): } -@pytest.mark.dask_client -def test_triton(tmp_path, dask_client): - _ = pytest.importorskip("tritonclient") +def _make_triton_wrapper(): + pytest.importorskip("tritonclient") if os.environ.get("TRITON_UNAVAILABLE") == "1": # Set by CI when the server image could not be pulled; a server that did # start is still expected to answer, so this does not mask real failures. @@ -97,15 +118,35 @@ def prepare_awkward(self, output_list, jets): "input_dict": common_prepare_awkward(jets), } - # Running the evaluation in lazy and non-lazy forms - tw = triton_wrapper_test( + return triton_wrapper_test( model_url="triton+grpc://localhost:8001/pn_test/1", client_args=dict( ssl=False ), # Solves SSL version mismatch for local inference server ) - ak_jets, dak_jets = prepare_jets_array(njets=256, tmp_path=tmp_path) + +def test_triton_eager(tmp_path): + tw = _make_triton_wrapper() + ak_jets = prepare_jets_array(njets=256) + + ak_res = tw(["output"], ak_jets) + for k in ak_res.keys(): + assert len(ak_res[k]) == len(ak_jets) + + # Length 0 tests + ak_res = tw(["output"], ak_jets[ak_jets.eta < 0]) + for k in ak_res.keys(): + assert len(ak_res[k]) == 0 + + +@pytest.mark.dask_client +def test_triton(tmp_path, dask_client): + dak = _dask_awkward() + tw = _make_triton_wrapper() + + ak_jets = prepare_jets_array(njets=256) + dak_jets = to_dask(ak_jets, str(tmp_path / "ml_tools.parquet")) # Vanilla awkward arrays ak_res = tw(["output"], ak_jets) @@ -113,17 +154,8 @@ def prepare_awkward(self, output_list, jets): for k in ak_res.keys(): assert ak.all(ak_res[k] == dak_res[k].compute()) - expected_columns = { - "eta", - "phi", - "pfcands.pt", - "pfcands.phi", - "pfcands.eta", - "pfcands.feat1", - "pfcands.feat2", - } columns = set(list(dak.necessary_columns(dak_res).values())[0]) - assert columns == expected_columns + assert columns == EXPECTED_JET_COLUMNS # Length 0 tests ak_res = tw(["output"], ak_jets[ak_jets.eta < 0]) @@ -132,9 +164,8 @@ def prepare_awkward(self, output_list, jets): assert len(ak_res[k]) == 0 and len(dak_res[k].compute()) == 0 -@pytest.mark.dask_client -def test_torch(tmp_path, dask_client): - _ = pytest.importorskip("torch") +def _make_torch_wrapper(**kwargs): + pytest.importorskip("torch") from coffea.ml_tools.torch_wrapper import torch_wrapper @@ -147,36 +178,48 @@ def prepare_awkward(self, jets): "mask": ak.values_astype(default["mask"], np.float16), } - tw = torch_wrapper_test("tests/samples/pn_demo.pt") - ak_jets, dak_jets = prepare_jets_array(njets=256, tmp_path=tmp_path) + return torch_wrapper_test("tests/samples/pn_demo.pt", **kwargs) + + +def test_torch_eager(tmp_path): + tw = _make_torch_wrapper() + ak_jets = prepare_jets_array(njets=256) + assert len(tw(ak_jets)) == len(ak_jets) + + # Length-0 testing + tw = _make_torch_wrapper(expected_output_shape=(None,)) + ak_jets = prepare_jets_array(njets=256) + ak_jets = ak_jets[ak_jets.eta < -100] # Mimicking a low efficiency selection + assert len(ak_jets) == 0 + assert len(tw(ak_jets)) == 0 + + +@pytest.mark.dask_client +def test_torch(tmp_path, dask_client): + dak = _dask_awkward() + tw = _make_torch_wrapper() + + ak_jets = prepare_jets_array(njets=256) + dak_jets = to_dask(ak_jets, str(tmp_path / "ml_tools.parquet")) ak_res = tw(ak_jets) dak_res = tw(dak_jets) assert np.all(np.isclose(ak_res, dak_res.compute())) - expected_columns = { - "eta", - "phi", - "pfcands.pt", - "pfcands.phi", - "pfcands.eta", - "pfcands.feat1", - "pfcands.feat2", - } columns = set(list(dak.necessary_columns(dak_res).values())[0]) - assert columns == expected_columns + assert columns == EXPECTED_JET_COLUMNS # Length-0 testing - tw = torch_wrapper_test("tests/samples/pn_demo.pt", expected_output_shape=(None,)) - ak_jets, dak_jets = prepare_jets_array(njets=256, tmp_path=tmp_path) + tw = _make_torch_wrapper(expected_output_shape=(None,)) + ak_jets = prepare_jets_array(njets=256) + dak_jets = to_dask(ak_jets, str(tmp_path / "ml_tools_len0.parquet")) ak_jets = ak_jets[ak_jets.eta < -100] # Mimicking a low efficiency selection dak_jets = dak_jets[dak_jets.eta < -100] ak_res, dak_res = tw(ak_jets), tw(dak_jets) assert len(ak_jets) == 0 and len(dak_res.compute()) == 0 -@pytest.mark.dask_client -def test_tensorflow(tmp_path, dask_client): - _ = pytest.importorskip("tensorflow") +def _make_tf_wrapper(): + pytest.importorskip("tensorflow") from coffea.ml_tools.tf_wrapper import tf_wrapper @@ -211,67 +254,109 @@ def postprocess_awkward(self, ret, jets): return ret # The tensorflow model here is used to classify jet constitutes - tfw = tf_wrapper_test("tests/samples/tf_model.keras") - ak_jets, dak_jets = prepare_jets_array(njets=256, tmp_path=tmp_path) + return tf_wrapper_test("tests/samples/tf_model.keras") - ak_res = tfw(ak_jets) - dak_res = tfw(dak_jets) - assert np.all(np.isclose(ak_res, dak_res.compute())) - expected_columns = {"ncands"} | {f"pfcands.feat{i}" for i in range(1, 19)} - columns = set(list(dak.necessary_columns(dak_res).values())[0]) - assert columns == expected_columns +def _make_tf_length0_wrapper(): + pytest.importorskip("tensorflow") + + from coffea.ml_tools.tf_wrapper import tf_wrapper # Length 0 testing. we cannot use the unflatten module in this case - class tf_wrapper_lenght0_test(tf_wrapper): + class tf_wrapper_length0_test(tf_wrapper): def prepare_awkward(self, arr): return [arr], {} - tfw_length0_tester = tf_wrapper_lenght0_test( + return tf_wrapper_length0_test( "tests/samples/tf_model.keras", skip_length_zero=True ) + +def test_tensorflow_eager(tmp_path): + tfw = _make_tf_wrapper() + ak_jets = prepare_jets_array(njets=256) + assert len(tfw(ak_jets)) == len(ak_jets) + + tfw_length0_tester = _make_tf_length0_wrapper() + # Making an explicit shape arr = ak.from_numpy(np.random.random(size=(10, 64, 18))) - tf_length10_path = str(tmp_path / "tf_length10.parquet") - ak.to_parquet(arr, tf_length10_path) - darr = dak.from_parquet(tf_length10_path) + assert len(tfw_length0_tester(arr)) == 10 + # Reducing the length 0 + arr = ak.from_numpy(np.zeros(shape=(0, 64, 18))) + tfw_length0_tester(arr) + + +@pytest.mark.dask_client +def test_tensorflow(tmp_path, dask_client): + dak = _dask_awkward() + tfw = _make_tf_wrapper() + + ak_jets = prepare_jets_array(njets=256) + dak_jets = to_dask(ak_jets, str(tmp_path / "ml_tools.parquet")) + + ak_res = tfw(ak_jets) + dak_res = tfw(dak_jets) + + assert np.all(np.isclose(ak_res, dak_res.compute())) + expected_columns = {"ncands"} | {f"pfcands.feat{i}" for i in range(1, 19)} + columns = set(list(dak.necessary_columns(dak_res).values())[0]) + assert columns == expected_columns + + tfw_length0_tester = _make_tf_length0_wrapper() + + # Making an explicit shape + arr = ak.from_numpy(np.random.random(size=(10, 64, 18))) + darr = to_dask(arr, str(tmp_path / "tf_length10.parquet")) ak_res = tfw_length0_tester(arr) dak_res = tfw_length0_tester(darr) assert np.all(np.isclose(ak_res, dak_res.compute())) # Reducing the length 0 arr = ak.from_numpy(np.zeros(shape=(0, 64, 18))) - tf_length0_path = str(tmp_path / "tf_length0.parquet") - ak.to_parquet(arr, tf_length0_path) - darr = dak.from_parquet(tf_length0_path) + darr = to_dask(arr, str(tmp_path / "tf_length0.parquet")) ak_res = tfw_length0_tester(arr) dak_res = tfw_length0_tester(darr) -@pytest.mark.dask_client -def test_xgboost(tmp_path, dask_client): - _ = pytest.importorskip("xgboost") +def _make_xgboost_wrapper(): + pytest.importorskip("xgboost") from coffea.ml_tools.xgboost_wrapper import xgboost_wrapper - feature_list = [f"feat{i}" for i in range(16)] - class xgboost_test(xgboost_wrapper): def prepare_awkward(self, events): ret = ak.concatenate( - [events[name][:, np.newaxis] for name in feature_list], axis=1 + [events[name][:, np.newaxis] for name in XGBOOST_FEATURES], axis=1 ) return [], dict(data=ret) - xgb_wrap = xgboost_test("tests/samples/xgboost_example.ubj") + return xgboost_test("tests/samples/xgboost_example.ubj") + - # Dummy 1000 event array with 20 feature branches - ak_events = ak.zip( - {f"feat{i}": ak.from_numpy(np.random.random(size=1_000)) for i in range(20)} +def _xgboost_events(nevents=1_000): + # Dummy event array with 20 feature branches + return ak.zip( + {f"feat{i}": ak.from_numpy(np.random.random(size=nevents)) for i in range(20)} ) - xgboost_path = str(tmp_path / "ml_tools.xgboost.parquet") - ak.to_parquet(ak_events, xgboost_path) - dak_events = dak.from_parquet(xgboost_path) + + +def test_xgboost_eager(tmp_path): + xgb_wrap = _make_xgboost_wrapper() + ak_events = _xgboost_events() + + assert len(xgb_wrap(ak_events)) == len(ak_events) + + # Length 0 testing, xgboost always handles 0-length arrays elegantly + assert len(xgb_wrap(ak_events[ak_events.feat0 < 0])) == 0 + + +@pytest.mark.dask_client +def test_xgboost(tmp_path, dask_client): + dak = _dask_awkward() + xgb_wrap = _make_xgboost_wrapper() + + ak_events = _xgboost_events() + dak_events = to_dask(ak_events, str(tmp_path / "ml_tools.xgboost.parquet")) ak_res = xgb_wrap(ak_events) dak_res = xgb_wrap(dak_events) @@ -281,9 +366,8 @@ def prepare_awkward(self, events): # Should only load required columns columns = set(list(dak.necessary_columns(dak_res).values())[0]) - assert columns == set(feature_list) + assert columns == set(XGBOOST_FEATURES) # Length 0 testing, xgboost always handles 0-length arrays elegantly ak_res = xgb_wrap(ak_events[ak_events.feat0 < 0]) dak_res = xgb_wrap(dak_events[dak_events.feat0 < 0]) - assert len(ak_res) == 0 and len(dak_res.compute()) == 0 From 3cedfdcdf1a5c746383be5160c7c5b0fdf38910a Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Mon, 27 Jul 2026 11:31:13 -0500 Subject: [PATCH 6/7] ci: apply the retry and skip flag to both triton server steps Both the with- and without-dask-client jobs run triton tests now, so both need the same tolerance for a registry outage rather than only the one that had it. Assisted-by: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index abe574126..aac5dc875 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -89,7 +89,23 @@ jobs: - name: Start triton server with example model if: startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != '3.14t' run: | - docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models + # Pulls from nvcr.io intermittently time out. Back off 15s, 1m, 5m, 20m + # between attempts: long enough to ride out a transient registry outage, + # bounded so a sustained one costs ~26 minutes. On exhaustion report the + # server as unavailable so the triton tests skip rather than failing the + # job. A server that starts but misbehaves still fails the tests. + backoff=(15 60 300 1200) + for attempt in 0 1 2 3 4; do + if docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models; then + exit 0 + fi + if [ "${attempt}" -lt 4 ]; then + echo "::warning::triton server start failed (attempt $((attempt + 1))/5); retrying in ${backoff[$attempt]}s" + sleep "${backoff[$attempt]}" + fi + done + echo "::warning::triton server unavailable after 5 attempts; triton tests will skip" + echo "TRITON_UNAVAILABLE=1" >> "$GITHUB_ENV" - name: Test with pytest (without dask Client) run: | python -m pytest --cov-report=xml --cov=coffea --ignore=tests/test_taskvine_dask.py --ignore=tests/test_taskvine_virtual.py -m "not dask_client" -n 4 @@ -166,12 +182,11 @@ jobs: - name: Start triton server with example model if: startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != '3.14t' run: | - # Pulls from nvcr.io intermittently time out. Retry, and on exhaustion - # report the server as unavailable so test_triton skips rather than - # failing the job on a registry outage. A server that starts but - # misbehaves still fails the test. - # Back off 15s, 1m, 5m, 20m between attempts: long enough to ride out a - # transient registry outage, bounded so a sustained one costs ~26 min. + # Pulls from nvcr.io intermittently time out. Back off 15s, 1m, 5m, 20m + # between attempts: long enough to ride out a transient registry outage, + # bounded so a sustained one costs ~26 minutes. On exhaustion report the + # server as unavailable so the triton tests skip rather than failing the + # job. A server that starts but misbehaves still fails the tests. backoff=(15 60 300 1200) for attempt in 0 1 2 3 4; do if docker run -d --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${{ github.workspace }}/tests/samples/triton_models_test:/models nvcr.io/nvidia/tritonserver:25.11-pyt-python-py3 tritonserver --model-repository=/models; then From ed33cb9659ad57bf5a5dccb376aee90eef5a7ead Mon Sep 17 00:00:00 2001 From: Nick Manganelli Date: Mon, 27 Jul 2026 15:33:10 -0500 Subject: [PATCH 7/7] test(ml_tools): name the wrapper tests by the array type they exercise _ak for the plain awkward path, _dak for the dask_awkward one, rather than leaving the dask variant unsuffixed. The pair is now symmetric and the array type under test is visible in the name. The _ak tests stay awkward-only. Coffea's virtual mode is a property of how a factory materializes buffers from a file, not something these directly built arrays can carry, and numpy_call_wrapper dispatches on dask_awkward alone, so a virtual array would take the same path as an eager one. Assisted-by: Claude Opus 4.7 (1M context) --- tests/test_ml_tools.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_ml_tools.py b/tests/test_ml_tools.py index 88100314c..7f3e9b9bf 100644 --- a/tests/test_ml_tools.py +++ b/tests/test_ml_tools.py @@ -126,7 +126,7 @@ def prepare_awkward(self, output_list, jets): ) -def test_triton_eager(tmp_path): +def test_triton_ak(tmp_path): tw = _make_triton_wrapper() ak_jets = prepare_jets_array(njets=256) @@ -141,7 +141,7 @@ def test_triton_eager(tmp_path): @pytest.mark.dask_client -def test_triton(tmp_path, dask_client): +def test_triton_dak(tmp_path, dask_client): dak = _dask_awkward() tw = _make_triton_wrapper() @@ -181,7 +181,7 @@ def prepare_awkward(self, jets): return torch_wrapper_test("tests/samples/pn_demo.pt", **kwargs) -def test_torch_eager(tmp_path): +def test_torch_ak(tmp_path): tw = _make_torch_wrapper() ak_jets = prepare_jets_array(njets=256) assert len(tw(ak_jets)) == len(ak_jets) @@ -195,7 +195,7 @@ def test_torch_eager(tmp_path): @pytest.mark.dask_client -def test_torch(tmp_path, dask_client): +def test_torch_dak(tmp_path, dask_client): dak = _dask_awkward() tw = _make_torch_wrapper() @@ -272,7 +272,7 @@ def prepare_awkward(self, arr): ) -def test_tensorflow_eager(tmp_path): +def test_tensorflow_ak(tmp_path): tfw = _make_tf_wrapper() ak_jets = prepare_jets_array(njets=256) assert len(tfw(ak_jets)) == len(ak_jets) @@ -288,7 +288,7 @@ def test_tensorflow_eager(tmp_path): @pytest.mark.dask_client -def test_tensorflow(tmp_path, dask_client): +def test_tensorflow_dak(tmp_path, dask_client): dak = _dask_awkward() tfw = _make_tf_wrapper() @@ -340,7 +340,7 @@ def _xgboost_events(nevents=1_000): ) -def test_xgboost_eager(tmp_path): +def test_xgboost_ak(tmp_path): xgb_wrap = _make_xgboost_wrapper() ak_events = _xgboost_events() @@ -351,7 +351,7 @@ def test_xgboost_eager(tmp_path): @pytest.mark.dask_client -def test_xgboost(tmp_path, dask_client): +def test_xgboost_dak(tmp_path, dask_client): dak = _dask_awkward() xgb_wrap = _make_xgboost_wrapper()