pyproject.toml

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

[project]
name = "marin-root"
version = "0.1.0"
description = "Marin workspace root and experiments"
license = { file = "LICENSE" }
requires-python = ">=3.11"
dependencies = [
    # workspaces
    "marin-iris",
    "marin-fray",
    "marin-haliax",
    "marin-levanter",
    "marin-core",
    "marin-rigging",
    "marin-zephyr",
    # top level deps. try to keep these small.
    "watchdog",
    # `>= X.Y.Z.dev0` opts uv into pre-releases for these native (maturin)
    # packages so the daily nightly wheels (`X.Y.Z.dev<YYYYMMDDhhmm>`) are
    # picked up by `uv lock` automatically. Cut a stable tag
    # (`dupekit-v<X.Y.Z>` / `finelog-v<X.Y.Z>`) when promoting; uv prefers the
    # higher version, so a stable supersedes any earlier dev of the same Z.
    # Published by .github/workflows/dupekit-release-wheels.yaml.
    "marin-dupekit >= 0.1.0.dev0",
    # Published by .github/workflows/finelog-release-wheels.yaml. marin-finelog
    # carries a native in-process server ext (finelog._native), so it ships as
    # a pre-built platform wheel rather than a pure-Python workspace member.
    # `scripts/rust_mode.py dev` swaps in an editable source for local hacking.
    "marin-finelog >= 0.2.0.dev0",
]

[tool.uv]
fork-strategy = "fewest"
# Allow pre-releases globally so `uv lock` resolves marin-dupekit to its
# latest nightly (`X.Y.Z.dev<timestamp>` published by the dupekit release
# workflow). uv has no per-package prerelease config; in practice almost
# none of our other dependencies publish pre-releases on PyPI, so this
# only affects dupekit and the already-pinned `omegaconf>=2.4.0.dev4`
# below.
prerelease = "allow"
override-dependencies = [
    "omegaconf>=2.4.0.dev4",
    "aiofiles>=24.1.0",
    "anthropic>=0.71.0",
    "antlr4-python3-runtime==4.11",
    "python-multipart>=0.0.27",
    "wheel>=0.46.2",
    "datasets>=3.1.0,<5.0.0",
    "equinox>=0.11.10",  # Override vizier's pin for modern JAX compatibility
    "authlib>=1.7.2",
    "cbor2>=6.0.1",
    "click>=8.3.1",
    "gitpython>=3.1.50",
    "gradio>=6.14.0",
    "jupyter-server>=2.18.2",
    "jupyterlab>=4.5.7",
    "litellm>=1.83.14",
    "lxml>=6.1.0",
    "mistune>=3.2.1",
    "multipart>=1.3.1",
    "nltk>=3.9.4",
    "notebook>=7.5.6",
    "orjson>=3.11.9",
    "pillow>=12.2.0",
    "pyasn1>=0.6.3",
    "python-dotenv>=1.1.1",
    "ray>=2.55.1",
    "rich>=14.1.0",
    "setuptools>=82.0.1",
    "tiktoken>=0.12.0",
    "tornado>=6.5.5",
]
conflicts = [
    # Temporary cross-package conflicts for the GPU JAX 0.10 / non-GPU JAX 0.9 split.
    # Remove these once CPU, TPU, vLLM, and Fray TPU test paths can share the GPU JAX line.
    [
        { package = "marin-levanter", extra = "gpu" },
        { package = "marin-core", extra = "cpu" },
    ],
    [
        { package = "marin-levanter", extra = "gpu" },
        { package = "marin-core", extra = "tpu" },
    ],
    [
        { package = "marin-levanter", extra = "gpu" },
        { package = "marin-core", extra = "vllm" },
    ],
    [
        { package = "marin-core", extra = "gpu" },
        { package = "marin-levanter", extra = "tpu" },
    ],
    [
        { package = "marin-core", extra = "gpu" },
        { package = "marin-fray", group = "fray_tpu_test" },
    ],
    [
        { package = "marin-levanter", extra = "gpu" },
        { package = "marin-fray", group = "fray_tpu_test" },
    ],
]

[tool.uv.workspace]
# marin-finelog is intentionally NOT a member: it is a native (maturin) package
# consumed as a pre-built platform wheel (see [project] dependencies and the
# RUST-DEV SOURCES block below), mirroring marin-dupekit.
members = [
    "lib/iris",
    "lib/fray",
    "lib/haliax",
    "lib/levanter",
    "lib/marin",
    "lib/rigging",
    "lib/zephyr",
]

[tool.uv.sources]
marin-iris = { workspace = true }
marin-fray = { workspace = true }
marin-haliax = { workspace = true }
marin-levanter = { workspace = true }
marin-core = { workspace = true }
marin-rigging = { workspace = true }
marin-zephyr = { workspace = true }
# NOTE: harbor is a pinned git dependency (not a workspace member).
# harbor lives in marin-community/harbor as a pinned git dependency.
harbor = { git = "https://github.com/marin-community/harbor.git", rev = "354692d9c0eab497b05f266aa0dff30e2a238d2e" }
# ### BEGIN RUST-DEV SOURCES ###
# ### END RUST-DEV SOURCES ###


