Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/strategies/launchers/multiprocessing.py (Python)

Code before:
            use_deterministic_algorithms_warn_only=torch.is_deterministic_algorithms_warn_only_enabled(),
            cudnn_benchmark=torch.backends.cudnn.benchmark,
            rng_states=_collect_rng_states(),
        )

    def restore(self) -> None:
        """Restores all globals to the values captured in the :meth:`capture` method."""
        torch.use_deterministic_algorithms(
            self.use_deterministic_algorithms, warn_only=self.use_deterministic_algorithms_warn_only
        )
        torch.backends.cudnn.benchmark = self.cudnn_benchmark
        _set_rng_states(self.rng_states)


def _check_bad_cuda_fork() -> None:
    """Checks whether it is safe to fork and initialize CUDA in the new processes, and raises an exception if not.

    The error message replaces PyTorch's 'Cannot re-initialize CUDA in forked subprocess' with helpful advice for
    Lightning users.

    """
    # Use PyTorch's internal check for bad fork state, which is more accurate than just checking if CUDA
    # is initialized. This allows passive CUDA initialization (e.g., from library imports or device queries)
    # while still catching actual problematic cases where CUDA context was created before forking.
    _is_in_bad_fork = getattr(torch.cuda, "_is_in_bad_fork", None)
    if _is_in_bad_fork is not None and _is_in_bad_fork():
        message = (
            "Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, "
            "you must use the 'spawn' start method or avoid CUDA initialization in the main process."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)

    # Fallback to the old check if _is_in_bad_fork is not available (older PyTorch versions)
    if _is_in_bad_fork is None and torch.cuda.is_initialized():
        message = (
            "Lightning can't create new processes if CUDA is already initialized. Did you manually call"
            " `torch.cuda.*` functions, have moved the model to the device, or allocated memory on the GPU any"
            " other way? Please remove any such calls, or change the selected strategy."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)


def _disable_module_memory_sharing(data: Any) -> Any:
    """Disables memory sharing on parameters and buffers of `nn.Module`s contained in the given collection.

    Note: This is only required when running on CPU.


Review comment (suggestion):
The code uses `getattr` to check if `_is_in_bad_fork` exists, but it doesn't verify if the returned value is callable. While PyTorch's `_is_in_bad_fork` is indeed a function, it's better practice to verify callability when using `getattr` on potentially undefined attributes, especially for internal/private APIs that could change.

Consider adding a callable check:
```python
_is_in_bad_fork = getattr(torch.cuda, "_is_in_bad_fork", None)
if _is_in_bad_fork is not None and callable(_is_in_bad_fork) and _is_in_bad_fork():
```
```suggestion
    if _is_in_bad_fork is not None and callable(_is_in_bad_fork) and _is_in_bad_fork():
```

Code after:
            use_deterministic_algorithms_warn_only=torch.is_deterministic_algorithms_warn_only_enabled(),
            cudnn_benchmark=torch.backends.cudnn.benchmark,
            rng_states=_collect_rng_states(),
        )

    def restore(self) -> None:
        """Restores all globals to the values captured in the :meth:`capture` method."""
        torch.use_deterministic_algorithms(
            self.use_deterministic_algorithms, warn_only=self.use_deterministic_algorithms_warn_only
        )
        torch.backends.cudnn.benchmark = self.cudnn_benchmark
        _set_rng_states(self.rng_states)


def _check_bad_cuda_fork() -> None:
    """Checks whether it is safe to fork and initialize CUDA in the new processes, and raises an exception if not.

    The error message replaces PyTorch's 'Cannot re-initialize CUDA in forked subprocess' with helpful advice for
    Lightning users.

    """
    # Use PyTorch's internal check for bad fork state, which is more accurate than just checking if CUDA
    # is initialized. This allows passive CUDA initialization (e.g., from library imports or device queries)
    # while still catching actual problematic cases where CUDA context was created before forking.
    _is_in_bad_fork = getattr(torch.cuda, "_is_in_bad_fork", None)
    if _is_in_bad_fork is not None and callable(_is_in_bad_fork) and _is_in_bad_fork():
        message = (
            "Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, "
            "you must use the 'spawn' start method or avoid CUDA initialization in the main process."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)

    # Fallback to the old check if _is_in_bad_fork is not available (older PyTorch versions)
    if _is_in_bad_fork is None and torch.cuda.is_initialized():
        message = (
            "Lightning can't create new processes if CUDA is already initialized. Did you manually call"
            " `torch.cuda.*` functions, have moved the model to the device, or allocated memory on the GPU any"
            " other way? Please remove any such calls, or change the selected strategy."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)


def _disable_module_memory_sharing(data: Any) -> Any:
    """Disables memory sharing on parameters and buffers of `nn.Module`s contained in the given collection.

    Note: This is only required when running on CPU.