1 files changed, 120 insertions, 0 deletions
diff --git a/flake.nix b/flake.nix
new file mode 100644
index 0000000..86d2968
--- /dev/null
+++ b/flake.nix
@@ -0,0 +1,120 @@
+{
+  description = "Poetry Python dev shell with CUDA + common native libs";
+
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+  };
+
+  outputs = { self, nixpkgs }:
+    let
+      systems = [
+        "x86_64-linux"
+        "aarch64-linux"
+      ];
+
+      forAllSystems = f:
+        nixpkgs.lib.genAttrs systems (system: f system);
+    in
+    {
+      devShells = forAllSystems (system:
+        let
+          pkgs = import nixpkgs {
+            inherit system;
+            config = {
+              allowUnfree = true;
+              cudaSupport = true;
+            };
+          };
+
+          python = pkgs.python312;
+
+          runtimeLibs = with pkgs; [
+            stdenv.cc.cc.lib
+
+            zlib
+            openssl
+            libffi
+            sqlite
+            xz
+            bzip2
+            readline
+            ncurses
+            expat
+
+            openblas
+
+            glib
+            libGL
+            libxkbcommon
+            xorg.libX11
+            xorg.libXext
+            xorg.libXrender
+            xorg.libXi
+            xorg.libXrandr
+            xorg.libXcursor
+            xorg.libXfixes
+            xorg.libXinerama
+            alsa-lib
+            ffmpeg
+
+            cudaPackages.cudatoolkit
+            cudaPackages.cudnn
+          ];
+
+          libPath = pkgs.lib.makeLibraryPath runtimeLibs;
+          dynamicLinker = pkgs.lib.fileContents "${pkgs.stdenv.cc}/nix-support/dynamic-linker";
+        in
+        {
+          default = pkgs.mkShell {
+            packages = with pkgs; [
+              python
+              poetry
+
+              # build tools for packages with native extensions
+              pkg-config
+              gcc
+              gfortran
+              gnumake
+              cmake
+              ninja
+              patchelf
+              git
+            ] ++ runtimeLibs;
+
+            NIX_LD = dynamicLinker;
+            NIX_LD_LIBRARY_PATH = libPath;
+
+            LD_LIBRARY_PATH = libPath + ":/run/opengl-driver/lib:/run/opengl-driver-32/lib";
+
+            CUDA_PATH = "${pkgs.cudaPackages.cudatoolkit}";
+            CUDA_HOME = "${pkgs.cudaPackages.cudatoolkit}";
+
+            POETRY_VIRTUALENVS_CREATE = "true";
+            POETRY_VIRTUALENVS_IN_PROJECT = "true";
+
+            shellHook = ''
+              export PATH="${python}/bin:$PATH"
+
+              if [ -f pyproject.toml ]; then
+                poetry env use ${python}/bin/python >/dev/null 2>&1 || true
+
+                # Uncomment this block if you want first entry into nix develop
+                # to automatically install the project deps.
+                #
+                # if [ ! -d .venv ]; then
+                #   echo "No .venv found; running poetry install..."
+                #   poetry install
+                # fi
+
+                echo
+                echo "Python: $(which python)"
+                echo "Poetry env: $(poetry env info --path 2>/dev/null || true)"
+                echo "Run once: poetry install"
+                echo "Then:     poetry run python3 main.py"
+                echo
+              fi
+            '';
+          };
+        });
+    };
+}