From 08676049056fd069b3b2f2f6220d78a2508ef1ae Mon Sep 17 00:00:00 2001 From: Meesum Qazalbash Date: Tue, 27 Jan 2026 17:22:36 +0500 Subject: [PATCH 1/4] CI(gh-2126): transition to uv for dependency management and testing --- .github/workflows/ci.yml | 135 +++++++++++++++++---------------------- 1 file changed, 58 insertions(+), 77 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 291be9f28..e9c825f3f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,54 +13,47 @@ env: jobs: prek: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v6 - name: prek check uses: j178/prek-action@v1 with: extra-args: --all-files --skip ruff --skip ruff-format --skip ty --skip mypy - lint: - runs-on: ubuntu-latest strategy: matrix: python-version: ["3.11", "3.13"] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + - uses: actions/checkout@v6 + - name: Install uv + uses: astral-sh/setup-uv@v7 with: + enable-cache: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | sudo apt install -y pandoc gsfonts - python -m pip install --upgrade pip - pip install jaxlib - pip install jax - pip install '.[doc,test]' - pip install https://github.com/pyro-ppl/funsor/archive/master.zip - pip install -r docs/requirements.txt - pip freeze + uv pip install --upgrade jaxlib jax + uv pip install --upgrade '.[doc,test]' + uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --upgrade -r docs/requirements.txt + uv pip freeze - name: Lint with mypy and ruff run: | - make lint + uv run make lint - name: Build documentation run: | - make docs + uv run make docs - name: Test documentation run: | - make doctest - python -m doctest -v README.md - + uv run make doctest + uv run python -m doctest -v README.md test-modeling: - runs-on: ubuntu-latest needs: [lint, prek] strategy: @@ -68,39 +61,37 @@ jobs: python-version: ["3.11", "3.13"] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + - uses: actions/checkout@v6 + - name: Install uv + uses: astral-sh/setup-uv@v7 with: + enable-cache: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | sudo apt install -y graphviz - python -m pip install --upgrade pip # Keep track of pyro-api master branch - pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip - pip install jaxlib - pip install jax - pip install https://github.com/pyro-ppl/funsor/archive/master.zip - pip install -e '.[dev,test]' - pip freeze + uv pip install --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip + uv pip install --upgrade jaxlib jax + uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --upgrade -e '.[dev,test]' + uv pip freeze - name: Test with pytest run: | - CI=1 pytest -vs -k "not test_example" --durations=100 --ignore=test/infer/ --ignore=test/contrib/ + CI=1 uv run pytest -vs -k "not test_example" --durations=100 --ignore=test/infer/ --ignore=test/contrib/ - name: Test x64 run: | - JAX_ENABLE_X64=1 pytest -vs test/test_distributions.py -k "powerLaw or Dagum" + JAX_ENABLE_X64=1 uv run pytest -vs test/test_distributions.py -k "powerLaw or Dagum" - name: Test tracer leak if: matrix.python-version == '3.13' env: JAX_CHECK_TRACER_LEAKS: 1 run: | - pytest -vs test/infer/test_mcmc.py::test_chain_inside_jit - pytest -vs test/infer/test_mcmc.py::test_chain_jit_args_smoke - pytest -vs test/infer/test_mcmc.py::test_reuse_mcmc_run - pytest -vs test/infer/test_mcmc.py::test_model_with_multiple_exec_paths - pytest -vs test/test_distributions.py::test_mean_var -k Gompertz - + uv run pytest -vs test/infer/test_mcmc.py::test_chain_inside_jit + uv run pytest -vs test/infer/test_mcmc.py::test_chain_jit_args_smoke + uv run pytest -vs test/infer/test_mcmc.py::test_reuse_mcmc_run + uv run pytest -vs test/infer/test_mcmc.py::test_model_with_multiple_exec_paths + uv run pytest -vs test/test_distributions.py::test_mean_var -k Gompertz - name: Coveralls if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13' uses: coverallsapp/github-action@v2 @@ -109,9 +100,7 @@ jobs: parallel: true flag-name: test-modeling - test-inference: - runs-on: ubuntu-latest needs: [lint, prek] strategy: @@ -119,41 +108,39 @@ jobs: python-version: ["3.11", "3.13"] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + - uses: actions/checkout@v6 + - name: Install uv + uses: astral-sh/setup-uv@v7 with: + enable-cache: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - python -m pip install --upgrade pip - # Keep track of pyro-api master branch - pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip - pip install jaxlib - pip install jax - pip install https://github.com/pyro-ppl/funsor/archive/master.zip - pip install -e '.[dev,test]' - pip freeze + uv pip install --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip + uv pip install --upgrade jaxlib jax + uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --upgrade -e '.[dev,test]' + uv pip freeze - name: Test with pytest run: | - pytest -vs --durations=20 test/infer/test_mcmc.py - pytest -vs --durations=20 test/infer --ignore=test/infer/test_mcmc.py --ignore=test/contrib/test_nested_sampling.py - pytest -vs --durations=20 test/contrib --ignore=test/contrib/stochastic_support/test_dcc.py + uv run pytest -vs --durations=20 test/infer/test_mcmc.py + uv run pytest -vs --durations=20 test/infer --ignore=test/infer/test_mcmc.py --ignore=test/contrib/test_nested_sampling.py + uv run pytest -vs --durations=20 test/contrib --ignore=test/contrib/stochastic_support/test_dcc.py - name: Test x64 run: | - JAX_ENABLE_X64=1 pytest -vs test/infer/test_mcmc.py -k x64 + JAX_ENABLE_X64=1 uv run pytest -vs test/infer/test_mcmc.py -k x64 - name: Test chains run: | - XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/infer/test_mcmc.py -k "chain or pmap or vmap" - XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/contrib/test_tfp.py -k "chain" - XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/contrib/stochastic_support/test_dcc.py - XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/infer/test_hmc_gibbs.py -k "chain" + XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_mcmc.py -k "chain or pmap or vmap" + XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/contrib/test_tfp.py -k "chain" + XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/contrib/stochastic_support/test_dcc.py + XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_hmc_gibbs.py -k "chain" - name: Test custom prng run: | - JAX_ENABLE_CUSTOM_PRNG=1 pytest -vs test/infer/test_mcmc.py + JAX_ENABLE_CUSTOM_PRNG=1 uv run pytest -vs test/infer/test_mcmc.py - name: Test nested sampling run: | - JAX_ENABLE_X64=1 pytest -vs test/contrib/test_nested_sampling.py + JAX_ENABLE_X64=1 uv run pytest -vs test/contrib/test_nested_sampling.py - name: Coveralls if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13' uses: coverallsapp/github-action@v2 @@ -162,9 +149,7 @@ jobs: parallel: true flag-name: test-inference - examples: - runs-on: ubuntu-latest needs: [lint, prek] strategy: @@ -172,22 +157,21 @@ jobs: python-version: ["3.13"] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + - uses: actions/checkout@v6 + - name: Install uv + uses: astral-sh/setup-uv@v7 with: + enable-cache: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - python -m pip install --upgrade pip - pip install jaxlib - pip install jax - pip install https://github.com/pyro-ppl/funsor/archive/master.zip - pip install -e '.[dev,examples,test]' - pip freeze + uv pip install --upgrade jaxlib jax + uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --upgrade -e '.[dev,examples,test]' + uv pip freeze - name: Test with pytest run: | - CI=1 XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs -k test_example + CI=1 XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs -k test_example - name: Coveralls if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13' uses: coverallsapp/github-action@v2 @@ -196,9 +180,7 @@ jobs: parallel: true flag-name: examples - finish: - needs: [test-modeling, test-inference, examples] runs-on: ubuntu-latest steps: @@ -208,4 +190,3 @@ jobs: github-token: ${{ secrets.GITHUB_TOKEN }} parallel-finished: true carryforward: "test-modeling,test-inference,examples" - From 5453f405f09c2292e950e35d62c0198c97966545 Mon Sep 17 00:00:00 2001 From: Meesum Qazalbash Date: Tue, 27 Jan 2026 17:38:44 +0500 Subject: [PATCH 2/4] fix: install packages in system --- .github/workflows/ci.yml | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e9c825f3f..e6dc21ffb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,10 +37,10 @@ jobs: - name: Install dependencies run: | sudo apt install -y pandoc gsfonts - uv pip install --upgrade jaxlib jax - uv pip install --upgrade '.[doc,test]' - uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip - uv pip install --upgrade -r docs/requirements.txt + uv pip install --system --upgrade jaxlib jax + uv pip install --system --upgrade '.[doc,test]' + uv pip install --system --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --system --upgrade -r docs/requirements.txt uv pip freeze - name: Lint with mypy and ruff run: | @@ -71,10 +71,10 @@ jobs: run: | sudo apt install -y graphviz # Keep track of pyro-api master branch - uv pip install --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip - uv pip install --upgrade jaxlib jax - uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip - uv pip install --upgrade -e '.[dev,test]' + uv pip install --system --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip + uv pip install --system --upgrade jaxlib jax + uv pip install --system --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --system --upgrade -e '.[dev,test]' uv pip freeze - name: Test with pytest run: | @@ -116,10 +116,10 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - uv pip install --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip - uv pip install --upgrade jaxlib jax - uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip - uv pip install --upgrade -e '.[dev,test]' + uv pip install --system --upgrade https://github.com/pyro-ppl/pyro-api/archive/master.zip + uv pip install --system --upgrade jaxlib jax + uv pip install --system --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --system --upgrade -e '.[dev,test]' uv pip freeze - name: Test with pytest run: | @@ -165,9 +165,9 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - uv pip install --upgrade jaxlib jax - uv pip install --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip - uv pip install --upgrade -e '.[dev,examples,test]' + uv pip install --system --upgrade jaxlib jax + uv pip install --system --upgrade https://github.com/pyro-ppl/funsor/archive/master.zip + uv pip install --system --upgrade -e '.[dev,examples,test]' uv pip freeze - name: Test with pytest run: | From 551d801448bcbee7fd2c6d2030c30b35d9880ae5 Mon Sep 17 00:00:00 2001 From: Meesum Qazalbash Date: Tue, 27 Jan 2026 17:45:39 +0500 Subject: [PATCH 3/4] fix: update path --- .github/workflows/ci.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e6dc21ffb..8ff31b38c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,6 +33,7 @@ jobs: uses: astral-sh/setup-uv@v7 with: enable-cache: true + update-path: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | @@ -65,6 +66,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v7 with: + update-path: true enable-cache: true python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -113,6 +115,7 @@ jobs: uses: astral-sh/setup-uv@v7 with: enable-cache: true + update-path: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | @@ -162,6 +165,7 @@ jobs: uses: astral-sh/setup-uv@v7 with: enable-cache: true + update-path: true python-version: ${{ matrix.python-version }} - name: Install dependencies run: | From 2fdb2b70f2458146887700f0c164164aae7b30b1 Mon Sep 17 00:00:00 2001 From: Meesum Qazalbash Date: Tue, 27 Jan 2026 17:48:13 +0500 Subject: [PATCH 4/4] fix: specify uv python version in enviornment variable --- .github/workflows/ci.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8ff31b38c..04588e191 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,6 +26,8 @@ jobs: strategy: matrix: python-version: ["3.11", "3.13"] + env: + UV_PYTHON: ${{ matrix.python-version }} steps: - uses: actions/checkout@v6 @@ -60,6 +62,8 @@ jobs: strategy: matrix: python-version: ["3.11", "3.13"] + env: + UV_PYTHON: ${{ matrix.python-version }} steps: - uses: actions/checkout@v6 @@ -108,6 +112,8 @@ jobs: strategy: matrix: python-version: ["3.11", "3.13"] + env: + UV_PYTHON: ${{ matrix.python-version }} steps: - uses: actions/checkout@v6 @@ -158,6 +164,8 @@ jobs: strategy: matrix: python-version: ["3.13"] + env: + UV_PYTHON: ${{ matrix.python-version }} steps: - uses: actions/checkout@v6