[tool.black]
line-length = 121
target-version = ["py310"]
preview = true

# Note :: Grow more strict over time!
extend-exclude = """
(
    scripts/
)
"""

[tool.ruff]
line-length = 121
target-version = "py310"

# Note :: Grow more strict over time!
# Most of `scripts/**` is excluded until cleaned up, but `scripts/workflows/**`
# is CI infrastructure and stays under ruff.
extend-exclude = ["scripts/"]
extend-include = ["scripts/workflows/**/*.py"]

[tool.ruff.lint]
# A,B,E,F,I,NPY,RUF,UP,W are the established families. The individually-listed
# codes below mechanically enforce CodeQL "Code quality" maintainability/reliability
# checks locally so they are caught at pre-commit time rather than only post-merge:
#   N805    -> "First parameter of a method is not named 'self'"
#   PLR0124 -> "Comparison of identical values" / "Redundant comparison"
#   PLR1722 -> "Use of exit() or quit()" (use sys.exit)
#   RET502  -> "Explicit returns mixed with implicit returns" (bare `return` value path)
#   RET503  -> "Explicit returns mixed with implicit returns" (missing explicit return)
#   SIM115  -> "File is not always closed" (open() without a context manager)
#   PLC0415 -> "import should be at the top-level of a file" (no function-local
#              imports except to break a circular import or guard an optional dep)
# B018 (already via "B") and F401/F841/F811 (via "F") cover the rest of the
# actionable CodeQL maintainability set; B018 correctly ignores `...`/docstrings.
select = [
    "A", "B", "E", "F", "I", "NPY", "RUF", "UP", "W",
    "N805", "PLR0124", "PLR1722", "RET502", "RET503", "SIM115", "PLC0415",
    # Local-import hygiene (issue #6219):
    #   TID252 -> "relative-imports" (prefer absolute over parent-relative imports)
    #   ICN001 -> "unconventional-import-alias" (e.g. `import numpy as np`)
    # lib/haliax and lib/levanter carry their own [tool.ruff.lint] configs; these
    # two rules are mirrored there via extend-select so they apply repo-wide.
    "TID252", "ICN001",
]
ignore = ["F722", "B008", "UP015", "A005", "E741"]

