tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6
7 # build-system
8 setuptools,
9
10 # dependencies
11 pytorch-lightning,
12 torch,
13
14 # tests
15 pytestCheckHook,
16 pythonAtLeast,
17}:
18
19buildPythonPackage (finalAttrs: {
20 pname = "finetuning-scheduler";
21 version = "2.10.0.post0";
22 pyproject = true;
23
24 src = fetchFromGitHub {
25 owner = "speediedan";
26 repo = "finetuning-scheduler";
27 tag = "v${finalAttrs.version}";
28 hash = "sha256-OeIpbxEjhvUzToy1jH9JcontSMfeozFjisTJCa0f4P0=";
29 };
30
31 build-system = [ setuptools ];
32
33 pythonRelaxDeps = [
34 "pytorch-lightning"
35 ];
36
37 dependencies = [
38 pytorch-lightning
39 torch
40 ];
41
42 # needed while lightning is installed as package `pytorch-lightning` rather than`lightning`:
43 env.PACKAGE_NAME = "pytorch";
44
45 nativeCheckInputs = [ pytestCheckHook ];
46 enabledTestPaths = [ "tests" ];
47 disabledTests = [
48 # AssertionError: assert 'lightning @ git+' in 'lightning>=2.5.0,<2.5.6'
49 "test_get_lightning_requirement"
50 ]
51 ++ lib.optionals (pythonAtLeast "3.14") [
52 # RuntimeError: torch.compile is not supported on Python 3.14+
53 "test_fts_dynamo_enforce_p0"
54 "test_fts_dynamo_resume"
55 ]
56 ++ lib.optionals (stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux) [
57 # slightly exceeds numerical tolerance on aarch64-linux:
58 "test_fts_frozen_bn_track_running_stats"
59 ];
60
61 pythonImportsCheck = [ "finetuning_scheduler" ];
62
63 __darwinAllowLocalNetworking = true;
64
65 meta = {
66 description = "PyTorch Lightning extension for foundation model experimentation with flexible fine-tuning schedules";
67 homepage = "https://finetuning-scheduler.readthedocs.io";
68 changelog = "https://github.com/speediedan/finetuning-scheduler/blob/${finalAttrs.src.tag}/CHANGELOG.md";
69 license = lib.licenses.asl20;
70 maintainers = with lib.maintainers; [ bcdarwin ];
71 };
72})