guberti commented on code in PR #13627:
URL: https://github.com/apache/tvm/pull/13627#discussion_r1058647242
##########
docs/conf.py:
##########
@@ -84,6 +86,148 @@ def git_describe_version(original_version):
version = git_describe_version(tvm.__version__)
release = version
+
+def monkey_patch(module_name, func_name):
+ """Helper function for monkey-patching library functions.
+
+ Used to modify a few sphinx-gallery behaviors to make the "Open in Colab"
+ button work correctly. Should be used as a decorator with arguments.
+ """
+ module = import_module(module_name)
+ original_func = getattr(module, func_name)
+
+ def decorator(function):
+ updated_func = partial(function, real_func=original_func)
+ setattr(module, func_name, updated_func)
+ return updated_func
+
+ return decorator
+
+
+# This header replaces the default sphinx-gallery one in
sphinx_gallery/gen_rst.py.
+COLAB_HTML_HEADER = """
+.. DO NOT EDIT. THIS FILE WAS AUTOMATICALLY GENERATED BY
+.. TVM'S MONKEY-PATCHED VERSION OF SPHINX-GALLERY. TO MAKE
+.. CHANGES, EDIT THE SOURCE PYTHON FILE:
+.. "{0}"
+
+.. only:: html
+
+ .. note::
+ :class: sphx-glr-download-link-note
+
+ This tutorial can be used interactively with Google Colab! You can
also click
+ :ref:`here <sphx_glr_download_{1}>` to run the Jupyter notebook
locally.
+
+ .. image::
https://raw.githubusercontent.com/apache/web-data/main/images/utilities/colab_button.svg
+ :align: center
+ :target: {2}
+ :width: 300px
+
+.. rst-class:: sphx-glr-example-title
+
+.. _sphx_glr_{1}:
+
+"""
+
+# Google Colab allows opening .ipynb files on GitHub by appending the GitHub
path to its own url.
+COLAB_URL_BASE = "https://colab.research.google.com/github";
+IPYTHON_GITHUB_BASE = "apache/tvm-site/blob/asf-site/docs/_downloads/"
+
+
+@monkey_patch("sphinx_gallery.gen_rst", "save_rst_example")
+def save_rst_example(example_rst, example_file, time_elapsed, memory_used,
gallery_conf, real_func):
+ """Monkey-patch save_rst_example to include the "Open in Colab" button."""
+
+ # The url is the md5 hash of the notebook path.
+ example_fname = os.path.relpath(example_file, gallery_conf["src_dir"])
+ ref_fname = example_fname.replace(os.path.sep, "_")
+ notebook_path = example_fname[:-2] + "ipynb"
+ digest = md5(notebook_path.encode()).hexdigest()
+
+ # Fixed documentation versions must link to different (earlier) .ipynb
notebooks.
+ colab_url = f"{COLAB_URL_BASE}/{IPYTHON_GITHUB_BASE}"
+ if "dev" not in version:
+ colab_url += version + "/"
+ colab_url += digest + "/" + os.path.basename(notebook_path)
+
+ new_header = COLAB_HTML_HEADER.format(example_fname, ref_fname, colab_url)
+ with patch("sphinx_gallery.gen_rst.EXAMPLE_HEADER", new_header):
+ real_func(example_rst, example_file, time_elapsed, memory_used,
gallery_conf)
+
+
+@monkey_patch("sphinx_gallery.notebook", "rst2md")
+def rst2md(text, gallery_conf, target_dir, heading_levels, real_func):
+ """Monkey-patch rst2md to add limited include directive support to
sphinx-gallery."""
+
+ include_re = re.compile(r"^([ \t]*)\.\. include::\s*(.+)\n", flags=re.M)
+
+ def load_include(match):
+ full_path = os.path.join(target_dir, match.group(2))
+ with open(full_path) as f:
+ lines = f.read()
+ indented = textwrap.indent(lines, match.group(1)) + "\n"
+ return indented
+
+ text = re.sub(include_re, load_include, text)
+
+ return real_func(text, gallery_conf, target_dir, heading_levels)
+
+
+INSTALL_TVM_DEV = f"""%%shell
+# Installs the latest dev build of TVM from PyPI. If you wish to build
+# from source, see https://tvm.apache.org/docs/install/from_source.html
+pip install apache-tvm --pre"""
+
+INSTALL_TVM_FIXED = f"""%%shell
+# Installs TVM version {version} from PyPI. If you wish to build
+# from source, see https://tvm.apache.org/docs/install/from_source.html
+pip install apache-tvm=={version}"""
+
+INSTALL_TVM_CUDA_DEV = f"""%%shell
+# Installs the latest dev build of TVM from PyPI, with CUDA enabled. To use
this,
+# you must request a Google Colab instance with a GPU by going to Runtime ->
+# Change runtime type -> Hardware accelerator -> GPU. If you wish to build from
+# source, see see https://tvm.apache.org/docs/install/from_source.html
+pip install tlcpack-nightly-cu113 --pre -f https://tlcpack.ai/wheels""";
+
+INSTALL_TVM_CUDA_FIXED = f"""%%shell
+# Installs TVM version {version} from PyPI, with CUDA enabled. To use this,
+# you must request a Google Colab instance with a GPU by going to Runtime ->
+# Change runtime type -> Hardware accelerator -> GPU. If you wish to build from
+# source, see see https://tvm.apache.org/docs/install/from_source.html
+pip install apache-tvm-cu113=={version} -f https://tlcpack.ai/wheels""";
+
+
+@monkey_patch("sphinx_gallery.gen_rst", "jupyter_notebook")
+def jupyter_notebook(script_blocks, gallery_conf, target_dir, real_func):
+ """Monkey-patch sphinx-gallery to add a TVM import block to each IPython
notebook.
+
+ If we had only one import block, we could skip the patching and just set
first_notebook_cell.
+ However, how we import TVM depends on if we are using a fixed or dev
version, and whether we
+ will use the GPU.
+
+ We install a CUDA version of TVM if and only if the tutorial calls "cuda"
as a function. It
+ would be cleaner specify this with a config flag, but sphinx-gallery does
not support this.
+ """
+
+ call_cuda_re = re.compile(r"cuda\([0-9]*\)", re.IGNORECASE)
Review Comment:
I agree this regex approach is gross. I've added a new sphinx-gallery flag
to determine whether a tutorial requires CUDA, and I've set it in all relevant
tutorials:
```python
# sphinx_gallery_requires_cuda = True
```
Unfortunately, monkey-patching the extra flag is a bit of a hack, but I
still think it is nicer than the regex approach.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
