pyproject.toml

[project]
authors = [{ name = "Sanggeun Park", email = "psk6950@snu.ac.kr" }]
name = "miniworld"
requires-python = ">= 3.10"
version = "2.0.0.dev0"
dependencies = [
  "torch>=2.8.0,<3",
  "wandb>=0.21.0",
  "omegaconf>=2.3.0,<3",
  "numpy>=2.2.5,<3",
  "einops>=0.8.1,<0.9",
  "jaxtyping>=0.3.2,<0.4",
  "beartype>=0.21.0,<0.22",
  "scipy>=1.15.3,<2",
  "lightning>=2.5.6,<3",
  "biomol>=1.0.0",
  "cuequivariance-torch>=0.12,<0.13",
  "pyyaml>=6,<7",
  "rdkit>=2024.3,<2027",
  # formerly provided through team-gm
  "pydantic>=2,<3",
  "lmdb>=1.7,<2",
  "pyarrow>=18,<19",
  "pandas>=2.3,<3",
  "biopython>=1.85,<2",
  "hydra-core>=1.3,<2",
  "tqdm>=4.67,<5",
  "matplotlib>=3.10,<4",
  "click>=8.3,<9",
  "ninja>=1.13,<2",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.setuptools]
include-package-data = true


[tool.pixi.workspace]
channels = ["conda-forge"]
platforms = ["linux-64"]

[tool.pixi.feature.cuda13.system-requirements]
cuda = "13"

[tool.pixi.feature.cuda13.pypi-dependencies]
torch = "==2.13.0"
cuequivariance-ops-torch-cu13 = "==0.12.*"
# FlashAttention-4 (CuTeDSL) for the SWA atom-attention flash backend.
# Supports Hopper (H100 / sm_90) and Blackwell (B200 / sm_100 / sm_110). FA2/3
# lack sm_100 support, so on Blackwell the SWA path otherwise falls back to
# O(S^2) SDPA.
# Prerelease CuTeDSL JIT wheel — kept OUT of the hard solve (don't perturb it);
# install manually into the env:
#   uv pip install --python .pixi/envs/default/bin/python --prerelease=allow "flash-attn-4[cu13]"

[tool.pixi.feature.cuda13.pypi-options]
extra-index-urls = ["https://download.pytorch.org/whl/cu130"]

# miniworld-engine v2.2.0 runtime: torch 2.13.0+cu129 / cuequivariance 0.12. CUDA 12.9 is
# the newest the cluster's 575 driver runs. The environment keeps its historical `cu128`
# name because run scripts address `.pixi/envs/cu128` and `-e cu128`.
[tool.pixi.feature.cuda129.system-requirements]
cuda = "12.9"

[tool.pixi.feature.cuda129.pypi-dependencies]
torch = "==2.13.0"
cuequivariance-ops-torch-cu12 = "==0.12.*"
# FlashAttention-4 (CuTeDSL), pinned: miniworld-engine 2.2.0 unpacks the four values flash_attn.cute returns in 4.0.0b32; a later
# beta returns more and the SWA atom attention fails with "too many values to unpack". The exact pre-release pin lets the
# solver take it; the cu12 extra picks the CUTLASS DSL built for CUDA 12.
flash-attn-4 = { version = "==4.0.0b32", extras = ["cu12"] }

# The H100 TriMul / TriangleAttention kernels JIT-compile against CUTLASS 4.2 headers: `pixi run -e cu128 setup-cutlass` fetches
# them once into ./.cutlass, and the environment points CUTLASS_PATH there.
[tool.pixi.feature.cuda129.tasks]
setup-cutlass = "test -f .cutlass/include/cute/tensor.hpp || git clone --depth 1 --branch v4.2.1 https://github.com/NVIDIA/cutlass.git .cutlass"
# Build every engine CUDA extension the production models use (all shapes, inference + training) and verify that a new
# process builds nothing, so the first model call never waits on nvcc. Run on a GPU node of each architecture after creating
# the environment or bumping the engine pin: `pixi run -e cu128 engine-prebuild [--jobs 8]` (H100: scripts/slurm/engine_prebuild_h100.sbatch).
engine-prebuild = "python scripts/prebuild_engine.py --mode all"

[tool.pixi.feature.cuda129.activation.env]
CUTLASS_PATH = "$PIXI_PROJECT_ROOT/.cutlass"

