.envrc.recommended

#!/usr/bin/env bash

set -euo pipefail

info() {
  DIRENV_LOG_FORMAT="envrc: %s" log_status "$@"
}

export PIXI_BETA_WARNING_OFF=true

if ! nvidia_smi="$(type -p "nvidia-smi")" || [[ -z "$nvidia_smi" ]]; then
  info "Did not detect nvidia-smi, using non-GPU environment"
  pixi_env="dev"
else
  # HACK: may be a bit fragile
  # Return the __cuda virtual system requirement version from pixi
  cuda_major_version=$(pixi info | awk 'BEGIN { FS = "[: \t\n=.]+"}; /__cuda/ { print $3}; ')
  info "Detected CUDA major version: $cuda_major_version"
  pixi_env="cuda${cuda_major_version}-dev"
fi

info "Using pixi environment: $pixi_env"

eval "$(pixi shell-hook -e "$pixi_env")"

if xhost_cmd="$(type -p "xhost")" && [[ -n "$xhost_cmd" ]]; then
  "$xhost_cmd" +si:localuser:"$(whoami)" >&/dev/null && {
    info "Display present, setting XLA_PYTHON_CLIENT_MEM_FRACTION to 40%"
    export XLA_PYTHON_CLIENT_MEM_FRACTION=".40"
  }
fi