[57] Pinning flash-attention and other dependencies (ecmwf#75)

* work

* work

* changes

* changes

* reset

* removing flash from install

Author: tjhunter

INSTALL.md

@@ -10,12 +10,6 @@
 %>cd WeatherGenerator
 %>uv sync
  
-# install flash_attn which is not uv compatible
-# (this step might take very long, it is faster when ninja is available but then
-# MAX_JOBS=4 should be set as environment variable; the step might need to be run
-# on a compute node to ensure the correct CUDA environment is available) 
-%>uv pip install torch
-%>uv pip install flash_attn --no-build-isolation
  
 %>uv run train
 ```

pyproject.toml

@@ -9,18 +9,21 @@ authors = [
 
 requires-python = ">=3.11,<3.13"
 # TODO: split the plotting dependencies into their own dep groups, they are not required.
-dependencies = [ 'torch',
-                 'numpy',
-                 'astropy_healpix',
-                 'zarr',
-                 'anemoi-datasets',
-                 'pandas',
-                 'pynvml',
-                 'tqdm',
-                 'matplotlib',
-                 'packaging',
-                 'wheel',
-                 'psutil']
+dependencies = [
+ 'torch==2.6.0',
+ 'numpy~=2.2',
+ 'astropy_healpix~=1.0',
+ 'zarr~=2.17',
+ 'anemoi-datasets~=0.5',
+ 'pandas~=2.2',
+ 'pynvml',
+ 'tqdm',
+ 'matplotlib',
+ 'packaging',
+ 'wheel',
+ 'psutil',
+ "flash-attn",
+]
 
 [project.urls]
 Homepage = "https://www.weathergenerator.eu"
@@ -91,4 +94,25 @@ ignore = [
   "F811",
   # To ignore, not relevant for us
   "E741",
-]
+]
+
+[tool.uv]
+# This guarantees that the build is deterministic and will not be impacted 
+# by future releases of dependencies or sub-dependencies. 
+# See https://docs.astral.sh/uv/reference/settings/#exclude-newer
+# TODO: pytorch does not publish valid release timestamps, so sadly it does not work.
+# exclude-newer = "2025-03-14T00:00:00Z"
+
+# Following the recommendations from https://docs.astral.sh/uv/guides/integration/pytorch
+[[tool.uv.index]]
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+# This URL was evaluated this way:
+# uv run ~/WeatherGenerator-private/hpc/hpc2020/ecmwf/get-flash-atten.sh
+flash-attn = { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl" }