feat: add nix build definitions #745

aws-nslick · 2024-12-09T23:02:19Z

Issue #, if available:

Description of changes:

This adds a nix flake for aws-ofi-nccl.

Usage:

To use the development shell, install direnv then create .envrc with contents:

watch_file nix/shell.nix
use_flake . --impure

This will create a .clangd file which can be used for IDE-like features atop clangd. Individual editor configuration varies. Example emacs configuration using doom emacs:

(use-package! eglot
  :config
  (add-to-list 'eglot-server-programs
               `((c-mode c-ts-mode c++-mode c++-ts-mode objc-mode)
                 . ("clangd"
                    "-j=8"
                    "-background-index-priority=normal"
                    "--enable-config"
                    "--all-scopes-completion"
                    "--header-insertion=iwyu"
                    "--header-insertion-decorators"
                    "--log=verbose"
                    "--clang-tidy"
                    "--background-index"
                    "--malloc-trim"
                    "--completion-style=detailed"
                    "--import-insertions"
                    "--pch-storage=memory"
                    "--function-arg-placeholders=1"))))

To build aws-ofi-nccl, libfabric, openmpi, etc.: nix build .#default.
To build wrappers for nccl-tests that run atop ubuntu, run: nix build .#ubuntu-test-runners. ./result/bin/ will contain shell scripts that LD_PRELOAD nvidia libraries from ubuntu expected locations, then forwards args to the nccl-test in question.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

bwbarrett

Nix isn't really my thing, but I have no objection to most of this patch (other than the Libfabric patch I called out below). Other than that, I'm not sure it's a competent review, but I'd be fine approving.

nix/overlays/libfabric/0001-cuda-dlopen-correctly.patch

Signed-off-by: Nicholas Sielicki <nslick@amazon.com>

aws-nslick · 2024-12-12T19:41:24Z

@ConnorBaker if you have any free time, any comments or criticisms are much appreciated.

ConnorBaker · 2024-12-17T07:04:33Z

nix/pkgs/aws-ofi-nccl/cleanSource.nix

+  dirs = {
+    third-party = ./../../../3rd-party;
+    docs = ./../../../doc;
+    headers = ./../../../include;
+    mfour = ./../../../m4;
+    nix = ./../../../nix;
+    tus = ./../../../src;
+    tests = ./../../../tests;
+    topologies = ./../../../topology;
+  };


Seems unused?

ConnorBaker · 2024-12-17T07:05:42Z

nix/pkgs/aws-ofi-nccl/cleanSource.nix

+  cleanRepo = traceVal (gitTracked ../../../.);
+  cleaned = x: intersection x (gitTracked ../../../.);
+  sourceFiles = cleaned (sourceFilter ../../../.);
+  buildFiles = cleaned (buildFileFilter ../../../.);
+  thirdPartyFiles = cleaned ../../../.;


Would recommend adding root = ../../.. to the let in block here and then using that instead of ../../../. (which is equivalent) everywhere. At the end of the let in you can then do inherit root; in the lib.fileset.toSource call.

ConnorBaker · 2024-12-17T07:08:02Z

nix/pkgs/aws-ofi-nccl/default.nix

+
+  cudaBuildDepsJoined = symlinkJoin {
+    name = "cuda-build-deps-joined";
+    paths = lib.optionals (cudaSupport) (


Suggested change

paths = lib.optionals (cudaSupport) (

paths = lib.optionals cudaSupport (

Unnecessary parenthesis.

ConnorBaker · 2024-12-17T07:09:41Z

nix/pkgs/aws-ofi-nccl/default.nix

+      [
+        (lib.getDev cudaPackages.cuda_nvcc)
+        cudaPackages.cuda_cudart.include
+      ]
+      ++ (
+        if effectiveStdenv.hostPlatform.isStatic then
+          [
+            (lib.getOutput "static" cudaPackages.cuda_cudart)
+          ]
+        else
+          [
+            (lib.getLib cudaPackages.cuda_cudart)
+          ]
+      )


Suggested change

[

(lib.getDev cudaPackages.cuda_nvcc)

cudaPackages.cuda_cudart.include

]

++ (

if effectiveStdenv.hostPlatform.isStatic then

[

(lib.getOutput "static" cudaPackages.cuda_cudart)

]

else

[

(lib.getLib cudaPackages.cuda_cudart)

]

)

[

(lib.getDev cudaPackages.cuda_nvcc)

cudaPackages.cuda_cudart.include

(

lib.getOutput

(if effectiveStdenv.hostPlatform.isStatic then "static" else "lib")

cudaPackages.cuda_cudart

)

]

ConnorBaker · 2024-12-17T07:22:35Z

flake.nix

+                (import ./nix/cudaPackagesExtensions/0001-add-latest-nccl.nix { inherit (pkgs) fetchFromGitHub; })
+                (import ./nix/cudaPackagesExtensions/0002-use-latest-nccl.nix)
+                (import ./nix/cudaPackagesExtensions/0003-nccl-tests-use-mpi.nix { inherit config; })
+                (import ./nix/cudaPackagesExtensions/0004-add-ncclAws.nix {
+                  inherit lib config;
+                  inherit (pkgs) symlinkJoin patchelf;
+                })
+                (import ./nix/cudaPackagesExtensions/0005-add-nccl-tests-aws.nix {
+                  inherit (pkgs) replaceDependency;
+                })


Why not push the arguments inside the extension? That is, instead of passing pkgs.fetchFromGitHub to the extension, you should be able to have something like

ffinal: let inherit (ffinal.pkgs) fetchFromGitHub; in pprev: ...

pkgs is a member of every cudaPackages package set; this would also prevent your pkgs from contaminating the overlays generated by overlayAttrs (because you're passing things from your pkgs to them directly instead of allowing the extension to retrieve them from whatever fixed point it is being evaluated within).

I'd recommend changing usages in the extensions of config.packages.blarg with ffinal.pkgs.blarg for the same reason; if you're using overlays properly, everything that you'd normally need to pass through directly will be available in the resulting Nixpkgs instantiation which is used to evaluate these extensions.

ConnorBaker · 2024-12-17T07:23:13Z

flake.nix

+              aws-ofi-nccl = (
+                pkgs.callPackage ./nix/pkgs/aws-ofi-nccl {
+                  inherit inputs self;
+                }
+              );


Suggested change

aws-ofi-nccl = (

pkgs.callPackage ./nix/pkgs/aws-ofi-nccl {

inherit inputs self;

}

);

aws-ofi-nccl = pkgs.callPackage ./nix/pkgs/aws-ofi-nccl {

inherit inputs self;

};

ConnorBaker · 2024-12-17T07:26:20Z

nix/cudaPackagesExtensions/0001-add-latest-nccl.nix

+      rev = "v2.23.4-1";
+      hash = "sha256-DlMxlLO2F079fBkhORNPVN/ASYiVIRfLJw7bDoiClHw=";
+    };
+    name = "cuda${ffinal.cudaMajorMinorPatchVersion}-nccl-2.23.4-1";


Is it necessary to change the name? Using backendStdenv, name should be the result of ${backendStdenv.cudaNamePrefix}-${pname}-${version} (https://github.com/ConnorBaker/cuda-packages/blob/bb6bf09b919e8e7dda006bcee9a48675d034850f/cuda-packages/common/backendStdenv.nix#L55) by default, so updating the version should be enough.

ConnorBaker · 2024-12-17T07:27:52Z

nix/cudaPackagesExtensions/0001-add-latest-nccl.nix

@@ -0,0 +1,13 @@
+{ fetchFromGitHub }:
+ffinal: pprev: {
+  nccl_latest = pprev.nccl.overrideAttrs (prevAttrs: {


It doesn't look like prevAttrs is being used -- remember that overrideAttrs takes three different types of argument:

attribute set (unconditional overrides)

function to attribute set (uses previous values)

function to function to attribute set (uses previous and final values)

ConnorBaker

Nix code looks good to me, left a few comments/nits -- do note that I'm not able to test it at the moment.

aws-nslick · 2024-12-17T17:44:06Z

Thank you for taking the time, I really appreciate it. Taking every change.

aws-nslick requested a review from a team as a code owner December 9, 2024 23:02

aws-nslick force-pushed the nix-for-master branch 3 times, most recently from 6bca13b to bdfd488 Compare December 12, 2024 17:14

bwbarrett requested changes Dec 12, 2024

View reviewed changes

nix/overlays/libfabric/0001-cuda-dlopen-correctly.patch Outdated Show resolved Hide resolved

feat: add nix build definitions

5e824fc

Signed-off-by: Nicholas Sielicki <nslick@amazon.com>

aws-nslick force-pushed the nix-for-master branch from bdfd488 to 5e824fc Compare December 12, 2024 19:37

aws-nslick requested a review from bwbarrett December 12, 2024 19:40

ConnorBaker reviewed Dec 17, 2024

View reviewed changes

ConnorBaker approved these changes Dec 17, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add nix build definitions #745

feat: add nix build definitions #745

aws-nslick commented Dec 9, 2024 •

edited

Loading

bwbarrett left a comment

aws-nslick commented Dec 12, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker Dec 17, 2024

ConnorBaker left a comment

aws-nslick commented Dec 17, 2024

	paths = lib.optionals (cudaSupport) (
	paths = lib.optionals cudaSupport (

feat: add nix build definitions #745

Are you sure you want to change the base?

feat: add nix build definitions #745

Conversation

aws-nslick commented Dec 9, 2024 • edited Loading

bwbarrett left a comment

Choose a reason for hiding this comment

aws-nslick commented Dec 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ConnorBaker left a comment

Choose a reason for hiding this comment

aws-nslick commented Dec 17, 2024

aws-nslick commented Dec 9, 2024 •

edited

Loading