Repository: metno/bris-inference
Pull request #63: 19 graph switch
Author: login: einrone
Created: 2025-03-06T10:33:24Z | Merged: 2025-03-07T12:50:12Z

Added functionality to change the internal graph in the model. This is done by providing a key called "graph" inside the config where the value of "graph" points to a graph file.

`graph: /foo/bar/baz/GRAPH.pt`

This triggers: changing the internal graph ->  building a new model with respect to the new graph -> reassigns the model weights to the new model. 


Diff:
diff --git a/bris/__main__.py b/bris/__main__.py
index f571bee9..24cab3b3 100644
--- a/bris/__main__.py
+++ b/bris/__main__.py
@@ -32,9 +32,11 @@ def main():
     checkpoint.set_encoder_decoder_num_chunks(
         getattr(config, "inference_num_chunks", 1)
     )
-    if hasattr(config.model, "graph"):
+    if hasattr(config, "graph"):
         LOGGER.info("Update graph is enabled. Proceeding to change internal graph")
-        checkpoint.update_graph(config.model.graph)  # Pass in a new graph if needed
+        # At the moment config.graph is only a POSIX path.
+        # TODO: In future a graph can be generated "on the fly" by providing a config 
+        checkpoint.update_graph(config.graph)  # Pass in a new graph if needed
 
     # Get timestep from checkpoint. Also store a version in seconds for local use.
     config.timestep = None
diff --git a/bris/checkpoint.py b/bris/checkpoint.py
index c6d4c804..83142c21 100644
--- a/bris/checkpoint.py
+++ b/bris/checkpoint.py
@@ -1,5 +1,7 @@
 import logging
 import os
+
+from copy import deepcopy
 from functools import cached_property
 from typing import Any, Optional
 
@@ -115,6 +117,27 @@ def graph(self) -> HeteroData:
             if hasattr(self._model_instance, "graph_data")
             else None
         )
