From c19b229532442cc0350d6c67eb07b51fecdc097c Mon Sep 17 00:00:00 2001
From: sasha0552 <admin@sasha0552.org>
Date: Tue, 3 Sep 2024 12:20:23 +0000
Subject: [PATCH] [misc] Use transient-package for triton installation

---
 airootfs/root/customize_airootfs.sh.jinja2        |  6 ++++++
 .../scripts/1000-vllm.sh.jinja2                   |  6 ------
 .../scripts/2000-replace-triton.sh.jinja2         | 15 +++++++++++++++
 3 files changed, 21 insertions(+), 6 deletions(-)
 create mode 100644 airootfs/root/customize_airootfs/scripts/2000-replace-triton.sh.jinja2

diff --git a/airootfs/root/customize_airootfs.sh.jinja2 b/airootfs/root/customize_airootfs.sh.jinja2
index f054639..c80cae9 100644
--- a/airootfs/root/customize_airootfs.sh.jinja2
+++ b/airootfs/root/customize_airootfs.sh.jinja2
@@ -8,6 +8,9 @@ export CUSTOMIZE_AIROOTFS=/root/customize_airootfs
 # home directory
 export HOME=/home/tori
 
+# path
+export PATH=$HOME/.local/bin:$PATH
+
 # directory with patches
 export TORI_PATCHES=$HOME/.local/share/tori/patches
 ##### ENVIRONMENT VARIABLES #####
@@ -34,6 +37,9 @@ systemctl enable nvidia-persistenced
 systemctl enable ly
 {% endif %}
 
+# install transient-package
+su tori -c "pipx install transient-package"
+
 {% if type != "empty" %}
 # enter home directory
 cd "/home/tori"
diff --git a/airootfs/root/customize_airootfs/scripts/1000-vllm.sh.jinja2 b/airootfs/root/customize_airootfs/scripts/1000-vllm.sh.jinja2
index 2e76906..ab5ee4e 100644
--- a/airootfs/root/customize_airootfs/scripts/1000-vllm.sh.jinja2
+++ b/airootfs/root/customize_airootfs/scripts/1000-vllm.sh.jinja2
@@ -7,10 +7,4 @@ export PIP_EXTRA_INDEX_URL="https://sasha0552.github.io/pascal-pkgs-ci/"
 
 # install patched vllm
 pipx install vllm-pascal
-
-# remove triton
-pipx runpip vllm-pascal uninstall -y triton
-
-# install patched triton
-pipx runpip vllm-pascal install triton-pascal
 {% endif %}
diff --git a/airootfs/root/customize_airootfs/scripts/2000-replace-triton.sh.jinja2 b/airootfs/root/customize_airootfs/scripts/2000-replace-triton.sh.jinja2
new file mode 100644
index 0000000..391e206
--- /dev/null
+++ b/airootfs/root/customize_airootfs/scripts/2000-replace-triton.sh.jinja2
@@ -0,0 +1,15 @@
+#!/bin/sh
+set -eu
+
+{% if platform == "cuda" %}
+# use repo with patched packages
+export PIP_EXTRA_INDEX_URL="https://sasha0552.github.io/pascal-pkgs-ci/"
+
+# Iterate over all virtual environments with triton
+for dir in $(find /home/tori -type d -wholename "*/site-packages/triton"); do
+  # Replace triton
+  transient-package install                     \
+    --interpreter "$dir/../../../../bin/python" \
+    --source triton                             \
+    --target triton-pascal
+{% endif %}