.env

export LIBRARY_PATH="$CONDA_PREFIX/lib"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib"
export CUDA_HOME="$CONDA_PREFIX"
export OMP_NUM_THREADS="1"
export NCCL_P2P_LEVEL="NVL"

# in current directory
export SLURM_OUTPUT_DIR=".cache/slurm_outputs"
export SLURM_SNAPSHOT_ROOT=".cache/slurm_snapshots"

mkdir -p "tango_workspace" ".cache/torchrunx_logs" ".cache/torch_extensions" ".cache/triton" ".cache/torchinductor"
export TANGO_WORKSPACE_DIR=$(realpath "tango_workspace")
export TORCHRUNX_LOG_DIR=$(realpath ".cache/torchrunx_logs")
export TORCH_EXTENSIONS_DIR=$(realpath ".cache/torch_extensions")
export TRITON_CACHE_DIR=$(realpath ".cache/triton")
export TORCHINDUCTOR_CACHE_DIR=$(realpath ".cache/torchinductor")

# in current & snapshot directories
export PYTHONPATH="."
export PYTHONNOUSERSITE="1"
export PYTHONPYCACHEPREFIX=".cache/pycache"