feat: modernize codebase for Python 3.9+ with pydantic validation by tcoroller · Pull Request #148 · Novartis/torchsurv

tcoroller · 2026-03-17T19:24:38Z

Drop Python 3.8 support (EOL Oct 2024), require Python >=3.9
Add Python 3.12 classifier
Add pydantic>=2.0 as required dependency
Bump torch requirement to >=2.8
Replace validate_data.py with new pydantic-based validators.py
Add tools/init.py module
Refactor all metrics and loss modules to use new validators
Add hypothesis-based property tests in tests/strategies.py
Add shared test fixtures in tests/conftest.py
Add tests/test_validators.py for new validation logic
Update dev/environment.yml and run-unittests.sh
Update ruff target-version to py39
Add docs/package_overview.md
Add .hypothesis/ to .gitignore

- Drop Python 3.8 support (EOL Oct 2024), require Python >=3.9 - Add Python 3.12 classifier - Add pydantic>=2.0 as required dependency - Bump torch requirement to >=2.8 - Replace validate_data.py with new pydantic-based validators.py - Add tools/__init__.py module - Refactor all metrics and loss modules to use new validators - Add hypothesis-based property tests in tests/strategies.py - Add shared test fixtures in tests/conftest.py - Add tests/test_validators.py for new validation logic - Update dev/environment.yml and run-unittests.sh - Update ruff target-version to py39 - Add docs/package_overview.md - Add .hypothesis/ to .gitignore Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

These packages are required by the modernized codebase but were missing from dev/environment.yml, causing ModuleNotFoundError in CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Remove the optional 'checks' / 'check' boolean parameter from all loss functions, metric classes, and stats functions. Validation via pydantic validators is now always performed (guarded only by the existing torch.jit.is_scripting/is_tracing check for JIT compatibility). Affected: - loss/cox.py: neg_partial_log_likelihood, baseline_survival_function - loss/weibull.py: neg_log_likelihood_weibull - loss/survival.py: neg_log_likelihood, survival_function - metrics/auc.py: Auc.__init__ - metrics/brier_score.py: BrierScore.__init__ - metrics/cindex.py: ConcordanceIndex.__init__ - stats/ipcw.py: get_ipcw - stats/kaplan_meier.py: KaplanMeierEstimator.__call__ Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The tests directory is not a package (no __init__.py), so pytest adds it to sys.path directly. Use 'from strategies import' consistent with how all other test files import from 'utils'. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…ions Replace O(K*N) Python loops with fully vectorized tensor operations for the standard (non-time-varying) Cox path: - H/R list comprehensions → scatter_add (events per time) + reverse cumsum (suffix sum of exp(log_hz) for risk sets), reducing cost from O(K*N) loop-based torch.where calls to O(N) tensor ops - Nested Efron for-loop (k, r) → padded (K, max_ties) tensor: compute all Efron denominator terms in parallel, mask padding with log(1)=0 - Breslow log_denominator stack → log(reverse_cumsum[first_idx]) - _cumulative_baseline_hazard: same reverse-cumsum trick The time-varying extended Cox path retains the existing list-comprehension approach (irregular 2-D indexing prevents general vectorization). Strata loops in the outer functions are kept as-is (structurally required; typically <=5 strata). Numerical equivalence verified against the reference implementation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

1. test_brier_score_simulated_data (flaky): sksurv requires eval times strictly within [test_min, test_max) — filter new_time to exclude test_time.max() before passing to sksurv; skip IBS comparison when time grids differ to avoid integrating over mismatched intervals. 2. test_cox_equivalence (TorchScript): torch.jit.script rejects PEP 604 'X | Y' union syntax — replace 'torch.Tensor | None' with Optional[torch.Tensor] in all three function signatures; add 'from typing import Optional'. Suppress ruff UP045 since Optional is required for TorchScript compatibility. 3. test_training (MPS device): torch.searchsorted is not supported on Apple MPS — add _searchsorted() helper that runs on CPU and moves the result back to the original device; use it everywhere in cox.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Each public class and function now has runnable Examples sections covering: - happy-path usage with exact tensor output - auto-coercion behaviour (e.g. int event → bool, model_type normalisation) - error cases showing the ValidationError / ValueError message Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…rict, ruff clean - validators.py: Pydantic v2 models for all public inputs; TimeVaryingCoxInputs, impute_missing_log_shape in ModelInputs, helpers, base class; log_params/log_hz allow inf (legitimate model outputs) - cox.py: Optional → X | None; TimeVaryingCoxInputs at call sites; fix baseline_survival return type; typed pll_tensor variable - momentum.py: module-level _SurvivalTuple; Callable[..., Any]; deque[Any]; type: ignore[misc] on torch decorator/subclass lines - metrics/auc.py, cindex.py, brier_score.py: __init__ return None, instance attrs typed as X | None, assert guards in private methods, full method signatures - stats/kaplan_meier.py: device: str | None, self.device: Any, return types, Iterator fix - __init__.py files: __all__ + re-exports for root, loss, metrics, stats, tools - conftest.py: collections.abc.Callable - tests: ruff-clean (B017 specific exceptions, E712 bool idioms, UP006/UP035, F401/F841) - pyproject.toml: follow_imports=skip for third-party; pandas-stubs; ruff test exclusion removed - dev/codeqc.sh: now runs ruff format + ruff check + mypy - dev/environment.yml: pandas-stubs via pip Quality gates: mypy strict 0 errors, ruff clean, 87 tests pass, 85% coverage Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Migrate development tooling from conda to uv - Make pyproject.toml the single source of truth for all dependencies - Add missing test deps (lifelines, scikit-survival, lightning, loguru, torchvision, scikit-learn, pandas) to dev dependency group - Add missing doc deps (nbsphinx, sphinxcontrib-bibtex, sphinx-math-dollar, ipython, matplotlib, pandas, lightning) to docs optional extra - Add publish dependency group (build, twine) - Migrate all CI workflows to use astral-sh/setup-uv + uv sync - Update codeqc.yml, docs.yml, build-release.yml, torch-compatibility.yml - Update CLAUDE.md, CONTRIBUTING.md, README.md, docs/devnotes.md - Deprecate dev/environment.yml with notice pointing to uv - Regenerate uv.lock with all new dependency groups Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix codeqc CI: run dev script via uv run The shell script calls ruff/mypy which live in .venv and need uv run to be on PATH. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix mypy type-ignore comments and pre-commit hook - Update type: ignore from [misc] to [untyped-decorator] to match actual mypy error codes from newer mypy + torch versions - Update pre-commit mypy hook: use 'uv run --no-sync mypy' with explicit 'src/' target. --no-sync prevents uv from rebuilding the editable install mid-hook (which confuses pre-commit). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * uv issue * Add missing sphinx-book-theme to docs dependencies The docs conf.py uses sphinx_book_theme but it was only in the legacy environment.yml, not in pyproject.toml docs extras. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add pandoc to CI workflows that build docs pandoc is a system dependency required by nbsphinx to convert notebooks. Previously installed implicitly via conda; now needs explicit apt-get install. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

