earth2studio/tox.ini at main · NVIDIA/earth2studio · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
[tox]
requires =
    tox >= 4.0
    tox-uv >= 1.28.0
envlist =
    test
    test-data
    test-perturb
    test-stats
    test-px-models
    test-dx-models
    test-da-models
    test-px-models-ace2

[testenv]
# Base configuration for all environments
# This installs the local package in editable mode
# plus any extras or groups defined below.
package = editable
basepython = python3.13
passenv =
    EARTH2STUDIO_CACHE
    EARTH2STUDIO_DATA_CACHE
    EARTH2STUDIO_MODEL_CACHE
    EARTH2STUDIO_PACKAGE_TIMEOUT
    TESTMON_DATAFILE
    UV_CACHE_DIR
    S3FS_ENDPOINT
    S3FS_KEY
    S3FS_SECRET
    HF_TOKEN
    CDSAPI_KEY
    CDSAPI_URL
    EUMETSAT_CONSUMER_KEY
    EUMETSAT_CONSUMER_SECRET
    EARTHMOVER_API_KEY
    EARTHMOVER_ORGANIZATION
    MAX_JOBS
setenv =
    # https://docs.jax.dev/en/latest/gpu_memory_allocation.html
    XLA_PYTHON_CLIENT_ALLOCATOR = platform
    XLA_PYTHON_CLIENT_PREALLOCATE = false
    # Performance optimization for Atlas model
    TORCH_ALLOW_TF32_CUBLAS_OVERRIDE = 1
    # Unique coverage file per tox env so parallel CI jobs don't overwrite each other
    COVERAGE_FILE = .coverage.{envname}

[testenv:test]
runner = uv-venv-lock-runner
description = Run tests that need no extra dependencies
commands =
    pytest {posargs:-s --cov --cov-append --slow --testmon} \
        test/io \
        test/run \
        test/utils/test_coords.py \
        test/utils/test_imports.py \
        test/models/test_auto_models.py \
        test/models/test_batch.py

[testenv:test-serve]
runner = uv-venv-lock-runner
description = Run tests for the serve
extras = serve
commands =
    pytest {posargs:-s --cov --cov-append --slow --testmon} \
        test/serve

[testenv:test-data]
runner = uv-venv-lock-runner
description = Run tests for data sources
extras = data,cbottle
commands =
    pytest {posargs: --cov --cov-append --slow --package --testmon} \
        test/data \
        test/lexicon

[testenv:test-perturb]
runner = uv-venv-lock-runner
description = Run tests for perturbation methods
extras = perturbation
commands =
    pytest {posargs:-s --cov --cov-append --slow --testmon} \
        test/perturbation

[testenv:test-stats]
runner = uv-venv-lock-runner
description = Run tests for statistical methods
extras = statistics,utils
commands =
    pytest {posargs:-s --cov --cov-append --slow --testmon} \
        test/statistics \
        test/utils/test_interp.py

# =====================
# TODO: NGeneva figure out how to add packages back
# =====================
[testenv:test-px-models]
runner = uv-venv-lock-runner
description = Run tests for px models using all env
extras = all
commands =
    pytest {posargs:-s --cov --cov-append --slow --testmon} \
        test/models/px

[testenv:test-dx-models]
runner = uv-venv-lock-runner
description = Run tests for dx models using all env
extras = all
commands =
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/dx

[testenv:test-da-models]
runner = uv-venv-lock-runner
description = Run tests for da models using all env
extras = all
commands =
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/da

# =====================
# Conflict environments
# =====================
[testenv:test-px-models-conflicts]
runner = uv-venv-lock-runner
description = Run tests for px models with conflicts
commands =
    uv sync --extra ace2
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_ace2.py

    uv sync --extra aifs
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_aifs.py

	uv sync --extra aifs2
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_aifs2.py

	uv sync --extra aifsens
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_aifsens.py

	uv sync --extra aifs2ens
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_aifs2ens.py

    uv sync --extra stormcast-conus
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/px/test_stormcastconus.py

    uv sync --extra da-cosmo
    pytest {posargs:-s --cov --cov-append --slow --package --testmon} \
        test/models/dx/test_corrdiff_cosmo_era5.py \
        test/models/da/test_da_sda_corrdiff_cosmo_era5.py