pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 fetchpatch,
7
8 # build-system
9 setuptools,
10
11 # dependencies
12 pytorch-lightning,
13 torch,
14
15 # tests
16 pythonOlder,
17 pythonAtLeast,
18 pytestCheckHook,
19}:
20
21buildPythonPackage rec {
22 pname = "finetuning-scheduler";
23 version = "2.5.3";
24 pyproject = true;
25
26 src = fetchFromGitHub {
27 owner = "speediedan";
28 repo = "finetuning-scheduler";
29 tag = "v${version}";
30 hash = "sha256-6WRKDYug7eVaTSY2R2jBcj9o/984mqKZZi36XRT7KyI=";
31 };
32
33 patches = [
34 (fetchpatch {
35 url = "https://github.com/speediedan/finetuning-scheduler/commit/78e6e225f353d1ba95db05d7fc6ff541859ed6a2.patch";
36 hash = "sha256-7mbtsaHrnHph8lvuwhBGqxPQimbZcbGeyBYXzApFPn4=";
37 })
38 ];
39
40 postPatch = ''
41 substituteInPlace pyproject.toml \
42 --replace-fail "setuptools<77.0.0" "setuptools"
43 '';
44
45 build-system = [ setuptools ];
46
47 pythonRelaxDeps = [
48 "pytorch-lightning"
49 ];
50
51 dependencies = [
52 pytorch-lightning
53 torch
54 ];
55
56 # needed while lightning is installed as package `pytorch-lightning` rather than`lightning`:
57 env.PACKAGE_NAME = "pytorch";
58
59 nativeCheckInputs = [ pytestCheckHook ];
60 enabledTestPaths = [ "tests" ];
61 disabledTests =
62 lib.optionals (pythonOlder "3.12") [
63 # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
64 # LoweringException: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler'
65 "test_fts_dynamo_enforce_p0"
66 "test_fts_dynamo_resume"
67 "test_fts_dynamo_intrafit"
68 ]
69 ++ lib.optionals (pythonAtLeast "3.13") [
70 # RuntimeError: Dynamo is not supported on Python 3.13+
71 "test_fts_dynamo_enforce_p0"
72 "test_fts_dynamo_resume"
73 ]
74 ++ lib.optionals (stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux) [
75 # slightly exceeds numerical tolerance on aarch64-linux:
76 "test_fts_frozen_bn_track_running_stats"
77 ];
78
79 pythonImportsCheck = [ "finetuning_scheduler" ];
80
81 __darwinAllowLocalNetworking = true;
82
83 meta = {
84 description = "PyTorch Lightning extension for foundation model experimentation with flexible fine-tuning schedules";
85 homepage = "https://finetuning-scheduler.readthedocs.io";
86 changelog = "https://github.com/speediedan/finetuning-scheduler/blob/v${version}/CHANGELOG.md";
87 license = lib.licenses.asl20;
88 maintainers = with lib.maintainers; [ bcdarwin ];
89 };
90}