tcoroller · 2026-03-26T19:36:01Z

PR Review Summary: Modernize Codebase + Migrate to uv

This PR has two major themes: (1) codebase modernization (Python 3.9+, Pydantic v2, typing, mypy strict) and (2) migration from conda to uv as the development toolchain. Here's a guided tour for reviewers.

🔴 High Attention — Review Carefully

1. Pydantic v2 Validators (`src/torchsurv/tools/validators.py` — NEW file, replaces `validate_data.py`)

This is the biggest behavioral change. The old validate_data.py (deleted) used manual checks; the new validators.py uses Pydantic v2 models for input validation throughout the library.

What to check: Validation logic, error messages, edge cases (empty tensors, NaNs, shape mismatches)
Impact: Every loss function and metric now runs inputs through Pydantic validators
Tests: tests/test_validators.py (new)

2. Cox Loss Vectorization (`src/torchsurv/loss/cox.py`)

The Cox partial log-likelihood computation was vectorized — replacing Python loops with tensor operations. This is a correctness-critical performance change.

What to check: Efron/Breslow tie handling, strata support, numerical equivalence with previous implementation
Tests: tests/test_cox.py (heavily modified — compare against lifelines/scikit-survival baselines)

3. Source Code — Type Annotations & `all` (all files under `src/torchsurv/`)

Every module got strict mypy typing and explicit __all__ exports. All __init__.py files now have public API definitions.

What to check: Return types, Optional vs Union usage, that __all__ exports match the intended public API
Quick scan OK for: Most changes are additive annotations on existing signatures

🟡 Medium Attention — Verify Correctness

4. CI Workflows (`.github/workflows/` — all 4 files)

All workflows migrated from conda → uv:

codeqc.yml: setup-miniconda → setup-uv + uv sync + uv run
docs.yml: Same pattern, adds apt-get install pandoc (system dep that conda handled implicitly)
build-release.yml: pip install build twine → uv build + uvx twine check
torch-compatibility.yml: pip → uv sync --override for PyTorch version matrix
What to check: That each workflow installs the right dependency groups (--group dev, --extra docs). The torch-compatibility override pattern (echo 'torch==X.Y.Z' > /tmp/torch-override.txt) is new.

5. `pyproject.toml`

Major changes:

requires-python raised from >=3.8 to >=3.9
pydantic>=2.0 added as a runtime dependency
[dependency-groups] dev expanded with test deps previously only in environment.yml (lifelines, scikit-survival, lightning, loguru, torchvision, etc.)
[project.optional-dependencies] docs expanded (sphinx-book-theme, nbsphinx, sphinxcontrib-bibtex, etc.)
New publish dependency group (build, twine)
What to check: That the dependency list is complete and version bounds are sensible. Pydantic as a runtime dep is a user-facing change.

6. Tests (all files under `tests/`)

New shared fixtures in conftest.py and strategies.py (Hypothesis strategies)
Test files updated for new validator API and type annotations
test_hypothesis.py heavily reworked
What to check: That test coverage is maintained. Spot-check a few test files against previous behavior.

🟢 Low Attention — Quick Scan

7. Documentation Updates

docs/devnotes.md: uv is now the primary dev setup path; conda moved to collapsed legacy section. All command examples updated.
CONTRIBUTING.md: Already was uv-first, removed conda fallback line.
README.md: Dev setup section in <details> updated to use uv sync. End-user install (pip/conda) unchanged.
CLAUDE.md: Internal dev guidelines, commands updated to uv run.
New docs: docs/benchmarks.md, docs/package_overview.md, docs/notebooks/non_medical_applications.ipynb

8. `.pre-commit-config.yaml`

mypy hook changed from entry: mypy to entry: uv run --no-sync mypy with explicit src/ target
This was needed because mypy is inside .venv, not on system PATH

9. `dev/environment.yml`

Added deprecation notice header pointing to pyproject.toml + uv sync
File kept for backward compatibility but no longer maintained

10. Housekeeping

.gitignore additions
CODEOWNERS.md added
communications/ folder added
dev/*.sh scripts: minor cleanups (set -euo pipefail, PYTHONPATH fixes)
uv.lock: New lockfile (auto-generated, ~9k lines — no need to read)