🍔 Use a larger runner for the LRE workflow

(draft) Rewrites large parts of the packaging logic.

The regular runner doesn't have enough memory.

Also implement some cursed hacks to make `mojo package` work in remote
execution.

Author: aaronmondal

.github/workflows/local-remote-execution.yaml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-24.04]
+        os: [large-ubuntu-22.04]
     name: ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 45
@@ -52,4 +52,8 @@ jobs:
       - name: Run Mojo tests against the cluster
         run: |
           nix develop --impure --command \
-          bash -c "lre-bazel test --jobs=4 @mojo//... --verbose_failures"
+          bash -c "lre-bazel test \
+            --jobs=$(nproc) \
+            --keep_going \
+            --verbose_failures \
+              @mojo//..."

README.md

@@ -31,7 +31,7 @@ Brings Mojo to all Linux distributions.
   Also, the LRE framework is still actively under development (in fact,
   `rules_mojo` acts as a validation case for LRE), so expect larger scale
   refactors.
-- At the moment there is no way to chose the Mojo version. `rules_mojo` uses
+- At the moment there is no way to choose the Mojo version. `rules_mojo` uses
   a pinned `nightly` toolchain that happened to not crash because stable
   crashed.
 

examples/BUILD.bazel

@@ -17,6 +17,7 @@ mojo_binary(
 )
 
 mojo_test(
-    name = "hello_world_test",
+    name = "hello_test",
     srcs = ["hello.mojo"],
+    deps = [":mypackage"],
 )

examples/hello.mojo

@@ -1,5 +1,6 @@
 from mypackage.mymodule import MyPair
 
+
 fn main():
     var mine = MyPair(2, 4)
     mine.dump()

flake.nix

@@ -81,59 +81,37 @@
           , ...
           }:
           let
+            tinfo6-bin = pkgs.callPackage ./tinfo6-bin.nix { };
+
+            ncurses6-bin = pkgs.callPackage ./ncurses6-bin.nix { };
+
             native-cli = inputs.nativelink.packages.${system}.native-cli;
 
