chore: optimize uv configuration

[0xaskr]

Summary

• Move pytest from dev extra to base dependencies so CI environments no longer need --extra dev
• Simplify optional dependency groups: remove unused torchvision, flash-linear-attention, and pin-free numpy
• Update project description to reflect broader scope beyond GLA
• Add uv settings in pyproject.toml (compile-bytecode, default-groups, python-downloads)
• Update CLAUDE.md install commands and K8s/CI templates to match new dependency structure

Test plan

• Verify uv sync installs pytest correctly
• Verify uv sync --extra gpu works without --extra dev
• Verify CI pipelines pass with updated config

🤖 Generated with Claude Code

Summary by CodeRabbit

• Documentation

  • Updated development setup instructions for clearer onboarding

• Chores

  • Simplified CI/build configuration for more efficient setup processes
  • Streamlined dependency management and removed unused optional features
  • Refreshed project description to reflect broader neural network architecture support
  • Updated GPU setup configuration for consistency

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request consolidates dependency management and CI setup to use uv more systematically. Changes include updating CI configuration to use uv sync and uv run instead of manual virtual environment creation, removing redundant dependencies (numpy, pytest from core), restructuring optional dependency extras, and updating project documentation and Kubernetes startup scripts accordingly.

Changes

Cohort / File(s)	Summary
CI and Setup Scripts .github/ci/gpu-tests.sky.yaml, scripts/launch-tpuv7.yml	Replaced manual Python virtual environment creation and activation with uv sync --extra gpu/tpu, and switched test execution to uv run python -m pytest for automated environment handling.
Documentation CLAUDE.md	Updated dependency installation instructions to reflect pytest inclusion by default and introduction of a new dev extra for development tools (ruff/pre-commit).
Package Configuration pyproject.toml	Broadened project description; removed numpy and pytest from base/dev dependencies; trimmed GPU extra specificity by removing torchvision and narrowing JAX/flash-attention extras; removed entire profile extra containing xprof.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• PR #102: Directly related as the original commit adding .github/ci/gpu-tests.sky.yaml CI configuration; this PR refactors that same file to use uv sync and uv run commands instead of manual venv setup.

Poem

🐰 With uv as our guide, dependencies align,
No more venvs to craft, no more scripts to twine,
Pytest joins the fold, extras trimmed so fine,
Our kernels run faster, architectures shine! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The title 'chore: optimize uv configuration' is vague and does not convey the specific, meaningful changes in the changeset. While technically related to dependency/configuration management, it fails to highlight the main changes: moving pytest to base dependencies, removing unused dependencies, updating project description, and consolidating uv setup across CI and K8s templates.	Consider a more specific title that captures the primary changes, such as 'chore: consolidate uv setup and simplify dependencies' or 'refactor: restructure dependencies and unify uv configuration across CI/K8s' to better communicate the scope of modifications.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch chore/optimize-uv-config

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request streamlines the project's dependency management and CI/CD workflows by adopting uv sync and uv run. Key changes include updating pyproject.toml to refine dependency groups, removing unused packages like torchvision, and updating the TPU launch script and GPU test configuration. Feedback highlights that moving pytest to the base dependencies is non-standard for a library. Additionally, removing the [gpu] extra from flash-linear-attention may cause issues with Triton-based kernels, and the TPU launch script needs to use uv run to correctly access the virtual environment created by uv sync.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

scripts/launch-tpuv7.yml (1)

36-39: The pip fallback may not work as expected after a failed uv sync.

If uv sync --extra tpu fails partway through (e.g., due to a dependency resolution error), it may have created a partial .venv directory. The subsequent pip install -e '.[tpu]' runs against the system Python, not the venv, which could lead to an inconsistent state where dependencies are split between system and venv.

♻️ Consider a cleaner fallback strategy

          # 4. Clone and setup project
          git clone git@github.com:primatrix/pallas-kernel.git ~/pallas-kernel
          cd ~/pallas-kernel
-          uv sync --extra tpu || pip install -e '.[tpu]'
+          if ! uv sync --extra tpu; then
+            echo "uv sync failed, falling back to pip"
+            rm -rf .venv  # Clean up partial state
+            pip install -e '.[tpu]'
+          fi

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@scripts/launch-tpuv7.yml` around lines 36 - 39, The fallback runs pip with
system Python after a possibly partially-created virtualenv from `uv sync
--extra tpu`; change the logic to check the exit status of `uv sync --extra tpu`
and, on failure, remove any partial `.venv`, create and use an explicit
virtualenv, and run the fallback install inside that venv (use `python -m venv
.venv` then the venv's pip or activate the venv) so `pip install -e '.[tpu]'`
installs into the venv rather than the system Python; reference the `uv sync
--extra tpu` command and the `pip install -e '.[tpu]'` fallback when making the
change.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@scripts/launch-tpuv7.yml`:
- Around line 33-39: The uv sync step currently creates a project virtualenv but
later verification runs use system Python; update the workflow so dependencies
are executed inside the venv: run uv sync --extra tpu (optionally add --no-dev
to avoid dev deps) and replace the standalone python verification command with
uv run python -c '...' (or ensure uv creates a predictable venv path and
activate it before running verification). Target the uv invocation lines (uv
sync and the subsequent verification python command) and ensure you either add
--no-dev to uv sync or use uv run so the verification uses the virtualenv
created by uv.

---

Nitpick comments:
In `@scripts/launch-tpuv7.yml`:
- Around line 36-39: The fallback runs pip with system Python after a possibly
partially-created virtualenv from `uv sync --extra tpu`; change the logic to
check the exit status of `uv sync --extra tpu` and, on failure, remove any
partial `.venv`, create and use an explicit virtualenv, and run the fallback
install inside that venv (use `python -m venv .venv` then the venv's pip or
activate the venv) so `pip install -e '.[tpu]'` installs into the venv rather
than the system Python; reference the `uv sync --extra tpu` command and the `pip
install -e '.[tpu]'` fallback when making the change.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f6d62ed4-3859-4157-92ce-8120faca40c4

📥 Commits

Reviewing files that changed from the base of the PR and between 3942b75 and e92eec7.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (4)

• .github/ci/gpu-tests.sky.yaml
• CLAUDE.md
• pyproject.toml
• scripts/launch-tpuv7.yml

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Consider adding --no-dev flag or ensure uv creates the venv in a predictable location.

The uv sync --extra tpu command will create a virtual environment (typically .venv/ in the project directory). However, the subsequent python -c 'import jax; ...' on line 41 runs without activating this venv or using uv run, so it will use the system Python which won't have the dependencies installed.

🐛 Proposed fix to use uv run for the verification command

          # 4. Clone and setup project
          git clone git@github.com:primatrix/pallas-kernel.git ~/pallas-kernel
          cd ~/pallas-kernel
          uv sync --extra tpu || pip install -e '.[tpu]'
 
-          python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
+          uv run python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@scripts/launch-tpuv7.yml` around lines 33 - 39, The uv sync step currently
creates a project virtualenv but later verification runs use system Python;
update the workflow so dependencies are executed inside the venv: run uv sync
--extra tpu (optionally add --no-dev to avoid dev deps) and replace the
standalone python verification command with uv run python -c '...' (or ensure uv
creates a predictable venv path and activate it before running verification).
Target the uv invocation lines (uv sync and the subsequent verification python
command) and ensure you either add --no-dev to uv sync or use uv run so the
verification uses the virtualenv created by uv.