From b86ba0c2bbbed323af578383b5e1e65f46f18efd Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Sun, 11 Aug 2024 18:05:10 +0800
Subject: [PATCH 1/9] Unique session in each process: version 1

---
 pygmt/__init__.py           |  3 ---
 pygmt/_state.py             | 10 ++++++++++
 pygmt/clib/session.py       | 14 ++++++++++++++
 pygmt/session_management.py | 11 +++--------
 4 files changed, 27 insertions(+), 11 deletions(-)
 create mode 100644 pygmt/_state.py

diff --git a/pygmt/__init__.py b/pygmt/__init__.py
index dbf292e4f1f..338f57fc68d 100644
--- a/pygmt/__init__.py
+++ b/pygmt/__init__.py
@@ -26,7 +26,6 @@
 from pygmt.accessors import GMTDataArrayAccessor
 from pygmt.figure import Figure, set_display
 from pygmt.io import load_dataarray
-from pygmt.session_management import begin as _begin
 from pygmt.session_management import end as _end
 from pygmt.src import (
     binstats,
@@ -66,7 +65,5 @@
     xyz2grd,
 )
 
-# Start our global modern mode session
-_begin()
 # Tell Python to run _end when shutting down
 _atexit.register(_end)
diff --git a/pygmt/_state.py b/pygmt/_state.py
new file mode 100644
index 00000000000..d8f941b3440
--- /dev/null
+++ b/pygmt/_state.py
@@ -0,0 +1,10 @@
+"""
+Private dictionary to keep tracking of current PyGMT state.
+
+The feature is only meant for internal use by PyGMT and is experimental!
+"""
+
+_STATE = {
+    "session_name": None,
+    "module_calls": None,
+}
diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index a0c11373d68..a70fc2a5f06 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -7,6 +7,7 @@
 
 import contextlib
 import ctypes as ctp
+import os
 import pathlib
 import sys
 import warnings
@@ -17,6 +18,7 @@
 import pandas as pd
 import xarray as xr
 from packaging.version import Version
