Skip to content

feat: allow addtion of python libs for testing in build config #150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from

Conversation

drbh
Copy link
Collaborator

@drbh drbh commented Jun 10, 2025

This PR adds the ability to include python libraries in the dev shell by specifying them in the build.toml file

Example usage:

[general]
name = "megablocks"
universal = false

[torch]
src = [
  "torch-ext/torch_binding.cpp",
  "torch-ext/torch_binding.h"
]

[kernel.megablocks]
backend = "cuda"
src = [
    "csrc/new_cumsum.h",
    "csrc/new_cumsum.cu",
    "csrc/new_histogram.h",
    "csrc/new_histogram.cu",
    "csrc/new_indices.h",
    "csrc/new_indices.cu",
    "csrc/new_replicate.cu",
    "csrc/new_replicate.h",
    "csrc/new_sort.h",
    "csrc/new_sort.cu",
]
depends = [ "torch", "cutlass_3_8" ]

[test]
python-git-packages = [
    { url = "https://github.com/stanford-futuredata/stk.git", rev = "7363137", sha256 = "0m6g5l9nlwaiwybg5j8dhnz159wdpabdnkzapnn3dsifxrsb59vz" }
]

This changes adds support for

  • python-git-packages for the addition of git repos
  • python_packages for libraries that are on nixpkgs

#[serde(deny_unknown_fields, rename_all = "kebab-case")]
pub struct Test {
#[serde(default)]
pub python_packages: Vec<String>,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I really like this idea!

Since this is only used in Nix-land, I think it would be nicer to pass optional Python test dependencies as an argument to genFlakeOutputs? Then also no special handling is needed for git packages, since it's Nix and you could plug arbitrary derivations in there.

I think build.toml isn't the best place, since it suggests generic Python test dependencies and there can be mismatches between Nix Python attribute names and PyPI package names, so they are not generally usable.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for the feedback I took a stab at handing all of this in nix land here: #153

@drbh
Copy link
Collaborator Author

drbh commented Jun 12, 2025

closing this PR in favor of #153

@drbh drbh closed this Jun 12, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants