Migrate tests to pytest + add GA

.github/workflows/pr-gpu.yaml

@@ -0,0 +1,46 @@
+name: PR GPU tests
+on:
+  push:
+    branches:
+    - main
+    - dev
+    - release/*
+    - eitan-tests  # todo: remove this branch
+  pull_request_target:
+    branches:
+    - main
+    - dev
+    - release/**
+    - eitan-tests  # todo: remove this branch
+  workflow_dispatch:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' && github.ref != 'refs/heads/dev' }}
+jobs:
+  gpu-pytests:
+    uses: mosaicml/ci-testing/.github/workflows/[email protected]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - name: "python3.11-pytorch2.3-gpus1"
+          gpu_num: 1
+          container: mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
+        - name: "python3.11-pytorch2.3-gpus8"
+          gpu_num: 8
+          container: mosaicml/pytorch:2.3.0_cu121-python3.11-ubuntu20.04
+    name: ${{ matrix.name }}
+    if: github.repository_owner == 'eitanturok'  # todo: replace with `databricks`
+    with:
+      container: ${{ matrix.container }}
+      git_repo: eitanturok/megablocks  # todo: replace with `databricks/megablocks`
+      mcloud-timeout: 3600
+      name: ${{ matrix.name }}
+      pip_deps: "[all]"
+      pytest-command: "coverage run -m pytest tests"
+      pytest-markers: "gpu"
+      python-version: 3.9
+      composer_package_name: mosaicml  # Required as Composer is built from source
+      gpu_num: ${{ matrix.gpu_num }}
+    secrets:
+      mcloud-api-key: ${{ secrets.MCLOUD_API_KEY }}

.gitignore

@@ -1,8 +1,165 @@
-*~
-gpt2-merges.txt
-gpt2-vocab.json
-core
-__pycache__
-*.egg-info
-dist
-build
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+docs/path/
+docs/source/_build/
+docs/source/api_reference/
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Virtual Env
+venv/
+
+# WandB
+wandb/
+
+# Neptune
+.neptune/
+
+# Spacemacs
+._#*
+.#*
+.vscode/
+
+# Vim
+*.swp
+
+# README backups
+README.md.bkp
+
+# OS X
+.DS_Store
+
+# Sphinx
+api_reference/
+
+# Node (for pyright)
+node_modules/
+package.json
+package-lock.json
+
+# pycharm
+.idea/

Makefile

@@ -1,11 +1,23 @@
-clean:
-	rm -rf dist/*
+# several pytest settings
+WORLD_SIZE ?= 1  # world size for launcher tests
+MASTER_PORT ?= 26000 # port for distributed tests
+PYTHON ?= python3  # Python command
+PYTEST ?= pytest  # Pytest command
+PYRIGHT ?= pyright  # Pyright command. Pyright must be installed seperately -- e.g. `node install -g pyright`
+EXTRA_ARGS ?=  # extra arguments for pytest
+EXTRA_LAUNCHER_ARGS ?= # extra arguments for the composer cli launcher
 
-dist: clean
-	python3 setup.py sdist
+test:
+	LOCAL_WORLD_SIZE=1 $(PYTHON) -m $(PYTEST) $(EXTRA_ARGS)
 
-upload: dist
-	twine upload dist/*
+test-gpu:
+	LOCAL_WORLD_SIZE=1 $(PYTHON) -m $(PYTEST) -m gpu $(EXTRA_ARGS)
 
-upload-test: dist
-	twine upload --repository-url https://test.pypi.org/legacy/ dist/*
+# runs tests with the launcher
+test-dist:
+	$(PYTHON) -m composer.cli.launcher -n $(WORLD_SIZE) --master_port $(MASTER_PORT) $(EXTRA_LAUNCHER_ARGS) -m $(PYTEST) $(EXTRA_ARGS)
+
+test-dist-gpu:
+	$(PYTHON) -m composer.cli.launcher -n $(WORLD_SIZE) --master_port $(MASTER_PORT) $(EXTRA_LAUNCHER_ARGS) -m $(PYTEST) -m gpu $(EXTRA_ARGS)
+
+.PHONY: test test-gpu test-dist test-dist-gpu

pyproject.toml

@@ -0,0 +1,48 @@
+# build requirements
+[build-system]
+requires = ["setuptools < 70.0.0", "torch >= 2.3.0, < 2.4"]
+build-backend = "setuptools.build_meta"
+
+# Pytest
+[tool.pytest.ini_options]
+# By default, skip gpu tests
+addopts = "--tb=short -m 'not gpu'"
+
+markers = [
+    # For distributed testing
+    "world_size(val)",
+    # Should be run during daily regression
+    "daily",
+    # Whether the test will be reading data from a remote source, and may require credentials
+    "remote",
+    # whether the test requires a gpu
+    "gpu",
+]
+
+filterwarnings = [
+    # "error",  # warnings should be treated like errors, but still need to fix some warnings
+    'ignore:ExtraArgumentWarning',  # extra arguments originate from pytest-specific CLI args
+    'ignore:DistributedDefaultValueWarning',  # default distributed values are fine
+    'ignore:NoDistributedWarning',  # running without distributed is fine
+    'ignore:Deterministic mode is activated:UserWarning',  # all tests run with deterministic mode
+    'ignore:SubsetNumBatchesWarning',  # different subsets OK for testing
+    'ignore:No optimizer:UserWarning',  # testing defaults
+    'ignore:No scheduler:UserWarning',  # testing defaults
+    'ignore::DeprecationWarning:tensorboard',  # ignore tensorboard
+]
+
+# Enable logging for pytest
+log_cli = true
+log_cli_level = "INFO"
+log_cli_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)"
+log_cli_date_format = "%Y-%m-%d %H:%M:%S"
+
+# Coverage
+[tool.coverage.run]
+parallel = true
+branch = true
+relative_files = true
+concurrency = ["thread"]
+include = [
+    "megablocks/*"
+]

setup.py

@@ -1,6 +1,7 @@
-from setuptools import setup, find_packages
 import os
+
 import torch
+from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 if os.environ.get("TORCH_CUDA_ARCH_LIST"):
@@ -28,19 +29,28 @@
     )
 ]
 
-install_requires=[
-    "triton>=2.1.0",
-    "stanford-stk==0.7.0",
+install_requires = [
+    'numpy>=1.21.5,<2.1.0',
+    'torch>=2.3.0,<2.4',
+    'triton>=2.1.0',
+    'stanford-stk @ git+https://[email protected]/stanford-futuredata/stk.git@a1ddf98466730b88a2988860a9d8000fd1833301',
+    'packaging>=21.3.0,<24.2',
 ]
 
 extra_deps = {}
 
 extra_deps["gg"] = [
-    "grouped_gemm==0.1.4",
+    'grouped_gemm @ git+https://[email protected]/tgale96/grouped_gemm.git@66c7195e35e8c4f22fa6a014037ef511bfa397cb',
 ]
 
-extra_deps["dev"] = [
-    "absl-py",
+extra_deps['dev'] = [
+    'absl-py',
+    'coverage[toml]==7.4.4',
+    'pytest_codeblocks>=0.16.1,<0.17',
+    'pytest-cov>=4,<5',
+    'pytest>=7.2.1,<8',
+    'pre-commit>=3.4.0,<4',
+    'mosaicml>=0.22.0',
 ]
 
 extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)