Setup Kaggle environment on extension load (#31)

* Automatically setup kaggle environment #29

* Update package lists before installing cuda toolkit

.gitignore

@@ -28,9 +28,6 @@ pip-delete-this-directory.txt
 nosetests.xml
 coverage.xml
 
-# Virtual Environment
-*env*
-
 # Misc
 .pytest_cache/
 .DS_Store

nvcc4jupyter/__init__.py

@@ -5,4 +5,4 @@
 from .parsers import Profiler, set_defaults  # noqa: F401
 from .plugin import NVCCPlugin, load_ipython_extension  # noqa: F401
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"

nvcc4jupyter/plugin.py

@@ -23,6 +23,7 @@
     get_parser_cuda_group_save,
 )
 from .path_utils import CUDA_SEARCH_PATHS, find_executable
+from .setup_env import setup_environment
 
 DEFAULT_EXEC_FNAME = "cuda_exec.out"
 SHARED_GROUP_NAME = "shared"
@@ -364,5 +365,6 @@ def load_ipython_extension(shell: InteractiveShell):
     """
     Method used by IPython to load the extension.
     """
+    setup_environment()
     nvcc_plugin = NVCCPlugin(shell)
     shell.register_magics(nvcc_plugin)

nvcc4jupyter/setup_env.py

@@ -0,0 +1,72 @@
+"""
+Setup steps for platforms such as Kaggle, Colab, etc. to allow our extension
+to work on them immediately after loading it.
+"""
+
+# pylint: disable=missing-function-docstring
+
+import os
+import traceback
+from subprocess import DEVNULL, STDOUT, check_call
+from typing import Optional
+
+PATH_PRIORITY_DIR = "/usr/bin/priority"
+KAGGLE_GCC_8_PATH = "/usr/bin/gcc-8"
+
+
+def print_platform(platform: str) -> None:
+    print(f'Detected platform "{platform}". Running its setup...')
+
+
+def kaggle_setup() -> None:
+    print("Updating the package lists...")
+    check_call(["/usr/bin/apt-get", "update"], stdout=DEVNULL, stderr=STDOUT)
+
+    print("Installing nvidia-cuda-toolkit, this may take a few minutes...")
+    check_call(
+        ["/usr/bin/apt-get", "install", "-y", "nvidia-cuda-toolkit"],
+        stdout=DEVNULL,
+        stderr=STDOUT,
+    )
+    os.makedirs(PATH_PRIORITY_DIR, exist_ok=True)
+
+    gcc_symlink_path = os.path.join(PATH_PRIORITY_DIR, "gcc")
+    if not os.path.exists(gcc_symlink_path):
+        os.symlink(KAGGLE_GCC_8_PATH, gcc_symlink_path)
+
+    if PATH_PRIORITY_DIR not in os.environ["PATH"].split(":"):
+        os.environ["PATH"] = f"{PATH_PRIORITY_DIR}:" + os.environ["PATH"]
+
+
+def colab_setup() -> None:
+    pass
+
+
+def setup_environment() -> None:
+    """
+    Detect the platform the extension was loaded on and run the necessary
+    steps (install dependencies, add executables to PATH, etc.) for the
+    extension to work.
+    """
+
+    if "NVCC4JUPYTER_NO_SETUP" in os.environ:
+        return
+
+    platform: Optional[str] = None
+    try:
+        if "KAGGLE_URL_BASE" in os.environ:
+            platform = "Kaggle"
+            print_platform(platform)
+            kaggle_setup()
+        elif "COLAB_RELEASE_TAG" in os.environ:
+            platform = "Colab"
+            print_platform(platform)
+            colab_setup()
+    except Exception:  # pylint: disable=broad-exception-caught
+        print(
+            f'Setup failed for detected platform "{platform}". Set the'
+            ' "NVCC4JUPYTER_NO_SETUP" environment variable to disable running'
+            " the setup on load. Please report the following error to"
+            " https://github.com/andreinechaev/nvcc4jupyter/issues:"
+            f" following error message:\n{traceback.format_exc()}"
+        )

pyproject.toml

@@ -88,6 +88,10 @@ enable-unstable-feature = ["string_processing"]
 
 [tool.coverage.run]
 branch = true
+omit = [
+    # cannot test installing dependencies on platforms such as kaggle
+    "nvcc4jupyter/setup_env.py",
+]
 
 [tool.pyright]
 include = ["src"]