From 2b47e104e006e898605984a0de0e1e2b7e61d72a Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Thu, 22 Aug 2024 11:06:40 -0400
Subject: [PATCH] fix requirements + aggregate setup & req. & config in
 pyproject (#453)

---
 .github/azure-gpu-tests.yml     |  2 +-
 .github/workflows/cpu-tests.yml | 37 +++++++++++++++++---------
 .gitignore                      |  4 ++-
 README.md                       |  2 +-
 howto/tpus.md                   |  2 +-
 pyproject.toml                  | 46 +++++++++++++++++++++++++++++++++
 requirements.txt                |  9 -------
 setup.py                        | 27 ++-----------------
 tests/test_model.py             |  2 +-
 9 files changed, 80 insertions(+), 51 deletions(-)
 create mode 100644 pyproject.toml
 delete mode 100644 requirements.txt

diff --git a/.github/azure-gpu-tests.yml b/.github/azure-gpu-tests.yml
index 944b91be..0389c957 100644
--- a/.github/azure-gpu-tests.yml
+++ b/.github/azure-gpu-tests.yml
@@ -50,7 +50,7 @@ jobs:
         python -c "import torch ; mgpu = torch.cuda.device_count() ; assert mgpu == 2, f'GPU: {mgpu}'"
       displayName: 'Image info & NVIDIA'
 
-    - script: pip install pytest -r requirements.txt
+    - script: pip install ".[all]" "pytest"
       displayName: 'Install dependencies'
 
     - bash: pytest -v --durations=10 --disable-pytest-warnings --strict-markers --color=yes
diff --git a/.github/workflows/cpu-tests.yml b/.github/workflows/cpu-tests.yml
index 09fcc930..6df676ac 100644
--- a/.github/workflows/cpu-tests.yml
+++ b/.github/workflows/cpu-tests.yml
@@ -15,15 +15,14 @@ defaults:
     shell: bash
 
 jobs:
-  cpu-tests:
+  pytester:
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        include:
-          - {os: "macOS-11", python-version: "3.10"}
-          - {os: "ubuntu-20.04", python-version: "3.10"}
-          - {os: "windows-2022", python-version: "3.10"}
+        os: ["ubuntu-22.04", "macos-13", "windows-2022"]
+        python-version: ["3.10"]
+        pkg-install: ["no", "yes"]
     timeout-minutes: 15
 
     steps:
@@ -35,18 +34,32 @@ jobs:
         python-version: ${{ matrix.python-version }}
         cache: 'pip'
         cache-dependency-path: |
-          requirements.txt
+          pyproject.toml
           setup.py
 
-    - name: Run tests without the package installed
+    - name: Install package & dependencies
       run: |
-        pip install pytest -r requirements.txt
+        pip install ".[all]" pytest
         pip list
 
-        pytest --disable-pytest-warnings --strict-markers --color=yes
+    - name: Drop package itself
+      if: matrix.pkg-install == 'no'
+      run: pip uninstall -y lit-llama
 
     - name: Run tests
-      run: |
-        pip install . --no-deps
+      run: pytest -v --durations=10
+
 
-        pytest -v --durations=10 --disable-pytest-warnings --strict-markers --color=yes
+  testing-guardian:
+    runs-on: ubuntu-latest
+    needs: pytester
+    if: always()
+    steps:
+      - run: echo "${{ needs.pytester.result }}"
+      - name: failing...
+        if: needs.pytester.result == 'failure'
+        run: exit 1
+      - name: cancelled or skipped...
+        if: contains(fromJSON('["cancelled", "skipped"]'), needs.pytester.result)
+        timeout-minutes: 1
+        run: sleep 90
diff --git a/.gitignore b/.gitignore
index e8438525..d0f346c9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,4 +13,6 @@ wandb
 
 # downloaded by our tests
 original_model.py
-original_adapter.py
\ No newline at end of file
+original_adapter.py
+
+.ruff_cache/
\ No newline at end of file
diff --git a/README.md b/README.md
index 6ba807f3..48d6faee 100644
--- a/README.md
+++ b/README.md
@@ -67,7 +67,7 @@ cd lit-llama
 install dependencies
 
 ```bash
-pip install -r requirements.txt
+pip install -e ".[all]"
 ```
 
 You are all set! 🎉
diff --git a/howto/tpus.md b/howto/tpus.md
index 55749d78..c3fdc365 100644
--- a/howto/tpus.md
+++ b/howto/tpus.md
@@ -14,7 +14,7 @@ Now that you are in the machine, let's clone the repository and install the depe
 ```shell
 git clone https://github.com/Lightning-AI/lit-llama
 cd lit-llama
-pip install -r requirements.txt
+pip install -e ".[all]"
 ```
 
 By default, computations will run using the new (and experimental) PjRT runtime. Still, it's recommended that you set the following environment variables
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..72c1aa8e
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "lit-llama"
+version = "0.1.0"
+description = "Implementation of the LLaMA language model"
+license = {text = "Apache-2.0"}
+authors = [
+    { name = "Lightning AI", email = "community@lightning.ai" }
+]
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.0.0",
+    "lightning @ git+https://github.com/Lightning-AI/lightning@master",
+    "sentencepiece",
+    "bitsandbytes",
+]
+classifiers = [
+    "Topic :: Text Processing"
+]
+
+[project.optional-dependencies]
+all = [
+    "tqdm",  # convert_checkpoint.py
+    "numpy <2.0",  # train.py dataset memmap
+    "jsonargparse[signatures]",  # generate.py, convert_checkpoint.py CLI
+    "datasets",  # evaluate.py
+    "zstandard",  # prepare_redpajama.py"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]  # list of folders that contain the packages (["."] by default)
+include = ["lit_llama"]  # package names should match these glob patterns (["*"] by default)
+exclude = []  # exclude packages matching these glob patterns (empty by default)
+namespaces = false  # to disable scanning PEP 420 namespaces (true by default)
+
+
+[tool.pytest.ini_options]
+addopts = [
+    "--strict-markers",
+    "--color=yes",
+    "--disable-pytest-warnings",
+]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index e8706b99..00000000
--- a/requirements.txt
+++ /dev/null
@@ -1,9 +0,0 @@
-torch>=2.0.0
-lightning @ git+https://github.com/Lightning-AI/lightning@master
-sentencepiece
-tqdm  # convert_checkpoint.py
-numpy  # train.py dataset memmap
-jsonargparse[signatures]  # generate.py, convert_checkpoint.py CLI
-bitsandbytes  # quantization.py
-datasets  # evaluate.py
-zstandard  # prepare_redpajama.py
diff --git a/setup.py b/setup.py
index 633eb01c..fd417ee2 100644
--- a/setup.py
+++ b/setup.py
@@ -1,28 +1,5 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-import os
+from setuptools import setup
 
-from setuptools import setup, find_packages
-
-
-_PATH_ROOT = os.path.dirname(__file__)
-
-with open(os.path.join(_PATH_ROOT, "README.md"), encoding="utf-8") as fo:
-    readme = fo.read()
-
-setup(
-    name='lit-llama',
-    version='0.1.0',
-    description='Implementation of the LLaMA language model',
-    author='Lightning AI',
-    url='https://github.com/lightning-AI/lit-llama',
-    install_requires=[
-        "torch>=2.0.0",
-        "lightning @ git+https://github.com/Lightning-AI/lightning@master",
-        "sentencepiece",
-        "bitsandbytes",
-    ],
-    packages=find_packages(),
-    long_description=readme,
-    long_description_content_type="text/markdown",
-)
+setup()
diff --git a/tests/test_model.py b/tests/test_model.py
index 22ad6e6e..43b6ff3d 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -212,7 +212,7 @@ def test_adapter_parity(orig_llama_adapter):
 
     expected = orig_llama_model(token_sample, 0)
     out = llama_model(token_sample)
-    assert torch.allclose(out, expected)
+    assert torch.allclose(out, expected, atol=1e-5, rtol=1e-5)
 
 
 @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="torch.compile not supported on this platform")