base_image: "anyscale/ray-ml:nightly-py37-gpu"

python:
  pip_packages:
    - boto3
    - pyarrow<7.0.0
    - tqdm
  conda_packages: []

post_build_cmds:
  - pip uninstall -y ray || true
  - pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
  # TODO (Alex): We need to do this because the ray-ml image pins
  # tensorflow=2.6, which requires numpy~=1.19.2. This is ok because the test
  # doesn't actually use tensorflow, but in the long term, but we should
  # consider upgrading to tensorflow 2.7 as a long term solution.
  - pip install -U numpy>=1.20
  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}