[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["E402", "F401"]
# Notebook cells legitimately interleave imports with computation as they are
# executed top-to-bottom; PLC0415's top-of-file rule does not apply to cells.
"*.ipynb" = ["PLC0415"]
# Bloom filter tests deliberately compare an object with itself to exercise its
# reflexive comparison operators (==, <=, >=, ...), so PLR0124 is expected here.
"lib/dupekit/tests/test_bloom.py" = ["PLR0124"]

[tool.ruff.lint.isort]
known-first-party = ["experiments"]
known-third-party = ["wandb"]


[tool.mypy]
python_version = "3.10"

# Note: Grow more strict over time!
ignore_missing_imports = true
exclude = ["marin/", "scripts/"]

# Pyrefly type checker configuration
[tool.pyrefly]
# Worktrees often live under hidden parent directories (e.g. `.codex/worktrees/*`).
# Disable pyrefly's built-in exclude heuristics so hidden ancestors do not exclude
# the entire project when running from those paths.
disable-project-excludes-heuristics = true

project-includes = [
    "lib/marin/src/**/*.py",
    "lib/levanter/src/**/*.py",
    "lib/haliax/src/**/*.py",
    "lib/fray/src/**/*.py",
    "lib/iris/src/**/*.py",
    "lib/rigging/src/**/*.py",
    "lib/zephyr/src/**/*.py",
    "scripts/workflows",
]

# Explicitly tell Pyrefly where our editable packages live so it resolves imports
# against the library sources instead of the top-level `src` directory, which only
# contains a handful of legacy modules.
search-path = [
    "lib/marin/src",
    "lib/levanter/src",
    "lib/haliax/src",
    "lib/rigging/src",
    "lib/zephyr/src",
    "lib/fray/src",
    "lib/iris/src",
]
disable-search-path-heuristics = true

# Type-check against real installed third-party types. CI's marin-lint job syncs
# this venv (uv sync --all-packages --extra=cpu --no-default-groups); locally it is
# your worktree .venv. Relative path keeps it worktree-local.
python-interpreter-path = ".venv/bin/python"
use-ignore-files = false

# Parse jaxtyping/native tensor shape annotations (e.g. Float[Array, "batch seq"])
# as shape specifications instead of forward-reference type expressions. Without this,
# Pyrefly reads the dimension tokens inside the shape string ("B", "N", "groups", ...)
# as undefined forward references and emits ~141 false `unknown-name` errors across
# levanter/grug and levanter/kernels. See docs/plans/pyrefly-tensor-shapes.md.
tensor-shapes = true

# Modules treated as Any when absent on disk. Pyrefly still uses real stubs when
# the dep is installed (so a deleted/misspelled internal module still trips
# missing-import). Covers deps no lint profile installs (vllm/lm_eval, resiliparse)
# and compute deps a local venv commonly skips (torch, sympy). List both `pkg` and
# `pkg.*`.
ignore-missing-imports = [
    "harbor",
    "harbor.*",
    "jax_triton",
    "jax_triton.*",
    "jaxlib",
    "jaxlib.*",
    "kubernetes",
    "kubernetes.*",
    "librosa",
    "librosa.*",
    "lm_eval",
    "lm_eval.*",
    "math_verify",
    "math_verify.*",
    "memray",
    "memray.*",
    "prompt_toolkit",
    "prompt_toolkit.*",
    "pylatexenc",
    "pylatexenc.*",
    "resiliparse",
    "resiliparse.*",
    "sympy",
    "sympy.*",
    "tensorboardX",
    "tensorboardX.*",
    "torch",
    "torch.*",
    "tpu_inference",
    "tpu_inference.*",
    "trackio",
    "trackio.*",
    "triton",
    "triton.*",
    "verifiers",
    "verifiers.*",
    "vllm",
    "vllm.*",
    "xprof",
    "xprof.*",
]

# Exclude non-production code from type checking
project-excludes = [
    "experiments/**",
    "scripts/canary/**",
    "scripts/datakit/**",
    "scripts/debug/**",
    "scripts/ferries/**",
    "scripts/gcs-reorganization/**",
    "scripts/gpu_eval/**",
    "scripts/iris/**",
    "scripts/migrations/**",
    "scripts/ops/**",
    "scripts/pm/**",
    "scripts/ray/**",
    "scripts/speedrun/**",
    "scripts/storage/**",
    "scripts/training/**",
    "scripts/*.py",
    "tests/**",
    "examples/**",           # Example code doesn't need strict typing
    "lib/**/crawl/**",       # Crawl scripts have library typing issues with smart_open
    "lib/iris/src/iris/rpc/*_pb2*",  # Generated protobuf files
    "lib/iris/src/iris/rpc/*_connect.py",  # Generated Connect RPC files
    # Keep key excludes that pyrefly normally adds heuristically.
    "**/node_modules",
    "**/__pycache__",
    "**/venv/**",
]

# Disable specific error codes that are primarily noise from missing type stubs.
# Enable gradually as the codebase improves.
[tool.pyrefly.errors]
# Untyped-library attribute access (jax/flax/equinox); sampling found ~no real bugs.
missing-attribute = false

# Real undefined-name bugs. The former jaxtyping shape-string false positives are now
# handled natively by tensor-shapes = true above; residual false positives are
# suppressed in .pyrefly-baseline.json.
unknown-name = true

# Complex overload resolution (stdlib/library overloads); sampling found ~no real bugs.
no-matching-overload = false

# Mostly third-party stub noise (smart_open, Flax, jax/numpy arrays), with a handful of
# real callsite mismatches buried in the volume.
bad-argument-type = false

# Real "code that can be better written" bugs mixed with control-flow false positives
# pyrefly cannot prove safe.
unbound-name = false

# Kept enabled (high signal): bad-override, bad-assignment, bad-return,
# unsupported-operation, not-callable. Pre-existing instances are suppressed via
# .pyrefly-baseline.json.

# jit_scheduler.py annotates eqx.Module fields with single-axis haxtyping shapes,
# e.g. `ht.i32[NamedArray, "position"]`. tensor-shapes = true parses multi-token
# shape strings ("seq page") correctly, but reads a lone bare-identifier axis
# ("position", "seq") as a forward reference, yielding false unknown-name errors.
# Scope the disable to this one file rather than scattering 17 inline ignores.
[[tool.pyrefly.sub_config]]
matches = "lib/levanter/src/levanter/inference/jit_scheduler.py"
[tool.pyrefly.sub_config.errors]
unknown-name = false

[tool.hatch.metadata]
allow-direct-references = true

[tool.hatch.build.targets.wheel]
packages = ["experiments"]

[tool.hatch.build.targets.sdist]
packages = ["experiments"]

[tool.pytest.ini_options]
timeout = 60
filterwarnings = ["ignore::DeprecationWarning"]
log_format = "%(asctime)s %(levelname)s %(message)s"
log_date_format = "%Y-%m-%d %H:%M:%S"
markers = [
    "slow: mark tests as slow for CI - use -m 'not slow'",
    "tpu_ci: mark tests that require a TPU",
    "integration: mark tests as integration tests that require cluster infrastructure",
    "data_integration: mark tests that require external data fixtures (e.g. HuggingFace datasets)",
    "requires_cluster: mark tests that need a running iris cluster",
]
testpaths = ["tests", "experiments"]

# Don't run TPU, slow, integration, data_integration, or cluster-bound tests by default.
# Integration tests require external infrastructure (Iris cluster, GCS,
# gated HF repos, etc.); data_integration tests fetch fixtures from HF;
# requires_cluster tests need a running iris cluster (local or remote).
# All are run from dedicated CI workflows.
addopts = "--session-timeout=600 -m 'not tpu_ci and not slow and not integration and not data_integration and not requires_cluster'"