Upgrade pytorch (meteofrance#86)

* upgrade pytorch and python to fit the docker image

* upgrade requirements.txt

* CI automatically gets the python version

* upgrade lint requirements

* upgrade python version in pyproject.toml

* lint

* add exception to bandit

Author: A669015

.github/workflows/tests.yml

@@ -18,10 +18,15 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - name: Set up Python 3.10
+
+    - name: Get python version from env.yaml
+      id: python_version
+      run: echo PYTHON_VERS=$(cat env.yaml | grep python= | sed 's/.*python=\([0-9]\.[0-9]*\.[0-9]*\).*/\1/') >> $GITHUB_OUTPUT
+
+    - name: Set up Python ${{ steps.python_version.outputs.PYTHON_VERS }}
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "${{ steps.python_version.outputs.PYTHON_VERS }}"
 
     - name: Get pytorch version from env.yaml
       id: pytorch_version

Dockerfile

@@ -2,7 +2,7 @@ ARG DOCKER_REGISTRY=docker.io
 ARG TORCH_VERS=2.2.2
 ARG CUDA_VERS=12.1
 
-FROM ${DOCKER_REGISTRY}/pytorch/pytorch:${TORCH_VERS}-cuda${CUDA_VERS}-cudnn8-devel
+FROM ${DOCKER_REGISTRY}/pytorch/pytorch:${TORCH_VERS}-cuda${CUDA_VERS}-cudnn9-devel
 
 ARG INJECT_MF_CERT
 

env.yaml

@@ -3,9 +3,9 @@ channels:
   - pytorch
   - nvidia
 dependencies:
-  - python=3.10.13          # The reference python version for the project
+  - python=3.11.9           # The reference python version for the project
   - pytorch-cuda=12.1       # The reference cuda version for the project
-  - pytorch==2.2.2          # The reference pytorch version for the project
+  - pytorch==2.4.1          # The reference pytorch version for the project
   - pip
   - pip:
      - -r requirements.txt

lint.sh

@@ -7,4 +7,4 @@ pwd
 flake8 --ignore E203,W503 --max-line-length 120 $1
 isort --profile black -c $1
 black --check $1
-bandit -ll -r --skip B402,B321 $1
+bandit -ll -r --skip B402,B321,B614 $1

py4cast/datasets/base.py

@@ -2,6 +2,7 @@
 Base classes defining our software components
 and their interfaces
 """
+
 import warnings
 from abc import ABC, abstractclassmethod, abstractmethod, abstractproperty
 from copy import deepcopy
@@ -298,9 +299,11 @@ def index_select_dim(
             return NamedTensor(
                 self.index_select_dim(dim_name, indices, bare_tensor=True),
                 self.names,
-                self.feature_names
-                if dim_name != self.feature_dim_name
-                else [self.feature_names[i] for i in indices],
+                (
+                    self.feature_names
+                    if dim_name != self.feature_dim_name
+                    else [self.feature_names[i] for i in indices]
+                ),
                 feature_dim_name=self.feature_dim_name,
             )
 
@@ -614,7 +617,9 @@ class DatasetInfo:
     units: Dict[str, str]  # d[shortname] = unit (str)
     weather_dim: int
     forcing_dim: int
-    step_duration: float  # Duration (in hour) of one step in the dataset. 0.25 means 15 minutes.
+    step_duration: (
+        float  # Duration (in hour) of one step in the dataset. 0.25 means 15 minutes.
+    )
     statics: Statics  # A lot of static variable
     stats: Stats
     diff_stats: Stats

py4cast/datasets/dummy.py

@@ -3,6 +3,7 @@
 Can be used as a starting point to implement your own dataset.
 inputs, outputs and forcing are filled with random tensors.
 """
+
 from dataclasses import dataclass
 from functools import cached_property
 from io import BytesIO

py4cast/datasets/poesy.py

@@ -241,7 +241,9 @@ class PoesySettings:
     term: dict
     num_input_steps: int  # = 2  # Number of input timesteps
     num_output_steps: int  # = 1  # Number of output timesteps (= 0 for inference)
-    num_inference_pred_steps: int = 0  # 0 in training config ; else used to provide future information about forcings
+    num_inference_pred_steps: int = (
+        0  # 0 in training config ; else used to provide future information about forcings
+    )
     standardize: bool = False
     members: Tuple[int] = (0,)
 

py4cast/datasets/smeagol.py

@@ -251,7 +251,9 @@ class SmeagolSettings:
     term: dict  # Terms used in this configuration. Should be present in nc files.
     num_input_steps: int  # = 2  # Number of input timesteps
     num_output_steps: int  # = 1  # Number of output timesteps (= 0 for inference)
-    num_inference_pred_steps: int = 0  # 0 in training config ; else used to provide future information about forcings
+    num_inference_pred_steps: int = (
+        0  # 0 in training config ; else used to provide future information about forcings
+    )
     standardize: bool = True  # Do we need to standardize our data ?
     members: Tuple[int] = (
         0,

py4cast/ideas/minimal_leak.py

@@ -4,6 +4,7 @@
 of Torch tensor on CPU.
 Using numpy seems to work fine.
 """
+
 import gc
 import json
 import os

py4cast/ideas/recursive_dict_register.py

@@ -4,7 +4,6 @@
 
 
 class RegisterDictMixin:
-
     """
     Register dictionnaries.
     Enable to recursively register dictionnary (or other object with a getitem method).

py4cast/lightning.py

@@ -477,9 +477,11 @@ def _common_step(
 
             if scale_y:
                 step_diff_std, step_diff_mean = self._step_diffs(
-                    self.output_feature_names
-                    if inference
-                    else batch.outputs.feature_names,
+                    (
+                        self.output_feature_names
+                        if inference
+                        else batch.outputs.feature_names
+                    ),
                     prev_states.device,
                 )
 

py4cast/models/base.py

@@ -2,6 +2,7 @@
 Abstract Base Class for all models
 Contains also a few functionnality used in various model.
 """
+
 from abc import ABC, abstractproperty
 from typing import Tuple
 

py4cast/models/nlam/create_mesh.py

@@ -107,7 +107,7 @@ def mk_2d_graph(xy, nx, ny):
 
     # turn into directed graph
     dg = networkx.DiGraph(g)
-    for (u, v) in g.edges():
+    for u, v in g.edges():
         d = np.sqrt(np.sum((g.nodes[u]["pos"] - g.nodes[v]["pos"]) ** 2))
         dg.edges[u, v]["len"] = d
         dg.edges[u, v]["vdiff"] = g.nodes[u]["pos"] - g.nodes[v]["pos"]

py4cast/models/vision/conv.py

@@ -2,6 +2,7 @@
 Convolutional neural network models
 for pn-ia.
 """
+
 from collections import OrderedDict
 from dataclasses import dataclass
 from functools import reduce

py4cast/models/vision/transformers.py

@@ -4,7 +4,6 @@
 and adapted to our needs (upsampler + extra settings)
 """
 
-
 from dataclasses import dataclass
 from functools import partial
 from math import sqrt
@@ -195,15 +194,15 @@ def forward(self, x, return_layer_outputs=False):
 
         layer_outputs = []
         i = 0
-        for (get_overlap_patches, overlap_embed, layers) in self.stages:
+        for get_overlap_patches, overlap_embed, layers in self.stages:
             x = get_overlap_patches(x)
 
             num_patches = x.shape[-1]
             ratio = int(sqrt((h * w) / num_patches))
             x = rearrange(x, "b c (h w) -> b c h w", h=h // ratio)
 
             x = overlap_embed(x)
-            for (attn, ff) in layers:
+            for attn, ff in layers:
                 x = attn(x) + x
                 x = ff(x) + x
 

py4cast/models/vision/unetrpp.py

@@ -676,9 +676,7 @@ def __init__(
         self.hidden_size = settings.hidden_size
         self.spatial_dims = settings.spatial_dims
         # Number of pixels after stem layer
-        no_pixels = (input_shape[0] * input_shape[1]) // (
-            settings.downsampling_rate**2
-        )
+        no_pixels = (input_shape[0] * input_shape[1]) // (settings.downsampling_rate**2)
         encoder_input_size = [
             no_pixels,
             no_pixels // 4,

py4cast_plugin_example.py

@@ -2,7 +2,6 @@
 A simple plugin example for py4cast model with a Identity model.
 """
 
-
 from dataclasses import dataclass
 
 import torch

pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 ]
 description = "Library to train a variety of Neural Network architectures on various weather forecasting datasets."
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: Apache 2",

requirements.txt

@@ -1,40 +1,30 @@
-numpy>=1.24.2
-wandb>=0.13.10
-matplotlib>=3.7.0
-scipy>=1.10.0
-pytorch-lightning>=2.1.2
-lightning==2.2.2
-shapely>=2.0.1
-networkx>=3.0
-Cartopy>=0.22.0
-pyproj>=3.4.1
-tueplots>=0.0.8
-plotly>=5.15.0
-tornado>=6.3.3
-cython>=3
-cfgrib==0.9.14.0
-dataclasses-json==0.6.4
-xarray==2024.7.0
+numpy==1.26.4
+matplotlib==3.9.2
+scipy==1.14.1
+lightning==2.4.0
+networkx==3.3
+Cartopy==0.24.1
+tueplots==0.0.17
+cfgrib==0.9.14.1
+dataclasses-json==0.6.7
+xarray==2024.9.0
 argparse-dataclass==2.0.0
-tensorboard==2.16.1
-typer==0.9.0
-netCDF4==1.6.5
+tensorboard==2.17.1
+typer==0.12.5
+netCDF4==1.7.1.post2
 tensorboard-plugin-profile==2.17.0
-torch-tb-profiler==0.4.1
-einops==0.7.0
+torch-tb-profiler==0.4.3
+einops==0.8.0
 torchinfo==1.8.0
 tabulate==0.9.0
-pytest==8.1.1
-coverage==7.6.1
-onnx==1.16.1
-onnxruntime==1.18.1
-onnxruntime-gpu==1.18.1
-onnxscript==0.1.0.dev20240905
-monai==1.3.1
+pytest==8.3.3
+coverage==7.6.3
+onnx==1.17.0
+onnxruntime-gpu==1.19.2
+onnxscript==0.1.0.dev20241018
+monai==1.4.0
 gif==23.3.0
 scikit-image==0.24.0
-typer==0.9.0
-transformers==4.44.2
-tensorflow==2.16.1
-py-spy
-torch-geometric==2.3.1
+typer==0.12.5
+transformers==4.45.2
+torch-geometric==2.6.1

requirements_lint.txt

@@ -1,5 +1,4 @@
-pylint
-bandit==1.7.4
-flake8==4.0.1
-black==22.3.0
-isort==5.11.2
+bandit==1.7.10
+flake8==7.1.1
+black==24.10.0
+isort==5.13.2

tests/test_datasets.py

@@ -1,6 +1,7 @@
 """
 Unit tests for datasets and NamedTensor.
 """
+
 import datetime
 
 import numpy as np

tests/test_io.py

@@ -1,6 +1,7 @@
 """
 Unit tests for datasets and NamedTensor.
 """
+
 from dataclasses import dataclass
 from functools import cached_property
 

tests/test_models.py

@@ -5,6 +5,7 @@
 3. onnx exported
 4. onnx loaded and used for inference
 """
+
 import tempfile
 from dataclasses import dataclass
 from pathlib import Path