Fix IDAKLU solver crash when running in multiple threads + when telemetry is prompted for

CHANGELOG.md

@@ -7,7 +7,7 @@
 - Adds an option "voltage as a state" that can be "false" (default) or "true". If "true" adds an explicit algebraic equation for the voltage. ([#4507](https://github.com/pybamm-team/PyBaMM/pull/4507))
 - Improved `QuickPlot` accuracy for simulations with Hermite interpolation. ([#4483](https://github.com/pybamm-team/PyBaMM/pull/4483))
 - Added Hermite interpolation to the (`IDAKLUSolver`) that improves the accuracy and performance of post-processing variables. ([#4464](https://github.com/pybamm-team/PyBaMM/pull/4464))
-- Added basic telemetry to record which functions are being run. See [Telemetry section in the User Guide](https://docs.pybamm.org/en/latest/source/user_guide/index.html#telemetry) for more information. ([#4441](https://github.com/pybamm-team/PyBaMM/pull/4441))
+- Added basic telemetry to record which functions are being run. See [Telemetry section in the User Guide](https://docs.pybamm.org/en/latest/source/user_guide/index.html#telemetry) for more information. ([#4441](https://github.com/pybamm-team/PyBaMM/pull/4441), [#4583](https://github.com/pybamm-team/PyBaMM/pull/4583))
 - Added `BasicDFN` model for sodium-ion batteries ([#4451](https://github.com/pybamm-team/PyBaMM/pull/4451))
 - Added sensitivity calculation support for `pybamm.Simulation` and `pybamm.Experiment` ([#4415](https://github.com/pybamm-team/PyBaMM/pull/4415))
 - Added OpenMP parallelization to IDAKLU solver for lists of input parameters ([#4449](https://github.com/pybamm-team/PyBaMM/pull/4449))

docs/source/user_guide/installation/index.rst

@@ -72,6 +72,7 @@ Package                                                             Minimum supp
 `pandas <https://pypi.org/project/pandas/>`__                       1.5.0
 `pooch <https://www.fatiando.org/pooch/>`__                         1.8.1
 `posthog <https://posthog.com/>`__                                  3.6.5
+`platformdirs <https://pypi.org/project/platformdirs/>`__           4.0.0
 =================================================================== ==========================
 
 .. _install.optional_dependencies:

noxfile.py

@@ -253,9 +253,11 @@ def run_tests(session):
     set_environment_variables(PYBAMM_ENV, session=session)
     session.install("setuptools", silent=False)
     session.install("-e", ".[all,dev,jax]", silent=False)
-    specific_test_files = session.posargs if session.posargs else []
     session.run(
-        "python", "-m", "pytest", *specific_test_files, "-m", "unit or integration"
+        "python",
+        "-m",
+        "pytest",
+        *(session.posargs if session.posargs else ["-m", "unit or integration"]),
     )
 
 

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
     "pandas>=1.5.0",
     "pooch>=1.8.1",
     "posthog",
+    "platformdirs",
 ]
 
 [project.urls]

src/pybamm/__init__.py

@@ -4,10 +4,10 @@
 demote_expressions_to_32bit = False
 
 # Utility classes and methods
-from .util import root_dir
 from .util import Timer, TimerTime, FuzzyDict
 from .util import (
     root_dir,
+    is_notebook,
     load,
     is_constant_and_can_evaluate,
 )
@@ -190,7 +190,7 @@
 from .plotting.dynamic_plot import dynamic_plot
 
 # Simulation
-from .simulation import Simulation, load_sim, is_notebook
+from .simulation import Simulation, load_sim
 
 # Batch Study
 from .batch_study import BatchStudy

src/pybamm/config.py

@@ -1,11 +1,14 @@
-import uuid
 import os
-import platformdirs
-from pathlib import Path
-import pybamm
 import sys
-import threading
 import time
+import uuid
+
+from pathlib import Path
+
+import pybamm
+import platformdirs
+
+from pybamm.util import is_notebook
 
 
 def is_running_tests():  # pragma: no cover
@@ -45,20 +48,120 @@ def is_running_tests():  # pragma: no cover
     return False
 
 
-def ask_user_opt_in(timeout=10):
+def get_input_or_timeout(timeout):  # pragma: no cover
     """
-    Ask the user if they want to opt in to telemetry.
+    Cross-platform input with timeout, using various methods depending on the
+    environment. Works in Jupyter notebooks, Windows, and Unix-like systems.
 
-    Parameters
-    ----------
-    timeout : float, optional
-        The timeout for the user to respond to the prompt. Default is 10 seconds.
+    Args:
+        timeout (float): Timeout in seconds
 
-    Returns
-    -------
-    bool
-        True if the user opts in, False otherwise.
+    Returns:
+        tuple: A tuple containing:
+            - str: The user input if received before the timeout, or None if the timeout was reached.
+            - bool: True if the timeout was reached, False otherwise.
     """
+    # Check for telemetry disable flag
+    if os.getenv("PYBAMM_DISABLE_TELEMETRY", "false").lower() != "false":
+        return None, True
+
+    if not (sys.stdin.isatty() or is_notebook()):
+        return None, True
+
+    # 1. Handling for Jupyter notebooks. This is a simplified
+    # implementation in comparison to widgets because notebooks
+    # are usually run interactively, and we don't need to worry
+    # about the event loop. The input is disabled when running
+    # tests due to the environment variable set.
+    if is_notebook():  # pragma: no cover
+        try:
+            from IPython.display import clear_output
+
+            user_input = input("Do you want to enable telemetry? (Y/n): ")
+            clear_output()
+
+            return user_input, False
+
+        except Exception:  # pragma: no cover
+            return None, True
+
+    # 2. Windows-specific handling
+    if sys.platform == "win32":
+        try:
+            import msvcrt
+
+            start_time = time.time()
+            input_chars = []
+            sys.stdout.write("Do you want to enable telemetry? (Y/n): ")
+            sys.stdout.flush()
+
+            while time.time() - start_time < timeout:
+                if msvcrt.kbhit():
+                    char = msvcrt.getwche()
+                    if char in ("\r", "\n"):
+                        sys.stdout.write("\n")
+                        return "".join(input_chars), False
+                    input_chars.append(char)
+                time.sleep(0.1)
+            return None, True
+        except Exception:
+            return None, True
+
+    # 3. POSIX-like systems will need to use termios
+    else:  # pragma: no cover
+        try:
+            import termios
+            import tty
+            import select
+
+            # Save terminal settings for later
+            old_settings = termios.tcgetattr(sys.stdin)
+            try:
+                # Set terminal to raw mode
+                tty.setraw(sys.stdin.fileno())
+
+                sys.stdout.write("Do you want to enable telemetry? (Y/n): ")
+                sys.stdout.flush()
+
+                input_chars = []
+                start_time = time.time()
+
+                while time.time() - start_time < timeout:
+                    rlist, _, _ = select.select([sys.stdin], [], [], 0.1)
+                    if rlist:
+                        char = sys.stdin.read(1)
+                        if char in ("\r", "\n"):
+                            sys.stdout.write("\n")
+                            return "".join(input_chars), False
+                        input_chars.append(char)
+                        sys.stdout.write(char)
+                        sys.stdout.flush()
+                return None, True
+
+            finally:
+                # Restore saved terminal settings
+                termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
+                sys.stdout.write("\n")
+                sys.stdout.flush()
+
+        except Exception:  # pragma: no cover
+            return None, True
+
+    return None, True
+
+
+def ask_user_opt_in(timeout=10):  # pragma: no cover
+    """
+    Ask the user if they want to opt in to telemetry.
+    """
+    # Check for telemetry disable flag first
+    if os.getenv("PYBAMM_DISABLE_TELEMETRY", "false").lower() != "false":
+        return False
+
+    # Skip telemetry prompt in non-interactive environments
+    if not (sys.stdin.isatty() or is_notebook()):
+        return False
+
     print(
         "PyBaMM can collect usage data and send it to the PyBaMM team to "
         "help us improve the software.\n"
@@ -69,44 +172,24 @@ def ask_user_opt_in(timeout=10):
         "For more information, see https://docs.pybamm.org/en/latest/source/user_guide/index.html#telemetry"
     )
 
-    def get_input():  # pragma: no cover
-        try:
-            user_input = (
-                input("Do you want to enable telemetry? (Y/n): ").strip().lower()
-            )
-            answer.append(user_input)
-        except Exception:
-            # Handle any input errors
-            pass
-
-    time_start = time.time()
-
-    while True:
-        if time.time() - time_start > timeout:
-            print("\nTimeout reached. Defaulting to not enabling telemetry.")
-            return False
-
-        answer = []
-        # Create and start input thread
-        input_thread = threading.Thread(target=get_input)
-        input_thread.daemon = True
-        input_thread.start()
-
-        # Wait for either timeout or input
-        input_thread.join(timeout)
-
-        if answer:
-            if answer[0] in ["yes", "y", ""]:
-                print("\nTelemetry enabled.\n")
-                return True
-            elif answer[0] in ["no", "n"]:
-                print("\nTelemetry disabled.\n")
-                return False
-            else:
-                print("\nInvalid input. Please enter 'yes/y' for yes or 'no/n' for no.")
-        else:
-            print("\nTimeout reached. Defaulting to not enabling telemetry.")
-            return False
+    user_input, timed_out = get_input_or_timeout(timeout)
+
+    if timed_out:
+        print("\nTimeout reached. Defaulting to not enabling telemetry.")
+        return False
+
+    if user_input is None or not user_input:  # Empty input should mean a yes
+        print("\nTelemetry enabled.\n")
+        return True
+    elif user_input.lower() in ["y", "yes"]:
+        print("\nTelemetry enabled.\n")
+        return True
+    elif user_input.lower() in ["n", "no"]:
+        print("\nTelemetry disabled.\n")
+        return False
+    else:
+        print("\nInvalid input. Defaulting to not enabling telemetry.")
+        return False
 
 
 def generate():
@@ -118,7 +201,7 @@ def generate():
         return
 
     # Ask the user if they want to opt in to telemetry
-    opt_in = ask_user_opt_in()
+    opt_in = ask_user_opt_in(timeout=10)
     config_file = Path(platformdirs.user_config_dir("pybamm")) / "config.yml"
     write_uuid_to_file(config_file, opt_in)
 

src/pybamm/simulation.py

@@ -16,24 +16,6 @@
 from pybamm.expression_tree.operations.serialise import Serialise
 
 
-def is_notebook():
-    try:
-        shell = get_ipython().__class__.__name__
-        if shell == "ZMQInteractiveShell":  # pragma: no cover
-            # Jupyter notebook or qtconsole
-            cfg = get_ipython().config
-            nb = len(cfg["InteractiveShell"].keys()) == 0
-            return nb
-        elif shell == "TerminalInteractiveShell":  # pragma: no cover
-            return False  # Terminal running IPython
-        elif shell == "Shell":  # pragma: no cover
-            return True  # Google Colab notebook
-        else:  # pragma: no cover
-            return False  # Other type (?)
-    except NameError:
-        return False  # Probably standard Python interpreter
-
-
 class Simulation:
     """A Simulation class for easy building and running of PyBaMM simulations.
 
@@ -141,7 +123,7 @@ def __init__(
         self._set_random_seed()
 
         # ignore runtime warnings in notebooks
-        if is_notebook():  # pragma: no cover
+        if pybamm.is_notebook():  # pragma: no cover
             import warnings
 
             warnings.filterwarnings("ignore")

src/pybamm/util.py

@@ -22,6 +22,24 @@ def root_dir():
     return str(pathlib.Path(pybamm.__path__[0]).parent.parent)
 
 
+def is_notebook():
+    try:
+        shell = get_ipython().__class__.__name__
+        if shell == "ZMQInteractiveShell":  # pragma: no cover
+            # Jupyter notebook or qtconsole
+            cfg = get_ipython().config
+            nb = len(cfg["InteractiveShell"].keys()) == 0
+            return nb
+        elif shell == "TerminalInteractiveShell":  # pragma: no cover
+            return False  # Terminal running IPython
+        elif shell == "Shell":  # pragma: no cover
+            return True  # Google Colab notebook
+        else:  # pragma: no cover
+            return False  # Other type (?)
+    except NameError:
+        return False  # Probably standard Python interpreter
+
+
 def get_git_commit_info():
     """
     Get the git commit info for the current PyBaMM version, e.g. v22.8-39-gb25ce8c41