path: root/nixos/tests/web-servers/llama-swap.nix

{
  pkgs,
  lib ? pkgs.lib,
  runTest,
  ...
}:

let
  makeLlamaSwapTest =
    name:
    {
      withUI ? true,
      # opt in for e2e test with a small LLM
      withLLM ? false,
    }:
    (
      let
        wrapSrc = attrs: pkgs.runCommand "${attrs.pname}-${attrs.version}" attrs "ln -s $src $out";

        smollm2-135m = wrapSrc rec {
          pname = "smollm2-135m";
          version = "9e6855bc4be717fca1ef21360a1db4b29d5c559a";
          src = pkgs.fetchurl {
            url = "https://huggingface.co/unsloth/SmolLM2-135M-Instruct-GGUF/resolve/${version}/SmolLM2-135M-Instruct-Q4_K_M.gguf";
            hash = "sha256-7V+jDEh7KC7BVsKQYvEiLlwgh1qUSsmCidvSQulH90c=";
          };

          meta.license = with lib.licenses; [
            asl20 # actual license of the model weights
            unfree # to force an opt-in - do not remove
          ];
        };

        # grab allowUnfreePredicate if it exists or default deny
        allowUnfreePredicate =
          if builtins.hasAttr "allowUnfreePredicate" pkgs.config then
            pkgs.config.allowUnfreePredicate
          else
            (_: false);

        # check if we can use smollm2-135m taking either globally allowUnfree or
        # explicit allow with predicate
        # or the test explicitly opts in
        useSmollm2-135m = withLLM || pkgs.config.allowUnfree || allowUnfreePredicate smollm2-135m;
      in
      {
        name = "llama-swap-${name}";
        meta.maintainers = with lib.maintainers; [
          jk
          podium868909
        ];

        nodes = {
          machine =
            { pkgs, ... }:
            {
              # running models can be memory intensive but
              # default `virtualisation.memorySize` is fine

              services.llama-swap = {
                enable = true;
                settings =
                  # config for basic tests
                  if !useSmollm2-135m then
                    { }
                  # config for extended tests using SmolLM2
                  else
                    let
                      llama-cpp = pkgs.llama-cpp;
                      llama-server = lib.getExe' llama-cpp "llama-server";
                    in
                    {
                      # more useful logging output
                      logLevel = "debug";
                      logToStdout = "both";

                      hooks.on_startup.preload = [
                        "smollm2"
                      ];
                      # temperature and top-k important for SmolLM2 performance/accuracy
                      models = {
                        "smollm2" = {
                          ttl = 10;
                          cmd = "${llama-server} --port \${PORT} -m ${smollm2-135m} --alias smollm2 --no-webui --temp 0.2 --top-k 9";
                        };
                        "smollm2-group-1" = {
                          cmd = "${llama-server} --port \${PORT} -m ${smollm2-135m} --alias smollm2-group-1 --no-webui --temp 0.2 --top-k 9 -c 1024";
                        };
                        "smollm2-group-2" = {
                          proxy = "http://127.0.0.1:5802";
                          cmd = "${llama-server} --port 5802 -m ${smollm2-135m} --alias smollm2-group-2 --no-webui --temp 0.2 --top-k 9 -c 1024";
                        };
                      };
                      groups = {
                        "standalone" = {
                          swap = true;
                          exclusive = true;
                          members = [
                            "smollm2"
                          ];
                        };
                        "group" = {
                          swap = false;
                          exclusive = true;
                          members = [
                            "smollm2-group-1"
                            "smollm2-group-2"
                          ];
                        };
                      };
                    };
              };
            };
        };

        testScript =
          { nodes, ... }:
          ''
            # core tests
            import json

            def get_json(route):
              args = [
                '-v',
                '-s',
                '--fail',
                '-H "Content-Type: application/json"'
              ]
              return json.loads(machine.succeed("curl {args} http://localhost:8080{route}".format(args=" ".join(args), route=route)))

            def post_json(route, data):
              args = [
                '-v',
                '-s',
                '--fail',
                '-H "Content-Type: application/json"',
                "-d '{d}'".format(d=json.dumps(data))
              ]
              return json.loads(machine.succeed('curl {args} http://localhost:8080{route}'.format(args=" ".join(args), route=route)))

            machine.wait_for_unit('llama-swap')
            machine.wait_for_open_port(8080)

          ''
          + lib.optionalString withUI ''
            with subtest('check is serving ui'):
              machine.succeed('curl --fail http:/localhost:8080/ui/')

          ''
          + ''
            with subtest('check is healthy'):
              response_healthy = machine.wait_until_succeeds('curl --silent --fail http://localhost:8080/health')
              assert 'OK' in response_healthy, '/health was not OK'

          ''
          + lib.optionalString useSmollm2-135m ''
            # extended tests using SmolLM2
            with subtest('check `/running` for preloaded smollm2'):
              machine.wait_until_succeeds('curl --silent --fail http://localhost:8080/running | grep "smollm2"')
              running_response = get_json('/running')
              assert len(running_response['running']) == 1, f'running doesn\'t have one entry as expected: {running_response}'
              running_model = running_response['running'][0]
              assert running_model['model'] == 'smollm2', f'running model is not smollm2: {running_response}'
              assert running_model['state'] == 'ready', f'running smollm2 is not ready: {running_response}'

            with subtest('runs smollm2'):
              response = None
              with subtest('send request to smollm2'):
                data = {
                  'model': 'smollm2',
                  'messages': [
                    {
                      'role': 'user',
                      'content': 'Say hello'
                    }
                  ]
                }
                response = post_json('/v1/chat/completions', data)

              with subtest('response is from smollm2'):
                assert response['model'] == 'smollm2', f'response was not from smollm2: {response['model']}'

              with subtest('response contains at least one item in "choices"'):
                assert len(response['choices']) >= 1, f'response had no choices: {response}'

              assistant_choices = None
              with subtest('response contains at least one "assistant" message'):
                assistant_choices = [c for c in response['choices'] if c['message']['role'] == 'assistant']
                assert len(assistant_choices) >= 1, f'response had no assistant message: {response}'

              with subtest('first message (lowercase) starts with "hello"'):
                assert assistant_choices[0]['message']['content'].lower()[:5] == 'hello', f'response didn\'t start with hello: {response}'

              with subtest('check `/running` for just smollm2'):
                running_response = get_json('/running')
                assert len(running_response['running']) == 1, f'running doesn\'t have one entry as expected: {running_response}'
                running_model = running_response['running'][0]
                assert running_model['model'] == 'smollm2', f'running model is not smollm2: {running_response}'
                assert running_model['state'] == 'ready', f'running smollm2 is not ready: {running_response}'

            with subtest('check `/running` for smollm2 to timeout'):
              machine.wait_until_succeeds('curl --silent --fail http://localhost:8080/running | grep -v "smollm2"', timeout=11)
              running_response = get_json('/running')
              assert len(running_response['running']) == 0, "smollm2 was still running after timeout"

            with subtest('runs smollm2-group-1 and smollm2-group-2'):
              response_1 = None
              with subtest('send request to smollm2-group-1'):
                data = {
                  'model': 'smollm2-group-1',
                  'messages': [
                    {
                      'role': 'user',
                      'content': 'Say hello'
                    }
                  ]
                }
                response_1 = post_json('/v1/chat/completions', data)

              with subtest('response 1 is from smollm2-group-1'):
                assert response_1['model'] == 'smollm2-group-1', f'response was not from smollm2-group-1: {response_1['model']}'

              with subtest('response 1 contains at least one item in "choices"'):
                assert len(response_1['choices']) >= 1, f'response had no choices: {response_1}'

              assistant_choices_1 = None
              with subtest('response 1 contains at least one "assistant" message'):
                assistant_choices_1 = [c for c in response_1['choices'] if c['message']['role'] == 'assistant']
                assert len(assistant_choices_1) >= 1, f'response had no assistant message: {response_1}'

              with subtest('first message (lowercase) in response 1 starts with "hello"'):
                assert assistant_choices_1[0]['message']['content'].lower()[:5] == 'hello', f'response didn\'t start with hello: {response_1}'

              with subtest('check `/running` for just smollm2-group-1'):
                running_response = get_json('/running')
                assert len(running_response['running']) == 1, f'running doesn\'t have one entry as expected: {running_response}'
                running_model = running_response['running'][0]
                assert running_model['model'] == 'smollm2-group-1', f'running model is not smollm2-group-1: {running_response}'
                assert running_model['state'] == 'ready', f'running smollm2-group-1 is not ready: {running_response}'

              response_2 = None
              with subtest('send request to smollm2-group-2'):
                data = {
                  'model': 'smollm2-group-2',
                  'messages': [
                    {
                      'role': 'user',
                      'content': 'Say hello'
                    }
                  ]
                }
                response_2 = post_json('/v1/chat/completions', data)

              with subtest('response 2 is from smollm2-group-2'):
                assert response_2['model'] == 'smollm2-group-2', f'response was not from smollm2-group-2: {response_2['model']}'

              with subtest('response 2 contains at least one item in "choices"'):
                assert len(response_2['choices']) >= 1, f'response had no choices: {response_2}'

              assistant_choices_2 = None
              with subtest('response 2 contains at least one "assistant" message'):
                assistant_choices_2 = [c for c in response_2['choices'] if c['message']['role'] == 'assistant']
                assert len(assistant_choices_2) >= 1, f'response had no assistant message: {response_2}'

              with subtest('first message (lowercase) in response 1 starts with "hello"'):
                assert assistant_choices_2[0]['message']['content'].lower()[:5] == 'hello', f'response didn\'t start with hello: {response_2}'

              with subtest('check `/running` for both smollm2-group-1 and smollm2-group-2'):
                running_response = get_json('/running')['running']
                assert len(running_response) == 2, 'expected both in smollm2 group to be running'
                assert len([
                  rm for rm in running_response
                  if rm['state'] == 'ready' and rm['model'] == 'smollm2-group-1'
                ]) == 1, f'smollm2-group-1 not in running group and ready {running_response}'
                assert len([
                  rm for rm in running_response
                  if rm['state'] == 'ready' and rm['model'] == 'smollm2-group-2'
                ]) == 1, f'smollm2-group-2 not in running group and ready {running_response}'
          '';
      }
    );
in
lib.recurseIntoAttrs (
  builtins.mapAttrs (k: v: runTest (makeLlamaSwapTest k v)) {
    full = { };
    minimal = {
      withUI = false;
    };

    # opt in to smollm2-135m
    full-with-llm = {
      withLLM = true;
    };
  }
)