From 5741c522547756ac4bb7a16df32106a15efb8a57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cosmin=20=C8=98tefan=20Ciocan?=
 <57830279+cosminc98@users.noreply.github.com>
Date: Wed, 3 Apr 2024 17:23:33 +0200
Subject: [PATCH] Setup Kaggle environment on extension load (#31)

* Automatically setup kaggle environment https://github.com/andreinechaev/nvcc4jupyter/issues/29

* Update package lists before installing cuda toolkit
---
 .gitignore                |  3 --
 nvcc4jupyter/__init__.py  |  2 +-
 nvcc4jupyter/plugin.py    |  2 ++
 nvcc4jupyter/setup_env.py | 72 +++++++++++++++++++++++++++++++++++++++
 pyproject.toml            |  4 +++
 5 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 nvcc4jupyter/setup_env.py

diff --git a/.gitignore b/.gitignore
index 3485fef..f253064 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,9 +28,6 @@ pip-delete-this-directory.txt
 nosetests.xml
 coverage.xml
 
-# Virtual Environment
-*env*
-
 # Misc
 .pytest_cache/
 .DS_Store
diff --git a/nvcc4jupyter/__init__.py b/nvcc4jupyter/__init__.py
index 5bf6155..4fa704e 100644
--- a/nvcc4jupyter/__init__.py
+++ b/nvcc4jupyter/__init__.py
@@ -5,4 +5,4 @@ nvcc4jupyter: CUDA C++ plugin for Jupyter Notebook
 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"
diff --git a/nvcc4jupyter/plugin.py b/nvcc4jupyter/plugin.py
index 612d321..6de1054 100644
--- a/nvcc4jupyter/plugin.py
+++ b/nvcc4jupyter/plugin.py
@@ -23,6 +23,7 @@ from .parsers import (
     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)
diff --git a/nvcc4jupyter/setup_env.py b/nvcc4jupyter/setup_env.py
new file mode 100644
index 0000000..ea19993
--- /dev/null
+++ b/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()}"
+        )
diff --git a/pyproject.toml b/pyproject.toml
index e70fe45..2d6ada1 100644
--- a/pyproject.toml
+++ b/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"]