+    
+    @cached_property
+    def _model_params(self) -> dict:
+        """
+            The state of model being cached in CPU memory.
+            Keep in mind this is only the model weights and its
+            layer names, i.e does not include optimizer state.
+            It also worth mentioining this model.named_parameters()
+            and not model.state_dict.
+
+            Args:
+                None
+            Return
+                torch dict containing the state of the model.
+                Keys: name of the layer
+                Value: The state for a given layer
+        """
+
+        _model_params = tuple(self._model_instance.named_parameters())
+        return deepcopy({layer_name : param for layer_name, param in _model_params})
+        
 
     def update_graph(self, path: Optional[str] = None) -> HeteroData:
         """
@@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
+        else:
+            if path and os.path.exists(path):
+                
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
                 self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
+                self._model_instance.config = self.config #conf 
+
+                LOGGER.info("Rebuilding layers to support new graph")
+                
+                try: 
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
+                
+                except Exception as e:
+                    raise RuntimeError("Failed to rebuild model with new graph.") from e
+
+                _model_params = self._model_params
+                
+                for layer_name, param in self._model_instance.named_parameters():
+                    param.data = _model_params[layer_name].data
+
                 LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
+                    "Successfully builded model with external graph and reassigning model weights!"
                 )
                 return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
-        else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+                
+            else:
+                # future implementation
+                # _graph = anemoi.graphs.create() <-- skeleton
+                # self._model_instance.graph_data = _graph <- update graph obj within inst
+                # return _graph <- return graph
+                raise NotImplementedError
 
     @cached_property
     def set_base_seed(self) -> None:


Reviews:
user: login: havardhhaugen
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T12:44:43Z
user: login: einrone
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T07:21:07Z
user: login: einrone
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T07:27:03Z
user: login: tnipen
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T07:36:57Z
user: login: tnipen
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T07:41:45Z
user: login: havardhhaugen
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T08:08:32Z

Review comments:
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
+
+                    _model_params = self._model_params
+
+                    for layer_name, param in self._model_instance.named_parameters():
+                        param.data = _model_params[layer_name].data
+
+                    LOGGER.info(
+                        "Successfully builded model with external graph and reassigning model weights!"
+                    )
+                    return self._model_instance.graph_data
+
+                except Exception as e:
+                    raise RuntimeError("Failed to update the graph.") from e
path: bris/checkpoint.py
user: login: havardhhaugen
body: I don't think the try except here adds any information since the function is already named update_graph
created_at: 2025-03-06T12:39:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
+
+                    _model_params = self._model_params
+
+                    for layer_name, param in self._model_instance.named_parameters():
+                        param.data = _model_params[layer_name].data
+
+                    LOGGER.info(
+                        "Successfully builded model with external graph and reassigning model weights!"
path: bris/checkpoint.py
user: login: havardhhaugen
body: builded
created_at: 2025-03-06T12:39:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
path: bris/checkpoint.py
user: login: havardhhaugen
body: What is the purpose of self.UPDATE_GRAPH ?
created_at: 2025-03-06T12:44:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
path: bris/checkpoint.py
user: login: einrone
body: This is enforces that the graph can only be updated once during the code execution.

`if self.UPDATE_GRAPH:
            raise RuntimeError(
                "Graph has already been updated. Mutliple updates is not allowed"
            )`
created_at: 2025-03-07T07:21:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
+
+                    _model_params = self._model_params
+
+                    for layer_name, param in self._model_instance.named_parameters():
+                        param.data = _model_params[layer_name].data
+
+                    LOGGER.info(
+                        "Successfully builded model with external graph and reassigning model weights!"
+                    )
+                    return self._model_instance.graph_data
+
+                except Exception as e:
+                    raise RuntimeError("Failed to update the graph.") from e
path: bris/checkpoint.py
user: login: einrone
body: I don't think I understood you correctly, could you please elaborate?

 I also see now the intended use were to have try: execpt around self._model_instance._build_model() only, but essentially idea is to provide a useful error message with the exception as to why the model rebuild failed.
created_at: 2025-03-07T07:27:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
path: bris/checkpoint.py
user: login: tnipen
body: Why can't the graph be updated multiple times?
created_at: 2025-03-07T07:36:57Z
author_association: MEMBER
diff_hunk: @@ -139,29 +162,42 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
+        else:
+            if path:
+                assert os.path.exists(
path: bris/checkpoint.py
user: login: tnipen
body: I wouldn't use an assertion for argument checking.
created_at: 2025-03-07T07:39:19Z
author_association: MEMBER
diff_hunk: @@ -139,29 +162,40 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
             raise RuntimeError(
                 "Graph has already been updated. Mutliple updates is not allowed"
             )
-
-        if path:
-            assert os.path.exists(
-                path
-            ), f"Cannot locate graph file. Got path: {path}"
-            external_graph = torch.load(path, map_location="cpu")
-            LOGGER.info("Loaded external graph from path")
-            try:
-                self._model_instance.graph_data = external_graph
-                self.UPDATE_GRAPH = True
-                LOGGER.info(
-                    "Successfully changed internal graph with external graph!"
-                )
-                return self._model_instance.graph_data
-
-            except Exception as e:
-                raise e  # RuntimeError("Failed to update the graph.") from e
         else:
-            # future implementation
-            # _graph = anemoi.graphs.create() <-- skeleton
-            # self._model_instance.graph_data = _graph <- update graph obj within inst
-            # return _graph <- return graph
-            raise NotImplementedError
+            if path:
+                assert os.path.exists(
+                    path
+                ), f"Cannot locate graph file. Got path: {path}"
+                external_graph = torch.load(path, map_location="cpu",weights_only=False)
+                LOGGER.info("Loaded external graph from path")
+                
+                try:
+                    self._model_instance.graph_data = external_graph
+                    self._model_instance.config = self.config #conf 
+
+                    LOGGER.info("Rebuilding layers to support new graph")
+                    self._model_instance._build_model()
+                    self.UPDATE_GRAPH = True
+
+                    _model_params = self._model_params
+
+                    for layer_name, param in self._model_instance.named_parameters():
+                        param.data = _model_params[layer_name].data
+
+                    LOGGER.info(
+                        "Successfully builded model with external graph and reassigning model weights!"
+                    )
+                    return self._model_instance.graph_data
+
+                except Exception as e:
+                    raise RuntimeError("Failed to update the graph.") from e
path: bris/checkpoint.py
user: login: havardhhaugen
body: I meant that you don't have to wrap code-blocks in a try except when the printed message does not add information about what goes wrong. In this instance the error message from _build_model() is more useful than "Failed to update the graph".
created_at: 2025-03-07T08:08:32Z
author_association: CONTRIBUTOR

PR comments:
user: login: ways
created_at: 2025-03-06T13:06:43Z
body: Can you supply a valid config example for this?

Also, at least for training, we have a config option called

```
graph:
  overwrite: True
```

Perhaps this should use somethink like:

```
graph:
  path: ...
```

author_association: CONTRIBUTOR
user: login: einrone
created_at: 2025-03-07T07:19:19Z
body: @ways I will documents this in the wiki, but as stated in the comment above the intended use (for now) is 
to provide a graph path, and this will trigger a rebuild of a new model where the internal graph has been switched with an external graph. 

> Can you supply a valid config example for this?
> 
> Also, at least for training, we have a config option called
> 
> ```
> graph:
>   overwrite: True
> ```
> 
> Perhaps this should use somethink like:
> 
> ```
> graph:
>   path: ...
> ```


author_association: CONTRIBUTOR
user: login: einrone
created_at: 2025-03-07T08:43:52Z
body: > Can you supply a valid config example for this?
> 
> Also, at least for training, we have a config option called
> 
> ```
> graph:
>   overwrite: True
> ```
> 
> Perhaps this should use somethink like:
> 
> ```
> graph:
>   path: ...
> ```

wiki: https://github.com/metno/bris-inference/wiki/Changing-the-grid
author_association: CONTRIBUTOR
user: login: ways
created_at: 2025-03-07T09:20:43Z
body: My suggestion: https://github.com/metno/bris-inference/tree/19-larsfp => https://github.com/metno/bris-inference/pull/67

Have tested a run as far as I can get on a laptop (which includes the replacing of graph).
author_association: CONTRIBUTOR