Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,6 @@ python -m apps.grpo.main

If you need to re-build the wheels for whatever reason, you can do so with:
```bash
conda create -n forge python=3.10
conda activate forge
./scripts/build_wheels.sh
```

Expand Down
47 changes: 42 additions & 5 deletions scripts/install.sh
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,13 @@ export CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
export PATH="${CUDA_HOME}/bin:$PATH"
export CUDA_INCLUDE_DIRS=$CUDA_HOME/include
export CUDA_CUDART_LIBRARY=$CUDA_HOME/lib64/libcudart.so
export LD_LIBRARY_PATH=${CONDA_PREFIX}/lib:/usr/local/cuda-12.9/compat:${LD_LIBRARY_PATH:-}
export LIBRARY_PATH=${CUDA_HOME}/lib64:${LIBRARY_PATH:-}

# Add only CUDA compat libs to LD_LIBRARY_PATH (safe for system tools)
if [ -n "${LD_LIBRARY_PATH:-}" ]; then
export LD_LIBRARY_PATH="/usr/local/cuda-${CUDA_VERSION}/compat:${LD_LIBRARY_PATH}"
else
export LD_LIBRARY_PATH="/usr/local/cuda-${CUDA_VERSION}/compat"
fi
EOF

# Create deactivation script to clean up
Expand All @@ -175,12 +180,41 @@ unset CUDA_NVCC_EXECUTABLE
unset CUDA_HOME
unset CUDA_INCLUDE_DIRS
unset CUDA_CUDART_LIBRARY
# Note: We don't unset PATH and LD_LIBRARY_PATH as they may have other content
# We intentionally do not mutate PATH or LD_LIBRARY_PATH here.
EOF

##########################################
# 2) Python-only LD_LIBRARY_PATH shim(s) #
##########################################
# These shell *functions* ensure that any `python`/`python3` invocation
# gets ${CONDA_PREFIX}/lib in its environment, without polluting the shell.
# This avoids OpenSSL/libcrypto collisions with system tools like /usr/bin/conda.
# TODO: Build Monarch with ABI3 to avoid this hack.
local py_shim_activate="${conda_env_dir}/etc/conda/activate.d/python_ld_shim.sh"
cat > "$py_shim_activate" << 'EOF'
# Define python wrappers that only set LD_LIBRARY_PATH for the launched process
python() { LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" command python "$@"; }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol

python3() { LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" command python3 "$@"; }

# Export functions to subshells when possible (best-effort, shell-dependent)
if [ -n "${BASH_VERSION:-}" ]; then
export -f python python3 2>/dev/null || true
elif [ -n "${ZSH_VERSION:-}" ]; then
typeset -fx python python3 >/dev/null 2>&1 || true
fi
EOF

# Source the activation script so it works in the current session.
log_info "Loading CUDA environment for current session..."
# Deactivation script to remove the function wrappers
cat > "${conda_env_dir}/etc/conda/deactivate.d/python_ld_shim.sh" << 'EOF'
unset -f python 2>/dev/null || true
unset -f python3 2>/dev/null || true
EOF

log_info "Loading CUDA env and python LD shim for current session..."
# shellcheck source=/dev/null
source "$cuda_activation_script"
# shellcheck source=/dev/null
source "$py_shim_activate"

# Test installation
log_info "Testing installation..."
Expand All @@ -197,6 +231,9 @@ EOF

echo ""
log_info "Installation completed successfully!"
echo ""
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should make it possible to source the install.sh script so we don't have to re-activate the conda environment (right now it kills ur terminal lol)

log_info "Re-activate the conda environment to make the changes take effect:"
log_info " conda activate $CONDA_DEFAULT_ENV"
}

main "$@"
Loading