CPU option in Docker file

Dockerfile

@@ -1,9 +1,7 @@
 FROM ubuntu:latest
 
-
 ENV CONDA_ENV_NAME=graph
-ENV PYTHON_VERSION=3.10
-
+ENV PYTHON_VERSION=3.12
 
 # Basic setup
 RUN apt update && apt install -y bash \
@@ -24,7 +22,15 @@ ENV PATH="/conda/bin:${PATH}"
 
 RUN git clone https://github.com/openclimatefix/graph_weather.git && mv graph_weather/ gw/ && cd gw/ && mv * .. && rm -rf gw/
 
-RUN conda update -n base -c defaults conda && conda env create -f environment.yml
+# Copy the appropriate environment file based on CUDA availability
+COPY environment_cpu.yml /tmp/environment_cpu.yml
+COPY environment_cuda.yml /tmp/environment_cuda.yml
+
+RUN conda update -n base -c defaults conda
+
+# Check if CUDA is available and accordingly choose env
+RUN cuda=$(command -v nvcc > /dev/null && echo "true" || echo "false") \
+    && if [ "$cuda" == "true" ]; then conda env create -f /tmp/environment_cuda.yml; else conda env create -f /tmp/environment_cpu.yml; fi
 
 # Switch to bash shell
 SHELL ["/bin/bash", "-c"]
@@ -35,8 +41,9 @@ RUN echo "source activate ${CONDA_ENV_NAME}" >> ~/.bashrc
 # Cp in the development directory and install
 RUN source activate ${CONDA_ENV_NAME} && pip install -e .
 
+
 # Make RUN commands use the new environment:
 SHELL ["conda", "run", "-n", "graph", "/bin/bash", "-c"]
 
-# Example commnad that can be used, need to set API_KEY, API_SECRET and SAVE_DIR
+# Example command that can be used, need to set API_KEY, API_SECRET and SAVE_DIR
 CMD ["conda", "run", "-n", "graph", "python", "-u", "train/pl_graph_weather.py", "--gpus", "16", "--hidden", "64", "--num-blocks", "3", "--batch", "16"]

environment_cpu.yml

@@ -0,0 +1,35 @@
+name: graph
+channels:
+  - pytorch
+  - pyg
+  - conda-forge
+  - defaults
+dependencies:
+  - numcodecs
+  - pandas
+  - pip
+  - pyg
+  - python=3.12
+  - pytorch
+  - cpuonly
+  - pytorch-cluster
+  - pytorch-scatter
+  - pytorch-sparse
+  - pytorch-spline-conv
+  - scikit-learn
+  - scipy
+  - torchvision
+  - tqdm
+  - xarray
+  - zarr
+  - h3-py
+  - numpy
+  - pip:
+      - datasets
+      - einops
+      - fsspec
+      - torch-geometric-temporal
+      - huggingface-hub
+      - pysolar
+      - pytorch-lightning
+      - click

environment.yml → environment_cuda.yml

@@ -11,7 +11,7 @@ dependencies:
   - pandas
   - pip
   - pyg
-  - python=3.11
+  - python=3.12
   - pytorch
   - pytorch-cluster
   - pytorch-scatter
@@ -24,6 +24,8 @@ dependencies:
   - xarray
   - zarr
   - h3-py
+  - numpy
+  - torch
   - pip:
       - datasets
       - einops
@@ -32,3 +34,4 @@ dependencies:
       - huggingface-hub
       - pysolar
       - pytorch-lightning
+      - click