-            lre-mojo-cluster = import ./local-remote-execution/lre-mojo-cluster.nix {
+            lre-mojo-cluster = pkgs.callPackage ./local-remote-execution/lre-mojo-cluster.nix {
               inherit native-cli;
-              inherit (pkgs)
-                curl
-                git
-                kubectl
-                kustomize
-                nix
-                writeShellScriptBin;
             };
 
             lre-cc = nativelink.packages.${system}.lre-cc;
             inherit (nix2container.packages.${system}.nix2container) buildImage;
 
-            mojo-sdk = import ./mojo-sdk.nix {
-              inherit (pkgs)
-                autoPatchelfHook
-                fetchurl
-                glibc
-                icu
-                libedit
-                libgcc
-                libxml2
-                ncurses
-                stdenv
-                zlib
-                zstd;
+            mojo-sdk = pkgs.callPackage ./mojo-sdk.nix {
+              inherit (pkgs.lib) makeBinPath makeLibraryPath;
+              ncurses = ncurses6-bin;
+              tinfo = tinfo6-bin;
             };
 
-            lre-mojo = import ./local-remote-execution/lre-mojo.nix {
-              inherit (pkgs)
-                gcc
-                lib
-                ncurses
-                python312
-                zlib;
-              inherit buildImage lre-cc mojo-sdk;
+            lre-mojo = pkgs.callPackage ./local-remote-execution/lre-mojo.nix {
+              # ncurses = ncurses6-bin;
+              # tinfo = ncurses6-bin;
+              mojoEnv = self.lib.defaultMojoEnv {
+                inherit pkgs mojo-sdk;
+              };
+              inherit buildImage lre-cc;
             };
 
-            lre-kill-the-mojo = import ./local-remote-execution/lre-kill-the-mojo.nix {
-              inherit (pkgs) docker findutils kind writeShellScriptBin;
-            };
+            lre-kill-the-mojo = pkgs.callPackage ./local-remote-execution/lre-kill-the-mojo.nix { };
 
-            createWorker = import ./local-remote-execution/create-worker.nix {
-              inherit (pkgs)
-                bash
-                buildEnv
-                coreutils
-                lib
-                runCommand
-                runtimeShell;
+            createWorker = pkgs.callPackage ./local-remote-execution/create-worker.nix {
               inherit buildImage self;
               nativelink = nativelink.packages.${system}.nativelink-debug;
             };
@@ -143,11 +121,7 @@
               exec ${pkgs.bazelisk}/bin/bazelisk "$@"
             '';
 
-            lre-bazel = import ./lre-bazel.nix {
-              inherit bazel;
-              inherit (pkgs) kubectl writeShellScriptBin;
-            };
-
+            lre-bazel = pkgs.callPackage ./lre-bazel.nix { inherit bazel; };
           in
           {
             _module.args.pkgs = import self.inputs.nixpkgs {
@@ -186,11 +160,12 @@
                 bazel
                 lre-bazel
                 pkgs.kubectl
-                pkgs.zlib
                 pkgs.python312
                 pkgs.tektoncd-cli
                 pkgs.kind
                 pkgs.vale
+                pkgs.jq
+                pkgs.dive
               ];
 
               shellHook = ''

local-remote-execution/lre-mojo.nix

@@ -1,32 +1,28 @@
 { buildImage
-, gcc
-, lib
 , lre-cc
-, mojo-sdk
-, ncurses
-, python312
-, zlib
+, mojoEnv
 , ...
 }:
 
 let
   # This environment is shared between toolchain autogen images and the final
   # toolchain image.
-  Env = [
-    # Add all tooling here so that the generated toolchains use `/nix/store/*`
-    # paths instead of `/bin` or `/usr/bin`. This way we're guaranteed to use
-    # binary identical toolchains during local and remote execution.
-    ("PATH="
-      + (lib.strings.concatStringsSep ":" [
-      "${gcc}/bin"
-      "${mojo-sdk}/bin"
-    ]))
-    "MODULAR_HOME=${mojo-sdk}"
-    "MOJO_COMPILER=${mojo-sdk}/bin/mojo"
-    "MOJO_CC_PATH=${gcc}/bin:${mojo-sdk}/bin"
-    "MOJO_LIBRARY_PATH=${zlib}/lib:${ncurses}/lib"
-    "MOJO_PYTHON_LIBRARY=${python312}/lib"
-  ];
+  Env = mojoEnv;
+  # [
+  #   # Add all tooling here so that the generated toolchains use `/nix/store/*`
+  #   # paths instead of `/bin` or `/usr/bin`. This way we're guaranteed to use
+  #   # binary identical toolchains during local and remote execution.
+  #   ("PATH="
+  #     + (lib.strings.concatStringsSep ":" [
+  #     "${gcc}/bin"
+  #     "${mojo-sdk}/bin"
+  #   ]))
+  #   "MODULAR_HOME=${mojo-sdk}"
+  #   "MOJO_COMPILER=${mojo-sdk}/bin/mojo"
+  #   "MOJO_CC_PATH=${gcc}/bin:${mojo-sdk}/bin"
+  #   "MOJO_LIBRARY_PATH=${zlib}/lib:${ncurses}/lib"
+  #   "MOJO_PYTHON_LIBRARY=${python312}/lib"
+  # ];
 in
 buildImage {
   name = "lre-mojo";

modules/defaultMojoEnv.nix

@@ -1,10 +1,14 @@
 { pkgs, mojo-sdk, ... }:
 
 [
+  # "PATH=${pkgs.lib.makeBinPath mojo-sdk.passthru.mojoBinPath}"
   "MODULAR_HOME=${mojo-sdk}"
   "MOJO_COMPILER=${mojo-sdk}/bin/mojo"
-  "MOJO_CC_PATH=${pkgs.gcc}/bin:${mojo-sdk}/bin"
-  "MOJO_LIBRARY_PATH=${pkgs.zlib}/lib:${pkgs.ncurses}/lib"
+  "MOJO_CC_PATH=${pkgs.lib.makeBinPath mojo-sdk.passthru.mojoBinPath}"
+  "MOJO_LIBRARY_PATH=${pkgs.lib.makeLibraryPath mojo-sdk.passthru.mojoLibraryPath}"
+  "LD_LIBRARY_PATH=${pkgs.lib.makeLibraryPath mojo-sdk.passthru.mojoLibraryPath}"
+
+  # "MOJO_LIBRARY_PATH=${pkgs.zlib-ng.override{withZlibCompat=true;}}/lib:${pkgs.ncurses}/lib"
 
   # TODO(aaronmondal): This needs to be set during runtime. Let's just add it to
   #                    a mojo wrapper script.

mojo-sdk.nix

@@ -5,12 +5,23 @@
 , libedit
 , libgcc
 , libxml2
+, makeLibraryPath
+, makeWrapper
+, mold
 , ncurses
 , stdenv
-, zlib
+, tinfo
+, uutils-coreutils-noprefix
+, zlib-ng
 , zstd
+, clang
 , ...
 }:
+let
+
+  # The zlib-ng library is a more modern implementation of Zlib and can act as drop-in replacement.
+  zlib = zlib-ng.override { withZlibCompat = true; };
+in
 
 stdenv.mkDerivation rec {
   pname = "mojo";
@@ -26,37 +37,96 @@ stdenv.mkDerivation rec {
   };
 
   buildInputs = [
+    glibc
+    icu
+    libedit
+    libgcc
     libxml2
+    mold
+    tinfo
     ncurses
-    libedit
-    icu
     zlib
-    glibc
-    libgcc
   ];
 
-  nativeBuildInputs = [ autoPatchelfHook zstd ];
+  nativeBuildInputs = [ autoPatchelfHook zstd makeWrapper ];
+
+  # nativeBuildInputs = [ autoPatchelfHook zstd makeWrapper ];
 
   unpackPhase = ''
     zstd -d $src -c | tar xvf -
   '';
 
-  installPhase = ''
-    cp -r . $out
-
+  patchPhase = ''
     # Nixpkgs uses `libedit.so.0.0.72` for the version of libedit
     # that is `libedit.so.2.0.72` in Ubuntu.
-    ln -s ${libedit}/lib/libedit.so.0 $out/lib/libedit.so.2
+    ln -s ${libedit}/lib/libedit.so.0 lib/libedit.so.2
+
+    # Brute-force copy tinfo into the same directory as the mojo sources.
+    # It appears that mojo can't figure out how to find libtinfo otherwise.
+    cp ${tinfo}/lib/libtinfo.so.6.4 lib/libtinfo.so.6.4
+    ln -s libtinfo.so.6.4 lib/libtinfo.so.6
 
     # The mojo command uses this config file to determine the
     # locations to bundled dependencies. Remap it to /nix/store.
-    sed -i "s|\$_self.install_path|$out|g" $out/modular.cfg
+    sed -i "s|\$_self.install_path|$out|g" modular.cfg
 
     # The nightly config has what looks like a typo, omitting the `_`.
-    sed -i "s|\$self.install_path|$out|g" $out/modular.cfg
+    sed -i "s|\$self.install_path|$out|g" modular.cfg
 
     # Nightly builds require settings under `[mojo-nightly]`
-    sed -i "s|\[mojo]|[mojo-nightly]|g" $out/modular.cfg
+    sed -i "s|\[mojo]|[mojo-nightly]|g" modular.cfg
+
+    # This evil hack is a ~95% linktime speedup by using mold instead of ld.
+    #
+    # The mojo compiler somehow behaves differently when running in Bazel. There
+    # it's necessary to have the setup below specified.
+    #
+    # To work around the fact that we can't resolve `-Wl,-rpath` properly
+    # because of unknown wrapping logic, we use`-Xlinker,-rpath` which has
+    # semantics like `-Xlinker <arg>`.
+    #
+    # WARNING: At the moment linking phase can't figure out how to link tinfo.
+    # To work around this we allow leaving the tinfo symbols undefined and
+    # unlinked. This is a very risky and dangerous practice that might trigger
+    # random segfaults during runtime.
+    sed -i "s|system_libs = -lrt,-ldl,-lpthread,-lm,-lz,-ltinfo;|system_libs = -fuse-ld=mold,-lrt,-ldl,-lpthread,-lm,-L,${zlib}/lib,-Xlinker,-rpath=${zlib}/lib,-Xlinker,-L,${tinfo}/lib,-Xlinker,-rpath=${tinfo}/lib,-Xlinker,-rpath=${stdenv.cc.cc.lib}/lib,--verbose,-Xlinker,--verbose,-Xlinker,-v,-Xlinker,--allow-shlib-undefined,-Xlinker,--unresolved-symbols=ignore-all;|g" modular.cfg
+  '';
+
+  installPhase = ''
+    cp -r . $out
+
+    # The autoPatchelfHook would run before the fixupPhase, so we need to call
+    # wrapping logic in a custom postInstall phase. This way we run patchelf
+    # after the wrappers have been created.
+    runHook postInstall
+  '';
 
+  # These fixups are not too pretty, but at the moment the `mojo` compiler
+  # doesn't respect standard environment variables like `CC` and `LD`. Instead,
+  # we have to add C++ tooling to the PATH. We also wrap certain executables
+  # with `MODULAR_HOME` so that they work in mkShell environemnts.
+  postInstall = ''
+    wrap_with_env() {
+      wrapProgram $1 \
+        --prefix PATH : ${clang}/bin:${mold}/bin:${uutils-coreutils-noprefix}/bin \
+        --prefix MODULAR_HOME : $out \
+        --prefix LIBRARY_PATH : ${makeLibraryPath [ tinfo ncurses zlib stdenv.cc.cc.lib  ]} \
+        --prefix LD_LIBRARY_PATH : ${makeLibraryPath [ tinfo ncurses zlib stdenv.cc.cc.lib  ]} \
+        --set LD_LIBRARY_PATH ${makeLibraryPath [ tinfo ncurses zlib stdenv.cc.cc.lib ]}
+    }
+
+    wrap_with_env $out/bin/mojo
+    wrap_with_env $out/bin/mojo-lldb
+    wrap_with_env $out/bin/mojo-lsp-server
+    wrap_with_env $out/bin/lldb-argdumper
+    wrap_with_env $out/bin/lldb-server
+    wrap_with_env $out/bin/mojo-lldb-dap
   '';
+
+  passthru = {
+    # These may be passed through to Bazel invocations and can be used to
+    # construct remote execution toolchains for Mojo.
+    mojoBinPath = [ clang mold uutils-coreutils-noprefix ];
+    mojoLibraryPath = [ tinfo ncurses zlib stdenv.cc.cc.lib ];
+  };
 }

mojo/BUILD.bazel

@@ -15,6 +15,8 @@ load("//mojo:toolchain.bzl", "mojo_toolchain")
         "MOJO_CC_PATH",
         "MOJO_LIBRARY_PATH",
         "MOJO_PYTHON_LIBRARY",
+        "LD_LIBRARY_PATH",
+        "MOJO_DYNAMIC_LINKER",
 
         # Unset values default to an empty string.
         "MOJO_UNSET",

mojo/mojo.bzl

@@ -129,10 +129,12 @@ def _mojo_library_impl(ctx):
         tools = [toolchain.symbolizer],
         use_default_shell_env = False,
         env = {
-            "PATH": toolchain.MOJO_CC_PATH,
-            "MODULAR_HOME": toolchain.MODULAR_HOME,
-            "MOJO_PYTHON_LIBRARY": toolchain.MOJO_PYTHON_LIBRARY,
-            "LLVM_SYMBOLIZER_PATH": toolchain.symbolizer.path,
+            # "PATH": toolchain.MOJO_CC_PATH,
+            # "MODULAR_HOME": toolchain.MODULAR_HOME,
+            # "LIBRARY_PATH": toolchain.MOJO_LIBRARY_PATH,
+            # "MOJO_PYTHON_LIBRARY": toolchain.MOJO_PYTHON_LIBRARY,
+            # "LLVM_SYMBOLIZER_PATH": toolchain.symbolizer.path,
+            # "LD_LIBRARY_PATH": toolchain.MOJO_LIBRARY_PATH,
         },
     )
 
@@ -197,11 +199,12 @@ def _mojo_binary_impl(ctx):
         tools = [toolchain.symbolizer],
         use_default_shell_env = False,
         env = {
-            "MODULAR_HOME": toolchain.MODULAR_HOME,
-            "PATH": toolchain.MOJO_CC_PATH,
-            "LIBRARY_PATH": toolchain.MOJO_LIBRARY_PATH,
-            "MOJO_PYTHON_LIBRARY": toolchain.MOJO_PYTHON_LIBRARY,
-            "LLVM_SYMBOLIZER_PATH": toolchain.symbolizer.path,
+            # "MODULAR_HOME": toolchain.MODULAR_HOME,
+            # "PATH": toolchain.MOJO_CC_PATH,
+            # "LIBRARY_PATH": toolchain.MOJO_LIBRARY_PATH,
+            # "LD_LIBRARY_PATH": toolchain.MOJO_LIBRARY_PATH,
+            # "MOJO_PYTHON_LIBRARY": toolchain.MOJO_PYTHON_LIBRARY,
+            # "LLVM_SYMBOLIZER_PATH": toolchain.symbolizer.path,
         },
     )