Fix install script for all dependencies without breaking conda (Real!) (#121)

Author: joecummings

README.md

@@ -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
 ```
 

scripts/install.sh

@@ -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
@@ -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  "$@"; }
+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..."
@@ -197,6 +231,9 @@ EOF
 
     echo ""
     log_info "Installation completed successfully!"
+    echo ""
+    log_info "Re-activate the conda environment to make the changes take effect:"
+    log_info "  conda activate $CONDA_DEFAULT_ENV"
 }
 
 main "$@"