+from pygmt._state import _STATE
 from pygmt.clib.conversion import (
     array_to_datetime,
     as_c_contiguous,
@@ -33,6 +35,7 @@
     data_kind,
     tempfile_from_geojson,
     tempfile_from_image,
+    unique_name,
 )
 
 FAMILIES = [
@@ -209,6 +212,10 @@ def __enter__(self):
 
         Calls :meth:`pygmt.clib.Session.create`.
         """
+        # This is the first time a Session object is created.
+        if _STATE["session_name"] is None:
+            # Set GMT_SESSION_NAME to a customized, unique value.
+            _STATE["session_name"] = os.environ["GMT_SESSION_NAME"] = unique_name()
         self.create("pygmt-session")
         return self
 
@@ -623,6 +630,12 @@ def call_module(self, module: str, args: str | list[str]):
         GMTCLibError
             If the returned status code of the function is non-zero.
         """
+        if _STATE["module_calls"] is None:
+            from pygmt.session_management import begin
+
+            _STATE["module_calls"] = []
+            begin()
+
         c_call_module = self.get_libgmt_func(
             "GMT_Call_Module",
             argtypes=[ctp.c_void_p, ctp.c_char_p, ctp.c_int, ctp.c_void_p],
@@ -651,6 +664,7 @@ def call_module(self, module: str, args: str | list[str]):
                 "'args' must be either a string or a list of strings."
             )
 
+        _STATE["module_calls"].append(module)
         status = c_call_module(self.session_pointer, module.encode(), mode, argv)
         if status != 0:
             raise GMTCLibError(
diff --git a/pygmt/session_management.py b/pygmt/session_management.py
index 87055bb44e8..d3234cbda63 100644
--- a/pygmt/session_management.py
+++ b/pygmt/session_management.py
@@ -2,11 +2,8 @@
 Modern mode session management modules.
 """
 
-import os
-import sys
-
+from pygmt._state import _STATE
 from pygmt.clib import Session
-from pygmt.helpers import unique_name
 
 
 def begin():
@@ -17,10 +14,6 @@ def begin():
 
     Only meant to be used once for creating the global session.
     """
-    # On Windows, need to set GMT_SESSION_NAME to a unique value
-    if sys.platform == "win32":
-        os.environ["GMT_SESSION_NAME"] = unique_name()
-
     prefix = "pygmt-session"
     with Session() as lib:
         lib.call_module(module="begin", args=[prefix])
@@ -39,3 +32,5 @@ def end():
     """
     with Session() as lib:
         lib.call_module(module="end", args=[])
+
+    _STATE["session_name"] = None  # Reset the sesion name to None

From be9a202ace7e4cdd7f895bb5eba52fa3870b6546 Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Tue, 13 Aug 2024 22:15:38 +0800
Subject: [PATCH 2/9] Unique session in each process: version 2

---
 pygmt/_state.py       |  1 -
 pygmt/clib/session.py | 16 +++++++++-------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/pygmt/_state.py b/pygmt/_state.py
index d8f941b3440..3586d759a54 100644
--- a/pygmt/_state.py
+++ b/pygmt/_state.py
@@ -6,5 +6,4 @@
 
 _STATE = {
     "session_name": None,
-    "module_calls": None,
 }
diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index a70fc2a5f06..0b6e69a088b 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -212,11 +212,20 @@ def __enter__(self):
 
         Calls :meth:`pygmt.clib.Session.create`.
         """
+        _init_cli_session = False
         # This is the first time a Session object is created.
         if _STATE["session_name"] is None:
             # Set GMT_SESSION_NAME to a customized, unique value.
             _STATE["session_name"] = os.environ["GMT_SESSION_NAME"] = unique_name()
+            # Need to initialize the GMT CLI session.
+            _init_cli_session = True
         self.create("pygmt-session")
+
+        if _init_cli_session:
+            self.call_module("begin", args=["pygmt-session"])
+            self.call_module(module="set", args=["GMT_COMPATIBILITY=6"])
+            del _init_cli_session
+
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -630,12 +639,6 @@ def call_module(self, module: str, args: str | list[str]):
         GMTCLibError
             If the returned status code of the function is non-zero.
         """
-        if _STATE["module_calls"] is None:
-            from pygmt.session_management import begin
-
-            _STATE["module_calls"] = []
-            begin()
-
         c_call_module = self.get_libgmt_func(
             "GMT_Call_Module",
             argtypes=[ctp.c_void_p, ctp.c_char_p, ctp.c_int, ctp.c_void_p],
@@ -664,7 +667,6 @@ def call_module(self, module: str, args: str | list[str]):
                 "'args' must be either a string or a list of strings."
             )
 
-        _STATE["module_calls"].append(module)
         status = c_call_module(self.session_pointer, module.encode(), mode, argv)
         if status != 0:
             raise GMTCLibError(

From 7753b9c111f41849bab794043d345697307c9f1d Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Sat, 24 Aug 2024 22:40:19 +0800
Subject: [PATCH 3/9] Fix test_gmt_compat_6_is_applied

---
 pygmt/tests/test_session_management.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/pygmt/tests/test_session_management.py b/pygmt/tests/test_session_management.py
index d949f1a51c0..289e4f15b02 100644
--- a/pygmt/tests/test_session_management.py
+++ b/pygmt/tests/test_session_management.py
@@ -36,10 +36,8 @@ def test_gmt_compat_6_is_applied(capsys):
     """
     end()  # Kill the global session
     try:
-        # Generate a gmt.conf file in the current directory
-        # with GMT_COMPATIBILITY = 5
-        with Session() as lib:
-            lib.call_module("gmtset", ["GMT_COMPATIBILITY=5"])
+        # Generate a gmt.conf file in the current directory with GMT_COMPATIBILITY = 5
+        Path("gmt.conf").write_text("GMT_COMPATIBILITY = 5", encoding="utf-8")
         begin()
         with Session() as lib:
             lib.call_module("basemap", ["-R10/70/-3/8", "-JX4i/3i", "-Ba"])

From 90ad065ca585f39ef4f21dadf991b71e4b77cc24 Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Sat, 24 Aug 2024 22:42:03 +0800
Subject: [PATCH 4/9] Set GMT_SESSION_NAME to the current process id

---
 pygmt/clib/session.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/pygmt/clib/session.py b/pygmt/clib/session.py
index 0b6e69a088b..80435d9d8a3 100644
--- a/pygmt/clib/session.py
+++ b/pygmt/clib/session.py
@@ -35,7 +35,6 @@
     data_kind,
     tempfile_from_geojson,
     tempfile_from_image,
-    unique_name,
 )
 
 FAMILIES = [
@@ -215,8 +214,8 @@ def __enter__(self):
         _init_cli_session = False
         # This is the first time a Session object is created.
         if _STATE["session_name"] is None:
-            # Set GMT_SESSION_NAME to a customized, unique value.
-            _STATE["session_name"] = os.environ["GMT_SESSION_NAME"] = unique_name()
+            # Set GMT_SESSION_NAME to the current process id.
+            _STATE["session_name"] = os.environ["GMT_SESSION_NAME"] = str(os.getpid())
             # Need to initialize the GMT CLI session.
             _init_cli_session = True
         self.create("pygmt-session")

From 91ffeb9ffd22038cf94b918c2de4ab31bc38f9ac Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Sat, 24 Aug 2024 22:46:43 +0800
Subject: [PATCH 5/9] No need to reload

---
 pygmt/tests/test_session_management.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/pygmt/tests/test_session_management.py b/pygmt/tests/test_session_management.py
index 289e4f15b02..6c3c52501d4 100644
--- a/pygmt/tests/test_session_management.py
+++ b/pygmt/tests/test_session_management.py
@@ -3,10 +3,10 @@
 """
 
 import multiprocessing as mp
-from importlib import reload
 from pathlib import Path
 
 import pytest
+from pygmt import Figure
 from pygmt.clib import Session
 from pygmt.session_management import begin, end
 
@@ -67,10 +67,7 @@ def _gmt_func_wrapper(figname):
     Currently, we have to import pygmt and reload it in each process. Workaround from
     https://github.com/GenericMappingTools/pygmt/issues/217#issuecomment-754774875.
     """
-    import pygmt
-
-    reload(pygmt)
-    fig = pygmt.Figure()
+    fig = Figure()
     fig.basemap(region=[10, 70, -3, 8], projection="X8c/6c", frame="afg")
     fig.savefig(figname)
 

From 66ee9242a74f433edcb0d50e54215777206ab43f Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Thu, 19 Sep 2024 08:27:54 +0800
Subject: [PATCH 6/9] Move multiprocessing tests into a separate test file

---
 pygmt/tests/test_multiprocessing.py    | 28 ++++++++++++++++++++++++++
 pygmt/tests/test_session_management.py | 25 -----------------------
 2 files changed, 28 insertions(+), 25 deletions(-)
 create mode 100644 pygmt/tests/test_multiprocessing.py

diff --git a/pygmt/tests/test_multiprocessing.py b/pygmt/tests/test_multiprocessing.py
new file mode 100644
index 00000000000..870db3878d8
--- /dev/null
+++ b/pygmt/tests/test_multiprocessing.py
@@ -0,0 +1,28 @@
+"""
+Test multiprocessing support.
+"""
+
+import multiprocessing as mp
+from pathlib import Path
+
+from pygmt import Figure
+
+
+def _func(figname):
+    """
+    A wrapper function for testing multiprocessing support.
+    """
+    fig = Figure()
+    fig.basemap(region=[10, 70, -3, 8], projection="X8c/6c", frame="afg")
+    fig.savefig(figname)
+
+
+def test_multiprocessing():
+    """
+    Make sure that multiprocessing is supported if pygmt is re-imported.
+    """
+    prefix = "test_session_multiprocessing"
+    with mp.Pool(2) as p:
+        p.map(_func, [f"{prefix}-1.png", f"{prefix}-2.png"])
+    Path(f"{prefix}-1.png").unlink()
+    Path(f"{prefix}-2.png").unlink()
diff --git a/pygmt/tests/test_session_management.py b/pygmt/tests/test_session_management.py
index 6c3c52501d4..1a811343b90 100644
--- a/pygmt/tests/test_session_management.py
+++ b/pygmt/tests/test_session_management.py
@@ -2,11 +2,9 @@
 Test the session management modules.
 """
 
-import multiprocessing as mp
 from pathlib import Path
 
 import pytest
-from pygmt import Figure
 from pygmt.clib import Session
 from pygmt.session_management import begin, end
 
@@ -58,26 +56,3 @@ def test_gmt_compat_6_is_applied(capsys):
         # Make sure no global "gmt.conf" in the current directory
         assert not Path("gmt.conf").exists()
         begin()  # Restart the global session
-
-
-def _gmt_func_wrapper(figname):
-    """
-    A wrapper for running PyGMT scripts with multiprocessing.
-
-    Currently, we have to import pygmt and reload it in each process. Workaround from
-    https://github.com/GenericMappingTools/pygmt/issues/217#issuecomment-754774875.
-    """
-    fig = Figure()
-    fig.basemap(region=[10, 70, -3, 8], projection="X8c/6c", frame="afg")
-    fig.savefig(figname)
-
-
-def test_session_multiprocessing():
-    """
-    Make sure that multiprocessing is supported if pygmt is re-imported.
-    """
-    prefix = "test_session_multiprocessing"
-    with mp.Pool(2) as p:
-        p.map(_gmt_func_wrapper, [f"{prefix}-1.png", f"{prefix}-2.png"])
-    Path(f"{prefix}-1.png").unlink()
-    Path(f"{prefix}-2.png").unlink()

From a392ee078fb1245d695db1cf45e5282f261a7d3f Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Thu, 19 Sep 2024 08:32:14 +0800
Subject: [PATCH 7/9] Add the old test back

---
 pygmt/tests/test_multiprocessing.py | 31 +++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/pygmt/tests/test_multiprocessing.py b/pygmt/tests/test_multiprocessing.py
index 870db3878d8..e00eb58a8dc 100644
--- a/pygmt/tests/test_multiprocessing.py
+++ b/pygmt/tests/test_multiprocessing.py
@@ -3,6 +3,7 @@
 """
 
 import multiprocessing as mp
+from importlib import reload
 from pathlib import Path
 
 from pygmt import Figure
@@ -26,3 +27,33 @@ def test_multiprocessing():
         p.map(_func, [f"{prefix}-1.png", f"{prefix}-2.png"])
     Path(f"{prefix}-1.png").unlink()
     Path(f"{prefix}-2.png").unlink()
+
+
+def _func_reload(figname):
+    """
+    A wrapper for running PyGMT scripts with multiprocessing.
+
+    Before the official multiprocessing support in PyGMT, we have to reload the
+    PyGMT library. Workaround from
+    https://github.com/GenericMappingTools/pygmt/issues/217#issuecomment-754774875.
+
+    This test makes sure that the old workaround still works.
+    """
+    import pygmt
+
+    reload(pygmt)
+    fig = pygmt.Figure()
+    fig.basemap(region=[10, 70, -3, 8], projection="X8c/6c", frame="afg")
+    fig.savefig(figname)
+
+
+def test_multiprocessing_reload():
+    """
+    Make sure that multiprocessing is supported if pygmt is re-imported.
+    """
+
+    prefix = "test_session_multiprocessing"
+    with mp.Pool(2) as p:
+        p.map(_func_reload, [f"{prefix}-1.png", f"{prefix}-2.png"])
+    Path(f"{prefix}-1.png").unlink()
+    Path(f"{prefix}-2.png").unlink()

From 8620c174d9839842173162fe1a869de5017746ed Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Thu, 19 Sep 2024 08:50:03 +0800
Subject: [PATCH 8/9] Add one more test

---
 pygmt/tests/test_multiprocessing.py | 32 ++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/pygmt/tests/test_multiprocessing.py b/pygmt/tests/test_multiprocessing.py
index e00eb58a8dc..01f4af77bf6 100644
--- a/pygmt/tests/test_multiprocessing.py
+++ b/pygmt/tests/test_multiprocessing.py
@@ -6,21 +6,22 @@
 from importlib import reload
 from pathlib import Path
 
-from pygmt import Figure
+import numpy.testing as npt
+import pygmt
 
 
 def _func(figname):
     """
     A wrapper function for testing multiprocessing support.
     """
-    fig = Figure()
+    fig = pygmt.Figure()
     fig.basemap(region=[10, 70, -3, 8], projection="X8c/6c", frame="afg")
     fig.savefig(figname)
 
 
 def test_multiprocessing():
     """
-    Make sure that multiprocessing is supported if pygmt is re-imported.
+    Test multiprocessing support for plotting figures.
     """
     prefix = "test_session_multiprocessing"
     with mp.Pool(2) as p:
@@ -29,6 +30,31 @@ def test_multiprocessing():
     Path(f"{prefix}-2.png").unlink()
 
 
+def _func_datacut(dataset):
+    """
+    A wrapper function for testing multiprocessing support.
+    """
+    xrgrid = pygmt.grdcut(dataset, region=[-10, 10, -5, 5])
+    return xrgrid
+
+
+def test_multiprocessing_data_processing():
+    """
+    Test multiprocessing support for data processing.
+    """
+    with mp.Pool(2) as p:
+        grids = p.map(_func_datacut, ["@earth_relief_01d_g", "@moon_relief_01d_g"])
+    assert len(grids) == 2
+    # The Earth relief dataset
+    assert grids[0].shape == (11, 21)
+    npt.assert_allclose(grids[0].min(), -5118.0, atol=0.5)
+    npt.assert_allclose(grids[0].max(), 680.5, atol=0.5)
+    # The Moon relief dataset
+    assert grids[1].shape == (11, 21)
+    npt.assert_allclose(grids[0].min(), -1122.0, atol=0.5)
+    npt.assert_allclose(grids[0].max(), 943.0, atol=0.5)
+
+
 def _func_reload(figname):
     """
     A wrapper for running PyGMT scripts with multiprocessing.

From 1d61c790e356e07cb565381f4c19b4827ad0e41e Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Thu, 19 Sep 2024 08:56:45 +0800
Subject: [PATCH 9/9] Fix a typo

---
 pygmt/tests/test_multiprocessing.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pygmt/tests/test_multiprocessing.py b/pygmt/tests/test_multiprocessing.py
index 01f4af77bf6..bfb3a3d3983 100644
--- a/pygmt/tests/test_multiprocessing.py
+++ b/pygmt/tests/test_multiprocessing.py
@@ -51,8 +51,8 @@ def test_multiprocessing_data_processing():
     npt.assert_allclose(grids[0].max(), 680.5, atol=0.5)
     # The Moon relief dataset
     assert grids[1].shape == (11, 21)
-    npt.assert_allclose(grids[0].min(), -1122.0, atol=0.5)
-    npt.assert_allclose(grids[0].max(), 943.0, atol=0.5)
+    npt.assert_allclose(grids[1].min(), -1122.0, atol=0.5)
+    npt.assert_allclose(grids[1].max(), 943.0, atol=0.5)
 
 
 def _func_reload(figname):