smartg 2.0.2__tar.gz → 2.0.4__tar.gz

{smartg-2.0.2/smartg.egg-info → smartg-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: smartg
-Version: 2.0.2
+Version: 2.0.4
 Summary: SMART-G (Speed-up Monte Carlo Advanced Radiative Transfer Code using GPU) is a radiative transfer code using a Monte-Carlo technique to simulate the propagation of the polarized light in the atmosphere and/or ocean, and using GPU acceleration.
 Author-email: Mustapha Moulana <mm@hygeos.com>, Didier Ramon <dr@hygeos.com>, François Steinmetz <fs@hygeos.com>, Dominique Jolivet <dj@hygeos.com>, Mathieu Compiègne <mc@hygeos.com>
 Project-URL: Source, https://github.com/hygeos/smartg
@@ -326,6 +326,11 @@ environment. Without them it skips itself with a message rather than
 failing, so run it with `pixi run -e extra pytest
 smartg/tests/test_smartg_jax.py`.
 
+`test_jax_context.py` needs JAX on a GPU, in the same environment, and
+a pytest process of its own: `test_smartg_jax.py` keeps JAX on the CPU
+for the whole process, and the test then skips itself. Run it with
+`pixi run -e extra pytest smartg/tests/test_jax_context.py`.
+
 To avoid repeating some pytest arguments, a `pytest.ini` file can be created at the root of the project. The following is an example of the contents of such a file:
 
 ```

{smartg-2.0.2 → smartg-2.0.4}/README.md

@@ -275,6 +275,11 @@ environment. Without them it skips itself with a message rather than
 failing, so run it with `pixi run -e extra pytest
 smartg/tests/test_smartg_jax.py`.
 
+`test_jax_context.py` needs JAX on a GPU, in the same environment, and
+a pytest process of its own: `test_smartg_jax.py` keeps JAX on the CPU
+for the whole process, and the test then skips itself. Run it with
+`pixi run -e extra pytest smartg/tests/test_jax_context.py`.
+
 To avoid repeating some pytest arguments, a `pytest.ini` file can be created at the root of the project. The following is an example of the contents of such a file:
 
 ```

{smartg-2.0.2 → smartg-2.0.4}/pyproject.toml

@@ -4,7 +4,7 @@ requires = ["setuptools>=61"]
 
 [project]
 name = "smartg"
-version = "2.0.2"
+version = "2.0.4"
 description = """
 SMART-G (Speed-up Monte Carlo Advanced Radiative Transfer Code using GPU) is a radiative \
 transfer code using a Monte-Carlo technique to simulate the propagation of the polarized \

{smartg-2.0.2 → smartg-2.0.4}/smartg/smartg.py

@@ -32,6 +32,8 @@ reduce_diff
     Post-process ALIS finite-difference runs into sensitivities.
 """
 
+import atexit
+import functools
 import os
 import subprocess
 import tempfile
@@ -402,7 +404,7 @@ class Alis:
     """
     The options of the ALIS method, for Smartg.run.
 
-    ALIS (Emde et al. 2010) computes the spectrum of a whole band from
+    ALIS (Emde et al. 2011) computes the spectrum of a whole band from
     a single set of photon paths, by weighting each path with the
     absorption of every wavelength. The Smartg object must have been
     compiled with alis=True, together with alt_pp=True or pp=False.
@@ -562,7 +564,10 @@ def _check_alis_kernel(
 
 
 def _check_objects_kernel(
-    my_objects: list | None, alis: bool, pp: bool, opt3d: bool
+    my_objects: Sequence[Entity | GroupE] | None,
+    alis: bool,
+    pp: bool,
+    opt3d: bool,
 ) -> None:
     """
     Refuse the 3D objects in a kernel whose move does not follow them.
@@ -850,6 +855,122 @@ class LocalEstimate:
         )
 
 
