1 files changed, 88 insertions, 27 deletions

diff --git a/pkgs/development/python-modules/ffcv/default.nix b/pkgs/development/python-modules/ffcv/default.nix
index a3dd531c1285..158050a9af77 100644
--- a/pkgs/development/python-modules/ffcv/default.nix
+++ b/pkgs/development/python-modules/ffcv/default.nix
@@ -1,20 +1,32 @@
-{ buildPythonPackage
-, fetchFromGitHub
-, lib
-, libjpeg
-, numba
-, opencv4
-, pandas
-, pkg-config
-, pytorch-pfn-extras
-, terminaltables
-, tqdm
+{
+  buildPythonPackage,
+  pythonAtLeast,
+  fetchFromGitHub,
+  setuptools,
+  lib,
+  libjpeg,
+  numba,
+  opencv-python,
+  pandas,
+  pkg-config,
+  pytorch-pfn-extras,
+  terminaltables,
+  tqdm,
+  pytestCheckHook,
+  assertpy,
+  psutil,
+  torchvision,
+  webdataset,
+  stdenv,
 }:
 
 buildPythonPackage rec {
   pname = "ffcv";
   version = "1.0.0";
-  format = "setuptools";
+  pyproject = true;
+
+  # version 1.0.0 uses distutils which was removed in Python 3.12
+  disabled = pythonAtLeast "3.12";
 
   src = fetchFromGitHub {
     owner = "libffcv";
@@ -26,29 +38,78 @@ buildPythonPackage rec {
   # See https://github.com/libffcv/ffcv/issues/159.
   postPatch = ''
     substituteInPlace setup.py \
-      --replace "'assertpy'," "" \
-      --replace "'fastargs'," "" \
-      --replace "'imgcat'," "" \
-      --replace "'matplotlib'," "" \
-      --replace "'psutil'," "" \
-      --replace "'sklearn'," "" \
-      --replace "'webdataset'," ""
+      --replace-fail "'assertpy'," "" \
+      --replace-fail "'fastargs'," "" \
+      --replace-fail "'psutil'," "" \
   '';
 
+  build-system = [ setuptools ];
   nativeBuildInputs = [ pkg-config ];
   buildInputs = [ libjpeg ];
-  propagatedBuildInputs = [ opencv4 numba pandas pytorch-pfn-extras terminaltables tqdm ];
-
-  # `ffcv._libffcv*.so` cannot be loaded in the nix build environment for some
-  # reason. See https://github.com/NixOS/nixpkgs/pull/160441#issuecomment-1045204722.
-  doCheck = false;
+  propagatedBuildInputs = [
+    opencv-python
+    numba
+    pandas
+    pytorch-pfn-extras
+    terminaltables
+    tqdm
+  ];
 
   pythonImportsCheck = [ "ffcv" ];
 
-  meta = with lib; {
+  # C/C++ python modules are only in the installed output and not in the build
+  # directory. Since tests are run from the build directory python prefers to
+  # import the local module first which does not contain the C/C++ python
+  # modules and results in an import error. By changing the directory to
+  # 'tests' the build directory is no long available and python will import
+  # from the installed output in the nix store which does contain the C/C++
+  # python modules.
+  preCheck = ''
+    cd tests
+  '';
+
+  nativeCheckInputs = [
+    assertpy
+    psutil
+    pytestCheckHook
+    torchvision
+    webdataset
+  ];
+
+  disabledTestPaths = [
+    # Tests require network access and do not work in the sandbox
+    "test_augmentations.py"
+    # Occasionally causes the testing phase to hang
+    "test_basic_pipeline.py"
+  ];
+
+  disabledTests = [
+    # Tests require network access and do not work in the sandbox
+    "test_cifar_subset"
+    # Requires CUDA which is unfree and unfree packages are not built by Hydra
+    "test_cuda"
+    "test_gpu_normalization"
+    # torch.multiprocessing.spawn.ProcessRaisedException
+    "test_traversal_sequential_2"
+    "test_traversal_sequential_3"
+    "test_traversal_sequential_4"
+    "test_traversal_random_2"
+    "test_traversal_random_3"
+    "test_traversal_random_4"
+    "test_traversal_sequential_distributed_with_indices"
+    "test_traversal_random_distributed_with_indices"
+  ];
+
+  meta = {
     description = "FFCV: Fast Forward Computer Vision";
     homepage = "https://ffcv.io";
-    license = licenses.asl20;
-    maintainers = with maintainers; [ samuela ];
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [
+      samuela
+      djacu
+    ];
+    # OSError: dlopen(libc.so.6, 0x0006): tried: '/usr/lib/libc.so.6' (no such file, not in dyld cache),
+    # 'libc.so.6' (no such file), '/usr/local/lib/libc.so.6' (no such file), '/usr/lib/libc.so.6' (no such file, not in dyld cache)
+    broken = stdenv.hostPlatform.isDarwin;
   };
 }