From 8340aa00d1102db13b1661c4aedf9a942dab211a Mon Sep 17 00:00:00 2001 From: MB <michael.blaschek@univie.ac.at> Date: Fri, 17 Nov 2023 10:57:53 +0100 Subject: [PATCH] added a mojo recipe --- .gitignore | 2 + definition-files/Makefile | 6 +- definition-files/README.md | 31 ++ .../Singularity.ubuntu.20.04.mojo | 116 ++++++ definition-files/src/mandelbrot.mojo | 101 +++++ definition-files/src/mojokernel.py | 387 ++++++++++++++++++ 6 files changed, 641 insertions(+), 2 deletions(-) create mode 100644 definition-files/README.md create mode 100644 definition-files/Singularity.ubuntu.20.04.mojo create mode 100644 definition-files/src/mandelbrot.mojo create mode 100644 definition-files/src/mojokernel.py diff --git a/.gitignore b/.gitignore index 102af62..7dfb411 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ *.sif models/WRF/sandbox.wrf.dev models/WRF/run +*/.ipynb_checkpoints/* + diff --git a/definition-files/Makefile b/definition-files/Makefile index 405d7b0..13d3f26 100644 --- a/definition-files/Makefile +++ b/definition-files/Makefile @@ -1,8 +1,10 @@ SFILES := $(wildcard Singularity.*) +BUILD_FLAGS := ${BUILD_FLAGS} +APP_FLAGS := ${APP_FLAGS} %.sif:% - sudo singularity build $@ $< + sudo singularity $(APP_FLAGS) build $(BUILD_FLAGS) $@ $< $(addsuffix .sif, $(SFILES)): $(SFILES) @@ -11,4 +13,4 @@ list: @LC_ALL=C $(MAKE) -pRrq -f $(firstword $(MAKEFILE_LIST)) : 2>/dev/null | awk -v RS= -F: '/(^|\n)# Files(\n|$$)/,/(^|\n)# Finished Make data base/ {if ($$1 !~ "^[#.]") {print $$1}}' | sort | grep -E -v -e '^[^[:alnum:]]' -e '^$@$$' clean: - rm -f *.sif \ No newline at end of file + rm -f *.sif diff --git a/definition-files/README.md b/definition-files/README.md new file mode 100644 index 0000000..720f34b --- /dev/null +++ b/definition-files/README.md @@ -0,0 +1,31 @@ + +# Development Notes + +Building a apptainer + +```sh +# just run the build script +./build.sh [recipe] +# build an example +./build.sh Singularity.almalinux.base +# run the final container +./Singularity.almalinux.base.sif +``` + + +## run scripts + +```sh title='Example run script to source environment vars' + +%runscript + for script in /.singularity.d/env/*.sh; do + if [ -f "$script" ]; then + . "$script" + fi + done + if [ $# -eq 0 ]; then + exec /bin/bash --norc "$@" + fi + exec "$@" + +``` \ No newline at end of file diff --git a/definition-files/Singularity.ubuntu.20.04.mojo b/definition-files/Singularity.ubuntu.20.04.mojo new file mode 100644 index 0000000..8abda9f --- /dev/null +++ b/definition-files/Singularity.ubuntu.20.04.mojo @@ -0,0 +1,116 @@ +Bootstrap: docker +From: ubuntu:20.04 + +%labels + maintainer IT-IMGW <it.img-wien@univie.ac.at> + +%files + ./runscript /.singularity.d/runscript + ./run-help /.singularity.d/runscript.help + ./src/mojokernel.py /modular/pkg/packages.modular.com_mojo/jupyter/kernel/mojokernel2.py + +%post + export DEFAULT_TZ=Vienna/Europe + export DEBIAN_FRONTEND='noninteractive' + apt-get update \ + && apt-get install -y \ + tzdata \ + vim \ + curl \ + wget \ + g++ \ + make \ + file \ + git \ + python3-venv \ + apt-transport-https \ + libedit2 + rm -rf /var/lib/apt/lists/* + mkdir -p /modular + + export MODULAR_HOME="/modular" + export PATH="$MODULAR_HOME/pkg/packages.modular.com_mojo/bin:$PATH" + + curl https://get.modular.com | sh - + apt-get install -y modular + modular --help + modular auth $MOJO_AUTH + modular install mojo + + # Cleanup + apt-get -y autoremove --purge + apt-get -y clean + + # add a environment script for initalizing + cat > /.singularity.d/env/99-mojo.sh<<EOF +mkdir -p \$HOME/.modular +cp -u /modular/modular.cfg \$HOME/.modular/ +if [ ! -d \$HOME/.modular/pkg ]; then + ln -s /modular/pkg \$HOME/.modular/pkg +fi +EOF + # create a jupyter kernel that can be used. + mkdir -p /modular/pkg/packages.modular.com_mojo/venv/share/jupyter/kernels/mojo-jupyter-kernel + cp /root/.local/share/jupyter/kernels/mojo-jupyter-kernel/* /modular/pkg/packages.modular.com_mojo/venv/share/jupyter/kernels/mojo-jupyter-kernel/ + cat > /modular/pkg/packages.modular.com_mojo/venv/share/jupyter/kernels/mojo-jupyter-kernel/kernel.json <<EOF +{ + "display_name": "Mojo", + "argv": [ + "python3", + "/modular/pkg/packages.modular.com_mojo/jupyter/kernel/mojokernel2.py", + "-f", + "{connection_file}" + ], + "language": "mojo", + "codemirror_mode": "mojo", + "language_info": { + "name": "mojo", + "mimetype": "text/x-mojo", + "file_extension": ".mojo", + "codemirror_mode": { + "name": "mojo" + } + } +} +EOF + # fix environment for additional packages + sed -i 's/false/true/' /modular/pkg/packages.modular.com_mojo/venv/pyvenv.cfg + # command prompt name + CNAME=u20.mojo + # does not work goes into /.singularity.d/env/91-environment.sh + echo "export PS1=\"[IMGW-$CNAME]\w\$ \"" >> /.singularity.d/env/99-zz-custom-env.sh + # add some labels + echo "libc $(ldd --version | head -n1 | cut -d' ' -f4)" >> "$SINGULARITY_LABELS" + echo "linux $(cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2)" >> "$SINGULARITY_LABELS" + +%environment + export LC_ALL=C + export LANG=C.UTF-8 + export LIBRARY=/opt/conda/lib:/usr/lib64:/lib64:/lib + export INCLUDE=/opt/conda/include:/usr/include + export MODULAR_HOME="$HOME/.modular" + export PATH="$MODULAR_HOME/pkg/packages.modular.com_mojo/bin:$MODULAR_HOME/pkg/packages.modular.com_mojo/venv/bin:/opt/conda/bin:$PATH" + export SHELL=/bin/bash + + +# DOCKERFILE: +# FROM ubuntu:20.04 + +# ENV DEBIAN_FRONTEND='noninteractive' +# RUN apt-get update \ +# && apt-get install -y curl g++ make file python3-venv apt-transport-https libedit2\ +# && mkdir -p /modular +# ENV LC_ALL=C +# ENV LANG=C.UTF-8 +# ENV LIBRARY=/usr/lib64:/lib64:/lib +# ENV INCLUDE=/usr/include +# ENV MODULAR_HOME="/modular" +# ENV PATH="$MODULAR_HOME/pkg/packages.modular.com_mojo/bin:$PATH" + +# RUN curl https://get.modular.com | sh - \ +# && apt-get install -y modular +# RUN modular auth $MOJO_AUTH \ +# && modular install mojo +# # Cleanup +# RUN apt-get -y autoremove --purge\ +# && apt-get -y clean diff --git a/definition-files/src/mandelbrot.mojo b/definition-files/src/mandelbrot.mojo new file mode 100644 index 0000000..beaaa9d --- /dev/null +++ b/definition-files/src/mandelbrot.mojo @@ -0,0 +1,101 @@ +# ===----------------------------------------------------------------------=== # +# Copyright (c) 2023, Modular Inc. All rights reserved. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions: +# https://llvm.org/LICENSE.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ===----------------------------------------------------------------------=== # + +import benchmark +from complex import ComplexSIMD, ComplexFloat64 +from math import iota +from python import Python +from runtime.llcl import num_cores +from algorithm import parallelize, vectorize +from tensor import Tensor +from utils.index import Index +from python import Python + +alias float_type = DType.float64 +alias simd_width = 2 * simdwidthof[float_type]() + +alias width = 960 +alias height = 960 +alias MAX_ITERS = 200 + +alias min_x = -2.0 +alias max_x = 0.6 +alias min_y = -1.5 +alias max_y = 1.5 + + +fn mandelbrot_kernel_SIMD[ + simd_width: Int +](c: ComplexSIMD[float_type, simd_width]) -> SIMD[float_type, simd_width]: + """A vectorized implementation of the inner mandelbrot computation.""" + let cx = c.re + let cy = c.im + var x = SIMD[float_type, simd_width](0) + var y = SIMD[float_type, simd_width](0) + var y2 = SIMD[float_type, simd_width](0) + var iters = SIMD[float_type, simd_width](0) + + var t: SIMD[DType.bool, simd_width] = True + for i in range(MAX_ITERS): + if not t.reduce_or(): + break + y2 = y * y + y = x.fma(y + y, cy) + t = x.fma(x, y2) <= 4 + x = x.fma(x, cx - y2) + iters = t.select(iters + 1, iters) + return iters + + +fn main() raises: + let p = Python.import_module("numpy") + let b = Python.import_module("numpy") + let t = Tensor[float_type](height, width) + + @parameter + fn worker(row: Int): + let scale_x = (max_x - min_x) / width + let scale_y = (max_y - min_y) / height + + @parameter + fn compute_vector[simd_width: Int](col: Int): + """Each time we operate on a `simd_width` vector of pixels.""" + let cx = min_x + (col + iota[float_type, simd_width]()) * scale_x + let cy = min_y + row * scale_y + let c = ComplexSIMD[float_type, simd_width](cx, cy) + t.data().simd_store[simd_width]( + row * width + col, mandelbrot_kernel_SIMD[simd_width](c) + ) + + # Vectorize the call to compute_vector where call gets a chunk of pixels. + vectorize[simd_width, compute_vector](width) + + @parameter + fn bench[simd_width: Int](): + for row in range(height): + worker(row) + + let vectorized = benchmark.run[bench[simd_width]]().mean() + print("Number of threads:", num_cores()) + print("Vectorized:", vectorized, "s") + + # Parallelized + @parameter + fn bench_parallel[simd_width: Int](): + parallelize[worker](height, height) + + let parallelized = benchmark.run[bench_parallel[simd_width]]().mean() + print("Parallelized:", parallelized, "s") + print("Parallel speedup:", vectorized / parallelized) + + _ = t # Make sure tensor isn't destroyed before benchmark is finished diff --git a/definition-files/src/mojokernel.py b/definition-files/src/mojokernel.py new file mode 100644 index 0000000..3057ae8 --- /dev/null +++ b/definition-files/src/mojokernel.py @@ -0,0 +1,387 @@ +#!/usr/bin/python +# ===----------------------------------------------------------------------=== # +# +# This file is Modular Inc proprietary. +# +# ===----------------------------------------------------------------------=== # +# +# This file contains an implementation of a Jupyter kernel for Mojo. It +# communicates to Mojo using the MojoJupyter API library. +# +# ===----------------------------------------------------------------------=== # + + +import argparse +import ctypes +import json +import os +import shutil +import time +import traceback +from configparser import ConfigParser +from enum import IntEnum +from pathlib import Path +from typing import Any, Dict, Optional + +from ipykernel.kernelapp import IPKernelApp +from ipykernel.kernelbase import Kernel + +# ===----------------------------------------------------------------------=== # +# OutputProcessor +# ===----------------------------------------------------------------------=== # + +# Special start and end output markers that denote a display message. +display_start = "%%%%%%%DISPLAY_START" +display_end = "%%%%%%%DISPLAY_END" + + +class ExecutionFinishedState(IntEnum): + """ + Copied from MojoJupyter/Kernel.cpp - models the possible states of a kernel execution. + """ + + NotFinished = 0 + FinishedSuccess = 1 + FinishedError = 2 + + +class OutputProcessor: + """ + Process the output coming from stdout/stderr to find special markers that + indicate specific messages need to be sent. + """ + + def __init__(self, kernel: Kernel): + # The kernel object that we are processing output for. + self.kernel = kernel + + # The current contents of a pending display message. + self.pending_display_message = None + + def send_message(self, name: str, text: str) -> None: + """ + Send a stream message to the client. `name` should be stderr or stderr. + """ + stream_content = { + "name": name, + "text": text, + } + self.kernel.log.info(stream_content) + self.kernel.send_response( + self.kernel.iopub_socket, "stream", stream_content + ) + + # Create the output callback function. This is called by the MojoJupyter + # library to send output back to the Jupyter client. + def process_output(self, name: bytes, msg: bytes) -> None: + """Process output from the Mojo kernel.""" + msgstr = msg.decode() + + def send_stream(text: str) -> None: + self.send_message(name.decode(), text) + + # Short-circuit stderr right away - we don't report images or anything + # through stderr, and we should always report it immediately. + if name.decode() == "stderr": + send_stream(msgstr) + return + + # If we don't have a pending display message, and we don't have a + # display start marker, then just send the output as a normal stream message. + display_marker_loc = msgstr.find(display_start) + if self.pending_display_message is None and display_marker_loc == -1: + send_stream(msgstr) + return + + while len(msgstr) > 0: + # Buffer the display message, sending the correct bits to the stdout. + if self.pending_display_message is None: + # We need this for the case where we've just come back around from sending a display message. + if display_marker_loc == -1: + send_stream(msgstr) + msgstr = "" + continue + + # Send the beginning of the message to the stream output. + if display_marker_loc > 0: + send_stream(msgstr[:display_marker_loc]) + # Meanwhile, set up for the display message. + self.pending_display_message = "" + msgstr = msgstr[display_marker_loc + len(display_start) :] + continue + # endif self.pending_display_message is None + + # Try to find the end marker. If we did find it, add it to the + # pending display message and send it out. + display_marker_loc = msgstr.find(display_end) + if display_marker_loc != -1: + self.pending_display_message += msgstr[:display_marker_loc] + self._send_display_message() + # Reset display_marker_loc to a potential display_start index + # and reset msgstr so things get sent properly. + msgstr = msgstr[display_marker_loc + len(display_end) :] + display_marker_loc = msgstr.find(display_start) + continue + # endif display_marker_loc != -1 + + self.pending_display_message += msgstr + # If we just completed the display end message, send it out. + display_marker_loc = self.pending_display_message.find(display_end) + if display_marker_loc == -1: + msgstr = "" + continue + self.pending_display_message = self.pending_display_message[ + :display_marker_loc + ] + # Discard the display end message. + msgstr = self.pending_display_message[ + display_marker_loc + len(display_end) : + ].rstrip("\r\n") + self._send_display_message() + + # Finally, flush the stream from the message. + send_stream(msgstr) + # endwhile len(msgstr) > 0 + + def _send_display_message(self): + """Send the current pending display message to the client.""" + + display_message = json.loads(self.pending_display_message) + self.kernel.send_response( + self.kernel.iopub_socket, + "display_data", + { + "data": display_message[0], + "metadata": display_message[1], + }, + ) + self.pending_display_message = None + + +# ===----------------------------------------------------------------------=== # +# MojoKernel +# ===----------------------------------------------------------------------=== # + + +class MojoKernel(Kernel): + """A Jupyter kernel for Mojo.""" + + def __init__(self, **kwargs): + """Initialize the Mojo kernel. + + This loads the MojoJupyter library and starts a kernel repl session. + """ + # Kernel Metadata. + self.implementation = "MojoKernel" + self.implementation_version = "0.1" + self.language = "mojo" + self.language_version = "0.1" + self.language_info = { + "name": "mojo", + "mimetype": "text/x-mojo", + "file_extension": ".mojo", + "codemirror_mode": {"name": "mojo"}, + } + self.banner = "" + self.auto_gen_cell_id_count = 0 + super(MojoKernel, self).__init__(**kwargs) + + # Load the MojoJupyter library, and initialize the result types of the + # functions we use. + self.lib_mojo_jupyter: ctypes.CDLL = self.load_mojo_lib() + self.lib_mojo_jupyter.initMojoKernel.restype = ctypes.c_void_p + self.lib_mojo_jupyter.checkMojoExecutionFinished.restype = ctypes.c_int + + # The type of the output callback function. It takes a name and a + # message. + self.output_callback_type: ctypes.CFUNCTYPE = ctypes.CFUNCTYPE( + None, + ctypes.c_char_p, + ctypes.c_char_p, + ) + self.output_processor = OutputProcessor(self) + self.output_callback = self.output_callback_type( + lambda name, msg: self.output_processor.process_output(name, msg) + ) + + self.mojo_kernel: ctypes.c_void_p = ( + self.lib_mojo_jupyter.initMojoKernel( + self.output_callback, + ctypes.c_char_p(self.mojoReplExe.encode("utf-8")), + ctypes.c_char_p(None), # lldbInitFile + ) + ) + if not self.mojo_kernel: + raise RuntimeError("Unable to initialize Mojo kernel.") + + def _send_internal_error_message(self): + self.output_processor.send_message( + "stderr", + ( + "Internal Mojo Kernel Error\n\nThe Jupyter Notebook encountered" + " an internal error and was unable to evaluate the provided" + " expression. Please report this issue.\n\nMore information" + " about this error can be found in the server error log." + ), + ) + + def __del__(self): + """Destroy the Mojo kernel.""" + self.lib_mojo_jupyter.destroyMojoKernel(self.mojo_kernel) + + def load_mojo_lib(self) -> ctypes.CDLL: + """Load the libMojoJupyter library. + + On success, this initializes `mojoReplExe` returns the loaded library. + """ + + # Grab the mojo repl executable from the config. + config = ConfigParser() + config.read(Path(os.environ["MODULAR_HOME"]) / "modular.cfg") + mojoReplExePath = Path( + config.get("mojo", "repl_entry_point").rstrip(";") + ) + if not mojoReplExePath.exists(): + raise RuntimeError( + "Unable to locate `mojo-repl-entry-point` executable." + ) + self.mojoReplExe = str(mojoReplExePath) + + # Make sure the lib directory is in the path. + libDir = mojoReplExePath.parent + os.environ["PATH"] += os.pathsep + str(libDir) + + # Load the MojoJupyter library. This library provides the internal + # implementation of the kernel. + mojoJupyterPath = Path(config.get("mojo", "jupyter_path").rstrip(";")) + if not mojoJupyterPath.exists(): + raise RuntimeError("Unable to locate `MojoJupyter` library.") + return ctypes.cdll.LoadLibrary(str(mojoJupyterPath)) + + def do_execute( + self, + code: str, + silent: bool = False, + store_history: bool = True, + user_expressions: Optional[Dict[str, Any]] = None, + allow_stdin: bool = False, + *, + cell_id: Optional[str] = None, + ): + """Execute a code cell.""" + # TODO: Better propagate errors from the kernel execution, process + # provided arguments, etc. + + # Wait for the currently running execution to finish. + def wait_for_execution() -> ExecutionFinishedState: + # Wait for the execution to finish. + while True: + # Sleep for a bit to avoid busy spinning while waiting for the + # execution to finish. + time.sleep(0.05) + + # Poll the kernel to see if the execution has finished. + result: int = self.lib_mojo_jupyter.checkMojoExecutionFinished( + ctypes.c_void_p(self.mojo_kernel), + ) + if result != ExecutionFinishedState.NotFinished: + return ExecutionFinishedState(result) + + try: + # jupyter on the cli doesn't provide a cell id, so we need to + # autogenerate one. + if cell_id is None: + cell_id = f"__autogen_cell_id_{self.auto_gen_cell_id_count}" + self.auto_gen_cell_id_count += 1 + + # Start execution of the expression. + finish_state = self.lib_mojo_jupyter.startMojoExecution( + ctypes.c_void_p(self.mojo_kernel), + ctypes.c_char_p(cell_id.encode("utf-8")), + ctypes.c_char_p(code.encode("utf-8")), + ctypes.c_int(store_history), + ) + if finish_state == ExecutionFinishedState.NotFinished: + finish_state = wait_for_execution() + if finish_state == ExecutionFinishedState.FinishedError: + return {"status": "error"} + + return { + "status": "ok", + "execution_count": self.execution_count, + "payload": [], + "user_expressions": {}, + } + except KeyboardInterrupt: + # Interrupt the current kernel execution. + self.lib_mojo_jupyter.interruptMojoExecution( + ctypes.c_void_p(self.mojo_kernel) + ) + wait_for_execution() + + # TODO: When Mojo actually has debug info again, we should emit the + # current stack frame here. + return { + "status": "error", + "execution_count": self.execution_count, + } + + except: + traceback.print_exc() + self._send_internal_error_message() + + return { + "status": "error", + "execution_count": self.execution_count, + } + + def do_complete(self, code: str, cursor_pos: int): + """Find code completions for the given code and cursor position.""" + + # The type of the completion function, it takes a completion label. + completion_callback_type: ctypes.CFUNCTYPE = ctypes.CFUNCTYPE( + None, ctypes.c_char_p + ) + + # Build the callback handler used to process completion results. + results = [] + completion_callback = completion_callback_type( + lambda result: results.append(result.decode()) + ) + + self.lib_mojo_jupyter.checkMojoCodeComplete( + ctypes.c_void_p(self.mojo_kernel), + ctypes.c_char_p(code.encode()), + ctypes.c_int(cursor_pos), + completion_callback, + ) + + return { + "matches": results, + "cursor_end": cursor_pos, + "cursor_start": cursor_pos, + "metadata": {}, + "status": "ok", + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument( + "--modular-home", + required=False, + help="The value of the env var MODULAR_HOME.", + default=None + ) + args, jupyter_args = parser.parse_known_args() + + if os.environ.get('MODULAR_HOME') is None: + if args.modular_home is None: + raise Exception("MODULAR_HOME required") + os.environ["MODULAR_HOME"] = args.modular_home + + # We pass the kernel name as a command-line arg, since Jupyter gives those + # highest priority (in particular overriding any system-wide config). + IPKernelApp.launch_instance( + argv=jupyter_args + ["--IPKernelApp.kernel_class=__main__.MojoKernel"] + ) -- GitLab