Repository: netket/netket

Issue:
TensorBoardLog with MPI
Not sure this can be solved in netket, but when using TensorBoardLog with MPI, I'm getting errors.

MWE (mpitest.py):
```python
import netket as nk
logger = nk.logging.TensorBoardLog("test_mpi_log")
```

When you run the above with
```console
mpirun -np 4 python ./mpitest.py
```

You'll see messages like:
```console
KeyError: 'test_mpi_log'

FileExistsError: [Errno 17] File exists: 'test_mpi_log'
Traceback (most recent call last):
  File "/usr/local/anaconda3/envs/localjw/lib/python3.8/site-packages/tensorboardX/record_writer.py", line 47, in directory_check
    factory = REGISTERED_FACTORIES[prefix]
```






Discussion:
We should initialize the internal logger structures only on rank 0.

As a workaround, create the logger only on rank 0 and leave it None elsewhere
This also affects the other loggers right? There is a reason our default logger initialization looks like this:
https://github.com/netket/netket/blob/bdf7071ff4000a31adce360a86f132708fc6db82/netket/driver/abstract_variational_driver.py#L235-L243
No, because the other loggers (`StateLog` and `JsonLog`) are lazily initialised: they don't create the output files until they actually have to write them (see [this for example](https://github.com/netket/netket/blob/bdf7071ff4000a31adce360a86f132708fc6db82/netket/logging/state_log.py#L165)). And since our drivers don't log on `rank != 0`, this is not an issue.

Basically, as long as a logger is lazily initialised, there are no issues.

The choice was made so that loggers are not MPI-aware (that's why they should be lazy)

Resolving patch:
diff --git a/netket/logging/tensorboard.py b/netket/logging/tensorboard.py
--- a/netket/logging/tensorboard.py
+++ b/netket/logging/tensorboard.py
@@ -114,13 +114,24 @@ def __init__(
         *args,
         **kwargs,
     ):
-        from tensorboardX import SummaryWriter
+        self._init_args = args
+        """Store the args for the lazily initialized SummaryWriter's constructor."""
+        self._init_kwargs = kwargs
+        """Store the kwargs for the lazily initialized SummaryWriter's constructor."""
 
-        self._writer = SummaryWriter(*args, **kwargs)
+        self._writer = None
+        """Lazily initialized summarywriter constructor"""
 
         self._old_step = 0
 
+    def _init_tensoboard(self):
+        from tensorboardX import SummaryWriter
+
+        self._writer = SummaryWriter(*self._init_args, **self._init_kwargs)
+
     def __call__(self, step, item, machine):
+        if self._writer is None:
+            self._init_tensoboard()
 
         data = []
         tree_log(item, "", data)
@@ -133,7 +144,8 @@ def __call__(self, step, item, machine):
         self._old_step = step
 
     def _flush_log(self):
-        self._writer.flush()
+        if self._writer is not None:
+            self._writer.flush()
 
     def _flush_params(self, _):
         return None


Test patch:
diff --git a/test/logging/test_state_log.py b/test/logging/test_state_log.py
--- a/test/logging/test_state_log.py
+++ b/test/logging/test_state_log.py
@@ -111,3 +111,13 @@ def test_dir(vstate, tmp_path):
 
     for file in files:
         assert file.endswith(".mpack")
+
+
+def test_lazy_init(tmp_path):
+    path = str(tmp_path) + "/dir1/dir2"
+
+    # check that overwriting works
+    log = nk.logging.StateLog(path, "w", tar=False, save_every=1)
+
+    files = glob.glob(path + "/*")
+    assert len(files) == 0
diff --git a/test/logging/test_tensorboard.py b/test/logging/test_tensorboard.py
new file mode 100644
--- /dev/null
+++ b/test/logging/test_tensorboard.py
@@ -0,0 +1,60 @@
+import pytest
+
+import tarfile
+import glob
+
+import netket as nk
+from jax.nn.initializers import normal
+from jax import numpy as jnp
+
+from .. import common
+
+pytestmark = common.skipif_mpi
+
+
+@pytest.fixture()
+def vstate(request):
+    N = 8
+    hi = nk.hilbert.Spin(1 / 2, N)
+
+    ma = nk.models.RBM(
+        alpha=1,
+        dtype=float,
+        hidden_bias_init=normal(),
+        visible_bias_init=normal(),
+    )
+
+    return nk.vqs.MCState(
+        nk.sampler.MetropolisLocal(hi),
+        ma,
+    )
+
+
+def test_tblog(vstate, tmp_path):
+    # skip test if tensorboardX not installed
+    pytest.importorskip("tensorboardX")
+
+    path = str(tmp_path) + "/dir1/dir2"
+
+    log = nk.logging.TensorBoardLog(path)
+
+    for i in range(10):
+        log(i, {"Energy": jnp.array(1.0), "complex": jnp.array(1.0 + 1j)}, vstate)
+
+    log.flush()
+    del log
+
+    files = glob.glob(path + "/*")
+    assert len(files) >= 1
+
+
+def test_lazy_init(tmp_path):
+    # skip test if tensorboardX not installed
+    pytest.importorskip("tensorboardX")
+
+    path = str(tmp_path) + "/dir1"
+
+    log = nk.logging.TensorBoardLog(path)
+
+    files = glob.glob(path + "/*")
+    assert len(files) == 0