Skip to content

installation.md

Latest commit

 

History

History
161 lines (107 loc) · 3.79 KB

installation.md

File metadata and controls

161 lines (107 loc) · 3.79 KB

Installation

This guide covers how to install NKIPy and its dependencies on an AWS Trainium instance.

Prerequisites

NKIPy requires a Trainium instance with the Neuron Driver and Runtime installed.

Install Neuron Driver and Runtime

If you are using a Neuron Multi-Framework DLAMI, the driver and runtime are already installed. You can skip to the next section.

Otherwise, follow the Neuron Setup Guide up to the "Install Drivers and Tools" section for your OS. Note that NKIPy does not require PyTorch, but it supports Torch tensors if available.

Installation with uv (Recommended)

uv is a fast Python package manager. It's the recommended way to install NKIPy.

Install uv

curl -LsSf https://astral.sh/uv/install.sh | sh

Clone and Install

git clone https://github.com/aws-neuron/nkipy.git
cd nkipy

# Install all packages (creates .venv automatically)
uv sync

This single command will:

  • Create a .venv virtual environment
  • Install nkipy and spike in editable mode
  • Install all dependencies including neuronx-cc from the Neuron repository
  • Install development tools (pytest, ruff, mypy)

Running Commands

You can run commands in two ways:

Option 1: Activate the virtual environment

# Activate the environment
source .venv/bin/activate

# Now run commands directly
python examples/playground/simple_nkipy_kernel.py
python -c "import nkipy; import spike"

# Deactivate when done
deactivate

Option 2: Use uv run without activation

# Run Python with the workspace environment
uv run python examples/playground/simple_nkipy_kernel.py 

# Start an interactive Python session
uv run python

Building Wheels

# Build nkipy
uv build --package nkipy

# Build spike
uv build --package spike

# Output will be in dist/
ls dist/

Alternative: Installation with pip

If you prefer using pip instead of uv, follow these instructions.

Create a Virtual Environment

python3.10 -m venv nkipy_venv
source nkipy_venv/bin/activate

Configure Neuron Repository

NKIPy depends on the Neuron Compiler (neuronx-cc). Configure pip to use the Neuron repository:

pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

Install Packages

NKIPy and Spike are not published on PyPI and must be installed from source using the local directory paths:

# Install both packages in editable mode
pip install -e ./nkipy -e ./spike

# Or install without editable mode
pip install ./nkipy ./spike

Building Wheels with pip

# Install build tool
pip install build

# Build wheel packages
python -m build nkipy --wheel --outdir dist
python -m build spike --wheel --outdir dist

# Install the built wheels
pip install dist/*.whl

Verifying Installation

After installation, verify that everything works:

# With uv
uv run python -c "import nkipy; import spike; print('Installation successful!')"

# With pip (after activating venv)
python -c "import nkipy; import spike; print('Installation successful!')"

Troubleshooting

neuronx-cc not found

If you see an error about neuronx-cc not being found:

With uv: The Neuron repository is already configured in pyproject.toml. Try running uv sync --refresh.

With pip: Make sure you've configured the extra index URL:

pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

spike build fails

The spike package contains C++ code that requires compilation. Ensure you have:

  • CMake installed (apt install cmake)
  • A C++ compiler (gcc/clang)
  • Python development headers (apt install python3-dev)
  • Neuron Runtime (libnrt) installed