[tool.pixi.feature.cuda129.pypi-options]
extra-index-urls = ["https://download.pytorch.org/whl/cu129"]

[tool.pixi.environments]
default = { features = ["cuda13"] }
cu128 = { features = ["cuda129"] }

[tool.pixi.tasks]
# Consolidated main contains the reviewed fixes; do not replay the old 1.x patches.
engine-setup = "python scripts/check_engine_install.py"

[tool.pixi.pypi-dependencies]
miniworld = { path = ".", editable = true }
# miniworld-engine is an upstream package. The blocks / layers / diffusion / training loop that used to live in team-gm are
# part of this package (miniworld.modules, miniworld.diffusion, miniworld.training); there is no team-gm dependency.
# Pinned to engine main; see docs/repositories.md.
miniworld-engine = { git = "ssh://git@github.com/SanggeunParrk/miniworld-engine.git", rev = "545b68508f4b7a9bfbf742633549079a723a427b" }

[tool.pixi.dependencies]
pytest = ">=8.3.5,<9"
python = "~=3.10.16"
ruff = ">=0.11.6,<0.12"
pyright = ">=1.1.406,<1.2"
ipykernel = ">=7.1.0,<8"
matplotlib = ">=3.10.7,<4"
pandas = ">=2.3.3,<3"
pyarrow = ">=18,<19"
tqdm = ">=4.67.1,<5"
click = ">=8.3.1,<9"
biopython = ">=1.87,<2"

[tool.ruff]
line-length = 89
target-version = "py310"
extend-exclude = ["libs", "docs", "runs", "benchmarks"]

[tool.ruff.lint]
# The gate: correctness (E/F/B), import order (I), modern syntax (UP), simplifications (SIM, C4, PIE).
# The broader `ALL` set was never enforceable here; the ignores below are known debt, not decisions --
# remove an entry when its violations are fixed.
select = ["E", "F", "W", "I", "UP", "B", "SIM", "C4", "PIE"]
ignore = [
  "E501",    # the formatter owns line length
  "E731",
  "E741",
  "F722",    # incompatible with jaxtyping shape strings
  "B905",
  "UP006",
  "UP007",
  "UP038",   # deprecated upstream
  "B023",    # debt: closures over loop variables in benchmark scripts
  "C408",    # debt: dict(...) / list() calls
  "UP031",   # debt: printf-style formatting
  "B007",
  "F841",
  "E402",    # scripts set sys.path / env before importing
  "SIM115",
  "SIM113",
  "SIM102",
  "C401",
  "B904",
  "PIE810",
  "SIM105",
  "SIM108",
  "SIM117",
]

[tool.ruff.lint.per-file-ignores]
"scripts/bench_token_dit.py" = ["F821"]
"__init__.py" = ["F401"]
"scripts/fix_cif.py" = ["W291"]  # trailing whitespace is inside string literals

[tool.pytest.ini_options]
testpaths = ["tests"]
pythonpath = ["tests"]
addopts = "-ra"

[tool.pyright]
include = ["src/miniworld"]
exclude = ["libs", "scratchpad_logs", "docs", "runs", "benchmarks", "tests/artifacts",
  # evaluation helpers that need joblib and a vendored DockQ; nothing imports them
  "src/miniworld/utils/structure/eval.py", "src/miniworld/utils/structure/eval2.py"]
enableReachabilityAnalysis = false
# Known debt, still reported as warnings: jaxtyping shapes and flag-guarded branches confuse the checker.
# New categories of error fail the run.
reportArgumentType = "warning"
reportAttributeAccessIssue = "warning"
reportCallIssue = "warning"
reportOptionalMemberAccess = "warning"
reportOptionalSubscript = "warning"
reportOptionalOperand = "warning"
reportOptionalIterable = "warning"
reportReturnType = "warning"
reportIndexIssue = "warning"
reportIncompatibleMethodOverride = "warning"
reportIncompatibleVariableOverride = "warning"
reportAssignmentType = "warning"
reportPossiblyUnboundVariable = "warning"
reportOptionalCall = "warning"
reportFunctionMemberAccess = "warning"
reportGeneralTypeIssues = "warning"
reportOperatorIssue = "warning"

[tool.hatch.build.targets.wheel]
packages = [
  "src/miniworld",
]