janet: Launch a jupyterlab session on Spartan

jlab.janet

#!/usr/bin/env janet

(use sh)
(import cmd)

# terminal formatting
(defn format [v]
  (def t {:green "\e[0;32m"
          :red "\e[0;31m"
          :yellow "\e[0;33m"
          :bold "\e[1m"
		  :grey "\e[0;37m"
          :reset "\e[0m"})
  (get t v))

(defmacro check-paths [& paths]
  ~(cond
     ,;(mapcat (fn [path]
                 [~(os/stat (string ,path "/bin/jupyter-lab"))
                  ~(string ,path "/bin/jupyter-lab")])
               paths)))

(defn find-jupyter-lab []
  (def git-root (let [devnull @""
                      stdout @""]
                  ($? git rev-parse --show-toplevel > ,stdout > [stderr devnull])
                  (string/trim stdout)))
  (def pwd (os/getenv "PWD"))
  (check-paths
    (os/getenv "VIRTUAL_ENV")
    (os/getenv "PIXI_ENV")
    (os/getenv "CONDA_PREFIX")
    (string pwd "/.venv")
    (string pwd "/venv")
    (string pwd "/.pixi/envs/default")
    (string git-root "/.venv")
    (string git-root "/venv")
    (string git-root "/.pixi/envs/default")))


(cmd/main
  (cmd/fn
    `
Launch an interactive Jupyter lab session.

Extra arguments, not listed below, will be passed to the srun command and override the defaults:
	-n 1 -p interactive
`
    [[jupyter-path --jupyter -j] (optional :string)
     [--port -p] (optional :number)
     [--time -t] (optional :string "02:00:00")
     [--mem -m] (optional :string "72G")
     [--cpus -c] (optional :string "4")
     [--job-name -J] (optional :string "jupyter-lab")
     extra (escape :string)]

    (def path (or jupyter-path (find-jupyter-lab)))

    (if (nil? path)
      (do (printf "%sFailed to find jupyter-lab executable!%s\n" (format :red) (format :reset)) (os/exit 1)))

    (printf "%sUsing jupyter-lab executable at %s%s" (format :green) path (format :reset))

    (def srun-opts (string/split " " (or (if-not (empty? extra) extra) "-n 1 -p interactive")))
    (array/push srun-opts (string/format "--time=%s" time))
    (array/push srun-opts (string/format "--mem=%s" mem))
    (array/push srun-opts (string/format "--cpus-per-task=%s" cpus))
    (array/push srun-opts (string/format "--job-name=%s" job-name))

    (printf "%sUsing srun options: %s%s\n" (format :grey) (string/join srun-opts " ") (format :reset))

    (def use-port (or port
                      (let [rand (scan-number ($< bash -c "echo -n $RANDOM"))]
                        (+ (mod rand (+ (- 4999 4001) 1)) 4001))))

    (def submit-script
      (string/format `
banner() {
    msg=" $2 "
    edge=$(echo "$msg" | sed 's/./─/g')
	echo "%s%s$1"
    echo "┌$edge┐"
    echo "│$msg│"
    echo "└$edge┘"
	echo "${3:-%s}"
}

echo ""
banner "Use this to connect:" "ssh -NL %d:localhost:%d -J spartan $(hostname -s)"
%s --no-browser --port %d
`
                     (format :yellow) (format :bold) (format :reset) use-port use-port path use-port))

    ($ srun ,;srun-opts bash -c ,submit-script)))