infrastructure

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README

ollama.nix (958B)

      1 {
      2   config,
      3   pkgs,
      4   pkgs-unstable,
      5   lib,
      6   ...
      7 }:
      8 
      9 let
     10   cfg = config.services.local.ollama;
     11 in
     12 {
     13   options.services.local.ollama = {
     14     enable = lib.mkEnableOption "Ollama with ROCm support";
     15 
     16     rocmOverrideGfx = lib.mkOption {
     17       type = lib.types.str;
     18       default = "11.0.1";
     19       description = "ROCm GFX version override";
     20     };
     21 
     22     visibleDevices = lib.mkOption {
     23       type = lib.types.str;
     24       default = "";
     25       description = "ROCR_VISIBLE_DEVICES value (GPU UUID)";
     26       example = "GPU-198836ed61e1b1cc";
     27     };
     28   };
     29 
     30   config = lib.mkIf cfg.enable {
     31     services.ollama = {
     32       enable = true;
     33       package = pkgs-unstable.ollama-rocm;
     34       acceleration = "rocm";
     35       rocmOverrideGfx = cfg.rocmOverrideGfx;
     36       environmentVariables = {
     37         HSA_OVERRIDE_GFX_VERSION = cfg.rocmOverrideGfx;
     38       } // lib.optionalAttrs (cfg.visibleDevices != "") {
     39         ROCR_VISIBLE_DEVICES = cfg.visibleDevices;
     40       };
     41     };
     42   };
     43 }