diff --git a/.gitignore b/.gitignore index 067ef89..b921639 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,46 @@ *.pyc *.egg-info -build/ \ No newline at end of file +build/ + +# Testing +.pytest_cache/ +.coverage +htmlcov/ +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.tox/ +.nox/ + +# Claude +.claude/* + +# Virtual environments +venv/ +ENV/ +env/ +.venv/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ +.DS_Store + +# Distribution / packaging +dist/ +*.egg +.eggs/ +*.whl + +# Cache +__pycache__/ +*.pyo +*.pyd +.Python + +# Logs +*.log \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..140ed72 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,81 @@ +[tool.poetry] +name = "jetmoe" +version = "0.1.0" +description = "JetMoE: Reaching LLaMA2 Performance with 0.1M Dollars" +authors = ["JetMoE Team"] +readme = "README.md" +license = "MIT" +homepage = "https://research.myshell.ai/jetmoe" +repository = "https://github.com/myshell-ai/JetMoE" +packages = [{include = "jetmoe"}] + +[tool.poetry.dependencies] +python = "^3.10.10" +torch = "^2.0.0" +transformers = "^4.30.0" +scattermoe = {git = "https://github.com/shawntan/scattermoe.git", branch = "main"} + +[tool.poetry.group.dev.dependencies] +pytest = "^7.4.0" +pytest-cov = "^4.1.0" +pytest-mock = "^3.11.1" + +[tool.poetry.scripts] +test = "pytest:main" +tests = "pytest:main" + +[tool.pytest.ini_options] +minversion = "7.0" +testpaths = ["tests"] +python_files = ["test_*.py", "*_test.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "-ra", + "--strict-markers", + "--cov=jetmoe", + "--cov-branch", + "--cov-report=term-missing:skip-covered", + "--cov-report=html", + "--cov-report=xml", + "--cov-fail-under=0", # Set to 80 once you have actual tests + "-v" +] +markers = [ + "unit: Unit tests", + "integration: Integration tests", + "slow: Slow tests" +] + +[tool.coverage.run] +source = ["jetmoe"] +omit = [ + "*/tests/*", + "*/__init__.py", + "*/conftest.py", + "setup.py" +] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "if __name__ == .__main__.:", + "raise AssertionError", + "raise NotImplementedError", + "if TYPE_CHECKING:", + "pass" +] +precision = 2 +show_missing = true +skip_covered = true + +[tool.coverage.html] +directory = "htmlcov" + +[tool.coverage.xml] +output = "coverage.xml" + +[build-system] +requires = ["poetry-core>=1.0.0"] +build-backend = "poetry.core.masonry.api" \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..b67e047 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,101 @@ +import pytest +import tempfile +import shutil +from pathlib import Path +from unittest.mock import Mock, MagicMock +import torch + + +@pytest.fixture +def temp_dir(): + """Create a temporary directory for test files.""" + temp_path = tempfile.mkdtemp() + yield Path(temp_path) + shutil.rmtree(temp_path) + + +@pytest.fixture +def mock_config(): + """Create a mock configuration object.""" + config = Mock() + config.hidden_size = 768 + config.num_attention_heads = 12 + config.num_hidden_layers = 12 + config.intermediate_size = 3072 + config.vocab_size = 50257 + config.max_position_embeddings = 1024 + config.num_experts = 8 + config.num_experts_per_tok = 2 + return config + + +@pytest.fixture +def mock_model(): + """Create a mock model object.""" + model = MagicMock() + model.config = mock_config() + model.device = torch.device("cpu") + return model + + +@pytest.fixture +def sample_tensor(): + """Create a sample tensor for testing.""" + return torch.randn(2, 10, 768) + + +@pytest.fixture +def sample_input_ids(): + """Create sample input IDs for testing.""" + return torch.randint(0, 50257, (2, 10)) + + +@pytest.fixture +def sample_attention_mask(): + """Create sample attention mask for testing.""" + return torch.ones(2, 10, dtype=torch.long) + + +@pytest.fixture(autouse=True) +def reset_torch_seed(): + """Reset PyTorch random seed for reproducibility.""" + torch.manual_seed(42) + yield + torch.manual_seed(42) + + +@pytest.fixture +def mock_transformers_model(): + """Create a mock transformers model.""" + model = MagicMock() + model.config = mock_config() + model.forward = MagicMock(return_value=MagicMock(logits=torch.randn(2, 10, 50257))) + return model + + +@pytest.fixture +def cleanup_cache(monkeypatch): + """Clean up any cache directories created during tests.""" + cache_dirs = [] + + def track_cache_dir(path): + cache_dirs.append(path) + return path + + monkeypatch.setattr("tempfile.mkdtemp", track_cache_dir) + + yield + + for cache_dir in cache_dirs: + if Path(cache_dir).exists(): + shutil.rmtree(cache_dir) + + +@pytest.fixture +def disable_gpu(): + """Force CPU usage for tests.""" + original_cuda_visible_devices = torch.cuda.CUDA_VISIBLE_DEVICES if hasattr(torch.cuda, 'CUDA_VISIBLE_DEVICES') else None + torch.cuda.CUDA_VISIBLE_DEVICES = "" + yield + if original_cuda_visible_devices is not None: + torch.cuda.CUDA_VISIBLE_DEVICES = original_cuda_visible_devices \ No newline at end of file diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_setup_validation.py b/tests/test_setup_validation.py new file mode 100644 index 0000000..ec5a8c0 --- /dev/null +++ b/tests/test_setup_validation.py @@ -0,0 +1,69 @@ +import pytest +import sys +from pathlib import Path + + +def test_pytest_installed(): + """Test that pytest is properly installed.""" + assert "pytest" in sys.modules or True + + +def test_project_structure(): + """Test that the project structure is set up correctly.""" + project_root = Path(__file__).parent.parent + + # Check main package exists + assert (project_root / "jetmoe").exists() + assert (project_root / "jetmoe" / "__init__.py").exists() + + # Check test structure + assert (project_root / "tests").exists() + assert (project_root / "tests" / "__init__.py").exists() + assert (project_root / "tests" / "conftest.py").exists() + assert (project_root / "tests" / "unit").exists() + assert (project_root / "tests" / "integration").exists() + + +def test_pyproject_toml_exists(): + """Test that pyproject.toml exists and is properly configured.""" + project_root = Path(__file__).parent.parent + pyproject_path = project_root / "pyproject.toml" + + assert pyproject_path.exists() + + # Check content + content = pyproject_path.read_text() + assert "[tool.poetry]" in content + assert "[tool.pytest.ini_options]" in content + assert "[tool.coverage.run]" in content + + +@pytest.mark.unit +def test_markers_work(): + """Test that custom markers are properly configured.""" + assert True + + +@pytest.mark.integration +def test_integration_marker(): + """Test integration marker.""" + assert True + + +@pytest.mark.slow +def test_slow_marker(): + """Test slow marker.""" + assert True + + +def test_fixtures_available(temp_dir, mock_config, sample_tensor): + """Test that fixtures from conftest.py are available.""" + assert temp_dir.exists() + assert hasattr(mock_config, 'hidden_size') + assert sample_tensor.shape == (2, 10, 768) + + +def test_coverage_configured(): + """Test that coverage is properly configured.""" + import coverage + assert coverage.__version__ \ No newline at end of file diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29