feat: add GPU support for NVIDIA & AMD GPUs in installation and Docker setup (ollama)

[psyray]

Fixes #227

Config for the 2 types of card has been added, but only NVIDIA GPU is effective and tested for the moment.
I don't have any AMD card to do testing.
Maybe later, but for now let's go with NVIDIA

Summary by Sourcery

Add GPU support for NVIDIA and AMD GPUs in the installation and Docker setup, including scripts for GPU detection and configuration, and Dockerfiles for specific GPU types.

New Features:

• Introduce GPU support for NVIDIA and AMD GPUs in the installation script and Docker setup, allowing users to leverage GPU capabilities for enhanced performance.

Enhancements:

• Add a new script, gpu_support.sh, to detect GPU type and configure the environment accordingly.
• Update install.sh to check for GPU support and install necessary dependencies based on the detected GPU type.
• Modify the Makefile to include GPU configuration options, enabling users to activate GPU support during Docker operations.

Build:

• Create separate Dockerfiles for NVIDIA and AMD GPU support under the docker/ollama directory.
• Add a new Docker Compose file, docker-compose.gpu.yml, to handle GPU-specific configurations.

Todo

• NVIDIA test
• AMD test

[sourcery-ai]

Reviewer's Guide by Sourcery

This PR implements GPU support for NVIDIA and AMD GPUs in the installation process and Docker setup. The implementation includes automatic GPU detection, installation of required drivers/toolkits, and Docker configuration for GPU support. The changes primarily focus on NVIDIA GPU support, with AMD support added but untested.

Sequence diagram for GPU detection and setup

sequenceDiagram
    participant User
    participant InstallScript as install.sh
    participant GPUSupportScript as gpu_support.sh
    participant Docker

    User->>InstallScript: Run installation
    InstallScript->>GPUSupportScript: Source gpu_support.sh
    InstallScript->>GPUSupportScript: Detect GPU type
    GPUSupportScript-->>InstallScript: Return GPU type (nvidia/amd/none)
    InstallScript->>InstallScript: Check GPU type
    alt NVIDIA GPU detected
        InstallScript->>Docker: Install NVIDIA Container Toolkit
        InstallScript->>Docker: Configure Docker runtime for NVIDIA
    else AMD GPU detected
        InstallScript->>Docker: Install ROCm
        InstallScript->>Docker: Configure user groups
    else No GPU detected
        InstallScript->>InstallScript: Continue with CPU-only setup
    end
    InstallScript-->>User: Installation complete

Loading

Class diagram for GPU support scripts and Dockerfiles

classDiagram
    class InstallScript {
        +check_gpu_support()
        +install_make()
        +main()
    }
    class GPUSupportScript {
        +detect_gpu_type()
        +debug_log()
    }
    class Makefile {
        +gpu_config()
    }
    class DockerComposeGPU {
        +services
        +profiles
    }
    class DockerfileNVIDIA {
        +NVIDIA_VISIBLE_DEVICES
        +NVIDIA_DRIVER_CAPABILITIES
    }
    class DockerfileAMD {
        +HSA_OVERRIDE_GFX_VERSION
        +ROCR_VISIBLE_DEVICES
    }
    InstallScript --> GPUSupportScript : uses
    InstallScript --> Makefile : modifies
    DockerComposeGPU --> DockerfileNVIDIA : extends
    DockerComposeGPU --> DockerfileAMD : extends

Loading

File-Level Changes

Change	Details	Files
Added GPU detection and configuration system	Created new script for automatic GPU detection (NVIDIA/AMD) Added GPU type validation and environment configuration Implemented interactive and non-interactive GPU setup options Added error handling for GPU detection and setup failures	scripts/gpu_support.sh install.sh
Enhanced Docker and Makefile system for GPU support	Added GPU-specific Docker Compose configuration Implemented GPU detection and runtime configuration in Makefile Added GPU support flags to build and deployment commands Created separate Dockerfiles for NVIDIA and AMD GPU configurations	Makefile docker/docker-compose.gpu.yml docker/ollama/nvidia.Dockerfile docker/ollama/amd.Dockerfile
Added NVIDIA Container Toolkit integration	Added NVIDIA repository and key configuration Implemented NVIDIA Container Toolkit installation Added Docker runtime configuration for NVIDIA GPUs Configured environment variables for NVIDIA GPU support	install.sh docker/ollama/nvidia.Dockerfile
Added AMD ROCm support framework	Added ROCm repository and installation configuration Implemented user group configuration for AMD GPU access Added ROCm-specific environment variables Created AMD-specific Docker configuration	install.sh docker/ollama/amd.Dockerfile

Possibly linked issues

• Release/2.1.0 #1: The PR addresses the issue by adding NVIDIA GPU support in the ollama Docker setup.

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time. You can also use
  this command to specify where the summary should be inserted.

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey @psyray - I've reviewed your changes - here's some feedback:

Overall Comments:

• The AMD GPU support is currently untested and should be clearly marked as experimental in the documentation. Consider adding warning logs when users attempt to use AMD GPU support to make it clear it's not production-ready.

Here's what I looked at during the review

• 🟡 General issues: 1 issue found
• 🟢 Security: all looks good
• 🟢 Testing: all looks good
• 🟢 Complexity: all looks good
• 🟢 Documentation: all looks good

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}