Drop Python 3.10 and restore JAX 0.9 XLA coverage

Author: Trinkle23897

README.md

@@ -43,7 +43,7 @@ Check out our [arXiv paper](https://arxiv.org/abs/2206.10558) for more details!
 
 ### PyPI
 
-EnvPool is currently hosted on [PyPI](https://pypi.org/project/envpool/). It supports Python 3.10-3.13.
+EnvPool is currently hosted on [PyPI](https://pypi.org/project/envpool/). It supports Python 3.11-3.13.
 
 You can simply install EnvPool with the following command:
 

docs/content/build.rst

@@ -79,7 +79,7 @@ EnvPool requires **GCC/G++ version >= 9.0** to build the source code. To install
     # to change the default cc to gcc-9:
     sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 60 --slave /usr/bin/g++ g++ /usr/bin/g++-9
 
-It also requires **Python version >= 3.10**:
+It also requires **Python version >= 3.11**:
 
 .. code-block:: bash
 

docs/content/new_env.rst

@@ -716,7 +716,7 @@ If we want to add ``tianshou`` as a build dependency, in ``setup.cfg``:
 
     [options]
     packages = find:
-    python_requires = >=3.10
+    python_requires = >=3.11
     install_requires =
         dm-env>=1.4
         gym>=0.18

docs/index.rst

@@ -47,7 +47,7 @@ Installation
 ------------
 
 EnvPool is currently hosted on `PyPI <https://pypi.org/project/envpool/>`_.
-It supports Python 3.10-3.13.
+It supports Python 3.11-3.13.
 
 You can install EnvPool with the following command:
 

envpool/python/xla_template.py

@@ -14,11 +14,10 @@
 """xla template on python side."""
 
 from collections import namedtuple
-from typing import Any, Callable, List, Tuple, Union, cast
+from typing import Any, Callable, List, Tuple, cast
 
 import numpy as np
 from jax import ShapeDtypeStruct, dtypes, ffi
-from jax import numpy as jnp
 
 
 def _normalize_specs(
@@ -71,7 +70,9 @@ def _make_xla_function(
       tuple(_layout(shape) for shape, _ in out_specs)
       if len(out_specs) > 1 else _layout(out_specs[0][0])
     ),
-    custom_call_api_version=0,
+    # JAX target registration uses api_version=0 for the legacy untyped
+    # handler, but StableHLO custom_call uses API_VERSION_ORIGINAL == 1.
+    custom_call_api_version=1,
     legacy_backend_config=cast(Any, handle),
   )
 

envpool/vizdoom/vizdoom_pretrain_test.py

@@ -186,7 +186,7 @@ def test_d3(self) -> None:
       try:
         # Full bazel runs may overlap this spawned evaluation with other
         # long-running pretrain tests on shared devbox CPUs.
-        reward, length = result_queue.get(timeout=240)
+        reward, length = result_queue.get(timeout=360)
       except queue.Empty:
         proc.terminate()
         proc.join(timeout=5)

gh/wf/release.yml

@@ -0,0 +1,53 @@
+name: Release PyPI Wheel
+
+# on: [push, pull_request]
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - v*
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    container: trinkle23897/envpool-release:2023-01-02-5f1a5fd
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        run: |
+          ln -sf /root/.cache $HOME/.cache
+          ln -sf /root/.pyenv $HOME/.pyenv
+          pyenv global ${{ matrix.python-version }}-dev
+      - name: Build
+        run: |
+          make pypi-wheel
+          pip3 install dist/*.whl --force-reinstall
+      - name: Test
+        run: |
+          make release-test
+      - name: Upload artifact
+        uses: actions/upload-artifact@main
+        with:
+          name: wheel
+          path: wheelhouse/
+
+  publish:
+    runs-on: ubuntu-latest
+    needs: [release]
+    steps:
+      - uses: actions/download-artifact@v3
+        with:
+          path: artifact
+      - name: Move files so the next action can find them
+        run: |
+          mkdir dist && mv artifact/wheel/* dist/
+          ls dist/
+      - name: Publish distribution to PyPI
+        if: startsWith(github.ref, 'refs/tags')
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

setup.cfg

@@ -11,7 +11,6 @@ project_urls =
     Bug Tracker = https://github.com/sail-sg/envpool/issues
 classifiers =
     Programming Language :: Python :: 3
-    Programming Language :: Python :: 3.10
     Programming Language :: Python :: 3.11
     Programming Language :: Python :: 3.12
     Programming Language :: Python :: 3.13
@@ -21,7 +20,7 @@ classifiers =
 
 [options]
 packages = find:
-python_requires = >=3.10
+python_requires = >=3.11
 install_requires =
     dm-env>=1.4
     gym>=0.18

third_party/pip_requirements/requirements-dev-lock.txt

@@ -24,10 +24,8 @@ gymnasium==1.2.3
 h5py==3.16.0
 idna==3.11
 imageio==2.37.3
-jax==0.6.2; python_version < "3.11"
-jax==0.9.2; python_version >= "3.11"
-jaxlib==0.6.2; python_version < "3.11"
-jaxlib==0.9.2; python_version >= "3.11"
+jax==0.9.2
+jaxlib==0.9.2
 jinja2==3.1.6
 labmaze==1.0.6
 llvmlite==0.46.0

third_party/pip_requirements/requirements-dev.txt

@@ -5,8 +5,7 @@ dm-env>=1.6
 gym>=0.26.2
 gymnasium>=1.2.3
 optree>=0.19.0
-jax[cpu]>=0.6.2,<0.7; python_version<'3.11'
-jax[cpu]>=0.9.2,<0.10; python_version>='3.11'
+jax[cpu]>=0.9.2,<0.10
 absl-py
 packaging>=26.0
 tqdm