+# The CUDA contexts that clear_context destroyed, kept referenced so
+# that no new context takes the place of one of them. The Smartg objects
+# of a process share the context of pycuda.autoinit: once one of them
+# clears it, the others must not make it current again, which crashes
+# the driver.
+_RELEASED_CONTEXTS: list[Any] = []
+
+
+def _released(ctx: Any) -> bool:
+    """Return whether ctx is a context that clear_context destroyed."""
+    return any(c is ctx for c in _RELEASED_CONTEXTS)
+
+
+def _make_current(ctx: Any) -> bool:
+    """Make the CUDA context ctx the current one, if it is not.
+
+    Another library of the process (JAX for instance) may have made its
+    own context current, where the kernels of SMART-G are unknown
+    ('invalid resource handle'). Returns True if ctx was pushed, to be
+    popped by the caller. Raises a RuntimeError if ctx was destroyed by
+    the clear_context of another Smartg object that shared it.
+    """
+    if ctx is None:
+        return False
+    if _released(ctx):
+        raise RuntimeError(
+            'The CUDA context of this Smartg object was destroyed by the '
+            'clear_context of another Smartg object that shared it: '
+            'create a new Smartg object.')
+    # PyCUDA only knows the contexts it pushed itself: ask the driver
+    if _driver_current_context() == ctx.handle:
+        return False
+    ctx.push()
+    return True
+
+
+def _autoinit_current_at_exit(sg: 'Smartg') -> None:
+    """At exit, make the pycuda.autoinit context of sg current again.
+
+    The exit handler of pycuda.autoinit pops its context, which fails
+    ('cannot pop non-current context') if another library (JAX) made
+    its own context current in the meantime.
+    """
+    import sys
+    mod = sys.modules.get('pycuda.autoinit')
+    if (sg.ctx is None or _released(sg.ctx) or mod is None
+            or getattr(mod, 'context', None) is not sg.ctx):
+        return
+    # set it current at the driver level: pushing it would leave it
+    # twice on the stack of PyCUDA, which aborts at the module cleanup
+    if _driver_current_context() == sg.ctx.handle:
+        return
+    try:
+        import ctypes
+        for name in ('libcuda.so.1', 'libcuda.so', 'nvcuda.dll'):
+            try:
+                set_current = ctypes.CDLL(name).cuCtxSetCurrent
+                break
+            except (OSError, AttributeError):
+                continue
+        else:
+            return
+        set_current(ctypes.c_void_p(sg.ctx.handle))
+    # an exit handler must not raise: the one of pycuda.autoinit, which
+    # runs next, reports a context that is still not current
+    except Exception:  # noqa: BLE001, S110
+        pass
+
+
+def _driver_current_context() -> int | None:
+    """Handle of the current CUDA context of the driver, or None."""
+    global _CU_CTX_GET_CURRENT
+    if _CU_CTX_GET_CURRENT is None:
+        import ctypes
+        import ctypes.util
+        names = ('libcuda.so.1', 'libcuda.so', 'nvcuda.dll',
+                 ctypes.util.find_library('cuda'))
+        for name in names:
+            if not name:
+                continue
+            try:
+                _CU_CTX_GET_CURRENT = ctypes.CDLL(name).cuCtxGetCurrent
+                break
+            except (OSError, AttributeError):
+                continue
+        else:
+            _CU_CTX_GET_CURRENT = False
+    if not _CU_CTX_GET_CURRENT:
+        return None
+    import ctypes
+    handle = ctypes.c_void_p()
+    if _CU_CTX_GET_CURRENT(ctypes.byref(handle)) != 0:
+        return None
+    return handle.value
+
+
+_CU_CTX_GET_CURRENT = None
+
+
+def _in_context(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Run a method of Smartg with its CUDA context current."""
+
+    @functools.wraps(method)
+    def wrapper(self: 'Smartg', *args: Any, **kwargs: Any) -> Any:
+        pushed = _make_current(self.ctx)
+        try:
+            return method(self, *args, **kwargs)
+        finally:
+            # run may have popped and released the context itself
+            # (autoinit=False, keep_context=False)
+            if pushed and self.ctx is not None:
+                cuda.Context.pop()  # pyright: ignore[reportAttributeAccessIssue]
+
+    return wrapper
+
+
 class Smartg:
     """Initialize the Smartg object.
 
@@ -872,7 +993,7 @@ class Smartg:
     double : bool, optional
         Accumulate photons table in double precision (default double).
     alis : bool, optional
-        Use the ALIS method (Emde et al. 2010) for treating gaseous
+        Use the ALIS method (Emde et al. 2011) for treating gaseous
         absorption and perturbed profile. Needs alt_pp=True or
         pp=False. The 3D objects are not available with it.
     back : bool, optional
@@ -906,7 +1027,8 @@ class Smartg:
         Include the Sun Induced Fluorescence, an experimental option.
         Needs alis=True, together with alt_pp=True or pp=False.
     thermal : bool, optional
-        Still in dev...
+        Still in dev... The forward mode (back=False) needs alt_pp=True
+        in the plane parallel geometry.
     rng : str, optional
         The pseudo-random number generator to use, only 2 choice:
         - PHILOX
@@ -1027,6 +1149,13 @@ class Smartg:
                 "geometry (pp=True): the spherical move does not follow "
                 "its cells"
             )
