pyproject.toml

[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"

[tool.setuptools.packages.find]
include = ["skyrl*"]

[project]
name = "skyrl"
dynamic = ["version"]
description = "Unified API for training and inference"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "datasets>=4.0.0",
    "pillow>=11.3.0",
    "rich>=14.1.0",
    "safetensors>=0.6.2",
    "tokenizers>=0.21.2",
    "transformers>=4.56.1,<5",
    "typer>=0.17.4",
    # "wandb>=0.22.0",
    "peft",
    "hf_transfer",
    "cloudpathlib>=0.23.0",
]

[project.optional-dependencies]
gpu = [
    "jax[cuda12]>=0.7.2; sys_platform == 'linux'",
]

tpu = [
    "jax[tpu]>=0.7.2; sys_platform == 'linux'",
]

tinker = [
    "tinker>=0.3.0",
    "fastapi[standard]",
    "sqlmodel",
    "sqlalchemy[asyncio]",
    "aiosqlite",
    "asyncpg",
    "psycopg2-binary",
]

aws = [
    "cloudpathlib[s3]",
]

gcp = [
    "cloudpathlib[gs]",
]

azure = [
    "cloudpathlib[azure]",
]

# The extras "jax", "fsdp", and "megatron" are the dependencies the
# engine needs for --backend="jax", --backend="fsdp", and --backend="megatron",
# respectively.

jax = [
    "jax>=0.8,<1.0",
    "flax>=0.12.2",
    "optax>=0.2.5",
]

skyrl-train = [
    "loguru",
    "tqdm",
    "ninja",
    "tensorboard",
    "func_timeout",
    "transformers>=4.51.0",
    "hydra-core==1.3.2",
    "accelerate",
    "torchdata",
    "omegaconf",
    "ray==2.51.1",
    "peft",
    "debugpy==1.8.0",
    "hf_transfer",
    "wandb",
    "datasets>=4.0.0",
    "tensordict",
    "jaxtyping",
    "skyrl-gym",
    "flash-attn; sys_platform == 'linux'",
    "polars",
    "s3fs",
    "fastapi",
    "uvicorn",
    "pybind11",
    "setuptools",
]

fsdp = [
    "skyrl[skyrl-train]",
    "vllm==0.16.0; sys_platform == 'linux'",
    "flash-attn==2.8.3; sys_platform == 'linux'",
    "torch==2.9.1; sys_platform == 'linux'",
    "flashinfer-python==0.6.3; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "flashinfer-jit-cache==0.6.3; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torchvision; sys_platform == 'linux'",
]

megatron = [
    "skyrl[skyrl-train]",
    "transformer-engine[pytorch]==2.10.0; sys_platform == 'linux'",
    "flash-attn==2.8.1; sys_platform == 'linux'",
    "vllm==0.16.0; sys_platform == 'linux'",
    "torch==2.9.1; sys_platform == 'linux'",
    "flashinfer-python==0.6.3; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torchvision; sys_platform == 'linux'",
    "megatron-bridge; sys_platform == 'linux'",
    "megatron-core==0.16.0; sys_platform == 'linux'",
    "flashinfer-jit-cache==0.6.3; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "nvidia-modelopt; sys_platform == 'linux'",
]

flashrl = [
    "skyrl[skyrl-train]",
    # NOTE: Custom vLLM wheel must be installed separately.
    # See examples/flash_rl/README.md for installation instructions.
    "flash-attn==2.8.3; sys_platform == 'linux'",
    "torch==2.7.0; sys_platform == 'linux'",
    "flashinfer-python; sys_platform == 'linux'",
    "torchvision; sys_platform == 'linux'",
]
miniswe = [
    "skyrl[skyrl-train]",
    # NOTE (sumanthrh): Needs to be a commit after https://github.com/SWE-agent/mini-swe-agent/commit/4f5d445e99d13b5482478c23508bf2fbf7c0670c
    "mini-swe-agent>=1.12.0",
    "litellm",
]
harbor = [
    "harbor; python_version >= '3.12'",
]

dev = [
    "mkdocs",
    "mkdocs-material",
    "mkdocstrings[python]>=0.24.0",
    "pymdown-extensions>=10.7",
    "pytest",
    "pytest-forked",
    "pytest-asyncio",
    "pre-commit",
    "litellm",
    "torch",
    "ty",
    "cloudpathlib[s3]",
    "alembic",
    "griffe2md",
]

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

[tool.setuptools.dynamic]
version = {attr = "skyrl.__version__"}

