From 88ff1cfbec1487f86ef2503413315c0fc38c292a Mon Sep 17 00:00:00 2001
From: Hugo Herter <git@hugoherter.com>
Date: Wed, 10 Sep 2025 10:36:53 +0200
Subject: [PATCH] python3Packages.unsloth: Add test using CUDA

Adds a test from the Unsloth repository that runs QLora training
and merging, ensuring basic functionality of the library.

Co-authored-by: Someone <else@someonex.net>
---
 .../python-modules/unsloth/default.nix        | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/pkgs/development/python-modules/unsloth/default.nix b/pkgs/development/python-modules/unsloth/default.nix
index dcb2ba5daf82..9b6fa8befe4b 100644
--- a/pkgs/development/python-modules/unsloth/default.nix
+++ b/pkgs/development/python-modules/unsloth/default.nix
@@ -27,8 +27,30 @@
   hf-transfer,
   diffusers,
   torchvision,
+
+  # tests
+  fetchFromGitHub,
+  cudaPackages,
+  python,
+  gcc,
 }:
 
+let
+  # Test files are absent from the PyPI package, so we fetch them separately.
+  testSrc = fetchFromGitHub {
+    owner = "unslothai";
+    repo = "unsloth";
+    rev = "cb78f0e83dc2d61fb1571b6e904eb2f064510d63";
+    hash = "sha256-0oR3m8jnjSdfjH+NslW6SsVj+0cQ4VUhKXZ38U/VBy0=";
+    # Keep only the tests directory, use the PyPI package for everything else.
+    postFetch = ''
+      mv $out/tests $TMPDIR/tests
+      rm -rf $out/*
+      mv $TMPDIR/tests $out/tests
+    '';
+  };
+in
+
 buildPythonPackage rec {
   pname = "unsloth";
   version = "2025.9.4";
@@ -85,6 +107,36 @@ buildPythonPackage rec {
   # NotImplementedError: Unsloth: No NVIDIA GPU found? Unsloth currently only supports GPUs!
   dontUsePythonImportsCheck = true;
 
+  passthru.tests = {
+    qlora-train-and-merge =
+      # FIXME: Replace python3.pkgs with python3Packages once possible, as to unbeak splicing.
+      # Cf. https://github.com/NixOS/nixpkgs/pull/394838#issuecomment-3319287038
+      cudaPackages.writeGpuTestPython.override { python3Packages = python.pkgs; }
+        {
+          libraries = ps: [
+            ps.unsloth
+            ps.unsloth-zoo
+          ];
+          # In-derivation test would require fetching the model from hugging-face which will not be trivial.
+          gpuCheckArgs.meta.broken = true;
+        }
+        # Triton JIT requires a C compiler at runtime and imports files from the test directory.
+        ''
+          import os, sys, runpy
+
+          os.environ["CC"]  = "${lib.getExe' gcc "cc"}";
+          os.environ["CXX"] = "${lib.getExe' gcc "cxx"}";
+
+          sys.path.insert(0, "${testSrc}")
+
+          # Execute the test file from the Nix store at runtime (no eval-time IFD).
+          runpy.run_path(
+              "${testSrc}/tests/qlora/test_unsloth_qlora_train_and_merge.py",
+              run_name="__main__"
+          )
+        '';
+  };
+
   meta = {
     description = "Finetune Llama 3.3, DeepSeek-R1 & Reasoning LLMs 2x faster with 70% less memory";
     homepage = "https://github.com/unslothai/unsloth";