diff --git a/pkgs/development/python-modules/nvdlfw-inspect/default.nix b/pkgs/development/python-modules/nvdlfw-inspect/default.nix new file mode 100644 index 000000000000..21a9be48d090 --- /dev/null +++ b/pkgs/development/python-modules/nvdlfw-inspect/default.nix @@ -0,0 +1,44 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # dependencies + pyyaml, + torch, + + # tests + pytestCheckHook, +}: + +buildPythonPackage (finalAttrs: { + pname = "nvdlfw-inspect"; + version = "0.2.2"; + pyproject = true; + + src = fetchFromGitHub { + owner = "NVIDIA"; + repo = "nvidia-dlfw-inspect"; + # No tag, found the commit using `git blame` + rev = "3a7314ac44f8479a3368a69ac4c7938e40e6104c"; + hash = "sha256-qLyPzdwIIXHUDeP226w4yo2B1fqOKZ3yeSXrG3pNQyk="; + }; + + dependencies = [ + pyyaml + torch + ]; + + pythonImportsCheck = [ "nvdlfw_inspect" ]; + + nativeCheckInputs = [ + pytestCheckHook + ]; + + meta = { + description = "Facilitates debugging convergence issues and testing new algorithms/recipes for training LLMs using Nvidia libraries"; + homepage = "https://github.com/NVIDIA/nvidia-dlfw-inspect"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 77401c0a50f6..759ae0c51e3e 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11074,6 +11074,8 @@ self: super: with self; { nvchecker = callPackage ../development/python-modules/nvchecker { }; + nvdlfw-inspect = callPackage ../development/python-modules/nvdlfw-inspect { }; + nvdlib = callPackage ../development/python-modules/nvdlib { }; nvidia-dlprof-pytorch-nvtx =