Repository: graphnet-team/graphnet
Issue: model.save_config() fails

**Describe the bug**
When I attempt to save the model config using `model.save_config` the procedure fails.

**To Reproduce**
Steps to reproduce the behavior:
1. make a copy of ` 01_train_model.py` training script
2. Add `model.save_config('test_config.yaml')` anywhere you'd like after line 52.
3. Run script
4. See error

**Expected behavior**
Config should save without error.

**Full traceback**
Please include the full error message to allow for debugging
```
graphnet [MainProcess] INFO     2023-03-08 12:22:55 - Logger.__init__ - Writing log to logs/graphnet_20230308-122255.log
wandb: (1) Create a W&B account
wandb: (2) Use an existing W&B account
wandb: (3) Don't visualize my results
wandb: Enter your choice: 3
wandb: You chose 'Don't visualize my results'
wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id 1s3nx38z.
wandb: Tracking run with wandb version 0.13.0rc2
wandb: W&B syncing is set to `offline` in this directory.  
wandb: Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
graphnet [MainProcess] WARNING  2023-03-08 12:23:04 - has_icecube_package - `icecube` not available. Some functionality may be missing.
graphnet [MainProcess] INFO     2023-03-08 12:23:05 - StringSelectionResolver.resolve - Resolving selection: event_no % 5 == 0
graphnet [MainProcess] INFO     2023-03-08 12:23:05 - StringSelectionResolver.resolve - Resolving selection: event_no % 5 == 1
graphnet [MainProcess] INFO     2023-03-08 12:23:05 - StringSelectionResolver.resolve - Resolving selection: event_no % 5 > 1
graphnet [MainProcess] INFO     2023-03-08 12:23:06 - StandardModel.save - Model saved to test_model.pth
Traceback (most recent call last):
  File "/home/iwsatlas1/oersoe/phd/tmp/model_config_debug.py", line 178, in <module>
    main(
  File "/home/iwsatlas1/oersoe/phd/tmp/model_config_debug.py", line 97, in main
    model.save_config('test_config.yaml')
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/configurable.py", line 35, in save_config
    self.config.dump(path)
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/base_config.py", line 33, in dump
    config_dict = self.as_dict()[self.__class__.__name__]
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/model_config.py", line 225, in as_dict
    config_dict["arguments"] = traverse_and_apply(
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 24, in traverse_and_apply
    return {
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 25, in <dictcomp>
    key: traverse_and_apply(val, fn, fn_kwargs)
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 22, in traverse_and_apply
    return [traverse_and_apply(elem, fn, fn_kwargs) for elem in obj]
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 22, in <listcomp>
    return [traverse_and_apply(elem, fn, fn_kwargs) for elem in obj]
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 31, in traverse_and_apply
    return fn(obj, **fn_kwargs)
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/model_config.py", line 193, in _serialise
    return obj.as_dict()
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/model_config.py", line 225, in as_dict
    config_dict["arguments"] = traverse_and_apply(
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 24, in traverse_and_apply
    return {
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 25, in <dictcomp>
    key: traverse_and_apply(val, fn, fn_kwargs)
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/parsing.py", line 31, in traverse_and_apply
    return fn(obj, **fn_kwargs)
  File "/home/iwsatlas1/oersoe/github/graphnet/src/graphnet/utilities/config/model_config.py", line 198, in _serialise
    return "!" + inspect.getsource(obj).split("=")[1].strip("\\n ,")
  File "/opt/anaconda3/lib/python3.9/inspect.py", line 1024, in getsource
    lines, lnum = getsourcelines(object)
  File "/opt/anaconda3/lib/python3.9/inspect.py", line 1006, in getsourcelines
    lines, lnum = findsource(object)
  File "/opt/anaconda3/lib/python3.9/inspect.py", line 835, in findsource
    raise OSError('could not get source code')
OSError: could not get source code
wandb: Waiting for W&B process to finish... (failed 1).
wandb:                                                                                
wandb: You can sync this run to the cloud by running:
wandb: wandb sync ./wandb/wandb/offline-run-20230308_122259-1s3nx38z
wandb: Find logs at: ./wandb/wandb/offline-run-20230308_122259-1s3nx38z/logs
```

**Additional context**
Add any other context about the problem here.


Changed files:
src/graphnet/utilities/config/model_config.py, tests/utilities/test_model_config.py

diff --git a/src/graphnet/utilities/config/model_config.py b/src/graphnet/utilities/config/model_config.py
index 52adcac1..d4733ca6 100644
--- a/src/graphnet/utilities/config/model_config.py
+++ b/src/graphnet/utilities/config/model_config.py
@@ -145,7 +145,15 @@ class ModelConfig(BaseConfig):
 
         elif isinstance(obj, str) and obj.startswith("!lambda"):
             if trust:
-                return eval(obj[1:])
+                source = obj[1:]
+                f = eval(source)
+
+                # Save a copy of the source code attached to the callable,
+                # since the `inspect` module is not able to get the source code
+                # for functions that are not defined on file.
+                # See `self._serialise`.
+                f._source = source
+                return f
             else:
                 raise ValueError(
                     "Constructing model containing a lambda function "
@@ -195,7 +203,14 @@ class ModelConfig(BaseConfig):
             return f"!class {obj.__module__} {obj.__name__}"
         elif isinstance(obj, Callable):  # type: ignore[arg-type]
             if hasattr(obj, "__name__") and obj.__name__ == "<lambda>":
-                return "!" + inspect.getsource(obj).split("=")[1].strip("\\n ,")
+                if hasattr(obj, "_source"):
+                    # If source code is set manually during deserialisation.
+                    # See `self._deserialise`.
+                    source = obj._source
+                else:
+                    source = inspect.getsource(obj).split("=")[1].strip("\\n ,")
+
+                return "!" + source
             else:
                 try:
                     source = inspect.getsource(obj)
diff --git a/tests/utilities/test_model_config.py b/tests/utilities/test_model_config.py
index fe605ae1..13370103 100644
--- a/tests/utilities/test_model_config.py
+++ b/tests/utilities/test_model_config.py
@@ -2,6 +2,7 @@
 
 import os.path
 
+import pytest
 import torch
 from torch.optim.adam import Adam
 
@@ -143,4 +144,19 @@ def test_complete_model_config(path: str = "/tmp/complete_model.yml") -> None:
     assert repr(constructed_model) == repr(model)
 
 
-test_complete_model_config()
+@pytest.mark.run(after="test_complete_model_config")
+def test_loading_and_saving_model_config(
+    load_path: str = "/tmp/complete_model.yml",
+    save_path: str = "/tmp/complete_model_duplicate.yml",
+) -> None:
+    """Test loading and then saving model config."""
+    # Load and construct model
+    loaded_config = ModelConfig.load(load_path)
+    model = Model.from_config(loaded_config, trust=True)
+
+    # Save config to file
+    model.save_config(save_path)
+
+    # Check that the configs agree
+    saved_config = ModelConfig.load(save_path)
+    assert saved_config == loaded_config