Installation
=============

Here we will outline how to install the ``zea`` package and its dependencies.
Besides installation through `PyPI <https://pypi.org/project/zea>`_, we also provide Docker images for easy development (see the `Docker`_ section below). For instructions on contributing and development installation, see :doc:`contributing`.

.. tab-set::

    .. tab-item:: uv

         Add ``zea`` to your project with `uv <https://docs.astral.sh/uv/>`_, which
         manages the environment and lockfile for you (run ``uv init`` first if you
         don't have a project yet):

         .. code-block:: shell

               uv add zea

         Or install it into a plain environment, like ``pip``:

         .. code-block:: shell

               uv venv
               uv pip install zea

    .. tab-item:: pip

         .. code-block:: shell

               pip install zea

    .. tab-item:: conda

         Use `conda <https://docs.conda.io/>`_ to create and manage the environment.

         .. code-block:: shell

               conda create -n zea python=3.12  # 3.10 - 3.13
               conda activate zea
               pip install zea

    .. tab-item:: Docker

         Pre-built images ship with the backends already installed (see :ref:`Docker <docker-information>`).

         .. code-block:: shell

               docker pull zeahub/all:latest
               docker run --gpus 'all' -it zeahub/all:latest

.. note::

   When installing from PyPI you still need a machine learning backend (JAX, PyTorch,
   or TensorFlow). See :ref:`Backend <backend-installation>` below. The Docker images
   already include the backends.


.. _backend-installation:

Backend
-------

.. include:: getting-started.rst
   :start-after: .. backend-installation-start
   :end-before: .. backend-installation-end

.. dropdown:: Notes for Windows users
   :name: windows-users
   :class-container: dropdown-accordion-windows-users

   ``zea`` installs and runs on Windows the same way as above, with a couple of caveats:

   - **GPU support is limited.** JAX and TensorFlow only support GPU acceleration on
     Windows through WSL2. Native Windows installs are CPU-only for those backends.
     PyTorch is the exception and does support CUDA natively on Windows. See the backend
     links above for each framework's current platform support. ``zea`` itself runs fine
     on CPU, but beamforming and model inference are much slower without a GPU, so one is
     recommended for anything beyond small experiments.
   - **Streaming reads** (``hf://`` sources, see :ref:`working-with-zea-data-files`) fall
     back to a lock-guarded read instead of ``zea``'s fully parallel path, since Windows
     lacks ``os.pread``. This is still faster than downloading whole files, just not as
     fast as on Linux/macOS.

   If you need GPU acceleration and are on Windows, running ``zea`` inside WSL2 (or the
   :ref:`Docker images <docker-information>` via WSL2's GPU passthrough) is the most
   reliable path.


.. _docker-information:

Docker
-------

This repository provides multiple `Docker images <https://hub.docker.com/r/zeahub>`_ built and tested in our CI pipeline.

Pre-built images
~~~~~~~~~~~~~~~~


The following images are available on Docker Hub:

- `zeahub/all <https://hub.docker.com/r/zeahub/all>`_: This image includes support for all machine learning backends (TensorFlow, PyTorch, and JAX).
- `zeahub/tensorflow <https://hub.docker.com/r/zeahub/tensorflow>`_: This image includes support for TensorFlow.
- `zeahub/torch <https://hub.docker.com/r/zeahub/torch>`_: This image includes support for PyTorch.
- `zeahub/jax <https://hub.docker.com/r/zeahub/jax>`_: This image includes support for JAX.

These images are uploaded to Docker Hub via the CI pipeline and can be used directly in your projects via:

.. code-block:: shell

   docker pull zeahub/all:latest


Build
~~~~~

One can build an image for a specific backend using the provided `Dockerfile` and build arguments.

.. code-block:: shell

   docker build --build-arg INSTALL_JAX=gpu . -t zeahub/jax:latest

To build the full image including all backends with GPU support:

.. code-block:: shell

   docker build --build-arg INSTALL_JAX=gpu --build-arg INSTALL_TORCH=gpu --build-arg INSTALL_TF=gpu . -t zeahub/all:latest

Note that these build arguments can be set to either `cpu` or `gpu` depending on your needs.

Run
~~~

Run a container with one of the built images. Ensure you mount your repository at ``/zea`` so that changes are reflected inside the container.
We recommend using `rootless docker <https://docs.docker.com/engine/security/rootless/>`_ to avoid permission issues.

.. code-block:: shell

   docker run --name {CONTAINER-NAME} --gpus 'all' \
     -v ~/zea:/zea \
     -d -it {IMAGE-NAME}:{IMAGE-TAG}

.. dropdown:: Docker run command flags explained

   - ``docker run``: create and run a new container from an image.
   - ``--name``: name the container.
   - ``--gpus``: specify GPU devices to add to the container ('all' to pass all GPUs).
   - ``-v`` or ``--volume``: bind mount a volume.
   - ``-d`` or ``--detach``: start the container as a background process.
   - ``-it``: start an interactive terminal session.
      - ``--interactive``: keep STDIN open.
      - ``--tty``: allocate a pseudo-TTY.
   - ``-m`` or ``--memory``: set a memory limit (use g for gigabytes).
   - ``--cpus``: specify the number of CPU cores to use.
   - ``-w`` or ``--workdir``: set the working directory inside the container.
   - ``--rm``: automatically remove the container when it *exits*.
   - ``--env-file``: load environment variables from a .env file.
   - ``--hostname``: set the container hostname (useful for ``users.yaml`` file).

.. important::

   Mount your ``zea`` repository to ``/zea`` inside the container so that changes are reflected in the ``zea`` installation inside the container.

.. tip::

   The Docker container sets a random hostname by default. You can set a hostname with the ``--hostname`` flag. This is useful for the ``users.yaml`` file. Alternatively, you can use the hostname wildcard in the ``users.yaml`` file.