[project.entry-points."vllm.general_plugins"]
skyrl_ipc_engine = "skyrl.backends.skyrl_train.weight_sync.vllm_ipc_engine:register_ipc_engine"

[project.scripts]
# The following is for supporting the skyrl-train dependency

[tool.uv]
# Resolve for both Linux (production) and macOS (dev)
required-environments = [
    "sys_platform == 'linux'",
    "sys_platform == 'darwin' and platform_machine == 'arm64'",
]

constraint-dependencies = [
    "flashinfer-jit-cache==0.6.3",
]
# each backend should have separate dependencies that can potentially clash
# megatron also clashes with the jax dependency from gpu and tpu extras
conflicts = [
    [
        { extra = "jax" },
        { extra = "megatron" },
        { extra = "fsdp" },
        { extra = "flashrl" },
    ],
    [
        { extra = "megatron" },
        { extra = "gpu" },
        { extra = "tpu" },
        { extra = "flashrl" },
        { extra = "miniswe" },
    ]
]
# disable build isolation for megatron related dependencies
no-build-isolation-package = [
    "transformer-engine-torch",
    "transformer-engine",
    "nv-grouped-gemm",
]
# override unnecessary dependencies and pin versions to override Megatron-Bridge
# unpinned dependencies.
override-dependencies = [
    "nvidia-resiliency-ext; sys_platform == 'never'",
    "mamba-ssm; sys_platform == 'never'",
    "causal-conv1d; sys_platform == 'never'",
    "transformer-engine[pytorch]==2.10.0; sys_platform == 'linux'",
    "megatron-core==0.16.0; sys_platform == 'linux'",
    "ml_dtypes>=0.5.0; sys_platform == 'linux'",
]

[tool.uv.extra-build-dependencies]
flash-attn = [{requirement = "torch", match-runtime = true}]
transformer-engine = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja"]
transformer-engine-torch = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja"]

[tool.uv.extra-build-variables]
flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE"}

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "jax-tpu"
url = "https://storage.googleapis.com/jax-releases/libtpu_releases.html"

[[tool.uv.index]]
name = "flashinfer-cu128"
url = "https://flashinfer.ai/whl/cu128"
explicit = true

[tool.uv.sources]
skyrl-gym = { path = "./skyrl-gym", editable = true }
# flashinfer wheels are only available from the custom cu128 index
flashinfer-jit-cache = { index = "flashinfer-cu128", marker = "sys_platform == 'linux'" }
# Use CUDA torch on Linux, CPU torch on macOS (must match skyrl-train config)
torch = [
    { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
    { index = "pytorch-cpu", marker = "sys_platform == 'darwin'" },
]
torchvision = [
    { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
    { index = "pytorch-cpu", marker = "sys_platform == 'darwin'" },
]
# pin megatron bridge commit to fix for MoE + LoRA merging. Update this when an official release is cut
megatron-bridge = {git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge", rev = "02b5fccab5e5b21856d36c2e357839e0123b4b8f", marker = "sys_platform == 'linux'"}
harbor = { git = "https://github.com/laude-institute/harbor", rev = "8c040e1bb010201fd3c75bee3dede2407b9f57cd" }

[tool.black]
line-length = 120
include = '\.pyi?$'
extend-exclude = '''
/(
  # directories
  \.eggs
  | \.git
  | \.hg
  | \.mypy_cache
  | \.tox
  | \.venv
  | build
  | dist
)/
'''

[tool.flake8]
max-line-length = 120
max-doc-length = 120
extend-ignore = [
    # Default ignored errors by flake8
    "E121", "E123", "E126", "E226", "E24", "E704",
    # F401 module imported but unused
    "F401",
    # E203 whitespace before ':' (conflict with black)
    "E203",
    # E231 missing whitespace after ',' (conflict with black)
    "E231",
    # E501 line too long (conflict with black)
    "E501",
    # E741 do not use variables named 'l', 'O', or 'I'
    "E741",
    # W503 line break before binary operator (conflict with black)
    "W503",
    # W504 line break after binary operator (conflict with black)
    "W504",
    # W505 doc line too long (conflict with black)
    "W505",
    # W605 invalid escape sequence 'x' (conflict with latex within docs)
    "W605",
]

[tool.ruff]
exclude = [
    "skyrl-agent/",
    "examples/",
]

[tool.ruff.lint]
extend-select = ["I"]
ignore = [
    "F722" # Syntax error in annotation - ignored because this doesn't play well with jaxtyping
]

[tool.ruff.lint.isort]
known-first-party = ["skyrl", "skyrl_gym"]