Commit b0234de4 authored by sebastianpetrovski's avatar sebastianpetrovski
Browse files

nixos/litellm: seed tokenizer cache via CUSTOM_TIKTOKEN_CACHE_DIR 

Ensure LiteLLM starts in offline NixOS tests by pre-seeding the cl100k tokenizer cache in the path LiteLLM actually maps to tiktoken at runtime.

🤖 Generated with [eca](https://eca.dev

)

Co-Authored-By: default avatareca <noreply@eca.dev>
parent fe524207

nixos/modules/services/misc/litellm.nix

+39 −0
Original line number Diff line number Diff line
@@ -9,6 +9,36 @@ let 


  cfg = config.services.litellm;

  settingsFormat = pkgs.formats.yaml { };



  tiktokenEncodings = {

    cl100k_base = {

      url = "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken";

      hash = "sha256-Ijkht27pm96ZW3/3OFE+7xAPtR0YyTWXoRO8/+hlsqc=";

    };

  };



  tiktokenCacheEntries = lib.mapAttrsToList (

    _: encoding:

    let

      cacheKey = builtins.hashString "sha1" encoding.url;

      sourceFile = pkgs.fetchurl {

        inherit (encoding) url hash;

      };

    in

    {

      inherit cacheKey sourceFile;

    }

  ) tiktokenEncodings;



  seedTiktokenCacheScript = pkgs.writeShellScript "litellm-seed-tiktoken-cache" ''

    set -eu



    mkdir -p "$CUSTOM_TIKTOKEN_CACHE_DIR"



    ${lib.concatMapStringsSep "\n" (entry: ''

      ln -sf ${entry.sourceFile} "$CUSTOM_TIKTOKEN_CACHE_DIR/${entry.cacheKey}"

    '') tiktokenCacheEntries}

  '';

in

{

  options = {
@@ -134,6 +164,7 @@ in 
  config = lib.mkIf cfg.enable {

    systemd.tmpfiles.rules = [

      "d '${cfg.stateDir}/ui' 0700 - - - -"

      "d '${cfg.stateDir}/tiktoken-cache' 0700 - - - -"

    ];



    systemd.services.litellm = {
@@ -147,6 +178,9 @@ in 
        # UI lives in the read-only Nix store, so point it at a writable runtime path.

        LITELLM_NON_ROOT = "true";

        LITELLM_UI_PATH = "${cfg.stateDir}/ui";



        # LiteLLM sets TIKTOKEN_CACHE_DIR internally from this variable.

        CUSTOM_TIKTOKEN_CACHE_DIR = "${cfg.stateDir}/tiktoken-cache";

      }

      // cfg.environment;
@@ -156,6 +190,10 @@ in 
        in

        {

          ExecStartPre = [

            # Seed tokenizer cache with fixed-output files so startup does not

            # depend on outbound network access.

            seedTiktokenCacheScript



            # LiteLLM may rewrite/copy UI assets with read-only permissions

            # during previous runs; normalize writability on each start.

            "${pkgs.runtimeShell} -euc 'chmod -R u+rwX ${cfg.stateDir}/ui'"
@@ -166,6 +204,7 @@ in 
          StateDirectory = [

            "litellm"

            "litellm/ui"

            "litellm/tiktoken-cache"

          ];

          RuntimeDirectory = "litellm";

          RuntimeDirectoryMode = "0755";