Draft: Refactor JobSet for Pathways

Dockerfile

@@ -93,7 +93,7 @@ RUN apt-get install -y google-perftools
 ENV PIP_FIND_LINKS=https://storage.googleapis.com/jax-releases/libtpu_releases.html
 # Ensure we install the TPU version, even if building locally.
 # Jax will fallback to CPU when run on a machine without TPU.
-RUN pip install .[core,tpu]
+RUN pip install .[core,tpu,pathways]
 RUN if [ -n "$EXTRAS" ]; then pip install .[$EXTRAS]; fi
 COPY . .