+        if thermal and not back and pp and not alt_pp:
+            raise ValueError(
+                "The forward thermal emission (thermal=True, back=False) "
+                "needs alt_pp=True in the plane parallel geometry: the "
+                "fast move does not start the emitted photons at their "
+                "altitude"
+            )
 
         if device is not None:
             env_modif = {'CUDA_DEVICE': str(device)}
@@ -1091,6 +1220,10 @@ class Smartg:
                 self.ctx = make_default_context()
 
         self.autoinit = autoinit
+        if autoinit:
+            # registered after the exit handler of pycuda.autoinit, so
+            # run before it: its context must be current to be popped
+            atexit.register(_autoinit_current_at_exit, self)
         self.pp = pp
         self.double = double
         self.alis = alis
@@ -1197,6 +1330,9 @@ class Smartg:
             # 3D Object mode
             options.append('-DOBJ3D')
         options.append('-D' + rng)
+        # CUDA 13 deprecates double4 (warning 1444) for double4_16a and
+        # double4_32a, which older toolkits lack
+        options.append('-diag-suppress=1444')
         # options.append('-lineinfo')
 
         #
@@ -1208,23 +1344,29 @@ class Smartg:
         with open(SRC_DEVICE, encoding='ascii', errors='ignore') as f:
             src_device_content = f.read()
 
-        # kernel compilation
-        self.mod = SourceModule(
-            src_device_content,
-            nvcc='nvcc',
-            options=options,
-            no_extern_c=True,
-            cache_dir=str(cache_dir),
-            include_dirs=[
-                str(DIR_SRC),
-                str(DIR_SRC / 'incRNGs' / 'Random123'),
-            ],
-        )
+        # kernel compilation, in the context of this object (another
+        # library may have made its own context current)
+        pushed = _make_current(self.ctx)
+        try:
+            self.mod = SourceModule(
+                src_device_content,
+                nvcc='nvcc',
+                options=options,
+                no_extern_c=True,
+                cache_dir=str(cache_dir),
+                include_dirs=[
+                    str(DIR_SRC),
+                    str(DIR_SRC / 'incRNGs' / 'Random123'),
+                ],
+            )
 
-        # load the kernel
-        self.kernel = self.mod.get_function('launchKernel')
-        # self.kernel2 = self.mod.get_function('launchKernel2')
-        self.kernel2 = self.mod.get_function('reduce_absorption_gpu')
+            # load the kernel
+            self.kernel = self.mod.get_function('launchKernel')
+            # self.kernel2 = self.mod.get_function('launchKernel2')
+            self.kernel2 = self.mod.get_function('reduce_absorption_gpu')
+        finally:
+            if pushed:
+                cuda.Context.pop()  # pyright: ignore[reportAttributeAccessIssue]
 
         #
         # common attributes
@@ -1280,25 +1422,34 @@ class Smartg:
         Notes
         -----
         Once this method has been called, the run method can no longer
-        be used: the Smartg object must be reinitialized.
+        be used: the Smartg object must be reinitialized. The Smartg
+        objects with autoinit share one context: once one of them has
+        cleared it, the others can no longer run, and their own
+        clear_context only prints that there is no context to clear.
         """
         try:
             assert self.ctx is not None
+            # another library (JAX) may have made its context current
+            _make_current(self.ctx)
             self.ctx.pop()
             self.ctx.detach()
+            _RELEASED_CONTEXTS.append(self.ctx)
             self.ctx = None
             from pycuda.tools import clear_context_caches
 
             clear_context_caches()
             if self.autoinit:
-                # In case of autoinit delete pycuda.autoinit
+                # With autoinit, delete pycuda.autoinit and its exit
+                # handler, which would pop the destroyed context
                 import pycuda.autoinit
 
+                atexit.unregister(pycuda.autoinit._finish_up)
                 del pycuda.autoinit
         # pycuda's own error classes again, nothing else to do here
         except Exception:  # noqa: BLE001
             print("There is no current context to clear.")
 
+    @_in_context
     def run(
         self,
         wavelength: NumericArrayLike | BandSet | Sequence[BandLike],

{smartg-2.0.2 → smartg-2.0.4}/smartg/src/communs.h

@@ -286,10 +286,6 @@ public:
     #endif
     /* ---------------------*/
 
-    #ifdef SPHERIQUE
-    float radius;
-    #endif
-
     #ifdef BACK
     // Cumulative Mueller Matrix
     float4x4 M;